# -*- coding: utf-8 -*-
import numpy as np
import pymc #2.3
import pylab
from scipy.special import binom

#Jaynes' PT:TLoS example 4.1

def db(x): # Decibel transform
    return 10.0*np.log10(x)

def evidence(x): # Equation 4.8
    return db(1E-100 + x/(1.0 - x + 1E-100))

def b(r, n, f): # Equation 3.86
    return binom(n, r)*np.power(f, r)*np.power(1.0 - f, n - r)

# Boxes prior - equation 4.31
priorA = 1.0/11*(1 - 1e-6)
priorB = 10.0/11*(1 - 1e-6)
priorC = 1.0 - priorA - priorB
sum_priors = priorA + priorB + priorC
print "Priors (A, B, C, Sum): ", priorA, priorB, priorC, sum_priors

# Boxes defect rate
boxA = 1.0/3.0
boxB = 1.0/6.0
boxC = 99.0/100.0

def real_posterior(m): # Equations 4.33 and 4.39
    plA = priorA*b(m, m, boxA)
    plB = priorB*b(m, m, boxB)
    plC = priorC*b(m, m, boxC)

    eA = evidence(priorA) + db(b(m, m, boxA)*(priorC + priorB)/(plC + plB))
    eB = evidence(priorB) + db(b(m, m, boxB)*(priorA + priorC)/(plA + plC))
    eC = evidence(priorC) + db(b(m, m, boxC)*(priorA + priorB)/(plA + plB))

    return eA, eB, eC

# PyMC binomial doesn't work properly for N = 0
# so the evidence vectors must be initialized here
eAplot = [evidence(priorA)]
eBplot = [evidence(priorB)]
eCplot = [evidence(priorC)]

realAplot = [evidence(priorA)]
realBplot = [evidence(priorB)]
realCplot = [evidence(priorC)]

minN = 1
assert(minN > 0)
maxN = 20
assert(maxN >= minN)

for N in range(minN, maxN + 1):
    # Observtion
    NBad = N # Number of bad ones
    assert(NBad <= N) # For when it's not always bad
    print "Number of tests:", N

    # Categorical choice (chosen box must be either A, B or C)
    box = pymc.Categorical('box', p = [priorA, priorB, priorC])

    # Each box has a probability of having a defective piece
    @pymc.deterministic
    def box_func(choice = box):
        if choice == 0: #A
            return boxA
        elif choice == 1: #B
            return boxB
        elif choice == 2: #C
            return boxC
        assert(False)

    # Total defective pieces found
    bad = pymc.Binomial('bad', n = N, p = box_func, observed = True, value = NBad)

    # MCMC sampling
    model = pymc.MCMC([box, box_func, bad])
    model.sample(iter=10000, burn=1000, thin=2)

    # Trace plot, autocorrelation and histogram
    if maxN == minN: # Debug only
        pymc.Matplot.plot(model)
    # Useful statistics
    print
    print "MCMC median, mean = ", np.median(model.trace('box')[:]), np.mean(model.trace('box')[:])

    hist = list(model.trace('box')[:])
    total = len(hist)
    # Posterior probabilities for each hypothesis (box)
    postA = float(hist.count(0))/total
    postB = float(hist.count(1))/total
    postC = float(hist.count(2))/total
    print "Posterior A, B, C = ", postA, postB, postC

    eA = evidence(postA)
    eB = evidence(postB)
    eC = evidence(postC)

    eAplot.append(eA)
    eBplot.append(eB)
    eCplot.append(eC)

    realA, realB, realC = real_posterior(N)

    realAplot.append(realA)
    realBplot.append(realB)
    realCplot.append(realC)

    print "Evidence (dB) A, B, C = ", eA, eB, eC
    print "Actual:", realA, realB, realC

fig = pylab.figure()
pylab.title('MCMC approximation')
pylab.plot(eAplot)
pylab.plot(eBplot)
pylab.plot(eCplot)
ax = fig.gca()
ax.set_xticks(np.arange(minN - 1, maxN + 1, 1))
ax.set_yticks(np.arange(-30, 20, 3))
pylab.ylabel('Evidence (dB)')
pylab.ylim([-30, 20])
pylab.grid()
pylab.legend(['A', 'B', 'C'])
pylab.show()

fig = pylab.figure()
pylab.title('Jaynes\' Analytical Answer')
pylab.plot(realAplot)
pylab.plot(realBplot)
pylab.plot(realCplot)
ax = fig.gca()
ax.set_xticks(np.arange(minN - 1, maxN + 1, 1))
ax.set_yticks(np.arange(-30, 20, 3))
pylab.ylabel('Evidence (dB)')
pylab.ylim([-30, 20])
pylab.grid()
pylab.legend(['A', 'B', 'C'])
pylab.show()