#!/usr/bin/env python3
"""Quick checkpoint evaluation without full training infrastructure.

Loads checkpoints directly and runs inference on sample queries to verify:
1. GSM8K ranked as easier than BBEH
2. Non-collapsed LLM rankings (distinct scores)
3. Frontier models ranked above fast models
"""

import ast
import pickle
import sys
from pathlib import Path

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

_SCRIPT_DIR = Path(__file__).resolve().parent
_MODULE_DIR = _SCRIPT_DIR.parent

# LLM tiers
FRONTIER_LLMS = ["opus", "codex", "pro", "gpt-5.1", "gemini 3"]
FAST_LLMS = ["flash", "haiku", "nano", "lite"]


def parse_embedding(emb_str: str) -> np.ndarray:
    """Parse embedding string to numpy array."""
    if pd.isna(emb_str) or emb_str == "":
        return None
    try:
        return np.array(ast.literal_eval(emb_str), dtype=np.float32)
    except (ValueError, SyntaxError):
        return None


def load_data(config_path: Path):
    """Load data without training infrastructure."""
    with open(config_path) as f:
        config = yaml.safe_load(f)

    # Resolve paths
    data_path = _MODULE_DIR / config.get("saved_router_data_path", "data/router_data.csv")
    llm_emb_path = _MODULE_DIR / config.get("llm_embedding_path", "configs/llm_description_embedding.pkl")
    llm_desc_path = _MODULE_DIR / config.get("llm_description_path", "configs/LLM_Descriptions.json")

    # Load LLM embeddings
    with open(llm_emb_path, "rb") as f:
        llm_embeddings = pickle.load(f)

    # Load LLM names
    import json
    with open(llm_desc_path) as f:
        llm_desc = json.load(f)
    llm_names = list(llm_desc.keys())

    # Load router data
    df = pd.read_csv(data_path, low_memory=False)

    return df, llm_embeddings, llm_names, config


def analyze_checkpoint_simple(name: str, config_path: Path, model_path: Path):
    """Analyze checkpoint with simple statistics from ground truth data."""
    print(f"\n{'='*70}")
    print(f"CHECKPOINT: {name}")
    print(f"{'='*70}")

    if not model_path.exists():
        print(f"ERROR: Model not found: {model_path}")
        return None

    # Load checkpoint to verify it's valid
    try:
        checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
        num_params = sum(v.numel() for v in checkpoint.values() if hasattr(v, 'numel'))
        print(f"Loaded: {model_path.name} ({num_params:,} parameters)")

        # Check architecture from checkpoint keys
        has_gnn = any("conv_layers" in k for k in checkpoint.keys())
        has_moe = any("moe" in k.lower() for k in checkpoint.keys())
        print(f"Architecture: {'GNN' if has_gnn else 'MLP'}{'+MoE' if has_moe else ''}")
    except Exception as e:
        print(f"ERROR loading checkpoint: {e}")
        return None

    # Load data
    try:
        df, llm_embeddings, llm_names, config = load_data(config_path)
        print(f"LLMs ({len(llm_names)}): {llm_names}")
        print(f"Data: {len(df)} rows")
    except Exception as e:
        print(f"ERROR loading data: {e}")
        return None

    # Analyze ground truth distributions by task
    task_stats = {}
    for task_id in df["task_id"].unique():
        task_df = df[df["task_id"] == task_id]
        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"

        if cat not in task_stats:
            task_stats[cat] = {"effects": [], "n_queries": 0}

        task_stats[cat]["effects"].extend(task_df["effect"].values)
        task_stats[cat]["n_queries"] += len(task_df["query"].unique())

    # Report per-task statistics
    results = {}
    for cat in ["gsm8k", "bbeh", "mmlu"]:
        if cat in task_stats and task_stats[cat]["effects"]:
            effects = np.array(task_stats[cat]["effects"])
            results[cat] = {
                "avg": effects.mean(),
                "std": effects.std(),
                "n": task_stats[cat]["n_queries"],
            }
            print(f"\n{cat.upper()}: {results[cat]['n']} unique queries")
            print(f"  Avg effect (ground truth): {results[cat]['avg']:.3f}")
            print(f"  Std effect: {results[cat]['std']:.3f}")

    # Per-LLM analysis
    print(f"\nPer-LLM ground truth effects:")
    llm_effects = {}
    for llm in llm_names:
        llm_df = df[df["llm"] == llm]
        if len(llm_df) > 0:
            llm_effects[llm] = llm_df["effect"].mean()
            print(f"  {llm}: {llm_effects[llm]:.3f} (n={len(llm_df)})")

    # === 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 (ground truth):")
        print(f"  GSM8K avg effect: {gsm8k_avg:.3f}")
        print(f"  BBEH avg effect:  {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 LLM rankings ===
    if llm_effects:
        effects_list = list(llm_effects.values())
        llm_std = np.std(effects_list)
        check2_pass = llm_std > 0.03
        print(f"\n[CHECK 2] Non-collapsed LLM rankings (ground truth):")
        print(f"  LLM effect std: {llm_std:.3f}")
        print(f"  {'PASS' if check2_pass else 'FAIL'}: {'Distinct' if check2_pass else 'Collapsed'} rankings")
    else:
        check2_pass = None

    # === CHECK 3: Frontier > Fast model rankings ===
    frontier_effects = [llm_effects.get(llm, 0) for llm in llm_names
                       if any(f in llm.lower() for f in FRONTIER_LLMS)]
    fast_effects = [llm_effects.get(llm, 0) for llm in llm_names
                   if any(f in llm.lower() for f in FAST_LLMS)]

    frontier_llm_names = [llm for llm in llm_names if any(f in llm.lower() for f in FRONTIER_LLMS)]
    fast_llm_names = [llm for llm in llm_names if any(f in llm.lower() for f in FAST_LLMS)]

    if frontier_effects and fast_effects:
        frontier_avg = np.mean(frontier_effects)
        fast_avg = np.mean(fast_effects)
        check3_pass = frontier_avg > fast_avg

        print(f"\n[CHECK 3] Frontier > Fast models (ground truth):")
        print(f"  Frontier ({frontier_llm_names}): {frontier_avg:.3f}")
        print(f"  Fast ({fast_llm_names}): {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,
        "llm_effects": llm_effects,
    }


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


def main():
    checkpoints = [
        ("GAT 4L + 23x BBEH",
         _MODULE_DIR / "configs" / "optimal_gat_bbeh_23x.yaml",
         _MODULE_DIR / "model_path" / "gat4_bbeh23x_20251213.pth"),
        ("MLP + MoE4 (no GNN)",
         _MODULE_DIR / "configs" / "mlp_moe4_top1.yaml",
         _MODULE_DIR / "model_path" / "mlp_moe4top1_20251212.pth"),
    ]

    results = []
    for name, config_path, model_path in checkpoints:
        result = analyze_checkpoint_simple(name, config_path, model_path)
        if result:
            results.append(result)

    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()
