import os
import uuid
import json
import glob
import torch
import boto3
import random
import resampy
import torchaudio
import numpy as np
import torch.nn as nn
import soundfile as sf
import concurrent.futures
import pyloudnorm as pyln
import multiprocessing as mp
import nnAudio.features as feat
import torch.nn.functional as F

from tqdm import tqdm
from time import perf_counter
from typing import Optional, List
from suno_utils.audio import Audio
from suno_utils.utils.text import read_jsonl

N_BINS = 240
N_BANDS = 8


def apply_log_filter(stft_output, filter_matrix):
    """
    Apply the logarithmic filter matrix to the Short-Time Fourier Transform (STFT) output.

    This function applies a precomputed logarithmic filter matrix to the STFT output of an audio signal
    to reduce its dimensionality and to capture the energy in logarithmically spaced frequency bands.

    Parameters
    ----------
    stft_output : torch.Tensor
        A tensor representing the STFT output with shape (batch_size, num_bins, num_frames), where
        num_bins is the number of frequency bins and num_frames is the number of time frames.
    filter_matrix : torch.Tensor
        A tensor representing the logarithmic filter matrix with shape (num_bands, num_bins), where
        num_bands is the number of logarithmically spaced frequency bands.

    Returns
    -------
    torch.Tensor
        A tensor representing the filtered STFT output with shape (batch_size, num_bands, num_frames).
        Each band contains the aggregated energy from the corresponding set of frequency bins.
    """
    stft_output_transposed = stft_output.transpose(1, 2)
    filtered_output_transposed = torch.matmul(stft_output_transposed, filter_matrix.T)
    filtered_output = filtered_output_transposed.transpose(1, 2)
    return filtered_output


def evaluate_bpm(model: torch.nn.Module, eval_audio: torch.Tensor, device: str):
    """
    Args:
        system (torch.nn.Module):
        eval_audio (torch.Tensor): Audio to evaluate with shape (bs, n_harmonics=6, n_bins, n_bands)

    """
    with torch.no_grad():
        eval_audio = eval_audio.to(device)
        outputs = model(eval_audio)
        probs = torch.softmax(outputs, dim=1)
        confs, preds = torch.max(probs, 1)
    return torch.tensor([class_to_bpm(pred) for pred in preds.tolist()]), torch.tensor(
        confs.tolist()
    )


def class_to_bpm(class_index, min_bpm=30, max_bpm=286, num_classes=256):
    """Map a class index back to a BPM value (to the center of the class interval)."""
    class_width = (max_bpm - min_bpm) / num_classes
    bpm = min_bpm + class_width * (class_index)
    return bpm


def compute_hcqm(y, stft_spec, band_filter, cqt_specs):
    """
    Compute the Harmonic Constant-Q Modulation (HCQM) for an input signal.

    As described by Foroughmand & Peeters in
    "Deep-Rhythm for Tempo Estimation and Rhythm Pattern Recognition", 2019

    Parameters:
    - y (Tensor): The input signal tensor of shape (batch_size, num_samples).
    - stft_spec (STFT object): An object to compute the Short-Time Fourier Transform (STFT).
    - band_filter (Tensor): A filter matrix of shape (num_bands, num_bins) to apply to the STFT.
    - cqt_specs (list of CQT objects): A list of Constant-Q Transform (CQT) objects for different harmonics / bands

    Returns:
    - hcqm (Tensor): The computed HCQM of shape (batch_size, N_BINS, N_BANDS, N_HARMONICS), where 6 corresponds to the number of different harmonics analyzed.
    """
    stft = stft_spec(y)
    stft_bands = apply_log_filter(stft, band_filter)
    stft_bands_flat = stft_bands.reshape(
        stft.size(0) * stft_bands.size(1), stft_bands.size(2)
    )
    osf_flat = onset_strength(y=stft_bands_flat)
    hcqm = torch.zeros((stft.size(0) * N_BANDS, N_BINS, 6))
    for h, spec in enumerate(cqt_specs):
        hcqm[:, :, h] = spec(osf_flat).mean(-1)
    hcqm = hcqm.reshape(stft_bands.size(0), N_BINS, N_BANDS, 6)
    return hcqm


