# Heavily influenced by https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/modules/conditioners.py
import os
import re

os.environ["TOKENIZERS_PARALLELISM"] = "false"

import torch
import logging, warnings
import random
import string
import typing as tp
import gc
import numpy as np

from .adp import NumberEmbedder
from ..inference.utils import set_audio_channels
from .factory import create_pretransform_from_config
from .pretransforms import Pretransform
from ..training.utils import copy_state_dict
from .transformer import ScaledSinusoidalEmbedding
from .utils import load_ckpt_state_dict

import torchaudio
from torch import nn
from suno_utils.utils.s3 import read_from_s3
from suno_utils.tasks.dac_2c_12cb import DAC
from suno_utils.models.musicfm.modeling_MusicFM import MusicFM_MERTLong
from suno_utils.models.dac.nn.quantize_2 import ResidualVectorQuantize


class Conditioner(nn.Module):
    def __init__(
        self,
        dim: int,
        output_dim: int,
        project_out: bool = False,
    ):
        super().__init__()

        self.dim = dim
        self.output_dim = output_dim
        self.proj_out = (
            nn.Linear(dim, output_dim)
            if (dim != output_dim or project_out)
            else nn.Identity()
        )

    def forward(self, x: tp.Any) -> tp.Any:
        raise NotImplementedError()


class IntConditioner(Conditioner):
    def __init__(self, output_dim: int, min_val: int = 0, max_val: int = 512):
        super().__init__(output_dim, output_dim)

        self.min_val = min_val
        self.max_val = max_val
        self.int_embedder = nn.Embedding(
            max_val - min_val + 1, output_dim
        ).requires_grad_(True)

    def forward(self, ints: tp.List[int], device=None) -> tp.Any:
        # self.int_embedder.to(device)

        ints = torch.tensor(ints).to(device)
        ints = ints.clamp(self.min_val, self.max_val)

        int_embeds = self.int_embedder(ints).unsqueeze(1)

        return [int_embeds, torch.ones(int_embeds.shape[0], 1).to(device)]


class NumberConditioner(Conditioner):
    """
    Conditioner that takes a list of floats, normalizes them for a given range, and returns a list of embeddings
    """

    def __init__(self, output_dim: int, min_val: float = 0, max_val: float = 1):
        super().__init__(output_dim, output_dim)

        self.min_val = min_val
        self.max_val = max_val

        self.embedder = NumberEmbedder(features=output_dim)

    def forward(self, floats: tp.List[float], device=None) -> tp.Any:
        # Cast the inputs to floats
        floats = [float(x) for x in floats]

        floats = torch.tensor(floats).to(device)

        floats = floats.clamp(self.min_val, self.max_val)

        normalized_floats = (floats - self.min_val) / (self.max_val - self.min_val)

        # Cast floats to same type as embedder
        embedder_dtype = next(self.embedder.parameters()).dtype
        normalized_floats = normalized_floats.to(embedder_dtype)

        float_embeds = self.embedder(normalized_floats).unsqueeze(1)

        return [float_embeds, torch.ones(float_embeds.shape[0], 1).to(device)]


# this is deprecated
class MERTConditioner(Conditioner):
    def __init__(
        self,
        output_dim: int,
        input_sample_rate: int = 48000,
    ):
        super().__init__(768, output_dim)
        self.input_sample_rate = input_sample_rate

        # Mert
        from suno_utils.tasks.mert_25 import (
            preload_models as preload_semantic_models,
            encode as semantic_encode,
        )

        _ = preload_semantic_models(
            checkpoint_filepath="s3://suno-data/georg/models/semantic/mert_25.pt",
            centroids_filepath="s3://suno-data/georg/models/semantic/mert_25_2x4k.npy",
            device="cuda",
        )

        self.embedding = torch.nn.Embedding(4096, 768)
        self.encode_fn = semantic_encode

    def forward(
        self,
        cond_dicts: tp.List[dict],
        device: tp.Any = "cuda",
    ):
        """
        If codes are not provided then use MERT to extract quantized embeddings (codebook indices).
        You can directly provide codebook indices, in which case audio is ignored. This is useful for inference.

        Args:
            audios (List[torch.Tensor]): List of audio tensors.
            codes (List[torch.Tensor]): List of codebook indices. Optional.
        """
        audios = []
        codes = []
        for cond_dict in cond_dicts:
            audios.append(cond_dict["audio"])
            codes.append(cond_dict["codes"])

        if None in codes:  # compute based on audio
            with torch.no_grad():
                with torch.cuda.amp.autocast(enabled=False):
                    # audios = audios.mean(dim=1, keepdim=True)  # make mono
                    audios = [audio.mean(dim=0, keepdim=True) for audio in audios]
                    audios = [
                        torchaudio.functional.resample(
                            audio, self.input_sample_rate, 24000
                        )
                        for audio in audios
                    ]
                    latents = self.encode_fn(audios)
                    latents = [
                        torch.from_numpy(latent[:, 0]).long() for latent in latents
                    ]
                    codes = torch.stack(latents).type_as(audios[0]).long()

        else:
            # stack codes into single tensor
            codes = torch.stack(codes, dim=0)

        # diffusion model expects: `[batch, sequence, channels]`.

        return [
            self.embedding(codes),
            torch.ones(codes.shape[0], 1).to(device),
        ]


