#!/usr/bin/env python3
"""Check if the router is collapsing to a single LLM selection."""

import logging
import sys
from pathlib import Path
from unittest.mock import MagicMock

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

_SCRIPT_DIR = Path(__file__).resolve().parent
_MODULE_DIR = _SCRIPT_DIR.parent
sys.path.insert(0, str(_MODULE_DIR.parent))

from embergraphrouter.model.graph_nn import GNNPredictor
from embergraphrouter.model.multi_task_graph_router import GraphRouterTrainer

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

# Load config
config_path = _MODULE_DIR / "registry" / "v7_gat4_needle_weighted.yaml"
model_path = _MODULE_DIR / "registry" / "v7_gat4_needle_weighted_best.pth"

with open(config_path) as f:
    config = yaml.safe_load(f)
config["model_path"] = str(model_path)


def skip_train_validate(self, data, data_validate, data_for_test):
    self.train_mask = data.train_mask.clone()
    self.valid_mask = data.valid_mask.clone()
    self.test_mask = data.test_mask.clone()
    checkpoint = torch.load(config["model_path"], map_location="cpu", weights_only=False)
    self.model.load_state_dict(checkpoint)
    self.model.eval()
    logger.info("Loaded checkpoint")


GNNPredictor.train_validate = skip_train_validate

# Initialize trainer (use 7-LLM config to match training data)
trainer = GraphRouterTrainer(
    router_data_path=str(_MODULE_DIR / "data" / "router_data_minilm_384.parquet"),
    llm_path=str(_MODULE_DIR / "configs" / "LLM_Descriptions_7llm.json"),
    llm_embedding_path=str(_MODULE_DIR / "configs" / "llm_description_embedding.pkl"),
    config=config,
    wandb=MagicMock(),
)

# Run test to get predictions
logger.info("Running test...")
metrics, _ = trainer.test()

# Analyze test metrics
print("\n" + "=" * 60)
print("V7 ROUTER TEST RESULTS")
print("=" * 60)
print()
print("Key Metrics:")
print(f"  frac_of_oracle: {metrics['frac_of_oracle']:.4f} ({metrics['frac_of_oracle']*100:.1f}%)")
print(f"  top1_acc: {metrics['top1_acc']:.4f} ({metrics['top1_acc']*100:.1f}%)")
print(f"  chosen_reward: {metrics['chosen_reward']:.4f}")
print(f"  oracle_reward: {metrics['oracle_reward']:.4f}")
print(f"  regret: {metrics['regret']:.4f}")
print()
print("Quartile Breakdown (frac_oracle):")
print(f"  q0 (easy): {metrics['frac_oracle_q0']:.1%}")
print(f"  q1: {metrics['frac_oracle_q1']:.1%}")
print(f"  q2: {metrics['frac_oracle_q2']:.1%}")
print(f"  q3 (hard): {metrics['frac_oracle_q3']:.1%}")
print()

# Interpret diversity from top1_acc
# With 7 LLMs, random selection would give ~14.3% top1_acc
random_acc = 1.0 / 7  # 0.143
actual_acc = metrics['top1_acc']

print("=" * 60)
print("DIVERSITY ANALYSIS")
print("=" * 60)
print(f"Random selection top1_acc: {random_acc:.1%}")
print(f"Router top1_acc: {actual_acc:.1%}")
print()

if actual_acc > 0.9:
    print("COLLAPSED: Router picks one LLM almost always")
elif actual_acc < random_acc * 1.5:
    print("DIVERSIFIED but POOR: Router makes varied selections but chooses poorly")
else:
    print("MODERATELY SELECTIVE: Router shows some preference but still diversifies")

print()
print("=" * 60)
print("BASELINE COMPARISON")
print("=" * 60)

# Load data for baseline comparison
df = pd.read_parquet(str(_MODULE_DIR / "data" / "router_data_minilm_384.parquet"))
df['group_key'] = df['query'].astype(str) + '|||' + df['task_id'].astype(str)
group_counts = df.groupby('group_key').size()
valid_groups = group_counts[group_counts == 7].index
df = df[df['group_key'].isin(valid_groups)]

# Compute single-model baselines
query_llm_rewards = df.pivot_table(index='group_key', columns='llm', values='effect', aggfunc='first')
oracle = query_llm_rewards.max(axis=1)
worst = query_llm_rewards.min(axis=1)

print()
print("Single-Model frac_oracle (on full dataset):")
for llm in sorted(query_llm_rewards.columns):
    chosen = query_llm_rewards[llm]
    valid_mask = oracle > worst
    frac = ((chosen[valid_mask] - worst[valid_mask]) / (oracle[valid_mask] - worst[valid_mask])).mean()
    marker = " <- BEST" if frac == max(((query_llm_rewards[l][valid_mask] - worst[valid_mask]) / (oracle[valid_mask] - worst[valid_mask])).mean() for l in query_llm_rewards.columns) else ""
    print(f"  {llm:25}: {frac:.1%}{marker}")

print()
print(f"Router v7 (test set): {metrics['frac_of_oracle']:.1%}")

