from sklearn.preprocessing import LabelEncoder
from huggingface_hub import hf_hub_download
import pandas as pd
import gradio as gr
import joblib


MODEL_NAME = "regressiontest"
HF_USER = "universalml"
REPO_ID = HF_USER + "/" + MODEL_NAME
MODEL = joblib.load(hf_hub_download(repo_id=REPO_ID, filename="model.joblib"))
SCALER = joblib.load(hf_hub_download(repo_id=REPO_ID, filename="scaler.joblib"))


def encode_categorical_columns(data_frame):
    label_encoder = LabelEncoder()
    ordinal_columns = data_frame.select_dtypes(include=['object']).columns

    for col in ordinal_columns:
        data_frame[col] = label_encoder.fit_transform(data_frame[col])

    nominal_columns = data_frame.select_dtypes(include=['object']).columns.difference(ordinal_columns)
    data_frame = pd.get_dummies(data_frame, columns=nominal_columns, drop_first=True)

    return data_frame


def prediction_function(*args):
    values_list = []

    for arg in args:
        values_list.append(int(arg))

    input_data_frame = pd.DataFrame([values_list], columns=MODEL.data)
    data_frame = encode_categorical_columns(input_data_frame)
    scaled_input = SCALER.transform(data_frame)
    prediction_result = MODEL.predict(scaled_input)[0]

    return prediction_result


def regression_inputs():
    input_labels = MODEL.data
    inputs = []

    for input_label in input_labels:
        value = gr.Textbox(label=input_label, type="text")
        inputs.append(value)

    return inputs


def regression_output():
    output_label = MODEL.target
    output = gr.Textbox(label=output_label, type="text")

    return output


def create_interface():
    interface = gr.Interface(
        fn=prediction_function,
        inputs=regression_inputs(),
        outputs=regression_output(),
        title=MODEL_NAME,
    )

    interface.launch(debug=True)


create_interface()