import os
import time
import modal
import torch
import funcy
import json
import random
import tempfile
import torchaudio
import numpy as np
import scipy.signal as signal
import pyloudnorm as pyln
from typing import Tuple

from suno_utils.audio import Audio
from suno_utils.worker.settings import s3_client
from suno_utils.worker.modal_base import MODAL_MOUNTS
from suno_utils.utils.text import read_jsonl

from suno_utils.tasks.mert_25 import (
    preload_models as preload_semantic_models,
    encode as encode_semantic,
)

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


def biquad(
    gain_db: float,
    cutoff_freq: float,
    q_factor: float,
    sample_rate: float,
    filter_type: str,
) -> Tuple[np.ndarray, np.ndarray]:
    """Use design parameters to generate coefficients for a specific filter type."""
    A = 10 ** (gain_db / 40.0)
    w0 = 2.0 * np.pi * (cutoff_freq / sample_rate)
    alpha = np.sin(w0) / (2.0 * q_factor)
    cos_w0 = np.cos(w0)
    sqrt_A = np.sqrt(A)

    if filter_type == "high_shelf":
        b0 = A * ((A + 1) + (A - 1) * cos_w0 + 2 * sqrt_A * alpha)
        b1 = -2 * A * ((A - 1) + (A + 1) * cos_w0)
        b2 = A * ((A + 1) + (A - 1) * cos_w0 - 2 * sqrt_A * alpha)
        a0 = (A + 1) - (A - 1) * cos_w0 + 2 * sqrt_A * alpha
        a1 = 2 * ((A - 1) - (A + 1) * cos_w0)
        a2 = (A + 1) - (A - 1) * cos_w0 - 2 * sqrt_A * alpha
    elif filter_type == "low_shelf":
        b0 = A * ((A + 1) - (A - 1) * cos_w0 + 2 * sqrt_A * alpha)
        b1 = 2 * A * ((A - 1) - (A + 1) * cos_w0)
        b2 = A * ((A + 1) - (A - 1) * cos_w0 - 2 * sqrt_A * alpha)
        a0 = (A + 1) + (A - 1) * cos_w0 + 2 * sqrt_A * alpha
        a1 = -2 * ((A - 1) + (A + 1) * cos_w0)
        a2 = (A + 1) + (A - 1) * cos_w0 - 2 * sqrt_A * alpha
    elif filter_type == "peaking":
        b0 = 1 + alpha * A
        b1 = -2 * cos_w0
        b2 = 1 - alpha * A
        a0 = 1 + alpha / A
        a1 = -2 * cos_w0
        a2 = 1 - alpha / A

    b = np.array([b0, b1, b2]) / a0
    a = np.array([1.0, a1 / a0, a2 / a0])
    return b, a


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


def apply_channel_imbalance(
    audio: torch.Tensor, sample_rate: float, imbalance: float = 0.0
):
    if not -1 <= imbalance <= 1 or audio.shape[-2] != 2:
        raise ValueError("Invalid input")
    out = audio.clone()
    l_gain, r_gain = (1.0 - imbalance, 1.0) if imbalance > 0 else (1.0, 1.0 + imbalance)
    out[0, :], out[1, :] = out[0, :] * l_gain, out[1, :] * r_gain
    return out


def apply_highpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float = 1000.0):
    return torchaudio.functional.highpass_biquad(audio, sample_rate, cutoff_hz)


def apply_lowpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float = 1000.0):
    return torchaudio.functional.lowpass_biquad(audio, sample_rate, cutoff_hz)


def apply_noise(
    audio: torch.Tensor,
    sample_rate: float,
    gain_db: float = 0.0,
    noise_type: str = "white",
):
    gain_lin = 10 ** (gain_db / 20.0)
    noise = torch.randn_like(audio)

    if noise_type == "white":
        return audio + gain_lin * noise
    elif noise_type == "pink":
        b = torch.tensor([0.049922035, -0.095993537, 0.050612699, -0.004408786])
        a = torch.tensor([1, -2.494956002, 2.017265875, -0.522189400])
        noise = torchaudio.functional.filtfilt(noise, a, b)
        noise /= noise.abs().max()
        return audio + gain_lin * noise
    else:
        raise ValueError(f"Invalid noise type: {noise_type}")


