# Self-made GM drum clips (channel 10) with exact ground truth. 2026-10-20.
import json,struct
TPQ=480
def vlq(n):
    b=[n&0x7f]; n>>=7
    while n: b.append((n&0x7f)|0x80); n>>=7
    return bytes(reversed(b))
def write(path,notes,bpm):
    spb=60/bpm; ev=[]
    for p,o,v in notes:
        ev.append((round(o/spb*TPQ),1,p,v)); ev.append((round((o+0.1)/spb*TPQ),0,p,0))
    ev.sort(key=lambda e:(e[0],e[1]))
    trk=bytearray(vlq(0)+b'\xff\x51\x03'+round(60e6/bpm).to_bytes(3,'big')); t=0
    for tick,on,p,v in ev:
        trk+=vlq(tick-t)+bytes([0x99 if on else 0x89,p,v]); t=tick
    trk+=vlq(0)+b'\xff\x2f\x00'
    open(path,'wb').write(b'MThd'+struct.pack('>IHHH',6,0,1,TPQ)+b'MTrk'+struct.pack('>I',len(trk))+trk)
K,S,CH,OH,CR,T1,T2,T3=36,38,42,46,49,50,47,43
clips={}
def grid(bpm,bars,fn,start=0.5):
    q=60/bpm; out=[]
    for bar in range(bars):
        for step in range(16):
            t=start+(bar*16+step)*q/4
            for p,v in fn(bar,step): out.append((p,round(t,4),v))
    return out
# 1 rock beat, 100 BPM, 8th-note closed hats, kick 1 & 3, snare 2 & 4, 4 bars
clips['1-rock-beat']=(100,grid(100,4,lambda b,s:([(CH,90)] if s%2==0 else [])+([(K,110)] if s in (0,8) else [])+([(S,110)] if s in (4,12) else [])))
# 2 kick + snare only (no cymbals), 90 BPM, syncopated kick
clips['2-kick-snare-only']=(90,grid(90,4,lambda b,s:([(K,110)] if s in (0,6,10) else [])+([(S,105)] if s in (4,12) else [])))
# 3 16th-note hats at 120 BPM with kick/snare
clips['3-sixteenth-hats']=(120,grid(120,4,lambda b,s:[(CH,80 if s%2 else 95)]+([(K,110)] if s in (0,7,8) else [])+([(S,110)] if s in (4,12) else [])))
# 4 hats only, 8ths, open hat on the "and" of 4
clips['4-hats-only']=(100,grid(100,4,lambda b,s:[(OH if s==14 else CH,90)] if s%2==0 else []))
# 5 tom fill: 16ths high->low tom, crash + kick on the next downbeat
def fill(b,s):
    if b<3: return ([(CH,90)] if s%2==0 else [])+([(K,110)] if s in (0,8) else [])+([(S,110)] if s in (4,12) else [])
    return [([T1,T1,T1,T1,T2,T2,T2,T2,T3,T3,T3,T3,S,S,S,S][s],105)]
n=grid(110,4,fill); n.append((CR,110,round(0.5+64*(60/110)/4,4))) if False else None
end=round(0.5+64*(60/110)/4,4); n+= [(CR,end,115),(K,end,115)]
clips['5-tom-fill']=(110,n)
# 6 ghost notes: snare backbeat 110 + ghost snares at velocity 35
clips['6-ghost-notes']=(95,grid(95,4,lambda b,s:([(CH,85)] if s%2==0 else [])+([(K,110)] if s in (0,8,10) else [])+([(S,110)] if s in (4,12) else [])+([(S,35)] if s in (7,9,15) else [])))
# 7 fast double-kick 16ths at 140 BPM, snare on 2 & 4
clips['7-double-kick']=(140,grid(140,2,lambda b,s:[(K,105)]+([(S,110)] if s in (4,12) else [])))
truth={}
for k,(bpm,notes) in clips.items():
    write(k+'.mid',notes,bpm); truth[k]={'bpm':bpm,'notes':[{'pitch':p,'onset':o,'vel':v} for p,o,v in sorted(notes,key=lambda x:(x[1],x[0]))]}
json.dump(truth,open('ground-truth.json','w'),indent=1)
for k,v in truth.items(): print(k,v['bpm'],len(v['notes']))
