## drawing balls from urns
## Stat 624 homework starter

import numpy as np
import pandas as pd

def new_urn(n_balls=2,proportion_black = .5):
  return pd.DataFrame([[n_balls*proportion_black, 
          n_balls*(1-proportion_black)]],
          columns = ["Black","White"])

def proportion(urn):
  return urn.Black[0]/np.sum(urn,axis=1)[0]

def draw_ball(urn):
  prop_black = proportion(urn)
  drawn_ball = ['White','Black'][((np.random.uniform(0,1,1) < prop_black)*1)[0]]
  urn[drawn_ball] = urn[drawn_ball] + 1
  return urn

# theta = "Proportion of Black Balls"
# Ho: theta = 0.5
# Ha: theta > 0.5

theta = 0.5
n_start = 10
n_draws = 100
precision = 1000

def doit():
  urn = new_urn(n_start,theta)
  for i in range(n_draws):
    urn = draw_ball(urn)
  return proportion(urn)

for i in range(n_draws):
    urn = draw_ball(urn)

draws = np.array([doit() for i in range(precision)])

observed = 0.75
p_value = np.mean(np.array(draws) > observed)
ci = p_value + np.array([-1,1])*1.96*np.sqrt(p_value*(1-p_value)/draws.size)   
ci2 = p_value + np.array([-1,1])*1.96*np.std(draws>observed)/np.sqrt(draws.size)