#!/usr/bin/env python
"""In-band triage classifier (plan v3 Phase 0a item 4): predict, from statement text alone (CPU baseline) or from
8B hidden states (GPU hook, computed later), whether the base model's pass count on a statement is IN-BAND
(1 <= k <= n-1 accepted samples out of n), i.e. neither never-solved nor always-solved.

numpy only: hashed char n-gram features + a handful of dense statement features -> L2-regularised logistic
regression fit by full-batch Adam. Evaluation: stratified k-fold CV AUC (rank-based, ties averaged) and a
held-out split chosen by a hash of problem_id (not by row order), with problem-level bootstrap intervals, and the
AUC of trivial baselines (statement length) so "the classifier is better than length" is a number, not a feeling.

Cautions recorded with the result (they are properties of the data, not of the code):
  * labels are noisy: k out of 8 samples, so a problem with true p=0.9 is labelled in-band only ~57% of the time;
  * the pool is one source (Goedel-Pset) filtered by the SAME model that labels it, and near-templated statements
    (lean_workbook / orca_math) can sit on both sides of a random split. CV AUC is therefore optimistic for
    OTHER sources: the real test is the transfer to the profiled sample of FormalMATH / Numina / Nemotron.

    .venv/bin/python scripts/p2/supply/triage.py train --problems .tmp/p2/problems/train_pool.jsonl \
        --verdicts .tmp/p2/r1/pool/ver_r0.jsonl --n 8 --out .tmp/p2/supply/triage
    .venv/bin/python scripts/p2/supply/triage.py score --model .tmp/p2/supply/triage/model.npz \
        --problems <set>.jsonl --out <set>.triage.jsonl
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import re
import sys
import zlib
from collections import defaultdict
from pathlib import Path
from typing import Iterable, Optional, Sequence

import numpy as np

REPO = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(REPO))

N_FEATURES = 1 << 14
DENSE_NAMES = ("log_chars", "log_tokens", "n_hyp", "n_digits", "n_quant", "n_real", "n_nat_int", "has_informal", "log_informal")


def _load_common():
    if "p2_common" in sys.modules:
        return sys.modules["p2_common"]
    spec = importlib.util.spec_from_file_location("p2_common", REPO / "scripts" / "p2" / "p2_common.py")
    mod = importlib.util.module_from_spec(spec)
    sys.modules["p2_common"] = mod
    spec.loader.exec_module(mod)
    return mod


# --------------------------------------------------------------------------- labels


def read_pass_counts(verdicts_path: str | Path, n: int = 8) -> dict[str, tuple[int, int]]:
    """problem_id -> (accepted, total) from a verdict JSONL, keeping ONLY problems with exactly `n` distinct
    samples (a problem still being generated has a partial count that would mislabel it). Tolerates a torn last
    line (the file may be appended to by a live chain). Duplicate (problem, sample_idx) keeps the last verdict."""
    seen: dict[str, dict[int, bool]] = defaultdict(dict)
    with open(verdicts_path, encoding="utf-8") as f:
        for line in f:
            try:
                r = json.loads(line)
            except json.JSONDecodeError:
                continue
            seen[r["problem_id"]][int(r["sample_idx"])] = bool(r["accepted"])
    return {p: (sum(v.values()), len(v)) for p, v in seen.items() if len(v) == n}


def in_band(k: int, n: int) -> int:
    return int(1 <= k <= n - 1)


# --------------------------------------------------------------------------- features


def statement_body(statement: str) -> str:
    s = re.sub(r"/-.*?-/", " ", statement, flags=re.S)
    s = re.sub(r"--[^\n]*", "", s)
    s = re.sub(r"^\s*(import|open|set_option)\b[^\n]*$", "", s, flags=re.M)
    s = re.sub(r"\btheorem\s+\S+", "theorem", s, count=1)
    return " ".join(s.replace(":= by sorry", "").split())


def _hash_ngrams(text: str, ns: Sequence[int], salt: str, vec: dict[int, float]) -> None:
    for n in ns:
        for i in range(len(text) - n + 1):
            h = zlib.crc32(f"{salt}{text[i : i + n]}".encode())
            idx, sign = h % N_FEATURES, 1.0 if (h >> 31) & 1 else -1.0
            vec[idx] = vec.get(idx, 0.0) + sign


def featurize(statement: str, informal: Optional[str] = None) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """-> (indices, values, dense). Signed hashing of char 2..4-grams of the statement body plus word unigrams and
    bigrams of the informal text; values are sublinear (sign*log1p|c|) and the sparse part is L2-normalised."""
    body = statement_body(statement)
    vec: dict[int, float] = {}
    _hash_ngrams(body, (2, 3, 4), "s", vec)
    iw = re.findall(r"[a-z0-9]+", (informal or "").lower())
    for i, w in enumerate(iw):
        for g in (w, w + " " + iw[i + 1] if i + 1 < len(iw) else None):
            if g:
                h = zlib.crc32(("w" + g).encode())
                idx = h % N_FEATURES
                vec[idx] = vec.get(idx, 0.0) + (1.0 if (h >> 31) & 1 else -1.0)
    idx = np.fromiter(vec.keys(), dtype=np.int64, count=len(vec))
    val = np.array([np.sign(c) * np.log1p(abs(c)) for c in vec.values()], dtype=np.float64)
    nrm = np.linalg.norm(val)
    if nrm > 0:
        val /= nrm
    toks = body.split()
    dense = np.array([
        np.log1p(len(body)), np.log1p(len(toks)), len(re.findall(r"\(h[\w₀-₉']*\s*:", body)),
        len(re.findall(r"\d", body)) / max(1, len(body)) * 10, len(re.findall(r"[∀∃]", body)),
        body.count("ℝ") + body.count("Real"), body.count("ℕ") + body.count("ℤ"),
        float(bool(informal)), np.log1p(len(iw)),
    ], dtype=np.float64)
    return idx, val, dense


def build_matrix(rows: Sequence[dict]) -> np.ndarray:
    """Dense float32 [n, N_FEATURES + len(DENSE_NAMES)]. 850 rows is ~56 MB; for large sets use `batches`."""
    X = np.zeros((len(rows), N_FEATURES + len(DENSE_NAMES)), dtype=np.float32)
    for r, row in enumerate(rows):
        idx, val, dense = featurize(row["formal_statement"], row.get("informal_prefix"))
        X[r, idx] = val
        X[r, N_FEATURES:] = dense
    return X


def batches(rows: Iterable[dict], size: int = 2048):
    buf: list[dict] = []
    for r in rows:
        buf.append(r)
        if len(buf) >= size:
            yield buf, build_matrix(buf)
            buf = []
    if buf:
        yield buf, build_matrix(buf)


# --------------------------------------------------------------------------- GPU hook


class HiddenStateFeatures:
    """Hook for 8B hidden-state features, computed LATER on GPU (not here: no GPU in this lane).

    Contract: a `.npz` with `ids` (str array of problem_id) and `X` (float32 [n, d]) (or a directory holding
    `ids.json` + `X.npy`, same content) -- e.g. the mean-pooled
    last-layer (or a middle-layer) hidden state of the Goedel-Prover-V2-8B prompt tokens of each statement, one
    forward pass, no sampling. `train --hidden-npz` aligns it to the labelled problems by id and fits the same
    logistic regression on it (alone, or concatenated with the text features via --combine)."""

    def __init__(self, path: str | Path):
        p = Path(path)
        if p.is_dir():
            ids, X = json.loads((p / "ids.json").read_text()), np.load(p / "X.npy", allow_pickle=False)
        else:
            z = np.load(p, allow_pickle=False)
            ids, X = z["ids"], z["X"]
        if len(ids) != len(X):
            raise ValueError(f"{len(ids)} ids but {len(X)} feature rows")
        self.index = {str(i): j for j, i in enumerate(ids)}
        self.X = X.astype(np.float32)

    def matrix(self, ids: Sequence[str]) -> np.ndarray:
        miss = [i for i in ids if i not in self.index]
        if miss:
            raise KeyError(f"{len(miss)} problems have no hidden-state row, e.g. {miss[:3]}")
        return self.X[[self.index[i] for i in ids]]

    @staticmethod
    def compute_on_gpu(*_a, **_k):  # pragma: no cover
        raise NotImplementedError("run on GPU in a heavy-lock window: forward the prompt through the 8B, pool, np.savez(ids=..., X=...)")


# --------------------------------------------------------------------------- model


def auc(y: np.ndarray, s: np.ndarray) -> float:
    """Rank AUC with average ranks for ties (== Mann-Whitney U / (n_pos n_neg))."""
    y = np.asarray(y).astype(bool)
    npos, nneg = int(y.sum()), int((~y).sum())
    if npos == 0 or nneg == 0:
        return float("nan")
    order = np.argsort(s, kind="mergesort")
    ss = np.asarray(s)[order]
    ranks = np.empty(len(s))
    i = 0
    while i < len(ss):
        j = i
        while j + 1 < len(ss) and ss[j + 1] == ss[i]:
            j += 1
        ranks[order[i : j + 1]] = (i + j) / 2 + 1
        i = j + 1
    return float((ranks[y].sum() - npos * (npos + 1) / 2) / (npos * nneg))


class Logistic:
    """L2 logistic regression, full-batch Adam. Features are standardised using training statistics (dense part
    only; the sparse part is already unit-norm). Intercept unpenalised."""

    def __init__(self, l2: float = 1e-2, iters: int = 400, lr: float = 0.05, seed: int = 0):
        self.l2, self.iters, self.lr, self.seed = l2, iters, lr, seed
        self.w: Optional[np.ndarray] = None
        self.b = 0.0
        self.mu: Optional[np.ndarray] = None
        self.sd: Optional[np.ndarray] = None

    def _prep(self, X: np.ndarray) -> np.ndarray:
        return (X - self.mu) / self.sd

    def fit(self, X: np.ndarray, y: np.ndarray, sample_weight: Optional[np.ndarray] = None) -> "Logistic":
        X = X.astype(np.float64)
        self.mu = X.mean(0)
        self.sd = X.std(0)
        self.sd[self.sd < 1e-8] = 1.0
        # sparse hashed columns are not standardised: centring a sparse signed-hash column adds noise
        self.mu[:N_FEATURES] = 0.0
        self.sd[:N_FEATURES] = 1.0
        if X.shape[1] <= N_FEATURES:  # a non-text feature matrix (hidden states): standardise everything
            self.mu, self.sd = X.mean(0), np.maximum(X.std(0), 1e-8)
        Z = self._prep(X)
        n, d = Z.shape
        sw = np.ones(n) if sample_weight is None else sample_weight / sample_weight.mean()
        w, b = np.zeros(d), float(np.log((y.mean() + 1e-3) / (1 - y.mean() + 1e-3)))
        m, v, mb, vb = np.zeros(d), np.zeros(d), 0.0, 0.0
        for t in range(1, self.iters + 1):
            p = 1 / (1 + np.exp(-(Z @ w + b)))
            r = (p - y) * sw
            g = Z.T @ r / n + self.l2 * w
            gb = r.mean()
            m, v = 0.9 * m + 0.1 * g, 0.999 * v + 0.001 * g * g
            mb, vb = 0.9 * mb + 0.1 * gb, 0.999 * vb + 0.001 * gb * gb
            mh, vh = m / (1 - 0.9**t), v / (1 - 0.999**t)
            w -= self.lr * mh / (np.sqrt(vh) + 1e-8)
            b -= self.lr * (mb / (1 - 0.9**t)) / (np.sqrt(vb / (1 - 0.999**t)) + 1e-8)
        self.w, self.b = w, b
        return self

    def decision(self, X: np.ndarray) -> np.ndarray:
        return self._prep(X.astype(np.float64)) @ self.w + self.b

    def proba(self, X: np.ndarray) -> np.ndarray:
        return 1 / (1 + np.exp(-self.decision(X)))


def stratified_folds(y: np.ndarray, k: int, seed: int) -> np.ndarray:
    rs = np.random.RandomState(seed)
    fold = np.zeros(len(y), dtype=int)
    for c in (0, 1):
        idx = np.where(y == c)[0]
        rs.shuffle(idx)
        fold[idx] = np.arange(len(idx)) % k
    return fold


def cross_val_auc(X: np.ndarray, y: np.ndarray, k: int = 5, seed: int = 0, **kw) -> dict:
    fold = stratified_folds(y, k, seed)
    oof = np.zeros(len(y))
    per = []
    for f in range(k):
        tr, te = fold != f, fold == f
        m = Logistic(**kw).fit(X[tr], y[tr])
        oof[te] = m.decision(X[te])
        per.append(auc(y[te], oof[te]))
    return {"fold_auc": [round(a, 4) for a in per], "mean": float(np.mean(per)), "std": float(np.std(per)), "oof_auc": auc(y, oof)}


def bootstrap_auc_ci(y: np.ndarray, s: np.ndarray, B: int = 1000, seed: int = 0) -> tuple[float, float]:
    rs = np.random.RandomState(seed)
    vals = []
    for _ in range(B):
        i = rs.randint(0, len(y), len(y))
        a = auc(y[i], s[i])
        if not np.isnan(a):
            vals.append(a)
    return float(np.percentile(vals, 2.5)), float(np.percentile(vals, 97.5))


def holdout_mask(ids: Sequence[str], frac: float = 0.25, salt: str = "triage-v1") -> np.ndarray:
    """Deterministic by problem_id hash, independent of row order."""
    return np.array([int(hashlib.sha256((salt + i).encode()).hexdigest()[:8], 16) / 0xFFFFFFFF < frac for i in ids])


# --------------------------------------------------------------------------- CLI


def cmd_train(a) -> int:
    C = _load_common()
    probs = {r["problem_id"]: r for r in C.iter_jsonl(a.problems)}
    counts = read_pass_counts(a.verdicts, a.n)
    ids = sorted(i for i in counts if i in probs)
    rows = [probs[i] for i in ids]
    k = np.array([counts[i][0] for i in ids])
    y = np.array([in_band(int(x), a.n) for x in k])
    print(f"labelled problems: {len(ids)} (of {len(counts)} with n={a.n} verdicts, {len(probs)} in pool); "
          f"in-band {int(y.sum())} ({y.mean():.1%}); k histogram {np.bincount(k, minlength=a.n + 1).tolist()}")
    Xt = build_matrix(rows)
    hidden = HiddenStateFeatures(a.hidden_npz).matrix(ids) if a.hidden_npz else None
    X = Xt if hidden is None else (np.hstack([Xt, hidden]) if a.combine else hidden)
    held = holdout_mask(ids)
    res: dict = {"n_labelled": len(ids), "in_band": int(y.sum()), "k_hist": np.bincount(k, minlength=a.n + 1).tolist(),
                 "features": "text_hash_char_ngram" + ("+hidden" if hidden is not None and a.combine else "_hidden" if hidden is not None else ""),
                 "n_train": int((~held).sum()), "n_heldout": int(held.sum()), "heldout_in_band": int(y[held].sum())}
    best = None
    for l2 in (1e-3, 1e-2, 1e-1, 1.0):  # chosen by CV on the TRAIN part only; the held-out part is touched once
        cv = cross_val_auc(X[~held], y[~held], k=5, seed=0, l2=l2)
        res.setdefault("cv_by_l2", {})[str(l2)] = round(cv["oof_auc"], 4)
        if best is None or cv["oof_auc"] > best[1]["oof_auc"]:
            best = (l2, cv)
    l2, cv = best
    res["chosen_l2"], res["train_cv"] = l2, cv
    m = Logistic(l2=l2).fit(X[~held], y[~held])
    s = m.decision(X[held])
    lo, hi = bootstrap_auc_ci(y[held], s)
    res["heldout_auc"] = round(auc(y[held], s), 4)
    res["heldout_auc_ci95"] = [round(lo, 4), round(hi, 4)]
    # baselines: statement length, and "informal text length" -- a classifier that cannot beat these has learnt nothing
    length = np.array([len(statement_body(r["formal_statement"])) for r in rows], dtype=float)
    res["baseline_heldout_auc"] = {"statement_length": round(auc(y[held], length[held]), 4),
                                   "neg_statement_length": round(auc(y[held], -length[held]), 4)}
    # secondary targets: never solved vs not, always solved vs not (what the band filter actually removes)
    for name, tgt in (("solved_at_least_once", (k >= 1).astype(int)), ("always_solved", (k == a.n).astype(int))):
        mm = Logistic(l2=l2).fit(X[~held], tgt[~held])
        res.setdefault("secondary_heldout_auc", {})[name] = round(auc(tgt[held], mm.decision(X[held])), 4)
    final = Logistic(l2=l2).fit(X, y)  # refit on all labels for scoring
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    np.savez(out / "model.npz", w=final.w, b=final.b, mu=final.mu, sd=final.sd, n_features=N_FEATURES)
    (out / "triage_report.json").write_text(json.dumps(res, indent=1))
    print(json.dumps(res, indent=1))
    return 0


def cmd_score(a) -> int:
    C = _load_common()
    z = np.load(a.model)
    m = Logistic()
    m.w, m.b, m.mu, m.sd = z["w"], float(z["b"]), z["mu"], z["sd"]
    n = 0
    with open(a.out, "w", encoding="utf-8") as f:
        for rows, X in batches(C.iter_jsonl(a.problems)):
            for r, p in zip(rows, m.proba(X)):
                f.write(json.dumps({"problem_id": r["problem_id"], "p_in_band": round(float(p), 5)}) + "\n")
            n += len(rows)
    print(f"scored {n} -> {a.out}")
    return 0


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    t = sub.add_parser("train")
    t.add_argument("--problems", required=True)
    t.add_argument("--verdicts", required=True)
    t.add_argument("--n", type=int, default=8)
    t.add_argument("--out", required=True)
    t.add_argument("--hidden-npz", help="GPU features (ids, X); see HiddenStateFeatures")
    t.add_argument("--combine", action="store_true", help="concatenate hidden-state and text features")
    t.set_defaults(fn=cmd_train)
    s = sub.add_parser("score")
    s.add_argument("--model", required=True)
    s.add_argument("--problems", required=True)
    s.add_argument("--out", required=True)
    s.set_defaults(fn=cmd_score)
    a = ap.parse_args(argv)
    return a.fn(a)


if __name__ == "__main__":
    raise SystemExit(main())