"""E2b (core) which parts of the fly's smell circuit matter? Remove one at a time.

    python experiments/e2b_ablation.py data/demo_calibration

Every row is the full mushroom body with ONE fly mechanism removed, scored on three tasks:
    標準辨識      grouped cross-validation by day (mid humidity, all concentrations)
    濃度外推      train at 1x on the other days, test other concentrations of the held-out day
    少樣本(1)     1 labelled sample per gas, 20 random draws, held-out last day
A mechanism "matters" when removing it lowers a score by more than the draw-to-draw
spread. Expected from biology: gain control -> concentration; adaptation -> accuracy and
few-shot; sparse KC code -> separating similar smells; separate dopamine compartments
-> no forgetting (see E6, not visible here).
"""
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.model_selection import LeaveOneGroupOut, cross_val_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.mushroom import MushroomBodyClassifier  # noqa: E402

VARIANTS: dict[str, dict] = {
    "完整果蠅模型": {},
    "拿掉增益控制（除總量＋APL）": {"normalize": False, "kc_mode": "threshold"},
    "拿掉 PN 適應": {"adaptation": False},
    "拿掉稀疏編碼（KC 50% 放電）": {"active_fraction": 0.5},
    "拿掉分區學習（跨區懲罰）": {"punishment": 1 / 3},
}
DRAWS: int = 20


def concentration_by_day(model, features: np.ndarray, labels: np.ndarray, groups: np.ndarray,
                         train_1x: np.ndarray) -> float:
    """Train at 1x on the other days, test the other concentrations of the held-out day."""
    correct = total = 0
    for day in np.unique(groups):
        train = train_1x & (groups != day)
        test = ~train_1x & (groups == day)
        if test.any() and len(np.unique(labels[train])) == len(np.unique(labels)):
            correct += int((clone(model).fit(features[train], labels[train]).predict(features[test]) == labels[test]).sum())
            total += int(test.sum())
    return correct / max(total, 1)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("folder", type=Path)
    parser.add_argument("--group", default="date")
    arguments = parser.parse_args()
    table, features, labels, _ = load_table(arguments.folder / "features.csv")
    keep = (table["rh_level"] == "mid").to_numpy() if "rh_level" in table else np.ones(len(table), bool)
    table, features, labels = table[keep].reset_index(drop=True), features[keep], labels[keep]
    groups = table[arguments.group].to_numpy()
    train_1x = (table["concentration"] == 1.0).to_numpy()
    test_day = groups == sorted(set(groups))[-1]
    pool = np.flatnonzero(~test_day)
    rng = np.random.default_rng(0)
    # The same 20 one-shot draws for every variant, so differences are not draw-to-draw noise.
    draws = [np.concatenate([rng.choice(pool[labels[pool] == label], 1) for label in np.unique(labels)])
             for _ in range(DRAWS)]

    rows = []
    for name, settings in VARIANTS.items():
        model = MushroomBodyClassifier(**settings)
        standard = cross_val_score(model, features, labels, groups=groups, cv=LeaveOneGroupOut()).mean()
        concentration = concentration_by_day(model, features, labels, groups, train_1x)
        one_shot = []
        for chosen in draws:
            few = clone(model).set_params(background=features[pool]).fit(features[chosen], labels[chosen])
            one_shot.append((few.predict(features[test_day]) == labels[test_day]).mean())
        rows.append({"模型": name, "標準辨識": standard, "濃度外推": concentration,
                     "少樣本(1)": float(np.mean(one_shot)), "少樣本標準差": float(np.std(one_shot))})
    result = pd.DataFrame(rows)
    results = output_dir(arguments.folder / "results")
    result.to_csv(results / "e2b_ablation.csv", index=False)
    print(result.round(3).to_string(index=False))

    figure, axis = plt.subplots(figsize=(9, 4))
    tasks = ["標準辨識", "濃度外推", "少樣本(1)"]
    width = 0.8 / len(result)
    for index, row in result.iterrows():
        axis.bar(np.arange(len(tasks)) + index * width, [row[task] for task in tasks], width, label=row["模型"])
    axis.set_xticks(np.arange(len(tasks)) + 0.4 - width / 2, tasks)
    axis.set_ylabel("準確率")
    axis.set_ylim(0, 1.05)
    axis.legend(fontsize=8, loc="lower right")
    axis.set_title("拿掉果蠅的哪一步，表現掉最多？")
    save_figure(figure, results / "e2b_ablation.png")


if __name__ == "__main__":
    main()