import os
import time
import h5py
import inspect
import warnings
import multiprocessing
import pandas as pd
import numpy as np
import scipy.stats
from lmfit import Model
from scipy.linalg import svd
from functools import partial
from scipy.stats import linregress
from collections import OrderedDict
from numpy.linalg import lstsq, pinv
from sklearn.linear_model import LinearRegression
import pynibs


class Element(object):
    """ 
    Fit Element object class
    """
    def __init__(self, x, y, ele_id, fun=pynibs.expio.fit_funs.sigmoid4, score_type="R2", select_signed_data=False, constants=None,
                 **kwargs):
        """
        Initializes Fit Element instance
        """
        self.x = x
        self.y = y
        self.y_passed = self.y
        self.fun = fun
        self.ele_id = ele_id
        self.init_vals = dict()
        self.random_vals_init_range = dict()
        self.limits = dict()
        self.status = True
        self.log_scale = False
        self.select_signed_data = select_signed_data
        self.score_type = score_type
        self.score = None
        self.best_values = None
        self.r2_lin = None
        self.gmodel = None
        self.fit = None
        self.residual = None
        self.var_y = np.var(self.y)
        self.var_y_passed = self.var_y

        self.norm_y = np.linalg.norm(self.y)
        self.param_names = None

        # Add additional parameters as object fields
        self.__dict__.update(kwargs)

        if constants is None:
            self.constants = {}

        # select whether the fit is performed for positive or negative data by running an initial linear fit
        # select the data, which yields a fit with a lower p-value; removes unused x and y data
        if self.select_signed_data:
            if not ((self.x < 0).all() or (self.x > 0).all()):
                self.run_select_signed_data()

        # set initial values of
        self.setup_model()

    def set_init_vals(self, value):
        """ Sets initial values in self.init_vals and gmodel instance """
        for key in value.keys():
            self.init_vals[key] = value[key]
            self.gmodel.param_hints[key]['value'] = value[key]

        self.gmodel.make_params()

    def set_limits(self, value):
        """ 
        Sets limits in self.limits and gmodel instance

        Parameters
        ----------
        value : dict
            Parameters (keys) to set in self as limits.
        """
        for key in value.keys():
            self.limits[key] = value[key]
            self.gmodel.param_hints[key]['min'], self.gmodel.param_hints[key]['max'] = value[key]

        self.gmodel.make_params()

    def set_constants(self, value):
        """ Sets constants in self.constants and gmodel instance """
        for key in value.keys():
            self.constants[key] = value[key]
            self.gmodel.param_hints[key]['value'] = value[key]
            self.gmodel.param_hints[key]['vary'] = False

        self.gmodel.make_params()

    def setup_model(self):
        """ Setup model parameters (limits, initial values, etc. ...) """
        x_min = np.min(self.x)
        x_max = np.max(self.x)
        y_min = np.min(self.y)
        y_max = np.max(self.y)

        # set up gmodel
        self.gmodel = Model(self.fun)
        self.param_names = self.gmodel.param_names

        # create the param_hints OrderedDict
        for p in self.gmodel.param_names:
            self.gmodel.param_hints[p] = OrderedDict()

        # if values are not already given by the configuration file, they are set here

        if self.fun == pynibs.expio.fit_funs.linear:
            # linear function starts with generic and same values for each element
            if self.limits == {}:
                self.set_limits({"m": [-100, 100], "n": [-100, 100]})
            else:
                self.set_limits(self.limits)
            if self.init_vals == {}:
                self.set_init_vals({"m": 0.3, "n": -1})
            else:
                self.set_init_vals(self.init_vals)
            if self.random_vals_init_range == {}:
                self.random_vals_init_range = {"m": [0, 100], "n": [0, .3]}

        elif self.fun == pynibs.expio.fit_funs.exp0:
            if self.limits == {}:
                self.set_limits({"x0": [0, 1000], "r": [1e-12, 100]})
            else:
                self.set_limits(self.limits)
            if self.init_vals == {}:
                self.set_init_vals({"x0": 10, "r": .1})
            else:
                self.set_init_vals(self.init_vals)
            if self.random_vals_init_range == {}:
                self.random_vals_init_range = {"x0": [0, 10], "r": [0, .2]}

        elif self.fun in [pynibs.expio.fit_funs.sigmoid, pynibs.expio.fit_funs.sigmoid_log,
                          pynibs.expio.fit_funs.sigmoid4, pynibs.expio.fit_funs.sigmoid4_log]:
            if self.fun in [pynibs.expio.fit_funs.sigmoid_log, pynibs.expio.fit_funs.sigmoid4_log]:
                self.log_scale = True
                self.y_passed = np.log10(self.y)
                self.var_y_passed = np.var(self.y_passed)

            if self.fun == pynibs.expio.fit_funs.sigmoid4_log and y_min <= 0:
                y0 = 1e-3
            elif (self.fun == pynibs.expio.fit_funs.sigmoid4_log and y_min > 0) or self.fun == pynibs.expio.fit_funs.sigmoid4:
                y0 = y_min
            else:
                y0 = 0

            if (self.x < 0).all():
                y1 = y_max
                y2 = y_min
                x1 = x_max
                x2 = x_min
            else:
                y1 = y_min
                y2 = y_max
                x1 = x_min
                x2 = x_max

            # set initial amp to maximum of y-data
            if "amp" not in self.init_vals:
                amp = y_max
            else:
                amp = y_max * self.init_vals["amp"]['m'] + self.init_vals["amp"]['t']

            # set initial x_0 to 3/4 of x-data range
            if "x0" not in self.init_vals:
                x0 = x1 + 3 / 4 * (x2 - x1)
            else:
                x0 = x1 + float(self.init_vals["x0"]['p']) * (x2 - x1)

            # set initial r slope to slope of tangent over middle of 25% e-range
            if "r" not in self.init_vals:
                r = 16 / (x_max - x_min) * (y2 - y1) / (amp - y0)
            else:
                r = np.square(self.init_vals["r"]['p']) / (x_max - x_min) * (y2 - y1) / (amp - y0)

            # choose factor to multiply initial values with to calculate limits
            if "limit_factor" not in self.limits:
                limit_factor = 100
            else:
                limit_factor = self.limits["limit_factor"]

            # choose factor to multiply initial values with to calculate initial value range during refitting
            if "range_factor" not in self.random_vals_init_range:
                range_factor = 3
            else:
                range_factor = self.random_vals_init_range["range_factor"]

            self.set_init_vals({"x0": x0,
                                "amp": amp,
                                "r": r})
            self.random_vals_init_range = {"x0": [np.min((0, x0 * range_factor)),
                                                  np.max((0, x0 * range_factor))],
                                           "amp": [np.min((0, amp * range_factor)),
                                                   np.max((0, amp * range_factor))],
                                           "r": [np.min((0, r * range_factor)),
                                                 np.max((0, r * range_factor))]}
            self.set_limits({"x0": [np.min((0, x0 * limit_factor)),
                                    np.max((0, x0 * limit_factor))],
                             "amp": [np.min((1e-12, amp * limit_factor)),
                                     np.max((1e-12, amp * limit_factor))],
                             "r": [np.min((1e-12, r * limit_factor)),
                                   np.max((1e-12, r * limit_factor))]})

            if self.fun in [pynibs.expio.fit_funs.sigmoid4, pynibs.expio.fit_funs.sigmoid4_log]:
                if "y0" not in self.init_vals:
                    self.set_init_vals({"y0": y0})
                else:
                    self.set_init_vals({"y0": y0 * self.init_vals["y0"]['m'] + self.init_vals["y0"]['t']})

                self.random_vals_init_range["y0"] = [1e-12,
                                                     self.init_vals["y0"] * range_factor]

                if "y0" not in self.limits:
                    self.set_limits({"y0": [1e-12, y_max]})
                else:
                    self.set_limits({"y0": [self.limits["y0"]['c'], y_max * self.limits["y0"]['m']]})

        elif self.fun == pynibs.expio.fit_funs.dummy_fun:
            if "a" not in self.limits:
                self.set_limits({"a": [0, 1]})
            else:
                self.set_limits(self.limits)

            if "a" not in self.init_vals:
                self.set_init_vals({"a": 1})
            else:
                self.set_init_vals(self.init_vals)

            if "a" not in self.random_vals_init_range:
                self.random_vals_init_range["a"] = [0, 1]
        else:
            raise NotImplementedError(self.fun)

    def set_random_init_vals(self):
        """ Set random initial values """
        init_vals_new = dict()

        for p in self.init_vals.keys():
            init_vals_new[p] = np.random.rand() * \
                               (self.random_vals_init_range[p][1] - self.random_vals_init_range[p][0]) \
                               + self.random_vals_init_range[p][0]

        self.set_init_vals(init_vals_new)

    def run_select_signed_data(self):
        """
        Selects positive or negative data by performing an initial linear fit by comparing the resulting p-values,
        slopes and R2 values. Either positive or negative data (w.r.t. x-axis) yielding a fit with a p-value < 0.05,
        a positive slope and the higher R2 value is used and the remaining data with the other sign is omitted
        from the analysis
        """
        print('Running run_select_signed_data now - not all data is used!')
        mask_pos = self.x > 0
        mask_neg = np.logical_not(mask_pos)

        # fit positive data (perform the regression when we have at least 20 data points to stabilize results)
        if np.sum(mask_pos) > 20:
            s_pos = scipy.stats.linregress(x=self.x[mask_pos], y=self.y[mask_pos])
            p_pos = s_pos.pvalue
            b_pos = s_pos.slope
            residual = s_pos.slope * self.x[mask_pos] + s_pos.intercept - self.y[mask_pos]
            r2_pos = 1 - np.var(residual) / np.var(self.y[mask_pos])
        else:
            p_pos = 1
            r2_pos = np.NaN
            b_pos = -1

        # fit negative data (perform the regression when we have at least 20 data points to stabilize results)
        if np.sum(mask_neg) > 20:
            s_neg = scipy.stats.linregress(x=self.x[mask_neg], y=self.y[mask_neg])
            p_neg = s_neg.pvalue
            b_neg = s_neg.slope
            residual = s_neg.slope * self.x[mask_neg] + s_neg.intercept - self.y[mask_neg]
            r2_neg = 1 - np.var(residual) / np.var(self.y[mask_neg])
        else:
            p_neg = 1
            r2_neg = np.NaN
            b_neg = 1

        # only use data with p < 0.001 and when slopes show an increase of y-data with increasing |x|-data otherwise,
        # set status to False to indicate that the fit is omitted and set score to NaN
        pos_valid = False
        neg_valid = False

        if (b_pos > 0) and (p_pos < 0.001):
            pos_valid = True
        if (b_neg < 0) and (p_neg < 0.001):
            neg_valid = True
        if pos_valid and not neg_valid:
            self.x = self.x[mask_pos]
            self.y = self.y[mask_pos]
            self.r2_lin = r2_pos
        elif neg_valid and not pos_valid:
            self.x = self.x[mask_neg]
            self.y = self.y[mask_neg]
            self.r2_lin = r2_neg
        elif pos_valid and neg_valid:
            if r2_pos > r2_neg:
                self.x = self.x[mask_pos]
                self.y = self.y[mask_pos]
                self.r2_lin = r2_pos
            else:
                self.x = self.x[mask_neg]
                self.y = self.y[mask_neg]
                self.r2_lin = r2_neg
        else:
            self.status = False
            self.score = np.NaN

        self.var_y = np.var(self.y)
        self.norm_y = np.linalg.norm(self.y)

        if self.log_scale:
            self.y_passed = np.log10(self.y)
        else:
            self.y_passed = self.y

    def run_fit(self, max_nfev=1000):
        """
        Perform data fit with lmfit.
        """
        if self.score_type != "rho":
            fit = self.gmodel.fit(self.y_passed, x=self.x,
                                  calc_covar=False, method="leastsq", max_nfev=max_nfev, scale_covar=False)
            self.best_values = fit.best_values
            self.residual = fit.residual
        self.calc_score()

    def calc_score(self):
        """
        Determine goodness-of-fit score.
        """
        # R2
        if self.score_type == "R2":
            # self.score = 1 - np.var(self.residual) / self.var_y_passed
            self.score = 1 - np.sum(self.residual**2) / (self.var_y_passed * len(self.residual))
        elif self.score_type == "R2_old":
            self.score = 1 - np.var(self.residual) / self.var_y_passed
        # Relative standard error of regression
        elif self.score_type == "SR":
            self.score = 1 - np.linalg.norm(self.residual) / self.norm_y
        elif self.score_type == "rho":
            self.score = scipy.stats.spearmanr(self.y_passed, b=self.x)
        else:
            raise NotImplementedError(f"{self.score_type} not implemented.")


def workhorse_element_run_fit(element, max_nfev=10):
    """ Workhorse to run single Element fit. If status is False, the element will not be fitted. """
    if element.status:
        element.run_fit(max_nfev=max_nfev * len(element.x))

    return element


def workhorse_element_init(ele_id, e_matrix, mep, fun, score_type, select_signed_data, constants, **kwargs):
    """ Workhorse to initialize Elements. """
    element = Element(x=e_matrix[:, ele_id],
                      y=mep,
                      ele_id=ele_id,
                      fun=fun,
                      score_type=score_type,
                      select_signed_data=select_signed_data,
                      constants=constants,
                      **kwargs)
    return element


