"""
PNASNet, adapted from
https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/pnasnet.py
and
https://github.com/tensorflow/models/blob/master/research/slim/nets/nasnet/pnasnet.py

Paper "Progressive Neural Architecture Search", https://arxiv.org/abs/1712.00559

Changes from original:
* Removed dynamic padding

Generated by gpt-5.3-codex xhigh.
"""

# Reference license: Apache-2.0 (both)

from functools import partial
from typing import Any
from typing import Optional

import torch
import torch.nn.functional as F
from torch import nn
from torchvision.ops import Conv2dNormActivation

from birder.model_registry import registry
from birder.net.base import BaseNet


def _calc_reduction_layers(num_cells: int, num_reduction_layers: int) -> list[int]:
    reduction_layers = []
    for pool_num in range(1, num_reduction_layers + 1):
        layer_num = int((pool_num / (num_reduction_layers + 1)) * num_cells)
        reduction_layers.append(layer_num)

    return reduction_layers


class SeparableConv2d(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int, padding: int) -> None:
        super().__init__()
        self.depthwise_conv2d = nn.Conv2d(
            in_channels,
            in_channels,
            kernel_size=(kernel_size, kernel_size),
            stride=(stride, stride),
            padding=(padding, padding),
            groups=in_channels,
            bias=False,
        )
        self.pointwise_conv2d = nn.Conv2d(
            in_channels, out_channels, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=False
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.depthwise_conv2d(x)
        x = self.pointwise_conv2d(x)
        return x


class BranchSeparable(nn.Module):
    def __init__(
        self, in_channels: int, out_channels: int, kernel_size: int, stride: int, stem_cell: bool = False
    ) -> None:
        super().__init__()
        padding = kernel_size // 2
        if stem_cell is True:
            middle_channels = out_channels
        else:
            middle_channels = in_channels

        self.act_1 = nn.ReLU()
        self.separable_1 = SeparableConv2d(
            in_channels,
            middle_channels,
            kernel_size,
            stride=stride,
            padding=padding,
        )
        self.bn_sep_1 = nn.BatchNorm2d(middle_channels, eps=0.001)

        self.act_2 = nn.ReLU()
        self.separable_2 = SeparableConv2d(
            middle_channels,
            out_channels,
            kernel_size,
            stride=1,
            padding=padding,
        )
        self.bn_sep_2 = nn.BatchNorm2d(out_channels, eps=0.001)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.act_1(x)
        x = self.separable_1(x)
        x = self.bn_sep_1(x)
        x = self.act_2(x)
        x = self.separable_2(x)
        x = self.bn_sep_2(x)

        return x


class ActConvBn(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int, padding: int) -> None:
        super().__init__()
        self.act = nn.ReLU()
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=(kernel_size, kernel_size),
            stride=(stride, stride),
            padding=(padding, padding),
            bias=False,
        )
        self.bn = nn.BatchNorm2d(out_channels, eps=0.001)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.act(x)
        x = self.conv(x)
        x = self.bn(x)
        return x


