#!/usr/bin/env python3
"""Difficulty prediction analysis for tiered routing.

This script analyzes how well the router's mean logit signal predicts query
difficulty, enabling production tiered routing decisions:

1. EASY queries: Use cheap/fast LLM (e.g., GPT-5 Nano, Gemini Flash Lite)
2. HARD queries: Use top-tier LLM (e.g., Claude Opus, GPT-5.1)
3. VERY HARD queries: Use ensemble of top-tier LLMs

The mean logit represents the model's belief about overall success probability:
- High mean logit → trivial/easy query → route to cheap LLM
- Low mean logit → very hard/impossible query → use ensemble or best-on-hard

Usage:
    python -m embergraphrouter.model.baselines.difficulty_prediction_analysis \
        --checkpoint_dir results/baselines_comprehensive \
        --data_path embergraphrouter/data/router_data_8llm_qwen3.parquet
"""

from __future__ import annotations

import argparse
import ast
import json
import logging
from collections import defaultdict
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from scipy import stats
from sklearn.metrics import roc_auc_score, roc_curve
from torch.utils.data import DataLoader, Dataset

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


# ==============================================================================
# Difficulty Categories for Tiered Routing
# ==============================================================================

ROUTING_TIERS = {
    "easy": {
        "n_correct_range": (7, 8),  # 7-8 models correct
        "description": "Most/all LLMs succeed - use cheapest",
        "recommended_action": "route_to_cheapest",
    },
    "medium": {
        "n_correct_range": (4, 6),  # 4-6 models correct
        "description": "Mixed success - use router's top pick",
        "recommended_action": "use_router_prediction",
    },
    "hard": {
        "n_correct_range": (1, 3),  # 1-3 models correct
        "description": "Few succeed - use top-tier LLM",
        "recommended_action": "route_to_top_tier",
    },
    "very_hard": {
        "n_correct_range": (0, 0),  # 0 models correct
        "description": "No model succeeds - use ensemble or escalate",
        "recommended_action": "use_ensemble_or_escalate",
    },
}


@dataclass
class DifficultyPredictionResult:
    """Results from difficulty prediction analysis."""

    # Mean logit statistics by tier
    mean_logit_by_tier: dict[str, float]
    mean_logit_std_by_tier: dict[str, float]

    # Separation quality
    easy_vs_hard_auc: float
    easy_vs_very_hard_auc: float
    hard_vs_very_hard_auc: float

    # Correlation with ground truth
    mean_logit_n_correct_corr: float
    mean_logit_pct_correct_corr: float

    # Recommended thresholds
    easy_threshold: float  # mean_logit > this → easy
    hard_threshold: float  # mean_logit < this → hard/very_hard
    optimal_separation_accuracy: float

    # LLM performance on hard queries
    best_llms_on_hard: list[tuple[str, float]]  # (llm_name, success_rate)
    best_llms_on_very_hard: list[tuple[str, float]]

    # Cost savings potential
    easy_query_pct: float
    potential_cost_savings: float  # If routing easy to cheapest

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass
class TieredRoutingRecommendation:
    """Production routing recommendation based on mean logit."""

    mean_logit: float
    predicted_tier: str
    routing_action: str
    confidence: float
    recommended_llm: str


# ==============================================================================
# Dataset (reuse from comprehensive_eval)
# ==============================================================================


def parse_embedding(emb_val):
    if isinstance(emb_val, str):
        return np.array(ast.literal_eval(emb_val), dtype=np.float32)
    elif isinstance(emb_val, (list, np.ndarray)):
        return np.array(emb_val, dtype=np.float32)
    else:
        raise ValueError(f"Unknown embedding type: {type(emb_val)}")


