# app.py (CPU-friendly, with docstrings)
# Gradio-приложение: мультимодальный инференс цены и калорий по фото подноса.
"""
Приложение Gradio для Hugging Face Spaces (CPU):
— Принимает фото подноса (вид сверху).
— Мультимодальная LLM Qwen2-VL-2B-Instruct в один прогон выдаёт JSON с {dish_id, qty}.
— Приложение маппит dish_id в каталог и считает итоговую стоимость и калории.

Ожидаемая структура репозитория:
    app.py
    requirements.txt
    config.yaml                # seed и опционально inference.categories_order
    data/
      catalog.csv              # столбцы: dish_id,name_en,category,price,calories

Особенности CPU-режима:
    - Принудительно device_map="cpu" и torch_dtype=torch.float32
    - low_cpu_mem_usage=True при загрузке модели
    - Ограничение потоков (NUM_THREADS, по умолчанию 2)
    - generate(..., max_new_tokens=128)
    - Gradio queue с concurrency_count=1
"""

from __future__ import annotations

import os
import re
import json
from pathlib import Path
from typing import List, Dict, Optional

# Тише и экономнее
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("TQDM_DISABLE", "1")
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")

import gradio as gr
import pandas as pd
import torch
import yaml
from PIL import Image
from transformers import AutoProcessor, AutoModelForImageTextToText
from transformers.utils import logging as tlogging

tlogging.set_verbosity_error()

MODEL_ID = "Qwen/Qwen2-VL-2B-Instruct"
DEFAULT_ORDER = ["soups", "drinks", "sides", "mains", "salads", "vegetables", "fruits"]

# ---------- загрузка конфигов/артефактов ----------

def project_root() -> Path:
    """
    Возвращает корневую папку проекта (где лежит app.py).

    Returns:
        Path: абсолютный путь к директории с текущим файлом.
    """
    return Path(__file__).resolve().parent


