from __future__ import annotations

import logging
from math import erf, exp

import awkward as ak
import numba
import numpy as np
import pint
import pyg4ometry
from lgdo import Array, VectorOfVectors
from numpy.typing import ArrayLike, NDArray

from .. import units
from ..units import ureg as u
from .utils import HPGePulseShapeLibrary, HPGeRZField

log = logging.getLogger(__name__)


def r90(edep: ak.Array, xloc: ak.Array, yloc: ak.Array, zloc: ak.Array) -> ak.Array:
    """R90 HPGe pulse shape heuristic.

    Parameters
    ----------
    edep
        array of energy.
    xloc
        array of x coordinate position.
    yloc
        array of y coordinate position.
    zloc
        array of z coordinate position.

    Returns
    -------
    calculated r90 for each hit
    """
    pos = [units.units_conv_ak(pos, "mm") for pos in [xloc, yloc, zloc]]

    edep = units.units_conv_ak(edep, "keV")

    tot_energy = ak.sum(edep, axis=-1, keepdims=True)

    def eweight_mean(field, energy):
        return ak.sum(energy * field, axis=-1, keepdims=True) / tot_energy

    # Compute distance of each edep to the weighted mean
    dist = np.sqrt(sum((p - eweight_mean(edep, p)) ** 2 for p in pos))

    # Sort distances and corresponding edep within each event
    sorted_indices = ak.argsort(dist, axis=-1)
    sorted_dist = dist[sorted_indices]
    sorted_edep = edep[sorted_indices]

    def _ak_cumsum(layout, **_kwargs):
        if layout.is_numpy:
            return ak.contents.NumpyArray(np.cumsum(layout.data))

        return None

    # Calculate the cumulative sum of energies for each event
    cumsum_edep = ak.transform(
        _ak_cumsum, sorted_edep
    )  # Implement cumulative sum over whole jagged array
    if len(edep) == 1:
        cumsum_edep_corrected = cumsum_edep
    else:
        cumsum_edep_corrected = (
            cumsum_edep[1:] - cumsum_edep[:-1, -1]
        )  # correct to get cumsum of each lower level array
        cumsum_edep_corrected = ak.concatenate(
            [
                cumsum_edep[:1],  # The first element of the original cumsum is correct
                cumsum_edep_corrected,
            ]
        )

    threshold = 0.9 * tot_energy
    r90_indices = ak.argmax(cumsum_edep_corrected >= threshold, axis=-1, keepdims=True)
    r90 = sorted_dist[r90_indices]

    return units.attach_units(ak.Array(ak.flatten(r90)), "mm")


def drift_time(
    xloc: ak.Array | VectorOfVectors,
    yloc: ak.Array | VectorOfVectors,
    zloc: ak.Array | VectorOfVectors,
    dt_map: HPGeRZField,
    coord_offset: pint.Quantity | pyg4ometry.gdml.Position = (0, 0, 0) * u.m,
) -> ak.Array:
    """Calculates drift times for each step (cluster) in an HPGe detector.

    Parameters
    ----------
    xloc
        awkward array of x coordinate position.
    yloc
        awkward array of y coordinate position.
    zloc
        awkward array of z coordinate position.
    dt_map
        the drift time map.
    coord_offset
        this `(x, y, z)` coordinates will be subtracted to (xloc, yloc, zloc)`
        before drift time computation. The length units must be the same as
        `xloc`, `yloc` and `zloc`.
    """
    # sanitize coord_offset
    coord_offset = units.pg4_to_pint(coord_offset)

    # unit handling (.r_units) for data in (xloc, yloc)]
    xu, yu = [units.units_convfact(data, dt_map.r_units) for data in (xloc, yloc)]
    zu = units.units_convfact(zloc, dt_map.z_units)

    # unwrap LGDOs
    xloc, yloc, zloc = [units.unwrap_lgdo(data)[0] for data in (xloc, yloc, zloc)]

    # awkward transform to apply the drift time map to the step coordinates
    def _ak_dt_map(layouts, **_kwargs):
        if layouts[0].is_numpy and layouts[1].is_numpy:
            return ak.contents.NumpyArray(
                dt_map.φ(np.stack([layouts[0].data, layouts[1].data], axis=1))
            )

        return None

    xloc = xu * xloc - coord_offset[0].to(dt_map.r_units).m
    yloc = yu * yloc - coord_offset[1].to(dt_map.r_units).m
    zloc = zu * zloc - coord_offset[2].to(dt_map.z_units).m

    # evaluate the drift time
    dt_values = ak.transform(
        _ak_dt_map,
        np.sqrt(xloc**2 + yloc**2),
        zloc,
    )
    return units.attach_units(ak.Array(dt_values), units.unit_to_lh5_attr(dt_map.φ_units))