def regress_data(e_matrix, mep, elm_idx_list=None, element_list=None, fun=pynibs.expio.fit_funs.sigmoid4, n_cpu=4, con=None,
                 n_refit=50, zap_idx=None, return_fits=False, score_type="R2",
                 verbose=False, pool=None, refit_discontinuities=True, select_signed_data=False,
                 mp_context="fork", **kwargs):
    """
    Mass-univariate nonlinear regressions on raw MEP_{AMP} ~ E.
    That is, for each element in elm_idx_list, it's E (mag | norm | tan) for each zap regressed on the raw MEP
    amplitude. An element wise R2 score is returned.
    The function reads the precomputed array of E-MEP data from an .hdf5 file.

    Parameters
    ----------
    e_matrix : np.ndarray of float
         (n_zaps, n_ele) Electric field matrix.
    mep : np.ndarray of float
        (n_zaps,) Motor evoked potential for each stimulation.
    elm_idx_list : np.ndarray of int or list of int
        (n_zaps,) List containing the element indices the fit is performed for.
    element_list : list of Element object instances, optional
        [n_ele] pynibs.Element objects ot skip initialization here.
    fun : pynibs.exp.Mep, default: pynibs.sigmoid4
        A pynibs.exp.Mep function (exp0, sigmoid, sigmoid4, ...).
    n_cpu : int, default: 4
        Number of threads, if n_cpu=1 no parallel pool will be opened and all calculations are done in serial.
    con : np.ndarray of float, optional
        (n_ele, 3 or 4) Connectivity matrix of ROI. Needed in case of refit for discontinuity checks.
    n_refit : int, default: 50
        Maximum number of refits of zero elements. No refit is applied in case of n_refit = 0.
    zap_idx : np.ndarray of int or list of int, optional
        Which e/mep pairs to use.
    return_fits : bool, optional
        Return fit objects containing the parameter estimates.
    score_type : str, default: "R2"
        Error measure of fit:

        * "R2": R2 score (Model variance / Total variance); linear fits: [0, 1], 1 ... perfect fit
        * "SR": Relative standard error of regression (1 - Error 2-norm / Data 2-norm); [-Inf, 1], 1 ... perfect fit
        * "rho": Spearman correlation coefficient [-1, 1]; finds any monotonous correlation (0 means no correlation)
    verbose : bool, default: False
        Plot output messages.
    pool : multiprocessing.Pool()
        pool instance to use.
    refit_discontinuities : bool, default: True
        Refit discontinuous elements. If True, provide _con_.
    mp_context : str, default: "fork"
        Controls the method the sub-processes of the multiprocessing pool (in case of n_cpu > 1) are launched.

        * fork: (only supported by Unix) mp processes diverge from the main process,
          the entire stack, variables and other resources are copied over.
          From the docs: "The child process, when it begins, is effectively identical to the parent process.
          All resources of the parent are inherited by the child process. Note that safely forking a
          multithreaded process is problematic."
        * spawn: (supported by Window and Unix) mp processes are launched in an entirely new Python interpreter
          as separate processes. Variables are copied other resources are freshly instantiated.
          From the docs: "In particular, unnecessary file descriptors and handles from the parent process
          will not be inherited. Starting a process using this method is rather slow compared to using
          fork or forkserver."
    **kwargs
        Passed on to pynibs.Element() to set fit parameters.

    Returns
    -------
    score : np.ndarray of float
        (n_roi, n_qoi) Score for each element.
    best_values : list of dict
        (n_ele) List of parameter fits. Only returned if return_fits=True.
    """
    if zap_idx is None:
        zap_idx = np.array(range(e_matrix.shape[0]))
    if isinstance(zap_idx, list):
        zap_idx = np.array(zap_idx)

    if refit_discontinuities:
        assert con is not None, f"Provide 'con' parameter to fit discontinuties"
    refit_thr = 1e-6
    constants = None

    if elm_idx_list is None:
        elm_idx_list = np.arange(e_matrix.shape[1])

    if fun == pynibs.expio.fit_funs.dummy_fun:
        c_all = np.random.random(len(elm_idx_list))

        if return_fits:
            best_values = [{"a": 1} for _ in range(len(elm_idx_list))]
            return c_all, best_values
        else:
            return c_all

    # shuffle elements because some of them need longer to compute
    # (in this way it is distributed more equally over all cores)
    np.random.shuffle(elm_idx_list)

    # Setting up parallelization
    ####################################################################
    if n_cpu > 1:
        if not pool:
            n_cpu_available = multiprocessing.cpu_count()
            n_cpu = min(n_cpu, n_cpu_available, len(elm_idx_list))
            pool = multiprocessing.get_context(mp_context).Pool(n_cpu)
            local_pool = True

            if verbose:
                print(" > Setting up multiprocessing using {}/{} cores".format(n_cpu, n_cpu_available))
        else:
            local_pool = False  # close pool only if created locally
            if verbose:
                print(" > Using provided pool object")

        # defining workhorses
        workhorse_partial = partial(workhorse_element_run_fit,
                                    max_nfev=10)

        workhorse_init_partial = partial(workhorse_element_init,
                                         e_matrix=e_matrix[zap_idx],
                                         mep=mep[zap_idx],
                                         fun=fun,
                                         score_type=score_type,
                                         select_signed_data=select_signed_data,
                                         constants=constants,
                                         **kwargs)
    else:
        local_pool = False
        if verbose:
            print(" > Running computations with n_cpu=1.")

    # initialize elements
    ####################################################################
    if element_list is None:
        start = time.time()
        if n_cpu <= 1:
            element_list = [Element(x=e_matrix[zap_idx][:, ele_id],
                                    y=mep[zap_idx],
                                    ele_id=ele_id,
                                    fun=fun,
                                    score_type=score_type,
                                    select_signed_data=select_signed_data,
                                    constants=constants,
                                    **kwargs
                                    ) for ele_id in elm_idx_list]

        else:
            element_list = pool.map(workhorse_init_partial, elm_idx_list)
        stop = time.time()
        if verbose:
            print(f"Initialized {len(elm_idx_list)} elements: {stop - start:2.2f} s")

    # run fit
    ####################################################################
    start = time.time()

    if n_cpu <= 1:
        for ele in element_list:
            ele.run_fit(max_nfev=10 * len(ele.x))
    else:
        element_list = pool.map(workhorse_partial, element_list)

    stop = time.time()

    if verbose:
        print(f"Determined scores: {stop - start:2.2f} s")

    if fun == pynibs.expio.fit_funs.linear:
        n_refit = 0
        refit_discontinuities = False
        print("Skipping refit for linear fits ...")

    # refit elements
    ####################################################################
    if n_refit > 0:

        # refit bad elements
        ####################################################################
        i_refit = 0
        while i_refit < n_refit:

            # get index in Element_list of elements where refit should be performed
            ele_idx_refit = [i_ele for i_ele, ele in enumerate(element_list) if ele.score < refit_thr]
            element_list_refit = [element_list[i_ele] for i_ele in ele_idx_refit]

            if len(element_list_refit) > 0:

                if verbose:
                    print(f" > Performing refit for {len(element_list_refit)} zero elements ...")

                # set random start values
                for ele in element_list_refit:
                    ele.set_random_init_vals()

                start = time.time()

                if n_cpu == 1:
                    for ele in element_list_refit:
                        ele.run_fit(max_nfev=10 * len(ele.x))
                else:
                    element_list_refit = pool.map(workhorse_partial, element_list_refit)

                stop = time.time()

                if verbose:
                    print(f"Determined scores: {stop - start:2.2f} s")

                # replace new fits if they have higher scores than the old ones
                for i_ele, ele_idx_re in enumerate(ele_idx_refit):
                    if element_list_refit[i_ele].score > element_list[ele_idx_re].score:
                        element_list[ele_idx_re] = element_list_refit[i_ele]

                i_refit += 1
            else:
                break

    # sort elements for discontinuity check
    element_list = [ele for _, ele in sorted(zip(elm_idx_list, element_list))]

    # find discontinuities and refit
    ##################################################################
    if refit_discontinuities:

        if len(element_list) > 1:
            score = np.array([ele.score for ele in element_list])
            not_fitted_elms = np.array([idx for idx, ele in enumerate(element_list) if np.isnan(ele.score)])
            idx_disc, idx_neighbor = pynibs.get_indices_discontinuous_data(data=score,
                                                                           con=con,
                                                                           neighbor=True,
                                                                           deviation_factor=2,
                                                                           not_fitted_elms=not_fitted_elms)
            element_list_disc = [element_list[i_ele] for i_ele in idx_disc]

            if len(idx_disc) > 0:
                if verbose:
                    print(f" > Performing refit for {len(idx_disc)} discontinuous elements ...")

                # refit for discontinuous elements
                if len(idx_disc) > 0:
                    # set start values from neighbors
                    for i_ele, idx_ne in zip(range(len(idx_disc)), idx_neighbor):
                        element_list_disc[i_ele].set_init_vals(element_list[idx_ne].best_values)

                    start = time.time()

                    if n_cpu == 1:
                        for ele in element_list_disc:
                            ele.run_fit(max_nfev=10 * len(ele.x))
                    else:
                        element_list_disc = pool.map(workhorse_partial, element_list_disc)

                    stop = time.time()

                    if verbose:
                        print(f"Determined scores: {stop - start:2.2f} s")

                    # replace new fits if they have higher scores than the old ones
                    for i_ele, ele_idx_re in enumerate(idx_disc):
                        if element_list_disc[i_ele].score > element_list[ele_idx_re].score:
                            element_list[ele_idx_re] = element_list_disc[i_ele]

    score = np.array([ele.score for ele in element_list])
    if local_pool:
        pool.close()
        pool.join()

    if return_fits:
        best_values = np.array([ele.best_values for ele in element_list])
        return score, best_values
    else:
        return score


def sing_elm_raw(elm_idx_list, mep_lst, mep_params, e, alpha=1000):
    """
    Mass-univariate ridge regressions on raw MEP_{AMP} ~ E.
    That is, for each element in elm_idx_list, it's E (mag | norm | tan) for each zap regressed on the raw MEP
    amplitude. An element wise sklearn.metrics.regression.r2_score is returned.

    elm_idx_list : np.ndarray
        (chunksize) List of element indices, the congruence factor is computed for.
    mep: list of Mep object instances
        (n_cond) List of fitted Mep object instances for all conditions (see exp.py for more information of Mep class).
    mep_params: np.ndarray of float
        (n_mep_params_total) List of all mep parameters of curve fits used to calculate the MEP (accumulated into 1
        array) (e.g.: [mep_#1_para_#1, mep_#1_para_#2, mep_#1_para_#3, mep_#2_para_#1, mep_#2_para_#1, ...])
    e: np.ndarray of float
        (n_elm, n_cond, n_qoi) array of the electric field to compute the r2 factor for, e.g. (e_mag, e_norm, e_tan).

    Returns
    -------
    r2: np.ndarray of float
        (n_roi, n_datasets) R^2 for each element in elm_idx_list.
    """
    from pandarallel import pandarallel
    pandarallel.initialize(verbose=0)

    def cartesian_product(*arrays):
        """
        Fast implementation to get cartesian product of two arrays.

        cartesian_product([a,b,c],[2,3]) =
             [a, 2
              a, 3
              b, 2
              b, 3
              c, 2
              c, 3]
        """
        la = len(arrays)
        dtype = np.result_type(*arrays)
        arr = np.empty([len(a) for a in arrays] + [la], dtype=dtype)
        for i, a in enumerate(np.ix_(*arrays)):
            arr[..., i] = a
        return arr.reshape(-1, la)

    n_eqoi = e.shape[2]
    n_cond = e.shape[1]
    n_elm = e.shape[0]
    assert n_cond == len(mep_lst)
    scores = None
    reg_r2 = np.empty((n_elm, n_eqoi))

    for qoi_idx in range(n_eqoi):
        # t_q = time.time()
        x = pd.DataFrame()
        index_shift = 0
        amplitudes = np.array(())

        start = time.time()
        for mep_i, mep in enumerate(mep_lst):
            # condition wise, as we stimulated with different intensities per conditions

            # for each element in roi, one datapoint for each zap.
            current = cartesian_product(e[:, mep_i, qoi_idx], mep.intensities)

            # index is iteration of zaps over all conditions
            index = cartesian_product(e[:, mep_i, qoi_idx], np.arange(len(mep.intensities)))[:, 1]
            index += index_shift
            index_shift = index[-1] + 1

            # el is range(n_elements) * n_zaps_in_condition
            el_idx = np.repeat(np.arange(e.shape[0]), len(mep.intensities))

            # intensity * e
            e_zap = np.multiply(current[:, 0], current[:, 1])

            # put all together
            x_cond = pd.DataFrame(data={"index": index.astype(int),
                                        "el": el_idx,
                                        "e": e_zap})
            amplitudes = np.append(amplitudes, mep.mep)
            # "current": current[:, 1],
            # "mep": mep[:,1]})

            # x_cond['condition'] = mep_i
            x = x.append(x_cond)

        stop = time.time()
        print(f"Prepare dataset t = {stop - start}")
        # x.shape is now (n_zaps*n_elms, 3)

        # reshape to (n_zaps, n_elms)
        start = time.time()
        x = x.pivot(index="index", columns="el", values="e")
        stop = time.time()
        print(f"Pivot t = {stop - start}")

        do_reg_poly = False
        do_reg_linear = True
        reg = LinearRegression()

        if do_reg_poly:
            # fn = "/data/pt_01756/tmp/reg/mean_data_roi_lasso.hdf5"
            # print "creating polynomial features"
            # poly = PolynomialFeatures(2, interaction_only=True, include_bias=False)
            # x_pol = poly.fit_transform(x.iloc[:,:-3])
            # print "fitting regressor"
            # reg.fit(x_pol, x['mep'])
            # with h5py.File(fn, 'a') as f:
            #     f.create_dataset(name='/data/tris/' + e,
            #                      data=reg.coef_[:E.shape[0]])
            #     e_poly = []
            #     # get interaction mapped back to elemens
            #     for el in range(E.shape[0]):
            #         print "getting indices for interactions"
            #         idx = [i for i, s in enumerate(poly.get_feature_names()) if 'x{} '.format(el + 1) in s]
            #         e_poly.append(np.sum(reg.coef_[idx]))
            #     f.create_dataset(name='/data/tris/' + e + '_poly',
            #                      data=e_poly)
            # data_qoi_tmp = e_poly
            raise NotImplementedError

        elif do_reg_linear:
            # Do one regression per element.
            # r_t = time.time()

            def get_score(x_i):
                """Helper function do be used by pd.apply() to speed up things.

                Paramsmeters
                ------------
                x_i: pd.Series
                 Column with e for a single elm.

                Returns
                -------
                r2 for amplitudes ~ E
                """
                # x_i = x_i.reshape(-1, 1)
                reg.fit(x_i.reshape(-1, 1), amplitudes)
                return reg.score(x_i.reshape(-1, 1), amplitudes)

            # apply get_score function column wise
            start = time.time()
            # scores = x.parallel_apply(get_score, axis=0, raw=True)
            scores = x.transpose().swifter.apply(get_score, axis=1, raw=True)
            stop = time.time()
            print(f"Fit and score t = {stop - start}")
            # print "all_reg: {}".format(time.time() - r_t)

        reg_r2[:, qoi_idx] = np.array(scores)
        # data.append(data_qoi_tmp)

        # print "qoi {}: {}".format(qoi_idx, time.time() - t_q)

    return reg_r2