class RouterDataset(Dataset):
    def __init__(
        self,
        parquet_path: str,
        embedding_dim: int = 512,
        embedding_prefix: str = "query_embedding_qwen3_instruct_router_v1",
        use_task_embedding: bool = True,
    ):
        df = pd.read_parquet(parquet_path)
        self.llm_names = sorted(df["llm"].unique().tolist())
        self.llm_to_idx = {name: i for i, name in enumerate(self.llm_names)}
        self.num_models = len(self.llm_names)

        query_emb_col = f"{embedding_prefix}_{embedding_dim}"
        task_emb_col = f"task_description_embedding_qwen3_instruct_router_v1_{embedding_dim}"

        grouped = df.groupby(["task_id", "query"])

        query_embs = []
        task_embs = []
        targets = []
        n_correct_list = []
        task_ids = []

        for (task_id, query), group in grouped:
            if len(group) != self.num_models:
                continue

            emb = parse_embedding(group[query_emb_col].iloc[0])
            query_embs.append(emb)

            if use_task_embedding and task_emb_col in df.columns:
                task_emb = parse_embedding(group[task_emb_col].iloc[0])
                task_embs.append(task_emb)

            target = np.zeros(self.num_models)
            for _, row in group.iterrows():
                llm_idx = self.llm_to_idx[row["llm"]]
                target[llm_idx] = row["effect"]
            targets.append(target)

            n_correct = (target > 0.5).sum()
            n_correct_list.append(n_correct)
            task_ids.append(task_id)

        self.query_embs = torch.tensor(np.array(query_embs), dtype=torch.float32)
        self.targets = torch.tensor(np.array(targets), dtype=torch.float32)
        self.n_correct = np.array(n_correct_list)
        self.task_ids = task_ids

        if task_embs:
            task_embs_tensor = torch.tensor(np.array(task_embs), dtype=torch.float32)
            self.query_embs = torch.cat([self.query_embs, task_embs_tensor], dim=-1)
            self.embedding_dim = self.query_embs.shape[-1]
        else:
            self.embedding_dim = embedding_dim

        torch.manual_seed(42)
        self.model_embs = F.normalize(
            torch.randn(self.num_models, self.embedding_dim), dim=-1
        )

        # Assign routing tiers
        self.tiers = self._assign_tiers()

    def _assign_tiers(self) -> list[str]:
        tiers = []
        for n in self.n_correct:
            if n >= 7:
                tiers.append("easy")
            elif n >= 4:
                tiers.append("medium")
            elif n >= 1:
                tiers.append("hard")
            else:
                tiers.append("very_hard")
        return tiers

    def __len__(self):
        return len(self.query_embs)

    def __getitem__(self, idx):
        return self.query_embs[idx], self.targets[idx]


# ==============================================================================
# Router Loading (simplified versions)
# ==============================================================================


def load_two_tower_router(checkpoint_path: str, device: str = "cpu"):
    """Load Two Tower router from checkpoint."""
    from .comprehensive_eval import TwoTowerRouter

    ckpt = torch.load(checkpoint_path, map_location=device)
    config = ckpt["config"]

    router = TwoTowerRouter(
        query_dim=config.get("query_dim", 1024),
        model_dim=config.get("model_dim", 1024),
        hidden_dim=384,
        num_layers=2,
        dropout=0.1,
    )
    router.load_state_dict(ckpt["state_dict"])
    router.to(device)
    router.eval()
    return router


def load_mlp_router(checkpoint_path: str, num_models: int, device: str = "cpu"):
    """Load Multi-Head MLP router from checkpoint."""
    from .comprehensive_eval import MultiHeadMLPRouter

    ckpt = torch.load(checkpoint_path, map_location=device)
    config = ckpt["config"]

    router = MultiHeadMLPRouter(
        query_dim=config.get("query_dim", 1024),
        num_models=num_models,
        hidden_dim=384,
        num_layers=2,
        dropout=0.1,
    )
    router.load_state_dict(ckpt["state_dict"])
    router.to(device)
    router.eval()
    return router


# ==============================================================================
# Difficulty Prediction Analysis
# ==============================================================================


def compute_mean_logit(logits: torch.Tensor) -> torch.Tensor:
    """Compute mean logit per query (difficulty signal).

    The mean logit represents the model's belief about average LLM success.
    Higher mean logit → easier query (more LLMs likely to succeed)
    Lower mean logit → harder query (fewer LLMs likely to succeed)
    """
    return logits.mean(dim=-1)


