#!/usr/bin/env python3
"""Run threshold calibration using GNN predictions and ground truth.

This script:
1. Loads GNN model predictions from a registry checkpoint
2. Extracts difficulty signals (mean_logit, max_prob, logit_margin, etc.)
3. Combines with ground truth effects from training cache
4. Runs threshold calibration to find optimal tier boundaries
5. Generates a deployment-ready config

Usage:
    python scripts/run_threshold_calibration.py --checkpoint registry/eb1-gnn-preview-12-19
    python scripts/run_threshold_calibration.py --checkpoint registry/eb1-gnn-preview-12-19 --strategy bayesian
    python scripts/run_threshold_calibration.py --checkpoint registry/eb1-gnn-preview-12-19 --iterations 100
"""

from __future__ import annotations

import argparse
import json
import logging
import pickle
import sys
from dataclasses import dataclass
from pathlib import Path

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

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

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


@dataclass
class GNNSignals:
    """GNN difficulty signals per query."""

    mean_logit: np.ndarray  # (N,) mean of logits
    max_prob: np.ndarray  # (N,) max probability
    logit_margin: np.ndarray  # (N,) difference between top1 and top2 logit
    logit_std: np.ndarray  # (N,) std of logits
    entropy: np.ndarray  # (N,) normalized entropy
    top1_idx: np.ndarray  # (N,) index of top1 model
    top1_prob: np.ndarray  # (N,) probability of top1 model


def extract_signals_from_logits(
    logits: np.ndarray,
    n_llms: int,
) -> GNNSignals:
    """Extract difficulty signals from raw logits.

    Args:
        logits: (N, n_llms) raw model logits
        n_llms: number of LLMs

    Returns:
        GNNSignals dataclass with per-query signals
    """
    # Probabilities via softmax
    probs = F.softmax(torch.from_numpy(logits).float(), dim=-1).numpy()

    # Mean logit (key difficulty signal)
    mean_logit = logits.mean(axis=1)

    # Max probability and top1
    max_prob = probs.max(axis=1)
    top1_idx = probs.argmax(axis=1)
    top1_prob = max_prob

    # Logit margin (top1 - top2)
    sorted_logits = np.sort(logits, axis=1)[:, ::-1]
    logit_margin = sorted_logits[:, 0] - sorted_logits[:, 1]

    # Logit std
    logit_std = logits.std(axis=1)

    # Normalized entropy
    entropy = -np.sum(probs * np.log(probs + 1e-8), axis=1)
    max_entropy = np.log(n_llms)
    normalized_entropy = entropy / max_entropy

    return GNNSignals(
        mean_logit=mean_logit,
        max_prob=max_prob,
        logit_margin=logit_margin,
        logit_std=logit_std,
        entropy=normalized_entropy,
        top1_idx=top1_idx,
        top1_prob=top1_prob,
    )


