#!/usr/bin/env python3
"""Evaluate v7 model on test set with correct 384-dim MiniLM embeddings.

The v7 checkpoint was trained with MiniLM 384-dim embeddings from
router_data_multi_embed.parquet, not the 256-dim Gemini embeddings
in router_data.csv.
"""

import ast
import json
import logging
import pickle
import random
import sys
from pathlib import Path

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

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

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)


def load_model(config_path: Path, model_path: Path, input_dim: int = 384):
    """Load trained model with correct input dimension."""
    from embergraphrouter.model.graph_nn import GATEncoderDecoderNet

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

    model = GATEncoderDecoderNet(
        query_feature_dim=input_dim,  # MiniLM 384-dim
        llm_feature_dim=input_dim,    # Same dim for LLM embeddings
        hidden_features=int(config.get("embedding_dim", 256)),
        in_edges=int(config.get("edge_dim", 6)),
        dropout=float(config.get("dropout", 0.25)),
        num_gnn_layers=int(config.get("num_gnn_layers", 4)),
        use_difficulty_head=False,
        num_attention_heads=int(config.get("num_attention_heads", 4)),
        attention_dropout=float(config.get("attention_dropout", 0.15)),
        hidden_multiplier=int(config.get("hidden_multiplier", 3)),
        query_mlp_layers=int(config.get("query_mlp_layers", 1)),
        llm_mlp_layers=int(config.get("llm_mlp_layers", 1)),
        input_dropout=float(config.get("input_dropout", 0.1)),
        use_residual=bool(config.get("use_residual", True)),
        use_swiglu=bool(config.get("use_swiglu", False)),
        use_variational=bool(config.get("use_variational", False)),
        use_contrastive=bool(config.get("use_contrastive", True)),
        contrastive_weight=float(config.get("contrastive_weight", 0.1)),
    )

    checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
    model.load_state_dict(checkpoint)
    model.eval()

    return model, config


def load_data_with_minilm_embeddings(embed_path: Path, rewards_path: Path, llm_embedding_path: Path):
    """Load router data with 384-dim MiniLM embeddings.

    The multi_embed parquet has embeddings but not LLM rewards.
    We join with router_data.csv to get the rewards.
    """
    logger.info(f"Loading embeddings from {embed_path}")

    # Load embeddings from multi_embed parquet
    embed_df = pd.read_parquet(
        embed_path,
        columns=[
            "task_id", "query",
            "query_embedding_minilm_384",
            "task_description_embedding_minilm_384",
        ]
    )
    logger.info(f"Loaded {len(embed_df)} unique queries with embeddings")

    # Load rewards from router_data.csv
    logger.info(f"Loading rewards from {rewards_path}")
    rewards_df = pd.read_csv(
        rewards_path,
        usecols=["task_id", "query", "llm", "effect", "cost", "latency"]
    )
    logger.info(f"Loaded {len(rewards_df)} reward rows")

    # Merge: join embeddings with rewards
    df = rewards_df.merge(
        embed_df,
        on=["task_id", "query"],
        how="inner"
    )
    logger.info(f"Merged data: {len(df)} rows")

    # Load LLM embeddings
    with open(llm_embedding_path, "rb") as f:
        llm_embeddings = pickle.load(f)

    logger.info(f"LLM embeddings: {llm_embeddings.shape}")

    return df, llm_embeddings


def parse_embedding(emb_str):
    """Parse embedding string to numpy array."""
    if pd.isna(emb_str) or emb_str == "" or emb_str == "[]":
        return None
    try:
        return np.array(ast.literal_eval(emb_str), dtype=np.float32)
    except:
        return None