def analyze_difficulty_prediction(
    logits: torch.Tensor,
    targets: torch.Tensor,
    n_correct: np.ndarray,
    tiers: list[str],
    llm_names: list[str],
) -> DifficultyPredictionResult:
    """Analyze how well mean logit predicts difficulty."""

    mean_logits = compute_mean_logit(logits).numpy()

    # Statistics by tier
    mean_logit_by_tier = {}
    mean_logit_std_by_tier = {}

    for tier in ["easy", "medium", "hard", "very_hard"]:
        mask = np.array([t == tier for t in tiers])
        if mask.sum() > 0:
            mean_logit_by_tier[tier] = float(mean_logits[mask].mean())
            mean_logit_std_by_tier[tier] = float(mean_logits[mask].std())
        else:
            mean_logit_by_tier[tier] = 0.0
            mean_logit_std_by_tier[tier] = 0.0

    # Binary separation AUCs
    easy_mask = np.array([t == "easy" for t in tiers])
    hard_mask = np.array([t == "hard" for t in tiers])
    very_hard_mask = np.array([t == "very_hard" for t in tiers])

    # Easy vs Hard
    if easy_mask.sum() > 0 and hard_mask.sum() > 0:
        combined_mask = easy_mask | hard_mask
        labels = easy_mask[combined_mask].astype(int)
        scores = mean_logits[combined_mask]
        easy_vs_hard_auc = roc_auc_score(labels, scores)
    else:
        easy_vs_hard_auc = 0.5

    # Easy vs Very Hard
    if easy_mask.sum() > 0 and very_hard_mask.sum() > 0:
        combined_mask = easy_mask | very_hard_mask
        labels = easy_mask[combined_mask].astype(int)
        scores = mean_logits[combined_mask]
        easy_vs_very_hard_auc = roc_auc_score(labels, scores)
    else:
        easy_vs_very_hard_auc = 0.5

    # Hard vs Very Hard
    if hard_mask.sum() > 0 and very_hard_mask.sum() > 0:
        combined_mask = hard_mask | very_hard_mask
        labels = hard_mask[combined_mask].astype(int)  # 1 for hard, 0 for very_hard
        scores = mean_logits[combined_mask]
        hard_vs_very_hard_auc = roc_auc_score(labels, scores)
    else:
        hard_vs_very_hard_auc = 0.5

    # Correlation with ground truth
    pct_correct = n_correct / len(llm_names)
    mean_logit_n_correct_corr = float(np.corrcoef(mean_logits, n_correct)[0, 1])
    mean_logit_pct_correct_corr = float(np.corrcoef(mean_logits, pct_correct)[0, 1])

    # Find optimal thresholds using ROC
    # Easy threshold: maximize separation between easy and non-easy
    non_easy_mask = ~easy_mask
    if easy_mask.sum() > 0 and non_easy_mask.sum() > 0:
        labels = easy_mask.astype(int)
        fpr, tpr, thresholds = roc_curve(labels, mean_logits)
        optimal_idx = np.argmax(tpr - fpr)
        easy_threshold = float(thresholds[optimal_idx])
    else:
        easy_threshold = float(np.percentile(mean_logits, 75))

    # Hard threshold: maximize separation between hard/very_hard and rest
    hard_or_very_hard_mask = hard_mask | very_hard_mask
    if hard_or_very_hard_mask.sum() > 0 and (~hard_or_very_hard_mask).sum() > 0:
        labels = hard_or_very_hard_mask.astype(int)
        fpr, tpr, thresholds = roc_curve(labels, -mean_logits)  # Negate for hard
        optimal_idx = np.argmax(tpr - fpr)
        hard_threshold = float(-thresholds[optimal_idx])
    else:
        hard_threshold = float(np.percentile(mean_logits, 25))

    # Compute optimal separation accuracy
    predicted_easy = mean_logits > easy_threshold
    predicted_hard = mean_logits < hard_threshold
    predicted_medium = ~predicted_easy & ~predicted_hard

    correct_predictions = 0
    total = len(tiers)
    for i, tier in enumerate(tiers):
        if tier == "easy" and predicted_easy[i]:
            correct_predictions += 1
        elif tier in ["hard", "very_hard"] and predicted_hard[i]:
            correct_predictions += 1
        elif tier == "medium" and predicted_medium[i]:
            correct_predictions += 1
    optimal_separation_accuracy = correct_predictions / total

    # LLM performance on hard queries
    targets_np = targets.numpy()

    hard_indices = np.where(hard_mask)[0]
    llm_success_on_hard = []
    for i, llm_name in enumerate(llm_names):
        if len(hard_indices) > 0:
            success_rate = (targets_np[hard_indices, i] > 0.5).mean()
            llm_success_on_hard.append((llm_name, float(success_rate)))
    llm_success_on_hard.sort(key=lambda x: x[1], reverse=True)
    best_llms_on_hard = llm_success_on_hard[:3]

    very_hard_indices = np.where(very_hard_mask)[0]
    llm_success_on_very_hard = []
    for i, llm_name in enumerate(llm_names):
        if len(very_hard_indices) > 0:
            success_rate = (targets_np[very_hard_indices, i] > 0.5).mean()
            llm_success_on_very_hard.append((llm_name, float(success_rate)))
    llm_success_on_very_hard.sort(key=lambda x: x[1], reverse=True)
    best_llms_on_very_hard = llm_success_on_very_hard[:3]

    # Cost savings potential
    easy_query_pct = easy_mask.mean()
    # Assume routing easy to cheapest saves 80% cost vs average
    potential_cost_savings = easy_query_pct * 0.8

    return DifficultyPredictionResult(
        mean_logit_by_tier=mean_logit_by_tier,
        mean_logit_std_by_tier=mean_logit_std_by_tier,
        easy_vs_hard_auc=easy_vs_hard_auc,
        easy_vs_very_hard_auc=easy_vs_very_hard_auc,
        hard_vs_very_hard_auc=hard_vs_very_hard_auc,
        mean_logit_n_correct_corr=mean_logit_n_correct_corr,
        mean_logit_pct_correct_corr=mean_logit_pct_correct_corr,
        easy_threshold=easy_threshold,
        hard_threshold=hard_threshold,
        optimal_separation_accuracy=optimal_separation_accuracy,
        best_llms_on_hard=best_llms_on_hard,
        best_llms_on_very_hard=best_llms_on_very_hard,
        easy_query_pct=easy_query_pct,
        potential_cost_savings=potential_cost_savings,
    )


