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


MODEL_NAME = "lkerf"
HF_USER = "universalml"
REPO_ID = HF_USER + "/" + MODEL_NAME
model = joblib.load(hf_hub_download(repo_id=REPO_ID, filename=MODEL_NAME))


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.input_feature_names)
    data_frame = encode_categorical_columns(input_data_frame)
    scaler = StandardScaler()
    # Scale the input data using the loaded scaler
    scaled_input = scaler.transform(data_frame)
    prediction_result = model.predict(scaled_input)[0]

    return prediction_result


def regression_inputs():
    input_labels = model.input_feature_names
    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.output_feature_name
    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()