def write_regression_hdf5(fn_exp_hdf5, fn_reg_hdf5, qoi_path_hdf5, qoi_phys, e_results_folder, qoi_e, roi_idx,
                          conds_exclude):
    """
    Reads the stimulation intensities from the experiment.hdf5 file.
    Reads the qoi from the experiment.hdf5 file.
    Reads the electric fields from the electric field folder.
    Weights the electric field voxel wise with the respective intensities
    and writes an .hdf5 file containing the preprocessed data (pandas dataframe).

    fn_exp_hdf5 : str
        Filename of the experiment.hdf5 file.
    fn_reg_hdf5 : str
        Filename of output regression.hdf5 file.
    qoi_path_hdf5 : str
        Path in experiment.hdf5 file pointing to the pandas dataframe containing the qoi.
        (e.g.: "phys_data/postproc/EMG")
    qoi : str
        Name of QOI the congruence factor is calculated with.
        (e.g.: "p2p")
    fn_e_results : str
        Folder containing the electric fields.
        (e.g.: "/data/pt_01756/probands/13061.30/results/electric_field/1")
    qoi_e : str or list of str
        Quantities of the electric field used to calculate the congruence factor (e.g. ["E", "E_norm", "E_tan"]
        Has to be included in e.hdf5 -> e.g.: "data/midlayer/roi_surface/1/E".
    roi_idx : int
        ROI index.
    conds_exclude : str or list of str
        Conditions to exclude.

    Returns
    -------
    <File>: .hdf5 file
        File containing the intensity (current) scaled E-fields of the conditions in the ROI.
        Saved in datasets with the same name as qoi_e ["E", "E_norm", "E_tan"]
    """
    def cartesian_product(*arrays):
        """
        Fast implementation to get cartesian product of two arrays.

        cartesian_product([a,b,c],[2,3]) =
             [a, 2
              a, 3
              b, 2
              b, 3
              c, 2
              c, 3]
        """
        la = len(arrays)
        dtype = np.result_type(*arrays)
        arr = np.empty([len(a) for a in arrays] + [la], dtype=dtype)
        for i, a in enumerate(np.ix_(*arrays)):
            arr[..., i] = a
        return arr.reshape(-1, la)

    if conds_exclude is not list:
        conds_exclude = [conds_exclude]

    if type(qoi_e) is not list:
        qoi_e = [qoi_e]

    # read dataframe of stimulation data
    df_stim_data = pd.read_hdf(fn_exp_hdf5, "stim_data")  # type: pd.DataFrame
    conds = np.unique(df_stim_data["condition"])
    conds = [c for c in conds if c not in conds_exclude]

    # read dataframe of postproc containing the qoi
    df_qoi = pd.read_hdf(fn_exp_hdf5, qoi_path_hdf5)  # type: pd.DataFrame

    n_conds = len(conds)
    n_qoi_e = len(qoi_e)
    n_ele = 0

    # read electric fields
    for i_c, c in enumerate(conds):

        # read electric field
        fn_e_hdf5 = os.path.join(e_results_folder, c, 'simulations', 'e.hdf5')
        print(" > Loading electric field from file: {}".format(fn_e_hdf5))
        f = h5py.File(fn_e_hdf5, 'r')

        # generate E array in first iteration
        if i_c == 0:
            n_ele = f[f"data/midlayer/roi_surface/{roi_idx}/{qoi_e[0]}"][:].shape[0]

            # np.array [n_ele x n_zaps x n_qoi]
            e = np.zeros((n_ele, n_conds, n_qoi_e))

        # midlayer E
        for i_q_e, q_e in enumerate(qoi_e):
            e[:, i_c, i_q_e] = f[f"data/midlayer/roi_surface/{roi_idx}/{q_e}"][:]

    print("\n")

    # scale electric fields with respective intensities
    for i_q_e, q_e in enumerate(qoi_e):
        print(f"Preparing regression datasets for {q_e}")
        print(f"========================================")
        x = pd.DataFrame()
        index_shift = 0
        qoi_amplitudes = np.array(())

        start = time.time()
        for i_c, c in enumerate(conds):
            # extract stimulation intensity and qoi amplitude for condition
            mask = df_stim_data["condition"].values == c
            stim_intensity = df_stim_data.loc[mask, "current"].values

            # for each element in roi, one datapoint for each zap.
            e_stim_intensity = cartesian_product(e[:, i_c, i_q_e], stim_intensity)

            # intensity * e
            e_scaled = np.multiply(e_stim_intensity[:, 0], e_stim_intensity[:, 1])

            # index is iteration of zaps over all conditions
            index = cartesian_product(e[:, i_c, i_q_e], np.arange(len(stim_intensity)))[:, 1]
            index += index_shift
            index_shift = index[-1] + 1

            # el is range(n_elements) * n_zaps_in_condition
            el_idx = np.repeat(np.arange(n_ele), len(stim_intensity))

            # put all together
            x_cond = pd.DataFrame(data={"index": index.astype(int),
                                        "el": el_idx,
                                        "e": e_scaled})

            qoi_amplitudes = np.append(qoi_amplitudes, df_qoi.loc[mask, qoi_phys].values)

            x = x.append(x_cond)

        stop = time.time()
        print(f"Prepare dataset: t = {stop - start}")
        # x.shape is now (n_zaps*n_elms, 3)

        # reshape to (n_zaps, n_elms)
        start = time.time()
        x = x.pivot(index="index", columns="el", values="e")
        x["qoi_amplitudes"] = qoi_amplitudes
        stop = time.time()
        print(f"Pivot: t = {stop - start}")

        start = time.time()
        x.to_hdf(fn_reg_hdf5, q_e)
        stop = time.time()
        print(f"Write hdf5: t = {stop - start}")
        print(f"\n")


def ridge_from_hdf5(elm_idx_list, fn_reg_hdf5, qoi_path_hdf5, zap_idx=None):
    """
    Mass-univariate ridge regressions on raw MEP_{AMP} ~ E.
    That is, for each element in elm_idx_list, it's E (mag | norm | tan) for each zap regressed on the raw MEP
    amplitude. An element wise sklearn.metrics.regression.r2_score is returned.
    The function reads the precomputed array of E-MEP data from an .hdf5 file.
    Always uses all cores of a machine!

    elm_idx_list : list of int
        List containing the element indices the fit is performed for.
    fn_hdf5 : str
        Filename (incl. path) containing the precomputed E-MEP dataframes.
    qoi_path_hdf5 : str
        Path in .hdf5 file to dataset of electric field qoi.
    zap_idx : np.ndarray, optional
        (n_zaps) Indices of zaps the congruence factor is calculated with (default: all).

    Returns
    -------
    r2 : np.ndarray of float
        (n_roi, n_datasets) R^2 for each element in elm_idx_list.
    """
    x = pd.read_hdf(fn_reg_hdf5, qoi_path_hdf5)  # type: pd.DataFrame

    # use all stimuli if zap_idx is not provided
    if zap_idx is None:
        zap_idx = np.arange(x.shape[0])

    n_zaps = x.shape[0]
    mask = np.zeros(n_zaps).astype(bool)
    mask[zap_idx] = True
    qoi_amplitudes = x.loc[mask, "qoi_amplitudes"].values
    x = x.drop("qoi_amplitudes", axis=1)

    def get_score(x_i):
        """Helper function do be used by pd.apply() to speed up things.

        Parameters
        ----------
        x_i : pd.Series
            Column with e for a single elm.

        Returns
        -------
        r2 for amplitudes ~ E
        """

        reg.fit(x_i.reshape(-1, 1), qoi_amplitudes)
        return reg.score(x_i.reshape(-1, 1), qoi_amplitudes)

    reg = LinearRegression()

    # apply get_score function column (element) wise
    # start = time.time()
    r2 = x.loc[mask, elm_idx_list].apply(get_score, axis=0, raw=True)
    # stop = time.time()
    # print(f"Fit and score t = {stop - start}")

    return np.array(r2)[:, np.newaxis]


def fit_elms(elm_idx_list, e_matrix, mep, zap_idx=None,
             fun=pynibs.expio.fit_funs.sigmoid4, init_vals=None, limits=None, log_scale=False,
             constants=None, max_nfev=None, bad_elm_idx=None, score_type="R2",  # mask_e_field=None,
             verbose=False):
    """
    Workhorse for Mass-univariate nonlinear regressions on raw MEP_{AMP} ~ E.
    That is, for each element in elm_idx_list, it's E (mag | norm | tan) for each zap regressed on the raw MEP
    amplitude. An element wise r2 score is returned.

    Parameters
    ----------
    elm_idx_list : list of int or np.ndarray
        List containing the element indices the fit is performed for.
    e_matrix : np.ndarray of float
        (n_zaps, n_ele) Electric field matrix.
    mep : np.ndarray of float
        (n_zaps) Motor evoked potentials for every stimulation.
    zap_idx : np.ndarray, optional
        (n_zaps) Indices of zaps the congruence factor is calculated with (default: all).
    fun : str
        A function name of pynibs.exp.Mep (exp0, sigmoid).
    init_vals : np.ndarray of dict
         (len(elm_idx_list)) Dictionary containing the initial values for each element as np.ndarray.
        The keys are the free parameters of fun, e.g. "x0", "amp", etc.
    limits : pd.DataFrame
        Dictionary containing the limits of each parameter for each element e.g.: limits["x0"][elm_idx] = [min, max].
    log_scale : bool, default: False
        Log-transform data before fit (necessary for functions defined in the log domain).
    constants : dict of <string>:<num>, optional
        key:value pair of model parameters not to optimize.
    max_nfev : int, optional
        Max fits, passed to model.fit() as max_nfev=max_nfev*len(x).
    bad_elm_idx : np.ndarray, optional
        Indices of elements not to fit, with indices corresponding to indices (not values) of elm_idx_list.
    score_type : str, default: "R2"
        Goodness of fit measure; Choose SR for nonlinear fits and R2 or SR for linear fits:

        * "R2": R2 score (Model variance / Total variance) [0, 1] for linear models; 0: bad fit; 1: perfect fit
        * "SR": Relative standard error of regression (1 - Error 2-norm / Data 2-norm) [-inf, 1]; 1: perfect fit
    verbose : bool, default: False
        Print verbosity information

    Returns
    -------
    r2 : np.ndarray of float
        (n_roi, 1) R2 for each element in elm_idx_list.
    best_values : np.ndarray of object
        Fit parameters returned from the optimizer.
    """
    str_pref = "Main"
    start = time.time()
    try:
        str_pref = f"{multiprocessing.current_process()._identity[0]:0>2} "
    except:
        pass

    # use all stimuli if zap_idx is not provided
    if zap_idx is not None:
        e_matrix = e_matrix[zap_idx, :]
        mep = mep.iloc[zap_idx]
    n_zaps = e_matrix.shape[0]

    if bad_elm_idx is None:
        bad_elm_idx = set(np.array([]))

    best_values = [0] * len(elm_idx_list)
    r2 = np.zeros((len(elm_idx_list),)) - 10

    if fun == pynibs.expio.fit_funs.dummy_fun:
        return r2, np.array(best_values)

    if max_nfev is not None:
        max_nfev = n_zaps * max_nfev

    # set up gmodel
    gmodel = Model(fun)

    # create the param_hints ordereddict only once.
    for p in gmodel.param_names:
        gmodel.param_hints[p] = OrderedDict()

    param_hints = gmodel.param_hints
    make_params = gmodel.make_params
    param_names = gmodel.param_names

    if limits is not None:
        limits_p = list(limits.columns)
    else:
        limits_p = set()

    if init_vals is not None:
        inits_p = list(init_vals.columns)
        if constants is not None:
            inits_p = np.setdiff1d(inits_p, list(constants.keys()))
    else:
        inits_p = set()

    constants_elmwise = {}
    if constants is not None:
        for p in param_names:
            if p in constants:
                param_hints[p]['vary'] = False

                try:
                    len(constants[p])
                    constants_elmwise[p] = constants[p]
                except TypeError:
                    param_hints[p]['value'] = constants[p]

    # transform mep to log domain
    if log_scale:
        qoi = np.log10(mep)
    else:
        qoi = mep

    qoi_norm = np.linalg.norm(qoi)
    qoi_var = np.var(qoi)
    best_values = {k: [] for k in init_vals.keys()}
    # loop over elements
    for i, elm_idx in enumerate(elm_idx_list):
        if i in bad_elm_idx:
            [best_values[k].append(0) for k, v in best_values.items()]
            continue
        for p in limits_p:
            # set limits
            # skip the set_param_hin function and do this by hand
            param_hints[p]['min'], param_hints[p]['max'] = limits.iloc[i][p]

        for p in inits_p:
            # set initial values
            param_hints[p]['value'] = init_vals.iloc[i][p]

        for p in constants_elmwise:
            param_hints[p]['value'] = constants_elmwise[p][i]
        make_params()

        # perform fit (for some reason, sigmoid4_log function generates NaN, which I can not reproduce)
        # I catch the fitting error (ValueError) and set an r2 score of -1, such that it will go into the refit

        e_elm = e_matrix[:, elm_idx]
        # set max_nfev to a reasonable range
        fit = gmodel.fit(qoi, x=e_elm,
                         calc_covar=False, method="leastsq", max_nfev=max_nfev, scale_covar=False)
        [best_values[k].append(v) for k, v in fit.best_values.items()]

        # calculate goodness-of-fit score
        if score_type == "R2":
            # "R2"
            r2[i] = 1 - np.var(fit.residual) / qoi_var
        elif score_type == "SR":
            # Relative standard error of regression
            r2[i] = 1 - np.linalg.norm(fit.residual) / qoi_norm  # TODO: what is this??
            # np.sqrt(np.sum(fit.residual**2) / len(mep)) # <- this should be the correct S
            # print(f"Diff: "
            #       f"{np.round((1 - np.linalg.norm(fit.residual) / qoi_norm) -
            #       np.sqrt(np.sum(fit.residual**2) / len(mep)),2)}")
        elif score_type == 'BIC':
            r2[i] = -fit.bic
        elif score_type == 'AIC':
            r2[i] = fit.aic
        else:
            raise ValueError(f"Error score '{score_type}' not implemented ... ")
    if verbose:
        if len(elm_idx_list) == 0:
            elm_time = 0.0
        else:
            elm_time = (time.time() - start) / len(elm_idx_list)

        print(f"Proc{str_pref}:    > fit_elms workhorse: done "
              f"({time.time() - start:.2f} s, {elm_time:2.4} s/elm, "
              f"mean R2: {np.mean(r2):2.2f})")
    return r2, pd.DataFrame().from_dict(best_values)


