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

from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput

try:
    from .configuration_rinalmo import RiNALMoConfig
except ImportError:
    from configuration_rinalmo import RiNALMoConfig


def _rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def _apply_rotary_pos_emb(q, k, cos, sin):
    cos = cos.to(q.dtype)
    sin = sin.to(q.dtype)
    return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin)


class RotaryPositionEmbedding(nn.Module):
    def __init__(self, dim: int, base: int = 10000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self._seq_len_cached = None
        self._cos_cached = None
        self._sin_cached = None

    def _update_cache(self, seq_len: int, device, dtype):
        if seq_len != self._seq_len_cached:
            self._seq_len_cached = seq_len
            t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1)
            self._cos_cached = emb.cos()[None, None, :, :]
            self._sin_cached = emb.sin()[None, None, :, :]

    def forward(self, q, k):
        self._update_cache(q.shape[-2], q.device, q.dtype)
        return _apply_rotary_pos_emb(q, k, self._cos_cached, self._sin_cached)


class RiNALMoAttention(nn.Module):
    def __init__(self, config: RiNALMoConfig):
        super().__init__()
        self.embed_dim = config.embed_dim
        self.num_heads = config.num_heads
        self.head_dim = config.embed_dim // config.num_heads

        self.qkv_proj = nn.Linear(config.embed_dim, 3 * config.embed_dim, bias=False)
        self.out_proj = nn.Linear(config.embed_dim, config.embed_dim, bias=False)
        self.attn_dropout = nn.Dropout(p=config.attention_dropout)

        if config.use_rot_emb:
            self.rotary_emb = RotaryPositionEmbedding(self.head_dim, base=config.rope_base)
        else:
            self.rotary_emb = None

    def forward(self, x, key_padding_mask=None, output_attentions=False):
        B, T, _ = x.shape

        qkv = self.qkv_proj(x)
        q, k, v = qkv.chunk(3, dim=-1)

        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)

        if self.rotary_emb is not None:
            q, k = self.rotary_emb(q, k)

        scale = math.sqrt(self.head_dim)
        attn = torch.matmul(q, k.transpose(-1, -2)) / scale

        if key_padding_mask is not None:
            attn = attn.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf"))

        attn = attn.softmax(dim=-1)
        attn_weights = attn if output_attentions else None
        attn = self.attn_dropout(attn)

        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, T, self.embed_dim)
        out = self.out_proj(out)
        return out, attn_weights


class RiNALMoSdpaAttention(RiNALMoAttention):
    def forward(self, x, key_padding_mask=None, output_attentions=False):
        if output_attentions:
            return super().forward(x, key_padding_mask, output_attentions=True)

        B, T, _ = x.shape

        qkv = self.qkv_proj(x)
        q, k, v = qkv.chunk(3, dim=-1)

        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)

        if self.rotary_emb is not None:
            q, k = self.rotary_emb(q, k)

        attn_mask = None
        if key_padding_mask is not None:
            attn_mask = torch.zeros(B, 1, 1, T, dtype=q.dtype, device=q.device)
            attn_mask = attn_mask.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf"))

        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0)
        out = out.transpose(1, 2).contiguous().view(B, T, self.embed_dim)
        out = self.out_proj(out)
        return out, None