def load_config(pdir: Path) -> dict:
    """
    Загружает конфигурацию из config.yaml (если есть), иначе возвращает дефолт.

    Args:
        pdir (Path): корневая директория проекта.

    Returns:
        dict: конфиг со структурой:
              {
                "seed": Optional[int],
                "inference": {"categories_order": List[str]}
              }
    """
    cfg_path = pdir / "config.yaml"
    if not cfg_path.exists():
        return {"seed": None, "inference": {"categories_order": DEFAULT_ORDER}}
    with cfg_path.open("r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def load_catalog(pdir: Path) -> pd.DataFrame:
    """
    Загружает каталог блюд с dish_id и характеристиками.

    Требуемые колонки: dish_id,name_en,category,price,calories.

    Args:
        pdir (Path): корневая директория проекта.

    Returns:
        pd.DataFrame: датафрейм каталога с dish_id типа int.

    Raises:
        FileNotFoundError: если data/catalog.csv отсутствует.
        ValueError: если не хватает необходимых колонок.
    """
    path = pdir / "data" / "catalog.csv"
    if not path.exists():
        raise FileNotFoundError(
            "Нет data/catalog.csv. Нужны колонки: dish_id,name_en,category,price,calories"
        )
    df = pd.read_csv(path)
    req = {"dish_id", "name_en", "category", "price", "calories"}
    missing = req - set(df.columns)
    if missing:
        raise ValueError(f"В catalog.csv отсутствуют колонки: {missing}")
    df["dish_id"] = df["dish_id"].astype(int)
    return df


def allowed_by_category(catalog_df: pd.DataFrame, order: Optional[List[str]]) -> Dict[str, List[Dict]]:
    """
    Строит карту {категория: [{dish_id, name_en}, ...]} в заданном порядке категорий.

    Args:
        catalog_df (pd.DataFrame): каталог блюд.
        order (Optional[List[str]]): требуемый порядок категорий; если None — порядок по csv.

    Returns:
        Dict[str, List[Dict]]: карта разрешённых блюд по категориям.
    """
    if order is None:
        order = list(dict.fromkeys(catalog_df["category"].tolist()))
    out: Dict[str, List[Dict]] = {}
    for cat in order:
        sub = catalog_df[catalog_df["category"] == cat]
        out[cat] = [{"dish_id": int(r["dish_id"]), "name_en": str(r["name_en"])} for _, r in sub.iterrows()]
    return out

# ---------- промпт и парсинг ----------

def make_system_prompt_idscheme(allowed_map: Dict[str, List[Dict]], order: List[str]) -> str:
    """
    Формирует строгий system-промпт.
    Требование: если на изображении НЕТ подноса/еды с вида сверху — вернуть [] и НИЧЕГО БОЛЬШЕ.
    Список ниже — это БЕЛЫЙ СПИСОК (whitelist), а НЕ подсказка к угадыванию.
    """
    lines = ["Allowed items by category with dish_id (use ONLY these dish_id values in JSON):"]
    for cat in order:
        items = allowed_map.get(cat, [])
        lines.append(f"{cat}:")
        for it in items:
            lines.append(f"- {it['dish_id']}: {it['name_en']}")
    allowed_block = "\n".join(lines)

    return (
        "You are given a single image.\n"
        "Task:\n"
        "1) First, decide whether the image shows a TOP-DOWN cafeteria TRAY that fills most of the frame and has FOOD items on it.\n"
        "2) If NOT (e.g., car, street, landscape, documents, people, animals, kitchen table without tray, random objects), "
        "   return ONLY an empty JSON array: []\n"
        "3) If YES, identify ONLY the items that are CLEARLY VISIBLE on the tray and are present in the ALLOWED LIST below. "
        "   Count integer quantities.\n\n"
        "OUTPUT: Return ONLY a valid JSON array on a single line (no prose, no markdown), "
        'exactly like: [{"dish_id": <integer>, "qty": <integer>}, ...]\n'
        "RULES:\n"
        "- The list below is a WHITELIST, not a hint. Do NOT guess from it.\n"
        "- If unsure about an item, OMIT it.\n"
        "- NEVER enumerate the whole catalog.\n"
        "- If there are no valid items, return [].\n"
        "- Return at most 12 objects.\n\n"
        f"{allowed_block}\n"
    )


def extract_json_array(text: str) -> list:
    """
    Пытается извлечь последний валидный JSON-массив из текста.

    Args:
        text (str): произвольный текст ответа модели.

    Returns:
        list: разобранный массив JSON или пустой список.
    """
    candidates = list(re.finditer(r"\[.*?\]", text, flags=re.S))
    if candidates:
        segment = candidates[-1].group(0)
        try:
            return json.loads(segment)
        except Exception:
            if "'" in segment and '"' not in segment:
                try:
                    return json.loads(segment.replace("'", '"'))
                except Exception:
                    pass
    return []


def extract_id_qty_objects(text: str) -> list:
    """
    Извлекает последовательность объектов {dish_id, qty} даже из «оборванного» ответа.

    Ищет все JSON-похожие подстроки вида { ... }, парсит их и достаёт пары (dish_id, qty),
    сохраняя порядок появления, фильтруя невалидные и дубликаты по dish_id.

    Args:
        text (str): сырой текст ответа модели.

    Returns:
        list: список словарей [{"dish_id": int, "qty": int}, ...] (до 32 элементов).
    """
    objs = []
    for m in re.finditer(r"\{[^{}]*\}", text, flags=re.S):
        segment = m.group(0)
        try:
            seg = segment.replace("'", '"')
            obj = json.loads(seg)
        except Exception:
            continue
        if isinstance(obj, dict) and "dish_id" in obj and "qty" in obj:
            try:
                did = int(obj["dish_id"])
                qty = int(obj["qty"])
            except Exception:
                continue
            objs.append({"dish_id": did, "qty": qty})

    # дедуп и фильтр
    seen = set()
    out = []
    for o in objs:
        if o["qty"] <= 0:
            continue
        if o["dish_id"] in seen:
            continue
        seen.add(o["dish_id"])
        out.append(o)
        if len(out) >= 32:
            break
    return out

# ---------- модель и данные (глобально, один раз) ----------

_processor = None
_model = None
_catalog_df = None
_order = None
_id2name = None
_valid_ids = None

def init_everything() -> None:
    """
    Инициализирует конфиг, каталожные данные и загружает модель на CPU.

    Действия:
        - Читает config.yaml (seed, порядок категорий).
        - Ограничивает численность потоков CPU (NUM_THREADS или 2).
        - Загружает catalog.csv и подготавливает маппинги id->name.
        - Загружает AutoProcessor и AutoModelForVision2Seq для Qwen2-VL на CPU.
    """
    global _processor, _model, _catalog_df, _order, _id2name, _valid_ids
    pdir = project_root()
    cfg = load_config(pdir)

    # Ограничение потоков для CPU Space
    num_threads = int(os.getenv("NUM_THREADS", "2"))
    torch.set_num_threads(num_threads)
    os.environ["OMP_NUM_THREADS"] = str(num_threads)

    # Seed (только CPU-пути)
    seed = cfg.get("seed", None)
    if seed is not None:
        import random
        import numpy as np
        os.environ["PYTHONHASHSEED"] = str(seed)
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)

    # Каталог и маппинги
    _catalog_df = load_catalog(pdir)
    _order = cfg.get("inference", {}).get("categories_order", DEFAULT_ORDER)
    _id2name = {int(r["dish_id"]): str(r["name_en"]) for _, r in _catalog_df.iterrows()}
    _valid_ids = set(_id2name.keys())

    # Модель на CPU (float32) с экономией памяти
    _processor = AutoProcessor.from_pretrained(MODEL_ID)
    _model = AutoModelForImageTextToText.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float32,
        device_map="cpu",
        low_cpu_mem_usage=True,
    )


def _decode_generated_only(output_ids, input_ids) -> str:
    """
    Декодирует только сгенерированную часть последовательности, без эхирования промпта.

    Args:
        output_ids (torch.Tensor): тензор токенов выхода generate().
        input_ids (torch.Tensor): тензор токенов входа (для определения длины префикса).

    Returns:
        str: текст сгенерированного продолжения.
    """
    gen_only = output_ids[:, input_ids.shape[1]:]
    return _processor.batch_decode(gen_only, skip_special_tokens=True)[0].strip()

