# SPDX-License-Identifier: MIT
"""Build the figures for YS-EN-0003 from results/ and write them into the note
between <!-- fig:name --> markers.

Usage: python build_figures.py   (after run_all.sh, baselines.py runs, ogle_experiment.py)
Needs only the standard library and NumPy (for the OGLE period histogram).
"""
from __future__ import annotations

import json
import math
import statistics as st
import sys
from collections import defaultdict
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent))
import svgchart as sc  # noqa: E402

PAGE = HERE.parent.parent / "ys-en-0003-light-curves-without-periods.html"
RES = HERE / "results"
CLASSES = ["Cepheid", "Eclipsing binary", "RR Lyrae"]  # colour follows the class: slots 1, 2, 3


def jsonl(name):
    p = RES / name
    return [json.loads(line) for line in p.read_text().splitlines() if line.strip()] if p.exists() else []


IMAGE = [d for d in jsonl("image_runs.jsonl") if d.get("dataset", "slc") == "slc"]
IMAGE_OGLE = [d for d in jsonl("image_runs.jsonl") if d.get("dataset") == "ogle"]
BASE = jsonl("baseline_runs.jsonl")
OGLE = jsonl("ogle_runs.jsonl")
PUB = json.loads((RES / "published.json").read_text())


def group(rows, key=lambda d: (d["run"], d["ntrain"])):
    g = defaultdict(list)
    for d in rows:
        g[key(d)].append(d)
    return g


def summary(rows):
    a = [d["acc"] for d in rows]
    return {"mean": st.mean(a), "sd": st.stdev(a) if len(a) > 1 else 0.0, "min": min(a), "max": max(a), "n": len(a),
            "conf": [[sum(d["confusion"][i][j] for d in rows) for j in range(3)] for i in range(3)]}


IMG = {k: summary(v) for k, v in group(IMAGE).items()}
BL = {k: summary(v) for k, v in group(BASE).items()}  # resamples 0-4: our check against the published means
BL3 = {k: summary(v) for k, v in group([d for d in BASE if d["resample"] <= 2]).items()}  # 0-2: matched to the image models


def pct(v, nd=2):
    return f"{100 * v:.{nd}f}%"


def errors_per_test_set(conf, n_sets):
    """Mean count per 8,236-curve test set of each error type."""
    return {f"{CLASSES[i]} → {CLASSES[j]}": conf[i][j] / n_sets for i in range(3) for j in range(3) if i != j}


# ---------------------------------------------------------------- generic dot plot
def dot_plot(rows, lo, hi, ticks, label, ref=None, ref_text="", width=720, left=250):
    """rows: dicts with label, mean, min, max, slot, note, n. Whiskers span min to max."""
    row_h = 28
    fr = sc.Frame(width, 40 + row_h * len(rows) + 40, label, margins=(left, 30, 30, 46))
    fr.x_scale(lo, hi)
    fr.y_scale(0, 1)
    fr.out.append('<g class="grid">')
    for t in ticks:
        fr.out.append(f'<line x1="{fr.x(t):.1f}" x2="{fr.x(t):.1f}" y1="{fr.mt}" y2="{fr.mt + fr.ph}"/>')
    fr.out.append("</g>")
    for t in ticks:
        fr.out.append(f'<text x="{fr.x(t):.1f}" y="{fr.mt + fr.ph + 18}" text-anchor="middle">{pct(t, 1)}</text>')
    if ref is not None:
        x = fr.x(ref)
        fr.out.append(f'<line class="axis" x1="{x:.1f}" x2="{x:.1f}" y1="{fr.mt}" y2="{fr.mt + fr.ph}"/>')
        fr.note(x + 5, fr.mt - 8, ref_text)
    fr.titles("Test accuracy on StarLightCurves (8,236 curves)")
    for i, r in enumerate(rows):
        cy = fr.mt + row_h * (i + 0.5) + 6
        fr.out.append(f'<text class="ink" x="{fr.ml - 12}" y="{cy + 4:.1f}" text-anchor="end">{sc.esc(r["label"])}</text>')
        fr.out.append(f'<line class="grid" x1="{fr.ml}" x2="{fr.ml + fr.pw}" y1="{cy:.1f}" y2="{cy:.1f}"/>')
        m = max(lo, min(hi, r["mean"]))
        if r.get("min") is not None and r["max"] > r["min"]:
            a, b = fr.x(max(lo, r["min"])), fr.x(min(hi, r["max"]))
            fr.out.append(f'<line class="s{r["slot"]}" x1="{a:.1f}" x2="{b:.1f}" y1="{cy:.1f}" y2="{cy:.1f}" stroke-width="2"/>')
        cx = fr.x(m)
        if r["mean"] < lo:  # off the left edge: arrow and value
            fr.out.append(f'<path class="s{r["slot"]}" d="M{cx:.1f},{cy:.1f} l9,-5 l0,10 Z" stroke="none"/>')
            fr.note(cx + 14, cy + 4, pct(r["mean"]), cls="ink")
        else:
            fr.out.append(f'<circle class="dot s{r["slot"]}" cx="{cx:.1f}" cy="{cy:.1f}" r="5"/>')
        spread = f" (range {pct(r['min'])}–{pct(r['max'])}, n = {r['n']})" if r.get("n", 1) > 1 else ""
        sc.point_hit(fr, cx, cy, i, r["label"], f'{r["slot"]}|{r["note"]}|{pct(r["mean"])}{spread}', r=13)
    return fr.svg()


