import gradio as gr
from typing import Dict, List, Union
import os
import pandas as pd
import pickle

class PreTrainedPipeline():
    def __init__(self, path=""):
        with open('iris_model.pickle', 'rb') as f:
            self.pipeline = pickle.load(f)

    def __call__(self, inputs):
        """
        Args:
            inputs (:obj:`dict`):
                a dictionary containing a key 'data' mapping to a dict in which
                the values represent each column.
        Return:
            A :obj:`list` of floats or strings: The classification output for each row.
        """
        # Convert the input dictionary to a pandas DataFrame
        data = inputs['data']  # Extracting the 'data' dictionary
        X = pd.DataFrame(data)  # Converting 'data' dictionary to DataFrame

        return self.pipeline.predict(X)
    
# Initialize your pre-trained pipeline
model = PreTrainedPipeline()

def predict(sepal_length, sepal_width, petal_length, petal_width):
    # Create a dictionary with the input data
    inputs = {'data': {
                'sepal_length': [sepal_length],
                'sepal_width': [sepal_width],
                'petal_length': [petal_length],
                'petal_width': [petal_width],
              }}
    # Use the model to predict the species
    prediction = model(inputs)
    # Assuming the model returns a list of predictions
    return prediction[0]  # Return the first (and only) prediction

# Define the Gradio interface with specific inputs for the Iris features
iface = gr.Interface(fn=predict,
                     inputs=[
                         gr.inputs.Number(label="Sepal Length (cm)"),
                         gr.inputs.Number(label="Sepal Width (cm)"),
                         gr.inputs.Number(label="Petal Length (cm)"),
                         gr.inputs.Number(label="Petal Width (cm)")
                     ],
                     outputs="text",
                     description="Predict the species of an Iris flower")

# Launch the app (when running locally)
if __name__ == "__main__":
    iface.launch()