"""Phase-2 measurement harness: base MiniCPM4.1-8B + bge-small embedder + rules vs Kim-verified
ground truth. Computes per-axis Cohen's kappa, per-category breakdowns (sparse → N/A), confusion
matrices, an OOD report, and embedder boundary-recall vs hand-identified Focus candidates.

Pure metrics (cohen_kappa, confusion) are dependency-free and unit-tested in tests/test_kappa.py.
8B calls go through a disk cache (eval/_cache/preds_8b.jsonl) keyed by sha1(prompt) + a thread pool,
so runs are resumable and Phase 3 (Critical) reuses every cached response.

Credentials (shared OpenBMB hackathon key) are read from the environment, never hardcoded:
  OPENBMB_BASE_URL, OPENBMB_TOKEN
Run:  python -m eval.kappa            # 4 Phase-2 axes (Focus, Technique, Interaction, Input Quality)
      python -m eval.kappa --limit 5 # quick probe on 5 convs (latency/parse check before full run)
"""
from __future__ import annotations

import concurrent.futures as cf
import hashlib
import json
import os
import sys
import threading

import numpy as np

from prompt_card.scoring import observable_axes as OA
from prompt_card.scoring.embedding import FastEmbedder

HERE = os.path.dirname(__file__)
ROOT = os.path.dirname(HERE)
VP = os.path.join(ROOT, "prompt_card", "training", "validation_pool")
CACHE = os.path.join(HERE, "_cache", "preds_8b.jsonl")

TECH = ["zero_shot_role", "few_shot", "thought_generation", "decomposition", "self_criticism", "flipped"]
IQ = ["goal_stated", "specific_specification"]
CE = ["skepticism", "rebuttal", "source_request", "independent_verification", "re_questioning"]


# ---------------- pure metrics (no deps) ----------------

def confusion(y_true, y_pred, labels):
    idx = {l: i for i, l in enumerate(labels)}
    m = [[0] * len(labels) for _ in labels]
    for t, p in zip(y_true, y_pred):
        m[idx[t]][idx[p]] += 1
    return m


def cohen_kappa(y_true, y_pred):
    """Cohen's kappa. Returns None when undefined (empty, or both raters use one identical label
    → chance agreement is 1, kappa degenerate). Caller reports those as N/A."""
    n = len(y_true)
    if n == 0:
        return None
    labels = sorted(set(y_true) | set(y_pred))
    if len(labels) == 1:
        return None  # degenerate: everything one class (e.g. an all-negative sparse category)
    m = confusion(y_true, y_pred, labels)
    po = sum(m[i][i] for i in range(len(labels))) / n
    rows = [sum(r) for r in m]
    cols = [sum(m[i][j] for i in range(len(labels))) for j in range(len(labels))]
    pe = sum((rows[i] / n) * (cols[i] / n) for i in range(len(labels)))
    if pe >= 1.0:
        return None
    return (po - pe) / (1 - pe)


def binary_counts(y_true, y_pred):
    tp = sum(1 for t, p in zip(y_true, y_pred) if t and p)
    fp = sum(1 for t, p in zip(y_true, y_pred) if not t and p)
    fn = sum(1 for t, p in zip(y_true, y_pred) if t and not p)
    tn = sum(1 for t, p in zip(y_true, y_pred) if not t and not p)
    return dict(tp=tp, fp=fp, fn=fn, tn=tn)


def prf(c):
    tp, fp, fn = c["tp"], c["fp"], c["fn"]
    p = tp / (tp + fp) if tp + fp else None
    r = tp / (tp + fn) if tp + fn else None
    f1 = (2 * p * r / (p + r)) if (p and r) else None
    return p, r, f1


# ---------------- cached concurrent 8B ----------------

