import sys, struct, math, os, time

sys.path.append("./autosort_neuron/")
import numpy as np
from tqdm import tqdm
from intanutil.read_header import read_header
from intanutil.get_bytes_per_data_block import get_bytes_per_data_block
from intanutil.read_one_data_block import read_one_data_block
from intanutil.notch_filter import notch_filter
from intanutil.data_to_result import data_to_result
from scipy import signal
import os
import spikeinterface

spikeinterface.__version__
import spikeinterface
import spikeinterface.extractors as se
import spikeinterface.sorters as ss
# import spikeinterface.toolkit as st
import spikeinterface.widgets as sw
import matplotlib.pyplot as plt
import numpy as np
import shutil
import pandas as pd
from spikeinterface.core.npzsortingextractor import NpzSortingExtractor
from pylab import *
import pylab
from probeinterface import Probe
from probeinterface.plotting import plot_probe
import sys
import warnings

warnings.filterwarnings("ignore")


def sorting_unit_show(we, recording_cmr, sorting, pack_folder, waveform_folder):
    # plot_probe(mesh_probe,with_channel_index=True)

    fig, axs = plt.subplots(
        int(np.ceil(len(sorting.unit_ids) / 4)),
        4,
        figsize=(20, 5 * np.ceil(len(sorting.unit_ids) / 4)),
    )
    sw.plot_unit_templates(we, unit_ids=sorting.unit_ids, axes=axs)
    plt.savefig(waveform_folder + "/templates_map.png", dpi=300)

    extremum_channels_ids = spikeinterface.core.get_template_extremum_channel(we, peak_sign="neg")

    colors = []
    cm = get_cmap("rainbow")
    NUM_COLORS = len(sorting.unit_ids)
    for i in range(NUM_COLORS):
        colors.append(cm(1.0 * i / NUM_COLORS))  # color will now be an RGBA tuple

    fig, axs = plt.subplots(
        int(np.ceil(len(sorting.unit_ids) / 4)),
        4,
        figsize=(20, 5 * np.ceil(len(sorting.unit_ids) / 4)),
    )

    for i, unit_id in enumerate(sorting.unit_ids):
        template = we.get_template(unit_id)
        ax = axs[int(np.floor(i / 4)), int(np.mod(i, 4))]
        ax.plot(
            template[:, extremum_channels_ids[unit_id]].T,
            lw=3,
            label=unit_id,
            color=colors[i],
        )
        ax.set_title(f"template{unit_id}")

    plt.savefig(waveform_folder + "/extremum_templates_map.png", dpi=300)