def apply_shelving_filter(
    audio: torch.Tensor,
    sample_rate: float,
    gain_db: float,
    cutoff_freq: float,
    q_factor: float,
    filter_type: str,
):
    # convert x to numpy
    audio = audio.numpy()
    b, a = biquad(
        gain_db,
        cutoff_freq,
        q_factor,
        sample_rate,
        filter_type,
    )
    x = signal.lfilter(b, a, audio).astype(np.float32)
    return torch.from_numpy(x)


# randomized corrputions


def apply_random_noise(audio: torch.Tensor, sample_rate: float):
    noise_type = random.choice(["white", "pink"])
    if noise_type == "white":
        noise_gain = random.uniform(-96, -48)
    else:
        noise_gain = random.uniform(-48, -12)
    return apply_noise(audio, sample_rate, noise_gain, noise_type)


def apply_random_stereo_to_mono(audio: torch.Tensor, sample_rate: float):
    return apply_stereo_to_mono(audio, sample_rate)


def apply_random_channel_imbalance(audio: torch.Tensor, sample_rate: float):
    imbalance = random.uniform(-1.0, 1.0)
    return apply_channel_imbalance(audio, sample_rate, imbalance)


def apply_random_filter(audio: torch.Tensor, sample_rate: float):
    filter_type = random.choice(["highpass", "lowpass", "high_shelf", "low_shelf"])
    if filter_type == "highpass":
        cutoff_freq = random.uniform(20, 4000)
        return apply_highpass(audio, sample_rate, cutoff_freq)
    elif filter_type == "lowpass":
        cutoff_freq = random.uniform(1000, 16000)
        return apply_lowpass(audio, sample_rate, cutoff_freq)
    else:
        gain_db = random.uniform(-12, 12)
        if filter_type == "high_shelf":
            cutoff_freq = random.uniform(6000, 20000)
        else:
            cutoff_freq = random.uniform(20, 2000)
        q_factor = random.uniform(0.1, 10.0)
        return apply_shelving_filter(
            audio, sample_rate, gain_db, cutoff_freq, q_factor, filter_type
        )


def corrupt(waveform_tensor, sample_rate):
    """
    Apply a random number of corruptions (at least one) to the input waveform.
    """
    # List of available random corruption functions
    corruption_fns = [
        apply_random_noise,
        apply_random_stereo_to_mono,
        apply_random_channel_imbalance,
        apply_random_filter,
    ]
    n_corr = random.randint(1, len(corruption_fns))  # at least one
    selected = random.sample(corruption_fns, n_corr)
    print(selected)
    out = waveform_tensor.clone()
    for fn in selected:
        out = fn(out, sample_rate)
    return torch.tanh(out)


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(
        "flashinfer-python", index_url="https://flashinfer.ai/whl/cu124/torch2.5/"
    )
    .pip_install_private_repos(
        "github.com/suno-ai/glockenspiel.git@a5dba4e50#subdirectory=descript-audio-codec&egg=descript-audio-codec",
        git_user="mcamac",
        secrets=[modal.Secret.from_name("victor-modal-github-token")],
    )
    .pip_install_private_repos(
        "github.com/suno-ai/neon.git@1c83548#subdirectory=hoot",
        git_user="mcamac",
        secrets=[modal.Secret.from_name("victor-modal-github-token")],
    )
    .pip_install(
        "boto3",
        "transformers",
        "tokenizers",
        "encodec",
        "ctc_segmentation",
        "psutil",
        "redis",
        "pydantic",
        "nnAudio",
        "rpyc",
        "biopython>=1.81",  # TODO: don't love this depdendency, for hoot
        "pynvml",  # for torch cuda utilization
        "torchsde",
        "ninja",
        "wheel",
    )
    .pip_install_from_pyproject(
        "/home/christian/code/glockenspiel/suno_utils/pyproject.toml",
    )
    .run_commands(  # This is really slow
        "git clone https://github.com/Dao-AILab/flash-attention.git",
        "cd flash-attention/hopper && python setup.py install",
        gpu="h100",
    )
    .apt_install(
        "libogg0",
        "libopus0",
        "opus-tools",
    )
    .pip_install("transformers==4.44.0", "wandb")
)


