"""Publishable figures + metrics (aggregates only - no rows leave the box).

Reads runs/baselines.json, runs/laya-*/eval.json and the saved test predictions, writes
results/metrics.json and results/*.png for the model card.
"""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from common import split

OUT = Path("results"); OUT.mkdir(exist_ok=True)
# Reference palette (dataviz skill), light surface. Validated: see docs/PROCESS.md.
SURFACE, INK, INK2, MUTED, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#8a8984", "#e6e5e1"
BLUE, ORANGE = "#2a78d6", "#eb6834"
plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 11, "axes.edgecolor": GRID,
                     "axes.labelcolor": INK2, "xtick.color": INK2, "ytick.color": INK2,
                     "figure.facecolor": SURFACE, "axes.facecolor": SURFACE, "savefig.facecolor": SURFACE})

base = json.loads(Path("runs/baselines.json").read_text())["paid_7d"]
evals = {p.parent.name: json.loads(p.read_text()) for p in sorted(Path("runs").glob("laya-*/eval.json"))}

MODELS = [  # (label, source, is_ours)
    ("Recency rule", base["recency_rule"], False),
    ("Logistic regression", base["logreg"], False),
    ("LightGBM (telemetry)", base["lightgbm"], False),
]
if "laya-text" in evals:
    MODELS.append(("x402-demand · listing only (open)", evals["laya-text"]["paid_7d"], True))
if "laya-full" in evals:
    MODELS.append(("x402-demand · listing + telemetry (API)", evals["laya-full"]["paid_7d"], True))
# Stacking pair: shown once, for the open listing-only model (the question it answers:
# does the free text score add anything on top of telemetry?).
st = evals.get("laya-text", {}).get("stacking_paid_7d", {})
if st:
    MODELS.append(("LightGBM, calibration hosts only*", st["gbm_calib_only"], False))
    MODELS.append(("  … + listing-only score*", st["gbm_calib+laya"], True))

SEGS = [("all", "All test hosts"), ("new_listings", "New listings (cold start)"),
        ("small_hosts(<=20)", "Hosts with ≤ 20 listings")]

# ---------- Figure 1: AUC by segment, dot plot small multiples ----------
fig, axes = plt.subplots(1, 3, figsize=(13, 0.52 * len(MODELS) + 1.9), sharey=True)
y = np.arange(len(MODELS))[::-1]
for ax, (seg, title) in zip(axes, SEGS):
    ax.axvline(0.5, color=MUTED, lw=1, ls=(0, (3, 3)), zorder=1)
    ax.text(0.505, y.min() - 0.55, "chance", color=MUTED, ha="left", fontsize=9)
    for yi, (label, src, ours) in zip(y, MODELS):
        m = src.get(seg)
        if not m or m.get("auc") is None:
            continue
        c = BLUE if ours else MUTED
        ax.hlines(yi, 0.5, m["auc"], color=c, lw=2, zorder=2)
        ax.scatter([m["auc"]], [yi], s=70, color=c, edgecolor=SURFACE, linewidth=2, zorder=3)
        ax.text(m["auc"] + 0.008, yi, f"{m['auc']:.3f}", va="center", fontsize=9.5, color=INK)
    ax.set_title(title, loc="left", color=INK, fontsize=12)
    ax.set_xlim(0.38, 0.85); ax.set_ylim(y.min() - 0.8, y.max() + 0.5)
    ax.grid(axis="x", color=GRID, lw=0.8); ax.set_axisbelow(True)
    ax.spines[["top", "right", "left"]].set_visible(False)
    ax.tick_params(axis="y", length=0)
    ax.set_xlabel("ROC AUC, paid within 7 days")
axes[0].set_yticks(y, [m[0] for m in MODELS])
for t, (_, _, ours) in zip(axes[0].get_yticklabels(), MODELS):
    t.set_color(INK if ours else INK2)
fig.suptitle("Held-out test: future dates × hosts never seen in training", x=0.01, ha="left",
             color=INK, fontsize=13, fontweight="bold")
note = "Blue = uses the x402-demand checkpoint. Grey = baseline."
if any("*" in m[0] for m in MODELS):
    note += "  * both trained only on calibration hosts; the gap isolates the listing-only score's contribution."
fig.text(0.01, 0.005, note, color=INK2, fontsize=9)
fig.tight_layout(rect=(0, 0.03, 1, 0.95))
fig.savefig(OUT / "auc_by_segment.png", dpi=160)
plt.close(fig)