def drift_time_heuristic(
    drift_time: ak.Array | VectorOfVectors,
    edep: ak.Array | VectorOfVectors,
) -> ak.Array:
    """HPGe drift-time-based pulse-shape heuristic.

    See :func:`_drift_time_heuristic_impl` for a description of the algorithm.

    Parameters
    ----------
    drift_time
        drift time of charges originating from steps/clusters. Can be
        calculated with :func:`drift_time`.
    edep
        energy deposited in step/cluster (same shape as `drift_time`).
    """
    # extract LGDO data and units
    drift_time, t_units = units.unwrap_lgdo(drift_time)
    edep, e_units = units.unwrap_lgdo(edep)

    # we want to attach the right units to the dt heuristic, if possible
    if t_units is not None and e_units is not None:
        unit = units.unit_to_lh5_attr(t_units / e_units)
        return units.attach_units(ak.Array(_drift_time_heuristic_impl(drift_time, edep)), unit)

    return ak.Array(_drift_time_heuristic_impl(drift_time, edep))


@numba.njit(cache=True)
def _drift_time_heuristic_impl(
    dt: ak.Array,
    edep: ak.Array,
) -> NDArray:
    r"""Low-level implementation of the HPGe drift-time-based pulse-shape heuristic.

    Accepts Awkward arrays and uses Numba to speed up the computation.

    For each hit (collection of steps), the drift times and corresponding
    energies are sorted in ascending order. The function finds the optimal
    split point :math:`m` that maximizes the *identification metric*:

    .. math::

       I = \frac{|T_1 - T_2|}{E_\text{s}(E_1, E_2)}

    where:

    .. math::

        T_1 = \frac{\sum_{i < m} t_i E_i}{\sum_{i < m} E_i}
        \quad \text{and} \quad
        T_2 = \frac{\sum_{i \geq m} t_i E_i}{\sum_{i \geq m} E_i}

    are the energy-weighted mean drift times of the two groups.

    .. math::

        E_\text{scale}(E_1, E_2) = \frac{1}{\sqrt{E_1 E_2}}

    is the scaling factor.

    The function iterates over all possible values of :math:`m` and selects the
    maximum `I` as the drift time heuristic value.
    """
    dt_heu = np.zeros(len(dt))

    # loop over hits
    for i in range(len(dt)):
        t = np.asarray(dt[i])
        e = np.asarray(edep[i])

        valid_idx = np.where(e > 0)[0]
        if len(valid_idx) < 2:
            continue

        t = t[valid_idx]
        e = e[valid_idx]

        sort_idx = np.argsort(t)
        t = t[sort_idx]
        e = e[sort_idx]

        max_id_metric = 0
        for j in range(1, len(t)):
            e1 = np.sum(e[:j])
            e2 = np.sum(e[j:])

            t1 = np.sum(t[:j] * e[:j]) / e1
            t2 = np.sum(t[j:] * e[j:]) / e2

            id_metric = abs(t1 - t2) * np.sqrt(e1 * e2)

            max_id_metric = max(max_id_metric, id_metric)

        dt_heu[i] = max_id_metric

    return dt_heu


