import sys, math, os, time
import os
from intanutil.read_header import read_header
from intanutil.get_bytes_per_data_block import get_bytes_per_data_block
from intanutil.read_one_data_block import read_one_data_block
from intanutil.notch_filter import notch_filter
from intanutil.data_to_result import data_to_result
import os
import spikeinterface
import spikeinterface.extractors as se
import spikeinterface.sorters as ss
import spikeinterface.widgets as sw
from spikeinterface.core.npzsortingextractor import NpzSortingExtractor
from pylab import *
ss.Kilosort3Sorter.set_kilosort3_path('/kilosort3')
import pylab
from probeinterface import Probe
import copy
import sys
sys.path.append("./autosort/util/") 
import scipy
import pickle
from collections import Counter
import numpy as np


def save_obj(objt, name):
    with(open(name+'.pkl','wb')) as f:
        pickle.dump(objt,f,pickle.HIGHEST_PROTOCOL)
        f.close()



def read_data(filename):
    """Reads Intan Technologies RHD2000 data file generated by evaluation board GUI.

    Data are returned in a dictionary, for future extensibility.
    """

    tic = time.time()
    fid = open(filename, 'rb')
    filesize = os.path.getsize(filename)

    header = read_header(fid)

    print('Found {} amplifier channel{}.'.format(header['num_amplifier_channels'],
                                                 plural(header['num_amplifier_channels'])))
    print('Found {} auxiliary input channel{}.'.format(header['num_aux_input_channels'],
                                                       plural(header['num_aux_input_channels'])))
    print('Found {} supply voltage channel{}.'.format(header['num_supply_voltage_channels'],
                                                      plural(header['num_supply_voltage_channels'])))
    print('Found {} board ADC channel{}.'.format(header['num_board_adc_channels'],
                                                 plural(header['num_board_adc_channels'])))
    print('Found {} board digital input channel{}.'.format(header['num_board_dig_in_channels'],
                                                           plural(header['num_board_dig_in_channels'])))
    print('Found {} board digital output channel{}.'.format(header['num_board_dig_out_channels'],
                                                            plural(header['num_board_dig_out_channels'])))
    print('Found {} temperature sensors channel{}.'.format(header['num_temp_sensor_channels'],
                                                           plural(header['num_temp_sensor_channels'])))
    print('')

    # Determine how many samples the data file contains.
    bytes_per_block = get_bytes_per_data_block(header)

    # How many data blocks remain in this file?
    data_present = False
    bytes_remaining = filesize - fid.tell()
    if bytes_remaining > 0:
        data_present = True

    if bytes_remaining % bytes_per_block != 0:
        raise Exception('Something is wrong with file size : should have a whole number of data blocks')

    num_data_blocks = int(bytes_remaining / bytes_per_block)

    num_amplifier_samples = header['num_samples_per_data_block'] * num_data_blocks
    num_aux_input_samples = int((header['num_samples_per_data_block'] / 4) * num_data_blocks)
    num_supply_voltage_samples = 1 * num_data_blocks
    num_board_adc_samples = header['num_samples_per_data_block'] * num_data_blocks
    num_board_dig_in_samples = header['num_samples_per_data_block'] * num_data_blocks
    num_board_dig_out_samples = header['num_samples_per_data_block'] * num_data_blocks

    record_time = num_amplifier_samples / header['sample_rate']

    if data_present:
        print('File contains {:0.3f} seconds of data.  Amplifiers were sampled at {:0.2f} kS/s.'.format(record_time,
                                                                                                        header[
                                                                                                            'sample_rate'] / 1000))
    else:
        print('Header file contains no data.  Amplifiers were sampled at {:0.2f} kS/s.'.format(
            header['sample_rate'] / 1000))

    if data_present:
        # Pre-allocate memory for data.
        print('')
        print('Allocating memory for data...')

        data = {}
        if (header['version']['major'] == 1 and header['version']['minor'] >= 2) or (header['version']['major'] > 1):
            data['t_amplifier'] = np.zeros(num_amplifier_samples, dtype=np.int)
        else:
            data['t_amplifier'] = np.zeros(num_amplifier_samples, dtype=np.uint)

        data['amplifier_data'] = np.zeros([header['num_amplifier_channels'], num_amplifier_samples], dtype=np.uint)
        data['aux_input_data'] = np.zeros([header['num_aux_input_channels'], num_aux_input_samples], dtype=np.uint)
        data['supply_voltage_data'] = np.zeros([header['num_supply_voltage_channels'], num_supply_voltage_samples],
                                               dtype=np.uint)
        data['temp_sensor_data'] = np.zeros([header['num_temp_sensor_channels'], num_supply_voltage_samples],
                                            dtype=np.uint)
        data['board_adc_data'] = np.zeros([header['num_board_adc_channels'], num_board_adc_samples], dtype=np.uint)

        # by default, this script interprets digital events (digital inputs and outputs) as booleans
        # if unsigned int values are preferred(0 for False, 1 for True), replace the 'dtype=np.bool' argument with 'dtype=np.uint' as shown
        # the commented line below illustrates this for digital input data; the same can be done for digital out

        # data['board_dig_in_data'] = np.zeros([header['num_board_dig_in_channels'], num_board_dig_in_samples], dtype=np.uint)
        data['board_dig_in_data'] = np.zeros([header['num_board_dig_in_channels'], num_board_dig_in_samples],
                                             dtype=np.bool)
        data['board_dig_in_raw'] = np.zeros(num_board_dig_in_samples, dtype=np.uint)

        data['board_dig_out_data'] = np.zeros([header['num_board_dig_out_channels'], num_board_dig_out_samples],
                                              dtype=np.bool)
        data['board_dig_out_raw'] = np.zeros(num_board_dig_out_samples, dtype=np.uint)

        # Read sampled data from file.
        print('Reading data from file...')

        # Initialize indices used in looping
        indices = {}
        indices['amplifier'] = 0
        indices['aux_input'] = 0
        indices['supply_voltage'] = 0
        indices['board_adc'] = 0
        indices['board_dig_in'] = 0
        indices['board_dig_out'] = 0

        print_increment = 10
        percent_done = print_increment
        for i in range(num_data_blocks):
            read_one_data_block(data, header, indices, fid)

            # Increment indices
            indices['amplifier'] += header['num_samples_per_data_block']
            indices['aux_input'] += int(header['num_samples_per_data_block'] / 4)
            indices['supply_voltage'] += 1
            indices['board_adc'] += header['num_samples_per_data_block']
            indices['board_dig_in'] += header['num_samples_per_data_block']
            indices['board_dig_out'] += header['num_samples_per_data_block']

            fraction_done = 100 * (1.0 * i / num_data_blocks)
            if fraction_done >= percent_done:
                print('{}% done...'.format(percent_done))
                percent_done = percent_done + print_increment

        # Make sure we have read exactly the right amount of data.
        bytes_remaining = filesize - fid.tell()
        if bytes_remaining != 0: raise Exception('Error: End of file not reached.')

    # Close data file.
    fid.close()

    if (data_present):
        print('Parsing data...')

        # Extract digital input channels to separate variables.
        for i in range(header['num_board_dig_in_channels']):
            data['board_dig_in_data'][i, :] = np.not_equal(
                np.bitwise_and(data['board_dig_in_raw'], (1 << header['board_dig_in_channels'][i]['native_order'])), 0)

        # Extract digital output channels to separate variables.
        for i in range(header['num_board_dig_out_channels']):
            data['board_dig_out_data'][i, :] = np.not_equal(
                np.bitwise_and(data['board_dig_out_raw'], (1 << header['board_dig_out_channels'][i]['native_order'])),
                0)

        # Scale voltage levels appropriately.
        data['amplifier_data'] = np.multiply(0.195,
                                             (data['amplifier_data'].astype(np.int32) - 32768))  # units = microvolts
        data['aux_input_data'] = np.multiply(37.4e-6, data['aux_input_data'])  # units = volts
        data['supply_voltage_data'] = np.multiply(74.8e-6, data['supply_voltage_data'])  # units = volts
        if header['eval_board_mode'] == 1:
            data['board_adc_data'] = np.multiply(152.59e-6,
                                                 (data['board_adc_data'].astype(np.int32) - 32768))  # units = volts
        elif header['eval_board_mode'] == 13:
            data['board_adc_data'] = np.multiply(312.5e-6,
                                                 (data['board_adc_data'].astype(np.int32) - 32768))  # units = volts
        else:
            data['board_adc_data'] = np.multiply(50.354e-6, data['board_adc_data'])  # units = volts
        data['temp_sensor_data'] = np.multiply(0.01, data['temp_sensor_data'])  # units = deg C

        # Check for gaps in timestamps.
        num_gaps = np.sum(np.not_equal(data['t_amplifier'][1:] - data['t_amplifier'][:-1], 1))
        if num_gaps == 0:
            print('No missing timestamps in data.')
        else:
            print('Warning: {0} gaps in timestamp data found.  Time scale will not be uniform!'.format(num_gaps))

        # Scale time steps (units = seconds).
        data['t_amplifier'] = data['t_amplifier'] / header['sample_rate']
        data['t_aux_input'] = data['t_amplifier'][range(0, len(data['t_amplifier']), 4)]
        data['t_supply_voltage'] = data['t_amplifier'][
            range(0, len(data['t_amplifier']), header['num_samples_per_data_block'])]
        data['t_board_adc'] = data['t_amplifier']
        data['t_dig'] = data['t_amplifier']
        data['t_temp_sensor'] = data['t_supply_voltage']

        # If the software notch filter was selected during the recording, apply the
        # same notch filter to amplifier data here.
        if header['notch_filter_frequency'] > 0 and header['version']['major'] < 3:
            print('Applying notch filter...')

            print_increment = 10
            percent_done = print_increment
            for i in range(header['num_amplifier_channels']):
                data['amplifier_data'][i, :] = notch_filter(data['amplifier_data'][i, :], header['sample_rate'],
                                                            header['notch_filter_frequency'], 10)

                fraction_done = 100 * (i / header['num_amplifier_channels'])
                if fraction_done >= percent_done:
                    print('{}% done...'.format(percent_done))
                    percent_done += print_increment
    else:
        data = [];

    # Move variables to result struct.
    result = data_to_result(header, data, data_present)

    print('Done!  Elapsed time: {0:0.1f} seconds'.format(time.time() - tic))
    return result



