from typing import Optional, Tuple

import jax
from flax import linen as nn
from flax.core import FrozenDict, unfreeze, freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax import numpy as jnp
from transformers import FlaxPreTrainedModel
from transformers.modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutput
from transformers.modeling_flax_utils import ACT2FN

from .configuration_retnet import RetNetConfig


def rotate_every_two(tensor):
    rotate_half_tensor = jnp.stack(
        (-tensor[:, :, :, 1::2], tensor[:, :, :, ::2]), axis=-1
    )
    rotate_half_tensor = rotate_half_tensor.reshape(
        rotate_half_tensor.shape[:-2] + (-1,)
    )
    return rotate_half_tensor


def theta_shift(x, sin, cos):
    return (x * cos) + (rotate_every_two(x) * sin)


class FlaxRetNetRelPos(nn.Module):
    config: RetNetConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self) -> None:
        angle = 1.0 / (
            10000
            ** jnp.linspace(
                0, 1, self.config.hidden_size // self.config.num_rettention_heads // 2
            )
        )
        self.angle = angle.repeat(2).flatten()
        self.decay = jnp.log(
            1
            - 2
            ** (-5 - jnp.arange(self.config.num_rettention_heads, dtype=jnp.float32))
        )
        self.recurrent_chunk_size = self.config.recurrent_chunk_size

    def __call__(
        self,
        slen: int,
        activate_recurrent: bool = False,
        chunkwise_recurrent: bool = False,
    ):
        if activate_recurrent:
            sin = jnp.sin(self.angle * (slen - 1))
            cos = jnp.cos(self.angle * (slen - 1))
            retention_rel_pos = ((sin, cos), jnp.exp(self.decay))
        elif chunkwise_recurrent:
            index = jnp.arange(slen)
            sin = jnp.sin(index[:, None] * self.angle[None, :])
            cos = jnp.cos(index[:, None] * self.angle[None, :])

            block_index = jnp.arange(self.recurrent_chunk_size)
            mask = jnp.tril(
                jnp.ones((self.recurrent_chunk_size, self.recurrent_chunk_size))
            )
            mask = jnp.where(
                ~mask.astype(jnp.bool_),
                float("inf"),
                block_index[:, None] - block_index[None, :],
            )
            mask = jnp.exp(mask * self.decay[:, None, None])
            mask = jnp.nan_to_num(mask)
            scale = jnp.sqrt(mask.sum(axis=-1, keepdims=True))
            mask = mask / scale

            cross_decay = jnp.exp(self.decay * self.recurrent_chunk_size)
            inner_decay = jnp.exp(self.decay[:, None] * (block_index + 1))
            cross_decay = cross_decay[:, None, None]
            inner_decay = inner_decay[:, :, None] / (scale / scale[:, -1, None])

            retention_rel_pos = ((sin, cos), (mask, cross_decay, inner_decay))
        else:
            index = jnp.arange(slen)
            sin = jnp.sin(index[:, None] * self.angle[None, :])
            cos = jnp.cos(index[:, None] * self.angle[None, :])
            mask = jnp.tril(jnp.ones((slen, slen)))
            mask = jnp.where(
                ~mask.astype(jnp.bool_), float("inf"), index[:, None] - index[None, :]
            )
            mask = jnp.exp(mask * self.decay[:, None, None])
            mask = jnp.nan_to_num(mask)
            mask = mask / jnp.sqrt(mask.sum(axis=-1, keepdims=True))
            retention_rel_pos = ((sin, cos), mask)

        return retention_rel_pos


class FlaxRetNetFeedForward(nn.Module):
    config: RetNetConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self) -> None:
        self.fc1 = nn.Dense(
            self.config.intermediate_size,
            kernel_init=nn.initializers.xavier_normal(),
            dtype=self.dtype,
        )
        self.fc2 = nn.Dense(
            self.config.hidden_size,
            kernel_init=nn.initializers.xavier_normal(),
            dtype=self.dtype,
        )
        self.activation_fn = ACT2FN[self.config.hidden_act]
        self.activation_dropout = nn.Dropout(rate=self.config.dropout)
        self.dropout = nn.Dropout(rate=self.config.dropout)

    def __call__(
        self,
        hidden_states: jnp.ndarray,
        deterministic: bool = True,
    ) -> jnp.ndarray:
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.activation_fn(hidden_states)
        hidden_states = self.activation_dropout(
            hidden_states, deterministic=deterministic
        )
        hidden_states = self.fc2(hidden_states)
        hidden_states = self.dropout(hidden_states, deterministic=deterministic)

        return hidden_states


