from __future__ import annotations

import numpy as np
import pylab as plt
import matplotlib.offsetbox

import scipy.optimize
import scipy.stats

import scikit_fluorescence as skf
import scikit_fluorescence.decay
import scikit_fluorescence.math


def is_linear_time_axis(
        time_axis: np.ndarray,
        rtol: float = 1.e-5,
        atol: float = 1.e-8
) -> bool:
    """True if the time axis is linear

    Parameters
    ----------
    time_axis: numpy.array
        A time axis

    Returns
    -------
    True if the all elements of the time axis derivative are identical

    """
    dt = np.diff(time_axis)
    dt0 = dt[0]
    for jdt in dt[1:]:
        if not np.isclose(jdt, dt0, rtol=rtol, atol=atol):
            return False
    return True


def split_lifetime(
        x: float,
        lt: float,
        lt_max: float,
        lt_min: float
):
    """Split a lifetime with associated amplitude into two lifetimes

    The fluorescence average lifetime of the new two lifetimes corresponds
    to the original lifetime

    (x1*lt1**2 + x2 * lt2**2)/(x1*lt1+x2*lt2) = lt

    and lt1 and lt2 are half the distance to lt in the ranges (lt_min, lt) and
    (lt, lt_max)

    |ltmin     lt1       lt   lt2    |ltmax

    Parameters
    ----------
    x: float
    lt : float
    lt_max : float
    lt_min : float

    Returns
    -------


    """
    lt2 = lt - (lt - lt_min) / 2.0
    lt1 = lt + (lt_max - lt) / 2.0
    x1 = (lt2 * (lt2 - lt)) / ((lt2 - lt1) * (-lt + lt1 + lt2))
    return (1-x1) * x, lt2, x1 * x, lt1


def split_largest_lifetime(
        lifetime_spectrum: np.ndarray,
        tau_max: float = 5.0,
        tau_min: float = 1.5,
        idx: int = -1,
        selection: str = 'smallest'
):
    lifetime_spectrum = np.array(lifetime_spectrum)
    xs, lts = skf.math.datatools.interleaved_to_two_columns(lifetime_spectrum)
    if idx < 0:
        if selection == 'smallest':
            idx = np.argmin(xs)
        else:
            idx = np.argmax(xs)
    xn = list()
    ltn = list()
    for i, x in enumerate(xs):
        if i != idx:
            xn.append(x)
            ltn.append(lts[i])
    lt = lts[idx]
    x = xs[idx]
    x1, l1, x2, l2 = split_lifetime(x, lt, tau_max, tau_min)
    ltn.append(l1)
    xn.append(x1)
    ltn.append(l2)
    xn.append(x2)
    new_lts = skf.math.datatools.two_column_to_interleaved(
        np.array(xn),
        np.array(ltn)
    )
    return new_lts


def split_lifetime_spectrum(
            lifetime_spectrum: np.ndarray,
            tau_max: float = 5.0,
            tau_min: float = 0.5
    ):
    new_lifetime_spectrum = np.empty(len(lifetime_spectrum) * 2, dtype=np.float64)
    for i, (x, lt) in enumerate(
        zip(lifetime_spectrum[::2], lifetime_spectrum[1::2])
    ):
        lt1 = lt - (lt - tau_min) / 2.0
        lt2 = lt + (tau_max - lt) / 2.0
        x1 = (lt2 * (lt2 - lt)) / ((lt2 - lt1) * (-lt + lt1 + lt2))
        new_lifetime_spectrum[i * 4 + 0] = x1 * x
        new_lifetime_spectrum[i * 4 + 1] = lt1
        new_lifetime_spectrum[i * 4 + 2] = (1. - x1) * x
        new_lifetime_spectrum[i * 4 + 3] = lt2
    return new_lifetime_spectrum


def normalize_lifetime_spectrum(
        decay_object: skf.decay.Decay,
        take_absolute_value_lifetime_spectrum: bool = True,
        normalize_lifetime_amplitudes: bool = True
):
    lt = np.copy(decay_object.lifetime_spectrum)
    if take_absolute_value_lifetime_spectrum:
        lt = np.abs(lt)
    if normalize_lifetime_amplitudes:
        lt[::2] /= np.sum(lt[::2])
    decay_object.lifetime_spectrum = lt


