#!/usr/bin/env python3
import os
import sys
from math import sqrt, nan
from random import randrange
from scipy import stats
import pandas as pd
import numpy as np

CURDIR = os.path.dirname(os.path.realpath(__file__))
sys.path.insert(0, os.path.join(CURDIR, "helpers"))

from pure_http_client import ClickHouseClient


# unpooled variance z-test for proportions of two samples
def twosample_proportion_ztest(s1, s2, t1, t2, alpha):
    if s1 == 0 or s2 == 0 or s1 > t1 or s2 > t2 or t1 + t2 == 0:
        return nan, nan, nan, nan

    p1 = s1 / t1
    p2 = s2 / t2
    se = sqrt(p1 * (1 - p1) / t1 + p2 * (1 - p2) / t2)
    if se == 0:
        return nan, nan, nan, nan
    z_stat = (p1 - p2) / se

    one_side = 1 - stats.norm.cdf(abs(z_stat))
    p_value = one_side * 2

    z = stats.norm.ppf(1 - 0.5 * alpha)
    ci_lower = (p1 - p2) - z * se
    ci_upper = (p1 - p2) + z * se

    return z_stat, p_value, ci_lower, ci_upper


def test_and_check(name, z_stat, p_value, ci_lower, ci_upper, precision=1e-2):
    client = ClickHouseClient()
    real = client.query_return_df(
        "SELECT roundBankers({}.1, 16) as z_stat, ".format(name)
        + "roundBankers({}.2, 16) as p_value, ".format(name)
        + "roundBankers({}.3, 16) as ci_lower, ".format(name)
        + "roundBankers({}.4, 16) as ci_upper ".format(name)
        + "FORMAT TabSeparatedWithNames;"
    )
    real_z_stat = real["z_stat"][0]
    real_p_value = real["p_value"][0]
    real_ci_lower = real["ci_lower"][0]
    real_ci_upper = real["ci_upper"][0]
    assert (np.isnan(real_z_stat) and np.isnan(z_stat)) or abs(
        real_z_stat - np.float64(z_stat)
    ) < precision, "clickhouse_z_stat {}, py_z_stat {}".format(real_z_stat, z_stat)
    assert (np.isnan(real_p_value) and np.isnan(p_value)) or abs(
        real_p_value - np.float64(p_value)
    ) < precision, "clickhouse_p_value {}, py_p_value {}".format(real_p_value, p_value)
    assert (np.isnan(real_ci_lower) and np.isnan(ci_lower)) or abs(
        real_ci_lower - np.float64(ci_lower)
    ) < precision, "clickhouse_ci_lower {}, py_ci_lower {}".format(
        real_ci_lower, ci_lower
    )
    assert (np.isnan(real_ci_upper) and np.isnan(ci_upper)) or abs(
        real_ci_upper - np.float64(ci_upper)
    ) < precision, "clickhouse_ci_upper {}, py_ci_upper {}".format(
        real_ci_upper, ci_upper
    )


def test_mean_ztest():
    counts = [0, 0]
    nobs = [0, 0]
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        counts[0], counts[1], nobs[0], nobs[1], 0.05
    )
    test_and_check(
        "proportionsZTest(%d, %d, %d, %d, 0.95, 'unpooled')"
        % (counts[0], counts[1], nobs[0], nobs[1]),
        z_stat,
        p_value,
        ci_lower,
        ci_upper,
    )
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        10, 10, 10, 10, 0.05
    )

    counts = [10, 10]
    nobs = [10, 10]
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        counts[0], counts[1], nobs[0], nobs[1], 0.05
    )
    test_and_check(
        "proportionsZTest(%d, %d, %d, %d, 0.95, 'unpooled')"
        % (counts[0], counts[1], nobs[0], nobs[1]),
        z_stat,
        p_value,
        ci_lower,
        ci_upper,
    )
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        10, 10, 10, 10, 0.05
    )

    counts = [16, 16]
    nobs = [16, 18]
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        counts[0], counts[1], nobs[0], nobs[1], 0.05
    )
    test_and_check(
        "proportionsZTest(%d, %d, %d, %d, 0.95, 'unpooled')"
        % (counts[0], counts[1], nobs[0], nobs[1]),
        z_stat,
        p_value,
        ci_lower,
        ci_upper,
    )

    counts = [10, 20]
    nobs = [30, 40]
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        counts[0], counts[1], nobs[0], nobs[1], 0.05
    )
    test_and_check(
        "proportionsZTest(%d, %d, %d, %d, 0.95, 'unpooled')"
        % (counts[0], counts[1], nobs[0], nobs[1]),
        z_stat,
        p_value,
        ci_lower,
        ci_upper,
    )

    counts = [20, 10]
    nobs = [40, 30]
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        counts[0], counts[1], nobs[0], nobs[1], 0.05
    )
    test_and_check(
        "proportionsZTest(%d, %d, %d, %d, 0.95, 'unpooled')"
        % (counts[0], counts[1], nobs[0], nobs[1]),
        z_stat,
        p_value,
        ci_lower,
        ci_upper,
    )

    counts = [randrange(10, 20), randrange(10, 20)]
    nobs = [
        randrange(counts[0] + 1, counts[0] * 2),
        randrange(counts[1], counts[1] * 2),
    ]
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        counts[0], counts[1], nobs[0], nobs[1], 0.05
    )
    test_and_check(
        "proportionsZTest(%d, %d, %d, %d, 0.95, 'unpooled')"
        % (counts[0], counts[1], nobs[0], nobs[1]),
        z_stat,
        p_value,
        ci_lower,
        ci_upper,
    )

    counts = [randrange(1, 100), randrange(1, 200)]
    nobs = [randrange(counts[0], counts[0] * 2), randrange(counts[1], counts[1] * 3)]
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        counts[0], counts[1], nobs[0], nobs[1], 0.05
    )
    test_and_check(
        "proportionsZTest(%d, %d, %d, %d, 0.95, 'unpooled')"
        % (counts[0], counts[1], nobs[0], nobs[1]),
        z_stat,
        p_value,
        ci_lower,
        ci_upper,
    )

    counts = [randrange(1, 200), randrange(1, 100)]
    nobs = [randrange(counts[0], counts[0] * 3), randrange(counts[1], counts[1] * 2)]
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        counts[0], counts[1], nobs[0], nobs[1], 0.05
    )
    test_and_check(
        "proportionsZTest(%d, %d, %d, %d, 0.95, 'unpooled')"
        % (counts[0], counts[1], nobs[0], nobs[1]),
        z_stat,
        p_value,
        ci_lower,
        ci_upper,
    )

    counts = [randrange(1, 1000), randrange(1, 1000)]
    nobs = [randrange(counts[0], counts[0] * 2), randrange(counts[1], counts[1] * 2)]
    z_stat, p_value, ci_lower, ci_upper = twosample_proportion_ztest(
        counts[0], counts[1], nobs[0], nobs[1], 0.05
    )
    test_and_check(
        "proportionsZTest(%d, %d, %d, %d, 0.95, 'unpooled')"
        % (counts[0], counts[1], nobs[0], nobs[1]),
        z_stat,
        p_value,
        ci_lower,
        ci_upper,
    )


if __name__ == "__main__":
    test_mean_ztest()
    print("Ok.")