app/src/feature_extractor.py

from __future__ import annotations

from pathlib import Path
from typing import Dict, Tuple

import numpy as np
from PIL import Image, ImageFilter, ImageStat, UnidentifiedImageError

from config import (
    ALMOST_BLACK_MEAN,
    ALMOST_WHITE_MEAN,
    LOW_STDDEV,
    MIN_IMAGE_SIZE,
)


def _laplacian_like_sharpness(gray: Image.Image) -> float:
    edges = gray.filter(ImageFilter.FIND_EDGES)
    arr = np.asarray(edges, dtype=np.float32)
    return float(arr.var())


def _edge_density(gray: Image.Image) -> float:
    edges = gray.filter(ImageFilter.FIND_EDGES)
    arr = np.asarray(edges, dtype=np.float32)
    threshold = arr.mean() + arr.std()
    if threshold <= 0:
        return 0.0
    density = (arr > threshold).mean()
    return float(density)


def _dominant_colors_count(img: Image.Image, max_colors: int = 12) -> int:
    small = img.convert("RGB").resize((64, 64))
    palette_img = small.quantize(colors=max_colors, method=Image.MEDIANCUT)
    colors = palette_img.getcolors()
    return len(colors) if colors else 0


def _whitespace_ratio(gray: Image.Image) -> float:
    arr = np.asarray(gray, dtype=np.float32)
    near_white = (arr > 240).mean()
    near_black = (arr < 15).mean()
    return float(max(near_white, near_black))


def _layout_density(gray: Image.Image) -> float:
    edges = gray.filter(ImageFilter.FIND_EDGES)
    arr = np.asarray(edges, dtype=np.float32)
    active = (arr > 30).mean()
    return float(active)


def _center_activity(gray: Image.Image) -> float:
    arr = np.asarray(gray.filter(ImageFilter.FIND_EDGES), dtype=np.float32)
    h, w = arr.shape
    y1, y2 = int(h * 0.25), int(h * 0.75)
    x1, x2 = int(w * 0.25), int(w * 0.75)
    center = arr[y1:y2, x1:x2]
    if center.size == 0:
        return 0.0
    return float((center > 30).mean())


def _grid_balance_3x3(gray: Image.Image) -> float:
    arr = np.asarray(gray.filter(ImageFilter.FIND_EDGES), dtype=np.float32)
    h, w = arr.shape
    ys = np.linspace(0, h, 4, dtype=int)
    xs = np.linspace(0, w, 4, dtype=int)

    cells = []
    for i in range(3):
        for j in range(3):
            cell = arr[ys[i]:ys[i + 1], xs[j]:xs[j + 1]]
            if cell.size == 0:
                cells.append(0.0)
            else:
                cells.append(float((cell > 30).mean()))

    mean_val = float(np.mean(cells))
    std_val = float(np.std(cells))
    balance = max(0.0, 1.0 - std_val / (mean_val + 1e-6))
    return balance


def inspect_image_content(image_path: Path) -> Tuple[bool, str]:
    try:
        with Image.open(image_path) as img:
            img.load()
            width, height = img.size

            if width < MIN_IMAGE_SIZE or height < MIN_IMAGE_SIZE:
                return False, f"too_small_{width}x{height}"

            gray = img.convert("L")
            extrema = gray.getextrema()
            if extrema is None:
                return False, "failed_extrema_check"

            if extrema[0] == extrema[1]:
                return False, "blank_uniform_image"

            stat = ImageStat.Stat(gray)
            mean_val = stat.mean[0]
            stddev = stat.stddev[0]

            if mean_val > ALMOST_WHITE_MEAN and stddev < LOW_STDDEV:
                return False, "almost_blank_white_image"

            if mean_val < ALMOST_BLACK_MEAN and stddev < LOW_STDDEV:
                return False, "almost_blank_black_image"

        return True, "ok"

    except UnidentifiedImageError:
        return False, "unidentified_image"
    except Exception as e:
        return False, f"image_inspection_error: {e}"


def extract_features(image_path: Path) -> Dict[str, float | int | bool | str]:
    has_content, reason = inspect_image_content(image_path)

    if not has_content:
        return {
            "content_present_rule": False,
            "blank_reason": reason,
            "mean_brightness": 0.0,
            "contrast": 0.0,
            "saturation_mean": 0.0,
            "dominant_colors_count": 0,
            "sharpness": 0.0,
            "edge_density": 0.0,
            "whitespace_ratio": 1.0,
            "layout_density": 0.0,
            "center_activity": 0.0,
            "grid_balance_3x3": 0.0,
        }

    with Image.open(image_path) as img:
        img = img.convert("RGB")

        gray = img.convert("L")
        hsv = img.convert("HSV")

        gray_stat = ImageStat.Stat(gray)
        hsv_stat = ImageStat.Stat(hsv)

        mean_brightness = float(gray_stat.mean[0]) / 255.0
        contrast = float(gray_stat.stddev[0]) / 64.0
        saturation_mean = float(hsv_stat.mean[1]) / 255.0

        dominant_colors_count = _dominant_colors_count(img)
        sharpness = _laplacian_like_sharpness(gray) / 1000.0
        edge_density = _edge_density(gray)
        whitespace_ratio = _whitespace_ratio(gray)
        layout_density = _layout_density(gray)
        center_activity = _center_activity(gray)
        grid_balance_3x3 = _grid_balance_3x3(gray)

        return {
            "content_present_rule": True,
            "blank_reason": "ok",
            "mean_brightness": mean_brightness,
            "contrast": contrast,
            "saturation_mean": saturation_mean,
            "dominant_colors_count": dominant_colors_count,
            "sharpness": sharpness,
            "edge_density": edge_density,
            "whitespace_ratio": whitespace_ratio,
            "layout_density": layout_density,
            "center_activity": center_activity,
            "grid_balance_3x3": grid_balance_3x3,
        }