cli.py

"""URA-SIM command-line interface for reduced-frame reconstruction.

Supports 2D-SIM, NL-SIM, and 3D-SIM. This is the package-owned entry point
and does not wrap ``tools.r5_recon.cli.run_single_stack``.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import tifffile

from . import (
    NLSimCompletionConfig,
    Sim3DCompletionConfig,
    cli_completion_config,
    complete,
    complete_nlsim,
    complete_sim3d,
    load_nlsim_stack,
    load_raw_cycle,
    load_sim3d_stack,
)
from .backends import (
    DEFAULT_CUDA_LIB_DIR,
    DEFAULT_PHASEGATE_BIN,
    Open3DSIMRunConfig,
    PhaseGateReconConfig,
    default_recon_config,
    estimate_illumination,
    phasegate_summary,
    reconstruct_nlsim,
    run_open3dsim,
    run_phasegate,
)
from .io.sim3d import write_sequential_stack
from .modalities.nlsim.config import BackendConfig, OpticalConfig, default_backend_config
from .modalities.sim2d.cli_parity import CLI_COMPLETION_MODES

PACKAGE_ROOT = Path(__file__).resolve().parent
DEFAULT_SOURCE_TIF = (
    PACKAGE_ROOT
    / "data"
    / "3-63倍1.47-cos7-pkmito-50mWb-50ms-3.35stripe-delay5s-49-first9-c0000.tif"
)
DEFAULT_OUT_DIR = PACKAGE_ROOT / "output"
PHASEGATE_ESTIMATORS = ("baseline", "apg", "cpg")
PHASEGATE_RECON_MODES = ("baseline", "hifi")


def parse_selection_text(text: str) -> list[int]:
    values = (
        [int(part.strip()) for part in text.split(",") if part.strip()]
        if "," in text
        else [int(ch) for ch in text.strip()]
    )
    if not values:
        raise ValueError("Expected a non-empty --selection.")
    if len(set(values)) != len(values):
        raise ValueError(f"--selection contains duplicates: {text}")
    return sorted(values)


def selection_token(selection: list[int]) -> str:
    return "".join(str(index) for index in selection)


def build_parser() -> argparse.ArgumentParser:
    defaults = default_recon_config()
    parser = argparse.ArgumentParser(
        prog="python -m ura_sim",
        description="Complete a reduced SIM stack and optionally reconstruct it with the bundled backend.",
    )
    parser.add_argument(
        "--modality",
        choices=("sim2d", "nlsim", "sim3d"),
        default="sim2d",
        help="Acquisition modality. Default: sim2d.",
    )
    parser.add_argument(
        "--source-tif",
        "--raw-tif",
        dest="source_tif",
        type=Path,
        default=DEFAULT_SOURCE_TIF,
        help="Input TIFF. 2D: 9-frame cycle. NL-SIM: 25-frame stack. 3D: sequential Open-3DSIM TIFF.",
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=DEFAULT_OUT_DIR,
        help="Output directory. Default: ura_sim/output",
    )
    parser.add_argument("--cycle-index", type=int, default=0)
    parser.add_argument(
        "--selection",
        type=str,
        required=True,
        help="Observed frame indices. 2D/NL-SIM are 0-based (01258 or 0,1,2,5,8). 3D-SIM is 1-based raw indices (1,2,3,6,7,8,11,12,13).",
    )
    parser.add_argument(
        "--completion-mode",
        choices=CLI_COMPLETION_MODES,
        default="baseline",
        help="Preset matching the legacy r5_recon CLI completion configs.",
    )
    parser.add_argument(
        "--completion-model",
        choices=("two_phase_prior", "harmonic", "nearest", "joint_five", "fixed_dc", "phase_smooth", "five_component"),
        default="two_phase_prior",
        help="2D: two_phase_prior/harmonic/nearest. NL-SIM: joint_five/fixed_dc. 3D: phase_smooth/five_component.",
    )
    parser.add_argument("--dc-scale", type=float)
    parser.add_argument("--singleton-gain", type=float)
    parser.add_argument(
        "--preserve-observed",
        action="store_true",
        help="Copy selected input frames back onto the completed stack.",
    )
    parser.add_argument(
        "--complete-only",
        action="store_true",
        help="Write the completed raw stack and skip reconstruction.",
    )
    parser.add_argument("--matlab-bin", type=Path, help="MATLAB executable for NL-SIM and 3D-SIM reconstruction.")
    parser.add_argument("--illumination-mat", type=Path, help="Existing NL-SIM illumination .mat; skip estimation if provided.")
    parser.add_argument("--phasegate-bin", type=Path, default=DEFAULT_PHASEGATE_BIN)
    parser.add_argument("--cuda-lib-dir", type=Path, default=DEFAULT_CUDA_LIB_DIR)
    parser.add_argument("--estimator", choices=PHASEGATE_ESTIMATORS, default=defaults.estimator)
    parser.add_argument("--recon-mode", choices=PHASEGATE_RECON_MODES, default=defaults.recon_mode)
    parser.add_argument("--lambda", dest="wavelength_nm", type=float, default=defaults.wavelength_nm)
    parser.add_argument("--na", type=float, default=defaults.na)
    parser.add_argument("--pixel-size", dest="pixel_size_um", type=float, default=defaults.pixel_size_um)
    parser.add_argument("--mag", type=float, default=defaults.mag)
    parser.add_argument("--contrast", type=float, default=defaults.contrast)
    parser.add_argument("--wotf-amp", type=float, default=defaults.wotf_amp)
    parser.add_argument("--wotf-sigma-factor", type=float, default=defaults.wotf_sigma_factor)
    parser.add_argument("--wiener-w1", type=float, default=defaults.wiener_w1)
    parser.add_argument("--wiener-w2", type=float, default=defaults.wiener_w2)
    parser.add_argument("--k-mag-ref", type=float, default=defaults.k_mag_ref)
    parser.add_argument("--wiener-secondary", type=float, default=0.1, help="NL-SIM JSFR secondary Wiener parameter.")
    parser.add_argument("--final-apodization-rad", type=float, default=0.0, help="NL-SIM JSFR final apodization radius.")
    return parser


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    return build_parser().parse_args(argv)


def recon_config_from_args(args: argparse.Namespace) -> PhaseGateReconConfig:
    return PhaseGateReconConfig(
        bin_path=args.phasegate_bin,
        estimator=args.estimator,
        recon_mode=args.recon_mode,
        wavelength_nm=args.wavelength_nm,
        na=args.na,
        pixel_size_um=args.pixel_size_um,
        mag=args.mag,
        contrast=args.contrast,
        wotf_amp=args.wotf_amp,
        wotf_sigma_factor=args.wotf_sigma_factor,
        wiener_w1=args.wiener_w1,
        wiener_w2=args.wiener_w2,
        k_mag_ref=args.k_mag_ref,
        cuda_lib_dir=args.cuda_lib_dir,
    )


def completion_config_from_args(args: argparse.Namespace):
    config = cli_completion_config(args.completion_mode)
    return type(config)(
        model=args.completion_model,
        dc_scale=config.dc_scale if args.dc_scale is None else float(args.dc_scale),
        singleton_gain=config.singleton_gain if args.singleton_gain is None else float(args.singleton_gain),
        prior_config=config.prior_config,
        preserve_observed=bool(args.preserve_observed),
    )


def _case_dir(args: argparse.Namespace, source_tif: Path, selection: list[int], mode_label: str) -> Path:
    token = selection_token(selection)
    label = f"{args.modality}_count{len(selection)}_sel{token}_{mode_label}"
    out_dir = args.out_dir.expanduser() / source_tif.stem / label
    out_dir.mkdir(parents=True, exist_ok=True)
    return out_dir


def run_sim2d(args: argparse.Namespace, source_tif: Path, selection: list[int]) -> dict[str, object]:
    completion_config = completion_config_from_args(args)
    raw, cycle_record = load_raw_cycle(source_tif, args.cycle_index)
    completed, manifest = complete(raw, selection, config=completion_config)
    out_dir = _case_dir(args, source_tif, selection, args.completion_mode)
    completed_tif = out_dir / "completed_stack.tif"
    tifffile.imwrite(str(completed_tif), completed)
    payload: dict[str, object] = {
        "modality": "sim2d",
        "source_tif": str(source_tif.resolve()),
        "out_dir": str(out_dir.resolve()),
        "label": out_dir.name,
        "cycle": cycle_record,
        "selection": selection,
        "completion_mode": args.completion_mode,
        "completion_model": args.completion_model,
        "completed_tif": str(completed_tif.resolve()),
        "completed_shape": [int(value) for value in completed.shape],
        "completion_manifest": manifest,
        "complete_only": bool(args.complete_only),
    }
    if not args.complete_only:
        recon_config = recon_config_from_args(args)
        sr_tif = out_dir / "sr_f32.tif"
        recon = run_phasegate(completed, sr_tif, out_dir, recon_config)
        payload["sr_tif"] = str(sr_tif.resolve())
        payload["sr_shape"] = recon["sr_shape"]
        payload["reconstruction"] = recon
        payload["reconstruction_config"] = phasegate_summary(recon_config)
    return payload


def run_nlsim(args: argparse.Namespace, source_tif: Path, selection: list[int]) -> dict[str, object]:
    raw, cycle_record = load_nlsim_stack(source_tif)
    model = args.completion_model if args.completion_model in {"joint_five", "fixed_dc"} else "joint_five"
    completed, manifest = complete_nlsim(raw, selection, config=NLSimCompletionConfig(model=model))
    out_dir = _case_dir(args, source_tif, selection, model)
    completed_tif = out_dir / "completed_stack.tif"
    tifffile.imwrite(str(completed_tif), completed)
    payload: dict[str, object] = {
        "modality": "nlsim",
        "source_tif": str(source_tif.resolve()),
        "out_dir": str(out_dir.resolve()),
        "label": out_dir.name,
        "cycle": cycle_record,
        "selection": selection,
        "completion_model": model,
        "completed_tif": str(completed_tif.resolve()),
        "completed_shape": [int(value) for value in completed.shape],
        "completion_manifest": manifest,
        "complete_only": bool(args.complete_only),
    }
    if not args.complete_only:
        backend = default_backend_config()
        backend = BackendConfig(
            matlab_bin=args.matlab_bin or backend.matlab_bin,
            jsfr_root=backend.jsfr_root,
            optics=OpticalConfig(
                na=args.na if args.na is not None else backend.optics.na,
                wavelength_um=args.wavelength_nm / 1000.0,
                magnification=args.mag,
                pixel_size_um=args.pixel_size_um,
            ),
            wiener_secondary=args.wiener_secondary,
            final_apodization_rad=args.final_apodization_rad,
        )
        illumination_mat = args.illumination_mat or (out_dir / "illumination.mat")
        if args.illumination_mat is None:
            estimate_illumination(completed_tif, illumination_mat, out_dir, backend)
        sr_tif = out_dir / "sr_jsfr.tif"
        recon = reconstruct_nlsim(completed_tif, sr_tif, illumination_mat, out_dir, backend)
        payload["sr_tif"] = str(sr_tif.resolve())
        payload["illumination_mat"] = str(Path(illumination_mat).resolve())
        payload["reconstruction"] = recon
    return payload


def run_sim3d(args: argparse.Namespace, source_tif: Path, selection: list[int]) -> dict[str, object]:
    raw, volume_record = load_sim3d_stack(source_tif)
    model = args.completion_model if args.completion_model in {"phase_smooth", "five_component"} else "phase_smooth"
    completed, manifest = complete_sim3d(raw, selection, config=Sim3DCompletionConfig(model=model))
    out_dir = _case_dir(args, source_tif, selection, model)
    completed_tif = out_dir / "completed_stack.tif"
    write_sequential_stack(completed_tif, completed)
    payload: dict[str, object] = {
        "modality": "sim3d",
        "source_tif": str(source_tif.resolve()),
        "out_dir": str(out_dir.resolve()),
        "label": out_dir.name,
        "volume": volume_record,
        "selection": selection,
        "completion_model": model,
        "completed_tif": str(completed_tif.resolve()),
        "completed_shape": [int(value) for value in completed.shape],
        "completion_manifest": manifest,
        "complete_only": bool(args.complete_only),
    }
    if not args.complete_only:
        sr_name = "sr_open3dsim.tif"
        recon = run_open3dsim(
            Open3DSIMRunConfig(
                dataset=completed_tif,
                output_dir=out_dir,
                output_name=sr_name,
                observed_raw_idx=tuple(range(1, 16)),
                matlab_bin=args.matlab_bin,
            )
        )
        payload["sr_tif"] = recon["sr_tif"]
        payload["reconstruction"] = recon
    return payload


def run(args: argparse.Namespace) -> dict[str, object]:
    source_tif = args.source_tif.expanduser()
    if not source_tif.is_file():
        raise FileNotFoundError(f"source TIFF not found: {source_tif}")
    selection = parse_selection_text(args.selection)
    if args.modality == "sim2d":
        payload = run_sim2d(args, source_tif, selection)
    elif args.modality == "nlsim":
        payload = run_nlsim(args, source_tif, selection)
    else:
        payload = run_sim3d(args, source_tif, selection)
    out_dir = Path(str(payload["out_dir"]))
    (out_dir / "manifest.json").write_text(
        json.dumps(payload, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    return payload


def main(argv: list[str] | None = None) -> int:
    payload = run(parse_args(argv))
    print(json.dumps(payload, indent=2, ensure_ascii=False))
    return 0


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