"""
Set of functions that allow the t-sne algorithm to use the result of the kilosort automatic spike clustering algorithm.
The functions assume file names in the base_folder (the folder that kilosort puts its results in) as generated by
kilosort.

Some functions from this collection are also useful together with the GUI that allows cleaning of the kilosort templates
and the GUI that allows manual curation of the kilosort results using the t-sne embedding (e.g. generate_spike_info)
"""


import numpy as np
from os import path
import pickle
import pandas as pd
from . import io_with_cpp as io


def spike_indices_of_template(spike_templates, clean_templates=None, clean_template_index=None):
    """
    Returns the indices of the spikes belonging to a template from the group of clean templates (after manual cleaning
    of the kilosort results). The manual cleaning results are pulled from the template_markings.npy file saved the
    cleaning GUI in the same folder as the kilosort results

    :param spike_templates: The vector in the spike_templates.npy file of the kilosort results
    :type spike_templates: int[:]
    :param clean_templates: The vector of indices in the vector of the original templates of the templates that are
    clean. E.g. If originally templates [1,2,4,5,13,22,23] where found but templates 2 and 13 where noise then
    clean_templates = [0,2,3,4,6,7]

    :type clean_templates: int[:]
    :param clean_template_index: The index in the vector of the clean_templates of the template whose spikes we need
    :type clean_template_index: int
    :return: a vector of the indices of spikes belonging to the template
    :rtype: int[:]
    """
    if clean_templates is None:
        original_template_index = spike_templates
    else:
        original_template_index = clean_templates[clean_template_index]
    spikes_indices = np.squeeze(np.argwhere(np.in1d(spike_templates, original_template_index)))
    return spikes_indices


def find_templates_with_number_of_spikes_under_threshold(spike_templates, clean_templates, threshold):
    """
    Finds all the templates that have fewer number of spikes than a given threshold

    :param spike_templates: the original spike_templates as generated by kilosort
    :type spike_templates: int[:]
    :param clean_templates: the vector of which template each spike is assigned to but only for the clean templates
    :type clean_templates: int[:]
    :param threshold: the number of spikes allowed in a chosen template
    :type threshold: int
    :return: a dictionary with keys the templates with fewer than threshold spikes (small templates) and values their
    spike indices, the total number of spikes in these small templates, a dictionary with keys the remaining (large)
    templates and values the indices of their spikes

    :rtype: dict, int , dict
    """
    small_clean_templates_with_indices = dict()
    large_clean_templates_with_indices = dict()
    num_of_spikes_in_small_templates = 0
    for i in np.arange(len(clean_templates)):
        spike_indices = spike_indices_of_template(spike_templates, clean_templates, i)
        num_of_spikes = spike_indices.size
        if num_of_spikes <= threshold:
            small_clean_templates_with_indices[i] = spike_indices
            num_of_spikes_in_small_templates += num_of_spikes
        else:
            large_clean_templates_with_indices[i] = spike_indices

    return small_clean_templates_with_indices, num_of_spikes_in_small_templates, large_clean_templates_with_indices


def get_template_marking(base_folder):
    """
    Returns the template_marking vector (if a manual curation session has been performed on the kilosort results) which
    assigns to each template one of 7 possibilities. Since 0 means noise and all other possibilities are positive
    numbers this vector can act as a mask to exclude spikes of templates that have been marked as noise.
    If no template_markings.npy file is found then the function returns a vector of ones as long as the spike_templates
    vector

    :param base_folder: the folder with the kilosort results
    :type base_folder: string
    :return: vector of markings for each spike according to the type its template has been manually assigned to
    :rtype: int[:]
    """

    if path.isfile(path.join(base_folder, 'template_marking.npy')):
        template_marking = np.load(path.join(base_folder, 'template_marking.npy'))
    else:
        templates = np.load(path.join(base_folder, 'templates.npy'))
        num_of_templates = templates.shape[2]
        template_marking = np.ones(num_of_templates)
        del templates
        print('No template_marking.npy found. Using all templates.')

    return template_marking


