try:
    import spaces
except ImportError:
    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
    class spaces:
        class GPU:
            def __init__(self, func=None, duration=60):
                self.func = func

            def __call__(self, *args, **kwargs):
                if self.func is not None:
                    return self.func(*args, **kwargs)
                func = args[0]
                return func

import sys
sys.stdout.reconfigure(line_buffering=True)

import tempfile
import threading

import gradio as gr
import numpy as np
import soundfile as sf
import torch
import torchaudio
from huggingface_hub import hf_hub_download
from omegaconf import OmegaConf
from pyharp import ModelCard, build_endpoint

# cqt_nsgt_pytorch clips an int-dtype array against np.inf with out=, which
# numpy's ufunc casting rules reject on any numpy version. Compute without
# `out` and write the result back instead of failing outright.
_np_clip = np.clip


def _clip_int_out_safe(a, a_min, a_max, out=None, **kwargs):
    if out is not None:
        out[...] = _np_clip(a, a_min, a_max).astype(out.dtype)
        return out
    return _np_clip(a, a_min, a_max, **kwargs)


np.clip = _clip_int_out_safe

from diff_params.edm import EDM
from networks.unet_cqt_oct_with_projattention_adaLN_2 import Unet_CQT_oct_with_attention
from testing.edm_sampler_inpainting import Sampler

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CHECKPOINT_REPO = "Eloimoliner/audio-inpainting-diffusion"

# Same network configs as conf/network/paper_1912_unet_cqt_oct_attention_adaLN_2.yaml
# and conf/network/paper_1912_unet_cqt_oct_attention_44k_2.yaml.
CHECKPOINTS = {
    "MAESTRO (solo piano)": {
        "filename": "maestro_22k_8s-750000.pt",
        "sample_rate": 22050,
        "audio_len": 184184,
        "network": {
            "use_fencoding": False,
            "use_norm": True,
            "depth": 7,
            "emb_dim": 256,
            "Ns": [64, 96, 96, 128, 128, 256, 256],
            "attention_layers": [0, 0, 0, 0, 1, 1, 1, 1],
            "Ss": [2, 2, 2, 2, 2, 2, 2],
            "num_dils": [2, 3, 4, 5, 6, 7, 7],
            "cqt": {"window": "kaiser", "beta": 1, "num_octs": 7, "bins_per_oct": 64},
            "bottleneck_type": "res_dil_convs",
            "num_bottleneck_layers": 1,
            "attention_dict": {
                "num_heads": 8, "attn_dropout": 0.0, "bias_qkv": False, "N": 0,
                "rel_pos_num_buckets": 32, "rel_pos_max_distance": 64,
                "use_rel_pos": False, "Nproj": 8,
            },
        },
    },
    "MusicNet (mixed classical instrumentation)": {
        "filename": "musicnet_44k_4s-560000.pt",
        "sample_rate": 44100,
        "audio_len": 184184,
        "network": {
            "use_fencoding": False,
            "use_norm": True,
            "depth": 8,
            "emb_dim": 256,
            "Ns": [64, 64, 96, 96, 128, 128, 256, 256],
            "attention_layers": [0, 0, 0, 0, 0, 1, 1, 1, 1],
            "Ss": [2, 2, 2, 2, 2, 2, 2],
            "num_dils": [2, 3, 4, 5, 6, 7, 8, 8],
            "cqt": {"window": "kaiser", "beta": 1, "num_octs": 8, "bins_per_oct": 64},
            "bottleneck_type": "res_dil_convs",
            "num_bottleneck_layers": 1,
            "attention_dict": {
                "num_heads": 8, "attn_dropout": 0.0, "bias_qkv": False, "N": 0,
                "rel_pos_num_buckets": 32, "rel_pos_max_distance": 64,
                "use_rel_pos": False, "Nproj": 8,
            },
        },
    },
}

