"""
scripts/compute_density_map.py
------------------------------
Phase 1 + 2: Collect per-head activations and compute capability density map.

Steps:
  1. Load GPT-2 Medium
  2. Load calibration data (WikiText-103-raw-v1, 128 seq × 512 tok)
  3. Run calibration data through model, collect per-head attention outputs
  4. For each of 384 heads: train a TopK SAE, compute β, H, ψ, δ
  5. Save density map to .npz

Runtime: ~15–20 min on T4 GPU (384 SAEs × 5 epochs each)

Usage:
    python scripts/compute_density_map.py --output results/density_map.npz
"""

import argparse
import random
import os
import sys

import numpy as np
import torch

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

# ── Reproducibility ───────────────────────────────────────────────────────────
# Note: original experiments were run without fixed seeds.
# This seed ensures approximate reproducibility across runs.
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)

from transformers import GPT2Model, GPT2Tokenizer
from datasets import load_dataset
from cgc.density import collect_head_activations, compute_density_map


def parse_args():
    p = argparse.ArgumentParser(description="Compute CGC capability density map")
    p.add_argument("--output",       type=str,   default="results/density_map.npz",
                   help="Output path for density map")
    p.add_argument("--n_sequences",  type=int,   default=128,
                   help="Number of calibration sequences (default: 128)")
    p.add_argument("--seq_len",      type=int,   default=512,
                   help="Tokens per sequence (default: 512)")
    p.add_argument("--batch_size",   type=int,   default=8,
                   help="Sequences per forward pass (default: 8)")
    return p.parse_args()


def load_calibration_data(tokenizer, n_sequences, seq_len, device):
    """
    Load WikiText-103-raw-v1 training split as non-overlapping token chunks.
    Same dataset and split used for calibration in the CGC v1 paper.
    """
    print("Loading calibration data (WikiText-103-raw-v1, train split)...")
    dataset   = load_dataset("Salesforce/wikitext",
                             "wikitext-103-raw-v1", split="train")
    full_text = " ".join([
        item["text"].strip()
        for item in dataset
        if len(item["text"].strip()) > 50
    ])
    all_tokens = tokenizer.encode(full_text)

    chunks = []
    for i in range(n_sequences):
        start = i * seq_len
        end   = start + seq_len
        if end > len(all_tokens):
            print(f"Warning: only {len(chunks)} sequences available.")
            break
        chunks.append(torch.tensor(all_tokens[start:end], dtype=torch.long))

    data = torch.stack(chunks).to(device)
    print(f"Calibration data: {data.shape}  "
          f"({data.shape[0] * data.shape[1]:,} tokens total)")
    return data


def main():
    args = parse_args()
    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    if device.type == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        vram = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"VRAM: {vram:.1f} GB")

    # Load GPT-2 Medium
    print("\nLoading GPT-2 Medium...")
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
    tokenizer.pad_token = tokenizer.eos_token
    model = GPT2Model.from_pretrained("gpt2-medium")
    model = model.to(device)
    model.eval()
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Loaded. Layers={model.config.n_layer}  "
          f"Heads={model.config.n_head}  "
          f"Params={n_params:,}")

    # Calibration data
    cal_data = load_calibration_data(
        tokenizer, args.n_sequences, args.seq_len, device
    )

    # Phase 1: collect per-head activations
    head_acts = collect_head_activations(
        model, cal_data, device, args.batch_size
    )

    # Phase 2: train SAEs and compute density
    density_map, _ = compute_density_map(
        head_acts, device, seq_len=args.seq_len
    )

    # Save
    density_map.save(args.output)
    print(f"\nDensity map saved to: {args.output}")
    print("Next: python scripts/run_compression.py "
          f"--density_map {args.output}")


if __name__ == "__main__":
    main()