#!/usr/bin/env python3
"""Run baseline comparison experiment.

This script trains and evaluates all three router architectures on the same
data split, providing a fair head-to-head comparison.

Usage:
    python -m embergraphrouter.model.baselines.run_comparison \
        --data_path data/router_data.csv \
        --output_dir results/baseline_comparison

    # Or with specific configs:
    python -m embergraphrouter.model.baselines.run_comparison \
        --two_tower_config configs/baselines/two_tower.yaml \
        --mlp_config configs/baselines/multi_head_mlp.yaml \
        --gnn_config configs/baselines/gnn.yaml
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from datetime import datetime
from pathlib import Path
from typing import Optional

import numpy as np
import torch
import yaml

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


def load_config(config_path: str) -> dict:
    """Load YAML configuration file."""
    with open(config_path) as f:
        return yaml.safe_load(f)


def create_synthetic_data(
    num_samples: int = 1000,
    query_dim: int = 512,
    model_dim: int = 512,
    num_models: int = 7,
    seed: int = 42,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Create synthetic data for testing.

    This generates plausible routing data where different query types
    have different optimal models.

    Args:
        num_samples: Number of query samples.
        query_dim: Query embedding dimension.
        model_dim: Model embedding dimension.
        num_models: Number of models.
        seed: Random seed.

    Returns:
        Tuple of (query_embs, model_embs, targets).
    """
    torch.manual_seed(seed)
    np.random.seed(seed)

    # Generate query embeddings with some cluster structure
    num_clusters = 5
    cluster_centers = torch.randn(num_clusters, query_dim)
    cluster_assignments = torch.randint(0, num_clusters, (num_samples,))

    query_embs = torch.zeros(num_samples, query_dim)
    for i in range(num_samples):
        query_embs[i] = cluster_centers[cluster_assignments[i]] + torch.randn(query_dim) * 0.3

    # Generate model embeddings
    model_embs = torch.randn(num_models, model_dim)

    # Generate targets based on query-model affinity
    # Each cluster has a preferred model
    cluster_model_affinity = torch.softmax(torch.randn(num_clusters, num_models), dim=-1)

    targets = torch.zeros(num_samples, num_models)
    for i in range(num_samples):
        cluster = cluster_assignments[i]
        # Base affinity from cluster
        base_affinity = cluster_model_affinity[cluster]
        # Add noise
        noise = torch.randn(num_models) * 0.1
        targets[i] = torch.sigmoid(base_affinity + noise)

    return query_embs, model_embs, targets


def run_comparison(
    query_embs: torch.Tensor,
    model_embs: torch.Tensor,
    targets: torch.Tensor,
    config: dict,
    output_dir: Path,
) -> dict:
    """Run the baseline comparison.

    Args:
        query_embs: Query embeddings.
        model_embs: Model embeddings.
        targets: Target scores.
        config: Configuration dictionary.
        output_dir: Directory for outputs.

    Returns:
        Dictionary of results.
    """
    from .base import RouterConfig
    from .evaluation import BaselineEvaluator
    from .gnn_baseline import SimplifiedGNNRouter
    from .multi_head_mlp import MultiHeadMLPRouter
    from .trainer import BaselineTrainer, RouterDataset, TrainerConfig
    from .two_tower import TwoTowerRouter

    # Split data
    num_samples = query_embs.shape[0]
    train_size = int(num_samples * 0.7)
    val_size = int(num_samples * 0.15)

    indices = torch.randperm(num_samples)
    train_idx = indices[:train_size]
    val_idx = indices[train_size:train_size + val_size]
    test_idx = indices[train_size + val_size:]

    train_data = RouterDataset(
        query_embs[train_idx],
        model_embs,
        targets[train_idx],
    )
    val_data = RouterDataset(
        query_embs[val_idx],
        model_embs,
        targets[val_idx],
    )
    test_data = RouterDataset(
        query_embs[test_idx],
        model_embs,
        targets[test_idx],
    )

    logger.info(f"Data split: train={len(train_data)}, val={len(val_data)}, test={len(test_data)}")

    # Router configuration
    router_config = RouterConfig(
        query_dim=query_embs.shape[1],
        model_dim=model_embs.shape[1],
        hidden_dim=config.get("hidden_dim", 384),
        num_models=model_embs.shape[0],
        num_layers=config.get("num_layers", 2),
        num_heads=config.get("num_heads", 4),
        dropout=config.get("dropout", 0.1),
    )

    # Trainer configuration
    trainer_config = TrainerConfig(
        batch_size=config.get("batch_size", 32),
        learning_rate=config.get("learning_rate", 1e-3),
        num_epochs=config.get("num_epochs", 50),
        early_stopping_patience=config.get("early_stopping_patience", 10),
        loss_type=config.get("loss_type", "hybrid"),
        use_wandb=config.get("use_wandb", False),
        checkpoint_dir=str(output_dir / "checkpoints"),
    )

    # Initialize routers
    routers = [
        ("Two Tower", TwoTowerRouter(router_config)),
        ("Multi-Head MLP", MultiHeadMLPRouter(router_config)),
        ("GNN (Simplified)", SimplifiedGNNRouter(router_config)),
    ]

    # Train each router
    trained_routers = []
    training_results = {}

    for name, router in routers:
        logger.info(f"\n{'='*60}")
        logger.info(f"Training {name}...")
        logger.info(f"Parameters: {router.count_parameters():,}")
        logger.info(f"{'='*60}")

        trainer = BaselineTrainer(router, trainer_config)
        metrics = trainer.train(train_data, val_data)

        training_results[name] = {
            "best_epoch": metrics.best_epoch,
            "best_frac_oracle": metrics.best_metrics().get("frac_oracle", 0),
            "training_time": metrics.training_time_seconds,
            "train_loss_history": metrics.train_loss,
            "val_loss_history": metrics.val_loss,
        }

        logger.info(f"Best epoch: {metrics.best_epoch}")
        logger.info(f"Best frac_oracle: {training_results[name]['best_frac_oracle']:.4f}")
        logger.info(f"Training time: {metrics.training_time_seconds:.1f}s")

        trained_routers.append((name, router))

    # Evaluate on test set
    logger.info(f"\n{'='*60}")
    logger.info("Evaluating on test set...")
    logger.info(f"{'='*60}")

    evaluator = BaselineEvaluator(test_data)
    eval_results = evaluator.compare(trained_routers)

    # Print comparison table
    logger.info("\n" + "=" * 70)
    logger.info("BASELINE COMPARISON RESULTS")
    logger.info("=" * 70)
    evaluator.print_comparison_table(eval_results)

    # Compile final results
    results = {
        "timestamp": datetime.now().isoformat(),
        "config": config,
        "data": {
            "num_samples": num_samples,
            "query_dim": query_embs.shape[1],
            "model_dim": model_embs.shape[1],
            "num_models": model_embs.shape[0],
            "train_size": len(train_data),
            "val_size": len(val_data),
            "test_size": len(test_data),
        },
        "training": training_results,
        "evaluation": {r.name: r.to_dict() for r in eval_results},
    }

    # Save results
    output_dir.mkdir(parents=True, exist_ok=True)
    results_path = output_dir / "comparison_results.json"
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2, default=str)

    logger.info(f"\nResults saved to {results_path}")

    return results


def main() -> None:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Run router baseline comparison",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        "--data_path",
        type=str,
        default=None,
        help="Path to router data CSV (default: use synthetic data)",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="results/baseline_comparison",
        help="Output directory for results",
    )
    parser.add_argument(
        "--config",
        type=str,
        default=None,
        help="Path to YAML configuration file",
    )
    parser.add_argument(
        "--num_samples",
        type=int,
        default=2000,
        help="Number of synthetic samples (if not using real data)",
    )
    parser.add_argument(
        "--num_epochs",
        type=int,
        default=50,
        help="Maximum training epochs",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed",
    )
    parser.add_argument(
        "--use_wandb",
        action="store_true",
        help="Enable Weights & Biases logging",
    )

    args = parser.parse_args()

    # Load or create config
    if args.config:
        config = load_config(args.config)
    else:
        config = {
            "hidden_dim": 384,
            "num_layers": 2,
            "num_heads": 4,
            "dropout": 0.1,
            "batch_size": 32,
            "learning_rate": 1e-3,
            "num_epochs": args.num_epochs,
            "early_stopping_patience": 10,
            "loss_type": "hybrid",
            "use_wandb": args.use_wandb,
            "seed": args.seed,
        }

    # Set seeds
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    # Load or generate data
    if args.data_path:
        logger.info(f"Loading data from {args.data_path}...")
        # TODO: Implement actual data loading
        raise NotImplementedError("Real data loading not yet implemented")
    else:
        logger.info(f"Generating synthetic data ({args.num_samples} samples)...")
        query_embs, model_embs, targets = create_synthetic_data(
            num_samples=args.num_samples,
            seed=args.seed,
        )

    # Run comparison
    output_dir = Path(args.output_dir)
    results = run_comparison(
        query_embs,
        model_embs,
        targets,
        config,
        output_dir,
    )

    # Print summary
    logger.info("\n" + "=" * 70)
    logger.info("SUMMARY")
    logger.info("=" * 70)

    for name, eval_result in results["evaluation"].items():
        logger.info(
            f"{name}: frac_oracle={eval_result['frac_oracle']:.4f}, "
            f"top1_acc={eval_result['top1_accuracy']:.4f}, "
            f"params={eval_result['parameter_count']:,}"
        )


if __name__ == "__main__":
    main()