class FactorizedReduction(nn.Module):
    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.act = nn.ReLU()

        self.path_1_avg_pool = nn.AvgPool2d(
            kernel_size=(1, 1),
            stride=(2, 2),
            count_include_pad=False,
        )
        self.path_1_conv = nn.Conv2d(
            in_channels,
            out_channels // 2,
            kernel_size=(1, 1),
            stride=(1, 1),
            padding=(0, 0),
            bias=False,
        )

        self.path_2_avg_pool = nn.AvgPool2d(
            kernel_size=(1, 1),
            stride=(2, 2),
            count_include_pad=False,
        )
        final_filter_size = (out_channels // 2) + (out_channels % 2)
        self.path_2_conv = nn.Conv2d(
            in_channels,
            final_filter_size,
            kernel_size=(1, 1),
            stride=(1, 1),
            padding=(0, 0),
            bias=False,
        )
        self.final_path_bn = nn.BatchNorm2d(out_channels, eps=0.001)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.act(x)

        x_path1 = self.path_1_avg_pool(x)
        x_path1 = self.path_1_conv(x_path1)

        x_path2 = x[:, :, 1:, 1:]
        x_path2 = F.pad(x_path2, pad=(0, 1, 0, 1))
        x_path2 = self.path_2_avg_pool(x_path2)
        x_path2 = self.path_2_conv(x_path2)

        out = self.final_path_bn(torch.concat((x_path1, x_path2), dim=1))

        return out


# pylint: disable=abstract-method
class CellBase(nn.Module):
    def cell_forward(self, x_left: torch.Tensor, x_right: torch.Tensor) -> torch.Tensor:
        x_comb_iter_0_left = self.comb_iter_0_left(x_left)
        x_comb_iter_0_right = self.comb_iter_0_right(x_left)
        x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right

        x_comb_iter_1_left = self.comb_iter_1_left(x_right)
        x_comb_iter_1_right = self.comb_iter_1_right(x_right)
        x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right

        x_comb_iter_2_left = self.comb_iter_2_left(x_right)
        x_comb_iter_2_right = self.comb_iter_2_right(x_right)
        x_comb_iter_2 = x_comb_iter_2_left + x_comb_iter_2_right

        x_comb_iter_3_left = self.comb_iter_3_left(x_comb_iter_2)
        x_comb_iter_3_right = self.comb_iter_3_right(x_right)
        x_comb_iter_3 = x_comb_iter_3_left + x_comb_iter_3_right

        x_comb_iter_4_left = self.comb_iter_4_left(x_left)
        x_comb_iter_4_right = self.comb_iter_4_right(x_right)
        x_comb_iter_4 = x_comb_iter_4_left + x_comb_iter_4_right

        x_out = torch.concat((x_comb_iter_0, x_comb_iter_1, x_comb_iter_2, x_comb_iter_3, x_comb_iter_4), dim=1)

        return x_out


class CellStem0(CellBase):
    def __init__(self, in_chs_left: int, out_chs_left: int, in_chs_right: int, out_chs_right: int) -> None:
        super().__init__()
        self.conv_1x1 = ActConvBn(in_chs_right, out_chs_right, kernel_size=1, stride=1, padding=0)

        self.comb_iter_0_left = BranchSeparable(in_chs_left, out_chs_left, kernel_size=5, stride=2, stem_cell=True)
        self.comb_iter_0_right = nn.Sequential(
            nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)),
            nn.Conv2d(in_chs_left, out_chs_left, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=False),
            nn.BatchNorm2d(out_chs_left, eps=0.001),
        )

        self.comb_iter_1_left = BranchSeparable(out_chs_right, out_chs_right, kernel_size=7, stride=2)
        self.comb_iter_1_right = nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))

        self.comb_iter_2_left = BranchSeparable(out_chs_right, out_chs_right, kernel_size=5, stride=2)
        self.comb_iter_2_right = BranchSeparable(out_chs_right, out_chs_right, kernel_size=3, stride=2)

        self.comb_iter_3_left = BranchSeparable(out_chs_right, out_chs_right, kernel_size=3, stride=1)
        self.comb_iter_3_right = nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))

        self.comb_iter_4_left = BranchSeparable(in_chs_right, out_chs_right, kernel_size=3, stride=2, stem_cell=True)
        self.comb_iter_4_right = ActConvBn(out_chs_right, out_chs_right, kernel_size=1, stride=2, padding=0)

    def forward(self, x_left: torch.Tensor) -> torch.Tensor:
        x_right = self.conv_1x1(x_left)
        x_out = self.cell_forward(x_left, x_right)

        return x_out


class Cell(CellBase):
    def __init__(
        self,
        in_chs_left: int,
        out_chs_left: int,
        in_chs_right: int,
        out_chs_right: int,
        is_reduction: bool,
        match_prev_layer_dims: bool,
    ) -> None:
        super().__init__()
        stride = 2 if is_reduction is True else 1

        if match_prev_layer_dims is True:
            self.conv_prev_1x1 = FactorizedReduction(in_chs_left, out_chs_left)
        else:
            self.conv_prev_1x1 = ActConvBn(in_chs_left, out_chs_left, kernel_size=1, stride=1, padding=0)

        self.conv_1x1 = ActConvBn(in_chs_right, out_chs_right, kernel_size=1, stride=1, padding=0)

        self.comb_iter_0_left = BranchSeparable(out_chs_left, out_chs_left, kernel_size=5, stride=stride)
        self.comb_iter_0_right = nn.MaxPool2d(kernel_size=(3, 3), stride=(stride, stride), padding=(1, 1))

        self.comb_iter_1_left = BranchSeparable(out_chs_right, out_chs_right, kernel_size=7, stride=stride)
        self.comb_iter_1_right = nn.MaxPool2d(kernel_size=(3, 3), stride=(stride, stride), padding=(1, 1))

        self.comb_iter_2_left = BranchSeparable(out_chs_right, out_chs_right, kernel_size=5, stride=stride)
        self.comb_iter_2_right = BranchSeparable(out_chs_right, out_chs_right, kernel_size=3, stride=stride)

        self.comb_iter_3_left = BranchSeparable(out_chs_right, out_chs_right, kernel_size=3, stride=1)
        self.comb_iter_3_right = nn.MaxPool2d(kernel_size=(3, 3), stride=(stride, stride), padding=(1, 1))

        self.comb_iter_4_left = BranchSeparable(out_chs_left, out_chs_left, kernel_size=3, stride=stride)
        if is_reduction is True:
            self.comb_iter_4_right = ActConvBn(out_chs_right, out_chs_right, kernel_size=1, stride=stride, padding=0)
        else:
            self.comb_iter_4_right = nn.Identity()

    def forward(self, x_left: torch.Tensor, x_right: torch.Tensor) -> torch.Tensor:
        x_left = self.conv_prev_1x1(x_left)
        x_right = self.conv_1x1(x_right)
        x_out = self.cell_forward(x_left, x_right)

        return x_out


