"""
Functions to optimize single channel currents for multichannel TMS arrays.
"""
import h5py
import math
import tqdm
import scipy
import warnings
import itertools
import numpy as np
import pandas as pd
from functools import partial
from scipy.optimize import minimize, basinhopping
from tqdm.contrib.concurrent import process_map  # or thread_map
try:
    import terminalplot
except ImportError:
    pass


def get_score_raw(x, e, n_stim, n_ele, n_channel,
                  x_opt=None, opt_target='elms', fun=None, fun_args=None, res_scale=.001,
                  metric='mean',
                  verbose=False):
    """
    Compute score for e-efield cross correlations.
    Non-normalized score is returned, so you need to do sth like

    score = 1 / ((n_ele ** 2 - n_ele) / 2) * get_score()

    Parameters
    ----------
    x : np.ndarray of float
        (n_channel * n_stim_opt, ) Vector to scale channels for each.
    e : np.ndarray of float
        (n_ele*3, n_channels) E-field.
    n_stim : int
        Number of stimulations to compute score for.
    n_ele : int
        Number of elements.
    n_channel : int
        Number of channels.
    x_opt :  np.ndarray of float, optinonal
        (n_pre_opt,) Previously optimized channel currents.
    opt_target : str, default: 'elms'
        Optimization target. 'elms' for optimizing decorrelations of elements, 'stims' for stimulations.
    fun : function, optional
        Function to scale the e-field.
    fun_args : tuple, optional
        Arguments for the function.
    res_scale : float, default: 1
        Scaling factor for the result.
    verbose : book, default: False
        Print out progress.

    Returns
    -------
    score : float
        non-normalized score np.nansum(np.abs(np.triu(np.corrcoef(e_mag), k=1)))
    """
    if x_opt is None:
        currents_all_zaps_channels = np.reshape(x, (n_channel, n_stim))
    else:
        currents_all_zaps_channels = np.hstack((x_opt, np.reshape(x, (n_channel, n_stim - x_opt.shape[1]))))
    # currents_all_zaps_channels.shape = (n_chans, n_zaps)

    # determine total electric field (vector form)
    e_vec = np.matmul(e, currents_all_zaps_channels)  # e_vec.shape = (n_elems*3, n_zaps)
    # determine magnitude
    # for e_ in e_vec.T:
    #     a = np.linalg.norm(np.reshape(e_, (n_ele, 3)), axis=1)
    #     # e_.shape = (n_ele=3,)
    #     # np.reshape(e_, (n_ele, 3)).shape = (n_ele,3)
    #     # a.shape = (n_ele,)
    if e.ndim == 2:
        e_mag = np.vstack([np.linalg.norm(np.reshape(e_, (n_ele, 3)), axis=1) for e_ in e_vec.T]).T
    elif e.ndim == 3:
        e_mag = np.linalg.norm(e_vec, axis=1)
    else:
        raise ValueError
    # e_mag.shape = (n_ele, n_zaps)

    # scale e_mag if necessary
    if fun is not None:
        if fun_args is None:
            fun_args = {}
        e_mag = fun(e_mag, **fun_args)

    # determine average correlation coefficient
    # r_avg = 1/((n_ele**2-n_ele)/2) * np.sum(np.abs(np.triu(np.corrcoef(e_mag), k=1)))

    if opt_target == 'elms':
        # r_sum = np.nansum(np.abs(np.triu(np.corrcoef(e_mag), k=1)))
        # r_sum = np.sum(np.abs(np.triu(np.corrcoef(e_mag), k=1)))
        # a = np.corrcoef(e_mag)
        # np.mean(a)
        # np.mean(np.triu(a, k=0))
        # c = np.tril(a, k=-1)
        # np.tril(a, k=-1).sum() == np.triu(a, k=1).sum()
        # np.mean(np.tril(a, k=0))
        if metric == "mean":
            r_sum = np.sum(np.triu(np.corrcoef(e_mag), k=1))
        elif metric =="max":
            r_sum = np.max(np.triu(np.abs(np.corrcoef(e_mag)), k=1))
        else:
            raise ValueError

        if np.isnan(r_sum):
            r_sum = n_ele**2
    elif opt_target == 'stims':
        # r_sum = np.nansum(np.abs(np.triu(np.corrcoef(e_mag.T), k=1)))
        # r_sum = np.sum(np.abs(np.triu(np.corrcoef(e_mag.T), k=1)))
        if metric == "mean":
            r_sum = np.sum(np.triu(np.corrcoef(e_mag.T), k=1))
        elif metric == "max":
            r_sum = np.max(np.triu(np.abs(np.corrcoef(e_mag.T)), k=1))
        else:
            raise ValueError
        if np.isnan(r_sum):
            r_sum = n_stim**2
    else:
        raise ValueError("opt_target has to be 'elms' or 'stims'")

    if verbose and np.random.randint(0,1000) > 998:
        bins=200
        y_hist, x_hist = np.histogram(e_mag, bins=bins)
        # y_hist, x_hist = np.histogram(x, bins=bins)
        terminalplot.plot(list(x_hist), list(y_hist), rows=50, columns=80)
        # currents_all_zaps_channels = np.reshape(x, (n_channel, n_stim))
        # didt_constr_res = (np.sum(currents_all_zaps_channels ** 2, axis=0) < 2*0.8**2).sum() - n_stim
        didt_constr_res = strain_constr(x, 0.64, 5, 20)
        # print(f"")
        print(f"didt_constr: {didt_constr_res}, score: {r_sum / res_scale}")
    return r_sum / res_scale


def strain_constr(x, max_strain, n_channel, n_stim):
    """
    Compute mechanical strain on 5 channel planar mTMS coil
    """
    # Ensure weights is a numpy array and take the absolute value
    currents = np.reshape(x, (n_stim, n_channel)).T
    pulse_strain = calculate_strain(currents)

    return -(pulse_strain > max_strain).sum()


def strain_constr_min(x, max_strain, n_channel, n_stim):
    """
    Compute mechanical strain on 5 channel planar mTMS coil
    """
    # Ensure weights is a numpy array and take the absolute value
    currents = np.reshape(x, (n_stim, n_channel)).T
    pulse_strain = calculate_strain(currents)

    return -(pulse_strain < max_strain).sum()


def calculate_strain(currents):
    """
    Computes planar mTMS strain from single channel currents.

    Written by Mikael Laine, 2025.

    Parameters
    ----------
    currents : list

    Returns
    -------
    pulse_strain : float

    """
    # Ensure weights is a numpy array and take the absolute value
    currents = np.abs(np.array(currents))

    # Calculate coil 5 strain
    term1_5 = np.sqrt(0.9 * currents[0] ** 2 + currents[1] ** 2) * 0.2
    term2_5 = np.sqrt(0.9 * currents[2] ** 2 + currents[3] ** 2)
    coil_5_strain = (term1_5 + term2_5) * currents[4]

    # Calculate coil 1 strain
    term1_1 = 0.9 * currents[1]
    term2_1 = 0.7 * np.sqrt(currents[2] ** 2 + 0.9 * currents[3] ** 2)
    term3_1 = 0.2 * currents[4]
    coil_1_strain = (term1_1 + term2_1 + term3_1) * currents[0]

    # Find the maximum of the two strain values
    pulse_strain = np.maximum(coil_5_strain, coil_1_strain)

    return pulse_strain

