"""Probing the model's internal "thoughts" (inspired by Anthropic's circuit work).

We can't do full attribution graphs cheaply, but two lightweight probes go a long
way on this architecture:

  logit_lens         apply the (tied) LM head to each Block-AttnRes block state to
                     see what the model "predicts" at each depth — the thought
                     sharpening from block to block.
  representation_report  per-block geometry of the concept space: rms norm,
                     anisotropy (mean pairwise cosine — high == degenerate), and
                     effective rank (how many dimensions are actually used).

Both consume the per-block states exposed by CortexLM._trunk(..., collect=True).
"""

from __future__ import annotations

import jax
import jax.numpy as jnp
import numpy as np


def block_states(model, tokens):
    _, states = model._trunk(jnp.asarray(tokens), collect=True)
    return states


def logit_lens(model, tokens, k: int = 5, pos: int = -1):
    """Top-k predictions + entropy at `pos` for every block state."""
    _, states = model._trunk(jnp.asarray(tokens), collect=True)
    out = []
    for i, h in enumerate(states):
        logits = model._head(h)[:, pos].astype(jnp.float32)        # [B, V]
        probs = jax.nn.softmax(logits, axis=-1)
        top_p, top_i = jax.lax.top_k(probs, k)
        ent = -jnp.sum(probs * jnp.log(probs + 1e-9), axis=-1)
        out.append({
            "stage": i,
            "top_ids": np.asarray(top_i),
            "top_probs": np.asarray(top_p),
            "entropy": np.asarray(ent),
        })
    return out


def representation_report(model, tokens, sample: int = 512, seed: int = 0):
    """Per-block geometry of the hidden states."""
    _, states = model._trunk(jnp.asarray(tokens), collect=True)
    rng = np.random.default_rng(seed)
    rows = []
    for i, h in enumerate(states):
        hf = np.asarray(h, dtype=np.float32).reshape(-1, h.shape[-1])
        if hf.shape[0] > sample:
            hf = hf[rng.choice(hf.shape[0], sample, replace=False)]
        rms = float(np.sqrt((hf ** 2).mean()))
        hn = hf / (np.linalg.norm(hf, axis=-1, keepdims=True) + 1e-8)
        sims = hn @ hn.T
        n = sims.shape[0]
        aniso = float((sims.sum() - np.trace(sims)) / (n * (n - 1)))   # mean off-diag cosine
        # effective rank = exp(entropy of normalized singular-value spectrum)
        s = np.linalg.svd(hf - hf.mean(0, keepdims=True), compute_uv=False)
        p = (s ** 2) / max((s ** 2).sum(), 1e-9)
        erank = float(np.exp(-(p * np.log(p + 1e-12)).sum()))
        rows.append({"stage": i, "rms": rms, "anisotropy": aniso, "eff_rank": erank})
    return rows


def format_report(rows) -> str:
    lines = [f"{'stage':>6} {'rms':>10} {'anisotropy':>12} {'eff_rank':>10}"]
    for r in rows:
        lines.append(f"{r['stage']:>6} {r['rms']:>10.3f} {r['anisotropy']:>12.3f} {r['eff_rank']:>10.1f}")
    return "\n".join(lines)