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.shimmerscore import shimmerscore, shimmerscore_url

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


def calculate_stereo_width(waveform):
    # Split into left and right channels
    left = waveform[0]
    right = waveform[1]

    # Compute mid/side representation
    mid = (left + right) / 2
    side = (left - right) / 2

    # Compute RMS energy of mid and side channels
    mid_energy = torch.sqrt(torch.mean(mid**2))
    side_energy = torch.sqrt(torch.mean(side**2))

    # Compute stereo width based on mid/side ratio
    # Normalize to range 0-1 using sigmoid-like function
    width_ratio = (side_energy / (mid_energy + 1e-8)).item()
    stereo_width = 2 * (1 / (1 + np.exp(-width_ratio)) - 0.5)

    return stereo_width


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

        print("Loading codec model...")
        _ = preload_codec_models(CODEC_FILEPATH)
        print("Loading ear model...")
        self.model = load_model(ear_model_filepath, compile=True)

        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"))
                shimmer_score = shimmerscore_url(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]

                # calculate stereo width
                stereo_width = calculate_stereo_width(audio_tensor)

                return meta["id"], audio_tensor, stereo_width, shimmer_score
            except Exception as e:
                print(f"Error processing {s3_filepath}: {e}")
                return meta["id"], None, None, 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, stereo_width, shimmer_score 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,
                "stereo_width": stereo_width,
                "shimmer_score": shimmer_score,
            }
            print(
                f"{meta_id}: {mean_score:0.3f} {stereo_width:0.3f} {shimmer_score:0.3f}"
            )

        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)


CODEC_FILEPATH = "s3://suno-data/minz/models/dac_vae_tuned_25hz.pth"


@app.local_entrypoint()
def main():

    dataset_name = "interesting_clips_ahi_d3_20250504"
    # load the main metas from the pretraining dataset
    # metas_s3_filepath = (
    #    f"s3://suno-data/datasets/bundles/v4/{dataset_name}/metas_v0.jsonl"
    # )
    # metas = read_from_s3(metas_s3_filepath, read_f=read_jsonl)
    filepath = "/home/christian/code/christian/metadata/dpo/interesting_clips_ahi_d3_20250504.jsonl"
    metas = read_jsonl(filepath)
    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(
        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")