qprep

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "typer",
#     "InquirerPy",
# ]
# ///

import re
import sys
import typer
from pathlib import Path
import os
from typing import Optional

# --- Cluster Configurations ---
CLUSTER_CONFIG = {
    "jupiter": {
        "scheduler": "pbs",
        "orca_path": "/home/Softwares/orca5.0.4/orca",
        "queues": {
            "small": {"max_procs": 8, "default_mem_per_proc": 1375},
            "fast": {"max_procs": 8, "default_mem_per_proc": 1375},
            "freq": {"max_procs": 8, "default_mem_per_proc": 3700},
            "big": {"max_procs": 16, "default_mem_per_proc": 1650},
        },
        "default_queue": "small",
        "walltime": "720:00:00",
        "submission_command": "qsub",
        "recli_remote": "jupiter",
    },
    "colaco": {
        "scheduler": "pbs",
        "orca_path": "/home/Softwares/orca5.0.4/orca",
        "queues": {
            "small": {"max_procs": 8, "default_mem_per_proc": 1375},
            "fast": {"max_procs": 8, "default_mem_per_proc": 1375},
            "freq": {"max_procs": 8, "default_mem_per_proc": 3700},
            "big": {"max_procs": 16, "default_mem_per_proc": 1650},
        },
        "default_queue": "small",
        "walltime": "720:00:00",
        "submission_command": "qsub",
        "recli_remote": "colaco",
    },
    "loboc": {
        "scheduler": "pbs",
        "orca_path": "/home/users/vport/bin/orca_5_0_4/orca",
        "default_walltime": "60:00:00",
        "submission_command": "qsub",
        "modules": ["gcc/11.2.0", "openmpi-gnu/4.1.1"],
    },
    "newton": {
        "scheduler": "slurm",
        "orca_paths": {
            "5": "/opt/orca/orca-5.0.4/orca",
            "6": "/opt/orca/orca-6.1.0/orca",
        },
        "default_orca_version": "5",
        "walltime": "4-00:00:00",
        "partition": "long",
        "submission_command": "sbatch",
        "modules": ["fermi/openmpi-4.1.6"],
    },
    "babel": {
        "default_orca_version": "6",
        "max_procs": 16,
        "submission_command": "source",
        "recli_remote": "babel",
    },
}

# --- Job Script Templates ---

PBS_JUPITER_TEMPLATE = """#!/bin/bash
#PBS -l nodes=1:ppn={nprocs}
#PBS -l walltime={walltime}
#PBS -l mem={memory_gb}gb
#PBS -V
#PBS -q {queue}

shopt -s extglob

JOB_ID=${{PBS_JOBID%.*}}

export RSH_COMMAND="/usr/bin/ssh -x"

# temporary directory
export TDIR=/rascunho/${{USER}}-${{PBS_JOBID}}
mkdir -p ${{TDIR}}

cp ${{PBS_O_WORKDIR}}/{input_filename} ${{TDIR}}
for f in ${{PBS_O_WORKDIR}}/*.gbw ${{PBS_O_WORKDIR}}/*xyz ${{PBS_O_WORKDIR}}/*.hess ${{PBS_O_WORKDIR}}/*.res.* ${{PBS_O_WORKDIR}}/*.mp2nat; do cp $f ${{TDIR}}; done

cd ${{TDIR}}

# Log the job start details
echo "start of job:" `date` > ${{PBS_O_WORKDIR}}/{job_name}.job.${{JOB_ID}}
echo "nboenv:" $NBOEXE > ${{PBS_O_WORKDIR}}/{job_name}.job.${{JOB_ID}}
echo "Job started from ${{PBS_O_HOST}}, running on $(hostname) in ${{PBS_O_WORKDIR}} " >> ${{PBS_O_WORKDIR}}/{job_name}.job.${{JOB_ID}}
echo "temporary directory: ${{TDIR}}" >> ${{PBS_O_WORKDIR}}/{job_name}.job.${{JOB_ID}}

# run
{orca_path} {input_filename} > ${{PBS_O_WORKDIR}}/{input_stem}.out || true

echo "final of job:" `date` >> $PBS_O_WORKDIR/{job_name}.job.$JOB_ID

# copy the archives from temporary directory to submit directory
mv /rascunho/${{USER}}-${{PBS_JOBID}}/!(*tmp*) ${{PBS_O_WORKDIR}}/
rm -rf /rascunho/${{USER}}-${{PBS_JOBID}}
"""