# this is deprecated
class DACConditioner(Conditioner):
    def __init__(
        self,
        output_dim: int,
        input_sample_rate: int,
        codec_ckpt_path: str,
        codebook_dropout: bool = False,
    ):
        super().__init__(128, output_dim)
        self.output_dim = output_dim
        self.input_sample_rate = input_sample_rate
        self.codec_ckpt_path = codec_ckpt_path
        self.codebook_dropout = codebook_dropout
        self.model = self.load_codec(codec_ckpt_path)

    def load_codec(self, codec_path: str):
        sd = torch.load(codec_path, map_location="cpu")
        model = DAC(**sd["metadata"]["kwargs"])
        model.load_state_dict(sd["state_dict"])
        model.eval()

        for param_name, param in model.named_parameters():
            param.requires_grad = False

        return model

    def forward(
        self,
        cond_dicts: tp.List[tuple],
        device: tp.Any = "cuda",
    ):
        """
        If codes are not provided then use DAC to extract quantized embeddings (codebook indices).
        You can directly provide codebook indices, in which case audio is ignored. This is useful for inference.

        Args:
            audios (List[torch.Tensor]): List of stereo audio tensors.
            codes (List[torch.Tensor]): List of codebook indices. Optional.
        """

        audios = []
        codes = []
        for cond_dict in cond_dicts:
            audios.append(cond_dict["audio"])
            codes.append(cond_dict["codes"])

        if None in codes:  # compute based on audio
            with torch.no_grad():
                with torch.cuda.amp.autocast(enabled=False):
                    audios = torch.stack(audios, dim=0)

                    if self.input_sample_rate != 48000:
                        audios = torchaudio.functional.resample(
                            audios, self.input_sample_rate, 48000
                        )

                    if self.codebook_dropout:
                        n_quantizers = np.random.randint(1, 13)
                    else:
                        n_quantizers = None

                    z, codes, latents, commitment_loss, codebook_loss = (
                        self.model.encode(audios, n_quantizers)
                    )
                    z = z.permute(0, 2, 1)
        else:
            # lookup codes to go back to continuous
            # codes = torch.stack(codes)
            # print(codes.shape)
            # z, _, _ = self.model.quantizer.from_codes(codes)  # b, n, t
            # z = z.permute(0, 2, 1)
            z = torch.stack(codes)

        return [
            self.proj_out(z),
            torch.ones(z.shape[0], 1).to(device),
        ]


class LatentContextConditioner(Conditioner):
    def __init__(self, io_channels: int, output_dim: int):
        super().__init__(output_dim, output_dim)
        self.proj_out = nn.Linear(io_channels, output_dim)
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)

    def forward(
        self,
        latent_context: tp.List[torch.Tensor],
        device: tp.Any = "cuda",
        mode: str = "val",
    ):
        # latent context has shape (io_channels, seq_len)
        latent_context = torch.stack(latent_context, dim=0)
        latent_context = self.proj_out(latent_context.permute(0, 2, 1))
        latent_context = latent_context + self.pos_embedding(latent_context)
        # outputs bs, seq_len, output_dim
        return [latent_context, torch.ones(latent_context.shape[0], 1).to(device)]


class SemanticConditioner(Conditioner):
    def __init__(
        self,
        output_dim: int,
        vocab_size: int = 4000,
        dropout_rate: float = 0.1,
    ):
        super().__init__(output_dim, output_dim)
        self.vocab_size = vocab_size
        self.embedding = torch.nn.Embedding(vocab_size + 1, output_dim)
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)
        self.dropout_rate = dropout_rate

    def forward(
        self,
        semantic_codes: tp.List[torch.Tensor],
        device: tp.Any = "cuda",
        mode: str = "val",
    ):
        # with torch.cuda.amp.autocast(enabled=True):
        semantic_codes = torch.stack(semantic_codes, dim=0)
        if mode == "train" and np.random.random() < self.dropout_rate:
            semantic_codes[:] = self.vocab_size
        semantic_embeds = self.embedding(semantic_codes)
        semantic_embeds = semantic_embeds + self.pos_embedding(semantic_embeds)
        semantic_embeds = self.proj_out(semantic_embeds)

        # diffusion model expects: `[batch, sequence, channels]`.

        return [
            semantic_embeds,
            torch.ones(semantic_embeds.shape[0], semantic_embeds.shape[1]).type_as(
                semantic_embeds
            ),
        ]


class CodecConditioner(Conditioner):
    def __init__(
        self,
        output_dim: int,
        vocab_size: int = 2048,
        dropout_rate: float = 0.5,
        use_partial_dropout: float = True,
    ):
        super().__init__(output_dim, output_dim)
        self.embeddings = torch.nn.ModuleList()
        self.vocab_size = vocab_size
        for _ in range(12):
            self.embeddings.append(torch.nn.Embedding(vocab_size + 1, output_dim))
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)
        self.dropout_rate = dropout_rate
        self.use_partial_dropout = use_partial_dropout

    def forward(
        self,
        codec_codes: tp.List[torch.Tensor],
        device: tp.Any = "cuda",
        mode: str = "val",
    ):
        # codec codes have shape T x 12, one set of indices for each codebook level
        # with torch.cuda.amp.autocast(enabled=True):
        # if training, dropout some codebooks
        keep_n_codebooks = 12
        if mode == "train":
            if random.random() < self.dropout_rate:
                keep_n_codebooks = 0
            elif self.use_partial_dropout:
                keep_n_codebooks = random.randint(1, 12)

        codec_codes = torch.stack(codec_codes, dim=0)  # stack along batch dim
        codec_codes[:, keep_n_codebooks:] = self.vocab_size
        codec_embeds = self.embeddings[0](codec_codes[:, :, 0])
        for codebook_idx in range(1, 12):
            codec_embeds += self.embeddings[codebook_idx](
                codec_codes[:, :, codebook_idx]
            )

        codec_embeds = codec_embeds + self.pos_embedding(codec_embeds)
        codec_embeds = self.proj_out(codec_embeds)

        # diffusion model expects: `[batch, sequence, channels]`.

        return [
            codec_embeds,
            torch.ones(codec_embeds.shape[0], codec_embeds.shape[1]).type_as(
                codec_embeds
            ),
        ]


