from __future__ import annotations
from scikit_fluorescence import typing

import numpy as np
import pylab as p

import scikit_fluorescence as skf
import scikit_fluorescence.math


class FluorescenceDecay(typing.TypedDict):
    time_axis: np.ndarray
    decay_histogram: np.ndarray
    counting_noise: np.ndarray


class TCSPCDataset(FluorescenceDecay):
    instrument_response_function: np.ndarray
    linearization: np.ndarray


def rebin_decay_dataset(
        decay_dataset: TCSPCDataset,
        min_dt_irf_fwhm_fraction: float = 0.25,
        coarsening_factor: int = None,
        verbose: bool = True,
        plot_rebined: bool = True
) -> TCSPCDataset:
    """

    Parameters
    ----------
    decay_dataset
    min_dt_irf_fwhm_fraction
    coarsening_factor
    verbose

    Returns
    -------

    """
    time_axis = decay_dataset['time_axis']
    if skf.decay.analysis.is_linear_time_axis(time_axis=time_axis):
        # compute rebin factor
        decay_histogram = decay_dataset['decay_histogram']
        linearization = decay_dataset['linearization']
        instrument_response_function = decay_dataset['instrument_response_function']
        counting_noise = decay_dataset['counting_noise']

        irf_fwhm, _, _ = skf.math.signal.calculate_fwhm(
            x_values=time_axis,
            y_values=instrument_response_function
        )
        if coarsening_factor is None:
            dt = time_axis[1] - time_axis[0]
            f1 = irf_fwhm * min_dt_irf_fwhm_fraction / dt
            coarsening_factor = int(2**int(np.log2(f1)))
        n_data_points = len(decay_histogram)
        if n_data_points % coarsening_factor != 0:
            n_loss = n_data_points % coarsening_factor
            n_data_points_new = n_data_points - n_loss
            Warning(
                "Decay with %s is coarsened by factor %s. Removing last %s channels." %
                (n_data_points, coarsening_factor, n_loss)
            )
            instrument_response_function = instrument_response_function[:n_data_points_new]
            linearization = linearization[:n_data_points_new]
            decay_histogram = decay_histogram[:n_data_points_new]
            time_axis = time_axis[:n_data_points_new]
            counting_noise = counting_noise[:n_data_points_new]
        new_number_of_channels = int(n_data_points / coarsening_factor)
        new_shape = [new_number_of_channels, coarsening_factor]

        rebined_dataset: TCSPCDataset = {
            'instrument_response_function': instrument_response_function.reshape(new_shape).sum(axis=1),
            'linearization': linearization.reshape(new_shape).mean(axis=1),
            'decay_histogram': decay_histogram.reshape(new_shape).sum(axis=1),
            'time_axis': time_axis.reshape(new_shape).mean(axis=1),
            'counting_noise': np.sqrt((counting_noise.reshape(new_shape)**2.0).sum(axis=1))
        }
    else:
        raise ValueError("Rebining TCSPCDataset only works with linear time axis.")
    if verbose:
        print("REBIN_DECAY_DATASET")
        print("-- coarsening_factor:", coarsening_factor)
        print('-- old_shape', decay_histogram.shape[0])
        print('-- new_shape', new_shape[0])
    if plot_rebined:
        fig, ax = p.subplots(3, 1, True, False)
        ax[0].plot(rebined_dataset['time_axis'], rebined_dataset['linearization'])
        ax[0].plot(decay_dataset['time_axis'], decay_dataset['linearization'])

        ax[1].semilogy(rebined_dataset['time_axis'], rebined_dataset['instrument_response_function'])
        ax[1].semilogy(decay_dataset['time_axis'], decay_dataset['instrument_response_function'])

        ax[1].semilogy(rebined_dataset['time_axis'], rebined_dataset['decay_histogram'])
        ax[1].semilogy(decay_dataset['time_axis'], decay_dataset['decay_histogram'])

        ax[2].semilogy(rebined_dataset['time_axis'], rebined_dataset['counting_noise'])
        ax[2].semilogy(decay_dataset['time_axis'], decay_dataset['counting_noise'])
        p.show()

    return rebined_dataset