class CachedClient:
    """Wraps any `generate(prompt)->str` client with a sha1-keyed disk cache + thread pool."""

    def __init__(self, base, workers=8):
        self.base = base
        self.workers = workers
        self._lock = threading.Lock()
        self.cache = {}
        if os.path.exists(CACHE):
            for line in open(CACHE):
                try:
                    d = json.loads(line)
                    self.cache[d["k"]] = d["v"]
                except Exception:
                    pass
        self.hits = 0
        self.misses = 0

    @staticmethod
    def _key(prompt):
        return hashlib.sha1(prompt.encode("utf-8")).hexdigest()

    def _persist(self, k, v):
        with self._lock:
            with open(CACHE, "a") as f:
                f.write(json.dumps({"k": k, "v": v}, ensure_ascii=False) + "\n")

    def run_all(self, prompts):
        """Return {prompt: response} for a list of prompts, using cache + concurrency."""
        uniq = list(dict.fromkeys(prompts))
        todo = [p for p in uniq if self._key(p) not in self.cache]
        self.hits += len(uniq) - len(todo)

        def work(p):
            v = self.base.generate(p)
            k = self._key(p)
            self.cache[k] = v
            self._persist(k, v)
            return p

        if todo:
            with cf.ThreadPoolExecutor(max_workers=self.workers) as ex:
                for i, _ in enumerate(ex.map(work, todo), 1):
                    self.misses += 1
                    if i % 50 == 0:
                        print(f"  ... {i}/{len(todo)} new 8B calls", flush=True)
        return {p: self.cache[self._key(p)] for p in uniq}


# ---------------- data loading ----------------

def load_convs():
    convs = {}
    for fn in ("pool.jsonl", "supplement.jsonl"):
        for line in open(os.path.join(VP, fn)):
            c = json.loads(line)
            convs[c["id"]] = c
    return convs


def user_turns(conv):
    return [t["text"] for t in conv["turns"] if t["role"] == "user"]


def load_gt():
    return [json.loads(l) for l in open(os.path.join(VP, "ground_truth.jsonl"))]


# ---------------- prediction assembly ----------------

def build_prompts(gt, convs, embedder, focus_tmax=0.70, builders=OA,
                  axes=("technique", "input_quality", "interaction", "focus")):
    """Return (prompts, plan, focus_geom). `builders` supplies build_*_prompt (swap for prompt versions);
    `axes` selects which axes to assemble (lets the cascade re-run only changed axes)."""
    prompts, plan = [], []
    focus_geom = {}  # conv_id -> list of (i, cosine) for every adjacent user-turn boundary
    for r in gt:
        cid = r["id"]
        ut = user_turns(convs[cid])
        if "technique" in axes:
            for row in r["technique"]:
                i = int(row["turn"][1:]) - 1
                p = builders.build_technique_prompt(ut[i])
                prompts.append(p); plan.append(("technique", cid, row["turn"], p))
        if "input_quality" in axes:
            for row in r["input_quality"]:
                i = int(row["turn"][1:]) - 1
                p = builders.build_input_quality_prompt(ut[i])
                prompts.append(p); plan.append(("input_quality", cid, row["turn"], p))
        if "interaction" in axes:
            for row in r["interaction"]:
                i = int(row["turn"][1:]) - 1
                p = builders.build_interaction_prompt(ut[i - 1], ut[i])
                prompts.append(p); plan.append(("interaction", cid, row["turn"], p))
        # focus: embedder geometry over ALL adjacent boundaries; 8B only on cos<focus_tmax
        if "focus" in axes and len(ut) >= 2:
            vecs = embedder.embed(ut)
            geom = []
            for i in range(len(ut) - 1):
                cos = float(np.dot(vecs[i], vecs[i + 1]))
                geom.append((i, cos))
                if cos < focus_tmax:
                    p = builders.build_focus_boundary_prompt(ut[i], ut[i + 1])
                    prompts.append(p); plan.append(("focus", cid, f"U{i+1}→U{i+2}", p))
            focus_geom[cid] = geom
    return prompts, plan, focus_geom


# ---------------- axis scoring ----------------

