#!/usr/bin/env python
"""
Generates an e-tree from a given model string using MPI to parallelize the task.
This utility can read in the new XML configuration format. If no file-format is specified,
then this utility will ask a series of prompts before generating the e-tree.

Copyright 2017 Southern California Earthquake Center

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
# Python Imports
import sys
import time
import os

os.environ["ucvm_has_bootstrapped"] = "Yes"

# Package Imports
import xmltodict
from mpi4py import MPI

# UCVM Imports
from ucvm.src.framework.ucvm import UCVM
from ucvm.src.framework.etree import ask_questions, etree_extract_mpi


def usage() -> None:
    """
    Displays the help text associated with this utility.

    Returns:
        Nothing
    """
    UCVM.print_with_replacements(
        "\n"
        "ucvm_etree_create_mpi - UCVM Version [version]\n"
        "\n"
        "Creates an e-tree or an e-tree configuration file with the given parameters. This generates\n"
        "a file that is compatible with wave propagation code such as Hercules. If your wave propagation\n"
        "code (such as AWP-ODC) uses the binary float mesh format, then you need to use the\n"
        "ucvm_mesh_create or ucvm_mesh_create_mpi utility.\n"
        "\n"
        "-c, --config-only c:   Generates the XML-style configuration file only. No mesh will be\n"
        "                       made at the end of the questions.\n"
        "-f, --file f:          Specifies the configuration file from which this utility should read.\n"
        "-r, --rows r:          Extracts one or more rows of the etree. If r is a number like '1'\n"
        "                       just that row will be extracted. If r is an interval like '1-5'\n"
        "                       those rows (1, 2, 3, 4, 5 e.g.) inclusive will be extracted.\n"
        "-i, --interval i:      Extracts exactly the interval desired. The parameter must be\n"
        "                       specified as row,column-row,column. E.g. the parameters\n"
        "                       1,5-10,12 would extract row 1, column 5 through row 10,\n"
        "                       column 12 inclusive."
    )


def main() -> int:
    """
    The main UCVM e-tree create MPI function.

    Returns:
         0 if successful. Raises an error code otherwise, if not.
    """
    start_time = time.time()

    try:
        options = UCVM.parse_options([
            {"short": "c", "long": "config-only", "value": True, "required": False},
            {"short": "f", "long": "file", "value": True, "required": False},
            {"short": "r", "long": "rows", "value": True, "required": False},
            {"short": "i", "long": "interval", "value": True, "required": False}
        ], usage)
    except ValueError as v_err:
        print("[ERROR]: " + str(v_err) + "\n")
        sys.exit(-1)

    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()

    should_exit = False

    if rank == 0:
        if options["config-only"] is not None:
            with open(options["config-only"], "w") as fd:
                fd.write(xmltodict.unparse({"root": ask_questions()}, pretty=True))
            print(
                "\n"
                "Your e-tree configuration file has been saved to " + options["config-only"] + ".\n"
                "\n"
                "To extract an e-tree with your new configuration, run\n"
                "ucvm_etree_create_mpi -f " + options["config-only"] + ""
                "\n"
            )
            should_exit = comm.bcast(True, root=0)

        if options["file"] is not None:
            with open(options["file"], "r") as fd:
                etree_information = xmltodict.parse(fd.read())["root"]
        else:
            etree_information = ask_questions()
    else:
        etree_information = None

    etree_information = comm.bcast(etree_information, root=0)

    # If we were just using the MPI utility to generate a configuration file, then exit.
    if should_exit:
        return 0

    comm.Barrier()

    etree_extract_mpi(etree_information, options["rows"], options["interval"])

    if rank == 0:
        print("\nE-tree extraction finished in %s seconds" % (time.time() - start_time), flush=True)

if __name__ == "__main__":
    sys.exit(main())
