"""
VideoClassifierBase - Abstract base class for video classifiers.

All video classifier implementations (V-JEPA2, VideoMAE, etc.) inherit from this class.
Shared logic lives here; model-specific logic is implemented in subclasses.
"""

from abc import ABC, abstractmethod
from typing import Optional

import torch
import torch.nn as nn

from badas.core.modules import (
    AttentionHead,
    AttentionProcessor,
    AttentiveProbe,
    LinearHead,
    LSTMProcessor,
    MaxPooling,
    MeanPooling,
    MLPHead,
)


class VideoClassifierBase(ABC, nn.Module):
    """
    Abstract base class for video classifiers.

    Subclasses must implement:
        - _load_backbone(): Load the model backbone
        - _extract_features(): Extract features from video input
        - feature_dim (property): Return the backbone's feature dimension

    Shared functionality:
        - Temporal processing (mean, max, LSTM, attention, probe)
        - Classification heads (linear, MLP, attention)
        - Forward pass orchestration
    """

    # Model type identifier for checkpoint saving/loading
    MODEL_TYPE: str = None

    def __init__(self, config):
        super().__init__()
        self.config = config
        self._feature_dim: Optional[int] = None

        # Components (initialized by subclass and shared methods)
        self.backbone = None
        self.temporal_processor = None
        self.classifier = None

        # Build model
        self._build_model()

    def _build_model(self):
        """Build the complete model pipeline."""
        self._load_backbone()
        self._setup_temporal_processing()
        self._setup_classification_head()

    @property
    @abstractmethod
    def feature_dim(self) -> int:
        """Return the backbone's output feature dimension."""
        pass

    @property
    def output_feature_dim(self) -> int:
        """Return the feature dimension after temporal processing."""
        return self._output_feature_dim

    @abstractmethod
    def _load_backbone(self):
        """Load the model backbone. Must set self.backbone."""
        pass

    @abstractmethod
    def _extract_features(self, x: torch.Tensor) -> torch.Tensor:
        """
        Extract features from video input.

        Args:
            x: Video tensor (B, T, C, H, W)

        Returns:
            Features tensor (B, seq_len, feature_dim) or (B, feature_dim)
        """
        pass

    def _setup_temporal_processing(self):
        """Setup temporal processing module based on config."""
        temporal_method = self.config.model.temporal_method.lower()
        dim = self.feature_dim

        # Track output dimension (may change for some processors)
        self._output_feature_dim = dim

        if temporal_method == "mean":
            self.temporal_processor = MeanPooling(dim)

        elif temporal_method == "max":
            self.temporal_processor = MaxPooling(dim)

        elif temporal_method == "lstm":
            hidden_dim = self.config.model.temporal_hidden_dim
            self.temporal_processor = LSTMProcessor(dim, hidden_dim)
            self._output_feature_dim = hidden_dim

        elif temporal_method == "attention":
            num_heads = self.config.model.temporal_num_heads
            self.temporal_processor = AttentionProcessor(dim, num_heads)

        elif temporal_method == "probe":
            num_queries = getattr(self.config.model, "temporal_num_queries", 8)
            hidden_dim = self.config.model.temporal_hidden_dim
            self.temporal_processor = AttentiveProbe(dim, hidden_dim, num_queries)
            self._output_feature_dim = num_queries * hidden_dim

        else:
            # Default to mean pooling
            self.temporal_processor = MeanPooling(dim)

    def _setup_classification_head(self):
        """Setup classification head based on config."""
        head_type = self.config.model.head_type.lower()
        input_dim = self._output_feature_dim

        num_classes = self.config.model.num_classes
        if num_classes == 2 and not getattr(self.config.model, "force_num_classes", False):
            num_classes = 1  # Binary classification

        dropout = self.config.model.head_dropout

        if head_type == "linear":
            self.classifier = LinearHead(input_dim, num_classes, dropout)

        elif head_type == "mlp":
            self.classifier = MLPHead(
                input_dim,
                num_classes,
                self.config.model.head_hidden_dim,
                self.config.model.head_num_layers,
                dropout,
            )

        elif head_type == "attention":
            self.classifier = AttentionHead(
                input_dim,
                num_classes,
                self.config.model.head_hidden_dim,
                self.config.model.head_num_heads,
                self.config.model.head_num_layers,
                dropout,
            )

        else:
            # Default to linear
            self.classifier = LinearHead(input_dim, num_classes, dropout)

    def _apply_temporal_processing(self, features: torch.Tensor) -> torch.Tensor:
        """
        Apply temporal processing to convert sequence to single vector.

        Args:
            features: (B, seq_len, dim) or (B, dim)

        Returns:
            (B, output_dim)
        """
        if len(features.shape) == 2:
            # Already a single vector per sample
            return features

        if self.temporal_processor is None:
            # Fallback to mean pooling
            return features.mean(dim=1)

        return self.temporal_processor(features)

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

        Args:
            x: Video tensor (B, T, C, H, W)

        Returns:
            Logits tensor (B,) for binary or (B, num_classes) for multiclass
        """
        # Step 1: Extract features
        features = self._extract_features(x)

        # Step 2: Temporal processing
        features = self._apply_temporal_processing(features)

        # Step 3: Classification
        logits = self.classifier(features)

        # Squeeze for binary classification
        if self.config.model.num_classes == 2:
            logits = logits.squeeze(-1)

        return logits

    def get_checkpoint_metadata(self) -> dict:
        """Return metadata to save in checkpoint for proper loading."""
        if self.MODEL_TYPE is None:
            raise ValueError(f"{self.__class__.__name__} must define MODEL_TYPE")
        return {"model_type": self.MODEL_TYPE}

    def _report_architecture(self):
        """Print model architecture summary."""
        total_params = sum(p.numel() for p in self.parameters())
        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)

        print(f"\nModel Summary ({self.__class__.__name__}):")
        print(f"  Model type: {self.MODEL_TYPE}")
        print(f"  Backbone: {self.config.model.model_name}")
        print(f"  Feature dim: {self.feature_dim}")
        print(f"  Temporal: {self.config.model.temporal_method}")
        print(f"  Head: {self.config.model.head_type}")
        print(f"  Total params: {total_params:,}")
        print(f"  Trainable: {trainable_params:,}")
