#!/usr/bin/env python3
"""Entry 019: prospectively specified focal-recording source extension."""
import argparse,hashlib,itertools,json,math,time
import os
from pathlib import Path
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
ROUND=Path(__file__).resolve().parents[1]
CACHE=Path(os.environ.get('BIO_DISCOVERY_CACHE','/mnt/data/research/bioinformatics-discovery/cache'));WORK=CACHE/'round-002'
PREFIX='019_focal'

def rank(value):return hashlib.sha256(f'019:20260919:{value}'.encode()).hexdigest()

def main():
    ap=argparse.ArgumentParser();ap.add_argument('--stage',choices=['eligibility','mixtures'],required=True);args=ap.parse_args();started=time.monotonic()
    model_path=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={x.split('_')[0]:i for i,x in enumerate(labels)}
    model=Interpreter(model_path=str(model_path),num_threads=1);model.allocate_tensors();inp=model.get_input_details()[0];out=model.get_output_details()[0]
    def predict(wave):
        assert len(wave)==144000 and np.max(np.abs(wave))<1
        model.set_tensor(inp['index'],wave[None,:]);model.invoke();return 1/(1+np.exp(-np.clip(model.get_tensor(out['index'])[0],-15,15)))
    def window(event,position):
        wave=np.zeros(144000,dtype='float32');start=[7200,79200][position];wave[start:start+57600]=clips[str(event)];return wave
    if args.stage=='eligibility':
        data=pd.read_csv(WORK/'019_train_xenocanto.csv');data['event_id']=np.arange(len(data));data=data[data.end_time>data.start_time].copy()
        options=[]
        for file,g in data.groupby('filepath'):
            info=sf.info(WORK/'019_focal_audio'/file);duration_file=info.frames/info.samplerate
            for e in g.itertuples():
                if e.label_status!='valid' or e.scientific_name not in mapping:continue
                start=(e.start_time+e.end_time)/2-.6;end=start+1.2
                if start<0 or end>duration_file:continue
                overlap=g[(g.start_time<end)&(g.end_time>start)]
                if set(overlap.scientific_name)!={e.scientific_name}:continue
                options.append(dict(event_id=int(e.event_id),filepath=file,species=e.scientific_name,start=start,end=end,selection_hash=rank(e.event_id),xc_id=e.xc_id,xc_recordist=e.xc_recordist,xc_url=e.xc_url,license=e.license))
        candidates=pd.DataFrame(options).sort_values('selection_hash').drop_duplicates(['species','filepath']).groupby('species',group_keys=False).head(12).sort_values('selection_hash').head(360)
        candidates.to_csv(ROUND/f'results/{PREFIX}_candidate_sample.csv',index=False)
        clips={}
        for file,g in candidates.groupby('filepath'):
            with sf.SoundFile(WORK/'019_focal_audio'/file) as source:
                rate=source.samplerate;gcd=math.gcd(rate,48000)
                for e in g.itertuples():
                    # Read a bounded excerpt with resampling margins, even for long source files.
                    native_start=max(0,math.floor((e.start-.1)*rate));native_start-=native_start%(rate//gcd)
                    source.seek(native_start);raw=source.read(math.ceil((e.end+.1)*rate)-native_start,dtype='float32',always_2d=True).mean(axis=1)
                    audio=resample_poly(raw,48000//gcd,rate//gcd).astype('float32');offset=round(e.start*48000)-native_start*48000//rate
                    clip=audio[offset:offset+57600].copy();assert len(clip)==57600
                    clip-=clip.mean();fade=np.linspace(0,1,240,dtype='float32');clip[:240]*=fade;clip[-240:]*=fade[::-1]
                    rms=float(np.sqrt(np.mean(clip**2)))
                    if rms<=0:raise ValueError(f'Silent annotated candidate {e.event_id}')
                    clip*=.03/rms;clips[str(e.event_id)]=clip
        scale=min(1.,.6/max(float(np.abs(c).max()) for c in clips.values()));clips={k:(v*scale*.5).astype('float32') for k,v in clips.items()}
        rows=[];baseline={}
        for e in candidates.itertuples():
            row=e._asdict();row.pop('Index',None)
            for position in [0,1]:
                scores=predict(window(e.event_id,position));baseline[f'{e.event_id}_{position}']=scores;row[f'score_{position}']=float(scores[mapping[e.species]])
            row['rms']=float(np.sqrt(np.mean(clips[str(e.event_id)]**2)));row['peak']=float(np.abs(clips[str(e.event_id)]).max());rows.append(row)
        table=pd.DataFrame(rows);table.to_csv(ROUND/f'results/{PREFIX}_all_baselines.csv',index=False)
        selected=table[(table.score_0>=.5)&(table.score_1>=.5)].sort_values('selection_hash').groupby('species',group_keys=False).head(4)
        counts=selected.groupby('species').size();selected=selected[selected.species.isin(counts.index[counts>=3])];species=sorted(selected.species.unique(),key=rank)[:16];selected=selected[selected.species.isin(species)]
        selected.to_csv(ROUND/f'results/{PREFIX}_selected_events.csv',index=False)
        ids=set(selected.event_id.astype(str));np.savez_compressed(WORK/f'{PREFIX}_clips.npz',**{k:v for k,v in clips.items() if k in ids});np.savez_compressed(WORK/f'{PREFIX}_baseline_scores.npz',**{k:v for k,v in baseline.items() if k.rsplit('_',1)[0] in ids});np.save(WORK/f'{PREFIX}_silence.npy',predict(np.zeros(144000,dtype='float32')))
        summary=dict(candidate_events=len(candidates),selected_events=len(selected),selected_species=len(species),selected_files=selected.filepath.nunique(),global_scale=scale,eligible=bool(len(species)>=8 and len(selected)>=24),runtime_seconds=time.monotonic()-started)
        (ROUND/f'results/{PREFIX}_eligibility.json').write_text(json.dumps(summary,indent=2)+'\n');print(json.dumps(summary,indent=2),flush=True);return
    eligibility=json.loads((ROUND/f'results/{PREFIX}_eligibility.json').read_text());assert eligibility['eligible'],'Frozen source-diversity gate failed; no mixture stage permitted.'
    selected=pd.read_csv(ROUND/f'results/{PREFIX}_selected_events.csv');clips=np.load(WORK/f'{PREFIX}_clips.npz');baseline=np.load(WORK/f'{PREFIX}_baseline_scores.npz')
    rows=[];pairs=list(itertools.combinations(selected.to_dict('records'),2))
    for pair_index,(a,b) in enumerate(pairs):
        if a['filepath']==b['filepath']:continue
        for condition,pa,pb in [('nonoverlap_AB',0,1),('nonoverlap_BA',1,0),('simultaneous',0,0)]:
            wa=window(a['event_id'],pa);wb=window(b['event_id'],pb);scores=predict(wa+wb)
            if condition.startswith('nonoverlap'):assert not np.any((wa!=0)&(wb!=0))
            for target,other,pos,otherpos in [(a,b,pa,pb),(b,a,pb,pa)]:
                ix=mapping[target['species']];rows.append(dict(event_id=target['event_id'],species=target['species'],file=target['filepath'],other_event=other['event_id'],other_species=other['species'],other_file=other['filepath'],same_species=target['species']==other['species'],condition=condition,baseline_score=float(baseline[f"{target['event_id']}_{pos}"][ix]),other_alone_score=float(baseline[f"{other['event_id']}_{otherpos}"][ix]),mixed_score=float(scores[ix]),position=pos))
        if pair_index%100==0:print(f'Mixture pairs {pair_index}/{len(pairs)}',flush=True)
    result=pd.DataFrame(rows);result.to_csv(ROUND/f'results/{PREFIX}_mixture_results.csv',index=False);summaries=[]
    for condition,g in result[~result.same_species].groupby('condition'):
        summaries.append(dict(condition=condition,target_appearances=len(g),loss_fraction=float((g.mixed_score<.5).mean()),median_score_change=float((g.mixed_score-g.baseline_score).median())))
    non=result[~result.same_species&result.condition.str.startswith('nonoverlap')];affected=non[non.mixed_score<.5].groupby('species').apply(lambda x:x[['file','other_file']].drop_duplicates().shape[0])
    summary=dict(**eligibility,conditions=summaries,advance=bool((non.mixed_score<.5).mean()>=.2 and (affected>=2).sum()>=4),species_affected_in_two_pairs=int((affected>=2).sum()),mixture_runtime_seconds=time.monotonic()-started,model_sha256=hashlib.sha256(model_path.read_bytes()).hexdigest())
    (ROUND/f'results/{PREFIX}_summary.json').write_text(json.dumps(summary,indent=2)+'\n');print(json.dumps(summary,indent=2))

if __name__=='__main__':main()