"""
J-lens: Jacobian Lens for Transformer Models

Replicates Anthropic's technique from:
"Verbalizable Representations Form a Global Workspace in Language Models"
https://transformer-circuits.pub/2026/workspace/index.html

Core idea: For each token in the vocabulary, compute the average gradient
of log p(token) with respect to the residual stream at each layer,
averaged over many contexts. This reveals which concepts are "verbalizable"
— readily available for the model to report on.

Usage:
    python jlens.py --model checkpoints/ckpt.pt --data data/shakespeare_char
"""

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import numpy as np
import argparse
import json
import os
import pickle
from pathlib import Path
from collections import defaultdict


def load_model(checkpoint_path, model_class, device='cuda'):
    """Load a trained nanoGPT model from checkpoint."""
    checkpoint = torch.load(checkpoint_path, map_location=device)
    # nanoGPT stores model args, state_dict + optimizer in checkpoint
    model_args = checkpoint['model_args']
    
    # Create model with saved config
    model = model_class(model_args)
    
    # Fix state dict keys (nanoGPT wraps in DataParallel)
    state_dict = checkpoint['model']
    unwanted_prefix = '_orig_mod.'
    for k in list(state_dict.keys()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model, model_args


def compute_jlens_single_token(model, token_id, dataloader, layer_idx, device='cuda'):
    """
    Compute J-lens vector for a single token at a specific layer.
    
    J_l(token, layer) = E_x [ ∇_{resid[layer]} log p(token | x) ]
    
    Where the expectation is taken over all positions in the corpus.
    """
    vectors = []
    
    with torch.no_grad():
        for batch_idx, (x, y) in enumerate(dataloader):
            x, y = x.to(device), y.to(device)
            B, T = x.shape
            
            # We need gradients, so we'll do forward passes with hooks
            # Strategy: use torch.autograd.grad on a forward pass
            # where we capture residual stream activations

            # Register hook to capture residual stream at target layer
            activations = {}
            
            def make_hook():
                def hook(module, input, output):
                    # output is (B, T, n_embd)
                    # Detach then require grad so we can compute gradient through it
                    activations['resid'] = output.detach().requires_grad_(True)
                    return activations['resid']
                return hook
            
            # Find the target layer
            target_block = model.transformer.h[layer_idx]
            # nanoGPT architecture: h = x + attn(ln1(x)), then x = h + mlp(ln2(h))
            # We want the residual stream AFTER the attention + MLP of this layer
            # which is the output of the block
            
            handle = target_block.register_forward_hook(make_hook())
            
            # Forward pass
            logits, loss = model(x, y)
            
            handle.remove()
            
            # Now compute gradient of log p(token_id) w.r.t. residual stream
            # log p(token_id) at each position = log_softmax(logits)[:, :, token_id]
            log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
            token_log_probs = log_probs[:, :, token_id].sum()  # sum over B, T
            
            # Gradient of this sum w.r.t. the captured activations
            grad = torch.autograd.grad(
                token_log_probs, 
                activations['resid'],
                retain_graph=False
            )[0]  # Shape: (B, T, n_embd)
            
            vectors.append(grad.detach().cpu())
            
            # Cleanup
            del logits, loss, log_probs, activations, grad
            torch.cuda.empty_cache()
    
    # Average over all positions in the corpus
    all_vectors = torch.cat([v.reshape(-1, v.shape[-1]) for v in vectors], dim=0)
    jlens_vector = all_vectors.mean(dim=0)  # Shape: (n_embd,)
    
    return jlens_vector


def compute_jlens_all_tokens(model, dataloader, layer_idx, vocab_size, device='cuda'):
    """
    Compute J-lens vectors for all tokens at a specific layer.
    
    Returns: dict mapping token_id -> jlens_vector (n_embd,)
    """
    jlens_vectors = {}
    
    for token_id in range(vocab_size):
        vec = compute_jlens_single_token(model, token_id, dataloader, layer_idx, device)
        jlens_vectors[token_id] = vec
        
        if (token_id + 1) % 10 == 0:
            print(f"  Token {token_id + 1}/{vocab_size} done")
    
    return jlens_vectors


def compute_jlens_all_layers(model, dataloader, n_layers, vocab_size, device='cuda', 
                              use_batched=True):
    """
    Compute J-lens vectors for all layers and all tokens.
    
    Uses batched approach: for each context, compute gradients for ALL tokens
    at once using vector-Jacobian products. Much faster than per-token.
    
    Returns: dict mapping layer_idx -> {token_id: jlens_vector}
    """
    all_layer_vectors = defaultdict(dict)
    
    if use_batched:
        # Optimized: compute all token J-lens vectors simultaneously
        # For each context position, the gradient of log p(token) w.r.t. resid
        # for all tokens is just the Jacobian of the unembedding layer
        # which equals W_U^T * (one_hot(token) - softmax(logits))
        # Wait, let me think about this more carefully...
        
        print("Using batched J-lens computation...")
        
        for layer_idx in range(n_layers):
            print(f"\nLayer {layer_idx}/{n_layers}...")
            
            layer_accum = torch.zeros(vocab_size, model.config.n_embd, device='cpu')
            token_count = torch.zeros(vocab_size, device='cpu')
            
            with torch.no_grad():
                for batch_idx, (x, y) in enumerate(dataloader):
                    x, y = x.to(device), y.to(device)
                    B, T = x.shape
                    
                    # Capture residual stream at target layer
                    resid_captured = {}
                    
                    def make_hook(resid_dict):
                        def hook(module, input, output):
                            resid_dict['val'] = output.detach().requires_grad_(True)
                            return resid_dict['val']
                        return hook
                    
                    target_block = model.transformer.h[layer_idx]
                    handle = target_block.register_forward_hook(make_hook(resid_captured))
                    
                    logits, loss = model(x, y)
                    handle.remove()
                    
                    # Now: for each token in vocab, we want d(logit_t)/d(resid)
                    # This is the Jacobian of unembedding w.r.t. residual stream
                    # Chain rule: d(logit_t)/d(resid) = W_U[t, :] * d(layer_out)/d(resid)
                    # where layer_out is the final layer output after all remaining layers
                    # plus the direct path through the residual stream.
                    
                    # Actually, since we captured resid at layer L, and the model
                    # applies remaining layers resid_L -> ... -> resid_final -> logits,
                    # the gradient d(logits)/d(resid_L) = d(logits)/d(resid_final) * d(resid_final)/d(resid_L)
                    # 
                    # We can compute this by:
                    # 1. Get logits
                    # 2. For EACH position, compute gradient of logit for EACH token
                    #    w.r.t. the captured residual stream
                    # 3. Average across positions
                    
                    # Vectorized approach: compute gradients for ALL tokens simultaneously
                    # using torch.autograd.grad with list of outputs
                    
                    # For efficiency, compute per position, then aggregate
                    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)  # (B, T, vocab)
                    
                    # For each position (b, t), we need jacobian of log_probs[b,t,:] w.r.t. resid[b,t,:]
                    # This is (vocab, n_embd) per position
                    # We can batch by computing gradient of sum_{tokens} a_i * log_p(token_i)
                    # where a_i cycles through standard basis vectors
                    
                    # Practical approach for small vocab (nanoGPT: 65 tokens):
                    # Just loop over tokens, compute gradient, and accumulate
                    
                    resid = resid_captured['val']  # (B, T, n_embd)
                    
                    for token_id in range(vocab_size):
                        # Gradient of log_p(token_id) summed over all positions
                        token_log_prob = log_probs[:, :, token_id].sum()
                        
                        grad = torch.autograd.grad(
                            token_log_prob, resid, retain_graph=(token_id < vocab_size - 1)
                        )[0]  # (B, T, n_embd)
                        
                        # Accumulate: sum of gradients across all positions
                        layer_accum[token_id] += grad.detach().cpu().reshape(-1, model.config.n_embd).sum(dim=0)
                        token_count[token_id] += B * T
                    
                    del logits, loss, log_probs, resid
                    del grad  # pyright: ignore[reportPossiblyUnboundVariable]
                    torch.cuda.empty_cache()
                    
                    if (batch_idx + 1) % 10 == 0:
                        print(f"  Batch {batch_idx + 1}/{len(dataloader)}")
            
            # Average: divide sum by count
            for token_id in range(vocab_size):
                if token_count[token_id] > 0:
                    all_layer_vectors[layer_idx][token_id] = layer_accum[token_id] / token_count[token_id]
                else:
                    all_layer_vectors[layer_idx][token_id] = torch.zeros(model.config.n_embd)
            
            print(f"  Layer {layer_idx} complete. Saved {vocab_size} token vectors.")
    
    return dict(all_layer_vectors)