def get_score_raw_single_channel(x, e, x_opt=None):
    """
    Compute score for e-efield cross correlations.
    Non-normalized score is returned, so you need to do sth like

    score = 1 / ((n_ele ** 2 - n_ele) / 2) * get_score()

    Parameters
    ----------
    x : np.ndarray of float
        (n_placements, ) selection of coil placements.
    e : np.ndarray of float
        (n_ele*3, n_channels) E-field.
    x_opt :  np.ndarray of float, optinonal
        (n_pre_opt,) Previously optimized channel currents.

    Returns
    -------
    score : float
        non-normalized score np.nansum(np.abs(np.triu(np.corrcoef(e_mag), k=1)))
    """
    # x = ((x / x.max()) * e.shape[0]-1).astype(int)
    if x_opt is not None:
        x = np.hstack((x_opt, x))
    x = x.astype(int)
    e_mag = e[x, :]

    # determine average correlation coefficient
    # r_avg = 1/((n_ele**2-n_ele)/2) * np.sum(np.abs(np.triu(np.corrcoef(e_mag), k=1)))
    # a = np.abs(np.triu(np.corrcoef(e_mag.T), k=1))

    r_avg = np.nansum(np.abs(np.triu(np.corrcoef(e_mag.T), k=1)))
    return r_avg


def optimize_currents(e, n_stim, currents_prev=None, seed=None,
                      maxiter=200, basiniter=50, method='SLSQP', opt_target='elms',
                      fun=None, fun_args=None, bounds=None, x0=None, constraints=None,
                      verbose=False):
    """
    Optimize the currents for a multichannel TMS array by minimizing e-fields cross-correlation.

    Parameters
    ----------
    e : np.ndarray of float
        (n_elms * 3, n_channel) or (n_elms, 3, n_channel). E in ROI for currents = 1.
    n_stim : int
        Number of stimulations.
    currents_prev : np.ndarray of float, optional
        (n_channels, n_stims_prev) Previous currents to append to.
    seed : int, optional
        Seed for random number generator.
    maxiter : int, default=200
        Max iterations of the optimization.
    method : str, default: 'SLSQP'
        Optimization method.
    opt_target : str, default: 'elms'
        Optimization target. 'elms' for optimizing decorrelations of elements, 'stims' for stimulations.
    fun : function, optional
        Function to scale the e-field.
    fun_args : tuple, optional
        Arguments for the fun function.
    bounds : tuple, optional
        Bounds of the optimization problem.
    x0 : np.ndarray of float, optional
        Initial guess for the optimization problem.
    constraints : list of dict, optional
        Constraints for the optimization problem.
    verbose : bool, default: False
        Print additional information.

    Returns
    -------
    currents : np.ndarray
        (n_channels, n_stims) The optimized currents to drive the multichannel array.
    score : float
        Final score of the solution.
    res
    """
    if e.ndim == 2:
        n_channel = e.shape[1]
        n_ele = int(e.shape[0] / 3)
    elif e.ndim == 3:
        n_channel = e.shape[2]
        n_ele = int(e.shape[0])
    else:
        raise ValueError

    if currents_prev is None:
        n_stim_opt = n_stim
    else:
        n_stim_opt = n_stim - currents_prev.shape[1]

        if n_stim_opt <= 0:
            raise ValueError(f"N_stim({n_stim}) has to be larger "
                             f"than already optimized optimal values ({currents_prev.shape[1]}).")

    # initial guess for currents for all channels * stimulations
    if seed is not None:
        np.random.seed(seed)
    if x0 is None:
        x0 = (np.random.rand(n_channel * n_stim_opt) * 2) - 1

    if verbose:
        print(f"n_ele: {n_ele}, n_channels: {n_channel}, n_stims: {n_stim}")

    print(len(x0))
    if bounds is None:
        bounds = [(-1, 1) for _ in range(len(x0))]
    else:
        assert len(bounds) == len(x0)

    res_scale = 1
    metric = 'mean'
    # optimization algorithm
    minimizer_kwargs = {'args': (e, n_stim, n_ele, n_channel, currents_prev, opt_target, fun, fun_args, res_scale, metric),
                        'method': method,
                        'options': {'disp': False, 'maxiter': maxiter},
                        'bounds': bounds,
                        'tol': 1e-4,
                        'constraints': constraints}

    # def print_fun(x, f, accepted):
    #     print("at minimum %.4f accepted %d" % (f, int(accepted)))

    print_fun = None

    res = basinhopping(get_score_raw, x0, minimizer_kwargs=minimizer_kwargs,
                       niter=basiniter, disp=True, callback=print_fun,
                       T=.25, interval=10, stepsize=.1, target_accept_rate=.5)
    print(res.fun, res.success, res.message)
    # res = minimize(get_score_raw,
    #                args=(e, n_stim, n_ele, n_channel, currents_prev, opt_target, fun, fun_args, res_scale),
    #                x0=x0,
    #                method=method,
    #                options={'disp': True, 'maxiter': maxiter},
    #                bounds=bounds,
    #                tol=1e-6,
    #                constraints=constraints)
    # print(res.fun, res.success, res.message)

    if currents_prev is None:
        currents = np.reshape(res.x, (n_channel, n_stim))
    else:
        currents = np.hstack((currents_prev, np.reshape(res.x, (n_channel, n_stim - currents_prev.shape[1]))))
    if opt_target == 'elms':
        score = 1 / ((n_ele ** 2 - n_ele) / 2) * res.fun * res_scale
    elif opt_target == 'stims':
        score = 1 / ((n_stim ** 2 - n_stim) / 2) * res.fun * res_scale
    else:
        raise ValueError("opt_target has to be 'elms' or 'stims'")

    return currents, score, res


