"""
Numba version of bias-corrected and accelerated (BCa) bootstrap.
"""
import numpy as np
from numba import njit, boolean
from itertools import product
import warnings
try:
    from scipy.special import ndtri, ndtr
except:
    warnings.warn('scipy not available. Numba BCa bootstrap will not work.', RuntimeWarning)
    
@njit
def resample_paired_nb(X, Y, func, output_len=1, R=int(1e6), seed=0):
    np.random.seed(seed)
    N = X.size    
    data_paired = np.vstack((X, Y)).T
    idxs_resampling = np.random.randint(low=0, high=N, size=R*N)
    data_resampled = data_paired[idxs_resampling].reshape(R, N, 2)    
    stat = func(X, Y)
    
    boot_sample = np.empty((R, output_len))
    for i, r in enumerate(data_resampled):
        x, y = r.T
        boot_sample[i] = func(x, y)
    return boot_sample

@njit
def resample_nb(X, func, output_len=1, R=int(1e6), seed=0):
    """X: array of shape (N_samples, n_vars)."""
    np.random.seed(seed)
    N = X.shape[0]
    idxs_resampling = np.random.randint(low=0, high=N, size=R*N)
    data_resampled = X[idxs_resampling].reshape(R, N, X.shape[1])    
    
    boot_sample = np.empty((R, output_len))
    for i, r in enumerate(data_resampled):
        boot_sample[i] = func(r)
    return boot_sample

@njit
def resample_block_nb(X, Y, func, output_len=1, R=int(1e5), seed=0):
    """
    X, Y:   ragged arrays or tuples. Each element is an array containing the data for a block. 
    func:   numba function f: X,Y  ->  Z,   Z: 1D array of size output_len.
    """
    np.random.seed(seed)
    def stack(arr_list):
        return np.array([a for arr in arr_list for a in arr])
    
    n_x = [len(x) for x in X]
    n_y = [len(y) for y in Y]
    idxs_resampling_x = [np.random.randint(low=0, high=n, size=R*n) for n in n_x]
    idxs_resampling_y = [np.random.randint(low=0, high=n, size=R*n) for n in n_y]
    X_resampled = [x[idxs_resampling].reshape(R, n) for x, n, idxs_resampling in zip(X, n_x, idxs_resampling_x)]
    Y_resampled = [y[idxs_resampling].reshape(R, n) for y, n, idxs_resampling in zip(Y, n_y, idxs_resampling_y)]
    
    boot_sample = np.empty((R, output_len))
    for i in range(R):
        Xi = stack([x[i] for x in X_resampled])
        Yi = stack([y[i] for y in Y_resampled])
        boot_sample[i] = func(Xi, Yi)
    return boot_sample

def resample(X, func, output_len=1, R=int(1e4), seed=0):
    """X: array of shape (N_samples, n_vars)."""
    np.random.seed(seed)
    N = X.shape[0]
    idxs_resampling = np.random.randint(low=0, high=N, size=R*N)
    data_resampled = X[idxs_resampling].reshape(R, N, X.shape[1])
    
    boot_sample = np.empty((R, output_len))
    for i, r in enumerate(data_resampled):
        boot_sample[i] = func(r)
    return boot_sample


def resample_block(X, Y, func, output_len=1, R=int(1e5), seed=0):
    """
    X, Y:   ragged arrays or tuples. Each element is an array containing the data for a block. 
    func:   numba function f: X,Y  ->  Z,   Z: 1D array of size output_len.
    """
    np.random.seed(seed)
    
    n_x = [len(x) for x in X]
    n_y = [len(y) for y in Y]
    idxs_resampling_x = [np.random.randint(low=0, high=n, size=R*n) for n in n_x]
    idxs_resampling_y = [np.random.randint(low=0, high=n, size=R*n) for n in n_y]
    X_resampled = [x[idxs_resampling].reshape(R, n) for x, n, idxs_resampling in zip(X, n_x, idxs_resampling_x)]
    Y_resampled = [y[idxs_resampling].reshape(R, n) for y, n, idxs_resampling in zip(Y, n_y, idxs_resampling_y)]
    
    boot_sample = np.empty((R, output_len))
    for i in range(R):
        Xi = np.hstack([x[i] for x in X_resampled])
        Yi = np.hstack([y[i] for y in Y_resampled])
        boot_sample[i] = func(Xi, Yi)
    return boot_sample

    
@njit
def jackknife_resampling(data):
    """Performs jackknife resampling on numpy arrays.

    Jackknife resampling is a technique to generate 'n' deterministic samples
    of size 'n-1' from a measured sample of size 'n'. The i-th
    sample  is generated by removing the i-th measurement
    of the original sample.
    """
    n = data.shape[0]
    if data.ndim > 1:
        resamples = np.empty((n, n - 1) + data.shape[1:])
        base_mask = np.ones((n), dtype=boolean)
        for i in range(n): # np.delete does not support 'axis' argument in numba.
            mask_i = base_mask.copy()
            mask_i[i] = False
            resamples[i] = data[mask_i]
    else:
        resamples = np.empty((n, n - 1))
        for i in range(n):
            resamples[i] = np.delete(data, i)
    return resamples

