"""
Contains frequently used workflows in dataset preprocessing.

The functions herein are a collection of simple, but frequent,
transformations of arrays of raw measurement data.
"""

import numpy as np
import copy
from collections.abc import MutableMapping
from contextlib import contextmanager
from itertools import chain
from multiprocessing import Pool
from numbers import Number
from pybamm import Scalar
from scipy.optimize import root_scalar
from scipy.stats import chi2, norm
from ep_bolfi.models.solversetup import solver_setup, simulation_setup
from ep_bolfi.utility.fitting_functions import (
    fit_OCV, inverse_OCV_fit_function, smooth_fit,
    verbose_spline_parameterization
)


class SubstitutionDict(MutableMapping):
    """
    A dictionary with some automatic substitutions.

    *substitutions* is a dictionary that extends *storage* with
    automatic substitution rules depending on its value types:
     - string, which serves the value of *storage* at that value.
     - callable which takes one parameter, which will get passed its
       ``SubstitutionDict`` instance and serves its return value.
     - any other type, which serves the value as-is.
    Assigning values to keys afterwards will overwrite substitutions.
    """

    def __init__(self, storage, substitutions={}):
        self._storage = storage
        self._substitutions = substitutions
        self._log = []
        """Used to log keys used in a "single" key evaluation."""
        self._log_switch = False
        """Used in a property lock on *_log*. See *log_lock*."""

    def __delitem__(self, key):
        """
        Deletes *key* from *self._storage* and *self._substitutions*.

        :param key:
            The key to delete, be it from *storage* or *substitutions*.
        """
        if key in self._storage.keys():
            self._storage.__delitem__(key)
        if key in self._substitutions.keys():
            self._substitutions.__delitem__(key)
        if (
            key not in self._storage.keys()

            and key not in self._substitutions.keys()
        ):
            raise KeyError(key)

    def __getitem__(self, key):
        """
        Gets *key*, as described in the class description.

        :param key:
            The key of the value to get; might cascade.
        """
        if self._log_switch:
            self._log.append(key)
        if key in self._substitutions.keys():
            storage_value = self._substitutions[key]
            if isinstance(storage_value, Number):
                return storage_value
            elif isinstance(storage_value, str):
                return self[storage_value]
            elif callable(storage_value):
                return storage_value(self)
            else:
                # fallback
                return storage_value
        elif key in self._storage.keys():
            return self._storage[key]
        else:
            raise KeyError(key)

    def __iter__(self):
        """
        Iterates through both *self._storage* and *self._substitutions*.
        """
        return chain(self._storage, self._substitutions)

    def __len__(self):
        """
        Length of *self._storage* and *self._substitutions* combined.

        :returns:
            The integer of the length.
        """
        return len(set(
            list(self._storage.keys()) + list(self._substitutions.keys())
        ))

    def __setitem__(self, key, value):
        """
        Sets *key*, deleting a substitution under it if there was one.

        :param key:
            The key under which to set the new value.
        :param value:
            The new value to set. Can not be a substitution rule.
        """
        if key in self._substitutions.keys():
            del self._substitutions[key]
        self._storage[key] = value

    def __str__(self):
        """
        The same string representation as if this were a normal dict.

        :returns:
            A string representation, with all substitutions applied.
        """
        return '{' + ', '.join(
            [str(k) + ': ' + str(v) for k, v in self.items()]
        ) + '}'

    def __repr__(self):
        """
        The same as the string represenation, as internal states are
        intentionally hidden and not properly displayable in all cases.

        :returns:
            A string representation, with all substitutions applied.
        """
        return self.__str__()

    @contextmanager
    def log_lock(self):
        """
        Modifies *__get_item__* to log the keys it touched and cleans up
        the *self._log* after each logging session.
        """
        self._log_switch = True
        self._log = []
        try:
            yield
        finally:
            self._log_switch = False
            self._log = []

    def dependent_variables(self, parameters):
        """
        Useful to determine the whole amount of effective keys.

        :param parameters:
            A list of keys.
        :returns:
            All keys that were touched during substitution rules to
            produce ``[self[p] for p in parameters]``.
        """
        with self.log_lock():
            for p in parameters:
                self[p]
            log = copy.deepcopy(self._log)
        for s in self._substitutions.keys():
            with self.log_lock():
                self[s]
                for p in parameters:
                    if p in self._log:
                        log.append(s)
                        break
        return log


def fix_parameters(parameters_to_be_fixed):
    """
    Returns a function which sets some parameters in advance.

    :param parameters_to_be_fixed:
        These parameters will at least be a part of the dictionary that
        the returned function returns.
    :returns:
        The function which adds additional parameters to a dictionary or
        replaces existing parameters with the new ones.
    """

    def return_all_parameters(free_parameters):
        """
        Adds the *free_parameters* to a pool of parameters.

        :param free_parameters:
            A dictionary which gets added to the fixed pool of
            parameters (see the function ``fix_parameters``).
        :returns:
            A dictionary containing free_parameters and some other
            key-value-pairs as defined by ``fix_parameters``. If a key
            is present in both the value from *free_parameters* is used.
        """

        return_dict = copy.deepcopy(parameters_to_be_fixed)
        return_dict.update(free_parameters)
        return return_dict

    return return_all_parameters


def combine_parameters_to_try(parameters, parameters_to_try_dict):
    """
    Give every combination as full parameter sets.

    Compatible with ``SubstitutionDict``, if *parameters* is one.

    :param parameters:
        The base full parameter set as a dictionary.
    :param parameters_to_try_dict:
        The keys of this dictionary correspond to the *parameters*' keys
        where different values are to be inserted. These are given by
        the tuples which are the values of this dictionary.
    :returns:
        A 2-tuple where the first item is the list of all parameter set
        combinations and the second the list of the combinations only.
    """

    # Fetch the parameters to try.
    parameters_to_try = []
    for key, value in parameters_to_try_dict.items():
        parameters_to_try.append((key, value))

    def recursive_combination(last_index):
        """
        Recursively get every combination of parameters to try.

        :param last_index:
            The recursion index; start off at ``len(dict) - 1``.
        :returns:
            At the end, returns the whole list of combinations.
        """
        return_list = []
        if last_index == 0:
            key = parameters_to_try[0][0]
            for value in parameters_to_try[0][1]:
                return_list.append({key: value})
        elif last_index > 0:
            key = parameters_to_try[last_index][0]
            list_of_combinations = recursive_combination(last_index - 1)
            for value in parameters_to_try[last_index][1]:
                for combination in list_of_combinations:
                    return_dict = {key: value}
                    return_dict.update(combination)
                    return_list.append(return_dict)
        return return_list

    combinations = recursive_combination(len(parameters_to_try) - 1)
    parameters_list = []
    for combination in combinations:
        parameters_list.append(fix_parameters(parameters)(combination))
    return (parameters_list, combinations)


