import spaces  # imported first, as ZeroGPU requires

try:  # torch>=2.6 flipped torch.load(weights_only) to True; legacy ckpts need False
    import torch as _torch

    if getattr(_torch.load, "__harp_compat__", False) is False:
        _torch_load_orig = _torch.load

        def _torch_load_compat(*args, **kwargs):
            # Lightning passes weights_only=None, which torch reads as True
            if kwargs.get("weights_only") is None:
                kwargs["weights_only"] = False
            return _torch_load_orig(*args, **kwargs)

        _torch_load_compat.__harp_compat__ = True
        _torch.load = _torch_load_compat
except Exception:  # torch not installed / unexpected API -- nothing to patch
    pass

import time

import numpy as np
import soundfile as sf
import audiotools as at
import torch
from pyharp.core import build_endpoint, ModelCard
from pyharp.media.audio import save_audio
from pyharp.tags import Subcategory, SampleRate, Channels

import gradio as gr
from vampnet.interface import Interface
from vampnet import mask as pmask

model_card = ModelCard(
    name="VampNet",
    description=(
        "Generates new audio by remixing and extending your input, filling in "
        "the blanks to create a coherent continuation in the same style. "
        "Besides the default, three fine-tuned models are available. cat10 was "
        "fine-tuned on expressive cat sounds, lazaro-ros-sep on Afro-Cuban "
        "percussion extracted from the Lázaro Ros discography, and supermetroid "
        "on Super Metroid (SNES) gameplay audio, for retro, game-style "
        "soundscapes. Generation runs in mono at the model's own rate. The "
        "result is resampled back to the input's rate and copied across its "
        "channels, so the output has the format it came in with but no stereo "
        "image. The code is MIT, but the "
        "pretrained weights are CC-BY-NC-SA-4.0, so results cannot be used "
        "commercially and must be shared alike. Repository: "
        "https://github.com/hugofloresgarcia/vampnet Paper: "
        "https://arxiv.org/abs/2307.04686"
    ),
    author="Hugo Flores Garcia, Prem Seetharaman, Rithesh Kumar, Bryan Pardo",
    tags=[Subcategory.INFILLING, SampleRate(44100), Channels(1), 'sound-effects', 'domain:music']
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
init_model_choice = open("DEFAULT_MODEL").read().strip()
MODEL_CHOICES = Interface.available_models()

# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the default model is built and placed there alone. Each weight set is about 2.4 GB,
# too much to hold all four together, so the others are only fetched here, and built inside the
# request that picks one. Their paths are kept, since load_finetuned would list the model
# repository over the network from inside that request's GPU reservation.
if __name__ != "__main__":
    from vampnet import download_default, download_finetuned

    interface = Interface.default()
    interface.load_finetuned(init_model_choice)
    interface.to(DEVICE)
    CHECKPOINTS = {
        name: download_default() if name == "default" else download_finetuned(name)
        for name in MODEL_CHOICES
    }
    print(f"Loaded initial model: {init_model_choice}")


from torch_pitch_shift import pitch_shift, get_fast_shifts
def shift_pitch(signal, interval: int):
    signal.samples = pitch_shift(
        signal.samples,
        shift=interval,
        sample_rate=signal.sample_rate
    )
    return signal


def _vamp_internal(
        seed, input_audio, model_choice,
        pitch_shift_amt, periodic_p,
        n_mask_codebooks, onset_mask_width,
        dropout, sampletemp, typical_filtering,
        typical_mass, typical_min_tokens, top_p,
        sample_cutoff, stretch_factor, sampling_steps, beat_mask_ms, num_feedback_steps
    ):
    t0 = time.time()
    _seed = seed if seed > 0 else None
    if _seed is None:
        _seed = int(torch.randint(0, 2**32, (1,)).item())
    at.util.seed(_seed)

    if input_audio is None:
        raise gr.Error("no input audio received!")
    sr, input_audio = input_audio
    input_audio = input_audio / np.iinfo(input_audio.dtype).max

    sig = at.AudioSignal(input_audio, sr).to_mono()

    loudness = sig.loudness()
    sig = interface._preprocess(sig)

    # Builds the chosen model's checkpoints, unless they are the ones already built
    interface.reload(*CHECKPOINTS[model_choice])

    if pitch_shift_amt != 0:
        sig = shift_pitch(sig, pitch_shift_amt)

    codes = interface.encode(sig)

    mask = interface.build_mask(
        codes,
        sig=sig,
        periodic_prompt=periodic_p,
        onset_mask_width=onset_mask_width,
        _dropout=dropout,
        upper_codebook_mask=n_mask_codebooks,
    )
    if beat_mask_ms > 0:
        mask = pmask.mask_and(
            mask, interface.make_beat_mask(
                sig, after_beat_s=beat_mask_ms/1000.,
            )
        )
        mask = pmask.codebook_mask(mask, n_mask_codebooks)

    interface.set_chunk_size(10.0)

    # lord help me
    if top_p is not None:
        if top_p > 0:
            pass
        else:
            top_p = None

    codes, mask_z = interface.vamp(
        codes, mask,
        batch_size=2,
        feedback_steps=num_feedback_steps,
        _sampling_steps=sampling_steps,
        time_stretch_factor=stretch_factor,
        return_mask=True,
        temperature=sampletemp,
        typical_filtering=typical_filtering,
        typical_mass=typical_mass,
        typical_min_tokens=typical_min_tokens,
        top_p=top_p,
        seed=_seed,
        sample_cutoff=sample_cutoff,
    )
    print(f"vamp took {time.time() - t0} seconds")

    sig = interface.decode(codes)
    sig = sig.normalize(loudness)

    return sig


# Generation covers the whole input, stretched, in 10 s chunks, so the cost grows with the
# output's 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_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, model_choice,
                stretch_factor):
    # Measured at about 15 s for 30 s of output and 19 s for 64 s, padded, with room for loading
    # a fine-tuned model
    return 15 + 0.4 * sf.info(input_audio).duration * max(1, int(stretch_factor))


