"""
J-lens v3: FAITHFUL implementation of Anthropic's Jacobian lens.

Paper: "Verbalizable Representations Form a Global Workspace in Language Models"
https://transformer-circuits.pub/2026/workspace/index.html

Definition (from the paper's Methods):
    J_l = E_{t, t' >= t, prompt}[ d h_final,t' / d h_l,t ]      (d_model x d_model per layer)
    J-lens vectors at layer l = rows of W_U * J_l
    lens(h_l) = softmax(W_U * norm(J_l * h_l))

The average is over source positions t, all future positions t' >= t, and a corpus
of prompts. The result is a single d x d matrix per layer that maps the residual
stream at layer l to the FINAL residual stream, read out through the model's own
unembedding (W_U = lm_head.weight in nanoGPT).

This is different from the earlier proxy (jlens_v2) which differentiated
log_softmax(logits) directly, folding the token-dependent softmax calibration
into the averaged object.

Implementation: basis-chunked vector-Jacobian products. For a chunk C of basis
vectors {e_c}, compute
    S_c = sum_{t'} (h_final,t' . e_c)
and backpropagate S_c to h_l (the captured residual stream at layer l).
Because h_final,t' does not depend on h_l,t for t' < t (causality), summing over
all t' yields exactly sum over valid pairs. One batched autograd call per chunk
(is_grads_batched=True) gives all rows in the chunk.

Also implements the BOTH-WAYS comparison: old proxy ||E[grad log p(k|x)]|| vs
faithful ||(W_U J_l)[k]||, each correlated against token frequency, per layer.

Usage (inside Docker on meru):
    python3 src/jlens_v3.py --checkpoint out-shakespeare-char/ckpt.pt \
        --data_dir data/shakespeare_char --n_prompts 20 --layers 2,3,4
"""

import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import argparse
import torch
import numpy as np
import pickle

import jlens_v2  # for load_model


def make_batches(train_data, block_size, batch_size, n_prompts, device, seed=1337):
    """Random contiguous prompts from train.bin (n_prompts prompts of batch_size)."""
    g = torch.Generator().manual_seed(seed)
    batches = []
    for _ in range(n_prompts):
        ix = torch.randint(len(train_data) - block_size, (batch_size,), generator=g)
        x = torch.stack([torch.from_numpy(
            train_data[i:i + block_size].astype(np.int64)) for i in ix])
        y = torch.stack([torch.from_numpy(
            train_data[i + 1:i + 1 + block_size].astype(np.int64)) for i in ix])
        batches.append((x.to(device), y.to(device)))
    return batches


def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32):
    """
    Compute the faithful J-lens VECTORS (rows of W_U * J_l) directly.

    (W_U J_l)[k] = E_{t, t'>=t, prompt}[ d(W_U[k] . h_final,t') / d h_l,t ]
                 = E[ grad of the k-th RAW LOGIT (pre-softmax) w.r.t. h_l ]

    Instead of materializing the d x d residual Jacobian (384 basis VJPs), we
    probe with the unembedding rows W_U[k] (V=65 probes) — mathematically
    identical, ~6x cheaper. One batched autograd call per prompt per layer
    (is_grads_batched=True) computes all V token vectors at once.
    """
    d = model.config.n_embd
    V = model.config.vocab_size
    n_layer = model.config.n_layer
    last_block = model.transformer.h[n_layer - 1]
    W_U = model.lm_head.weight.detach().float()  # (V, d), nanoGPT weight tying

    accum = torch.zeros(V, d, device='cpu')
    n_pairs = 0

    for (x, y) in batches:
        B, T = x.shape
        resid_captured = {}

        def hook(module, inp, out):
            resid_captured['h_l'] = out

        def hook_final(module, inp, out):
            resid_captured['h_final'] = out

        h1 = model.transformer.h[layer_idx].register_forward_hook(hook)
        h2 = last_block.register_forward_hook(hook_final)
        logits, loss = model(x, y)           # graph connected
        h1.remove()
        h2.remove()

        h_l = resid_captured['h_l']          # (B, T, d)
        h_final = resid_captured['h_final']  # (B, T, d)
        # sum over all future positions t' (causal: t' < t contributes zero grad)
        h_final_sum = h_final.sum(dim=1)     # (B, d)

        # grad_outputs[v, b, :] = W_U[v] -> batched VJPs, chunked over vocab
        # (is_grads_batched with all V=65 probes at once OOMs the 4GB K2200)
        for v0 in range(0, V, chunk):
            C = min(chunk, V - v0)
            grad_outputs = (W_U[v0:v0 + C].to(device).unsqueeze(1)
                            .expand(C, B, d).contiguous())
            grads = torch.autograd.grad(
                h_final_sum, h_l, grad_outputs=grad_outputs,
                is_grads_batched=True,
                retain_graph=(v0 + C < V))[0]  # (C, B, T, d)
            accum[v0:v0 + C] += grads.detach().cpu().reshape(C, -1, d).sum(dim=1)
            del grads, grad_outputs
            torch.cuda.empty_cache()

        n_pairs += B * T * (T + 1) // 2      # valid (t, t' >= t) pairs
        del logits, loss, h_l, h_final, h_final_sum
        torch.cuda.empty_cache()

    accum /= n_pairs
    return accum  # (V, d): faithful J-lens vectors, rows of W_U * J_l


