"""Gradient-side per-module v_hack extraction (spec.md §B, top-k variant).

We sample the per-module GRPO update direction on labeled (hack, clean) pairs.
For a pair with advantages (adv_h=+1, adv_c=-1) the Dr.GRPO single-step grad
`-adv_h * grad_logp(hack) - adv_c * grad_logp(clean)` algebraically equals
`grad_NLL(hack) - grad_NLL(clean)`, so we compute it by the simpler path:
forward each completion, take mean-NLL on completion tokens, backward, and
capture `delta_S.grad` per AntiPaSTO-wrapped Linear. Naming the steps NLL is
an implementation detail; the *meaning* is "the GRPO update on this pair."

Then per module, with D = [g_hack_i - g_clean_i for each pair] in R^{n_pairs x r}:
    SVD(D) = U Σ Vh
    v_hack[name] = top_k rows of Vh, each oriented so mean(D @ v_i) > 0

This generalizes mean-diff (which corresponds to top-1 PC of paired diffs under
isotropic covariance) to a rank-k hack subspace, motivated by CHaRS (Abdullaev
2025 -- see docs/paper_chars.md): hack signal is multi-modal across hack flavors
(weak tests, hardcode, persona, ...), so a single global direction is brittle.

Orientation matters because proj.py applies a per-direction one-sided gate
(only subtracts <g, v_i> when positive). +v_i must point hack-ward.

Saves `out/v_hack.safetensors` = dict[name -> Tensor[k, r]] (cpu fp32, rows
unit-norm + orthonormal from SVD) with header {"model": str, "dtype": str,
"top_k": str(k)}.

Run: uv run python -m vgrout.extract_vhack_grad
"""
from __future__ import annotations

import sys
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path

import torch
import tyro
from jaxtyping import Float
from loguru import logger
from safetensors import safe_open
from safetensors.torch import save_file
from tabulate import tabulate
from transformers import AutoModelForCausalLM, AutoTokenizer

from .antipasto import wrap_model_with_antipasto
from .pairs import PAIRS
from .pairs_from_pool import load_pairs_json


CACHE_ROOT = Path("svd_cache")
OUT_DIR = Path("out")


@dataclass
class Config:
    model: str = "Qwen/Qwen3-4B"
    dtype: str = "bf16"  # must match train.py, else SVD basis cache can differ silently
    out_path: Path = OUT_DIR / "vhack" / "v_hack.safetensors"
    train_grads_path: Path = OUT_DIR / "vhack_grads" / "vhack_grads_train.safetensors"
    n_heldout: int = 2  # last n pairs reserved for held-out validation
    # top_k=12 = max(n_train_pairs after n_heldout=2 from N=14 pairs). Extract once
    # at max rank; train.py slices via --v-hack-k for k-ablation without re-extract.
    top_k: int = 12
    # tau_axis: zero rows where S_i/S_0 < tau_axis. Diagnostic -- projection along
    # noise-direction unit vectors removes only ~||g||/sqrt(r) ≈ 2% of grad
    # magnitude on r=2560 modules, so this rarely changes effect size; it does
    # make k-ablations honest (axes 4-5 might be pure noise on N=12 pairs).
    tau_axis: float = 0.0
    # Override the hand-crafted PAIRS list with pool-derived pairs (see
    # pairs_from_pool.py). Path to a JSON file with list[HackPair-as-dict].
    # When set, hand-crafted PAIRS are NOT loaded -- this lets us extract
    # v_hack from a half-A-only set of hacks to test cross-mechanism
    # generalisation (docs/spec/20260528_cross_mechanism_v_hack.md).
    pairs_from_pool: Path | None = None