def optimize_currents_single_channel(e, n_stim, currents_prev=None, seed=None,
                                     maxiter=200, method='SLSQP', verbose=False):
    """
    Optimize the coil placement selection for a single channel e-field set minimizing e-fields cross-correlation.

    Parameters
    ----------
    e : np.ndarray of float
        (n_elms3, n_placements).
    n_stim : int
        Number of stimulations.
    currents_prev : np.ndarray of float, optional
        (n_channels, n_stims_prev) Previous currents to append to.
    seed : int, optional
        Seed for random number generator.
    maxiter : int, default=200
        Max iterations of the optimization.
    method : str, default: 'SLSQP'
        Optimization method.
    verbose : bool, default: False
        Print additional information.

    Returns
    -------
    currents :  np.ndarray
        (n_channels, n_stims) The optimized currents to drive the multichannel array.
    score : float
        Final score of the solution.
    """
    n_placements = e.shape[0]
    n_ele = e.shape[1]

    if currents_prev is None:
        n_stim_opt = n_stim
    else:
        n_stim_opt = n_stim - currents_prev.shape[0]

        if n_stim_opt <= 0:
            raise ValueError("N_stim has to be larger than already optimized optimal values!")

    # initial guess for currents for all channels * stimulations
    if seed is not None:
        np.random.seed(seed)
    x0 = np.random.randint(0, n_placements, n_stim_opt)
    if verbose:
        print(f"n_ele: {n_ele}, n_placements: {n_placements}, n_stims: {n_stim}")

    # optimization algorithm
    res = minimize(get_score_raw_single_channel,
                   args=(e, currents_prev),
                   x0=x0,
                   method=method,
                   options={'disp': False, 'maxiter': maxiter},
                   bounds=[(0, n_placements) for _ in range(len(x0))])

    # if currents_prev is None:
    #     currents = np.reshape(res.x, (n_placements, n_stim))
    # else:
    #     currents = np.hstack((currents_prev, np.reshape(res.x, (n_placements, n_stim - currents_prev.shape[1]))))

    score = 1 / ((n_ele ** 2 - n_ele) / 2) * res.fun
    # score = res.fun

    return res.x, score


def get_score(x, e, n_stim, n_ele, n_channel, x_opt=None, opt_target='elms', fun=None, fun_args=None):
    """
    Normalize the score by the number of elements.

    Parameters
    ----------
    x : np.ndarray of float
        (n_channel * n_stim_opt, ) Vector to scale channels for each.
    e : np.ndarray of float
        (n_ele*3, n_channels) E-field.
    n_stim : int
        Number of stimulations to compute score for.
    n_ele : int
        Number of elements.
    n_channel : int
        Number of channels.
    x_opt :  np.ndarray of float, optinonal
        (n_pre_opt,) Previously optimized channel currents.

    Returns
    -------
    score : float
        The normalized score.
    """
    score_raw = get_score_raw(x, e, n_stim, n_ele, n_channel, x_opt, opt_target=opt_target, fun=fun, fun_args=fun_args)
    return 1 / ((n_ele ** 2 - n_ele) / 2) * score_raw


def get_current_init(n_channel, n_stim, bounds=None, only_pos=None, seed=None):
    """
    Get initial guess for the optimization problem.

    Example:
    --------
    # To set channels 3 and 4 to init values within (-.7, -1) and (0.7, 1):
    get_current_init(n_channels=5, n_stim=150,
                     bounds = [(1,-1), (1,-1), (.7,1), (.7,-1), (1,-1)],
                     only_pos = [True, True, False, False])

    Parameters
    ----------
    n_channel : int
        Number of channels.
    n_stim : int
        Number of stimulations.
    bounds : List of tuples
        Bounds of the optimization problem.
    only_pos : List of bool
        If True, only positive values are used from within bounds.

    Returns
    -------
    x0 : np.ndarray of float
        Initial guess for the optimization problem.
    """
    if bounds is None:
        bounds = [(1, -1) for _ in range(n_channel)]
    if only_pos is None:
        only_pos = [False for _ in range(n_channel)]
    assert len(bounds) == n_channel
    assert len(only_pos) == n_channel
    x0 = (np.random.rand(n_channel * n_stim) * 2) - 1
    for i in range(n_channel):
        a = np.random.rand(n_stim)
        max_bound, min_bound = max(bounds[i]), min(bounds[i])
        x0[n_stim * i: n_stim * (i + 1)] = (((a - a.min()) * (max_bound - min_bound)) / (a.max() - a.min()) + min_bound)
        if not only_pos[i]:
            x0[n_stim * i: n_stim * (i + 1)] *= np.random.choice([-1, 1], n_stim)
    return x0


def constraint_min_current(x, n_stim, channels=(2, 3), min_current=0.7, frac=0.95, factor=100., verbose=False):
    """
    This function constrains the currents to be above a certain threshold.
    Should be used with the 'SLSQP' method of scipy.optimize.minimize as an inequality constraint.

    Parameters
    ----------
    x : np.ndarray of float
        (n_channel * n_stim_opt, ) Vector to scale channels for each.
    n_stim : int
        Number of stimulations to compute score for.
    channels : tuple of int, default: (2, 3)
        Channels to constrain.
    min_current : float, default: 0.7
        Minimum current.
    frac : float, default: 0.95
        Fraction of channels that must be within the bounds.
    factor : float, default: 100.
        Factor to scale the constraint.
    verbose : bool, default: False
        Print additional information

    Returns
    -------
    res : float
        >=0 if the constraint is met, <0 otherwise
    """
    np.random.seed(n_stim)
    res = (x[n_stim * channels[0]: n_stim * (channels[0] + 1)] ** 2 > min_current ** 2)
    for i in range(1, len(channels)):
        res += (x[n_stim * channels[i]: n_stim * (channels[i] + 1)] ** 2 > min_current ** 2)
    res = (res.sum() - n_stim * frac) / factor

    # res = (((x[n_stim * 2: n_stim * 3] ** 2 > 0.5) + (
    #             x[n_stim * 3: n_stim * 4] ** 2 > 0.5)).sum() - n_stim * 0.95) / 100.

    if verbose:
        print(f"Bad_currents: {res}")
    return res


def calc_ll(l):
    """
    Calculate matrix multiplication of Lead-filed * transpose(Lead-field).

    Parameters
    ----------
    l : np.ndarray
        [n_elm*3, n_coils] Flattened lead-field matrix. x,y,z direction in each point as consecutive rows.

    Returns
    -------
    LL : np.ndarray
        [n_coils, n_coils] Resulting square matrix LL.
    """
    return l.T @ l


def calc_e(l, j):
    """
    Calculate E-field. Multiplication of lead-field and coil currents.

    Parameters
    ----------
    l : np.ndarray
        [n_elm*3, n_coils] Flattened lead-field matrix. x,y,z direction in each point as consecutive rows.
    j : np.ndarray
        [n_coils]

    Returns
    -------
    E : np.ndarray of dimenstion
        [n_elm, 3] Resulting e-field for all elements in ROI.
    """
    n_k = int(len(l) / 3)  # cound of elements in ROI
    e = l @ j
    return e.reshape((n_k, 3))


