scripts/train_math_ink_06_p_boundary_auxiliary.py

"""승인 paired trajectory만으로 shared encoder용 boundary auxiliary head를 smoke 학습한다."""

from __future__ import annotations

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

import numpy as np
import torch
from sklearn.metrics import f1_score, roc_auc_score
from torch import nn
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.math_ink_06 import MathInk06Model, boundary_auxiliary_loss06
from scripts.train_math_ink_06_skeleton_adapter import _build_adapter06, _resolve_device06


def _parse_args() -> argparse.Namespace:
    """필요 변수: P-track feature cache·base/adapter checkpoint. 작동 원리: synthetic boundary smoke CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Train Math Ink 0.6 P-track boundary auxiliary head")
    parser.add_argument(
        "--training-cache", type=Path,
        default=Path(r"D:\Aiflow-CUDA\ink06_feature_cache\paired-training-e811a2cfb9871e990f87.pt"),
    )
    parser.add_argument(
        "--validation-cache", type=Path,
        default=Path(r"D:\Aiflow-CUDA\ink06_feature_cache\paired-validation-b98e59caaacf15025b4f.pt"),
    )
    parser.add_argument(
        "--test-cache", type=Path,
        default=Path(r"D:\Aiflow-CUDA\ink06_feature_cache\paired-paired-test-26bee2f320c6f0a7eca3.pt"),
    )
    parser.add_argument(
        "--base-checkpoint", type=Path,
        default=PROJECT_ROOT / "research/runs/math_ink_06_federated_virtual_ce025_family010_seed17_20260723/math_ink_06_candidate.pt",
    )
    parser.add_argument(
        "--adapter-checkpoint", type=Path,
        default=PROJECT_ROOT / "research/runs/math_ink_06_online_casecontext_refined_seed17_20260723/skeleton_adapter.pt",
    )
    parser.add_argument("--samples-per-class", type=int, default=1200)
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--learning-rate", type=float, default=1e-3)
    parser.add_argument("--seed", type=int, default=17)
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def _file_sha25606(path: Path) -> str:
    """필요 변수: artifact 경로. 작동 원리: lineage 검증용 SHA-256을 streaming 계산한다."""

    digest = sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _load_feature_cache06(path: Path) -> tuple[torch.Tensor, torch.Tensor, str]:
    """필요 변수: paired feature cache. 작동 원리: mmap으로 첫 online 가설과 label만 읽어 메모리 복제를 막는다."""

    payload = torch.load(path, map_location="cpu", weights_only=True, mmap=True)
    features = payload["features"]
    targets = payload["targets"]
    if features.ndim != 4 or features.shape[1:] != (4, 128, 19):
        raise ValueError(f"paired feature cache shape가 다릅니다: {tuple(features.shape)}")
    return features[:, 0], targets, str(payload["cache_key"])


def _resample_valid06(sequence: torch.Tensor, maximum: int) -> torch.Tensor:
    """필요 변수: padding 포함 128×19 sequence·점 상한. 작동 원리: 시작/끝을 포함한 균등 index로 유효 타점만 축약한다."""

    valid = sequence[sequence[:, 8] >= 0].clone()
    if len(valid) <= maximum:
        return valid
    indices = torch.linspace(0, len(valid) - 1, maximum).round().long()
    return valid[indices]


def _transform_half06(points: torch.Tensor, *, x_offset: float, y_scale: float, y_offset: float) -> torch.Tensor:
    """필요 변수: 한 기호 타점·배치 scale/offset. 작동 원리: 원 필순을 보존해 수식 내 좌·우 또는 첨자 위치로 옮긴다."""

    output = points.clone()
    output[:, 2] = output[:, 2].clamp(0.0, 1.0) * 0.40 + x_offset
    output[:, 3] = output[:, 3].clamp(0.0, 1.0) * y_scale + y_offset
    # 비등방 변환 뒤 방향 벡터만 다시 정규화한다.
    direction = output[:, 4:6] * torch.tensor([0.40, y_scale])
    output[:, 4:6] = direction / direction.square().sum(dim=1, keepdim=True).sqrt().clamp_min(1e-6)
    return output


def _merge_candidate06(first: torch.Tensor, second: torch.Tensor, variant: int) -> torch.Tensor:
    """필요 변수: 서로 다른 두 P-track 기호 sequence·배치 variant. 작동 원리: 행·첨자·분수 슬롯 경계 침범 후보를 합성한다."""

    first_points = _resample_valid06(first, 64)
    second_points = _resample_valid06(second, 64)
    layout = variant % 5
    if layout == 0:
        first_points = _transform_half06(first_points, x_offset=0.05, y_scale=0.80, y_offset=0.10)
        second_points = _transform_half06(second_points, x_offset=0.55, y_scale=0.80, y_offset=0.10)
    elif layout == 1:
        first_points = _transform_half06(first_points, x_offset=0.05, y_scale=0.70, y_offset=0.25)
        second_points = _transform_half06(second_points, x_offset=0.55, y_scale=0.45, y_offset=0.05)
    elif layout == 2:
        first_points = _transform_half06(first_points, x_offset=0.05, y_scale=0.70, y_offset=0.10)
        second_points = _transform_half06(second_points, x_offset=0.55, y_scale=0.45, y_offset=0.50)
    elif layout == 3:
        # 분수선 후보가 분자·분모를 함께 먹는 상황을 근사하는 수직 슬롯 배치다.
        first_points = _transform_half06(first_points, x_offset=0.30, y_scale=0.34, y_offset=0.05)
        second_points = _transform_half06(second_points, x_offset=0.30, y_scale=0.34, y_offset=0.61)
    else:
        # 등호·중위연산자 양쪽의 완성 기호를 하나로 합치는 넓은 행 후보를 만든다.
        first_points = _transform_half06(first_points, x_offset=0.01, y_scale=0.72, y_offset=0.14)
        second_points = _transform_half06(second_points, x_offset=0.59, y_scale=0.72, y_offset=0.14)
    second_points[0, 7] = 1.0
    valid = torch.cat((first_points, second_points), dim=0)[:128]
    minimum = valid[:, 2:4].amin(dim=0)
    span = (valid[:, 2:4].amax(dim=0) - minimum).clamp_min(1e-6)
    valid[:, 0:2] = (valid[:, 2:4] - minimum) / span
    valid[:, 9] = span[0] / span[1]
    valid[:, 10] = minimum[1]
    valid[:, 11] = minimum[1] + span[1]
    valid[:, 12] = span[1]
    valid[:, 13] = minimum[1] + span[1] * 0.5
    valid[:, 14] = 1.0
    valid[:, 18] = 0.0
    output = torch.zeros(128, 19, dtype=valid.dtype)
    output[:, 8] = -1.0
    output[:len(valid)] = valid
    return output


def _balanced_boundary_set06(
    features: torch.Tensor,
    labels: torch.Tensor,
    *,
    samples_per_class: int,
    seed: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """필요 변수: writer-disjoint paired feature·label. 작동 원리: 실제 다획 단일기호와 합성 두기호 후보를 같은 수로 만든다."""

    if samples_per_class <= 0 or len(features) < samples_per_class * 2:
        raise ValueError("boundary smoke 표본 상한이 cache 크기와 맞지 않습니다.")
    generator = torch.Generator().manual_seed(seed)
    indices = torch.randperm(len(features), generator=generator)[:samples_per_class * 2]
    negative_indices = indices[:samples_per_class]
    first_indices = indices[:samples_per_class]
    second_indices = indices[samples_per_class:]
    # 같은 label도 실제로는 경계일 수 있지만 smoke에서는 서로 다른 label을 우선해 target 모호성을 줄인다.
    collision = labels[first_indices] == labels[second_indices]
    if collision.any():
        second_indices[collision] = second_indices[collision].roll(1)
    negatives = features[negative_indices].clone()
    positives = torch.stack([
        _merge_candidate06(features[int(first)], features[int(second)], variant)
        for variant, (first, second) in enumerate(zip(first_indices, second_indices, strict=True))
    ])
    output = torch.cat((negatives, positives), dim=0)
    targets = torch.cat((torch.zeros(samples_per_class), torch.ones(samples_per_class)))
    permutation = torch.randperm(len(output), generator=generator)
    return output[permutation], targets[permutation]


def _load_encoder06(
    base_path: Path, adapter_path: Path, device: torch.device,
) -> tuple[MathInk06Model, nn.Module, dict, dict]:
    """필요 변수: base·online adapter checkpoint. 작동 원리: 기존 exact/family weight를 고정하고 새 boundary head만 초기화한다."""

    base = torch.load(base_path, map_location="cpu", weights_only=False)
    adapter_payload = torch.load(adapter_path, map_location="cpu", weights_only=False)
    model = MathInk06Model(
        exact_classes=len(base["exact_labels"]),
        family_classes=len(base["family_labels"]),
        hidden_size=int(base["hidden_size"]),
        hypotheses=int(base["hypotheses"]),
        raster_architecture=str(base["raster_architecture"]),
        use_boundary_head=True,
    )
    incompatible = model.load_state_dict(base["state_dict"], strict=False)
    allowed_missing = {
        "boundary_head.weight", "boundary_head.bias",
        "raster_encoder.fine_projection.weight", "raster_encoder.fine_projection.bias",
        "raster_encoder.pointer_projection.weight", "raster_encoder.pointer_projection.bias",
        "virtual_decoder.pointer_query.weight", "virtual_decoder.pointer_key.weight",
    }
    if set(incompatible.missing_keys) != allowed_missing or incompatible.unexpected_keys:
        raise ValueError(f"base checkpoint 호환 오류: {incompatible}")
    adapter = _build_adapter06(str(adapter_payload["adapter_architecture"]))
    adapter.load_state_dict(adapter_payload["state_dict"])
    for parameter in model.parameters():
        parameter.requires_grad = False
    if model.boundary_head is None:
        raise RuntimeError("boundary head 초기화에 실패했습니다.")
    for parameter in model.boundary_head.parameters():
        parameter.requires_grad = True
    for parameter in adapter.parameters():
        parameter.requires_grad = False
    return model.to(device), adapter.to(device), base, adapter_payload


def _embeddings06(
    model: MathInk06Model,
    adapter: nn.Module,
    features: torch.Tensor,
    *,
    device: torch.device,
    batch_size: int,
) -> torch.Tensor:
    """필요 변수: 고정 encoder/adapter·candidate feature. 작동 원리: head 반복학습 전에 embedding을 한 번만 GPU 계산한다."""

    rows = []
    model.eval()
    adapter.eval()
    with torch.inference_mode():
        for start in range(0, len(features), batch_size):
            batch = features[start:start + batch_size].to(device)
            rows.append(model.encode_trajectory(adapter(batch)).cpu())
    return torch.cat(rows)


def _metrics06(logits: torch.Tensor, targets: torch.Tensor, *, threshold: float = 0.5) -> dict[str, float]:
    """필요 변수: boundary logit·정답·threshold. 작동 원리: balanced accuracy/F1/AUC와 class recall을 계산한다."""

    probability = logits.sigmoid().numpy()
    truth = targets.numpy().astype(np.int64)
    predicted = (probability >= threshold).astype(np.int64)
    return {
        "threshold": threshold,
        "accuracy": float((predicted == truth).mean()),
        "f1": float(f1_score(truth, predicted)),
        "roc_auc": float(roc_auc_score(truth, probability)),
        "single_symbol_recall": float(((predicted == 0) & (truth == 0)).sum() / max((truth == 0).sum(), 1)),
        "cross_boundary_recall": float(((predicted == 1) & (truth == 1)).sum() / max((truth == 1).sum(), 1)),
    }


def _select_threshold06(
    logits: torch.Tensor,
    targets: torch.Tensor,
    *,
    minimum_single_recall: float = 0.95,
    minimum_boundary_recall: float = 0.85,
) -> tuple[dict[str, float], list[dict[str, float]]]:
    """필요 변수: validation logit·두 recall floor. 작동 원리: 보존 제약을 만족하는 F1 최대 threshold를 선택한다."""

    trials = [
        _metrics06(logits, targets, threshold=float(value))
        for value in np.linspace(0.05, 0.95, 37)
    ]
    eligible = [
        row for row in trials
        if (
            row["single_symbol_recall"] >= minimum_single_recall
            and row["cross_boundary_recall"] >= minimum_boundary_recall
        )
    ]
    if eligible:
        winner = max(eligible, key=lambda row: (
            row["f1"], row["accuracy"], row["roc_auc"], row["cross_boundary_recall"],
        ))
        return {**winner, "recall_gate_passed": True}, trials
    preservation_candidates = [
        row for row in trials if row["single_symbol_recall"] >= minimum_single_recall
    ]
    fallback = max(
        preservation_candidates or trials,
        key=lambda row: (
            row["single_symbol_recall"] >= minimum_single_recall,
            row["cross_boundary_recall"], row["f1"], row["accuracy"],
        ),
    )
    return {
        **fallback,
        "recall_gate_passed": False,
        "required_single_symbol_recall": minimum_single_recall,
        "required_cross_boundary_recall": minimum_boundary_recall,
    }, trials


def main() -> None:
    """필요 변수: P-track train/validation cache. 작동 원리: shared encoder를 동결하고 boundary head만 학습·선택·저장한다."""

    args = _parse_args()
    torch.manual_seed(args.seed)
    device = _resolve_device06(args.device)
    train_features, train_labels, train_cache_key = _load_feature_cache06(args.training_cache)
    validation_features, validation_labels, validation_cache_key = _load_feature_cache06(args.validation_cache)
    test_features, test_labels, test_cache_key = _load_feature_cache06(args.test_cache)
    train_x, train_y = _balanced_boundary_set06(
        train_features, train_labels, samples_per_class=args.samples_per_class, seed=args.seed,
    )
    validation_x, validation_y = _balanced_boundary_set06(
        validation_features, validation_labels,
        samples_per_class=min(args.samples_per_class, len(validation_features) // 2),
        seed=args.seed + 1,
    )
    test_x, test_y = _balanced_boundary_set06(
        test_features, test_labels,
        samples_per_class=min(args.samples_per_class, len(test_features) // 2),
        seed=args.seed + 2,
    )
    model, adapter, base, adapter_payload = _load_encoder06(
        args.base_checkpoint, args.adapter_checkpoint, device,
    )
    train_embedding = _embeddings06(
        model, adapter, train_x, device=device, batch_size=args.batch_size,
    )
    validation_embedding = _embeddings06(
        model, adapter, validation_x, device=device, batch_size=args.batch_size,
    )
    test_embedding = _embeddings06(
        model, adapter, test_x, device=device, batch_size=args.batch_size,
    )
    if model.boundary_head is None:
        raise RuntimeError("boundary head가 없습니다.")
    head = model.boundary_head.cpu()
    optimizer = torch.optim.AdamW(head.parameters(), lr=args.learning_rate, weight_decay=1e-3)
    loader = DataLoader(
        TensorDataset(train_embedding, train_y), batch_size=args.batch_size,
        shuffle=True, generator=torch.Generator().manual_seed(args.seed),
    )
    best = None
    history = []
    exact_probe = model.exact_head(validation_embedding[:32].to(device)).detach().cpu()
    family_probe = model.family_head(validation_embedding[:32].to(device)).detach().cpu()
    for epoch in range(1, args.epochs + 1):
        head.train()
        losses = []
        for embedding, target in loader:
            optimizer.zero_grad(set_to_none=True)
            loss = boundary_auxiliary_loss06(head(embedding).squeeze(-1), target)
            loss.backward()
            optimizer.step()
            losses.append(float(loss.detach()))
        head.eval()
        with torch.inference_mode():
            validation_logits = head(validation_embedding).squeeze(-1)
        metrics = _metrics06(validation_logits, validation_y)
        row = {"epoch": epoch, "loss": float(np.mean(losses)), **metrics}
        history.append(row)
        if best is None or (row["f1"], row["roc_auc"], row["accuracy"]) > (
            best["f1"], best["roc_auc"], best["accuracy"],
        ):
            best = {**row, "state_dict": {key: value.detach().clone() for key, value in head.state_dict().items()}}
    if best is None:
        raise RuntimeError("boundary 학습 결과가 없습니다.")
    head.load_state_dict(best.pop("state_dict"))
    head.eval()
    with torch.inference_mode():
        validation_logits = head(validation_embedding).squeeze(-1)
        test_logits = head(test_embedding).squeeze(-1)
    selected_threshold, threshold_trials = _select_threshold06(validation_logits, validation_y)
    paired_test = _metrics06(
        test_logits, test_y, threshold=float(selected_threshold["threshold"]),
    )
    exact_after = model.exact_head(validation_embedding[:32].to(device)).detach().cpu()
    family_after = model.family_head(validation_embedding[:32].to(device)).detach().cpu()
    non_regression = {
        "exact_logit_max_abs": float((exact_after - exact_probe).abs().max()),
        "family_logit_max_abs": float((family_after - family_probe).abs().max()),
    }
    args.output.mkdir(parents=True, exist_ok=True)
    checkpoint_path = args.output / "boundary_auxiliary_head.pt"
    torch.save({
        "schema": "aiflow-math-ink-06-p-boundary-auxiliary-v1",
        "state_dict": head.state_dict(),
        "input_embedding": int(train_embedding.shape[1]),
        "selected_epoch": int(best["epoch"]),
        "synthetic_layouts": ["same_row", "superscript", "subscript", "fraction_slots", "wide_infix_sides"],
        "base_checkpoint_sha256": _file_sha25606(args.base_checkpoint),
        "adapter_checkpoint_sha256": _file_sha25606(args.adapter_checkpoint),
        "training_cache_key": train_cache_key,
        "validation_cache_key": validation_cache_key,
        "test_cache_key": test_cache_key,
        "threshold": float(selected_threshold["threshold"]),
        "track": "P_with_obligations",
        "product_validation": False,
    }, checkpoint_path)
    report = {
        "experiment": "P-MATH-INK-06-BOUNDARY-AUXILIARY-SMOKE-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,
        "base_model_version": base["model_version"],
        "adapter_model_version": adapter_payload["model_version"],
        "training_samples": len(train_y),
        "validation_samples": len(validation_y),
        "selected": best,
        "threshold_selection": selected_threshold,
        "threshold_trials": threshold_trials,
        "paired_test": paired_test,
        "history": history,
        "non_regression": non_regression,
        "checkpoint": str(checkpoint_path),
        "checkpoint_sha256": _file_sha25606(checkpoint_path),
        "interpretation_limit": (
            "승인 paired 고립기호를 합성 배치한 boundary proxy smoke이며 실제 연속식 writer/device 제품 gate가 아니다."
        ),
        "track": "P_with_obligations",
        "product_validation": False,
    }
    (args.output / "report.json").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 not in {"history", "threshold_trials"}
    }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()