import os
import json
import glob
import torch
import torchaudio
import numpy as np

from tqdm import tqdm
from typing import List

from suno_utils.utils.s3 import read_from_s3
from suno_utils.models.dac.nn.quantize_2 import ResidualVectorQuantize

# ---------------- functional corruptions ----------------


def apply_stereo_to_mono(audio: torch.tensor, sample_rate: float):
    return audio.mean(dim=0, keepdims=True).repeat(2, 1)


def apply_overlay_copy(audio: torch.tensor, sample_rate: float):
    copy_size = np.random.randint(audio.shape[-1] // 4, audio.shape[-1])
    start_idx = np.random.randint(0, audio.shape[-1] - copy_size)
    audio_copy = audio[:, start_idx : start_idx + copy_size]
    copy_start_idx = np.random.randint(0, audio.shape[-1] - copy_size)
    gain_db = np.random.rand() * -8
    gain_lin = 10 ** (gain_db / 20.0)
    audio_out = audio.clone()
    audio_out[:, copy_start_idx : copy_start_idx + copy_size] += gain_lin * audio_copy
    return audio_out


def apply_drop_random_samples(audio: torch.tensor, sample_rate: float):
    p = torch.rand(1).item()
    mask = torch.bernoulli(torch.full(audio.size(), p))
    return audio * mask


def apply_random_tanh_distortion(audio: torch.tensor, sample_rate: float):
    gain_db = np.random.rand() * 18
    gain_lin = 10 ** (gain_db / 20.0)
    return torch.tanh(audio * gain_lin)


def apply_random_clipping_distortion(audio: torch.tensor, sample_rate: float):
    gain_db = np.random.rand() * 18
    gain_lin = 10 ** (gain_db / 20.0)
    return (audio * gain_lin).clamp(-1, 1)


def apply_random_white_noise(audio: torch.Tensor, sample_rate: float):
    gain_db = (np.random.rand() * -48) - 18
    gain_lin = 10 ** (gain_db / 20.0)
    noise = gain_lin * torch.randn(audio.shape)
    return noise + audio


# ---------------- ffmpeg filter corruptions -------------


def get_random_bitcrusher_filter():
    bit_depths = [1, 2, 3, 4, 5, 6]
    aas = [0, 1]
    modes = ["lin", "log"]

    mix = np.random.rand() / 2
    bits = np.random.choice(bit_depths)
    mode = np.random.choice(modes)
    aa = np.random.choice(aas)

    return (
        f"acrusher=level_in=1:level_out=0.8:bits={bits}:mode={mode}:mix={mix:0.3f}:aa=1"
    )


def get_random_exciter_filter():
    amounts = np.arange(1, 64)
    drives = np.linspace(0.1, 10.0, num=50)
    blends = np.linspace(-10.0, 10.0, num=50)
    freqs = np.linspace(2000.0, 12000.0, num=50)
    ceils = np.linspace(9999.9, 20000.0, num=50)

    amount = np.random.choice(amounts)
    drive = np.random.choice(drives)
    blend = np.random.choice(blends)
    freq = np.random.choice(freqs)
    ceil = np.random.choice(ceils)

    return f"aexciter=level_in=1:level_out=0.8:amount={amount}:drive={drive}:blend={blend}:freq={freq}:ceil={ceil}"


def get_random_psyclip_filter():
    clips = np.linspace(-24.0, -4.0, num=24)
    clips = np.linspace(0.5, 1.0, num=25)
    diffs = [0, 1]
    adaptives = np.linspace(0.0, 1.0, num=10)

    clip = np.random.choice(clips)
    diff = np.random.choice(diffs)
    adaptive = np.random.choice(adaptives)

    return (
        f"apsyclip=level_in=1:level_out=0.8:clip={clip}:diff={diff}:adaptive={adaptive}"
    )


def get_random_pulsator_filter():
    modes = ["sine", "triangle", "square", "sawup", "sawdown"]
    amounts = np.linspace(0.0, 1.0, num=20)
    bpms = np.linspace(30, 300, num=100)
    widths = np.linspace(0.0, 2.0, num=20)

    mode = np.random.choice(modes)
    amount = np.random.choice(amounts)
    width = np.random.choice(widths)
    bpm = np.random.choice(bpms)

    return f"apulsator=level_in=1:level_out=1.0:mode={mode}:amount={amount}:width={width}:timing=bpm:bpm={bpm}"


def get_random_denoiser_filter():
    pass


def get_random_softclip_filter():
    clip_types = [
        "hard",
        "tanh",
        "atan",
        "cubic",
        "exp",
        "alg",
        "quintic",
        "sin",
        "erf",
    ]
    thresholds = np.linspace(0.125, 1.0, num=25)

    clip_type = np.random.choice(clip_types)
    threshold = np.random.choice(thresholds)

    return f"asoftclip=type={clip_type}:threshold={threshold}"


def get_random_bandpass_filter():
    frequencies = np.linspace(100.0, 10000.0, num=1000)
    width_type = "q"
    widths = [0.707, 1.0, 2.0, 4.0, 8.0, 10.0]

    frequency = np.random.choice(frequencies)
    width = np.random.choice(widths)

    return f"bandpass=frequency={frequency}:width_type={width_type}:width={width}"


def get_random_bandreject_filter():
    frequencies = np.linspace(100.0, 10000.0, num=1000)
    width_type = "q"
    widths = [0.707, 1.0, 2.0, 4.0, 8.0, 10.0]

    frequency = np.random.choice(frequencies)
    width = np.random.choice(widths)

    return f"bandreject=frequency={frequency}:width_type={width_type}:width={width}"


def get_random_crystalizer_filter():
    i = -(np.random.rand() * 10)

    return f"crystalizer=i={i}"


def get_random_dcshift_filter():
    shift = (np.random.rand() * 2) - 1
    limitergain = np.random.rand() * 0.1

    return f"dcshift=shift={shift}:limitergain={limitergain}"


def get_random_deesser_filter():
    i = np.random.rand()
    m = np.random.rand()
    f = np.random.rand()

    return f"deesser=i={i}:m={m}:f={f}"


def get_random_highpass_filter():
    frequencies = np.linspace(100.0, 10000.0, num=100)
    poless = [1, 2]
    width_type = "q"
    widths = [0.707, 1.0, 2.0, 4.0, 8.0, 10.0]

    frequency = np.random.choice(frequencies)
    poles = np.random.choice(poless)
    width = np.random.choice(widths)

    return f"highpass=frequency={frequency}:poles={poles}:width_type={width_type}:width={width}"


def get_random_lowpass_filter():
    frequencies = np.logspace(1, np.log10(4000), num=100)
    poless = [1, 2]
    width_type = "q"
    widths = [0.707, 1.0, 2.0, 4.0, 8.0, 10.0]

    frequency = np.random.choice(frequencies)
    poles = np.random.choice(poless)
    width = np.random.choice(widths)

    return f"lowpass=frequency={frequency}:poles={poles}:width_type={width_type}:width={width}"


def get_4k_lowpass_filter():
    poless = [1, 2]
    width_type = "q"
    widths = [0.707, 1.0, 2.0, 4.0, 8.0, 10.0]

    frequency = 4000.0
    poles = np.random.choice(poless)
    width = np.random.choice(widths)

    return f"lowpass=frequency={frequency}:poles={poles}:width_type={width_type}:width={width}"


def get_6k_lowpass_filter():
    poless = [1, 2]
    width_type = "q"
    widths = [0.707, 1.0, 2.0, 4.0, 8.0, 10.0]

    frequency = 6000.0
    poles = np.random.choice(poless)
    width = np.random.choice(widths)

    return f"lowpass=frequency={frequency}:poles={poles}:width_type={width_type}:width={width}"


def apply_audio_filters(
    audio: torch.tensor,
    sample_rate: float,
    filter_string: str,
    apply_mp3_codec: bool = False,
):
    """Apply FFmpeg audio filters to audio tensor via torchaudio `AudioEffector` interface.

    Args:
        audio (torch.Tensor): Stereo audio tensor (2, samples).
        sample_rate (float): Audio sample rate.
        filter_string (str): String defining a series of FFmpeg audio filters.
        apply_mp3_codec (bool): Apply mp3 codec to output audio.

    Returns:
        torch.Tensor: Processed stereo audio tensor (2, samples).
    """
    if apply_mp3_codec:
        bit_rates = [
            8_000,
            16_000,
            24_000,
            32_000,
            40_000,
            48_000,
            64_000,
            80_000,
            96_000,
            112_000,
            128_000,
        ]
        bit_rate = np.random.choice(bit_rates)
        effector = torchaudio.io.AudioEffector(
            filter_string,
            format="mp3",
            codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),
        )
    else:
        effector = torchaudio.io.AudioEffector(filter_string)

    return effector.apply(audio.T, sample_rate).T


