src/ink06_federation.py

"""Math Ink 0.6의 상용 허용 trajectory 출처를 공통 paired record로 결합한다."""

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass
from hashlib import sha256
import json
from pathlib import Path
import statistics
from typing import Iterable, Sequence

import numpy as np
import torch
from torch.utils.data import Dataset, WeightedRandomSampler

from .external_corpus import read_jsonl
from .ink06_canonical import canonicalize_ink06, render_canonical_ink
from .ink06_source_registry import approved_training_source_ids06, load_source_registry06
from .trajectory_sequence import shape_family


@dataclass(frozen=True, slots=True)
class FederationSource06:
    """필요 변수: source 권리·분할 metadata. 작동 원리: 한 출처의 학습 허용 근거와 record 수를 고정한다."""

    source_id: str
    license_id: str
    approval_id: str | None
    records: tuple[dict, ...]


def resolve_training_device06(requested: str) -> str:
    """필요 변수: auto·cpu·cuda device 문자열. 작동 원리: CUDA 가능 여부를 확인하고 CPU 묵시적 fallback을 차단한다."""

    value = requested.strip().lower()
    if value == "auto":
        return "cuda" if torch.cuda.is_available() else "cpu"
    if value.startswith("cuda") and not torch.cuda.is_available():
        raise ValueError("CUDA 학습을 요청했지만 현재 PyTorch가 CUDA device를 찾지 못했습니다.")
    if value == "cpu" or value.startswith("cuda"):
        return value
    raise ValueError(f"지원하지 않는 학습 device입니다: {requested}")


def _read_registry_status(registry_path: Path) -> dict[str, dict]:
    """필요 변수: UTF-8 dataset registry. 작동 원리: dataset ID별 P-track 권리 상태를 반환한다."""

    payload = json.loads(registry_path.read_text(encoding="utf-8"))
    return {str(row["id"]): row for row in payload["datasets"]}


def _verify_hwrt_approval(approval_path: Path) -> dict:
    """필요 변수: 프로젝트 소유자 승인 JSON. 작동 원리: HWRT model_training scope와 고정 curation ID를 검사한다."""

    approval = json.loads(approval_path.read_text(encoding="utf-8"))
    if not approval.get("approved") or "model_training" not in approval.get("approved_scopes", []):
        raise ValueError("HWRT model_training 승인이 없습니다.")
    if approval.get("dataset_id") != "hwrt" or approval.get("curation_id") != "OPEN-HWRT-TRAJECTORY-001":
        raise ValueError("HWRT 승인 대상과 현재 curation이 다릅니다.")
    return approval


def _stable_writer_split06(records: Sequence[dict]) -> list[dict]:
    """필요 변수: writer_key가 있는 HWRT 승인 train 표본. 작동 원리: writer 전체를 80/10/10 train·validation·test로 고정 분리한다."""

    writers = {
        str(record.get("writer_key") or record.get("writer_id") or "").strip()
        for record in records
    }
    writers.discard("")
    writers.discard("missing")
    if len(writers) < 3:
        raise ValueError("HWRT writer-disjoint 분할에는 식별 가능한 writer가 최소 3명 필요합니다.")
    ordered = sorted(
        writers,
        key=lambda writer: sha256(f"math-ink-06:{writer}".encode("utf-8")).digest(),
    )
    holdout_count = max(1, round(len(ordered) * 0.1))
    if holdout_count * 2 >= len(ordered):
        holdout_count = 1
    test_writers = set(ordered[:holdout_count])
    validation_writers = set(ordered[holdout_count:holdout_count * 2])
    output: list[dict] = []
    for record in records:
        writer = str(record.get("writer_key") or record.get("writer_id") or "").strip()
        if not writer or writer == "missing":
            raise ValueError("HWRT train 표본에 writer ID가 없어 writer-disjoint 분할을 만들 수 없습니다.")
        split = "test" if writer in test_writers else "validation" if writer in validation_writers else "train"
        value = dict(record)
        value["split"] = split
        value["eligible_for_training"] = split == "train"
        value["split_contract"] = "aiflow_writer_disjoint_80_10_10_v1"
        output.append(value)
    return output


