exoskeleton/code/analysis/make_paper_artifacts.py

"""Generate versioned trial metrics and paper source-data tables."""

from __future__ import annotations

import argparse
import csv
import json
import os
import tempfile
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

from experiments.hashing import file_sha256, stable_hash, to_jsonable
from experiments.io import atomic_write_json, atomic_write_jsonl, utc_now
from experiments.plan import load_document
from experiments.schema import SCHEMA_VERSION, STORAGE_FORMATS
from experiments.validate import require_valid_batch

from .metrics import derive_trial_metrics


PROVENANCE_IDENTITY_FIELDS = (
    "h2_data_group_id",
    "h2_data_seed_record_hash",
    "bilateral_data_group_id",
    "bilateral_data_seed_record_hash",
    "network_pair_group_id",
    "network_profile_id",
)


def _method_selected(
    family: str,
    family_configuration: Mapping[str, Any],
    method_id: str,
) -> bool:
    """Apply explicit per-family method filters without silent typos."""
    include = family_configuration.get("include_methods")
    exclude = family_configuration.get("exclude_methods")
    for name, value in (
        ("include_methods", include),
        ("exclude_methods", exclude),
    ):
        if value is not None and (
            not isinstance(value, list)
            or any(not isinstance(item, str) or not item for item in value)
        ):
            raise ValueError(f"{family}.{name} must be a list of method IDs")
    if include is not None and method_id not in include:
        return False
    if exclude is not None and method_id in exclude:
        return False
    return True


def _trial_metric_configuration(
    configuration: Mapping[str, Any],
    method_id: str,
) -> dict[str, Any] | None:
    enabled = configuration.get("enabled")
    if not isinstance(enabled, list) or not enabled:
        raise ValueError("metric configuration needs a non-empty enabled list")
    selected = []
    trial_configuration: dict[str, Any] = {}
    for family in enabled:
        if not isinstance(family, str) or not family:
            raise ValueError("enabled metric family names must be non-empty strings")
        family_configuration = configuration.get(family, {})
        if not isinstance(family_configuration, Mapping):
            raise ValueError(f"{family} metric configuration must be a mapping")
        if _method_selected(family, family_configuration, method_id):
            selected.append(family)
            trial_configuration[family] = dict(family_configuration)
    if not selected:
        return None
    trial_configuration["enabled"] = selected
    return trial_configuration


def _scalar_cell(value: Any) -> Any:
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    return json.dumps(
        to_jsonable(value),
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    )