def read_data(filename):
    """Reads Intan Technologies RHD2000 data file generated by evaluation board GUI.

    Data are returned in a dictionary, for future extensibility.
    """

    tic = time.time()
    fid = open(filename, "rb")
    filesize = os.path.getsize(filename)

    header = read_header(fid)

    print(
        "Found {} amplifier channel{}.".format(
            header["num_amplifier_channels"], plural(header["num_amplifier_channels"])
        )
    )
    print(
        "Found {} auxiliary input channel{}.".format(
            header["num_aux_input_channels"], plural(header["num_aux_input_channels"])
        )
    )
    print(
        "Found {} supply voltage channel{}.".format(
            header["num_supply_voltage_channels"],
            plural(header["num_supply_voltage_channels"]),
        )
    )
    print(
        "Found {} board ADC channel{}.".format(
            header["num_board_adc_channels"], plural(header["num_board_adc_channels"])
        )
    )
    print(
        "Found {} board digital input channel{}.".format(
            header["num_board_dig_in_channels"],
            plural(header["num_board_dig_in_channels"]),
        )
    )
    print(
        "Found {} board digital output channel{}.".format(
            header["num_board_dig_out_channels"],
            plural(header["num_board_dig_out_channels"]),
        )
    )
    print(
        "Found {} temperature sensors channel{}.".format(
            header["num_temp_sensor_channels"],
            plural(header["num_temp_sensor_channels"]),
        )
    )
    print("")

    # Determine how many samples the data file contains.
    bytes_per_block = get_bytes_per_data_block(header)

    # How many data blocks remain in this file?
    data_present = False
    bytes_remaining = filesize - fid.tell()
    if bytes_remaining > 0:
        data_present = True

    if bytes_remaining % bytes_per_block != 0:
        raise Exception(
            "Something is wrong with file size : should have a whole number of data blocks"
        )

    num_data_blocks = int(bytes_remaining / bytes_per_block)

    num_amplifier_samples = header["num_samples_per_data_block"] * num_data_blocks
    num_aux_input_samples = int(
        (header["num_samples_per_data_block"] / 4) * num_data_blocks
    )
    num_supply_voltage_samples = 1 * num_data_blocks
    num_board_adc_samples = header["num_samples_per_data_block"] * num_data_blocks
    num_board_dig_in_samples = header["num_samples_per_data_block"] * num_data_blocks
    num_board_dig_out_samples = header["num_samples_per_data_block"] * num_data_blocks

    record_time = num_amplifier_samples / header["sample_rate"]

    if data_present:
        print(
            "File contains {:0.3f} seconds of data.  Amplifiers were sampled at {:0.2f} kS/s.".format(
                record_time, header["sample_rate"] / 1000
            )
        )
    else:
        print(
            "Header file contains no data.  Amplifiers were sampled at {:0.2f} kS/s.".format(
                header["sample_rate"] / 1000
            )
        )

    if data_present:
        # Pre-allocate memory for data.
        print("")
        print("Allocating memory for data...")

        data = {}
        if (header["version"]["major"] == 1 and header["version"]["minor"] >= 2) or (
            header["version"]["major"] > 1
        ):
            data["t_amplifier"] = np.zeros(num_amplifier_samples, dtype=np.int)
        else:
            data["t_amplifier"] = np.zeros(num_amplifier_samples, dtype=np.uint)

        data["amplifier_data"] = np.zeros(
            [header["num_amplifier_channels"], num_amplifier_samples], dtype=np.uint
        )
        data["aux_input_data"] = np.zeros(
            [header["num_aux_input_channels"], num_aux_input_samples], dtype=np.uint
        )
        data["supply_voltage_data"] = np.zeros(
            [header["num_supply_voltage_channels"], num_supply_voltage_samples],
            dtype=np.uint,
        )
        data["temp_sensor_data"] = np.zeros(
            [header["num_temp_sensor_channels"], num_supply_voltage_samples],
            dtype=np.uint,
        )
        data["board_adc_data"] = np.zeros(
            [header["num_board_adc_channels"], num_board_adc_samples], dtype=np.uint
        )

        # by default, this script interprets digital events (digital inputs and outputs) as booleans
        # if unsigned int values are preferred(0 for False, 1 for True), replace the 'dtype=np.bool' argument with 'dtype=np.uint' as shown
        # the commented line below illustrates this for digital input data; the same can be done for digital out

        # data['board_dig_in_data'] = np.zeros([header['num_board_dig_in_channels'], num_board_dig_in_samples], dtype=np.uint)
        data["board_dig_in_data"] = np.zeros(
            [header["num_board_dig_in_channels"], num_board_dig_in_samples],
            dtype=np.bool,
        )
        data["board_dig_in_raw"] = np.zeros(num_board_dig_in_samples, dtype=np.uint)

        data["board_dig_out_data"] = np.zeros(
            [header["num_board_dig_out_channels"], num_board_dig_out_samples],
            dtype=np.bool,
        )
        data["board_dig_out_raw"] = np.zeros(num_board_dig_out_samples, dtype=np.uint)

        # Read sampled data from file.
        print("Reading data from file...")

        # Initialize indices used in looping
        indices = {}
        indices["amplifier"] = 0
        indices["aux_input"] = 0
        indices["supply_voltage"] = 0
        indices["board_adc"] = 0
        indices["board_dig_in"] = 0
        indices["board_dig_out"] = 0

        print_increment = 10
        percent_done = print_increment
        for i in range(num_data_blocks):
            read_one_data_block(data, header, indices, fid)

            # Increment indices
            indices["amplifier"] += header["num_samples_per_data_block"]
            indices["aux_input"] += int(header["num_samples_per_data_block"] / 4)
            indices["supply_voltage"] += 1
            indices["board_adc"] += header["num_samples_per_data_block"]
            indices["board_dig_in"] += header["num_samples_per_data_block"]
            indices["board_dig_out"] += header["num_samples_per_data_block"]

            fraction_done = 100 * (1.0 * i / num_data_blocks)
            if fraction_done >= percent_done:
                print("{}% done...".format(percent_done))
                percent_done = percent_done + print_increment

        # Make sure we have read exactly the right amount of data.
        bytes_remaining = filesize - fid.tell()
        if bytes_remaining != 0:
            raise Exception("Error: End of file not reached.")

    # Close data file.
    fid.close()

    if data_present:
        print("Parsing data...")

        # Extract digital input channels to separate variables.
        for i in range(header["num_board_dig_in_channels"]):
            data["board_dig_in_data"][i, :] = np.not_equal(
                np.bitwise_and(
                    data["board_dig_in_raw"],
                    (1 << header["board_dig_in_channels"][i]["native_order"]),
                ),
                0,
            )

        # Extract digital output channels to separate variables.
        for i in range(header["num_board_dig_out_channels"]):
            data["board_dig_out_data"][i, :] = np.not_equal(
                np.bitwise_and(
                    data["board_dig_out_raw"],
                    (1 << header["board_dig_out_channels"][i]["native_order"]),
                ),
                0,
            )

        # Scale voltage levels appropriately.
        data["amplifier_data"] = np.multiply(
            0.195, (data["amplifier_data"].astype(np.int32) - 32768)
        )  # units = microvolts
        data["aux_input_data"] = np.multiply(
            37.4e-6, data["aux_input_data"]
        )  # units = volts
        data["supply_voltage_data"] = np.multiply(
            74.8e-6, data["supply_voltage_data"]
        )  # units = volts
        if header["eval_board_mode"] == 1:
            data["board_adc_data"] = np.multiply(
                152.59e-6, (data["board_adc_data"].astype(np.int32) - 32768)
            )  # units = volts
        elif header["eval_board_mode"] == 13:
            data["board_adc_data"] = np.multiply(
                312.5e-6, (data["board_adc_data"].astype(np.int32) - 32768)
            )  # units = volts
        else:
            data["board_adc_data"] = np.multiply(
                50.354e-6, data["board_adc_data"]
            )  # units = volts
        data["temp_sensor_data"] = np.multiply(
            0.01, data["temp_sensor_data"]
        )  # units = deg C

        # Check for gaps in timestamps.
        num_gaps = np.sum(
            np.not_equal(data["t_amplifier"][1:] - data["t_amplifier"][:-1], 1)
        )
        if num_gaps == 0:
            print("No missing timestamps in data.")
        else:
            print(
                "Warning: {0} gaps in timestamp data found.  Time scale will not be uniform!".format(
                    num_gaps
                )
            )

        # Scale time steps (units = seconds).
        data["t_amplifier"] = data["t_amplifier"] / header["sample_rate"]
        data["t_aux_input"] = data["t_amplifier"][range(0, len(data["t_amplifier"]), 4)]
        data["t_supply_voltage"] = data["t_amplifier"][
            range(0, len(data["t_amplifier"]), header["num_samples_per_data_block"])
        ]
        data["t_board_adc"] = data["t_amplifier"]
        data["t_dig"] = data["t_amplifier"]
        data["t_temp_sensor"] = data["t_supply_voltage"]

        # If the software notch filter was selected during the recording, apply the
        # same notch filter to amplifier data here.
        if header["notch_filter_frequency"] > 0 and header["version"]["major"] < 3:
            print("Applying notch filter...")

            print_increment = 10
            percent_done = print_increment
            for i in range(header["num_amplifier_channels"]):
                data["amplifier_data"][i, :] = notch_filter(
                    data["amplifier_data"][i, :],
                    header["sample_rate"],
                    header["notch_filter_frequency"],
                    10,
                )

                fraction_done = 100 * (i / header["num_amplifier_channels"])
                if fraction_done >= percent_done:
                    print("{}% done...".format(percent_done))
                    percent_done += print_increment
    else:
        data = []

    # Move variables to result struct.
    result = data_to_result(header, data, data_present)

    print("Done!  Elapsed time: {0:0.1f} seconds".format(time.time() - tic))
    return result


