#!/usr/bin/env python3

# Copyright (c) 2023 CINN Authors. All Rights Reserved.
#
# Licensed 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 unittest
import numpy as np
from op_test import OpTest, OpTestTool
import paddle
import cinn
from cinn.frontend import *
from cinn.common import *


def cinn_dtype_convert(dtype_str):
    if dtype_str == "float32":
        return Float(32)
    elif dtype_str == "int64":
        return Int(64)
    elif dtype_str == "bool":
        return Bool()
    else:
        print("Datatype %s has not been supported yet", dtype_str)


##################################
####  TestElementwiseAddGrad  ####
##################################
# 1) x is 0D, y is 0D
@OpTestTool.skip_if(not is_compiled_with_cuda(),
                    "x86 test will be skipped due to timeout.")
class TestElementwiseAddGrad(OpTest):
    def setUp(self):
        np.random.seed(2023)
        self.init_input()

    def init_input(self):
        self.inputs = {
            "x": np.random.randint(-10, 10, []).astype("float32"),
            "y": np.random.randint(-10, 10, []).astype("float32"),
            "dout": np.random.randint(-10, 10, []).astype("float32"),
        }

    def build_paddle_program(self, target):
        x = paddle.to_tensor(self.inputs["x"], stop_gradient=False)
        y = paddle.to_tensor(self.inputs["y"], stop_gradient=False)
        out = paddle.add(x, y)

        self.paddle_outputs = [out]
        self.paddle_grads = self.get_paddle_grads([out], [x, y],
                                                  [self.inputs["dout"]])

    def build_cinn_program(self, target):
        builder = NetBuilder("add")
        x = builder.create_input(Float(32), self.inputs["x"].shape, "x")
        y = builder.create_input(Float(32), self.inputs["y"].shape, "y")
        # Test elementwise_add here, next unittest tests add, actually these two APIs are same.
        out = builder.elementwise_add(x, y)

        dout = builder.create_input(
            Float(32), self.inputs["dout"].shape, "dout")
        x_grad, y_grad = builder.elementwise_add_grad(dout, x, y)

        prog = builder.build()
        res = self.get_cinn_output(
            prog, target, [x, y, dout],
            [self.inputs["x"], self.inputs["y"], self.inputs["dout"]],
            [out, x_grad, y_grad])

        out, x_grad, y_grad = res
        self.cinn_outputs = [out]
        self.cinn_grads = [x_grad, y_grad]
        self.assertEqual(out.shape, self.inputs["dout"].shape)
        self.assertEqual(x_grad.shape, self.inputs["x"].shape)
        self.assertEqual(y_grad.shape, self.inputs["y"].shape)

    def test_check_results(self):
        self.check_outputs_and_grads()


# 2) x is ND, y is 0D
# NOTE: CINN only supports x's rank >= y's rank, hence no need to test next scenario: `3) x is 0D, y is ND`
class TestElementwiseAddGrad1(TestElementwiseAddGrad):
    def init_input(self):
        self.inputs = {
            "x": np.random.randint(-10, 10, [3, 5]).astype("float32"),
            "y": np.random.randint(-10, 10, []).astype("float32"),
            "dout": np.random.randint(-10, 10, [3, 5]).astype("float32"),
        }


##################################
#### TestElementwiseBinaryOp  ####
##################################
@OpTestTool.skip_if(not is_compiled_with_cuda(),
                    "x86 test will be skipped due to timeout.")
