"""Running a model and reading the cone off its own trajectory.

    rt = ulens.Runner("qwen3.5-4b")
    tr = rt.run("Is Hamlet's madness feigned?", max_new_tokens=600, seed=42)

    tr.u_lens                     # the score
    tr.field                      # [n_blocks, T, k] signed projections: depth x token x axis
    tr.a_cone_series()            # A_cone per token at the readout block
    tr.surface()                  # A_cone over depth x token

Two passes: `generate` with the model's own sampling settings, then one teacher-forced pass
over prompt + generation that reads every lens block through a forward hook and the per-token
entropies through the unembedding. The second pass is what makes the field exact — a
generation pass sees each token once, with the state it had at the time, which is the same
thing only when nothing is being intervened on.

Alignment, which decides whether the score matches the published one. `entropy[i]` is the
entropy of the distribution that **produced** token `i`, read from the logits at `i - 1`, so
the thinking-block mean is `entropy[prompt_len : eot + 1].mean()` and includes the
distribution that produced the end-of-think token itself. `A_cone` is read at the end-of-think
token's own state. Both match `pipeline/04_forward.py` of the reproduction code.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass
from dataclasses import field as _field
from pathlib import Path

import numpy as np

from . import axes as _axes
from . import basis as _basis
from . import readout as _readout
from ._deps import require, require_extra
from .models import ModelSpec
from .models import get as _model

SEG_PROMPT, SEG_THINK, SEG_ANSWER = "prompt", "think", "answer"

#: positions per chunk when pushing hidden states through the unembedding
LOGIT_CHUNK = 256


@dataclass
class Trace:
    """One generation, with the cone read off every depth and every token."""
    model: str
    prompt: str
    tokens: list[int]
    token_strs: list[str]
    prompt_len: int
    eot: int                         # absolute index of the end-of-think token
    entropy: np.ndarray              # [T], of the distribution that produced each token
    field: np.ndarray                # [n_blocks, T, k] signed projections
    blocks: list[int]
    names: list[str]
    readout_block: int
    generated: str
    truncated: bool                  # the thinking block was never closed
    config: dict = _field(default_factory=dict)

    # ------------------------------------------------------------------ geometry
    @property
    def n_tokens(self) -> int:
        return len(self.tokens)

    @property
    def n_axes(self) -> int:
        return int(self.field.shape[2])

    @property
    def think_len(self) -> int:
        return self.eot + 1 - self.prompt_len

    @property
    def segments(self) -> np.ndarray:
        """`[T]` of 'prompt' / 'think' / 'answer'."""
        s = np.array([SEG_ANSWER] * self.n_tokens, dtype=object)
        s[: self.prompt_len] = SEG_PROMPT
        s[self.prompt_len: self.eot + 1] = SEG_THINK
        return s

    def block_index(self, block: int | None = None) -> int:
        b = self.readout_block if block is None else block
        if b not in self.blocks:
            raise KeyError(f"block {b} was not captured; have {self.blocks}")
        return self.blocks.index(b)

    def projections(self, block: int | None = None) -> np.ndarray:
        """`[T, k]` signed projections at one block."""
        return self.field[self.block_index(block)]

    def components(self, block: int | None = None) -> np.ndarray:
        """`[T, k]` rectified components at one block."""
        return _readout.rectify(self.projections(block))

    def a_cone_series(self, block: int | None = None) -> np.ndarray:
        """`[T]` A_cone per token at one block."""
        return _readout.a_cone_of(self.projections(block))

    def surface(self) -> np.ndarray:
        """`[n_blocks, T]` A_cone over depth and token: the field, as one number per cell."""
        return _readout.a_cone_of(self.field)

    # ------------------------------------------------------------------ the score
    @property
    def a_cone(self) -> float:
        """A_cone at the end-of-think token, at the readout block."""
        return float(self.a_cone_series()[self.eot])

    @property
    def think_entropy(self) -> float:
        """Mean entropy of the distributions that produced the thinking block's tokens;
        NaN when there is no thinking block (thinking off), without numpy's warning."""
        e = self.entropy[self.prompt_len: self.eot + 1]
        return float(e.mean()) if e.size else float("nan")

    @property
    def u_lens(self) -> float:
        """The score. Higher means more likely to be wrong.

        Check `readout_valid` before comparing this with a published number.
        """
        return self.a_cone * self.think_entropy

    @property
    def readout_valid(self) -> bool:
        """Whether the score was read where the method says to read it.

        False when the thinking block was never closed: the readout then falls on the last
        token of a budget-truncated block rather than on an end-of-think marker, and the
        answer is empty. The number is still computed — the field is real — but it is not
        the published quantity and must not be compared with one.
        """
        return not self.truncated and self.think_len > 0

    def summary(self) -> str:
        t = ("TRUNCATED, never closed - readout is not the published one"
             if self.truncated else "closed")
        return (f"{self.model}: {self.n_tokens} tokens "
                f"(prompt {self.prompt_len}, think {self.think_len} [{t}], "
                f"answer {self.n_tokens - self.eot - 1}) · "
                f"{len(self.blocks)} blocks [{self.blocks[0]}-{self.blocks[-1]}] · "
                f"readout block {self.readout_block}\n"
                f"  A_cone {self.a_cone:.4f} × think entropy {self.think_entropy:.4f} "
                f"= U-Lens {self.u_lens:.4f}")

    def axis_table(self, block: int | None = None) -> str:
        """Per-axis mean rectified component over the thinking block."""
        C = self.components(block)[self.prompt_len: self.eot + 1]
        tot = float((C ** 2).sum()) or 1.0
        b = self.readout_block if block is None else block
        out = [f"axis contributions over the thinking block (block {b})"]
        for i, n in enumerate(self.names):
            out.append(f"  {n:22s} mean {C[:, i].mean():.4f}  "
                       f"share {100 * (C[:, i] ** 2).sum() / tot:5.1f}%  "
                       f"argmax {100 * (C.argmax(1) == i).mean():5.1f}%")
        return "\n".join(out)

    # ------------------------------------------------------------------ persistence
    def _meta(self) -> dict:
        return {"model": self.model, "prompt": self.prompt, "prompt_len": self.prompt_len,
                "eot": self.eot, "blocks": list(self.blocks), "names": list(self.names),
                "readout_block": self.readout_block, "generated": self.generated,
                "truncated": self.truncated, "token_strs": self.token_strs,
                "config": self.config}

    def save(self, path: str | Path) -> Path:
        """Write `<path>.npz` plus a readable `<path>.json` sidecar."""
        path = Path(path)
        if path.suffix != ".npz":
            path = path.with_suffix(".npz")
        path.parent.mkdir(parents=True, exist_ok=True)
        np.savez_compressed(path, field=self.field.astype(np.float32),
                            entropy=self.entropy.astype(np.float32),
                            tokens=np.array(self.tokens),
                            meta=np.array(json.dumps(self._meta())))
        side = dict(self._meta())
        # a trace without a thinking block has no thinking-block entropy and no score; the
        # sidecar says null, which is JSON, rather than NaN, which is not
        side["scores"] = {k: (float(v) if np.isfinite(v) else None) for k, v in
                          (("a_cone", self.a_cone), ("think_entropy", self.think_entropy),
                           ("u_lens", self.u_lens))}
        path.with_suffix(".json").write_text(json.dumps(side, indent=1, allow_nan=False)
                                             + "\n")
        return path


