###############################################################################
# Code needed 
###############################################################################

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





##################################################################
# Stat inference
##################################################################


# NA people see ad A
# nA people click on ad A
#
# Bernoulli trial:
# nA ~= Normal( NA*pA, sqrt( NA*pA*(1-pA) )
# nA/NA ~= Normal( pA, sqrt( pA*(1-pA)/NA )


# NB people see ad B
# nB people click on ad B
#
# Bernoulli trial:
# nB ~= Normal( NB*pB, sqrt( NB*pB*(1-pB) )
# nB/NB ~= Normal( pB, sqrt( pB*(1-pB)/NB )


# Compute Normal distibution parameters for given N and n
#
# Sample usage:
#
#   pA, sigmaA = estimated_parameters(NA, nA)

def estimated_parameters(N, n):
    p = n / N
    sigma = math.sqrt(p * (1 - p) / N)
    return p, sigma

################################################################
# If we assume those two normal distributions are independent 
# (which seems reasonable, since the individual Bernoulli trials 
# ought to be), then their **difference** is also be normal with
#
#  nA   nB
#  -- - --:
#  NA   NB
#
#      mean                = pA - pB                     and 
#      standard deviation  = sqrt(sigmaA^2 + sigmaB^2)
#
# PS: this is just an APPROXIMATION by summing 2 Normal distr
#     It's the sum of 2 normal distr vars 
#
# To transform (nA/NA - nB/NB) into a Normal(0, 1) distribution,
# we use:
#
#         (X - mu)/sigma

# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
# We test: ad A is the same as ad B
#
# I.e., H0 =     "pA == pB"
#          =     pA - pB = 0

# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
# approximately be a standard normal: 

def a_b_test_statistic(N_A, n_A, N_B, n_B):
    p_A, sigma_A = estimated_parameters(N_A, n_A)
    p_B, sigma_B = estimated_parameters(N_B, n_B)

    # We use: mu = 0, and X = pA - pB to transform to Normal(0,1)
    return (p_B - p_A) / math.sqrt(sigma_A ** 2 + sigma_B ** 2)


# Example: A gets 200 clicks out of 1,000 views and 
#          B gets 180 clicks out of 1,000 views, 
#          the statistic equals:

z = a_b_test_statistic(1000, 200, 1000, 180) # -1.14
print(z)

# The probability of seeing such a (large) difference 
# if the means were actually equal would be:

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)


p_z = two_sided_p_value(z)
print(p_z)



z = a_b_test_statistic(1000, 200, 1000, 150) # -2.94
print(z)

p_z = two_sided_p_value(z) # 0.003
print(p_z)