scripts/train_math_ink_06_boundary_behavior_guard.py

"""Grouping 이전 geometry로 기호 경계 침범 확률을 학습하고 최종 selector에서 검증한다."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import json
from pathlib import Path
import sys
from typing import Any

import joblib
import numpy as np
from sklearn.ensemble import HistGradientBoostingClassifier

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.aiflow_ocr05 import AIFlowOCR05
from math_grid_drawer.research.cross_visual import CrossVisualModel
from math_grid_drawer.research.equality_visual import EqualityVisualModel
from math_grid_drawer.research.segmentation_lattice import (
    LATTICE_FEATURE_NAMES,
    lattice_candidate_features,
)
from scripts.crohme_lattice_common import load_cache_for_samples, load_cached_split, writer_fit_validation
from scripts.evaluate_crohme_gt_free_grouping import _truth_partition
from scripts.evaluate_crohme_lattice_ocr_fusion import _fit_geometry
from scripts.evaluate_crohme_tray_joint_selector import _prepared_signals, _weighted
from scripts.sweep_math_ink_06_multistroke_family_guard import _family_metrics06
from scripts.sweep_math_ink_06_x_grouping_guard import _target_metrics06
from scripts.train_crohme_segmentation_lattice_joint_selector import _metrics


def _parse_args() -> argparse.Namespace:
    """필요 변수: CROHME split·cache·head 경로. 작동 원리: writer-disjoint boundary 학습/검증 CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Train Math Ink 0.6 boundary behavior guard")
    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(
        "--cache-dir", type=Path,
        default=PROJECT_ROOT / "research/runs/crohme_lattice_ocr_cache_v2_20260722",
    )
    parser.add_argument(
        "--bundle", type=Path,
        default=Path(r"research\runs\aiflow_ocr_05_dual_trajectory_3seed_20260720\bundle.manifest.json"),
    )
    parser.add_argument(
        "--cross-model", type=Path,
        default=PROJECT_ROOT / "research/runs/crohme_cross_visual_loop3_polyline_20260722/cross_visual.json",
    )
    parser.add_argument(
        "--equality-model", type=Path,
        default=PROJECT_ROOT / "research/runs/crohme_equality_visual_loop1_20260722/equality_visual.json",
    )
    parser.add_argument("--profile", default="median_height_32")
    parser.add_argument("--maximum-x-regression-pp", type=float, default=1.0)
    parser.add_argument("--maximum-family-regression-pp", type=float, default=2.0)
    parser.add_argument("--maximum-pair-f1-regression-pp", type=float, default=0.25)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def _boundary_training_rows(samples: list[dict[str, Any]]) -> tuple[np.ndarray, np.ndarray]:
    """필요 변수: fit writer 수식·truth partition. 작동 원리: 둘 이상의 정답 기호를 침범한 다획 후보를 boundary=1로 학습한다."""

    feature_rows: list[np.ndarray] = []
    targets: list[bool] = []
    for sample in samples:
        strokes = sample["profiled_strokes"]
        candidates = AIFlowOCR05.build_segmentation_lattice(strokes)
        features = lattice_candidate_features(candidates, strokes)
        truth_groups, _labels = _truth_partition(sample, "aiflow_geometry")
        for candidate, feature in zip(candidates, features, strict=True):
            group = frozenset(int(value) for value in candidate["source_indices"])
            if len(group) <= 1:
                continue
            feature_rows.append(feature)
            targets.append(sum(bool(group & truth) for truth in truth_groups) > 1)
    return np.asarray(feature_rows, dtype=np.float32), np.asarray(targets, dtype=np.int64)