def img_row(run, label, slot, note, ntrain=1000):
    s = IMG.get((run, ntrain))
    return None if s is None else {"label": label, **s, "slot": slot, "note": note}


def pub_row(key, label, slot=2):
    p = PUB[key]
    return {"label": label, "mean": p["mean"], "min": p["mean"] - p["std"], "max": p["mean"] + p["std"], "n": p["n"],
            "slot": slot, "note": "published, mean ± 1 SD of 30 resamples"}


# ---------------------------------------------------------------- figure: encodings
def fig_encodings():
    names = [("stack", "Stack (GASF, GADF, MTF)"), ("rp", "Recurrence plot"), ("gasf", "GASF"), ("gadf", "GADF"),
             ("line", "Line plot"), ("mtf", "Markov transition field")]
    rows = [r for e, lab in names if (r := img_row(f"{e}-resnet18-pt", lab, 1, "ResNet-18, ImageNet weights, fine-tuned"))]
    rows.sort(key=lambda r: -r["mean"])
    mr = BL3[("MiniROCKET", 1000)]
    plot = dot_plot(rows, 0.965, 0.985, [0.965, 0.97, 0.975, 0.98, 0.985],
                    "Accuracy of a fine-tuned ImageNet ResNet-18 on six encodings of StarLightCurves",
                    ref=mr["mean"], ref_text=f"MiniROCKET, same splits: {pct(mr['mean'])}")
    tab = [[r["label"], pct(r["mean"]), pct(r["min"]), pct(r["max"]), str(r["n"])] for r in rows]
    tab.append(["MiniROCKET (time-series, for reference)", pct(mr["mean"]), pct(mr["min"]), pct(mr["max"]), str(mr["n"])])
    return sc.figure(
        code="FIG. 2 — WHICH PICTURE?", fid="fig-encodings", plot=plot,
        caption=("<b>Figure 2.</b> The same network and recipe on six ways of drawing a light curve. Dots are the mean "
                 "over resamples 0–2; whiskers span the lowest and highest run. The vertical line is MiniROCKET on "
                 "resamples 0–2. Every encoding except the Markov transition field lands within a few tenths of a "
                 "point of it, and none is above it."),
        data_html=sc.table("Accuracy by encoding (ResNet-18, ImageNet weights)", ["Encoding", "Mean", "Min", "Max", "Runs"], tab),
    )


