"""End-to-end demo of clfkit on built-in / synthetic data (no downloads).

Run:  python examples/demo_end_to_end.py [--out examples/output]

Parts
  A. Breast cancer (binary, tabular)    - calibration, conformal, abstention, certificates
  B. Digits (10 classes, tabular)       - class-conditional coverage gap and Mondrian fix;
                                          the same data loaded as "embeddings" (.npy)
  C. Support tickets (synthetic text)   - full CLI round trip + three drift scenarios
  D. Embedding similarity               - why cosine 0.2 can be significant, and why you
                                          still need an empirical null
"""

from __future__ import annotations

import argparse
import time
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 import similarity as sim
from clfkit.cli import main as cli
from clfkit.data import detect_column_types, load_embeddings
from clfkit.datasets import make_anisotropic_embeddings, make_text_tickets
from clfkit.report import write_report


def banner(t):
    print("\n" + "=" * 78 + f"\n{t}\n" + "=" * 78)


def part_a(out: Path):
    banner("A. Breast cancer (binary): calibration / conformal / abstention / certificate")
    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")
    b = pipeline.train(df[types.all()], df["label"].to_numpy(), types, model="auto", seed=0,
                       calib_size=0.3)
    res = pipeline.evaluate(b, alpha=0.05, target_risk=0.15, delta=0.05)
    pipeline.save(b, out / "breast_cancer", res)
    write_report(res, b, out / "breast_cancer" / "report.html")
    m = res["test_metrics"]
    print(f"test acc {m['accuracy']:.3f}, ECE {m['ece']:.3f}")
    for k, v in res["conformal"].items():
        print(f"  {k:13s} coverage {v['coverage']:.3f} (target 0.95)  avg set {v['avg_set_size']:.2f}")
    s = res["selective"]
    if s["threshold"] is not None:
        print(f"  abstain below conf {s['threshold']:.4f}: test answered {s['test_coverage']:.1%}, "
              f"error among answered {s['test_selective_risk']:.4f} (certified <= 0.15 w.p. 0.95)")
    else:
        print("  abstention: target 15% risk not certifiable with this little held-out data")
    # How much data would a stricter target need? Even with ZERO errors, the bound is
    # 1 - (delta/m)^(1/n), so certifying risk r needs n >= ln(m/delta)/r answered points.
    tight = pipeline.choose_threshold(b.P_conf.max(1), b.P_conf.argmax(1) == b.y_conf, 0.02, 0.05)
    need = int(np.ceil(np.log(tight["grid_size"] / 0.05) / 0.02))
    print(f"  stricter target 2%: {'certified' if tight['threshold'] is not None else 'NOT certifiable'} "
          f"with n={len(b.y_conf)}; even error-free it needs >= {need} answered held-out points")
    c = res["certified"]
    print(f"  logreg certified: clean acc {c['clean_accuracy']:.3f}, median L2 radius {c['median_radius']:.3f} "
          "(standardised-feature units)")


def part_b(out: Path):
    banner("B. Digits (10 classes): class-conditional coverage gap -> Mondrian conformal")
    X, y = load_digits(return_X_y=True)
    # Save as embeddings to exercise the .npy loader (any model's embeddings work the same way).
    np.save(out / "digits_emb.npy", X.astype(np.float32))
    np.save(out / "digits_labels.npy", y.astype(str))
    Xdf, ys, types = load_embeddings(out / "digits_emb.npy", out / "digits_labels.npy")
    # A bigger calibration split (40%) so each digit has ~36 conformal calibration points;
    # Mondrian needs n_k >= 1/alpha - 1 per class or its threshold is infinite.
    b = pipeline.train(Xdf, ys, types, model="logreg", compare=False, seed=0, verbose=False,
                       calib_size=0.4)
    res = pipeline.evaluate(b, alpha=0.1)
    pipeline.save(b, out / "digits", res)
    write_report(res, b, out / "digits" / "report.html")
    print(f"test acc {res['test_metrics']['accuracy']:.3f}; calibration T = "
          f"{[r for r in res['calibration'] if r['method'] == 'temperature'][0]['T']:.2f}")
    for k in ("lac", "lac_mondrian", "aps", "aps_mondrian"):
        v = res["conformal"][k]
        worst = min(v["per_class_coverage"], key=v["per_class_coverage"].get)
        print(f"  {k:13s} coverage {v['coverage']:.3f}  avg set {v['avg_set_size']:.2f}  "
              f"worst class '{b.classes[worst]}' {v['worst_class_coverage']:.3f}")
    print(f"  (target 0.90; calib_conf has {len(b.y_conf)} points, ~{len(b.y_conf) // 10} per class, and the")
    print("   test split has ~36 per class, so per-class numbers are noisy; Mondrian's guarantee is per")
    print("   class on average over calibration draws, not for every finite test sample.)")
    # Average over 50 random re-splits of the pooled held-out points to see the
    # systematic effect rather than one noisy draw.
    from clfkit.conformal import SplitConformal, evaluate_sets

    P = np.vstack([b.P_conf, b.P_test]); yy = np.r_[b.y_conf, b.y_test]
    rng = np.random.default_rng(0)
    acc = {False: [], True: []}
    for r in range(50):
        perm = rng.permutation(len(yy)); c, t = perm[:360], perm[360:]
        for mond in (False, True):
            cp = SplitConformal(0.1, "lac", mond, seed=r).fit(P[c], yy[c])
            acc[mond].append(list(evaluate_sets(cp.predict_sets(P[t]), yy[t])["per_class_coverage"].values()))
    for mond in (False, True):
        pc = np.mean(acc[mond], axis=0)
        print(f"  LAC{' Mondrian' if mond else '         '} averaged over 50 re-splits: worst class "
              f"'{b.classes[pc.argmin()]}' {pc.min():.3f}, best class {pc.max():.3f}")


