src/p_formula_dataset06.py

"""P Formula v1의 실제 symbol group을 formula-relative 128×19 학습 tensor로 만든다."""

from __future__ import annotations

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

import torch
from torch import Tensor

from .ink06_canonical import CanonicalInk06, canonicalize_ink06
from .trajectory_sequence import visual_label_family


@dataclass(frozen=True, slots=True)
class PFormulaTensorBatch06:
    """필요 변수: symbol tensor·정답·identity·결측 flag. 작동 원리: metric 분모를 잃지 않는 split batch를 보존한다."""

    features: Tensor
    truths: tuple[str, ...]
    writer_ids: tuple[str, ...]
    device_ids: tuple[str, ...]
    source_ids: tuple[str, ...]
    timestamp_missing: Tensor
    pressure_missing: Tensor


def _formula_box06(record: dict[str, Any]) -> tuple[float, float, float, float]:
    """필요 변수: 한 P Formula record. 작동 원리: 모든 symbol stroke의 실제 ink bbox를 계산한다."""

    points = [
        point
        for symbol in record["symbols"]
        for stroke in symbol["strokes"]
        for point in stroke
    ]
    if not points:
        raise ValueError(f"formula {record.get('formula_id')}에 point가 없습니다.")
    try:
        xs = [float(point["x"]) for point in points]
        ys = [float(point["y"]) for point in points]
    except (KeyError, TypeError, ValueError) as error:
        raise ValueError("P Formula point x/y가 유효하지 않습니다.") from error
    if not all(math.isfinite(value) for value in (*xs, *ys)):
        raise ValueError("P Formula point에 비유한 좌표가 있습니다.")
    return min(xs), min(ys), max(xs), max(ys)


def p_formula_symbol_ink06(
    symbol: dict[str, Any],
    *,
    formula_box: tuple[float, float, float, float],
) -> CanonicalInk06:
    """필요 변수: 정답 symbol·전체 formula bbox. 작동 원리: 원본 필순과 formula-relative 위치를 canonical ink로 보존한다."""

    left, top, right, bottom = formula_box
    width = max(right - left, 1e-5)
    height = max(bottom - top, 1e-5)
    strokes = []
    observed_timestamps = True
    for order, raw_stroke in enumerate(symbol["strokes"]):
        if not isinstance(raw_stroke, list) or not raw_stroke:
            raise ValueError("P Formula symbol stroke가 비어 있습니다.")
        points = []
        for point in raw_stroke:
            timestamp = point.get("t")
            observed_timestamps = observed_timestamps and timestamp is not None
            points.append({
                "x": float(point["x"]) - left,
                "y": float(point["y"]) - top,
                "t_ms": None if timestamp is None else float(timestamp),
            })
        strokes.append({"stroke_id": order, "order": order, "points": points})
    return canonicalize_ink06(
        strokes,
        canvas_width=width,
        canvas_height=height,
        source_modality="online",
        trust_timestamps=observed_timestamps,
    )


def p_formula_symbol_feature06(
    symbol: dict[str, Any],
    *,
    formula_box: tuple[float, float, float, float],
) -> Tensor:
    """필요 변수: 정답 symbol·전체 formula bbox. 작동 원리: canonical ink의 128×19 feature를 학습 tensor로 변환한다."""

    ink = p_formula_symbol_ink06(symbol, formula_box=formula_box)
    return torch.from_numpy(ink.features)


def materialize_p_formula_split06(
    records: Sequence[dict[str, Any]],
    *,
    allowed_labels: Sequence[str],
) -> PFormulaTensorBatch06:
    """필요 변수: 한 split P Formula records·378 vocabulary. 작동 원리: truth group별 tensor와 identity/결측 slice를 함께 만든다."""

    allowed = set(str(label) for label in allowed_labels)
    if len(allowed) != len(allowed_labels):
        raise ValueError("P Formula vocabulary에 중복 label이 있습니다.")
    features: list[Tensor] = []
    truths: list[str] = []
    writers: list[str] = []
    devices: list[str] = []
    sources: list[str] = []
    timestamp_missing: list[bool] = []
    pressure_missing: list[bool] = []
    for record in records:
        formula_box = _formula_box06(record)
        writer = str(record["writer_id"])
        device = str(record["device_id"])
        source = str(record["source_id"])
        for symbol in record["symbols"]:
            label = str(symbol["token"])
            if label not in allowed:
                raise ValueError(f"0.6 vocabulary 밖 P Formula token입니다: {label}")
            points = [point for stroke in symbol["strokes"] for point in stroke]
            features.append(
                p_formula_symbol_feature06(symbol, formula_box=formula_box),
            )
            truths.append(label)
            writers.append(writer)
            devices.append(device)
            sources.append(source)
            timestamp_missing.append(any(point.get("t") is None for point in points))
            pressure_missing.append(any(point.get("pressure") is None for point in points))
    if not features:
        raise ValueError("P Formula split에 지원 symbol이 없습니다.")
    return PFormulaTensorBatch06(
        features=torch.stack(features),
        truths=tuple(truths),
        writer_ids=tuple(writers),
        device_ids=tuple(devices),
        source_ids=tuple(sources),
        timestamp_missing=torch.tensor(timestamp_missing, dtype=torch.bool),
        pressure_missing=torch.tensor(pressure_missing, dtype=torch.bool),
    )