# ---------- Figure 2: reliability diagram + Figure 3: lift by decile ----------
df = pd.read_parquet("data/rows.parquet")
_, _, te = split(df, 7)
yt = te.paid_7d.to_numpy()
curves = [("LightGBM (telemetry)", np.load("runs/lightgbm_paid7d_test_pred.npy"), ORANGE)]
AQUA = "#1baf7a"  # palette slot 3; below 3:1 on light surface -> legend labels carry identity
for name, lab, c in (("laya-text", "x402-demand · listing only (open)", AQUA),
                     ("laya-full", "x402-demand · listing + telemetry (API, recalibrated)", BLUE)):
    p = Path(f"runs/{name}/test_pred_paid7d_isotonic.npy")
    if not p.exists():
        p = Path(f"runs/{name}/test_pred_paid7d.npy")
    if p.exists():
        curves.append((lab, np.load(p), c))

fig, ax = plt.subplots(figsize=(6.2, 5.6))
ax.plot([0, 1], [0, 1], color=MUTED, lw=1, ls=(0, (3, 3)))
ax.text(0.83, 0.77, "perfect", color=MUTED, rotation=45, fontsize=9)
for label, p, c in curves:
    q = np.quantile(p, np.linspace(0, 1, 11)); q[-1] += 1e-9
    b = np.clip(np.digitize(p, q) - 1, 0, 9)
    xs = [p[b == i].mean() for i in range(10) if (b == i).any()]
    ys = [yt[b == i].mean() for i in range(10) if (b == i).any()]
    ax.plot(xs, ys, color=c, lw=2, marker="o", ms=8, mec=SURFACE, mew=2, label=label)
ax.set_xlim(0, 1); ax.set_ylim(0, 1)
ax.set_xlabel("Predicted probability (decile mean)"); ax.set_ylabel("Observed share paid within 7 days")
ax.set_title("Calibration on the held-out test set", loc="left", color=INK, fontsize=12)
ax.grid(color=GRID, lw=0.8); ax.set_axisbelow(True); ax.spines[["top", "right"]].set_visible(False)
ax.legend(frameon=False, loc="upper left", fontsize=9.5, labelcolor=INK)
fig.tight_layout(); fig.savefig(OUT / "calibration.png", dpi=160); plt.close(fig)

best_label = "x402-demand · listing + telemetry"
best_p = np.load("runs/laya-full/test_pred_paid7d.npy")
q = np.quantile(best_p, np.linspace(0, 1, 11)); q[-1] += 1e-9
b = np.clip(np.digitize(best_p, q) - 1, 0, 9)
rates = [yt[b == i].mean() for i in range(10)]
fig, ax = plt.subplots(figsize=(7.5, 4.4))
bars = ax.bar(np.arange(1, 11), rates, width=0.78, color=BLUE, edgecolor=SURFACE, linewidth=2)
ax.axhline(yt.mean(), color=MUTED, lw=1, ls=(0, (3, 3)))
ax.text(0.6, yt.mean() + 0.012, f"base rate {yt.mean():.0%}", color=INK2, fontsize=9)
for r, bar in zip(rates, bars):
    if r in (rates[0], rates[-1]) or bar is bars[4]:
        ax.text(bar.get_x() + bar.get_width() / 2, r + 0.01, f"{r:.0%}", ha="center", fontsize=9.5, color=INK)
ax.set_xticks(range(1, 11)); ax.set_xlabel(f"Decile of predicted probability ({best_label})")
ax.set_ylabel("Share paid within 7 days"); ax.set_ylim(0, max(rates) * 1.18)
ax.set_title("Who actually gets paid, by predicted decile (test set)", loc="left", color=INK, fontsize=12)
ax.grid(axis="y", color=GRID, lw=0.8); ax.set_axisbelow(True); ax.spines[["top", "right"]].set_visible(False)
fig.tight_layout(); fig.savefig(OUT / "lift_by_decile.png", dpi=160); plt.close(fig)

metrics = {"task": "paid_7d / paid_14d", "test": "anchors 17-21 (7d) / 15 (14d), hosts in buckets 80-99",
           "baselines": json.loads(Path("runs/baselines.json").read_text()),
           "checkpoints": evals,
           "lift_by_decile_paid7d": {"model": best_label, "share_paid": [round(float(r), 4) for r in rates],
                                     "base_rate": round(float(yt.mean()), 4)}}
(OUT / "metrics.json").write_text(json.dumps(metrics, indent=1))
print("wrote", sorted(p.name for p in OUT.iterdir()))