"""E6 continual learning: learn 3 classes, then only the 4th. Are the first 3 forgotten?

    python experiments/e6_continual.py data/demo_banana --group batch --new-class rotten

Compared:
    蕈狀體                      MushroomBodyClassifier.partial_fit; every MBON has its own dopamine
                                compartment (punishment=0), so learning class 4 leaves the
                                synapses of classes 1-3 untouched. Hypothesis H5.
    蕈狀體（跨區懲罰）          punishment=1/3: new learning also weakens other MBONs (ablation)
    邏輯回歸（逐步更新）        SGDClassifier(loss="log_loss").partial_fit, the usual
                                way to update a linear model with new data only
Both see exactly the same data in the same order. Scored on a held-out group.
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.linear_model import SGDClassifier
from sklearn.metrics import accuracy_score
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.mushroom import MushroomBodyClassifier  # noqa: E402

EPOCHS: int = 5  # passes over each phase's data for the SGD model


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("folder", type=Path)
    parser.add_argument("--group", default="batch")
    parser.add_argument("--new-class", default="rotten")
    arguments = parser.parse_args()
    table, features, labels, _ = load_table(arguments.folder / "features.csv")
    groups = table[arguments.group].to_numpy()
    test = groups == sorted(set(groups))[-1]  # sort real values: batch 10 comes after batch 9
    classes = np.unique(labels)
    old = labels != arguments.new_class
    phase_one = ~test & old
    phase_two = ~test & ~old

    scaler = StandardScaler().fit(features[~test])
    mushroom = MushroomBodyClassifier(intensity_pns=8)
    crosstalk = MushroomBodyClassifier(intensity_pns=8, punishment=1 / 3)
    linear = SGDClassifier(loss="log_loss", random_state=0)
    rows = []

    def score(stage: str) -> None:
        for name, model, data in (("蕈狀體", mushroom, features), ("蕈狀體（跨區懲罰）", crosstalk, features),
                                  ("邏輯回歸（逐步更新）", linear, scaler.transform(features))):
            predicted = model.predict(data[test])
            for subset, mask in (("舊類別", old[test]), ("新類別", ~old[test])):
                if mask.any():
                    rows.append({"stage": stage, "method": name, "subset": subset,
                                 "accuracy": accuracy_score(labels[test][mask], predicted[mask])})

    for model in (mushroom, crosstalk):
        model.partial_fit(features[phase_one], labels[phase_one], classes=classes[classes != arguments.new_class])
    for _ in range(EPOCHS):
        linear.partial_fit(scaler.transform(features[phase_one]), labels[phase_one], classes=classes)
    score("學完舊 3 類")
    for model in (mushroom, crosstalk):
        model.partial_fit(features[phase_two], labels[phase_two])
    for _ in range(EPOCHS):
        linear.partial_fit(scaler.transform(features[phase_two]), labels[phase_two])
    score("再學新類別")

    result = pd.DataFrame(rows)
    results = output_dir(arguments.folder / "results")
    result.to_csv(results / "e6_continual.csv", index=False)
    pivot = result.pivot_table(index=["method", "subset"], columns="stage", values="accuracy")
    print(pivot.round(3).to_string())

    figure, axes = plt.subplots(1, 3, figsize=(14, 3.8), sharey=True)
    for axis, (name, group) in zip(axes, result.groupby("method")):
        for subset, part in group.groupby("subset"):
            axis.plot(part["stage"], part["accuracy"], marker="o", label=subset)
        axis.set_title(name)
        axis.legend()
    axes[0].set_ylabel("準確率")
    save_figure(figure, results / "e6_continual.png")


if __name__ == "__main__":
    main()