class SemanticConditionerV2(Conditioner):
    """
    Takes advantage of the k-means centroids instead of using trainable embeddings.
    """

    def __init__(
        self,
        dim: int = 768,
        output_dim: int = 768,
        project_out: bool = False,
        vocab_size: int = 4000,
        dropout_rate: float = 0.1,
        centroid_path: str = "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy",
    ):
        super().__init__(output_dim, output_dim)
        self.vocab_size = vocab_size
        centroids = read_from_s3(centroid_path, read_f=np.load)[0]

        # create a new tensor with an extra row for the dropout token
        centroids = np.concatenate(
            [centroids, np.random.randn(1, centroids.shape[1])], axis=0
        )

        # initialize the embedding with the centroids
        self.embedding = torch.nn.Embedding(centroids.shape[0], centroids.shape[1])
        self.embedding.weight.data.copy_(torch.tensor(centroids, dtype=torch.float32))

        self.proj_out = (
            nn.Linear(dim, output_dim)
            if (dim != output_dim or project_out)
            else nn.Identity()
        )
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)
        self.dropout_rate = dropout_rate

    def forward(
        self,
        semantic_codes: tp.List[torch.Tensor],
        device: tp.Any = "cuda",
        mode: str = "val",
    ):
        # replace semantic codes with the vocab size with probability `dropout_rate`
        # each item in the list is a tensor of shape (batch, sequence)
        for i, semantic_code in enumerate(semantic_codes):
            if mode == "train" and np.random.random() < self.dropout_rate:
                semantic_codes[i] = torch.full_like(semantic_code, self.vocab_size)

        semantic_codes = torch.stack(semantic_codes, dim=0)
        semantic_embeds = self.embedding(semantic_codes)
        semantic_embeds = self.proj_out(semantic_embeds)
        semantic_embeds = semantic_embeds + self.pos_embedding(semantic_embeds)

        # diffusion model expects: `[batch, sequence, channels]`.

        return [
            semantic_embeds,
            torch.ones(semantic_embeds.shape[0], semantic_embeds.shape[1]).type_as(
                semantic_embeds
            ),
        ]


class DiscreteVAEConditioner(Conditioner):
    def __init__(
        self,
        dim: int = 128,
        output_dim: int = 768,
        vocab_size: int = 32768,
        centroid_path: str = "s3://suno-data/christian/vae_100hz_32768.npy",
    ):
        super().__init__(output_dim, output_dim)
        self.vocab_size = vocab_size
        centroids = read_from_s3(centroid_path, read_f=np.load)
        print(centroids.shape)
        self.embedding = torch.nn.Embedding(centroids.shape[0], centroids.shape[1])
        self.embedding.weight.data.copy_(torch.tensor(centroids, dtype=torch.float32))
        self.embedding.weight.requires_grad = False
        self.proj_out = (
            nn.Linear(dim, output_dim) if (dim != output_dim) else nn.Identity()
        )

        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)

    def forward(self, discrete_codes: tp.List[torch.Tensor], device: tp.Any = "cuda"):
        discrete_codes = torch.stack(discrete_codes, dim=0)
        vae_embeds = self.embedding(discrete_codes)
        vae_embeds = self.proj_out(vae_embeds)
        vae_embeds = vae_embeds + self.pos_embedding(vae_embeds)
        # diffusion model expects: `[batch, sequence, channels]`.

        return [
            vae_embeds,
            torch.ones(vae_embeds.shape[0], vae_embeds.shape[1]).type_as(vae_embeds),
        ]


class CodecConditionerV2(Conditioner):
    def __init__(
        self,
        dim: int = 128,
        output_dim: int = 768,
        project_out: bool = False,
        keep_n_codebooks: tp.Optional[int] = None,
        vocab_size: int = 2048,
        dropout_rate: float = 0.5,
        use_partial_dropout: float = True,
        codec_ckpt_path: str = "s3://suno-data/georg/models/codec/dac_2c_25x12.pt",
        codec_input_dim: int = 128,
        codec_n_codebooks: int = 12,
        codec_codebook_size: int = 2048,
        codec_codebook_dim: int = 8,
        codec_quantizer_dropout: float = 0.0,
    ):
        """
        Note: if `keep_n_codebooks` is specified then dropout is not used during training.
        """
        super().__init__(output_dim, output_dim)

        self.codec_input_dim = codec_input_dim
        self.codec_n_codebooks = codec_n_codebooks
        self.codec_codebook_size = codec_codebook_size
        self.codec_codebook_dim = codec_codebook_dim
        self.codec_quantizer_dropout = codec_quantizer_dropout

        self.embeddings = torch.nn.ModuleList()
        self.vocab_size = vocab_size
        for _ in range(12):
            self.embeddings.append(torch.nn.Embedding(vocab_size + 1, output_dim))
        self.proj_out = (
            nn.Linear(dim, output_dim)
            if (dim != output_dim or project_out)
            else nn.Identity()
        )
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)
        self.dropout_rate = dropout_rate
        self.keep_n_codebooks = keep_n_codebooks
        self.use_partial_dropout = use_partial_dropout
        self.rvq = self.load_rvq(codec_ckpt_path)

    def load_rvq(self, codec_path: str):
        sd = read_from_s3(codec_path, read_f=torch.load)
        model = ResidualVectorQuantize(
            input_dim=self.codec_input_dim,
            n_codebooks=self.codec_n_codebooks,
            codebook_size=self.codec_codebook_size,
            codebook_dim=self.codec_codebook_dim,
            quantizer_dropout=self.codec_quantizer_dropout,
        )

        model.load_state_dict(
            {k[10:]: v for k, v in sd["state_dict"].items() if k.startswith("quantize")}
        )
        model.eval()

        for param_name, param in model.named_parameters():
            param.requires_grad = False

        return model

    def decode_vq(self, codes, n_quantizers):
        z_q = 0
        for i, quantizer in enumerate(self.rvq.quantizers[:n_quantizers]):
            _z_q = quantizer.embed_code(codes[:, :, i]).transpose(1, 2)
            _z_q = quantizer.out_proj(_z_q)
            z_q += _z_q.transpose(1, 2)
        return z_q

    def forward(
        self,
        codec_codes: tp.List[torch.Tensor],
        device: tp.Any = "cuda",
        keep_n_codebooks: int = None,
        force_n_codebooks: int = None,
        mode: str = "val",
    ):
        # codec codes have shape T x 12, one set of indices for each codebook level
        # with torch.cuda.amp.autocast(enabled=True):
        if force_n_codebooks:
            # ignore partial dropout and force a certain number of codebooks
            keep_n_codebooks = force_n_codebooks
        # if training, dropout some codebooks
        elif mode == "train":
            if self.use_partial_dropout and keep_n_codebooks is None:
                keep_n_codebooks = random.randint(1, 12)
        else:
            keep_n_codebooks = keep_n_codebooks

        # stack along batch dim (batch, time, 12)
        codec_codes = torch.stack(codec_codes, dim=0)
        z_q = self.decode_vq(codec_codes, keep_n_codebooks)
        codec_embeds = self.proj_out(z_q)
        codec_embeds = codec_embeds + self.pos_embedding(codec_embeds)

        # diffusion model expects: `[batch, sequence, channels]`.

        return [
            codec_embeds,
            torch.ones(codec_embeds.shape[0], codec_embeds.shape[1]).type_as(
                codec_embeds
            ),
        ]


