"""Calibrate the 7-detector AI ensemble on HC3 (human vs ChatGPT answers).

Pipeline:
  1. sample balanced human/AI texts from data/calibration/hc3_all.jsonl
  2. score all 7 detectors on each text
  3. train a logistic meta-classifier (standardized features)
  4. isotonic calibration on a held-out calib split -> honest probabilities
  5. split-conformal abstain band: t_hi = 95th percentile of p over calib
     HUMAN texts (≤5% of humans exceed it), t_lo = 5th percentile over calib
     AI texts — between the two the verdict is INCONCLUSIVE, not a guess
  6. save models/ai_meta.json + report held-out test metrics

Run:  python scripts/calibrate_ensemble.py [n_per_class]   (default 300)
"""
import json
import os
import random
import sys
import time

import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from plagdetect.aidetect import FEATURE_ORDER, detect_ai  # noqa: E402
from plagdetect.normalize import deobfuscate  # noqa: E402
from plagdetect.textutils import sentences  # noqa: E402

HC3 = os.path.join(ROOT, "data", "calibration", "hc3_all.jsonl")
OUT = os.path.join(ROOT, "models", "ai_meta.json")
SCORED = os.path.join(ROOT, "data", "calibration", "hc3_scored.jsonl")
MIN_CHARS, MIN_SENTS = 700, 8


def load_samples(n_per_class, seed=13):
    rng = random.Random(seed)
    human, ai = [], []
    with open(HC3, "r", encoding="utf-8") as f:
        lines = f.readlines()
    rng.shuffle(lines)
    for ln in lines:
        try:
            rec = json.loads(ln)
        except json.JSONDecodeError:
            continue
        for pool, key in ((human, "human_answers"), (ai, "chatgpt_answers")):
            if len(pool) >= n_per_class:
                continue
            for ans in rec.get(key) or []:
                ans = (ans or "").strip()
                if len(ans) >= MIN_CHARS and len(sentences(ans)) >= MIN_SENTS:
                    pool.append(ans)
                    break
        if len(human) >= n_per_class and len(ai) >= n_per_class:
            break
    return human, ai


def score_all(texts, label):
    rows = []
    t0 = time.time()
    for i, t in enumerate(texts):
        t = deobfuscate(t)[0]
        try:
            r = detect_ai(t)
        except Exception as exc:
            print(f"  [skip] {exc}")
            continue
        det = {d["name"]: d["score"] for d in r["detectors"]}
        if not all(k in det for k in FEATURE_ORDER):
            continue
        rows.append({"y": label, "x": [det[k] for k in FEATURE_ORDER]})
        if (i + 1) % 25 == 0:
            rate = (i + 1) / (time.time() - t0)
            print(f"  scored {i+1}/{len(texts)} (label={label}, "
                  f"{rate:.1f}/s, eta {(len(texts)-i-1)/rate:.0f}s)", flush=True)
    return rows


def main(n_per_class=300, refit=False):
    if refit and os.path.exists(SCORED):
        rows = [json.loads(ln) for ln in open(SCORED, encoding="utf-8")]
        print(f"REFIT: loaded {len(rows)} cached scores from {SCORED} "
              "(skipping detector scoring)")
    else:
        print("loading HC3 samples...")
        human, ai = load_samples(n_per_class)
        print(f"human={len(human)} ai={len(ai)}")
        rows = score_all(human, 0) + score_all(ai, 1)
        with open(SCORED, "w", encoding="utf-8") as f:
            for r in rows:
                f.write(json.dumps(r) + "\n")
        print(f"scored {len(rows)} samples -> {SCORED}")

    X = np.array([r["x"] for r in rows], dtype=float)
    y = np.array([r["y"] for r in rows], dtype=int)
    rng = np.random.RandomState(13)
    idx = rng.permutation(len(y))
    n_tr, n_cal = int(0.6 * len(y)), int(0.2 * len(y))
    tr, cal, te = idx[:n_tr], idx[n_tr:n_tr + n_cal], idx[n_tr + n_cal:]

    from sklearn.isotonic import IsotonicRegression
    from sklearn.linear_model import LogisticRegression
    from sklearn.metrics import roc_auc_score
    from sklearn.preprocessing import StandardScaler

    sc = StandardScaler().fit(X[tr])
    clf = LogisticRegression(max_iter=2000, C=1.0).fit(sc.transform(X[tr]), y[tr])
    p_cal = clf.predict_proba(sc.transform(X[cal]))[:, 1]
    p_te_raw = clf.predict_proba(sc.transform(X[te]))[:, 1]

    iso = IsotonicRegression(out_of_bounds="clip", y_min=0.001, y_max=0.999)
    iso.fit(p_cal, y[cal])
    p_cal_c = iso.predict(p_cal)
    p_te = iso.predict(p_te_raw)

    # split-conformal abstain band (class-conditional, alpha=0.05):
    #   t_lo = 95th pct of HUMAN probs  -> below it, human is safe to call
    #   t_hi =  5th pct of AI    probs  -> above it, AI is safe to call
    #   between them: INCONCLUSIVE (statistically honest abstain)
    t_lo = float(np.quantile(p_cal_c[y[cal] == 0], 0.95))
    t_hi = float(np.quantile(p_cal_c[y[cal] == 1], 0.05))
    if t_lo >= t_hi:
        # classes so separable the thresholds cross — keep a minimum honest
        # abstain band around the midpoint so borderline real text still
        # gets INCONCLUSIVE rather than a false-confident call
        mid = 0.5 * (t_lo + t_hi)
        t_lo, t_hi = min(t_lo, mid - 0.08), max(t_hi, mid + 0.08)
        t_lo, t_hi = max(0.05, t_lo), min(0.95, t_hi)

    acc = float(((p_te >= 0.5) == y[te]).mean())
    auc = float(roc_auc_score(y[te], p_te))
    decided = (p_te <= t_lo) | (p_te >= t_hi)
    abstain = float(1 - decided.mean())
    err_dec = float(((p_te[decided] >= 0.5) != y[te][decided]).mean()) if decided.any() else 0.0
    print(f"\nTEST: acc={acc:.3f} auc={auc:.3f} | conformal band "
          f"[{t_lo:.3f}, {t_hi:.3f}] abstain={abstain:.1%} "
          f"error-when-decided={err_dec:.3%}")
    print("feature coefs:", dict(zip(FEATURE_ORDER,
                                     np.round(clf.coef_[0], 3).tolist())))

    # piecewise isotonic curve for runtime interpolation
    gx = np.linspace(0, 1, 101)
    meta = {
        "features": FEATURE_ORDER,
        "mean": sc.mean_.tolist(), "scale": sc.scale_.tolist(),
        "coef": clf.coef_[0].tolist(), "intercept": float(clf.intercept_[0]),
        "iso_x": gx.tolist(), "iso_y": iso.predict(gx).tolist(),
        "t_lo": round(t_lo, 4), "t_hi": round(t_hi, 4),
        "test_accuracy": round(acc, 4), "test_auc": round(auc, 4),
        "abstain_rate": round(abstain, 4),
        "error_when_decided": round(err_dec, 5),
        "n_samples": len(rows), "dataset": "HC3",
        "trained": time.strftime("%Y-%m-%d %H:%M"),
    }
    with open(OUT, "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=1)
    print("saved", OUT)


if __name__ == "__main__":
    args = [a for a in sys.argv[1:] if a != "--refit"]
    main(int(args[0]) if args else 300, refit="--refit" in sys.argv)