from typing import Union
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from math import log2, ceil

from dac.nn.layers import WNConv1d
from dac.nn.vae import VAEBottleneck


def log(t, eps=1e-5):
    return t.clamp(min=eps).log()


def entropy(prob):
    return (-prob * log(prob)).sum(dim=-1)


class VectorQuantize(nn.Module):
    """
    Implementation of VQ similar to Karpathy's repo:
    https://github.com/karpathy/deep-vector-quantization
    Additionally uses following tricks from Improved VQGAN
    (https://arxiv.org/pdf/2110.04627.pdf):
        1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
            for improved codebook usage
        2. l2-normalized codes: Converts euclidean distance to cosine similarity which
            improves training stability
    """

    def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):
        super().__init__()
        self.codebook_size = codebook_size
        self.codebook_dim = codebook_dim

        self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)
        self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)
        self.codebook = nn.Embedding(codebook_size, codebook_dim)

    def forward(self, z):
        """Quantized the input tensor using a fixed codebook and returns
        the corresponding codebook vectors

        Parameters
        ----------
        z : Tensor[B x D x T]

        Returns
        -------
        Tensor[B x D x T]
            Quantized continuous representation of input
        Tensor[1]
            Commitment loss to train encoder to predict vectors closer to codebook
            entries
        Tensor[1]
            Codebook loss to update the codebook
        Tensor[B x T]
            Codebook indices (quantized discrete representation of input)
        Tensor[B x D x T]
            Projected latents (continuous representation of input before quantization)
        """

        # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
        z_e = self.in_proj(z)  # z_e : (B x D x T)
        z_q, indices = self.decode_latents(z_e)

        commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
        codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])

        z_q = (
            z_e + (z_q - z_e).detach()
        )  # noop in forward pass, straight-through gradient estimator in backward pass

        z_q = self.out_proj(z_q)

        # orthogonal_loss
        b, n, d = z_q.shape
        normed_codes = F.normalize(z_q, dim=-1)
        cosine_sim = normed_codes @ normed_codes.transpose(1, 2)
        orthogonal_loss = (cosine_sim**2).sum() / (b * n**2) - (1 / n)

        return z_q, commitment_loss, codebook_loss, orthogonal_loss, indices, z_e

    def embed_code(self, embed_id):
        return F.embedding(embed_id, self.codebook.weight)

    def decode_code(self, embed_id):
        return self.embed_code(embed_id).transpose(1, 2)

    def decode_latents(self, latents):
        encodings = rearrange(latents, "b d t -> (b t) d")
        codebook = self.codebook.weight  # codebook: (N x D)

        # L2 normalize encodings and codebook (ViT-VQGAN)
        encodings = F.normalize(encodings)
        codebook = F.normalize(codebook)

        # Compute euclidean distance with codebook
        dist = (
            encodings.pow(2).sum(1, keepdim=True)
            - 2 * encodings @ codebook.t()
            + codebook.pow(2).sum(1, keepdim=True).t()
        )
        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
        z_q = self.decode_code(indices)
        return z_q, indices