# def nl_hdf5_workhorse_idx(elm_idx_list, e_matrix, mep,
#                           fun=sigmoid4, zap_idx=None, init_vals=None, limits=None, log_scale=False,
#                           constants=None, max_nfev=None,
#                           verbose=False):
#     """
#     Workhorse for Mass-univariate nonlinear regressions on raw MEP_{AMP} ~ E.
#     That is, for each element in elm_idx_list, it's E (mag | norm | tan) for each zap regressed on the raw MEP
#     amplitude. An element wise r2 score is returned.
#     The function reads the precomputed array of E-MEP data from an .hdf5 file.
#
#     Parameters
#     ----------
#     elm_idx_list : list of int
#         List containing the element indices the fit is performed for
#     e_matrix : np.ndarray of float [n_zaps x n_ele]
#         Electric field matrix
#     mep : np.ndarray of float [n_zaps]
#         Motor evoked potentials for every stimulation
#     zap_idx : np.array [n_zaps], default: None
#         Indices of zaps the congruence factor is calculated with (default: all)
#     fun : str
#         A function name of pynibs.exp.Mep (exp0, sigmoid)
#     init_vals : dict
#         Dictionary containing the initial values for each element as np.ndarray [len(elm_idx_list)].
#         The keys are the free parameters of fun, e.g. "x0", "amp", etc
#     limits : dict
#         Dictionary containing the limits of each parameter for each element e.g.: limits["x0"][elm_idx] = [min, max]
#     log_scale : bool, default: False
#         Log-transform data before fit (necessary for functions defined in the log domain)
#     constants : dict of <string>:<num>, default: None
#         key:value pair of model parameters not to optimize.
#     bad_elm_idx : np.ndarray
#         Indices of elements not to fit.
#     max_nfev : int, default: None
#         Max fits, passed to model.fit() as max_nfev=max_nfev*len(x).
#     verbose : bool, default: False
#         Print verbosity information
#
#     Returns
#     -------
#     r2 : np.ndarray of float [n_roi, 1]
#         R2 for each element in elm_idx_list
#     fit : fit objects
#         Fit objects returned from the optimizers
#     """
#     # from matplotlib import pyplot as plt
#     str_pref = "Main"
#     start = time.time()
#     try:
#         str_pref = f"{multiprocessing.current_process()._identity[0]:0>2} "
#     except:
#         pass
#
#     # use all stimuli if zap_idx is not provided
#     if zap_idx is not None:
#         zap_idx = np.arange(e_matrix.shape[0])
#         mask = np.zeros(e_matrix.shape[0]).astype(bool)
#         mask[zap_idx] = True
#         e_matrix = e_matrix[mask]
#         mep = mep[mask]
#     n_zaps = e_matrix.shape[0]
#
#     # if verbose:
#     #     print(f"Proc{str_pref}:    > regression_nl_hdf5_workhorse: "
#     #           f"starting ({e_matrix.shape[1]} elms / {n_zaps} zaps)")
#
#     best_values = [0] * e_matrix.shape[1]
#     r2 = np.zeros((e_matrix.shape[1],)) - 10
#
#     if fun == dummy_fun:
#         return r2, best_values
#
#     if max_nfev is not None:
#         max_nfev = n_zaps * max_nfev
#
#     if log_scale:
#         mep = np.log10(mep)
#
#     # set up gmodel
#     gmodel = Model(fun)
#
#     # create the param_hints ordereddict only once.
#     for p in gmodel.param_names:
#         gmodel.param_hints[p] = OrderedDict()
#
#     # compute variance only once
#     var_qoi = np.var(mep)
#
#     param_hints = gmodel.param_hints
#     make_params = gmodel.make_params
#     param_names = gmodel.param_names
#     if limits is not None:
#         limits_p = param_names
#     else:
#         limits_p = set()
#
#     if init_vals is not None :
#         inits_p = list(init_vals.keys())
#         if constants is not None:
#             inits_p = np.setdiff1d(inits_p, list(constants.keys()))
#     else:
#         inits_p = set()
#
#     if constants is not None:
#         for p in param_names:
#             if p in constants:
#                 param_hints[p]['value'] = constants[p]
#                 param_hints[p]['vary'] = False
#
#     #         # set_param_hint(p, value=constants[p], vary=False)
#     for i, elm in enumerate(elm_idx_list):
#         for p in limits_p:
#             # set limits
#             # skip the set_param_hint function and do this by hand
#             param_hints[p]['min'], param_hints[p]['max'] = limits[p][elm]
#
#         for p in inits_p:
#             # set initial values
#             param_hints[p]['value'] = init_vals[p][elm]
#             # set_param_hint(p, value=init_vals[p][elm_idx])
#
#         make_params()
#
#         # perform fit (for some reason, sigmoid4_log function generates NaN, which I can not reproduce)
#         # I catch the fitting error (ValueError) and set an r2 score of -1, such that it will go into the refit
#
#         try:
#             # set max_nfev to a reasonable range
#             fit = gmodel.fit(mep, x=e_matrix[:, i],
#                              calc_covar=False, method="leastsq", max_nfev=max_nfev, scale_covar=False)
#             best_values[i] = fit.best_values
#
#             # calculate R2 score
#             # times1.append(fit.nfev)
#             # mid2 = time.time()
#             if log_scale:
#                 # this could be optimized by pow(10)
#                 residual = 10 ** fit.best_fit - 10 ** qoi_amplitudes
#                 r2[i] = 1 - np.var(residual) / np.var(10 ** qoi_amplitudes)
#
#             else:
#                 r2[i] = 1 - np.var(fit.residual) / var_qoi
#
#         except ValueError:
#             print(f"value error: {i}: {elm}")
#             pass
#
#     if verbose:
#         print(f"Proc{str_pref}:    > regression_nl_hdf5_workhorse: done "
#               f"({time.time() - start:.2f} s, "
#               f"{(time.time() - start) / len(elm_idx_list):2.4} s/elm, "
#               f"mean R2: {np.mean(r2):2.2f})")
#     # f"Done at {time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime())}, pid: {os.getpid()}")
#     return r2, np.array(best_values)


def nl_hdf5(elm_idx_list=None, fn_reg_hdf5=None, qoi_path_hdf5=None, e_matrix=None, mep=None, zap_idx=None,
            fun=pynibs.expio.fit_funs.sigmoid4, n_cpu=4, con=None, n_refit=50, return_fits=False, score_type="R2",
            verbose=False, pool=None, refit_discontinuities=True):
    """
    Mass-univariate nonlinear regressions on raw MEP_{AMP} ~ E.
    That is, for each element in elm_idx_list, it's E (mag | norm | tan) for each zap regressed on the raw MEP
    amplitude. An element wise r2 score is returned.
    The function reads the precomputed array of E-MEP data from an .hdf5 file.

    Parameters
    ----------
    elm_idx_list : np.ndarray of int, optional
        List containing the element indices the fit is performed for, if not all.
    fn_reg_hdf5 : str, optional
        Filename (incl. path) containing the precomputed E-MEP dataframes.
    qoi_path_hdf5: Union[str, list[str]], optional
        Path in .hdf5 file to dataset of electric field qoi e.g.: ["E", "E_norm", "E_tan"]
    e_matrix : np.ndarray of float, optional
        (n_zaps, n_ele) Electric field matrix.
    mep : np.ndarray of float
        (n_zaps) Motor evoked potentials for every stimulation.
    zap_idx : np.ndarray, optional
        (n_used_zaps) Indices of zaps the congruence factor is calculated with, if not all.
    fun : pynibs.exp.Mep function, default: sigmoid4
        A function of pynibs.exp.Mep (exp0, sigmoid).
    n_cpu : int, default: 4
        Number of threads to use.
    con : np.ndarray of float, optional
        (n_roi, 3 or 4) Connectivity matrix of ROI (needed in case of refit because of discontinuity check).
    n_refit : int, default: 50
        Maximum number of refits of zero elements. No refit is applied in case of n_refit = 0.
    return_fits : bool, default: False
        Return fit objects containing the parameter estimates.
    score_type : str, default: "R2"
        Error measure of fit:

        * "R2": R2 score (Model variance / Total variance); linear fits: [0, 1], 1 ... perfect fit
        * "SR": Relative standard error of regression (1 - Error 2-norm / Data 2-norm); [-Inf, 1], 1 ... perfect fit
    verbose : bool, default: False
        Plot output messages.
    pool : multiprocessing.Pool(), optional
        Pool instance to use.
    refit_discontinuities : bool, default: True
        Run refit for discontinuous elements at the end.

    Returns
    -------
    r2 : np.ndarray of float
        (n_roi, n_qoi) R2 for each element in elm_idx_list.
    """
    refit_thr = 1e-6

    if elm_idx_list is None:
        elm_idx_list = np.arange(e_matrix.shape[1])

    if fun == pynibs.expio.fit_funs.dummy_fun:
        c_all = np.random.random(len(elm_idx_list))

        if return_fits:
            best_values = [{"a": 1} for _ in range(len(elm_idx_list))]
            return c_all, best_values
        else:
            return c_all

    if qoi_path_hdf5 is None:
        qoi_path_hdf5 = ["e_matrix"]

    if fn_reg_hdf5 is not None and qoi_path_hdf5 is not None:
        df_reg = pd.read_hdf(fn_reg_hdf5, qoi_path_hdf5)  # type: pd.DataFrame
        e_matrix = df_reg.values[:, :-1]
        mep = df_reg.loc[:, "qoi_amplitudes"].values

    elif e_matrix is None or mep is None:
        raise ValueError("Please provide e_matrix and mep or fn_reg_hdf5 and qoi_path_hdf5!")

    if n_refit > 0 and con is None:
        raise ValueError("Please provide connectivity matrix (con) in case of refit!")

    # shuffle elements because some of them need longer to compute
    # (in this way it is distributed more equally over all cores)
    np.random.shuffle(elm_idx_list)
    if not isinstance(elm_idx_list, np.ndarray):
        elm_idx_list = np.array(elm_idx_list)
    elm_idx_list_shuffle_idx = np.argsort(elm_idx_list).astype(int)
    elm_idx_list = elm_idx_list.tolist()

    # Setting up parallelization
    if not pool:
        n_cpu_available = multiprocessing.cpu_count()
        n_cpu = min(n_cpu, n_cpu_available, len(elm_idx_list))
        pool = multiprocessing.Pool(n_cpu)
        local_pool = True

        if verbose:
            print(" > Setting up multiprocessing using {}/{} cores".format(n_cpu, n_cpu_available))
    else:
        local_pool = False  # close pool only if created locally
        if verbose:
            print(" > Using provided pool object")
    if n_cpu > 1:
        elm_idx_list_chunks = pynibs.compute_chunks(elm_idx_list, n_cpu)
    elif len(elm_idx_list) == 1:
        elm_idx_list_chunks = [elm_idx_list]
    elif n_cpu == 1:
        elm_idx_list_chunks = [elm_idx_list]
    else:
        raise ValueError(f'n_cpu={n_cpu} is invalid. ')

    # setting up initial values and limits
    if verbose:
        print(f" > Setting up initial values and limits for {fun.__name__} function (from {fun.__module__})!")

    log_scale, limits, init_vals, max_vals_refit = get_model_init_values(fun=fun,
                                                                         elm_idx_list=elm_idx_list,
                                                                         e_matrix=e_matrix,
                                                                         mep=mep,
                                                                         mask_e_field=None)

    if verbose:
        print(" > Calculating congruence factor ...")

    if type(qoi_path_hdf5) is not list:
        qoi_path_hdf5 = [qoi_path_hdf5]

    c_all = np.zeros((len(elm_idx_list), len(qoi_path_hdf5)))
    best_values = None
    # loop over electric field QOIs
    for i_q, q in enumerate(qoi_path_hdf5):
        workhorse_partial = partial(fit_elms,
                                    e_matrix=e_matrix,
                                    mep=mep,
                                    fun=fun,
                                    zap_idx=zap_idx,
                                    init_vals=init_vals,
                                    limits=limits,
                                    log_scale=log_scale,
                                    max_nfev=10,
                                    score_type=score_type,
                                    verbose=verbose)

        start = time.time()
        res = pool.map(workhorse_partial, elm_idx_list_chunks, chunksize=1)
        stop = time.time()
        if verbose:
            # print(len(elm_idx_list_chunks))
            print(f"Determine c-factors for {q} / {qoi_path_hdf5}: {stop - start:2.2f} s")

        best_values = []
        c = None
        for i in range(len(res)):
            if i == 0:
                c = res[i][0]
                best_values = res[i][1]
            else:
                if c.ndim < 2:
                    c = c[:, np.newaxis]
                if res[i][0].ndim < 2:
                    c = np.vstack((c, res[i][0][:, np.newaxis]))
                else:
                    c = np.vstack((c, res[i][0]))
                best_values += res[i][1]

        # resort c values
        c = c[elm_idx_list_shuffle_idx]
        best_values = [best_values[i] for i in elm_idx_list_shuffle_idx]
        elm_idx_list = np.array(elm_idx_list)[elm_idx_list_shuffle_idx].astype(list)

        # refit elements
        ####################################################################
        if n_refit > 0:
            params = inspect.getfullargspec(fun).args[1:]

            # refit bad elements
            ####################################################################
            i_refit = 0
            while i_refit < n_refit:
                idx_refit = np.where(c < refit_thr)[0]

                if len(idx_refit) > 0:
                    if verbose:
                        print(f" > Performing refit for {len(idx_refit)} zero elements ...")

                    # set random start values
                    if len(idx_refit) > 0:
                        for p in params:
                            for idx_re in idx_refit:
                                init_vals[p][idx_re] = max_vals_refit[p][idx_re] * np.random.rand()
                                # init_vals[p][idx_re] = init_vals[p][idx_re] + \
                                #                        max_vals_refit[p][idx_re] * (np.random.rand() - 0.5)

                        if n_cpu > 1:
                            elm_idx_list_chunks_zero = pynibs.compute_chunks(np.array(elm_idx_list)[idx_refit].tolist(),
                                                                             n_cpu)
                        elif len(elm_idx_list[idx_refit]) == 1:
                            elm_idx_list_chunks_zero = [elm_idx_list[idx_refit]]
                        else:  # n_cpu == 1:
                            elm_idx_list_chunks_zero = [elm_idx_list[idx_refit]]

                        workhorse_partial = partial(fit_elms,
                                                    e_matrix=e_matrix,
                                                    mep=mep,
                                                    fun=fun,
                                                    zap_idx=zap_idx,
                                                    init_vals=init_vals,
                                                    limits=limits,
                                                    log_scale=log_scale,
                                                    max_nfev=100,
                                                    score_type=score_type,
                                                    verbose=verbose)
                        start = time.time()
                        res_refit = pool.map(workhorse_partial, elm_idx_list_chunks_zero)
                        stop = time.time()

                        if verbose:
                            print(f"Determine c-factors (refit) for {q} / {qoi_path_hdf5}: {stop - start} s")

                        best_values_refit = []
                        c_refit = None
                        for i in range(len(res_refit)):
                            if i == 0:
                                c_refit = res_refit[i][0]
                                best_values_refit = res_refit[i][1]
                            else:
                                if c_refit.ndim < 2:
                                    c_refit = c_refit[:, np.newaxis]
                                if res_refit[i][0].ndim < 2:
                                    c_refit = np.vstack((c_refit, res_refit[i][0][:, np.newaxis]))
                                else:
                                    c_refit = np.vstack((c_refit, res_refit[i][0]))
                                best_values_refit += res_refit[i][1]

                        # overwrite old values with refitted ones if r2/sr score was higher,
                        # keep old value otherwise
                        for i_c_re, c_re in enumerate(c_refit):
                            if c_re < c[idx_refit[i_c_re]]:
                                c[idx_refit[i_c_re]] = c_re
                                best_values[idx_refit[i_c_re]] = best_values_refit[i_c_re]

                        i_refit += 1
                else:
                    break

            # find discontinuities and refit
            ##################################################################
            if refit_discontinuities and len(c) > 1:
                idx_disc, idx_neighbor = pynibs.get_indices_discontinuous_data(data=c,
                                                                               con=con,
                                                                               neighbor=True,
                                                                               deviation_factor=2)
                idx_disc = np.array(idx_disc)

                if len(idx_disc) > 0:
                    if verbose:
                        print(f" > Performing refit for {len(idx_disc)} discontinuous elements ...")

                    # refit for discontinuous elements
                    if len(idx_disc) > 0:
                        # set start values from neighbors
                        for p in params:
                            for idx_re, idx_ne in zip(idx_disc, idx_neighbor):
                                init_vals[p][idx_re] = best_values[idx_ne][p]

                        if n_cpu > 1:
                            elm_idx_list_chunks_disc = pynibs.compute_chunks(np.array(elm_idx_list)[idx_disc].tolist(),
                                                                             n_cpu)
                        elif len(elm_idx_list[idx_disc]) == 1:
                            elm_idx_list_chunks_disc = [elm_idx_list[idx_disc]]
                        else:  # n_cpu == 1:
                            elm_idx_list_chunks_disc = [elm_idx_list[idx_disc]]

                        start = time.time()
                        res_refit = pool.map(workhorse_partial, elm_idx_list_chunks_disc)
                        stop = time.time()

                        if verbose:
                            print(f"Determined c-factors (discontinuous refit) for "
                                  f"{q} / {qoi_path_hdf5}: {stop - start} s")

                        best_values_refit = []
                        c_refit = None
                        for i in range(len(res_refit)):
                            if i == 0:
                                c_refit = res_refit[i][0]
                                best_values_refit = res_refit[i][1]
                            else:
                                if c_refit.ndim < 2:
                                    c_refit = c_refit[:, np.newaxis]
                                if res_refit[i][0].ndim < 2:
                                    c_refit = np.vstack((c_refit, res_refit[i][0][:, np.newaxis]))
                                else:
                                    c_refit = np.vstack((c_refit, res_refit[i][0]))
                                best_values_refit += res_refit[i][1]

                        c[idx_disc] = c_refit

                        for j, i in enumerate(idx_disc):
                            best_values[i] = best_values_refit[j]

        c_all[:, i_q] = c.flatten()

    if local_pool:
        pool.close()
        pool.join()

    if return_fits:
        return c_all, best_values
    else:
        return c_all