def corrupt_audio(
    audio: torch.Tensor,
    sample_rate: int,
    corrupts: List[str],
    corrupt_probs: List[float],
    ffmpeg_filters: List[str],
    ffmpeg_filter_probs: List[float],
    mp3_codec_prob: float,
):
    assert len(corrupts) == len(corrupt_probs)
    assert len(ffmpeg_filters) == len(ffmpeg_filter_probs)

    # apply corruptions in sequence to audio
    if len(corrupts) > 0:
        # create a random ordering of the funcs
        func_order = torch.randperm(len(corrupts))
        for func_idx in func_order:
            # sample whether each func is active or not based on prob
            if torch.rand(1).item() < corrupt_probs[func_idx]:
                corrupt_func = CORRUPTS[corrupts[func_idx]]
                audio = corrupt_func(audio, sample_rate)

    filter_string = ""
    if len(ffmpeg_filters) > 0:
        # create a random ordering of the filters
        func_order = torch.randperm(len(ffmpeg_filters))
        for idx, func_idx in enumerate(func_order):
            # sample whether each func is active or not based on prob
            if torch.rand(1).item() < ffmpeg_filter_probs[func_idx]:
                ffmpeg_filter_func = FFMPEG_FILTERS[ffmpeg_filters[func_idx]]
                filter_string += ffmpeg_filter_func()
                if (idx + 1) != len(ffmpeg_filters):
                    filter_string += ","

    filter_string = filter_string.strip(",")  # remove trailing comma

    if filter_string == "":  # if no filters, set to None
        filter_string = None

    # apply the composite filters (if any) and mp3 codec
    audio = apply_audio_filters(
        audio,
        sample_rate,
        filter_string,
        apply_mp3_codec=True if np.random.rand() < mp3_codec_prob else False,
    )

    if audio.abs().max() > 1.0:
        audio /= audio.abs().max()

    return audio