def process_fn(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, model_choice,
               stretch_factor):
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if model_choice not in MODEL_CHOICES:
        raise gr.Error(f"Unknown model: {model_choice}")
    if not input_audio:
        raise gr.Error("Provide the audio to rework.")
    try:
        duration = sf.info(input_audio).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.")
    output_seconds = duration * max(1, int(stretch_factor))
    if gpu_seconds(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, model_choice,
                   stretch_factor) > MAX_GPU_SECONDS:
        raise gr.Error(
            f"That would generate {output_seconds:.0f} s of audio, more than one request can in "
            f"the {MAX_GPU_SECONDS // 60} minutes it may run. Use a shorter input or a smaller "
            "stretch factor."
        )

    return harp_vamp(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks,
                     model_choice, stretch_factor)


@spaces.GPU(duration=gpu_seconds)
def harp_vamp(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, model_choice, stretch_factor):
    sig = at.AudioSignal(input_audio)

    # Captured before the downmix, because generation runs in mono at VampNet's
    # own rate and both have to be put back afterwards
    original_sr = sig.sample_rate
    original_channels = sig.num_channels

    sig = sig.to_mono()

    input_audio = sig.cpu().detach().numpy()[0][0]
    input_audio = input_audio * np.iinfo(np.int16).max
    input_audio = input_audio.astype(np.int16)
    input_audio = input_audio.reshape(1, -1)
    input_audio = (sig.sample_rate, input_audio)

    sig =  _vamp_internal(
        seed=0,
        input_audio=input_audio,
        model_choice=model_choice,
        pitch_shift_amt=0,
        periodic_p=int(periodic_p),
        n_mask_codebooks=int(n_mask_codebooks),
        onset_mask_width=0,
        dropout=dropout,
        sampletemp=sampletemp,
        typical_filtering=False,
        typical_mass=0.15,
        typical_min_tokens=1,
        top_p=None,
        sample_cutoff=1.0,
        stretch_factor=stretch_factor,
        sampling_steps=36,
        beat_mask_ms=int(0),
        num_feedback_steps=1,
    )

    sig = sig.detach().cpu()

    if sig.sample_rate != original_sr:
        sig = sig.resample(original_sr)

    # Copies of one generated channel, not a stereo image, which the card says
    if original_channels > 1 and sig.num_channels == 1:
        sig.audio_data = sig.audio_data.repeat(1, original_channels, 1)

    return save_audio(sig)


if __name__ == "__main__":
    with gr.Blocks() as demo:
        harp_input_components = [
            gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
                "The audio to rework. The model masks out parts of this recording "
                "and fills them back in, so the result keeps its character."
            )),
            gr.Slider(label="Randomness", minimum=0.1, maximum=10.0, value=1.0, step=0.001,
                  info="Controls sampling diversity- higher values yield more novel variations."),
            gr.Slider(label="Rhythmic Reference", minimum=0, maximum=13, step=1, value=7,
                      info="Determines how frequently rhythmic cues are repeated in generation."),
            gr.Slider(label="Variation Amount", minimum=0.0, maximum=1.0, step=0.01, value=0.0,
                      info="How much of the input the model changes. 0 = close copy, higher = greater remix."),
            gr.Slider(label="Timbral Reference", value=3, minimum=1, maximum=14, step=1,
                      info="Controls the richness of timbre retained from the source audio. Lower = simpler, higher = richer timbre."),
            gr.Dropdown(label="Model Choice", choices=list(MODEL_CHOICES), value=init_model_choice,
                        info="Select which fine-tuned VampNet model to use."),
            gr.Slider(label="Time Stretch Factor", minimum=1, maximum=8, step=1, value=1,
                      info="Scales the output's length relative to the input. 1 = the same length, 2 = twice as long."),
        ]

        harp_output_components = [
            gr.Audio(type="filepath", label="Generated Audio").set_info((
                "The reworked audio, at the same sample rate as the input."
            )),
        ]

        build_endpoint(
            model_card=model_card,
            input_components=harp_input_components,
            output_components=harp_output_components,
            process_fn=process_fn,
        )

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