# AUTOGENERATED! DO NOT EDIT! File to edit: 14_synchro.nested_stims.ipynb (unless otherwise specified).

__all__ = ['get_stim_ids', 'get_synced_file', 'get_synced_file_precompile', 'get_stim_dict', 'cluster_by_list',
           'find_stim_onsets', 'append_tread_movement', 'twoP_dataChunks_2', 'append_to_reM', 'master_loop']

# Cell
from .io import *
from .processing import *
from .extracting import *

from ..core import *
from ..utils import *
from ..database import *
from ..processing import *
from ..plotting import *

import numpy as np
import glob, os, datetime
import ast

def get_stim_ids(log, record_time, lastmodif_time):
    """Obtain the correct name of the stimulus, based on the name that appears in the QDSpy log

    This is needed if no central database of compiled stimuli is used, but instead stimuli are compiled in QDSpy
    every time before they are presented.

    params:
        - log: log object generated by get_QDSpy_logs
        - record_time: the stimulus presentation time, as calculated by subtracting stimulus duration from lastmodif_time
        - lastmodif_time: the timestamp of the .npy file (time of compilation, assumed to be only seconds before record_time)
    returns:
        - stim_ids: Correct stim_id for the stimulus presented after the current log stimulus start
    """
    stim_ids = []
    photo_early_end = []
    for i, stim in enumerate(log.stimuli):
        if stim.stop_time is None:
            stim.stop_time = log.stimuli[i + 1].start_time
        if (stim.start_time > record_time) & (stim.stop_time < lastmodif_time):
            stim_ids.append(stim)
        elif (stim.start_time > record_time):
            photo_early_end.append(stim)
    if len(stim_ids)==0:
        print("Check the timestamp of the corresponding Photodiode file. The diode seems to have stopped before the stim ended.")
        print("No matching stimulus. Taking first stim after start of photodiode. Check timestamps carefully.")
        stim_ids = [photo_early_end[0]]
    return stim_ids

def get_synced_file(stim_list_dir, stim_id):
    '''
        Find the stimulus in the stimulus list directory that is temporally closest to the
        stimulus in the log. Works based on the modification time of the stimulus (i.e.
        expects stimulus to be compiled shortly before display).

        params:
            - stim_list_dir: fullpath to stimuli, string
            - stim_id: stimulus read from QDSpy log, theonerig.synchro.extracting.Stimulus
                        object
        returns:
            - stim: filename of the stimulus that needs loading, str
    '''
    stims = {"stim_name": [], "stim_delta": []}
    for stim_list in os.listdir(stim_list_dir):
        stim_load = datetime.datetime.fromtimestamp(int(os.path.getmtime(os.path.join(stim_list_dir, stim_list))))
        stim_present = stim_id.start_time
        # If the stimulus was compiled before display calculate difference, otherwise set to max
        stim_delta = stim_present - stim_load if stim_present > stim_load else datetime.timedelta.max
        stims["stim_name"].append(stim_list)
        stims["stim_delta"].append(stim_delta)
    # Obtain the index of the compiletime closest to the stimulus display time
    closest_stim_idx = stims["stim_delta"].index(min(stims["stim_delta"]))
    stim_fn = stims["stim_name"][closest_stim_idx]
    stim_path = os.path.join(stim_list_dir, stim_fn)
    # Sanity check
    if not stim_id.name in stim_path:
        print("Compiled stimulus file not matching this name")
        print("stim_id: {}".format(stim_id.name))
        print("stimulus file: {}".format(os.path.basename(stim_path)))
        stim_path = os.path.join(stim_list_dir, os.path.basename(stim_path))
        # Clean up: Ugly 1-time fix
        if "npz" in stim_path:
            stim = np.load(stim_path)
            stim = stim['stim' + stim_id.filename[3:6]].flatten()
        else:
            stim = np.load(stim_path)
    else:
        stim = np.load(stim_path)

    # Some of the stimuli have a shape of repetition_number x stim_onset:
    if len(stim.shape) > 1:
        stim = stim.flatten()
    return stim, stim_path

