import librosa
import numpy as np
import torch
from tqdm import tqdm

from tools import VAE_out_put_to_spc, rms_normalize, nnData2Audio
from model.DiffSynthSampler import DiffSynthSampler

def sample_pipeline(device, uNet, VAE, MMM, CLAP_tokenizer,
                    positive_prompts, negative_prompts, batchsize, sample_steps, CFG, seed=None, duration=3.0,
                    freq_resolution=512, time_resolution=256, channels=4, VAE_scale=4, timesteps=1000, noise_strategy="repeat", sampler="ddim", return_latent=True):

    height = int(freq_resolution/VAE_scale)
    width = int(time_resolution/VAE_scale)
    VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder

    text2sound_embedding = \
        MMM.get_text_features(**CLAP_tokenizer([positive_prompts], padding=True, return_tensors="pt"))[0].to(device)
    negative_condition = \
        MMM.get_text_features(**CLAP_tokenizer([negative_prompts], padding=True, return_tensors="pt"))[0].to(device)

    mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy, mute=True)
    mySampler.activate_classifier_free_guidance(CFG, negative_condition)

    mySampler.respace(list(np.linspace(0, timesteps - 1, sample_steps, dtype=np.int32)))

    condition = text2sound_embedding.repeat(batchsize, 1)

    latent_representations, initial_noise = \
    mySampler.sample(model=uNet, shape=(batchsize, channels, height, width), seed=seed,
                      return_tensor=True, condition=condition, sampler=sampler)

    latent_representations = latent_representations[-1]

    quantized_latent_representations, _, (_, _, _) = VAE_quantizer(latent_representations)

    if return_latent:
        return quantized_latent_representations.detach()
    reconstruction_batch = VAE_decoder(quantized_latent_representations).to("cpu").detach().numpy()
    time_resolution = int(time_resolution * ((duration+1) / 4))

    rec_signals = nnData2Audio(reconstruction_batch, resolution=(freq_resolution, time_resolution))
    rec_signals = [rms_normalize(rec_signal) for rec_signal in rec_signals]

    return quantized_latent_representations.detach(), reconstruction_batch, rec_signals

def sample_pipeline_GAN(device, gan_generator, VAE, MMM, CLAP_tokenizer,
                    positive_prompts, negative_prompts, batchsize, sample_steps, CFG, seed=None, duration=3.0,
                    freq_resolution=512, time_resolution=256, channels=4, VAE_scale=4, timesteps=1000, noise_strategy="repeat", sampler="ddim", return_latent=True):

    height = int(freq_resolution/VAE_scale)
    width = int(time_resolution/VAE_scale)
    VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder

    text2sound_embedding = \
        MMM.get_text_features(**CLAP_tokenizer([positive_prompts], padding=True, return_tensors="pt"))[0].to(device)

    condition = text2sound_embedding.repeat(batchsize, 1)

    noise = torch.randn(batchsize, channels, height, width).to(device)
    latent_representations = gan_generator(noise, condition)

    quantized_latent_representations, _, (_, _, _) = VAE_quantizer(latent_representations)

    if return_latent:
        return quantized_latent_representations.detach()
    reconstruction_batch = VAE_decoder(quantized_latent_representations).to("cpu").detach().numpy()
    time_resolution = int(time_resolution * ((duration+1) / 4))

    rec_signals = nnData2Audio(reconstruction_batch, resolution=(freq_resolution, time_resolution))
    rec_signals = [rms_normalize(rec_signal) for rec_signal in rec_signals]

    return quantized_latent_representations.detach(), reconstruction_batch, rec_signals

def inpaint_pipeline(device, uNet, VAE, MMM, CLAP_tokenizer, use_dynamic_mask, noising_strength, guidance,
                    positive_prompts, negative_prompts, batchsize, sample_steps, CFG, seed=None, duration=3.0, mask_flexivity=0.99,
                    freq_resolution=512, time_resolution=256, channels=4, VAE_scale=4, timesteps=1000, noise_strategy="repeat", sampler="ddim", return_latent=True):

    height = int(freq_resolution/VAE_scale)
    width = int(time_resolution * ((duration + 1) / 4) / VAE_scale)
    VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder


    text2sound_embedding = \
        MMM.get_text_features(**CLAP_tokenizer([positive_prompts], padding=True, return_tensors="pt"))[0]
    negative_condition = \
        MMM.get_text_features(**CLAP_tokenizer([negative_prompts], padding=True, return_tensors="pt"))[0]


    mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy, mute=True)
    mySampler.activate_classifier_free_guidance(CFG, negative_condition)
    mySampler.respace(list(np.linspace(0, timesteps - 1, sample_steps, dtype=np.int32)))

    condition = text2sound_embedding.repeat(batchsize, 1)
    guidance = guidance.repeat(batchsize, 1, 1, 1).to(device)

    # mask = 1, freeze
    latent_mask = torch.zeros((batchsize, 1, height, width), dtype=torch.float32).to(device)
    latent_mask[:, :, :, -int(time_resolution * (1 / 4) / VAE_scale):] = 1.0

    latent_representations, initial_noise = \
        mySampler.inpaint_sample(model=uNet, shape=(batchsize, channels, height, width),
                                  noising_strength=noising_strength,
                                  guide_img=guidance, mask=latent_mask, return_tensor=True,
                                  condition=condition, sampler=sampler,
                                  use_dynamic_mask=use_dynamic_mask,
                                  end_noise_level_ratio=0.0,
                                  mask_flexivity=mask_flexivity)

    latent_representations = latent_representations[-1]

    quantized_latent_representations, _, (_, _, _) = VAE_quantizer(latent_representations)

    if return_latent:
        return quantized_latent_representations.detach()
    reconstruction_batch = VAE_decoder(quantized_latent_representations).to("cpu").detach().numpy()
    time_resolution = int(time_resolution * ((duration+1) / 4))

    rec_signals = nnData2Audio(reconstruction_batch, resolution=(freq_resolution, time_resolution))
    rec_signals = [rms_normalize(rec_signal) for rec_signal in rec_signals]

    return quantized_latent_representations.detach(), reconstruction_batch, rec_signals


def generate_audios_with_diffuSynth_sample(device, uNet, VAE, MMM, CLAP_tokenizer, num_batches, positive_prompts, negative_prompts="", CFG=6, sample_steps=10):
    diffuSynth_signals = []
    for _ in tqdm(range(num_batches)):
        _, _, signals = sample_pipeline(device, uNet, VAE, MMM, CLAP_tokenizer,
                                        positive_prompts=positive_prompts, negative_prompts=negative_prompts,
                                                      batchsize=16, sample_steps=sample_steps, CFG=CFG, seed=None, return_latent=False)
        diffuSynth_signals.extend(signals)
    return np.array(diffuSynth_signals)


def generate_audios_with_diffuSynth_inpaint(device, uNet, VAE, MMM, CLAP_tokenizer, num_batches, guidance, duration, use_dynamic_mask, noising_strength, positive_prompts, negative_prompts="", CFG=6, sample_steps=10):

    diffuSynth_signals = []
    for _ in tqdm(range(num_batches)):
        _, _, signals = inpaint_pipeline(device, uNet, VAE, MMM, CLAP_tokenizer,
                                         use_dynamic_mask=use_dynamic_mask, noising_strength=noising_strength, guidance=guidance,
                    positive_prompts=positive_prompts, negative_prompts=negative_prompts, batchsize=16, sample_steps=sample_steps, CFG=CFG, seed=None, duration=duration, mask_flexivity=0.999,
                    return_latent=False)
        diffuSynth_signals.extend(signals)
    return np.array(diffuSynth_signals)