import spaces
from transkun.ModelTransformer import TransKun, writeMidi
import numpy as np
import torch
import librosa

import tempfile
import os

import soundfile as sf

# setuptools dropped this module in version 81, and it is only needed here
# to locate data files inside the installed transkun package, which
# importlib.resources does from the standard library.
from importlib.resources import files
import moduleconf

from pyharp.core import ModelCard, build_endpoint
from pyharp.media.audio import load_audio
from pyharp.tags import Channels, SampleRate, Subcategory
import gradio as gr

model_card = ModelCard(
    name="Transkun",
    description=(
        "Transcribes solo piano performance into MIDI notation, with onsets, offsets, and "
        "velocities. The input is mixed to mono and resampled to 44.1 kHz internally. "
        "Repository: https://github.com/Yujia-Yan/Transkun Paper: "
        "https://arxiv.org/abs/2404.09466"
    ),
    author="Yujia Yan, Zhiyao Duan",
    tags=[Subcategory.MUSIC_ANALYSIS, SampleRate(44100), Channels(1), 'domain:music']
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def load_model(device: str):
    defaultWeight = str(files("transkun") / "pretrained" / "2.0.pt")
    defaultConf = str(files("transkun") / "pretrained" / "2.0.conf")

    confManager = moduleconf.parseFromFile(defaultConf)
    conf = confManager["Model"].config

    checkpoint = torch.load(defaultWeight, map_location="cpu")
    model = TransKun(conf=conf)

    # Checkpoints keep their weights under either key
    if "best_state_dict" in checkpoint:
        model.load_state_dict(checkpoint["best_state_dict"], strict=False)
    else:
        model.load_state_dict(checkpoint["state_dict"], strict=False)

    model.eval()
    return model.to(device)

# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the model is built and placed there alone, rather than inside every GPU call
if __name__ != "__main__":
    model = load_model(DEVICE)

# The recording is transcribed in overlapping segments, so the cost grows with its length
# PyHARP stops a job that runs past its timeout_s, 15 minutes by default, so a request whose
# reservation would exceed that is refused rather than started. A shorter one is left to
# ZeroGPU, which admits it when the caller's own allowance covers it.
MAX_GPU_SECONDS = 900


def gpu_seconds(input_file):
    # Measured at about 2 to 8 s for 30 s of piano and 10 to 15 s for 300 s, upload included,
    # padded. The model is built when the worker starts, not inside the call.
    return 8 + 0.05 * sf.info(input_file).duration


@spaces.GPU(duration=gpu_seconds)
def transcribe(input_file):
    signal = load_audio(input_file)
    waveform = np.asarray(signal.audio_data)
    sr = int(signal.sample_rate)

    waveform = np.squeeze(waveform)

    # If 2D, assume (channels, samples). Make mono.
    if waveform.ndim == 2:
        if waveform.shape[0] > 1:
            waveform = waveform.mean(axis=0)
        else:
            waveform = waveform.reshape(-1)

    if sr != model.fs:
        waveform = librosa.resample(waveform.astype(np.float32), orig_sr=sr, target_sr=model.fs)
        sr = model.fs

    x = torch.from_numpy(waveform.reshape(-1, 1)).to(DEVICE)

    notesEst = model.transcribe(x, discardSecondHalf=False)

    outputMidi = writeMidi(notesEst)

    out_fd, out_path = tempfile.mkstemp(suffix=".mid")
    os.close(out_fd)
    outputMidi.write(out_path)
    return out_path

def process_fn(input_audio_path: str) -> str:
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not input_audio_path:
        raise gr.Error("Provide the piano recording to transcribe.")
    try:
        duration = sf.info(input_audio_path).duration
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if duration == 0:
        raise gr.Error("The audio file is empty.")
    if gpu_seconds(input_audio_path) > MAX_GPU_SECONDS:
        raise gr.Error(
            f"The recording is {duration:.0f} s long, more than one request can transcribe in "
            f"the {MAX_GPU_SECONDS // 60} minutes it may run. Trim it."
        )
    midi_path = transcribe(input_audio_path)
    return midi_path

if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_audio = gr.Audio(label="Upload Solo Piano Audio", type="filepath").harp_required(True).set_info((
            "A solo piano recording. Trained on piano alone, so other "
            "instruments in the recording will confuse it."
        ))

        output_midi = gr.File(
            label="Transcription",
            file_types=[".mid", ".midi"],
            type="filepath"
        ).set_info("The transcribed notes, with onsets, offsets and velocities.")


        app = build_endpoint(
            model_card=model_card,
            input_components=[input_audio],
            output_components=[output_midi],
            process_fn=process_fn
        )

    demo.queue().launch(pwa=True, show_error=True, ssr_mode=False)