@numba.njit(cache=True)
def _njit_erf(x: ArrayLike) -> NDArray:
    """Error function that can take in a numpy array."""
    out = np.empty_like(x)
    for i in range(x.size):
        out[i] = erf(x[i])
    return out


@numba.njit(cache=True)
def _current_pulse_model(
    times: ArrayLike,
    amax: float,
    mu: float,
    sigma: float,
    tail_fraction: float,
    tau: float,
    high_tail_fraction: float = 0,
    high_tau: float = 0,
) -> NDArray:
    r"""Analytic model for the current pulse in a Germanium detector.

    Consists of a Gaussian, a high side exponential tail and a low side tail:

    .. math::

      \begin{align}
      A(t) = \; &A_\text{max} \times (1-p-p_h) \times \text{Gauss}(t;\mu,\sigma) \\
        &+ A \times p \; \left(1 - \text{erf}\left(\frac{t-\mu}{\sigma_i}\right)\right) \times \frac{e^{t/\tau}}{2e^{\mu/\tau}} \\
        &+ A \times p_h \; \left(1 - \text{erf}\left(-\frac{t-\mu}{\sigma_i}\right)\right) \times \frac{1}{2}e^{-t/\tau}
      \end{align}

    Parameters
    ----------
    times
        Array of times to compute current for.
    amax
        Maximum current for the template
    mu
        Time of the maximum current.
    sigma
        Width of the current pulse
    tail_fraction
        Fraction of the tail in the pulse.
    tau
        Time constant of the low time tail.
    high__tail_fraction
        Fraction of the high tail in the pulse.
    high_tau
        Time constant of the high time tail.

    Returns
    -------
    The predicted current waveform for this energy deposit.
    """
    norm = 2 * exp(mu / tau)
    norm_high = 2

    dx = times - mu
    term1 = (
        amax * (1 - tail_fraction - high_tail_fraction) * np.exp(-(dx * dx) / (2 * sigma * sigma))
    )
    term2 = amax * tail_fraction * (1 - _njit_erf(dx / sigma)) * np.exp(times / tau) / norm
    term3 = (
        amax
        * high_tail_fraction
        * (1 - _njit_erf(-dx / sigma))
        * np.exp(-(times - mu) / high_tau)
        / norm_high
    )

    return term1 + term2 + term3


@numba.njit(cache=True)
def _interpolate_pulse_model(
    template: Array, time: float, start: float, end: float, dt: float, mu: float
) -> NDArray:
    """Interpolate to extract the pulse model given a particular mu."""
    local_time = time - mu - start

    if (local_time < start) or (int(local_time) > end):
        return 0

    sample = int(local_time / dt)
    A_before = template[sample]
    A_after = template[sample + 1]

    frac = (local_time - int(local_time)) / dt
    return A_before + frac * (A_after - A_before)


def make_convolved_surface_library(bulk_template: np.array, surface_library: np.array) -> NDArray:
    """Make the convolved surface library out of the template.

    This convolves every row of the surface_library with the template and reshapes the output
    to match the initial template. It returns a 2D array with one more row than the surface_library
    and each row the same length as the template. The final row is the bulk_template for easier interpolation.

    Parameters
    ----------
    bulk_template
        The template for the bulk response
    surface_library
        The 2D array of the surface library.

    Returns
    -------
    2D array of the surface library convolved with the bulk response.
    """
    # force surface library to be 2D
    if surface_library.ndim == 1:
        surface_library = surface_library.reshape((-1, 1))

    templates = np.zeros((len(bulk_template), np.shape(surface_library)[1] + 1))

    for i in range(np.shape(surface_library)[1]):
        templates[:, i] = convolve_surface_response(
            surface_library[1:, i] - surface_library[:-1, i], bulk_template
        )[: len(bulk_template)]

    templates[:, -1] = bulk_template

    return templates