def load_model_and_predict(
    checkpoint_path: Path,
    device: str = "cpu",
) -> tuple[GNNSignals, np.ndarray, list[str], np.ndarray]:
    """Load GNN model and generate predictions.

    Uses the same approach as eval_difficulty_aware.py for consistency.

    Args:
        checkpoint_path: Path to registry checkpoint
        device: Device for inference

    Returns:
        Tuple of (signals, rewards_matrix, llm_names, test_indices)
    """
    # Load training graph cache
    cache_path = checkpoint_path / "training_graph_cache.pkl"
    if not cache_path.exists():
        raise FileNotFoundError(f"Training graph cache not found: {cache_path}")

    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"
    if config_path.exists():
        with open(config_path) as f:
            raw_config = yaml.safe_load(f)
        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,
        }
    else:
        config = {
            "hidden_features": 384,
            "num_gnn_layers": 3,
            "num_attention_heads": 4,
            "dropout": 0.4,
            "use_difficulty_head": True,
            "hidden_multiplier": 2,
        }

    # Load model weights and detect architecture
    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 from state_dict
    if "model_align.query_transform.weight" in state_dict:
        config["hidden_features"] = state_dict["model_align.query_transform.weight"].shape[0]
        inferred_query_dim = state_dict["model_align.query_transform.weight"].shape[1]
        config["inferred_query_dim"] = inferred_query_dim

    if "model_align.llm_transform.weight" in state_dict:
        inferred_llm_dim = state_dict["model_align.llm_transform.weight"].shape[1]
        config["inferred_llm_dim"] = inferred_llm_dim

    logger.info(
        "Model config: hidden=%d, layers=%d, heads=%d, contrastive=%s, uncertainty=%s",
        config["hidden_features"],
        config["num_gnn_layers"],
        config["num_attention_heads"],
        config["use_contrastive"],
        config["use_uncertainty_head"],
    )

    # 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(
        "Loaded cache: %d queries, %d LLMs, %d edges (%d test)",
        num_queries,
        num_llms,
        len(source_nodes),
        test_mask.sum().item(),
    )

    # Build edge_index: source is query node (0..num_queries-1),
    # dest is LLM node (num_queries..num_queries+num_llms-1)
    edge_index = torch.stack([
        source_nodes,
        dest_nodes + num_queries,  # Offset LLM indices
    ])

    # Build edge features: combined_edge has (perf, cost, latency)
    # Model expects 6D: (perf, cost, latency, w_perf, w_cost, w_latency)
    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

    cache_query_dim = query_embeddings.shape[1]
    cache_llm_dim = llm_embeddings.shape[1]
    query_dim = config.get("inferred_query_dim", cache_query_dim)
    llm_dim = config.get("inferred_llm_dim", cache_llm_dim)

    # Handle dimension mismatch
    if cache_query_dim != query_dim:
        logger.warning(
            "Cache query_dim=%d != checkpoint query_dim=%d, adjusting embeddings",
            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:
        logger.warning(
            "Cache llm_dim=%d != checkpoint llm_dim=%d, adjusting embeddings",
            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)

    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"],
    )

    # Load weights
    logger.info("Loading model weights from %s", model_path)
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    if unexpected:
        logger.info("Ignored deprecated weights: %s", unexpected)
    model.eval()
    model.to(device)

    # Move tensors to device
    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)

    # For inference, edge_can_see = all edges (full graph visibility)
    edge_can_see = torch.ones(n_edges, dtype=torch.bool, device=device)

    # Run inference
    logger.info("Running inference on test set...")
    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 edge_logits from (num_test_edges,) to (num_test_queries, num_llms)
    test_source_nodes = source_nodes[test_mask.cpu()].numpy()
    test_dest_nodes = dest_nodes[test_mask.cpu()].numpy()

    # Auto-detect effect column: the effect column has high variance and mean near 0.5
    # Different checkpoints may use column 0 or column 1 for effects
    combined_edge_np = combined_edge.numpy()
    col0_std = combined_edge_np[:, 0].std()
    col1_std = combined_edge_np[:, 1].std()
    effect_col = 0 if col0_std > col1_std else 1
    logger.info(
        "Auto-detected effect column: %d (col0_std=%.4f, col1_std=%.4f)",
        effect_col, col0_std, col1_std
    )
    test_perf = combined_edge[test_mask.cpu(), effect_col].numpy()

    # Get unique test queries and their order
    test_query_indices = np.unique(test_source_nodes)
    num_test_queries = len(test_query_indices)

    logger.info(
        "Reshaping %d test edges to (%d queries, %d LLMs)",
        edge_logits.shape[0],
        num_test_queries,
        num_llms,
    )

    # Build per-query logit matrix and rewards matrix
    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()

    # Map test edges to matrix positions
    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]
        col = dst  # dest_nodes are already 0-indexed LLM indices
        logits_matrix[row, col] = edge_logits_np[i]
        rewards_matrix[row, col] = test_perf[i]

    logger.info(
        "Inference complete: logits=%s, rewards=%s",
        logits_matrix.shape,
        rewards_matrix.shape,
    )

    # Extract signals
    signals = extract_signals_from_logits(logits_matrix, num_llms)

    return signals, rewards_matrix, llm_order, test_query_indices


