"""
Controlled frequency experiment for J-lens.

Tests causality: if we artificially double the frequency of a character,
does its J-lens norm predictably drop?

Hypothesis (from chain rule):
  J-lens norm ∝ (1 - p_avg), where p_avg is average predicted probability.
  Doubling frequency → model learns higher p → J-lens norm drops.

Control: train two identical models on:
  A) Original Shakespeare (baseline)
  B) Modified Shakespeare with 'q' frequency doubled
"""

import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def create_modified_dataset(input_path, output_path, target_char='q', factor=2.0):
    """Double the frequency of target_char in the text.
    Inserts target_char at random positions until frequency is doubled.
    """
    import random
    random.seed(42)
    
    with open(input_path, 'r') as f:
        text = f.read()
    
    target_count = text.count(target_char)
    target_freq = target_count / len(text)
    print(f"Original: {target_count} occurrences of '{target_char}', "
          f"frequency = {target_freq:.4%}")
    
    # Insert additional copies at random positions
    extra_needed = int(target_count * (factor - 1))
    positions = sorted(random.sample(range(len(text)), extra_needed))
    
    modified = list(text)
    for i, pos in enumerate(positions):
        modified.insert(pos + i, target_char)  # offset by previous insertions
    
    modified_text = ''.join(modified)
    new_count = modified_text.count(target_char)
    new_freq = new_count / len(modified_text)
    print(f"Modified: {new_count} occurrences of '{target_char}', "
          f"frequency = {new_freq:.4%}")
    
    with open(output_path, 'w') as f:
        f.write(modified_text)
    
    return modified_text


def prepare_data(input_path, output_dir):
    """Run nanoGPT's prepare.py on a text file."""
    import subprocess
    # Write a temp prepare script
    import os
    os.makedirs(output_dir, exist_ok=True)
    
    # Read the text
    with open(input_path, 'r') as f:
        text = f.read()
    
    chars = sorted(list(set(text)))
    vocab_size = len(chars)
    stoi = {ch: i for i, ch in enumerate(chars)}
    itos = {i: ch for i, ch in enumerate(chars)}
    
    # Encode
    import numpy as np
    data = np.array([stoi[ch] for ch in text], dtype=np.uint16)
    n = int(0.9 * len(data))
    train_data = data[:n]
    val_data = data[n:]
    
    train_data.tofile(os.path.join(output_dir, 'train.bin'))
    val_data.tofile(os.path.join(output_dir, 'val.bin'))
    
    import pickle
    with open(os.path.join(output_dir, 'meta.pkl'), 'wb') as f:
        pickle.dump({'stoi': stoi, 'itos': itos, 'vocab_size': vocab_size}, f)
    
    # Also save the raw text
    with open(os.path.join(output_dir, 'input.txt'), 'w') as f:
        f.write(text)
    
    print(f"Prepared {output_dir}: {len(train_data):,} train, "
          f"{len(val_data):,} val, {vocab_size} vocab")
    return stoi, itos, vocab_size


