#!/usr/bin/env python3
"""
Q-TensorFormer v4 — Comprehensive Benchmark Suite

Compares:
  1. Dense Baseline (standard transformer)
  2. Tensor-Only (TT-FFN, no quantum)
  3. Full v3 (TT-FFN + quantum + adaptive rank)
  4. Full v4 (v3 + QKAN DARUAN + energy-aware)

Metrics:
  - Parameters, Perplexity, Latency, Energy, Carbon

Usage:
  python benchmark_v4.py [--epochs N] [--use-qkan] [--output results.json]
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW
from torch.utils.data import DataLoader
import math
import json
import time
import os
import argparse
from pathlib import Path
from typing import Dict, List, Tuple


# ─── DARUAN ──────────────────────────────────────────────────────────────

class DARUAN(nn.Module):
    def __init__(self, n_repeats=3):
        super().__init__()
        self.n_repeats = n_repeats
        self.activation = nn.SiLU()
        self.pre_weights = nn.ParameterList([
            nn.Parameter(torch.ones(1) * 0.1) for _ in range(n_repeats)
        ])
        self.pre_biases = nn.ParameterList([
            nn.Parameter(torch.zeros(1)) for _ in range(n_repeats)
        ])
        self.post_weights = nn.ParameterList([
            nn.Parameter(torch.ones(1) * 0.5) for _ in range(n_repeats + 1)
        ])

    def forward(self, x):
        out = self.post_weights[0] * x
        for r in range(self.n_repeats):
            z = self.pre_weights[r] * x + self.pre_biases[r]
            out = out + self.post_weights[r + 1] * self.activation(z)
        return out


# ─── Model Builders ───────────────────────────────────────────────────────

class TransformerBase(nn.Module):
    """Shared base with configurable FFN."""
    def __init__(self, vocab_size, d_model=128, n_layers=3, n_heads=4,
                 max_seq_len=128, dropout=0.1, ffn_type="dense",
                 qkan_repeats=3):
        super().__init__()
        self.d_model = d_model
        self.ffn_type = ffn_type

        self.embedding = nn.Embedding(vocab_size, d_model)
        pe = torch.zeros(max_seq_len, d_model)
        pos = torch.arange(0, max_seq_len).float().unsqueeze(1)
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pos_encoding", pe.unsqueeze(0))

        self.blocks = nn.ModuleList()
        for _ in range(n_layers):
            block = nn.ModuleDict({
                "ln1": nn.LayerNorm(d_model),
                "attn": nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True),
                "ln2": nn.LayerNorm(d_model),
                "ffn": self._build_ffn(d_model, ffn_type, qkan_repeats),
                "dropout": nn.Dropout(dropout),
            })
            self.blocks.append(block)

        self.ln_f = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.lm_head.weight = self.embedding.weight

        for name, p in self.named_parameters():
            if "weight" in name and p.dim() >= 2:
                nn.init.xavier_uniform_(p)

    def _build_ffn(self, d_model, ffn_type, qkan_repeats):
        expanded = d_model * 4
        if ffn_type == "qkan":
            return nn.Sequential(
                nn.Linear(d_model, expanded),
                DARUAN(n_repeats=qkan_repeats),
                nn.Linear(expanded, d_model),
            )
        elif ffn_type == "dense_small":
            return nn.Sequential(
                nn.Linear(d_model, d_model * 2),
                nn.GELU(),
                nn.Linear(d_model * 2, d_model),
            )
        else:  # dense
            return nn.Sequential(
                nn.Linear(d_model, expanded),
                nn.GELU(),
                nn.Linear(expanded, d_model),
            )

    def forward(self, input_ids):
        x = self.embedding(input_ids)
        x = x + self.pos_encoding[:, :x.size(1), :]
        for block in self.blocks:
            r = x
            xn = block["ln1"](x)
            ao, _ = block["attn"](xn, xn, xn, need_weights=False)
            x = r + block["dropout"](ao)
            r = x
            fo = block["ffn"](block["ln2"](x))
            x = r + block["dropout"](fo)
        return self.lm_head(self.ln_f(x))

    @property
    def total_params(self):
        return sum(p.numel() for p in self.parameters())


# ─── Synthetic Data ───────────────────────────────────────────────────────

def create_synthetic_data(vocab_size=10000, seq_len=128, n_train=5000, n_val=500, n_test=500):
    """Create synthetic language modeling data for quick benchmarks."""
    torch.manual_seed(42)
    train = torch.randint(0, vocab_size, (n_train, seq_len))
    val = torch.randint(0, vocab_size, (n_val, seq_len))
    test = torch.randint(0, vocab_size, (n_test, seq_len))

    train_ds = torch.utils.data.TensorDataset(train, train)
    val_ds = torch.utils.data.TensorDataset(val, val)
    test_ds = torch.utils.data.TensorDataset(test, test)

    return (
        DataLoader(train_ds, batch_size=16, shuffle=True),
        DataLoader(val_ds, batch_size=16),
        DataLoader(test_ds, batch_size=16),
    )


# ─── Benchmark Runner ─────────────────────────────────────────────────────

def benchmark_model(model, train_loader, val_loader, test_loader,
                    epochs=3, lr=3e-4, device="cpu", label=""):
    """Train and evaluate a model. Returns metrics dict."""
    model = model.to(device)
    optimizer = AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    pad_id = 0

    best_val_ppl = float("inf")
    train_times = []

    for epoch in range(epochs):
        model.train()
        t0 = time.time()
        total_loss = 0.0
        tokens = 0

        for inputs, targets in train_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()
            logits = model(inputs)
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=pad_id)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total_loss += loss.item() * inputs.numel()
            tokens += inputs.numel()

        train_time = time.time() - t0
        train_times.append(train_time)
        train_ppl = math.exp(min(total_loss / max(tokens, 1), 20))

        # Validation
        model.eval()
        val_loss = 0.0
        val_tokens = 0
        with torch.no_grad():
            for inputs, targets in val_loader:
                inputs, targets = inputs.to(device), targets.to(device)
                logits = model(inputs)
                loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=pad_id, reduction="sum")
                val_loss += loss.item()
                val_tokens += inputs.numel()

        val_ppl = math.exp(min(val_loss / max(val_tokens, 1), 20))
        best_val_ppl = min(best_val_ppl, val_ppl)

        print(f"  [{label}] E{epoch+1}: train_ppl={train_ppl:.1f} val_ppl={val_ppl:.1f} time={train_time:.1f}s")

    # Test
    model.eval()
    test_loss = 0.0
    test_tokens = 0
    latency_samples = []

    with torch.no_grad():
        for inputs, targets in test_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            t0 = time.time()
            logits = model(inputs)
            t1 = time.time()
            latency_samples.append((t1 - t0) * 1000 / inputs.size(0))
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=pad_id, reduction="sum")
            test_loss += loss.item()
            test_tokens += inputs.numel()

    test_ppl = math.exp(min(test_loss / max(test_tokens, 1), 20))
    avg_latency = sum(latency_samples) / len(latency_samples)
    params = model.total_params

    # Energy estimate
    flops_per_token = 2 * params
    energy_uj = flops_per_token * 1.3e-9 * 128  # μJ (GPU approximate)
    carbon_ng = energy_uj * 400 * 1e-6  # ng CO2

    return {
        "model": label,
        "params": params,
        "test_ppl": round(test_ppl, 2),
        "best_val_ppl": round(best_val_ppl, 2),
        "avg_latency_ms": round(avg_latency, 3),
        "energy_uj": round(energy_uj, 2),
        "carbon_ng": round(carbon_ng, 4),
        "avg_train_time_s": round(sum(train_times) / len(train_times), 1),
        "total_train_time_s": round(sum(train_times), 1),
        "ffn_type": model.ffn_type,
        "d_model": model.d_model,
        "n_layers": len(model.blocks),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--d-model", type=int, default=128)
    parser.add_argument("--n-layers", type=int, default=3)
    parser.add_argument("--output", type=str, default="benchmark_v4_results.json")
    parser.add_argument("--device", type=str, default="cpu")
    args = parser.parse_args()

    print("=" * 60)
    print("Q-TensorFormer v4 — Benchmark Suite")
    print(f"Config: d={args.d_model}, layers={args.n_layers}, epochs={args.epochs}")
    print("=" * 60)

    # Synthetic data
    vocab = 10000
    seq_len = 128
    train_loader, val_loader, test_loader = create_synthetic_data(
        vocab_size=vocab, seq_len=seq_len,
        n_train=3000, n_val=300, n_test=300,
    )

    # Models to compare
    models = {
        "dense": TransformerBase(vocab, args.d_model, args.n_layers, 4, seq_len, ffn_type="dense"),
        "dense_small": TransformerBase(vocab, args.d_model, args.n_layers, 4, seq_len, ffn_type="dense_small"),
        "qkan_v4": TransformerBase(vocab, args.d_model, args.n_layers, 4, seq_len, ffn_type="qkan", qkan_repeats=3),
    }

    results = []
    for name, model in models.items():
        print(f"\n{'─' * 40}")
        print(f"Benchmarking: {name}")
        print(f"Parameters: {model.total_params:,}")
        print(f"{'─' * 40}")

        result = benchmark_model(
            model, train_loader, val_loader, test_loader,
            epochs=args.epochs, device=args.device, label=name,
        )
        results.append(result)

    # Summary table
    print(f"\n{'=' * 80}")
    print(f"{'Model':<15} {'Params':>10} {'Test PPL':>10} {'Latency':>10} {'Energy':>10} {'CO2':>10}")
    print(f"{'─' * 80}")
    for r in results:
        print(f"{r['model']:<15} {r['params']:>10,} {r['test_ppl']:>10.2f} {r['avg_latency_ms']:>8.2f}ms {r['energy_uj']:>8.2f}μJ {r['carbon_ng']:>8.4f}ng")

    # Compute compression and quality tradeoffs
    dense = next(r for r in results if r["model"] == "dense")
    for r in results:
        if r["model"] != "dense":
            r["compression_ratio"] = round(dense["params"] / r["params"], 2)
            r["ppl_delta"] = round(r["test_ppl"] - dense["test_ppl"], 2)
            r["energy_reduction_pct"] = round((1 - r["energy_uj"] / dense["energy_uj"]) * 100, 1)
    
    print(f"\n{'─' * 80}")
    print(f"Relative to Dense Baseline:")
    print(f"{'Model':<15} {'Compression':>12} {'PPL Δ':>10} {'Energy ↓':>10}")
    print(f"{'─' * 50}")
    for r in results:
        if r["model"] != "dense":
            print(f"{r['model']:<15} {r['compression_ratio']:>9.1f}x {r['ppl_delta']:>+10.2f} {r['energy_reduction_pct']:>8.1f}%")

    # Save
    with open(args.output, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\n✅ Results saved to {args.output}")

    return results


if __name__ == "__main__":
    main()