def find_spike_indices_for_representative_tsne(base_folder, save_to_folder, threshold, total_spikes_required):
    """
    The function is used to create a subset of spikes using the following criteria. 1) The total number of spikes in the
    subset has to be smaller than the total_spikes_required value. 2) All spikes belonging to templates with threshold
    or fewer number of spikes should be included. If the threshold is too big and the total_spikes_required too small
    a number then these might not be mutually achievable criteria. If this is not the case then all the spikes belonging
    to templates with fewer than threshold spikes will be included. That will give a total of spikes smaller than the
    total_spikes_required value. The difference between the total_spikes_required and the total spikes in the small
    templates will be filled by spikes from the remaining large templates (the ones with higher number of spikes than
    threshold). The algorithm will calculate how many spikes these large templates have and use this to calculate a
    percentage of spikes for each large template that will bring the number of spikes in the subset equal to the
    total_spikes_required value. In this way the subset will be a full representation of the small templates (which
    will be the large majority of templates) and a random representation of part of the larger templates (which will be
    the small minority). Each large spike will contribute a specific percentage of its own total spikes.

    This technique can be used to create subsets of spikes whose total number of spikes does not exceed the t-sne
    capabilities. Use the indices_of_spikes_used returned vector to pass to the spikes_used_with_clean_indexing
    parameter of the calculate_template_features_matrix_for_tsne function.
    After running t-sne on this set any large templates that are still of interest (have not been
    delegated to noise and after any merging and splitting) can be run by themselves on a separate t-sne.

    :param base_fofder: the folder with the kilosort results
    :type base_folder: string
    :param save_to_folder: the function except returning the results will also save them as files to this folder. It
    also save dictionaries that make searching through the small and large templates easier

    :type save_to_folder: string
    :param threshold: the number of spikes that full included templated must be under
    :type threshold: int
    :param total_spikes_required: the total number of spikes in the resulting subset
    :type total_spikes_required: int
    :return: a vector with the indices of all the spikes in the subset,
    a dictionary with keys the small templates and values their spike indices,
    a dictionary with keys the large templates and values their spike indices

    :rtype: int[:], dict, dict
    """
    spike_templates = np.load(path.join(base_folder, 'spike_templates.npy'))
    template_marking = get_template_marking(base_folder)

    clean_templates = np.argwhere(template_marking)

    small_clean_templates_with_spike_indices, num_of_spikes_in_small_templates, large_clean_templates_with_spike_indices = \
        find_templates_with_number_of_spikes_under_threshold(spike_templates, clean_templates, threshold)

    extra_spikes_required = total_spikes_required - num_of_spikes_in_small_templates

    spikes_clean_index = np.squeeze(np.argwhere(np.in1d(spike_templates, clean_templates)))
    percentage_of_kept_spikes_in_large_templates = extra_spikes_required / (spikes_clean_index.size -
                                                                            num_of_spikes_in_small_templates)

    spikes_chosen = []
    num_of_spikes_in_large_templates = 0
    for large_template_index in large_clean_templates_with_spike_indices.keys():
        spike_indices = large_clean_templates_with_spike_indices[large_template_index]
        num_of_spikes = spike_indices.size
        num_of_spikes_in_large_templates += num_of_spikes
        chosen_num_of_spikes = int(num_of_spikes * percentage_of_kept_spikes_in_large_templates)
        chosen_spike_indices = np.random.choice(spike_indices, chosen_num_of_spikes, replace=False)
        spikes_chosen.append(chosen_spike_indices)

    num_of_spikes_in_small_templates = 0
    for small_template_index in small_clean_templates_with_spike_indices.keys():
        spike_indices = small_clean_templates_with_spike_indices[small_template_index]
        spikes_chosen.append(spike_indices)
        num_of_spikes_in_small_templates += spike_indices.size

    print('{} templates with more than {} spikes (total spikes in those = {}). \n{} templates with less '
          'than {} spikes (total spikes in those = {})'.format(len(large_clean_templates_with_spike_indices),
                                                               threshold, num_of_spikes_in_large_templates,
                                                               len(small_clean_templates_with_spike_indices),
                                                               threshold, num_of_spikes_in_small_templates))

    spikes_chosen_flat = []
    for sublist in spikes_chosen:
        if np.size(sublist) > 1:
            for item in sublist:
                spikes_chosen_flat.append(item)
        else:
            spikes_chosen_flat.append(sublist)

    indices_of_spikes_used = np.array(spikes_chosen_flat)
    small_clean_templates_indices = np.fromiter(small_clean_templates_with_spike_indices.keys(), int,
                                                len(small_clean_templates_with_spike_indices))
    large_clean_tempalates_indices = np.fromiter(large_clean_templates_with_spike_indices.keys(), int,
                                                 len(large_clean_templates_with_spike_indices))
    pickle.dump(small_clean_templates_with_spike_indices,
                open(path.join(save_to_folder, "small_clean_templates_with_spike_indices.pkl"), "wb"))
    pickle.dump(large_clean_templates_with_spike_indices,
                open(path.join(save_to_folder, "large_clean_templates_with_spike_indices.pkl"), "wb"))
    np.save(path.join(save_to_folder, 'indices_of_spikes_used'), indices_of_spikes_used)
    np.save(path.join(save_to_folder, 'indices_of_small_templates'), small_clean_templates_indices)
    np.save(path.join(save_to_folder, 'indices_of_large_templates'), large_clean_tempalates_indices)

    return indices_of_spikes_used, small_clean_templates_with_spike_indices, large_clean_templates_with_spike_indices


