from tqdm import tqdm
import soundfile as sf
import pandas as pd
import matplotlib.pyplot as plt
import os

os.system('rm predpngs/*')

folder = '/nfs/NASDELL/SABIOD/SITE/KM3Net/DATA_WAV/'

df = pd.read_pickle('preds/KM3Net_dw_64_32kHz.preds')
sub = df[df.pred>15].sample(500)
#df['pos'] = (( df.pred > 20)&( df.pred < 30))
#sub = df.groupby(df.fn.str[:-5]).sum()
#sub = sub[((sub.pos>50)&(sub.pos<100))]
#sub = df.groupby('fn').quantile(.9)
#sub = sub[((sub.pred > 20 )&(sub.pred < 30))]
#sub = df[((df.pred>20)&(df.pred<25))].sample(100)

fs = 195312
print('n files : ', len(sub))
for i, row in tqdm(sub.iterrows(), total=len(sub)):
#    sub = df[df.fn.str.startswith(i)]
   # r = sub.iloc[sub.pred.argmax()]
#    r = sub[((sub.pred>20)&(sub.pred<30))].iloc[0]
    r = row
    sig, fs = sf.read(folder+r.fn, start=int(fs*r.offset), stop=int(fs*(r.offset+20)))
    plt.figure(figsize=[20, 10])
    plt.specgram(sig, NFFT=2048, Fs=fs, noverlap=1536)
    plt.ylim(3000, 30000)
    plt.title(r.fn+' '+str(r.offset))
    plt.savefig('predpngs/'+str(r.name))
    plt.close()