import numpy as np
import scipy.stats as st
import matplotlib.pyplot as plt
import math
from tools import *

np.random.seed(42)

#Defines the KDE function
def kde(X, data, h = None):
	N = data.shape[0]
	if h is None:
		h = 1.06 * np.std(data) * N ** (-1. / 5)
	pdf = np.zeros(X.shape[0])
	for i in range(X.shape[0]):
		pdf[i] = np.sum(np.exp( - ((X[i] - data) / h)** 2 / 2.) / np.sqrt(2*math.pi)) / (N * h)
	return pdf
# Determines some parameters
c, d = 2, 4
N = 1000
#generates samples form the Bur XII distr
X = st.burr12(c, d).rvs(size = (N,))

#plots density and KDE
x = np.linspace(0, 3, 1001)
plt.plot(x, st.burr12(c,d).pdf(x), linewidth = 1.)
plt.plot(x, kde(x, X), linewidth = 1.)
plt.hist(X, bins = 100, density = True)
plt.legend(['Burr XII','KDE','Histogram'])

plt.show()

#Prepares to plot
n = [100 * 10**i for i in range(4)]
delta = np.array(n) ** (-1. / 5.)
print(n)
#plots
test_vals = np.zeros(len(n))
fig = plt.figure(figsize = (14,3))
for i in range(len(n)):
	X = st.burr12(c,d).rvs(size = (n[i],))
	h = np.std(X) * delta[i]
	w = np.random.randint(0, high = n[i], size = 20000)
	X_kde = np.random.normal(loc = X[w], scale = h)
	ax = fig.add_subplot(1,4,i+1)
	ax.plot(x, st.burr12(c,d).pdf(x), linewidth = 1.)
	ax.plot(x, kde(x, X, h = h), linewidth = 1.)
	ax.hist(X_kde, bins = 100, density = True)
	ax.set_title('n = ' + str(n[i]))
	mess, stat = KSTest(X_kde, alpha = 0.1)
	test_vals[i] = stat['Statistic']
	print(mess)


handles, labels = ax.get_legend_handles_labels()
fig.legend(['Burr XII', 'KDE', 'Histogram'],
           loc='lower center',
           ncol=3,
           bbox_to_anchor=(0.5, 0.02))
plt.tight_layout(rect=[0, 0.05, 1, 1])
plt.show()


plt.loglog(10.0**np.arange(4),test_vals, '-x')
plt.title('KS values')
plt.show()