def build_calibration_dataset(
    signals: GNNSignals,
    rewards_matrix: np.ndarray,
    llm_names: list[str],
    test_indices: np.ndarray,
) -> "OfflineDataset":
    """Build calibration dataset from signals and rewards.

    Args:
        signals: GNN difficulty signals
        rewards_matrix: (N, n_llms) ground truth effects
        llm_names: List of LLM names
        test_indices: Indices of test queries

    Returns:
        OfflineDataset for calibration
    """
    from embergraphrouter.calibration.dataset import CalibrationSample, OfflineDataset

    samples = []
    for i in range(len(test_indices)):
        # Build ground truth effects dict
        effects = {name: float(rewards_matrix[i, j]) for j, name in enumerate(llm_names)}

        # Find oracle
        oracle_idx = rewards_matrix[i].argmax()
        oracle_model = llm_names[oracle_idx]
        oracle_effect = float(rewards_matrix[i, oracle_idx])

        # Top1 from GNN
        top1_idx = int(signals.top1_idx[i])
        top1_model = llm_names[top1_idx]

        sample = CalibrationSample(
            query_id=str(test_indices[i]),
            query_text=f"query_{test_indices[i]}",
            task_id="eval",
            mean_logit=float(signals.mean_logit[i]),
            max_prob=float(signals.max_prob[i]),
            logit_margin=float(signals.logit_margin[i]),
            logit_std=float(signals.logit_std[i]),
            entropy=float(signals.entropy[i]),
            top1_model=top1_model,
            top1_prob=float(signals.top1_prob[i]),
            ground_truth_effects=effects,
            oracle_model=oracle_model,
            oracle_effect=oracle_effect,
        )
        samples.append(sample)

    logger.info("Built calibration dataset with %d samples", len(samples))

    return OfflineDataset(
        samples=samples,
        model_names=llm_names,
        task_ids=["eval"],
        metadata={"source": "gnn_predictions"},
    )


def run_calibration(
    dataset: "OfflineDataset",
    strategy: str = "bayesian",
    n_iterations: int = 50,
    optimization_target: str = "balanced",
) -> "CalibrationResult":
    """Run threshold calibration.

    Args:
        dataset: Offline calibration dataset
        strategy: Search strategy ("grid", "random", "bayesian")
        n_iterations: Number of search iterations
        optimization_target: Optimization target ("quality", "cost", "balanced", "pareto")

    Returns:
        CalibrationResult with optimal thresholds
    """
    from embergraphrouter.calibration.config import CalibrationConfig, OptimizationTarget
    from embergraphrouter.calibration.optimizer import ThresholdCalibrator

    target_map = {
        "quality": OptimizationTarget.QUALITY,
        "cost": OptimizationTarget.COST,
        "latency": OptimizationTarget.LATENCY,
        "balanced": OptimizationTarget.BALANCED,
        "pareto": OptimizationTarget.PARETO,
    }

    config = CalibrationConfig(
        optimization_target=target_map.get(optimization_target, OptimizationTarget.BALANCED),
        n_search_iterations=n_iterations,
        n_cv_folds=3,
        quality_weight=0.5,
        cost_weight=0.3,
        latency_weight=0.2,
    )

    calibrator = ThresholdCalibrator(config, search_strategy=strategy)
    result = calibrator.calibrate(dataset)

    return result


def print_results(result: "CalibrationResult") -> None:
    """Print calibration results."""
    from embergraphrouter.calibration.report import CalibrationReport

    report = CalibrationReport.from_result(result)
    print(report.to_text())

    # Print YAML config for deployment
    print("\n" + "=" * 70)
    print("DEPLOYMENT CONFIG (copy to config.yaml)")
    print("=" * 70)
    print(report.to_yaml_config())


