"""Export held-out model outputs for the browser demo on the project page.

Run:  python examples/export_demo_data.py [--out demo/data]

For each demo model (the same three as examples/demo_end_to_end.py, same seeds and
settings) this trains with clfkit and writes one JSON file with the *uncalibrated*
probabilities and labels of the three held-out splits:

    calib_fit   the browser fits the temperature on these
    calib_conf  conformal quantiles and the abstention threshold
    test        coverage, set sizes and error are measured here

plus the APS tie-breaking uniforms clfkit draws (so the browser's APS sets match the
Python ones exactly) and something to show for each test item (ticket text, the
digit's pixels, a few tumour measurements).

For the support tickets it also writes tickets_drift.json: the drift reference (the
test split) and pools of fresh tickets at four new-vocabulary rates, each row with the
model's predicted label and calibrated confidence and the per-row values clfkit's drift tests look at.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.datasets import load_breast_cancer, load_digits

from clfkit import pipeline
from clfkit.data import (
    ColumnTypes, clean_text_columns, detect_column_types, load_csv, make_split,
)
from clfkit.datasets import make_text_tickets
from clfkit.models import full_proba

CHANNELS = ["chat", "email", "phone"]
SIG = 7  # significant digits kept for probabilities


def r(a, sig=SIG):
    """Round to significant digits (keeps tiny probabilities, unlike fixed decimals)."""
    return [float(f"{x:.{sig}g}") for x in np.asarray(a, float).ravel()]


def rows(P, sig=SIG):
    return [r(p, sig) for p in np.asarray(P, float)]


def held_out(b: pipeline.Bundle, X: pd.DataFrame, y: np.ndarray, calib_size: float):
    """Raw probabilities on the three held-out splits, recomputed from the bundle's model."""
    split = make_split(y, 0.2, calib_size, b.seed)  # the same call pipeline.train makes
    assert np.array_equal(y[split.calib_conf], b.y_conf) and np.array_equal(y[split.test], b.y_test)
    K = b.n_classes
    P_conf_raw = full_proba(b.model, X.iloc[split.calib_conf], K)
    assert np.allclose(b.calibrator.transform(P_conf_raw), b.P_conf)
    # The uniforms SplitConformal draws: seed for calibration, seed + 7919 for test.
    u_cal = np.random.default_rng(b.seed).uniform(size=len(b.y_conf))
    u_test = np.random.default_rng(b.seed + 7919).uniform(size=len(b.y_test))
    out = {
        "classes": b.classes.tolist(),
        "model": b.model_name,
        "split_sizes": b.split_sizes,
        "calib_fit": {"P": rows(b.P_calfit_raw), "y": b.y_calfit.tolist()},
        "calib_conf": {"P": rows(P_conf_raw), "y": b.y_conf.tolist(), "u": r(u_cal, 12)},
        "test": {"P": rows(b.P_test_raw), "y": b.y_test.tolist(), "u": r(u_test, 12)},
        "python": {"T": getattr(b.calibrator, "T_", None)},
    }
    return out, split


