# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import contextlib
import functools
import itertools
import os
from typing import Callable, Tuple, Union

import numpy as np

import megengine._internal as mgb
from megengine._internal.plugin import CompGraphProfiler

from ..core import Tensor, graph, tensor
from .sublinear_memory_config import SublinearMemoryConfig


def sideeffect(f):
    # during eager tracing, wrapped function is called with proxy inputs
    # during static tracing, wrapped function will not be called at all
    @functools.wraps(f)
    def wrapper(*args, **kwargs):  # pylint: disable=inconsistent-return-statements
        if not trace._active_instance:
            return f(*args, **kwargs)

        tensors = {}
        for i, x in itertools.chain(enumerate(args), kwargs.items()):
            if isinstance(x, Tensor):
                tensors[i] = x
        if tensors:
            _keys, tensors = zip(*tensors.items())
        else:
            _keys, tensors = (), ()

        def callback(*tensors, f=f, keys=_keys, args=args, kwargs=kwargs):
            replace = dict(zip(keys, tensors))
            args = tuple(replace.get(i, x) for i, x in enumerate(args))
            kwargs = {i: replace.get(i, x) for i, x in kwargs.items()}
            if f(*args, **kwargs) is not None:
                raise TypeError("a sideeffect function should return None")
            # TODO: clear memory

        trace._active_instance._register_callback(callback, tensors)

    return wrapper


def mark_impure(x):
    if not trace._active_instance:
        return x
    return trace._active_instance._mark_impure(x)


def barrier(x):
    if not trace._active_instance:
        return x
    return trace._active_instance._insert_barrier(x)


def _dummy():
    return mgb.make_immutable(*graph._use_default_if_none(None, None), 0)


class unset:
    pass


class trace:
    """
    Wrap a callable and provide:

    * tracing via :meth:`.trace` and :meth:`.dump`
    * accelerated evalutaion via :meth:`.__call__`

    :param func: Positional only argument.
    :param symbolic: Whether to use symbolic tensor. Default: False
    :param opt_level: Optimization level for compiling trace.
    :param log_level: Log level.
    :param sublinear_memory_config: Configuration for sublinear memory optimization.
        If not None, it enables sublinear memory optimization with given setting.
    :param allreduce_pack_max_size: Maximum size of an allreduce pack in MB.
        If not None, multiple gradients will be packed and synchronized together
    :param profiling: Whether to profile compiled trace. Default: False
    """

    _active_instance = None
    enabled = not os.getenv("MGE_DISABLE_TRACE")

    _UNSTARTED = "unstarted"
    _STARTED = "started"
    _FINISHED = "finished"

    def __new__(cls, *args, **kwargs):
        if not args:
            return functools.partial(cls, **kwargs)
        return super().__new__(cls)

    def __init__(
        self,
        func: Callable[..., Union[None, Tensor, Tuple[Tensor]]],
        *,
        symbolic: bool = False,
        opt_level: int = None,
        log_level: int = None,
        sublinear_memory_config: SublinearMemoryConfig = None,
        allreduce_pack_max_size: int = None,
        profiling: bool = False
    ):
        self.__wrapped__ = func
        self._symbolic = symbolic
        self._graph_opt_level = opt_level
        self._log_level = log_level
        self._sublinear_memory_config = sublinear_memory_config
        self._allreduce_pack_max_size = allreduce_pack_max_size
        self._status = self._UNSTARTED
        self._args = None
        self._kwargs = None
        self._outputs = unset
        self._sym_outputs = unset
        self._outspec = None
        self._checkpoint = None
        self._compiled_func = None
        self._profiling = profiling
        self._profiler = None

    @property
    def _active(self):
        c1 = self._status == self._STARTED
        c2 = type(self)._active_instance is self
        assert c1 == c2
        return c1

    def _register_callback(self, f, args=()):
        assert self._active
        assert isinstance(args, (tuple, list))
        proxies = self._make_proxies(args)
        self._forward(args, proxies, checkpoint=True)
        # NOTE: under eager graph callback will fire immediately
        job = mgb.opr.callback_injector(
            self._insert_barrier(_dummy()), lambda _: f(*proxies)
        )
        self._insert_checkpoint(job)
        self._outspec.append(job)

    def _insert_barrier(self, x):
        assert self._active
        if self._checkpoint is None:
            return x
        if isinstance(x, Tensor):
            x = x._symvar
            wrap = True
        else:
            wrap = False
        if not isinstance(x, mgb.SymbolVar):
            raise TypeError
        x = mgb.opr.virtual_dep([x, self._checkpoint])
        if wrap:
            x = Tensor(x)
        return x

    def _insert_checkpoint(self, *args, no_barrier=False):
        assert self._active
        if not args:
            return
        args = tuple(x._symvar if isinstance(x, Tensor) else x for x in args)
        for x in args:
            if not isinstance(x, mgb.SymbolVar):
                raise TypeError
        if not no_barrier and self._checkpoint is not None:
            # normally no need to _insert_barrier here, but if
            # someone forget to call _insert_barrier beforehand,
            # this can make things less broken
            args += (self._checkpoint,)
        if len(args) == 1:
            self._checkpoint = args[0]
        else:
            self._checkpoint = mgb.opr.virtual_dep(args)

    def _mark_impure(self, x):
        assert self._active
        ret = x
        if isinstance(x, Tensor):
            x = x._symvar
        if not isinstance(x, mgb.SymbolVar):
            raise TypeError
        self._outspec.append(x)
        self._insert_checkpoint(x)
        return ret

    def _make_proxies(self, args):
        assert isinstance(args, (tuple, list))
        for x in args:
            assert isinstance(x, Tensor)
        return tuple(tensor(dtype=x.dtype, device=x.device) for x in args)

    def _forward(self, srcs, dests, checkpoint=True):
        # pseudo-op: does not run under static graph; traced
        # TODO: use shared memory
        assert len(srcs) == len(dests)
        if not self._active:
            for s, d in zip(srcs, dests):
                d.set_value(s, share=False)
            return
        jobs = []
        for s, d in zip(srcs, dests):

            def callback(value, dest=d):
                dest.set_value(value, share=False)

            s = self._insert_barrier(s._symvar)
            # NOTE: callback immediately fire in eager graph
            jobs.append(
mgb.opr.callback_injector(s, callback)