def convolve_surface_response(surf_current: np.ndarray, bulk_pulse: np.ndarray) -> NDArray:
    """Convolve the surface response pulse with the bulk current pulse.

    This combines the current induced on the edge of the FCCD region with the bulk response
    on the p+ contact.

    Parameters
    ----------
    surf_current
        array of the current induced via diffusion against time.
    bulk_pulse
        the pulse template to convolve the surface current with.

    Returns
    -------
    the current waveform after convolution.
    """
    return np.convolve(surf_current, bulk_pulse, mode="full")[: len(surf_current)]


@numba.njit(cache=True)
def get_current_waveform(
    edep: ak.Array,
    drift_time: ak.Array,
    template: ArrayLike,
    start: float,
    dt: float,
    range_t: tuple,
) -> tuple(NDArray, NDArray):
    r"""Estimate the current waveform.

    Based on modelling the current as a sum over the current pulse model defined by
    the template.

    .. math::
        A(t) = \sum_i E_i \times N f(t, dt_i, \vec{\theta})

    Where:
        - :math:`f(t)` is the template
        - :math`\vec{\theta}` are the parameters :math:`(\sigma, p, \tau)`
        - :math:`E_i` and :math:`dt_i` are the deposited energy and drift time.
        - :math:`N` is a normalisation term

    Parameters
    ----------
    edep
        Array of energies for each step
    drift_time
        Array of drift times for each step
    template
        array of the template for the current waveforms
    start
        first time value of the template
    dt
        timestep (in ns) for the template.
    range_t
        a range of times to search around

    Returns
    -------
    A tuple of the time and current for the current waveform for this event.
    """
    n = len(template)

    times = np.arange(n) * dt + start
    y = np.zeros_like(times, dtype=np.float64)

    for j in range(n):
        time = start + dt * j
        if (time < range_t[0]) or (time > (range_t[1] - dt)):
            continue
        y[j] = _get_waveform_value(j, edep, drift_time, template, start, dt, range_t)

    return times, y


@numba.njit(cache=True)
def _get_waveform_value_surface(
    idx: int,
    edep: NDArray,
    drift_time: np.array,
    dist_to_nplus: np.array,
    bulk_template: ArrayLike,
    templates_surface: ArrayLike,
    activeness_surface: ArrayLike,
    distance_step_in_um: float,
    fccd: float,
    start: float,
    dt: float,
) -> tuple[float, float]:
    """Get the value of the waveform at a certain index."""
    n = len(bulk_template)
    out = 0
    etmp = 0
    time = start + dt * idx

    for i in range(len(edep)):
        E = edep[i]
        mu = drift_time[i]
        dist = dist_to_nplus[i]

        if dist < fccd:
            dist_bin = int(dist / distance_step_in_um)

            # get two values (to allow linear interpolation)
            value_low = _interpolate_pulse_model(
                templates_surface[dist_bin], time, start, start + dt * n, dt, mu
            )
            value_high = _interpolate_pulse_model(
                templates_surface[dist_bin + 1], time, start, start + dt * n, dt, mu
            )

            # interpolate between distance bins
            diff = dist / distance_step_in_um - dist_bin
            out += E * (value_low + diff * (value_high - value_low))

            act_low = activeness_surface[dist_bin]
            act_high = activeness_surface[dist_bin + 1]
            etmp += (act_low + diff * (act_high - act_low)) * E

        else:
            out += E * _interpolate_pulse_model(bulk_template, time, start, start + dt * n, dt, mu)
            etmp += E

    return out, etmp


@numba.njit(cache=True)
def _get_waveform_value(
    idx: int,
    edep: ak.Array,
    drift_time: ak.Array,
    template: ArrayLike,
    start: float,
    dt: float,
) -> float:
    """Get the value of the waveform at a certain index."""
    out = 0
    time = start + dt * idx

    for i in range(len(edep)):
        n = len(template)

        E = edep[i]
        mu = drift_time[i]

        out += E * _interpolate_pulse_model(template, time, start, start + dt * n, dt, mu)

    return out