class LFQuantize(nn.Module):
    """
    Implementation of VQ similar to Karpathy's repo:
    https://github.com/karpathy/deep-vector-quantization
    Additionally uses following tricks from Improved VQGAN
    (https://arxiv.org/pdf/2110.04627.pdf):
        1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
            for improved codebook usage
        2. l2-normalized codes: Converts euclidean distance to cosine similarity which
            improves training stability
    """

    def __init__(self, input_dim: int, codebook_size: int):
        super().__init__()
        assert codebook_size == 2 ** ceil(
            log2(codebook_size)
        ), "Codebook size must be a power of 2"
        self.codebook_size = codebook_size
        self.codebook_dim = int(log2(codebook_size))
        self.register_buffer("mask", 2 ** torch.arange(self.codebook_dim))

        self.in_proj = WNConv1d(input_dim, self.codebook_dim, kernel_size=1)
        self.out_proj = WNConv1d(self.codebook_dim, input_dim, kernel_size=1)

    def forward(self, z):
        """Quantized the input tensor using a fixed codebook and returns
        the corresponding codebook vectors

        Parameters
        ----------
        z : Tensor[B x D x T]

        Returns
        -------
        Tensor[B x D x T]
            Quantized continuous representation of input
        Tensor[1]
            Commitment loss to train encoder to predict vectors closer to codebook
            entries
        Tensor[1]
            Codebook loss to update the codebook
        Tensor[B x T]
            Codebook indices (quantized discrete representation of input)
        Tensor[B x D x T]
            Projected latents (continuous representation of input before quantization)
        """

        # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
        z_e = self.in_proj(z)  # z_e : (B x D x T)
        z_q = torch.where(z_e > 0, 1.0, -1.0).detach()
        indices = self.bits_to_index(z_q)
        commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])

        z_q = (
            z_e + (z_q - z_e).detach()
        )  # noop in forward pass, straight-through gradient estimator in backward pass

        z_q = self.out_proj(z_q)

        # orthogonal_loss
        b, n, d = z_q.shape
        normed_codes = F.normalize(z_q, dim=-1)
        cosine_sim = normed_codes @ normed_codes.transpose(1, 2)
        orthogonal_loss = (cosine_sim**2).sum() / (b * n**2) - (1 / n)

        return (
            z_q,
            commitment_loss,
            commitment_loss * 0.0,
            orthogonal_loss,
            indices,
            z_e,
        )

    @torch.no_grad()
    def bits_to_index(self, bits: torch.Tensor):
        """
        Convert bits to index

        Parameters
        ----------
        bits : Tensor[B x N x T]
            Quantized discrete representation of input"""

        b, n, t = bits.shape
        bits = rearrange(bits, "b n t -> (b t) n")
        bits = (bits > 0).float()
        indices = bits.float() @ self.mask.float()
        return rearrange(indices, "(b t) -> b t", b=b).int().detach()


class PassthroughQuantize(nn.Module):
    """
    Dont quantize the input, just pass it through
    """

    def __init__(self):
        super().__init__()

    def forward(self, z, *args, **kwargs):
        """Quantized the input tensor using a fixed codebook and returns
        the corresponding codebook vectors
        """

        # return Return(z, None, z, 0, 0)
        return {
            "z": z,
            "latents": z,
        }


class ResidualOp(nn.Module):
    def forward(self, z, z_q):
        return z - z_q


class MLPResidualOp(nn.Module):
    def __init__(self, input_dim: int):
        super().__init__()
        hidden_dim = input_dim * 4
        self.mlp = nn.Sequential(
            nn.Linear(input_dim * 2, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, input_dim),
        )

    def forward(self, z, z_q):
        t = self.mlp(torch.cat([z, z_q], dim=1).transpose(1, 2)).transpose(1, 2)
        assert t.shape == z.shape
        return z - t


