%matplotlib inline
import sympy
sympy.init_printing(use_latex = True)

def density(Sxx, Sx, n, mu, tau):
    """Return the density (2\pi\sigma^2)^{-n/2} \exp(-\sum (x_i-\mu)^2 / 2\sigma^2)
    = (2\pi)^{-n/2} \tau^{n/2} \exp(-\tau (S_{xx} - 2\mu S_x + n\mu^2)/2)"""
    return (2 * sympy.pi)**(-n/2) * tau**(n/2) * sympy.exp( - tau * (Sxx - 2*mu*Sx + n*mu**2) / 2 )

n = sympy.symbols("n", integer = True, positive = True)
Sx, mu = sympy.symbols("Sx mu", real = True)
Sxx, tau = sympy.symbols("Sxx tau", real = True, positive = True)

density(Sxx, Sx, n, mu, tau)

def phi(a, b):
    def func(Sxx, Sx, n):
        return b*b*Sxx + 2*a*b*Sx + n*a*a, b*Sx + n*a, n
    return func
    
a, b = sympy.symbols("a b", real = True)
density(*phi(a,b)(Sxx,Sx,n),mu=mu,tau=tau)

def psi(a, b):
    def func(mu, tau):
        return b*mu+a, b**(-2) * tau
    return func

density(Sxx,Sx,n,*psi(a,b)(mu,tau))

rhs = density(*(phi(a,b)(Sxx,Sx,n) + psi(a,b)(mu,tau))) * sympy.Abs(b)**n
sympy.simplify(rhs)

lhs = density(Sxx,Sx,n,mu,tau)
lhs

rhs == lhs

rhs - lhs == 0

sympy.simplify(rhs - lhs)

sympy.simplify(rhs - lhs) == 0

posterior = density(Sxx, Sx, n, mu, tau) * tau**(-1/2)
posterior

int1 = sympy.integrate(posterior, (mu, -sympy.oo, sympy.oo))
int1 = sympy.simplify(int1)
int1

sympy.integrate(int1, (tau, 1, sympy.oo))

alpha = sympy.symbols("alpha", real = True, positive = True)
int2 = (1/sympy.sqrt(n)) * (2*sympy.pi)**((1-n)/2) * tau**((n-2)/2) * sympy.exp(-tau * alpha / 2)
int2

sympy.simplify( sympy.integrate(int2, (tau,0,sympy.oo)) )

def normalise_my_prior(alpha, n):
    return 2 * sympy.sqrt(2) * sympy.pi**((1-n)/2) * sympy.gamma((n+2)/2) / ( alpha**(n/2) * n * sympy.sqrt(n) )

normalise_my_prior(alpha, n)

sympy.simplify( int1 / normalise_my_prior(alpha, n) )

posterior = density(Sxx, Sx, n, mu, tau) * tau**(-1)
sympy.simplify( sympy.integrate(posterior, (mu,-sympy.oo, sympy.oo)))

int1 = n**(-1/2) * (2*sympy.pi)**((1-n)/2) * tau**((n-3)/2) * sympy.exp(-tau * alpha / 2)
sympy.integrate(int1, (tau,0,sympy.oo))

def normalise_other_prior(alpha, n):
    return ( 2 / (alpha * n**(1/2)) ) * (2*sympy.pi)**((1-n)/2) * (alpha/2)**((3-n)/2) * sympy.gamma((n-1)/2)

normalise_other_prior(alpha, n)

sympy.simplify( int1 / normalise_other_prior(alpha, n) )

import numpy as np
%matplotlib inline
import matplotlib.pyplot as plt

def make_example(n):
    x = np.random.normal(size = n)
    Sxx = np.sum(x*x)
    Sx = np.sum(x)

    def posterior(mu, tau):
        return ( (tau / (2*sympy.pi))**(n/2) * sympy.exp( - tau * (Sxx - 2*mu*Sx + n*mu*mu) / 2 ) * tau**(-1/2)
                / normalise_my_prior(Sxx-Sx*Sx/n, n) )

    def posterior1(mu, tau):
        return ( (tau / (2*sympy.pi))**(n/2) * sympy.exp( - tau * (Sxx - 2*mu*Sx + n*mu*mu) / 2 ) * tau**(-1)
                / normalise_other_prior(Sxx-Sx*Sx/n, n) )

    return posterior, posterior1

posterior, posterior1 = make_example(10)
mubins = np.linspace(-1, 1, 50)
taubins = np.linspace(0.1, 2, 50)

# Slow as our functions use sympy primitives not numpy primitives.
Z = np.zeros([len(mubins),len(taubins)])
Z1 = np.zeros([len(mubins),len(taubins)])
for x, mux in enumerate(mubins):
    for y, tauy in enumerate(taubins):
        Z[x,y] = posterior(mux, tauy)
        Z1[x,y] = posterior1(mux, tauy)

plt.contour(mubins, taubins, Z.T)

plt.contour(mubins, taubins, Z1.T)

posterior = (alpha * tau / 2)**(n/2) / (tau * sympy.gamma(n/2)) * sympy.exp(-alpha * tau / 2)
posterior

posterior1 = (alpha * tau / 2)**((n-1)/2) / (tau * sympy.gamma((n-1)/2)) * sympy.exp(-alpha * tau / 2)
posterior1

from scipy.special import gamma

def make_example(n):
    x = np.random.normal(size=n)
    alpha = np.sum(x*x) - np.sum(x)**2/n

    def p1(tau):
        return (alpha * tau / 2)**(n/2) / (tau * gamma(n/2)) * np.exp(-alpha * tau / 2)
    def p2(tau):
        return (alpha * tau / 2)**((n-1)/2) / (tau * gamma((n-1)/2)) * np.exp(-alpha * tau / 2)
    
    return p1, p2

p1, p2 = make_example(10)
x = np.linspace(0.1, 2, 50)
y1 = p1(x)
y2 = p2(x)
plt.plot(x, y1)
plt.plot(x, y2)