def calculate_means_and_standard_deviations(
    mean,
    covariance,
    free_parameters_names,
    transform_parameters={},
    bounds_in_standard_deviations=1,
    **kwargs
):
    """
    Calculate means and standard deviations.

    Please note that standard deviations translate differently into
    confidence regions in different dimensions. For the confidence
    region, use ``approximate_confidence_ellipsoid``.

    :param mean:
        The mean of the uncertain parameters as a dictionary.
    :param covariance:
        The covariance of the uncertain parameters as a two-dimensional
        NumPy array.
    :param free_parameters_names:
        The names of the parameters that are uncertain as a list. This
        parameter maps the order of parameters in *covariance*.
    :param transform_parameters:
        Optional transformations between the parameter space that is
        used for searching for optimal parameters and the model
        parameters. Any missing free parameter is not transformed.
        The values are 2-tuples. The first entry is a function taking
        the search space parameter and returning the model parameter.
        The second entry is the inverse function.
        For convenience, any value may also be one of the following:
         - 'none' => (identity, identity)
         - 'log' => (exp, log)
    :param bounds_in_standard_deviations:
        Sets how many standard deviations in each direction the returned
        error bounds are. These are first applied and then transformed.
    :param **kwargs:
        Keyword arguments for ``scipy.integrate.quad``, which is used to
        numerically calculate mean and variance.
    :returns:
        A 3-tuple with three dictionaries. Their keys are the free
        parameters' names as keys and their values are those parameters'
        means, standard deviations, and error bounds.
    """

    # Substitute transformations given by name.
    if transform_parameters is not {}:
        for name, function in transform_parameters.items():
            if type(function) is str:
                if function == 'none':
                    transform_parameters[name] = (
                        lambda s: s, lambda b: b
                    )
                elif function == 'log':
                    transform_parameters[name] = (
                        lambda s: np.exp(s), lambda b: np.log(b)
                    )
    # Fill-in 'none' transformations.
    for name in free_parameters_names:
        if name not in transform_parameters.keys():
            transform_parameters[name] = (
                lambda s: s, lambda b: b
            )

    means = {}
    standard_deviations = {}
    error_bounds = {}

    for i, name in enumerate(free_parameters_names):
        mean_internal = transform_parameters[name][1](mean[name])
        variance_internal = covariance[i][i]
        mean_actual = norm.expect(
            lambda x: (transform_parameters[name][0](
                mean_internal + x * np.sqrt(variance_internal)
            )),
            **kwargs
        )
        variance_actual = norm.expect(
            lambda x: (transform_parameters[name][0](
                mean_internal + x * np.sqrt(variance_internal)
            ))**2,
            **kwargs
        ) - mean_actual**2
        means[name] = mean_actual
        standard_deviations[name] = np.sqrt(variance_actual)
        error_bounds[name] = (
            transform_parameters[name][0](
                mean_internal
                - bounds_in_standard_deviations * np.sqrt(variance_internal)
            ),
            transform_parameters[name][0](
                mean_internal
                + bounds_in_standard_deviations * np.sqrt(variance_internal)
            ),
        )

    return (means, standard_deviations, error_bounds)


def approximate_confidence_ellipsoid(
    parameters,
    free_parameters_names,
    covariance,
    mean=None,
    transform_parameters={},
    refinement=True,
    confidence=0.95
):
    """
    Approximate a confidence ellipsoid.

    Compatible with SubstitutionDict, if *parameters* is one.
    The geometric approximation is a refinement of the polytope with
    nodes on the semiaxes of the confidence ellipsoid. The refinement
    step adds a node for each face, i.e., each sub-polytope with
    dimension smaller by 1. This node is centered on that face and
    projected onto the confidence ellipsoid.

    :param parameters:
        The base full parameter set as a dictionary.
    :param free_parameters_names:
        The names of the parameters that are uncertain as a list. This
        parameter has to match the order of parameters in *covariance*.
    :param covariance:
        The covariance of the uncertain parameters as a two-dimensional
        NumPy array.
    :param mean:
        The mean of the uncertain parameters as a dictionary. If not
        set, the values from *parameters* will be used.
    :param transform_parameters:
        Optional transformations between the parameter space that is
        used for searching for optimal parameters and the model
        parameters. Any missing free parameter is not transformed.
        The values are 2-tuples. The first entry is a function taking
        the search space parameter and returning the model parameter.
        The second entry is the inverse function.
        For convenience, any value may also be one of the following:
         - 'none' => (identity, identity)
         - 'log' => (exp, log)
    :param confidence:
        The confidence within the ellipsoid. Defaults to 0.95, i.e., the
        95% confidence ellipsoid.
    :param refinement:
        If False, only the nodes on the semiaxes get returned. If True,
        the nodes centered on the faces get returned as well.
    :returns:
        A 2-tuple where the first item is the list of all parameter set
        combinations and the second the ellipsoid nodes only as a
        two-dimensional numpy array with each node in on row.
    """

    if mean is None:
        mean = parameters
    # Substitute transformations given by name.
    if transform_parameters is not {}:
        for name, function in transform_parameters.items():
            if type(function) is str:
                if function == 'none':
                    transform_parameters[name] = (
                        lambda s: s, lambda b: b
                    )
                elif function == 'log':
                    transform_parameters[name] = (
                        lambda s: np.exp(s), lambda b: np.log(b)
                    )
    # Fill-in 'none' transformations.
    for name in free_parameters_names:
        if name not in transform_parameters.keys():
            transform_parameters[name] = (
                lambda s: s, lambda b: b
            )

    mahalanobis_squared = chi2(len(free_parameters_names))
    standard_deviations = np.sqrt(mahalanobis_squared.ppf(confidence))

    semiaxes_length_squared, semiaxes_normed = np.linalg.eigh(covariance)
    semiaxes_length = np.sqrt(semiaxes_length_squared)
    confidence_semiaxes_length = semiaxes_length * standard_deviations

    # These calculations happen in the eigenspace of the covariance.
    semiaxis_nodes = []
    for dim, csl in enumerate(confidence_semiaxes_length):
        semiaxis_nodes.append(
            [0.0] * dim
            + [-csl]
            + [0.0] * (len(free_parameters_names) - 1 - dim)
        )
        semiaxis_nodes.append(
            [0.0] * dim
            + [csl]
            + [0.0] * (len(free_parameters_names) - 1 - dim)
        )

    if refinement:
        # [aᵢ / √n]ᵢⁿ fulfills ∑ᵢⁿ (aᵢ / √n)² / aᵢ² = 1.
        diagonal = np.array(
            [-confidence_semiaxes_length,
             confidence_semiaxes_length]
        ) / np.sqrt(len(free_parameters_names))

        def recursive_combination(length, index=0):
            """
            Recursively get every combination of node edges.

            :param length:
                The recursion depth; start off at ``len(list) - 1``.
            :param index:
                The recursion index; start off at 0.
            :returns:
                At the end, returns the whole list of nodes.
            """
            return_list = []
            if index == length - 1:
                for entry in diagonal.T[index]:
                    return_list.append([entry])
            elif index < length - 1:
                list_of_combinations = recursive_combination(length, index + 1)
                for entry in diagonal.T[index]:
                    for combination in list_of_combinations:
                        node = [entry]
                        node.extend(combination)
                        return_list.append(node)
            return return_list

        combinations = recursive_combination(len(free_parameters_names))
    else:
        combinations = []

    transformed_mean = np.array([
        transform_parameters[name][1](mean[name])
        for name in free_parameters_names
    ])
    transformed_ellipsoid = transformed_mean + np.array([
        (semiaxes_normed @ node)
        for node in [*semiaxis_nodes, *combinations]
    ])
    ellipsoid = np.array([
        [
            transform_parameters[name][0](v[i])
            for i, name in enumerate(free_parameters_names)
        ]
        for v in transformed_ellipsoid
    ])
    parameters_list = []
    for node in ellipsoid:
        node_parameters = {
            name: node[i] for i, name in enumerate(free_parameters_names)
        }
        parameters_list.append(fix_parameters(parameters)(node_parameters))
    return (parameters_list, ellipsoid)