def tickets(out: Path, tmp: Path):
    data = tmp / "tickets.csv"
    make_text_tickets(3000, seed=0).to_csv(data, index=False)
    X, y_str, types = load_csv(data, "label", ["text"])  # what `clfkit train` does
    b = pipeline.train(X, y_str, types, target="label", seed=0, verbose=False)
    _, y = np.unique(y_str, return_inverse=True)
    d, split = held_out(b, X, y, 0.2)
    d.update(name="tickets", title="Support tickets", alpha=0.1, risk=0.05, delta=0.05,
             about="3,000 synthetic support tickets (free text, channel, past tickets), "
                   "4 teams, gradient-boosted trees picked by clfkit train.")
    d["items"] = {"kind": "text", "text": X["text"].iloc[split.test].tolist()}
    (out / "tickets.json").write_text(json.dumps(d, separators=(",", ":")))

    # ---- drift: reference = test split (what pipeline.drift uses), pools of new tickets
    def feats(df: pd.DataFrame) -> dict:
        toks = df["text"].fillna("").astype(str).str.lower().str.split()
        return {
            "n_prev_tickets": df["n_prev_tickets"].astype(int).tolist(),
            "channel": [CHANNELS.index(c) for c in df["channel"].astype(str)],
            "length": [len(t) for t in toks],
            "oov": [sum(w not in b.vocab for w in t) for t in toks],  # oov_rate = oov / length
        }

    def scored(df: pd.DataFrame) -> dict:
        # What pipeline.drift uses: the calibrated model's predicted label and confidence.
        P = b.calibrator.transform(full_proba(b.model, df[types.all()], b.n_classes))
        return {"pred": P.argmax(1).tolist(), "conf": r(P.max(1), 6)}

    ref = b.X_ref
    drift = {
        "classes": b.classes.tolist(),
        "channels": CHANNELS,
        "ref": {**scored(ref), "y": b.y_ref.tolist(), **feats(ref)},
        "pools": [],
    }
    assert np.array_equal(b.y_ref, b.y_test)
    for seed, rate in ((101, 0.0), (102, 0.05), (103, 0.10), (104, 0.25)):
        df = make_text_tickets(2400, seed=seed, new_vocab_rate=rate)
        Xn = clean_text_columns(df, types)
        idx = {c: i for i, c in enumerate(b.classes)}
        drift["pools"].append({"new_vocab_rate": rate, "seed": seed, **scored(Xn),
                               "y": [idx[c] for c in df["label"]], **feats(Xn)})
    (out / "tickets_drift.json").write_text(json.dumps(drift, separators=(",", ":")))


def breast_cancer(out: Path):
    ds = load_breast_cancer(as_frame=True)
    df = ds.frame.rename(columns={"target": "label"})
    df["label"] = np.where(df["label"] == 1, "benign", "malignant")
    types = detect_column_types(df, target="label")
    X = df[types.all()]
    b = pipeline.train(X, df["label"].to_numpy(), types, model="auto", seed=0, calib_size=0.3,
                       verbose=False)
    _, y = np.unique(df["label"].astype(str).to_numpy(), return_inverse=True)
    d, split = held_out(b, X, y, 0.3)
    d.update(name="breast_cancer", title="Breast cancer", alpha=0.05, risk=0.15, delta=0.05,
             about="569 tumour biopsies (30 measurements each), benign or malignant, "
                   "a random forest picked by clfkit train. Small data on purpose.")
    cols = ["mean radius", "mean texture", "worst concave points"]
    d["items"] = {"kind": "table", "cols": cols,
                  "X": [r(v, 4) for v in X[cols].iloc[split.test].to_numpy()]}
    (out / "breast_cancer.json").write_text(json.dumps(d, separators=(",", ":")))


def digits(out: Path):
    Xa, ya = load_digits(return_X_y=True)
    cols = [f"e{i}" for i in range(Xa.shape[1])]
    X = pd.DataFrame(Xa.astype(np.float32), columns=cols)  # what load_embeddings gives
    y_str = ya.astype(str)
    b = pipeline.train(X, y_str, ColumnTypes(numeric=cols), model="logreg", compare=False, seed=0,
                       verbose=False, calib_size=0.4)
    _, y = np.unique(y_str, return_inverse=True)
    d, split = held_out(b, X, y, 0.4)
    d.update(name="digits", title="Handwritten digits", alpha=0.1, risk=0.05, delta=0.05,
             about="1,797 handwritten digits as 8×8 pixel images, 10 classes, "
                   "a logistic regression loaded through clfkit's embeddings path.")
    d["items"] = {"kind": "pixels", "px": Xa[split.test].astype(int).tolist()}
    (out / "digits.json").write_text(json.dumps(d, separators=(",", ":")))


if __name__ == "__main__":
    import tempfile

    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=str(Path(__file__).resolve().parent.parent / "demo" / "data"))
    a = ap.parse_args()
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory() as tmp:
        tickets(out, Path(tmp))
    breast_cancer(out)
    digits(out)
    for f in sorted(out.glob("*.json")):
        print(f"{f}  {f.stat().st_size / 1024:.0f} KB")