import numpy as np
import os
from collections import namedtuple
from copy import deepcopy
from . import utility as ut
from . import plotting as pl

Event = namedtuple('Event', ['time','name'])

class session:
    'Class containing data from a single session.'
    def __init__(self, file_name, data_path, IDs, file_type):
        print(file_name)
        self.file_name = file_name
        self.subject_ID =  int(file_name.split('-',1)[0][1:])
        self.date = file_name.split('-',1)[1].split('.')[0]
        self.IDs = IDs
        self.file_type = file_type

        # -------------------------------------------------------------------------------------------
        # Import data, extract timestamps and event codes.
        # -------------------------------------------------------------------------------------------

        with open(os.path.join(data_path, file_name), 'r') as data_file:
            split_lines = [line.strip().split(' ') for line in data_file]

        data_lines = [[int(i) for i in line] for line in split_lines if
                       len(line) > 1 and all([len(i)>0 and i[-1].isdigit() for i in line])]

        if self.file_type == 'pyControl_1': # Remove extra 0 from data lines to match format of other file types.
            [line.pop(1) for line in data_lines]
            
        event_lines       = [line for line in data_lines if line[1] in self.IDs.values()]
        block_start_lines = [line for line in data_lines if line[1] == -1] 
        stim_lines        = [line for line in data_lines if line[1] == -2] 

        #  Convert lines to numpy arrays of timestamps (in seconds) and event IDs.

        raw_time_stamps = np.array([line[0] for line in event_lines])
        event_codes     = np.array([line[1] for line in event_lines])

        if self.file_type == 'Arduino':

            if 'start_stop' in list(IDs.keys()):
                start_stop_inds = np.where(event_codes == IDs['start_stop'])[0]
                start_ind = start_stop_inds[0]
                stop_ind  = start_stop_inds[1]
            else:
                start_ind = np.where(event_codes == IDs['session_start'])[0][0]
                stop_ind  = np.where(event_codes == IDs['session_stop' ])[0][0]

            raw_time_stamps = raw_time_stamps[start_ind:stop_ind]
            time_stamps = (raw_time_stamps - raw_time_stamps[0])/1000
            event_codes = event_codes[start_ind:stop_ind]

        else: # File generated by pyControl behaviour system.
            time_stamps = raw_time_stamps / 1000.
            event_codes = event_codes

        ID2event = {v: k for k, v in self.IDs.items()} # Inverse of IDs dict.
        events = [Event(ts, ID2event[ec]) for (ts, ec) in zip(time_stamps, event_codes)]
        self.duration = time_stamps[-1]

        # -------------------------------------------------------------------------------------------
        # Make times dictionary: {'event_name': times_array}
        #--------------------------------------------------------------------------------------------

        self.times = {event_type: time_stamps[event_codes == self.IDs[event_type]]    
                      for event_type in self.IDs.keys()} # Dictionary of event names and times.

        self.times['choice'] = self.ordered_times(['left_active', 'right_active'])

        TsRwWp = [ev for ev in events if ev.name in 
                  ['trial_start', 'left_reward',
                   'right_reward', 'wait_for_poke_out']] + [Event(-1,'')]
        
        self.times['outcome'] =  np.array([TsRwWp[i+1].time for i,ev in 
                                 enumerate(TsRwWp[:-1]) if ev.name == 'trial_start'])

        # -------------------------------------------------------------------------------------------
        # Make dictionary of choices, transitions, second steps and outcomes on each trial.
        #--------------------------------------------------------------------------------------------
        
        second_steps = [ev.name == 'left_active' for ev in events
                        if ev.name in ['left_active', 'right_active']]
        
        outcomes = [TsRwWp[i+1].name in ['left_reward', 'right_reward'] for
                    (i,ev) in enumerate(TsRwWp[:-1]) if ev.name == 'trial_start']

        ChSs = [ev.name for ev in events if ev.name in 
                ['high_poke', 'low_poke', 'left_active', 'right_active']]

        choices = [ChSs[i] == 'high_poke' for (i,ev) in enumerate(ChSs[1:])
                   if ev in ['left_active', 'right_active']]

        self.trial_data = {'choices'     : np.array(choices                    , int), # 1 if high poke, 0 if low poke.
                           'second_steps': np.array(second_steps[:len(choices)], int), # 1 if left, 0 if right.
                           'outcomes'    : np.array(outcomes[:len(choices)]    , int)} # 1 if rewarded, 0 if not.

        self.trial_data['transitions'] = ((self.trial_data['choices'] ==               # 1 if high --> left or low --> right,
                                         self.trial_data['second_steps']).astype(int)) # 0 if high --> right or low --> left.
        
        self.n_trials = len(choices)
        self.rewards  = sum(outcomes)
        self.fraction_rewarded = self.rewards/self.n_trials

        #--------------------------------------------------------------------------------------------
        # Extract block information.
        #--------------------------------------------------------------------------------------------
   
        if len(block_start_lines) > 0:
            start_times = [(line[0]- raw_time_stamps[0])/1000 
                           for line in block_start_lines]
            start_trials = [np.searchsorted(self.times['trial_start'],st)
                            for st in start_times]
            reward_states     = np.array([line[-2] for line in block_start_lines])
            transition_states = np.array([line[-1] for line in block_start_lines])
            end_trials = start_trials[1:] + [self.n_trials]

            if self.file_type == 'Arduino': 
                reward_states   = 2 - reward_states # Arduino data used oposite coding of reward state.
                start_trials[0] = 0  # Timestamp for first block info follows first trial start in Arduino data files.

            trial_trans_state = np.zeros(self.n_trials, dtype = bool) # Boolean array indicating state of tranistion matrix for each trial.
            trial_rew_state   = np.zeros(self.n_trials, dtype = int)  # Integer array indicating state of rewared probabilities for each trial.
            end_trials = start_trials[1:] + [self.n_trials]
            for start_trial,end_trial, trans_state, reward_state in \
                    zip(start_trials, end_trials, transition_states, reward_states):
                trial_trans_state[start_trial:end_trial] = trans_state   
                trial_rew_state[start_trial:end_trial]   = reward_state   

            self.blocks = {'start_times'       : start_times,
                           'start_trials'      : start_trials, # index of first trial of blocks, first trial of session is trial 0. 
                           'end_trials'        : end_trials,
                           'reward_states'     : reward_states,      # 0 for left good, 1 for neutral, 2 for right good.
                           'transition_states' : transition_states,  # 1 for high --> left common, 0 for high --> right common.
                           'trial_trans_state' : trial_trans_state,
                           'trial_rew_state'   : trial_rew_state}   

        #--------------------------------------------------------------------------------------------
        # Extract stim information.    
        #--------------------------------------------------------------------------------------------

        if len(stim_lines) > 0:
            stim_timestamps = (np.array([line[ 0] for line in stim_lines])) / 1000
            stim_state =                [line[-1] for line in stim_lines] # 1: onset, 0: offset.

            self.times['stim_on' ] = np.array([stim_timestamps[i] for i,st in 
                                               enumerate(stim_state) if st == 1])
            self.times['stim_off'] = np.array([stim_timestamps[i] for i,st in 
                                               enumerate(stim_state) if st == 0])
            
            stim_trials = np.zeros(self.n_trials + 1, bool) 
            stim_trials[self.times['choice'].searchsorted(self.times['stim_on']
                        + 0.005)] = True
            self.stim_trials = stim_trials[:-1] # Boolean array indicating for each choice whether stim occured since previous choice.

    #------------------------------------------------------------------------------------------------

    def select_trials(self, selection_type, select_n = 20, first_n_mins = False,
                      block_type = 'all'):
        ''' Select specific trials for analysis.  

        The first selection step is specified by selection_type:

        'end' : Only final select_n trials of each block are selected.

        'xtr' : Select all trials except select_n trials following transition reversal.

        'all' : All trials are included.

        The first_n_mins argument can be used to select only trials occuring within
        a specified number of minutes of the session start.

        The block_type argument allows additional selection for only 'neutral' or 'non_neutral' blocks.
        '''

        assert selection_type in ['end', 'xtr', 'all'], 'Invalid trial select type.'

        if selection_type == 'xtr': # Select all trials except select_n following transition reversal.
            trials_to_use = np.ones(self.n_trials, dtype = bool)
            trans_change = np.hstack((
                False, ~np.equal(self.blocks['transition_states'][:-1],
                                 self.blocks['transition_states'][1:])))
            start_trials = (self.blocks['start_trials'] + 
                            [self.blocks['end_trials'][-1] + select_n])
            for i in range(len(trans_change)):
                if trans_change[i]:
                    trials_to_use[start_trials[i]:start_trials[i] + select_n] = False

        elif selection_type == 'end': # Select only select_n trials before block transitions.
            trials_to_use = np.zeros(self.n_trials, dtype = bool)
            for b in self.blocks['start_trials'][1:]:
                trials_to_use[b - 1 - select_n:b -1] = True

        elif selection_type == 'all': # Use all trials.
            trials_to_use = np.ones(self.n_trials, dtype = bool)
            
        if first_n_mins:  #  Restrict analysed trials to only first n minutes. 
            time_selection = self.times['trial_start'][:self.n_trials] < (60*first_n_mins)
            trials_to_use = trials_to_use & time_selection

        if not block_type == 'all': #  Restrict analysed trials to blocks of certain types.
            if block_type == 'neutral':       # Include trials only from neutral blocks.
                block_selection = self.blocks['trial_rew_state'] == 1
            elif block_type == 'non_neutral': # Include trials only from non-neutral blocks.
                block_selection = self.blocks['trial_rew_state'] != 1
            trials_to_use = trials_to_use & block_selection

        return trials_to_use

    #------------------------------------------------------------------------------------------------

    def plot(self, fig_no = 1): pl.plot_session(self, fig_no)

    def ordered_times(self,event_types):
        'Return a array of specified event times in order of occurence.'
        return np.sort(np.hstack([self.times[event_type] for event_type in event_types]))

    def get_IDs(self, event_list):
        return [self.IDs[val] for val in event_list]

    def unpack_trial_data(self, order = 'CTSO', dtype = int):
        'Return elements of trial_data dictionary in specified order and data type.'
        o_dict = {'C': 'choices', 'T': 'transitions', 'S': 'second_steps', 'O': 'outcomes'}
        if dtype == int:
            return [self.trial_data[o_dict[i]] for i in order]
        else:
            return [self.trial_data[o_dict[i]].astype(dtype) for i in order]



