#!/usr/bin/env python3 """Small genome data preparation substeps. This module exposes reusable pieces of run_genome_prepare.py as independent stages so agents can call them directly from natural language. """ from __future__ import annotations import argparse import json import shutil import subprocess import sys from datetime import datetime from pathlib import Path from typing import Any PROJECT_ROOT = Path(__file__).resolve().parents[2] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from cli.run_genome_prepare import ( MultiSpeciesDetector, MultiSpeciesGenomePreparer, compute_file_hash, generate_species_config, get_tool_versions, merge_species_files, organize_files_by_config, prepare_directories, ) STAGE_TO_OPERATION = { "prepare_organize": "organize", "prepare_normalize_gff": "normalize_gff", "prepare_longest_transcript": "longest_transcript", "prepare_extract_pep": "extract_pep", "prepare_extract_cds": "extract_cds", "prepare_filter_fasta": "filter_fasta", "prepare_simplify_ncbi": "simplify_ncbi", } OPERATION_TO_STAGE = {value: key for key, value in STAGE_TO_OPERATION.items()} def _artifact(path: Path, name: str, artifact_type: str = "file") -> dict[str, Any]: return { "artifact": name, "path": str(path.name), "full_path": str(path.absolute()), "type": artifact_type, "size": path.stat().st_size if path.exists() else 0, "sha256": compute_file_hash(path) if path.exists() else "", "mtime": path.stat().st_mtime if path.exists() else None, } def _write_manifest( workspace: Path, input_dir: Path, prefix: str, stage: str, operation: str, start_time: datetime, artifacts: dict[str, dict[str, Any]], metrics: dict[str, Any], parameters: dict[str, Any], status: str = "success", error: str = "", ) -> Path: final_dir = workspace / "final" final_dir.mkdir(parents=True, exist_ok=True) end_time = datetime.now() manifest = { "manifest_version": "1.0", "stage": stage, "status": status, "prefix": prefix, "operation": operation, "timestamp": end_time.isoformat(), "start_time": start_time.isoformat(), "end_time": end_time.isoformat(), "elapsed_sec": round((end_time - start_time).total_seconds(), 2), "workspace_root": str(workspace.resolve()), "workspace": str(final_dir.resolve()), "input_dir": str(input_dir.resolve()), "parameters": parameters, "metrics": metrics, "artifacts": artifacts, "tool_versions": get_tool_versions(), } if error: manifest["error"] = error manifest_path = final_dir / f"{prefix}.{stage}.manifest.json" manifest["artifacts"]["manifest"] = { "artifact": "manifest", "path": f"final/{manifest_path.name}", "full_path": str(manifest_path.absolute()), "type": "json", "size": 0, } manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8") input_final = input_dir / "final" input_final.mkdir(parents=True, exist_ok=True) try: shutil.copy2(manifest_path, input_final / manifest_path.name) except Exception: pass return manifest_path def _preparer(input_dir: Path, workspace: Path) -> MultiSpeciesGenomePreparer: paths = prepare_directories(input_dir) paths["work_dir"] = input_dir / "genome_prepare" paths["work_dir"].mkdir(exist_ok=True) return MultiSpeciesGenomePreparer(input_dir, workspace, Path("scripts"), paths) def _iter_species(input_dir: Path): return MultiSpeciesDetector.scan_input_dir(input_dir).values() def _run_organize(input_dir: Path, workspace: Path) -> tuple[dict, dict, dict]: org = organize_files_by_config(input_dir) artifacts: dict[str, dict[str, Any]] = {} metrics = { "organized": bool(org.get("organized")), "moved_files": len(org.get("moved_files", [])), "missing_files": len(org.get("missing_files", [])), "target_species": org.get("target_species", []), "ref_species": org.get("ref_species", []), } genome_prepare_dir = input_dir / "genome_prepare" genome_prepare_dir.mkdir(exist_ok=True) if org.get("organized"): merge = merge_species_files(input_dir, genome_prepare_dir, org) cfg = generate_species_config(input_dir, genome_prepare_dir, org) artifacts["species_config"] = _artifact(cfg, "species_config", "json") for key, value in merge.get("merged_files", {}).items(): path = Path(value) if path.exists(): artifacts[key] = _artifact(path, key, "fasta") metrics["merged_files"] = list(merge.get("merged_files", {})) return artifacts, metrics, org def _run_normalize_gff(input_dir: Path, workspace: Path) -> tuple[dict, dict]: prep = _preparer(input_dir, workspace) artifacts: dict[str, dict[str, Any]] = {} count = 0 for species in _iter_species(input_dir): if not species.gff: continue gff = prep.decompress_if_needed(species.gff) db_type = prep.detect_db_type_from_gff(gff) out = input_dir / f"{species.species_code}.normalized.gff3" prep._preprocess_gff(gff, out, db_type) artifacts[f"{species.species_code}_normalized_gff"] = _artifact(out, f"{species.species_code}_normalized_gff", "gff3") count += 1 return artifacts, {"normalized_gff": count} def _run_longest_transcript(input_dir: Path, workspace: Path) -> tuple[dict, dict]: prep = _preparer(input_dir, workspace) longest_script = prep._ensure_script("gff_filter_longest.py") artifacts: dict[str, dict[str, Any]] = {} count = 0 for species in _iter_species(input_dir): if not species.gff: continue gff = prep.decompress_if_needed(species.gff) db_type = prep.detect_db_type_from_gff(gff) gff_tmp = prep.work_dir / f"{species.species_code}.gff3.tmp" prep._preprocess_gff(gff, gff_tmp, db_type) table = prep.work_dir / f"{species.species_code}_gene_mrna_cds.ids" out_gff = input_dir / f"{species.species_code}.longest.gff3" rc, _, stderr = prep.run_command( [sys.executable, str(longest_script), str(gff_tmp.resolve()), str(table.resolve()), str(out_gff.resolve())], f"Select longest transcript for {species.species_code}", ) if rc != 0: raise RuntimeError(f"Longest transcript failed for {species.species_code}: {stderr}") ids = input_dir / f"{species.species_code}.longest.ids" prep._write_ids_from_longest_table(table, ids) artifacts[f"{species.species_code}_longest_gff"] = _artifact(out_gff, f"{species.species_code}_longest_gff", "gff3") artifacts[f"{species.species_code}_longest_ids"] = _artifact(ids, f"{species.species_code}_longest_ids", "txt") count += 1 return artifacts, {"species_with_longest_transcript": count} def _run_gffread(input_dir: Path, workspace: Path, mode: str) -> tuple[dict, dict]: prep = _preparer(input_dir, workspace) prep._ensure_tool("gffread") artifacts: dict[str, dict[str, Any]] = {} count = 0 flag = "-y" if mode == "pep" else "-x" suffix = "pep.fasta" if mode == "pep" else "cds.fasta" for species in _iter_species(input_dir): if not species.gff or not species.genome: continue gff = prep.decompress_if_needed(species.gff) genome = prep.decompress_if_needed(species.genome) out = input_dir / f"{species.species_code}.{suffix}" rc, _, stderr = prep.run_command( ["gffread", flag, str(out), "-g", str(genome), str(gff)], f"Extract {mode.upper()} for {species.species_code}", ) if rc != 0 or not out.exists() or out.stat().st_size == 0: raise RuntimeError(f"gffread {mode} extraction failed for {species.species_code}: {stderr}") artifacts[f"{species.species_code}_{mode}"] = _artifact(out, f"{species.species_code}_{mode}", "fasta") count += 1 return artifacts, {f"extracted_{mode}_species": count} def _run_filter_fasta(input_dir: Path, workspace: Path) -> tuple[dict, dict]: prep = _preparer(input_dir, workspace) prep._ensure_tool("seqtk") artifacts: dict[str, dict[str, Any]] = {} count = 0 for ids in sorted(input_dir.glob("*.ids")) + sorted(input_dir.glob("*_mRNA.id")): species_code = ids.name.split(".", 1)[0].split("_", 1)[0] candidates = sorted(input_dir.glob(f"{species_code}*.fasta")) + sorted(input_dir.glob(f"{species_code}*.fa")) for fasta in candidates: if ids.resolve() == fasta.resolve() or ".subset." in fasta.name: continue out = input_dir / f"{fasta.stem}.subset.fasta" rc, _, stderr = prep.run_command_to_file(["seqtk", "subseq", str(fasta), str(ids)], out, f"Filter {fasta.name} by {ids.name}") if rc != 0: raise RuntimeError(f"seqtk subseq failed for {fasta.name}: {stderr}") artifacts[f"{out.stem}"] = _artifact(out, out.stem, "fasta") count += 1 return artifacts, {"filtered_fasta_files": count} def _run_simplify_ncbi(input_dir: Path, workspace: Path) -> tuple[dict, dict]: prep = _preparer(input_dir, workspace) artifacts: dict[str, dict[str, Any]] = {} count = 0 for species in _iter_species(input_dir): if species.gff: gff = prep.decompress_if_needed(species.gff) out_gff = input_dir / f"{species.species_code}.ncbi_simplified.gff3" shutil.copy2(gff, out_gff) prep._simplify_ncbi_gff_attrs(out_gff) artifacts[f"{species.species_code}_ncbi_simplified_gff"] = _artifact(out_gff, f"{species.species_code}_ncbi_simplified_gff", "gff3") count += 1 if species.cds: cds = prep.decompress_if_needed(species.cds) out_cds = input_dir / f"{species.species_code}.ncbi_normalized.cds.fasta" prep._normalize_ncbi_cds_headers(cds, out_cds) artifacts[f"{species.species_code}_ncbi_normalized_cds"] = _artifact(out_cds, f"{species.species_code}_ncbi_normalized_cds", "fasta") return artifacts, {"simplified_ncbi_species": count} def run_operation(operation: str, input_dir: Path, workspace: Path) -> tuple[dict, dict, dict]: if operation == "organize": artifacts, metrics, params = _run_organize(input_dir, workspace) elif operation == "normalize_gff": artifacts, metrics = _run_normalize_gff(input_dir, workspace) params = {} elif operation == "longest_transcript": artifacts, metrics = _run_longest_transcript(input_dir, workspace) params = {} elif operation == "extract_pep": artifacts, metrics = _run_gffread(input_dir, workspace, "pep") params = {} elif operation == "extract_cds": artifacts, metrics = _run_gffread(input_dir, workspace, "cds") params = {} elif operation == "filter_fasta": artifacts, metrics = _run_filter_fasta(input_dir, workspace) params = {} elif operation == "simplify_ncbi": artifacts, metrics = _run_simplify_ncbi(input_dir, workspace) params = {} else: raise ValueError(f"Unknown operation: {operation}") return artifacts, metrics, params def main() -> int: parser = argparse.ArgumentParser(description="Run a single genome data preparation substep") parser.add_argument("--input-dir", required=True, type=Path) parser.add_argument("--workspace", required=True, type=Path) parser.add_argument("--prefix", default="data_prep") parser.add_argument("--stage", choices=sorted(STAGE_TO_OPERATION), default="") parser.add_argument("--operation", choices=sorted(OPERATION_TO_STAGE), default="") parser.add_argument("--cpu", default=4) args = parser.parse_args() operation = args.operation or STAGE_TO_OPERATION.get(args.stage) if not operation: parser.error("--operation or --stage is required") stage = args.stage or OPERATION_TO_STAGE[operation] start = datetime.now() try: args.input_dir.mkdir(parents=True, exist_ok=True) args.workspace.mkdir(parents=True, exist_ok=True) artifacts, metrics, params = run_operation(operation, args.input_dir, args.workspace) params.update({"operation": operation, "cpu": args.cpu}) manifest = _write_manifest(args.workspace, args.input_dir, args.prefix, stage, operation, start, artifacts, metrics, params) print(f"Completed {stage}; manifest: {manifest}") return 0 except Exception as exc: manifest = _write_manifest( args.workspace, args.input_dir, args.prefix, stage, operation, start, {}, {}, {"operation": operation, "cpu": args.cpu}, status="failed", error=str(exc), ) print(f"Failed {stage}: {exc}", file=sys.stderr) print(f"Manifest: {manifest}", file=sys.stderr) return 1 if __name__ == "__main__": raise SystemExit(main())