#!/usr/bin/env python
"""
Generates a binary float mesh 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 mesh.

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 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.awp_mesh import ask_questions, mesh_extract_mpi
from ucvm.src.framework.mesh_common import InternalMesh


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

    Returns:
        Nothing
    """
    UCVM.print_with_replacements(
        "\n"
        "ucvm_mesh_create_mpi - UCVM Version [version]\n"
        "\n"
        "Creates a mesh or a mesh configuration file with the given parameters. This generates\n"
        "a binary float mesh that is compatible with wave propagation code such as AWP-ODC. If\n"
        "your wave propagation code (such as Hercules) uses the e-tree format, then you need\n"
        "to use the ucvm_etree_create or ucvm_etree_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\n"
        "                       read. Note that this auto-detects a legacy (before UCVM 15.10.0)\n"
        "                       style configuration file vs. the new XML format.\n"
    )


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

    Returns:
         0 if successful. Raises an error code otherwise, if not.
    """
    try:
        options = UCVM.parse_options([
            {"short": "c", "long": "config-only", "value": True, "required": False},
            {"short": "f", "long": "file", "value": True, "required": False},
            {"short": "s", "long": "slices", "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()
    size = comm.Get_size()

    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 mesh configuration file has been saved to " + options["config-only"] + ".\n"
                "\n"
                "To extract a mesh with your new configuration, run\n"
                "ucvm_mesh_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:
                mesh_information = xmltodict.parse(fd.read())["root"]
        else:
            mesh_information = ask_questions()
    else:
        mesh_information = None

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

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

    i_mesh = InternalMesh(mesh_information)
    max_points_per_cpu = i_mesh.get_max_points_extract(2)

    # Check to see if we are doing an interval or if we are doing one or more slices.
    if options["slices"] is not None:
        i_mesh.do_slices(options["slices"])
    elif options["interval"] is not None:
        i_mesh.do_interval(options["interval"])

    # All processes that are not rank 0 should send the number of computing CPUs that they have
    # available to them.
    if rank != 0:
        comm.send({
            "max_per_cpu": max_points_per_cpu
        }, dest=0)

    # Rank zero has the task of compiling the topology of CPUs available to us.
    if rank == 0:
        # Make each node responsible for its share of the mesh.
        print("Starting MPI mesh extraction for " + mesh_information["mesh_name"] + "\n")

        points_per_process = int(i_mesh.total_size / size) + 1

        for i in range(1, size):
            if i == size - 1:
                comm.send((i * points_per_process + i_mesh.start_point, i_mesh.total_size + i_mesh.start_point), dest=i)
            else:
                comm.send((i * points_per_process + i_mesh.start_point, (i + 1) * points_per_process +
                           i_mesh.start_point), dest=i)

        (start_point, end_point) = (i_mesh.start_point, points_per_process + i_mesh.start_point)
    else:
        (start_point, end_point) = comm.recv(source=0)

    sys.stdout.flush()

    comm.Barrier()

    mesh_extract_mpi(mesh_information, (start_point, end_point))

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