def goal(phi, l, ll, target_stim=80, u=0, s=0):
    """
    Goal Function for focality optimization.

    Parameters
    ----------
    phi : np.ndarray
        [n_coils] Angles of hyperellipsoid cylinder.
    l:  np.ndarray of float
        [n_elm*3, n_coils] Flattened lead-field matrix. x,y,z direction in each point as consecutive rows.
    ll : np.ndarray of float
        [n_coils, n_coils] Result of matrix multiplication of lead-field matrix * transpose(lead-field matrix)
    target_stim : int, default: 80
        Desired target stimulation. scaling value for unit current.
    u : np.ndarray of float
        [n_coils, n_coils] Left singular vectors of lead-field matrix.
    s : np.ndarray of float
        [n_coils] Singular values of lead-field matrix.

    Returns
    -------
    eps : int
        Result value of goal function for given phi.
    J : np.ndarray of dimension
        [n_coils] Resulting currents for given phi.
    """
    # Count of Coils
    n_c = len(l[0])
    n_elm = l.shape[0] / 100
    # phi /= 1000
    # Local lead-field: lead-field at target position k
    # start =
    # end =
    # phi *= np.pi
    # Check for very small singular values and avoid division by zero
    s[s < 1e-10] = 1e-10
    # reciprocal of sqrt to normalize singular values. transforms them into set of scaling factors
    e = 1. / np.sqrt((np.array(s[0:3])))

    # hyperellipsoid cylinder
    x = np.zeros(n_c)
    # for i in range(n_c):
    #     if i == 0:
    x[0] = e[0] * math.cos(phi[0])
    # elif i == 1:
    x[1] = e[1] * math.sin(phi[0]) * math.cos(phi[1])
    #     elif i == 2:
    x[2] = e[2] * math.sin(phi[0]) * math.sin(phi[1])
    # else:
    for i in range(3, n_c):
        x[i] = phi[i - 1]

    j = u @ x
    # eps =
    # print(f"GOAL phi: {np.round(phi,4)}: error: {np.round(eps/n_elm, 4)}, j:{j}")
    j = j * target_stim
    return get_focality_score(j, ll) / n_elm, j


def get_focality_score(j, ll):
    """
    Calculate the focality score.
    """
    # Calculate the focality score
    eps = j @ ll @ j.T
    return eps


def foc_opt(l, ll, k, target_stim=.1, current_constraint=None, verbose=False):
    """
    Focality optimization routine. Minimizing goal function value via gradient descend of scipy.minimize().
    The goal function is the squared local lead-field matrix at target position k.
    The optimization is done in the space of the hyperellipsoid cylinder.

    1) Initially, an unconstrained optimization is performed.
    2) In case current_constraints are set and the resulting currents exceed the constraints,
        a constrained optimization is performed.
        If the constrained optimization fails, the function returns None for e, eps, j, and phi_opt.

    ``l`` is ``np.vstack([E_chan1.flatten(), E_chan2.flatten(),...]).T``
    with ``E_chan1.shape = (n_elms, 3)``

    ``ll`` is ``L.T @ L``

    Parameters
    ----------
    l : np.ndarray of float
        [n_elm*3, n_coils] Flattened lead-field matrix. x,y,z direction in each point as consecutive rows.
    ll : np.ndarray of float
        [n_coils, n_coils] Result of matrix multiplication of lead-field matrix * transpose(lead-field matrix)
    k : int
        index of optimization target element in ``l``.
    target_stim : int, default = 80
        desired target stimulation. scaling value for unit current.
    verbose : bool, default = False
        If True, displays optimization progress and debug information. This parameter controls the verbosity of the
        optimization process.
        If set to False, it suppresses the output.
    current_constraint : int, optional
        If set, a current constraint is applied for the optimization, ensuring the resulting currents don't exceed
        this value.

    Returns
    -------
    e : np.ndarray
        [n_elm, 3] The resulting electric field from the optimization process.
    eps : float
        The result of the goal function for the given phi.
    j : np.ndarray
        [n_coils] The optimal currents for the given phi.
    phi_opt : np.ndarray of dimension
        [n_coils] The resulting optimal angles of the hyperellipsoid.
    constraint_opt : str
        The type of constraint optimization used. 'foc_opt' for unconstrained, 'foc_opt_constrained' for constrained.
    eps_raw : float
        The result of the goal function for the given phi.
    """
    l_k = l[3 * k:3 * k + 3, :]
    # singular value decomposition of squared local lead-field + perturbation to improve numeric stability
    [u, s, _] = np.linalg.svd(l_k.T @ l_k + np.eye(l_k.shape[1]) * 1e-9)
    # l_k.shape

    def current_constraint_fun(x):
        """
        Constraint function to ensure that the currents do not exceed a certain threshold.
        """
        eps, j = goal(x, l, ll, target_stim, u, s)
        res = (current_constraint - np.abs(j)).min() + 0.00001
        # print(f"target stim {np.round(target_stim, 2)} "
        #       f"j: {j}, phi: {x}, res={res}, res: {np.round((res + + 0.0001) / 10000, 2)}")
        # print(f"CONS phi: {np.round(x, 4)}: ret: {np.round(res, 4)}, j:{j}")
        # return (res+ + 0.0001) / 10000
        return res

    if not current_constraint:
        current_constraint = None
    elif current_constraint is True:
        raise ValueError("current_constraint has to be an integer value.")
    assert l.shape[1] == ll.shape[0] == ll.shape[1]
    constraint_opt = 'foc_opt'
    n_chans = l.shape[1]

    # check if target strength can be achieved with constrained
    if current_constraint is not None:
        e_unflattened = l.T.reshape(n_chans, -1, 3)
        j_max, e_max = get_maximality_currents(e_unflattened[:, k, :], current_constraint)

        if e_max < target_stim:
            print(f"Maximality currents {j_max} yield e-field {e_max} < target_stim {target_stim}.")
            eps = get_focality_score(j_max, ll)
            e = calc_e(l, j_max)
            return e, eps, j_max, [-1] * n_chans, "Maximality", np.inf

    # do an initial optimization without constraints
    constraint = {'type': 'ineq', 'fun': current_constraint_fun}

    epss = [20, 50, 100, 200]
    ftols = [1e-8, 1e-10, 1e-12, 1e-14, 1e-16, 1e-18, 1e-20, 1e-22]
    tols = [1e-4, 1e-5, 1e-6]
    steps = [0.01, 0.03, 0.05, 0.07, 0.09]
    eps_goal = np.inf
    best_params = {}
    for eps in epss:
        for ftol in ftols:
            for step in steps:
                for tol in tols:
                    for i in range(10):
                        phi = np.random.rand(len(l.T)) * np.pi
                        opt_slsqp_unconstrained = scipy.optimize.minimize(
                            lambda phi: goal(phi, l, ll, target_stim, u, s)[0],
                            phi,
                            method='SLSQP',
                            tol=tol,
                            options={
                                "disp": verbose > 1,
                                "maxiter": 1000,
                                'ftol': ftol,
                                'eps': eps,
                                'finite_diff_rel_step': step,
                            },
                            constraints=None,
                            bounds=[(0, np.pi) for _ in range(n_chans)]  # (-np.pi/2, np.pi/2),(-np.pi/2, np.pi/2))
                        )

                        if opt_slsqp_unconstrained.success and opt_slsqp_unconstrained.fun < eps_goal:
                            eps_goal = opt_slsqp_unconstrained.fun
                            phi_unconstrained = opt_slsqp_unconstrained.x
                            best_params['eps'] = eps
                            best_params['ftol'] = ftol
                            best_params['step'] = step
                            best_params['tol'] = tol
                            # print(f"{np.round(eps_goal)}: {np.round(np.abs(goal(phi_unconstrained, l, ll, target_stim, u, s)[1]).max())}")

    eps_raw_unconstrained, j_unconstrained = goal(phi_unconstrained, l, ll, target_stim, u, s)
    eps_unconstrained = get_focality_score(j_unconstrained, ll)
    phi_unconstrained = phi_unconstrained

    if current_constraint is not None and np.any(np.abs(j_unconstrained) * .99 > current_constraint):
        if verbose:
            print(f"Unconstrained optimization yields j_max: {np.max(abs(j_unconstrained))}. "
                  f"Performing constrained optimization.")
        # perform constrained optimization
        n_iter = 0
        n_success = 0
        eps_final = np.inf
        while n_success <= 10 and n_iter < 1000:
            phi = np.random.random(n_chans) * np.pi
            opt_slsqp_constrained = scipy.optimize.minimize(
                lambda phi: goal(phi, l, ll, target_stim, u, s)[0],
                phi,
                method='SLSQP',
                tol=best_params['tol'] + np.random.rand(1) * best_params['tol'] - best_params['tol']/2,
                options={
                    "disp": verbose > 1,
                    "maxiter": 1000,
                    'ftol': best_params['ftol'] + np.random.rand(1) * best_params['ftol'] - best_params['ftol']/2,
                    'eps': best_params['eps'] + np.random.rand(1) * best_params['eps'] - best_params['eps']/2,
                    'finite_diff_rel_step': best_params['step'] +
                                            np.random.rand(1) * best_params['step'] - best_params['step']/2,
                },
                constraints=constraint,
                bounds=[(0, np.pi) for _ in range(n_chans)]
            )

            if opt_slsqp_constrained.success and opt_slsqp_constrained.fun < eps_final:
                eps_final = opt_slsqp_constrained.fun
                phi_opt = opt_slsqp_constrained.x
                # print(f"n_iter: {n_iter} -- "
                #       f"{np.round(opt_slsqp_constrained.fun, 4)}, "
                #       f"{opt_slsqp_constrained.message} {opt_slsqp_constrained.fun < eps_final} n_succes={n_success}")
            if opt_slsqp_constrained.success:
                n_success += 1
            n_iter += 1

        if n_success == 0:
            print(f"Constrained optimization failed after {n_iter} iterations.")
            e = calc_e(l, j_max)
            eps = get_focality_score(j_max, ll)
            return e, eps, j_max, [-1] * n_chans, "Failed", np.inf, eps, j_max, [-1] * n_chans, np.inf

        # print(f"{k} n_iter: {n_iter}, n_success: {n_success}, emax: {np.round(e_max, 2)}")
        constraint_opt = 'foc_opt_constrained'
        eps_raw_opt, j_opt = goal(phi_opt, l, ll, target_stim, u, s)
        eps_opt = get_focality_score(j_opt, ll)

    else:
        eps_opt = eps_unconstrained
        j_opt = j_unconstrained
        eps_raw_opt = eps_raw_unconstrained
        phi_opt = phi_unconstrained

    e = calc_e(l, j_opt)
    return (e, eps_opt, j_opt, phi_opt, constraint_opt, eps_raw_opt,
            eps_unconstrained, j_unconstrained, phi_unconstrained, eps_raw_unconstrained)


