import glob
import os
import pynibs
import datetime
import numpy as np


def get_mep_elements(session_path, chan_idx=None, recording_id=None,
                     tms_port_id=2, mep_start_ms=16, mep_end_ms=40, plot_folder=None, verbose=False):
    """
    This function calculates the MEP elements for a given session and channel for Neurone Tesla EEG data.

    Parameters
    ----------
    session_path : str
        Path to the session folder.
    mep_start_ms : int
        Start time of the MEP in ms.
    mep_end_ms : int
        End time of the MEP in ms.
    chan_idx : int, optional
        Channel index. Default: All channels.
    recording_id : int, optional
        Recording ID. Default: All recordings
    tms_port_id : int
        TMS port ID for the pulse trigger.
    verbose : bool
        Print verbose output.
    plot_folder : str, optional
        Folder to store the plots.

    Returns
    -------
    p2p_array : np.ndarray of float
        (N_stim) Peak to peak values of N sweeps.
    time_mep_lst : list of datetime.timedelta
        MEP-timestamps
    mep_raw_data : np.ndarray of float
        (N_channel, N_stim, N_samples) Raw (unfiltered) MEP data.
    mep_filt_data : np.ndarray of float
        (N_channel, N_stim, N_samples) Filtered MEP data (Butterworth lowpass filter).
    time : np.ndarray of float
        (N_samples) Time axis corresponding to MEP data.
    mep_onset_array : np.ndarray of float
        (S_samples) MEP onset after TMS pulse.
    """
    from neurone_tools import neurone_tools
    # read all recordings in the session
    if recording_id is None:
        subdirs = glob.glob(f"{session_path}/*")
        n_subdirs = len([os.path.split(subdir)[1] for subdir in subdirs if os.path.isdir(subdir)])
        if verbose:
            print(f"Found {n_subdirs} recordings in session {session_path}")
        recording_ids = list(range(1, n_subdirs+1))
    else:
        recording_ids = [recording_id]

    # read all channels in the recording
    nt = neurone_tools(session_path, recording_ids[0], channels='all')
    if chan_idx is None:
        chan_ids = range(len(nt.get_channel_info()))
        if verbose:
            print(f"Found {len(chan_ids)} channels in recording {recording_id}")
    else:
        chan_ids = [chan_idx]

    p2p_arr = []
    mep_raw_arr = []
    mep_filt_arr = []
    mep_onset_arr = []
    time_tms_pulse_list_chan = []
    time_axes_chan = []

    for chan_idx in chan_ids:
        p2p_arr_chan = []
        time_tms_pulse_list_chan = []
        mep_raw_arr_chan = []
        mep_filt_arr_chan = []
        time_axes_chan = []
        mep_onset_arr_chan = []

        for recording_id in recording_ids:
            nt = neurone_tools(session_path, recording_id, channels='all')
            fsample = nt.get_sampling_rate()

            if nt.get_channel_info()[chan_idx]['Unit'] == 'µV':
                fac = 1/1000
            else:
                raise ValueError(f"Unknown unit {nt.get_channel_info()[chan_idx]['Unit']}")

            # mep_start_sample = int(fsample * mep_start_ms/1000)
            mep_end_sample = int(fsample * mep_end_ms*2/1000)

            events = nt.load_events()
            tms_pulse_idx = events['start_sample_index'][events['source_port'] == tms_port_id].values
            data, chan_names = nt.load_data()
            if not len(tms_pulse_idx):
                print(f"No TMS pulses found in session {recording_id}")
                continue
            timedelta = datetime.timedelta(milliseconds=0)
            for pulse_id, pulse in enumerate(tms_pulse_idx):
                sweep = data[pulse:int(mep_end_sample + pulse), chan_idx]
                sweep *= fac
                sweep -= sweep.mean()

                if plot_folder is not None:
                    fn_plot = f"{plot_folder}/ses{recording_id:0>2}_chan{chan_idx}_{pulse_id:0>3}_mep.png"
                else:
                    fn_plot = None
                p2p, sweep_filt, onset = pynibs.expio.calc_p2p(sweep, tms_pulse_time=0,
                                                               start_mep=mep_start_ms,
                                                               end_mep=mep_end_ms,
                                                               measurement_start_time=0,
                                                               sampling_rate=fsample,
                                                               cutoff_freq=500,
                                                               fn_plot=fn_plot)
                p2p_arr_chan.append(p2p)

                time_tms_pulse_list_chan.append(timedelta)
                timedelta += datetime.timedelta(milliseconds=1000 * (pulse / fsample))
                mep_filt_arr_chan.append(sweep_filt)
                mep_onset_arr_chan.append(onset)
                mep_raw_arr_chan.append(sweep)
                time_axes_chan.append(np.arange(sweep.shape[0]) / fsample)

            print(f"channel {chan_idx} session {recording_id: >2}: {(events['source_port'] == tms_port_id).sum()}")
        mep_raw_arr_chan = np.vstack(mep_raw_arr_chan)
        mep_filt_arr_chan = np.vstack(mep_filt_arr_chan)

        p2p_arr.append(p2p_arr_chan)
        mep_filt_arr.append(mep_filt_arr_chan)
        mep_onset_arr.append(mep_onset_arr_chan)
        mep_raw_arr.append(mep_raw_arr_chan)
    p2p_arr = np.vstack(p2p_arr)
    mep_raw_arr = np.stack(mep_raw_arr, 0)
    mep_filt_arr = np.stack(mep_filt_arr, 0)
    mep_onset_arr = np.vstack(mep_onset_arr)

    return p2p_arr, time_tms_pulse_list_chan, mep_raw_arr, mep_filt_arr, time_axes_chan, mep_onset_arr
