"""OpenHTR visual encoder with local and global feature mixing."""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.init import kaiming_normal_, ones_, trunc_normal_, zeros_

from .layers import DropPath, Identity

__all__ = ['OpenHTR_Encoder']


class ConvBNLayer(nn.Module):

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size=3,
        stride=1,
        padding=0,
        bias=False,
        groups=1,
        act=nn.GELU,
    ):
        super().__init__()
        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            groups=groups,
            bias=bias,
        )
        self.norm = nn.BatchNorm2d(out_channels)
        self.act = act()

    def forward(self, inputs):
        out = self.conv(inputs)
        out = self.norm(out)
        out = self.act(out)
        return out


class LastStage(nn.Module):

    def __init__(self, in_channels, out_channels, last_drop, out_char_num=0):
        super().__init__()
        self.last_conv = nn.Linear(in_channels, out_channels, bias=False)
        self.hardswish = nn.Hardswish()
        self.dropout = nn.Dropout(p=last_drop)

    def forward(self, x, sz):
        x = x.reshape(-1, sz[0], sz[1], x.shape[-1])
        x = x.mean(1)
        x = self.last_conv(x)
        x = self.hardswish(x)
        x = self.dropout(x)
        return x, [1, sz[1]]


class Feat2D(nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x, sz):
        C = x.shape[-1]
        x = x.transpose(1, 2).reshape(-1, C, sz[0], sz[1])
        return x, sz


def build_rope_2d(head_dim, H, W, device, dtype=torch.float32):
    """2D RoPE angles for a h-major flattened grid (H*W, w fastest).

    Returns cos/sin of shape [N, head_dim // 2] for the interleaved
    rotation convention (token dims are rotated in (2i, 2i+1) pairs).
    The head dim is split evenly: first half of the rotation frequencies
    encode the H axis, second half the W axis.
    """
    half = head_dim // 2
    assert half % 2 == 0, 'head_dim must be divisible by 4 for 2D RoPE'
    inv_freq = 1.0 / (10000.0**(torch.arange(0, half, 2, device=device,
                                             dtype=dtype) / half))
    h_pos = torch.arange(H, device=device, dtype=dtype)
    w_pos = torch.arange(W, device=device, dtype=dtype)
    ang_h = h_pos[:, None] * inv_freq[None, :]  # [H, half//2]
    ang_w = w_pos[:, None] * inv_freq[None, :]  # [W, half//2]
    ang_h = ang_h.repeat_interleave(W, dim=0)  # [H*W, half//2] (h-major)
    ang_w = ang_w.repeat(H, 1)  # [H*W, half//2]
    angles = torch.cat([ang_h, ang_w], dim=1)  # [N, half]
    return angles.cos(), angles.sin()


def rotate_interleaved(x, cos, sin):
    """RoPE rotation, interleaved-pair convention.

    x: [B, heads, N, D]; cos/sin: [N, D//2] (broadcastable).
    """
    x1 = x[..., 0::2]
    x2 = x[..., 1::2]
    out1 = x1 * cos - x2 * sin
    out2 = x1 * sin + x2 * cos
    return torch.stack([out1, out2], dim=-1).flatten(-2)


class RMSNorm(nn.Module):

    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        dtype = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * x.to(dtype)


class SwiGLU(nn.Module):

    def __init__(self, dim, hidden_dim, drop=0.0):
        super().__init__()
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)  # gate
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)  # up
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)  # down
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        return self.drop(self.w2(F.silu(self.w1(x)) * self.w3(x)))


class ConvNeXtBlock(nn.Module):
    """ConvNeXt block operating on [B, C, H, W] feature maps."""

    def __init__(self, dim, kernel_size=7, mlp_ratio=4, drop_path=0.0):
        super().__init__()
        self.dwconv = nn.Conv2d(dim,
                                dim,
                                kernel_size,
                                1,
                                kernel_size // 2,
                                groups=dim)
        self.norm = nn.LayerNorm(dim, eps=1e-6)
        hidden = int(dim * mlp_ratio)
        self.pw1 = nn.Linear(dim, hidden)
        self.act = nn.GELU()
        self.pw2 = nn.Linear(hidden, dim)
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()

    def forward(self, x):
        B, C, H, W = x.shape
        y = self.dwconv(x).permute(0, 2, 3, 1)  # B H W C
        y = self.norm(y)
        y = self.pw2(self.act(self.pw1(y)))
        y = y.permute(0, 3, 1, 2)
        return x + self.drop_path(y)


