File size: 23,729 Bytes
abde881
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
"""
pmi_core.py — the PMI main line for HUMAN-REVISED MIDI.

revised MIDI -> notes (s,e,p) -> K-S key estimate (round-1 LOCKED) + snapping
            -> transpose to C  -> symbol string
            -> Needleman-Wunsch (affine gap) alignment -> PMI %

Key estimation uses estimate_key_locked(): the FIRST K-S estimate (made on the
original, information-complete notes) is locked and used to drive conservative
snapping of out-of-key slips. Later rounds clean the absolute pitch but do NOT
re-estimate the key — except in the rare 'wildly-off' case where a re-estimate is wildly
different (tonic far away AND the locked key's correlation has collapsed), when the
new key is adopted once. This matches the observation that round-1 is the most
trustworthy estimate and re-estimating after snapping tends to drift.

Ported (same constants/profiles/scoring) from the batch pipeline so numbers match.
Self-contained: depends only on numpy + pretty_midi.
"""
import numpy as np
import pretty_midi

# ---- Krumhansl-Kessler profiles + alphabet (verbatim from batch pipeline) ----
_MAJOR = np.array([6.35,2.23,3.48,2.33,4.38,4.09,2.52,5.19,2.39,3.66,2.29,2.88])
_MINOR = np.array([6.33,2.68,3.52,5.38,2.60,3.53,2.54,4.75,3.98,2.69,3.34,3.17])
NOTE_NAMES   = ['C','C#','D','D#','E','F','F#','G','G#','A','A#','B']
PC_TO_LETTER = ['C','d','D','e','E','F','g','G','a','A','b','B']
_MAJOR_INTERVALS = (0,2,4,5,7,9,11)
_MINOR_INTERVALS = (0,2,3,5,7,8,10)

MIN_DUR = 0.08
MARGIN_THRESH = 0.05

# scale-degree symbols (used only if encoding="scale_degree")
_DEG_MAJOR = {0:'1', 2:'2', 4:'3', 5:'4', 7:'5', 9:'6', 11:'7'}
_DEG_MINOR = {0:'1', 2:'2', 3:'3', 5:'4', 7:'5', 8:'6', 10:'7'}
_CHROMA_MAJOR = {1:'a', 3:'b', 6:'c', 8:'d', 10:'e'}
_CHROMA_MINOR = {1:'a', 4:'f', 6:'c', 9:'h', 11:'g'}


# ===========================================================================
# key estimation (plain K-S, NO snapping)
# ===========================================================================
def _pc(note): return int(note[2]) % 12

def scale_for(tonic, mode):
    iv = _MAJOR_INTERVALS if mode == 'major' else _MINOR_INTERVALS
    return {(tonic + i) % 12 for i in iv}

def note_weight(s, e, conf=None, min_dur=MIN_DUR):
    dur = max(e - s, 1e-3)
    w_dur = min(dur / min_dur, 1.0)
    w_conf = float(np.clip(conf, 0.0, 1.0)) if conf is not None else 1.0
    return dur * w_dur * w_conf

def pitch_class_histogram(notes, weighted=True, confs=None):
    hist = np.zeros(12)
    for i, note in enumerate(notes):
        conf = (float(confs[i]) if confs is not None
                else (float(note[3]) if len(note) > 3 else None))
        w = note_weight(note[0], note[1], conf) if weighted else max(note[1]-note[0], 1e-3)
        hist[_pc(note)] += w
    return hist

def _correlate_all_keys(hist):
    corrs = []
    for mode, prof in (('major', _MAJOR), ('minor', _MINOR)):
        for tonic in range(12):
            corrs.append((np.corrcoef(hist, np.roll(prof, tonic))[0, 1], tonic, mode))
    corrs.sort(reverse=True)
    return corrs

