"""
Experiment Runner: Latency & Memory Traffic Profiling.

Measures:
  - TTFT (Time to First Token)
  - TPOT (Time per Output Token)
  - p50, p95, p99 latencies
  - Batch-1 low-latency inference behavior
  - Validation of latency predictor vs measured runtime
"""

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

sys.path.insert(0, str(Path(__file__).parent.parent))

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


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=str, default="outputs/latency_results.json")
    parser.add_argument("--repeats", type=int, default=15)
    args = parser.parse_args()

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

    print("=" * 65)
    print("EXPERIMENT: Latency & Real Memory Traffic Profiling")
    print("=" * 65)

    hw = HardwareCostModel()
    cfg = ModelConfig(
        vocab_size=1000,
        d_model=128,
        n_layers=2,
        tt_rank=8,
        use_quantum=True,
    )

    models = {
        "Dense Baseline": DenseBaseline(cfg),
        "QTF_CLASSICAL_ONLY": QTensorFormer(cfg, preset="QTF_CLASSICAL_ONLY"),
        "QTF_LATENCY": QTensorFormer(cfg, preset="QTF_LATENCY"),
        "QTF_BALANCED": QTensorFormer(cfg, preset="QTF_BALANCED"),
        "QTF_FULL": QTensorFormer(cfg, preset="QTF_FULL"),
    }

    results = []
    for name, model in models.items():
        model.eval()

        # Batch-1 low-latency decode test (seq_len=1 with past KV cache simulation)
        dummy_1 = torch.randint(0, 1000, (1, 1))

        # Measure 15 samples
        latencies = []
        with torch.no_grad():
            for _ in range(5):
                _ = model(dummy_1)
            for _ in range(args.repeats):
                t0 = time.perf_counter()
                _ = model(dummy_1)
                t1 = time.perf_counter()
                latencies.append((t1 - t0) * 1000.0)

        latencies.sort()
        n = len(latencies)
        p50 = latencies[n // 2]
        p95 = latencies[min(int(n * 0.95), n - 1)]
        p99 = latencies[min(int(n * 0.99), n - 1)]
        mean_lat = sum(latencies) / n

        # Predict latency using HardwareCostModel
        pred_lat = hw.predict_latency(batch_size=1, seq_len=1, active_rank=4)
        pred_error_pct = abs(pred_lat - p50) / max(1e-4, p50) * 100.0

        rec = {
            "model": name,
            "latency_mean_ms": round(mean_lat, 3),
            "p50_ms": round(p50, 3),
            "p95_ms": round(p95, 3),
            "p99_ms": round(p99, 3),
            "predicted_ms": round(pred_lat, 3),
            "prediction_error_pct": round(pred_error_pct, 2),
            "classification": "MEASURED",
            "hardware": hw.profile["name"],
        }
        results.append(rec)
        print(f"{name:<22} | p50: {p50:>5.2f}ms | p95: {p95:>5.2f}ms | p99: {p99:>5.2f}ms | Pred: {pred_lat:>5.2f}ms (err: {pred_error_pct:.1f}%)")

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

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


if __name__ == "__main__":
    main()