#!/usr/bin/env python3
"""Standalone vocal / instrumental separation with UVR MDX-Net on LiteRT.

This is a self-contained, runnable reference for the "STFT outside the graph" design used by the
MusicStemSeparation Android app. The `.tflite` graph is the **learned core only** — a convolutional
U-Net that maps a spectrogram to the vocal spectrogram. The STFT, chunking, overlap-add and inverse
STFT are ordinary CPU code here (numpy), exactly as they are ordinary Kotlin code in the app. Per chunk:

        STFT (numpy) -> inference (LiteRT) -> iSTFT (numpy) -> weighted overlap-add

Everything below mirrors the app's DSP line-for-line so the output matches; the Kotlin sources are cited
inline (`CacDsp.kt`, `MdxSeparationPipeline.kt`, `MdxLiteRtEngine.kt`, `MdxTypes.kt`, audio/*). The whole
pipeline is a faithful reimplementation of UVR's `mdx.py` demix (with a 10% crossfade instead of hard
tiling — set OVERLAP = 0 to get the classic hard concat).

THE MODEL I/O CONTRACT (the one thing to get right):
    input  = OUTPUT = float32 tensor  [1, 4, dim_f, 256]  (NCHW)
    the 4 channels are complex-as-channels, plane order  [L_re, L_im, R_re, R_im]
    flat row-major index = ((plane * dim_f) + bin) * 256 + frame
    9482 : dim_f 2048, n_fft 4096      Voc FT : dim_f 3072, n_fft 6144      both hop 1024
A wrong packing (swapped re/im planes, kept Nyquist bin, a stray 1/sqrt(n)) yields plausible-but-wrong
audio, not an error — so this is where to look first if a stem sounds off.

NOTE ON PRECISION / SPEED: this runs the model on the CPU via ai-edge-litert (XNNPACK). The `.tflite`
carries `reduced_precision_support=fp16accfp16`, which only engages fp16 kernels on ARMv8.2-FP16 cores;
on a desktop it runs fp32 (correct, just not the ~2x mobile speed). The real CPU->GPU->NPU accelerator
ladder is Android-only — see the README's "Android on-device inference" section and the app's
`MdxLiteRtEngine.kt`.

Usage:
    python separate.py song.wav                      # 9482 model, writes song_vocals.wav + song_instrumental.wav
    python separate.py song.wav --model voc_ft       # higher-quality (larger) model
    python separate.py song.wav --out-dir out/ --denoise
Input must be a WAV (convert other formats first, e.g. `ffmpeg -i in.mp3 in.wav`). Feed 44.1 kHz for
exact UVR parity; other rates are linearly resampled (matching the app, not librosa/soxr).
"""
import argparse
import os
import sys
import wave

import numpy as np
from ai_edge_litert.interpreter import Interpreter

# ── Fixed constants (MdxTypes.kt) ────────────────────────────────────────────────────────────────
SAMPLE_RATE = 44100      # MDX_SAMPLE_RATE — MDX-Net operates at 44.1 kHz
DIM_T = 256              # native trained segment (~5.92 s @ hop 1024); static in the shipped .tflite
HOP = 1024               # STFT hop, both models
OVERLAP = 0.10           # MDX_OVERLAP — fixed 10% chunk crossfade (0 -> classic hard-tiled concat)

# The .tflite assets live beside this script. Each entry carries its STFT framing (MdxModel enum).
HERE = os.path.dirname(os.path.abspath(__file__))
MODELS = {
    "9482":   dict(tflite="UVR_MDXNET_9482.fp16acc.tflite",     n_fft=4096, dim_f=2048, compensation=1.0),
    "voc_ft": dict(tflite="UVR-MDX-NET-Voc_FT.fp16acc.tflite",  n_fft=6144, dim_f=3072, compensation=1.0),
}


# ══════════════════════════════════════════════════════════════════════════════════════════════════
#  Host DSP — a numpy port of CacDsp.kt (CacStft / CacIstft / reflect pad / crossfade window).
#  Mirrors UVR's STFT/ISTFT: periodic Hann, center=True reflect pad, normalized=False (no 1/sqrt(n)),
#  Nyquist bin dropped (dim_f = n_fft/2).
# ══════════════════════════════════════════════════════════════════════════════════════════════════

def hann_periodic(n):
    """torch.hann_window(n, periodic=True) = 0.5*(1 - cos(2*pi*k/n)), k=0..n-1 (divisor n)."""
    k = np.arange(n, dtype=np.float64)
    return (0.5 * (1.0 - np.cos(2.0 * np.pi * k / n))).astype(np.float32)