def calculate_template_features_matrix_for_tsne(base_folder, save_to_folder, spikes_used_with_original_indexing=None,
                                                spikes_used_with_clean_indexing=None):
    """
    Using the kilosort results, this function creates a matrix of samples x elements that can be used as an input to
    the t-sne algorithm. Each sample represents a spike and each element the distance of this spike to a specific
    template.

    :param base_folder:
    :type base_folder:
    :param save_to_folder:
    :type save_to_folder:
    :param spikes_used_with_original_indexing:
    :type spikes_used_with_original_indexing:
    :param spikes_used_with_clean_indexing:
    :type spikes_used_with_clean_indexing:
    :return:
    :rtype:
    """
    spike_templates = np.load(path.join(base_folder, 'spike_templates.npy'))
    template_features = np.load(path.join(base_folder, 'template_features.npy'))
    template_features_ind = np.load(path.join(base_folder, 'template_feature_ind.npy'))

    template_marking = get_template_marking(base_folder)

    clean_templates = np.argwhere(template_marking)
    spikes_clean_index = np.squeeze(np.argwhere(np.in1d(spike_templates, clean_templates)))

    if spikes_used_with_original_indexing is not None and spikes_used_with_clean_indexing is not None:
        print('Use one of the spikes_used_... variable')
        return None
    elif spikes_used_with_original_indexing is None and spikes_used_with_clean_indexing is None:
        pass
    elif spikes_used_with_original_indexing is None and spikes_used_with_clean_indexing is not None:
        spikes_clean_index = spikes_clean_index[spikes_used_with_clean_indexing]
    elif spikes_used_with_original_indexing is not None and spikes_used_with_clean_indexing is None:
        spikes_clean_index = spikes_used_with_original_indexing

    clean_templates = np.unique(spike_templates[spikes_clean_index])

    template_features_sparse_clean = np.zeros((spikes_clean_index.size, clean_templates.size))
    s = 0
    for spike in spikes_clean_index:
        cluster = spike_templates[spike][0]
        indices = template_features_ind[cluster, :]
        if s % 5000 == 0:
            print('Spikes completed: {}'.format(s))
        s+=1
        for i in np.arange(len(indices)):
            template_features_sparse_clean[np.argwhere(spikes_clean_index == spike),
                                           np.argwhere(clean_templates == indices[i])] = template_features[spike, i]

    np.save(path.join(save_to_folder, 'data_to_tsne_' + str(template_features_sparse_clean.shape) + '.npy'),
            template_features_sparse_clean)

    return template_features_sparse_clean


def load_template_features_matrix_for_tsne(save_to_folder, shape):
    """
    Loads the npy array that carries the sparse matrix of spikes x distances to templates

    :param save_to_folder: the folder the matrix has been saved to
    :type save_to_folder: string
    :param shape: the shape of the matrix
    :type shape: tuple
    :return: the sparse matrix
    :rtype: float32[shape]
    """
    template_features_sparse_clean = np.load(path.join(save_to_folder, r'data_to_tsne_' + str(shape) + '.npy'))
    return template_features_sparse_clean


