# post_process.py

## Vibe-coded with Z AI

### Project Los Angeles
### Tegridy Code 2026

### License CC BY-NC-SA 4.0

import math

# --------------------------------------------------------------------------- #
#  Helpers
# --------------------------------------------------------------------------- #
def _ioi_mode(onsets, res_ms=5.0):
    """Most common inter-onset interval (binned), ignoring <40 ms gaps."""
    us = sorted(set(onsets))
    iois = [us[i + 1] - us[i] for i in range(len(us) - 1) if us[i + 1] - us[i] > 1e-6]
    if not iois:
        return None
    bins = {}
    for d in iois:
        b = int(round(d / res_ms)) * res_ms
        if b >= 40.0:
            bins[b] = bins.get(b, 0) + 1
    if not bins:
        return None
    return max(bins, key=bins.get)


def _detect_beat_period_ms(onsets, res_ms=5.0, lo_bpm=50.0, hi_bpm=210.0):
    """
    Estimate the quarter-note period (ms) from absolute note-onset times
    using the autocorrelation of an onset impulse train. Peak is searched
    inside the 70-180 BPM preferred range and parabolically refined.
    """
    if len(onsets) < 2:
        return 500.0
    t_max = max(onsets)
    if t_max <= 0:
        return 500.0

    n = int(t_max / res_ms) + 2
    env = [0.0] * n
    for t in onsets:
        idx = int(round(t / res_ms))
        if 0 <= idx < n:
            env[idx] += 1.0

    lag_min = max(1, int(round((60000.0 / hi_bpm) / res_ms)))
    lag_max = max(lag_min + 1, int(round((60000.0 / lo_bpm) / res_ms)))
    acf = [0.0] * (lag_max + 2)
    for lag in range(lag_min, lag_max + 1):
        s = 0.0
        for i in range(n - lag):
            s += env[i] * env[i + lag]
        acf[lag] = s

    # preferred tempo window: 70-180 BPM -> 333-858 ms
    pref_lo = max(lag_min, int(round((60000.0 / 180.0) / res_ms)))
    pref_hi = min(lag_max, int(round((60000.0 / 70.0) / res_ms)))
    best_lag, best_val = None, -1.0
    for lag in range(pref_lo, pref_hi + 1):
        if acf[lag] > best_val:
            best_val, best_lag = acf[lag], lag
    if best_lag is None or best_val <= 0:                      # fallback: global max
        for lag in range(lag_min, lag_max + 1):
            if acf[lag] > best_val:
                best_val, best_lag = acf[lag], lag
    if best_lag is None:
        return 500.0

    # parabolic interpolation around the peak
    y0 = acf[best_lag - 1] if best_lag - 1 >= 0 else acf[best_lag]
    y1 = acf[best_lag]
    y2 = acf[best_lag + 1] if best_lag + 1 < len(acf) else acf[best_lag]
    denom = (y0 - 2 * y1 + y2)
    offset = 0.0
    if abs(denom) > 1e-12:
        offset = max(-0.5, min(0.5, 0.5 * (y0 - y2) / denom))
    beat = (best_lag + offset) * res_ms

    # light octave correction toward the preferred range
    while beat < 333.0:
        beat *= 2.0
    while beat > 858.0:
        beat /= 2.0
    return beat


def _level_of(g, spacings_asc):
    """Coarsest grid value v whose spacing divides g (within tol)."""
    for v, s in spacings_asc:                       # coarse (small v) first
        r = g / s
        if abs(r - round(r)) * s < 1e-3 + 1e-3 * s:
            return v
    return spacings_asc[-1][0]


def _duration_level(dur, spacings_asc):
    """Grid value v whose spacing makes the note last ~1-4 units."""
    best_v, best_err = None, None
    for v, s in spacings_asc:
        if s <= 0:
            continue
        units = dur / s
        k = max(1, min(4, int(round(units))))
        err = abs(dur - k * s) / max(dur, s)
        if best_err is None or err < best_err:
            best_err, best_v = err, v
    return best_v


def _level_distance(lv, v_star, spacings_asc):
    vs = [v for v, _ in spacings_asc]
    try:
        return abs(vs.index(lv) - vs.index(v_star))
    except ValueError:
        return 99


def _snap_one(t, finest, spacings_asc, back_window, forward_bonus,
              duration_hint, dur=None):
    """Forward-biased nearest-grid snap for a single (chord-reference) onset."""
    if finest <= 0:
        return t
    k0 = math.floor(t / finest + 1e-9)
    g0 = k0 * finest
    g1 = (k0 + 1) * finest
    if abs(t - g0) < 1e-6:                           # already on grid
        return g0

    d_bwd = t - g0                                   # > 0
    d_fwd = g1 - t                                   # > 0

    v_star = None
    if duration_hint and dur is not None and dur > 0:
        v_star = _duration_level(dur, spacings_asc)

    best, best_score = None, None
    for g, d, tag in ((g0, d_bwd, 'back'), (g1, d_fwd, 'fwd')):
        if tag == 'back' and d > back_window:        # backward only when mild
            continue
        score = -d
        if tag == 'fwd':
            score += forward_bonus * finest          # forward preferred on ties
        if v_star is not None:
            lv = _level_of(g, spacings_asc)
            if lv == v_star:
                score += 0.30 * finest
            elif _level_distance(lv, v_star, spacings_asc) == 1:
                score += 0.10 * finest
        if best_score is None or score > best_score + 1e-9:
            best_score, best = score, g

    if best is None:                                 # fallback (shouldn't happen)
        best = g0 if d_bwd <= d_fwd else g1
    return best