def resolve_dtype(s: str) -> torch.dtype:
    return {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[s]


def completion_nll(model, tokenizer, prompt: str, completion: str, device) -> torch.Tensor:
    """Mean NLL over completion tokens only (length-normalized)."""
    prompt_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    full_ids = tokenizer(prompt + completion, return_tensors="pt").input_ids.to(device)
    n_prompt = prompt_ids.shape[1]
    logits = model(full_ids).logits[:, :-1]            # [1, L-1, V]
    targets = full_ids[:, 1:]                          # [1, L-1]
    logp = torch.nn.functional.log_softmax(logits.float(), dim=-1)
    nll = -logp.gather(-1, targets.unsqueeze(-1)).squeeze(-1)  # [1, L-1]
    # mask: positions whose target is a completion token (i.e. index >= n_prompt in full_ids)
    pos = torch.arange(full_ids.shape[1] - 1, device=device).unsqueeze(0)
    mask = (pos >= (n_prompt - 1)).float()
    return (nll * mask).sum() / mask.sum().clamp_min(1.0)


def extract_v_hack(
    model,
    tokenizer,
    wrappers: dict,
    pairs: list,
    top_k: int,
    tau_axis: float,
    n_heldout: int,
    device,
) -> tuple[
    dict[str, Float[torch.Tensor, "k r"]],
    dict[str, Float[torch.Tensor, "k"]],
    dict[str, Float[torch.Tensor, "n_pairs r"]],
    list[dict],
]:
    """Run pair-grads + per-module SVD on D = g_hack - g_clean, return v_hack.

    Pure function -- caller owns model loading, wrapping, and saving. train.py
    calls this on its already-wrapped model when v_hack cache is missing, so
    we don't pay the cost of a second model load.

    Returns:
        v_hack: dict[name -> Tensor[k, r]] (cpu fp32), top-k right singular
            vectors of D per module, oriented so mean(D @ v_i) > 0. If
            tau_axis > 0, rows where S_i/S_0 < tau_axis are zeroed.
        v_sv: dict[name -> Tensor[k]] (cpu fp32), singular values matching v_hack.
            Saved alongside V under `_sv/{name}` keys so load-time noise-floor
            filtering works without re-extracting.
        raw_grads: dict["hack/name"|"clean/name" -> Tensor[n_pairs, r]] for
            offline analysis (verify_vhack_heldout reads this).
        diag_rows: per-module diagnostic dicts (sv_top frac, ||D||, etc.).
    """
    train_pairs = pairs[:-n_heldout] if n_heldout > 0 else pairs
    n_pairs = len(train_pairs)

    grads_hack: dict[str, list[torch.Tensor]] = defaultdict(list)
    grads_clean: dict[str, list[torch.Tensor]] = defaultdict(list)

    for pi, pair in enumerate(train_pairs):
        for label, completion in (("hack", pair.hack), ("clean", pair.clean)):
            model.zero_grad(set_to_none=True)
            loss = completion_nll(model, tokenizer, pair.prompt, completion, device)
            if not torch.isfinite(loss):
                # Skip-on-nonfinite would silently leave G_h and G_c with mismatched
                # lengths and explode later in D = G_h - G_c. Fail fast.
                raise RuntimeError(f"non-finite loss at pair={pi} label={label}: {loss.item()}")
            loss.backward()
            bucket = grads_hack if label == "hack" else grads_clean
            for name, info in wrappers.items():
                layer = info["layer"]
                if getattr(layer, "_lora_grad_probe", False) and layer._lora_h is not None:
                    # LoRA-frozen-B: the routing handle is the r-bottleneck gradient
                    # g_h = B^T δ_y (B frozen -> static path), not A.grad. Sum over (batch,
                    # tokens) to mirror how AntiPaSTO's δS.grad accumulates over positions.
                    gh = layer._lora_h.grad
                    if gh is None:
                        raise RuntimeError(f"no bottleneck grad on {name}; aborting LoRA extract")
                    g = gh.sum(dim=tuple(range(gh.dim() - 1)))   # [r]
                else:
                    g = info["delta_S"].grad
                    if g is None:
                        raise RuntimeError(f"no grad on {name}; aborting extract")
                bucket[name].append(g.detach().float().cpu().clone())
        if (pi + 1) % 5 == 0:
            logger.info(f"  pair {pi+1}/{n_pairs}  loss={loss.item():.3f}")
    model.zero_grad(set_to_none=True)  # leave caller with clean state

    raw_grads = {
        **{f"hack/{n}": torch.stack(gs) for n, gs in grads_hack.items()},
        **{f"clean/{n}": torch.stack(gs) for n, gs in grads_clean.items()},
    }

    # Per module: D = g_hack - g_clean (paired diff cancels prompt-specific noise).
    # SVD(D) gives orthonormal right singular vectors = principal axes of variation
    # of the hack-clean axis. Top-k generalizes mean-diff (which is the k=1 case).
    v_hack: dict[str, torch.Tensor] = {}
    v_sv:   dict[str, torch.Tensor] = {}
    rows = []
    n_zero = 0
    k = min(top_k, n_pairs)
    n_axes_kept_total = 0
    for name in grads_hack:
        G_h = torch.stack(grads_hack[name])   # [n_pairs, r]
        G_c = torch.stack(grads_clean[name])
        D = G_h - G_c

        U_d, S_d, Vh_d = torch.linalg.svd(D, full_matrices=False)
        V = Vh_d[:k]   # [k, r], rows orthonormal in R^r
        # Orient by per-pair majority vote: for each axis i, count pairs where
        # d_p @ v_i > 0; if strict majority disagree with current SVD sign, flip.
        # More outlier-robust than sign(mean): one extreme pair can't flip a
        # consensus direction. Matches repeng's _orient_svd convention.
        proj_per_pair = D @ V.T                           # [n_pairs, k]
        n_pos = (proj_per_pair > 0).float().sum(0)        # [k]
        flip = torch.where(n_pos < n_pairs / 2, -torch.ones(k), torch.ones(k))
        V = V * flip.unsqueeze(1)

        # tau_axis: zero rows where S_i/S_0 < tau_axis (diagnostic; see Config comment).
        n_axes_kept = k
        if tau_axis > 0 and S_d[0] > 1e-12:
            ratios = S_d[:k] / S_d[0]
            keep = (ratios >= tau_axis).float()
            V = V * keep.unsqueeze(1)
            n_axes_kept = int(keep.sum())
        n_axes_kept_total += n_axes_kept

        nrm = D.norm()
        if nrm < 1e-12:
            n_zero += 1
            v_hack[name] = torch.zeros((k, D.shape[1]), dtype=V.dtype).contiguous()
        else:
            v_hack[name] = V.contiguous()
        # Record singular values so the load-time noise-floor filter has the
        # extraction-time S_i per axis without re-extracting. Saved under
        # `_sv/{name}` keys in the safetensors file (combined at save site).
        v_sv[name] = S_d[:k].clone().contiguous()
        sv_top = S_d[:k]
        sv_total = S_d.sum().clamp_min(1e-12)
        rows.append({
            "module": name.split(".")[-1],
            "r": D.shape[1],
            "||D||": f"{nrm:.2e}",
            "sv_0": f"{S_d[0].item():.2e}" if S_d.numel() else "-",
            f"sv_top{k}_frac": f"{(sv_top.sum() / sv_total).item():.2f}",
            "sv_ratio_0/1": ("-" if S_d.numel() < 2
                             else f"{(S_d[0] / S_d[1].clamp_min(1e-12)).item():.2f}"),
            "axes_kept": n_axes_kept,
        })
    n_modules = len(grads_hack)
    logger.info(
        f"v_hack: modules={n_modules}  k_max={k}  zero-||D||={n_zero}  "
        f"axes_kept_avg={n_axes_kept_total/max(1,n_modules):.1f} (tau_axis={tau_axis})"
    )
    return v_hack, v_sv, raw_grads, rows


def main(cfg: Config) -> int:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dtype = resolve_dtype(cfg.dtype)
    if cfg.pairs_from_pool is not None:
        pairs = load_pairs_json(cfg.pairs_from_pool)
        logger.info(f"pairs source: pool-derived ({cfg.pairs_from_pool}) -> {len(pairs)} pairs")
    else:
        pairs = list(PAIRS)
        logger.info(f"pairs source: hand-crafted vgrout.pairs.PAIRS ({len(pairs)} pairs)")
    logger.info(
        f"device={device}  model={cfg.model}  dtype={cfg.dtype}  "
        f"N_pairs={len(pairs)}  heldout={cfg.n_heldout}  top_k={cfg.top_k}  tau_axis={cfg.tau_axis}"
    )

    tokenizer = AutoTokenizer.from_pretrained(cfg.model)
    model = AutoModelForCausalLM.from_pretrained(
        cfg.model, dtype=dtype, attn_implementation="sdpa"
    ).to(device)
    model.eval()  # disable dropout; gradients still flow through delta_S
    wrappers = wrap_model_with_antipasto(
        model, model_name=cfg.model, cache_root=CACHE_ROOT, svd_device=device,
    )
    n_mod = len(wrappers)
    n_delta = sum(info["delta_S"].numel() for info in wrappers.values())
    logger.info(f"wrapped {n_mod} modules; total delta_S scalars = {n_delta:,}")

    train_pairs = pairs[:-cfg.n_heldout] if cfg.n_heldout > 0 else pairs
    logger.info(f"train pairs: {len(train_pairs)}  held: {cfg.n_heldout}")

    v_hack, v_sv, raw_grads, rows = extract_v_hack(
        model, tokenizer, wrappers, pairs,
        top_k=cfg.top_k, tau_axis=cfg.tau_axis,
        n_heldout=cfg.n_heldout, device=device,
    )
    n_zero = sum(1 for v in v_hack.values() if v.norm() < 1e-12)
    k = min(cfg.top_k, len(train_pairs))

    cfg.out_path.parent.mkdir(parents=True, exist_ok=True)
    cfg.train_grads_path.parent.mkdir(parents=True, exist_ok=True)
    save_file(raw_grads, str(cfg.train_grads_path),
              metadata={"model": cfg.model, "dtype": cfg.dtype})
    # v_hack file layout: bare `{name}` keys hold V[k, r]; `_sv/{name}` keys
    # hold S[k]. Loader at train.py:load_v_hack splits them back apart.
    save_payload = {**v_hack, **{f"_sv/{n}": s for n, s in v_sv.items()}}
    save_file(save_payload, str(cfg.out_path),
              metadata={"model": cfg.model, "dtype": cfg.dtype, "top_k": str(k),
                        "tau_axis": str(cfg.tau_axis), "schema": "v2_with_sv"})

    # summary: aggregate by suffix -- track top-k energy concentration
    by_suffix: dict[str, list] = defaultdict(list)
    for r in rows:
        by_suffix[r["module"]].append(float(r[f"sv_top{k}_frac"]))
    agg_rows = []
    for suf, vals in sorted(by_suffix.items()):
        agg_rows.append({
            "suffix": suf,
            "n": len(vals),
            f"mean_sv_top{k}_frac": f"{sum(vals)/len(vals):.2f}",
            f"min_sv_top{k}_frac": f"{min(vals):.2f}",
            f"max_sv_top{k}_frac": f"{max(vals):.2f}",
        })

    # Final tail: BLUF -- what an agent reads first should be result + interp.
    mean_frac = sum(float(r[f"sv_top{k}_frac"]) for r in rows) / max(len(rows), 1)
    cue = "🟢" if (mean_frac > 0.5 and n_zero == 0) else ("🟡" if n_zero == 0 else "🔴")

    print(f"\nSHOULD: mean_sv_top{k}_frac > 0.5 per suffix (subspace captures most energy). "
          f"zero-||D||==0 (else grad flow broken).\n")
    print(tabulate(agg_rows, headers="keys", tablefmt="tsv", floatfmt=".2f"))
    print()
    print(f"out: {cfg.out_path}")
    print(f"argv: extract_vhack_grad --model={cfg.model} --top-k={k} --n-heldout={cfg.n_heldout}")
    print(f"main metric: mean_sv_top{k}_frac={mean_frac:.2f}  [modules={len(v_hack)} zero-||D||={n_zero}]")
    print(f"{cue}  k={k}  pairs={len(train_pairs)}/{len(pairs)}  modules={len(v_hack)}  "
          f"mean_top{k}_frac={mean_frac:.2f}  zero={n_zero}")

    if n_zero > 0:
        logger.error(f"FAIL: {n_zero}/{len(v_hack)} modules have zero ||D|| -- gradient flow broken")
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main(tyro.cli(Config)))


