import os
import glob
import time
import modal
import torch
import pathlib
import funcy
import pandas as pd
import numpy as np
import tempfile
import uuid
import json
import torchaudio

from suno_utils.audio import Audio
from suno_utils.utils.s3 import read_from_s3
from suno_utils.utils.text import read_jsonl
from suno_utils.worker.settings import s3_client
from suno_utils.worker.modal_base import MODAL_MOUNTS
from suno_utils.utils.s3 import read_from_s3, upload_s3_files, list_s3_dir
from suno_utils.tasks.ear import load_model

from suno_utils.tasks.dac_vae_fixed_25hz import (
    preload_models as preload_codec_models,
    decode as codec_decode,
    encode as codec_encode,
    decode_stream_to_full_audio,
)

MOUNT_PATH = "/suno/models"

aws_secret = modal.Secret.from_name("studio-aws")
SECRETS = [
    aws_secret,
    modal.Secret.from_dict(
        {
            "SUNO_ASSETS_PATH": "/suno/models/assets",
            "XDG_CACHE_HOME": "/suno/models/",
        }
    ),
    modal.Secret.from_name("api-callback-token"),
    modal.Secret.from_name("datadog-metrics"),
]

base_image = (
    modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.10")
    .apt_install(
        "curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3", "zlib1g-dev", "git", "clang"
    )
    .run_commands(
        [
            'curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"',
            "unzip -q awscliv2.zip",
            "./aws/install",
        ]
    )
    .dockerfile_commands(
        [
            "COPY --from=datadog/serverless-init:1.2.1 /datadog-init /app/datadog-init",
            'ENTRYPOINT ["/app/datadog-init"]',
        ]
    )
    .pip_install(
        "torch==2.5.1",
        "torchaudio==2.5.1",
    )
    .pip_install(
        "boto3",
        # "transformers",
        # "tokenizers",
        # "ctc_segmentation",
        # "psutil",
        # "redis",
        # "gradio",
        # "pydantic",
        "nnAudio",
        # "rpyc",
        # "biopython>=1.81",  # TODO: don't love this depdendency, for hoot
        # "pynvml",  # for torch cuda utilization
        # "torchsde",
        # "funcy",
        # "wandb",
    )
    .pip_install_from_pyproject(
        str(
            "/home/christian/code/glockenspiel/suno_utils/pyproject.toml",
        ),
        # force_build=True,
    )
    .pip_install(
        "transformers==4.44.0",
        "numpy==1.26.4",
    )
)

import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import nnAudio.features as feat
import torchaudio.transforms

from suno_utils.utils.s3 import read_from_s3

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()
    ckpt = read_from_s3(model_path, read_f=torch.load)
    model.load_state_dict(ckpt)
    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)
    if audio_filepath.startswith("s3://"):
        audio, sample_rate = read_from_s3(audio_filepath, read_f=torchaudio.load)
    else:
        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)


