import os
import glob
import torch
import shutil
import torchaudio
import numpy as np
import pyloudnorm as pyln


def loudness_normalize(
    audio: torch.Tensor, sample_rate: int, target_lufs_db: float = -16.0
):
    meter = pyln.Meter(sample_rate)
    input_loudness_lufs_db = meter.integrated_loudness(audio.permute(1, 0).numpy())
    delta_loudness_lufs_db = target_lufs_db - input_loudness_lufs_db
    return 10 ** (delta_loudness_lufs_db / 20.0) * audio


if __name__ == "__main__":

    num_examples = 10
    seconds = 10.0
    target_lufs_db = -16.0

    shutil.rmtree("listening-test")
    os.makedirs("listening-test")

    random_filepaths = glob.glob(
        os.path.join("/app/suno/christian/data/suno_generations+postprocess", "*.wav")
    )
    print(
        f"Found {len(random_filepaths)} random filepaths. Selecting {num_examples} random examples."
    )
    random_filepaths = [fp for fp in random_filepaths if "output" not in fp]
    random_subset_filepaths = np.random.choice(
        random_filepaths, num_examples, replace=False
    )

    # now get random filepaths from trending list
    trending_filepaths = glob.glob(
        os.path.join("/app/suno/christian/data/trending-05062024+postprocess", "*.wav")
    )
    print(
        f"Found {len(trending_filepaths)} trending filepaths. Selecting {num_examples} random examples."
    )
    trending_filepaths = [fp for fp in trending_filepaths if "output" not in fp]
    trending_subset_filepaths = np.random.choice(
        trending_filepaths, num_examples, replace=False
    )

    filepath_lists = [random_subset_filepaths, trending_subset_filepaths]

    # load the input and output filepaths
    for subset_name, filepath_list in zip(["random", "trending"], filepath_lists):
        for idx, filepath in enumerate(filepath_list):
            output_filepath = filepath.replace(".wav", "-output.wav")

            x_in, sr = torchaudio.load(filepath)
            x_out, sr = torchaudio.load(output_filepath)

            # duration
            num_frames = x_in.shape[-1]
            mid_frame = num_frames // 2

            start_sample = mid_frame
            end_sample = start_sample + int(seconds * sr)

            # take the first seconds
            x_in_crop = x_in[:, start_sample:end_sample]
            x_out_crop = x_out[:, start_sample:end_sample]

            # loudness normalize both clips
            x_in_crop = loudness_normalize(x_in_crop, sr, target_lufs_db)
            x_out_crop = loudness_normalize(x_out_crop, sr, target_lufs_db)

            in_filepath = f"listening-test/{subset_name}-{idx+1:03d}-input.wav"
            out_filepath = f"listening-test/{subset_name}-{idx+1:03d}-output.wav"

            torchaudio.save(in_filepath, x_in_crop, sr)
            torchaudio.save(out_filepath, x_out_crop, sr)