class ResidualVectorQuantize(nn.Module):
    """
    Introduced in SoundStream: An end2end neural audio codec
    https://arxiv.org/abs/2107.03312
    """

    def __init__(
        self,
        input_dim: int = 512,
        n_codebooks: int = 9,
        codebook_size: int = 1024,
        codebook_dim: Union[int, list] = 8,
        quantizer_dropout: float = 0.0,
        use_vae: bool = False,
        residual_type: str = "subtract",
        use_lfq: bool = False,
    ):
        super().__init__()
        if isinstance(codebook_dim, int):
            codebook_dim = [codebook_dim for _ in range(n_codebooks)]

        self.n_codebooks = n_codebooks
        self.codebook_dim = codebook_dim
        self.codebook_size = codebook_size
        self.use_vae = use_vae
        if use_vae:
            self.vae = VAEBottleneck(input_dim, input_dim)

        if use_lfq:
            self.quantizers = nn.ModuleList(
                [LFQuantize(input_dim, codebook_size) for _ in range(n_codebooks)]
            )
        else:
            self.quantizers = nn.ModuleList(
                [
                    VectorQuantize(input_dim, codebook_size, codebook_dim[i])
                    for i in range(n_codebooks)
                ]
            )
        self.quantizer_dropout = quantizer_dropout

        self.residual_type = residual_type
        if residual_type == "mlp":
            self.residuals = nn.ModuleList(
                [MLPResidualOp(input_dim) for _ in range(n_codebooks - 1)]
            )
        elif residual_type == "subtract":
            self.residuals = nn.ModuleList(
                [ResidualOp() for _ in range(n_codebooks - 1)]
            )

    def forward(self, z, n_quantizers: int = None):
        """Quantized the input tensor using a fixed set of `n` codebooks and returns
        the corresponding codebook vectors
        Parameters
        ----------
        z : Tensor[B x D x T]
        n_quantizers : int, optional
            No. of quantizers to use
            (n_quantizers < self.n_codebooks ex: for quantizer dropout)
            Note: if `self.quantizer_dropout` is True, this argument is ignored
                when in training mode, and a random number of quantizers is used.
        Returns
        -------
        dict
            A dictionary with the following keys:

            "z" : Tensor[B x D x T]
                Quantized continuous representation of input
            "codes" : Tensor[B x N x T]
                Codebook indices for each codebook
                (quantized discrete representation of input)
            "latents" : Tensor[B x N*D x T]
                Projected latents (continuous representation of input before quantization)
            "vq/commitment_loss" : Tensor[1]
                Commitment loss to train encoder to predict vectors closer to codebook
                entries
            "vq/codebook_loss" : Tensor[1]
                Codebook loss to update the codebook
        """

        out = {}
        if self.use_vae:
            vae = self.vae(z)
            z = vae["z"]
            out.update(vae)

        z_q = 0
        residual = z
        commitment_loss = 0
        codebook_loss = 0
        orthogonal_loss = 0

        codebook_indices = []
        latents = []

        if n_quantizers is None:
            n_quantizers = self.n_codebooks
        if self.training:
            n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1
            dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))
            n_dropout = int(z.shape[0] * self.quantizer_dropout)
            n_quantizers[:n_dropout] = dropout[:n_dropout]
            n_quantizers = n_quantizers.to(z.device)

        for i, quantizer in enumerate(self.quantizers):
            if self.training is False and i >= n_quantizers:
                break

            (
                z_q_i,
                commitment_loss_i,
                codebook_loss_i,
                orthogonal_loss_i,
                indices_i,
                z_e_i,
            ) = quantizer(residual)

            # Create mask to apply quantizer dropout
            mask = (
                torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
            )
            z_q = z_q + z_q_i * mask[:, None, None]

            if i < len(self.residuals):
                residual = self.residuals[i](residual, z_q_i)

            # Sum losses
            commitment_loss += (commitment_loss_i * mask).mean()
            codebook_loss += (codebook_loss_i * mask).mean()
            orthogonal_loss += (orthogonal_loss_i * mask).mean()

            codebook_indices.append(indices_i)
            latents.append(z_e_i)

        codes = torch.stack(codebook_indices, dim=1)
        latents = torch.cat(latents, dim=1)

        out.update(
            {
                "z": z_q,
                "codes": codes,
                "latents": latents,
                "vq/commitment_loss": commitment_loss,
                "vq/codebook_loss": codebook_loss,
                "vq/orthogonal_loss": orthogonal_loss,
            }
        )
        return out

    def from_codes(self, codes: torch.Tensor):
        """Given the quantized codes, reconstruct the continuous representation
        Parameters
        ----------
        codes : Tensor[B x N x T]
            Quantized discrete representation of input
        Returns
        -------
        Tensor[B x D x T]
            Quantized continuous representation of input
        """
        z_q = 0.0
        z_p = []
        n_codebooks = codes.shape[1]
        for i in range(n_codebooks):
            z_p_i = self.quantizers[i].decode_code(codes[:, i, :])
            z_p.append(z_p_i)

            z_q_i = self.quantizers[i].out_proj(z_p_i)
            z_q = z_q + z_q_i
        return z_q, torch.cat(z_p, dim=1), codes

    def from_latents(self, latents: torch.Tensor):
        """Given the unquantized latents, reconstruct the
        continuous representation after quantization.

        Parameters
        ----------
        latents : Tensor[B x N x T]
            Continuous representation of input after projection

        Returns
        -------
        Tensor[B x D x T]
            Quantized representation of full-projected space
        Tensor[B x D x T]
            Quantized representation of latent space
        """
        z_q = 0
        z_p = []
        codes = []
        dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])

        n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[
            0
        ]
        for i in range(n_codebooks):
            j, k = dims[i], dims[i + 1]
            z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])
            z_p.append(z_p_i)
            codes.append(codes_i)

            z_q_i = self.quantizers[i].out_proj(z_p_i)
            z_q = z_q + z_q_i

        return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)


