#!/usr/bin/env python3
"""
Evaluate Baseline (Traditional) vs Tiered (Fused) HRM models on the specified dataset.
Both models are evaluated using the EXACT SAME weights to verify correctness and compare speed.
"""

import os
import sys
import yaml
import time
import argparse

import torch
import numpy as np
from safetensors.torch import load_file
from omegaconf import OmegaConf

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

from pretrain import PretrainConfig, create_dataloader
from models.hrm.hrm_tiered import HRM_Tiered
from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1
from models.losses import ACTLossHead
from models.memory_tier import MemoryTierManager

def build_model(arch, config_dict, device):
    if arch == "tiered":
        mm = MemoryTierManager(device=device, enable_tracking=True)
        model = HRM_Tiered(config_dict, memory_manager=mm)
    else:
        model = HierarchicalReasoningModel_ACTV1(config_dict)
    
    # Wrap in ACTLossHead exactly as pretrain.py does
    loss_head = ACTLossHead(model, loss_type="stablemax_cross_entropy").to(device)
    loss_head.eval()
    return loss_head, (mm if arch == "tiered" else None)

@torch.no_grad()
def benchmark_model(model_name, model, dataloader, metadata, device):
    print(f"\n[{model_name}] Starting Evaluation on dataset...")
    
    model.eval()
    all_metrics = []
    
    start_time = time.perf_counter()
    total_samples = 0
    
    # We will accumulate the exact accuracy matching evaluate.py
    total_accuracy = 0
    total_exact_accuracy = 0
    total_count = 0
    
    for set_name, batch, batch_size in dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        # ACTLossHead wraps initial_carry
        with torch.device(device):
            carry = model.initial_carry(batch)
        
        while True:
            carry, loss, metrics, _, all_finish = model(carry=carry, batch=batch, return_keys=[])
            if all_finish:
                break
                
        total_accuracy += metrics["accuracy"].item()
        total_exact_accuracy += metrics["exact_accuracy"].item()
        total_count += metrics["count"].item()
        total_samples += batch_size

    # Synchronize GPU to ensure timing is correct
    if torch.cuda.is_available():
        torch.cuda.synchronize()
        
    end_time = time.perf_counter()
    duration = end_time - start_time
    
    acc = total_accuracy / max(total_count, 1)
    exact_acc = total_exact_accuracy / max(total_count, 1)
    throughput = total_samples / duration
    
    print(f"[{model_name}] Results:")
    print(f"  Duration:   {duration:.2f}s")
    print(f"  Throughput: {throughput:.1f} samples/sec")
    print(f"  Token Acc:  {acc*100:.2f}%")
    print(f"  Exact Acc:  {exact_acc*100:.2f}%")
    
    return {
        "duration_s": duration,
        "throughput": throughput,
        "token_acc": acc,
        "exact_acc": exact_acc
    }

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from  matplotlib.gridspec import GridSpec
import multiprocessing as mp

def run_evaluation(arch, model_cfg, state_dict, data_path, global_batch_size, gpu_id, result_queue):
    device = torch.device(f"cuda:{gpu_id}")
    torch.cuda.set_device(device)
    
    # Needs to recreate dataloader per process
    cfg_container = {
        "arch": model_cfg,
        "data_path": data_path,
        "global_batch_size": global_batch_size,
        "epochs": 1, "lr": 7e-5, "lr_min_ratio": 1.0, "lr_warmup_steps": 2000,
        "weight_decay": 1.0, "beta1": 0.9, "beta2": 0.95,
        "puzzle_emb_lr": 7e-5, "puzzle_emb_weight_decay": 1.0,
        "seed": 0
    }
    config = PretrainConfig(**cfg_container)
    eval_loader, eval_metadata = create_dataloader(
        config, "test", test_set_mode=True, epochs_per_iter=1, 
        global_batch_size=config.global_batch_size, rank=0, world_size=1
    )

    model_cfg = model_cfg.copy()
    model_cfg.update({
        "batch_size": global_batch_size,
        "vocab_size": eval_metadata.vocab_size,
        "seq_len": eval_metadata.seq_len,
        "num_puzzle_identifiers": eval_metadata.num_puzzle_identifiers,
        "causal": False
    })
    
    model_name = "Traditional HRM" if arch == "baseline" else "Fused Tiered HRM"
    model, _ = build_model(arch, model_cfg, device)
    
    try:
        model.load_state_dict(state_dict, strict=True)
    except:
        model.load_state_dict(state_dict, strict=False)
        
    res = benchmark_model(model_name, model, eval_loader, eval_metadata, device)
    res["arch"] = arch
    result_queue.put(res)