def p_formula_release_metrics06(
    exact_logits: Tensor,
    exact_targets: Tensor,
    *,
    labels: Sequence[str],
    writer_ids: Sequence[str],
    source_ids: Sequence[str],
    timestamp_missing: Tensor,
    pressure_missing: Tensor,
) -> dict[str, Any]:
    """필요 변수: exact logit/정답·identity·결측 mask. 작동 원리: top-k·macro-F1·writer/source floor·결측 하락을 같은 분모로 계산한다."""

    samples = len(exact_targets)
    if (
        exact_logits.ndim != 2
        or exact_logits.shape[0] != samples
        or len(writer_ids) != samples
        or len(source_ids) != samples
        or len(timestamp_missing) != samples
        or len(pressure_missing) != samples
    ):
        raise ValueError("P Formula release metric 분모가 서로 다릅니다.")
    top_k = min(5, exact_logits.shape[1])
    top = exact_logits.topk(top_k, dim=1).indices
    prediction = top[:, 0]
    correct = prediction.eq(exact_targets)
    top5_correct = top.eq(exact_targets.unsqueeze(1)).any(dim=1)
    visual_correct = torch.tensor([
        visual_label_family(str(labels[int(predicted)]))
        == visual_label_family(str(labels[int(truth)]))
        for predicted, truth in zip(
            prediction.tolist(),
            exact_targets.tolist(),
            strict=True,
        )
    ])
    supported = sorted(set(exact_targets.tolist()))
    f1_values = []
    for target in supported:
        truth_mask = exact_targets == target
        predicted_mask = prediction == target
        true_positive = int((truth_mask & predicted_mask).sum())
        false_positive = int((~truth_mask & predicted_mask).sum())
        false_negative = int((truth_mask & ~predicted_mask).sum())
        precision = true_positive / max(true_positive + false_positive, 1)
        recall = true_positive / max(true_positive + false_negative, 1)
        f1_values.append(
            2 * precision * recall / max(precision + recall, 1e-12),
        )

    def identity_slice06(values: Sequence[str]) -> dict[str, Any]:
        """필요 변수: writer 또는 source ID. 작동 원리: identity별 top-1과 floor/p10을 집계한다."""

        indices: dict[str, list[int]] = {}
        for index, value in enumerate(values):
            indices.setdefault(str(value), []).append(index)
        accuracies = {
            identity: float(correct[rows].float().mean())
            for identity, rows in indices.items()
        }
        ordered = sorted(accuracies.values())
        p10_index = min(int(len(ordered) * 0.10), max(len(ordered) - 1, 0))
        return {
            "count": len(accuracies),
            "floor": min(ordered) if ordered else 0.0,
            "p10": ordered[p10_index] if ordered else 0.0,
            "accuracies": dict(sorted(accuracies.items())),
        }

    overall = float(correct.float().mean())

    def missing_slice06(mask: Tensor) -> dict[str, float | int | None]:
        """필요 변수: symbol별 결측 mask. 작동 원리: 결측 분모와 전체 대비 top-1 하락을 분리한다."""

        count = int(mask.sum())
        if not count:
            return {"samples": 0, "top1": None, "drop_pp": None}
        accuracy = float(correct[mask].float().mean())
        return {
            "samples": count,
            "top1": accuracy,
            "drop_pp": (overall - accuracy) * 100.0,
        }

    return {
        "samples": samples,
        "exact_top1": overall,
        "exact_top5": float(top5_correct.float().mean()),
        "visual_family_top1": float(visual_correct.float().mean()),
        "macro_f1": sum(f1_values) / max(len(f1_values), 1),
        "supported_labels": len(supported),
        "writer": identity_slice06(writer_ids),
        "source": identity_slice06(source_ids),
        "missing_slices": {
            "timestamp": missing_slice06(timestamp_missing),
            "pressure": missing_slice06(pressure_missing),
        },
    }


def p_formula_seed_gate06(
    metrics: dict[str, Any],
    *,
    top1_minimum: float = 0.92,
    top5_minimum: float = 0.99,
    macro_f1_minimum: float = 0.90,
    writer_floor_minimum: float = 0.75,
    missing_drop_maximum_pp: float = 3.0,
) -> dict[str, Any]:
    """필요 변수: test release metric·명세 threshold. 작동 원리: 모든 정확도와 결측 slice를 AND gate로 판정한다."""

    missing_pass = all(
        row["drop_pp"] is None
        or float(row["drop_pp"]) <= missing_drop_maximum_pp
        for row in metrics["missing_slices"].values()
    )
    checks = {
        "exact_top1": float(metrics["exact_top1"]) >= top1_minimum,
        "exact_top5": float(metrics["exact_top5"]) >= top5_minimum,
        "macro_f1": float(metrics["macro_f1"]) >= macro_f1_minimum,
        "writer_floor": float(metrics["writer"]["floor"]) >= writer_floor_minimum,
        "missing_slices": missing_pass,
    }
    return {
        "checks": checks,
        "passed": all(checks.values()),
        "thresholds": {
            "exact_top1": top1_minimum,
            "exact_top5": top5_minimum,
            "macro_f1": macro_f1_minimum,
            "writer_floor": writer_floor_minimum,
            "missing_drop_maximum_pp": missing_drop_maximum_pp,
        },
    }