def create_log_filter(num_bins, num_bands):
    log_bins = (
        np.logspace(np.log10(1), np.log10(num_bins), num=num_bands + 1, base=10.0) - 1
    )
    log_bins = np.unique(np.round(log_bins).astype(int))
    filter_matrix = torch.zeros(num_bands, num_bins)
    for i in range(num_bands):
        if i < num_bands - 1:
            start_bin, end_bin = log_bins[i], log_bins[i + 1]
        else:
            start_bin, end_bin = log_bins[i], num_bins
        filter_matrix[i, start_bin:end_bin] = 1 / (end_bin - start_bin)

    return filter_matrix


def load_tempo_model(model_path: str):
    model = DeepRhythmModel()
    model.load_state_dict(torch.load(model_path))
    model.cuda()
    model.eval()
    return model


def make_kernels(len_audio=22050 * 8, sr=22050):
    n_fft = 2048
    hop = 512
    n_fft_bins = int(1 + n_fft / 2)
    band_filter = create_log_filter(n_fft_bins, N_BANDS)
    stft_spec = feat.stft.STFT(
        sr=sr, n_fft=n_fft, hop_length=hop, output_format="Magnitude", verbose=False
    )
    cqt_specs = []
    for h in [1 / 2, 1, 2, 3, 4, 5]:
        # Convert from BPM to Hz
        fmin = (32.7 * h) / 60
        sr_cqt = len_audio // (hop * 8)
        fmax = sr_cqt / 2
        num_octaves = np.log2(fmax / fmin)
        bins_per_octave = N_BINS / num_octaves
        cqt_spec = feat.cqt.CQT(
            sr=sr_cqt,
            hop_length=len_audio // hop,
            n_bins=N_BINS,
            bins_per_octave=bins_per_octave,
            fmin=fmin,
            output_format="Magnitude",
            verbose=False,
            pad_mode="constant",
        )
        cqt_specs.append(cqt_spec)
    return stft_spec, band_filter, cqt_specs


def onset_strength(
    y=None,
    n_fft=2048,
    hop_length=512,
    lag=1,
    ref=None,
    detrend=False,
    center=True,
    aggregate=None,
):
    """
    Compute the onset strength of an audio signal or a spectrogram.

    The onset strength is a measure of the increase in energy of an audio signal.

    Parameters
    ----------
    y : torch.Tensor, optional
        The raw audio waveform, expected to be a 2D tensor of shape (batch_size, time_samples).
        If provided, it will be used to compute the spectrogram internally. Default is None.
    n_fft : int, optional
        The number of FFT components. Default is 2048.
    hop_length : int, optional
        The number of samples between successive frames. Default is 512.
    lag : int, optional
        The lag between frames for computing the difference in energy. Default is 1.
    ref : torch.Tensor, optional
        The reference spectrogram to which the energy difference is computed. If None, the
        spectrogram provided by `S` or computed from `y` is used as the reference. Default is None.
    detrend : bool, optional
        If True, remove the mean from the onset envelope. Default is False.
    center : bool, optional
        If True, pad the time dimension of the onset envelope so that frames are centered around
        their timestamps. Default is True.
    aggregate : callable, optional
        A function to aggregate the channels dimension (e.g., torch.mean, torch.sum). If None,
        the mean is used. Default is None.

    Returns
    -------
    torch.Tensor
        The onset strength envelope, a 2D tensor of shape (batch_size, time_frames).

    """
    # Ensure y is reshaped to (batch, channels, time) if it's not already
    if y is not None and y.dim() == 2:
        y = y.unsqueeze(1)

    S = torchaudio.transforms.AmplitudeToDB(top_db=80)(y)
    ref = S

    # Compute difference to reference, spaced by lag
    onset_env = S[..., lag:] - ref[..., :-lag]
    onset_env = torch.clamp(onset_env, min=0.0)  # Discard negatives

    if aggregate is None:
        aggregate = torch.mean
    if callable(aggregate):
        onset_env = aggregate(onset_env, dim=-2)

    # Padding and detrending
    pad_width = lag
    if center:
        pad_width += n_fft // (2 * hop_length)
    onset_env = F.pad(onset_env, (pad_width, 0), "constant", 0)

    if detrend:
        onset_env -= onset_env.mean(dim=-1, keepdim=True)

    if center:
        onset_env = onset_env[..., : S.shape[-1]]
    return onset_env


