app.py

import os
os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"

from pathlib import Path
from collections import Counter
from datetime import datetime
import base64
import io
import math
import time

import cv2
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from PIL import Image, ImageOps

import gradio as gr
from ultralytics import YOLO


# ============================================================
# PATHS
# ============================================================
# Nouveau modèle empty-space : dans le même dossier que app.py
EMPTY_WEIGHTS = Path("best.pt")

# Ancien modèle produit : conservé uniquement pour détection produit/KPI
YOLO_WEIGHTS = Path("product_detector.pt")

# CLIP/facing : conservé mais désactivé par défaut
LN_CKPT_PATH = Path("facing_model.pth")
LN_DB_PATH = Path("ln_database.pth")


# ============================================================
# FAST CONFIG FOR HUGGING FACE CPU BASIC
# ============================================================
EMPTY_CONF = 0.25
YOLO_CONF = 0.25

# 768 = plus rapide sur CPU Hugging Face
MAX_IMAGE_SIDE = 768

MAX_PRODUCTS_ANALYZED = 80
MAX_PRODUCTS_FOR_CLIP = 12

DEFAULT_ENABLE_FACING_CLIP = False
DEFAULT_ENABLE_GROUPING_CLIP = False
DEFAULT_ENABLE_PRICE_MATCHING = False

# "errors_only" : trace seulement espaces vides + produits back si CLIP activé
# "all" : trace aussi tous les produits
ANNOTATION_MODE = "errors_only"

MIN_CROP_PX = 10
SIM_THRESHOLD = 0.65
EMPTY_SPACE_PONDERATION = 1.5

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

try:
    torch.set_num_threads(max(1, min(4, os.cpu_count() or 1)))
except Exception:
    pass


# ============================================================
# SAFE LOAD
# ============================================================
def safe_torch_load(path, map_location):
    try:
        return torch.load(path, map_location=map_location, weights_only=False)
    except TypeError:
        return torch.load(path, map_location=map_location)


# ============================================================
# LOAD MODELS
# ============================================================
if not EMPTY_WEIGHTS.exists():
    raise FileNotFoundError(
        f"best.pt introuvable : {EMPTY_WEIGHTS.resolve()}. "
        "Ajoute best.pt dans le même dossier que app.py."
    )

print("Loading empty-space model best.pt...")
empty_model = YOLO(str(EMPTY_WEIGHTS))

if YOLO_WEIGHTS.exists():
    print("Loading product detector...")
    yolo_model = YOLO(str(YOLO_WEIGHTS))
else:
    print("product_detector.pt not found. Product detection disabled.")
    yolo_model = None

try:
    empty_model.fuse()
except Exception:
    pass

if yolo_model is not None:
    try:
        yolo_model.fuse()
    except Exception:
        pass


# ============================================================
# LAZY CLIP GLOBALS
# ============================================================
_clip_loaded = False
_clip_model = None
_clip_preprocess = None
_tokenizer = None
_LABELS = None
_TEXT_EMB = None

_db_loaded = False
_db_available = False
_db_emb = None
_meta = []
_db_names = []
_db_fams = []
_pkey = None


# ============================================================
# HELPERS
# ============================================================
def resize_for_inference(pil_img, max_side=768):
    w, h = pil_img.size
    scale = max_side / max(w, h)

    if scale >= 1.0:
        return pil_img, 1.0

    new_w = int(w * scale)
    new_h = int(h * scale)

    resized = pil_img.resize((new_w, new_h))
    return resized, scale


def polygon_area(poly):
    x = poly[:, 0]
    y = poly[:, 1]

    return 0.5 * abs(
        np.dot(x, np.roll(y, -1)) -
        np.dot(y, np.roll(x, -1))
    )


