import gradio as gr
import pickle
import pandas as pd


PARAMS_NAME = [
    "Age",
    "Class",
    "Wifi",
    "Booking",
    "Seat",
    "Checkin"
]


with open("model/rf.pkl", "rb") as f:
    model = pickle.load(f)

COLUMNS_PATH = "model/categories_ohe.pickle"
with open(COLUMNS_PATH, 'rb') as handle:
    ohe_tr = pickle.load(handle)


def predict(*args):
    answer_dict = {}

    for i in range(len(PARAMS_NAME)):
        answer_dict[PARAMS_NAME[i]] = [args[i]]

    single_instance = pd.DataFrame.from_dict(answer_dict)

    # Reformat columns
    single_instance_ohe = pd.get_dummies(single_instance).reindex(columns = ohe_tr).fillna(0)
    
    prediction = model.predict(single_instance_ohe)

    response = int(prediction[0])

    if response == 0:
        response = "This flight was a really hell!!"
    if response == 1:
        response = "I have touch the sky with my hand, what a lovely flight!"

    return response

with gr.Blocks() as demo:
        gr.Markdown(
        '''
        # Flight satisfaction 🛩
        '''
        )

        with gr.Row():
            with gr.Column():
                
                gr.Markdown(
                    '''
                        ## Input 🛫

                    '''
                )

                Age = gr.Slider(label="Age", minimum=6, maximum=120, step=1, randomize=True)

                Class = gr.Radio(
                    label="Class",
                    choices=["Business", "Eco", "Eco Plus"],
                    value="Eco Plus"
                )

                Wifi = gr.Slider(label="Wifi", minimum=1, maximum=5, step=1, randomize=True)

                Booking = gr.Slider(label="Booking", minimum=1, maximum=5, step=1, randomize=True)

                Seat = gr.Slider(label="Seat", minimum=1, maximum=5, step=1, randomize=True)

                Checkin = gr.Slider(label="Checkin", minimum=1, maximum=5, step=1, randomize=True)

            with gr.Column():

                gr.Markdown(
                    '''
                        ## Prediction 🛬

                    '''
                )

                label = gr.Label(label="Satisfaction")
                predict_btn = gr.Button(value="Shoot")
                predict_btn.click(
                    predict,
                    inputs=[
                        Age,
                        Class,
                        Wifi,
                        Booking,
                        Seat,
                        Checkin,
                    ],
                    outputs=[label],
                    api_name="Flight satisfaction"
                )
        gr.Markdown(
            '''

            <p style='text-align:center'>
                <a href='https://www.escueladedatosvivos.ai/cursos/bootcamp-de-data-science'
                    target='_blank'>Estudia con Carlos Bustillo en Escuela de Datos Vivos haciendo click aqui y hace muchas de estas APIS 😎 !
                </a>
            </p>
            '''
        )

demo.launch()