# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

import logging
import inspect

import mxnet as mx
from mxnet import nd
from mxnet import np

from .ndarray_utils import get_mx_ndarray, nd_forward_and_profile, nd_forward_backward_and_profile
from .common_utils import merge_map_list
from .op_registry_utils import prepare_op_inputs
from benchmark.opperf.rules.default_params import PARAMS_OF_TYPE_NDARRAY, PARAMS_OF_TYPE_NP_ARRAY
from .profiler_utils import cpp_profile, python_profile

no_backward = {'gather_nd', 'softmax_cross_entropy', 'linalg_gelqf', 'linalg_slogdet', 'moments', 'SequenceLast', 'Embedding'}

def _prepare_op_inputs(inputs, run_backward, dtype, ctx, module):
    mx.random.seed(41)
    kwargs_list = []
    if module == 'mxnet.numpy_extension' or module == 'mxnet.numpy':
        PARAMS_TYPE = PARAMS_OF_TYPE_NP_ARRAY
        get_array_fn = get_mx_np_ndarray
    else:
        PARAMS_TYPE = PARAMS_OF_TYPE_NDARRAY
        get_array_fn = get_mx_ndarray

    for inp in inputs:
        kwargs = {}
        for key, value in inp.items():
            if key in PARAMS_TYPE:
                kwargs[key] = get_array_fn(ctx=ctx, in_tensor=value,
                                           dtype=dtype,
                                           initializer=nd.normal,
                                           attach_grad=run_backward)
            else:
                kwargs[key] = value
        kwargs_list.append(kwargs)
    return kwargs_list

def get_mx_np_ndarray(ctx, in_tensor, dtype, initializer, attach_grad=True):
    """Helper function to prepare a MXNet Numpy NDArray tensor in given Context (ctx) of type (dtype).
    You can get a new Tensor by providing only "Shape" or "Numpy NDArray" or another MXNet NDArray as
    "in_tensor".

    NOTE: This is a sync call and waits for the Tensor to be created.

    Parameters
    ----------
    ctx: mx.ctx, default mx.cpu()
        Context of the new MXNet NDArray Tensor.
    in_tensor: Numpy NDArray or MXNet NDArray or Tuple of shape
        Can be a tuple of shape or Numpy NDArray or MXNet NDArray.
    dtype: str
        Precision or Dtype of the expected Tensor. Ex: "float32", "Int64"
    initializer:
        Function reference to the initialize to use. Ex: mx.nd.random.normal, mx.nd.zeros
    attach_grad: Boolean, default True
        To attach a gradient for the Tensor. Default is True.

    Returns
    -------
    MXNet NDArray Tensor.
    """
    if isinstance(in_tensor, int) or isinstance(in_tensor, float):
        return in_tensor

    if isinstance(in_tensor, tuple):
        nd_ndarray = get_mx_ndarray(ctx=ctx, in_tensor=in_tensor,
                                             dtype="float32",
                                             initializer=initializer,
                                             attach_grad=attach_grad)
        tensor = nd_ndarray.as_np_ndarray().astype(dtype=dtype)
    elif isinstance(in_tensor, list):
        tensor = np.array(in_tensor, ctx=ctx)
    elif isinstance(in_tensor, nd.NDArray):
        tensor = in_tensor.as_np_ndarray()
    elif isinstance(in_tensor, np.ndarray):
        tensor = in_tensor.as_in_context(ctx)
    else:
        raise ValueError("Invalid input type for creating input tensor. Input can be tuple() of shape or Numpy Array or"
                         " MXNet NDArray. Given - ", in_tensor)
    if attach_grad:
        tensor.attach_grad()

    tensor.wait_to_read()
    return tensor

def parse_input_ndarray(input_dict):
    """Parse input for ndarray and extract array shape for better readability

    Parameters
    ----------
    input_dict : dict
         Dictionary of input

    Input Dictionary

    'inputs': {'weight':
    [[ 2.2122064   0.7740038   1.0434405   1.1839255   1.8917114 ]
     [-1.2347414  -1.771029   -0.45138445  0.57938355 -1.856082  ]
     [-1.9768796  -0.20801921  0.2444218  -0.03716067 -0.48774993]
     [-0.02261727  0.57461417  1.4661262   0.6862904   0.35496104]
     [ 1.0731696   0.12017461 -0.9711102  -0.77569664 -0.7882176 ]]
    <NDArray 5x5 @cpu(0)>, 'grad':
    [[ 0.7417728  -1.4734439  -1.0730928  -1.0424827  -1.3278849 ]
     [-1.4749662  -0.52414197  1.2662556   0.8950642  -0.6015945 ]
     [ 1.2040559  -0.9712193  -0.58256227  0.3717077   0.9300072 ]
     [-1.4225755  -0.5176199   2.0088325   0.2863085   0.5604595 ]
     [ 0.96975976 -0.52853745 -1.88909     0.65479124 -0.45481315]]
    <NDArray 5x5 @cpu(0)>, 'mean':
    [[ 0.32510808 -1.3002341   0.3679345   1.4534262   0.24154152]
     [ 0.47898006  0.96885103 -1.0218245  -0.06812762 -0.31868345]
     [-0.17634277  0.35655284  0.74419165  0.7787424   0.6087823 ]
     [ 1.0741756   0.06642842  0.8486986  -0.8003802  -0.16882208]
     [ 0.93632793  0.357444    0.77932847 -1.0103073  -0.39157307]]
    <NDArray 5x5 @cpu(0)>, 'var':
    [[ 1.3166187  -0.43292624  0.71535987  0.9254156  -0.90495086]
     [-0.074684    0.82254    -1.8785107   0.8858836   1.9118724 ]
     [ 0.33342266  0.11883813 -1.9198899  -0.67558455  1.007749  ]
     [-0.35391203  1.6323917  -0.33354783 -1.7378405   0.7737382 ]
     [ 0.89126545  3.2904532  -1.1976235   1.8938874  -0.5669272 ]]
    <NDArray 5x5 @cpu(0)>, 't': 1, 'wd': 0.1}

    Output
    {'inputs': {'weight': '<NDArray 5x5 @cpu(0)>', 'grad': '<NDArray 5x5 @cpu(0)>', 'mean': '<NDArray 5x5 @cpu(0)>', 'var': '<NDArray 5x5 @cpu(0)>', 't': 1, 'wd': 0.1}
    """
    no_new_line_input_dict=dict()
    for key,value in input_dict.items():
        if isinstance(value,nd.NDArray):
            # if value in input is NDArray then extract last line only
            val = str(value).split('\n')[-1]
            no_new_line_input_dict[key]=val
        else:
            no_new_line_input_dict[key]=value
    return no_new_line_input_dict


def _run_operator_performance_test(op, inputs, run_backward, warmup, runs, kwargs_list, profiler):
    if profiler == 'native':
        if run_backward:
            benchmark_helper_func = cpp_profile(nd_forward_backward_and_profile)
        else:
            benchmark_helper_func = cpp_profile(nd_forward_and_profile)
    elif profiler == 'python':
        if run_backward:
            benchmark_helper_func = python_profile(nd_forward_backward_and_profile)
        else: