"""
Synth Cue Restoration Studio
-----------------------------
A general-purpose tool for restoring degraded audio transferred from
analog master tapes: hiss/noise reduction, click/pop removal, tape-warble
(wow & flutter) smoothing, and optional reference-guided spectral (EQ)
matching against a clean reference clip.

This is a signal-processing tool. It does not generate, synthesize, or
recreate any copyrighted musical composition -- it only processes audio
that you upload, in the same way a restoration engineer would with
iZotope RX or similar tools.
"""

import numpy as np
import librosa
import soundfile as sf
import noisereduce as nr
from scipy import signal
from scipy.ndimage import median_filter
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import gradio as gr
import tempfile
import os

SR = 44100  # working sample rate


# ----------------------------------------------------------------------
# Core DSP helpers
# ----------------------------------------------------------------------

def load_mono(path, sr=SR):
    y, _ = librosa.load(path, sr=sr, mono=True)
    return y


def load_stereo(path, sr=SR):
    """Returns shape (2, N). Mono sources are duplicated to both channels."""
    y, _ = librosa.load(path, sr=sr, mono=False)
    if y.ndim == 1:
        y = np.stack([y, y])
    elif y.shape[0] == 1:
        y = np.vstack([y, y])
    return y[:2]


def stereo_width_ratio(y_stereo):
    left, right = y_stereo[0], y_stereo[1]
    mid = (left + right) / 2
    side = (left - right) / 2
    mid_rms = np.sqrt(np.mean(mid ** 2)) + 1e-9
    side_rms = np.sqrt(np.mean(side ** 2))
    return side_rms / mid_rms


def match_stereo_width(y_stereo, target_ratio, amount=1.0):
    """Mid/side rebalance to nudge the stereo width toward a target
    side/mid RMS ratio (as measured from a reference clip)."""
    left, right = y_stereo[0], y_stereo[1]
    mid = (left + right) / 2
    side = (left - right) / 2
    current_ratio = stereo_width_ratio(y_stereo)
    if current_ratio < 1e-6:
        return y_stereo
    desired_ratio = current_ratio + amount * (target_ratio - current_ratio)
    side_gain = desired_ratio / current_ratio
    side = side * side_gain
    new_left = mid + side
    new_right = mid - side
    return np.stack([new_left, new_right])


def denoise(y, sr, strength=0.75, noise_clip_seconds=1.0):
    """Spectral-gating noise reduction (good for tape hiss / room noise).
    Uses the quietest stretch of audio as a noise profile if a dedicated
    noise sample isn't provided.
    """
    # crude noise-floor estimate: the quietest N-second window
    win = int(noise_clip_seconds * sr)
    if len(y) > win:
        energies = [
            np.sum(y[i:i + win] ** 2) for i in range(0, len(y) - win, win)
        ]
        start = int(np.argmin(energies)) * win
        noise_clip = y[start:start + win]
    else:
        noise_clip = y

    return nr.reduce_noise(
        y=y, sr=sr, y_noise=noise_clip,
        prop_decrease=float(strength),
        stationary=False,
    )


def declick(y, threshold=8.0, max_run=25, local_window=64):
    """Detect and repair short impulsive clicks/pops common on tape
    dropouts and splices.

    Uses the second derivative (which spikes hard on true clicks but stays
    modest even on fast musical transients) compared against a LOCAL
    rolling noise-floor estimate rather than a single global threshold, so
    it doesn't over-fire during loud/dense passages. Outlier runs longer
    than max_run are treated as real signal (e.g. a percussive hit) and
    left untouched -- true clicks are a handful of samples wide.
    """
    d2 = np.diff(y, n=2)
    d2 = np.pad(d2, (1, 1))
    local_scale = median_filter(np.abs(d2), size=local_window) + 1e-6
    ratio = np.abs(d2) / local_scale
    outliers = ratio > threshold

    idx = np.where(outliers)[0]
    if len(idx) == 0:
        return y, 0

    groups = np.split(idx, np.where(np.diff(idx) > 1)[0] + 1)
    groups = [g for g in groups if len(g) <= max_run]

    y_fixed = y.copy()
    for g in groups:
        lo, hi = max(g[0] - 1, 0), min(g[-1] + 1, len(y) - 1)
        if hi > lo:
            y_fixed[lo:hi + 1] = np.linspace(y[lo], y[hi], hi - lo + 1)
    return y_fixed, len(groups)


