import math
from typing import Optional

import torch
from einops import rearrange
from torch import nn
from torch.nn import functional as F

try:
    from xformers.ops import memory_efficient_attention
except ImportError as e:
    memory_efficient_attention = None


class AlibiPostionEmbedding(nn.Module):
    def __init__(self, nheads, maxpos):
        super().__init__()

        context_position = torch.arange(maxpos)[:, None]
        memory_position = torch.arange(maxpos)[None, :]
        relative_position = memory_position - context_position
        relative_position = (
            torch.abs(relative_position).unsqueeze(0).expand(nheads, -1, -1)
        )
        self.slopes = torch.Tensor(self.get_slopes(nheads)) * -1
        alibi = self.slopes.unsqueeze(1).unsqueeze(1) * relative_position
        alibi = alibi.view(nheads, maxpos, maxpos)

        self.register_buffer("alibi", alibi)

    @staticmethod
    def get_slopes_power_of_2(n):
        start = 2 ** (-(2 ** -(math.log2(n) - 3)))
        ratio = start
        return [start * ratio**i for i in range(n)]

    def get_slopes(self, n):
        if math.log2(n).is_integer():
            return self.get_slopes_power_of_2(n)

        closest_power_of_2 = 2 ** math.floor(math.log2(n))
        return (
            self.get_slopes_power_of_2(closest_power_of_2)
            + self.get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2]
        )

    def __call__(self, x):
        # N, T, C
        return self.alibi[:, : x.size(1), : x.size(1)].to(x.device)


class KVCache(nn.Module):
    def __init__(
        self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16
    ):
        super().__init__()
        cache_shape = (max_batch_size, max_seq_length, n_heads * head_dim)
        self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype))
        self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype))

    def update(self, input_pos, k_val, v_val):
        assert input_pos is not None, "input_pos should not be None"

        k_out = self.k_cache
        v_out = self.v_cache
        k_out[:, input_pos] = k_val
        v_out[:, input_pos] = v_val

        return k_out, v_out


class MultiheadAttention(nn.Module):
    def __init__(self, d_model, nhead, dropout=0.1):
        super().__init__()
        assert d_model % nhead == 0
        self.nhead = nhead
        self.d_model = d_model
        self.head_dim = d_model // nhead

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        self.kv_cache = None

    def forward(
        self,
        q,
        k,
        v,
        attn_mask=None,
        key_padding_mask=None,
        attn_bias=None,
        return_weights=False,
        input_pos=None,
    ):
        # (B, T, C)
        batch_size = q.size(0)
        q_length = q.size(1)

        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)

        if self.kv_cache is not None:
            k, v = self.kv_cache.update(input_pos, k, v)

        k_length = k.size(1)

        if attn_bias is not None:
            assert attn_bias.size() == (
                self.nhead,
                q_length,
                k_length,
            ), f"Should be {(self.nhead, q_length, k_length)}. Got {attn_bias.size()}"

            attn_bias = attn_bias.unsqueeze(0).expand(batch_size, -1, -1, -1)

        if attn_mask is not None:
            assert attn_mask.size() == (
                q_length,
                k_length,
            ), f"Should be {(q_length, k_length)}. Got {attn_mask.size()}"
            assert attn_mask.dtype == torch.bool
            attn_mask = attn_mask.unsqueeze(0).expand(batch_size * self.nhead, -1, -1)

        if key_padding_mask is not None:
            assert key_padding_mask.size() == (
                batch_size,
                k_length,
            ), f"Should be {(batch_size, k_length)}. Got {key_padding_mask.size()}"
            assert key_padding_mask.dtype == torch.bool
            key_padding_mask = (
                key_padding_mask.unsqueeze(1)
                .unsqueeze(1)
                .expand(-1, self.nhead, -1, -1)
            )
            key_padding_mask = key_padding_mask.reshape(
                batch_size * self.nhead, 1, k_length
            )
            if attn_mask is None:
                attn_mask = key_padding_mask.expand(-1, q.size(1), -1)
            else:
                attn_mask = attn_mask.logical_or(key_padding_mask)

        if (
            return_weights is False
            and memory_efficient_attention is not None
            and q.device.type == "cuda"
        ):
            # (-> b, t,. n, d)
            q = rearrange(q, "b t (n d) -> b t n d", n=self.nhead)
            k = rearrange(k, "b t (n d) -> b t n d", n=self.nhead)
            v = rearrange(v, "b t (n d) -> b t n d", n=self.nhead)

            if attn_mask is not None:
                attn_mask = rearrange(attn_mask, "(b n) q k -> b n q k", n=self.nhead)

                if attn_bias is None:
                    attn_bias = torch.zeros_like(
                        attn_mask, dtype=q.dtype, device=q.device
                    )
                attn_bias = attn_bias.masked_fill(attn_mask, float("-inf"))

            if attn_bias is not None:
                attn_bias = attn_bias.to(q.dtype)

            attn_output = memory_efficient_attention(
                q,
                k,
                v,
                attn_bias=attn_bias,
                scale=self.head_dim**-0.5,
                p=self.dropout.p,
            )
            attn_output = rearrange(attn_output, "b t n d -> b t (n d)", n=self.nhead)

            returned_weights = None
        else:
            q = rearrange(q, "b t (n d) -> (b n) t d", n=self.nhead)
            k = rearrange(k, "b t (n d) -> (b n) t d", n=self.nhead)
            v = rearrange(v, "b t (n d) -> (b n) t d", n=self.nhead)

            attn_weights = torch.bmm(q, k.mT) * (self.head_dim**-0.5)
            assert attn_weights.size() == (
                batch_size * self.nhead,
                q.size(1),
                k.size(1),
            )

            if attn_bias is not None:
                attn_bias = rearrange(attn_bias, "b n q k -> (b n) q k")
                attn_weights = attn_weights + attn_bias

            if attn_mask is not None:
                attn_weights = attn_weights.masked_fill(attn_mask, float("-inf"))

            attn_weights = F.softmax(attn_weights, dim=-1, dtype=attn_weights.dtype)
            returned_weights = attn_weights.view(
                batch_size, self.nhead, q.size(1), k.size(1)
            )

            attn_probs = self.dropout(attn_weights)
            attn_output = torch.bmm(attn_probs, v)
            attn_output = rearrange(attn_output, "(b n) t d -> b t (n d)", n=self.nhead)

        attn_output = self.out_proj(attn_output)
        return attn_output, returned_weights


