Source code for coseda.ici.no_run_prep_singlelist

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
no_run_prep_singlelist.py
Step 1 with "defaults-if-no-args":
- If launched with NO arguments, uses hardcoded defaults for root/geom/cell/h5 and a standard flag set.
- If launched WITH arguments, behaves as a normal CLI and appends any flags after `--`.
- Always writes:
    <run-root>/runs_ts/run_000/lst_000.lst
    <run-root>/runs_ts/run_000/sh_000.sh
- Does NOT execute indexamajig; only writes the exact command to sh_000.sh.

Default dataset (used only when no CLI args are provided):
    default_root = "/Users/xiaodong/Desktop/simulations/MFM300-VIII_tI/sim_004"
    default_geom = default_root + "/MFM300-VIII.geom"
    default_cell = default_root + "/MFM300-VIII.cell"
    default_h5   = default_root + "/sim.h5"
"""

from __future__ import annotations
from coseda.logging_utils import log_print
import argparse, os, sys, shlex
from pathlib import Path
from typing import List
import h5py

IMAGES_DS = "/entry/data/images"

def _abs(p: str) -> str:
    return os.path.abspath(os.path.expanduser(p))

def _gather_h5(paths: List[str]) -> List[str]:
    out = set()
    for p in paths:
        ap = _abs(p)
        if os.path.isfile(ap) and ap.lower().endswith(".h5"):
            out.add(ap)
        elif os.path.isdir(ap):
            for root, _, files in os.walk(ap):
                for fn in files:
                    if fn.lower().endswith(".h5"):
                        out.add(os.path.join(root, fn))
    return sorted(out)

def _count_images(h5_path: str) -> int:
    with h5py.File(h5_path, "r") as f:
        if IMAGES_DS not in f:
            raise KeyError(f"{IMAGES_DS} not found in {h5_path}")
        return int(f[IMAGES_DS].shape[0])
    
# --- add this new helper (replace _write_combined_lst) ---
def _write_per_image_lists(run_dir: str, h5_files: List[str]) -> int:
    """
    Create run_000/event_XXXXXX/ folders and write a one-line lst in each:
      <run_dir>/event_000000/lst_000000.lst  ->  "<h5> //0"
      <run_dir>/event_000001/lst_000001.lst  ->  "<h5> //1"
      ...
    Returns total number of images across all HDF5s.
    """
    Path(run_dir).mkdir(parents=True, exist_ok=True)
    total = 0
    ev = 0
    for h5 in h5_files:
        n = _count_images(h5)
        for i in range(n):
            ev_str = f"{ev:06d}"
            ev_dir = os.path.join(run_dir, f"event_{ev_str}")
            Path(ev_dir).mkdir(parents=True, exist_ok=True)
            lst_path = os.path.join(ev_dir, f"lst_{ev_str}.lst")
            with open(lst_path, "w", encoding="utf-8") as f:
                f.write(f"{h5} //{i}\n")
            ev += 1
        total += n
    return total




def build_argparser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(
        description="Create a single combined .lst and a sh_000.sh with the exact indexamajig command (no execution).",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    ap.add_argument("--run-root", help="Root folder under which 'run_000' will be created.")
    ap.add_argument("--geom", help="Path to .geom file.")
    ap.add_argument("--cell", help="Path to .cell file.")
    ap.add_argument("sources", nargs="*", help="One or more .h5 files or directories to search recursively for .h5.")
    return ap





def main(argv: List[str]) -> int:
    ours = argv
    ap = build_argparser()
    args = ap.parse_args(ours)

    # using_defaults = (len(ours) == 0)

    # Validate provided arguments
    if not args.run_root or not args.geom or not args.cell or not args.sources:
        log_print("ERROR: Missing required arguments. Provide --run-root, --geom, --cell and at least one source, or run with NO args to use defaults.", file=sys.stderr)
        return 2
    run_root = _abs(args.run_root)
    geom = _abs(args.geom)
    cell = _abs(args.cell)
    h5s = _gather_h5(args.sources)

    if not h5s:
        log_print("No .h5 sources found.", file=sys.stderr)
        return 2

    run_dir = os.path.join(run_root, "run_000")
    os.makedirs(run_dir, exist_ok=True)

    lst_path = os.path.join(run_dir, "lst_000.lst")

    # print("=== Step 1 (single-list): plan indexamajig run (no execution) ===")
    # print(f"Mode:        {'DEFAULTS' if using_defaults else 'CLI'}")
    # print(f"Run dir:     {run_dir}")
    # print(f"Geom:        {geom}")
    # print(f"Cell:        {cell}")
    # print(f"HDF5 files:  {len(h5s)}")
    # print("===============================================================")

    total_images = _write_per_image_lists(run_dir, h5s)
    # print(f"Wrote {total_images} one-line .lst files under {run_dir}/event_*/")
    # print("Per-image shell scripts will be generated by create_run_sh.py.")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))