import math
from typing import Optional

import torch
from torch import Tensor
from torch.nn import (
    BCEWithLogitsLoss,
    Conv1d,
    CrossEntropyLoss,
    Dropout,
    Embedding,
    Linear,
    Module,
    ModuleList,
    MSELoss,
    Parameter,
    SiLU,
    functional as F,
)
from transformers import PreTrainedModel
from transformers.modeling_outputs import (
    BaseModelOutput,
    MaskedLMOutput,
    SequenceClassifierOutput,
    TokenClassifierOutput,
)
from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS

try:
    from .iii_configuration import IIIBERTConfig
except ImportError:
    from iii_configuration import IIIBERTConfig

_CONFIG_FOR_DOC = "IIIBERTConfig"


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

    def forward(self, tensor):
        input_dtype = tensor.dtype
        tensor = tensor.to(torch.float32)
        variance = tensor.pow(2).mean(-1, keepdim=True)
        tensor = tensor * torch.rsqrt(variance + self.epsilon)
        return self.weight * tensor.to(input_dtype)


ALL_LAYERNORM_LAYERS.append(IIIRMSNorm)


class IIIMLP(Module):
    def __init__(self, config, dim=None, intermediate_dim=None):
        super().__init__()
        self.config = config
        self.dim = config.hidden_dim if dim is None else dim
        self.intermediate_dim = (
            config.intermediate_dim if intermediate_dim is None else intermediate_dim
        )

        self.gate = Linear(self.dim, self.intermediate_dim, bias=False)
        self.up = Linear(self.dim, self.intermediate_dim, bias=False)
        self.down = Linear(self.intermediate_dim, self.dim, bias=False)
        self.activation = SiLU()

    def forward(self, x):
        value = self.up(x)
        gate = self.activation(self.gate(x))
        return self.down(gate * value)


class IIIAttention(Module):
    def __init__(self, config: IIIBERTConfig):
        super().__init__()
        self.config = config

        self.hidden_dim = config.hidden_dim
        self.lora_dim = config.lora_dim
        self.num_heads = config.n_heads

        if self.lora_dim % self.num_heads != 0:
            raise ValueError(
                f"lora_dim ({self.lora_dim}) must be divisible by "
                f"n_heads ({self.num_heads})."
            )

        self.head_dim = self.lora_dim // self.num_heads

        self.qkv = Linear(self.hidden_dim, self.lora_dim * 3, bias=config.attention_bias)
        self.o = Linear(self.lora_dim, self.hidden_dim, bias=config.attention_bias)

        self.q_norm = IIIRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = IIIRMSNorm(self.head_dim, eps=config.rms_norm_eps)

    def _prepare_attention_mask(self, attention_mask, device):
        if attention_mask is None:
            return None

        mask = attention_mask.to(device)

        if mask.dim() == 2:
            mask = mask[:, None, None, :]
        elif mask.dim() == 3:
            mask = mask[:, None, :, :]
        elif mask.dim() != 4:
            raise ValueError("attention_mask must be 2D, 3D, or 4D.")

        if not torch.is_floating_point(mask) and mask.dtype != torch.bool:
            mask = mask != 0

        return mask

    def forward(self, hidden_states, attention_mask=None, output_attentions=False):
        batch_size, seq_len, _ = hidden_states.shape

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

        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

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

        mask = self._prepare_attention_mask(attention_mask, hidden_states.device)

        if mask is not None and mask.dtype != torch.bool:
            mask = mask.to(q.dtype)

        attention_weights = None

        if output_attentions:
            scale = self.head_dim ** -0.5
            scores = torch.matmul(q, k.transpose(-2, -1)) * scale

            if mask is not None:
                if mask.dtype == torch.bool:
                    scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
                else:
                    scores = scores + mask.to(scores.dtype)

            attention_weights = torch.softmax(scores, dim=-1)
            attn_output = torch.matmul(attention_weights, v)
        else:
            attn_output = F.scaled_dot_product_attention(
                q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
            )

        attn_output = (
            attn_output.transpose(1, 2)
            .contiguous()
            .view(batch_size, seq_len, self.lora_dim)
        )

        return self.o(attn_output), attention_weights