import os


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 BPMWorker:
    def __init__(
        self,
        tempo_model_filepath: str,
        output_path: str,
    ):
        self.output_path = output_path
        start_time = time.time()
        print("Start loading models")

        # load gpt model
        num_gpus = torch.cuda.device_count()
        cuda_device = torch.cuda.current_device()
        print(f"Found {num_gpus} GPUs. Using GPU {cuda_device}.")

        print("Loading tempo model...")
        self.model = load_tempo_model(tempo_model_filepath)

        print(
            f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds"
        )

    def score(self, work_items):
        from joblib import Parallel, delayed

        # split work_items into index and metas
        metas, index = work_items
        print(f"Scoring {len(metas)} metas...")

        # first download, load, and resample all audio in parallel
        def process_filepath(meta):
            try:
                s3_filepath = meta.get("audio_filepath", meta.get("s3_filepath"))
                audio_tensor = prepare_audio(s3_filepath)
                return meta["id"], audio_tensor
            except Exception as e:
                print(f"Error processing {s3_filepath}: {e}")
                return meta["id"], None

        # Process all filepaths in parallel
        results = Parallel(n_jobs=-1)(delayed(process_filepath)(meta) for meta in metas)
        print(f"Processed {len(results)} filepaths")

        # filter out None results
        results = [result for result in results if result[1] is not None]
        print(f"Successfully processed {len(results)} filepaths")

        example_scores = {}
        for meta_id, audio_tensor in results:

            bpm, confidence = evaluate_bpm(self.model, audio_tensor, "cuda")
            example_scores[meta_id] = {
                "bpm": bpm.item(),
                "confidence": confidence.item(),
            }
            print(f"{meta_id}: {bpm} {confidence}")

        print(f"Saving scores for {len(example_scores)} examples")
        # save the example_scores to a json file
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_dir = pathlib.Path(temp_dir)
            json_filepath = os.path.join(temp_dir, f"{index:06d}_bpm.json")
            with open(json_filepath, "w") as f:
                json.dump(example_scores, f)
            # upload the json file to s3
            upload_s3_files(
                json_filepath,
                f"s3://suno-data/{self.output_path}/{index:06d}_bpm.json",
            )


# prod diffusion model
# DIT_MODEL_FILEPATH = "s3://suno-data/georg/tmp/2b_prefix_ft.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/tony/tmp/diff/dit_v3_dpo_t10_3k_5e6_b100_t25.pt"


def download_model_wrapper_d():
    # this print is necessary to have modal rerun this when MODEL changes
    # Modal tracks referenced global variables
    # Change the name of the function to force a rerun
    print("Downloading model for history encoder")  #
    # EarWorker.download_models(EAR_MODEL_FILEPATH)


image = base_image.run_function(download_model_wrapper_d, secrets=SECRETS)
APP_NAME = f"batch-score-bpm"
app = modal.App(APP_NAME, image=image, secrets=SECRETS)

N_MAX_REPLICAS = 64
TEMPO_MODEL_FILEPATH = "s3://suno-data/christian/deeprhythm-0.5.pth"


@app.cls(
    gpu=modal.gpu.A10G(count=1),
    cpu=4,
    secrets=SECRETS,
    timeout=2 * 60 * 60,
    container_idle_timeout=240,
    mounts=MODAL_MOUNTS,
    memory=15000,
    concurrency_limit=N_MAX_REPLICAS,
)
class GenerateStub:
    def __init__(self, tempo_model_filepath: str, output_path: str):
        import torch

        num_gpus = torch.cuda.device_count()
        print(f"Found {num_gpus} GPUs.")
        self.worker = BPMWorker(tempo_model_filepath, output_path)

    @modal.method()
    def generate(self, work_item: list[dict]):
        return self.worker.score(work_item)


@app.local_entrypoint()
def main():

    dataset_name = "discogs_subset"
    # load the main metas from the pretraining dataset
    filepath = f"s3://suno-data/datasets/bundles/v4/{dataset_name}/metas_v0.jsonl"
    metas = read_from_s3(filepath, read_f=read_jsonl)
    print(len(metas))

    chunksize = 128  # num of prompts per worker
    # split the filepaths into chunks of chunksize and include an index
    work_items = [
        (metas[i : i + chunksize], i // chunksize)
        for i in range(0, len(metas), chunksize)
    ]
    print(len(work_items), "work items")

    worker = GenerateStub(
        TEMPO_MODEL_FILEPATH,
        f"christian/outputs/bpm/{dataset_name}",
    )

    if False:
        print("Testing inference...")
        t0 = time.time()
        for work_item in work_items[:1]:
            _ = worker.generate.remote(work_item)
        print(f"{int(round(time.time()-t0))}s for test")

    if True:
        print("Running batch inference...")
        t0 = time.time()
        _ = list(worker.generate.map(work_items))
        print(round((time.time() - t0) / 60 / 60), "h total for batch generation")