#!/usr/bin/env python3
"""Train baseline routers on actual router_data.

This script loads the router_data parquet file and trains all three
baseline architectures for fair comparison.

Usage:
    python -m embergraphrouter.model.baselines.train_on_router_data \
        --data_path embergraphrouter/data/router_data_8llm_qwen3.parquet \
        --embedding_dim 512 \
        --num_epochs 100
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset, random_split

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


class RouterParquetDataset(Dataset):
    """Dataset that loads router data from parquet file.

    Expects data format where each row is (query, llm) pair with:
    - query_embedding_{dim}: Query embedding as list/array
    - task_description_embedding_{dim}: Task embedding (optional)
    - effect: Ground truth performance (0-1)
    - llm: Model name
    """

    def __init__(
        self,
        parquet_path: str,
        embedding_dim: int = 512,
        embedding_prefix: str = "query_embedding_qwen3_instruct_router_v1",
        use_task_embedding: bool = True,
    ) -> None:
        """Initialize dataset from parquet file.

        Args:
            parquet_path: Path to parquet file.
            embedding_dim: Dimension of embeddings to use (256, 512, or 1024).
            embedding_prefix: Prefix for embedding column names.
            use_task_embedding: Whether to concatenate task embeddings.
        """
        logger.info(f"Loading data from {parquet_path}...")
        df = pd.read_parquet(parquet_path)

        # Get unique LLMs and create mapping
        self.llm_names = sorted(df["llm"].unique().tolist())
        self.llm_to_idx = {name: i for i, name in enumerate(self.llm_names)}
        self.num_models = len(self.llm_names)
        logger.info(f"Found {self.num_models} LLMs: {self.llm_names}")

        # Get embedding column names
        query_emb_col = f"{embedding_prefix}_{embedding_dim}"
        task_emb_col = f"task_description_embedding_qwen3_instruct_router_v1_{embedding_dim}"

        if query_emb_col not in df.columns:
            raise ValueError(f"Embedding column {query_emb_col} not found. Available: {df.columns.tolist()}")

        # Group by query to get unique queries
        # Use (task_id, query) as unique identifier
        logger.info("Grouping by query...")
        grouped = df.groupby(["task_id", "query"])

        # Extract embeddings and targets
        query_embs = []
        task_embs = []
        targets = []

        def parse_embedding(emb_val):
            """Parse embedding from string or array."""
            if isinstance(emb_val, str):
                import ast
                return np.array(ast.literal_eval(emb_val), dtype=np.float32)
            elif isinstance(emb_val, (list, np.ndarray)):
                return np.array(emb_val, dtype=np.float32)
            else:
                raise ValueError(f"Unknown embedding type: {type(emb_val)}")

        for (task_id, query), group in grouped:
            if len(group) != self.num_models:
                continue  # Skip incomplete groups

            # Get query embedding (same for all rows in group)
            emb = parse_embedding(group[query_emb_col].iloc[0])
            query_embs.append(emb)

            # Get task embedding if available
            if use_task_embedding and task_emb_col in df.columns:
                task_emb = parse_embedding(group[task_emb_col].iloc[0])
                task_embs.append(task_emb)

            # Get effect for each LLM
            target = np.zeros(self.num_models)
            for _, row in group.iterrows():
                llm_idx = self.llm_to_idx[row["llm"]]
                target[llm_idx] = row["effect"]
            targets.append(target)

        self.query_embs = torch.tensor(np.array(query_embs), dtype=torch.float32)
        self.targets = torch.tensor(np.array(targets), dtype=torch.float32)

        # Concatenate task embeddings if used
        if task_embs:
            task_embs_tensor = torch.tensor(np.array(task_embs), dtype=torch.float32)
            self.query_embs = torch.cat([self.query_embs, task_embs_tensor], dim=-1)
            self.embedding_dim = self.query_embs.shape[-1]
        else:
            self.embedding_dim = embedding_dim

        # Create model embeddings from average query-model affinities
        # This is a simple approach; could also use LLM description embeddings
        self.model_embs = self._compute_model_embeddings()

        logger.info(f"Loaded {len(self)} queries with embedding dim {self.embedding_dim}")
        logger.info(f"Target shape: {self.targets.shape}")
        logger.info(f"Model embeddings shape: {self.model_embs.shape}")

    def _compute_model_embeddings(self) -> torch.Tensor:
        """Compute model embeddings from data statistics.

        For each model, compute embedding as weighted average of queries
        where that model performs well.
        """
        # Simple approach: use random embeddings of same dimension as queries
        # In production, would use LLM description embeddings
        torch.manual_seed(42)
        model_embs = torch.randn(self.num_models, self.embedding_dim)
        return F.normalize(model_embs, dim=-1)

    def __len__(self) -> int:
        return len(self.query_embs)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        return self.query_embs[idx], self.model_embs, self.targets[idx]


def collate_fn(batch):
    """Collate function for DataLoader."""
    query_embs = torch.stack([x[0] for x in batch])
    model_embs = batch[0][1]
    targets = torch.stack([x[2] for x in batch])
    return query_embs, model_embs, targets


def train_baseline(
    name: str,
    router,
    train_loader: DataLoader,
    val_loader: DataLoader,
    config: dict,
    device: torch.device,
) -> dict:
    """Train a single baseline router.

    Args:
        name: Router name for logging.
        router: Router model.
        train_loader: Training data loader.
        val_loader: Validation data loader.
        config: Training configuration.
        device: Device to train on.

    Returns:
        Dictionary of training results.
    """
    from torch.optim import AdamW
    from torch.optim.lr_scheduler import CosineAnnealingLR

    router.to(device)
    router.train()

    optimizer = AdamW(
        router.parameters(),
        lr=config["learning_rate"],
        weight_decay=config["weight_decay"],
    )
    scheduler = CosineAnnealingLR(
        optimizer,
        T_max=config["num_epochs"],
        eta_min=config["learning_rate"] * 0.01,
    )

    best_frac_oracle = -float("inf")
    best_epoch = 0
    history = {"train_loss": [], "val_loss": [], "frac_oracle": [], "top1_acc": []}

    start_time = time.time()

    for epoch in range(config["num_epochs"]):
        # Training
        router.train()
        epoch_loss = 0.0
        num_batches = 0

        for query_embs, model_embs, targets in train_loader:
            query_embs = query_embs.to(device)
            model_embs = model_embs.to(device)
            targets = targets.to(device)

            optimizer.zero_grad()

            output = router(query_embs, model_embs)

            # Hybrid loss: contrastive + KL
            best_idx = targets.argmax(dim=-1)
            scaled_scores = output.scores / config["temperature"]
            log_probs = F.log_softmax(scaled_scores, dim=-1)
            contrastive_loss = -log_probs.gather(1, best_idx.unsqueeze(1)).squeeze(1).mean()

            soft_targets = F.softmax(targets / 1.5, dim=-1).clamp_min(1e-8)
            kl_loss = F.kl_div(F.log_softmax(output.scores, dim=-1), soft_targets, reduction="batchmean")

            loss = 0.5 * contrastive_loss + 0.5 * kl_loss

            loss.backward()
            torch.nn.utils.clip_grad_norm_(router.parameters(), 1.0)
            optimizer.step()

            epoch_loss += loss.item()
            num_batches += 1

        scheduler.step()
        avg_loss = epoch_loss / max(num_batches, 1)
        history["train_loss"].append(avg_loss)

        # Validation
        router.eval()
        val_scores = []
        val_targets = []

        with torch.no_grad():
            for query_embs, model_embs, targets in val_loader:
                query_embs = query_embs.to(device)
                model_embs = model_embs.to(device)

                output = router(query_embs, model_embs)
                val_scores.append(output.scores.cpu())
                val_targets.append(targets)

        val_scores = torch.cat(val_scores, dim=0)
        val_targets = torch.cat(val_targets, dim=0)

        # Compute metrics
        pred_best = val_scores.argmax(dim=-1)
        true_best = val_targets.argmax(dim=-1)
        top1_acc = (pred_best == true_best).float().mean().item()

        oracle_perf = val_targets.max(dim=-1).values
        random_perf = val_targets.mean(dim=-1)
        selected_perf = val_targets.gather(1, pred_best.unsqueeze(1)).squeeze(1)
        denom = (oracle_perf - random_perf).clamp_min(1e-8)
        frac_oracle = ((selected_perf - random_perf) / denom).clamp(0, 1).mean().item()

        # Val loss
        best_idx = val_targets.argmax(dim=-1)
        scaled_scores = val_scores / config["temperature"]
        log_probs = F.log_softmax(scaled_scores, dim=-1)
        val_loss = -log_probs.gather(1, best_idx.unsqueeze(1)).squeeze(1).mean().item()

        history["val_loss"].append(val_loss)
        history["frac_oracle"].append(frac_oracle)
        history["top1_acc"].append(top1_acc)

        if frac_oracle > best_frac_oracle:
            best_frac_oracle = frac_oracle
            best_epoch = epoch

        if (epoch + 1) % 10 == 0 or epoch == 0:
            logger.info(
                f"[{name}] Epoch {epoch+1}/{config['num_epochs']} - "
                f"Loss: {avg_loss:.4f}, Val Loss: {val_loss:.4f}, "
                f"Frac Oracle: {frac_oracle:.4f}, Top1: {top1_acc:.4f}"
            )

        # Early stopping
        if epoch - best_epoch > config.get("patience", 20):
            logger.info(f"[{name}] Early stopping at epoch {epoch+1}")
            break

    training_time = time.time() - start_time

    return {
        "name": name,
        "best_frac_oracle": best_frac_oracle,
        "best_top1_acc": history["top1_acc"][best_epoch],
        "best_epoch": best_epoch,
        "training_time": training_time,
        "history": history,
        "params": router.count_parameters(),
    }


def main():
    parser = argparse.ArgumentParser(description="Train baseline routers on router data")
    parser.add_argument("--data_path", type=str, required=True, help="Path to parquet file")
    parser.add_argument("--embedding_dim", type=int, default=512, help="Embedding dimension")
    parser.add_argument("--hidden_dim", type=int, default=384, help="Hidden dimension")
    parser.add_argument("--num_epochs", type=int, default=100, help="Training epochs")
    parser.add_argument("--batch_size", type=int, default=64, help="Batch size")
    parser.add_argument("--learning_rate", type=float, default=1e-3, help="Learning rate")
    parser.add_argument("--output_dir", type=str, default="results/baselines", help="Output directory")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    args = parser.parse_args()

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

    # Device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Using device: {device}")
    if torch.cuda.is_available():
        logger.info(f"GPU: {torch.cuda.get_device_name(0)}")

    # Load data
    dataset = RouterParquetDataset(
        args.data_path,
        embedding_dim=args.embedding_dim,
    )

    # Split data
    train_size = int(0.7 * len(dataset))
    val_size = int(0.15 * len(dataset))
    test_size = len(dataset) - train_size - val_size

    train_dataset, val_dataset, test_dataset = random_split(
        dataset,
        [train_size, val_size, test_size],
        generator=torch.Generator().manual_seed(args.seed),
    )

    logger.info(f"Split: train={train_size}, val={val_size}, test={test_size}")

    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=4,
        pin_memory=True,
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=collate_fn,
        num_workers=4,
        pin_memory=True,
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=collate_fn,
        num_workers=4,
        pin_memory=True,
    )

    # Import baselines
    import sys
    sys.path.insert(0, ".")
    from embergraphrouter.model.baselines import (
        TwoTowerRouter,
        MultiHeadMLPRouter,
        SimplifiedGNNRouter,
        RouterConfig,
    )

    # Config
    config = RouterConfig(
        query_dim=dataset.embedding_dim,
        model_dim=dataset.embedding_dim,
        hidden_dim=args.hidden_dim,
        num_models=dataset.num_models,
        num_layers=2,
        num_heads=4,
        dropout=0.1,
    )

    train_config = {
        "num_epochs": args.num_epochs,
        "learning_rate": args.learning_rate,
        "weight_decay": 1e-4,
        "temperature": 0.07,
        "patience": 20,
    }

    # Train all baselines
    results = []

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

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

        result = train_baseline(name, router, train_loader, val_loader, train_config, device)
        results.append(result)

        # Test evaluation
        router.eval()
        test_scores = []
        test_targets = []

        with torch.no_grad():
            for query_embs, model_embs, targets in test_loader:
                query_embs = query_embs.to(device)
                model_embs = model_embs.to(device)
                output = router(query_embs, model_embs)
                test_scores.append(output.scores.cpu())
                test_targets.append(targets)

        test_scores = torch.cat(test_scores, dim=0)
        test_targets = torch.cat(test_targets, dim=0)

        # Test metrics
        pred_best = test_scores.argmax(dim=-1)
        true_best = test_targets.argmax(dim=-1)
        test_top1 = (pred_best == true_best).float().mean().item()

        oracle_perf = test_targets.max(dim=-1).values
        random_perf = test_targets.mean(dim=-1)
        selected_perf = test_targets.gather(1, pred_best.unsqueeze(1)).squeeze(1)
        denom = (oracle_perf - random_perf).clamp_min(1e-8)
        test_frac_oracle = ((selected_perf - random_perf) / denom).clamp(0, 1).mean().item()

        result["test_frac_oracle"] = test_frac_oracle
        result["test_top1_acc"] = test_top1

        logger.info(f"[{name}] TEST - Frac Oracle: {test_frac_oracle:.4f}, Top1: {test_top1:.4f}")

    # Print summary
    logger.info("\n" + "=" * 70)
    logger.info("FINAL RESULTS")
    logger.info("=" * 70)
    logger.info(f"{'Router':<20} {'Params':>12} {'Test FracOracle':>16} {'Test Top1':>12} {'Time (s)':>10}")
    logger.info("-" * 70)
    for r in results:
        logger.info(
            f"{r['name']:<20} {r['params']:>12,} {r['test_frac_oracle']:>16.4f} "
            f"{r['test_top1_acc']:>12.4f} {r['training_time']:>10.1f}"
        )

    # Save results
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    results_path = output_dir / f"baseline_comparison_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    with open(results_path, "w") as f:
        # Convert non-serializable items
        serializable_results = []
        for r in results:
            sr = {k: v for k, v in r.items() if k != "history"}
            sr["history"] = {k: [float(x) for x in v] for k, v in r["history"].items()}
            serializable_results.append(sr)
        json.dump(serializable_results, f, indent=2)

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


if __name__ == "__main__":
    main()
