import numpy as np
import gradio as gr
from sklearn.tree import DecisionTreeRegressor
import pickle
# from joblib import load

tree_model = pickle.load(open('tree_model.pkl', 'rb'))
# tree_model = load('tree_model.joblib')


def predict(IW, IF, VW, FP):
    params = [[float(IW), float(IF), float(VW), float(FP)]]
    prediction = np.round(tree_model.predict(params), 2)
    depth, width = prediction[0]
    print(depth, width)
    return depth, width

input1 = gr.inputs.Slider(minimum=0, maximum=100, label="Input 1")
input2 = gr.inputs.Slider(minimum=0, maximum=100, label="Input 2")
input3 = gr.inputs.Slider(minimum=0, maximum=100, label="Input 3")
input4 = gr.inputs.Slider(minimum=0, maximum=100, label="Input 4")

output1 = gr.outputs.Textbox(label="Output")
output2 = gr.outputs.Textbox(label="Output")
interface = gr.Interface(fn=predict, inputs=[input1, input2, input3, input4], outputs=[output1, output2])
interface.launch(share=True)