#!/usr/bin/env python3
"""Generic manifest, execution, recovery, and result CLI for SpectraREML."""

from __future__ import annotations

import argparse
import csv
import gzip
import hashlib
import json
import os
import shutil
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Mapping, NoReturn, Sequence

for _name in (
    "OMP_NUM_THREADS",
    "OPENBLAS_NUM_THREADS",
    "MKL_NUM_THREADS",
    "NUMEXPR_NUM_THREADS",
):
    os.environ.setdefault(_name, "1")

try:
    import numpy as np
except ImportError as exc:  # pragma: no cover
    raise SystemExit("NumPy is required: {}".format(exc))


MANIFEST_FORMAT = "spectra-reml-manifest-v2"
BLOCK_FORMAT = "spectra-reml-block-v4"
RUN_SIGNATURE_FORMAT = "spectra-reml-run-signature-v3"
FINALIZE_FORMAT = "spectra-reml-finalize-v3"
TASK_HEADER = (
    "task_index",
    "task_id",
    "phenotype_row",
    "n_extra_covariates",
)
SUMMARY_HEADER = (
    "task_index",
    "task_id",
    "status",
    "n_fixed",
    "n_extra_covariates",
    "beta_offset",
    "cov_offset",
    "sigma_g2",
    "sigma_e2",
    "h2",
    "logL",
    "iterations",
    "line_search_steps",
    "grad_inf",
    "fixed_test_method",
    "fixed_test_status",
    "fixed_test_offset",
    "contrast_test_offset",
    "extra_joint_num_df",
    "extra_joint_den_df",
    "extra_joint_f",
    "extra_joint_p",
    "fixed_test_error",
    "error",
)
SUCCESS_STATUSES = frozenset(("converged", "converged_boundary"))
BLOCK_PATTERNS = (
    "block_*.summary.tsv",
    "block_*.beta.f64.bin",
    "block_*.cov.f64.bin",
    "block_*.fixed_se.f64.bin",
    "block_*.fixed_stat.f64.bin",
    "block_*.fixed_ddf.f64.bin",
    "block_*.fixed_p.f64.bin",
    "block_*.contrast_estimate.f64.bin",
    "block_*.contrast_se.f64.bin",
    "block_*.contrast_stat.f64.bin",
    "block_*.contrast_numdf.f64.bin",
    "block_*.contrast_ddf.f64.bin",
    "block_*.contrast_p.f64.bin",
    "block_*.complete",
)


def fail(message: str) -> NoReturn:
    raise RuntimeError(message)


def log(message: str) -> None:
    print(time.strftime("%Y-%m-%d %H:%M:%S"), "|", message, flush=True)


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def resolved(path: Path) -> Path:
    return path.expanduser().resolve()


