import os
import glob
import torch
import torchaudio
import numpy as np
import multiprocessing as mp

from tqdm import tqdm

# in this script we will load a series of audio files from disk
# for each audio file we will extract a melspectrogram as well as VAE latents from pretrained model
# we will then save each of these to disk as a numpy memmap file

USE_VAL = True

# VAE config
EMBED_DIM = 128  # VAE embed dim
SEQ_LEN = 10.0  # seconds
RATE_HZ = 100  # VAE frame rate
VAE_SAMPLE_RATE = 48_000  # audio sample rate
NUM_LATENTS = int(RATE_HZ * SEQ_LEN)

# melspec config
N_MELS = 64  # number of mel bands
HOP_LEN = 961
N_FFT = 1920
STD_NORM = True  # whether to standardize the melspecs
MELSPEC_SAMPLE_RATE = 24_000

# load the pretrained VAE model
from suno_utils.tasks.dac_vae_peaq import (
    preload_models as preload_vae_models,
    load_model as load_vae_model,
    encode as vae_encode,
    decode as vae_decode,
)

vae_model = load_vae_model(
    checkpoint_filepath="s3://suno-data/christian/mw_vae_peaq_128_fix.pth",
    device="cuda",
)

# load the pretrained semantic model
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",
)


# load the melspec transform
melspec_encode = torchaudio.transforms.MelSpectrogram(
    sample_rate=MELSPEC_SAMPLE_RATE,
    n_fft=N_FFT,  # Length of the FFT window
    win_length=None,  # Window size
    hop_length=HOP_LEN,  # Number of samples between successive frames
    n_mels=N_MELS,  # Number of Mel bands
    center=True,  # Whether the t-th frame is centered at t*hop_length
    pad_mode="reflect",  # Padding mode
    power=2.0,  # Power of the norm
)


if __name__ == "__main__":

    # set up paths
    root_dir = "/app/suno/data/audio_2ch_48khz_lg"
    if USE_VAL:
        root_dir = os.path.join(root_dir, "val")
    else:
        root_dir = os.path.join(root_dir, "train")
    subset_dirs = glob.glob(os.path.join(root_dir, "*"))
    audio_files = []
    for subset_dir in subset_dirs:
        audio_files += glob.glob(os.path.join(subset_dir, "*.wav"))
    print(f"Found {len(audio_files)} audio files")

    num_examples = len(audio_files)

    # set up memmaps with train and val sets
    # compute melspec time frames

    # melspec_memmap = np.memmap(
    #    "melspecs.bin", dtype=np.float32, mode="w+", shape=(num_examples, 128, 128)
    # )
    latents_memmap = np.memmap(
        "vae.bin",
        dtype=np.float32,
        mode="w+",
        shape=(num_examples, EMBED_DIM, NUM_LATENTS),
    )

    # process files
    for audio_file in tqdm(audio_files):
        # load audio (assume input is 48khz)
        load_frames = int(SEQ_LEN * 48_000)
        frame_offset = np.random.randint(
            0, torchaudio.info(audio_file).num_frames - load_frames
        )
        audio, sample_rate = torchaudio.load(
            audio_file, frame_offset=frame_offset, num_frames=load_frames
        )

        if sample_rate != VAE_SAMPLE_RATE:
            audio_vae = torchaudio.functional.resample(
                audio, sample_rate, VAE_SAMPLE_RATE
            )
        else:
            audio_vae = audio

        if sample_rate != MELSPEC_SAMPLE_RATE:
            audio_melspec = torchaudio.functional.resample(
                audio, sample_rate, MELSPEC_SAMPLE_RATE
            )
        else:
            audio_melspec = audio

        # crop audio to target length
        audio_vae = audio_vae[:, : int(SEQ_LEN * VAE_SAMPLE_RATE)]
        audio_melspec = audio_melspec[:, : int(SEQ_LEN * MELSPEC_SAMPLE_RATE)]

        with torch.no_grad():
            # get melspec of mono audio
            melspec = melspec_encode(audio_melspec.mean(dim=0, keepdim=True))
            # torch.Size([1, 64, 250])

            # get VAE latents
            vae_latents = (
                vae_model(audio_vae.unsqueeze(0).cuda())["z"].detach().cpu().numpy()
            )
            # (1, 128, 250)