scripts/audit_math_ink_06_p_proxy_shift.py

"""승인 P 합성 배치와 실제 연속식 행동 feature의 분포 차이를 정량 감사한다."""

from __future__ import annotations

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

import numpy as np
import torch
from torch import Tensor

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,
)
from scripts.audit_math_ink_06_case_context import _load_model06
from scripts.crohme_lattice_common import writer_fit_validation
from scripts.train_math_ink_06_behavior_role import (
    _materialize_product_proxy06,
    _materialize_split06,
    _product_proxy_records06,
)


AUDIT_FEATURES06 = (
    "teacher_family_mass",
    "teacher_top1",
    "teacher_entropy",
    "bbox_width",
    "bbox_height",
    "local_height_ratio",
    "local_width_ratio",
    "left_gap",
    "right_gap",
)


def distribution_shift06(reference: Tensor, candidate: Tensor) -> dict[str, float | int]:
    """필요 변수: 기준·후보 1차원 값. 작동 원리: 평균·표준편차·분위수 거리와 표준화 평균차를 계산한다."""

    reference = reference.detach().float().flatten()
    candidate = candidate.detach().float().flatten()
    if not len(reference) or not len(candidate):
        raise ValueError("분포 비교에는 양쪽 표본이 모두 필요합니다.")
    reference_mean = float(reference.mean())
    candidate_mean = float(candidate.mean())
    reference_std = float(reference.std(unbiased=False))
    candidate_std = float(candidate.std(unbiased=False))
    pooled_std = max(
        ((reference_std ** 2 + candidate_std ** 2) * 0.5) ** 0.5,
        1e-6,
    )
    quantiles = torch.linspace(0.0, 1.0, 101)
    quantile_distance = float(
        (torch.quantile(reference, quantiles) - torch.quantile(candidate, quantiles))
        .abs()
        .mean()
    )
    return {
        "reference_samples": len(reference),
        "candidate_samples": len(candidate),
        "reference_mean": reference_mean,
        "candidate_mean": candidate_mean,
        "reference_std": reference_std,
        "candidate_std": candidate_std,
        "standardized_mean_difference": (candidate_mean - reference_mean) / pooled_std,
        "mean_absolute_quantile_distance": quantile_distance,
    }


def compare_contexts06(
    reference_context: Tensor,
    reference_target: Tensor,
    candidate_context: Tensor,
    candidate_target: Tensor,
    *,
    feature_names: Sequence[str] = AUDIT_FEATURES06,
) -> dict[str, object]:
    """필요 변수: 실제/합성 context와 역할 target. 작동 원리: 공통 역할별 feature shift를 누수 없이 비교한다."""

    feature_index = {
        name: index for index, name in enumerate(BEHAVIOR_CONTEXT_FEATURES06)
    }
    unknown = sorted(set(feature_names) - set(feature_index))
    if unknown:
        raise ValueError(f"알 수 없는 행동 feature입니다: {unknown}")
    roles: dict[str, object] = {}
    for role_index, role in enumerate(BEHAVIOR_ROLE_LABELS06):
        reference_mask = reference_target == role_index
        candidate_mask = candidate_target == role_index
        if not reference_mask.any() or not candidate_mask.any():
            continue
        features = {
            name: distribution_shift06(
                reference_context[reference_mask, feature_index[name]],
                candidate_context[candidate_mask, feature_index[name]],
            )
            for name in feature_names
        }
        ranked = sorted(
            (
                {
                    "feature": name,
                    "absolute_standardized_mean_difference": abs(
                        float(values["standardized_mean_difference"])
                    ),
                }
                for name, values in features.items()
            ),
            key=lambda row: row["absolute_standardized_mean_difference"],
            reverse=True,
        )
        roles[role] = {
            "reference_samples": int(reference_mask.sum()),
            "candidate_samples": int(candidate_mask.sum()),
            "features": features,
            "largest_shifts": ranked[:5],
        }
    return roles


def _parse_args() -> argparse.Namespace:
    """필요 변수: 제품 teacher·CROHME train·P proxy 설정. 작동 원리: 재현 가능한 분포 감사 CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Audit P synthetic behavior proxy shift")
    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("--profile", default="median_height_32")
    parser.add_argument("--seed", type=int, default=17)
    parser.add_argument("--teacher-batch-size", type=int, default=256)
    parser.add_argument("--product-proxy-per-label", type=int, default=100)
    parser.add_argument("--product-proxy-target-bases", default="cosuvwxz")
    parser.add_argument("--product-proxy-lowercase-ratio", type=float, default=1.0)
    parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda")
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    """필요 변수: CLI 인자. 작동 원리: 실제 validation과 P 합성 proxy를 역할별로 비교해 잘못된 배치 가정을 찾는다."""

    args = _parse_args()
    device = torch.device(args.device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA 감사를 요청했지만 사용할 수 없습니다.")
    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)
    validation, validation_counts = _materialize_split06(
        validation_samples,
        engine,
        adapter,
        device=device,
        teacher_batch_size=args.teacher_batch_size,
    )
    selection_args = argparse.Namespace(**vars(args))
    # 기존 fail-closed P source loader가 요구하는 경로와 split 계약을 그대로 사용한다.
    selection_args.data = (
        PROJECT_ROOT / "research/data/open_pretrain/hwrt_expanded_v2/hwrt_expanded.jsonl.gz"
    )
    selection_args.commercial_paired = (
        PROJECT_ROOT / "research/data/external_trajectory_v1/commercial_ccby4.jsonl.gz"
    )
    selection_args.dataset_registry = PROJECT_ROOT / "research/dataset_registry.json"
    selection_args.source_registry = PROJECT_ROOT / "research/math_ink_06_source_registry.json"
    selection_args.hwrt_approval = (
        PROJECT_ROOT / "research/approvals/HWRT-ODBL-USE-APPROVAL-v1.json"
    )
    product_records = _product_proxy_records06(selection_args, engine.labels)
    proxy, 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,
    )
    role_shift = compare_contexts06(
        validation.tensors[1],
        validation.tensors[2],
        proxy.tensors[1],
        proxy.tensors[2],
    )
    layout_features = {"bbox_height", "local_height_ratio", "bbox_width", "local_width_ratio"}
    maximum_layout_shift = max(
        (
            float(row["absolute_standardized_mean_difference"])
            for role in role_shift.values()
            for row in role["largest_shifts"]
            if row["feature"] in layout_features
        ),
        default=0.0,
    )
    report = {
        "experiment": "R-MATH-INK-06-P-PROXY-SHIFT-001",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "seed": args.seed,
        "reference": "CROHME writer-validation truth groups",
        "candidate": "approved P isolated trajectories with synthetic row layout",
        "validation_role_counts": validation_counts,
        "proxy": proxy_info,
        "role_shift": role_shift,
        "maximum_layout_absolute_smd": maximum_layout_shift,
        "decision": {
            "layout_distribution_compatible": maximum_layout_shift <= 0.50,
            "threshold_absolute_smd": 0.50,
            "expand_to_three_seeds": False,
            "product_validation": False,
        },
        "track": "diagnostic_R_reference_plus_P_proxy",
        "product_validation": False,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(report["decision"], ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()