#!/usr/bin/env python3
"""Figures for the second round (MP3 cells, pitch-following TV). Run after run2.py has written its results.

  python figures2.py diag     # cell-position histogram, containment by bitrate, TV contrast  (needs data/, data_tune/)
  python figures2.py results  # ablation bars and bitrate sweep                               (needs results/round2_*.json)
"""
import json
import os
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from figures import ACCENT, GOLD, INK, MUTED, TEAL, OUT  # noqa: E402  (same style as the first round)

RES = os.path.join(HERE, "results")


def diag():
    import run as R
    import run2 as R2
    import unmix as U
    import voiceline as VL
    import cqt as CQ
    from mp3 import mp3model as MM
    rows, hist = [], None
    for br in ["mix.mp3", "mix_64k.mp3", "mix_32k.mp3"]:
        X = R2.setup(os.path.join(HERE, "data"), br, seconds=20)
        bs = X["bs"]
        S = MM.Synthesis(bs)
        clean = sum(X["truth"].values()) / X["gain"]                      # undo the gain the setup applied
        full = np.zeros_like(bs["pcm"]); full[:, 576 * X["g0"]:576 * X["g1"]] = clean   # excerpt only is enough
        import soundfile as sf
        st = [sf.read(os.path.join(HERE, "data", "stems", k + ".wav"))[0].T for k in ["voice", "chords", "bass", "drums"]]
        cl = sum(st); d = X["delay"]
        xs = np.zeros_like(bs["pcm"]); n = min(cl.shape[1], xs.shape[1] - d); xs[:, d:d + n] = cl[:, :n]
        A = S.analysis(xs) * X["gain"]
        ok = MM.reliable(bs)
        res = {"bitrate": br.replace("mix.mp3", "128k").replace("mix_", "").replace(".mp3", ""), "reliable": float(ok[bs["step"] > 0].mean())}
        for sl in [0.0, 0.5]:
            lo, hi = MM.cells(bs, slack=sl)
            res[f"inside_{sl}"] = float(((A >= lo) & (A <= hi))[ok].mean())
        p = bs["pcm"][:, :xs.shape[1]]
        res["codec_snr"] = float(10 * np.log10(np.sum((X["gain"] * xs) ** 2) / np.sum((p - X["gain"] * xs) ** 2)))
        rows.append(res)
        if hist is None:
            q = np.abs(bs["q"]).astype(float); lv = lambda m: np.maximum(m, 0) ** (4 / 3)
            sel = ok & (q >= 2) & (np.sign(A) == np.sign(bs["q"]))
            a = (lv(q - 1) + lv(q)) / 2; b = (lv(q) + lv(q + 1)) / 2
            hist = ((np.abs(A) / np.where(bs["step"] > 0, bs["step"], 1)) - a)[sel] / (b - a)[sel]
        print(res, flush=True)
    json.dump(rows, open(os.path.join(RES, "round2_cells.json"), "w"), indent=1)

    fig, ax = plt.subplots(figsize=(7.5, 3.2))
    ax.hist(np.clip(hist, -0.5, 1.5), bins=80, color=ACCENT, alpha=0.9)
    ax.axvline(0, color=INK, lw=1); ax.axvline(1, color=INK, lw=1)
    ax.set_xlabel("where the clean mix falls in its quantiser cell (0 = lower edge, 1 = upper edge)")
    ax.set_ylabel("coefficients"); ax.grid(axis="x", visible=False)
    ax.set_title(f"128 kbps, |q| ≥ 2: {np.mean((hist >= 0) & (hist <= 1)) * 100:.1f}% inside", fontsize=11, loc="left", color=INK)
    fig.savefig(os.path.join(OUT, "cells.png")); plt.close(fig)

    # TV contrast on the tuning song
    X = R2.setup(os.path.join(HERE, "data_tune"), seconds=12)
    g = X["g"]; T = g.shape[1]; F = U.Frame(T); ft = F.frame_times()
    f0t = R.oracle_f0(X["f0_samples"], F); f0e = VL.track(g, F)
    Gm = np.abs(F.fwd(g.mean(0))); sc = float(np.median(Gm[Gm > 0]))
    vest = U.separate2(g, f0e, a_H=3 * sc, a_P=9 * sc * np.sqrt(1024) ** 0.5, a_V=0.3 * sc, iters=150)["v"].astype(np.float64)
    fr = CQ.CQTFrame(T)
    voice = X["truth"]["voice"][0]; acc = (X["truth"]["chords"] + X["truth"]["bass"] + X["truth"]["drums"]).mean(0)
    variants = [("time only,\nphase from mix", f0e, fr.fwd(g.mean(0)), False),
                ("follows pitch,\nharmonic phase", f0e, None, True),
                ("follows pitch,\nphase from mix", f0e, fr.fwd(g.mean(0)), True),
                ("follows pitch,\nphase from voice\nestimate (loop)", f0e, fr.fwd(vest), True),
                ("follows true pitch,\nphase from\ntrue voice (bound)", f0t, fr.fwd(voice), True)]
    vals = []
    for name, f0, ref, fol in variants:
        D = CQ.DirectionalTV(fr, ft, f0, follow=fol, phase_ref=ref)
        tv = lambda x: np.abs(D(fr.fwd(x))).sum() / np.abs(fr.fwd(x)).sum()
        vals.append((name, tv(voice), tv(acc)))
        print(name.replace("\n", " "), round(vals[-1][1], 3), round(vals[-1][2], 3), flush=True)
    json.dump([{"variant": n, "voice": v, "accompaniment": a} for n, v, a in vals],
              open(os.path.join(RES, "round2_tv_contrast.json"), "w"), indent=1)
    plot_tv()


