#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/pp_ocrv6_small_det/modular_pp_ocrv6_small_det.py.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          modular_pp_ocrv6_small_det.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2026 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

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

from ...activations import ACT2FN
from ...backbone_utils import load_backbone
from ...modeling_layers import GradientCheckpointingLayer
from ...modeling_outputs import BaseModelOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...processing_utils import Unpack
from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
from .configuration_pp_ocrv6_small_det import PPOCRV6SmallDetConfig


class PPOCRV6SmallDetConvBatchnormLayer(nn.Module):
    """
    A basic wrapper for Convolution-BatchNorm-Activation, typically used for head components.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = 1,
        padding: int | str = 1,
        groups: int = 1,
        activation: str = "relu",
        bias: bool = False,
        convolution_transpose: bool = False,
    ):
        super().__init__()
        if convolution_transpose:
            self.convolution = nn.ConvTranspose2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                stride=stride,
            )
        else:
            self.convolution = nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                groups=groups,
                bias=bias,
            )

        self.norm = nn.BatchNorm2d(out_channels)
        self.act_fn = nn.Identity() if activation is None else ACT2FN[activation]

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.convolution(hidden_states)
        hidden_states = self.norm(hidden_states)
        hidden_states = self.act_fn(hidden_states)
        return hidden_states


class PPOCRV6SmallDetHead(nn.Module):
    """
    Standard segmentation head for generating probability maps. It uses transposed
    convolution to upsample the feature map back to the original image size.
    """

    def __init__(
        self,
        config: PPOCRV6SmallDetConfig,
    ):
        super().__init__()

        in_channels = config.neck_out_channels
        kernel_list = config.kernel_list
        self.conv_down = PPOCRV6SmallDetConvBatchnormLayer(
            in_channels=in_channels,
            out_channels=in_channels // 4,
            kernel_size=kernel_list[0],
            padding=int(kernel_list[0] // 2),
        )
        self.conv_up = PPOCRV6SmallDetConvBatchnormLayer(
            in_channels=in_channels // 4,
            out_channels=in_channels // 4,
            kernel_size=kernel_list[1],
            stride=2,
            convolution_transpose=True,
        )

        self.conv_final = nn.ConvTranspose2d(
            in_channels=in_channels // 4,
            out_channels=1,
            kernel_size=kernel_list[2],
            stride=2,
        )

    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        hidden_states = self.conv_down(hidden_states)
        hidden_states = self.conv_up(hidden_states)
        hidden_states = self.conv_final(hidden_states)
        hidden_states = torch.sigmoid(hidden_states)
        return hidden_states


class PPOCRV6SmallDetSqueezeExcitationModule(nn.Module):
    """
    Simplified Squeeze-and-Excitation (SE) Module for the neck network.
    Applies channel-wise recalibration with a clamped activation to stabilize training.
    """

    def __init__(self, in_channels, reduction, activation="relu"):
        super().__init__()

        self.avg_pool = nn.AdaptiveAvgPool2d(output_size=1)
        self.conv1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=in_channels // reduction,
            kernel_size=1,
            stride=1,
            padding=0,
        )
        self.conv2 = nn.Conv2d(
            in_channels=in_channels // reduction,
            out_channels=in_channels,
            kernel_size=1,
            stride=1,
            padding=0,
        )
        self.act_fn = ACT2FN[activation]

    def forward(self, hidden_states):
        residual = hidden_states
        hidden_states = self.avg_pool(hidden_states)
        hidden_states = self.conv2(self.act_fn(self.conv1(hidden_states)))
        hidden_states = torch.clamp(0.2 * hidden_states + 0.5, min=0.0, max=1.0)
        return residual * hidden_states


class PPOCRV6SmallDetDepthwiseSeparableConvLayer(GradientCheckpointingLayer):
    """
    The differences from PPLCNetDepthwiseSeparableConvLayer are:
    1. Uses standard 2D convolutions instead of the original custom convolution layer.
    2. Has slightly different default settings.
    3. Adds a residual connection at the end.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        reduction,
    ):
        super().__init__()
        self.depthwise_convolution = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=1,
            padding=kernel_size // 2,
            groups=in_channels,
            bias=True,
        )
        self.squeeze_excitation_module = PPOCRV6SmallDetSqueezeExcitationModule(out_channels // 4, reduction)
        self.pointwise_convolution = nn.Conv2d(
            in_channels=out_channels,
            out_channels=out_channels // 4,
            kernel_size=1,
            bias=False,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.depthwise_convolution(hidden_states)
        hidden_states = self.pointwise_convolution(hidden_states)
        hidden_states = hidden_states + self.squeeze_excitation_module(hidden_states)

        return hidden_states


class PPOCRV6SmallDetResidualSqueezeExcitationLayer(nn.Module):
    """
    Residual Squeeze-and-Excitation (RSE) Layer for the neck network.
    Combines a 1x1/3x3 convolution with an SE Module.
    """

    def __init__(self, in_channels, out_channels, kernel_size, reduction):
        super().__init__()
        self.in_conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            padding=int(kernel_size // 2),
            bias=False,
        )
        self.squeeze_excitation_block = PPOCRV6SmallDetSqueezeExcitationModule(out_channels, reduction)

    def forward(self, hidden_states):
        hidden_states = self.in_conv(hidden_states)
        hidden_states = hidden_states + self.squeeze_excitation_block(hidden_states)

        return hidden_states


class PPOCRV6SmallDetNeck(nn.Module):
    """
    The only difference from PPOCRV5MobileDetNeck is the module used by input_conv.
    """

    def __init__(self, config: PPOCRV6SmallDetConfig):
        super().__init__()
        self.interpolate_mode = config.interpolate_mode

        self.insert_conv = nn.ModuleList()
        self.input_conv = nn.ModuleList()
        for i in range(len(config.layer_list_out_channels)):
            self.insert_conv.append(
                PPOCRV6SmallDetResidualSqueezeExcitationLayer(
                    config.layer_list_out_channels[i],
                    config.neck_out_channels,
                    1,
                    config.reduction,
                )
            )
            self.input_conv.append(
                PPOCRV6SmallDetDepthwiseSeparableConvLayer(
                    config.neck_out_channels, config.neck_out_channels, config.dilated_kernel_size, config.reduction
                )
            )

    def forward(self, feature_maps):
        fused = []
        for conv, feature in zip(self.insert_conv, feature_maps):  # [p2, p3, p4, p5]
            hidden_states = conv(feature)
            fused.append(hidden_states)

        for i in range(2, -1, -1):  # p4 -> p3-> p2
            fused[i] = fused[i] + F.interpolate(fused[i + 1], scale_factor=2, mode=self.interpolate_mode)

        features = []
        for conv, feat in zip(self.input_conv, [fused[0], fused[1], fused[2], fused[3]]):
            features.append(conv(feat))

        processed = []
        upsample_scales = [1, 2, 4, 8]  # p2, p3, p4, p5
        for feat, scale in zip(features, upsample_scales):
            if scale != 1:
                hidden_states = F.interpolate(feat, scale_factor=scale, mode=self.interpolate_mode)
            else:
                hidden_states = feat
            processed.append(hidden_states)

        fused_feature_map = torch.cat(processed[::-1], dim=1)  # [p5, p4, p3, p2]
        return fused_feature_map


class PPOCRV6SmallDetPreTrainedModel(PreTrainedModel):
    """
    Base class for all PPOCRV5 Server Det pre-trained models. Handles model initialization,
    configuration, and loading of pre-trained weights, following the Transformers library conventions.
    """

    config: PPOCRV6SmallDetConfig
    base_model_prefix = "pp_ocrv6_small_det"
    main_input_name = "pixel_values"
    input_modalities = ("image",)
    _can_compile_fullgraph = True


@auto_docstring
class PPOCRV6SmallDetModel(PPOCRV6SmallDetPreTrainedModel):
    def __init__(self, config: PPOCRV6SmallDetConfig):
        super().__init__(config)
        self.backbone = load_backbone(config)
        self.neck = PPOCRV6SmallDetNeck(config)
        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.FloatTensor] | BaseModelOutputWithNoAttention:
        backbone_outputs = self.backbone(pixel_values, **kwargs)
        hidden_state = backbone_outputs.feature_maps
        hidden_state = self.neck(hidden_state)

        return BaseModelOutputWithNoAttention(
            last_hidden_state=hidden_state,
            hidden_states=backbone_outputs.hidden_states,
        )


@auto_docstring(custom_intro="PPOCR6SmallRec model for text recognition tasks.")
class PPOCRV6SmallDetForObjectDetection(PPOCRV6SmallDetPreTrainedModel):
    _keys_to_ignore_on_load_missing = ["num_batches_tracked"]

    def __init__(self, config: PPOCRV6SmallDetConfig):
        super().__init__(config)
        self.model = PPOCRV6SmallDetModel(config)
        self.head = PPOCRV6SmallDetHead(config)
        self.post_init()

    @can_return_tuple
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.FloatTensor] | BaseModelOutputWithNoAttention:
        outputs = self.model(pixel_values, **kwargs)
        logits = self.head(outputs.last_hidden_state)

        return BaseModelOutputWithNoAttention(
            last_hidden_state=logits,
            hidden_states=outputs.hidden_states,
        )


__all__ = ["PPOCRV6SmallDetForObjectDetection", "PPOCRV6SmallDetModel", "PPOCRV6SmallDetPreTrainedModel"]