#--------------------------------------------------------------------------------------------
# Concatenated session class.  
#--------------------------------------------------------------------------------------------

class concatenated_session:
    '''Concatinates a set of consecutive sessions into a single 
    long session'''

    def __init__(self, subject_sessions):
        assert len(set([s.subject_ID for s in subject_sessions])) == 1, \
            'All sessions must be from a single subject to concatenate.'
        subject_sessions = sorted(subject_sessions, key = lambda s:s.day)
        choices, transitions, second_steps, outcomes, trial_trans_state, trial_rew_state = ([],[],[],[],[],[]) 
        self.n_trials = sum([s.n_trials for s in subject_sessions])
        session_start_trials = np.zeros(self.n_trials, bool)
        session_start_trial = 0
        for i, session in enumerate(subject_sessions):
                choices      += session.trial_data['choices'].tolist()
                transitions  += session.trial_data['transitions'].tolist()
                second_steps += session.trial_data['second_steps'].tolist()
                outcomes     += session.trial_data['second_steps'].tolist()
                trial_trans_state += session.blocks['trial_trans_state'].tolist() 
                trial_rew_state   += session.blocks['trial_rew_state'].tolist()            
                session_start_trials[session_start_trial] = True
                session_start_trial += session.n_trials

        self.trial_data = {'choices'      : np.array(choices),
                     'transitions'  : np.array(transitions),
                     'second_steps' : np.array(second_steps),
                     'outcomes'     : np.array(outcomes)
                     }

        self.blocks = {'trial_trans_state': np.array(trial_trans_state),
                       'trial_rew_state'  : np.array(trial_rew_state),
                       'session_start_trials': session_start_trials}

        self.IDs = subject_sessions[0].IDs
        self.subject_ID = subject_sessions[0].subject_ID
        self.n_trials = sum([s.n_trials for s in subject_sessions])

