import numpy as np
import pandas as pd
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from io import BytesIO
import json
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)


# таблица с данными
detections = []


# =========================== ДЕТЕКЦИЯ ПЛОДА ============================

def detect_image(image_path: str, weight: float):
    # Считываем изображение
    image = cv2.imread(image_path)

    # Выполняем детекцию
    detections = model.predict(source=image_path, conf=0.5)
    result_np_image = detections[0].plot()
    result_np_image = cv2.cvtColor(result_np_image, cv2.COLOR_BGR2RGB)

    detected_fruit = None
    
    for det in detections:
        label = model.names[int(det[5])]  # название фрукта
        if label in fruits_data:
            detected_fruit = label
            break

    if not detected_fruit:
        return result_np_image, "Фрукт не обнаружен", None, None
    
    fruit_name = fruits_data[detected_fruit]["name"]
    price_per_kg = fruits_data[detected_fruit]["price_per_kg"]
    total_price = round(price_per_kg * weight, 2)
    
    return result_np_image, fruit_name, weight, total_price


# =========================== ЧЕК ============================

# функция для добавления позиции в чек
def create_receipt(fruit_name, weight, total_price):
    # белый фон для чека
    receipt_img = Image.new("RGB", (300, 200), color="white")
    draw = ImageDraw.Draw(receipt_img)
    
    # шрифт
    try:
        font = ImageFont.truetype("arial.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"Цена за кг: {fruits_data[detected_fruit]['price_per_kg']} руб.", fill="black", font=font)
    draw.text((10, 140), f"Сумма: {total_price} руб.", fill="black", font=font)
    
    return receipt_img


# ======================= ИНТЕРФЕЙС ============================

def gradio_interface(image, weight):
    image, fruit_name, weight, total_price = detect_fruit(image, weight)
    if not fruit_name:
        return image, None  # Вернуть пустой чек
    
    receipt = create_receipt(fruit_name, weight, total_price)
    return image, receipt

image_input = gr.Image(label="Изображение")
weight_input = gr.Number(label="Вес (кг)")
image_output = gr.Image(label="Распознанный фрукт")
receipt_output = gr.Image(label="Чек")

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