import sys
import os
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
from thunderfish.dataloader import open_data
from thunderfish.eodanalysis import eod_waveform
from IPython import embed


def load_tracker_data(foldername):
    '''
    Load data
    :param foldername: where are the data saved
    :return:    ident_v:    identity vector
                power:      power vector
                freq_v:     frequency vector
                timeidx:    time vector in indices
                realtime:   time vector in relative real time to the beginning
    '''
    ident_v = np.load('../../../data/' + foldername + '/all_ident_v.npy', allow_pickle=True)
    power = np.load('../../../data/' + foldername + '/all_sign_v.npy', allow_pickle=True)
    freq_v = np.load('../../../data/' + foldername + '/all_fund_v.npy', allow_pickle=True)
    timeidx = np.load('../../../data/' + foldername + '/all_idx_v.npy', allow_pickle=True)
    realtime = np.load('../../../data/' + foldername + '/all_times.npy', allow_pickle=True)
    aifl = np.load('../data/' + foldername + '/aifl2.npy', allow_pickle=True)

    return ident_v, power, freq_v, timeidx, realtime, aifl


def extract_times_freq_for_plotting(aifl, ident_v, freq_v, power_v, timeidx, realtime, fish_number):

    t = []
    f = []
    time_res = np.diff(realtime)[0]
    dt = 10  # IN SECONDS
    for channel in tqdm(range(16)):
        for id in aifl[channel,fish_number][~np.isnan(aifl[channel, fish_number])]:
            i_df_power_channel = []
            itr = np.arange(len(ident_v[channel]))[ident_v[channel] == id]
            for i0 in itr[np.array(np.round(np.linspace(0, len(itr)*.9-1, 100)), dtype=int)]:
                freqs_of_interest = freq_v[channel][(np.arange(len(ident_v[channel])) >= i0) &
                                           (np.arange(len(ident_v[channel])) <= i0 + dt / time_res) &
                                           (ident_v[channel] == id)]

                df = np.max(freqs_of_interest) - np.min(freqs_of_interest)
                mean_f = np.mean(freqs_of_interest)

                # if df <= 0.5:
                sign_of_interest = power_v[channel][(np.arange(len(ident_v[channel])) >= i0) &
                                                  (np.arange(len(ident_v[channel])) <= i0 + dt / time_res) &
                                                  (ident_v[channel] == id)]
                max_channel = np.argmax(sign_of_interest, axis=1)
                power = list(map(lambda x: np.max(x), sign_of_interest))

                i_df_power_channel.append([i0, mean_f, np.min(power), max_channel[0], df])

            i0, mean_f, power, c, df = i_df_power_channel[np.nanargmax(np.array(i_df_power_channel)[:, 4])]
            t.append(realtime[timeidx[channel][i0]])
            f.append(mean_f)
    return t, f


def main(filename):
    '''
    Extracts the time points and corresponding frequencies where the EOD waveforms should be calculated

    :param filename:
    :return:    all_t:      list; all time points of each fish where the eod should be extracted
                all_f:      list; all frequencies of each fish where the eod should be extracted
                all_fish:   list; fish_number of the aifl
    '''
    # load data
    ident_v, power_v, freq_v, timeidx, realtime, aifl = load_tracker_data(filename)

    fish_in_aifl = list(np.unique(np.where(~np.isnan(aifl[:, :, 0]))[1]))
    all_t = []
    all_f = []
    all_fish = []
    for fish_number in fish_in_aifl:
        print(fish_number)
        # if 400 <= all_q10[fish_number, 2] <= 650 and power_means[fish_number] >= -90.0:
        times, freqs = extract_times_freq_for_plotting(aifl, ident_v, freq_v, power_v, timeidx, realtime, fish_number)

        all_t.append(times)
        all_f.append(freqs)
        all_fish.append(fish_number)

    return all_t, all_f, all_fish



if __name__ == '__main__':

    for index, filename_idx in enumerate([0, 2, 3, 5]):
        filename = sorted(os.listdir('../../../data/mount_data/sanmartin/softgrid_1x16/'))[filename_idx]
        print('new file: '+ filename)
        all_t, all_f, all_fish = main(filename)

        #############################################################################################################
        # save data
        if filename not in os.listdir('../data/'):
            os.mkdir('../data/' + filename)

        np.save('../data/' + filename + '/eod_times_new_new.npy', all_t)
        np.save('../data/' + filename + '/eod_freq_new_new.npy', all_f)
        np.save('../data/' + filename + '/eod_fishnumber_new_new.npy', all_fish)