class_routing.py

"""
Class Routing - YOLO vs VLM
------------------------------
Two separate mappings needed for Agent 2's laptop damage detection:

1. TAXONOMY_MAP: every one of the 18 fine-grained Roboflow/YOLO classes maps
   to one of the shared 9-category taxonomy used throughout Module 2
   (screen_display, structural_body, port_connector, missing_component,
   liquid_moisture, input_control, power_boot_failure, cosmetic_wear,
   other_unclassified). Needed regardless of the routing decision below,
   since Agent 4/blockchain/mobile all expect the shared taxonomy, not
   Roboflow's granular labels.

2. YOLO_TRAINED_CLASSES / VLM_FALLBACK_CLASSES: which of the 18 classes the
   trained YOLO model is actually responsible for (had enough real source
   data - >=100 instances) vs which get routed to the VLM classifier
   instead (too little data to trust a trained detector on, per real counts
   found by inspection). This does NOT mean those damage types are lost -
   it means Agent 2 asks the VLM about them instead of the trained model,
   the same way printer/projector already work VLM-only.
"""

# ----- 1. Fine-grained YOLO class -> shared taxonomy category -----
TAXONOMY_MAP = {
    "scratch": "cosmetic_wear",
    "chip": "cosmetic_wear",
    "adhesive_residue": "cosmetic_wear",

    "crack": "structural_body",
    "broken_other": "structural_body",
    "broken_lock": "structural_body",
    "disassembled": "structural_body",
    "depression": "structural_body",

    "damaged_screen": "screen_display",
    "dead_pixel": "screen_display",
    "display_lines": "screen_display",
    "display_spot": "screen_display",
    "display_fade": "screen_display",

    "missing_button": "missing_component",
    "missing_screw": "missing_component",

    "keyboard_issue": "input_control",
    "broken_button": "input_control",

    "normal": None,  # not a damage type - "no damage detected here" reference class
}

# ----- 2. Which classes the trained YOLO model actually handles -----
# Based on real per-class instance counts in the training data (before
# oversampling - duplicated copies don't add real information, so the
# threshold is judged against genuine source image counts).
#
# NOTE: "scratch" was originally here (21,656 real instances - by far the
# most data of any class) but is EXCLUDED despite that, moved to
# VLM_FALLBACK_CLASSES instead. Confirmed via real testing across TWO
# model generations (YOLOv8: 0.191 mAP50, YOLO26: 0.206 mAP50) that more
# data does not fix it - scratches are thin, low-contrast, and densely
# packed, which is a detection-difficulty problem, not a data-volume one.
# The VLM's own visual reasoning ("is there a scratch, how bad") doesn't
# need precise bounding boxes the way object detection does, and is
# expected to handle this better.
YOLO_TRAINED_CLASSES = {
    "broken_other", "chip", "missing_button", "keyboard_issue",
    "missing_screw", "dead_pixel", "damaged_screen", "crack", "display_lines",
    "normal",  # kept for its negative/reference role, not as a "damage type"
}

VLM_FALLBACK_CLASSES = {
    "scratch", "broken_lock", "depression", "disassembled", "display_spot",
    "adhesive_residue", "display_fade", "broken_button",
}


def get_shared_category(fine_grained_class: str) -> str:
    """Map a fine-grained YOLO class name to the shared 9-category taxonomy."""
    return TAXONOMY_MAP.get(fine_grained_class, "other_unclassified")


def yolo_can_handle(fine_grained_class: str) -> bool:
    """True if the trained YOLO model has enough real data to be trusted on this class."""
    return fine_grained_class in YOLO_TRAINED_CLASSES


def route_detection(yolo_detections: list, confidence_threshold: float = 0.4) -> dict:
    """
    Given YOLO's raw detections for one crop (list of {"class_name": str,
    "confidence": float}), decide whether to trust YOLO's answer or fall
    back to the VLM classifier.

    Trusts YOLO only if its TOP detection is (a) a class it has enough real
    data for, and (b) above the confidence threshold. Otherwise signals a
    VLM fallback - covering both "YOLO found nothing confident" and
    "YOLO's best guess was one of the data-insufficient classes."
    """
    if not yolo_detections:
        return {"use_yolo": False, "reason": "no_detections", "fine_grained_class": None}

    top = max(yolo_detections, key=lambda d: d["confidence"])

    if top["confidence"] < confidence_threshold:
        return {"use_yolo": False, "reason": "low_confidence", "fine_grained_class": top["class_name"]}

    if not yolo_can_handle(top["class_name"]):
        return {
            "use_yolo": False,
            "reason": "class_needs_more_training_data",
            "fine_grained_class": top["class_name"],
        }

    return {
        "use_yolo": True,
        "reason": "trained_class_confident",
        "fine_grained_class": top["class_name"],
        "shared_category": get_shared_category(top["class_name"]),
        "confidence": top["confidence"],
    }