"""
LUNA 100M — Validate Pretrained + Quantization Benchmark
=========================================================
1. Load pretrained base model (latest.pt — auto-downloads from HF)
2. Run eval prompts with the base (F32) model
3. Simulate quantisation at each level (F16, Q8_0, Q4_K_M) IN PYTORCH
4. Run the SAME eval prompts with each quantised copy
5. Compute precision metrics (cosine-sim of logits, perplexity delta)
6. Export all GGUF files
7. Print comparison report + pick the best quantisation

Usage:
    python validate_and_quantize.py
    python validate_and_quantize.py --ckpt Base/out/pretrain/luna_100m/latest.pt
    python validate_and_quantize.py --skip_gguf   # skip GGUF export
"""

import os, sys, copy, math, json, argparse, struct, time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path

# ─── Model (identical to train.py / sft_train.py) ────────────────────────────

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=1024):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        t = torch.arange(max_seq_len).float()
        freqs = torch.einsum("i,j->ij", t, inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer("cos_cached", emb.cos())
        self.register_buffer("sin_cached", emb.sin())

    def forward(self, seq_len):
        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)

def apply_rotary(x, cos, sin):
    c = cos.unsqueeze(0).unsqueeze(0)
    s = sin.unsqueeze(0).unsqueeze(0)
    return x * c + rotate_half(x) * s

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, rotary_pct=0.25):
        super().__init__()
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.rot_dim = int(self.head_dim * rotary_pct)
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=True)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=True)
        self.rotary = RotaryEmbedding(self.rot_dim, block_size)

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x).reshape(B, T, 3, self.n_head, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        cos, sin = self.rotary(T)
        q = torch.cat([apply_rotary(q[..., :self.rot_dim], cos, sin), q[..., self.rot_dim:]], dim=-1)
        k = torch.cat([apply_rotary(k[..., :self.rot_dim], cos, sin), k[..., self.rot_dim:]], dim=-1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.c_proj(y.transpose(1, 2).contiguous().view(B, T, C))

class MLP(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.fc = nn.Linear(n_embd, 4 * n_embd, bias=True)
        self.gelu = nn.GELU()
        self.proj = nn.Linear(4 * n_embd, n_embd, bias=True)
    def forward(self, x):
        return self.proj(self.gelu(self.fc(x)))

class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = MLP(n_embd)
    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

class LUNAModel(nn.Module):
    def __init__(self, vocab_size, block_size, n_layer, n_embd, n_head):
        super().__init__()
        self.block_size = block_size
        self.wte = nn.Embedding(vocab_size, n_embd)
        self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.lm_head.weight = self.wte.weight
    def forward(self, idx):
        x = self.wte(idx)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        return self.lm_head(x)
    @property
    def num_params(self):
        return sum(p.numel() for p in self.parameters()) - self.wte.weight.numel()


# ─── Quantise-and-dequantise in PyTorch (simulates precision loss) ────────────

BLOCK_SIZE = 32

def _sim_q8_0(tensor: torch.Tensor) -> torch.Tensor:
    """Simulate Q8_0: blockwise int8 quantise → dequantise."""
    orig_shape = tensor.shape
    flat = tensor.flatten().float()
    pad = (-len(flat)) % BLOCK_SIZE
    if pad:
        flat = F.pad(flat, (0, pad))
    blocks = flat.view(-1, BLOCK_SIZE)
    scales = blocks.abs().max(dim=1, keepdim=True).values / 127.0
    scales = scales.clamp(min=1e-8)
    q = (blocks / scales).round().clamp(-128, 127)
    deq = (q * scales).flatten()[:tensor.numel()]
    return deq.view(orig_shape).to(tensor.dtype)

def _sim_q4_k_m(tensor: torch.Tensor) -> torch.Tensor:
    """Simulate Q4_K_M: blockwise 4-bit quantise → dequantise."""
    orig_shape = tensor.shape
    flat = tensor.flatten().float()
    pad = (-len(flat)) % BLOCK_SIZE
    if pad:
        flat = F.pad(flat, (0, pad))
    blocks = flat.view(-1, BLOCK_SIZE)
    abs_max = blocks.abs().max(dim=1, keepdim=True).values
    scales = abs_max / 7.0
    scales = scales.clamp(min=1e-8)
    q = ((blocks / scales) + 8).round().clamp(0, 15)
    deq = ((q - 8) * scales).flatten()[:tensor.numel()]
    return deq.view(orig_shape).to(tensor.dtype)

# Which params get quantised (biases + norms stay F32)
_QUANT_PARAM_SUFFIXES = (".weight",)
_SKIP_QUANT = ("ln1.", "ln2.", "ln_f.")

def apply_simulated_quant(model: LUNAModel, quant: str):
    """Apply simulated quantisation to model weights (in-place). Returns model."""
    if quant == "F32":
        return model
    for name, p in model.named_parameters():
        if not any(name.endswith(s) for s in _QUANT_PARAM_SUFFIXES):
            continue
        if any(skip in name for skip in _SKIP_QUANT):
            continue
        if quant == "F16":
            p.data = p.data.half().float()
        elif quant == "Q8_0":
            p.data = _sim_q8_0(p.data)
        elif quant == "Q4_K_M":
            p.data = _sim_q4_k_m(p.data)
    return model


# ─── Generation ───────────────────────────────────────────────────────────────

@torch.no_grad()
def generate(model, input_ids, max_new_tokens=100, temperature=0.7, top_k=40):
    """Greedy/sampling generation."""
    device = input_ids.device
    for _ in range(max_new_tokens):
        idx_cond = input_ids[:, -model.block_size:]
        logits = model(idx_cond)
        logits = logits[:, -1, :] / max(temperature, 1e-8)
        if top_k > 0:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = float("-inf")
        probs = F.softmax(logits, dim=-1)
        nxt = torch.multinomial(probs, num_samples=1)
        input_ids = torch.cat([input_ids, nxt], dim=1)
        if nxt.item() == 0:  # EOS
            break
    return input_ids

@torch.no_grad()
def get_logits(model, input_ids):
    """Get full logits for a sequence (for precision comparison)."""
    return model(input_ids[:, -model.block_size:])

@torch.no_grad()
def compute_perplexity(model, input_ids):
    """Compute perplexity of the model on a token sequence."""
    if input_ids.size(1) < 2:
        return float("inf")
    logits = model(input_ids[:, -model.block_size:])
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = input_ids[:, 1:].contiguous()
    loss = F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1)
    )
    return math.exp(loss.item())


# ─── Eval prompts ─────────────────────────────────────────────────────────────

EVAL_PROMPTS = [
    # Identity
    "Who are you?",
    "Who created you?",
    "What is your name?",
    # Knowledge
    "The capital of France is",
    "Water boils at a temperature of",
    "The largest planet in our solar system is",
    "Albert Einstein is famous for",
    # English comprehension
    "The quick brown fox jumps over the lazy",
    "In a groundbreaking study, researchers found that",
    "The most important thing about education is",
    "Once upon a time, in a land far away,",
    "The future of artificial intelligence will",
    # Reasoning / grammar
    "If it rains tomorrow, I will",
    "She went to the store because she needed to buy",
    "The difference between a cat and a dog is that",
]

# Reference sentences for perplexity measurement (well-formed English)
PERPLEXITY_TEXTS = [
    "The quick brown fox jumps over the lazy dog and then runs into the forest.",
    "Artificial intelligence has transformed the way we interact with technology in recent years.",
    "Education is the most powerful weapon which you can use to change the world.",
    "The sun rises in the east and sets in the west, a cycle that has continued for billions of years.",
    "Water is composed of two hydrogen atoms and one oxygen atom, making it essential for all life.",
]


# ─── Main ─────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="LUNA 100M — Validate & Quantize Benchmark")
    parser.add_argument("--ckpt", default="Base/out/pretrain/luna_100m/latest.pt",
                        help="Path to latest.pt checkpoint")
    parser.add_argument("--hf_repo", default="ASTERIZER/LUNA-100M",
                        help="HF model repo to download from if ckpt not found")
    parser.add_argument("--tok_dir", default="Base/checkpoints/EleutherAI/pythia-160m",
                        help="Tokenizer directory")
    parser.add_argument("--max_tokens", type=int, default=80,
                        help="Max tokens to generate per prompt")
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--top_k", type=int, default=40)
    parser.add_argument("--skip_gguf", action="store_true",
                        help="Skip GGUF export (just do the PyTorch comparison)")
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\n{'='*70}")
    print(f"  LUNA 100M — Validate & Quantize Benchmark")
    print(f"  Device: {device}")
    print(f"{'='*70}")

    # ── 1. Load tokenizer ─────────────────────────────────────────────────────
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(args.tok_dir)
    print(f"\n  Tokenizer: {args.tok_dir}  (vocab={tok.vocab_size})")

    # ── 2. Load checkpoint ────────────────────────────────────────────────────
    ckpt_path = Path(args.ckpt)
    if not ckpt_path.exists():
        print(f"\n  Checkpoint not found locally: {ckpt_path}")
        print(f"  Downloading from HuggingFace: {args.hf_repo}")
        from huggingface_hub import hf_hub_download
        ckpt_path.parent.mkdir(parents=True, exist_ok=True)
        hf_hub_download(
            repo_id=args.hf_repo,
            filename="latest.pt",
            local_dir=str(ckpt_path.parent),
            token=os.environ.get("HF_TOKEN"),
        )
        print(f"  Downloaded to: {ckpt_path}")

    print(f"\n  Loading checkpoint: {ckpt_path}")
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
    # Handle both formats: {"model": sd, "step": ...} or raw state_dict
    if isinstance(ckpt, dict) and "model" in ckpt:
        state = ckpt["model"]
        step = ckpt.get("step", "?")
        tokens_seen = ckpt.get("tokens_seen", 0)
    else:
        state = ckpt
        step = "final"
        tokens_seen = 0
    print(f"  Pretrained @ step {step}, tokens seen: {tokens_seen:,}")

    # ── 3. Build model ────────────────────────────────────────────────────────
    model = LUNAModel(
        vocab_size=50304, block_size=1024,
        n_layer=10, n_embd=768, n_head=12,
    )
    model.load_state_dict(state, strict=True)
    model = model.to(device).eval()
    print(f"  Parameters: {model.num_params:,}")
    del ckpt, state

    # Save original F32 weights for restoring after each quant
    original_sd = {k: v.clone() for k, v in model.state_dict().items()}

    # ── 4. Run benchmark across all quant levels ──────────────────────────────
    quant_levels = ["F32", "F16", "Q8_0", "Q4_K_M"]
    all_results = {}      # quant -> {prompt: generated_text}
    all_ppls = {}         # quant -> avg perplexity
    logit_cosine = {}     # quant -> avg cosine similarity vs F32
    base_logits = {}      # prompt -> F32 logits (for comparison)

    for qi, quant in enumerate(quant_levels):
        # Restore original weights
        model.load_state_dict(original_sd, strict=True)

        # Apply simulated quantisation
        apply_simulated_quant(model, quant)

        print(f"\n{'='*70}")
        print(f"  [{qi+1}/{len(quant_levels)}]  {quant}")
        print(f"{'='*70}")

        # ── Generate from eval prompts ────────────────────────────────────────
        results = {}
        cosines = []

        for prompt in EVAL_PROMPTS:
            ids = tok.encode(prompt, return_tensors="pt").to(device)
            out_ids = generate(model, ids, max_new_tokens=args.max_tokens,
                               temperature=args.temperature, top_k=args.top_k)
            text = tok.decode(out_ids[0], skip_special_tokens=True)
            results[prompt] = text

            # Compute logit similarity vs F32
            cur_logits = get_logits(model, ids)
            if quant == "F32":
                base_logits[prompt] = cur_logits.cpu()
            else:
                bl = base_logits[prompt].to(device)
                min_len = min(cur_logits.size(1), bl.size(1))
                cos = F.cosine_similarity(
                    cur_logits[:, :min_len, :].flatten().unsqueeze(0),
                    bl[:, :min_len, :].flatten().unsqueeze(0),
                ).item()
                cosines.append(cos)

            print(f"\n  Prompt: \"{prompt}\"")
            print(f"  Output: {text}")

        all_results[quant] = results

        # ── Perplexity on reference English text ──────────────────────────────
        ppls = []
        for ref in PERPLEXITY_TEXTS:
            ref_ids = tok.encode(ref, return_tensors="pt").to(device)
            ppl = compute_perplexity(model, ref_ids)
            ppls.append(ppl)
        avg_ppl = sum(ppls) / len(ppls)
        all_ppls[quant] = avg_ppl
        print(f"\n  Avg Perplexity: {avg_ppl:.2f}")

        if cosines:
            avg_cos = sum(cosines) / len(cosines)
            logit_cosine[quant] = avg_cos
            print(f"  Logit Cosine Sim vs F32: {avg_cos:.6f}")

    # ── 5. Comparison Report ──────────────────────────────────────────────────
    print(f"\n\n{'='*70}")
    print(f"  QUANTISATION COMPARISON REPORT")
    print(f"{'='*70}")
    print(f"\n  {'Quant':<10} {'Avg PPL':>10} {'Cosine vs F32':>15} {'PPL Delta':>12}")
    print(f"  {'-'*50}")

    base_ppl = all_ppls["F32"]
    scores = {}
    for quant in quant_levels:
        ppl = all_ppls[quant]
        cos = logit_cosine.get(quant, 1.0)
        delta = ppl - base_ppl
        scores[quant] = (cos, delta)
        cos_str = f"{cos:.6f}" if quant != "F32" else "1.000000 (ref)"
        delta_str = f"+{delta:.2f}" if delta >= 0 else f"{delta:.2f}"
        if quant == "F32":
            delta_str = "— (ref)"
        print(f"  {quant:<10} {ppl:>10.2f} {cos_str:>15} {delta_str:>12}")

    # Pick best non-F32 quant
    best_quant = None
    best_score = -1
    for q in ["F16", "Q8_0", "Q4_K_M"]:
        cos, delta = scores[q]
        # Score: high cosine + low ppl delta = good
        score = cos - (abs(delta) / max(base_ppl, 1)) * 0.1
        if score > best_score:
            best_score = score
            best_quant = q

    print(f"\n  Best quantisation: {best_quant}")
    print(f"  (highest logit fidelity with minimal perplexity increase)")

    # ── 6. Side-by-side output comparison ─────────────────────────────────────
    print(f"\n\n{'='*70}")
    print(f"  SIDE-BY-SIDE: F32 (base) vs {best_quant}")
    print(f"{'='*70}")
    for prompt in EVAL_PROMPTS:
        f32_out = all_results["F32"][prompt]
        best_out = all_results[best_quant][prompt]
        match = "MATCH" if f32_out.strip() == best_out.strip() else "DIFFER"
        print(f"\n  Prompt: \"{prompt}\"")
        print(f"  F32  : {f32_out}")
        print(f"  {best_quant:<5}: {best_out}")
        print(f"  [{match}]")

    # ── 7. English Understanding Validation ───────────────────────────────────
    print(f"\n\n{'='*70}")
    print(f"  ENGLISH UNDERSTANDING VALIDATION")
    print(f"{'='*70}")

    english_tests = [
        ("Completion", "The capital of the United Kingdom is"),
        ("Grammar", "She has been working at the company for five"),
        ("Reasoning", "If a train travels at 60 miles per hour for 2 hours, it covers"),
        ("Vocab", "The opposite of hot is"),
        ("Context", "Doctors work in hospitals, and teachers work in"),
        ("Fluency", "In the year 2025, technology has advanced to the point where"),
    ]

    for quant_test in ["F32", best_quant]:
        model.load_state_dict(original_sd, strict=True)
        apply_simulated_quant(model, quant_test)
        print(f"\n  --- {quant_test} ---")
        for label, prompt in english_tests:
            ids = tok.encode(prompt, return_tensors="pt").to(device)
            out_ids = generate(model, ids, max_new_tokens=50,
                               temperature=0.3, top_k=10)
            text = tok.decode(out_ids[0], skip_special_tokens=True)
            print(f"  [{label:>10}] {text}")

    # ── 8. Export GGUF files ──────────────────────────────────────────────────
    if not args.skip_gguf:
        print(f"\n\n{'='*70}")
        print(f"  EXPORTING GGUF FILES")
        print(f"{'='*70}")
        gguf_script = Path("quantisations/convert_to_gguf.py")
        if gguf_script.exists():
            import subprocess
            cmd = [
                sys.executable, str(gguf_script),
                "--ckpt", str(args.ckpt),
                "--tok_dir", str(args.tok_dir),
                "--quant", "all",
            ]
            print(f"  Running: {' '.join(cmd)}")
            subprocess.run(cmd, check=True)
        else:
            print(f"  WARNING: {gguf_script} not found — skipping GGUF export")
    else:
        print(f"\n  (GGUF export skipped)")

    # ── 9. Final Summary ──────────────────────────────────────────────────────
    print(f"\n\n{'='*70}")
    print(f"  FINAL SUMMARY")
    print(f"{'='*70}")
    print(f"  Pretrained step: {step}  |  Tokens seen: {tokens_seen:,}")
    print(f"  Base F32 perplexity: {base_ppl:.2f}")
    print(f"  Best quantisation: {best_quant}")
    print(f"    Cosine similarity vs F32: {logit_cosine.get(best_quant, 1.0):.6f}")
    print(f"    Perplexity: {all_ppls[best_quant]:.2f} (Δ {all_ppls[best_quant] - base_ppl:+.2f})")
    print(f"\n  Recommendation:")
    print(f"    Use {best_quant} for deployment — best precision/size tradeoff.")
    if not args.skip_gguf:
        print(f"    GGUF file: quantisations/LUNA-100M-{best_quant}.gguf")
    print(f"\n{'='*70}")
    print(f"  Done!")
    print(f"{'='*70}\n")


if __name__ == "__main__":
    main()