revision/scripts/run_single_neural_split.py

"""Run one frozen neural outer split so long training is restart-safe."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd


PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from revision.scripts.reanalysis_neural import run_neural_stack
from revision.scripts.reanalysis_pipeline import (
    _extract_features,
    _write_sha256_manifest,
    validate_config,
)
from revision.scripts.run_revision_reanalysis import configure_utf8_console


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Train GAT, GCN, FPNN, and the fixed stack for one frozen outer split."
    )
    parser.add_argument(
        "--config",
        default=str(PROJECT_ROOT / "revision" / "config" / "reanalysis.json"),
    )
    parser.add_argument(
        "--artifacts-root",
        default=str(PROJECT_ROOT / "revision" / "artifacts" / "reviewer_requested_reanalysis_v4"),
    )
    parser.add_argument("--strategy", required=True)
    parser.add_argument("--seed", required=True, type=int)
    return parser.parse_args()


def main() -> int:
    configure_utf8_console()
    args = parse_args()
    config = validate_config(
        json.loads(Path(args.config).resolve().read_text(encoding="utf-8")),
        PROJECT_ROOT,
    )
    if args.strategy not in config["split_strategies"]:
        raise ValueError(f"Unknown strategy: {args.strategy}")
    if args.seed not in config["outer_seeds"]:
        raise ValueError(f"Seed was not predeclared: {args.seed}")

    artifacts_root = Path(args.artifacts_root).resolve()
    expected_root = Path(config["output_root"]).resolve()
    if artifacts_root != expected_root:
        raise ValueError(
            f"Artifacts root must match the frozen config: {expected_root}"
        )
    if not (artifacts_root / "SHA256SUMS.txt").is_file():
        raise FileNotFoundError("Completed classical artifacts and SHA256SUMS.txt are required first.")

    annotated = pd.read_csv(artifacts_root / "record_identity_manifest.csv")
    descriptor_matrix, fingerprint_matrix = _extract_features(
        annotated,
        config["descriptor_features"],
        config["fingerprint"],
    )
    split_dir = artifacts_root / args.strategy / f"seed_{args.seed}"
    split_data = np.load(split_dir / "split_indices.npz", allow_pickle=True)
    development = np.asarray(split_data["development_indices"], dtype=int)
    test = np.asarray(split_data["test_indices"], dtype=int)
    inner_train = split_data["inner_train_indices"]
    inner_validation = split_data["inner_validation_indices"]
    inner_folds = [
        (np.asarray(train, dtype=int), np.asarray(validation, dtype=int))
        for train, validation in zip(inner_train, inner_validation)
    ]

    output = run_neural_stack(
        config,
        annotated=annotated,
        descriptor_matrix=descriptor_matrix,
        fingerprint_matrix=fingerprint_matrix,
        train_indices=development,
        test_indices=test,
        inner_folds=inner_folds,
        output_dir=split_dir / "neural_stack",
        seed=args.seed,
        smoke_only=False,
    )
    _write_sha256_manifest(artifacts_root)
    print(f"Completed frozen neural split: {output}")
    return 0


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