def pad_to_square_avg(img: Image.Image) -> Image.Image:
    w, h = img.size

    if w == h:
        return img

    arr = np.asarray(img)
    avg = tuple(int(c) for c in arr.reshape(-1, arr.shape[-1]).mean(axis=0))

    side = max(w, h)
    canvas = Image.new("RGB", (side, side), avg)
    canvas.paste(img, ((side - w) // 2, (side - h) // 2))

    return canvas


def pil_to_base64(img):
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img)

    buffer = io.BytesIO()
    img.save(buffer, format="JPEG", quality=80)

    encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return f"data:image/jpeg;base64,{encoded}"


def sanitize_for_json(obj):
    if isinstance(obj, dict):
        return {k: sanitize_for_json(v) for k, v in obj.items()}

    if isinstance(obj, (list, tuple)):
        return [sanitize_for_json(v) for v in obj]

    if isinstance(obj, np.integer):
        return int(obj)

    if isinstance(obj, np.floating):
        f = float(obj)
        return None if math.isnan(f) or math.isinf(f) else f

    if isinstance(obj, np.ndarray):
        return sanitize_for_json(obj.tolist())

    if isinstance(obj, float):
        return None if math.isnan(obj) or math.isinf(obj) else obj

    return obj


def load_trainable_state_dict(model, trainable_sd, device):
    full_sd = model.state_dict()
    loaded = 0
    ignored = []

    for k, v in trainable_sd.items():
        if k in full_sd:
            full_sd[k] = v.to(device)
            loaded += 1
        else:
            ignored.append(k)

    model.load_state_dict(full_sd)
    return loaded, ignored


# ============================================================
# FACING PROMPTS
# ============================================================
PROMPTS = {
    "front": [
        "a product crop showing a large brand logo and product name",
        "packaging with bold colorful artwork and a readable brand name",
        "the marketing label of a product with a hero photo of the contents",
        "a crop of branded packaging with large flavor name and product imagery",
        "saturated colorful product packaging with a prominent logo",
        "a product label with a big photograph of food or the product itself",
        "consumer packaging artwork dominated by brand name and visuals",
        "the branded printed surface of a product, glossy and graphic-heavy",
    ],
    "back": [
        "a nutrition facts table with rows of small black text",
        "an ingredients list in tiny dense print",
        "a barcode printed on packaging",
        "the back of a package, mostly plain with small text blocks",
        "a side panel with stacked multilingual text",
        "the top of a container showing a lid or cap",
        "the bottom of a container with small printed codes",
        "a plain white or single-color panel with regulatory text",
        "a recycling symbol and disposal instructions on packaging",
        "the unbranded side of a product with no logo or hero image",
    ],
}


@torch.no_grad()
def build_text_embeddings(prompts_dict):
    global _clip_model, _tokenizer

    class_embeds = {}

    for label, prompts in prompts_dict.items():
        tokens = _tokenizer(prompts).to(DEVICE)

        e = _clip_model.encode_text(tokens)
        e = F.normalize(e.float(), dim=-1).mean(dim=0)
        e = F.normalize(e, dim=-1)

        class_embeds[label] = e

    labels = list(class_embeds.keys())
    matrix = torch.stack([class_embeds[l] for l in labels])

    return labels, matrix


def lazy_load_clip():
    global _clip_loaded, _clip_model, _clip_preprocess, _tokenizer, _LABELS, _TEXT_EMB

    if _clip_loaded:
        return

    print("Lazy loading CLIP ViT-B/16...")
    import open_clip

    _clip_model, _, _clip_preprocess = open_clip.create_model_and_transforms(
        "ViT-B-16",
        pretrained="laion2b_s34b_b88k",
    )

    _tokenizer = open_clip.get_tokenizer("ViT-B-16")
    _clip_model = _clip_model.to(DEVICE).eval()

    for p in _clip_model.parameters():
        p.requires_grad = False

    if LN_CKPT_PATH.exists():
        try:
            print("Loading LN-tuned checkpoint...")
            ckpt = safe_torch_load(LN_CKPT_PATH, DEVICE)

            if isinstance(ckpt, dict):
                ln_sd = ckpt.get("ln_state_dict") or ckpt.get("trainable_state_dict") or ckpt
            else:
                ln_sd = ckpt

            loaded_ln, ignored_ln = load_trainable_state_dict(_clip_model, ln_sd, DEVICE)
            print(f"Loaded LN tensors: {loaded_ln}")
            print(f"Ignored LN tensors: {len(ignored_ln)}")
        except Exception as e:
            print(f"Could not load LN checkpoint: {e}")
    else:
        print("facing_model.pth not found. Using base CLIP.")

    _LABELS, _TEXT_EMB = build_text_embeddings(PROMPTS)
    _clip_loaded = True


def find_price_key(metadata):
    if not metadata:
        return None

    possible_keys = [
        "price",
        "prix",
        "price_mad",
        "price_dh",
        "unit_price",
        "selling_price",
    ]

    all_keys = set()

    for m in metadata[:500]:
        if isinstance(m, dict):
            all_keys.update(m.keys())

    for k in possible_keys:
        if k in all_keys:
            return k

    for k in all_keys:
        kl = k.lower()
        if "price" in kl or "prix" in kl:
            return k

    return None


def lazy_load_database():
    global _db_loaded, _db_available, _db_emb, _meta, _db_names, _db_fams, _pkey

    if _db_loaded:
        return

    _db_loaded = True

    if not LN_DB_PATH.exists():
        print("ln_database.pth not found. Price/value computation disabled.")
        return

    try:
        print("Loading ln_database.pth...")
        db = safe_torch_load(LN_DB_PATH, DEVICE)

        _db_emb = F.normalize(db["embeddings"].to(DEVICE), dim=-1)
        _meta = db["metadata"]

        _db_names = [m.get("product_name", "unknown") for m in _meta]
        _db_fams = [m.get("family", "unknown") for m in _meta]
        _pkey = find_price_key(_meta)

        _db_available = True

        print(f"LN database loaded: {_db_emb.shape[0]} entries")
        print(f"Detected pkey: {_pkey}")

    except Exception as e:
        print(f"Could not load ln_database.pth: {e}")
        _db_available = False


def get_price_from_meta(db_idx):
    if _pkey is None:
        return None

    if db_idx is None or db_idx < 0 or db_idx >= len(_meta):
        return None

    value = _meta[db_idx].get(_pkey)

    if value is None:
        return None

    try:
        return float(value)
    except Exception:
        s = str(value).strip().replace(",", ".")
        s = "".join(ch for ch in s if ch.isdigit() or ch == ".")

        try:
            return float(s) if s else None
        except Exception:
            return None


def prepare_clip_images(crops):
    processed = []

    for crop in crops:
        padded = pad_to_square_avg(crop)
        processed.append(_clip_preprocess(padded))

    return torch.stack(processed).to(DEVICE)


@torch.no_grad()
def embed_crops_batch(crops, batch_size=8):
    lazy_load_clip()

    if len(crops) == 0:
        return torch.zeros(0, 512)

    all_embs = []

    for i in range(0, len(crops), batch_size):
        batch = crops[i:i + batch_size]
        x = prepare_clip_images(batch)

        e = _clip_model.encode_image(x)
        e = F.normalize(e.float(), dim=-1)

        all_embs.append(e.cpu())

    return torch.cat(all_embs, dim=0)


@torch.no_grad()
def classify_crops_batch(crop_embs):
    if crop_embs.shape[0] == 0:
        return [], []

    e = crop_embs.to(DEVICE)
    logits = (e @ _TEXT_EMB.T) * _clip_model.logit_scale.exp()
    probs = logits.softmax(dim=-1).cpu().numpy()

    preds = []
    scores = []

    for row in probs:
        idx = int(row.argmax())
        preds.append(_LABELS[idx])
        scores.append({l: float(p) for l, p in zip(_LABELS, row)})

    return preds, scores


# ============================================================
# EMPTY DETECTION USING best.pt
# ============================================================
def detect_empty_spaces_best(image_np, conf, imgsz):
    """
    best.pt is a standard YOLO detection model.
    We read result.boxes.xyxy, exactly like the previous FastAPI code.
    Then we convert each rectangle into a 4-point polygon for area/KPI compatibility.
    """
    results = empty_model.predict(
        image_np,
        conf=conf,
        imgsz=int(imgsz),
        verbose=False,
    )

    result = results[0]
    boxes = result.boxes

    image_h, image_w = image_np.shape[:2]

    empty_polys_list = []
    empty_confs_list = []

    if boxes is not None and boxes.xyxy is not None:
        xyxy = boxes.xyxy.cpu().numpy()
        confs = boxes.conf.cpu().numpy() if boxes.conf is not None else np.ones(len(xyxy))

        for box, c in zip(xyxy, confs):
            x1, y1, x2, y2 = box.tolist()

            x1 = max(0.0, min(float(image_w - 1), float(x1)))
            y1 = max(0.0, min(float(image_h - 1), float(y1)))
            x2 = max(0.0, min(float(image_w), float(x2)))
            y2 = max(0.0, min(float(image_h), float(y2)))

            if x2 <= x1 or y2 <= y1:
                continue

            poly = np.array(
                [
                    [x1, y1],
                    [x2, y1],
                    [x2, y2],
                    [x1, y2],
                ],
                dtype=np.float32,
            )

            empty_polys_list.append(poly)
            empty_confs_list.append(float(c))

    if empty_polys_list:
        empty_polys = np.stack(empty_polys_list, axis=0)
        empty_confs = np.array(empty_confs_list, dtype=np.float32)
    else:
        empty_polys = np.zeros((0, 4, 2), dtype=np.float32)
        empty_confs = np.zeros((0,), dtype=np.float32)

    return empty_polys, empty_confs


# ============================================================
# CLEAN ANNOTATION — STYLE ANCIEN FASTAPI
# ============================================================
def draw_result_image(image_np, empty_polys, empty_confs, boxes, facing_preds, empty_df):
    """
    Tracé des espaces vides comme dans l'ancien FastAPI :
    - rectangle rouge transparent
    - contour rouge
    - label "Vide XX% (YY.Y%)"
    - bannière en haut avec remplissage/vide/nombre de zones
    """
    out = image_np.copy()
    image_h, image_w = image_np.shape[:2]

    # OpenCV sur array RGB ici, donc rouge = (255, 0, 0), vert = (0, 200, 0)
    red = (255, 0, 0)
    white = (255, 255, 255)
    dark = (40, 40, 40)
    gray = (100, 100, 100)
    green = (0, 200, 0)
    blue_red = (200, 0, 0)

    # Boxes espaces vides
    for i, poly in enumerate(empty_polys):
        pts = poly.astype(np.int32)

        x1 = int(pts[:, 0].min())
        y1 = int(pts[:, 1].min())
        x2 = int(pts[:, 0].max())
        y2 = int(pts[:, 1].max())

        conf = float(empty_confs[i]) if i < len(empty_confs) else 0.0

        if empty_df is not None and not empty_df.empty and i < len(empty_df):
            zone_pct = float(empty_df.iloc[i].get("area_image_percent", 0.0))
        else:
            zone_pct = 0.0

        overlay = out.copy()
        cv2.rectangle(overlay, (x1, y1), (x2, y2), red, -1)
        out = cv2.addWeighted(overlay, 0.30, out, 0.70, 0)

        cv2.rectangle(out, (x1, y1), (x2, y2), red, 2)

        label = f"Vide {conf:.0%} ({zone_pct:.1f}%)"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)

        label_y1 = max(0, y1 - th - 10)
        label_y2 = max(th + 8, y1)

        cv2.rectangle(
            out,
            (x1, label_y1),
            (min(image_w - 1, x1 + tw + 5), label_y2),
            red,
            -1,
        )
        cv2.putText(
            out,
            label,
            (x1 + 2, max(th + 2, y1 - 5)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            white,
            2,
            cv2.LINE_AA,
        )

    # Produits : conservés, mais invisibles par défaut sauf back/side avec CLIP
    # Si tu veux voir tous les produits, mets ANNOTATION_MODE = "all".
    for i, box in enumerate(boxes):
        facing = facing_preds[i] if i < len(facing_preds) else "unknown"

        if ANNOTATION_MODE == "errors_only" and facing != "back":
            continue

        x1, y1, x2, y2 = [int(v) for v in box]

        if facing == "back":
            color = (255, 60, 60)
            thickness = 3
        elif facing == "front":
            color = (0, 210, 80)
            thickness = 2
        else:
            color = (255, 140, 0)
            thickness = 2

        cv2.rectangle(out, (x1, y1), (x2, y2), color, thickness)

    # Bannière style ancien FastAPI
    if empty_df is not None and not empty_df.empty:
        pct_empty = float(empty_df["area_image_percent"].sum())
    else:
        pct_empty = 0.0

    pct_empty = max(0.0, min(100.0, pct_empty))
    pct_merchandise = 100.0 - pct_empty
    n_empty = len(empty_polys)

    banner_h = 80
    cv2.rectangle(out, (0, 0), (image_w, banner_h), dark, -1)

    score_text = f"Remplissage: {pct_merchandise:.1f}% | Vide: {pct_empty:.1f}% | {n_empty} zone(s)"
    cv2.putText(
        out,
        score_text,
        (15, 35),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.9,
        white,
        2,
        cv2.LINE_AA,
    )

    # Barre de progression
    bar_x, bar_y, bar_w, bar_h = 15, 50, max(10, image_w - 30), 20
    cv2.rectangle(out, (bar_x, bar_y), (bar_x + bar_w, bar_y + bar_h), gray, -1)

    fill_w = int(bar_w * pct_merchandise / 100.0)
    cv2.rectangle(out, (bar_x, bar_y), (bar_x + fill_w, bar_y + bar_h), green, -1)

    if fill_w < bar_w:
        cv2.rectangle(
            out,
            (bar_x + fill_w, bar_y),
            (bar_x + bar_w, bar_y + bar_h),
            blue_red,
            -1,
        )

    return out


# ============================================================
# CORE ANALYSIS
# ============================================================
def run_analysis(
    image,
    enable_facing_clip=False,
    enable_grouping_clip=False,
    enable_price_matching=False,
    max_image_side=MAX_IMAGE_SIDE,
    max_products_for_clip=MAX_PRODUCTS_FOR_CLIP,
):
    t0 = time.perf_counter()

    enable_facing_clip = bool(enable_facing_clip)
    enable_grouping_clip = bool(enable_grouping_clip)
    enable_price_matching = bool(enable_price_matching)

    if enable_price_matching:
        enable_grouping_clip = True

    original_img = ImageOps.exif_transpose(image).convert("RGB")
    shelf_img, resize_scale = resize_for_inference(original_img, int(max_image_side))

    image_np = np.array(shelf_img)
    image_h, image_w = image_np.shape[:2]
    image_area = image_h * image_w

    # -----------------------------
    # Empty-space detection with best.pt
    # -----------------------------
    t_empty0 = time.perf_counter()

    empty_polys, empty_confs = detect_empty_spaces_best(
        image_np=image_np,
        conf=EMPTY_CONF,
        imgsz=int(max_image_side),
    )

    t_empty = time.perf_counter() - t_empty0

    n_empty = len(empty_polys)
    empty_surface = float(sum(polygon_area(p) for p in empty_polys))

    empty_rows = []

    for i, poly in enumerate(empty_polys):
        area = float(polygon_area(poly))
        xs = poly[:, 0]
        ys = poly[:, 1]

        empty_rows.append({
            "zone_id": i + 1,
            "confidence": float(empty_confs[i]),
            "area_px2": area,
            "area_image_percent": (area / image_area) * 100 if image_area > 0 else 0.0,
            "x_min": float(xs.min()),
            "y_min": float(ys.min()),
            "x_max": float(xs.max()),
            "y_max": float(ys.max()),
            "width_px": float(xs.max() - xs.min()),
            "height_px": float(ys.max() - ys.min()),
            "center_x": float(xs.mean()),
            "center_y": float(ys.mean()),
        })

    empty_df = pd.DataFrame(empty_rows)

    if not empty_df.empty:
        empty_df = empty_df.sort_values(by="area_px2", ascending=False).reset_index(drop=True)
        empty_df["zone_id"] = np.arange(1, len(empty_df) + 1)
        empty_df["share_of_empty_surface_percent"] = (
            empty_df["area_px2"] / max(empty_surface, 1e-9)
        ) * 100

        # Rebuild polygons sorted same as DataFrame
        sorted_polys = []
        sorted_confs = []
        for _, row in empty_df.iterrows():
            x1, y1, x2, y2 = row["x_min"], row["y_min"], row["x_max"], row["y_max"]
            sorted_polys.append(np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]], dtype=np.float32))
            sorted_confs.append(float(row["confidence"]))
        empty_polys = np.stack(sorted_polys, axis=0) if sorted_polys else np.zeros((0, 4, 2), dtype=np.float32)
        empty_confs = np.array(sorted_confs, dtype=np.float32)

    # -----------------------------
    # Product detection
    # -----------------------------
    t_yolo0 = time.perf_counter()

    if yolo_model is not None:
        yolo_out = yolo_model.predict(
            image_np,
            conf=YOLO_CONF,
            imgsz=int(max_image_side),
            max_det=int(MAX_PRODUCTS_ANALYZED),
            verbose=False,
        )[0]

        if yolo_out.boxes is not None and yolo_out.boxes.xyxy is not None:
            all_boxes = yolo_out.boxes.xyxy.cpu().numpy().astype(int)
            all_confs = yolo_out.boxes.conf.cpu().numpy()
        else:
            all_boxes = np.zeros((0, 4), dtype=int)
            all_confs = np.zeros((0,))
    else:
        all_boxes = np.zeros((0, 4), dtype=int)
        all_confs = np.zeros((0,))

    t_yolo = time.perf_counter() - t_yolo0

    raw_products_detected = len(all_boxes)

    if len(all_boxes) > int(MAX_PRODUCTS_ANALYZED):
        order = np.argsort(-all_confs)[:int(MAX_PRODUCTS_ANALYZED)]
        all_boxes = all_boxes[order]
        all_confs = all_confs[order]

    boxes = []
    crops = []
    det_confs = []

    for box, dc in zip(all_boxes, all_confs):
        x1, y1, x2, y2 = box.tolist()

        x1 = max(0, min(image_w - 1, x1))
        y1 = max(0, min(image_h - 1, y1))
        x2 = max(0, min(image_w, x2))
        y2 = max(0, min(image_h, y2))

        if x2 <= x1 or y2 <= y1:
            continue

        crop = shelf_img.crop((x1, y1, x2, y2))

        if min(crop.size) < MIN_CROP_PX:
            continue

        boxes.append([x1, y1, x2, y2])
        crops.append(crop)
        det_confs.append(float(dc))

    boxes = np.array(boxes, dtype=float)
    n_products_detected_after_filter = len(boxes)

    product_surface = float(
        sum((b[2] - b[0]) * (b[3] - b[1]) for b in boxes)
    )

    total_detected_surface = empty_surface + product_surface

    empty_ratio = (
        empty_surface / total_detected_surface
        if total_detected_surface > 0
        else 0.0
    )

    product_ratio = (
        product_surface / total_detected_surface
        if total_detected_surface > 0
        else 0.0
    )

    empty_image_ratio = empty_surface / image_area if image_area > 0 else 0.0
    product_image_ratio = product_surface / image_area if image_area > 0 else 0.0

    # -----------------------------
    # Optional CLIP
    # -----------------------------
    clip_used = enable_facing_clip or enable_grouping_clip or enable_price_matching

    clip_indices = []
    crop_embs = torch.zeros(0, 512)
    t_clip = 0.0

    if clip_used and n_products_detected_after_filter > 0:
        order = np.argsort(-np.array(det_confs))[:int(max_products_for_clip)]
        clip_indices = order.astype(int).tolist()
        clip_crops = [crops[i] for i in clip_indices]

        t_clip0 = time.perf_counter()
        crop_embs = embed_crops_batch(clip_crops, batch_size=8)
        t_clip = time.perf_counter() - t_clip0

    facing_preds = ["unknown"] * n_products_detected_after_filter
    facing_scores = [
        {"front": None, "back": None}
        for _ in range(n_products_detected_after_filter)
    ]

    if enable_facing_clip and len(clip_indices) > 0:
        clip_facing_preds, clip_facing_scores = classify_crops_batch(crop_embs)

        for local_i, global_i in enumerate(clip_indices):
            facing_preds[global_i] = clip_facing_preds[local_i]
            facing_scores[global_i] = clip_facing_scores[local_i]

    n_front = sum(1 for p in facing_preds if p == "front")
    n_back = sum(1 for p in facing_preds if p == "back")
    back_ratio = n_back / n_products_detected_after_filter if n_products_detected_after_filter > 0 else 0.0

    # -----------------------------
    # Optional grouping
    # -----------------------------
    group_ids = np.array([-1] * n_products_detected_after_filter)
    n_visual_clusters = 0
    n_groups = 0
    counts = Counter()

    if enable_grouping_clip and len(clip_indices) > 0:
        if len(clip_indices) == 1:
            group_ids[clip_indices[0]] = 0
            n_visual_clusters = 1
            n_groups = 1
            counts = Counter([0])
        else:
            try:
                from scipy.cluster.hierarchy import linkage, fcluster
                from scipy.spatial.distance import squareform

                sim_matrix = (crop_embs @ crop_embs.T).numpy()
                dist_matrix = np.clip(1.0 - sim_matrix, 0, None)
                np.fill_diagonal(dist_matrix, 0.0)

                z_matrix = linkage(squareform(dist_matrix, checks=False), method="average")

                cluster_ids = fcluster(
                    z_matrix,
                    t=1.0 - SIM_THRESHOLD,
                    criterion="distance",
                )

                unique_clusters = sorted(set(cluster_ids))
                cluster_to_gid = {cid: idx for idx, cid in enumerate(unique_clusters)}

                local_group_ids = np.array([cluster_to_gid[cid] for cid in cluster_ids])

                for local_i, global_i in enumerate(clip_indices):
                    group_ids[global_i] = int(local_group_ids[local_i])

                n_visual_clusters = len(unique_clusters)
                n_groups = len(unique_clusters)
                counts = Counter(local_group_ids.tolist())

            except Exception as e:
                print(f"Grouping disabled because scipy/grouping failed: {e}")

    # -----------------------------
    # Optional price matching
    # -----------------------------
    top1_idx = [-1] * n_products_detected_after_filter
    top1_names = ["disabled"] * n_products_detected_after_filter
    top1_fams = ["disabled"] * n_products_detected_after_filter
    top1_sims = [0.0] * n_products_detected_after_filter
    unit_prices = [None] * n_products_detected_after_filter

    has_money_value = False

    if enable_price_matching and len(clip_indices) > 0:
        lazy_load_database()
        has_money_value = bool(_db_available and _pkey is not None)

        if has_money_value:
            embs_device = crop_embs.to(DEVICE)
            sims_matrix = embs_device @ _db_emb.T
            top_sims, top_idxs = sims_matrix.max(dim=1)

            local_top1_idx = top_idxs.cpu().numpy().astype(int).tolist()
            local_top1_sims = top_sims.cpu().numpy().astype(float).tolist()

            for local_i, global_i in enumerate(clip_indices):
                db_idx = local_top1_idx[local_i]
                top1_idx[global_i] = db_idx
                top1_sims[global_i] = local_top1_sims[local_i]
                top1_names[global_i] = _db_names[db_idx]
                top1_fams[global_i] = _db_fams[db_idx]
                unit_prices[global_i] = get_price_from_meta(db_idx)

    # -----------------------------
    # Products table
    # -----------------------------
    product_rows = []

    for i in range(n_products_detected_after_filter):
        x1, y1, x2, y2 = boxes[i]
        width = float(x2 - x1)
        height = float(y2 - y1)
        area = width * height
        gid = int(group_ids[i]) if i < len(group_ids) else -1

        product_rows.append({
            "product_id": i + 1,
            "clip_analyzed": bool(i in clip_indices),
            "group_id": gid,
            "detection_confidence": det_confs[i],
            "facing": facing_preds[i] if i < len(facing_preds) else "unknown",
            "front_score": facing_scores[i].get("front") if i < len(facing_scores) else None,
            "back_score": facing_scores[i].get("back") if i < len(facing_scores) else None,
            "matched_product_name": top1_names[i],
            "matched_family": top1_fams[i],
            "matched_similarity": top1_sims[i],
            "matched_db_idx": top1_idx[i],
            "price_key": _pkey,
            "unit_price": unit_prices[i],
            "x_min": float(x1),
            "y_min": float(y1),
            "x_max": float(x2),
            "y_max": float(y2),
            "width_px": width,
            "height_px": height,
            "area_px2": area,
            "area_image_percent": (area / image_area) * 100 if image_area > 0 else 0.0,
        })

    products_df = pd.DataFrame(product_rows)

    # -----------------------------
    # Groups table
    # -----------------------------
    group_rows = []
    total_shelf_value = 0.0

    for gid in sorted(counts.keys()):
        members = (
            products_df[products_df["group_id"] == gid]
            if not products_df.empty
            else pd.DataFrame()
        )

        if not members.empty:
            avg_conf = float(members["detection_confidence"].mean())
            total_area = float(members["area_px2"].sum())
            front_count = int((members["facing"] == "front").sum())
            back_count = int((members["facing"] == "back").sum())
            facings_count = len(members)

            if has_money_value:
                product_name = members["matched_product_name"].mode().iloc[0]
                family = members["matched_family"].mode().iloc[0]
                prices = members["unit_price"].dropna().astype(float).tolist()
                unit_price = float(np.median(prices)) if prices else None
                group_value = unit_price * facings_count if unit_price is not None else None
            else:
                product_name = "disabled"
                family = "disabled"
                unit_price = None
                group_value = None

            if group_value is not None:
                total_shelf_value += group_value

        else:
            avg_conf = None
            total_area = 0.0
            front_count = 0
            back_count = 0
            facings_count = 0
            product_name = "unknown"
            family = "unknown"
            unit_price = None
            group_value = None

        group_rows.append({
            "group_id": int(gid),
            "product_name": product_name,
            "family": family,
            "facings_count": int(facings_count),
            "front_count": front_count,
            "back_count": back_count,
            "unit_price": unit_price,
            "group_value": group_value,
            "avg_detection_confidence": avg_conf,
            "total_area_px2": total_area,
            "total_area_image_percent": (total_area / image_area) * 100 if image_area > 0 else 0.0,
        })

    groups_df = pd.DataFrame(group_rows)

    if not groups_df.empty:
        groups_df = groups_df.sort_values(
            by=["facings_count", "total_area_px2"],
            ascending=[False, False],
        )

    # -----------------------------
    # Business KPI
    # -----------------------------
    shelf_loss_pct = (empty_ratio + back_ratio) * 100
    shelf_profitability_pct = max(0.0, 100 - shelf_loss_pct)

    weighted_loss_pct = (
        (EMPTY_SPACE_PONDERATION * empty_ratio) + back_ratio
    ) * 100

    weighted_profitability_pct = max(0.0, 100 - weighted_loss_pct)

    if has_money_value:
        shelf_realised_value = float(total_shelf_value)
        shelf_unrealised_value = shelf_realised_value * (
            1 + (EMPTY_SPACE_PONDERATION * empty_ratio) + back_ratio
        )
        shelf_loss_value = shelf_unrealised_value - shelf_realised_value
    else:
        shelf_realised_value = None
        shelf_unrealised_value = None
        shelf_loss_value = None

    if weighted_profitability_pct >= 85:
        status = "Bon"
        severity = "low"
        recommendation = "Le rayon est globalement correct. Vérifier seulement les petites anomalies."
    elif weighted_profitability_pct >= 65:
        status = "Moyen"
        severity = "medium"
        recommendation = "Des actions de réassort ou de correction facing sont recommandées."
    else:
        status = "Critique"
        severity = "high"
        recommendation = "Priorité élevée : corriger les ruptures visibles et les produits mal orientés."

    annotated = draw_result_image(
        image_np=image_np,
        empty_polys=empty_polys,
        empty_confs=empty_confs,
        boxes=boxes,
        facing_preds=facing_preds,
        empty_df=empty_df,
    )

    total_time = time.perf_counter() - t0

    payload = {
        "ok": True,
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "status": status,
        "severity": severity,
        "recommendation": recommendation,
        "mode": "best_pt_empty_fastapi_style_boxes",
        "config": {
            "MAX_IMAGE_SIDE": int(max_image_side),
            "MAX_PRODUCTS_ANALYZED": int(MAX_PRODUCTS_ANALYZED),
            "MAX_PRODUCTS_FOR_CLIP": int(max_products_for_clip),
            "ENABLE_FACING_CLIP": bool(enable_facing_clip),
            "ENABLE_GROUPING_CLIP": bool(enable_grouping_clip),
            "ENABLE_PRICE_MATCHING": bool(enable_price_matching),
            "ANNOTATION_MODE": ANNOTATION_MODE,
            "EMPTY_CONF": EMPTY_CONF,
            "YOLO_CONF": YOLO_CONF,
            "SIM_THRESHOLD": SIM_THRESHOLD,
            "DEVICE": DEVICE,
        },
        "timing_seconds": {
            "total": float(total_time),
            "empty_model_best_pt": float(t_empty),
            "product_detector": float(t_yolo),
            "clip": float(t_clip),
        },
        "models": {
            "empty_space_model": str(EMPTY_WEIGHTS),
            "product_detection_model": str(YOLO_WEIGHTS) if yolo_model is not None else None,
            "ln_clip_model": str(LN_CKPT_PATH),
            "clip_loaded": bool(_clip_loaded),
            "number_of_models_loaded_at_startup": 2 if yolo_model is not None else 1,
        },
        "reference_database": {
            "loaded": bool(_db_available),
            "path": str(LN_DB_PATH),
            "is_model": False,
            "pkey": _pkey,
            "entries": int(_db_emb.shape[0]) if _db_available else 0,
            "price_matching_enabled": bool(enable_price_matching),
        },
        "counts": {
            "empty_spaces": int(n_empty),
            "raw_products_detected": int(raw_products_detected),
            "products_analyzed": int(n_products_detected_after_filter),
            "products_clip_analyzed": int(len(clip_indices)),
            "front_products": int(n_front),
            "back_products": int(n_back),
            "visual_clusters": int(n_visual_clusters),
            "product_groups": int(n_groups),
        },
        "surfaces": {
            "image_area_px2": float(image_area),
            "empty_surface_px2": float(empty_surface),
            "product_surface_px2": float(product_surface),
            "empty_ratio_notebook": float(empty_ratio),
            "empty_ratio_notebook_percent": float(empty_ratio * 100),
            "product_ratio_notebook": float(product_ratio),
            "product_ratio_notebook_percent": float(product_ratio * 100),
            "empty_ratio_image_percent": float(empty_image_ratio * 100),
            "product_ratio_image_percent": float(product_image_ratio * 100),
        },
        "facing": {
            "enabled": bool(enable_facing_clip),
            "front_products": int(n_front),
            "back_products": int(n_back),
            "back_ratio": float(back_ratio),
            "back_ratio_percent": float(back_ratio * 100),
        },
        "business": {
            "money_value_available": bool(has_money_value),
            "price_key_used": _pkey,
            "shelf_realised_value": shelf_realised_value,
            "shelf_unrealised_value": shelf_unrealised_value,
            "shelf_loss_value": shelf_loss_value,
            "shelf_loss_percent_notebook": float(shelf_loss_pct),
            "shelf_profitability_percent_notebook": float(shelf_profitability_pct),
            "weighted_loss_percent": float(weighted_loss_pct),
            "weighted_profitability_percent": float(weighted_profitability_pct),
        },
        "annotated_image_base64": pil_to_base64(annotated),
    }

    payload = sanitize_for_json(payload)

    return {
        "annotated": annotated,
        "payload": payload,
        "empty_df": empty_df,
        "products_df": products_df,
        "groups_df": groups_df,
    }


