"""
LIT v1 Tiny, adapted from
https://github.com/ziplab/LIT/blob/main/classification/code_for_lit_ti/lit.py

Paper "Less is More: Pay Less Attention in Vision Transformers", https://arxiv.org/abs/2105.14217

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 MLP
from birder.net.lit_v1 import DeformablePatchMerging
from birder.net.lit_v1 import IdentityDownsample
from birder.net.vit import adjust_position_embedding


class MLPBlock(nn.Module):
    def __init__(self, dim: int, mlp_ratio: float, drop_path: float) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim, eps=1e-6)
        self.mlp = MLP(dim, int(dim * mlp_ratio))
        self.drop_path = StochasticDepth(drop_path, mode="row")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.drop_path(self.mlp(self.norm(x)))


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.size()
        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)

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

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)

        return x


class ViTBlock(nn.Module):
    def __init__(self, dim: int, num_heads: int, mlp_ratio: float, drop_path: float) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
        self.attn = Attention(dim, num_heads)
        self.drop_path1 = StochasticDepth(drop_path, mode="row")

        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
        self.mlp = MLP(dim, int(dim * mlp_ratio))
        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 LITStage(nn.Module):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        input_resolution: tuple[int, int],
        depth: int,
        num_heads: int,
        mlp_ratio: float,
        has_msa: bool,
        downsample: bool,
        use_cls_token: bool,
        drop_path: list[float],
    ) -> None:
        super().__init__()
        self.dynamic_size = False
        self.input_resolution = input_resolution
        self.downsample: nn.Module
        if downsample is True:
            self.downsample = DeformablePatchMerging(in_dim, out_dim)
        else:
            self.downsample = IdentityDownsample()

        blocks: list[nn.Module] = []
        for i in range(depth):
            if has_msa is True:
                blocks.append(ViTBlock(out_dim, num_heads, mlp_ratio, drop_path[i]))
            else:
                blocks.append(MLPBlock(out_dim, mlp_ratio, drop_path[i]))

        self.blocks = nn.ModuleList(blocks)

        num_tokens = input_resolution[0] * input_resolution[1]
        if use_cls_token is True:
            self.cls_token = nn.Parameter(torch.zeros(1, 1, out_dim))
            nn.init.trunc_normal_(self.cls_token, std=0.02)
            num_tokens += 1
        else:
            self.cls_token = None

        self.pos_embed = nn.Parameter(torch.zeros(1, num_tokens, out_dim))
        nn.init.trunc_normal_(self.pos_embed, std=0.02)

    def set_dynamic_size(self, dynamic_size: bool = True) -> None:
        self.dynamic_size = dynamic_size

    def _get_pos_embed(self, H: int, W: int) -> torch.Tensor:
        if self.dynamic_size is False or (H == self.input_resolution[0] and W == self.input_resolution[1]):
            return self.pos_embed

        if self.cls_token is not None:
            num_prefix_tokens = 1
        else:
            num_prefix_tokens = 0

        return adjust_position_embedding(
            self.pos_embed, self.input_resolution, (H, W), num_prefix_tokens=num_prefix_tokens
        )

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

        if self.cls_token is not None:
            cls_tokens = self.cls_token.expand(x.size(0), -1, -1)
            x = torch.concat((cls_tokens, x), dim=1)

        x = x + self._get_pos_embed(H, W)

        for block in self.blocks:
            x = block(x)

        return (x, H, W)


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

    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
        stage_dims: list[int] = self.config["stage_dims"]
        depths: list[int] = self.config["depths"]
        num_heads: list[int] = self.config["num_heads"]
        mlp_ratios: list[float] = self.config["mlp_ratios"]
        has_msa: list[bool] = self.config["has_msa"]
        drop_path_rate: float = self.config["drop_path_rate"]

        num_stages = len(depths)

        self.stem = nn.Sequential(
            nn.Conv2d(
                self.input_channels,
                stage_dims[0],
                kernel_size=(patch_size, patch_size),
                stride=(patch_size, patch_size),
                padding=(0, 0),
            ),
            Permute([0, 2, 3, 1]),
            nn.LayerNorm(stage_dims[0], eps=1e-6),
        )

        # 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] = []
        resolution = (self.size[0] // patch_size, self.size[1] // patch_size)

        for i in range(num_stages):
            if i > 0:
                resolution = (resolution[0] // 2, resolution[1] // 2)

            stage = LITStage(
                stage_dims[i - 1] if i > 0 else stage_dims[0],
                stage_dims[i],
                input_resolution=resolution,
                depth=depths[i],
                num_heads=num_heads[i],
                mlp_ratio=mlp_ratios[i],
                has_msa=has_msa[i],
                downsample=i > 0,
                use_cls_token=i == num_stages - 1,
                drop_path=dpr[i],
            )
            stages[f"stage{i + 1}"] = stage
            return_channels.append(stage_dims[i])

        self.body = nn.ModuleDict(stages)
        self.norm = nn.LayerNorm(stage_dims[-1], eps=1e-6)
        self.return_channels = return_channels
        self.embedding_size = stage_dims[-1]
        self.classifier = self.create_classifier()
        self.patch_size = patch_size

        # 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)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                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:
                if stage.cls_token is not None:
                    spatial_x = x[:, 1:]
                else:
                    spatial_x = x

                out[name] = spatial_x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()

        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)
        x = self.norm(x)
        return x[:, 0]

    def set_dynamic_size(self, dynamic_size: bool = True) -> None:
        super().set_dynamic_size(dynamic_size)
        for stage in self.body.values():
            stage.set_dynamic_size(dynamic_size)

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

        super().adjust_size(new_size)

        new_patches_resolution = (new_size[0] // self.patch_size, new_size[1] // self.patch_size)

        h, w = new_patches_resolution
        for stage in self.body.values():
            if not isinstance(stage.downsample, IdentityDownsample):
                h = h // 2
                w = w // 2

            out_resolution = (h, w)
            if out_resolution == stage.input_resolution:
                continue

            if stage.cls_token is not None:
                num_prefix_tokens = 1
            else:
                num_prefix_tokens = 0

            with torch.no_grad():
                pos_embed = adjust_position_embedding(
                    stage.pos_embed,
                    stage.input_resolution,
                    out_resolution,
                    num_prefix_tokens=num_prefix_tokens,
                )

            stage.input_resolution = out_resolution
            stage.pos_embed = nn.Parameter(pos_embed)


registry.register_model_config(
    "lit_v1_t",
    LIT_v1_Tiny,
    config={
        "stage_dims": [64, 128, 320, 512],
        "depths": [3, 4, 6, 3],
        "num_heads": [1, 2, 5, 8],
        "mlp_ratios": [8.0, 8.0, 4.0, 4.0],
        "has_msa": [False, False, True, True],
        "drop_path_rate": 0.1,
    },
)

registry.register_weights(
    "lit_v1_t_il-common",
    {
        "description": "LIT v1 Tiny model trained on the il-common dataset",
        "resolution": (256, 256),
        "formats": {
            "pt": {
                "file_size": 75.2,
                "sha256": "93813b2716eb9f33e06dc15ab2ba335c6d219354d2983bbc4f834f8f4e688e5c",
            }
        },
        "net": {"network": "lit_v1_t", "tag": "il-common"},
    },
)
