#!/usr/bin/env python3
"""Properly evaluate v7 model using GraphRouterTrainer's test mechanism.

The GraphRouter requires the full graph context (all queries) for inference.
Single-query inference breaks the GNN message passing design.

This script:
1. Prepares data with 384-dim MiniLM embeddings (matching v7 training)
2. Uses the training framework's test() method for proper evaluation
"""

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

import pandas as pd
import torch

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

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 prepare_minilm_data(source_parquet: Path, output_parquet: Path) -> None:
    """Create a data file with MiniLM 384-dim embeddings renamed for trainer compatibility.

    The v7 model was trained with MiniLM 384-dim embeddings, but the trainer expects
    columns named 'query_embedding' and 'task_description_embedding'.
    """
    logger.info("Preparing data with MiniLM 384-dim embeddings...")

    # Load the multi-embed parquet
    df = pd.read_parquet(source_parquet)

    # Rename MiniLM columns to standard names
    df = df.rename(columns={
        "query_embedding_minilm_384": "query_embedding",
        "task_description_embedding_minilm_384": "task_description_embedding",
    })

    # Keep only columns needed by the trainer
    # Need: task_id, query, ground_truth, metric, task_description, query_embedding, task_description_embedding
    # Plus: LLM response columns
    cols_to_keep = [
        "task_id", "query", "ground_truth", "metric", "task_description",
        "query_embedding", "task_description_embedding",
    ]

    # Check what LLM/effect columns are available
    available_cols = df.columns.tolist()
    logger.info(f"Available columns: {len(available_cols)}")

    # The trainer also needs the performance data - check for router_data.csv columns
    # Load router_data.csv to get the LLM performance columns
    router_csv = _MODULE_DIR / "data" / "router_data.csv"
    if router_csv.exists():
        logger.info("Loading LLM performance data from router_data.csv...")
        router_df = pd.read_csv(router_csv, usecols=lambda c: c not in ['query_embedding', 'task_description_embedding'])

        # Merge on query + task_id
        logger.info(f"Router CSV shape: {router_df.shape}")
        logger.info(f"Multi-embed parquet shape: {df.shape}")

        # Get the columns from router_csv that aren't in multi-embed
        extra_cols = [c for c in router_df.columns if c not in df.columns]
        logger.info(f"Extra columns from router_csv: {extra_cols}")

        # Merge
        df = df.merge(router_df[['query', 'task_id'] + extra_cols], on=['query', 'task_id'], how='inner')
        logger.info(f"Merged shape: {df.shape}")

    # Save
    df.to_parquet(output_parquet, index=False)
    logger.info(f"Saved prepared data to {output_parquet}")
    return output_parquet


def main():
    # Load v7 config
    config_path = _MODULE_DIR / "registry" / "v7_gat4_needle_weighted.yaml"
    model_path = _MODULE_DIR / "registry" / "v7_gat4_needle_weighted_best.pth"

    logger.info(f"Loading config from {config_path}")
    with open(config_path) as f:
        config = yaml.safe_load(f)

    # Override model path to use the v7 checkpoint
    config["model_path"] = str(model_path)
    logger.info(f"Model path: {config['model_path']}")

    # Use zero-padded Gemini 256->384 query embeddings (matching v7 training)
    prepared_data_path = _MODULE_DIR / "data" / "router_data_gemini_padded_384.parquet"
    logger.info(f"Data path: {prepared_data_path}")

    # Create mock wandb to avoid logging
    wandb_mock = MagicMock()

    # Import the trainer class and GNN predictor
    from embergraphrouter.model.multi_task_graph_router import GraphRouterTrainer
    from embergraphrouter.model.graph_nn import GNNPredictor

    # Monkey-patch to skip training but build data
    original_train_validate = GNNPredictor.train_validate

    def skip_train_validate(self, data, data_validate, data_for_test):
        """Skip training, just load the checkpoint."""
        logger.info("SKIPPING training - loading checkpoint instead...")

        # Store masks for test
        self.train_mask = data.train_mask.clone()
        self.valid_mask = data.valid_mask.clone()
        self.test_mask = data.test_mask.clone()

        # Load checkpoint
        checkpoint = torch.load(config["model_path"], map_location="cpu", weights_only=False)
        self.model.load_state_dict(checkpoint)
        self.model.eval()
        logger.info(f"Loaded checkpoint from {config['model_path']}")

    GNNPredictor.train_validate = skip_train_validate

    logger.info("Initializing GraphRouterTrainer with MiniLM 384-dim embeddings...")

    # Resolve paths relative to module dir
    def _resolve(path_str: str) -> str:
        candidate = Path(path_str)
        if candidate.is_absolute():
            return str(candidate)
        return str((_MODULE_DIR / candidate).resolve())

    # Use 7-LLM config (excludes Gemini 3 Flash Preview which was added later)
    llm_path_7 = _MODULE_DIR / "configs" / "LLM_Descriptions_7llm.json"

    # Use zero-padded Gemini embeddings (256 -> 384) to match v7 training
    llm_embedding_path = _MODULE_DIR / "configs" / "llm_description_embedding_gemini_padded.pkl"
    logger.info(f"LLM embedding path: {llm_embedding_path}")

    try:
        trainer = GraphRouterTrainer(
            router_data_path=str(prepared_data_path),  # Use MiniLM data
            llm_path=str(llm_path_7),
            llm_embedding_path=str(llm_embedding_path),
            config=config,
            wandb=wandb_mock,
        )

        # Now run test with proper graph context
        logger.info("Running test with full graph context...")
        metrics, loss = trainer.test()

        print("\n" + "=" * 70)
        print("V7 MODEL EVALUATION (PROPER GRAPH CONTEXT)")
        print("=" * 70)
        print("\nMetrics:")
        for key, value in metrics.items():
            if isinstance(value, float):
                print(f"  {key}: {value:.4f}")
            else:
                print(f"  {key}: {value}")

        print(f"\nTest loss: {loss:.4f}")

    finally:
        # Restore original method
        GNNPredictor.train_validate = original_train_validate


if __name__ == "__main__":
    main()