# ---------- основной инференс ----------

def infer_tray(image: Image.Image):
    """
    Основная функция инференса для Gradio.

    Принимает изображение, формирует промпт-справочник с dish_id,
    вызывает Qwen2-VL-2B-Instruct, разбирает ответ (устойчиво к «оборванным» массивам),
    валидирует dish_id, агрегирует qty и считает итог (цена/калории).

    Args:
        image (PIL.Image.Image): загруженное пользователем изображение подноса (вид сверху).

    Returns:
        Tuple[pd.DataFrame, str, str]:
            - таблица с позициями (name_en, category, qty, price, calories, line_price, line_calories)
            - текст с итогами (цена и калории)
            - JSON-строка с полным ответом
    """
    if image is None:
        return pd.DataFrame(), "Загрузите изображение.", "{}"

    # Промпт-справочник
    allowed_map = allowed_by_category(_catalog_df, _order)
    sys_prompt = make_system_prompt_idscheme(allowed_map, _order)
    messages = [
        {"role": "system", "content": [{"type": "text", "text": sys_prompt}]},
        {"role": "user", "content": [
            {"type": "image"},
            {"type": "text", "text":
                "If the image does NOT contain a top-down cafeteria tray with food items from the allowed list, return [] and nothing else. "
                "Otherwise, return ONLY the JSON array with objects {\"dish_id\": <int>, \"qty\": <int>}."
            }
        ]},
    ]
    chat = _processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
    inputs = _processor(text=chat, images=[image.convert("RGB")], return_tensors="pt")  # CPU

    # Один прогон модели
    with torch.no_grad():
        output = _model.generate(**inputs, max_new_tokens=128, do_sample=False)

    # Разбор ответа (валидный JSON или «оборванный»)
    raw = _decode_generated_only(output, inputs["input_ids"])
    data = extract_json_array(raw)
    if not isinstance(data, list) or not data:
        data = extract_id_qty_objects(raw)
        if not data:
            empty = {"items": [], "total_price": 0.0, "total_calories": 0.0}
            return pd.DataFrame(), "Не удалось распознать блюда на подносе.", json.dumps(empty, ensure_ascii=False)

    # Агрегация по валидным dish_id
    agg: Dict[int, int] = {}
    for it in data:
        try:
            did = int(it.get("dish_id", -1))
            qty = int(it.get("qty", 1))
        except Exception:
            continue
        if did not in _valid_ids or qty <= 0:
            continue
        agg[did] = agg.get(did, 0) + qty

    if not agg:
        empty = {"items": [], "total_price": 0.0, "total_calories": 0.0}
        return pd.DataFrame(), "Не найдено валидных блюд.", json.dumps(empty, ensure_ascii=False)

    # Подсчёт цены/ккал
    items = [{"name_en": _id2name[did], "qty": qty} for did, qty in sorted(agg.items())]
    merged = pd.DataFrame(items).merge(_catalog_df, on="name_en", how="left")
    merged["line_price"] = merged["qty"] * merged["price"]
    merged["line_calories"] = merged["qty"] * merged["calories"]
    total_price = float(merged["line_price"].sum())
    total_calories = float(merged["line_calories"].sum())

    # Выводы
    table = merged[["name_en", "category", "qty", "price", "calories", "line_price", "line_calories"]]
    totals = f"ИТОГО: {total_price:.2f} ₽ | {int(total_calories)} ккал"

    json_out = {
        "items": table.to_dict("records"),
        "total_price": total_price,
        "total_calories": total_calories
    }
    return table, totals, json.dumps(json_out, ensure_ascii=False, indent=2)

# ---------- UI ----------

with gr.Blocks(theme=gr.themes.Default(), title="Tray Price & Calories (Qwen2-VL, CPU)") as demo:
    """
    Конструктор Gradio UI:
        - левая колонка: загрузка изображения и кнопка "Рассчитать"
        - правая колонка: таблица распознанных позиций, итоги и JSON-вывод
    """
    gr.Markdown("## Подсчёт цены и калорий по фото подноса  \nМодель: **Qwen2-VL-2B-Instruct** · CPU · один проход · закрытый словарь с `dish_id`")
    with gr.Row():
        with gr.Column(scale=1):
            img = gr.Image(type="pil", label="Загрузите фото подноса (вид сверху)")
            btn = gr.Button("Рассчитать")
        with gr.Column(scale=2):
            df = gr.Dataframe(
                headers=["name_en", "category", "qty", "price", "calories", "line_price", "line_calories"],
                label="Распознанные позиции",
                interactive=False
            )
            totals = gr.Markdown()
            jtxt = gr.Code(label="JSON", language="json")

    btn.click(fn=infer_tray, inputs=img, outputs=[df, totals, jtxt])

# Инициализация при старте
init_everything()

if __name__ == "__main__":
    # Очередь и 1 параллельный запрос — дружелюбно к CPU Space
    # demo.queue(concurrency_count=1, max_size=8).launch()
    demo.queue(max_size=8).launch()