#!/usr/bin/env python3
"""Score a pianoify analysis.json against a make_test_song truth.json."""
import json
import sys

import numpy as np


def beat_to_time(beat_times, origin, b):
    i = origin + b
    i0 = int(np.clip(np.floor(i), 0, len(beat_times) - 2))
    return beat_times[i0] + (i - i0) * (beat_times[i0 + 1] - beat_times[i0])


def main(analysis_path, truth_path):
    an, tr = json.load(open(analysis_path)), json.load(open(truth_path))
    bt, org, shift = an["beat_times"], an["origin_beat"], an.get("transpose", 0)
    step = 0.02
    end = max(c["end"] for c in tr["chords"])

    # chords: fraction of time with the correct label (root + maj/min)
    det = [(beat_to_time(bt, org, c["start_beat"]), beat_to_time(bt, org, c["start_beat"] + c["beats"]),
            c["name"]) for c in an["chords"]]
    ok = tot = 0
    for t in np.arange(0, end, step):
        truth = next((c["chord"] for c in tr["chords"] if c["start"] <= t < c["end"]), "N")
        guess = next((n for a, b, n in det if a <= t < b), "N.C.")
        if truth == "N":
            continue
        tot += 1
        ok += guess == truth
    chord_acc = ok / tot

    # melody: frame-level pitch accuracy over truth-voiced time
    mel = [(beat_to_time(bt, org, n["start_beat"]), beat_to_time(bt, org, n["start_beat"] + n["beats"]),
            n["midi"] - shift) for n in an["melody"]]
    exact = chroma = voiced = cov = 0
    for t in np.arange(0, end, step):
        truth = next((n["midi"] for n in tr["notes"] if n["start"] <= t < n["end"]), None)
        guess = next((m for a, b, m in mel if a <= t < b), None)
        if truth is None:
            continue
        voiced += 1
        if guess is not None:
            cov += 1
            exact += guess == truth
            chroma += (guess - truth) % 12 == 0
    # note onsets (±80 ms)
    det_on = np.array([a for a, _, _ in mel])
    hits = sum(np.any(np.abs(det_on - n["start"]) < 0.08) for n in tr["notes"]) if len(det_on) else 0
    prec = hits / max(len(det_on), 1)
    print(f"tempo        : {an['bpm']} (truth {tr['bpm']})")
    print(f"key          : {an['key']} (truth {tr['key']})")
    print(f"chords       : {chord_acc:.1%} of time correct")
    print(f"melody cover : {cov / voiced:.1%} of sung time has a note")
    print(f"melody pitch : {exact / voiced:.1%} exact, {chroma / voiced:.1%} ignoring octave")
    print(f"note onsets  : recall {hits / len(tr['notes']):.1%}, precision {min(prec, 1):.1%} "
          f"({len(det_on)} detected / {len(tr['notes'])} true)")


if __name__ == "__main__":
    main(*sys.argv[1:3])
