"""Measure Fast-DetectGPT accuracy on HC3 and fit its logistic calibration.

This is the provenance for the constants in models/text_calib/fastdetectgpt.json.
Re-run it to reproduce or refresh them.

    python scripts/eval_text_detection.py results.json
    python scripts/eval_text_detection.py results.json --write-calibration

Reports AUROC first, because it is threshold-free and calibration-independent:
it answers "does the curvature signal separate human from machine text at all?"
Accuracy at a fixed threshold conflates that question with how good the
calibration is, so it is reported second.

FPR (humans wrongly flagged) is reported separately and broken out by document
length. For an academic-integrity tool a false accusation is the expensive
error, so a single averaged accuracy figure hides the number that matters.

Requires: datasets, scikit-learn (both in requirements.txt). Downloads the HC3
dataset by streaming and the scoring model (default distilgpt2) on first run.
"""

from __future__ import annotations

import argparse
import json
import random
import sys
import time
from pathlib import Path

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score, roc_curve

# Allow `python scripts/eval_text_detection.py` from the repo root: Python puts
# this file's directory on sys.path, not the cwd, so the package is not found.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from text_detection.calibration import sigmoid  # noqa: E402
from text_detection.config import settings  # noqa: E402
from text_detection.detectors.fastdetectgpt import (  # noqa: E402
    _DEFAULT_A,
    _DEFAULT_B,
    FastDetectGPTDetector,
)

MIN_WORDS = 50           # matches the liability firewall's reliability floor
N_PER_CLASS = 300
SEED = 0
TRAIN_FRACTION = 0.5
FPR_TARGETS = (0.0, 0.005, 0.01, 0.02, 0.05)
LENGTH_BUCKETS = ((50, 100), (100, 200), (200, 10**9))


def collect_samples() -> list[tuple[str, int]]:
    """Balanced HC3 sample: label 1 = ChatGPT, 0 = human.

    Takes at most one answer of each class per row, so the human and machine
    answers are responses to the same questions. That pairing controls for
    topic: any separation measured is register, not subject matter.

    Works with datasets>=3 (script-based) or loads raw JSONL files directly
    as fallback (compatible with datasets v4+/v5+).
    """
    try:
        from datasets import load_dataset

        ds = load_dataset("Hello-SimpleAI/HC3", "all", split="train", streaming=True)
        return collect_from_records(ds)
    except (ImportError, RuntimeError):
        pass

    # Fallback: load JSONL files directly from HuggingFace Hub
    import json

    import httpx

    BASE = "https://huggingface.co/datasets/Hello-SimpleAI/HC3/resolve/main"
    JSONL_FILES = [
        "all.jsonl",
        "finance.jsonl",
        "medicine.jsonl",
        "open_qa.jsonl",
        "reddit_eli5.jsonl",
        "wiki_csai.jsonl",
    ]
    rows: list[dict] = []
    for fname in JSONL_FILES:
        url = f"{BASE}/{fname}"
        resp = httpx.get(url, follow_redirects=True, timeout=120)
        resp.raise_for_status()
        for line in resp.text.strip().splitlines():
            line = line.strip()
            if line:
                rows.append(json.loads(line))

    return collect_from_records(rows)


