import torch
import torch.nn.functional as F
from torch.optim import Adam
from pathlib import Path
from datasets import load_dataset
from transformer_lens import HookedTransformer
from sae import SparseAutoencoder

MODEL_NAME      = "meta-llama/Llama-3.2-1B-Instruct"
SAE_LAYER       = 8
N_FEATURES      = 16384
D_MODEL         = 2048
L1_COEFF        = 10.0
LR              = 1e-4
BATCH_SIZE      = 2048
N_TOKENS        = 2_000_000
SEQ_LEN         = 64
SAVE_EVERY      = 2000
SAE_SAVE_PATH   = Path.home() / "inspection-backend" / "sae" / "sae_layer8.pt"
ACTS_CACHE_PATH = Path.home() / "inspection-backend" / "sae" / "acts_cache.pt"
NORM_CACHE_PATH = Path.home() / "inspection-backend" / "sae" / "acts_norm.pt"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def collect_activations(model, n_tokens, seq_len):
    if ACTS_CACHE_PATH.exists():
        print("[collect] loading cache...", flush=True)
        acts = torch.load(ACTS_CACHE_PATH, map_location="cpu")
        print(f"[collect] loaded {acts.shape[0]:,} activations", flush=True)
        return acts

    print(f"[collect] collecting {n_tokens:,} activations at layer {SAE_LAYER}...", flush=True)
    ACTS_CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)

    dataset = load_dataset("Skylion007/openwebtext", split="train", streaming=True)
    tokenizer = model.tokenizer
    tokenizer.padding_side = "right"

    chunk_size = 50_000
    chunk, chunk_files = [], []
    total, chunk_idx, buf = 0, 0, []

    model.eval()
    with torch.no_grad():
        for sample in dataset:
            if total >= n_tokens:
                break
            ids = tokenizer.encode(sample["text"][:1000], add_special_tokens=False)
            buf.extend(ids)

            while len(buf) >= seq_len and total < n_tokens:
                tokens = torch.tensor([buf[:seq_len]], device=DEVICE)
                buf = buf[seq_len:]
                _, cache = model.run_with_cache(
                    tokens,
                    names_filter=f"blocks.{SAE_LAYER}.hook_resid_post",
                    return_type=None,
                )
                chunk.append(cache[f"blocks.{SAE_LAYER}.hook_resid_post"][0].cpu())
                total += seq_len

                if len(chunk) >= chunk_size:
                    p = ACTS_CACHE_PATH.parent / f"chunk_{chunk_idx}.pt"
                    torch.save(torch.cat(chunk, dim=0), p)
                    chunk_files.append(p)
                    chunk, chunk_idx = [], chunk_idx + 1
                    print(f"[collect] {total:,} / {n_tokens:,} tokens", flush=True)

    if chunk:
        p = ACTS_CACHE_PATH.parent / f"chunk_{chunk_idx}.pt"
        torch.save(torch.cat(chunk, dim=0), p)
        chunk_files.append(p)

    print("[collect] merging chunks...", flush=True)
    all_acts = torch.cat([torch.load(p) for p in chunk_files], dim=0)[:n_tokens]
    torch.save(all_acts, ACTS_CACHE_PATH)
    for p in chunk_files:
        p.unlink()
    print(f"[collect] saved {all_acts.shape[0]:,} activations", flush=True)
    return all_acts


def normalize_activations(acts):
    if NORM_CACHE_PATH.exists():
        d = torch.load(NORM_CACHE_PATH)
        mean, std = d["mean"], d["std"]
    else:
        print("[norm] computing mean/std...", flush=True)
        mean = acts.mean(0)
        std = acts.std(0).clamp(min=1e-6)
        torch.save({"mean": mean, "std": std}, NORM_CACHE_PATH)
    normalized = (acts - mean) / std
    return normalized, mean, std


def train(acts):
    print(f"[train] normalizing activations...", flush=True)
    acts, mean, std = normalize_activations(acts)
    print(f"[train] acts mean={acts.mean():.4f} std={acts.std():.4f}", flush=True)
    print(f"[train] training on {acts.shape[0]:,} activations", flush=True)
    SAE_SAVE_PATH.parent.mkdir(parents=True, exist_ok=True)

    sae = SparseAutoencoder(d_model=D_MODEL, n_features=N_FEATURES).to(DEVICE)
    opt = Adam(sae.parameters(), lr=LR)
    n, step = acts.shape[0], 0

    with torch.no_grad():
        sae.b_pre.data = acts[:10000].to(DEVICE).mean(0)

    for epoch in range(1, 11):
        acts_shuffled = acts[torch.randperm(n)]
        for start in range(0, n - BATCH_SIZE, BATCH_SIZE):
            batch = acts_shuffled[start:start + BATCH_SIZE].to(DEVICE)
            f, x_hat = sae(batch)
            loss = F.mse_loss(x_hat, batch) + L1_COEFF * f.abs().mean()
            opt.zero_grad()
            loss.backward()
            opt.step()
            sae._normalise_decoder()
            step += 1

            if step % 200 == 0:
                dead = (f.max(0).values == 0).sum().item()
                l0 = (f > 0).float().sum(-1).mean().item()
                recon = F.mse_loss(x_hat, batch).item()
                print(f"[train] step={step} epoch={epoch} recon={recon:.4f} L0={l0:.1f} dead={dead}/{N_FEATURES}", flush=True)

            if step % SAVE_EVERY == 0:
                sae.save(SAE_SAVE_PATH)

        sae.save(SAE_SAVE_PATH)
        print(f"[train] epoch {epoch} done", flush=True)

    print("[train] complete", flush=True)
    return sae


if __name__ == "__main__":
    if not torch.cuda.is_available():
        raise RuntimeError("No GPU detected.")
    print(f"[main] loading {MODEL_NAME}...", flush=True)
    model = HookedTransformer.from_pretrained(MODEL_NAME)
    model.eval()
    model.to(DEVICE)
    acts = collect_activations(model, N_TOKENS, SEQ_LEN)
    del model
    torch.cuda.empty_cache()
    sae = train(acts)
    print("[main] done.", flush=True)
