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

"""Automate generation of parametrized launchers."""

__author__ = "Lluís Vilanova"
__copyright__ = "Copyright 2008-2024, Lluís Vilanova"
__license__ = "GPL version 3 or later"


import collections
import contextlib
import functools
import itertools
import operator
import os
import shutil
import six
import warnings

from sciexp2.common import parallel as parallel_pkg
from sciexp2.common import progress
from sciexp2.common import utils
from sciexp2.common.instance import InstanceGroup, Instance
from sciexp2.common.filter import *
from sciexp2.common.utils import execute_with_sigint
from sciexp2.common import pp
from sciexp2.common import text
from . import templates
from . import launcher


class _sorted_dict(object):
    """Helper class to print dicts with sorted keys."""
    def __init__(self, object):
        self._object = object
    def __repr__(self):
        keys = sorted(self._object.keys())
        return "{%s}" % ", ".join("%r: %r" % (key, self._object[key])
                                  for key in keys)

######################################################################
# Experiments

def _is_dummy(experiments):
    if isinstance(experiments, Experiments):
        return (len(experiments._experiments) == 1 and
                len(experiments._experiments.get_index(0)) == 0)
    else:
        return len(experiments) == 1 and len(experiments[0]) == 0


def _val_is_translate(val):
    return (isinstance(val, six.string_types) and
            len(text.get_variables(val)) > 0) or \
            isinstance(val, _WithExp) or \
            isinstance(val, _WithExpTpl)


def _get_translate_exp(exp):
    res = utils.OrderedSet()
    for key, vals in exp.items():
        if isinstance(vals, six.string_types):
            if _val_is_translate(vals):
                res.add(key)
        elif isinstance(vals, collections.abc.Iterable):
            if any(_val_is_translate(val) for val in vals):
                res.add(key)
    return res

def _get_translate(experiments):
    translate = utils.OrderedSet()
    for exp in experiments:
        translate |= _get_translate_exp(exp)
    return translate


def _do_execute(experiments, cmd, **kwargs):
    if len(experiments) == 0:
        experiments_list = [Instance()]
    else:
        experiments_list = experiments

    if isinstance(cmd, six.string_types):
        is_shell = True
        cmd = [cmd]
    else:
        is_shell = False

    parallel = kwargs.pop("parallel", None)
    repeat = kwargs.pop("repeat", "warn")
    if repeat not in ["ignore", "warn", "error"]:
        raise ValueError("invalid repeat value: %r" % repeat)
    std_args = {
        'stdin':  kwargs.pop("stdin",  None),
        'stdout': kwargs.pop("stdout", None),
        'stderr': kwargs.pop("stderr", None),
    }
    utils.assert_kwargs(kwargs)

    def std_translator(name):
        arg = std_args[name]
        if isinstance(arg, six.string_types):
            return text.Translator(arg)
        else:
            return None
    std_translator = {
        'stdin':  std_translator('stdin'),
        'stdout': std_translator('stdout'),
        'stderr': std_translator('stderr'),
    }


    # flatten command to get faster translation
    cmd_arg_sep = "|;argsep;|"
    cmd_flat = cmd_arg_sep.join(cmd)
    def get_cmd(cmd):
        return cmd.split(cmd_arg_sep)

    translator = text.Translator(cmd_flat)
    exp_cmds = translator.translate_many(experiments_list, with_envs=True)

    def execute(exp_cmd):
        cmd = get_cmd(exp_cmd[0])
        exp = exp_cmd[1][0]
        exp_rest = exp_cmd[1][1:]

        progress.verbose(" ".join(cmd))
        if len(exp_rest) > 0:
            if repeat == "warn":
                warnings.warn(
                    "Repeated command translation for experiments %r and %r" % (
                        exp, exp_rest[0]))
            elif repeat == "error":
                raise RuntimeError(
                    "repeated command translation for experiments %r and %r" % (
                        exp, exp_rest[0]))


        def get_stdfile(name, env):
            @contextlib.contextmanager
            def as_context(arg):
                yield arg
            translator = std_translator[name]
            if translator is None:
                return as_context(std_args[name])
            else:
                with env.auto_translate():
                    path = translator.translate(env)
                return open(path, "w")

        with get_stdfile("stdin",  exp) as stdin,\
             get_stdfile("stdout", exp) as stdout,\
             get_stdfile("stderr", exp) as stderr:
            res = execute_with_sigint(
                *cmd,
                shell=is_shell,
                stdin=stdin,
                stdout=stdout,
                stderr=stderr)
        if res != 0:
            raise Exception("Command did not finish correctly: "
                            "%s -> %d" % (" ".join(cmd), res))

        return len(exp_cmd[1])

    # use progr() explicitly to ensure progress is made after each command
    with progress.get(experiments_list, msg="Executing commands...") as progr:
        if parallel is False:
            for exp in exp_cmds:
                exp_cmd = exp[0]
                exps = exp[1]
                count = execute(exp)
                progr(count)
        else:
            if parallel is True:
                parallel = None
            for count in parallel_pkg.t_imap_unordered(
                    exp_cmds, execute, parallelism=parallel):
                progr(count)


