src/segmentation_lattice.py

"""획 grouping을 조기 확정하지 않고 복수 symbol 후보로 보존하는 연구용 lattice다."""

from __future__ import annotations

from collections import defaultdict
from typing import Any, Sequence
import math

import numpy as np


def _stroke_box(stroke: dict[str, Any]) -> tuple[float, float, float, float]:
    """필요 변수: point가 있는 stroke. 작동 원리: 후보 거리·정렬에 사용할 bbox를 계산한다."""

    points = stroke.get("points") or []
    xs = [float(point.get("x", 0.0) if isinstance(point, dict) else point[0]) for point in points]
    ys = [float(point.get("y", 0.0) if isinstance(point, dict) else point[1]) for point in points]
    if not xs:
        raise ValueError("빈 stroke는 segmentation lattice에 넣을 수 없습니다.")
    return min(xs), min(ys), max(xs), max(ys)


def _group_box(indices: frozenset[int], boxes: Sequence[tuple[float, float, float, float]]) -> dict[str, float]:
    """필요 변수: stroke index 집합·개별 bbox. 작동 원리: group 전체 bbox를 직렬화 가능한 값으로 합친다."""

    return {
        "left": min(boxes[index][0] for index in indices),
        "top": min(boxes[index][1] for index in indices),
        "right": max(boxes[index][2] for index in indices),
        "bottom": max(boxes[index][3] for index in indices),
    }


def _box_gap(first: tuple[float, float, float, float], second: tuple[float, float, float, float]) -> float:
    """필요 변수: 두 bbox. 작동 원리: 겹치면 0, 아니면 유클리드 외곽 간격을 반환한다."""

    dx = max(first[0] - second[2], second[0] - first[2], 0.0)
    dy = max(first[1] - second[3], second[1] - first[3], 0.0)
    return math.hypot(dx, dy)


def build_group_candidate_lattice(
    strokes: Sequence[dict[str, Any]],
    seed_partitions: Sequence[Sequence[dict[str, Any]]],
    *,
    temporal_window: int = 4,
    spatial_neighbors: int = 3,
) -> list[dict[str, Any]]:
    """필요 변수: 원 stroke·여러 geometry partition. 작동 원리: singleton/시간창/공간이웃/기존 group을 중복 제거해 보존한다."""

    usable = [(index, stroke) for index, stroke in enumerate(strokes) if stroke.get("points")]
    if not usable:
        return []
    source_indices = [index for index, _stroke in usable]
    stroke_by_index = {index: stroke for index, stroke in usable}
    boxes_by_source = {index: _stroke_box(stroke) for index, stroke in usable}
    dense_boxes = [boxes_by_source[index] for index in source_indices]
    evidence: dict[frozenset[int], set[str]] = defaultdict(set)

    def add(indices: frozenset[int], reason: str) -> None:
        if indices and indices.issubset(stroke_by_index):
            evidence[indices].add(reason)

    for index in source_indices:
        add(frozenset({index}), "singleton")
    for partition_index, partition in enumerate(seed_partitions):
        for segment in partition:
            add(frozenset(int(index) for index in segment["source_indices"]), f"geometry:{partition_index}")

    # 필기 순서가 인접한 다획 기호를 복구하되 지수 폭발을 막기 위해 짧은 연속 창만 만든다.
    for start in range(len(source_indices)):
        for length in range(2, min(temporal_window, len(source_indices) - start) + 1):
            add(frozenset(source_indices[start:start + length]), f"temporal:{length}")

    # 실제 기호 획이 시간상 끼어 쓰였을 때를 위해 각 획의 가까운 공간 이웃 pair를 추가한다.
    for dense_index, source_index in enumerate(source_indices):
        neighbor_order = sorted(
            (other for other in range(len(source_indices)) if other != dense_index),
            key=lambda other: (_box_gap(dense_boxes[dense_index], dense_boxes[other]), other),
        )[:spatial_neighbors]
        for other in neighbor_order:
            add(frozenset({source_index, source_indices[other]}), "spatial_pair")

    candidates = []
    for indices, reasons in evidence.items():
        box = _group_box(indices, boxes_by_source)
        candidates.append({
            "source_indices": sorted(indices),
            "strokes": [stroke_by_index[index] for index in sorted(indices)],
            "box": box,
            "evidence": sorted(reasons),
            "candidate_id": "g:" + ",".join(str(index) for index in sorted(indices)),
        })
    candidates.sort(key=lambda row: (row["box"]["left"], len(row["source_indices"]), row["source_indices"]))
    return candidates