class PNASNet(BaseNet):
    default_size = (331, 331)

    def __init__(
        self,
        input_channels: int,
        num_classes: int,
        *,
        config: Optional[dict[str, Any]] = None,
        size: Optional[tuple[int, int]] = None,
    ) -> None:
        super().__init__(input_channels, num_classes, config=config, size=size)
        assert self.config is not None, "must set config"

        num_reduction_layers = 2
        num_cells: int = self.config["num_cells"]
        num_conv_filters: int = self.config["num_conv_filters"]
        stem_multiplier: float = self.config["stem_multiplier"]
        dropout_rate: float = self.config["dropout_rate"]

        reduction_indices = _calc_reduction_layers(num_cells, num_reduction_layers)

        stem_channels = int(32 * stem_multiplier)
        stem0_filters = int(num_conv_filters * 0.25)
        stem1_filters = int(num_conv_filters * 0.5)
        norm_layer = partial(nn.BatchNorm2d, eps=0.001, momentum=0.1)
        self.stem = Conv2dNormActivation(
            self.input_channels,
            stem_channels,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(0, 0),
            bias=False,
            norm_layer=norm_layer,
            activation_layer=None,
        )
        self.cell_stem_0 = CellStem0(
            in_chs_left=stem_channels,
            out_chs_left=stem0_filters,
            in_chs_right=stem_channels,
            out_chs_right=stem0_filters,
        )
        self.cell_stem_1 = Cell(
            in_chs_left=stem_channels,
            out_chs_left=stem1_filters,
            in_chs_right=5 * stem0_filters,
            out_chs_right=stem1_filters,
            is_reduction=True,
            match_prev_layer_dims=True,
        )

        filter_scaling = 1.0
        left_channels = 5 * stem0_filters
        right_channels = 5 * stem1_filters
        cells: list[nn.Module] = []
        for cell_idx in range(num_cells):
            is_reduction = cell_idx in reduction_indices
            if is_reduction is True:
                filter_scaling *= 2.0

            filters = int(num_conv_filters * filter_scaling)
            match_prev_layer_dims = (cell_idx == 0) or ((cell_idx - 1) in reduction_indices)
            cells.append(
                Cell(
                    in_chs_left=left_channels,
                    out_chs_left=filters,
                    in_chs_right=right_channels,
                    out_chs_right=filters,
                    is_reduction=is_reduction,
                    match_prev_layer_dims=match_prev_layer_dims,
                )
            )
            left_channels = right_channels
            right_channels = 5 * filters

        self.cells = nn.ModuleList(cells)
        self.act = nn.ReLU()
        self.features = nn.Sequential(
            nn.AdaptiveAvgPool2d(output_size=(1, 1)),
            nn.Flatten(1),
            nn.Dropout(p=dropout_rate, inplace=True),
        )
        self.embedding_size = right_channels
        self.classifier = self.create_classifier()

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        x_conv_0 = self.stem(x)
        x_stem_0 = self.cell_stem_0(x_conv_0)
        prev = x_stem_0
        curr = self.cell_stem_1(x_conv_0, x_stem_0)

        for cell in self.cells:
            out = cell(prev, curr)
            prev = curr
            curr = out

        return self.act(curr)

    def embedding(self, x: torch.Tensor) -> torch.Tensor:
        x = self.forward_features(x)
        return self.features(x)


registry.register_model_config(
    "pnasnet_mobile",
    PNASNet,
    config={"num_cells": 9, "num_conv_filters": 54, "stem_multiplier": 1.0, "dropout_rate": 0.0},
)
registry.register_model_config(
    "pnasnet_large",
    PNASNet,
    config={"num_cells": 12, "num_conv_filters": 216, "stem_multiplier": 3.0, "dropout_rate": 0.5},
)
