import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_probability as tfp
import pandas as pd
tfd = tfp.distributions
tfl = tfp.layers

st.title("1 dimensional normal distribution")
mean = st.slider('Mean', -5, 5, 0)
std = st.slider('Scale', 0, 5, 1)

p = tfd.Normal(2, 1)

z = f"""\\begin{{array}}{{cc}}
  \mu & {mean} \\\\
  \sigma & {std}
\\end{{array}}
"""

st.latex(z)


q = tfd.Normal(mean, std)
z_values = tf.linspace(-5, 5, 200)
z_values = tf.cast(z_values, tf.float32)
prob_values_p = p.prob(z_values)
prob_values_q = q.prob(z_values)

fig, ax = plt.subplots()
ax.plot(z_values, prob_values_p, label=r'p', linestyle='--', lw=5, alpha=0.5)
ax.plot(z_values, prob_values_q, label=r'q')

ax.set_xlabel("x")
ax.set_ylabel("PDF(x)")
ax.legend()
ax.set_ylim((0, 1))


kl = tfd.kl_divergence(q, p)
st.latex(f"D_{{KL}}(q||p) \\text{{  is : }}{kl:0.2f}")

ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)

# Only show ticks on the left and bottom spines
ax.yaxis.set_ticks_position('left')
ax.xaxis.set_ticks_position('bottom')

st.pyplot(fig)
hide_streamlit_style = """
            <style>
            #MainMenu {visibility: hidden;}
            footer {visibility: hidden;}
            </style>
            """
st.markdown(hide_streamlit_style, unsafe_allow_html=True)