"""
Lightning Module for Video Collision Anticipation
=================================================

Lightning module implementation that wraps video classifiers
and provides standardized training, validation, and testing loops.
"""

from typing import Any, Dict

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torchmetrics import (
    AUROC,
    Accuracy,
    AveragePrecision,
    Precision,
    Recall,
    Specificity,
)

from badas.training.config import Config
from badas.core import create_model, detect_model_type


class CollisionAnticipationLoss(nn.Module):
    """Loss function for collision anticipation"""

    def __init__(
        self, num_classes: int, temperature: float = None, label_smoothing: float = 0.0
    ):
        super().__init__()
        self.num_classes = num_classes
        self.temperature = temperature if temperature is not None else 1.0
        self.label_smoothing = label_smoothing
        self.ce_loss = (
            nn.CrossEntropyLoss(label_smoothing=label_smoothing)
            if num_classes > 2
            else nn.BCEWithLogitsLoss()
        )

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        if self.temperature is not None and self.num_classes > 2:
            logits = logits / self.temperature
        if self.num_classes == 2:
            targets = targets.float()
            if self.label_smoothing > 0:
                targets = (
                    targets * (1.0 - self.label_smoothing) + self.label_smoothing * 0.5
                )
        return self.ce_loss(logits, targets)


class FocalLoss(nn.Module):
    """Focal Loss for handling class imbalance and improving calibration.

    Focal loss down-weights easy examples and focuses on hard ones.
    FL(p_t) = -alpha_t * (1 - p_t)^gamma * log(p_t)

    Supports both binary and multi-class classification.
    """

    def __init__(
        self,
        gamma: float = 2.0,
        alpha: float = 0.25,
        num_classes: int = 2,
        label_smoothing: float = 0.0,
    ):
        super().__init__()
        self.gamma = gamma
        self.alpha = alpha
        self.num_classes = num_classes
        self.label_smoothing = label_smoothing

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        if self.num_classes == 2:
            return self._binary_focal_loss(logits, targets)
        else:
            return self._multiclass_focal_loss(logits, targets)

    def _binary_focal_loss(
        self, logits: torch.Tensor, targets: torch.Tensor
    ) -> torch.Tensor:
        """Binary focal loss."""
        targets = targets.float()
        if self.label_smoothing > 0:
            targets = (
                targets * (1.0 - self.label_smoothing) + self.label_smoothing * 0.5
            )
        probs = torch.sigmoid(logits)

        # Compute focal weights
        p_t = probs * targets + (1 - probs) * (1 - targets)
        alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets)
        focal_weight = alpha_t * (1 - p_t) ** self.gamma

        # BCE loss
        bce = nn.functional.binary_cross_entropy_with_logits(
            logits, targets, reduction="none"
        )

        return (focal_weight * bce).mean()

    def _multiclass_focal_loss(
        self, logits: torch.Tensor, targets: torch.Tensor
    ) -> torch.Tensor:
        """Multi-class focal loss."""
        probs = torch.softmax(logits, dim=1)

        # Get probability of true class
        p_t = probs.gather(1, targets.unsqueeze(1)).squeeze(1)

        # Compute focal weight
        focal_weight = (1 - p_t) ** self.gamma

        # Cross entropy loss (unreduced)
        ce_loss = nn.functional.cross_entropy(
            logits, targets, reduction="none", label_smoothing=self.label_smoothing
        )

        return (focal_weight * ce_loss).mean()


class NegativeMarginLoss(nn.Module):
    """Loss that pushes negative predictions toward 0 with a margin.

    For positives: standard BCE loss
    For negatives: penalize predictions above the margin (e.g., 0.3)
    This encourages the model to output low values for negatives, reducing noise floor.
    """

    def __init__(self, margin: float = 0.3, label_smoothing: float = 0.0):
        super().__init__()
        self.margin = margin
        self.label_smoothing = label_smoothing
        self.bce = nn.BCEWithLogitsLoss(reduction="none")

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        targets = targets.float()
        is_negative = (targets == 0).float()
        if self.label_smoothing > 0:
            targets = (
                targets * (1.0 - self.label_smoothing) + self.label_smoothing * 0.5
            )
        probs = torch.sigmoid(logits)

        # Standard BCE for all samples
        bce_loss = self.bce(logits, targets)

        # Additional penalty for negatives above margin
        above_margin = torch.relu(probs - self.margin)
        margin_penalty = is_negative * (above_margin**2)

        # Combined loss
        return bce_loss.mean() + margin_penalty.mean()


