"""Loading and post-processing the extracted hack-direction basis (v_hack).

v_hack is a per-module set of top-k right singular vectors of the labeled-pair
GRPO gradient, saved by extract_vhack_grad. Here we load it for a wrapped model
(checking the model/dtype/rank all match) and apply the top-k slice plus the
global noise-floor filter. The same post-processing serves both the init-time
load and the in-loop refresh.
"""
from __future__ import annotations

from pathlib import Path

import torch
from jaxtyping import Float
from loguru import logger
from safetensors import safe_open


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`, collects every S_i across every module and drops
    the bottom-fraction by global quantile. Modules whose every axis is below
    the global threshold get filtered out of the returned dict (projection on
    those modules becomes a no-op — they didn't carry hack signal anywhere).
    """
    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]; delta_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: collect every S_i across every module, drop the bottom
    `drop_bottom_frac` by quantile. A module whose every axis falls below the
    global threshold is removed entirely — projection iterates v_hack so it
    becomes a no-op for that module. Threshold recomputes per call (tracks
    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