class VAEConditioner(Conditioner):
    def __init__(
        self,
        dim: int = 128,
        output_dim: int = 768,
        project_out: bool = False,
        dropout_rate: float = 0.5,
        use_partial_dropout: float = True,
    ):
        """
        Note: if `keep_n_codebooks` is specified then dropout is not used during training.
        """
        super().__init__(output_dim, output_dim)

        self.proj_out = (
            nn.Linear(dim, output_dim)
            if (dim != output_dim or project_out)
            else nn.Identity()
        )
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)
        self.dropout_rate = dropout_rate
        self.use_partial_dropout = use_partial_dropout

    def forward(
        self,
        latents: tp.List[torch.Tensor],
        device: tp.Any = "cuda",
    ):
        # stack along batch dim (batch, time, embed_dim)
        latents = torch.stack(latents, dim=0)
        print(latents.shape)

        latents = self.proj_out(latents)
        latents = latents + self.pos_embedding(latents)

        # diffusion model expects: `[batch, sequence, channels]`.
        return [
            latents,
            torch.ones(latents.shape[0], latents.shape[1]).type_as(latents),
        ]


class PhonemeConditionerV2(Conditioner):
    def __init__(
        self,
        output_dim: int = 768,
        dropout_rate: float = 0.1,
        max_length: int = 2560,
    ):
        super().__init__(output_dim, output_dim)
        self.dropout_rate = dropout_rate
        self.tokenizer = {
            " ": 0,
            "(": 1,
            ")": 2,
            ".": 3,
            "1": 4,
            "a": 5,
            "b": 6,
            "c": 7,
            "d": 8,
            "e": 9,
            "f": 10,
            "g": 11,
            "h": 12,
            "i": 13,
            "j": 14,
            "k": 15,
            "l": 16,
            "m": 17,
            "n": 18,
            "o": 19,
            "p": 20,
            "q": 21,
            "r": 22,
            "s": 23,
            "t": 24,
            "u": 25,
            "v": 26,
            "w": 27,
            "x": 28,
            "y": 29,
            "z": 30,
            "æ": 31,
            "ç": 32,
            "ð": 33,
            "ŋ": 34,
            "ɐ": 35,
            "ɑ": 36,
            "ɒ": 37,
            "ɔ": 38,
            "ɕ": 39,
            "ɖ": 40,
            "ə": 41,
            "ɚ": 42,
            "ɛ": 43,
            "ɜ": 44,
            "ɟ": 45,
            "ɡ": 46,
            "ɣ": 47,
            "ɨ": 48,
            "ɪ": 49,
            "ɫ": 50,
            "ɬ": 51,
            "ɭ": 52,
            "ɯ": 53,
            "ɲ": 54,
            "ɳ": 55,
            "ɹ": 56,
            "ɻ": 57,
            "ɾ": 58,
            "ʀ": 59,
            "ʁ": 60,
            "ʂ": 61,
            "ʃ": 62,
            "ʈ": 63,
            "ʉ": 64,
            "ʊ": 65,
            "ʋ": 66,
            "ʌ": 67,
            "ʐ": 68,
            "ʑ": 69,
            "ʒ": 70,
            "ʔ": 71,
            "ʰ": 72,
            "ʲ": 73,
            "ː": 74,
            "̃": 75,
            "̩": 76,
            "θ": 77,
            "χ": 78,
            "ᵐ": 79,
            "ᵑ": 80,
            "ᵻ": 81,
            "ⁿ": 82,
        }
        self.vocab_size = len(self.tokenizer)
        self.max_length = max_length
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)
        self.embedding = nn.Embedding(self.vocab_size + 1, output_dim)

    def forward(
        self, phonemes: tp.List[str], device: tp.Any = "cuda", mode: str = "val"
    ):
        # loop through each phoneme sequence and tokenize
        tokenized_phonemes = torch.ones(len(phonemes), self.max_length)
        tokenized_phonemes = tokenized_phonemes * self.vocab_size  # fill with pad token
        tokenized_phonemes = tokenized_phonemes.long().to(device)  # move to device

        for i, phoneme in enumerate(phonemes):
            if mode == "train" and random.random() < self.dropout_rate:
                continue
            else:
                tokenized_phoneme_seq = [
                    self.tokenizer.get(p, len(self.tokenizer)) for p in phoneme
                ]
                # truncate to max length
                tokenized_phoneme_seq = tokenized_phoneme_seq[: self.max_length]
                tokenized_phoneme_seq = torch.tensor(tokenized_phoneme_seq).to(device)
                tokenized_phonemes[i, : len(tokenized_phoneme_seq)] = (
                    tokenized_phoneme_seq
                )

        phoneme_embeds = self.embedding(tokenized_phonemes)
        phoneme_embeds = phoneme_embeds + self.pos_embedding(phoneme_embeds)
        return [
            phoneme_embeds,
            torch.ones(phoneme_embeds.shape[0], phoneme_embeds.shape[1]).type_as(
                phoneme_embeds
            ),
        ]


