from pymatgen.io.abinit.abiobjects import *
#from pymatgen.io.abinit.calculations import *
from pymatgen.io.abinit.events import EventsParser, autodoc_event_handlers
#from pymatgen.io.abinit.qadapters import *
from pymatgen.io.abinit.qadapters import show_qparams, all_qtypes
from pymatgen.io.abinit.netcdf import NetcdfReader
from pymatgen.io.abinit.launcher import PyFlowScheduler, PyLauncher
from pymatgen.io.abinit.pseudos import Pseudo, PseudoTable, PseudoParser
from pymatgen.io.abinit.wrappers import Mrgscr, Mrgddb, Mrggkk, Cut3D
from pymatgen.io.abinit.nodes import Status
from pymatgen.io.abinit.tasks import *
from pymatgen.io.abinit.works import *
from pymatgen.io.abinit.flows import (Flow, G0W0WithQptdmFlow, bandstructure_flow, PhononFlow,
    g0w0_flow, phonon_flow, phonon_conv_flow, nonlinear_coeff_flow)
from pymatgen.io.abinit.abitimer import AbinitTimerParser, AbinitTimerSection
from pymatgen.io.abinit.abiinspect import GroundStateScfCycle, D2DEScfCycle


def flow_main(main):
    """
    This decorator is used to decorate main functions producing `Flows`.
    It adds the initialization of the logger and an argument parser that allows one to select
    the loglevel, the workdir of the flow as well as the YAML file with the parameters of the `TaskManager`.
    The main function shall have the signature:

        main(options)

    where options in the container with the command line options generated by `ArgumentParser`.

    Args:
        main: main function.
    """
    from functools import wraps

    @wraps(main)
    def wrapper(*args, **kwargs):
        parser = build_flow_main_parser()
        options = parser.parse_args()

        # loglevel is bound to the string value obtained from the command line argument.
        # Convert to upper case to allow the user to specify --loglevel=DEBUG or --loglevel=debug
        import logging
        numeric_level = getattr(logging, options.loglevel.upper(), None)
        if not isinstance(numeric_level, int):
            raise ValueError('Invalid log level: %s' % options.loglevel)
        logging.basicConfig(level=numeric_level)

        # Istantiate the manager.
        options.manager = TaskManager.as_manager(options.manager)

        def execute():
            """This is the function that performs the work depending on options."""
            flow = main(options)

            if options.scheduler:
                flow.rmtree()
                return flow.make_scheduler().start()

            elif options.batch:
                flow.rmtree()
                flow.build_and_pickle_dump()
                return flow.batch()

            return 0

        if options.prof:
            # Profile execute
            import pstats, cProfile
            cProfile.runctx("execute()", globals(), locals(), "Profile.prof")
            s = pstats.Stats("Profile.prof")
            s.strip_dirs().sort_stats("time").print_stats()
            return 0
        else:
            return execute()

    return wrapper

def build_flow_main_parser():
    """Build the parser used in the abipy/data/runs scripts."""
    import argparse
    parser = argparse.ArgumentParser()

    parser.add_argument('--loglevel', default="ERROR", type=str,
                        help="set the loglevel. Possible values: CRITICAL, ERROR (default), WARNING, INFO, DEBUG")

    parser.add_argument("-w", '--workdir', default="", type=str, help="Working directory of the flow.")

    parser.add_argument("-m", '--manager', default=None,
                        help="YAML file with the parameters of the task manager. "
                             "Default None i.e. the manager is read from standard locations: "
                             "working directory first then ~/.abinit/abipy/manager.yml.")

    parser.add_argument("-s", '--scheduler', action="store_true", default=False,
                        help="Run the flow with the scheduler")

    parser.add_argument("-b", '--batch', action="store_true", default=False,
                        help="Run the flow in batch mode")

    parser.add_argument("-r", "--remove", default=False, action="store_true", help="Remove old flow workdir")

    parser.add_argument("--prof", action="store_true", default=False, help="Profile code wth cProfile ")
    return parser
