import math
import os
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import CrossEntropyLoss, BCEWithLogitsLoss, MSELoss

from transformers.modeling_outputs import (
    BaseModelOutput,
    CausalLMOutput,
    SequenceClassifierOutput
)

from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging

from .rita_configuration import RITAConfig
import torch.nn.functional as F
logger = logging.get_logger(__name__)

@torch.jit.script
def RITA_gelu(hidden_states):
    return hidden_states * 0.5 * (1.0 + torch.tanh(0.79788456 * hidden_states * (1 + 0.044715 * hidden_states * hidden_states)))

class RITAGELU(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, hidden_states):
        return RITA_gelu(hidden_states)

def rotate_half(x):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=x1.ndim - 1)

class RotaryEmbedding(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.d_model % config.num_heads == 0
        
        self.d_model = config.d_model
        self.num_heads = config.num_heads
        self.max_seq_len = config.max_seq_len
        
        head_dim = self.d_model // self.num_heads
        inv_freq = 1.0 / (10000 ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer('inv_freq', inv_freq)
        self.seq_len_cached = None
        self.cos_cached = None
        self.sin_cached = None
    
    def forward(self, x: torch.FloatTensor, seq_dim=1) -> torch.FloatTensor:
        seq_len = x.shape[seq_dim]
        if seq_len != self.seq_len_cached:
            self.seq_len_cached = seq_len
            t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq)
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            self.cos_cached = emb.cos()[None, None, :, :]
            self.sin_cached = emb.sin()[None, None, :, :]
        return self.cos_cached, self.sin_cached
    
    def apply_rotary_pos_emb(self, q, k, cos, sin):
        return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)

    
class SelfAttention(nn.Module):
    """Implementation of MultiHeadAttention following `Karpathy's MinGPT <https://github.com/karpathy/minGPT>`_.
    modified to use rotary embeddings.
    
    Parameters
    ----------
    d_model: int,
         total dimension of the model.
    num_heads: int,
        number of parallel attention heads.
    num_layers: int,
        number of layers in the model, used for the Megatron-like init.
    rotaty_embedding: Optional[Block], default None,
        a RotaryEmbedding Block to add positionnal information in Queries and Keys
    dropout: float, default 0.1,
        amount of dropout on the attention weights.
    sigma: float, default 0.02,
        standard deviation used for the init.
    trainable: bool, default True,
        if False, the Module parameters will be hidden from the optimizer.
    """

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        num_layers: int,
        rotary_embedding= None,
        dropout: float = 0.1,
        sigma=0.02,
        use_cache: bool = False,
        bias=True,
    ):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = self.d_model // self.num_heads
        self.num_layers = num_layers
        self.dropout = dropout
        self.sigma = sigma
        self.bias = bias

        # key, query, value projections for all heads
        self.key = nn.Linear(d_model, d_model, bias=bias)
        self.query = nn.Linear(d_model, d_model, bias=bias)
        self.value = nn.Linear(d_model, d_model, bias=bias)
        # regularization
        self.attn_drop = nn.Dropout(dropout)
        self.resid_drop = nn.Dropout(dropout)
        # output projection
        self.proj = nn.Linear(d_model, d_model, bias=bias)

        self.rotary_embedding = rotary_embedding
        self.layer_id = None  # will be set by the Transformer itself
        self.use_cache = use_cache
        self.qkv = None
        self.bias = bias

    def forward(
        self,
        x,
        causal_mask: Optional[torch.BoolTensor] = None,
        attention_mask: Optional[torch.BoolTensor] = None,
    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:

        N, L, D = x.size()  # Batch_size, Context_size, d_model

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        k = (
            self.key(x).view(N, L, self.num_heads, D // self.num_heads).transpose(1, 2)
        )  # (N, nh, L, hs)
        q = (
            self.query(x).view(N, L, self.num_heads, D // self.num_heads).transpose(1, 2)
        )  # (N, nh, L, hs)
        v = (
            self.value(x).view(N, L, self.num_heads, D // self.num_heads).transpose(1, 2)
        )  # (N, nh, L, hs)
        
        if self.rotary_embedding is not None:
            cos, sin = self.rotary_embedding(x)
            q, k = self.rotary_embedding.apply_rotary_pos_emb(q, k, cos, sin)

        # causal self-attention; Self-attend: (N, nh, L, hs) x (N, nh, hs, L) -> (N, nh, L, L)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        
        if causal_mask is not None:
            att[:,:,-L:, -L: ].masked_fill_(causal_mask.view(1, 1, L, L), float("-inf"))
            
        att = (
            att.transpose(0, 2)
            .masked_fill(attention_mask.view(1, 1, N, L)==0, float("-inf"))
            .transpose(0, 2)
            if attention_mask is not None
            else att
        )
        
        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)
        y = att @ v  # (N, nh, L, L) x (N, nh, L, hs) -> (N, nh, L, hs)
        y = (
            y.transpose(1, 2).contiguous().view(N, L, D)
        )  # re-assemble all head outputs side by side

        # output projection
        y = self.resid_drop(self.proj(y))
        return y

class DecoderLayer(nn.Module):
    """Transformer block containing the self-attention module and the feedfoward module."""

    def __init__(
        self, config
    ):
        super().__init__()
        self.self_attention = SelfAttention(config.d_model, config.num_heads, config.dropout, rotary_embedding=RotaryEmbedding(config))
        self.attn_norm = nn.LayerNorm(config.d_model)
        self.attn_dropout = nn.Dropout(config.dropout)

        self.mlp = nn.Sequential(
            nn.Linear(config.d_model, config.d_feedforward, bias=True),
            RITAGELU(),
            nn.Linear(config.d_feedforward, config.d_model, bias=True),
        )
        self.mlp_norm = nn.LayerNorm(config.d_model)
        self.mlp_dropout = nn.Dropout(config.dropout)
        
    def forward(
        self,
        x: torch.FloatTensor,
        causal_mask: torch.BoolTensor,
        attention_mask: Optional[torch.BoolTensor] = None,
    ) -> torch.FloatTensor:
        y = self.attn_norm(x)
        y = self.self_attention(y, causal_mask=causal_mask, attention_mask=attention_mask)
        x = x + self.attn_dropout(y)

        y = self.mlp_norm(x)
        y = self.mlp(y)
        x = x + self.mlp_dropout(y)
        return x

class RITAModel(PreTrainedModel):
    config_class = RITAConfig
    base_model_prefix = "transformer"
    is_parallelizable = False
    
    def __init__(
        self,
        config
    ):
        super().__init__(config)
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_layers)])
        self.final_norm = nn.LayerNorm(config.d_model)

    def forward(
        self,
        input_ids=None,
        past_key_values=None,  # NOT USED
        attention_mask=None,
        causal_mask=None,
        token_type_ids=None, # NOT USED
        position_ids=None, # NOT USED
        head_mask=None, # NOT USED
        inputs_embeds=None,
        encoder_hidden_states=None,  # NOT USED
        encoder_causal_mask=None, # NOT USED
        labels=None,
        use_cache=None, # NOT USED
        output_attentions=None, # NOT USED
        output_hidden_states=None, # NOT USED
        return_dict=None # NOT USED
        ) -> torch.FloatTensor:
        if inputs_embeds == None:
            x = self.embedding(input_ids)  # N x L x D
        else:
            x = inputs_embeds
        if causal_mask == None:
            causal_mask = (torch.triu(torch.ones(input_ids.size(1), input_ids.size(1))) == 0).transpose(0, 1).contiguous().to(input_ids.device)
        for layer in self.layers:
            x = layer(x, causal_mask=causal_mask, attention_mask=attention_mask)
        x = self.final_norm(x)  # N x L x D

        return BaseModelOutput(
            hidden_states=x,
        )

    #Some common HF functions.
    def get_input_embeddings(self):
        return self.embedding

    def set_input_embeddings(self, new_embeddings):
        self.embedding = new_embeddings

    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


