scripts/train_math_ink_06_behavior_role.py

"""CROHME 연속 수식에서 AIFlow 0.6 행동 의미 역할 head를 GPU 학습한다."""

from __future__ import annotations

import argparse
from collections import Counter
from copy import deepcopy
from datetime import datetime, timezone
import json
import math
from pathlib import Path
import random
import sys
from typing import Any, Sequence

import numpy as np
import torch
from torch import Tensor
from torch.utils.data import DataLoader, TensorDataset

PROJECT_ROOT = Path(__file__).parents[1]
SOURCE_ROOT = PROJECT_ROOT / "src"
for path in (PROJECT_ROOT, SOURCE_ROOT):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from math_grid_drawer.research.behavior_role_head06 import (
    BEHAVIOR_CONTEXT_FEATURES06,
    BEHAVIOR_ROLE_LABELS06,
    BehaviorRoleHead06,
    behavior_role_index06,
    validate_behavior_context06,
)
from math_grid_drawer.research.cross_visual import CROSS_FEATURE_NAMES, cross_pair_feature_rows
from math_grid_drawer.research.ink06_canonical import canonicalize_ink06
from scripts.crohme_lattice_common import writer_fit_validation
from scripts.train_crohme_segmentation_lattice_selector import _samples
from scripts.train_math_ink_06_skeleton_adapter import (
    _fused_exact_family_logits06,
    _paired_record_feature06,
)
from scripts.audit_math_ink_06_case_context import _load_model06
from scripts.materialize_math_ink_06_paired_feature_cache import _selected_records


TARGET_LABELS06 = frozenset({"c", "C", "x", "X", "z", "Z", r"\times"})
OPERATOR_LABELS06 = frozenset({
    "+", "-", "=", "<", ">", "/", "*", r"\times", r"\div", r"\pm", r"\cdot",
})


def _parse_args() -> argparse.Namespace:
    """필요 변수: adapter·공식 CROHME root·GPU 설정. 작동 원리: seed별 재현 가능한 행동학습 CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Train Math Ink 0.6 behavior role head")
    parser.add_argument("--adapter", type=Path, required=True)
    parser.add_argument(
        "--train-root", type=Path,
        default=PROJECT_ROOT / "research/data/R_noncommercial/ICFHR_package/CROHME2012_data/trainData",
    )
    parser.add_argument(
        "--test-root", type=Path,
        default=PROJECT_ROOT / "research/data/R_noncommercial/ICFHR_package/CROHME2012_data/testDataGT",
    )
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--seed", type=int, required=True)
    parser.add_argument("--epochs", type=int, default=60)
    parser.add_argument("--patience", type=int, default=14)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--teacher-batch-size", type=int, default=256)
    parser.add_argument("--learning-rate", type=float, default=1e-3)
    parser.add_argument("--weight-decay", type=float, default=1e-3)
    parser.add_argument("--hidden", type=int, default=48)
    parser.add_argument("--dropout", type=float, default=0.15)
    parser.add_argument("--profile", default="median_height_32")
    parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda")
    parser.add_argument("--product-proxy-per-label", type=int, default=0)
    parser.add_argument("--product-proxy-loss-weight", type=float, default=0.35)
    parser.add_argument("--product-proxy-target-bases", default="cosuvwxz")
    parser.add_argument("--product-proxy-lowercase-ratio", type=float, default=1.0)
    parser.add_argument(
        "--data", type=Path,
        default=PROJECT_ROOT / "research/data/open_pretrain/hwrt_expanded_v2/hwrt_expanded.jsonl.gz",
    )
    parser.add_argument(
        "--commercial-paired", type=Path,
        default=PROJECT_ROOT / "research/data/external_trajectory_v1/commercial_ccby4.jsonl.gz",
    )
    parser.add_argument(
        "--dataset-registry", type=Path, default=PROJECT_ROOT / "research/dataset_registry.json",
    )
    parser.add_argument(
        "--source-registry", type=Path,
        default=PROJECT_ROOT / "research/math_ink_06_source_registry.json",
    )
    parser.add_argument(
        "--hwrt-approval", type=Path,
        default=PROJECT_ROOT / "research/approvals/HWRT-ODBL-USE-APPROVAL-v1.json",
    )
    return parser.parse_args()


def _seed06(seed: int) -> None:
    """필요 변수: 실험 seed. 작동 원리: Python·NumPy·PyTorch 초기화를 함께 고정한다."""

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _group_box06(group: frozenset[int], strokes: Sequence[dict[str, Any]]) -> dict[str, float]:
    """필요 변수: 정답 stroke group·원본 획. 작동 원리: 수식 좌표계의 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]
    left, top, right, bottom = min(xs), min(ys), max(xs), max(ys)
    return {
        "left": left, "top": top, "right": right, "bottom": bottom,
        "width": max(right - left, 1e-5), "height": max(bottom - top, 1e-5),
        "cx": (left + right) * 0.5, "cy": (top + bottom) * 0.5,
    }