def print_difficulty_analysis(result: DifficultyPredictionResult, model_name: str):
    """Print difficulty prediction analysis."""
    logger.info(f"\n{'='*70}")
    logger.info(f"DIFFICULTY PREDICTION ANALYSIS: {model_name}")
    logger.info(f"{'='*70}")

    logger.info("\n--- Mean Logit by Routing Tier ---")
    for tier in ["easy", "medium", "hard", "very_hard"]:
        mean = result.mean_logit_by_tier.get(tier, 0)
        std = result.mean_logit_std_by_tier.get(tier, 0)
        logger.info(f"  {tier:12s}: mean_logit = {mean:.4f} ± {std:.4f}")

    logger.info("\n--- Separation Quality (AUC) ---")
    logger.info(f"  Easy vs Hard:      {result.easy_vs_hard_auc:.4f}")
    logger.info(f"  Easy vs Very Hard: {result.easy_vs_very_hard_auc:.4f}")
    logger.info(f"  Hard vs Very Hard: {result.hard_vs_very_hard_auc:.4f}")

    logger.info("\n--- Correlation with Ground Truth ---")
    logger.info(f"  Mean Logit ↔ N Correct:   {result.mean_logit_n_correct_corr:.4f}")
    logger.info(f"  Mean Logit ↔ % Correct:   {result.mean_logit_pct_correct_corr:.4f}")

    logger.info("\n--- Recommended Thresholds ---")
    logger.info(f"  Easy threshold:  mean_logit > {result.easy_threshold:.4f}")
    logger.info(f"  Hard threshold:  mean_logit < {result.hard_threshold:.4f}")
    logger.info(f"  Separation accuracy: {result.optimal_separation_accuracy:.2%}")

    logger.info("\n--- Best LLMs on Hard Queries ---")
    for llm, rate in result.best_llms_on_hard:
        logger.info(f"  {llm}: {rate:.2%} success")

    logger.info("\n--- Best LLMs on Very Hard Queries ---")
    for llm, rate in result.best_llms_on_very_hard:
        logger.info(f"  {llm}: {rate:.2%} success")

    logger.info("\n--- Cost Optimization ---")
    logger.info(f"  Easy queries: {result.easy_query_pct:.1%} of total")
    logger.info(f"  Potential cost savings: {result.potential_cost_savings:.1%}")


def get_routing_recommendation(
    mean_logit: float,
    result: DifficultyPredictionResult,
    cheap_llm: str = "GPT-5 Nano",
) -> TieredRoutingRecommendation:
    """Get production routing recommendation for a query."""

    if mean_logit > result.easy_threshold:
        return TieredRoutingRecommendation(
            mean_logit=mean_logit,
            predicted_tier="easy",
            routing_action="route_to_cheapest",
            confidence=min(1.0, (mean_logit - result.easy_threshold) / 0.5),
            recommended_llm=cheap_llm,
        )
    elif mean_logit < result.hard_threshold:
        # Use best LLM on hard queries
        best_llm = result.best_llms_on_hard[0][0] if result.best_llms_on_hard else "Claude Opus 4.5"
        return TieredRoutingRecommendation(
            mean_logit=mean_logit,
            predicted_tier="hard",
            routing_action="route_to_top_tier",
            confidence=min(1.0, (result.hard_threshold - mean_logit) / 0.5),
            recommended_llm=best_llm,
        )
    else:
        return TieredRoutingRecommendation(
            mean_logit=mean_logit,
            predicted_tier="medium",
            routing_action="use_router_prediction",
            confidence=0.5,
            recommended_llm="<use_router_top1>",
        )


