"""E2 classifier comparison with grouped cross-validation, plus a PCA picture.

    python experiments/e2_compare_methods.py data/demo_calibration --group date
    python experiments/e2_compare_methods.py data/demo_banana --group batch
    python experiments/e2_compare_methods.py data/demo_calibration --connectome data/connectome/kc_wiring.npy

--group decides what is held out together: a whole measurement day (E1-E4 data)
or a whole banana batch (E7 data), so the test data never comes from the same day
or batch as the training data. Hypothesis H1 (non-inferiority): the best fly model is
within 5 percentage points of the best classic method. Also used for the optional
extensions E9 (fruit vinegar) and E10 (other fruit): point it at that data folder.
"""
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.decomposition import PCA
from sklearn.metrics import ConfusionMatrixDisplay, accuracy_score
from sklearn.model_selection import LeaveOneGroupOut
from sklearn.preprocessing import StandardScaler

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


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("folder", type=Path)
    parser.add_argument("--group", default="date", help="labels.csv column to hold out together")
    parser.add_argument("--all-humidity", action="store_true", help="also use low/high humidity rows")
    parser.add_argument("--connectome", type=Path, help="kc_wiring.npy from connectome_wiring.py")
    arguments = parser.parse_args()
    table, features, labels, names = load_table(arguments.folder / "features.csv")
    if "rh_level" in table and not arguments.all_humidity:
        keep = (table["rh_level"] == "mid").to_numpy()
        table, features, labels = table[keep], features[keep], labels[keep]
    results = output_dir(arguments.folder / "results")
    connectivity = np.load(arguments.connectome) if arguments.connectome else None
    methods = make_methods(names, len(np.unique(labels)), connectivity)
    groups = table[arguments.group].to_numpy()
    classes = np.unique(labels)

    rows, predictions = [], {}
    for name, method in methods.items():
        predicted = np.empty(len(labels), dtype=object)
        fold_scores = []
        for train, test in LeaveOneGroupOut().split(features, labels, groups):
            model = clone(method).fit(features[train], labels[train])
            predicted[test] = model.predict(features[test])
            fold_scores.append(accuracy_score(labels[test], predicted[test]))
        predictions[name] = predicted
        rows.append({"方法": name, "平均準確率": np.mean(fold_scores), "標準差": np.std(fold_scores),
                     "折數": len(fold_scores)})
    summary = pd.DataFrame(rows).sort_values("平均準確率", ascending=False)
    summary.to_csv(results / "e2_accuracy.csv", index=False)
    print(summary.round(3).to_string(index=False))

    figure, axis = plt.subplots(figsize=(5, 4.4))
    ConfusionMatrixDisplay.from_predictions(labels, predictions["蕈狀體"], labels=classes, ax=axis,
                                            xticks_rotation=45, colorbar=False)
    axis.set_title("蕈狀體的混淆矩陣")
    save_figure(figure, results / "e2_confusion_mushroom.png")

    points = PCA(n_components=2).fit_transform(StandardScaler().fit_transform(features))
    figure, axis = plt.subplots(figsize=(5.5, 4.4))
    for label in classes:
        mask = labels == label
        axis.scatter(points[mask, 0], points[mask, 1], s=14, label=label)
    axis.set_xlabel("PC1")
    axis.set_ylabel("PC2")
    axis.legend(fontsize=8)
    axis.set_title("PCA：各類氣味分得開嗎？")
    save_figure(figure, results / "e2_pca.png")


if __name__ == "__main__":
    main()