# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/API/confint_2group_diff.ipynb.

# %% auto 0
__all__ = ['create_jackknife_indexes', 'create_repeated_indexes', 'compute_meandiff_jackknife', 'compute_bootstrapped_diff',
           'compute_delta2_bootstrapped_diff', 'compute_meandiff_bias_correction', 'compute_interval_limits',
           'calculate_group_var', 'calculate_weighted_delta']

# %% ../../nbs/API/confint_2group_diff.ipynb 4
import numpy as np
from numpy import arange, delete, errstate
from numpy import mean as npmean
from numpy import sum as npsum
from numpy.random import PCG64, RandomState
import pandas as pd
from scipy.stats import norm
from numpy import isnan

# %% ../../nbs/API/confint_2group_diff.ipynb 5
def create_jackknife_indexes(data):
    """
    Given an array-like, creates a jackknife bootstrap.

    For a given set of data Y, the jackknife bootstrap sample J[i]
    is defined as the data set Y with the ith data point deleted.

    Keywords
    --------
    data: array-like

    Returns
    -------
    Generator that yields all jackknife bootstrap samples.
    """

    index_range = arange(0, len(data))
    return (delete(index_range, i) for i in index_range)


def create_repeated_indexes(data):
    """
    Convenience function. Given an array-like with length N,
    returns a generator that yields N indexes [0, 1, ..., N].
    """

    index_range = arange(0, len(data))
    return (index_range for i in index_range)


def _create_two_group_jackknife_indexes(x0, x1, is_paired):
    """Creates the jackknife bootstrap for 2 groups."""

    if is_paired and len(x0) == len(x1):
        out = list(
            zip(
                [j for j in create_jackknife_indexes(x0)],
                [i for i in create_jackknife_indexes(x1)],
            )
        )
    else:
        jackknife_c = list(
            zip(
                [j for j in create_jackknife_indexes(x0)],
                [i for i in create_repeated_indexes(x1)],
            )
        )

        jackknife_t = list(
            zip(
                [i for i in create_repeated_indexes(x0)],
                [j for j in create_jackknife_indexes(x1)],
            )
        )
        out = jackknife_c + jackknife_t
        del jackknife_c
        del jackknife_t

    return out


def compute_meandiff_jackknife(x0, x1, is_paired, effect_size):
    """
    Given two arrays, returns the jackknife for their effect size.
    """
    from . import effsize as __es

    jackknives = _create_two_group_jackknife_indexes(x0, x1, is_paired)

    out = []

    for j in jackknives:
        x0_shuffled = x0[j[0]]
        x1_shuffled = x1[j[1]]

        es = __es.two_group_difference(x0_shuffled, x1_shuffled, is_paired, effect_size)
        out.append(es)

    return out


def _calc_accel(jack_dist):
    """
    Given the Jackknife distribution, calculates the acceleration factor.
    """
    jack_mean = npmean(jack_dist)

    numer = npsum((jack_mean - jack_dist) ** 3)
    denom = 6.0 * (npsum((jack_mean - jack_dist) ** 2) ** 1.5)

    with errstate(invalid="ignore"):
        # does not raise warning if invalid division encountered.
        return numer / denom


def compute_bootstrapped_diff(
    x0, x1, is_paired, effect_size, resamples=5000, random_seed=12345
):
    """Bootstraps the effect_size for 2 groups."""

    from . import effsize as __es

    rng = RandomState(PCG64(random_seed))

    out = np.repeat(np.nan, resamples)
    x0_len = len(x0)
    x1_len = len(x1)

    for i in range(int(resamples)):
        if is_paired:
            if x0_len != x1_len:
                raise ValueError("The two arrays do not have the same length.")
            random_idx = rng.choice(x0_len, x0_len, replace=True)
            x0_sample = x0[random_idx]
            x1_sample = x1[random_idx]
        else:
            x0_sample = rng.choice(x0, x0_len, replace=True)
            x1_sample = rng.choice(x1, x1_len, replace=True)

        out[i] = __es.two_group_difference(x0_sample, x1_sample, is_paired, effect_size)

    return out


