import pandas as pd
import numpy as np
import gradio as gr
from sklearn.ensemble import RandomForestRegressor
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

# Load and prepare data
df = pd.read_csv("supply_chain_data.csv")
df = df.dropna()

# Features and targets
X = df.drop(columns=["Revenue generated", "Number of products sold", "Stock levels"])
y_revenue = df["Revenue generated"]
y_sold = df["Number of products sold"]
y_stock = df["Stock levels"]

categorical_cols = X.select_dtypes(include=["object"]).columns.tolist()
numerical_cols = X.select_dtypes(include=["number"]).columns.tolist()

# Preprocessing pipeline
preprocessor = ColumnTransformer([
    ("cat", OneHotEncoder(handle_unknown="ignore"), categorical_cols),
    ("num", StandardScaler(), numerical_cols)
])

# Train model function
def train_model(target):
    pipeline = Pipeline(steps=[
        ("preprocessor", preprocessor),
        ("regressor", RandomForestRegressor(random_state=42))
    ])
    pipeline.fit(X, target)
    return pipeline

# Train all three models
model_revenue = train_model(y_revenue)
model_sold = train_model(y_sold)
model_stock = train_model(y_stock)

# Single prediction
def predict_single(product_type, price, availability, customer_demographics, lead_times, order_quantities, location, production_volumes):
    input_dict = {
        "Product type": [product_type],
        "Price": [float(price)],
        "Availability": [int(availability)],
        "Customer demographics": [customer_demographics],
        "Lead times": [int(lead_times)],
        "Order quantities": [int(order_quantities)],
        "Location": [location],
        "Production volumes": [int(production_volumes)],
    }
    input_df = pd.DataFrame(input_dict)

    pred_revenue = model_revenue.predict(input_df)[0]
    pred_sold = model_sold.predict(input_df)[0]
    pred_stock = model_stock.predict(input_df)[0]

    return f"Revenue: {pred_revenue:.2f} | Sold: {pred_sold:.0f} | Stock: {pred_stock:.0f}"

# Batch prediction from uploaded CSV

def predict_from_csv(file):
    import matplotlib.pyplot as plt
    import seaborn as sns
    from PIL import Image
    import io

    input_df = pd.read_csv(file.name)
    pred_revenue = model_revenue.predict(input_df)
    pred_sold = model_sold.predict(input_df)
    pred_stock = model_stock.predict(input_df)

    input_df["Predicted Revenue"] = pred_revenue
    input_df["Predicted Products Sold"] = pred_sold
    input_df["Predicted Stock Level"] = pred_stock

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

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

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

    plt.tight_layout()

    # Save to buffer and return PIL image
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)
    image = Image.open(buf)

    return input_df, image

    input_df = pd.read_csv(file.name)
    pred_revenue = model_revenue.predict(input_df)
    pred_sold = model_sold.predict(input_df)
    pred_stock = model_stock.predict(input_df)

    input_df["Predicted Revenue"] = pred_revenue
    input_df["Predicted Products Sold"] = pred_sold
    input_df["Predicted Stock Level"] = pred_stock

    return input_df

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

    with gr.Tab("Input Manual"):
        with gr.Row():
            product_type = gr.Dropdown(choices=["haircare", "skincare"], label="Product Type")
            price = gr.Number(label="Price")
            availability = gr.Number(label="Availability")
            customer_demographics = gr.Dropdown(choices=["Male", "Female", "Non-binary", "Unknown"], label="Customer Demographics")
        with gr.Row():
            lead_times = gr.Number(label="Lead Time")
            order_quantities = gr.Number(label="Order Quantities")
            location = gr.Dropdown(choices=["Mumbai", "Kolkata", "Delhi", "Chennai"], label="Location")
            production_volumes = gr.Number(label="Production Volumes")
        
        output = gr.Textbox(label="Hasil Prediksi")
        btn = gr.Button("Submit")
        btn.click(fn=predict_single, inputs=[
            product_type, price, availability, customer_demographics,
            lead_times, order_quantities, location, production_volumes
        ], outputs=output)

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

demo.launch()