class VideoLightningModule(pl.LightningModule):
    """Lightning module for video collision anticipation"""

    def __init__(self, config: Config):
        super().__init__()
        self.config = config

        # Detect model type for registry
        model_type = detect_model_type(config.model.model_name)

        # Save hyperparameters for checkpointing (include model_type)
        from dataclasses import asdict

        hparams = asdict(config)
        hparams["model_type"] = model_type
        self.save_hyperparameters(hparams)

        # Create model using registry
        self.model = create_model(config, model_type=model_type)

        # Initialize loss function based on config
        loss_type = getattr(config.training, "loss_type", "bce")
        num_classes = config.model.num_classes
        label_smoothing = getattr(config.training, "label_smoothing", 0.0)

        if loss_type == "focal":
            gamma = getattr(config.training, "focal_gamma", 2.0)
            alpha = getattr(config.training, "focal_alpha", 0.25)
            self.criterion = FocalLoss(
                gamma=gamma,
                alpha=alpha,
                num_classes=num_classes,
                label_smoothing=label_smoothing,
            )
            print(
                f"Using Focal Loss (gamma={gamma}, alpha={alpha}, num_classes={num_classes}, label_smoothing={label_smoothing})"
            )
        elif loss_type == "negative_margin":
            if num_classes > 2:
                print(
                    f"WARNING: NegativeMarginLoss only supports binary classification. "
                    f"Falling back to CrossEntropyLoss for {num_classes} classes."
                )
                self.criterion = CollisionAnticipationLoss(
                    num_classes=num_classes, label_smoothing=label_smoothing
                )
            else:
                margin = getattr(config.training, "negative_margin", 0.3)
                self.criterion = NegativeMarginLoss(
                    margin=margin, label_smoothing=label_smoothing
                )
                print(
                    f"Using Negative Margin Loss (margin={margin}, label_smoothing={label_smoothing})"
                )
        else:
            self.criterion = CollisionAnticipationLoss(
                num_classes=config.model.num_classes,
                temperature=(
                    config.model.temperature
                    if config.model.use_temperature_scaling
                    else None
                ),
                label_smoothing=label_smoothing,
            )
            print(
                f"Using standard BCE Loss{f' (label_smoothing={label_smoothing})' if label_smoothing > 0 else ''}"
            )

        task = "binary" if config.model.num_classes == 2 else "multiclass"
        num_classes = config.model.num_classes if config.model.num_classes > 2 else None

        self.train_accuracy = Accuracy(task=task, num_classes=num_classes)

        self.val_accuracy = Accuracy(task=task, num_classes=num_classes)
        self.val_precision = Precision(task=task, num_classes=num_classes)
        self.val_recall = Recall(task=task, num_classes=num_classes)
        self.val_specificity = Specificity(
            task=task, num_classes=num_classes
        )  # TNR = 1 - FPR
        self.val_auc = AUROC(task=task, num_classes=num_classes)
        self.val_ap = AveragePrecision(task=task, num_classes=num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass"""
        return self.model(x)

    def training_step(self, batch: tuple, batch_idx: int) -> torch.Tensor:
        """Training step"""
        videos, labels = batch

        outputs = self(videos)
        loss = self.criterion(outputs, labels)

        if self.config.model.num_classes == 2:
            preds = (outputs > 0.0).long()
        else:
            preds = torch.argmax(outputs, dim=1)

        self.train_accuracy(preds, labels)

        self.log("train_loss", loss, prog_bar=True)
        self.log("train_acc", self.train_accuracy, on_epoch=True, prog_bar=True)

        return loss

    def validation_step(self, batch: tuple, batch_idx: int) -> torch.Tensor:
        """Validation step"""
        videos, labels = batch

        outputs = self(videos)
        loss = self.criterion(outputs, labels)

        if self.config.model.num_classes == 2:
            preds = (outputs > 0.0).long()
            probs = torch.sigmoid(outputs)
        else:
            preds = torch.argmax(outputs, dim=1)
            probs = torch.softmax(outputs, dim=1)

        self.val_accuracy(preds, labels)
        self.val_precision(preds, labels)
        self.val_recall(preds, labels)
        self.val_specificity(preds, labels)
        self.val_auc(probs, labels)
        self.val_ap(probs, labels)

        self.log("val_loss", loss, on_epoch=True, prog_bar=True, sync_dist=True)
        self.log(
            "val_acc", self.val_accuracy, on_epoch=True, prog_bar=True, sync_dist=True
        )
        self.log("val_precision", self.val_precision, on_epoch=True, sync_dist=True)
        self.log("val_recall", self.val_recall, on_epoch=True, sync_dist=True)
        self.log(
            "val_specificity", self.val_specificity, on_epoch=True, sync_dist=True
        )  # TNR
        self.log("val_auc", self.val_auc, on_epoch=True, sync_dist=True)
        self.log("val_ap", self.val_ap, on_epoch=True, sync_dist=True)

        return loss

    def configure_optimizers(self) -> Dict[str, Any]:
        """Configure optimizers and learning rate schedulers."""
        optimizer = optim.AdamW(
            self.parameters(),
            lr=self.config.training.learning_rate,
            weight_decay=self.config.training.weight_decay,
        )

        scheduler = CosineAnnealingLR(
            optimizer,
            T_max=self.config.training.scheduler_T_max,
            eta_min=self.config.training.scheduler_eta_min,
        )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": 1,
            },
        }

    def on_train_epoch_start(self) -> None:
        if self.current_epoch == 0:
            self.model._report_architecture()