def plural(n):
    """Utility function to optionally pluralize words based on the value of n.
    """

    if n == 1:
        return ''
    else:
        return 's'



def sorting_day_split(sorting, date_id_all, day_length, pack_folder,
                      sorting_save_name='firings_inlier'):
    sampling_freq = sorting.get_sampling_frequency()

    fig, ax = plt.subplots(1, 1, figsize=(20, 40))
    sw.plot_rasters(sorting, time_range=(0, np.sum(day_length)), ax=ax)

    colors = []
    cm = pylab.get_cmap('rainbow')
    NUM_COLORS = len(day_length)
    for i in range(NUM_COLORS):
        colors.append(cm(1. * i / NUM_COLORS))  # color will now be an RGBA tuple

    for i in range(len(day_length)):
        ax.axvspan(np.sum(day_length[:i]) / sampling_freq, np.sum(day_length[:(i + 1)]) / sampling_freq,
                   facecolor=colors[i], alpha=0.1)
    plt.show()
    plt.savefig(pack_folder + '/sorting/' + sorting_save_name + '.png', dpi=300)

    for i in range(len(day_length)):
        pack_folder_i = pack_folder + '/' + date_id_all[i] + '/'

        if os.path.exists(pack_folder_i) == False:
            os.mkdir(pack_folder_i)

        start_frame = np.sum(day_length[:i])
        end_frame = np.sum(day_length[:(i + 1)])

        sub_sorting = sorting.frame_slice(start_frame, end_frame)

        keep_unit_ids = []
        for unit_id in sub_sorting.unit_ids:
            spike_train = sub_sorting.get_unit_spike_train(unit_id=unit_id)
            n = spike_train.size
            if (n > 20):
                keep_unit_ids.append(unit_id)

        curated_sub_sorting = sub_sorting.select_units(unit_ids=keep_unit_ids, renamed_unit_ids=None)

        save_path = pack_folder_i + '/sorting/' + sorting_save_name + '.npz'
        if os.path.exists(pack_folder_i + '/sorting/') == False:
            os.mkdir(pack_folder_i + '/sorting/')
        NpzSortingExtractor.write_sorting(curated_sub_sorting, save_path)