# conf/diff_params/edm.yaml -- the training-time noise schedule, read by EDM.__init__.
BASE_DIFF_PARAMS = {
    "sigma_data": 0.063,
    "sigma_min": 1e-5,
    "sigma_max": 10,
    "P_mean": -1.2,
    "P_std": 1.2,
    "ro": 13,
    "ro_train": 10,
    "Schurn": 5,
    "Snoise": 1,
    "Stmin": 0,
    "Stmax": 50,
    "aweighting": {"use_aweighting": False, "ntaps": 101},
}

# conf/tester/inpainting_tester.yaml -- sampling-time overrides, read by Sampler.
# This is the same config the repo's own notebooks/demo_inpainting_spectrogram.ipynb
# uses for both checkpoints (it only swaps exp= and network=).
TESTER_DIFF_PARAMS = {
    "same_as_training": False,
    "sigma_data": 0.063,
    "sigma_min": 1e-4,
    "sigma_max": 1,
    "P_mean": -1.2,
    "P_std": 1.2,
    "ro": 13,
    "ro_train": 13,
    "Schurn": 10,
    "Snoise": 1.0,
    "Stmin": 0,
    "Stmax": 50,
}

_model_cache = {}
_model_locks = {name: threading.Lock() for name in CHECKPOINTS}


def build_args(checkpoint_name, steps, guidance_strength):
    """Assemble the config tree the network, EDM and Sampler classes expect, mirroring the repo's hydra config composition for one checkpoint."""
    ckpt_cfg = CHECKPOINTS[checkpoint_name]
    return OmegaConf.create({
        "exp": {
            "sample_rate": ckpt_cfg["sample_rate"],
            "audio_len": ckpt_cfg["audio_len"],
        },
        "network": ckpt_cfg["network"],
        "diff_params": BASE_DIFF_PARAMS,
        "tester": {
            "T": steps,
            "order": 2,
            "filter_out_cqt_DC_Nyq": True,
            "posterior_sampling": {"xi": guidance_strength, "norm": 2, "smoothl1_beta": 1},
            "data_consistency": {"use": True, "type": "always", "smooth": True, "hann_size": 50},
            "diff_params": TESTER_DIFF_PARAMS,
        },
    })


def get_network(checkpoint_name):
    """Build the network directly on DEVICE and load its checkpoint, caching per checkpoint name.

    Building fresh on DEVICE (rather than building on CPU and calling .to(DEVICE)
    later) matters here: the network's CQT transform allocates its kernels for a
    specific device at construction time, so a later .to() would leave it stale.
    """
    with _model_locks[checkpoint_name]:
        if checkpoint_name not in _model_cache:
            ckpt_cfg = CHECKPOINTS[checkpoint_name]
            args = build_args(checkpoint_name, steps=35, guidance_strength=0.25)  # unused by the network
            network = Unet_CQT_oct_with_attention(args, DEVICE).to(DEVICE)
            ckpt_path = hf_hub_download(repo_id=CHECKPOINT_REPO, filename=ckpt_cfg["filename"])
            state_dict = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
            network.load_state_dict(state_dict["ema"])
            network.eval()
            _model_cache[checkpoint_name] = network
    return _model_cache[checkpoint_name]


def prepare_segment(audio, sr_in, target_sr, target_len):
    """Resample to the checkpoint's rate, then center-crop or zero-pad to its fixed segment length.

    Adapted from Tester.resample_audio in testing/tester_inpainting.py, extended
    to pad shorter clips: the repo only ever sees pre-cropped test-set segments,
    but an interactive demo can be handed audio of any length.
    """
    if sr_in != target_sr:
        audio = torchaudio.functional.resample(audio, sr_in, target_sr)
    n = audio.shape[-1]
    if n > target_len:
        start = (n - target_len) // 2
        audio = audio[..., start:start + target_len]
    elif n < target_len:
        audio = torch.nn.functional.pad(audio, (0, target_len - n))
    return audio