def _do_recombine(experiments, exp_filters, variables, var_filters, append):
    exp_filters = and_filters(*exp_filters)
    var_filters = and_filters(*var_filters)

    # normalize variable value types
    for key, val in variables.items():
        if not isinstance(val, collections.abc.Iterable) or \
           isinstance(val, six.string_types):
            variables[key] = [val]
        elif isinstance(val, collections.abc.Sized):
            variables[key] = list(val)

    # calculate various properties about variable values
    var_translate = _get_translate_exp(variables)
    var_from_func = [var for var, vals in variables.items()
                     if any(isinstance(val, _FromFunctionKey) for val in vals)]
    var_size = functools.reduce(operator.mul,
                                (len(vals) for vals in variables.values()),
                                1)

    # validate variable filters
    var_filters.validate((set(experiments.variables()) | set(variables.keys())) - var_translate)

    new_experiments = InstanceGroup()
    new_dups = set()

    total_size = len(experiments) * var_size
    if append:
        total_size += var_size

    def cartesian_product(base, variables_missing=None):
        """Yield cartesian product of `variables` and overlay each on `base`"""
        if variables_missing is None:
            variables_missing = sorted(variables.keys())

        if len(variables_missing) == 0:
            yield {}

        else:
            key = variables_missing[0]
            variables_missing = variables_missing[1:]

            for val in variables[key]:
                for exp in cartesian_product(base, variables_missing):
                    new_exp = Instance(base)
                    new_exp.update(exp)
                    new_exp[key] = val
                    yield new_exp

    def cartesian_product_ff(base, ff_info, ffs_missing=None):
        """Yield cartesian product of `ff_info` and overlay each on `base`"""
        if ffs_missing is None:
            ffs_missing = sorted(ff_info.keys())

        if len(ffs_missing) == 0:
            yield {}

        else:
            ff = ffs_missing[0]
            ffs_missing = ffs_missing[1:]

            for val in ff_info[ff][0]:
                for exp in cartesian_product_ff(base, ff_info, ffs_missing):
                    new_exp = Instance(base)
                    new_exp.update(exp)
                    for var in ff_info[ff][1]:
                        new_exp[var[0]] = val[var[1]]
                    yield new_exp

    def merge_old_experiment(exp, force=False):
        exp_dups = tuple(exp.items())
        if exp_dups not in new_dups:
            new_experiments.add(exp)
            new_dups.add(exp_dups)
        elif force:
            new_experiments.add(exp)

    def merge_new_experiment_noff(progr, exp):
        if var_filters.match(exp):
            merge_old_experiment(exp)
        progr()

    def merge_new_experiment(progr, exp):
        ff_vars = dict(((var, exp[var]) for var in var_from_func
                        if isinstance(exp[var], _FromFunctionKey)))

        if len(ff_vars) == 0:
            merge_new_experiment_noff(progr, exp)

        else:
            exp = dict((key, val) for key, val in exp.items()
                       if key not in ff_vars)

            ff_info = {}
            progr_size_add = 1
            for var, val in ff_vars.items():
                val = val._from_func
                if val not in ff_info:
                    exp_ff = dict(exp)
                    vals = val._get_results(exp)
                    if len(vals) == 0:
                        # experiment dropped, but still increase progress
                        progr()
                        return
                    ff_info[val] = (vals, [])
                    progr_size_add *= len(vals)
                ff_info[val][1].append((var, ff_vars[var]._key_name))

            assert progr_size_add > 0
            progr += progr_size_add

            for exp_prod in cartesian_product_ff(exp, ff_info):
                merge_new_experiment_noff(progr, exp_prod)


    with progress.get(total_size,
                      msg="recombining variables...") as progr:

        if append:
            for exp in experiments:
                merge_old_experiment(exp, force=True)
                progr(var_size)

            for exp_prod in cartesian_product({}):
                merge_new_experiment(progr, exp_prod)
        else:
            for exp_base in experiments:
                try:
                    is_match = exp_filters.match(exp_base)
                except NameError:
                    is_match = False

                if is_match:
                    for exp_prod in cartesian_product(exp_base):
                        merge_new_experiment(progr, exp_prod)
                else:
                    merge_old_experiment(exp_base)
                    progr(var_size)

    return new_experiments, var_translate


def _do_pack(experiments, template_from, template_to, dereference):
    if len(experiments) == 0:
        experiments_list = [Instance()]
    else:
        experiments_list = experiments
    if dereference is None:
        dereference = experiments.dereference

    translate_from = text.Translator(template_from)
    translate_to = text.Translator(template_to)

    done_first = {}
    done_exp = {}
    with progress.get(len(experiments_list), msg="packing files...") as progr:
        for exp in experiments_list:
            with exp.auto_translate():
                path_from = utils.get_path(translate_from.translate(exp))
                path_to = utils.get_path(translate_to.translate(exp))
                path_to = os.path.join(experiments.out, path_to)

                done_idx = path_from + "+" + path_to
                if done_idx in done_first:
                    if done_first[done_idx]:
                        done_first[done_idx] = False
                        progress.info(
                            "Skipping already packed file: %s -> %s",
                            path_from, path_to)
                    progress.verbose(
                        "Conflicting experiments during pack:\n%s\n%s",
                        done_exp[done_idx], exp)
                else:
                    done_first[done_idx] = True
                    done_exp[done_idx] = exp

                    progress.debug("Packing: %s -> %s", path_from, path_to)
                    utils.copy_path(path_from, path_to)

                progr()


def _do_translate(experiments, exp_filters, exp_translate,
                  templates, with_exps, with_unique):
    exp_filters = and_filters(*exp_filters)
    parsed_templates = [text.Translator(template)
                        for template in templates]

    if with_exps:
        if with_unique:
            translations = collections.OrderedDict()
            if len(templates) == 1:
                def add(translated, exp):
                    translated = translated[0]
                    if translated not in translations:
                        translations[translated] = []
                    translations[translated].append(dict(exp))
            else:
                def add(translated, exp):
                    if translated not in translations:
                        translations[translated] = []
                    translations[translated].append(dict(exp))
            def result():
                return translations.items()
        else:
            translations = []
            if len(templates) == 1:
                def add(translated, exp):
                    translations.append((translated[0], [dict(exp)]))
            else:
                def add(translated, exp):
                    translations.append((translated, [dict(exp)]))
            def result():
                return translations
    else:
        if with_unique:
            translations = utils.OrderedSet()
            if len(templates) == 1:
                def add(translated, exp):
                    translations.add(translated[0])
            else:
                def add(translated, exp):
                    translations.add(tuple(translated))
            def result():
                return translations
        else:
            translations = []
            if len(templates) == 1:
                def add(translated, exp):
                    translations.append(translated[0])
            else:
                def add(translated, exp):
                    translations.append(tuple(translated))
            def result():
                return translations

    for exp in experiments.select(exp_filters):
        with exp.auto_translate():
            try:
                is_match = exp_filters.match(exp)
            except NameError:
                is_match = False

            if is_match:
                _WithExp._EXPERIMENT = exp
                _WithExpTpl._EXPERIMENT = exp
                translated = [template.translate(exp)
                              for template in parsed_templates]
                add(translated, exp)

    return list(result())


