""" Python script for reducing ASCII data generated by the UniverseMachine.
The desired subvolumes and columns are all specied at the command line.
The script then iterates over each subvolume and creates a memory-mapped array
for each requested column, storing the results in a directory tree structure
that has been standardized to simplify parallel I/O.
"""
import os
import argparse
from time import time
from umachine_pyio.process_ascii_into_memmap import write_ascii_to_memmap_tree
from umachine_pyio.directory_tree_utils import sf_history_ascii_fname_iterator
from umachine_pyio.sf_history_header_processing import retrieve_requested_colnames

################################################################################
if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("input_dirname", help="Directory of the ascii history files")
    parser.add_argument(
        "column_info_fname",
        help="User-created ASCII file used to interpret the " "UniverseMachine outputs",
    )
    parser.add_argument(
        "output_dirname", help="Directory to store the Numpy binary files"
    )
    parser.add_argument("scale_factor", help="Scale factor of the history files.")

    parser.add_argument(
        "-fname_prefix_pattern",
        default="sfh_catalog_",
        help="Beginning of substring used to identify the ascii files in the input drn."
        "Default is ``sfh_catalog_``",
    )
    parser.add_argument(
        "-fname_suffix_pattern",
        default=".txt",
        help="Concluding substring used to identify ascii files in the input drn."
        "Default is ``.txt``",
    )
    parser.add_argument(
        "-subvolume_labels",
        type=int,
        nargs="+",
        default="-1",
        help="Sequence of integers of subvolumes to process."
        "Default is to process all.",
    )
    parser.add_argument(
        "-galaxy_colnames",
        type=str,
        nargs="+",
        default=["all"],
        help="Sequence of strings of galaxy properties to reduce to binaries. "
        "Each string must appear in the first column of ``column_info_fname``. "
        "Default behavior is to process all columns.",
    )

    args = parser.parse_args()
    ################################################################################

    print(
        "\n...creating binary reductions for scale factor = {0}".format(
            args.scale_factor
        )
    )
    # Iterate over the requested filenames and memory-map each requested column

    requested_colnames = retrieve_requested_colnames(
        args.galaxy_colnames, args.column_info_fname
    )

    fname_iter = sf_history_ascii_fname_iterator(
        args.subvolume_labels,
        args.input_dirname,
        args.fname_prefix_pattern + args.scale_factor,
        args.fname_suffix_pattern,
    )

    scale_factor_subdrname = "a_" + args.scale_factor
    output_dirname = os.path.join(args.output_dirname, scale_factor_subdrname)
    os.makedirs(output_dirname, exist_ok=True)
    start = time()
    print("...beginning loop over files")
    for subvol_index, ascii_fname in fname_iter:
        output_subdir = "subvol_" + str(subvol_index)
        subvol_output_dirname = os.path.join(output_dirname, output_subdir)

        start1 = time()
        write_ascii_to_memmap_tree(
            ascii_fname,
            args.column_info_fname,
            subvol_output_dirname,
            requested_colnames,
        )
        end1 = time()
        runtime1 = end1 - start1
        msg = "Runtime to reduce {0} = {1:.1f} seconds".format(output_subdir, runtime1)
        print(msg)

    end = time()
    print("Total runtime = {0:.2f} seconds\n".format((end - start)))
