from __future__ import annotations
from scikit_fluorescence import typing

import numpy as np
import numba as nb
import scikit_fluorescence as skf
import scikit_fluorescence.math


@nb.jit(nopython=True)
def rate_constant_to_lifetime(
        rate_constant: float,
        lifetime: float
):
    """

    Parameters
    ----------
    rate_constant : float
    lifetime : float

    Returns
    -------

    """
    return 1. / (1. / lifetime + rate_constant)


def fretrate_to_distance(fretrate, forster_radius, tau0, kappa2=2. / 3.):
    """Calculates the distance given a FRET-rate

    :param fretrate: FRET-rate
    :param forster_radius: Forster radius
    :param tau0: lifetime of the donor
    :param kappa2: orientation factor
    :return:
    """
    return forster_radius * (fretrate * tau0/kappa2 * 2./3.)**(-1./6)


@nb.jit(nopython=True)
def distance_to_fret_rate_constant(
        r: np.ndarray,
        forster_radius: float,
        tau0: float,
        kappa2: float = 2./3.
) -> np.ndarray:
    """Compute FRET rate constant for distances

    Parameters
    ----------
    r : array-like
        donor-acceptor distance
    forster_radius : float
        Forster-radius
    tau0 : float
        Fluorescence lifetime of the donor that corresponds to the fluorescence
        quantum yield of the donor used in the calculation of the Forster radius
        provided by the parameter `forster_radius`.
    kappa2 : float
        Orientation factor (default is 2/3)

    Returns
    -------
    FRET rate constants

    """
    return 3. / 2. * kappa2 * 1. / tau0 * (forster_radius / r) ** 6.0


def combine_interleaved_spectra(
        lifetime_spectra: typing.List[np.ndarray],
        fractions: typing.List[float] = None,
        normalize_fractions: bool = True,
        normalize_spectra: bool = False
) -> np.ndarray:
    """Combines a list of lifetime spectra in a joint spectrum

    Takes an array of lifetime spectra and an array of fractions and returns an
    mixed array of lifetimes whereas the amplitudes are multiplied by the
    fractions. `normalize_fractions` is True the fractions are normalized to one.

    Parameters
    ----------
    lifetime_spectra: a list of lifetime spectra
    fractions: numpy-array
        amplitudes / fraction that by which the spectra will be multiplied
    normalize_fractions: bool
        If set to True the fractions will be normalized to one
    normalize_spectra: bool
        If set to True the amplitudes of the spectra will be normalized to one
        before they are scaled by the fractions and combined.

    Returns
    -------
    lifetime spectrum: numpy-array

    See Also
    --------

    scikit_fluorescence.math.datatools.two_column_to_interleaved :

    """
    if fractions is None:
        fractions = np.ones(len(lifetime_spectra), dtype=np.float64)
    if normalize_fractions:
        fn = np.array(fractions, dtype=np.float64) / np.sum(fractions)
    else:
        fn = fractions
    scale = 1.0
    re = list()
    for i, ls in enumerate(lifetime_spectra):
        ls = np.copy(ls)
        if normalize_spectra:
            scale = np.sum(ls[::2])
        if scale > 0.0:
            ls[::2] = ls[::2] * fn[i] / scale
        re.append(ls)
    return np.hstack(re)


def bin_lifetime_spectrum(
        lifetime_spectrum: np.array,
        n_lifetime_bins: int,
        lifetime_range: (float, float) = (0.001, 10.0)
) -> np.ndarray:
    """Takes a interleaved lifetime spectrum

    :param lifetime_spectrum: interleaved lifetime spectrum
    :param n_lifetime_bins:
    :return: lifetime_spectrum
    """
    amplitudes, lifetimes = skf.math.datatools.interleaved_to_two_columns(
        lifetime_spectrum,
        sort=False
    )
    lt, am = skf.math.datatools.histogram1D(
        values=lifetimes,
        weights=amplitudes,
        n_bins=n_lifetime_bins,
        tth_min=lifetime_range[0],
        tth_max=lifetime_range[1]
    )
    binned_lifetime_spectrum = skf.math.datatools.two_column_to_interleaved(
        x=am,
        t=lt
    )
    return binned_lifetime_spectrum


