"""
Stores the PDF for Gaussian on uniform background, parametrized in number of decays. This model lacks any information on the efficiency and exposure.
"""
import numba as nb
import numpy as np

import freqfit.models.constants as constants
from freqfit.utils import inspectparameters

nb_kwd = {
    "nopython": True,
    "parallel": False,
    "nogil": True,
    "cache": True,
    "fastmath": True,
    "inline": "always",
}

QBB = constants.QBB
N_A = constants.NA
M_A = constants.MDET

# default analysis window and width
# window
#     uniform background regions to pull from, must be a 2D array of form e.g. `np.array([[0,1],[2,3]])`
#     where edges of window are monotonically increasing (this is not checked), in keV.
#     Default is typical analysis window.
WINDOW = np.array(constants.WINDOW)

WINDOWSIZE = 0.0
for i in range(len(WINDOW)):
    WINDOWSIZE += WINDOW[i][1] - WINDOW[i][0]

SEED = 42  # set the default random seed


@nb.jit(**nb_kwd)
def nb_pdf(
    Es: np.array,
    n_S: float,
    n_B: float,
    delta: float,
    sigma: float,
    check_window: bool = False,
) -> np.array:
    """
    Parameters
    ----------
    Es
        Energies at which this function is evaluated, in keV
    n_S
        The signal rate, in units of counts
    n_B
        The background index rate, in counts
    delta
        Systematic energy offset from QBB, in keV
    sigma
        The energy resolution at QBB, in keV
    check_window
        whether to check if the passed Es fall inside of the window. Default is False and assumes that the passed Es
        all fall inside the window (for speed)

    Notes
    -----
    This function computes the following:
    mu_S = n_S
    mu_B = n_B
    pdf(E) = 1/(mu_S+mu_B) * [mu_S * norm(E_j, QBB - delta, sigma) + mu_B/windowsize]
    """

    # Precompute the signal and background counts
    mu_S = n_S
    mu_B = n_B

    # Precompute the prefactors so that way we save multiplications in the for loop
    B_amp = mu_B / WINDOWSIZE
    S_amp = mu_S / (np.sqrt(2 * np.pi) * sigma)

    # Initialize and execute the for loop
    y = np.empty_like(Es, dtype=np.float64)
    for i in nb.prange(Es.shape[0]):
        y[i] = (1 / (mu_S + mu_B)) * (
            S_amp * np.exp(-((Es[i] - QBB + delta) ** 2) / (2 * sigma**2)) + B_amp
        )

    if check_window:
        for i in nb.prange(Es.shape[0]):
            inwindow = False
            for j in range(len(WINDOW)):
                if WINDOW[j][0] <= Es[i] <= WINDOW[j][1]:
                    inwindow = True
            if not inwindow:
                y[i] = 0.0

    return y


@nb.jit(**nb_kwd)
def nb_density(
    Es: np.array,
    n_S: float,
    n_B: float,
    delta: float,
    sigma: float,
) -> np.array:
    """
    Parameters
    ----------
    Es
        Energies at which this function is evaluated, in keV
    n_S
        The signal rate, in units of counts
    n_B
        The background index rate, in counts
    delta
        Systematic energy offset from QBB, in keV
    sigma
        The energy resolution at QBB, in keV

    Notes
    -----
    This function computes the following, faster than without a numba wrapper:
    mu_S = n_S
    mu_B = n_B
    CDF(E) = mu_S + mu_B
    pdf(E) =[mu_S * norm(E_j, QBB - delta, sigma) + mu_B/windowsize]
    """

    mu_S = n_S
    mu_B = n_B

    if sigma == 0:
        return np.inf, np.full_like(Es, np.inf, dtype=np.float64)

    # Precompute the prefactors so that way we save multiplications in the for loop
    B_amp = mu_B / WINDOWSIZE
    S_amp = mu_S / (np.sqrt(2 * np.pi) * sigma)

    # Initialize and execute the for loop
    y = np.empty_like(Es, dtype=np.float64)
    for i in nb.prange(Es.shape[0]):
        y[i] = S_amp * np.exp(-((Es[i] - QBB + delta) ** 2) / (2 * sigma**2)) + B_amp

    return mu_S + mu_B, y


@nb.jit(**nb_kwd)
def nb_log_density(
    Es: np.array,
    n_S: float,
    n_B: float,
    delta: float,
    sigma: float,
) -> np.array:
    """
    Parameters
    ----------
    Es
        Energies at which this function is evaluated, in keV
    n_S
        The signal rate, in units of counts
    n_B
        The background index rate, in counts
    delta
        Systematic energy offset from QBB, in keV
    sigma
        The energy resolution at QBB, in keV
    Notes
    -----
    This function computes the following, faster than without a numba wrapper:
    mu_S = n_S
    mu_B = n_B
    CDF(E) = mu_S + mu_B
    pdf(E) =[mu_S * norm(E_j, QBB - delta, sigma) + mu_B/windowsize]
    """
    raise NotImplementedError("not yet implemented , sorry!")


