#!/usr/bin/env python3
"""Held-out transitive binding probe for ConeML checkpoints.

The existing chat activation probe used the same small name pool and a near
identical prompt template as the focused SFT retention rows. This probe splits
that apart by evaluating train-template/new-name, unseen-query/new-name,
unseen-relation, and non-name entity variants under raw and chat surfaces.
"""

from __future__ import annotations

import os

# WSL ROCm: force the HSA /dev/dxg detection path so the 7900 XT is visible when
# this script is run via .venv/bin/python directly (bypassing venv activate).
# Must be set before `import torch` initializes HIP. See CLAUDE.md.
os.environ.setdefault("HSA_ENABLE_DXG_DETECTION", "1")

import argparse
import json
import random
import re
import time
from collections import Counter, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch
from transformers import AutoTokenizer

import sys

sys.path.insert(0, "scripts")
sys.path.insert(0, "src")

from pretrain_corpus import PretrainConfig, build_model  # noqa: E402
from probe_full_corpus_cpu import gen, wilson  # noqa: E402


TRAIN_NAMES = ["Anna", "Ben", "Cara", "Dana", "Evan", "Finn", "Gina", "Hugo"]
HELDOUT_NAMES = ["Isabel", "Jonas", "Keira", "Liam", "Maya", "Noah", "Olivia", "Priya", "Quinn", "Rosa", "Sofia", "Theo"]
HELDOUT_ENTITIES = [
    "red card",
    "blue card",
    "green card",
    "silver card",
    "black card",
    "white card",
    "gold card",
    "gray card",
    "orange card",
    "purple card",
]


@dataclass(frozen=True)
class Suite:
    name: str
    pool: tuple[str, ...]
    relation: str
    high_word: str
    low_word: str
    high_query: str
    low_query: str
    answer_label: str
    completion_template: str
    sft_like: bool = False


SUITES = [
    Suite(
        name="sft_template_heldout_names",
        pool=tuple(HELDOUT_NAMES),
        relation="is taller than",
        high_word="tallest",
        low_word="shortest",
        high_query="{chain} Of all of them, the tallest is who? Return only the name.",
        low_query="{chain} Of all of them, the shortest is who? Return only the name.",
        completion_template="{chain} Of all of them, the {target} is",
        answer_label="name",
        sft_like=True,
    ),
    Suite(
        name="unseen_query_heldout_names",
        pool=tuple(HELDOUT_NAMES),
        relation="is taller than",
        high_word="tallest",
        low_word="shortest",
        high_query="Given these facts: {chain} Which person is highest in the height order? Answer with only the name.",
        low_query="Given these facts: {chain} Which person is lowest in the height order? Answer with only the name.",
        completion_template="Given these facts: {chain} The person highest in the height order is",
        answer_label="name",
    ),
    Suite(
        name="older_relation_heldout_names",
        pool=tuple(HELDOUT_NAMES),
        relation="is older than",
        high_word="oldest",
        low_word="youngest",
        high_query="{chain} Which person is oldest? Return only the name.",
        low_query="{chain} Which person is youngest? Return only the name.",
        completion_template="{chain} The {target} person is",
        answer_label="name",
    ),
    Suite(
        name="before_relation_entities",
        pool=tuple(HELDOUT_ENTITIES),
        relation="comes before",
        high_word="first",
        low_word="last",
        high_query="{chain} Which item comes first? Return only the item.",
        low_query="{chain} Which item comes last? Return only the item.",
        completion_template="{chain} The item that comes {target} is",
        answer_label="item",
    ),
]


def load_model(ckpt: Path, config: Path, tokenizer: str, device: str):
    cfg_d = json.load(config.open("r", encoding="utf-8"))
    cfg = PretrainConfig(**{k: v for k, v in cfg_d.items() if k in PretrainConfig.__dataclass_fields__})
    model = build_model(cfg, device)
    payload = torch.load(ckpt, map_location="cpu", weights_only=False)
    model.load_state_dict(payload["model"] if "model" in payload else payload)
    model.to(device)
    model.eval()
    tok = AutoTokenizer.from_pretrained(tokenizer)
    return model, tok, payload


def chat_prompt(user: str) -> str:
    return f"User:\n{user.strip()}\nAssistant:\n"


def make_chain(items: list[str], relation: str) -> str:
    return " ".join(f"{items[i]} {relation} {items[i + 1]}." for i in range(len(items) - 1))


def first_choice(generated: str, choices: list[str]) -> str:
    text = generated.lower()
    hits = []
    for choice in sorted(choices, key=len, reverse=True):
        m = re.search(rf"(?<![a-z]){re.escape(choice.lower())}(?![a-z])", text)
        if m:
            hits.append((m.start(), choice))
    if hits:
        return sorted(hits)[0][1]
    return ""


def answer_anywhere(generated: str, gold: str) -> bool:
    return re.search(rf"(?<![a-z]){re.escape(gold.lower())}(?![a-z])", generated.lower()) is not None