class FlaxRetNetRetention(nn.Module):
    config: RetNetConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self) -> None:
        self.factor = 2
        self.embed_dim = self.config.hidden_size
        self.num_heads = self.config.num_rettention_heads
        self.head_dim = self.embed_dim * self.factor // self.num_heads
        self.key_dim = self.embed_dim // self.num_heads
        self.scaling = self.key_dim**-0.5

        self.q_proj = nn.Dense(
            self.embed_dim,
            use_bias=True,
            kernel_init=jax.nn.initializers.xavier_normal(),
            dtype=self.dtype,
        )
        self.k_proj = nn.Dense(
            self.embed_dim,
            use_bias=True,
            kernel_init=jax.nn.initializers.xavier_normal(),
            dtype=self.dtype,
        )
        self.v_proj = nn.Dense(
            self.embed_dim * self.factor,
            use_bias=True,
            kernel_init=jax.nn.initializers.xavier_normal(),
            dtype=self.dtype,
        )
        self.g_proj = nn.Dense(
            self.embed_dim * self.factor,
            use_bias=True,
            kernel_init=nn.initializers.xavier_normal(),
            dtype=self.dtype,
        )

        self.out_proj = nn.Dense(
            self.embed_dim,
            use_bias=True,
            kernel_init=jax.nn.initializers.xavier_normal(),
            dtype=self.dtype,
        )

        self.group_norm = nn.LayerNorm(epsilon=1e-6, dtype=self.dtype)

    def parallel_forward(self, qr, kr, v, mask):
        bsz, tgt_len, embed_dim = v.shape

        vr = v.reshape(bsz, tgt_len, self.num_heads, self.head_dim).transpose(
            (0, 2, 1, 3)
        )

        qk_mat = qr @ kr.transpose((0, 1, 3, 2))
        qk_mat = qk_mat * mask
        qk_mat /= jnp.abs(
            jax.lax.stop_gradient(qk_mat).sum(axis=-1, keepdims=True)
        ).clip(min=1)
        output = jnp.matmul(qk_mat, vr)
        output = output.transpose((0, 2, 1, 3))

        return output

    def chunk_recurrent_forward(self, qr, kr, v, inner_mask):
        mask, cross_decay, inner_decay = inner_mask
        bsz, tgt_len, embed_dim = v.shape
        chunk_len = mask.shape[1]
        num_chunks = tgt_len // chunk_len

        assert tgt_len % chunk_len == 0

        qr = qr.reshape(
            bsz, self.num_heads, num_chunks, chunk_len, self.key_dim
        ).transpose((0, 2, 1, 3, 4))
        kr = kr.reshape(
            bsz, self.num_heads, num_chunks, chunk_len, self.key_dim
        ).transpose((0, 2, 1, 3, 4))
        v = v.reshape(
            bsz, num_chunks, chunk_len, self.num_heads, self.head_dim
        ).transpose((0, 1, 3, 2, 4))

        kr_t = kr.transpose((0, 1, 2, 4, 3))

        qk_mat = qr @ kr_t
        qk_mat = qk_mat
        inner_scale = jnp.abs(
            jax.lax.stop_gradient(qk_mat).sum(axis=-1, keepdims=True)
        ).clip(min=1)
        qk_mat = qk_mat / inner_scale
        inner_output = jnp.matmul(qk_mat, v)

        kv = kr_t @ v
        kv = kv.reshape(bsz, num_chunks, self.num_heads, self.key_dim, self.head_dim)

        kv_recurrent = []
        cross_scale = []
        kv_state = jnp.zeros((bsz, self.num_heads, self.key_dim, self.head_dim))
        kv_scale = jnp.ones((bsz, self.num_heads, 1, 1))

        for i in range(num_chunks):
            kv_recurrent.append(kv_state / kv_scale)
            cross_scale.append(kv_scale)

            kv_state = kv_state * cross_decay + kv[:, i]
            kv_scale = (
                jnp.abs(jax.lax.stop_gradient(kv_state).sum(axis=-2, keepdims=True))
                .max(axis=-1, keepdims=True)
                .clip(min=1)
            )

        kv_recurrent = jnp.stack(kv_recurrent, axis=1)
        cross_scale = jnp.stack(cross_scale, axis=1)

        all_scale = jnp.maximum(inner_scale, cross_scale)
        align_inner_scale = all_scale / inner_scale
        align_cross_scale = all_scale / cross_scale

        cross_output = (qr * inner_decay) @ kv_recurrent
        output = inner_output / align_inner_scale + cross_output / align_cross_scale

        output = output.transpose((0, 2, 1, 3, 4))
        return output

    def __call__(
        self,
        hidden_states: jnp.ndarray,
        rel_pos: Optional[jnp.ndarray] = None,
        chunkwise_recurrent: bool = True,
        incremental_state=None,
    ) -> jnp.ndarray:
        bsz, tgt_len, _ = hidden_states.shape
        (sin, cos), inner_mask = rel_pos

        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        v = self.v_proj(hidden_states)
        g = self.g_proj(hidden_states)

        k *= self.scaling
        q = q.reshape(bsz, tgt_len, self.num_heads, self.key_dim).transpose(
            (0, 2, 1, 3)
        )
        k = k.reshape(bsz, tgt_len, self.num_heads, self.key_dim).transpose(
            (0, 2, 1, 3)
        )

        qr = theta_shift(q, sin, cos)
        kr = theta_shift(k, sin, cos)

        if incremental_state is not None:
            raise NotImplementedError
        elif self.config.attention_type == "chunkwise_recurrent":
            output = self.chunk_recurrent_forward(qr, kr, v, inner_mask=inner_mask)
        else:
            output = self.parallel_forward(qr, kr, v, inner_mask)

        output = self.group_norm(output)
        output = output.reshape(bsz, tgt_len, -1)

        output = nn.swish(g) * output
        output = self.out_proj(output)

        return output