class GatedCNN(Module):
    def __init__(self, config: IIIBERTConfig, layer_id):
        super().__init__()
        self.config = config
        self.hidden_dim = config.hidden_dim
        self.window_size = config.layers_kernel_size[layer_id]
        self.dilation = config.layers_diliation_size[layer_id]

        if self.window_size <= 0:
            raise ValueError("GatedCNN requires window_size > 0.")
        if self.dilation <= 0:
            raise ValueError("GatedCNN requires dilation > 0.")

        self.gate_proj = Linear(self.hidden_dim, self.hidden_dim * 2, bias=False)

        self.conv = Conv1d(
            in_channels=self.hidden_dim,
            out_channels=self.hidden_dim,
            kernel_size=self.window_size,
            dilation=self.dilation,
            padding="same",
            groups=self.hidden_dim,
        )

        self.out_proj = Linear(self.hidden_dim, self.hidden_dim)

    def _gate(self, hidden_states):
        value, score = self.gate_proj(hidden_states).chunk(2, dim=-1)
        return value * torch.sigmoid(score)

    @staticmethod
    def _keep_mask(attention_mask):
        if attention_mask is None or attention_mask.dim() != 2:
            return None
        if torch.is_floating_point(attention_mask):
            keep = attention_mask >= 0
        else:
            keep = attention_mask != 0
        return keep.unsqueeze(-1)

    def forward(self, hidden_states, attention_mask=None):
        gated = self._gate(hidden_states)

        keep = self._keep_mask(attention_mask)
        if keep is not None:
            gated = gated * keep.to(gated.dtype)

        out = gated.transpose(1, 2)
        out = self.conv(out)
        out = out.transpose(1, 2)

        if keep is not None:
            out = out * keep.to(out.dtype)

        return self.out_proj(out)


class IIILayer(Module):
    def __init__(self, config: IIIBERTConfig, layer_id):
        super().__init__()
        self.config = config
        self.layer_id = layer_id

        self.window_size = config.layers_kernel_size[layer_id]
        self.dilation = config.layers_diliation_size[layer_id]

        self.input_layernorm = IIIRMSNorm(config.hidden_dim, eps=config.rms_norm_eps)

        if self.window_size == 0:
            self.is_attention = True
            self.mixer = IIIAttention(config)
        else:
            self.is_attention = False
            self.mixer = GatedCNN(config, layer_id)

        self.post_mixer_layernorm = IIIRMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
        self.mlp = IIIMLP(config)

    def forward(self, hidden_states, attention_mask=None, output_attentions=False):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        if self.is_attention:
            mixer_output, attention_weights = self.mixer(
                hidden_states,
                attention_mask=attention_mask,
                output_attentions=output_attentions,
            )
        else:
            mixer_output = self.mixer(hidden_states, attention_mask=attention_mask)
            attention_weights = None

        hidden_states = residual + mixer_output

        residual = hidden_states
        hidden_states = self.post_mixer_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states, attention_weights


class IIIBERTLayerGroup(Module):
    def __init__(self, config: IIIBERTConfig):
        super().__init__()
        self.layers = ModuleList(
            [IIILayer(config, i) for i in range(config.n_sublayers)]
        )
        self.norm = IIIRMSNorm(config.hidden_dim, eps=config.rms_norm_eps)

    def forward(self, hidden_states, attention_mask=None, output_attentions=False):
        all_attentions = []

        for layer in self.layers:
            hidden_states, attention_weights = layer(
                hidden_states,
                attention_mask=attention_mask,
                output_attentions=output_attentions,
            )
            if attention_weights is not None:
                all_attentions.append(attention_weights)

        return hidden_states, tuple(all_attentions)


class SinusiodalPositionalEncoding(Module):
    def __init__(self, config: IIIBERTConfig):
        super().__init__()

        self.dropout = Dropout(config.dropout)

        pe = torch.zeros(config.max_sequence_length, config.embedding_dim)
        position = torch.arange(config.max_sequence_length, dtype=torch.float32).unsqueeze(1)

        div_term = torch.exp(
            torch.arange(0, config.embedding_dim, 2, dtype=torch.float32)
            * (-math.log(config.pe_theta) / config.embedding_dim)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term[: pe[:, 1::2].shape[1]])

        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x: Tensor):
        seq_len = x.size(1)

        if seq_len > self.pe.size(1):
            raise ValueError("Sequence length exceeds max_sequence_length")

        x = x + self.pe[:, :seq_len, :].to(dtype=x.dtype)
        return self.dropout(x)


class IIIBERTPreTrainedModel(PreTrainedModel):
    config_class = IIIBERTConfig
    base_model_prefix = "iii_bert"
    supports_gradient_checkpointing = False
    _no_split_modules = ["IIILayer"]

    def _init_weights(self, module):
        std = self.config.initializer_range

        if isinstance(module, Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, Conv1d):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, IIIRMSNorm):
            module.weight.data.fill_(1.0)