def get_model_init_values(fun, elm_idx_list, e_matrix, mep, mask_e_field=None,
                          rem_empty_hints=True):
    """
    Calc appropriate init, limit, and max values for models fits depending on the data. If negative and positive x-data
    is present in case of e.g. normal component values are set according to the side (positive or negative)
    where more values are present. When more positive x-axis values are present, negative x-axis values will be ignored.
    When more negative x-axis values are present, the absolute values will be taken and the positive values are ignored.
    Only parameters for sigmoid* are optimized.

    Parameters
    ----------
    fun : pyfempp.exp.Mep
        IO curve function object.
    elm_idx_list : np.ndarray of int
        (n_used_elms) Array containing the element indices the fit is performed for.
    e_matrix : np.ndarray of float
        (n_zaps, n_ele) Electric field matrix.
    mep : np.ndarray of float
        (n_zaps) Motor evoked potentials for every stimulation.
    mask_e_field : np.ndarray of bool, optional
        (n_zaps, n_ele) Mask indicating for which e-field (and mep) values the fit is performed for.
        Changes for normal component in each element because of the sign and p-values.
        If None, all data is used in each element.
    rem_empty_hints: bool, default: True
        Remove any non-filled param hints from limits dict.

    Returns
    -------
    log_scale : bool
        Log scale.
    limits : dict of list
        (n elm_index_list) Element-wise limit values for function fitting.
    init_vals : dict of list
        (n elm_index_list) Element-wise init values for function fitting.
    max_vals_refit : dict of list
        (n elm_index_list) Element-wise perturbation range for refitting function.
    """
    init_vals = dict()
    max_vals_refit = dict()
    limits = dict()

    if mask_e_field is None:
        mask_e_field = np.ones(e_matrix.shape).astype(bool)

    # get functions-specific argument names
    params = inspect.getfullargspec(fun).args[1:]
    for p in params:
        init_vals[p] = []
        max_vals_refit[p] = []
        limits[p] = []

    if fun == pynibs.expio.fit_funs.linear:
        # linear function starts with generic and same values for each element
        for _ in range(len(elm_idx_list)):
            limits["m"].append([-100, 100])
            limits["n"].append([-100, 100])

            init_vals["m"].append(0.3)
            init_vals["n"].append(-1)

            max_vals_refit["m"].append(100)
            max_vals_refit["n"].append(.3)

    elif fun == pynibs.expio.fit_funs.exp0:
        for _ in range(len(elm_idx_list)):
            limits["x0"].append([0, 1000])
            limits["r"].append([1e-12, 100])

            # init_vals["x0"].append(40)
            init_vals["x0"].append(1)
            init_vals["r"].append(.1)

            # max_vals_refit["x0"].append(100)
            max_vals_refit["x0"].append(10)
            max_vals_refit["r"].append(.2)

    elif fun == pynibs.expio.fit_funs.sigmoid or fun == pynibs.expio.fit_funs.sigmoid_log:
        # limits["x0"] = [[0, 1000] for _ in range(len(elm_idx_list))]
        # limits["amp"] = [[1e-12, 1000] for _ in range(len(elm_idx_list))]
        # limits["r"] = [[1e-12, 100] for _ in range(len(elm_idx_list))]

        for _, elm in enumerate(elm_idx_list):
            e_elm = np.abs(e_matrix[mask_e_field[:, elm], elm])

            # x0 first guess: center between e_min, e_max
            e_min = np.min(e_elm)
            e_max = np.max(e_elm)

            x_0 = e_min + np.max(e_elm) / 2
            init_vals["x0"].append(x_0)

            amp = np.max(mep)
            init_vals["amp"].append(amp)  # largest MEP

            # r0 initial: slope from max_mep/min_mep over middle 50% of e-range
            # r = 8 / np.max(e_elm) * ((np.max(mep) - np.min(mep)) / (e_max - e_min))
            r = (np.max(mep) - np.min(mep)) / (e_max - e_min)
            init_vals["r"].append(r)

            # max_vals_refit["x0"].append(200)
            max_vals_refit["x0"].append(3 * x_0)
            max_vals_refit["amp"].append(3 * amp)
            max_vals_refit["r"].append(3 * r)

        # set upper bound of limits in relation to init_vals
        factor = 100
        limits["x0"] = [[0, init_vals["x0"][i] * 3] for i in range(len(elm_idx_list))]
        limits["amp"] = [[1e-12, init_vals["amp"][i] * 3] for i in range(len(elm_idx_list))]
        limits["r"] = [[1e-12, init_vals["r"][i] * factor] for i in range(len(elm_idx_list))]

    elif fun == pynibs.expio.fit_funs.sigmoid4 or fun == pynibs.expio.fit_funs.sigmoid4_log:
        for _, elm in enumerate(elm_idx_list):
            e_elm = np.abs(e_matrix[mask_e_field[:, elm], elm])

            # x0 first guess: center between e_min, e_max
            e_min = np.min(e_elm)
            e_max = np.max(e_elm)

            x_0 = e_min + np.max(e_elm) / 2
            init_vals["x0"].append(x_0)

            amp = np.max(mep)
            init_vals["amp"].append(amp)  # largest MEP

            # r0 initial: slope from max_mep/min_mep over middle 50% of e-range
            # r = 8 / np.max(e_elm) * ((np.max(mep) - np.min(mep)) / (e_max - e_min))
            r = (np.max(mep) - np.min(mep)) / (e_max - e_min)
            init_vals["r"].append(r)

            # max_vals_refit["x0"].append(200)
            max_vals_refit["x0"].append(3 * x_0)
            max_vals_refit["amp"].append(3 * amp)
            max_vals_refit["r"].append(3 * r)

            init_vals["y0"].append(np.min(mep))

            # if init_vals["y0"][-1] <= 0:
            #     init_vals["y0"][-1] = 1e-3

            max_vals_refit["y0"].append(np.min(mep) * 3)

        # set upper bound of limits in relation to init_vals
        factor = 100
        limits["x0"] = [[0, init_vals["x0"][i] * 3] for i in range(len(elm_idx_list))]
        limits["amp"] = [[1e-12, init_vals["amp"][i] * 3] for i in range(len(elm_idx_list))]
        limits["r"] = [[1e-12, init_vals["r"][i] * factor] for i in range(len(elm_idx_list))]
        y0_init = init_vals["y0"][0]
        limits["y0"] = [[np.abs(y0_init) * -2, np.abs(y0_init) * 2] for _ in range(len(elm_idx_list))]

    elif fun == pynibs.expio.fit_funs.dummy_fun:
        for _ in range(len(elm_idx_list)):
            limits["a"].append([0, 1])
            init_vals["a"].append(1)
            max_vals_refit["a"].append(1)
    else:
        raise NotImplementedError(fun)

    if fun == pynibs.expio.fit_funs.sigmoid_log or fun == pynibs.expio.fit_funs.sigmoid4_log:
        log_scale = True
    else:
        log_scale = False

    if rem_empty_hints:
        # remove params hints for params that are not taken care of here
        keys_2_remove = [k for k, v in limits.items() if not v]
        for k in keys_2_remove:
            del limits[k]
            del init_vals[k]
            del max_vals_refit[k]

    return log_scale, limits, init_vals, max_vals_refit


