"""Model architecture for SignSegmenterV2.

Ported from /dev/4TB_SSD/segmentor_v2/models_Wilor.py
"""

from __future__ import annotations

import math

import torch
import torch.nn as nn
from torch import Tensor


class TemporalConv(nn.Module):
    """1D temporal convolution with batch normalization and ReLU.

    Applies convolution across the temporal dimension to extract
    local temporal features from the input sequence.
    """

    def __init__(self, input_size: int, hidden_size: int, conv_type: int = 2):
        """Initialize temporal convolution layer.

        Args:
            input_size: Input feature dimension
            hidden_size: Output feature dimension
            conv_type: Kernel size (2, 3, or 5). Default 2 = simple linear projection.
        """
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.conv_type = conv_type

        # If conv_type=2, just use linear projection (no temporal convolution)
        # Otherwise use actual convolution
        if self.conv_type == 2:
            # Simple linear projection to match expected output size
            # Use LayerNorm instead of BatchNorm for 3D tensors
            self.temporal_conv = nn.Sequential(
                nn.Linear(input_size, hidden_size),
                nn.LayerNorm(hidden_size),
                nn.ReLU(inplace=True),
            )
        else:
            if self.conv_type == 3:
                self.kernel_size = ["K3"]
            elif self.conv_type == 5:
                self.kernel_size = ["K5"]
            else:
                self.kernel_size = []

            modules = []
            for layer_idx, ks in enumerate(self.kernel_size):
                input_sz = self.input_size if layer_idx == 0 else self.hidden_size
                if ks[0] == "P":
                    modules.append(nn.MaxPool1d(kernel_size=int(ks[1]), ceil_mode=False))
                elif ks[0] == "K":
                    modules.append(
                        nn.Conv1d(
                            input_sz, self.hidden_size, kernel_size=int(ks[1]), stride=1, padding=0
                        )
                    )
                    modules.append(nn.BatchNorm1d(self.hidden_size))
                    modules.append(nn.ReLU(inplace=True))
            self.temporal_conv = nn.Sequential(*modules)

    def forward(self, x: Tensor) -> Tensor:
        """Apply temporal convolution.

        Args:
            x: Input tensor of shape (batch, seq_len, features)

        Returns:
            Output tensor of shape (batch, seq_len, hidden_size)
        """
        if self.conv_type == 2:
            # Linear projection, no permutation needed
            return self.temporal_conv(x)
        else:
            # 1D convolution, need to permute
            # Change to (batch, channels, seq_len)
            x = self.temporal_conv(x.permute(0, 2, 1))
            # Change back to (batch, seq_len, features)
            return x.permute(0, 2, 1)


class PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding.

    Pre-computes position encodings and adds them to input embeddings
    in the forward pass.

    Implementation based on OpenNMT-py.
    """

    def __init__(self, size: int = 0, max_len: int = 5000):
        """Initialize positional encoding.

        Args:
            size: Embedding dimension size (must be even)
            max_len: Maximum sequence length

        Raises:
            ValueError: If size is odd
        """
        if size % 2 != 0:
            raise ValueError(
                f"Cannot use sin/cos positional encoding with odd dim (got dim={size})"
            )

        super().__init__()

        pe = torch.zeros(max_len, size)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, size, 2, dtype=torch.float) * -(math.log(10000.0) / size)
        )
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)  # shape: (1, max_len, size)

        self.register_buffer("pe", pe)
        self.dim = size

    def forward(self, emb: Tensor) -> Tensor:
        """Add positional encodings to embeddings.

        Args:
            emb: Input embeddings of shape (batch, seq_len, dim)

        Returns:
            Positionally encoded embeddings
        """
        return emb + self.pe[:, : emb.size(1)]


class TransformerEncoderV2(nn.Module):
    """Sign segmentation transformer encoder (v2 architecture).

    Uses WiLoR features (192-dim) split into left/right hands,
    applies temporal convolutions and hand-specific MLPs, then
    processes through a transformer encoder.
    """

    def __init__(
        self,
        input_size: int = 192,
        hidden_size: int = 512,
        num_layers: int = 4,
        num_heads: int = 4,
        ff_size: int = 2048,
        dropout: float = 0.1,
        temp_conv_kernel: int = 2,
        num_classes: int = 3,
    ):
        """Initialize transformer encoder.

        Args:
            input_size: Input feature dimension (default 192 for WiLoR)
            hidden_size: Transformer hidden dimension
            num_layers: Number of transformer encoder layers
            num_heads: Number of attention heads
            ff_size: Feed-forward layer size
            dropout: Dropout probability
            temp_conv_kernel: Temporal convolution kernel size
            num_classes: Number of output classes (0=out, 1=in, 2=begin)
        """
        super().__init__()

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.temp_conv_kernel = temp_conv_kernel

        # Adapter size is half of hidden size (left + right = hidden_size)
        self.adapter_size = hidden_size // 2

        # Split input into two hands (each hand is 96 dims for WiLoR)
        two_stream_size = input_size // 2

        # Temporal convolutions for each hand
        self.left_conv = TemporalConv(
            two_stream_size, self.adapter_size, conv_type=temp_conv_kernel
        )
        self.right_conv = TemporalConv(
            two_stream_size, self.adapter_size, conv_type=temp_conv_kernel
        )

        # Hand-specific MLPs (3 linear layers with ReLU)
        self.left_hand = nn.Sequential(
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
        )

        self.right_hand = nn.Sequential(
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
        )

        # Positional encoding
        self.pe = PositionalEncoding(self.adapter_size)

        # Transformer encoder
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=num_heads,
            dim_feedforward=ff_size,
            dropout=dropout,
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        # Output projection to class predictions
        self.sign_layer = nn.Linear(hidden_size, num_classes)

    def forward(
        self,
        src: Tensor,
        src_length: Tensor,
        src_mask: Tensor | None = None,
    ) -> Tensor:
        """Forward pass through the model.

        Args:
            src: Input features of shape (batch, seq_len+2*padding, 192)
            src_length: Actual sequence lengths (batch,)
            src_mask: Attention mask (batch, 1, seq_len)

        Returns:
            Class predictions of shape (batch, seq_len, num_classes)
        """
        # Split into left/right hands: (batch, seq_len+2*pad, 2, 96)
        # The input src includes padding added before this call
        src_dual_hand = src.reshape(src.shape[0], src.shape[1], 2, 96)

        # Apply temporal convolutions to each hand
        # For conv_type=3: kernel=3, padding=0 reduces length by 2, compensating for the +2 padding
        # For conv_type=2: linear layer, we need to trim the padding manually after
        temp_lh = self.left_conv(src_dual_hand[:, :, 0, :])
        temp_rh = self.right_conv(src_dual_hand[:, :, 1, :])

        # For linear projection (conv_type=2), trim padding manually
        temp_pad = (self.temp_conv_kernel - 1) // 2
        if self.temp_conv_kernel == 2 and temp_pad > 0:
            temp_lh = temp_lh[:, temp_pad:-temp_pad, :]
            temp_rh = temp_rh[:, temp_pad:-temp_pad, :]

        # Apply hand-specific MLPs
        lh = self.left_hand(temp_lh)  # (batch, seq_len, adapter_size)
        rh = self.right_hand(temp_rh)  # (batch, seq_len, adapter_size)

        # Add positional encoding to each hand separately
        lh = self.pe(lh)
        rh = self.pe(rh)

        # Concatenate left and right hands
        all_h = torch.cat((lh, rh), dim=-1)  # (batch, seq_len, hidden_size)

        # Create attention mask for transformer
        # Convert src_mask from (batch, 1, seq_len) to (batch, seq_len)
        if src_mask is not None:
            attn_mask = ~src_mask.squeeze(1)  # Invert: True = masked
        else:
            attn_mask = None

        # Pass through transformer encoder
        encoder_output = self.encoder(all_h, src_key_padding_mask=attn_mask)

        # Project to class predictions
        sign_prediction = self.sign_layer(encoder_output)

        return sign_prediction