def _formula_box06(strokes: Sequence[dict[str, Any]]) -> dict[str, float]:
    """필요 변수: 한 수식의 전체 획. 작동 원리: canvas 정규화에 필요한 원점·크기를 반환한다."""

    group = frozenset(range(len(strokes)))
    return _group_box06(group, strokes)


def _canonical_group06(
    group: frozenset[int],
    strokes: Sequence[dict[str, Any]],
    formula_box: dict[str, float],
) -> Tensor:
    """필요 변수: symbol group·수식 bbox. 작동 원리: 원본 필순을 보존하며 formula-relative 128×19 입력을 만든다."""

    selected = []
    for order, stroke_index in enumerate(sorted(group)):
        stroke = strokes[stroke_index]
        selected.append({
            "order": order,
            "stroke_id": order,
            "points": [
                [
                    float(point[0]) - formula_box["left"],
                    float(point[1]) - formula_box["top"],
                    point[2] if len(point) > 2 else None,
                ]
                for point in stroke["points"]
            ],
        })
    ink = canonicalize_ink06(
        selected,
        canvas_width=formula_box["width"],
        canvas_height=formula_box["height"],
        source_modality="online",
        trust_timestamps=False,
    )
    return torch.from_numpy(ink.features)


def _teacher_role_indices06(labels: Sequence[str]) -> dict[str, list[int]]:
    """필요 변수: 378 vocabulary. 작동 원리: 이웃의 예측 확률을 digit/identifier/operator 역할로 합칠 index를 만든다."""

    digit = [index for index, label in enumerate(labels) if str(label).isdigit()]
    operator = [index for index, label in enumerate(labels) if str(label) in OPERATOR_LABELS06]
    identifier = [
        index for index, label in enumerate(labels)
        if (
            (len(str(label)) == 1 and str(label).isalpha())
            or (str(label).startswith("\\") and str(label) not in OPERATOR_LABELS06)
        )
    ]
    return {"digit": digit, "identifier": identifier, "operator": operator}


def _role_mass06(probability: Tensor, indices: dict[str, list[int]]) -> list[float]:
    """필요 변수: 한 기호 teacher 확률·역할 index. 작동 원리: 네 coarse role 확률을 반환한다."""

    digit = float(probability[indices["digit"]].sum())
    identifier = float(probability[indices["identifier"]].sum())
    operator = float(probability[indices["operator"]].sum())
    return [digit, identifier, operator, max(0.0, 1.0 - digit - identifier - operator)]


def _neighbor06(
    boxes: Sequence[dict[str, float]],
    index: int,
    *,
    direction: str,
    reference_height: float,
) -> int | None:
    """필요 변수: group bbox·기준 위치·방향. 작동 원리: 같은 로컬 행에 가까운 좌우 기호를 선택한다."""

    target = boxes[index]
    candidates = []
    for other_index, other in enumerate(boxes):
        if other_index == index:
            continue
        horizontal = other["cx"] < target["cx"] if direction == "left" else other["cx"] > target["cx"]
        if not horizontal:
            continue
        vertical = abs(other["cy"] - target["cy"]) / max(reference_height, target["height"], 1e-5)
        if vertical > 0.85:
            continue
        gap = (
            max(0.0, target["left"] - other["right"])
            if direction == "left"
            else max(0.0, other["left"] - target["right"])
        )
        candidates.append((gap + vertical * reference_height * 0.35, other_index))
    return min(candidates)[1] if candidates else None


def _family_log_probability06(
    log_probability: Tensor,
    labels: Sequence[str],
    base: str,
) -> tuple[float, float, float, float]:
    """필요 변수: teacher log 확률·case base. 작동 원리: lower/upper/times와 family 총질량을 공통 feature로 만든다."""

    label_to_index = {str(label): index for index, label in enumerate(labels)}
    members = (base, base.upper(), r"\times")
    values = [
        float(log_probability[label_to_index[label]]) if label in label_to_index else -20.0
        for label in members
    ]
    available = [label_to_index[label] for label in members if label in label_to_index]
    mass = float(log_probability[available].logsumexp(dim=0).exp()) if available else 0.0
    return values[0], values[1], values[2], mass