CORRUPTS = {
    "stereo_to_mono": apply_stereo_to_mono,
    "drop_random_samples": apply_drop_random_samples,
    "tanh_distortion": apply_random_tanh_distortion,
    "clipping_distortion": apply_random_clipping_distortion,
    "white_noise": apply_random_white_noise,
    "overlay_copy": apply_overlay_copy,
}

# mapping from strings to func
FFMPEG_FILTERS = {
    "bitcrusher": get_random_bitcrusher_filter,
    "exciter": get_random_exciter_filter,
    "pulsator": get_random_pulsator_filter,
    "bandpass": get_random_bandpass_filter,
    "bandreject": get_random_bandreject_filter,
    "crystalizer": get_random_crystalizer_filter,
    "dcshift": get_random_dcshift_filter,
    "highpass": get_random_highpass_filter,
    "deesser": get_random_deesser_filter,
    "lowpass": get_random_lowpass_filter,
    "softclip": get_random_softclip_filter,
    "4k_lowpass": get_4k_lowpass_filter,
}


class AudioCorruptionDataset(torch.utils.data.Dataset):
    def __init__(self, filepaths: List[str], num_frames: int, chunks_per_file: int):
        self.filepaths = filepaths
        self.num_frames = num_frames
        self.chunks_per_file = chunks_per_file
        self.examples = []
        # create examples by chunking into num_frames chunks
        print(f"Creating file chunks of {num_frames} frames.")
        for filepath in tqdm(filepaths):
            # file_num_frames = torchaudio.info(filepath).num_frames
            # num_chunks = file_num_frames // num_frames
            # frame_offsets = [
            #    (chunk_idx * num_frames) for chunk_idx in range(num_chunks)
            # ]
            # for frame_offset in frame_offsets:
            #    self.examples.append((filepath, frame_offset))
            for _ in range(chunks_per_file):
                self.examples.append(filepath)
        print(f"Total of {len(self.examples)} chunks.")

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, idx):
        filepath = self.examples[idx]
        valid = True

        file_num_frames = torchaudio.info(filepath).num_frames

        if file_num_frames < self.num_frames:
            valid = False
            audio = torch.zeros(2, self.num_frames)
            corrupted_audio = torch.zeros(2, self.num_frames)
        else:
            frame_offset = np.random.randint(0, file_num_frames - self.num_frames - 1)

            # read random 10 sec chunk from disk
            audio, sample_rate = torchaudio.load(
                filepath, frame_offset=frame_offset, num_frames=self.num_frames
            )

            # run through randomized corruption pipeline
            corrupted_audio = corrupt_audio(
                audio,
                sample_rate,
                corrupts,
                corrupt_probs,
                ffmpeg_filters,
                ffmpeg_filter_probs,
                mp3_codec_prob,
            )

        return audio, corrupted_audio, valid


