#!/usr/bin/env python3
"""Run the separation experiments on a stems folder made by make_stems.py and print/save a results table.

  python run.py data [--seconds 20] [--iters 300] [--tune data_tune]

Methods
  mixture         the mix itself as the voice estimate (floor)
  mid             (L+R)/2: the karaoke trick inverted, keeps everything centred
  repet-sim       librosa's nearest-neighbour foreground/background soft mask (REPET-SIM style), on mid
  comb-mask       soft mask of the mid STFT to harmonics of the tracked pitch (the support alone, no solver)
  pdhg            the proposed constrained reconstruction (stereo, comb support, split fidelity)
  pdhg-mono       same, ignoring stereo
  pdhg-oracle-f0  same, with the true pitch (upper bound for the pitch tracker)
  pdhg-naive-f0   same, with the plain Pianoify pitch tracker (no centre/vibrato/register cues)
  pdhg-tolerance  same, with 2% per-band data tolerance instead of exact equality
Metrics: SDR and SI-SDR (dB) of the voice and of the accompaniment, against the true stems,
and raw pitch accuracy (share of voiced frames within 50 cents) of a pitch track run on each voice estimate.
"""
import argparse
import json
import os
import subprocess
import sys
import time

import numpy as np
import soundfile as sf

sys.path.insert(0, os.path.dirname(__file__))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "pianoify"))
import unmix as U  # noqa: E402
import voiceline as VL  # noqa: E402

SR = U.SR


def load_mp3(path):
    raw = subprocess.run(["ffmpeg", "-v", "error", "-i", path, "-ac", "2", "-ar", str(SR), "-f", "f32le", "-"],
                         capture_output=True, check=True).stdout
    return np.frombuffer(raw, np.float32).reshape(-1, 2).T.copy()


def sdr(ref, est):
    return 10 * np.log10(np.sum(ref ** 2) / (np.sum((ref - est) ** 2) + 1e-12))


def si_sdr(ref, est):
    r, e = ref.ravel(), est.ravel()
    a = (e @ r) / (r @ r + 1e-12)
    return 10 * np.log10(np.sum((a * r) ** 2) / (np.sum((a * r - e) ** 2) + 1e-12))


def track_f0(mono, n_frames, frame_times):
    """Pianoify's predominant-pitch tracker (harmonic salience + Viterbi) → Hz per STFT frame."""
    import librosa
    import pianoify as PZ
    y = librosa.resample(mono.astype(np.float64), orig_sr=SR, target_sr=PZ.SR)
    t, midi, _ = PZ.melody_contour(y, monophonic=False)
    hz = np.where(np.isfinite(midi), 440 * 2 ** ((midi - 69) / 12), 0.0)
    idx = np.clip(np.searchsorted(t, frame_times), 0, len(t) - 1)
    return hz[idx]


def track_f0_mono(mono, frame_times):
    import librosa
    import pianoify as PZ
    y = librosa.resample(mono.astype(np.float64), orig_sr=SR, target_sr=PZ.SR)
    t, midi, _ = PZ.melody_contour(y, monophonic=True)
    hz = np.where(np.isfinite(midi), 440 * 2 ** ((midi - 69) / 12), 0.0)
    idx = np.clip(np.searchsorted(t, frame_times), 0, len(t) - 1)
    return hz[idx]


def rpa(est_hz, true_hz):
    v = true_hz > 0
    ok = (est_hz > 0) & v & (np.abs(1200 * np.log2(np.maximum(est_hz, 1e-9) / np.maximum(true_hz, 1e-9))) < 50)
    return ok.sum() / max(v.sum(), 1)


def oracle_f0(f0_samples, F):
    out = np.zeros(F.nF)
    for n in range(F.nF):
        a = n * F.hop
        seg = f0_samples[a:a + F.N]
        seg = seg[seg > 0]
        out[n] = np.median(seg) if seg.size > F.N // 2 else 0.0
    return out


def comb_mask_baseline(mid, F, f0):
    supp = U.comb_support(F, f0)
    X = F.fwd(mid)
    # soft mask: keep supported bins, attenuate the rest
    return F.adj(np.where(supp, X, 0.05 * X))