def get_synced_file_precompile(stim_list_dir, stim_id):
    for stim_list in os.listdir(stim_list_dir):
        if stim_id.md5 in stim_list:
            stim_path = os.path.join(stim_list_dir, stim_list)
            stim = np.load(stim_path)
    # Some of the stimuli have a shape of repetition_number x stim_onset:
    if len(stim.shape) > 1:
        stim = stim.flatten()
    return stim, stim_path

def get_stim_dict(stim_dict_path, stim_filename):
    """Get the correct dictionary file that decodes the ints that encode the stimulus sequence.
    Applicable to nested stimuli, whose _intensity_.npy file contains only a list of integers.

    params:
        - stim_dict_path: [string], the path to the directory where all dicts are stored
        - stim_filename: [string], the name of the stimulus, make sure that stim_id from the QDSpy log and the name of the dictionary match

    """
    for dict in os.listdir(stim_dict_path):
        if stim_filename[0:6] in dict:
            with open(os.path.join(stim_dict_path, dict),'r') as inf:
                dict_from_file = []
                for line in inf:
                    dict_from_file.append(ast.literal_eval(line))
            return dict_from_file

def cluster_by_list(data, frame_timepoints, frame_signals, stim_list):
    """Assign the stimulus identity values from stim_list to the frames in data. stim_list contains only the
        sequence of stimuli. Those need to be expanded. Opposed to cluster_frame_signals and cluster_by_epochs no
        AUC operation is performed.
        params:
            - data: raw data used to compute the stimulus times
            - frame_timepoints: timepoints delimiting each frame
            - frame_signals: binary 1-D numpy array indicating if high_threshold was passed in 'detect_frames'
            - stim_list: 1-D numpy array containing the sequence of the stimuli presented
        returns:
            - frame_signals: [same size as frame_timepoints] stim_signals list containing the correct value from
                stim_list at every entry"""
    epoch_end, stim_change = find_stim_onsets(frame_signals)
    # Split into on times & values vs off times & values
    stim_ons = stim_change[0::2]
    #stim_ons_idx = stim_idx[0::2]
    stim_offs = stim_change[1::2]
    #stim_offs_idx = stim_idx[1::2]

    # Replace the frame_signal entries with the stimulus codes
    frame_signals[frame_signals == 0] = -1 # To avoid confusion with the '0' stimulus code
    for i,stim_type in enumerate(stim_list):
        if i >= len(stim_ons):
            print("WARNING: Not all stimuli of the stim_list were actually presented. Recording ended early.")
            frame_signals[stim_ons[i - 1:]] = -1
            return frame_signals, stim_ons, stim_offs, epoch_end
        frame_signals[stim_ons[i]:stim_offs[i]] = stim_type

    return frame_signals, stim_ons, stim_offs, epoch_end

def find_stim_onsets(frame_signals):
    """Get the frame timepoints during which the stimuli started and ended.
    Crop out the last two ones that are currently just an offset marker.

    params:
        - frame_signals: processed signal from the Photodiode files, containing the marker information
    returns:
        - epoch_end: the frame timepoints of the offset signal
        - stim_change: the frame timepoints of every stimulus on and offset
    """
    stim_change = np.where(frame_signals[:-1] != frame_signals[1:])[0]
    stim_change = stim_change + 1 # since I need to compare to [1:] all values are shifted by 1
    #stim_idx = frame_timepoints[stim_change]

    # QDSpy currently is set to emit a short peak at the end to indicate the end of the stimulus presentation
    # This peak needs to be ignored
    epoch_end = stim_change[-2:] # return it for future analysis
    stim_change = stim_change[:-2] #add it to the no stimulus category
    return epoch_end, stim_change

def append_tread_movement(reM, epoch_counter, treadmill_data, ref_timepoints):
    '''
    Append treadmill data to the record Master
    '''
    tread_tp = np.arange(len(treadmill_data), dtype="int")
    tread_downsamp = resample_to_timepoints(tread_tp, treadmill_data, ref_timepoints)
    reM[epoch_counter - 1]["treadmill"] = tread_downsamp
    return reM

