"""
Tools for general statistical hypothesis testing
"""
#***************************************************************************************************
# Copyright 2015, 2019, 2025 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
# in this software.
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License.  You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
#***************************************************************************************************

from typing import (
    Optional,
    Literal,
)
import numpy as _np


def bonferroni_correction(significance: float, numtests: int) -> float:
    """
    Calculates the standard Bonferroni correction.

    This is used for reducing the "local" significance for > 1 statistical
    hypothesis test to guarantee maintaining a "global" significance (i.e., a
    family-wise error rate) of `significance`.

    Parameters
    ----------
    significance : float
        Significance of each individual test.

    numtests : int
        The number of hypothesis tests performed.

    Returns
    -------
    The Boferroni-corrected local significance, given by
    `significance` / `numtests`.
    """
    local_significance = significance / numtests
    return local_significance


def sidak_correction(significance: float, numtests: int) -> float:
    """
    Sidak correction.

    Convert a significance level for an individual test repeated independently, numtest times
    into a significance level for the whole family of tests.

    https://en.wikipedia.org/wiki/sidak_correction

    Parameters
    ----------
    significance : float
        Significance of each individual test.

    numtests : int
        The number of hypothesis tests performed.

    Returns
    -------
    float
    """
    adjusted_significance = 1 - (1 - significance)**(1 / numtests)
    return adjusted_significance


def generalized_bonferroni_correction(significance: float, weights: _np.ndarray, numtests: Optional[int] = None,
                                      nested_method: Literal["bonferroni", "sidak"] = 'bonferroni',
                                      tol: float = 1e-10) -> float:
    """
    Generalized Bonferroni correction.

    Parameters
    ----------
    significance : float
        Significance of each individual test.

    weights : array-like
        An array of non-negative floating-point weights, one per individual test,
        that sum to 1.0.

    numtests : int
        The number of hypothesis tests performed.

    nested_method : {'bonferroni', 'sidak'}
        Which method is used to find the significance of the composite test.

    tol : float, optional
        Tolerance when checking that the weights add to 1.0.

    Returns
    -------
    float
    """
    weights = _np.array(weights)
    assert (_np.allclose(_np.abs(weights), weights)), "Invalid weighting! The weights must be non-negative."
    assert (_np.abs(_np.sum(weights) - 1.) < tol), "Invalid weighting! The weights must add up to 1."

    adjusted_significance = _np.zeros(len(weights), float)
    adjusted_significance = significance * weights

    if numtests is not None:

        assert (len(numtests) == len(weights)), "The number of tests must be specified for each weight!"
        for i in range(0, len(weights)):

            if nested_method == 'bonferroni':
                adjusted_significance[i] = bonferroni_correction(adjusted_significance[i], numtests[i])

            if nested_method == 'sidak':
                adjusted_significance[i] = sidak_correction(adjusted_significance[i], numtests[i])

    return adjusted_significance
