#!/usr/bin/env python
# coding: utf-8

# In[48]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt
import numpy as np
save = False


# ## 確率分布
# 
# - ### $P(x)$
#     $\mu=0$
#     
#     $\sigma^2=2.5$
# - ### $Q(x)$
#     $\mu=3$
#     
#     $\sigma^2=2$

# In[49]:


mean_p, vari_p= 0, 2.5
mean_q, vari_q= 3, 2
P = lambda x: (np.e**-(((x-mean_p)**2)/(2*vari_p)))/(2*np.pi*vari_p)**0.5
Q = lambda x: (np.e**-(((x-mean_q)**2)/(2*vari_q)))/(2*np.pi*vari_q)**0.5


# In[50]:


x_min, x_max = -6, 7.5
x = np.arange(x_min,x_max,0.1)
plt.plot(x,P(x),label="P(x)")
plt.plot(x,Q(x),label="Q(x)")
plt.legend()
plt.xlim(x_min,x_max)
plt.xlabel("x",fontsize=15)
plt.ylabel("probability",fontsize=15)
if save: plt.savefig("prob")


# # KL Divergence
# ### $KL(P||Q) = E_{x \sim P(x)} log \cfrac{P(x)}{Q(x)}$

# ## 情報量
# ### $I_p=-log P(x)$
# ### $I_q=-log Q(x)$

# In[51]:


I = lambda p: -np.log(p)


# In[52]:


# x_min, x_max = -4, 6.5
x = np.arange(x_min,x_max,0.1)
plt.plot(x,I(P(x)),label="I_p(x)")
plt.plot(x,I(Q(x)),label="I_q(x)")
plt.legend()
plt.xlim(x_min,x_max)
plt.xlabel("x",fontsize=15)
plt.ylabel("Information",fontsize=15)
if save: plt.savefig("info")


#  ## 差分
#  ### $D(x) = I_q-I_p$

# In[53]:


# x_min, x_max = -4, 6.5
x = np.arange(x_min,x_max,0.1)
plt.plot(x,I(Q(x))-I(P(x)))
plt.xlim(x_min,x_max)
plt.xlabel("x",fontsize=15)
plt.ylabel("difference",fontsize=15)
if save: plt.savefig("diff")


# ## 重み付き差分
#    ### $WD(x) = P(x)(I_q-I_p)$
#    
#    #### ※ $KL(P||Q) = \int WD(x) dx$

# In[54]:


# x_min, x_max = -4, 6.5
x = np.arange(x_min,x_max,0.1)
plt.plot(x,P(x)*(I(Q(x))-I(P(x))),label="WD")
plt.plot(x,P(x),label="P(x)")
plt.plot(x,Q(x),label="Q(x)")
plt.legend()
plt.xlim(x_min,x_max)
plt.xlabel("x",fontsize=15)
plt.ylabel("probability",fontsize=15)
if save: plt.savefig("wd_p")


# ### $WD_q(x)=Q(x)(I_p-I_q)$

# In[55]:


# x_min, x_max = -4, 6.5
x = np.arange(x_min,x_max,0.1)
plt.plot(x,Q(x)*(I(P(x))-I(Q(x))),label="WD")
plt.plot(x,P(x),label="P(x)")
plt.plot(x,Q(x),label="Q(x)")
plt.legend()
plt.xlim(x_min,x_max)
plt.xlabel("x",fontsize=15)
plt.ylabel("probability",fontsize=15)
if save: plt.savefig("wd_q")


# # JS Divergence
# ### $JS(P||Q) = \cfrac{1}{2}KL\Big(P || \cfrac{P+Q}{2}\Big) + \cfrac{1}{2}KL\Big(Q || \cfrac{P+Q}{2}\Big)$

# ## Mean

# In[58]:


x = np.arange(x_min,x_max,0.1)
plt.plot(x,(P(x)+Q(x))/2)
plt.xlim(x_min,x_max)
plt.xlabel("x",fontsize=15)
plt.ylabel("probability",fontsize=15)
plt.title("Mean")
if save: plt.savefig("mean")


# In[ ]: