#!/usr/bin/env python3
"""
eval_dummy.py — Evaluate Tiered vs Baseline HRM using dummy (random) tensors.

No datasets or checkpoints needed. This measures pure model throughput,
latency, and memory usage on synthetic inputs.

Usage:
    source venv/bin/activate
    python eval_dummy.py
    python eval_dummy.py --batch-size 64 --seq-len 256 --hidden-size 1024
    python eval_dummy.py --iterations 50 --plot
"""

import argparse
import json
import os
import sys
import time
from dataclasses import dataclass, asdict

import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from models.memory_tier import MemoryTierManager
from models.hrm.hrm_tiered import HRM_Tiered
from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1


# =====================================================================
#  Dummy batch generator
# =====================================================================

def make_batch(batch_size, seq_len, vocab_size, device):
    return {
        "inputs": torch.randint(0, vocab_size, (batch_size, seq_len), device=device),
        "labels": torch.randint(0, vocab_size, (batch_size, seq_len), device=device),
        "puzzle_identifiers": torch.arange(batch_size, device=device),
    }


# =====================================================================
#  Single-model evaluation
# =====================================================================

@dataclass
class EvalResult:
    model_name: str
    batch_size: int
    seq_len: int
    hidden_size: int
    param_count: int
    latency_mean_ms: float
    latency_std_ms: float
    throughput_sps: float
    gpu_memory_mb: float


def eval_model(model, model_name, batch, device, warmup=5, iterations=20):
    """Run inference on dummy data and collect timing stats."""
    model.eval()
    bs = batch["inputs"].shape[0]
    sl = batch["inputs"].shape[1]
    param_count = sum(p.numel() for p in model.parameters())

    # Warmup
    with torch.no_grad():
        for _ in range(warmup):
            carry = model.initial_carry(batch)
            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
            carry.steps = carry.steps.to(device)
            carry.halted = carry.halted.to(device)
            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
            model(carry, batch)

    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats(device)
        torch.cuda.synchronize()

    # Timed iterations
    latencies = []
    with torch.no_grad():
        for _ in range(iterations):
            carry = model.initial_carry(batch)
            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
            carry.steps = carry.steps.to(device)
            carry.halted = carry.halted.to(device)
            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}

            if torch.cuda.is_available():
                start = torch.cuda.Event(enable_timing=True)
                end = torch.cuda.Event(enable_timing=True)
                start.record()

            model(carry, batch)

            if torch.cuda.is_available():
                end.record()
                torch.cuda.synchronize()
                latencies.append(start.elapsed_time(end))
            else:
                pass  # CPU fallback handled by perf_counter

    mean_ms = sum(latencies) / len(latencies)
    std_ms = (sum((x - mean_ms)**2 for x in latencies) / max(len(latencies)-1, 1)) ** 0.5
    throughput = bs / (mean_ms / 1000) if mean_ms > 0 else 0
    gpu_mem = torch.cuda.max_memory_allocated(device) / (1024**2) if torch.cuda.is_available() else 0

    return EvalResult(
        model_name=model_name,
        batch_size=bs, seq_len=sl,
        hidden_size=model.config.hidden_size,
        param_count=param_count,
        latency_mean_ms=mean_ms,
        latency_std_ms=std_ms,
        throughput_sps=throughput,
        gpu_memory_mb=gpu_mem,
    )


# =====================================================================
#  Main comparison
# =====================================================================