# ============================================================
# GRADIO WRAPPERS
# ============================================================
def analyze_shelf(
    image,
    enable_facing_clip,
    enable_grouping_clip,
    enable_price_matching,
    max_image_side,
    max_products_for_clip,
):
    if image is None:
        return (
            None,
            "Upload une image.",
            {},
            pd.DataFrame(),
            pd.DataFrame(),
            pd.DataFrame(),
        )

    result = run_analysis(
        image=image,
        enable_facing_clip=enable_facing_clip,
        enable_grouping_clip=enable_grouping_clip,
        enable_price_matching=enable_price_matching,
        max_image_side=max_image_side,
        max_products_for_clip=max_products_for_clip,
    )

    payload = result["payload"]

    def money(v):
        if v is None:
            return "N/A"
        return f"{v:.2f} MAD"

    summary = f"""
# ShelfGuide — Résumé

## Décision

| Indicateur | Valeur |
|---|---:|
| Statut | **{payload["status"]}** |
| Sévérité | **{payload["severity"]}** |
| Recommandation | {payload["recommendation"]} |

## Indicateurs clés

| KPI | Valeur |
|---|---:|
| Zones vides | {payload["counts"]["empty_spaces"]} |
| Produits détectés brut | {payload["counts"]["raw_products_detected"]} |
| Produits analysés | {payload["counts"]["products_analyzed"]} |
| Produits analysés par CLIP | {payload["counts"]["products_clip_analyzed"]} |
| Produits back/side | {payload["counts"]["back_products"]} |
| Groupes produits | {payload["counts"]["product_groups"]} |
| Empty ratio | {payload["surfaces"]["empty_ratio_notebook_percent"]:.1f} % |
| Back/side ratio | {payload["facing"]["back_ratio_percent"]:.1f} % |
| Profitabilité pondérée | {payload["business"]["weighted_profitability_percent"]:.1f} % |

## Temps d'exécution

| Étape | Secondes |
|---|---:|
| Total | {payload["timing_seconds"]["total"]:.2f} |
| Empty model best.pt | {payload["timing_seconds"]["empty_model_best_pt"]:.2f} |
| Product detector | {payload["timing_seconds"]["product_detector"]:.2f} |
| CLIP | {payload["timing_seconds"]["clip"]:.2f} |

## Valeur monétaire

| Indicateur | Valeur |
|---|---:|
| Price matching | {enable_price_matching} |
| pkey | {_pkey} |
| Shelf realised value | {money(payload["business"]["shelf_realised_value"])} |
| Shelf unrealised value | {money(payload["business"]["shelf_unrealised_value"])} |
| Shelf loss value | {money(payload["business"]["shelf_loss_value"])} |

## Configuration

- Empty-space model : `best.pt`
- Annotation empty-space : style ancien FastAPI, rectangle rouge transparent
- Product detector : `{payload["models"]["product_detection_model"]}`
- Mode : `{payload["mode"]}`
- Max image side : `{payload["config"]["MAX_IMAGE_SIDE"]}`
- Max products CLIP : `{payload["config"]["MAX_PRODUCTS_FOR_CLIP"]}`
- Device : `{DEVICE}`
- CLIP loaded : `{payload["models"]["clip_loaded"]}`
"""

    return (
        result["annotated"],
        summary,
        payload,
        result["empty_df"],
        result["products_df"],
        result["groups_df"],
    )


