#!/usr/bin/env python3
"""Frozen BirdNET pilot: paired song/contact-call windows, one pair/file/species.

No model fitting. All sampling uses SHA256 seed 20260919, before predictions.
Audio is mono float32 at 48 kHz, default sigmoid sensitivity 1, no location prior.
Species identity is mapped exactly to the model's Latin-name labels.
"""
import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import time

os.environ.setdefault('OPENBLAS_NUM_THREADS','1')
os.environ.setdefault('OMP_NUM_THREADS','1')
import numpy as np
import pandas as pd
import soundfile as sf
from scipy.signal import resample_poly
from ai_edge_litert.interpreter import Interpreter

ROOT=Path(__file__).resolve().parents[1]
CACHE=Path(os.environ.get('BIO_DISCOVERY_CACHE','/mnt/data/research/bioinformatics-discovery/cache'))
SEED=20260919

def rank(text): return hashlib.sha256(f'{SEED}:{text}'.encode()).hexdigest()

def union_duration(intervals):
    total=0.; right=-math.inf
    for lo,hi in sorted(intervals):
        total+=max(0.,hi-max(lo,right));right=max(right,hi)
    return total

def candidates():
    data=pd.read_csv(ROOT/'evidence/ceb_v1_1_test_soundscape.csv')
    data['event_id']=np.arange(len(data))
    data['voc']=data['ebird#voc_type'].str.split('#').str[-1]
    valid=data[(data.label_status=='valid') & (data.end_time>data.start_time)]
    rows=[]
    for (species,file),events in valid.groupby(['scientific_name','filepath']):
        available={}
        for voc in ['song','contact call']:
            options=[]
            for e in events[events.voc==voc].itertuples():
                center=(e.start_time+e.end_time)/2
                start=center-1.5;end=center+1.5
                # Avoid padding and all other focal-species behaviors in the input.
                if start<0 or end>180: continue
                overlap=events[(events.start_time<end)&(events.end_time>start)]
                if (overlap.voc!=voc).any(): continue
                intervals=[(max(start,x.start_time),min(end,x.end_time)) for x in overlap.itertuples()]
                all_overlap=valid[(valid.filepath==file)&(valid.start_time<end)&(valid.end_time>start)]
                options.append(dict(species=species,filepath=file,voc=voc,event_id=int(e.event_id),
                    start=start,end=end,event_duration=e.end_time-e.start_time,
                    focal_occupancy=union_duration(intervals),focal_events=len(overlap),
                    other_species=int(all_overlap.scientific_name.nunique()-1),
                    selection_hash=rank(str(e.event_id))))
            if options: available[voc]=min(options,key=lambda e:e['selection_hash'])
        if len(available)==2:
            rows.extend(available.values())
    out=pd.DataFrame(rows)
    out['file_hash']=out.filepath.map(rank)
    out['phase']=out.file_hash.map(lambda h:'pilot' if int(h[:8],16)%3==0 else 'validation')
    return out

def main():
    parser=argparse.ArgumentParser();parser.add_argument('--phase',choices=['plan','pilot','validation'],default='plan')
    args=parser.parse_args()
    plan=candidates()
    plan.to_csv(ROOT/'results/003_eligible_pairs.csv',index=False)
    counts=plan.groupby(['species','phase']).filepath.nunique().unstack(fill_value=0)
    counts.to_csv(ROOT/'results/003_sampling_counts.csv')
    print(counts.to_string(),flush=True)
    if args.phase=='plan': return
    eligible=counts.index[(counts.sum(axis=1)>=3)&(counts.get('pilot',0)>=1)&(counts.get('validation',0)>=1)]
    sample=plan[(plan.species.isin(eligible))&(plan.phase==args.phase)]
    if args.phase=='pilot':
        # At most three paired files per species; selection is independent of scores.
        sample=sample.sort_values('file_hash').groupby(['species','voc'],group_keys=False).head(3)
    sample=sample.sort_values(['filepath','species','voc']).copy()
    sample.to_csv(ROOT/f'results/003_{args.phase}_sample.csv',index=False)
    model=CACHE/'v1_5_1_BirdNET_GLOBAL_6K_V2.4_Model_FP32.tflite'
    labels=(CACHE/'v1_5_1_BirdNET_GLOBAL_6K_V2.4_Labels_en_uk.txt').read_text().splitlines()
    mapping={label.split('_')[0]:i for i,label in enumerate(labels)}
    missing=set(sample.species)-mapping.keys()
    if missing: raise ValueError(f'Unmapped species: {missing}')
    interpreter=Interpreter(model_path=str(model),num_threads=1);interpreter.allocate_tensors()
    inp=interpreter.get_input_details()[0];out=interpreter.get_output_details()[0]
    assert tuple(inp['shape'])==(1,144000) and int(out['shape'][1])==len(labels)
    rows=[];started=time.monotonic()
    for file,events in sample.groupby('filepath',sort=True):
        path=CACHE/'ceb_test_audio'/file
        audio,rate=sf.read(path,dtype='float32',always_2d=True)
        audio=audio.mean(axis=1)
        if rate!=48000:
            gcd=math.gcd(rate,48000)
            audio=resample_poly(audio,48000//gcd,rate//gcd).astype(np.float32)
        for e in events.to_dict('records'):
            start=round(e['start']*48000);chunk=audio[start:start+144000]
            if len(chunk)!=144000: raise ValueError(f'Invalid window in {file}')
            interpreter.set_tensor(inp['index'],chunk[None,:]);interpreter.invoke()
            logits=interpreter.get_tensor(out['index'])[0]
            scores=1/(1+np.exp(-np.clip(logits,-15,15)))
            ix=mapping[e['species']]
            e.update(score=float(scores[ix]),logit=float(logits[ix]),
                rank=int(1+(logits>logits[ix]).sum()),top_label=labels[int(logits.argmax())],
                peak=float(np.abs(chunk).max()),rms=float(np.sqrt(np.mean(chunk**2))),source_rate=rate)
            rows.append(e)
        print(f'{len(rows)}/{len(sample)} windows; {time.monotonic()-started:.1f}s',flush=True)
    result=pd.DataFrame(rows);result.to_csv(ROOT/f'results/003_{args.phase}_predictions.csv',index=False)
    summaries=[]
    for species,group in result.groupby('species'):
        row=dict(species=species,paired_files=group.filepath.nunique())
        for threshold in [.1,.5]:
            recall=group.assign(detected=group.score>=threshold).groupby('voc').detected.mean()
            row.update({f'song_recall_{threshold}':float(recall['song']),f'call_recall_{threshold}':float(recall['contact call']),f'gap_{threshold}':float(recall['song']-recall['contact call'])})
        summaries.append(row)
    summary=pd.DataFrame(summaries);summary.to_csv(ROOT/f'results/003_{args.phase}_summary.csv',index=False)
    print(summary.to_string(index=False),flush=True)
    metrics=dict(phase=args.phase,species=len(summary),windows=len(result),paired_species_files=len(result)//2,
        median_gap_01=float(summary['gap_0.1'].median()),median_gap_05=float(summary['gap_0.5'].median()),
        runtime_seconds=time.monotonic()-started,model_sha256=hashlib.sha256(model.read_bytes()).hexdigest(),
        note='Descriptive paired-file pilot; no independent-bird or population-level inference.')
    (ROOT/f'results/003_{args.phase}_metrics.json').write_text(json.dumps(metrics,indent=2));print(json.dumps(metrics,indent=2))

if __name__=='__main__':main()