"""
Paper "DeiT III: Revenge of the ViT", https://arxiv.org/abs/2204.07118
and
Paper "Vision Transformers Need Registers", https://arxiv.org/abs/2309.16588
"""

import math
from functools import partial
from typing import Any
from typing import Literal
from typing import Optional

import torch
from torch import nn

from birder.common.masking import mask_tensor
from birder.model_registry import registry
from birder.net._vit_configs import BASE
from birder.net._vit_configs import HUGE
from birder.net._vit_configs import LARGE
from birder.net._vit_configs import MEDIUM
from birder.net._vit_configs import SMALL
from birder.net._vit_configs import TINY
from birder.net.base import DetectorBackbone
from birder.net.base import MaskedTokenOmissionMixin
from birder.net.base import MaskedTokenRetentionMixin
from birder.net.base import PreTrainEncoder
from birder.net.base import TokenOmissionResultType
from birder.net.base import TokenRetentionResultType
from birder.net.base import normalize_out_indices
from birder.net.vit import Encoder
from birder.net.vit import EncoderBlock
from birder.net.vit import PatchEmbed
from birder.net.vit import adjust_position_embedding


# pylint: disable=too-many-instance-attributes
class DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTokenRetentionMixin):
    block_group_regex = r"encoder\.block\.(\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"

        image_size = self.size
        attention_dropout = 0.0
        dropout = 0.0
        patch_size: int = self.config["patch_size"]
        num_layers: int = self.config["num_layers"]
        num_heads: int = self.config["num_heads"]
        hidden_dim: int = self.config["hidden_dim"]
        mlp_dim: int = self.config["mlp_dim"]
        layer_scale_init_value: Optional[float] = self.config.get("layer_scale_init_value", 1e-5)
        num_reg_tokens: int = self.config.get("num_reg_tokens", 0)
        out_indices: Optional[list[int]] = self.config.get("out_indices", None)
        drop_path_rate: float = self.config["drop_path_rate"]

        torch._assert(image_size[0] % patch_size == 0, "Input shape indivisible by patch size!")
        torch._assert(image_size[1] % patch_size == 0, "Input shape indivisible by patch size!")
        torch._assert(hidden_dim % num_heads == 0, "Hidden dim indivisible by num heads!")
        self.patch_size = patch_size
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        self.num_reg_tokens = num_reg_tokens
        self.num_special_tokens = 1 + self.num_reg_tokens
        self.out_indices = normalize_out_indices(out_indices, num_layers)
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]  # Stochastic depth decay rule

        self.conv_proj = nn.Conv2d(
            self.input_channels,
            hidden_dim,
            kernel_size=(patch_size, patch_size),
            stride=(patch_size, patch_size),
            padding=(0, 0),
        )
        self.patch_embed = PatchEmbed()

        seq_length = (image_size[0] // patch_size) * (image_size[1] // patch_size)

        # Add a class token
        self.class_token = nn.Parameter(torch.zeros(1, 1, hidden_dim))

        # Add optional register tokens
        if self.num_reg_tokens > 0:
            self.reg_tokens = nn.Parameter(torch.zeros(1, self.num_reg_tokens, hidden_dim))
        else:
            self.reg_tokens = None

        # Add positional embedding
        self.pos_embedding = nn.Parameter(torch.empty(1, seq_length, hidden_dim).normal_(std=0.02))  # from BERT

        self.encoder = Encoder(
            num_layers,
            num_heads,
            hidden_dim,
            mlp_dim,
            dropout,
            attention_dropout,
            dpr,
            layer_scale_init_value=layer_scale_init_value,
        )
        self.norm = nn.LayerNorm(hidden_dim, eps=1e-6)

        num_return_stages = len(self.out_indices) if self.out_indices is not None else 1
        self.return_stages = [f"stage{stage_idx + 1}" for stage_idx in range(num_return_stages)]
        self.return_channels = [hidden_dim] * num_return_stages
        self.embedding_size = hidden_dim
        self.classifier = self.create_classifier()

        self.max_stride = patch_size
        self.stem_stride = patch_size
        self.stem_width = hidden_dim
        self.encoding_size = hidden_dim
        self.decoder_block = partial(
            EncoderBlock,
            16,
            mlp_dim=None,
            dropout=0,
            attention_dropout=0,
            drop_path=0,
            activation_layer=nn.GELU,
        )

        # Weight initialization
        if isinstance(self.conv_proj, nn.Conv2d):
            # Init the patchify stem
            fan_in = self.conv_proj.in_channels * self.conv_proj.kernel_size[0] * self.conv_proj.kernel_size[1]
            nn.init.trunc_normal_(self.conv_proj.weight, std=math.sqrt(1 / fan_in))
            if self.conv_proj.bias is not None:
                nn.init.zeros_(self.conv_proj.bias)

        if isinstance(self.classifier, nn.Linear):
            nn.init.zeros_(self.classifier.weight)
            nn.init.zeros_(self.classifier.bias)

    def _get_pos_embed(self, H: int, W: int) -> torch.Tensor:
        if self.dynamic_size is False:
            return self.pos_embedding

        if H == self.size[0] and W == self.size[1]:
            return self.pos_embedding

        return adjust_position_embedding(
            self.pos_embedding,
            (self.size[0] // self.patch_size, self.size[1] // self.patch_size),
            (H // self.patch_size, W // self.patch_size),
            0,
            antialias=False,
        )

    def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        H, W = x.shape[-2:]

        x = self.conv_proj(x)
        x = self.patch_embed(x)

        # Expand special tokens to batch size and prepend in order [REG..., CLS, PATCH...]
        special_tokens: list[torch.Tensor] = []
        if self.reg_tokens is not None:
            special_tokens.append(self.reg_tokens.expand(x.size(0), -1, -1))

        special_tokens.append(self.class_token.expand(x.size(0), -1, -1))
        x = x + self._get_pos_embed(H, W)
        x = torch.concat(special_tokens + [x], dim=1)

        if self.out_indices is None:
            xs = [self.encoder(x)]
        else:
            xs = self.encoder.forward_features(x, out_indices=self.out_indices)

        out: dict[str, torch.Tensor] = {}
        for stage_name, stage_x in zip(self.return_stages, xs, strict=True):
            stage_x = stage_x[:, self.num_special_tokens :]
            stage_x = stage_x.permute(0, 2, 1)
            B, C, _ = stage_x.size()
            stage_x = stage_x.reshape(B, C, H // self.patch_size, W // self.patch_size)
            out[stage_name] = stage_x

        return out

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

        self.pos_embedding.requires_grad_(False)

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

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

    def transform_to_backbone(self) -> None:
        super().transform_to_backbone()
        self.norm = nn.Identity()

    def set_causal_attention(self, is_causal: bool = True) -> None:
        self.encoder.set_causal_attention(is_causal)

    def masked_encoding_omission(
        self,
        x: torch.Tensor,
        ids_keep: Optional[torch.Tensor] = None,
        return_all_features: bool = False,
        return_keys: Literal["all", "tokens", "embedding"] = "tokens",
    ) -> TokenOmissionResultType:
        H, W = x.shape[-2:]

        # Reshape and permute the input tensor
        x = self.conv_proj(x)
        x = self.patch_embed(x)

        # Add pos embedding without special tokens
        pos_embedding = self._get_pos_embed(H, W)
        x = x + pos_embedding

        # Mask tokens
        if ids_keep is not None:
            x = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, x.size(2)))

        # Expand special tokens to batch size and prepend in order [REG..., CLS, PATCH...]
        special_tokens: list[torch.Tensor] = []
        if self.reg_tokens is not None:
            special_tokens.append(self.reg_tokens.expand(x.size(0), -1, -1))

        special_tokens.append(self.class_token.expand(x.size(0), -1, -1))

        x = torch.concat(special_tokens + [x], dim=1)

        # Apply transformer
        if return_all_features is True:
            xs = self.encoder.forward_features(x)
            xs[-1] = self.norm(xs[-1])
            x = torch.stack(xs, dim=-1)
        else:
            x = self.encoder(x)
            x = self.norm(x)

        result: TokenOmissionResultType = {}
        if return_keys in ("all", "tokens"):
            result["tokens"] = x

        if return_keys in ("all", "embedding"):
            if return_all_features is True:
                x = x[..., -1]

            result["embedding"] = x[:, self.num_reg_tokens]

        return result

    def masked_encoding_retention(
        self,
        x: torch.Tensor,
        mask: torch.Tensor,
        mask_token: Optional[torch.Tensor] = None,
        return_keys: Literal["all", "features", "embedding"] = "features",
    ) -> TokenRetentionResultType:
        H, W = x.shape[-2:]

        x = self.conv_proj(x)
        x = mask_tensor(x, mask, mask_token=mask_token, patch_factor=self.max_stride // self.stem_stride)

        # Reshape and permute the input tensor
        x = self.patch_embed(x)

        # Expand special tokens to batch size and prepend in order [REG..., CLS, PATCH...]
        special_tokens: list[torch.Tensor] = []
        if self.reg_tokens is not None:
            special_tokens.append(self.reg_tokens.expand(x.size(0), -1, -1))

        special_tokens.append(self.class_token.expand(x.size(0), -1, -1))

        x = x + self._get_pos_embed(H, W)
        x = torch.concat(special_tokens + [x], dim=1)

        x = self.encoder(x)
        x = self.norm(x)

        result: TokenRetentionResultType = {}
        if return_keys in ("all", "features"):
            features = x[:, self.num_special_tokens :]
            features = features.permute(0, 2, 1)
            B, C, _ = features.size()
            features = features.reshape(B, C, H // self.patch_size, W // self.patch_size)
            result["features"] = features

        if return_keys in ("all", "embedding"):
            result["embedding"] = x[:, self.num_reg_tokens]

        return result

    def forward_features(self, x: torch.Tensor, return_input_embedding: bool = False) -> torch.Tensor:
        H, W = x.shape[-2:]

        # Reshape and permute the input tensor
        x = self.conv_proj(x)
        x = self.patch_embed(x)

        # Expand special tokens to batch size and prepend in order [REG..., CLS, PATCH...]
        special_tokens: list[torch.Tensor] = []
        if self.reg_tokens is not None:
            special_tokens.append(self.reg_tokens.expand(x.size(0), -1, -1))

        special_tokens.append(self.class_token.expand(x.size(0), -1, -1))

        patch_embedding = x
        x = x + self._get_pos_embed(H, W)
        x = torch.concat(special_tokens + [x], dim=1)

        if return_input_embedding is True:
            input_embedding = torch.concat(special_tokens + [patch_embedding], dim=1)
        else:
            input_embedding = None  # For TorchScript compatibility

        x = self.encoder(x)
        x = self.norm(x)

        if return_input_embedding is True and input_embedding is not None:
            return torch.stack([input_embedding, x], dim=-1)

        return x

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

        return x

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

        assert new_size[0] % self.patch_size == 0, "Input shape indivisible by patch size!"
        assert new_size[1] % self.patch_size == 0, "Input shape indivisible by patch size!"

        old_size = self.size
        super().adjust_size(new_size)

        # Add back class tokens
        with torch.no_grad():
            pos_embedding = adjust_position_embedding(
                self.pos_embedding,
                (old_size[0] // self.patch_size, old_size[1] // self.patch_size),
                (new_size[0] // self.patch_size, new_size[1] // self.patch_size),
                0,
            )
        self.pos_embedding = nn.Parameter(pos_embedding)


registry.register_model_config(
    "deit3_t16",
    DeiT3,
    config={"patch_size": 16, **TINY},
)
registry.register_model_config(
    "deit3_t14",
    DeiT3,
    config={"patch_size": 14, **TINY},
)
registry.register_model_config(
    "deit3_s16",
    DeiT3,
    config={"patch_size": 16, **SMALL, "drop_path_rate": 0.05},
)
registry.register_model_config(
    "deit3_s14",
    DeiT3,
    config={"patch_size": 14, **SMALL, "drop_path_rate": 0.05},
)
registry.register_model_config(
    "deit3_m16",
    DeiT3,
    config={"patch_size": 16, **MEDIUM, "drop_path_rate": 0.1},
)
registry.register_model_config(
    "deit3_m14",
    DeiT3,
    config={"patch_size": 14, **MEDIUM, "drop_path_rate": 0.1},
)
registry.register_model_config(
    "deit3_b16",
    DeiT3,
    config={"patch_size": 16, **BASE, "drop_path_rate": 0.2},
)
registry.register_model_config(
    "deit3_b14",
    DeiT3,
    config={"patch_size": 14, **BASE, "drop_path_rate": 0.2},
)
registry.register_model_config(
    "deit3_l16",
    DeiT3,
    config={"patch_size": 16, **LARGE, "drop_path_rate": 0.45},
)
registry.register_model_config(
    "deit3_l14",
    DeiT3,
    config={"patch_size": 14, **LARGE, "drop_path_rate": 0.45},
)
registry.register_model_config(
    "deit3_h16",
    DeiT3,
    config={"patch_size": 16, **HUGE, "drop_path_rate": 0.55},
)
registry.register_model_config(
    "deit3_h14",
    DeiT3,
    config={"patch_size": 14, **HUGE, "drop_path_rate": 0.55},
)

# With registers
####################

registry.register_model_config(
    "deit3_reg4_t16",
    DeiT3,
    config={"patch_size": 16, **TINY, "num_reg_tokens": 4},
)
registry.register_model_config(
    "deit3_reg4_t14",
    DeiT3,
    config={"patch_size": 14, **TINY, "num_reg_tokens": 4},
)
registry.register_model_config(
    "deit3_reg4_s16",
    DeiT3,
    config={"patch_size": 16, **SMALL, "num_reg_tokens": 4, "drop_path_rate": 0.05},
)
registry.register_model_config(
    "deit3_reg4_s14",
    DeiT3,
    config={"patch_size": 14, **SMALL, "num_reg_tokens": 4, "drop_path_rate": 0.05},
)
registry.register_model_config(
    "deit3_reg4_m16",
    DeiT3,
    config={"patch_size": 16, **MEDIUM, "num_reg_tokens": 4, "drop_path_rate": 0.1},
)
registry.register_model_config(
    "deit3_reg4_m14",
    DeiT3,
    config={"patch_size": 14, **MEDIUM, "num_reg_tokens": 4, "drop_path_rate": 0.1},
)
registry.register_model_config(
    "deit3_reg4_b16",
    DeiT3,
    config={"patch_size": 16, **BASE, "num_reg_tokens": 4, "drop_path_rate": 0.2},
)
registry.register_model_config(
    "deit3_reg4_b14",
    DeiT3,
    config={"patch_size": 14, **BASE, "num_reg_tokens": 4, "drop_path_rate": 0.2},
)
registry.register_model_config(
    "deit3_reg4_l16",
    DeiT3,
    config={"patch_size": 16, **LARGE, "num_reg_tokens": 4, "drop_path_rate": 0.45},
)
registry.register_model_config(
    "deit3_reg4_l14",
    DeiT3,
    config={"patch_size": 14, **LARGE, "num_reg_tokens": 4, "drop_path_rate": 0.45},
)
registry.register_model_config(
    "deit3_reg4_h16",
    DeiT3,
    config={"patch_size": 16, **HUGE, "num_reg_tokens": 4, "drop_path_rate": 0.55},
)
registry.register_model_config(
    "deit3_reg4_h14",
    DeiT3,
    config={"patch_size": 14, **HUGE, "num_reg_tokens": 4, "drop_path_rate": 0.55},
)

registry.register_weights(
    "deit3_t16_il-common",
    {
        "description": "DeiT3 tiny model trained on the il-common dataset",
        "resolution": (256, 256),
        "formats": {
            "pt": {
                "file_size": 21.5,
                "sha256": "a04141c7f6c459ae075a48ccdee5b82d191bbaa82337673140c06ef82f0a8dc5",
            }
        },
        "net": {"network": "deit3_t16", "tag": "il-common"},
    },
)
registry.register_weights(
    "deit3_reg4_t16_il-common",
    {
        "description": "DeiT3 reg4 tiny model trained on the il-common dataset",
        "resolution": (256, 256),
        "formats": {
            "pt": {
                "file_size": 21.5,
                "sha256": "d26320462da64df6d62b307f7fb35d09c86a5f073002dfb24a51f014074e65c3",
            }
        },
        "net": {"network": "deit3_reg4_t16", "tag": "il-common"},
    },
)