def capacity(parameters, electrode="positive"):
    """
    Convenience function for calculating the capacity.

    :param parameters:
        A parameter file as defined by ``models.standard_parameters``.
    :param electrode:
        The prefix of the electrode to use for capacity calculation.
        Change to "negative" to use the one with the lower OCP.
    :returns:
        The capacity of the parameterized battery in C.
    """

    return (
        parameters[
            electrode.capitalize()
            + " electrode active material volume fraction"
        ]
        * parameters[electrode.capitalize() + " electrode thickness [m]"]
        * parameters["Current collector perpendicular area [m2]"]
        * parameters[
            "Maximum concentration in " + electrode + " electrode [mol.m-3]"
        ]
        * 96485.33212
    )


def calculate_SOC(timepoints, currents, initial_SOC=0, sign=1, capacity=1):
    """
    Transforms applied current over time into SOC.

    :param timepoints:
        Array of the timepoint segments.
    :param currents:
        Array of the current segments.
    :param initial_SOC:
        The SOC value to start accumulating from.
    :param sign:
        The value by which to multiply the current.
    :param capacity:
        A scaling by which to convert from C to dimensionless SOC.
    :returns:
        An array of the same shape describing SOC in C.
    """

    complete_SOC = []
    for t, I in zip(timepoints, currents):
        Δt = np.array([t1 - t0 for (t0, t1) in zip(t[:-1], t[1:])])
        try:
            I_int = np.array([0.5 * (I1 + I0)
                             for (I0, I1) in zip(I[:-1], I[1:])])
        except TypeError:
            I_int = np.atleast_1d(I)
        complete_SOC.append(np.append(
            initial_SOC,
            initial_SOC + sign * np.cumsum(I_int * Δt) / capacity
        ))
        initial_SOC = complete_SOC[-1][-1]
    return complete_SOC


def calculate_both_SOC_from_OCV(
    parameters,
    negative_SOC_from_cell_SOC,
    positive_SOC_from_cell_SOC,
    OCV
):
    """
    Calculates the SOC of both electrodes from their OCV.

    The SOCs are substitued in the given *parameters*. The SOC of the
    cell as a whole gets returned in case it is needed.

    :param parameters:
        The parameters of the battery as used for the PyBaMM simulations
        (see ``models.standard_parameters``).
    :param negative_SOC_from_cell_SOC:
        A function that takes the SOC of the cell and returns the SOC of
        the negative electrode.
    :param positive_SOC_from_cell_SOC:
        A function that takes the SOC of the cell and returns the SOC of
        the positive electrode.
    :param OCV:
        The OCV for which the SOCs shall be calculated.
    :returns:
        The SOC of the cell as a whole.
    """

    positive_OCV = parameters["Positive electrode OCP [V]"]
    negative_OCV = parameters["Negative electrode OCP [V]"]
    cell_SOC = root_scalar(
        lambda cell_SOC: positive_OCV(positive_SOC_from_cell_SOC(cell_SOC))
        - negative_OCV(negative_SOC_from_cell_SOC(cell_SOC)) - OCV,
        method='toms748', bracket=[0, 1], x0=0.5
    ).root
    positive_SOC = positive_SOC_from_cell_SOC(cell_SOC)
    negative_SOC = negative_SOC_from_cell_SOC(cell_SOC)
    parameters["Initial concentration in positive electrode [mol.m-3]"] = (
        positive_SOC
        * parameters["Maximum concentration in positive electrode [mol.m-3]"]
    )
    parameters["Initial concentration in negative electrode [mol.m-3]"] = (
        negative_SOC
        * parameters["Maximum concentration in negative electrode [mol.m-3]"]
    )

    return cell_SOC


