app/services/fracture_model.py

"""Bone fracture detection using YOLOv8.

Uses the Ultralytics YOLOv8 model for real-time object detection
and localization of fractures in X-ray images.
"""

from __future__ import annotations
from pathlib import Path
import numpy as np
import logging

logger = logging.getLogger(__name__)

# Lazy-loaded model singleton
_model = None


def _get_model():
    """Load the YOLOv8 model (lazy singleton)."""
    global _model
    if _model is None:
        logger.info("Loading YOLOv8 model...")
        try:
            from ultralytics import YOLO
            from app.config import get_settings

            settings = get_settings()
            weights = settings.yolo_weights_path
            weights_name = Path(weights).name.lower()
            generic_weights = {"yolov8n.pt", "yolov8s.pt", "yolov8m.pt", "yolov8l.pt", "yolov8x.pt"}
            if weights_name in generic_weights and not settings.allow_generic_yolo_weights:
                raise RuntimeError(
                    "Generic YOLO weights are not valid for medical fracture delivery. "
                    "Set YOLO_WEIGHTS_PATH to fracture-trained weights, or set "
                    "ALLOW_GENERIC_YOLO_WEIGHTS=true only for demo mode."
                )
            if not Path(weights).exists() and weights_name not in generic_weights:
                raise FileNotFoundError(
                    f"Fracture YOLO weights not found at '{weights}'. "
                    "Place fracture-trained weights there or update YOLO_WEIGHTS_PATH."
                )
            _model = YOLO(weights)
            logger.info(f"YOLOv8 loaded from: {weights}")
        except Exception as e:
            logger.error(f"Failed to load YOLOv8: {e}")
            raise
    return _model


def predict_fractures(image: np.ndarray,
                      confidence_threshold: float = 0.15) -> list[dict]:
    """Run fracture detection inference.

    Args:
        image: BGR numpy array (original size, YOLO handles resizing).
        confidence_threshold: minimum confidence to include detections.

    Returns:
        List of findings with bounding boxes, confidence, and severity.
    """
    model = _get_model()

    results = model(image, conf=confidence_threshold, imgsz=960, verbose=False)

    findings = []
    img_h, img_w = image.shape[:2]

    for result in results:
        boxes = result.boxes
        if boxes is None:
            continue

        for box in boxes:
            # Get bounding box coordinates (xyxy format)
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
            conf = float(box.conf[0].cpu().numpy())
            cls_id = int(box.cls[0].cpu().numpy())
            cls_name = str(result.names.get(cls_id, f"class_{cls_id}"))
            cls_key = cls_name.lower().replace(" ", "_")

            # Convert to percentage-based coordinates for the frontend
            bbox = {
                "x": round(float(x1 / img_w * 100), 1),
                "y": round(float(y1 / img_h * 100), 1),
                "w": round(float((x2 - x1) / img_w * 100), 1),
                "h": round(float((y2 - y1) / img_h * 100), 1),
            }

            confidence = round(conf * 100, 1)

            # Map raw class names to clean medical labels
            label = _clean_class_name(cls_name)

            # Skip very low-confidence "Not_Fracture" detections — not useful to display
            if cls_key in ("not_fracture", "normal", "negative"):
                continue

            # Determine severity
            if confidence >= 80:
                severity = "high"
                color = "destructive"
            elif confidence >= 60:
                severity = "moderate"
                color = "warning"
            else:
                severity = "low"
                color = "info"

            icd = "S02-S92" if cls_key == "fracture" or cls_key.endswith("_fracture") else ""

            findings.append({
                "name": label,
                "confidence": confidence,
                "severity": severity,
                "model": "YOLOv8",
                "region": _infer_region(bbox),
                "icd_code": icd,
                "bbox": bbox,
                "color": color,
            })

    # Sort by confidence descending
    findings.sort(key=lambda x: x["confidence"], reverse=True)

    # If no fractures detected, return a "clear" finding
    if not findings:
        findings.append({
            "name": "No fracture box localized",
            "confidence": 0.0,
            "severity": "low",
            "model": "YOLOv8",
            "region": "Full image",
            "icd_code": "",
            "color": "warning",
        })

    return findings


def _clean_class_name(cls_name: str) -> str:
    """Map raw Roboflow class names to clean medical labels."""
    mapping = {
        "fracture":      "Fracture Detected",
        "not fracture":  "No Fracture Detected",
        "not_fracture":  "No Fracture Detected",
        "normal":        "No Fracture Detected",
        "negative":      "No Fracture Detected",
        "boneanomaly":   "Bone Anomaly",
        "bonelesion":    "Bone Lesion",
        "foreignbody":   "Foreign Body",
        "metal":         "Metallic Implant",
        "periostealreaction": "Periosteal Reaction",
        "pronationsign": "Pronation Sign",
        "softtissue":    "Soft Tissue Finding",
        "hardware":      "Surgical Hardware",
    }
    return mapping.get(cls_name.lower().replace(" ", "_"), cls_name.replace("_", " ").title())


def _infer_region(bbox: dict) -> str:
    """Infer the anatomical region based on bounding box position."""
    cx = bbox["x"] + bbox["w"] / 2
    cy = bbox["y"] + bbox["h"] / 2

    if cy < 30:
        return "Upper extremity"
    elif cy < 60:
        return "Mid-body / Torso"
    else:
        return "Lower extremity"