src/structure_relations.py

"""기호 인식과 독립적으로 수식의 2D 배치 관계를 추론하는 연구용 레이어다."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Sequence

import numpy as np

from .math_tray import infer_math_trays


RELATION_TYPES = ("RIGHT", "SUPERSCRIPT", "SUBSCRIPT", "ABOVE", "BELOW", "CONTAINS")
STRUCTURAL_RELATION_TYPES = frozenset({"SUPERSCRIPT", "SUBSCRIPT", "ABOVE", "BELOW", "CONTAINS"})


@dataclass(frozen=True, slots=True)
class RelationConfig:
    """필요 변수: 첨자 기하 임계값. 작동 원리: writer validation에서 고를 수 있도록 관계 판정값을 외부화한다."""

    script_min_height_ratio: float = 0.18
    script_max_height_ratio: float = 0.82
    script_vertical_shift_ratio: float = 0.18
    script_max_horizontal_gap_ratio: float = 0.75


def _label(segment: dict[str, Any]) -> str:
    """필요 변수: 인식된 segment. 작동 원리: 구조 힌트를 우선하고 첫 전문가의 top-1을 보조 라벨로 읽는다."""

    hint = segment.get("structural_hint")
    if hint:
        return str(hint)
    for result in (segment.get("results") or {}).values():
        candidates = result.get("candidates") or []
        if candidates:
            return str(candidates[0]["label"])
    return ""


def _geometry(segment: dict[str, Any]) -> dict[str, float]:
    """필요 변수: segment bbox. 작동 원리: 관계 판정에 필요한 폭·높이·중심 좌표를 한 번 계산한다."""

    box = segment["box"]
    left, top = float(box["left"]), float(box["top"])
    right, bottom = float(box["right"]), float(box["bottom"])
    return {
        "left": left, "top": top, "right": right, "bottom": bottom,
        "width": max(right - left, 1.0), "height": max(bottom - top, 1.0),
        "cx": (left + right) * 0.5, "cy": (top + bottom) * 0.5,
    }


def infer_spatial_relations(
    segments: Sequence[dict[str, Any]], config: RelationConfig | None = None,
    *, trays: Sequence[dict[str, Any]] | None = None,
) -> list[dict[str, Any]]:
    """필요 변수: geometry group과 선택적 OCR 결과. 작동 원리: 각 기호를 한 번만 자식으로 소비해 2D 관계 그래프를 만든다."""

    relation_config = config or RelationConfig()
    if len(segments) < 2:
        return []
    boxes = [_geometry(segment) for segment in segments]
    labels = [_label(segment) for segment in segments]
    maximum_height = max(box["height"] for box in boxes)
    body_heights = [box["height"] for box in boxes if box["height"] >= maximum_height * 0.25]
    reference_height = float(np.median(body_heights))
    relations: list[dict[str, Any]] = []
    assigned: set[int] = set()
    structural_parents: set[int] = set()

    # 근호는 내부 기호를 symbol group에 합치지 않고 CONTAINS edge로만 소유한다.
    for parent, (box, label) in enumerate(zip(boxes, labels, strict=True)):
        root_like = label in {r"\sqrt", r"\sqrt{}"}
        if not root_like:
            continue
        structural_parents.add(parent)
        for child, inner in enumerate(boxes):
            if child == parent or child in assigned:
                continue
            inside_x = box["left"] + box["width"] * 0.20 <= inner["cx"] <= box["right"] + reference_height * 0.15
            inside_y = inner["cy"] >= box["top"] and inner["top"] <= box["bottom"] + reference_height * 0.35
            if inside_x and inside_y:
                relations.append({"parent": parent, "child": child, "type": "CONTAINS", "confidence": 0.95})
                assigned.add(child)

    # 분수는 bar 단독이 아니라 분자·분모가 모두 점유된 원자적 Tray일 때만 관계를 만든다.
    tray_rows = list(trays) if trays is not None else infer_math_trays(segments)
    fraction_trays = [tray for tray in tray_rows if tray["type"] == "FRACTION"]
    for tray in fraction_trays:
        parent = int(tray["anchor"])
        structural_parents.add(parent)
        for slot, relation in (("numerator", "ABOVE"), ("denominator", "BELOW")):
            for child in tray["slots"].get(slot, []):
                if child in assigned:
                    continue
                relations.append({
                    "parent": parent, "child": child, "type": relation,
                    "confidence": float(tray["constraint_score"]),
                })
                assigned.add(child)

    # 적분·합·곱·극한의 상하한은 일반 최근접 첨자보다 수학 Tray의 anchor 소유권을 우선한다.
    operator_trays = [tray for tray in tray_rows if tray["type"] in {"INTEGRAL", "SUM", "PRODUCT", "LIM"}]
    for tray in operator_trays:
        parent = int(tray["anchor"])
        structural_parents.add(parent)
        for slot, relation in (("lower", "SUBSCRIPT"), ("upper", "SUPERSCRIPT")):
            for child in tray["slots"].get(slot, []):
                if child in assigned:
                    continue
                relations.append({
                    "parent": parent, "child": int(child), "type": relation,
                    "confidence": float(tray["constraint_score"]),
                })
                assigned.add(int(child))

    # 남은 기호는 가장 가까운 왼쪽 기호를 기준으로 첨자 또는 RIGHT 관계를 갖는다.
    for child, box in enumerate(boxes):
        if child in assigned or child in structural_parents:
            continue
        candidates = [index for index, other in enumerate(boxes) if index != child and other["cx"] < box["cx"]]
        if not candidates:
            continue
        parent = min(candidates, key=lambda index: max(0.0, box["left"] - boxes[index]["right"]))
        base = boxes[parent]
        horizontal_gap = max(0.0, box["left"] - base["right"])
        script_sized = (
            reference_height * relation_config.script_min_height_ratio
            <= box["height"]
            <= base["height"] * relation_config.script_max_height_ratio
        )
        script_near = horizontal_gap <= reference_height * relation_config.script_max_horizontal_gap_ratio
        vertical_shift = base["height"] * relation_config.script_vertical_shift_ratio
        if script_sized and script_near and box["cy"] < base["cy"] - vertical_shift:
            relation = "SUPERSCRIPT"
            confidence = 0.82
        elif script_sized and script_near and box["cy"] > base["cy"] + vertical_shift:
            relation = "SUBSCRIPT"
            confidence = 0.82
        else:
            relation = "RIGHT"
            confidence = 0.75
        relations.append({"parent": parent, "child": child, "type": relation, "confidence": confidence})
        assigned.add(child)

    return sorted(relations, key=lambda item: (item["child"], item["parent"], item["type"]))


def serialize_relation_graph(
    labels: Sequence[str], segments: Sequence[dict[str, Any]], relations: Sequence[dict[str, Any]],
) -> str:
    """필요 변수: group별 라벨·bbox·관계 edge. 작동 원리: 구조 자식을 한 번만 소비해 2D graph를 LaTeX로 직렬화한다."""

    if len(labels) != len(segments):
        raise ValueError("관계 serializer의 label과 segment 수가 다릅니다.")
    children: dict[int, dict[str, list[int]]] = {}
    structural_children: set[int] = set()
    for relation in relations:
        relation_type = str(relation["type"])
        if relation_type not in STRUCTURAL_RELATION_TYPES:
            continue
        parent, child = int(relation["parent"]), int(relation["child"])
        if parent < 0 or child < 0 or parent >= len(labels) or child >= len(labels) or parent == child:
            raise ValueError("관계 graph에 유효하지 않은 node index가 있습니다.")
        children.setdefault(parent, {}).setdefault(relation_type, []).append(child)
        structural_children.add(child)

    def left(index: int) -> float:
        """필요 변수: node index. 작동 원리: 같은 구조 슬롯 안의 기호를 원래 x 순서로 정렬한다."""

        return float(segments[index]["box"]["left"])

    def render_sequence(indices: Sequence[int], active: frozenset[int]) -> str:
        """필요 변수: node index 열·재귀 경로. 작동 원리: 같은 슬롯의 node를 좌→우로 한 번씩 출력한다."""

        return "".join(render_node(index, active) for index in sorted(set(indices), key=left))

    def render_node(index: int, active: frozenset[int]) -> str:
        """필요 변수: 현재 node·재귀 경로. 작동 원리: root/fraction/script 구조를 중첩 가능한 LaTeX node로 만든다."""

        if index in active:
            raise ValueError("관계 graph에 cycle이 있습니다.")
        next_active = active | {index}
        slots = children.get(index, {})
        above = slots.get("ABOVE", [])
        below = slots.get("BELOW", [])
        contained = slots.get("CONTAINS", [])
        if above and below:
            base = rf"\frac{{{render_sequence(above, next_active)}}}{{{render_sequence(below, next_active)}}}"
        elif contained and labels[index] in {r"\sqrt", r"\sqrt{}"}:
            base = rf"\sqrt{{{render_sequence(contained, next_active)}}}"
        else:
            base = labels[index]
        superscript = slots.get("SUPERSCRIPT", [])
        subscript = slots.get("SUBSCRIPT", [])
        if subscript:
            base += rf"_{{{render_sequence(subscript, next_active)}}}"
        if superscript:
            base += rf"^{{{render_sequence(superscript, next_active)}}}"
        return base

    roots = [index for index in range(len(labels)) if index not in structural_children]
    return render_sequence(roots, frozenset())