class GroupedResidualVectorQuantize(nn.Module):
    """
    Introduced in SoundStream: An end2end neural audio codec
    https://arxiv.org/abs/2107.03312
    """

    def __init__(
        self,
        input_dim: int = 512,
        n_codebooks_per_level: list[int] = [2, 2, 2, 2],
        codebook_size: int = 1024,
        codebook_dim: int = 8,
        quantizer_dropout: float = 0.0,
    ):
        super().__init__()

        self.n_codebooks_per_level = n_codebooks_per_level
        self.input_dim = input_dim
        self.codebook_dim = codebook_dim
        self.codebook_size = codebook_size

        for n_codebooks in n_codebooks_per_level:
            assert (
                input_dim % n_codebooks == 0
            ), f"input_dim ({input_dim}) must be divisible by n_codebooks ({n_codebooks})"

        self.quantizers = nn.ModuleList(
            [
                nn.ModuleList(
                    [
                        VectorQuantize(
                            input_dim // n_codebooks, codebook_size, codebook_dim
                        )
                        for _ in range(n_codebooks)
                    ]
                )
                for n_codebooks in n_codebooks_per_level
            ]
        )
        self.quantizer_dropout = quantizer_dropout

    @property
    def n_codebooks(self):
        return sum(self.n_codebooks_per_level)

    @property
    def n_levels(self):
        return len(self.n_codebooks_per_level)

    def forward(self, z, n_quantizers: int = None):
        """Quantized the input tensor using a fixed set of `n` codebooks and returns
        the corresponding codebook vectors
        Parameters
        ----------
        z : Tensor[B x D x T]
        n_quantizers : int, optional
            No. of quantizers to use
            (n_quantizers < self.n_codebooks ex: for quantizer dropout)
            Note: if `self.quantizer_dropout` is True, this argument is ignored
                when in training mode, and a random number of quantizers is used.
        Returns
        -------
        dict
            A dictionary with the following keys:

            "z" : Tensor[B x D x T]
                Quantized continuous representation of input
            "codes" : Tensor[B x N x T]
                Codebook indices for each codebook
                (quantized discrete representation of input)
            "latents" : Tensor[B x N*D x T]
                Projected latents (continuous representation of input before quantization)
            "vq/commitment_loss" : Tensor[1]
                Commitment loss to train encoder to predict vectors closer to codebook
                entries
            "vq/codebook_loss" : Tensor[1]
                Codebook loss to update the codebook
        """
        z_q = torch.zeros_like(z)
        # residual = z
        commitment_loss = 0
        codebook_loss = 0
        orthogonal_loss = 0

        codebook_indices = []
        latents = []

        if n_quantizers is None:
            n_quantizers = self.n_codebooks
        if self.training:
            n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1
            dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))
            n_dropout = int(z.shape[0] * self.quantizer_dropout)
            n_quantizers[:n_dropout] = dropout[:n_dropout]
            n_quantizers = n_quantizers.to(z.device)

        for i, quantizers in enumerate(self.quantizers):
            if self.training is False and i >= n_quantizers:
                break
            split_size = z.shape[1] // len(quantizers)
            for j, quantizer in enumerate(quantizers):
                split_slice = slice(j * split_size, (j + 1) * split_size)
                (
                    z_q_i,
                    commitment_loss_i,
                    codebook_loss_i,
                    orthogonal_loss_i,
                    indices_i,
                    z_e_i,
                ) = quantizer(z[:, split_slice, :] - z_q[:, split_slice, :])

                # Create mask to apply quantizer dropout
                quantizer_idx = sum(self.n_codebooks_per_level[:i]) + j
                mask = (
                    torch.full((z.shape[0],), fill_value=quantizer_idx, device=z.device)
                    < n_quantizers
                )
                z_q[:, split_slice, :] = (
                    z_q[:, split_slice, :] + z_q_i * mask[:, None, None]
                )

                # Sum losses
                commitment_loss += (commitment_loss_i * mask).mean()
                codebook_loss += (codebook_loss_i * mask).mean()
                orthogonal_loss += (orthogonal_loss_i * mask).mean()

                codebook_indices.append(indices_i)
                latents.append(z_e_i)

        codes = torch.stack(codebook_indices, dim=1)
        latents = torch.cat(latents, dim=1)

        return {
            "z": z_q,
            "codes": codes,
            "latents": latents,
            "vq/commitment_loss": commitment_loss,
            "vq/codebook_loss": codebook_loss,
            "vq/orthogonal_loss": orthogonal_loss,
        }

    def from_codes(self, codes: torch.Tensor):
        """Given the quantized codes, reconstruct the continuous representation
        Parameters
        ----------
        codes : Tensor[B x N x T]
            Quantized discrete representation of input
        Returns
        -------
        Tensor[B x D x T]
            Quantized continuous representation of input
        """
        B, N, T = codes.shape
        z_q = torch.zeros(B, self.input_dim, T, device=codes.device)
        z_p = []
        cb_idx = 0
        for i in range(self.n_levels):
            z_p_i = []
            split_size = self.input_dim // self.n_codebooks_per_level[i]
            for j, quantizer in enumerate(self.quantizers[i]):
                z_p_i_j = quantizer.decode_code(codes[:, cb_idx, :])
                z_p_i.append(z_p_i_j)

                z_q_i = quantizer.out_proj(z_p_i_j)
                z_q[:, split_size * j : split_size * (j + 1), :] += z_q_i
                cb_idx += 1
            z_p.append(torch.cat(z_p_i, dim=1))

        return z_q, torch.cat(z_p, dim=1), codes

    def from_latents(self, latents: torch.Tensor):
        """Given the unquantized latents, reconstruct the
        continuous representation after quantization.

        Parameters
        ----------
        latents : Tensor[B x N x T]
            Continuous representation of input after projection

        Returns
        -------
        Tensor[B x D x T]
            Quantized representation of full-projected space
        Tensor[B x D x T]
            Quantized representation of latent space
        """
        raise NotImplementedError
        z_q = 0
        z_p = []
        codes = []
        dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])

        n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[
            0
        ]
        for i in range(n_codebooks):
            j, k = dims[i], dims[i + 1]
            z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])
            z_p.append(z_p_i)
            codes.append(codes_i)

            z_q_i = self.quantizers[i].out_proj(z_p_i)
            z_q = z_q + z_q_i

        return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)