PBS_LOBOC_TEMPLATE = """#!/bin/bash
#PBS -l select=1:ncpus={nprocs}:mpiprocs={nprocs}
#PBS -l walltime={walltime}
#PBS -V

module load {modules}

JOB_ID=${{PBS_JOBID%.*}}

# Log the job start details
echo "start of job:" `date` > ${{PBS_O_WORKDIR}}/{job_name}.job.${{JOB_ID}}
echo "Job started from ${{PBS_O_HOST}}, running on $(hostname) in ${{PBS_O_WORKDIR}} " >> ${{PBS_O_WORKDIR}}/{job_name}.job.${{JOB_ID}}

cd ${{PBS_O_WORKDIR}}

# Run the Orca job
{orca_path} $PBS_O_WORKDIR/{input_filename} > $PBS_O_WORKDIR/{input_stem}.out || true

echo "final of job:" `date` >> ${{PBS_O_WORKDIR}}/{job_name}.job.${{JOB_ID}}
"""

SLURM_NEWTON_TEMPLATE = """#!/bin/bash
#SBATCH --job-name={job_name}
#SBATCH -o slurm.%j.out
#SBATCH -e slurm.%j.err
#SBATCH -N 1
#SBATCH --ntasks={nprocs}
#SBATCH --mem={memory_mb}
#SBATCH -t {walltime}
#SBATCH --no-requeue
#SBATCH -p {partition}

module load {modules}

echo "########################################################"
echo "## ORCA Job Submission Script"
echo "# Job name: {job_name}"
echo "# ORCA version: {orca_version_label}"
echo "# ORCA executable: {orca_exec}"
echo "# Number of cores: {nprocs}"
echo "# Total memory: {memory_mb} MB"
echo "# Job started at: $(date)"
echo "########################################################"

{orca_exec} {input_filename} > {input_stem}.out

echo "########################################################"
echo "# ORCA job finished at: $(date)"
echo "########################################################"
"""

SCRIPT_EXT = {
    "jupiter": ".pbs",
    "colaco": ".pbs",
    "loboc": ".pbs",
    "newton": ".slurm",
    "babel": ".pueue",
}

def _complete_remote(incomplete: str) -> list[str]:
    return [r for r in CLUSTER_CONFIG if r.startswith(incomplete)]

def _complete_queue(incomplete: str) -> list[str]:
    return [q for q in ["small", "fast", "freq", "big"] if q.startswith(incomplete)]

def _complete_strategy(incomplete: str) -> list[str]:
    return [s for s in ["basename", "script", "directory"] if s.startswith(incomplete)]

def _complete_orca_version(incomplete: str) -> list[str]:
    return [v for v in ["5", "6"] if v.startswith(incomplete)]

app = typer.Typer(add_completion=True)


