import numpy as np
import scipy.special as sps
import matplotlib.pyplot as plt
from pandas import Series
import statsmodels.graphics.tsaplots as sm
import seaborn as sns
sns.set(color_codes=True)# Defines the pdf 
sns.set(font_scale=2) #fontsize in plots
sns.set_style("white")

np.random.seed(42)

def f(x,a,b):
    return b**a*x**(a-1)*np.exp(-b*x)/sps.gamma(a)*1*(x>=0)

def Cfunc(alpha, a, b):
    return ((alpha - a)/(1-b))**(alpha-a) * np.exp(a-alpha) * sps.gamma(a)/sps.gamma(alpha)/b**a

#### Some parameters
alpha=4.85 
a = np.floor(alpha)
b = np.floor(alpha)/alpha
C = Cfunc(alpha,a,b)

#### Does the accept reject part
n = 5000
Aar = 0
Xar = np.zeros(n)
i = 0
p = 0
while i<n:
    xi = np.random.gamma(shape = a, scale = 1/b, size=1)
    ratio = f(xi,alpha,1)/(f(xi,a,b)*C)
    p+=1
    if np.random.random(1) < ratio:
        Xar[i]=xi
        i+=1
Aar=n/p # acceptance probability

### Does the Metropolis Hastings part
Xmh = np.zeros(n)
Xmh[0] = 4
Amh=0
for i in range(n-1):
    #samples proposal
    y = np.random.gamma(shape = a, scale = 1/b, size=1)
    #computes MH ratio
    py=f(y,alpha,1)
    px=f(Xmh[i],alpha,1)
    qyx=f(Xmh[i],a,b)
    qxy=f(y,a,b)
    
    ratio=(py*qyx )/(px*qxy)
    if np.random.random(1) < ratio:
        Xmh[i+1]=y
        Amh+=1
    else:
        Xmh[i+1]=Xmh[i]
Amh=Amh/n        

print('---------------------------------------')
print('AR acceptance rate '+str(Aar))
print('MH acceptance rate '+str(Amh))
print('---------------------------------------')
print('AR mean '+str(np.mean(Xar)))
print('MH mean '+str(np.mean(Xmh)))
print('---------------------------------------')


xx=np.linspace(0,20,1000)
plt.plot(xx,C*f(xx,a,b))
plt.plot(xx,f(xx,alpha,1))
plt.legend(["Cg(x)","f(x)"] )
plt.show()

nn=np.arange(1,n+1)
plt.plot(nn,np.cumsum(Xar)/nn)
plt.plot(nn,np.cumsum(Xmh)/nn)            
plt.hlines(alpha,0,n,color='black')
plt.legend(["AR","MH","True mean"] )
plt.show()

sm.plot_acf(Xar,lags=20)
plt.gca().set_rasterized(True)
plt.legend(["AR"] )
plt.show()

sm.plot_acf(Xmh,lags=20)
plt.gca().set_rasterized(True)
plt.legend(["MH"] )
plt.show()

sns.kdeplot(Xar, shade=True, color="r")
plt.gca().set_rasterized(True)
x = np.linspace(0,20,100)
plt.plot(x, f(x,alpha,1),color='black')
plt.title('KDE for AR')
plt.show()

sns.kdeplot(Xmh, shade=True, color="r")
plt.gca().set_rasterized(True)
plt.plot(x, f(x,alpha,1),color='black')
plt.title('KDE for MH')
plt.show()