from scipy import stats
import matplotlib.pyplot as plt
import numpy as np

tot = 0
if __name__ == "__main__":

    n = 5000
    trial = np.cumsum([1.0] * n)
    fig = plt.figure(1)

    N = 100
    for i in range(N):
        rateA = np.cumsum(stats.binom.rvs(n = 1, p = .03, size = n))/trial
        rateB = np.cumsum(stats.binom.rvs(n = 1, p = .03, size = n))/trial
        pvals = 2*(1-stats.norm.cdf(np.abs(rateA-rateB), loc = 0, scale = (rateA*(1-rateA)/trial + rateB*(1-rateB)/trial)**.5))[100:]
        plt.subplot(121)
        plt.plot(2*trial, rateA-rateB,alpha=.3)
        if np.min(pvals) < .01:
            plt.subplot(122)
            plt.plot(2*trial[100:], pvals,alpha=.3)
            tot += 1

    plt.subplot(121)
    plt.ylim([-.025,.025])
    plt.ylabel("difference")
    plt.xlabel("trial")
    plt.subplot(122)
    plt.ylim([0,1])
    plt.plot([0,2*n],[.05]*2,'k--')
    plt.text(n/2,.9,str(tot)+"/"+str(N)+" significant @ alpha 0.02 (two-tailed)")
    plt.ylabel("p-values")
    plt.xlabel("trial")

    print str(tot)+"/"+str(N)+" significant @ 0.05 (two-tailed)"
    plt.show()


    fig = plt.figure(2)

    critical_values = stats.norm.ppf(0.95, loc = 0, scale = (2*.05*.95/trial)**.5)[:5000]
    plt.plot(2*trial[:5000], 1-stats.norm.cdf(critical_values, .02, (2*.05*.95/trial[:5000])**.5))
    plt.plot([0,10000],[.9]*2,'k--')
    plt.plot([4000]*2,[0,1],'k--')
    plt.ylabel("power")
    plt.xlabel("n + n")
    plt.show()