def plot_tv():
    vals = [(r["variant"], r["voice"], r["accompaniment"]) for r in json.load(open(os.path.join(RES, "round2_tv_contrast.json")))]
    fig, ax = plt.subplots(figsize=(9.5, 3.9))
    x = np.arange(len(vals)); w = 0.38
    ax.bar(x - w / 2, [v for _, v, _ in vals], w, color=ACCENT, label="true voice")
    ax.bar(x + w / 2, [a for _, _, a in vals], w, color=GOLD, label="true accompaniment")
    for i, (_, v, a) in enumerate(vals):
        ax.text(i, max(v, a) + 0.005, f"{a / v:.2f}×", ha="center", fontsize=9.5, color=INK)
    ax.set_xticks(x, [n for n, _, _ in vals], fontsize=8.5)
    ax.set_ylabel("TV per unit of signal\n(lower = smoother)")
    ax.set_ylim(0, max(max(v, a) for _, v, a in vals) * 1.18)
    ax.legend(frameon=False, loc="upper center", bbox_to_anchor=(0.5, 1.16), ncol=2)
    ax.grid(axis="x", visible=False)
    fig.savefig(os.path.join(OUT, "tv_contrast.png")); plt.close(fig)


def results():
    r = {br: json.load(open(os.path.join(RES, f"round2_{br}.json")))["results"] for br in ["128k", "64k", "32k"]}
    a = r["128k"]
    variants = [("fixed", "Voice TV, time direction"), ("follow", "Voice TV following the pitch\n(harmonic phase)"),
                ("loop", "… with outer loop\n(phase from voice estimate)"), ("mp3l1", "MP3 cells instead of equality"),
                ("mp3loop", "MP3 cells + outer-loop TV"), ("gate", "Range gate on the pitch"),
                ("bound", "Ceiling: TV with the true\nvoice's phase (not a method)")]
    key = {"fixed": ("eq+fixed", None), "follow": ("eq+follow", None), "loop": ("eq+loop", "eq+loop+oracle"),
           "mp3l1": ("mp3+l1", "mp3+l1+oracle"), "mp3loop": ("mp3+loop", "mp3+loop+oracle"), "gate": ("eq+l1+gate", None),
           "bound": ("eq+bound", "eq+bound+oracle")}
    fig, ax = plt.subplots(1, 2, figsize=(12, 4.6), sharey=True)
    y = np.arange(len(variants))[::-1]
    for axx, (m, title) in zip(ax, [("voice_si_sdr", "Voice, change in SI-SDR (dB)"), ("acc_sdr", "Accompaniment, change in SDR (dB)")]):
        for off, (which, base, col, lab) in zip([0.19, -0.19], [(0, "eq+l1+refit", ACCENT, "tracked pitch"), (1, "eq+l1+oracle+refit", TEAL, "true pitch")]):
            vals = []
            for v, _ in variants:
                k = key[v][which]
                vals.append(a[k + "+refit"][m] - a[base][m] if k and (k + "+refit") in a else np.nan)
            vals = np.array(vals)
            axx.barh(y + off, np.nan_to_num(vals), height=0.36, color=col, label=lab)
            for yi, v in zip(y + off, vals):
                if np.isfinite(v):
                    axx.text(v + (0.04 if v >= 0 else -0.04), yi, f"{v:+.2f}", va="center", ha="left" if v >= 0 else "right", fontsize=8.5, color=INK)
                else:
                    axx.text(0.04, yi, "not run", va="center", fontsize=8, color=MUTED)
        axx.axvline(0, color=INK, lw=0.8)
        axx.set_title(title, loc="left", fontsize=12, color=INK); axx.grid(axis="y", visible=False)
        lim = 2.2 if m == "voice_si_sdr" else 1.2
        axx.set_xlim(-0.8, lim)
    ax[0].set_yticks(y, [n for _, n in variants], fontsize=9.5)
    h, l = ax[1].get_legend_handles_labels()
    fig.legend(h, l, frameon=False, loc="upper center", bbox_to_anchor=(0.62, 1.04), ncol=2)
    fig.savefig(os.path.join(OUT, "round2_results.png")); plt.close(fig)

    brs = ["32k", "64k", "128k"]; xs = [32, 64, 128]
    fig, ax = plt.subplots(1, 2, figsize=(10.5, 3.6))
    for axx, key, title in [(ax[0], "voice_si_sdr", "Voice, SI-SDR (dB)"), (ax[1], "acc_sdr", "Accompaniment, SDR (dB)")]:
        for k, lab, col, ls in [("eq+l1+refit", "equality, tracked pitch", ACCENT, "--"), ("mp3+l1+refit", "MP3 cells, tracked pitch", ACCENT, "-"),
                                ("eq+l1+oracle+refit", "equality, true pitch", TEAL, "--"), ("mp3+l1+oracle+refit", "MP3 cells, true pitch", TEAL, "-")]:
            axx.plot(xs, [r[b][k][key] for b in brs], ls, color=col, marker="o", label=lab)
        axx.set_xscale("log", base=2); axx.set_xticks(xs, [f"{x} kbps" for x in xs]); axx.minorticks_off()
        axx.set_title(title, loc="left", fontsize=12, color=INK)
    ax[1].legend(frameon=False, fontsize=9)
    fig.savefig(os.path.join(OUT, "bitrate.png")); plt.close(fig)


if __name__ == "__main__":
    os.makedirs(RES, exist_ok=True)
    {"diag": diag, "results": results, "tv": plot_tv}[sys.argv[1] if len(sys.argv) > 1 else "results"]()