def estimate_key(notes, margin_thresh=MARGIN_THRESH, use_out_of_key_tiebreak=True,
                 weighted=True, confs=None):
    """Plain K-S key estimate (no snapping). Accepts (s,e,p[,conf]) notes.
    Returns dict(tonic, mode, scale, margin, confidence, out_of_key, hist, corrs) or None."""
    hist = pitch_class_histogram(notes, weighted=weighted, confs=confs)
    if hist.sum() == 0:
        return None
    corrs = _correlate_all_keys(hist)
    top1, top2 = corrs[0], corrs[1]
    margin = top1[0] - top2[0]
    def ook(tonic, mode):
        sc = scale_for(tonic, mode)
        return sum(hist[i] for i in range(12) if i not in sc)
    if use_out_of_key_tiebreak and margin < margin_thresh:
        _, c, tonic, mode = sorted((ook(t, m), c, t, m) for c, t, m in (top1, top2))[0]
    else:
        c, tonic, mode = top1
    return dict(tonic=tonic, mode=mode, scale=scale_for(tonic, mode), margin=margin,
                confidence=c, out_of_key=ook(tonic, mode), hist=hist, corrs=corrs)


# ===========================================================================
# snapping (verbatim from batch pipeline) + iterative estimator
# ===========================================================================
SNAP_THRESH = 0.60
KEEP_THRESH = 0.30
MAX_ROUNDS  = 4

def _conf(note, confs, idx):
    if confs is not None:
        return float(confs[idx])
    return float(note[3]) if len(note) > 3 else None

def slip_likelihood(idx, notes, key, confs=None, min_dur=MIN_DUR):
    if _pc(notes[idx]) in key['scale']:
        return 0.0
    s, e, p = notes[idx][0], notes[idx][1], _pc(notes[idx])
    dur = max(e - s, 1e-3)
    conf = _conf(notes[idx], confs, idx)
    e_conf = (1.0 - float(np.clip(conf, 0.0, 1.0))) if conf is not None else 0.0
    e_short = 1.0 - min(dur / min_dur, 1.0)
    pcs = [_pc(n) for n in notes]
    recur = pcs.count(p)
    left_in  = (idx == 0)            or (pcs[idx-1] in key['scale'])
    right_in = (idx == len(notes)-1) or (pcs[idx+1] in key['scale'])
    e_iso = 1.0 if (recur == 1 and left_in and right_in) else 0.0
    e_recur = 1.0 - 1.0 / recur
    score = (0.5*e_conf + 0.25*e_short + 0.25*e_iso) * (1.0 - 0.5*e_recur)
    return float(np.clip(score, 0.0, 1.0))

def is_minor_mode_member(p, key, notes):
    if key['mode'] != 'minor':
        return False
    raised = {(key['tonic']+11) % 12, (key['tonic']+9) % 12}
    if p % 12 not in raised:
        return False
    return [_pc(n) for n in notes].count(p % 12) >= 2

def _snap_target(p, scale):
    for cand in ((p-1) % 12, (p+1) % 12):
        if cand in scale:
            return cand
    return p

def classify_notes(notes, key, confs=None, snap_thresh=SNAP_THRESH, keep_thresh=KEEP_THRESH):
    scale = key['scale']
    out = []
    for i, note in enumerate(notes):
        p = _pc(note)
        if p in scale:
            out.append(dict(idx=i, label='diatonic', out_pc=p, confidence=1.0)); continue
        if is_minor_mode_member(p, key, notes):
            out.append(dict(idx=i, label='intentional_chromatic', out_pc=p, confidence=1.0)); continue
        sc = slip_likelihood(i, notes, key, confs)
        if sc >= snap_thresh:
            out.append(dict(idx=i, label='corrected', out_pc=_snap_target(p, scale), confidence=sc))
        elif sc <= keep_thresh:
            out.append(dict(idx=i, label='intentional_chromatic', out_pc=p, confidence=1.0-sc))
        else:
            out.append(dict(idx=i, label='uncertain', out_pc=p, confidence=sc))
    return out

def corrected_pitches(notes, classification):
    out = [_pc(n) for n in notes]
    for rec in classification:
        out[rec['idx']] = rec['out_pc']
    return out


# tonic distance on the circle of pitch classes (0..6 semitones)
def _tonic_dist(a, b):
    d = abs(int(a) - int(b)) % 12
    return min(d, 12 - d)