def run_interactive_prompts(
    remote: Optional[str],
    queue: Optional[str],
    orca_version: Optional[str],
) -> tuple[str, Optional[str], Optional[str]]:
    """Fill in unspecified options via InquirerPy prompts."""
    from InquirerPy import inquirer
    from InquirerPy.base.control import Choice

    if not remote:
        remote = inquirer.select(
            message="Remote:",
            choices=list(CLUSTER_CONFIG.keys()),
        ).execute()

    if remote == "jupiter" and not queue:
        jupiter_queues = CLUSTER_CONFIG["jupiter"]["queues"]
        choices = [Choice(value=None, name="auto  (nprocs=16 → big, else → small)")]
        for qname, qcfg in sorted(jupiter_queues.items()):
            choices.append(Choice(value=qname, name=f"{qname}  (max {qcfg['max_procs']} procs)"))
        queue = inquirer.select(
            message="Queue:",
            choices=choices,
        ).execute()

    if remote in ("newton", "babel") and not orca_version:
        if remote == "newton":
            versions = list(CLUSTER_CONFIG["newton"]["orca_paths"].keys())
            default = CLUSTER_CONFIG["newton"]["default_orca_version"]
        else:
            versions = ["5", "6"]
            default = CLUSTER_CONFIG["babel"]["default_orca_version"]
        orca_version = inquirer.select(
            message="ORCA version:",
            choices=versions,
            default=default,
        ).execute()

    return remote, queue, orca_version


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def parse_orca_input(content: str) -> dict:
    """Parses ORCA input file content to extract nprocs and maxcore."""
    nprocs, maxcore = None, None

    pal_match = re.search(r"^\s*!\s*.*PAL(\d+)", content, re.IGNORECASE | re.MULTILINE)
    nprocs_match = re.search(
        r"^\s*%pal\s*\n\s*nprocs\s+(\d+)", content, re.IGNORECASE | re.MULTILINE
    )
    nprocs_block_match = re.search(r"nprocs\s+(\d+)", content, re.IGNORECASE)
    maxcore_match = re.search(r"maxcore\s+(\d+)", content, re.IGNORECASE)

    if pal_match:
        nprocs = int(pal_match.group(1))
    elif nprocs_match:
        nprocs = int(nprocs_match.group(1))
    elif nprocs_block_match:
        nprocs = int(nprocs_block_match.group(1))

    if maxcore_match:
        maxcore = int(maxcore_match.group(1))

    return {"nprocs": nprocs, "maxcore": maxcore}


def rewrite_nprocs_in_content(content: str, new_nprocs: int) -> str:
    """Rewrite the nprocs value in an ORCA input string."""
    if re.search(r"\bPAL\d+\b", content, re.IGNORECASE):
        return re.sub(r"\bPAL\d+\b", f"PAL{new_nprocs}", content, flags=re.IGNORECASE)
    if re.search(r"\bnprocs\s+\d+", content, re.IGNORECASE):
        return re.sub(r"(\bnprocs\s+)\d+", rf"\g<1>{new_nprocs}", content, flags=re.IGNORECASE)
    return content


def extract_orca_tags(content: str) -> list[str]:
    """Derive recli tags from the ORCA input keywords."""
    tags = []

    kw_line = " ".join(
        line.strip().lower()
        for line in content.splitlines()
        if line.strip().startswith("!")
    )

    # Job type — specific types before generic ones
    if "neb-ts" in kw_line:
        tags.append("neb-ts")
    elif "neb" in kw_line:
        tags.append("neb")
    elif "optts" in kw_line:
        tags.append("optts")
    elif "irc" in kw_line:
        tags.append("irc")
    elif "optfreq" in kw_line or ("opt" in kw_line and ("freq" in kw_line or "numfreq" in kw_line)):
        tags.append("optfreq")
    elif "opt" in kw_line:
        if re.search(r"^\s*scan\b", content, re.IGNORECASE | re.MULTILINE):
            tags.append("scan")
        else:
            tags.append("opt")
    elif "numfreq" in kw_line or "freq" in kw_line:
        tags.append("freq")
    elif "nmr" in kw_line:
        tags.append("nmr")
    elif "eda" in kw_line:
        tags.append("eda")
    elif "nbo" in kw_line:
        tags.append("nbo")
    elif "goat" in kw_line:
        tags.append("goat")
    else:
        tags.append("sp")

    # Method — only tag non-DFT methods worth distinguishing
    if "xtb" in kw_line:
        tags.append("xtb")
    elif "dlpno" in kw_line:
        tags.append("dlpno")
    elif any(m in kw_line for m in ["ccsd", "cisd", "fci"]):
        tags.append("wf")
    elif "mp2" in kw_line:
        tags.append("mp2")
    elif "hf" in kw_line and not any(
        x in kw_line for x in ["b3lyp", "pbe", "tpss", "m06", "cam-b3lyp", "wb97"]
    ):
        tags.append("hf")

    return tags