#--------------------------------------------------------------------------------------------
# Trim sessions
#--------------------------------------------------------------------------------------------

def trim_sessions(sessions, start_time = 0., end_time = np.inf):
    '''Return copy of sessions where each session has been trimmed such that only 
    trials occuring after start_time (minutes) and before end_time are included.
    Only session.trial_data and session.blocks are modified by trim_sessions().'''
    if type(sessions) is not list: # Single session passed in, convert to list.
        sessions = [sessions]
    trimmed_sessions = deepcopy(sessions)
    for session in trimmed_sessions:
        start_trial = sum(session.times['trial_start']<(60*start_time))
        end_trial   = sum(session.times['trial_start']<(60*end_time  ))
        for key in list(session.trial_data.keys()):
            session.trial_data[key] = session.trial_data[key][start_trial:end_trial]
        session.n_trials = len(session.trial_data['choices'])
        if hasattr(session,'blocks'): 
            start_block = sum(np.array(session.blocks['end_trials']) < start_trial)
            end_block   = sum(np.array(session.blocks['end_trials']) < end_trial)+1
            session.blocks = {
              'start_trials' : [t-start_trial for t in 
                                session.blocks['start_trials']][start_block:end_block],
              'end_trials'   : [t-start_trial for t in
                                session.blocks['end_trials']][start_block:end_block],
              'reward_states': session.blocks['reward_states'][start_block:end_block],
              'transition_states': session.blocks['transition_states'][
                                       start_block:end_block],
              'trial_rew_state'  : session.blocks['trial_rew_state'][
                                       start_trial:end_trial],
              'trial_trans_state': session.blocks['trial_trans_state'][
                                       start_trial:end_trial]}
            session.blocks['start_trials'][0] = 0
            session.blocks['end_trials'][-1]  = session.n_trials
    if len(trimmed_sessions) == 1: # Return session if session passed in.
        return trimmed_sessions[0]
    else: # Return list of sessions.
        return trimmed_sessions
