%matplotlib inline

from pathlib import Path
import os
import shutil
from pprint import pprint
import getpass


import numpy as np
import matplotlib.pyplot as plt

import MEArec as mr
import neo
import quantities as pq


import spikeinterface.extractors  as se
import spikeinterface.widgets  as sw
import spikeinterface.sorters  as ss

from spikeinterface.comparison import GroundTruthStudy

basedir = '/mnt/data/sam/DataSpikeSorting/'

basedir = Path(basedir)

workdir = basedir / 'mearec_bench_2021'

study_folder = workdir /'study_mearec_march_2021'

tmp_folder = workdir / 'tmp'
tmp_folder.mkdir(parents=True, exist_ok=True)

template_filename = workdir / 'templates_Neuronexus-32_100.h5'
probe = 'Neuronexus-32'
n_cell = 15
duration = 10 * 60.

recording_filename = workdir /  f'recordings_{n_cell}cells_{probe}_{duration:0.0f}s.h5'


fs = 30000.


#~ spgen = mr.SpikeTrainGenerator()
rec_params = mr.get_default_recordings_params()

rec_params['recordings']['fs'] = fs
rec_params['recordings']['sync_rate'] = None
rec_params['recordings']['sync_jitter'] = 5
rec_params['recordings']['noise_level'] = 5
rec_params['recordings']['filter'] = False
rec_params['recordings']['chunk_duration'] = 10.
rec_params['spiketrains']['duration'] = duration
rec_params['spiketrains']['n_exc'] = n_cell
rec_params['spiketrains']['n_inh'] = 0
rec_params['templates']['n_overlap_pairs'] = None
rec_params['templates']['min_dist'] = 0

recgen = mr.gen_recordings(params=rec_params, #spgen=spgen, 
            templates=template_filename, verbose=True,
            n_jobs=1, tmp_mode='memmap',
            tmp_folder=str(tmp_folder))

mr.save_recording_generator(recgen, filename=recording_filename)

user = getpass.getuser()

kilosort_path = f'/home/{user}/Documents/SpikeInterface/code_sorters/KiloSort1'
ss.KilosortSorter.set_kilosort_path(kilosort_path)

kilosort2_path = f'/home/{user}/Documents/SpikeInterface/code_sorters/Kilosort2'
ss.Kilosort2Sorter.set_kilosort2_path(kilosort2_path)

kilosort2_5_path = f'/home/{user}/Documents/SpikeInterface/code_sorters/Kilosort2.5'
ss.Kilosort2_5Sorter.set_kilosort2_5_path(kilosort2_path)

kilosort3_path = f'/home/{user}/Documents/SpikeInterface/code_sorters/Kilosort3'
ss.Kilosort3Sorter.set_kilosort3_path(kilosort3_path)

ironclust_path = f'/home/{user}/Documents/SpikeInterface/code_sorters/ironclust/'
ss.IronClustSorter.set_ironclust_path(ironclust_path)

mearec_filename = workdir / 'recordings_15cells_Neuronexus-32_600s.h5'

if study_folder.is_dir():
    shutil.rmtree(study_folder)

rec  = se.MEArecRecordingExtractor(mearec_filename)
sorting_gt = se.MEArecSortingExtractor(mearec_filename)
print(rec)
print(sorting_gt)

gt_dict = {'rec0' : (rec, sorting_gt) }

study = GroundTruthStudy.create(study_folder, gt_dict)

study = GroundTruthStudy(study_folder)
rec = study.get_recording()
probe = rec.get_probe()
print(probe)
from probeinterface.plotting import plot_probe
plot_probe(probe)

sorter_list = ['spykingcircus', 'kilosort2', 'kilosort3', 'tridesclous']
study = GroundTruthStudy(study_folder)
study.run_sorters(sorter_list, mode_if_folder_exists='overwrite', verbose=False)
study.copy_sortings()

study = GroundTruthStudy(study_folder)
study.copy_sortings()


study.run_comparisons(exhaustive_gt=True, delta_time=1.5)


comparisons = study.comparisons
dataframes = study.aggregate_dataframes()

for (rec_name, sorter_name), comp in comparisons.items():
    print()
    print('*'*20)
    print(rec_name, sorter_name)
    print(comp.count_score)

for (rec_name, sorter_name), comp in comparisons.items():
    fig, ax = plt.subplots()
    sw.plot_agreement_matrix(comp, ax=ax)
    fig.suptitle(rec_name+'   '+ sorter_name)

from pathlib import Path
import shutil
import time
import matplotlib.pyplot as plt

base_folder = Path('/mnt/data/sam/DataSpikeSorting/eduarda_arthur') 
data_folder = base_folder / 'raw_awake'

import spikeextractors as se
import spiketoolkit as st

print('spikeextractors version', se.__version__)
print('spiketoolkit version', st.__version__)

# step 1: open
file_path = data_folder / 'raw_awake_01_g0_t0.imec0.ap.bin'
recording = se.SpikeGLXRecordingExtractor(file_path)

# step 2: lazy filter
rec_filtered = st.preprocessing.bandpass_filter(recording,  freq_min=300. freq_max=6000.)
print(rec_filtered)

save_folder = base_folder / 'raw_awake_filtered_old'
if save_folder.is_dir():
    shutil.rmtree(save_folder)
save_folder.mkdir()

save_file = save_folder / 'filetred_recording.dat'
dump_file = save_folder / 'filetred_recording.json'

