# SPDX-License-Identifier: MIT
"""YS-EN-0003 · time-series-native baselines on the same splits as the image models.

  python baselines.py runs      # MiniROCKET + MultiROCKET-Hydra, resamples 0-4, full + learning-curve sizes
  python baselines.py split70   # MiniROCKET on a pooled, stratified 70/30 split (seeds 0-2)
  python baselines.py published # published accuracies (30 resamples) from timeseriesclassification.com via aeon

Results: results/baseline_runs.jsonl and results/published.json
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

import numpy as np

import lcdata

HERE = Path(__file__).resolve().parent
OUT = HERE / "results"
SIZES = [25, 50, 100, 250, 500, 1000]


def runs(resamples=range(5)):
    from aeon.classification.convolution_based import MiniRocketClassifier, MultiRocketHydraClassifier

    OUT.mkdir(exist_ok=True)
    (OUT / "preds").mkdir(exist_ok=True)
    log = OUT / "baseline_runs.jsonl"
    done = set()
    if log.exists():  # resumable: skip runs already recorded
        for line in log.read_text().splitlines():
            d = json.loads(line)
            done.add((d["run"], d["resample"], d["ntrain"]))
    for r in resamples:
        X, y, Xt, yt = lcdata.load(r)
        for n in SIZES:
            idx = lcdata.subsample(y, n, seed=r) if n < len(y) else np.arange(len(y))
            for name, cls in (("MiniROCKET", MiniRocketClassifier), ("MultiROCKET-Hydra", MultiRocketHydraClassifier)):
                if name == "MultiROCKET-Hydra" and n < 100:
                    continue  # Hydra's per-class groups need more than a handful of examples per class
                if (name, r, int(len(idx))) in done:
                    continue
                t0 = time.time()
                clf = cls(random_state=r, n_jobs=-1).fit(X[idx][:, None, :], y[idx])
                # predict in chunks: the full test feature matrix doesn't fit in WSL's 7 GB
                pred = np.concatenate([clf.predict(Xt[i:i + 1000, None, :]) for i in range(0, len(Xt), 1000)])
                cm = np.zeros((3, 3), int)
                for a, b in zip(yt, pred):
                    cm[a, b] += 1
                rec = {"run": name, "resample": r, "ntrain": int(len(idx)), "acc": float((pred == yt).mean()),
                       "confusion": cm.tolist(), "seconds": round(time.time() - t0, 1)}
                if n == 1000:
                    np.savez_compressed(OUT / "preds" / f"{name}-r{r}-n1000.npz", pred=pred.astype(np.int8), y=yt)
                with open(OUT / "baseline_runs.jsonl", "a") as f:
                    f.write(json.dumps(rec) + "\n")
                print(json.dumps({k: rec[k] for k in ("run", "resample", "ntrain", "acc", "seconds")}), flush=True)


def split70(seeds=range(3)):
    """MiniROCKET on a pooled 70/30 split of all 9,236 curves."""
    from aeon.classification.convolution_based import MiniRocketClassifier
    from sklearn.model_selection import train_test_split

    X, y, Xt, yt = lcdata.load(0)
    XA, yA = np.concatenate([X, Xt]), np.concatenate([y, yt])
    for s in seeds:
        tr, te = train_test_split(np.arange(len(yA)), test_size=0.3, stratify=yA, random_state=s)
        t0 = time.time()
        clf = MiniRocketClassifier(random_state=s, n_jobs=-1).fit(XA[tr][:, None, :], yA[tr])
        pred = np.concatenate([clf.predict(XA[te][i:i + 1000, None, :]) for i in range(0, len(te), 1000)])
        rec = {"run": "MiniROCKET-split70", "seed": s, "ntrain": len(tr), "ntest": len(te),
               "acc": float((pred == yA[te]).mean()), "seconds": round(time.time() - t0, 1)}
        with open(OUT / "split70_runs.jsonl", "a") as f:
            f.write(json.dumps(rec) + "\n")
        print(json.dumps(rec), flush=True)


def published():
    from aeon.benchmarking.results_loaders import get_available_estimators, get_estimator_results

    names = list(get_available_estimators(task="classification").iloc[:, 0])
    res = get_estimator_results(estimators=names, datasets=["StarLightCurves"], num_resamples=30,
                                task="classification", measure="accuracy")
    out = {}
    for k, v in res.items():
        a = np.asarray(v.get("StarLightCurves"), dtype=float)
        if a.size and np.isfinite(a).any():
            a = a[np.isfinite(a)]
            out[k] = {"mean": float(a.mean()), "std": float(a.std(ddof=1)) if a.size > 1 else 0.0,
                      "n": int(a.size), "resample0": float(a[0])}
    OUT.mkdir(exist_ok=True)
    (OUT / "published.json").write_text(json.dumps(out, indent=1))
    for k, v in sorted(out.items(), key=lambda kv: -kv[1]["mean"]):
        print(f"{k:18s} {v['mean']:.4f} ± {v['std']:.4f}  (n={v['n']}, r0={v['resample0']:.4f})")


if __name__ == "__main__":
    {"runs": runs, "split70": split70, "published": published}[sys.argv[1]]()