# -*- coding: utf-8 -*-
"""
2019-06-13
@author: Maxime Beau, Neural Computations Lab, University College London

Dataset merger class: represents a collection of Neuropixels Datasets (no theoretical upper limit on number of probes)
as a networkx network
with nodes being characterized as cell types
and edges as putative connections.

The way to use it is as follow:

# initiation
dp_dic = {'probe1': 'path/to/dataset1', 'probe2':'path/to/dataset2', ...}
merged = Merger(dp_dic)

# Connect the graph with putative monosynaptic connections
merged.connect_graph()

# Plot the graph
merged.plot_graph()

# Get N of putative connections of a given node spotted on the graph
merged.get_node_edges(node)

# Only keep a set of relevant nodes or edges and plot it again
merged.keep_edges(edges_list)
merged.keep_nodes_list(nodes_list)
merged.plot_graph()

# every graph operation of dataset merger can be performed on external networkx graphs
# provided with the argument 'src_graph'
g=merged.get_graph_copy(mergerGraph='undigraph')
merged.keep_nodes_list(nodes_list, src_graph=g) # g itself will be modified, not need to do g=...
merged.plot_graph(graph_src=g)


"""


import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

import os
import os.path as op; opj=op.join
from pathlib import Path

from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import numpy as np
import pandas as pd

from npyx.utils import npa, align_timeseries, assert_float

from npyx.io import get_npix_sync
from npyx.gl import get_units, load_units_qualities


