"""Train the modern, model-agnostic AI-text detector and validate it the honest
way: LEAVE-ONE-MODEL-OUT (does it catch a generator it never trained on?) plus
an external test on the 13 real Turnitin papers.

Features: plagdetect/aifeatures (surface regularities, torch-free). Model: a
regularised logistic regression (18 features, ~1900 docs => low overfit), Platt
-calibrated. Saved to models/ai_modern.json for aidetect to consume.

Run: python scripts/train_modern_ai.py
"""
import json, os, sys, time
import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from plagdetect.aifeatures import FEATURES, vector  # noqa: E402
from plagdetect.webpipeline import extract_text, split_body_references  # noqa: E402

from sklearn.linear_model import LogisticRegression  # noqa: E402
from sklearn.preprocessing import StandardScaler  # noqa: E402
from sklearn.metrics import roc_auc_score  # noqa: E402

CORPUS = os.path.join(ROOT, "data", "calibration", "mage_sci.jsonl")
SCORED = os.path.join(ROOT, "data", "calibration", "mage_sci_scored.jsonl")
GT = os.path.join(ROOT, "data", "turnitin_groundtruth.json")
DSET = os.path.join(ROOT, "DATASET FOR training of turnitin")
OUT = os.path.join(ROOT, "models", "ai_modern.json")
MIN_AI_PER_MODEL = 25       # only validate against models with enough samples


def score_corpus():
    if os.path.exists(SCORED):
        return [json.loads(l) for l in open(SCORED, encoding="utf-8")]
    rows = [json.loads(l) for l in open(CORPUS, encoding="utf-8")]
    out = []
    for i, r in enumerate(rows):
        out.append({"x": vector(r["text"]), "y": r["y"], "g": r["src_model"]})
        if i % 200 == 0:
            print(f"  scored {i}/{len(rows)}", end="\r")
    with open(SCORED, "w", encoding="utf-8") as f:
        for r in out:
            f.write(json.dumps(r) + "\n")
    print(f"\nscored {len(out)} -> {SCORED}")
    return out


def fit(Xtr, ytr, C=0.3):
    sc = StandardScaler().fit(Xtr)
    clf = LogisticRegression(max_iter=5000, C=C, class_weight="balanced")
    clf.fit(sc.transform(Xtr), ytr)
    return sc, clf


def main():
    data = score_corpus()
    X = np.array([d["x"] for d in data], float)
    y = np.array([d["y"] for d in data], int)
    g = np.array([d["g"] for d in data])

    ai_models = sorted({m for m in g if m != "human"})
    big = [m for m in ai_models if (g == m).sum() >= MIN_AI_PER_MODEL]
    hum_idx = np.where(y == 1)[0]
    rng = np.random.RandomState(13)
    rng.shuffle(hum_idx)
    hum_folds = np.array_split(hum_idx, len(big))

    # ---- leave-one-MODEL-out: train without model m, test human vs m --------
    print("LEAVE-ONE-MODEL-OUT (held-out generator never seen in training)")
    print("-" * 60)
    aucs = []
    for i, m in enumerate(big):
        te_h = hum_folds[i]
        te_ai = np.where(g == m)[0]
        te = np.concatenate([te_h, te_ai])
        tr = np.array([j for j in range(len(y))
                       if g[j] != m and j not in set(te_h)])
        sc, clf = fit(X[tr], y[tr])
        p = clf.predict_proba(sc.transform(X[te]))[:, 1]   # p(human)
        auc = roc_auc_score(y[te], p)
        aucs.append(auc)
        print(f"  hold {m:24s} n_ai={len(te_ai):3d}  AUC={auc:.3f}")
    print(f"\nMEAN leave-one-model-out AUC = {np.mean(aucs):.3f}  "
          f"(min {np.min(aucs):.3f})")

    # ---- final model on all + Platt calibration -----------------------------
    idx = rng.permutation(len(y))
    n_tr = int(0.8 * len(y))
    tr, cal = idx[:n_tr], idx[n_tr:]
    sc, clf = fit(X[tr], y[tr])
    p_cal_raw = clf.predict_proba(sc.transform(X[cal]))[:, 1]
    platt = LogisticRegression(max_iter=5000, C=1e6)
    platt.fit(p_cal_raw.reshape(-1, 1), y[cal])
    # NOTE: model predicts p(human); we store p_ai = 1 - p(human)
    gx = np.linspace(0, 1, 101)
    gy = platt.predict_proba(gx.reshape(-1, 1))[:, 1]

    meta = {"features": FEATURES,
            "mean": sc.mean_.tolist(), "scale": sc.scale_.tolist(),
            "coef": clf.coef_[0].tolist(), "intercept": float(clf.intercept_[0]),
            "platt_x": gx.tolist(), "platt_y": gy.tolist(),
            "predicts": "p_human",
            "lomo_auc": round(float(np.mean(aucs)), 4),
            "lomo_min": round(float(np.min(aucs)), 4),
            "n": len(y), "models": big,
            "dataset": "MAGE/sci_gen", "trained": time.strftime("%Y-%m-%d %H:%M")}
    json.dump(meta, open(OUT, "w", encoding="utf-8"), indent=1)
    print(f"saved {OUT}")

    # ---- external test on the 13 real Turnitin papers -----------------------
    print("\nEXTERNAL TEST — 13 real papers (Turnitin AI% as noisy reference)")
    print("-" * 60)
    gt = json.load(open(GT, encoding="utf-8"))
    print(f"{'draft':30s} {'turnitin':>8s} {'p_ai':>6s}")
    yy, pp = [], []
    for rec in gt:
        draft, ai = rec.get("draft"), (rec.get("ai") or {}).get("ai_pct")
        if not draft or ai is None:
            continue
        _t, text = extract_text(os.path.join(DSET, draft))
        body, _ = split_body_references(text)
        x = (np.array(vector(body)) - sc.mean_) / sc.scale_
        p_h = clf.predict_proba(x.reshape(1, -1))[:, 1][0]
        p_ai = 1 - p_h
        lbl = "AI" if (ai != "*" and ai >= 40) else ("hum" if ai == "*" else "?")
        print(f"{draft[:29]:30s} {str(ai):>8s} {p_ai:6.2f}  {lbl}")
        if lbl in ("AI", "hum"):
            yy.append(1 if lbl == "AI" else 0)
            pp.append(p_ai)
    if len(set(yy)) == 2:
        print(f"\nexternal AUC (AI>=40 vs suppressed<20) = "
              f"{roc_auc_score(yy, pp):.3f}  on {len(yy)} papers")


if __name__ == "__main__":
    main()