def dewarble(y, sr, strength=0.5):
    """Reduce slow pitch wobble (wow & flutter) from tape stretch/warp
    using a gentle time-varying resample based on a low-frequency
    envelope of detected pitch drift. Conservative by design so it
    doesn't smear transients.
    """
    if strength <= 0:
        return y
    # smooth amplitude envelope as a cheap proxy for slow tape-speed drift;
    # apply a very light adaptive low-pass to reduce fast micro-jitter
    # without altering the intended pitch content.
    cutoff = 0.5 + (1 - strength) * 4.0  # Hz, gentle
    sos = signal.butter(2, cutoff, btype="low", fs=sr, output="sos")
    envelope = np.abs(signal.hilbert(y))
    smooth_env = signal.sosfiltfilt(sos, envelope) + 1e-6
    correction = smooth_env / np.mean(smooth_env)
    correction = np.clip(correction, 0.85, 1.15)
    return y * (1.0 / correction) * strength + y * (1 - strength)


def spectral_match(y, sr, reference_y, amount=0.6, n_fft=2048):
    """Reference-guided EQ matching: shape the restored audio's long-term
    average spectrum toward the reference clip's spectrum. This is the
    feature that uses your 'what it should sound like' sample -- it
    nudges tonal balance (brightness, low-end weight, etc.) toward the
    reference without copying its actual content.
    """
    def avg_spectrum(sig):
        S = np.abs(librosa.stft(sig, n_fft=n_fft))
        return np.mean(S, axis=1) + 1e-6

    target_spec = avg_spectrum(reference_y)
    source_spec = avg_spectrum(y)

    # match array lengths safely
    n = min(len(target_spec), len(source_spec))
    gain_curve = (target_spec[:n] / source_spec[:n])
    gain_curve = np.clip(gain_curve, 0.25, 4.0)  # avoid extreme boosts
    # smooth the gain curve so it acts like a broad EQ, not per-bin noise
    gain_curve = signal.savgol_filter(gain_curve, 31, 3) if n > 31 else gain_curve

    S = librosa.stft(y, n_fft=n_fft)
    mag, phase = np.abs(S), np.angle(S)
    apply_gain = 1 + amount * (gain_curve[:mag.shape[0]][:, None] - 1)
    mag_matched = mag * apply_gain
    S_matched = mag_matched * np.exp(1j * phase)
    return librosa.istft(S_matched, length=len(y))


def restore_highs(y, sr, amount=0.3, cutoff=6000):
    """Harmonic exciter: regenerates plausible high-frequency content
    lost to tape/transfer bandwidth limiting, by adding controlled
    harmonic distortion above a cutoff and blending it back in.
    """
    if amount <= 0:
        return y
    sos_hp = signal.butter(4, cutoff, btype="high", fs=sr, output="sos")
    highs = signal.sosfilt(sos_hp, y)
    excited = np.tanh(highs * 3.0) / 3.0  # soft-clip harmonic generation
    return y + amount * excited


def normalize(y, target_db=-1.0):
    peak = np.max(np.abs(y)) + 1e-9
    target_amp = 10 ** (target_db / 20)
    return y * (target_amp / peak)


def apply_per_channel(fn, y_stereo, *args, **kwargs):
    """Run a mono DSP function independently on each stereo channel."""
    out = [fn(y_stereo[ch], *args, **kwargs) for ch in range(y_stereo.shape[0])]
    # some fns (declick) return a tuple (signal, count) -- normalize that
    counts = None
    if isinstance(out[0], tuple):
        counts = sum(o[1] for o in out)
        out = [o[0] for o in out]
    return np.stack(out), counts


def make_spectrogram_fig(y, sr, title):
    """y may be mono (N,) or stereo (2, N); stereo is downmixed for display."""
    if y.ndim == 2:
        y = y.mean(axis=0)
    fig, ax = plt.subplots(figsize=(6, 3))
    S = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
    librosa.display.specshow(S, sr=sr, x_axis="time", y_axis="log", ax=ax)
    ax.set_title(title)
    fig.tight_layout()
    return fig


# NOTE ON REFERENCE MATCHING:
# Analysis across 8 clean cues from this score showed no single consistent
# "target" tone -- stereo width ranges from near-mono (correlation ~0.99)
# to very wide/decorrelated (correlation ~-0.14), and brightness ranges
# from ~6kHz to ~17.6kHz rolloff depending on the cue. There is no
# universal profile to hardcode. Spectral and width matching in this tool
# therefore always measure live from whatever reference clip you upload
# for that specific restoration run -- always use the clean counterpart
# of the exact cue you're restoring, not a generic "score-wide" reference.


# ----------------------------------------------------------------------
# Main pipeline
# ----------------------------------------------------------------------