def load_v_hack(
    path: Path, model_name: str, wrappers: dict,
    k_use: int | None = None, drop_bottom_frac: float = 0.0,
) -> dict[str, Float[torch.Tensor, "k r"]]:
    """Load v_hack (top-k directions) for this wrapped model.

    File schema (v2): bare `{name}` keys hold V[k_max, r]; `_sv/{name}` keys hold
    S[k_max]. v_hack is model-specific because module names and per-module SVD
    ranks depend on the exact checkpoint; a smoke (Qwen3.5-0.8B) v_hack must
    not be reused for a full (Qwen3-4B) run.

    If `k_use` is given, slices V (and S) to top-k_use rows. Errors if
    k_use > k_max saved (re-extract with a higher top_k).

    If `drop_bottom_frac > 0`, drops the bottom-fraction of singular values Sᵢ by
    global quantile; a module with every axis below the threshold is dropped from
    the returned dict (projection no-ops there -- no hack signal).
    """
    with safe_open(str(path), framework="pt", device="cpu") as f:
        meta = f.metadata() or {}
        saved_model = meta.get("model")
        saved_dtype = meta.get("dtype")
        if saved_model is None or saved_dtype is None:
            raise ValueError(
                f"{path} has no model/dtype header metadata. "
                f"Re-extract with `uv run python -m vgrout.extract_vhack_grad "
                f"--model={model_name} --dtype=bf16 --out-path={path}`."
            )
        if saved_model != model_name:
            raise ValueError(f"v_hack model mismatch: {path} has {saved_model}, run uses {model_name}")
        # dtype mismatch: cross-dtype SVD bases can diverge silently, so error
        # unless the saved dtype matches what train.py uses on this device.
        # CPU runs in fp32, CUDA runs in bf16 (see model-load site above).
        expected_dtype = "fp32" if torch.cuda.is_available() is False else "bf16"
        if saved_dtype != expected_dtype:
            raise ValueError(
                f"v_hack dtype/SVD-basis mismatch: {path} was extracted with dtype={saved_dtype}; "
                f"this run loads models in {expected_dtype}. Re-extract with `--dtype={expected_dtype}`."
            )
        v_hack = {k: f.get_tensor(k) for k in f.keys() if not k.startswith("_sv/")}
        v_sv   = {k[len("_sv/"):]: f.get_tensor(k) for k in f.keys() if k.startswith("_sv/")}

    wrapper_keys = set(wrappers)
    vhack_keys = set(v_hack)
    missing = sorted(wrapper_keys - vhack_keys)
    extra = sorted(vhack_keys - wrapper_keys)
    # v_hack[name] is [k_max, r]; δS is [r]. Check last-dim match (rank r).
    rank_bad = [
        (name, tuple(v_hack[name].shape), tuple(wrappers[name]["delta_S"].shape))
        for name in sorted(wrapper_keys & vhack_keys)
        if v_hack[name].ndim != 2 or v_hack[name].shape[-1] != wrappers[name]["delta_S"].shape[0]
    ]
    if missing or extra or rank_bad:
        raise ValueError(
            "v_hack incompatible with wrapped model: "
            f"missing={len(missing)} examples={missing[:5]}  "
            f"extra={len(extra)} examples={extra[:5]}  "
            f"rank_bad={len(rank_bad)} examples={rank_bad[:5]}. "
            "Extract a fresh v_hack with `uv run python -m vgrout.extract_vhack_grad "
            f"--model={model_name} --out-path={path}`."
        )

    v_hack = postprocess_v_hack(
        v_hack, v_sv, k_use=k_use, drop_bottom_frac=drop_bottom_frac, source=str(path),
    )
    return v_hack