def plural(n):
    """Utility function to optionally pluralize words based on the value of n."""

    if n == 1:
        return ""
    else:
        return "s"


def sorting_day_split(
    sorting, date_id_all, day_length, pack_folder, sorting_save_name="firings_inlier"
):
    sampling_freq = sorting.get_sampling_frequency()

    fig, ax = plt.subplots(1, 1, figsize=(20, 40))
    sw.plot_rasters(sorting, time_range=(0, np.sum(day_length)), ax=ax)

    colors = []
    cm = pylab.get_cmap("rainbow")
    NUM_COLORS = len(day_length)
    for i in range(NUM_COLORS):
        colors.append(cm(1.0 * i / NUM_COLORS))  # color will now be an RGBA tuple

    for i in range(len(day_length)):
        ax.axvspan(
            np.sum(day_length[:i]) / sampling_freq,
            np.sum(day_length[: (i + 1)]) / sampling_freq,
            facecolor=colors[i],
            alpha=0.1,
        )
    plt.show()
    plt.savefig(pack_folder + "/sorting/" + sorting_save_name + ".png", dpi=300)

    for i in range(len(day_length)):
        pack_folder_i = pack_folder + "/" + date_id_all[i] + "/"

        if os.path.exists(pack_folder_i) == False:
            os.mkdir(pack_folder_i)

        start_frame = np.sum(day_length[:i])
        end_frame = np.sum(day_length[: (i + 1)])

        sub_sorting = sorting.frame_slice(start_frame, end_frame)

        keep_unit_ids = []
        for unit_id in sub_sorting.unit_ids:
            spike_train = sub_sorting.get_unit_spike_train(unit_id=unit_id)
            n = spike_train.size
            if n > 20:
                keep_unit_ids.append(unit_id)

        curated_sub_sorting = sub_sorting.select_units(
            unit_ids=keep_unit_ids, renamed_unit_ids=None
        )

        save_path = pack_folder_i + "/sorting/" + sorting_save_name + ".npz"
        if os.path.exists(pack_folder_i + "/sorting/") == False:
            os.mkdir(pack_folder_i + "/sorting/")
        NpzSortingExtractor.write_sorting(curated_sub_sorting, save_path)


