"""

"""
from __future__ import annotations

from math import exp, ceil
import numba as nb
import numpy as np


@nb.jit(nopython=True, nogil=True)
def convolve_lifetime_spectrum(
        decay: np.array,
        lifetime_spectrum: np.array,
        irf: np.array,
        stop: int = -1,
        time_axis: np.array = None,
) -> None:
    """Fills the pre-allocated array with a fluorescence intensity decay convolved (non-periodically)
    with an instrumental response function.


    This function calculates a fluorescence intensity decay that is convolved with an instrumnet response
    function (IRF). The fluorescence intensity decay is specified by its fluorescence lifetime spectrum, i.e.,
    an interleaved array containing fluorescence lifetimes with corresponding amplitudes.

    :param decay: numpy-array This array filled with the values of the computed fluorescence intensity decay
    :param lifetime_spectrum: interleaved vector of amplitudes and fluorescence lifetimes
    :param irf: the instrument response function
    :param stop: convolution stop channel (the index on the time-axis)
    :param time_axis: the time-axis of the decay

    Example
    -------

    >>> import scipy.stats
    >>> time_axis = np.linspace(0, 50, 2048)
    >>> irf_position = 5.0
    >>> irf_width = 1.0
    >>> irf = scipy.stats.norm.pdf(time_axis, loc=irf_position, scale=irf_width)
    >>> lifetime_spectrum = np.array([0.8, 1.1, 0.2, 4.0])
    >>> decay = np.zeros_like(time_axis)
    >>> convolve_lifetime_spectrum(decay, lifetime_spectrum=lifetime_spectrum, irf=irf, time_axis=time_axis)

    """
    number_of_exponentials = lifetime_spectrum.shape[0] // 2
    if stop <= 0:
        stop = decay.shape[0]
    for i in range(stop):
        decay[i] = 0.0
    for ne in range(number_of_exponentials):
        current_amplitude = lifetime_spectrum[2 * ne]
        if current_amplitude == 0.0:
            continue
        current_lifetime = (lifetime_spectrum[2 * ne + 1])
        if current_lifetime == 0.0:
            continue
        current_model_value = 0.0

        for i in range(1, stop):
            dt = (time_axis[i] - time_axis[i - 1])
            dt_2 = dt / 2.0
            current_exponential = exp(-dt / current_lifetime)
            current_model_value = (current_model_value + dt_2 * irf[i-1]) * current_exponential + dt_2 * irf[i]
            decay[i] += current_model_value * current_amplitude


@nb.jit(nopython=True, nogil=True)
def convolve_lifetime_spectrum_periodic(
        decay: np.array,
        lifetime_spectrum: np.array,
        irf: np.array,
        start: int,
        stop: int,
        n_points: int,
        period: float,
        dt: float,
        conv_stop: int
):
    """

    :param decay: array of doubles
        Here the convolved fit is stored
    :param lifetime_spectrum: array of doubles
        Lifetime-spectrum of the form (amplitude, lifetime, amplitude, lifetime, ...)
    :param irf: array-doubles
        The instrument response function
    :param start: int
        Start channel of convolution (position in array of IRF)
    :param stop: int
        Stop channel of convolution (position in array of IRF)
    :param n_points: int
        Number of points in fit and lamp
    :param period: double
        Period of repetition in nano-seconds
    :param dt: double
        Channel-width in nano-seconds
    :param conv_stop: int
        Stopping channel of convolution

    Example
    -------

    >>> import scipy.stats
    >>> n_points = 2048
    >>> time_axis = np.linspace(0, 16, n_points)
    >>> irf_position = 5.0
    >>> irf_width = 0.5
    >>> dt = time_axis[1] - time_axis[0]
    >>> irf = scipy.stats.norm.pdf(time_axis, loc=irf_position, scale=irf_width)
    >>> lifetime_spectrum = np.array([0.8, 1.1, 0.2, 4.0])
    >>> decay = np.zeros_like(time_axis)
    >>> convolve_lifetime_spectrum_periodic(decay, lifetime_spectrum=lifetime_spectrum, irf=irf, start=0, stop=n_points, n_points=n_points, period=16, conv_stop=n_points, dt=dt)

    """
    stop = min(stop, n_points - 1)
    start = max(start, 0)

    n_exp = lifetime_spectrum.shape[0] // 2
    period_n = ceil(period / dt - 0.5)

    for i in range(start, stop):
        decay[i] = 0

    stop1 = min(n_points, period_n)
    dt_2 = dt * 0.5

    # convolution
    for ne in range(n_exp):

        x_curr = lifetime_spectrum[2 * ne]
        lt_curr = lifetime_spectrum[2 * ne + 1]

        tail_a = 1./(1.-exp(-period/lt_curr))
        exp_curr = exp(-dt/lt_curr)

        fit_curr = 0.
        decay[0] += dt_2 * irf[0] * (exp_curr + 1.) * x_curr

        for i in range(conv_stop):
            fit_curr = (fit_curr + dt_2 * irf[i - 1]) * exp_curr + dt_2 * irf[i]
            decay[i] += fit_curr * x_curr

        for i in range(conv_stop, stop1):
            fit_curr *= exp_curr
            decay[i] += fit_curr * x_curr

        fit_curr *= exp(-(period_n - stop1) * dt / lt_curr)
        for i in range(stop):
            fit_curr *= exp_curr
            decay[i] += fit_curr * x_curr * tail_a

