import gradio as gr
import numpy as np
from PIL import Image
import requests
import xgboost
import pandas as pd

from app_funcs import *

import hopsworks
import joblib

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

X_columns_to_drop = ["percentage_future", "cohort_first_month", "month",  
                     "cohort_first_product"]

feature_view = fs.get_feature_view(name="cohorts_fv", version=1)

mr = project.get_model_registry()
model = mr.get_model("cohort_model", version=1)
model_dir = model.download()
model = joblib.load(model_dir + "/xgb_.pkl")


def cohort_predict(cohort_start, start_date, product):

    # Get feature view in a dataframe
    data, labels = feature_view.get_training_data(training_dataset_version=1)
    data["percentage_future"] = labels

    # Convert it to pandas datetime
    data["cohort_first_month"] = pd.to_datetime(data["cohort_first_month"])
    data["month"] = pd.to_datetime(data["month"])

    # Sort and assert
    data = data.sort_values(by=["cohort_first_product", "cohort_first_month", "month"])
    
    new_seq = generate_new_data(data, cohort_start, start_date, product, model, X_columns_to_drop, 12)
    hist_seq = get_sequence(data, cohort_start, product)
    print(hist_seq.head())

    fig = plot_example_from_case(hist_seq, new_seq, 25, product)
    fig.canvas.draw()

    arr = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    arr = arr.reshape(fig.canvas.get_width_height()[::-1] + (3,))    
    
    return Image.fromarray(arr)
        
demo = gr.Interface(
    fn=cohort_predict,
    title="Cohort Active Percentage Prediction",
    description="Predicts active user percentage in a future month for a cohort that started in specific date with specific product",
    allow_flagging="never",
    inputs=[
        gr.Textbox(default='2022-04-01', label="Cohort Start Date"),
        gr.Textbox(default='2022-04-01', label="Prediction Start Date"),
        gr.Textbox(default="3m", label="Product (1m, 3m, 4m)"),
        ],
    outputs=gr.Image(type="pil"))

demo.launch()