def create_comparison_plots(base_res, tier_res, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    
    c_base, c_tier = "#4A90D9", "#E85D75"
    bg, text, grid = "#1a1a2e", "#e0e0e0", "#333355"

    plt.rcParams.update({
        "figure.facecolor": bg, "axes.facecolor": "#16213e",
        "axes.edgecolor": grid, "axes.labelcolor": text,
        "text.color": text, "xtick.color": text, "ytick.color": text,
        "grid.color": grid, "grid.alpha": 0.3,
        "font.family": "sans-serif", "font.size": 11,
    })

    fig = plt.figure(figsize=(15, 6))
    fig.suptitle("Sudoku Extreme: Traditional vs Fused Tiered HRM", fontsize=16, fontweight="bold", y=0.98)
    gs = GridSpec(1, 3, figure=fig, wspace=0.3)
    labels = ["Traditional (v1)", "Fused Tiered"]

    def bar_plot(ax, title, ylabel, vals, fmt=".2f", is_percent=False):
        bars = ax.bar(labels, vals, color=[c_base, c_tier], edgecolor="white", width=0.5)
        ax.set_title(title, fontweight="bold")
        ax.set_ylabel(ylabel)
        for b, v in zip(bars, vals):
            val_str = f"{v*100:{fmt}}%" if is_percent else f"{v:{fmt}}"
            ax.text(b.get_x() + b.get_width()/2, b.get_height() * 1.02,
                    val_str, ha="center", fontsize=11, color=text, fontweight="bold")
        ax.grid(axis="y")
        if is_percent: ax.set_ylim(0, 1.1)

    # 1. Throughput
    bar_plot(fig.add_subplot(gs[0, 0]), "Inference Throughput", "Samples / Second",
             [base_res["throughput"], tier_res["throughput"]], fmt=".1f")

    # 2. Token Accuracy
    bar_plot(fig.add_subplot(gs[0, 1]), "Token Accuracy", "Accuracy",
             [base_res["token_acc"], tier_res["token_acc"]], is_percent=True)

    # 3. Exact Match Accuracy
    bar_plot(fig.add_subplot(gs[0, 2]), "Exact Puzzle Accuracy", "Accuracy",
             [base_res["exact_acc"], tier_res["exact_acc"]], is_percent=True)

    path = os.path.join(output_dir, "eval_fused_vs_v1_comparison.png")
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close()
    print(f"\n  Plot saved → {path}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--weights", type=str, default="hf_upload/tiered_hrm_sram_dram/model.safetensors")
    parser.add_argument("--output-dir", type=str, default="benchmark_results")
    args = parser.parse_args()
    
    # Base Configuration
    model_cfg = {
        "name": "hrm.hrm_tiered@HRM_Tiered",
        "loss": {"name": "losses@ACTLossHead", "loss_type": "stablemax_cross_entropy"},
        "hidden_size": 512, "num_heads": 8, "expansion": 4,
        "H_layers": 4, "L_layers": 4, "H_cycles": 2, "L_cycles": 2,
        "halt_max_steps": 16, "halt_exploration_prob": 0.1,
        "pos_encodings": "rope", "puzzle_emb_ndim": 512,
        "batch_size": 384, "vocab_size": 32, "seq_len": 81,
        "num_puzzle_identifiers": 384, "causal": False
    }
    data_path = "data/sudoku-extreme-1k-aug-1000"
    
    print(f"Loading weights from: {args.weights}")
    try:
        state_dict = load_file(args.weights)
    except Exception as e:
        print(f"Could not load as safetensors, falling back to torch.load... ({e})")
        raw_state_dict = torch.load(args.weights, map_location="cpu", weights_only=True)
        # Strip torch.compile prefix just in case as evaluate.py does
        state_dict = {k.removeprefix("_orig_mod."): v for k, v in raw_state_dict.items()}
        
    if "model.inner.embed_tokens.embedding_weight" not in state_dict and "inner.embed_tokens.embedding_weight" in state_dict:
        state_dict = {f"model.{k}": v for k, v in state_dict.items()}

    mp.set_start_method('spawn', force=True)
    ctx = mp.get_context('spawn')
    queue = ctx.Queue()

    print("\nStarting Parallel Evaluation on 2 GPUs...")
    print("  Traditional HRM -> GPU 0")
    print("  Fused Tiered HRM  -> GPU 1")
    
    # Launch parallel processes
    p1 = ctx.Process(target=run_evaluation, args=("baseline", model_cfg.copy(), state_dict, data_path, 384, 0, queue))
    p2 = ctx.Process(target=run_evaluation, args=("tiered", model_cfg.copy(), state_dict, data_path, 384, 1, queue))
    
    p1.start()
    p2.start()
    
    p1.join()
    p2.join()

    # Collect Results
    results = {}
    while not queue.empty():
        res = queue.get()
        results[res["arch"]] = res

    if "baseline" in results and "tiered" in results:
        base_res = results["baseline"]
        tier_res = results["tiered"]

        print("\n" + "="*50)
        print("      FINAL COMPARISON OVERVIEW")
        print("="*50)
        print(f"Token Accuracy:  Traditional {base_res['token_acc']*100:.2f}%  vs  Fused {tier_res['token_acc']*100:.2f}%")
        print(f"Exact Accuracy:  Traditional {base_res['exact_acc']*100:.2f}%  vs  Fused {tier_res['exact_acc']*100:.2f}%")
        print(f"Throughput:      Traditional {base_res['throughput']:.1f} samp/s  vs  Fused {tier_res['throughput']:.1f} samp/s")
        print(f"Speedup Margin:  {tier_res['throughput'] / base_res['throughput']:.2f}x")
        
        # Plotting
        create_comparison_plots(base_res, tier_res, args.output_dir)
    else:
        print("\nEvaluation failed. One or both models did not return results.")
    
if __name__ == "__main__":
    main()