"""Unified trainer for router baseline comparison.

This module provides a consistent training and evaluation framework for all
router architectures, ensuring fair comparison with identical data splits,
loss functions, and evaluation metrics.

Usage:
    from embergraphrouter.model.baselines import (
        TwoTowerRouter,
        MultiHeadMLPRouter,
        SimplifiedGNNRouter,
        RouterConfig,
    )
    from embergraphrouter.model.baselines.trainer import BaselineTrainer

    config = RouterConfig(query_dim=512, model_dim=512, hidden_dim=384)
    router = TwoTowerRouter(config)
    trainer = BaselineTrainer(router, config=TrainerConfig(...))
    metrics = trainer.train(train_data, val_data)

Design Principles:
    - Identical data loading and preprocessing across architectures
    - Same loss function options (contrastive, listwise KL, hybrid)
    - Unified evaluation metrics (frac_oracle, top1_acc, NDCG, ECE)
    - Clean experiment tracking via optional W&B integration
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
from torch.utils.data import DataLoader, Dataset, random_split

from .base import BaseRouter, DifficultySignals, RouterConfig

logger = logging.getLogger(__name__)


@dataclass
class TrainerConfig:
    """Configuration for baseline training.

    Attributes:
        batch_size: Training batch size.
        learning_rate: Initial learning rate.
        weight_decay: L2 regularization weight.
        num_epochs: Maximum training epochs.
        warmup_epochs: Linear warmup epochs before cosine decay.
        early_stopping_patience: Epochs without improvement before stopping.
        gradient_clip_norm: Maximum gradient norm (0 = disabled).
        loss_type: Loss function type (contrastive, kl, hybrid).
        loss_temperature: Temperature for soft targets.
        contrastive_temperature: Temperature for contrastive loss.
        val_fraction: Fraction of training data for validation.
        eval_every: Evaluate every N epochs.
        save_best: Whether to save best model checkpoint.
        checkpoint_dir: Directory for model checkpoints.
        use_wandb: Whether to log to Weights & Biases.
        wandb_project: W&B project name.
        wandb_run_name: W&B run name.
        seed: Random seed for reproducibility.
    """

    batch_size: int = 32
    learning_rate: float = 1e-3
    weight_decay: float = 1e-4
    num_epochs: int = 100
    warmup_epochs: int = 5
    early_stopping_patience: int = 20
    gradient_clip_norm: float = 1.0
    loss_type: str = "hybrid"
    loss_temperature: float = 1.5
    contrastive_temperature: float = 0.07
    val_fraction: float = 0.1
    eval_every: int = 1
    save_best: bool = True
    checkpoint_dir: str = "checkpoints"
    use_wandb: bool = False
    wandb_project: str = "router-baselines"
    wandb_run_name: Optional[str] = None
    seed: int = 42


@dataclass
class TrainingMetrics:
    """Metrics collected during training.

    Attributes:
        train_loss: Training loss per epoch.
        val_loss: Validation loss per epoch.
        frac_oracle: Fraction of oracle performance achieved.
        top1_accuracy: Top-1 routing accuracy.
        ndcg_at_3: NDCG@3 ranking quality.
        mean_reciprocal_rank: MRR for ranking evaluation.
        expected_calibration_error: ECE for probability calibration.
        best_epoch: Epoch with best validation performance.
        training_time_seconds: Total training time.
    """

    train_loss: list[float] = field(default_factory=list)
    val_loss: list[float] = field(default_factory=list)
    frac_oracle: list[float] = field(default_factory=list)
    top1_accuracy: list[float] = field(default_factory=list)
    ndcg_at_3: list[float] = field(default_factory=list)
    mean_reciprocal_rank: list[float] = field(default_factory=list)
    expected_calibration_error: list[float] = field(default_factory=list)
    best_epoch: int = 0
    training_time_seconds: float = 0.0

    def best_metrics(self) -> dict:
        """Get metrics at best epoch."""
        if not self.frac_oracle:
            return {}
        idx = self.best_epoch
        return {
            "frac_oracle": self.frac_oracle[idx] if idx < len(self.frac_oracle) else 0,
            "top1_accuracy": self.top1_accuracy[idx] if idx < len(self.top1_accuracy) else 0,
            "ndcg_at_3": self.ndcg_at_3[idx] if idx < len(self.ndcg_at_3) else 0,
            "mrr": self.mean_reciprocal_rank[idx] if idx < len(self.mean_reciprocal_rank) else 0,
            "ece": self.expected_calibration_error[idx] if idx < len(self.expected_calibration_error) else 0,
            "best_epoch": self.best_epoch,
            "training_time": self.training_time_seconds,
        }


class RouterDataset(Dataset):
    """Dataset for router training.

    Each sample contains:
        - query_emb: Query embedding (query_dim,)
        - model_embs: Model embeddings (num_models, model_dim)
        - targets: Ground truth effects (num_models,)
    """

    def __init__(
        self,
        query_embs: torch.Tensor,
        model_embs: torch.Tensor,
        targets: torch.Tensor,
    ) -> None:
        """Initialize dataset.

        Args:
            query_embs: Query embeddings, shape (num_samples, query_dim).
            model_embs: Model embeddings, shape (num_models, model_dim).
            targets: Target effects, shape (num_samples, num_models).
        """
        self.query_embs = query_embs
        self.model_embs = model_embs
        self.targets = targets

    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_router_batch(
    batch: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Collate function for RouterDataset.

    Args:
        batch: List of (query_emb, model_embs, targets) tuples.

    Returns:
        Batched tensors: (query_embs, model_embs, targets).
    """
    query_embs = torch.stack([x[0] for x in batch])
    model_embs = batch[0][1]  # Same for all samples
    targets = torch.stack([x[2] for x in batch])
    return query_embs, model_embs, targets


