"""E1c combinatorial coding: does the e-nose get better with more kinds of "receptors"?

    python experiments/e1c_receptor_diversity.py data/demo_calibration

A fly has ~50 receptor types and recognises odours by the COMBINATION that responds.
Our "receptor types" are the 5 BME688 heater temperatures plus SGP30 ethanol, SGP30 H2
and SCD40 CO2. This script scores the mushroom body with growing subsets of them
(grouped cross-validation by day, mid humidity, all concentrations).
Each KC samples about a third of the PNs, like a real KC samples a small share of the
glomeruli (median 5 in the male-cns connectome), so the KC wiring scales with the subset.
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd
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

SUBSETS: tuple[tuple[str, tuple[str, ...]], ...] = (
    ("只有 SGP30 乙醇", ("ethanol",)),
    ("BME 1 段", ("bme2",)),
    ("BME 1 段＋SGP30＋CO2", ("bme2", "ethanol", "h2", "co2")),
    ("BME 3 段＋SGP30＋CO2", ("bme0", "bme2", "bme4", "ethanol", "h2", "co2")),
    ("BME 5 段＋SGP30＋CO2（全部）", ("bme0", "bme1", "bme2", "bme3", "bme4", "ethanol", "h2", "co2")),
)


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, _, labels, _ = load_table(arguments.folder / "features.csv")
    if "rh_level" in table:
        keep = (table["rh_level"] == "mid").to_numpy()
        table, labels = table[keep], labels[keep]

    rows = []
    for name, columns in SUBSETS:
        pn_count = 2 * len(columns)
        model = MushroomBodyClassifier(inputs_per_kc=max(1, pn_count // 3))
        scores = cross_val_score(model, table[list(columns)].to_numpy(dtype=float), labels,
                                 groups=table[arguments.group], cv=LeaveOneGroupOut())
        rows.append({"受器組合": name, "受器種類數": len(columns), "平均準確率": scores.mean(), "標準差": scores.std()})
    result = pd.DataFrame(rows)
    results = output_dir(arguments.folder / "results")
    result.to_csv(results / "e1c_receptor_diversity.csv", index=False)
    print(result.round(3).to_string(index=False))

    figure, axis = plt.subplots(figsize=(7, 3.8))
    axis.bar(range(len(result)), result["平均準確率"], yerr=result["標準差"], color="#5eead4", edgecolor="black",
             capsize=3)
    axis.set_xticks(range(len(result)), [f"{n}\n({k} 種)" for n, k in zip(result["受器組合"], result["受器種類數"])],
                    fontsize=8)
    axis.axhline(1 / len(np.unique(labels)), color="gray", linestyle=":", label="亂猜")
    axis.set_ylabel("蕈狀體準確率")
    axis.set_ylim(0, 1.05)
    axis.legend(fontsize=8)
    axis.set_title("受器種類越多，組合編碼越好？")
    save_figure(figure, results / "e1c_receptor_diversity.png")


if __name__ == "__main__":
    main()