def sphere(points, connections, radius, center):
    """
    Define target-sphere around M1-center with given radius.
    All triangle surface elements included that contain minimum one point inside sphere.

    Parameters
    ----------
    points : np.ndarray
        [n_points x 3] Point elements for describing triangle surface. x,y,z coordinates.
    connections : np.ndarray
         [n_elm x 3] Triangle connectivity, indices of ``points``.
    radius : int
        Sphere radius.
    center : list of int
         [3] Center [x, y, z] of sphere.

    Returns
    -------
    valid_p : np.ndarray
        All points included in the sphere.
    idx_c : np.ndarray
        Triangle surface indices included in the sphere.
    """
    assert points.shape[1] == 3
    distances = np.linalg.norm(points - center, axis=1)
    valid_p = np.argwhere(distances < radius)
    valid_p = valid_p.flatten().tolist()
    # for point in points:
    #     distance = np.sqrt(np.sum((point - center) ** 2))
    #     if distance <= radius:
    #         valid_p.append(i)
    idx_c = []
    for count, c in tqdm.tqdm(enumerate(connections), total=connections.shape[0], desc='Finding sphere indices'):
        # c = connections[0,:]
        if np.in1d(valid_p, c).any():
            idx_c.append(count)

        # if count > 1000:
        #     break
        # if c[0] in valid_p or c[1] in valid_p or c[2] in valid_p:
        #     idx_c.append(c)
    # valid_p in connections[:,0]
    # len(valid_p)
    # np.max(valid_p)
    # idx_c = [i for i, row in enumerate(
    #     connections) if any(point in valid_p for point in row)]

    return valid_p, idx_c


def optimization_all_k(l, points, connections, radius, center, ll, current_constraint=None, target_stim=80):
    """
    Run optimization on all target points inside sphere with certain radius.
    Save results as hdf5.

    Parameters
    ----------
    l : np.ndarray of dimension
         [n_elm*3, n_coils] Flattened lead-field matrix. x,y,z direction in each point as consecutive rows.
    points : np.ndarray of dimension
         [n_points, 3] Point elements for describing triangle surface. x,y,z coordinates.
    connections : np.ndarray of dimension
        [n_elm x 3] Triangle surface elements. defined by three point element indices.
    radius : int
        Sphere radius.
    center : list of float
        [3] Center [x, y, z] of target sphere.
    ll : np.ndarray
        [n_coils, n_coils] Result of matrix multiplication of lead-field matrix * transpose(lead-field matrix).
    current_constraint: int, optional
        If set, applies a current constraint during the optimization, ensuring the resulting currents
        don't exceed a certain threshold (100 in this case).
    target_stim: int
        desired target stimulation value.

    Returns
    -------
    ee : np.ndarray of dimension [n_targets, n_elm, 3]
        Electric fields computed for each target in the optimization process. Each row represents the electric
        field values across the ROI for a specific target.
    eps : list of int
        Result values of the goal function for each target. This represents how well the optimization achieved
        the target stimulation for each point.
    jj : np.ndarray of dimension [n_points x n_coils]
        Currents for each point in the sphere. Each row corresponds to the resulting currents for a specific point.
    phi : list of np.ndarray
        Optimal angles of the hyperellipsoid for each target point, resulting from the optimization.
    curr_cons : list of float
        The maximum stimulation for each target after applying current constraints (if any).
    idx : list of int
        Indices of the points inside the target sphere, selected based on the radius and center.
    """
    eps_all = []
    jj = []
    phi_all = []
    ee = []
    curr_cons = []
    targets = []
    if radius <= 0:
        print(f"Computing for all points in ROI")
        idx = np.arange(len(points))
    else:
        print("calculating target sphere")
        idx = sphere(points, connections, radius, center)[1]
    c = 0
    t = tqdm.tqdm(idx)
    e = None
    for i in t:
        k = i
        c += 1
        targets.append(i)
        # print(f"target number {c} of total {len(idx)}, with index {i}")

        e, eps, j, phi, max_stim = foc_opt(l, ll, k, target_stim, current_constraint, False, t)
        e = np.linalg.norm(e, axis=1)
        eps_all.append(eps)
        jj.append(j)
        phi_all.append(phi)
        ee.append(e)
        curr_cons.append(max_stim)

    # save EE as k x 3*ROI dataframe
    print(f"E.shape: {e.shape}")
    print(f"np.array(ee)")
    ee = np.array(ee)
    # reshaped_EE = EE.reshape(len(EE), -1)
    # df_EE = pd.DataFrame(reshaped_EE)

    return ee, eps_all, jj, phi_all, curr_cons, idx


