#!/usr/bin/env python3
"""Analyze tier distribution by dataset (GSM8K, BBEH, etc.).

This script:
1. Loads GNN predictions and computes mean_logit per query
2. Maps queries back to their task_id (dataset)
3. Shows tier distribution by dataset to validate threshold calibration

Usage:
    python scripts/analyze_tier_by_dataset.py \
        --checkpoint registry/anticollapse_ablation_baseline

    # With custom thresholds
    python scripts/analyze_tier_by_dataset.py \
        --checkpoint registry/anticollapse_ablation_baseline \
        --trivial-threshold 0.1 \
        --hard-threshold -0.2
"""

from __future__ import annotations

import argparse
import logging
import pickle
import sys
from collections import defaultdict
from pathlib import Path

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

# Ensure imports work
_eb1_dir = Path(__file__).resolve().parent.parent.parent
if str(_eb1_dir) not in sys.path:
    sys.path.insert(0, str(_eb1_dir))

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


def classify_tier(
    mean_logit: float,
    trivial_threshold: float,
    hard_threshold: float,
) -> str:
    """Classify query into tier based on mean_logit."""
    if mean_logit > trivial_threshold:
        return "trivial"
    elif mean_logit > hard_threshold:
        return "easy"
    else:
        return "hard"


def load_and_analyze(
    checkpoint_path: Path,
    data_path: Path,
    trivial_threshold: float,
    hard_threshold: float,
    device: str = "cpu",
) -> dict[str, dict]:
    """Load GNN, run inference, and analyze tier distribution by dataset.

    Returns:
        Dict mapping task_id -> {tier_counts, mean_logit_stats, oracle_accuracy}
    """
    # Load training graph cache
    cache_path = checkpoint_path / "training_graph_cache.pkl"
    logger.info("Loading training graph cache from %s", cache_path)
    with open(cache_path, "rb") as f:
        cache = pickle.load(f)

    # Load config
    config_path = checkpoint_path / "config.yaml"
    with open(config_path) as f:
        raw_config = yaml.safe_load(f)

    # Extract config
    config = {
        "hidden_features": raw_config.get("hidden_dim", 384),
        "num_gnn_layers": raw_config.get("num_gat_layers", 3),
        "num_attention_heads": raw_config.get("num_gat_heads", 4),
        "dropout": raw_config.get("dropout", 0.4),
        "use_difficulty_head": raw_config.get("use_difficulty_head", True),
        "hidden_multiplier": 2,
    }

    # Load model
    model_path = checkpoint_path / "model.pth"
    state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
    state_keys = set(state_dict.keys())

    config["use_contrastive"] = any(k.startswith("contrastive_encoder.") for k in state_keys)
    config["use_uncertainty_head"] = any(k.startswith("uncertainty_proj.") for k in state_keys)

    # Infer dimensions
    if "model_align.query_transform.weight" in state_dict:
        config["hidden_features"] = state_dict["model_align.query_transform.weight"].shape[0]
        query_dim = state_dict["model_align.query_transform.weight"].shape[1]
    else:
        query_dim = 256

    if "model_align.llm_transform.weight" in state_dict:
        llm_dim = state_dict["model_align.llm_transform.weight"].shape[1]
    else:
        llm_dim = 256

    # Extract cache contents
    query_embeddings = torch.from_numpy(cache["query_embeddings"]).float()
    task_embeddings = torch.from_numpy(cache["task_embeddings"]).float()
    llm_embeddings = torch.from_numpy(cache["llm_embeddings"]).float()
    combined_edge = torch.from_numpy(cache["combined_edge"]).float()
    source_nodes = torch.from_numpy(cache["source_nodes"]).long()
    dest_nodes = torch.from_numpy(cache["dest_nodes"]).long()
    test_mask = torch.from_numpy(cache["test_mask"]).bool()

    num_queries = cache["num_queries"]
    num_llms = cache["num_llms"]
    llm_order = cache.get("llm_order", [f"LLM_{i}" for i in range(num_llms)])

    logger.info(
        "Cache: %d queries, %d LLMs, %d edges (%d test)",
        num_queries,
        num_llms,
        len(source_nodes),
        test_mask.sum().item(),
    )

    # Handle dimension alignment
    cache_query_dim = query_embeddings.shape[1]
    cache_llm_dim = llm_embeddings.shape[1]

    if cache_query_dim != query_dim:
        if cache_query_dim > query_dim:
            query_embeddings = query_embeddings[:, :query_dim]
            task_embeddings = task_embeddings[:, :query_dim]
        else:
            pad = torch.zeros(query_embeddings.shape[0], query_dim - cache_query_dim)
            query_embeddings = torch.cat([query_embeddings, pad], dim=1)
            task_pad = torch.zeros(task_embeddings.shape[0], query_dim - cache_query_dim)
            task_embeddings = torch.cat([task_embeddings, task_pad], dim=1)

    if cache_llm_dim != llm_dim:
        if cache_llm_dim > llm_dim:
            llm_embeddings = llm_embeddings[:, :llm_dim]
        else:
            pad = torch.zeros(llm_embeddings.shape[0], llm_dim - cache_llm_dim)
            llm_embeddings = torch.cat([llm_embeddings, pad], dim=1)

    # Build edge_index
    edge_index = torch.stack([
        source_nodes,
        dest_nodes + num_queries,
    ])

    # Build edge features (6D)
    n_edges = combined_edge.shape[0]
    if combined_edge.shape[1] == 3:
        w_perf = torch.ones(n_edges, 1) * 0.7
        w_cost = torch.ones(n_edges, 1) * 0.2
        w_latency = torch.ones(n_edges, 1) * 0.1
        edge_weight = torch.cat([combined_edge, w_perf, w_cost, w_latency], dim=1)
    else:
        edge_weight = combined_edge

    # Initialize model
    from embergraphrouter.model.graph_nn import GATEncoderDecoderNet

    model = GATEncoderDecoderNet(
        query_feature_dim=query_dim,
        llm_feature_dim=llm_dim,
        hidden_features=config["hidden_features"],
        in_edges=6,
        dropout=config["dropout"],
        num_gnn_layers=config["num_gnn_layers"],
        use_difficulty_head=config["use_difficulty_head"],
        use_uncertainty_head=config["use_uncertainty_head"],
        num_attention_heads=config["num_attention_heads"],
        hidden_multiplier=config["hidden_multiplier"],
        use_contrastive=config["use_contrastive"],
    )

    model.load_state_dict(state_dict, strict=False)
    model.eval()
    model.to(device)

    # Move tensors
    query_embeddings = query_embeddings.to(device)
    task_embeddings = task_embeddings.to(device)
    llm_embeddings = llm_embeddings.to(device)
    edge_index = edge_index.to(device)
    edge_weight = edge_weight.to(device)
    test_mask = test_mask.to(device)
    edge_can_see = torch.ones(n_edges, dtype=torch.bool, device=device)

    # Run inference
    logger.info("Running inference...")
    with torch.no_grad():
        edge_logits = model(
            task_features=task_embeddings,
            query_features=query_embeddings,
            llm_features=llm_embeddings,
            edge_index=edge_index,
            edge_mask=test_mask,
            edge_can_see=edge_can_see,
            edge_weight=edge_weight,
        )

    # Reshape to per-query logits matrix
    test_source_nodes = source_nodes[test_mask.cpu()].numpy()
    test_dest_nodes = dest_nodes[test_mask.cpu()].numpy()
    test_perf = combined_edge[test_mask.cpu(), 1].numpy()  # effect column

    test_query_indices = np.unique(test_source_nodes)
    num_test_queries = len(test_query_indices)

    logits_matrix = np.zeros((num_test_queries, num_llms))
    rewards_matrix = np.zeros((num_test_queries, num_llms))
    edge_logits_np = edge_logits.cpu().numpy()

    query_to_row = {q: i for i, q in enumerate(test_query_indices)}
    for i, (src, dst) in enumerate(zip(test_source_nodes, test_dest_nodes)):
        row = query_to_row[src]
        logits_matrix[row, dst] = edge_logits_np[i]
        rewards_matrix[row, dst] = test_perf[i]

    # Compute mean_logit per query
    mean_logits = logits_matrix.mean(axis=1)

    logger.info("Computed mean_logit for %d test queries", num_test_queries)
    logger.info(
        "mean_logit stats: min=%.2f, max=%.2f, mean=%.2f, std=%.2f",
        mean_logits.min(),
        mean_logits.max(),
        mean_logits.mean(),
        mean_logits.std(),
    )

    # Load original data to get task_id mapping
    logger.info("Loading original data from %s", data_path)
    df = pd.read_parquet(data_path)

    # Get embedding column suffix from config
    emb_suffix = raw_config.get("embedding_column_suffix", "qwen3_instruct_router_v1_256")
    emb_col = f"query_embedding_{emb_suffix}"

    # We need to map test query indices back to task_ids
    # The cache stores queries by index, we need to reconstruct mapping

    # Get unique queries and their task_ids
    unique_df = df.drop_duplicates(subset=["query"]).reset_index(drop=True)
    logger.info("Unique queries in data: %d", len(unique_df))

    # Match test queries to task_ids using embeddings
    # First, get test query embeddings from cache
    test_query_embs = cache["query_embeddings"][test_query_indices]

    # Parse embeddings from data
    def parse_emb(x):
        if isinstance(x, np.ndarray):
            return x
        elif isinstance(x, str):
            import ast
            return np.array(ast.literal_eval(x), dtype=np.float32)
        return np.array(x, dtype=np.float32)

    # Build mapping from embedding hash to task_id
    emb_to_task = {}
    for idx, row in unique_df.iterrows():
        try:
            emb = parse_emb(row[emb_col])
            # Use first 8 floats as a hash key
            emb_key = tuple(np.round(emb[:8], 4))
            emb_to_task[emb_key] = row["task_id"]
        except Exception:
            continue

    logger.info("Built embedding-to-task mapping for %d queries", len(emb_to_task))

    # Map test queries to task_ids
    test_task_ids = []
    for i, qidx in enumerate(test_query_indices):
        emb = cache["query_embeddings"][qidx]
        emb_key = tuple(np.round(emb[:8], 4))
        task_id = emb_to_task.get(emb_key, "unknown")
        test_task_ids.append(task_id)

    # Analyze by dataset
    results = defaultdict(lambda: {
        "trivial": 0,
        "easy": 0,
        "hard": 0,
        "mean_logits": [],
        "gnn_correct": 0,
        "total": 0,
    })

    for i, (task_id, mean_logit) in enumerate(zip(test_task_ids, mean_logits)):
        tier = classify_tier(mean_logit, trivial_threshold, hard_threshold)
        results[task_id][tier] += 1
        results[task_id]["mean_logits"].append(mean_logit)
        results[task_id]["total"] += 1

        # Check if GNN argmax matches oracle
        gnn_choice = logits_matrix[i].argmax()
        oracle_choice = rewards_matrix[i].argmax()
        if gnn_choice == oracle_choice:
            results[task_id]["gnn_correct"] += 1

    # Compute stats
    for task_id, stats in results.items():
        mean_logits_arr = np.array(stats["mean_logits"])
        stats["mean_logit_avg"] = float(mean_logits_arr.mean())
        stats["mean_logit_std"] = float(mean_logits_arr.std())
        stats["gnn_accuracy"] = stats["gnn_correct"] / stats["total"] if stats["total"] > 0 else 0
        del stats["mean_logits"]

    return dict(results)


