"""
LIT v2, adapted from
https://github.com/ziplab/LITv2/blob/main/classification/models/litv2.py

Paper "Fast Vision Transformers with HiLo Attention", https://arxiv.org/abs/2205.13213

Generated by Claude Code Opus 4.5
"""

# Reference license: Apache-2.0

import math
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 Permute
from torchvision.ops import StochasticDepth

from birder.model_registry import registry
from birder.net.base import DetectorBackbone
from birder.net.lit_v1 import DeformablePatchMerging
from birder.net.lit_v1 import IdentityDownsample


class DepthwiseMLP(nn.Module):
    def __init__(self, in_features: int, hidden_features: int) -> None:
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.dwconv = nn.Conv2d(
            hidden_features, hidden_features, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=hidden_features
        )
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden_features, in_features)

    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
        x = self.fc1(x)

        B, N, C = x.size()
        x = x.reshape(B, H, W, C).permute(0, 3, 1, 2).contiguous()
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1).reshape(B, N, C)
        x = self.act(x)
        x = self.fc2(x)

        return x


class DepthwiseMLPBlock(nn.Module):
    def __init__(self, dim: int, mlp_ratio: float, drop_path: float) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.mlp = DepthwiseMLP(dim, int(dim * mlp_ratio))
        self.drop_path = StochasticDepth(drop_path, mode="row")

    def forward(self, x: torch.Tensor, resolution: tuple[int, int]) -> torch.Tensor:
        H, W = resolution
        return x + self.drop_path(self.mlp(self.norm(x), H, W))


class HiLoAttention(nn.Module):
    """
    HiLo Attention: High-frequency local attention + Low-frequency global attention

    Hi-Fi (High frequency): Local window attention
    Lo-Fi (Low frequency): Global attention with average pooling
    """

    def __init__(
        self,
        dim: int,
        num_heads: int,
        window_size: int,
        alpha: float,
    ) -> None:
        super().__init__()
        assert dim % num_heads == 0, "dim must be divisible by num_heads"

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

        # Split heads between Lo-Fi (global) and Hi-Fi (local)
        self.l_heads = int(num_heads * alpha)  # Lo-Fi heads
        self.h_heads = num_heads - self.l_heads  # Hi-Fi heads
        self.l_dim = self.l_heads * self.head_dim
        self.h_dim = self.h_heads * self.head_dim
        self.head_dim = self.head_dim

        # ws == 1 is equal to standard multi-head self-attention
        if window_size == 1:
            self.h_heads = 0
            self.h_dim = 0
            self.l_heads = num_heads
            self.l_dim = dim

        # Lo-Fi: Global attention with pooling
        if self.l_heads > 0:
            if window_size > 1:
                self.sr = nn.AvgPool2d(kernel_size=(window_size, window_size), stride=(window_size, window_size))
            else:
                self.sr = nn.Identity()

            self.l_q = nn.Linear(dim, self.l_dim)
            self.l_kv = nn.Linear(dim, self.l_dim * 2)
            self.l_proj = nn.Linear(self.l_dim, self.l_dim)
        else:
            self.l_q = nn.Identity()
            self.l_kv = nn.Identity()
            self.l_proj = nn.Identity()

        # Hi-Fi: Local window attention
        if self.h_heads > 0:
            self.h_qkv = nn.Linear(dim, self.h_dim * 3)
            self.h_proj = nn.Linear(self.h_dim, self.h_dim)
        else:
            self.h_qkv = nn.Identity()
            self.h_proj = nn.Identity()

    def _lofi(self, x: torch.Tensor) -> torch.Tensor:
        B, H, W, C = x.size()

        q = self.l_q(x).reshape(B, H * W, self.l_heads, self.head_dim).permute(0, 2, 1, 3)

        # Spatial reduction for k, v
        if self.window_size > 1:
            x = x.permute(0, 3, 1, 2)
            x = self.sr(x).reshape(B, C, -1).permute(0, 2, 1)
            kv = self.l_kv(x).reshape(B, -1, 2, self.l_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        else:
            kv = self.l_kv(x).reshape(B, -1, 2, self.l_heads, self.head_dim).permute(2, 0, 3, 1, 4)

        k, v = kv.unbind(0)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = F.softmax(attn, dim=-1)

        x = (attn @ v).transpose(1, 2).reshape(B, H, W, self.l_dim)
        x = self.l_proj(x)

        return x

    def _hifi(self, x: torch.Tensor) -> torch.Tensor:
        B, H, W, _ = x.size()
        ws = self.window_size

        # Pad if needed
        pad_h = (ws - H % ws) % ws
        pad_w = (ws - W % ws) % ws
        if pad_h > 0 or pad_w > 0:
            x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))

        _, h_pad, w_pad, _ = x.size()
        h_groups = h_pad // ws
        w_groups = w_pad // ws
        total_groups = h_groups * w_groups

        x = x.reshape(B, h_groups, ws, w_groups, ws, -1).transpose(2, 3)

        qkv = self.h_qkv(x).reshape(B, total_groups, -1, 3, self.h_heads, self.head_dim).permute(3, 0, 1, 4, 2, 5)
        q, k, v = qkv.unbind(0)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = F.softmax(attn, dim=-1)

        x = (attn @ v).transpose(2, 3).reshape(B, h_groups, w_groups, ws, ws, self.h_dim)
        x = x.transpose(2, 3).reshape(B, h_pad, w_pad, self.h_dim)
        x = self.h_proj(x)

        # Remove padding
        if pad_h > 0 or pad_w > 0:
            x = x[:, :H, :W, :].contiguous()

        return x

    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
        B, N, C = x.size()
        x = x.reshape(B, H, W, C)

        if self.h_heads == 0:
            x = self._lofi(x)
            return x.reshape(B, N, C)

        if self.l_heads == 0:
            x = self._hifi(x)
            return x.reshape(B, N, C)

        # Process both branches and concatenate
        hifi_out = self._hifi(x)
        lofi_out = self._lofi(x)

        x = torch.concat((hifi_out, lofi_out), dim=-1)
        return x.reshape(B, N, C)


class HiLoBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        window_size: int,
        alpha: float,
        mlp_ratio: float,
        drop_path: float,
    ) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = HiLoAttention(dim, num_heads, window_size, alpha)
        self.drop_path1 = StochasticDepth(drop_path, mode="row")
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = DepthwiseMLP(dim, int(dim * mlp_ratio))
        self.drop_path2 = StochasticDepth(drop_path, mode="row")

    def forward(self, x: torch.Tensor, resolution: tuple[int, int]) -> torch.Tensor:
        H, W = resolution
        x = x + self.drop_path1(self.attn(self.norm1(x), H, W))
        x = x + self.drop_path2(self.mlp(self.norm2(x), H, W))
        return x


class LITStage(nn.Module):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        resolution: tuple[int, int],
        depth: int,
        num_heads: int,
        window_size: int,
        alpha: float,
        mlp_ratio: float,
        downsample: bool,
        drop_path: list[float],
    ) -> None:
        super().__init__()
        if downsample is True:
            self.downsample = DeformablePatchMerging(in_dim, out_dim)
            resolution = (resolution[0] // 2, resolution[1] // 2)
        else:
            self.downsample = IdentityDownsample()

        blocks: list[nn.Module] = []
        for i in range(depth):
            if window_size > 0:
                blocks.append(HiLoBlock(out_dim, num_heads, window_size, alpha, mlp_ratio, drop_path[i]))
            else:
                blocks.append(DepthwiseMLPBlock(out_dim, mlp_ratio, drop_path[i]))

        self.blocks = nn.ModuleList(blocks)

    def forward(self, x: torch.Tensor, input_resolution: tuple[int, int]) -> tuple[torch.Tensor, int, int]:
        x, H, W = self.downsample(x, input_resolution)
        for block in self.blocks:
            x = block(x, (H, W))

        return (x, H, W)


# pylint: disable=invalid-name
class LIT_v2(DetectorBackbone):
    block_group_regex = r"body\.stage(\d+)\.blocks\.(\d+)"

    # 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 = 4
        embed_dim: int = self.config["embed_dim"]
        depths: list[int] = self.config["depths"]
        num_heads: list[int] = self.config["num_heads"]
        local_ws: list[int] = self.config["local_ws"]
        alpha: float = self.config["alpha"]
        drop_path_rate: float = self.config["drop_path_rate"]

        num_stages = len(depths)

        self.stem = nn.Sequential(
            nn.Conv2d(
                self.input_channels,
                embed_dim,
                kernel_size=(patch_size, patch_size),
                stride=(patch_size, patch_size),
                padding=(0, 0),
            ),
            Permute([0, 2, 3, 1]),
            nn.LayerNorm(embed_dim),
        )

        # Stochastic depth
        dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(depths)).split(depths)]

        stages: OrderedDict[str, nn.Module] = OrderedDict()
        return_channels: list[int] = []
        prev_dim = embed_dim
        resolution = (self.size[0] // patch_size, self.size[1] // patch_size)
        for i_stage in range(num_stages):
            in_dim = prev_dim
            out_dim = in_dim * 2 if i_stage > 0 else in_dim
            stage = LITStage(
                in_dim,
                out_dim,
                resolution,
                depth=depths[i_stage],
                num_heads=num_heads[i_stage],
                window_size=local_ws[i_stage],
                alpha=alpha,
                mlp_ratio=4.0,
                downsample=i_stage > 0,
                drop_path=dpr[i_stage],
            )
            stages[f"stage{i_stage + 1}"] = stage

            if i_stage > 0:
                resolution = (resolution[0] // 2, resolution[1] // 2)

            prev_dim = out_dim
            return_channels.append(out_dim)

        num_features = embed_dim * (2 ** (num_stages - 1))
        self.body = nn.ModuleDict(stages)
        self.features = nn.Sequential(
            nn.LayerNorm(num_features),
            Permute([0, 2, 1]),
            nn.AdaptiveAvgPool1d(output_size=1),
            nn.Flatten(1),
        )
        self.return_channels = return_channels
        self.embedding_size = num_features
        self.classifier = self.create_classifier()

        # Weight initialization
        for name, m in self.named_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.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Conv2d):
                if name.endswith("offset_conv") is True:
                    continue

                fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                fan_out //= m.groups
                nn.init.normal_(m.weight, mean=0.0, std=math.sqrt(2.0 / fan_out))
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        x = self.stem(x)
        B, H, W, C = x.size()
        x = x.reshape(B, H * W, C)

        out = {}
        for name, stage in self.body.items():
            x, H, W = stage(x, (H, W))
            if name in self.return_stages:
                features = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
                out[name] = features

        return out

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

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

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

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        B, H, W, C = x.size()
        x = x.reshape(B, H * W, C)
        for stage in self.body.values():
            x, H, W = stage(x, (H, W))

        return x

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


registry.register_model_config(
    "lit_v2_s",
    LIT_v2,
    config={
        "embed_dim": 96,
        "depths": [2, 2, 6, 2],
        "num_heads": [3, 6, 12, 24],
        "local_ws": [0, 0, 2, 1],
        "alpha": 0.9,
        "drop_path_rate": 0.2,
    },
)
registry.register_model_config(
    "lit_v2_m",
    LIT_v2,
    config={
        "embed_dim": 96,
        "depths": [2, 2, 18, 2],
        "num_heads": [3, 6, 12, 24],
        "local_ws": [0, 0, 2, 1],
        "alpha": 0.9,
        "drop_path_rate": 0.3,
    },
)
registry.register_model_config(
    "lit_v2_b",
    LIT_v2,
    config={
        "embed_dim": 128,
        "depths": [2, 2, 18, 2],
        "num_heads": [4, 8, 16, 32],
        "local_ws": [0, 0, 2, 1],
        "alpha": 0.9,
        "drop_path_rate": 0.5,
    },
)