def _space_repl(m):
    s = m.group()
    n_newline = s.count("\n")
    if n_newline >= 2:
        return "\n\n"
    elif n_newline == 1:
        return "\n"
    return " "


def _simplify_whitespace(text, retain_newlines=True):
    """simplify while respecting up to 2 newlines"""
    if retain_newlines:
        text = re.sub(r"\s+", _space_repl, text).strip()
    else:
        text = re.sub(r"\s+", " ", text).strip()
    return text


CASE_AUGMENT_FUNCS = [
    str.upper,
    str.lower,
    str.capitalize,
    str.title,
]


def _augment_tag(s):
    # case augment
    if random.random() >= 0.8:
        s = random.choice(CASE_AUGMENT_FUNCS)(s)
    # other misc formatting
    if random.random() >= 0.5:
        s = s.replace("-", " ").strip()
    return s


def _clean_tag(s, retain_newlines=False):
    s = s.replace("[", " ").replace("]", " ")
    return _simplify_whitespace(s, retain_newlines=retain_newlines)


MAX_TAG_LEN = 256
MAX_TOT_TAGS_LEN = 512


class TagsAndLyricsConditioner(Conditioner):
    def __init__(
        self,
        tokenizer_file: str = "/app/suno/data/chirp_v4/multi/tokenizer_60k.json",
        output_dim: int = 768,
        max_length: str = 2560,
        dropout_rate: float = 0.1,
        use_legacy_format: bool = False,
    ):
        super().__init__(768, output_dim)
        self.max_length = max_length
        self.dropout_rate = dropout_rate
        self.use_legacy_format = use_legacy_format

        # load model and tokenizer
        from transformers import PreTrainedTokenizerFast
        from tokenizers import AddedToken

        self.tokenizer = PreTrainedTokenizerFast(
            tokenizer_file=tokenizer_file,
            unk_token="[UNK]",
            pad_token="[PAD]",
        )

        if use_legacy_format:
            special_tokens = {
                "additional_special_tokens": [
                    AddedToken("<tags>"),
                    AddedToken("</tags>"),
                    AddedToken("<lyrics>"),
                    AddedToken("</lyrics>"),
                    AddedToken("\n"),
                ]
            }
        else:
            special_tokens = {
                "additional_special_tokens": [
                    AddedToken("\n"),
                ]
            }
        self.tokenizer.add_special_tokens(special_tokens)
        n_vocab = self.tokenizer.vocab_size + len(
            special_tokens["additional_special_tokens"]
        )
        self.n_vocab = n_vocab
        # create learnable embedding
        self.embedding = torch.nn.Embedding(n_vocab, output_dim)
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)

    def forward(
        self,
        tags_and_lyrics_text: tp.List[str],
        device: tp.Any = "cuda",
        mode: str = "val",
    ):
        concatenated_texts = []
        # Loop through each pair of tags and lyrics
        for tags, lyrics in tags_and_lyrics_text:
            if self.use_legacy_format:
                # Concatenate with special tokens
                concatenated_text = f"<tags>{tags}</tags><lyrics>{lyrics}</lyrics>"
                concatenated_texts.append(concatenated_text)
            else:
                # Augment tags
                if mode == "train" and (random.random() < self.dropout_rate):
                    concatenated_texts.append("")
                elif mode == "train":
                    # augment tags
                    if random.random() < 0.5:
                        random.shuffle(tags)
                        if len(tags) > 0:
                            tags = tags[: random.randint(1, len(tags))]
                        tags = [_augment_tag(tag) for tag in tags]
                    tags = [
                        clean_tag
                        for tag in tags
                        if len(clean_tag := _clean_tag(tag, retain_newlines=False)) > 0
                    ]
                    merge_char = random.choice([", ", " ", "; ", ",", ";"])
                    tags_str = f"{merge_char.join([tag[:MAX_TAG_LEN] for tag in tags])[:MAX_TOT_TAGS_LEN]}"
                    lyrics = _simplify_whitespace(lyrics, retain_newlines=True)
                    if random.random() < 0.05:
                        lyrics = lyrics.lower()
                    if random.random() < 0.05:
                        lyrics = re.sub(r"\n+", " ", lyrics)
                    merge_char = random.choice(["\n\n", "\n", " "])
                    concatenated_texts.append(
                        (f"[{tags_str}]" + merge_char + lyrics).strip()
                    )
                else:
                    tags = [
                        clean_tag
                        for tag in tags
                        if len(clean_tag := _clean_tag(tag, retain_newlines=False)) > 0
                    ]
                    tags_str = f"{', '.join([tag[:MAX_TAG_LEN] for tag in tags])[:MAX_TOT_TAGS_LEN]}"
                    lyrics = _simplify_whitespace(lyrics, retain_newlines=True)
                    concatenated_texts.append(
                        (f"[{tags_str}]" + "\n\n" + lyrics).strip()
                    )

        # Tokenize the entire sequence
        inputs = self.tokenizer(
            concatenated_texts,
            truncation=True,
            return_tensors="pt",
            padding="max_length",
            max_length=self.max_length,
        )

        # move input to device
        inputs = {k: v.to(device) for k, v in inputs.items()}

        # get embeddings
        embeddings = self.embedding(inputs["input_ids"])
        embeddings = embeddings + self.pos_embedding(embeddings)

        # diffusion model expects: `[batch, sequence, channels]`.
        return [
            embeddings,
            torch.ones(embeddings.shape[0], embeddings.shape[1]).type_as(embeddings),
        ]