def create_mesh_probe(n):
    positions = np.zeros((n, 2))
    for i in range(n):
        x = i // 8
        y = i % 8
        positions[i] = x, y
    positions *= 300

    test = [positions[i] for i in range(8, 16)]
    test1 = copy.deepcopy(test[::-1])
    for idx, i in enumerate(range(8, 16)):
        positions[i] = test1[idx]

    test = [positions[i] for i in range(24, 32)]
    test1 = copy.deepcopy(test[::-1])
    for idx, i in enumerate(range(24, 32)):
        positions[i] = test1[idx]

    mesh_probe = Probe(ndim=2, si_units='um')
    mesh_probe.set_contacts(positions=positions, shapes='circle', shape_params={'radius': 5})

    ant = {'first_index': 0}
    mesh_probe.annotate(**ant)
    channel_indices_raw = np.arange(32)
    channel_indices = [i for i in channel_indices_raw]
    mesh_probe.set_device_channel_indices(channel_indices)
    return mesh_probe



def stack_recordings(pack_folder_pre, mesh_probe, trigger_val=2):

    cont_data_all = []
    cont_trigger_all = []
    for dirpath, dirname, filenames in os.walk(pack_folder_pre):
        for i in filenames:

            if '.rhd' in i:
                print(dirpath+'/'+i)

                raw_data=read_data(dirpath+'/'+i)

                sampling_freq = raw_data['frequency_parameters']['amplifier_sample_rate']
                test=raw_data['board_adc_data']
