#!/usr/bin/env python3
"""Evaluate trained model and extract actual predictions.

This script properly constructs the heterogeneous graph and runs actual
model inference to get real router selections (not simulated).

Source: Actual model inference via heterogeneous graph pipeline.
"""

import argparse
import logging
import sys
from pathlib import Path

import numpy as np
import torch
import yaml

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

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


def main():
    parser = argparse.ArgumentParser(
        description="Evaluate model with actual predictions (not simulated)"
    )
    parser.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to config YAML",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Path to save predictions (optional)",
    )
    parser.add_argument(
        "--skip-training",
        action="store_true",
        help="Skip training (load existing weights only)",
    )
    args = parser.parse_args()

    config_path = Path(args.config)
    with open(config_path) as f:
        config = yaml.safe_load(f)

    # Set epochs to 0 to skip training if requested
    if args.skip_training:
        config["train_epoch"] = 0

    base_dir = config_path.parent.parent
    router_data_path = config.get("saved_router_data_path", "data/router_data.parquet")
    llm_path = config.get("llm_description_path", "configs/LLM_Descriptions.json")
    llm_embedding_path = config.get(
        "llm_embedding_path", "configs/llm_description_embedding.pkl"
    )

    if not Path(router_data_path).is_absolute():
        router_data_path = str(base_dir / router_data_path)
    if not Path(llm_path).is_absolute():
        llm_path = str(base_dir / llm_path)
    if not Path(llm_embedding_path).is_absolute():
        llm_embedding_path = str(base_dir / llm_embedding_path)

    logger.info("Config: %s", args.config)
    logger.info("Router data: %s", router_data_path)
    logger.info("LLM descriptions: %s", llm_path)
    logger.info("LLM embeddings: %s", llm_embedding_path)

    # Import here to avoid issues
    from embergraphrouter.model.multi_task_graph_router import GraphRouterTrainer

    logger.info("Initializing GraphRouterTrainer...")
    trainer = GraphRouterTrainer(
        router_data_path=router_data_path,
        llm_path=llm_path,
        llm_embedding_path=llm_embedding_path,
        config=config,
        wandb=None,
    )

    # Load best weights
    checkpoint_path = config["model_path"]
    if not Path(checkpoint_path).exists():
        checkpoint_path = str(base_dir / checkpoint_path)

    logger.info("Loading checkpoint: %s", 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()

    # Run test evaluation
    logger.info("Running test evaluation with actual model inference...")
    metrics, loss = trainer.test()

    # Get actual predictions from test data
    test_data = trainer._test_data
    mask = test_data.edge_mask.clone()
    edge_can_see = trainer.gnn.valid_mask | trainer.gnn.train_mask

    wp = config.get("inference_weights", {}).get("perf", 0.7)
    wc = config.get("inference_weights", {}).get("cost", 0.2)
    wl = config.get("inference_weights", {}).get("latency", 0.1)

    with torch.no_grad():
        augmented_edge = trainer.gnn._edge_features_with_pref(
            test_data.combined_edge, wp, wc, wl
        )
        logits = trainer.gnn.model(
            task_features=test_data.task_id,
            query_features=test_data.query_features,
            llm_features=test_data.llm_features,
            edge_index=test_data.edge_index,
            edge_mask=mask,
            edge_can_see=edge_can_see,
            edge_weight=augmented_edge,
        )

    # Reshape to (num_queries, num_llms)
    scores = logits.reshape(-1, trainer.num_llms)
    router_choices = scores.argmax(dim=1).cpu().numpy()

    # Get ground truth for comparison
    rewards = trainer.gnn._rewards_from_pref(test_data.combined_edge, wp, wc, wl)
    rewards_grouped = rewards[mask].reshape(-1, trainer.num_llms)
    oracle_choices = rewards_grouped.argmax(dim=1).cpu().numpy()

    # Analyze selection diversity
    print("\n" + "=" * 70)
    print("ACTUAL MODEL PREDICTIONS (Source: model inference)")
    print("=" * 70)

    llm_names = trainer.llm_names
    num_test = len(router_choices)

    print(f"\nTest queries: {num_test}")
    print(f"LLMs: {llm_names}")

    # Selection distribution
    router_counts = np.bincount(router_choices, minlength=trainer.num_llms)
    oracle_counts = np.bincount(oracle_choices, minlength=trainer.num_llms)

    print("\n--- Selection Distribution ---")
    print(f"{'Model':<30} {'Oracle %':>12} {'Router %':>12} {'Delta':>12}")
    print("-" * 70)

    for i, name in enumerate(llm_names):
        oracle_pct = 100 * oracle_counts[i] / num_test
        router_pct = 100 * router_counts[i] / num_test
        delta = router_pct - oracle_pct
        print(f"{name:<30} {oracle_pct:>11.1f}% {router_pct:>11.1f}% {delta:>+11.1f}%")

    # Entropy analysis
    def compute_entropy(counts):
        probs = counts / counts.sum()
        probs = probs[probs > 0]
        return -np.sum(probs * np.log(probs))

    max_entropy = np.log(trainer.num_llms)
    oracle_entropy = compute_entropy(oracle_counts)
    router_entropy = compute_entropy(router_counts)

    print("\n--- Diversity Analysis ---")
    print(f"Oracle entropy:  {oracle_entropy:.3f} ({100*oracle_entropy/max_entropy:.1f}% of max)")
    print(f"Router entropy:  {router_entropy:.3f} ({100*router_entropy/max_entropy:.1f}% of max)")
    print(f"Max entropy:     {max_entropy:.3f} (uniform)")

    # Agreement
    agreement = (router_choices == oracle_choices).mean()
    print(f"\nRouter-Oracle agreement: {100*agreement:.1f}%")

    # Print metrics
    print("\n--- Test Metrics ---")
    for k, v in sorted(metrics.items()):
        if isinstance(v, float):
            print(f"{k}: {v:.4f}")

    # Save predictions if requested
    if args.output:
        np.savez(
            args.output,
            router_choices=router_choices,
            oracle_choices=oracle_choices,
            scores=scores.cpu().numpy(),
            rewards=rewards_grouped.cpu().numpy(),
            llm_names=llm_names,
        )
        logger.info("Saved predictions to %s", args.output)

    return metrics, router_choices, oracle_choices


if __name__ == "__main__":
    main()
