from __future__ import annotations
from scikit_fluorescence import typing

import numpy as np
import numba as nb


@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
def distance_to_fret_rate_constant(
        r: np.ndarray,
        forster_radius: float,
        tau0: float,
        kappa2: float = 2./3.
) -> np.ndarray:
    """ Converts the DA-distance to a FRET-rate

    :param r: donor-acceptor distance
    :param forster_radius: Forster-radius
    :param tau0: lifetime
    :param kappa2: orientation factor
    :return:
    """
    return 3. / 2. * kappa2 * 1. / tau0 * (forster_radius / r) ** 6.0


def combine_lifetime_spectra(
        lifetime_spectra: typing.List[np.ndarray],
        fractions: typing.List[float] = None,
        normalize_fractions: bool = True
) -> np.ndarray:
    """Joins a list of lifetime spectra

    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.

    :return: numpy-array

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