@numba.njit(cache=True)
def _get_waveform_value_pulse_shape_library(
    idx: int,
    edep: ak.Array,
    drift_time: ak.Array,
    r: ak.Array,
    z: ak.Array,
    pulse_shape_library: tuple[np.array, np.array, np.array],
    start: float,
    dt: float,
) -> float:
    """Get the value of the waveform at a certain index, using the pulse shape library."""
    out = 0
    time = start + dt * idx

    for i in range(len(edep)):
        ri, zi = _get_template_idx(r[i], z[i], pulse_shape_library[1], pulse_shape_library[2])

        n = len(pulse_shape_library[0][ri][zi])

        E = edep[i]
        mu = drift_time[i]

        out += E * _interpolate_pulse_model(
            pulse_shape_library[0][ri][zi], time, start, start + dt * n, dt, mu
        )

    return out


@numba.njit(cache=True)
def _get_template_idx(
    r: float,
    z: float,
    r_grid: np.array,
    z_grid: np.array,
) -> tuple[int, int]:
    """Extract the closest template to a given (r,z) point with uniform grid, apart from the first and last point."""
    if r < r_grid[1]:
        ri = 0
    elif r > r_grid[-2]:
        ri = len(r_grid) - 1
    else:
        dr = r_grid[2] - r_grid[1]
        ri = int((r - r_grid[1]) / dr) + 1

    if z < z_grid[1]:
        zi = 0
    elif z > z_grid[-2]:
        zi = len(z_grid) - 1
    else:
        dz = z_grid[2] - z_grid[1]
        zi = int((z - z_grid[1]) / dz) + 1

    return ri, zi


def get_current_template(
    low: float = -1000, high: float = 4000, step: float = 1, mean_aoe: float = 1, **kwargs
) -> tuple[NDArray, NDArray]:
    """Build the current template from the analytic model, defined by :func:`_current_pulse_model`.

    Parameters
    ----------
    low
        start of the template
    high
        end of the template
    step
        time-step, this should divide high-low
    mean_aoe
        The mean AoE value for this detector (to normalise current pulses).
    **kwargs
        Other keyword arguments passed to :func:`_current_pulse_model`.

    Returns
    -------
    tuple of the (template,times)
    """
    if int((high - low) / step) != (high - low) / step:
        msg = "Time template is not a multiple of the time-step."
        raise ValueError(msg)

    x = np.linspace(low, high, int((high - low) / step) + 1)
    template = _current_pulse_model(x, **kwargs)
    template /= np.max(template)
    template *= mean_aoe

    return template, x


@numba.njit(cache=True)
def _get_waveform_maximum_impl(
    t: ArrayLike,
    e: ArrayLike,
    dist: ArrayLike,
    r: ArrayLike,
    z: ArrayLike,
    template: ArrayLike,
    pulse_shape_library: tuple[np.array, np.array, np.array],
    templates_surface: ArrayLike,
    activeness_surface: ArrayLike,
    tmin: float,
    tmax: float,
    start: float,
    fccd: float,
    n: int,
    time_step: int,
    surface_step_in_um: float,
    include_surface_effects: bool,
    use_library: bool,
):
    """Basic implementation to get the maximum of the waveform."""
    max_a = 0
    max_t = 0
    energy = np.sum(e)

    for j in range(0, n, time_step):
        time = start + j

        # skip anything not in the range tmin to tmax (for surface affects this can be later)
        has_surface_hit = include_surface_effects

        if time < tmin or (time > (tmax + time_step)):
            continue

        if not has_surface_hit and (not use_library):
            val_tmp = _get_waveform_value(j, e, t, template, start=start, dt=1.0)
        elif use_library:
            val_tmp = _get_waveform_value_pulse_shape_library(
                j, e, t, r, z, pulse_shape_library, start=start, dt=1.0
            )
        else:
            val_tmp, energy = _get_waveform_value_surface(
                j,
                e,
                t,
                dist,
                template,
                templates_surface.T,
                activeness_surface,
                distance_step_in_um=surface_step_in_um,
                fccd=fccd,
                start=start,
                dt=1.0,
            )

        if val_tmp > max_a:
            max_t = time
            max_a = val_tmp

    return max_t, max_a, energy


