"""
GPT-2 Small J-lens — efficient sampling approach.

Instead of computing J-lens for all 50257 tokens, sample:
- 100 most common + 100 rarest = 200 tokens for frequency test
- 1000 random tokens for effective rank estimation

Each backward pass takes ~50ms on K2200.
200 tokens × 3 batches × 3 layers = 1800 backward passes ≈ 90s
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch, numpy as np, pickle, time, subprocess

try:
    import transformers
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "transformers==4.44.0", "accelerate", "requests"])
    import transformers
import requests

device = 'cuda'

def main():
    print("=" * 60)
    print("GPT-2 SMALL J-LENS — Frequency + Dimensionality Tests")
    print("=" * 60)
    
    # Load model
    model = transformers.GPT2LMHeadModel.from_pretrained(
        "openai-community/gpt2", torch_dtype=torch.float32).to(device)
    model.eval()
    tokenizer = transformers.GPT2Tokenizer.from_pretrained("openai-community/gpt2")
    
    d_model = model.config.n_embd
    n_layers = model.config.n_layer
    vocab_size = model.config.vocab_size
    print(f"GPT-2 Small: {n_layers} layers, d={d_model}, V={vocab_size}, V/d={vocab_size/d_model:.1f}x")
    
    # Load corpus from local Shakespeare file
    with open('data/shakespeare_char/input.txt') as f:
        text = f.read()[:500000]
    tokens = tokenizer(text, return_tensors='np', truncation=True, max_length=2000)['input_ids'][0]
    print(f"Corpus: {len(tokens)} tokens, {len(set(tokens))} unique")
    
    # Estimate token frequencies in our corpus
    from collections import Counter
    freq = Counter(tokens.tolist())
    total = len(tokens)
    sorted_by_freq = sorted(freq.items(), key=lambda x: x[1], reverse=True)
    
    # Pick: 50 most common, 50 rarest
    sample_tokens = [t for t, _ in sorted_by_freq[:50]] + [t for t, _ in sorted_by_freq[-50:]]
    sample_tokens = list(set(sample_tokens))  # dedupe
    print(f"Sampling {len(sample_tokens)} tokens for frequency test")
    
    # J-lens computation
    n_batches = 3
    seq_len = 32
    layers_to_test = [3, 6, 9]  # early, middle, late
    
    # Accumulators: {layer: {token_id: [norm_sum, count]}}
    accum = {layer: {tid: [0.0, 0] for tid in sample_tokens} for layer in layers_to_test}
    
    for batch_i in range(n_batches):
        ix = torch.randint(0, len(tokens) - seq_len - 1, (1,))
        x = torch.from_numpy(tokens[ix[0]:ix[0]+seq_len].astype(np.int64)).unsqueeze(0).to(device)
        
        for layer_idx in layers_to_test:
            # Capture residual at this layer
            resid_captured = {}
            def hook(module, inp, out):
                h = out[0] if isinstance(out, tuple) else out
                resid_captured['val'] = h
            
            target = model.transformer.h[layer_idx]
            handle = target.register_forward_hook(hook)
            
            # Forward (no no_grad — we need gradients)
            result = model(x)
            handle.remove()
            
            logits = result.logits  # (1, seq_len, V)
            resid = resid_captured['val']  # (1, seq_len, d)
            
            for tid in sample_tokens:
                token_logprob = torch.nn.functional.log_softmax(logits, dim=-1)[:, :, tid].sum()
                try:
                    grad = torch.autograd.grad(token_logprob, resid, retain_graph=True)[0]
                    norm = grad.norm().item()
                    accum[layer_idx][tid][0] += norm
                    accum[layer_idx][tid][1] += 1
                except:
                    pass
            
            del logits, result, resid
            torch.cuda.empty_cache()
        
        print(f"  Batch {batch_i+1}/{n_batches} done")
    
    # Results
    print(f"\n{'='*60}")
    print("FREQUENCY vs J-LENS NORM (GPT-2 Small)")
    print(f"{'='*60}")
    
    for layer_idx in layers_to_test:
        print(f"\n--- Layer {layer_idx} ---")
        
        norms = {}
        for tid in sample_tokens:
            s, c = accum[layer_idx][tid]
            if c > 0:
                norms[tid] = s / c
        
        # Token frequencies
        freqs = {tid: freq.get(tid, 0)/total*100 for tid in norms}
        
        # Correlation
        n_arr = np.array(list(norms.values()))
        f_arr = np.array([freqs[t] for t in norms])
        corr = np.corrcoef(n_arr, f_arr)[0, 1]
        
        # Top/bottom by norm
        sorted_tokens = sorted(norms.items(), key=lambda x: x[1], reverse=True)
        print(f"  Top 5 by J-lens norm:")
        for tid, n in sorted_tokens[:5]:
            tok_str = tokenizer.decode([tid]).replace('\n', '\\n')
            print(f"    '{tok_str}' (freq={freqs[tid]:.3f}%): norm={n:.4f}")
        print(f"  Bottom 5 by J-lens norm:")
        for tid, n in sorted_tokens[-5:]:
            tok_str = tokenizer.decode([tid]).replace('\n', '\\n')
            print(f"    '{tok_str}' (freq={freqs[tid]:.3f}%): norm={n:.4f}")
        print(f"  Pearson r(norm, freq): {corr:.3f}")
    
    # Summary
    avg_corr = np.mean([np.corrcoef(
        np.array(list(accum[l][t][0]/max(accum[l][t][1],1) for t in sample_tokens)),
        np.array([freq.get(t,0)/total*100 for t in sample_tokens])
    )[0,1] for l in layers_to_test])
    
    print(f"\n{'='*60}")
    print("RESULTS SUMMARY")
    print(f"{'='*60}")
    print(f"  nanoGPT (V/d=0.2x): full-rank J-space, r=-0.65 freq correlation")
    print(f"  GPT-2   (V/d=65x):  avg r={avg_corr:.3f}")
    print(f"  HYPOTHESIS: frequency anti-correlation persists at scale")


if __name__ == '__main__':
    main()