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)


lo, hi = normal_two_sided_bounds(0.95, mu_0, sigma_0)	# (469, 531)

print("\nnormal_two_sided_bounds(0.95):")
print( lo, hi )

# What is power ????
# Power = 1 - P[false positive]
#       = 1 - P[accept false H0]

# let’s check what happens if p is really 0.55, so that the coin 
# is slightly biased toward heads.

# actual mu and sigma based on p = 0.55
mu_1, sigma_1 = normal_approximation_to_binomial(1000, 0.55)

# we can calculate the power of the test with:

type_2_probability = normal_probability_between(lo, hi, mu_1, sigma_1)
# a type 2 error means we ACCEPT the null hypothesis (lo < outcome < hi)
# while H0 is FALSE (use mu_1 and sigma_1)

power = 1 - type_2_probability 		# 0.887
print("Power = ", power)

# Note: Power should be >= 0.8