def nl_hdf5_single_core_write(i, elm_idx_list, fn_reg_hdf5=None, qoi_path_hdf5=None, e_matrix=None,
                              mep=None,
                              fun=pynibs.expio.fit_funs.sigmoid4,
                              con=None, n_refit=50, return_fits=False, constants=None, verbose=False,
                              seed=None, stepdown=False, score_type='R2', return_progress=False, geo=None):
    """
    Perform single-core processing for non-linear optimization and write results to an HDF5 file.

    Parameters
    ----------
    i : int
        The index of the subset of data to process.
    elm_idx_list : list of int
        List of element indices.
    fn_reg_hdf5 : str, optional
        Path to the registration HDF5 file.
    qoi_path_hdf5 : str, optional
        Path to the HDF5 file containing the quantity of interest (QOI) data.
    e_matrix : np.ndarray, optional
        The electromagnetic forward matrix.
    mep : pandas.DataFrame, optional
        The motor evoked potential (MEP) data.
    fun : function, default: sigmoid4
        The non-linear optimization function to use (default: pynibs.sigmoid4).
    con : object, optional
        Constraints for optimization, if applicable.
    n_refit : int, default: 50
        Number of refitting iterations.
    return_fits : bool, default: False
        If True, return fits alongside the coefficients.
    constants : dict, optional
        Constants used in the optimization function.
    verbose : bool, default: False
        If True, print verbose messages.
    seed : int, optional
        The random seed for optimization.
    stepdown : bool, default: False
        If True, use a stepdown approach for optimization.
    score_type : str, default: 'R2'
        The type of score to use for optimization.
    return_progress : bool, default: False
        If True, return progress data.
    geo : object, optional
        Geometry data.

    Returns
    -------
    dict
        A dictionary containing the following elements:

        * 'progress_data': Progress data if 'return_progress' is True.
        * 'best_values': Best optimization values if 'return_fits' is True.
    """
    # get string prefix for useful multiprocessing logging
    str_pref = "Main"
    try:
        str_pref = f"{multiprocessing.current_process()._identity[0]:0>2} "
    except Exception:
        pass

    # compute c
    if stepdown:
        res = stepdown_approach(zap_idx=z[i],
                                elm_idx_list=elm_idx_list,
                                fn_reg_hdf5=fn_reg_hdf5,
                                qoi_path_hdf5=qoi_path_hdf5,
                                e_matrix=e_matrix[z[i]],
                                mep=mep.iloc[z[i]],
                                fun=fun,
                                con=con,
                                n_refit=n_refit, return_fits=return_fits,
                                constants=constants,
                                verbose=verbose, seed=seed,
                                score_type=score_type,
                                return_progress=return_progress,
                                geo=geo)
    else:
        res = nl_hdf5_single_core(zap_idx=z[i],
                                  elm_idx_list=elm_idx_list,
                                  fn_reg_hdf5=fn_reg_hdf5,
                                  qoi_path_hdf5=qoi_path_hdf5,
                                  e_matrix=e_matrix,
                                  mep=mep,
                                  fun=fun,
                                  con=con,
                                  n_refit=n_refit, return_fits=return_fits, constants=constants,
                                  verbose=verbose, seed=seed)

    c = res['c']
    best_values = res['best_values'] if 'best_values' in res else None
    progress_data = res['progress'] if 'progress' in res else {}
    stats = res['stats'] if 'stats' in res else {}

    # write to hdf5, threadsafe
    if verbose:
        print(f"{str_pref}: > Writing results for {len(z[i])} zaps to {fn}.")
    with lock:
        with h5py.File(fn, 'a') as f:
            f.create_dataset(f'c/{len(z[i])}', data=c)
            for key in stats:
                stat = stats[key]
                if not isinstance(stat, np.ndarray):
                    stat = np.array([stat])
                f.create_dataset(f'{key}/{len(z[i])}', data=stat)

    if verbose:
        print(f"{str_pref}: > Writing results for {len(z[i])} zaps done.")

    return {'progress_data': progress_data,
            'best_values': best_values}


def get_bad_elms(x, y, method="lstsq", verbose=False):
    """
    This does an element-wise fast linear regression fit to identify bad elements.
    Bad is defined here as a negative slope.

    x : np.ndarray of float
        (n_zaps, n_ele) Electric field matrix.
    y : np.ndarray of float
        (n_zaps) Motor evoked potentials for every stimulation.
    method : str, default: "lstsq"
        Which method to use. (numpy.linalg.)lstsq, (scipy.stats.)linregress, or pinv
    verbose : bool, default: False
        Indicating verbosity of messages.

    Returns
    -------
    idx: np.ndarray
        Indices of bad elements.
    """
    str_pref = "Main"
    try:
        str_pref = f"{multiprocessing.current_process()._identity[0]:0>2} "
    except:
        pass
    start = time.time()
    if method == "linregress":
        idx = [i for i in range(x.shape[1]) if
               linregress(x[:, i], y).slope < 0]
    elif method == "lstsq":
        idx = [i for i in range(x.shape[1]) if
               lstsq(np.vstack([x[:, i], np.ones(x.shape[0])]).T, y, rcond=None)[0][0] < 0]
    elif method == 'pinv':
        warnings.warn("pinv method is untested.")
        idx = [i for i in range(x.shape[1]) if
               pinv(x[:, i].reshape(x.shape[0], 1)).dot(y)[0] < 0]
    else:
        raise NotImplementedError(f'Method {method} unknown.')

    if verbose:
        print(
                f"Proc{str_pref}:  > {len(idx)}/{x.shape[1]} bad elms removed from fitting ({time.time() - start:2.4} s).")

    return idx


def nl_hdf5_single_core(zap_idx, elm_idx_list, fn_reg_hdf5=None, qoi_path_hdf5=None, e_matrix=None, mep=None,
                        fun=pynibs.expio.fit_funs.sigmoid4,
                        con=None, n_refit=50, return_fits=False, constants=None, verbose=False, seed=None,
                        rem_bad_elms=True, return_e_field_stats=True):
    """
    Mass-univariate nonlinear regressions on raw MEP_{AMP} ~ E.
    That is, for each element in elm_idx_list, it's E (mag | norm | tan) for each zap regressed on the raw MEP
    amplitude. An element wise r2 score is returned.
    The function reads the precomputed array of E-MEP data from an .hdf5 file.

    Parameters
    ----------
    zap_idx : np.ndarray of int
        (n_zaps_used) Indices of zaps the congruence factor is calculated with.
    elm_idx_list : np.ndarray of int
        List containing the element indices the fit is performed for.
    fn_reg_hdf5 : str, optional
        Filename (incl. path) containing the precomputed E-MEP dataframes
    qoi_path_hdf5: str or list of str, optional
        Path in .hdf5 file to dataset of electric field qoi e.g.: ["E", "E_norm", "E_tan"].
    e_matrix : np.ndarray of float, optional
        (n_zaps, n_ele) Electric field matrix.
    mep : np.ndarray of float, optional
        (n_zaps) Motor evoked potentials for every stimulation.
    fun : function object, default: sigmoid4
        A function of pynibs.exp.Mep (exp0, sigmoid).
    con : np.ndarray of float, optional
         (n_roi, 3 or 4) Connectivity matrix of ROI (needed in case of refit because of discontinuity check).
    n_refit : int, default: 50
        Maximum number of refits of zero elements. No refit is applied in case of n_refit = 0.
    return_fits : bool, default: False
        Return fit objects containing the parameter estimates
    constants : dict of <string>:<num>, optional
        key:value pair of model parameters not to optimize.
    verbose : bool, default: False
        Plot output messages.
    seed: int, optional
        Seed to use.
    rem_bad_elms: bool, default: True
        Remove elements based on a fast linear regression slope estimation.
    return_e_field_stats : bool, default: True
        Return some stats on the efield variance

    Returns
    -------
    r2 : np.ndarray of float
        (n_roi, n_qoi) R2 for each element in elm_idx_list.
    """
    starttime = time.time()
    str_pref = "Main"
    try:
        str_pref = f"{multiprocessing.current_process()._identity[0]:0>2} "
    except:
        pass
    best_values = None

    if qoi_path_hdf5 is None:
        qoi_path_hdf5 = ["e_matrix"]

    if fn_reg_hdf5 is not None and qoi_path_hdf5 is not None:
        df_reg = pd.read_hdf(fn_reg_hdf5, qoi_path_hdf5)
        e_matrix = df_reg.values[:, :-1]
        mep = df_reg.loc[:, "qoi_amplitudes"].values

    elif e_matrix is None or mep is None:
        raise ValueError(f"Proc{str_pref}: Please provide e_matrix and mep or fn_reg_hdf5 and qoi_path_hdf5!")

    if n_refit > 0 and con is None:
        raise ValueError(f"Proc{str_pref}: Please provide connectivity matrix (con) in case of refit!")

    # shuffle elements because some of them need longer to compute
    # (in this way it is distributed more equally over all cores)
    if seed:
        np.random.seed(seed)
    np.random.shuffle(elm_idx_list)
    if not isinstance(elm_idx_list, np.ndarray):
        elm_idx_list = np.array(elm_idx_list)
    elm_idx_list_shuffle_idx = np.argsort(elm_idx_list)

    # setting up initial values and limits
    if verbose:
        print(
                f"Proc{str_pref}: > Setting up initial values and limits "
                f"for {fun.__name__} function (from {fun.__module__}).")

    log_scale, limits, init_vals, max_vals_refit = get_model_init_values(fun,
                                                                         elm_idx_list,
                                                                         e_matrix,
                                                                         mep)

    if verbose:
        print(f"Proc{str_pref}: > c-map for {len(zap_idx)}: starting.")

    n_elm = len(elm_idx_list)
    if type(qoi_path_hdf5) is not list:
        qoi_path_hdf5 = [qoi_path_hdf5]

    c_all = np.zeros((len(elm_idx_list), len(qoi_path_hdf5)))

    if rem_bad_elms:
        bad_elm_idx = get_bad_elms(e_matrix[zap_idx], mep.iloc[zap_idx].values, method='lstsq', verbose=True)
    else:
        bad_elm_idx = np.empty((0,))

    # get e field stats
    e_stats_dicts = None
    if return_e_field_stats:
        mc = pynibs.mutual_coherence(e_matrix.transpose())
        _, sv_rat, _ = svd(e_matrix)
        sv_rat = np.max(sv_rat) / np.min(sv_rat)
        e_stats_dicts = {'mc': mc,
                         'sv_rat': sv_rat}

    # fast return for testruns
    if fun == pynibs.expio.fit_funs.dummy_fun:
        c_all = np.random.random(n_elm)
        ret = (c_all,)
        if return_fits:
            best_values = [{"a": 1} for _ in range(n_elm)]
            ret += (best_values,)
        if return_e_field_stats:
            ret += (e_stats_dicts,)
        return ret

    # loop over electric field QOIs
    for i_q, q in enumerate(qoi_path_hdf5):

        start = time.time()
        res = fit_elms(
                elm_idx_list=elm_idx_list,
                e_matrix=e_matrix,
                mep=mep,
                fun=fun,
                zap_idx=zap_idx,
                init_vals=init_vals,
                limits=limits,
                log_scale=log_scale,
                constants=constants,
                max_nfev=10,
                verbose=verbose,
                bad_elm_idx=bad_elm_idx)
        stop = time.time()

        if verbose:
            print(f"Proc{str_pref}:  > Determine c-factors done. ({stop - start:2.2f} s)")

        c = res[0]
        best_values = res[1]

        # resort c values
        c = c[elm_idx_list_shuffle_idx]
        best_values = np.array([best_values[i] for i in elm_idx_list_shuffle_idx])
        elm_idx_list = elm_idx_list[elm_idx_list_shuffle_idx]

        # refit elements
        ####################################################################
        if n_refit > 0:
            params = inspect.getfullargspec(fun).args[1:]

            # refit zero elements
            ####################################################################
            i_refit_zero = 0
            while i_refit_zero < n_refit:
                idx_zero = np.where(c < 1e-6)[0]
                idx_zero = np.setdiff1d(idx_zero, bad_elm_idx)

                # refit zero elements
                if len(idx_zero) > 0:
                    if verbose:
                        print(f"Proc{str_pref}:   > Zero refit {i_refit_zero}/{n_refit} "
                              f"({len(idx_zero)} elements): starting.")

                    # set random start values
                    for p in params:
                        for idx_ze in idx_zero:
                            init_vals[p][idx_ze] = max_vals_refit[p][idx_ze] * np.random.rand()

                    start = time.time()
                    c_refit, best_values_refit = fit_elms(
                            elm_idx_list=elm_idx_list[idx_zero],
                            e_matrix=e_matrix,
                            mep=mep,
                            fun=fun,
                            zap_idx=zap_idx,
                            init_vals=init_vals,
                            limits=limits,
                            log_scale=log_scale,
                            max_nfev=10,
                            verbose=verbose)
                    stop = time.time()

                    if verbose:
                        print(
                                f"Proc{str_pref}:   > Zero refit {i_refit_zero}/{n_refit} "
                                f"({len(idx_zero)} elements): done. "
                                f"({stop - start:.2f} s, "
                                f"{(stop - start) / len(idx_zero):.2f} / elm, "
                                f"{np.sum(c_refit > 1e-6)} > 0 )")

                    # overwrite old values with refitted ones if r2 score was higher, keep old value otherwise
                    for idx_c_ref, c_ref in enumerate(c_refit):
                        if c_ref > c[idx_zero[idx_c_ref]]:
                            # print(f"{idx_zero[idx_c_ref]}: {c[idx_zero[idx_c_ref]]} -> {c_ref}")
                            c[idx_zero[idx_c_ref]] = c_ref
                            best_values[idx_zero[idx_c_ref]] = best_values_refit[idx_c_ref]

                    i_refit_zero += 1

                # find discontinuities and refit

                idx_disc, idx_neighbor = pynibs.get_indices_discontinuous_data(data=c, con=con, neighbor=True,
                                                                               deviation_factor=2, min_val=1e-12,
                                                                               not_fitted_elms=bad_elm_idx)
                idx_disc = np.setdiff1d(idx_disc, bad_elm_idx)

                if len(idx_disc) > 0:

                    if verbose:
                        print(f"Proc{str_pref}:   > Discontinuous refit ({len(idx_disc)} elements): starting.")

                    # set start values from neighbors
                    for p in params:
                        for idx_re, idx_ne in zip(idx_disc, idx_neighbor):
                            init_vals[p][idx_re] = best_values[idx_ne][p]

                    start = time.time()
                    c_refit, best_values_refit = fit_elms(
                            elm_idx_list=elm_idx_list[idx_disc],
                            e_matrix=e_matrix,
                            mep=mep,
                            fun=fun,
                            zap_idx=zap_idx,
                            init_vals=init_vals,
                            limits=limits,
                            log_scale=log_scale,
                            max_nfev=10,
                            verbose=verbose)

                    stop = time.time()

                    if verbose:
                        print(
                                f"Proc{str_pref}:   > Discontinuous refit ({len(idx_disc)} elements): "
                                f"done ({stop - start:.2f} s)")

                    # overwrite old values with refitted ones
                    c[idx_disc] = c_refit
                    best_values[idx_disc] = best_values_refit

        c_all[:, i_q] = c.flatten()

    endtime = time.time()
    print(f"Proc{str_pref}:  > c-map for {len(zap_idx)} zaps done ({endtime - starttime:2.2f}s). ")

    ret = (c_all,)
    if return_fits:
        best_values = [{"a": 1} for _ in range(n_elm)]
        ret += (best_values,)
    if return_e_field_stats:
        ret += (e_stats_dicts,)
    return ret