def _do_generate(experiments, template_from, template_to,
                 pre_exp_fun, post_exp_fun):
    translate_from = text.Translator(template_from)
    translate_to = text.Translator(template_to)

    done_first = {}
    done_exp = {}
    with progress.get(len(experiments), msg="generating files...") as progr:
        for exp in experiments:
            with exp.auto_translate():
                _WithExp._EXPERIMENT = exp
                _WithExpTpl._EXPERIMENT = exp

                path_from = utils.get_path(translate_from.translate(exp))
                path_to = utils.get_path(translate_to.translate(exp))
                path_to = os.path.join(experiments.out, path_to)

                done_idx = path_from + "+" + path_to
                if done_idx in done_first:
                    if done_first[done_idx]:
                        done_first[done_idx] = False
                        progress.info(
                            "Skipping already generated file: %s -> %s",
                            path_from, path_to)
                    progress.verbose(
                        "Conflicting experiments during generation:\n%s\n%s",
                        done_exp[done_idx], exp)
                else:
                    done_first[done_idx] = True
                    done_exp[done_idx] = exp

                    progress.debug("Generating: %s -> %s", path_from, path_to)

                    exp, contents_to = pre_exp_fun(exp, path_from, path_to)

                    utils.assert_dir(os.path.dirname(path_to))
                    if os.path.isfile(path_to):
                        with open(path_to, "r") as f:
                            contents_to_old = f.read()
                    else:
                        contents_to_old = ""

                    if contents_to != contents_to_old:
                        with open(path_to, "w") as f:
                            f.write(contents_to)
                    shutil.copymode(path_from, path_to)

                    post_exp_fun(exp, path_to)

                progr()

def _do_generate_user(experiments, template_from, template_to):
    def pre_exp_fun(exp, path_from, path_to):
        with open(path_from, "r") as f:
            contents_to = f.read()
        return exp, text.translate(contents_to, exp)
    def post_exp_fun(exp, path_to):
        pass
    _do_generate(experiments, template_from, template_to,
                 pre_exp_fun, post_exp_fun)

def _do_generate_jobs(experiments, system, template_to, export, depends,
                      submit_args, job_desc):
    template = templates.get(system)

    depends = set(depends) | set(["{{LAUNCHER}}"])

    # export all variables in depends
    export = set(export)
    for elem in depends:
        if not isinstance(elem, six.string_types):
            raise TypeError("depends element is not a string: %r" % elem)
        for elem_var in text.get_variables(elem, nested="all"):
            if elem_var == "LAUNCHER":
                continue
            export.add(elem_var)

    if isinstance(experiments, ExperimentsView):
        experiments_vars = list(experiments._base._experiments.variables())
    else:
        experiments_vars = list(experiments._experiments.variables())
    for export_elem in export:
        if not isinstance(export_elem, six.string_types):
            raise TypeError("export element is not a string: %r" % export_elem)
        if export_elem not in experiments_vars:
            raise ValueError("missing variable to export: %s" % export_elem)
    export = export | set(template.system.assumes())

    out = os.path.abspath(utils.get_path(experiments.out))
    job_desc = os.path.join(out, job_desc)
    job_desc = os.path.abspath(utils.get_path(job_desc))
    if not job_desc.startswith(out):
        raise ValueError("job descriptor file must be inside output directory")
    job_desc_base = os.path.dirname(job_desc[len(out)+1:])

    # load template file and apply overrides
    path_from = template.template_path
    with open(path_from, "r") as f:
        template_from = f.read()
    overrides = dict(template.overrides)
    for var in text.get_variables(template_from):
        if var not in overrides:
            overrides[var] = "{{%s}}" % var
    template_from = text.translate(template_from, overrides,
                                   recursive=False)
    template_from_vars = text.get_variables(template_from)
    translate_from = text.Translator(template_from)

    def pre_exp_fun(exp, path_from, path_to):
        exp = dict(exp)

        if "LAUNCHER" in exp:
            raise ValueError("experiment has reserved variable 'LAUNCHER'")
        exp["LAUNCHER"] = path_to[len(out)+1:]

        for key, val in template.defaults.items():
            if key not in exp:
                exp[key] = val

        contents_to = translate_from.translate(exp)

        return exp, contents_to

    experiments_generated = InstanceGroup()

    def post_exp_fun(exp, path_to):
        for var in export:
            if var not in exp:
                continue
            val = exp[var]
            if not _val_is_translate(val):
                continue
            exp[var] = text.translate(val, exp)
        experiments_generated.add(Instance(exp))

    _do_generate(experiments, template.template_path, template_to,
                 pre_exp_fun, post_exp_fun)

    launcher.save(job_desc, job_desc_base,
                  template.system, experiments_generated,
                  export, depends, submit_args)


