import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import seaborn as sb

# one variable
N = 20; z = 14
a = 1; b = 1
sigma = 0.2
nIter = 100000
theta = np.random.rand(1)
thetaHist = np.zeros(nIter)
for i_iter in range(nIter):
    proposed = theta + np.random.normal(0,sigma,1)
    while proposed > 1 or proposed < 0:
        proposed = theta + np.random.normal(1)*sigma
    vProp = (proposed**z * (1-proposed)**(N-z) * proposed**(a-1) * (1-proposed)**(b-1) )
    vCurr = (theta**z * (1-theta)**(N-z) * theta**(a-1) * (1-theta)**(b-1) )
    pmove = min(1,vProp/vCurr)
    if np.random.rand(1) < pmove:
        theta = proposed
    thetaHist[i_iter] = theta 
sb.distplot( thetaHist, bins=100 )

# two variable
N1 = 8; z1 = 6
N2 = 7; z2 = 2
a1 = 2; b1 = 2
a2 = 2; b2 = 2

sigma = 0.2
nIter = 1000000
theta1 = np.random.rand(1)
theta2 = np.random.rand(1)
thetaHist = np.zeros([nIter,2])
for i_iter in range(nIter):
    
    prop1 = theta1 + np.random.normal(0,sigma,1)
    while prop1 > 1 or prop1 < 0:
        prop1 = theta1 + np.random.normal(0,sigma,1)
    
    prop2 = theta2 + np.random.normal(0,sigma,1)
    while prop2 > 1 or prop2 < 0:
        prop2 = theta2 + np.random.normal(0,sigma,1)


    vProp = (prop1**z1* (1-prop1)**(N1-z1)* prop2**z2 *(1-prop2)**(N2-z2) 
             *prop1**(a1-1) *(1-prop1)**(b1-1) *prop2**(a2-1) *(1-prop2)**(b2-1))
    vCurr = (theta1**z1*(1-theta1)**(N1-z1)*theta2**z2*(1-theta2)**(N2-z2)
             *theta1**(a1-1)*(1-theta1)**(b1-1)*theta2**(a2-1)*(1-theta2)**(b2-1))
    pmove = min(1,vProp/vCurr)
    if np.random.rand(1) < pmove:
        theta1 = prop1
        theta2 = prop2
    thetaHist[i_iter,...] = [theta1,theta2] 
print(np.mean(thetaHist[...,1]),np.mean(thetaHist[...,0]))
sb.jointplot(x=thetaHist[...,0], y=thetaHist[...,1], kind='kde',space=0, ratio=4)