def estimate_key_locked(notes, max_rounds=MAX_ROUNDS, snap_thresh=SNAP_THRESH,
                        keep_thresh=KEEP_THRESH, confs=None,
                        outlier_tonic_dist=3, outlier_corr_drop=0.15, **est_kw):
    """Iterative key + snapping, but the ROUND-1 key is LOCKED and used to drive
    all snapping. Subsequent rounds snap notes (cleaning absolute pitch) but do
    NOT overwrite the key — UNLESS a re-estimate is wildly different from round 1
    (the 'wildly-off' guard): tonic more than `outlier_tonic_dist` semitones away AND the
    round-1 key's correlation has dropped by more than `outlier_corr_drop`. In that
    rare case the new key is adopted once and re-locked.

    Rationale: round-1 is estimated on the original, information-complete notes and
    is the most trustworthy; re-estimating after snapping tends to drift. Returns
    (key, classification). Caller's notes are never mutated.
    """
    work = [list(n) for n in notes]
    orig_pc = [_pc(n) for n in notes]
    corrected = {}; frozen = set()

    key = estimate_key(work, confs=confs, **est_kw)          # ROUND-1 key
    if key is None:
        return None, []
    locked = key                                            # <-- lock it

    cls = classify_notes(work, locked, confs, snap_thresh, keep_thresh)
    for _ in range(max_rounds):
        changed = False
        for rec in cls:
            i = rec['idx']
            if rec['label'] == 'corrected' and i not in frozen and work[i][2] % 12 != rec['out_pc']:
                work[i][2] = rec['out_pc']
                corrected[i] = (rec['out_pc'], rec['confidence'])
                frozen.add(i)
                changed = True

        # re-estimate ONLY to check for the 'wildly-off' (wildly-off) case; do not adopt by default
        check = estimate_key(work, confs=confs, **est_kw)
        if check is not None:
            far_tonic = _tonic_dist(check['tonic'], locked['tonic']) > outlier_tonic_dist
            # correlation the locked key now gets on the (snapped) histogram
            locked_corr = next((c for c, t, m in check['corrs']
                                if t == locked['tonic'] and m == locked['mode']), None)
            corr_drop = (locked_corr is not None) and (check['confidence'] - locked_corr > outlier_corr_drop)
            if far_tonic and corr_drop:
                locked = check                              # adopt + re-lock (rare)

        # keep snapping against the LOCKED key
        cls = classify_notes(work, locked, confs, snap_thresh, keep_thresh)
        if not changed:
            break

    for rec in cls:
        i = rec['idx']
        rec['orig_pc'] = orig_pc[i]
        if i in corrected:
            rec['label'] = 'corrected'
            rec['out_pc'], rec['confidence'] = corrected[i]
    return locked, cls


# ===========================================================================
# transpose to C + encode
# ===========================================================================
def transpose_to_C(notes, tonic):
    """Every note's pitch class shifted so the estimated tonic maps to C (pc 0).
    Input notes are (s,e,p[,...]); output is a time-ordered pitch-class list."""
    return [(int(p) - int(tonic)) % 12 for _, _, p, *_ in sorted(notes)]

def encode(pcs):
    return "".join(PC_TO_LETTER[p % 12] for p in pcs)

def encode_scale_degree(pcs, mode):
    diatonic, chroma = (_DEG_MAJOR, _CHROMA_MAJOR) if mode == 'major' else (_DEG_MINOR, _CHROMA_MINOR)
    return ''.join(diatonic.get(pc % 12) or chroma.get(pc % 12, 'x') for pc in pcs)


# ===========================================================================
# alignment + PMI  (Needleman-Wunsch affine gap; Savage GOP=12, GEP=6)
# ===========================================================================
def nw_align(a, b, match=2, mismatch=-2, gop=12, gep=6):
    n, m = len(a), len(b); NEG = -1e9
    M = np.full((n+1, m+1), NEG); Ix = np.full((n+1, m+1), NEG); Iy = np.full((n+1, m+1), NEG); M[0,0] = 0
    for i in range(1, n+1): Ix[i,0] = -gop - (i-1)*gep
    for j in range(1, m+1): Iy[0,j] = -gop - (j-1)*gep
    for i in range(1, n+1):
        ai = a[i-1]
        for j in range(1, m+1):
            sc = match if ai == b[j-1] else mismatch
            M[i,j]  = max(M[i-1,j-1], Ix[i-1,j-1], Iy[i-1,j-1]) + sc
            Ix[i,j] = max(M[i-1,j] - gop, Ix[i-1,j] - gep)
            Iy[i,j] = max(M[i,j-1] - gop, Iy[i,j-1] - gep)
    i, j = n, m; al = []; bl = []
    st = max((M[n,m], 'M'), (Ix[n,m], 'X'), (Iy[n,m], 'Y'))[1]
    while i > 0 or j > 0:
        if i == 0: st = 'Y'
        elif j == 0: st = 'X'
        if st == 'M':
            al.append(a[i-1]); bl.append(b[j-1])
            st = max((M[i-1,j-1], 'M'), (Ix[i-1,j-1], 'X'), (Iy[i-1,j-1], 'Y'))[1]; i -= 1; j -= 1
        elif st == 'X':
            al.append(a[i-1]); bl.append('-')
            st = 'X' if (i > 1 and abs(Ix[i,j] - (Ix[i-1,j] - gep)) < 1e-6) else 'M'; i -= 1
        else:
            al.append('-'); bl.append(b[j-1])
            st = 'Y' if (j > 1 and abs(Iy[i,j] - (Iy[i,j-1] - gep)) < 1e-6) else 'M'; j -= 1
    return "".join(reversed(al)), "".join(reversed(bl))

