# Copyright (c) 2023-2024 Simons Observatory.
# Full license can be found in the top level "LICENSE" file.

import argparse
from functools import wraps

import toast
import toast.mpi
import toast.traits


def workflow_timer(f):
    """Decorator for timing workflow functions.

    This requires that the workflow function conform to the standard
    function signature.

    """
    fname = f.__qualname__

    @wraps(f)
    def df(job, otherargs, runargs, *args):
        log = toast.utils.Logger.get()
        timer = toast.timing.Timer()
        timer.start()

        data = None
        comm = None
        for arg in args:
            if isinstance(arg, toast.Data):
                data = arg
                comm = data.comm.comm_world
            elif isinstance(arg, toast.mpi.MPI_Comm):
                comm = arg

        log.info_rank(f"Running {fname}...", comm=comm)

        result = f(job, otherargs, runargs, *args)

        log.info_rank(f"Finished {fname} in", comm=comm, timer=timer)
        if data is not None:
            job.operators.mem_count.prefix = f"After {fname}"
            job.operators.mem_count.apply(data)
        return result

    return df


def setup_job(
    parser=None,
    operators=list(),
    templates=list(),
    opts=None,
    config_files=None,
    argparse_opts=None,
    verbose=False,
):
    """Generate configuration for a workflow and parse options.

    Given a list of operators and templates to be used in a workflow, build the
    default configuration.  Then parse any option overrides from the commandline
    OR from the combination of argparse option dictionary and config files.  If
    an existing ArgumentParser is given, then it is used.  Otherwise a new empty
    parser is created.

    This function is designed to work from a script running in a batch job as well
    as interactively from a notebook.  The final configuration can optionally be
    written out for reference.

    For interactive use or ad-hoc scripts, the `opts` dictionary can be passed with
    options matching typical commandline syntax without the leading dashes.  A list
    of config files to load can also be supplied with `config_files`.

    For non-interactive scripts that need to parse commandline arguments, leave
    `opts` and `config_files` at None.  Or the commandline can be replaced with a
    list of argparse-style options in `argparse_opts`.

    Args:
        parser (ArgumentParser):  The input parser to modify.  If None, a new one
            is created.
        operators (list):  The list of operator instances to configure.
        templates (list):  The list of template instances to configure.
        opts (dict):  The dictionary of interactive config options.
        config_files (list):  A list of configuration files to load.
        argparse_opts (list):  When running in batch mode, an argparse list of
            command line arguments to use instead of parsing sys.argv.
        verbose (bool):  If True, dump out the final config to stdout.

    Returns:
        (tuple):  The (job namespace, config dictionary, other args, runtime args)
            from the result of parsing opts + config_files or command line and then
            configuring all operators and templates.

    """
    # We set up the utility memory tracking operator here, so it is always available
    # to any of the workflows.
    operators.append(toast.ops.MemoryCounter(name="mem_count", enabled=False))

    # Check for consistent arguments
    interactive = opts is not None or config_files is not None
    if interactive and argparse_opts is not None:
        msg = "Either opts/config_files should be used OR argparse_opts"
        raise ValueError(msg)

    if interactive:
        # Build the equivalent list of argparse options
        argparse_opts = list()
        if opts is not None:
            for k, v in opts.items():
                optkey = f"--{k}"
                if isinstance(v, bool):
                    # This is just a switch
                    if v:
                        argparse_opts.extend([optkey])
                    else:
                        # Invert
                        kparts = k.split(".")
                        kparts[-1] = f"no_{kparts[-1]}"
                        argparse_opts.extend([f"--{'.'.join(kparts)}"])
                else:
                    argparse_opts.extend([optkey, str(v)])
        if config_files is not None:
            argparse_opts.append("--config")
            argparse_opts.extend(config_files)

    if parser is None:
        # Create a parser
        parser = argparse.ArgumentParser(description="Simons Obs TOAST Workflow")

    parser.add_argument(
        "--full_pointing",
        required=False,
        default=False,
        action="store_true",
        help="Save detector pointing rather than regenerating to save memory.",
    )
    parser.add_argument(
        "--dry_run",
        required=False,
        default=False,
        action="store_true",
        help="Exit after setting up job options.",
    )

    # The "config" is the internal representation of all config parameters
    # merged from defaults, config files, and command line arguments.  Otherargs
    # are anything else parsed from opts / CLI.  Runargs contains runtime parameters
    # like procs per node, groupsize, etc.
    config, otherargs, runargs = toast.config.parse_config(
        parser,
        operators=operators,
        templates=templates,
        prefix="",
        opts=argparse_opts,
    )

    # Instantiate operators and templates
    job = toast.traits.create_from_config(config)

    if verbose:
        out = ""
        for objtype, objlist in zip(
            ["Operators", "Templates"],
            [job.operators, job.templates],
        ):
            out += f"{objtype}:\n"
            for objname, obj in vars(objlist).items():
                out += f"  {objname}:\n"
                for trait_name, trait in obj.traits().items():
                    if trait_name == "API":
                        continue
                    if trait_name == "kernel_implementation":
                        continue
                    if trait_name == "name":
                        continue
                    out += f"    {trait_name} = {trait.get(obj)} # {trait.help}\n"
        print(out)

    return job, config, otherargs, runargs


def reduction_group_size(job, runargs, comm):
    log = toast.utils.Logger.get()
    if runargs.group_size is not None:
        msg = f"Using user-specifed process group size of {runargs.group_size}"
        log.info_rank(msg, comm=comm)
        group_size = runargs.group_size
    else:
        if job.operators.mlmapmaker.enabled:
            msg = f"ML mapmaker is enabled, forcing process group size to 1"
            log.info_rank(msg, comm=comm)
            group_size = 1
        else:
            msg = f"Using default process group size"
            log.info_rank(msg, comm=comm)
            if comm is None:
                group_size = 1
            else:
                group_size = comm.size
    return group_size
