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

import mlx.core as mx
import mlx.nn as nn

from .base import BaseModelArgs


@dataclass
class ModelArgs(BaseModelArgs):
    model_type: str
    add_bias_linear: bool = False
    add_qkv_bias: bool = True
    apply_query_key_layer_scaling: bool = True
    apply_residual_connection_post_layernorm: bool = False
    attention_dropout: float = 0.0
    attention_softmax_in_fp32: bool = True
    bias_dropout_fusion: bool = True
    ffn_hidden_size: int = 13696
    fp32_residual_connection: bool = False
    hidden_dropout: float = 0.0
    hidden_size: int = 4096    
    kv_channels: int = 128
    layernorm_epsilon: float = 1.5625e-07
    multi_query_attention: bool = True
    multi_query_group_num: int = 2
    num_attention_heads: int = 32
    num_hidden_layers: int = 40
    num_layers: int = 40
    rope_ratio: int = 500
    original_rope: bool = True
    padded_vocab_size: int = 151552
    post_layer_norm: bool = True
    rmsnorm: bool = True
    seq_length: int = 131072
    use_cache: bool = True
    torch_dtype: str = "bfloat16"
    tie_word_embeddings: bool = False

    def __post_init__(self):
        pass

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, rope_ratio=1, original_impl=False, dtype=None):
        super().__init__()
        # inv_freq = 1.0 / (10000 ** (mx.arange(0, dim, 2, dtype=dtype) / dim))
        # self.register_buffer("inv_freq", inv_freq)
        # self.inv_freq = mx.array(inv_freq, dtype=dtype)
        self.inv_freq_type = dtype
        self.dim = dim
        self.original_impl = original_impl
        self.rope_ratio = rope_ratio

    def forward_impl(
            self, seq_len: int, n_elem: int, dtype: mx.Dtype, base: int = 10000
    ):
        """Enhanced Transformer with Rotary Position Embedding.
        Derived from:https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
        transformers/rope/__init__.py. MIT License:
        https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
        """
        # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
        base = base * self.rope_ratio
        theta = 1.0 / (base ** (mx.arange(0, n_elem, 2, dtype=mx.float16) / n_elem))

        # Create position indexes `[0, 1, ..., seq_len - 1]`
        seq_idx = mx.arange(seq_len, dtype=mx.float16)

        # Calculate the product of position index and $\theta_i$
        idx_theta = mx.outer(seq_idx, theta).astype(mx.float16)

        cache = mx.stack([mx.cos(idx_theta), mx.sin(idx_theta)], axis=-1)

        # this is to mimic the behaviour of complex32, else we will get different results
        if dtype in (mx.float16, mx.bfloat16, mx.int8):
            cache = cache.astype(mx.bfloat16) if dtype == mx.bfloat16 else cache.astype(mx.float16)
        return cache
    
    def __call__(self, max_seq_len, offset=0):
        return self.forward_impl(
            max_seq_len, self.dim, dtype=self.inv_freq_type,
        )

