#!/usr/bin/env python


"""Set up and run calcdb.x.

Run calcdb.x to compute ground magnetic field perturbations for a MAGE
magnetosphere simulation. This is all done by a linked set of PBS jobs running
in the MAGE results directory.

NOTE: Before running this script, the user must run the setupEnvironment.sh
scripts for kaipy and kaiju.

NOTE: If the calcdb.x binary was built with a module set other than those
listed below, change the module set in the PBS scripts appropriately.

Author
------
Eric Winter (eric.winter@jhuapl.edu)

"""


# Standard modules.
import argparse
import copy
import os
import pathlib
import re
import shutil
import sys

# Third-party modules.
from jinja2 import Template

# Kaipy modules.
from kaipy import kaiH5
from kaipy import kaiTools


# Program constants and defaults

# Program description.
DESCRIPTION = "Run calcdb.x to compute ground delta-B values."

# Default values for command-line arguments.
DEFAULT_ARGUMENTS = {
    "calcdb": "calcdb.x",
    "debug": False,
    "dt": 60.0,
    "hpc": "pleiades",
    "parintime": 1,
    "pbs_account": None,
    "verbose": False,
    "mage_results_path": None,
}

# Valid values for command-line arguments.
VALID_HPC = ["derecho", "pleiades"]

# Directory containing this script.
SCRIPT_DIRECTORY = pathlib.Path(__file__).parent.resolve()
TEMPLATE_DIRECTORY = os.path.join(SCRIPT_DIRECTORY, "templates")

# Location of template XML file for calcdb.x.
CALCDB_XML_TEMPLATE = os.path.join(TEMPLATE_DIRECTORY, "calcdb-template.xml")

# # Locations of template PBS file for running calcdb.x.
CALCDB_PBS_TEMPLATE = os.path.join(TEMPLATE_DIRECTORY, "calcdb-template.pbs")

# Location of template PBS file for pitmerge.py.
PITMERGE_PBS_TEMPLATE = os.path.join(
    TEMPLATE_DIRECTORY, "pitmerge-template.pbs"
)


def create_command_line_parser():
    """Create the command-line parser.

    Create the parser for the command line.

    Parameters
    ----------
    None

    Returns
    -------
    parser : argparse.ArgumentParser
        Command-line parser for this script.

    Raises
    ------
    None
    """
    parser = argparse.ArgumentParser(description=DESCRIPTION)
    parser.add_argument(
        "--calcdb", type=str, default=DEFAULT_ARGUMENTS["calcdb"],
        help="Path to calcdb.x binary (default: %(default)s)."
    )
    parser.add_argument(
        "--debug", "-d", default=DEFAULT_ARGUMENTS["debug"],
        action="store_true",
        help="Print debugging output (default: %(default)s)."
    )
    parser.add_argument(
        "--dt", type=float, default=DEFAULT_ARGUMENTS["dt"],
        help="Time interval for delta B computation (simulated seconds) "
             "(default: %(default)s)."
    )
    parser.add_argument(
        "--hpc", type=str, default=DEFAULT_ARGUMENTS["hpc"],
        choices=VALID_HPC,
        help="HPC system to run analysis (default: %(default)s)."
    )
    parser.add_argument(
        "--parintime", type=int, default=DEFAULT_ARGUMENTS["parintime"],
        help="Split the calculation into this many parallel chunks of kaiju "
        "simulation steps, one chunk per node (default: %(default)s)."
    )
    parser.add_argument(
        "--pbs_account", type=str, default=DEFAULT_ARGUMENTS["pbs_account"],
        help="PBS account to use for job accounting (default: %(default)s)."
    )
    parser.add_argument(
        "--verbose", "-v", default=DEFAULT_ARGUMENTS["verbose"],
        action="store_true",
        help="Print verbose output (default: %(default)s)."
    )
    parser.add_argument(
        "mage_results_path", type=str,
        default=DEFAULT_ARGUMENTS["mage_results_path"],
        help="Path to a result file for a MAGE magnetosphere run."
    )
    return parser


