import os
import json
import glob
import torch
import torchaudio
from stable_audio_tools.inference.generation import (
    upsample_diffusion,
    upsample_diffusion_from_codes,
)
from stable_audio_tools.interface.gradio import load_model

from stable_audio_tools.models.utils import apply_normalization


from suno_utils.tasks.dac_2c_12cb import DAC

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

# VAE
from suno_utils.tasks.dac_vae_peaq import (
    preload_models as preload_vae_models,
    encode as vae_encode,
    decode as vae_decode,
)

# DAC
from suno_utils.tasks.dac_2c_12cb import (
    preload_models as preload_codec_models,
    encode as codec_encode,
)

# load MERT semantic
_ = 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",
)


# load VAE
_ = preload_vae_models(
    "/app/suno/checkpoints/dac_mw/mw_vae_peaq_128_fix/best/dac/weights.pth"
)

# load DAC codec
_ = preload_codec_models(
    checkpoint_filepath="s3://suno-data/georg/models/codec/dac_2c_25x12.pt",
)

if __name__ == "__main__":

    num_steps = 1000
    # ckpt_path = "/home/christian/christian/stable-audio-tools/checkpoints/upsample_model_mert_unwrap-epoch=32-step=270000.ckpt"
    # ckpt_path = "/home/christian/christian/stable-audio-tools/checkpoints/upsample_model_codec_unwrap-epoch=45-step=370000.ckpt"
    ckpt_path = "/home/christian/christian/stable-audio-tools/checkpoints/upsample_model_semantic+codec_unwrap-epoch=0-step=10000.ckpt"

    if "semantic+codec" in ckpt_path:
        model_type = "semantic+codec"
        model_config_path = "/home/christian/christian/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_semantic+codec_48khz.json"
    elif "codec" in ckpt_path:
        model_type = "codec"
        model_config_path = "/home/christian/christian/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_codec_48khz.json"
    elif "semantic" in ckpt_path:
        model_type = "semantic"
        model_config_path = "/home/christian/christian/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_semantic_48khz.json"

    # load model from checkpoint
    if model_config_path is not None:
        # Load config from json file
        with open(model_config_path) as f:
            model_config = json.load(f)
    else:
        model_config = None

    # model_config["model"]["diffusion"]["config"]["use_checkpointing"] = True

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model, model_config = load_model(
        model_config,
        ckpt_path,
        # pretrained_name=pretrained_name,
        # pretransform_ckpt_path=pretransform_ckpt_path,
        # model_half=model_half,
        device="cuda",
    )

    # file_ext = ".mp3.wav"
    # audio_paths = glob.glob("/app/suno/christian/data/trending-05062024-wav/*.mp3.wav")
    audio_paths = glob.glob("/app/suno/christian/reference-audio-wav/*.wav")
    # audio_paths = glob.glob("/app/suno/christian/data/bad_audio/*.wav")
    file_ext = ".wav"

    for audio_path in audio_paths:
        print(audio_path)

        audio_in, sr = torchaudio.load(audio_path)

        if sr != 48000:
            audio_in = torchaudio.functional.resample(audio_in, sr, 48000)

        start_frame = audio_in.shape[-1] // 2
        seconds_total = audio_in.shape[-1] / 48000
        end_frame = start_frame + int(10 * 48000)

        audio_in = audio_in[:, start_frame:end_frame]
        # audio_in = apply_normalization(audio_in, 48000, target_loudness_lufs_db=-16.0)
        audio_in /= audio_in.abs().max()

        semantic_codes = None
        codec_codes = None

        if "semantic" in model_type:
            audio_in_encode = torchaudio.functional.resample(audio_in, 48000, 24000)
            latents = semantic_encode([audio_in_encode.mean(dim=0, keepdim=True)])
            latents = [torch.from_numpy(latent[:, 0]).long() for latent in latents]
            semantic_codes = torch.cat(latents).long().cuda()
            print("semantic_codes", semantic_codes.shape)

        if "codec" in model_type:
            audio_in_encode = audio_in
            latents = codec_encode([audio_in_encode])
            latents = [torch.from_numpy(latent) for latent in latents]
            codec_codes = torch.cat(latents).long().cuda()
            print("codec_codes", codec_codes.shape)

        # apply lowpass
        # audio_in = torchaudio.functional.lowpass_biquad(audio_in, 48000, 500.0)

        for cfg in [1.0]:
            with torch.no_grad():
                upsampled_latents = upsample_diffusion_from_codes(
                    model,
                    semantic_codes=semantic_codes,
                    codec_codes=codec_codes,
                    steps=num_steps,
                    cfg_scale=cfg,
                    start_sample=start_frame,
                    seconds_total=seconds_total,
                    sample_size=semantic_codes.shape[-1],
                    sample_rate=48000,
                )
            upsampled_latents = upsampled_latents.cpu().squeeze().permute(1, 0)
            print(upsampled_latents.shape)

            # convert upsampled latents back to audio
            upsampled_audio = vae_decode(upsampled_latents)
            upsampled = torch.from_numpy(upsampled_audio.array_float)
            upsampled /= upsampled.abs().max()
            print(upsampled.shape)

            in_path = os.path.basename(audio_path).replace(file_ext, "-input.wav")
            out_path = os.path.basename(audio_path).replace(
                file_ext, f"-upsampled-{model_type}-cfg={cfg}-steps={num_steps}.wav"
            )
            in_path = os.path.join("upsampled", "20240604", in_path)
            out_path = os.path.join("upsampled", "20240604", out_path)

            if not os.path.isfile(in_path):
                torchaudio.save(in_path, audio_in.cpu(), 48000)

            torchaudio.save(out_path, upsampled.squeeze(0), 48000)