from __future__ import annotations
from scikit_fluorescence import typing

import numpy as np
import scikit_fluorescence as skf
import scikit_fluorescence.decay


@skf.memory.cache()
def compute_lifetime_spectra(
        atoms: np.ndarray,
        donor_av_parameter: typing.Dict,
        acceptor_av_parameter: typing.Dict,
        donor_lifetime_spectrum: np.ndarray,
        forster_radius: float,
        tau0: float,
        kappa2: float = 0.6667,
        kappa2_model: str = 'dynamic',
        donor_only_fraction: float = 0.0,
        bin_lifetime_spectrum: bool = True,
        n_lifetime_bins = 101,
        lifetime_range = (0.001, 5)
) -> typing.Dict[str, np.ndarray]:
    """
    Compute AVs
    """
    av1 = skf.modeling.label.av.LabelDistributionAV(
        atoms=atoms,
        **donor_av_parameter
    )
    av2 = skf.modeling.label.av.LabelDistributionAV(
        atoms=atoms,
        **acceptor_av_parameter
    )
    donor_lifetime_spectrum_with_fret = av1.get_donor_fret_lifetime_spectrum(
        av=av2,
        forster_radius=forster_radius,
        kappa2=kappa2,
        kappa2_model=kappa2_model,
        tau0=tau0,
        donor_lifetime_spectrum=donor_lifetime_spectrum
    )

    """
    Add donor only
    """
    fret_fraction = 1. - donor_only_fraction
    complete_lifetime_spectrum = skf.decay.rate_spectra.combine_interleaved_spectra(
        lifetime_spectra=[
            donor_lifetime_spectrum,
            donor_lifetime_spectrum_with_fret
        ],
        fractions=[donor_only_fraction, fret_fraction],
        normalize_fractions=True,
        normalize_spectra=True
    )
    re = dict()
    if bin_lifetime_spectrum:
        re['donor'] = skf.decay.rate_spectra.bin_lifetime_spectrum(
            lifetime_spectrum=complete_lifetime_spectrum,
            n_lifetime_bins=n_lifetime_bins,
            lifetime_range=lifetime_range
        )
    else:
        re['donor'] = complete_lifetime_spectrum
    return re

