import joblib
import os
import warnings

import gradio as gr

from huggingface_hub import hf_hub_download

warnings.filterwarnings("ignore")
TOKEN = os.environ['TOKEN']
REPO_ID = "SimonRaviv/wine-quality"
MODEL_FILENAME = "sklearn_model.joblib"

model_file_path = hf_hub_download(REPO_ID, MODEL_FILENAME, use_auth_token=TOKEN)
model = joblib.load(model_file_path)


def predict(data):
    prediction = model.predict(data.to_numpy()).tolist()
    prediction = [[p] for p in prediction]
    return prediction


headers = [
    "fixed acidity",
    "volatile acidity",
    "citric acid",
    "residual sugar",
    "chlorides",
    "free sulfur dioxide",
    "total sulfur dioxide",
    "density",
    "pH",
    "sulphates",
    "alcohol",
]
default = [
    [7.4, 0.7, 0, 1.9, 0.076, 11, 34, 0.9978, 3.51, 0.56, 9.4],
    [7.8, 0.88, 0, 2.6, 0.098, 25, 67, 0.9968, 3.2, 0.68, 9.8],
    [7.8, 0.76, 0.04, 2.3, 0.092, 15, 54, 0.997, 3.26, 0.65, 9.8],
]

inputs = gr.inputs.Dataframe(
    row_count=(3, "dynamic"),
    col_count=(11, "dynamic"),
    headers=headers,
    default=default)
outputs = gr.outputs.Dataframe(type="numpy", headers=["Quality"])

interface = gr.Interface(
    fn=predict,
    title="Wine Quality predictor with SKLearn",
    inputs=inputs,
    outputs=outputs
)
interface.launch()