# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/02_hypothesis_testing.ipynb (unless otherwise specified).

__all__ = ['get_mean_test_statistic', 'get_boot_test_statistics', 'get_p_value',
           'plot_hist_with_kde_for_hypothesis_test', 'get_chi2_test_statistic', 'get_spearmanr_test_statistic']

# Cell
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from scipy import stats
import pandas as pd

from thinkstats import basic_stats
from thinkstats import nsfg

# Cell
def get_mean_test_statistic(a:np.ndarray, b: np.ndarray) -> float: return a.mean() - b.mean()

# Cell
def get_boot_test_statistics(a:np.ndarray, b: np.ndarray, n_tests:int=100,
                             test_statistic_fun:callable=get_mean_test_statistic,
                             mode="merge+shuffle"):
    "Computes bootstrapped test statistics using shuffling"
    test_statistic_boot = np.zeros(n_tests)
    if mode == "merge+shuffle":
        samples = np.hstack((a, b))
        n = len(a)
        for i in range(n_tests):
            np.random.shuffle(samples)
            _a, _b = samples[:n], samples[n:]
            test_statistic_boot[i] = test_statistic_fun(_a, _b)
    elif mode == "shuffle":
        for i in range(n_tests):
            _a, _b = np.random.choice(a, size=len(a)), np.random.choice(b, size=len(b))
            test_statistic_boot[i] = test_statistic_fun(_a, _b)
    else:
        raise NotImplementedError("`mode` needs to be 'merge+shuffle' or 'shuffle'.")
    return test_statistic_boot

# Cell
def get_p_value(x:np.ndarray, test_statistic:float):
    "Computes the p-value using `x` as generated from `get_boot_test_statistics` and the original `test_statistic`"
    kde = stats.gaussian_kde(x)
    p_val = kde.integrate_box_1d(-np.inf, test_statistic)
    if p_val > .5:
        p_val = 1. - p_val
    return p_val

# Cell
def plot_hist_with_kde_for_hypothesis_test(test_statistic_boot:np.ndarray, p_val:float,
                                           left_ext:float=.5, right_ext:float=.5):
    ax = basic_stats.plot_hist_with_kde(test_statistic_boot,
                                   plot_params={"xlabel": "test statistic", "ylabel": "frequency",
                                                "title": f"Sampled test statistics: p-value = {p_val:.4f}"},
                                   fig_params={"figsize": (7,7)}, left_ext=left_ext, right_ext=right_ext)
    ax.plot([test_statistic, test_statistic], ax.get_ylim(), marker="o", color="red", label="actual test statistic")
    ax.legend(loc="best")
    return ax

# Cell
def get_chi2_test_statistic(a:np.ndarray, b: np.ndarray) -> float:
    "`b` is the reference"
    freq_e = pd.Series(b).value_counts(normalize=True).sort_index()
    freq_o = pd.Series(a).value_counts(normalize=True).sort_index()
    return ((freq_o - freq_e)**2/freq_e).sum()

# Cell
def get_spearmanr_test_statistic(a:np.ndarray, b:np.ndarray) -> float:
    "Computes the ranked Spearman correlation"
    assert len(a) == len(b)
    return stats.spearmanr(a, b).correlation