def part_c(out: Path):
    banner("C. Synthetic support tickets (text + categorical + numeric) via the CLI")
    data = out / "tickets.csv"
    make_text_tickets(3000, seed=0).to_csv(data, index=False)
    batches = {
        "same_distribution": make_text_tickets(600, seed=11),
        "label_shift": make_text_tickets(600, seed=12, class_probs=[0.55, 0.15, 0.15, 0.15]),
        "new_vocabulary": make_text_tickets(600, seed=13, new_vocab_rate=0.25),
    }
    for name, df in batches.items():
        df.to_csv(out / f"batch_{name}.csv", index=False)
    mdir = out / "tickets_model"
    cli(["train", str(data), "--target", "label", "--text-col", "text", "--out", str(mdir)])
    print("\n$ clfkit predict ... batch_same_distribution.csv --alpha 0.1")
    cli(["predict", str(mdir), str(out / "batch_same_distribution.csv"), "--alpha", "0.1", "--show", "3"])
    for name in batches:
        print(f"\n$ clfkit drift ... batch_{name}.csv   (true label proportions for label_shift: "
              "billing .55, others .15)" if name == "label_shift" else f"\n$ clfkit drift ... batch_{name}.csv")
        cli(["drift", str(mdir), str(out / f"batch_{name}.csv"), "--out", str(mdir / f"drift_{name}.json")])
    # The report picks up MODEL_DIR/drift.json: show the vocabulary-shift scenario in it.
    (mdir / "drift.json").write_text((mdir / "drift_new_vocabulary.json").read_text())
    cli(["report", str(mdir)])


def part_d(out: Path):
    banner("D. Embedding similarity: null distributions and a precision-calibrated threshold")
    for d in (64, 384, 768, 1536):
        print(f"  isotropic null d={d:5d}: sd {sim.isotropic_null_sd(d):.4f}, "
              f"P(cos >= 0.2) = {sim.isotropic_null_pvalue(0.2, d):.2e}")
    E, y = make_anisotropic_embeddings(n=1500, d=384, n_classes=10, seed=0)
    np.save(out / "aniso_emb.npy", E)
    np.save(out / "aniso_labels.npy", y.astype(str))
    pairs, rel = sim.sample_pairs(y, 6000, seed=0)
    for center in (False, True):
        cos = sim.pair_cosines(sim.normalize(E, center=center), pairs)
        ns = sim.null_summary(cos[rel == 0], E.shape[1])
        r = sim.threshold_for_precision(cos, rel, 0.95, 0.05)
        c = r["certified"]
        print(f"  synthetic anisotropic embeddings (d=384, centred={center}): unrelated-pair cosine "
              f"mean {ns['empirical_mean']:.3f} sd {ns['empirical_sd']:.3f}; "
              f"{ns['empirical_frac_above_0.2']:.0%} of UNRELATED pairs exceed 0.2")
        if c:
            print(f"    threshold for precision >= 0.95 (95% conf.): cos >= {c['threshold']:.2f} "
                  f"(empirical precision {c['precision']:.3f}, recall {c['recall']:.3f})")
        else:
            print("    no cosine threshold certifies precision >= 0.95 (the shared direction swamps the signal)")
    print("\n$ clfkit similarity aniso_emb.npy aniso_labels.npy --center")
    cli(["similarity", str(out / "aniso_emb.npy"), str(out / "aniso_labels.npy"), "--center"])


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=str(Path(__file__).parent / "output"))
    args = ap.parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    t0 = time.perf_counter()
    for part in (part_a, part_b, part_c, part_d):
        t = time.perf_counter()
        part(out)
        print(f"[{part.__name__} took {time.perf_counter() - t:.1f}s]")
    print(f"\nAll done in {time.perf_counter() - t0:.1f}s. Reports: {out}/*/report.html")