def _fit_boundary_model(samples: list[dict[str, Any]]) -> tuple[HistGradientBoostingClassifier, dict[str, Any]]:
    """필요 변수: fit writer 후보. 작동 원리: class-balanced gradient boosting으로 기호 경계 침범 확률을 학습한다."""

    features, targets = _boundary_training_rows(samples)
    positives = max(int(targets.sum()), 1)
    negatives = max(len(targets) - positives, 1)
    sample_weight = np.where(
        targets == 1, len(targets) / (2 * positives), len(targets) / (2 * negatives),
    )
    model = HistGradientBoostingClassifier(
        learning_rate=0.06,
        max_iter=220,
        max_leaf_nodes=31,
        l2_regularization=2.0,
        min_samples_leaf=40,
        random_state=17,
    ).fit(features, targets, sample_weight=sample_weight)
    return model, {
        "candidate_rows": len(targets),
        "boundary_rows": int(targets.sum()),
        "boundary_rate": float(targets.mean()),
    }


def _boundary_probabilities(
    prepared: list[dict[str, Any]], model: HistGradientBoostingClassifier,
) -> list[np.ndarray]:
    """필요 변수: 평가 후보 geometry·학습 모델. 작동 원리: singleton은 0, 다획 후보만 경계 침범 확률을 반환한다."""

    output = []
    for row in prepared:
        probabilities = model.predict_proba(row["features"][:, :len(LATTICE_FEATURE_NAMES)])[:, 1]
        multistroke = np.asarray(
            [len(candidate["source_indices"]) > 1 for candidate in row["candidates"]],
            dtype=bool,
        )
        output.append(np.where(multistroke, probabilities, 0.0).astype(np.float32))
    return output


def _evaluate06(
    samples: list[dict[str, Any]],
    prepared: list[dict[str, Any]],
    probabilities: list[np.ndarray],
    *,
    threshold: float,
    weight: float,
) -> dict[str, Any]:
    """필요 변수: selector row·boundary 확률·threshold/weight. 작동 원리: 확률 초과분만 logit에서 빼고 보호 지표를 계산한다."""

    weighted = _weighted(
        prepared, tray_weight=4.0, symbol_weight=4.0,
        fraction_weight=8.0, infix_weight=8.0,
    )
    guarded = []
    for row, probability in zip(weighted, probabilities, strict=True):
        penalty = np.maximum(probability - threshold, 0.0) / max(1.0 - threshold, 1e-6)
        guarded.append({**row, "logits": row["logits"] - weight * penalty})
    return {
        "threshold": threshold,
        "weight": weight,
        "global": _metrics(guarded, -2.0),
        "behavior_targets": _target_metrics06(samples, guarded, group_bias=-2.0),
        "families": _family_metrics06(samples, guarded),
    }


def _delta_pp06(candidate: float, reference: float) -> float:
    """필요 변수: 후보·기준 비율. 작동 원리: 채택 판단용 percentage-point 차이를 반환한다."""

    return (candidate - reference) * 100.0


