src/inference/age_progression.py

"""
Age-to-70 transformation — texture-overlay approach.

Key insight: we NEVER replace the person's face.
We only ADD aging artifacts on top:

  1. Real wrinkle texture extracted from old UTKFace images (high-pass filter)
     → adds actual wrinkle lines without changing face shape/colour
  2. Strong hair graying (top of frame)
  3. Natural age spots (small, skin-colour based)
  4. Under-eye darkening
  5. Mild skin desaturation

Speed: ~0.1 s on CPU (pure NumPy/OpenCV).
"""

from __future__ import annotations

import random
from pathlib import Path
from typing import Optional

import cv2
import numpy as np

_RNG = random.Random(42)


# ── skin mask ─────────────────────────────────────────────────────────────

def _skin_mask(bgr: np.ndarray) -> np.ndarray:
    """Float [0,1] soft mask for skin pixels."""
    ycrcb = cv2.cvtColor(bgr, cv2.COLOR_BGR2YCrCb)
    mask  = cv2.inRange(ycrcb,
                        np.array([0,  133, 77],  np.uint8),
                        np.array([255, 173, 127], np.uint8))
    hsv   = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
    mask2 = cv2.inRange(hsv,
                        np.array([0,  15,  50],  np.uint8),
                        np.array([35, 230, 255], np.uint8))
    combined = cv2.bitwise_or(mask, mask2)
    return cv2.GaussianBlur(combined, (21, 21), 0).astype(np.float32) / 255.0


# ── wrinkle texture extraction ────────────────────────────────────────────

def _extract_wrinkle_texture(old_bgr: np.ndarray, target_hw: tuple) -> np.ndarray:
    """
    Extract high-frequency wrinkle lines from an old face image.
    Returns a float32 array shaped (H, W) in range [-1, 1].
    Positive values = darker (shadow/crease), negative = lighter (ridge).
    """
    gray    = cv2.cvtColor(old_bgr, cv2.COLOR_BGR2GRAY).astype(np.float32)
    # Resize texture source to target size
    gray    = cv2.resize(gray, (target_hw[1], target_hw[0]))
    # High-pass: subtract blurred version → keeps only wrinkle lines
    blurred = cv2.GaussianBlur(gray, (0, 0), sigmaX=4)
    hp      = gray - blurred          # range roughly -80 to +80
    # Normalise to [-1, 1]
    norm    = hp / (np.abs(hp).max() + 1e-6)
    return norm.astype(np.float32)


# ── reference cache ───────────────────────────────────────────────────────
_ref_cache: dict = {}   # gender → list of bgr arrays

def _load_refs(gender: int, utkface_dir: str, n: int = 40) -> list:
    if gender in _ref_cache:
        return _ref_cache[gender]

    # Try bundled old_faces first (works on HF Spaces without full dataset)
    base = Path(utkface_dir).parent
    candidates = [
        base / "old_faces",
        Path("data/old_faces"),
        Path(utkface_dir),
    ]

    refs = []
    for search_dir in candidates:
        if not search_dir.exists():
            continue
        for p in sorted(search_dir.glob("*.jpg")):
            parts = p.stem.split("_")
            if len(parts) < 3:
                continue
            try:
                age, g = int(parts[0]), int(parts[1])
            except ValueError:
                continue
            if age < 63 or g != gender:
                continue
            bgr = cv2.imread(str(p))
            if bgr is not None:
                refs.append(bgr)
            if len(refs) >= n:
                break
        if refs:
            break

    _ref_cache[gender] = refs
    return refs


# ── forehead lines ────────────────────────────────────────────────────────