#                 test[test<trigger_val]=0
#                 test[test>=trigger_val]=1
                cont_trigger_all.append(test)
                cont_data_all.append(raw_data['amplifier_data'].T)

    cont_data_all = np.vstack(cont_data_all)
    cont_trigger_all = np.concatenate(cont_trigger_all,axis=1)

    np_mean = np.mean(cont_data_all)
    np_std = np.std(cont_data_all)
    threshold_neg = -300
    threshold_pos = 300

    neg_indices = np.where(cont_data_all<threshold_neg)
    pos_indices = np.where(cont_data_all>threshold_pos)


    remove_indices_neg_0 = [neg_indices[0]+i for i in np.arange(-10,10)]
    remove_indices_neg_1 = [neg_indices[1] for i in np.arange(-10,10)]

    remove_indices_pos_0 = [pos_indices[0]+i for i in np.arange(-10,10)]
    remove_indices_pos_1 = [pos_indices[1] for i in np.arange(-10,10)]

    remove_indices_0 = np.array(np.append(remove_indices_neg_0,remove_indices_pos_0)).flatten()
    remove_indices_1 = np.array(np.append(remove_indices_neg_1,remove_indices_pos_1)).flatten()


    remove_indices_0_temp = np.where(remove_indices_0>0, remove_indices_0, 0)
    remove_indices_0 = np.where(remove_indices_0_temp<cont_data_all.shape[0]-1, remove_indices_0_temp, cont_data_all.shape[0]-1)

    cont_data_all[remove_indices_0,remove_indices_1] = np_mean


    recording = se.NumpyRecording(traces_list=cont_data_all,
                              sampling_frequency=sampling_freq)

    recording.set_probe(mesh_probe, in_place=True)

    return recording, cont_data_all,cont_trigger_all


def read_intan_data_from_dir(date_id_all, pack_folder_ix=['./', '22_C6']):
    recording_traces = []
    session_length_concat = []
    day_length = []
    cont_trigger_all_all = []
    save_folder_name = '_'.join(date_id_all)
    data_folder_all = f'./processed_data/Ephys_concat_{save_folder_name}/'

    if os.path.exists(data_folder_all):
        recording_concat = spikeinterface.core.base.BaseExtractor.load_from_folder(data_folder_all)
        session_length_concat = np.load(data_folder_all + 'session_length.npy')
        day_length = np.load(data_folder_all + 'day_length.npy')
        cont_trigger_all_all = np.load(data_folder_all + 'cont_trigger_all_all.npy')
    else:
        for date_id in date_id_all:
            pack_folder_pre = pack_folder_ix[0] + date_id + pack_folder_ix[1]

            data_folder_pre = f'./processed_data/Ephys_{date_id}/'

            if os.path.exists(data_folder_pre):
                recording = spikeinterface.core.base.BaseExtractor.load_from_folder(data_folder_pre)
                session_length = np.load(data_folder_pre + 'session_length.npy')
                cont_trigger_all = np.load(data_folder_pre + 'cont_trigger_all.npy')
            else:
                mesh_probe = create_mesh_probe(32)
                # plot_probe(mesh_probe,with_channel_index=True)

                recording, session_length, cont_trigger_all = stack_recordings(pack_folder_pre, mesh_probe,
                                                                               trigger_val=3)
                #             recording.set_probe(mesh_probe, in_place=True)

                recording.set_probe(mesh_probe, in_place=True)

                recording = recording.save(folder=data_folder_pre)
                np.save(data_folder_pre + 'session_length.npy', session_length)
                np.save(data_folder_pre + 'cont_trigger_all.npy', cont_trigger_all)

            sampling_freq = recording.get_sampling_frequency()
            recording_trace = recording.get_traces()
            recording_traces.append(recording_trace)
            session_length_concat.append(session_length)
            cont_trigger_all_all.append(cont_trigger_all)
            print(recording_trace.shape)
            day_length.append(recording_trace.shape[0])

        recording_traces = np.vstack(recording_traces)
        session_length_concat = np.vstack(session_length_concat)

        recording_concat = se.NumpyRecording(traces_list=recording_traces,
                                             sampling_frequency=sampling_freq)
        recording_concat.set_probe(mesh_probe, in_place=True)

        recording_concat = recording_concat.save(folder=data_folder_all)

        np.save(data_folder_all + 'session_length.npy', session_length_concat)
        np.save(data_folder_all + 'day_length.npy', day_length)
        np.save(data_folder_all + 'cont_trigger_all_all.npy', cont_trigger_all_all)

    print(recording_concat)
    print('Num. channels = {}'.format(len(recording_concat.get_channel_ids())))
    print('Sampling frequency = {} Hz'.format(recording_concat.get_sampling_frequency()))
    print('Num. timepoints seg0= {}'.format(recording_concat.get_num_segments()))

    return recording_concat, session_length_concat, day_length, cont_trigger_all_all




def read_ns4datafile(datafile):
    brpylib_ver_req = "1.3.1"
    if brpylib_ver.split('.') < brpylib_ver_req.split('.'):
        raise Exception("requires brpylib " + brpylib_ver_req + " or higher, please use latest version")
    elec_ids = 'all'  # 'all' is default for all (1-indexed)                                        
    nsx_file = NsxFile(datafile)
    # Extract data - note: data will be returned based on *SORTED* elec_ids, see cont_data['elec_ids']
    cont_data = nsx_file.getdata(elec_ids)
    unit = nsx_file.extended_headers[1]['Units']
    # Close the nsx file now that all data is out
    nsx_file.close()
    return cont_data, unit