def subtract_OCV_curve_from_cycles(
    dataset,
    parameters,
    starting_SOC=None,
    starting_OCV=None,
    electrode="positive",
    current_sign=0,
    voltage_sign=0,
):
    """
    Removes the OCV curve from a cycling measurement.

    :param dataset:
        A ``Cycling_Information`` object of the measurement.
    :param parameters:
        The parameters of the battery as used for the PyBaMM simulations
        (see ``models.standard_parameters``).
    :param starting_SOC:
        The SOC at the beginning of the measurement. If not given, the
        OCV curve will be inverted to determine the initial SOC.
    :param starting_OCV:
        The OCV at the beginning of the measurement. If not given and
        *starting_SOC* is also not given, the first entry of voltages is
        used for this. If not given, but *starting_SOC* is, the OCP
        function will be evaluated at *starting_SOC* to get the OCV.
    :param electrode:
        "positive" (default) or "negative" for current sign correction
        and capacity calculation. "positive" adds SOC with positive
        current and vice versa. The sign corrections can be overwritten
        with *x_sign*.
    :param current_sign:
        1 adds SOC, -1 subtracts it, 0 follows the default behaviour
        above.
    :param voltage_sign:
        1 subtracts the OCP, -1 adds it, 0 follows the default behaviour
        above.
    :returns:
        2-tuple. First entry are the voltages minus the OCV as estimated
        for each data point. These are structured in exactly the same
        way as in the *dataset*. Second entry are the electrode SOCs as
        Coulomb-counted in the data.
    """

    if current_sign == 0:
        if electrode == "positive":
            current_sign = 1
        else:
            current_sign = -1
    if voltage_sign == 0:
        if electrode == "positive":
            voltage_sign = 1
        else:
            voltage_sign = -1

    OCV_function = parameters[electrode.capitalize() + " electrode OCP [V]"]
    if type(OCV_function(0.5)) is Scalar:
        def OCV_function(s): return (
            parameters[electrode.capitalize() + " electrode OCP [V]"](s).value
        )
    if starting_SOC is None:
        starting_OCV = starting_OCV or dataset.voltages[0][0]
        initial_SOC = root_scalar(
            lambda s: OCV_function(s) - starting_OCV,
            method='toms748',
            bracket=[0, 1],
            x0=0.5
        ).root
    else:
        initial_SOC = starting_SOC
        starting_OCV = starting_OCV or OCV_function(initial_SOC)
    C = capacity(parameters, electrode)
    SOCs = calculate_SOC(
        dataset.timepoints,
        dataset.currents,
        initial_SOC=initial_SOC,
        sign=current_sign,
        capacity=C,
    )

    returned_SOCs = copy.deepcopy(SOCs)

    # Skip over segments where nothing happens anyway.
    SOC_skip_indices = []
    for i, current in enumerate(dataset.currents):
        if not any(current):  # Read "if current is perfectly 0".
            SOC_skip_indices.append(i)
    for i in reversed(SOC_skip_indices):
        SOCs.pop(i)
    open_circuit_voltages = []
    for soc in SOCs:
        open_circuit_voltages.append(list(OCV_function(soc)))
    if 0 in SOC_skip_indices:
        open_circuit_voltages.insert(
            0, [starting_OCV] * len(dataset.timepoints[0])
        )
        SOC_skip_indices.pop(0)
    for i in SOC_skip_indices:
        open_circuit_voltages.insert(
            i, [open_circuit_voltages[i - 1][-1]] * len(dataset.timepoints[i])
        )
    overpotentials = [
        [
            U_entry - voltage_sign * OCV_entry
            for U_entry, OCV_entry in zip(U, OCV)
        ]
        for U, OCV in zip(dataset.voltages, open_circuit_voltages)
    ]
    return overpotentials, returned_SOCs


def subtract_both_OCV_curves_from_cycles(
    dataset,
    parameters,
    negative_SOC_from_cell_SOC,
    positive_SOC_from_cell_SOC,
    starting_SOC=None,
    starting_OCV=None,
):
    """
    Removes the OCV curve from a single cycle.

    :param dataset:
        A ``Cycling_Information`` object of the measurement.
    :param parameters:
        The parameters of the battery as used for the PyBaMM simulations
        (see ``models.standard_parameters``).
    :param negative_SOC_from_cell_SOC:
        A function that takes the SOC of the cell and returns the SOC of
        the negative electrode.
    :param positive_SOC_from_cell_SOC:
        A function that takes the SOC of the cell and returns the SOC of
        the positive electrode.
    :param starting_SOC:
        The SOC at the beginning of the measurement. If not given, the
        OCV curves will be inverted to determine the initial SOC.
    :param starting_OCV:
        The OCV at the beginning of the measurement. If not given, the
        first entry of voltages is used for this.
    :returns:
        2-tuple. First entry are the voltages minus the OCV as estimated
        for each data point. These are structured in exactly the same
        way as in the *dataset*. Second entry are the electrode SOCs as
        counted in the data.
    """

    positive_OCV = parameters["Positive electrode OCP [V]"]
    if type(positive_OCV(0.5)) is Scalar:
        def positive_OCV(s): return (
            parameters["Positive electrode OCP [V]"](s).value
        )
    negative_OCV = parameters["Negative electrode OCP [V]"]
    if type(negative_OCV(0.5)) is Scalar:
        def negative_OCV(s): return (
            parameters["Negative electrode OCP [V]"](s).value
        )
    if starting_SOC is None:
        starting_OCV = starting_OCV or dataset.voltages[0][0]
        cell_SOC_tracker = root_scalar(
            lambda cell_SOC: (
                positive_OCV(positive_SOC_from_cell_SOC(cell_SOC))
                - negative_OCV(negative_SOC_from_cell_SOC(cell_SOC))
                - starting_OCV
            ),
            method='toms748',
            bracket=[0, 1],
            x0=0.5
        ).root
    else:
        cell_SOC_tracker = starting_SOC
        starting_OCV = starting_OCV or (
            positive_OCV(positive_SOC_from_cell_SOC(cell_SOC_tracker))
            - negative_OCV(negative_SOC_from_cell_SOC(cell_SOC_tracker))
        )
    C = capacity(parameters)
    returned_SOCs = []
    voltages = []
    for t, I, U in zip(dataset.timepoints, dataset.currents, dataset.voltages):
        Δt = [t1 - t0 for (t0, t1) in zip(t[:-1], t[1:])]
        try:
            I_int = [0.5 * (I0 + I1) for (I0, I1) in zip(I[:-1], I[1:])]
        except TypeError:
            I_int = np.atleast_1d(I)
        cell_SOC = cell_SOC_tracker + np.cumsum(
            np.array(I_int) * np.array(Δt)
        ) / C
        voltages.append([U[0] - starting_OCV] + list(
            np.array(U[1:])
            - np.array(positive_OCV(positive_SOC_from_cell_SOC(cell_SOC)))
            + np.array(negative_OCV(negative_SOC_from_cell_SOC(cell_SOC)))
        ))
        returned_SOCs.append(cell_SOC)
        cell_SOC_tracker = cell_SOC[-1]
        starting_OCV = (
            positive_OCV(positive_SOC_from_cell_SOC(cell_SOC_tracker))
            - negative_OCV(negative_SOC_from_cell_SOC(cell_SOC_tracker))
        )
    return voltages, returned_SOCs