def _context_row06(
    formula: dict[str, Any],
    index: int,
    boxes: Sequence[dict[str, float]],
    teacher_logits: Tensor,
    labels: Sequence[str],
    role_indices: dict[str, list[int]],
) -> list[float]:
    """필요 변수: 수식·target group·teacher logits. 작동 원리: 누수 없는 예측 이웃과 실제 배치·획 feature를 49차로 만든다."""

    target_label = str(formula["truth_labels"][index])
    base = "x" if target_label == r"\times" else target_label.lower()
    probability = teacher_logits.softmax(dim=1)
    log_probability = teacher_logits.log_softmax(dim=1)
    target_probability = probability[index]
    lower, upper, times, family_mass = _family_log_probability06(
        log_probability[index], labels, base,
    )
    entropy = float(
        -(target_probability * target_probability.clamp_min(1e-9).log()).sum()
        / max(math.log(len(labels)), 1.0)
    )
    heights = [box["height"] for box in boxes]
    reference_height = float(np.median(heights))
    widths = [box["width"] for box in boxes]
    reference_width = float(np.median(widths))
    left_index = _neighbor06(boxes, index, direction="left", reference_height=reference_height)
    right_index = _neighbor06(boxes, index, direction="right", reference_height=reference_height)
    left_role = _role_mass06(probability[left_index], role_indices) if left_index is not None else [0.0] * 4
    right_role = _role_mass06(probability[right_index], role_indices) if right_index is not None else [0.0] * 4
    target = boxes[index]
    left = boxes[left_index] if left_index is not None else None
    right = boxes[right_index] if right_index is not None else None
    formula_box = _formula_box06(formula["strokes"])
    group = formula["truth_groups"][index]
    target_strokes = [formula["strokes"][stroke_index] for stroke_index in sorted(group)]
    cross = np.zeros(len(CROSS_FEATURE_NAMES), dtype=np.float32)
    if len(target_strokes) == 2:
        cross = cross_pair_feature_rows([(0, 1)], target_strokes)[0]
    point_count = sum(len(stroke.get("points") or []) for stroke in target_strokes)
    values = [
        lower, upper, times, family_mass,
        float(target_probability.max()), entropy,
        *left_role, *right_role,
        max(0.0, target["left"] - left["right"]) / reference_height if left else 2.0,
        max(0.0, right["left"] - target["right"]) / reference_height if right else 2.0,
        (target["cy"] - left["cy"]) / reference_height if left else 0.0,
        (right["cy"] - target["cy"]) / reference_height if right else 0.0,
        (target["cx"] - formula_box["left"]) / formula_box["width"],
        (target["cy"] - formula_box["top"]) / formula_box["height"],
        target["width"] / formula_box["width"],
        target["height"] / formula_box["height"],
        target["height"] / max(reference_height, 1e-5),
        target["width"] / max(reference_width, 1e-5),
        min(len(group) / 8.0, 1.0),
        min(point_count / 256.0, 1.0),
        float(base == "c"), float(base == "x"), float(base == "z"),
        *cross.tolist(),
        float(len(target_strokes) == 2),
    ]
    validate_behavior_context06(values)
    return values