#
# @nb.jit(nopython=True, nogil=True)
# def fconv_per(
#         decay: np.array,
#         lifetime_spectrum: np.array,
#         irf: np.array,
#         start: int,
#         stop: int,
#         n_points: int,
#         period: float,
#         dt: float
# ):
#     dt_half = dt * 0.5
#     decay *= 0.0
#
#     x = lifetime_spectrum
#     n_exp = x.shape[0] // 2
#
#     period_n = int(ceil(period/dt-0.5))
#     stop1 = min(period_n, n_points - 1)
#
#     for ne in range(n_exp):
#         x_curr = lifetime_spectrum[2 * ne]
#         if x_curr == 0.0: continue
#         lt_curr = lifetime_spectrum[2 * ne + 1]
#         if lt_curr == 0.0: continue
#
#         exp_curr = exp(-dt / lt_curr)
#         tail_a = 1./(1.-exp(-period/lt_curr))
#
#         fit_curr = 0.0
#         for i in range(stop1):
#             fit_curr = (fit_curr + dt_half * irf[i-1]) * exp_curr + dt_half * irf[i]
#             decay[i] += fit_curr * x_curr
#         fit_curr *= exp(-(period_n - stop1 + start)*dt/lt_curr)
#
#         for i in range(start, stop):
#             fit_curr *= exp_curr
#             decay[i] += fit_curr * x_curr * tail_a
#     return 0
#
#
# @nb.jit(nopython=True, nogil=True)
# def fconv_per_dt(
#         decay: np.array,
#         lifetime_spectrum: np.array,
#         irf: np.array,
#         start: int,
#         stop: int,
#         n_points: int,
#         period: float,
#         time: np.array
# ):
#     # TODO: in future adaptive time-axis with increasing bin size
#     x = lifetime_spectrum
#     n_exp = x.shape[0]/2
#
#     #period_n = int(ceil(period/dt-0.5))
#
#
#     irf_start = 0
#     while irf[irf_start] == 0:
#         irf_start += 1
#
#     decay *= 0.0
#     #for i in range(stop):
#     #    decay[i] = 0.0
#
#     #stop1 = n_points - 1 if period_n + irf_start > n_points - 1 else period_n + irf_start
#     #if period_n + irf_start > n_points - 1:
#     #    stop1 = n_points - 1
#     #else:
#     #    stop1 = period_n + irf_start
#
#     for ne in range(n_exp):
#         x_curr = lifetime_spectrum[2 * ne]
#         if x_curr == 0.0:
#             continue
#         lt_curr = lifetime_spectrum[2 * ne + 1]
#         if lt_curr == 0.0:
#             continue
#
#         #tail_a = 1./(1.-exp(-period/lt_curr))
#
#         fit_curr = 0.0
#         for i in range(1, n_points):
#             dt = (time[i] - time[i-1])
#             dt_half = dt * 0.5
#             exp_curr = exp(-dt / lt_curr)
#             fit_curr = (fit_curr + dt_half * irf[i-1]) * exp_curr + dt_half * irf[i]
#             decay[i] += fit_curr * x_curr
#         #fit_curr *= exp(-(period_n - stop1 + start)*dt/lt_curr)
#
#         #for i in range(start, stop):
#         #    fit_curr *= exp_curr
#         #    decay[i] += fit_curr * x_curr * tail_a
#     return 0
#


@nb.jit(nopython=True, nogil=True)
def convolve_decay(
        decay,
        irf: np.array,
        start: int,
        stop: int,
        dt: float
) -> np.ndarray:
    """ Convolves a fluorescence decay with an instrument response function.

    This function convolves a fluorescence decay, that is provided as a numpy array,
    with an experimental response function.

    :param decay: fluorescence decay
    :param irf: instrument response function
    :param start: convolution start
    :param stop: convolution stop
    :param dt: bin-width of the fluorescence decay
    :return: convolved fluorescence decay as numpy array.

    Example
    -------

    >>> import scipy.stats
    >>> n_points = 2048
    >>> time_axis = np.linspace(0, 16, n_points)
    >>> irf_position = 5.0
    >>> irf_width = 0.5
    >>> dt = time_axis[1] - time_axis[0]
    >>> irf = scipy.stats.norm.pdf(time_axis, loc=irf_position, scale=irf_width)
    >>> decay_u = np.exp(- time_axis / 4.1)
    >>> decay = convolve_decay(decay_u, irf=irf, start=0, stop=n_points, dt=dt)

    """
    decay_out = np.empty_like(decay)
    for i in range(start, stop):
        decay_out[i] = 0.5 * irf[0] * decay[i]
        for j in range(1, i):
            decay_out[i] += irf[j] * decay[i - j]
        decay_out[i] += 0.5 * irf[i] * decay[0]
        decay_out[i] *= dt
    decay_out[0] = 0
    return decay_out