def laplace_transform(x, y, s):
    """
    Performs a basic Laplace transformation.

    :param x:
        The independent variable.
    :param y:
        The dependent variable.
    :param s:
        The (possibly complex) frequencies for which to perform the
        transform.
    :returns:
        The evaluation of the laplace transform at s.
    """

    x = np.array(x)
    y = np.array(y)
    Δx = x[1:] - x[:-1]
    x_int = 0.5 * (x[1:] + x[:-1])
    y_int = 0.5 * (y[1:] + y[:-1])
    return np.sum(y_int * np.exp(-(s[:, None] + 0.0j) * x_int) * Δx, axis=1)


def find_occurrences(sequence, value):
    """
    Gives indices in *sequence* where it is closest to *value*.

    :param sequence:
        A list that represents a differentiable function.
    :param value:
        The value that is searched for in *sequence*. Also, crossings of
        consecutive values in *sequence* with *value* are searched for.
    @return
        A list of indices in *sequence* in ascending order where *value*
        or a close match for *value* was found.
    """
    root_finder = np.array(sequence) - value
    crossings = (
        np.sign(root_finder[1:]) - np.sign(root_finder[:-1])
    ).nonzero()[0]
    nearest_indices = np.where(
        np.abs(root_finder[crossings]) < np.abs(root_finder[crossings + 1]),
        crossings, crossings + 1
    )
    if len(nearest_indices) == 0:
        nearest_indices = (
            [0]
            if np.abs(root_finder[0]) < np.abs(root_finder[-1]) else
            [len(sequence) - 1]
        )
    return nearest_indices


def OCV_from_CC_CV(
    charge,
    cv,
    discharge,
    name,
    phases,
    eval_points=200,
    spline_SOC_range=(0.01, 0.99),
    spline_order=2,
    spline_smoothing=2e-3,
    spline_print=None,
    parameters_print=False
):
    """
    Tries to extract the OCV curve from CC-CV cycling data.

    :param charge:
        A ``Cycling_Information`` object containing the constant charge
        cycle(s). If more than one CC-CV-cycle shall be analyzed, please
        make sure that the order of this, cv and discharge align.
    :param cv:
        A ``Cycling_Information`` object containing the constant voltage
        part between charge and discharge cycle(s).
    :param discharge:
        A ``Cycling_Information`` object containing the constant
        discharge cycle(s). These occur after each cv cycle.
    :param name:
        Name of the material for which the CC-CV-cycling was measured.
    :param phases:
        Number of phases in the ``fitting_functions.OCV_fit_function``
        as an int. The higher it is, the more (over-)fitted the model
        becomes.
    :param eval_points:
        The number of points for plotting of the OCV curves.
    :param spline_SOC_range:
        2-tuple giving the SOC range in which the inverted
        fitting_functions.OCV_fit_function will be interpolated by a
        smoothing spline. Outside of this range the spline is used for
        extrapolation. Use this to fit the SOC range of interest more
        precisely, since a fit of the whole range usually fails due to
        the singularities at SOC 0 and 1. Please note that this range
        considers the 0-1-range in which the given SOC lies and not the
        linear transformation of it from the fitting process.
    :param spline_order:
        Order of this smoothing spline. If it is set to 0, this only
        calculates and plots the ``fitting_functions.OCV_fit_function``.
    :param spline_smoothing:
        Smoothing factor for this smoothing spline. Default: 2e-3. Lower
        numbers give more precision, while higher numbers give a simpler
        spline that smoothes over steep steps in the fitted OCV curve.
    :param spline_print:
        If set to either 'python' or 'matlab', a string representation
        of the smoothing spline is printed in the respective format.
    :param parameters_print:
        Set to True if the fit parameters should be printed to console.
    :returns:
        An 8-tuple consisting of the following:
        0: OCV_fits
            The fitted OCV curve parameters for each CC-CV cycle as
            returned by ``fitting_functions.fit_OCV``.
        1: I_mean
            The currents assigned to each CC-CV cycle (without CV).
        2: C_charge
            The moved capacities during the charge segment(s). This is
            a list of the same length as charge, cv or discharge.
        3: U_charge
            The voltages during the charge segment(s). Length: same.
        4: C_discharge
            The moved capacities during the discharge segment(s).
            Length: same.
        5: U_discharge
            The voltages during the discharge segment(s). Length: same.
        6: C_evals
            Structurally the same as C_charge or C_discharge, this
            contains the moved capacities that were assigned to the mean
            voltages of charge and discharge cycle(s).
        7: U_means
            The mean voltages of each charge and discharge cycle.
    """

    I_mean = [0.0] * len(charge.timepoints)
    C_min = [0.0] * len(charge.timepoints)
    C_max = [0.0] * len(charge.timepoints)
    C_charge = []
    U_charge = []
    charge_splines = []
    corr = []
    C_discharge = []
    U_discharge = []
    discharge_splines = []
    OCV_fits = []
    C_evals = []
    U_means = []

    for i, (t, I, U) in enumerate(zip(cv.timepoints, cv.currents,
                                      cv.voltages)):
        Δt = np.array([t1 - t0 for (t0, t1) in zip(t[:-1], t[1:])])
        I_int = np.array([0.5 * (I0 + I1) for (I0, I1) in zip(I[:-1], I[1:])])
        C = [0.0] + list(np.cumsum(Δt * I_int) / 3600.0)
        corr.append(np.abs(C[-1]))

    for i, (t, I, U) in enumerate(zip(charge.timepoints, charge.currents,
                                      charge.voltages)):
        I_mean[i] += 0.5 * np.mean(I)
        Δt = np.array([t1 - t0 for (t0, t1) in zip(t[:-1], t[1:])])
        C = [0.0] + list(np.cumsum(Δt * np.array(I[1:])) / 3600.0)
        C = np.array(C) - np.min(C) + corr[i]
        C_charge.append(C)
        U_charge.append(U)
        C_min[i] = np.min(C)
        C_max[i] = np.max(C)
        charge_splines.append(smooth_fit(C, U,
                                         s=spline_smoothing))

    for i, (t, I, U) in enumerate(zip(discharge.timepoints, discharge.currents,
                                      discharge.voltages)):
        I_mean[i] -= 0.5 * np.mean(I)
        Δt = np.array([t1 - t0 for (t0, t1) in zip(t[:-1], t[1:])])
        C = [0.0] + list(np.cumsum(Δt * np.array(I[1:])) / 3600.0)
        C = np.array(C) - np.min(C) - corr[i]
        U = np.array(U)[C >= 0]
        C = C[C >= 0]
        C_discharge.append(C)
        U_discharge.append(U)
        C_min[i] = np.max([C_min[i], np.min(C)])
        C_max[i] = np.min([C_max[i], np.max(C)])
        discharge_splines.append(smooth_fit(C, U,
                                            s=spline_smoothing))

    for i, (min, max, c, d) in enumerate(zip(C_min, C_max, charge_splines,
                                             discharge_splines)):
        C_eval = np.linspace(min, max, eval_points)
        C_evals.append(C_eval)
        U_mean = 0.5 * (c(C_eval) + d(C_eval))
        U_means.append(U_mean)
        # diff = 0.5 * (c(C_eval) - d(C_eval))
        dummy_SOC = np.linspace(0.0, 1.0, eval_points)
        OCV_model = fit_OCV(dummy_SOC, U_mean, N=phases)
        OCV_fits.append(OCV_model)
        fit_SOC = np.linspace(*spline_SOC_range, eval_points)
        spline_OCV = smooth_fit(fit_SOC, inverse_OCV_fit_function(fit_SOC,
                                *OCV_model.fit, inverted=True), spline_order,
                                s=spline_smoothing)
        if parameters_print:
            print("Parameters of OCV fit function (" + str(I_mean[i]) + " A):")
            print("SOC range of data: " + repr(OCV_model.SOC_range))
            print("E₀: [" + ", ".join([str(x) for x in OCV_model.E_0]) + "]")
            print("a: [" + ", ".join([str(x) for x in OCV_model.a]) + "]")
            print("Δx: [" + ", ".join([str(x) for x in OCV_model.Δx]) + "]")
        if spline_order > 0:
            if parameters_print:
                print("Knots of interpolating spline:")
                print(spline_OCV.get_knots())
                print("Coefficients of this spline:")
                print(spline_OCV.get_coeffs())
            if spline_print is not None:
                print(verbose_spline_parameterization(
                    spline_OCV.get_coeffs(), spline_OCV.get_knots(),
                    spline_order, function_name=name, format=spline_print,
                    derivatives=1
                ))

    return (OCV_fits, I_mean, C_charge, U_charge, C_discharge, U_discharge,
            C_evals, U_means)


