import gradio as gr
import pandas as pd
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
from PIL import Image
import io

# Load model
model_revenue = joblib.load("model_revenue.pkl")
model_sold = joblib.load("model_sold.pkl")
model_stock = joblib.load("model_stock.pkl")

# Preprocessing input
def preprocess_input(product_type, price, availability, customer, lead, order, location, production):
    product_dict = {"haircare": 0, "skincare": 1, "fragrance": 2}
    gender_dict = {"Male": 0, "Female": 1}
    location_dict = {"Chennai": 0, "Mumbai": 1, "Delhi": 2}

    return pd.DataFrame([{
        "Product_Type": product_dict.get(product_type, 0),
        "Price": price,
        "Availability": availability,
        "Customer_Demographics": gender_dict.get(customer, 0),
        "Lead_Time": lead,
        "Order_Quantities": order,
        "Location": location_dict.get(location, 0),
        "Production_Volumes": production
    }])

# Manual prediction
def predict_manual(product_type, price, availability, customer, lead, order, location, production):
    df = preprocess_input(product_type, price, availability, customer, lead, order, location, production)
    revenue = model_revenue.predict(df)[0]
    sold = model_sold.predict(df)[0]
    stock = model_stock.predict(df)[0]
    return f"Revenue: {revenue:.2f}, Sold: {sold:.0f}, Stock: {stock:.0f}"

# CSV prediction
def predict_from_csv(file):
    df = pd.read_csv(file.name)

    # Rename columns
    df = df.rename(columns={
        "Product type": "Product_Type",
        "Price": "Price",
        "Availability": "Availability",
        "Customer demographics": "Customer_Demographics",
        "Lead time": "Lead_Time",
        "Order quantities": "Order_Quantities",
        "Location": "Location",
        "Production volumes": "Production_Volumes"
    })

    df['Product_Type'] = df['Product_Type'].astype('category').cat.codes
    df['Customer_Demographics'] = df['Customer_Demographics'].astype('category').cat.codes
    df['Location'] = df['Location'].astype('category').cat.codes

    X = df[["Product_Type", "Price", "Availability", "Customer_Demographics", 
            "Lead_Time", "Order_Quantities", "Location", "Production_Volumes"]]

    df["Predicted Revenue"] = model_revenue.predict(X)
    df["Predicted Products Sold"] = model_sold.predict(X)
    df["Predicted Stock Level"] = model_stock.predict(X)

    # Plot
    fig, axes = plt.subplots(3, 1, figsize=(10, 12))
    sns.barplot(x=df.index, y="Predicted Revenue", data=df, ax=axes[0])
    axes[0].set_title("Predicted Revenue")

    sns.lineplot(x=df.index, y="Predicted Products Sold", data=df, ax=axes[1], marker="o")
    axes[1].set_title("Predicted Products Sold")

    sns.lineplot(x=df.index, y="Predicted Stock Level", data=df, ax=axes[2], marker="o", color="green")
    axes[2].set_title("Predicted Stock Level")

    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(buf, format="png")
    plt.close()
    buf.seek(0)
    img = Image.open(buf)

    return df, img

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Supply Chain Prediction App")
    gr.Markdown("Masukkan data manual atau upload file CSV untuk prediksi pendapatan, produk terjual, dan stok.")

    with gr.Tab("Input Manual"):
        with gr.Row():
            product_type = gr.Dropdown(["haircare", "skincare", "fragrance"], label="Product Type")
            price = gr.Number(label="Price")
            availability = gr.Number(label="Availability")
            customer = gr.Dropdown(["Male", "Female"], label="Customer Demographics")
        with gr.Row():
            lead = gr.Number(label="Lead Time")
            order = gr.Number(label="Order Quantities")
            location = gr.Dropdown(["Chennai", "Mumbai", "Delhi"], label="Location")
            production = gr.Number(label="Production Volumes")
        output_manual = gr.Textbox(label="Hasil Prediksi")
        btn = gr.Button("Submit")
        btn.click(predict_manual, inputs=[product_type, price, availability, customer, lead, order, location, production], outputs=output_manual)

    with gr.Tab("Upload CSV"):
        file_input = gr.File(label="Upload CSV File")
        output_table = gr.Dataframe()
        output_plot = gr.Image(label="Visualisasi Prediksi")
        file_input.change(fn=predict_from_csv, inputs=file_input, outputs=[output_table, output_plot])

demo.launch()