import gradio as gr
import pandas as pd
import os, shutil
from PIL import Image
from ultralyticsplus import YOLO


# Gradio Theme
theme = gr.themes.Soft(
    primary_hue="yellow",
    secondary_hue="blue",
    neutral_hue="gray",
    font=[gr.themes.GoogleFont('Inter'), 'ui-sans-serif', 'system-ui', 'sans-serif'],
    font_mono=[gr.themes.GoogleFont('Inter'), 'ui-monospace', 'Consolas', 'monospace'],
).set(
    background_fill_primary='*neutral_100',
)


# Bread Prices
bread_types = {
    "baguette": {"name": "Baguette", "price": 108},
    "binangkal": {"name": "Binangkal", "price": 11},
    "bonete": {"name": "Bonete", "price": 8},
    "cornbread": {"name": "Cornbread", "price": 55},
    "croissant": {"name": "Croissant", "price": 75},
    "ensaymada": {"name": "Ensaymada", "price": 14},
    "flatbread": {"name": "Flatbread", "price": 17},
    "kalihim": {"name": "Kalihim", "price": 15},
    "monay": {"name": "Monay", "price": 6},
    "pandesal": {"name": "Pandesal", "price": 3},
    "sourdough": {"name": "Sourdough", "price": 150},
    "spanish-bread": {"name": "Spanish Bread", "price": 14},
    "wheat-bread": {"name": "Wheat Bread", "price": 8},
    "white-bread": {"name": "White Bread", "price": 4},
    "whole-grain-bread": {"name": "Whole Grain Bread", "price": 10},
}


# Instantiate the model
model = YOLO("best.pt")


# Converts image input into a list
def preprocess_image(image):
    img_list = []
    
    # Preprocess the image and add it to the list 
    for im in image:
        image = Image.open(im.name)
        resize_img = image.resize((640, 640))
        img_list.append(resize_img)
    
    return img_list


# Gets all output images
def get_predictions(directory):
    allowed_extensions = ('.png', '.jpg', '.jpeg', '.gif', '.bmp')
    return [
        os.path.join(root, file)
        for root, _, files in os.walk(directory)
        for file in files
        if file.lower().endswith(allowed_extensions)
    ]


# Clear output from previous detection
def clear_output():
    shutil.rmtree('output/', ignore_errors=True)


# Bread Prediction function
def detect_bread(image):
    clear_output()
    image_list = preprocess_image(image)
    results = model.predict(image_list, conf=0.4, save=True, hide_conf=True, project = "output", name = "results")

    detected_classes = []

    for result in results:
        for cls in result.boxes.cls:                                # Stores all detected classes in the list
            detected_classes.append(result.names[int(cls)])   

    receipt = generate_receipt(detected_classes)

    return get_predictions(f'output/results'), receipt


# Generate Receipt function
def generate_receipt(detected_classes):
    detected_items = []
    counts = {}     # Dictionary to store bread type counts
 
    for item_class in detected_classes:                         # Counts the quantity of each class
        counts[item_class] = counts.get(item_class, 0) + 1      

    for item_class, count in counts.items():                    # Gets the name and price of each class
        bread_info = bread_types.get(item_class, {})
        item_name = bread_info.get("name", "Unknown Bread")
        price = bread_info.get("price", 0)
        detected_items.append({"item": item_name, "quantity": count, "price": price})

    total_amount = sum(item["quantity"] * item["price"] for item in detected_items)

    # Creates the receipt dictionary 
    data = {"Item": [], "Quantity": [], "Price": [], "Amount": []}

    for item_info in detected_items:
        item = item_info["item"]
        quantity = item_info["quantity"]
        price = item_info["price"]
        total_item_amount = quantity * price

        data["Item"].append(item)
        data["Quantity"].append(quantity)
        data["Price"].append(price)
        data["Amount"].append(total_item_amount)

    # Appends the last row of the dataframe for the total amount
    data["Item"].append("TOTAL")
    data["Quantity"].append("")
    data["Price"].append("")
    data["Amount"].append(total_amount)

    df = pd.DataFrame(data)

    return df


# Export to CSV function
def export_csv(df):
    df.to_csv("receipt.csv", index=False)
    return gr.File.update(value="receipt.csv", visible=True)


# Export to JSON function
def export_json(df):
    df.to_json("receipt.json")
    return gr.File.update(value="receipt.json", visible=True)


# Select image from Files
def preview(files, sd: gr.SelectData):
    prev = files[sd.index].name
    return gr.Image.update(value=prev, visible=True)


# Gradio Interface
with gr.Blocks(theme=theme) as demo:
    
    clear_output()
    gr.Markdown("# Bread Detector w/ POS")
    gr.Markdown("An application that detects different types of bread and calculates the total price.")
    gr.Markdown("**Bread types include:** baguette, binangkal, bonete, cornbread, croissant, ensaymada, flatbread, kalihim, monay, pandesal, sourdough, spanish bread, wheat bread, white bread, and whole grain bread.")

    with gr.Row():
        with gr.Column():
                fn = detect_bread
                img_input = gr.Files(file_types=["image"], label="Input Image")
                img_preview = gr.Image(label="Preview Image", interactive=False, visible=False)
                detect_btn = gr.Button(variant="primary", value="Detect")
        with gr.Column():
            img_output = gr.Gallery(label='Output Image')
            receipt_output = gr.Dataframe(
                    headers=["Item", "Quantity", "Price", "Amount"],
                    datatype=["str", "number", "number", "number"],
                    label="Receipt",
                    interactive=False,
                )
            with gr.Row():
                clear_btn = gr.ClearButton()
                export_csv_btn = gr.Button(variant="primary", value="Export as CSV")
                export_json_btn = gr.Button(variant="primary", value="Export as JSON")
            with gr.Row():
                csv = gr.File(interactive=False, visible=False)           
 

    # Gradio Buttons
    img_input.select(preview, img_input, img_preview)        
    detect_btn.click(detect_bread, inputs=img_input, outputs=[img_output, receipt_output])
    export_csv_btn.click(export_csv, receipt_output, csv)
    export_json_btn.click(export_json, receipt_output, csv)
    clear_btn.click(lambda: [None, None, None, gr.File.update(visible=False), gr.Image.update(visible=False)],
        outputs=[img_input, img_output, receipt_output, csv, img_preview]
        )


demo.queue()
demo.launch()