class GlobalAttentionBlock(nn.Module):
    """Pre-norm attention block with 2D RoPE, RMSNorm and SwiGLU.

    Input/output: [B, C, H, W]; attention runs on the flattened (h-major)
    sequence so global mixing covers the whole line.
    """

    def __init__(self,
                 dim,
                 num_heads,
                 mlp_ratio=2.67,
                 drop=0.0,
                 attn_drop=0.0,
                 drop_path=0.0,
                 qkv_bias=True):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        assert self.head_dim * num_heads == dim
        assert self.head_dim % 4 == 0, 'head_dim must be divisible by 4'
        self.norm1 = RMSNorm(dim)
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(drop)
        self.attn_drop = attn_drop
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
        self.norm2 = RMSNorm(dim)
        hidden = max(64, int(dim * mlp_ratio) // 64 * 64)
        self.mlp = SwiGLU(dim, hidden, drop=drop)

    def forward(self, x):
        # x: [B, C, H, W]
        B, C, H, W = x.shape
        y = x.flatten(2).transpose(1, 2)  # B N C
        y2 = self.norm1(y)
        qkv = self.qkv(y2).reshape(B, -1, 3, self.num_heads,
                                   self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        cos, sin = build_rope_2d(self.head_dim, H, W, x.device,
                                 q.dtype)
        cos = cos.to(q.dtype)[None, None]  # 1 1 N hd/2
        sin = sin.to(q.dtype)[None, None]
        q = rotate_interleaved(q, cos, sin)
        k = rotate_interleaved(k, cos, sin)
        y = F.scaled_dot_product_attention(
            q, k, v, dropout_p=self.attn_drop if self.training else 0.0)
        y = y.transpose(1, 2).reshape(B, -1, C)
        y = self.proj_drop(self.proj(y))
        x = x.flatten(2).transpose(1, 2) + self.drop_path(y)
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x.transpose(1, 2).reshape(B, C, H, W)


class EncoderStage(nn.Module):

    def __init__(self,
                 dim=64,
                 out_dim=256,
                 depth=3,
                 mixer=['Conv'] * 3,
                 kernel_sizes=[7] * 3,
                 sub_k=[2, 1],
                 num_heads=2,
                 mlp_ratio=4,
                 attn_mlp_ratio=2.67,
                 qkv_bias=True,
                 drop_rate=0.0,
                 attn_drop_rate=0.0,
                 drop_path=[0.1] * 3,
                 downsample=None,
                 **kwargs):
        super().__init__()
        self.blocks = nn.ModuleList()
        for i in range(depth):
            if mixer[i] == 'Conv':
                self.blocks.append(
                    ConvNeXtBlock(dim=dim,
                                  kernel_size=kernel_sizes[i],
                                  mlp_ratio=mlp_ratio,
                                  drop_path=drop_path[i]))
            else:
                self.blocks.append(
                    GlobalAttentionBlock(
                        dim=dim,
                        num_heads=num_heads,
                        mlp_ratio=attn_mlp_ratio,
                        drop=drop_rate,
                        attn_drop=attn_drop_rate,
                        drop_path=drop_path[i],
                        qkv_bias=qkv_bias,
                    ))

        if downsample:
            self.downsample = EncoderDownsample(dim, out_dim, stride=sub_k)
        else:
            self.downsample = None

    def forward(self, x, sz):
        for blk in self.blocks:
            x = blk(x)
        if self.downsample is not None:
            x, sz = self.downsample(x, sz)
        return x, sz


class EncoderDownsample(nn.Module):
    """Strided 3x3 conv + RMSNorm, operating on [B, C, H, W] (stays 4D)."""

    def __init__(self, in_channels, out_channels, stride=[2, 1]):
        super().__init__()
        self.conv = nn.Conv2d(in_channels,
                              out_channels,
                              kernel_size=3,
                              stride=stride,
                              padding=1)
        self.norm = RMSNorm(out_channels)

    def forward(self, x, sz):
        x = self.conv(x)
        C, H, W = x.shape[1:]
        y = self.norm(x.permute(0, 2, 3, 1))  # B H W C
        return y.permute(0, 3, 1, 2).contiguous(), [H, W]


class OpenHTR_Encoder(nn.Module):
    """Encode local strokes and global context as a sequence of visual tokens."""

    def __init__(self,
                 max_sz=[32, 128],
                 in_channels=3,
                 out_channels=192,
                 depths=[3, 6, 3],
                 dims=[64, 128, 256],
                 mixer=[['Conv'] * 3, ['Conv'] * 3 + ['Global'] * 3,
                        ['Global'] * 3],
                 sub_k=[[1, 1], [2, 1], [1, 1]],
                 num_heads=[2, 4, 8],
                 mlp_ratio=4,
                 attn_mlp_ratio=2.67,
                 qkv_bias=True,
                 drop_rate=0.0,
                 last_drop=0.1,
                 attn_drop_rate=0.0,
                 drop_path_rate=0.1,
                 last_stage=False,
                 feat2d=False,
                 num_convs=None,  # unused, kept for config compat
                 kernel_sizes=[[7] * 3, [7] * 3 + [7] * 3, [7] * 3],
                 **kwargs):
        super().__init__()
        num_stages = len(depths)
        self.num_features = dims[-1]

        # Two stride-2 convolutions reduce the input resolution by a factor of four.
        self.pope = nn.Sequential(
            ConvBNLayer(in_channels=in_channels,
                        out_channels=dims[0] // 2,
                        kernel_size=3,
                        stride=2,
                        padding=1,
                        act=nn.GELU),
            ConvBNLayer(in_channels=dims[0] // 2,
                        out_channels=dims[0],
                        kernel_size=3,
                        stride=2,
                        padding=1,
                        act=nn.GELU),
        )

        dpr = np.linspace(0, drop_path_rate, sum(depths))

        self.stages = nn.ModuleList()
        for i_stage in range(num_stages):
            stage = EncoderStage(
                dim=dims[i_stage],
                out_dim=dims[i_stage + 1] if i_stage < num_stages - 1 else 0,
                depth=depths[i_stage],
                mixer=mixer[i_stage],
                kernel_sizes=kernel_sizes[i_stage] if len(
                    kernel_sizes[i_stage]) == len(mixer[i_stage]) else [7] *
                len(mixer[i_stage]),
                sub_k=sub_k[i_stage],
                num_heads=num_heads[i_stage],
                mlp_ratio=mlp_ratio,
                attn_mlp_ratio=attn_mlp_ratio,
                qkv_bias=qkv_bias,
                drop=drop_rate,
                attn_drop=attn_drop_rate,
                drop_path=dpr[sum(depths[:i_stage]):sum(depths[:i_stage + 1])],
                downsample=False if i_stage == num_stages - 1 else True,
            )
            self.stages.append(stage)

        self.out_channels = self.num_features
        self.last_stage = last_stage
        self.last_stage_layer = None
        self.feat2d_layer = None
        if last_stage:
            self.out_channels = out_channels
            self.last_stage_layer = LastStage(self.num_features, out_channels,
                                              last_drop)
        if feat2d:
            self.feat2d_layer = Feat2D()
        self.apply(self._init_weights)

    def _init_weights(self, m: nn.Module):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, mean=0, std=0.02)
            if m.bias is not None:
                zeros_(m.bias)
        if isinstance(m, nn.LayerNorm):
            zeros_(m.bias)
            ones_(m.weight)
        if isinstance(m, nn.Conv2d):
            kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'patch_embed', 'downsample', 'pos_embed'}

    def forward(self, x):
        if len(x.shape) == 5:
            x = x.flatten(0, 1)
        sz = x.shape[2:]
        x = self.pope(x)
        sz = [sz[0] // 4, sz[1] // 4]
        for stage in self.stages:
            x, sz = stage(x, sz)
        # All EncoderStage blocks keep [B, C, H, W]; convert once at the end.
        if self.last_stage_layer is not None:
            x, sz = self.last_stage_layer(x.flatten(2).transpose(1, 2), sz)
        elif self.feat2d_layer is not None:
            pass  # already [B, C, H, W]
        else:
            x = x.flatten(2).transpose(1, 2)  # B N C
        return x