class RITAModelForCausalLM(PreTrainedModel):
    config_class = RITAConfig
    base_model_prefix = "transformer"
    is_parallelizable = False

    def __init__(
        self,
        config
    ):
        super().__init__(config)
        self.transformer = RITAModel(config)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

    def forward(
        self,
        input_ids=None,
        past_key_values=None,  # NOT USED
        attention_mask=None,
        causal_mask=None,
        token_type_ids=None, # NOT USED
        position_ids=None, # NOT USED
        head_mask=None, # NOT USED
        inputs_embeds=None,
        encoder_hidden_states=None,  # NOT USED
        encoder_causal_mask=None, # NOT USED
        labels=None,
        use_cache=None, # NOT USED
        output_attentions=None, # NOT USED
        output_hidden_states=None, # NOT USED
        return_dict=None # NOT USED
        ) -> torch.FloatTensor:

        transformer_outputs = self.transformer(
            input_ids,
            past_key_values=past_key_values,
            causal_mask=causal_mask,
            attention_mask = attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        
        logits = self.lm_head(transformer_outputs.hidden_states)
        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

        return CausalLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=transformer_outputs.hidden_states,
        )

    #Some common HF functions.
    def get_input_embeddings(self):
        return self.transformer.embedding

    def set_input_embeddings(self, new_embeddings):
        self.transformer.embedding = new_embeddings

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, lm_head):
        self.lm_head = lm_head

    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


class RITAModelForSequenceClassification(PreTrainedModel):
    config_class = RITAConfig
    base_model_prefix = "transformer"
    is_parallelizable = False

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.transformer = RITAModel(config)
        self.score = nn.Linear(config.d_model, self.num_labels, bias=False)

    def forward(
        self,
        input_ids=None,
        past_key_values=None,
        attention_mask=None,
        causal_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        transformer_outputs = self.transformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            causal_mask=causal_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0]
        logits = self.score(hidden_states)

        if input_ids is not None:
            batch_size, sequence_length = input_ids.shape[:2]
        else:
            batch_size, sequence_length = inputs_embeds.shape[:2]

        assert (
            self.config.pad_token_id is not None or batch_size == 1
        ), "Cannot handle batch sizes > 1 if no padding token is defined."
        if self.config.pad_token_id is None:
            sequence_lengths = -1
        else:
            if input_ids is not None:
                sequence_lengths = torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1
            else:
                sequence_lengths = -1
                logger.warning(
                    f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
                    f"unexpected if using padding tokens in conjunction with `inputs_embeds.`"
                )

        pooled_logits = logits[torch.arange(batch_size, device=self.device), sequence_lengths]

        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 self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(pooled_logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(pooled_logits, labels)
        if not return_dict:
            output = (pooled_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=pooled_logits,
        )
        
    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)