def collect_from_records(rows) -> list[tuple[str, int]]:
    """Pull HC3 answers from dataset rows, balanced by class."""
    import datasets as _ds

    human: list[str] = []
    ai: list[str] = []

    for row in rows:
        if len(human) >= N_PER_CLASS and len(ai) >= N_PER_CLASS:
            break
        if isinstance(row, dict):
            ha = row.get("human_answers") or []
            ca = row.get("chatgpt_answers") or []
        else:
            ha = row["human_answers"] if hasattr(row, "get") else []
            ca = row["chatgpt_answers"] if hasattr(row, "get") else []
        if not isinstance(ha, list):
            try:
                ha = list(ha) if hasattr(ha, "__iter__") else [str(ha)]
            except TypeError:
                ha = []
        if not isinstance(ca, list):
            try:
                ca = list(ca) if hasattr(ca, "__iter__") else [str(ca)]
            except TypeError:
                ca = []
        for txt in ha:
            if isinstance(txt, str) and txt.strip():
                if len(human) < N_PER_CLASS and len(txt.split()) >= MIN_WORDS:
                    human.append(txt.strip())
                    break
        for txt in ca:
            if isinstance(txt, str) and txt.strip():
                if len(ai) < N_PER_CLASS and len(txt.split()) >= MIN_WORDS:
                    ai.append(txt.strip())
                    break

    if len(human) < N_PER_CLASS or len(ai) < N_PER_CLASS:
        print(
            f"  collected {len(human)} human / {len(ai)} AI samples "
            f"(target {N_PER_CLASS} each) — fewer than requested",
            flush=True,
        )

    samples = [(t, 0) for t in human] + [(t, 1) for t in ai]
    random.Random(SEED).shuffle(samples)
    return samples


def metrics(labels: np.ndarray, preds: np.ndarray) -> dict:
    tp = int(((preds == 1) & (labels == 1)).sum())
    tn = int(((preds == 0) & (labels == 0)).sum())
    fp = int(((preds == 1) & (labels == 0)).sum())
    fn = int(((preds == 0) & (labels == 1)).sum())
    n = len(labels)
    return {
        "accuracy": round((tp + tn) / n, 4) if n else 0.0,
        "fpr_humans_wrongly_flagged": round(fp / (fp + tn), 4) if (fp + tn) else 0.0,
        "recall_ai_caught": round(tp / (tp + fn), 4) if (tp + fn) else 0.0,
        "precision": round(tp / (tp + fp), 4) if (tp + fp) else 0.0,
        "tp": tp, "tn": tn, "fp": fp, "fn": fn,
    }


def score_samples(
    samples: list[tuple[str, int]]
) -> tuple[np.ndarray, np.ndarray, np.ndarray, float]:
    """Score every sample, returning (scores, labels, word_counts, seconds)."""
    det = FastDetectGPTDetector()
    scores: list[float] = []
    labels: list[int] = []
    words: list[int] = []

    t0 = time.perf_counter()
    for i, (text, y) in enumerate(samples):
        try:
            res = det._score(text)
        except Exception as exc:
            print(f"  scoring failed on {i}: {exc}", flush=True)
            continue
        scores.append(float(res.raw["discrepancy"]))
        labels.append(y)
        words.append(len(text.split()))
        if (i + 1) % 100 == 0:
            print(f"  scored {i + 1}/{len(samples)} "
                  f"({time.perf_counter() - t0:.0f}s)", flush=True)

    return (
        np.asarray(scores),
        np.asarray(labels),
        np.asarray(words),
        time.perf_counter() - t0,
    )


def operating_points(y: np.ndarray, s: np.ndarray) -> list[dict]:
    """Recall achievable at each target FPR ceiling.

    The decision-relevant view for an integrity tool: pick the false-accusation
    rate you can defend, then read off how much AI text you still catch.
    """
    fpr, tpr, thr = roc_curve(y, s)
    out = []
    for target in FPR_TARGETS:
        i = max(int(np.searchsorted(fpr, target, side="right")) - 1, 0)
        out.append({
            "target_fpr": target,
            "actual_fpr": round(float(fpr[i]), 4),
            "recall_ai_caught": round(float(tpr[i]), 4),
            "raw_threshold": round(float(thr[i]), 4),
        })
    return out


