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

project = hopsworks.login()
fs = project.get_feature_store()

dataset_api = project.get_dataset_api()

dataset_api.download("Resources/images/wine/df_recent.png")
dataset_api.download("Resources/images/wine/confusion_matrix.png")

#split one column as targets
def feature_target_split(df, column_name):
    features = df.drop(column_name, axis=1).reset_index(drop=True)
    targets = df[column_name].reset_index(drop=True)
    return features, targets

# One-hot encode the categorical feature
def one_hot_encoder(df, column_name):
    import pandas as pd
    from sklearn.preprocessing import OneHotEncoder

    encoder = OneHotEncoder(drop='first',sparse=False)
    encoded_features = encoder.fit_transform(df[[column_name]])
    encoded_df = pd.DataFrame(encoded_features, columns=encoder.get_feature_names_out([column_name]))
    encoded_df = pd.concat([df.drop(column_name, axis=1), encoded_df], axis=1)
    return encoded_df

def find_thresholds(train, test, column_name):
    import pandas as pd

    df = pd.concat([train, test], axis=0)
    column_list = list(df[column_name])
    column_list.sort()
    lower_threshold_idx = int(len(column_list)/3)
    upper_threshold_idx = lower_threshold_idx * 2
    lower_threshold = column_list[lower_threshold_idx]
    upper_threshold = column_list[upper_threshold_idx]
    return lower_threshold, upper_threshold

def quality2class(series, lower_threshold, upper_threshold):
    import pandas as pd

    quality_list = list(series)
    class_list = []
    for i in range(len(quality_list)):
        if quality_list[i] < lower_threshold:
            class_list.append('poor')
        elif quality_list[i] > upper_threshold:
            class_list.append('excellent')
        else:
            class_list.append('fair')
    class_series= pd.Series(class_list, name = 'class')

    return class_series

# Read the latest wine added to the feature store
wine_fg = fs.get_feature_group(name="wine", version=1)
wine_df = wine_fg.read(read_options={"use_hive": True})
latest_wine_df = wine_df.tail(1)
latest_wine_string = latest_wine_df.to_string(index=False)

# Predict the wine class
mr = project.get_model_registry()
model = mr.get_model("wine_class_model", version=1)
model_dir = model.download()
model = joblib.load(model_dir + "/wine_class_model.pkl")
batch_data, batch_label = feature_target_split(wine_df, 'quality')
batch_data = one_hot_encoder(batch_data, 'type')
y_pred = model.predict(batch_data.tail(1))

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            gr.Label("Most Recent Wine Added")
            out = gr.Textbox(latest_wine_string, readonly=True)
        with gr.Column():
            gr.Label("Predicted Wine Class")
            out = gr.Textbox(y_pred[0], readonly=True)
    with gr.Row():
      with gr.Column():
          gr.Label("Recent Prediction History")
          input_img = gr.Image("df_recent.png", elem_id="recent-predictions")
      with gr.Column():          
          gr.Label("Confusion Matrix with Historical Prediction Performance")
          input_img = gr.Image("confusion_matrix.png", elem_id="confusion-matrix")        

demo.launch()