def build_rows(suite: Suite, depth: int, n: int, seed: int) -> list[dict[str, str]]:
    rng = random.Random(seed)
    rows = []
    for _ in range(max(1, n // 2)):
        items = rng.sample(list(suite.pool), depth + 1)
        chain = make_chain(items, suite.relation)
        rows.append({
            "type": "high",
            "target_word": suite.high_word,
            "chat_user": suite.high_query.format(chain=chain),
            "completion": suite.completion_template.format(chain=chain, target=suite.high_word),
            "gold": items[0],
            "chain": chain,
        })
        rows.append({
            "type": "low",
            "target_word": suite.low_word,
            "chat_user": suite.low_query.format(chain=chain),
            "completion": suite.completion_template.format(chain=chain, target=suite.low_word),
            "gold": items[-1],
            "chain": chain,
        })
    return rows[:n]


@torch.no_grad()
def probe_one(model, tok, suite: Suite, n_per_depth: int, max_new: int, seed: int, progress: int) -> dict[str, Any]:
    report: dict[str, Any] = {}
    for depth in (1, 2, 3, 4, 5):
        rows = build_rows(suite, depth, n_per_depth, seed + depth * 1009)
        by_surface: dict[str, Counter] = defaultdict(Counter)
        by_surface_type: dict[str, dict[str, Counter]] = defaultdict(lambda: defaultdict(Counter))
        examples: dict[str, list[dict[str, Any]]] = {"raw_completion": [], "chat": []}
        for row_i, row in enumerate(rows, start=1):
            prompts = {
                "raw_completion": row["completion"],
                "chat": chat_prompt(row["chat_user"]),
            }
            for surface, prompt in prompts.items():
                stop = ["\n", ".", "User:", "Assistant:"] if surface == "chat" else ["\n", "."]
                generated = gen(model, tok, prompt, max_new=max_new, temp=0.0, stop=stop)
                pred = first_choice(generated, list(suite.pool))
                first_ok = pred == row["gold"]
                anywhere = answer_anywhere(generated, row["gold"])
                by_surface[surface]["N"] += 1
                by_surface[surface]["first_ok"] += int(first_ok)
                by_surface[surface]["anywhere"] += int(anywhere)
                by_surface_type[surface][row["type"]]["N"] += 1
                by_surface_type[surface][row["type"]]["first_ok"] += int(first_ok)
                by_surface_type[surface][row["type"]]["anywhere"] += int(anywhere)
                if len(examples[surface]) < 8:
                    examples[surface].append({
                        "type": row["type"],
                        "prompt": prompt[:320],
                        "gold": row["gold"],
                        "generated": generated[:160],
                        "first_choice": pred,
                        "first_ok": first_ok,
                        "answer_anywhere": anywhere,
                    })
            if progress and (row_i % progress == 0 or row_i == len(rows)):
                print(f"[heldout-transitive] {suite.name} depth_{depth} {row_i}/{len(rows)}", flush=True)

        depth_out: dict[str, Any] = {
            "N": len(rows),
            "chance": 1 / (depth + 1),
            "sft_like_template": suite.sft_like,
            "surfaces": {},
        }
        for surface, counts in by_surface.items():
            n = counts["N"]
            by_type = {}
            for typ, typ_counts in by_surface_type[surface].items():
                tn = typ_counts["N"]
                by_type[typ] = {
                    "N": tn,
                    "first_choice_accuracy": typ_counts["first_ok"] / max(1, tn),
                    "answer_anywhere_rate": typ_counts["anywhere"] / max(1, tn),
                    "ci95_first_choice": wilson(typ_counts["first_ok"], tn),
                }
            depth_out["surfaces"][surface] = {
                "N": n,
                "first_choice_accuracy": counts["first_ok"] / max(1, n),
                "answer_anywhere_rate": counts["anywhere"] / max(1, n),
                "ci95_first_choice": wilson(counts["first_ok"], n),
                "by_type": by_type,
                "examples": examples[surface],
            }
        report[f"depth_{depth}"] = depth_out
    return report


def summarize_suite(suite_report: dict[str, Any]) -> dict[str, Any]:
    summary = {}
    for surface in ("raw_completion", "chat"):
        vals = []
        for depth, row in suite_report.items():
            surf = row["surfaces"][surface]
            vals.append((depth, surf["first_choice_accuracy"], surf["answer_anywhere_rate"]))
        summary[surface] = {
            depth: {"first_choice_accuracy": acc, "answer_anywhere_rate": anyr}
            for depth, acc, anyr in vals
        }
    return summary


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", type=Path, action="append", required=True)
    ap.add_argument("--label", action="append", default=[])
    ap.add_argument("--config", type=Path, default=Path("config.json"))
    ap.add_argument("--tokenizer", default="models/tokenizers/v9_67m_32k")
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--n-per-depth", type=int, default=128)
    ap.add_argument("--max-new", type=int, default=16)
    ap.add_argument("--progress-every", type=int, default=0)
    args = ap.parse_args()

    labels = args.label or [p.stem for p in args.ckpt]
    if len(labels) != len(args.ckpt):
        raise SystemExit("--label count must match --ckpt count")

    report = {
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "probe": "heldout_transitive_probe",
        "config": str(args.config),
        "tokenizer": args.tokenizer,
        "n_per_depth": args.n_per_depth,
        "suites": [s.name for s in SUITES],
        "checkpoints": {},
    }
    for label, ckpt in zip(labels, args.ckpt):
        model, tok, payload = load_model(ckpt, args.config, args.tokenizer, args.device)
        ckpt_report = {
            "ckpt": str(ckpt),
            "step": payload.get("step"),
            "device": args.device,
            "suites": {},
            "summary": {},
        }
        for suite_i, suite in enumerate(SUITES):
            suite_report = probe_one(
                model,
                tok,
                suite,
                args.n_per_depth,
                args.max_new,
                seed=20260623 + suite_i * 10000,
                progress=args.progress_every,
            )
            ckpt_report["suites"][suite.name] = suite_report
            ckpt_report["summary"][suite.name] = summarize_suite(suite_report)
        report["checkpoints"][label] = ckpt_report
        del model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(report, indent=2, ensure_ascii=False, sort_keys=True) + "\n", encoding="utf-8")
    compact = {
        "out": str(args.out),
        "checkpoints": {
            label: ckpt_report["summary"] for label, ckpt_report in report["checkpoints"].items()
        },
    }
    print(json.dumps(compact, indent=2, ensure_ascii=False, sort_keys=True))


if __name__ == "__main__":
    main()