def foc_opt_worker_fun(idx, l, ll, points_wb, connections_wb, l_wb, current_constraint, target_stim):
    """
    Worker fun to be used with multiprocessing.
    """
    # perform optimization
    res = foc_opt(l, ll, idx, target_stim, current_constraint, False)
    e, eps, j, phi, constraint_type, eps_raw, eps_unc, j_unc, phi_unc, eps_raw_unc = res
    e = np.linalg.norm(e, axis=1)
    max_e = e[idx]
    percentiles = np.linspace(0.5, 1.0, 6)
    # eps_all.append(eps)
    # j_all.append(j)
    # phi_all.append(phi)
    # constraint_types.append(constraint_type)
    # eps_raw_all.append(eps_raw)
    eval_res = {}
    eval_res['idx'] = idx
    eval_res['eps'] = eps
    eval_res['eps_cun'] = eps_unc
    eval_res['eps_raw'] = eps_raw
    eval_res['eps_raw_unc'] = eps_raw_unc
    for chan in range(len(j)):
        eval_res[f"j_{chan:0>2}"] = j[chan]
        eval_res[f"j_unc_{chan:0>2}"] = j_unc[chan]
    eval_res['max_e'] = max_e
    eval_res['constraint_type'] = constraint_type

    # get wb e-field
    e_wb = calc_e(l_wb, j)
    e_wb_mag = np.linalg.norm(e_wb, axis=1)

    # evaluate
    t2m = target_to_max_ratio(e_wb_mag, k=None, e_target=max_e)
    eval_res['T2M'] = t2m
    for percentile in percentiles:
        res_oa = overstimulated_area(e_wb_mag, None, connections_wb,
                                     points_wb, percentile, e_target=max_e)
        eva_rel, eva_abs, total_area_cm2 = res_oa
        eval_res[f'OA_eva_rel_{percentile}'] = eva_rel
        eval_res[f'OA_eva_abs_{percentile}'] = eva_abs
        eval_res[f'OA_total_area_cm2_{percentile}'] = total_area_cm2

    # compute f = e^2 / sum(e_wb_mag^2)
    e_wb_mag[idx] = 0
    f = target_stim ** 2 / np.sum(e_wb_mag ** 2)
    eval_res['f'] = f

    return eval_res


def optimization_roi(l, points_wb, connections_wb, l_wb, opt_indices=None,
                     current_constraint=None, target_stim=100, hdf5_file_path=None, n_cpus=1):
    """
    Run optimization on all target points inside sphere with certain radius.
    Save results as hdf5.

    Parameters
    ----------
    l : np.ndarray of dimension
         [n_elm*3, n_coils] Flattened lead-field matrix. x,y,z direction in each point as consecutive rows.
    points_wb : np.ndarray of float
         [n_points, 3] Point elements for describing triangle surface. x,y,z coordinates.
    connections_wb : np.ndarray of float
        [n_elm x 3] Triangle surface elements. defined by three point element indices.
    l_wb : np.ndarray of float
        [n_elm*3, n_coils] Flattened lead-field matrix. x,y,z direction in each point as consecutive rows.
    opt_indices : list of int, optional
        [n_points] Indices of the elements to optimize.
    current_constraint: int, optional
        If set, applies a current constraint during the optimization, ensuring the resulting currents
        don't exceed a certain threshold (100 in this case).
    target_stim: int
        desired target stimulation value
    hdf5_file_path: str, optional
        Path to save the results in HDF5 format.
    n_cpus: int, default: 1
        Number of CPUs to use for parallel processing.

    Returns
    -------
    eps : list of int
        Result values of the goal function for each target. This represents how well the optimization achieved
        the target stimulation for each point.
    j : np.ndarray of dimension [n_points x n_coils]
        Currents for each point in the sphere. Each row corresponds to the resulting currents for a specific point.
    phi : list of np.ndarray
        Optimal angles of the hyperellipsoid for each target point, resulting from the optimization.
    max_e : list of float
        The maximum stimulation for each target after applying current constraints (if any).
    constraint_types : list of string
        The type of constraint optimization used. 'foc_opt' for unconstrained, 'foc_opt_constrained' for constrained.
    eps_raw : list of float
        The result of the goal function for the given phi.
    eval_res : dict
        Dictionary containing various evaluation metrics for the optimization process, including
        overstimulated area and target-to-max ratio.
    """
    ll = l.T @ l
    eps_all = []
    eps_raw_all = []
    j_all = []
    phi_all = []
    max_e = []
    constraint_types = []
    if opt_indices is None:
        opt_indices = np.arange(int(l.shape[0] / 3))

    _fun = partial(foc_opt_worker_fun, l=l, ll=ll, points_wb=points_wb, connections_wb=connections_wb, l_wb=l_wb,
                   current_constraint=current_constraint, target_stim=target_stim)
    dict_list = process_map(_fun, opt_indices, max_workers=n_cpus)
    eval_res = {k: [] for k in dict_list[0].keys()}
    for d in dict_list:
        for k, v in d.items():
            eval_res[k].append(v)

    if hdf5_file_path is not None:
        with h5py.File(hdf5_file_path, 'w') as file:
            for k, d in eval_res.items():
                data = np.zeros((int(l.shape[0] / 3)))
                try:
                    data[opt_indices] = d
                except ValueError:
                    d = np.unique(d, return_inverse=True)[1]
                    data[opt_indices] = d
                file.create_dataset(k, data=data)

    return eps_all, j_all, phi_all, max_e, constraint_types, eps_raw_all, eval_res