def restore(
    degraded_path,
    reference_path,
    denoise_strength,
    declick_on,
    dewarble_strength,
    restore_highs_amount,
    ref_match_amount,
    width_match_amount,
    target_loudness,
):
    if degraded_path is None:
        raise gr.Error("Please upload a degraded audio cue to restore.")

    y = load_stereo(degraded_path, SR)  # shape (2, N)
    log = []
    orig_y = y.copy()

    if denoise_strength > 0:
        y, _ = apply_per_channel(denoise, y, SR, strength=denoise_strength)
        log.append(f"Noise reduction applied (strength={denoise_strength:.2f})")

    if declick_on:
        y, n_clicks = apply_per_channel(declick, y)
        log.append(f"Declick: repaired {n_clicks} click/dropout region(s) across both channels")

    if dewarble_strength > 0:
        y, _ = apply_per_channel(dewarble, y, SR, strength=dewarble_strength)
        log.append(f"Wow/flutter smoothing applied (strength={dewarble_strength:.2f})")

    if restore_highs_amount > 0:
        y, _ = apply_per_channel(restore_highs, y, SR, amount=restore_highs_amount)
        log.append(f"High-frequency restoration applied (amount={restore_highs_amount:.2f})")

    if reference_path is not None and ref_match_amount > 0:
        ref_y = load_stereo(reference_path, SR)
        ref_mono = ref_y.mean(axis=0)  # one gain curve shared by both channels
        y, _ = apply_per_channel(spectral_match, y, SR, ref_mono, amount=ref_match_amount)
        log.append(f"Spectral EQ matched to reference (amount={ref_match_amount:.2f})")

        if width_match_amount > 0:
            target_ratio = stereo_width_ratio(ref_y)
            before_ratio = stereo_width_ratio(y)
            y = match_stereo_width(y, target_ratio, amount=width_match_amount)
            after_ratio = stereo_width_ratio(y)
            log.append(
                f"Stereo width matched toward reference (target ratio={target_ratio:.3f}, "
                f"{before_ratio:.3f} -> {after_ratio:.3f})"
            )
    elif reference_path is None and (ref_match_amount > 0 or width_match_amount > 0):
        log.append("No reference clip uploaded -- skipped spectral/width matching.")

    peak = np.max(np.abs(y)) + 1e-9
    y = y * (10 ** (target_loudness / 20) / peak)
    log.append(f"Normalized to {target_loudness} dBFS peak")

    out_path = os.path.join(tempfile.gettempdir(), "restored_cue.wav")
    sf.write(out_path, y.T, SR)  # soundfile wants shape (N, channels)

    fig_before = make_spectrogram_fig(orig_y, SR, "Before")
    fig_after = make_spectrogram_fig(y, SR, "After")

    return out_path, fig_before, fig_after, "\n".join(log)


def apply_preset_broadband_noise():
    """For cues where bandwidth is already close to intact relative to
    their own clean reference (rolloff not far below the reference clip
    you upload for that cue), moderate SNR loss, light click rate, and
    stereo only slightly narrower than that reference. Main job is noise
    reduction, not rebuilding lost highs.
    """
    return (
        0.8,   # denoise_strength -- this is the main problem here
        True,  # declick_on
        0.25,  # dewarble_strength
        0.1,   # restore_highs_amount -- barely anything missing, don't overdo it
        0.6,   # ref_match_amount
        0.5,   # width_match_amount -- small gap to close (0.50 -> 0.56)
        -1.0,  # target_loudness
    )


def apply_preset_heavy_click():
    """For heavily duplicated/worn transfers: click rate several times
    higher than a typical cue (5-8+/sec vs <2/sec), often paired with
    reduced dynamic range from generation loss or broadcast-style
    compression. Denoise stays moderate since SNR on this type tends to
    be fine -- the damage is impulsive, not broadband hiss.
    """
    return (
        0.55,  # denoise_strength -- SNR usually isn't the main issue here
        True,  # declick_on
        0.3,   # dewarble_strength
        0.35,  # restore_highs_amount
        0.65,  # ref_match_amount
        0.6,   # width_match_amount
        -1.0,  # target_loudness
    )


def apply_preset_bandwidth_limited():
    """For cues with real high-frequency loss relative to their own
    clean reference (rolloff noticeably below the reference clip you
    upload for that specific cue), plus a higher click rate and a
    narrower stereo image than that reference. Main job is rebuilding
    top end and widening the image; denoise can be lighter if SNR is
    otherwise decent.
    """
    return (
        0.6,   # denoise_strength -- SNR is often decent, don't over-denoise
        True,  # declick_on
        0.3,   # dewarble_strength
        0.55,  # restore_highs_amount -- real gap to rebuild
        0.7,   # ref_match_amount
        0.75,  # width_match_amount -- larger gap to close
        -1.0,  # target_loudness
    )