def pmi(seqA, seqB, **kw):
    if not seqA or not seqB:
        return dict(PMI=float('nan'), ID=0, alnA='', alnB='')
    A, B = nw_align(seqA, seqB, **kw)
    ID = sum(1 for x, y in zip(A, B) if x == y and x != '-')
    return dict(PMI=100.0 * ID / ((len(seqA) + len(seqB)) / 2.0), ID=ID, alnA=A, alnB=B)


# ===========================================================================
# MIDI -> notes, and the per-song + pairwise drivers
# ===========================================================================
def notes_from_midi(midi_path, instrument_index=None):
    """Revised MIDI -> time-sorted [(start, end, pitch, conf), ...].

    The editor encodes confidence as velocity (conf = velocity/127); hand-edited /
    verified notes are written at velocity 127 (conf=1). Reading it back lets the
    key estimator treat verified notes as high-confidence (never snapped) and only
    let low-confidence auto-transcribed leftovers be eligible for snapping.
    """
    pm = pretty_midi.PrettyMIDI(midi_path)
    insts = [i for i in pm.instruments if not i.is_drum]
    if not insts:
        raise ValueError("No non-drum instrument tracks in MIDI.")
    if instrument_index is not None:
        insts = [pm.instruments[instrument_index]]
    notes = [(float(n.start), float(n.end), int(n.pitch), float(n.velocity) / 127.0)
             for inst in insts for n in inst.notes]
    if not notes:
        raise ValueError("MIDI contains no notes.")
    notes.sort(key=lambda t: (t[0], t[1], t[2]))
    return notes

def midi_to_symbols(midi_path, encoding="pitch_class", weighted=True,
                    instrument_index=None, snap=True):
    """revised MIDI -> dict(notes, key, transposed, symbols).

    snap=True (default): use estimate_key_locked — the round-1 K-S key is locked
    and drives conservative snapping of out-of-key slips; the cleaned (snapped)
    absolute pitches are what get transposed to C. snap=False: plain single K-S.
    """
    notes = notes_from_midi(midi_path, instrument_index=instrument_index)

    if snap:
        key, classification = estimate_key_locked(notes, confs=None, weighted=weighted)
        if key is None:
            raise ValueError("Key estimation returned None (empty pitch-class histogram).")
        # transpose the FINAL (snapped) pitches, in note order
        final_pcs = corrected_pitches([list(n) for n in notes], classification)
        transposed = [(pc - key["tonic"]) % 12 for pc in final_pcs]
    else:
        key = estimate_key(notes, weighted=weighted)
        if key is None:
            raise ValueError("Key estimation returned None (empty pitch-class histogram).")
        classification = []
        transposed = transpose_to_C(notes, key["tonic"])

    if encoding == "pitch_class":
        symbols = encode(transposed)
    elif encoding == "scale_degree":
        symbols = encode_scale_degree(transposed, key["mode"])
    else:
        raise ValueError(f"Unknown encoding: {encoding}")
    n_corr = sum(1 for r in classification if r.get('label') == 'corrected')
    return dict(notes=notes, key=key, transposed=transposed, symbols=symbols,
                n_corrected=n_corr)

def key_name(key):
    return f"{NOTE_NAMES[key['tonic']]} {key['mode']}"

def key_name_tm(tonic, mode):
    return f"{NOTE_NAMES[int(tonic) % 12]} {mode}"