def optimize_decay_lifetimes(
        decay_object: skf.decay.Decay,
        fixed: skf.typing.List[bool] = [True, True, True, True],
        normalize_lifetime_amplitudes: bool = True,
        take_absolute_value_lifetime_spectrum: bool = True,
        score_type: str = 'poisson',
        verbose: bool = False,
        plot_optimized_decays: bool = False,
        method: str = 'leastsqbound',
        use_bounds: bool = False,
        fix_lifetimes = False
) -> float:
    """

    Initial values for the fit are take from the skf.decay.Decay object passed
    by the parameter decay_object.

    Parameters
    ----------
    decay_object
    constant_offset
    irf_shift
    irf_background_counts
    scatter_fraction
    n_lifetimes
    score_range
    normalize_amplitudes
    verbose
    plot_optimized_decays
    method

    Returns
    -------

    """

    def objective_function(
            x: np.array,
            fixed: skf.typing.Tuple[bool],
            decay_object: skf.decay.Decay,
            return_type: str = 'score',
            score_type: str = 'poisson'
    ):
        lt = x[4:]
        if not fixed[0]:
            decay_object.constant_offset = x[0]
        if not fixed[1]:
            decay_object.irf_shift_channels = x[1]
        if not fixed[2]:
            decay_object.irf_background_counts = x[2]
        if not fixed[3]:
            decay_object.scatter_fraction = x[3]
        decay_object.lifetime_spectrum = lt
        decay_object.update()
        score = decay_object.get_score(score_type=score_type)
        if verbose:
            print("x:", x)
            print("fixed:", fixed)
            print("Score:", score)
        if return_type == 'residuals':
            i_min, i_max = decay_object.score_range
            wres = decay_object.weighted_residuals
            return np.copy(wres)
        else:
            return score

    x0 = np.hstack(
        [
            decay_object.constant_offset,
            decay_object.irf_shift_channels,
            decay_object.irf_background_counts,
            decay_object.scatter_fraction,
            np.copy(decay_object.lifetime_spectrum)
        ]
    )
    while len(fixed) < len(x0):
        fixed += [fix_lifetimes, ]
    n_lifetime_spectrum = len(x0) - 4
    if use_bounds:
        bounds = [
                     (0, 1e6),  # offset
                     (None, None),  # irf shift
                     (0, 1e6),  # irf bg
                     (0, 1.0)  # scatter
                 ] + [(None, None)] * n_lifetime_spectrum
    else:
        bounds = None
    if method == 'L-BFGS-B':
        fit_result = scipy.optimize.minimize(
            objective_function,
            x0=x0,
            args=(fixed, decay_object, 'score', score_type, ),
            bounds=bounds,
            method=method
        )
        x = fit_result.x
    elif method == 'leastsqbound':
        fit_result = skf.math.optimize.leastsqbound.leastsqbound(
            objective_function,
            x0=x0,
            full_output=True,
            bounds=bounds,
            args=(fixed, decay_object, 'residuals',)
        )
        x = fit_result[0]
    else:
        fit_result = scipy.optimize.minimize(
            objective_function,
            x0=x0,
            bounds=bounds,
            args=(fixed, decay_object, 'score', score_type, )
        )
        x = fit_result.x

    # update the decay object
    score = objective_function(x, fixed, decay_object, 'score', score_type)
    normalize_lifetime_spectrum(
        decay_object,
        normalize_lifetime_amplitudes=normalize_lifetime_amplitudes,
        take_absolute_value_lifetime_spectrum=take_absolute_value_lifetime_spectrum
    )
    if verbose:
        print('FIT_DECAY_LIFETIMES')
        print("-- mean_lifetime computed from data:", decay_object.get_mean_lifetime(decay_object._data))
        print('-- score_range:', decay_object.score_range)
        print('-- normalize_decay:', normalize_lifetime_amplitudes)
        print('-- score:', score)
    if plot_optimized_decays:
        i_min, i_max = decay_object.score_range
        fig, ax = plt.subplots(2, 1, True, False)
        ax[1].semilogy(decay_object.time_axis, decay_object.corrected_irf)
        ax[1].semilogy(decay_object.time_axis, decay_object.data)
        ax[1].semilogy(decay_object.time_axis[i_min: i_max], decay_object.model[i_min: i_max])
        ax[0].plot(
            decay_object.time_axis[i_min:i_max],
            decay_object.weighted_residuals
        )
        ax[0].text(0, 0, "score: {:.2f}".format(score))
        plt.show()
    return score


def shrink_analysis_range(
        decay_object: skf.decay.Decay,
        irf_fwhm_range: skf.typing.Tuple[int, int],
        increment: int = 4,
        irf_fwhm_added_range_offset: int = 0,
        verbose: bool = False,
        plot_shrink_analysis_range: bool = False
):
    original_score_range = decay_object.score_range
    upper = original_score_range[1]
    upperbound_score_range, _ = irf_fwhm_range
    upperbound_score_range += irf_fwhm_added_range_offset
    original_score = decay_object.get_score()
    scores = list()
    scan_range = range(original_score_range[0], upperbound_score_range, increment)
    # constant offset, irf shift, irf background counts, scatter fraction
    # fixed = [True, False, True, False]
    # fixed += [False] * len(decay_object.lifetime_spectrum)
    for lower in scan_range:
        decay_object.score_range = lower, upper
        score = optimize_decay_lifetimes(decay_object) #, fixed=fixed)
        scores.append(score)
    if plot_shrink_analysis_range:
        plt.plot(scan_range, scores)
        plt.show()
    best_score_range = scan_range[np.argmin(scores)], upper
    decay_object.score_range = best_score_range
    best_score = optimize_decay_lifetimes(decay_object)
    if verbose:
        print("RESTICT_FITTING_RANGE")
        print("-- original_score_range:", original_score_range)
        print("-- original score:", original_score)
        print("-- best_score_range:", best_score_range)
        print("-- best_score:", best_score)


