"""
Export-friendly model wrappers for BADAS.

Provides wrappers that make models compatible with ONNX/TensorRT export
by removing dynamic operations and torch.compile.
"""

import torch
import torch.nn as nn
from typing import Optional

from badas.core.base import VideoClassifierBase


class ExportWrapper(nn.Module):
    """
    Generic export wrapper for BADAS video classifiers.

    Wraps a VideoClassifierBase model for ONNX/TensorRT export by:
    - Removing torch.compile (if present)
    - Simplifying the forward pass
    - Optionally including/excluding the predictor module

    Args:
        model: A VideoClassifierBase model to wrap.
        include_predictor: Whether to include future prediction.
            Defaults to False for simpler, more reliable export.

    Example:
        >>> from badas.inference.checkpoint import load_model_from_checkpoint
        >>> from badas.export.wrappers import ExportWrapper
        >>>
        >>> model, config = load_model_from_checkpoint("model.ckpt")
        >>> export_model = ExportWrapper(model)
        >>> torch.onnx.export(export_model, dummy_input, "model.onnx")
    """

    def __init__(
        self,
        model: VideoClassifierBase,
        include_predictor: bool = False,
    ):
        super().__init__()

        # Unwrap torch.compile if present
        if hasattr(model, '_orig_mod'):
            model = model._orig_mod

        self.backbone = model.backbone
        self.temporal_processor = model.temporal_processor
        self.classifier = model.classifier

        # Predictor handling
        self.include_predictor = include_predictor and hasattr(model, 'predictor')
        if self.include_predictor:
            self.predictor = model.predictor
            self.has_predictor = getattr(model, 'has_predictor', False)
        else:
            self.has_predictor = False

    def forward(self, pixel_values_videos: torch.Tensor) -> torch.Tensor:
        """
        Forward pass optimized for export.

        Args:
            pixel_values_videos: Input video tensor (B, T, C, H, W).

        Returns:
            Logits tensor (B,) for binary classification.
        """
        # Extract encoder features
        output = self.backbone(
            pixel_values_videos=pixel_values_videos,
            skip_predictor=True,
        )
        features = output.last_hidden_state

        # Temporal pooling
        features = self.temporal_processor(features)

        # Classification
        logits = self.classifier(features)

        # Squeeze for binary classification
        if logits.dim() > 1 and logits.size(-1) == 1:
            logits = logits.squeeze(-1)

        return logits


class ExportWrapperWithPredictor(nn.Module):
    """
    Export wrapper that includes the V-JEPA2 predictor module.

    This wrapper is more complex and may have export compatibility issues
    with some ONNX/TensorRT versions. Use ExportWrapper for simpler export.

    Args:
        model: A VideoClassifierBase model with predictor.
        frame_count: Number of input frames.
        future_seconds: Seconds to predict into the future.
        original_fps: Original video FPS.

    Example:
        >>> wrapper = ExportWrapperWithPredictor(
        ...     model,
        ...     frame_count=16,
        ...     future_seconds=1.0,
        ...     original_fps=8,
        ... )
    """

    def __init__(
        self,
        model: VideoClassifierBase,
        frame_count: int = 16,
        future_seconds: float = 1.0,
        original_fps: float = 8.0,
    ):
        super().__init__()

        if hasattr(model, '_orig_mod'):
            model = model._orig_mod

        self.backbone = model.backbone
        self.temporal_processor = model.temporal_processor
        self.classifier = model.classifier

        if not hasattr(model, 'predictor') or model.predictor is None:
            raise ValueError("Model does not have a predictor module")

        self.predictor = model.predictor
        self.frame_count = frame_count
        self.future_seconds = future_seconds
        self.original_fps = original_fps

    def forward(self, pixel_values_videos: torch.Tensor) -> torch.Tensor:
        """Forward pass with future prediction."""
        # Extract encoder features
        output = self.backbone(
            pixel_values_videos=pixel_values_videos,
            skip_predictor=True,
        )
        features = output.last_hidden_state  # (B, seq_len, D)

        # Future prediction
        features = self._predict_future(features)

        # Temporal pooling
        features = self.temporal_processor(features)

        # Classification
        logits = self.classifier(features)

        if logits.dim() > 1 and logits.size(-1) == 1:
            logits = logits.squeeze(-1)

        return logits

    def _predict_future(self, encoder_features: torch.Tensor) -> torch.Tensor:
        """Vectorized future prediction for export."""
        B, seq_len, D = encoder_features.shape
        device = encoder_features.device
        dtype = encoder_features.dtype

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

        # Create masks
        context_mask = torch.arange(seq_len, device=device).unsqueeze(0)
        target_mask = torch.arange(
            seq_len, seq_len + future_tokens, device=device
        ).unsqueeze(0)

        # Expand features with zero padding
        zeros_padding = torch.zeros(B, future_tokens, D, device=device, dtype=dtype)
        expanded_features = torch.cat([encoder_features, zeros_padding], dim=1)

        # Call predictor
        predicted_output = self.predictor(
            expanded_features,
            context_mask=[context_mask],
            target_mask=[target_mask],
        )
        future_features = predicted_output.last_hidden_state

        # Combine present and future
        return torch.cat([encoder_features, future_features], dim=1)


def create_export_wrapper(
    model: VideoClassifierBase,
    include_predictor: bool = False,
    frame_count: Optional[int] = None,
    future_seconds: Optional[float] = None,
    original_fps: Optional[float] = None,
) -> nn.Module:
    """
    Factory function to create the appropriate export wrapper.

    Args:
        model: The model to wrap.
        include_predictor: Whether to include predictor (default: False).
        frame_count: Required if include_predictor=True.
        future_seconds: Required if include_predictor=True.
        original_fps: Required if include_predictor=True.

    Returns:
        Export-ready nn.Module.

    Example:
        >>> wrapper = create_export_wrapper(model)  # Simple wrapper
        >>> wrapper = create_export_wrapper(       # With predictor
        ...     model,
        ...     include_predictor=True,
        ...     frame_count=16,
        ...     future_seconds=1.0,
        ...     original_fps=8,
        ... )
    """
    if not include_predictor:
        return ExportWrapper(model, include_predictor=False)

    if frame_count is None or future_seconds is None or original_fps is None:
        raise ValueError(
            "frame_count, future_seconds, and original_fps are required "
            "when include_predictor=True"
        )

    return ExportWrapperWithPredictor(
        model,
        frame_count=frame_count,
        future_seconds=future_seconds,
        original_fps=original_fps,
    )


__all__ = [
    'ExportWrapper',
    'ExportWrapperWithPredictor',
    'create_export_wrapper',
]