# ---------------------------------------------------------------------------
# transpose a song to C using a GIVEN key (tonic+mode), with snapping vs that key.
# Used when the user picks a candidate / overrides the key.
# ---------------------------------------------------------------------------
def symbols_for_given_key(notes, tonic, mode, encoding="pitch_class", weighted=True):
    """Snap notes against the given (tonic,mode), transpose to C, encode.
    Verified (high-confidence) notes stay protected; only low-confidence
    out-of-key slips are snapped."""
    key = dict(tonic=int(tonic) % 12, mode=mode, scale=scale_for(int(tonic) % 12, mode))
    cls = classify_notes([list(n) for n in notes], key)
    final_pcs = corrected_pitches([list(n) for n in notes], cls)
    transposed = [(pc - key["tonic"]) % 12 for pc in final_pcs]
    sym = encode(transposed) if encoding == "pitch_class" else encode_scale_degree(transposed, mode)
    n_oot = sum(1 for pc in transposed if pc not in _MAJOR_INTERVALS) if mode == 'major' \
            else sum(1 for pc in transposed if pc not in _MINOR_INTERVALS)
    return dict(tonic=key["tonic"], mode=mode, transposed=transposed, symbols=sym,
                n_out_of_key=n_oot)


def top_candidate_keys(notes, k=2, weighted=True):
    """Round-1 K-S top-k (tonic,mode) candidates, most likely first.
    The first is the locked estimate; the second is the runner-up used for the
    'listen and pick' A/B comparison."""
    hist = pitch_class_histogram(notes, weighted=weighted)
    if hist.sum() == 0:
        return []
    corrs = _correlate_all_keys(hist)   # already sorted, best first
    return [(t, m, float(c)) for c, t, m in corrs[:k]]


# ---------------------------------------------------------------------------
# absolute-pitch similarity (NO transposition) — used only to decide whether to
# OFFER soft key-unification. Never part of the main PMI.
# ---------------------------------------------------------------------------
def absolute_symbols(notes):
    """Encode notes by raw pitch class (no transposition)."""
    return encode([_pc(n) for n in sorted(notes, key=lambda t: (t[0], t[1], t[2]))])

def absolute_similarity(notesA, notesB):
    return pmi(absolute_symbols(notesA), absolute_symbols(notesB))["PMI"]


# ---------------------------------------------------------------------------
# per-song analysis for the /analyze endpoint
# ---------------------------------------------------------------------------
def analyze_one_midi(midi_path, encoding="pitch_class"):
    """One revised MIDI -> locked key estimate + top-2 candidates, each candidate
    transposed-to-C so the front-end can play them back for the listen-and-pick.
    Also returns the absolute (untransposed) pitch-class sequence for soft-unify."""
    notes = notes_from_midi(midi_path)
    base = midi_to_symbols(midi_path, encoding=encoding)        # locked estimate + snap
    cands = top_candidate_keys(notes, k=2)
    cand_out = []
    for t, m, corr in cands:
        s = symbols_for_given_key(notes, t, m, encoding=encoding)
        cand_out.append(dict(tonic=t, mode=m, corr=corr, name=key_name_tm(t, m),
                             symbols=s["symbols"], transposed=s["transposed"],
                             n_out_of_key=s["n_out_of_key"]))
    return dict(
        n_notes=len(notes),
        key=dict(tonic=base["key"]["tonic"], mode=base["key"]["mode"],
                 name=key_name(base["key"]), margin=float(base["key"]["margin"]),
                 confidence=float(base["key"]["confidence"])),
        symbols=base["symbols"],
        transposed=base["transposed"],
        abs_pcs=[_pc(n) for n in notes],          # for soft-unify similarity
        candidates=cand_out,
        encoding=encoding,
    )


