#!/usr/bin/env python3
"""Compare checkpoints on key sanity metrics.

Checks:
1. GSM8K ranked as easier than BBEH
2. Non-collapsed rankings (distinct LLM scores)
3. Frontier models (opus, codex, pro) ranked above fast models (flash, haiku, nano)

Uses the full training infrastructure for proper model loading.
"""

import sys
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import wandb
import yaml

_SCRIPT_DIR = Path(__file__).resolve().parent
_MODULE_DIR = _SCRIPT_DIR.parent
if str(_MODULE_DIR.parent) not in sys.path:
    sys.path.insert(0, str(_MODULE_DIR.parent))

from embergraphrouter.model.multi_task_graph_router import GraphRouterTrainer

# LLM tiers based on typical performance
FRONTIER_LLMS = ["opus", "codex", "pro", "gpt-4", "o1", "sonnet"]
FAST_LLMS = ["flash", "haiku", "nano", "lite", "mini", "gpt-3.5"]


def load_trainer_and_model(config_path: Path, model_path: Path):
    """Load model using full training infrastructure."""
    with open(config_path) as f:
        config = yaml.safe_load(f)

    # Resolve relative paths
    for key in ("llm_description_path", "llm_embedding_path", "saved_router_data_path",
                "unified_qa_data_path", "model_path"):
        if key in config and not Path(config[key]).is_absolute():
            config[key] = str(_MODULE_DIR / config[key])

    # Override model path to the checkpoint we want to evaluate
    config["model_path"] = str(model_path)

    # Initialize wandb in disabled mode
    wandb.init(mode="disabled")

    # Create trainer (loads data and initializes model)
    trainer = GraphRouterTrainer(
        router_data_path=config["saved_router_data_path"],
        llm_path=config["llm_description_path"],
        llm_embedding_path=config["llm_embedding_path"],
        config=config,
        wandb=wandb,
    )

    # Load checkpoint weights
    checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
    trainer.gnn.model.load_state_dict(checkpoint)
    trainer.gnn.model.eval()

    wandb.finish()
    return trainer, config