def mage_filename_to_runid(filename: str):
    """Parse the runid from a MAGE results file name.

    Parse the runid from a MAGE results file name.

    For a result file from an MPI run, the runid is all text before the
    underscore before the set of 6 underscore-separated 4-digit sequences at
    the end of the name and the terminal .gam.h5 string.

    For a result file from a serial run, the runid is the name of the file,
    less the .gam.h5 extension.

    Parameters
    ----------
    filename : str
        Name of MAGE results file.

    Returns
    -------
    runid : str
        The MAGE runid for the file.

    Raises
    ------
    None
    """
    # Check to see if the result file is for an MPI run. If not, it must be
    # for a serial run.
    mpi_pattern = (
        r"^(.+)_(\d{4})_(\d{4})_(\d{4})_(\d{4})_(\d{4})_(\d{4})\.gam.h5$"
    )
    serial_pattern = r"^(.+)\.gam.h5$"
    mpi_re = re.compile(mpi_pattern)
    serial_re = re.compile(serial_pattern)
    m = mpi_re.match(filename)
    if not m:
        m = serial_re.match(filename)
    runid = m.groups()[0]
    return runid


def create_calcdb_xml_file(runid: str, args: dict):
    """Create the XML input file for calcdb.x from a template.

    Create the XML input file for calcdb.x from a template. The file is
    created in the current directory, which much contain the specified MAGE
    results file. Note that the contents of the file depends on whether
    parintime is 1 (use a single multi-threaded job to run calcdb.x on all
    simulation steps) or > 1 (use multiple jobs to split the calcdb.x work
    into an array of multi-threaded jobs which run on multiple nodes in
    parallel).

    Parameters
    ----------
    runid : str
        Run ID for MAGE results file to use in calcdb.x calculations.
    args : dict
        Dictionary of command-line options.

    Returns
    -------
    xml_file : str
        Name of XML file.

    Raises
    ------
    AssertionError
        If the MAGE result file contains no time steps.
    """
    # Local convenience variables.
    debug = args["debug"]
    verbose = args["verbose"]

    # Fetch run information from the MAGE result file.
    if verbose:
        print(f"Fetching MAGE run information for runid {runid}.")
    filename, isMPI, Ri, Rj, Rk = kaiTools.getRunInfo(".", runid)
    if debug:
        print(f"filename = {filename}")
        print(f"isMPI = {isMPI}")
        print(f"Ri = {Ri}")
        print(f"Rj = {Rj}")
        print(f"Rk = {Rk}")

    # Get the number of steps and the step IDs from the MAGE results file.
    if verbose:
        print(f"Counting time steps for run {runid}.")
    nSteps, sIds = kaiH5.cntSteps(filename)
    if debug:
        print(f"nSteps = {nSteps}")
        print(f"sIds = {sIds}")
    assert nSteps > 0

    # Find the time for the last step.
    if verbose:
        print(f"Finding time for last step for run {runid}.")
    tFin = kaiH5.tStep(filename, sIds[-1], aID="time")
    if debug:
        print(f"tFin = {tFin}")

    # Read the template XML file.
    if verbose:
        print("Reading calcdb XML template.")
    with open(CALCDB_XML_TEMPLATE, "r", encoding="utf-8") as f:
        template_content = f.read()
    if debug:
        print(f"template_content = {template_content}")
    template = Template(template_content)
    if debug:
        print(f"template = {template}")

    # Fill in the template options.
    options = {
        "runid": runid,
        "dt": args["dt"],
        "tFin": tFin,
        "ebfile": runid,
        "isMPI": isMPI,
        "Ri": Ri,
        "Rj": Rj,
        "Rk": Rk,
        "parintime": args["parintime"],
    }
    if debug:
        print(f"options = {options}")

    # Render the template.
    xml_file = f"calcdb-{runid}.xml"
    if verbose:
        print("Rendering template.")
    if verbose:
        print(f"Creating {xml_file}.")
    xml_content = template.render(options)
    with open(xml_file, "w", encoding="utf-8") as f:
        f.write(xml_content)

    # Return the name of the XML file.
    return xml_file


