"""Backbone feature extraction: one forward pass per example, hidden states captured by hooks."""
from __future__ import annotations

import argparse
import json
import math
import random
import re
import time
from pathlib import Path

import torch

from .cache import CacheWriter, Decision, FeatureCache, Record
from .data import read_manifest
from .prompting import alias_token_ids, canvas_ids, shuffled_order, tokenize_prompt
from .schema import Example, argmax

MODEL_ID = "google/diffusiongemma-26B-A4B-it"
PROBE_STATE = "A short probe state."
RANDOM_DRAWS = 4
SEED_MODES = ("pad", "alias", "random")
ZERO_SHOT_LOG = "configs/session1_zero_shot.jsonl"


class NormCapture:
    """Registers forward hooks on final-norm modules (and any modules named in `extra`) and keeps their latest
    outputs. With no names given every norm outside the layer stacks is hooked (discovery, see probe)."""

    def __init__(self, model, encoder_name: str | None = None, decoder_name: str | None = None,
                 extra: tuple[str, ...] | list[str] = ()):
        self.outputs: dict[str, torch.Tensor] = {}
        self.handles = []
        extra = set(extra)
        for name, mod in model.named_modules():
            if name in extra:
                self.handles.append(mod.register_forward_hook(self._hook(name)))
                continue
            if ".layers." in f".{name}." or "Norm" not in type(mod).__name__:
                continue
            if encoder_name is not None and name not in (encoder_name, decoder_name):
                continue
            self.handles.append(mod.register_forward_hook(self._hook(name)))

    def _hook(self, name):
        def fn(_m, _i, out):
            self.outputs[name] = out[0] if isinstance(out, tuple) else out
        return fn

    def remove(self):
        for h in self.handles:
            h.remove()


def unit_rms(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    """Scale every row to unit RMS (a parameter-free RMSNorm). Intermediate layers are residual-stream states
    of arbitrary magnitude; this keeps them on the final norm's scale and inside float16 for the cache."""
    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)


_LAYER_RE = re.compile(r"(.*\blayers)\.(\d+)")


def layer_stacks(model) -> dict[str, int]:
    """Layer-stack prefixes and their depths, from module names of the form `<prefix>.layers.<n>`."""
    stacks: dict[str, int] = {}
    for name, _ in model.named_modules():
        m = _LAYER_RE.fullmatch(name)
        if m:
            stacks[m.group(1)] = max(stacks.get(m.group(1), 0), int(m.group(2)) + 1)
    return stacks


def select_layers(n_layers: int, fracs) -> list[int]:
    """Zero-based layer indices at the given depth fractions: 0.5 of 30 layers is layer 14, 0.75 is 22."""
    out = set()
    for f in fracs:
        if not 0.0 < float(f) <= 1.0:
            raise ValueError(f"layer fraction must be in (0, 1], got {f}")
        out.add(max(0, min(n_layers - 1, math.ceil(float(f) * n_layers) - 1)))
    return sorted(out)


def resolve_extra_layers(model, encoder_norm: str, decoder_norm: str, fracs) -> dict[str, list[str]]:
    """Module names of the extra layers to capture: for each final norm, the layer stack beside it (same parent,
    else the shallowest stack below that parent) at `fracs` of its depth."""
    stacks = layer_stacks(model)

    def parent(name: str) -> str:
        return name.rsplit(".", 1)[0] if "." in name else ""

    def stack_for(norm: str) -> str:
        p = parent(norm)
        beside = [s for s in stacks if parent(s) == p]
        below = [s for s in stacks if p and s.startswith(p + ".")]
        cands = beside or below
        if not cands:
            raise RuntimeError(f"no layer stack found beside {norm!r}; stacks: {sorted(stacks)}")
        return min(cands, key=len)

    enc, dec = stack_for(encoder_norm), stack_for(decoder_norm)
    if enc == dec:
        raise RuntimeError(f"encoder and decoder norms resolve to the same layer stack {enc!r}")
    return {"encoder": [f"{enc}.{i}" for i in select_layers(stacks[enc], fracs)],
            "decoder": [f"{dec}.{i}" for i in select_layers(stacks[dec], fracs)]}


def extra_layer_names(info: dict) -> list[str]:
    extra = info.get("extra_layers") or {}
    return list(extra.get("encoder", [])) + list(extra.get("decoder", []))