def repet_sim(mid):
    import librosa
    S = librosa.stft(mid.astype(np.float64), n_fft=1024, hop_length=256)
    M, P = np.abs(S), np.angle(S)
    filt = librosa.decompose.nn_filter(M, aggregate=np.median, metric="cosine",
                                       width=int(librosa.time_to_frames(2, sr=SR, hop_length=256)))
    filt = np.minimum(M, filt)
    mask_v = librosa.util.softmask(M - filt, 10 * filt, power=2)
    return librosa.istft(mask_v * M * np.exp(1j * P), hop_length=256, length=len(mid)).astype(np.float32)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("data")
    ap.add_argument("--start", type=float, default=2.5)
    ap.add_argument("--seconds", type=float, default=20.0)
    ap.add_argument("--iters", type=int, default=300)
    ap.add_argument("--aH", type=float, default=None)
    ap.add_argument("--aP", type=float, default=None)
    ap.add_argument("--aV", type=float, default=None)
    ap.add_argument("--eps", type=float, default=0.0)
    ap.add_argument("--only", default="")
    ap.add_argument("--out", default=None)
    ap.add_argument("--merge", action="store_true", help="update an existing results file instead of replacing it")
    a = ap.parse_args()

    g_all = load_mp3(os.path.join(a.data, "mix.mp3"))
    stems = {k: sf.read(os.path.join(a.data, "stems", k + ".wav"), dtype="float32")[0].T for k in ["voice", "chords", "bass", "drums"]}
    f0_true_s = np.load(os.path.join(a.data, "f0_true.npy"))
    # the MP3 decoder adds a small delay: align it to the clean mix by cross-correlation
    clean = sum(stems.values())
    lag = int(np.argmax(np.correlate(g_all[0, :SR * 4], clean[0, :SR * 2], mode="valid")))
    g_all = g_all[:, lag:lag + clean.shape[1]]
    hop = 256
    s0 = int(a.start * SR)
    T = (int(a.seconds * SR) // hop) * hop
    cut = lambda x: x[..., s0:s0 + T]
    g = cut(g_all)
    voice, acc = cut(stems["voice"]), cut(stems["chords"] + stems["bass"] + stems["drums"])
    harm, drums = cut(stems["chords"] + stems["bass"]), cut(stems["drums"])
    F = U.Frame(T)
    ft = F.frame_times()
    mid = g.mean(axis=0)
    f0_true = oracle_f0(cut(f0_true_s), F)
    t0 = time.time()
    f0_naive = track_f0(mid, F.nF, ft)
    f0_est = VL.track(g, F)
    print(f"pitch of the voice from the mix: plain Pianoify tracker {rpa(f0_naive, f0_true):.3f}, "
          f"centre+vibrato+register {rpa(f0_est, f0_true):.3f}  ({time.time() - t0:.1f}s); "
          f"voice type {VL.voice_type(f0_est)[0]}")
    results_f0 = {"naive": rpa(f0_naive, f0_true), "voiceline": rpa(f0_est, f0_true), "voice_type": VL.voice_type(f0_est)[0]}
    print(f"codec noise of the MP3 vs clean mix: SNR {sdr(cut(clean), g):.1f} dB")

    # default weights: scale-free choice relative to the mixture's coefficient magnitudes
    Gm = np.abs(F.fwd(mid))
    scale = float(np.median(Gm[Gm > 0]))
    # chosen on the separate tuning song (data_tune), never on the evaluation song
    aH = a.aH if a.aH is not None else 3.0 * scale
    aP = a.aP if a.aP is not None else 9.0 * scale * np.sqrt(F.N) ** 0.5
    aV = a.aV if a.aV is not None else 0.3 * scale

    results = {}

    def score(name, v_est, acc_est=None, h_est=None, p_est=None, extra=None):
        v2 = np.stack([v_est, v_est]) if v_est.ndim == 1 else v_est
        acc_est = acc_est if acc_est is not None else g - v2
        r = {"voice_sdr": sdr(voice, v2), "voice_si_sdr": si_sdr(voice, v2),
             "acc_sdr": sdr(acc, acc_est), "acc_si_sdr": si_sdr(acc, acc_est)}
        if h_est is not None:
            r["harm_sdr"] = sdr(harm, h_est); r["drums_sdr"] = sdr(drums, p_est)
        r["voice_rpa"] = rpa(track_f0_mono(v2.mean(axis=0), ft), f0_true)
        if extra: r.update(extra)
        results[name] = r
        print(f"{name:16s} voice SDR {r['voice_sdr']:6.2f}  SI-SDR {r['voice_si_sdr']:6.2f}  | acc SDR {r['acc_sdr']:6.2f}"
              + (f"  | harm {r['harm_sdr']:6.2f} drums {r['drums_sdr']:6.2f}" if h_est is not None else "")
              + f"  | pitch acc {r['voice_rpa']:.3f}", flush=True)

    want = lambda n: not a.only or n in a.only.split(",")
    if want("mixture"): score("mixture", g)
    if want("mid"): score("mid", mid)
    if want("repet-sim"): score("repet-sim", repet_sim(mid))
    if want("comb-mask"): score("comb-mask", comb_mask_baseline(mid, F, f0_est))
    if want("comb-mask-oracle"): score("comb-mask-oracle", comb_mask_baseline(mid, F, f0_true))

    def run_pdhg(name, f0, stereo=True, eps=a.eps):
        t = time.time()
        out = U.separate(g, f0, a_H=aH, a_P=aP, a_V=aV, eps_rel=eps, iters=a.iters, stereo=stereo, verbose=False)
        score(name, out["v"], out["h"] + out["p"], out["h"], out["p"], {"seconds": time.time() - t})
        if stereo:
            ref = U.refit(g, out)
            score(name + "+refit", ref["v"], ref["h"] + ref["p"], ref["h"], ref["p"])
            out = ref
        return out

    best = None
    if want("pdhg"): best = run_pdhg("pdhg", f0_est)
    if want("pdhg-mono"): run_pdhg("pdhg-mono", f0_est, stereo=False)
    if want("pdhg-oracle-f0"): run_pdhg("pdhg-oracle-f0", f0_true)
    if want("pdhg-naive-f0"): run_pdhg("pdhg-naive-f0", f0_naive)
    if want("pdhg-tolerance"): run_pdhg("pdhg-tolerance", f0_est, eps=0.02)

    out = a.out or os.path.join(a.data, "results.json")
    if a.merge and os.path.exists(out):
        old = json.load(open(out))["results"]
        old.update(results)
        results.clear(); results.update(old)
    clean_json = lambda d: {k: (clean_json(v) if isinstance(v, dict) else (float(v) if isinstance(v, (np.floating, float)) else v)) for k, v in d.items()}
    json.dump(clean_json({"weights": {"aH": aH, "aP": aP, "aV": aV, "eps": a.eps, "iters": a.iters, "seconds": a.seconds},
               "pitch_tracking": results_f0, "results": results}), open(out, "w"), indent=1)
    if best is not None:
        sf.write(os.path.join(a.data, "est_voice.wav"), best["v"], SR)
        sf.write(os.path.join(a.data, "est_accomp.wav"), (best["h"] + best["p"]).T, SR)
    print("saved", out)


if __name__ == "__main__":
    main()