def get_sorting_info(
    data_folder_all,
    pack_folder,
    waveform_save_folder,
    input_state="_merged",
    we_load_if_exists=True,
    freq_max=3000,
    freq_min=300,
):
    recording = spikeinterface.core.base.BaseExtractor.load_from_folder(data_folder_all)
    recording_f = spikeinterface.preprocessing.bandpass_filter(recording, freq_min=freq_min, freq_max=freq_max)
    recording_cmr = spikeinterface.preprocessing.common_reference(
        recording_f, reference="global", operator="average"
    )

    sorting = se.NpzSortingExtractor(
        pack_folder + "/sorting/firings" + input_state + ".npz"
    )

    we = spikeinterface.extract_waveforms(
        recording_cmr,
        sorting,
        waveform_save_folder,
        load_if_exists=we_load_if_exists,
        overwrite=False,
        ms_before=1,
        ms_after=2,
        max_spikes_per_unit=1000000,
        n_jobs=1,
        chunk_size=30000,
    )
    return recording_cmr, sorting, we


def recording_shank_split(
    recording,
    sorting,
    we,
    pack_folder,
    input_state="inlier",
    we_load_if_exists=True,
    probe_groups=None,
):

    extremum_channels_ids = spikeinterface.core.get_template_extremum_channel(we, peak_sign="neg")

    if probe_groups is None:
        probe_groups = np.arange(0, 30)

    NUmShanks = np.shape(np.unique(probe_groups))[0]

    recording.set_property("group", probe_groups, ids=None)
    slice_recording = recording.split_by()
    slice_sorting = []
    slice_we = []
    slice_unit_ids = []

    for shank_id in np.arange(NUmShanks):
        shank_channel_ids = np.where(probe_groups == shank_id)[0]
        shank_unit_ids = sorting.unit_ids[
            np.where(
                np.isin(
                    np.array(list(extremum_channels_ids.values())), shank_channel_ids
                )
            )[0]
        ]
        print(f"Shank:{shank_id+1}, units: {shank_unit_ids}")
        slice_unit_ids.append(shank_unit_ids)

        shank_recording = slice_recording[shank_id]
        shank_sorting = sorting.select_units(
            unit_ids=shank_unit_ids, renamed_unit_ids=None
        )
        slice_sorting.append(shank_sorting)

        shank_waveform_folder = (
            pack_folder + f"/waveforms_shank{shank_id+1}_{input_state}"
        )
        if we_load_if_exists == False:
            if os.path.exists(shank_waveform_folder):
                shutil.rmtree(shank_waveform_folder)

        shank_we = spikeinterface.extract_waveforms(
            shank_recording,
            shank_sorting,
            shank_waveform_folder,
            load_if_exists=we_load_if_exists,
            ms_before=1,
            ms_after=2.0,
            max_spikes_per_unit=10000,
            n_jobs=1,
            chunk_size=300,
        )

        slice_we.append(shank_we)

    return slice_recording, slice_sorting, slice_we


