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, 
                           figsize=(20, 5 * np.ceil(len(sorting.unit_ids) / 4)))
    plt.savefig(waveform_folder + "/templates_map.png", dpi=300)

    extremum_channels_ids = st.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=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=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=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 = st.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 = st.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
