#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/pp_ocrv6_medium_det/modular_pp_ocrv6_medium_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_medium_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_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_medium_det import PPOCRV6MediumDetConfig


class PPOCRV6MediumDetConvBatchnormLayer(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 PPOCRV6MediumDetHead(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: PPOCRV6MediumDetConfig,
    ):
        super().__init__()

        in_channels = config.neck_out_channels
        kernel_list = config.kernel_list
        self.conv_down = PPOCRV6MediumDetConvBatchnormLayer(
            in_channels=in_channels,
            out_channels=in_channels // 4,
            kernel_size=kernel_list[0],
            padding=int(kernel_list[0] // 2),
        )
        self.conv_up = PPOCRV6MediumDetConvBatchnormLayer(
            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 PPOCRV6MediumDetIntraclassBlock(nn.Module):
    """
    Intra-Class Relationship Block. It uses multi-scale convolution (7x7, 5x5, 3x3)
    and asymmetric kernels (e.g., 7x1, 1x7) to capture long-range spatial dependencies
    within text regions.
    """

    def __init__(
        self,
        intraclass_block_config: dict | None = None,
        in_channels: int = 96,
        reduce_factor: int = 4,
    ):
        super().__init__()

        reduced_channels = in_channels // reduce_factor

        self.conv_reduce_channel = nn.Conv2d(in_channels, reduced_channels, *intraclass_block_config["reduce_channel"])

        self.vertical_long_to_small_conv_longratio = nn.Conv2d(
            reduced_channels, reduced_channels, *intraclass_block_config["vertical_long_to_small_conv_longratio"]
        )
        self.vertical_long_to_small_conv_midratio = nn.Conv2d(
            reduced_channels, reduced_channels, *intraclass_block_config["vertical_long_to_small_conv_midratio"]
        )
        self.vertical_long_to_small_conv_shortratio = nn.Conv2d(
            reduced_channels, reduced_channels, *intraclass_block_config["vertical_long_to_small_conv_shortratio"]
        )

        self.horizontal_small_to_long_conv_longratio = nn.Conv2d(
            reduced_channels, reduced_channels, *intraclass_block_config["horizontal_small_to_long_conv_longratio"]
        )
        self.horizontal_small_to_long_conv_midratio = nn.Conv2d(
            reduced_channels, reduced_channels, *intraclass_block_config["horizontal_small_to_long_conv_midratio"]
        )
        self.horizontal_small_to_long_conv_shortratio = nn.Conv2d(
            reduced_channels, reduced_channels, *intraclass_block_config["horizontal_small_to_long_conv_shortratio"]
        )

        self.symmetric_conv_long_longratio = nn.Conv2d(
            reduced_channels, reduced_channels, *intraclass_block_config["symmetric_conv_long_longratio"]
        )
        self.symmetric_conv_long_midratio = nn.Conv2d(
            reduced_channels, reduced_channels, *intraclass_block_config["symmetric_conv_long_midratio"]
        )
        self.symmetric_conv_long_shortratio = nn.Conv2d(
            reduced_channels, reduced_channels, *intraclass_block_config["symmetric_conv_long_shortratio"]
        )

        self.conv_final = PPOCRV6MediumDetConvBatchnormLayer(
            in_channels=reduced_channels,
            out_channels=in_channels,
            kernel_size=intraclass_block_config["return_channel"][0],
            stride=intraclass_block_config["return_channel"][1],
            padding=intraclass_block_config["return_channel"][2],
            bias=True,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.conv_reduce_channel(hidden_states)

        hidden_states = (
            self.symmetric_conv_long_longratio(hidden_states)
            + self.vertical_long_to_small_conv_longratio(hidden_states)
            + self.horizontal_small_to_long_conv_longratio(hidden_states)
        )
        hidden_states = (
            self.symmetric_conv_long_midratio(hidden_states)
            + self.vertical_long_to_small_conv_midratio(hidden_states)
            + self.horizontal_small_to_long_conv_midratio(hidden_states)
        )
        hidden_states = (
            self.symmetric_conv_long_shortratio(hidden_states)
            + self.vertical_long_to_small_conv_shortratio(hidden_states)
            + self.horizontal_small_to_long_conv_shortratio(hidden_states)
        )

        hidden_states = self.conv_final(hidden_states)

        return residual + hidden_states


class PPOCRV6MediumDetNeck(nn.Module):
    """
    Large Kernel Path Aggregation Network (Neck) for PPOCRV5 Server Detection.
    Fuses multi-scale features from backbone stages (stage 2 to stage 5) via top-down and bottom-up paths,
    enhanced with large kernel convolution for better spatial dependency modeling.
    """

    def __init__(self, config):
        super().__init__()
        self.interpolate_mode = config.interpolate_mode
        self.scale_factor_list = config.scale_factor_list
        self.num_backbone_stages = len(config.backbone_config.stage_out_channels)

        self.input_channel_adjustment_convolution = nn.ModuleList()
        self.input_feature_projection_convolution = nn.ModuleList()
        self.path_aggregation_head_convolution = nn.ModuleList()
        self.path_aggregation_lateral_convolution = nn.ModuleList()

        backbone_stage_output_channels = config.backbone_config.stage_out_channels

        for backbone_stage_index in range(len(backbone_stage_output_channels)):
            channel_adjustment_convolution = nn.Conv2d(
                in_channels=backbone_stage_output_channels[backbone_stage_index],
                out_channels=config.neck_out_channels,
                kernel_size=1,
                bias=False,
            )
            self.input_channel_adjustment_convolution.append(channel_adjustment_convolution)

            # [Key Change] bias=True
            feature_projection_convolution = nn.Conv2d(
                in_channels=config.neck_out_channels,
                out_channels=config.neck_out_channels // 4,
                kernel_size=9,
                padding=4,
                bias=True,
            )
            self.input_feature_projection_convolution.append(feature_projection_convolution)

            if backbone_stage_index > 0:
                pan_head_convolution = nn.Conv2d(
                    in_channels=config.neck_out_channels // 4,
                    out_channels=config.neck_out_channels // 4,
                    kernel_size=3,
                    padding=1,
                    stride=2,
                    bias=False,
                )
                self.path_aggregation_head_convolution.append(pan_head_convolution)

            # [Key Change] bias=True
            pan_lateral_convolution = nn.Conv2d(
                in_channels=config.neck_out_channels // 4,
                out_channels=config.neck_out_channels // 4,
                kernel_size=9,
                padding=4,
                bias=True,
            )
            self.path_aggregation_lateral_convolution.append(pan_lateral_convolution)

        self.intraclass_blocks = nn.ModuleList()
        for _ in range(config.intraclass_block_number):
            self.intraclass_blocks.append(
                PPOCRV6MediumDetIntraclassBlock(
                    config.intraclass_block_config, config.neck_out_channels // 4, reduce_factor=config.reduce_factor
                )
            )

    def forward(self, backbone_stage_feature_maps: list[torch.Tensor], **kwargs) -> torch.Tensor:
        channel_adjusted = []
        for i, feature_map in enumerate(backbone_stage_feature_maps):
            hidden_states = self.input_channel_adjustment_convolution[i](feature_map)
            channel_adjusted.append(hidden_states)

        top_down = [None] * self.num_backbone_stages
        top_down[3] = channel_adjusted[3]
        for i in range(self.num_backbone_stages - 2, -1, -1):
            top_down[i] = channel_adjusted[i] + F.interpolate(
                top_down[i + 1], scale_factor=2, mode=self.interpolate_mode
            )

        projected = []
        for i in range(self.num_backbone_stages):
            hidden_states = top_down[i] if i < self.num_backbone_stages - 1 else channel_adjusted[-1]
            hidden_states = self.input_feature_projection_convolution[i](hidden_states)
            projected.append(hidden_states)

        bottom_up = [None] * self.num_backbone_stages
        bottom_up[0] = projected[0]
        for i in range(1, self.num_backbone_stages):
            bottom_up[i] = projected[i] + self.path_aggregation_head_convolution[i - 1](bottom_up[i - 1])

        lateral_refined = []
        for i in range(self.num_backbone_stages):
            hidden_states = projected[0] if i == 0 else bottom_up[i]
            hidden_states = self.path_aggregation_lateral_convolution[i](hidden_states)
            lateral_refined.append(hidden_states)

        intraclass_refined = [block(feature) for block, feature in zip(self.intraclass_blocks, lateral_refined)]

        upsampled = []
        for feature, scale_factor in zip(intraclass_refined, self.scale_factor_list):
            if scale_factor > 1:
                hidden_states = F.interpolate(feature, scale_factor=scale_factor, mode=self.interpolate_mode)
            else:
                hidden_states = feature
            upsampled.append(hidden_states)

        upsampled = [
            F.interpolate(feature, scale_factor=scale_factor, mode=self.interpolate_mode)
            if scale_factor > 1
            else feature
            for feature, scale_factor in zip(intraclass_refined, self.scale_factor_list)
        ]

        return torch.cat(upsampled[::-1], dim=1)


class PPOCRV6MediumDetPreTrainedModel(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: PPOCRV6MediumDetConfig
    base_model_prefix = "pp_ocrv6_medium_det"
    main_input_name = "pixel_values"
    input_modalities = ("image",)
    _can_compile_fullgraph = True


@auto_docstring
class PPOCRV6MediumDetModel(PPOCRV6MediumDetPreTrainedModel):
    def __init__(self, config: PPOCRV6MediumDetConfig):
        super().__init__(config)
        self.backbone = load_backbone(config)
        self.neck = PPOCRV6MediumDetNeck(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="PPOCRV6MediumDet model for text detection tasks.")
class PPOCRV6MediumDetForObjectDetection(PPOCRV6MediumDetPreTrainedModel):
    _keys_to_ignore_on_load_missing = ["num_batches_tracked"]

    def __init__(self, config: PPOCRV6MediumDetConfig):
        super().__init__(config)
        self.model = PPOCRV6MediumDetModel(config)
        self.head = PPOCRV6MediumDetHead(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__ = ["PPOCRV6MediumDetForObjectDetection", "PPOCRV6MediumDetModel", "PPOCRV6MediumDetPreTrainedModel"]