# ---------------------------------------------------------------- figure: scoreboard
def fig_board():
    rows = []
    for key, lab in (("HC2", "HIVE-COTE 2"), ("MiniROCKET", "MiniROCKET"), ("MR-Hydra", "MultiROCKET-Hydra"),
                     ("InceptionTime", "InceptionTime"), ("ResNet", "ResNet (1D)")):
        rows.append(pub_row(key, lab + " · published"))
    for key, lab in (("MiniROCKET", "MiniROCKET · ours"), ("MultiROCKET-Hydra", "MultiROCKET-Hydra · ours")):
        rows.append({"label": lab, **BL3[(key, 1000)], "slot": 2, "note": "aeon, resamples 0–2, same splits as our image models"})
    for run, lab, note in (
        ("stack-resnet18-pt", "ResNet-18 · ImageNet", "stack encoding, fine-tuned, resamples 0–2"),
        ("stack-resnet50-pt", "ResNet-50 · ImageNet", "stack encoding, fine-tuned"),
        ("stack-convnext_t-pt", "ConvNeXt-T · ImageNet", "stack encoding, fine-tuned"),
        ("stack-efficientnet_b0-pt", "EfficientNet-B0 · ImageNet", "stack encoding, fine-tuned"),
        ("stack-vit_s16-pt", "ViT-S/16 · ImageNet-21k", "stack encoding, fine-tuned"),
        ("stack-dinov2_s-pt", "DINOv2 ViT-S/14", "stack encoding, fine-tuned"),
        ("stack-resnet18-scratch", "ResNet-18 · random init", "stack encoding, trained from scratch"),
        ("stack-resnet50-probe", "ResNet-50 · frozen + linear", "frozen ImageNet features, logistic regression"),
        ("stack-dinov2_s-probe", "DINOv2 · frozen + linear", "frozen features, logistic regression"),
        ("original-vit-rainbow-gasf", "Our 2023 sketch, re-created", "rainbow GASF, ViT from scratch, 12 epochs"),
    ):
        r = img_row(run, lab, 1 if "scratch" not in run and "original" not in run else 0, note)
        if r:
            rows.append(r)
    for lab, v, note in (("Timage, ResNet-50 (2019)", 0.9806, "Wenninger et al., one run, official split"),
                         ("TiViT, frozen ViT-H (2025)", 0.973, "Roschmann et al., mean of 3 seeds")):
        rows.append({"label": lab, "mean": v, "min": None, "max": v, "n": 1, "slot": 3, "note": note})
    plot = dot_plot(rows, 0.955, 0.99, [0.96, 0.97, 0.98, 0.99],
                    "Accuracy on StarLightCurves: time-series methods, our image models and published image models",
                    ref=PUB["HC2"]["mean"], ref_text="HC2, best published", left=262)
    leg = ('<div class="legend"><span class="k2"><i></i>Time-series method</span>'
           '<span class="k1"><i></i>Pretrained image model (ours)</span><span class="k0"><i></i>Image model without pretraining (ours)</span>'
           '<span class="k3"><i></i>Image model (published)</span></div>')
    tab = [[r["label"], pct(r["mean"]), "—" if r.get("min") is None else f"{pct(r['min'])}–{pct(r['max'])}", str(r["n"]), r["note"]] for r in rows]
    return sc.figure(
        code="FIG. 3 — THE SCOREBOARD", fid="fig-board", plot=plot, legend_html=leg,
        caption=("<b>Figure 3.</b> Our runs (mean over resamples 0–2, whiskers min to max; the image models and our "
                 "two ROCKET rows share those splits) against published results (mean ± one standard deviation over 30 "
                 "resamples). An arrow means the value is off the "
                 "left edge. One percentage point is 82 test curves."),
        data_html=sc.table("Accuracy on StarLightCurves", ["Model", "Accuracy", "Spread", "Runs", "Notes"], tab),
    )


