"""E7 banana ripening and spoilage: how many days earlier does the e-nose notice?

    python experiments/e7_banana.py data/demo_banana
    python experiments/e7_banana.py data/demo_banana --intensity-pns 0     (pattern only)

The e-nose model is the mushroom body WITH intensity PNs (default 8): spoilage is mostly
"stronger ethanol and CO2", which the pattern-only model throws away (hypothesis H6).

labels.csv columns (one row per banana per day):
    sample_id, batch, banana, storage (room/fridge), day, label (fresh/ripe/overripe/rotten),
    brown_pct, brix, ph, weight_loss_pct, inedible_votes (0-3, from 3 blinded classmates)

Steps
 1. Leave-one-batch-out: train on the other batches, predict every measurement of this batch.
 2. e-nose rotten day = first day the model says "rotten" with confidence >= 0.7 two days in a row.
 3. visual rotten day = first day at least 2 of 3 classmates vote "inedible".
 4. lead = visual day - e-nose day (positive = the e-nose was earlier).
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.metrics import accuracy_score

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

from enose.common import STAGE_NAMES, STAGE_ORDER, load_table, output_dir, plt, save_figure  # noqa: E402
from enose.mushroom import MushroomBodyClassifier  # noqa: E402

CONFIDENCE: float = 0.7
CONSECUTIVE_DAYS: int = 2


def first_run_day(days: np.ndarray, flags: np.ndarray, run: int) -> float:
    """First day that starts `run` consecutive True flags on consecutive calendar days
    (days must be sorted; a missing day breaks the run)."""
    for start in range(len(flags) - run + 1):
        if flags[start:start + run].all() and days[start + run - 1] - days[start] == run - 1:
            return float(days[start])
    return float("nan")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("folder", type=Path)
    parser.add_argument("--intensity-pns", type=int, default=8)
    arguments = parser.parse_args()
    table, features, labels, _ = load_table(arguments.folder / "features.csv")
    results = output_dir(arguments.folder / "results")

    table["predicted"] = ""
    table["rotten_confidence"] = np.nan
    for batch in sorted(table["batch"].unique()):
        test = (table["batch"] == batch).to_numpy()
        model = MushroomBodyClassifier(intensity_pns=arguments.intensity_pns).fit(features[~test], labels[~test])
        table.loc[test, "predicted"] = model.predict(features[test])
        if "rotten" in model.classes_:
            rotten_column = list(model.classes_).index("rotten")
            table.loc[test, "rotten_confidence"] = model.predict_proba(features[test])[:, rotten_column]
        else:
            print(f"注意：訓練資料（第 {batch} 批以外）沒有腐敗樣本，這批無法判定腐敗日")
            table.loc[test, "rotten_confidence"] = 0.0
        print(f"第 {batch} 批（用其他批訓練）準確率：{accuracy_score(labels[test], table.loc[test, 'predicted']):.3f}")
    table.to_csv(results / "e7_predictions.csv", index=False)

    rows = []
    for banana, group in table.sort_values("day").groupby("banana"):
        days = group["day"].to_numpy()
        nose_flags = ((group["predicted"] == "rotten") & (group["rotten_confidence"] >= CONFIDENCE)).to_numpy()
        eye_flags = (group["inedible_votes"] >= 2).to_numpy()
        nose_day = first_run_day(days, nose_flags, CONSECUTIVE_DAYS)
        eye_day = first_run_day(days, eye_flags, 1)
        rows.append({"banana": banana, "batch": group["batch"].iloc[0], "storage": group["storage"].iloc[0],
                     "enose_day": nose_day, "visual_day": eye_day, "lead_days": eye_day - nose_day})
    lead = pd.DataFrame(rows)
    lead.to_csv(results / "e7_lead_days.csv", index=False)
    judged = lead.dropna(subset=["lead_days"])
    print(f"\n兩種判定都有結果的香蕉：{len(judged)} / {len(lead)}（其餘在追蹤期內沒腐敗或電子鼻沒判出）")
    for storage, group in judged.groupby("storage"):
        print(f"{storage}: 提早天數中位數 {group['lead_days'].median():.1f} 天，"
              f"提早 ≥ 1 天的比例 {(group['lead_days'] >= 1).mean():.0%}（n={len(group)}）")

    figure, axes = plt.subplots(1, 3, figsize=(15, 4))
    for storage, group in table.groupby("storage"):
        daily = group.groupby("day")
        axes[0].plot(daily["brown_pct"].mean(), marker="o", label=f"{storage} 褐斑 %")
        axes[1].plot(daily["ethanol"].mean().abs(), marker="o", label=f"{storage} 乙醇訊號")
        axes[1].plot(daily["co2"].mean(), marker="s", linestyle="--", label=f"{storage} CO2（千 ppm）")
    axes[0].set_title("腐敗曲線：褐斑面積")
    axes[1].set_title("腐敗曲線：氣味訊號")
    for axis in axes[:2]:
        axis.set_xlabel("天數")
        axis.legend(fontsize=8)
    for index, storage in enumerate(sorted(judged["storage"].unique())):
        values = judged[judged["storage"] == storage]["lead_days"]
        axes[2].scatter(np.full(len(values), index) + np.random.default_rng(0).uniform(-0.1, 0.1, len(values)),
                        values, alpha=0.7)
        axes[2].hlines(values.median(), index - 0.25, index + 0.25, color="black")
    axes[2].axhline(0, color="gray", linestyle=":")
    axes[2].set_xticks(range(len(judged["storage"].unique())), sorted(judged["storage"].unique()))
    axes[2].set_ylabel("肉眼日 − 電子鼻日（天）")
    axes[2].set_title("每根香蕉的提早天數（黑線＝中位數）")
    save_figure(figure, results / "e7_spoilage.png")

    stage_names = [STAGE_NAMES[stage] for stage in STAGE_ORDER]
    counts = pd.crosstab(table["label"], table["predicted"]).reindex(index=STAGE_ORDER, columns=STAGE_ORDER,
                                                                    fill_value=0)
    print("\n混淆矩陣（列＝標準答案、欄＝電子鼻）：")
    print(counts.set_axis(stage_names, axis=0).set_axis(stage_names, axis=1).to_string())


if __name__ == "__main__":
    main()