def compute_linearization_table(
        data: np.ndarray,
        window_length: int,
        window_function_type: str,
        x_min: int,
        x_max: int,
        fill_value: float = 1.0,
        verbose: bool = False,
        plot_linearization: bool = False
) -> np.ndarray:
    """Calculate a table to account for differential non-linearities.

    This function can be used to smooth a a measurement of uncorrelated light.
    The resulting table can be used to account for differential non-linearities
    (DNL) in TCSPC measurements. The function reduces the noise. Thus, the
    resulting table can be multiplied with a model function without increasing
    the noise.

    The function uses a smoothing window function function to reduce the high
    frequency noise.

    Parameters
    ----------
    data : numpy-array
        The experimental data that is used to compute a linearization table.
    window_length : int
        The length of the averaging window
    window_function_type : str
        The type of the window function that is used (either 'flat', 'hanning',
        'hamming', 'bartlett', 'blackman').
    x_min : int
        The smallest data point that is considered. Data points smaller than
        this value will be set to the parameter *fill_value*.
    x_max : int
        The largest data point that is considered. Data points larger than
        this value will be set to the parameter *fill_value*.
    fill_value : float
        Values outside of the the range [xmin, xmax] are set to this value
        (default: 1.0).

    Returns
    -------
    numpy-array
        The returned numpy array can be used as a table to perturb a model
        function by DNLs.

    Examples
    --------
    >>> import numpy as np
    >>> import scikit_fluorescence as skf
    >>> import scikit_fluorescence.decay
    >>> x = np.linspace(0, 40, 128)
    >>> dnl_fraction = 0.01
    >>> counts = 10000
    >>> mean = np.sin(x) * dnl_fraction  * counts + (1 - dnl_fraction) * counts
    >>> np.random.seed(0)
    >>> data = np.random.poisson(mean).astype(np.float64)
    >>> skf.decay.tcspc.compute_linearization_table(data, 12, "hanning", 10, 90)
    array([1.        , 1.        , 1.        , 1.        , 1.        ,
           1.        , 0.99980049, 0.99922759, 0.99823618, 0.99681489,
           0.99518504, 0.99393009, 0.99341097, 0.99356403, 0.99429588,
           0.99518651, 0.99570281, 0.99588077, 0.99593212, 0.99620644,
           0.99725223, 0.99907418, 1.00121338, 1.00320785, 1.00486314,
           1.00613769, 1.00711954, 1.00770543, 1.00739009, 1.00586985,
           1.00286934, 0.99849134, 0.99377306, 0.98981455, 0.98744986,
           0.98683355, 0.9874091 , 0.98868764, 0.99063644, 0.99339915,
           0.99669779, 1.00031741, 1.00394849, 1.00699941, 1.00914148,
           1.01040307, 1.01119998, 1.01183167, 1.01202324, 1.01120659,
           1.00905205, 1.0055979 , 1.0012921 , 0.99701193, 0.99370229,
           0.99214721, 0.99253283, 0.99448162, 0.99720137, 0.99972508,
           1.001355  , 1.00164798, 1.00115969, 1.00083382, 1.00120239,
           1.0025449 , 1.00450071, 1.00633953, 1.00702856, 1.00606332,
           1.00373424, 1.00058981, 0.99734557, 0.99442657, 0.9925278 ,
           0.99219378, 0.99332066, 0.9954457 , 0.99805603, 1.00079518,
           1.00285328, 1.00367601, 1.00341471, 1.00277917, 1.00262715,
           1.00308127, 1.00383407, 1.00447143, 1.00466959, 1.00415533,
           1.00289327, 1.00144407, 1.00036365, 0.99984595, 0.99978234,
           0.99991273, 1.        , 1.        , 1.        , 1.        ,
           1.        , 1.        , 1.        , 1.        , 1.        ,
           1.        , 1.        , 1.        , 1.        , 1.        ,
           1.        , 1.        , 1.        , 1.        , 1.        ,
           1.        , 1.        , 1.        , 1.        , 1.        ,
           1.        , 1.        , 1.        , 1.        , 1.        ,
           1.        , 1.        , 1.        ])

    """
    if verbose:
        print("COMPUTE_LINEARIZATION_TABLE")
        print("-- window_length:", window_length)
        print("-- window_function_type:", window_function_type)
        print("-- x_min:", x_min)
        print("-- x_max:", x_max)
        print("-- fill_value:", fill_value)
    x2 = data.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 = fill_value
    mnx /= mnx.mean()
    yn = mnx.filled()
    lin_table = skf.math.signal.window(
        data=yn,
        window_len=window_length,
        window_function_type=window_function_type
    )
    if plot_linearization:
        p.plot(lin_table)
        p.show()
    return lin_table


