#!/usr/bin/env python3

"""
Generated the spherical graph that can then be used to navigate.
The graph gets serialized to disk at the end of the calculation.

Examples:
    Generate a graph and visualize it. Because it makes using it in search algorithms faster, we also include
    the neighbor table. the also prunes by default using the *Earth2014* dataset (variant *TBI*, 1 arc-min
    resolution).

    .. code-block:: bash

       ./scripts/create_earth_graph.py 500000 earth_graph_500_km.hdf5

"""

# Standard library
from argparse import ArgumentDefaultsHelpFormatter
from argparse import ArgumentParser
import os.path

# Data set access
from pyrate.common.raster_datasets import DataSetAccess
from pyrate.common.raster_datasets import transformers_concrete

# Graph generation
from pyrate.plan.graph import create_earth_graph
from pyrate.plan.graph import min_required_frequency

# Script visualize_earth_graph
try:
    from visualize_earth_graph import dump_2d_plots
except ImportError:
    # add scripts folder
    import sys

    sys.path.append(os.path.abspath(os.path.dirname(__file__)))
    del sys

    from visualize_earth_graph import dump_2d_plots


def calculate_and_save(  # pylint: disable=too-many-arguments
    requested_distance: float,
    out_graph_file: str,
    prune_land_areas: bool,
    bathymetric_dataset: str,
    generate_neighbors: bool,
    dump_plots: bool,
    out_visualization_directory: str,
) -> None:
    """Calculates and saves the graph with the given node distance while performing some logging.

    Args:
        requested_distance: the maximum distance between two neighboring nodes, in meters
        out_graph_file: the target file where to save the graph to; usually end in ``.hdf5``
        prune_land_areas: whether to prune by land areas
        bathymetric_dataset: the path to the bathymetric dataset if parameter ``prune_land_areas`` is set to
                             ``True``; e.g. the Earth2014 dataset with depth in meters
        generate_neighbors: whether to generate and serialize all neighbors too
        dump_plots: whether to dump all plots with the default config using
                    :mod:`scripts.visualize_earth_graph` after completing the graph generation
        out_visualization_directory: the target directory (may not yet exist) where to save the visualizations
                                     to if parameter ``dump_plots`` is set to ``True``
    """

    print("Starting generation of earth graph")
    graph = create_earth_graph(min_required_frequency(requested_distance, in_meters=True))

    if prune_land_areas:
        print("Pruning graph")

        # generate the "keep condition" and then remove the property afterwards
        data_set = DataSetAccess(bathymetric_dataset)
        mode = transformers_concrete.BathymetricTransformer.Modes.FRACTION_NAVIGABLE
        graph.append_properties(
            [transformers_concrete.BathymetricTransformer(data_set, [mode])], show_progress=True
        )
        # keep all nodes that have more than 0% (i.e. that have any) navigable locations
        keep_condition = graph.node_properties[mode.column_name] >= 0.0
        graph.clear_node_properties()

        graph.prune_nodes(keep_condition.to_numpy())

    if generate_neighbors:
        print("Generating neighbor table")
        _ = graph.neighbors

    print("Completed generation of earth graph")

    print(f'Serializing to disk: "{out_graph_file}"')
    os.makedirs(os.path.dirname(out_graph_file) or ".", exist_ok=True)
    graph.to_disk(out_graph_file)

    if dump_plots:
        print(f"Dumping visualizations: {out_visualization_directory}")
        dump_2d_plots(graph, out_visualization_directory)


def _main() -> None:
    """The main function."""
    parser = ArgumentParser(
        description="Create and serialize a graph of the earth. "
        "Optionally perform pruning of land area, neighbor discovery and "
        "dumping of visualizations.",
        formatter_class=ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("requested_distance", type=float, help="The max. distance between nodes in meters")
    parser.add_argument("out_graph_file", type=str)
    parser.add_argument("--prune_land_areas", type=bool, default=True)
    parser.add_argument(
        "--bathymetric_dataset",
        type=str,
        default="../data/topography/earth2014/Earth2014.TBI2014.1min.geod.geo.tif",
    )
    parser.add_argument("--generate_neighbors", type=bool, default=True)
    parser.add_argument("--dump_plots", type=bool, default=True)
    parser.add_argument(
        "--visualization_directory",
        type=str,
        default=None,
        help='default: "dirname(out_graph_file)/visualization/"',
    )
    args = parser.parse_args()

    out_visualization_directory = (
        args.visualization_directory
        if args.visualization_directory is not None
        else os.path.join(os.path.dirname(args.out_graph_file), "visualization")
    )

    calculate_and_save(
        requested_distance=args.requested_distance,
        out_graph_file=args.out_graph_file,
        prune_land_areas=args.prune_land_areas,
        bathymetric_dataset=args.bathymetric_dataset,
        generate_neighbors=args.generate_neighbors,
        dump_plots=args.dump_plots,
        out_visualization_directory=out_visualization_directory,
    )


if __name__ == "__main__":
    _main()