def analyze_shelf_api(image):
    """
    Mobile API rapide:
    - best.pt pour espaces vides
    - product_detector.pt pour produits si présent
    - CLIP désactivé
    """
    if image is None:
        return {
            "ok": False,
            "error": "No image provided",
        }

    try:
        result = run_analysis(
            image=image,
            enable_facing_clip=False,
            enable_grouping_clip=False,
            enable_price_matching=False,
            max_image_side=MAX_IMAGE_SIDE,
            max_products_for_clip=0,
        )
        return sanitize_for_json(result["payload"])
    except Exception as e:
        print("analyze_shelf_api error:", e)
        return {
            "ok": False,
            "error": str(e),
        }


# ============================================================
# GRADIO APP ONLY
# ============================================================
with gr.Blocks(title="ShelfGuide") as demo:
    gr.Markdown("# ShelfGuide")
    gr.Markdown(
        "Analyse rapide des rayons : espaces vides avec `best.pt`, "
        "annotation rouge style ancien FastAPI, produits, surfaces et profitabilité."
    )

    with gr.Row():
        input_image = gr.Image(type="pil", label="Image du rayon")
        output_image = gr.Image(type="numpy", label="Image annotée")

    with gr.Accordion("Options avancées", open=False):
        enable_facing_clip = gr.Checkbox(
            value=DEFAULT_ENABLE_FACING_CLIP,
            label="Activer facing avec CLIP — plus lent",
        )
        enable_grouping_clip = gr.Checkbox(
            value=DEFAULT_ENABLE_GROUPING_CLIP,
            label="Activer grouping visuel avec CLIP — plus lent",
        )
        enable_price_matching = gr.Checkbox(
            value=DEFAULT_ENABLE_PRICE_MATCHING,
            label="Activer price matching avec base LN — très lent sur CPU",
        )
        max_image_side = gr.Slider(
            minimum=512,
            maximum=1280,
            value=MAX_IMAGE_SIDE,
            step=64,
            label="Taille max image",
        )
        max_products_for_clip = gr.Slider(
            minimum=0,
            maximum=40,
            value=MAX_PRODUCTS_FOR_CLIP,
            step=1,
            label="Nombre max de produits envoyés à CLIP",
        )

    analyze_button = gr.Button("Analyser", variant="primary")

    output_summary = gr.Markdown()
    output_json = gr.JSON(label="Résultat JSON complet")

    with gr.Tab("Zones vides"):
        output_empty_df = gr.Dataframe(label="Détails des zones vides")

    with gr.Tab("Produits"):
        output_products_df = gr.Dataframe(label="Détails des produits analysés")

    with gr.Tab("Groupes produits"):
        output_groups_df = gr.Dataframe(label="Groupes produits")

    analyze_button.click(
        fn=analyze_shelf,
        inputs=[
            input_image,
            enable_facing_clip,
            enable_grouping_clip,
            enable_price_matching,
            max_image_side,
            max_products_for_clip,
        ],
        outputs=[
            output_image,
            output_summary,
            output_json,
            output_empty_df,
            output_products_df,
            output_groups_df,
        ],
        api_name="analyze",
    )

    # Hidden mobile endpoint.
    mobile_input = gr.Image(type="pil", visible=False)
    mobile_output = gr.JSON(visible=False)
    mobile_button = gr.Button("Mobile API", visible=False)

    mobile_button.click(
        fn=analyze_shelf_api,
        inputs=mobile_input,
        outputs=mobile_output,
        api_name="mobile_analyze",
    )

demo.queue(max_size=8)
demo.launch()