#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import logging
import argparse
import pathlib

import numpy as np

# global settings
from . import config
from .core.register import registers, import_all_modules_for_register
from .utils.command import parse_input_file, dict2str


def main():
    # - register
    import_all_modules_for_register()

    description = "gdpx: Generating Deep Potential with Python\n"

    # - arguments
    parser = argparse.ArgumentParser(prog="gdp", description=description)

    parser.add_argument(
        "-rs", "--random_seed", default=None, type=int, help="global random seed"
    )

    parser.add_argument(
        "-d", "--directory", default=pathlib.Path.cwd(), help="working directory"
    )

    # the workflow tracker
    parser.add_argument(
        "-p",
        "--potential",
        default=None,
        help="target potential related configuration (json/yaml)",
    )

    parser.add_argument(
        "-nj", "--n_jobs", default=1, type=int, help="number of processors"
    )

    parser.add_argument(
        "--debug", action="store_true", help="debug mode that gives more information"
    )

    parser.add_argument("--log", default="gdp.out", help="logging output file")

    # subcommands in the entire workflow
    subparsers = parser.add_subparsers(
        title="available subcommands", dest="subcommand", help="sub-command help"
    )

    # - run session
    parser_session = subparsers.add_parser(
        "session",
        help="run gdpy session",
        description=str(registers.variable) + "\n" + str(registers.operation),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser_session.add_argument(
        "SESSION", help="session configuration file (json/yaml)"
    )
    parser_session.add_argument(
        "--feed", default=None, nargs="+", help="session placeholders"
    )
    parser_session.add_argument(
        "--timewait",
        default=-1,
        type=float,
        help="waiting time between repeated running",
    )

    # - build structures
    parser_build = subparsers.add_parser(
        "build",
        help="build structures",
        description=str(registers.builder),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser_build.add_argument("CONFIG", help="builder configuration file (json/yaml)")
    parser_build.add_argument(
        "-s",
        "--substrates",
        default=None,
        help="file that stores substrates (e.g. *.xyz)",
    )
    parser_build.add_argument(
        "-n", "--number", default=1, type=int, help="number of structures to build"
    )

    # - convert dataset format
    parser_convert = subparsers.add_parser(
        "convert",
        help="convert dataset formats",
    )
    parser_convert.add_argument("INPUT", help="path of the input dataset")

    # - automatic training
    parser_train = subparsers.add_parser(
        "train",
        help="automatic training utilities",
        description=str(registers.trainer) + "\n" + str(registers.dataloader),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser_train.add_argument("CONFIG", help="training configuration file (json/yaml)")

    # --- compute interface
    parser_compute = subparsers.add_parser(
        "compute",
        help="compute structures with basic methods (MD, MIN, and ...)",
        description=str(registers.manager),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser_compute.add_argument(
        "STRUCTURE",
        nargs="*",
        help="a structure file that stores one or more structures",
    )
    parser_compute.add_argument(
        "-b",
        "--batch",
        default=None,
        type=int,
        help="run selected batch number (useful when queue run)",
    )
    parser_compute.add_argument(
        "--spawn",
        action="store_true",
        help="If the computation is spawned, it will not save results until all jobs are finished.",
    )
    parser_compute.add_argument(
        "--archive",
        action="store_true",
        help="whether archive computation folders when retrieve",
    )

    # --- expedition interface
    parser_explore = subparsers.add_parser(
        "explore",
        help="explore structures with advanced methods (GA, MC, and ...)",
        description=str(registers.expedition),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser_explore.add_argument(
        "CONFIG", help="json/yaml file that stores parameters for a task"
    )
    parser_explore.add_argument(
        "--wait", default=None, type=float, help="wait time after each run"
    )

    # selection
    parser_select = subparsers.add_parser(
        "select",
        help="apply various selection operations",
        description=str(registers.selector),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser_select.add_argument("CONFIG", help="selection configuration file")
    parser_select.add_argument(
        "-s", "--structure", required=True, help="structure generator"
    )

    # --- validation
    parser_validation = subparsers.add_parser(
        "valid",
        help="validate properties with trained models",
        description=str(registers.validator),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser_validation.add_argument("CONFIG", help="validation configuration file")

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

    # - update global configuration
    if args.debug:
        config.logger.setLevel(logging.DEBUG)

    curr_wdir = pathlib.Path(args.directory)
    if not curr_wdir.exists():
        curr_wdir.mkdir(parents=True)

    if args.log:
        logfpath = curr_wdir / args.log
        if logfpath.exists():
            fh = logging.FileHandler(logfpath, mode="a")
        else:
            fh = logging.FileHandler(logfpath, mode="w")
        fh.setFormatter(config.formatter)
        config.logger.addHandler(fh)

    # -- set LOGO
    for line in config.LOGO_LINES:
        config._print(line)

    # -- set njobs
    config.NJOBS = args.n_jobs
    if config.NJOBS != 1:
        config._print(f"Use {config.NJOBS} processors.")

    # -- set rng
    # TODO: load random state from a file???
    random_seed = args.random_seed
    if random_seed is None:
        # NOTE: np.random should only be called here once...
        random_seed = np.random.randint(0, 1e8)
    else:
        ...

    config.GRNG = np.random.Generator(np.random.PCG64(random_seed))

    config._print(f"GLOBAL RANDOM SEED : {random_seed}")
    rng_state = config.GRNG.bit_generator.state
    for l in dict2str(rng_state).split("\n"):
        config._print(l)

    # - potential
    if args.potential:
        # a worker or a List of worker
        from .cli.compute import convert_config_to_potter

        potter = convert_config_to_potter(args.potential)
    else:
        potter = [None]

    # - use subcommands
    if args.subcommand == "session":
        from gdpx.core.session import run_session

        run_session(args.SESSION, args.feed, args.timewait, args.directory)
    elif args.subcommand == "convert":
        from gdpx.data import convert_dataset

        convert_dataset(args.INPUT)
    elif args.subcommand == "train":
        from gdpx.trainer import run_newtrainer

        run_newtrainer(args.CONFIG, args.directory)
    elif args.subcommand == "build":
        build_config = parse_input_file(args.CONFIG)
        from .cli.build import build_structures

        build_structures(build_config, args.substrates, args.number, args.directory)
    elif args.subcommand == "select":
        from gdpx.selector.interface import run_selection

        run_selection(args.CONFIG, args.structure, args.directory)
    elif args.subcommand == "compute":
        if isinstance(potter, list):
            first_worker = potter[0]
        else:
            # FIXME: This is for reactor but we need unify this with computer.
            first_worker = potter
            potter = [potter]
        config._print(f"{first_worker =}")
        if first_worker is not None:
            # For compatibility, the classic mode
            # `gdp -p ./worker.yaml compute structures.xyz`
            from .cli.compute import run_worker

            run_worker(
                args.STRUCTURE,
                potter,
                batch=args.batch,
                spawn=args.spawn,
                archive=args.archive,
                directory=args.directory,
            )
        else:  # Use GridWorker here!!
            from .cli.compute import run_grid_worker

            run_grid_worker(
                parse_input_file(args.STRUCTURE[0]),
                batch=args.batch,
                spawn=args.spawn,
                directory=args.directory,
            )
    elif args.subcommand == "explore":
        from .cli.explore import run_expedition

        params = parse_input_file(args.CONFIG)
        run_expedition(params, args.wait, args.directory, potter[0])
    elif args.subcommand == "valid":
        from gdpx.validator import run_validation

        params = parse_input_file(args.CONFIG)
        run_validation(params, args.directory, potter[0])
    else:
        ...

    # - report the end random state
    config._print(f"GLOBAL RANDOM SEED : {random_seed}")
    rng_state = config.GRNG.bit_generator.state
    for l in dict2str(rng_state).split("\n"):
        config._print(l)

    return


if __name__ == "__main__":
    main()
