"""Reward-shape sanity check.

Goal: verify the hypothesis that stage2_v7a's CoT collapse is driven by the
*shape* of the reward stack (format=0.4 baseline + all-or-nothing CoT rewards),
not by any reward function being buggy.

Method:
  - Build 4 synthetic completion archetypes per real GT sample:
      A) PERFECT_COT  : <think> with correct <timestep> + Observed/If authentic/
                         Ruled out anchors; <answer> matches GT exactly.
      B) COT_WRONG_ANS: <think> looks right but answer drifts off GT.
      C) CHEAT_NOCOT  : no <think>, <answer> matches GT exactly. ← the trap
      D) NOCOT_WRONG  : no <think>, answer drifts.
  - Also run the rewards on real eval outputs:
      stage1            (short answer, no CoT — looks like archetype C/D)
      stage2_v5         (full CoT — looks like archetype A/B)
  - For each archetype/source, print per-reward stats AND the summed score.
  - Simulate GRPO advantage: treat {A,B,C,D} as 4 rollouts of one prompt,
    normalise, and check whether C gets a positive advantage (= collapse trap).

Usage:
  python reward_sanity.py [--n 200]
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Callable

import numpy as np

REPO = Path("/mnt/local-fast/zhangt/forensics_grpo")
sys.path.insert(0, str(REPO / "src"))

from open_r1.reward import (
    forensics_iou_reward,
    forensics_format_reward,
    forensics_cot_gt_alignment_hungarian_reward,
    forensics_cot_consistency_reward,
    forensics_k_precision_reward,
    forensics_k_recall_reward,
)

# Order matches what stage2_v7a actually used.
REWARDS: list[tuple[str, Callable]] = [
    ("iou",        forensics_iou_reward),
    ("format",     forensics_format_reward),
    ("cot_align",  forensics_cot_gt_alignment_hungarian_reward),
    ("cot_consis", forensics_cot_consistency_reward),
    ("k_prec",     forensics_k_precision_reward),
    ("k_recall",   forensics_k_recall_reward),
]


# ---------------------------------------------------------------------------
# Synthetic completion builders
# ---------------------------------------------------------------------------

_ANCHOR_BLOCK = (
    "Observed: temporal inconsistency detected.\n"
    "If authentic: motion would be smooth.\n"
    "Ruled out: artefacts indicate manipulation."
)


def _think_block(intervals: list[tuple[float, float]]) -> str:
    body = "Reasoning over candidate manipulated regions.\n\n"
    for s, e in intervals:
        body += f"<timestep>{s:.2f} to {e:.2f}</timestep>\n{_ANCHOR_BLOCK}\n\n"
    return f"<think>\n{body}</think>\n"


def _answer_block(intervals: list[tuple[float, float]]) -> str:
    body = "; ".join(f"{s:.2f} to {e:.2f}" for s, e in intervals)
    return f"<answer>{body}</answer>"


def perfect_cot(gt: list[tuple[float, float]]) -> str:
    return _think_block(gt) + _answer_block(gt)


def cot_wrong_ans(gt: list[tuple[float, float]]) -> str:
    # Think has GT; answer shifts by +5s, clamped to >=0
    drifted = [(max(0.0, s - 5.0), max(0.5, e - 5.0)) for s, e in gt]
    return _think_block(gt) + _answer_block(drifted)


def cheat_nocot(gt: list[tuple[float, float]]) -> str:
    return _answer_block(gt)


def nocot_wrong(gt: list[tuple[float, float]]) -> str:
    drifted = [(s + 7.0, e + 7.0) for s, e in gt]
    return _answer_block(drifted)


ARCHETYPES = [
    ("PERFECT_COT",   perfect_cot),
    ("COT_WRONG_ANS", cot_wrong_ans),
    ("CHEAT_NOCOT",   cheat_nocot),
    ("NOCOT_WRONG",   nocot_wrong),
]


# ---------------------------------------------------------------------------
# Scoring
# ---------------------------------------------------------------------------

def score_batch(completions: list[str], solutions: list[list[tuple[float, float]]]):
    """Return dict[reward_name] -> np.ndarray[n_samples] of per-sample scores."""
    out = {}
    # The reward fns themselves do verbose printing — silence them.
    devnull = open(os.devnull, "w")
    old_stdout = sys.stdout
    sys.stdout = devnull
    try:
        for name, fn in REWARDS:
            kwargs = {"solution": solutions} if "solution" in fn.__code__.co_varnames else {}
            if name == "iou":
                kwargs["generator"] = [None] * len(completions)
            vals = fn(completions, **kwargs)
            out[name] = np.array(vals, dtype=float)
    finally:
        sys.stdout = old_stdout
        devnull.close()
    return out


def summarise(name: str, scores: dict[str, np.ndarray]):
    print(f"\n=== {name}  (n={len(next(iter(scores.values())))}) ===")
    header = f"{'reward':<12s} {'mean':>7s} {'std':>7s} {'p10':>7s} {'p50':>7s} {'p90':>7s} {'>0.5':>6s}"
    print(header)
    print("-" * len(header))
    totals = np.zeros_like(next(iter(scores.values())))
    for n, _ in REWARDS:
        v = scores[n]
        totals += v
        print(f"{n:<12s} {v.mean():7.3f} {v.std():7.3f} "
              f"{np.quantile(v,0.1):7.3f} {np.quantile(v,0.5):7.3f} {np.quantile(v,0.9):7.3f} "
              f"{(v>0.5).mean()*100:5.1f}%")
    print(f"{'SUM':<12s} {totals.mean():7.3f} {totals.std():7.3f} "
          f"{np.quantile(totals,0.1):7.3f} {np.quantile(totals,0.5):7.3f} {np.quantile(totals,0.9):7.3f}")
    return totals


# ---------------------------------------------------------------------------
# Real-eval loaders
# ---------------------------------------------------------------------------

def load_eval(eval_dir: Path, limit: int):
    samples = []
    for f in sorted(eval_dir.glob("rank_*.jsonl")):
        with open(f) as fh:
            for line in fh:
                d = json.loads(line)
                if d.get("parse_failed"):
                    continue
                gt = [tuple(g) for g in d["gt"]]
                if not gt:
                    continue
                samples.append((d["output_text"], gt))
                if len(samples) >= limit:
                    return samples
    return samples


# ---------------------------------------------------------------------------
# GRPO advantage simulation
# ---------------------------------------------------------------------------

def grpo_advantage_sim(per_archetype_totals: dict[str, np.ndarray]):
    """For each prompt i, take the 4 archetype totals as 4 rollouts.
    Compute advantage_a = (r_a - mean) / (std + eps). Report mean advantage per
    archetype. The one with the highest mean advantage is what GRPO will
    reinforce — if it's CHEAT_NOCOT, the trap is confirmed.
    """
    eps = 1e-4
    names = [n for n, _ in ARCHETYPES]
    R = np.stack([per_archetype_totals[n] for n in names], axis=1)  # (n_prompts, 4)
    mean = R.mean(axis=1, keepdims=True)
    std = R.std(axis=1, keepdims=True)
    A = (R - mean) / (std + eps)
    print("\n=== GRPO ADVANTAGE SIMULATION (4 archetypes as 4 rollouts/prompt) ===")
    print(f"{'archetype':<14s} {'mean_total':>10s} {'mean_adv':>10s} {'win_rate':>10s}")
    print("-" * 50)
    wins = R.argmax(axis=1)  # which archetype wins each prompt
    for k, n in enumerate(names):
        win_rate = (wins == k).mean() * 100
        print(f"{n:<14s} {R[:, k].mean():10.3f} {A[:, k].mean():10.3f} {win_rate:9.1f}%")

    # The smoking-gun number: what is the advantage gap CHEAT - PERFECT?
    cheat_idx = names.index("CHEAT_NOCOT")
    perf_idx  = names.index("PERFECT_COT")
    gap = (A[:, perf_idx] - A[:, cheat_idx])
    print(f"\nAdvantage(PERFECT_COT) - Advantage(CHEAT_NOCOT):")
    print(f"  mean = {gap.mean():+.3f}   <0 on {(gap<0).mean()*100:.1f}% of prompts")
    print(f"  (<0 means GRPO prefers cheating on that prompt)")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=200, help="samples to draw from each source")
    args = ap.parse_args()

    # Pull GT from stage1 eval (we only need the GT, output_text we'll replace).
    raw = load_eval(REPO / "eval_stage1", args.n)
    if not raw:
        raise SystemExit("No samples found in eval_stage1")
    print(f"loaded {len(raw)} GT samples from eval_stage1\n")

    gts = [g for _, g in raw]

    # 1. Score each synthetic archetype
    archetype_totals = {}
    for name, builder in ARCHETYPES:
        comps = [builder(g) for g in gts]
        scores = score_batch(comps, gts)
        totals = summarise(f"SYNTHETIC  {name}", scores)
        archetype_totals[name] = totals

    # 2. Score real stage1 (short answers, no CoT)
    stage1 = raw[: args.n]
    s1_scores = score_batch([c for c, _ in stage1], [g for _, g in stage1])
    summarise("REAL  stage1 (eval_stage1)", s1_scores)

    # 3. Score real stage2_v5 (CoT with anchors)
    v5_raw = load_eval(REPO / "eval_stage2_counterfactual_v5", args.n)
    if v5_raw:
        v5_scores = score_batch([c for c, _ in v5_raw], [g for _, g in v5_raw])
        summarise("REAL  stage2_v5 (eval_stage2_counterfactual_v5)", v5_scores)

    # 4. GRPO advantage simulation
    grpo_advantage_sim(archetype_totals)

    # 5. REALISTIC head-to-head: pair real stage1 (no-CoT) and real stage2_v5
    # (CoT) outputs that share a sample_id. This is the actual dynamic that
    # decides v7a's collapse: same prompt, two policies, which gets higher
    # advantage in a GRPO group?
    if v5_raw:
        head_to_head(REPO / "eval_stage1", REPO / "eval_stage2_counterfactual_v5",
                     limit=args.n)

    # 6. THE SMOKING-GUN CHECK: v5's reward stack vs v7a's reward stack,
    # applied to the SAME paired outputs.
    weighted_stack_compare(REPO / "eval_stage1", REPO / "eval_stage2_counterfactual_v5",
                           limit=args.n)


def _weighted_sum(scores: dict[str, np.ndarray], weights: dict[str, float]) -> np.ndarray:
    total = np.zeros_like(next(iter(scores.values())))
    for k, w in weights.items():
        total += w * scores[k]
    return total


def weighted_stack_compare(stage1_dir: Path, v5_dir: Path, limit: int):
    """Apply v5's reward stack and v7a's reward stack to the SAME (no-CoT, CoT)
    paired outputs. If the v7a stack flips the sign of (CoT - no-CoT) margin,
    the diagnosis is confirmed.
    """
    # Same loading as head_to_head
    s1, v5 = {}, {}
    for f in sorted(stage1_dir.glob("rank_*.jsonl")):
        for line in open(f):
            d = json.loads(line); s1[d["sample_id"]] = d
    for f in sorted(v5_dir.glob("rank_*.jsonl")):
        for line in open(f):
            d = json.loads(line); v5[d["sample_id"]] = d
    shared = list(set(s1) & set(v5))[:limit]
    if not shared:
        return
    s1_comps = [s1[k]["output_text"] for k in shared]
    v5_comps = [v5[k]["output_text"] for k in shared]
    gts      = [[tuple(g) for g in s1[k]["gt"]] for k in shared]

    s1_scores = score_batch(s1_comps, gts)
    v5_scores = score_batch(v5_comps, gts)

    # cot_consis is computed with weight=1.0 by default in score_batch (env not
    # set); v7a runs with FORENSICS_COT_CONSIS_WEIGHT=0.3 — so multiply post-hoc.
    STACKS = {
        "v5_stack    (iou+format+cot_align+cot_consis@1.0)": {
            "iou": 1.0, "format": 1.0, "cot_align": 1.0, "cot_consis": 1.0,
            "k_prec": 0.0, "k_recall": 0.0,
        },
        "v7a_stack   (iou+format+cot_align+cot_consis@0.3+k_prec+k_recall)": {
            "iou": 1.0, "format": 1.0, "cot_align": 1.0, "cot_consis": 0.3,
            "k_prec": 1.0, "k_recall": 1.0,
        },
        "FIX_A simple revert (v5 stack)": {
            "iou": 1.0, "format": 1.0, "cot_align": 1.0, "cot_consis": 1.0,
            "k_prec": 0.0, "k_recall": 0.0,
        },
        "FIX_B keep k_*, but gate them on <timestep> in <think>": {
            # Simulated as: k_prec/k_recall × cot_align_indicator
            # (cheat mode has no <think>→cot_align=0→k_*=0)
            "iou": 1.0, "format": 1.0, "cot_align": 1.0, "cot_consis": 1.0,
            "k_prec": 0.0, "k_recall": 0.0,  # see special handling below
        },
    }
    print(f"\n=== SMOKING GUN: v5 reward stack vs v7a reward stack ===")
    print(f"(same {len(shared)} paired outputs; only the reward weights differ)")
    for name, w in STACKS.items():
        s1_tot = _weighted_sum(s1_scores, w)
        v5_tot = _weighted_sum(v5_scores, w)
        diff = v5_tot - s1_tot
        cot_wins = (diff > 0).mean() * 100

        # Simulate the early-v7a GRPO group [CoT, CoT, CoT, cheat]
        R = np.stack([v5_tot, v5_tot, v5_tot, s1_tot], axis=1)
        A = (R - R.mean(axis=1, keepdims=True)) / (R.std(axis=1, keepdims=True) + 1e-4)
        cot_adv  = A[:, 0].mean()
        chea_adv = A[:, 3].mean()

        print(f"\n  {name}")
        print(f"    SUM  no-CoT mean: {s1_tot.mean():.3f}")
        print(f"    SUM  CoT    mean: {v5_tot.mean():.3f}")
        print(f"    margin (CoT - no-CoT): {diff.mean():+.3f}   CoT wins {cot_wins:.1f}%")
        print(f"    GRPO advantage in [CoT,CoT,CoT,cheat] group:")
        print(f"      advantage(CoT)    = {cot_adv:+.3f}")
        print(f"      advantage(cheat)  = {chea_adv:+.3f}")

    # --- The v7a step-1 reality check ----------------------------------
    # Per-reward means measured directly from v7a log at step 1 (CoT mode,
    # model just loaded from stage1_decomp_boundary which is no-CoT-trained):
    v7a_step1_cot = {"iou": 0.231, "format": 0.796, "cot_align": 0.154,
                     "cot_consis": 0.201, "k_prec": 0.439, "k_recall": 0.823}
    # And the no-CoT cheat mode, measured from real stage1 outputs:
    cheat_means = {k: float(s1_scores[k].mean()) for k in v7a_step1_cot}
    print(f"\n  === v7a step-1 reality check (per-reward means) ===")
    print(f"  reward       v7a-step1-CoT    stage1-cheat")
    for k in v7a_step1_cot:
        print(f"  {k:<12s} {v7a_step1_cot[k]:14.3f}    {cheat_means[k]:.3f}")
    for name, w in STACKS.items():
        cot_sum   = sum(w[k] * v7a_step1_cot[k] for k in v7a_step1_cot)
        cheat_sum = sum(w[k] * cheat_means[k]   for k in cheat_means)
        # FIX_B special: keep k_prec/k_recall (weight 1.0) but gated by
        # cot_align>0 (= <timestep> in <think>). CoT mode passes gate, cheat fails.
        if name.startswith("FIX_B"):
            cot_sum   += 1.0 * v7a_step1_cot["k_prec"] + 1.0 * v7a_step1_cot["k_recall"]
            # cheat has cot_align=0 → gate fails → k_*=0; no contribution
        winner = "CoT" if cot_sum > cheat_sum else "CHEAT"
        margin = cot_sum - cheat_sum
        # crude "advantage strength" — margin / typical SUM std (~0.6)
        flag = "  ★ STRONG" if margin > 0.4 else ("  ok" if margin > 0 else "  ⚠ COLLAPSE")
        print(f"  {name[:48]:<48s}  CoT={cot_sum:.3f}  cheat={cheat_sum:.3f}  "
              f"margin={margin:+.3f}{flag}")


def head_to_head(stage1_dir: Path, v5_dir: Path, limit: int):
    """Pair real outputs by sample_id; simulate a 4-rollout GRPO group as
    [stage1-style, stage1-style, v5-style, v5-style] (the policy in transition).
    """
    s1 = {}
    for f in sorted(stage1_dir.glob("rank_*.jsonl")):
        for line in open(f):
            d = json.loads(line)
            s1[d["sample_id"]] = d
    v5 = {}
    for f in sorted(v5_dir.glob("rank_*.jsonl")):
        for line in open(f):
            d = json.loads(line)
            v5[d["sample_id"]] = d

    shared = list(set(s1) & set(v5))[:limit]
    print(f"\n=== HEAD-TO-HEAD: paired real outputs on {len(shared)} shared prompts ===")
    if not shared:
        print("  no shared sample_ids")
        return

    s1_comps = [s1[k]["output_text"] for k in shared]
    v5_comps = [v5[k]["output_text"] for k in shared]
    gts      = [[tuple(g) for g in s1[k]["gt"]] for k in shared]

    s1_scores = score_batch(s1_comps, gts)
    v5_scores = score_batch(v5_comps, gts)

    s1_total = sum(s1_scores.values())
    v5_total = sum(v5_scores.values())

    diff = v5_total - s1_total
    print(f"\n  stage1 (no-CoT)    SUM: mean={s1_total.mean():.3f}  std={s1_total.std():.3f}")
    print(f"  stage2_v5 (CoT)    SUM: mean={v5_total.mean():.3f}  std={v5_total.std():.3f}")
    print(f"  CoT - no-CoT per prompt: mean={diff.mean():+.3f}  std={diff.std():.3f}")
    print(f"  CoT wins on {(diff>0).mean()*100:.1f}% of prompts")

    # Simulate 4-rollout GRPO group: 2 of each mode (representative of a
    # policy still exploring both).
    R = np.stack([s1_total, s1_total, v5_total, v5_total], axis=1)
    A = (R - R.mean(axis=1, keepdims=True)) / (R.std(axis=1, keepdims=True) + 1e-4)
    print(f"\n  GRPO group [s1, s1, v5, v5]:")
    print(f"    advantage(no-CoT) mean = {A[:, 0].mean():+.3f}")
    print(f"    advantage(CoT)    mean = {A[:, 2].mean():+.3f}")

    # Now the killer test: when std is tiny (rollouts all on same mode),
    # what does a single deviant rollout do?
    # Group of [s1*3, v5*1] — 3 cheaters + 1 CoT writer
    R2 = np.stack([s1_total, s1_total, s1_total, v5_total], axis=1)
    A2 = (R2 - R2.mean(axis=1, keepdims=True)) / (R2.std(axis=1, keepdims=True) + 1e-4)
    print(f"  GRPO group [s1, s1, s1, v5] (3 cheaters + 1 CoT):")
    print(f"    advantage(no-CoT) mean = {A2[:, 0].mean():+.3f}")
    print(f"    advantage(CoT)    mean = {A2[:, 3].mean():+.3f}")

    # And the reverse: 3 CoT writers + 1 cheater (= early v7a state)
    R3 = np.stack([v5_total, v5_total, v5_total, s1_total], axis=1)
    A3 = (R3 - R3.mean(axis=1, keepdims=True)) / (R3.std(axis=1, keepdims=True) + 1e-4)
    print(f"  GRPO group [v5, v5, v5, s1] (3 CoT + 1 cheater — early v7a state):")
    print(f"    advantage(CoT)    mean = {A3[:, 0].mean():+.3f}")
    print(f"    advantage(no-CoT) mean = {A3[:, 3].mean():+.3f}")


if __name__ == "__main__":
    main()