def score_binary_axis(gt, responses, plan_index, axis, fields):
    """For a per-turn multi-binary axis (technique/input_quality): return per-field (yt,yp) + parse fails."""
    per = {f: ([], []) for f in fields}
    parse_fail = 0
    gtmap = {r["id"]: r for r in gt}
    key = "types" if axis in ("technique", "critical") else "features"
    for (ax, cid, turn, prompt) in plan_index[axis]:
        row = next(rr for rr in gtmap[cid][axis] if rr["turn"] == turn)
        pred = OA.parse(responses[prompt], fields)
        if pred is None:
            parse_fail += 1
            pred = {f: False for f in fields}
        for f in fields:
            per[f][0].append(int(f in row[key]))
            per[f][1].append(int(bool(pred.get(f))))
    return per, parse_fail


def score_interaction(gt, responses, plan_index):
    yt, yp = [], []
    parse_fail = 0
    gtmap = {r["id"]: r for r in gt}
    for (ax, cid, turn, prompt) in plan_index["interaction"]:
        row = next(rr for rr in gtmap[cid]["interaction"] if rr["turn"] == turn)
        pred = OA.parse(responses[prompt], ("refinement_attempt",))
        if pred is None:
            parse_fail += 1
            pred = {}
        yt.append(int(bool(row["refinement"])))
        yp.append(int(bool(pred.get("refinement_attempt"))))
    return yt, yp, parse_fail


def score_focus(gt, responses, plan_index, focus_geom, threshold):
    """Binary shift vs not over EVERY adjacent boundary; embedder gates 8B at `threshold`.
    Also returns embedder recall on GT topic_shift boundaries and a 3-class tally on co-identified."""
    rel_pred = {(cid, turn): OA.parse(responses[p], ("relation",)) for (ax, cid, turn, p) in plan_index["focus"]}
    gtmap = {r["id"]: r for r in gt}
    yt, yp = [], []
    gt_shift_total = recall_hit = 0
    cand_total = 0
    three = {}  # (gt_rel, pred_rel) -> count on co-identified candidates
    for r in gt:
        cid = r["id"]
        gt_rel = {f"U{c['a'][1:]}→U{c['b'][1:]}" if False else f"{c['a']}→{c['b']}": c["relation"] for c in r["focus"]}
        for (i, cos) in focus_geom.get(cid, []):
            bkey = f"U{i+1}→U{i+2}"
            gtr = gt_rel.get(bkey)
            is_gt_shift = (gtr == "topic_shift")
            yt.append(int(is_gt_shift))
            is_cand = cos < threshold
            cand_total += int(is_cand)
            pred_shift = 0
            if is_cand:
                pr = rel_pred.get((cid, bkey)) or {}
                prel = pr.get("relation")
                pred_shift = int(prel == "topic_shift")
                if gtr is not None:  # co-identified: both GT and embedder flagged this boundary
                    three[(gtr, prel)] = three.get((gtr, prel), 0) + 1
            yp.append(pred_shift)
            if is_gt_shift:
                gt_shift_total += 1
                recall_hit += int(is_cand)
    recall = recall_hit / gt_shift_total if gt_shift_total else None
    return yt, yp, dict(recall=recall, gt_shift=gt_shift_total, cand_total=cand_total, three=three)


# ---------------- OOD ----------------

def ood_report(gt, convs, embedder):
    ids = [r["id"] for r in gt]
    centroids = []
    for cid in ids:
        ut = user_turns(convs[cid])
        v = embedder.embed(ut)
        c = v.mean(axis=0)
        centroids.append(c / (np.linalg.norm(c) + 1e-9))
    M = np.array(centroids)
    glob = M.mean(axis=0); glob /= np.linalg.norm(glob) + 1e-9
    sims = M @ glob
    mu, sd = float(sims.mean()), float(sims.std())
    flagged = [(ids[i], float(sims[i])) for i in range(len(ids)) if sims[i] < mu - 2 * sd]
    return dict(mu=mu, sd=sd, flagged=sorted(flagged, key=lambda x: x[1]))


# ---------------- main ----------------

def _fmt_k(k):
    return "N/A" if k is None else f"{k:+.3f}"


