import gradio as gr
from load_bnn_model import load_bnn_model
import numpy as np
import tensorflow as tf
from scipy.stats import norm
import matplotlib.pyplot as plt

#load our model
model = load_bnn_model()

#create ranges for our input variables
alcohol= [8, 15.0]
chlorides= [0.014, 0.22]
citric_acid= [0.0, 1.7]
density= [0.98, 1.00]
fixed_acidity= [4.7, 10.2]
free_sulfur_dioxide= [3.0, 128.0]
pH= [2.7, 3.8]
residual_sugar= [0.6, 26.]
sulphates= [0.2, 1.1]
total_sulfur_dioxide= [28.0, 294.0]
volatile_acidity= [0.07, 1.]

names = ['alcohol',
'chlorides',
'citric acid',
'density',
'fixed acidity',
'free sulfur dioxide',
'pH',
'residual sugar',
'sulphates',
'total sulfur dioxide',
'volatile acidity']

demo = gr.Blocks()

with demo:
    with gr.Row():
        gr.Markdown("""
        # Probabilistic Bayesian Neural Networks
        This space is a demo of the Bayesian network created in [this](https://keras.io/examples/keras_recipes/bayesian_neural_networks/) 
        keras example. The model takes in 11 physicochemical wine features and looks to predict wine quality, 
        which is a score from 0-10 determined by wine experts. However, the output of the model will be different from the typical
        regression paradigm. Rather than predicting a single point estimate for a given input, we produce a distribution. In this
        example, that distribution is Gaussian, with learnable mean and variance.
        
        While the model output is a distirbution, that is not all. We actually allow our model weights to be represented as distributions
        rather than single point estimates. Any time we call the model (even with the same inputs), the weights are sampled according to 
        their learned distribution. This means you will get a different result when calling the model several times on the same input data!
        
        
        Full credits for this model & example go to [Khalid Salama](https://www.linkedin.com/in/khalid-salama-24403144/).<br>
        Demo by [Brenden Connors](https://www.linkedin.com/in/brenden-connors-6a0512195).""")
    with gr.Row():
        
        with gr.Column():
            gr.Markdown("## Wine Features 🍷")
            with gr.Row():
                alcohol_slider = gr.Slider(minimum=alcohol[0], maximum=alcohol[1], label='alcohol', value=np.mean(alcohol))
                chlorides_slider = gr.Slider(minimum=chlorides[0], maximum=chlorides[1], label='chlorides', value=np.mean(chlorides))
            with gr.Row():
                citric_acid_slider = gr.Slider(minimum=citric_acid[0], maximum=citric_acid[1], label='citric_acid', value=np.mean(citric_acid))
                density_slider = gr.Slider(minimum=density[0], maximum=density[1], label='density', value=np.mean(density))
            with gr.Row():
                fixed_acidity_slider = gr.Slider(minimum=fixed_acidity[0], maximum=fixed_acidity[1], label='fixed_acidity', value=np.around(np.mean(fixed_acidity), 2))
                free_sulfur_dioxide_slider = gr.Slider(minimum=free_sulfur_dioxide[0], maximum=free_sulfur_dioxide[1], label='free_sulfur_dioxide', value=np.mean(free_sulfur_dioxide))
            with gr.Row():
                pH_slider = gr.Slider(minimum=pH[0], maximum=pH[1], label='pH', value=np.mean(pH))
                residual_sugar_slider = gr.Slider(minimum=residual_sugar[0], maximum=residual_sugar[1], label='residual_sugar', value=np.mean(residual_sugar))  
            with gr.Row():
                sulphates_slider = gr.Slider(minimum=sulphates[0], maximum=sulphates[1], label='sulphates', value=np.mean(sulphates))
                total_sulfur_dioxide_slider = gr.Slider(minimum=total_sulfur_dioxide[0], maximum=total_sulfur_dioxide[1], label='total_sulfur_dioxide', value=np.mean(total_sulfur_dioxide))
            with gr.Row():
                volatile_acidity_slider = gr.Slider(minimum=volatile_acidity[0], maximum=volatile_acidity[1], label='volatile_acidity', value=np.mean(volatile_acidity))
                number_samples = gr.Dropdown(choices=[1,2,3,4,5], label='number of samples')
            b = gr.Button()
        with gr.Column():
            gr.Markdown("## Output Distribution")
            with gr.Row():
                out = gr.Plot(None)
     
    
    all_sliders = [alcohol_slider,
                    chlorides_slider,
                    citric_acid_slider,
                    density_slider,
                    alcohol_slider,
                    chlorides_slider,
                    fixed_acidity_slider,
                    free_sulfur_dioxide_slider,
                    pH_slider,
                    residual_sugar_slider,
                    sulphates_slider,
                    total_sulfur_dioxide_slider,
                    volatile_acidity_slider]
    fig = gr.Variable(None)
    std = gr.Variable(None)
    
    def predict(n, *args):
        fig = plt.figure()
        plt.title('Wine Quality Distribution')
        plt.xlabel('Wine Quality Score')
        plt.ylabel('Density')
        for i in range(n):
            out = (model({n:tf.constant([a]) for n,a in zip(names,args)}))
            mean =(out.mean().numpy()[0][0])
            std = (out.stddev().numpy()[0][0])
            x_axis = np.arange(mean - 4*std, mean+4*std, 0.01)
            plt.plot(x_axis, norm.pdf(x_axis, mean, std))
        
        
        return fig

        
    b.click(fn=predict, inputs=[number_samples] + all_sliders, outputs=out)
    
demo.launch()