def shank_waveform_show(slice_sorting, slice_we, waveform_show=False, save_folder=None):
    NumShanks = len(slice_sorting)
    for shank_id in np.arange(NumShanks):
        shank_sorting = slice_sorting[shank_id]
        shank_we = slice_we[shank_id]

        # waveform
        if shank_sorting is not None:
            if len(shank_sorting.unit_ids) > 0:
                fig, axs = plt.subplots(
                    int(np.ceil(len(shank_sorting.unit_ids) / 8)),
                    8,
                    figsize=(20, 3 * np.ceil(len(shank_sorting.unit_ids) / 8)),
                )
                sw.plot_unit_templates(
                    shank_we, unit_ids=shank_sorting.unit_ids, axes=axs
                )
                fig.suptitle(f"Shank{shank_id+1}: templates")
                fig.tight_layout()
                if save_folder is not None:
                    plt.savefig(
                        save_folder + f"/template_map_shank{shank_id+1}.png", dpi=300
                    )

                if waveform_show == True:
                    fig, axs = plt.subplots(
                        int(np.ceil(len(shank_sorting.unit_ids) / 8)),
                        8,
                        figsize=(20, 3 * np.ceil(len(shank_sorting.unit_ids) / 8)),
                    )
                    sw.plot_unit_waveforms(
                        shank_we, unit_ids=shank_sorting.unit_ids, axes=axs
                    )
                    fig.suptitle(f"Shank{shank_id+1}: waveforms")
                    fig.tight_layout()
                    if save_folder is not None:
                        plt.savefig(
                            save_folder
                            + f"/extremum_waveforms_map_shank{shank_id+1}.png",
                            dpi=300,
                        )


def shank_ISI_show(slice_sorting, slice_we, save_folder=None):
    NumShanks = len(slice_sorting)

    unit_ids = [list(slice_sorting[shank_id].unit_ids) for shank_id in range(NumShanks)]
    NumUnits = np.sum([len(unit_ids[shank_id]) for shank_id in range(NumShanks)])
    Max_NumUnits = np.max([len(unit_ids[shank_id]) for shank_id in range(NumShanks)])

    # Set colors
    cm = pylab.get_cmap("rainbow")
    colors = []
    for shank_id in np.arange(NumShanks):
        unit_shank_ids = unit_ids[shank_id]
        colors_shank = []
        for j in range(len(unit_shank_ids)):
            idx = j * NumShanks + shank_id
            colors_shank.append(cm(1.0 * idx / (Max_NumUnits * NumShanks)))
        colors.append(colors_shank)

    # Plot ISI histograms
    for shank_id in np.arange(NumShanks):
        shank_sorting = slice_sorting[shank_id]
        shank_we = slice_we[shank_id]
        unit_shank_ids = unit_ids[shank_id]
        if int(np.ceil(len(unit_shank_ids) / 8)) > 0:
            fig, axs = plt.subplots(
                int(np.ceil(len(unit_shank_ids) / 8)),
                8,
                figsize=(20, 2.5 * np.ceil(len(unit_shank_ids) / 8)),
            )

            for i, unit_id in enumerate(unit_shank_ids):
                if int(np.ceil(len(unit_shank_ids) / 8)) == 1:
                    ax = axs[i]
                else:
                    ax = axs[int(np.floor(i / 8)), int(np.mod(i, 8))]

                ax.set_title(f"unit{unit_id}")
                sw.plot_isi_distribution(
                    shank_sorting, unit_ids=np.array(unit_id, ndmin=1), axes=ax
                )

            fig.suptitle(f"Shank{shank_id+1}: ISI Histograms")
            fig.tight_layout()

        if save_folder is not None:
            plt.savefig(save_folder + f"/ISI_histograms_shank{shank_id+1}.png", dpi=300)