def detect_spike(trace0_car,thr_min = 5, distance=30, ch_max_simul_firing = 3):
#     noise_std_detect = np.median(abs(trace0_car - np.mean(trace0_car, axis=0)[None, :]) / 0.6745, axis=0)
    noise_std_detect = np.median(abs(trace0_car)/0.6745, axis=0)

    thr = thr_min * noise_std_detect
    thrmax = 30 * noise_std_detect

    spikes = np.zeros(trace0_car.shape)
    for i in range(noise_std_detect.shape[0]):
        peaks, _ = scipy.signal.find_peaks(-trace0_car[:, i], thr[i], distance=distance)
        spikes[peaks, i] = 1

    print(np.sum(spikes == 1))
    # larger value no more than thrmax
    points = trace0_car.shape[0]
    spike_coord = np.argwhere(spikes == 1)
    for i in range(spike_coord.shape[0]):
        near_start = spike_coord[i, 0] - 5
        near_end = spike_coord[i, 0] + 5
        if near_start < 0:
            near_start = 0
        if near_end >= points:
            near_end = points - 1
        if np.any(np.max(trace0_car[near_start:near_end, :], axis=0) >= thrmax):
            spikes[spike_coord[i, 0], spike_coord[i, 1]] = 0

    # no simultanous firing!!!!
    thres_cross = ch_max_simul_firing
    spikes[np.sum(spikes, axis=1) > thres_cross, :] = 0
    print(np.sum(spikes == 1))

    return spikes



def list_substraction(a, b):

    ca = Counter(a)
    cb = Counter(b)

    result_b = sorted((cb - ca).elements())
    return result_b





def offline_smoother(yIn, kernSD, stepSize):
    causal = False;

    if (kernSD == 0) or (yIn.shape[1]==1):
        yOut = yIn
        return yOut


    # Filter half length
    # Go 3 standard deviations out
    fltHL = math.ceil(3 * kernSD / stepSize)

    # Length of flt is 2*fltHL + 1
    flt = scipy.stats.norm.pdf(np.arange(-fltHL*stepSize ,fltHL*stepSize,stepSize), loc=0, scale = kernSD)


    [yDim, T] = shape(yIn)
    yOut      = np.zeros((yDim, T))

    nm = scipy.signal.convolve(flt.reshape(1,-1), np.ones((1,T)), mode='full', method='auto')

    for i in range(yDim):

        ys = np.divide(scipy.signal.convolve(flt.reshape(1,-1),yIn[i,:].reshape(1,-1)) , nm)
    #     # Cut off edges so that result of convolution is same length
    #     # as original data
        yOut[i,:] = ys[0,fltHL:ys.shape[1]-fltHL+1]
    return yOut




class FSM2:
    def __init__(self, k, d, tau0=None, Minv0=None, W0=None, learning_rate=None):
        self.k = k
        self.d = d
        if W0 is None:
            self.W = np.eye(k, d) / d
        else:
            self.W = W0;
        self.M = np.eye(k)
        if Minv0 is None:
            self.Minv = np.eye(k)
        else:
            self.Minv = Minv0
        if tau0 is None:
            self.tau = 0.5
        else:
            self.tau = tau0
        if learning_rate is None:
            self.lr = lambda x: 1.0 / (2 * (x + 100) + 5)
        else:
            self.lr = lambda x: learning_rate

        self.t = 0
        self.outer_W = np.zeros(shape(self.W))
        self.outer_Minv = np.zeros(shape(self.Minv))
        self.y = []

    def fit_next(self, x):
        step = self.lr(self.t)
        #         print('step',step)

        z = np.dot(self.W, x.T)
        y = np.zeros(shape(z))
        #         print('z',z)
        #         print('y',y)

        # The loop below solves the linear system My = Wx
        y[0] = z[0] / self.M[0, 0]
        #         print('y[0]',y[0])
        for i in np.arange(1, self.k):
            y[i] = z[i]
            #             print('i',i,'y[i]',y[i])
            for j in np.arange(0, i):
                y[i] = y[i] - y[j] * self.M[i, j]
            #                 print('i',i,'j',j,'y[i]',y[i])

            y[i] = y[i] / self.M[i, i]
        #             print('i',i,'y[i]',y[i])

        y_up = np.tril(
            np.dot(y.reshape(y.shape[0], 1), y.reshape(y.shape[0], 1).T))  # extract lower triangular component of y
        #         print('y_up',y_up)
        self.W = self.W + step * (y.reshape(y.shape[0], 1) * x.reshape(x.shape[0], 1).T - self.W)  # update W matrix
        #         print('W',self.W)
        self.M = self.M + step * (y_up - self.M)  # update M matrix
        #         print('M',self.M)
        self.t = self.t + 1;
        #         print('t',self.t)
        return y