def save_jlens(jlens_data, output_path, metadata=None):
    """Save J-lens vectors to disk."""
    output = {
        'metadata': metadata or {},
        'vectors': {
            str(layer): {
                str(token_id): vec.numpy() for token_id, vec in tokens.items()
            }
            for layer, tokens in jlens_data.items()
        }
    }
    
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with open(output_path, 'wb') as f:
        pickle.dump(output, f)
    
    print(f"Saved J-lens data to {output_path}")


def load_jlens(path):
    """Load saved J-lens vectors."""
    with open(path, 'rb') as f:
        data = pickle.load(f)
    
    # Convert back to tensors
    jlens = {}
    for layer_str, tokens in data['vectors'].items():
        layer = int(layer_str)
        jlens[layer] = {}
        for token_id_str, vec in tokens.items():
            jlens[layer][int(token_id_str)] = torch.from_numpy(vec)
    
    return jlens, data['metadata']


def analyze_jlens(jlens_data, itos, n_layers, output_dir='outputs'):
    """Analyze and visualize J-lens vectors."""
    os.makedirs(output_dir, exist_ok=True)
    vocab_size = len(itos)
    
    print(f"\n{'='*60}")
    print("J-LENS ANALYSIS")
    print(f"{'='*60}")
    
    for layer_idx in range(n_layers):
        if layer_idx not in jlens_data:
            continue
        
        layer_vectors = jlens_data[layer_idx]
        
        # Compute norm of each token's J-lens vector
        norms = {}
        for token_id, vec in layer_vectors.items():
            norms[token_id] = vec.norm().item()
        
        # Sort by norm (most "verbalizable" tokens first)
        sorted_tokens = sorted(norms.items(), key=lambda x: x[1], reverse=True)
        
        print(f"\n--- Layer {layer_idx} ---")
        print(f"Top 10 most verbalizable tokens:")
        for token_id, norm in sorted_tokens[:10]:
            token_str = itos[token_id].replace('\n', '\\n')
            print(f"  '{token_str}': norm={norm:.4f}")
        
        print(f"Bottom 5 least verbalizable tokens:")
        for token_id, norm in sorted_tokens[-5:]:
            token_str = itos[token_id].replace('\n', '\\n')
            print(f"  '{token_str}': norm={norm:.4f}")
    
    # Compute J-space "capacity" — how many tokens have significant norm?
    print(f"\n--- J-space Capacity ---")
    for layer_idx in range(n_layers):
        if layer_idx not in jlens_data:
            continue
        layer_vectors = jlens_data[layer_idx]
        norms = torch.tensor([v.norm().item() for v in layer_vectors.values()])
        
        # Count "active" tokens (norm > median * 2)
        threshold = norms.median() * 2
        active = (norms > threshold).sum().item()
        print(f"  Layer {layer_idx}: {active}/{vocab_size} tokens active (threshold={threshold:.4f})")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='J-lens: Jacobian Lens for nanoGPT')
    parser.add_argument('--checkpoint', type=str, required=True, 
                        help='Path to model checkpoint')
    parser.add_argument('--data_dir', type=str, default='data/shakespeare_char',
                        help='Path to data directory')
    parser.add_argument('--output_dir', type=str, default='outputs/jlens',
                        help='Directory for saving outputs')
    parser.add_argument('--batch_size', type=int, default=32,
                        help='Batch size for processing')
    parser.add_argument('--max_batches', type=int, default=100,
                        help='Max batches to process (limit for speed)')
    parser.add_argument('--layers', type=str, default=None,
                        help='Comma-separated layer indices (default: all)')
    parser.add_argument('--device', type=str, default='cuda',
                        help='Device to use')
    
    args = parser.parse_args()
    
    # Import model from project root
    import sys
    # jlens.py is in src/, model.py is in repo root
    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    sys.path.insert(0, project_root)
    from model import GPT, GPTConfig
    
    # Load model
    print(f"Loading model from {args.checkpoint}")
    model, model_args = load_model(args.checkpoint, GPT, args.device)
    print(f"Model: {model_args.n_layer} layers, {model_args.n_embd} dim, "
          f"{model_args.n_head} heads, {model_args.vocab_size} vocab")
    
    # Load data
    data_dir = Path(args.data_dir)
    train_data = np.memmap(data_dir / 'train.bin', dtype=np.uint16, mode='r')
    val_data = np.memmap(data_dir / 'val.bin', dtype=np.uint16, mode='r')
    
    # Load vocab mappings
    meta_path = data_dir / 'meta.pkl'
    if meta_path.exists():
        with open(meta_path, 'rb') as f:
            meta = pickle.load(f)
        itos = meta['itos']
        stoi = meta['stoi']
    else:
        # Default char-level vocab
        chars = sorted(list(set(open(data_dir / 'input.txt').read())))
        stoi = {ch: i for i, ch in enumerate(chars)}
        itos = {i: ch for i, ch in enumerate(chars)}
    
    print(f"Vocabulary size: {len(itos)}")
    print(f"Train data: {len(train_data):,} tokens")
    
    # Create dataloader
    def get_batch(split):
        data = train_data if split == 'train' else val_data
        block_size = model_args.block_size
        ix = torch.randint(len(data) - block_size, (args.batch_size,))
        x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) 
                         for i in ix])
        y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) 
                         for i in ix])
        return x, y
    
    class SimpleDataset(torch.utils.data.IterableDataset):
        def __iter__(self):
            while True:
                yield get_batch('train')
    
    dataset = SimpleDataset()
    dataloader = DataLoader(dataset, batch_size=None, num_workers=0)
    
    # Limit to max_batches
    limited_dataloader = []
    for i, batch in enumerate(dataloader):
        if i >= args.max_batches:
            break
        limited_dataloader.append(batch)
    
    print(f"Processing {len(limited_dataloader)} batches of size {args.batch_size}")
    
    # Determine layers to process
    if args.layers:
        layers_to_process = [int(l) for l in args.layers.split(',')]
    else:
        layers_to_process = list(range(model_args.n_layer))
    
    print(f"Computing J-lens for layers: {layers_to_process}")
    
    # Compute J-lens for selected layers
    jlens_data = {}
    for layer_idx in layers_to_process:
        print(f"\nComputing J-lens for layer {layer_idx}...")
        layer_vectors = compute_jlens_all_tokens(
            model, limited_dataloader, layer_idx, 
            model_args.vocab_size, args.device
        )
        jlens_data[layer_idx] = layer_vectors
    
    # Save results
    save_path = os.path.join(args.output_dir, 'jlens_vectors.pkl')
    metadata = {
        'model_args': vars(model_args),
        'num_batches': len(limited_dataloader),
        'batch_size': args.batch_size,
        'layers_processed': layers_to_process,
        'vocab_size': model_args.vocab_size,
    }
    save_jlens(jlens_data, save_path, metadata)
    
    # Analyze
    analyze_jlens(jlens_data, itos, model_args.n_layer, args.output_dir)
    
    print(f"\nDone! Results saved to {args.output_dir}")