@numba.njit(cache=True)
def _estimate_current_impl(
    edep: ak.Array,
    dt: ak.Array,
    dist_to_nplus: ak.Array,
    r: ak.Array,
    z: ak.Array,
    template: np.array,
    pulse_shape_library: tuple[np.array, np.array, np.array],
    times: np.array,
    include_surface_effects: bool,
    use_library: bool,
    fccd: float,
    templates_surface: np.array,
    activeness_surface: np.array,
    surface_step_in_um: float,
) -> tuple[NDArray, NDArray, NDArray]:
    """Estimate the maximum current that would be measured in the HPGe detector.

    This is based on extracting a waveform with :func:`get_current_waveform` and finding the maxima of it.
    """
    A = np.zeros(len(dt))
    maximum_t = np.zeros(len(dt))
    energy = np.zeros(len(dt))

    time_step = 1
    n = len(times)
    start = times[0]

    if include_surface_effects:
        offsets = times[np.argmax(templates_surface, axis=0)]

    # make the convolved surface library
    if include_surface_effects and np.diff(times)[0] != 1.0:
        msg = "The surface convolution requires a template with 1 ns binning"
        raise ValueError(msg)

    for i in range(len(dt)):
        t = np.asarray(dt[i])
        e = np.asarray(edep[i])
        r_tmp = np.asarray(r[i])
        z_tmp = np.asarray(z[i])

        dist = np.asarray(dist_to_nplus[i])

        # get the expected maximum
        tmax = float(np.max(t))
        tmin = float(np.min(t))

        # correct the maximum expected time for surface sims
        if include_surface_effects:
            ncols = templates_surface.shape[1]

            for j, d in enumerate(dist):
                dtmp = int(d / surface_step_in_um)

                use_offset = dtmp <= ncols
                offset_val = offsets[dtmp] if use_offset else 0.0
                time_tmp = t[j] + offset_val * use_offset

                tmax = max(tmax, time_tmp)

        for time_step in [20, 1]:
            if time_step == 1:
                tmin = int(maximum_t[i] - 50)
                tmax = int(maximum_t[i] + 50)

            # get the value
            maximum_t[i], A[i], energy[i] = _get_waveform_maximum_impl(
                t,
                e,
                dist,
                r=r_tmp,
                z=z_tmp,
                template=template,
                pulse_shape_library=pulse_shape_library,
                templates_surface=templates_surface,
                activeness_surface=activeness_surface,
                tmin=tmin,
                tmax=tmax,
                start=start,
                fccd=fccd,
                n=n,
                time_step=time_step,
                surface_step_in_um=surface_step_in_um,
                include_surface_effects=include_surface_effects,
                use_library=use_library,
            )

    return A, maximum_t, energy


def prepare_surface_inputs(
    dist_to_nplus: ak.Array,
    edep: ak.Array,
    templates_surface: ArrayLike,
    activeness_surface,
    template: ArrayLike,
) -> tuple:
    """Prepare the inputs needed for surface sims."""
    include_surface_effects = False

    # prepare surface templates
    if templates_surface is not None:
        if dist_to_nplus is None:
            msg = "Surface effects requested but distance not provided"
            raise ValueError(msg)

        include_surface_effects = True
    else:
        # convert types to keep numba happy
        templates_surface = np.zeros((1, len(template)))
        dist_to_nplus = ak.full_like(edep, np.nan)

    # convert types for numba
    if activeness_surface is None:
        activeness_surface = np.zeros(len(template))

    return include_surface_effects, dist_to_nplus, templates_surface, activeness_surface