def _probe_example() -> Example:
    from .schema import Question
    q = Question.from_jev("ok", {"type": "noul", "instructions": "Is the probe fine?"})
    return Example(id="probe", source="probe", state=PROBE_STATE, questions=[q], gold={"ok": None})


@torch.no_grad()
def _forward(model, prompt_ids: list[int], canvas: list[int]):
    dev = model.device
    input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=dev)
    return model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids),
                 decoder_input_ids=torch.tensor([canvas], dtype=torch.long, device=dev))


def probe(model, tok, hidden_size: int) -> dict:
    ex = _probe_example()
    tp = tokenize_prompt(ex, tok)
    for canvas_length in (64, 256):
        canvas, _ = canvas_ids(ex, tok, slot_token_id=0, canvas_length=canvas_length)
        cap = NormCapture(model)
        try:
            _forward(model, tp.ids, canvas)
        except Exception as err:  # noqa: BLE001 - the whole point is to learn which length works
            cap.remove()
            print(f"canvas_length {canvas_length} failed: {type(err).__name__}: {err}")
            continue
        cap.remove()
        enc = [n for n, o in cap.outputs.items() if o.shape[-2] == len(tp.ids) and o.shape[-1] == hidden_size]
        dec = [n for n, o in cap.outputs.items() if o.shape[-2] == canvas_length and o.shape[-1] == hidden_size]
        if not enc or not dec:
            shapes = {k: tuple(v.shape) for k, v in cap.outputs.items()}
            raise RuntimeError(f"could not find final norms; captured {shapes}")
        if enc[-1] == dec[-1]:
            raise RuntimeError(f"encoder and decoder norms resolved to the same module: {enc[-1]}")
        if len(enc) > 1 or len(dec) > 1:
            print(f"warning: several final-norm candidates, using the last; encoder {enc}, decoder {dec}")
        return {"canvas_length": canvas_length, "encoder_norm": enc[-1], "decoder_norm": dec[-1],
                "hidden_size": hidden_size}
    raise RuntimeError("neither canvas length 64 nor 256 ran")


def _slot_token(seed_mode: str, alias_ids: list[int], rng: random.Random, vocab: int) -> list[int]:
    if seed_mode == "pad":
        return [0]
    if seed_mode == "alias":
        return [alias_ids[0]]
    if seed_mode == "random":
        return [rng.randrange(10, vocab) for _ in range(RANDOM_DRAWS)]
    raise ValueError(f"unknown seed_mode {seed_mode}")


def _rows(cap: "NormCapture", name: str, n_rows: int, what: str) -> torch.Tensor:
    """A captured layer's rows for this forward (batch element 0), checked against the expected row count."""
    out = cap.outputs[name][0].float()
    if out.shape[0] != n_rows:
        raise RuntimeError(f"{name} output has shape {tuple(out.shape)}, expected {n_rows} {what} rows")
    return out


