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
import polars as pl

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.mert_25 import (
    preload_models as preload_semantic_models,
    load_model as load_semantic_model,
    encode as encode_semantic,
    EMBEDDING_RATE as SEMANTIC_HZ,
)

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 random
import torchaudio


def apply_tanh_distortion(
    audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0
):
    gain_lin = 10 ** (gain_db / 20.0)
    return torch.tanh(audio * gain_lin)


def apply_clipping_distortion(
    audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0
):
    gain_lin = 10 ** (gain_db / 20.0)
    return (audio * gain_lin).clamp(-1, 1)


def apply_audio_codec(
    audio: torch.Tensor,
    sample_rate: float,
    bit_rate: int = 16000,
    n_passes: int = 1,
    format_str: str = "mp3",
):
    for _ in range(n_passes):
        effector = torchaudio.io.AudioEffector(
            format=format_str,
            codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),
        )
        audio = effector.apply(audio.T, sample_rate).T
    return audio


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_highpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float):
    return torchaudio.functional.highpass_biquad(audio, sample_rate, cutoff_hz)


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


def apply_random_tanh_distortion(audio: torch.Tensor, sample_rate: float):
    gain_db = random.uniform(0, 12)
    return apply_tanh_distortion(audio, sample_rate, gain_db)


def apply_random_clipping_distortion(audio: torch.Tensor, sample_rate: float):
    gain_db = random.uniform(0, 12)
    return apply_clipping_distortion(audio, sample_rate, gain_db)


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


def apply_random_audio_codec(audio: torch.Tensor, sample_rate: float):
    format_str = "mp3"
    bit_rate = random.choice([8000, 16000, 32000, 64000, 128000])
    n_passes = random.choice([1, 2])
    return apply_audio_codec(audio, sample_rate, bit_rate, n_passes, format_str)


def apply_random_highpass(audio: torch.Tensor, sample_rate: float):
    cutoff_hz = random.uniform(20, 6000)
    return apply_highpass(audio, sample_rate, cutoff_hz)


def apply_random_lowpass(audio: torch.Tensor, sample_rate: float):
    cutoff_hz = random.uniform(100, 12000)
    return apply_lowpass(audio, sample_rate, cutoff_hz)


def corrupt_audio(filepath: str, p=0.33):
    audio, sample_rate = torchaudio.load(filepath)

    if random.random() < p:
        audio = apply_random_highpass(audio, sample_rate)
    if random.random() < p:
        audio = apply_random_lowpass(audio, sample_rate)
    if random.random() < p:
        audio = apply_random_noise(audio, sample_rate)
    if random.random() < p:
        audio = apply_random_tanh_distortion(audio, sample_rate)
    if random.random() < p:
        audio = apply_random_clipping_distortion(audio, sample_rate)
    if random.random() < p:
        audio = apply_random_audio_codec(audio, sample_rate)

    return audio


class EarWorker:
    def __init__(
        self,
        ear_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}.")

        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"
        )
        _ = preload_semantic_models(semantic_model_filepath, semantic_clusters_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, sr = read_from_s3(s3_filepath, read_f=torchaudio.load)
                if sr != 48000:
                    audio_tensor = torchaudio.functional.resample(
                        audio_tensor, sr, 48000
                    )
                # ensure stereo (2 channels)
                if audio_tensor.ndim < 2:
                    audio_tensor = audio_tensor.unsqueeze(0)  # Add channel dimension

                # Check channel dimension and ensure it's exactly 2 channels
                if audio_tensor.shape[0] == 1:
                    # Convert mono to stereo by duplicating the channel
                    audio_tensor = audio_tensor.repeat(2, 1)
                elif audio_tensor.shape[0] > 2:
                    # If more than 2 channels, keep only the first 2
                    audio_tensor = audio_tensor[:2]

                # make sure the audio tensor is at least 60s
                if audio_tensor.shape[1] < 60 * 48000:
                    return meta["id"], None

                # Verify we have exactly 2 channels
                assert (
                    audio_tensor.shape[0] == 2
                ), f"Expected 2 channels, got {audio_tensor.shape[0]}"

                # crop to 4 min
                # ensure stereo
                audio_tensor = audio_tensor[:, : 4 * 60 * 48000]
                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:
            scores, mean_score = self.model.get_score(
                audio_tensor, sample_rate=48000, return_scores=True
            )
            example_scores[meta_id] = {
                "mean_score": mean_score,
                "scores": scores,
            }
            print(f"{meta_id}: {mean_score}")

        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}_ear_scores.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}_ear_scores.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-ear"
app = modal.App(APP_NAME, image=image, secrets=SECRETS)

N_MAX_REPLICAS = 32
EAR_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/ear/ear_v2_s3080.pt"


@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, ear_ckpt: str, output_path: str):
        import torch

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

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


@app.local_entrypoint()
def main():

    dataset_name = "imslp"

    # load the base metas
    metas_s3_filepath = "s3://suno-data/datasets/bundles/v4/genius/metas_v0.jsonl"
    metas = read_from_s3(metas_s3_filepath, read_f=pl.read_ndjson)

    # load the data cut
    data_cut = "/home/christian/code/christian/metadata/ear/genius_t6.json"
    with open(data_cut, "r") as f:
        info = json.load(f)
    print(len(info["genius"]))

    # filter the metas based on the data cut
    metas = metas.filter(pl.col("id").is_in(info["genius"]))
    print(len(metas))

    # set the seed
    import random

    random.seed(42)

    # select some random tracks
    rand_metas = metas.sample(1000)
    print(len(rand_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(
        EAR_MODEL_FILEPATH,
        f"christian/outputs/ear_scores/{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")