def compute_mean_fret(
        lifetime_spectrum: np.ndarray,
        lifetime_spectrum_donor_only: np.ndarray,
        forster_radius: float,
        fluorescence_lifetime_donor: float,
        min_Rapp_sd: float = 1.0,
        use_longest_donor_lifetime: bool = False,
        computed_value: str = 'Rapp',
        verbose: int = False
) -> typing.Tuple[float, float]:
    """Computes a mean FRET rate constant with corresponding width that can be
    represented as apparent distance with corresponding width.

    This function computes for the fluorescence lifetime spectrum of the donor
    in the presence and in the absence of FRET an average distance with corresponding
    width. The average distance is computed by first calcaulting the mean (species)
    average fluorescence lifetime of the donor in the absence of FRET, :math:`<\tau_{D0}>_x`

    :math:`<tau_{D0}>_x = \sum x_i \tau_{D0,i}`.

    Next, for all fluorescence lifetimes in the fluorescence lifetime spectrum of
    the donor in the presence of FRET the corresponding FRET rate constants
    :math:`k_{RET, i}` are computed:

    :math:`k_{RET,i} = 1/ \tau_{DA,i} - 1/<\tau_{D0}>_x`.

    The FRET rate constants :math:`k_{RET,i}` are expressed as apparent
    distances :math:`R_{app,i}` using the Forster radius and the corresponding
    fluorescence lifetime of the donor

    :math:`R_i = R_0(k_{RET} \cdot \tau_0)^{-1/6}`.

    Finally, the apparent distances with corresponding fractions :math:`x_i` are
    used to compute an average apparent distance

    :math:`<R> = \sum x_i R_i`

    and an standard deviation :math:`\sigma_{DA}` of the average apparent distance

    :math:` \sigma_{DA} = \sum x_i R_i^{2} - <R>^2`

    Fluorescence lifetimes that are longer than :math:`<tau_{D0}>_x` are ignored.

    Parameters
    ----------
    lifetime_spectrum: numpy-array
        An interleaved fluorescence lifetime spectrum of the donor in the presence
        of FRET
    lifetime_spectrum_donor_only: numpy-array
        Fluorescence lifetime spectrum of the donor in the absence of FRET
    forster_radius: float
        The forster radius that corresponds to the donor fluorescence lifetime
        in the absence of FRET

    Notes
    -----
    The returned mean and standard deviation are coarse approximations.

    Returns
    -------
    A mean apparent distance with corresponding distribution width (sigma)

    Examples
    >>> import scikit_fluorescence as skf
    >>> import scikit_fluorescence.data
    >>> decay_object = skf.data.decay_object()
    >>> lifetime_spectrum = decay_object.lifetime_spectrum


    """
    if not use_longest_donor_lifetime:
        kD0 = 1. / skf.decay.lifetime.species_averaged_lifetime(lifetime_spectrum_donor_only)
    else:
        _, lifetimes = skf.math.datatools.interleaved_to_two_columns(lifetime_spectrum_donor_only)
        kD0 = 1. / np.max(lifetimes)
    value = list()
    amplitude = list()
    if verbose > 0:
        print("COMPUTE_MEAN_FRET")
        print("-- max(tauD0): {:.2f}".format(1./kD0))
    amplitudes, lifetimes = skf.math.datatools.interleaved_to_two_columns(lifetime_spectrum)
    for a, lifetime in zip(amplitudes, lifetimes):
        kDA = 1. / lifetime
        if kDA > kD0:
            kRET = kDA - kD0
            if computed_value == "Rapp":
                apparent_distance = forster_radius * (kRET * fluorescence_lifetime_donor)**(-1./6.)
                value.append(apparent_distance)
            else:
                value.append(kRET)
            amplitude.append(a)
    value = np.array(value)
    amplitudes = np.array(amplitude)
    amplitudes /= sum(amplitude)
    mean_distance = float(np.dot(amplitudes, value))
    mean_distance2 = np.dot(amplitudes, value**2.0)
    sd_distance = float(np.sqrt(mean_distance2 - mean_distance**2.0))
    if verbose > 0:
        print("-- mean_distance: {:.2f}".format(mean_distance))
        print("-- sd_distance: {:.2f}".format(sd_distance))
    if computed_value == 'Rapp':
        if sd_distance < min_Rapp_sd:
            UserWarning("SD of distance smaller then minimum SD")
            sd_distance = min_Rapp_sd
    return mean_distance, sd_distance

