src/aiflow_math_ink_05/grid.py

"""온라인 필기 stroke를 모델 독립적인 2차원 수식 셀로 묶는다."""

from __future__ import annotations

from dataclasses import dataclass
from statistics import median
from typing import Any, Sequence

from PIL import Image, ImageDraw


@dataclass(frozen=True, slots=True)
class GridConfig:
    """필요 변수: 기하 임계값과 입력 상한이다.
    작동 원리: 고정 최소값과 대표 획 높이 비율을 함께 사용해 기기 좌표계 변화에 대응한다.
    """

    baseline_tolerance: float = 24.0
    expression_gap: float = 72.0
    structure_gap: float = 52.0
    padding: float = 12.0
    max_strokes: int = 2_000
    max_points_per_stroke: int = 4_096
    max_total_points: int = 200_000
    coordinate_limit: float = 20_000.0


@dataclass(frozen=True, slots=True)
class StrokeBox:
    """필요 변수: 좌·상·우·하 경계다.
    작동 원리: 획과 수식 셀의 축 정렬 경계 및 병합 연산을 제공한다.
    """

    left: float
    top: float
    right: float
    bottom: float

    @property
    def width(self) -> float:
        """필요 변수: 좌우 좌표. 작동 원리: 음수가 아닌 폭을 반환한다."""
        return max(0.0, self.right - self.left)

    @property
    def height(self) -> float:
        """필요 변수: 상하 좌표. 작동 원리: 음수가 아닌 높이를 반환한다."""
        return max(0.0, self.bottom - self.top)

    @property
    def center_y(self) -> float:
        """필요 변수: 상하 좌표. 작동 원리: 같은 기준선 판정용 중심을 반환한다."""
        return (self.top + self.bottom) / 2.0

    def union(self, other: "StrokeBox") -> "StrokeBox":
        """필요 변수: 다른 경계. 작동 원리: 두 경계를 포함하는 최소 경계를 만든다."""
        return StrokeBox(
            min(self.left, other.left),
            min(self.top, other.top),
            max(self.right, other.right),
            max(self.bottom, other.bottom),
        )


@dataclass(frozen=True, slots=True)
class Stroke:
    """필요 변수: 안정적인 획 ID·순서 있는 점·선 굵기다.
    작동 원리: 선 반지름을 포함한 경계를 계산해 필기 폭을 보존한다.
    """

    stroke_id: int
    points: tuple[tuple[float, float], ...]
    width: float = 3.0

    @property
    def box(self) -> StrokeBox:
        """필요 변수: 필기점·굵기. 작동 원리: 선 반지름까지 포함한 경계를 계산한다."""
        radius = self.width / 2.0
        return StrokeBox(
            min(point[0] for point in self.points) - radius,
            min(point[1] for point in self.points) - radius,
            max(point[0] for point in self.points) + radius,
            max(point[1] for point in self.points) + radius,
        )


@dataclass(frozen=True, slots=True)
class FormulaGrid:
    """필요 변수: 셀 순번·경계·소속 획 ID다.
    작동 원리: 인식기와 무관한 결정적 수식 셀 계약을 표현한다.
    """

    index: int
    box: StrokeBox
    stroke_ids: tuple[int, ...]


class _DisjointSet:
    """필요 변수: 획 개수. 작동 원리: 경로 압축과 크기 기준 병합으로 연결 성분을 관리한다."""

    def __init__(self, size: int) -> None:
        """필요 변수: 원소 수. 작동 원리: 각 획을 독립 집합으로 초기화한다."""
        self.parent = list(range(size))
        self.size = [1] * size

    def find(self, value: int) -> int:
        """필요 변수: 획 인덱스. 작동 원리: 대표를 찾는 동안 경로를 압축한다."""
        while self.parent[value] != value:
            self.parent[value] = self.parent[self.parent[value]]
            value = self.parent[value]
        return value

    def union(self, first: int, second: int) -> None:
        """필요 변수: 두 획 인덱스. 작동 원리: 작은 트리를 큰 트리 아래에 결합한다."""
        root_first = self.find(first)
        root_second = self.find(second)
        if root_first == root_second:
            return
        if self.size[root_first] < self.size[root_second]:
            root_first, root_second = root_second, root_first
        self.parent[root_second] = root_first
        self.size[root_first] += self.size[root_second]


