src/fitzpatrick_estimator.py

"""Fitzpatrick skin type estimation via ITA (Individual Typology Angle).

Calibrated on 61 DFU images with expert ground truth.
Validation: 86.9% exact match, 98.4% adjacent, r=0.975.

Includes:
  - Shadow/wound contamination filtering via L* outlier rejection
  - Adaptive trimming (tighter when ITA variance is high)
  - Lighting quality assessment to flag poor illumination
"""
import numpy as np
import cv2
from dataclasses import dataclass
from typing import Optional

# Calibrated ITA thresholds for DFU clinical photography (61-image validation)
ITA_THRESHOLDS = {
    "I":   (46.86, float("inf")),
    "II":  (34.25, 46.86),
    "III": (20.87, 34.25),
    "IV":  (3.57, 20.87),
    "V":   (-28.38, 3.57),
    "VI":  (float("-inf"), -28.38),
}

FITZPATRICK_LABELS = {
    "I": "Very Light", "II": "Light", "III": "Intermediate",
    "IV": "Tan", "V": "Brown", "VI": "Dark",
}

# Lighting quality thresholds (L* scale 0-100)
L_SCENE_MIN = 35.0
L_SCENE_LOW = 50.0
L_SKIN_MIN = 25.0
L_SKIN_SUSPICIOUS = 40.0


@dataclass
class FitzpatrickResult:
    fitzpatrick_type: str
    fitzpatrick_int: int
    fitzpatrick_label: str
    ita_angle: float
    ita_std: float
    l_skin_mean: float
    b_skin_mean: float
    healthy_pixels: int
    healthy_ratio: float
    confidence: float
    l_scene_mean: float = 0.0
    lighting_quality: str = "good"
    lighting_warning: str = ""


def filter_shadow_pixels(l_values: np.ndarray, b_values: np.ndarray) -> tuple:
    """Remove shadow/contamination pixels from healthy skin sample.

    Strategy: In a well-sampled skin region, L* follows a unimodal distribution.
    Shadow contamination creates a low-L* tail. We detect this by checking if the
    distribution is bimodal or has excessive spread, and keep only the main mode.

    Returns filtered (l_values, b_values).
    """
    if len(l_values) < 100:
        return l_values, b_values

    # Compute L* statistics
    l_median = np.median(l_values)
    l_std = np.std(l_values)

    # If L* spread is reasonable (std < 12), no filtering needed
    if l_std < 12:
        return l_values, b_values

    # High spread: likely shadow contamination. Use IQR-based filtering.
    q1 = np.percentile(l_values, 25)
    q3 = np.percentile(l_values, 75)
    iqr = q3 - q1

    # Keep pixels within [Q1 - 0.5*IQR, Q3 + 1.5*IQR]
    # Asymmetric: more aggressive on the low end (shadows) than high end (specular)
    lo = q1 - 0.5 * iqr
    hi = q3 + 1.5 * iqr
    keep = (l_values >= lo) & (l_values <= hi)

    if np.sum(keep) < 50:
        # Fallback: just use upper half of L* values
        keep = l_values >= l_median

    return l_values[keep], b_values[keep]


def compute_ita(l_values: np.ndarray, b_values: np.ndarray) -> tuple:
    """Compute ITA angle from L* and b* values with adaptive trimming.

    ITA = arctan((L* - 50) / b*) * (180 / pi)
    Higher ITA = lighter skin, lower ITA = darker skin.

    Uses adaptive percentile trimming: starts at 10-90, tightens if variance is high.
    """
    ita_per_pixel = np.degrees(np.arctan2(l_values - 50.0, b_values))

    # First pass: 10th-90th percentile (tighter than original 5-95)
    p10, p90 = np.percentile(ita_per_pixel, [10, 90])
    trimmed = ita_per_pixel[(ita_per_pixel >= p10) & (ita_per_pixel <= p90)]

    if len(trimmed) < 10:
        trimmed = ita_per_pixel

    first_std = float(np.std(trimmed))

    # If still high variance, tighten to 25-75 (IQR)
    if first_std > 20 and len(ita_per_pixel) > 200:
        p25, p75 = np.percentile(ita_per_pixel, [25, 75])
        iqr_trimmed = ita_per_pixel[(ita_per_pixel >= p25) & (ita_per_pixel <= p75)]
        if len(iqr_trimmed) >= 50:
            trimmed = iqr_trimmed

    return float(np.mean(trimmed)), float(np.std(trimmed))


def classify_fitzpatrick(ita: float) -> tuple:
    """Classify ITA angle into Fitzpatrick type using calibrated DFU thresholds."""
    for ftype, (lo, hi) in ITA_THRESHOLDS.items():
        if lo <= ita < hi:
            idx = list(ITA_THRESHOLDS.keys()).index(ftype) + 1
            return ftype, idx
    return "III", 3