def apply_rotary_pos_emb(x: mx.array, rope_cache: mx.array) -> mx.array:
    # x: [b, np, sq, hn]
    b, np, sq, hn = x.shape[0], x.shape[1], x.shape[2], x.shape[3]
    rot_dim = rope_cache.shape[-2] * 2
    x, x_pass = x[..., :rot_dim], x[..., rot_dim:]
    # truncate to support variable sizes
    rope_cache = rope_cache[:, :sq]
    xshaped = x.reshape(b, np, sq, rot_dim // 2, 2)
    rope_cache = rope_cache.reshape(-1, 1, sq, xshaped.shape[3], 2)
    x_out2 = mx.stack(
        [
            xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
            xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
        ],
        -1,
    )
    x_out2 = x_out2.flatten(3)
    return mx.concatenate((x_out2, x_pass), axis=-1)

# class RMSNorm(nn.Module):
#     def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):
#         super().__init__()
#         self.weight = nn.empty(normalized_shape, device=device, dtype=dtype)
#         self.eps = eps

#     def __call__(self, hidden_states: mx.array):
#         input_dtype = hidden_states.dtype
#         variance = hidden_states.astype("float32").power(2).mean(-1, keepdims=True)
#         hidden_states = hidden_states * variance.rsqrt()

#         return (self.weight * hidden_states).astype(input_dtype)
    

class CoreAttention(nn.Module):
    def __init__(self, args: ModelArgs, layer_number):
        super().__init__()

        self.apply_query_key_layer_scaling = args.apply_query_key_layer_scaling
        self.attention_softmax_in_fp32 = args.attention_softmax_in_fp32
        if self.apply_query_key_layer_scaling:
            self.attention_softmax_in_fp32 = True
        self.layer_number = max(1, layer_number)

        projection_size = args.kv_channels * args.num_attention_heads

        # Per attention head and per partition values.
        self.hidden_size_per_partition = projection_size
        self.hidden_size_per_attention_head = projection_size // args.num_attention_heads
        self.num_attention_heads_per_partition = args.num_attention_heads

        coeff = None
        self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
        if self.apply_query_key_layer_scaling:
            coeff = self.layer_number
            self.norm_factor *= coeff
        self.coeff = coeff

        self.attention_dropout = nn.Dropout(args.attention_dropout)

    def __call__(self, query_layer, key_layer, value_layer, attention_mask):
        # scale_factor = 1 / math.sqrt(query_layer.shape[-1])
        scale_factor = query_layer.shape[-1] ** -0.5
        # if self.layer_number == 1:
        #     print(f"== |{self.layer_number}| query_layer:{query_layer.shape} key_layer:{key_layer.shape} value_layer:{value_layer.shape}")
        if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
            attention_mask = nn.MultiHeadAttention.create_additive_causal_mask(query_layer.shape[2]).astype(query_layer.dtype)
            context_layer = mx.fast.scaled_dot_product_attention(query_layer, key_layer, value_layer, scale=scale_factor,mask=attention_mask)
        else:
            if attention_mask is not None:
                attention_mask = ~attention_mask
            context_layer = mx.fast.scaled_dot_product_attention(query_layer, key_layer, value_layer, scale=scale_factor, mask=attention_mask)
        context_layer = context_layer.transpose((0,2,1,3)) 
        new_context_layer_shape = context_layer.shape[:-2] + (self.hidden_size_per_partition,)
        context_layer = context_layer.reshape(*new_context_layer_shape)

        return context_layer

class SelfAttention(nn.Module):
    def __init__(self, args: ModelArgs, layer_number):
        super(SelfAttention, self).__init__()
        self.layer_number = max(1, layer_number)

        self.projection_size = args.kv_channels * args.num_attention_heads

        # Per attention head and per partition values.
        self.hidden_size_per_attention_head = self.projection_size // args.num_attention_heads
        self.num_attention_heads_per_partition = args.num_attention_heads
        self.multi_query_attention = args.multi_query_attention
        self.qkv_hidden_size = 3 * self.projection_size
        if self.multi_query_attention:
            self.num_multi_query_groups_per_partition = args.multi_query_group_num
            self.qkv_hidden_size = (
                    self.projection_size + 2 * self.hidden_size_per_attention_head * args.multi_query_group_num
            )
        self.query_key_value = nn.Linear(args.hidden_size, self.qkv_hidden_size,
                                         bias=args.add_bias_linear or args.add_qkv_bias)

        self.core_attention = CoreAttention(args, self.layer_number)

        # Output.
        self.dense = nn.Linear(self.projection_size, args.hidden_size, bias=args.add_bias_linear)

    def __call__(self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True):
        # hidden_states: [b, sq, h]

        # =================================================
        # Pre-allocate memory for key-values for inference.
        # =================================================
        # =====================
        # Query, Key, and Value
        # =====================

        # Attention heads [b, sq, h] --> [b, sq, (np * 3 * hn)]
        mixed_x_layer = self.query_key_value(hidden_states)

        if self.multi_query_attention:
            q_k_v_len = [
                    self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
                    self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
                    self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
            ]
            mixs = mixed_x_layer.split([
                    q_k_v_len[0],
                    q_k_v_len[0]+q_k_v_len[1],
                    q_k_v_len[0]+q_k_v_len[1]+q_k_v_len[2],
                ],
                axis=-1,
            )

            query_layer, key_layer, value_layer = mixs[0], mixs[1], mixs[2]
            query_layer = query_layer.reshape(
                query_layer.shape[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
            )
            key_layer = key_layer.reshape( key_layer.shape[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head))
            value_layer = value_layer.reshape(
                value_layer.shape[:-1]
                + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
            )
        else:
            new_tensor_shape = mixed_x_layer.shape[:-1] + \
                               (self.num_attention_heads_per_partition,
                                3 * self.hidden_size_per_attention_head)
            mixed_x_layer = mixed_x_layer.reshape(*new_tensor_shape)

            # [b, sq, np, 3 * hn] --> 3 [b, sq, np, hn]
            (query_layer, key_layer, value_layer) = mx.split_along_last_dim(mixed_x_layer, 3)

        # [b, sq, np, hn] -> [b, np, sq, hn]
        query_layer, key_layer, value_layer = [k.transpose((0,2,1,3)) for k in [query_layer, key_layer, value_layer]]

        # apply relative positional encoding (rotary embedding)
        if rotary_pos_emb is not None:
            query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
            key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)


        # adjust key and value for inference        
        if use_cache:            
            key_layer, value_layer = kv_cache.update_and_fetch(key_layer, value_layer)                
        else:
            kv_cache = None

        # if self.multi_query_attention:
        #     # key_layer = key_layer.unsqueeze(2)
        #     key_layer = mx.expand_dims(key_layer,2)
        #     key_layer_shape = key_layer.shape
        #     key_layer = mx.broadcast_to(key_layer,[
        #         key_layer_shape[0], key_layer_shape[1], self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, key_layer_shape[3], key_layer_shape[4]]
        #     )
        #     key_layer = key_layer.reshape(
        #         key_layer.shape[:1] + (self.num_attention_heads_per_partition,) + key_layer.shape[3:]
        #     )

        #     # value_layer = value_layer.unsqueeze(2)
        #     value_layer = mx.expand_dims(value_layer,2)
        #     value_layer_shape = value_layer.shape
        #     value_layer = mx.broadcast_to(value_layer,[
        #         value_layer_shape[0], value_layer_shape[1], self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, value_layer_shape[3], value_layer_shape[4]]
        #     )
        #     value_layer = value_layer.reshape(
        #         value_layer.shape[:1] + (self.num_attention_heads_per_partition,) + value_layer.shape[3:]
        #     )

        # ==================================
        # core attention computation
        # ==================================
        context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)

        # =================
        # Output. [sq, b, h]
        # =================

        output = self.dense(context_layer)

        return output