def prepare_pulse_shape_library(
    template: ArrayLike | HPGePulseShapeLibrary,
    times: ArrayLike,
    edep: ak.Array,
    r: ak.Array,
    z: ak.Array,
):
    """Prepare the inputs for the full pulse shape library."""
    use_library = False
    if isinstance(template, HPGePulseShapeLibrary):
        # convert to a form we can use
        times = template.t
        pulse_shape_library = (template.waveforms, template.r, template.z)
        template = np.zeros_like(template.waveforms[0][0])
        use_library = True

    else:
        pulse_shape_library = (np.zeros((1, 1, len(template))), np.zeros(1), np.zeros(1))
        r = ak.full_like(edep, np.nan)
        z = ak.full_like(edep, np.nan)

    return use_library, pulse_shape_library, template, times, r, z


def maximum_current(
    edep: ArrayLike,
    drift_time: ArrayLike,
    dist_to_nplus: ArrayLike | None = None,
    r: ArrayLike | None = None,
    z: ArrayLike | None = None,
    *,
    template: np.array | HPGePulseShapeLibrary,
    times: np.array,
    fccd_in_um: float = 0,
    templates_surface: ArrayLike | None = None,
    activeness_surface: ArrayLike | None = None,
    surface_step_in_um: float = 10,
    return_mode: str = "current",
) -> ak.Array:
    """Estimate the maximum current in the HPGe detector based on :func:`_estimate_current_impl`.

    Parameters
    ----------
    edep
       Energies for each step.
    drift_time
        Drift times for each step.
    dist_to_nplus
        Distance to n-plus electrode, only needed if surface heuristics are enabled.
    r
        Radial coordinate (only needed if a full PSS library is used)
    z
        z coordinate (only needed if a full PSS library is used).
    template
        Array of the bulk pulse template
    times
        time-stamps for the bulk pulse template
    fccd_in_um
        Value of the full-charge-collection depth, if `None` no surface corrections are performed.
    templates_surface
        2D array (distance, time) of the rate of charge arriving at the p-n junction. Each row
        should be an array of length 10000 giving the charge arriving at the p-n junction for each timestep
        (in ns). This is produced by :func:`.hpge.surface.get_surface_response` or other libraries.
    activeness_surface
        An array of the activeness at each surface point.
    surface_step_in_um
        Distance step for the surface library.
    return_mode
        either current, energy or max_time

    Returns
    -------
    An Array of the maximum current/ time / energy for each hit.
    """
    # convert to target units
    drift_time = units.units_conv_ak(drift_time, "ns")
    edep = units.units_conv_ak(edep, "keV")
    dist_to_nplus = units.units_conv_ak(dist_to_nplus, "mm")
    r = units.units_conv_ak(r, "mm")
    z = units.units_conv_ak(z, "mm")

    # prepare inputs for surface sims
    include_surface_effects, dist_to_nplus, templates_surface, activeness_surface = (
        prepare_surface_inputs(dist_to_nplus, edep, templates_surface, activeness_surface, template)
    )

    # and for the full PS library
    use_library, pulse_shape_library, template, times, r, z = prepare_pulse_shape_library(
        template, times, edep, r, z
    )

    # and now compute the current
    curr, time, energy = _estimate_current_impl(
        ak.values_astype(ak.Array(edep), np.float64),
        ak.values_astype(ak.Array(drift_time), np.float64),
        ak.values_astype(ak.Array(dist_to_nplus), np.float64),
        r=ak.values_astype(ak.Array(r), np.float64),
        z=ak.values_astype(ak.Array(z), np.float64),
        template=template,
        pulse_shape_library=pulse_shape_library,
        times=times,
        fccd=fccd_in_um,
        include_surface_effects=include_surface_effects,
        use_library=use_library,
        templates_surface=templates_surface,
        activeness_surface=activeness_surface,
        surface_step_in_um=surface_step_in_um,
    )

    # return
    if return_mode == "max_time":
        return units.attach_units(ak.Array(time), "ns")
    if return_mode == "current":
        # current has no unit (depends on the template)
        return ak.Array(curr)
    if return_mode == "energy":
        return units.attach_units(ak.Array(energy), "keV")

    msg = f"Return mode {return_mode} is not implemented."
    raise ValueError(msg)
