import numpy as np
import os
import pandas as pd
import warnings

import nelpy as nel

warnings.filterwarnings("ignore")

"""
Explanation: Score multi-session events using the replay score from Davidson et al.
End of explanation
"""


datadirs = ['/Users/ckemere/Development/Data/Buzsaki']

fileroot = next( (dir for dir in datadirs if os.path.isdir(dir)), None)
# conda install pandas=0.19.2
if fileroot is None:
    raise FileNotFoundError('datadir not found')

load_from_nel = True

# load from nel file:
if load_from_nel:
    jar = nel.load_pkl(os.path.join(fileroot,'gor01vvp01_processed_speed.nel'))
    exp_data = jar.exp_data
    aux_data = jar.aux_data
    del jar
    
    with pd.HDFStore(os.path.join(fileroot,'DibaMetadata.h5')) as store:
        df = store.get('Session_Metadata')
        df2 = store.get('Subset_Metadata')

"""
Explanation: Load experimental data
End of explanation
"""


# restrict sessions to explore to a smaller subset
min_n_placecells = 16
min_n_PBEs = 27 # 27 total events ==> minimum 21 events in training set

df2_subset = df2[(df2.n_PBEs >= min_n_PBEs) & (df2.n_placecells >= min_n_placecells)]

sessions = df2_subset['time'].values.tolist()
segments = df2_subset['segment'].values.tolist()

print('Evaluating subset of {} sessions'.format(len(sessions)))

df2_subset.sort_values(by=['n_PBEs', 'n_placecells'], ascending=[0,0])

"""
Explanation: Define subset of sessions to score
End of explanation
"""


n_jobs = 20 # set this equal to number of cores
n_shuffles = 100 # 5000
n_samples = 35000 # 35000
w=3 # single sided bandwidth (0 means only include bin who's center is under line, 3 means a total of 7 bins)

import matplotlib.pyplot as plt
%matplotlib inline

# Parallelize by EVENT
import dask
import distributed.joblib
from joblib import Parallel, delayed 
from joblib import parallel_backend

# A function that can be called to do work:
def work_events(arg):    

    # Split the list to individual variables:
    session, segment, ii, bst, tc = arg
    scores, shuffled_scores, percentiles = nel.analysis.replay.score_Davidson_final_bst_fast(bst=bst,
                                                                                        tuningcurve=tc,
                                                                                        w=w,
                                                                                        n_shuffles=n_shuffles,
                                                                                        n_samples=n_samples)

    return (session, segment, ii, scores, shuffled_scores, percentiles)

# List of instances to pass to work():
# unroll all events:
parallel_events = []
for session, segment in zip(sessions, segments):
    for nn in range(aux_data[session][segment]['PBEs'].n_epochs):
        parallel_events.append((session, segment, nn, aux_data[session][segment]['PBEs'][nn], 
                                aux_data[session][segment]['tc']))

#parallel_results = list(map(work_events, parallel_events))

with parallel_backend('dask.distributed', scheduler_host='35.184.42.12:8786'):
    # Anything returned by work() can be stored:
    parallel_results = Parallel(n_jobs=n_jobs, verbose=1)(map(delayed(work_events), parallel_events))




# standardize parallel results
bdries_ = [aux_data[session][segment]['PBEs'].n_epochs for session, segment in zip(sessions, segments) ]
bdries = np.cumsum(np.insert(bdries_,0,0))
bdries

sessions_ = np.array([result[0] for result in parallel_results])
segments_ = np.array([result[1] for result in parallel_results])
idx = [result[2] for result in parallel_results]

scores_bayes_evt = np.array([float(result[3]) for result in parallel_results])
scores_bayes_shuffled_evt = np.array([result[4].squeeze() for result in parallel_results])
scores_bayes_percentile_evt = np.array([float(result[5]) for result in parallel_results])

results = {}
for nn in range(len(bdries)-1):
    session = np.unique(sessions_[bdries[nn]:bdries[nn+1]])
    if len(session) > 1:
        raise ValueError("parallel results in different format / order than expected!")
    session = session[0]
    segment = np.unique(segments_[bdries[nn]:bdries[nn+1]])
    if len(segment) > 1:
        raise ValueError("parallel results in different format / order than expected!")
    segment = segment[0]
    try:
        results[session][segment]['scores_bayes'] = scores_bayes_evt[bdries[nn]:bdries[nn+1]]
    except KeyError:
        try:
            results[session][segment] = dict()
            results[session][segment]['scores_bayes'] = scores_bayes_evt[bdries[nn]:bdries[nn+1]]
        except KeyError:
            results[session] = dict()
            results[session][segment] = dict()
            results[session][segment]['scores_bayes'] = scores_bayes_evt[bdries[nn]:bdries[nn+1]]

    results[session][segment]['scores_bayes_shuffled'] = scores_bayes_shuffled_evt[bdries[nn]:bdries[nn+1]]
    results[session][segment]['scores_bayes_percentile'] = scores_bayes_percentile_evt[bdries[nn]:bdries[nn+1]]

print('done packing results')

"""
Explanation: Parallel scoring
NOTE: it is relatively easy (syntax-wise) to score each session as a parallel task, but since the Bayesian scoring takes such a long time to compute, we can be more efficient (higher % utilization) by further parallelizing over events, and not just over sessions. This further level of parallelization makes the bookkeeping a little ugly, so I provide the code for both approaches here.
End of explanation
"""


jar = nel.ResultsContainer(results=results, description='gor01 and vvp01 speed restricted results for best 20 candidate sessions')
jar.save_pkl('score_bayes_all_sessions.nel')

"""
Explanation: Save results to disk
End of explanation
"""


