import gradio as gr
import numpy as np
import pandas as pd
from app_utils import annotate_planogram_compliance, do_sorting, xml_to_csv
from inference import run
import json
import os
from tempfile import NamedTemporaryFile


# Target names list
target_names = [
    "Bottle,100PLUS ACTIVE 1.5L",
    "Bottle,100PLUS ACTIVE 500ML",
    "Bottle,100PLUS LEMON LIME 1.5L",
    # Add all other target names here
]

# Define the function to run planogram compliance check
def planogram_compliance_check(planogram_image, master_planogram_image, annotation_file):
    # Convert uploaded images to numpy arrays
    planogram_img = np.array(planogram_image)
    master_planogram_img = np.array(master_planogram_image)

    # Perform inference on planogram image
    result_list = run(
        weights="base_line_best_model_exp5.pt",
        source=planogram_img,
        imgsz=[640, 640],
        conf_thres=0.6,
        iou_thres=0.6,
    )

    # Load annotation file and convert to DataFrame
    if annotation_file is not None:
        annotation_df = xml_to_csv(annotation_file)
        sorted_xml_df = do_sorting(annotation_df)
    else:
        sorted_xml_df = None

    # Run planogram compliance check
    compliance_score, annotated_image = run_compliance_check(
        planogram_img, master_planogram_img, sorted_xml_df, result_list
    )

    return compliance_score, annotated_image

def run_compliance_check(planogram_img, master_planogram_img, sorted_xml_df, result_list):
    # Placeholder for actual score calculation
    compliance_score = 0.0

    # Placeholder for annotated image
    annotated_image = planogram_img.copy()

    if sorted_xml_df is not None:
        annotate_df = sorted_xml_df[["xmin", "ymin", "xmax", "ymax", "line_number", "cls"]].astype(int)
    else:
        annotate_df = None

    mask = master_table != non_null_product
    m_detected_table = np.ma.masked_array(master_table, mask=mask)
    m_annotated_table = np.ma.masked_array(detected_table, mask=mask)

    # wrong_indexes = np.ravel_multi_index(master_table*mask != detected_table*mask, master_table.shape)
    wrong_indexes = np.where(master_table != detected_table)
    correct_indexes = np.where(master_table == detected_table)

    # Annotate planogram compliance on the image
    annotated_image = annotate_planogram_compliance(
        annotated_image, annotate_df, correct_indexes, wrong_indexes, target_names
    )

    # Calculate compliance score
    correct_matches = (np.ma.masked_equal(master_table, non_null_product) == detected_table).sum()
    total_products = (master_table != non_null_product).sum()
    if total_products != 0:
        compliance_score = correct_matches / total_products

    return compliance_score, annotated_image


# Gradio interface
planogram_check_interface = gr.Interface(
    fn=planogram_compliance_check,
    inputs=[
        gr.inputs.Image(label="Planogram Image"),
        gr.inputs.Image(label="Master Planogram Image"),
        gr.inputs.Dataframe(label="Annotation File (XML)")
    ],
    outputs=[
        gr.outputs.Textbox(label="Compliance Score"),
        gr.outputs.Image(label="Annotated Planogram Image"),
    ],
    title="Planogram Compliance Checker",
    description="Upload planogram image, master planogram image, and annotation file (if available) to check compliance."
)


planogram_check_interface.launch()