"""OpenHTR autoregressive decoder with visual cross-attention and KV caching."""

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

__all__ = ['OpenHTR_Decoder']


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)))


def build_rope_1d(head_dim, positions, device, dtype=torch.float32,
                  base=10000.0):
    """RoPE cos/sin for the given absolute positions (interleaved pairs).

    positions: LongTensor [T] or an int for a single position.
    Returns cos/sin of shape [T, head_dim // 2].
    """
    half = head_dim // 2
    if not torch.is_tensor(positions):
        positions = torch.tensor([positions], device=device)
    inv_freq = 1.0 / (base**(torch.arange(0, head_dim, 2, device=device,
                                          dtype=dtype) / head_dim))
    angles = positions.to(dtype)[:, None] * inv_freq[None, :]  # T, half
    return angles.cos(), angles.sin()


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

    x: [B, H, T, D]; cos/sin: [T, 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 NeoSelfAttention(nn.Module):
    """Fused-QKV self-attention with RoPE and KV-cache support."""

    def __init__(self, dim, num_heads, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        assert self.head_dim * num_heads == dim
        self.qkv = nn.Linear(dim, dim * 3, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.dropout = dropout

    def forward(self, x, cos=None, sin=None, kv_cache=None, is_causal=False):
        B, T, D = x.shape
        qkv = self.qkv(x).reshape(B, T, 3, self.num_heads,
                                  self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        if cos is not None:
            cos = cos.to(q.dtype)[None, None]  # 1 1 T hd/2
            sin = sin.to(q.dtype)[None, None]
            q = rotate_interleaved(q, cos, sin)
            k = rotate_interleaved(k, cos, sin)
        if kv_cache is not None and kv_cache[0] is not None:
            k = torch.cat([kv_cache[0], k], dim=2)
            v = torch.cat([kv_cache[1], v], dim=2)
        y = F.scaled_dot_product_attention(
            q,
            k,
            v,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=is_causal,
        )
        y = y.transpose(1, 2).reshape(B, T, D)
        return self.out_proj(y), k, v


class NeoCrossAttention(nn.Module):
    """Cross-attention against a fixed memory (K/V precomputed once)."""

    def __init__(self, dim, num_heads, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.embed_dim = dim
        self.q = nn.Linear(dim, dim, bias=False)
        self.kv = nn.Linear(dim, dim * 2, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.dropout = dropout

    def project_kv(self, memory):
        """memory: [B, N, D] -> (k, v) each [B, H, N, head_dim]."""
        B, N, _ = memory.shape
        kv = self.kv(memory).reshape(B, N, 2, self.num_heads,
                                     self.head_dim).permute(2, 0, 3, 1, 4)
        return kv.unbind(0)

    def forward(self, x, mem_kv):
        B, T, D = x.shape
        q = self.q(x).reshape(B, T, self.num_heads,
                              self.head_dim).transpose(1, 2)
        k, v = mem_kv
        y = F.scaled_dot_product_attention(
            q,
            k,
            v,
            dropout_p=self.dropout if self.training else 0.0,
        )
        y = y.transpose(1, 2).reshape(B, T, D)
        return self.out_proj(y)


class NeoDecoderBlock(nn.Module):
    """Pre-norm block: RoPE self-attn -> cross-attn -> SwiGLU FFN."""

    def __init__(self,
                 d_model,
                 nhead,
                 mlp_ratio=2.67,
                 attention_dropout_rate=0.0,
                 residual_dropout_rate=0.1,
                 rope=True):
        super().__init__()
        hidden = max(64, int(d_model * mlp_ratio) // 64 * 64)
        self.self_attn = NeoSelfAttention(d_model,
                                          nhead,
                                          dropout=attention_dropout_rate)
        self.norm1 = RMSNorm(d_model)
        self.cross_attn = NeoCrossAttention(d_model,
                                            nhead,
                                            dropout=attention_dropout_rate)
        self.norm2 = RMSNorm(d_model)
        self.ffn = SwiGLU(d_model, hidden, drop=residual_dropout_rate)
        self.norm3 = RMSNorm(d_model)
        self.dropout = nn.Dropout(residual_dropout_rate)
        self.rope = rope

    def _rope(self, positions, device, dtype):
        if not self.rope:
            return None, None
        return build_rope_1d(self.self_attn.head_dim, positions, device,
                             dtype)

    def forward_train(self, x, mem_kv):
        """Full-sequence causal pass. x: [B, T, D]."""
        T = x.shape[1]
        pos = torch.arange(T, device=x.device)
        cos, sin = self._rope(pos, x.device, x.dtype)
        a, _, _ = self.self_attn(self.norm1(x), cos, sin, is_causal=T > 1)
        x = x + self.dropout(a)
        x = x + self.dropout(self.cross_attn(self.norm2(x), mem_kv))
        x = x + self.dropout(self.ffn(self.norm3(x)))
        return x

    def step(self, x, mem_kv, kv_cache, pos):
        """Single-token incremental pass.

        x: [B, 1, D]; kv_cache: [k, v] list updated in place;
        pos: absolute RoPE position of the incoming token.
        """
        cos, sin = self._rope(pos, x.device, x.dtype)
        a, k, v = self.self_attn(self.norm1(x), cos, sin, kv_cache=kv_cache)
        kv_cache[0], kv_cache[1] = k, v
        x = x + self.dropout(a)
        x = x + self.dropout(self.cross_attn(self.norm2(x), mem_kv))
        x = x + self.dropout(self.ffn(self.norm3(x)))
        return x


class NeoEncoderBlock(nn.Module):
    """Pre-norm RoPE self-attention + SwiGLU block (optional memory side)."""

    def __init__(self,
                 d_model,
                 nhead,
                 mlp_ratio=2.67,
                 attention_dropout_rate=0.0,
                 residual_dropout_rate=0.1):
        super().__init__()
        hidden = max(64, int(d_model * mlp_ratio) // 64 * 64)
        self.self_attn = NeoSelfAttention(d_model,
                                          nhead,
                                          dropout=attention_dropout_rate)
        self.norm1 = RMSNorm(d_model)
        self.ffn = SwiGLU(d_model, hidden, drop=residual_dropout_rate)
        self.norm2 = RMSNorm(d_model)
        self.dropout = nn.Dropout(residual_dropout_rate)

    def forward(self, x):
        T = x.shape[1]
        pos = torch.arange(T, device=x.device)
        cos, sin = build_rope_1d(self.self_attn.head_dim, pos, x.device,
                                 x.dtype)
        a, _, _ = self.self_attn(self.norm1(x), cos, sin, is_causal=False)
        x = x + self.dropout(a)
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x


class OpenHTR_Decoder(nn.Module):
    """Decode text from visual tokens with causal self-attention and cross-attention."""

    def __init__(
        self,
        in_channels,
        out_channels,
        nhead=None,
        num_encoder_layers=0,
        beam_size=0,  # unused, greedy decoding only; kept for config compat
        num_decoder_layers=2,
        max_len=25,
        attention_dropout_rate=0.0,
        residual_dropout_rate=0.1,
        mlp_ratio=2.67,
        **kwargs,
    ):
        super().__init__()
        self.out_channels = out_channels
        self.ignore_index = out_channels - 1
        self.bos = out_channels - 2
        self.eos = 0
        self.max_len = max_len
        d_model = in_channels
        nhead = nhead if nhead is not None else d_model // 64
        assert d_model % nhead == 0, \
            f'd_model ({d_model}) must be divisible by nhead ({nhead})'

        self.embedding = nn.Embedding(out_channels, d_model, padding_idx=0)

        if num_encoder_layers is not None and num_encoder_layers > 0:
            self.encoder = nn.ModuleList([
                NeoEncoderBlock(
                    d_model,
                    nhead,
                    mlp_ratio=mlp_ratio,
                    attention_dropout_rate=attention_dropout_rate,
                    residual_dropout_rate=residual_dropout_rate,
                ) for _ in range(num_encoder_layers)
            ])
        else:
            self.encoder = None

        self.decoder = nn.ModuleList([
            NeoDecoderBlock(
                d_model,
                nhead,
                mlp_ratio=mlp_ratio,
                attention_dropout_rate=attention_dropout_rate,
                residual_dropout_rate=residual_dropout_rate,
            ) for _ in range(num_decoder_layers)
        ])

        self.norm_out = RMSNorm(d_model)
        self.tgt_word_prj = nn.Linear(d_model, self.out_channels - 2,
                                      bias=False)
        self.d_model = d_model
        self.apply(self._init_weights)
        nn.init.zeros_(self.embedding.weight.data[0])  # padding idx

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def _encode(self, src):
        if self.encoder is not None:
            for layer in self.encoder:
                src = layer(src)
        return src

    def forward_train(self, src, tgt):
        tgt = tgt[:, :-1]
        memory = self._encode(src)
        mem_kv = [layer.cross_attn.project_kv(memory)
                  for layer in self.decoder]
        x = self.embedding(tgt)
        for i, layer in enumerate(self.decoder):
            x = layer.forward_train(x, mem_kv[i])
        x = self.norm_out(x)
        return self.tgt_word_prj(x)

    def forward_test(self, src):
        bs = src.shape[0]
        device = src.device
        memory = self._encode(src)
        mem_kv = [layer.cross_attn.project_kv(memory)
                  for layer in self.decoder]

        dec_seq = torch.full((bs, self.max_len + 1),
                             self.ignore_index,
                             dtype=torch.int64,
                             device=device)
        dec_seq[:, 0] = self.bos

        logits = []
        kv_caches = [[None, None] for _ in self.decoder]

        with torch.no_grad():
            for step in range(0, self.max_len):
                tok = dec_seq[:, step:step + 1]  # B 1
                x = self.embedding(tok)
                for i, layer in enumerate(self.decoder):
                    x = layer.step(x, mem_kv[i], kv_caches[i], pos=step)
                x = self.norm_out(x)
                word_prob = F.softmax(self.tgt_word_prj(x[:, -1, :]),
                                      dim=-1).unsqueeze(1)  # B 1 V
                logits.append(word_prob)
                dec_seq[:, step + 1] = word_prob.squeeze(1).argmax(-1)
                # stop once every sequence has emitted at least one EOS
                if (dec_seq == self.eos).any(dim=-1).all():
                    break

        logits = torch.cat(logits, dim=1)
        return logits

    def forward(self, src, data=None):
        if self.training:
            max_len = data[1].max()
            tgt = data[0][:, :2 + max_len]
            return self.forward_train(src, tgt)
        return self.forward_test(src)