class GluMLP(nn.Module):
    def __init__(self, hidden_size=1024, intermediate_size=None, activation=nn.SiLU):
        super().__init__()

        if intermediate_size is None:
            intermediate_size = hidden_size * (11 / 3)
            intermediate_size = round(intermediate_size / 8) * 8

        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size

        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = activation()

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        RMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()

        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

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

        return self.weight * hidden_states.to(input_dtype)


class CrossAttentionLayer(nn.Module):
    def __init__(self, hidden_size=1024, intermediate_size=None, dropout=0.1):
        super().__init__()

        self.attn = MultiheadAttention(hidden_size, 1, dropout=dropout)
        self.mlp = GluMLP(hidden_size=hidden_size, intermediate_size=intermediate_size)
        self.input_layernorm_q = RMSNorm(hidden_size, eps=1e-6)
        self.input_layernorm_kv = RMSNorm(hidden_size, eps=1e-6)
        self.post_attention_layernorm = RMSNorm(hidden_size, eps=1e-6)

    def forward(
        self,
        tgt,
        memory,
        memory_key_padding_mask=None,
        input_pos=None,
    ):
        residual = tgt
        tgt, memory = self.input_layernorm_q(tgt), self.input_layernorm_kv(memory)
        x, attn_weights = self.attn(
            tgt,
            memory,
            memory,
            key_padding_mask=memory_key_padding_mask,
            return_weights=True,
            input_pos=input_pos,
        )
        residual = x + residual

        x = self.post_attention_layernorm(residual)
        x = self.mlp(x)
        x = x + residual

        return x, attn_weights


