sample_extractor.py

#!/usr/bin/env python3
"""
Sample Extractor v9 — Production quality.

Major improvements:
  - SuperFlux onset detection (Böck & Widmer 2013) tuned for 140+ BPM
  - Transient-only NCC: compare first 25ms (attack) not full clip
  - Mel-spectrogram cosine pre-filter: eliminates 70% of NCC pairs
  - Singleton cluster merging into nearest neighbor
  - MIDI quantization to BPM grid (16th/32nd note snap)
  - Parallel NCC computation via ThreadPoolExecutor
  - All previous features: caching, target range, parameter locking
"""

import argparse, json, os, sys, warnings, hashlib, itertools
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
from concurrent.futures import ThreadPoolExecutor

import librosa, numpy as np, soundfile as sf, torch
from scipy.signal import fftconvolve
from scipy.spatial.distance import cdist

warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)

@dataclass
class Hit:
    audio: np.ndarray; sr: int; onset_time: float; duration: float; index: int
    rms_energy: float = 0.0; spectral_centroid: float = 0.0
    label: str = ""; cluster_id: int = -1
    def save(self, path): sf.write(path, self.audio, self.sr, subtype='PCM_24')

@dataclass
class Cluster:
    cluster_id: int; label: str; hits: list = field(default_factory=list)
    best_hit_idx: int = 0; synthesized: Optional[np.ndarray] = None; midi_note: int = 60
    @property
    def best_hit(self): return self.hits[self.best_hit_idx]
    @property
    def count(self): return len(self.hits)

DEMUCS_MODELS = ["htdemucs","htdemucs_ft","htdemucs_6s","mdx","mdx_extra","mdx_extra_q"]
DEMUCS_STEMS = {"htdemucs":["drums","bass","other","vocals"],"htdemucs_ft":["drums","bass","other","vocals"],
    "htdemucs_6s":["drums","bass","other","vocals","guitar","piano"],
    "mdx":["drums","bass","other","vocals"],"mdx_extra":["drums","bass","other","vocals"],
    "mdx_extra_q":["drums","bass","other","vocals"]}

# ─── Cache ────────────────────────────────────────────────────────────────────
_cache = {}
def _ahash(a): return hashlib.md5(a[:4000].tobytes()).hexdigest()
def cache_get(k): return _cache.get(k)
def cache_set(k,v): _cache[k]=v; return v
def cache_clear(): _cache.clear()

# ─── Stage 1: Stem ───────────────────────────────────────────────────────────
def extract_stem(audio_path,stem="drums",device="cpu",model_name="htdemucs_ft",shifts=1,overlap=0.25,progress_cb=None):
    if stem=="all":
        y,sr=librosa.load(audio_path,sr=44100,mono=True)
        if progress_cb:
            progress_cb({"stage":"stem","completed_units":1,"total_units":1,"fraction":1.0,"detail":"loaded full mix"})
        return y.astype(np.float32),sr
    with open(audio_path,'rb') as f: fh=hashlib.md5(f.read(200000)).hexdigest()
    ck=("stem",fh,stem,model_name,shifts,overlap); c=cache_get(ck)
    if c:
        print(f"[Stem] cached")
        if progress_cb:
            progress_cb({"stage":"stem","completed_units":1,"total_units":1,"fraction":1.0,"detail":"memory-cache hit"})
        return c
    from demucs.pretrained import get_model; from demucs.apply import apply_model
    print(f"[Stem] {stem} via {model_name}...")
    model=get_model(model_name); model.eval().to(device); sr=model.samplerate
    if stem not in model.sources: raise ValueError(f"'{stem}' not in {model.sources}")
    a,_=librosa.load(audio_path,sr=sr,mono=False)
    if a.ndim==1: a=np.stack([a,a])
    elif a.shape[0]>2: a=a[:2]
    elif a.shape[0]==1: a=np.concatenate([a,a],axis=0)
    wav=torch.from_numpy(a).float().unsqueeze(0).to(device)

    class _CompletedFuture:
        def __init__(self, value): self._value = value
        def result(self): return self._value

    class _ProgressPool:
        def __init__(self, total_units):
            self.total_units=max(1,int(total_units))
            self.completed_units=0
        def submit(self, fn, *args, **kwargs):
            value=fn(*args, **kwargs)
            self.completed_units += 1
            if progress_cb:
                progress_cb({
                    "stage":"stem",
                    "completed_units":self.completed_units,
                    "total_units":self.total_units,
                    "fraction":min(1.0,self.completed_units/self.total_units),
                    "detail":f"Demucs chunk {self.completed_units}/{self.total_units}",
                })
            return _CompletedFuture(value)

    segment=float(getattr(model,"segment",8.0) or 8.0)
    segment_length=max(1,int(sr*segment))
    stride=max(1,int((1-float(overlap))*segment_length))
    chunk_count=max(1,len(range(0,wav.shape[-1],stride)))
    shift_runs=max(1,int(shifts) if int(shifts)>0 else 1)
    progress_pool=_ProgressPool(chunk_count*shift_runs) if progress_cb else None
    with torch.no_grad():
        src=apply_model(model,wav,device=device,shifts=shifts,split=True,overlap=overlap,pool=progress_pool)
    r=src[0,model.sources.index(stem)].mean(dim=0).cpu().numpy()
    if progress_cb:
        progress_cb({"stage":"stem","completed_units":chunk_count*shift_runs,"total_units":chunk_count*shift_runs,"fraction":1.0,"detail":"Demucs stem complete"})
    print(f"  ✓ {len(r)/sr:.1f}s"); return cache_set(ck,(r.astype(np.float32),sr))