def main(limit=None, workers=8):
    base_url = os.environ.get("OPENBMB_BASE_URL")
    token = os.environ.get("OPENBMB_TOKEN")
    if not base_url or not token:
        print("ERROR: set OPENBMB_BASE_URL and OPENBMB_TOKEN in the environment (shared hackathon key).",
              file=sys.stderr)
        sys.exit(2)
    from prompt_card.llm.minicpm import MiniCPMClient
    base = MiniCPMClient(base_url, token)

    gt = load_gt()
    if limit:
        gt = gt[:limit]
    convs = load_convs()
    embedder = FastEmbedder()
    print(f"[measure] {len(gt)} convs; embedding + assembling prompts ...", flush=True)
    prompts, plan, focus_geom = build_prompts(gt, convs, embedder)
    plan_index = {}
    for item in plan:
        plan_index.setdefault(item[0], []).append(item)
    print(f"[measure] {len(prompts)} prompts "
          f"(tech {len(plan_index.get('technique',[]))}, iq {len(plan_index.get('input_quality',[]))}, "
          f"inter {len(plan_index.get('interaction',[]))}, focus {len(plan_index.get('focus',[]))})", flush=True)

    client = CachedClient(base, workers=workers)
    responses = client.run_all(prompts)
    print(f"[measure] 8B done (cache hits {client.hits}, new {client.misses})", flush=True)

    # axis scoring
    tech_per, tech_fail = score_binary_axis(gt, responses, plan_index, "technique", TECH)
    iq_per, iq_fail = score_binary_axis(gt, responses, plan_index, "input_quality", IQ)
    inter_yt, inter_yp, inter_fail = score_interaction(gt, responses, plan_index)
    # focus threshold sweep
    sweep = {}
    for T in [0.40, 0.45, 0.50, 0.55, 0.60, 0.65, 0.70]:
        yt, yp, info = score_focus(gt, responses, plan_index, focus_geom, T)
        sweep[T] = (cohen_kappa(yt, yp), info["recall"], info["cand_total"], yt, yp, info)
    # pick T maximizing kappa (tie-break higher recall)
    bestT = max(sweep, key=lambda T: (sweep[T][0] if sweep[T][0] is not None else -9, sweep[T][1] or 0))
    fk, frec, fcand, fyt, fyp, finfo = sweep[bestT]

    write_results(gt, tech_per, tech_fail, iq_per, iq_fail, inter_yt, inter_yp, inter_fail,
                  sweep, bestT, ood_report(gt, convs, embedder))
    print("[measure] wrote eval/measurement_results.md", flush=True)


