"""Import functions for CED Signal data."""
import re
import os
import json
import h5py
import datetime
import warnings
import numpy as np
import pandas as pd
from scipy import io as spio
import pynibs


def cfs2hdf5(fn_cfs, fn_hdf5=None):
    """
    Converts EMG data included in .cfs file to .hdf5 format.

    Parameters
    ----------
    fn_cfs : str
        Filename of .cfs file.
    fn_hdf5 : str, optional
        Filename of .hdf5 file (if not provided, a file with same name as fn_cfs will be created with .hdf5 extension).

    Returns
    -------
    <file> : .hdf5 File
        File containing:

        * EMG data in f["emg"][:]
        * Time axis in f["time"][:]
    """
    try:
        import biosig
    except ImportError:
        ImportError("Please install biosig from pynibs/pkg/biosig folder!")

    if fn_hdf5 is None:
        fn_hdf5 = os.path.splitext(fn_cfs)[0] + ".hdf5"

    # load header and data
    cfs_header = biosig.header(fn_cfs)
    emg = biosig.data(fn_cfs)[:, 0]

    sweep_index = cfs_header.find('NumberOfSweeps')
    comma_index = cfs_header.find(',', sweep_index)
    sweeps = int(cfs_header[sweep_index + 18:comma_index])
    records = emg.shape[0]
    samples = int(records / sweeps)
    sampling_rate = pynibs.expio.get_mep_sampling_rate(fn_cfs)
    emg = np.reshape(emg, (sweeps, samples))
    time = np.linspace(0, samples, samples) / sampling_rate

    with h5py.File(fn_hdf5, "w") as f:
        f.create_dataset("emg", data=emg)
        f.create_dataset("time", data=time)
        f.create_dataset("sampling_rate", data=np.array([sampling_rate]))


def read_biosig_emg_data(fn_data, include_first_trigger=False, filetype="cfs"):
    """
    Reads EMG data from a biosig file.

    Parameters
    ----------
    fn_data : str
        Path to the biosig file.
    include_first_trigger : bool, default: False
        Whether to include the first trigger event in the data (default: False).
    filetype : str, default: 'cfs'
        Type of the biosig file.

    Returns
    -------
    emg_data : np.ndarray
        (num_sweeps, num_channels, samples_per_sweep) EMG data with shape.
    time_diff_list : list
        Time differences between trigger events in seconds.
    num_sweeps : int
        Number of sweeps in the EMG data.
    num_channels : int
        Number of channels in the EMG data.
    samples_per_sweep : int
        Number of samples per sweep in the EMG data.
    sampling_rate : int
        Sampling rate of the EMG data.
    """
    try:
        import biosig
    except ImportError:
        ImportError("Please install biosig from pynibs/pkg/biosig folder!")
        return

    if type == "cfs":  # TODO: also move the TXT reader here?
        cfs_fn = fn_data
        cfs_header = json.loads(biosig.header(cfs_fn))
        cfs_emg = biosig.data(cfs_fn)

        num_sweeps = cfs_header["NumberOfSweeps"]
        num_channels = cfs_emg.shape[1]

        total_num_samples = cfs_header["NumberOfSamples"]
        samples_per_sweep = int(total_num_samples / num_sweeps)
        sampling_rate = int(cfs_header["Samplingrate"])

        # get timestamps
        tms_pulse_timedelta = datetime.timedelta()
        # get hour, minute and second
        time_mep_list = []
        time_diff_list = []
        trigger_event_idcs = []
        if include_first_trigger:
            trigger_event_idcs.append(0)
            time_diff_list.append(0)
            time_mep_list.append(
                datetime.datetime.strptime(
                    cfs_header["EVENT"][0]["TimeStamp"], '%Y-%b-%d %H:%M:%S'
                )
                -
                datetime.timedelta(
                    seconds=float(cfs_header["EVENT"][0]["POS"])
                )
            )

        # convert time string into integer
        for event in cfs_header["EVENT"]:
            date = datetime.datetime.strptime(event["TimeStamp"], '%Y-%b-%d %H:%M:%S')

            # we are interested in the tms pulse time, so add it to ts
            date += tms_pulse_timedelta
            time_mep_list.append(date)
            time_diff_list.append((date - time_mep_list[0]).total_seconds())

            # compute indices in data block corresponding to the events
            if event["TYP"] == "0x7ffe":
                trigger_event_idcs.append(
                    round(event["POS"] * cfs_header["Samplingrate"])
                )

        num_sweeps = min(num_sweeps, len(trigger_event_idcs))

        emg_data = np.zeros((num_sweeps, num_channels, samples_per_sweep), dtype=np.float32)

        for c_idx in range(num_channels):
            # Use emg data starting from the index of the first trigger event
            # assumptions:
            # - after an initial offset all emg data were captured consecutively
            # - the first emg data frame may be captured without an explicit TMS
            #   tigger (e.g. by checking the "write to disk" option)
            # - if we had dropouts in between the emg data block (not just at the
            #   beginning) we would need to adhere to the entire trigger_event_indices
            #   list.
            emg_data[:, c_idx, :] = np.reshape(
                cfs_emg[trigger_event_idcs[0]:, c_idx],
                (num_sweeps, samples_per_sweep)
            )

        return emg_data, time_diff_list, num_sweeps, num_channels, samples_per_sweep, sampling_rate


