#!/usr/bin/env python3
"""Analyze whether router predictions can identify hard vs easy queries.

This validates the cost optimization use case: if prediction entropy correlates
with difficulty, we can route uncertain queries to premium models and confident
queries to cheap models.
"""

import sys
from pathlib import Path

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

sys.path.insert(0, str(Path(__file__).parent.parent))


def main():
    # Load config
    config_path = Path(__file__).parent.parent / "configs/hybrid_lambda_rank_openai512.yaml"
    with open(config_path) as f:
        config = yaml.safe_load(f)

    # Load data
    data_path = Path(__file__).parent.parent / "data/router_data_sweep.parquet"
    df = pd.read_parquet(data_path, columns=["query", "task_id", "llm", "effect"])

    # Filter to 7 LLMs
    llm_names = df["llm"].unique().tolist()[:7]
    df = df[df["llm"].isin(llm_names)]

    # Build query-level data
    df["group_key"] = df["query"].astype(str) + "|||" + df["task_id"].astype(str)
    group_sizes = df.groupby("group_key").size()
    valid_groups = group_sizes[group_sizes == 7].index
    df = df[df["group_key"].isin(valid_groups)]

    llm_to_idx = {name: i for i, name in enumerate(llm_names)}
    ground_truth_list = []
    task_ids_list = []

    for group_key, group_df in df.groupby("group_key", sort=False):
        gt_row = np.zeros(7)
        for _, row in group_df.iterrows():
            llm_idx = llm_to_idx.get(row["llm"])
            if llm_idx is not None:
                gt_row[llm_idx] = row["effect"]
        ground_truth_list.append(gt_row)
        task_ids_list.append(group_df["task_id"].iloc[0])

    ground_truth = np.array(ground_truth_list)
    task_ids = np.array(task_ids_list)

    num_correct = (ground_truth > 0.5).sum(axis=1)

    print("=" * 70)
    print("ROUTER DIFFICULTY PREDICTION VALIDATION")
    print("=" * 70)
    print("Question: Can router predictions identify hard vs easy queries?")
    print(f"LLMs: {llm_names}")

    # Load trained model and get predictions
    from embergraphrouter.model.multi_task_graph_router import GraphRouterTrainer

    trainer = GraphRouterTrainer(
        router_data_path=str(data_path),
        llm_path=str(Path(__file__).parent.parent / "configs/LLM_Descriptions.json"),
        llm_embedding_path=str(Path(__file__).parent.parent / "configs/llm_description_embedding.pkl"),
        config=config,
        wandb=None,
    )

    # Load best weights
    checkpoint_path = Path(__file__).parent.parent / "model_path/hybrid_lambda_rank_openai512_best.pth"
    if not checkpoint_path.exists():
        # Try alternative path
        checkpoint_path = Path(__file__).parent.parent / "model_path/sweep_concat_gemini768_openai256_best.pth"

    print(f"\nLoading checkpoint: {checkpoint_path}")
    state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    trainer.gnn.model.load_state_dict(state_dict)
    trainer.gnn.model.eval()

    # Get test set predictions
    test_indices = trainer.gnn.test_indices
    print(f"Test set size: {len(test_indices)}")

    # Get predictions
    predictions = []
    gt_test = []

    with torch.no_grad():
        for idx in test_indices:
            # Get logits from model
            pred_logits = trainer.gnn.predict_single(idx)
            gt = trainer.gnn.labels[idx].cpu().numpy()
            predictions.append(pred_logits)
            gt_test.append(gt)

    predictions = np.array(predictions)
    gt_test = np.array(gt_test)

    print(f"Predictions shape: {predictions.shape}")
    print(f"Ground truth shape: {gt_test.shape}")

    # Compute prediction entropy as difficulty proxy
    probs = torch.softmax(torch.tensor(predictions), dim=1).numpy()
    entropy = -np.sum(probs * np.log(probs + 1e-8), axis=1)
    max_entropy = np.log(7)
    norm_entropy = entropy / max_entropy

    # Max probability (confidence)
    max_prob = probs.max(axis=1)

    # Ground truth difficulty
    test_num_correct = (gt_test > 0.5).sum(axis=1)
    test_variance = gt_test.var(axis=1)

    print("\n" + "=" * 70)
    print("PREDICTION ENTROPY vs GROUND TRUTH DIFFICULTY")
    print("=" * 70)

    corr_entropy_correct = np.corrcoef(norm_entropy, test_num_correct)[0, 1]
    corr_conf_correct = np.corrcoef(max_prob, test_num_correct)[0, 1]

    print(f"Correlation(prediction_entropy, num_models_correct): {corr_entropy_correct:.3f}")
    print(f"Correlation(max_confidence, num_models_correct):     {corr_conf_correct:.3f}")
    print()
    print("Expected: Higher entropy -> fewer models correct (harder queries)")
    print("Expected: Higher confidence -> more models correct (easier queries)")

    # Bucket by prediction entropy
    q25, q75 = np.percentile(norm_entropy, [25, 75])

    print("\n--- Queries Bucketed by Prediction Entropy ---")
    buckets = [
        ("Low entropy (confident)", 0, q25),
        ("Medium entropy", q25, q75),
        ("High entropy (uncertain)", q75, 1.01),
    ]

    for label, lo, hi in buckets:
        mask = (norm_entropy >= lo) & (norm_entropy < hi)
        n = mask.sum()
        if n == 0:
            continue
        avg_correct = test_num_correct[mask].mean()
        oracle_acc = gt_test[mask].max(axis=1).mean() * 100
        random_acc = gt_test[mask].mean() * 100
        print(f"  {label}")
        print(f"    n={n:,}, avg_models_correct={avg_correct:.1f}")
        print(f"    Oracle acc: {oracle_acc:.1f}%, Random: {random_acc:.1f}%")

    # Cost optimization simulation
    print("\n" + "=" * 70)
    print("COST OPTIMIZATION USING PREDICTION ENTROPY")
    print("=" * 70)
    print("Strategy: Route uncertain (high entropy) to premium, confident to cheap")

    cheap_name = "GPT-5 Nano"
    premium_name = "Claude Opus 4.5"

    cheap_idx = llm_names.index(cheap_name) if cheap_name in llm_names else 0
    premium_idx = llm_names.index(premium_name) if premium_name in llm_names else 5

    cheap_cost = 0.45  # input + output per 1M tokens
    premium_cost = 30.0

    print(f"\nCheap model: {llm_names[cheap_idx]} @ ${cheap_cost}")
    print(f"Premium model: {llm_names[premium_idx]} @ ${premium_cost}")

    print("\n--- Testing Entropy Thresholds ---")
    print("-" * 70)

    for thresh in [0.3, 0.4, 0.5, 0.6, 0.7]:
        uncertain = norm_entropy >= thresh
        n_uncertain = uncertain.sum()
        n_confident = (~uncertain).sum()

        if n_uncertain < 10 or n_confident < 10:
            continue

        # Accuracy: cheap for confident, premium for uncertain
        acc_confident = gt_test[~uncertain, cheap_idx].mean()
        acc_uncertain = gt_test[uncertain, premium_idx].mean()
        blend_acc = (n_confident * acc_confident + n_uncertain * acc_uncertain) / len(norm_entropy)

        # Cost
        blend_cost = (n_confident * cheap_cost + n_uncertain * premium_cost) / len(norm_entropy)

        # Baseline: always premium
        always_premium_acc = gt_test[:, premium_idx].mean()

        pct_uncertain = 100 * n_uncertain / len(norm_entropy)
        pct_confident = 100 * n_confident / len(norm_entropy)
        cost_savings = 100 * (1 - blend_cost / premium_cost)
        acc_delta = (blend_acc - always_premium_acc) * 100

        print(f"Threshold >= {thresh:.1f}:")
        print(f"  Uncertain: {n_uncertain:,} ({pct_uncertain:.0f}%) -> Premium")
        print(f"  Confident: {n_confident:,} ({pct_confident:.0f}%) -> Cheap")
        print(f"  Blend accuracy: {blend_acc*100:.1f}% (vs always-premium {always_premium_acc*100:.1f}%: {acc_delta:+.1f}%)")
        print(f"  Blend cost: ${blend_cost:.1f} ({cost_savings:.0f}% savings)")

    # Summary
    print("\n" + "=" * 70)
    print("CONCLUSIONS")
    print("=" * 70)

    if corr_entropy_correct < -0.1:
        print("✓ Prediction entropy NEGATIVELY correlates with num_correct")
        print("  -> High entropy queries ARE harder (fewer models succeed)")
        print("  -> Router predictions CAN be used for cost optimization")
    elif corr_entropy_correct > 0.1:
        print("✗ Unexpected: entropy POSITIVELY correlates with num_correct")
        print("  -> Need to investigate model behavior")
    else:
        print("⚠ Weak correlation between entropy and difficulty")
        print("  -> Cost optimization may not work reliably")


if __name__ == "__main__":
    main()