def process_file(
    input_file: Path,
    remote: str,
    config: dict,
    queue: Optional[str],
    orca_version: Optional[str],
    use_dir_name: bool,
    extra_tags: Optional[str],
    strategy: Optional[str],
):
    if use_dir_name:
        current_dir = Path.cwd()
        job_name = f"{current_dir.parent.name}_{current_dir.name}"
    else:
        job_name = input_file.stem

    input_filename = input_file.name
    input_stem = input_file.stem
    content = input_file.read_text()

    parsed = parse_orca_input(content)
    nprocs = parsed.get("nprocs")
    maxcore = parsed.get("maxcore")

    if remote != "babel" and not nprocs:
        eprint(f"Error [{input_file.name}]: could not find nprocs or PAL in input file.")
        return

    job_script_content = ""
    resolved_queue = queue  # track for tag generation

    if remote in ("jupiter", "colaco"):
        if not queue:
            resolved_queue = "big" if nprocs == 16 else "small"
        else:
            resolved_queue = queue

        if resolved_queue not in config["queues"]:
            eprint(
                f"Error [{input_file.name}]: queue '{resolved_queue}' not valid for Jupiter. "
                f"Options: {list(config['queues'].keys())}"
            )
            return

        max_procs = config["queues"][resolved_queue]["max_procs"]

        if nprocs > max_procs:
            eprint(
                f"Error [{input_file.name}]: {nprocs} procs requested but '{resolved_queue}' allows {max_procs}."
            )
            return

        if nprocs < max_procs:
            eprint(
                f"Warning [{input_file.name}]: input has {nprocs} procs but '{resolved_queue}' supports up to {max_procs}."
            )
            if typer.confirm(f"  Upgrade nprocs {nprocs} → {max_procs} in {input_file.name}?"):
                content = rewrite_nprocs_in_content(content, max_procs)
                input_file.write_text(content)
                eprint(f"  Upgraded nprocs: {nprocs} → {max_procs}")
                nprocs = max_procs

        if maxcore:
            memory_gb = (maxcore * nprocs) // 1000
        else:
            memory_gb = (config["queues"][resolved_queue]["default_mem_per_proc"] * nprocs) // 1000

        log = f"queue={resolved_queue}  nprocs={nprocs}  mem={memory_gb}gb"

        job_script_content = PBS_JUPITER_TEMPLATE.format(
            nprocs=nprocs,
            walltime=config["walltime"],
            memory_gb=memory_gb,
            queue=resolved_queue,
            job_name=job_name,
            orca_path=config["orca_path"],
            input_filename=input_filename,
            input_stem=input_stem,
        )

    elif remote == "loboc":
        log = f"nprocs={nprocs}"
        job_script_content = PBS_LOBOC_TEMPLATE.format(
            nprocs=nprocs,
            walltime=config["default_walltime"],
            job_name=job_name,
            orca_path=config["orca_path"],
            modules=" ".join(config["modules"]),
            input_filename=input_filename,
            input_stem=input_stem,
        )

    elif remote == "newton":
        if not maxcore:
            maxcore = 3750
            eprint(
                f"Warning [{input_file.name}]: maxcore not found, using default {maxcore} MB/core"
            )

        version = orca_version or config["default_orca_version"]
        if version not in config["orca_paths"]:
            eprint(
                f"Error [{input_file.name}]: ORCA version '{version}' not supported on Newton. "
                f"Options: {list(config['orca_paths'].keys())}"
            )
            return

        total_mem_mb = nprocs * maxcore
        log = f"nprocs={nprocs}  mem={total_mem_mb}mb  orca={version}"

        job_script_content = SLURM_NEWTON_TEMPLATE.format(
            job_name=job_name,
            nprocs=nprocs,
            memory_mb=total_mem_mb,
            walltime=config["walltime"],
            partition=config["partition"],
            modules=" ".join(config["modules"]),
            orca_version_label=f"ORCA {version}",
            orca_exec=config["orca_paths"][version],
            input_filename=input_filename,
            input_stem=input_stem,
        )

    elif remote == "babel":
        version = orca_version or config["default_orca_version"]
        log = f"orca={version}"
        job_script_content = f"/home/vport/scripts/job -v {version} {input_filename}"

    job_script_path = Path(f"{job_name}{SCRIPT_EXT[remote]}")
    job_script_path.write_text(job_script_content)
    eprint(f"[{input_file.name}]  {log}  →  {job_script_path}")

    recli_remote = config.get("recli_remote")
    if recli_remote:
        auto_tags = extract_orca_tags(content)
        if extra_tags:
            auto_tags += [t.strip() for t in extra_tags.split(",") if t.strip()]
        cmd = ["recli", "submit", str(job_script_path), "-r", recli_remote]
        if strategy:
            cmd += ["--strategy", strategy]
        if auto_tags:
            cmd += ["--tags"] + auto_tags
        print(" ".join(cmd))