class CLAPTextConditioner(Conditioner):
    def __init__(
        self,
        output_dim: int,
        clap_ckpt_path,
        use_text_features=False,
        feature_layer_ix: int = -1,
        audio_model_type="HTSAT-base",
        enable_fusion=True,
        project_out: bool = False,
        finetune: bool = False,
    ):
        super().__init__(
            768 if use_text_features else 512, output_dim, project_out=project_out
        )

        self.use_text_features = use_text_features
        self.feature_layer_ix = feature_layer_ix
        self.finetune = finetune

        # Suppress logging from transformers
        previous_level = logging.root.manager.disable
        logging.disable(logging.ERROR)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            try:
                import laion_clap
                from laion_clap.clap_module.factory import (
                    load_state_dict as clap_load_state_dict,
                )

                model = laion_clap.CLAP_Module(
                    enable_fusion=enable_fusion, amodel=audio_model_type, device="cpu"
                )

                if self.finetune:
                    self.model = model
                else:
                    self.__dict__["model"] = model

                state_dict = clap_load_state_dict(clap_ckpt_path)
                self.model.model.load_state_dict(state_dict, strict=False)

                if self.finetune:
                    self.model.model.text_branch.requires_grad_(True)
                    self.model.model.text_branch.train()
                else:
                    self.model.model.text_branch.requires_grad_(False)
                    self.model.model.text_branch.eval()

            finally:
                logging.disable(previous_level)

        del self.model.model.audio_branch

        gc.collect()
        torch.cuda.empty_cache()

    def get_clap_features(self, prompts, layer_ix=-2, device: tp.Any = "cuda"):
        prompt_tokens = self.model.tokenizer(prompts)
        attention_mask = prompt_tokens["attention_mask"].to(
            device=device, non_blocking=True
        )
        prompt_features = self.model.model.text_branch(
            input_ids=prompt_tokens["input_ids"].to(device=device, non_blocking=True),
            attention_mask=attention_mask,
            output_hidden_states=True,
        )["hidden_states"][layer_ix]

        return prompt_features, attention_mask

    def forward(self, texts: tp.List[str], device: tp.Any = "cuda") -> tp.Any:
        self.model.to(device)

        if self.use_text_features:
            if len(texts) == 1:
                text_features, text_attention_mask = self.get_clap_features(
                    [texts[0], ""], layer_ix=self.feature_layer_ix, device=device
                )
                text_features = text_features[:1, ...]
                text_attention_mask = text_attention_mask[:1, ...]
            else:
                text_features, text_attention_mask = self.get_clap_features(
                    texts, layer_ix=self.feature_layer_ix, device=device
                )
            return [self.proj_out(text_features), text_attention_mask]

        # Fix for CLAP bug when only one text is passed
        if len(texts) == 1:
            text_embedding = self.model.get_text_embedding(
                [texts[0], ""], use_tensor=True
            )[:1, ...]
        else:
            text_embedding = self.model.get_text_embedding(texts, use_tensor=True)

        text_embedding = text_embedding.unsqueeze(1).to(device)

        return [
            self.proj_out(text_embedding),
            torch.ones(text_embedding.shape[0], 1).to(device),
        ]


class CLAPAudioConditioner(Conditioner):
    def __init__(
        self,
        output_dim: int,
        clap_ckpt_path,
        audio_model_type="HTSAT-base",
        enable_fusion=True,
        project_out: bool = False,
    ):
        super().__init__(512, output_dim, project_out=project_out)

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Suppress logging from transformers
        previous_level = logging.root.manager.disable
        logging.disable(logging.ERROR)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            try:
                import laion_clap
                from laion_clap.clap_module.factory import (
                    load_state_dict as clap_load_state_dict,
                )

                model = laion_clap.CLAP_Module(
                    enable_fusion=enable_fusion, amodel=audio_model_type, device="cpu"
                )

                if self.finetune:
                    self.model = model
                else:
                    self.__dict__["model"] = model

                state_dict = clap_load_state_dict(clap_ckpt_path)
                self.model.model.load_state_dict(state_dict, strict=False)

                if self.finetune:
                    self.model.model.audio_branch.requires_grad_(True)
                    self.model.model.audio_branch.train()
                else:
                    self.model.model.audio_branch.requires_grad_(False)
                    self.model.model.audio_branch.eval()

            finally:
                logging.disable(previous_level)

        del self.model.model.text_branch

        gc.collect()
        torch.cuda.empty_cache()

    def forward(
        self,
        audios: tp.Union[torch.Tensor, tp.List[torch.Tensor], tp.Tuple[torch.Tensor]],
        device: tp.Any = "cuda",
    ) -> tp.Any:
        self.model.to(device)

        if isinstance(audios, list) or isinstance(audios, tuple):
            audios = torch.cat(audios, dim=0)

        # Convert to mono
        mono_audios = audios.mean(dim=1)

        with torch.cuda.amp.autocast(enabled=False):
            audio_embedding = self.model.get_audio_embedding_from_data(
                mono_audios.float(), use_tensor=True
            )

        audio_embedding = audio_embedding.unsqueeze(1).to(device)

        return [
            self.proj_out(audio_embedding),
            torch.ones(audio_embedding.shape[0], 1).to(device),
        ]


