"""
VOLO, adapted from
https://github.com/sail-sg/volo

"VOLO: Vision Outlooker for Visual Recognition", https://arxiv.org/abs/2106.13112

Changes from original:
* Removed aux head and token labeling (mix_token)

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

# Reference license: Apache-2.0

import math
from typing import Any
from typing import Optional

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

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


class OutlookAttention(nn.Module):
    def __init__(self, dim: int, num_heads: int, kernel_size: int, padding: int, stride: int) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.kernel_size = kernel_size
        self.padding = padding
        self.stride = stride

        head_dim = dim // num_heads
        self.scale = head_dim**-0.5

        self.v = nn.Linear(dim, dim, bias=False)
        self.attn = nn.Linear(dim, kernel_size**4 * num_heads)
        self.proj = nn.Linear(dim, dim)

        self.unfold = nn.Unfold(kernel_size=kernel_size, padding=padding, stride=stride)
        self.pool = nn.AvgPool2d(kernel_size=stride, stride=stride, ceil_mode=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, H, W, C = x.shape
        k = self.kernel_size
        h = math.ceil(H / self.stride)
        w = math.ceil(W / self.stride)

        # Compute values and unfold to local windows
        v = self.v(x).permute(0, 3, 1, 2)  # (B, C, H, W)
        v = self.unfold(v)  # (B, C*k*k, h*w)
        v = v.reshape(B, self.num_heads, C // self.num_heads, k * k, h * w)
        v = v.permute(0, 1, 4, 3, 2)  # (B, heads, h*w, k*k, head_dim)

        # Generate attention weights from pooled features
        attn = self.pool(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)  # (B, h, w, C)
        attn = self.attn(attn).reshape(B, h * w, self.num_heads, k * k, k * k)
        attn = attn.permute(0, 2, 1, 3, 4)  # (B, heads, h*w, k*k, k*k)
        attn = attn * self.scale
        attn = attn.softmax(dim=-1)

        # Apply attention and fold back
        x = (attn @ v).permute(0, 1, 4, 3, 2)
        x = x.reshape(B, C * k * k, h * w)
        x = F.fold(x, output_size=(H, W), kernel_size=k, padding=self.padding, stride=self.stride)

        x = self.proj(x.permute(0, 2, 3, 1))

        return x


class Attention(nn.Module):
    def __init__(self, dim: int, num_heads: int) -> None:
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim**-0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=False)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, H, W, C = x.shape
        N = H * W
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)

        x = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, scale=self.scale)  # pylint: disable=not-callable
        x = x.transpose(1, 2).reshape(B, H, W, C)
        x = self.proj(x)

        return x


class OutlookerBlock(nn.Module):
    def __init__(
        self, dim: int, num_heads: int, kernel_size: int, padding: int, stride: int, mlp_ratio: float, drop_path: float
    ) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = OutlookAttention(dim, num_heads, kernel_size, padding, stride)
        self.drop_path1 = StochasticDepth(drop_path, mode="row")
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = MLP(dim, [int(dim * mlp_ratio), dim], activation_layer=nn.GELU)
        self.drop_path2 = StochasticDepth(drop_path, mode="row")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.drop_path1(self.attn(self.norm1(x)))
        x = x + self.drop_path2(self.mlp(self.norm2(x)))
        return x


class TransformerBlock(nn.Module):
    def __init__(self, dim: int, num_heads: int, mlp_ratio: float, drop_path: float) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = Attention(dim, num_heads)
        self.drop_path1 = StochasticDepth(drop_path, mode="row")
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = MLP(dim, [int(dim * mlp_ratio), dim], activation_layer=nn.GELU)
        self.drop_path2 = StochasticDepth(drop_path, mode="row")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.drop_path1(self.attn(self.norm1(x)))
        x = x + self.drop_path2(self.mlp(self.norm2(x)))
        return x


class ClassAttention(nn.Module):
    def __init__(self, dim: int, num_heads: int) -> None:
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim**-0.5

        self.q = nn.Linear(dim, dim, bias=False)
        self.kv = nn.Linear(dim, dim * 2, bias=False)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.shape
        q = self.q(x[:, :1]).reshape(B, 1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
        kv = self.kv(x).reshape(B, N, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        k, v = kv.unbind(0)

        x_cls = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, scale=self.scale)  # pylint: disable=not-callable
        x_cls = x_cls.transpose(1, 2).reshape(B, 1, C)
        x_cls = self.proj(x_cls)

        return x_cls


class ClassBlock(nn.Module):
    def __init__(self, dim: int, num_heads: int, mlp_ratio: float, drop_path: float) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = ClassAttention(dim, num_heads)
        self.drop_path1 = StochasticDepth(drop_path, mode="row")
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = MLP(dim, [int(dim * mlp_ratio), dim], activation_layer=nn.GELU)
        self.drop_path2 = StochasticDepth(drop_path, mode="row")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        cls_embed = x[:, :1]
        cls_embed = cls_embed + self.drop_path1(self.attn(self.norm1(x)))
        cls_embed = cls_embed + self.drop_path2(self.mlp(self.norm2(cls_embed)))
        return torch.concat([cls_embed, x[:, 1:]], dim=1)


class Stem(nn.Module):
    def __init__(self, in_channels: int, hidden_dim: int, embed_dim: int, stem_stride: int, patch_size: int) -> None:
        super().__init__()
        proj_kernel = patch_size // stem_stride
        self.conv = nn.Sequential(
            Conv2dNormActivation(
                in_channels,
                hidden_dim,
                kernel_size=(7, 7),
                stride=(stem_stride, stem_stride),
                padding=(3, 3),
                bias=False,
            ),
            Conv2dNormActivation(hidden_dim, hidden_dim, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
            Conv2dNormActivation(hidden_dim, hidden_dim, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
        )
        self.proj = nn.Conv2d(
            hidden_dim, embed_dim, kernel_size=(proj_kernel, proj_kernel), stride=(proj_kernel, proj_kernel)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv(x)
        x = self.proj(x)
        x = x.permute(0, 2, 3, 1)  # NCHW -> NHWC

        return x


class Downsample(nn.Module):
    def __init__(self, in_dim: int, out_dim: int) -> None:
        super().__init__()
        self.proj = nn.Conv2d(in_dim, out_dim, kernel_size=(2, 2), stride=(2, 2))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.permute(0, 3, 1, 2)  # NHWC -> NCHW
        x = self.proj(x)
        x = x.permute(0, 2, 3, 1)  # NCHW -> NHWC

        return x


class VOLO(BaseNet):
    # pylint: disable=too-many-locals
    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"

        patch_size = 8
        stem_stride = 2
        outlook_stride = 2
        outlook_kernel = 3
        num_post_layers = 2
        layers: list[int] = self.config["layers"]
        embed_dims: list[int] = self.config["embed_dims"]
        num_heads: list[int] = self.config["num_heads"]
        mlp_ratio: float = self.config["mlp_ratio"]
        stem_hidden_dim: int = self.config["stem_hidden_dim"]
        drop_path_rate: float = self.config["drop_path_rate"]

        self.stem = Stem(self.input_channels, stem_hidden_dim, embed_dims[0], stem_stride, patch_size)

        self.post_patch_h = self.size[0] // patch_size // 2
        self.post_patch_w = self.size[1] // patch_size // 2
        self.pos_embed = nn.Parameter(torch.zeros(1, self.post_patch_h, self.post_patch_w, embed_dims[1]))

        total_blocks = sum(layers)
        dpr = torch.linspace(0, drop_path_rate, total_blocks).tolist()

        # Outlooker stage
        outlook_padding = outlook_kernel // 2
        outlooker_blocks = nn.ModuleList()
        for i in range(layers[0]):
            outlooker_blocks.append(
                OutlookerBlock(
                    dim=embed_dims[0],
                    num_heads=num_heads[0],
                    kernel_size=outlook_kernel,
                    padding=outlook_padding,
                    stride=outlook_stride,
                    mlp_ratio=mlp_ratio,
                    drop_path=dpr[i],
                )
            )

        # Downsample
        self.downsample = Downsample(embed_dims[0], embed_dims[1])

        # Transformer stages
        transformer_blocks: list[nn.ModuleList] = []
        block_idx = layers[0]
        for stage_idx in range(1, len(layers)):
            stage_blocks = nn.ModuleList()
            for i in range(layers[stage_idx]):
                stage_blocks.append(
                    TransformerBlock(
                        dim=embed_dims[stage_idx],
                        num_heads=num_heads[stage_idx],
                        mlp_ratio=mlp_ratio,
                        drop_path=dpr[block_idx],
                    )
                )
                block_idx += 1

            transformer_blocks.append(stage_blocks)

        self.network = nn.ModuleList([outlooker_blocks, *transformer_blocks])

        # Class attention post layers
        self.post_network = nn.ModuleList()
        for _ in range(num_post_layers):
            self.post_network.append(ClassBlock(embed_dims[-1], num_heads[-1], mlp_ratio, drop_path=0.0))

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dims[-1]))
        self.norm = nn.LayerNorm(embed_dims[-1])

        self.embedding_size = embed_dims[-1]
        self.classifier = self.create_classifier()

        # Weight initialization
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.LayerNorm):
                nn.init.zeros_(m.bias)
                nn.init.ones_(m.weight)

        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)

    def _get_pos_embed(self, h: int, w: int) -> torch.Tensor:
        if self.dynamic_size is False:
            return self.pos_embed

        if h == self.post_patch_h and w == self.post_patch_w:
            return self.pos_embed

        pos_embed = self.pos_embed.permute(0, 3, 1, 2).float()
        pos_embed = F.interpolate(pos_embed, size=(h, w), mode="bicubic", antialias=False)
        pos_embed = pos_embed.permute(0, 2, 3, 1).to(self.pos_embed.dtype)

        return pos_embed

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

        for idx, stage in enumerate(self.network):
            if idx == 0:
                # Outlooker stage
                for blk in stage:
                    x = blk(x)

                # Downsample and add position embedding
                x = self.downsample(x)
                x = x + self._get_pos_embed(x.size(1), x.size(2))
            else:
                # Transformer stages
                for blk in stage:
                    x = blk(x)

        # Reshape to sequence: (B, H, W, C) -> (B, N, C)
        x = x.reshape(x.size(0), -1, x.size(-1))

        # Prepend class token and run class attention
        cls_tokens = self.cls_token.expand(x.size(0), -1, -1)
        x = torch.concat([cls_tokens, x], dim=1)
        for blk in self.post_network:
            x = blk(x)

        x = self.norm(x)

        return x

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

    def adjust_size(self, new_size: tuple[int, int]) -> None:
        if new_size == self.size:
            return

        old_h = self.post_patch_h
        old_w = self.post_patch_w
        super().adjust_size(new_size)

        self.post_patch_h = new_size[0] // 16  # patch_size=8, downsample=2
        self.post_patch_w = new_size[1] // 16
        if (old_h, old_w) != (self.post_patch_h, self.post_patch_w):
            with torch.no_grad():
                pos_embed = self.pos_embed.permute(0, 3, 1, 2).float()
                pos_embed = F.interpolate(
                    pos_embed,
                    size=(self.post_patch_h, self.post_patch_w),
                    mode="bicubic",
                    antialias=True,
                )
                pos_embed = pos_embed.permute(0, 2, 3, 1).to(self.pos_embed.dtype)

            self.pos_embed = nn.Parameter(pos_embed)


registry.register_model_config(
    "volo_d1",
    VOLO,
    config={
        "layers": [4, 4, 8, 2],
        "embed_dims": [192, 384, 384, 384],
        "num_heads": [6, 12, 12, 12],
        "mlp_ratio": 3.0,
        "stem_hidden_dim": 64,
        "drop_path_rate": 0.1,
    },
)
registry.register_model_config(
    "volo_d2",
    VOLO,
    config={
        "layers": [6, 4, 10, 4],
        "embed_dims": [256, 512, 512, 512],
        "num_heads": [8, 16, 16, 16],
        "mlp_ratio": 3.0,
        "stem_hidden_dim": 64,
        "drop_path_rate": 0.2,
    },
)
registry.register_model_config(
    "volo_d3",
    VOLO,
    config={
        "layers": [8, 8, 16, 4],
        "embed_dims": [256, 512, 512, 512],
        "num_heads": [8, 16, 16, 16],
        "mlp_ratio": 3.0,
        "stem_hidden_dim": 64,
        "drop_path_rate": 0.5,
    },
)
registry.register_model_config(
    "volo_d4",
    VOLO,
    config={
        "layers": [8, 8, 16, 4],
        "embed_dims": [384, 768, 768, 768],
        "num_heads": [12, 16, 16, 16],
        "mlp_ratio": 3.0,
        "stem_hidden_dim": 64,
        "drop_path_rate": 0.5,
    },
)
registry.register_model_config(
    "volo_d5",
    VOLO,
    config={
        "layers": [12, 12, 20, 4],
        "embed_dims": [384, 768, 768, 768],
        "num_heads": [12, 16, 16, 16],
        "mlp_ratio": 4.0,
        "stem_hidden_dim": 128,
        "drop_path_rate": 0.75,
    },
)