def calculate_desired_voltage(
    solution,
    t_eval,
    voltage_scale,
    overpotential,
    three_electrode=None,
    dimensionless_reference_electrode_location=0.5,
    parameters={},
):
    """
    Takes a ``pybamm.Solution`` object and postprocesses its voltage
    information, to give overpotentials and/or three-electrode setups.

    :param solution:
        The ``pybamm.Solution`` object to calculate the voltage from.
    :param t_eval:
        The times at which to evaluate the *solution*.
    :param voltage_scale:
        The returned voltage gets divided by this value. For example,
        1e-3 would produce a plot in [mV], and -1 would flip the sign.
    :param overpotential:
        If True, only the overpotential of *solutions* gets plotted.
        Otherwise, the cell voltage (OCV + overpotential) is plotted.
    :param three_electrode:
        With None, does nothing (i.e., cell potentials are used). If
        set to either 'positive' or 'negative', instead of cell
        potentials, the base for the displayed voltage will be the
        potential of the 'positive' or 'negative' electrode against a
        reference electrode. For placement of said reference electrode,
        please refer to *dimensionless_reference_electrode_location*.
    :param dimensionless_reference_electrode_location:
        The location of the reference electrode, given as a scalar
        between 0 (placed at the point where negative electrode and
        separator meet) and 1 (placed at the point where positive
        electrode and separator meet). Defaults to 0.5 (in the middle).
    :param parameters:
        The parameter dictionary that was used for the simulation. Only
        needed for a three-electrode output for the cell geometry.
    :returns:
        The array of the specified voltages over time.
    """

    if three_electrode:
        if parameters is None:
            raise ValueError(
                "To calculate the potentials for a three-electrode setup, "
                "the parameter dictionary of the simulation is needed. "
                "Give it as the keyword argument 'parameters'."
            )
        if three_electrode not in ["positive", "negative"]:
            raise ValueError(
                "'three_electrode' has to be either None, 'positive', or "
                "'negative'."
            )
        # When this issue gets completed, use it instead:
        # https://github.com/pybamm-team/PyBaMM/issues/2188
        L_n = parameters["Negative electrode thickness [m]"]
        L_s = parameters["Separator thickness [m]"]
        L_p = parameters["Positive electrode thickness [m]"]
        dimensional_location = (
            L_n + dimensionless_reference_electrode_location * L_s
        )
        x_working_electrode = (
            0 if three_electrode == "negative" else L_n + L_s + L_p
        )
        reference_electrode_potential = (
            solution["Electrolyte potential [V]"](
                t_eval, x=dimensional_location
            )
        )
        if overpotential:
            ocp = np.array([
                parameters[
                    three_electrode.capitalize() + " electrode OCP [V]"
                ](soc) for soc in solution[
                    "Average "
                    + three_electrode
                    + " particle concentration"
                ](t_eval)
            ])
            U = (
                solution[
                    three_electrode.capitalize()
                    + " electrode potential [V]"
                ](t_eval, x=x_working_electrode)
                - ocp
                - reference_electrode_potential
            ) / voltage_scale
        else:
            U = (
                solution[
                    three_electrode.capitalize()
                    + " electrode potential [V]"
                ](t_eval, x=x_working_electrode)
                - reference_electrode_potential
            ) / voltage_scale
    else:
        if overpotential:
            U = (
                solution["Voltage [V]"](t_eval)
                - solution["Bulk open-circuit voltage [V]"](t_eval)
            ) / voltage_scale
        else:
            U = solution["Voltage [V]"](t_eval) / voltage_scale
    return U


