import os
import pandas as pd
import numpy as np
from chirpdetection import chirpdetection
from IPython import embed

# check rec ../data/mount_data/2020-03-25-10_00/ starting at 3175


def get_valid_datasets(dataroot):

    datasets = sorted([name for name in os.listdir(dataroot) if os.path.isdir(
        os.path.join(dataroot, name))])

    valid_datasets = []
    for dataset in datasets:

        path = os.path.join(dataroot, dataset)
        csv_name = '-'.join(dataset.split('-')[:3]) + '.csv'

        if os.path.exists(os.path.join(path, csv_name)) is False:
            continue

        if os.path.exists(os.path.join(path, 'ident_v.npy')) is False:
            continue

        ident = np.load(os.path.join(path, 'ident_v.npy'))
        number_of_fish = len(np.unique(ident[~np.isnan(ident)]))
        if number_of_fish != 2:
            continue

        valid_datasets.append(dataset)

    datapaths = [os.path.join(dataroot, dataset) +
                 '/' for dataset in valid_datasets]

    return datapaths, valid_datasets


def main(datapaths):

    for path in datapaths:
        chirpdetection(path, plot='show')


if __name__ == '__main__':

    dataroot = '../data/mount_data/'


    datapaths, valid_datasets= get_valid_datasets(dataroot)

    recs = pd.DataFrame(columns=['recording'], data=valid_datasets)
    recs.to_csv('../recs.csv', index=False)
    # datapaths = ['../data/mount_data/2020-03-25-10_00/']
    main(datapaths)

# window 1524 + 244 in dataset index 4 is nice example