def print_analysis(
    results: dict[str, dict],
    trivial_threshold: float,
    hard_threshold: float,
) -> None:
    """Print tier analysis by dataset."""
    print("\n" + "=" * 90)
    print("TIER DISTRIBUTION BY DATASET")
    print(f"Thresholds: trivial > {trivial_threshold}, easy > {hard_threshold}, else hard")
    print("=" * 90)

    # Group by dataset category
    gsm8k = {k: v for k, v in results.items() if "GSM8K" in k.upper() or "gsm" in k.lower()}
    bbeh = {k: v for k, v in results.items() if "bbeh" in k.lower()}
    mmlu = {k: v for k, v in results.items() if "mmlu" in k.lower()}
    other = {k: v for k, v in results.items() if k not in gsm8k and k not in bbeh and k not in mmlu}

    def print_group(name: str, group: dict):
        if not group:
            return

        print(f"\n{'─' * 90}")
        print(f" {name}")
        print(f"{'─' * 90}")
        print(
            f"{'Dataset':<35} {'Total':>7} {'Trivial':>10} {'Easy':>10} {'Hard':>10} "
            f"{'mean_logit':>12} {'GNN Acc':>10}"
        )
        print("-" * 90)

        total_trivial = total_easy = total_hard = total_n = 0
        for task_id, stats in sorted(group.items()):
            n = stats["total"]
            triv = stats["trivial"]
            easy = stats["easy"]
            hard = stats["hard"]
            avg = stats["mean_logit_avg"]
            acc = stats["gnn_accuracy"]

            total_n += n
            total_trivial += triv
            total_easy += easy
            total_hard += hard

            print(
                f"{task_id:<35} {n:>7} {triv:>6} ({triv/n:.0%}) {easy:>6} ({easy/n:.0%}) "
                f"{hard:>6} ({hard/n:.0%}) {avg:>+8.2f}     {acc:>6.1%}"
            )

        if len(group) > 1:
            print("-" * 90)
            print(
                f"{'SUBTOTAL':<35} {total_n:>7} "
                f"{total_trivial:>6} ({total_trivial/total_n:.0%}) "
                f"{total_easy:>6} ({total_easy/total_n:.0%}) "
                f"{total_hard:>6} ({total_hard/total_n:.0%})"
            )

    print_group("GSM8K (Math)", gsm8k)
    print_group("BBEH (BIG-Bench Extra Hard)", bbeh)
    print_group("MMLU & MMLU-Pro", mmlu)
    print_group("Other Datasets", other)

    # Overall summary
    print("\n" + "=" * 90)
    print("OVERALL SUMMARY")
    print("=" * 90)

    total_n = sum(s["total"] for s in results.values())
    total_triv = sum(s["trivial"] for s in results.values())
    total_easy = sum(s["easy"] for s in results.values())
    total_hard = sum(s["hard"] for s in results.values())

    print(f"Total queries: {total_n}")
    print(f"  Trivial: {total_triv} ({total_triv/total_n:.1%}) - GNN skipped, use cheap model")
    print(f"  Easy:    {total_easy} ({total_easy/total_n:.1%}) - GNN routes directly")
    print(f"  Hard:    {total_hard} ({total_hard/total_n:.1%}) - Conductor orchestrates")

    # Recommendations
    print("\n" + "=" * 90)
    print("THRESHOLD RECOMMENDATIONS")
    print("=" * 90)

    # Check if GSM8K and easy datasets have high trivial rates
    easy_datasets = ["mmlu", "arc", "truthfulqa", "hellaswag"]
    hard_datasets = ["bbeh", "agieval"]

    easy_trivial_rate = 0
    easy_count = 0
    hard_trivial_rate = 0
    hard_count = 0

    for task_id, stats in results.items():
        is_easy = any(e in task_id.lower() for e in easy_datasets)
        is_hard = any(h in task_id.lower() for h in hard_datasets)

        if is_easy:
            easy_trivial_rate += stats["trivial"]
            easy_count += stats["total"]
        elif is_hard:
            hard_trivial_rate += stats["trivial"]
            hard_count += stats["total"]

    if easy_count > 0:
        print(f"Easy datasets (MMLU, ARC, etc.): {easy_trivial_rate/easy_count:.1%} trivial")
    if hard_count > 0:
        print(f"Hard datasets (BBEH, AGIEval): {hard_trivial_rate/hard_count:.1%} trivial")

    if easy_count > 0 and easy_trivial_rate / easy_count < 0.2:
        print("\n⚠️  Consider LOWERING trivial threshold to capture more easy queries")
        print(f"   Current: {trivial_threshold}, Try: {trivial_threshold - 0.1:.1f}")
    elif hard_count > 0 and hard_trivial_rate / hard_count > 0.05:
        print("\n⚠️  Consider RAISING trivial threshold to avoid misrouting hard queries")
        print(f"   Current: {trivial_threshold}, Try: {trivial_threshold + 0.1:.1f}")


def main():
    parser = argparse.ArgumentParser(description="Analyze tier distribution by dataset")
    parser.add_argument(
        "--checkpoint",
        type=str,
        default="embergraphrouter/registry/anticollapse_ablation_baseline",
        help="Path to GNN checkpoint",
    )
    parser.add_argument(
        "--data",
        type=str,
        default="embergraphrouter/data/router_data_qwen3_instruct_router_v1.parquet",
        help="Path to router data",
    )
    parser.add_argument(
        "--trivial-threshold",
        type=float,
        default=0.3,
        help="Mean logit threshold for trivial tier",
    )
    parser.add_argument(
        "--hard-threshold",
        type=float,
        default=-0.1,
        help="Mean logit threshold for hard tier",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cpu",
        help="Device for inference",
    )

    args = parser.parse_args()

    results = load_and_analyze(
        checkpoint_path=Path(args.checkpoint),
        data_path=Path(args.data),
        trivial_threshold=args.trivial_threshold,
        hard_threshold=args.hard_threshold,
        device=args.device,
    )

    print_analysis(results, args.trivial_threshold, args.hard_threshold)


if __name__ == "__main__":
    main()