def compute_delta2_bootstrapped_diff(
    x1: np.ndarray,  # Control group 1
    x2: np.ndarray,  # Test group 1
    x3: np.ndarray,  # Control group 2
    x4: np.ndarray,  # Test group 2
    is_paired: str = None,
    resamples: int = 5000,  # The number of bootstrap resamples to be taken for the calculation of the confidence interval limits.
    random_seed: int = 12345,  # `random_seed` is used to seed the random number generator during bootstrap resampling. This ensures that the confidence intervals reported are replicable.
) -> (
    tuple
):  # bootstraped result and empirical result of deltas' g, and the bootstraped result of delta-delta
    """
    Bootstraps the effect size deltas' g.

    """

    rng = RandomState(PCG64(random_seed))

    x1, x2, x3, x4 = map(np.asarray, [x1, x2, x3, x4])

    # Calculating pooled sample standard deviation
    stds = [np.std(x) for x in [x1, x2, x3, x4]]
    ns = [len(x) for x in [x1, x2, x3, x4]]

    sd_numerator = sum((n - 1) * s**2 for n, s in zip(ns, stds))
    sd_denominator = sum(n - 1 for n in ns)

    # Avoid division by zero
    if sd_denominator == 0:
        raise ValueError("Insufficient data to compute pooled standard deviation.")

    pooled_sample_sd = np.sqrt(sd_numerator / sd_denominator)

    # Ensure pooled_sample_sd is not NaN or zero (to avoid division by zero later)
    if np.isnan(pooled_sample_sd) or pooled_sample_sd == 0:
        raise ValueError("Pooled sample standard deviation is NaN or zero.")

    out_delta_g = np.empty(resamples)
    deltadelta = np.empty(resamples)

    # Bootstrapping
    for i in range(resamples):
        # Paired or unpaired resampling
        if is_paired:
            if len(x1) != len(x2) or len(x3) != len(x4):
                raise ValueError("Each control group must have the same length as its corresponding test group in paired analysis.")
            indices_1 = rng.choice(len(x1), len(x1), replace=True)
            indices_2 = rng.choice(len(x3), len(x3), replace=True)

            x1_sample, x2_sample = x1[indices_1], x2[indices_1]
            x3_sample, x4_sample = x3[indices_2], x4[indices_2]
        else:
            x1_sample = rng.choice(x1, len(x1), replace=True)
            x2_sample = rng.choice(x2, len(x2), replace=True)
            x3_sample = rng.choice(x3, len(x3), replace=True)
            x4_sample = rng.choice(x4, len(x4), replace=True)

        # Calculating deltas
        delta_1 = np.mean(x2_sample) - np.mean(x1_sample)
        delta_2 = np.mean(x4_sample) - np.mean(x3_sample)
        delta_delta = delta_2 - delta_1

        deltadelta[i] = delta_delta
        out_delta_g[i] = delta_delta / pooled_sample_sd

    # Empirical delta_g calculation
    delta_g = ((np.mean(x4) - np.mean(x3)) - (np.mean(x2) - np.mean(x1))) / pooled_sample_sd

    return out_delta_g, delta_g, deltadelta


def compute_meandiff_bias_correction(
    bootstraps,  # An numerical iterable, comprising bootstrap resamples of the effect size.
    effsize,  # The effect size for the original sample.
):  # The bias correction value for the given bootstraps and effect size.
    """
    Computes the bias correction required for the BCa method
    of confidence interval construction.

    Returns
    -------
    bias: numeric
        The bias correction value for the given bootstraps
        and effect size.

    """

    B = np.array(bootstraps)
    prop_less_than_es = sum(B < effsize) / len(B)

    return norm.ppf(prop_less_than_es)


def _compute_alpha_from_ci(ci):
    if ci < 0 or ci > 100:
        raise ValueError("`ci` must be a number between 0 and 100.")

    return (100.0 - ci) / 100.0


def _compute_quantile(z, bias, acceleration):
    numer = bias + z
    denom = 1 - (acceleration * numer)

    return bias + (numer / denom)


def compute_interval_limits(bias, acceleration, n_boots, ci=95):
    """
    Returns the indexes of the interval limits for a given bootstrap.

    Supply the bias, acceleration factor, and number of bootstraps.
    """

    alpha = _compute_alpha_from_ci(ci)

    alpha_low = alpha / 2
    alpha_high = 1 - (alpha / 2)

    z_low = norm.ppf(alpha_low)
    z_high = norm.ppf(alpha_high)

    kws = {"bias": bias, "acceleration": acceleration}
    low = _compute_quantile(z_low, **kws)
    high = _compute_quantile(z_high, **kws)

    if isnan(low) or isnan(high):
        return low, high

    
    low = int(norm.cdf(low) * n_boots)
    high = int(norm.cdf(high) * n_boots)
    return low, high


def calculate_group_var(control_var, control_N, test_var, test_N):
    return control_var / control_N + test_var / test_N


def calculate_weighted_delta(group_var, differences):
    """
    Compute the weighted deltas.
    """

    weight = 1 / group_var
    denom = np.sum(weight)
    num = np.sum(weight[i] * differences[i] for i in range(0, len(weight)))

    return num / denom