def create_calcdb_pbs_script(runid: str, calcdb_xml_file: str, args: dict):
    """Create the PBS script to run calcdb.x from a template.

    Create the PBS script to run calcdb.x from a template. The PBS script is
    created in the current directory, which must contain the MAGE results.

    Parameters
    ----------
    runid : str
        Run ID for MAGE results to use in calcdb.x calculations.
    calcdb_xml_file : str
        Path to XML input file for calcdb.x.
    args : dict
        Dictionary of command-line and other options.

    Returns
    -------
    calcdb_pbs_script : str
        Path to PBS script.

    Raises
    ------
    None
    """
    # Local convenience variables.
    debug = args["debug"]
    verbose = args["verbose"]

    # Copy the calcdb.x binary to the results directory, then make it
    # executable.
    if verbose:
        print(f"Copying {args['calcdb']} to results directory.")
    shutil.copyfile(args["calcdb"], "calcdb.x")
    os.chmod("calcdb.x", 0o755)

    # Read the PBS script template for running calcdb.x.
    if verbose:
        print("Reading calcdb.x PBS script template.")
    with open(CALCDB_PBS_TEMPLATE, "r", encoding="utf-8") as f:
        template_content = f.read()
    if debug:
        print(f"template_content = {template_content}")
    template = Template(template_content)
    if debug:
        print(f"template = {template}")

    # Fill in the template options.
    options = {
        "hpc": args["hpc"],
        "job_name": f"calcdb-{runid}",
        "pbs_account": args["pbs_account"],
        "conda_prefix": os.environ["CONDA_PREFIX"],
        "kaipyhome": os.environ["KAIPYHOME"],
        "kaijuhome": os.environ["KAIJUHOME"],
        "parintime": args["parintime"],
        "calcdb_xml_file": calcdb_xml_file,
    }
    if debug:
        print(f"options = {options}")

    # Render the template.
    if verbose:
        print("Rendering template.")
    calcdb_pbs_script = f"calcdb-{runid}.pbs"
    calcdb_pbs_content = template.render(options)
    with open(calcdb_pbs_script, "w", encoding="utf-8") as f:
        f.write(calcdb_pbs_content)

    # Return the name of the PBS script.
    return calcdb_pbs_script


def create_pitmerge_pbs_script(runid: str, args: dict):
    """Create the PBS script for pitmerge.py to stitch calcdb output.

    Create the PBS script for stitching together calcdb output using the
    pitmerge.py script.

    Parameters
    ----------
    runid : str
        Run ID for MAGE results used in calcdb.x calculations.
    args : dict
        Dictionary of command-line and other options.

    Returns
    -------
    stitching_pbs_script : str
        Path to PBS script.

    Raises
    ------
    None
    """
    # Local convenience variables.
    debug = args["debug"]
    verbose = args["verbose"]

    # Read the PBS script template for pitmerge.py.
    if verbose:
        print("Reading pitmerge.py PBS template.")
    with open(PITMERGE_PBS_TEMPLATE, "r", encoding="utf-8") as f:
        template_content = f.read()
    if debug:
        print(f"template_content = {template_content}")
    template = Template(template_content)
    if debug:
        print(f"template = {template}")

    # Fill in the template options.
    options = {
        "hpc": args["hpc"],
        "job_name": f"pitmerge-{runid}",
        "pbs_account": args["pbs_account"],
        "conda_prefix": os.environ["CONDA_PREFIX"],
        "kaipyhome": os.environ["KAIPYHOME"],
        "kaijuhome": os.environ["KAIJUHOME"],
        "runid": runid,
    }
    if debug:
        print(f"options = {options}")

    # Render the template.
    if verbose:
        print("Rendering template.")
    pitmerge_pbs_script = f"pitmerge-{runid}.pbs"
    pitmerge_pbs_content = template.render(options)
    with open(pitmerge_pbs_script, "w", encoding="utf-8") as f:
        f.write(pitmerge_pbs_content)

    # Return the name of the PBS script.
    return pitmerge_pbs_script


def create_submit_script(
        runid: str, calcdb_pbs_script: str, pitmerge_pbs_script: str,
        args: dict):
    """Create the bash script to submit all of the PBS jobs.

    Create the bash script to submit all of the PBS jobs. The submit script
    submits the job to run calcdb.x, then the job to run pitmerge.py on the
    results (if needed), then the analysis job. Each job only runs if the
    previous job completes successfully.

    Parameters
    ----------
    runid : str
        Run ID for MAGE results used in calcdb.x calculations.
    calcdb_pbs_script : str
        Path to calcdb.x PBS script.
    pitmerge_pbs_script : str
        Path to pitmerge.py PBS script (or None if not needed).
    args : dict
        Dictionary of command-line and other options.

    Returns
    -------
    submit_script : str
        Path to bash script to submit PBS jobs.

    Raises
    ------
    None
    """
    # Submit the scripts in dependency order.
    submit_script = f"submit-{runid}.sh"
    with open(submit_script, "w", encoding="utf-8") as f:

        # calcdb.x job
        if args["parintime"] > 1:
            cmd = (f"job_id=`qsub -J 1-{args['parintime']} "
                   f"{calcdb_pbs_script}`\n")
        else:
            cmd = f"job_id=`qsub {calcdb_pbs_script}`\n"
        f.write(cmd)
        cmd = "echo $job_id\n"
        f.write(cmd)
        cmd = "old_job_id=$job_id\n"
        f.write(cmd)

        # pitmerge.py job (if needed).
        if pitmerge_pbs_script is not None:
            cmd = (
                "job_id=`qsub -W depend=afterok:$old_job_id "
                f"{pitmerge_pbs_script}`\n"
            )
            f.write(cmd)
            cmd = "echo $job_id\n"
            f.write(cmd)
            cmd = "old_job_id=$job_id\n"
            f.write(cmd)

    # Return the name of the PBS script.
    return submit_script