class TestElementwiseBinaryOp_0DTo0D(OpTest):
    def setUp(self):
        np.random.seed(2023)
        self.init_dtype()
        self.init_input()

    def init_dtype(self):
        self.dtype = "float32"

    def init_input(self):
        self.inputs = {
            "x": np.random.randint(-10, 10, []).astype(self.dtype),
            "y": np.random.randint(-10, 10, []).astype(self.dtype),
        }
        self.target_shape = ()

    def paddle_func(self, x, y):
        return paddle.add(x, y)

    def cinn_func(self, builder, x, y):
        return builder.add(x, y)

    def build_paddle_program(self, target):
        x = paddle.to_tensor(self.inputs["x"], stop_gradient=False)
        y = paddle.to_tensor(self.inputs["y"], stop_gradient=False)
        out = self.paddle_func(x, y)

        self.paddle_outputs = [out]

    def build_cinn_program(self, target):
        builder = NetBuilder("binary_op")
        x = builder.create_input(
            cinn_dtype_convert(self.dtype), self.inputs["x"].shape, "x")
        y = builder.create_input(
            cinn_dtype_convert(self.dtype), self.inputs["y"].shape, "y")
        out = self.cinn_func(builder, x, y)

        prog = builder.build()
        res = self.get_cinn_output(prog, target, [x, y],
                                   [self.inputs["x"], self.inputs["y"]], [out])

        self.cinn_outputs = res
        self.assertEqual(res[0].shape, self.target_shape)

    def test_check_results(self):
        self.check_outputs_and_grads()


class TestElementwiseBinaryOp_NdTo0d(TestElementwiseBinaryOp_0DTo0D):
    def init_input(self):
        self.inputs = {
            "x": np.random.randint(-10, 10, [3, 5]).astype(self.dtype),
            "y": np.random.randint(-10, 10, []).astype(self.dtype),
        }
        self.target_shape = (3, 5)


def create_unit_test(parent,
                     test_name,
                     fn_paddle,
                     fn_cinn,
                     dtype="float32",
                     **kwargs):
    @OpTestTool.skip_if(not is_compiled_with_cuda(),
                        "x86 test will be skipped due to timeout.")
    class TestClass(parent):
        def setUp(self):
            super().setUp()
            for k, v in kwargs.items():
                setattr(self, k, v)

        def init_dtype(self):
            self.dtype = dtype

        def paddle_func(self, *args):
            return fn_paddle(*args)

        def cinn_func(self, builder, *args):
            return eval(fn_cinn)(*args)

    cls_name = "{}_{}".format(parent.__name__, test_name)
    TestClass.__name__ = cls_name
    globals()[cls_name] = TestClass


# NOTE: CINN only supports x's rank >= y's rank, hence no need to test scenario: x is 0D, y is ND
create_unit_test(TestElementwiseBinaryOp_0DTo0D, "sub", paddle.subtract,
                 "builder.subtract")
create_unit_test(TestElementwiseBinaryOp_NdTo0d, "sub", paddle.subtract,
                 "builder.subtract")
create_unit_test(TestElementwiseBinaryOp_0DTo0D, "mul1", paddle.multiply,
                 "builder.multiply")
create_unit_test(TestElementwiseBinaryOp_NdTo0d, "mul1", paddle.multiply,
                 "builder.multiply")
create_unit_test(TestElementwiseBinaryOp_0DTo0D, "mul2", paddle.multiply,
                 "builder.elementwise_mul")
create_unit_test(TestElementwiseBinaryOp_NdTo0d, "mul2", paddle.multiply,
                 "builder.elementwise_mul")
create_unit_test(TestElementwiseBinaryOp_0DTo0D, "div", paddle.divide,
                 "builder.divide")
create_unit_test(TestElementwiseBinaryOp_NdTo0d, "div", paddle.divide,
                 "builder.divide")
# # Paddle'atan2 only supports 0D + 0D -> 0D
create_unit_test(TestElementwiseBinaryOp_0DTo0D, "atan2", paddle.atan2,
                 "builder.atan2")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "floor_divide",
    paddle.floor_divide,
    "builder.floor_divide",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "floor_divide",
    paddle.floor_divide,
    "builder.floor_divide",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "mod",
    paddle.mod,
    "builder.mod",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "mod",
    paddle.mod,
    "builder.mod",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "remainder",
    paddle.remainder,
    "builder.remainder",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "remainder",
    paddle.remainder,
    "builder.remainder",
    dtype="int64")