def run_budget_optimization(
    dataset: "OfflineDataset",
    output_path: Path | None = None,
) -> None:
    """Run budget-constrained optimization and save Pareto frontier.

    Args:
        dataset: Calibration dataset
        output_path: Path to save budget cache (defaults to registry)
    """
    from embergraphrouter.calibration.budget_optimizer import BudgetOptimizer, TierCosts

    # Configure tier costs (these should match your production costs)
    tier_costs = TierCosts(
        trivial_cost=0.0001,  # $0.10 per 1K (cheap model like Flash)
        discriminative_cost=0.001,  # $1.00 per 1K (Pro model)
        uncertain_cost=0.01,  # $10.00 per 1K (EB1 ensemble)
        trivial_latency_ms=100.0,
        discriminative_latency_ms=200.0,
        uncertain_latency_ms=1500.0,
    )

    logger.info("Running budget-constrained optimization...")
    optimizer = BudgetOptimizer.from_calibration_data(dataset, tier_costs)

    # Compute Pareto frontier
    frontier = optimizer.compute_pareto_frontier(n_points=50, grid_resolution=15)

    print("\n" + "=" * 70)
    print("PARETO FRONTIER (Quality vs Cost)")
    print("=" * 70)
    print(f"{'Budget ($/1K)':<15} {'Quality':<10} {'Latency (ms)':<15} {'Tier Distribution'}")
    print("-" * 70)
    for p in frontier.points[:10]:  # Show top 10 points
        dist_str = f"T:{p.tier_distribution['TRIVIAL']:.0%} D:{p.tier_distribution['DISCRIMINATIVE']:.0%} U:{p.tier_distribution['UNCERTAIN']:.0%}"
        print(f"${p.cost_per_1k:<14.2f} {p.quality:<10.1%} {p.latency_ms:<15.0f} {dist_str}")

    # Show optimal for different budgets
    print("\n" + "=" * 70)
    print("OPTIMAL THRESHOLDS BY BUDGET")
    print("=" * 70)
    for budget in [1.0, 2.0, 5.0, 10.0]:
        thresholds = frontier.get_thresholds_for_budget(budget)
        if thresholds:
            result = optimizer.optimize_for_budget(cost_budget_per_1k=budget)
            print(f"\nBudget: ${budget}/1K → Quality: {result.quality:.1%}")
            print(f"  trivial_mean_logit: {thresholds.trivial_mean_logit:.3f}")
            print(f"  uncertain_mean_logit: {thresholds.uncertain_mean_logit:.3f}")
            print(f"  uncertain_logit_margin: {thresholds.uncertain_logit_margin:.3f}")

    # Save cache
    if output_path is None:
        output_path = Path("embergraphrouter/registry/eb1-gnn-preview-12-19/budget_cache.json")
    optimizer.save_optimization_cache(output_path)
    print(f"\nSaved budget optimization cache to: {output_path}")


def main():
    parser = argparse.ArgumentParser(description="Run threshold calibration")
    parser.add_argument(
        "--checkpoint",
        type=str,
        default="embergraphrouter/registry/eb1-gnn-preview-12-19",
        help="Path to registry checkpoint",
    )
    parser.add_argument(
        "--strategy",
        type=str,
        default="bayesian",
        choices=["grid", "random", "bayesian"],
        help="Search strategy",
    )
    parser.add_argument(
        "--iterations",
        type=int,
        default=50,
        help="Number of search iterations",
    )
    parser.add_argument(
        "--target",
        type=str,
        default="balanced",
        choices=["quality", "cost", "latency", "balanced", "pareto"],
        help="Optimization target",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cpu",
        help="Device for inference",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output path for calibration results (JSON)",
    )
    parser.add_argument(
        "--budget-optimize",
        action="store_true",
        help="Run budget-constrained optimization and save Pareto frontier",
    )
    parser.add_argument(
        "--budget-cache",
        type=str,
        default=None,
        help="Output path for budget optimization cache (JSON)",
    )

    args = parser.parse_args()

    checkpoint_path = Path(args.checkpoint)
    if not checkpoint_path.exists():
        logger.error("Checkpoint not found: %s", checkpoint_path)
        sys.exit(1)

    # Load model and predict
    logger.info("Loading model from %s", checkpoint_path)
    signals, rewards, llm_names, test_indices = load_model_and_predict(
        checkpoint_path, device=args.device
    )

    # Build calibration dataset
    dataset = build_calibration_dataset(signals, rewards, llm_names, test_indices)

    # Print dataset summary
    summary = dataset.summary()
    print("\n" + "=" * 70)
    print("DATASET SUMMARY")
    print("=" * 70)
    print(f"Samples: {summary['n_samples']}")
    print(f"Models: {summary['n_models']}")
    print("Signal stats:")
    for sig, stats in summary["signal_stats"].items():
        print(f"  {sig}: mean={stats['mean']:.3f}, std={stats['std']:.3f}")
    print("Difficulty distribution:")
    for diff, pct in summary["difficulty_distribution"].items():
        print(f"  {diff}: {pct:.1f}%")
    print()

    # Run calibration
    logger.info(
        "Running calibration with %s strategy, %d iterations",
        args.strategy,
        args.iterations,
    )
    result = run_calibration(
        dataset,
        strategy=args.strategy,
        n_iterations=args.iterations,
        optimization_target=args.target,
    )

    # Print results
    print_results(result)

    # Save results
    if args.output:
        output_path = Path(args.output)
        with open(output_path, "w") as f:
            json.dump(result.to_dict(), f, indent=2)
        logger.info("Results saved to %s", output_path)

    # Run budget optimization if requested
    if args.budget_optimize:
        budget_cache_path = Path(args.budget_cache) if args.budget_cache else None
        run_budget_optimization(dataset, budget_cache_path)


if __name__ == "__main__":
    main()