def load_rvq(
    codec_path: str,
    input_dim: int,
    n_codebooks: int,
    codebook_size: int,
    codebook_dim: int,
    quantizer_dropout: int,
):
    sd = read_from_s3(codec_path, read_f=torch.load)
    model = ResidualVectorQuantize(
        input_dim=input_dim,
        n_codebooks=n_codebooks,
        codebook_size=codebook_size,
        codebook_dim=codebook_dim,
        quantizer_dropout=quantizer_dropout,
    )

    model.load_state_dict(
        {k[10:]: v for k, v in sd["state_dict"].items() if k.startswith("quantize")}
    )
    model.eval()

    for param_name, param in model.named_parameters():
        param.requires_grad = False

    return model


# convert codes to latent embedding
def decode_vq(rvq_model, codes, n_quantizers):
    z_q = 0
    for i, quantizer in enumerate(rvq_model.quantizers[:n_quantizers]):
        _z_q = quantizer.embed_code(codes[:, :, i]).transpose(1, 2)
        _z_q = quantizer.out_proj(_z_q)
        z_q += _z_q.transpose(1, 2)
    return z_q


if __name__ == "__main__":
    # given a directory of audio files
    # load pretrained embedding model
    # read all audio files from disk
    # apply random corruptions to the audio
    # embed both input and corrupted version with model
    # store each seq of embedding to separate memmap with shared index

    use_val = False
    batch_size = 32
    chunks_per_file = 1
    num_frames = 480000
    model_type = "dac_2c_25x12"

    if model_type == "dac_100hz_12cb":
        n_codebooks = 12
        n_codes = 1000
        embed_dim = 128
        example_size = int(n_codebooks * n_codes)

        from suno_utils.tasks.dac_peaq100 import (
            preload_models as preload_vae_models,
            encode as vae_encode,
            decode as vae_decode,
        )

        # load VAE
        model = preload_vae_models(
            checkpoint_filepath="s3://suno-data/minz/models/dac_100hz_12cb.pth"
        )
    elif model_type == "mw_vae_peaq_128_fix":
        n_codebooks = 12
        n_codes = 250
        embed_dim = 128
        example_size = int(embed_dim * n_codes)

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

        # load VAE
        model = preload_vae_models(
            checkpoint_filepath="s3://suno-data/christian/mw_vae_peaq_128_fix.pth"
        )
    elif model_type == "dac_2c_25x12":
        n_codebooks = 12
        n_codes = 250
        embed_dim = 128
        example_size = int(n_codebooks * n_codes)

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

        # load VAE
        model = preload_vae_models(
            checkpoint_filepath="s3://suno-data/georg/models/codec/dac_2c_25x12.pt"
        )
    else:
        raise ValueError(f"Invalid model_type: {model_type}")

    root_dir = "/app/suno/data/audio_2ch_48khz_lg/"
    root_dir = root_dir + "val" if use_val else root_dir + "train"
    mm_out_dir = f"/app/suno/christian/data/enhance_{model_type}/"
    os.makedirs(mm_out_dir, exist_ok=True)

    if use_val:
        input_mm_filepath = os.path.join(mm_out_dir, "input_val.bin")
        corrupt_mm_filepath = os.path.join(mm_out_dir, "corrupt_val.bin")
    else:
        input_mm_filepath = os.path.join(mm_out_dir, "input_tr.bin")
        corrupt_mm_filepath = os.path.join(mm_out_dir, "corrupt_tr.bin")

    subsets = [
        "genius_hq",
        "imslp",
        "jamendo",
        # "podcasts",
        # "pond5_music",
        # "pond5_sfx",
        # "shutter_music",
        "spot_genres",
        # "tency",
        # "youtube_music",
    ]

    # setup corruption pipeline configuration
    corrupts = [
        "stereo_to_mono",
        "drop_random_samples",
        "tanh_distortion",
        "clipping_distortion",
        "white_noise",
        "overlay_copy",
    ]
    corrupt_probs = [0.25, 0.2, 0.3, 0.3, 0.3, 0.1]

    ffmpeg_filters = [
        "bitcrusher",
        "exciter",
        "bandpass",
        "highpass",
        "deesser",
        "lowpass",
    ]
    ffmpeg_filter_probs = [0.1, 0.1, 0.2, 0.2, 0.2, 0.2]
    mp3_codec_prob = 0.9

    # load audio
    subset_dirs = glob.glob(os.path.join(root_dir, "*"))

    # search for audio
    filepaths = []
    for subset_dir in subset_dirs:
        subset_name = os.path.basename(subset_dir)
        if subset_name not in subsets:
            continue
        print(subset_dir)
        subset_filepaths = glob.glob(os.path.join(subset_dir, "*.wav"))
        print(f"Found {len(subset_filepaths)} in {subset_name} (use_val={use_val}).")
        filepaths.extend(subset_filepaths)

    # create dataloader for multiprocessing
    dataset = AudioCorruptionDataset(filepaths, num_frames, chunks_per_file)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=batch_size,
        shuffle=False,
    )

    # create memmap filepointers
    num_examples = len(dataset)

    input_mm = np.memmap(
        input_mm_filepath,
        dtype=np.uint16,
        mode="w+",
        shape=(num_examples * example_size),
    )
    corrupt_mm = np.memmap(
        corrupt_mm_filepath,
        dtype=np.uint16,
        mode="w+",
        shape=(num_examples * example_size),
    )

    mm_write_idx = 0
    for batch_idx, batch in enumerate(tqdm(dataloader)):
        input_audios, corrupt_audios, valid = batch

        # check for valid audios
        input_audios = input_audios[valid, ...]
        corrupt_audios = corrupt_audios[valid, ...]

        input_audios_list = []
        corrupt_audios_list = []
        for elem_idx in range(input_audios.shape[0]):
            input_audios_list.append(input_audios[elem_idx, ...])
            corrupt_audios_list.append(corrupt_audios[elem_idx, ...])

        # embed input/output with VAE
        input_embeds = vae_encode(input_audios_list)
        corrupt_embeds = vae_encode(corrupt_audios_list)

        # write codes to the respective memmaps
        for input_embed, corrupt_embed in zip(input_embeds, corrupt_embeds):

            # check for nan in embeds
            # if np.any(np.isnan(input_embed)):
            #    print("Warning: found NaN in input_embed. Skipping..")
            #    continue

            # if np.any(np.isnan(corrupt_embed)):
            #    print("Warning: found NaN in corrupt_embed. Skipping...")
            #    continue

            input_mm[mm_write_idx : mm_write_idx + example_size] = np.reshape(
                input_embed, (-1)
            )
            corrupt_mm[mm_write_idx : mm_write_idx + example_size] = np.reshape(
                corrupt_embed, (-1)
            )
            mm_write_idx += example_size