import numpy as np
from scipy.misc import factorial

N = 45
z = 3
theta = 1/6

def binomial(theta, N, z):
    coef = factorial(N) / factorial(N-z) / factorial(z)
    p = coef * theta**z * (1 - theta)**(N-z)
    return p

tail = np.arange(z+1)
tail

p = binomial(theta, N, tail).sum() * 2 # left and right tail probability
p

"""
Explanation: Chapter 11. Null Hypothesis Significance Testing

Exercise 11.1
Exercise 11.2
Exercise 11.3

Exercise 11.1
Purpose: To compute p values for stopping at fixed N and fixed z.
Part A
fixed N
End of explanation
"""


right_tail = np.arange(z, N)
p_right = z / right_tail * binomial(theta, right_tail, z)
p = (1 - p_right.sum()) * 2
p

"""
Explanation: Part B
fixed z
End of explanation
"""


N = 45
z = 3

"""
Explanation: Exercise 11.2
Purpose: To determine NHST CIs, and notice that they depend on the experimenter’s intention.
End of explanation
"""


left_tail = np.arange(z+1)
theta = np.arange(0.170, 0.190, 0.001)

p = map(lambda t: binomial(t, N, left_tail).sum()*2, theta)
p = list(p)

list(zip(theta, p))

p = np.array(p)
p_idx = np.nonzero(p > 0.05)[0][-1]
theta1 = theta[p_idx]
p[p_idx], theta1

right_tail = np.arange(z, N)
theta = np.arange(0.005, 0.020, 0.001)

p = map(lambda t: binomial(t, N, right_tail).sum()*2, theta)
p = list(p)

p = np.array(p)
p_idx = np.nonzero(p > 0.05)[0][0]
theta2 = theta[p_idx]
p[p_idx], theta2

theta2, theta1

"""
Explanation: Part A
fixed N
End of explanation
"""


theta = np.arange(0.150, 0.160, 0.001)

low_tail = np.arange(z, N)

def p_greated_than(theta):
    p_right = z / low_tail * binomial(theta, low_tail, z)
    p = (1 - p_right.sum()) * 2
    return p

p = map(p_greated_than, theta)
p = list(p)

list(zip(theta, p))

p_idx = np.nonzero(np.array(p) > 0.05)[0][-1]
theta1 = theta[p_idx]
p[p_idx], theta1

theta = np.arange(0.005, 0.020, 0.001)
high_tail = np.arange(z+1)

high_tail = np.arange(z, N+1)

def p_less_than(theta):
    p = z / high_tail * binomial(theta, high_tail, z)
    p = 2 * p.sum()
    return p

p = map(p_less_than, theta)
p = list(p)

p = np.array(p)
p_idx = np.nonzero(p > 0.05)[0][0]
theta2 = theta[p_idx]
p[p_idx], theta2

theta2, theta1

"""
Explanation: Part B
fixed z
End of explanation
"""


N = 45
z = 3
theta = 1/6

Ns = np.arange(40, 51)
p_N = np.ones_like(Ns) / len(Ns)

p_total = 0
for i, n in enumerate(Ns):
    # For the current `n`, determine the max z that is in the low tail:
    z_max = np.arange(0, n+1) / n
    z_max = np.nonzero(z_max <= z/N)[0][-1]
    low_tail = np.arange(0, z_max+1)
    p = 2*binomial(theta, n, low_tail).sum()
    p_total += p_N[i] * p
    print(n, p)

p_total

"""
Explanation: Exercise 11.3
Purpose: To determine the p value when data collection stops at a fixed duration.
End of explanation
"""


