import numpy as np
from scipy.special import gammaln

outcomes = np.zeros((10001, 3), int)
outcome_to_index = {'W': 0, 'B': 1, 'D': 2}
for game, line in enumerate(open('data/chess_outcomes.txt')):
    outcome = line.strip()
    index = outcome_to_index[outcome]
    outcomes[game + 1][index] = 1
    
counts_up_to = outcomes.cumsum(axis=0)

def compute_prior(w, b):
    if 0 < w < 1 and \
       0 < b < 1 and \
       w + b < 1:
        return 1. # the factor of 2 is taken care of by the Bayes denominator
    else:
        return 0.

def compute_log_evidence_factor(w, b, W, B, D):
    return W * np.log(w) + B * np.log(b) + D * np.log(1 - w - b)

def compute_log_denominator(W, B, D):
    return gammaln(W + 1) + gammaln(B + 1) + gammaln(D + 1) - gammaln(W + B + D + 3)

def compute_p(w, b, W, B, D):
    prior = compute_prior(w, b)
    if prior == 0.:
        p = 0.
    else:
        log_evidence_factor = compute_log_evidence_factor(w, b, W, B, D)
        log_bayes_denominator = compute_log_denominator(W, B, D)
        p = np.exp(log_evidence_factor - log_bayes_denominator) * prior
    return p

compute_p = np.vectorize(compute_p)

n = 200  # Number of grid points on each axis
w, w_step = np.linspace(0, 1., num=n, retstep=True)
b, b_step = np.linspace(0, 1., num=n, retstep=True)
ww, bb = np.meshgrid(w, b)

def get_p(num_games):
    W, B, D = counts_up_to[num_games]
    p = compute_p(ww, bb, W, B, D)
    return p

ns = range(10) + range(10, 100, 10) + range(100, 1000, 100) + range(1000, 10001, 1000)

%time ps = [get_p(n) for n in ns]

import matplotlib.pyplot as plt
from matplotlib import animation
from JSAnimation import IPython_display

fig, ax = plt.subplots()

im = ax.imshow(get_p(0), extent=[0, 1, 0, 1], origin='lower')
ax.set_xlabel('w')
ax.set_ylabel('b')
fig.colorbar(im)

def animate(i):
    im.set_array(ps[i])
    im.autoscale()
    ax.set_title('Posterior after {0:5d} games'.format(ns[i]))
    return [im]

animation.FuncAnimation(fig, animate, frames=len(ns), interval=400)