import itertools

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from mpl_toolkits.mplot3d import Axes3D
from sklearn.manifold import Isomap, MDS, TSNE, LocallyLinearEmbedding
from sklearn.decomposition import PCA, FactorAnalysis, FastICA
from one.api import ONE

from iblutil.numerical import bincount2D


def find_nearest(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return idx


def filliti(v):
    """
    Sets all time bins within a trial to the same choice value v
    """
    for x in range(len(v[0])):
        if v[0, x] == 0:
            v[0, x] = v[0, x - 1]
    return v


def bin_types(spikes, trials, t_bin, clusters):
    """
    This creates a dictionary of binned time series,
    all having the same number of observations.

    INPUT:

    spikes: alf.io.load_object(alf_path, 'spikes')
    trials: alf.io.load_object(alf_path, '_ibl_trials')
    t_bin: float, time bin in sec
    clusters: alf.io.load_object(alf_path, 'clusters')

    OUTPUT:

    binned_data['summed_spike_amps']: channels x observations
    binned_data['reward']: observations
    binned_data['choice']: observations
    binned_data['trial_number']: observations

    """

    # TO GET MEAN: bincount2D(..., weight=positions) / bincount2D(...,
    # weight=None)

    # Bin spikes (summing together, i.e. not averaging per bin)
    R1, times1, _ = bincount2D(
        spikes['times'], spikes['clusters'], t_bin, weights=spikes['amps'])

    # Get choice per bin
    R6, times6, _ = bincount2D(trials['goCue_times'], trials['choice'], t_bin)
    # Flatten choice -1 Left, 1  Right
    R6 = np.sum(R6 * np.array([[-1], [1]]), axis=0)
    R6 = np.expand_dims(R6, axis=0)
    # Fill 0 between trials with choice outcome of trial
    R6 = filliti(R6)
    R6[R6 == -1] = 0
    R6 = R6[0]

    # Get reward per bin
    R7, times7, _ = bincount2D(
        trials['goCue_times'], trials['feedbackType'], t_bin)
    # Flatten reward -1 error, 1  reward
    R7 = np.sum(R7 * np.array([[-1], [1]]), axis=0)
    R7 = np.expand_dims(R7, axis=0)
    # Fill 0 between trials with reward outcome of trial
    R7 = filliti(R7)
    R7[R7 == -1] = 0
    R7 = R7[0]

    # restrict each time series to the same time bins
    start = max([x for x in [times1[0], times6[0], times7[0]]])
    stop = min([x for x in [times1[-1], times6[-1], times7[-1]]])

    time_points = np.linspace(start, stop, int((stop - start) / t_bin))

    binned_data = {}
    binned_data['summed_spike_amps'] = R1[:, find_nearest(
        times1, start):find_nearest(times1, stop)]
    binned_data['choice'] = R6[find_nearest(
        times6, start):find_nearest(times6, stop)]
    binned_data['reward'] = R7[find_nearest(
        times7, start):find_nearest(times7, stop)]
    binned_data['trial_number'] = np.digitize(
        time_points, trials['goCue_times'])

    # check lengths again for potential jumps
    chans, obs = binned_data['summed_spike_amps'].shape
    l_choice = len(binned_data['choice'])
    l_reward = len(binned_data['reward'])
    l_trial = len(binned_data['trial_number'])

    MIN = min([obs, l_choice, l_reward, l_trial])

    w = binned_data['summed_spike_amps'][:, :MIN]
    binned_data['summed_spike_amps'] = w
    binned_data['reward'] = binned_data['reward'][:MIN]
    binned_data['choice'] = binned_data['choice'][:MIN]
    binned_data['trial_number'] = binned_data['trial_number'][:MIN]

    print('Range of trials: ',
          [binned_data['trial_number'][0],
           binned_data['trial_number'][-1]])

    return binned_data


def color_3D_projection(
        data_projection,
        variable_data,
        title,
        color_map='jet'):
    """
    Plot a 3d scatter plot, each point being neural activity at a certain
    time bin, colored by the corresponding behavioral variable
    """
    x, y, z = np.split(data_projection, 3, axis=1)
    fig = plt.figure(title[:3])
    ax = Axes3D(fig)
    p = ax.scatter(x, y, z, s=20, alpha=0.25, c=variable_data, cmap=color_map)
    fig.colorbar(p)
    ax.set_title(title, fontsize=18)
    plt.show()
    return ax


if __name__ == "__main__":
    """
    Starting from an alf folder, neural IBL data is read in,
    restricted to specific trials, restricted to a specific brain region.
    These many channels of neural activity are dimensionally
    reduced to 3 and plotted as a scatter plot, colored either
    by "choice" or "reward". Several dimensionalty-reduction
    methods are used and the scatter plots displayed.
    """
    one = ONE()

    # read in the alf objects [brain region info only for code camp data now]
    session_path = 'ZM_1735/2019-08-01/001'
    probe_label = 'probe00'
    # can be addressed as spikes['time'] or spikes.time
    spikes = one.load_object(session_path, 'spikes', collection=f'alf{probe_label}')
    clusters = one.load_object(session_path, 'clusters', collection=f'alf{probe_label}')
    channels = one.load_object(session_path, 'channels', collection=f'alf{probe_label}')
    trials = one.load_object(session_path, 'trials', collection='alf')

    # Print number of clusters for each brain region
    locDict_bothProbes = clusters['brainAcronyms']['brainAcronyms'].to_dict()
    cluster_idx_probe1 = np.unique(spikes['clusters'])
    locDict = {}
    for i in locDict_bothProbes:
        if i in cluster_idx_probe1:
            locDict[i] = locDict_bothProbes[i]
    print([(k, len(list(v))) for k, v in itertools.groupby(sorted(locDict.values()))])

    # set key parameters
    T_BIN = 0.1
    TRIALS_TO_PLOT = np.arange(20, 23)  # use the real trial numbers
    PROJECTED_DIMENSIONS = 3
    BEHAVIORAL_VARIABLE = 'choice'
    BRAIN_AREA = 'MB'  # that depends on the dataset

    # Reduce neural data to region of interest
    if BRAIN_AREA:
        locations = clusters['brainAcronyms']
        spikes = pd.DataFrame.from_dict(spikes)
        loc_idx = locations.loc[(
            locations['brainAcronyms'] == BRAIN_AREA)].index
        spikes = spikes[np.isin(spikes['clusters'], loc_idx)]

    # Bin neural data and align to behavioral data
    binned_data = bin_types(spikes, trials, T_BIN, clusters)

    # Temporally restrict neural and behavioral time series
    # to sevral requested consecutive trials.
    bin_index = np.isin(binned_data['trial_number'], TRIALS_TO_PLOT)
    neural_data = binned_data['summed_spike_amps'][:, bin_index].T
    variable_data = binned_data[BEHAVIORAL_VARIABLE][bin_index]

    # Reduce neural dimensions and plot colored by behavioral_variable
    Title = "ZM_1735/2019-08-01, \n area %s; behav.  %s" % (
        BRAIN_AREA, BEHAVIORAL_VARIABLE)

    # PCA

    SVD_SOLVER = 'full'
    pca_projected_data = PCA(
        n_components=PROJECTED_DIMENSIONS,
        svd_solver=SVD_SOLVER).fit_transform(neural_data)
    color_3D_projection(pca_projected_data, variable_data, 'PCA; ' + Title)

    # Factor analysis

    fa_projected_data = FactorAnalysis(
        n_components=PROJECTED_DIMENSIONS).fit_transform(neural_data)
    color_3D_projection(fa_projected_data, variable_data, 'FA; ' + Title)

    # ICA

    ICA_projected_data = FastICA(
        n_components=PROJECTED_DIMENSIONS).fit_transform(neural_data)
    color_3D_projection(ICA_projected_data, variable_data, 'ICA; ' + Title)

    # Isomap

    N_NEIGHBORS = 30
    Isomap_projected_data = Isomap(
        n_components=PROJECTED_DIMENSIONS,
        n_neighbors=N_NEIGHBORS).fit_transform(neural_data)
    color_3D_projection(
        Isomap_projected_data,
        variable_data,
        'Isomap; ' + Title)

    # tSNE

    PERPLEXITY = 30  # normally ranges 5-50
    TSNE_projected_data = TSNE(
        n_components=PROJECTED_DIMENSIONS,
        perplexity=PERPLEXITY).fit_transform(neural_data)
    color_3D_projection(TSNE_projected_data, variable_data, 'tSNE; ' + Title)

    # Multidimensional scaling

    MDS_projected_data = MDS(
        n_components=PROJECTED_DIMENSIONS).fit_transform(neural_data)
    color_3D_projection(MDS_projected_data, variable_data, 'MS; ' + Title)

    # Locally Linear Embedding

    N_NEIGHBORS = 30
    LLE_projected_data = LocallyLinearEmbedding(
        n_components=PROJECTED_DIMENSIONS,
        n_neighbors=N_NEIGHBORS).fit_transform(neural_data)
    color_3D_projection(LLE_projected_data, variable_data, 'LLE; ' + Title)
