from pathlib import Path

import numpy as np

from iblutil.util import Bunch
from ibllib.io.raw_data_loaders import load_data, load_settings

"""
cf. extractors to get the trigger times from bpod
-   In version 5 and above,  stim_trigger < stim_onset ~= cue_trigger < cue_onset ~= close_loop
-   stim_trigger - stim_onset : should max at 100ms. NaN if not detected. Cap number of Nans
for a session.
-   gocue_trigger - gocue_onset: get delay and count Nans
-   there should be exactly 1 step audio for a good trial, 2 steps. We know it can double up.
Having less is even more problematic

-   External data: Camera: count frames: video, timestamps and bpod
-   External data: wheel: check closed loop
"""


def qc_behaviour_bpod_session(ses_path):
    ses_path = Path(ses_path)
    raw_trials = load_data(ses_path)
    settings = load_settings(ses_path)
    if not raw_trials:
        return
    n_trials = len(raw_trials)
    # getting the task protocol will be useful at some point
    settings['IBLRIG_VERSION_TAG']
    settings['PYBPOD_PROTOCOL']

    # important stuff is the relationship between GoCue Audio and the stim onset

    # init the QC dictionary: one row per trial
    qc_trials = Bunch({
        'xor': np.zeros(n_trials, bool),  # a trial is either an error, correct or a no-go
        'correct_rewarded': np.zeros(n_trials) * np.nan,  # a correct trial needs to be rewarded
        'n_bnc1_high': np.zeros(n_trials) * np.nan,  # number of bnc1 fronts
        'n_bnc2_high': np.zeros(n_trials) * np.nan,  # number of bnc2 fronts
        'n_bnc1_low': np.zeros(n_trials) * np.nan,  # number of bnc1 fronts
        'n_bnc2_low': np.zeros(n_trials) * np.nan,  # number of bnc2 fronts
        'stim2cue': np.zeros(n_trials) * np.nan,  # time elapsed between stim onset and audio cue

    })

    for m, raw_trial in enumerate(raw_trials):
        states = raw_trial['behavior_data']['States timestamps']
        #  a trial is either an error, correct or a no-go. No in-between
        qc_trials.xor[m] = np.sum(int(np.all(np.isnan(states['error'][0]))) +
                                  int(np.all(np.isnan(states['correct'][0]))) +
                                  int(np.all(np.isnan(states['no_go'][0])))) == 2

        qc_trials.correct_rewarded[m] = (int(np.all(np.isnan(states['correct'][0]))) ==
                                         int(np.all(np.isnan(states['reward'][0]))))

        timestamps = raw_trial['behavior_data']['Events timestamps']
        qc_trials.n_bnc1_high[m] = len(timestamps['BNC1High'])
        qc_trials.n_bnc2_high[m] = len(timestamps['BNC2High'])
        qc_trials.n_bnc1_low[m] = len(timestamps['BNC1Low'])
        qc_trials.n_bnc2_low[m] = len(timestamps['BNC2Low'])

        qc_trials.stim2cue[m] = timestamps['BNC2High'][0] - states['stim_on'][0][0]  # first BNC1


def get_session_flatiron():
    from oneibl.one import ONE
    one = ONE()
    ses = one.search(subjects='CSHL_003', date_range=['2019-04-17'])  # session ok
    ses = one.search(subjects='CSHL_003', date_range=['2019-04-18'])  # session has wrong reaction
    one.load(ses, dataset_types=['_iblrig_taskData.raw', '_iblrig_taskSettings.raw'],
             download_only=True)


if __name__ == "__main__":
    ses_path = '/datadisk/Data/Subjects/ZM_1374/2019-03-24/001'
    ses_path = '/datadisk/FlatIron/churchlandlab/Subjects/CSHL_003/2019-04-17/001/'
    ses_path = '/datadisk/FlatIron/churchlandlab/Subjects/CSHL_003/2019-04-18/001/'
    qc_behaviour_bpod_session(ses_path)