def get_analysis_range(
        fluorescence_decay: np.ndarray,
        count_threshold: float = 10.0,
        area: float = 0.95,
        start_at_peak: bool = True,
        start_fraction: float = 0.8,
        skip_first_channels: int = 0,
        skip_last_channels: int = 0,
        verbose: bool = False
) -> skf.typing.Tuple[int, int]:
    """Determines a fitting range based on the total number of photons to be
    fitted (fitting area).

    Parameters
    ----------
    fluorescence_decay : numpy-array
        a numpy array containing the photon counts
    count_threshold : float
        a threshold value. Lowest index of the fitting range is the first
        encountered bin with a photon count higher than the threshold.
    area : float
        The area which should be considered for fitting. Here 1.0 corresponds
        to all measured photons. 0.9 to 90% of the total measured photons.
    count_threshold : float
         The minimum number of photons in a channel (default value = 10.0)
    area: float
         (Default value = 0.999)
    start_at_peak: bool
        If set to True (default is False) sets the initial value of the fit range
        to min(start_fraction * max(fluorescence_decay))
    start_fraction: float
        Is the fraction of the maximum of the fluorescence decay at which the
        fit range starts (if start_at_peak is set to True)
    verbose: bool
        If set to True (default is False) prints the analysis range to the
        standard output.

    Returns
    -------
        A tuple of integers (start, stop)
    """
    if start_at_peak:
        x_min = np.where(fluorescence_decay > start_fraction * np.max(fluorescence_decay))[0][0]
    else:
        x_min = np.where(fluorescence_decay > count_threshold)[0][0]
    f_min = fluorescence_decay[x_min:]
    cumsum = np.cumsum(f_min, dtype=np.float64)
    s = np.sum(f_min, dtype=np.float64)
    n_decay = len(fluorescence_decay)
    x_max = np.where(cumsum >= s * area)[0][0] + x_min
    if fluorescence_decay[x_max] < count_threshold:
        x_max = np.searchsorted(fluorescence_decay[::-1], [count_threshold])[-1]
    x_max = min(n_decay, x_max)
    while True:
        if x_max >= n_decay:
            break
        if fluorescence_decay[x_max] < count_threshold:
            break
        else:
            x_max += 1
    a_min = x_min + skip_first_channels
    a_max = min(x_max, len(fluorescence_decay) - skip_last_channels - 1)

    if verbose:
        print("GET_ANALYSIS_RANGE")
        print("-- count_threshold:", count_threshold)
        print("-- area:", area)
        print("-- start_at_peak:", start_at_peak)
        print("-- start_fraction:", start_fraction)
        print("-- analysis range: ", a_min, a_max)
    return a_min, a_max


def add_lifetime(
        decay_object: skf.decay.Decay,
        add_new_method: str = 'split',
        minimum_initial_lifetime: float = 0.25,
        maximum_initial_lifetime: float = None,
        verbose: bool = False,
        **kwargs
):
    if verbose:
        print('ADD_LIFETIME')
        print('-- lifetime_spectrum:', decay_object.lifetime_spectrum)
        print('-- minimum_initial_lifetime:', minimum_initial_lifetime)
        print('-- maximum_initial_lifetime:', maximum_initial_lifetime)
        print(decay_object)
    if maximum_initial_lifetime is None:
        maximum_initial_lifetime = decay_object.mean_lifetime
    if len(decay_object.lifetime_spectrum) == 0:
        decay_object.lifetime_spectrum = [1.0, decay_object.mean_lifetime]
        return
    if add_new_method == 'split':
        # split lifetime spectrum for next iteration
        lifetime_spectrum = split_largest_lifetime(
            np.copy(decay_object.lifetime_spectrum),
            tau_min=minimum_initial_lifetime,
            tau_max=maximum_initial_lifetime
        )
    else:
        mean_lifetime = decay_object.mean_lifetime
        if add_new_method == 'delta_mean':
            tauf = skf.decay.lifetime.fluorescence_averaged_lifetime(decay_object.lifetime_spectrum)
            new_lifetime = mean_lifetime * np.random.uniform(minimum_initial_lifetime, tauf)
            if verbose:
                print('-- tauf:', tauf)
        elif add_new_method == 'short':
            lts = decay_object.lifetime_spectrum[1::2]
            new_lifetime = min(lts) * 0.9
        else:
            new_lifetime = mean_lifetime * np.random.uniform(0.8, 1.2)
        if verbose:
            print('-- new_lifetime:', new_lifetime)
        new_amplitude = np.mean(decay_object.lifetime_spectrum[::2])
        lifetime_spectrum = list(decay_object.lifetime_spectrum) + \
                            [new_amplitude, new_lifetime]
    decay_object.lifetime_spectrum = lifetime_spectrum


