import math
import torch
from torch import nn
from .fvq import FactorizedVectorQuantize


class ResidualVQ(nn.Module):
    """Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf"""

    def __init__(self, *, num_quantizers, codebook_size, **kwargs):
        super().__init__()
        VQ = FactorizedVectorQuantize
        if type(codebook_size) == int:
            codebook_size = [codebook_size] * num_quantizers
        self.layers = nn.ModuleList(
            [VQ(codebook_size=2**size, **kwargs) for size in codebook_size]
        )
        self.num_quantizers = num_quantizers
        self.quantizer_dropout = kwargs.get("quantizer_dropout", 0.0)
        self.dropout_type = kwargs.get("dropout_type", None)

    def forward(self, x, n_quantizers=None):
        quantized_out = 0.0
        residual = x

        all_losses = []
        all_indices = []
        all_quantized = []

        if n_quantizers is None:
            n_quantizers = self.num_quantizers
        if self.training:
            n_quantizers = torch.ones((x.shape[0],)) * self.num_quantizers + 1
            if self.dropout_type == "linear":
                dropout = torch.randint(1, self.num_quantizers + 1, (x.shape[0],))
            elif self.dropout_type == "exp":
                dropout = torch.randint(
                    1, int(math.log2(self.num_quantizers)), (x.shape[0],)
                )
                dropout = torch.pow(2, dropout)
            n_dropout = int(x.shape[0] * self.quantizer_dropout)
            n_quantizers[:n_dropout] = dropout[:n_dropout]
            n_quantizers = n_quantizers.to(x.device)

        for idx, layer in enumerate(self.layers):
            if not self.training and idx >= n_quantizers:
                break
            quantized, indices, loss = layer(residual)

            mask = (
                torch.full((x.shape[0],), fill_value=idx, device=x.device)
                < n_quantizers
            )

            residual = residual - quantized

            quantized_out = quantized_out + quantized * mask[:, None, None]

            # loss
            loss = (loss * mask).mean()

            all_indices.append(indices)
            all_losses.append(loss)
            all_quantized.append(quantized)
        all_losses, all_indices, all_quantized = map(
            torch.stack, (all_losses, all_indices, all_quantized)
        )
        return quantized_out, all_indices, all_losses, all_quantized

    def vq2emb(self, vq):
        # vq: [n_quantizers, B, T]
        quantized_out = 0.0
        for idx, layer in enumerate(self.layers):
            quantized = layer.vq2emb(vq[idx])
            quantized_out += quantized
        return quantized_out

    def get_emb(self):
        embs = []
        for idx, layer in enumerate(self.layers):
            embs.append(layer.get_emb())
        return embs