import os
import tempfile
import torch
import soundfile as sf
import numpy as np
from pathlib import Path
from pyharp import ModelCard, build_endpoint, SampleRate, Channels
from pyharp.tags import Subcategory
import gradio as gr

from nam.train.lightning_module import LightningModule as _LightningModule
import nam.data as data

os.environ["TOKENIZERS_PARALLELISM"] = "true"

model = None
CKPT_PATH = Path(__file__).parent / "demo_ckpt.ckpt"
MODEL_CONFIG_PATH = Path(__file__).parent / "default_config_files/models/wavenet.json"
if not CKPT_PATH.exists():
    print("Downloading checkpoint...", flush=True)
    import urllib.request
    urllib.request.urlretrieve(
        "https://raw.githubusercontent.com/ETH-DISCO/PANAMA/main/demo_ckpt.ckpt",
        CKPT_PATH
    )
    print("Checkpoint downloaded.", flush=True)

def get_model():
    global model
    if model is None:
        print("Loading PANAMA model...", flush=True)
        import json
        with open(MODEL_CONFIG_PATH, "r") as fp:
            model_config = json.load(fp)
        model = _LightningModule.load_from_checkpoint(
            CKPT_PATH,
            **_LightningModule.parse_config(model_config)
        )
        model.eval()
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device)
        print("Model loaded.", flush=True)
    return model

model_card = ModelCard(
    name="PANAMA",
    description=(
        "Parametric neural amp modeling. Apply amp simulation with 6 "
        "controllable knobs. Amp modeling is mono: a stereo file is downmixed "
        "before the amp and the single result is copied back across its "
        "channels, so the output has the channel count it came in with but not "
        "its stereo image. The input's sample rate is restored. Repository: "
        "https://github.com/ETH-DISCO/PANAMA Paper: "
        "https://arxiv.org/abs/2509.26564"
    ),
    author="Florian Grötschla, Longxiang Jiao, Luca A. Lanzendörfer, Roger Wattenhofer",
    tags=[Subcategory.NEURAL_ANALOG_EFFECTS, SampleRate(48000), Channels(1), 'amp-modeling', 'guitar'],
)

def process_fn(input_audio_path: str, knob1: float, knob2: float, knob3: float, knob4: float, knob5: float, knob6: float) -> str:
    m = get_model()
    device = next(m.parameters()).device

    # nam's loader rejects anything but mono, and an amp model is mono by
    # nature, so the downmix happens here and the channel count is restored
    # after inference rather than refusing the file
    import librosa
    import soundfile

    with soundfile.SoundFile(input_audio_path) as handle:
        original_sr = handle.samplerate
        original_channels = handle.channels

    input_audio_np, _ = librosa.load(input_audio_path, sr=None, mono=True)

    if original_sr != m.net.sample_rate:
        input_audio_np = librosa.resample(
            input_audio_np, orig_sr=original_sr, target_sr=m.net.sample_rate
        )

    input_audio = torch.tensor(input_audio_np, dtype=torch.float32)

    input_audio = input_audio.to(device)
    g_tensor = torch.tensor([[knob1, knob2, knob3, knob4, knob5, knob6]], dtype=torch.float32).to(device)

    print(f"Processing with g-vector: {g_tensor}", flush=True)

    with torch.no_grad():
        output = m.forward(input_audio, g_tensor)

    output_np = output.cpu().numpy()

    # Back to the rate the session uses, before the channels are rebuilt, so
    # the resampler only ever sees one channel
    if original_sr != m.net.sample_rate:
        output_np = librosa.resample(
            output_np, orig_sr=m.net.sample_rate, target_sr=original_sr
        )

    # Two identical channels play back the same as a mono file on a stereo
    # track. It does not make the result stereo, which the model card says.
    if original_channels > 1:
        output_np = np.tile(output_np[:, None], (1, original_channels))

    output_path = tempfile.mktemp(suffix=".wav")
    sf.write(output_path, output_np, original_sr)
    return output_path

if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
                "A dry guitar or bass recording. Feeding it an already-amped tone "
                "stacks one amp on another."
            )),
            gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Gain",
            info="Preamp drive. Raise it for more distortion and compression."),
            gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Bass",
            info="Low-end tone control."),
            gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Mid",
            info=(
                "Midrange tone control. Cutting it gives the scooped sound heavy "
                "styles use."
            )),
            gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Treble",
            info="High-end tone control."),
            gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Master",
            info=(
                "Power amp level, which adds its own saturation as it rises rather "
                "than only getting louder."
            )),
            gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Presence",
            info="Upper-midrange bite, which helps a part cut through a mix."),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Output Audio").set_info("Amp-processed audio."),
        ]
        app = build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )

    print("Launching Gradio...", flush=True)
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True)