def _materialize_split06(
    samples: Sequence[dict[str, Any]],
    engine,
    adapter: torch.nn.Module,
    *,
    device: torch.device,
    teacher_batch_size: int,
    return_metadata: bool = False,
) -> tuple[TensorDataset, dict[str, int]] | tuple[TensorDataset, dict[str, int], list[dict[str, Any]]]:
    """필요 변수: 수식 split·frozen teacher. 작동 원리: 모든 이웃은 teacher 예측으로 만들고 target stroke만 corpus에 보존한다."""

    labels = tuple(str(label) for label in engine.labels)
    role_indices = _teacher_role_indices06(labels)
    sequence_rows: list[Tensor] = []
    context_rows: list[list[float]] = []
    target_rows: list[int] = []
    metadata_rows: list[dict[str, Any]] = []
    counts: Counter[str] = Counter()
    engine.model.eval()
    adapter.eval()
    with torch.inference_mode():
        for formula in samples:
            groups = formula["truth_groups"]
            boxes = [_group_box06(group, formula["strokes"]) for group in groups]
            formula_box = _formula_box06(formula["strokes"])
            sequences = [
                _canonical_group06(group, formula["strokes"], formula_box)
                for group in groups
            ]
            logits_parts = []
            for start in range(0, len(sequences), teacher_batch_size):
                batch = torch.stack(sequences[start:start + teacher_batch_size]).to(device)
                four_hypotheses = batch.unsqueeze(1).expand(-1, 4, -1, -1).contiguous()
                logits, _family = _fused_exact_family_logits06(engine, adapter, four_hypotheses)
                logits_parts.append(logits.cpu())
            teacher_logits = torch.cat(logits_parts)
            for index, label_value in enumerate(formula["truth_labels"]):
                label = str(label_value)
                if label not in TARGET_LABELS06:
                    continue
                target = behavior_role_index06(label)
                if target is None:
                    continue
                sequence_rows.append(sequences[index])
                context_rows.append(_context_row06(
                    formula, index, boxes, teacher_logits, labels, role_indices,
                ))
                target_rows.append(target)
                counts[BEHAVIOR_ROLE_LABELS06[target]] += 1
                if return_metadata:
                    teacher_index = int(teacher_logits[index].argmax())
                    metadata_rows.append({
                        "sample_id": str(formula["sample_id"]),
                        "truth_label": label,
                        "teacher_label": labels[teacher_index],
                        "teacher_probability": float(
                            teacher_logits[index].softmax(dim=0)[teacher_index]
                        ),
                    })
    if not sequence_rows:
        raise ValueError("행동 role 학습 표본이 없습니다.")
    dataset = TensorDataset(
        torch.stack(sequence_rows),
        torch.tensor(context_rows, dtype=torch.float32),
        torch.tensor(target_rows, dtype=torch.long),
    )
    if return_metadata:
        return dataset, dict(counts), metadata_rows
    return dataset, dict(counts)


def _synthetic_scale06(label: str, rng: np.random.Generator) -> float:
    """필요 변수: case label·seed RNG. 작동 원리: 실제 행의 겹치는 크기 분포를 포함한 수식 배치 높이를 만든다."""

    ambiguous = bool(rng.random() < 0.20)
    if label.isupper():
        value = rng.normal(0.84 if ambiguous else 0.96, 0.055)
    else:
        value = rng.normal(0.84 if ambiguous else 0.68, 0.070)
    return float(np.clip(value, 0.48, 1.05))


def _place_synthetic_row06(feature: Tensor, scale: float) -> Tensor:
    """필요 변수: 실제 glyph 128×19 feature·합성 높이. 작동 원리: 필순은 보존하고 row-relative canvas/baseline만 배치한다."""

    output = feature.clone()
    top = (1.0 - scale) * 0.5
    bottom = top + scale
    output[:, 3] = top + output[:, 1] * scale
    output[:, 10] = top
    output[:, 11] = bottom
    output[:, 12] = scale
    output[:, 13] = 0.5
    output[:, 14] = 1.0
    return output


def _product_proxy_records06(args: argparse.Namespace, labels: Sequence[str]) -> list[dict[str, Any]]:
    """필요 변수: fail-closed registry·teacher vocabulary. 작동 원리: 승인 training split만 읽어 대소문자 proxy 원천을 만든다."""

    selection_args = argparse.Namespace(
        data=args.data,
        commercial_paired=args.commercial_paired,
        dataset_registry=args.dataset_registry,
        source_registry=args.source_registry,
        hwrt_approval=args.hwrt_approval,
        paired_source=["uci-uji-pen-v1", "uci-uji-pen-v2"],
        split="training",
        maximum_paired_train_per_label=args.product_proxy_per_label,
        maximum_paired_eval_per_label=0,
    )
    records, _fingerprint = _selected_records(selection_args, list(labels))
    return records


