# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_power.ipynb.

# %% auto 0
__all__ = ['simulate_power_binary', 'sample_size_binary', 'simulate_power_continuous', 'sample_size_continuous']

# %% ../nbs/01_power.ipynb 6
import numpy as np
import pandas as pd

from scipy import interpolate

import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.figure_factory as ff

import statsmodels.stats.api as sms
from statsmodels.stats.proportion import proportions_ztest
from scipy.stats import chi2_contingency, ttest_ind, fisher_exact, binom, norm


from ab_test_toolkit.generator import (
    generate_binary_data,
    data_to_contingency,
    generate_continuous_data,
    generate_contingency
)

from bayesian_testing.experiments import BinaryDataTest, NormalDataTest

# %% ../nbs/01_power.ipynb 7
def _compute_sample_size_from_powerline(sizes, powerline, power):
    """
    Helper function to interpolate power line and return sample size
    """
    yreduced = np.array(powerline) - power
    freduced = interpolate.UnivariateSpline(sizes, yreduced, s=0)
    freduced_root = freduced.roots()
    if len(freduced_root > 0):
        sample_size = freduced_root[0]
    else:
        sample_size = np.nan
    return sample_size

# %% ../nbs/01_power.ipynb 8
def simulate_power_binary(
    sizes=np.arange(10000, 100001, 10000),
    cr0=0.010,
    cr1=0.012,
    realizations=1000,
    alpha=0.05,
    one_sided=True,
    power=0.8,
    fast=True
):
    """
    Simulates the power using two test: Bayesian testing (Probability B>A) and fisher test (p-value).
    Return total sample size (e.g. we need to devide by 2 for each variant sample size given 50/50 split)
    Fast = True: Contigency table drawn from binom distribution
    Fast = False: Contigency table built form data generated at individual level
    """
    if one_sided == True:
        alternative = "greater"
        alternative_z = "larger"
        print("One sided.")
    else:
        alternative = "two-sided"
        print("Two sided tests, except for bayesian which is one sided.")
    approaches = ["fisher", "z", "bayes"]
    pvs = dict()
    pvs_fisher = dict()
    pvs_z = dict()
    # run simulations for each sample size
    for idx, size in enumerate(sizes):
        pvs[idx] = []
        pvs_fisher[idx] = []
        pvs_z[idx] = []
        for realization in range(0, realizations):
            if fast==False:
                df_binary = generate_binary_data(N=size, cr0=cr0, cr1=cr1)
                df_contingency = data_to_contingency(df_binary)
            elif fast==True:
                df_contingency = generate_contingency(N=size, cr0=cr0, cr1=cr1,exact=False)
            else:
                raise Exception("Invalid input for fast parameter.")
            test = BinaryDataTest()
            for group in [0, 1]:
                nUser = df_contingency.loc[group].users
                nConverted = df_contingency.loc[group].converted
                test.add_variant_data_agg(
                    str(group), totals=nUser, positives=nConverted
                )
            results = test.evaluate(sim_count=5000)
            pvs[idx].append(results[1]["prob_being_best"])
            _, p_fischer = fisher_exact(
                np.array(df_contingency[["not_converted", "converted"]]),
                alternative=alternative,
            )
            pvs_fisher[idx].append(p_fischer)
            _, p_z = proportions_ztest(
                [
                    df_contingency.loc[1].converted,
                    df_contingency.loc[0].converted,
                ],
                nobs=size,
                alternative=alternative_z,
            )
            pvs_z[idx].append(p_z)
        print(f"{idx} done")
    pvalues_results = {
        "fisher": pvs_fisher,
        "bayes": pvs,
        "z": pvs_z,
    }
    # compute power lines
    power_lines = dict()
    power_lines["bayes"] = [
        (np.array(pvalues_results["bayes"][i]) > 1.0 - alpha).sum()
        / realizations
        for i in range(0, len(sizes))
    ]
    for approach in [x for x in approaches if x != "bayes"]:
        power_lines[approach] = [
            (np.array(pvalues_results[approach][i]) < alpha).sum()
            / realizations
            for i in range(0, len(sizes))
        ]
    # get sample size for 0.8 power:
    sample_sizes = dict()
    for approach in approaches:
        sample_sizes[approach] = _compute_sample_size_from_powerline(
            sizes, power_lines[approach], power
        )
    # prepare output
    out = dict()
    out["sizes"] = sizes
    out["approaches"] = approaches
    for approach in approaches:
        out[f"power_{approach}"] = power_lines[approach]
        out[f"sample_size_{approach}"] = sample_sizes[approach]
    return out

