"""ARCHON Phase A integration test: load pretrain ckpt + wire M7 DoLa + bench.

Runs on V100 in /workspace/archon_sft_v2/.
"""
from __future__ import annotations

import sys
import time
import json
from pathlib import Path

import torch
import torch.nn.functional as F

# Add source/ to path for model imports
ROOT = Path("/workspace/archon_sft_v2")
sys.path.insert(0, str(ROOT / "source"))

from config import ArchonBrainConfig
from model import ArchonBrain

# Phase A modules
sys.path.insert(0, str(ROOT))
import m7_dola  # noqa: E402


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[ARCHON Phase A] device={device}")

    # === Build config matching the SFT v2 (vocab 32006 + 6 ChatML extension) ===
    cfg = ArchonBrainConfig()
    cfg.vocab_size = 32006  # tokenizer v3 with ChatML
    cfg.max_seq_len = 4096
    print(f"[ARCHON Phase A] vocab={cfg.vocab_size} layers={cfg.num_layers} dim={cfg.hidden_dim} MTP={cfg.mtp_heads}")

    # === Build model ===
    model = ArchonBrain(cfg).to(device).to(torch.bfloat16)
    print(f"[ARCHON Phase A] params={sum(p.numel() for p in model.parameters())/1e6:.1f}M")

    # === Load ckpt — surgered = vocab 32006 with ChatML extension ===
    ckpt_path = ROOT / "ckpts" / "step_259567_v3_surgered.pt"
    print(f"[ARCHON Phase A] loading {ckpt_path}")
    sd = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    if "model" in sd:
        sd = sd["model"]
    # Auto-detect vocab size from ckpt to handle either 32000 or 32006
    embed_key = next((k for k in sd if k.endswith("embed.weight")), None)
    if embed_key is not None:
        ckpt_vocab = sd[embed_key].shape[0]
        if ckpt_vocab != cfg.vocab_size:
            print(f"[ARCHON Phase A] ckpt vocab {ckpt_vocab} != model {cfg.vocab_size}, rebuilding")
            cfg.vocab_size = ckpt_vocab
            model = ArchonBrain(cfg).to(device).to(torch.bfloat16)
    missing, unexpected = model.load_state_dict(sd, strict=False)
    print(f"[ARCHON Phase A] missing={len(missing)} unexpected={len(unexpected)}")
    if missing:
        print(f"  first missing: {missing[:5]}")
    # Re-tie embed <-> lm_head
    model.lm_head.weight = model.embed.weight

    # === Bench baseline forward (no M7) ===
    seq_len = 256
    input_ids = torch.randint(0, cfg.vocab_size, (1, seq_len), device=device)
    model.eval()
    with torch.no_grad():
        # Warmup
        _ = model(input_ids)
        torch.cuda.synchronize()
        t0 = time.time()
        N = 10
        for _ in range(N):
            logits, _, _ = model(input_ids)
        torch.cuda.synchronize()
        t1 = time.time()
    baseline_ms = (t1 - t0) / N * 1000
    print(f"[BASELINE] forward {seq_len}t × {N} runs = {baseline_ms:.2f}ms/run, logits {logits.shape}")

    # === Patch with M7 DoLa: capture L6 + L18 hiddens during forward ===
    captured = {}

    original_forward = model.forward

    def forward_capturing(self, input_ids, targets=None):
        nonlocal captured
        B, T = input_ids.shape
        h = self.embed(input_ids)
        captured["L0"] = h.detach()
        for li, layer in enumerate(self.layers):
            h = layer(h)
            if li + 1 in (6, 18):
                captured[f"L{li+1}"] = h.detach()
        h = self.norm(h)
        logits = self.lm_head(h)
        return logits, None, []

    import types
    model.forward = types.MethodType(forward_capturing, model)

    # === Bench M7 DoLa overhead ===
    cfg_dola = m7_dola.DoLaConfig()
    with torch.no_grad():
        _ = model(input_ids)  # warmup capture
        torch.cuda.synchronize()
        t0 = time.time()
        for _ in range(N):
            logits, _, _ = model(input_ids)
            # Apply DoLa contrast on captured hiddens
            dola = m7_dola.dola_logits(
                captured["L18"], captured["L6"], model.embed.weight, cfg_dola
            )
        torch.cuda.synchronize()
        t1 = time.time()
    dola_ms = (t1 - t0) / N * 1000
    overhead = (dola_ms - baseline_ms) / baseline_ms * 100
    print(f"[M7 DoLa] forward+contrast = {dola_ms:.2f}ms/run, overhead {overhead:+.1f}%")
    print(f"[M7 DoLa] dola logits shape: {dola.shape}")

    # === Verdict ===
    if overhead < 25.0:
        print("[M7 DoLa] WIRE OK — overhead <25%, ready Phase A deploy")
    else:
        print(f"[M7 DoLa] WARNING — overhead {overhead:.1f}% > 25%, optimize before deploy")

    # === Top-5 token diff baseline vs DoLa (sanity check) ===
    with torch.no_grad():
        baseline_logits = original_forward(input_ids)[0][:, -1, :]
        top5_base = baseline_logits.topk(5).indices[0].tolist()
        top5_dola = dola[:, -1, :].topk(5).indices[0].tolist()
    print(f"[sanity] top-5 baseline last token: {top5_base}")
    print(f"[sanity] top-5 DoLa last token:     {top5_dola}")
    overlap = len(set(top5_base) & set(top5_dola))
    print(f"[sanity] overlap baseline∩DoLa = {overlap}/5 (expect 1-4: contrast reorders)")

    # === Save results ===
    results = {
        "device": str(device),
        "vocab_size": cfg.vocab_size,
        "params_M": sum(p.numel() for p in model.parameters()) / 1e6,
        "seq_len": seq_len,
        "baseline_ms_per_forward": baseline_ms,
        "m7_dola_ms_per_forward": dola_ms,
        "m7_overhead_pct": overhead,
        "wire_status": "OK" if overhead < 25.0 else "WARN",
        "top5_baseline": top5_base,
        "top5_dola": top5_dola,
        "overlap": overlap,
    }
    out_path = ROOT / "phase_a_integration_results.json"
    out_path.write_text(json.dumps(results, indent=2))
    print(f"[ARCHON Phase A] results -> {out_path}")


if __name__ == "__main__":
    main()