def counting_noise(
        decay: np.ndarray,
        treat_zeros: bool = True,
        zero_value: float = 1.0
) -> np.array:
    """Calculated Poisson noise (sqrt of counts) for TCSPC fluorescence decays

    Parameters
    ----------
    decay : numpy-array
        the photon counts
    treat_zeros: bool
        If treat zeros is True (default) a number specified by `zero_value`
        will be assigned for cases the decay is zero.
    zero_value : float
        The number that will be assigned cases where the value of the decay is
        zero.

    Returns
    -------
    numpy-array
        an array containing noise estimates of the decay that can ve used in
        the data analysis.

    """
    w = np.array(decay, dtype=np.float64)
    if treat_zeros:
        w[w <= 0.0] = zero_value
    return np.sqrt(w)


def combined_counting_noise_parallel_perpendicular(
        parallel: np.ndarray,
        perpendicular: np.ndarray,
        g_factor: float,
        treat_zeros: bool = True,
        zero_value: float = 1.0
) -> np.ndarray:
    """Computes the combined counting weights (1/noise) for two counting
    channels that are added with a scale parameter.

    The weight is calculated a combined decay total that was computed using two
    counting channels `parallel` and `perpendicular` that were added using
    a scale parameter `scale_perpendicular`

    .. math::

         f(t)_{combined} = f(t)_{parallel} + 2.0 * g * f(t)_{perpendicular}

    The noise of the combined decay is computed assuming that in the individual
    noise is Poissonian counting noise

    .. math::

         \Delta(f(t)_{parallel}) = \sqrt{f(t)_{parallel}}
         \Delta(f(t)_{parallel}) = \sqrt{f(t)_{parallel}}

    Parameters
    ----------
    parallel : numpy-array
        parallel model_decay
    perpendicular : numpy-array
        perpendicular model_decay
    g_factor : float
        weight of perpendicular model_decay
    zero_value : float
    treat_zeros : True

    Returns
    -------
    numpy-array
        The weights of the channels of a decay combined.

    """
    scale_perpendicular = 2.0 * g_factor
    vp = counting_noise(
        decay=parallel,
        treat_zeros=treat_zeros,
        zero_value=zero_value
    )
    vs = counting_noise(
        decay=perpendicular,
        treat_zeros=treat_zeros,
        zero_value=zero_value
    )
    vt = np.sqrt(vp ** 2 + scale_perpendicular ** 2 * vs ** 2)
    if treat_zeros:
        vt = np.maximum(vt, zero_value)
    return vt


def combine_parallel_perpendicular(
        parallel: np.ndarray,
        perpendicular: np.ndarray,
        g_factor: float
) -> np.ndarray:
    """Combine a parallel and a perpendicular decay to yield a anisotropy
    free total decay.

    The parallel decay defined by the parameter `parallel` and the perpendicular
    decay specified by the parameter `perpendicular` are added weighted by a correction
    factor that accounts for the detection efficiencies of for the two decays to
    yield a total fluorescence decay that is polarization free.

    .. math::

         \Delta(f(t)_{total}) = f(t)_{parallel} + 2 * g * f(t)_{perpendicular}

    The combined decay :math:`f(t)_{total}` is returned.

    Parameters
    ----------
    parallel : numpy-array
        The fluorescence in the parallel detection channel
    perpendicular : numpy-array
        The fluorescence in the perpendicular detection channel
    g_factor : float
        A parameter accounting for the detection efficiency of the parallel
        and the perpendicular detection channel.

    Returns
    -------


    """
    scale_perpendicular = 2.0 * g_factor
    return parallel + perpendicular * scale_perpendicular