class Stft:
    """Forward STFT: stereo chunk (chunk_size samples/channel) -> CaC tensor [1, 4, dim_f, 256].

    Port of CacStft.transform. For each channel: reflect-pad by n_fft/2 (center=True), slide a
    hop-spaced periodic-Hann window over DIM_T frames, rfft each frame, keep bins [0, dim_f) (dropping
    the Nyquist bin since dim_f = n_fft/2), and store real/imag into planes [L_re, L_im, R_re, R_im].
    """

    def __init__(self, n_fft, dim_f):
        self.n = n_fft
        self.dim_f = dim_f
        self.pad = n_fft // 2                       # centerPad
        self.chunk_size = HOP * (DIM_T - 1)
        self.hann = hann_periodic(n_fft)

    def _channel_spec(self, samples):
        # reflectPad(x, left=pad, right=pad): mirror without repeating the edge == np.pad(mode="reflect").
        padded = np.pad(samples, (self.pad, self.pad), mode="reflect")
        # DIM_T frames starting at f*hop, each n_fft long (padded length = chunk_size + n_fft).
        frames = np.lib.stride_tricks.sliding_window_view(padded, self.n)[::HOP]
        assert frames.shape[0] == DIM_T, (frames.shape, DIM_T)
        frames = frames * self.hann                 # analysis window
        spec = np.fft.rfft(frames, axis=1)          # [DIM_T, n/2+1], unnormalized (normalized=False)
        spec = spec[:, : self.dim_f]                # crop to dim_f -> drops the Nyquist bin
        # -> [dim_f, DIM_T] planes; rfft already gives DC imag == 0.
        return spec.real.T.astype(np.float32), spec.imag.T.astype(np.float32)

    def __call__(self, left, right):
        l_re, l_im = self._channel_spec(left)
        r_re, r_im = self._channel_spec(right)
        out = np.empty((1, 4, self.dim_f, DIM_T), dtype=np.float32)
        out[0, 0], out[0, 1], out[0, 2], out[0, 3] = l_re, l_im, r_re, r_im   # [L_re, L_im, R_re, R_im]
        return out