class Experiments(pp.Pretty):
    """Define and generate experiment sets.

    Each experiment corresponds to a `dict` that maps variable names to their
    values, and can be accessed as if `Experiments` was a list.

    Parameters
    ----------
    experiments : sequence of dict, optional
        Initial experiment set.
    out : optional
        Initial value for the `out` attribute.
    dereference : bool, optional
        Initial value for the `dereference` attribute.

    Attributes
    ----------
    out : str
        Output directory for all non-absolute paths.
    dereference : bool
        Whether to dereference source symlinks during `pack`.

    See also
    --------
    with_exp, with_exp_tpl
        Use functions as variable values.
    from_function, from_find_files
        Use functions to create new variable values.
    sciexp2.common.text.translate
        Used to interpret template arguments in the `Experiments` methods, used
        with the variable maps of each experiment.
    sciexp2.common.text.extract
        Used to interpret template arguments in the `Experiments` methods, used
        to build the variable maps of new experiments.
    sciexp2.common.filter.Filter
        Values accepted for the filter arguments of the methods in
        `Experiments`. A list of strings will be used as an argument to
        `sciexp2.common.filter.and_filters`.

    """

    #: Output directory for all non-absolute paths.
    out = "./out"

    #: Whether to dereference source symlinks during pack.
    dereference = False

    def __init__(self, experiments=None, **kwargs):
        out = kwargs.pop("out", None)
        dereference = kwargs.pop("dereference", None)
        utils.assert_kwargs(kwargs)

        if experiments is None:
            experiments = InstanceGroup()
            # add dummy experiment
            experiments += Instance()
        if isinstance(experiments, Experiments) or\
           isinstance(experiments, ExperimentsView):
            if dereference is None:
                dereference = experiments.dereference
            if out is None:
                out = experiments.out
            self._experiments = experiments._experiments
            self._translate = utils.OrderedSet(experiments._translate)
        else:
            self._experiments = InstanceGroup()
            dups = set()
            for exp in experiments:
                exp_dups = tuple(exp.items())
                if exp_dups not in dups:
                    self._experiments.add(Instance(exp))
                    dups.add(exp_dups)
            self._translate = _get_translate(self._experiments)

        if out is not None:
            self.out = out
        if dereference is not None:
            self.dereference = dereference


    def _repr_pretty_(self, p, cycle):
        with self.pformat(p, cycle):
            have_prev = False
            if len(self) > 0:
                p.pretty(list(map(lambda elem: _sorted_dict(elem), self)))
                have_prev = True
            if self.out != Experiments.out:
                if have_prev:
                    p.text(",")
                    p.breakable()
                p.text("out=%r" % self.out)
                have_prev = True
            if self.dereference != Experiments.dereference:
                if have_prev:
                    p.text(",")
                    p.breakable()
                p.text("dereference=%r" % self.dereference)


    def __iter__(self):
        if _is_dummy(self):
            return iter([])
        else:
            return iter(self._experiments)

    def __len__(self):
        if _is_dummy(self):
            return 0
        else:
            return len(self._experiments)

    def __getitem__(self, index):
        if _is_dummy(self):
            return list()[index]
        else:
            return self._experiments.get_index(index)

    def __eq__(self, other):
        return isinstance(other, Experiments) and \
            self.out == other.out and \
            self.dereference == other.dereference and \
            list(self._experiments) == list(other._experiments) and \
            list(self._translate) == list(other._translate)

    def view(self, *filters):
        """Get a view to a portion of this of this experiment set.

        Any changes to the experiments on a view will be reflected on the base
        experiment set and all its views. Adding new experiments to a view will
        maintain the relative order between experiments on the base experiment
        set.

        Parameters
        ----------
        filters : list of filters
            Filters to select a matching subset of the experiments.

        Examples
        --------
        A view shows a subset of the experiments:

        >>> e = Experiments()
        >>> e.params(a=range(2), b=range(2))
        >>> e
        Experiments([{'a': 0, 'b': 0},
                     {'a': 0, 'b': 1},
                     {'a': 1, 'b': 0},
                     {'a': 1, 'b': 1}])
        >>> v = e.view("a == b")
        >>> v
        ExperimentsView([{'a': 0, 'b': 0}, {'a': 1, 'b': 1}])

        All operations are performed only on selected elements, but apply to the
        underlying experiment set as if they happened "in-place":

        >>> v = e.view("a == b")
        >>> v.params(c=range(2))
        >>> v
        ExperimentsView([{'a': 0, 'b': 0, 'c': 0},
                         {'a': 0, 'b': 0, 'c': 1},
                         {'a': 1, 'b': 1, 'c': 0},
                         {'a': 1, 'b': 1, 'c': 1}])
        >>> e
        Experiments([{'a': 0, 'b': 0, 'c': 0},
                     {'a': 0, 'b': 0, 'c': 1},
                     {'a': 0, 'b': 1},
                     {'a': 1, 'b': 0},
                     {'a': 1, 'b': 1, 'c': 0},
                     {'a': 1, 'b': 1, 'c': 1}])

        For convenience, you can also get the inverse of a view, which is a view
        with the negated filters:

        >>> i = v.view_inverse()
        >>> i
        ExperimentsView([{'a': 0, 'b': 1}, {'a': 1, 'b': 0}])
        >>> i.params(d=range(2))
        >>> e
        Experiments([{'a': 0, 'b': 0, 'c': 0},
                     {'a': 0, 'b': 0, 'c': 1},
                     {'a': 0, 'b': 1, 'd': 0},
                     {'a': 0, 'b': 1, 'd': 1},
                     {'a': 1, 'b': 0, 'd': 0},
                     {'a': 1, 'b': 0, 'd': 1},
                     {'a': 1, 'b': 1, 'c': 0},
                     {'a': 1, 'b': 1, 'c': 1}])

        Finally, you can also create new views from other views:

        >>> n = v.view("c == 0")
        >>> n.params(e=[10, 20])
        >>> n.view_inverse().params(e=[30, 40])
        >>> e
        Experiments([{'a': 0, 'b': 0, 'c': 0, 'e': 10},
                     {'a': 0, 'b': 0, 'c': 0, 'e': 20},
                     {'a': 0, 'b': 0, 'c': 1, 'e': 30},
                     {'a': 0, 'b': 0, 'c': 1, 'e': 40},
                     {'a': 0, 'b': 1, 'd': 0},
                     {'a': 0, 'b': 1, 'd': 1},
                     {'a': 1, 'b': 0, 'd': 0},
                     {'a': 1, 'b': 0, 'd': 1},
                     {'a': 1, 'b': 1, 'c': 0, 'e': 10},
                     {'a': 1, 'b': 1, 'c': 0, 'e': 20},
                     {'a': 1, 'b': 1, 'c': 1, 'e': 30},
                     {'a': 1, 'b': 1, 'c': 1, 'e': 40}])

        """
        return ExperimentsView(self, Filter("True"), Filter(*filters))

    def execute(self, cmd, **kwargs):
        """Execute one program for each experiment on the set.

        If `cmd` is a single string, it's interpreted as a shell command,
        otherwise it's interpreted as a program name and its argument list.

        Arguments `stdin`, `stdout`, `stderr`, and all elements in `cmd` are
        interpreted as templates to translate with each experiment.

        Parameters
        ----------
        cmd : list of str
            Command to execute.
        repeat : str, optional
            What to do when more than one experiment is translated into the same
            command. Can be ``"ignore"``, ``"warn"`` (default), or ``"error"``.
        stdin, stdout, stderr : file or str, optional
            Paths for the standard input/output/error for the command (defaults
            to None).
        parallel : optional
            Whether to execute commands in parallel using separate processes
            (default is `False`).

        Notes
        -----
        Command execution is internally handled by `subproces.Popen` and thus
        the `stdin`, `stdout` and `stderr` arguments can also have any value
        accepted by it.

        See also
        --------
        sciexp2.common.parallel.get_parallelism
            Used to interpret argument `parallel`.

        """
        _do_execute(self, cmd, **kwargs)


    def params(self, *filters, **variables):
        """Add new parameter permutations to experiment set.

        Performs the cartesian product between the current experiment set and
        all permutations of the given variable values (i.e., a parameter can
        have a sequence of values).

        Parameters
        ----------
        filters : list of filters, optional
            Ensures new experiments match the given filters (default is to
            accept all).
        variables : dict
            A dictionary of variable name / values pairs to perform
            permutations.

        Notes
        -----
        Values in `variables` can be templates referencing other variables
        (either existing in the experiment set or provided in `variables`), and
        will thus be translated accordingly.

        Providing a variable in `params` that already exists in the experiment
        set will override its value.

        Examples
        --------
        Can easily create value permutations that match the given filters:
        >>> e = Experiments()
        >>> e.params("a % 2 == 0",
        ...          a=range(4),
        ...          b=["foo", "bar"])
        >>> e
        Experiments([{'a': 0, 'b': 'foo'},
                     {'a': 0, 'b': 'bar'},
                     {'a': 2, 'b': 'foo'},
                     {'a': 2, 'b': 'bar'}])

        New variables can be appended to the experiment set:

        >>> e.params(c=3)
        >>> e
        Experiments([{'a': 0, 'b': 'foo', 'c': 3},
                     {'a': 0, 'b': 'bar', 'c': 3},
                     {'a': 2, 'b': 'foo', 'c': 3},
                     {'a': 2, 'b': 'bar', 'c': 3}])

        Including ones with multiple values:

        >>> e.params(d=[1, 2])
        >>> e
        Experiments([{'a': 0, 'b': 'foo', 'c': 3, 'd': 1},
                     {'a': 0, 'b': 'foo', 'c': 3, 'd': 2},
                     {'a': 0, 'b': 'bar', 'c': 3, 'd': 1},
                     {'a': 0, 'b': 'bar', 'c': 3, 'd': 2},
                     {'a': 2, 'b': 'foo', 'c': 3, 'd': 1},
                     {'a': 2, 'b': 'foo', 'c': 3, 'd': 2},
                     {'a': 2, 'b': 'bar', 'c': 3, 'd': 1},
                     {'a': 2, 'b': 'bar', 'c': 3, 'd': 2}])

        New values override old ones, even if that means eliminating some experiments:

        >>> e.params(b="baz")
        >>> e
        Experiments([{'a': 0, 'b': 'baz', 'c': 3, 'd': 1},
                     {'a': 0, 'b': 'baz', 'c': 3, 'd': 2},
                     {'a': 2, 'b': 'baz', 'c': 3, 'd': 1},
                     {'a': 2, 'b': 'baz', 'c': 3, 'd': 2}])

        """
        new_experiments = _do_recombine(self._experiments, [],
                                        variables, filters,
                                        False)
        self._experiments = new_experiments[0]
        self._translate |= new_experiments[1]


    def params_append(self, *filters, **variables):
        """Same as `params`, but appends the new permutations.

        Examples
        --------
        >>> e = Experiments()
        >>> e.params(a=range(2))
        >>> e.params_append(b=range(2))
        >>> e
        Experiments([{'a': 0},
                     {'a': 1},
                     {'b': 0},
                     {'b': 1}])

        """
        new_experiments = _do_recombine(self._experiments, [],
                                        variables, filters,
                                        True)
        self._experiments = new_experiments[0]
        self._translate |= new_experiments[1]


    def pack(self, template_from, template_to, dereference=None):
        """Copy files, accepting templates for both source and destination paths.

        Parameters
        ----------
        template_from, template_to : str
            Source/destination file templates.
        dereference : bool, optional
            Whether to dereference symlinks in source files (defaults to
            ``self.dereference``).

        Notes
        -----
        The copy is performed only if the source is newer than the destination,
        or if the destination does not exist.

        There can be cases where multiple source files correspond to a single
        destination file (after translating the templates). Such cases are
        logged by a message.

        """
        _do_pack(self, template_from, template_to, dereference)

    def generate(self, template_from, template_to):
        """Generate files from other templated files.

        Like `pack`, the paths to source and destination files are
        templates. Every source file is, in addition, considered a template that
        will be translated with each of the experiments.

        Parameters
        ----------
        template_from, template_to : str
            Templates for Source/destination file paths.

        Notes
        -----
        Like `pack`, if multiple source files correspond to a single destination
        file (after translating the templates), the case is logged by a message
        and only the first generated file is kept.

        """
        return _do_generate_user(self, template_from, template_to)

    def generate_jobs(self, system, template_to, export=[], depends=[],
                      submit_args=[], job_desc="jobs.jd"):
        """Generate job files for given job system.

        Works similar to `generate`, but takes a predefined source file template
        that is identified by `system`. It also creates a job descriptor file
        that can be later used by the `launcher` program.

        Parameters
        ----------
        system : str
            Job system name.
        template_to : str
            Template to path of destination job scripts.
        export : list of str, optional
            List of variable names to export into the job descriptor file.
        depends : list of str, optional
            File path templates for job dependencies.
        submit_args : list of str, optional
            Templates for additional arguments to the job submission command.
        job_desc : str, optional
            Name of the output job descriptor file.

        All job systems implicitly define the *LAUNCHER* variable, which has the
        path of the to the generated script to run that. All job systems also
        define *DONE* and *FAIL* as file paths to identify the state of each
        job.

        The variables selected by `export` will be later available when using
        `sciexp2.expdef.launcher` with this job descriptor. Any variable
        directly or indirectly used in `depends` will be implicitly added to
        `export`. Variable *LAUNCHER* is also implicitly added.

        The templates specified by `depends` will be translated for each job
        info file paths, and each file path will be considered a dependency of
        that job. The path specified by *LAUNCHER* is implicitly added as a
        dependency.

        The state for a given job is decided as follows:

        - Not run: Neither the *DONE* nor the *FAIL* files exist.
        - Failed: The *FAIL* file exists.
        - Outdated: The *DONE* file exists, but is older than any of its
          dependencies.
        - Done: The *DONE* file exists, and is newer than any of its
          dependencies.
        - Running: Some job systems can track when a job is currently executing.


        Notes
        -----

        None of the experiments can define the *LAUNCHER* variable. Some job
        systems also define their own set of reserved variables and default
        values to others (in case your experiments don't specify them); see the
        job system's documentation for more information.

        """
        _do_generate_jobs(self, system, template_to, export, depends,
                          submit_args, job_desc)


    def translate(self, template, with_exps=False, with_unique=True):
        """Translate the given templates with each experiment.

        Parameters
        ----------
        template : str
            Template to translate.
        with_exps : bool, optional
            Return the list of experiments used to get each translation.
        with_unique : bool, optional
            Only return unique translations.

        Returns
        -------
        A list of the template translations (``with_exps=False``), or a list of
        tuples with each translation and the corresponding experiments
        (``with_exps=True``).

        Examples
        --------
        Translations are unique by default:
        >>> e = Experiments()
        >>> e.params(a=range(2), b=range(2), c=range(2))
        >>> e.translate("{{a}}-{{b}}")
        ['0-0', '0-1', '1-0', '1-1']
        >>> e.translate("{{a}}-{{b}}", with_unique=False)
        ['0-0', '0-0', '0-1', '0-1', '1-0', '1-0', '1-1', '1-1']

        And can be easily extended with the experiments that led to each
        translation:
        >>> e.translate("{{a}}-{{b}}", with_exps=True) == [
        ...     ('0-0', [{'a': 0, 'b': 0, 'c': 0}, {'a': 0, 'b': 0, 'c': 1}]),
        ...     ('0-1', [{'a': 0, 'b': 1, 'c': 0}, {'a': 0, 'b': 1, 'c': 1}]),
        ...     ('1-0', [{'a': 1, 'b': 0, 'c': 0}, {'a': 1, 'b': 0, 'c': 1}]),
        ...     ('1-1', [{'a': 1, 'b': 1, 'c': 0}, {'a': 1, 'b': 1, 'c': 1}])]
        True
        >>> e.translate("{{a}}-{{b}}", with_exps=True, with_unique=False) == [
        ...     ('0-0', [{'a': 0, 'b': 0, 'c': 0}]),
        ...     ('0-0', [{'a': 0, 'b': 0, 'c': 1}]),
        ...     ('0-1', [{'a': 0, 'b': 1, 'c': 0}]),
        ...     ('0-1', [{'a': 0, 'b': 1, 'c': 1}]),
        ...     ('1-0', [{'a': 1, 'b': 0, 'c': 0}]),
        ...     ('1-0', [{'a': 1, 'b': 0, 'c': 1}]),
        ...     ('1-1', [{'a': 1, 'b': 1, 'c': 0}]),
        ...     ('1-1', [{'a': 1, 'b': 1, 'c': 1}])]
        True

        """
        return _do_translate(self._experiments, [], self._translate,
                             [template], with_exps, with_unique)