@torch.no_grad()
def encode_example(model, tok, ex: Example, cap: "NormCapture", info: dict, *, slot_tokens: list[int],
                   alias_ids: list[int], max_tokens: int = 1024, option_order=None) -> tuple[torch.Tensor, Record]:
    """One example through the backbone: features averaged over slot tokens, plus a Record with priors.

    The Record's path is left empty; the caller names the file.
    """
    canvas_length = info["canvas_length"]
    tp = tokenize_prompt(ex, tok, max_tokens=max_tokens, option_order=option_order)
    feats_sum, prob_sum, slots = None, None, None
    for slot_tok in slot_tokens:
        canvas, slots = canvas_ids(ex, tok, slot_token_id=slot_tok, canvas_length=canvas_length)
        cap.outputs.clear()
        out = _forward(model, tp.ids, canvas)
        if out.logits.shape[1] != canvas_length:
            raise RuntimeError(f"logits cover {out.logits.shape[1]} positions, expected canvas_length {canvas_length}")
        extra = info.get("extra_layers") or {}
        wanted = [info["encoder_norm"], info["decoder_norm"], *extra_layer_names(info)]
        missing = [n for n in wanted if n not in cap.outputs]
        if missing:
            raise RuntimeError(f"norm outputs missing after forward: {missing}; captured {sorted(cap.outputs)}")
        hp = cap.outputs[info["encoder_norm"]][0].float()
        hc = cap.outputs[info["decoder_norm"]][0].float()
        if hp.shape[0] != len(tp.ids):
            raise RuntimeError(f"encoder norm output has shape {tuple(hp.shape)}, expected {len(tp.ids)} prompt rows")
        if hc.shape[0] != canvas_length:
            raise RuntimeError(f"decoder norm output has shape {tuple(hc.shape)}, expected canvas_length {canvas_length}")
        enc_extra = [unit_rms(_rows(cap, name, len(tp.ids), "prompt")) for name in extra.get("encoder", [])]
        dec_extra = [unit_rms(_rows(cap, name, canvas_length, "canvas")) for name in extra.get("decoder", [])]
        feats = torch.cat([torch.cat([hp, *enc_extra], dim=1), torch.cat([hc, *dec_extra], dim=1)], dim=0)
        probs = torch.softmax(out.logits[0].float(), dim=-1)[slots]  # (n_q, vocab)
        feats_sum = feats if feats_sum is None else feats_sum + feats
        prob_sum = probs if prob_sum is None else prob_sum + probs
    feats = feats_sum / len(slot_tokens)
    logp = (prob_sum / len(slot_tokens)).clamp_min(1e-12).log()
    decisions = []
    for qi, q in enumerate(ex.questions):
        ids = torch.tensor([alias_ids[a] for a in tp.option_alias[qi]], device=logp.device)
        prior = logp[qi, ids]
        prior = (prior - torch.logsumexp(prior, dim=0)).tolist()
        decisions.append(Decision(qid=q.qid, type=q.type, slot=len(tp.ids) + slots[qi],
                                  spans=[list(s) for s in tp.option_spans[qi]],
                                  alias=list(tp.option_alias[qi]), prior=prior, gold=ex.gold.get(q.qid)))
    rec = Record(id=ex.id, source=ex.source, path="", n_prompt=len(tp.ids),
                 n_canvas=canvas_length, decisions=decisions)
    return feats, rec


def extract(model, tok, examples: list[Example], out_dir, *, seed_mode: str, canvas_length: int,
            alias_ids: list[int], hidden_size: int, max_tokens: int = 1024, shuffle_rng=None,
            limit: int | None = None, vocab_size: int = 262144, log_every: int = 50,
            info: dict | None = None, layer_fracs=()) -> dict:
    """Write a feature cache. With `layer_fracs` (depth fractions, e.g. (0.5, 0.75)) every row also carries the
    unit-RMS outputs of the encoder (prompt rows) or decoder (canvas rows) layers at those depths, appended to the
    final-norm block, so the feature dim is hidden_size * (1 + len(layers))."""
    info = dict(probe(model, tok, hidden_size) if info is None else info)
    if info["canvas_length"] != canvas_length:
        raise ValueError(f"requested canvas {canvas_length} but probe says {info['canvas_length']}")
    if layer_fracs:
        info["extra_layers"] = resolve_extra_layers(model, info["encoder_norm"], info["decoder_norm"], layer_fracs)
    writer = CacheWriter(out_dir)  # opened before the hooks go on, so a folder that cannot be written leaks nothing
    cap = NormCapture(model, info["encoder_norm"], info["decoder_norm"], extra=extra_layer_names(info))
    rng = random.Random(1234)
    n_dec, t0 = 0, time.time()
    examples = examples[:limit] if limit else examples
    try:
        for i, ex in enumerate(examples):
            order = shuffled_order(ex, shuffle_rng) if shuffle_rng is not None else None
            draws = _slot_token(seed_mode, alias_ids, rng, vocab_size)
            feats, rec = encode_example(model, tok, ex, cap, info, slot_tokens=draws, alias_ids=alias_ids,
                                        max_tokens=max_tokens, option_order=order)
            rec.path = f"{ex.id.replace('/', '_')}.safetensors"
            writer.write(rec, feats)
            n_dec += len(rec.decisions)
            if (i + 1) % log_every == 0:
                print(f"{i + 1}/{len(examples)} examples, {time.time() - t0:.0f}s", flush=True)
    finally:  # close the manifest and drop the hooks even when an example raises
        writer.close()
        cap.remove()
    n_extra = len((info.get("extra_layers") or {}).get("encoder", []))
    summary = {"examples": len(examples), "decisions": n_dec, "seed_mode": seed_mode,
               "canvas_length": canvas_length, "seconds": round(time.time() - t0, 1),
               "feature_dim": hidden_size * (1 + n_extra), **info}
    Path(out_dir, "summary.json").write_text(json.dumps(summary, indent=2))
    return summary


