#!/usr/bin/env python3
"""Experiments for the three extensions: MP3 quantiser cells, pitch-following directional TV, range gate.

  python run2.py data --seconds 20 --iters 300 --methods eq,eq+fixed,eq+follow,mp3+follow,... [--mp3 mix_32k.mp3]

Works on the decoder's own timeline: the mixture is minimp3's output for the MP3, the excerpt is a
whole number of granules, and the true stems are shifted by the codec delay and scaled by the encoder's
input gain (both estimated from the clean mix, for scoring only).

Method names: <data>+<voice prior>[+oracle][+gate]
  data:  eq   exact equality with the decoded mixture (as in the first prototype)
         mp3  the mixture's MP3 analysis must lie in the quantiser cells the file specifies
         mp3h hybrid: quantiser cells where the file is informative, the decoded value elsewhere
  voice: l1      sparse comb coefficients only (first prototype)
         fixed   + phase-predicted TV along time in a constant-Q frame (fixed direction, mixture phase)
         follow  + TV along the pitch slope with harmonic phase prediction
         loop    + outer loop: re-solve with TV following the voice estimate's pitch and phase
         loopfixed  same, fixed direction
         bound   diagnostic ceiling: TV with the true voice's phase and pitch (not a method)
  oracle: use the true pitch;  gate: drop pitch frames outside the estimated voice type's range
"""
import argparse
import json
import os
import sys
import time

import numpy as np
import soundfile as sf

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import run as R        # noqa: E402  (metrics, oracle pitch helpers)
import unmix as U      # noqa: E402
import voiceline as VL # noqa: E402
from mp3 import mp3model as MM  # noqa: E402

SR = 16000


def setup(data, mp3name="mix.mp3", start=2.5, seconds=20.0):
    bs = MM.read_bitstream(os.path.join(data, mp3name))
    pcm = bs["pcm"]
    stems = {k: sf.read(os.path.join(data, "stems", k + ".wav"), dtype="float64")[0].T for k in ["voice", "chords", "bass", "drums"]}
    clean = sum(stems.values())
    seg = slice(20000, 200000)
    d = max(range(0, 2400), key=lambda d: float(np.dot(pcm[0, seg], clean[0, seg.start - d:seg.stop - d])))
    shift = lambda x: np.pad(x, ((0, 0), (d, 0)))[:, :pcm.shape[1]] if x.shape[1] + d >= pcm.shape[1] else np.pad(x, ((0, 0), (d, pcm.shape[1] - x.shape[1] - d)))
    cl = shift(clean)
    # encoder input gain, fitted in the coefficient domain on well-quantised values (|q| >= 2), so the
    # encoder's lowpass doesn't bias it
    A = MM.Synthesis(bs).analysis(cl)
    sel = (bs["step"] > 0) & (np.abs(bs["q"]) >= 2)
    gam = float(np.sum(A[sel] * bs["x"][sel]) / np.sum(A[sel] ** 2))
    n = int(round(seconds * SR / 576)) // 4 * 4                     # granules; T multiple of 256
    g0 = int(np.ceil((start * SR + d) / 576))
    g1 = g0 + n
    sl = slice(576 * g0, 576 * g1)
    g = pcm[:, sl].astype(np.float32)
    truth = {k: gam * shift(v)[:, sl] for k, v in stems.items()}
    f0s = np.pad(np.load(os.path.join(data, "f0_true.npy")), (d, 0))[: pcm.shape[1]][sl]
    return dict(bs=bs, g=g, truth=truth, f0_samples=f0s, g0=g0, g1=g1, delay=d, gain=gam)