class GenerateWorker:
    def __init__(
        self,
        output_path: str,
    ):
        self.output_path = output_path
        start_time = time.time()
        print("Start loading models")
        num_gpus = torch.cuda.device_count()
        cuda_device = torch.cuda.current_device()
        print(f"Found {num_gpus} GPUs. Using GPU {cuda_device}.")

        semantic_model_filepath = "s3://suno-data/georg/models/semantic/mert_25.pt"
        semantic_clusters_filepath = (
            "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy"
        )
        codec_filepath = CODEC_FILEPATH

        _ = preload_semantic_models(semantic_model_filepath, semantic_clusters_filepath)
        _ = preload_codec_models(codec_filepath)

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

    @staticmethod
    def download_models(dir_path=MOUNT_PATH):
        """Download diffusion models."""
        print("Start downloading models")

        print("Finish downloading models")

    def generate(self, work_items):
        """Generate audio from a work item."""
        for item in work_items:
            item_id = item["id"]

            s3_filepath = item["s3_filepath"]
            # read audio from s3 and then semantic encode
            audio = Audio.from_s3(s3_filepath, n_channels=2)

            # also encode semantic codes
            print("Encoding semantic codes...")
            codes = encode_semantic(
                audio.convert(sample_rate=24_000, byte_width=2, n_channels=1)
            ).astype(np.int64)
            print("semantic codes shape:", codes.shape)

            audio = audio.convert(
                sample_rate=48_000, byte_width=2, n_channels=2
            ).normalize_volume(target_db=-16)
            print("audio shape:", audio.array_float.shape)

            # corrupt the audio
            audio_corrupted = corrupt(
                torch.from_numpy(audio.array_float), audio.sample_rate
            )
            audio_corrupted_obj = Audio.from_array_float(
                audio_corrupted.numpy(), audio.sample_rate
            )
            print("audio corrupted shape:", audio_corrupted_obj.array_float.shape)

            # encode the vae latents
            print("Encoding vae latents...")
            vae_latents = codec_encode(audio)
            vae_latents_corrupted = codec_encode(audio_corrupted_obj)

            print("vae latents shape:", vae_latents.shape)
            print("vae latents corrupted shape:", vae_latents_corrupted.shape)

            with tempfile.TemporaryDirectory() as td:
                vae_latents_path = os.path.join(td, f"{item_id}_vae.npz")
                np.savez(vae_latents_path, vae_latents=vae_latents)

                s3_filepath = os.path.join(
                    self.output_path,
                    f"{item_id}_vae.npz",
                )
                s3_client.upload_file(
                    vae_latents_path,
                    "suno-data",
                    s3_filepath,
                    ExtraArgs={"ContentType": "application/octet-stream"},
                )

                vae_latents_corrupted_path = os.path.join(
                    td, f"{item_id}_corrupted_vae.npz"
                )
                np.savez(vae_latents_corrupted_path, vae_latents=vae_latents_corrupted)

                s3_filepath = os.path.join(
                    self.output_path,
                    f"{item_id}_corrupted_vae.npz",
                )
                s3_client.upload_file(
                    vae_latents_corrupted_path,
                    "suno-data",
                    s3_filepath,
                    ExtraArgs={"ContentType": "application/octet-stream"},
                )

                codes_path = os.path.join(td, f"{item_id}_semantic.npz")
                np.savez(codes_path, codes=codes)

                s3_filepath = os.path.join(
                    self.output_path,
                    f"{item_id}_semantic.npz",
                )
                s3_client.upload_file(
                    codes_path,
                    "suno-data",
                    s3_filepath,
                    ExtraArgs={"ContentType": "application/octet-stream"},
                )


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")
    GenerateWorker.download_models()


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

N_MAX_REPLICAS = 32
GPU_TYPE = "A10G"
CODEC_FILEPATH = "s3://suno-data/minz/models/dac_vae_tuned_25hz.pth"


@app.cls(
    gpu=modal.gpu.A10G(count=1) if GPU_TYPE == "A10G" else modal.gpu.H100(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, output_path: str):
        import torch

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

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


@app.local_entrypoint()
def main():
    work_items = read_jsonl(
        "/home/christian/code/christian/metadata/ear/genius_t6_sampled_10k.jsonl"
    )

    print(f"Total work items: {len(work_items)}")
    chunksize = 128  # num of prompts per worker
    work_items = list(funcy.chunks(chunksize, work_items))
    print(f"Chunksize: {chunksize}, total chunks: {len(work_items)}")

    worker = GenerateStub(f"christian/outputs/corrupt/genius_t6_sampled_10k")

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

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