from __future__ import annotations
from scikit_fluorescence import typing

import pathlib
import json
import numpy as np
import numba as nb
import fit2x

import scikit_fluorescence as skf
import scikit_fluorescence.decay.rate_spectra
import scikit_fluorescence.decay.convolve
import scikit_fluorescence.decay.corrections
import scikit_fluorescence.decay.lifetime
import scikit_fluorescence.decay.tcspc
import scikit_fluorescence.decay.analysis
import scikit_fluorescence.modeling
import scikit_fluorescence.math


class Decay(fit2x.Decay):

    def optimize(self, **kwargs):
        skf.decay.analysis.optimize_decay_lifetimes(
            self,
            **kwargs
        )

def scale_decay_model_to_experiment(
        model_decay: np.ndarray,
        experimental_decay: np.ndarray,
        weights: np.ndarray = None,
        start: int = 0,
        stop: int = None,
        experimental_background: float = 0.0,
        scale_inplace: bool = True
) -> float:
    """Scales a model decay to an experimental decay considering the noise and
    a constant background in experimental data for a specified range.

    A scaling factor for a model decay and an experimental decay is computed such,
    that the number of photons in the model decay match the number of photons
    in the experimental decay considering the uncertainty of the experimental
    decay that is provided as parameter.

    If the weights of the experimental decay are not provided as an input the
    weights are calculated assuming Poisson distributed counts in the
    experimental decay.

    The scaling factor is computed for a range specified by `start` and `stop`.

    Parameters
    ----------
    model_decay : numpy.array
        Model decay for that a scaling factor is computed.
    experimental_decay : numpy.array
        Experimental decay to which the model decay provided by `model_decay`
        is scaled by the returned floating number
    weights : numpy.array
        Weights of the experimental decay that are used scale the model to
        the experiment. In case no weights are provided, the weights are
        computed assuming a Poisson noise in the experimental data.
    experimental_background : float
        Constant offset in the experimental data that is subtracted from the
        experimental decay
    start : int
        Start index that defines the range in which the model decay is scaled
        to the experimental decay
    stop : int
        Stop index that defines the range in which the model decay is scaled
        to the experimental decay.
    scale_inplace : bool
        If True (default) the experimental decay provided by the parameter
        `experimental_decay` is scaled inplace (modified) to the experimental
        decay.

    Returns
    -------
    float
        The scaling factor that was used to scale the model function to the
        experimental decay.

    Examples
    --------

    """
    if not scale_inplace:
        model = np.copy(model_decay)
    else:
        model = model_decay
    if stop is None:
        stop = len(model_decay) - 1
    if weights is None and experimental_background == 0.0:
        scale = fit2x.rescale(
            fit=model,
            decay=experimental_decay.astype(np.float64),
            start=start,
            stop=stop
        )
    else:
        if weights is None:
            weights = np.ones_like(experimental_decay)
        scale = fit2x.rescale_w_bg(
            fit=model,
            decay=experimental_decay.astype(np.float64),
            w_sq=weights**2,
            bg=experimental_background,
            start=start,
            stop=stop
        )
    return scale


@nb.jit(nopython=True)
def fret_induced_donor_decay(
        fd0: np.ndarray,
        fda: np.ndarray
) -> np.ndarray:
    """Calculates the FRET induced donor decay

    :param fd0: the fluorescence decay of the donor in the absence of FRET
    :param fda: the fluorescence decay of the donor in the presence of FRET
    :return:
    """
    return fda / fd0


class LifetimeBinParameter(typing.TypedDict):
    bin_lifetime_spectrum: bool
    n_lifetime_bins: int
    lifetime_range: typing.Tuple[float, float]


class DecayParameter(typing.TypedDict):

    lifetime_spectrum: np.ndarray
    excitation_period: float
    scale_model_to_data: bool
    time_axis: np.ndarray
    irf_histogram: np.ndarray
    decay_histogram: np.ndarray
    convolution_range: typing.Tuple[int, int]
    use_amplitude_threshold: bool
    amplitude_threshold: float
    add_pile_up: bool
    dt: float
    irf_background_counts: float
    scatter_fraction: float
    constant_offset: float
    score_range: typing.Tuple[int, int]
    instrument_dead_time: float
    use_corrected_irf_as_scatter: bool
    acquisition_time: float
    number_of_photons: float
    irf_shift_channels: float


class FRETParameter(typing.TypedDict):
    forster_radius: float
    tau0: float
    kappa2: float
    kappa2_model: str  # either static of dynamic


class DonorReference(FRETParameter):
    lifetime_spectrum: skf.typing.List[float]