def load(path: str | Path) -> Trace:
    """Read a trace back. No model and no GPU needed — every view works from the field."""
    z = np.load(Path(path), allow_pickle=False)
    m = json.loads(str(z["meta"]))
    return Trace(model=m["model"], prompt=m["prompt"], tokens=[int(t) for t in z["tokens"]],
                 token_strs=m["token_strs"], prompt_len=m["prompt_len"], eot=m["eot"],
                 entropy=np.asarray(z["entropy"], np.float32),
                 field=np.asarray(z["field"], np.float32), blocks=m["blocks"],
                 names=m["names"], readout_block=m["readout_block"],
                 generated=m["generated"], truncated=m["truncated"],
                 config=m.get("config", {}))


class Runner:
    """A loaded model, its cone at every lens depth, and the pass that reads them together."""

    def __init__(self, model: str | ModelSpec, device: str = "cuda", dtype=None,
                 blocks: list[int] | None = None, attn_implementation: str = "sdpa",
                 device_map=None, verbose: bool = True):
        require_extra("run")
        torch = require("torch")
        self.spec = _model(model)
        self.device = device
        self.verbose = verbose
        self._log(f"building cones at every lens block for {self.spec.key} …")
        self.stack: _axes.ConeStack = _basis.field(self.spec, blocks=blocks,
                                                   verbose=verbose)
        self._log(f"loading {self.spec.hf} …")
        AutoModelForCausalLM = require("transformers").AutoModelForCausalLM
        AutoTokenizer = require("transformers").AutoTokenizer
        kw = {"revision": self.spec.revision} if self.spec.revision else {}
        self.tok = AutoTokenizer.from_pretrained(self.spec.hf, **kw)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.spec.hf, dtype=dtype or torch.bfloat16,
            device_map=device_map if device_map is not None else device,
            attn_implementation=attn_implementation, **kw).eval()
        self._layers = _decoder_blocks(self.model)
        self.open_id, self.close_id = _think_ids(self.tok, self.spec)
        self._U = {b: torch.tensor(self.stack.U[i], device=self._param_device())
                   for i, b in enumerate(self.stack.blocks)}
        self._log(f"ready: {len(self.stack.blocks)} blocks, {self.stack.n_axes} axes "
                  f"({', '.join(self.stack.names)}), think ids "
                  f"open={self.open_id} close={self.close_id}")

    def _log(self, msg: str) -> None:
        if self.verbose:
            print(msg, flush=True)

    def _param_device(self):
        return next(self.model.parameters()).device

    # ------------------------------------------------------------------ prompting
    def prompt_ids(self, prompt: str, thinking: bool = True) -> list[int]:
        """Prompt token ids through the model's own chat template."""
        msgs = [{"role": "user", "content": prompt}]
        for kw in ({"enable_thinking": bool(thinking)}, {}):
            try:
                ids = self.tok.apply_chat_template(msgs, tokenize=True,
                                                   add_generation_prompt=True,
                                                   return_dict=False, **kw)
            except TypeError:
                continue
            return _flatten(ids)
        raise RuntimeError("the chat template rejected both call forms")

    # ------------------------------------------------------------------ the passes
    def generate(self, ids: list[int], max_new_tokens: int = 1024, seed: int = 0,
                 **sampling) -> list[int]:
        """Sample a continuation. Returns prompt + generation."""
        torch = require("torch")
        cfg = dict(self.spec.sampling)
        cfg.update(sampling)
        x = torch.tensor([ids], device=self._param_device())
        torch.manual_seed(int(seed))
        with torch.no_grad():
            out = self.model.generate(x, attention_mask=torch.ones_like(x),
                                      max_new_tokens=int(max_new_tokens), do_sample=True,
                                      pad_token_id=self.tok.eos_token_id, **cfg)
        return [int(t) for t in out[0]]

    def capture(self, ids: list[int], prompt_len: int, prompt: str = "",
                config: dict | None = None) -> Trace:
        """One teacher-forced pass: the field at every lens block, and every entropy."""
        torch = require("torch")
        dev = self._param_device()
        proj: dict[int, np.ndarray] = {}
        hooks = []

        def mk(b):
            def hook(_m, _i, o):
                h = (o[0] if isinstance(o, tuple) else o)[0].float()
                h = h / (h.norm(dim=-1, keepdim=True) + 1e-9)
                proj[b] = (h @ self._U[b]).cpu().numpy().astype(np.float32)
            return hook

        for b in self.stack.blocks:
            hooks.append(self._layers[b].register_forward_hook(mk(b)))
        x = torch.tensor([ids], device=dev)
        try:
            with torch.no_grad():
                h = _decoder(self.model)(input_ids=x, use_cache=False).last_hidden_state[0]
                ent = self._entropies(h)
        finally:
            for hk in hooks:
                hk.remove()
        del h
        torch.cuda.empty_cache() if dev.type == "cuda" else None

        F = np.stack([proj[b] for b in self.stack.blocks], 0)       # [L, T, k]
        eot, truncated = _segment(ids, prompt_len, self.open_id, self.close_id)
        strs = [self.tok.decode([t]) for t in ids]
        return Trace(model=self.spec.key, prompt=prompt, tokens=list(ids), token_strs=strs,
                     prompt_len=prompt_len, eot=eot, entropy=ent, field=F,
                     blocks=list(self.stack.blocks), names=list(self.stack.names),
                     readout_block=self.spec.readout_block,
                     generated=self.tok.decode(ids[prompt_len:]), truncated=truncated,
                     config=dict(config or {}))

    def _entropies(self, h) -> np.ndarray:
        """`[T]` entropy of the distribution that produced each token; position 0 is NaN.

        The distribution for token `i` lives at position `i - 1`, so the array is shifted by
        one and the first token, which nothing predicted, has no entropy.
        """
        torch = require("torch")
        T = int(h.shape[0])
        out = np.full(T, np.nan, np.float32)
        head = self.model.get_output_embeddings()
        cap = getattr(_text_config(self.model), "final_logit_softcapping", None)
        with torch.no_grad():
            for i in range(0, T - 1, LOGIT_CHUNK):
                j = min(i + LOGIT_CHUNK, T - 1)
                z = head(h[i:j].to(self.model.dtype)).float()
                if cap:
                    z = torch.tanh(z / cap) * cap
                g = torch.log_softmax(z, dim=-1)
                out[i + 1:j + 1] = (-(g.exp() * g).sum(-1)).cpu().numpy()
        return out

    def run(self, prompt: str, max_new_tokens: int = 1024, seed: int = 0,
            thinking: bool = True, **sampling) -> Trace:
        """Generate from `prompt`, then capture the field over the whole sequence."""
        pre = self.prompt_ids(prompt, thinking)
        self._log(f"generating up to {max_new_tokens} tokens from {len(pre)} prompt tokens …")
        ids = self.generate(pre, max_new_tokens=max_new_tokens, seed=seed, **sampling)
        self._log(f"teacher-forced pass over {len(ids)} tokens × "
                  f"{len(self.stack.blocks)} blocks …")
        cfg = dict(self.spec.sampling)
        cfg.update(sampling)
        return self.capture(ids, len(pre), prompt=prompt,
                            config={"seed": seed, "max_new_tokens": max_new_tokens,
                                    "thinking": thinking, "sampling": cfg,
                                    "hf": self.spec.hf, "revision": self.spec.revision})