def stepdown_approach(zap_idx, elm_idx_list, fn_reg_hdf5=None, qoi_path_hdf5=None, e_matrix=None,
                      mep=None,
                      fun=pynibs.expio.fit_funs.sigmoid4,
                      con=None, n_refit=50, return_fits=False, constants=None, verbose=False,
                      seed=None,
                      rem_bad_elms=True, return_e_field_stats=True,
                      score_type='R2', return_progress=False, smooth_data=True, geo=None):
    """
    Mass-univariate nonlinear regressions on raw MEP_{AMP} ~ E in a stepdown manner to speed up computation.

    Initially, one set of fits is done for the complete dataset. Afterwards, the best 1% of the elements are used
    as initial fitting parameters for their neighboring elements. Then, neighboring elements are fitted accordingly
    and iteratively.
    Finally, discontinuous elements are refitted until a smooth map is found or n_refit is hit.
    Can be sped up with rem_bad_elms that computes a fast linear fit to identify elements with a negative slope.
    The function reads the precomputed array of E-MEP data from an .hdf5 file.

    Parameters
    ----------
    elm_idx_list : np.ndarray of int
        List containing the element indices the fit is performed for.
    fn_reg_hdf5 : str
        Filename (incl. path) containing the precomputed E-MEP dataframes.
    qoi_path_hdf5: str or list of str, optional
        Path in .hdf5 file to dataset of electric field qoi e.g.: ["E", "E_norm", "E_tan"].
    e_matrix : np.ndarray of float
        (n_zaps, n_ele) Electric field matrix.
    mep : np.ndarray of float
        (n_zaps) Motor evoked potentials for every stimulation.
    zap_idx : np.array, optional
        (n_zaps) Indices of zaps the congruence factor is calculated with (default: all).
    fun : function object
        A function of pynibs.exp.Mep (exp0, sigmoid).
    con : np.ndarray of float, optional
        (n_elm_roi, 3 or 4) Connectivity matrix of ROI (needed in case of refit because of discontinuity check)
    n_refit : int, default: 50
        Maximum number of refits of zero elements. No refit is applied in case of n_refit = 0.
    return_fits : bool, default: False
        Return fit objects containing the parameter estimates
    constants : dict of <string>:<num>, optional
        key:value pair of model parameters not to optimize.
    verbose : bool, default: False
        Plot output messages.
    seed: int, optional
        Seed to use.
    rem_bad_elms: bool, default: True
        Remove elements based on a fast linear regression slope estimation.
    return_e_field_stats : bool, default: True
        Return some stats on the efield variance
    score_type : str, default: "R2"
        Error measure of fit:

        * "R2": R2 score (Model variance / Total variance); linear fits: [0, 1], 1 ... perfect fit
        * "SR": Relative standard error of regression (1 - Error 2-norm / Data 2-norm); [-Inf, 1], 1 ... perfect fit
        * "rho": Spearman correlation coefficient [-1, 1]; finds any monotonous correlation (0 means no correlation)
    return_progress : bool, default: False
        Return c maps for all steps to allow visualization over e-fitting over timesteps.
    smooth_data : bool, default: False
        Smooth c-map as final step.
    geo : object, optional
        Geometry data.

    Returns
    -------
    dict:

        * r2 : np.ndarray of float
            (n_roi, n_qoi) R2 for each element in elm_idx_list.
        * best_values: list of dict
            Fit information, if wanted.
        * stats : dict
            If wanted:
            'mc': float
                Mutual coherence for e fields.
            'sv_rat' : float
                SVD singular value ratio.
        * progress : cmaps for each step.
    """
    starttime = time.time()

    # get string prefix for useful multiprocessing logging
    str_pref = "Main"
    try:
        str_pref = f"{multiprocessing.current_process()._identity[0]:0>2} "
    except:
        pass

    if qoi_path_hdf5 is None:
        qoi_path_hdf5 = ["e_matrix"]

    if fn_reg_hdf5 is not None and qoi_path_hdf5 is not None:
        df_reg = pd.read_hdf(fn_reg_hdf5, qoi_path_hdf5)
        e_matrix = df_reg.values[:, :-1]
        mep = df_reg.loc[:, "qoi_amplitudes"].values

    elif e_matrix is None or mep is None:
        raise ValueError(f"Proc{str_pref}: Please provide e_matrix and mep or fn_reg_hdf5 and qoi_path_hdf5!")

    if n_refit > 0 and con is None:
        raise ValueError(f"Proc{str_pref}: Please provide connectivity matrix (con) in case of refit!")

    if not isinstance(elm_idx_list, np.ndarray):
        elm_idx_list = np.array(elm_idx_list)

    # setting up initial values and limits
    if verbose:
        print(
                f"Proc{str_pref}: > Setting up initial values and limits "
                f"for {fun.__name__} function (from {fun.__module__}).")

    log_scale, limits, init_vals, max_vals_refit = get_model_init_values(fun,
                                                                         elm_idx_list,
                                                                         e_matrix,
                                                                         mep)
    # convert dict-of-lists to dict-of-dicts

    init_vals = pd.DataFrame().from_dict(init_vals)
    limits = pd.DataFrame().from_dict(limits)

    n_elm = len(elm_idx_list)
    if type(qoi_path_hdf5) is not list:
        qoi_path_hdf5 = [qoi_path_hdf5]
    c_all = np.zeros((n_elm, len(qoi_path_hdf5)))

    if e_matrix.shape[0] != len(zap_idx):
        e_matrix = e_matrix[zap_idx]
        mep = mep[zap_idx]

    # get bad elms by checking their linear slope fit
    if rem_bad_elms:
        bad_elm_idx = set(get_bad_elms(e_matrix[:, elm_idx_list], mep.values, method='lstsq', verbose=verbose))
    else:
        bad_elm_idx = set()

    # get e field stats
    e_stats_dicts = None
    if return_e_field_stats:
        stats_start = time.time()
        mc = pynibs.mutual_coherence(e_matrix.transpose())

        sv_rat = svd(e_matrix, check_finite=False, compute_uv=False)
        sv_rat = np.max(sv_rat) / np.min(sv_rat)
        e_stats_dicts = {'mc': mc,
                         'sv_rat': sv_rat}
        stats_end = time.time()
        if verbose:
            print(f"Proc{str_pref} : > Efield stats computation took {stats_end - stats_start:2.2f} s.")

    # fast return for testruns
    if fun == pynibs.expio.fit_funs.dummy_fun:
        c_all = np.random.random(n_elm)
        ret = (c_all,)
        if return_fits:
            best_values = [{"a": 1} for _ in range(n_elm)]
            ret += (best_values,)
        if return_e_field_stats:
            ret += (e_stats_dicts,)
        return ret

    if verbose:
        print(f"Proc{str_pref}: > c-map for {e_matrix.shape[0]}: starting.")
    if return_progress:
        return_progress_lst = []
    # loop over electric field QOIs
    for i_q, q in enumerate(qoi_path_hdf5):

        start = time.time()
        # get initial c-map
        start_sample = np.random.choice(elm_idx_list, int(.1 * len(elm_idx_list)), replace=False)
        # c, best_values = np.zeros(elm_idx_list.shape), np.zeros(elm_idx_list.shape).astype(object)

        c, best_values = np.zeros(elm_idx_list.shape), init_vals.copy()
        if return_progress:
            return_progress_lst.append(c.copy())
        c[start_sample], best_values.iloc[start_sample] = fit_elms(
                elm_idx_list=start_sample,
                e_matrix=e_matrix,
                mep=mep,
                fun=fun,
                init_vals=init_vals,
                limits=limits,
                log_scale=log_scale,
                constants=constants,
                max_nfev=10,
                verbose=verbose,
                bad_elm_idx=bad_elm_idx,
                score_type=score_type)
        if return_progress:
            return_progress_lst.append(c.copy())
        stop = time.time()

        if verbose:
            print(f"Proc{str_pref}:  > Initial c-factor map done. ({stop - start:2.2f} s)")

        # pick the best elements and compute their neibors' fits
        n_top = int(len(start_sample) * .2)

        elms_done = start_sample[np.argpartition(c[start_sample], -n_top)[-n_top:]]
        mask = np.ones(c.shape, np.bool)
        mask[elms_done] = 0
        c[mask] = 0
        if return_progress:
            return_progress_lst.append(c.copy())
        elms_done = elms_done[c[elms_done] > 0]
        elms_seed = elms_done.copy()
        elms_done = set(elms_done)

        params = inspect.getfullargspec(fun).args[1:]
        i_step = 0
        start = time.time()
        while True:
            # reorder to highest c first
            elms_seed = elms_seed[np.argsort(-c[elms_seed])]
            elm_to_compute = set()

            # compute last 5% in one batch
            if geo is not None and len(elms_done) > 0.95 * len(elm_idx_list):
                elm_to_compute = list(set(elm_idx_list).difference(elms_done))
                print(f"last 5% {len(elms_done)} -> {len(elm_to_compute)}")
                if not elm_to_compute:
                    break

                for elm in elm_to_compute:  # elm = list(elm_to_compute)[0]
                    # use nearest element's params to fit this element
                    nearest_idx = np.argmin(np.linalg.norm(np.mean(geo[con][c > 0], axis=1) -
                                                           np.mean(geo[con[elm]], axis=0), axis=1))
                    nearest_idx = np.argwhere(np.sum(con == con[c > 0][nearest_idx], axis=1) == 3)[0][0]
                    for p in params:
                        init_vals.iloc[elm][p] = best_values.iloc[nearest_idx][p]

            else:
                # increase spread with increasing steps
                for _ in range(i_step + 1):
                    for elm_done_i in elms_seed:
                        # find neighbors for the already computed element
                        mask = np.sum(np.isin(con, con[elm_done_i, :]), axis=1)
                        neighbors = set(np.where((0 < mask) & (mask < 3))[0])

                        # remove elms we already have computed or don't want to compute
                        neighbors.difference_update(elms_done, bad_elm_idx, elm_to_compute)
                        neighbors.intersection_update(elm_idx_list)
                        if not neighbors:
                            continue
                        elm_to_compute.update(neighbors)
                        for neigh in neighbors:
                            # use params of done neighbors as init vals for their neighbors
                            for p in params:
                                init_vals[p][neigh] = best_values.iloc[elm_done_i][p]
                            best_values.iloc[neigh] = best_values.iloc[elm_done_i]

                    print(f"{i_step:0>3}: Seed {len(elms_seed)} -> {len(elm_to_compute)}")

                    # refit already here some elements
                    not_fitted_elms = c[np.array(list(elms_done))] <= 0
                    # get random neighbor to refit
                    idx_disc, idx_neighbor = pynibs.get_indices_discontinuous_data(data=c[np.array(list(elms_done))],
                                                                                   con=con[np.array(list(elms_done))],
                                                                                   neighbor=True,
                                                                                   deviation_factor=3, min_val=1e-12,
                                                                                   not_fitted_elms=not_fitted_elms,
                                                                                   crit='randmax',
                                                                                   neigh_style='point')
                    idx_disc = np.setdiff1d(idx_disc, bad_elm_idx)
                    # set init values from neighbors best values
                    for p in params:
                        for idx_refit, idx_neighbor in zip(idx_disc, idx_neighbor):
                            init_vals[p][idx_refit] = best_values.iloc[idx_neighbor][p]
                    elm_to_compute.update(idx_disc)
                    # save the elms that will be computed now for the next iteration as seed elements
                    elms_seed = elm_to_compute.copy()
                elm_to_compute = list(elm_to_compute)
                elms_seed = np.array(list(elms_seed))

            if not elm_to_compute:
                # end loop
                break

            # start = time.time()
            c[elm_to_compute], best_values.iloc[elm_to_compute] = fit_elms(
                    elm_idx_list=elm_to_compute,
                    e_matrix=e_matrix,
                    mep=mep,
                    fun=fun,
                    init_vals=init_vals,
                    limits=limits,
                    log_scale=log_scale,
                    constants=constants,
                    max_nfev=1,
                    verbose=verbose,
                    score_type=score_type)
            if return_progress:
                return_progress_lst.append(c.copy())
            i_step += 1
            elms_done.update(elm_to_compute)
        stop = time.time()

        if verbose:
            print(
                    f"Proc{str_pref}:   > Stepdown iterations {i_step:0>2} "
                    f"({len(elms_done)}/{n_elm - len(bad_elm_idx)} ({len(bad_elm_idx)} excluded) elements): done. "
                    f"({stop - start:.2f} s "
                    f"{np.sum(c[elm_to_compute] > 1e-6)} > 0)")

        # refit the discontinues elemens
        not_fitted_elms = c == 0
        idx_disc, idx_neighbor = pynibs.get_indices_discontinuous_data(data=c, con=con[elm_idx_list], neighbor=True,
                                                                       deviation_factor=3, min_val=1e-12,
                                                                       not_fitted_elms=not_fitted_elms, crit='max',
                                                                       neigh_style='point')
        idx_disc = np.setdiff1d(idx_disc, bad_elm_idx)

        disc_fit_i, last_disc_idx = 1, idx_disc
        # only repeat until no improvement
        while 0 < len(idx_disc) and disc_fit_i <= n_refit:

            # set init values from neighbors best values
            for p in params:
                for idx_refit, idx_neighbor in zip(idx_disc, idx_neighbor):
                    init_vals[p][idx_refit] = best_values.iloc[idx_neighbor][p]

            start = time.time()
            c[idx_disc], best_values.iloc[idx_disc] = fit_elms(
                    elm_idx_list=elm_idx_list[idx_disc],
                    e_matrix=e_matrix,
                    mep=mep,
                    fun=fun,
                    init_vals=init_vals,
                    limits=limits,
                    log_scale=log_scale,
                    constants=constants,
                    max_nfev=int(np.log2(disc_fit_i)),  # inrease number of nfev each time
                    verbose=verbose,
                    score_type=score_type)
            if return_progress:
                return_progress_lst.append(c.copy())

            # prepare the next iteration
            last_disc_n = len(idx_disc)
            idx_disc, idx_neighbor = pynibs.get_indices_discontinuous_data(data=c, con=con[elm_idx_list], neighbor=True,
                                                                           deviation_factor=3, min_val=1e-12,
                                                                           not_fitted_elms=not_fitted_elms, crit='max',
                                                                           neigh_style='point')
            idx_disc = np.setdiff1d(idx_disc, bad_elm_idx)
            if np.all(last_disc_idx == idx_disc):
                break
            stop = time.time()

            if verbose:
                print(
                        f"Proc{str_pref}:   > Discontinuous refit {disc_fit_i:0>2} ({last_disc_n} elements): "
                        f"done ({stop - start:.2f} s)")
            disc_fit_i += 1

        # smooth data
        if smooth_data:
            smoothed_c = np.zeros_like(c)
            for i, c_dat in np.ndenumerate(c):
                mask = np.sum(np.isin(con, con[i, :]), axis=1)
                mask = np.logical_and(0 < mask, mask < 3)
                # if c_dat <= 0:
                if np.sum(c[mask] == 0) != 0:
                    # if all elements are < 0
                    smoothed_c[i] = 0.3 * np.mean(c[mask][c[mask] == 0]) + 0.7 * c_dat
                else:
                    # otherwise just use good elements
                    smoothed_c[i] = c_dat
                # else:
                # don't touch element if it's the only one with positive values
                # if len(c[mask][c[mask] > 0]) > 1:
                #     smoothed_c[i] = 0.3 * np.mean(c[mask][c[mask] > 0]) + 0.7 * c_dat
                # else:
                #     smoothed_c[i] = c_dat
            c = smoothed_c
        if return_progress:
            return_progress_lst.append(c.copy())
        c_all[:, i_q] = c.flatten()

    endtime = time.time()
    print(f"Proc{str_pref}:  > c-map for {len(zap_idx)} zaps done ({endtime - starttime:2.2f}s). ")
    if verbose:
        bad_elm_idx = list(bad_elm_idx)
        print(f"Proc{str_pref}:  >  min | med | max: "
              f"{np.min(np.delete(c, bad_elm_idx)):2.2f} | "
              f"{np.median(np.delete(c, bad_elm_idx)):2.2f} | "
              f"{np.max(np.delete(c, bad_elm_idx)):2.2f}. "
              f"{np.sum(c > 0)} elms > 0 | "
              f"{len(c)} n_elms")
    ret = {'c': c_all}
    if return_fits:
        ret['best_values'] = best_values
    if return_e_field_stats:
        ret['stats'] = e_stats_dicts
    if return_progress:
        ret['progress'] = return_progress_lst
    return ret