class T5Conditioner(Conditioner):
    T5_MODELS = [
        "t5-small",
        "t5-base",
        "t5-large",
        "t5-3b",
        "t5-11b",
        "google/flan-t5-small",
        "google/flan-t5-base",
        "google/flan-t5-large",
        "google/flan-t5-xl",
        "google/flan-t5-xxl",
    ]

    T5_MODEL_DIMS = {
        "t5-small": 512,
        "t5-base": 768,
        "t5-large": 1024,
        "t5-3b": 1024,
        "t5-11b": 1024,
        "t5-xl": 2048,
        "t5-xxl": 4096,
        "google/flan-t5-small": 512,
        "google/flan-t5-base": 768,
        "google/flan-t5-large": 1024,
        "google/flan-t5-3b": 1024,
        "google/flan-t5-11b": 1024,
        "google/flan-t5-xl": 2048,
        "google/flan-t5-xxl": 4096,
    }

    def __init__(
        self,
        output_dim: int,
        t5_model_name: str = "t5-base",
        max_length: str = 128,
        enable_grad: bool = False,
        project_out: bool = False,
    ):
        assert (
            t5_model_name in self.T5_MODELS
        ), f"Unknown T5 model name: {t5_model_name}"
        super().__init__(
            self.T5_MODEL_DIMS[t5_model_name], output_dim, project_out=project_out
        )

        from transformers import T5EncoderModel, AutoTokenizer

        self.max_length = max_length
        self.enable_grad = enable_grad

        # Suppress logging from transformers
        previous_level = logging.root.manager.disable
        logging.disable(logging.ERROR)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            try:
                # self.tokenizer = T5Tokenizer.from_pretrained(t5_model_name, model_max_length = max_length)
                # model = T5EncoderModel.from_pretrained(t5_model_name, max_length=max_length).train(enable_grad).requires_grad_(enable_grad)
                self.tokenizer = AutoTokenizer.from_pretrained(t5_model_name)
                model = (
                    T5EncoderModel.from_pretrained(t5_model_name)
                    .train(enable_grad)
                    .requires_grad_(enable_grad)
                    .to(torch.float16)
                )
            finally:
                logging.disable(previous_level)

        if self.enable_grad:
            self.model = model
        else:
            self.__dict__["model"] = model

    def forward(
        self, texts: tp.List[str], device: tp.Union[torch.device, str]
    ) -> tp.Tuple[torch.Tensor, torch.Tensor]:
        self.model.to(device)
        self.proj_out.to(device)

        encoded = self.tokenizer(
            texts,
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt",
        )

        input_ids = encoded["input_ids"].to(device)
        attention_mask = encoded["attention_mask"].to(device).to(torch.bool)

        self.model.eval()

        with torch.cuda.amp.autocast(dtype=torch.float16) and torch.set_grad_enabled(
            self.enable_grad
        ):
            embeddings = self.model(input_ids=input_ids, attention_mask=attention_mask)[
                "last_hidden_state"
            ]

        embeddings = self.proj_out(embeddings.float())

        embeddings = embeddings * attention_mask.unsqueeze(-1).float()

        return embeddings, attention_mask


class PhonemeConditioner(Conditioner):
    """
    A conditioner that turns text into phonemes and embeds them using a lookup table
    Only works for English text

    Args:
        output_dim: the dimension of the output embeddings
        max_length: the maximum number of phonemes to embed
        project_out: whether to add another linear projection to the output embeddings
    """

    def __init__(
        self,
        output_dim: int,
        max_length: int = 1024,
        project_out: bool = False,
    ):
        super().__init__(output_dim, output_dim, project_out=project_out)

        from g2p_en import G2p

        self.max_length = max_length

        self.g2p = G2p()

        # Reserving 0 for padding, 1 for ignored
        self.phoneme_embedder = nn.Embedding(len(self.g2p.phonemes) + 2, output_dim)

    def forward(
        self, texts: tp.List[str], device: tp.Union[torch.device, str]
    ) -> tp.Tuple[torch.Tensor, torch.Tensor]:
        self.phoneme_embedder.to(device)
        self.proj_out.to(device)

        batch_phonemes = [
            self.g2p(text) for text in texts
        ]  # shape [batch_size, length]

        phoneme_ignore = [" ", *string.punctuation]

        # Remove ignored phonemes and cut to max length
        batch_phonemes = [
            [p if p not in phoneme_ignore else "_" for p in phonemes]
            for phonemes in batch_phonemes
        ]

        # Convert to ids
        phoneme_ids = [
            [self.g2p.p2idx[p] + 2 if p in self.g2p.p2idx else 1 for p in phonemes]
            for phonemes in batch_phonemes
        ]

        # Pad to match longest and make a mask tensor for the padding
        longest = max([len(ids) for ids in phoneme_ids])
        phoneme_ids = [ids + [0] * (longest - len(ids)) for ids in phoneme_ids]

        phoneme_ids = torch.tensor(phoneme_ids).to(device)

        # Convert to embeddings
        phoneme_embeds = self.phoneme_embedder(phoneme_ids)

        phoneme_embeds = self.proj_out(phoneme_embeds)

        return phoneme_embeds, torch.ones(
            phoneme_embeds.shape[0], phoneme_embeds.shape[1]
        ).to(device)


class TokenizerLUTConditioner(Conditioner):
    """
    A conditioner that embeds text using a lookup table on a pretrained tokenizer's vocabulary

    Args:
        tokenizer_name: the name of the tokenizer from the Hugging Face transformers library
        output_dim: the dimension of the output embeddings
        max_length: the maximum length of the text to embed
        project_out: whether to add another linear projection to the output embeddings
    """

    def __init__(
        self,
        tokenizer_name: str,  # Name of a tokenizer from the Hugging Face transformers library
        output_dim: int,
        max_length: int = 1024,
        project_out: bool = False,
    ):
        super().__init__(output_dim, output_dim, project_out=project_out)

        from transformers import AutoTokenizer

        # Suppress logging from transformers
        previous_level = logging.root.manager.disable
        logging.disable(logging.ERROR)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            try:
                self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
            finally:
                logging.disable(previous_level)

        self.max_length = max_length

        self.token_embedder = nn.Embedding(len(self.tokenizer), output_dim)

    def forward(
        self, texts: tp.List[str], device: tp.Union[torch.device, str]
    ) -> tp.Tuple[torch.Tensor, torch.Tensor]:
        self.proj_out.to(device)

        encoded = self.tokenizer(
            texts,
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt",
        )

        input_ids = encoded["input_ids"].to(device)
        attention_mask = encoded["attention_mask"].to(device).to(torch.bool)

        embeddings = self.token_embedder(input_ids)

        embeddings = self.proj_out(embeddings)

        embeddings = embeddings * attention_mask.unsqueeze(-1).float()

        return embeddings, attention_mask


