#!/usr/bin/python3
########################################################################################
# hpc_outbox_assembly.py - CLOVER optimisation outbox assembly script.                 #
#                                                                                      #
# Authors: Ben Winchester                                                              #
# Copyright: Ben Winchester, 2022                                                      #
# Date created: 13/07/2021                                                             #
# License: Open source                                                                 #
#                                                                                      #
# For more information, please email:                                                  #
#   benedict.winchester@gmail.com                                                      #
########################################################################################
"""
hpc_outbox_assembly.py - CLOVER optimisation outbox assembly script.

This script assembles output files from optimisations that were carried out on Imperial
College London's High-Performance Computer(s) (HPC) into a single "outbox directory",
from where they can be easily copied from the HPC for futher analysis.

"""

import argparse
import os
import shutil
import subprocess
import sys

from typing import Any, List

# Destination directory:
#   THe name of the outbox directory into which to save optimisation output files.
DESTINATION_DIRECTORY: str = "optimisation_outbox"

# Optimisation name:
#   The name to use when finding optimisation output files.
OPTIMISATION_NAME: str = "optimisation_output_*.json"

# Simulation name:
#   The name to use when finding simulation output files.
SIMULATION_NAME: str = "simulation_output_*.csv"


def _parse_outbox_args(args: List[Any]) -> argparse.Namespace:
    """
    Parses the input arguments to the hpc outbox assembly script to determine the flow.

    Inputs:
        - args:
            Unparsed command-line arguments.

    Outputs:
        - The parsed HPC arguments.

    """

    parser = argparse.ArgumentParser()

    # Optimisation outputs:
    #   Whether to assemble optimisation outputs.
    parser.add_argument(
        "--optimisation",
        "-opt",
        action="store_true",
        default=False,
        help="Whether to include optimisation output files.",
    )

    # Optimisation outputs:
    #   Whether to assemble optimisation outputs.
    parser.add_argument(
        "--simulation",
        "-sim",
        action="store_true",
        default=False,
        help="Whether to include simulation output files.",
    )

    return parser.parse_args(args)


def main(args: List[Any]) -> None:
    """Main entry point for the CLOVER HPC outbox assembly script."""

    # Parse the command-line arguments.
    parsed_args = _parse_outbox_args(args)

    # Raise an error if no assembly was requested.
    if not parsed_args.simulation or parsed_args.optimisation:
        raise Exception(
            "Outbox assembly script must be called to either assemble optimisation or "
            "simulation output files or both. Neither argument was specified. Run with "
            "`--help` for more information."
        )

    filenames: List[str] = []

    if parsed_args.optimisation:
        # Determine the output files that have been generated by running an optimisation.
        optimisation_filenames = subprocess.run(
            ["find", "locations", "-name", OPTIMISATION_NAME],
            stdout=subprocess.PIPE,
            check=True,
        ).stdout.decode()
        optimisation_filenames_list = optimisation_filenames.split("\n")
        optimisation_filenames_list.remove("")
        filenames += optimisation_filenames_list

    if parsed_args.simulation:
        # Determine the output files that have been generated by running an optimisation.
        simulation_filenames = subprocess.run(
            ["find", "locations", "-name", SIMULATION_NAME],
            stdout=subprocess.PIPE,
            check=True,
        ).stdout.decode()
        simulation_filenames_list = simulation_filenames.split("\n")
        simulation_filenames_list.remove("")
        filenames += simulation_filenames_list

    # Copy each of these files to the outbox directory, including their directory tree
    # structure.
    for source_filename in filenames:
        destination_filename = os.path.join(DESTINATION_DIRECTORY, source_filename)

        # Make the directory structure within the outbox if it doesn't already exist
        if not os.path.isdir(os.path.dirname(destination_filename)):
            os.makedirs(os.path.dirname(destination_filename))

        # Copy across our file
        shutil.copy2(source_filename, destination_filename)


if __name__ == "__main__":
    main(sys.argv[1:])
