src/tray_joint.py

"""OCR top-label과 Tray 슬롯 계약을 lattice candidate score로 변환한다."""

from __future__ import annotations

from dataclasses import dataclass
import math
from typing import Any, Sequence

import numpy as np

from .cross_visual import CrossVisualModel, cross_pair_feature_rows
from .expression_contracts import horizontal_sides, segments_complete
from .equality_visual import EqualityVisualModel, equality_pair_feature_rows
from .math_tray import infer_math_trays
from .segmentation_lattice import LATTICE_FEATURE_NAMES
from .structure_relations import STRUCTURAL_RELATION_TYPES, infer_spatial_relations


ANCHOR_TYPES = {
    r"\sqrt": "ROOT", r"\sqrt{}": "ROOT",
    r"\int": "INTEGRAL", r"\oint": "INTEGRAL", r"\iint": "INTEGRAL", r"\iiint": "INTEGRAL",
    r"\sum": "LARGE_OPERATOR", r"\prod": "LARGE_OPERATOR", r"\lim": "LARGE_OPERATOR",
}
RELIABLE_MULTISTROKE = {"i", r"\pi", r"\rightarrow", r"\pm", r"\leq", r"\geq"}
MULTISTROKE_FAMILY_LABELS = frozenset({
    r"\sum", r"\Sigma", r"\pi", r"\Pi", r"\rightarrow", r"\shortrightarrow",
    r"\longrightarrow", r"\Rightarrow", r"\neq",
    r"\not\equiv", r"\pm",
})
MULTISTROKE_FAMILY_NAMES = frozenset({
    r"\pi", r"\rightarrow", r"\shortrightarrow", r"\longrightarrow",
    r"\Rightarrow", r"\neq", r"\not\equiv", r"\pm",
})


def _multistroke_family_geometry_guard(
    label: str,
    family: str,
    feature_row: np.ndarray,
) -> bool:
    """필요 변수: 합친 OCR label/family·lattice geometry. 작동 원리: 완성 단일기호 열을 먹는 넓은 alias 후보를 차단한다."""

    width_ref = float(feature_row[2])
    height_ref = float(feature_row[3])
    pair_gap_max = float(feature_row[12])
    if label == r"\Sigma" and width_ref < 2.0:
        return False
    if (label == r"\pm" or family == r"\pm") and height_ref > 2.0:
        return False
    arrow_family = {
        r"\rightarrow", r"\shortrightarrow", r"\longrightarrow", r"\Rightarrow",
    }
    if (label in arrow_family or family in arrow_family) and pair_gap_max > 0.50:
        return False
    return True


@dataclass(frozen=True, slots=True)
class TrayJointWeights:
    """필요 변수: validation 선택 가중치. 작동 원리: 구조 보너스와 두 must-not-link 감점을 한 계약으로 고정한다."""

    tray: float = 4.0
    symbol: float = 4.0
    fraction: float = 8.0
    infix: float = 8.0
    competition: float = 0.0
    local_baseline: float = 0.0
    group_bias: float = -2.0


def component_competition_penalties(
    candidates: Sequence[dict[str, Any]],
    features: np.ndarray,
    *,
    mode: str = "joint",
) -> np.ndarray:
    """필요 변수: lattice 후보·OCR merge gain. 작동 원리: 합친 OCR이 구성획보다 약한 다획 후보에 must-not-link 값을 준다."""

    if len(candidates) != len(features):
        raise ValueError("component competition candidate/feature 수가 다릅니다.")
    if mode not in {"top1", "joint"}:
        raise ValueError(f"지원하지 않는 component competition mode입니다: {mode}")
    merge_index = len(LATTICE_FEATURE_NAMES) + 6
    entropy_index = len(LATTICE_FEATURE_NAMES) + 7
    penalties = np.zeros(len(candidates), dtype=np.float32)
    for index, candidate in enumerate(candidates):
        if len(candidate["source_indices"]) <= 1:
            continue
        merge_gain = float(features[index][merge_index])
        deficit = -merge_gain
        if mode == "joint":
            deficit -= 0.5 * float(features[index][entropy_index])
        penalties[index] = min(1.0, max(0.0, deficit))
    return penalties


