%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt
import scipy as sp
import scipy.stats as stats
import random

random.seed(12345)
np.random.seed(5345436)

def two_samples(difference, N=6500, delta_variance=0.):
    As = np.random.normal(6., size=N)
    Bs = np.random.normal(6. + difference, scale=1+delta_variance, size=N)
    return As, Bs

a = plt.axes()
As, Bs = two_samples(0., N=100)
_=a.hist(As, bins=30, range=(2,10), alpha=0.6)
_=a.hist(Bs, bins=30, range=(2,10), alpha=0.6)
print "Mean for sample A: %.3f and for sample B: %.3f"%(np.mean(As), np.mean(Bs))

def one_sided_ttest(A, B, equal_var=True):
    t,p = stats.ttest_ind(A, B, equal_var=equal_var)
    # the t-test implemented in scipy is two sided, but we are interested
    # in the one sided p-value, hence this if statement and the divide by two.
    if t < 0:
        p /= 2.
    else:
        p = 1- p/2.
    print "P-value: %.5f, the smaller the less likely it is that the means are the same"%(p)
        
one_sided_ttest(As, Bs)

As2, Bs2 = two_samples(0., N=100)
one_sided_ttest(As2, Bs2)

def repeat_experiment(repeats=10000, diff=0.):
    p_values = []
    for i in xrange(repeats):
        A,B = two_samples(diff, N=100)
        t,p = stats.ttest_ind(A, B, equal_var=True)
        if t < 0:
            p /= 2.
        else:
            p = 1 - p/2.
            
        p_values.append(p)
    plt.hist(p_values, range=(0,1.), bins=20)
    plt.axvspan(0., 0.1, facecolor="red", alpha=0.5)
    
repeat_experiment()

repeat_experiment(diff=0.05)

def keep_or_not(improvement, threshold=0.05, N=100, repeats=1000):
    keep = 0
    for i in xrange(repeats):
        A,B = two_samples(improvement, N=N)
        t,p = stats.ttest_ind(A, B, equal_var=True)
        if t < 0:
            p /= 2.
        else:
            p = 1 - p/2.
        
        if p <= threshold:
            keep += 1
            
    return float(keep)/repeats

improvement = 0.05
thresholds = (0.01, 0.05, 0.1, 0.15, 0.2, 0.25)
for thresh in thresholds:
    kept = keep_or_not(improvement, thresh)*100
    plt.plot(thresh, kept, "bo")

plt.ylim((0, 45))
plt.xlim((0, thresholds[-1]*1.1))
plt.grid()
plt.xlabel("p-value threshold")
plt.ylabel("% cases correctly accepted")

improvements = np.linspace(0., 0.4, 9)
for improvement in improvements:
    kept = keep_or_not(improvement)*100
    plt.plot(improvement, kept, "bo")
    
plt.ylim((0, 100))
plt.xlim((0, improvements[-1]*1.1))
plt.grid()
plt.xlabel("Size of the improvement")
plt.ylabel("% cases correctly accepted")
plt.axhline(5)

improvements = (0.005, 0.05, 0.1, 0.3)
markers = ("ro", "gv", "b^", "ms")
for improvement, marker in zip(improvements, markers):
    sample_size = np.linspace(10, 5000, 10)
    kept = [keep_or_not(improvement, N=size, repeats=10000)*100 for size in sample_size]
    plt.plot(sample_size, kept, marker, label="improvement=%g%%"%(improvement*100))
    
plt.legend(loc='best')
plt.ylim((0, 100))
plt.xlim((0, sample_size[-1]*1.1))
plt.grid()
plt.xlabel("Sample size")
plt.ylabel("% cases correctly accepted")