def main():
    config_path = _MODULE_DIR / "registry" / "v7_gat4_needle_weighted.yaml"
    model_path = _MODULE_DIR / "registry" / "v7_gat4_needle_weighted_best.pth"
    data_path = _MODULE_DIR / "data" / "router_data_multi_embed.parquet"
    llm_embedding_path = _MODULE_DIR / "configs" / "llm_description_embedding.pkl"

    # Check if multi_embed parquet exists
    if not data_path.exists():
        logger.error(f"Required data file not found: {data_path}")
        logger.error("The v7 model requires router_data_multi_embed.parquet with MiniLM 384-dim embeddings")
        sys.exit(1)

    logger.info("Loading model...")
    model, config = load_model(config_path, model_path, input_dim=384)

    logger.info("Loading data with MiniLM 384-dim embeddings...")
    rewards_path = _MODULE_DIR / "data" / "router_data.csv"
    df, llm_embeddings = load_data_with_minilm_embeddings(
        embed_path=data_path,
        rewards_path=rewards_path,
        llm_embedding_path=llm_embedding_path
    )

    # Check LLM embedding dimension
    if llm_embeddings.shape[1] != 384:
        logger.warning(
            f"LLM embeddings are {llm_embeddings.shape[1]}-dim but model expects 384-dim. "
            "Need to regenerate LLM embeddings with MiniLM."
        )
        # For now, pad or truncate
        if llm_embeddings.shape[1] < 384:
            pad_width = 384 - llm_embeddings.shape[1]
            llm_embeddings = np.pad(llm_embeddings, ((0, 0), (0, pad_width)))
            logger.info(f"Padded LLM embeddings to 384-dim")
        else:
            llm_embeddings = llm_embeddings[:, :384]
            logger.info(f"Truncated LLM embeddings to 384-dim")

    # Get unique LLMs
    llm_names = df["llm"].unique().tolist()
    num_llms = len(llm_names)
    llm_to_idx = {llm: i for i, llm in enumerate(llm_names)}
    logger.info(f"LLMs ({num_llms}): {llm_names}")

    # Group by query to get per-query data
    # Each query should have exactly num_llms rows
    query_groups = df.groupby(["task_id", "query"])
    valid_groups = [(name, group) for name, group in query_groups if len(group) == num_llms]
    logger.info(f"Valid query groups: {len(valid_groups)}")

    # Parse embeddings and prepare data
    logger.info("Parsing embeddings...")
    queries_data = []
    for (task_id, query_text), group in valid_groups:
        query_emb = parse_embedding(group["query_embedding_minilm_384"].iloc[0])
        task_emb = parse_embedding(group["task_description_embedding_minilm_384"].iloc[0])

        if query_emb is None or task_emb is None:
            continue

        # Get rewards per LLM (sorted by llm_to_idx order)
        rewards = []
        for llm in llm_names:
            row = group[group["llm"] == llm]
            if len(row) > 0:
                rewards.append(row["effect"].iloc[0])
            else:
                rewards.append(0.0)

        queries_data.append({
            "task_id": task_id,
            "query": query_text,
            "query_emb": query_emb,
            "task_emb": task_emb,
            "rewards": np.array(rewards, dtype=np.float32),
        })

    logger.info(f"Prepared {len(queries_data)} queries with valid embeddings")

    # Create train/val/test split (same as training)
    random.seed(config.get("seed", 1881))
    random.shuffle(queries_data)

    n = len(queries_data)
    split_ratio = config.get("split_ratio", [0.7, 0.1, 0.2])
    train_end = int(n * split_ratio[0])
    val_end = train_end + int(n * split_ratio[1])

    test_data = queries_data[val_end:]
    logger.info(f"Test set: {len(test_data)} queries")

    # Run inference
    logger.info("Running inference on test set...")

    llm_features = torch.tensor(llm_embeddings, dtype=torch.float32)

    correct = 0
    total = 0
    chosen_rewards = []
    oracle_rewards = []
    frac_oracle_scores = []
    chosen_llm_counts = {i: 0 for i in range(num_llms)}
    oracle_llm_counts = {i: 0 for i in range(num_llms)}

    # Debug first query
    first_debug = True

    with torch.no_grad():
        for item in test_data:
            query_emb = torch.tensor(item["query_emb"], dtype=torch.float32).unsqueeze(0)
            task_emb = torch.tensor(item["task_emb"], dtype=torch.float32).unsqueeze(0)
            rewards = item["rewards"]

            # Create edge index for single query -> all LLMs
            edge_index = torch.tensor([
                [0] * num_llms,  # source: query 0
                list(range(1, num_llms + 1))  # dest: LLMs 1 to num_llms
            ], dtype=torch.long)

            # Edge features (placeholder zeros)
            edge_weight = torch.zeros(num_llms, 6)

            # Edge masks for inference
            # edge_mask = True means we want predictions for these edges
            # edge_can_see = False means NO message passing (critical for single-query inference)
            # With only 1 query, GNN message passing makes all LLM representations identical
            edge_mask = torch.ones(num_llms, dtype=torch.bool)
            edge_can_see = torch.zeros(num_llms, dtype=torch.bool)  # No message passing!

            if first_debug:
                print("\n--- DEBUG: First query ---")
                print(f"  query_emb shape: {query_emb.shape}, range: [{query_emb.min():.3f}, {query_emb.max():.3f}]")
                print(f"  task_emb shape: {task_emb.shape}")
                print(f"  llm_features shape: {llm_features.shape}")
                print(f"  edge_index: {edge_index}")
                print(f"  edge_weight: {edge_weight.shape}")

                # Debug: Check if LLM embeddings are different
                print(f"\n  LLM embedding norms:")
                for i in range(num_llms):
                    norm = llm_features[i].norm().item()
                    print(f"    LLM {i} ({llm_names[i][:20]}): norm={norm:.4f}")

                # Debug: Check FeatureAlign output
                x_initial = model.model_align(task_emb, query_emb, llm_features)
                print(f"\n  FeatureAlign output shape: {x_initial.shape}")
                print(f"  Query repr (row 0) norm: {x_initial[0].norm().item():.4f}")
                for i in range(num_llms):
                    print(f"  LLM {i} repr (row {i+1}) norm: {x_initial[i+1].norm().item():.4f}")

            # Forward pass
            logits = model(
                task_features=task_emb,
                query_features=query_emb,
                llm_features=llm_features,
                edge_index=edge_index,
                edge_mask=edge_mask,
                edge_can_see=edge_can_see,
                edge_weight=edge_weight,
            )

            # Get predictions
            probs = torch.sigmoid(logits.squeeze())
            chosen_idx = probs.argmax().item()
            oracle_idx = rewards.argmax()

            if first_debug:
                print(f"  logits shape: {logits.shape}")
                print(f"  logits: {logits.squeeze()}")
                print(f"  probs: {probs}")
                print(f"  chosen_idx: {chosen_idx}, oracle_idx: {oracle_idx}")
                print(f"  rewards: {rewards}")
                first_debug = False

            chosen_reward = rewards[chosen_idx]
            oracle_reward = rewards[oracle_idx]
            min_reward = rewards.min()

            chosen_rewards.append(chosen_reward)
            oracle_rewards.append(oracle_reward)

            if chosen_idx == oracle_idx:
                correct += 1
            total += 1

            # Track LLM selection distribution
            chosen_llm_counts[chosen_idx] += 1
            oracle_llm_counts[oracle_idx] += 1

            # Frac of oracle with correct tie handling.
            # For ties (all-correct or all-wrong), any choice = oracle -> frac = 1.0.
            span = oracle_reward - min_reward
            if span < 1e-6:
                frac = 1.0  # Tie: any choice matches oracle
            else:
                frac = (chosen_reward - min_reward) / span
            frac_oracle_scores.append(frac)

    # Compute metrics
    accuracy = correct / total if total > 0 else 0
    avg_chosen = np.mean(chosen_rewards)
    avg_oracle = np.mean(oracle_rewards)
    frac_oracle = np.mean(frac_oracle_scores)

    # Print LLM selection distribution
    print("\n--- LLM Selection Distribution ---")
    print(f"  {'LLM':<30} {'Chosen':<10} {'Oracle':<10}")
    print("  " + "-" * 50)
    for i, llm in enumerate(llm_names):
        print(f"  {llm:<30} {chosen_llm_counts[i]:<10} {oracle_llm_counts[i]:<10}")

    print("\n" + "=" * 70)
    print("V7 MODEL - ACTUAL EVALUATION WITH MINILM 384-DIM EMBEDDINGS")
    print("=" * 70)
    print(f"\nTest set: {total} queries")
    print(f"\n--- Core Metrics ---")
    print(f"  Top-1 Accuracy:  {accuracy*100:.2f}%")
    print(f"  Frac of Oracle:  {frac_oracle*100:.2f}%")
    print(f"  Avg Chosen:      {avg_chosen:.4f}")
    print(f"  Avg Oracle:      {avg_oracle:.4f}")

    print("\n" + "=" * 70)
    print("COMPARISON WITH BASELINES")
    print("=" * 70)
    print(f"\n  v7 Router:         {frac_oracle*100:>6.2f}% frac_oracle")
    print(f"  Best single model:   53.80% frac_oracle (Opus 4.5)")
    print(f"  Random:              42.90% frac_oracle")
    print(f"\n  Improvement over best single: {(frac_oracle - 0.538)*100:>+6.2f}%")


if __name__ == "__main__":
    main()