def solve_all_parameter_combinations(
    model,
    t_eval,
    parameters,
    parameters_to_try,
    submesh_types,
    var_pts,
    spatial_methods,
    full_factorial=True,
    **kwargs
):
    """
    Creates permutations of *parameters_to_try* and solves them.

    :param model:
        The PyBaMM battery model that is to be solved.
    :param t_eval:
        The timepoints in s at which this model is to be solved.
    :param parameters:
        The model parameters as a dictionary.
    :param parameters_to_try:
        A dictionary with the names of the model parameters as keys and
        lists of the values that are to be tried out for them as values.
    :param submesh_types:
        The submeshes for discretization. See
        ``solversetup.spectral_mesh_pts_and_method``.
    :param var_pts:
        The number of discretization points. See
        ``solversetup.spectral_mesh_pts_and_method``.
    :param spatial_methods:
        The spatial methods for discretization. See
        ``solversetup.spectral_mesh_pts_and_method``.
    :param full_factorial:
        If True, all parameter combinations are tried out. If False,
        only each parameter is varied with the others staying fixed.
    :param **kwargs:
        The optional parameters for ``solversetup.solver_setup``.
    :returns:
        A 2-tuple with the model solution for *parameters* as 1ˢᵗ entry.
        The second entry mimics *parameters_to_try* with each entry in
        their lists replaced by the model solution for the corresponding
        parameter substitution. The second entry has one additional key
        "all parameters", where all *parameters_to_try* combinations are
        the value.
    """

    solutions = {}
    errorbars = {}
    free_parameters = list(parameters_to_try.keys())
    solver = solver_setup(
        model, parameters, submesh_types, var_pts, spatial_methods,
        free_parameters=free_parameters, **kwargs
    )
    input_parameters = {name: parameters[name] for name in free_parameters}
    solutions[model.name] = solver(t_eval, inputs=input_parameters)

    errorbars = {name: [
        solver(t_eval, inputs=p)
        for p in combine_parameters_to_try(input_parameters, {name: limits})[0]
    ] for name, limits in parameters_to_try.items()}

    errorbars["all parameters"] = [solutions[model.name]]
    for e in errorbars.values():
        errorbars["all parameters"].extend(e)
    if full_factorial:
        errorbars["all parameters"].extend([
            solver(t_eval, inputs=p)
            for p in combine_parameters_to_try(
                input_parameters, parameters_to_try
            )[0]
        ])

    return (solutions, errorbars)


def prepare_parameter_combinations(
    parameters,
    parameters_to_try,
    covariance,
    order_of_parameter_names,
    transform_parameters,
    confidence
):
    """
    Calculates all permutations of the parameter boundaries.

    :param parameters:
        The model parameters as a dictionary.
    :param parameters_to_try:
        A dictionary with the names of the model parameters as keys and
        lists of the values that are to be tried out for them as values.
        Mutually exclusive to *covariance*.
    :param covariance:
        A covariance matrix describing an estimation result of model
        parameters. Will be used to calculate parameters to try that
        together approximate the confidence ellipsoid. This confidence
        ellipsoid will be centered on *parameters*.
        Mutually exclusive to *parameters_to_try*.
    :param order_of_parameter_names:
        A list of names from *parameters* that correspond to the order
        these parameters appear in the rows and columns of *covariance*.
        Only needed when *covariance* is set.
    :param transform_parameters:
        Optional transformations between the parameter space that is
        used for searching for optimal parameters and the model
        parameters. Any missing free parameter is not transformed.
        The values are 2-tuples. The first entry is a function taking
        the search space parameter and returning the model parameter.
        The second entry is the inverse function.
        For convenience, any value may also be one of the following:
         - 'none' => (identity, identity)
         - 'log' => (exp, log)
    :param confidence:
        The confidence within the ellipsoid. Defaults to 0.95, i.e., the
        95% confidence ellipsoid.
    :returns:
        A 2-tuple with the individual parameter variations and then all
        permutations of them.
    """

    if not ((parameters_to_try is None) ^ (covariance is None)):
        raise ValueError(
            "Please set either 'parameters_to_try' or 'covariance'."
        )

    if parameters_to_try is not None:
        individual_bounds = {
            name: combine_parameters_to_try(
                parameters, {name: limits}
            )[0]
            for name, limits in parameters_to_try.items()
        }
        combinations = combine_parameters_to_try(
            parameters, parameters_to_try
        )[0]
    elif covariance is not None:
        individual_bounds = {
            "confidence semiaxes": approximate_confidence_ellipsoid(
                parameters,
                list(order_of_parameter_names),
                covariance,
                transform_parameters=transform_parameters,
                refinement=False,
                confidence=confidence
            )[0]
        }
        combinations = approximate_confidence_ellipsoid(
            parameters,
            list(order_of_parameter_names),
            covariance,
            transform_parameters=transform_parameters,
            refinement=True,
            confidence=confidence
        )[0][len(individual_bounds):]

    return individual_bounds, combinations


def parallel_simulator_with_setup(
    model,
    current_input,
    parameters,
    submesh_types,
    var_pts,
    spatial_methods,
    calc_esoh,
    inputs,
    t_eval,
    voltage_scale,
    overpotential,
    three_electrode,
    dimensionless_reference_electrode_location,
    kwargs
):
    """
    A ``multiprocessing``-compatible part of
    ``simulate_all_parameter_combinations``. See there for details.
    """
    par_var = copy.deepcopy(parameters)
    par_var.update(inputs)
    solver, callback = simulation_setup(
        model, current_input, par_var,
        submesh_types, var_pts, spatial_methods,
        **kwargs
    )
    solution = solver(calc_esoh=calc_esoh, callbacks=callback)
    variable = calculate_desired_voltage(
        solution,
        t_eval,
        voltage_scale,
        overpotential,
        three_electrode,
        dimensionless_reference_electrode_location,
        parameters
    )
    return variable


