"""Matched audio measurements on full mixes and Demucs htdemucs stems (same algorithm for all three).
Estimates, not grades. Outputs analysis/audio-level.json and traces/<id>.npz"""
import json,sys,math; from pathlib import Path
import numpy as np, librosa, soundfile as sf, pyloudnorm as pyln
S=Path(sys.argv[1]); OUT=S/'analysis'; TR=S/'traces'; TR.mkdir(exist_ok=True)
B=Path('/workspaces/UGS/AdFactorySongFormat/trees/b614aa3b3e/runs/targ-music-gap-b614aa3b3e-20260915/audio/stems/htdemucs')
TRACKS={'targ':(S/'audio/targ.wav',B/'targ'),'baseline':(S/'audio/baseline.wav',B/'gopure'),'new':(S/'audio/new.wav',S/'stems/htdemucs/new')}
SR=22050; HOP=512; FR=HOP/SR
text={r['id']:r for r in json.load(open(OUT/'text-level.json'))}; blocks=json.load(open(OUT/'blocks.json'))
def load(p): return librosa.load(p,sr=SR,mono=True)[0]
def db(x): return 20*np.log10(np.maximum(1e-9,x))
def frames_rms(y): return librosa.feature.rms(y=y,frame_length=2048,hop_length=HOP)[0]
res={}
for tid,(mixp,stemdir) in TRACKS.items():
    print('==',tid,flush=True)
    mix=load(mixp); voc=load(stemdir/'vocals.wav'); dr=load(stemdir/'drums.wav'); ba=load(stemdir/'bass.wav'); ot=load(stemdir/'other.wav')
    n=min(map(len,[mix,voc,dr,ba,ot])); mix,voc,dr,ba,ot=[x[:n] for x in (mix,voc,dr,ba,ot)]; music=dr+ba+ot; dur=n/SR
    # loudness on original 44.1k stereo
    data,srr=sf.read(mixp); meter=pyln.Meter(srr); integrated=meter.integrated_loudness(data)
    st=[]; win=int(3*srr); step=int(1*srr)
    for i in range(0,len(data)-win,step):
        st.append(meter.integrated_loudness(data[i:i+win]))
    st=np.array(st); st_valid=st[np.isfinite(st)&(st>-70)]
    lra=float(np.percentile(st_valid,95)-np.percentile(st_valid,10))
    # frame features
    vr=frames_rms(voc); mr=frames_rms(music); drr=frames_rms(dr); t=librosa.frames_to_time(np.arange(len(vr)),sr=SR,hop_length=HOP)
    vdb=db(vr); mdb=db(mr); ddb=db(drr)
    vad=vdb>(vdb.max()-32)
    # smooth vad (min 100ms)
    k=int(0.1/FR); vad_s=np.convolve(vad.astype(float),np.ones(k)/k,'same')>0.5
    # pyin in 30s chunks
    f0=np.full(len(t),np.nan); prob=np.zeros(len(t)); vflag=np.zeros(len(t),bool)
    for start in range(0,int(dur)+1,30):
        a=start*SR; b=min(n,(start+30)*SR); c=voc[a:b]
        if len(c)<4096: continue
        ff,vv,pp=librosa.pyin(c,sr=SR,fmin=65.4,fmax=1046.5,frame_length=2048,hop_length=HOP)
        i0=int(round(a/HOP)); m=min(len(ff),len(t)-i0)
        f0[i0:i0+m]=ff[:m]; prob[i0:i0+m]=pp[:m]; vflag[i0:i0+m]=vv[:m]
        print(tid,'pyin',start,flush=True)
    midi=librosa.hz_to_midi(f0); good=np.isfinite(midi)&vflag&(prob>=0.5)
    # note segmentation
    notes=[]; i=0
    while i<len(t):
        if not good[i]: i+=1; continue
        j=i; vals=[midi[i]]
        while j+1<len(t) and good[j+1] and abs(midi[j+1]-np.median(vals))<0.6: j+=1; vals.append(midi[j])
        d=(j-i+1)*FR
        if d>=0.08: notes.append({'t':float(t[i]),'d':float(d),'m':float(np.median(vals)),'sd':float(np.std(vals))})
        i=j+1
    voiced_time=float(good.sum()*FR); note_time=sum(x['d'] for x in notes)
    stable=[x for x in notes if x['sd']<0.35]
    ivs=[notes[k+1]['m']-notes[k]['m'] for k in range(len(notes)-1) if notes[k+1]['t']-(notes[k]['t']+notes[k]['d'])<0.5]
    aiv=np.abs(np.array(ivs)) if ivs else np.array([])
    hist={'repeat(<0.75st)':float((aiv<0.75).mean()) if len(aiv) else None,'step(0.75-2.5)':float(((aiv>=0.75)&(aiv<2.5)).mean()) if len(aiv) else None,'third(2.5-4.5)':float(((aiv>=2.5)&(aiv<4.5)).mean()) if len(aiv) else None,'leap(>=4.5)':float((aiv>=4.5).mean()) if len(aiv) else None}
    sustained=[x for x in notes if x['d']>=0.5]; sus_long=[x for x in notes if x['d']>=0.8]
    hc=midi[good]
    # per-phrase contour
    phr=[]
    for p in text[tid]['phrasesList']:
        a=int(p['t']/FR); b=int(p['e']/FR); seg=midi[a:b]; g=good[a:b]
        if g.sum()<5: phr.append({'t':p['t'],'text':p['text'],'n':0}); continue
        idx=np.where(g)[0]; vals=seg[g]; peak=idx[np.argmax(vals)]
        # end slope: last 40% of voiced frames linear fit
        q=idx[int(len(idx)*0.6):]; slope=float(np.polyfit(t[a+q],seg[q],1)[0]) if len(q)>=4 else None
        nn=[x for x in notes if p['t']<=x['t']<p['e']]
        phr.append({'t':p['t'],'e':p['e'],'text':p['text'],'n':int(g.sum()),'range':float(vals.max()-vals.min()),'mean':float(vals.mean()),'peakPos':float(peak/max(1,(b-a))),'peakMidi':float(vals.max()),'endSlopeStPerSec':slope,'notes':len(nn),'notesPerSec':round(len(nn)/max(0.1,p['e']-p['t']),2),'maxNoteDur':max([x['d'] for x in nn],default=0),'syl':p['syl']})
    valid=[x for x in phr if x['n']>0]
    # vocal timbre & articulation
    Sv=np.abs(librosa.stft(voc,n_fft=2048,hop_length=HOP)); freqs=librosa.fft_frequencies(sr=SR,n_fft=2048)
    cent=librosa.feature.spectral_centroid(S=Sv,sr=SR)[0]; hf=(Sv[freqs>=4000]**2).sum(0)/np.maximum(1e-12,(Sv**2).sum(0))
    active=vad_s&(vdb>vdb.max()-25)
    onset_env=librosa.onset.onset_strength(y=voc,sr=SR,hop_length=HOP); onsets=librosa.onset.onset_detect(onset_envelope=onset_env,sr=SR,hop_length=HOP,units='time',backtrack=False)
    onsets=np.array([o for o in onsets if active[min(len(active)-1,int(o/FR))]])
    # beat grid from drums
    tempo,beats=librosa.beat.beat_track(y=dr,sr=SR,hop_length=HOP,trim=False); bt=librosa.frames_to_time(beats,sr=SR,hop_length=HOP)
    tempo=float(np.atleast_1d(tempo)[0])
    grid=[]
    for k in range(len(bt)-1): grid+= [bt[k], (bt[k]+bt[k+1])/2]
    grid=np.array(grid)
    dev=np.array([np.min(np.abs(grid-o)) for o in onsets if len(grid)]) if len(onsets) else np.array([])
    beat_ivs=np.diff(bt); tempo_cv=float(beat_ivs.std()/beat_ivs.mean()) if len(beat_ivs)>2 else None
    # music dynamics: dips and instrumental-only spans
    m05=[]; T05=[]
    w=int(0.5/FR)
    for i in range(0,len(mdb)-w,w): m05.append(float(np.median(mdb[i:i+w]))); T05.append(float(t[i]))
    m05=np.array(m05); med=np.median(m05[m05>m05.max()-40]); dips=[(T05[i],float(m05[i]-med)) for i in range(len(m05)) if m05[i]<med-6]
    # instrumental-only spans (no vocal >=1.0s)
    spans=[]; i=0
    while i<len(vad_s):
        if not vad_s[i]:
            j=i
            while j<len(vad_s) and not vad_s[j]: j+=1
            if (j-i)*FR>=1.0: spans.append((round(float(t[i]),2),round(float((j-i)*FR),2)))
            i=j
        else: i+=1
    # block-wise balances
    blk=[]
    for b_ in blocks[tid]:
        a=int(b_['start']/FR); e=int(b_['end']/FR); e=min(e,len(t))
        if e<=a: continue
        vb=db(np.sqrt(np.mean(vr[a:e]**2))); mb=db(np.sqrt(np.mean(mr[a:e]**2))); dbb=db(np.sqrt(np.mean(drr[a:e]**2)))
        seg_on=[o for o in onsets if b_['start']<=o<b_['end']]
        gm=midi[a:e][good[a:e]]
        blk.append({'id':b_['id'],'vocalDb':round(float(vb),1),'musicDb':round(float(mb),1),'drumsDb':round(float(dbb),1),'vocMinusMusic':round(float(vb-mb),1),'vocalOnsetsPerSec':round(len(seg_on)/(b_['end']-b_['start']),2),'pitchMedian':round(float(np.median(gm)),1) if len(gm) else None,'pitchP95':round(float(np.percentile(gm,95)),1) if len(gm) else None,'centroidHz':round(float(np.mean(cent[a:e][active[a:e]])),0) if active[a:e].any() else None})
    res[tid]={'duration':round(dur,3),'integratedLufs':round(float(integrated),2),'shortTermLra_LU':round(lra,2),'shortTermP10P50P95':[round(float(np.percentile(st_valid,q)),1) for q in (10,50,95)],
      'tempoBpm':round(tempo,2),'beatIntervalCV':round(tempo_cv,4) if tempo_cv else None,
      'voicedTimeHighConf':round(voiced_time,1),'voicedShareHighConf':round(voiced_time/dur,3),'vocalActiveShare':round(float(vad_s.mean()),3),
      'pitch':{'p5':round(float(np.percentile(hc,5)),1),'p50':round(float(np.percentile(hc,50)),1),'p95':round(float(np.percentile(hc,95)),1),'rangeP5P95':round(float(np.percentile(hc,95)-np.percentile(hc,5)),1)},
      'notes':{'count':len(notes),'perVoicedSec':round(len(notes)/max(1,voiced_time),2),'meanDur':round(float(np.mean([x['d'] for x in notes])),3),'medianDur':round(float(np.median([x['d'] for x in notes])),3),'stableShareOfNotes':round(len(stable)/max(1,len(notes)),3),'noteTimeShareOfVoiced':round(note_time/max(1,voiced_time),3),
         'sustained>=0.5s':len(sustained),'sustained>=0.8s':len(sus_long),'sustainedTimeShare':round(sum(x['d'] for x in sustained)/max(1,voiced_time),3),'longest':[{'t':round(x['t'],2),'d':round(x['d'],2),'midi':round(x['m'],1)} for x in sorted(notes,key=lambda x:-x['d'])[:15]],'intervalHist':{k:(round(v,3) if v is not None else None) for k,v in hist.items()},'meanAbsInterval':round(float(aiv.mean()),2) if len(aiv) else None},
      'phrases':{'count':len(valid),'meanRangeSt':round(float(np.mean([x['range'] for x in valid])),2),'medianRangeSt':round(float(np.median([x['range'] for x in valid])),2),'peakPosMean':round(float(np.mean([x['peakPos'] for x in valid])),3),'peakInFirstHalfShare':round(float(np.mean([x['peakPos']<0.5 for x in valid])),3),'fallingEndShare':round(float(np.mean([ (x['endSlopeStPerSec'] or 0)<-1 for x in valid])),3),'risingEndShare':round(float(np.mean([ (x['endSlopeStPerSec'] or 0)>1 for x in valid])),3),'notesPerSecMean':round(float(np.mean([x['notesPerSec'] for x in valid])),2),'maxNoteDurMean':round(float(np.mean([x['maxNoteDur'] for x in valid])),3),'list':phr},
      'timbre':{'centroidHzActive':round(float(np.mean(cent[active])),0),'hfRatio4k':round(float(np.mean(hf[active])),4),'vocalOnsetsPerActiveSec':round(len(onsets)/max(1,float(active.sum()*FR)),2),'onsetStrengthMeanAtOnsets':round(float(np.mean([onset_env[min(len(onset_env)-1,int(o/FR))] for o in onsets])),2) if len(onsets) else None},
      'grid':{'onsetDevMs':{'median':round(float(np.median(dev)*1000),0),'p75':round(float(np.percentile(dev,75)*1000),0),'within40ms':round(float((dev<0.04).mean()),3),'within60ms':round(float((dev<0.06).mean()),3)} if len(dev) else None,'eighthGridMs':round(float(30/tempo*1000),0)},
      'music':{'dips>6dB_0.5s':len(dips),'dipList':[(round(a,1),round(b,1)) for a,b in dips[:20]],'instrumentalSpans>=1s':len(spans),'instrumentalSpanTotal':round(sum(d for _,d in spans),1),'instrumentalSpans':spans[:20]},
      'blocks':blk}
    np.savez_compressed(TR/f'{tid}.npz',t=t,midi=midi,prob=prob,good=good,vdb=vdb,mdb=mdb,ddb=ddb,cent=cent,active=active,onsets=onsets,beats=bt,notes_t=np.array([x['t'] for x in notes]),notes_d=np.array([x['d'] for x in notes]),notes_m=np.array([x['m'] for x in notes]))
    json.dump(res,open(OUT/'audio-level.json','w'),ensure_ascii=False,indent=1)
print('DONE')