import gradio as gr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from io import BytesIO
from PIL import Image

PRIMARY_COLOR = "#0f2c59"


# -------------------------
# BUSINESS LOGIC
# -------------------------
def compute_slotting(df):
    """
    Determines aisle & rack suggestions based on SKU velocity.
    """

    results = []
    for _, row in df.iterrows():
        sku = row["SKU"]
        vel = row["Velocity"]
        freq = int(row["Frequency"])

        if vel.lower() == "fast":
            aisle = 3
            rack = 14
            reason = "Fast-moving SKU — Placed close to dispatch for quicker picking."
        elif vel.lower() == "medium":
            aisle = 6
            rack = 20
            reason = "Medium-moving SKU — Positioned in central aisles to balance travel distance."
        else:
            aisle = 20
            rack = 6
            reason = "Slow-moving SKU — Moved to back aisles to avoid congestion."

        results.append([sku, vel, freq, aisle, rack, reason])

    return pd.DataFrame(
        results,
        columns=["SKU", "Velocity", "Frequency", "Suggested Aisle", "Suggested Rack", "Reason"]
    )

def generate_heatmap(slotting_df):
    """
    Generates a heatmap and returns a PIL Image (required by Gradio).
    """

    aisles = slotting_df["Suggested Aisle"].astype(int)
    racks = slotting_df["Suggested Rack"].astype(int)

    grid = np.zeros((25, 25))
    for a, r in zip(aisles, racks):
        if a < 25 and r < 25:
            grid[a, r] = 1

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.imshow(grid, cmap="Oranges", origin="lower")
    ax.set_title("Warehouse Slotting Heatmap")
    ax.set_xlabel("Rack Number")
    ax.set_ylabel("Aisle Number")

    buf = BytesIO()
    plt.savefig(buf, format="png", dpi=120, bbox_inches="tight")
    plt.close(fig)
    buf.seek(0)

    # convert to PIL Image (this is what Gradio expects)
    return Image.open(buf)



def business_summary(df):
    fast = sum(df["Velocity"].str.lower() == "fast")
    med = sum(df["Velocity"].str.lower() == "medium")
    slow = sum(df["Velocity"].str.lower() == "slow")

    summary = f"""
### 📊 Business Insight Summary

- **Fast Movers**: {fast} SKUs placed near dispatch → reduces pick time.
- **Medium Movers**: {med} SKUs placed in central aisles → balances travel.
- **Slow Movers**: {slow} SKUs placed in far aisles → lowers congestion.

#### This Improves:
- 🚶‍♂️ Reduction in walking distance  
- 🚚 Faster order fulfillment  
- 🏭 Better warehouse space utilization  
- 🔄 Reduced aisle congestion  
"""
    return summary


# -------------------------
# GRADIO UI
# -------------------------
def process_slotting(input_df):
    try:
        slotting = compute_slotting(input_df)
        heatmap = generate_heatmap(slotting)
        insights = business_summary(slotting)
        return slotting, heatmap, insights
    except Exception as e:
        return None, None, f"❌ Error: {e}"


def build_ui():
    with gr.Blocks() as demo:

        gr.Markdown(
            "<h1 style='color:#FF6A00'>Procelevate Inventory Slotting Optimizer</h1>"
            "AI-powered SKU placement engine to reduce picking time & congestion."
        )

        with gr.Tab("Optimized Slotting"):
            df_input = gr.DataFrame(
                headers=["SKU", "Velocity", "Frequency"],
                value=[
                    ["A123", "Fast", 120],
                    ["B555", "Medium", 60],
                    ["C888", "Slow", 5],
                ],
                label="SKU Velocity Table",
                interactive=True
            )

            run_btn = gr.Button("Optimize Slotting", variant="primary")

            slotting_table = gr.DataFrame(label="Optimized Slotting", interactive=False)
            heatmap_output = gr.Image(type="pil", label="Heatmap")
            insights_output = gr.Markdown()

            run_btn.click(
                process_slotting,
                inputs=df_input,
                outputs=[slotting_table, heatmap_output, insights_output]
            )

    return demo


demo = build_ui()

if __name__ == "__main__":
    demo.launch()