class TransformerEncoderLayer(nn.Module):
    def __init__(self, hidden_size=1024, intermediate_size=None, nhead=16, dropout=0.1):
        super().__init__()

        self.attn = MultiheadAttention(hidden_size, nhead, dropout=dropout)
        self.mlp = GluMLP(hidden_size=hidden_size, intermediate_size=intermediate_size)
        self.input_layernorm = RMSNorm(hidden_size, eps=1e-6)
        self.post_attention_layernorm = RMSNorm(hidden_size, eps=1e-6)

    def forward(
        self, x, attn_bias=None, key_padding_mask=None, tgt_mask=None, input_pos=None
    ):
        residual = x
        x = self.input_layernorm(x)
        x, _ = self.attn(
            x,
            x,
            x,
            attn_bias=attn_bias,
            key_padding_mask=key_padding_mask,
            attn_mask=tgt_mask,
            return_weights=False,
            input_pos=input_pos,
        )
        residual = x + residual

        x = self.post_attention_layernorm(residual)
        x = self.mlp(x)
        x = x + residual

        return x


class FishSpeechTransformer(nn.Module):
    def __init__(
        self,
        vocab_size,
        codebook_size,
        num_codebooks,
        hidden_size=1024,
        intermediate_size=None,
        nhead=16,
        num_encoder_layers=12,
        num_decoder_layers=12,
        dropout=0.1,
        alignment_position=-2,
        max_position=8192,
    ):
        super().__init__()

        self.encoder_embedding = nn.Embedding(vocab_size, hidden_size)
        self.decoder_embeddings = nn.ModuleList(
            [nn.Embedding(codebook_size, hidden_size) for _ in range(num_codebooks)]
        )
        self.decoder_head = nn.Linear(hidden_size, codebook_size * num_codebooks)
        self.codebook_size = codebook_size
        self.num_codebooks = num_codebooks

        self.encoder = nn.ModuleList(
            [
                TransformerEncoderLayer(
                    hidden_size=hidden_size,
                    intermediate_size=intermediate_size,
                    nhead=nhead,
                    dropout=dropout,
                )
                for _ in range(num_encoder_layers)
            ]
        )

        self.alignment = CrossAttentionLayer(
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            dropout=dropout,
        )

        if alignment_position < 0:
            alignment_position = num_decoder_layers + alignment_position

        self.alignment_position = alignment_position
        assert 0 <= alignment_position < num_decoder_layers

        self.decoder = nn.ModuleList(
            [
                TransformerEncoderLayer(
                    hidden_size=hidden_size,
                    intermediate_size=intermediate_size,
                    nhead=nhead,
                    dropout=dropout,
                )
                for _ in range(num_decoder_layers)
            ]
        )

        self.alibi = AlibiPostionEmbedding(nhead, max_position)
        self.register_buffer(
            "causual_mask",
            torch.triu(torch.ones(max_position, max_position), diagonal=1).bool(),
        )

        self.max_batch_size = -1
        self.max_seq_length = -1

    def setup_kv_caches(self, max_batch_size, max_seq_length):
        if (
            self.max_seq_length >= max_seq_length
            and self.max_batch_size >= max_batch_size
        ):
            return

        if max_seq_length % 8 != 0:
            max_seq_length = max_seq_length + (8 - max_seq_length % 8)

        self.max_seq_length = max_seq_length
        self.max_batch_size = max_batch_size

        for b in self.decoder:
            b.attn.kv_cache = KVCache(
                max_batch_size, max_seq_length, b.attn.nhead, b.attn.head_dim
            )

    def forward(self, inputs, codes, input_mask=None, codes_mask=None):
        # x: (B, T)
        # y: (B, C, T)
        inputs = self.encoder_embedding(inputs)
        codes = rearrange(codes, "b c t -> c b t")
        codes = torch.stack(
            [emb(code) for emb, code in zip(self.decoder_embeddings, codes)], dim=0
        )
        codes = torch.mean(codes, dim=0)  # (B, T)

        attn_bias = self.alibi(inputs)
        for layer in self.encoder:
            inputs = layer(inputs, attn_bias=attn_bias, key_padding_mask=input_mask)

        attn_bias = self.alibi(codes)
        causual_mask = self.causual_mask[: codes.shape[1], : codes.shape[1]]

        for idx, layer in enumerate(self.decoder):
            if idx == self.alignment_position:
                codes, _ = self.alignment(
                    codes, inputs, memory_key_padding_mask=input_mask
                )

            codes = layer(
                codes,
                attn_bias=attn_bias,
                key_padding_mask=codes_mask,
                tgt_mask=causual_mask,
            )

        codes = self.decoder_head(codes)
        codes = rearrange(
            codes, "b t (c d) -> b c t d", c=self.num_codebooks, d=self.codebook_size
        )

        return codes

    def sample_decoder(
        self,
        x: torch.Tensor,
        context: torch.Tensor,
        input_pos: torch.Tensor,
        **sampling_kwargs,
    ):
        attn_bias = self.alibi.alibi[:, input_pos, : self.max_seq_length]
        causual_mask = self.causual_mask[input_pos, : self.max_seq_length]

        x = rearrange(x, "b c t -> c b t")
        x = torch.stack(
            [emb(code) for emb, code in zip(self.decoder_embeddings, x)], dim=0
        )
        x = torch.mean(x, dim=0)  # (B, T)

        for idx, layer in enumerate(self.decoder):
            if idx == self.alignment_position:
                x, _ = self.alignment(x, context)

            x = layer(
                x, attn_bias=attn_bias, input_pos=input_pos, tgt_mask=causual_mask
            )

        x = self.decoder_head(x)
        x = rearrange(
            x, "b t (c d) -> b c t d", c=self.num_codebooks, d=self.codebook_size
        )

        # Never predict EOS or BOS for sub-codebooks
        x[:, 1:, :2] = -float("Inf")

        next_token, probs = [], []
        for i in range(self.num_codebooks):
            next_token_i, probs_i = self.sample(x[:, i], **sampling_kwargs)
            next_token.append(next_token_i)
            probs.append(probs_i)

        return torch.stack(next_token, dim=0), torch.stack(probs, dim=0)

    @staticmethod
    def multinomial_sample_one_no_sync(
        probs_sort,
    ):  # Does multinomial sampling without a cuda synchronization
        q = torch.empty_like(probs_sort).exponential_(1)
        return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)

    @staticmethod
    def logits_to_probs(
        logits,
        temperature: float = 1.0,
        top_p: Optional[int] = None,
        top_k: Optional[int] = None,
    ):
        if top_p is not None:
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            cum_probs = torch.cumsum(
                torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1
            )
            sorted_indices_to_remove = cum_probs > top_p
            sorted_indices_to_remove[0] = False  # keep at least one option
            indices_to_remove = sorted_indices_to_remove.scatter(
                dim=0, index=sorted_indices, src=sorted_indices_to_remove
            )
            logits = logits.masked_fill(indices_to_remove, -float("Inf"))

        logits = logits / max(temperature, 1e-5)

        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            pivot = v.select(-1, -1).unsqueeze(-1)
            logits = torch.where(logits < pivot, -float("Inf"), logits)

        probs = torch.nn.functional.softmax(logits, dim=-1)
        return probs

    def sample(
        self,
        logits,
        temperature: float = 1.0,
        top_p: Optional[int] = None,
        top_k: Optional[int] = None,
    ):
        probs = self.logits_to_probs(logits[0, -1], temperature, top_p, top_k)
        idx_next = self.multinomial_sample_one_no_sync(probs)
        return idx_next, probs

    def decode_n_tokens(
        self,
        cur_token: torch.Tensor,
        context: torch.Tensor,
        input_pos: torch.Tensor,
        num_new_tokens: int,
        callback=lambda _: _,
        **sampling_kwargs,
    ):
        new_tokens, new_probs = [], []
        # Sliding context window
        batch_size = 1
        back_map = torch.zeros(
            [batch_size, 1], device=cur_token.device, dtype=torch.long
        )

        for i in range(num_new_tokens):
            next_token, next_prob = self.sample_decoder(
                cur_token, context, input_pos, **sampling_kwargs
            )

            # index_map = torch.arange(6, device=cur_token.device)
            # index_map = back_map[:, -1:] + index_map.repeat(batch_size, 1)
            # add = torch.arange(batch_size, device=index_map.device).unsqueeze(1) #N, 1
            # index_map = index_map + add * t_length

            input_pos += 1
            new_tokens.append(next_token.clone())
            callback(new_tokens[-1])
            new_probs.append(next_prob.clone())

            if next_token[0, 0] == 1:
                break

            cur_token = next_token.view(1, self.num_codebooks, -1)

        return new_tokens, new_probs

    def compile(self):
        self.sampler_decoder = torch.compile(
            self.sample_decoder, mode="reduce-overhead", fullgraph=True
        )

    @torch.no_grad()
    def inference(self, inputs, prompt=None, max_new_tokens=1024, **sampling_kwargs):
        # inputs: (B, T)
        # prompt: (B, C, T)

        assert inputs.size(0) == 1, "Only support batch size 1 for now"

        if prompt is None:
            prompt = torch.tensor(
                [[[0]] * self.num_codebooks], device=inputs.device, dtype=torch.long
            )

        T = prompt.size(2)
        T_new = T + max_new_tokens

        # Encode Features
        inputs = self.encoder_embedding(inputs)
        attn_bias = self.alibi(inputs)
        for layer in self.encoder:
            inputs = layer(inputs, attn_bias=attn_bias)

        device, dtype = inputs.device, inputs.dtype

        # Decode
        with torch.device(inputs.device):
            self.setup_kv_caches(max_batch_size=1, max_seq_length=T_new)

        # create an empty tensor of the expected final shape and fill in the current tokens
        empty = torch.empty(
            (1, self.num_codebooks, T_new), dtype=torch.long, device=device
        )
        empty[:, :, :T] = prompt
        seq = empty
        input_pos = torch.arange(0, T, device=device)

        # prefill
        next_token, _ = self.sample_decoder(
            prompt.view(1, self.num_codebooks, -1), inputs, input_pos, **sampling_kwargs
        )
        seq[:, :, T] = next_token

        # create an empty tensor of the expected final shape and fill in the current tokens
        input_pos = torch.tensor([T], device=device, dtype=torch.long)
        generated_tokens, _ = self.decode_n_tokens(
            next_token.view(1, self.num_codebooks, -1),
            context=inputs,
            input_pos=input_pos,
            num_new_tokens=max_new_tokens - 1,
            **sampling_kwargs,
        )

        generated_tokens = torch.stack(generated_tokens, dim=-1)
        seq = seq[:, :, : T + 1 + generated_tokens.size(-1)]
        seq[:, :, T + 1 :] = generated_tokens

        return seq