######################################################################
# Views

class ExperimentsView (pp.Pretty):
    """Proxy to a subset of elements in a `Experiments` instance.

    Notes
    -----
    You can use a view with Python's ``with`` statement to improve code
    readability:

    >>> e = Experiments()
    >>> e.params(a=range(2), b=range(2))
    >>> with e.view("a != b") as v, v.view_inverse() as i:
    ...     v.params(c=1)
    ...     i.params(c=2)
    >>> e
    Experiments([{'a': 0, 'b': 0, 'c': 2},
                 {'a': 0, 'b': 1, 'c': 1},
                 {'a': 1, 'b': 0, 'c': 1},
                 {'a': 1, 'b': 1, 'c': 2}])

    """

    def __init__(self, experiments, filter_base, filter_top,
                 **kwargs):
        assert len(kwargs) == 0
        self._base = experiments
        self._filters = [filter_base, filter_top]
        for filter_ in self._filters:
            filter_.validate(set(self._base._experiments.variables()))

    def _get_experiments(self):
        filter_ = self._as_filter()
        return self._base._experiments.select(filter_, allow_unknown=True)

    def _repr_pretty_(self, p, cycle):
        with self.pformat(p, cycle):
            have_prev = False
            if len(self) > 0:
                p.pretty(list(map(lambda elem: _sorted_dict(elem), self._get_experiments())))
                have_prev = True
            if self.out != Experiments.out:
                if have_prev:
                    p.text(",")
                    p.breakable()
                p.text("out=%r" % self.out)
                have_prev = True
            if self.dereference != Experiments.dereference:
                if have_prev:
                    p.text(",")
                    p.breakable()
                p.text("dereference=%r" % self.dereference)

    def __iter__(self):
        return self._get_experiments()

    def __len__(self):
        return len(list(self._get_experiments()))

    def __getitem__(self, index):
        return list(self._get_experiments())[index]

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass

    def _as_filter(self):
        return and_filters(*self._filters)

    @property
    def out(self):
        return self._base.out

    @property
    def dereference(self):
        return self._base.dereference

    def view(self, *filters):
        """Create a new view from this one.

        See also
        --------
        Experiments.view

        """
        return ExperimentsView(self._base,
                               self._as_filter(),
                               and_filters(*filters))

    def view_inverse(self):
        """Get the inverse view of this one.

        If using nested views, this will invert the filter of this one, not all
        filters for parent views.

        See also
        --------
        Experiments.view

        """
        return self._base.view(and_filters(self._filters[0],
                                           "not (%s)" % self._filters[1]))

    def execute(self, *args, **kwargs):
        """Same as `Experiments.execute`"""
        return _do_execute(self, *args, **kwargs)

    def params(self, *filters, **variables):
        """Same as `Experiments.params`"""
        new_experiments = _do_recombine(self._base._experiments, self._filters,
                                        variables, filters,
                                        False)
        self._base._experiments = new_experiments[0]
        self._base._translate |= new_experiments[1]

    def params_append(self, *filters, **variables):
        """Same as `Experiments.params_append`"""
        new_experiments = _do_recombine(self._base._experiments, [],
                                        variables, filters,
                                        True)
        self._base._experiments = new_experiments[0]
        self._base._translate |= new_experiments[1]

    def pack(self, template_from, template_to, dereference=None):
        """Same as `Experiments.pack`"""
        _do_pack(self, template_from, template_to, dereference=None)

    def generate(self, template_from, template_to):
        """Same as `Experiments.generate`"""
        _do_generate_user(self, template_from, template_to)

    def generate_jobs(self, system, template_to, export=[], depends=[],
                      submit_args=[], job_desc="jobs.jd"):
        """Same as `Experiments.generate_jobs`"""
        _do_generate_jobs(self, system, template_to, export, depends,
                      submit_args, job_desc)

    def translate(self, template, with_exps=False, with_unique=True):
        """Same as `Experiments.translate`"""
        return _do_translate(self._base._experiments, self._filters, self._translate,
                             [template], with_exps, with_unique)


    def __getattr__(self, name):
        if name in ["_base", "_filters"]:
            object.__getattribute__(self, name)
        else:
            attr = getattr(self._base, name)
            if callable(attr):
                if attr.__name__ in VIEW_SUPPORTED:
                    attr = getattr(self._base,
                                   VIEW_SUPPORTED[attr.__name__])

                def func(*args, **kwargs):
                    return attr(self._as_filter(), *args, **kwargs)
                return func
            else:
                return attr

    def __setattr__(self, name, value):
        if name in ["_base", "_filters"]:
            object.__setattr__(self, name, value)
        else:
            attr = getattr(self._base, name)
            assert not callable(attr)
            setattr(self._base, name, value)