def scd(connections, points, idx, mesh_skin, center_skin=None):
    """
    Calculate skin-cortex distance between midlayer target points and skin surface

    Parameters
    ----------
    connections : np.ndarray of dimension [n_ROI/3 x 3]
        Triangle surface elements, defined by three point element indices for each triangle.
    points : np.ndarray of dimension [n_points x 3]
        Point elements describing the triangle surface. Each point has x, y, z coordinates.
    idx : list or int
        Indices of the target points or a single index. If scalar, it is converted to a list of one element.
    mesh_skin : simnibs.mesh_io.Mesh
        A SimNIBS skin surface mesh.
    center_skin : np.ndarray of dimension, default: None
        (n_elements, 3) The centers of the skin mesh elements. If not provided, it is calculated using the
        `elements_baricenters()` method of the `mesh_skin`.

    Returns
    -------
    target_idx : np.ndarray of dimension [n_targetsphere]
        indices of targets in sphere
    scd : np.ndarray of dimension [n_targetsphere]
        skin cortex distance, or rather skin gm midlayer distance for each triangle mesh in target sphere
    """
    warnings.warn("This function is deprecated and will be removed in the future. "
                  "Use `pynibs.utils.scd` instead.",
                  DeprecationWarning)
    scd = []
    centers = []

    if center_skin is None:
        center_skin = mesh_skin.elements_baricenters()[:]

    # Check if idx is scalar or iterable
    if isinstance(idx, (list, np.ndarray)):
        target_idx = idx
    else:
        target_idx = [idx]  # Convert scalar to list for iteration

    for idx in target_idx:
        connection = connections[idx]
        pts = points[connection]
        center = np.mean(pts, axis=0)
        centers.append(center)

    for idx, target in zip(target_idx, centers):
        distances = np.linalg.norm(center_skin - target, axis=1)
        scd = distances.min()

    return target_idx, scd


def target_to_max_ratio(e, k, e_target=None):
    """
    Evaluation-metric: E_target/E_max.
    Quantifies the difference between target ``|E|`` and ``max(|E|)``.

    Parameters
    ----------
    e : np.ndarray
        [n_elm, 1] or [n_elm, 3] E-Field.
    k : int
        Index of optimized target.
    e_target : float, optional
        E-field at the target position. If not provided, it is calculated from ``e``.

    Returns
    -------
    t2m : int
        e_target/e_max Metric
        if max_stim = 1: e_target == e_max
        if max_stim < 1: e_target << e_max
        -> values < 1 indicate worse focality
    """
    assert k is None or e_target is None, "Either k or e_at_target should be None, not both."
    if e_target is None:
        e_target = e[k]
    if len(e.shape) == 2 and e.shape[1] == 3:
        e = np.linalg.norm(e, axis=1)
    t2m = e_target / e.max()
    return t2m


def calculate_triangle_area(p1, p2, p3):
    """
    Calculate the area of a triangle given its three vertices. v

    Parameters
    ----------
    p1, p2, p3 : np.ndarray
        (n, 3) Coordinates of the triangle's vertices.

    Returns
    -------
    area : np.ndarray
        (n, ) Area of the triangles.
    """
    # Using the cross product to find the area of the triangle
    vec1 = p2 - p1
    vec2 = p3 - p1
    cross_product = np.cross(vec1, vec2)
    area = np.linalg.norm(cross_product, axis=1) / 2.0
    return area


def overstimulated_area(e, k, connections, points, percentile=1, e_target=None, e_in_tris=True):
    """
    Calculate the prevalence of the target field and transform it into a measure in cm^2.

    Parameters
    ----------
    e : np.ndarray
        [n_nodes,] E-Field magnitudes (in triangles or nodes) .
    k : int
        Index of the optimized target.
    connections : np.ndarray
        [n_connections, 3] Triangle surface elements defined by three point element indices.
    points : np.ndarray of dimension
        [n_points, 3] Coordinates of the mesh points.
    percentile: int, default = 1
        percentile*e_target >= E, range 0-1.
    e_target : float, optional
        E-field at the target position. If not provided, it is calculated from ``e``.
    e_in_tris : bool, default: True
        Is E provided in tris (True) or nodes (False)

    Returns
    -------
    eva_rel : float
        Relative area for which \|E\|\ :sub:`target` >= \|E\|\ :sub:`offtarget`\.
    eva_abs : int
        Absolute count of the indices where \|E\|\ :sub:`target` >= \|E\|\ :sub:`offtarget`\.
    total_area_cm2 : float
        Total area in cm² for which \|E\|\ :sub:`target` >= \|E\|\ :sub:`offtarget`\.
    """
    import trimesh
    assert k is None or e_target is None, "Either k or e_at_target should be None, not both."
    if e_target is None:
        e_target = e[k]
    # print(percentile)
    # print(e_target)
    # print(e)
    if not e_in_tris:
        raise NotImplementedError
    else:
        assert e.shape[0] == connections.shape[0], f"Wrong shape: {e.shape[0]} != {connections.shape[0]}"
    over_idx = e_target * percentile <= e
    eva_abs = np.sum(over_idx)

    # mask = (percentile * e_target <= e)
    # indices = np.where(mask)[0]
    # eva_abs = len(indices)
    eva_rel = eva_abs / e.shape[0]
    # log_eva = math.log(eva_rel) if eva_rel > 0 else float('-inf')

    # Calculate the total area for the indices where |E|_target >= |E|_offtarget
    # triangles_mask = np.isin(connections, indices).any(axis=1)
    selected_triangles = connections[over_idx] # this only works if E is in tris
    total_area = trimesh.triangles.area(points[selected_triangles]).sum()
    # p1 = points[selected_triangles[:, 0]]
    # p2 = points[selected_triangles[:, 1]]
    # p3 = points[selected_triangles[:, 2]]
    #
    # total_area = np.sum(calculate_triangle_area(p1, p2, p3))

    # Convert area from mm² to cm²
    total_area_cm2 = total_area / 100.0

    return eva_rel, eva_abs, total_area_cm2


def foc_targets(r, oa, thr_oa, thr_r, result_path):
    """
    Calculate the prevalence of the target field and transform it into a measure in cm^2.

    Parameters
    ----------
    r:  np.ndarray of dimension
        [n_nodes,] Target-to-max Ratio focality measure evaluated at all nodes.
    oa : np.ndarray of dimension
        [n_nodes,] Overstimulated Area focality measure evaluated all nodes.
    thr_oa: int
        threshold for focality mask: OA <= thr_OA.
    thr_r : float
        threshold for focality mask: R >= thr_R.
    result_path : string
        path to results folder.

    Returns
    -------
    valid_array : np.ndarray of dimension
        [n_nodes] Array where valid_indices = 1, invalid_indices = 0, non-evaluated = nan.
    valid_indices : np.ndarray
        Indices for which OA <= thr_OA and R >= thr_R is true.
    """
    mask = ~np.isnan(r) & ~np.isnan(oa)
    r_clean, oa_clean = r[mask], oa[mask]

    valid_mask = (r_clean >= thr_r) & (oa_clean <= thr_oa)
    valid_indices = np.where(valid_mask)[0]
    valid_array = np.full(len(r), np.nan)

    # Set valid entries to 1 where the mask is true
    valid_array[mask][valid_mask] = 1

    # Set invalid entries to 0 where the mask is false
    valid_array[mask][~valid_mask] = 0

    # Save valid indices and the valid array to an HDF5 file
    with h5py.File(result_path, 'w') as f:
        f.create_dataset('valid_idx', data=valid_indices)
        f.create_dataset('valid_array', data=valid_array)

    return valid_array, valid_indices