# ---------------------------------------------------------------------------
# pairwise PMI from two analyses (+ optional user-forced keys + soft-unify)
# ---------------------------------------------------------------------------
def pmi_pair(notesA, notesB, encoding="pitch_class",
             forced_A=None, forced_B=None,
             soft_unify_threshold=70.0,
             labelA="Song A", labelB="Song B"):
    """Compute the main PMI (each song its own key, locked estimate or user-forced),
    plus a SEPARATE soft-unified PMI when the two melodies are already similar in
    ABSOLUTE pitch (>= threshold). forced_A/forced_B are optional (tonic, mode)."""
    # --- per-song key + symbols (main line) ---
    if forced_A is not None:
        A = symbols_for_given_key(notesA, forced_A[0], forced_A[1], encoding=encoding)
        keyA = dict(tonic=A["tonic"], mode=A["mode"], margin=None,
                    confidence=None, scale=scale_for(A["tonic"], A["mode"]))
    else:
        kA, clsA = estimate_key_locked(notesA)
        keyA = kA
        pcsA = corrected_pitches([list(n) for n in notesA], clsA)
        A = dict(tonic=kA["tonic"], mode=kA["mode"],
                 transposed=[(pc - kA["tonic"]) % 12 for pc in pcsA])
        A["symbols"] = encode(A["transposed"]) if encoding == "pitch_class" \
                       else encode_scale_degree(A["transposed"], kA["mode"])

    if forced_B is not None:
        B = symbols_for_given_key(notesB, forced_B[0], forced_B[1], encoding=encoding)
        keyB = dict(tonic=B["tonic"], mode=B["mode"], margin=None,
                    confidence=None, scale=scale_for(B["tonic"], B["mode"]))
    else:
        kB, clsB = estimate_key_locked(notesB)
        keyB = kB
        pcsB = corrected_pitches([list(n) for n in notesB], clsB)
        B = dict(tonic=kB["tonic"], mode=kB["mode"],
                 transposed=[(pc - kB["tonic"]) % 12 for pc in pcsB])
        B["symbols"] = encode(B["transposed"]) if encoding == "pitch_class" \
                       else encode_scale_degree(B["transposed"], kB["mode"])

    main = pmi(A["symbols"], B["symbols"])

    result = {
        "PMI": main["PMI"], "ID": main["ID"], "alnA": main["alnA"], "alnB": main["alnB"],
        "symbolsA": A["symbols"], "symbolsB": B["symbols"],
        "keyA": key_name_tm(A["tonic"], A["mode"]),
        "keyB": key_name_tm(B["tonic"], B["mode"]),
        "n_notes_A": len(notesA), "n_notes_B": len(notesB),
        "lenA": len(A["symbols"]), "lenB": len(B["symbols"]),
        "encoding": encoding, "labelA": labelA, "labelB": labelB,
        "forcedA": forced_A is not None, "forcedB": forced_B is not None,
    }

    # --- soft key-unification (SEPARATE, only OFFERED when abs-similarity high) ---
    abs_sim = absolute_similarity(notesA, notesB)
    result["abs_similarity"] = abs_sim
    result["soft_unify_threshold"] = soft_unify_threshold
    result["soft_unify_applicable"] = abs_sim >= soft_unify_threshold
    if result["soft_unify_applicable"]:
        # weight by margin (how decisively each song's own key beat the runner-up);
        # the more decisive key wins as the shared key. forced keys count as fully decisive.
        mA = 1.0 if forced_A is not None else float(keyA.get("margin") or 0.0)
        mB = 1.0 if forced_B is not None else float(keyB.get("margin") or 0.0)
        if mA >= mB:
            uni_t, uni_m, src = A["tonic"], A["mode"], labelA
        else:
            uni_t, uni_m, src = B["tonic"], B["mode"], labelB
        UA = symbols_for_given_key(notesA, uni_t, uni_m, encoding=encoding)
        UB = symbols_for_given_key(notesB, uni_t, uni_m, encoding=encoding)
        uni = pmi(UA["symbols"], UB["symbols"])
        result["soft_unify"] = {
            "shared_key": key_name_tm(uni_t, uni_m),
            "shared_key_source": src,
            "PMI": uni["PMI"], "ID": uni["ID"],
            "alnA": uni["alnA"], "alnB": uni["alnB"],
            "symbolsA": UA["symbols"], "symbolsB": UB["symbols"],
            "weight_margin_A": mA, "weight_margin_B": mB,
        }
    return result


def pmi_from_two_midis(midi_path_A, midi_path_B, encoding="pitch_class",
                       forced_A=None, forced_B=None, soft_unify_threshold=70.0,
                       labelA="Song A", labelB="Song B"):
    """Convenience: read two MIDIs and run pmi_pair."""
    notesA = notes_from_midi(midi_path_A)
    notesB = notes_from_midi(midi_path_B)
    return pmi_pair(notesA, notesB, encoding=encoding,
                    forced_A=forced_A, forced_B=forced_B,
                    soft_unify_threshold=soft_unify_threshold,
                    labelA=labelA, labelB=labelB)