def _materialize_product_proxy06(
    records: Sequence[dict[str, Any]],
    engine,
    adapter: torch.nn.Module,
    *,
    seed: int,
    device: torch.device,
    teacher_batch_size: int,
    target_bases: frozenset[str],
    lowercase_ratio: float,
) -> tuple[TensorDataset, dict[str, Any]]:
    """필요 변수: 승인 고립 trajectory·frozen teacher. 작동 원리: 실제 stroke를 확률적 수식 행에 배치해 case 행동 proxy를 만든다."""

    labels = tuple(str(label) for label in engine.labels)
    role_indices = _teacher_role_indices06(labels)
    rng = np.random.default_rng(seed + 600)
    if not target_bases or not target_bases <= frozenset({"c", "o", "s", "u", "v", "w", "x", "z"}):
        raise ValueError("제품 proxy target base가 지원 case family와 다릅니다.")
    if not 0.0 <= lowercase_ratio <= 1.0:
        raise ValueError("제품 proxy lowercase ratio는 0~1이어야 합니다.")
    targets = [
        record for record in records
        if (
            len(str(record["label"])) == 1
            and str(record["label"]).lower() in target_bases
        )
    ]
    targets = [
        record for index, record in enumerate(targets)
        if (
            str(record["label"]).isupper()
            or (index * 2654435761 + seed) % 10_000 < int(lowercase_ratio * 10_000)
        )
    ]
    anchors = [record for record in records if str(record["label"]).isdigit()]
    if not targets or not anchors:
        raise ValueError("제품 proxy에 case target 또는 숫자 anchor가 없습니다.")
    anchors_by_writer: dict[tuple[str, str], list[dict[str, Any]]] = {}
    anchors_by_source: dict[str, list[dict[str, Any]]] = {}
    for record in anchors:
        key = (str(record.get("source")), str(record.get("writer_key")))
        anchors_by_writer.setdefault(key, []).append(record)
        anchors_by_source.setdefault(str(record.get("source")), []).append(record)
    sequences: list[Tensor] = []
    triplet_features: list[Tensor] = []
    metadata: list[dict[str, Any]] = []
    for index, record in enumerate(targets):
        key = (str(record.get("source")), str(record.get("writer_key")))
        pool = anchors_by_writer.get(key) or anchors_by_source.get(str(record.get("source"))) or anchors
        left_record = pool[(index * 2) % len(pool)]
        right_record = pool[(index * 2 + 1) % len(pool)]
        label = str(record["label"])
        scale = _synthetic_scale06(label, rng)
        target_feature = _place_synthetic_row06(_paired_record_feature06(record)[0], scale)
        left_scale = float(np.clip(rng.normal(0.96, 0.04), 0.82, 1.05))
        right_scale = float(np.clip(rng.normal(0.96, 0.04), 0.82, 1.05))
        left_feature = _place_synthetic_row06(_paired_record_feature06(left_record)[0], left_scale)
        right_feature = _place_synthetic_row06(_paired_record_feature06(right_record)[0], right_scale)
        sequences.append(target_feature)
        triplet_features.extend((target_feature, left_feature, right_feature))
        metadata.append({
            "record": record,
            "label": label,
            "scale": scale,
            "left_scale": left_scale,
            "right_scale": right_scale,
            "gap_left": float(rng.uniform(0.12, 0.55)),
            "gap_right": float(rng.uniform(0.12, 0.55)),
        })
    teacher_rows = []
    engine.model.eval()
    adapter.eval()
    with torch.inference_mode():
        for start in range(0, len(triplet_features), teacher_batch_size):
            batch = torch.stack(triplet_features[start:start + teacher_batch_size]).to(device)
            hypotheses = batch.unsqueeze(1).expand(-1, 4, -1, -1).contiguous()
            logits, _family = _fused_exact_family_logits06(engine, adapter, hypotheses)
            teacher_rows.append(logits.cpu())
    teacher_logits = torch.cat(teacher_rows).reshape(len(targets), 3, -1)
    context_rows: list[list[float]] = []
    target_rows: list[int] = []
    counts: Counter[str] = Counter()
    ambiguous_rows = 0
    for index, row in enumerate(metadata):
        label = row["label"]
        target = behavior_role_index06(label)
        if target is None:
            continue
        probability = teacher_logits[index].softmax(dim=1)
        log_probability = teacher_logits[index, 0].log_softmax(dim=0)
        base = label.lower()
        lower, upper, times, family_mass = _family_log_probability06(
            log_probability, labels, base,
        )
        target_probability = probability[0]
        entropy = float(
            -(target_probability * target_probability.clamp_min(1e-9).log()).sum()
            / max(math.log(len(labels)), 1.0)
        )
        left_role = _role_mass06(probability[1], role_indices)
        right_role = _role_mass06(probability[2], role_indices)
        record = row["record"]
        target_strokes = list(record["strokes"])
        cross = np.zeros(len(CROSS_FEATURE_NAMES), dtype=np.float32)
        if len(target_strokes) == 2:
            cross = cross_pair_feature_rows([(0, 1)], target_strokes)[0]
        point_count = sum(len(stroke.get("points") or []) for stroke in target_strokes)
        scale = float(row["scale"])
        reference_height = (float(row["left_scale"]) + float(row["right_scale"])) * 0.5
        aspect = float(sequences[index][0, 9])
        values = [
            lower, upper, times, family_mass,
            float(target_probability.max()), entropy,
            *left_role, *right_role,
            float(row["gap_left"]), float(row["gap_right"]), 0.0, 0.0,
            0.5, 0.5, min(0.24 * max(aspect, 0.2), 0.45), scale,
            scale / max(reference_height, 1e-5),
            max(aspect, 1e-3),
            min(len(target_strokes) / 8.0, 1.0),
            min(point_count / 256.0, 1.0),
            float(base == "c"), float(base == "x"), float(base == "z"),
            *cross.tolist(),
            float(len(target_strokes) == 2),
        ]
        validate_behavior_context06(values)
        context_rows.append(values)
        target_rows.append(target)
        counts[BEHAVIOR_ROLE_LABELS06[target]] += 1
        ambiguous_rows += int(
            (label.isupper() and scale < 0.90)
            or (label.islower() and scale > 0.78)
        )
    return TensorDataset(
        torch.stack(sequences),
        torch.tensor(context_rows, dtype=torch.float32),
        torch.tensor(target_rows, dtype=torch.long),
    ), {
        "samples": len(target_rows),
        "role_counts": dict(counts),
        "ambiguous_size_samples": ambiguous_rows,
        "sources": dict(Counter(str(record.get("source")) for record in targets)),
        "track": "P_approved_synthetic_layout_proxy",
        "actual_continuous_formula": False,
        "target_bases": sorted(target_bases),
        "lowercase_ratio": lowercase_ratio,
    }