def _identity_row(
    trial: Mapping[str, Any],
    trial_metadata: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    trajectory = trial["trajectory"]
    method = trial["method"]
    row = {
        "trial_id": trial["trial_id"],
        "pair_id": trial["pair_id"],
        "study_id": trial["study_id"],
        "split": trial["split"],
        "method_id": method["method_id"],
        "trajectory_id": trajectory["trajectory_id"],
        "trajectory_family": trajectory.get("family", ""),
        "replicate": trial["replicate"],
    }
    for name, value in sorted(trial.get("factors", {}).items()):
        row[f"factor_{name}"] = _scalar_cell(value)
    metadata = {} if trial_metadata is None else trial_metadata
    for name in PROVENANCE_IDENTITY_FIELDS:
        value = metadata.get(name)
        if value is not None:
            row[name] = _scalar_cell(value)
    if (
        "bilateral_data_seed_record_hash" not in row
        and metadata.get("bilateral_data_seed_record") is not None
    ):
        row["bilateral_data_seed_record_hash"] = stable_hash(
            metadata["bilateral_data_seed_record"],
            prefix="bilateral-data-seed-record",
        )
    if "network_profile_id" not in row:
        profile = trial.get("factors", {}).get("network_profile")
        if isinstance(profile, Mapping) and profile.get("profile_id") is not None:
            row["network_profile_id"] = _scalar_cell(profile["profile_id"])
    return row


def _atomic_write_csv(path: Path, rows: Sequence[Mapping[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    identity_order = [
        "trial_id",
        "pair_id",
        "study_id",
        "split",
        "method_id",
        "trajectory_id",
        "trajectory_family",
        "replicate",
        *PROVENANCE_IDENTITY_FIELDS,
    ]
    all_fields = {key for row in rows for key in row}
    fieldnames = [name for name in identity_order if name in all_fields]
    fieldnames.extend(sorted(all_fields - set(fieldnames)))
    fd, temporary_name = tempfile.mkstemp(
        prefix=f".{path.name}.", suffix=".tmp", dir=path.parent
    )
    temporary = Path(temporary_name)
    try:
        with os.fdopen(fd, "w", encoding="utf-8", newline="") as stream:
            writer = csv.DictWriter(stream, fieldnames=fieldnames)
            writer.writeheader()
            for row in rows:
                writer.writerow(
                    {name: _scalar_cell(row.get(name)) for name in fieldnames}
                )
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def _finite_metric(row: Mapping[str, Any], name: str) -> float:
    if name not in row:
        raise ValueError(
            f"network paired gates require metric {name!r}"
        )
    try:
        value = float(row[name])
    except (TypeError, ValueError) as error:
        raise ValueError(
            f"network paired metric {name!r} must be numeric"
        ) from error
    if not np.isfinite(value):
        raise ValueError(
            f"network paired metric {name!r} must be finite"
        )
    return value


def _boolean_metric(row: Mapping[str, Any], name: str) -> bool:
    if name not in row:
        raise ValueError(
            f"network paired gates require metric {name!r}"
        )
    value = row[name]
    if isinstance(value, (bool, np.bool_)):
        return bool(value)
    if isinstance(value, (int, float, np.integer, np.floating)):
        numeric = float(value)
        if np.isfinite(numeric) and numeric in (0.0, 1.0):
            return bool(numeric)
    raise ValueError(
        f"network paired metric {name!r} must be boolean or 0/1"
    )


def _apply_network_paired_gates(
    rows: list[dict[str, Any]],
    metric_configuration: Mapping[str, Any],
) -> None:
    network_configuration = metric_configuration.get("network")
    if not isinstance(network_configuration, Mapping):
        return
    paired = network_configuration.get("paired_gates")
    if paired is None:
        return
    if not isinstance(paired, Mapping):
        raise ValueError("network.paired_gates must be a mapping")
    nominal_profile_id = paired.get("nominal_profile_id", "nominal")
    if not isinstance(nominal_profile_id, str) or not nominal_profile_id:
        raise ValueError(
            "network.paired_gates.nominal_profile_id must be non-empty"
        )
    maximum_tracking_delta = paired.get(
        "maximum_tracking_rmse_delta_vs_nominal_rad",
        paired.get("max_tracking_delta"),
    )
    maximum_contact_relative_change = paired.get(
        "maximum_abs_contact_rms_relative_change_vs_nominal",
        paired.get("max_abs_contact_relative_change"),
    )
    if maximum_tracking_delta is None:
        raise ValueError(
            "network paired gates require "
            "maximum_tracking_rmse_delta_vs_nominal_rad"
        )
    if maximum_contact_relative_change is None:
        raise ValueError(
            "network paired gates require "
            "maximum_abs_contact_rms_relative_change_vs_nominal"
        )
    maximum_tracking_delta = float(maximum_tracking_delta)
    maximum_contact_relative_change = float(
        maximum_contact_relative_change
    )
    contact_denominator_epsilon_N = float(
        paired.get("contact_rms_denominator_epsilon_N", 1e-12)
    )
    if (
        not np.isfinite(maximum_tracking_delta)
        or maximum_tracking_delta < 0.0
        or not np.isfinite(maximum_contact_relative_change)
        or maximum_contact_relative_change < 0.0
        or not np.isfinite(contact_denominator_epsilon_N)
        or contact_denominator_epsilon_N <= 0.0
    ):
        raise ValueError(
            "network paired-gate thresholds must be finite/non-negative "
            "and the contact denominator epsilon must be positive"
        )

    network_rows = [
        row for row in rows if "network_local_full_gate_pass" in row
    ]
    groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {}
    for row in network_rows:
        missing_identity = [
            name
            for name in ("network_pair_group_id", "network_profile_id")
            if row.get(name) in (None, "")
        ]
        if missing_identity:
            raise ValueError(
                "network paired gates require identity fields "
                f"{missing_identity} for trial {row.get('trial_id')}"
            )
        group_key = (
            row["network_pair_group_id"],
            row["method_id"],
            row["trajectory_id"],
            row["replicate"],
        )
        groups.setdefault(group_key, []).append(row)

    for group_key, group_rows in groups.items():
        nominal_rows = [
            row
            for row in group_rows
            if row["network_profile_id"] == nominal_profile_id
        ]
        if len(nominal_rows) != 1:
            reason = "missing" if not nominal_rows else "duplicate"
            raise ValueError(
                f"{reason} nominal network profile for group {group_key}; "
                f"expected exactly one {nominal_profile_id!r}"
            )
        nominal = nominal_rows[0]
        nominal_tracking = _finite_metric(
            nominal,
            "network_slave_zero_delay_tracking_rmse_rad",
        )
        nominal_feedback_lag = _finite_metric(
            nominal,
            "network_return_feedback_lag_rmse_Nm",
        )
        contact_expected = _boolean_metric(
            nominal,
            "network_contact_expected",
        )
        nominal_contact_rms: float | None = None
        contact_denominator: float | None = None
        if contact_expected:
            nominal_contact_rms = _finite_metric(
                nominal,
                "network_contact_force_rms_N",
            )
            contact_denominator = max(
                abs(nominal_contact_rms),
                contact_denominator_epsilon_N,
            )
        for row in group_rows:
            if (
                _boolean_metric(row, "network_contact_expected")
                != contact_expected
            ):
                raise ValueError(
                    "network paired gates require a consistent "
                    f"contact condition for group {group_key}"
                )
            if row is nominal:
                tracking_delta = 0.0
                feedback_lag_delta = 0.0
                contact_relative_change = (
                    0.0 if contact_expected else None
                )
            else:
                tracking_delta = (
                    _finite_metric(
                        row,
                        "network_slave_zero_delay_tracking_rmse_rad",
                    )
                    - nominal_tracking
                )
                feedback_lag_delta = (
                    _finite_metric(
                        row,
                        "network_return_feedback_lag_rmse_Nm",
                    )
                    - nominal_feedback_lag
                )
                contact_relative_change = (
                    (
                        _finite_metric(
                            row,
                            "network_contact_force_rms_N",
                        )
                        - nominal_contact_rms
                    )
                    / contact_denominator
                    if contact_expected
                    else None
                )
            tracking_gate = bool(
                tracking_delta <= maximum_tracking_delta
            )
            contact_gate = (
                bool(
                    abs(contact_relative_change)
                    <= maximum_contact_relative_change
                )
                if contact_expected
                else None
            )
            paired_gate = bool(
                tracking_gate
                and (contact_gate if contact_expected else True)
            )
            row.update(
                {
                    (
                        "network_slave_zero_delay_tracking_rmse_"
                        "delta_vs_nominal_rad"
                    ): tracking_delta,
                    (
                        "network_zero_delay_tracking_rmse_"
                        "delta_vs_nominal_rad"
                    ): tracking_delta,
                    (
                        "network_return_feedback_lag_rmse_"
                        "delta_vs_nominal_Nm"
                    ): feedback_lag_delta,
                    (
                        "network_contact_rms_relative_change_"
                        "vs_nominal"
                    ): contact_relative_change,
                    (
                        "network_contact_force_rms_relative_change_"
                        "vs_nominal"
                    ): contact_relative_change,
                    "network_paired_tracking_gate_pass": tracking_gate,
                    "network_paired_contact_gate_pass": contact_gate,
                    "network_paired_gate_pass": paired_gate,
                    "network_full_gate_pass": bool(
                        row["network_local_full_gate_pass"]
                        and paired_gate
                    ),
                }
            )


def generate_paper_source_data(
    batch_dir: Path,
    metric_configuration: Mapping[str, Any],
    *,
    output_root: Path | None = None,
) -> dict[str, Any]:
    """Recompute all configured metrics and emit traceable source-data tables."""
    batch_dir = Path(batch_dir)
    require_valid_batch(batch_dir, require_complete=True)
    plan = load_document(batch_dir / "plan.json")
    output_root = batch_dir if output_root is None else Path(output_root)
    derived_dir = output_root / "derived"
    source_dir = output_root / "paper" / "source_data"
    derived_dir.mkdir(parents=True, exist_ok=True)
    source_dir.mkdir(parents=True, exist_ok=True)

    rows: list[dict[str, Any]] = []
    input_files: dict[str, str] = {}
    for trial in plan["trials"]:
        trial_dir = batch_dir / "raw" / trial["trial_id"]
        sample_path = trial_dir / "samples.npz"
        trial_manifest_path = trial_dir / "trial_manifest.json"
        input_files[str(sample_path.relative_to(batch_dir))] = file_sha256(sample_path)
        input_files[
            str(trial_manifest_path.relative_to(batch_dir))
        ] = file_sha256(trial_manifest_path)
        trial_manifest = load_document(trial_manifest_path)
        trial_metadata = trial_manifest.get("metadata", {})
        if not isinstance(trial_metadata, Mapping):
            raise ValueError(
                f"{trial_manifest_path} metadata must be a mapping"
            )
        with np.load(sample_path, allow_pickle=False) as archive:
            samples = {name: archive[name] for name in archive.files}
        method_id = trial["method"]["method_id"]
        trial_configuration = _trial_metric_configuration(
            metric_configuration,
            method_id,
        )
        metrics = (
            {}
            if trial_configuration is None
            else derive_trial_metrics(samples, trial_configuration)
        )
        rows.append(
            {
                **_identity_row(trial, trial_metadata),
                **metrics,
            }
        )

    _apply_network_paired_gates(rows, metric_configuration)
    metric_path = derived_dir / "trial_metrics.jsonl"
    atomic_write_jsonl(metric_path, rows)
    table_files: dict[str, str] = {}
    family_row_counts: dict[str, int] = {}
    for family in metric_configuration["enabled"]:
        family_rows: list[dict[str, Any]] = []
        prefix = f"{family}_"
        for row in rows:
            provenance_identity_fields = set(PROVENANCE_IDENTITY_FIELDS)
            if not any(
                key.startswith(prefix)
                and key not in provenance_identity_fields
                for key in row
            ):
                continue
            selected = {
                key: value
                for key, value in row.items()
                if key in {
                    "trial_id",
                    "pair_id",
                    "study_id",
                    "split",
                    "method_id",
                    "trajectory_id",
                    "trajectory_family",
                    "replicate",
                    *PROVENANCE_IDENTITY_FIELDS,
                }
                or key.startswith("factor_")
                or key.startswith(prefix)
            }
            family_rows.append(selected)
        if not family_rows:
            raise ValueError(
                f"metric family {family!r} selected no methods in this plan"
            )
        table_path = source_dir / f"{family}.csv"
        _atomic_write_csv(table_path, family_rows)
        table_files[str(table_path.relative_to(output_root))] = file_sha256(table_path)
        family_row_counts[family] = len(family_rows)

    manifest = {
        "kind": "paper_artifact_manifest",
        "schema_version": SCHEMA_VERSION,
        "created_utc": utc_now(),
        "source_batch": str(batch_dir),
        "source_plan_hash": plan["plan_hash"],
        "metric_configuration": to_jsonable(metric_configuration),
        "metric_configuration_hash": stable_hash(
            metric_configuration, prefix="metric-configuration"
        ),
        "row_count": len(rows),
        "family_row_counts": family_row_counts,
        "row_hash": stable_hash(rows, prefix="trial-metric-rows"),
        "storage_formats": {
            "trial_metrics": STORAGE_FORMATS["trial_metrics"],
            "paper_source_data": STORAGE_FORMATS["paper_source_data"],
        },
        "inputs": input_files,
        "outputs": {
            str(metric_path.relative_to(output_root)): file_sha256(metric_path),
            **table_files,
        },
    }
    manifest_path = output_root / "paper" / "artifact_manifest.json"
    atomic_write_json(manifest_path, manifest)
    return manifest


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Recompute trial metrics and paper source-data tables"
    )
    parser.add_argument("--batch-dir", type=Path, required=True)
    parser.add_argument("--metric-config", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, default=None)
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    configuration = load_document(args.metric_config)
    manifest = generate_paper_source_data(
        args.batch_dir,
        configuration,
        output_root=args.output_root,
    )
    print(json.dumps(manifest, indent=2, ensure_ascii=False, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())