def prepare_audio(
    audio_filepath: str,
    start_s: float = None,
    end_s: float = None,
):
    # audio = Audio.from_s3(s3_filepath)
    # sample_rate = audio.sample_rate
    # audio = torch.from_numpy(audio.array_float)
    audio, sample_rate = torchaudio.load(audio_filepath)
    audio = audio.mean(dim=0)

    # crop audio based on metadata example
    if start_s is not None and end_s is not None:
        start_frame = int(start_s * sample_rate)
        end_frame = int(end_s * sample_rate)
        audio = audio[start_frame:end_frame]

    # downmix and resample decoded audio to appropriate sr, also set num_frames
    audio = torchaudio.functional.resample(audio, sample_rate, 22_050)
    num_frames = int(8 * 22_050)
    # clip 8 seconds (but lagged to avoid intros)
    if (
        audio.shape[-1] > num_frames and audio.shape[-1] < 15 * 22_050 + num_frames
    ):  # if longer than 8 but shorter than 23 seconds
        audio = audio[audio.shape[-1] - num_frames :]  # take last 8 seconds
    elif audio.shape[-1] > num_frames:  # if "normal" (> 23 seconds)
        start_ix = 15 * 22_050  # take 00:15 to 00:23
        audio = audio[start_ix : start_ix + num_frames]
    elif (
        audio.shape[-1] < num_frames
    ):  # pad by repeating the signal if shorter than window
        pad_size = num_frames - audio.shape[-1]
        audio = torch.tensor(
            np.pad(audio.detach().cpu().numpy(), (0, pad_size), "wrap")
        )
    audio = preprocess_tempo_audio(audio)

    return audio


def preprocess_tempo_audio(audio: torch.Tensor):
    stft_spec, band_filter, cqt_specs = make_kernels()
    input = torch.unsqueeze(audio, 0)
    preprocessed_audio = compute_hcqm(input, stft_spec, band_filter, cqt_specs).permute(
        0, 3, 1, 2
    )
    return preprocessed_audio


class DeepRhythmModel(nn.Module):
    def __init__(self, num_classes=256):
        super(DeepRhythmModel, self).__init__()
        # input shape is (6, 240, 8)
        self.num_classes = num_classes
        self.conv1 = nn.Conv2d(
            in_channels=6, out_channels=128, kernel_size=(4, 6), padding="same"
        )
        self.bn1 = nn.BatchNorm2d(128)
        self.conv2 = nn.Conv2d(
            in_channels=128, out_channels=64, kernel_size=(4, 6), padding="same"
        )
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(
            in_channels=64, out_channels=64, kernel_size=(4, 6), padding="same"
        )
        self.bn3 = nn.BatchNorm2d(64)
        self.conv4 = nn.Conv2d(
            in_channels=64, out_channels=32, kernel_size=(4, 6), padding="same"
        )
        self.bn4 = nn.BatchNorm2d(32)
        self.conv5 = nn.Conv2d(in_channels=32, out_channels=8, kernel_size=(120, 6))
        self.bn5 = nn.BatchNorm2d(8)
        self.fc1 = nn.Linear(2904, 256)
        self.elu = nn.ELU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(256, num_classes)
        self._initialize_weights()

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = F.relu(self.bn4(self.conv4(x)))
        x = F.relu(self.bn5(self.conv5(x)))
        x = x.reshape(x.size(0), -1)
        x = self.dropout(self.elu(self.fc1(x)))
        x = self.fc2(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                nn.init.constant_(m.bias, 0)


def download_audio(s3_filepath: str, example_id: str, tmp_dir: str):
    filename = os.path.basename(s3_filepath)
    out_filepath = os.path.join(tmp_dir, f"{example_id}-{filename}")
    # only download the file if its not already downloaded
    if not os.path.isfile(out_filepath):
        os.system(f"aws s3 cp {s3_filepath} {out_filepath} > /dev/null 2>&1")
    return out_filepath


class AudioMetadataDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        metas: List[dict],
        num_frames: int,
        tmp_dir: str = "/mnt/localdisk/cjs",
        extra_metas_map: dict = None,
    ):
        # we need to filter metas to only use ones with s3_filepath
        # but, we have to be careful to maintain the index of the original metadata
        filtered_metas = []
        for meta_idx, meta in enumerate(metas):
            if "audio_filepath" in meta:
                meta["orig_idx"] = meta_idx
                filtered_metas.append(meta)

        num_filtered = len(filtered_metas)
        num_original = len(metas)
        percent_remaining = (num_filtered / num_original) * 100
        print(
            f"{num_filtered}/{num_original} ({percent_remaining:0.2f}%) examples have s3_filepath."
        )
        self.metas = filtered_metas
        self.num_frames = num_frames
        self.tmp_dir = tmp_dir
        self.extra_metas_map = extra_metas_map

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

    def __getitem__(self, idx):
        meta = self.metas[idx]

        if meta["id"] in self.extra_metas_map:
            return (
                torch.zeros(1, self.num_frames),
                meta["id"],
                meta["audio_filepath"],
                torch.tensor(False),
            )

        filepath = download_audio(meta["audio_filepath"], meta["id"], self.tmp_dir)
        try:
            audio = prepare_audio(filepath, None, None)
        except Exception as e:
            print(f"Error: {e}")
            print(f"Failed to process: {filepath}")
            return (
                torch.zeros(1, self.num_frames),
                meta["id"],
                meta["audio_filepath"],
                torch.tensor(False),
            )

        # randomly sample N chunks
        # max_chunks = min(8, audio.shape[0])
        # chunk_idx = torch.randint(0, audio.shape[0], [max_chunks])

        # audio = audio[chunk_idx, ...]
        # os.remove(filepath)  # delete audio file from tmp directory
        # orig_idx = torch.tensor(meta["orig_idx"])

        # delete audio file
        os.remove(filepath)

        return audio, meta["id"], meta["audio_filepath"], torch.tensor(True)