def generate_spike_info(kilosort_folder, tsne_folder, tsne_filename='result.dat'):
    """
    spike_info.df is a pandas DataFrame that collates all the information about a group of spikes for sorting purposes.
    It is also the data frame that the GUI to manually curate kilosorted and t-sne data uses to load and save its
    results.

    :param kilosort_folder: the folder of the kilosort results
    :type kilosort_folder: string
    :param tsne_folder: the folder with the t-sne results
    :type tsne_folder: string
    :param tsne_filename: the filename of the t-sne results numpy array
    :type tsne_filename: string
    :return: saves and returns the spike_info dataframe
    :rtype: pandas.DataFrame
    """
    tsne = io.load_tsne_result(tsne_folder, tsne_filename)

    partial_tsne = False

    if path.isfile(path.join(tsne_folder, 'indices_of_spikes_used.npy')):
        spikes_used = np.load(path.join(tsne_folder, 'indices_of_spikes_used.npy'))
        partial_tsne = True
    else:
        spikes_used = np.arange(tsne.shape[0])

    if path.isfile(path.join(tsne_folder, 'indices_of_small_templates.npy')) and partial_tsne:
        indices_of_small_templates = np.load(path.join(tsne_folder, 'indices_of_small_templates.npy'))
    else:
        indices_of_small_templates = spikes_used

    if path.isfile(path.join(tsne_folder, 'weighted_template_positions.npy')):
        weighted_template_positions = np.load(path.join(tsne_folder, 'weighted_template_positions.npy'))
    else:
        weighted_template_positions = None

    if path.isfile(path.join(tsne_folder, 'weighted_spike_positions.npy')):
        weighted_spike_positions = np.load(path.join(tsne_folder, 'weighted_spike_positions.npy'))
    else:
        weighted_spike_positions = None

    if path.isfile(path.join(tsne_folder, 'template_marking.npy')):
        template_marking = np.load(path.join(kilosort_folder, 'template_marking.npy'))
    else:
        template_marking = np.ones(tsne.shape[0]) * 5 # Set it to Unspecified_1

    spike_templates = np.load(path.join(kilosort_folder, 'spike_templates.npy'))[spikes_used]
    spike_times = np.load(path.join(kilosort_folder, 'spike_times.npy'))[spikes_used]

    types = {0: 'Noise', 1: 'SS', 2: 'SS_Contaminated', 3: 'SS_Putative', 4: 'MUA', 5: 'Unspesified_1',
             6: 'Unspecified_2',
             7: 'Unspecified_3'}

    columns = ['original_index', 'times', 'template_after_cleaning', 'type_after_cleaning', 'template_after_sorting',
               'type_after_sorting', 'template_with_all_spikes_present', 'tsne_x', 'tsne_y', 'probe_position_x',
               'probe_position_z']

    spike_info = pd.DataFrame(index=np.arange(spikes_used.size), columns=columns)

    spike_info['original_index'] = spikes_used
    spike_info['times'] = spike_times
    spike_info['template_after_cleaning'] = spike_templates
    spike_info['type_after_cleaning'] = [types[int(template_marking[i])] for i in spike_templates]
    spike_info['template_after_sorting'] = spike_info['template_after_cleaning']
    spike_info['type_after_sorting'] = spike_info['type_after_cleaning']
    if partial_tsne:
        spike_info['template_with_all_spikes_present'] = [bool(np.in1d(spike_template, indices_of_small_templates))
                                                          for spike_template in spike_templates]
    else:
        spike_info['template_with_all_spikes_present'] = [True for spike_template in spike_templates]

    if weighted_spike_positions is not None:
        spike_info['probe_position_x'] = weighted_spike_positions[:, 0]
        spike_info['probe_position_y'] = weighted_spike_positions[:, 1]

    if weighted_spike_positions is None and weighted_template_positions is not None:
        spike_info['probe_position_x'] = [weighted_template_positions[spike_template, 0]
                                          for spike_template in spike_templates]
        spike_info['probe_position_y'] = [weighted_template_positions[spike_template, 1]
                                          for spike_template in spike_templates]

    spike_info['tsne_x'] = tsne[:, 0]
    spike_info['tsne_y'] = tsne[:, 1]

    spike_info.to_pickle(path.join(tsne_folder, 'spike_info.df'))

    return spike_info
