from __future__ import annotations

import gradio as gr

from pyharp import *


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):
            kwargs.setdefault("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 soundfile as sf
import pyloudnorm as pyln
import numpy as np


model_card = ModelCard(
    name="pyloudnorm",
    description="Flexible audio loudness meter in Python. Implementation of ITU-R BS.1770-4 for integrated loudness and EBU Tech 3342 for loudness range.",
    author="csteinmetz1",
    tags=["audio-analysis", "loudness", "metering"],
)


def process_fn(input_audio, analysis_type, filter_class, block_size):
    data, rate = sf.read(input_audio)

    meter = pyln.Meter(rate, filter_class=filter_class, block_size=block_size)

    results = {}
    if analysis_type == "Integrated Loudness (LUFS)":
        loudness = meter.integrated_loudness(data)
        results["integrated_loudness_lufs"] = float(loudness)
    elif analysis_type == "Loudness Range (LU)":
        lra = meter.loudness_range(data)
        results["loudness_range_lu"] = float(lra)

    return results


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info("Upload an audio file to analyze."),
        gr.Dropdown(choices=["Integrated Loudness (LUFS)", "Loudness Range (LU)"], value="Integrated Loudness (LUFS)", label="Analysis Type", info="Choose between measuring integrated loudness (ITU-R BS.1770-4) or loudness range (EBU Tech 3342)."),
        gr.Dropdown(choices=["K-weighting", "Fenton/Lee 1", "Fenton/Lee 2", "Dash et al.", "DeMan"], value="K-weighting", label="Weighting Filter Class", info="Class of weighting filter used for loudness measurement. 'K-weighting' is the default for ITU-R BS.1770-4."),
        gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.4, label="Gating Block Size (seconds)", info="The duration of the gating block in seconds. Standard is 0.400s (400ms). The input audio length must be greater than this value."),
    ]
    output_components = [
        gr.JSON(label="Analysis Results"),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

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