def postprocess_v_hack(
    v_hack: dict[str, Float[torch.Tensor, "k r"]],
    v_sv:   dict[str, Float[torch.Tensor, "k"]],
    k_use: int | None,
    drop_bottom_frac: float,
    source: str = "<refresh>",
) -> dict[str, Float[torch.Tensor, "k r"]]:
    """Apply k_use slice + global noise-floor filter.

    Shared between `load_v_hack` (init-time, reading from safetensors) and the
    in-loop refresh hook (where we hand in fresh `extract_v_hack` outputs).
    Mutates neither input dict; returns a fresh filtered dict.

    Global noise floor: drop the bottom `drop_bottom_frac` of singular values Sᵢ
    by quantile across all modules. A module with every axis below the threshold
    is removed (projection iterates v_hack, so it no-ops there). Threshold
    recomputes per call (tracks the current S distribution).
    """
    k_max = next(iter(v_hack.values())).shape[0]
    if k_use is not None:
        if k_use > k_max:
            raise ValueError(f"requested k_use={k_use} exceeds k_max={k_max} (source={source})")
        v_hack = {n: v[:k_use].contiguous() for n, v in v_hack.items()}
        v_sv   = {n: s[:k_use].contiguous() for n, s in v_sv.items()}
    n_dropped_modules = 0
    n_axes_before = sum(v.shape[0] for v in v_hack.values())
    threshold = None
    if drop_bottom_frac > 0 and v_sv:
        all_S = torch.cat([v_sv[n].float() for n in v_hack])
        threshold = torch.quantile(all_S, drop_bottom_frac).item()
        filtered: dict[str, torch.Tensor] = {}
        for name, V in v_hack.items():
            keep = v_sv[name].float() >= threshold
            if keep.any():
                filtered[name] = V[keep].contiguous()
            else:
                n_dropped_modules += 1
        v_hack = filtered
    n_axes_after = sum(v.shape[0] for v in v_hack.values())
    logger.info(
        f"postprocess_v_hack({source}): modules={len(v_hack)} (dropped {n_dropped_modules}); "
        f"k_use={k_use or k_max}/k_max={k_max}; axes={n_axes_after}/{n_axes_before} kept "
        f"(drop_bottom_frac={drop_bottom_frac}, threshold={threshold})"
    )
    return v_hack