def train_model(data_dir, out_dir, device='cuda'):
    """Train nanoGPT on a prepared dataset."""
    import subprocess
    
    # Use the config but override dataset path and output
    import torch
    from model import GPT, GPTConfig
    
    # Load and train
    script = f"""
import sys
sys.path.insert(0, '.')
import torch
import numpy as np
import pickle
import os
from model import GPT, GPTConfig

# Load data
data_dir = '{data_dir}'
train_data = np.memmap(f'{{data_dir}}/train.bin', dtype=np.uint16, mode='r')
val_data = np.memmap(f'{{data_dir}}/val.bin', dtype=np.uint16, mode='r')
with open(f'{{data_dir}}/meta.pkl', 'rb') as f:
    meta = pickle.load(f)

# Config — same as train_shakespeare_char but smaller for speed
out_dir = '{out_dir}'
eval_interval = 500
eval_iters = 100
log_interval = 100
always_save_checkpoint = False
wandb_log = False
dataset = 'custom'
gradient_accumulation_steps = 1
batch_size = 32
block_size = 128
n_layer = 6
n_head = 6
n_embd = 384
dropout = 0.2
learning_rate = 1e-3
max_iters = 5000
lr_decay_iters = 5000
min_lr = 1e-4
beta2 = 0.99
warmup_iters = 100
dtype = 'float32'
flash = False
device = '{device}'
compile = False
vocab_size = meta['vocab_size']

# Build model
model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd,
                  block_size=block_size, bias=False, vocab_size=vocab_size,
                  dropout=dropout)
gptconf = GPTConfig(**model_args)
model = GPT(gptconf)
model.to(device)

# Optimizer
optimizer = model.configure_optimizers(weight_decay=0.1, learning_rate=learning_rate,
                                        betas=(0.9, beta2), device_type=device)
scaler = torch.amp.GradScaler('cuda', enabled=False)

def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (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])
    x, y = x.to(device), y.to(device)
    return x, y

@torch.no_grad()
def estimate_loss():
    out = {{}}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

print(f"Training model: {{model_args}}")
print(f"Parameters: {{sum(p.numel() for p in model.parameters())/1e6:.2f}}M")

os.makedirs(out_dir, exist_ok=True)
best_val_loss = 1e9

for iter_num in range(max_iters):
    if iter_num % eval_interval == 0:
        losses = estimate_loss()
        print(f"step {{iter_num}}: train loss {{losses['train']:.4f}}, val loss {{losses['val']:.4f}}")
        if losses['val'] < best_val_loss:
            best_val_loss = losses['val']
            checkpoint = {{
                'model': model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'model_args': model_args,
                'iter_num': iter_num,
                'best_val_loss': best_val_loss,
            }}
            torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))
    
    X, Y = get_batch('train')
    with torch.amp.autocast(device_type=device, dtype=torch.float32):
        logits, loss = model(X, Y)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)
    
    if iter_num % log_interval == 0:
        print(f"iter {{iter_num}}: loss {{loss.item():.4f}}")

print(f"Training complete. Best val loss: {{best_val_loss:.4f}}")
"""
    
    import tempfile
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
        f.write(script)
        script_path = f.name
    
    result = subprocess.run(['python3', script_path], capture_output=True, text=True)
    os.unlink(script_path)
    print(result.stdout)
    if result.returncode != 0:
        print("STDERR:", result.stderr)
    return result.returncode == 0


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', default='data/shakespeare_char/input.txt')
    parser.add_argument('--target_char', default='q')
    parser.add_argument('--factor', type=float, default=2.0)
    parser.add_argument('--skip_train', action='store_true')
    parser.add_argument('--device', default='cuda')
    args = parser.parse_args()
    
    # 1. Create modified dataset
    print("=" * 60)
    print("STEP 1: Creating modified dataset")
    print("=" * 60)
    modified_input = f'data/freq_experiment/modified_{args.target_char}.txt'
    os.makedirs('data/freq_experiment', exist_ok=True)
    create_modified_dataset(args.input, modified_input, args.target_char, args.factor)
    
    # 2. Prepare both datasets
    print()
    print("=" * 60)
    print("STEP 2: Preparing datasets")
    print("=" * 60)
    
    # Control: already prepared as data/shakespeare_char/
    # Modified: prepare from modified text
    mod_dir = f'data/freq_experiment/modified_{args.target_char}'
    stoi_mod, itos_mod, vocab_mod = prepare_data(modified_input, mod_dir)
    
    # 3. Train both models
    if not args.skip_train:
        print()
        print("=" * 60)
        print("STEP 3: Training CONTROL model (original Shakespeare)")
        print("=" * 60)
        # Control already trained as out-shakespeare-char/
        
        print()
        print("=" * 60)
        print(f"STEP 4: Training MODIFIED model ({args.target_char} x{args.factor})")
        print("=" * 60)
        train_model(mod_dir, f'out-freq-{args.target_char}', args.device)
    
    # 4. Run J-lens on both
    print()
    print("=" * 60)
    print("STEP 5: Running J-lens on both models")
    print("=" * 60)
    from jlens_v2 import compute_jlens_layer, load_model
    import numpy as np
    
    def run_jlens(ckpt_path, data_dir, label):
        print(f"\n  J-lens on {label}...")
        model, config = load_model(ckpt_path, args.device)
        
        # Build batches
        train_data = np.memmap(f'{data_dir}/train.bin', dtype=np.uint16, mode='r')
        n_batches = 10
        b_size = 16
        block_size = config.block_size
        batches = []
        for _ in range(n_batches):
            ix = torch.randint(len(train_data) - block_size, (b_size,))
            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, y))
        
        # Compute for middle layers only (speed)
        results = {}
        for layer_idx in [2, 3, 4]:
            results[layer_idx] = compute_jlens_layer(model, layer_idx, batches, args.device)
        
        return results
    
    # Control
    ctrl_jlens = run_jlens('out-shakespeare-char/ckpt.pt', 
                           'data/shakespeare_char', 'CONTROL')
    
    # Modified
    mod_jlens = run_jlens(f'out-freq-{args.target_char}/ckpt.pt',
                          mod_dir, 'MODIFIED')
    
    # 5. Compare
    print()
    print("=" * 60)
    print(f"STEP 6: COMPARISON — '{args.target_char}' J-lens norm change")
    print("=" * 60)
    
    with open('data/shakespeare_char/meta.pkl', 'rb') as f:
        import pickle
        ctrl_meta = pickle.load(f)
    ctrl_itos = ctrl_meta['itos']
    
    with open(f'{mod_dir}/meta.pkl', 'rb') as f:
        mod_meta = pickle.load(f)
    mod_itos = mod_meta['itos']
    
    target_id_ctrl = ctrl_itos.index(args.target_char)
    target_id_mod = mod_itos.index(args.target_char)
    
    for layer_idx in sorted(ctrl_jlens.keys()):
        ctrl_norm = ctrl_jlens[layer_idx][target_id_ctrl].norm().item()
        mod_norm = mod_jlens[layer_idx][target_id_mod].norm().item()
        change = (mod_norm - ctrl_norm) / ctrl_norm * 100
        direction = "↓" if change < 0 else "↑"
        print(f"  Layer {layer_idx}: control={ctrl_norm:.4f} → modified={mod_norm:.4f} "
              f"({direction}{abs(change):.1f}%)")
    
    print()
    print("Hypothesis confirmed!" if change < 0 else "Hypothesis NOT confirmed — unexpected result.")


if __name__ == '__main__':
    main()