from tqdm import tqdm
import numpy as np
import pandas as pd
import soundfile as sf
from scipy import signal


df = pd.read_pickle('annot_all.pkl')
df = df[df.annot=='cachcach']
fs = 50000
HP = signal.butter(3, 2000*2/fs, 'highpass', output='sos')
BP = signal.butter(3, [7000*2/fs, 13000*2/fs], 'bandpass', output='sos')
df['SNR'] = 0.0

for i, r in tqdm(df.iterrows(), total=len(df)):
    sig, fs = sf.read('/nfs/NASDELL/SABIOD/SITE/BOMBYX/'+r.wavpath, start=int(r.time*fs)-62500, stop=int(r.time*fs)+62500)
    sig = sig[:,0]
    sigg = signal.sosfiltfilt(BP, sig)
    si = signal.sosfiltfilt(HP, sig)
    clickpos = np.argmax(sigg)
    Es = np.sum(np.square(si[clickpos-50:clickpos+50]))/100
    En = np.sum(np.square(si))/len(si)
    snr = 10 * np.log10(Es/En)
    df.at[i, 'SNR'] = snr
df.to_pickle('annot_all_snr.pkl')