@app.command()
def main(
    input_files: list[Path] = typer.Argument(
        ..., help="ORCA input file(s) (.inp)."
    ),
    remote: Optional[str] = typer.Option(
        None, "--remote", "-r",
        help=f"Target remote. Available: {list(CLUSTER_CONFIG.keys())}",
        autocompletion=_complete_remote,
    ),
    queue: Optional[str] = typer.Option(
        None, "--queue", "-q",
        help="Queue override (Jupiter: small/fast/freq/big). Auto-selected if omitted.",
        autocompletion=_complete_queue,
    ),
    orca_version: Optional[str] = typer.Option(
        None, "--orca-version", help="ORCA version to use (Newton/babel).",
        autocompletion=_complete_orca_version,
    ),
    use_dir_name: bool = typer.Option(
        False, "--use-dir-name",
        help="Use 'parent_dir_current_dir' as job name instead of the input file stem.",
    ),
    extra_tags: Optional[str] = typer.Option(
        None, "--tags",
        help="Extra tags (comma-separated) appended to auto-generated ones.",
    ),
    strategy: Optional[str] = typer.Option(
        None, "--strategy",
        help="recli file upload strategy (basename, script, directory). Uses recli's configured default if omitted.",
        autocompletion=_complete_strategy,
    ),
    interactive: bool = typer.Option(
        False, "--interactive", "-i",
        help="Prompt for remote, queue, and ORCA version interactively.",
    ),
):
    """Generate ORCA job scripts for HPC clusters. Prints recli submit commands to stdout."""
    if interactive:
        remote, queue, orca_version = run_interactive_prompts(remote, queue, orca_version)

    if not remote:
        eprint("Error: --remote/-r is required (or use --interactive).")
        raise typer.Exit(code=1)

    if remote not in CLUSTER_CONFIG:
        eprint(
            f"Error: remote '{remote}' not recognized. "
            f"Available: {list(CLUSTER_CONFIG.keys())}"
        )
        raise typer.Exit(code=1)

    config = CLUSTER_CONFIG[remote]

    if use_dir_name and len(input_files) > 1:
        eprint("Warning: --use-dir-name with multiple files assigns the same job name to all.")

    for input_file in input_files:
        if not input_file.exists():
            eprint(f"Error: file not found: {input_file}")
            continue
        process_file(input_file, remote, config, queue, orca_version, use_dir_name, extra_tags, strategy)


if __name__ == "__main__":
    app()