def _group_box(group: frozenset[int], strokes: Sequence[dict[str, Any]]) -> dict[str, float]:
    """필요 변수: stroke group·좌표. 작동 원리: candidate와 context symbol의 bbox를 계산한다."""

    points = [point for index in group for point in strokes[index]["points"]]
    xs, ys = [float(point[0]) for point in points], [float(point[1]) for point in points]
    return {"left": min(xs), "top": min(ys), "right": max(xs), "bottom": max(ys)}


def _recognized_segment(box: dict[str, float], label: str, score: float, family: str = "") -> dict[str, Any]:
    """필요 변수: 후보 bbox·top label/family·확률. 작동 원리: Tray와 부분식 parser용 최소 OCR segment를 만든다."""

    return {"box": box, "results": {"expanded": {"candidates": [{
        "label": label, "score": score, "visual_family": family,
    }]}}}


def local_baseline_boundary_penalties(
    strokes: Sequence[dict[str, Any]],
    candidates: Sequence[dict[str, Any]],
    ocr_labels: Sequence[str],
    features: np.ndarray,
    base_partition: Sequence[frozenset[int]],
    *,
    ocr_families: Sequence[str] | None = None,
) -> np.ndarray:
    """필요 변수: 초기 partition·OCR segment·공간 관계. 작동 원리: 첨자·분수 등 구조 경계 양쪽을 동시에 먹는 병합만 감점한다."""

    count = len(candidates)
    if len(ocr_labels) != count or len(features) != count:
        raise ValueError("local baseline candidate/label/feature 수가 다릅니다.")
    families = list(ocr_families or [""] * count)
    if len(families) != count:
        raise ValueError("local baseline candidate/family 수가 다릅니다.")
    index_by_group = {
        frozenset(int(value) for value in candidate["source_indices"]): index
        for index, candidate in enumerate(candidates)
    }
    partition_groups = [frozenset(group) for group in base_partition if frozenset(group) in index_by_group]
    top1_index = len(LATTICE_FEATURE_NAMES)
    segments = []
    for group in partition_groups:
        index = index_by_group[group]
        segments.append(_recognized_segment(
            _group_box(group, strokes), ocr_labels[index],
            float(features[index][top1_index]), families[index],
        ))
    trays = infer_math_trays(segments)
    relations = infer_spatial_relations(segments, trays=trays)
    structural_edges = [
        (
            partition_groups[int(relation["parent"])],
            partition_groups[int(relation["child"])],
            float(relation["confidence"]),
        )
        for relation in relations
        if str(relation["type"]) in STRUCTURAL_RELATION_TYPES
    ]
    penalties = np.zeros(count, dtype=np.float32)
    for index, candidate in enumerate(candidates):
        group = frozenset(int(value) for value in candidate["source_indices"])
        if len(group) <= 1:
            continue
        penalties[index] = max(
            (
                confidence for parent, child, confidence in structural_edges
                if group & parent and group & child
            ),
            default=0.0,
        )
    return penalties


def _candidate_box(
    candidate: dict[str, Any], group: frozenset[int], strokes: Sequence[dict[str, Any]],
) -> dict[str, float]:
    """필요 변수: lattice 후보·원 stroke. 작동 원리: cache bbox를 우선 재사용하고 최소 입력은 계산으로 호환한다."""

    box = candidate.get("box")
    return dict(box) if isinstance(box, dict) else _group_box(group, strokes)


def _point_xy(point: Any) -> tuple[float, float]:
    """필요 변수: 배열 또는 객체형 point. 작동 원리: 획 형태 계산용 좌표를 공통 형식으로 읽는다."""

    if isinstance(point, dict):
        return float(point.get("x", 0.0)), float(point.get("y", 0.0))
    return float(point[0]), float(point[1])


