import math, random

# Find mu,sigma for Normal Distr that approximate Bin(n,p)
#
#  mu = n*p
#  sigma = sqrt( n*p*(1-p) )

def normal_approximation_to_binomial(n, p):
    """finds mu and sigma corresponding to a Binomial(n, p)"""
    mu = p * n
    sigma = math.sqrt(p * (1 - p) * n)
    return mu, sigma

from scratch.probability import normal_cdf, inverse_normal_cdf

# ####################################################
# P[Z <= p]  ** This is just the CDF

normal_probability_below = normal_cdf

# ####################################################
# P[Z >= p] = 1 - P[Z <= p]

def normal_probability_above(lo, mu=0, sigma=1):
    return 1 - normal_cdf(lo, mu, sigma)

# ####################################################
# P[ lo <= Z <= hi] = P[Z <= hi] - P[Z <= lo]

def normal_probability_between(lo, hi, mu=0, sigma=1):
    return normal_cdf(hi, mu, sigma) - normal_cdf(lo, mu, sigma)

# ####################################################
# P[ Z <= lo || Z >= hi] = 1 - (P[Z <= hi] - P[Z <= lo])

def normal_probability_outside(lo, hi, mu=0, sigma=1):
    return 1 - normal_probability_between(lo, hi, mu, sigma)



# ######################################################
# normal_upper_bound(p) = z where P[ Z <= z ] = p

def normal_upper_bound(probability, mu=0, sigma=1):
    """
    returns the z for which P(Z <= z) = probability
    """
    return inverse_normal_cdf(probability, mu, sigma)

# ######################################################
# normal_upper_bound(p) = z where P[ Z >= z ] = p

def normal_lower_bound(probability, mu=0, sigma=1):
    """
    returns the z for which P(Z >= z) = probability
    """
    return inverse_normal_cdf(1 - probability, mu, sigma)

# ######################################################
# normal_two_sided_bounds(p) = a,b where P[ a <= Z <= b ] = p
#
# a, b are equi-distance from mean (mu)

def normal_two_sided_bounds(probability, mu=0, sigma=1):
    """
    returns the symmetric (about the mean) bounds
    that contain the specified probability
    """

    tail_probability = (1 - probability) / 2
    # upper bound should have tail_probability above it
    upper_bound = normal_lower_bound(tail_probability, mu, sigma)

    # lower bound should have tail_probability below it
    lower_bound = normal_upper_bound(tail_probability, mu, sigma)
    return lower_bound, upper_bound



mu_0, sigma_0 = normal_approximation_to_binomial(1000, 0.5)
print("Mu and sigma to approximate bin(1000,0.5) = ", mu_0, sigma_0)

def two_sided_p_value(x, mu=0, sigma=1):
    if x >= mu:
        # if x is greater than the mean, the tail is what's greater than x
        return 2 * normal_probability_above(x, mu, sigma)
    else:
        # if x is less than the mean, the tail is what's less than x
        return 2 * normal_probability_below(x, mu, sigma)

# Suppose we observed 530 heads out of 1000, the p-value of this event is:

p_value = two_sided_p_value(529.5, mu_0, sigma_0)
print("p-value(530 out of 1000) = ", p_value)

print("\nSimulation result:")

extreme_value_count = 0

# Perform the experiment 100000 times:
for _ in range(100000):
    num_heads = sum(1 if random.random() < 0.5 else 0    # Count # of heads
                    for _ in range(1000))                # in 1000 flips

    if num_heads >= 530 or num_heads <= 470:  # Count how often
        extreme_value_count += 1              # the num_of_heads is 'extreme'

print(extreme_value_count / 100000)   # 0.062

# Note: see 01-two-sided-hypo-testing1.py
#
# lo, hi = normal_two_sided_bounds(0.95, mu_0, sigma_0)   # (469, 531)
#
#     530 falls WITHIN the acceptable range
#
# Similarly, p-value = 0.062 which is MORE likely than the 0.05 threshold
# set. Because p-value = 0.062 > 0.05, we "accept" the H0 hypothesis

#########################################################
# When to reject H0:
#
# lo, hi = normal_two_sided_bounds(0.95, mu_0, sigma_0)   # (469, 531)
#
# Observing 532 heads will cause us to reject H0
#
# The p-value will also tell the story:

p_value2 = two_sided_p_value(531.5, mu_0, sigma_0)  # 0.0463
print("p-value(532 out of 1000) = ", p_value2)

# Because p-value = 0.0463 is LESS likely (less probabe) than the acceptable
# 0.05 threshold, (i.e., the result is MORE extreme), we REJECT hypothesis H0