class Istft:
    """Inverse STFT: CaC tensor [1, 4, dim_f, 256] -> stereo chunk (chunk_size samples/channel).

    Port of CacIstft.inverse. Rebuilds each frame's one-sided spectrum (Nyquist bin zero-padded back),
    irfft (includes 1/n), applies the synthesis Hann window, overlap-adds, divides by the running
    window-sum-of-squares envelope (+1e-8), and trims the n_fft/2 center pad each side.
    """

    def __init__(self, n_fft, dim_f):
        self.n = n_fft
        self.dim_f = dim_f
        self.pad = n_fft // 2
        self.chunk_size = HOP * (DIM_T - 1)
        self.hann = hann_periodic(n_fft)
        self.ola_len = (DIM_T - 1) * HOP + n_fft
        self.hann_sq = (self.hann * self.hann).astype(np.float32)

    def _channel_wave(self, re, im):
        # Full one-sided spectrum [DIM_T, n/2+1]; bins >= dim_f (i.e. the Nyquist bin) stay zero.
        full = np.zeros((DIM_T, self.n // 2 + 1), dtype=np.complex128)
        full[:, : self.dim_f] = re.T + 1j * im.T
        frames = np.fft.irfft(full, n=self.n, axis=1).astype(np.float32)   # [DIM_T, n], includes 1/n
        frames = frames * self.hann                                        # synthesis window
        ola = np.zeros(self.ola_len, dtype=np.float32)
        env = np.zeros(self.ola_len, dtype=np.float32)
        for t in range(DIM_T):
            s = t * HOP
            ola[s : s + self.n] += frames[t]
            env[s : s + self.n] += self.hann_sq
        seg = slice(self.pad, self.pad + self.chunk_size)
        return ola[seg] / (env[seg] + 1e-8)

    def __call__(self, spec):
        left = self._channel_wave(spec[0, 0], spec[0, 1])
        right = self._channel_wave(spec[0, 2], spec[0, 3])
        return left, right


def build_crossfade_window(gen_size, overlap_len):
    """Trapezoid: ramp up over first overlap_len, flat 1, ramp down over last overlap_len. Strictly
    positive so the envelope divide is well-defined at the un-overlapped first/last edges. Port of
    CacDsp.buildCrossfadeWindow (overlap_len == 0 -> all ones = hard tile)."""
    if overlap_len <= 0:
        return np.ones(gen_size, dtype=np.float32)
    denom = float(overlap_len + 1)
    k = np.arange(gen_size, dtype=np.float32)
    up = (k + 1.0) / denom
    down = (gen_size - k) / denom
    return np.minimum(1.0, np.minimum(up, down)).astype(np.float32)


# ══════════════════════════════════════════════════════════════════════════════════════════════════
#  Inference seam — a numpy port of the MdxEngine seam (MdxLiteRtEngine.kt). One run() method:
#  CaC spectrogram in -> vocal CaC spectrogram out. This is the whole "how to call the model" story.
# ══════════════════════════════════════════════════════════════════════════════════════════════════

class MdxLiteRtEngine:
    """LiteRT inference stage. On the CPU this is the classic Interpreter path (== InterpreterRunner in
    the app; the GPU/NPU rungs use the CompiledModel API and are Android-only)."""

    def __init__(self, tflite_path):
        self.it = Interpreter(model_path=tflite_path)   # LiteRT memory-maps the .tflite
        self.it.allocate_tensors()
        self._in = self.it.get_input_details()[0]
        self._out = self.it.get_output_details()[0]

    def _infer(self, x):
        self.it.set_tensor(self._in["index"], np.ascontiguousarray(x, dtype=np.float32))
        self.it.invoke()
        return self.it.get_tensor(self._out["index"]).copy()   # copy: the tensor is reused across calls

    def run(self, spec, denoise=False):
        out = self._infer(spec)
        if denoise:
            # UVR shift trick: spec_pred = 0.5*model(spec) - 0.5*model(-spec). 2x inference cost.
            out = 0.5 * out - 0.5 * self._infer(-spec)
        return out


# ══════════════════════════════════════════════════════════════════════════════════════════════════
#  Pipeline — a numpy port of MdxSeparationPipeline.separate: chunk, STFT -> infer -> iSTFT, weighted
#  overlap-add, envelope normalize, volume compensation, and the free instrumental residual.
# ══════════════════════════════════════════════════════════════════════════════════════════════════

def separate(engine, cfg, left, right, denoise=False):
    n_fft, dim_f, comp = cfg["n_fft"], cfg["dim_f"], cfg["compensation"]
    trim = n_fft // 2                              # STFT center trim / per-side overlap
    chunk_size = HOP * (DIM_T - 1)
    gen_size = chunk_size - 2 * trim               # useful samples a chunk contributes
    n_sample = min(len(left), len(right))
    if n_sample <= 0:
        raise ValueError("empty audio")

    # Overlap geometry (MdxSeparationPipeline.kt).
    stride = int(round(gen_size * (1.0 - OVERLAP)))
    stride = max(1, min(gen_size, stride))
    overlap_len = gen_size - stride
    window = build_crossfade_window(gen_size, overlap_len)
    n_chunks = 1 if n_sample <= gen_size else (n_sample - gen_size + stride - 1) // stride + 1
    out_len = (n_chunks - 1) * stride + gen_size

    # mixture = [trim zeros | audio | right-pad zeros | trim zeros] so every window is in-bounds.
    mix_len = trim + out_len + trim
    mix_l = np.zeros(mix_len, dtype=np.float32); mix_l[trim : trim + n_sample] = left[:n_sample]
    mix_r = np.zeros(mix_len, dtype=np.float32); mix_r[trim : trim + n_sample] = right[:n_sample]

    acc_l = np.zeros(out_len, dtype=np.float32)
    acc_r = np.zeros(out_len, dtype=np.float32)
    env = np.zeros(out_len, dtype=np.float32)

    stft = Stft(n_fft, dim_f)
    istft = Istft(n_fft, dim_f)
    for i in range(n_chunks):
        s = i * stride
        spec = stft(mix_l[s : s + chunk_size], mix_r[s : s + chunk_size])   # STFT
        voc = engine.run(spec, denoise)                                     # inference
        w_l, w_r = istft(voc)                                               # iSTFT -> chunk_size samples
        # Window the central gen_size samples (drop trim each side) and overlap-add.
        acc_l[s : s + gen_size] += window * w_l[trim : trim + gen_size]
        acc_r[s : s + gen_size] += window * w_r[trim : trim + gen_size]
        env[s : s + gen_size] += window
        print(f"  chunk {i + 1}/{n_chunks}", end="\r", flush=True)
    print()

    # Normalize by the crossfade envelope, crop to length, apply volume compensation.
    voc_l = (acc_l[:n_sample] / (env[:n_sample] + 1e-8) * comp).astype(np.float32)
    voc_r = (acc_r[:n_sample] / (env[:n_sample] + 1e-8) * comp).astype(np.float32)
    # Instrumental is the free residual: original mix - vocals.
    ins_l = (left[:n_sample] - voc_l).astype(np.float32)
    ins_r = (right[:n_sample] - voc_r).astype(np.float32)
    return voc_l, voc_r, ins_l, ins_r


# ══════════════════════════════════════════════════════════════════════════════════════════════════
#  Audio IO — stdlib `wave`, matching audio/AudioDecoder, AudioResampler, WavWriter (16-bit, /32768 in,
#  x32767 out). No external audio deps; WAV in / WAV out.
# ══════════════════════════════════════════════════════════════════════════════════════════════════

def read_wav_stereo_44100(path):
    with wave.open(path, "rb") as w:
        ch, width, sr, nframes = w.getnchannels(), w.getsampwidth(), w.getframerate(), w.getnframes()
        raw = w.readframes(nframes)
    if width == 2:
        data = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
    elif width == 1:                                  # WAV 8-bit is unsigned
        data = (np.frombuffer(raw, dtype=np.uint8).astype(np.float32) - 128.0) / 128.0
    elif width == 4:                                  # assume 32-bit signed PCM
        data = np.frombuffer(raw, dtype="<i4").astype(np.float32) / 2147483648.0
    else:
        raise ValueError(f"unsupported WAV sample width {width} bytes (use 16-bit PCM WAV)")
    data = data.reshape(-1, ch)
    if ch == 1:
        left = right = data[:, 0]                     # mono -> duplicate (AudioDecoder)
    else:
        left, right = data[:, 0], data[:, 1]          # keep first two channels
    if sr != SAMPLE_RATE:                             # resample only if needed (AudioResampler.resampleLinear)
        left = resample_linear(left, sr, SAMPLE_RATE)
        right = resample_linear(right, sr, SAMPLE_RATE)
    return np.ascontiguousarray(left), np.ascontiguousarray(right)


def resample_linear(src, src_sr, dst_sr):
    """Naive linear-interpolation resampler — a port of AudioResampler.resampleLinear (NOT librosa/soxr,
    so it matches the app, not UVR's reference; feed 44.1 kHz for exact parity)."""
    if src_sr == dst_sr or src.size < 2:
        return src.astype(np.float32)
    dst_len = max(1, int(src.size * dst_sr // src_sr))
    pos = np.arange(dst_len, dtype=np.float64) * (src_sr / dst_sr)
    i0 = np.floor(pos).astype(np.int64)
    frac = (pos - i0).astype(np.float32)
    i0 = np.clip(i0, 0, src.size - 1)
    i1 = np.clip(i0 + 1, 0, src.size - 1)
    return (src[i0] + (src[i1] - src[i0]) * frac).astype(np.float32)


def write_wav_stereo_16(path, left, right, sr=SAMPLE_RATE):
    n = min(len(left), len(right))
    l = np.clip(left[:n], -1.0, 1.0); r = np.clip(right[:n], -1.0, 1.0)
    inter = np.empty(n * 2, dtype="<i2")
    inter[0::2] = (l * 32767.0).astype(np.int16)      # x32767 truncate, matching WavWriter.toPcm16
    inter[1::2] = (r * 32767.0).astype(np.int16)
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    with wave.open(path, "wb") as w:
        w.setnchannels(2); w.setsampwidth(2); w.setframerate(sr)
        w.writeframes(inter.tobytes())


# ══════════════════════════════════════════════════════════════════════════════════════════════════

def main():
    ap = argparse.ArgumentParser(description="UVR MDX-Net vocal separation on LiteRT (STFT outside the graph).")
    ap.add_argument("input", help="input WAV (any sample rate; 44.1 kHz for exact parity)")
    ap.add_argument("--model", choices=list(MODELS), default="9482", help="which MDX-Net model (default: 9482)")
    ap.add_argument("--out-dir", default=".", help="directory for the output stems (default: .)")
    ap.add_argument("--denoise", action="store_true", help="UVR shift trick (2x inference, slightly cleaner)")
    args = ap.parse_args()

    cfg = MODELS[args.model]
    tflite = os.path.join(HERE, cfg["tflite"])
    if not os.path.exists(tflite):
        sys.exit(f"model not found: {tflite}\nDownload it from the repo (it is tracked with Git LFS).")

    print(f"model      : {args.model}  ({cfg['tflite']}, dim_f={cfg['dim_f']}, n_fft={cfg['n_fft']})")
    print(f"loading    : {args.input}")
    left, right = read_wav_stereo_44100(args.input)
    print(f"samples    : {len(left)} per channel ({len(left)/SAMPLE_RATE:.1f} s @ {SAMPLE_RATE} Hz)")

    engine = MdxLiteRtEngine(tflite)
    print("separating :")
    voc_l, voc_r, ins_l, ins_r = separate(engine, cfg, left, right, denoise=args.denoise)

    stem = os.path.splitext(os.path.basename(args.input))[0]
    voc_path = os.path.join(args.out_dir, f"{stem}_vocals.wav")
    ins_path = os.path.join(args.out_dir, f"{stem}_instrumental.wav")
    write_wav_stereo_16(voc_path, voc_l, voc_r)
    write_wav_stereo_16(ins_path, ins_l, ins_r)
    print(f"wrote      : {voc_path}")
    print(f"           : {ins_path}")


if __name__ == "__main__":
    main()