class RiNALMoFlashAttention2(RiNALMoAttention):
    def forward(self, x, key_padding_mask=None, output_attentions=False):
        if output_attentions:
            return super().forward(x, key_padding_mask, output_attentions=True)

        try:
            from flash_attn import flash_attn_func, flash_attn_varlen_func
            from flash_attn.bert_padding import pad_input, unpad_input
        except ImportError as e:
            raise ImportError(
                "flash_attn is required for attn_implementation='flash_attention_2'. "
                "Install with: pip install flash-attn --no-build-isolation"
            ) from e

        B, T, _ = x.shape

        qkv = self.qkv_proj(x)
        q, k, v = qkv.chunk(3, dim=-1)

        q = q.view(B, T, self.num_heads, self.head_dim)
        k = k.view(B, T, self.num_heads, self.head_dim)
        v = v.view(B, T, self.num_heads, self.head_dim)

        if self.rotary_emb is not None:
            q_t = q.transpose(1, 2)
            k_t = k.transpose(1, 2)
            q_t, k_t = self.rotary_emb(q_t, k_t)
            q = q_t.transpose(1, 2)
            k = k_t.transpose(1, 2)

        orig_dtype = q.dtype
        if q.dtype not in (torch.float16, torch.bfloat16):
            q = q.to(torch.bfloat16)
            k = k.to(torch.bfloat16)
            v = v.to(torch.bfloat16)

        if key_padding_mask is not None and key_padding_mask.any():
            attend_mask = ~key_padding_mask
            q_unpad, indices, cu_seqlens, max_seqlen, _ = unpad_input(q, attend_mask)
            k_unpad, *_ = unpad_input(k, attend_mask)
            v_unpad, *_ = unpad_input(v, attend_mask)

            out_unpad = flash_attn_varlen_func(
                q_unpad, k_unpad, v_unpad,
                cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
                max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
                causal=False,
            )
            out = pad_input(out_unpad.view(-1, self.embed_dim), indices, B, T)
        else:
            out = flash_attn_func(q, k, v, causal=False)
            out = out.view(B, T, self.embed_dim)

        out = out.to(orig_dtype)
        out = self.out_proj(out)
        return out, None


RINALMO_ATTENTION_CLASSES = {
    "eager": RiNALMoAttention,
    "sdpa": RiNALMoSdpaAttention,
    "flash_attention_2": RiNALMoFlashAttention2,
}


class RiNALMoSwiGLU(nn.Module):
    def __init__(self, embed_dim: int, ffn_dim: int):
        super().__init__()
        self.linear = nn.Linear(embed_dim, ffn_dim, bias=True)
        self.linear_gate = nn.Linear(embed_dim, ffn_dim, bias=True)
        self.beta = nn.Parameter(torch.ones(1))

    def forward(self, x):
        gate = self.linear(x)
        swish = gate * torch.sigmoid(self.beta * gate)
        return swish * self.linear_gate(x)


class TokenDropout(nn.Module):
    def __init__(self, active: bool, mask_ratio: float, mask_tkn_prob: float,
                 mask_idx: int, padding_idx: int):
        super().__init__()
        self.active = active
        self.mask_ratio_train = mask_ratio * mask_tkn_prob
        self.mask_idx = mask_idx
        self.padding_idx = padding_idx

    def forward(self, x, tokens):
        if not self.active:
            return x
        pad_mask = tokens.eq(self.padding_idx)
        src_lens = (~pad_mask).sum(dim=-1).to(x.dtype)
        x = torch.where((tokens == self.mask_idx).unsqueeze(-1), torch.zeros_like(x), x)
        mask_ratio_obs = (tokens == self.mask_idx).sum(dim=-1).to(x.dtype) / src_lens
        scale = (1.0 - self.mask_ratio_train) / (1.0 - mask_ratio_obs)
        x = x * scale[:, None, None]
        return x


class RiNALMoLayer(nn.Module):
    def __init__(self, config: RiNALMoConfig):
        super().__init__()
        ffn_dim = int(2 / 3 * config.transition_factor * config.embed_dim)
        attn_cls = RINALMO_ATTENTION_CLASSES[getattr(config, "_attn_implementation", "eager")]

        self.attn_layer_norm = nn.LayerNorm(config.embed_dim)
        self.attn = attn_cls(config)
        self.out_layer_norm = nn.LayerNorm(config.embed_dim)
        self.ffn = RiNALMoSwiGLU(config.embed_dim, ffn_dim)
        self.ffn_dropout = nn.Dropout(p=config.transition_dropout)
        self.ffn_down = nn.Linear(ffn_dim, config.embed_dim, bias=True)
        self.residual_dropout_1 = nn.Dropout(p=config.residual_dropout)
        self.residual_dropout_2 = nn.Dropout(p=config.residual_dropout)

    def forward(self, x, key_padding_mask=None, output_attentions=False):
        x = self.attn_layer_norm(x)
        attn_out, attn_weights = self.attn(x, key_padding_mask=key_padding_mask,
                                            output_attentions=output_attentions)
        x = x + self.residual_dropout_1(attn_out)

        residual = x
        x = self.out_layer_norm(x)
        x = residual + self.residual_dropout_2(self.ffn_down(self.ffn_dropout(self.ffn(x))))

        return x, attn_weights


