import gradio as gr
from PIL import Image
import requests
import hopsworks
import joblib
import pandas as pd

# Settings
# - - Names
# - - - Models
model_wine_dir = "wine_models"
model_red_wine_name = "wine_red_model"
model_red_wine_version = 1
model_white_wine_name = "wine_white_model"
model_white_wine_version = 1

def wine_quality(wine_type, fixed_acidity, volatile_acidity, citric_acid, residual_sugar, chlorides,
                 free_sulfur_dioxide, total_sulfur_dioxide, density, ph, sulphates, alcohol):
    print("Calling wine_quality() function")
    df = pd.DataFrame([[fixed_acidity, volatile_acidity, citric_acid, residual_sugar, chlorides,
                        free_sulfur_dioxide, total_sulfur_dioxide, density, ph, sulphates, alcohol]],
                      columns=['fixed_acidity', 'volatile_acidity', 'citric_acid', 'residual_sugar',
                               'chlorides', 'free_sulfur_dioxide', 'total_sulfur_dioxide', 'density',
                               'ph', 'sulphates', 'alcohol'])
    print("Predicting...")
    print(df)

    if wine_type.lower() == 'red':
        quality = red_wine_local_model.predict(df)
    elif wine_type.lower() == 'white':
        quality = white_wine_local_model.predict(df)
    else:
        quality = "Unknown wine type. Only 'red' and 'white' are known."

    return quality


print("Logging in to Hopsworks...")
project = hopsworks.login()

print("Getting feature store...")
fs = project.get_feature_store()

print("Getting model registry...")
mr = project.get_model_registry()

print("Getting model for red wine: ...")
model_red_wine = mr.get_model(model_red_wine_name, version=model_red_wine_version)

print("Downloading model for red wine...")
model_red_dir = model_red_wine.download()

print("Initializing model red wine locally...")
red_wine_local_model = joblib.load(model_red_dir + f"/{model_red_wine_name}.pkl")

print("Getting model for white wine: ...")
model_white_wine = mr.get_model(model_white_wine_name, version=model_white_wine_version)

print("Downloading model for white wine...")
model_white_dir = model_white_wine.download()

print("Initializing model white wine locally...")
white_wine_local_model = joblib.load(model_white_dir + f"/{model_white_wine_name}.pkl")


# Define the Gradio interface
iface = gr.Interface(fn=wine_quality,
                     inputs=["text", "number", "number", "number", "number", "number",
                             "number", "number", "number", "number", "number", "number"],
                     outputs="text",
                     examples=[['red', 7.4, 0.70, 0.00, 1.9, 0.076, 11.0, 34.0, 0.9978, 3.51, 0.56, 9.4],
                               ['white', 7.0, 0.27, 0.36, 20.7, 0.045, 45.0, 170.0, 1.0010, 3.00, 0.45, 8.8]])

iface.launch()