def main() -> None:
    """필요 변수: fit/validation/test writer split. 작동 원리: fit 학습·validation 선택 후 고정 winner만 official test에 적용한다."""

    args = _parse_args()
    fit, validation = writer_fit_validation(args.train_root, args.profile)
    boundary_model, training = _fit_boundary_model(fit)
    geometry_model = _fit_geometry(fit)
    equality_model = EqualityVisualModel.load(args.equality_model)
    cross_model = CrossVisualModel.load(args.cross_model)
    validation_cache = load_cache_for_samples(
        validation, args.cache_dir, split="validation", profile=args.profile,
        bundle=args.bundle, version=2,
    )
    validation_prepared = _prepared_signals(
        validation, validation_cache, geometry_model,
        equality_model=equality_model, cross_model=cross_model,
        cross_gap_ratio=0.40, multistroke_family_boost=6.0,
    )
    validation_probability = _boundary_probabilities(validation_prepared, boundary_model)
    trials = [
        _evaluate06(
            validation, validation_prepared, validation_probability,
            threshold=threshold, weight=weight,
        )
        for threshold in (0.30, 0.50, 0.65, 0.80, 0.90)
        for weight in (0.0, 1.0, 2.0, 4.0, 6.0, 8.0, 12.0)
    ]
    reference = trials[0]
    minimum_x = reference["behavior_targets"]["x"]["grouping_recall"] - args.maximum_x_regression_pp / 100.0
    minimum_family = reference["families"]["grouping_recall"] - args.maximum_family_regression_pp / 100.0
    minimum_pair_f1 = reference["global"]["pair_f1"] - args.maximum_pair_f1_regression_pp / 100.0
    eligible = [
        row for row in trials
        if (
            row["behavior_targets"]["x"]["grouping_recall"] >= minimum_x
            and row["families"]["grouping_recall"] >= minimum_family
            and row["global"]["pair_f1"] >= minimum_pair_f1
        )
    ]
    winner = max(eligible, key=lambda row: (
        row["global"]["exact_partition"], row["global"]["pair_f1"],
        row["global"]["exact_group_recall"], -row["weight"], row["threshold"],
    ))
    test, test_cache = load_cached_split(
        args.test_root, args.cache_dir, split="official_test", profile=args.profile,
        bundle=args.bundle, version=2,
    )
    test_prepared = _prepared_signals(
        test, test_cache, geometry_model,
        equality_model=equality_model, cross_model=cross_model,
        cross_gap_ratio=0.40, multistroke_family_boost=6.0,
    )
    test_probability = _boundary_probabilities(test_prepared, boundary_model)
    official_reference = _evaluate06(
        test, test_prepared, test_probability, threshold=0.30, weight=0.0,
    )
    official_winner = _evaluate06(
        test, test_prepared, test_probability,
        threshold=float(winner["threshold"]), weight=float(winner["weight"]),
    )
    deltas = {
        "exact_partition_pp": _delta_pp06(
            official_winner["global"]["exact_partition"], official_reference["global"]["exact_partition"],
        ),
        "pair_f1_pp": _delta_pp06(
            official_winner["global"]["pair_f1"], official_reference["global"]["pair_f1"],
        ),
        "x_grouping_pp": _delta_pp06(
            official_winner["behavior_targets"]["x"]["grouping_recall"],
            official_reference["behavior_targets"]["x"]["grouping_recall"],
        ),
        "family_grouping_pp": _delta_pp06(
            official_winner["families"]["grouping_recall"], official_reference["families"]["grouping_recall"],
        ),
    }
    adopted = bool(
        float(winner["weight"]) > 0.0
        and deltas["exact_partition_pp"] > 0.0
        and deltas["pair_f1_pp"] >= -args.maximum_pair_f1_regression_pp
        and deltas["x_grouping_pp"] >= -args.maximum_x_regression_pp
        and deltas["family_grouping_pp"] >= -args.maximum_family_regression_pp
    )
    artifact_path = args.output.parent / "boundary_behavior_guard.joblib"
    artifact_path.parent.mkdir(parents=True, exist_ok=True)
    joblib.dump({
        "schema": "aiflow-math-ink-06-boundary-behavior-v1",
        "feature_names": list(LATTICE_FEATURE_NAMES),
        "model": boundary_model,
        "threshold": float(winner["threshold"]) if adopted else None,
        "weight": float(winner["weight"]) if adopted else 0.0,
        "track": "R_noncommercial_only",
    }, artifact_path)
    report = {
        "experiment": "R-MATH-INK-06-BOUNDARY-BEHAVIOR-GUARD-001",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "training": training,
        "feature_names": list(LATTICE_FEATURE_NAMES),
        "reference_validation": reference,
        "winner_validation": winner,
        "trials": trials,
        "official_test_reference": official_reference,
        "official_test_winner": official_winner,
        "official_test_deltas": deltas,
        "decision": {
            "adopted": adopted,
            "selected_threshold": float(winner["threshold"]) if adopted else None,
            "selected_weight": float(winner["weight"]) if adopted else 0.0,
        },
        "artifact": str(artifact_path),
        "track": "R_noncommercial_only",
        "product_validation": False,
    }
    args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({key: value for key, value in report.items() if key != "trials"}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()