import gradio as gr
import pandas as pd
from ultralytics import YOLO
from datetime import datetime


categories = {
    0: "Beans",
    1: "Butter",
    2: "Cheese",
    3: "Drink",
    4: "Jogurt",
    5: "Milk",
    6: "Pasta",
    7: "Rice",
    8: "Snacks",
    9: "Tuna"
}

# inventory_separate_entries = pd.DataFrame(columns=["Product", "Time added"])
inventory_separate_entries = pd.read_excel("file.xlsx", sheet_name='History of entries')
print(inventory_separate_entries)


model = YOLO("../object_detection/runs/detect/train7/weights/best.pt")


def get_value_counts_df(inventory_separate_entries):
    value_counts = inventory_separate_entries['Product'].value_counts(dropna=True, sort=True)
    df_val_counts = pd.DataFrame(value_counts)
    df_value_counts_reset = df_val_counts.reset_index()
    df_value_counts_reset.columns = ['Product', 'Quantity']
    return df_value_counts_reset


def add_item(item_name):
    inventory_separate_entries.loc[-1] = [item_name, datetime.today().strftime('%Y-%m-%d %H:%M:%S')]
    inventory_separate_entries.index = inventory_separate_entries.index + 1  # shifting index
    inventory_separate_entries.sort_index(inplace=True)


def detect_objects(input_img):
    result_img = model([input_img])
    print("result", result_img[0].boxes)
    detected_categories = []
    for r in result_img[0].boxes:
        for element in r.data:
            detected_category = int(element[-1])
            confidence = element[-2]
            print(r.data)
            detected_categories.append((detected_category, confidence))
    print(detected_categories)
    sorted_detected_categories = sorted(detected_categories, key=lambda x: x[1], reverse=True)
    if len(sorted_detected_categories) > 0:
        detected_category = categories[sorted_detected_categories[0][0]]
    else:
        detected_category = "Nothing detected"

    label = gr.Label(detected_category, visible=True)

    result_img = result_img[0].plot()
    return gr.Image(result_img, visible=True), label#, add_item_btn # detection_interface


def refresh_inventory():
    value_counts_df = get_value_counts_df(inventory_separate_entries)

    with pd.ExcelWriter("file.xlsx") as writer:
        value_counts_df.to_excel(writer, sheet_name='Total inventory', index=False)
        inventory_separate_entries.to_excel(writer, sheet_name='History of entries', index=False)

    return value_counts_df, inventory_separate_entries


with gr.Blocks() as inventory_interface:
    with gr.Row():
        refresh_btn = gr.Button("Refresh Inventory")
        export_btn = gr.DownloadButton("Export Inventory", value="/file=file.xlsx")

    value_counts_df = get_value_counts_df(inventory_separate_entries)
    inventory_title = gr.Label("Inventory Summary", color="grey")
    inventory_table = gr.DataFrame(value_counts_df)

    history_title = gr.Label("History of added items", color="grey")
    history_table = gr.DataFrame(inventory_separate_entries)

    output_file = gr.File(visible=False)
    refresh_btn.click(refresh_inventory, outputs=[inventory_table, history_table])


with gr.Blocks() as scan_interface:
    with gr.Row(equal_height=False):
        input = gr.Image()
        greet_btn = gr.Button("Detect items")

    with gr.Row() as row:
        output = gr.Image()
        with gr.Column():
            label = gr.Label("Nothing detected", visible=True)
            add_item_btn = gr.Button("Add item", visible=True)

    add_item_btn.click(add_item, inputs=label)
    greet_btn.click(detect_objects, inputs=input, outputs=[output, label])


demo = gr.TabbedInterface([scan_interface, inventory_interface], ["Scan items", "Inventory"])

demo.launch(share=True, allowed_paths=["file.xlsx"])