class PretransformConditioner(Conditioner):
    """
    A conditioner that uses a pretransform's encoder for conditioning

    Args:
        pretransform: an instantiated pretransform to use for conditioning
        output_dim: the dimension of the output embeddings
    """

    def __init__(self, pretransform: Pretransform, output_dim: int):
        super().__init__(pretransform.encoded_channels, output_dim)

        self.pretransform = pretransform

    def forward(
        self,
        audio: tp.Union[torch.Tensor, tp.List[torch.Tensor], tp.Tuple[torch.Tensor]],
        device: tp.Union[torch.device, str],
    ) -> tp.Tuple[torch.Tensor, torch.Tensor]:
        self.pretransform.to(device)
        self.proj_out.to(device)

        if isinstance(audio, list) or isinstance(audio, tuple):
            audio = torch.cat(audio, dim=0)

        # Convert audio to pretransform input channels
        audio = set_audio_channels(audio, self.pretransform.io_channels)

        latents = self.pretransform.encode(audio)

        latents = self.proj_out(latents)

        return [
            latents,
            torch.ones(latents.shape[0], latents.shape[2]).to(latents.device),
        ]


class MultiConditioner(nn.Module):
    """
    A module that applies multiple conditioners to an input dictionary based on the keys

    Args:
        conditioners: a dictionary of conditioners with keys corresponding to the keys of the conditioning input dictionary (e.g. "prompt")
        default_keys: a dictionary of default keys to use if the key is not in the input dictionary (e.g. {"prompt_t5": "prompt"})
    """

    def __init__(
        self,
        conditioners: tp.Dict[str, Conditioner],
        default_keys: tp.Dict[str, str] = {},
    ):
        super().__init__()

        self.conditioners = nn.ModuleDict(conditioners)
        self.default_keys = default_keys

    def forward(
        self,
        batch_metadata: tp.List[tp.Dict[str, tp.Any]],
        device: tp.Union[torch.device, str],
        mode: str = "val",
    ) -> tp.Dict[str, tp.Any]:
        output = {}

        for key, conditioner in self.conditioners.items():
            condition_key = key

            conditioner_inputs = []

            for x in batch_metadata:
                if condition_key not in x:
                    if condition_key in self.default_keys:
                        condition_key = self.default_keys[condition_key]
                    else:
                        raise ValueError(
                            f"Conditioner key {condition_key} not found in batch metadata"
                        )

                # Unwrap the condition info if it's a single-element list or tuple, this is to support collation functions that wrap everything in a list
                if (
                    isinstance(x[condition_key], list)
                    or isinstance(x[condition_key], tuple)
                    and len(x[condition_key]) == 1
                ):
                    conditioner_inputs.append(x[condition_key])
                else:
                    conditioner_inputs.append(x[condition_key])

            output[key] = conditioner(conditioner_inputs, device, mode)

        return output


def create_multi_conditioner_from_conditioning_config(
    config: tp.Dict[str, tp.Any],
) -> MultiConditioner:
    """
    Create a MultiConditioner from a conditioning config dictionary

    Args:
        config: the conditioning config dictionary
        device: the device to put the conditioners on
    """
    conditioners = {}
    cond_dim = config["cond_dim"]

    default_keys = config.get("default_keys", {})

    for conditioner_info in config["configs"]:
        id = conditioner_info["id"]

        conditioner_type = conditioner_info["type"]

        conditioner_config = {"output_dim": cond_dim}

        conditioner_config.update(conditioner_info["config"])

        if conditioner_type == "t5":
            conditioners[id] = T5Conditioner(**conditioner_config)
        elif conditioner_type == "clap_text":
            conditioners[id] = CLAPTextConditioner(**conditioner_config)
        elif conditioner_type == "codec":
            conditioners[id] = CodecConditioner(**conditioner_config)
        elif conditioner_type == "mert":
            conditioners[id] = MERTConditioner(**conditioner_config)
        elif conditioner_type == "semantic":
            conditioners[id] = SemanticConditioner(**conditioner_config)
        elif conditioner_type == "semantic_v2":
            conditioners[id] = SemanticConditionerV2(**conditioner_config)
        elif conditioner_type == "codec_v2":
            conditioners[id] = CodecConditionerV2(**conditioner_config)
        elif conditioner_type == "vae":
            conditioners[id] = VAEConditioner(**conditioner_config)
        elif conditioner_type == "discrete_vae":
            conditioners[id] = DiscreteVAEConditioner(**conditioner_config)
        elif conditioner_type == "tags_and_lyrics":
            conditioners[id] = TagsAndLyricsConditioner(**conditioner_config)
        elif conditioner_type == "phoneme_v2":
            conditioners[id] = PhonemeConditionerV2(**conditioner_config)
        elif conditioner_type == "clap_audio":
            conditioners[id] = CLAPAudioConditioner(**conditioner_config)
        elif conditioner_type == "int":
            conditioners[id] = IntConditioner(**conditioner_config)
        elif conditioner_type == "number":
            conditioners[id] = NumberConditioner(**conditioner_config)
        elif conditioner_type == "phoneme":
            conditioners[id] = PhonemeConditioner(**conditioner_config)
        elif conditioner_type == "lut":
            conditioners[id] = TokenizerLUTConditioner(**conditioner_config)
        elif conditioner_type == "latent_context":
            conditioners[id] = LatentContextConditioner(**conditioner_config)
        elif conditioner_type == "pretransform":
            sample_rate = conditioner_config.pop("sample_rate", None)
            assert (
                sample_rate is not None
            ), "Sample rate must be specified for pretransform conditioners"

            pretransform = create_pretransform_from_config(
                conditioner_config.pop("pretransform_config"), sample_rate=sample_rate
            )

            if conditioner_config.get("pretransform_ckpt_path", None) is not None:
                pretransform.load_state_dict(
                    load_ckpt_state_dict(
                        conditioner_config.pop("pretransform_ckpt_path")
                    )
                )

            conditioners[id] = PretransformConditioner(
                pretransform, **conditioner_config
            )
        else:
            raise ValueError(f"Unknown conditioner type: {conditioner_type}")

    return MultiConditioner(conditioners, default_keys=default_keys)