@nb.jit(**nb_kwd)
def nb_density_gradient(
    Es: np.array,
    n_S: float,
    n_B: float,
    delta: float,
    sigma: float,
) -> np.array:
    """
    Parameters
    ----------
    Es
        Energies at which this function is evaluated, in keV
    n_S
        The signal rate, in units of counts
    n_B
        The background index rate, in counts
    delta
        Systematic energy offset from QBB, in keV
    sigma
        The energy resolution at QBB, in keV

    Notes
    -----
    This function computes the gradient of the density function and returns a tuple where the first element is the gradient of the CDF, and the second element is the gradient of the PDF.
    The first element has shape (K,) where K is the number of parameters, and the second element has shape (K,N) where N is the length of Es.
    mu_S = n_S
    mu_B = n_B
    pdf(E) = [mu_S * norm(E_j, QBB + delta, sigma) + mu_B/windowsize]
    cdf(E) = mu_S + mu_B
    """
    raise NotImplementedError("not yet implemented, sorry!")


@nb.jit(**nb_kwd)
def nb_logpdf(
    Es: np.array,
    n_S: float,
    n_B: float,
    delta: float,
    sigma: float,
) -> np.array:
    """
    Parameters
    ----------
    Es
        Energies at which this function is evaluated, in keV
    n_S
        The signal rate, in units of counts
    n_B
        The background index rate, in counts
    delta
        Systematic energy offset from QBB, in keV
    sigma
        The energy resolution at QBB, in keV

    Notes
    -----
    This function computes the following:
    mu_S = n_S
    mu_B = n_B
    logpdf(E) = log(1/(mu_S+mu_B) * [mu_S * norm(E_j, QBB - delta, sigma) + mu_B/windowsize])
    """
    raise NotImplementedError("not yet implemented!")


@nb.jit(nopython=True, fastmath=True, cache=True, error_model="numpy")
def nb_rvs(
    n_sig: int,
    n_bkg: int,
    delta: float,
    sigma: float,
    seed: int = SEED,
) -> np.array:
    """
    Parameters
    ----------
    n_sig
        Number of signal events to pull from
    n_bkg
        Number of background events to pull from
    delta
        Systematic energy offset from QBB, in keV
    sigma
        The energy resolution at QBB, in keV
    seed
        specify a seed, otherwise uses default seed

    Notes
    -----
    This function pulls from a Gaussian for signal events and from a uniform distribution for background events
    in the provided windows, which may be discontinuous.
    """

    np.random.seed(seed)

    # Get energy of signal events from a Gaussian distribution
    # preallocate for background draws
    Es = np.append(np.random.normal(QBB - delta, sigma, size=n_sig), np.zeros(n_bkg))

    # Get background events from a uniform distribution
    bkg = np.random.uniform(0, 1, n_bkg)

    breaks = np.zeros(shape=(len(WINDOW), 2))
    for i in range(len(WINDOW)):
        thiswidth = WINDOW[i][1] - WINDOW[i][0]

        if i > 0:
            breaks[i][0] = breaks[i - 1][1]

        if i < len(WINDOW):
            breaks[i][1] = breaks[i][0] + thiswidth / WINDOWSIZE

        for j in range(len(bkg)):
            if breaks[i][0] <= bkg[j] <= breaks[i][1]:
                Es[n_sig + j] = (bkg[j] - breaks[i][0]) * thiswidth / (
                    breaks[i][1] - breaks[i][0]
                ) + WINDOW[i][0]

    return Es


@nb.jit(nopython=True, fastmath=True, cache=True, error_model="numpy")
def nb_extendedrvs(
    n_S: float,
    n_B: float,
    delta: float,
    sigma: float,
    seed: int = SEED,
) -> np.array:
    """
    Parameters
    ----------
    n_S
        expected rate of signal events in events
    n_B
        rate of background events in events
    delta
        Systematic energy offset from QBB, in keV
    sigma
        The energy resolution at QBB, in keV
    seed
        specify a seed, otherwise uses default seed

    Notes
    -----
    This function pulls from a Gaussian for signal events and from a uniform distribution for background events
    in the provided windows, which may be discontinuous.
    """

    np.random.seed(seed)

    n_sig = n_S
    n_bkg = n_B

    # Get energy of signal events from a Gaussian distribution
    # preallocate for background draws
    Es = np.append(np.random.normal(QBB - delta, sigma, size=n_sig), np.zeros(n_bkg))

    # Get background events from a uniform distribution
    bkg = np.random.uniform(0, 1, n_bkg)

    breaks = np.zeros(shape=(len(WINDOW), 2))
    for i in range(len(WINDOW)):
        thiswidth = WINDOW[i][1] - WINDOW[i][0]

        if i > 0:
            breaks[i][0] = breaks[i - 1][1]

        if i < len(WINDOW):
            breaks[i][1] = breaks[i][0] + thiswidth / WINDOWSIZE

        for j in range(len(bkg)):
            if breaks[i][0] <= bkg[j] <= breaks[i][1]:
                Es[n_sig + j] = (bkg[j] - breaks[i][0]) * thiswidth / (
                    breaks[i][1] - breaks[i][0]
                ) + WINDOW[i][0]

    return Es, (n_sig, n_bkg)


