"""E4 (core) early decision: can the e-nose decide before the sensors have settled?

    python experiments/e4_early_decision.py data/demo_calibration

No new measurements are needed: the same raw.csv is cut after t seconds of the sample
phase. The "steady value" is then the mean of the last 15 s before the cut (one full
BME688 heater cycle). Models are trained on the full 240 s measurements of the training
days and tested on the cut measurements of the held-out days (--test-days, default last 2).

Why the fly should be fast: while a smell is still building up, its STRENGTH keeps
changing but the RATIO between channels is already there. The fly's gain control only
looks at the ratio. Hypothesis H3: at 30-60 s 蕈狀體 is clearly better than the classic
methods on raw features, so a measurement could be shortened from 4 minutes to ~1.
Watch out: the three sensors react at different speeds (SCD40 updates every 5 s, the
BME688 heater cycle takes 15 s), so the ratio at 30 s may still be off on real hardware.
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.base import clone

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from enose import features as feature_module  # noqa: E402
from enose.common import output_dir, plt, save_figure  # noqa: E402
from enose.features import build_feature_table, odor_columns  # noqa: E402
from enose.methods import make_methods  # noqa: E402

CUTS: tuple[int, ...] = (30, 45, 60, 90, 120, 240)
EARLY_WINDOW_SECONDS: int = 15
SHOWN: tuple[str, ...] = ("蕈狀體", "kNN（k=1）", "邏輯回歸", "kNN＋果蠅正規化", "邏輯回歸＋果蠅正規化")


def cut_table(raw: pd.DataFrame, labels: pd.DataFrame, seconds: int, heater_steps: int) -> pd.DataFrame:
    kept = raw[(raw["phase"] != "sample") | (raw["phase_second"] < seconds)]
    full_window = feature_module.WINDOW_SECONDS
    feature_module.WINDOW_SECONDS = EARLY_WINDOW_SECONDS if seconds < 240 else full_window
    try:
        return build_feature_table(kept, labels, heater_steps)
    finally:
        feature_module.WINDOW_SECONDS = full_window


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("folder", type=Path)
    parser.add_argument("--heater-steps", type=int, default=5)
    parser.add_argument("--test-days", type=int, default=2)
    arguments = parser.parse_args()
    raw = pd.read_csv(arguments.folder / "raw.csv")
    labels = pd.read_csv(arguments.folder / "labels.csv")
    if "rh_level" in labels:
        labels = labels[labels["rh_level"] == "mid"]
    full = pd.read_csv(arguments.folder / "features.csv")
    full = full[full["sample_id"].isin(labels["sample_id"])]
    names = odor_columns(full)
    test_days = sorted(labels["date"].unique())[-arguments.test_days:]
    train = full[~full["date"].isin(test_days)].dropna(subset=names)
    methods = {name: method for name, method in make_methods(names, full["label"].nunique()).items() if name in SHOWN}
    fitted = {name: clone(method).fit(train[names].to_numpy(), train["label"].to_numpy())
              for name, method in methods.items()}

    rows = []
    test_labels = labels[labels["date"].isin(test_days)]
    for seconds in CUTS:
        table = cut_table(raw[raw["sample_id"].isin(test_labels["sample_id"])], test_labels, seconds,
                          arguments.heater_steps)
        features = np.nan_to_num(table[names].to_numpy(dtype=float))
        for name, model in fitted.items():
            rows.append({"seconds": seconds, "method": name,
                         "accuracy": float((model.predict(features) == table["label"].to_numpy()).mean())})
    result = pd.DataFrame(rows)
    results = output_dir(arguments.folder / "results")
    result.to_csv(results / "e4_early_decision.csv", index=False)
    print("樣本放入後 t 秒就判斷的準確率：")
    print(result.pivot(index="method", columns="seconds", values="accuracy").reindex(list(methods)).round(2).to_string())

    figure, axis = plt.subplots(figsize=(7, 4))
    for name in methods:
        part = result[result["method"] == name]
        axis.plot(part["seconds"], part["accuracy"], marker="o", linewidth=3 if name == "蕈狀體" else 1.5, label=name)
    axis.set_xlabel("樣本放入後幾秒就判斷")
    axis.set_ylabel("準確率")
    axis.set_ylim(0, 1.05)
    axis.set_xticks(CUTS)
    axis.legend(fontsize=8)
    axis.set_title("不等感測器穩定，能多早判斷？")
    save_figure(figure, results / "e4_early_decision.png")


if __name__ == "__main__":
    main()