# ==============================================================================
# Main Analysis Runner
# ==============================================================================


def run_analysis(
    checkpoint_dir: str,
    data_path: str,
    embedding_dim: int = 512,
    device: str = "cuda",
) -> dict:
    """Run difficulty prediction analysis on trained models."""

    from .comprehensive_eval import (
        MultiHeadMLPRouter,
        SimplifiedGNNRouter,
        TwoTowerRouter,
    )

    checkpoint_dir = Path(checkpoint_dir)

    # Load dataset
    logger.info(f"Loading dataset with embedding_dim={embedding_dim}...")
    dataset = RouterDataset(data_path, embedding_dim=embedding_dim)
    logger.info(f"Loaded {len(dataset)} queries")

    # Split (use same seed for consistency)
    n = len(dataset)
    torch.manual_seed(42)
    indices = torch.randperm(n).tolist()
    test_indices = indices[int(0.85 * n):]

    test_embs = dataset.query_embs[test_indices].to(device)
    test_targets = dataset.targets[test_indices]
    test_n_correct = dataset.n_correct[test_indices]
    test_tiers = [dataset.tiers[i] for i in test_indices]
    model_embs = dataset.model_embs.to(device)

    results = {}

    # Analyze each model
    for model_type in ["two_tower", "multi-head_mlp", "gnn_(simplified)"]:
        checkpoint_path = checkpoint_dir / f"{model_type}_dim{embedding_dim}.pt"

        if not checkpoint_path.exists():
            logger.warning(f"Checkpoint not found: {checkpoint_path}")
            continue

        logger.info(f"\nAnalyzing {model_type}...")

        # Load model
        ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)

        if model_type == "two_tower":
            router = TwoTowerRouter(
                query_dim=dataset.embedding_dim,
                model_dim=dataset.embedding_dim,
                hidden_dim=384,
                num_layers=2,
                dropout=0.1,
            )
        elif model_type == "multi-head_mlp":
            router = MultiHeadMLPRouter(
                query_dim=dataset.embedding_dim,
                num_models=dataset.num_models,
                hidden_dim=384,
                num_layers=2,
                dropout=0.1,
            )
        else:  # GNN
            router = SimplifiedGNNRouter(
                query_dim=dataset.embedding_dim,
                model_dim=dataset.embedding_dim,
                hidden_dim=384,
                num_layers=2,
                num_heads=4,
                dropout=0.1,
            )

        router.load_state_dict(ckpt["state_dict"])
        router.to(device)
        router.eval()

        # Get predictions
        with torch.no_grad():
            logits = router(test_embs, model_embs).cpu()

        # Analyze difficulty prediction
        result = analyze_difficulty_prediction(
            logits=logits,
            targets=test_targets,
            n_correct=test_n_correct,
            tiers=test_tiers,
            llm_names=dataset.llm_names,
        )

        print_difficulty_analysis(result, model_type)
        results[model_type] = result.to_dict()

    return results


def main():
    parser = argparse.ArgumentParser(description="Difficulty prediction analysis")
    parser.add_argument(
        "--checkpoint_dir",
        type=str,
        default="results/baselines_comprehensive",
        help="Directory with model checkpoints",
    )
    parser.add_argument(
        "--data_path",
        type=str,
        default="embergraphrouter/data/router_data_8llm_qwen3.parquet",
        help="Path to router data",
    )
    parser.add_argument(
        "--embedding_dim",
        type=int,
        default=512,
        help="Embedding dimension to analyze",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output JSON path",
    )

    args = parser.parse_args()

    results = run_analysis(
        checkpoint_dir=args.checkpoint_dir,
        data_path=args.data_path,
        embedding_dim=args.embedding_dim,
        device=args.device,
    )

    if args.output:
        with open(args.output, "w") as f:
            json.dump(results, f, indent=2)
        logger.info(f"Saved results to {args.output}")

    # Print summary
    logger.info("\n" + "=" * 70)
    logger.info("SUMMARY: BEST MODEL FOR DIFFICULTY PREDICTION")
    logger.info("=" * 70)

    best_model = None
    best_auc = 0
    for model_name, result in results.items():
        auc = result["easy_vs_very_hard_auc"]
        if auc > best_auc:
            best_auc = auc
            best_model = model_name

    if best_model:
        logger.info(f"Best model: {best_model}")
        logger.info(f"Easy vs Very Hard AUC: {best_auc:.4f}")
        logger.info(f"Mean Logit ↔ N Correct Correlation: {results[best_model]['mean_logit_n_correct_corr']:.4f}")


if __name__ == "__main__":
    main()
