"""
VideoMAEv2Classifier - Video classifier using VideoMAEv2 backbone.

Supports:
    - VideoMAEv2 models (Base, Large, etc.)
    - Handles input format differences (expects B, C, T, H, W)
    - Automatic resizing to model's expected resolution
"""

import torch
import torch.nn.functional as F
from transformers import AutoConfig, AutoModel

from badas.core.base import VideoClassifierBase


class VideoMAEv2Classifier(VideoClassifierBase):
    """
    Video classifier using VideoMAEv2 backbone.

    Note: VideoMAEv2 expects different input format than V-JEPA2:
        - V-JEPA2: (B, T, C, H, W)
        - VideoMAEv2: (B, C, T, H, W)

    This class handles the conversion automatically.
    """

    MODEL_TYPE = "videomae"

    # Expected input resolution for VideoMAEv2
    EXPECTED_SIZE = 224

    def __init__(self, config):
        self._cached_feature_dim = None
        super().__init__(config)

    @property
    def feature_dim(self) -> int:
        if self._cached_feature_dim is not None:
            return self._cached_feature_dim
        raise RuntimeError("feature_dim accessed before backbone loaded")

    def _load_backbone(self):
        """Load VideoMAEv2 backbone."""
        model_name = self.config.model.model_name

        model_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)

        self.backbone = AutoModel.from_pretrained(
            model_name,
            config=model_config,
            trust_remote_code=True,
        )

        # Extract feature dimension from config
        self._cached_feature_dim = self._get_hidden_size(model_config)

        # Freeze backbone if configured
        if self.config.model.freeze_backbone:
            for param in self.backbone.parameters():
                param.requires_grad = False

    def _get_hidden_size(self, model_config) -> int:
        """Extract hidden size from VideoMAEv2 config."""
        # VideoMAEv2 has nested config structure
        if hasattr(model_config, "model_config"):
            inner = model_config.model_config
            return inner.get("embed_dim", 768)

        # Try common attribute names
        if hasattr(model_config, "hidden_size"):
            return model_config.hidden_size
        if hasattr(model_config, "embed_dim"):
            return model_config.embed_dim

        return 768  # Default for Base model

    def _extract_features(self, x: torch.Tensor) -> torch.Tensor:
        """
        Extract features from VideoMAEv2 backbone.

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

        Returns:
            Features (B, feature_dim) - VideoMAEv2 returns pooled features
        """
        # Convert to VideoMAEv2 format: (B, T, C, H, W) -> (B, C, T, H, W)
        x = x.permute(0, 2, 1, 3, 4)

        # Resize if needed (VideoMAEv2 expects 224x224)
        x = self._resize_if_needed(x)

        # Forward through backbone
        output = self.backbone(pixel_values=x)

        # VideoMAEv2 returns tensor directly (already pooled)
        if isinstance(output, torch.Tensor):
            return output

        # Or extract from output object
        if hasattr(output, "last_hidden_state"):
            features = output.last_hidden_state
            # Apply mean pooling if we get a sequence
            if len(features.shape) == 3:
                return features.mean(dim=1)
            return features

        return output

    def _resize_if_needed(self, x: torch.Tensor) -> torch.Tensor:
        """
        Resize video frames to expected size if needed.

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

        Returns:
            Resized tensor
        """
        B, C, T, H, W = x.shape

        if H == self.EXPECTED_SIZE and W == self.EXPECTED_SIZE:
            return x

        # Reshape for interpolation: (B, C, T, H, W) -> (B*C*T, 1, H, W)
        x = x.permute(0, 2, 1, 3, 4)  # (B, T, C, H, W)
        x = x.reshape(B * T, C, H, W)

        # Resize
        x = F.interpolate(
            x,
            size=(self.EXPECTED_SIZE, self.EXPECTED_SIZE),
            mode="bilinear",
            align_corners=False,
        )

        # Reshape back: (B*T, C, H, W) -> (B, C, T, H, W)
        x = x.reshape(B, T, C, self.EXPECTED_SIZE, self.EXPECTED_SIZE)
        x = x.permute(0, 2, 1, 3, 4)

        return x

    def _apply_temporal_processing(self, features: torch.Tensor) -> torch.Tensor:
        """
        Apply temporal processing.

        VideoMAEv2 typically returns already-pooled features (B, D),
        so temporal processing may not be needed.
        """
        if len(features.shape) == 2:
            # Already pooled, skip temporal processing
            return features

        return super()._apply_temporal_processing(features)