def assess_lighting(img_bgr: np.ndarray, l_skin_mean: float, ita_std: float) -> tuple:
    """Assess scene lighting quality for reliable Fitzpatrick estimation.

    Returns:
        (l_scene_mean, quality, warning, confidence_penalty)
    """
    lab_full = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2Lab).astype(np.float32)
    l_scene = float(np.mean(lab_full[:, :, 0]) * (100.0 / 255.0))

    warnings = []

    if l_scene < L_SCENE_MIN or l_skin_mean < L_SKIN_MIN:
        quality = "insufficient"
        warnings.append(
            f"Iluminacion insuficiente (L* escena={l_scene:.0f}, L* piel={l_skin_mean:.0f}). "
            "El tipo Fitzpatrick puede estar sobreestimado (piel aparenta mas oscura). "
            "Recapture con mejor iluminacion."
        )
        penalty = 0.15
    elif l_scene < L_SCENE_LOW or l_skin_mean < L_SKIN_SUSPICIOUS:
        quality = "low"
        warnings.append(
            f"Iluminacion suboptima (L* escena={l_scene:.0f}, L* piel={l_skin_mean:.0f}). "
            "El tipo Fitzpatrick podria estar 1-2 niveles sobreestimado."
        )
        penalty = 0.4
    else:
        quality = "good"
        penalty = 1.0

    # High ITA std indicates contaminated sample (shadows, wound edges)
    if ita_std > 20:
        if quality == "good":
            quality = "low"
        warnings.append(
            f"Alta variabilidad en la muestra de piel (ITA std={ita_std:.1f}). "
            "Posible contaminacion por sombras o bordes de herida."
        )
        penalty *= 0.5
    elif ita_std > 15:
        warnings.append(
            f"Variabilidad moderada (ITA std={ita_std:.1f}). "
            "Resultado puede tener +/- 1 nivel de error."
        )
        penalty *= 0.7

    warning = " ".join(warnings)
    return l_scene, quality, warning, penalty


def estimate_fitzpatrick(
    img_bgr: np.ndarray,
    masks: dict,
    periulcer_dilation_px: int = 60,
) -> FitzpatrickResult:
    """Estimate Fitzpatrick type from a DFU image using segmentation masks.

    Strategy:
        1. Healthy skin = foot region - dilated(perilesion + ulcer)
        2. Filter shadow/contamination pixels via L* outlier rejection
        3. Compute ITA with adaptive trimming
        4. Assess lighting quality and adjust confidence

    Args:
        img_bgr: BGR image (H, W, 3)
        masks: dict with keys 'foot', 'perilesion', 'ulcer' (bool arrays H, W)
        periulcer_dilation_px: Extra dilation around wound for safety margin (default 60)
    """
    h, w = img_bgr.shape[:2]
    foot = masks.get("foot", np.ones((h, w), dtype=bool))
    peri = masks.get("perilesion", np.zeros((h, w), dtype=bool))
    ulcer = masks.get("ulcer", np.zeros((h, w), dtype=bool))

    # Dilate ulcer+perilesion for safety margin (increased from 40 to 60)
    exclusion = (peri | ulcer).astype(np.uint8)
    if periulcer_dilation_px > 0:
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (periulcer_dilation_px, periulcer_dilation_px))
        exclusion = cv2.dilate(exclusion, kernel)
    exclusion = exclusion.astype(bool)

    # Healthy skin = foot minus exclusion zone
    healthy = foot & ~exclusion
    healthy_pixels = int(np.sum(healthy))

    if healthy_pixels < 100:
        healthy = foot & ~ulcer
        healthy_pixels = int(np.sum(healthy))

    if healthy_pixels < 50:
        return FitzpatrickResult(
            fitzpatrick_type="III", fitzpatrick_int=3,
            fitzpatrick_label="Intermediate",
            ita_angle=0.0, ita_std=0.0,
            l_skin_mean=0.0, b_skin_mean=0.0,
            healthy_pixels=healthy_pixels,
            healthy_ratio=healthy_pixels / (h * w),
            confidence=0.0,
            l_scene_mean=0.0,
            lighting_quality="insufficient",
            lighting_warning="Insuficientes pixeles de piel sana para estimar Fitzpatrick.",
        )

    # Convert to L*a*b*
    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2Lab).astype(np.float32)
    l_values = lab[healthy, 0] * (100.0 / 255.0)
    b_values = lab[healthy, 2] - 128.0

    # Filter shadow/contamination pixels BEFORE computing ITA
    l_filtered, b_filtered = filter_shadow_pixels(l_values, b_values)

    # Compute ITA with adaptive trimming
    ita_mean, ita_std = compute_ita(l_filtered, b_filtered)
    ftype, fint = classify_fitzpatrick(ita_mean)

    l_skin_mean = float(np.mean(l_filtered))
    b_skin_mean = float(np.mean(b_filtered))
    filtered_pixels = len(l_filtered)

    # Lighting quality assessment (now also considers ITA std)
    l_scene, lighting_quality, lighting_warning, lighting_penalty = assess_lighting(
        img_bgr, l_skin_mean, ita_std
    )

    # Confidence: pixel count + ITA consistency + coverage + lighting
    pixel_conf = min(filtered_pixels / 5000.0, 1.0)
    ita_conf = max(0.0, 1.0 - (ita_std / 25.0))
    coverage_conf = min((filtered_pixels / (h * w)) / 0.15, 1.0)
    base_confidence = pixel_conf * 0.3 + ita_conf * 0.4 + coverage_conf * 0.3
    confidence = base_confidence * lighting_penalty

    return FitzpatrickResult(
        fitzpatrick_type=ftype,
        fitzpatrick_int=fint,
        fitzpatrick_label=FITZPATRICK_LABELS[ftype],
        ita_angle=round(ita_mean, 2),
        ita_std=round(ita_std, 2),
        l_skin_mean=round(l_skin_mean, 2),
        b_skin_mean=round(b_skin_mean, 2),
        healthy_pixels=filtered_pixels,
        healthy_ratio=round(filtered_pixels / (h * w), 4),
        confidence=round(confidence, 3),
        l_scene_mean=round(l_scene, 2),
        lighting_quality=lighting_quality,
        lighting_warning=lighting_warning,
    )