import matplotlib.pyplot as plt 
import numpy as np
import scipy.stats as st
import statsmodels.api as sm

np.random.seed(42)

def f(x, y):
    return np.exp(-(1-x)**2 - (y-x**2)**2) + np.exp(-(1+x)**2 - (y+3+x**2)**2)

def g(x, y, xm, ym, w, sigma):
    return w * K1(x,y) + (1-w) * K2(x,y,xm,ym,sigma)

def K1(x,y):
    return 0.5 * st.multivariate_normal.pdf([x,y],mean=(1,1),cov=0.5*np.eye(2)) \
        + 0.5 * st.multivariate_normal.pdf([x,y],mean=(-1,-3),cov=0.5*np.eye(2))

def K2(x,y,xm,ym,sigma):
    return st.multivariate_normal.pdf([x,y],mean=(xm,ym),cov=sigma*np.eye(2))

# Plot density
x = np.linspace(-3,3,101)
y = np.linspace(-10,7,101)
X,Y = np.meshgrid(x,y)
Z = f(X,Y)
plt.contourf(X,Y,Z)

# Parameters
alpha = 0.05
coeff = st.norm.ppf(1-alpha/2)
w = 0.7
sigma = 0.15

# Begin chain
Xn = [0,0]
Xchain = [Xn]
n = 1
n_acc = 0
while True:
    print(n)
    u1 = st.uniform.rvs()
    if u1 <= w:
        u2 = st.uniform.rvs()
        if u2<=0.5:
            Xnew = [1 + np.sqrt(0.5)*st.norm.rvs(), 1 + np.sqrt(0.5)*st.norm.rvs()]
        else:
            Xnew = [-1 + np.sqrt(0.5)*st.norm.rvs(), -3 + np.sqrt(0.5)*st.norm.rvs()]
    else:
        Xnew = [Xn[0]+sigma*st.norm.rvs(), Xn[1]+sigma*st.norm.rvs()]

    ratio = ( f(Xnew[0], Xnew[1]) * g(Xn[0], Xn[1], Xnew[0], Xnew[1], w, sigma) ) / ( f(Xn[0], Xn[1]) * g(Xnew[0], Xnew[1], Xn[0], Xn[1], w, sigma) )
    if st.uniform.rvs() <= np.min([1,ratio]):
        Xn = Xnew
        n_acc +=1
    Xchain.append(Xn)
    n = n + 1
    if n==1000:
        break
    
print(n_acc/n)
plt.plot([X[0] for X in Xchain], [X[1] for X in Xchain],'+')
plt.show()
