import numpy as np
import scipy.stats as st
import scipy.optimize as opt
import matplotlib.pyplot as plt
from plotstyle import *

fig, (ax1, ax2) = plt.subplots(1, 2)
fig.subplots_adjust(**adjust_fs(fig, right=1.0))

# the data:
n = 100
shape = 2.0
scale = 1.0
rng = np.random.RandomState(4637281)
xd = rng.gamma(shape, scale, n)

# true pdf:
xx = np.arange(0.0, 10.1, 0.01)
rv = st.gamma(shape)
yy = rv.pdf(xx)

# mle fit:
a = st.gamma.fit(xd, 5.0)
yf = st.gamma.pdf(xx, *a)

# plot it:
ax1.set_xlim(0, 10.0)
ax1.set_ylim(0.0, 0.42)
ax1.set_xticks(np.arange(0, 11, 2))
ax1.set_yticks(np.arange(0, 0.42, 0.1))
ax1.set_xlabel('x')
ax1.set_ylabel('Prob. density')
ax1.plot(xx, yy, label='pdf', **lsB)
ax1.plot(xx, yf, label='mle', **lsCm)
kernel = st.gaussian_kde(xd)
x = kernel(xd)
x /= np.max(x)
sigma = 0.07
ax1.plot(xd, sigma*x*(rng.rand(len(xd))-0.5)+sigma, zorder=10, **psAm)
ax1.legend(loc='upper right')

# histogram:
h,b = np.histogram(xd, np.arange(0, 8.4, 0.5), density=True)

# fit histogram:
def gammapdf(x, n, l, s) :
    return st.gamma.pdf(x, n, l, s)
popt, pcov = opt.curve_fit(gammapdf, b[:-1]+0.5*(b[1]-b[0]), h)
yc = st.gamma.pdf(xx, *popt)

# plot it:
ax2.set_xlim(0, 10.0)
ax2.set_xticks(np.arange(0, 11, 2))
ax2.set_xlabel('x')
ax2.set_ylim(0.0, 0.42)
ax2.set_yticks(np.arange(0, 0.42, 0.1))
ax2.set_ylabel('Prob. density')
ax2.plot(xx, yy, label='pdf', **lsB)
ax2.plot(xx, yc, label='fit', **lsCm)
ax2.bar(b[:-1], h, np.diff(b), **fsA)
ax2.legend(loc='upper right')

plt.savefig('mlepdf.pdf')