def run_calcdb(args: dict):
    """Compute ground delta-B values for a MAGE magnetosphere run.

    Compute ground delta-B values for a MAGE magnetosphere run.

    Parameters
    ----------
    args: dict
        Dictionary of command-line options.

    Returns
    -------
    int 0 on success

    Raises
    ------
    AssertionError
        If an invalid argument is provided.
    """
    # Set defaults for command-line options, then update with values passed
    # from the caller.
    local_args = copy.deepcopy(DEFAULT_ARGUMENTS)
    local_args.update(args)
    args = local_args

    # Local convenience variables.
    debug = args["debug"]
    verbose = args["verbose"]

    # Validate arguments.
    assert len(args["calcdb"]) > 0
    assert args["dt"] > 0.0
    assert args["hpc"] in VALID_HPC
    assert args["parintime"] > 0
    assert len(args["mage_results_path"]) > 0

    # ------------------------------------------------------------------------

    # Split the MAGE results path into a directory and a file.
    (mage_results_dir, mage_results_file) = os.path.split(
        args["mage_results_path"]
    )
    if debug:
        print(f"mage_results_dir = {mage_results_dir}")
        print(f"mage_results_file = {mage_results_file}")

    # Move to the results directory.
    if verbose:
        print(f"Moving to MAGE results directory {mage_results_dir}.")
    os.chdir(mage_results_dir)

    # Compute the runid from the file name.
    if verbose:
        print(f"Computing runid from MAGE results file {mage_results_file}")
    runid = mage_filename_to_runid(mage_results_file)
    if debug:
        print(f"runid = {runid}")

    # Create the XML input file for calcdb.x.
    if verbose:
        print(f"Creating XML file for calcdb.x for runid {runid}.")
    calcdb_xml_file = create_calcdb_xml_file(runid, args)
    if debug:
        print(f"calcdb_xml_file = {calcdb_xml_file}")

    # Create the PBS script to run calcdb.x.
    if verbose:
        print(f"Creating PBS script to run calcdb.x for run {runid}.")
    calcdb_pbs_script = create_calcdb_pbs_script(runid, calcdb_xml_file, args)
    if debug:
        print(f"calcdb_pbs_script = {calcdb_pbs_script}")

    # Create the PBS script to stitch together the output from calcdb.x
    # using pitmerge.py. This script is only needed if parintime > 1.
    if args["parintime"] > 1:
        if verbose:
            print("Creating PBS script to stitch together the calcdb.x output "
                  f" for MAGE runid {runid}")
        pitmerge_pbs_script = create_pitmerge_pbs_script(runid, args)
    else:
        pitmerge_pbs_script = None
    if debug:
        print(f"pitmerge_pbs_script = {pitmerge_pbs_script}")

    # Create the bash script to submit the PBS scripts in the proper order.
    if verbose:
        print("Creating bash script to submit the PBS jobs.")
    submit_script = create_submit_script(
        runid, calcdb_pbs_script, pitmerge_pbs_script, args
    )
    if debug:
        print(f"submit_script = {submit_script}")

    if verbose:
        print(f"Please run {submit_script} (in the MAGE result directory) to "
              "submit the PBS jobs to perform the MAGE-SuperMag comparison. "
              "If you are not using PBS, then you can manually run the "
              f"scripts individually in the order listed in {submit_script}.")

    # Return normally.
    return 0


def main():
    """Driver for command-line version of code."""
    # Set up the command-line parser.
    parser = create_command_line_parser()

    # Parse the command-line arguments.
    args = parser.parse_args()
    if args.debug:
        print(f"args = {args}")

    # Convert the arguments from Namespace to dict.
    args = vars(args)

    # Pass the command-line arguments to the main function as a dict.
    return_code = run_calcdb(args)
    sys.exit(return_code)


if __name__ == "__main__":
    main()