def twoP_dataChunks_2(ref_timepoints, frame_timepoints, len_epochs, cursor, skip=None, popped_epoch=None, *args):
    """
    Factory function for two photon data.

    params:
        - ref_timepoints: Reference timepoints to create the DataChunk
        - frame_timepoints: List of frame timepoints for each sequence of two photon frame recorded.
        - len_epochs: Lenght of the recorded epochs (<= than the corresponding frame_timepoints). Int or list
        - args: matrices of all frames detected by CaImAn. (give as many as you want to synchronise)

    return:
        - tuple containing the synchronised matrices in the order it was given
    """
    assert len(args)>=1, "no matrix to be synchronised was given"
    res_l = [[] for i in range(len(args))] # A list containing 1 entry per Ca-matrix to be processed
    if cursor is None:
        cursor = 0 # counter of from where to start indexing into the Ca-matrices
    if isinstance(len_epochs, int):
        len_epochs = [len_epochs]
    if isinstance(popped_epoch, int):
        popped_epoch = [popped_epoch]
    # For every recording block (defined by len_epochs, that counts the number of frames within
    # the recording), find the index (for the ref_timepoints) of the first and last frame of
    # that block.
    for i, len_epoch in enumerate(len_epochs):
        print(i)
        if i == 0:
            set_trace()
        first_ca = frame_timepoints[i][0]
        if len_epoch > len(frame_timepoints[i]):
            last_ca = frame_timepoints[i][-1]
            print(i, "Ca data longer than sync timepoints!")
        else:
            last_ca = frame_timepoints[i][len_epoch-1]
        if isinstance(ref_timepoints, Data_Pipe):
            start_idx = np.argmax(ref_timepoints[i]['main_tp']>first_ca)
            stop_idx  = np.argmax(ref_timepoints[i]['main_tp']>last_ca)
        else:
            start_idx = np.argmax(ref_timepoints>first_ca)
            stop_idx  = np.argmax(ref_timepoints>last_ca)
        for k, matrix in enumerate(args):
            # Slice the Ca-matrix in the time dimension to the duration of the recording block
            sub_mat = matrix.T[cursor:cursor+len_epoch]
            # Generate a linear interpolation function of the Ca-matrix values over the time of the
            # recording block, axis 0 is axis of the Ca-intensity values over time
            f = interpolate.interp1d(range(len_epoch), sub_mat, axis=0)
            # stop_idx-start_idx defines the number of timepoints for which values need to be
            # interpolated (i.e resolution). The block duration is again defined by len_epoch.
            # This interpolated data will be inserted as a datachunk object starting at the
            # correct start_idx (indexing ref_timpoints)
            res_l[k].append(DataChunk(data=f(np.linspace(0,len_epoch-1,stop_idx-start_idx)),
                                           idx=start_idx,
                                           group="cell"))
        if (skip is not None) and (i+1 in [skip]):
            cursor += len_epoch + popped_epoch[0] # skip the 3000 water frames
            del popped_epoch[0]
        else:
            cursor += len_epoch # Skip indexing to the start of the next recording epoch
    return tuple(res_l)

# Cell
from .io import *
from .processing import *
from .extracting import *

from ..core import *
from ..utils import *
from ..database import *
from ..processing import *
from ..plotting import *

import numpy as np
import glob, os, datetime
def append_to_reM(epoch_counter, reM, retain_template, stim_list_dir, stim_ids, record_time,
                                photo_data, stim_code_dict, frame_tp, ref_timepoints,
                               ref_signals, tom_stim_dir):
    for stim in stim_ids:
    # Can't: Lerion didn't save it there. obtaining stimulus parameters from the database
    #     stim_params = stim_param_to_dict(stim_params_df, stim.md5)
    #     first_fr, last_fr = stim_params["first_frame_id"], stim_params["last_frame_id"]

        # One of Tom's stimuli??
        if stim.name in retain_template:
#             tom_stim_dir = os.path.join("/media/asari/Tristan/01_Code/013_Leiron_Sweep/Day_1")
            # unpacked is a tuple of (frames presented, REF and shader)
            unpacked = unpack_stim_npy(tom_stim_dir, stim.md5)
            estimate_start = get_position_estimate(stim.start_time, record_time, 10000)
            if estimate_start<0:
                print(stim, "Started before the record, passing")
                continue
            # Needs to be adjusted manually, open photo_data plots and read out first
            # and last frames. Find based on ref_tp alignment (REC) with unpacked[1] (REF)
            starting_pos = input("This stimulus needs to be synchronized. Please enter the starting ref_tp: ")
            starting_pos = int(starting_pos)
            signals = ref_signals[starting_pos:starting_pos+len(unpacked[1])]
            stim_tuple, shift_log, frame_replacement  = frame_error_correction(signals,
                                                                                unpacked,
                                                                                algo="conv")
            print("\n",stim.filename, "match at",starting_pos)
            display_match(starting_pos, unpacked[1], ref_signals, stim_tuple[1])

            # For Tom's stimuli: The first presentations are short burst encoding the stimulus
            # & the last one is a peak shorter than the rest. I'm trying to crop those out below
            # This needs to be verified later
#             epoch_end, stim_changes = find_stim_onsets(ref_signals)
#             last_frame =  epoch_end[0]

#             stim_duration = np.diff(stim_changes)
            first_frame = input("Please enter the first ref_tp of the first stimulus after the prelude code: ")
            first_frame = int(first_frame)
            last_frame = input("Please enter the last ref_tp of the first stimulus before the end-of-stimulus epoch: ")
            last_frame = int(last_frame)
            stim_tuple, shift_log, frame_replacement = chop_stim_edges(first_frame, last_frame,
                                                   stim_tuple, shift_log, frame_replacement)

            idx_stim_val = 0
            # If the photodiode data was presented in the lower right corner of the screen set '0'
            ud_inv = lr_inv = 0
#             if stim_tuple[2] is not None:
#                 idx_stim_val = 2 #If there is a shader, we take the shader instead of the intensity
            if stim.name == "moving_gratings":
                flipped_inten = flip_gratings(stim_tuple[idx_stim_val], ud_inv, lr_inv)
            else:
                flipped_inten = flip_stimulus(stim_tuple[idx_stim_val], ud_inv, lr_inv)
            stim_datachunk = stim_to_dataChunk(flipped_inten,
                                               starting_pos, #We sliced out the synchro elements of the stim
                                               ref_signals)

    #         stim_datachunk.attrs.update(stim_params) # only works if database data present
            stim_datachunk.attrs["signal_shifts"]     = shift_log
            stim_datachunk.attrs["frame_replacement"] = frame_replacement
            stim_datachunk.attrs["md5"]  = stim.md5
            stim_datachunk.attrs["name"] = stim.filename #Storing the name of the stim in the datachunck aswell

            stim_key = stim.filename
            reM[epoch_counter - 1][stim_key] = stim_datachunk


        # One of Leiron's stimuli
        else:
            # 1) Get the stimulus list from the correct .py QDSpy file
            stim_list, stim_path = get_synced_file_precompile(stim_list_dir, stim)
            # 2) Expand the stimulus values to the actually presented frames
            frame_signals_clustered, stim_ons, stim_offs, epoch_end = \
                    cluster_by_list(photo_data, ref_timepoints, ref_signals, stim_list)
            i = 0
            for onset, offset in zip(stim_ons, stim_offs):

                # 3) Are not read out from the log, but from the photodiode data in this case
                # first_fr is the first actual frame of stimulus presentation, in this case the stimulus starts immediately
                first_fr = onset
                last_fr = offset

                signals = np.array(ref_signals[first_fr:first_fr+len(
                                    frame_signals_clustered[first_fr:last_fr])])
                # Error correction compares the data from the ref_signals with the assigned 'cluster'
                # values.
                stim_tuple, shift_log, frame_replacement  = frame_error_correction(signals,
                                            [np.empty_like(signals),
                                             frame_signals_clustered[first_fr:last_fr],
                                            np.empty_like(signals)],
                                                                            algo="conv")