def select_best_n_lifetimes(
        score_range: skf.typing.Tuple[int, int],
        scores: skf.typing.List[float],
        n_lifetimes: skf.typing.List[int],
        maximum_number_of_lifetimes: int,
        prob_threshold: float,
        selection_mode: str = 'upper', # if upper selects from the right (more components) else, selects from left
        plot_probabilities: bool = False,
        plot_to_files: bool = False,
        plot_fn_prefix: bool = False,
        plot_path: pathlib.Path = None,
) -> skf.typing.Tuple[int, int, skf.typing.List[float]]:
    print("SELECT_BEST_N_LIFETIMES")
    print("-- score_range: ", score_range)
    print("-- scores: ", scores)
    print("-- n_lifetimes: ", n_lifetimes)
    print("-- maximum_number_of_lifetimes: ", maximum_number_of_lifetimes)
    print("-- prob_threshold: ", prob_threshold)
    print("-- plot_probabilities: ", plot_probabilities)
    # select best number of lifetimes
    nu_ = (score_range[1] - score_range[0]) / 10
    scores = np.array(scores)
    # Compute probability that next is
    probs = np.zeros_like(scores)
    probs[0] = 1.0
    for idx, i in enumerate(range(1, maximum_number_of_lifetimes)):
        f_value = scores[i - 1] / scores[i]
        probs[idx + 1] = scipy.stats.f.cdf(f_value, nu_, nu_)

    best_idx = max(0, np.where(probs < prob_threshold)[0][0] - 1)

    search = range(1, maximum_number_of_lifetimes)
    if selection_mode == 'upper':
        search = search[::-1]
        for i in search:
            if (probs[i] > prob_threshold) and (scores[i] < scores[i - i]):
                best_idx = i
                break
    else:
        for i in search:
            if (probs[i] < prob_threshold) and (scores[i] < scores[i - i]):
                best_idx = i - 1
                break

    best_n_lifetimes = n_lifetimes[best_idx]
    if plot_probabilities:
        fig, ax = plt.subplots(2, 1, sharex=True, sharey=False)
        ax[0].semilogy(n_lifetimes, scores, 'o', label='scores')
        ax[0].legend()
        ax[1].hlines(prob_threshold, xmin=0, xmax=maximum_number_of_lifetimes)
        ax[1].plot(n_lifetimes, probs, 'o', label='p(p(i+1) > p(i))')
        ax[1].plot(n_lifetimes[best_idx], probs[best_idx], 'ro', label='selected')
        ax[1].legend()
        if plot_to_files:
            fn = plot_fn_prefix + '_model_selection.png'
            plt.savefig(plot_path / fn)
        else:
            plt.show()
    return best_n_lifetimes, best_idx, probs


def reset_parameters_out_of_bound(
        decay_object: skf.decay.Decay,
        original_parameters: skf.typing.Dict,
        maximum_initial_lifetime: float,
        minimum_initial_lifetime: float,
        maximal_relative_difference: float = 0.2
):
    n = decay_object.parameter
    for k in original_parameters:
        p = original_parameters[k]
        if isinstance(k, float):
            if abs((p - n[k]) / p) > maximal_relative_difference:
                n[k] = p
    p = np.copy(decay_object.lifetime_spectrum)
    for i, l in enumerate(p[1::2]):
        if l > maximum_initial_lifetime or l < minimum_initial_lifetime:
            p[i * 2 + 1] = np.random.uniform(minimum_initial_lifetime, maximum_initial_lifetime)
    n['lifetime_spectrum'] = p
    decay_object.set(**n)