def merge_datasets(datapaths):
    '''
    Merges datasets together and aligns data accross probes, modelling drift as a affine function.

    Any npyx function will run on a merged dataset just as on a regular dataset directory.

    If datasets have already been merged, will simply return the right datapath and dataset table.

    Parameters:
    - datapaths: dict of structure
        {'name_probe_1':'path/to/kilosort/output1',
         'name_probe_2':'path/to/kilosort/output2', ...}

    Returns:
        - dp_merged: str, path to merged dataset
        - ds_table: pandas dataframe, containing datasets names, paths, probe names
    '''

    # Handle datapaths format
    val_e=ValueError('''
    Datapath should be a dict of kilosort paths which exist:
    {'name_probe_1':'path/to/kilosort/output1', ...}''')
    if type(datapaths) is dict:
        for i, v in datapaths.items():
            datapaths[i]=Path(v)
            if not datapaths[i].exists():
                raise val_e
    else:
        raise val_e

    # Make an alphabetically indexed table of all fed datasets (if several)
    ds_table=pd.DataFrame(columns=['dataset_name', 'dp', 'probe'])
    for i, (prb, dp) in enumerate(datapaths.items()):
        ds_table.loc[i, 'dataset_name']=op.basename(dp)
        ds_table.loc[i, 'dp']=dp
        ds_table.loc[i, 'probe']=prb
    ds_table.insert(0, 'dataset_i', ds_table['dataset_name'].argsort()) # sort table by dataset_name alphabetically
    ds_table.sort_values('dataset_i', axis=0, inplace=True)
    ds_table.set_index('dataset_i', inplace=True)

    n_datasets=len(ds_table.index)
    mess_prefix="\033[34;1m--- "
    mess_suffix="\033[0m"

    # Instanciate datasets and define the merger path, where the 'merged' data will be saved,
    # together with the dataset table, holding information about the datasets (indices, names, source paths etc)
    ds_names=''
    predirname=op.dirname(ds_table['dp'][0])
    for ds_i in ds_table.index:
        dp, prb = ds_table.loc[ds_i, 'dp'], ds_table.loc[ds_i, 'probe']
        ds_names+='_'+ds_table.loc[ds_i, 'dataset_name']
        if predirname!=op.dirname(dp):
            print('WARNING: all your datasets are not stored in the same pre-directory - {} is picked anyway.'.format(predirname))
    dp_merged = Path(predirname) / ('merged'+ds_names) #Path(predirname, 'merger'+ds_names)
    old_prophyler_path = Path(predirname) / ('prophyler'+ds_names)
    if old_prophyler_path.exists():
        print("Backward compatibility - prophyler renamed according to merger convention.")
        os.rename(str(old_prophyler_path), str(dp_merged))
    if not op.isdir(dp_merged):
        print(f'\n{mess_prefix}Merged data (from {n_datasets} dataset(s)) will be saved here: {dp_merged}.{mess_suffix}')
        os.mkdir(dp_merged)
    else:
        print(f'\n{mess_prefix}Merged dataset found at {dp_merged}.{mess_suffix}')
        ds_table_old=pd.read_csv(Path(dp_merged, 'datasets_table.csv'), index_col='dataset_i')
        old_dataset_names = list(ds_table_old.loc[:, 'dataset_name'].values)
        for ds_i in ds_table.index:
            assert ds_table.loc[ds_i, 'dataset_name'] in old_dataset_names,\
            f'''WARNING you ran dataset merger on these {n_datasets} datasets in the past
                but used the dataset names {old_dataset_names}!!'''

    ds_table.to_csv(Path(dp_merged, 'datasets_table.csv'))

    # Load and save units qualities
    # + check whether datasets have been re-spike sorted if not first instanciation
    re_spksorted=False
    qualities=pd.DataFrame(columns=['cluster_id', 'group'])
    for ds_i in ds_table.index:
        cl_grp = load_units_qualities(ds_table.loc[ds_i, 'dp'])
        cl_grp['cluster_id']=cl_grp['cluster_id']+1e-1*ds_i
        qualities=qualities.append(cl_grp, ignore_index=True)
    qualities_dp=Path(dp_merged, 'cluster_group.tsv')
    if op.exists(qualities_dp):
        qualities_old=pd.read_csv(qualities_dp, sep='	')
        # only consider re-spike sorted if cluster indices have been changed, do not if only qualities were changed (spike times are unimpacted by that)
        if not np.all(np.isin(qualities_old.loc[:, 'cluster_id'], qualities.loc[:, 'cluster_id'])):
            re_spksorted=True
            print(f'\n{mess_prefix}New spike-sorting detected.{mess_suffix}')
    qualities.to_csv(qualities_dp, sep='	', index=False)

    # Merge spike times (or copy them if only 1 dataset)
    # Only if merged_clusters_times does not exist already or does but re-spikesorting has been detected
    merge_fname_times='spike_times'
    merge_fname_clusters='spike_clusters'
    if (not op.exists(Path(dp_merged, merge_fname_times+'.npy'))) or re_spksorted:
        print(f"\n{mess_prefix}Loading spike trains of {n_datasets} datasets...{mess_suffix}")
        spike_times, spike_clusters, sync_signals = [], [], []
        for ds_i in ds_table.index:
            dp=ds_table.loc[ds_i,'dp']
            spike_times.append(np.load(Path(dp, 'spike_times.npy')).flatten())
            spike_clusters.append(np.load(Path(dp, 'spike_clusters.npy')).flatten())
            if n_datasets>1: # sync signals only needed if alignment necessary
                ons, offs = get_npix_sync(dp, output_binary = False, sourcefile='ap', unit='samples')
                syncchan=ask_syncchan(ons)
                if syncchan=='q':
                    print('Aborted. Returning Nones.')
                    return None, None
                sync_signals.append(ons[syncchan])
        assert all(len(x) == len(sync_signals[0]) for x in sync_signals), 'WARNING different number of events on sync channels of both probes! Try again.'
        NspikesTotal=0
        for i in range(len(spike_times)): NspikesTotal+=len(spike_times[i])

        # If several datasets are fed to the merger, align their spike times.
        # Merged dataset is saved as two arrays:
        # spike_times and spike_clusters, as in a regular dataset,
        # but where spike_clusters is a larger int (1 is unit 1 from dataset 0, 1,000,000,001 is unit 1 from dataset 1)
        if n_datasets>1:
            print(f"\n{mess_prefix}Aligning spike trains of {n_datasets} datasets...{mess_suffix}")
            spike_times = align_timeseries(spike_times, sync_signals, 30000)
        # Now save the new array
        merged_spike_clusters=npa(zeros=(NspikesTotal), dtype=np.float64)
        merged_spike_times=npa(zeros=(NspikesTotal), dtype=np.uint64)
        cum_Nspikes=0
        for ds_i in ds_table.index:
            Nspikes=len(spike_times[ds_i])
            merged_spike_clusters[cum_Nspikes:cum_Nspikes+Nspikes]=spike_clusters[ds_i]+1e-1*ds_i
            merged_spike_times[cum_Nspikes:cum_Nspikes+Nspikes]=spike_times[ds_i]
            cum_Nspikes+=Nspikes
        merged_spike_clusters=merged_spike_clusters[np.argsort(merged_spike_times)]
        merged_spike_times=np.sort(merged_spike_times)
        np.save(dp_merged/(merge_fname_clusters+'.npy'), merged_spike_clusters)
        np.save(dp_merged/(merge_fname_times+'.npy'), merged_spike_times)
        sync_dir = dp_merged/'sync_chan'
        sync_dir.mkdir(exist_ok=True)
        sync_file = sync_dir/f"merged_ref_{ds_table.loc[0, 'dataset_name']}.ap_sync_on_samples.npy"
        np.save(sync_file, sync_signals[0])
        print(f'\n{mess_prefix}Merged {merge_fname_times} and {merge_fname_clusters} saved at {dp_merged}.{mess_suffix}')

        success_message = '\n--> Merge successful! Use a float u.x in any npyx function to call unit u from dataset x:'
        for ds_i in ds_table.index:
            success_message+=f"\n- u.{ds_i} for dataset {ds_table.loc[ds_i,'dataset_name']},"
        success_message='\033[92;1m'+success_message[:-1]+'.\033[0m'
        print(success_message)

    return dp_merged, ds_table