def _hold_out_training_writers06(records: Sequence[dict]) -> list[dict]:
    """필요 변수: 공식 train/test와 writer_key. 작동 원리: 공식 test를 보존하면서 train writer의 10%를 validation으로 격리한다."""

    train_writers = {
        str(record.get("writer_key") or record.get("writer_id") or "").strip()
        for record in records
        if record.get("split") == "train"
    }
    train_writers.discard("")
    train_writers.discard("missing")
    if len(train_writers) < 2:
        return [dict(record) for record in records]
    ordered = sorted(
        train_writers,
        key=lambda writer: sha256(f"math-ink-06-validation:{writer}".encode("utf-8")).digest(),
    )
    validation_count = min(max(1, round(len(ordered) * 0.1)), len(ordered) - 1)
    validation_writers = set(ordered[:validation_count])
    output: list[dict] = []
    for record in records:
        value = dict(record)
        writer = str(value.get("writer_key") or value.get("writer_id") or "").strip()
        if value.get("split") == "train" and writer in validation_writers:
            value["split"] = "validation"
            value["eligible_for_training"] = False
            value["split_contract"] = "official_test_plus_writer_validation_v1"
        output.append(value)
    return output


def _deduplicate_federation_records06(
    grouped: dict[str, list[dict]],
    ordered_ids: Sequence[str],
) -> dict[str, list[dict]]:
    """필요 변수: 출처별 trajectory와 canonical 우선순위. 작동 원리: label-aware 정규화 hash 중복을 한 번만 남기고 충돌 label은 격리한다."""

    signature_labels: dict[str, set[str]] = {}
    for records in grouped.values():
        for record in records:
            signature = _trajectory_signature06(record)
            if signature is not None:
                signature_labels.setdefault(signature, set()).add(str(record.get("label")))
    conflicting_signatures = {
        signature for signature, labels in signature_labels.items() if len(labels) > 1
    }
    seen: set[tuple[str, str]] = set()
    cleaned: dict[str, list[dict]] = {source_id: [] for source_id in grouped}
    for source_id in ordered_ids:
        for record in grouped[source_id]:
            signature = _trajectory_signature06(record)
            if signature is None:
                cleaned[source_id].append(record)
                continue
            if signature in conflicting_signatures:
                continue
            key = (str(record.get("label")), signature)
            if key in seen:
                continue
            seen.add(key)
            cleaned[source_id].append(record)
    return cleaned


def load_product_federation06(
    *, registry_path: Path, commercial_path: Path, hwrt_path: Path, approval_path: Path,
    allowed_labels: Sequence[str], source_registry_path: Path | None = None,
) -> tuple[FederationSource06, ...]:
    """필요 변수: registry·P shard·승인·378 vocabulary. 작동 원리: 승인되고 로컬 전처리된 trajectory source만 fail-closed로 적재한다."""

    registry = _read_registry_status(registry_path)
    label_set = set(allowed_labels)
    if source_registry_path is None:
        source_ids = ("uci-pendigits", "uci-uji-pen-v2", "hwrt")
    else:
        source_ids = approved_training_source_ids06(load_source_registry06(source_registry_path))
    if not source_ids:
        raise ValueError("승인되고 로컬 전처리된 supervised trajectory source가 없습니다.")
    grouped: dict[str, list[dict]] = {source_id: [] for source_id in source_ids}
    for source_id in source_ids:
        if source_id == "hwrt":
            continue
        row = registry.get(source_id)
        if row is None or "P" not in row.get("allowed_tracks", []) or row.get("status") != "product_with_obligations":
            raise ValueError(f"{source_id}가 P-track allowlist를 통과하지 못했습니다.")
    for record in read_jsonl(commercial_path):
        source_id = str(record.get("source"))
        if source_id in grouped and source_id != "hwrt" and str(record.get("label")) in label_set:
            grouped[source_id].append(dict(record))

    approval = None
    if "hwrt" in grouped:
        hwrt_registry = registry.get("hwrt")
        approval = _verify_hwrt_approval(approval_path)
        if hwrt_registry is None or "P" not in hwrt_registry.get("allowed_tracks", []):
            raise ValueError("HWRT가 registry P-track allowlist에 없습니다.")
        approved_train = [
            dict(record)
            for record in read_jsonl(hwrt_path)
            if str(record.get("label")) in label_set and record.get("split") == "train"
        ]
        # 공식 test는 제품 경로에서 완전히 제외하고 승인 train writer만 독립 분할한다.
        for value in _stable_writer_split06(approved_train):
            value["approval_id"] = approval["approval_id"]
            grouped["hwrt"].append(value)

    sources = []
    # UJI v2가 v1 전체를 포함하므로 같은 family에서는 더 큰 v2를 canonical 원본으로 둔다.
    deduplication_order = ("uci-pendigits", "uci-uji-pen-v2", "uci-uji-pen-v1", "hwrt")
    ordered_ids = sorted(
        source_ids,
        key=lambda value: (
            deduplication_order.index(value)
            if value in deduplication_order
            else len(deduplication_order),
            value,
        ),
    )
    grouped = _deduplicate_federation_records06(grouped, ordered_ids)
    for source_id in ordered_ids:
        if source_id != "hwrt":
            grouped[source_id] = _hold_out_training_writers06(grouped[source_id])
    for source_id in ordered_ids:
        if not grouped[source_id]:
            # 승인된 mirror가 canonical source에 완전히 흡수된 경우 별도 source로 세지 않는다.
            continue
        license_id = str(grouped[source_id][0].get("license_id") or registry[source_id]["license"])
        sources.append(FederationSource06(
            source_id=source_id, license_id=license_id,
            approval_id=approval["approval_id"] if source_id == "hwrt" and approval is not None else None,
            records=tuple(grouped[source_id]),
        ))
    return tuple(sources)