######################################################################
# Function wrappers

class _WithExp(pp.Pretty):

    _EXPERIMENT = None

    def __init__(self, func):
        self._func = func

    def __call__(self, *args, **kwargs):
        return self._func(self._EXPERIMENT, *args, **kwargs)

    def _repr_pretty_(self, p, cycle):
        p.text("with_exp(")
        p.pretty(self._func)
        p.text(")")

    def __eq__(self, other):
        return isinstance(other, _WithExp) and self._func == other._func

    def __hash__(self):
        return hash((_WithExp, self._func))


def with_exp(func):
    """Wrap function to get experiment as first argument during translation.

    Examples
    --------
    >>> e = Experiments()
    >>> e.params(a=range(2), b=range(2),
    ...          c=with_exp(lambda exp: exp['a']+exp['b']))
    >>> e.translate("{{c}}")
    ['0', '1', '2']

    """
    return _WithExp(func)

class _WithExpTpl(pp.Pretty):

    _EXPERIMENT = None

    def __init__(self, func, template):
        self._func = func
        self._template = template
        self._translator = text.Translator(self._template)

    def __call__(self, *args, **kwargs):
        translated = self._translator.translate(self._EXPERIMENT)
        return self._func(self._EXPERIMENT, translated, *args, **kwargs)

    def _repr_pretty_(self, p, cycle):
        p.text("with_exp_tpl(")
        p.pretty(self._func)
        p.text(", ")
        p.pretty(self._template)
        p.text(")")

    def __eq__(self, other):
        return (isinstance(other, _WithExpTpl) and
                self._func == other._func and
                self._template == other._template)

    def __hash__(self):
        return hash((_WithExp, self._func, self._template))