def score_structure(
        atoms: np.ndarray,
        donor_only_fraction: float,
        donor_lifetime_spectrum: np.ndarray,
        donor_av_parameter: typing.Dict,
        acceptor_av_parameter: typing.Dict,
        decay_parameter: skf.decay.DecayParameter,
        fret_parameter: skf.decay.FRETParameter,
        lifetime_bin_parameter: skf.decay.LifetimeBinParameter,
        binned_lifetime_spectrum: np.ndarray = None,
        dye: str = 'donor',
        increment_lifetime_spectrum: bool = False
) -> typing.Tuple[float, skf.decay.Decay]:
    if increment_lifetime_spectrum and binned_lifetime_spectrum is None:
        raise ValueError("For increment_lifetime_spectrum binned_lifetime_spectrum "
                         "needs to be provided.")
    new_lifetime_spectrum = skf.modeling.structure.compute_lifetime_spectra(
        donor_av_parameter=donor_av_parameter,
        acceptor_av_parameter=acceptor_av_parameter,
        atoms=atoms,
        donor_lifetime_spectrum=donor_lifetime_spectrum,
        donor_only_fraction=donor_only_fraction,
        **fret_parameter,
        **lifetime_bin_parameter
    )[dye]
    if increment_lifetime_spectrum:
        binned_lifetime_spectrum += new_lifetime_spectrum
    else:
        binned_lifetime_spectrum = new_lifetime_spectrum
    decay_parameter['lifetime_spectrum'] = binned_lifetime_spectrum
    decay_object = skf.decay.Decay(**decay_parameter)
    return decay_object.get_score(), decay_object


def compute_fret_rate_constants(
        p_rda: skf.typing.List[float],
        rda_axis: skf.typing.List[float],
        forster_radius: float,
        tau0: float,
        kappa2: float = 2. / 3.,
        kappa2_model: str = 'dynamic',
        kappa2_min: float = 0.1,
        kappa2_max: float = 4,
        kappa2_bins: int = 32
) -> np.ndarray:
    """Compute FRET rate constant spectra for distance distribution

    Parameters
    ----------
    p_rda : array-like
        Population of the distances (frequency)
    rda_axis : array-like
        Distance axis corresponding to the populations
    forster_radius : float
        The Forster radius used to convert the distance axis to a FRET rate constant
        axis
    tau0 : float
        The fluorescence lifetime of the donor that corresponds to the fluorescence
        quantum yiel of the donor that was used to compute the Forster radius
    kappa2 : float
        The mean orientation factor that is used to compute the FRET rate constants.
        This parameter is used if for the dynamic orientation averaging kappa2 model
        was selected
    kappa2_model : str
        A parameter to select the orientation factor model. The default selection
        is the `dynamic` averaging. When dynamic averaging is used a single mean
        orientation factor is used to compute the distribution of FRET rate constants.
        Alternatively, *static* averaging can be used. In case of static averaging
        an orientation factor distribution is computed and the FRET rate constants
        are computed for all kappa2 of the distribution and weighted according to
        the frequency of a kappa2 in the kappa2 distribution.
    kappa2_min : float
        The minimum kappa2 value in the kappa2 distribution computed for static
        averaging.
    kappa2_max : float
        The maximum kappa2 value in the kappa2 distribution computed for static
        averaging.
    kappa2_bins : int
        The number of bins in the kappa2 distribution.

    Returns
    -------
    FRET rate constant spectrum (amplitude, FRET rate constant, ...)

    """
    # Make sure the inputs for the later numba functions are numpy arrays
    if not isinstance(p_rda, np.ndarray):
        p_rda = np.array(p_rda, dtype=np.float64)
    if not isinstance(rda_axis, np.ndarray):
        rda_axis = np.array(rda_axis, dtype=np.float64)

    # compute FRET rate constants that belong to the distances
    if kappa2_model == 'static':
        k2_range = np.linspace(kappa2_min, kappa2_max, kappa2_bins)
        p_k2 = skf.modeling.kappa2.p_isotropic_orientation_factor(k2=k2_range)
        return skf.decay.rate_spectra.combine_interleaved_spectra(
            lifetime_spectra=[
                skf.math.datatools.two_column_to_interleaved(
                    p_rda, skf.decay.rate_spectra.distance_to_fret_rate_constant(
                        r=rda_axis,
                        forster_radius=forster_radius,
                        tau0=tau0,
                        kappa2=k2
                    )
                ) for k2 in k2_range
            ],
            fractions=list(p_k2)
        )
    else:
        return skf.math.datatools.two_column_to_interleaved(
            p_rda, skf.decay.rate_spectra.distance_to_fret_rate_constant(
                r=rda_axis,
                forster_radius=forster_radius,
                tau0=tau0,
                kappa2=kappa2
            )
        )


