#!/usr/bin/env python3
"""Evaluate the best trained model on the test set.

Properly constructs the heterogeneous graph and computes all metrics.
"""

import argparse
import logging
import sys
from pathlib import Path

import torch
import yaml

# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from embergraphrouter.model.multi_task_graph_router import GraphRouterTrainer

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 trained model on test set")
    parser.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to config YAML (e.g., configs/hybrid_lambda_rank_openai512.yaml)",
    )
    parser.add_argument(
        "--checkpoint",
        type=str,
        default=None,
        help="Path to model checkpoint (defaults to model_path in config)",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use",
    )
    args = parser.parse_args()

    # Load config
    config_path = Path(args.config)
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    # Override model path if specified
    if args.checkpoint:
        config["model_path"] = args.checkpoint

    # Resolve paths relative to config directory
    base_dir = config_path.parent.parent  # embergraphrouter/
    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("Checkpoint: %s", config.get("model_path"))
    logger.info("Router data: %s", router_data_path)
    logger.info("Device: %s", args.device)
    logger.info("LLM count: %s", config.get("llm_num", 7))

    # Initialize the router (this loads data and constructs the graph)
    logger.info("Initializing GraphRouterTrainer...")
    router = GraphRouterTrainer(
        router_data_path=router_data_path,
        llm_path=llm_path,
        llm_embedding_path=llm_embedding_path,
        config=config,
        wandb=None,
    )

    # Load trained weights into the GNN
    checkpoint_path = config["model_path"]
    logger.info("Loading checkpoint: %s", checkpoint_path)
    state_dict = torch.load(checkpoint_path, map_location=args.device, weights_only=True)
    router.gnn.model.load_state_dict(state_dict)
    router.gnn.model.eval()

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

    logger.info("Inference weights: perf=%.2f, cost=%.2f, latency=%.2f", wp, wc, wl)

    # Run test evaluation using the built-in method
    logger.info("Running test evaluation...")
    metrics = router.test_with(perf=wp, cost=wc, latency=wl)

    # Print results
    print("\n" + "=" * 60)
    print("TEST SET EVALUATION RESULTS")
    print("=" * 60)
    print(f"Config: {args.config}")
    print(f"Checkpoint: {checkpoint_path}")
    print(f"LLMs: {router.num_llms}")
    print(f"Test queries: {len(router.gnn.test_indices)}")
    print(f"Preferences: perf={wp}, cost={wc}, latency={wl}")
    print("-" * 60)
    print(f"frac_oracle:      {metrics.get('frac_of_oracle', 0):.4f} ({metrics.get('frac_of_oracle', 0)*100:.2f}%)")
    print(f"top1_acc:         {metrics.get('top1_acc', 0):.4f}")
    print(f"regret:           {metrics.get('regret', 0):.4f}")
    print(f"ndcg@1:           {metrics.get('ndcg@1', 0):.4f}")
    print(f"ndcg@3:           {metrics.get('ndcg@3', 0):.4f}")
    print(f"ndcg@All:         {metrics.get('ndcg@All', 0):.4f}")
    print(f"ECE:              {metrics.get('ece', 0):.4f}")
    print(f"avg_chosen:       {metrics.get('avg_chosen_reward', 0):.4f}")
    print(f"avg_oracle:       {metrics.get('avg_oracle_reward', 0):.4f}")
    print("=" * 60)

    # Also compare against always choosing best single model
    print("\n--- Comparison vs Single Model Baselines ---")
    print("(Always choosing one LLM regardless of query)")
    # This would require computing per-LLM performance, which the trainer may support

    return metrics


if __name__ == "__main__":
    main()