def parse_writing_events(value: Any, config: GridConfig | None = None) -> list[Stroke]:
    """필요 변수: JSON 유사 writing events와 선택 설정이다.
    작동 원리: 유한한 좌표·점·획만 상한 안에서 받아 안전한 불변 Stroke로 정규화한다.
    """
    settings = config or GridConfig()
    if not isinstance(value, list):
        return []
    total_points = 0
    strokes: list[Stroke] = []
    for fallback_id, item in enumerate(value[: settings.max_strokes]):
        if not isinstance(item, dict) or not isinstance(item.get("points"), list):
            continue
        points: list[tuple[float, float]] = []
        for point in item["points"][: settings.max_points_per_stroke]:
            if not isinstance(point, dict):
                continue
            try:
                x = float(point["x"])
                y = float(point["y"])
            except (KeyError, TypeError, ValueError):
                continue
            if abs(x) <= settings.coordinate_limit and abs(y) <= settings.coordinate_limit:
                points.append((x, y))
        if len(points) < 2:
            continue
        if total_points + len(points) > settings.max_total_points:
            break
        try:
            stroke_id = int(item.get("stroke_id", fallback_id))
            width = min(32.0, max(0.5, float(item.get("width", 3.0))))
        except (TypeError, ValueError):
            continue
        strokes.append(Stroke(stroke_id, tuple(points), width))
        total_points += len(points)
    return strokes


def build_formula_grids(
    strokes: Sequence[Stroke],
    config: GridConfig | None = None,
) -> list[FormulaGrid]:
    """필요 변수: 정규화된 획과 선택 설정이다.
    작동 원리: 기준선·근접 부착·구조선 관계 그래프의 연결 성분을 2차원 수식 셀로 만든다.
    """
    if not strokes:
        return []
    settings = config or GridConfig()
    boxes = [stroke.box for stroke in strokes]
    typical_height = median(box.height for box in boxes)
    baseline = max(settings.baseline_tolerance, typical_height * 0.50)
    expression_gap = max(settings.expression_gap, typical_height * 1.25)
    structure_gap = max(settings.structure_gap, typical_height * 0.80)
    attachment_gap = max(18.0, typical_height * 0.18)
    groups = _DisjointSet(len(strokes))

    for first in range(len(strokes)):
        for second in range(first + 1, len(strokes)):
            horizontal_gap = _axis_gap(
                boxes[first].left,
                boxes[first].right,
                boxes[second].left,
                boxes[second].right,
            )
            vertical_gap = _axis_gap(
                boxes[first].top,
                boxes[first].bottom,
                boxes[second].top,
                boxes[second].bottom,
            )
            center_gap = abs(boxes[first].center_y - boxes[second].center_y)
            same_line = center_gap <= baseline and horizontal_gap <= expression_gap
            attachment = (
                horizontal_gap <= attachment_gap
                and vertical_gap <= attachment_gap
                and center_gap
                <= max(boxes[first].height, boxes[second].height) * 1.35 + attachment_gap
            )
            if same_line or attachment:
                groups.union(first, second)

    _join_structural_bridges(boxes, groups, structure_gap)
    members: dict[int, list[int]] = {}
    for index in range(len(strokes)):
        members.setdefault(groups.find(index), []).append(index)

    raw: list[tuple[StrokeBox, tuple[int, ...]]] = []
    for indices in members.values():
        box = boxes[indices[0]]
        for index in indices[1:]:
            box = box.union(boxes[index])
        padded = StrokeBox(
            box.left - settings.padding,
            box.top - settings.padding,
            box.right + settings.padding,
            box.bottom + settings.padding,
        )
        raw.append((padded, tuple(sorted(strokes[index].stroke_id for index in indices))))

    raw.sort(key=lambda item: (round(item[0].top / max(1.0, baseline)), item[0].left))
    return [
        FormulaGrid(index=index, box=box, stroke_ids=stroke_ids)
        for index, (box, stroke_ids) in enumerate(raw)
    ]