class IIIBERTModel(IIIBERTPreTrainedModel):
    def __init__(self, config: IIIBERTConfig):
        super().__init__(config)

        self.n_layers = config.n_layers

        self.embed_tokens = Embedding(config.vocab_size, config.embedding_dim)
        self.pe = SinusiodalPositionalEncoding(config)
        self.norm = IIIRMSNorm(config.embedding_dim, eps=config.rms_norm_eps)
        self.linear_proj = Linear(config.embedding_dim, config.hidden_dim, bias=False)

        self.model = IIIBERTLayerGroup(config)

        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def forward(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        token_type_ids: Optional[Tensor] = None,
        inputs_embeds: Optional[Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        output_attentions = (
            output_attentions if output_attentions is not None else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")

        embed = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
        embed = self.pe(embed)
        embed = self.norm(embed)
        embed = self.linear_proj(embed)

        all_hidden_states = (embed,) if output_hidden_states else None
        all_attentions = () if output_attentions else None

        for i in range(self.n_layers):
            embed, attentions = self.model(
                embed,
                attention_mask=attention_mask,
                output_attentions=output_attentions,
            )

            if i == self.n_layers - 1:
                embed = self.model.norm(embed)

            if output_hidden_states:
                all_hidden_states += (embed,)
            if output_attentions:
                all_attentions += attentions

        if not return_dict:
            return tuple(
                v for v in (embed, all_hidden_states, all_attentions) if v is not None
            )

        return BaseModelOutput(
            last_hidden_state=embed,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )


class IIIBERTForMaskedLM(IIIBERTPreTrainedModel):
    def __init__(self, config: IIIBERTConfig):
        super().__init__(config)

        self.iii_bert = IIIBERTModel(config)
        self.linear_out = Linear(config.hidden_dim, config.embedding_dim, bias=False)

        self.post_init()

    def forward(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        token_type_ids: Optional[Tensor] = None,
        inputs_embeds: Optional[Tensor] = None,
        labels: Optional[Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.iii_bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )

        logits = F.linear(
            self.linear_out(outputs.last_hidden_state),
            self.iii_bert.embed_tokens.weight,
        )

        loss = None
        if labels is not None:
            loss = CrossEntropyLoss()(
                logits.view(-1, self.config.vocab_size), labels.view(-1)
            )

        if not return_dict:
            output = (logits,) + tuple(
                v for v in (outputs.hidden_states, outputs.attentions) if v is not None
            )
            return ((loss,) + output) if loss is not None else output

        return MaskedLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class IIIBERTForSequenceClassification(IIIBERTPreTrainedModel):
    def __init__(self, config: IIIBERTConfig):
        super().__init__(config)

        self.num_labels = config.num_labels

        self.iii_bert = IIIBERTModel(config)
        self.dropout = Dropout(config.dropout)
        self.classifier = Linear(config.hidden_dim, config.num_labels)

        self.post_init()

    def forward(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        token_type_ids: Optional[Tensor] = None,
        inputs_embeds: Optional[Tensor] = None,
        labels: Optional[Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.iii_bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )

        hidden = outputs.last_hidden_state

        if attention_mask is not None and attention_mask.dim() == 2:
            keep = (attention_mask != 0).unsqueeze(-1).to(hidden.dtype)
            pooled = (hidden * keep).sum(dim=1) / keep.sum(dim=1).clamp(min=1.0)
        else:
            pooled = hidden.mean(dim=1)

        logits = self.classifier(self.dropout(pooled))

        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif labels.dtype in (torch.long, torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                if self.num_labels == 1:
                    loss = MSELoss()(logits.squeeze(), labels.squeeze())
                else:
                    loss = MSELoss()(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss = CrossEntropyLoss()(
                    logits.view(-1, self.num_labels), labels.view(-1)
                )
            else:
                loss = BCEWithLogitsLoss()(logits, labels)

        if not return_dict:
            output = (logits,) + tuple(
                v for v in (outputs.hidden_states, outputs.attentions) if v is not None
            )
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class IIIBERTForTokenClassification(IIIBERTPreTrainedModel):
    def __init__(self, config: IIIBERTConfig):
        super().__init__(config)

        self.num_labels = config.num_labels

        self.iii_bert = IIIBERTModel(config)
        self.dropout = Dropout(config.dropout)
        self.classifier = Linear(config.hidden_dim, config.num_labels)

        self.post_init()

    def forward(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        token_type_ids: Optional[Tensor] = None,
        inputs_embeds: Optional[Tensor] = None,
        labels: Optional[Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.iii_bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )

        logits = self.classifier(self.dropout(outputs.last_hidden_state))

        loss = None
        if labels is not None:
            loss = CrossEntropyLoss()(
                logits.view(-1, self.num_labels), labels.view(-1)
            )

        if not return_dict:
            output = (logits,) + tuple(
                v for v in (outputs.hidden_states, outputs.attentions) if v is not None
            )
            return ((loss,) + output) if loss is not None else output

        return TokenClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


IIIBERTConfig.register_for_auto_class()
IIIBERTModel.register_for_auto_class("AutoModel")
IIIBERTForMaskedLM.register_for_auto_class("AutoModelForMaskedLM")
IIIBERTForSequenceClassification.register_for_auto_class("AutoModelForSequenceClassification")
IIIBERTForTokenClassification.register_for_auto_class("AutoModelForTokenClassification")