class FlaxRetNetLayer(nn.Module):
    config: RetNetConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self) -> None:
        self.retention = FlaxRetNetRetention(self.config, dtype=self.dtype)
        self.retention_layer_norm = nn.LayerNorm(
            epsilon=self.config.layer_norm_eps, dtype=self.dtype
        )

        self.ffn = FlaxRetNetFeedForward(self.config, dtype=self.dtype)
        self.final_layer_norm = nn.LayerNorm(
            epsilon=self.config.layer_norm_eps, dtype=self.dtype
        )

        self.dropout_module = nn.Dropout(rate=self.config.dropout)

    def __call__(
        self,
        hidden_states: jnp.ndarray,
        retention_rel_pos: Optional[tuple] = None,
        deterministic: bool = True,
    ) -> jnp.ndarray:
        residual = hidden_states
        hidden_states = self.retention_layer_norm(hidden_states)
        hidden_states = self.retention(hidden_states, rel_pos=retention_rel_pos)
        hidden_states = self.dropout_module(hidden_states, deterministic=deterministic)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states = self.ffn(hidden_states, deterministic=deterministic)
        hidden_states = residual + hidden_states

        return hidden_states


class FlaxRetNetLayerCollection(nn.Module):
    config: RetNetConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self) -> None:
        self.layers = [
            FlaxRetNetLayer(self.config, dtype=self.dtype)
            for _ in range(self.config.num_hidden_layers)
        ]

    def __call__(
        self,
        hidden_states: jnp.ndarray,
        retention_rel_pos: tuple = None,
        deterministic: bool = True,
        output_retentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ) -> jnp.ndarray:
        all_hidden_states = () if output_hidden_states else None
        all_retentions = () if output_retentions else None

        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            layer_outputs = layer(
                hidden_states,
                retention_rel_pos=retention_rel_pos,
                deterministic=deterministic,
            )
            hidden_states = layer_outputs

        outputs = (hidden_states, all_hidden_states, all_retentions)
        return outputs


