#!/usr/bin/env python3
"""Enrich router_data.csv with interpretable labels.

This script adds derived columns that make the dataset self-documenting:
- Absolute cost/latency values (not just normalized)
- Model metadata (tier, provider, pricing)
- Task capability taxonomy
- Query difficulty scores
- Response type classification (success/refusal/error)

Usage:
    python -m embergraphrouter.data_processing.enrich_router_data

Output:
    data/router_data_enriched.csv - Full enriched dataset
    data/router_data_enriched.parquet - Parquet format for fast analysis
"""

import hashlib
import json
import logging
from pathlib import Path

import pandas as pd

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

_MODULE_DIR = Path(__file__).resolve().parent.parent


# =============================================================================
# CAPABILITY TAXONOMY (Hassabis-style decomposition)
# =============================================================================

TASK_TAXONOMY = {
    # Knowledge & Reasoning (Multiple Choice)
    "mmlu": {
        "capability_primary": "knowledge",
        "domain": "general",
        "capability_tags": ["multi_choice", "factual_recall"],
    },
    "mmlu_pro_chemistry": {
        "capability_primary": "knowledge",
        "domain": "science",
        "capability_tags": ["multi_choice", "domain_expertise", "scientific_reasoning"],
    },
    "mmlu_pro_biology": {
        "capability_primary": "knowledge",
        "domain": "science",
        "capability_tags": ["multi_choice", "domain_expertise", "scientific_reasoning"],
    },
    "mmlu_pro_law": {
        "capability_primary": "reasoning",
        "domain": "law",
        "capability_tags": ["multi_choice", "domain_expertise", "logical_reasoning"],
    },
    "mmlu_pro_psychology": {
        "capability_primary": "knowledge",
        "domain": "social_science",
        "capability_tags": ["multi_choice", "domain_expertise"],
    },
    "mmlu_pro_business": {
        "capability_primary": "reasoning",
        "domain": "business",
        "capability_tags": ["multi_choice", "domain_expertise", "quantitative"],
    },
    "mmlu_pro_history": {
        "capability_primary": "knowledge",
        "domain": "humanities",
        "capability_tags": ["multi_choice", "domain_expertise", "temporal_reasoning"],
    },
    "mmlu_pro_other": {
        "capability_primary": "knowledge",
        "domain": "general",
        "capability_tags": ["multi_choice", "domain_expertise"],
    },
    # Commonsense Reasoning
    "arc": {
        "capability_primary": "reasoning",
        "domain": "science",
        "capability_tags": ["multi_choice", "commonsense", "scientific_reasoning"],
    },
    "hellaswag": {
        "capability_primary": "reasoning",
        "domain": "general",
        "capability_tags": ["multi_choice", "commonsense", "narrative_completion"],
    },
    "winogrande": {
        "capability_primary": "reasoning",
        "domain": "language",
        "capability_tags": ["multi_choice", "coreference", "commonsense"],
    },
    "agieval": {
        "capability_primary": "reasoning",
        "domain": "general",
        "capability_tags": ["multi_choice", "logical_reasoning", "reading_comprehension"],
    },
    "truthfulqa": {
        "capability_primary": "knowledge",
        "domain": "general",
        "capability_tags": ["multi_choice", "factual_accuracy", "bias_resistance"],
    },
    # Math
    "GSM8K": {
        "capability_primary": "reasoning",
        "domain": "math",
        "capability_tags": ["arithmetic", "word_problems", "multi_step"],
    },
    "mathvista": {
        "capability_primary": "reasoning",
        "domain": "math",
        "capability_tags": ["visual_reasoning", "mathematical"],
    },
    "mathvqa": {
        "capability_primary": "reasoning",
        "domain": "math",
        "capability_tags": ["visual_reasoning", "mathematical"],
    },
    # Generation
    "alpaca_data": {
        "capability_primary": "generation",
        "domain": "general",
        "capability_tags": ["instruction_following", "open_ended", "helpfulness"],
    },
    "SQUAD": {
        "capability_primary": "reasoning",
        "domain": "language",
        "capability_tags": ["reading_comprehension", "extractive_qa"],
    },
    "multi_news": {
        "capability_primary": "generation",
        "domain": "language",
        "capability_tags": ["summarization", "long_context"],
    },
    # Code
    "mbpp": {
        "capability_primary": "code",
        "domain": "programming",
        "capability_tags": ["code_generation", "python"],
    },
    "livecodebench": {
        "capability_primary": "code",
        "domain": "programming",
        "capability_tags": ["code_generation", "algorithms"],
    },
    # BBEH (BIG-Bench Extra Hard)
    "bbeh_multistep_arithmetic": {
        "capability_primary": "reasoning",
        "domain": "math",
        "capability_tags": ["arithmetic", "multi_step", "hard"],
    },
    "bbeh_buggy_tables": {
        "capability_primary": "reasoning",
        "domain": "logic",
        "capability_tags": ["table_reasoning", "error_detection", "hard"],
    },
    "bbeh_geometric_shapes": {
        "capability_primary": "reasoning",
        "domain": "math",
        "capability_tags": ["spatial_reasoning", "geometry", "hard"],
    },
    "bbeh_hyperbaton": {
        "capability_primary": "reasoning",
        "domain": "language",
        "capability_tags": ["syntax", "linguistic", "hard"],
    },
    "bbeh_linguini": {
        "capability_primary": "reasoning",
        "domain": "language",
        "capability_tags": ["linguistic", "hard"],
    },
}