def find_optimal_lifetime_spectrum(
        decay_object: skf.decay.Decay,
        optimize_decay_parameter: dict,
        initial_settings: dict,
        maximum_number_of_lifetimes: int = 6,
        prob_threshold: float = 0.68,
        do_reset_parameter_before_fit: bool = True,
        do_fit_without_nusiance_first: bool = True,
        verbose: bool = True,
        plot_probabilities: bool = True,
        plot_weighted_residuals: bool = True,
        plot_to_files: bool = False,
        plot_fn_prefix: str = "",
        plot_path: pathlib.Path = None
):
    """Find an optimal number of fluorescence lifetimes

    This function determines an optimal number of fluorescence lifetimes for an
    experimental fluorescence decay histogram. The returned optimal number of
    fluorescence lifetimes is determined by optimizing an increasing set of
    lifetimes to the decay. An information criterion is used to determine an
    optimal number of lifetimes (change in chi2).

    Parameters
    ----------
    time_axis: numpy.array
        The time axis of the fluorescence decay histogram
    irf_histogram: numpy.array
        A histogram of the instrument response function
    decay_histogram: numpy.array
        The experimental fluorescence decay histogram
    linearization: numpy.array
        An array that is multiplied with the model function to correct for
        differential non-linearity of the dataf
    score_range: int, int
        The range on which the score of the computed fluorescence decays are
        evaluated on
    maximum_number_of_lifetimes: int
        Maximum number of lifetimes
    verbose: bool
        If set to True (the default value is False) the function will print
        the score for each lifetime to the standard output

    Returns
    -------

    """
    scores = list()
    if verbose:
        print("Lifetimes\tScore")
    n_lifetimes = np.arange(1, maximum_number_of_lifetimes + 1, dtype=np.int)
    lifetime_spectra = list()
    if plot_weighted_residuals:
        fig, ax = plt.subplots(len(n_lifetimes), 1, sharex=True, sharey=False)
    x_min, x_max = decay_object.score_range
    original_decay_parameters = decay_object.parameter
    fitted_parameters = list()
    for j, nlt in enumerate(n_lifetimes):
        if do_reset_parameter_before_fit:
            reset_parameters_out_of_bound(
                decay_object=decay_object,
                original_parameters=original_decay_parameters,
                maximum_initial_lifetime=initial_settings['maximum_initial_lifetime'],
                minimum_initial_lifetime=initial_settings['minimum_initial_lifetime'],
                maximal_relative_difference=initial_settings['maximal_relative_difference']
            )
        if do_fit_without_nusiance_first:
            # First fit with fixed nuisance parameters
            fixed = [True, True, True, True]  # offset , irf, irf background, scatter
            fixed += [False] * 2 * nlt
            _ = optimize_decay_lifetimes(
                decay_object=decay_object,
                fixed=fixed,
                **optimize_decay_parameter
            )
        # Second fit with free nuisance parameters
        fixed = [False, False, False, False]  # offset , irf, irf background, scatter
        fixed += [False] * 2 * nlt
        score = optimize_decay_lifetimes(
            decay_object=decay_object,
            fixed=fixed,
            **optimize_decay_parameter
        )
        fitted_parameters.append(decay_object.parameter)
        lifetime_spectra.append(decay_object.lifetime_spectrum)
        if verbose:
            print("%s\t%s" % (nlt, score))
        scores.append(score)
        if plot_weighted_residuals:
            ax[j].hlines(0, xmin=0, xmax=x_max-x_min)
            ax[j].plot(decay_object.weighted_residuals, alpha=0.5)
            ax[j].set_title(nlt)
        add_lifetime(
            decay_object=decay_object,
            maximum_initial_lifetime=initial_settings['maximum_initial_lifetime'],
            minimum_initial_lifetime=initial_settings['minimum_initial_lifetime'],
            add_new_method=initial_settings['add_new_method']
        )
    if plot_weighted_residuals:
        if plot_to_files:
            fn = plot_fn_prefix + '_wres.png'
            plt.savefig(plot_path / fn)
        else:
            plt.show()

    best_n_lifetimes, best_idx, probs = select_best_n_lifetimes(
        score_range=decay_object.score_range,
        scores=scores,
        n_lifetimes=n_lifetimes,
        maximum_number_of_lifetimes=maximum_number_of_lifetimes,
        prob_threshold=prob_threshold,
        plot_probabilities=plot_probabilities,
        plot_path=plot_path,
        plot_fn_prefix=plot_fn_prefix,
        plot_to_files=plot_to_files
    )
    if verbose:
        print('GET_LIFETIME_SPECTRUM')
        print('-- best_n_lifetimes:', best_n_lifetimes)
        print('-- lifetime_spectrum:', decay_object.lifetime_spectrum)
        print('-- n_lifetimes:', n_lifetimes)
        print('-- scores:', scores)
        print(decay_object)
    return {
        'best_number_of_lifetimes': best_n_lifetimes,
        'n_lifetimes': n_lifetimes,
        'scores': scores,
        'lifetime_spectrum': np.copy(decay_object.lifetime_spectrum),
        'prob': probs
    }


def estimate_irf_shift(
        decay_object: skf.decay.Decay,
        irf_fwhm_range: skf.typing.Tuple[int, int],
        irf_time_shift_scan_range: skf.typing.Tuple[float, float] = (-8.0, 8.0),
        irf_time_shift_scan_n_steps: int = 20,
        shift_score_range: int = 0,
        verbose: int = 2,
        make_plot: bool = True
):
    lb_i, ub_i = irf_fwhm_range
    lb_i += shift_score_range
    ub_i += shift_score_range
    if verbose:
        print('ESTIMATE_IRF_SHIFT')
        print('-- Analysis range:', lb_i, ub_i)
        print('-- mean lifetime:', decay_object.get_mean_lifetime(decay_object._data))
    scan = np.linspace(*irf_time_shift_scan_range, num=irf_time_shift_scan_n_steps)
    scores = list()
    if make_plot:
        fig, ax = plt.subplots(2, 1, False, False)
    if verbose > 1:
        print("Shift\tScore")
    for shift in scan:
        decay_object.irf_shift_channels = shift
        score = decay_object.get_score(lb_i, ub_i)
        scores.append(score)
        if make_plot:
            ax[0].plot(decay_object.weighted_residuals, label=shift)
            ax[0].legend()
        if verbose > 1:
            print("%.2f\t%.1f" % (shift, score))
    best_shift = scan[np.argmin(scores)]
    if verbose:
        print("-- shift with minimum score: ", best_shift)
    decay_object.irf_shift_channels = best_shift
    if make_plot:
        ax[1].plot(scan, scores)
        plt.show()
    if make_plot:
        plt.semilogy(decay_object.time_axis, decay_object.irf, label='IRF')
        plt.semilogy(decay_object.time_axis, decay_object.corrected_irf, label='Corrected IRF')
        plt.semilogy(decay_object.time_axis, decay_object.data, label='Data')
        plt.semilogy(decay_object.time_axis[lb_i:ub_i], decay_object.model[lb_i:ub_i], label='Model')
        plt.legend()
        plt.show()