def render_grid_images(
    image: Image.Image,
    strokes: Sequence[Stroke],
    grids: Sequence[FormulaGrid],
) -> list[Image.Image]:
    """필요 변수: 기준 캔버스·원본 획·수식 셀이다.
    작동 원리: 각 셀의 소속 획만 흰 배경에 재렌더링하고 셀 경계로 잘라 RGB 이미지로 반환한다.
    """
    stroke_by_id = {stroke.stroke_id: stroke for stroke in strokes}
    results: list[Image.Image] = []
    for grid in grids:
        canvas = Image.new("RGB", image.size, "white")
        draw = ImageDraw.Draw(canvas)
        for stroke_id in grid.stroke_ids:
            stroke = stroke_by_id.get(stroke_id)
            if stroke is not None:
                draw.line(
                    stroke.points,
                    fill="black",
                    width=max(1, round(stroke.width)),
                    joint="curve",
                )
        results.append(_crop(canvas, grid.box))
    return results


def _join_structural_bridges(
    boxes: Sequence[StrokeBox],
    groups: _DisjointSet,
    structure_gap: float,
) -> None:
    """필요 변수: 획 경계·집합 구조·구조 거리다.
    작동 원리: 긴 가로선은 분수선, 긴 세로선은 괄호·행렬 경계 후보로 보고 인접 획을 연결한다.
    """
    for bridge_index, bridge in enumerate(boxes):
        horizontal_bridge = bridge.width >= 30 and bridge.width >= bridge.height * 4
        vertical_bridge = bridge.height >= 45 and bridge.height >= bridge.width * 2.5
        if not horizontal_bridge and not vertical_bridge:
            continue
        for other_index, other in enumerate(boxes):
            if other_index == bridge_index:
                continue
            if horizontal_bridge:
                overlap = _overlap(bridge.left, bridge.right, other.left, other.right)
                denominator = max(1.0, min(bridge.width, other.width))
                if (
                    overlap / denominator >= 0.20
                    and _axis_gap(bridge.top, bridge.bottom, other.top, other.bottom)
                    <= structure_gap
                ):
                    groups.union(bridge_index, other_index)
            elif (
                bridge.top - 18 <= other.center_y <= bridge.bottom + 18
                and _axis_gap(bridge.left, bridge.right, other.left, other.right)
                <= structure_gap
            ):
                groups.union(bridge_index, other_index)


def _crop(image: Image.Image, box: StrokeBox) -> Image.Image:
    """필요 변수: 이미지와 셀 경계. 작동 원리: 경계를 이미지 안 정수 좌표로 제한해 안전하게 자른다."""
    left = max(0, min(image.width - 1, int(box.left)))
    top = max(0, min(image.height - 1, int(box.top)))
    right = max(left + 1, min(image.width, int(box.right + 0.999)))
    bottom = max(top + 1, min(image.height, int(box.bottom + 0.999)))
    return image.crop((left, top, right, bottom))


def _axis_gap(
    first_start: float,
    first_end: float,
    second_start: float,
    second_end: float,
) -> float:
    """필요 변수: 한 축의 두 구간. 작동 원리: 겹치면 0, 아니면 최근접 끝점 거리를 반환한다."""
    return max(0.0, max(first_start, second_start) - min(first_end, second_end))


def _overlap(
    first_start: float,
    first_end: float,
    second_start: float,
    second_end: float,
) -> float:
    """필요 변수: 한 축의 두 구간. 작동 원리: 두 구간의 교집합 길이를 반환한다."""
    return max(0.0, min(first_end, second_end) - max(first_start, second_start))