def metrics_cal(slice_we, isi_threshold_ms=1.5):
    NumShanks = len(slice_we)

    metrics_all = []
    slice_curated_ids = [None] * NumShanks
    slice_curated_sorting = [None] * NumShanks
    slice_curated_we = [None] * NumShanks

    for shank_id in np.arange(NumShanks):
        shank_we = slice_we[shank_id]
        if shank_we is not None:

            pc = spikeinterface.postprocessing.compute_principal_components(
                shank_we, load_if_exists=True, n_components=3, mode="by_channel_local"
            )
            metrics = spikeinterface.qualitymetrics.compute_quality_metrics(
                shank_we,
                metric_names=[
                    "snr",
                    "isi_violation",
                    "nearest_neighbor",
                    "isolation_distance",
                    "l_ratio",
                    "firing_rate",
                ],
            )
            isi_violations_rate = compute_isi_violations(
                shank_we, isi_threshold_ms=isi_threshold_ms
            )

            metrics["shank"] = shank_id + 1
            metrics["unit_id"] = shank_we.sorting.unit_ids
            metrics["isi_violation"] = list(isi_violations_rate.values())
            metrics_all.append(metrics)

    metrics_all = pd.concat(metrics_all)

    return metrics_all


def compute_isi_violations(
    waveform_extractor, isi_threshold_ms=1.5, min_isi_ms=0, **kwargs
):
    recording = waveform_extractor.recording
    sorting = waveform_extractor.sorting
    unit_ids = sorting.unit_ids
    num_segs = sorting.get_num_segments()
    fs = recording.get_sampling_frequency()

    seg_durations = [recording.get_num_samples(i) / fs for i in range(num_segs)]
    total_duration = np.sum(seg_durations)

    isi_threshold_s = isi_threshold_ms / 1000
    min_isi_s = min_isi_ms / 1000
    isi_threshold_samples = int(isi_threshold_s * fs)

    isi_violations_rate = {}
    isi_violations_count = {}
    isi_violations_ratio = {}

    # all units converted to seconds
    for unit_id in unit_ids:
        num_violations = 0
        num_spikes = 0
        for segment_index in range(num_segs):
            spike_train = sorting.get_unit_spike_train(
                unit_id=unit_id, segment_index=segment_index
            )
            isis = np.diff(spike_train)
            num_spikes += len(spike_train)
            num_violations += np.sum(isis < isi_threshold_samples)

        isi_violations_rate[unit_id] = num_violations / num_spikes
    return isi_violations_rate


def sorting_shank_split(
    pack_folder=None,
    waveform_folder=None,
    waveform_save_folder=None,
    merge_unit_ids_pack=None,
    we_load_if_exists=False,
    output_state="merged",
    waveform_show=False,
    isi_threshold_ms=1.5,
):
    """After the sorting is done, this plots the template map, waveforms, templates, and autocorrelograms.
    Calculates metrics.
    """
    method_save_path = pack_folder
    curation_save_folder = method_save_path + f"/curation_result_{output_state}/"
    if os.path.exists(curation_save_folder) == False:
        os.mkdir(curation_save_folder)

    recording_cmr, sorting, we = get_sorting_info(
        data_folder_all,
        pack_folder,
        waveform_save_folder,
        input_state="_merged",
        we_load_if_exists=we_load_if_exists,
        freq_max=3000,
        freq_min=300,
    )

    # If applicable, merge units
    if merge_unit_ids_pack is not None:
        merged_we, merged_sorting = units_merge(
            recording_cmr,
            sorting,
            merge_unit_ids_pack,
            method_save_path,
            we_load_if_exists=we_load_if_exists,
        )
        sorting = merged_sorting
        we = merged_we

    # Plot template maps
    slice_recording, slice_sorting, slice_we = recording_shank_split(
        recording_cmr,
        sorting,
        we,
        method_save_path,
        input_state=output_state,
        we_load_if_exists=we_load_if_exists,
    )
    fig, axs = plt.subplots(
        int(np.ceil(len(sorting.unit_ids) / 5)),
        5,
        figsize=(20, 12.5 * np.ceil(len(sorting.unit_ids) / 5)),
    )
    sw.plot_unit_templates(we, unit_ids=sorting.unit_ids, axes=axs)
    plt.savefig(curation_save_folder + "/template_map.png", dpi=100)

    colors = []
    cm = pylab.get_cmap("rainbow")
    NUM_COLORS = len(sorting.unit_ids)
    for i in range(NUM_COLORS):
        colors.append(cm(1.0 * i / NUM_COLORS))  # color will now be an RGBA tuple

    # Plot extremum waveforms
    if waveform_show == True:
        extremum_channels_ids = spikeinterface.core.get_template_extremum_channel(we, peak_sign="neg")
        fig, axs = plt.subplots(
            int(np.ceil(len(sorting.unit_ids) / 5)),
            5,
            figsize=(20, 4 * np.ceil(len(sorting.unit_ids) / 5)),
        )
        for i, unit_id in enumerate(sorting.unit_ids):
            template = we.get_waveforms(unit_id)
            ax = axs[int(np.floor(i / 5)), int(np.mod(i, 5))]
            ax.plot(
                template[:, :, extremum_channels_ids[unit_id]].T,
                lw=0.3,
                label=unit_id,
                color=colors[i],
            )
            ax.set_title(f"template{unit_id}")
        plt.savefig(curation_save_folder + "/extremum_waveforms_map.png", dpi=100)

    shank_waveform_show(
        slice_sorting,
        slice_we,
        waveform_show=waveform_show,
        save_folder=curation_save_folder,
    )

    # Plot ISI histograms

    shank_ISI_show(slice_sorting, slice_we, save_folder=curation_save_folder)

    # Calculate metrics
    metrics_all = metrics_cal(slice_we, isi_threshold_ms=isi_threshold_ms)
    bad_mask = (
        (metrics_all["snr"] < 6)
        & (metrics_all["nn_hit_rate"] < 0.80)
        & (metrics_all["firing_rate"] < 0.10)
    )
    shank_qc_bad_unit_ids = bad_mask[bad_mask].index.values
    print(f"qc bad units: {shank_qc_bad_unit_ids}")
    metrics_all
    return (
        recording_cmr,
        sorting,
        we,
        slice_recording,
        slice_sorting,
        slice_we,
        metrics_all,
    )


