import os
import glob
import json
import torch
import torchaudio

from stable_audio_tools.interface.gradio import load_model
from stable_audio_tools.models.utils import apply_normalization


if __name__ == "__main__":
    ckpt_path = "/home/christian/christian/stable-audio-tools/checkpoints/vae/vae_model_unwrap-epoch=61-step=1350000.ckpt"
    model_config_path = "/home/christian/christian/stable-audio-tools/stable_audio_tools/configs/model_configs/autoencoders/stable_audio_2_0_vae_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

    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",
    )

    audio_paths = glob.glob(
        os.path.join("/app/suno/christian/data/minz-codec-comparison", "*.wav")
    )
    print(audio_paths)

    for audio_path in audio_paths:
        audio_in, sr = torchaudio.load(audio_path)

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

        # loudness normalize to -16 dB
        # audio_in = apply_normalization(audio_in, 48000, target_loudness_lufs_db=-16.0)

        # start_frame = 524288
        # audio_in = audio_in[:, start_frame : start_frame + 524288]
        audio_in = audio_in.repeat(2, 1)
        audio_in = audio_in.cuda()
        audio_in /= audio_in.abs().max()

        with torch.no_grad():
            latents, encoder_info = model.encode(audio_in, return_info=True)
            decoded = model.decode(latents)

        print(decoded.abs().max())
        decoded /= decoded.abs().max()
        audio_in /= audio_in.abs().max()

        print(decoded.shape)

        basename = os.path.basename(audio_path).replace(".wav", "")
        input_filepath = os.path.join(
            "minz-codec-comparison-cycled", basename + "-input.wav"
        )
        output_filepath = os.path.join(
            "minz-codec-comparison-cycled", basename + "-decoded.wav"
        )

        torchaudio.save(input_filepath, audio_in.cpu().squeeze(0), 48000)
        torchaudio.save(output_filepath, decoded.cpu().squeeze(0), 48000)