def fpr_by_length(y: np.ndarray, p: np.ndarray, w: np.ndarray) -> list[dict]:
    """FPR among humans only, bucketed by document length."""
    out = []
    for lo, hi in LENGTH_BUCKETS:
        m = (y == 0) & (w >= lo) & (w < hi)
        if not m.sum():
            continue
        out.append({
            "words": f"{lo}-{'+' if hi >= 10**9 else hi}",
            "n_human": int(m.sum()),
            "fpr": round(float((p[m] >= settings.ai_threshold).mean()), 4),
        })
    return out


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("out", nargs="?", default="eval_results.json",
                    help="where to write the JSON report")
    ap.add_argument("--write-calibration", action="store_true",
                    help=f"also write the refitted (a, b) to "
                         f"{settings.calib_dir}/fastdetectgpt.json")
    args = ap.parse_args()

    print("collecting HC3 samples...", flush=True)
    samples = collect_samples()
    n_human = sum(1 for _, y in samples if y == 0)
    print(f"  {len(samples)} samples ({n_human} human / "
          f"{len(samples) - n_human} AI)", flush=True)

    s, y, w, elapsed = score_samples(samples)

    out: dict = {
        "n": int(len(y)),
        "n_human": int((y == 0).sum()),
        "n_ai": int((y == 1).sum()),
        "scoring_model": settings.scoring_model,
        "ai_threshold": settings.ai_threshold,
        "seconds": round(elapsed, 1),
        "ms_per_doc": round(1000 * elapsed / max(len(y), 1), 1),
        "raw_score": {
            "human_mean": round(float(s[y == 0].mean()), 4),
            "ai_mean": round(float(s[y == 1].mean()), 4),
            "human_std": round(float(s[y == 0].std()), 4),
            "ai_std": round(float(s[y == 1].std()), 4),
        },
        # Threshold-free and calibration-independent - report this first.
        "auroc": round(float(roc_auc_score(y, s)), 4),
        "operating_points": operating_points(y, s),
    }

    # Shipped calibration, whatever is currently in effect.
    p_default = np.array([sigmoid(_DEFAULT_A * v + _DEFAULT_B) for v in s])
    out["shipped_calibration"] = {
        "a": _DEFAULT_A, "b": _DEFAULT_B,
        **metrics(y, (p_default >= settings.ai_threshold).astype(int)),
    }

    # Refit on a train split, measure on held-out test. The split is what makes
    # the threshold metrics honest; AUROC would be unchanged by any refit.
    idx = np.arange(len(y))
    np.random.default_rng(SEED).shuffle(idx)
    cut = int(len(idx) * TRAIN_FRACTION)
    tr, te = idx[:cut], idx[cut:]

    lr = LogisticRegression()
    lr.fit(s[tr].reshape(-1, 1), y[tr])
    a_fit = float(lr.coef_[0][0])
    b_fit = float(lr.intercept_[0])

    p_test = np.array([sigmoid(a_fit * v + b_fit) for v in s[te]])
    out["refitted_calibration"] = {
        "a": round(a_fit, 4), "b": round(b_fit, 4),
        "n_train": int(len(tr)), "n_test": int(len(te)),
        "at_ai_threshold": metrics(
            y[te], (p_test >= settings.ai_threshold).astype(int)
        ),
    }

    p_all = np.array([sigmoid(a_fit * v + b_fit) for v in s])
    out["fpr_by_length"] = fpr_by_length(y, p_all, w)

    print("\n=== RESULTS ===")
    print(json.dumps(out, indent=2))
    Path(args.out).write_text(json.dumps(out, indent=2))
    print(f"\nwrote {args.out}")

    if args.write_calibration:
        path = Path(settings.calib_dir) / "fastdetectgpt.json"
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps({"a": round(a_fit, 4), "b": round(b_fit, 4)}))
        print(f"wrote {path}")

    print(
        "\nNOTE: HC3 is an easy benchmark - unedited 2022-era ChatGPT prose against\n"
        "native-English Reddit/expert answers. It does not measure over-flagging of\n"
        "non-native-English writing, nor paraphrased or human-edited AI text.\n"
        "Treat these numbers as an upper bound, not a production expectation."
    )


if __name__ == "__main__":
    main()