@torch.no_grad()
@torch.inference_mode()
def main():
    rvq = ResidualVectorQuantize(quantizer_dropout=True).eval()
    x = torch.randn(16, 512, 80)
    y = rvq(x)
    print(y["latents"].shape)
    assert x.shape == y["z"].shape, f"{x.shape} != {y['z'].shape}"

    from_codes = rvq.from_codes(y["codes"])
    assert x.shape == from_codes[0].shape
    assert x.shape == y["z"].shape, f"{x.shape} != {y['z'].shape}"
    assert torch.allclose(from_codes[0], y["z"], atol=1e-2), from_codes[0] - y["z"]

    from_latents = rvq.from_latents(y["latents"])
    assert x.shape == from_latents[0].shape
    assert torch.allclose(from_latents[0], y["z"], atol=1e-2), from_latents[0] - y["z"]

    grvq = GroupedResidualVectorQuantize(quantizer_dropout=True).eval()
    y = grvq(x)
    print(y["latents"].shape)
    from_codes = grvq.from_codes(y["codes"])
    assert x.shape == from_codes[0].shape
    assert x.shape == y["z"].shape, f"{x.shape} != {y['z'].shape}"
    assert torch.allclose(from_codes[0], y["z"], atol=1e-2), from_codes[0] - y["z"]

    # from_latents = grvq.from_latents(y["latents"])
    # assert x.shape == from_latents[0].shape
    # assert torch.allclose(from_latents[0], y["z"], atol=1e-2), from_latents[0] - y["z"]


if __name__ == "__main__":
    main()