def gate_range(f0, margin=1.25):
    """Range consistency: drop pitch frames outside the line's own register (5th-95th pct, widened)."""
    v = f0[f0 > 0]
    if v.size < 20:
        return f0
    lo, hi = np.percentile(v, 5) / margin, np.percentile(v, 95) * margin
    return np.where((f0 >= lo) & (f0 <= hi), f0, 0.0)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("data")
    ap.add_argument("--mp3", default="mix.mp3")
    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("--methods", default="eq+l1,eq+fixed,eq+follow,mp3+l1,mp3+follow")
    ap.add_argument("--aH", type=float, default=3.0)
    ap.add_argument("--aP", type=float, default=9.0)
    ap.add_argument("--aV", type=float, default=0.3)
    ap.add_argument("--aT", type=float, default=1.0, help="voice TV weight, in units of the mixture's median coefficient")
    ap.add_argument("--slack", type=float, default=0.5)
    ap.add_argument("--out", required=True)
    ap.add_argument("--merge", action="store_true")
    a = ap.parse_args()

    X = setup(a.data, a.mp3, a.start, a.seconds)
    g, truth = X["g"], X["truth"]
    T = g.shape[1]
    F = U.Frame(T)
    ft = F.frame_times()
    f0_true = R.oracle_f0(X["f0_samples"], F)
    f0_est = VL.track(g, F)
    voice, acc = truth["voice"], truth["chords"] + truth["bass"] + truth["drums"]
    Gm = np.abs(F.fwd(g.mean(axis=0)))
    sc = float(np.median(Gm[Gm > 0]))
    print(f"{a.mp3}: excerpt granules {X['g0']}-{X['g1']} ({T / SR:.2f} s), codec delay {X['delay']}, encoder gain {X['gain']:.4f}; "
          f"pitch accuracy tracked {R.rpa(f0_est, f0_true):.3f}", flush=True)
    mp3d = U.MP3Data(X["bs"], X["g0"], X["g1"], slack=a.slack)
    mp3h = U.MP3Data(X["bs"], X["g0"], X["g1"], slack=a.slack, unreliable="decoded")
    tsum = (voice + acc).astype(np.float64)
    inside = (mp3d.fwd(tsum) >= mp3d.lo) & (mp3d.fwd(tsum) <= mp3d.hi)
    print(f"true mix inside its reliable cells: {inside[mp3d.ok].mean():.4f}  (reliable share {mp3d.ok.mean():.3f})", flush=True)

    results = {}
    if a.merge and os.path.exists(a.out):
        results = json.load(open(a.out)).get("results", {})

    def score(name, v, h, p, secs):
        v2 = np.stack([v, v]); acc_est = h + p
        r = {"voice_sdr": R.sdr(voice, v2), "voice_si_sdr": R.si_sdr(voice, v2), "acc_sdr": R.sdr(acc, acc_est),
             "acc_si_sdr": R.si_sdr(acc, acc_est), "voice_rpa": R.rpa(R.track_f0_mono(v2.mean(0), ft), f0_true), "seconds": secs}
        results[name] = {k: float(x) for k, x in r.items()}
        print(f"{name:26s} voice SDR {r['voice_sdr']:6.2f} SI-SDR {r['voice_si_sdr']:6.2f} | acc SDR {r['acc_sdr']:6.2f} | pitch {r['voice_rpa']:.3f} | {secs:5.0f}s", flush=True)

    for m in a.methods.split(","):
        parts = m.split("+")
        data, vprior = parts[0], parts[1]
        f0 = f0_true if "oracle" in parts else f0_est
        if "gate" in parts:
            f0 = gate_range(f0)
        t0 = time.time()
        kw = dict(a_H=a.aH * sc, a_P=a.aP * sc * np.sqrt(F.N) ** 0.5, a_V=a.aV * sc, mp3={"mp3": mp3d, "mp3h": mp3h}.get(data), iters=a.iters)
        if vprior in ("l1",):
            out = U.separate2(g, f0, **kw)
        elif vprior in ("fixed", "follow"):
            out = U.separate2(g, f0, a_T=a.aT * sc, voice_tv=vprior, **kw)
        elif vprior in ("loop", "loopfixed"):              # outer loop: phase (and pitch) from the voice estimate
            out = U.separate_outer(g, f0, outer=2, voice_tv=("follow" if vprior == "loop" else "fixed"), a_T=a.aT * sc, **kw)
        elif vprior == "bound":                             # diagnostic ceiling: phase from the TRUE voice
            out = U.separate2(g, f0, a_T=a.aT * sc, voice_tv="follow", phase_from=voice[0], f0_hz_cqt=f0_true, **kw)
        secs = time.time() - t0
        score(m, out["v"], out["h"], out["p"], secs)
        ref = U.refit(g.astype(np.float64), out)
        score(m + "+refit", ref["v"], ref["h"], ref["p"], secs)
        json.dump({"args": vars(a), "delay": X["delay"], "gain": X["gain"], "results": results}, open(a.out, "w"), indent=1)


if __name__ == "__main__":
    main()