# ---------------------------------------------------------------- figure: learning curve
def fig_curve():
    def by_n(rows_by_key, run):
        pts = sorted((n, s["mean"]) for (r, n), s in rows_by_key.items() if r == run)
        return pts

    def spread(rows_by_key, run, n):
        s = rows_by_key.get((run, n))
        return "—" if s is None else f"{pct(s['mean'])} ({pct(s['min'])}–{pct(s['max'])})"

    series = []
    for slot, name, pts in ((2, "MiniROCKET", by_n(BL3, "MiniROCKET")),
                            (1, "ResNet-18, ImageNet weights", by_n(IMG, "stack-resnet18-pt")),
                            (0, "ResNet-18, from scratch", by_n(IMG, "stack-resnet18-scratch"))):
        if pts:
            series.append({"slot": slot, "name": name, "pts": dict(pts)})
    xs = sorted({n for s in series for n in s["pts"]})
    for s in series:
        s["values"] = [s["pts"].get(n) for n in xs]
    fr = sc.Frame(720, 360, "Accuracy against number of training curves for MiniROCKET and ResNet-18 with and without pretraining")
    fr.x_scale(20, 1200, log=True)
    fr.y_scale(0.80, 1.0)
    fr.grid_y([0.8, 0.85, 0.9, 0.95, 1.0], fmt=lambda v: pct(v, 0))
    fr.grid_x([25, 50, 100, 250, 500, 1000], fmt=lambda v: f"{v:,.0f}", lines=False)
    fr.baseline()
    fr.titles("Training curves (log scale; stratified subsets of the 1,000)", "Test accuracy")
    fr.cursor()
    # draw_lines needs complete series on shared x; draw each on its own points
    for s in series:
        pts = [(fr.x(n), fr.y(max(0.8, v))) for n, v in zip(xs, s["values"]) if v is not None]
        fr.out.append(f'<path class="line s{s["slot"]}" d="M' + " L".join(f"{a:.1f},{b:.1f}" for a, b in pts) + '"/>')
        ex, ey = pts[-1]
        fr.out.append(f'<circle class="dot s{s["slot"]}" cx="{ex:.1f}" cy="{ey:.1f}" r="4"/>')
    sc.column_hits(fr, xs, series, lambda i, x: f"{x:,} training curves", lambda v: pct(v))
    src = {"MiniROCKET": (BL3, "MiniROCKET"), "ResNet-18, ImageNet weights": (IMG, "stack-resnet18-pt"),
           "ResNet-18, from scratch": (IMG, "stack-resnet18-scratch")}
    rows = [[f"{n:,}"] + [spread(*src[s["name"]], n) for s in series] for n in xs]
    return sc.figure(
        code="FIG. 4 — LEARNING CURVE", fid="fig-curve", plot=fr.svg(), legend_html=sc.legend(series),
        caption=("<b>Figure 4.</b> Accuracy as the training set shrinks: means over resamples 0–2, the same "
                 "subsets for every model (the small subsets are drawn per class in proportion, so 50 becomes 51). "
                 "The table view gives each mean with its lowest and highest run."),
        data_html=sc.table("Accuracy by training-set size", ["Training curves", *(s["name"] for s in series)], rows),
    )


# ---------------------------------------------------------------- table: where the errors are
def tab_errors():
    rows = []
    for lab, s in (("MiniROCKET", BL3.get(("MiniROCKET", 1000))),
                   ("MultiROCKET-Hydra", BL3.get(("MultiROCKET-Hydra", 1000))),
                   ("ResNet-18 · ImageNet, stack", IMG.get(("stack-resnet18-pt", 1000))),
                   ("ResNet-18 · ImageNet, recurrence plot", IMG.get(("rp-resnet18-pt", 1000)))):
        if not s:
            continue
        e = errors_per_test_set(s["conf"], s["n"])
        tot = sum(e.values())
        rows.append([lab, f"{tot:.1f}", f'{e["Cepheid → RR Lyrae"]:.1f}', f'{e["RR Lyrae → Cepheid"]:.1f}',
                     f'{tot - e["Cepheid → RR Lyrae"] - e["RR Lyrae → Cepheid"]:.1f}',
                     pct((e["Cepheid → RR Lyrae"] + e["RR Lyrae → Cepheid"]) / tot, 1)])
    return sc.table("Errors per 8,236-curve test set, mean over resamples 0–2",
                    ["Model", "All errors", "Cepheid → RR Lyrae", "RR Lyrae → Cepheid", "Anything else", "Share Cep ↔ RRL"], rows)


