"""
Face Mask Detection and Inpainting Module.

Detects if a person is wearing a face mask and removes it via inpainting
to reveal the underlying face for recognition.

Mask detection uses a heuristic based on lower-face analysis.
Inpainting uses OpenCV's TELEA algorithm on the masked region.
"""

from typing import Tuple, Optional
import numpy as np
import cv2


class MaskProcessor:
    """
    Detects face masks and performs inpainting to reveal the face.

    Uses color/texture analysis in the lower half of the face to detect masks.
    Inpaints using cv2.INPAINT_TELEA for CPU-friendly operation.
    """

    def __init__(self, mask_threshold: float = 0.5):
        """
        Initialize the mask processor.

        Args:
            mask_threshold: Confidence threshold for mask detection.
        """
        self.mask_threshold = mask_threshold

        # Skin color ranges in HSV for lower-face analysis
        self.lower_skin = np.array([0, 20, 70], dtype=np.uint8)
        self.upper_skin = np.array([20, 150, 255], dtype=np.uint8)

        # Mask color ranges (blue/green surgical masks, black cloth masks)
        self.mask_colors = {
            'blue': (np.array([90, 50, 50], dtype=np.uint8), np.array([120, 255, 255], dtype=np.uint8)),
            'green': (np.array([40, 50, 50], dtype=np.uint8), np.array([80, 255, 255], dtype=np.uint8)),
            'white': (np.array([0, 0, 180], dtype=np.uint8), np.array([180, 30, 255], dtype=np.uint8)),
            'black': (np.array([0, 0, 0], dtype=np.uint8), np.array([180, 255, 60], dtype=np.uint8)),
        }

    def detect_mask(self, face_crop: np.ndarray) -> Tuple[bool, float]:
        """
        Detect if a face crop contains a mask.

        Analyzes the lower half of the face for mask-like features.
        Uses color segmentation and texture analysis.

        Args:
            face_crop: Cropped face image (BGR).

        Returns:
            Tuple of (is_masked: bool, confidence: float).
        """
        if face_crop is None or face_crop.size == 0:
            return False, 0.0

        h, w = face_crop.shape[:2]
        if h < 20 or w < 20:
            return False, 0.0

        # Extract lower half of face (where masks are worn)
        lower_half = face_crop[h // 2:, :, :]

        # Convert to HSV for color analysis
        hsv = cv2.cvtColor(lower_half, cv2.COLOR_BGR2HSV)

        # Check for skin pixels in lower half
        skin_mask = cv2.inRange(hsv, self.lower_skin, self.upper_skin)
        skin_ratio = np.sum(skin_mask > 0) / skin_mask.size

        # Check for mask color pixels
        mask_pixel_ratio = 0.0
        for color_name, (lower, upper) in self.mask_colors.items():
            color_mask = cv2.inRange(hsv, lower, upper)
            ratio = np.sum(color_mask > 0) / color_mask.size
            mask_pixel_ratio = max(mask_pixel_ratio, ratio)

        # Edge detection: masks have fewer edges than skin (smooth surface)
        gray = cv2.cvtColor(lower_half, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 50, 150)
        edge_density = np.sum(edges > 0) / edges.size

        # Decision logic:
        # Mask is likely if: low skin ratio + high mask color ratio + low edge density
        is_masked = False
        confidence = 0.0

        if skin_ratio < 0.2 and mask_pixel_ratio > 0.15:
            # Strong mask color match
            is_masked = True
            confidence = min(1.0, mask_pixel_ratio * 2.0)
        elif skin_ratio < 0.15 and edge_density < 0.05:
            # Low skin and low edges -> likely mask
            is_masked = True
            confidence = 0.7
        elif skin_ratio < 0.3 and mask_pixel_ratio > 0.1:
            # Moderate evidence
            is_masked = True
            confidence = 0.6
        elif edge_density < 0.02 and mask_pixel_ratio > 0.05:
            # Very smooth surface + some mask color
            is_masked = True
            confidence = 0.55
        else:
            # Likely no mask - skin visible
            confidence = max(0, min(0.5, skin_ratio))

        return is_masked, float(confidence)

    def remove_mask(
        self,
        face_crop: np.ndarray,
        mask_region: Optional[np.ndarray] = None
    ) -> np.ndarray:
        """
        Remove mask from face via inpainting.

        Creates a mask of the lower half of the face and inpaints
        using OpenCV's TELEA algorithm for smooth results.

        Args:
            face_crop: Cropped face image (BGR).
            mask_region: Optional custom mask region. If None, uses lower half.

        Returns:
            Inpainted face crop.
        """
        if face_crop is None or face_crop.size == 0:
            return face_crop

        h, w = face_crop.shape[:2]
        if h < 20 or w < 20:
            return face_crop

        if mask_region is not None:
            inpaint_mask = mask_region
        else:
            # Create mask for lower half of face (below nose area)
            inpaint_mask = np.zeros((h, w), dtype=np.uint8)
            # Lower 45% of face (chin + mouth area)
            mask_start = int(h * 0.45)
            inpaint_mask[mask_start:, :] = 255

            # Also mask the sides (cheeks) for more natural blending
            side_width = int(w * 0.08)
            inpaint_mask[mask_start:, :side_width] = 255
            inpaint_mask[mask_start:, w - side_width:] = 255

        # Apply inpainting
        inpainted = cv2.inpaint(
            face_crop,
            inpaint_mask,
            inpaintRadius=3,
            flags=cv2.INPAINT_TELEA
        )

        return inpainted


if __name__ == "__main__":
    # Quick test
    mp = MaskProcessor()
    test_face = np.random.randint(0, 255, (200, 150, 3), dtype=np.uint8)
    is_masked, conf = mp.detect_mask(test_face)
    print(f"Mask detection test: is_masked={is_masked}, confidence={conf:.3f}")
    inpainted = mp.remove_mask(test_face)
    print(f"Inpainting test: output shape={inpainted.shape}")
    print("MaskProcessor OK!")