def estimate_background(
        decay,
        irf_fwhm_range: skf.typing.Tuple[int, int],
        average_window: int = 10,
        plot_background_estimate: bool = True,
        verbose: bool = True
):
    lb_i, ub_i = irf_fwhm_range
    n_count = 0
    s_count = 0.0
    for i in range(0, lb_i):
        if np.mean(decay[i: i + average_window]) < decay[i]:
            n_count += 1
            s_count += decay[i]
    irf_counts = s_count / n_count
    if plot_background_estimate:
        plt.semilogy(decay)
        plt.hlines(irf_counts, xmin=0, xmax=len(decay))
        plt.show()
    if verbose:
        print('ESTIMATE_BACKGROUND')
        print('-- estimated counts / channel:', irf_counts)
    return irf_counts


def check_linearization_quality(
        decay_object: skf.decay.Decay,
        do_try_reversed_linearization: bool = True,
        relative_tolerance_score: float = 0.01,
        do_refit: bool = True,
        verbose: bool = False
) -> bool:
    """

    Parameters
    ----------
    decay_object
    do_try_reversed_linearization
    verbose

    Returns
    -------
    True if linearization is good otherwise False
    """
    initial_use = decay_object.use_linearization
    lin_table_backup = np.copy(decay_object.linearization)
    decay_object.use_linearization = False
    if do_refit:
        optimize_decay_lifetimes(decay_object)
    score_without = decay_object.get_score()

    decay_object.use_linearization = True
    if do_refit:
        optimize_decay_lifetimes(decay_object)
    score_with = decay_object.get_score()

    score_with_reversed = 1e12
    if do_try_reversed_linearization:
        decay_object.linearization = lin_table_backup[::-1]
        if do_refit:
            optimize_decay_lifetimes(decay_object)
        score_with_reversed = decay_object.get_score()

    if verbose:
        print('TRY_WITHOUT_LINEARIZATION')
        print('-- score with:', score_with)
        print('-- score without:', score_without)
        if do_try_reversed_linearization:
            print('-- score lin reversed:', score_with_reversed)

    rel_without = (score_with - score_without) / score_with
    rel_reversed = (score_with - score_with_reversed) / score_with

    if rel_without > relative_tolerance_score or rel_reversed > relative_tolerance_score:
        is_good = False
        if rel_without < rel_reversed:
            decay_object.use_linearization = True
            decay_object.linearization = lin_table_backup[::-1]
            if do_refit:
                optimize_decay_lifetimes(decay_object)
            if verbose:
                print('-- Linearization is bad!')
                print('-- Using reversed linearization!')
        else:
            decay_object.use_linearization = False
            decay_object.linearization = np.ones_like(lin_table_backup)
            if do_refit:
                optimize_decay_lifetimes(decay_object)
            if verbose:
                print('-- Linearization is bad!')
                print('-- Using no linearization!')
    else:
        is_good = True
        optimize_decay_lifetimes(decay_object)
        decay_object.linearization = lin_table_backup
        decay_object.use_linearization = True
        if do_refit:
            optimize_decay_lifetimes(decay_object)
        if verbose:
            print('-- Linearization is good.')
    return is_good


def estimate_initial_parameter(
        decay_object: skf.decay.Decay,
        do_decay_background_estimation: bool,
        do_irf_background_estimation: bool,
        do_irf_shift_estimation: bool,
        irf_fwhm_range: skf.typing.Tuple[int, int],
        estimate_lifetime_spectrum: bool = False,
        estimate_background_parameter: skf.typing.Dict = None,
        estimate_irf_shift_parameters: skf.typing.Dict = None,
        verbose: bool = False
):
    if estimate_lifetime_spectrum:
        # do a scan, increment the number of lifetimes
        decay_object.lifetime_spectrum = [1., decay_object.mean_lifetime]
    decay_object.lifetime_spectrum = [1., decay_object.mean_lifetime]
    # Estimate initial backgrounds
    if do_irf_background_estimation:
        assert estimate_background_parameter is not None
        decay_object.irf_background_counts = estimate_background(
            decay=decay_object.irf,
            irf_fwhm_range=irf_fwhm_range,
            **estimate_background_parameter
        )
        if verbose:
            print('-- Estimated irf background:', decay_object.irf_background_counts)
    if do_decay_background_estimation:
        decay_object.constant_offset = estimate_background(
            decay=decay_object.data,
            irf_fwhm_range=irf_fwhm_range,
            **estimate_background_parameter
        )
        if verbose:
            print('-- Estimated decay background:', decay_object.constant_offset)
    # Estimate initial time shift
    if do_irf_shift_estimation:
        assert estimate_irf_shift_parameters is not None
        assert estimate_lifetime_spectrum is True
        estimate_irf_shift(
            decay_object=decay_object,
            irf_fwhm_range=irf_fwhm_range,
            **estimate_irf_shift_parameters
        )


