"""
SHViT, adapted from
https://github.com/ysj9909/SHViT/blob/main/model/shvit.py

Paper "SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design",
https://arxiv.org/abs/2401.16456

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

# Reference license: MIT

from collections import OrderedDict
from typing import Any
from typing import Optional

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

from birder.model_registry import registry
from birder.net.base import DetectorBackbone
from birder.net.base import make_divisible


class Conv2dBN(nn.Sequential):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: tuple[int, int],
        stride: tuple[int, int],
        padding: tuple[int, int],
        groups: int = 1,
        bn_weight_init: float = 1.0,
        reparameterized: bool = False,
    ) -> None:
        super().__init__()
        self.c: nn.Module
        self.reparameterized = reparameterized
        self.add_module(
            "c",
            nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                groups=groups,
                bias=reparameterized,
            ),
        )
        if reparameterized is False:
            self.add_module("bn", nn.BatchNorm2d(out_channels))
            nn.init.constant_(self.bn.weight, bn_weight_init)
            nn.init.zeros_(self.bn.bias)

    def reparameterize(self) -> None:
        if self.reparameterized is True:
            return

        c, bn = self._modules.values()
        w = bn.weight / (bn.running_var + bn.eps) ** 0.5
        w = c.weight * w[:, None, None, None]
        b = bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5
        self.c = nn.Conv2d(
            w.size(1) * self.c.groups,
            w.size(0),
            kernel_size=w.shape[2:],
            stride=self.c.stride,
            padding=self.c.padding,
            dilation=self.c.dilation,
            groups=self.c.groups,
            device=c.weight.device,
            dtype=c.weight.dtype,
        )
        self.c.weight.data.copy_(w)
        self.c.bias.data.copy_(b)

        # Delete un-used branches
        for param in self.parameters():
            param.detach_()

        del self.bn
        self.reparameterized = True


class NormLinear(nn.Sequential):
    def __init__(self, in_dim: int, out_dim: int, reparameterized: bool) -> None:
        super().__init__()
        self.li: nn.Module
        self.reparameterized = reparameterized
        if reparameterized is False:
            self.add_module("bn", nn.BatchNorm1d(in_dim))

        self.add_module("li", nn.Linear(in_dim, out_dim))
        nn.init.trunc_normal_(self.li.weight, std=0.02)
        nn.init.zeros_(self.li.bias)

    def reparameterize(self) -> None:
        if self.reparameterized is True:
            return

        bn, li = self._modules.values()
        w = bn.weight / (bn.running_var + bn.eps) ** 0.5
        b = bn.bias - self.bn.running_mean * self.bn.weight / (bn.running_var + bn.eps) ** 0.5
        w = li.weight * w[None, :]
        if li.bias is None:
            b = b @ self.li.weight.T
        else:
            b = (li.weight @ b[:, None]).view(-1) + self.li.bias

        self.li = nn.Linear(w.size(1), w.size(0), device=li.weight.device, dtype=li.weight.dtype)
        self.li.weight.data.copy_(w)
        self.li.bias.data.copy_(b)

        # Delete un-used branches
        for param in self.parameters():
            param.detach_()

        del self.bn
        self.reparameterized = True


class Residual(nn.Module):
    def __init__(self, module: nn.Module, reparameterized: bool) -> None:
        super().__init__()
        self.m = module
        self.reparameterized = reparameterized
        self.fused = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.fused is True:
            return self.m(x)

        return x + self.m(x)

    def reparameterize(self) -> None:
        if self.reparameterized is True:
            return

        if hasattr(self.m, "reparameterize") is True:
            self.m.reparameterize()

        if isinstance(self.m, Conv2dBN):
            conv = self.m.c

            identity = torch.ones(
                conv.weight.size(0), conv.weight.size(1), 1, 1, device=conv.weight.device, dtype=conv.weight.dtype
            )
            identity = F.pad(
                identity,
                [
                    conv.kernel_size[1] // 2,
                    conv.kernel_size[1] // 2,
                    conv.kernel_size[0] // 2,
                    conv.kernel_size[0] // 2,
                ],
            )
            conv.weight.data.add_(identity)

            # Delete un-used branches
            for param in self.parameters():
                param.detach_()

            self.m = conv
            self.fused = True

        self.reparameterized = True


class PatchMerging(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, reparameterized: bool) -> None:
        super().__init__()
        hidden_dim = int(in_dim * 4)
        self.conv1 = Conv2dBN(
            in_dim,
            hidden_dim,
            kernel_size=(1, 1),
            stride=(1, 1),
            padding=(0, 0),
            reparameterized=reparameterized,
        )
        self.act1 = nn.ReLU()
        self.conv2 = Conv2dBN(
            hidden_dim,
            hidden_dim,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(1, 1),
            groups=hidden_dim,
            reparameterized=reparameterized,
        )
        self.act2 = nn.ReLU()
        self.se = SqueezeExcitation(hidden_dim, make_divisible(hidden_dim // 4, 8))
        self.conv3 = Conv2dBN(
            hidden_dim,
            out_dim,
            kernel_size=(1, 1),
            stride=(1, 1),
            padding=(0, 0),
            reparameterized=reparameterized,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv1(x)
        x = self.act1(x)
        x = self.conv2(x)
        x = self.act2(x)
        x = self.se(x)
        x = self.conv3(x)

        return x


class FFN(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, reparameterized: bool) -> None:
        super().__init__()
        self.pw1 = Conv2dBN(
            dim,
            hidden_dim,
            kernel_size=(1, 1),
            stride=(1, 1),
            padding=(0, 0),
            reparameterized=reparameterized,
        )
        self.act = nn.ReLU()
        self.pw2 = Conv2dBN(
            hidden_dim,
            dim,
            kernel_size=(1, 1),
            stride=(1, 1),
            padding=(0, 0),
            bn_weight_init=0.0,
            reparameterized=reparameterized,
        )

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

        return x


class SHSA(nn.Module):
    """
    Single-Head Self-Attention
    """

    def __init__(self, dim: int, qk_dim: int, partial_dim: int, reparameterized: bool) -> None:
        super().__init__()
        self.scale = qk_dim**-0.5
        self.qk_dim = qk_dim
        self.dim = dim
        self.partial_dim = partial_dim
        self.pre_norm = nn.GroupNorm(1, partial_dim)
        self.qkv = Conv2dBN(
            partial_dim,
            qk_dim * 2 + partial_dim,
            kernel_size=(1, 1),
            stride=(1, 1),
            padding=(0, 0),
            reparameterized=reparameterized,
        )
        self.proj = nn.Sequential(
            nn.ReLU(),
            Conv2dBN(
                dim,
                dim,
                kernel_size=(1, 1),
                stride=(1, 1),
                padding=(0, 0),
                bn_weight_init=0.0,
                reparameterized=reparameterized,
            ),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, _, h, w = x.shape
        x1, x2 = torch.split(x, [self.partial_dim, self.dim - self.partial_dim], dim=1)
        x1 = self.pre_norm(x1)

        qkv = self.qkv(x1)
        q, k, v = torch.split(qkv, [self.qk_dim, self.qk_dim, self.partial_dim], dim=1)
        q = q.flatten(2)
        k = k.flatten(2)
        v = v.flatten(2)

        attn = (q.transpose(-2, -1) @ k) * self.scale
        attn = attn.softmax(dim=-1)
        x1 = (v @ attn.transpose(-2, -1)).reshape(b, self.partial_dim, h, w)
        x = self.proj(torch.concat([x1, x2], dim=1))

        return x


class BasicBlock(nn.Module):
    def __init__(self, dim: int, qk_dim: int, partial_dim: int, use_shsa: bool, reparameterized: bool) -> None:
        super().__init__()
        self.conv = Residual(
            Conv2dBN(
                dim,
                dim,
                kernel_size=(3, 3),
                stride=(1, 1),
                padding=(1, 1),
                groups=dim,
                bn_weight_init=0.0,
                reparameterized=reparameterized,
            ),
            reparameterized,
        )
        if use_shsa is True:
            self.mixer = Residual(SHSA(dim, qk_dim, partial_dim, reparameterized), reparameterized)
        else:
            self.mixer = nn.Identity()

        self.ffn = Residual(FFN(dim, int(dim * 2), reparameterized), reparameterized)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv(x)
        x = self.mixer(x)
        x = self.ffn(x)

        return x


class StageBlock(nn.Module):
    def __init__(
        self, prev_dim: int, dim: int, qk_dim: int, partial_dim: int, use_shsa: bool, depth: int, reparameterized: bool
    ) -> None:
        super().__init__()
        if prev_dim != dim:
            self.downsample = nn.Sequential(
                Residual(
                    Conv2dBN(
                        prev_dim,
                        prev_dim,
                        kernel_size=(3, 3),
                        stride=(1, 1),
                        padding=(1, 1),
                        groups=prev_dim,
                        reparameterized=reparameterized,
                    ),
                    reparameterized,
                ),
                Residual(FFN(prev_dim, int(prev_dim * 2), reparameterized), reparameterized),
                PatchMerging(prev_dim, dim, reparameterized),
                Residual(
                    Conv2dBN(
                        dim,
                        dim,
                        kernel_size=(3, 3),
                        stride=(1, 1),
                        padding=(1, 1),
                        groups=dim,
                        reparameterized=reparameterized,
                    ),
                    reparameterized,
                ),
                Residual(FFN(dim, int(dim * 2), reparameterized), reparameterized),
            )
        else:
            self.downsample = nn.Identity()

        self.blocks = nn.Sequential(
            *[BasicBlock(dim, qk_dim, partial_dim, use_shsa, reparameterized) for _ in range(depth)]
        )

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

        return x


class SHViT(DetectorBackbone):
    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"
        assert "head_bias" not in self.config, "head_bias customization is not supported"
        assert "mlp_head" not in self.config, "mlp_head customization is not supported"

        self.reparameterized = False
        qk_dims = [16, 16, 16]
        use_shsa = [False, True, True]
        embed_dims: list[int] = self.config["embed_dims"]
        depths: list[int] = self.config["depths"]
        partial_dims: list[int] = self.config["partial_dims"]

        stem_dim = embed_dims[0]

        self.stem = nn.Sequential(
            Conv2dBN(
                self.input_channels,
                stem_dim // 8,
                kernel_size=(3, 3),
                stride=(2, 2),
                padding=(1, 1),
                reparameterized=self.reparameterized,
            ),
            nn.ReLU(),
            Conv2dBN(
                stem_dim // 8,
                stem_dim // 4,
                kernel_size=(3, 3),
                stride=(2, 2),
                padding=(1, 1),
                reparameterized=self.reparameterized,
            ),
            nn.ReLU(),
            Conv2dBN(
                stem_dim // 4,
                stem_dim // 2,
                kernel_size=(3, 3),
                stride=(2, 2),
                padding=(1, 1),
                reparameterized=self.reparameterized,
            ),
            nn.ReLU(),
            Conv2dBN(
                stem_dim // 2,
                stem_dim,
                kernel_size=(3, 3),
                stride=(2, 2),
                padding=(1, 1),
                reparameterized=self.reparameterized,
            ),
        )

        stages: OrderedDict[str, nn.Module] = OrderedDict()
        return_channels: list[int] = []
        prev_dim = stem_dim
        for idx, (dim, depth, partial_dim, qk_dim, use_shsa_stage) in enumerate(
            zip(embed_dims, depths, partial_dims, qk_dims, use_shsa, strict=True)
        ):
            stages[f"stage{idx + 1}"] = StageBlock(
                prev_dim=prev_dim,
                dim=dim,
                qk_dim=qk_dim,
                partial_dim=partial_dim,
                use_shsa=use_shsa_stage,
                depth=depth,
                reparameterized=self.reparameterized,
            )
            return_channels.append(dim)
            prev_dim = dim

        self.body = nn.Sequential(stages)
        self.features = nn.Sequential(
            nn.AdaptiveAvgPool2d(output_size=(1, 1)),
            nn.Flatten(1),
        )
        self.return_stages = self.return_stages[: len(depths)]
        self.return_channels = return_channels
        self.embedding_size = embed_dims[-1]
        self.classifier = self.create_classifier()

    def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        x = self.stem(x)

        out = {}
        for name, module in self.body.named_children():
            x = module(x)
            if name in self.return_stages:
                out[name] = x

        return out

    def freeze_stages(self, up_to_stage: int) -> None:
        for param in self.stem.parameters():
            param.requires_grad_(False)

        for idx, module in enumerate(self.body.children()):
            if idx >= up_to_stage:
                break

            for param in module.parameters():
                param.requires_grad_(False)

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        return self.body(x)

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

    def create_classifier(
        self, embed_dim: Optional[int] = None, head_bias: Optional[bool] = None, mlp_head: Optional[bool] = None
    ) -> nn.Module:
        assert head_bias is None, "head_bias customization is not supported"
        assert mlp_head is None, "mlp_head customization is not supported"

        if self.num_classes == 0:
            return nn.Identity()

        if embed_dim is None:
            embed_dim = self.embedding_size

        return NormLinear(embed_dim, self.num_classes, self.reparameterized)

    @torch.no_grad()  # type: ignore[untyped-decorator]
    def reparameterize_model(self) -> None:
        for module in self.modules():
            if hasattr(module, "reparameterize") is True:
                module.reparameterize()

        self.reparameterized = True


registry.register_model_config(
    "shvit_s1",
    SHViT,
    config={"embed_dims": [128, 224, 320], "depths": [2, 4, 5], "partial_dims": [32, 48, 68]},
)
registry.register_model_config(
    "shvit_s2",
    SHViT,
    config={"embed_dims": [128, 308, 448], "depths": [2, 4, 5], "partial_dims": [32, 66, 96]},
)
registry.register_model_config(
    "shvit_s3",
    SHViT,
    config={"embed_dims": [192, 352, 448], "depths": [3, 5, 5], "partial_dims": [48, 75, 96]},
)
registry.register_model_config(
    "shvit_s4",
    SHViT,
    config={"embed_dims": [224, 336, 448], "depths": [4, 7, 6], "partial_dims": [48, 72, 96]},
)
