"""E3 (core) concentration extrapolation: learn one concentration, recognise all the others.

    python experiments/e3_concentration.py data/demo_calibration

A fly that learned "vinegar" on one visit still recognises it much weaker or much stronger.
Leave-one-day-out: for each measurement day, train on the --train-concentration
(default 1x) samples of the OTHER days and test that day's other concentrations
(mid humidity), so day-to-day drift never helps.

What the rows mean:
    蕈狀體                   the fly model (PN gain control + APL)
    蕈狀體（無增益控制）      ablation: no PN normalisation, no APL -> is the fly mechanism the cause?
    kNN / 邏輯回歸           classic methods on the raw features
    kNN／邏輯回歸＋果蠅正規化  classic methods with ONLY the fly's normalisation added
                             -> is the fly principle portable to other methods?
Hypothesis H2: 蕈狀體 stays >= 80% from 0.25x to 4x while the classic methods on raw
features drop below it; and the ablation drops.
"""
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.common import load_table, output_dir, plt, save_figure  # noqa: E402
from enose.methods import make_methods  # noqa: E402

SHOWN: tuple[str, ...] = ("蕈狀體", "蕈狀體（無增益控制）", "kNN（k=1）", "邏輯回歸", "kNN＋果蠅正規化", "邏輯回歸＋果蠅正規化")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("folder", type=Path)
    parser.add_argument("--train-concentration", type=float, default=1.0)
    parser.add_argument("--group", default="date", help="column holding the measurement day")
    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].reset_index(drop=True), features[keep], labels[keep]
    concentrations = sorted(table["concentration"].unique())
    train_level = (table["concentration"] == arguments.train_concentration).to_numpy()
    days = table[arguments.group].to_numpy()
    methods = {name: method for name, method in make_methods(names, len(np.unique(labels))).items() if name in SHOWN}

    predictions = {name: np.full(len(labels), None, dtype=object) for name in methods}
    for day in np.unique(days):
        train = train_level & (days != day)
        test = ~train_level & (days == day)
        if not test.any() or len(np.unique(labels[train])) < len(np.unique(labels)):
            continue
        for name, method in methods.items():
            predictions[name][test] = clone(method).fit(features[train], labels[train]).predict(features[test])
    rows = []
    for name in methods:
        for concentration in concentrations:
            test = (table["concentration"] == concentration).to_numpy() & (predictions[name] != None)  # noqa: E711
            if concentration == arguments.train_concentration or not test.any():
                continue
            rows.append({"method": name, "concentration": concentration, "n": int(test.sum()),
                         "accuracy": float((predictions[name][test] == labels[test]).mean())})
    result = pd.DataFrame(rows)
    results = output_dir(arguments.folder / "results")
    result.to_csv(results / "e3_concentration.csv", index=False)
    pivot = result.pivot(index="method", columns="concentration", values="accuracy").reindex(list(methods))
    print(f"訓練濃度 {arguments.train_concentration}×，測試其他濃度的準確率：")
    print(pivot.round(2).to_string())

    figure, axis = plt.subplots(figsize=(7.5, 4.2))
    styles = {"蕈狀體": ("#0f766e", "-", 3), "蕈狀體（無增益控制）": ("#0f766e", ":", 1.5),
              "kNN（k=1）": ("#6b7280", "-", 1.5), "邏輯回歸": ("#9ca3af", "-", 1.5),
              "kNN＋果蠅正規化": ("#7c3aed", "--", 1.5), "邏輯回歸＋果蠅正規化": ("#a78bfa", "--", 1.5)}
    for name in methods:
        part = result[result["method"] == name]
        colour, line, width = styles[name]
        axis.plot(part["concentration"], part["accuracy"], marker="o", color=colour, linestyle=line,
                  linewidth=width, label=name)
    axis.axvline(arguments.train_concentration, color="#f59e0b", linestyle=":", label="訓練濃度")
    axis.set_xscale("log")
    axis.set_xlabel("測試濃度（倍數，對數刻度）")
    axis.set_ylabel("準確率")
    axis.set_ylim(0, 1.05)
    axis.legend(fontsize=8, loc="lower center", ncol=2)
    axis.set_title("只學一種濃度，能認出多濃／多淡的同一種氣味嗎？")
    save_figure(figure, results / "e3_concentration.png")


if __name__ == "__main__":
    main()