"""Test statistic for differences in the mean and/or variance for multivariate data."""

__author__ = ["johannvk"]

import numpy as np
from numba import njit

from skchange.utils.numba.stats import log_det_covariance


@njit(cache=True)
def half_integer_digamma(twice_n: int) -> float:
    """Calculate the digamma function for half integer values, i.e. twice_n/2.

    The digamma function is the logarithmic derivative of the gamma function.
    This function is capable of calculating the
    digamma function for half integer values.

    Source: https://en.wikipedia.org/wiki/Digamma_function

    Parameters
    ----------
    twice_n : int
        Twice the integer value.

    Returns
    -------
    res : float
        Value of the digamma function for the half integer value.
    """
    assert twice_n > 0, "n must be a positive integer."

    if twice_n % 2 == 0:
        # Even integer: twice_n = 2n
        res = -np.euler_gamma
        n = twice_n // 2
        for k in range(0, n - 1):
            res += 1.0 / (k + 1.0)
    else:
        res = -2 * np.log(2) - np.euler_gamma
        # Odd integer: twice_n = 2n + 1
        n = (twice_n - 1) // 2
        for k in range(1, n + 1):
            res += 2.0 / (2.0 * k - 1.0)

    return res


@njit(cache=True)
def likelihood_ratio_expected_value(
    sequence_length: int, cut_point: int, dimension: int
) -> float:
    """Calculate the expected value of twice the negative log likelihood ratio.

    We check that the cut point is within the sequence length, and that both 'k' and 'n'
    are large enough relative to the dimension 'p', to ensure that the expected
    value is finite.
    Should at least have 'p+1' points on each side of a split, for 'p' dimensional data.

    Parameters
    ----------
    sequence_length : int
        Length of the sequence.
    cut_point : int
        Cut point of the sequence.
    dimension : int
        Dimension of the data.

    Returns
    -------
    g_k_n : float
        Expected value of twice the negative log likelihood ratio.

    """
    n, k, p = sequence_length, cut_point, dimension

    assert 0 < k < n, "Cut point 'k' must be within the sequence length 'n'."
    assert p > 0, "Dimension 'p' must be a positive integer."
    assert k > (p + 1), "Cut point 'k' must be larger than the dimension + 1."
    assert n - k > (
        p + 1
    ), "Run length after cut point 'n - k' must be larger than dimension + 1."

    g_k_n = p * (
        np.log(2)
        + (n - 1) * np.log(n - 1)
        - (n - k - 1) * np.log(n - k - 1)
        - (k - 1) * np.log(k - 1)
    )

    for j in range(1, p + 1):
        g_k_n += (
            (n - 1) * half_integer_digamma(n - j)
            - (k - 1) * half_integer_digamma(k - j)
            - (n - k - 1) * half_integer_digamma(n - k - j)
        )

    return g_k_n


@njit(cache=True)
def bartlett_correction(twice_negated_log_lr, sequence_length, cut_point, dimension):
    """Calculate the Bartlett correction for the twice negated log likelihood ratio.

    Parameters
    ----------
    twice_negated_log_lr : float
        Twice the negative log likelihood ratio.
    sequence_length : int
        Length of the sequence.
    cut_point : int
        Cut point of the sequence.
    dimension : int
        Dimension of the data.

    Returns
    -------
    bartlett_corr_log_lr : float
    """
    g_k_n = likelihood_ratio_expected_value(
        sequence_length=sequence_length, cut_point=cut_point, dimension=dimension
    )
    bartlett_corr_log_lr = dimension * (dimension + 3.0) * twice_negated_log_lr / g_k_n

    return bartlett_corr_log_lr


@njit(cache=True)
def init_mean_cov_score(X: np.ndarray) -> np.ndarray:
    """Initialize the precision matrix change point detection."""
    return X


@njit(cache=True)
def _mean_cov_log_det_term(X: np.ndarray, start: int, end: int) -> float:
    """Compute the contribution of a segment to the mean_cov score.

    Parameters
    ----------
    X : np.ndarray
        Data matrix. Rows are observations and columns are variables.
    start : int
        Start index of the interval to test for a change in the covariance matrix.
        (Inclusive)
    end : int
        End index of the interval to test for a change in the covariance matrix.
        (Inclusive)

    Returns
    -------
    contribution : float
        Contribution of the segment to the mean and covariance cost.
    """
    log_det_cov = log_det_covariance(X[start : end + 1, :])
    if np.isnan(log_det_cov):
        raise RuntimeError(
            f"The covariance matrix of X[{start}:{end + 1}] is not positive definite."
            + " Quick and dirty fix: Add a tiny amount of random noise to the data."
        )

    return (end - start + 1) * log_det_cov


@njit(cache=True)
def _mean_cov_score(
    X: np.ndarray,
    start: int,
    end: int,
    split: int,
) -> float:
    """Calculate the score for a change in mean and variance.

    The score is computed as the Bartlett corrected log likelihood
    ratio for a test of change in mean and/or variance at the split
    point under a multivariate Gaussian model.

    Parameters
    ----------
    X : np.ndarray
        Data matrix. Rows are observations and columns are variables.
    start : int
        Start index of the interval to test for a change in the precision matrix.
        (Inclusive)
    end : int
        End index of the interval to test for a change in the precision matrix.
        (Inclusive)
    split : int
        Split index of the interval to test for a change in the precision matrix.
        Include the element at the index 'split' in the first segment.

    Returns
    -------
    score : float
        Score (twice negative log likelihood ratio) for the change in mean and variance.
    """
    full_span_loss = _mean_cov_log_det_term(X, start, end)
    pre_split_loss = _mean_cov_log_det_term(X, start, split)
    post_split_loss = _mean_cov_log_det_term(X, split + 1, end)
    twice_negated_log_lr = full_span_loss - pre_split_loss - post_split_loss

    score = bartlett_correction(
        twice_negated_log_lr,
        sequence_length=end + 1 - start,
        cut_point=split - start,
        dimension=X.shape[1],
    )
    return score


@njit(cache=True)
def mean_cov_score(
    precomputed_params: np.ndarray,
    starts: np.ndarray,
    ends: np.ndarray,
    splits: np.ndarray,
) -> np.ndarray:
    """Calculate likelihood ratio scores for a change in mean and covariance[1]_.

    References
    ----------
    .. [1] Zamba, K. D., & Hawkins, D. M. (2009). A Multivariate Change-Point Model\
    for Change in Mean Vector and/or Covariance Structure. \
    Journal of Quality Technology, 41(3), 285-303.

    Parameters
    ----------
    precomputed_params : Tuple[np.ndarray, np.ndarray, np.ndarray]
        Precomputed parameters from init_covariance_score.
    starts : np.ndarray
        Start indices of the intervals to test for a change in the covariance.
    ends : np.ndarray
        End indices of the intervals to test for a change in the covariance.
    splits : np.ndarray
        Split indices of the intervals to test for a change in the covariance.

    Returns
    -------
    scores : np.ndarray
        Scores for each start, end, split combination.
    """
    X = precomputed_params
    num_splits = len(splits)

    scores = np.zeros(num_splits, dtype=np.float64)
    for split_idx in range(num_splits):
        scores[split_idx] = _mean_cov_score(
            X, starts[split_idx], ends[split_idx], splits[split_idx]
        )

    return scores