class BaselineTrainer:
    """Unified trainer for router baselines.

    Provides consistent training loop, loss computation, and evaluation
    across all router architectures.
    """

    def __init__(
        self,
        router: BaseRouter,
        config: TrainerConfig,
        device: Optional[torch.device] = None,
    ) -> None:
        """Initialize trainer.

        Args:
            router: Router model to train.
            config: Training configuration.
            device: Device to train on (default: auto-detect).
        """
        self.router = router
        self.config = config
        self.device = device or torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )

        self.router.to(self.device)
        self._setup_logging()

    def _setup_logging(self) -> None:
        """Setup experiment logging."""
        if self.config.use_wandb:
            try:
                import wandb
                wandb.init(
                    project=self.config.wandb_project,
                    name=self.config.wandb_run_name,
                    config={
                        "router": self.router.get_architecture_name(),
                        "trainer": self.config.__dict__,
                        "router_config": self.router.config.__dict__,
                    },
                )
            except ImportError:
                logger.warning("wandb not installed, disabling W&B logging")
                self.config.use_wandb = False

    def _compute_loss(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
        query_repr: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Compute training loss.

        Args:
            scores: Predicted scores, shape (batch, num_models).
            targets: Ground truth effects, shape (batch, num_models).
            query_repr: Query representations for uniformity loss (optional).

        Returns:
            Scalar loss tensor.
        """
        loss_type = self.config.loss_type

        if loss_type == "contrastive":
            # InfoNCE: pull toward best model
            best_idx = targets.argmax(dim=-1)
            scaled_scores = scores / self.config.contrastive_temperature
            log_probs = F.log_softmax(scaled_scores, dim=-1)
            positive_log_prob = log_probs.gather(1, best_idx.unsqueeze(1)).squeeze(1)
            return -positive_log_prob.mean()

        elif loss_type == "kl":
            # KL divergence to soft targets
            soft_targets = F.softmax(
                targets / self.config.loss_temperature, dim=-1
            ).clamp_min(1e-8)
            log_probs = F.log_softmax(scores, dim=-1)
            return F.kl_div(log_probs, soft_targets, reduction="batchmean")

        elif loss_type == "hybrid":
            # Combination of contrastive and KL
            contrastive_weight = 0.5
            kl_weight = 0.5

            # Contrastive term
            best_idx = targets.argmax(dim=-1)
            scaled_scores = scores / self.config.contrastive_temperature
            log_probs = F.log_softmax(scaled_scores, dim=-1)
            positive_log_prob = log_probs.gather(1, best_idx.unsqueeze(1)).squeeze(1)
            contrastive_loss = -positive_log_prob.mean()

            # KL term
            soft_targets = F.softmax(
                targets / self.config.loss_temperature, dim=-1
            ).clamp_min(1e-8)
            log_probs_kl = F.log_softmax(scores, dim=-1)
            kl_loss = F.kl_div(log_probs_kl, soft_targets, reduction="batchmean")

            return contrastive_weight * contrastive_loss + kl_weight * kl_loss

        else:
            raise ValueError(f"Unknown loss type: {loss_type}")

    def _evaluate(
        self,
        router: BaseRouter,
        dataloader: DataLoader,
    ) -> dict[str, float]:
        """Evaluate router on a dataset.

        Args:
            router: Router model to evaluate.
            dataloader: Evaluation data loader.

        Returns:
            Dictionary of evaluation metrics.
        """
        router.eval()

        all_scores = []
        all_targets = []

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

                output = router(query_embs, model_embs)
                all_scores.append(output.scores.cpu())
                all_targets.append(targets.cpu())

        scores = torch.cat(all_scores, dim=0)
        targets = torch.cat(all_targets, dim=0)

        return self._compute_metrics(scores, targets)

    def _compute_metrics(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
    ) -> dict[str, float]:
        """Compute evaluation metrics.

        Args:
            scores: Predicted scores, shape (num_samples, num_models).
            targets: Ground truth effects, shape (num_samples, num_models).

        Returns:
            Dictionary of metrics.
        """
        num_samples = scores.shape[0]

        # Top-1 accuracy: does the model correctly identify the best model?
        pred_best = scores.argmax(dim=-1)
        true_best = targets.argmax(dim=-1)
        top1_accuracy = (pred_best == true_best).float().mean().item()

        # Fraction of oracle: weighted performance achieved
        # Oracle selects best model, random selects uniformly
        oracle_perf = targets.max(dim=-1).values
        random_perf = targets.mean(dim=-1)
        selected_perf = targets.gather(1, pred_best.unsqueeze(1)).squeeze(1)

        # Avoid division by zero
        denom = (oracle_perf - random_perf).clamp_min(1e-8)
        frac_oracle_per_sample = (selected_perf - random_perf) / denom
        frac_oracle = frac_oracle_per_sample.clamp(0, 1).mean().item()

        # NDCG@3: ranking quality for top 3
        ndcg_at_3 = self._compute_ndcg(scores, targets, k=3)

        # MRR: Mean Reciprocal Rank
        mrr = self._compute_mrr(scores, targets)

        # ECE: Expected Calibration Error
        probs = F.softmax(scores, dim=-1)
        ece = self._compute_ece(probs, targets)

        # Loss
        loss = self._compute_loss(scores, targets).item()

        return {
            "loss": loss,
            "frac_oracle": frac_oracle,
            "top1_accuracy": top1_accuracy,
            "ndcg_at_3": ndcg_at_3,
            "mrr": mrr,
            "ece": ece,
        }

    def _compute_ndcg(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
        k: int = 3,
    ) -> float:
        """Compute NDCG@k."""
        # DCG@k for predictions
        _, pred_order = scores.sort(dim=-1, descending=True)
        pred_order_k = pred_order[:, :k]

        # Gather relevance scores in predicted order
        relevance = targets.gather(1, pred_order_k)

        # DCG with log2 discount
        positions = torch.arange(1, k + 1, device=scores.device, dtype=torch.float)
        discounts = torch.log2(positions + 1)
        dcg = (relevance / discounts).sum(dim=-1)

        # Ideal DCG (sorted by true relevance)
        ideal_relevance, _ = targets.sort(dim=-1, descending=True)
        ideal_relevance_k = ideal_relevance[:, :k]
        idcg = (ideal_relevance_k / discounts).sum(dim=-1)

        # NDCG
        ndcg = dcg / idcg.clamp_min(1e-8)
        return ndcg.mean().item()

    def _compute_mrr(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
    ) -> float:
        """Compute Mean Reciprocal Rank."""
        # Rank of the best model according to predictions
        true_best = targets.argmax(dim=-1)  # (batch,)
        _, pred_order = scores.sort(dim=-1, descending=True)

        # Find rank of true best in predicted order
        ranks = (pred_order == true_best.unsqueeze(1)).nonzero(as_tuple=True)[1] + 1
        reciprocal_ranks = 1.0 / ranks.float()
        return reciprocal_ranks.mean().item()

    def _compute_ece(
        self,
        probs: torch.Tensor,
        targets: torch.Tensor,
        num_bins: int = 10,
    ) -> float:
        """Compute Expected Calibration Error."""
        # Get predicted class probabilities
        max_probs, pred_classes = probs.max(dim=-1)
        true_best = targets.argmax(dim=-1)
        correctness = (pred_classes == true_best).float()

        # Bin by confidence
        bin_boundaries = torch.linspace(0, 1, num_bins + 1, device=probs.device)
        ece = 0.0
        total_samples = probs.shape[0]

        for i in range(num_bins):
            in_bin = (max_probs > bin_boundaries[i]) & (max_probs <= bin_boundaries[i + 1])
            bin_size = in_bin.sum().item()

            if bin_size > 0:
                bin_accuracy = correctness[in_bin].mean().item()
                bin_confidence = max_probs[in_bin].mean().item()
                ece += (bin_size / total_samples) * abs(bin_accuracy - bin_confidence)

        return ece

    def train(
        self,
        train_dataset: RouterDataset,
        val_dataset: Optional[RouterDataset] = None,
    ) -> TrainingMetrics:
        """Train the router model.

        Args:
            train_dataset: Training dataset.
            val_dataset: Validation dataset (optional, will split from train if None).

        Returns:
            TrainingMetrics with results.
        """
        # Set random seed
        torch.manual_seed(self.config.seed)
        np.random.seed(self.config.seed)

        # Split validation if not provided
        if val_dataset is None:
            val_size = int(len(train_dataset) * self.config.val_fraction)
            train_size = len(train_dataset) - val_size
            train_dataset, val_dataset = random_split(
                train_dataset,
                [train_size, val_size],
                generator=torch.Generator().manual_seed(self.config.seed),
            )

        # Create data loaders
        train_loader = DataLoader(
            train_dataset,
            batch_size=self.config.batch_size,
            shuffle=True,
            collate_fn=collate_router_batch,
            drop_last=True,
        )
        val_loader = DataLoader(
            val_dataset,
            batch_size=self.config.batch_size,
            shuffle=False,
            collate_fn=collate_router_batch,
        )

        # Optimizer and scheduler
        optimizer = AdamW(
            self.router.parameters(),
            lr=self.config.learning_rate,
            weight_decay=self.config.weight_decay,
        )

        warmup_scheduler = LinearLR(
            optimizer,
            start_factor=0.1,
            end_factor=1.0,
            total_iters=self.config.warmup_epochs,
        )
        cosine_scheduler = CosineAnnealingLR(
            optimizer,
            T_max=self.config.num_epochs - self.config.warmup_epochs,
            eta_min=self.config.learning_rate * 0.01,
        )
        scheduler = SequentialLR(
            optimizer,
            schedulers=[warmup_scheduler, cosine_scheduler],
            milestones=[self.config.warmup_epochs],
        )

        # Training loop
        metrics = TrainingMetrics()
        best_frac_oracle = -float("inf")
        epochs_without_improvement = 0
        start_time = time.time()

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

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

                optimizer.zero_grad()

                output = self.router(query_embs, model_embs)
                loss = self._compute_loss(output.scores, targets, output.query_repr)

                loss.backward()

                if self.config.gradient_clip_norm > 0:
                    nn.utils.clip_grad_norm_(
                        self.router.parameters(),
                        self.config.gradient_clip_norm,
                    )

                optimizer.step()

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

            scheduler.step()
            avg_train_loss = epoch_loss / max(num_batches, 1)
            metrics.train_loss.append(avg_train_loss)

            # Evaluation
            if (epoch + 1) % self.config.eval_every == 0:
                val_metrics = self._evaluate(self.router, val_loader)

                metrics.val_loss.append(val_metrics["loss"])
                metrics.frac_oracle.append(val_metrics["frac_oracle"])
                metrics.top1_accuracy.append(val_metrics["top1_accuracy"])
                metrics.ndcg_at_3.append(val_metrics["ndcg_at_3"])
                metrics.mean_reciprocal_rank.append(val_metrics["mrr"])
                metrics.expected_calibration_error.append(val_metrics["ece"])

                # Early stopping check
                if val_metrics["frac_oracle"] > best_frac_oracle:
                    best_frac_oracle = val_metrics["frac_oracle"]
                    metrics.best_epoch = epoch
                    epochs_without_improvement = 0

                    if self.config.save_best:
                        self._save_checkpoint(epoch, val_metrics)
                else:
                    epochs_without_improvement += 1

                # Logging
                logger.info(
                    f"Epoch {epoch + 1}/{self.config.num_epochs} - "
                    f"Train Loss: {avg_train_loss:.4f}, "
                    f"Val Loss: {val_metrics['loss']:.4f}, "
                    f"Frac Oracle: {val_metrics['frac_oracle']:.4f}, "
                    f"Top1 Acc: {val_metrics['top1_accuracy']:.4f}"
                )

                if self.config.use_wandb:
                    import wandb
                    wandb.log({
                        "epoch": epoch + 1,
                        "train_loss": avg_train_loss,
                        "val_loss": val_metrics["loss"],
                        "frac_oracle": val_metrics["frac_oracle"],
                        "top1_accuracy": val_metrics["top1_accuracy"],
                        "ndcg_at_3": val_metrics["ndcg_at_3"],
                        "mrr": val_metrics["mrr"],
                        "ece": val_metrics["ece"],
                        "lr": scheduler.get_last_lr()[0],
                    })

                # Early stopping
                if epochs_without_improvement >= self.config.early_stopping_patience:
                    logger.info(
                        f"Early stopping after {epoch + 1} epochs "
                        f"(no improvement for {epochs_without_improvement} epochs)"
                    )
                    break

        metrics.training_time_seconds = time.time() - start_time

        if self.config.use_wandb:
            import wandb
            wandb.log({"best_metrics": metrics.best_metrics()})
            wandb.finish()

        return metrics

    def _save_checkpoint(
        self,
        epoch: int,
        metrics: dict[str, float],
    ) -> None:
        """Save model checkpoint."""
        checkpoint_dir = Path(self.config.checkpoint_dir)
        checkpoint_dir.mkdir(parents=True, exist_ok=True)

        checkpoint_path = checkpoint_dir / f"{self.router.get_architecture_name()}_best.pt"

        torch.save({
            "epoch": epoch,
            "model_state_dict": self.router.state_dict(),
            "router_config": self.router.config.__dict__,
            "trainer_config": self.config.__dict__,
            "metrics": metrics,
        }, checkpoint_path)

        logger.info(f"Saved checkpoint to {checkpoint_path}")

    def load_checkpoint(self, checkpoint_path: str) -> dict:
        """Load model from checkpoint.

        Args:
            checkpoint_path: Path to checkpoint file.

        Returns:
            Checkpoint metadata dict.
        """
        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        self.router.load_state_dict(checkpoint["model_state_dict"])
        return checkpoint