# ─── Stage 2: SuperFlux Onset Detection ──────────────────────────────────────
def detect_onsets(y,sr,pre_pad=0.003,min_dur=0.02,max_dur=1.5,min_gap=0.03,
                  energy_threshold_db=-35.0,mode="auto",onset_delta=0.12,
                  hop_length=220,n_fft=1024,max_size=3,lag=1,backtrack=True):
    """SuperFlux-based onset detection. hop_length=220 (~5ms) for 140+ BPM."""
    print(f"[Onsets] mode={mode}, delta={onset_delta}, hop={hop_length}")

    if mode == "percussive":
        S = librosa.feature.melspectrogram(y=y,sr=sr,n_fft=n_fft,hop_length=hop_length,
                                            n_mels=128,fmax=8000,power=1)
        oe = librosa.onset.onset_strength(S=librosa.power_to_db(S,ref=np.max),sr=sr,
                                           hop_length=hop_length,lag=lag,max_size=max_size)
    elif mode == "harmonic":
        yh,_=librosa.effects.hpss(y)
        S = librosa.feature.melspectrogram(y=yh,sr=sr,n_fft=n_fft,hop_length=hop_length,
                                            n_mels=128,fmax=8000,power=1)
        oe = librosa.onset.onset_strength(S=librosa.power_to_db(S,ref=np.max),sr=sr,
                                           hop_length=hop_length,lag=2,max_size=5)
    elif mode == "broadband":
        oe = librosa.onset.onset_strength(y=y,sr=sr,hop_length=hop_length)
    else:  # auto: SuperFlux multi-band
        yh,_=librosa.effects.hpss(y)
        def _sf(audio,fmin=None,fmax=None,ms=3,l=1):
            S=librosa.feature.melspectrogram(y=audio,sr=sr,n_fft=n_fft,hop_length=hop_length,
                n_mels=128,fmin=fmin or 20,fmax=fmax or sr//2,power=1)
            return librosa.onset.onset_strength(S=librosa.power_to_db(S,ref=np.max),
                sr=sr,hop_length=hop_length,lag=l,max_size=ms)
        def _n(x): m=x.max(); return x/m if m>0 else x
        oe = np.maximum.reduce([_n(_sf(y,20,300)), _n(_sf(y,300,4000)),
                                 _n(_sf(y,4000,16000)), _n(_sf(yh,l=2,ms=5))])

    wait = max(1, int(min_gap * sr / hop_length))
    fr = librosa.onset.onset_detect(onset_envelope=oe,sr=sr,hop_length=hop_length,
        wait=wait,pre_avg=3,post_avg=3,pre_max=3,post_max=5,
        delta=onset_delta,backtrack=backtrack,units='frames')
    times = librosa.frames_to_time(fr,sr=sr,hop_length=hop_length)
    print(f"  Raw: {len(times)}")

    thr=10**(energy_threshold_db/20); hits=[]
    for i,t in enumerate(times):
        s=max(0,int((t-pre_pad)*sr))
        e=min(int(times[i+1]*sr) if i+1<len(times) else len(y), s+int(max_dur*sr))
        seg=y[s:e]
        if len(seg)<int(min_dur*sr): continue
        rms=np.sqrt(np.mean(seg**2))
        if rms<thr: continue
        fl=min(int(0.003*sr),len(seg)//4)
        if fl>0: seg=seg.copy(); seg[-fl:]*=np.linspace(1,0,fl)
        sc=float(librosa.feature.spectral_centroid(y=seg,sr=sr).mean())
        hits.append(Hit(audio=seg,sr=sr,onset_time=t,duration=len(seg)/sr,
                        index=len(hits),rms_energy=float(rms),spectral_centroid=sc))
    print(f"  ✓ {len(hits)} hits"); return hits

# ─── Stage 3: Classification ─────────────────────────────────────────────────
LABEL_RULES=[("kick",lambda lr,mr,hr,c,zcr,d:lr>0.5 and c<800),
    ("hihat_closed",lambda lr,mr,hr,c,zcr,d:hr>0.35 and c>4000 and d<0.15),
    ("hihat_open",lambda lr,mr,hr,c,zcr,d:hr>0.35 and c>4000 and d>=0.15),
    ("cymbal",lambda lr,mr,hr,c,zcr,d:hr>0.25 and c>3000),
    ("snare",lambda lr,mr,hr,c,zcr,d:mr>0.4 and zcr>0.1 and c>1000),
    ("tom",lambda lr,mr,hr,c,zcr,d:lr>0.3 and mr>0.3 and c<1500),
    ("bass",lambda lr,mr,hr,c,zcr,d:lr>0.6 and c<400 and d>0.2),
    ("vocal",lambda lr,mr,hr,c,zcr,d:mr>0.5 and 500<c<3000 and zcr<0.15),
    ("bright",lambda lr,mr,hr,c,zcr,d:c>2500),("mid",lambda lr,mr,hr,c,zcr,d:c>800)]
def classify_hit(hit):
    y,sr=hit.audio,hit.sr; D=np.abs(librosa.stft(y,n_fft=2048))
    f=librosa.fft_frequencies(sr=sr,n_fft=2048)
    le=np.sum(D[(f>=20)&(f<200)]**2); me=np.sum(D[(f>=200)&(f<4000)]**2); he=np.sum(D[(f>=4000)]**2)
    t=le+me+he+1e-10; lr,mr,hr=le/t,me/t,he/t
    zcr=float(librosa.feature.zero_crossing_rate(y=y).mean())
    for n,fn in LABEL_RULES:
        if fn(lr,mr,hr,hit.spectral_centroid,zcr,hit.duration): return n
    return "other"
def classify_hits(hits):
    print(f"[Classify] {len(hits)} hits...")
    for h in hits: h.label=classify_hit(h)
    c=defaultdict(int)
    for h in hits: c[h.label]+=1
    for l,n in sorted(c.items(),key=lambda x:-x[1]): print(f"    {l}: {n}")
    return hits

# ─── Stage 4: Two-stage clustering (mel pre-filter + transient NCC) ──────────

def _extract_transient(y, sr, onset_time, attack_ms=25.0):
    """Extract the 25ms attack transient — the perceptually defining part."""
    n = int(attack_ms * sr / 1000)
    s = max(0, int(onset_time * sr) - int(0.002 * sr))  # 2ms pre-onset
    return y[s:s+n] if s+n <= len(y) else np.pad(y[s:], (0, s+n-len(y)))

def _transient_ncc(a, b):
    """NCC on equal-length transients. Fast: ~0.02ms per pair for 25ms clips."""
    n = min(len(a), len(b)); a,b = a[:n].copy(), b[:n].copy()
    a -= a.mean(); b -= b.mean()
    na, nb = np.linalg.norm(a), np.linalg.norm(b)
    if na < 1e-8 or nb < 1e-8: return 0.0
    a /= na; b /= nb
    cc = np.fft.irfft(np.fft.rfft(a, n=2*n) * np.conj(np.fft.rfft(b, n=2*n)))
    return float(np.max(np.abs(cc)))

def _mel_fingerprint(y, sr, onset_time, window_ms=100, n_mels=64):
    """Compact mel fingerprint for cosine pre-filtering."""
    s = max(0, int(onset_time * sr))
    seg = y[s:s+int(window_ms/1000*sr)]
    if len(seg) < 512: seg = np.pad(seg, (0, 512-len(seg)))
    S = librosa.feature.melspectrogram(y=seg,sr=sr,n_fft=512,hop_length=256,n_mels=n_mels,power=2)
    v = librosa.power_to_db(S, ref=np.max).flatten().astype(np.float32)
    norm = np.linalg.norm(v)
    return v / norm if norm > 1e-8 else v

def build_similarity_matrix(hits, sr, audio, attack_ms=25.0, mel_threshold=0.75):
    """Two-stage: mel cosine pre-filter → transient NCC only on candidates."""
    N = len(hits)
    key = ("sim_matrix", tuple(_ahash(h.audio) for h in hits), attack_ms)
    c = cache_get(key)
    if c is not None: print(f"  Cached similarity matrix"); return c

    print(f"  Stage A: Mel fingerprints ({N} hits)...")
    fps = [_mel_fingerprint(audio, sr, h.onset_time) for h in hits]
    # Pad fingerprints to same length
    maxl = max(len(f) for f in fps)
    fps_padded = np.array([np.pad(f, (0, maxl-len(f))) if len(f)<maxl else f for f in fps])
    mel_sim = 1.0 - cdist(fps_padded, fps_padded, metric='cosine')  # (N,N)

    print(f"  Stage B: Transient NCC on {int((mel_sim > mel_threshold).sum()/2)} candidates...")
    transients = [_extract_transient(audio, sr, h.onset_time, attack_ms) for h in hits]
    sim = np.zeros((N, N), dtype=np.float32)
    np.fill_diagonal(sim, 1.0)

    # Only compute NCC for pairs that pass mel pre-filter
    pairs = [(i,j) for i in range(N) for j in range(i+1,N) if mel_sim[i,j] > mel_threshold]
    print(f"  {len(pairs)} NCC pairs (of {N*(N-1)//2} total, saved {100-len(pairs)*100/max(1,N*(N-1)//2):.0f}%)")

    # Parallel NCC for large pair counts
    if len(pairs) > 500:
        def _compute_chunk(chunk):
            return [(i,j,_transient_ncc(transients[i],transients[j])) for i,j in chunk]
        chunks = [pairs[k:k+64] for k in range(0,len(pairs),64)]
        with ThreadPoolExecutor() as ex:
            for res in ex.map(_compute_chunk, chunks):
                for i,j,v in res: sim[i,j]=sim[j,i]=v
    else:
        for i,j in pairs: v=_transient_ncc(transients[i],transients[j]); sim[i,j]=sim[j,i]=v

    return cache_set(key, sim)

def _labels_to_clusters(labels, hits):
    cm=defaultdict(list)
    for i,l in enumerate(labels): cm[l].append(i)
    clusters=[]
    for _,idx in sorted(cm.items()):
        v=defaultdict(int)
        for i in idx: v[hits[i].label]+=1
        maj=max(v,key=v.get); ex=sum(1 for c in clusters if c.label.rsplit('_',1)[0]==maj)
        clusters.append(Cluster(cluster_id=len(clusters),label=f"{maj}_{ex}",hits=[hits[i] for i in idx]))
    clusters.sort(key=lambda c:c.count,reverse=True)
    for i,c in enumerate(clusters): c.cluster_id=i
    return clusters

def _merge_singletons(clusters, sim_matrix, hits, merge_ratio=2.0):
    """Merge singleton clusters into nearest multi-hit cluster if close enough."""
    if len(clusters) < 2: return clusters
    multi = [c for c in clusters if c.count > 1]
    singles = [c for c in clusters if c.count == 1]
    if not multi or not singles: return clusters

    # For each singleton, find nearest multi-hit cluster
    merged = 0
    for sc in singles:
        si = sc.hits[0].index  # hit index in original list
        best_dist, best_cluster = float('inf'), None
        for mc in multi:
            dists = [1.0 - sim_matrix[si, h.index] for h in mc.hits]
            avg_dist = np.mean(dists)
            # Cluster internal radius = mean pairwise distance
            if len(mc.hits) >= 2:
                internal = np.mean([1.0 - sim_matrix[mc.hits[a].index, mc.hits[b].index]
                                     for a in range(len(mc.hits)) for b in range(a+1,len(mc.hits))])
            else:
                internal = 0.3  # default
            if avg_dist < internal * merge_ratio and avg_dist < best_dist:
                best_dist, best_cluster = avg_dist, mc
        if best_cluster is not None:
            best_cluster.hits.append(sc.hits[0])
            merged += 1
        else:
            multi.append(sc)  # keep as separate cluster

    if merged > 0: print(f"  Merged {merged} singletons")
    # Re-sort and re-index
    multi.sort(key=lambda c: c.count, reverse=True)
    for i,c in enumerate(multi): c.cluster_id = i
    return multi


def _cosine(a, b):
    """Fast cosine similarity for normalized or unnormalized one-dimensional vectors."""
    n = min(len(a), len(b))
    if n <= 0:
        return 0.0
    av = a[:n]
    bv = b[:n]
    denom = float(np.linalg.norm(av) * np.linalg.norm(bv))
    if denom < 1e-8:
        return 0.0
    return float(np.dot(av, bv) / denom)


def _retitle_clusters(clusters):
    """Sort, re-index, and make labels stable after incremental assignment."""
    clusters.sort(key=lambda c: c.count, reverse=True)
    seen = defaultdict(int)
    for i, c in enumerate(clusters):
        c.cluster_id = i
        majority = defaultdict(int)
        for hit in c.hits:
            majority[hit.label] += 1
        base = max(majority, key=majority.get) if majority else c.label.rsplit('_', 1)[0]
        suffix = seen[base]
        seen[base] += 1
        c.label = f"{base}_{suffix}"
    return clusters


def cluster_hits_online(hits, audio=None, sr=44100, ncc_threshold=0.72, attack_ms=25.0,
                        mel_threshold=0.62, target_min=0, target_max=0,
                        max_clusters=0):
    """Prototype-based online clustering for near-realtime previews.

    The batch algorithm builds an all-pairs matrix and then runs agglomerative
    clustering. This mode instead processes hits in onset order and compares
    each new hit only against current cluster prototypes. Complexity is roughly
    O(number_of_hits × number_of_clusters), so it can update progressively while
    audio is being analyzed. It is intentionally a preview/final-fast algorithm,
    not a replacement for the highest-quality batch pass.
    """
    if not hits:
        return []
    if len(hits) == 1:
        return [Cluster(cluster_id=0, label=f"{hits[0].label}_0", hits=[hits[0]])]
    if audio is None:
        audio = np.concatenate([h.audio for h in hits])

    cap = int(max_clusters or target_max or 0)
    if cap <= 0:
        cap = max(1, min(len(hits), int(target_min or 16)))
    cap = max(1, min(cap, len(hits)))

    print(f"[Cluster:online] {len(hits)} hits, cap={cap}, attack={attack_ms}ms")
    ordered = sorted(hits, key=lambda h: h.onset_time)
    clusters = []
    proto_fp = []
    proto_tr = []
    proto_energy = []

    for hit in ordered:
        fp = _mel_fingerprint(audio, sr, hit.onset_time)
        tr = _extract_transient(audio, sr, hit.onset_time, attack_ms)
        best_idx = -1
        best_score = -1.0
        best_mel = 0.0
        best_ncc = 0.0
        for i, cluster in enumerate(clusters):
            # Prefer same broad class when possible, but do not make it mandatory.
            label_bonus = 0.05 if cluster.label.startswith(hit.label + "_") else 0.0
            mel = _cosine(fp, proto_fp[i])
            if mel < mel_threshold:
                continue
            ncc = _transient_ncc(tr, proto_tr[i])
            score = (0.45 * mel) + (0.55 * ncc) + label_bonus
            if score > best_score:
                best_idx, best_score, best_mel, best_ncc = i, score, mel, ncc

        should_create = best_idx < 0 or (best_ncc < ncc_threshold and best_score < ncc_threshold)
        if should_create and len(clusters) < cap:
            cluster = Cluster(cluster_id=len(clusters), label=f"{hit.label}_{len(clusters)}", hits=[hit])
            clusters.append(cluster)
            proto_fp.append(fp)
            proto_tr.append(tr)
            proto_energy.append(max(hit.rms_energy, 1e-8))
            continue

        if best_idx < 0:
            # Cap reached and no good match: assign to the nearest prototype by mel.
            similarities = [_cosine(fp, existing) for existing in proto_fp]
            best_idx = int(np.argmax(similarities))
        cluster = clusters[best_idx]
        cluster.hits.append(hit)

        # Energy-weighted rolling prototype update; keeps loud clean hits dominant.
        w_old = proto_energy[best_idx]
        w_new = max(hit.rms_energy, 1e-8)
        total = w_old + w_new
        max_len = max(len(proto_fp[best_idx]), len(fp))
        old_fp = np.pad(proto_fp[best_idx], (0, max_len - len(proto_fp[best_idx])))
        new_fp = np.pad(fp, (0, max_len - len(fp)))
        proto_fp[best_idx] = ((old_fp * w_old) + (new_fp * w_new)) / total
        max_tr = max(len(proto_tr[best_idx]), len(tr))
        old_tr = np.pad(proto_tr[best_idx], (0, max_tr - len(proto_tr[best_idx])))
        new_tr = np.pad(tr, (0, max_tr - len(tr)))
        proto_tr[best_idx] = ((old_tr * w_old) + (new_tr * w_new)) / total
        proto_energy[best_idx] = total

    clusters = _retitle_clusters(clusters)
    for c in clusters:
        print(f"    {c.label}: {c.count}")
    return clusters

def cluster_hits(hits, audio=None, sr=44100, ncc_threshold=0.80, attack_ms=25.0,
                 mel_threshold=0.75, target_min=0, target_max=0,
                 linkage='average', merge_singletons=True):
    """Two-stage clustering: mel pre-filter → transient NCC → agglomerative."""
    from sklearn.cluster import AgglomerativeClustering as AC
    if not hits: return []
    N=len(hits)
    if N==1: return [Cluster(cluster_id=0,label=f"{hits[0].label}_0",hits=[hits[0]])]

    # Use the source audio for transient extraction if available
    if audio is None:
        # Fallback: concatenate hit audio (less accurate)
        audio = np.concatenate([h.audio for h in hits])

    print(f"[Cluster] {N} hits, attack={attack_ms}ms, mel_thresh={mel_threshold}")
    sim = build_similarity_matrix(hits, sr, audio, attack_ms, mel_threshold)
    D = np.clip(1.0 - sim, 0, 2).astype(np.float32)  # distance matrix

    use_t = target_min>0 and target_max>0 and target_max>=target_min
    tmin,tmax = max(1,min(target_min or 1,N)), max(max(1,target_min or 1),min(target_max or N,N))

    if use_t:
        print(f"  Target: {tmin}–{tmax}")
        lo,hi=0.001,1.5; bl,bn=None,-1
        for _ in range(30):
            mid=(lo+hi)/2
            lb=AC(n_clusters=None,distance_threshold=max(0.001,mid),metric='precomputed',linkage=linkage).fit_predict(D)
            n=len(set(lb))
            if tmin<=n<=tmax: bl,bn=lb,n; break
            elif n>tmax: lo=mid
            else: hi=mid
            if bl is None or abs(n-(tmin+tmax)/2)<abs(bn-(tmin+tmax)/2): bl,bn=lb,n
        if bn<tmin or bn>tmax:
            tm=min((tmin+tmax)//2,N-1)
            try: bl=AC(n_clusters=tm,metric='precomputed',linkage=linkage).fit_predict(D); bn=tm
            except: pass
        labels=bl; print(f"  → {bn} clusters")
    else:
        dt=max(0.001,1.0-ncc_threshold)
        labels=AC(n_clusters=None,distance_threshold=dt,metric='precomputed',linkage=linkage).fit_predict(D)

    print(f"  ✓ {len(set(labels))} clusters")
    clusters = _labels_to_clusters(labels, hits)

    if merge_singletons:
        clusters = _merge_singletons(clusters, sim, hits)

    for c in clusters: print(f"    {c.label}: {c.count}")
    return clusters

# ─── Stage 5: Quality scoring ────────────────────────────────────────────────
def sample_quality_score(y,sr,label="other"):
    """Fast deterministic quality score for ranking candidate samples.

    This intentionally avoids heavy onset re-detection inside every candidate so results can
    appear quickly in the web UI while extraction is still finishing.
    """
    y=np.asarray(y,dtype=np.float32)
    if y.size==0:
        return {'total':0.0,'completeness':0.0,'cleanness':0.0,'onset_quality':0.0}
    frame=max(64,int(sr*0.006)); hop=max(32,frame//2)
    if len(y)<frame:
        frames=y.reshape(1,-1)
    else:
        count=1+(len(y)-frame)//hop
        frames=np.stack([y[i*hop:i*hop+frame] for i in range(count)])
    rms=np.sqrt(np.mean(frames**2,axis=1)+1e-12)
    peak_idx=int(np.argmax(rms)); peak=float(rms[peak_idx])
    tail=rms[peak_idx:]
    tail_floor=float(np.mean(tail[-max(1,len(tail)//5):])) if tail.size else 0.0
    completeness=float(np.clip(1.0-tail_floor/(peak+1e-9)*3.5,0,1))
    pre=y[:max(1,min(len(y),int(sr*0.01)))]
    body=y[peak_idx*hop:min(len(y),peak_idx*hop+max(1,int(sr*0.08)))]
    pre_power=float(np.mean(pre**2)+1e-12)
    body_power=float(np.mean(body**2)+1e-12)
    cleanness=float(np.clip((-10*np.log10(pre_power/body_power)-4)/32,0,1))
    attack=float(np.max(np.abs(np.diff(y[:max(3,min(len(y),int(sr*0.035)))])))) if len(y)>2 else 0.0
    onset_quality=float(np.clip((attack/(np.mean(np.abs(y))+1e-8)-0.5)/8,0,1))
    total=float((completeness*0.34+cleanness*0.42+onset_quality*0.18+0.06)*100)
    return {'total':total,'completeness':completeness,'cleanness':cleanness,'onset_quality':onset_quality}

def select_best(clusters):
    print(f"[Select] best per cluster...")
    for c in clusters:
        if c.count<=1: c.best_hit_idx=0; continue
        c.best_hit_idx=int(np.argmax([sample_quality_score(h.audio,h.sr,c.label.rsplit('_',1)[0])['total'] for h in c.hits]))

# ─── Stage 6: Synthesis ──────────────────────────────────────────────────────
def synthesize_from_cluster(cluster):
    if cluster.count<2: return None
    tl=int(np.median([len(h.audio) for h in cluster.hits])); al,wt=[],[]; pp=None
    for i,h in enumerate(cluster.hits):
        a=h.audio.copy(); p=np.argmax(np.abs(a))
        if pp is None: pp=p
        s=pp-p
        if s>0: a=np.pad(a,(s,0))
        elif s<0: a=a[-s:]
        a=a[:tl] if len(a)>=tl else np.pad(a,(0,tl-len(a)))
        pk=np.abs(a).max()
        if pk>0: a/=pk
        al.append(a); wt.append(2.0 if i==cluster.best_hit_idx else 1.0)
    al=np.array(al); w=np.array(wt); w/=w.sum()
    sy=np.average(al,axis=0,weights=w); pk=np.abs(sy).max()
    return (sy*0.95/pk).astype(np.float32) if pk>0 else sy.astype(np.float32)

# ─── Stage 7: MIDI with quantization ─────────────────────────────────────────
def quantize_time(t, bpm, subdivision=16):
    """Snap time to nearest grid position. subdivision: 4=quarter, 8=eighth, 16=sixteenth, 32=thirty-second."""
    grid = 60.0 / bpm / (subdivision / 4)
    return round(t / grid) * grid

def build_midi(clusters, bpm=120.0, quantize=True, subdivision=16):
    import pretty_midi
    pm=pretty_midi.PrettyMIDI(initial_tempo=bpm)
    for i,c in enumerate(clusters): c.midi_note=min(36+i,127)
    inst=pretty_midi.Instrument(program=0,is_drum=True,name='Samples'); pm.instruments.append(inst)
    for c in clusters:
        for h in c.hits:
            t = quantize_time(h.onset_time, bpm, subdivision) if quantize else h.onset_time
            v=max(1,min(127,int(h.rms_energy/0.3*127)))
            inst.notes.append(pretty_midi.Note(velocity=v,pitch=c.midi_note,
                start=max(0,t),end=max(0,t)+max(h.duration,0.05)))
    inst.notes.sort(key=lambda n:n.start); return pm
def _midi_varlen(value):
    value=max(0,int(value)); buffer=value & 0x7f
    value >>= 7
    while value:
        buffer <<= 8
        buffer |= ((value & 0x7f) | 0x80)
        value >>= 7
    out=[]
    while True:
        out.append(buffer & 0xff)
        if buffer & 0x80:
            buffer >>= 8
        else:
            break
    return bytes(out)

def _write_fallback_midi(clusters,path,bpm=120.0,quantize=True,subdivision=16):
    """Write a small standards-compliant drum MIDI file when pretty_midi is unavailable."""
    import struct
    ticks_per_beat=480
    bpm=float(bpm or 120.0)
    us_per_quarter=max(1,int(60_000_000/bpm))
    events=[]
    for i,c in enumerate(clusters):
        c.midi_note=min(36+i,127)
    for c in clusters:
        for h in c.hits:
            start=quantize_time(h.onset_time,bpm,subdivision) if quantize else h.onset_time
            end=max(start+max(h.duration,0.05),start+0.03)
            velocity=max(1,min(127,int(h.rms_energy/0.3*127)))
            st=int(round(start*bpm/60*ticks_per_beat)); et=int(round(end*bpm/60*ticks_per_beat))
            events.append((st,0,bytes([0x99,int(c.midi_note),velocity])))
            events.append((et,1,bytes([0x89,int(c.midi_note),0])))
    events.sort(key=lambda item:(item[0],item[1]))
    track=bytearray()
    track.extend(b'\x00\xff\x51\x03'+us_per_quarter.to_bytes(3,'big'))
    last=0
    for tick,_,payload in events:
        track.extend(_midi_varlen(tick-last)); track.extend(payload); last=tick
    track.extend(b'\x00\xff\x2f\x00')
    data=bytearray()
    data.extend(b'MThd'+struct.pack('>IHHH',6,0,1,ticks_per_beat))
    data.extend(b'MTrk'+struct.pack('>I',len(track))+track)
    with open(path,'wb') as handle: handle.write(data)
    print(f"  ✓ MIDI fallback: {len(events)//2} notes" + (" (quantized)" if quantize else ""))
    return None

def export_midi(clusters,path,bpm=120.0,quantize=True,subdivision=16):
    try:
        pm=build_midi(clusters,bpm,quantize,subdivision); pm.write(path)
        print(f"  ✓ MIDI: {len(pm.instruments[0].notes)} notes" + (" (quantized)" if quantize else "")); return pm
    except ModuleNotFoundError as exc:
        if exc.name != 'pretty_midi':
            raise
        return _write_fallback_midi(clusters,path,bpm,quantize,subdivision)
def detect_bpm(y,sr):
    ck=("bpm",_ahash(y),sr); c=cache_get(ck)
    if c: return c
    try:
        oe=librosa.onset.onset_strength(y=y,sr=sr,aggregate=np.median)
        if len(oe) == 0 or float(np.max(oe)) <= 1e-9:
            return cache_set(ck,120.0)
        bpm=float(librosa.feature.tempo(onset_envelope=oe,sr=sr).item())
        try:
            _,beats=librosa.beat.beat_track(onset_envelope=oe,sr=sr,units='time')
        except Exception:
            beats=[]
        if len(beats)>2:
            ibi=60.0/float(np.median(np.diff(beats)))
            for c in [bpm,ibi]:
                if 70<=c<=200: bpm=c; break
            else:
                if bpm<70: bpm*=2
                elif bpm>200: bpm/=2
        if not np.isfinite(bpm) or bpm <= 0:
            bpm=120.0
        return cache_set(ck,round(float(bpm),1))
    except Exception:
        return cache_set(ck,120.0)
def render_midi_with_samples(clusters,sr=44100):
    me=max((h.onset_time+h.duration for c in clusters for h in c.hits),default=1.0)
    buf=np.zeros(int((me+1.0)*sr),dtype=np.float64)
    for c in clusters:
        s=c.best_hit.audio.astype(np.float64); re=c.best_hit.rms_energy if c.best_hit.rms_energy>0 else 0.1
        for h in c.hits:
            vs=min(2.0,h.rms_energy/(re+1e-8))**0.5; i=int(h.onset_time*sr); e=i+len(s)
            if e>len(buf): buf=np.concatenate([buf,np.zeros(e-len(buf))])
            buf[i:e]+=s*vs
    pk=np.abs(buf).max(); return (buf/pk*0.9).astype(np.float32) if pk>1e-8 else buf.astype(np.float32)
def build_sample_map(clusters):
    return {c.midi_note:{'label':c.label,'count':c.count,'duration_ms':int(c.best_hit.duration*1000)} for c in clusters}
def build_archive(clusters,bpm,sr,midi_path=None,rendered_audio=None,target_rendered_audio=None):
    import zipfile,tempfile,io; zp=tempfile.mktemp(suffix='.zip')
    idx={'bpm':round(bpm,1),'sample_rate':sr,'total_clusters':len(clusters),'total_hits':sum(c.count for c in clusters),'samples':{}}
    if rendered_audio is not None:
        idx['reproduction_file']='rendered_reproduction_full_mix.wav'
    if target_rendered_audio is not None:
        idx['target_reconstruction_file']='rendered_reconstruction_target_stem.wav'
    with zipfile.ZipFile(zp,'w',compression=zipfile.ZIP_STORED) as zf:
        for c in clusters:
            b=c.best_hit; fn=f"samples/{c.label}.wav"; buf=io.BytesIO()
            sf.write(buf,b.audio,sr,format='WAV',subtype='PCM_24'); zf.writestr(fn,buf.getvalue())
            ot=sorted([h.onset_time for h in c.hits])
            idx['samples'][c.label]={'file':fn,'classification':c.label.rsplit('_',1)[0],'midi_note':c.midi_note,
                'occurrences':c.count,'onset_times_sec':[round(t,4) for t in ot],'duration_sec':round(b.duration,4),
                'rms_energy':round(b.rms_energy,6),'spectral_centroid_hz':round(b.spectral_centroid,1)}
            if c.synthesized is not None:
                sf2=f"samples/{c.label}__synth.wav"; b2=io.BytesIO()
                sf.write(b2,c.synthesized,sr,format='WAV',subtype='PCM_24'); zf.writestr(sf2,b2.getvalue())
                idx['samples'][c.label]['synthesized_file']=sf2
        zf.writestr('index.json',json.dumps(idx,indent=2))
        if midi_path and os.path.exists(midi_path): zf.write(midi_path,'reconstruction.mid')
        if rendered_audio is not None:
            rb=io.BytesIO(); sf.write(rb,rendered_audio,sr,format='WAV',subtype='PCM_16')
            zf.writestr('rendered_reproduction_full_mix.wav',rb.getvalue())
            zf.writestr('rendered_reconstruction.wav',rb.getvalue())
        if target_rendered_audio is not None:
            tb=io.BytesIO(); sf.write(tb,target_rendered_audio,sr,format='WAV',subtype='PCM_16')
            zf.writestr('rendered_reconstruction_target_stem.wav',tb.getvalue())
    return zp

# ─── Auto-tuner with locking ─────────────────────────────────────────────────
def _reconstruction_score(orig,rend,sr):
    n=min(len(orig),len(rend))
    if n<4096: return 0.0
    eo=np.abs(librosa.stft(orig[:n],n_fft=4096,hop_length=2048)).mean(axis=1)
    er=np.abs(librosa.stft(rend[:n],n_fft=4096,hop_length=2048)).mean(axis=1)
    sc=float(np.corrcoef(eo,er)[0,1]) if eo.std()>1e-10 and er.std()>1e-10 else 0.0
    ro=librosa.feature.rms(y=orig[:n],hop_length=1024)[0]
    rr=librosa.feature.rms(y=rend[:n],hop_length=1024)[0]
    n2=min(len(ro),len(rr)); ro,rr=ro[:n2],rr[:n2]
    rc=float(np.corrcoef(ro,rr)[0,1]) if n2>4 and ro.std()>1e-10 and rr.std()>1e-10 else 0.0
    lr=min(len(rend),len(orig))/(max(len(rend),len(orig))+1)
    return max(0.0,(sc*0.5+rc*0.4+lr*0.1)*100)

def auto_tune(stem_audio, sr, mode="auto", locks=None, log_fn=None):
    locks=locks or {}; log=[]
    def _log(m): log.append(m); (log_fn(m) if log_fn else None); print(m)
    _log(f"[Auto-tune] {len(stem_audio)/sr:.1f}s, locked={list(locks.keys()) or 'none'}")
    base_ocs=[{"onset_delta":0.08,"energy_threshold_db":-40,"min_gap":0.02},
        {"onset_delta":0.10,"energy_threshold_db":-35,"min_gap":0.025},
        {"onset_delta":0.12,"energy_threshold_db":-35,"min_gap":0.03},
        {"onset_delta":0.15,"energy_threshold_db":-30,"min_gap":0.04},
        {"onset_delta":0.20,"energy_threshold_db":-28,"min_gap":0.05},
        {"onset_delta":0.25,"energy_threshold_db":-25,"min_gap":0.06}]
    seen=set(); ocs=[]
    for oc in base_ocs:
        for k in ['onset_delta','energy_threshold_db','min_gap']:
            if k in locks: oc[k]=locks[k]
        key=tuple(sorted(oc.items()))
        if key not in seen: seen.add(key); ocs.append(oc)
    if 'target_min' in locks and 'target_max' in locks:
        tmn,tmx=int(locks['target_min']),int(locks['target_max'])
        cts=sorted(set([tmn,tmx,(tmn+tmx)//2,tmn+(tmx-tmn)//3,tmn+2*(tmx-tmn)//3]))
        cts=[t for t in cts if tmn<=t<=tmx]
    elif 'target_min' in locks: cts=[t for t in [3,5,8,10,15,20,30] if t>=int(locks['target_min'])]
    elif 'target_max' in locks: cts=[t for t in [3,5,8,10,15,20,30] if t<=int(locks['target_max'])]
    else: cts=[3,5,8,10,15,20,30]
    best_score,best_params=-1,{}
    for oi,oc in enumerate(ocs):
        _log(f"  [{oi+1}/{len(ocs)}] delta={oc['onset_delta']}, E={oc['energy_threshold_db']}dB")
        hits=detect_onsets(stem_audio,sr,mode=mode,**oc)
        if len(hits)<2: continue
        hits=classify_hits(hits)
        for nt in cts:
            if nt>=len(hits): continue
            cl=cluster_hits(hits,audio=stem_audio,sr=sr,target_min=nt,target_max=nt)
            if not cl: continue
            for c in cl:
                if c.count<=1: c.best_hit_idx=0
                else: c.best_hit_idx=int(np.argmax([h.rms_energy for h in c.hits]))
            rend=render_midi_with_samples(cl,sr=sr)
            score=_reconstruction_score(stem_audio,rend,sr)
            if score>best_score:
                best_score=score; best_params={**oc,'target_min':nt,'target_max':nt,'n_clusters':len(cl)}
                _log(f"    ★ n={nt}→{len(cl)}cl, score={score:.1f} BEST")
            else: _log(f"    n={nt}→{len(cl)}cl, score={score:.1f}")
    if best_params:
        bt=best_params.get('target_min',10); foc={k:best_params[k] for k in ['onset_delta','energy_threshold_db','min_gap']}
        _log(f"  Fine-tune ±3 around {bt}...")
        hits=detect_onsets(stem_audio,sr,mode=mode,**foc)
        if len(hits)>=2:
            hits=classify_hits(hits)
            lo=max(2,int(locks.get('target_min',bt-3))); hi=min(len(hits)-1,int(locks.get('target_max',bt+3)))
            for ft in range(lo,hi+1):
                cl=cluster_hits(hits,audio=stem_audio,sr=sr,target_min=ft,target_max=ft)
                if not cl: continue
                for c in cl:
                    if c.count<=1: c.best_hit_idx=0
                    else: c.best_hit_idx=int(np.argmax([h.rms_energy for h in c.hits]))
                rend=render_midi_with_samples(cl,sr=sr); score=_reconstruction_score(stem_audio,rend,sr)
                if score>best_score:
                    best_score=score; best_params={**foc,'target_min':ft,'target_max':ft,'n_clusters':len(cl)}
                    _log(f"    ★ n={ft}→{len(cl)}cl, score={score:.1f} BEST")
    _log(f"\n[Auto-tune] Best: score={best_score:.1f}, {best_params}")
    return best_params, best_score, log