dcaspt2_input

#!/usr/bin/env python3

import argparse
import glob
import os
import shutil
import sys
import tempfile
from datetime import datetime
from pathlib import Path
from subprocess import PIPE, Popen
from typing import Set, Tuple


class PrintVersionAction(argparse.Action):
    """Print the version of the dirac-caspt2 program"""

    def __init__(self, option_strings, dest=argparse.SUPPRESS, default=argparse.SUPPRESS, help="Show the version of dirac-caspt2 program", **kwargs):
        super(PrintVersionAction, self).__init__(option_strings=option_strings, dest=dest, default=default, nargs=0, help=help)

    def __call__(self, parser, namespace, values, option_string=None):
        binary_dir = Path(__file__).resolve().parent
        version = get_version(binary_dir)
        print(f"dirac-caspt2 version (commit hash): {version}")
        parser.exit()


class ExitWithHelpArgumentParser(argparse.ArgumentParser):
    """A custom ArgumentParser class that exits with help message when the unrecognized arguments are specified"""

    def error(self, message):
        print(f"\nError: Invalid arguments : {message}\nPlease check the following help message and try again.\n")
        self.print_help(sys.stderr)
        self.exit(2, f"\nError: Invalid argument : {message}\n")


def parse_args() -> "argparse.Namespace":
    """Parse user input arguments

    Returns:
        argparse.Namespace: User input arguments (argv)
    """
    desc = "\
======================================== dcaspt2 ========================================\n\
dirac-caspt2 wrapper script.\n\
You can can calculate the CASPT2 energy of a molecule using this script.\n\
You need to prepare the DIRAC (http://diracprogram.org) 1-2 integrals file in advance.\n\
==========================================================================================\n"

    parser = ExitWithHelpArgumentParser(description=desc, formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("-i", "--inp", type=str, dest="input", help="Input file [default: active.inp]")
    parser.add_argument("-o", "--out", type=str, dest="output", help="Specify the name of output file [default: dirac_caspt2.out]")
    parser.add_argument("--int", dest="integrals", help="Specify the directory where 1-2 integrals files are stored.\n[default: current directory]", metavar="1-2_INTEGRALS_DIR")
    parser.add_argument("--mpi", dest="mpi", type=int, default=1, help="Specify the number of MPI process. [default: 1]", metavar="NUM_MPI_PROCESS")
    parser.add_argument(
        "--omp",
        default=-1,  # minus number or 0 means that this script doesn't override the OMP_NUM_THREADS environment variable
        dest="threads",
        type=int,
        help="Specify the number of OpenMP threads per (MPI) process.\n\
This option overrides the OMP_NUM_THREADS environment variable.\n\
[default: $OMP_NUM_THREADS]",
        metavar="NUM_OMP_THREADS",
    )
    parser.add_argument(
        "--get",
        dest="get",
        type=str,
        help='Specify the name of the files to be copied from the scratch directory.\n\
Wildcard is available.\n\
(e.g.) --get "MDCINTNEW*" => Copy MDCINTNEW, MDCINTNEW1, MDCINTNEW2, ...\n\
[default: None]\n\
(Note: The file name must be enclosed in double quotes or single quotes to prevent shell expansion.)\n',
        metavar='"FILES"',
    )
    parser.add_argument(
        "--scratch",
        default="~/dcaspt2_scratch",
        dest="scratchdir",
        type=str,
        help="Set the scratch directory. Add the subdirectory to the specified directory.\n\
(e.g.) --scratch /home/username/scratch \n\
=> Create /home/username/scratch/active_2020-01-01_12-34-56 directory.\n\
[default: ~/dcaspt2_scratch]",
    )
    parser.add_argument("--scratchfull", dest="scratchfull", type=str, help="Set the scratch directory. Use the specified directory as the scratch directory.\n[default: None]")
    parser.add_argument("-s", "--save", "--keep_scratch", dest="save", action="store_true", help="Save the scratch directory after calculation. [default: False]")
    parser.add_argument("-v", "--version", action=PrintVersionAction, dest="version", help="Show the version of dirac-caspt2 program")

    # If version is specified, print the version and exit

    return parser.parse_args()


def create_scratch_dir(args: "argparse.Namespace", input_file_path: Path, user_submitted_dir: Path) -> Path:
    """Create a temporary directory and return the path to the directory

    Args:
        args ("argparse.Namespace"): User input arguments (argv)
        input_file_path (Path): Path to the input file
        user_submitted_dir (Path): Path to the directory where this script is used

    Returns:
        scratch_dir (Path): Path to the temporary directory
    """

    # ================
    #  Sub functions
    # ================

    def permission_denied_error_message(scratch_dir: Path) -> None:
        """Print error message when permission denied

        Args:
            scratch_dir (Path): Path to the temporary directory
        """
        print(f"Permission denied: You do not have permission to create the directory '{scratch_dir}'.")
        print("Please specify the directory that you have permission to create.")
        print("Error: Failed to create the temporary directory.\nExiting...")

    def validate_scratchfull_option(scratch_full_dir: Path) -> None:
        """--scratchfull is not allowed to specify ~ or upper paths.
        So we need to validate whether the user specified --scratchfull option is valid
        """

        def create_banned_dirs() -> Set[Path]:
            home_dir = Path("~").expanduser().resolve()
            banned_dirs: Set[Path] = set()
            tmp_dir = home_dir
            while True:
                banned_dirs.add(tmp_dir)
                if tmp_dir == tmp_dir.parent:
                    break
                tmp_dir = tmp_dir.parent
            return banned_dirs

        banned_dirs = create_banned_dirs()
        if scratch_full_dir in banned_dirs:
            # Validation error
            msg = "Invalid --scratchfull. The --scratchfull option can't be set for the $HOME directory and its parent directories.\n\
Because if these directories are removed after calcuation, it can cause serious problems such as the OS not booting.\n\
Please specify the valid path and try again."
            raise ValueError(msg)

    # ================
    #  Main process
    # ================

    # If --scratchfull is specified, Use the specified directory as the scratch directory
    if args.scratchfull is not None:
        scratch_dir = Path(args.scratchfull).expanduser().resolve()
        validate_scratchfull_option(scratch_dir)
        try:
            os.makedirs(scratch_dir, exist_ok=True)
        except PermissionError:
            permission_denied_error_message(scratch_dir)
            raise
        return scratch_dir

    # --scratch option, Create the root scratch directory path
    root_scr_dir = Path(args.scratchdir).expanduser().resolve()

    try:
        os.makedirs(root_scr_dir, exist_ok=True)
    except PermissionError:
        permission_denied_error_message(root_scr_dir)
        raise

    # Preprocess to create scratch dir name
    input_file_name = input_file_path.stem  # Remove the extension (e.g. /path/to/active.inp => active)
    if input_file_name == "":
        input_file_name = "input"
    time_string: str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

    # Set the temporary directory's prefix using the input file name and the current time (e.g. input_2020-01-01_12-34-56_)
    prefix: str = f"{input_file_name}_{time_string}_"

    # Create a temporary directory (e.g. /home/username/dcaspt2_scratch/input_2020-01-01_12-34-56_XXXXXX)
    scratch_dir = tempfile.mkdtemp(dir=root_scr_dir, prefix=prefix)
    os.chmod(scratch_dir, 0o755)  # tempfile.mkdtemp creates a directory with 0o700 permission. So we change the permission to 0o755.
    scratch_dir = Path(scratch_dir).resolve()
    return scratch_dir


def copy_essential_files_to_scratch_dir(args: "argparse.Namespace", input_file_path: Path, integrals_dir: Path, scratch_dir: Path) -> None:
    """Copy input file and integral files to scratch directory
    Args:
        input_file_path (Path): Path to the input file
        integrals_dir (Path): Directory that 1-2 integrals files are stored
        scratch_dir (Path): Path to the temporary directory

    Returns:
        None
    """

    input_file_dir = input_file_path.parent  # (e.g. /home/username/active.inp => /home/username)
    if input_file_dir != scratch_dir:
        scratch_input_file_path = scratch_dir / "active.inp"  # The name of input file is fixed to active.inp in the scratch directory.
        scratch_input_file_path.symlink_to(input_file_path)
        restart_files = input_file_dir / "caspt2_restart*"
        for restart_file in glob.glob(str(restart_files)):
            scratch_restart_file_path = scratch_dir / Path(restart_file).name
            scratch_restart_file_path.symlink_to(restart_file)
    if integrals_dir != scratch_dir:
        # Create a symbolic link to the 1-2 integrals file
        scratch_mrconee_path = scratch_dir / "MRCONEE"
        scratch_mrconee_path.symlink_to(integrals_dir / "MRCONEE")
        dfpcmo_file_path = integrals_dir / "DFPCMO"
        if dfpcmo_file_path:
            scratch_dfpcmo_path = scratch_dir / "DFPCMO"
            scratch_dfpcmo_path.symlink_to(integrals_dir / "DFPCMO")
        mdcints = integrals_dir / "MDCIN*"
        for file in glob.glob(str(mdcints)):
            file_name = Path(file).name
            scratch_mdcint_path = Path(scratch_dir, file_name)
            scratch_mdcint_path.symlink_to(integrals_dir / file_name)
    return None


def set_integrals_dir(args: "argparse.Namespace", user_submitted_dir: Path) -> Path:
    """Set the directory to the 1-2 integrals file

    default: user_submitted_dir
    otherwise: args.integrals

    Args:
        args ("argparse.Namespace"):User input arguments (argv)
        user_submitted_dir (Path): Path to the directory where this script is used

    Returns:
        integrals_dir (Path): Directory that 1-2 integrals files are stored
    """
    if args.integrals is None:
        integrals_dir = user_submitted_dir.resolve()
    else:
        integrals_dir = Path(args.integrals).expanduser().resolve()
    return integrals_dir


def check_integral_files(integrals_dir: Path) -> None:
    """Check if the 1-2 integral files exist"""
    files = ["MDCINT", "MRCONEE"]
    for file in files:
        integral_path = (integrals_dir / file).resolve()
        if not integral_path.exists():
            raise Exception(
                f"\n{integral_path} file is not found under {integrals_dir}.\n\
Please specify the directory where 1-2 integrals files are stored with the -int or --integrals option.\n\
(e.g.) If you saved a MDCINT file at /home/username/H2O/MDCINT, use the following command: dcaspt2 -i input_file_path --int /home/username/H2O/.\n\
For more information, please use the --help option ( dcaspt2 --help )."
            )
    return


def create_command(args: "argparse.Namespace", binary_dir: Path) -> str:
    """Create a command to run dcaspt2

    Args:
        args ("argparse.Namespace"): User input arguments (argv)
        binary_dir (Path): Path to the directory where the compiled binaries are stored

    Returns:
        str: Created command
    """
    # --omp option
    if args.threads < 1:
        pass  # This script don't set OMP_NUM_THREADS
    else:
        os.environ["OMP_NUM_THREADS"] = str(args.threads)

    # --mpi option
    mpi = f"mpirun -np {args.mpi}" if args.mpi > 1 else ""
    bin_path = binary_dir / "r4dcaspt2exe"
    command = f"{mpi} {bin_path}"

    print(f"Created command : {command}")
    return command


def run_command(args: "argparse.Namespace", command: str, scratch_dir: Path, user_submitted_dir: Path, output_file_path: Path, version: str) -> None:
    """Run the dcaspt2 command

    Args:
        args ("argparse.Namespace"): User input arguments (argv)
        command (str): Command to run dcaspt2
        scratch_dir (Path): Path to the scratch directory
        user_submitted_dir (Path): Path to the directory where this script is used
        output_file_path (Path): Path to the output file
        version (str): Version of the dirac-caspt2 binary
    """

    def delete_scratch_dir():
        # Delete the scratch directory if the user didn't specify the -s or --save option
        if not args.save and scratch_dir.exists():
            shutil.rmtree(scratch_dir)

    def copy_specified_files_to_user_submitted_dir() -> None:
        """Copy the specified files from the scratch directory to the user submitted directory.
        If --save option is not specified, move the files instead of copying them.
        """

        if args.get is None:
            return  # No files to copy

        if user_submitted_dir == scratch_dir:
            print("The scratch directory is the user submitted directory. No need to copy files.")
            return  # No need to copy files

        print(args.get)
        get_files = str(args.get).split()  # Split the string by space
        for file_glob in get_files:  # glob is allowed (e.g.) MDCINTNEW* => MDCINTNEW, MDCINTNEW1, MDCINTNEW2, ...)
            for file in glob.glob(file_glob):
                scratch_dir_file = scratch_dir / file
                print(f"Copying {scratch_dir_file} to {user_submitted_dir}")
                if scratch_dir_file.exists():
                    if args.save:
                        try:
                            shutil.copy(scratch_dir_file, user_submitted_dir)
                        except shutil.SameFileError:
                            print(f"File {scratch_dir_file} already exists in {user_submitted_dir}")
                        except PermissionError:
                            print(f"Permission denied: {scratch_dir_file}")
                    else:
                        try:
                            shutil.move(scratch_dir_file, user_submitted_dir)
                        except shutil.Error:
                            print(f"File {scratch_dir_file} already exists in {user_submitted_dir}")
                        except PermissionError:
                            print(f"Permission denied: {scratch_dir_file}")
                else:
                    print(f"File {scratch_dir_file} is not found")

    def get_elapsed_time(start_time: datetime, end_time: datetime) -> Tuple[int, int, int, int, int]:
        """Get the elapsed time

        Args:
            start_time (datetime): Start time
            end_time (datetime): End time

        Returns:
            Tuple[int, int, int, int, int]: Elapsed time (days, hours, minutes, seconds, milliseconds)
        """
        diff_time = end_time - start_time
        time_hours = diff_time.seconds // 3600
        time_minutes = (diff_time.seconds % 3600) // 60
        time_seconds = diff_time.seconds % 60
        time_milliseconds = diff_time.microseconds // 1000
        return diff_time.days, time_hours, time_minutes, time_seconds, time_milliseconds

    def generate_end_message(p: Popen[str], start_time: datetime, end_time: datetime) -> str:
        """Generate the end message

        Args:
            p (Popen): Popen object

        Returns:
            str: End message
        """

        def generate_err_message() -> str:
            """Generate the error message"""
            err_msg = "\n================= Standard error =================\n"
            if p.stderr is None:
                return err_msg
            else:
                return f"{err_msg}{p.stderr.read()}"

        def generate_finished_message(start_time: datetime, end_time: datetime) -> str:
            """Generate the finished message"""
            start_time_str = start_time.strftime("%Y-%m-%d %H:%M:%S")
            end_time_str = end_time.strftime("%Y-%m-%d %H:%M:%S")
            user_command = " ".join(sys.argv)
            days, hours, mins, secs, millisecs = get_elapsed_time(start_time, end_time)
            finished_message = f"\n================= Calculation finished ================\n\
User Command : {user_command}\n\
Auto-created Command : {command}\n\
Scratch directory : {scratch_dir}\n\
Output file : {output_file_path}\n\
Calculation started at : {start_time_str}\nCalculation finished at : {end_time_str}\n\
Elapsed time (sec) : {(end_time - start_time).total_seconds():.4f} sec\n\
Elapsed time : {days} day {hours} hour {mins} min {secs} sec {millisecs:03} millisecond\n\
dirac-caspt2 version (commit hash) : {version}\n"
            return finished_message

        def generate_returncode_condition_message() -> str:
            """Generate the return code condition message"""
            if p.returncode == 0:
                return "NORMAL END OF dirac-caspt2 CALCULATION\n"
            return f"ERROR: dirac-caspt2 calculation failed with return code {p.returncode}\n"

        end_message = generate_err_message() + generate_finished_message(start_time, end_time) + generate_returncode_condition_message()

        return end_message

    os.chdir(scratch_dir)

    with open(output_file_path, "w") as f:
        start_time = datetime.now()

        try:
            p = Popen(command, shell=True, stdout=f, stderr=PIPE, text=True, encoding="utf-8")
            p.wait()
        except KeyboardInterrupt:
            delete_scratch_dir()
            end_message = "\nProcess is interrupted by the user.(ctrl+c, SIGINT)\nExiting..."
            print(end_message)
            f.write(end_message)
            raise  # Raise the KeyboardInterrupt to exit the script

        end_time = datetime.now()
        end_message = generate_end_message(p, start_time, end_time)

        # Write the end message to the output file and stdout
        print(end_message)
        f.write(end_message)

        copy_specified_files_to_user_submitted_dir()
        delete_scratch_dir()

        # When the calculation failed, exit with the return code of the calculation
        if p.returncode != 0:
            sys.exit(p.returncode)

        return None


def get_version(binary_dir: Path) -> str:
    """Get the version of the dirac-caspt2 binary

    Args:
        binary_dir (Path): Path to the directory where the compiled binaries are stored

    Returns:
        str: Version of the dirac-caspt2 binary
    """
    version_file = binary_dir / ".commit_hash"
    if version_file.exists():
        versions = version_file.read_text(encoding="utf-8").splitlines()
        return "unknown" if len(versions) == 0 else versions[0].strip()
    return "unknown"


def set_input_file_path(args: "argparse.Namespace", user_submitted_dir: Path) -> Path:
    """Set the path to the input file

    Args:
        args ("argparse.Namespace"): User input arguments (argv)
        user_submitted_dir (Path): Path to the directory where this script is used

    Returns:
        Path: Path to the input file
    """
    if args.input is None:
        input_file_path = (user_submitted_dir / "active.inp").resolve()
    else:
        input_file_path = Path(args.input).expanduser().resolve()

    # Check if the input file exists
    if not input_file_path.exists():
        msg = f"Input file {input_file_path} does not exist.\nPlease specify the input file path with -i or --input option."
        raise FileNotFoundError(msg)

    return input_file_path


def set_ouput_file_path(args: "argparse.Namespace", user_submitted_dir: Path) -> Path:
    """Set the path to the output file

    Args:
        args ("argparse.Namespace"): User input arguments (argv)
        user_submitted_dir (Path): Path to the directory where this script is used

    Returns:
        Path: Path to the output file
    """
    if args.output is None:
        output_file_path = (user_submitted_dir / "dirac_caspt2.out").resolve()
    else:
        output_file_path = Path(args.output).expanduser().resolve()

    # Does output_file_path have w+ permission?
    try:
        with open(output_file_path, "w"):
            pass
    except PermissionError:
        msg = f"Permission denied: You do not have write permission to the file '{output_file_path}'.\n\
Please specify the output file that you have permission to write.\nExiting..."
        print(msg)
        raise

    return output_file_path


def main() -> None:
    user_submitted_dir: Path = Path.cwd().resolve()
    script_path = Path(__file__).resolve()
    binary_dir = script_path.parent
    args: "argparse.Namespace" = parse_args()
    print("Script path : ", script_path)
    print("Command submitted directory : ", user_submitted_dir)

    input_file_path = set_input_file_path(args, user_submitted_dir)
    output_file_path = set_ouput_file_path(args, user_submitted_dir)
    version = get_version(binary_dir)
    print(f"dirac-caspt2 version (commit hash): {version}")

    # 1-2 integrals directory
    integrals_dir = set_integrals_dir(args, user_submitted_dir)
    check_integral_files(integrals_dir)

    # Set and create scratch directory
    scratch_dir = create_scratch_dir(args, input_file_path, user_submitted_dir)
    print(f"Scratch directory : {scratch_dir}")
    copy_essential_files_to_scratch_dir(args, input_file_path, integrals_dir, scratch_dir)

    # Create command to execute and run it
    command = create_command(args, binary_dir)
    run_command(args, command, scratch_dir, user_submitted_dir, output_file_path, version)


if __name__ == "__main__":
    main()