from typing import List, Optional, Tuple

import torch
from torch import nn
from modules.wenet_extractor.utils.common import get_activation, get_rnn


def ApplyPadding(input, padding, pad_value) -> torch.Tensor:
    """
    Args:
        input:   [bs, max_time_step, dim]
        padding: [bs, max_time_step]
    """
    return padding * pad_value + input * (1 - padding)


class PredictorBase(torch.nn.Module):
    # NOTE(Mddct): We can use ABC abstract here, but
    # keep this class simple enough for now
    def __init__(self) -> None:
        super().__init__()

    def init_state(
        self, batch_size: int, device: torch.device, method: str = "zero"
    ) -> List[torch.Tensor]:
        _, _, _ = batch_size, method, device
        raise NotImplementedError("this is a base precictor")

    def batch_to_cache(self, cache: List[torch.Tensor]) -> List[List[torch.Tensor]]:
        _ = cache
        raise NotImplementedError("this is a base precictor")

    def cache_to_batch(self, cache: List[List[torch.Tensor]]) -> List[torch.Tensor]:
        _ = cache
        raise NotImplementedError("this is a base precictor")

    def forward(
        self,
        input: torch.Tensor,
        cache: Optional[List[torch.Tensor]] = None,
    ):
        (
            _,
            _,
        ) = (
            input,
            cache,
        )
        raise NotImplementedError("this is a base precictor")

    def forward_step(
        self, input: torch.Tensor, padding: torch.Tensor, cache: List[torch.Tensor]
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        (
            _,
            _,
            _,
        ) = (
            input,
            padding,
            cache,
        )
        raise NotImplementedError("this is a base precictor")


class RNNPredictor(PredictorBase):
    def __init__(
        self,
        voca_size: int,
        embed_size: int,
        output_size: int,
        embed_dropout: float,
        hidden_size: int,
        num_layers: int,
        bias: bool = True,
        rnn_type: str = "lstm",
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.n_layers = num_layers
        self.hidden_size = hidden_size
        # disable rnn base out projection
        self.embed = nn.Embedding(voca_size, embed_size)
        self.dropout = nn.Dropout(embed_dropout)
        # NOTE(Mddct): rnn base from torch not support layer norm
        # will add layer norm and prune value in cell and layer
        # ref: https://github.com/Mddct/neural-lm/blob/main/models/gru_cell.py
        self.rnn = get_rnn(rnn_type=rnn_type)(
            input_size=embed_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            bias=bias,
            batch_first=True,
            dropout=dropout,
        )
        self.projection = nn.Linear(hidden_size, output_size)

    def forward(
        self,
        input: torch.Tensor,
        cache: Optional[List[torch.Tensor]] = None,
    ) -> torch.Tensor:
        """
        Args:
            input (torch.Tensor): [batch, max_time).
            padding (torch.Tensor): [batch, max_time]
            cache : rnn predictor cache[0] == state_m
                    cache[1] == state_c
        Returns:
            output: [batch, max_time, output_size]
        """

        # NOTE(Mddct): we don't use pack input format
        embed = self.embed(input)  # [batch, max_time, emb_size]
        embed = self.dropout(embed)
        states: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
        if cache is None:
            state = self.init_state(batch_size=input.size(0), device=input.device)
            states = (state[0], state[1])
        else:
            assert len(cache) == 2
            states = (cache[0], cache[1])
        out, (m, c) = self.rnn(embed, states)
        out = self.projection(out)

        # NOTE(Mddct): Although we don't use staate in transducer
        # training forward, we need make it right for padding value
        # so we create forward_step for infering, forward for training
        _, _ = m, c
        return out

    def batch_to_cache(self, cache: List[torch.Tensor]) -> List[List[torch.Tensor]]:
        """
        Args:
           cache: [state_m, state_c]
               state_ms: [1*n_layers, bs, ...]
               state_cs: [1*n_layers, bs, ...]
        Returns:
           new_cache: [[state_m_1, state_c_1], [state_m_2, state_c_2]...]
        """
        assert len(cache) == 2
        state_ms = cache[0]
        state_cs = cache[1]

        assert state_ms.size(1) == state_cs.size(1)

        new_cache: List[List[torch.Tensor]] = []
        for state_m, state_c in zip(
            torch.split(state_ms, 1, dim=1), torch.split(state_cs, 1, dim=1)
        ):
            new_cache.append([state_m, state_c])
        return new_cache

    def cache_to_batch(self, cache: List[List[torch.Tensor]]) -> List[torch.Tensor]:
        """
        Args:
            cache : [[state_m_1, state_c_1], [state_m_1, state_c_1]...]

        Returns:
            new_caceh: [state_ms, state_cs],
                state_ms: [1*n_layers, bs, ...]
                state_cs: [1*n_layers, bs, ...]
        """
        state_ms = torch.cat([states[0] for states in cache], dim=1)
        state_cs = torch.cat([states[1] for states in cache], dim=1)
        return [state_ms, state_cs]

    def init_state(
        self,
        batch_size: int,
        device: torch.device,
        method: str = "zero",
    ) -> List[torch.Tensor]:
        assert batch_size > 0
        # TODO(Mddct): xavier init method
        _ = method
        return [
            torch.zeros(1 * self.n_layers, batch_size, self.hidden_size, device=device),
            torch.zeros(1 * self.n_layers, batch_size, self.hidden_size, device=device),
        ]

    def forward_step(
        self, input: torch.Tensor, padding: torch.Tensor, cache: List[torch.Tensor]
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """
        Args:
            input (torch.Tensor): [batch_size, time_step=1]
            padding (torch.Tensor): [batch_size,1], 1 is padding value
            cache : rnn predictor cache[0] == state_m
                    cache[1] == state_c
        """
        assert len(cache) == 2
        state_m, state_c = cache[0], cache[1]
        embed = self.embed(input)  # [batch, 1, emb_size]
        embed = self.dropout(embed)
        out, (m, c) = self.rnn(embed, (state_m, state_c))

        out = self.projection(out)
        m = ApplyPadding(m, padding.unsqueeze(0), state_m)
        c = ApplyPadding(c, padding.unsqueeze(0), state_c)

        return (out, [m, c])


class EmbeddingPredictor(PredictorBase):
    """Embedding predictor

    Described in:
    https://arxiv.org/pdf/2109.07513.pdf

    embed-> proj -> layer norm -> swish
    """

    def __init__(
        self,
        voca_size: int,
        embed_size: int,
        embed_dropout: float,
        n_head: int,
        history_size: int = 2,
        activation: str = "swish",
        bias: bool = False,
        layer_norm_epsilon: float = 1e-5,
    ) -> None:
        super().__init__()
        # multi head
        self.num_heads = n_head
        self.embed_size = embed_size
        self.context_size = history_size + 1
        self.pos_embed = torch.nn.Linear(
            embed_size * self.context_size, self.num_heads, bias=bias
        )
        self.embed = nn.Embedding(voca_size, self.embed_size)
        self.embed_dropout = nn.Dropout(p=embed_dropout)
        self.ffn = nn.Linear(self.embed_size, self.embed_size)
        self.norm = nn.LayerNorm(self.embed_size, eps=layer_norm_epsilon)
        self.activatoin = get_activation(activation)

    def init_state(
        self, batch_size: int, device: torch.device, method: str = "zero"
    ) -> List[torch.Tensor]:
        assert batch_size > 0
        _ = method
        return [
            torch.zeros(
                batch_size, self.context_size - 1, self.embed_size, device=device
            ),
        ]

    def batch_to_cache(self, cache: List[torch.Tensor]) -> List[List[torch.Tensor]]:
        """
        Args:
            cache : [history]
                history: [bs, ...]
        Returns:
            new_ache : [[history_1], [history_2], [history_3]...]
        """
        assert len(cache) == 1
        cache_0 = cache[0]
        history: List[List[torch.Tensor]] = []
        for h in torch.split(cache_0, 1, dim=0):
            history.append([h])
        return history

    def cache_to_batch(self, cache: List[List[torch.Tensor]]) -> List[torch.Tensor]:
        """
        Args:
            cache : [[history_1], [history_2], [history3]...]

        Returns:
            new_caceh: [history],
                history: [bs, ...]
        """
        history = torch.cat([h[0] for h in cache], dim=0)
        return [history]

    def forward(self, input: torch.Tensor, cache: Optional[List[torch.Tensor]] = None):
        """forward for training"""
        input = self.embed(input)  # [bs, seq_len, embed]
        input = self.embed_dropout(input)
        if cache is None:
            zeros = self.init_state(input.size(0), device=input.device)[0]
        else:
            assert len(cache) == 1
            zeros = cache[0]

        input = torch.cat(
            (zeros, input), dim=1
        )  # [bs, context_size-1 + seq_len, embed]

        input = input.unfold(1, self.context_size, 1).permute(
            0, 1, 3, 2
        )  # [bs, seq_len, context_size, embed]
        # multi head pos: [n_head, embed, context_size]
        multi_head_pos = self.pos_embed.weight.view(
            self.num_heads, self.embed_size, self.context_size
        )

        # broadcast dot attenton
        input_expand = input.unsqueeze(2)  # [bs, seq_len, 1, context_size, embed]
        multi_head_pos = multi_head_pos.permute(
            0, 2, 1
        )  # [num_heads, context_size, embed]

        # [bs, seq_len, num_heads, context_size, embed]
        weight = input_expand * multi_head_pos
        weight = weight.sum(dim=-1, keepdim=False).unsqueeze(
            3
        )  # [bs, seq_len, num_heads, 1, context_size]
        output = weight.matmul(input_expand).squeeze(
            dim=3
        )  # [bs, seq_len, num_heads, embed]
        output = output.sum(dim=2)  # [bs, seq_len, embed]
        output = output / (self.num_heads * self.context_size)

        output = self.ffn(output)
        output = self.norm(output)
        output = self.activatoin(output)
        return output

    def forward_step(
        self,
        input: torch.Tensor,
        padding: torch.Tensor,
        cache: List[torch.Tensor],
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """forward step for inference
        Args:
            input (torch.Tensor): [batch_size, time_step=1]
            padding (torch.Tensor): [batch_size,1], 1 is padding value
            cache: for embedding predictor, cache[0] == history
        """
        assert input.size(1) == 1
        assert len(cache) == 1
        history = cache[0]
        assert history.size(1) == self.context_size - 1
        input = self.embed(input)  # [bs, 1, embed]
        input = self.embed_dropout(input)
        context_input = torch.cat((history, input), dim=1)
        input_expand = context_input.unsqueeze(1).unsqueeze(
            2
        )  # [bs, 1, 1, context_size, embed]

        # multi head pos: [n_head, embed, context_size]
        multi_head_pos = self.pos_embed.weight.view(
            self.num_heads, self.embed_size, self.context_size
        )

        multi_head_pos = multi_head_pos.permute(
            0, 2, 1
        )  # [num_heads, context_size, embed]
        # [bs, 1, num_heads, context_size, embed]
        weight = input_expand * multi_head_pos
        weight = weight.sum(dim=-1, keepdim=False).unsqueeze(
            3
        )  # [bs, 1, num_heads, 1, context_size]
        output = weight.matmul(input_expand).squeeze(dim=3)  # [bs, 1, num_heads, embed]
        output = output.sum(dim=2)  # [bs, 1, embed]
        output = output / (self.num_heads * self.context_size)

        output = self.ffn(output)
        output = self.norm(output)
        output = self.activatoin(output)
        new_cache = context_input[:, 1:, :]
        # TODO(Mddct): we need padding new_cache in future
        # new_cache = ApplyPadding(history, padding, new_cache)
        return (output, [new_cache])


class ConvPredictor(PredictorBase):
    def __init__(
        self,
        voca_size: int,
        embed_size: int,
        embed_dropout: float,
        history_size: int = 2,
        activation: str = "relu",
        bias: bool = False,
        layer_norm_epsilon: float = 1e-5,
    ) -> None:
        super().__init__()

        assert history_size >= 0
        self.embed_size = embed_size
        self.context_size = history_size + 1
        self.embed = nn.Embedding(voca_size, self.embed_size)
        self.embed_dropout = nn.Dropout(p=embed_dropout)
        self.conv = nn.Conv1d(
            in_channels=embed_size,
            out_channels=embed_size,
            kernel_size=self.context_size,
            padding=0,
            groups=embed_size,
            bias=bias,
        )
        self.norm = nn.LayerNorm(embed_size, eps=layer_norm_epsilon)
        self.activatoin = get_activation(activation)

    def init_state(
        self, batch_size: int, device: torch.device, method: str = "zero"
    ) -> List[torch.Tensor]:
        assert batch_size > 0
        assert method == "zero"
        return [
            torch.zeros(
                batch_size, self.context_size - 1, self.embed_size, device=device
            )
        ]

    def cache_to_batch(self, cache: List[List[torch.Tensor]]) -> List[torch.Tensor]:
        """
        Args:
            cache : [[history_1], [history_2], [history3]...]

        Returns:
            new_caceh: [history],
                history: [bs, ...]
        """
        history = torch.cat([h[0] for h in cache], dim=0)
        return [history]

    def batch_to_cache(self, cache: List[torch.Tensor]) -> List[List[torch.Tensor]]:
        """
        Args:
            cache : [history]
                history: [bs, ...]
        Returns:
            new_ache : [[history_1], [history_2], [history_3]...]
        """
        assert len(cache) == 1
        cache_0 = cache[0]
        history: List[List[torch.Tensor]] = []
        for h in torch.split(cache_0, 1, dim=0):
            history.append([h])
        return history

    def forward(self, input: torch.Tensor, cache: Optional[List[torch.Tensor]] = None):
        """forward for training"""
        input = self.embed(input)  # [bs, seq_len, embed]
        input = self.embed_dropout(input)
        if cache is None:
            zeros = self.init_state(input.size(0), device=input.device)[0]
        else:
            assert len(cache) == 1
            zeros = cache[0]

        input = torch.cat(
            (zeros, input), dim=1
        )  # [bs, context_size-1 + seq_len, embed]
        input = input.permute(0, 2, 1)
        out = self.conv(input).permute(0, 2, 1)
        out = self.activatoin(self.norm(out))
        return out

    def forward_step(
        self, input: torch.Tensor, padding: torch.Tensor, cache: List[torch.Tensor]
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """forward step for inference
        Args:
            input (torch.Tensor): [batch_size, time_step=1]
            padding (torch.Tensor): [batch_size,1], 1 is padding value
            cache: for embedding predictor, cache[0] == history
        """
        assert input.size(1) == 1
        assert len(cache) == 1
        history = cache[0]
        assert history.size(1) == self.context_size - 1
        input = self.embed(input)  # [bs, 1, embed]
        input = self.embed_dropout(input)
        context_input = torch.cat((history, input), dim=1)
        input = context_input.permute(0, 2, 1)
        out = self.conv(input).permute(0, 2, 1)
        out = self.activatoin(self.norm(out))

        new_cache = context_input[:, 1:, :]
        # TODO(Mddct): apply padding in future
        return (out, [new_cache])