#!/usr/bin/env python3
"""Evaluate intelligent routing with ACTUAL model predictions.

This script loads a trained checkpoint and runs inference-only evaluation
on the validation set using the actual GNN model architecture.
"""

import json
import logging
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
import yaml

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

sys.path.insert(0, str(_MODULE_DIR.parent))

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class MockWandb:
    """Minimal mock for wandb when not training."""

    def log(self, *args, **kwargs) -> None:
        pass

    def config(self) -> dict:
        return {}


def compute_entropy(probs: np.ndarray) -> np.ndarray:
    """Compute normalized entropy."""
    n_llms = probs.shape[1]
    entropy = -np.sum(probs * np.log(probs + 1e-8), axis=1)
    return entropy / np.log(n_llms)


def evaluate_strategy(name, choices, ground_truth, premium_idx, routable_mask=None):
    """Evaluate a routing strategy.

    Args:
        name: Strategy name
        choices: Array of LLM indices chosen per query
        ground_truth: [n_queries, n_llms] reward matrix
        premium_idx: Index of premium LLM
        routable_mask: Boolean mask for queries where routing matters (LLMs differ).
            Used only for reporting accuracy on routable subset; frac_oracle includes
            all queries with correct tie handling.
    """
    n = len(choices)
    oracle_choices = ground_truth.argmax(axis=1)
    oracle_rewards = ground_truth.max(axis=1)
    min_rewards = ground_truth.min(axis=1)

    chosen_rewards = np.array([ground_truth[i, choices[i]] for i in range(n)])

    # Accuracy on routable queries (where routing matters)
    if routable_mask is not None:
        choices_r = choices[routable_mask]
        oracle_choices_r = oracle_choices[routable_mask]
        accuracy = (choices_r == oracle_choices_r).mean() if len(choices_r) > 0 else 0.0
    else:
        accuracy = (choices == oracle_choices).mean()

    # frac_oracle on ALL queries with correct tie handling.
    # For ties (all-correct or all-wrong), any choice = oracle -> frac = 1.0.
    # This is the Dean/Hassabis principle: don't filter, measure correctly.
    span = oracle_rewards - min_rewards
    is_tie = span < 1e-6
    per_query_frac = np.where(
        is_tie,
        1.0,  # Ties: any choice matches oracle
        (chosen_rewards - min_rewards) / np.maximum(span, 1e-6),
    )
    frac_oracle = per_query_frac.mean()

    return {
        "name": name,
        "accuracy": accuracy,
        "frac_oracle": frac_oracle,
        "avg_reward": chosen_rewards.mean(),
        "premium_pct": (choices == premium_idx).mean(),
    }