def _stroke_direction(stroke: dict[str, Any]) -> tuple[float, float]:
    """필요 변수: 두 점 이상인 stroke. 작동 원리: 첫·끝점 벡터를 방향 정규화해 평행·교차 형태를 판정한다."""

    points = stroke.get("points") or []
    if len(points) < 2:
        return 0.0, 0.0
    start, end = _point_xy(points[0]), _point_xy(points[-1])
    dx, dy = end[0] - start[0], end[1] - start[1]
    length = math.hypot(dx, dy)
    return (dx / length, dy / length) if length > 1e-6 else (0.0, 0.0)


def _completed_operand_context_score(
    group: frozenset[int], group_box: dict[str, float], context: Sequence[tuple[frozenset[int], dict[str, Any]]],
) -> float:
    """필요 변수: 중위 후보 bbox·초기 partition. 작동 원리: 좌변·우변이 모두 완료된 표현일 때만 구조 점수를 연다."""

    usable = [(other_group, segment) for other_group, segment in context if not other_group & group]
    left, right = horizontal_sides(group_box, usable)
    return 1.0 if segments_complete(left) and segments_complete(right) else 0.0


def _occupied_operand_context_score(
    group: frozenset[int], group_box: dict[str, float], context: Sequence[tuple[frozenset[int], dict[str, Any]]],
) -> float:
    """필요 변수: 교차선 후보·현재 partition. 작동 원리: 곱셈은 OCR 문법을 강제하지 않고 같은 행의 양쪽 점유만 확인한다."""

    usable = [(other_group, segment) for other_group, segment in context if not other_group & group]
    left, right = horizontal_sides(group_box, usable)
    return 1.0 if left and right else 0.0


def _infix_shape_score(
    group: frozenset[int], strokes: Sequence[dict[str, Any]], stroke_boxes: Sequence[dict[str, float]],
) -> tuple[str, float]:
    """필요 변수: 두 획 후보·사전 계산 bbox. 작동 원리: 평행선 equality와 교차선 multiply family를 분리해 반환한다."""

    if len(group) != 2:
        return "", 0.0
    first_index, second_index = sorted(group)
    first, second = strokes[first_index], strokes[second_index]
    first_direction, second_direction = _stroke_direction(first), _stroke_direction(second)
    horizontal = min(abs(first_direction[0]), abs(second_direction[0]))
    vertical_leak = max(abs(first_direction[1]), abs(second_direction[1]))
    first_box, second_box = stroke_boxes[first_index], stroke_boxes[second_index]
    first_width = max(first_box["right"] - first_box["left"], 1e-6)
    second_width = max(second_box["right"] - second_box["left"], 1e-6)
    x_overlap = max(0.0, min(first_box["right"], second_box["right"]) - max(first_box["left"], second_box["left"]))
    overlap_ratio = x_overlap / min(first_width, second_width)
    length_ratio = min(first_width, second_width) / max(first_width, second_width)
    first_y = (first_box["top"] + first_box["bottom"]) * 0.5
    second_y = (second_box["top"] + second_box["bottom"]) * 0.5
    separation_ratio = abs(first_y - second_y) / max((first_width + second_width) * 0.5, 1e-6)
    parallel = abs(first_direction[0] * second_direction[0] + first_direction[1] * second_direction[1])
    if (
        horizontal >= 0.92 and vertical_leak <= 0.38 and parallel >= 0.94
        and overlap_ratio >= 0.60 and length_ratio >= 0.60 and 0.04 <= separation_ratio <= 0.65
    ):
        return "equality", min(1.0, 0.35 + 0.25 * overlap_ratio + 0.20 * length_ratio + 0.20 * parallel)

    diagonal = min(abs(first_direction[0] * first_direction[1]), abs(second_direction[0] * second_direction[1])) * 2.0
    opposite_slopes = first_direction[0] * first_direction[1] * second_direction[0] * second_direction[1] < 0.0
    x_intersects = max(first_box["left"], second_box["left"]) <= min(first_box["right"], second_box["right"])
    y_intersects = max(first_box["top"], second_box["top"]) <= min(first_box["bottom"], second_box["bottom"])
    if opposite_slopes and x_intersects and y_intersects and diagonal >= 0.55:
        return "multiply", min(1.0, diagonal)
    return "", 0.0


