#!/usr/bin/env python3
"""Candidate 011: compare solo scores with batched identical waveforms."""
from ceb_behavior_pilot import ROOT,CACHE,Interpreter,resample_poly,sf,np,pd,math
import json

model=Interpreter(model_path=str(CACHE/'v1_5_1_BirdNET_GLOBAL_6K_V2.4_Model_FP32.tflite'),num_threads=1)
samples=pd.read_csv(ROOT/'results/003_pilot_sample.csv').head(4)
chunks=[]
for e in samples.itertuples():
    audio,rate=sf.read(CACHE/'ceb_test_audio'/e.filepath,dtype='float32',always_2d=True)
    gcd=math.gcd(rate,48000)
    audio=resample_poly(audio.mean(axis=1),48000//gcd,rate//gcd).astype(np.float32)
    start=round(e.start*48000);chunks.append(audio[start:start+144000])
input_detail=model.get_input_details()[0]
def predict(batch):
    model.resize_tensor_input(input_detail['index'],batch.shape,strict=False);model.allocate_tensors()
    model.set_tensor(input_detail['index'],batch);model.invoke()
    logits=model.get_tensor(model.get_output_details()[0]['index'])
    return 1/(1+np.exp(-np.clip(logits,-15,15)))

rows=[]
for i,chunk in enumerate(chunks):
    solo=predict(chunk[None,:])[0]
    for name,other in [('silence',np.zeros_like(chunk)),('other',chunks[(i+1)%len(chunks)]),('loud_other',10*chunks[(i+1)%len(chunks)])]:
        batched=predict(np.stack([chunk,other]))[0]
        rows.append(dict(sample=i,distractor=name,max_score_difference=float(np.max(np.abs(batched-solo))),threshold_01_flips=int(((batched>=.1)!=(solo>=.1)).sum())))
result=dict(input_shape_signature=input_detail['shape_signature'].tolist(),comparisons=rows)
(ROOT/'results/011_batch_invariance.json').write_text(json.dumps(result,indent=2)+'\n')
print(json.dumps(result,indent=2))