# pylint: disable=cell-var-from-loop,no-else-return
"""Helper utilities to check functions and their gradients."""
from __future__ import absolute_import as _abs

import logging
import numpy as np

import tvm
from tvm.contrib import graph_runtime
from tvm.testing import check_numerical_grads
from tvm import relay

import nnvm
from nnvm.compiler import graph_util
from nnvm.compiler.graph_attr import TCODE_TO_DTYPE, DTYPE_TO_TCODE
from nnvm.to_relay import to_relay
from .config import ctx_list

def infer_shapes_dtypes(graph, shape=None, dtype=None, fallback_dtype=None):
    """Runs dtype and shape inference passes on a graph and returns the resulting graph
    along with the inferred information.

    Parameters
    ----------
    graph : nnvm.graph.Graph
        A graph we want to run inference on.

    shape : Dict[str, Tuple[int]] or Tuple[int], optional
        A dict mapping input variable names to shapes.
        By default shapes will be inferred from variables' attributes.
        Note that this parameter takes precedence over variables' attributes.

    dtype : Dict[str, str] or str, optional
        A dict mapping input variable names to dtypes, or just a single dtype.
        By default dtypes will be inferred from variables' attributes.
        Note that this parameter takes precedence over variables' attributes.

    fallback_dtype : str, optional
        A dtype that will be used for variables whose dtype can't be inferred from other
        variables' dtypes.

    Returns
    -------
    graph : nnvm.graph.Graph
        The resulting graph with dtype and shape information on its nodes.

    input_shapes : Dict[str, Tuple[int]]
        The inferred shapes of input variables merged with the `shape` dictionary.

    input_dtypes : Dict[str, str]
        The inferred dtypes of input variables merged with the `dtype` dictionary.

    output_shapes : List[Tuple[int]]
        The inferred shapes of outputs.

    output_dtypes : List[str]
        The inferred dtypes of outputs.
    """
    # Preprocess input parameters
    if shape is None:
        provided_shapes = {}
    elif isinstance(shape, dict):
        provided_shapes = shape
    else:
        provided_shapes = {x: shape for x in graph.symbol.list_input_variables()}

    if dtype is None:
        provided_dtypes = {}
    elif isinstance(dtype, dict):
        provided_dtypes = dtype
    else:
        provided_dtypes = {x: dtype for x in graph.symbol.list_input_variables()}

    provided_shapes = _dict_var_to_dict_str(provided_shapes)
    provided_dtypes = _dict_var_to_dict_str(provided_dtypes)

    # The graph may already contain shape and dtype info, so extract it and merge with
    # the user-specified shapes and dtypes (use the user-specified one on contradiction)
    preexisting_shapes = graph.json_attr('shape')
    preexisting_dtypes = graph.json_attr('dtype')

    if preexisting_shapes:
        for x in graph.index.input_names:
            if x not in provided_shapes:
                x_shape = tuple(preexisting_shapes[graph.index.entry_id(x)])
                provided_shapes[x] = x_shape

    if preexisting_dtypes:
        for x in graph.index.input_names:
            if x not in provided_dtypes:
                x_dtype = TCODE_TO_DTYPE[preexisting_dtypes[graph.index.entry_id(x)]]
                provided_dtypes[x] = x_dtype

    # Perform inference
    nnvm.compiler.graph_attr.set_shape_inputs(graph, provided_shapes)
    nnvm.compiler.graph_attr.set_dtype_inputs(graph, provided_dtypes)

    graph = graph.apply('InferShape').apply('InferType')

    inferred_shapes = graph.json_attr('shape')
    inferred_dtypes = graph.json_attr('dtype')

    index = graph.index

    output_shapes = [tuple(inferred_shapes[index.entry_id(entry)])
                     for entry in index.output_entries]
    output_dtypes = [TCODE_TO_DTYPE[inferred_dtypes[index.entry_id(entry)]]
                     for entry in index.output_entries]

    # Postprocess the results
    input_shapes = provided_shapes.copy()
    input_dtypes = provided_dtypes.copy()

    for x in graph.symbol.list_input_variables():
        x_name = x.attr('name')
        x_entry_id = graph.index.entry_id(x_name)
        input_shapes[x_name] = tuple(inferred_shapes[x_entry_id])
        input_dtypes[x_name] = TCODE_TO_DTYPE[inferred_dtypes[x_entry_id]]

    # Merge the original user-specified shapes in case some of them are specified for non-existing
    # variables
    for x_name, x_shape in provided_shapes.items():
        x_shape = tuple(x_shape)
        if input_shapes.get(x_name, x_shape) != x_shape:
            raise RuntimeError("Inferred shape differs from the provided shape.\n"
                               "Provided shapes: {}\nInferred shapes: {}"
                               .format(provided_shapes, input_shapes))
        else:
            input_shapes[x_name] = x_shape

    # Merge the original user-specified dtypes
    for x_name, x_dtype in provided_dtypes.items():
        if not isinstance(x_dtype, str):
            x_dtype = TCODE_TO_DTYPE[x_dtype]
        if input_dtypes.get(x_name, x_dtype) != x_dtype:
            raise RuntimeError("Inferred dtype differs from the provided dtype.\n"
                               "Provided dtypes: {}\nInferred dtypes: {}"
                               .format(provided_dtypes, input_dtypes))
        else:
            input_dtypes[x_name] = x_dtype

    # If some dtypes weren't inferred and there is a fallback dtype, assign it to those varibles
    # and repeat the inference
    if fallback_dtype is not None and not all(input_dtypes.values()):
        input_dtypes = {x: input_dtypes[x] if input_dtypes[x] else fallback_dtype
                        for x in input_dtypes}
        return infer_shapes_dtypes(graph, input_shapes, input_dtypes, fallback_dtype=None)

    return graph, input_shapes, input_dtypes, output_shapes, output_dtypes

