#Stolen shamelessly from http://bit.ly/1xXTU7t.

#The code below can be passed over, as it is currently not important, plus it
# uses advanced topics we have not covered yet. 
%matplotlib inline
from IPython.core.pylabtools import figsize
import numpy as np
from matplotlib import pyplot as plt
figsize(12,10)

import scipy.stats as stats

#beta dist -- P(x, alpha, beta) = const. * x**(alpha - 1) (1 - x)**(beta - 1)
dist = stats.beta
orange_data = np.array([0, 1, 1, 0, 1, 1, 1, 0, 1, 0, 1])
yellow_data = np.array([0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0])
brown_data =  np.array([1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0])
n_trials = [0, 1, 3, 5, 10]
x = np.linspace(0, 1, 100)

for k, N in enumerate(n_trials):
    sx = plt.subplot(3, 2, k + 1)

    plt.setp(sx.get_yticklabels(), visible=False)
    oranges = orange_data[:N].sum()
    yellows = yellow_data[:N].sum()
    browns = brown_data[:N].sum()    
    
    #Hex codes for the colors here: http://www.colourlovers.com/palette/1119759/Reeses_Pieces
    
    #Given a frequency x, what probability to select this many browns?
    y = dist.pdf(x, 1 + browns, 1 + N - browns)
    plt.plot(x, y, color="black")
    plt.fill_between(x, 0, y, color="#47250A")
    
    y = dist.pdf(x, 1 + oranges, 1 + N - oranges)
    plt.plot(x, y, color="black")
    plt.fill_between(x, 0, y, color="#FF6511")
   
    y = dist.pdf(x, 1 + yellows, 1 + N - yellows)
    plt.plot(x, y, color="black")
    plt.fill_between(x, 0, y, color="#FCCF20", alpha=0.8)

    if(k == 0):
        plt.text(0.7, 3.5, "$n_{\\rm trials} =$ %i" % N, fontsize=24)
    else:
        plt.text(0.9, 3.5, "%i" % N, fontsize=24)
    
    plt.vlines(0.3, 0, 4, color="black", linestyles="--", lw=1)    
#    leg = plt.legend()
#    leg.get_frame().set_alpha(0.4)
    plt.autoscale(tight=True)

    plt.tick_params(which='both', labelsize=18)

plt.xlabel("$f_{\\rm color}$, frequency of color", fontsize=24)

plt.suptitle("Bayesian updating of posterior probabilities",
             y=1.02,
             fontsize=24)

plt.tight_layout()
plt.savefig("reeses-pieces_bayesian-inference.png", bbox_inches="tight")