def analyze_checkpoint(name: str, config_path: Path, model_path: Path):
    """Analyze a checkpoint on key metrics."""
    print(f"\n{'='*70}")
    print(f"CHECKPOINT: {name}")
    print(f"{'='*70}")

    try:
        trainer, config = load_trainer_and_model(config_path, model_path)
        print(f"Loaded: {model_path.name}")
        print(f"Config: num_gnn_layers={config.get('num_gnn_layers', 'N/A')}, "
              f"use_query_moe={config.get('use_query_moe', False)}, "
              f"task_weights={config.get('task_weights', 'none')}")
    except Exception as e:
        import traceback
        print(f"ERROR loading model: {e}")
        traceback.print_exc()
        return None

    # Get LLM names
    llm_names = trainer.llm_names
    print(f"\nLLMs: {llm_names}")

    # Load full data to get task info
    df = pd.read_csv(config["saved_router_data_path"], low_memory=False)

    # Get test set indices and run inference
    test_indices = trainer.test_query_indices
    print(f"Test set: {len(test_indices)} queries")

    # Group test queries by task type
    task_predictions = {"gsm8k": [], "bbeh": [], "mmlu": [], "other": []}

    model = trainer.gnn.model
    device = next(model.parameters()).device

    with torch.no_grad():
        for idx in test_indices[:500]:  # Limit for speed
            # Get the query's task
            group_rows = trainer.group_indices[idx]
            task_id = trainer.query_task_labels[idx]

            # Determine task category
            task_lower = str(task_id).lower()
            if "gsm8k" in task_lower:
                cat = "gsm8k"
            elif "bbeh" in task_lower:
                cat = "bbeh"
            elif "mmlu" in task_lower:
                cat = "mmlu"
            else:
                cat = "other"

            # Get predictions
            query_emb = torch.tensor(trainer.query_embeddings[idx:idx+1], device=device)
            llm_embs = torch.tensor(trainer.llm_embeddings, device=device)

            # Forward pass - simplified, may need adjustment based on actual model API
            try:
                # Get effect/cost/latency targets for this query
                effects = trainer.data_df.iloc[group_rows]["effect"].values
                costs = trainer.data_df.iloc[group_rows]["cost"].values
                latencies = trainer.data_df.iloc[group_rows]["latency"].values

                # Use form_data to create proper input
                data = trainer.form_data.form_data(
                    query_emb.numpy(),
                    llm_embs.numpy(),
                    np.array([effects]),
                    np.array([costs]),
                    np.array([latencies]),
                    np.array([0]),  # task_id placeholder
                )

                # Run model
                out = model(data)
                if hasattr(out, 'quality_logits'):
                    preds = torch.sigmoid(out.quality_logits).cpu().numpy().flatten()
                else:
                    preds = effects  # Fallback to ground truth

                task_predictions[cat].append(preds)
            except Exception as e:
                continue

    # Aggregate by task
    results = {}
    for cat in ["gsm8k", "bbeh", "mmlu"]:
        if task_predictions[cat]:
            preds = np.array(task_predictions[cat])
            results[cat] = {
                "avg": preds.mean(),
                "per_llm": preds.mean(axis=0),
                "std_across_llm": preds.mean(axis=0).std(),
                "n": len(preds),
            }
            print(f"\n{cat.upper()}: {results[cat]['n']} samples")
            print(f"  Avg confidence: {results[cat]['avg']:.3f}")
            print(f"  Per-LLM avg: {[f'{x:.3f}' for x in results[cat]['per_llm']]}")
            print(f"  LLM std: {results[cat]['std_across_llm']:.3f}")

    if not results:
        print("ERROR: No predictions obtained")
        return None

    # === CHECK 1: GSM8K easier than BBEH ===
    if "gsm8k" in results and "bbeh" in results:
        gsm8k_avg = results["gsm8k"]["avg"]
        bbeh_avg = results["bbeh"]["avg"]
        check1_pass = gsm8k_avg > bbeh_avg
        print(f"\n[CHECK 1] GSM8K easier than BBEH:")
        print(f"  GSM8K avg: {gsm8k_avg:.3f}")
        print(f"  BBEH avg:  {bbeh_avg:.3f}")
        print(f"  {'PASS' if check1_pass else 'FAIL'}: GSM8K {'>' if check1_pass else '<='} BBEH")
    else:
        check1_pass = None
        print("\n[CHECK 1] SKIP: Missing GSM8K or BBEH data")

    # === CHECK 2: Non-collapsed rankings ===
    all_stds = [results[k]["std_across_llm"] for k in results if "std_across_llm" in results[k]]
    avg_llm_std = np.mean(all_stds) if all_stds else 0
    check2_pass = avg_llm_std > 0.03
    print(f"\n[CHECK 2] Non-collapsed LLM rankings:")
    print(f"  Avg LLM std: {avg_llm_std:.3f}")
    print(f"  {'PASS' if check2_pass else 'FAIL'}: {'Distinct' if check2_pass else 'Collapsed'} rankings")

    # === CHECK 3: Frontier > Fast model rankings ===
    frontier_indices = [i for i, llm in enumerate(llm_names)
                       if any(f in llm.lower() for f in FRONTIER_LLMS)]
    fast_indices = [i for i, llm in enumerate(llm_names)
                   if any(f in llm.lower() for f in FAST_LLMS)]

    if frontier_indices and fast_indices and results:
        all_per_llm = np.mean([results[k]["per_llm"] for k in results], axis=0)
        frontier_avg = all_per_llm[frontier_indices].mean()
        fast_avg = all_per_llm[fast_indices].mean()
        check3_pass = frontier_avg > fast_avg

        print(f"\n[CHECK 3] Frontier > Fast models:")
        print(f"  Frontier ({[llm_names[i] for i in frontier_indices]}): {frontier_avg:.3f}")
        print(f"  Fast ({[llm_names[i] for i in fast_indices]}): {fast_avg:.3f}")
        print(f"  {'PASS' if check3_pass else 'FAIL'}: Frontier {'>' if check3_pass else '<='} Fast")
    else:
        check3_pass = None
        print("\n[CHECK 3] SKIP: Can't identify frontier/fast LLMs")

    print(f"\n{'='*70}")
    print(f"SUMMARY: {name}")
    print(f"  Check 1 (Difficulty): {_fmt(check1_pass)}")
    print(f"  Check 2 (Non-collapse): {_fmt(check2_pass)}")
    print(f"  Check 3 (LLM ranking): {_fmt(check3_pass)}")
    print(f"{'='*70}")

    return {
        "name": name,
        "check1_pass": check1_pass,
        "check2_pass": check2_pass,
        "check3_pass": check3_pass,
        "results": results,
    }


def _fmt(val):
    if val is None:
        return "N/A"
    return "PASS" if val else "FAIL"


def main():
    # Checkpoints to compare
    checkpoints = [
        ("GAT 4L + 23x BBEH",
         _MODULE_DIR / "configs" / "optimal_gat_bbeh_23x.yaml",
         _MODULE_DIR / "model_path" / "optimal_gat_bbeh_23x_20251213.pth"),
        ("MLP + MoE4 (no GNN)",
         _MODULE_DIR / "configs" / "mlp_moe4_top1.yaml",
         _MODULE_DIR / "model_path" / "canonical_mlp_moe4_top1_20251212.pth"),
    ]

    results = []
    for name, config_path, model_path in checkpoints:
        if model_path.exists() and config_path.exists():
            result = analyze_checkpoint(name, config_path, model_path)
            if result:
                results.append(result)
        else:
            print(f"\nSKIPPING {name}: missing files")
            print(f"  Config: {config_path} exists={config_path.exists()}")
            print(f"  Model:  {model_path} exists={model_path.exists()}")

    # Final comparison
    if len(results) >= 2:
        print(f"\n{'='*70}")
        print("FINAL COMPARISON")
        print(f"{'='*70}")
        for r in results:
            checks = sum([1 for c in [r["check1_pass"], r["check2_pass"], r["check3_pass"]] if c])
            total = sum([1 for c in [r["check1_pass"], r["check2_pass"], r["check3_pass"]] if c is not None])
            print(f"{r['name']}: {checks}/{total} checks passed")


if __name__ == "__main__":
    main()