def sing_elm_fitted(elm_idx_list, mep_lst, mep_params, e, alpha=1000, n_samples=100):
    """
    Mass-univariate ridge regressions on fitted MEP_{AMP} ~ E.
    That is, for each element in elm_idx_list, it's E (mag | norm | tan) for each zap regressed on the raw MEP
    amplitude. An element wise sklearn.metrics.regression.r2_score is returned.

    Parameters
    ----------
        elm_idx_list : np.ndarray
            (n_used_ele) List of element indices, the congruence factor is computed for.
        mep_lst : list of Mep object instances
            (n_conds) List of fitted Mep object instances for all conditions (see exp.py for more information).
        mep_params : np.ndarray of float
            (n_mep_params_total) List of all mep parameters of curve fits used to calculate the MEP accumulated into
            one array.(e.g.: [mep_#1_para_#1, mep_#1_para_#2, mep_#1_para_#3, mep_#2_para_#1, mep_#2_para_#1, ...])
        e : np.ndarray of float
            (n_elm, n_cond, n_qoi) Electric field to compute the r2 factor for, e.g. (e_mag, e_norm, e_tan).
        n_samples : int, default=100
            Number of data points to generate discrete mep and e curves.

    Returns
    -------
    r2 : np.ndarray of float
        (n_roi, n_datasets) R^2 for each element in elm_idx_list.
    """

    def cartesian_product(*arrays):
        """
        Fast implementation to get cartesian product of two arrays.

        cartesian_product([a,b,c],[2,3]) =
             [a, 2
              a, 3
              b, 2
              b, 3
              c, 2
              c, 3]
        """
        la = len(arrays)
        dtype = np.result_type(*arrays)
        arr = np.empty([len(a) for a in arrays] + [la], dtype=dtype)
        for i, a in enumerate(np.ix_(*arrays)):
            arr[..., i] = a
        return arr.reshape(-1, la)

    n_eqoi = e.shape[2]
    n_cond = e.shape[1]
    n_elm = e.shape[0]
    assert n_cond == len(mep_lst)
    mep_params = np.array(mep_params).flatten()
    mep_params_cond = []
    start_idx = 0
    for i_cond in range(n_cond):
        mep_params_cond.append(mep_params[start_idx:(start_idx + mep_lst[i_cond].popt.size)])
        start_idx = start_idx + mep_lst[i_cond].popt.size

    del start_idx

    intensities = []
    amplitudes = []

    reg_r2 = np.empty((n_elm, n_eqoi))

    # get amplitudes from fitted meps
    for i_cond in range(n_cond):
        intensities.append(np.linspace(mep_lst[i_cond].x_limits[0], mep_lst[i_cond].x_limits[1], n_samples))
        amplitudes.append(mep_lst[i_cond].eval(intensities[-1], mep_params_cond[i_cond]))
        # intensities_min.append(mep_lst[i_cond].x_limits[0])
        # intensities_max.append(mep_lst[i_cond].x_limits[1])
    amplitudes = np.array(amplitudes).flatten()

    for qoi_idx in range(n_eqoi):
        # t_q = time.time()
        x = pd.DataFrame()
        index_shift = 0

        for mep_i, mep_lst in enumerate(mep_lst):
            # condition wise, as we stimulated with different intensities per conditions

            # for each element in roi, one datapoint for each zap.
            current = cartesian_product(e[:, mep_i, qoi_idx], intensities[mep_i])

            # index is iteration of zaps over all conditions
            index = cartesian_product(e[:, mep_i, qoi_idx], np.arange(n_samples))[:, 1]
            index += index_shift
            index_shift = index[-1] + 1

            # el is range(n_elements) * n_zaps_in_condition
            el_idx = np.repeat(np.arange(e.shape[0]), n_samples)

            # intensity * e
            e_zap = np.multiply(current[:, 0], current[:, 1])

            # put all together
            x_cond = pd.DataFrame(data={"index": index.astype(int),
                                        "el": el_idx,
                                        "e": e_zap})
            # amplitudes = np.append(amplitudes, mep_lst.mep)
            # "current": current[:, 1],
            # "mep": mep[:,1]})

            # x_cond['condition'] = mep_i
            x = x.append(x_cond)
        # x.shape is now (n_zaps*n_elms, 3)

        # reshape to (n_zaps, n_elms)
        x = x.pivot(index="index", columns="el", values="e")  # this is pretty slow

        do_reg_poly = False
        do_reg_linear = True

        # reg = Ridge(alpha=alpha)
        reg = LinearRegression(normalize=True)
        if do_reg_poly:
            raise NotImplementedError

        elif do_reg_linear:
            # Do one regression per element.
            r_t = time.time()

            def get_score(x_i):
                """Helper function do be used by pd.apply() to speed up things.

                Parameters
                ----------
                x_i: pd.Series
                    Column with e for a single elm.

                Returns
                -------
                r2 for amplitudes ~ E
                """
                # x_i = x_i.reshape(-1, 1)
                reg.fit(x_i.reshape(-1, 1), amplitudes)
                return reg.score(x_i.reshape(-1, 1), amplitudes)

            scores = x.apply(get_score, axis=0, raw=True)
            # coefs = [get_score(x.iloc[:, el_idx]) for el_idx in range(n_elm)]
            print("all_reg: {}".format(time.time() - r_t))

        reg_r2[:, qoi_idx] = np.array(scores)
        # data.append(data_qoi_tmp)

        # print "qoi {}: {}".format(qoi_idx, time.time() - t_q)

    return reg_r2


def logistic_regression():
    """
    Some ideas on how to improve regression approach

    1. De-log data

    Data range has to be transformed to a reasonable range. For a full sigmoid, -10:10 looks ok

    .. code-block:: python

       sig <- function(z) {
         return( 1 / (1 + exp(-z)))
       }

    .. code-block:: python

       desig <- function(x) {
         return(- log((1/x) - 1))
       }

    This might be a reasonable fast approach, but the parameter range has to be estimated. Maybe remove some outliters?

    2. fit logistic regression to raw data
    scipy.optimize provides fit_curve(), which does OLS-ish fitting to a given function
    https://stackoverflow.com/questions/54376900/fit-sigmoid-curve-in-python

    I expect this to be rather slow.

    3. Use the sklearn logistic_regression classifyer and access raw fit data
    The logistic_regression is implemented as a classifyer, maybe it's possible to use its regression fit results.
    Implementation should be pretty fast.
    """
    raise NotImplementedError


def init(l, zap_lists, res_fn):
    """
    Pool init function to use with regression_nl_hdf5_single_core_write().

    Parameters
    ----------
    l : multiprocessing.Lock()
    zap_lists : list of list of int
        Which zaps to compute.
    res_fn : str
        .hdf5 fn
    """

    global lock, z, fn
    lock = l
    z = zap_lists
    fn = res_fn


def single_fit(x, y, fun):
    """
    Performs a single fit and returns fit object.

    Parameters
    ----------
    x : ndarray of float
        x-values.
    y : ndarray of float
        y-values.
    fun : function
        Function for fitting.

    Returns
    -------
    fit : gmodel fit object
        Fit object.
    """
    params = inspect.getfullargspec(fun).args[1:]

    limits = dict()
    init_vals = dict()
    log_scale = False
    if fun == pynibs.expio.fit_funs.linear:
        log_scale = False
        limits["m"] = [-100, 100]
        limits["n"] = [-100, 100]

        init_vals["m"] = 0.3
        init_vals["n"] = -1

    elif fun == pynibs.expio.fit_funs.linear_log:
        log_scale = True
        limits["m"] = [-100, 100]
        limits["n"] = [-100, 100]

        init_vals["m"] = 0.3
        init_vals["n"] = -1

    elif fun == pynibs.expio.fit_funs.exp0:
        log_scale = False
        limits["x0"] = [0, 1000]
        limits["r"] = [1e-12, 5]

        init_vals["x0"] = np.mean(x)  # 40
        init_vals["r"] = 10 / np.max(x)

    elif fun == pynibs.expio.fit_funs.exp:
        log_scale = False
        limits["x0"] = [0, 1000]
        limits["r"] = [1e-12, 5]
        limits['y0'] = [-.1, 5]

        init_vals["x0"] = np.mean(x)  # 40
        init_vals["r"] = 10 / np.max(x)
        init_vals["y0"] = 0

    elif fun == pynibs.expio.fit_funs.exp0_log:
        log_scale = True
        limits["x0"] = [1e-12, 1000]
        limits["r"] = [1e-12, 5]
        limits['y0'] = [-.1, 5]

        init_vals["x0"] = np.mean(x)  # 40
        init_vals["r"] = 10 / np.max(x)
        init_vals["y0"] = 0

    elif fun == pynibs.expio.fit_funs.exp_log:
        log_scale = True
        limits["x0"] = [1e-12, 1000]
        limits["r"] = [1e-12, 5]
        limits['y0'] = [-.1, 5]

        init_vals["x0"] = np.mean(x)  # 40
        init_vals["r"] = 10 / np.max(x)
        init_vals["y0"] = 0

    elif fun == pynibs.expio.fit_funs.sigmoid:
        log_scale = False

        limits["x0"] = [0, 1000]
        limits["amp"] = [1e-12, 1000]
        limits["r"] = [1e-12, 100]

        init_vals["x0"] = np.mean(x)  # 70
        init_vals["amp"] = np.max(y)
        init_vals["r"] = 10 / np.max(x)

    elif fun == pynibs.expio.fit_funs.sigmoid4:
        log_scale = False

        limits["x0"] = [0, 1000]
        limits["amp"] = [1e-12, 1000]
        limits["r"] = [1e-12, 100]
        limits["y0"] = [1e-12, 10]

        init_vals["x0"] = np.mean(x)  # 70
        init_vals["amp"] = np.max(y)
        init_vals["r"] = 10 / np.max(x)
        init_vals["y0"] = 1e-2

    elif fun == pynibs.expio.fit_funs.sigmoid_log:
        log_scale = True

        limits["x0"] = [0, 1000]
        limits["amp"] = [1e-12, 1000]
        limits["r"] = [1e-12, 100]

        init_vals["x0"] = np.mean(x)  # 50
        init_vals["amp"] = np.max(y)
        init_vals["r"] = 10 / np.max(x)

    elif fun == pynibs.expio.fit_funs.sigmoid4_log:
        log_scale = True

        limits["x0"] = [0, 1000]
        limits["amp"] = [1e-12, 100]
        limits["r"] = [1e-12, 10]
        limits["y0"] = [1e-6, 1e-1]

        init_vals["x0"] = np.mean(x)  # 50
        init_vals["amp"] = np.max(y)
        init_vals["r"] = 10 / np.max(x)
        init_vals["y0"] = 1e-2
    else:
        raise NotImplementedError(f"Function {fun} not implemented.")
    if log_scale:
        y = np.log10(y)

    # set up gmodel
    gmodel = Model(fun)

    for p in params:
        gmodel.set_param_hint(p, value=init_vals[p], min=limits[p][0], max=limits[p][1])

    gmodel.make_params()

    # perform fit
    fit = gmodel.fit(y, x=x)

    return fit