# ---------------------------------------------------------------------- helpers

def _flatten(ids) -> list[int]:
    if hasattr(ids, "input_ids"):
        ids = ids.input_ids
    ids = list(ids)
    if ids and isinstance(ids[0], (list, tuple)):
        ids = list(ids[0])
    return [int(i) for i in ids]


def _text_config(model):
    return getattr(model.config, "text_config", model.config)


def _decoder(model):
    return getattr(model, "model", model)


def _decoder_blocks(model) -> dict[int, object]:
    """`{index: block}` for the deepest `.layers.N` group, which is the decoder stack."""
    groups: dict[str, dict[int, object]] = {}
    for name, mod in model.named_modules():
        m = re.match(r"(.*)\.layers\.(\d+)$", name)
        if m:
            groups.setdefault(m.group(1), {})[int(m.group(2))] = mod
    if not groups:
        raise RuntimeError("no `.layers.N` modules found; cannot locate the decoder blocks")
    return groups[max(groups, key=lambda k: len(groups[k]))]


def _think_ids(tok, spec: ModelSpec) -> tuple[int | None, int | None]:
    """(open, close) token ids of the thinking channel, or None when not single tokens."""
    def one(s: str) -> int | None:
        ids = tok.encode(s, add_special_tokens=False)
        return int(ids[0]) if len(ids) == 1 else None
    return one(spec.think_start), one(spec.think_end)


def _segment(ids: list[int], prompt_len: int, open_id: int | None, close_id: int | None,
             ) -> tuple[int, bool]:
    """(absolute index of the end-of-think token, whether the block was never closed).

    The chat template may already have opened the thinking channel at the end of the prompt,
    which is the usual case; otherwise the opener is looked for in the generation. A block
    that is opened and never closed runs to the last token, and the whole generation is
    thinking — the case a token budget produces, and the one that must not be read as a
    zero-token answer.
    """
    head, gen = ids[:prompt_len], ids[prompt_len:]
    if close_id is None:
        return len(ids) - 1, True
    last_open = max((i for i, t in enumerate(head) if t == open_id), default=-1)
    last_close = max((i for i, t in enumerate(head) if t == close_id), default=-1)
    if open_id is not None and last_open > last_close:
        start = 0                                       # already open at the prompt's end
    else:
        start = next((i + 1 for i, t in enumerate(gen) if t == open_id), None)
        if start is None:
            return max(prompt_len - 1, 0), False        # no block at all: it is all answer
    rest = gen[start:]
    if close_id in rest:
        return prompt_len + start + rest.index(close_id), False
    return len(ids) - 1, True