class MLP(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.add_bias = args.add_bias_linear

        # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
        self.dense_h_to_4h = nn.Linear(
            args.hidden_size,
            args.ffn_hidden_size * 2,
            bias=self.add_bias,
        )

        def swiglu(x):
            x = mx.split(x, 2, axis=-1)
            return nn.silu(x[0]) * x[1]

        self.activation_func = swiglu

        # Project back to h.
        self.dense_4h_to_h = nn.Linear(
            args.ffn_hidden_size,
            args.hidden_size,
            bias=self.add_bias,
        )

    def __call__(self, hidden_states):
        # [s, b, 4hp]
        intermediate_parallel = self.dense_h_to_4h(hidden_states)
        intermediate_parallel = self.activation_func(intermediate_parallel)
        # [s, b, h]
        output = self.dense_4h_to_h(intermediate_parallel)
        return output
    

class GLMBlock(nn.Module):
    def __init__(self, args: ModelArgs, layer_number):
        super(GLMBlock, self).__init__()
        self.layer_number = layer_number

        self.apply_residual_connection_post_layernorm = args.apply_residual_connection_post_layernorm

        self.fp32_residual_connection = args.fp32_residual_connection

        LayerNormFunc = nn.RMSNorm if args.rmsnorm else nn.LayerNorm
        # Layernorm on the input data.
        self.input_layernorm = LayerNormFunc(args.hidden_size, eps=args.layernorm_epsilon)

        # Self attention.
        self.self_attention = SelfAttention(args, layer_number)
        self.hidden_dropout = args.hidden_dropout

        self.dropout = nn.Dropout(self.hidden_dropout)

        # Layernorm on the attention output
        self.post_attention_layernorm = LayerNormFunc(args.hidden_size, eps=args.layernorm_epsilon)

        # MLP
        self.mlp = MLP(args)

    def __call__(
            self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,
    ):
        # hidden_states: [s, b, h]

        # Layer norm at the beginning of the transformer layer.
        layernorm_output = self.input_layernorm(hidden_states)
        # Self attention.
        attention_output = self.self_attention(
            layernorm_output,
            attention_mask,
            rotary_pos_emb,
            kv_cache=kv_cache,
            use_cache=use_cache
        )

        # Residual connection.
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = hidden_states

        layernorm_input = self.dropout(attention_output)
        layernorm_input = residual + layernorm_input

        # Layer norm post the self attention.
        layernorm_output = self.post_attention_layernorm(layernorm_input)

        # MLP.
        mlp_output = self.mlp(layernorm_output)

        # Second residual connection.
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = layernorm_input

        output = self.dropout(mlp_output)
        output = residual + output

        return output

class GLMTransformer(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.fp32_residual_connection = args.fp32_residual_connection
        self.post_layer_norm = args.post_layer_norm

        # Number of layers.
        self.num_layers = args.num_layers

        # Transformer layers.
        def build_layer(layer_number):
            return GLMBlock(args, layer_number)

        self.layers = [build_layer(i + 1) for i in range(self.num_layers)]

        if self.post_layer_norm:
            LayerNormFunc = nn.RMSNorm if args.rmsnorm else nn.LayerNorm
            # Final layer norm before output.
            self.final_layernorm = LayerNormFunc(args.hidden_size, eps=args.layernorm_epsilon)

        self.gradient_checkpointing = False

    def _get_layer(self, layer_number):
        return self.layers[layer_number]

    def __call__(
            self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,
            use_cache: Optional[bool] = True,
    ):
        if not kv_caches:
            kv_caches = [None for _ in range(self.num_layers)]

        for index in range(self.num_layers):
            layer = self._get_layer(index)            
            layer_ret = layer(
                hidden_states,
                attention_mask,
                rotary_pos_emb,
                kv_cache=kv_caches[index],
                use_cache=use_cache
            )
            hidden_states = layer_ret

        # Final layer norm.
        if self.post_layer_norm:
            hidden_states = self.final_layernorm(hidden_states)

        return hidden_states

class Embedding(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.hidden_size = args.hidden_size
        # Word embeddings (parallel).
        self.word_embeddings = nn.Embedding(
            args.padded_vocab_size,
            self.hidden_size,
        )
        self.fp32_residual_connection = args.fp32_residual_connection

    def __call__(self, input_ids):
        # Embeddings.
        words_embeddings = self.word_embeddings(input_ids)
        embeddings = words_embeddings
        # If the input flag for fp32 residual connection is set, convert for float.
        if self.fp32_residual_connection:
            embeddings = embeddings.float()
        return embeddings


class ChatGLMModel(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.embedding = Embedding(args)
        self.num_layers = args.num_layers
        self.multi_query_group_num = args.multi_query_group_num

        self.kv_channels = args.kv_channels
        self.use_cache = args.use_cache
        self.use_return_dict = False
        self.output_hidden_states = False

        # Rotary positional embeddings
        self.seq_length = args.seq_length
        rotary_dim = (
            args.hidden_size // args.num_attention_heads if args.kv_channels is None else args.kv_channels
        )

        self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, rope_ratio=args.rope_ratio, original_impl=args.original_rope,dtype=args.torch_dtype)
        self.encoder = GLMTransformer(args)
        self.output_layer = nn.Linear(args.hidden_size, args.padded_vocab_size, bias=False)

        self.new_position_id = None
        self.is_first_forward = True

    def get_input_embeddings(self):
        return self.embedding.word_embeddings

    def set_input_embeddings(self, value):
        self.embedding.word_embeddings = value

    def get_masks(self, input_ids, past_key_values, padding_mask=None):
        batch_size, seq_length = input_ids.shape
        full_attention_mask = mx.ones((batch_size, seq_length, seq_length), dtype=input_ids.dtype)
        full_attention_mask = mx.tril(full_attention_mask)
        past_length = 0
        if past_key_values and past_key_values[0].keys is not None:
            past_length = past_key_values[0].offset
        if past_length:
            full_attention_mask = mx.concatenate((mx.ones((batch_size, seq_length, past_length), dtype=input_ids.dtype),
                                             full_attention_mask), axis=-1)
        if padding_mask is not None:
            full_attention_mask = full_attention_mask * mx.expand_dims(padding_mask,1)
        if not past_length and padding_mask is not None:
            full_attention_mask -= mx.expand_dims(padding_mask,-1) - 1
        full_attention_mask = (full_attention_mask < 0.5)
        full_attention_mask = mx.expand_dims(full_attention_mask,1)
        return full_attention_mask

    def get_position_ids(self, input_ids):
        batch_size, seq_length = input_ids.shape
        position_ids = mx.arange(seq_length, dtype=mx.int32)
        position_ids = mx.broadcast_to(position_ids, (batch_size, seq_length))
        return position_ids

    def __call__(
            self,
            input_ids,
            position_ids: Optional[mx.array] = None,
            attention_mask: Optional[mx.array] = None,
            full_attention_mask: Optional[mx.array] = None,
            past_key_values: Optional[Tuple[Tuple[mx.array, mx.array], ...]] = None,
            inputs_embeds: Optional[mx.array] = None,
            use_cache: Optional[bool] = None,
    ):
        
        # prepare_inputs_for_generation
        if self.new_position_id is None:
            position_ids = self.get_position_ids(input_ids)
        else:
            position_ids = self.new_position_id
        
        new_position_id = position_ids[..., -1:]
        # print(f"== new_position_id:{new_position_id}")
        new_position_id += 1
        # print(f"== new_position_id:{new_position_id}")
        new_position_id = mx.concatenate(
            [position_ids, new_position_id], axis=-1
        )
        # print(f"== new_position_id:{new_position_id}")
        self.new_position_id = new_position_id

        if past_key_values and past_key_values[0].offset > 0: # TODO: check pre_seq
            position_ids = position_ids[..., -1:]
            input_ids = input_ids[:, -1:]

        # print(f"== position_ids:{position_ids} input_ids:{input_ids}")
        batch_size, seq_length = input_ids.shape

        if inputs_embeds is None:
            inputs_embeds = self.embedding(input_ids)
        
        # Rotary positional embeddings
        rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
        if position_ids is not None:
            rotary_pos_emb = rotary_pos_emb[position_ids]
        else:
            rotary_pos_emb = rotary_pos_emb[None, :seq_length]
        # print(f"== rotary_pos_emb:{rotary_pos_emb.shape}")

        # Run encoder.
        hidden_states  = self.encoder(
            inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
            kv_caches=past_key_values, use_cache=use_cache
        )

        return hidden_states
    
    
class Model(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.args = args
        self.model_type = args.model_type
        self.transformer = ChatGLMModel(args)

    def __call__(
        self,
        inputs: mx.array,
        cache=None,
    ):  
        out = self.transformer(inputs, None, None, None, cache, None, True)
        if self.args.tie_word_embeddings:
            out = self.model.embedding.as_linear(out)
        else:
            out = self.model.output_layer(out)
        return out

    def sanitize(self, weights):
        # Remove unused precomputed rotary freqs
        return {
            k: v for k, v in weights.items() if "transformer.rotary_pos_emb.inv_freq" not in k
        }
        # return weights

    @property
    def layers(self):
        return self.model.encoder.layers

    @property
    def head_dim(self):
        return self.args.hidden_size // self.args.num_attention_heads

    @property
    def n_kv_heads(self):
        return self.args.multi_query_group_num
    
    @property
    def model(self):
        return self.transformer