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

import torch.nn as nn
import nnAudio.features as feat
import torch.nn.functional as F
from suno_utils.audio import Audio

N_BINS = 240
N_BANDS = 8
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
device = torch.device('cuda') 

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
):
    """
    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(
    s3_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)

    # 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)

if __name__ == "__main__":
    use_val = True

    # load pretrained tempo model
    tempo_model_path = '/app/suno/christian_c/label_bpm/deeprhythm-0.5.pth'

    tempo_model = load_tempo_model(tempo_model_path)

    # load s3 file
    assert len(sys.argv) == 2, "Please specify the s3_filepath as an command-line argument."
    s3_filepath = sys.argv[1]
    # print(s3_filepath)

    result = {}

    audio = prepare_audio(s3_filepath)
    output = evaluate_bpm(tempo_model, audio)
    
    # print(bpm)
    result['s3_filepath'] = s3_filepath
    result['bpm'] = output[0].item()
    result['bpm_conf'] = output[1].item()

    with open('test_bpm.json', 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=4)