def create_mesh_probe(positions, num_all_channels):
    mesh_probe = Probe(ndim=2, si_units="um")
    mesh_probe.set_contacts(
        positions=positions, shapes="circle", shape_params={"radius": 5}
    )

    ant = {"first_index": 0}
    mesh_probe.annotate(**ant)
    channel_indices_raw = np.arange(num_all_channels)
    channel_indices = [i for i in channel_indices_raw]
    mesh_probe.set_device_channel_indices(channel_indices)
    return mesh_probe


def stack_recordings(pack_folder_pre, mesh_probe, trigger_val=3.5, trigger=True):
    cont_data_all = []
    cont_trigger_all = []
    for dirpath, dirname, filenames in os.walk(pack_folder_pre):
        for i in sorted(filenames):
            if ".rhd" in i:
                print(dirpath + "/" + i)

                raw_data = read_data(dirpath + "/" + i)

                sampling_freq = raw_data["frequency_parameters"][
                    "amplifier_sample_rate"
                ]
                keep_list = np.arange(32)
                keep_list = np.delete(keep_list, 23)
                keep_list = np.delete(keep_list, 24)

                data = raw_data["amplifier_data"].T
                data_trigger = raw_data["board_dig_in_data"].flatten()
                if sampling_freq != 10000:
                    print("downsampling to 10000")
                    orig_sr = sampling_freq  # 10 kHz
                    new_sr = 10000  # 1 kH
                    resample_factor = orig_sr / new_sr

                    new_data = []
                    for j in tqdm(range(data.shape[1])):
                        new_data_j = signal.resample(
                            data[:, j], int(data.shape[0] / resample_factor)
                        )
                        new_data.append(new_data_j)
                    new_data = np.vstack(new_data)
                    new_data_trigger = data_trigger[:: int(resample_factor)]
                    new_data = new_data.T

                    data = new_data
                    data_trigger = new_data_trigger

                cont_data_all.append(data[:, list(keep_list)])
                if trigger:
                    cont_trigger_all.append(data_trigger)

    cont_data_all = np.vstack(cont_data_all)
    if trigger:
        cont_trigger_all = np.hstack(cont_trigger_all)

    recording = se.NumpyRecording(traces_list=cont_data_all, sampling_frequency=10000)

    recording.set_probe(mesh_probe, in_place=True)
    if trigger:
        return recording, cont_data_all, cont_trigger_all
    else:
        return recording, cont_data_all