# Default for unknown tasks
DEFAULT_TAXONOMY = {
    "capability_primary": "unknown",
    "domain": "unknown",
    "capability_tags": [],
}


# =============================================================================
# MODEL METADATA
# =============================================================================

MODEL_TIERS = {
    "GPT-5 Nano": "budget",
    "Gemini 2.5 Flash Lite": "budget",
    "Claude Haiku 4.5": "standard",
    "GPT-5.1": "standard",
    "GPT-5.1 Codex": "standard",
    "Gemini 3 Pro Preview": "standard",
    "Claude Opus 4.5": "premium",
}

MODEL_PROVIDERS = {
    "GPT-5 Nano": "openai",
    "GPT-5.1": "openai",
    "GPT-5.1 Codex": "openai",
    "Gemini 2.5 Flash Lite": "google",
    "Gemini 3 Pro Preview": "google",
    "Claude Haiku 4.5": "anthropic",
    "Claude Opus 4.5": "anthropic",
}


def _query_hash(query: str, llm_name: str) -> str:
    """Generate cache key for query+llm pair."""
    return hashlib.sha256(f"{query}||{llm_name}".encode()).hexdigest()[:16]


def enrich_router_data(
    router_path: Path,
    cache_path: Path,
    llm_desc_path: Path,
    output_path: Path,
) -> pd.DataFrame:
    """Enrich router_data.csv with interpretable labels."""

    logger.info("Loading data...")
    df = pd.read_csv(router_path)
    cache = json.loads(cache_path.read_text())
    llm_desc = json.loads(llm_desc_path.read_text())

    logger.info(f"  Router data: {len(df):,} rows")
    logger.info(f"  Cache entries: {len(cache):,}")

    # =========================================================================
    # 1. ABSOLUTE LATENCY (from cache)
    # =========================================================================
    logger.info("Adding absolute latency from cache...")

    def get_latency_sec(row):
        key = _query_hash(row["query"], row["llm"])
        entry = cache.get(key, {})
        return entry.get("latency", 0.0)

    df["latency_sec"] = df.apply(get_latency_sec, axis=1)

    # =========================================================================
    # 2. RESPONSE TYPE (success/refusal/empty)
    # =========================================================================
    logger.info("Classifying response types...")

    def get_response_type(row):
        key = _query_hash(row["query"], row["llm"])
        entry = cache.get(key, {})
        if entry.get("refusal"):
            return "refusal"
        response = entry.get("response", "")
        if not response or not response.strip():
            return "empty"
        return "success"

    df["response_type"] = df.apply(get_response_type, axis=1)

    # =========================================================================
    # 3. MODEL METADATA (tier, provider, pricing)
    # =========================================================================
    logger.info("Adding model metadata...")

    df["model_tier"] = df["llm"].map(MODEL_TIERS).fillna("unknown")
    df["model_provider"] = df["llm"].map(MODEL_PROVIDERS).fillna("unknown")

    # Add pricing from LLM descriptions
    df["input_price_per_m"] = df["llm"].map(
        lambda x: llm_desc.get(x, {}).get("input_price", 0.0)
    )
    df["output_price_per_m"] = df["llm"].map(
        lambda x: llm_desc.get(x, {}).get("output_price", 0.0)
    )

    # =========================================================================
    # 4. TASK CAPABILITY TAXONOMY
    # =========================================================================
    logger.info("Adding capability taxonomy...")

    def get_taxonomy_field(task_id: str, field: str):
        taxonomy = TASK_TAXONOMY.get(task_id, DEFAULT_TAXONOMY)
        return taxonomy.get(field, DEFAULT_TAXONOMY.get(field))

    df["capability_primary"] = df["task_id"].apply(
        lambda x: get_taxonomy_field(x, "capability_primary")
    )
    df["domain"] = df["task_id"].apply(
        lambda x: get_taxonomy_field(x, "domain")
    )
    df["capability_tags"] = df["task_id"].apply(
        lambda x: ",".join(get_taxonomy_field(x, "capability_tags"))
    )

    # =========================================================================
    # 5. QUERY DIFFICULTY (avg effect across all models)
    # =========================================================================
    logger.info("Computing query difficulty scores...")

    # Group by query to get difficulty
    query_difficulty = df.groupby("query")["effect"].mean()
    df["query_difficulty"] = df["query"].map(query_difficulty)

    # Model advantage: how much better/worse than average on this query
    df["model_advantage"] = df["effect"] - df["query_difficulty"]

    # =========================================================================
    # 6. TASK-LEVEL AGGREGATES
    # =========================================================================
    logger.info("Computing task-level aggregates...")

    # Task difficulty (avg effect across all models and queries in task)
    task_difficulty = df.groupby("task_id")["effect"].mean()
    df["task_difficulty"] = df["task_id"].map(task_difficulty)

    # Model-task affinity: how much better is this model on this task vs overall
    model_task_effect = df.groupby(["llm", "task_id"])["effect"].mean()
    model_overall_effect = df.groupby("llm")["effect"].mean()

    def get_model_task_affinity(row):
        model_on_task = model_task_effect.get((row["llm"], row["task_id"]), 0)
        model_overall = model_overall_effect.get(row["llm"], 0)
        return model_on_task - model_overall

    df["model_task_affinity"] = df.apply(get_model_task_affinity, axis=1)

    # =========================================================================
    # SAVE ENRICHED DATA
    # =========================================================================
    logger.info(f"Saving enriched data to {output_path}...")

    # Reorder columns for clarity
    column_order = [
        # Identifiers
        "task_id", "query", "llm",
        # Task taxonomy
        "capability_primary", "domain", "capability_tags",
        # Model metadata
        "model_tier", "model_provider", "input_price_per_m", "output_price_per_m",
        # Outcomes
        "effect", "response_type",
        # Absolute measurements
        "latency_sec",
        # Normalized (for training)
        "cost", "latency",
        # Derived insights
        "query_difficulty", "model_advantage",
        "task_difficulty", "model_task_affinity",
        # Original fields
        "ground_truth", "metric", "task_description",
        # Embeddings (at end - large)
        "query_embedding", "task_description_embedding",
    ]

    # Keep only columns that exist
    column_order = [c for c in column_order if c in df.columns]
    df = df[column_order]

    # Save as CSV
    df.to_csv(output_path, index=False)
    logger.info(f"  Saved CSV: {output_path}")

    # Save as Parquet for faster analysis
    parquet_path = output_path.with_suffix(".parquet")
    df.to_parquet(parquet_path, index=False)
    logger.info(f"  Saved Parquet: {parquet_path}")

    # =========================================================================
    # PRINT SUMMARY STATISTICS
    # =========================================================================
    print("\n" + "=" * 70)
    print("ENRICHED DATASET SUMMARY")
    print("=" * 70)

    print(f"\nShape: {df.shape[0]:,} rows × {df.shape[1]} columns")

    print("\n--- Response Types ---")
    print(df["response_type"].value_counts().to_string())

    print("\n--- Model Tiers ---")
    tier_perf = df.groupby("model_tier")["effect"].agg(["mean", "count"])
    print(tier_perf.to_string())

    print("\n--- Capability Performance ---")
    cap_perf = df.groupby("capability_primary")["effect"].agg(["mean", "count"])
    cap_perf = cap_perf.sort_values("mean", ascending=False)
    print(cap_perf.to_string())

    print("\n--- Domain Performance ---")
    domain_perf = df.groupby("domain")["effect"].agg(["mean", "count"])
    domain_perf = domain_perf.sort_values("mean", ascending=False)
    print(domain_perf.to_string())

    print("\n--- Query Difficulty Distribution ---")
    print(f"  Mean: {df['query_difficulty'].mean():.3f}")
    print(f"  Std:  {df['query_difficulty'].std():.3f}")
    print(f"  Min:  {df['query_difficulty'].min():.3f}")
    print(f"  Max:  {df['query_difficulty'].max():.3f}")

    # Hardest and easiest queries
    unique_queries = df.groupby("query").first().reset_index()
    print("\n--- Hardest Queries (all models fail) ---")
    hardest = unique_queries.nsmallest(3, "query_difficulty")
    for _, row in hardest.iterrows():
        print(f"  [{row['task_id']}] difficulty={row['query_difficulty']:.3f}")
        print(f"    {row['query'][:80]}...")

    print("\n--- Easiest Queries (all models succeed) ---")
    easiest = unique_queries.nlargest(3, "query_difficulty")
    for _, row in easiest.iterrows():
        print(f"  [{row['task_id']}] difficulty={row['query_difficulty']:.3f}")
        print(f"    {row['query'][:80]}...")

    return df


if __name__ == "__main__":
    data_dir = _MODULE_DIR / "data"

    router_path = data_dir / "router_data.csv"
    cache_path = data_dir / ".router_cache" / "response_cache.json"
    llm_desc_path = _MODULE_DIR / "configs" / "LLM_Descriptions.json"
    output_path = data_dir / "router_data_enriched.csv"

    if not router_path.exists():
        raise FileNotFoundError(f"Router data not found: {router_path}")
    if not cache_path.exists():
        raise FileNotFoundError(f"Cache not found: {cache_path}")

    enrich_router_data(router_path, cache_path, llm_desc_path, output_path)