# %% ../nbs/01_power.ipynb 9
def sample_size_binary(
    cr0=0.010, cr1=0.012, alpha=0.05, power=0.8, one_sided=True
):
    """
    Calculates the sample size for chi2 using sns package
    """
    alternative = "larger" if one_sided == True else "two-sided"
    effect_size = sms.proportion_effectsize(cr1, cr0)

    required_n = sms.NormalIndPower().solve_power(
        effect_size, power=power, alpha=alpha, ratio=1, alternative=alternative
    )

    return required_n

# %% ../nbs/01_power.ipynb 11
def simulate_power_continuous(
    sizes=np.arange(1000, 20001, 2000),
    effect=0.05,
    realizations=1000,
    alpha=0.05,
    one_sided=True,
    power=0.8,
):
    """
    Simulates the power using t-test and bayesian testing (only one-sided)
    Return total sample size (e.g. we need to devide by 2 for each variant sample size given 50/50 split)
    """
    if one_sided == True:
        alternative = "greater"
        print("One sided.")
    else:
        alternative = "two-sided"
        print("Two sided tests, except for bayesian which is one sided.")
    #     alternative_z = "larger" if one_sided == True else "two-sided"
    approaches = ["t", "bayes"]
    pvs_bayes = dict()
    pvs_t = dict()
    # run simulations for each sample size
    for idx, size in enumerate(sizes):
        pvs_bayes[idx] = []
        pvs_t[idx] = []
        for realization in range(0, realizations):
            df = generate_continuous_data(N=size, effect=effect)
            a = df[df["group"] == 0]["target"].values
            b = df[df["group"] == 1]["target"].values
            p_t = ttest_ind(b, a, alternative=alternative).pvalue
            pvs_t[idx].append(p_t)
            test = NormalDataTest()
            test.add_variant_data("A", a)
            test.add_variant_data("B", b)
            results_bayes = test.evaluate(sim_count=5000)
            p_b = results_bayes[1]["prob_being_best"]
            pvs_bayes[idx].append(p_b)
        print(f"{idx} done")
    pvalues_results = {
        "bayes": pvs_bayes,
        "t": pvs_t,
    }
    # compute power lines
    power_lines = dict()
    power_lines["bayes"] = [
        (np.array(pvalues_results["bayes"][i]) > 1.0 - alpha).sum()
        / realizations
        for i in range(0, len(sizes))
    ]
    power_lines["t"] = [
        (np.array(pvalues_results["t"][i]) < alpha).sum() / realizations
        for i in range(0, len(sizes))
    ]
    # get sample size for 0.8 power:
    sample_sizes = dict()
    for approach in approaches:
        sample_sizes[approach] = _compute_sample_size_from_powerline(
            sizes, power_lines[approach], power
        )
    # prepare output
    out = dict()
    out["sizes"] = sizes
    out["approaches"] = approaches
    for approach in approaches:
        out[f"power_{approach}"] = power_lines[approach]
        out[f"sample_size_{approach}"] = sample_sizes[approach]
    return out

# %% ../nbs/01_power.ipynb 14
def sample_size_continuous(
    mu1=5.00, mu2=5.05, sigma=1, alpha=0.05, power=0.8, one_sided=True
):
    """
    forumla: https://towardsdatascience.com/required-sample-size-for-a-b-testing-6f6608dd330a
    web calulator: https://www.stat.ubc.ca/~rollin/stats/ssize/n2.html
    returns sample size per variant
    two-sided test
    """
    assert mu2 >= mu1, "B needs to be greater or equal to A"

    if one_sided == True:
        z1 = norm.ppf(1.0 - alpha)
    else:
        z1 = norm.ppf(1.0 - alpha / 2.0)

    z2 = norm.ppf(power)
    return (2 * sigma**2) * ((z1 + z2) ** 2) / ((mu2 - mu1) ** 2)