class number_of_counts_0vbb_gen:
    def __init__(self):
        self.parameters = inspectparameters(self.density)
        pass

    def pdf(
        self,
        Es: np.array,
        n_S: float,
        n_B: float,
        delta: float,
        sigma: float,
        check_window: bool = False,
    ) -> np.array:
        return nb_pdf(Es, n_S, n_B, delta, sigma, check_window)

    def logpdf(
        self,
        Es: np.array,
        n_S: float,
        n_B: float,
        delta: float,
        sigma: float,
    ) -> np.array:
        return nb_logpdf(Es, n_S, n_B, delta, sigma)

    # for iminuit ExtendedUnbinnedNLL
    def density(
        self,
        Es: np.array,
        n_S: float,
        n_B: float,
        delta: float,
        sigma: float,
    ) -> np.array:
        return nb_density(Es, n_S, n_B, delta, sigma)

    # for iminuit ExtendedUnbinnedNLL
    def density_gradient(
        self,
        Es: np.array,
        n_S: float,
        n_B: float,
        delta: float,
        sigma: float,
    ) -> np.array:
        return nb_density_gradient(Es, n_S, n_B, delta, sigma)

    # for iminuit ExtendedUnbinnedNLL
    def log_density(
        self,
        Es: np.array,
        n_S: float,
        n_B: float,
        delta: float,
        sigma: float,
    ) -> np.array:
        return nb_log_density(Es, n_S, n_B, delta, sigma)

    # should we have an rvs method for drawing a random number of events?
    # `extendedrvs`
    # needs to use same parameters as the rest of the functions...
    def rvs(
        self,
        n_sig: int,
        n_bkg: int,
        delta: float,
        sigma: float,
        seed: int = SEED,
    ) -> np.array:
        return nb_rvs(n_sig, n_bkg, delta, sigma, seed=seed)

    def extendedrvs(
        self,
        n_S: float,
        n_B: float,
        delta: float,
        sigma: float,
        seed: int = SEED,
    ) -> np.array:
        return nb_extendedrvs(n_S, n_B, delta, sigma, seed=seed)

    def plot(
        self,
        Es: np.array,
        n_S: float,
        n_B: float,
        delta: float,
        sigma: float,
    ) -> None:
        y = nb_pdf(Es, n_S, n_B, delta, sigma)

        import matplotlib.pyplot as plt

        plt.step(Es, y)
        plt.show()

    # function call needs to take the same parameters as the other function calls, in the same order repeated twice
    # this is intended only for empty datasets
    # returns `None` if we couldn't combine the datasets (a dataset was not empty)
    def combine(
        self,
        a_Es: np.array,
        a_n_S: float,
        a_n_B: float,
        a_delta: float,
        a_sigma: float,
        b_Es: np.array,
        b_n_S: float,
        b_n_B: float,
        b_delta: float,
        b_sigma: float,
    ) -> list | None:
        # datasets must be empty to be combined
        if len(a_Es) != 0 or len(b_Es) != 0:
            return None

        Es = np.array([])  # both of these datasets are empty
        n_S = 0.0  # this should be overwritten in the fit later
        # n_B = 0.0  # this should be overwritten in the fit later

        # exp = 1 + 1  # total exposure

        # # exposure weighted fixed parameters (important to calculate correctly)
        # sigma = (a_sigma + b_exp * b_sigma) / exp
        # eff = (a_exp * a_eff + b_exp * b_eff) / exp
        # delta = (a_exp * a_delta + b_exp * b_delta) / exp

        # # these are fully correlated in this model so the direct sum is appropriate
        # # (maybe still appropriate even if not fully correlated?)
        # effunc = (a_exp * a_effunc + b_exp * b_effunc) / exp

        # effuncscale = 0.0  # this should be overwritten in the fit later

        # return [Es, S, BI, delta, sigma, eff, effunc, effuncscale, exp]
        return None

    def can_combine(
        self,
        a_Es: np.array,
        a_n_S: float,
        a_n_B: float,
        a_delta: float,
        a_sigma: float,
    ) -> bool:
        """
        This sets an arbitrary rule if this dataset can be combined with other datasets.
        In this case, if the dataset contains no data, then it can be combined, but more complex rules can be imposed.
        """
        if len(a_Es) == 0:
            return True
        else:
            return False


number_of_counts_0vbb = number_of_counts_0vbb_gen()
