from __future__ import annotations

from functools import partial
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F


class Fp32LayerNorm(nn.LayerNorm):
    """LayerNorm computed in fp32 for mixed-precision stability."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = F.layer_norm(
            x.float(),
            self.normalized_shape,
            self.weight.float() if self.weight is not None else None,
            self.bias.float() if self.bias is not None else None,
            self.eps,
        )
        return y.to(dtype=x.dtype)


def drop_path(
    x: torch.Tensor, drop_prob: float = 0.0, training: bool = False
) -> torch.Tensor:
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1.0 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()
    return x.div(keep_prob) * random_tensor


class DropPath(nn.Module):
    def __init__(self, drop_prob: float = 0.0):
        super().__init__()
        self.drop_prob = float(drop_prob)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return drop_path(x, self.drop_prob, self.training)


class Mlp(nn.Module):
    def __init__(
        self,
        in_features: int,
        hidden_features: Optional[int] = None,
        out_features: Optional[int] = None,
        act_layer: type[nn.Module] = nn.GELU,
        drop: float = 0.0,
        bias: bool = True,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
        self.drop = nn.Dropout(drop)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.drop(self.fc2(self.act(self.fc1(x))))


class GatedMlp(nn.Module):
    def __init__(
        self,
        in_features: int,
        hidden_features: Optional[int] = None,
        out_features: Optional[int] = None,
        act_layer: type[nn.Module] = nn.SiLU,
        drop: float = 0.0,
        bias: bool = True,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = int(2 * (hidden_features or in_features) / 3)
        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
        self.fc3 = nn.Linear(in_features, hidden_features, bias=bias)
        self.drop = nn.Dropout(drop)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.drop(self.fc2(self.act(self.fc1(x)) * self.fc3(x)))


class SelfAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        qkv_bias: bool = False,
        proj_bias: bool = False,
        proj_drop: float = 0.0,
        qk_norm: bool = True,
        norm_layer: partial = partial(
            Fp32LayerNorm, bias=False, elementwise_affine=False
        ),
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads

        self.wq = nn.Linear(dim, dim, bias=qkv_bias)
        self.wk = nn.Linear(dim, dim, bias=qkv_bias)
        self.wv = nn.Linear(dim, dim, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim, bias=proj_bias)
        self.proj_drop = nn.Dropout(proj_drop)

        if qk_norm:
            self.q_norm = norm_layer(self.head_dim)
            self.k_norm = norm_layer(self.head_dim)
        else:
            self.q_norm = nn.Identity()
            self.k_norm = nn.Identity()

    def forward(
        self, x: torch.Tensor, block_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        b, n, d = x.shape
        q = self.wq(x).reshape(b, n, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).reshape(b, n, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).reshape(b, n, self.num_heads, self.head_dim).transpose(1, 2)

        q = self.q_norm(q)
        k = self.k_norm(k)

        y = F.scaled_dot_product_attention(q, k, v, attn_mask=block_mask)
        y = y.transpose(1, 2).reshape(b, n, d)
        return self.proj_drop(self.proj(y))


class Block(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: Optional[int] = None,
        head_dim: Optional[int] = None,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = False,
        proj_bias: bool = False,
        mlp_bias: bool = False,
        drop: float = 0.0,
        drop_path_rate: float = 0.0,
        act_layer: type[nn.Module] = nn.SiLU,
        norm_layer: partial = partial(
            Fp32LayerNorm, bias=False, elementwise_affine=False
        ),
        gated_mlp: bool = True,
        qk_norm: bool = True,
    ):
        super().__init__()
        nh = num_heads or (dim // int(head_dim))
        self.norm1 = norm_layer(dim)
        self.attn = SelfAttention(
            dim=dim,
            num_heads=nh,
            qkv_bias=qkv_bias,
            proj_bias=proj_bias,
            proj_drop=drop,
            qk_norm=qk_norm,
            norm_layer=norm_layer,
        )
        self.drop_path = (
            DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
        )
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        mlp_layer = GatedMlp if gated_mlp else Mlp
        self.mlp = mlp_layer(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            bias=mlp_bias,
            drop=drop,
        )

    def forward(
        self, x: torch.Tensor, block_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        x = x + self.drop_path(self.attn(self.norm1(x), block_mask=block_mask))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


class Transformer(nn.Module):
    """OAT-style transformer stack used by the register encoder."""

    def __init__(
        self,
        dim: int,
        depth: int,
        head_dim: int = 64,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = False,
        proj_bias: bool = False,
        mlp_bias: bool = False,
        drop: float = 0.0,
        drop_path_rate: float = 0.0,
        act_layer: type[nn.Module] = nn.SiLU,
        norm_layer: partial = partial(
            Fp32LayerNorm, bias=False, elementwise_affine=False
        ),
        gated_mlp: bool = True,
        qk_norm: bool = True,
        weight_init_style: str = "xavier",
        zero_init_query_proj: bool = False,
    ):
        super().__init__()
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
        self.blocks = nn.ModuleList(
            [
                Block(
                    dim=dim,
                    head_dim=head_dim,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    proj_bias=proj_bias,
                    mlp_bias=mlp_bias,
                    drop=drop,
                    drop_path_rate=dpr[i],
                    act_layer=act_layer,
                    norm_layer=norm_layer,
                    gated_mlp=gated_mlp,
                    qk_norm=qk_norm,
                )
                for i in range(depth)
            ]
        )
        self.weight_init_style = weight_init_style
        self.zero_init_query_proj = zero_init_query_proj
        self._init_weights()

    def _init_weights(self) -> None:
        for name, m in self.named_modules():
            if isinstance(m, nn.Linear):
                if "wq" in name and self.zero_init_query_proj:
                    nn.init.constant_(m.weight, 0)
                elif self.weight_init_style == "xavier":
                    nn.init.xavier_uniform_(m.weight)
                elif self.weight_init_style == "trunc_normal":
                    nn.init.trunc_normal_(m.weight, std=0.02)
                else:
                    raise ValueError(
                        f"Unsupported weight init: {self.weight_init_style}"
                    )
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.LayerNorm):
                if m.weight is not None:
                    nn.init.constant_(m.weight, 1.0)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(
        self, x: torch.Tensor, block_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        for block in self.blocks:
            x = block(x, block_mask=block_mask)
        return x
