scripts/evaluate_math_ink_06_p_boundary_device_stress.py

"""채택된 P boundary joint delta를 미관측 device/sampling stress에서 평가한다."""

from __future__ import annotations

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

import torch

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 scripts.train_math_ink_06_p_boundary_auxiliary import (
    _balanced_boundary_set06,
    _load_encoder06,
    _load_feature_cache06,
    _metrics06,
)
from scripts.train_math_ink_06_p_boundary_joint import (
    _authentic_metrics06,
    _boundary_logits06,
    _family_target_map06,
)
from scripts.train_math_ink_06_skeleton_adapter import _resolve_device06


def _parse_args() -> argparse.Namespace:
    """필요 변수: seed별 base/adapter/head/delta. 작동 원리: 고정 3-seed stress 평가 CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Evaluate Math Ink 0.6 P boundary device stress")
    parser.add_argument(
        "--test-cache", type=Path,
        default=Path(r"D:\Aiflow-CUDA\ink06_feature_cache\paired-paired-test-26bee2f320c6f0a7eca3.pt"),
    )
    parser.add_argument("--samples-per-class", type=int, default=1200)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
    parser.add_argument(
        "--boundary-template", type=str, default="",
        help="PROJECT_ROOT 기준 checkpoint template이며 {seed}를 치환한다.",
    )
    parser.add_argument(
        "--no-joint-delta", action="store_true",
        help="채택 delta 없이 정정된 base+shared adapter+auxiliary head를 평가한다.",
    )
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def _sparse_sequence06(sequence: torch.Tensor) -> torch.Tensor:
    """필요 변수: 128×19 canonical sequence. 작동 원리: raw event 절반을 제거한 뒤 6Hz canonical timeline으로 재보간한다."""

    valid = sequence[sequence[:, 8] >= 0]
    if len(valid) <= 2:
        return sequence.clone()
    anchors = torch.nonzero(valid[:, 7] > 0.5, as_tuple=False).flatten()
    indices = torch.unique(torch.cat((
        torch.arange(0, len(valid), 2), anchors, torch.tensor([0, len(valid) - 1]),
    )), sorted=True)
    kept = valid[indices].clone()
    # 좌표·방향·시간 등 연속 channel은 희소 raw event에서 canonical tick으로 선형 복원한다.
    restored = torch.nn.functional.interpolate(
        kept.transpose(0, 1).unsqueeze(0), size=len(valid),
        mode="linear", align_corners=True,
    ).squeeze(0).transpose(0, 1)
    # Pen state와 modality/missing 계약은 추정 interpolation 값이 아니라 명시적 anchor를 유지한다.
    restored[:, 7] = 0.0
    restored[torch.nonzero(valid[:, 7] > 0.5, as_tuple=False).flatten(), 7] = 1.0
    restored[:, 17] = valid[:, 17]
    restored[:, 18] = valid[:, 18]
    restored[:, 8] = valid[:, 8]
    output = torch.zeros_like(sequence)
    output[:, 8] = -1.0
    output[:len(restored)] = restored
    return output


def _stress06(features: torch.Tensor, mode: str, *, seed: int) -> torch.Tensor:
    """필요 변수: canonical feature batch·stress mode. 작동 원리: label을 보지 않고 device/sampling 변동만 적용한다."""

    if mode == "clean":
        return features.clone()
    output = features.clone()
    valid = output[:, :, 8] >= 0
    if mode == "coordinate_jitter":
        generator = torch.Generator().manual_seed(seed)
        noise = torch.randn(output.shape[:2] + (2,), generator=generator) * 0.008
        for columns in ((0, 1), (2, 3)):
            values = output[:, :, list(columns)]
            values[valid] = (values[valid] + noise[valid]).clamp(0.0, 1.0)
            output[:, :, list(columns)] = values
    elif mode == "timestamp_missing":
        output[:, :, 15:17] = 0.0
        output[:, :, 17] = 1.0
    elif mode == "sparse_sampling":
        output = torch.stack([_sparse_sequence06(sequence) for sequence in output])
    elif mode == "affine_device":
        x = ((output[:, :, 2] - 0.5) * 1.12 + 0.5).clamp(0.0, 1.0)
        y = ((output[:, :, 3] - 0.5) * 0.88 + 0.5).clamp(0.0, 1.0)
        output[:, :, 2] = torch.where(valid, x, output[:, :, 2])
        output[:, :, 3] = torch.where(valid, y, output[:, :, 3])
        output[:, :, 9] = output[:, :, 9] * (1.12 / 0.88)
    else:
        raise ValueError(f"지원하지 않는 stress mode입니다: {mode}")
    return output


def _seed_paths06(seed: int, boundary_template: str = "") -> dict[str, Path]:
    """필요 변수: seed·선택 boundary template. 작동 원리: base·adapter·head·legacy delta 경로를 결정한다."""

    boundary_run = (
        "math_ink_06_p_boundary_auxiliary_layout5_20260724"
        if seed == 17 else f"math_ink_06_p_boundary_auxiliary_layout5_seed{seed}_20260724"
    )
    boundary_path = (
        PROJECT_ROOT / boundary_template.format(seed=seed)
        if boundary_template
        else PROJECT_ROOT / f"research/runs/{boundary_run}/boundary_auxiliary_head.pt"
    )
    return {
        "base": PROJECT_ROOT / f"research/runs/math_ink_06_federated_virtual_ce025_family010_seed{seed}_20260723/math_ink_06_candidate.pt",
        "adapter": PROJECT_ROOT / f"research/runs/math_ink_06_online_casecontext_refined_seed{seed}_20260723/skeleton_adapter.pt",
        "boundary": boundary_path,
        "joint": PROJECT_ROOT / f"research/runs/math_ink_06_p_boundary_joint_seed{seed}_20260724/boundary_joint_delta.pt",
    }


def main() -> None:
    """필요 변수: paired test와 3-seed delta. 작동 원리: clean 대비 stress 회귀를 seed별로 계산하고 최저 gate를 판정한다."""

    args = _parse_args()
    device = _resolve_device06(args.device)
    features, targets, cache_key = _load_feature_cache06(args.test_cache)
    modes = ("clean", "coordinate_jitter", "timestamp_missing", "sparse_sampling", "affine_device")
    seed_reports = []
    for seed in (17, 31, 47):
        paths = _seed_paths06(seed, args.boundary_template)
        model, adapter, base, _adapter_payload = _load_encoder06(paths["base"], paths["adapter"], device)
        boundary = torch.load(paths["boundary"], map_location="cpu", weights_only=True)
        if model.boundary_head is None:
            raise RuntimeError("boundary head가 없습니다.")
        model.boundary_head.load_state_dict(boundary["state_dict"])
        threshold = float(boundary["threshold"])
        if not args.no_joint_delta:
            joint = torch.load(paths["joint"], map_location="cpu", weights_only=True)
            if not joint["adopted"]:
                raise ValueError(f"채택되지 않은 joint delta입니다: seed {seed}")
            model.load_state_dict(joint["state_dict"], strict=False)
            threshold = float(joint["threshold"])
        family_map = _family_target_map06(base)
        boundary_x, boundary_y = _balanced_boundary_set06(
            features, targets, samples_per_class=args.samples_per_class, seed=seed + 2,
        )
        rows = {}
        for mode_index, mode in enumerate(modes):
            authentic_stressed = _stress06(features, mode, seed=seed * 100 + mode_index)
            boundary_stressed = _stress06(boundary_x, mode, seed=seed * 1000 + mode_index)
            authentic_metrics = _authentic_metrics06(
                model, adapter, authentic_stressed, targets, family_map,
                device=device, batch_size=args.batch_size,
            )
            boundary_logits = _boundary_logits06(
                model, adapter, boundary_stressed, device=device, batch_size=args.batch_size,
            )
            boundary_metrics = _metrics06(
                boundary_logits, boundary_y, threshold=threshold,
            )
            rows[mode] = {"authentic": authentic_metrics, "boundary": boundary_metrics}
        clean = rows["clean"]
        for mode in modes[1:]:
            rows[mode]["deltas"] = {
                "exact_top1_pp": (
                    rows[mode]["authentic"]["exact_top1"] - clean["authentic"]["exact_top1"]
                ) * 100.0,
                "family_top1_pp": (
                    rows[mode]["authentic"]["family_top1"] - clean["authentic"]["family_top1"]
                ) * 100.0,
            }
            rows[mode]["gate_passed"] = bool(
                rows[mode]["deltas"]["exact_top1_pp"] >= -3.0
                and rows[mode]["deltas"]["family_top1_pp"] >= -3.0
                and rows[mode]["boundary"]["single_symbol_recall"] >= 0.90
                and rows[mode]["boundary"]["cross_boundary_recall"] >= 0.90
            )
        seed_reports.append({
            "seed": seed, "threshold": threshold, "modes": rows,
            "all_stress_gates_passed": all(rows[mode]["gate_passed"] for mode in modes[1:]),
        })
    failures = [
        {"seed": row["seed"], "mode": mode}
        for row in seed_reports for mode in modes[1:]
        if not row["modes"][mode]["gate_passed"]
    ]
    report = {
        "experiment": "P-MATH-INK-06-BOUNDARY-DEVICE-STRESS-001",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "device": str(device),
        "cuda_device": torch.cuda.get_device_name(device) if device.type == "cuda" else None,
        "test_cache_key": cache_key,
        "stress_contract": {
            "model_variant": (
                "base_plus_shared_adapter_plus_auxiliary_head"
                if args.no_joint_delta else "legacy_joint_delta"
            ),
            "modes": list(modes),
            "maximum_exact_regression_pp": 3.0,
            "maximum_family_regression_pp": 3.0,
            "minimum_single_symbol_recall": 0.90,
            "minimum_cross_boundary_recall": 0.90,
        },
        "seeds": seed_reports,
        "decision": {
            "all_seed_stress_gates_passed": not failures,
            "failures": failures,
            "product_validation": False,
        },
        "track": "P_with_obligations",
        "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, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()