#%%
import numpy as np
from scipy.stats import norm, poisson

from  matplotlib import pyplot as plt
import tensorflow as tf
import seaborn as sns
sns.set()


def kl(p, q):
  return np.sum(np.where(p!=0, p*np.log(p/q),0))


# example
x = np.arange(-10,10,0.001)
p = norm.pdf(x,0,1)
q = norm.pdf(x,1,2)

plt.title(f'KL(P||Q) = {kl(p,q):.2f}')
plt.plot(x,p)
plt.plot(x,q, color="red")
plt.show()
#%%
# poisson distribution
from scipy.stats import poisson as poi
r = poisson.cdf(x, 0.9)
plt.plot(x, r, color="green")
plt.show()

x = np.arange(20)
lambda_ = [1.5, 4.25]
colors = ["#348ABD", "#A60628"]

# lambda = 1.5
plt.bar(x, poi.pmf(x, 1.5), color="red", label="$\lambda = 1.5$", alpha=0.60, lw="3")

plt.bar(x, poi.pmf(x, 4.5), color="blue", label="$\lambda = 4.5$", alpha=0.60, lw="3")

plt.xticks(x + 0.4, x)
plt.ylabel("Probability of $k$")
plt.xlabel("$k$")
plt.show()

#%%
# exponentional dist (EXP)
from scipy.stats import expon as expo

x = np.linspace(0,4,100)
colors = ['red','blue']
lambdas = [0.5, 1.0]

for l, c in zip(lambdas, colors):
  y = expo.pdf(x, scale=1.0/l)
  plt.plot(x, y, lw=2, color=c)