from __future__ import annotations

import numba as nb
import numpy as np
#import fit2x

import scikit_fluorescence as skf
import scikit_fluorescence.math.signal


def compute_linearization_table(
        y: np.array,
        window_length: int,
        window_function_type: str,
        x_min: int,
        x_max: int
) -> np.array:
    """
    This function calculates a liberalization table for differential
    non-linearities given a measurement of uncorrelated light. The
    linearization table is smoothed to reduce the noise contained in the
    liberalization measurements.

    :param y:
    :param window_length:
    :param window_function_type:
    :param x_min:
    :param x_max:
    :return:
    """
    x2 = y.copy()
    x2 /= x2[x_min:x_max].mean()
    mnx = np.ma.array(x2)
    mask2 = np.array([i < x_min or i > x_max for i in range(len(x2))])
    mnx.mask = mask2
    mnx.fill_value = 1.0
    mnx /= mnx.mean()
    yn = mnx.filled()
    return skf.math.signal.window(
        yn,
        window_length,
        window_function_type
    )


def add_pile_up_to_model(
        data: np.ndarray,
        model: np.ndarray,
        repetition_rate: float,
        dead_time: float,
        measurement_time: float,
        pile_up_model: str = "coates"
) -> None:
    """Add pile up to a model function

    This function adds pile up to a model function.

    Attention: This changes the scaling of the model function. Thus, the pile up
    should be added before scaling the model function to the data.

    :param repetition_rate: float
        The repetition-rate in MHz
    :param dead_time: float
        The dead-time of the system in nanoseconds
    :param measurement_time: float
        The measurement time in seconds
    :param data: numpy-array
        The array containing the experimental decay
    :param model: numpy-array
        The array containing the model function
    :param pile_up_model: The used pile up model (for future use) currently only
    the "coates" model is supported.

    References
    ----------

    .. [1]  Coates, P.B., A fluorimetric attachment for an
            atomic-absorption spectrophotometer
            1968 J. Phys. E: Sci. Instrum. 1 878

    .. [2]  Walker, J.G., Iterative correction for pile-up in
            single-photon lifetime measurement
            2002 Optics Comm. 201 271-277

    """
    fit2x.add_pile_up(
        model=model,
        data=data,
        repetition_rate=repetition_rate,
        dead_time=dead_time,
        measurement_time=measurement_time,
        pile_up_model=pile_up_model
    )