#                 print("\n",stim.filename, "match at",first_fr)
#                 display_match(first_fr, frame_signals_clustered[first_fr:last_fr],
#                               ref_signals, corrected=None)

                stim_datachunk = stim_to_dataChunk(frame_signals_clustered[first_fr:last_fr],
                                                   first_fr, #No syncro elements in this data
                                                   ref_signals)

                # 8) Add additional info from log file
                stim_datachunk.attrs["md5"]  = stim.md5
                stim_datachunk.attrs["signal_shifts"]     = shift_log
                stim_datachunk.attrs["frame_replacement"] = frame_replacement

                #Loading the stim dictionary and putting the correct names
                stim_idx = np.unique(stim_datachunk.data)[0]
                if stim_idx == -1:
                    stim_datachunk.attrs["name"] = "not recorded, photodiode ended early"
                else:
                    stim_datachunk.attrs["name"] = stim_code_dict[stim_idx] #Storing the name of the stim in the datachunck aswell


                # 9) Update the reM
                stim_key =  stim_datachunk.attrs["name"]
                # Based on stim_key read meaning of stim_key from dictionary
                reM[epoch_counter - 1][stim_key] = stim_datachunk
                i += 1
    return reM

# Cell
from .io import *
from .processing import *
from .extracting import *

from ..core import *
from ..utils import *
from ..database import *
from ..processing import *
from ..plotting import *

import numpy as np
import glob, os, datetime

def master_loop(log, retain_template, tom_stim_dir, sync_dir, calcium_dir, stim_list_dir,
                stim_dict_dir, stack_info_dir):
#     sync_dir = os.path.join(working_dir, "sync/Session01")
    epoch_counter = 0
    sync_files = np.sort(glob.glob(sync_dir + "/Photo*"))
    l_epochs = stack_len_extraction(stack_info_dir)
#     l_epochs[-1] = l_epochs[0] = 9033
    # Files should be numbered from 00 to 99, extract that number
    assert os.path.isdir(tom_stim_dir),      "precompiled stim dir path does not exist"
    assert os.path.isdir(sync_dir),          "sync dir path does not exist"
    assert os.path.isdir(calcium_dir),       "calcium_data dir path does not exist"
    assert os.path.isdir(stim_list_dir),     "stim_list dir path does not exist"
    assert os.path.isdir(stim_dict_dir),     "stim_dict dir path does not exist"
    assert len(sync_files) > 0,              "sync dir doesn't contain any Photodiode files"
    print("len sync_files: {}".format(len(sync_files)))
    print("len l_epochs: {}".format(len(l_epochs)))
#     assert len(sync_files) == len(l_epochs), "number of sync files and recording epochs don't align"

#     # For this particular record
    skip = [12,13]
    skipped_epoch = []
    for skips in reversed(skip):
#     skip = 12 #water sync broken
        skipped_epoch.append(l_epochs.pop(skips-1))

#     del l_epochs[skip]

    twoP_timepoints = []
    for record_file in sync_files:
        record_file_number = record_file.split('_')[-1]
        epoch_counter_current = int(record_file_number)

        if (skip is not None) and (epoch_counter_current == skip):
            print("Skipping recording #{}".format(skip))
            continue

        # Maybe list all files using glob glob instead
        if epoch_counter_current < epoch_counter:
            break
        else:
            # 0) Update file counter
            epoch_counter = epoch_counter_current
            # 1) Load the data
            photo_path = os.path.join(sync_dir,"Photodiode_" + record_file_number)
            # Should be global
#             stim_list_dir = os.path.join("/media/asari/Tristan/01_Code/013_Leiron_Sweep/Day_1")
            print("Processing: " + photo_path)
            photo_data = load_sync_raw(photo_path)
            cam_data = load_sync_raw(os.path.join(sync_dir,"CamTrigs_" + record_file_number))
            treadmill_data = load_sync_raw(os.path.join(sync_dir,"Treadmill_" +
                                                        record_file_number))
            scanning2p_data = load_sync_raw(os.path.join(sync_dir,"SyncPulses_" +
                                                         record_file_number))