def read_data_folder(
    data_folder_all, date_id_all, raw_data_path, mesh_probe, trigger=True
):

    recording_traces = []
    session_length_concat = []
    day_length = []
    cont_trigger_all_all = []
    if os.path.exists(data_folder_all):
        recording_concat = spikeinterface.core.base.BaseExtractor.load_from_folder(
            data_folder_all
        )
        session_length_concat = np.load(data_folder_all + "session_length.npy")
        day_length = np.load(data_folder_all + "day_length.npy")
        if trigger:
            cont_trigger_all_all = np.load(data_folder_all + "cont_trigger_all_all.npy")
    else:
        for date_id in date_id_all:
            pack_folder_pre = raw_data_path + date_id
            data_folder_pre = f"../processed_data/Ephys_{date_id}/"
            if os.path.exists(data_folder_pre):
                recording = spikeinterface.core.base.BaseExtractor.load_from_folder(
                    data_folder_pre
                )
                session_length = np.load(data_folder_pre + "session_length.npy")
                if trigger:
                    cont_trigger_all = np.load(data_folder_pre + "cont_trigger_all.npy")
            else:
                if trigger:
                    recording, session_length, cont_trigger_all = stack_recordings(
                        pack_folder_pre, mesh_probe, trigger_val=3
                    )
                else:
                    recording, session_length = stack_recordings(
                        pack_folder_pre, mesh_probe, trigger_val=3, trigger=trigger
                    )

                recording.set_probe(mesh_probe, in_place=True)

                recording = recording.save(folder=data_folder_pre)
                np.save(data_folder_pre + "session_length.npy", session_length)
                if trigger:
                    np.save(data_folder_pre + "cont_trigger_all.npy", cont_trigger_all)

            sampling_freq = recording.get_sampling_frequency()
            recording_trace = recording.get_traces()
            recording_traces.append(recording_trace)
            session_length_concat.append(session_length)
            if trigger:
                cont_trigger_all_all.append(cont_trigger_all)
            day_length.append(recording_trace.shape[0])
        recording_traces = np.vstack(recording_traces)
        session_length_concat = np.vstack(session_length_concat)
        recording_concat = se.NumpyRecording(
            traces_list=recording_traces, sampling_frequency=sampling_freq
        )
        recording_concat.set_probe(mesh_probe, in_place=True)
        recording_concat = recording_concat.save(folder=data_folder_all)
        np.save(data_folder_all + "session_length.npy", session_length_concat)
        np.save(data_folder_all + "day_length.npy", day_length)
        if trigger:
            cont_trigger_all_all = np.hstack(cont_trigger_all_all)
            np.save(data_folder_all + "cont_trigger_all_all.npy", cont_trigger_all_all)
    print(recording_concat)
    print("Num. channels = {}".format(len(recording_concat.get_channel_ids())))
    print(
        "Sampling frequency = {} Hz".format(recording_concat.get_sampling_frequency())
    )
    print("Num. timepoints seg0= {}".format(recording_concat.get_num_segments()))
    return recording_concat, day_length


def units_merge(
    recording_cmr,
    sorting,
    merge_unit_ids_pack,
    delete_unit_ids_pack,
    pack_folder,
    we_load_if_exists=True,
):
    S = sorting._sorting_segments[0]
    merged_sorting = sorting
    remove_ids = []

    for idx in range(len(merge_unit_ids_pack)):
        merge_unit_ids = merge_unit_ids_pack[idx]

        for unit_id_id, unit_id in enumerate(merge_unit_ids):
            S.spike_labels[S.spike_labels == unit_id] = merge_unit_ids[0]

        merged_sorting._sorting_segments[0] = S

        remove_ids.extend(merge_unit_ids[1:])

    remove_ids += delete_unit_ids_pack

    keep_ids = merged_sorting.unit_ids[~np.isin(merged_sorting.unit_ids, remove_ids)]

    merged_sorting = merged_sorting.select_units(
        unit_ids=keep_ids, renamed_unit_ids=None
    )

    waveform_save_folder = pack_folder + "/waveforms_merged"
    if we_load_if_exists == False:
        if os.path.exists(waveform_save_folder):
            shutil.rmtree(waveform_save_folder)
    save_sorting_merged_path = pack_folder + f"/sorting/firings_merged.npz"
    NpzSortingExtractor.write_sorting(merged_sorting, save_sorting_merged_path)
    merged_sorting = se.NpzSortingExtractor(save_sorting_merged_path)
    merged_we = spikeinterface.extract_waveforms(
        recording_cmr,
        merged_sorting,
        waveform_save_folder,
        load_if_exists=we_load_if_exists,
        overwrite=False,
        ms_before=1,
        ms_after=2,
        max_spikes_per_unit=1000000,
        n_jobs=1,
        chunk_size=30000,
    )
    sorting = merged_sorting
    we = merged_we
    return sorting, we