def zero_shot_accuracy(cache_dir) -> float:
    cache = FeatureCache(cache_dir)
    hits = total = 0
    for rec in cache.records:
        for d in rec.decisions:
            if d.gold is None:
                continue
            total += 1
            hits += int(argmax(d.prior) == argmax(d.gold))
    return hits / max(total, 1)


def load_backbone(model_id: str = MODEL_ID, load_in_4bit: bool = False):
    from transformers import AutoTokenizer, DiffusionGemmaForBlockDiffusion
    kwargs = {"device_map": "auto", "dtype": torch.bfloat16}
    if load_in_4bit:
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
    model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, **kwargs).eval()
    return model, AutoTokenizer.from_pretrained(model_id)


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("probe")
    p.add_argument("--model", default=MODEL_ID)
    e = sub.add_parser("extract")
    e.add_argument("--model", default=MODEL_ID)
    e.add_argument("--manifest", required=True)
    e.add_argument("--out", required=True)
    e.add_argument("--seed-mode", default="pad",
                   help="pad, alias or random, or a comma-separated list of them; a list writes <out>_<mode> per mode")
    e.add_argument("--canvas", type=int, default=64)
    e.add_argument("--limit", type=int, default=None)
    e.add_argument("--shuffle-seed", type=int, default=None, help="shuffle choice options (training augmentation)")
    e.add_argument("--load-in-4bit", action="store_true")
    e.add_argument("--zero-shot-log", action="store_true", help=f"append each cache's zero-shot accuracy to {ZERO_SHOT_LOG}")
    for sp in (p, e):
        sp.add_argument("--layers", default=None,
                        help="depth fractions of extra encoder/decoder layers to append to the features, e.g. 0.5,0.75 "
                             "(feature dim becomes hidden_size * (1 + count)); probe records their module names")
    z = sub.add_parser("zero-shot")
    z.add_argument("--cache", required=True)
    args = ap.parse_args()
    fracs = tuple(float(x) for x in args.layers.split(",")) if getattr(args, "layers", None) else ()
    if args.cmd == "zero-shot":
        print(json.dumps({"cache": args.cache, "zero_shot_accuracy": zero_shot_accuracy(args.cache)}))
        return
    if args.cmd == "extract":
        modes = args.seed_mode.split(",")
        if any(m not in SEED_MODES for m in modes) or len(set(modes)) != len(modes):
            e.error(f"--seed-mode takes distinct values from {', '.join(SEED_MODES)}; got {args.seed_mode!r}")
    model, tok = load_backbone(args.model, getattr(args, "load_in_4bit", False))
    hidden = model.config.text_config.hidden_size
    if args.cmd == "probe":
        info = probe(model, tok, hidden)
        if fracs:
            info["extra_layers"] = resolve_extra_layers(model, info["encoder_norm"], info["decoder_norm"], fracs)
            info["feature_dim"] = hidden * (1 + len(info["extra_layers"]["encoder"]))
        Path("configs").mkdir(exist_ok=True)
        Path("configs/backbone.json").write_text(json.dumps({**info, "model_id": args.model}, indent=2))
        print(json.dumps(info, indent=2))
        return
    alias = alias_token_ids(tok, 64)
    vocab = getattr(model.config.text_config, "vocab_size", 262144)
    examples = read_manifest(args.manifest)
    info = probe(model, tok, hidden) if len(modes) > 1 else None  # one probe shared by every mode
    for mode in modes:
        out = f"{args.out}_{mode}" if len(modes) > 1 else args.out
        # a fresh rng per mode, so each cache matches what a separate run of that mode would write
        rng = random.Random(args.shuffle_seed) if args.shuffle_seed is not None else None
        summary = extract(model, tok, examples, out, seed_mode=mode, canvas_length=args.canvas, alias_ids=alias,
                          hidden_size=hidden, shuffle_rng=rng, limit=args.limit, vocab_size=vocab, info=info,
                          layer_fracs=fracs)
        print(json.dumps(summary, indent=2))
        acc = zero_shot_accuracy(out)
        print(f"zero_shot {mode} {acc:.4f}", flush=True)
        if args.zero_shot_log:
            log = Path(ZERO_SHOT_LOG)
            log.parent.mkdir(parents=True, exist_ok=True)
            with log.open("a", encoding="utf-8") as f:
                f.write(json.dumps({"cache": out, "seed_mode": mode, "zero_shot_accuracy": acc}) + "\n")


if __name__ == "__main__":
    main()