"""
VJEPA2Classifier - Video classifier using V-JEPA2 backbone.

Supports:
    - V-JEPA2 models (ViT-B, ViT-L, ViT-H, ViT-G)
    - Future prediction using the predictor module
"""

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

from badas.core.base import VideoClassifierBase


class VJEPA2Classifier(VideoClassifierBase):
    """
    Video classifier using V-JEPA2 backbone.

    V-JEPA2 models include a predictor for future frame prediction,
    which can be used to enhance classification with temporal context.
    """

    MODEL_TYPE = "vjepa2"

    # Dimension mapping for V-JEPA2 model sizes
    DIMENSION_MAP = {
        "vits": 384,
        "vitb": 768,
        "vitl": 1024,
        "vith": 1280,
        "vitg": 1408,
    }

    def __init__(self, config):
        # V-JEPA2 specific state
        self.predictor = None
        self.has_predictor = False
        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 V-JEPA2 backbone."""
        model_name = self.config.model.model_name

        model_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
        model_config.drop_path_rate = self.config.model.drop_path_rate

        from_pretrained_kwargs = {'trust_remote_code': True}
        attn_impl = getattr(self.config.model, 'attn_implementation', None)
        if attn_impl:
            from_pretrained_kwargs['attn_implementation'] = attn_impl

        self.backbone = AutoModel.from_pretrained(
            model_name,
            config=model_config,
            **from_pretrained_kwargs,
        )

        # Calculate feature dimension
        self._cached_feature_dim = self._get_dimension_from_name(model_name)

        # Enable gradient checkpointing if requested (only useful for full fine-tuning;
        # frozen backbone uses torch.no_grad() in _extract_features instead)
        freeze_backbone = getattr(self.config.model, 'freeze_backbone', False)
        if getattr(self.config.model, 'gradient_checkpointing', False) and not freeze_backbone:
            self.backbone.gradient_checkpointing_enable(
                gradient_checkpointing_kwargs={"use_reentrant": False}
            )
            print("Gradient checkpointing enabled")

        # Setup predictor for future prediction
        self._setup_predictor()

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

            # Optionally unfreeze the last N transformer blocks
            n_unfreeze = getattr(self.config.model, 'unfreeze_last_n_blocks', 0)
            if n_unfreeze > 0:
                # Locate transformer block list (handles different HF architectures)
                encoder_layers = None
                if hasattr(self.backbone, 'encoder') and hasattr(self.backbone.encoder, 'layer'):
                    encoder_layers = self.backbone.encoder.layer
                elif hasattr(self.backbone, 'encoder') and hasattr(self.backbone.encoder, 'layers'):
                    encoder_layers = self.backbone.encoder.layers
                elif hasattr(self.backbone, 'blocks'):
                    encoder_layers = self.backbone.blocks

                if encoder_layers is not None:
                    total_blocks = len(encoder_layers)
                    n_unfreeze = min(n_unfreeze, total_blocks)
                    for block in encoder_layers[total_blocks - n_unfreeze:]:
                        for param in block.parameters():
                            param.requires_grad = True
                    print(f"Unfroze last {n_unfreeze}/{total_blocks} transformer blocks")

                # Always unfreeze final layer norm (output of last block)
                for attr in ('layernorm', 'norm', 'final_layer_norm'):
                    if hasattr(self.backbone, attr):
                        for param in getattr(self.backbone, attr).parameters():
                            param.requires_grad = True
                        break

    def _get_dimension_from_name(self, model_name: str) -> int:
        """Get feature dimension from model name."""
        name_lower = model_name.lower()
        for key, dim in self.DIMENSION_MAP.items():
            if key in name_lower:
                return dim
        return 1408  # Default to ViT-G

    def _setup_predictor(self):
        """Setup V-JEPA2 predictor for future prediction."""
        if not self.config.model.use_future_prediction:
            # Remove predictor to save memory
            if hasattr(self.backbone, "predictor"):
                del self.backbone.predictor
            return

        if not hasattr(self.backbone, "predictor"):
            self.config.model.use_future_prediction = False
            return

        self.predictor = self.backbone.predictor
        self.has_predictor = True

        # Register fusion weights here so the parameter is always visible to
        # the optimizer and torch.compile before the first forward pass.
        if getattr(self.config.model, 'predictor_combination_method', 'concat') == 'weighted_sum':
            self._fusion_weights = nn.Parameter(torch.tensor([0.7, 0.3]))

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

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

        Returns:
            Features (B, seq_len, feature_dim)
        """
        extraction_layer = getattr(self.config.model, 'extraction_layer', None)
        freeze_backbone = getattr(self.config.model, 'freeze_backbone', False)
        # When backbone is frozen, use no_grad to prevent activation storage (saves ~40+ GB on ViT-G)
        backbone_ctx = torch.no_grad() if freeze_backbone else torch.enable_grad()

        if extraction_layer is not None:
            # Extract from a specific intermediate encoder layer.
            # Future-prediction is skipped: the predictor was trained on the final
            # encoder output and would produce nonsense from intermediate layers.
            with backbone_ctx:
                output = self.backbone(
                    pixel_values_videos=x,
                    skip_predictor=True,
                    output_hidden_states=True,
                )
            return output.hidden_states[extraction_layer]

        # Default path: use last_hidden_state with optional future prediction.
        with backbone_ctx:
            output = self.backbone(pixel_values_videos=x, skip_predictor=True)

        if hasattr(output, "last_hidden_state"):
            features = output.last_hidden_state
        else:
            features = output

        # Apply future prediction if enabled
        if self.has_predictor and self.config.model.use_future_prediction:
            future_features = self._predict_future(features)
            features = self._combine_features(features, future_features)

        return features

    def _predict_future(self, encoder_features: torch.Tensor) -> torch.Tensor:
        """
        Predict future features using V-JEPA2 predictor.

        Args:
            encoder_features: (B, seq_len, D)

        Returns:
            Future features (B, future_tokens, D)
        """
        B, seq_len, D = encoder_features.shape

        # Calculate future tokens needed
        future_seconds = self.config.model.future_prediction_seconds
        original_fps = self.config.model.original_fps
        present_seconds = self.config.data.frame_count / original_fps
        future_tokens = max(1, int(seq_len * (future_seconds / present_seconds)))

        device = encoder_features.device
        dtype = encoder_features.dtype
        expanded_len = seq_len + future_tokens

        # Create masks
        context_indices = torch.arange(seq_len, device=device).unsqueeze(0)
        target_indices = torch.arange(seq_len, expanded_len, device=device).unsqueeze(0)

        # Process each sample (predictor expects specific format)
        results = []
        for i in range(B):
            sample = encoder_features[i:i+1]
            padding = torch.zeros(1, future_tokens, D, device=device, dtype=dtype)
            expanded = torch.cat([sample, padding], dim=1)

            output = self.predictor(
                expanded,
                context_mask=[context_indices],
                target_mask=[target_indices],
            )
            results.append(output.last_hidden_state)

        return torch.cat(results, dim=0)

    def _combine_features(
        self,
        present: torch.Tensor,
        future: torch.Tensor,
    ) -> torch.Tensor:
        """
        Combine present and future features.

        Args:
            present: (B, seq_len, D)
            future: (B, future_tokens, D)

        Returns:
            Combined features
        """
        method = self.config.model.predictor_combination_method

        if method == "concat":
            return torch.cat([present, future], dim=1)

        elif method == "weighted_sum":
            min_len = min(present.size(1), future.size(1))
            present = present[:, :min_len, :]
            future = future[:, :min_len, :]

            weights = torch.softmax(self._fusion_weights, dim=0)
            return weights[0] * present + weights[1] * future

        else:
            # Default to concat
            return torch.cat([present, future], dim=1)