# ---------------------------------------------------------------- figure: OGLE periods
def fig_periods():
    import numpy as np

    d = np.load(Path.home() / "ys-work" / "ogle" / "ogle_slc.npz")
    lp, y = np.log10(d["period"]), d["y"]
    edges = np.arange(-1.2, 3.01, 0.1)
    mids = (edges[:-1] + edges[1:]) / 2
    # share of each class per bin, so the 5,327 RR Lyrae do not flatten the other two classes
    series = [{"slot": c + 1, "name": CLASSES[c],
               "values": (np.histogram(lp[y == c], edges)[0] / (y == c).sum()).tolist()} for c in range(3)]
    top = max(max(s["values"]) for s in series)
    fr = sc.Frame(720, 360, "Histogram of catalogue periods in OGLE-SLC by class, log scale")
    fr.x_scale(edges[0], edges[-1])
    ymax = math.ceil(top * 20) / 20
    fr.y_scale(0, ymax)
    fr.grid_y([k / 20 for k in range(0, int(round(ymax * 20)) + 1)], fmt=lambda v: pct(v, 0))
    lab = {-1: "0.1 d", 0: "1 d", 1: "10 d", 2: "100 d", 3: "1,000 d"}
    fr.grid_x(list(lab), fmt=lambda v: lab[v], lines=False)
    fr.baseline()
    fr.titles("Period (log scale)", "Share of the class per 0.1 dex")
    fr.vline(0.0, "1 day")
    fr.cursor()
    for s in series:  # step outline per class
        pts = []
        for i, v in enumerate(s["values"]):
            pts += [(fr.x(edges[i]), fr.y(v)), (fr.x(edges[i + 1]), fr.y(v))]
        fr.out.append(f'<path class="line s{s["slot"]}" d="M' + " L".join(f"{a:.1f},{b:.1f}" for a, b in pts) + '"/>')
    sc.column_hits(fr, mids.tolist(), series, lambda i, x: f"{10 ** edges[i]:.3g}–{10 ** edges[i + 1]:.3g} d",
                   lambda v: pct(v, 1), hover_marks=False)
    stats = []
    for c in range(3):
        p = d["period"][y == c]
        stats.append([CLASSES[c], f"{(y == c).sum():,}", f"{p.min():.3g} d", f"{np.median(p):.3g} d", f"{p.max():.4g} d",
                      pct((p < 1).mean(), 1)])
    return sc.figure(
        code="FIG. 5 — WHAT THE BENCHMARK THREW AWAY", fid="fig-periods", plot=fr.svg(), legend_html=sc.legend(series),
        caption=("<b>Figure 5.</b> Catalogue periods of the 9,236 stars in OGLE-SLC, each class as a share of its own "
                 "stars. RR Lyrae sit below a day and nearly all Cepheids above it; eclipsing binaries span four decades and overlap both."),
        data_html=sc.table("Periods by class", ["Class", "Stars", "Shortest", "Median", "Longest", "Below 1 day"], stats),
    )


# ---------------------------------------------------------------- figure: OGLE shape vs period
OGLE_ITEMS = [  # label, run key (None = the image model), slot, note
    ("Period only", "period-only", 0, "gradient-boosted trees on log P"),
    ("Shape · MiniROCKET", "shape-minirocket", 2, "folded curve, no period"),
    ("Shape · ResNet-18", None, 1, "stack encoding, ImageNet weights, fine-tuned"),
    ("Shape + period", "shape+period", 2, "MiniROCKET scores + log P, gradient-boosted trees"),
]
OGLE_CONTROLS = [
    ("Shape · random phase", "shape-randphase", "every curve rotated by a random phase"),
    ("Shape · 250 epochs", "shape-e250", "each star cut to 250 random epochs"),
]


def ogle_summary(key):
    if key is None:
        return summary(IMAGE_OGLE) if IMAGE_OGLE else None
    rows = [d for d in OGLE if d["run"] == key]
    return summary(rows) if rows else None


def ogle_row(lab, s, note):
    e = errors_per_test_set(s["conf"], s["n"])
    return [lab, pct(s["mean"]), f"{pct(s['min'])}–{pct(s['max'])}", str(s["n"]), f"{sum(e.values()):.1f}",
            f'{e["Cepheid → RR Lyrae"]:.1f}', f'{e["RR Lyrae → Cepheid"]:.1f}',
            f'{sum(v for k, v in e.items() if "Eclipsing" in k):.1f}', note]


