## @package task
# Module caffe2.python.task

from caffe2.python import core, context
from caffe2.python.schema import Field, from_blob_list
from collections import defaultdict
from copy import copy
from future.utils import viewitems


def _merge_node_kwargs(a, b):
    # TODO(azzolini): consistency checks
    if a is None:
        return b
    if b is None:
        return a
    c = copy(a)
    c.update(b)
    return c


class Cluster(context.DefaultManaged):
    """
    Context that keeps track of all the node names used.
    Users shouldn't have to use them directly, since a Cluster is automatically
    generated at the first usage of 'Node'.
    """

    def __init__(self):
        # list instead of set to keep order
        self._nodes = []
        self._node_kwargs = {}

    def add_node(self, node):
        if str(node) not in self._nodes:
            self._nodes.append(str(node))
        self._node_kwargs[str(node)] = _merge_node_kwargs(
            node.kwargs(),
            self._node_kwargs.get(str(node)))

    def nodes(self):
        """
        Returns the list of unique node names used within this context.
        """
        return self._nodes

    def node_kwargs(self):
        return self._node_kwargs

    def __repr__(self):
        return "Cluster(nodes={}, node_kwargs={})".format(
            self.nodes(), self.node_kwargs())


class Node(context.DefaultManaged):
    """
    A Node context is used to indicate that all Tasks instantiated within will
    run on the given node name. (Only the name of the node actually counts.)
    Example:

        with TaskGroup() as tg:
            with Node('node1'):
                s1 = execution_step(...)
                Task(step=s1)
            with Node('node2'):
                s2 = execution_step(...)
            with Node('node1'):
                s3 = execution_step(...)

        In this example, all three execution steps will run in parallel.
        Moreover, s1 and s3 will run on the same node, and can see each
        others blobs.

        Additionally, a Node can be passed implementation-specific kwargs,
        in order to specify properties of the node.
    """

    def __init__(self, node='local', **kwargs):
        self._name = str(node)
        self._kwargs = kwargs
        Cluster.current().add_node(self)

    def __str__(self):
        return self._name

    def __repr__(self):
        return "Node(name={}, kwargs={})".format(self._name, self._kwargs)

    def kwargs(self):
        return self._kwargs


class WorkspaceType(object):
    """
    Determines whether tasks of a TaskGroup will run directly at the global
    workspace, which is kept alive across runs, or whether a new child
    workspace will be created for the run and destroyed afterwards.
    """
    PRIVATE = 'private'
    GLOBAL = 'global'


def get_setup_nets(key, steps_or_nets, target):
    init_net = core.Net(key + '/init')
    exit_net = core.Net(key + '/exit')
    init_nets = []
    exit_nets = []
    objs = []
    for step_or_net in steps_or_nets:
        if hasattr(step_or_net, 'get_all_attributes'):
            objs += step_or_net.get_all_attributes(key)
        elif hasattr(step_or_net, 'get_attributes'):
            objs += step_or_net.get_attributes(key)
    for obj in objs:
        # these are needed in order to allow nesting of TaskGroup, which
        # is a feature not yet implemented.
        if hasattr(obj, '_setup_used') and obj._setup_used:
            continue
        if hasattr(obj, '_setup_target') and obj._setup_target != target:
            continue
        if hasattr(obj, 'setup'):
            nets = obj.setup(init_net)
            if isinstance(nets, (list, tuple)):
                init_nets += nets
            elif isinstance(nets, (core.Net, core.ExecutionStep)):
                init_nets.append(nets)
            elif nets is not None:
                raise TypeError('Unsupported type for setup: %s' % type(nets))
            obj._setup_used = True
        if hasattr(obj, 'exit'):
            nets = obj.exit(exit_net)
            if isinstance(nets, (list, tuple)):
                exit_nets += nets
            elif isinstance(nets, (core.Net, core.ExecutionStep)):
                exit_nets.append(nets)
            elif nets is not None:
                raise TypeError('Unsupported type for setup: %s' % type(nets))
            obj._setup_used = True

    if len(init_net.Proto().op) > 0:
        init_nets.insert(0, init_net)
    if len(exit_net.Proto().op) > 0:
        exit_nets.insert(0, exit_net)
    return init_nets, exit_nets


def add_setup_steps(step, init_nets, exit_nets, name):
    if not init_nets and not exit_nets:
        return step
    steps = []
    if init_nets:
        steps.append(core.execution_step('%s:init' % name, init_nets))
    steps.append(step)
    if len(exit_nets) > 0:
        steps.append(core.execution_step('%s:exit' % name, exit_nets))
    return core.execution_step(name, steps)