class AudioFileDataset(torch.utils.data.Dataset):
    def __init__(self, audio_files: List[str], num_frames: int):
        self.audio_files = audio_files
        self.num_frames = num_frames

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

    def __getitem__(self, idx):
        audio_file = self.audio_files[idx]
        audio = prepare_audio(audio_file, self.num_frames, None, None)
        return audio_file, audio


if __name__ == "__main__":
    num_compare = 5
    num_frames = 131072
    use_val = False

    # load pretrained tempo model
    tempo_model_path = "checkpoints/deeprhythm-0.5.pth"
    if not os.path.isfile(tempo_model_path):
        os.system("aws s3 cp s3://suno-data/christian/deeprhythm-0.5.pth checkpoints/")
    tempo_model = load_tempo_model(tempo_model_path)

    # load metadata
    meta_filepath = "metadata/genius_hq_metas.jsonl"
    metas = read_jsonl(meta_filepath, progress=True)

    # output metadata
    tempo_metas_map = {}
    tempo_metas_filepath = "metadata/genius_hq_metas_tempo.json"

    if os.path.isfile(tempo_metas_filepath):
        with open(tempo_metas_filepath, "r") as f:
            tempo_metas_map = json.load(f)
        print(f"Loaded {len(tempo_metas_map)} tempo metas.")

    # load metadata dataset
    dataset = AudioMetadataDataset(metas, num_frames, extra_metas_map=tempo_metas_map)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=32)

    # label
    for batch_idx, batch in enumerate(tqdm(dataloader)):
        audios, example_id, audio_filepath, valid = batch

        if valid.item() is False:
            continue

        example_id = example_id[0]
        audio_filepath = audio_filepath[0]

        audios = audios.squeeze(0)

        # run tempo estimation
        output = evaluate_bpm(tempo_model, audios, "cuda")

        # add this meta to existing new meta
        new_meta = {
            "audio_filepath": audio_filepath,
            "tempo": {
                "bpm": f"{output[0].item():0.0f}",
                "confidence": f"{output[1].item():0.2f}",
            },
        }
        tempo_metas_map[example_id] = new_meta

        # save every 1000 iterations
        if batch_idx % 10000 == 0:
            print(f"Saving tempo metas... (step={batch_idx})")
            with open(tempo_metas_filepath, "w") as f:
                json.dump(tempo_metas_map, f, indent=4)

    # final save
    with open(tempo_metas_filepath, "w") as f:
        json.dump(tempo_metas_map, f, indent=4)