def main():
    parser = argparse.ArgumentParser(description="HRM Dummy Tensor Evaluation")
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--seq-len", type=int, default=81, help="Sudoku=81, or any length")
    parser.add_argument("--hidden-size", type=int, default=512)
    parser.add_argument("--num-heads", type=int, default=8)
    parser.add_argument("--H-cycles", type=int, default=2)
    parser.add_argument("--L-cycles", type=int, default=2)
    parser.add_argument("--H-layers", type=int, default=4)
    parser.add_argument("--L-layers", type=int, default=4)
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--iterations", type=int, default=20)
    parser.add_argument("--output", type=str, default="benchmark_results/eval_dummy.json")
    parser.add_argument("--plot", action="store_true")
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    vocab_size = 32

    config_dict = {
        "batch_size": args.batch_size,
        "seq_len": args.seq_len,
        "puzzle_emb_ndim": 0,
        "num_puzzle_identifiers": args.batch_size,
        "vocab_size": vocab_size,
        "H_cycles": args.H_cycles,
        "L_cycles": args.L_cycles,
        "H_layers": args.H_layers,
        "L_layers": args.L_layers,
        "hidden_size": args.hidden_size,
        "expansion": 4.0,
        "num_heads": args.num_heads,
        "pos_encodings": "rope",
        "halt_max_steps": 1,
        "halt_exploration_prob": 0.0,
    }

    batch = make_batch(args.batch_size, args.seq_len, vocab_size, device)

    print(f"\n{'='*64}")
    print(f"  HRM Dummy Tensor Evaluation")
    print(f"  Device:      {device} ({torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'})")
    print(f"  Batch:       {args.batch_size}")
    print(f"  Seq Len:     {args.seq_len}")
    print(f"  Hidden:      {args.hidden_size}")
    print(f"  H/L cycles:  {args.H_cycles}/{args.L_cycles}")
    print(f"  H/L layers:  {args.H_layers}/{args.L_layers}")
    print(f"  Iterations:  {args.iterations} (warmup: {args.warmup})")
    print(f"{'='*64}")

    results = []

    # ── Baseline ──
    print(f"\n  → Evaluating Baseline (hrm_act_v1)...")
    baseline_model = HierarchicalReasoningModel_ACTV1(config_dict).to(device)
    r_baseline = eval_model(baseline_model, "Baseline", batch, device, args.warmup, args.iterations)
    results.append(r_baseline)
    del baseline_model
    torch.cuda.empty_cache()

    # ── Tiered ──
    print(f"  → Evaluating Tiered (SRAM/DRAM Triton)...")
    mem_mgr = MemoryTierManager(device=device, enable_tracking=True)
    tiered_model = HRM_Tiered(config_dict, memory_manager=mem_mgr).to(device)
    r_tiered = eval_model(tiered_model, "Tiered", batch, device, args.warmup, args.iterations)

    # Grab tier-specific stats
    tier_timing = tiered_model.get_timing_stats()
    tier_mem = mem_mgr.get_stats()
    results.append(r_tiered)
    del tiered_model
    torch.cuda.empty_cache()

    # ── Print comparison table ──
    print(f"\n{'='*64}")
    print(f"  {'Model':<12} {'Params':>10} {'Latency(ms)':>13} {'±σ':>8} {'Throughput':>12} {'GPU MB':>8}")
    print(f"  {'-'*58}")
    for r in results:
        print(f"  {r.model_name:<12} {r.param_count/1e6:>9.1f}M {r.latency_mean_ms:>13.2f} {r.latency_std_ms:>8.2f} {r.throughput_sps:>12.1f} {r.gpu_memory_mb:>8.1f}")

    # Speedup
    speedup = r_baseline.latency_mean_ms / r_tiered.latency_mean_ms if r_tiered.latency_mean_ms > 0 else 0
    mem_diff = r_baseline.gpu_memory_mb - r_tiered.gpu_memory_mb
    print(f"\n  Speedup:        {speedup:.2f}x")
    print(f"  Memory saving:  {mem_diff:.1f} MB")

    # Tier-specific stats
    if tier_timing:
        l_us = tier_timing.get("L_forward_us", {}).get("mean_us", 0)
        h_us = tier_timing.get("H_forward_us", {}).get("mean_us", 0)
        ratio = h_us / l_us if l_us > 0 else 0
        print(f"\n  L-level (SRAM): {l_us:.1f} μs")
        print(f"  H-level (DRAM): {h_us:.1f} μs")
        print(f"  H/L ratio:      {ratio:.2f}x")

    if tier_mem:
        print(f"  SRAM hit rate:  {tier_mem['sram']['hit_rate']:.2%}")
    print(f"{'='*64}\n")

    # ── Save ──
    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
    with open(args.output, "w") as f:
        json.dump([asdict(r) for r in results], f, indent=2)
    print(f"  Results saved → {args.output}")

    # ── Plots ──
    if args.plot:
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt

            fig, axes = plt.subplots(1, 3, figsize=(15, 5))
            fig.suptitle("HRM Dummy Eval: Baseline vs Tiered", fontweight="bold")
            names = [r.model_name for r in results]
            colors = ["#e74c3c", "#2ecc71"]

            # Latency
            axes[0].bar(names, [r.latency_mean_ms for r in results], color=colors)
            axes[0].set_ylabel("Latency (ms)")
            axes[0].set_title("Inference Latency")
            axes[0].grid(axis="y", alpha=0.3)

            # Throughput
            axes[1].bar(names, [r.throughput_sps for r in results], color=colors)
            axes[1].set_ylabel("Samples/sec")
            axes[1].set_title("Throughput")
            axes[1].grid(axis="y", alpha=0.3)

            # Memory
            axes[2].bar(names, [r.gpu_memory_mb for r in results], color=colors)
            axes[2].set_ylabel("GPU Memory (MB)")
            axes[2].set_title("Peak Memory")
            axes[2].grid(axis="y", alpha=0.3)

            plt.tight_layout()
            plot_dir = os.path.dirname(args.output) or "benchmark_results"
            plot_path = os.path.join(plot_dir, "eval_dummy_comparison.png")
            plt.savefig(plot_path, dpi=150)
            plt.close()
            print(f"  Plot saved → {plot_path}")
        except ImportError:
            print("  (matplotlib not found — skipping plots)")

    print("  Done!\n")


if __name__ == "__main__":
    main()