import os
import numpy as np
import concurrent.futures
import tqdm
from analysis_utils import makePSTH_numba
from scipy.interpolate import interp1d
from functools import partial


DEFAULT_CACHE_DIR = '/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/vbn_s3_cache/'
_CACHE_BY_DIR = {}


def get_cache(cache_dir=None):
    cache_dir = cache_dir or os.environ.get('VBN_CACHE_DIR', DEFAULT_CACHE_DIR)
    if cache_dir not in _CACHE_BY_DIR:
        from allensdk.brain_observatory.behavior.behavior_project_cache.\
            behavior_neuropixels_project_cache \
            import VisualBehaviorNeuropixelsProjectCache

        _CACHE_BY_DIR[cache_dir] = VisualBehaviorNeuropixelsProjectCache.from_s3_cache(
            cache_dir=cache_dir
        )
    return _CACHE_BY_DIR[cache_dir]


def get_session(session_id, cache_dir=None):

    return get_cache(cache_dir).get_ecephys_session(session_id)


def find_running_acceleration_deceleration_times(session, stimulus_block=5):
    running = session.running_speed
    running.loc[running['speed']<0, 'speed'] = 0
    
    # Calculate rolling means for the speed column
    rolling_mean_before = running['speed'].rolling(window=30).mean().shift(1)
    rolling_mean_after = running['speed'].rolling(window=30).mean().shift(-29)

    # Find indices where the acceleration conditions are met
    condition = (rolling_mean_before < 1) & (rolling_mean_after > 5)
    indices = np.where(condition)[0]

    acceleration_times = []
    for ir in indices:
        if (ir > 30) and (ir < len(running) - 31):
            window_indices = running.iloc[ir-30:ir+30].index.values
            min_index = running.loc[window_indices].idxmin()['speed']
            window_diffs = running.loc[min_index:window_indices[-1]].diff()
            try:
                max_diff_index = window_diffs.idxmax()['speed']
                if running.loc[max_diff_index]['speed'] < 1:
                    last_point_below_threshold = running.loc[max_diff_index]['timestamps']
                else:
                    last_point_below_threshold = np.where(running.loc[min_index:max_diff_index]<1)[0][-1]
                    last_point_below_threshold = running.loc[min_index+last_point_below_threshold]['timestamps']
                if len(acceleration_times) > 0:
                    if last_point_below_threshold - acceleration_times[-1] < 0.5:
                        continue
                acceleration_times.append(last_point_below_threshold)
            except Exception as exc:
                sess_id = session.metadata['ecephys_session_id']
                print(f'{sess_id} generated an exception: {exc}')
                continue

    stims = session.stimulus_presentations
    passive_start = stims[stims['stimulus_block']==stimulus_block]['start_time'].iloc[0]
    passive_end = stims[stims['stimulus_block']==stimulus_block]['end_time'].iloc[-1]

    acceleration_times=np.array(acceleration_times)
    passive_acceleration_times = acceleration_times[(acceleration_times>passive_start)&(acceleration_times<passive_end)]


    # Find indices where the deceleration conditions are met
    condition = (rolling_mean_before > 5) & (rolling_mean_after < 1)
    indices = np.where(condition)[0]

    deceleration_times = []
    for ir in indices:
        if (ir > 30) and (ir < len(running) - 31):
            window_indices = running.iloc[ir-30:ir+30].index.values
            max_index = running.loc[window_indices].idxmax()['speed']
            min_index = running.loc[max_index:window_indices[-1]].idxmin()['speed']
            # window_diffs = running.loc[max_index:window_indices[-1]].diff()
            window_diffs = running.loc[window_indices[0]:min_index].diff()
            min_diff_index = window_diffs.idxmin()['speed']
            max_decel_point = running.loc[min_diff_index]['timestamps']
            # if running.loc[min_diff_index]['speed'] > 5:
            #     last_point_above_threshold = running.loc[min_diff_index]['timestamps']
            # else:
            #     last_point_above_threshold = np.where(running.loc[max_index:min_diff_index]>5)[0][-1]
            #     last_point_above_threshold = running.loc[max_index+last_point_above_threshold]['timestamps']
            if len(deceleration_times) > 0:
                if max_decel_point - deceleration_times[-1] < 0.5:
                    continue
            deceleration_times.append(max_decel_point)

    deceleration_times=np.array(deceleration_times)
    passive_deceleration_times = deceleration_times[(deceleration_times>passive_start)&(deceleration_times<passive_end)]


    return passive_acceleration_times, passive_deceleration_times


def unit_peth(session_id, alignment_func, time_before, time_after, binsize, cache_dir=None):

    session = get_session(session_id, cache_dir=cache_dir)
    alignment_times1, alignment_times2 = alignment_func(session)

    units = session.get_units()
    units = units[units['quality']=='good']

    total_time = time_before+time_after
    time_bins = int(total_time/binsize)

    condition_peths = []
    for alignment_times in [alignment_times1, alignment_times2]:
        if len(alignment_times)<5:
            condition_peths.append(np.full((len(units), time_bins), np.nan))
            unit_ids = units.index.values
            continue

        peths = []
        unit_ids = []
        for unit in units.index.values:

            spike_times = session.spike_times[unit]
            peth, time = makePSTH_numba(spike_times, alignment_times-time_before, 
                                    total_time, binSize=binsize)
            
            peths.append(peth[:time_bins])
            unit_ids.append(unit)
    
        condition_peths.append(peths)

    return np.array(condition_peths[0]), np.array(condition_peths[1]), unit_ids


def unit_averaged_psth_time_aligned(session_list, alignment_time_func = 'running',
                                    time_before=0.5, time_after=0.5, binsize=0.001, stimulus_block=None,
                                    cache_dir=None):

    if alignment_time_func == 'running':
        if stimulus_block is None:
            stimulus_block = 5
        alignment_time_func = partial(find_running_acceleration_deceleration_times, stimulus_block=stimulus_block)

    pool = concurrent.futures.ProcessPoolExecutor(max_workers=20)        
    future_to_session = {}
    for session in session_list:

        fut = pool.submit(unit_peth,
                            session,
                            alignment_time_func,
                            time_before,
                            time_after,
                            binsize,
                            cache_dir)

        future_to_session[fut] = session

    session_data_1 = []
    session_data_2 = []
    unit_ids = []
    session_ids = []
    for future in tqdm.tqdm(concurrent.futures.as_completed(future_to_session), total=len(future_to_session), leave=True):
    #for future in concurrent.futures.as_completed(future_to_session):

        session = future_to_session[future]
        try:
            data = future.result()
            session_data_1.append(data[0])
            session_data_2.append(data[1])
            unit_ids.append(data[2])
            session_ids.append(session)

        except Exception as exc:
            print(f'{session} generated an exception: {exc}')

    return session_data_1, session_data_2, unit_ids


def resample_df_to_times(df, time_column, val_column, new_times):

    timestamps = df[time_column].values*1000
    vals = df[val_column].values
    interpolator = interp1d(timestamps, vals, kind='linear', bounds_error=False)#, fill_value=np.nan)
    new_values = interpolator(new_times)

    return new_values, new_times