def main():
    # Import here after sys.path is set up
    from embergraphrouter.model.graph_nn import FormData, GNNPredictor
    from embergraphrouter.model.multi_task_graph_router import GraphRouterTrainer

    # Config
    config_path = _MODULE_DIR / "configs/config.yaml"
    with open(config_path) as f:
        config = yaml.safe_load(f)

    llm_desc_path = _MODULE_DIR / config.get(
        "llm_description_path", "configs/LLM_Descriptions.json"
    )
    llm_emb_path = _MODULE_DIR / config.get(
        "llm_embedding_path", "configs/llm_description_embedding.pkl"
    )
    # Use Qwen3 data for evaluation (most recent, largest dataset)
    qwen3_data = _MODULE_DIR / "data/router_data_qwen3.parquet"
    csv_data = _MODULE_DIR / config.get("saved_router_data_path", "data/router_data.csv")
    data_path = qwen3_data if qwen3_data.exists() else csv_data

    # Find checkpoint matching the data
    if "qwen3" in str(data_path):
        # Qwen3 checkpoints with 256-dim embeddings (fastest to load/run)
        checkpoints = [
            "sweep_qwen3_256_best.pth",  # 27MB = 256-dim (fastest)
            "sweep_qwen3_512_best.pth",  # 61MB = 512-dim
        ]
        # Override embedding config for Qwen3 256-dim
        config["embedding_column_suffix"] = "qwen3_256"
    else:
        # Standard 256-dim checkpoints
        checkpoints = [
            "best_model.pth",
            "sweep_gemini_256_best.pth",
            "sweep_openai_large_256_best.pth",
        ]

    checkpoint_path = None
    for name in checkpoints:
        cp = _MODULE_DIR / "model_path" / name
        if cp.exists():
            checkpoint_path = cp
            break

    if not checkpoint_path:
        print("ERROR: No checkpoint found in model_path/")
        return

    print("=" * 80)
    print("EVALUATING WITH ACTUAL MODEL PREDICTIONS")
    print("=" * 80)
    print(f"Config: {config_path}")
    print(f"Data: {data_path}")
    print(f"Checkpoint: {checkpoint_path}")

    # Override GraphRouterTrainer to skip training
    class InferenceOnlyTrainer(GraphRouterTrainer):
        """GraphRouterTrainer that skips training for inference-only mode."""

        def _train(self) -> None:
            """Override _train to skip training in inference mode."""
            logger.info("Inference-only mode: skipping training")
            return

    print("\nInitializing trainer (inference-only mode)...")
    trainer = InferenceOnlyTrainer(
        router_data_path=str(data_path),
        llm_path=str(llm_desc_path),
        llm_embedding_path=str(llm_emb_path),
        config=config,
        wandb=MockWandb(),
    )

    # Load checkpoint (allow missing keys for architecture mismatch)
    print(f"\nLoading checkpoint: {checkpoint_path.name}")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    state_dict = torch.load(checkpoint_path, map_location=device, weights_only=True)
    missing, unexpected = trainer.gnn.model.load_state_dict(state_dict, strict=False)
    if missing:
        print(f"Missing keys (using default init): {len(missing)} keys")
        for k in missing[:3]:
            print(f"  - {k}")
        if len(missing) > 3:
            print(f"  ... and {len(missing) - 3} more")
    if unexpected:
        print(f"Unexpected keys (ignored): {len(unexpected)} keys")
    trainer.gnn.model.eval()
    print("Checkpoint loaded successfully!")

    # Get dimensions
    n_queries = trainer.num_queries
    n_llms = trainer.num_llms
    llm_names = trainer.llm_names
    print(f"\nDataset: {n_queries} queries, {n_llms} LLMs")
    print(f"LLMs: {llm_names}")

    # Get validation mask indices
    valid_mask = trainer.valid_mask  # [num_edges]

    # Create the full Data object for inference
    data = trainer.form_data.formulation(
        task_features=trainer.task_embeddings,
        query_features=trainer.query_embeddings,
        llm_features=trainer.llm_embeddings,
        source_nodes=trainer.source_nodes,
        dest_nodes=trainer.dest_nodes,
        edge_features=trainer.perf_list,  # Performance features per edge
        labels=trainer.labels,
        edge_mask=valid_mask,  # We'll run on validation set
        combined_edge=trainer.combined_edge,
        train_mask=trainer.train_mask,
        valid_mask=valid_mask,
        test_mask=trainer.test_mask,
        query_task_labels=trainer.query_task_labels,
    )

    # Run forward pass on validation set
    print("\nRunning inference on validation set...")

    with torch.no_grad():
        # Prepare edge features (pad to expected dimension)
        edge_weight = torch.zeros(
            data.combined_edge.shape[0], config["edge_dim"], device=device
        )
        edge_weight[:, :3] = data.combined_edge[:, :3]  # cost, quality, latency

        # Edge mask for validation
        mask = valid_mask.to(device)

        # Edge visibility: model can see train edges during validation
        edge_can_see = trainer.train_mask.to(device)

        # Forward pass
        logits = trainer.gnn.model(
            task_features=data.task_id.to(device),
            query_features=data.query_features.to(device),
            llm_features=data.llm_features.to(device),
            edge_index=data.edge_index.to(device),
            edge_mask=mask,
            edge_can_see=edge_can_see,
            edge_weight=edge_weight,
        )

    # Reshape logits to [num_val_queries, num_llms]
    num_val_edges = mask.sum().item()
    num_val_queries = num_val_edges // n_llms
    logits_grouped = logits.cpu().reshape(num_val_queries, n_llms)

    print(f"Validation queries: {num_val_queries}")
    print(f"Logits shape: {logits_grouped.shape}")

    # Get ground truth labels for validation
    # combined_edge[:, 1] is quality (ground truth)
    val_quality = data.combined_edge[mask, 1].cpu().reshape(num_val_queries, n_llms)
    ground_truths = val_quality.numpy()

    # Convert logits to probabilities
    probs = torch.sigmoid(logits_grouped).numpy()
    entropy = compute_entropy(probs)
    max_prob = probs.max(axis=1)
    argmax_choices = probs.argmax(axis=1)

    # Find premium model (Claude Opus 4.5)
    premium_idx = next((i for i, n in enumerate(llm_names) if "Opus" in n), n_llms - 1)
    print(f"\nPremium model: {llm_names[premium_idx]} (idx={premium_idx})")

    # Compute routable mask (Hassabis: separate dataset difficulty from router quality)
    oracle_rewards = ground_truths.max(axis=1)
    min_rewards = ground_truths.min(axis=1)
    routable_mask = (oracle_rewards - min_rewards) > 1e-6  # LLMs have different outcomes

    n_trivial = ((ground_truths.min(axis=1) == 1.0)).sum()  # All correct
    n_impossible = ((ground_truths.max(axis=1) == 0.0)).sum()  # All wrong
    n_routable = routable_mask.sum()

    print(f"\nDataset composition:")
    print(f"  Routable (LLMs differ):    {n_routable:5d} ({100*n_routable/num_val_queries:.1f}%)")
    print(f"  Trivial (all correct):     {n_trivial:5d} ({100*n_trivial/num_val_queries:.1f}%)")
    print(f"  Impossible (all wrong):    {n_impossible:5d} ({100*n_impossible/num_val_queries:.1f}%)")
    print(f"\n  [Metrics below computed on ROUTABLE queries only]")

    results = []

    # Oracle (best possible) - should be 100% on routable queries
    oracle_choices = ground_truths.argmax(axis=1)
    results.append(evaluate_strategy("Oracle", oracle_choices, ground_truths, premium_idx, routable_mask))

    # Random baseline
    np.random.seed(1881)
    random_choices = np.random.randint(0, n_llms, size=num_val_queries)
    results.append(evaluate_strategy("Random", random_choices, ground_truths, premium_idx, routable_mask))

    # Best single model (always pick one)
    single_avg = ground_truths.mean(axis=0)
    best_single_idx = single_avg.argmax()
    always_best = np.full(num_val_queries, best_single_idx)
    results.append(
        evaluate_strategy(
            f"Always {llm_names[best_single_idx][:20]}", always_best, ground_truths, premium_idx, routable_mask
        )
    )

    # Always premium
    always_premium = np.full(num_val_queries, premium_idx)
    results.append(evaluate_strategy("Always Premium", always_premium, ground_truths, premium_idx, routable_mask))

    # Argmax (simple router using model predictions)
    results.append(
        evaluate_strategy("Argmax P(correct)", argmax_choices, ground_truths, premium_idx, routable_mask)
    )

    # Entropy cascade strategies
    for thresh in [0.3, 0.5, 0.7, 0.9]:
        cascade = np.where(entropy > thresh, premium_idx, argmax_choices)
        results.append(
            evaluate_strategy(f"Entropy Cascade (t={thresh})", cascade, ground_truths, premium_idx, routable_mask)
        )

    # Confidence gate strategies
    for thresh in [0.3, 0.5, 0.7, 0.9]:
        gated = np.where(max_prob > thresh, argmax_choices, premium_idx)
        results.append(
            evaluate_strategy(f"Confidence Gate (>{thresh})", gated, ground_truths, premium_idx, routable_mask)
        )

    # Print results
    print("\n" + "=" * 100)
    print("INTELLIGENT ROUTING EVALUATION (ACTUAL MODEL)")
    print("=" * 100)

    print(
        f"\n{'Strategy':<35} {'Accuracy':>10} {'FracOracle':>12} {'AvgReward':>12} {'Premium%':>10}"
    )
    print("-" * 100)

    sorted_results = sorted(results, key=lambda r: -r["frac_oracle"])
    for r in sorted_results:
        marker = " ★" if r["name"] == "Argmax P(correct)" else ""
        print(
            f"{r['name']:<35} {r['accuracy']*100:>9.1f}% "
            f"{r['frac_oracle']*100:>11.1f}% {r['avg_reward']:>11.4f} "
            f"{r['premium_pct']*100:>9.1f}%{marker}"
        )

    # Key insights
    print("\n" + "-" * 100)
    print("KEY INSIGHTS:")
    print("-" * 100)

    oracle = next(r for r in results if r["name"] == "Oracle")
    argmax = next(r for r in results if r["name"] == "Argmax P(correct)")
    baselines = [r for r in results if "Always" in r["name"]]
    best_baseline = max(baselines, key=lambda r: r["frac_oracle"])
    model_strats = [
        r
        for r in results
        if r["name"] not in ["Oracle", "Random"] and "Always" not in r["name"]
    ]
    best_model = max(model_strats, key=lambda r: r["frac_oracle"])

    print(f"  Oracle:                   {oracle['frac_oracle']*100:.1f}% frac_oracle")
    print(f"  Best single model:        {best_baseline['frac_oracle']*100:.1f}% ({best_baseline['name']})")
    print(f"  Simple argmax:            {argmax['frac_oracle']*100:.1f}%")
    print(f"  Best intelligent routing: {best_model['frac_oracle']*100:.1f}% ({best_model['name']})")

    impr_baseline = best_model["frac_oracle"] - best_baseline["frac_oracle"]
    impr_argmax = best_model["frac_oracle"] - argmax["frac_oracle"]

    print()
    if impr_baseline > 0:
        print(f"  ✓ Best strategy beats single model by {impr_baseline*100:+.1f}%")
    else:
        print(f"  ✗ No improvement over single model ({impr_baseline*100:+.1f}%)")

    if impr_argmax > 0 and best_model["name"] != "Argmax P(correct)":
        print(f"  ✓ Entropy/confidence improves over argmax by {impr_argmax*100:+.1f}%")

    # Entropy analysis
    print("\n" + "=" * 100)
    print("PREDICTION ENTROPY ANALYSIS")
    print("=" * 100)
    print(f"\nEntropy: min={entropy.min():.3f}, mean={entropy.mean():.3f}, max={entropy.max():.3f}")
    print(f"Confidence: min={max_prob.min():.3f}, mean={max_prob.mean():.3f}, max={max_prob.max():.3f}")

    # Model prediction distribution
    choices_dist = np.bincount(argmax_choices, minlength=n_llms)
    print(f"\nModel prediction distribution:")
    for i, name in enumerate(llm_names):
        pct = 100 * choices_dist[i] / num_val_queries
        bar = "█" * int(pct / 2)
        print(f"  {name:<25} {pct:5.1f}% {bar}")


if __name__ == "__main__":
    main()