@njit
def jackknife_stat_nb(data, statistic):
    resamples = jackknife_resampling(data)
    stats = np.array([statistic(r) for r in resamples])
    return stats

def jackknife_stat_two_samples(data, data2, statistic):
    jk_X = jackknife_resampling(data)
    jk_Y = jackknife_resampling(data2)
    jk_XY = [*product(jk_X, jk_Y)]
    stats = np.array([statistic(*r) for r in jk_XY])
    return stats

def jackknife_stat_(data, statistic):
    resamples = jackknife_resampling(data)
    stats = np.array([statistic(r) for r in resamples])
    return stats

def _percentile_of_score(a, score, axis, account_equal=False):
    """Vectorized, simplified `scipy.stats.percentileofscore`.

    Unlike `stats.percentileofscore`, the percentile returned is a fraction
    in [0, 1].
    """
    B = a.shape[axis]
    if account_equal:
        return ((a < score).sum(axis=axis) + (a <= score).sum(axis=axis)) / (2 * B)
    else:
        return (a < score).sum(axis=axis) / B
    
def CI_bca(data, statistic, data2=None, alternative='two-sided', alpha=0.05, R=int(2e5), account_equal=False, use_numba=True, **kwargs):
    """If data2 is provided, assumes a block resampling and statistic takes two arguments."""
    if alternative == 'two-sided':
        probs = np.array([alpha/2, 1 - alpha/2])
    elif alternative == 'less':
        probs = np.array([0, 1-alpha])
    elif alternative == 'greater':
        probs = np.array([alpha, 1])
    else:
        raise ValueError(f"alternative '{alternative}' not valid. Available: 'two-sided', 'less', 'greater'.")
    
    if data2 is None:
        resample_func = resample_nb if use_numba else resample
        theta_hat_b = resample_func(data[:,None] if data.ndim == 1 else data,
                                    statistic, R=R, **kwargs).squeeze()
    else:
        resample_func = resample_block_nb if use_numba else resample_block
        theta_hat_b = resample_func(data, data2, statistic, R=R, **kwargs).squeeze()
        data = np.hstack(data)
        data2 = np.hstack(data2)
        
    alpha_bca = _bca_interval(data, data2, statistic, probs, theta_hat_b, account_equal, use_numba)[0]
    
    if np.isnan(alpha_bca).all(): 
        warnings.warn('CI shows there is only one value. Check data.', RuntimeWarning)
        sample_stat = statistic(data)
        return np.array([sample_stat, sample_stat])
    elif alternative == 'two-sided':
        return  np.percentile(theta_hat_b, alpha_bca*100, axis=0)
    elif alternative == 'less':
         return np.array([-np.inf, np.percentile(theta_hat_b, alpha_bca[0]*100, axis=0)])
    elif alternative == 'greater':
        return np.array([np.percentile(theta_hat_b, alpha_bca[0]*100, axis=0), np.inf])
    
def _bca_interval(data, data2, statistic, probs, theta_hat_b, account_equal, use_numba):
    """Bias-corrected and accelerated interval."""
    # calculate z0_hat
    if data2 is None:
        theta_hat = statistic(data)
    else:
        theta_hat = statistic(data, data2)
    percentile = _percentile_of_score(theta_hat_b, theta_hat, axis=-1, account_equal=account_equal)
    z0_hat = ndtri(percentile)

    # calculate a_hat
    if data2 is None:
        jackknife_computer = jackknife_stat_nb if use_numba else jackknife_stat
        theta_hat_jk = jackknife_computer(data, statistic)  # jackknife resample
    else:
        theta_hat_jk = jackknife_stat_two_samples(data, data2, statistic)
    n = theta_hat_jk.shape[0]
    theta_hat_jk_dot = theta_hat_jk.mean(axis=0) 

    U = (n - 1) * (theta_hat_jk_dot - theta_hat_jk)
    num = (U**3).sum(axis=0) / n**3
    den = (U**2).sum(axis=0) / n**2
    a_hat = 1/6 * num / (den**(3/2))

    # calculate alpha_1, alpha_2
    def compute_alpha(p):
        z_alpha = ndtri(p)
        num = z0_hat + z_alpha
        return ndtr(z0_hat + num/(1 - a_hat*num))
    alpha_bca = compute_alpha(probs[(probs != 0) & (probs != 1)])
    if (alpha_bca > 1).any() or (alpha_bca < 0).any():
        warnings.warn('percentiles must be in [0, 1]. bca percentiles: {}\nForcing percentiles in [0,1]...'.format(alpha_bca), RuntimeWarning)
        alpha_bca = np.clip(alpha_bca, 0, 1)
    return alpha_bca, a_hat  # return a_hat for testing