"""
This module contains functions for constructing SGE or SLURM submission scripts
for running Gaussian calculations
"""

import os
import hpc_suite as hpc


def create_submission_script(f_name, n_cores, memory, machine, gaussian_module,
                             gaussian_scratch, gaussian_chk, gaussian_version,
                             gen_fchk=True):
    """
    Creates Gaussian submission script for HPC systems

    Parameters
    ----------
        f_name : str
            File name of .data file generated by create_data_file function
        n_cores : int
            Number of cores (threads) to request
        memory : int
            Total memory in gigabytes
        machine : str {"SLURM", "SGE"}
            Type of HPC system
        gaussian_module : str
            Path to Gaussian module for module load command
        gaussian_scratch : str
            Path to Gaussian scratch folder
        gaussian_chk : str
            Filename (with path to $CurrDir for CSF3) of checkpoint file
        gaussian_version : str
            version of Gaussian
        gen_fchk : bool, default True
            If true, generates formatted checkpoint file using `formchk`
            utility

    Returns
    -------
        None
    """

    f_head = os.path.splitext(f_name)[0]

    args = hpc.read_args(['generate_job'])

    args.profile = machine
    args.omp = n_cores

    # Set environmental variables
    args.env_vars["CurrDir"] = "$(pwd -P)"
    args.env_vars["JNAME"] = f_head
    args.env_vars["GAUSS_SCRDIR"] = gaussian_scratch
    args.env_vars["GAUSS_MDEF"] = "{:d}GB".format(memory)
    args.env_vars["GAUSS_YDEF"] = gaussian_chk
    args.env_vars["GAUSS_PDEF"] = n_cores

    # Set gaussian module
    if gaussian_module:
        args.modules.append(gaussian_module)

    # Set job, input, output and error names
    args.job_name = f_head

    body = "mkdir -p $GAUSS_SCRDIR\n"

    # Define call to gaussian program
    gaussian_args = "{} < $JNAME.com > $JNAME.log".format(
        gaussian_version
    )
    body += "{}\n".format(gaussian_args)

    if gen_fchk:
        body += "formchk {}\n".format(gaussian_chk)

    args.body = body

    # Generate job submission script
    hpc.generate_job_func(args)

    return