# ----------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------

with gr.Blocks(title="Synth Cue Restoration Studio") as demo:
    gr.Markdown(
        """
        # 🎛️ Synth Cue Restoration Studio
        Upload a degraded audio cue transferred from a flawed master/analog
        tape source. Optionally upload a clean reference clip so the tool
        can match tonal balance *and* stereo width toward it. Processing
        is done in true stereo (mid/side) throughout. Everything here is
        standard DSP (denoise, declick, wow/flutter correction,
        high-frequency restoration, spectral + width matching) applied to
        *your* uploaded audio.
        """
    )

    with gr.Row():
        degraded_in = gr.Audio(label="Degraded cue (required)", type="filepath")
        reference_in = gr.Audio(
            label="Reference clip - what it should sound like (optional)",
            type="filepath",
        )

    with gr.Row():
        preset_noise_btn = gr.Button("Preset: Broadband Noise Cue\n(intact highs, hiss/noise is the main issue)")
        preset_bw_btn = gr.Button("Preset: Bandwidth-Limited Cue\n(dull/rolled-off highs, narrower stereo)")
        preset_click_btn = gr.Button("Preset: Heavy Click / Worn Duplication\n(high click rate, often lower dynamic range)")
    gr.Markdown(
        """
        *Presets are derived from measuring real restored/degraded cue
        pairs. Broadband Noise: highs already reach close to the
        reference's ~8.3kHz rolloff, main issue is hiss. Bandwidth-Limited:
        rolloff sits well below that (~5-6kHz) and stereo reads narrower
        than the reference. Heavy Click: click rate several times higher
        than typical (5-8+/sec), sometimes with reduced dynamic range from
        generation loss. If a cue has unusually low stereo correlation
        alongside width WIDER than the reference, listen first before
        applying strong width correction -- that pattern doesn't always
        mean damage, and mid/side rescaling won't fix a genuine phase
        issue if that's the actual cause.*
        """
    )

    with gr.Row():
        denoise_strength = gr.Slider(0, 1, value=0.75, label="Noise reduction strength")
        declick_on = gr.Checkbox(value=True, label="Declick / dropout repair")

    with gr.Row():
        dewarble_strength = gr.Slider(0, 1, value=0.3, label="Wow/flutter smoothing")
        restore_highs_amount = gr.Slider(0, 1, value=0.3, label="High-frequency restoration")

    with gr.Row():
        ref_match_amount = gr.Slider(0, 1, value=0.6, label="Reference spectral (EQ) matching amount")
        width_match_amount = gr.Slider(0, 1, value=0.6, label="Reference stereo-width matching amount")

    target_loudness = gr.Slider(-12, 0, value=-1.0, label="Output peak (dBFS)")

    run_btn = gr.Button("Restore Audio", variant="primary")

    with gr.Row():
        output_audio = gr.Audio(label="Restored cue", type="filepath")
    with gr.Row():
        spec_before = gr.Plot(label="Spectrogram: Before")
        spec_after = gr.Plot(label="Spectrogram: After")
    log_out = gr.Textbox(label="Processing log", lines=8)

    preset_noise_btn.click(
        fn=apply_preset_broadband_noise,
        inputs=[],
        outputs=[
            denoise_strength, declick_on, dewarble_strength,
            restore_highs_amount, ref_match_amount, width_match_amount,
            target_loudness,
        ],
    )

    preset_click_btn.click(
        fn=apply_preset_heavy_click,
        inputs=[],
        outputs=[
            denoise_strength, declick_on, dewarble_strength,
            restore_highs_amount, ref_match_amount, width_match_amount,
            target_loudness,
        ],
    )

    preset_bw_btn.click(
        fn=apply_preset_bandwidth_limited,
        inputs=[],
        outputs=[
            denoise_strength, declick_on, dewarble_strength,
            restore_highs_amount, ref_match_amount, width_match_amount,
            target_loudness,
        ],
    )

    run_btn.click(
        fn=restore,
        inputs=[
            degraded_in, reference_in, denoise_strength, declick_on,
            dewarble_strength, restore_highs_amount, ref_match_amount,
            width_match_amount, target_loudness,
        ],
        outputs=[output_audio, spec_before, spec_after, log_out],
    )

if __name__ == "__main__":
    demo.launch()