class FlaxRetNetPretrainedModel(FlaxPreTrainedModel):
    config_class = RetNetConfig
    base_model_prefix = "transformer"
    main_input_name = "input_ids"
    module_class: nn.Module = None

    def __init__(
        self,
        config: RetNetConfig,
        input_shape: Tuple = (1, 1),
        seed: int = 0,
        dtype: jnp.dtype = jnp.float32,
        _do_init: bool = True,
        **kwargs
    ):
        module = self.module_class(config, dtype=dtype, **kwargs)
        super().__init__(
            config,
            module,
            input_shape=input_shape,
            seed=seed,
            dtype=dtype,
            _do_init=_do_init,
        )

    def init_weights(
        self,
        rng: jax.random.PRNGKey,
        input_shape: Tuple,
        params: FrozenDict = None,
    ) -> FrozenDict:
        input_ids = jnp.zeros(input_shape, dtype="i4")
        attention_mask = jnp.ones_like(input_ids)
        params_rng, dropout_rng = jax.random.split(rng)
        rngs = {"params": params_rng, "dropout": dropout_rng}

        module_init_outputs = self.module.init(
            rngs, input_ids, attention_mask, return_dict=False
        )

        random_params = module_init_outputs["params"]

        if params is not None:
            random_params = flatten_dict(unfreeze(random_params))
            params = flatten_dict(unfreeze(params))
            for missing_key in self._missing_keys:
                params[missing_key] = random_params[missing_key]
            self._missing_keys = []
            return freeze(unflatten_dict(params))
        else:
            return random_params

    def __call__(
        self,
        input_ids: jnp.ndarray,
        attention_mask: Optional[jnp.ndarray] = None,
        params: dict = None,
        dropout_rng: jnp.ndarray = None,
        train: bool = False,
        output_retentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ):
        output_retentions = (
            output_retentions
            if output_retentions is not None
            else self.config.output_retentions
        )
        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.return_dict
        )

        batch_size, sequence_length = input_ids.shape

        if attention_mask is None:
            attention_mask = jnp.ones((batch_size, sequence_length))

        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        inputs = {"params": params or self.params}

        outputs = self.module.apply(
            inputs,
            jnp.array(input_ids, dtype="i4"),
            jnp.array(attention_mask, dtype="i4"),
            not train,
            output_retentions,
            output_hidden_states,
            return_dict,
            rngs=rngs,
        )

        return outputs


class FlaxRetNetModule(nn.Module):
    config: RetNetConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self) -> None:
        self.embed_tokens = nn.Embed(
            self.config.vocab_size,
            self.config.hidden_size,
            embedding_init=jax.nn.initializers.xavier_normal(),
            dtype=self.dtype,
        )
        self.retnet_rel_pos = FlaxRetNetRelPos(self.config, dtype=self.dtype)

        self.layers = FlaxRetNetLayerCollection(self.config, dtype=self.dtype)

    def __call__(
        self,
        input_ids: jnp.ndarray,
        attention_mask: Optional[jnp.ndarray] = None,
        deterministic: bool = True,
        output_retentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ):
        input_embeds = self.embed_tokens(input_ids)

        batch_size, sequence_length = input_embeds.shape[:2]
        retention_rel_pos = self.retnet_rel_pos(
            sequence_length,
            activate_recurrent=False,
            chunkwise_recurrent=self.config.attention_type == "chunkwise_recurrent",
        )

        outputs = self.layers(
            input_embeds,
            retention_rel_pos=retention_rel_pos,
            deterministic=deterministic,
            output_retentions=output_retentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        if not return_dict:
            return tuple(v for v in outputs if v is not None)

        return FlaxBaseModelOutput(
            last_hidden_state=outputs[0],
            hidden_states=outputs[1],
            attentions=outputs[-1],
        )


class FlaxRetNetModel(FlaxRetNetPretrainedModel):
    module_class = FlaxRetNetModule


class FlaxRetNetForCausalLMModule(nn.Module):
    config: RetNetConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self) -> None:
        self.transformer = FlaxRetNetModule(self.config, dtype=self.dtype)

        self.lm_head = nn.Dense(
            self.config.vocab_size,
            use_bias=False,
            kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
            dtype=self.dtype,
        )

    def __call__(
        self,
        input_ids: jnp.ndarray,
        attention_mask: Optional[jnp.ndarray] = None,
        deterministic: bool = True,
        output_retentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ):
        outputs = self.transformer(
            input_ids,
            attention_mask=attention_mask,
            deterministic=deterministic,
            output_retentions=output_retentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = outputs[0]

        lm_logits = self.lm_head(hidden_states)

        if not return_dict:
            return (lm_logits,) + outputs[1:]

        return FlaxCausalLMOutput(
            logits=lm_logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class FlaxRetNetForCausalLM(FlaxRetNetPretrainedModel):
    module_class = FlaxRetNetForCausalLMModule