"""A range of functions to compute bootstraps for the mean difference"""

# 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', 'bootstrap_indices',
           'compute_bootstrapped_diff', 'delta2_bootstrap_loop', 'compute_delta2_bootstrapped_diff',
           'compute_meandiff_bias_correction', 'compute_interval_limits', 'calculate_group_var',
           'calculate_bootstraps_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
from numba import njit, prange
from scipy.stats import norm
from numpy import isnan

# %% ../../nbs/API/confint_2group_diff.ipynb 5
@njit(cache=True, parallel=True)
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.
    """

    n = len(data)
    indexes = np.empty((n, n - 1), dtype=np.int64)
    for i in prange(n):
        indexes[i] = np.concatenate((np.arange(i), np.arange(i + 1, n)))
    return indexes


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

    n = len(data)
    indexes = np.empty((n, n), dtype=np.int64)  # Pre-allocate the output array
    for i in prange(n):
        indexes[i, :] = np.arange(n)  # Fill each row with the full index range
    return indexes


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


@njit(cache=True) # parallelization must be turned off for random number generation
def bootstrap_indices(is_paired, x0_len, x1_len, resamples, random_seed):
    np.random.seed(random_seed)
    indices = np.empty((resamples, x0_len if is_paired else x0_len + x1_len), dtype=np.int64)
    
    for i in range(resamples):
        if is_paired:
            indices[i, :x0_len] = np.random.choice(x0_len, x0_len)
        else:  
            indices[i, :x0_len] = np.random.choice(x0_len, x0_len)
            indices[i, x0_len:x0_len+x1_len] = np.random.choice(x1_len, x1_len)
    return indices


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

    x0_len, x1_len = len(x0), len(x1)
    indices = bootstrap_indices(is_paired, x0_len, x1_len, resamples, random_seed)
    out = np.empty(resamples, dtype=np.float64)

    for i in range(resamples):
        if is_paired:
            x0_sample = x0[indices[i, :x0_len]]
            x1_sample = x1[indices[i, :x0_len]]
        else:
            x0_sample = x0[indices[i, :x0_len]]
            x1_sample = x1[indices[i, x0_len:x0_len+x1_len]]

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

    return out


@njit(cache=True)
def delta2_bootstrap_loop(x1, x2, x3, x4, resamples, pooled_sd, rng_seed, is_paired, proportional=False):
    """
    Compute bootstrapped differences for delta-delta, handling both regular and proportional data
    """
    np.random.seed(rng_seed)
    deltadelta = np.empty(resamples)
    out_delta_g = np.empty(resamples)
    
    n1, n2, n3, n4 = len(x1), len(x2), len(x3), len(x4)
    if is_paired and (n1 != n2 or n3 != n4):
        raise ValueError("Each control group must have the same length as its corresponding test group in paired analysis.")

    # Bootstrapping
    for i in range(resamples):
        # Paired or unpaired resampling
        if is_paired:
            indices_1 = np.random.choice(len(x1), len(x1))
            indices_2 = np.random.choice(len(x3), len(x3))
            x1_sample, x2_sample = x1[indices_1], x2[indices_1]
            x3_sample, x4_sample = x3[indices_2], x4[indices_2]
        else:
            indices_1 = np.random.randint(0, len(x1), len(x1))
            indices_2 = np.random.randint(0, len(x2), len(x2))
            indices_3 = np.random.randint(0, len(x3), len(x3))
            indices_4 = np.random.randint(0, len(x4), len(x4))
            x1_sample, x2_sample = x1[indices_1], x2[indices_2]
            x3_sample, x4_sample = x3[indices_3], x4[indices_4]

        # Calculate 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 if proportional else delta_delta/pooled_sd

    return out_delta_g, deltadelta


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,
    random_seed: int = 12345,
    proportional: bool = False
) -> tuple:
    """
    Bootstraps the effect size deltas' g or proportional delta-delta
    """
    x1, x2, x3, x4 = map(np.asarray, [x1, x2, x3, x4])
    
    if proportional:
        # For proportional data, pass 1.0 as dummy pooled_sd (won't be used)
        out_delta_g, deltadelta = delta2_bootstrap_loop(
            x1, x2, x3, x4, resamples, 1.0, random_seed, is_paired, proportional=True
        )
        # For proportional data, delta_g is the empirical delta-delta
        delta_g = ((np.mean(x4) - np.mean(x3)) - (np.mean(x2) - np.mean(x1)))
    else:
        # Calculate pooled sample standard deviation for non-proportional data
        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)
        
        if sd_denominator == 0:
            raise ValueError("Insufficient data to compute pooled standard deviation.")
            
        pooled_sample_sd = np.sqrt(sd_numerator / sd_denominator)
        
        if np.isnan(pooled_sample_sd) or pooled_sample_sd == 0:
            raise ValueError("Pooled sample standard deviation is NaN or zero.")
            
        out_delta_g, deltadelta = delta2_bootstrap_loop(
            x1, x2, x3, x4, resamples, pooled_sample_sd, random_seed, is_paired, proportional=False
        )
        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


@njit(cache=True)
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


@njit(cache=True)
def calculate_group_var(control_var, control_N, test_var, test_N):
    
    pooled_var = ((control_N - 1) * control_var + (test_N - 1) * test_var) / (control_N + test_N - 2) 
    
    return pooled_var

def calculate_bootstraps_var(bootstraps):

    bootstraps_var_list = [np.var(x, ddof=1) for x in bootstraps]
    bootstraps_var_array = np.array(bootstraps_var_list)
    
    return bootstraps_var_array
    


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

    weight = np.true_divide(1, bootstrap_dist_var)
    denom = np.sum(weight)
    num = 0.0
    for i in range(len(weight)):
        num += weight[i] * differences[i]
    return np.true_divide(num, denom)