def federation_audit06(sources: Sequence[FederationSource06]) -> dict:
    """필요 변수: 승인 source 묶음. 작동 원리: origin·정규화 trajectory·writer/device split 누수를 함께 보고한다."""

    origins_by_source: dict[str, set[str]] = {}
    signatures_by_source: dict[str, set[str]] = {}
    identity_splits: dict[str, dict[str, set[str]]] = {
        "origin": {},
        "writer": {},
        "device": {},
    }
    rows = {}
    for source in sources:
        origins = {str(record.get("origin_id") or record["sample_id"]) for record in source.records}
        origins_by_source[source.source_id] = origins
        signature_values = [
            _trajectory_signature06(record)
            for record in source.records
        ]
        valid_signatures = [
            signature for signature in signature_values if signature is not None
        ]
        signatures = set(valid_signatures)
        signatures_by_source[source.source_id] = signatures
        for record in source.records:
            split = str(record.get("split") or "missing")
            origin = str(record.get("origin_id") or record.get("sample_id") or "").strip()
            writer = str(record.get("writer_key") or record.get("writer_id") or "").strip()
            device = str(record.get("device_id") or "").strip()
            if origin:
                identity_splits["origin"].setdefault(origin, set()).add(split)
            if writer and writer.lower() != "missing":
                identity_splits["writer"].setdefault(
                    f"{source.source_id}:{writer}",
                    set(),
                ).add(split)
            if device and device.lower() != "missing":
                identity_splits["device"].setdefault(
                    f"{source.source_id}:{device}",
                    set(),
                ).add(split)
        rows[source.source_id] = {
            "records": len(source.records),
            "training_records": sum(bool(record.get("eligible_for_training")) for record in source.records),
            "evaluation_only_records": sum(not bool(record.get("eligible_for_training")) for record in source.records),
            "labels": len({str(record["label"]) for record in source.records}),
            "writers": len({str(record.get("writer_key") or "missing") for record in source.records}),
            "splits": dict(Counter(str(record.get("split") or "missing") for record in source.records)),
            "trajectory_signatures": len(signatures),
            "trajectory_signature_missing": sum(
                signature is None for signature in signature_values
            ),
            "trajectory_signature_duplicates_within_source": (
                len(valid_signatures) - len(signatures)
            ),
            "license_id": source.license_id, "approval_id": source.approval_id,
        }
    overlaps = {}
    signature_overlaps = {}
    for first_index, first in enumerate(sources):
        for second in sources[first_index + 1:]:
            key = f"{first.source_id}|{second.source_id}"
            overlaps[key] = len(origins_by_source[first.source_id] & origins_by_source[second.source_id])
            signature_overlaps[key] = len(
                signatures_by_source[first.source_id]
                & signatures_by_source[second.source_id]
            )
    split_leakage = {
        kind: {
            identity: sorted(splits)
            for identity, splits in values.items()
            if len(splits) > 1
        }
        for kind, values in identity_splits.items()
    }
    return {
        "sources": rows, "source_count": len(sources), "origin_overlap": overlaps,
        "origin_overlap_total": sum(overlaps.values()),
        "trajectory_signature_overlap": signature_overlaps,
        "trajectory_signature_overlap_total": sum(signature_overlaps.values()),
        "split_identity_leakage": split_leakage,
        "split_identity_leakage_total": sum(
            len(values) for values in split_leakage.values()
        ),
    }