def fig_ogle():
    cats, vals, slots, rows = [], [], [], []
    for lab, key, slot, note in OGLE_ITEMS:
        s = ogle_summary(key)
        if s is None:
            continue
        cats.append(lab)
        vals.append(s["mean"])
        slots.append(slot)
        rows.append(ogle_row(lab, s, note))
    for lab, key, note in OGLE_CONTROLS:
        s = ogle_summary(key)
        if s:
            rows.append(ogle_row(lab, s, note))
    fr = sc.Frame(720, 330, "Accuracy on OGLE-SLC with period only, shape only, and both", margins=(58, 30, 30, 46))
    fr.y_scale(0.80, 1.0)
    fr.grid_y([0.80, 0.85, 0.90, 0.95, 1.0], fmt=lambda v: pct(v, 0))
    fr.baseline()
    fr.titles("", "Test accuracy (8,236 curves; axis starts at 80%)")
    sc.columns(fr, cats, vals, fmt=lambda v: pct(v), slots=slots, metric="Accuracy")
    return sc.figure(
        code="FIG. 6 — SHAPE, PERIOD, BOTH", fid="fig-ogle", plot=fr.svg(),
        caption=("<b>Figure 6.</b> The same task on OGLE-SLC, where every curve keeps its period. Means over "
                 "resamples (five for the tree models and MiniROCKET, three for the image model). The axis starts at "
                 "80%. The table view adds the two controls and the error counts."),
        data_html=sc.table("OGLE-SLC accuracy and mean errors per test set",
                           ["Model", "Accuracy", "Range", "Runs", "Errors", "Cep → RRL", "RRL → Cep", "Involving an EB", "Notes"], rows),
    )


# ---------------------------------------------------------------- table: which Cepheids stay wrong
def tab_modes():
    import numpy as np

    ogle = Path.home() / "ys-work" / "ogle"
    d = np.load(ogle / "ogle_slc.npz")
    period = dict(zip(d["ids"], d["period"]))
    o1 = {line.split()[0] for line in open(ogle / "cep_cep1O.dat")}

    def kind(i):
        return "Fundamental mode" if i not in o1 else ("First overtone, P < 1 d" if period[i] < 1 else "First overtone, P ≥ 1 d")

    kinds = ["Fundamental mode", "First overtone, P ≥ 1 d", "First overtone, P < 1 d"]
    test = {k: 0 for k in kinds}
    wrong = {run: {k: 0 for k in kinds} for run in ("shape-minirocket", "shape+period")}
    for run in wrong:
        for r in range(5):
            p = np.load(RES / "preds" / f"ogle-{run}-r{r}.npz")
            for i, yy, pp in zip(p["ids"], p["y"], p["pred"]):
                if yy != 0:
                    continue
                if run == "shape-minirocket":
                    test[kind(i)] += 1
                if pp == 2:
                    wrong[run][kind(i)] += 1
    rows = [[k, f"{test[k] / 5:.0f}", f'{wrong["shape-minirocket"][k] / 5:.1f}', f'{wrong["shape+period"][k] / 5:.1f}'] for k in kinds]
    rows.append(["All Cepheids", f"{sum(test.values()) / 5:.0f}", f'{sum(wrong["shape-minirocket"].values()) / 5:.1f}',
                 f'{sum(wrong["shape+period"].values()) / 5:.1f}'])
    return sc.table("Test Cepheids misread as RR Lyrae on OGLE-SLC, mean per test set over 5 resamples",
                    ["Cepheid type", "In test set", "Shape only", "Shape + period"], rows)


if __name__ == "__main__":
    for name, fn in [("encodings", fig_encodings), ("board", fig_board), ("curve", fig_curve), ("errors", tab_errors),
                     ("periods", fig_periods), ("ogle", fig_ogle), ("modes", tab_modes)]:
        try:
            sc.inject(PAGE, name, fn())
            print("built", name)
        except SystemExit as e:
            print("skipped", name, "-", e)