if __name__ == "__main__":
    # mha = MultiheadAttention(512, 8, dropout=0)
    # mha.eval()
    # mha.cuda()

    # q, k, v = torch.randn(3, 10, 16, 512)
    # q, k, v = q.cuda(), k.cuda(), v.cuda()
    # alibi = AlibiPostionEmbedding(8, 1024)

    # mha.bfloat16()
    # q, k, v = q.bfloat16(), k.bfloat16(), v.bfloat16()
    # bias = alibi(q).bfloat16()

    # # Causual mask
    # attn_mask = torch.triu(torch.ones(16, 16), diagonal=1).bool().cuda()
    # o, w = mha(q, k, v, return_weights=True, attn_bias=bias, attn_mask=attn_mask)

    # print(o.size())
    # print(w.size())

    # o1, w = mha(q, k, v, return_weights=False, attn_bias=bias, attn_mask=attn_mask)
    # print(o1.size())

    # print(o[0], o1.float()[0])

    # assert torch.allclose(o.float(), o1.float(), atol=1e-2, rtol=1e-2)
    # print("ok")

    # cross = CrossAttentionLayer(512, 1024, dropout=0)
    # cross.eval()
    # cross.cuda()

    # tgt = torch.randn(3, 10, 512).cuda()
    # memory = torch.randn(3, 20, 512).cuda()
    # o, w = cross(tgt, memory)

    # print(o.size())
    # print(w.size())

    # ten = TransformerEncoderLayer(512, 1024, 8, dropout=0)
    # ten.eval()
    # ten.cuda()

    # tgt = torch.randn(3, 10, 512).cuda()
    # o = ten(tgt)
    # print(o.size())

    trans = (
        FishSpeechTransformer(
            vocab_size=30000,
            codebook_size=120,
            num_codebooks=4,
            hidden_size=1024,
            intermediate_size=None,
            nhead=16,
            num_encoder_layers=12,
            num_decoder_layers=12,
        )
        .bfloat16()
        .cuda()
    )
    # Print n param
    print("Total params:", sum(i.numel() for i in trans.parameters()) / 1024 / 1024)
    inputs = torch.randint(0, 1000, (1, 16)).cuda()
    codes = torch.randint(0, 120, (1, 4, 128)).cuda()
    print(trans(inputs, codes).size())

    r = trans.inference(inputs, max_new_tokens=1024, top_k=5, temperature=0.3)
    print(r)