def simulate_all_parameter_combinations(
    model,
    current_input,
    submesh_types,
    var_pts,
    spatial_methods,
    parameters,
    parameters_to_try=None,
    covariance=None,
    order_of_parameter_names=None,
    additional_input_parameters=[],
    transform_parameters={},
    confidence=0.95,
    full_factorial=True,
    calc_esoh=False,
    voltage_scale=1.0,
    overpotential=False,
    three_electrode=None,
    dimensionless_reference_electrode_location=0.5,
    t_eval=None,
    **kwargs
):
    """
    Creates permutations of *parameters_to_try* and simulates them.

    :param model:
        The PyBaMM battery model that is to be solved.
    :param current_input:
        The list of battery operation conditions, in the format of
        ``pybamm.Simulation``.
    :param submesh_types:
        The submeshes for discretization. See
        ``solversetup.spectral_mesh_pts_and_method``.
    :param var_pts:
        The number of discretization points. See
        ``solversetup.spectral_mesh_pts_and_method``.
    :param spatial_methods:
        The spatial methods for discretization. See
        ``solversetup.spectral_mesh_pts_and_method``.
    :param parameters:
        The model parameters as a dictionary.
    :param parameters_to_try:
        A dictionary with the names of the model parameters as keys and
        lists of the values that are to be tried out for them as values.
        Mutually exclusive to *covariance*.
    :param covariance:
        A covariance matrix describing an estimation result of model
        parameters. Will be used to calculate parameters to try that
        together approximate the confidence ellipsoid. This confidence
        ellipsoid will be centered on *parameters*.
        Mutually exclusive to *parameters_to_try*.
    :param order_of_parameter_names:
        A list of names from *parameters* that correspond to the order
        these parameters appear in the rows and columns of *covariance*.
        Only needed when *covariance* is set.
    :param additional_input_parameters:
        A list of the parameter names that are changed by any of the
        variable parameters, if *parameters* is a ``SubstitutionDict``.
        Use its ``dependent_variables`` method to obtain this list.
    :param transform_parameters:
        Optional transformations between the parameter space that is
        used for searching for optimal parameters and the model
        parameters. Any missing free parameter is not transformed.
        The values are 2-tuples. The first entry is a function taking
        the search space parameter and returning the model parameter.
        The second entry is the inverse function.
        For convenience, any value may also be one of the following:
         - 'none' => (identity, identity)
         - 'log' => (exp, log)
    :param confidence:
        The confidence within the ellipsoid. Defaults to 0.95, i.e., the
        95% confidence ellipsoid.
    :param full_factorial:
        When *parameters_to_try* is set:
         - If True, all parameter combinations are tried out. If False,
           only each parameter is varied with the others staying fixed.
        When "covariance" is set:
         - If False, only the points on the semiaxes of the confidence
           ellipsoid constitute the parameters to try. If True, the
           centres of the faces of the polytope of these points get
           added to the parameters to try, projected onto the surface
           of the confidence ellipsoid.
    :param calc_esoh:
        Passed on to ``pybamm.Simulator``, see there.
    :param t_eval:
        The timepoints at which the "errorbars" shall be evaluated in s.
        If None are given, the timepoints of the solutions will be
        chosen.
    :param **kwargs:
        The optional parameters for ``solversetup.solver_setup``.
    :returns:
        A 2-tuple with the model solution for *parameters* as 1ˢᵗ entry.
        The second entry mimics *parameters_to_try* with each entry in
        their lists replaced by the model solution for the corresponding
        parameter substitution. The second entry has one additional key
        "all parameters", where all *parameters_to_try* combinations are
        the value.
    """

    individual_bounds, combinations = prepare_parameter_combinations(
        parameters,
        parameters_to_try,
        covariance,
        order_of_parameter_names,
        transform_parameters,
        confidence
    )

    if parameters_to_try is not None:
        if isinstance(parameters, SubstitutionDict):
            free_parameters = (
                list(parameters_to_try.keys())
                + list(additional_input_parameters)
            )
        else:
            free_parameters = list(parameters_to_try.keys())
    elif covariance is not None:
        if isinstance(parameters, SubstitutionDict):
            free_parameters = (
                list(order_of_parameter_names)
                + list(additional_input_parameters)
            )
        else:
            free_parameters = list(order_of_parameter_names)

    individual_bounds = {
        name: [
            {
                free_name: parameters_variation[free_name]
                for free_name in free_parameters
            }
            for parameters_variation in bounds
        ]
        for name, bounds in individual_bounds.items()
    }
    combinations = [
        {
            free_name: parameters_variation[free_name]
            for free_name in free_parameters
        }
        for parameters_variation in combinations
    ]

    solutions = {}
    errorbars = {}

    solver, callback = simulation_setup(
        copy.deepcopy(model), current_input, copy.deepcopy(parameters),
        submesh_types, var_pts, spatial_methods,
        **kwargs
    )
    solutions[model.name] = solver(calc_esoh=calc_esoh, callbacks=callback)
    if t_eval is None:
        t_eval = solutions[model.name].t

    parallel_arguments = []
    indices_of_individual_bounds = {}
    index_counter = 0
    for name in individual_bounds:
        new_parallel_arguments = [
            (
                copy.deepcopy(model),
                current_input,
                copy.deepcopy(parameters),
                submesh_types,
                var_pts,
                spatial_methods,
                calc_esoh,
                ib,
                t_eval,
                voltage_scale,
                overpotential,
                three_electrode,
                dimensionless_reference_electrode_location,
                kwargs,
            )
            for ib in individual_bounds[name]
        ]
        indices_of_individual_bounds[name] = [i for i in range(
            index_counter, index_counter + len(new_parallel_arguments)
        )]
        index_counter += len(new_parallel_arguments)
        parallel_arguments.extend(new_parallel_arguments)
    with Pool() as p:
        all_errorbars = p.starmap(
            parallel_simulator_with_setup, parallel_arguments
        )
    for name in individual_bounds:
        errorbars[name] = [
            all_errorbars[i] for i in indices_of_individual_bounds[name]
        ]

    errorbars["all parameters"] = [solutions[model.name]]
    for e in errorbars.values():
        errorbars["all parameters"].extend(e)
    if full_factorial:
        parallel_arguments = [
            (
                copy.deepcopy(model),
                current_input,
                copy.deepcopy(parameters),
                submesh_types,
                var_pts,
                spatial_methods,
                calc_esoh,
                c,
                t_eval,
                voltage_scale,
                overpotential,
                three_electrode,
                dimensionless_reference_electrode_location,
                kwargs,
            )
            for c in combinations
        ]
        with Pool() as p:
            errorbars["all parameters"].extend(p.starmap(
                parallel_simulator_with_setup, parallel_arguments
            ))

    return (solutions, errorbars)
