import gradio as gr
import tensorflow as tf
import numpy as np
import os
import PIL
import PIL.Image

# Create a Gradio App using Blocks    
with gr.Blocks() as demo:
    gr.Markdown(
    """
    # AI/ML Playground
    """
    )
    with gr.Accordion("Click for Instructions:"):
            gr.Markdown(
    """
    * uploading an image will setup, train, and evaluate the base model
    """)

    # Train, evaluate and test a ML
    # image classification model for
    # clothes images
    def modelTraining(img):
        # clothing dataset
        mnist = tf.keras.datasets.mnist

        #split the training data in to a train/test sets
        (x_train, y_train), (x_test, y_test) = mnist.load_data()
        x_train, x_test = x_train / 255.0, x_test / 255.0

        # create the neural net layers
        model = tf.keras.models.Sequential([
          tf.keras.layers.Flatten(input_shape=(28, 28)),
          tf.keras.layers.Dense(128, activation='relu'),
          tf.keras.layers.Dropout(0.2),
          tf.keras.layers.Dense(10)
        ])

        #make a post-training predition on the 
        #training set data
        predictions = model(x_train[:1]).numpy()

        # converts the logits into a probability
        tf.nn.softmax(predictions).numpy()

        #create and train the loss function
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
        loss_fn(y_train[:1], predictions).numpy()

        # compile the model with the loss function
        model.compile(optimizer='adam',
                      loss=loss_fn,
                      metrics=['accuracy'])
        
        # train the model - 5 runs
        # evaluate the model on the test set
        model.fit(x_train, y_train, epochs=5)
        test_loss, test_acc = model.evaluate(x_test,  y_test, verbose=2)

        print("Test accuracy: ", test_acc)

        # Define any necessary preprocessing steps for the image input here
        # the image can be passed as a PIL or numpy
        # create the final model for production
        probability_model = tf.keras.Sequential([model, tf.keras.layers.Softmax()])
        
        # Make a prediction using the model
        prediction = probability_model.predict(img)

        # Postprocess the prediction and return it
        return np.argmax(predictions[0])
        

    # Creates the Gradio interface objects
    with gr.Row():
        with gr.Column(scale=2):
            image_data = gr.Image(label="Upload Image", type="numpy")
        with gr.Column(scale=1):
            model_performance = gr.Text(label="Model Performance", interactive=False)
            model_prediction = gr.Text(label="Model Prediction", interactive=False)
    image_data.change(modelTraining, image_data, model_prediction)
    
    
# creates a local web server
# if share=True creates a public
# demo on huggingface.c
demo.launch(share=False)