def graph_to_function(graph, target, ctx, shape=None, dtype=None):
    """Convert a graph to a function taking a keyword args and returning a list of results
    (both args and results are numpy arrays).

    Example::

        fun = graph_to_function(graph, llvm, cpu(0))
        [res1, res2] = fun(x=np.zeros((1,2)), y=np.zeros((1,)))

    Parameters
    ----------
    graph : nnvm.graph.Graph
        A graph we want to convert to a function.

    target : str or :any:`tvm.target.Target`
        The build target

    ctx : TVMContext
        The context to deploy the module.

    shape : Dict[str, Tuple[int]], optional
        A dict mapping input variable names to shapes.
        By default shapes will be inferred from variables' attributes.
        Note that this parameter takes precedence over variables' attributes.

    dtype : Dict[str, str] or str, optional
        A dict mapping input variable names to dtypes, or just a single dtype.
        By default dtypes will be inferred from variables' attributes.
        Note that this parameter takes precedence over variables' attributes.

    Returns
    -------
    function : Callable[..., List[numpy.ndarray]]
    """
    # Infer missing shapes and dtypes
    graph, shape, dtype, output_shapes, output_dtypes = \
        infer_shapes_dtypes(graph, shape=shape, dtype=dtype)

    if None in dtype.values():
        raise ValueError("Input variables with no type: {}".format(dtype))

    if not all(shape.values()):
        raise ValueError("Input variables with no shape: {}".format(shape))

    compute_graph, lib, params = nnvm.compiler.build(graph, target, shape=shape, dtype=dtype)
    module = graph_runtime.create(compute_graph, lib, ctx)

    if params:
        module.set_inputs(**params)

    def run(**kwargs):
        module.run(**kwargs)
        res = []
        for i, (o_shape, o_dtype) in enumerate(zip(output_shapes, output_dtypes)):
            res.append(module.get_output(i, tvm.nd.empty(o_shape, o_dtype)).asnumpy())
        return res

    return run

def _dict_var_to_dict_str(dictionary):
    """Convert a Dict[nnvm.Symbol, T] to Dict[str, T]"""
    if isinstance(dictionary, dict):
        return {s.attr('name') if isinstance(s, nnvm.symbol.Symbol) else s:
                dictionary[s] for s in dictionary}
    else:
        return dictionary