class RiNALMoPreTrainedModel(PreTrainedModel):
    config_class = RiNALMoConfig
    base_model_prefix = "model"
    _supports_sdpa = True
    _supports_flash_attn_2 = True

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


class RiNALMoModel(RiNALMoPreTrainedModel):
    def __init__(self, config: RiNALMoConfig):
        super().__init__(config)
        self.embedding = nn.Embedding(config.vocab_size, config.embed_dim, padding_idx=config.padding_idx)
        self.token_dropout = TokenDropout(
            active=config.token_dropout_active,
            mask_ratio=config.mask_ratio,
            mask_tkn_prob=config.mask_tkn_prob,
            mask_idx=config.mask_idx,
            padding_idx=config.padding_idx,
        )
        self.layers = nn.ModuleList([RiNALMoLayer(config) for _ in range(config.num_layers)])
        self.final_layer_norm = nn.LayerNorm(config.embed_dim)
        self.post_init()

    def forward(
        self,
        input_ids,
        attention_mask=None,
        output_hidden_states=None,
        output_attentions=None,
        return_dict=None,
    ):
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if attention_mask is not None:
            key_padding_mask = attention_mask.eq(0)
        else:
            key_padding_mask = input_ids.eq(self.config.padding_idx)

        x = self.embedding(input_ids)
        x = self.token_dropout(x, input_ids)

        all_hidden_states = []
        all_attentions = []
        if output_hidden_states:
            all_hidden_states.append(x)

        for layer in self.layers:
            x, attn_weights = layer(x, key_padding_mask=key_padding_mask,
                                    output_attentions=output_attentions)
            if output_hidden_states:
                all_hidden_states.append(x)
            if output_attentions:
                all_attentions.append(attn_weights)

        x = self.final_layer_norm(x)

        return BaseModelOutput(
            last_hidden_state=x,
            hidden_states=tuple(all_hidden_states) if output_hidden_states else None,
            attentions=tuple(all_attentions) if output_attentions else None,
        )


class RiNALMoForMaskedLM(RiNALMoPreTrainedModel):
    def __init__(self, config: RiNALMoConfig):
        super().__init__(config)
        self.model = RiNALMoModel(config)
        self.lm_head = RiNALMoLMHead(config)
        self.post_init()

    def forward(
        self,
        input_ids,
        attention_mask=None,
        labels=None,
        output_hidden_states=None,
        output_attentions=None,
        return_dict=None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        out = self.model(input_ids, attention_mask=attention_mask,
                         output_hidden_states=output_hidden_states,
                         output_attentions=output_attentions, return_dict=return_dict)
        logits = self.lm_head(out.last_hidden_state)
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size),
                                   labels.view(-1), ignore_index=-100)
        return MaskedLMOutput(loss=loss, logits=logits,
                              hidden_states=out.hidden_states,
                              attentions=out.attentions)


class RiNALMoLMHead(nn.Module):
    def __init__(self, config: RiNALMoConfig):
        super().__init__()
        self.linear1 = nn.Linear(config.embed_dim, config.embed_dim)
        self.layer_norm = nn.LayerNorm(config.embed_dim)
        self.linear2 = nn.Linear(config.embed_dim, config.vocab_size)

    def forward(self, x):
        x = self.linear1(x)
        x = F.gelu(x)
        x = self.layer_norm(x)
        x = self.linear2(x)
        return x