# Copyright (c) 2010- The University of Notre Dame.
# This software is distributed under the GNU General Public License.
# See the file COPYING for details.

""" Weaver nest module """

from weaver.abstraction import SENTINEL
from weaver.compat      import map
from weaver.data        import parse_input_list
from weaver.engine      import Makeflow
from weaver.logger      import D_NEST, debug
from weaver.stack       import CurrentNest, CurrentScript, WeaverNests, \
                               stack_context_manager
from weaver.util        import Stash, groups, make_directory, parse_string_list

import collections
import os


# Nest class

@stack_context_manager(D_NEST, WeaverNests)
class Nest(Makeflow):
    """ Weaver Nest class.

    Fundamentally, a :class:`Nest` serves as an execution context.  Every task
    generated by Weaver is associated with a :class:`Nest`.  Likewise, every
    :class:`~weaver.file.File` and :class:`~weaver.abstraction.Abstraction` is
    also associated with a particular :class:`Nest`.

    Internally, a :class:`Nest` consists of a directory that serves as the
    workflow workspace ...

    **Internal attributes**:

    - work_dir  -- Working directory for the target workflow.
    - stash     -- Stash for intermediate output files.
    - futures
    """

    def __init__(self, work_dir=None, dag_path=None, stash=None, barrier=None,
        wrapper=None, track_imports=True, track_exports=True):
        self.work_dir = work_dir or '.'
        self.tasks    = []
        self.parent   = CurrentNest()
        if self.parent:
            self.work_dir = os.path.join(self.parent.work_dir, self.work_dir)
        self.stash    = stash or Stash(root=os.path.join(self.work_dir, '_Stash'))

        if not os.path.exists(self.work_dir):
            make_directory(self.work_dir)

        Makeflow.__init__(self, wrapper=wrapper,
            track_imports=track_imports, track_exports=track_exports)

        self.dag_path = dag_path or os.path.join(self.work_dir, 'Makeflow')
        self.dag_file = open(self.dag_path, 'w')
        self.includes.add(self.dag_path)
        # TODO: fix work_dir so it can be translated by makeflow_link

        if barrier:
            self.includes.update(parse_input_list(barrier))

        # Since Abstractions and SubNests are not compiled immediately, these
        # objects must regster with their parent Nest, who will compile them in
        # the order that they are registered to ensure proper semantics.
        self.futures  = []

        if self.parent:
            debug(D_NEST, 'Register child {0} with parent {1}'.format(
                self, self.parent))
            self.parent.futures.append((self, True))

        debug(D_NEST, 'Created {0}'.format(self))

    def __str__(self):
        return 'Nest({0})'.format(self.work_dir)

    def _optimize_nested_abstractions(self):
        """ Internally, we perform inline abstractions optimization as we build
        the DAG, so we should only execute the body of this method if we want
        to automatically nest abstractions after the fact.
        """
        if not CurrentScript().nested_abstractions:
            return

        debug(D_NEST, 'Inlining Abstractions for {0}'.format(self))

        # Group tasks into bins based on Abstractions.
        task_dict = collections.defaultdict(list)
        for task in self.tasks:
            abstraction = task[0]
            task_dict[abstraction].append(task)

        # For each Abstraction, create InlineNest and schedule tasks to be
        # executed there; only do this if we have more than one Abstraction.
        self.tasks = []
        if len(task_dict.keys()) > 1:
            for abstraction, tasks in task_dict.items():
                # For tasks scheduled directly by a Function (Abstraction is
                # None), then simply schedule for execution in current Nest.
                if abstraction is SENTINEL:
                    self.tasks.extend(tasks)
                    continue

                # Otherwise, create a new InlineNest and then schedule tasks to
                # run in this new Nest.
                with InlineNest() as inline_nest:
                    for task in tasks:
                        inline_nest.schedule(*task)
                    inline_nest.compile()

                # Engine is also a Function, so call it to schedule the task
                # responsible for InlineNest to run in the current Nest.
                with abstraction.options:
                    inline_nest()
        else:
            # Copy tasks from Abstractions to Nest task list.
            for abstraction, tasks in task_dict.items():
                for task in tasks:
                    self.tasks.append(task)

    def _optimize_inline_tasks(self):
        """ Group tasks by abstraction and function and then break them into
        sub-groups and schedule the sub-groups has sub DAGs.
        """
        if CurrentScript().inline_tasks <= 1:
            return

        debug(D_NEST, 'Inlining tasks for {0}'.format(self))

        # Group tasks into bins based on Function.
        task_dict = collections.defaultdict(list)
        for task in self.tasks:
            abstraction = task[0]
            function    = task[1]
            task_dict[(abstraction, function)].append(task)

        # For each set of tasks, split the set into small sub-groups; for each
        # sub-group, create a new InlineNest and schedule the tasks there.
        self.tasks = []
        for (abstraction, function), tasks in task_dict.items():
            inline_tasks = max(CurrentScript().inline_tasks, abstraction.group)
            if inline_tasks < len(tasks):
                for group in groups(tasks, inline_tasks):
                    with InlineNest() as inline_nest:
                        for task in group:
                            inline_nest.schedule(*task)
                        inline_nest.compile()
                    with abstraction.options:
                        inline_nest()
            else:
                for task in tasks:
                    self.tasks.append(task)

    def compile(self):
        """ Compile Abstractions to generate tasks and output file lists. """
        # Compile Abstractions and SubNests to ensure they have generated
        # tasks.
        debug(D_NEST, 'Compiling Abstractions and SubNests for {0}'.format(self))
        for future, is_nest in self.futures:
            if is_nest:
                with future:
                    future.compile()
                future()
            else:
                future.compile()

        # Perform optimizations.
        debug(D_NEST, 'Optimizing tasks for {0}'.format(self))
        self._optimize_nested_abstractions()
        self._optimize_inline_tasks()

        # Emit stored tasks to workflow DAG using engine.
        debug(D_NEST, 'Emitting tasks for {0}'.format(self))
        for task in self.tasks:
            self.emit_task(*task)

        # Emit variables and exports
        debug(D_NEST, 'Emitting variables for {0}'.format(self))
        self.emit_variables()

        debug(D_NEST, 'Emitting exports for {0}'.format(self))
        self.emit_exports()

    def schedule(self, abstraction, function, command, inputs, outputs,
        options):
        """ Schedule task for execution. """
        debug(D_NEST,
            'Scheduling task({0}, {1}, {2}, [{3}], [{4}], {5}) for {6}'.format(
            abstraction, function, command,
            ', '.join(map(str, inputs)), ', '.join(map(str, outputs)),
            options, self))

        if abstraction is None:
            abstraction = SENTINEL
        self.tasks.append(
            [abstraction, function, command, inputs, outputs, options])

    @property
    def root(self):
        """ Return the root Nest. """
        if self.parent is None:
            return self

        root = self.parent
        while root.parent is not None:
            root = root.parent
        return root


class InlineNest(Nest):
    """ Weaver InlineNest class. """
    def __init__(self):
        Nest.__init__(self, work_dir='.', dag_path=next(CurrentNest().stash),
            stash=CurrentNest().stash)


# Define and export functions

def Define(key, value, export=False):
    if export:
        Export(key)
    CurrentNest().variables[key] = value

def Export(variables):
    CurrentNest().exports.update(parse_string_list(variables))

# vim: set sts=4 sw=4 ts=8 expandtab ft=python:
