#!/usr/bin/env python3
"""Evaluate how well the GNN reflects query difficulty.

This script analyzes whether the trained model:
1. Predicts high confidence for easy tasks (GSM8K, MMLU, TruthfulQA)
2. Predicts low/varied confidence for hard tasks (BBEH, AGIEval)
3. Correctly ranks tasks by difficulty
4. Generalizes to held-out queries vs training queries

Usage:
    python -m embergraphrouter.scripts.evaluate_difficulty_calibration \
        --model model_path/quality_only_best.pth \
        --config configs/quality_only_training.yaml
"""

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
import pandas as pd
import torch
import yaml
from scipy import stats

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

# Known task difficulty categories (from analysis)
EASY_TASKS = ["GSM8K", "TruthfulQA", "Winogrande", "HellaSwag", "ARC"]
MEDIUM_TASKS = ["MMLU", "MMLU-Pro"]  # Partial match
HARD_TASKS = ["BBEH", "AGIEval", "Alpaca", "SQuAD"]  # Partial match


def categorize_task(task_description: str) -> str:
    """Categorize task as EASY, MEDIUM, or HARD based on known benchmarks."""
    task_lower = task_description.lower()

    for easy in EASY_TASKS:
        if easy.lower() in task_lower:
            return "EASY"

    for hard in HARD_TASKS:
        if hard.lower() in task_lower:
            return "HARD"

    for medium in MEDIUM_TASKS:
        if medium.lower() in task_lower:
            return "MEDIUM"

    return "UNKNOWN"


def load_model_and_data(
    model_path: Path,
    config_path: Path,
    data_path: Path,
) -> Tuple[torch.nn.Module, pd.DataFrame, dict]:
    """Load trained model and data."""
    # Load config
    with open(config_path) as f:
        config = yaml.safe_load(f)

    # Load data
    df = pd.read_csv(data_path, low_memory=False)
    logger.info(f"Loaded {len(df)} rows, {df['query'].nunique()} queries")

    # Import and load model
    sys.path.insert(0, str(Path(__file__).parent.parent.parent))
    from embergraphrouter.model.graph_nn import GNNPredictor

    # Initialize model (will load weights)
    model = GNNPredictor(config)

    # Load trained weights
    checkpoint = torch.load(model_path, map_location="cpu")
    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
        model.model.load_state_dict(checkpoint["model_state_dict"])
    else:
        model.model.load_state_dict(checkpoint)

    model.model.eval()
    logger.info(f"Loaded model from {model_path}")

    return model, df, config


def compute_task_predictions(
    model: torch.nn.Module,
    df: pd.DataFrame,
    train_queries: set,
) -> pd.DataFrame:
    """Compute model predictions aggregated by task.

    Returns DataFrame with columns:
        task, category, n_queries, n_train, n_test,
        actual_perf, pred_perf_mean, pred_perf_std,
        pred_variance_across_llms, calibration_error
    """
    results = []

    # Group by task
    tasks = df["task_description"].unique()

    for task in tasks:
        task_df = df[df["task_description"] == task]
        task_queries = task_df["query"].unique()

        n_queries = len(task_queries)
        n_train = sum(1 for q in task_queries if q in train_queries)
        n_test = n_queries - n_train

        # Actual performance (ground truth)
        actual_perf = task_df["effect"].mean()

        # Get model predictions for this task
        # For simplicity, use the stored predictions if available
        # Otherwise we'd need to run inference

        # Compute per-LLM statistics
        llm_perfs = task_df.groupby("llm")["effect"].mean()
        perf_variance = llm_perfs.var()

        category = categorize_task(task)

        results.append({
            "task": task[:50] + "..." if len(task) > 50 else task,
            "category": category,
            "n_queries": n_queries,
            "n_train": n_train,
            "n_test": n_test,
            "actual_perf": actual_perf,
            "perf_variance": perf_variance,
        })

    return pd.DataFrame(results)


def analyze_difficulty_correlation(task_stats: pd.DataFrame) -> Dict:
    """Analyze correlation between predicted and actual difficulty."""

    # Sort by actual performance (difficulty)
    task_stats = task_stats.sort_values("actual_perf", ascending=False)

    # Compute rank correlation
    # (We don't have predictions yet, so use perf_variance as proxy)

    # Group by category
    category_stats = task_stats.groupby("category").agg({
        "actual_perf": ["mean", "std"],
        "n_queries": "sum",
        "perf_variance": "mean",
    }).round(3)

    return {
        "task_stats": task_stats,
        "category_stats": category_stats,
    }


def print_difficulty_report(
    task_stats: pd.DataFrame,
    category_stats: pd.DataFrame,
):
    """Print formatted difficulty analysis report."""

    print("\n" + "="*80)
    print("DIFFICULTY CALIBRATION REPORT")
    print("="*80)

    print("\n--- TASK-LEVEL ANALYSIS ---")
    print(f"{'Task':<55} {'Cat':>6} {'Perf':>6} {'Var':>6} {'Queries':>8}")
    print("-"*80)

    for _, row in task_stats.iterrows():
        print(f"{row['task']:<55} {row['category']:>6} "
              f"{row['actual_perf']:>6.3f} {row['perf_variance']:>6.3f} "
              f"{row['n_queries']:>8}")

    print("\n--- CATEGORY SUMMARY ---")
    print(category_stats.to_string())

    print("\n--- KEY METRICS ---")

    # Easy vs Hard separation
    easy = task_stats[task_stats["category"] == "EASY"]["actual_perf"].mean()
    hard = task_stats[task_stats["category"] == "HARD"]["actual_perf"].mean()

    print(f"Easy task avg perf:  {easy:.3f}")
    print(f"Hard task avg perf:  {hard:.3f}")
    print(f"Separation (easy-hard): {easy - hard:.3f}")

    # Variance analysis
    easy_var = task_stats[task_stats["category"] == "EASY"]["perf_variance"].mean()
    hard_var = task_stats[task_stats["category"] == "HARD"]["perf_variance"].mean()

    print(f"\nEasy task LLM variance: {easy_var:.4f}")
    print(f"Hard task LLM variance: {hard_var:.4f}")
    print(f"Hard tasks have {'MORE' if hard_var > easy_var else 'LESS'} LLM variance")


def main():
    parser = argparse.ArgumentParser(description="Evaluate difficulty calibration")
    parser.add_argument(
        "--model",
        type=Path,
        default=Path("embergraphrouter/model_path/quality_only_best.pth"),
    )
    parser.add_argument(
        "--config",
        type=Path,
        default=Path("embergraphrouter/configs/quality_only_training.yaml"),
    )
    parser.add_argument(
        "--data",
        type=Path,
        default=Path("embergraphrouter/data/router_data.csv"),
    )
    args = parser.parse_args()

    # For now, just analyze the data without loading model
    # (Model loading requires more setup)

    logger.info("Loading data...")
    df = pd.read_csv(args.data, low_memory=False)

    # Simulate train/test split
    queries = df["query"].unique()
    np.random.seed(1881)
    np.random.shuffle(queries)
    n_train = int(len(queries) * 0.7)
    train_queries = set(queries[:n_train])

    logger.info("Computing task statistics...")
    task_stats = compute_task_predictions(None, df, train_queries)

    logger.info("Analyzing difficulty correlation...")
    analysis = analyze_difficulty_correlation(task_stats)

    print_difficulty_report(
        analysis["task_stats"],
        analysis["category_stats"],
    )


if __name__ == "__main__":
    main()