def initialize_score_range(
        decay_object: skf.decay.Decay,
        analysis_range_parameter: skf.typing.Dict
):
    score_range = skf.decay.analysis.get_analysis_range(
        fluorescence_decay=decay_object.data,
        **analysis_range_parameter
    )
    decay_object.score_range = score_range


def initialize_decay_object(
        decay_object: skf.decay.Decay,
        tcspc_data_set: skf.decay.tcspc.TCSPCDataset,
        rebin_parameters: dict,
        decay_parameters: skf.typing.Dict,
        analysis_range_parameter: skf.typing.Dict,
        linearization_parameter: skf.typing.Dict,
        estimate_background_parameter: skf.typing.Dict,
        estimate_irf_shift_parameters: skf.typing.Dict,
        do_rebining: bool = False,
        do_smooth_linearization: bool = True,
        do_initial_irf_background_estimation: bool = True,
        do_initial_irf_shift_estimation: bool = False,
        do_initial_decay_background_estimation: bool = True,
        verbose: bool = False
) -> skf.typing.NoReturn:
    """ Initialize a scikit_fluorescence.decay.Decay object

    The initialization sets

        1.  time axis, data, instrument response function.
        2.  based on data defines the score range (fit range)
        3.  reduces noise in linearization and applies linearization table
        4.  estimates initial parameters for background, lifetime spectrum,
            irf background, irf time-shift

    :param decay_object:
    :param tcspc_data_set:
    :param rebin_parameters:
    :param decay_parameters:
    :param analysis_range_parameter:
    :param linearization_parameter:
    :param estimate_background_parameter:
    :param estimate_irf_shift_parameters:
    :param do_rebining:
    :param do_smooth_linearization:
    :param do_initial_irf_background_estimation:
    :param do_initial_irf_shift_estimation:
    :param do_initial_decay_background_estimation:
    :param verbose:
    :return:
    """
    print("-- do_rebining:", do_rebining)
    print("-- do_smooth_linearization:", do_smooth_linearization)
    print("-- do_initial_irf_background_estimation:", do_initial_irf_background_estimation)
    print("-- do_initial_irf_shift_estimation:", do_initial_irf_shift_estimation)
    print("-- do_initial_decay_background_estimation:", do_initial_decay_background_estimation)

    # decay_object.parameter = decay_parameters
    if do_rebining:
        tcspc_data_set = skf.decay.tcspc.rebin_decay_dataset(
            tcspc_data_set,
            **rebin_parameters
        )
    decay_object.time_axis = tcspc_data_set['time_axis']
    decay_object.data = tcspc_data_set['decay_histogram']
    decay_object.irf = tcspc_data_set['instrument_response_function']

    initialize_score_range(
        decay_object=decay_object,
        analysis_range_parameter=analysis_range_parameter
    )
    if do_smooth_linearization:
        linearization_parameter.update(
            {
                'x_min': decay_object.score_range[0],
                'x_max': decay_object.score_range[1]
            }
        )
        lin_table = skf.decay.tcspc.compute_linearization_table(
            tcspc_data_set['linearization'],
            **linearization_parameter
        )
    else:
        lin_table = tcspc_data_set['linearization']
    decay_object.linearization = lin_table
    # Estimate initial parameters
    _, irf_fwhm_range, _ = skf.math.signal.calculate_fwhm(
        x_values=tcspc_data_set['time_axis'],
        y_values=tcspc_data_set['instrument_response_function']
    )
    estimate_initial_parameter(
        decay_object=decay_object,
        do_decay_background_estimation=do_initial_decay_background_estimation,
        do_irf_shift_estimation=do_initial_irf_shift_estimation,
        do_irf_background_estimation=do_initial_irf_background_estimation,
        irf_fwhm_range=irf_fwhm_range,
        estimate_background_parameter=estimate_background_parameter,
        estimate_irf_shift_parameters=estimate_irf_shift_parameters,
        estimate_lifetime_spectrum=True
    )