def federation_provenance06(
    sources: Sequence[FederationSource06],
    source_registry_path: Path,
) -> dict:
    """필요 변수: 실제 적재 source와 registry 파일. 작동 원리: checkpoint가 학습 출처·독립 그룹·registry byte hash를 스스로 증명하게 한다."""

    entries = load_source_registry06(source_registry_path)
    by_id = {entry.source_id: entry for entry in entries}
    source_ids = tuple(sorted(source.source_id for source in sources))
    missing = [source_id for source_id in source_ids if source_id not in by_id]
    if missing:
        raise ValueError(f"registry에 없는 실제 학습 source입니다: {missing}")
    groups = sorted({
        by_id[source_id].independent_source_group
        for source_id in source_ids
        if by_id[source_id].independent_source_group
    })
    return {
        "training_source_ids": list(source_ids),
        "training_independent_source_groups": groups,
        "source_registry_sha256": sha256(source_registry_path.read_bytes()).hexdigest(),
    }


def _trajectory_signature06(record: dict) -> str | None:
    """필요 변수: trajectory record. 작동 원리: 이동·크기 차이를 제거한 stroke별 좌표를 hash해 미러 중복을 찾는다."""

    raw_strokes = record.get("strokes")
    if not isinstance(raw_strokes, list) or not raw_strokes:
        return None
    parsed: list[list[tuple[float, float]]] = []
    try:
        for raw_stroke in raw_strokes:
            points_value = (
                raw_stroke.get("points")
                if isinstance(raw_stroke, dict)
                else raw_stroke
            )
            if not isinstance(points_value, list) or not points_value:
                return None
            points = []
            for point in points_value:
                if isinstance(point, dict):
                    points.append((float(point["x"]), float(point["y"])))
                else:
                    points.append((float(point[0]), float(point[1])))
            parsed.append(points)
    except (KeyError, TypeError, ValueError, IndexError):
        return None
    all_points = [point for stroke in parsed for point in stroke]
    left = min(point[0] for point in all_points)
    top = min(point[1] for point in all_points)
    width = max(max(point[0] for point in all_points) - left, 1e-9)
    height = max(max(point[1] for point in all_points) - top, 1e-9)
    normalized = [
        [
            (round((x - left) / width, 4), round((y - top) / height, 4))
            for x, y in stroke
        ]
        for stroke in parsed
    ]
    payload = json.dumps(normalized, separators=(",", ":")).encode("utf-8")
    return sha256(payload).hexdigest()


def _canvas_for_record(record: dict) -> tuple[float, float]:
    """필요 변수: normalized external record. 작동 원리: source 좌표 계약에 맞는 canvas를 반환한다."""

    canvas = record.get("canvas") or {}
    return float(canvas.get("width", 768.0)), float(canvas.get("height", 128.0))


class FederatedPairedInk06Dataset(Dataset):
    """필요 변수: source record·label index. 작동 원리: 모든 P source를 같은 128 raster와 19채널 trajectory로 변환한다."""

    def __init__(self, records: Sequence[dict], exact_to_index: dict[str, int], family_to_index: dict[str, int]) -> None:
        self.records = tuple(records)
        self.exact_to_index = exact_to_index
        self.family_to_index = family_to_index

    def __len__(self) -> int:
        """필요 변수: record 목록. 작동 원리: federation 표본 수를 반환한다."""

        return len(self.records)

    def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]:
        """필요 변수: 표본 index. 작동 원리: source 시간 추정을 observed로 위장하지 않고 paired tensor를 만든다."""

        record = self.records[index]
        width, height = _canvas_for_record(record)
        ink = canonicalize_ink06(
            record["strokes"], canvas_width=width, canvas_height=height, trust_timestamps=False,
        )
        raster = 1.0 - np.asarray(render_canonical_ink(ink), dtype=np.float32) / 255.0
        features = ink.features
        coordinates = features[:, 2:4].copy()
        valid = features[:, 8] >= 0
        states = np.full(len(features), 2, dtype=np.int64)
        states[valid] = 0
        states[np.logical_and(valid, features[:, 7] > 0.5)] = 1
        states[-1] = 2
        label = str(record["label"])
        return (
            torch.from_numpy(features), torch.from_numpy(raster).unsqueeze(0),
            torch.from_numpy(coordinates), torch.from_numpy(states),
            torch.tensor(self.exact_to_index[label]),
            torch.tensor(self.family_to_index[shape_family(label)]), str(record["source"]),
        )


def source_label_balanced_sampler06(records: Sequence[dict], *, seed: int, samples: int) -> WeightedRandomSampler:
    """필요 변수: source·label record와 seed. 작동 원리: source와 label 빈도의 역수로 각 batch의 편향을 줄인다."""

    source_label_counts = Counter((str(row["source"]), str(row["label"])) for row in records)
    labels_per_source = Counter()
    for source_id, _label in source_label_counts:
        labels_per_source[source_id] += 1
    weights = [
        1.0 / labels_per_source[str(row["source"])]
        / source_label_counts[(str(row["source"]), str(row["label"]))]
        for row in records
    ]
    generator = torch.Generator().manual_seed(seed)
    return WeightedRandomSampler(
        torch.tensor(weights, dtype=torch.double), num_samples=samples, replacement=True, generator=generator,
    )