def file_identity(path: Path) -> dict[str, Any]:
    path = resolved(path)
    stat = path.stat()
    return {
        "path": str(path),
        "size_bytes": int(stat.st_size),
        "mtime_ns": int(stat.st_mtime_ns),
    }


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def canonical_sha256(value: Any) -> str:
    encoded = json.dumps(
        value,
        ensure_ascii=False,
        allow_nan=False,
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def atomic_json(path: Path, value: Mapping[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(path.name + ".tmp.{}".format(os.getpid()))
    try:
        with temporary.open("w", encoding="utf-8", newline="") as handle:
            json.dump(value, handle, ensure_ascii=False, allow_nan=False, indent=2)
            handle.write("\n")
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, path)
    finally:
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass


def read_json(path: Path) -> dict[str, Any]:
    try:
        with path.open("r", encoding="utf-8-sig") as handle:
            value = json.load(handle)
    except (OSError, json.JSONDecodeError) as exc:
        fail("Cannot read JSON {}: {}".format(path, exc))
    if not isinstance(value, dict):
        fail("JSON root is not an object: {}".format(path))
    return value


def require_file(path: Path, label: str, allow_empty: bool = False) -> None:
    if not path.is_file():
        fail("{} is missing: {}".format(label, path))
    if not allow_empty and path.stat().st_size == 0:
        fail("{} is empty: {}".format(label, path))


def require_size(path: Path, expected: int, label: str) -> None:
    require_file(path, label, allow_empty=(expected == 0))
    observed = path.stat().st_size
    if observed != expected:
        fail("{} size mismatch: observed {}; expected {}: {}".format(
            label, observed, expected, path
        ))


def manifest_paths(manifest: Mapping[str, Any]) -> dict[str, Path | None]:
    raw = manifest.get("paths")
    if not isinstance(raw, dict):
        fail("Manifest has no paths object.")
    required = (
        "grm_bin",
        "base_x_f64",
        "phenotype_f64",
        "tasks_tsv",
        "extra_offsets_i64",
        "extra_indices_i32",
        "output_directory",
        "contrast_matrix_f64",
        "contrast_metadata_tsv",
    )
    missing = [name for name in required if name not in raw]
    if missing:
        fail("Manifest paths are missing: {}".format(", ".join(missing)))
    result: dict[str, Path | None] = {}
    for name in required + ("grm_id", "extra_covariate_f32"):
        value = raw.get(name)
        result[name] = Path(str(value)) if value not in (None, "") else None
    return result


def read_tasks(path: Path) -> list[dict[str, str]]:
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        reader = csv.DictReader(handle, delimiter="\t")
        if tuple(reader.fieldnames or ()) != TASK_HEADER:
            fail("Task header must be exactly: {}".format("\t".join(TASK_HEADER)))
        return list(reader)


def read_contrast_metadata(path: Path) -> list[dict[str, str]]:
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        reader = csv.DictReader(handle, delimiter="\t")
        expected = ("contrast_index", "contrast_id", "rhs")
        if tuple(reader.fieldnames or ()) != expected:
            fail("Contrast metadata header must be exactly: {}".format("\t".join(expected)))
        rows = list(reader)
    seen: set[str] = set()
    for expected_index, row in enumerate(rows):
        try:
            index = int(row["contrast_index"])
            rhs = float(row["rhs"])
        except ValueError as exc:
            fail("Invalid contrast metadata at row {}: {}".format(expected_index, exc))
        if index != expected_index:
            fail("Contrast indices must be consecutive and zero based.")
        if not row["contrast_id"] or row["contrast_id"] in seen:
            fail("Contrast IDs must be nonempty and unique.")
        if not np.isfinite(rhs):
            fail("Contrast rhs must be finite.")
        seen.add(row["contrast_id"])
    return rows


def validate_manifest(path: Path) -> dict[str, Any]:
    path = resolved(path)
    manifest = read_json(path)
    if manifest.get("format") != MANIFEST_FORMAT:
        fail("Unsupported manifest format: {!r}".format(manifest.get("format")))
    dimensions = manifest.get("dimensions")
    if not isinstance(dimensions, dict):
        fail("Manifest has no dimensions object.")
    names = (
        "sample_count",
        "base_covariate_count",
        "phenotype_row_count",
        "extra_covariate_row_count",
        "task_count",
        "extra_index_count",
        "contrast_count",
    )
    parsed: dict[str, int] = {}
    for name in names:
        value = dimensions.get(name)
        if isinstance(value, bool) or not isinstance(value, int) or value < 0:
            fail("Invalid dimension {}={!r}".format(name, value))
        parsed[name] = value
    if min(parsed["sample_count"], parsed["base_covariate_count"], parsed["task_count"]) <= 0:
        fail("Sample, base-covariate, and task counts must be positive.")

    paths = manifest_paths(manifest)
    n = parsed["sample_count"]
    require_size(paths["grm_bin"], n * (n + 1) // 2 * 4, "GRM")  # type: ignore[arg-type]
    require_size(paths["base_x_f64"], n * parsed["base_covariate_count"] * 8, "base design")  # type: ignore[arg-type]
    require_size(paths["phenotype_f64"], n * parsed["phenotype_row_count"] * 8, "phenotypes")  # type: ignore[arg-type]
    require_size(paths["extra_offsets_i64"], (parsed["task_count"] + 1) * 8, "extra offsets")  # type: ignore[arg-type]
    require_size(paths["extra_indices_i32"], parsed["extra_index_count"] * 4, "extra indices")  # type: ignore[arg-type]
    if parsed["extra_covariate_row_count"]:
        if paths["extra_covariate_f32"] is None:
            fail("extra_covariate_f32 is required when its row count is nonzero.")
        require_size(paths["extra_covariate_f32"], n * parsed["extra_covariate_row_count"] * 4, "extra covariates")
    if paths["grm_id"] is not None:
        require_file(paths["grm_id"], "GRM IDs")
        with paths["grm_id"].open("r", encoding="utf-8-sig") as handle:
            if sum(bool(line.strip()) for line in handle) != n:
                fail("GRM ID row count differs from sample_count.")
    contrast_paths = (paths["contrast_matrix_f64"], paths["contrast_metadata_tsv"])
    if parsed["contrast_count"]:
        if any(path is None for path in contrast_paths):
            fail("Both contrast_matrix_f64 and contrast_metadata_tsv are required.")
        require_size(
            paths["contrast_matrix_f64"],
            parsed["contrast_count"] * parsed["base_covariate_count"] * 8,
            "contrast matrix",
        )
        contrast_rows = read_contrast_metadata(paths["contrast_metadata_tsv"])
        if len(contrast_rows) != parsed["contrast_count"]:
            fail("Contrast metadata row count differs from contrast_count.")
        contrast_matrix = np.fromfile(paths["contrast_matrix_f64"], dtype="<f8")
        contrast_matrix = contrast_matrix.reshape(
            parsed["contrast_count"], parsed["base_covariate_count"]
        )
        if not np.all(np.isfinite(contrast_matrix)):
            fail("Contrast matrix contains a non-finite value.")
        if np.any(np.all(contrast_matrix == 0.0, axis=1)):
            fail("Every contrast row must contain at least one nonzero weight.")
    elif any(path is not None for path in contrast_paths):
        fail("Contrast paths must be omitted when contrast_count is zero.")

    sources = manifest.get("source_files")
    if not isinstance(sources, dict) or not sources:
        fail("Manifest has no source identities.")
    for label, expected in sources.items():
        if not isinstance(expected, dict):
            fail("Malformed source identity: {}".format(label))
        source = Path(str(expected.get("path", "")))
        if not source.is_file():
            fail("Source is missing: {}: {}".format(label, source))
        observed = source.stat()
        if (
            observed.st_size != int(expected.get("size_bytes", -1))
            or observed.st_mtime_ns != int(expected.get("mtime_ns", -1))
        ):
            fail("Source changed since manifest creation: {}: {}".format(label, source))

    controls = {
        "base_x_f64": paths["base_x_f64"],
        "tasks_tsv": paths["tasks_tsv"],
        "extra_offsets_i64": paths["extra_offsets_i64"],
        "extra_indices_i32": paths["extra_indices_i32"],
    }
    if parsed["contrast_count"]:
        controls["contrast_matrix_f64"] = paths["contrast_matrix_f64"]
        controls["contrast_metadata_tsv"] = paths["contrast_metadata_tsv"]
    hashes = manifest.get("control_sha256")
    if not isinstance(hashes, dict) or set(hashes) != set(controls):
        fail("control_sha256 does not cover all control files.")
    for name, control in controls.items():
        assert control is not None
        if sha256_file(control) != hashes[name]:
            fail("Control checksum mismatch: {}".format(control))

    tasks = read_tasks(paths["tasks_tsv"])  # type: ignore[arg-type]
    if len(tasks) != parsed["task_count"]:
        fail("Task count differs from manifest.")
    offsets = np.fromfile(paths["extra_offsets_i64"], dtype="<i8")
    indices = np.fromfile(paths["extra_indices_i32"], dtype="<i4")
    if offsets[0] != 0 or np.any(offsets[1:] < offsets[:-1]) or int(offsets[-1]) != indices.size:
        fail("Extra-covariate CSR offsets are invalid.")
    if indices.size and (
        int(indices.min()) < 0
        or int(indices.max()) >= parsed["extra_covariate_row_count"]
    ):
        fail("An extra-covariate index is outside its matrix.")
    seen_ids: set[str] = set()
    for expected, row in enumerate(tasks):
        try:
            task_index = int(row["task_index"])
            phenotype_row = int(row["phenotype_row"])
            extra_count = int(row["n_extra_covariates"])
        except ValueError as exc:
            fail("Invalid task integer at row {}: {}".format(expected, exc))
        if task_index != expected:
            fail("Task indices must be consecutive and zero based.")
        if not row["task_id"] or row["task_id"] in seen_ids:
            fail("Task IDs must be nonempty and unique.")
        seen_ids.add(row["task_id"])
        if phenotype_row < 0 or phenotype_row >= parsed["phenotype_row_count"]:
            fail("Task phenotype row is outside its matrix.")
        if extra_count != int(offsets[expected + 1] - offsets[expected]):
            fail("Task extra count disagrees with CSR offsets.")
    return manifest


def make_manifest(args: argparse.Namespace) -> Path:
    target = resolved(args.manifest)
    paths: dict[str, Path | None] = {
        "grm_bin": resolved(args.grm_bin),
        "grm_id": resolved(args.grm_id) if args.grm_id else None,
        "base_x_f64": resolved(args.base_x),
        "phenotype_f64": resolved(args.phenotypes),
        "extra_covariate_f32": resolved(args.extra_covariates) if args.extra_covariates else None,
        "tasks_tsv": resolved(args.tasks),
        "extra_offsets_i64": resolved(args.extra_offsets),
        "extra_indices_i32": resolved(args.extra_indices),
        "contrast_matrix_f64": resolved(args.contrast_matrix) if args.contrast_matrix else None,
        "contrast_metadata_tsv": resolved(args.contrast_metadata) if args.contrast_metadata else None,
        "output_directory": resolved(args.output_dir),
    }
    tasks = read_tasks(paths["tasks_tsv"])  # type: ignore[arg-type]
    extra_index_count = paths["extra_indices_i32"].stat().st_size // 4  # type: ignore[union-attr]
    source_names = ("grm_bin", "phenotype_f64")
    sources = {name: file_identity(paths[name]) for name in source_names}  # type: ignore[arg-type]
    for name in ("grm_id", "extra_covariate_f32"):
        if paths[name] is not None:
            sources[name] = file_identity(paths[name])  # type: ignore[arg-type]
    controls = ["base_x_f64", "tasks_tsv", "extra_offsets_i64", "extra_indices_i32"]
    if args.n_contrasts:
        if paths["contrast_matrix_f64"] is None or paths["contrast_metadata_tsv"] is None:
            fail("--contrast-matrix and --contrast-metadata are required when --n-contrasts is positive.")
        controls.extend(("contrast_matrix_f64", "contrast_metadata_tsv"))
    elif paths["contrast_matrix_f64"] is not None or paths["contrast_metadata_tsv"] is not None:
        fail("Contrast paths require a positive --n-contrasts.")
    manifest = {
        "format": MANIFEST_FORMAT,
        "created_utc": utc_now(),
        "dimensions": {
            "sample_count": args.n_samples,
            "base_covariate_count": args.n_base_covariates,
            "phenotype_row_count": args.n_phenotype_rows,
            "extra_covariate_row_count": args.n_extra_covariate_rows,
            "task_count": len(tasks),
            "extra_index_count": extra_index_count,
            "contrast_count": args.n_contrasts,
        },
        "paths": {name: str(value) if value is not None else None for name, value in paths.items()},
        "source_files": sources,
        "control_sha256": {name: sha256_file(paths[name]) for name in controls},  # type: ignore[arg-type]
    }
    atomic_json(target, manifest)
    validate_manifest(target)
    log("manifest created: {} tasks -> {}".format(len(tasks), target))
    return target


def resolve_engine(value: str) -> Path:
    candidate = Path(value).expanduser()
    if candidate.is_file():
        return candidate.resolve()
    found = shutil.which(value)
    if found:
        return Path(found).resolve()
    fail("SpectraREML engine was not found: {}".format(value))


def build_engine_command(engine: Path, manifest: Mapping[str, Any], args: argparse.Namespace) -> list[str]:
    paths = manifest_paths(manifest)
    dims = manifest["dimensions"]
    command = [
        str(engine),
        "--grm-bin", str(paths["grm_bin"]),
        "--base-x", str(paths["base_x_f64"]),
        "--phenotypes", str(paths["phenotype_f64"]),
        "--tasks", str(paths["tasks_tsv"]),
        "--extra-offsets", str(paths["extra_offsets_i64"]),
        "--extra-indices", str(paths["extra_indices_i32"]),
        "--out-dir", str(paths["output_directory"]),
        "--n-samples", str(dims["sample_count"]),
        "--n-base-covariates", str(dims["base_covariate_count"]),
        "--n-phenotype-rows", str(dims["phenotype_row_count"]),
        "--n-extra-covariate-rows", str(dims["extra_covariate_row_count"]),
        "--n-contrasts", str(dims["contrast_count"]),
        "--block-size", str(args.block_size),
        "--threads", str(args.threads),
        "--fixed-effect-test", args.fixed_effect_test,
    ]
    if paths["grm_id"] is not None:
        command.extend(("--grm-id", str(paths["grm_id"])))
    if int(dims["extra_index_count"]):
        command.extend(("--extra-covariates", str(paths["extra_covariate_f32"])))
    if int(dims["contrast_count"]):
        command.extend(("--contrast-matrix", str(paths["contrast_matrix_f64"])))
        command.extend(("--contrast-metadata", str(paths["contrast_metadata_tsv"])))
    option_map = (
        ("max_iterations", "--max-iterations"),
        ("line_search_max_evals", "--line-search-max-evals"),
        ("line_search_max_zoom", "--line-search-max-zoom"),
        ("gradient_abs_tol", "--gradient-abs-tol"),
        ("gradient_rel_tol", "--gradient-rel-tol"),
        ("step_rel_tol", "--step-rel-tol"),
        ("likelihood_rel_tol", "--likelihood-rel-tol"),
        ("wolfe_c1", "--wolfe-c1"),
        ("wolfe_c2", "--wolfe-c2"),
        ("initial_step", "--initial-step"),
        ("maximum_step", "--maximum-step"),
        ("line_search_expansion", "--line-search-expansion"),
        ("zoom_safeguard", "--zoom-safeguard"),
        ("rank_tol", "--rank-tol"),
        ("covariance_floor", "--covariance-floor"),
        ("boundary_h2_trigger", "--boundary-h2-trigger"),
        ("boundary_score_tol", "--boundary-score-tol"),
        ("boundary_logl_rel_tol", "--boundary-logl-rel-tol"),
    )
    for attribute, option in option_map:
        command.extend((option, repr(getattr(args, attribute))))
    for attribute, option in (("initial_sigma_e", "--initial-sigma-e"), ("initial_sigma_g", "--initial-sigma-g")):
        value = getattr(args, attribute)
        if value is not None:
            command.extend((option, repr(value)))
    if args.resume:
        command.append("--resume")
    elif args.force:
        command.append("--overwrite")
    return command


def engine_identity(engine: Path) -> dict[str, Any]:
    value = file_identity(engine)
    value["sha256"] = sha256_file(engine)
    return value


def signature_payload(manifest: Mapping[str, Any], engine: Path, args: argparse.Namespace) -> dict[str, Any]:
    excluded = {"action", "manifest", "engine", "dry_run", "output"}
    options = {
        key: value for key, value in vars(args).items()
        if key not in excluded and key not in {"resume", "force"}
    }
    return {
        "manifest_sha256": canonical_sha256(manifest),
        "engine": engine_identity(engine),
        "options": options,
    }


def make_signature(manifest: Mapping[str, Any], engine: Path, args: argparse.Namespace) -> dict[str, Any]:
    payload = signature_payload(manifest, engine, args)
    return {
        "format": RUN_SIGNATURE_FORMAT,
        "created_utc": utc_now(),
        "signature_sha256": canonical_sha256(payload),
        "payload": payload,
    }


def read_signature(path: Path) -> dict[str, Any]:
    value = read_json(path)
    payload = value.get("payload")
    if value.get("format") != RUN_SIGNATURE_FORMAT or not isinstance(payload, dict):
        fail("Invalid run signature: {}".format(path))
    if value.get("signature_sha256") != canonical_sha256(payload):
        fail("Run signature integrity check failed: {}".format(path))
    return value


def block_files(output: Path) -> list[Path]:
    values: set[Path] = set()
    for pattern in BLOCK_PATTERNS:
        values.update(path for path in output.glob(pattern) if path.is_file())
    return sorted(values)


def validate_resume_layout(output: Path, task_count: int, block_size: int) -> None:
    for marker in output.glob("block_*.complete"):
        digits = marker.name[len("block_"):-len(".complete")]
        if len(digits) != 6 or not digits.isdigit():
            fail("Invalid block marker: {}".format(marker))
        index = int(digits)
        expected = min(block_size, max(0, task_count - index * block_size))
        values: dict[str, str] = {}
        with marker.open("r", encoding="utf-8") as handle:
            for line in handle:
                fields = line.rstrip("\r\n").split("\t", 1)
                if len(fields) != 2:
                    fail("Malformed marker: {}".format(marker))
                values[fields[0]] = fields[1]
        if expected <= 0 or values.get("format") != BLOCK_FORMAT or int(values.get("tasks", -1)) != expected:
            fail("Block marker is incompatible with this manifest/block size: {}".format(marker))


def read_marker(path: Path) -> dict[str, str]:
    values: dict[str, str] = {}
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            fields = line.rstrip("\r\n").split("\t", 1)
            if len(fields) != 2 or not fields[0] or fields[0] in values:
                fail("Malformed completion marker: {}".format(path))
            values[fields[0]] = fields[1]
    return values


def run_engine(manifest: Mapping[str, Any], args: argparse.Namespace) -> None:
    if args.dry_run:
        print(json.dumps(build_engine_command(Path(args.engine), manifest, args), ensure_ascii=False))
        return
    engine = resolve_engine(args.engine)
    expected_signature = make_signature(manifest, engine, args)
    paths = manifest_paths(manifest)
    output = paths["output_directory"]
    assert output is not None
    output.mkdir(parents=True, exist_ok=True)
    signature_path = output / "run.signature.json"
    markers = list(output.glob("block_*.complete"))
    if args.force:
        try:
            signature_path.unlink()
        except FileNotFoundError:
            pass
        for path in block_files(output):
            path.unlink()
        atomic_json(signature_path, expected_signature)
    elif args.resume:
        if markers:
            if not signature_path.is_file():
                fail("Completed blocks have no run signature; use --force.")
            observed = read_signature(signature_path)
            if observed["signature_sha256"] != expected_signature["signature_sha256"]:
                fail("Run signature differs from engine, manifest, or options; use --force.")
            validate_resume_layout(output, int(manifest["dimensions"]["task_count"]), args.block_size)
        else:
            atomic_json(signature_path, expected_signature)
    else:
        if block_files(output):
            fail("Block outputs already exist; use --resume or --force.")
        atomic_json(signature_path, expected_signature)

    command = build_engine_command(engine, manifest, args)
    environment = os.environ.copy()
    environment.update({
        "OMP_NUM_THREADS": str(args.threads),
        "MKL_NUM_THREADS": str(args.blas_threads),
        "OPENBLAS_NUM_THREADS": str(args.blas_threads),
        "OMP_DYNAMIC": "FALSE",
    })
    log("starting SpectraREML engine")
    with (output / "engine.log").open("a" if args.resume else "w", encoding="utf-8") as handle:
        handle.write("# {}\n# command={}\n".format(utc_now(), json.dumps(command, ensure_ascii=False)))
        handle.flush()
        completed = subprocess.run(command, stdout=handle, stderr=subprocess.STDOUT, env=environment, check=False)
    atomic_json(output / "run.json", {
        "finished_utc": utc_now(),
        "return_code": completed.returncode,
        "command": command,
        "run_signature_sha256": expected_signature["signature_sha256"],
    })
    if completed.returncode:
        fail("SpectraREML failed with exit code {}; see engine.log".format(completed.returncode))


def finalizable_signature(manifest: Mapping[str, Any]) -> dict[str, Any]:
    output = manifest_paths(manifest)["output_directory"]
    assert output is not None
    path = output / "run.signature.json"
    if not path.is_file():
        fail("Cannot finalize without a run signature.")
    value = read_signature(path)
    if value["payload"].get("manifest_sha256") != canonical_sha256(manifest):
        fail("Run signature does not belong to the current manifest.")
    return value


def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
    output = manifest_paths(manifest)["output_directory"]
    assert output is not None
    task_count = int(manifest["dimensions"]["task_count"])
    results: dict[int, dict[str, Any]] = {}
    for marker in sorted(output.glob("block_*.complete")):
        stem = marker.name[:-len(".complete")]
        summary_path = output / (stem + ".summary.tsv")
        beta_path = output / (stem + ".beta.f64.bin")
        cov_path = output / (stem + ".cov.f64.bin")
        fixed_se_path = output / (stem + ".fixed_se.f64.bin")
        fixed_stat_path = output / (stem + ".fixed_stat.f64.bin")
        fixed_ddf_path = output / (stem + ".fixed_ddf.f64.bin")
        fixed_p_path = output / (stem + ".fixed_p.f64.bin")
        contrast_estimate_path = output / (stem + ".contrast_estimate.f64.bin")
        contrast_se_path = output / (stem + ".contrast_se.f64.bin")
        contrast_stat_path = output / (stem + ".contrast_stat.f64.bin")
        contrast_numdf_path = output / (stem + ".contrast_numdf.f64.bin")
        contrast_ddf_path = output / (stem + ".contrast_ddf.f64.bin")
        contrast_p_path = output / (stem + ".contrast_p.f64.bin")
        marker_values = read_marker(marker)
        block_format = marker_values.get("format")
        try:
            declared_tasks = int(marker_values["tasks"])
            beta_elements = int(marker_values["beta_elements"])
            cov_elements = int(marker_values["cov_elements"])
            fixed_test_elements = int(marker_values["fixed_test_elements"])
            contrast_test_elements = int(marker_values["contrast_test_elements"])
        except (KeyError, ValueError) as exc:
            fail("Malformed completion counts in {}: {}".format(marker, exc))
        if block_format != BLOCK_FORMAT or min(
            declared_tasks, beta_elements, cov_elements, fixed_test_elements,
            contrast_test_elements,
        ) < 0:
            fail("Invalid completion marker: {}".format(marker))
        require_size(beta_path, beta_elements * 8, "block beta")
        require_size(cov_path, cov_elements * 8, "block covariance")
        require_size(fixed_se_path, fixed_test_elements * 8, "fixed-effect SE")
        require_size(fixed_stat_path, fixed_test_elements * 8, "fixed-effect statistic")
        require_size(fixed_ddf_path, fixed_test_elements * 8, "fixed-effect denominator df")
        require_size(fixed_p_path, fixed_test_elements * 8, "fixed-effect p-value")
        for path, label in (
            (contrast_estimate_path, "contrast estimate"),
            (contrast_se_path, "contrast SE"),
            (contrast_stat_path, "contrast statistic"),
            (contrast_numdf_path, "contrast numerator df"),
            (contrast_ddf_path, "contrast denominator df"),
            (contrast_p_path, "contrast p-value"),
        ):
            require_size(path, contrast_test_elements * 8, label)
        with summary_path.open("r", encoding="utf-8-sig", newline="") as handle:
            reader = csv.DictReader(handle, delimiter="\t")
            if tuple(reader.fieldnames or ()) != SUMMARY_HEADER:
                fail("Unexpected summary header: {}".format(summary_path))
            rows = list(reader)
        if len(rows) != declared_tasks:
            fail("Summary row count disagrees with marker: {}".format(marker))
        beta = np.fromfile(beta_path, dtype="<f8")
        cov = np.fromfile(cov_path, dtype="<f8")
        fixed_se = np.fromfile(fixed_se_path, dtype="<f8")
        fixed_stat = np.fromfile(fixed_stat_path, dtype="<f8")
        fixed_ddf = np.fromfile(fixed_ddf_path, dtype="<f8")
        fixed_p = np.fromfile(fixed_p_path, dtype="<f8")
        contrast_estimate = np.fromfile(contrast_estimate_path, dtype="<f8")
        contrast_se = np.fromfile(contrast_se_path, dtype="<f8")
        contrast_stat = np.fromfile(contrast_stat_path, dtype="<f8")
        contrast_numdf = np.fromfile(contrast_numdf_path, dtype="<f8")
        contrast_ddf = np.fromfile(contrast_ddf_path, dtype="<f8")
        contrast_p = np.fromfile(contrast_p_path, dtype="<f8")
        next_beta_offset = 0
        next_cov_offset = 0
        next_fixed_test_offset = 0
        next_contrast_test_offset = 0
        q = int(manifest["dimensions"]["contrast_count"])
        for row in rows:
            index = int(row["task_index"])
            if index in results or index < 0 or index >= task_count:
                fail("Duplicate/out-of-range task index: {}".format(index))
            p = int(row["n_fixed"])
            beta_offset, cov_offset = int(row["beta_offset"]), int(row["cov_offset"])
            fixed_test_offset = int(row["fixed_test_offset"])
            contrast_test_offset = int(row["contrast_test_offset"])
            if beta_offset < 0 or cov_offset < 0:
                if beta_offset != cov_offset:
                    fail("Only one output offset is negative for task {}".format(index))
                beta_values: list[float] = []
                cov_values: list[float] = []
            else:
                packed = p * (p + 1) // 2
                if beta_offset != next_beta_offset or cov_offset != next_cov_offset:
                    fail("Non-contiguous output offsets for task {}".format(index))
                if beta_offset + p > beta.size or cov_offset + packed > cov.size:
                    fail("Output offsets exceed binary arrays for task {}".format(index))
                beta_values = beta[beta_offset:beta_offset + p].tolist()
                cov_values = cov[cov_offset:cov_offset + packed].tolist()
                next_beta_offset += p
                next_cov_offset += packed
            if fixed_test_offset < 0:
                fixed_se_values: list[float] = []
                fixed_stat_values: list[float] = []
                fixed_ddf_values: list[float] = []
                fixed_p_values: list[float] = []
            else:
                if (fixed_test_offset != next_fixed_test_offset or
                        fixed_test_offset + p > fixed_se.size):
                    fail("Invalid fixed-effect test offset for task {}".format(index))
                fixed_se_values = fixed_se[fixed_test_offset:fixed_test_offset + p].tolist()
                fixed_stat_values = fixed_stat[fixed_test_offset:fixed_test_offset + p].tolist()
                fixed_ddf_values = fixed_ddf[fixed_test_offset:fixed_test_offset + p].tolist()
                fixed_p_values = fixed_p[fixed_test_offset:fixed_test_offset + p].tolist()
                next_fixed_test_offset += p
            if contrast_test_offset < 0:
                contrast_estimate_values: list[float] = []
                contrast_se_values: list[float] = []
                contrast_stat_values: list[float] = []
                contrast_numdf_values: list[float] = []
                contrast_ddf_values: list[float] = []
                contrast_p_values: list[float] = []
            else:
                if (contrast_test_offset != next_contrast_test_offset or
                        contrast_test_offset + q > contrast_estimate.size):
                    fail("Invalid contrast test offset for task {}".format(index))
                section = slice(contrast_test_offset, contrast_test_offset + q)
                contrast_estimate_values = contrast_estimate[section].tolist()
                contrast_se_values = contrast_se[section].tolist()
                contrast_stat_values = contrast_stat[section].tolist()
                contrast_numdf_values = contrast_numdf[section].tolist()
                contrast_ddf_values = contrast_ddf[section].tolist()
                contrast_p_values = contrast_p[section].tolist()
                next_contrast_test_offset += q
            result: dict[str, Any] = dict(row)
            result["beta_json"] = json.dumps(beta_values, separators=(",", ":"))
            result["covariance_packed_lower_json"] = json.dumps(cov_values, separators=(",", ":"))
            result["fixed_effect_se_json"] = json.dumps(fixed_se_values, separators=(",", ":"))
            result["fixed_effect_statistic_json"] = json.dumps(fixed_stat_values, separators=(",", ":"))
            result["fixed_effect_denominator_df_json"] = json.dumps(fixed_ddf_values, separators=(",", ":"))
            result["fixed_effect_p_value_json"] = json.dumps(fixed_p_values, separators=(",", ":"))
            result["contrast_estimate_json"] = json.dumps(contrast_estimate_values, separators=(",", ":"))
            result["contrast_standard_error_json"] = json.dumps(contrast_se_values, separators=(",", ":"))
            result["contrast_statistic_json"] = json.dumps(contrast_stat_values, separators=(",", ":"))
            result["contrast_numerator_df_json"] = json.dumps(contrast_numdf_values, separators=(",", ":"))
            result["contrast_denominator_df_json"] = json.dumps(contrast_ddf_values, separators=(",", ":"))
            result["contrast_p_value_json"] = json.dumps(contrast_p_values, separators=(",", ":"))
            results[index] = result
        if (next_beta_offset != beta_elements or next_cov_offset != cov_elements or
                next_fixed_test_offset != fixed_test_elements or
                next_contrast_test_offset != contrast_test_elements):
            fail("Block binary arrays contain unused elements: {}".format(stem))
    missing = sorted(set(range(task_count)).difference(results))
    if missing:
        fail("Completed blocks are missing {} tasks.".format(len(missing)))
    return [results[index] for index in range(task_count)]


def finalize(manifest: Mapping[str, Any], output: Path) -> Path:
    signature = finalizable_signature(manifest)
    rows = read_completed(manifest)
    output = resolved(output)
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_name(output.name + ".tmp.{}".format(os.getpid()))
    opener = gzip.open if output.suffix == ".gz" else open
    fields = list(SUMMARY_HEADER) + [
        "beta_json", "covariance_packed_lower_json", "fixed_effect_se_json",
        "fixed_effect_statistic_json", "fixed_effect_denominator_df_json",
        "fixed_effect_p_value_json",
        "contrast_estimate_json", "contrast_standard_error_json",
        "contrast_statistic_json", "contrast_numerator_df_json",
        "contrast_denominator_df_json", "contrast_p_value_json",
    ]
    try:
        with opener(temporary, "wt", encoding="utf-8", newline="") as handle:
            writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n")
            writer.writeheader()
            writer.writerows(rows)
        os.replace(temporary, output)
    finally:
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass
    counts: dict[str, int] = {}
    for row in rows:
        counts[row["status"]] = counts.get(row["status"], 0) + 1
    atomic_json(output.with_suffix(output.suffix + ".json"), {
        "format": FINALIZE_FORMAT,
        "finished_utc": utc_now(),
        "manifest_sha256": canonical_sha256(manifest),
        "run_signature_sha256": signature["signature_sha256"],
        "tasks": len(rows),
        "status_counts": counts,
    })
    log("finalized {} tasks -> {}".format(len(rows), output))
    return output


def status(manifest: Mapping[str, Any]) -> None:
    output = manifest_paths(manifest)["output_directory"]
    assert output is not None
    counts: dict[str, int] = {}
    completed = 0
    for marker in output.glob("block_*.complete"):
        summary = output / (marker.name[:-len(".complete")] + ".summary.tsv")
        with summary.open("r", encoding="utf-8-sig", newline="") as handle:
            for row in csv.DictReader(handle, delimiter="\t"):
                completed += 1
                counts[row.get("status", "unknown")] = counts.get(row.get("status", "unknown"), 0) + 1
    print(json.dumps({
        "tasks": manifest["dimensions"]["task_count"],
        "completed": completed,
        "status_counts": counts,
        "output_directory": str(output),
    }, ensure_ascii=False, indent=2))


def positive_integer(value: str) -> int:
    parsed = int(value)
    if parsed <= 0:
        raise argparse.ArgumentTypeError("must be positive")
    return parsed


def add_run_options(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--engine", default="spectra_reml")
    parser.add_argument("--threads", type=positive_integer, default=1)
    parser.add_argument("--blas-threads", type=positive_integer, default=1)
    parser.add_argument(
        "--fixed-effect-test",
        choices=("satterthwaite", "kenward-roger", "none"),
        default="satterthwaite",
    )
    parser.add_argument("--block-size", type=positive_integer, default=256)
    parser.add_argument("--max-iterations", type=positive_integer, default=100)
    parser.add_argument("--line-search-max-evals", type=positive_integer, default=48)
    parser.add_argument("--line-search-max-zoom", type=positive_integer, default=48)
    parser.add_argument("--gradient-abs-tol", type=float, default=1e-7)
    parser.add_argument("--gradient-rel-tol", type=float, default=1e-8)
    parser.add_argument("--step-rel-tol", type=float, default=1e-9)
    parser.add_argument("--likelihood-rel-tol", type=float, default=1e-11)
    parser.add_argument("--wolfe-c1", type=float, default=1e-4)
    parser.add_argument("--wolfe-c2", type=float, default=0.9)
    parser.add_argument("--initial-step", type=float, default=1.0)
    parser.add_argument("--maximum-step", type=float, default=64.0)
    parser.add_argument("--line-search-expansion", type=float, default=1.618)
    parser.add_argument("--zoom-safeguard", type=float, default=0.02)
    parser.add_argument("--rank-tol", type=float, default=1e-10)
    parser.add_argument("--covariance-floor", type=float, default=1e-12)
    parser.add_argument("--boundary-h2-trigger", type=float, default=1e-3)
    parser.add_argument("--boundary-score-tol", type=float, default=1e-8)
    parser.add_argument("--boundary-logl-rel-tol", type=float, default=1e-12)
    parser.add_argument("--initial-sigma-e", type=float)
    parser.add_argument("--initial-sigma-g", type=float)
    mode = parser.add_mutually_exclusive_group()
    mode.add_argument("--resume", action="store_true")
    mode.add_argument("--force", action="store_true")
    parser.add_argument("--dry-run", action="store_true")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="action", required=True)
    make = subparsers.add_parser("make-manifest")
    make.add_argument("--manifest", type=Path, required=True)
    make.add_argument("--grm-bin", type=Path, required=True)
    make.add_argument("--grm-id", type=Path)
    make.add_argument("--base-x", type=Path, required=True)
    make.add_argument("--phenotypes", type=Path, required=True)
    make.add_argument("--extra-covariates", type=Path)
    make.add_argument("--tasks", type=Path, required=True)
    make.add_argument("--extra-offsets", type=Path, required=True)
    make.add_argument("--extra-indices", type=Path, required=True)
    make.add_argument("--contrast-matrix", type=Path)
    make.add_argument("--contrast-metadata", type=Path)
    make.add_argument("--output-dir", type=Path, required=True)
    make.add_argument("--n-samples", type=positive_integer, required=True)
    make.add_argument("--n-base-covariates", type=positive_integer, required=True)
    make.add_argument("--n-phenotype-rows", type=positive_integer, required=True)
    make.add_argument("--n-extra-covariate-rows", type=int, default=0)
    make.add_argument("--n-contrasts", type=int, default=0)
    for action in ("validate", "status"):
        command = subparsers.add_parser(action)
        command.add_argument("--manifest", type=Path, required=True)
    finalize_parser = subparsers.add_parser("finalize")
    finalize_parser.add_argument("--manifest", type=Path, required=True)
    finalize_parser.add_argument("--output", type=Path, required=True)
    for action in ("run", "all"):
        command = subparsers.add_parser(action)
        command.add_argument("--manifest", type=Path, required=True)
        if action == "all":
            command.add_argument("--output", type=Path, required=True)
        add_run_options(command)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    if args.action == "make-manifest":
        make_manifest(args)
        return 0
    manifest = validate_manifest(args.manifest)
    if args.action == "validate":
        print(json.dumps({
            "manifest": str(resolved(args.manifest)),
            "manifest_sha256": canonical_sha256(manifest),
            "dimensions": manifest["dimensions"],
        }, ensure_ascii=False, indent=2))
    elif args.action == "status":
        status(manifest)
    elif args.action == "finalize":
        finalize(manifest, args.output)
    elif args.action == "run":
        run_engine(manifest, args)
    elif args.action == "all":
        run_engine(manifest, args)
        if not args.dry_run:
            finalize(manifest, args.output)
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        raise SystemExit(130)
    except Exception as exc:
        print("ERROR: {}".format(exc), file=sys.stderr)
        raise SystemExit(1)