"""E5 few-shot learning under real-life variation: 1, 3, 5, 10 samples per class, 20 draws each.

    python experiments/e5_few_shot.py data/demo_calibration

A family teaching an e-nose will give it a few samples, and the next smell will not be
exactly as strong. So the training samples are drawn from every concentration of the
training days, and the fixed test set (one held-out day, --test-group) also contains
every concentration. Hypothesis H4: with 1-3 samples per class 蕈狀體 beats kNN and
logistic regression on raw features; the "+果蠅正規化" rows show whether the fly's
normalisation alone explains the difference.
Output: e5_few_shot.csv and e5_learning_curve.png.
"""
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
from sklearn.metrics import accuracy_score

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

from enose.common import load_table, output_dir, plt, save_figure  # noqa: E402
from enose.methods import make_methods  # noqa: E402

SHOTS: tuple[int, ...] = (1, 3, 5, 10)
DRAWS: int = 20


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("folder", type=Path)
    parser.add_argument("--group", default="date")
    parser.add_argument("--test-group", help="which day/batch is the fixed test set (default: the last one)")
    parser.add_argument("--shots", type=int, nargs="+", default=list(SHOTS), help="e.g. --shots 1 3 5")
    arguments = parser.parse_args()
    table, features, labels, names = load_table(arguments.folder / "features.csv")
    if "rh_level" in table:
        keep = (table["rh_level"] == "mid").to_numpy()
        table, features, labels = table[keep], features[keep], labels[keep]
    groups = table[arguments.group].to_numpy()
    # sort real values (batch 10 after batch 9); --test-group is compared as text
    test_group = arguments.test_group or sorted(set(groups))[-1]
    test = table[arguments.group].astype(str).to_numpy() == str(test_group)
    pool = np.flatnonzero(~test)
    classes = np.unique(labels)
    # Every method may use the unlabelled measurements of the training days for its scaling /
    # adaptation (a fly has smelled many things without reward); only the few samples are labelled.
    methods = {name: method for name, method in make_methods(names, len(classes), background=features[pool]).items()
               if name in ("蕈狀體", "kNN（k=1）", "邏輯回歸", "kNN＋果蠅正規化", "邏輯回歸＋果蠅正規化")}
    rng = np.random.default_rng(0)

    rows = []
    for shots in arguments.shots:
        for draw in range(DRAWS):
            chosen = []
            for label in classes:
                candidates = pool[labels[pool] == label]
                if len(candidates) < shots:
                    raise SystemExit(f"類別 {label} 在訓練池只有 {len(candidates)} 筆，不夠抽 {shots} 筆")
                chosen.extend(rng.choice(candidates, size=shots, replace=False))
            for name, method in methods.items():
                model = clone(method).fit(features[chosen], labels[chosen])
                rows.append({"shots": shots, "draw": draw, "method": name,
                             "accuracy": accuracy_score(labels[test], model.predict(features[test]))})
    table_out = pd.DataFrame(rows)
    table_out.to_csv(output_dir(arguments.folder / "results") / "e5_few_shot.csv", index=False)
    curve = table_out.groupby(["method", "shots"])["accuracy"].agg(["mean", "std"]).reset_index()
    print(curve.pivot(index="shots", columns="method", values="mean").round(3).to_string())

    figure, axis = plt.subplots(figsize=(6, 4))
    for name, group in curve.groupby("method"):
        axis.errorbar(group["shots"], group["mean"], yerr=group["std"], marker="o", capsize=3, label=name)
    axis.axhline(1 / len(classes), color="gray", linestyle=":", label="亂猜")
    axis.set_xlabel("每類訓練樣本數")
    axis.set_ylabel("測試準確率")
    axis.set_xticks(arguments.shots)
    axis.legend(fontsize=8)
    axis.set_title(f"少樣本學習曲線（{DRAWS} 次抽樣，平均 ± 標準差）")
    save_figure(figure, output_dir(arguments.folder / "results") / "e5_learning_curve.png")


if __name__ == "__main__":
    main()