# --------------------------------------------------------------------------- #
#  Main function
# --------------------------------------------------------------------------- #
def snap_score_to_grid(
    score,
    grid_values,
    beat_ms=None,
    tempo_bpm=None,
    chord_tol_ms=None,
    back_factor=0.40,
    max_backward_ms=45.0,
    forward_bonus=0.10,
    duration_hint=True,
    preserve_order=True,
    resort=True,
    verbose=False,
):
    """
    Snap MIDI.py note onsets (millisecond, single flat track) to a musical grid.

    score         flat list of MIDI.py events in ms absolute time
                  (midi2ms_score output, ticks removed). Note events:
                  ['note', start_ms, dur_ms, channel, note, velocity, extra]
                  The 7th element (index 6) is a no-op and is passed through.
                  Non-note events are preserved unchanged.
    grid_values   note-value denominators, e.g. [1,2,4,8,16,32]
                  (1=whole, 2=half, 4=quarter, 8=eighth, 16=sixteenth,
                   32=thirty-second). Spacing for value v = (4*beat_ms)/v.
    beat_ms       quarter-note period in ms (auto-detected if None).
    tempo_bpm     overrides beat_ms  (beat_ms = 60000/tempo_bpm).
    chord_tol_ms  max onset spread to group asynchronous chord notes
                  (default min(45.0, finest*0.5)).
    back_factor   backward snaps allowed only within back_factor*finest
                  of a grid line (default 0.40).
    max_backward_ms  absolute cap on any backward shift (default 45.0).
    forward_bonus    fraction of finest added to forward candidate score
                  so forward wins near-ties (default 0.10).
    duration_hint    use note durations to bias toward the correct grid level.
    preserve_order   enforce non-decreasing snapped onsets.
    resort           re-sort output events by time.
    verbose          print detected tempo + statistics.

    Returns: new score (same structure), onsets snapped, durations + 7th
             element preserved.
    """
    if score is None:
        return []

    # tolerate a leading ticks integer, just in case it wasn't stripped
    leading_ticks = None
    events = score
    if len(events) > 0 and isinstance(events[0], int) and not isinstance(events[0], bool):
        leading_ticks = events[0]
        events = events[1:]

    # collect note events
    note_idxs = [i for i, e in enumerate(events)
                 if isinstance(e, list) and len(e) >= 6 and e[0] == 'note']
    if not note_idxs:
        return list(score)

    notes = [events[i] for i in note_idxs]
    onsets = [float(n[1]) for n in notes]
    durs = [float(n[2]) if len(n) > 2 else 0.0 for n in notes]

    # ---- 1. beat period -------------------------------------------------- #
    if tempo_bpm is not None:
        beat = 60000.0 / float(tempo_bpm)
    elif beat_ms is not None:
        beat = float(beat_ms)
    else:
        beat = _detect_beat_period_ms(onsets)

    # ---- 2. grid --------------------------------------------------------- #
    gvs = sorted(set(int(v) for v in grid_values if int(v) > 0))
    whole = 4.0 * beat
    spacings_asc = [(v, whole / v) for v in gvs]            # coarse (small v) first
    spacings_by_v = dict(spacings_asc)
    finest = spacings_by_v[gvs[-1]]
    if finest <= 0:
        finest = beat

    if chord_tol_ms is None:
        chord_tol_ms = min(45.0, finest * 0.5)
    back_window = min(max_backward_ms, finest * back_factor)

    # ---- 3. chord grouping (async onsets -> one reference) --------------- #
    order = sorted(range(len(notes)), key=lambda k: (onsets[k], k))
    groups, cur = [], [order[0]]
    for k in order[1:]:
        if onsets[k] - onsets[cur[0]] <= chord_tol_ms:
            cur.append(k)
        else:
            groups.append(cur); cur = [k]
    groups.append(cur)

    group_ref = []
    for g in groups:
        gs = sorted(onsets[k] for k in g)
        group_ref.append(gs[len(gs) // 2])                 # median, robust

    # ---- 4. snap each group reference ------------------------------------ #
    snapped = [None] * len(notes)
    prev_snapped = -float('inf')
    n_back = n_fwd = n_exact = 0
    for gi, g in enumerate(groups):
        t = group_ref[gi]
        # duration hint = longest note in the chord (defines its rhythmic level)
        gdur = max((durs[k] for k in g), default=None)
        st = _snap_one(t, finest, spacings_asc, back_window,
                       forward_bonus, duration_hint, dur=gdur)
        if preserve_order and st < prev_snapped:
            st = prev_snapped                              # never slip backwards
        if abs(st - t) < 1e-6:
            n_exact += 1
        elif st < t:
            n_back += 1
        else:
            n_fwd += 1
        group_ref[gi] = st
        prev_snapped = st
        for k in g:
            snapped[k] = st

    # ---- 5. rebuild events ---------------------------------------------- #
    new_events = [list(e) if isinstance(e, list) else e for e in events]
    for pos, nk in enumerate(note_idxs):
        new_events[nk][1] = int(round(snapped[pos]))       # onset only
        # duration (idx 2), channel/note/velocity, and the 7th element (idx 6)
        # are left untouched.

    if resort:
        def time_key(e):
            if isinstance(e, list) and len(e) >= 2:
                return e[1]
            return 0
        indexed = list(enumerate(new_events))
        indexed.sort(key=lambda p: (time_key(p[1]), p[0]))
        new_events = [e for _, e in indexed]

    if verbose:
        print(f"[snap_score_to_grid] beat={beat:.2f} ms "
              f"({60000.0/beat:.1f} BPM), finest grid={finest:.2f} ms, "
              f"chord_tol={chord_tol_ms:.2f} ms, back_window={back_window:.2f} ms")
        print(f"[snap_score_to_grid] snaps: exact={n_exact}, "
              f"forward={n_fwd}, backward={n_back}, groups={len(groups)}")

    if leading_ticks is not None:
        return [leading_ticks] + new_events
    return new_events