def with_exp_tpl(func, template):
    """Similar to `with_exp`, but passes translated template as second argument.

    This is a convenience shortcut for using `text.translate` in `with_exp`.

    Examples
    --------
    In its basic form, we can combine translated templates with other experiment
    values:

    >>> e = Experiments()
    >>> def fun(exp, tpl):
    ...     return tpl + " == " + str(exp['a']+exp['b'])
    >>> e.params(a=range(2), b=range(2),
    ...          f=with_exp_tpl(fun, "{{a}}+{{b}}"))
    >>> e.translate("{{f}}")
    ['0+0 == 0', '0+1 == 1', '1+0 == 1', '1+1 == 2']

    But this comes in handy when we want to pull the contents of per-experiment
    files into new experiment values. Let's start by creating some files:

    >>> import tempfile
    >>> tmp_dir = tempfile.TemporaryDirectory()
    >>> with open(os.path.join(tmp_dir.name, "file-0"), "w") as f:
    ...     f.write("000")
    3
    >>> with open(os.path.join(tmp_dir.name, "file-1"), "w") as f:
    ...     f.write("111")
    3

    And now we can use them to populate experiment values:

    >>> def get_file(exp, path):
    ...    return open(path, "r").read()
    >>> v = e.view("b == 0")
    >>> v.params(f=with_exp_tpl(get_file, os.path.join(tmp_dir.name, "file-1")))
    >>> v.translate("{{f}}")
    ['111']
    >>> v.params(f=with_exp_tpl(get_file, os.path.join(tmp_dir.name, "file-{{a}}")))
    >>> v.translate("{{f}}")
    ['000', '111']

    """
    return _WithExpTpl(func, template)


