import numpy as np
import pandas as pd
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from io import BytesIO
import json
import cv2
from ultralytics import YOLO

# ======================= МОДЕЛЬ ===================================
model = YOLO("yolov11m_best.pt")

# ================== ЧТЕНИЕ НАЗВАНИЙ И ЦЕН =======================
with open('Fruit_Veggies_Price.json', 'r', encoding='utf-8') as file:
    fruits_data = json.load(file)

# ===================== ДЕТЕКЦИЯ ПЛОДА ============================
def detect_fruit(image):
    # считываем изображение
    # предполагается, что image - это объект PIL Image
    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    
    # детекция
    detections = model.predict(source=image_cv, conf=0.5)
    
    # проверка на наличие детекций
    if len(detections[0].boxes.cls) == 0:
        return image_cv, None

    result_np_image = detections[0].plot()
    result_np_image = cv2.cvtColor(result_np_image, cv2.COLOR_BGR2RGB)
    detected_fruit = model.names[int(detections[0].boxes.cls[0])]

    return result_np_image, detected_fruit

# =========================== ЧЕК ================================
def create_receipt(detected_fruit, weight):
    data = fruits_data[detected_fruit]
    fruit_name = data['name']
    price = data['price_per_kg']
    total_price = round(price * weight, 2)
    receipt_img = Image.new("RGB", (300, 200), color="white")
    draw = ImageDraw.Draw(receipt_img)

    try:
        font = ImageFont.truetype("DejaVuSans.ttf", 18)
    except IOError:
        font = ImageFont.load_default()

    draw.text((10, 10), "Чек", fill="black", font=font)
    draw.text((10, 50), f"Продукт: {fruit_name}", fill="black", font=font)
    draw.text((10, 80), f"Вес: {weight} кг", fill="black", font=font)
    draw.text((10, 110), f"Цена за кг: {price} руб.", fill="black", font=font)
    draw.text((10, 140), f"Сумма: {total_price} руб.", fill="black", font=font)

    return receipt_img

# ======================= ИНТЕРФЕЙС ===============================
def gradio_interface(image, weight):
    if weight <= 0:
        gr.Info('Укажите вес товара')
        return None, None

    result_np_image, detected_fruit = detect_fruit(image)
    if detected_fruit == None:
        gr.Info('Не удалось определить товар')
        return None, None
        
    receipt = create_receipt(detected_fruit, weight)
    return result_np_image, receipt

image_input = gr.Image(
    label="Изображение",
    width=640,
    height=380
)
weight_input = gr.Number(label="Вес (кг)")
image_output = gr.Image(
    label="Распознанный товар",
    type="numpy",
    width=640,
    height=380
)
receipt_output = gr.Image(
    label="Чек",
    type="pil",
    width=400,
    height=400
)

gr.Interface(
    fn=gradio_interface,
    inputs=[image_input, weight_input],
    outputs=[image_output, receipt_output],
    title="Определение товара и создание чека",
    description="Загрузите изображение, введите вес и получите чек"
).launch()