create_unit_test(TestElementwiseBinaryOp_0DTo0D, "max", paddle.maximum,
                 "builder.max")
create_unit_test(TestElementwiseBinaryOp_NdTo0d, "max", paddle.maximum,
                 "builder.max")
create_unit_test(TestElementwiseBinaryOp_0DTo0D, "min", paddle.minimum,
                 "builder.min")
create_unit_test(TestElementwiseBinaryOp_NdTo0d, "min", paddle.minimum,
                 "builder.min")
create_unit_test(TestElementwiseBinaryOp_0DTo0D, "pow", paddle.pow,
                 "builder.pow")
create_unit_test(TestElementwiseBinaryOp_NdTo0d, "pow", paddle.pow,
                 "builder.pow")

create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "logical_and",
    paddle.logical_and,
    "builder.logical_and",
    dtype="bool")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "logical_and",
    paddle.logical_and,
    "builder.logical_and",
    dtype="bool")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "logical_or",
    paddle.logical_or,
    "builder.logical_or",
    dtype="bool")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "logical_or",
    paddle.logical_or,
    "builder.logical_or",
    dtype="bool")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "logical_xor",
    paddle.logical_xor,
    "builder.logical_xor",
    dtype="bool")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "logical_xor",
    paddle.logical_xor,
    "builder.logical_xor",
    dtype="bool")

create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "bitwise_and",
    paddle.bitwise_and,
    "builder.bitwise_and",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "bitwise_and",
    paddle.bitwise_and,
    "builder.bitwise_and",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "bitwise_or",
    paddle.bitwise_or,
    "builder.bitwise_or",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "bitwise_or",
    paddle.bitwise_or,
    "builder.bitwise_or",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "bitwise_xor",
    paddle.bitwise_xor,
    "builder.bitwise_xor",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "bitwise_xor",
    paddle.bitwise_xor,
    "builder.bitwise_xor",
    dtype="int64")

create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "equal",
    paddle.equal,
    "builder.equal",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "equal",
    paddle.equal,
    "builder.equal",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "not_equal",
    paddle.not_equal,
    "builder.not_equal",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "not_equal",
    paddle.not_equal,
    "builder.not_equal",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "greater_than",
    paddle.greater_than,
    "builder.greater_than",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "greater_than",
    paddle.greater_than,
    "builder.greater_than",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "greater_equal",
    paddle.greater_equal,
    "builder.greater_equal",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "greater_equal",
    paddle.greater_equal,
    "builder.greater_equal",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "less_than",
    paddle.less_than,
    "builder.less_than",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "less_than",
    paddle.less_than,
    "builder.less_than",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_0DTo0D,
    "less_equal",
    paddle.less_equal,
    "builder.less_equal",
    dtype="int64")
create_unit_test(
    TestElementwiseBinaryOp_NdTo0d,
    "less_equal",
    paddle.less_equal,
    "builder.less_equal",
    dtype="int64")


######################
#### TestUnaryOp  ####
######################
@OpTestTool.skip_if(not is_compiled_with_cuda(),
                    "x86 test will be skipped due to timeout.")
class TestUnaryOp(OpTest):
    def setUp(self):
        np.random.seed(2023)
        self.init_dtype()
        self.init_input()

    def init_dtype(self):
        self.dtype = "float32"

    def init_input(self):
        self.inputs = {
            "x": np.random.uniform(0.0, 1.0, []).astype(self.dtype),
        }
        self.target_shape = ()

    def paddle_func(self, x):
        return paddle.sqrt(x)

    def cinn_func(self, builder, x):
        return builder.sqrt(x)

    def build_paddle_program(self, target):
        x = paddle.to_tensor(self.inputs["x"], stop_gradient=False)
        out = self.paddle_func(x)

        self.paddle_outputs = [out]

    def build_cinn_program(self, target):
        builder = NetBuilder("unary_op")
        x = builder.create_input(
            cinn_dtype_convert(self.dtype), self.inputs["x"].shape, "x")
        out = self.cinn_func(builder, x)

        prog = builder.build()
        res = self.get_cinn_output(prog, target, [x], [self.inputs["x"]],
                                   [out])

        self.cinn_outputs = res
        self.assertEqual(res[0].shape, self.target_shape)

    def test_check_results(self):
        self.check_outputs_and_grads()


create_unit_test(TestUnaryOp, "tanh", paddle.tanh, "builder.tanh")
create_unit_test(TestUnaryOp, "relu", paddle.nn.functional.relu,
                 "builder.relu")
create_unit_test(TestUnaryOp, "gelu", paddle.nn.functional.gelu,
                 "builder.gelu")
create_unit_test(TestUnaryOp, "sigmoid", paddle.nn.functional.sigmoid,
                 "builder.sigmoid")
create_unit_test(TestUnaryOp, "exp", paddle.exp, "builder.exp")
create_unit_test(TestUnaryOp, "erf", paddle.erf, "builder.erf")
create_unit_test(TestUnaryOp, "rsqrt", paddle.rsqrt, "builder.rsqrt")
create_unit_test(TestUnaryOp, "log", paddle.log, "builder.log")
create_unit_test(TestUnaryOp, "log2", paddle.log2, "builder.log2")
create_unit_test(TestUnaryOp, "log10", paddle.log10, "builder.log10")
create_unit_test(TestUnaryOp, "floor", paddle.floor, "builder.floor")
create_unit_test(TestUnaryOp, "ceil", paddle.ceil, "builder.ceil")
create_unit_test(TestUnaryOp, "round", paddle.round, "builder.round")
create_unit_test(TestUnaryOp, "trunc", paddle.trunc, "builder.trunc")
create_unit_test(TestUnaryOp, "sin", paddle.sin, "builder.sin")
create_unit_test(TestUnaryOp, "cos", paddle.cos, "builder.cos")
create_unit_test(TestUnaryOp, "tan", paddle.tan, "builder.tan")
create_unit_test(TestUnaryOp, "sinh", paddle.sinh, "builder.sinh")
create_unit_test(TestUnaryOp, "cosh", paddle.cosh, "builder.cosh")
create_unit_test(TestUnaryOp, "asin", paddle.asin, "builder.asin")
create_unit_test(TestUnaryOp, "acos", paddle.acos, "builder.acos")
create_unit_test(TestUnaryOp, "atan", paddle.atan, "builder.atan")
create_unit_test(TestUnaryOp, "asinh", paddle.asinh, "builder.asinh")
create_unit_test(TestUnaryOp, "atanh", paddle.atanh, "builder.atanh")
create_unit_test(TestUnaryOp, "isnan", paddle.isnan, "builder.is_nan")
create_unit_test(TestUnaryOp, "isfinite", paddle.isfinite, "builder.is_finite")
create_unit_test(TestUnaryOp, "isinf", paddle.isinf, "builder.is_inf")
create_unit_test(
    TestUnaryOp,
    "logical_not",
    paddle.logical_not,
    "builder.logical_not",
    dtype="bool")
create_unit_test(
    TestUnaryOp,
    "bitwise_not",
    paddle.bitwise_not,
    "builder.bitwise_not",
    dtype="int64")
create_unit_test(TestUnaryOp, "negative", paddle.neg, "builder.negative")
create_unit_test(TestUnaryOp, "sign", paddle.sign, "builder.sign")
create_unit_test(TestUnaryOp, "abs", paddle.abs, "builder.abs")
create_unit_test(TestUnaryOp, "reciprocal", paddle.reciprocal,
                 "builder.reciprocal")


# acosh requires input value > 1.0, specific init_input instead of using create_unit_test
class TestUnaryOp_acosh(TestUnaryOp):
    def init_input(self):
        self.inputs = {
            "x": np.random.uniform(1.0, 10.0, []).astype(self.dtype),
        }
        self.target_shape = ()

    def paddle_func(self, x):
        return paddle.acosh(x)

    def cinn_func(self, builder, x):
        return builder.acosh(x)