def compute_proxy_norms(model, layer_idx, batches, device, vocab_size, chunk=16):
    """
    Old proxy (jlens_v2 quantity): || E_x[ grad log_softmax(logits)[:, :, k] ] ||
    averaged over positions — norm of the mean gradient, per token.
    """
    d = model.config.n_embd
    accum = torch.zeros(vocab_size, d, device='cpu')
    counts = torch.zeros(vocab_size, device='cpu')
    for (x, y) in batches:
        B, T = x.shape
        resid_captured = {}

        def hook(module, inp, out):
            resid_captured['val'] = out

        handle = model.transformer.h[layer_idx].register_forward_hook(hook)
        logits, loss = model(x, y)
        handle.remove()
        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
        # sum over positions -> (V,); batched VJP over vocab, chunked for GPU memory
        lp_sum = log_probs.sum(dim=(0, 1))
        for v0 in range(0, vocab_size, chunk):
            C = min(chunk, vocab_size - v0)
            grad_outputs = torch.eye(vocab_size, device=device)[v0:v0 + C]
            grad = torch.autograd.grad(
                lp_sum, resid_captured['val'], grad_outputs=grad_outputs,
                is_grads_batched=True,
                retain_graph=(v0 + C < vocab_size))[0]  # (C, B, T, d)
            accum[v0:v0 + C] += grad.detach().cpu().reshape(C, -1, d).sum(dim=1)
            del grad, grad_outputs
            torch.cuda.empty_cache()
        counts += B * T
        del logits, loss, log_probs, lp_sum
        torch.cuda.empty_cache()
    norms = {}
    for token_id in range(vocab_size):
        norms[token_id] = (accum[token_id] / counts[token_id]).norm().item()
    return norms


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--checkpoint', required=True)
    ap.add_argument('--data_dir', default='data/shakespeare_char')
    ap.add_argument('--n_prompts', type=int, default=20)
    ap.add_argument('--batch_size', type=int, default=16)
    ap.add_argument('--layers', default='2,3,4', help='comma-separated layer indices')
    ap.add_argument('--device', default='cuda')
    ap.add_argument('--chunk', type=int, default=32)
    args = ap.parse_args()

    device = args.device
    model, config = jlens_v2.load_model(args.checkpoint, device)
    model.eval()
    d = config.n_embd
    V = config.vocab_size
    n_layer = config.n_layer
    print(f"Model: {n_layer} layers, d={d}, V={V}, V/d={V/d:.2f}x")

    train_data = np.memmap(f'{args.data_dir}/train.bin', dtype=np.uint16, mode='r')
    with open(f'{args.data_dir}/meta.pkl', 'rb') as f:
        meta = pickle.load(f)
    itos = meta['itos']

    # token frequencies
    counts = np.bincount(train_data, minlength=V).astype(float)
    freq = counts / counts.sum() * 100  # percent

    batches = make_batches(train_data, config.block_size, args.batch_size,
                           args.n_prompts, device)

    layers = [int(l) for l in args.layers.split(',')]
    W_U = model.lm_head.weight.detach().float()  # (V, d) — nanoGPT weight tying

    print(f"\n{'='*78}")
    print("BOTH-WAYS COMPARISON: old proxy vs faithful J-lens (rows of W_U * J_l)")
    print(f"{'='*78}")

    for layer_idx in layers:
        print(f"\n--- Layer {layer_idx} ---")
        print("  Computing faithful J-lens vectors (W_U-probed VJPs)...")
        faithful_vecs = compute_faithful_jlens(model, layer_idx, batches, device,
                                               args.chunk)  # (V, d)
        faithful_norms = {k: faithful_vecs[k].norm().item() for k in range(V)}

        if layer_idx == n_layer - 1:
            # Validation: J_{L-1} should be identity, so faithful vectors == W_U rows
            sims = torch.nn.functional.cosine_similarity(
                faithful_vecs.float(), W_U.float(), dim=1)
            print(f"  [validation] last layer: mean cos-sim(faithful, W_U rows) = "
                  f"{sims.mean().item():.4f} (expect ~1.0 if J=identity)")

        print("  Computing old proxy norms...")
        proxy_norms = compute_proxy_norms(model, layer_idx, batches, device, V,
                                          args.chunk)

        # correlations with frequency
        f_arr = np.array([freq[k] for k in range(V)])
        p_arr = np.array([proxy_norms[k] for k in range(V)])
        fa_arr = np.array([faithful_norms[k] for k in range(V)])
        r_proxy = np.corrcoef(p_arr, f_arr)[0, 1]
        r_faith = np.corrcoef(fa_arr, f_arr)[0, 1]

        print(f"  freq corr | old proxy:    r = {r_proxy:+.3f}")
        print(f"  freq corr | faithful:     r = {r_faith:+.3f}")

        # top/bottom by faithful norm
        srt = sorted(faithful_norms.items(), key=lambda kv: kv[1], reverse=True)
        def esc(s):
            return s.replace(chr(10), '\\n')
        print("  Top 5 by faithful norm:")
        for tid, n in srt[:5]:
            print(f"    '{esc(itos[tid])}' freq={freq[tid]:.3f}% norm={n:.4f}")
        print("  Bottom 5 by faithful norm:")
        for tid, n in srt[-5:]:
            print(f"    '{esc(itos[tid])}' freq={freq[tid]:.3f}% norm={n:.4f}")

        torch.save({'faithful_vecs': faithful_vecs, 'faithful_norms': faithful_norms,
                    'proxy_norms': proxy_norms},
                   f'outputs/jlens_v3_layer{layer_idx}.pt')
        print(f"  Saved outputs/jlens_v3_layer{layer_idx}.pt")

    print("\nDONE.")


if __name__ == '__main__':
    main()