from typing import Optional, Tuple
import torch

from wenet.ssl.bestrq.mask import compute_mask_indices
from wenet.utils.mask import make_pad_mask


class BestRQModel(torch.nn.Module):
    def __init__(
        self,
        encoder: torch.nn.Module,
        input_dim: int = 256,
        embedding_dim: int = 256,
        num_embeddings: int = 8192,
        num_codebooks: int = 1,
        dropout_rate: float = 0.1,
        mask_prob: float = 0.01,
        mask_length: int = 10,
        min_masks: int = 2,
        layer_norm_epsilon=1e-5,
    ) -> None:
        super().__init__()

        assert mask_prob > 0.0

        self.mask_prob = mask_prob
        # NOTE: should filter audio less than mask_length
        self.mask_length = mask_length
        self.min_masks = min_masks

        self.input_dropout = torch.nn.Dropout(dropout_rate)

        # [embedding_dim, num_embeddings]
        random_embedding_weight = torch.empty(
            num_codebooks, embedding_dim, num_embeddings, requires_grad=False
        )
        self.embeddings = torch.nn.init.normal_(random_embedding_weight)

        random_projection_weight = torch.empty(
            input_dim, embedding_dim, requires_grad=False
        )
        self.projection = torch.nn.init.xavier_normal_(random_projection_weight)

        mask_emb_weight = torch.Tensor(input_dim)
        mask_emb_weight.requires_grad = True
        self.mask_emb = torch.nn.init.normal_(mask_emb_weight, mean=0, std=0.1)

        self.input_layer_norm = torch.nn.LayerNorm(input_dim, layer_norm_epsilon)
        self.encoder = encoder
        self.encoder_top_n_out = torch.nn.parameter.Parameter(
            torch.Tensor(num_codebooks, self.encoder.output_size(), num_embeddings)
        )

    def forward(
        self,
        xs: torch.Tensor,
        xs_lens: torch.Tensor,
        text: Optional[torch.Tensor] = None,
        text_length: Optional[torch.Tensor] = None,
    ):
        # should support nonstreamming and streamming
        # TODO(Mddct): streamming future
        # eg: full attenton and chunk or  dynamic chunk training
        # 1 forward subsampling
        xs, pos_emb, masks = self._forward_subsampling(xs, xs_lens)
        unmasked_xs = xs
        # 2 mask features
        # 2.0 apply mask
        masked_xs, masked_masks = self._apply_mask(xs)
        # 2.1 get nearest embedding
        target_ids = self._nearest_embedding_idx(unmasked_xs)
        # 3 forward xxx-formaer block
        out, out_mask = self._forward_encoder_blocks(masked_xs, masks, pos_emb, masks)
        # 4 get logits
        out = out.unsqueeze(1)  # [B, 1, T', dim]
        top_n_out = self.encoder_top_n_out.unsqueeze(
            0
        )  # [num_codebooks, dim, num_embeddings]
        out = torch.matmul(out, top_n_out)  # [B, num_codebooks, T', num_embeddings]

        # 5 compute loss
        loss = self._compute_loss(out, target_ids, out_mask.squeeze(1) * masked_masks)
        return {"loss": loss}

    def _compute_loss(
        self, input: torch.Tensor, target: torch.Tensor, mask: torch.Tensor
    ):
        input = input.transpose(1, 3)  # [B, num_embeddings, T' num_codebooks]
        entropy = torch.nn.functional.cross_entropy(
            input, target, reduction="none"
        )  # [B, T', num_codebooks]
        # stop gradient for non mask area
        loss = entropy * mask.unsqueeze(2)
        return loss.sum() / (mask.sum() * loss.size(2))

    def _forward_encoder_blocks(
        self,
        xs: torch.Tensor,
        xs_masks: torch.Tensor,
        pos_emb: torch.Tensor,
        mask_pad: torch.Tensor,
    ):
        masks = xs_masks
        for layer in self.encoder.encoders:
            xs, masks, _, _ = layer(xs, xs_masks, pos_emb, mask_pad)
        if self.encoder.normalize_before:
            xs = self.encoder.after_norm(xs)
        # Here we assume the mask is not changed in encoder layers, so just
        # return the masks before encoder layers, and the masks will be used
        # for cross attention with decoder later
        return xs, masks

    def _nearest_embedding_idx(self, xs: torch.Tensor) -> torch.Tensor:
        xs = self.input_layer_norm(xs)
        xs = self.input_dropout(xs)
        xs = torch.matmul(xs, self.projection.to(xs.device))

        B, T, C = xs.size()
        flattened_input = xs.view(-1, C)
        embeddings = self.embeddings.to(
            xs.device
        )  # [num_codebooks, embedding_dim, num_embeddings]
        # [num_codebooks, B*T, num_embeddings]
        distance = (
            torch.sum(flattened_input**2, dim=1, keepdim=True).unsqueeze(0)
            + torch.sum(embeddings**2, dim=1, keepdim=True)
            - 2 * torch.matmul(flattened_input.unsqueeze(0), embeddings)
        )

        out = torch.argmin(distance, dim=-1)  # [num_codebooks, B*T]
        out = out.transpose(0, 1)  # [B*T, num_codebooks]
        return out.reshape(B, T, -1)  # [B, T, num_codebooks]

    def _apply_mask(self, xs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        masks = compute_mask_indices(
            xs.size()[:-1],
            self.mask_prob,
            self.mask_length,
            self.min_masks,
            device=xs.device,
        )
        masks_expand = masks.unsqueeze(-1)  # [B, T, 1]

        mask_emb = self.mask_emb.to(xs.device).view(1, 1, -1)
        xs = torch.where(masks_expand, mask_emb, xs)
        return xs, masks

    def _forward_subsampling(
        self, xs: torch.Tensor, xs_lens: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        T = xs.size(1)
        masks = ~make_pad_mask(xs_lens, T).unsqueeze(1)  # (B, 1, T)
        if self.encoder.global_cmvn is not None:
            xs = self.encoder.global_cmvn(xs)
        xs, pos_emb, masks = self.encoder.embed(xs, masks)
        return xs, pos_emb, masks