"""Define new Ops from existing Ops"""
from __future__ import absolute_import, print_function, division
from functools import reduce, partial
from collections import OrderedDict

import theano
from theano import gof
from theano.compat import izip
from theano.compile.function_module import orig_function
from theano.compile import SharedVariable, rebuild_collect_shared, optdb
from theano.gof import ops_with_inner_function
from theano.gof.graph import io_connection_pattern
from theano.gof.null_type import NullType
from theano.gradient import DisconnectedType


class OpFromGraph(gof.Op):
    """
    This creates an ``Op`` from inputs and outputs lists of variables.
    The signature is similar to :func:`theano.function <theano.function>`
    and the resulting ``Op``'s perform will do the same operation as::

        orig_function(inputs, outputs, **kwargs)

    Currently does not support ``updates`` or ``givens`` argument.

    Parameters
    ----------

    inputs: list of :class:`Variable <theano.gof.Variable>`

    outputs: list of :class:`Variable <theano.gof.Variable>`

    inline: bool, optional
        Defaults to ``False``

        ``True`` : Cause the Op's original graph being used during
        compilation, the Op will not be visible in the compiled
        graph but rather its internal graph.

        ``False`` : will use a pre-compiled function inside.

    grad_overrides : single or list of {0, None, Ellipsis, OpFromGraph, callable}, optional
        Defaults to ``None``.

        ``None`` : No value, gives NullType()
        ``0`` : zero value, gives DisconnectedType()
        ``...`` : Do not override, use default grad() result

        OpFromGraph instance : Override with another OpFromGraph, should
        accept inputs as the same order and types of "inputs" and "output_grads"
        arguments as one would specify in grad() method.

        callable : similar to OpFromGraph instance, must return list of
        :class:`Variable <theano.gof.Variable>`.

        list: Each OpFromGraph/callable must return a single
        :class:`Variable <theano.gof.Variable>`. Each list element corresponds to gradient of
        a specific input, length of list must be equal to number of inputs.

    rop_overrides : single or list of {0, None, Ellipsis, OpFromGraph, callable}, optional
        Defaults to ``None``.

        ``None`` : No value, gives NullType()
        ``0`` : zero value, gives zeros_like(...)
        ``...`` : Do not override, use default R_op() result

        OpFromGraph instance : Override with another OpFromGraph, should
        accept inputs as the same order and types of "inputs" and "output_grads"
        arguments as one would specify in grad() method.

        callable : similar to OpFromGraph instance, must return list of
        :class:`Variable <theano.gof.Variable>`.

        list: Each OpFromGraph/callable must return a single
        :class:`Variable <theano.gof.Variable>`. Each list element corresponds
        to a specific output of R_op, length of list must be equal to number of outputs.

    name : string, optional
        A name for debugging purposes

    \*\*kwargs : optional
        Check
        :func:`orig_function <theano.compile.function_module.orig_function>`
        for more arguments, only works when not inline.


    .. TODO:
        - examples for a multi-layer mlp. where?
        - __hash__, __eq__ otherwise won't merge, try
          gof.opt.is_same_graph_with_merge(op1.local_outputs, op2,
          local_outputs)
        - c_code() to remove the double overhead?
        - grad() make it support DisconnectedType and the new interface
        - extend grad() to L_op
        - add support for NullType and DisconnectedType when R_op supports them
        - check how it works with updates.
        - add test with constant as input or inside the inner graph.
        - Add support for the GPU? Probably just need an opt to remove transfer
        - Add support to pickle this Op.
        - Add support/test with random generator
        - Add optimization to removing unused inputs/outputs
        - Add optimization to work inplace on inputs when not inline

    Notes
    -----
    - We support shared variables in the inner graph. This is automatic
      and invisible to the user. They can be as input to the node or in
      the inner graph.
    - We support unused inputs. This is needed for the grad.
    - We support nested OpFromGraph.
    - ``inline=True`` will cause better runtime optimization at the cost
      of compilation time. Currently only works with ``fast_compile`` or
      ``fast_run`` mode.
    - It's recommanded to provide pure functions (no side effects like
      setting global variable) as callable(s). The callable(s) supplied
      for overriding gradient/rop will be called only once at the first
      call to grad/R_op, and will be converted to OpFromGraph instances.

    Examples
    --------

    Example 1:

    .. code-block:: python

        from theano import function, OpFromGraph, tensor
        x, y, z = tensor.scalars('xyz')
        e = x + y * z
        op = OpFromGraph([x, y, z], [e])
        # op behaves like a normal theano op
        e2 = op(x, y, z) + op(z, y, x)
        fn = function([x, y, z], [e2])

    Example 2 with shared variable:

    .. code-block:: python

        import numpy as np
        import theano
        from theano import config, function, OpFromGraph, tensor
        x, y, z = tensor.scalars('xyz')
        s = theano.shared(np.random.rand(2, 2).astype(config.floatX))
        e = x + y * z + s
        op = OpFromGraph([x, y, z], [e])
        # op behaves like a normal theano op
        e2 = op(x, y, z) + op(z, y, x)
        fn = function([x, y, z], [e2])

    Example 3 override gradient

    .. code-block:: python

        from theano import function, OpFromGraph, tensor, grad
        x, y, z = tensor.scalars('xyz')
        e = x + y * z
        def rescale_dy(inps, grads):
            x, y, z = inps
            g, = grads
            return z*2
        op = OpFromGraph(
            [x, y, z], [e], grad_overrides=[Ellipsis, rescale_dy, Ellipsis]
        e2 = op(x, y, z)
        dx, dy, dz = grad(e2, [x, y, z])
        fn = function([x, y, z], [dx, dy, dz])
        # the gradient wrt y is now doubled
        fn(2., 3., 4.) # [1., 8., 3.]

    """
    ofg_null_t = NullType(why_null='ofg_overridden')
    ofg_discon_t = DisconnectedType()

    @staticmethod
    def _filter_grad_var(grad, inp):
        # Returns (filtered_var, overrider_var)
        # Args:
        #     grad: gradient Variable
        #     inp: the corresponding input of gradient Variable
        #
        # a grad() call could return instance of NullType() or DisconnectedType()
        # which cannot be directly used in OfG
        #
        # Since we always use an OfG instance as self._grad_op, the current
        # workaround is to "remember" the special cases of the gradient and
        # replace them after self._grad_op is called.
        #
        # This helper function changes invalid types into a filtered_var,
        # and provides a overrider_var to be replaced at grad() call
        #
        # For now, this converts NullType or DisconnectedType into zeros_like.
        # other types are unmodified: overrider_var -> None
        if isinstance(grad.type, (NullType, DisconnectedType)):
            return inp.zeros_like(), grad
        else:
            return grad, None

    @staticmethod
    def _filter_rop_var(inpJ, out):
        # mostly similar to _filter_grad_var
        if isinstance(inpJ.type, NullType):
            return out.zeros_like(), inpJ
        if isinstance(inpJ.type, DisconnectedType):
            # since R_op does not have DisconnectedType yet, we will just
            # make them zeros.
            return out.zeros_like(), None
        else:
            return inpJ, None

    def __init__(
            self, inputs, outputs,
            inline=False,
            grad_overrides=Ellipsis, rop_overrides=Ellipsis,
            name=None, **kwargs
    ):
        if not isinstance(outputs, list):
            raise TypeError('outputs must be list, got %s' % type(outputs))
        for i in inputs + outputs:
            if not isinstance(i, gof.Variable):
                raise TypeError(
                    'inputs and outputs must be Variable instances', i)
        if 'updates' in kwargs or 'givens' in kwargs:
            raise TypeError('updates and givens are not allowed here')
        self.is_inline = inline
        # To correctly support shared variables the inner fct should
        # not see them. Otherwise there is a problem with the gradient.
        self.shared_inputs = [var for var in gof.graph.inputs(outputs)
                              if isinstance(var, SharedVariable)]
        shared_vars = [var.type() for var in self.shared_inputs]

        new = rebuild_collect_shared(outputs, inputs=inputs + shared_vars,
                                     replace=dict(izip(
                                         self.shared_inputs, shared_vars)),
                                     copy_inputs_over=False)
        (local_inputs, local_outputs,
         [clone_d, update_d, update_expr, shared_inputs]) = new
        assert len(local_inputs) == len(inputs) + len(self.shared_inputs)
        assert len(local_outputs) == len(outputs)
        assert not update_d
        assert not update_expr
        assert not shared_inputs

        self.local_inputs = local_inputs
        self.local_outputs = local_outputs
        self.inputs = inputs
        self.outputs = outputs
        self.kwargs = kwargs
        self.input_types = [inp.type for inp in inputs]
        self.output_types = [out.type for out in outputs]
        self.set_grad_overrides(grad_overrides)
        self.set_rop_overrides(rop_overrides)

        if name is not None:
            assert isinstance(name, str), 'name must be None or string object'
        self.name = name

    def __eq__(self, other):
        # TODO: recognize a copy
        return self is other

    def __hash__(self):
        # TODO: use internal variables in hash
        return hash(type(self))

    def __str__(self):
        name = self.__class__.__name__ if self.name is None else self.name
        is_inline = self.is_inline
        return '%(name)s{inline=%(is_inline)s}' % locals()

    def _recompute_grad_op(self):
        local_inputs = self.local_inputs
        local_outputs = self.local_outputs
        inp_len = len(local_inputs)
        grad_op = self._grad_op

        if isinstance(grad_op, OpFromGraph):
            self._grad_op_is_cached = True
            self._grad_op_overrides_l = [None] * inp_len
            return

        output_grads = [out_t() for out_t in self.output_types]
        fn_grad = partial(
            theano.gradient.grad,
            cost=None,
            disconnected_inputs='ignore',
            return_disconnected='Disconnected',
            null_gradients='return',
            known_grads=OrderedDict(izip(local_outputs, output_grads)))

        TYPE_ERR_MSG = 'Gradient override should be (single or list of)' \
            'OpFromGraph | Ellipsis | None | 0 | callable, got %s'
        # we need to convert _grad_op into an OfG instance
        if grad_op is Ellipsis:
            gdefaults_l = fn_grad(wrt=local_inputs)
            all_grads_ov_l = [None] * inp_len
            all_grads_l, all_grads_ov_l = izip(
                *[OpFromGraph._filter_grad_var(grad, inp) for grad, inp in izip(gdefaults_l, local_inputs)])
        elif grad_op is None:
            all_grads_l = [inp.zeros_like() for inp in local_inputs]
            all_grads_ov_l = [self.ofg_null_t()] * inp_len
        elif type(grad_op) is int and grad_op == 0:
            all_grads_l = [inp.zeros_like() for inp in local_inputs]
            all_grads_ov_l = [self.ofg_discon_t()] * inp_len
        elif isinstance(grad_op, list):
            goverrides_l = grad_op
            if len(goverrides_l) != inp_len:
                raise ValueError(
                    'Need to override %d gradients, got %d' % (
                        inp_len, len(goverrides_l)), goverrides_l)
            # compute non-overriding downsteam grads from upstreams grads
            # it's normal some input may be disconnected, thus the 'ignore'
            wrt_l = [lin for lin, gov in izip(
                local_inputs, goverrides_l) if gov is Ellipsis]
            gdefaults = iter(fn_grad(wrt=wrt_l) if wrt_l else [])
            # combine overriding gradients
            all_grads_l = []
            all_grads_ov_l = []
            for inp, fn_gov in izip(local_inputs, goverrides_l):
                if fn_gov is Ellipsis:
                    gnext, gnext_ov = OpFromGraph._filter_grad_var(
                        next(gdefaults), inp)
                    all_grads_l.append(gnext)
                    all_grads_ov_l.append(gnext_ov)
                elif fn_gov is 0:
                    all_grads_l.append(inp.zeros_like())
                    all_grads_ov_l.append(self.ofg_discon_t())
                elif fn_gov is None:
                    all_grads_l.append(inp.zeros_like())
                    all_grads_ov_l.append(self.ofg_null_t())
                else:
                    if not hasattr(fn_gov, '__call__'):
                        raise TypeError(TYPE_ERR_MSG % fn_gov)
                    gov, gov_ov = OpFromGraph._filter_grad_var(
                        fn_gov(local_inputs, output_grads), inp)
                    all_grads_l.append(gov)
                    all_grads_ov_l.append(gov_ov)
        else:
            # callable case
            if not hasattr(grad_op, '__call__'):
                raise TypeError(TYPE_ERR_MSG % grad_op)
            goverrides_l = grad_op(local_inputs, output_grads)
            if not isinstance(goverrides_l, list):
                raise TypeError(
                    'Gradient overriding function should return a list, '
                    'got "%s"' % type(goverrides_l))
            all_grads_l, all_grads_ov_l = izip(
                *[OpFromGraph._filter_grad_var(
                    grad, inp) for grad, inp in izip(goverrides_l, local_inputs)])
            if len(all_grads_l) != len(local_inputs):
                raise ValueError(
                    'Gradient overriding function should return list of '
                    '%d outputs, got %d' % (inp_len, len(all_grads_l)))
        all_grads_l = list(all_grads_l)
        all_grads_ov_l = list(all_grads_ov_l)
        self._grad_op = type(self)(
            inputs=local_inputs + output_grads,
            outputs=all_grads_l,
            inline=self.is_inline,
            name=(None if self.name is None else self.name + '_grad'),
            on_unused_input='ignore')
        self._grad_op_overrides_l = all_grads_ov_l
        self._grad_op_is_cached = True

    def _recompute_rop_op(self):
        local_inputs = self.local_inputs
        local_outputs = self.local_outputs
        out_len = len(local_outputs)
        rop_op = self._rop_op

        if isinstance(self._rop_op, OpFromGraph):
            self._rop_op_is_cached = True
            self._rop_op_overrides_l = [None] * out_len
            return

        eval_points = [inp_t() for inp_t in self.input_types]
        fn_rop = partial(
            theano.gradient.Rop,
            wrt=local_inputs,
            eval_points=eval_points)
        TYPE_ERR_MSG = 'R_op overrides should be (single or list of)' \
            'OpFromGraph | Ellipsis | None | 0 | callable, got %s'
        if rop_op is Ellipsis:
            all_rops_l = fn_rop(f=local_outputs)
            all_rops_ov_l = [None] * out_len
        elif rop_op is None:
            all_rops_l = [out.zeros_like() for out in local_outputs]
            all_rops_ov_l = [self.ofg_null_t()] * out_len
        elif rop_op is 0:
            all_rops_l = [out.zeros_like() for out in local_outputs]
            all_rops_ov_l = [None] * out_len
        elif isinstance(rop_op, list):
            roverrides_l = rop_op
            if len(roverrides_l) != out_len:
                raise ValueError(
                    'Need to override %d Rop, got %d' % (
                        out_len, len(roverrides_l)), roverrides_l)
            # get outputs that does not have Rop override
            odefaults_l = [
                lo for lo, rov in izip(local_outputs, roverrides_l)
                if rov is Ellipsis]
            rdefaults_l = fn_rop(f=odefaults_l)
            rdefaults = iter(rdefaults_l if odefaults_l else [])
            # combine overriding Rops
            all_rops_l = []
            all_rops_ov_l = []
            for out, fn_rov in izip(local_outputs, roverrides_l):
                if fn_rov is Ellipsis:
                    rnext, rnext_ov = OpFromGraph._filter_rop_var(
                        next(rdefaults), out)
                    all_rops_l.append(rnext)
                    all_rops_ov_l.append(rnext_ov)
                elif fn_rov is 0:
                    all_rops_l.append(out.zeros_like())
                    all_rops_ov_l.append(None)
                elif fn_rov is None:
                    all_rops_l.append(out.zeros_like())
                    all_rops_ov_l.append(self.ofg_null_t())
                else:
                    if not hasattr(fn_rov, '__call__'):
                        raise TypeError(TYPE_ERR_MSG % fn_rov)
                    rov, rov_ov = OpFromGraph._filter_rop_var(
                        fn_rov(local_inputs, eval_points), out)
                    all_rops_l.append(rov)
                    all_rops_ov_l.append(rov_ov)
        else:
            if not hasattr(rop_op, '__call__'):
                raise TypeError(TYPE_ERR_MSG % rop_op)
            roverrides_l = rop_op(local_inputs, eval_points)
            if not isinstance(roverrides_l, list):
                raise TypeError(
                    'Rop overriding function should return a list, '
                    'got "%s"' % type(roverrides_l))
            all_rops_l, all_rops_ov_l = izip(
                *[OpFromGraph._filter_rop_var(
                    rop, out) for rop, out in izip(roverrides_l, local_outputs)])
            if len(all_rops_l) != out_len:
                raise ValueError(
                    'Rop overriding function %s should return list of '
                    '%d outputs, got %d' % (
                        self._rop_op, out_len,
                        len(all_rops_l)), rop_op)
            all_rops_l = list(all_rops_l)
            all_rops_ov_l = list(all_rops_ov_l)
        self._rop_op = type(self)(
            inputs=local_inputs + eval_points,
            outputs=all_rops_l,
            inline=self.is_inline,
            name=(None if self.name is None else self.name + '_rop'),
            on_unused_input='ignore')
        self._rop_op_overrides_l = all_rops_ov_l
        self._rop_op_is_cached = True

    def get_grad_op(self):
        """
        getter method for self._grad_op
        """
        if not self._grad_op_is_cached:
            self._recompute_grad_op()
        return self._grad_op

    def get_rop_op(self):
        """
        getter method for self._rop_op
        """
        if not self._rop_op_is_cached:
            self._recompute_rop_op()
        return self._rop_op

    def set_rop_overrides(self, rop_overrides):
        """
        Set R_op overrides, see help(theano.OpFromGraph) for syntax
        This will completely remove any previously set R_op overrides

        """
        self._rop_op = rop_overrides
        self._rop_op_is_cached = False

    def set_grad_overrides(self, grad_overrides):
        """
        Set gradient overrides, see help(theano.OpFromGraph) for syntax
        This will completely remove any previously set gradient overrides

        """
        self._grad_op = grad_overrides
        self._grad_op_is_cached = False

    def R_op(self, inputs, eval_points):
        if not self._rop_op_is_cached:
            self._recompute_rop_op()
        ret_ofg_l = self._rop_op(
            *(list(inputs) + list(eval_points)), return_list=True)
        ret_l = [
            ret_ofg if ov is None else ov for ret_ofg, ov in izip(
                ret_ofg_l, self._rop_op_overrides_l)]
        return ret_l

    def grad(self, inputs, output_grads):
        if not self._grad_op_is_cached:
            self._recompute_grad_op()
        ret_ofg_l = self._grad_op(
            *(list(inputs) + list(output_grads)), return_list=True)
        ret_l = [
            ret_ofg if ov is None else ov for ret_ofg, ov in izip(
                ret_ofg_l, self._grad_op_overrides_l)]
        return ret_l

    def make_node(self, *inputs):
        num_expected_inps = len(self.local_inputs) - len(self.shared_inputs)
        if len(inputs) != num_expected_inps:
            raise ValueError("Expected %d inputs, got %d" % (num_expected_inps, len(inputs)))
        inputs = [inp_t.filter_variable(inp) for inp, inp_t in izip(inputs, self.input_types)]
        apply_node = gof.Apply(
            self, list(inputs) + self.shared_inputs,
            [type() for type in self.output_types])
        apply_node.local_inputs = self.local_inputs
        apply_node.local_outputs = self.local_outputs
        return apply_node

    def connection_pattern(self, node):
        """
        Return connection pattern of subfgraph defined by inputs and outputs.

        """
        return io_connection_pattern(
            self.local_inputs, self.local_outputs)

    def infer_shape(self, node, shapes):
        out_shp = theano.scan_module.scan_utils.infer_shape(
            self.local_outputs,
            self.local_inputs,
            shapes)

        # Clone the output shape so that shape are computed from outer inputs.
        # Note:
        # Here we can do it more simply like:
        #      ret = [theano.clone(shp, replace=repl) for shp in out_shp]
        # But  doing it multiple time could duplicate common subgraph between
        # each shape call. Theano optimizer will clean this up later, but this
        # will ask extra work to the optimizer.
        repl = dict(zip(self.local_inputs, node.inputs))
        cloned = theano.clone(reduce(tuple.__add__, out_shp), replace=repl)
        ret = []
        used = 0
        for i in range(len(out_shp)):
            nb = len(out_shp[i])
            ret.append(cloned[used: used + nb])
            used += nb

        return ret

    def prepare_node(self, node, storage_map, compute_map, impl):
        if not hasattr(self, "fn") and impl == 'py':
            self.fn = orig_function(self.local_inputs,
                                    self.local_outputs,
                                    **self.kwargs)
            self.fn.trust_input = True

    def perform(self, node, inputs, outputs):
        variables = self.fn(*inputs)
        assert len(variables) == len(outputs)
        for output, variable in izip(outputs, variables):
            # TODO: when function's output-borrowing semantics are correct,
            # we wont need this copy anymore
            output[0] = variable.copy()


@gof.local_optimizer([OpFromGraph])
def inline_ofg_expansion(node):
    """
    This optimization expands internal graph of OpFromGraph.
    Only performed if node.op.is_inline == True
    Doing so can improve optimization at the cost of compilation speed.
    """
    op = node.op
    if not isinstance(op, OpFromGraph):
        return False
    if not op.is_inline:
        return False
    return theano.clone(
        op.local_outputs, {
            u: v for u, v in izip(
                node.op.local_inputs, node.inputs)})

# We want to run this before the first merge optimizer
# and before the first scan optimizer.
optdb.register(
    'inline_ofg_expansion',
    gof.opt.in2out(inline_ofg_expansion),
    -0.01, 'fast_compile', 'fast_run')

# Since OpFromGraph contains a Theano compiled function,
# we should let DebugMode know about it
ops_with_inner_function[OpFromGraph] = 'fn'