def lattice_partition_oracle(
    candidates: Sequence[dict[str, Any]], truth_groups: Sequence[frozenset[int]],
) -> dict[str, Any]:
    """필요 변수: 후보 lattice·정답 partition. 작동 원리: 점수와 무관한 후보 coverage 상한을 계산한다."""

    candidate_groups = {frozenset(int(index) for index in row["source_indices"]) for row in candidates}
    matched = [group for group in truth_groups if group in candidate_groups]
    missing = [sorted(group) for group in truth_groups if group not in candidate_groups]
    return {
        "truth_groups": len(truth_groups),
        "matched_groups": len(matched),
        "group_recall": len(matched) / max(len(truth_groups), 1),
        "exact_partition_recoverable": not missing,
        "missing_groups": missing,
    }


LATTICE_FEATURE_NAMES = (
    "stroke_count", "point_count_log", "width_ref", "height_ref", "aspect_log",
    "temporal_span", "temporal_contiguous", "singleton", "geometry_votes",
    "has_temporal", "has_spatial_pair", "pair_gap_mean", "pair_gap_max",
    "stroke_width_mean", "stroke_width_std", "stroke_height_mean", "stroke_height_std",
)

LATTICE_OCR_FEATURE_NAMES = (
    "ocr_top1", "ocr_margin", "ocr_top5_mass", "ocr_entropy", "family_top1",
    "family_margin", "merge_top1_gain", "merge_entropy_gain", "label_alnum",
    "label_operator", "label_container", "label_delimiter", "label_dot_like",
)


def lattice_candidate_features(
    candidates: Sequence[dict[str, Any]], strokes: Sequence[dict[str, Any]],
) -> np.ndarray:
    """필요 변수: lattice 후보·원 stroke. 작동 원리: label 없이 공동 selector가 사용할 scale-normalized geometry feature를 만든다."""

    boxes = [_stroke_box(stroke) for stroke in strokes]
    nonzero_heights = [box[3] - box[1] for box in boxes if box[3] - box[1] > 1e-6]
    reference_height = float(np.median(nonzero_heights)) if nonzero_heights else 1.0
    rows = []
    for candidate in candidates:
        indices = [int(index) for index in candidate["source_indices"]]
        selected = [boxes[index] for index in indices]
        box = candidate["box"]
        width = float(box["right"] - box["left"])
        height = float(box["bottom"] - box["top"])
        gaps = [_box_gap(selected[first], selected[second]) / reference_height
                for first in range(len(selected)) for second in range(first + 1, len(selected))]
        widths = [(value[2] - value[0]) / reference_height for value in selected]
        heights = [(value[3] - value[1]) / reference_height for value in selected]
        evidence = set(candidate.get("evidence") or [])
        rows.append([
            len(indices), math.log1p(sum(len(strokes[index].get("points") or []) for index in indices)),
            width / reference_height, height / reference_height,
            math.log(max(width, 1e-6) / max(height, 1e-6)),
            max(indices) - min(indices) + 1, float(max(indices) - min(indices) + 1 == len(indices)),
            float(len(indices) == 1), sum(value.startswith("geometry:") for value in evidence),
            float(any(value.startswith("temporal:") for value in evidence)),
            float("spatial_pair" in evidence), float(np.mean(gaps)) if gaps else 0.0,
            max(gaps, default=0.0), float(np.mean(widths)), float(np.std(widths)),
            float(np.mean(heights)), float(np.std(heights)),
        ])
    return np.asarray(rows, dtype=np.float32)


def lattice_ocr_features(candidates: Sequence[dict[str, Any]]) -> np.ndarray:
    """필요 변수: OCR summary가 부착된 lattice. 작동 원리: confidence와 singleton 대비 결합 이득을 label-coarse feature로 바꾼다."""

    singleton = {
        int(candidate["source_indices"][0]): candidate["ocr_summary"]
        for candidate in candidates if len(candidate["source_indices"]) == 1
    }
    operators = {"+", "-", "=", "/", r"\times", r"\div", r"\pm", r"\neq", "<", ">"}
    containers = {r"\sqrt{}", r"\sqrt", r"\sum", r"\int", r"\prod"}
    delimiters = {"(", ")", "[", "]", r"\{", r"\}", r"\langle", r"\rangle", "|"}
    dot_like = {".", ",", r"\dots", r"\dotsc", r"\cdot"}
    rows = []
    for candidate in candidates:
        summary = candidate.get("ocr_summary")
        if not isinstance(summary, dict):
            raise ValueError("lattice OCR feature에는 ocr_summary가 필요합니다.")
        indices = [int(index) for index in candidate["source_indices"]]
        components = [singleton[index] for index in indices if index in singleton]
        component_top1 = float(np.mean([row["top1"] for row in components])) if components else 0.0
        component_entropy = float(np.mean([row["entropy"] for row in components])) if components else 0.0
        label = str(summary["top_label"])
        rows.append([
            float(summary["top1"]), float(summary["top1"] - summary["top2"]),
            float(summary["top5_mass"]), float(summary["entropy"]), float(summary["family_top1"]),
            float(summary["family_top1"] - summary["family_top2"]),
            float(summary["top1"] - component_top1) if len(indices) > 1 else 0.0,
            float(component_entropy - summary["entropy"]) if len(indices) > 1 else 0.0,
            float(label.isalnum()), float(label in operators), float(label in containers),
            float(label in delimiters), float(label in dot_like),
        ])
    return np.asarray(rows, dtype=np.float32)


def select_lattice_partition(
    candidates: Sequence[dict[str, Any]], scores: Sequence[float], stroke_count: int,
    *, beam_width: int = 128, group_bias: float = 0.0, options_per_stroke: int = 32,
    transition_scores: np.ndarray | None = None,
) -> list[frozenset[int]]:
    """필요 변수: 후보·logit score·획 수. 작동 원리: 모든 획을 정확히 한 번 덮는 최고 점수 partition을 beam exact-cover로 찾는다."""

    if len(candidates) != len(scores):
        raise ValueError("lattice 후보와 score 수가 다릅니다.")
    if stroke_count < 1:
        return []
    if transition_scores is not None and transition_scores.shape != (len(candidates), len(candidates)):
        raise ValueError("candidate transition score 행렬 크기가 다릅니다.")
    full_mask = (1 << stroke_count) - 1
    prepared = []
    by_stroke: dict[int, list[int]] = defaultdict(list)
    for candidate_index, (candidate, score) in enumerate(zip(candidates, scores, strict=True)):
        indices = tuple(int(index) for index in candidate["source_indices"])
        if not indices or any(index < 0 or index >= stroke_count for index in indices):
            continue
        mask = sum(1 << index for index in indices)
        prepared.append((candidate_index, mask, frozenset(indices), float(score) + group_bias))
        for index in indices:
            by_stroke[index].append(len(prepared) - 1)
    for index in by_stroke:
        by_stroke[index].sort(key=lambda row: prepared[row][3], reverse=True)
        by_stroke[index] = by_stroke[index][:options_per_stroke]
    if transition_scores is None:
        beams: dict[Any, tuple[float, tuple[frozenset[int], ...]]] = {0: (0.0, ())}
    else:
        beams = {(0, -1): (0.0, ())}
    for _step in range(stroke_count):
        expanded: dict[Any, tuple[float, tuple[frozenset[int], ...]]] = {}
        for state_key, (total_score, groups) in beams.items():
            used_mask, previous_candidate = (state_key, -1) if transition_scores is None else state_key
            if used_mask == full_mask:
                expanded[state_key] = max(expanded.get(state_key, (-math.inf, ())), (total_score, groups), key=lambda row: row[0])
                continue
            first_uncovered = next(index for index in range(stroke_count) if not used_mask & (1 << index))
            for prepared_index in by_stroke.get(first_uncovered, []):
                candidate_index, mask, group, score = prepared[prepared_index]
                if used_mask & mask:
                    continue
                new_mask = used_mask | mask
                transition = 0.0 if transition_scores is None or previous_candidate < 0 else float(
                    transition_scores[previous_candidate, candidate_index]
                )
                proposal = (total_score + score + transition, groups + (group,))
                new_key: Any = new_mask if transition_scores is None else (new_mask, candidate_index)
                if new_key not in expanded or proposal[0] > expanded[new_key][0]:
                    expanded[new_key] = proposal
        if not expanded:
            raise ValueError("모든 stroke를 덮는 lattice partition을 찾지 못했습니다.")
        beams = dict(sorted(expanded.items(), key=lambda row: row[1][0], reverse=True)[:beam_width])
        complete_keys = [key for key in beams if (key if transition_scores is None else key[0]) == full_mask]
        if complete_keys and all((key if transition_scores is None else key[0]) == full_mask for key in beams):
            break
    complete = [value for key, value in beams.items() if (key if transition_scores is None else key[0]) == full_mask]
    if not complete:
        raise ValueError("beam width 안에서 완전 partition을 찾지 못했습니다.")
    return list(max(complete, key=lambda row: row[0])[1])