def compute_donor_lifetime_spectrum_with_fret(
        donor_lifetime_spectrum: skf.typing.List[float],
        input: skf.typing.List[float],
        axis: skf.typing.List[float],
        forster_radius: float,
        tau0: float,
        input_type: str = "rda",
        verbose: bool = False,
        **kwargs
) -> np.ndarray:
    """Compute donor fluorescence lifetimes in the presence of FRET for distance
    distribution

    Parameters
    ----------
    donor_lifetime_spectrum : array_like
        The parameter `donor_lifetime_spectrum` is the interleaved lifetime
        spectrum (amplitude, lifetime, ...) of the donor in the absence of FRET.
    input : array_like
        The probability of a certain distance
    axis : array_like
        The distance axis
    forster_radius : float
        The Forster radius of the dye pair
    tau0 : float
        The fluorescence lifetime of the donor in the absence of FRET that was
        that corresponds to the fluorescence quantum yield that was used to
        compute the Forster radius provided by the parameter `forster_radius`.
    kwargs : dict
        Parameters that are passed to `scikit_fluorescence.decay.compute_fret_rate_constants`

    Returns
    -------
    Fluorescence lifetime spectrum of the donor in the presence of FRET (amplitdude, lifetime, ...)

    Examples
    --------
    >>> import scikit_fluorescence as skf
    >>> import scikit_fluorescence.decay
    >>> import numpy as np
    >>> donor_lifetime_spectrum = [1., 4.]
    >>> forster_radius = 52.0
    >>> tau0 = 4.0
    >>> rda_axis = [43.3, 52.0, 62.4]
    >>> p_rda = [0.5, 0.5, 0.5]
    >>> skf.decay.compute_donor_lifetime_spectrum_with_fret(
    ...   donor_lifetime_spectrum=donor_lifetime_spectrum,
    ...   tau0=tau0, forster_radius=forster_radius,
    ...   axis=rda_axis, input=p_rda, input_type='rda'
    ... )
    array([0.5       , 1.00004934, 0.5       , 2.        , 0.5       ,
       2.99648368])
    
    >>> kfret_axis = [0.1, 0.5, 10]
    >>> kfret = [0.5, 0.5, 0.5]
    >>> skf.decay.compute_donor_lifetime_spectrum_with_fret(
    ...   donor_lifetime_spectrum=donor_lifetime_spectrum,
    ...   tau0=tau0, forster_radius=forster_radius,
    ...   axis=rda_axis, input=p_rda, input_type='kfret' # 'rda' is default
    ... )
    
    
    See Also
    --------
    scikit_fluorescence.decay.compute_fret_rate_constants : Computing FRET rate
        constants from distance distributions.

    """
    if verbose:
        print("COMPUTE_DONOR_LIFETIME_SPECTRUM_WITH_FRET")
        print("-- rda_axis:", axis)
        print("-- p_rda:", input)
        print("-- forster_radius:", forster_radius)
        print("-- tau0:", tau0)
    if not isinstance(donor_lifetime_spectrum, np.ndarray):
        donor_lifetime_spectrum = np.array(donor_lifetime_spectrum, dtype=np.float64)
    if input_type == "rda":
        fret_rate_spectrum = compute_fret_rate_constants(
                p_rda=input,
                rda_axis=axis,
                forster_radius=forster_radius,
                tau0=tau0,
                **kwargs
            )
    elif input_type == "kfret":
        fret_rate_spectrum = skf.math.datatools.two_column_to_interleaved(input, axis)
    lt = skf.math.datatools.invert_interleaved(
        skf.math.datatools.ere2(
            fret_rate_spectrum,
            skf.math.datatools.invert_interleaved(donor_lifetime_spectrum)
        )
    )
    return lt


def mix_lifetime_spectra(
        decay_object: skf.decay.Decay,
        lifetime_spectra: skf.typing.List[np.ndarray],
        fractions: skf.typing.List[float],
        fraction_of_existing_spectrum: float = 0.0,
        normalize_fractions: bool = True,
        normalize_spectra: bool = True
):
    if fraction_of_existing_spectrum > 0.0:
        lifetime_spectra.append(np.copy(decay_object.lifetime_spectrum))
        fractions.append(fraction_of_existing_spectrum)
    decay_object.lifetime_spectrum = skf.decay.rate_spectra.combine_interleaved_spectra(
            lifetime_spectra=lifetime_spectra,
            fractions=fractions,
            normalize_fractions=normalize_fractions,
            normalize_spectra=normalize_spectra
    )

