"""Small, timm-compatible EfficientNetV2-S inference implementation.

The Apache-2.0 timm implementation was reduced to the layers exercised by
``tf_efficientnetv2_s.in21k_ft_in1k`` and modified to remove its runtime
dependency. Module names intentionally match timm so the official UTMOS v2
state dictionary loads; see LICENSE.
"""

from __future__ import annotations

import math

import torch
import torch.nn.functional as F
from torch import nn


class Conv2dSame(nn.Conv2d):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        height, width = x.shape[-2:]
        stride_h, stride_w = self.stride
        kernel_h, kernel_w = self.kernel_size
        dilation_h, dilation_w = self.dilation
        output_h = math.ceil(height / stride_h)
        output_w = math.ceil(width / stride_w)
        pad_h = max(
            (output_h - 1) * stride_h
            + (kernel_h - 1) * dilation_h
            + 1
            - height,
            0,
        )
        pad_w = max(
            (output_w - 1) * stride_w
            + (kernel_w - 1) * dilation_w
            + 1
            - width,
            0,
        )
        if pad_h or pad_w:
            x = F.pad(
                x,
                (
                    pad_w // 2,
                    pad_w - pad_w // 2,
                    pad_h // 2,
                    pad_h - pad_h // 2,
                ),
            )
        return F.conv2d(
            x,
            self.weight,
            self.bias,
            self.stride,
            0,
            self.dilation,
            self.groups,
        )


class BatchNormAct2d(nn.BatchNorm2d):
    def __init__(self, channels: int, *, activate: bool) -> None:
        super().__init__(channels, eps=1e-3, momentum=0.1)
        self.activate = activate

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = super().forward(x)
        return F.silu(x, inplace=True) if self.activate else x


def conv3x3(
    in_channels: int,
    out_channels: int,
    stride: int,
    *,
    groups: int = 1,
) -> nn.Conv2d:
    if stride == 2:
        return Conv2dSame(
            in_channels,
            out_channels,
            3,
            stride=2,
            groups=groups,
            bias=False,
        )
    return nn.Conv2d(
        in_channels,
        out_channels,
        3,
        stride=1,
        padding=1,
        groups=groups,
        bias=False,
    )


class ConvBnAct(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv = conv3x3(24, 24, 1)
        self.bn1 = BatchNormAct2d(24, activate=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.bn1(self.conv(x))
        return x + residual


class EdgeResidual(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        expansion: int,
        stride: int,
    ) -> None:
        super().__init__()
        expanded_channels = in_channels * expansion
        self.conv_exp = conv3x3(in_channels, expanded_channels, stride)
        self.bn1 = BatchNormAct2d(expanded_channels, activate=True)
        self.conv_pwl = nn.Conv2d(expanded_channels, out_channels, 1, bias=False)
        self.bn2 = BatchNormAct2d(out_channels, activate=False)
        self.has_residual = stride == 1 and in_channels == out_channels

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.bn1(self.conv_exp(x))
        x = self.bn2(self.conv_pwl(x))
        return x + residual if self.has_residual else x


class SqueezeExcite(nn.Module):
    def __init__(self, expanded_channels: int, reduced_channels: int) -> None:
        super().__init__()
        self.conv_reduce = nn.Conv2d(expanded_channels, reduced_channels, 1)
        self.conv_expand = nn.Conv2d(reduced_channels, expanded_channels, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        scale = x.mean((2, 3), keepdim=True)
        scale = F.silu(self.conv_reduce(scale), inplace=True)
        return x * torch.sigmoid(self.conv_expand(scale))


class InvertedResidual(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        expansion: int,
        stride: int,
    ) -> None:
        super().__init__()
        expanded_channels = in_channels * expansion
        self.conv_pw = nn.Conv2d(in_channels, expanded_channels, 1, bias=False)
        self.bn1 = BatchNormAct2d(expanded_channels, activate=True)
        self.conv_dw = conv3x3(
            expanded_channels,
            expanded_channels,
            stride,
            groups=expanded_channels,
        )
        self.bn2 = BatchNormAct2d(expanded_channels, activate=True)
        self.se = SqueezeExcite(expanded_channels, in_channels // 4)
        self.conv_pwl = nn.Conv2d(expanded_channels, out_channels, 1, bias=False)
        self.bn3 = BatchNormAct2d(out_channels, activate=False)
        self.has_residual = stride == 1 and in_channels == out_channels

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.bn1(self.conv_pw(x))
        x = self.bn2(self.conv_dw(x))
        x = self.se(x)
        x = self.bn3(self.conv_pwl(x))
        return x + residual if self.has_residual else x


def make_stage(
    block_type: type[EdgeResidual | InvertedResidual],
    in_channels: int,
    out_channels: int,
    expansion: int,
    repeats: int,
    stride: int,
) -> nn.Sequential:
    blocks = [block_type(in_channels, out_channels, expansion, stride)]
    blocks.extend(
        block_type(out_channels, out_channels, expansion, 1)
        for _ in range(repeats - 1)
    )
    return nn.Sequential(*blocks)


class EfficientNetV2S(nn.Module):
    """Feature-only TF EfficientNetV2-S with timm-compatible parameter names."""

    def __init__(self) -> None:
        super().__init__()
        self.conv_stem = Conv2dSame(3, 24, 3, stride=2, bias=False)
        self.bn1 = BatchNormAct2d(24, activate=True)
        self.blocks = nn.Sequential(
            nn.Sequential(ConvBnAct(), ConvBnAct()),
            make_stage(EdgeResidual, 24, 48, 4, 4, 2),
            make_stage(EdgeResidual, 48, 64, 4, 4, 2),
            make_stage(InvertedResidual, 64, 128, 4, 6, 2),
            make_stage(InvertedResidual, 128, 160, 6, 9, 1),
            make_stage(InvertedResidual, 160, 256, 6, 15, 2),
        )
        self.conv_head = nn.Conv2d(256, 1280, 1, bias=False)
        self.bn2 = BatchNormAct2d(1280, activate=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.bn1(self.conv_stem(x))
        x = self.blocks(x)
        x = self.bn2(self.conv_head(x))
        return x