# Copyright (c) 2024 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from concurrent.futures import ALL_COMPLETED
import numpy as np
import torch
import torch.nn as nn

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

from models.codec.amphion_codec.quantize import ResidualVQ
from models.codec.kmeans.vocos import VocosBackbone


def init_weights(m):
    if isinstance(m, nn.Conv1d):
        nn.init.trunc_normal_(m.weight, std=0.02)
        nn.init.constant_(m.bias, 0)
    if isinstance(m, nn.Linear):
        nn.init.trunc_normal_(m.weight, std=0.02)
        nn.init.constant_(m.bias, 0)


def compute_codebook_perplexity(indices, codebook_size):
    indices = indices.flatten()
    prob = torch.bincount(indices, minlength=codebook_size).float() / indices.size(0)
    perp = torch.exp(-torch.sum(prob * torch.log(prob + 1e-10)))
    return perp


class RepCodec(nn.Module):
    def __init__(
        self,
        codebook_size=8192,
        hidden_size=1024,
        codebook_dim=8,
        vocos_dim=384,
        vocos_intermediate_dim=2048,
        vocos_num_layers=12,
        num_quantizers=1,
        downsample_scale=1,
        cfg=None,
    ):
        super().__init__()
        codebook_size = (
            cfg.codebook_size
            if cfg is not None and hasattr(cfg, "codebook_size")
            else codebook_size
        )
        codebook_dim = (
            cfg.codebook_dim
            if cfg is not None and hasattr(cfg, "codebook_dim")
            else codebook_dim
        )
        hidden_size = (
            cfg.hidden_size
            if cfg is not None and hasattr(cfg, "hidden_size")
            else hidden_size
        )
        vocos_dim = (
            cfg.vocos_dim
            if cfg is not None and hasattr(cfg, "vocos_dim")
            else vocos_dim
        )
        vocos_intermediate_dim = (
            cfg.vocos_intermediate_dim
            if cfg is not None and hasattr(cfg, "vocos_dim")
            else vocos_intermediate_dim
        )
        vocos_num_layers = (
            cfg.vocos_num_layers
            if cfg is not None and hasattr(cfg, "vocos_dim")
            else vocos_num_layers
        )
        num_quantizers = (
            cfg.num_quantizers
            if cfg is not None and hasattr(cfg, "num_quantizers")
            else num_quantizers
        )
        downsample_scale = (
            cfg.downsample_scale
            if cfg is not None and hasattr(cfg, "downsample_scale")
            else downsample_scale
        )

        self.codebook_size = codebook_size
        self.codebook_dim = codebook_dim
        self.hidden_size = hidden_size
        self.vocos_dim = vocos_dim
        self.vocos_intermediate_dim = vocos_intermediate_dim
        self.vocos_num_layers = vocos_num_layers
        self.num_quantizers = num_quantizers
        self.downsample_scale = downsample_scale

        if self.downsample_scale != None and self.downsample_scale > 1:
            self.down = nn.Conv1d(
                self.hidden_size, self.hidden_size, kernel_size=3, stride=2, padding=1
            )
            self.up = nn.Conv1d(
                self.hidden_size, self.hidden_size, kernel_size=3, stride=1, padding=1
            )

        self.encoder = nn.Sequential(
            VocosBackbone(
                input_channels=self.hidden_size,
                dim=self.vocos_dim,
                intermediate_dim=self.vocos_intermediate_dim,
                num_layers=self.vocos_num_layers,
                adanorm_num_embeddings=None,
            ),
            nn.Linear(self.vocos_dim, self.hidden_size),
        )
        self.decoder = nn.Sequential(
            VocosBackbone(
                input_channels=self.hidden_size,
                dim=self.vocos_dim,
                intermediate_dim=self.vocos_intermediate_dim,
                num_layers=self.vocos_num_layers,
                adanorm_num_embeddings=None,
            ),
            nn.Linear(self.vocos_dim, self.hidden_size),
        )

        self.quantizer = ResidualVQ(
            input_dim=hidden_size,
            num_quantizers=num_quantizers,
            codebook_size=codebook_size,
            codebook_dim=codebook_dim,
            quantizer_type="fvq",
            quantizer_dropout=0.0,
            commitment=0.15,
            codebook_loss_weight=1.0,
            use_l2_normlize=True,
        )

        self.reset_parameters()

    def forward(self, x):

        # downsample
        if self.downsample_scale != None and self.downsample_scale > 1:
            x = x.transpose(1, 2)
            x = self.down(x)
            x = F.gelu(x)
            x = x.transpose(1, 2)

        # encoder
        x = self.encoder(x.transpose(1, 2)).transpose(1, 2)

        # vq
        (
            quantized_out,
            all_indices,
            all_commit_losses,
            all_codebook_losses,
            _,
        ) = self.quantizer(x)

        # decoder
        x = self.decoder(quantized_out)

        # up
        if self.downsample_scale != None and self.downsample_scale > 1:
            x = x.transpose(1, 2)
            x = F.interpolate(x, scale_factor=2, mode="nearest")
            x_rec = self.up(x).transpose(1, 2)

        codebook_loss = (all_codebook_losses + all_commit_losses).mean()
        all_indices = all_indices

        return x_rec, codebook_loss, all_indices

    def quantize(self, x):

        if self.downsample_scale != None and self.downsample_scale > 1:
            x = x.transpose(1, 2)
            x = self.down(x)
            x = F.gelu(x)
            x = x.transpose(1, 2)

        x = self.encoder(x.transpose(1, 2)).transpose(1, 2)

        (
            quantized_out,
            all_indices,
            all_commit_losses,
            all_codebook_losses,
            _,
        ) = self.quantizer(x)

        if all_indices.shape[0] == 1:
            return all_indices.squeeze(0), quantized_out.transpose(1, 2)
        return all_indices, quantized_out.transpose(1, 2)

    def reset_parameters(self):
        self.apply(init_weights)


if __name__ == "__main__":
    repcodec = RepCodec(vocos_dim=1024, downsample_scale=2)
    print(repcodec)
    print(sum(p.numel() for p in repcodec.parameters()) / 1e6)
    x = torch.randn(5, 10, 1024)
    x_rec, codebook_loss, all_indices = repcodec(x)
    print(x_rec.shape, codebook_loss, all_indices.shape)
    vq_id, emb = repcodec.quantize(x)
    print(vq_id.shape, emb.shape)