def build_gap_mask(audio_len, gap_ms, position_ms, sample_rate, device):
    """Binary mask with a single gap, matching Tester.prepare_mask's 'long' mode. position_ms of -1 centers the gap, per the repo's own 'start_gap_idx: None' convention."""
    gap = int(gap_ms * sample_rate / 1000)
    gap = max(1, min(gap, audio_len - 1))
    if position_ms < 0:
        start = audio_len // 2 - gap // 2
    else:
        start = int(position_ms * sample_rate / 1000)
        start = max(0, min(start, audio_len - gap))
    mask = torch.ones((1, audio_len), device=device)
    mask[..., start:start + gap] = 0
    return mask


@spaces.GPU
def process_fn(input_audio_path, checkpoint_choice, gap_length_s, gap_position_s, steps, guidance_strength):
    """Fill a gap in the input audio using the selected diffusion model.

    Not wrapped in torch.inference_mode(): reconstruction guidance (xi > 0,
    the default) needs torch.autograd.grad() inside the sampler, which
    inference_mode tensors can never support. The sampler applies its own
    torch.no_grad() around the branches that don't need gradients.
    """
    ckpt_cfg = CHECKPOINTS[checkpoint_choice]
    sample_rate = ckpt_cfg["sample_rate"]
    audio_len = ckpt_cfg["audio_len"]

    audio_np, sr_in = sf.read(input_audio_path, dtype="float32", always_2d=False)
    if audio_np.ndim > 1:
        audio_np = audio_np.mean(axis=-1)  # the model is mono-only
    audio = torch.from_numpy(audio_np).unsqueeze(0).to(DEVICE)

    segment = prepare_segment(audio, sr_in, sample_rate, audio_len)
    gap_position_ms = -1 if gap_position_s < 0 else gap_position_s * 1000
    mask = build_gap_mask(audio_len, gap_length_s * 1000, gap_position_ms, sample_rate, DEVICE)

    network = get_network(checkpoint_choice)
    args = build_args(checkpoint_choice, steps, guidance_strength)
    diff_params = EDM(args)
    sampler = Sampler(network, diff_params, args)

    y_masked = segment * mask
    pred = sampler.predict_inpainting(y_masked, mask)

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        out_path = f.name
    sf.write(out_path, pred.squeeze(0).cpu().numpy(), sample_rate)
    return out_path


model_card = ModelCard(
    name="Diffusion-Based Audio Inpainting",
    description=(
        "Fills a gap in an audio clip using a diffusion model trained on either solo "
        "piano (MAESTRO) or mixed classical instrumentation (MusicNet). The input is "
        "resampled to the chosen model's rate and center-cropped or zero-padded to its "
        "fixed segment length (about 8.4s for MAESTRO, 4.2s for MusicNet) before the gap "
        "is filled."
    ),
    author="Eloi Moliner, Vesa Välimäki",
    tags=["audio-inpainting", "diffusion", "restoration"],
)


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
        gr.Dropdown(
            choices=list(CHECKPOINTS.keys()),
            value="MAESTRO (solo piano)",
            label="Model",
            info="Which checkpoint to inpaint with.",
        ),
        gr.Slider(minimum=0.05, maximum=3.0, step=0.05, value=1.5, label="Gap Length (s)",
                  info="Length of the gap to fill (default: 1.5s, per repo config)"),
        gr.Number(value=-1, label="Gap Position (s)",
                  info="Where the gap starts, from the beginning of the processed segment. -1 = centered (default, per repo config)"),
        gr.Slider(minimum=10, maximum=150, step=1, value=35, label="Processing Steps",
                  info="Diffusion sampling steps (default: T=35, per repo config)"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.25, label="Reconstruction Guidance",
                  info="How strongly the output is pulled toward the surrounding audio (default: ξ=0.25, per repo config)"),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Output Audio").set_info(
            "The processed segment with the gap filled in."
        ),
    ]

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

if __name__ == "__main__":
    demo.queue().launch(pwa=True)