#             calcium_dir = os.path.join(working_dir, "calcium_data")
            twoP_timepoints.append(detect_calcium_frames(scanning2p_data)[0]) #MUST BE APPEND TO LIST
            # 2) Get the timestamps for alignment
            lastmodif_time = datetime.datetime.fromtimestamp(int(os.path.getmtime(photo_path)))
            time_lenfile   = datetime.timedelta(0, int(len(photo_data)/10000))
            record_time = lastmodif_time-time_lenfile
            date_str = str(record_time).split(" ")[0]
            # 3) Obtain the stimuli presented during the epoch
            stim_ids = get_stim_ids(log, record_time, lastmodif_time)
            # 4) Obtain the encoding of the stimuli
            print(stim_ids[0].filename)
            if stim_ids[0].name in retain_template:
                stim_code_dict = None
            else:
                stim_code_dict = get_stim_dict(stim_dict_dir,
                                           stim_ids[0].filename)[0]
            # Thresholds and sampling rate might need to be adjusted. Check photo_data
            low_th = 0.1
            if stim_ids[0].name == "water":
                high_th = 0.6
            else:
                high_th = 1.2
            frame_tp, frame_signals = detect_frames(photo_data, low_th, high_th,
                                                    increment=int(10000/60))

            ref_timepoints, ref_signals   = extend_sync_timepoints(frame_tp, frame_signals,
                                                                   up_bound=len(photo_data))
            if epoch_counter == 1:
                reM = RecordMaster([(ref_timepoints, ref_signals)])
            else:
                reM.append(ref_timepoints, ref_signals)
            # 6) Add all stimuli to the record Master
            reM = append_to_reM(epoch_counter, reM, retain_template, stim_list_dir, stim_ids, record_time,
                                photo_data, stim_code_dict, frame_tp, ref_timepoints,
                               ref_signals, tom_stim_dir)
            # 7) Add treadmill data to record Master
            reM = append_tread_movement(reM, epoch_counter, treadmill_data, ref_timepoints)

    # 8) Add the calcium recording data
    C_matrix = np.load(os.path.join(calcium_dir,"C.npy"))
    S_matrix = np.load(os.path.join(calcium_dir,"S.npy"))
    pipe = Data_Pipe(reM, ["main_tp"])
    pipe += "main_tp"
    ca2p_C_datachunks_l = []
    ca2p_S_datachunks_l = []
    ca2p_C_datachunks, ca2p_S_datachunks = twoP_dataChunks_2(pipe,
                                                               twoP_timepoints,
                                                               l_epochs, # compare to other script, shoul dbe - 1?
                                                               0, #Grey was recorded, but excluded from C matrix generation
                                                               skip,
                                                               skipped_epoch,
                                                               C_matrix,
                                                               S_matrix)
    for i, slc in enumerate(pipe._slices):
        # Add 1) Ca-data 2) movement info to every slice (i.e. Contiguous Record)
        reM[slc[0]]["C_matrix"] = ca2p_C_datachunks[i]
        reM[slc[0]]["S_matrix"] = ca2p_S_datachunks[i]

    # 9) Determine still and movement epochs
    # Slice only the parts of the treadmill record when stimuli were presented
    pipe = Data_Pipe(reM, ["treadmill"])
    pipe += "stim"


    movement_threshold = 0.2
    for data_d, slc in zip(pipe, pipe._slices):
        treadmill_dc = data_d["treadmill"]
        tread_mean_speed = np.array(treadmill_dc).mean() # Needs type conversion

        if np.abs(tread_mean_speed) > movement_threshold:
            reM[slc[0]]["moving"] = DataChunk(np.array(treadmill_dc), idx=treadmill_dc.idx,
                                             group="data", fill=treadmill_dc.fill)
        elif np.abs(tread_mean_speed) <= movement_threshold:
            reM[slc[0]]["still"] = DataChunk(np.array(treadmill_dc), idx=treadmill_dc.idx,
                                             group="data", fill=treadmill_dc.fill)

    return reM