def ask_syncchan(ons):
    chan_len=''.join([f'chan {k} ({len(v)} events).\n' for k,v in ons.items()])
    syncchan=None
    while (syncchan is None):
        syncchan=input(f'Data found on sync channels:\n{chan_len}Which channel shall be used to synchronize probes? (q to restart and pick other channel on previous probe) >>> ')
        try:
            syncchan=int(syncchan)
            if syncchan not in ons.keys():
                print(f'!! You need to feed an integer amongst {list(ons.keys())}!')
                syncchan=None
        except:
            if syncchan=='q':
                print('Aborting...')
                return syncchan
            print('!! You need to feed an integer!')
            syncchan=None
        if syncchan is not None:
            if not any(ons[syncchan]):
                print('!! This sync channel does not have any events! Pick another one!')
                syncchan=None
    return syncchan


def get_ds_table(dp):
    ds_table = pd.read_csv(Path(dp, 'datasets_table.csv'), index_col='dataset_i')
    for dp in ds_table['dp']:
        assert op.exists(dp), \
            f"WARNING you have instanciated this merged dataset from paths of which at least one doesn't exist anymore:{dp}!"
    return ds_table

def get_dataset_id(u):
    '''
    Parameters:
        - u: float, of format u.ds_i
    Returns:
        - u: int, unit index
        - ds_i: int, dataset index
    '''
    assert assert_float(u), "Seems like the unit passed isn't a float - \
        calling this function on a merged dataset is meaningless \
        (cannot tell which dataset the unit belongs to!)."
    return int(round(u%1*10)), int(u)

def assert_same_dataset(U):
    '''Asserts if all provided units belong to the same dataset.
    '''
    return all(get_dataset_id(U[0])[0] == get_dataset_id(u)[0] for u in U[1:])

def assert_multi(dp):
    return 'merged' in op.basename(dp)

def get_ds_ids(U):
    return (U%1*10).round(0).astype(int)

def get_dataset_ids(dp_pro):
    '''
    Parameters:
        - dp_pro: str, path to merged dataset
    Returns:
        - dataset_ids: np array of shape (N_spikes,), indices of dataset of origin for all spikes
    '''
    assert assert_multi(dp_pro)
    return get_ds_ids(get_units(dp_pro))

def get_source_dp_u(dp, u):
    '''If dp is a merged datapath, returns datapath from source dataset and unit as integer.
       Else, returns dp and u as they are.
    '''
    if assert_multi(dp):
        ds_i, u = get_dataset_id(u)
        ds_table=get_ds_table(dp)
        dp=ds_table['dp'][ds_i]
    return dp, u