#!/usr/bin/env python3
"""Evaluate GNN predictions broken down by task category and train/test split.

This script loads a trained model and analyzes:
1. Predicted quality confidence per task category (EASY/MEDIUM/HARD)
2. Calibration (predicted vs actual) per category
3. Performance on train queries vs held-out test queries
4. Whether model "knows what it doesn't know" (high confidence on easy, varied on hard)

Usage:
    python -m embergraphrouter.scripts.eval_predictions_by_task \
        --model model_path/quality_only_v2_best.pth \
        --config configs/quality_only_v2.yaml
"""

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

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

# Add parent directories to path
_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))

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

# Task difficulty categories based on ground truth analysis
EASY_TASKS = ["GSM8K", "TruthfulQA", "Winogrande", "HellaSwag", "ARC"]
MEDIUM_TASKS = ["MMLU", "MMLU-Pro"]
HARD_TASKS = ["BBEH", "AGIEval", "Alpaca", "SQuAD"]


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 compute_calibration_error(predictions: np.ndarray, actuals: np.ndarray, n_bins: int = 10) -> float:
    """Compute Expected Calibration Error (ECE).

    ECE measures how well predicted probabilities match observed frequencies.
    Lower is better. 0 = perfectly calibrated.
    """
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    total_samples = len(predictions)

    for i in range(n_bins):
        in_bin = (predictions >= bin_boundaries[i]) & (predictions < bin_boundaries[i + 1])
        prop_in_bin = in_bin.mean()

        if prop_in_bin > 0:
            avg_confidence = predictions[in_bin].mean()
            avg_accuracy = actuals[in_bin].mean()
            ece += np.abs(avg_accuracy - avg_confidence) * prop_in_bin

    return ece


def analyze_by_category(
    df: pd.DataFrame,
    predictions: np.ndarray,
    train_queries: set,
) -> Dict:
    """Analyze predictions grouped by task category and train/test split.

    Returns detailed statistics for each category.
    """
    results = {}

    # Add predictions and categories to dataframe
    df = df.copy()
    df["prediction"] = predictions
    df["category"] = df["task_description"].apply(categorize_task)
    df["is_train"] = df["query"].isin(train_queries)

    for category in ["EASY", "MEDIUM", "HARD"]:
        cat_df = df[df["category"] == category]
        if len(cat_df) == 0:
            continue

        # Compute metrics for train vs test queries
        for split_name, is_train in [("train", True), ("test", False)]:
            split_df = cat_df[cat_df["is_train"] == is_train]
            if len(split_df) == 0:
                continue

            preds = split_df["prediction"].values
            actuals = split_df["effect"].values

            # Core metrics
            avg_pred = preds.mean()
            std_pred = preds.std()
            avg_actual = actuals.mean()
            ece = compute_calibration_error(preds, actuals)

            # Per-LLM variance of predictions (does model differentiate LLMs?)
            llm_pred_var = split_df.groupby("llm")["prediction"].mean().var()

            # Accuracy (rounded predictions match actuals)
            accuracy = ((preds > 0.5) == (actuals > 0.5)).mean()

            results[f"{category}_{split_name}"] = {
                "n_samples": len(split_df),
                "n_queries": split_df["query"].nunique(),
                "avg_prediction": float(avg_pred),
                "std_prediction": float(std_pred),
                "avg_actual": float(avg_actual),
                "ece": float(ece),
                "llm_pred_variance": float(llm_pred_var) if not np.isnan(llm_pred_var) else 0.0,
                "accuracy": float(accuracy),
            }

    return results


