# SPDX-License-Identifier: MIT
"""YS-EN-0003 · what is the period worth? Shape vs period on OGLE-SLC.

OGLE-SLC (build_ogle.py) has StarLightCurves' size, length and class mix, but
every curve keeps its catalogue period. On the same 1,000-curve training sets
(resamples 0-4) we compare:

  period only        gradient-boosted trees on log10 P alone
  shape only         MiniROCKET on the folded curve, with aeon's MiniRocketClassifier head
  shape + period     MiniROCKET's ridge scores (out-of-fold on the training set)
                     plus log10 P, combined by the same gradient-boosted trees
  shape, random phase  as shape only, but every curve rotated by a random phase
                     (StarLightCurves' phase alignment is approximate)
  shape, 250 epochs  as shape only, on ogle_slc_e250.npz: each star cut to a
                     random 250 of its OGLE-IV epochs (OGLE-II had ~200-300)

  python ogle_experiment.py      # -> results/ogle_runs.jsonl, results/preds/ogle-*.npz
  python ogle_experiment.py cut  # -> results/period_cut.json: the best single period cut
                                 #    between RR Lyrae and Cepheids in the full LMC catalogue
"""
from __future__ import annotations

import json
import time
from pathlib import Path

import numpy as np

import lcdata

HERE = Path(__file__).resolve().parent
OUT = HERE / "results" / "ogle_runs.jsonl"
PREDS = HERE / "results" / "preds"


def confusion(yt, pred):
    cm = np.zeros((3, 3), int)
    for a, b in zip(yt, pred):
        cm[a, b] += 1
    return cm.tolist()


def main(resamples=range(5)):
    from aeon.transformations.collection.convolution_based import MiniRocket
    from sklearn.ensemble import HistGradientBoostingClassifier
    from sklearn.linear_model import RidgeClassifierCV
    from sklearn.model_selection import StratifiedKFold
    from sklearn.pipeline import make_pipeline
    from sklearn.preprocessing import StandardScaler

    PREDS.mkdir(parents=True, exist_ok=True)
    done = set()
    if OUT.exists():
        done = {(d["run"], d["resample"]) for d in map(json.loads, OUT.read_text().splitlines())}

    def record(run, r, yt, pred, ids, t0, **extra):
        rec = {"run": run, "resample": r, "acc": float((pred == yt).mean()), "confusion": confusion(yt, pred),
               "seconds": round(time.time() - t0, 1), **extra}
        np.savez_compressed(PREDS / f"ogle-{run}-r{r}.npz", pred=pred.astype(np.int8), y=yt, ids=ids)
        with open(OUT, "a") as f:
            f.write(json.dumps(rec) + "\n")
        print(json.dumps({k: rec[k] for k in ("run", "resample", "acc")}), flush=True)

    def combiner(r):
        return HistGradientBoostingClassifier(max_depth=3, max_iter=200, learning_rate=0.05, random_state=r)

    def ridge():  # exactly aeon's MiniRocketClassifier head, as used on StarLightCurves
        return make_pipeline(StandardScaler(with_mean=False), RidgeClassifierCV(alphas=np.logspace(-3, 3, 10)))

    def minirocket(X, Xt, r):
        mr = MiniRocket(random_state=r, n_jobs=-1).fit(X[:, None, :])
        F = mr.transform(X[:, None, :])
        Ft = np.concatenate([mr.transform(Xt[i:i + 1000, None, :]) for i in range(0, len(Xt), 1000)])
        return F, Ft

    for r in resamples:
        X, y, Xt, yt, P, Pt, ids = lcdata.load_ogle(r, with_period=True, with_ids=True)
        lp, lpt = np.log10(P)[:, None], np.log10(Pt)[:, None]

        if ("period-only", r) not in done:
            t0 = time.time()
            record("period-only", r, yt, combiner(r).fit(lp, y).predict(lpt), ids, t0)

        if ("shape-minirocket", r) not in done or ("shape+period", r) not in done:
            t0 = time.time()
            F, Ft = minirocket(X, Xt, r)
            shape = ridge().fit(F, y)
            if ("shape-minirocket", r) not in done:
                record("shape-minirocket", r, yt, shape.predict(Ft), ids, t0)
            # Out-of-fold ridge scores on the training set, so the combiner never sees
            # scores from a model that was fitted on the same curve.
            oof = np.zeros((len(y), 3))
            for tr, va in StratifiedKFold(5, shuffle=True, random_state=r).split(F, y):
                oof[va] = ridge().fit(F[tr], y[tr]).decision_function(F[va])
            comb = combiner(r).fit(np.hstack([oof, lp]), y)
            record("shape+period", r, yt, comb.predict(np.hstack([shape.decision_function(Ft), lpt])), ids, t0)

        if ("shape-randphase", r) not in done:
            t0 = time.time()
            rng = np.random.default_rng(500 + r)
            roll = lambda A: np.stack([np.roll(a, rng.integers(A.shape[1])) for a in A])  # noqa: E731
            F, Ft = minirocket(roll(X), roll(Xt), r)
            record("shape-randphase", r, yt, ridge().fit(F, y).predict(Ft), ids, t0)

        if ("shape-e250", r) not in done and (lcdata.WORK / "ogle" / "ogle_slc_e250.npz").exists():
            t0 = time.time()
            X2, y2, Xt2, yt2 = lcdata.load_ogle(r, variant="_e250")
            assert (y2 == y).all() and (yt2 == yt).all()
            F, Ft = minirocket(X2, Xt2, r)
            record("shape-e250", r, yt, ridge().fit(F, y).predict(Ft), ids, t0)


def period_cut():
    """Best single cut on log P between all OGLE-IV LMC RR Lyrae (ab + c) and classical
    Cepheids, with and without the first-overtone Cepheids."""
    import sys

    sys.path.insert(0, str(HERE))
    from build_ogle import read_table

    rr = np.log10([p for t in ("rrlyr_RRab.dat", "rrlyr_RRc.dat") for _, p, _ in read_table(t)])
    out = {}
    for name, tables in (("fundamental", ["cep_cepF.dat"]), ("fundamental+first_overtone", ["cep_cepF.dat", "cep_cep1O.dat"])):
        cep = np.log10([p for t in tables for _, p, _ in read_table(t)])
        cuts = np.sort(np.concatenate([rr, cep]))
        # correct = RR Lyrae below the cut + Cepheids at or above it
        correct = np.searchsorted(np.sort(rr), cuts, side="left") + (len(cep) - np.searchsorted(np.sort(cep), cuts, side="left"))
        k = int(np.argmax(correct))
        out[name] = {"n_rrlyrae": len(rr), "n_cepheids": len(cep), "n": len(rr) + len(cep),
                     "cut_days": float(10 ** cuts[k]), "accuracy": float(correct[k] / (len(rr) + len(cep))),
                     "misclassified": int(len(rr) + len(cep) - correct[k])}
    (HERE / "results" / "period_cut.json").write_text(json.dumps(out, indent=1))
    print(json.dumps(out, indent=1))


if __name__ == "__main__":
    import sys

    period_cut() if sys.argv[1:] == ["cut"] else main()