inspection/surface_gate.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import cv2
import numpy as np

from core.config import settings
from core.logger import logger

try:
    from ultralytics import YOLO
except Exception:  # pragma: no cover - optional import guard
    YOLO = None


@dataclass(frozen=True)
class SurfaceGateResult:
    is_steel: bool
    confidence: float
    mode: str
    label: str
    reason: str
    metrics: dict[str, float]
    roi_bbox: tuple[int, int, int, int] | None = None
    roi_area_ratio: float | None = None

    def to_metadata(self) -> dict[str, Any]:
        return {
            "passed": self.is_steel,
            "confidence": round(self.confidence, 3),
            "mode": self.mode,
            "label": self.label,
            "reason": self.reason,
            "metrics": self.metrics,
            "roi_bbox": list(self.roi_bbox) if self.roi_bbox else None,
            "roi_area_ratio": round(self.roi_area_ratio, 4) if self.roi_area_ratio is not None else None,
        }


class SteelSurfaceGate:
    def __init__(self) -> None:
        self.detector_model = self._load_model(settings.SURFACE_DETECTOR_PATH, "detector")
        self.classifier_model = self._load_model(settings.SURFACE_CLASSIFIER_PATH, "classifier")

    def _load_model(self, model_path: str | None, kind: str):
        if not model_path:
            return None

        resolved_path = Path(model_path)
        if not resolved_path.exists():
            logger.warning("Surface %s path does not exist: %s", kind, resolved_path)
            return None

        if YOLO is None:
            logger.warning("Ultralytics is unavailable for the surface %s model", kind)
            return None

        try:
            model = YOLO(str(resolved_path))
            logger.info("Surface %s model loaded from: %s", kind, resolved_path)
            return model
        except Exception as exc:  # pragma: no cover - defensive runtime fallback
            logger.warning("Surface %s model failed to load: %s", kind, exc)
            return None

    def evaluate(self, image: np.ndarray) -> SurfaceGateResult:
        if self.detector_model is not None:
            return self._evaluate_with_detector(image)

        if self.classifier_model is not None:
            return self._evaluate_with_classifier(image)

        return self._evaluate_with_heuristic(image)

    def _evaluate_with_detector(self, image: np.ndarray) -> SurfaceGateResult:
        results = self.detector_model(
            image,
            conf=settings.SURFACE_DETECTOR_CONFIDENCE,
            imgsz=settings.SURFACE_DETECTOR_IMAGE_SIZE,
            verbose=False,
        )
        result = results[0]
        boxes = getattr(result, "boxes", None)
        names = getattr(result, "names", None) or getattr(self.detector_model, "names", {})

        if boxes is None or len(boxes) == 0:
            return SurfaceGateResult(
                is_steel=False,
                confidence=0.99,
                mode="detector",
                label="non_steel",
                reason=(
                    "Frame skipped because no steel-surface region was localized. Aim the camera closer to the sheet or coil."
                ),
                metrics={"boxes_detected": 0.0},
            )

        height, width = image.shape[:2]
        image_area = max(height * width, 1)
        candidates: list[dict[str, float | tuple[int, int, int, int] | str]] = []

        for raw_box in boxes:
            xyxy = raw_box.xyxy[0].tolist()
            x1, y1, x2, y2 = [int(round(value)) for value in xyxy]
            x1 = max(0, min(x1, width - 1))
            y1 = max(0, min(y1, height - 1))
            x2 = max(x1 + 1, min(x2, width))
            y2 = max(y1 + 1, min(y2, height))

            box_width = x2 - x1
            box_height = y2 - y1
            area_ratio = (box_width * box_height) / image_area

            cls_tensor = getattr(raw_box, "cls", None)
            cls_index = int(cls_tensor[0].item()) if cls_tensor is not None else 0
            label = self._resolve_label(names, cls_index)
            confidence = float(raw_box.conf[0].item())

            if not self._detector_label_matches(label, cls_index, names):
                continue

            expand_ratio = settings.SURFACE_DETECTOR_EXPAND_RATIO
            pad_x = int(box_width * expand_ratio)
            pad_y = int(box_height * expand_ratio)
            expanded_x1 = max(0, x1 - pad_x)
            expanded_y1 = max(0, y1 - pad_y)
            expanded_x2 = min(width, x2 + pad_x)
            expanded_y2 = min(height, y2 + pad_y)

            expanded_bbox = (
                expanded_x1,
                expanded_y1,
                expanded_x2 - expanded_x1,
                expanded_y2 - expanded_y1,
            )
            expanded_area_ratio = (expanded_bbox[2] * expanded_bbox[3]) / image_area
            selector_score = (confidence * 0.8) + (min(expanded_area_ratio / 0.55, 1.0) * 0.2)

            candidates.append(
                {
                    "confidence": confidence,
                    "label": label,
                    "bbox": expanded_bbox,
                    "area_ratio": expanded_area_ratio,
                    "selector_score": selector_score,
                }
            )

        if not candidates:
            return SurfaceGateResult(
                is_steel=False,
                confidence=0.99,
                mode="detector",
                label="non_steel",
                reason=(
                    "Frame skipped because the localized objects did not match the expected steel-surface class."
                ),
                metrics={"boxes_detected": float(len(boxes))},
            )

        best_candidate = max(
            candidates,
            key=lambda candidate: (
                float(candidate["selector_score"]),
                float(candidate["confidence"]),
                float(candidate["area_ratio"]),
            ),
        )

        roi_area_ratio = float(best_candidate["area_ratio"])
        if roi_area_ratio < settings.SURFACE_DETECTOR_MIN_AREA_RATIO:
            return SurfaceGateResult(
                is_steel=False,
                confidence=float(best_candidate["confidence"]),
                mode="detector",
                label=str(best_candidate["label"]),
                reason=(
                    "Frame skipped because the localized steel region is too small for reliable defect inspection. Move closer to the material."
                ),
                metrics={
                    "boxes_detected": float(len(candidates)),
                    "roi_area_ratio": round(roi_area_ratio, 4),
                },
                roi_bbox=best_candidate["bbox"],
                roi_area_ratio=roi_area_ratio,
            )

        roi_x, roi_y, roi_width, roi_height = best_candidate["bbox"]
        roi_image = image[roi_y:roi_y + roi_height, roi_x:roi_x + roi_width]
        heuristic_result = self._evaluate_with_heuristic(roi_image)
        if not heuristic_result.is_steel:
            return SurfaceGateResult(
                is_steel=False,
                confidence=float(best_candidate["confidence"]),
                mode="detector+heuristic",
                label=str(best_candidate["label"]),
                reason=(
                    "Frame skipped because the localized ROI did not pass the steel-surface texture validation step. "
                    "Reduce background content and center the actual material."
                ),
                metrics={
                    "boxes_detected": float(len(candidates)),
                    "roi_area_ratio": round(roi_area_ratio, 4),
                    "selector_score": round(float(best_candidate["selector_score"]), 4),
                    **heuristic_result.metrics,
                },
                roi_bbox=best_candidate["bbox"],
                roi_area_ratio=roi_area_ratio,
            )

        return SurfaceGateResult(
            is_steel=True,
            confidence=float(best_candidate["confidence"]),
            mode="detector+heuristic",
            label=str(best_candidate["label"]),
            reason=(
                f"Steel-surface detector localized an inspection ROI with {float(best_candidate['confidence']):.0%} confidence, "
                "and the ROI passed surface-texture validation."
            ),
            metrics={
                "boxes_detected": float(len(candidates)),
                "roi_area_ratio": round(roi_area_ratio, 4),
                "selector_score": round(float(best_candidate["selector_score"]), 4),
                **heuristic_result.metrics,
            },
            roi_bbox=best_candidate["bbox"],
            roi_area_ratio=roi_area_ratio,
        )

    def _evaluate_with_classifier(self, image: np.ndarray) -> SurfaceGateResult:
        results = self.classifier_model(image, verbose=False)
        result = results[0]
        probs = getattr(result, "probs", None)
        names = getattr(result, "names", None) or getattr(self.classifier_model, "names", {})

        if probs is None:
            raise RuntimeError("Classification model returned no probabilities")

        top_index = int(getattr(probs, "top1", 0))
        top_confidence_raw = getattr(probs, "top1conf")
        top_confidence = float(
            top_confidence_raw.item() if hasattr(top_confidence_raw, "item") else top_confidence_raw
        )
        label = self._resolve_label(names, top_index)

        steel_label = settings.SURFACE_CLASSIFIER_STEEL_LABEL.strip().lower()
        normalized_label = label.strip().lower()
        is_steel = (
            normalized_label == steel_label
            or steel_label in normalized_label
            or normalized_label in steel_label
        )
        meets_threshold = top_confidence >= settings.SURFACE_MIN_STEEL_CONFIDENCE

        if is_steel and meets_threshold:
            reason = (
                f"Steel surface classifier accepted the frame with {top_confidence:.0%} confidence."
            )
        elif is_steel:
            reason = (
                "Frame resembles steel, but the classifier confidence is too low for reliable defect analysis."
            )
        else:
            reason = (
                f'Classifier labeled the frame as "{label}" instead of steel, so the defect model was skipped.'
            )

        return SurfaceGateResult(
            is_steel=bool(is_steel and meets_threshold),
            confidence=top_confidence,
            mode="classifier",
            label=label,
            reason=reason,
            metrics={},
            roi_bbox=(0, 0, image.shape[1], image.shape[0]) if is_steel and meets_threshold else None,
            roi_area_ratio=1.0 if is_steel and meets_threshold else None,
        )

    def _evaluate_with_heuristic(self, image: np.ndarray) -> SurfaceGateResult:
        features = self._extract_features(image)

        checks = {
            "gray_ratio": features["gray_ratio"] >= settings.SURFACE_MIN_GRAY_RATIO,
            "low_sat_ratio": features["low_saturation_ratio"] >= settings.SURFACE_MIN_LOW_SAT_RATIO,
            "mean_saturation": features["mean_saturation"] <= settings.SURFACE_MAX_MEAN_SATURATION,
            "colorfulness": features["colorfulness"] <= settings.SURFACE_MAX_COLORFULNESS,
            "skin_ratio": features["skin_ratio"] <= settings.SURFACE_MAX_SKIN_RATIO,
            "texture_variance": features["texture_variance"] >= settings.SURFACE_MIN_TEXTURE_VARIANCE,
        }

        weights = {
            "gray_ratio": 0.24,
            "low_sat_ratio": 0.22,
            "mean_saturation": 0.16,
            "colorfulness": 0.16,
            "skin_ratio": 0.12,
            "texture_variance": 0.10,
        }

        score = sum(weights[name] for name, passed in checks.items() if passed)
        chroma_gate = checks["gray_ratio"] and checks["low_sat_ratio"]
        is_steel = score >= 0.72 and chroma_gate and checks["skin_ratio"]

        failed_checks = [name for name, passed in checks.items() if not passed]
        confidence = score if is_steel else min(0.99, max(0.55, 1.0 - score + 0.08 * len(failed_checks)))

        if is_steel:
            reason = (
                f"Frame passed the steel-surface gate with {confidence:.0%} confidence and proceeded to defect segmentation."
            )
        else:
            reason = self._build_failure_reason(features, failed_checks)

        return SurfaceGateResult(
            is_steel=is_steel,
            confidence=confidence,
            mode="heuristic",
            label="steel" if is_steel else "non_steel",
            reason=reason,
            metrics={key: round(value, 4) for key, value in features.items()},
            roi_bbox=(0, 0, image.shape[1], image.shape[0]) if is_steel else None,
            roi_area_ratio=1.0 if is_steel else None,
        )

    def _resolve_label(self, names: Any, index: int) -> str:
        if isinstance(names, dict):
            return str(names.get(index, index))
        if isinstance(names, list) and 0 <= index < len(names):
            return str(names[index])
        return str(index)

    def _detector_label_matches(self, label: str, cls_index: int, names: Any) -> bool:
        target = settings.SURFACE_DETECTOR_CLASS_NAME.strip().lower()
        normalized_label = label.strip().lower()

        if not target:
            return True

        if normalized_label == target or target in normalized_label or normalized_label in target:
            return True

        if isinstance(names, dict) and len(names) == 1 and cls_index == 0:
            return True

        if isinstance(names, list) and len(names) == 1 and cls_index == 0:
            return True

        return False

    def _extract_features(self, image: np.ndarray) -> dict[str, float]:
        height, width = image.shape[:2]
        target_width = min(320, max(96, width))
        target_height = max(96, int(height * target_width / max(width, 1)))
        resized = cv2.resize(image, (target_width, target_height))

        hsv = cv2.cvtColor(resized, cv2.COLOR_BGR2HSV)
        gray = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY)
        ycrcb = cv2.cvtColor(resized, cv2.COLOR_BGR2YCrCb)

        b_channel, g_channel, r_channel = [channel.astype(np.float32) for channel in cv2.split(resized)]
        saturation = hsv[:, :, 1].astype(np.float32)

        gray_delta = settings.SURFACE_GRAY_DELTA
        gray_mask = (
            (np.abs(r_channel - g_channel) <= gray_delta)
            & (np.abs(r_channel - b_channel) <= gray_delta)
            & (np.abs(g_channel - b_channel) <= gray_delta)
        )

        rg = np.abs(r_channel - g_channel)
        yb = np.abs(0.5 * (r_channel + g_channel) - b_channel)
        colorfulness = (
            np.sqrt(float(rg.std()) ** 2 + float(yb.std()) ** 2)
            + 0.3 * np.sqrt(float(rg.mean()) ** 2 + float(yb.mean()) ** 2)
        )

        luminance, cr_channel, cb_channel = cv2.split(ycrcb)
        skin_mask = (
            (cr_channel > 135)
            & (cr_channel < 180)
            & (cb_channel > 85)
            & (cb_channel < 135)
            & (luminance > 60)
        )

        return {
            "gray_ratio": float(gray_mask.mean()),
            "low_saturation_ratio": float(
                (saturation <= settings.SURFACE_LOW_SAT_PIXEL_THRESHOLD).mean()
            ),
            "mean_saturation": float(saturation.mean()),
            "colorfulness": float(colorfulness),
            "skin_ratio": float(skin_mask.mean()),
            "texture_variance": float(cv2.Laplacian(gray, cv2.CV_32F).var()),
        }

    def _build_failure_reason(self, features: dict[str, float], failed_checks: list[str]) -> str:
        if features["skin_ratio"] > settings.SURFACE_MAX_SKIN_RATIO:
            return (
                "Frame skipped because prominent skin-tone regions were detected. Aim the camera only at the steel surface."
            )
        if features["gray_ratio"] < settings.SURFACE_MIN_GRAY_RATIO:
            return (
                "Frame skipped because it contains too much color variation to match the expected steel surface appearance."
            )
        if features["low_saturation_ratio"] < settings.SURFACE_MIN_LOW_SAT_RATIO:
            return (
                "Frame skipped because the image is too saturated. Move closer to the metal surface and reduce background content."
            )
        if features["texture_variance"] < settings.SURFACE_MIN_TEXTURE_VARIANCE:
            return (
                "Frame skipped because the visible area is too flat or out of focus for reliable steel-surface validation."
            )
        if features["colorfulness"] > settings.SURFACE_MAX_COLORFULNESS:
            return (
                "Frame skipped because the scene looks like a general object view instead of a steel inspection close-up."
            )

        failed_text = ", ".join(failed_checks) if failed_checks else "multiple surface validation checks"
        return (
            f"Frame skipped because it did not pass the steel-surface gate ({failed_text}). Reposition the camera toward the material and retry."
        )