def interpolate_state_dict06(
    baseline: dict[str, torch.Tensor], candidate: dict[str, torch.Tensor], *, alpha: float,
) -> dict[str, torch.Tensor]:
    """필요 변수: 학습 전후 state dict·보간율. 작동 원리: float weight만 선형 보간해 기존 분포 회귀를 제한한다."""

    if not 0.0 <= alpha <= 1.0:
        raise ValueError("state dict 보간 alpha는 0과 1 사이여야 합니다.")
    if baseline.keys() != candidate.keys():
        raise ValueError("보간할 state dict key가 다릅니다.")
    output = {}
    for key, base_value in baseline.items():
        candidate_value = candidate[key]
        if base_value.shape != candidate_value.shape or base_value.dtype != candidate_value.dtype:
            raise ValueError(f"보간할 tensor 계약이 다릅니다: {key}")
        output[key] = torch.lerp(base_value, candidate_value, alpha) if torch.is_floating_point(base_value) else base_value.clone()
    return output


def summarize_federated_seeds06(runs: Sequence[dict]) -> dict:
    """필요 변수: seed별 학습·strict 보고서. 작동 원리: 각 seed의 회귀와 절대 release gate를 분리해 distillation 가능 여부를 계산한다."""

    if not runs:
        raise ValueError("요약할 federation seed가 없습니다.")
    required_metrics = ("online_top1", "online_top5", "raster_top1", "raster_top5")
    rows = []
    raster_top5_failure_sets: list[set[str]] = []
    failure_truth: dict[str, str] = {}
    seen: set[int] = set()
    for run in sorted(runs, key=lambda value: int(value["seed"])):
        seed = int(run["seed"])
        if seed in seen:
            raise ValueError(f"federation seed가 중복입니다: {seed}")
        seen.add(seed)
        training = run["training"]
        strict_report = run["strict"]
        strict = strict_report["aiflow_math_ink_06"]
        strict_rows = strict_report.get("rows")
        if isinstance(strict_rows, list):
            failed = {str(item["sample_id"]) for item in strict_rows if not bool(item.get("raster_top5"))}
            raster_top5_failure_sets.append(failed)
            failure_truth.update({str(item["sample_id"]): str(item["truth"]) for item in strict_rows})
        deltas = training["test_delta"]
        holdout_nonregression = all(
            float(source[metric]) >= -1e-12 for source in deltas.values() for metric in required_metrics
        )
        gates = {
            "online_top1_92": float(strict["online_top1"]) >= 0.92,
            "online_top5_99": float(strict["online_top5"]) >= 0.99,
            "raster_top1_90": float(strict["raster_top1"]) >= 0.90,
            "model_under_25mb": int(strict["checkpoint_bytes"]) <= 25 * 1024 * 1024,
        }
        rows.append({
            "seed": seed, "selected_epoch": int(training["selected_epoch"]),
            "selected_interpolation_alpha": float(training.get("selected_interpolation_alpha", 0.0)),
            "holdout_nonregression": holdout_nonregression,
            "strict": {metric: float(strict[metric]) for metric in required_metrics},
            "gates": gates,
            "individual_release_gate_passed": holdout_nonregression and all(gates.values()),
        })
    aggregates = {}
    for metric in required_metrics:
        values = [row["strict"][metric] for row in rows]
        aggregates[metric] = {
            "mean": statistics.fmean(values), "population_std": statistics.pstdev(values),
            "minimum": min(values), "maximum": max(values),
        }
    all_passed = len(rows) == 3 and {row["seed"] for row in rows} == {17, 31, 47} and all(
        row["individual_release_gate_passed"] for row in rows
    )
    consensus_failures = set.intersection(*raster_top5_failure_sets) if len(raster_top5_failure_sets) == len(rows) else set()
    return {
        "required_seeds": [17, 31, 47], "runs": rows, "strict_aggregate": aggregates,
        "selected_update_seeds": [row["seed"] for row in rows if row["selected_epoch"] > 0],
        "strict_consensus_raster_top5_failures": [
            {"sample_id": sample_id, "truth": failure_truth[sample_id]} for sample_id in sorted(consensus_failures)
        ],
        "all_individual_release_gates_passed": all_passed,
        "student_distillation_allowed": all_passed,
        "product_validation": False,
    }