def write_results(gt, tech_per, tech_fail, iq_per, iq_fail, inter_yt, inter_yp, inter_fail,
                  sweep, bestT, ood):
    L = ["# Phase-2 measurement — base MiniCPM4.1-8B + bge-small + rules vs ground truth", ""]
    L.append(f"Validation set: {len(gt)} conversations. Metric: Cohen's κ (per axis & per category). "
             "Sparse categories (no positives) → κ undefined → **N/A**. Confusion as [TN FP / FN TP].")
    L.append("")

    def axis_block(title, scope, per, fail, fields):
        L.append(f"## {title}")
        L.append(f"_{scope}_  ·  parse failures: {fail}")
        L.append("| category | κ | TN | FP | FN | TP | precision | recall | f1 |")
        L.append("|---|---|---|---|---|---|---|---|---|")
        for f in fields:
            yt, yp = per[f]
            k = cohen_kappa(yt, yp); c = binary_counts(yt, yp); p, r, f1 = prf(c)
            note = "  ← N/A (0 positives)" if sum(yt) == 0 else ""
            L.append(f"| {f} | {_fmt_k(k)}{note} | {c['tn']} | {c['fp']} | {c['fn']} | {c['tp']} | "
                     f"{'-' if p is None else f'{p:.2f}'} | {'-' if r is None else f'{r:.2f}'} | "
                     f"{'-' if f1 is None else f'{f1:.2f}'} |")
        # axis-level "any" binary
        anyt = [int(any(per[f][0][j] for f in fields)) for j in range(len(per[fields[0]][0]))]
        anyp = [int(any(per[f][1][j] for f in fields)) for j in range(len(per[fields[0]][1]))]
        L.append(f"\n**Axis-level (any {title.split()[0].lower()} present): κ = {_fmt_k(cohen_kappa(anyt, anyp))}** "
                 f"(n={len(anyt)} turns, {sum(anyt)} positive).")
        L.append("")

    axis_block("Technique (6 binary)", "first ≤3 user turns/conv", tech_per, tech_fail, TECH)
    axis_block("Input Quality (2 binary)", "first ≤3 user turns/conv", iq_per, iq_fail, IQ)

    L.append("## Interaction (refinement_attempt)")
    L.append(f"_all follow-up turns_  ·  parse failures: {inter_fail}")
    k = cohen_kappa(inter_yt, inter_yp); c = binary_counts(inter_yt, inter_yp); p, r, f1 = prf(c)
    L.append(f"- **κ = {_fmt_k(k)}**  ·  n={len(inter_yt)}, positives={sum(inter_yt)}")
    L.append(f"- confusion [TN {c['tn']} · FP {c['fp']} / FN {c['fn']} · TP {c['tp']}]  ·  "
             f"P {'-' if p is None else f'{p:.2f}'} R {'-' if r is None else f'{r:.2f}'} F1 {'-' if f1 is None else f'{f1:.2f}'}")
    L.append("")

    L.append("## Focus (topic_shift detection, embedder-gated)")
    L.append("Binary shift-vs-not over **every** adjacent user-turn boundary; the embedder gates which "
             "boundaries the 8B classifies (cosine < T). Threshold swept; κ-maximizing T selected.")
    L.append("| T | κ (shift) | embedder recall@T (on GT shifts) | candidates flagged |")
    L.append("|---|---|---|---|")
    for T in sorted(sweep):
        k, rec, cand, *_ = sweep[T]
        mark = " ← selected" if T == bestT else ""
        L.append(f"| {T:.2f} | {_fmt_k(k)} | {'-' if rec is None else f'{rec:.2f}'} | {cand} |{mark}")
    fk, frec, fcand, fyt, fyp, finfo = sweep[bestT]
    cc = binary_counts(fyt, fyp)
    rec_str = "-" if frec is None else f"{frec:.2f}"
    hits = "" if frec is None else f" ({int(round(frec * finfo['gt_shift']))}/{finfo['gt_shift']} GT shifts flagged)"
    L.append("")
    L.append(f"**Selected T={bestT:.2f}: κ = {_fmt_k(fk)}**, embedder boundary-recall = {rec_str}{hits}. "
             f"Confusion [TN {cc['tn']} · FP {cc['fp']} / FN {cc['fn']} · TP {cc['tp']}].")
    if finfo["three"]:
        L.append("\n3-class agreement on co-identified candidates (GT rel → 8B rel):")
        for (g, p2), n in sorted(finfo["three"].items(), key=lambda x: -x[1]):
            L.append(f"- {g} → {p2}: {n}")
    L.append("\n_Embedder boundary-recall is the model-agnostic metric Kim requested: of the boundaries Kim "
             "hand-identified as topic_shift, how many the production embedder surfaces as candidates._")
    L.append("")

    L.append("## OOD report (validation-set self-check)")
    L.append(f"Per-conv mean-user-turn embedding vs the set centroid: μ_sim={ood['mu']:.3f}, σ={ood['sd']:.3f}. "
             f"Flagged (sim < μ−2σ, i.e. atypical conversations): {len(ood['flagged'])}.")
    for cid, s in ood["flagged"][:15]:
        L.append(f"- `{cid[:12]}`  sim={s:.3f}")
    L.append("\n_In production this same distance flags user uploads far from the validation distribution "
             "(low-confidence / out-of-scope scoring)._")
    L.append("")

    with open(os.path.join(HERE, "measurement_results.md"), "w") as f:
        f.write("\n".join(L) + "\n")


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--workers", type=int, default=8)
    args = ap.parse_args()
    main(limit=args.limit, workers=args.workers)