def candidate_signals(
    strokes: Sequence[dict[str, Any]], candidates: Sequence[dict[str, Any]],
    ocr_labels: Sequence[str], features: np.ndarray, base_partition: Sequence[frozenset[int]],
    *, ocr_families: Sequence[str] | None = None, strict_equality: bool = False,
    equality_model: EqualityVisualModel | None = None, cross_model: CrossVisualModel | None = None,
    cross_gap_ratio: float = 0.40,
    multistroke_family_boost: float = 6.0,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """필요 변수: 원 stroke·후보/label/feature·cross/family 설정. 작동 원리: Tray·신뢰 label·중위식 evidence를 계산한다."""

    count = len(candidates)
    if len(ocr_labels) != count or len(features) != count:
        raise ValueError("Tray joint candidate/label/feature 수가 다릅니다.")
    if cross_gap_ratio < 0.0:
        raise ValueError("cross gap ratio는 0 이상이어야 합니다.")
    if multistroke_family_boost < 0.0:
        raise ValueError("multistroke family boost는 0 이상이어야 합니다.")
    families = list(ocr_families or [""] * count)
    if len(families) != count:
        raise ValueError("Tray joint candidate/family 수가 다릅니다.")
    tray_signal = np.zeros(count, dtype=np.float32)
    symbol_signal = np.zeros(count, dtype=np.float32)
    infix_signal = np.zeros(count, dtype=np.float32)
    top1 = features[:, len(LATTICE_FEATURE_NAMES)]
    stroke_boxes = [_group_box(frozenset({index}), strokes) for index in range(len(strokes))]
    index_by_group = {
        frozenset(int(value) for value in candidate["source_indices"]): index
        for index, candidate in enumerate(candidates)
    }
    pair_groups = [
        frozenset(int(value) for value in candidate["source_indices"])
        for candidate in candidates if len(candidate["source_indices"]) == 2
    ]
    equality_probability: dict[frozenset[int], float] = {}
    cross_probability: dict[frozenset[int], float] = {}
    if equality_model is not None and pair_groups:
        ordered_pairs = [tuple(sorted(group)) for group in pair_groups]
        pair_rows = equality_pair_feature_rows(ordered_pairs, strokes)
        equality_probability = {
            group: equality_model.probability(row) for group, row in zip(pair_groups, pair_rows, strict=True)
        }
    if cross_model is not None and pair_groups:
        nonzero_heights = [box["bottom"] - box["top"] for box in stroke_boxes if box["bottom"] - box["top"] > 1e-5]
        reference = float(np.median(nonzero_heights)) if nonzero_heights else 1.0
        plausible_groups = []
        for group in pair_groups:
            first_index, second_index = sorted(group)
            first_box, second_box = stroke_boxes[first_index], stroke_boxes[second_index]
            x_gap = max(first_box["left"] - second_box["right"], second_box["left"] - first_box["right"], 0.0)
            y_gap = max(first_box["top"] - second_box["bottom"], second_box["top"] - first_box["bottom"], 0.0)
            if x_gap <= reference * cross_gap_ratio and y_gap <= reference * cross_gap_ratio:
                plausible_groups.append(group)
        ordered_pairs = [tuple(sorted(group)) for group in plausible_groups]
        pair_rows = cross_pair_feature_rows(ordered_pairs, strokes)
        cross_probability = {
            group: cross_model.probability(row) for group, row in zip(plausible_groups, pair_rows, strict=True)
        }
    context = [
        (
            group,
            _recognized_segment(
                _candidate_box(candidates[index_by_group[group]], group, strokes),
                ocr_labels[index_by_group[group]], float(top1[index_by_group[group]]), families[index_by_group[group]],
            ),
        )
        for group in base_partition
    ]
    for index, candidate in enumerate(candidates):
        group = frozenset(int(value) for value in candidate["source_indices"])
        label = ocr_labels[index]
        group_box = _candidate_box(candidate, group, strokes)
        shape_family, shape_score = _infix_shape_score(group, strokes, stroke_boxes)
        learned_equality = equality_probability.get(group, 0.0)
        learned_cross = cross_probability.get(group, 0.0)
        if equality_model is not None and learned_equality >= equality_model.threshold:
            shape_family, shape_score = "equality", learned_equality
        if cross_model is not None and learned_cross >= cross_model.threshold and learned_cross > learned_equality:
            # threshold 직상단의 불확실한 pair가 큰 병합 보너스를 받지 않도록 초과 margin만 사용한다.
            margin = (learned_cross - cross_model.threshold) / max(1.0 - cross_model.threshold, 1e-6)
            shape_family, shape_score = "cross_visual", min(1.0, max(0.0, margin))
        if shape_family == "equality":
            # 등호 인식은 수식 완성 여부와 독립적이다. parser는 shadow 분석에서만 gating한다.
            context_score = _completed_operand_context_score(group, group_box, context) if strict_equality else 1.0
        elif shape_family == "cross_visual":
            context_score = 1.0
        else:
            context_score = _occupied_operand_context_score(group, group_box, context)
        infix_signal[index] = shape_score * context_score
        if len(group) > 1:
            if label in RELIABLE_MULTISTROKE:
                symbol_signal[index] = float(top1[index])
            if (
                label in MULTISTROKE_FAMILY_LABELS
                or families[index] in MULTISTROKE_FAMILY_NAMES
            ) and _multistroke_family_geometry_guard(
                label, families[index], features[index],
            ):
                # 합친 OCR family가 이미 살아 있을 때만 추가한다. 인접 획이라는 이유만으로 병합하지 않는다.
                symbol_signal[index] += float(top1[index]) * multistroke_family_boost
        tray_type = ANCHOR_TYPES.get(label)
        if tray_type is None:
            continue
        anchor = _recognized_segment(_group_box(group, strokes), label, float(top1[index]), families[index])
        segments = [anchor, *(segment for other_group, segment in context if not other_group & group)]
        matches = (
            tray for tray in infer_math_trays(segments)
            if tray["anchor"] == 0 and tray["type"] == tray_type and tray["required_satisfied"]
        )
        tray_signal[index] = max((float(tray["constraint_score"]) for tray in matches), default=0.0)
    return tray_signal, symbol_signal, infix_signal


def adjusted_logits(
    base_logits: np.ndarray, tray_signal: np.ndarray, symbol_signal: np.ndarray,
    fraction_penalty: np.ndarray, weights: TrayJointWeights, infix_signal: np.ndarray | None = None,
    competition_penalty: np.ndarray | None = None,
    local_baseline_penalty: np.ndarray | None = None,
) -> np.ndarray:
    """필요 변수: base logit·구조 signal·세 penalty. 작동 원리: joint selector 목적함수의 candidate logit을 만든다."""

    effective_infix = np.zeros_like(base_logits) if infix_signal is None else infix_signal
    effective_competition = (
        np.zeros_like(base_logits) if competition_penalty is None else competition_penalty
    )
    effective_local_baseline = (
        np.zeros_like(base_logits) if local_baseline_penalty is None else local_baseline_penalty
    )
    arrays = (
        base_logits, tray_signal, symbol_signal, fraction_penalty,
        effective_infix, effective_competition, effective_local_baseline,
    )
    if len({len(values) for values in arrays}) != 1:
        raise ValueError("Tray joint logit과 signal 길이가 다릅니다.")
    return (
        base_logits + weights.tray * tray_signal + weights.symbol * symbol_signal
        + weights.infix * effective_infix - weights.fraction * fraction_penalty
        - weights.competition * effective_competition
        - weights.local_baseline * effective_local_baseline
    )