"""
Experiment Runner: Baselines Comparison.

Compares:
  - Dense Baseline Transformer
  - Static TT-compressed Transformer (fixed rank 4)
  - Quantized Baseline (INT8 / INT4 simulation)
  - Dynamic Depth (Layer skipping)
  - Adaptive Q-TensorFormer
"""

import sys
import os
import time
import math
import json
import argparse
from pathlib import Path

# Add project root
sys.path.insert(0, str(Path(__file__).parent.parent))

import torch
from src.config import ModelConfig
from src.models import QTensorFormer, DenseBaseline
from src.data import load_synthetic_data, load_wikitext2
from src.hardware_cost_model import HardwareCostModel


def evaluate_ppl(model, data_loader, device="cpu", max_batches=20):
    model.eval()
    model.to(device)
    total_loss, total_tokens = 0.0, 0
    with torch.no_grad():
        for i, (inp, tgt) in enumerate(data_loader):
            if i >= max_batches:
                break
            inp, tgt = inp.to(device), tgt.to(device)
            logits = model(inp)
            loss = torch.nn.functional.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                tgt.reshape(-1),
                ignore_index=0,
                reduction="sum",
            )
            total_loss += loss.item()
            total_tokens += inp.numel()
    avg = total_loss / max(1, total_tokens)
    return math.exp(min(avg, 20.0))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--d-model", type=int, default=128)
    parser.add_argument("--n-layers", type=int, default=2)
    parser.add_argument("--vocab-size", type=int, default=1000)
    parser.add_argument("--seq-len", type=int, default=32)
    parser.add_argument("--device", type=str, default="cpu")
    parser.add_argument("--output", type=str, default="outputs/baseline_results.json")
    args = parser.parse_args()

    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
    device = torch.device(args.device)

    print("=" * 65)
    print("EXPERIMENT: Standardized Baseline Comparison")
    print("=" * 65)

    loader = load_synthetic_data(vocab_size=args.vocab_size, seq_len=args.seq_len, n_samples=500, batch_size=16)
    hw = HardwareCostModel()

    cfg = ModelConfig(
        vocab_size=args.vocab_size,
        d_model=args.d_model,
        n_layers=args.n_layers,
        max_seq_len=args.seq_len,
        tt_rank=8,
        use_quantum=True,
    )

    models = {
        "Dense Baseline": DenseBaseline(cfg),
        "Static TT (r=4)": QTensorFormer(cfg, preset="QTF_CLASSICAL_ONLY"),
        "Dynamic Depth": QTensorFormer(cfg, preset="QTF_LATENCY"),
        "Q-TensorFormer (Full)": QTensorFormer(cfg, preset="QTF_FULL"),
        "Q-TensorFormer (Balanced)": QTensorFormer(cfg, preset="QTF_BALANCED"),
    }

    results = []
    for name, model in models.items():
        if name == "Static TT (r=4)":
            for b in model.blocks:
                b.set_rank(4)

        model.to(device)
        total_p = model.total_params
        meas = hw.profile_execution(model, batch_size=1, seq_len=args.seq_len, n_repeats=5, device=device)
        ppl = evaluate_ppl(model, loader, device=device, max_batches=10)

        record = {
            "model": name,
            "params": total_p,
            "latency_p50_ms": meas.latency_ms,
            "ttft_ms": meas.ttft_ms,
            "tpot_ms": meas.tpot_ms,
            "throughput_tokens_sec": meas.throughput_tokens_sec,
            "traffic_bytes_per_token": meas.memory_traffic_bytes_per_token,
            "peak_memory_mb": meas.peak_memory_mb,
            "joules_per_token": meas.joules_per_token,
            "energy_level": meas.energy_level,
            "perplexity": round(ppl, 2),
            "classification": "MEASURED",
            "hardware": meas.device_name,
        }
        results.append(record)

        print(f"{name:<26} | Params: {total_p:>8,} | Lat: {meas.latency_ms:>6.2f}ms | J/tok: {meas.joules_per_token*1e6:>6.2f}uJ | PPL: {ppl:>6.2f}")

    with open(args.output, "w") as f:
        json.dump(results, f, indent=2)

    print(f"\nResults saved to {args.output}")


if __name__ == "__main__":
    main()