def _add_forehead_lines(bgr: np.ndarray, skin: np.ndarray, alpha: float) -> np.ndarray:
    """
    Draw horizontal forehead wrinkle lines using the face's OWN skin colour
    (darkened) so they always look natural regardless of skin tone.
    """
    h, w  = bgr.shape[:2]
    result = bgr.copy().astype(np.float32)

    # Sample average forehead skin colour
    fh_patch = bgr[int(h*0.18):int(h*0.33), int(w*0.25):int(w*0.75)]
    if fh_patch.size == 0:
        return bgr
    base_color = fh_patch.reshape(-1, 3).mean(axis=0).astype(np.float32)
    # Crease colour = 28-35 % darker than skin
    crease = base_color * _RNG.uniform(0.62, 0.70)

    n_lines = 4
    for i in range(n_lines):
        y_frac = 0.20 + i * 0.036
        cy     = int(h * y_frac)
        line_layer = result.copy()
        pts = []
        for x in range(int(w * 0.12), int(w * 0.88), max(1, w // 24)):
            wave = int(h * 0.006 * np.sin(x / w * np.pi * 2.5 + i * 0.9))
            pts.append((x, cy + wave))
        for j in range(len(pts) - 1):
            cv2.line(line_layer.astype(np.uint8),
                     pts[j], pts[j+1],
                     (int(crease[0]), int(crease[1]), int(crease[2])),
                     1, cv2.LINE_AA)
        # Blend line into result weighted by skin mask and alpha
        # blur line layer
        blurred = cv2.GaussianBlur(line_layer, (5, 5), 0)
        line_weight = alpha * 0.60 * (0.6 + 0.1 * i)
        result = result * (1 - line_weight) + blurred * line_weight

    # Crow's feet (lateral eye corners)
    eye_y  = int(h * 0.43)
    for side, ox in [(-1, int(w * 0.14)), (1, int(w * 0.86))]:
        for angle_deg in np.linspace(155, 205, 4) if side == -1 else np.linspace(-25, 25, 4):
            rad    = np.deg2rad(angle_deg)
            length = int(w * 0.09 * _RNG.uniform(0.7, 1.2))
            x2     = int(ox + length * np.cos(rad))
            y2     = int(eye_y + length * np.sin(rad))
            x2, y2 = np.clip(x2, 0, w-1), np.clip(y2, 0, h-1)
            cv2.line(result.astype(np.uint8), (ox, eye_y), (x2, y2),
                     (int(crease[0]), int(crease[1]), int(crease[2])), 1, cv2.LINE_AA)

    result = cv2.GaussianBlur(result.astype(np.uint8), (3, 3), 0).astype(np.float32)
    return np.clip(result, 0, 255).astype(np.uint8)


# ── step 1: wrinkle overlay ───────────────────────────────────────────────

def _apply_wrinkles(face_bgr: np.ndarray, gender: int,
                    utkface_dir: str, alpha: float) -> np.ndarray:
    """
    Blend real wrinkle texture from old faces onto the input face.
    Only the high-frequency (line) information is transferred —
    the face colour and structure are completely preserved.
    """
    refs = _load_refs(gender, utkface_dir)
    if not refs:
        return face_bgr.copy()

    h, w = face_bgr.shape[:2]

    # Average 3 random old-face textures → smoother, less person-specific
    textures = [_extract_wrinkle_texture(_RNG.choice(refs), (h, w))
                for _ in range(min(3, len(refs)))]
    texture  = np.mean(textures, axis=0)

    skin  = _skin_mask(face_bgr)
    # Focus on forehead + cheeks; exclude eyes (landmarks 36-47) and mouth
    # Heuristic: blank out lower 15 % (chin) and top 15 % (hair)
    roi = np.ones((h, w), dtype=np.float32)
    roi[:int(h * 0.14), :] = 0     # hair top
    roi[int(h * 0.85):, :] = 0     # chin
    # eye strip — less wrinkle blending (lids distort texture)
    roi[int(h * 0.36):int(h * 0.52), int(w * 0.18):int(w * 0.82)] *= 0.3

    blend_weight = skin * roi * alpha * 0.75

    result = face_bgr.astype(np.float32)
    for c in range(3):
        # Darken crease shadows, brighten ridges (realistic depth)
        result[:, :, c] -= texture * blend_weight * 48
    result = np.clip(result, 0, 255).astype(np.uint8)

    # ── add explicit forehead horizontal lines ──────────────────────────
    result = _add_forehead_lines(result, skin, alpha)
    return result


# ── step 2: hair graying ──────────────────────────────────────────────────

def _gray_hair(bgr: np.ndarray, alpha: float) -> np.ndarray:
    h, w  = bgr.shape[:2]
    band  = np.zeros((h, w), dtype=np.float32)
    hair_h = int(h * 0.22)
    band[:hair_h, :] = 1.0
    fade = np.linspace(1.0, 0.0, max(1, int(h * 0.07)), dtype=np.float32)
    if len(fade) <= hair_h:
        band[hair_h - len(fade):hair_h, :] *= fade[:, None]

    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
    hsv[:, :, 1] *= 1.0 - band * alpha * 0.95          # desaturate
    # Lift brightness toward gray but cap at 210 (avoids over-brightening)
    v_target = np.minimum(hsv[:, :, 2] + band * alpha * 60, 210)
    hsv[:, :, 2] = hsv[:, :, 2] * (1 - band * alpha) + v_target * (band * alpha)
    return cv2.cvtColor(np.clip(hsv, 0, 255).astype(np.uint8), cv2.COLOR_HSV2BGR)


# ── step 3: age spots ─────────────────────────────────────────────────────

def _add_spots(bgr: np.ndarray, skin: np.ndarray, alpha: float) -> np.ndarray:
    h, w   = bgr.shape[:2]
    result = bgr.astype(np.float32)
    n_spots = int(10 * alpha)
    for _ in range(n_spots):
        cy = int(_RNG.uniform(0.28, 0.78) * h)
        cx = int(_RNG.uniform(0.10, 0.90) * w)
        if cy >= h or cx >= w or skin[cy, cx] < 0.25:
            continue
        # Spot colour: 25-35 % darker than local skin
        local = result[max(0, cy-3):cy+3, max(0, cx-3):cx+3].mean(axis=(0, 1))
        spot_col = np.clip(local * _RNG.uniform(0.62, 0.74), 0, 255)
        rx = int(_RNG.uniform(2, 6));  ry = int(_RNG.uniform(1, 4))
        strength = _RNG.uniform(0.25, 0.55) * alpha
        m  = np.zeros((h, w), dtype=np.float32)
        cv2.ellipse(m, (cx, cy), (rx, ry), _RNG.uniform(0, 180), 0, 360, 1.0, -1)
        m  = cv2.GaussianBlur(m, (5, 5), 0)
        for c in range(3):
            result[:, :, c] = result[:, :, c] * (1 - m * strength) + spot_col[c] * m * strength
    return np.clip(result, 0, 255).astype(np.uint8)


# ── step 4: skin aging ────────────────────────────────────────────────────

def _age_skin(bgr: np.ndarray, skin: np.ndarray, alpha: float) -> np.ndarray:
    """Desaturate + warm-yellow shift + slight value drop."""
    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
    hsv[:, :, 1] *= 1.0 - skin * alpha * 0.35
    hsv[:, :, 2] *= 1.0 - skin * alpha * 0.06
    bgr2 = cv2.cvtColor(np.clip(hsv, 0, 255).astype(np.uint8), cv2.COLOR_HSV2BGR)
    # Warm yellow cast
    f = bgr2.astype(np.float32)
    f[:, :, 2] = np.clip(f[:, :, 2] + skin * alpha * 10, 0, 255)   # R up
    f[:, :, 0] = np.clip(f[:, :, 0] - skin * alpha * 5,  0, 255)   # B down
    return np.clip(f, 0, 255).astype(np.uint8)


# ── step 5: under-eye bags ────────────────────────────────────────────────

def _undereye(bgr: np.ndarray, alpha: float) -> np.ndarray:
    h, w = bgr.shape[:2]
    ys   = np.linspace(0, 1, h)[:, None]
    xs   = np.linspace(0, 1, w)[None, :]
    mask = (np.exp(-(((ys-0.52)/0.055)**2 + ((xs-0.30)/0.12)**2))
           +np.exp(-(((ys-0.52)/0.055)**2 + ((xs-0.70)/0.12)**2)))
    mask = np.clip(mask, 0, 1).astype(np.float32) * alpha * 0.35
    f    = bgr.astype(np.float32)
    f[:, :, 0] -= mask * 25   # darken all channels; more on B for purple tinge
    f[:, :, 1] -= mask * 20
    f[:, :, 2] -= mask * 12
    return np.clip(f, 0, 255).astype(np.uint8)


# ── main ─────────────────────────────────────────────────────────────────

def age_to_70(face_rgb: np.ndarray,
              current_age: float = 30.0,
              gender: int = 0,
              utkface_dir: str = "data/UTKFace") -> np.ndarray:
    """
    Age face_rgb (H×W×3 uint8 RGB) toward age 70.
    Preserves the person's identity — only adds aging texture/colour on top.
    """
    target = 70.0
    delta  = max(0.0, target - current_age)
    alpha  = float(np.clip(delta / 38.0, 0.0, 1.0))

    if alpha < 0.05:
        return face_rgb.copy()

    orig_hw = face_rgb.shape[:2]

    # Work at ≥ 256 px
    h, w = orig_hw
    if min(h, w) < 256:
        s = 256 / min(h, w)
        face_rgb = cv2.resize(face_rgb, (int(w * s), int(h * s)))
        h, w = face_rgb.shape[:2]

    bgr  = cv2.cvtColor(face_rgb, cv2.COLOR_RGB2BGR)
    skin = _skin_mask(bgr)

    # Pipeline
    bgr = _apply_wrinkles(bgr, gender, utkface_dir, alpha)
    bgr = _gray_hair(bgr, alpha)
    bgr = _add_spots(bgr, skin, alpha)
    bgr = _age_skin(bgr, skin, alpha)
    bgr = _undereye(bgr, alpha)

    if bgr.shape[:2] != orig_hw:
        bgr = cv2.resize(bgr, (orig_hw[1], orig_hw[0]))

    return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)