def get_mep_elements(mep_fn, tms_pulse_time, drop_mep_idx=None, cfs_data_column=0, channels=None, time_format="delta",
                     plot=False, start_mep=18, end_mep=35):
    """
    Read EMG data from CED .cfs or .txt file and returns MEP amplitudes.

    Parameters
    ----------
    mep_fn : string
        path to .cfs-file or .txt file (Signal export).
    tms_pulse_time : float
        Time in [s] of TMS pulse as specified in signal.
    drop_mep_idx : List of int or None, optional
        Which MEPs to remove before matching.
    cfs_data_column : int or list of int, default: 0
        Column(s) of dataset in cfs file. +1 for .txt.
    channels : list of str, optional
        Channel names.
    time_format : str, default: "delta"
        Format of mep time stamps in time_mep_lst to return.

        * ``"delta"`` returns list of datetime.timedelta in seconds.
        * ``"hms"`` returns datetime.datetime(year, month, day, hour, minute, second, microsecond).

    plot : bool, default: False
        Plot MEPs.
    start_mep : float, default: 18
        Start of time frame after TMS pulse where p2p value is evaluated (in ms).
    end_mep : float, default: 35
        End of time frame after TMS pulse where p2p value is evaluated (in ms).

    Returns
    -------
    p2p_arr : np.ndarray of float
        (N_channel, N_stim) Peak to peak values of N sweeps.
    time_pulse_lst : list of datetime.timedelta
        TMS pulse 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).
    emg_time_axes : np.ndarray of float
        (N_samples) Time axis corresponding to EMG data.
    mep_onset_arr : np.ndarray of float
        (S_samples) MEP onset after TMS pulse.
    """
    # convert pulse time to datetime object in case of "delta"
    if time_format == "delta":
        tms_pulse_timedelta = datetime.timedelta(milliseconds=tms_pulse_time * 1000)
    elif time_format == "hms":
        tms_pulse_timedelta = datetime.timedelta()
    else:
        raise NotImplementedError("Specified time_format not implemented yet...")

    if mep_fn.endswith('.cfs'):
        # get data from cfs file
        import biosig
        mep_raw_data_tmp = biosig.data(mep_fn)
        mep_raw_data_tmp = mep_raw_data_tmp[:, cfs_data_column]  # get first channel

        # get header from cfs file
        cfs_header = biosig.header(mep_fn)

        # get timestamps
        # get all indices of timestamps from cfs header
        ts_mep_lst = [timestamp.start() for timestamp in re.finditer('TimeStamp', cfs_header)]
        # get hour, minute and second
        time_pulse_lst = []
        # convert time string into integer
        for index in ts_mep_lst:
            hour = int(cfs_header[index + 26:index + 28])
            minute = int(cfs_header[index + 29:index + 31])
            second = int(cfs_header[index + 32:index + 34])
            # fix bug with second 60
            if second == 60:
                ts = datetime.datetime(1900, 1, 1, hour, minute, 59)
                ts += datetime.timedelta(seconds=1)
            else:
                ts = datetime.datetime(1900, 1, 1, hour, minute, second)

            # we are interested in the tms pulse time, so add it to ts
            ts += tms_pulse_timedelta
            time_pulse_lst.append(ts)

        if time_format == "delta":
            time_pulse_lst = [time_pulse_lst[i] - time_pulse_lst[0] for i in range(len(time_pulse_lst))]
        if time_format == "hms":
            pass

        # add first timestamp (not saved by signal) and shift other by isi
        time_pulse_lst = [datetime.timedelta(seconds=0)] + [
            t + time_pulse_lst[1] - time_pulse_lst[0] for t in time_pulse_lst]

        # get peak-to-peak values
        # get the ratio of samples per sweep
        sweep_index = cfs_header.find('NumberOfSweeps')
        comma_index = cfs_header.find(',', sweep_index)
        n_sweeps = int(cfs_header[sweep_index + 18:comma_index])
        record_index = cfs_header.find('NumberOfRecords')
        comma_index = cfs_header.find(',', record_index)
        records = int(cfs_header[record_index + 19:comma_index])
        n_samples = int(records / n_sweeps)
        if not isinstance(n_samples, int):
            print('Warning: Number of samples is not an integer.')
            # TODO: Correct get_mep_elements() sample number check. This does not work as expected (from Ole)

        # reshape numpy array
        mep_raw_arr = np.zeros((len(cfs_data_column), n_sweeps, n_samples))

        for i in cfs_data_column:
            mep_raw_arr[i, :, :] = np.reshape(mep_raw_data_tmp[:, i], (n_sweeps, n_samples))

        sampling_rate = pynibs.expio.get_mep_sampling_rate(mep_fn)

    elif mep_fn.endswith('.mat'):
        mep_data = spio.loadmat(mep_fn, struct_as_record=False, squeeze_me=True)

        # find data
        for k in mep_data.keys():
            if isinstance(mep_data[k], spio.matlab.mio5_params.mat_struct):
                mep_data = mep_data[k].__dict__
                break

        n_samples = mep_data['points']
        mep_raw_arr = mep_data['values'].transpose(1, 2, 0)
        time_pulse_lst = [datetime.timedelta(seconds=f.__dict__['start']) for f in mep_data['frameinfo']]
        sampling_rate = pynibs.expio.get_mep_sampling_rate(mep_fn)

    elif mep_fn.endswith('.txt'):
        warnings.warn(".txt import is deprecated - use .mat or .cfs.", DeprecationWarning)
        print("Reading MEP from .txt file")
        # The Signal text export looks like this:
        #
        # "s"\t"ADC 0"\t"ADC 1"
        # 0.00000000\t-0.066681\t-0.047607
        # 0.00025000\t-0.066376\t-0.049286
        # 0.00050000\t-0.066528\t-0.056610
        #
        # "s"\t"ADC 0"\t"ADC 1"
        # 0.00000000\t-0.066681\t-0.047607
        # 0.00025000\t-0.066376\t-0.049286
        # 0.00050000\t-0.066528\t-0.056610
        #
        # With first column = time, second = 1st electrode, ...
        # This is an example of 2 sweeps, 3 samples each, sampling rate = 4000

        # Find number of samples per sweep
        pattern = '"s"'
        with open(mep_fn, 'r') as f:
            for line_nr, line in enumerate(f):
                print(f'{line_nr}: {line}')
                if pattern in line and line_nr > 0:
                    # find second occurance of "s" -> end of first sweep
                    n_samples = line_nr
                    print(f'{line_nr}: {line}')
                if line != '\n':
                    last_sample_time = line

        # extract time (first column) of last samples
        last_sample_time = float(last_sample_time[0:last_sample_time.find('\t')])

        # subtract 2 because first row is header ("s"\t"ADC 0"\t"ADC 1") and last row is blank
        n_samples = n_samples - 2

        df_mep = pd.read_csv(mep_fn,
                             delimiter="\t",
                             skip_blank_lines=True,
                             skiprows=lambda x: x % (n_samples + 2) == 0 and x > 0)

        n_sweeps = int(df_mep.shape[0] / n_samples)
        mep_raw_arr = np.zeros((len(cfs_data_column), n_sweeps, n_samples))

        for i in range(n_sweeps):
            mep_raw_arr[:, i, :] = df_mep.iloc[i * n_samples:(i + 1) * n_samples, 1:].transpose()

        # get sampling rate by dividing number of sweeps by timing
        sampling_rate = int(mep_raw_arr.shape[2] - 1) / last_sample_time

        # build time_mep_list
        # we only have information about the single mep timings, so let's assume signal sticks strictly to the protocol
        sample_len = last_sample_time + 1 / sampling_rate

        # TODO: The ISI is missing here, do we want to add it to the subject object?
        time_pulse_lst = [datetime.timedelta(seconds=i * sample_len) +
                          tms_pulse_timedelta for i in range(mep_raw_arr.shape[1])]

    else:
        raise ValueError("Unknown MEP file extension. Use .cfs or .txt.")

    # get peak to peak value of every sweep and plot results in mep/plots/channels
    if channels is None:
        channels = [str(i) for i in cfs_data_column]

    tmp = np.zeros((mep_raw_arr.shape[0], mep_raw_arr.shape[1], 3)).astype(object)
    for i_channel in range(mep_raw_arr.shape[0]):
        print(f"Calculating p2p values for channel: {channels[i_channel]}")

        for i_zap in range(mep_raw_arr.shape[1]):
            tmp[i_channel, i_zap, 0], \
                tmp[i_channel, i_zap, 1], \
                tmp[i_channel, i_zap, 2] = pynibs.expio.calc_p2p(sweep=mep_raw_arr[i_channel, i_zap, :],
                                                                 tms_pulse_time=tms_pulse_time,
                                                                 sampling_rate=sampling_rate,
                                                                 fn_plot=None,
                                                                 start_mep=start_mep,
                                                                 end_mep=end_mep)

    p2p_arr = np.zeros((tmp.shape[0], tmp.shape[1]))
    mep_onset_arr = np.zeros((tmp.shape[0], tmp.shape[1]))
    mep_filt_arr = np.zeros(mep_raw_arr.shape)

    emg_time_axes = np.arange(mep_raw_arr.shape[2]) / sampling_rate

    for idx_channel in cfs_data_column:
        for i, t in enumerate(tmp[idx_channel, :, :]):
            p2p_arr[idx_channel, i] = tmp[idx_channel, i, 0]
            mep_onset_arr[idx_channel, i] = tmp[idx_channel, i, 2]
            mep_filt_arr[idx_channel, i, :] = tmp[idx_channel, i, 1]

    if time_format == "delta":
        time_pulse_lst = [time_pulse_lst[i] - time_pulse_lst[0] for i in range(len(time_pulse_lst))]
    elif time_format == "hms":
        pass

    # remove MEPs according to drop_mep_idx and reset time
    if drop_mep_idx is not None:
        p2p_arr = np.delete(p2p_arr, drop_mep_idx)
        mep_onset_arr = np.delete(mep_onset_arr, drop_mep_idx)
        time_pulse_lst = np.delete(time_pulse_lst, drop_mep_idx)

    keep_mep_idx = [i for i in range(mep_raw_arr.shape[1]) if i not in np.array(drop_mep_idx)]
    mep_raw_arr = mep_raw_arr[:, keep_mep_idx, :]
    mep_filt_arr = mep_filt_arr[:, keep_mep_idx, :]

    return p2p_arr, time_pulse_lst, mep_raw_arr, mep_filt_arr, emg_time_axes, mep_onset_arr
