"""
Training loop for personality classifier.

Provides a trainer class with:
- Training/validation split
- Early stopping
- Learning rate scheduling
- Checkpoint saving
- Logging
"""

from __future__ import annotations

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

import torch
from torch import nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader, Dataset, random_split

from .model import PersonalityClassifier


logger = logging.getLogger(__name__)


@dataclass
class TrainingConfig:
    """Configuration for training the personality classifier."""

    # Optimization
    learning_rate: float = 1e-3
    weight_decay: float = 0.01
    batch_size: int = 64
    epochs: int = 50

    # Validation
    val_split: float = 0.1
    eval_every: int = 1  # Evaluate every N epochs

    # Early stopping
    early_stopping_patience: int = 5
    min_delta: float = 1e-4  # Minimum improvement to count as progress

    # Checkpoints
    checkpoint_dir: str = "./checkpoints"
    save_best: bool = True

    # Reproducibility
    seed: int = 42

    # Device
    device: str | None = None  # Auto-detect if None


@dataclass
class TrainingResult:
    """Results from training."""

    best_val_loss: float
    final_train_loss: float
    final_epoch: int
    total_epochs: int
    best_epoch: int
    history: dict = field(default_factory=dict)


class PersonalityTrainer:
    """
    Trainer for the personality classifier.

    Handles the training loop with validation, early stopping,
    and checkpoint management.
    """

    def __init__(
        self,
        model: PersonalityClassifier,
        train_dataset: Dataset,
        val_dataset: Dataset | None = None,
        config: TrainingConfig | None = None,
    ) -> None:
        """
        Initialize trainer.

        Args:
            model: PersonalityClassifier model to train
            train_dataset: Training dataset
            val_dataset: Validation dataset (if None, split from train)
            config: Training configuration
        """
        self.config = config or TrainingConfig()
        self.model = model

        # Set random seed for reproducibility
        torch.manual_seed(self.config.seed)

        # Determine device
        if self.config.device:
            self.device = torch.device(self.config.device)
        else:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        logger.info(f"Using device: {self.device}")
        self.model.to(self.device)

        # Split train/val if no validation set provided
        if val_dataset is None and self.config.val_split > 0:
            val_size = int(len(train_dataset) * self.config.val_split)
            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
        self.train_loader = DataLoader(
            train_dataset,
            batch_size=self.config.batch_size,
            shuffle=True,
            num_workers=0,  # Avoid issues with multiprocessing
        )

        self.val_loader = None
        if val_dataset is not None and len(val_dataset) > 0:
            self.val_loader = DataLoader(
                val_dataset,
                batch_size=self.config.batch_size,
                shuffle=False,
                num_workers=0,
            )

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

        self.scheduler = CosineAnnealingLR(
            self.optimizer,
            T_max=self.config.epochs,
            eta_min=self.config.learning_rate / 100,
        )

        # Loss function (MSE for regression)
        self.criterion = nn.MSELoss()

        # Training state
        self.current_epoch = 0
        self.best_val_loss = float("inf")
        self.best_epoch = 0
        self.patience_counter = 0
        self.history: dict[str, list] = {
            "train_loss": [],
            "val_loss": [],
            "learning_rate": [],
        }

    def train(
        self,
        progress_callback: Callable[[int, int, float], None] | None = None,
    ) -> TrainingResult:
        """
        Run the full training loop.

        Args:
            progress_callback: Optional callback(epoch, total_epochs, loss)

        Returns:
            TrainingResult with training metrics
        """
        logger.info(f"Starting training for {self.config.epochs} epochs")
        logger.info(f"Training samples: {len(self.train_loader.dataset)}")
        if self.val_loader:
            logger.info(f"Validation samples: {len(self.val_loader.dataset)}")

        for epoch in range(self.config.epochs):
            self.current_epoch = epoch + 1

            # Train one epoch
            train_loss = self._train_epoch()
            self.history["train_loss"].append(train_loss)
            self.history["learning_rate"].append(self.scheduler.get_last_lr()[0])

            # Validate
            val_loss = None
            if self.val_loader and (epoch + 1) % self.config.eval_every == 0:
                val_loss = self._validate()
                self.history["val_loss"].append(val_loss)

                # Early stopping check
                if val_loss < self.best_val_loss - self.config.min_delta:
                    self.best_val_loss = val_loss
                    self.best_epoch = epoch + 1
                    self.patience_counter = 0

                    if self.config.save_best:
                        self._save_checkpoint("best_model.pt")
                else:
                    self.patience_counter += 1

                logger.info(
                    f"Epoch {epoch + 1}/{self.config.epochs}: "
                    f"train_loss={train_loss:.4f}, val_loss={val_loss:.4f}, "
                    f"lr={self.scheduler.get_last_lr()[0]:.2e}"
                )

                # Early stopping
                if self.patience_counter >= self.config.early_stopping_patience:
                    logger.info(f"Early stopping at epoch {epoch + 1}")
                    break
            else:
                logger.info(
                    f"Epoch {epoch + 1}/{self.config.epochs}: "
                    f"train_loss={train_loss:.4f}, "
                    f"lr={self.scheduler.get_last_lr()[0]:.2e}"
                )

            # Progress callback
            if progress_callback:
                progress_callback(epoch + 1, self.config.epochs, train_loss)

            # Update learning rate
            self.scheduler.step()

        # Load best model if we saved checkpoints
        if self.config.save_best and self.best_epoch > 0:
            self._load_checkpoint("best_model.pt")
            logger.info(f"Loaded best model from epoch {self.best_epoch}")

        return TrainingResult(
            best_val_loss=self.best_val_loss,
            final_train_loss=self.history["train_loss"][-1],
            final_epoch=self.current_epoch,
            total_epochs=self.config.epochs,
            best_epoch=self.best_epoch,
            history=self.history,
        )

    def _train_epoch(self) -> float:
        """Train for one epoch. Returns average loss."""
        self.model.train()
        total_loss = 0.0
        num_batches = 0

        for embeddings, scores in self.train_loader:
            embeddings = embeddings.to(self.device)
            scores = scores.to(self.device)

            # Forward pass
            self.optimizer.zero_grad()
            predictions = self.model(embeddings)
            loss = self.criterion(predictions, scores)

            # Backward pass
            loss.backward()
            self.optimizer.step()

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

        return total_loss / max(num_batches, 1)

    def _validate(self) -> float:
        """Validate on validation set. Returns average loss."""
        if self.val_loader is None:
            return 0.0

        self.model.eval()
        total_loss = 0.0
        num_batches = 0

        with torch.no_grad():
            for embeddings, scores in self.val_loader:
                embeddings = embeddings.to(self.device)
                scores = scores.to(self.device)

                predictions = self.model(embeddings)
                loss = self.criterion(predictions, scores)

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

        return total_loss / max(num_batches, 1)

    def _save_checkpoint(self, filename: str) -> Path:
        """Save model checkpoint."""
        checkpoint_dir = Path(self.config.checkpoint_dir)
        checkpoint_dir.mkdir(parents=True, exist_ok=True)

        checkpoint_path = checkpoint_dir / filename
        torch.save(
            {
                "epoch": self.current_epoch,
                "model_state_dict": self.model.state_dict(),
                "optimizer_state_dict": self.optimizer.state_dict(),
                "scheduler_state_dict": self.scheduler.state_dict(),
                "best_val_loss": self.best_val_loss,
                "config": self.model.get_config(),
            },
            checkpoint_path,
        )

        return checkpoint_path

    def _load_checkpoint(self, filename: str) -> None:
        """Load model checkpoint."""
        checkpoint_path = Path(self.config.checkpoint_dir) / filename
        if not checkpoint_path.exists():
            logger.warning(f"Checkpoint not found: {checkpoint_path}")
            return

        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        self.model.load_state_dict(checkpoint["model_state_dict"])
        logger.info(f"Loaded checkpoint from {checkpoint_path}")

    def evaluate(self) -> dict:
        """
        Evaluate model on validation set.

        Returns:
            Dictionary with evaluation metrics
        """
        if self.val_loader is None:
            return {"error": "No validation set available"}

        self.model.eval()
        total_loss = 0.0
        all_preds = []
        all_targets = []

        with torch.no_grad():
            for embeddings, scores in self.val_loader:
                embeddings = embeddings.to(self.device)
                scores = scores.to(self.device)

                predictions = self.model(embeddings)
                loss = self.criterion(predictions, scores)

                total_loss += loss.item()
                all_preds.append(predictions.cpu())
                all_targets.append(scores.cpu())

        all_preds = torch.cat(all_preds, dim=0)
        all_targets = torch.cat(all_targets, dim=0)

        # Per-dimension MAE
        mae_per_dim = (all_preds - all_targets).abs().mean(dim=0)

        return {
            "val_loss": total_loss / len(self.val_loader),
            "mae": float((all_preds - all_targets).abs().mean()),
            "mae_per_dimension": mae_per_dim.tolist(),
            "num_samples": len(all_targets),
        }
