#!/usr/bin/env python3
"""
keys_to_notes.py: turn per-frame lit-key data from a "score + keyboard" video into
note events, then into a quantized two-staff MusicXML (hand = staff, from key colour).

Input (JSON): {"fps": 24, "notes": [[on_frame, off_frame, midi, "L"|"R"], ...],
               "bpm": 108, "beat0_frame": <frame of beat 1 of bar 1>, "key": -4, "beats": 4}
"""
import json, sys
from fractions import Fraction
from pathlib import Path
from music21 import stream, note, chord, meter, key, clef, tempo, layout, pitch, duration as m21dur

def quantize(x, grid):
    return round(x / grid) * grid

def main(src, out):
    d = json.loads(Path(src).read_text())
    fps, bpm = d["fps"], d["bpm"]
    fpb = fps * 60 / bpm                          # frames per beat
    b0 = d["beat0_frame"]
    ks = key.KeySignature(d.get("key", 0))
    grid = Fraction(1, 4)                         # 16th-note grid (in quarter notes)
    parts = {}
    for hand in ("R", "L"):
        ev = [n for n in d["notes"] if n[3] == hand]
        groups = {}
        for on, off, m, _ in ev:
            qs = Fraction(quantize(Fraction((on - b0) / fpb).limit_denominator(48), grid))
            qe = Fraction(quantize(Fraction((off - b0) / fpb).limit_denominator(48), grid))
            if qe <= qs:
                qe = qs + grid
            groups.setdefault(qs, []).append((m, qe - qs))
        p = stream.Part()
        p.insert(0, clef.TrebleClef() if hand == "R" else clef.BassClef())
        p.insert(0, ks)
        p.insert(0, meter.TimeSignature(f'{d.get("beats",4)}/4'))
        for qs in sorted(groups):
            pitches, durs = zip(*groups[qs])
            dur = min(durs)
            if len(pitches) == 1:
                el = note.Note(pitches[0])
            else:
                el = chord.Chord(sorted(set(pitches)))
            el.quarterLength = dur
            if qs >= 0:
                p.insert(float(qs), el)
        parts[hand] = p
    s = stream.Score()
    s.insert(0, parts["R"]); s.insert(0, parts["L"])
    s.insert(0, layout.StaffGroup([parts["R"], parts["L"]], symbol="brace"))
    # spell every pitch with the key signature's preference
    flat_key = ks.sharps < 0
    for n in s.recurse().notes:
        for pp in (n.pitches if n.isChord else [n.pitch]):
            if pp.accidental is not None and pp.accidental.alter != 0:
                if flat_key and pp.accidental.alter > 0:
                    pp.getEnharmonic(inPlace=True)
                elif not flat_key and pp.accidental.alter < 0:
                    pp.getEnharmonic(inPlace=True)
    s = s.makeNotation()
    s.write("musicxml", fp=out)
    print("wrote", out)

if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])