def _with_weight06(dataset: TensorDataset, weight: float) -> TensorDataset:
    """필요 변수: dataset·loss weight. 작동 원리: source별 기여도를 각 행에 명시적으로 붙인다."""

    return TensorDataset(
        *dataset.tensors,
        torch.full((len(dataset),), float(weight), dtype=torch.float32),
    )


def _combine_training06(parts: Sequence[TensorDataset]) -> TensorDataset:
    """필요 변수: 동일 계약의 weighted dataset. 작동 원리: source 순서를 보존해 한 학습 tensor로 결합한다."""

    return TensorDataset(*(
        torch.cat([dataset.tensors[index] for dataset in parts], dim=0)
        for index in range(len(parts[0].tensors))
    ))


def _normalize_context06(
    training: TensorDataset,
    validation: TensorDataset,
    testing: TensorDataset,
    *,
    statistics_reference: TensorDataset | None = None,
) -> tuple[TensorDataset, TensorDataset, TensorDataset, Tensor, Tensor]:
    """필요 변수: 세 split context. 작동 원리: fit 통계만 사용해 validation/test 누수를 차단한다."""

    reference = statistics_reference or training
    mean = reference.tensors[1].mean(dim=0)
    scale = reference.tensors[1].std(dim=0).clamp_min(1e-5)

    def normalized(dataset: TensorDataset) -> TensorDataset:
        """필요 변수: 원본 dataset. 작동 원리: sequence/target은 보존하고 context만 표준화한다."""

        return TensorDataset(
            dataset.tensors[0],
            (dataset.tensors[1] - mean) / scale,
            *dataset.tensors[2:],
        )

    return normalized(training), normalized(validation), normalized(testing), mean, scale


def _metrics06(logits: Tensor, target: Tensor) -> dict[str, Any]:
    """필요 변수: N×3 logits·정답. 작동 원리: accuracy·macro-F1·class recall·ECE를 같은 분모로 계산한다."""

    prediction = logits.argmax(dim=1)
    probability = logits.softmax(dim=1)
    confusion = torch.zeros((3, 3), dtype=torch.long)
    for truth, selected in zip(target.tolist(), prediction.tolist(), strict=True):
        confusion[truth, selected] += 1
    f1_values, recalls = [], {}
    for index, label in enumerate(BEHAVIOR_ROLE_LABELS06):
        true_positive = int(confusion[index, index])
        false_positive = int(confusion[:, index].sum()) - true_positive
        false_negative = int(confusion[index, :].sum()) - true_positive
        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))
        recalls[label] = recall
    confidence = probability.max(dim=1).values
    correct = prediction.eq(target)
    ece = 0.0
    for lower in torch.arange(0.0, 1.0, 0.1):
        selected = (confidence >= lower) & (confidence < lower + 0.1)
        if selected.any():
            ece += float(selected.float().mean()) * abs(
                float(correct[selected].float().mean()) - float(confidence[selected].mean())
            )
    return {
        "samples": len(target),
        "accuracy": float(correct.float().mean()),
        "macro_f1": float(np.mean(f1_values)),
        "recall": recalls,
        "confusion": confusion.tolist(),
        "ece": ece,
    }