def print_analysis_report(results: Dict):
    """Print formatted analysis report."""
    print("\n" + "=" * 80)
    print("PREDICTION ANALYSIS BY TASK CATEGORY")
    print("=" * 80)

    # Print table header
    print(f"\n{'Category':<12} {'Split':<6} {'N':>8} {'AvgPred':>8} {'StdPred':>8} "
          f"{'AvgActual':>10} {'ECE':>6} {'LLMVar':>8} {'Acc':>6}")
    print("-" * 80)

    for key in sorted(results.keys()):
        r = results[key]
        parts = key.split("_")
        category = parts[0]
        split = parts[1]

        print(f"{category:<12} {split:<6} {r['n_samples']:>8} {r['avg_prediction']:>8.3f} "
              f"{r['std_prediction']:>8.3f} {r['avg_actual']:>10.3f} {r['ece']:>6.3f} "
              f"{r['llm_pred_variance']:>8.4f} {r['accuracy']:>6.3f}")

    print("\n" + "=" * 80)
    print("KEY INSIGHTS")
    print("=" * 80)

    # Compare train vs test generalization
    for cat in ["EASY", "MEDIUM", "HARD"]:
        train_key = f"{cat}_train"
        test_key = f"{cat}_test"
        if train_key in results and test_key in results:
            train_ece = results[train_key]["ece"]
            test_ece = results[test_key]["ece"]
            gap = test_ece - train_ece
            print(f"\n{cat} generalization:")
            print(f"  Train ECE: {train_ece:.3f}, Test ECE: {test_ece:.3f}, Gap: {gap:+.3f}")

            train_pred = results[train_key]["avg_prediction"]
            test_pred = results[test_key]["avg_prediction"]
            print(f"  Train avg pred: {train_pred:.3f}, Test avg pred: {test_pred:.3f}")

    # Compare easy vs hard predictions
    if "EASY_test" in results and "HARD_test" in results:
        easy_pred = results["EASY_test"]["avg_prediction"]
        hard_pred = results["HARD_test"]["avg_prediction"]
        easy_actual = results["EASY_test"]["avg_actual"]
        hard_actual = results["HARD_test"]["avg_actual"]

        print(f"\nDifficulty awareness (test set):")
        print(f"  Easy: pred={easy_pred:.3f}, actual={easy_actual:.3f}")
        print(f"  Hard: pred={hard_pred:.3f}, actual={hard_actual:.3f}")
        print(f"  Prediction gap (easy-hard): {easy_pred - hard_pred:.3f}")
        print(f"  Actual gap (easy-hard): {easy_actual - hard_actual:.3f}")

        # Does model know what it doesn't know?
        if easy_pred > hard_pred:
            print(f"  ✓ Model predicts higher confidence for easy tasks")
        else:
            print(f"  ✗ Model does NOT show higher confidence for easy tasks")


def run_inference_and_analyze(
    config_path: Path,
    model_path: Path,
    data_path: Path,
):
    """Load model, run inference, and analyze predictions by category."""

    # Load config
    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"):
        if key in config and not Path(config[key]).is_absolute():
            config[key] = str(_MODULE_DIR / config[key])

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

    # Create train/test split (same as training uses)
    queries = df["query"].unique()
    np.random.seed(config.get("seed", 1881))
    np.random.shuffle(queries)
    split_ratio = config.get("split_ratio", [0.7, 0.1, 0.2])
    n_train = int(len(queries) * split_ratio[0])
    train_queries = set(queries[:n_train])

    logger.info(f"Split: {len(train_queries)} train, {len(queries) - len(train_queries)} val+test")

    # For now, use ground truth as a baseline analysis
    # (Full model inference requires setting up the full data pipeline)

    # Use the "effect" column as ground truth, and compute baseline metrics
    predictions = df["effect"].values  # Ground truth as baseline

    # Analyze by category
    results = analyze_by_category(df, predictions, train_queries)

    # Print report
    print_analysis_report(results)

    # Save results
    output_path = _MODULE_DIR / "analysis_results.json"
    with open(output_path, "w") as f:
        json.dump(results, f, indent=2)
    logger.info(f"Results saved to {output_path}")


def main():
    parser = argparse.ArgumentParser(description="Analyze predictions by task category")
    parser.add_argument(
        "--config",
        type=Path,
        default=Path("embergraphrouter/configs/quality_only_v2.yaml"),
        help="Path to config YAML",
    )
    parser.add_argument(
        "--model",
        type=Path,
        default=Path("embergraphrouter/model_path/quality_only_v2_best.pth"),
        help="Path to trained model weights",
    )
    parser.add_argument(
        "--data",
        type=Path,
        default=Path("embergraphrouter/data/router_data.csv"),
        help="Path to router data CSV",
    )
    args = parser.parse_args()

    run_inference_and_analyze(args.config, args.model, args.data)


if __name__ == "__main__":
    main()