def get_maximality_currents(e, current_constraint=100):
    """
    Compute the current set that maximizes \|E\|

    Parameters
    ----------
    e : np.ndarray of float
        (n_chans, 3) Leadfield in target element.
    current_constraint : float, default=100
        Maximum current constraint.

    Returns
    -------
    j_max : list of float
        [n_chans] Current-set to maximize E.
    e_max : Maximum \|E\| for given ``current_constraint``.
    """
    n_chans = e.shape[0]
    combs = np.array(list(itertools.product([-1, 1], repeat=n_chans)))
    combs = combs[:int(combs.shape[0] / 2)]

    e_max_candidates = []
    for i_step in range(combs.shape[0]):
        a = combs[i_step, :][:, np.newaxis] * current_constraint * e
        a = np.linalg.norm(np.sum(a, axis=0))
        e_max_candidates.append(a)

    e_max_candidates = np.array(e_max_candidates)
    max_ids = np.where(e_max_candidates == np.max(e_max_candidates))[0][0]
    j_max = combs[max_ids].copy()
    e_target_max = e_max_candidates[max_ids]

    return j_max, e_target_max


def maximality_focality_optimization(j_foc, flat_e_data, target_stim=80, current_constraint=100, steps=100):
    """
    Find balanced current solution between maximality and focality that yields a target stimulation
    while keeping within a current constraint.

    Parameters
    ----------
    j_foc : np.ndarray of float
        Focality current set.
    flat_e_data : np.ndarray of float
        (n_chans, 3) Flattened leadfield matrix for the element.
    target_stim : float, default=80
        Desired target stimulation.
    current_constraint : float, default=100
        Maximum current constraint.
    steps : int, default=100
        How many steps to consider for the balancing. Basically the resolution of the optimization.

    Returns
    -------
    j_inter : np.ndarray of float
        Current set that balances maximality and focality.
    dist_foc_int : float
        Distance between focality current and the intersection current.
    dist_max_int : float
        Distance between maximality current and the intersection current.
    """
    j_max = get_maximality_currents(flat_e_data, current_constraint)[0]

    # normalize j_foc
    j_foc = j_foc.copy()
    j_foc_max = max(j_foc, key=abs)
    j_foc /= j_foc_max

    # find distances to the maximality current
    foc_max_dists = [np.linalg.norm(j_foc - j_max), np.linalg.norm(j_foc - j_max * -1)]
    if foc_max_dists[0] > foc_max_dists[1]:
        j_foc *= -1

    stepsize = (j_foc - j_max) / steps
    j_inter = None
    for i_step in range(steps):
        j_inter = (j_foc - i_step * stepsize) * current_constraint
        # foc_str = np.array2string(j_foc, formatter={'float_kind': lambda x: f"{np.round(x,2): >4}"})
        # inter_str = np.array2string(j_inter, formatter={'float_kind': lambda x: f"{np.round(x/current_constraint,2): >4}"})
        # max_str = np.array2string(j_max, formatter={'float_kind': lambda x: f"{np.round(x,2): >4}"})
        # print(f"{i_step: >2}: {foc_str} {inter_str} {max_str}")
        e_inter = np.linalg.norm(
            np.sum(
                j_inter[:, np.newaxis] * flat_e_data, axis=0))
        if e_inter >= target_stim and (j_inter <= current_constraint).all():
            break
    print(f"{i_step} used")
    # compute distance of found solution to focality and maximality solutoin
    dist_foc_int = np.round(np.linalg.norm((j_inter / current_constraint) - j_foc), 2)
    dist_max_int = np.round(np.linalg.norm((j_inter / current_constraint) - j_max), 2)

    return j_inter, dist_foc_int, dist_max_int


def write_magventure_xls(opt_currents_fn, max_mso, xls_fn=None, rep=1):
    """
    Writes .xls file for MagVenture mTMS systems.
    """
    assert opt_currents_fn.endswith(".npy"), "opt_currents_fn must be a .npy file"
    currents_opt = np.load(opt_currents_fn)
    n_chans, n_zaps = currents_opt.shape
    current_scale_factor_opt = max_mso / np.sqrt((currents_opt ** 2).max())
    currents_opt *= current_scale_factor_opt

    data = pd.DataFrame(
        columns=["Trigger type", "Line delay [us]", "Trigout start [us]", "Trigout length [us]", "ch1 delay [us]",
                 "ch1 amplitude [%]", "ch2 delay [us]", "ch2 amplitude [%]", "ch3 delay [us]", "ch3 amplitude [%]",
                 "ch4 delay [us]", "ch4 amplitude [%]", "ch5 delay [us]", "ch5 amplitude [%]", "ch6 delay [us]",
                 "ch6 amplitude [%]", "ch7 delay [us]", "ch7 amplitude [%]", "ch8 delay [us]", "ch8 amplitude [%]",
                 "ch9 delay [us]", "ch9 amplitude [%]", "ch10 delay [us]", "ch10 amplitude [%]", "ch11 delay [us]",
                 "ch11 amplitude [%]", "ch12 delay [us]", "ch12 amplitude [%]", "ch13 delay [us]",
                 "ch13 amplitude [%]", "ch14 delay [us]", "ch14 amplitude [%]", "ch15 delay [us]",
                 "ch15 amplitude [%]", "ch16 delay [us]", "ch16 amplitude [%]"])

    data["Trigger type"] = [1] * n_zaps*rep
    data["Line delay [us]"] = [4000000] * n_zaps*rep
    data["Trigout start [us]"] = [5000] * n_zaps*rep
    data["Trigout length [us]"] = [5000] * n_zaps*rep

    for chan_i in range(1, 17):
        data[f"ch{chan_i} delay [us]"] = ['']* n_zaps*rep
        # data[f"ch{chan_i} amplitude [%]"] = currents_opt[chan_i - 1, :] * 100
        data[f"ch{chan_i} amplitude [%]"] = ''

    for chan_i in range(1, 7):
        data[f"ch{chan_i} delay [us]"] = [0]* n_zaps*rep
        data[f"ch{chan_i} amplitude [%]"] = np.repeat(currents_opt[chan_i - 1, :], rep)

    if xls_fn is None:
        xls_fn = opt_currents_fn.replace(".npy", f"_maxmso{max_mso}.xlsx")
    data.to_excel(xls_fn, index=False, engine='openpyxl')