def check_function(symbol, forward=None, backward=None, grad_input_vars=None,
                   shape=None, dtype=None, in_range=None, values=None,
                   exclude_targets=None, only_targets=None,
                   additional_params=None,
                   numerical_grads=None, numerical_grads_params=None,
                   atol=1e-5, rtol=1e-5, quiet=False):
    """Compute the function and/or its gradients on a random input and raise
    an exception if the result doesn't match the reference implementation.

    Parameters
    ----------
    symbol : nnvm.Symbol
        A symbol representing the output.

    forward : Callable[..., List[numpy.ndarray]], optional
        A reference implementation to compare with.

    backward : Callable[..., List[numpy.ndarray] or Dict[str, numpy.ndarray]], optional
        A reference implementation of gradients. Should also accept head_grads besides
        normal inputs which is a list of gradients of some scalar wrt the outputs or just a
        single gradient if there are multiple outputs.
        Should return either a dict mapping input variable names to the respective
        gradients or a list of gradients wrt variables from grad_input_vars in
        exactly the same order (in alphabetical order by default).

    grad_input_vars : List[nnvm.Symbol or str], optional
        A list of variables with respect to which the gradients will be computed.
        None (default) means that all input variables will be used in an alphabetical order.

    shape : Dict[nnvm.Symbol or str, Tuple[int]] or Tuple[int], optional
        A dict mapping input variable names to shapes, or just a single shape.
        By default shapes will be inferred from variables' attributes (see the Examples).
        Note that this parameter takes precedence over variables' attributes.

    dtype : Dict[nnvm.Symbol or str, str] or str, optional
        A dict mapping input variable names to dtypes, or just a single dtype.
        By default dtypes will be inferred from variables' attributes (see the Examples).
        If dtypes cannot be inferred for some variables then float32 will be used as a fallback.
        Note that this parameter takes precedence over variables' attributes.

    in_range : Dict[nnvm.Symbol or str, (float, float)] or (float, float), optional
        A dict mapping input variable names to ranges or just a single range
        (the same for all variables). Input values will be generated from
        uniform distributions on these ranges. `head_grads` can also be
        assigned a range this way.

    values : Dict[nnvm.Symbol or str, numpy.ndarray], optional
        A dict explicitly providing values for some variables instead of random generation.

    exclude_targets : Set[str], optional
        Skip compiling and running anything for these targets.

    only_targets : Set[str], optional
        Test only for those targets from `ctx_list()` that are also in this set.

    additional_params : dict, optional
        A dict of additional parameters which will be passed to forward and backward.

    numerical_grads : bool or 'if_possible', optional
        Whether to additionally check against numerically computed gradients. If 'if_possible' or
        None is passed (which is the default) then it will try to create a gradient computation
        graph and then check gradients numerically only if this graph can be created (i.e. if there
        are some operations with unimplemented gradients, it will just issue a warning).
        Checking against numerical gradients is done via the `check_numerical_grads` function.

    numerical_grads_params : dict, optional
        Additional parameters for `check_numerical_grads`.

    atol : float, optional
        Absolute tolerance for `tvm.testing.assert_allclose`. NOT used for numerical gradients.

    rtol : float, optional
        Relative tolerance for `tvm.testing.assert_allclose`. NOT used for numerical gradients.

    quiet : bool, optional
        Don't dump additional information to stdout on failure.

    Examples
    --------
    .. code-block:: python

        x = sym.Variable("x", shape=(1, 2))
        y = sym.Variable("y", shape=(1, 2))

        # check the function and its gradients both numerically and using a reference function
        check_function(x + 2*y,
                       lambda x, y: x + 2*y,
                       lambda x, y, head_grads: {'x': head_grads, 'y': 2*head_grads})

        # just check gradients numerically
        check_function(x + 2*y, numerical_grads=True)

        # just check the forward computation
        check_function(x + 2*y, lambda x, y: x + 2*y, numerical_grads=False)

        # specifying dtype
        check_function(x + 2*y, lambda x, y: x + 2*y, dtype='float64')

        # dtypes can also be specified during variable creation with dtype codes
        x = sym.Variable("x", dtype=0)
        check_function(x + 1, shape=(2, 2), numerical_grads=True)
    """
    # validate and preprocess the input params
    if numerical_grads is None and forward is None and backward is None:
        raise ValueError("No reference function was passed to check_function. If you only want to "
                         "check gradients numerically, pass numerical_grads=True explicitly.")

    if numerical_grads is None:
        numerical_grads = 'if_possible'

    if numerical_grads not in [False, True, 'if_possible']:
        raise ValueError("numerical_grads must be a bool or 'if_possible', not {}"
                         .format(numerical_grads))

    if additional_params is None:
        additional_params = {}

    input_vars = symbol.list_input_variables()
    input_dict = {x.attr('name'): x for x in input_vars}

    if grad_input_vars is None:
        grad_input_vars = sorted(input_vars, key=lambda x: x.attr('name'))
    else:
        grad_input_vars = [input_dict[x] if isinstance(x, str) else x for x in grad_input_vars]

    in_range = _dict_var_to_dict_str(in_range)
    values = _dict_var_to_dict_str(values)

    out_len = len(symbol.list_output_names())

    # Infer the output shapes and dtypes, and preprocess the shape and dtype params
    forward_graph, shape, dtype, out_shapes, out_dtypes = \
        infer_shapes_dtypes(nnvm.graph.create(symbol), shape=shape, dtype=dtype,
                            fallback_dtype='float32')

    if not all(out_shapes) or not all(out_dtypes):
        if not quiet:
            print(forward_graph.ir(join_node_attrs=['shape', 'dtype']))
        raise ValueError("Could not infer shapes or dtypes for outputs.\n"
                         "out_shapes = {}\nout_dtypes = {}".format(out_shapes, out_dtypes))

    backward_graph = None

    # If we want gradients, we have to recreate the graph, but now with gradient computations
    # Note that here we need out_shapes for defining the shape of head grads, so we have to
    # create the graph twice
    if backward is not None or numerical_grads:
        try:
            head_grads_symbols = [nnvm.symbol.Variable("head_grads_" + str(i),
                                                       shape=out_shapes[i],
                                                       dtype=DTYPE_TO_TCODE[out_dtypes[i]])
                                  for i in range(out_len)]
            grad_symbols = graph_util.gradients([symbol], grad_input_vars,
                                                grad_ys=head_grads_symbols)
            # Sometimes grads do not depend on head_grads, so head_grads does not appear
            # in the variable list; adding it manually prevents this, making things a bit easier
            backward_graph = \
                nnvm.graph.create(nnvm.symbol.Group([symbol] + grad_symbols + head_grads_symbols))

            backward_graph, shape, dtype, out_shapes, out_dtypes = \
                infer_shapes_dtypes(backward_graph, shape=shape, dtype=dtype,
                                    fallback_dtype='float32')
        except nnvm._base.NNVMError as err:
            if backward is None and numerical_grads == "if_possible":
                logging.warning("Won't check gradients because: %s", str(err).split('\n', 1)[0])
                numerical_grads = False
                backward_graph = None
            else:
                raise

    main_graph = backward_graph if backward_graph is not None else forward_graph

    # Generate random data for inputs (including head_grads)

    np_inputs = {}

    for x in main_graph.symbol.list_input_variables():
        x_name = x.attr('name')
        x_shape = shape[x_name]
        x_dtype = dtype[x_name]

        if values is not None and x_name in values:
            np_inputs[x_name] = values[x_name].astype(x_dtype)
            continue

        low = -1.0
        high = 1.0
        if in_range is not None:
            if isinstance(in_range, dict):
                if x_name in in_range:
                    low = in_range[x_name][0]
                    high = in_range[x_name][1]
            else:
                low = in_range[0]
                high = in_range[1]

        np_inputs[x_name] = np.random.uniform(size=x_shape, low=low, high=high).astype(x_dtype)

    np_inputs_without_head_grads = {k: np_inputs[k] for k in np_inputs
                                    if not k.startswith('head_grads_')}

    nothing_was_done = True

    # Compute and compare the results
    for target, ctx in ctx_list():
        if exclude_targets is not None:
            if target in exclude_targets or str(target) in exclude_targets:
                logging.info("Skipping target = %s, ctx = %s", target, ctx)
                continue
        if only_targets is not None:
            if target not in only_targets and str(target) not in only_targets:
                logging.info("Skipping target = %s, ctx = %s", target, ctx)
                continue

        logging.info("Checking computation on target = %s, ctx = %s", target, ctx)

        debug_stage = None

        try:
            nnvm_res = None

            debug_stage = "compiling"
            main_function = graph_to_function(main_graph, target, ctx)

            # nnvm_res contains the output and gradients (if they are needed)
            debug_stage = "running"
            nnvm_res = main_function(**np_inputs)

            try:
                logging.debug("checking to_relay conversion")
                inputs = np_inputs_without_head_grads.copy()
                func, inputs = to_relay(main_graph, shape, dtype, params=inputs)
                with relay.build_config(opt_level=3):
                    graph, lib, params = relay.build(func, target=target)
                m = graph_runtime.create(graph, lib, ctx)
                m.set_input(**inputs)
                m.set_input(**params)
                m.run()
                for i in range(out_len):
                    relay_out = m.get_output(i).asnumpy()
                    tvm.testing.assert_allclose(nnvm_res[i], relay_out, atol=atol, rtol=rtol)
            except NotImplementedError as err:
                # the NNVM operator is not supported yet
                logging.warning(err)

            if backward_graph is not None:
                grad_var_names = [x.attr('name') for x in grad_input_vars]
                nnvm_grads = {x: v for x, v in zip(grad_var_names, nnvm_res[out_len:])}

            if forward is not None:
                nothing_was_done = False
                debug_stage = "checking forward computation"
                logging.debug(debug_stage)

                params = {}
                params.update(np_inputs_without_head_grads)
                params.update(additional_params)
                numpy_res = forward(**params)

                if isinstance(numpy_res, tuple):
                    numpy_res = list(numpy_res)

                if not isinstance(numpy_res, list):
                    numpy_res = [numpy_res]

                if len(numpy_res) != out_len:
                    raise ValueError("Forward function returned {} values, but "
                                     "the nnvm graph returns {} values"
                                     .format(len(numpy_res), out_len))

                for i in range(out_len):
                    tvm.testing.assert_allclose(nnvm_res[i], numpy_res[i], atol=atol, rtol=rtol)

            if backward is not None:
                nothing_was_done = False
                debug_stage = "checking gradients"
                logging.debug(debug_stage)

                np_head_grads = [np_inputs["head_grads_" + str(i)] for i in range(out_len)]

                if out_len == 1:
                    np_head_grads = np_head_grads[0]

                params = {'head_grads': np_head_grads}
                params.update(np_inputs_without_head_grads)
                params.update(additional_params)
                numpy_grads = backward(**params)

                if not isinstance(numpy_grads, dict):
                    if isinstance(numpy_grads, tuple):
                        numpy_grads = list(numpy_grads)
                    if not isinstance(numpy_grads, list):
                        numpy_grads = [numpy_grads]
                    numpy_grads = {x: v for x, v in zip(grad_var_names, numpy_grads)}
                    if len(numpy_grads) != len(grad_var_names):
                        raise ValueError("The backward function returns a list of gradients which "
                                         "does not contain gradients for these variables: {}"
                                         .format(set(grad_var_names) - set(numpy_grads)))

                for x_name in numpy_grads:
                    tvm.testing.assert_allclose(nnvm_grads[x_name], numpy_grads[x_name],
                                                atol=atol, rtol=rtol)

            if numerical_grads:
                nothing_was_done = False
                debug_stage = "checking gradients numerically"
                logging.debug(debug_stage)

                forward_function = graph_to_function(forward_graph, target, ctx)

                # Since the result may be non-scalar, we have to put another operation on the top,
                # so we just multiple by the randomly generated head_grads and then sum everything.
                # This way we can reuse the gradient values which has been already computed.
                def scalar_function(**kwargs):
                    res = forward_function(**kwargs)
                    return np.sum([np.dot(np_inputs['head_grads_' + str(i)].ravel(), res[i].ravel())
                                   for i in range(out_len)])

                if numerical_grads_params is None:
                    numerical_grads_params = {}

                check_numerical_grads(
                    scalar_function,
                    input_values=np_inputs_without_head_grads,
                    grad_values=nnvm_grads,
                    **numerical_grads_params)

        except:
            if not quiet:
                print("\ncheck_function failed while {}, here is the main graph"
                      .format(debug_stage))
                print(main_graph.ir(join_node_attrs=['shape', 'dtype']))
                if nnvm_res is not None:
                    print("Generated inputs:")
                    print(np_inputs)
                    print()
            raise

    if nothing_was_done:
        logging.warning("Nothing was done in check_function. Check ctx_list().")
