# 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 copy

import numpy as np

import megengine.functional as F
from megengine.core import Function, tensor
from megengine.jit import trace
from megengine.test import assertTensorClose


def test_a_plus_b():
    data_shape = (1, 9, 2, 6)
    av = np.random.random(data_shape).astype(np.float32)
    bv = np.random.random(data_shape).astype(np.float32)
    a = tensor(av)
    b = tensor(bv)

    class MulFunc(Function):
        def forward(self, a, b):
            return a * b

        def backward(self, grad_o):
            return (grad_o * b * 2, grad_o * a * 3)

    c = MulFunc()(a, b).sum()
    assertTensorClose(c.numpy(), (av * bv).sum())
    assertTensorClose(F.grad(c, a, use_virtual_grad=False).numpy(), bv * 2)
    assertTensorClose(F.grad(c, b, use_virtual_grad=False).numpy(), av * 3)


def test_skip_invalid_grad():
    data_shape = (1, 9, 2, 6)
    av = np.random.random(data_shape).astype(np.float32)
    bv = np.random.random(data_shape).astype(np.float32)
    a = tensor(av)
    b = tensor(bv)
    cookie = tensor(np.random.random(data_shape).astype(np.float32))

    class EqWithFakeGrad(Function):
        def forward(self, a, b):
            return a == b

        def backward(self, grad_o):
            _ = grad_o
            return cookie, cookie

    c = EqWithFakeGrad()(a, b).sum()
    assertTensorClose(c.numpy(), (av == bv).sum().astype(np.float32))
    assertTensorClose(F.grad(c, a, use_virtual_grad=False).numpy(), cookie)
    assertTensorClose(F.grad(c, b, use_virtual_grad=False).numpy(), cookie)


def test_ste():
    class STE(Function):
        def forward(self, x):
            maxv, minv = x.max(), x.min()
            scale = F.maximum(maxv, -minv) / 127
            return F.round(x / scale) * scale

        def backward(self, grad_y):
            return grad_y

    data_shape = (1, 9, 2, 6)
    av = np.random.random(data_shape).astype(np.float32)
    a = tensor(av)
    q = STE()(a)
    q_2 = (q * 2.0).sum()
    assertTensorClose(
        
F.grad(q_2, a, use_virtual_grad=False)