class TaskGroup(context.Managed):
    """
    Context that gathers tasks which will run concurrently, potentially on
    multiple nodes. All tasks in the same node will share the same workspace
    and thus can share blobs, while tasks running in different nodes won't
    be able to directly share data.

    All tasks of the task group will start concurrently, and the task group
    will finish execution when the last task of the group finishes.

    Example:
        # suppose that s1 ... s5 are execution steps or nets.
        with TaskGroup() as tg:
            # these tasks go to default node 'local'
            Task(step=s1)
            Task(step=s2)

            with Node('n2'):
                Task(step=s3)
            with Node('n1'):
                Task(step=s4)
            with Node('n2'):
                Task(step=s5)

        # this will run all steps in parallel.
        # s1 and s2 will run at default node 'local'
        # s3 and s5 will run at node 'n2'
        # s4 will run at node 'n1'
        session.run(tg)
    """
    LOCAL_SETUP = 'local_setup'

    def __init__(self, workspace_type=None):
        self._plan_cache = None
        self._tasks = []
        self._already_used = False
        self._prev_active = None
        self._tasks_to_add = []
        self._report_nets = {}
        self._report_steps = []
        self._workspace_type = workspace_type
        self._tasks_by_node = None
        self._remote_nets = []

    def add_remote_net(self, net):
        self._remote_nets.append(net)

    def remote_nets(self):
        return self._remote_nets

    def add(self, task):
        assert not self._already_used, (
            'Cannot add Task to an already used TaskGroup.')
        assert (
            self._workspace_type is None or
            task._workspace_type is None or
            self._workspace_type == task._workspace_type)
        if task._workspace_type is None:
            task._workspace_type = (
                self._workspace_type or WorkspaceType.PRIVATE)
        if self._workspace_type is None:
            self._workspace_type = task._workspace_type
        task._notify_used()
        self._tasks.append(task)

    def tasks(self):
        for task in self._tasks_to_add:
            self.add(task)
        self._tasks_to_add = []
        self._already_used = True
        return self._tasks

    def num_registered_tasks(self):
        return len(self._tasks_to_add) + len(self._tasks)

    def used_nodes(self):
        # use list to keep order
        used = []
        for task in self._tasks + self._tasks_to_add:
            if task.node not in used:
                used.append(task.node)
        return used

    def report_step(self, step=None, node=None, interval_ms=1000):
        """
        Add a "report step" to this TaskGroup. This step will run repeatedly
        every `interval_ms` milliseconds for the duration of the TaskGroup
        execution on each of the nodes. It is guaranteed that this step
        will be run at least once after every Task in the node has finished.
        """
        step = core.to_execution_step(step)
        step.RunEveryMillis(interval_ms)
        self._report_steps.append((str(node or Node.current(node)), step))

    def report_net(self, net=None, node=None, report_interval=5):
        """
        DEPRECATED. Use report_step instead.
        """
        node = str(node or Node.current(node))
        assert net is None or node not in self._report_nets
        if node not in self._report_nets:
            self._report_nets[node] = (
                net if net else core.Net('%s/reporter' % node),
                report_interval)
        return self._report_nets[node][0]

    def tasks_by_node(self, node_remap=None):
        # tasks_by_node can't be called twice because the setup won't
        # work properly a second time.
        node_map = {}
        for task in self.tasks():
            node_map[task.node] =\
                node_remap(task.node) if node_remap else task.node
        if self._tasks_by_node is not None:
            tasks_by_node, prev_node_map = self._tasks_by_node
            assert prev_node_map == node_map, (
                'Cannot call tasks_by_node multiple times.')
            return tasks_by_node

        # now we have report_steps. report_net is deprecated
        for node, (net, interval) in viewitems(self._report_nets):
            self.report_step(net, node=node, interval_ms=interval * 1000)
        self._report_nets = {}

        tasks_by_node = defaultdict(list)
        for task in self.tasks():
            mapped_node = node_map[task.node]
            tasks_by_node[mapped_node].append(task)

        report_steps_by_node = defaultdict(list)
        for original_node, step in self._report_steps:
            report_steps_by_node[node_map[original_node]].append(step)

        grouped_by_node = TaskGroup()
        for node, tasks in viewitems(tasks_by_node):
            report_steps = report_steps_by_node[node]
            node_inits, node_exits = get_setup_nets(
                TaskGroup.LOCAL_SETUP,
                [t.get_step() for t in tasks] + report_steps,
                self)
            # shortcut for single task with no queue
            steps = report_steps
            outputs = []
            grouped_workspace_type = WorkspaceType.PRIVATE
            for task in tasks:
                step = task.get_step()
                step.SetCreateWorkspace(
                    task.workspace_type() == WorkspaceType.PRIVATE)
                if step is not None:
                    steps.append(step)
                outputs += task.outputs()
                # If any of the tasks in the node uses the global workspace,
                # then set the grouped task to use the global workspace as well
                if task.workspace_type() == WorkspaceType.GLOBAL:
                    grouped_workspace_type = WorkspaceType.GLOBAL
            if len(steps) == 0:
                steps.append(core.execution_step('empty', []))
            if len(steps) == 1:
                step = steps[0]
            else:
                step = core.execution_step(
                    '%s:body' % node, steps, concurrent_substeps=True)
            if len(node_inits) > 0 or len(node_exits) > 0:
                steps = []
                if len(node_inits) > 0:
                    steps.append(
                        core.execution_step('%s:init' % node, node_inits))
                steps.append(step)
                if len(node_exits) > 0:
                    steps.append(
                        core.execution_step('%s:exit' % node, node_exits))
                step = core.execution_step(node, steps)
            Task(
                node=node, step=step, outputs=outputs,
                name='grouped_by_node',
                group=grouped_by_node, workspace_type=grouped_workspace_type)
        self._tasks_by_node = (grouped_by_node, node_map)
        return grouped_by_node

    def to_task(self, node=None):
        node = str(Node.current(node))
        tasks = self.tasks_by_node(lambda x: node).tasks()
        if len(tasks) == 0:
            return Task()
        return tasks[0]

    def workspace_type(self):