# step 3: cache
t0 = time.perf_counter()
cached = se.CacheRecordingExtractor(rec_filtered, chunk_mb=50, n_jobs=10, 
    save_path=save_file)
t1 = time.perf_counter()
run_time_filter_old = t1-t0
print('Old spikeextractors cache', run_time_filter_old)

# step : dump
cached.dump_to_json(dump_file)

 
import spikeinterface as si

import spikeinterface.extractors as se
import spikeinterface.toolkit as st
print('spikeinterface version', si.__version__)

# step 1: open
recording = se.SpikeGLXRecordingExtractor(data_folder)
print(recording)

# step 2: lazy filter
rec_filtered =st.bandpass_filter(recording,  freq_min=300., freq_max=6000.)
print(rec_filtered)


filter_path = base_folder / 'raw_awake_filtered'
if filter_path.is_dir():
    shutil.rmtree(filter_path)

# step 3 : compute and save with 10 workers
t0 = time.perf_counter()
cached = rec_filtered.save(folder=filter_path,
    format='binary', dtype='int16',
    n_jobs=10,  total_memory="50M", progress_bar=True)
t1 = time.perf_counter()
run_time_filter_new = t1 -t0
print('New spikeinterface filter + save binary', run_time_filter_new)

from spikeextractors.baseextractor import BaseExtractor
import spikeextractors as se
import spiketoolkit as st
print('spikeextractors version', se.__version__)
print('spiketoolkit version', st.__version__)

save_folder = base_folder / 'raw_awake_filtered_old'
dump_file = save_folder / 'filetred_recording.json'
recording = BaseExtractor.load_extractor_from_json(dump_file)

sorting_KS3 = se.KiloSortSortingExtractor(base_folder / 'output_kilosort3')
waveform_folder = base_folder / 'waveforms_extractor_old_1_job'
if waveform_folder.is_dir():
    shutil.rmtree(waveform_folder)
waveform_folder.mkdir()
sorting_KS3.set_tmp_folder(waveform_folder)

t0 = time.perf_counter()
wf, indexes, channels = st.postprocessing.get_unit_waveforms(recording, sorting_KS3,
            max_spikes_per_unit=500, return_idxs=True, chunk_mb=50, n_jobs=1,
            memmap=True)
t1 = time.perf_counter()
run_time_waveform_old_1jobs = t1 - t0
print('OLD API get_unit_waveforms 1 jobs', run_time_waveform_old_1jobs)

save_folder = base_folder / 'raw_awake_filtered_old'
dump_file = save_folder / 'filetred_recording.json'
recording = BaseExtractor.load_extractor_from_json(dump_file)

sorting_KS3_bis = se.KiloSortSortingExtractor(base_folder / 'output_kilosort3')
waveform_folder = base_folder / 'waveforms_extractor_old_10_jobs_'
if waveform_folder.is_dir():
    shutil.rmtree(waveform_folder)
waveform_folder.mkdir()
sorting_KS3_bis.set_tmp_folder(waveform_folder)

t0 = time.perf_counter()
wf, indexes, channels = st.postprocessing.get_unit_waveforms(recording, sorting_KS3_bis,
            max_spikes_per_unit=500, return_idxs=True, chunk_mb=500, n_jobs=10,
            memmap=True, verbose=True)
t1 = time.perf_counter()
run_time_waveform_old_10jobs = t1 - t0
print('OLD API get_unit_waveforms 10 jobs', run_time_waveform_old_10jobs)

import spikeinterface.extractors as se
from spikeinterface import WaveformExtractor, load_extractor
print('spikeinterface version', si.__version__)

filter_path = base_folder / 'raw_awake_filtered'
filered_recording = load_extractor(filter_path)

sorting_KS3 = se.KiloSortSortingExtractor(base_folder / 'output_kilosort3')
print(sorting_KS3)

# 1 worker
waveform_folder = base_folder / 'waveforms_extractor_1_job_new_'
if waveform_folder.is_dir():
    shutil.rmtree(waveform_folder)
we = WaveformExtractor.create(filered_recording, sorting_KS3, waveform_folder)

t0 = time.perf_counter()
we.set_params(ms_before=3., ms_after=4., max_spikes_per_unit=500)
we.run(n_jobs=1, total_memory="50M", progress_bar=True)
t1 = time.perf_counter()
run_time_waveform_new_1jobs = t1 - t0
print('New WaveformExtractor 1 jobs',run_time_waveform_new_1jobs)

# 1 worker
waveform_folder = base_folder / 'waveforms_extractor_10_job_new_'
if waveform_folder.is_dir():
    shutil.rmtree(waveform_folder)
we = WaveformExtractor.create(filered_recording, sorting_KS3, waveform_folder)

t0 = time.perf_counter()
we.set_params(ms_before=3., ms_after=4., max_spikes_per_unit=500)
we.run(n_jobs=10, total_memory="500M", progress_bar=True)
t1 = time.perf_counter()
run_time_waveform_new_10jobs = t1 - t0
print('New WaveformExtractor 10 jobs', run_time_waveform_new_10jobs)

speedup_filter = run_time_filter_old / run_time_filter_new
print('speedup filter', speedup_filter)

speedup_waveform_1jobs = run_time_waveform_old_1jobs / run_time_waveform_new_1jobs
print('speedup waveforms 1 jobs', speedup_waveform_1jobs)

speedup_waveform_10jobs = run_time_waveform_old_10jobs / run_time_waveform_new_10jobs
print('speedup waveformd 10jobs', speedup_waveform_10jobs)