exoskeleton/code/experiments/validate.py

"""Independent validation of plans, manifests, and per-trial artifacts."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Mapping

import numpy as np

from .hashing import file_sha256, stable_hash
from .plan import load_document, validate_trial_plan
from .schema import ContractError, require_schema


def validate_trial_directory(
    trial_dir: Path,
    trial: Mapping[str, Any],
) -> list[str]:
    errors: list[str] = []
    trial_dir = Path(trial_dir)
    manifest_path = trial_dir / "trial_manifest.json"
    if not manifest_path.exists():
        return [f"{trial_dir}: missing trial_manifest.json"]
    try:
        manifest = load_document(manifest_path)
        require_schema(manifest, "trial_manifest")
    except Exception as exc:
        return [f"{trial_dir}: invalid manifest: {exc}"]
    if manifest.get("trial_id") != trial.get("trial_id"):
        errors.append(f"{trial_dir}: trial_id mismatch")
    if manifest.get("pair_id") != trial.get("pair_id"):
        errors.append(f"{trial_dir}: pair_id mismatch")
    if manifest.get("trial_spec_hash") != trial.get("trial_spec_hash"):
        errors.append(f"{trial_dir}: trial_spec_hash mismatch")
    if manifest.get("status") != "completed":
        errors.append(f"{trial_dir}: status is not completed")

    declared_files = manifest.get("files", {})
    for relative, expected_hash in declared_files.items():
        path = trial_dir / relative
        if not path.is_file():
            errors.append(f"{trial_dir}: missing {relative}")
        elif file_sha256(path) != expected_hash:
            errors.append(f"{trial_dir}: hash mismatch for {relative}")

    sample_path = trial_dir / "samples.npz"
    if sample_path.exists():
        try:
            with np.load(sample_path, allow_pickle=False) as archive:
                expected_count = int(manifest.get("sample_count", -1))
                declared_fields = manifest.get("sample_fields", {})
                if set(archive.files) != set(declared_fields):
                    errors.append(f"{trial_dir}: NPZ fields differ from manifest")
                for name in archive.files:
                    array = archive[name]
                    if array.ndim == 0 or array.shape[0] != expected_count:
                        errors.append(
                            f"{trial_dir}: invalid sample axis for field {name}"
                        )
                    expected = declared_fields.get(name, {})
                    if list(array.shape) != expected.get("shape"):
                        errors.append(f"{trial_dir}: shape mismatch for field {name}")
                    if str(array.dtype) != expected.get("dtype"):
                        errors.append(f"{trial_dir}: dtype mismatch for field {name}")
        except Exception as exc:
            errors.append(f"{trial_dir}: cannot read samples.npz: {exc}")

    event_path = trial_dir / "events.jsonl"
    if event_path.exists():
        try:
            with event_path.open("r", encoding="utf-8") as stream:
                for line_number, line in enumerate(stream, start=1):
                    if line.strip():
                        value = json.loads(line)
                        if not isinstance(value, dict):
                            raise ValueError(f"line {line_number} is not an object")
        except Exception as exc:
            errors.append(f"{trial_dir}: invalid events.jsonl: {exc}")
    return errors


def validate_batch(
    batch_dir: Path,
    *,
    require_complete: bool = True,
) -> dict[str, Any]:
    batch_dir = Path(batch_dir)
    errors: list[str] = []
    warnings: list[str] = []
    try:
        plan = load_document(batch_dir / "plan.json")
        validate_trial_plan(plan)
    except Exception as exc:
        return {
            "valid": False,
            "errors": [f"invalid plan: {exc}"],
            "warnings": [],
            "expected_trials": 0,
            "completed_trials": 0,
        }
    try:
        manifest = load_document(batch_dir / "batch_manifest.json")
        require_schema(manifest, "batch_manifest")
        if manifest.get("plan_hash") != plan.get("plan_hash"):
            errors.append("batch manifest plan_hash mismatch")
        hash_basis = dict(manifest)
        recorded_manifest_hash = hash_basis.pop("manifest_hash", None)
        hash_basis.pop("created_utc", None)
        expected_manifest_hash = stable_hash(
            hash_basis,
            prefix="batch-manifest",
        )
        if recorded_manifest_hash != expected_manifest_hash:
            errors.append("batch manifest_hash mismatch")
    except Exception as exc:
        errors.append(f"invalid batch manifest: {exc}")

    completed = 0
    raw_dir = batch_dir / "raw"
    for trial in plan["trials"]:
        trial_dir = raw_dir / trial["trial_id"]
        if not trial_dir.exists():
            if require_complete:
                errors.append(f"missing trial {trial['trial_id']}")
            continue
        trial_errors = validate_trial_directory(trial_dir, trial)
        errors.extend(trial_errors)
        if not trial_errors:
            completed += 1
    if raw_dir.exists():
        expected = {trial["trial_id"] for trial in plan["trials"]}
        for entry in raw_dir.iterdir():
            if entry.name.startswith("."):
                warnings.append(f"temporary trial artifact present: {entry.name}")
            elif entry.is_dir() and entry.name not in expected:
                warnings.append(f"unexpected trial directory: {entry.name}")
    return {
        "valid": not errors,
        "errors": errors,
        "warnings": warnings,
        "expected_trials": plan["trial_count"],
        "completed_trials": completed,
    }


def require_valid_batch(batch_dir: Path, *, require_complete: bool = True) -> None:
    report = validate_batch(batch_dir, require_complete=require_complete)
    if not report["valid"]:
        raise ContractError("; ".join(report["errors"]))