class _FromFunction(object):
    def __init__(self, func, keys):
        self._func = func
        self._keys = keys
        self._results = None

    def param(self, key_name):
        if key_name not in self._keys:
            raise ValueError("unknown key name: %s" % key_name)
        if self._results is not None:
            raise Exception("cannot get parameter after getting function results")
        return _FromFunctionKey(self, key_name)

    def _get_results(self, exp):
        results = self._func(exp)
        if not isinstance(results, collections.abc.Sized):
            results = list(results)
        return results


class _FromFunctionKey(object):
    def __init__(self, from_function, key_name):
        self._from_func = from_function
        self._key_name = key_name


def from_function(fun, keys):
    """Parameter value to add given parameter permutations to every experiment.

    For every experiment, it will permute it with the values returned by
    `fun`. Note that `fun` is expected to return entries with multiple variable
    names (specified by `keys`), so you must call ``param`` on the result of
    `from_function` to select the specific variable name.

    Parameters
    ----------
    fun : callable
        Function to call on every experiment, receives the experiment at hand
        and must return a list of dict.
    keys : list of str
        Variable names in the elements returned by `fun`.

    Examples
    --------
    The result can use any of the values available on each experiment; in this
    case ``b`` is two and three times the value of ``a`` (note how `from_func`
    and other values can be mixed):

    >>> def func(exp):
    ...     return [{'b': exp["a"]*2}, {'b': exp["a"]*3}]
    >>> e = Experiments()
    >>> e.params(a=range(2), b=from_function(func, ["b"]).param("b"))
    >>> e
    Experiments([{'a': 0, 'b': 0}, {'a': 1, 'b': 2}, {'a': 1, 'b': 3}])

    See also
    --------
    Experiments.params

    """
    return _FromFunction(fun, keys)


def from_find_files(template, **kwargs):
    """Helper to use `sciexp2.common.utils.find_files` with `from_function`.

    Will provide all variable names that are tags in the template. If argument
    `path` is provided, will also provide that variable name (with the full path
    to the found file).

    Parameters
    ----------
    template : str
        Template of file paths to find.
    kwargs
        Arguments to `sciexp2.common.utils.find_files`.

    Examples
    --------
    After creating some example files:

    >>> import tempfile
    >>> tmp = tempfile.TemporaryDirectory()
    >>> _ = open(os.path.join(tmp.name, "0.txt"), "w")
    >>> _ = open(os.path.join(tmp.name, "1.txt"), "w")

    Now we can use the variables on the file template (``a``) and its full path
    (``PATH``) as new variables on the experiments (``var`` and ``path``):

    >>> e = Experiments()
    >>> f = from_find_files(os.path.join(tmp.name, "{{a}}.txt"), path="PATH")
    >>> e.params(var=f.param("a"), path=f.param("PATH"))
    >>> e                               # doctest:+ELLIPSIS
    Experiments([{'path': '/tmp/.../0.txt', 'var': 0},
                 {'path': '/tmp/.../1.txt', 'var': 1}])

    """
    template_vars = text.get_variables(template)
    keys = template_vars

    path = kwargs.get("path", None)
    if path is not None:
        if not isinstance(path, six.string_types):
            raise TypeError("path argument must be a string")
        keys.append(path)

    results = list(utils.find_files(template, **kwargs))

    def fun(experiment):
        experiment = dict((key, val) for key, val in experiment.items()
                          if key in template_vars)
        if len(experiment) > 0:
            filter_ = Filter(experiment)
        else:
            filter_ = Filter()
        for result in results:
            if filter_.match(result):
                yield result

    return from_function(fun, keys)


__all__ = [
    "Experiments", "ExperimentsView",
    "with_exp", "with_exp_tpl",
    "from_function", "from_find_files",
]