#######################
#### TestSundryOp  ####
#######################
@OpTestTool.skip_if(not is_compiled_with_cuda(),
                    "x86 test will be skipped due to timeout.")
class TestScaleOp(OpTest):
    def setUp(self):
        np.random.seed(2023)
        self.dtype = "float32"
        self.init_input()

    def init_input(self):
        self.inputs = {
            "x": np.random.randint(-10, 10, []).astype(self.dtype),
        }
        self.target_shape = ()

    def build_paddle_program(self, target):
        x = paddle.to_tensor(self.inputs["x"], stop_gradient=False)
        out = paddle.scale(x, scale=2.0, bias=1.0)

        self.paddle_outputs = [out]

    def build_cinn_program(self, target):
        builder = NetBuilder("reduce_op")
        x = builder.create_input(
            cinn_dtype_convert(self.dtype), self.inputs["x"].shape, "x")
        out = builder.scale(x, 2.0, 1.0)

        prog = builder.build()
        res = self.get_cinn_output(prog, target, [x], [self.inputs["x"]],
                                   [out])

        self.cinn_outputs = res
        self.assertEqual(res[0].shape, self.target_shape)

    def test_check_results(self):
        self.check_outputs_and_grads()


@OpTestTool.skip_if(not is_compiled_with_cuda(),
                    "x86 test will be skipped due to timeout.")
class TestSumOp(OpTest):
    def setUp(self):
        np.random.seed(2023)
        self.dtype = "float32"
        self.init_input()

    def init_input(self):
        self.inputs = {
            "x": np.random.randint(-10, 10, []).astype(self.dtype),
            "y": np.random.randint(-10, 10, []).astype(self.dtype),
        }
        self.target_shape = ()

    def build_paddle_program(self, target):
        x = paddle.to_tensor(self.inputs["x"], stop_gradient=False)
        y = paddle.to_tensor(self.inputs["y"], stop_gradient=False)
        out = x + y

        self.paddle_outputs = [out]

    def build_cinn_program(self, target):
        builder = NetBuilder("reduce_op")
        x = builder.create_input(
            cinn_dtype_convert(self.dtype), self.inputs["x"].shape, "x")
        y = builder.create_input(
            cinn_dtype_convert(self.dtype), self.inputs["y"].shape, "y")
        out = builder.sum([x, y])

        prog = builder.build()
        res = self.get_cinn_output(prog, target, [x, y],
                                   [self.inputs["x"], self.inputs["y"]], [out])

        self.cinn_outputs = res
        self.assertEqual(res[0].shape, self.target_shape)

    def test_check_results(self):
        self.check_outputs_and_grads()


@OpTestTool.skip_if(not is_compiled_with_cuda(),
                    "x86 test will be skipped due to timeout.")
class TestDropoutOp(OpTest):
    def setUp(self):
        np.random.seed(2023)
        self.dtype = "float32"
        self.init_input()

    def init_input(self):
        self.inputs = {
            "x": np.random.randint(-10, 10, []).astype(self.dtype),
        }
        self.target_shape = ()

    def build_paddle_program(self, target):
        x = paddle.to_tensor(self.inputs["x"], stop_gradient=False)
        out = paddle.nn.functional.dropout(x, 1.0)

        self.paddle_outputs = [out]

    def build_cinn_program(self, target):
        builder = NetBuilder("reduce_op")
        x = builder.create_input(
            cinn_dtype_convert(self.dtype), self.inputs["x"].shape, "x")
        out = builder.dropout_infer(x, 1.0)

        prog = builder.build()
        res = self.get_cinn_output(prog, target, [x], [self.inputs["x"]],
                                   [out])

        self.cinn_outputs = res
        self.assertEqual(res[0].shape, self.target_shape)

    def test_check_results(self):
        self.check_outputs_and_grads()


if __name__ == "__main__":
    unittest.main()