def _evaluate06(
    model: BehaviorRoleHead06,
    dataset: TensorDataset,
    *,
    batch_size: int,
    device: torch.device,
) -> tuple[dict[str, Any], Tensor, Tensor]:
    """필요 변수: behavior head·split. 작동 원리: 순서를 보존해 전체 logits와 정답을 모은다."""

    model.eval()
    logits_rows, target_rows = [], []
    with torch.inference_mode():
        for batch in DataLoader(dataset, batch_size=batch_size, shuffle=False):
            sequence, context, target = batch[:3]
            logits_rows.append(model(sequence.to(device), context.to(device)).cpu())
            target_rows.append(target)
    logits, targets = torch.cat(logits_rows), torch.cat(target_rows)
    return _metrics06(logits, targets), logits, targets


def _teacher_baseline06(dataset: TensorDataset) -> dict[str, Any]:
    """필요 변수: 정규화 전 context의 teacher logp 3개. 작동 원리: 행동 head가 실제로 teacher family를 개선했는지 비교한다."""

    logits = dataset.tensors[1][:, :3]
    return _metrics06(logits, dataset.tensors[2])


def main() -> None:
    """필요 변수: seed별 adapter와 CROHME 공식 split. 작동 원리: validation 선택 후 official test를 한 번 평가한다."""

    args = _parse_args()
    if args.train_root.name.casefold() != "traindata" or args.test_root.name.casefold() != "testdatagt":
        raise ValueError("CROHME2012 공식 trainData/testDataGT 조합만 허용합니다.")
    device = torch.device(args.device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA 학습을 요청했지만 사용할 수 없습니다.")
    _seed06(args.seed)
    adapter_payload = torch.load(args.adapter, map_location="cpu", weights_only=False)
    base_checkpoint = Path(str(adapter_payload["base_checkpoint"]))
    if not base_checkpoint.is_absolute():
        base_checkpoint = PROJECT_ROOT / base_checkpoint
    engine, adapter = _load_model06(base_checkpoint, args.adapter, device)
    fit_samples, validation_samples = writer_fit_validation(args.train_root, args.profile)
    test_samples = _samples(args.test_root, args.profile)
    training_raw, training_counts = _materialize_split06(
        fit_samples, engine, adapter, device=device, teacher_batch_size=args.teacher_batch_size,
    )
    validation_raw, validation_counts = _materialize_split06(
        validation_samples, engine, adapter, device=device, teacher_batch_size=args.teacher_batch_size,
    )
    testing_raw, testing_counts = _materialize_split06(
        test_samples, engine, adapter, device=device, teacher_batch_size=args.teacher_batch_size,
    )
    proxy_info: dict[str, Any] | None = None
    training_weighted = _with_weight06(training_raw, 1.0)
    if args.product_proxy_per_label > 0:
        product_records = _product_proxy_records06(args, engine.labels)
        proxy_raw, proxy_info = _materialize_product_proxy06(
            product_records,
            engine,
            adapter,
            seed=args.seed,
            device=device,
            teacher_batch_size=args.teacher_batch_size,
            target_bases=frozenset(args.product_proxy_target_bases),
            lowercase_ratio=args.product_proxy_lowercase_ratio,
        )
        training_weighted = _combine_training06((
            training_weighted,
            _with_weight06(proxy_raw, args.product_proxy_loss_weight),
        ))
    del engine, adapter
    if device.type == "cuda":
        torch.cuda.empty_cache()
    teacher_baseline = {
        "validation": _teacher_baseline06(validation_raw),
        "official_test": _teacher_baseline06(testing_raw),
    }
    training, validation, testing, mean, scale = _normalize_context06(
        training_weighted, validation_raw, testing_raw,
        statistics_reference=training_raw,
    )
    model = BehaviorRoleHead06(hidden=args.hidden, dropout=args.dropout).to(device)
    counts = torch.bincount(training.tensors[2], minlength=3).float()
    class_weight = (counts.sum() / counts.clamp_min(1.0)).sqrt()
    class_weight = (class_weight / class_weight.mean()).to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay,
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
    loader = DataLoader(
        training, batch_size=args.batch_size, shuffle=True,
        generator=torch.Generator().manual_seed(args.seed),
    )
    best_key = (-1.0, -1.0)
    best_state: dict[str, Tensor] | None = None
    best_epoch = 0
    history = []
    stale = 0
    for epoch in range(1, args.epochs + 1):
        model.train()
        total_loss = 0.0
        samples = 0
        for sequence, context, target, sample_weight in loader:
            sequence, context = sequence.to(device), context.to(device)
            target, sample_weight = target.to(device), sample_weight.to(device)
            optimizer.zero_grad(set_to_none=True)
            logits = model(sequence, context)
            losses = torch.nn.functional.cross_entropy(
                logits, target, weight=class_weight, reduction="none",
            )
            loss = (losses * sample_weight).sum() / sample_weight.sum().clamp_min(1e-8)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 2.0)
            optimizer.step()
            total_loss += float(loss.detach()) * len(target)
            samples += len(target)
        scheduler.step()
        validation_metrics, _logits, _targets = _evaluate06(
            model, validation, batch_size=args.batch_size, device=device,
        )
        row = {
            "epoch": epoch,
            "training_loss": total_loss / max(samples, 1),
            "learning_rate": optimizer.param_groups[0]["lr"],
            "validation": validation_metrics,
        }
        history.append(row)
        key = (validation_metrics["macro_f1"], validation_metrics["accuracy"])
        if key > best_key:
            best_key = key
            best_epoch = epoch
            best_state = deepcopy({key: value.detach().cpu() for key, value in model.state_dict().items()})
            stale = 0
        else:
            stale += 1
        if epoch == 1 or epoch % 5 == 0:
            print(json.dumps({
                "seed": args.seed, "epoch": epoch, "loss": row["training_loss"],
                "validation_accuracy": validation_metrics["accuracy"],
                "validation_macro_f1": validation_metrics["macro_f1"],
            }, ensure_ascii=False), flush=True)
        if stale >= args.patience:
            break
    if best_state is None:
        raise RuntimeError("behavior role checkpoint가 선택되지 않았습니다.")
    model.load_state_dict(best_state)
    validation_metrics, _validation_logits, _validation_targets = _evaluate06(
        model, validation, batch_size=args.batch_size, device=device,
    )
    test_metrics, _test_logits, _test_targets = _evaluate06(
        model, testing, batch_size=args.batch_size, device=device,
    )
    args.output.mkdir(parents=True, exist_ok=True)
    checkpoint = args.output / "behavior_role_head.pt"
    torch.save({
        "schema": "aiflow-math-ink-06-behavior-role-v1",
        "state_dict": best_state,
        "context_mean": mean,
        "context_scale": scale,
        "context_features": BEHAVIOR_CONTEXT_FEATURES06,
        "role_labels": BEHAVIOR_ROLE_LABELS06,
        "sequence_channels": 19,
        "hidden": args.hidden,
        "dropout": args.dropout,
        "seed": args.seed,
        "teacher_adapter": str(args.adapter),
        "selected_epoch": best_epoch,
        "track": "R_noncommercial_only",
        "product_validation": False,
    }, checkpoint)
    report = {
        "experiment": "R-MATH-INK-06-BEHAVIOR-ROLE-001",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "seed": args.seed,
        "device": str(device),
        "cuda_device": torch.cuda.get_device_name(0) if device.type == "cuda" else None,
        "teacher_adapter": str(args.adapter),
        "base_checkpoint": str(base_checkpoint),
        "split_contract": "CROHME2012 trainData writer fit/validation; testDataGT official held-out",
        "formulas": {
            "fit": len(fit_samples), "validation": len(validation_samples), "official_test": len(test_samples),
        },
        "role_samples": {
            "fit": training_counts, "validation": validation_counts, "official_test": testing_counts,
        },
        "product_proxy": proxy_info,
        "product_proxy_loss_weight": (
            args.product_proxy_loss_weight if proxy_info is not None else 0.0
        ),
        "teacher_baseline": teacher_baseline,
        "selected_epoch": best_epoch,
        "validation": validation_metrics,
        "official_test": test_metrics,
        "history": history,
        "checkpoint": checkpoint.name,
        "checkpoint_bytes": checkpoint.stat().st_size,
        "track": "R_noncommercial_only",
        "product_validation": False,
        "interpretation_limit": (
            "CROHME 연구용 연속 수식 행동 head이며 상용 checkpoint에 병합할 수 없다. "
            "P-track writer/device-disjoint 재학습이 필요하다."
        ),
    }
    (args.output / "report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps({
        "seed": args.seed,
        "selected_epoch": best_epoch,
        "teacher_test": teacher_baseline["official_test"],
        "behavior_test": test_metrics,
        "checkpoint": str(checkpoint),
        "product_validation": False,
    }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()