# Copyright (c) 2016-present, Facebook, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##############################################################################

## @package session
# Module caffe2.python.session
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals


from caffe2.python import core, workspace
from caffe2.python.task import Cluster, Task, TaskGroup, WorkspaceType


class CompiledRunnable(object):
    """ Wrapper for compiled runnable returned from session.compile() """
    def __init__(self, obj, session_class):
        self.obj = obj
        self.session_class = session_class


class Session(object):
    """
    Allows to run Nets, ExecutionSteps, Plans, Tasks and TaskGroups.
    A session can potentially run in multiple nodes concurrently.


    Example:
        from core import Net
        from caffe2.python.task import Task, TaskGroup, WorkspaceType

        net = Net('test1')
        net.Add([net.Const(1), net.Const(2)])

        net2 = net.Clone()
        step = core.execution_step('step1', [net2])

        with TaskGroup(WorkspaceType.GLOBAL) as init_tg:
            with Node('node1'):
                n1setup = net.Net('n1setup')
                n1msg = n1setup.Const('Hello from node 1.')
                Task(step=n1setup)

        with TaskGroup() as private_tg:
            with Node('node1'):
                n1 = net.Net('n1')
                n1.Print(n1msg, 0)
                Task(step=n1)
            with Node('node2'):
                n2 = net.Net('n2')
                n2.Print(n2.Const('Hello from node 2.'), 0)
                Task(step=n2)

        session = LocalSession()
        session.run(net)
        session.run(step)
        session.run(init_tg)
        session.run(private_tg)


    Global Workspace:
        At the beggining of the session, a global workspace is created and kept
        alive for the duration of the session.


    Private Workspace:
        Tasks can be run either directly on the global workspace, or they can
        instantiate a private child workspace that is released after each run.

    Blob visibility:
        Tasks running in different nodes in parallel will always run under
        different workspaces, so it must be assumed that they won't be able to
        access each other's blobs. Tasks running on the same node will follow
        Workspace hierarchy rules: tasks running on separate private workspaces
        will only be able to share blobs defined on a common parent Workspace.
    """

    _compiled_cache = {}

    def __init__(self):
        self._open = True

    def is_open(self):
        return self._open

    @classmethod
    def compile(cls, runnable, workspace_type=None, setup_net_list=None):
        if isinstance(runnable, CompiledRunnable):
            assert cls == runnable.session_class, (
                'Runnable was compiled for different session type. ' +
                'Need: %s, got: %s' % (
                    cls.__name__, runnable.session_class.__name__))
            return runnable

        if runnable in cls._compiled_cache:
            return cls._compiled_cache[runnable]

        if isinstance(runnable, TaskGroup):
            if workspace_type:
                if runnable.workspace_type():
                    assert runnable.workspace_type() == workspace_type, \
                        "Require {} but already have {}".format(
                            workspace_type, runnable.workspace_type())
                else:
                    runnable._workspace_type = workspace_type
            tg = runnable
        else:
            if workspace_type is None:
                workspace_type = WorkspaceType.GLOBAL
            tg = TaskGroup(workspace_type=workspace_type)
            if isinstance(runnable, Task):
                tg.add(runnable)