from sympy import *
from sympy.stats import *

h = Symbol("h")
f = Beta(h, 10, 10)
fig = plot(density(f)(h), (h, 0, 1), title="Probability density function of the probability of heads", xlabel="$h$", ylabel="$f(h)$")

c = Symbol("c")
bias = sample(f)
C = Coin(c, bias)
print density(C)
print [sample(C) for i in range(10)]

f_exact = DiracDelta(h-bias)
# Note: DiracDelta is not a Sympy distribution, so we can not plot the density function or sample it.
integrate(f_exact, (h, 0, 1))

f_uniform = Uniform(h, 0, 1)
sample(f_uniform)
# This craps out on me.
# plot(density(f_uniform)(h), (h, 0, 1))

f1 = 1

f2 = h * f1 / integrate(h * f1, (h, 0, 1))
plot(f2, (h, 0, 1), title="Posterior distribution after observing $c=(\mathrm{H})$.", xlabel="$h$", ylabel="$f_2(h)$")
f2

plot(f2.subs(h, 1-h), (h, 0, 1), title="Posterior distribution after observing $c=(\mathrm{T})$.", xlabel="$h$", ylabel="$f_2(h)$")
f2.subs(h,1-h)

f3 = h * f2 / integrate(h * f2, (h, 0, 1))
fig = plot(f3, (h, 0, 1), title="Posterior distribution after observing $c=(\mathrm{HH})$.", xlabel="$h$", ylabel="$f_3(h)$")
f3

f3_ = (1-h) * f2 / integrate((1-h) * f2, (h, 0, 1))
plot(f3_, (h, 0, 1), title="Posterior distribution after observing $c=(\mathrm{HT})$.", xlabel="$h$", ylabel="$f_3(h)$")
f3_

f4 = h * f3 / integrate(h * f3, (h, 0, 1))
fig = plot(f4, (h, 0, 1), title="Posterior distribution after observing $c=(\mathrm{HHH})$.", xlabel="$h$", ylabel="$f_4(h)$")
f4

f4_ = (1-h) * f3 / integrate((1-h) * f3, (h, 0, 1))
fig = plot(f4_, (h, 0, 1), title="Posterior distribution after observing $c=(\mathrm{HHT})$.", xlabel="$h$", ylabel="$f_4(h)$")
f4_

H = Symbol("H", positive=True, integer=True)
T = Symbol("T", positive=True, integer=True)
# Warning, slow!
#expr = integrate(h**H*(1-h)**T, (h, 0, 1))
expr = gamma(H + 1)*hyper((-T, H + 1), (H + 2,), exp_polar(2*I*pi))/gamma(H + 2)
beta = expr.simplify()
beta

f = (1/beta) * h**H * (1-h)**T
fig = plot(f.subs({H:2,T:1}), (h, 0, 1), title="Posterior distribution after observing $c=(\mathrm{HHT})$.", xlabel="$h$", ylabel="$f(h)$")
fig = plot(f.subs({H:10,T:1}), (h, 0, 1), title="Posterior distribution after observing $c=(H=10,T)$.", xlabel="$h$", ylabel="$f(h)$")
fig = plot(f.subs({H:10,T:10}), (h, 0, 1), title="Posterior distribution after observing $c=(H=10,T=10)$.", xlabel="$h$", ylabel="$f(h)$")
fig = plot(f.subs({H:100,T:50}), (h, 0, 1), title="Posterior distribution after observing $c=(H=100,T=50)$.", xlabel="$h$", ylabel="$f(h)$")