import models
import os
import pandas as pd
import numpy as np
from run_model import run
import soundfile as sf
import argparse

stdc='models/stft_depthwise_ovs_128_k7_r1.stdc'
folder = '/nfs/NASDELL/SABIOD/SITE/BOMBYX/'
passages = pd.read_pickle(folder+'manip_2021/passages.pkl')
passages.dateS -= pd.Timedelta('3h')
passages.dateE += pd.Timedelta('3h')

model = models.get['stft_depthwise_ksize'](128, 7)

for i, p in passages.iterrows():
    outfn = folder+'manip_2021/passagesPreds/passage_'+str(i)+'_stft_depthwise_64_r0.pkl'
    if os.path.isfile(outfn):
        continue
    print('starting '+str(i))
    fns = []
    for day in pd.date_range(p.dateS.floor('d'), p.dateE, freq='d'):
        dayfolder = day.strftime('YEAR_MONTH_DAY/%Y/%m/%d/')
        if not os.path.isdir(folder+dayfolder):
            continue
        for f in os.listdir(day.strftime(folder+dayfolder)):
            if f > p.dateS.strftime('BOMBYX_%Y%m%d_%H%M%S.wav') and f < p.dateE.strftime('BOMBYX_%Y%m%d_%H%M%S.wav') and f.endswith('wav'):
                fns.append(dayfolder+f)
    try:
        preds = run(fns, stdc, model, folder, pool=False, batch_size=64)
        preds.to_pickle(outfn)
    except:
        print('FAILED preds with passage '+str(i))
        continue