def get_decay_analyzed_by_lifetimes(
        tcspc_data_set: skf.decay.tcspc.TCSPCDataset,
        initialize_decay_parameters: skf.typing.Dict,
        lifetime_fit_parameter: skf.typing.Dict,
        linearization_quality_check_parameter: skf.typing.Dict,
        do_shrink_analysis_range: bool = False,
        do_fit_polish: bool = False,
        do_try_linearization: bool = True,
        redo_analysis_if_linearization_bad: bool = True,
        verbose: bool = False,
        plot_resulting_fit: bool = True,
        plot_label: str = "",
        plot_to_files: bool = False,
        plot_fn_prefix: str = "", # the prefix put in front of the plot files
        plot_path: pathlib.Path = None
) -> skf.decay.Decay:
    """Finds an optimal fluorescence lifetime spectrum for the data

    This function determines an optimal fluorescence lifetime spectrum for the
    input data. The function first determines an optimal analysis / fit range.
    Next, for that fit range the number of fluorescence lifetimes are determined.
    The number of fluorescence lifetimes are determined by increasing the number
    an fluorescence lifetimes and fitting the fluorescence lifetime spectrum to
    the data. An information criterion is used to determine an optimal number of
    lifetimes.

    Parameters
    ----------
    tcspc_data_set
    rebin_parameters
    analysis_range_parameter
    linearization_parameter
    estimate_irf_shift_parameters
    decay_parameters
    estimate_background_parameter
    optimize_decay_parameter
    prob_threshold
    do_initial_irf_background_estimation
    do_initial_irf_shift_estimation
    do_initial_decay_background_estimation
    maximum_number_of_lifetimes
    verbose
    plot_probabilities
    plot_weighted_residuals
    plot_resulting_fit

    Returns
    -------
    skf.decay.Decay:
        A skf.decay.Decay object with an optimized number of fluorescence
        lifetimes and parameters

    """
    print("GET_ANALYZED_DECAY")
    print("-- do_shrink_analysis_range:", do_shrink_analysis_range)
    print("-- do_fit_polish:", do_fit_polish)
    print("-- redo_analysis_if_linearization_bad:", redo_analysis_if_linearization_bad)
    decay_object = skf.decay.Decay()
    initialize_decay_object(
        decay_object=decay_object,
        tcspc_data_set=tcspc_data_set,
        **initialize_decay_parameters,
        verbose=verbose
    )

    skf.decay.analysis.find_optimal_lifetime_spectrum(
        decay_object=decay_object,
        **lifetime_fit_parameter,
        plot_path=plot_path,
        plot_fn_prefix=plot_fn_prefix,
        plot_to_files=plot_to_files
    )

    if do_try_linearization:
        inital_lin = decay_object.use_linearization
        # check if lin is good and if reverse of lin is better
        linearlization_good = check_linearization_quality(
            decay_object=decay_object,
            **linearization_quality_check_parameter
        )
        if not linearlization_good and redo_analysis_if_linearization_bad and inital_lin is True:
            # do a scan, increment the number of lifetimes
            if verbose:
                print("-- Linearization bad. Redoing analysis...")
            decay_object.lifetime_spectrum = [1., decay_object.mean_lifetime]
            decay_object.linearization = np.ones_like(decay_object.linearization)
            decay_object.use_linearization = False
            skf.decay.analysis.find_optimal_lifetime_spectrum(
                decay_object=decay_object,
                **lifetime_fit_parameter,
            )
    _, irf_fwhm_range, _ = skf.math.signal.calculate_fwhm(
        x_values=tcspc_data_set['time_axis'],
        y_values=tcspc_data_set['instrument_response_function']
    )
    # optimize the fit range - go closer to the start of the irf and refit
    if do_shrink_analysis_range:
        shrink_analysis_range(
            decay_object=decay_object,
            irf_fwhm_range=irf_fwhm_range
        )
    imin, imax = decay_object.score_range
    if plot_resulting_fit:
        fig, ax = plt.subplots(2, 1, sharex=True, sharey=False)
        ax[0].plot(
            decay_object.time_axis[imin:imax],
            decay_object.weighted_residuals,
            label='Original'
        )
        ax[0].set_title(plot_label)
    if do_fit_polish:
        print("Polishing fit...")
        estimate_initial_parameter(
            decay_object=decay_object,
            do_decay_background_estimation=True,
            do_irf_shift_estimation=False,
            do_irf_background_estimation=True,
            irf_fwhm_range=irf_fwhm_range,
            estimate_background_parameter=initialize_decay_parameters['estimate_background_parameter'],
            estimate_irf_shift_parameters=initialize_decay_parameters['estimate_irf_shift_parameters'],
            estimate_lifetime_spectrum=False,
            verbose=verbose
        )
        optimize_decay_lifetimes(
            decay_object=decay_object,
            fixed=[False, False, False, False]
        )
        if plot_resulting_fit:
            ax[0].plot(
                decay_object.time_axis[imin:imax],
                decay_object.weighted_residuals,
                label='Polished'
            )
    if plot_resulting_fit:
        ax[0].legend()
        ax[1].semilogy(decay_object.time_axis, decay_object.corrected_irf)
        nu = len(decay_object.lifetime_spectrum) + 4
        chi2 = decay_object.get_score(score_type='normal', start=imin, stop=imax) / (imax - imin - nu)
        s = "Chi2=%.3f\n" % chi2
        for i in range(len(decay_object.lifetime_spectrum) // 2):
            s += '%.2f %.2f\n' % (decay_object.lifetime_spectrum[i * 2], decay_object.lifetime_spectrum[i * 2 + 1])
        anchored_text = matplotlib.offsetbox.AnchoredText(s, loc=1)
        ax[1].add_artist(anchored_text)
        ax[1].semilogy(decay_object.time_axis, decay_object.data)
        ax[1].semilogy(decay_object.time_axis[imin:imax], decay_object.model[imin:imax])
        if plot_to_files:
            fn = plot_fn_prefix + "_final_fit.png"
            plt.savefig(plot_path / fn)
        else:
            plt.show()
    return decay_object


