# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 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 itertools
import platform
from functools import partial

import numpy as np
import pytest
from utils import opr_test

import megengine.amp as amp
import megengine.config as config
import megengine.core.ops.builtin as builtin
import megengine.core.tensor.dtype as dtype
import megengine.functional as F
import megengine.jit as jit
from megengine import Parameter, Tensor, is_cuda_available, tensor
from megengine.core._trace_option import use_symbolic_shape
from megengine.core.autodiff.grad import Grad
from megengine.core.tensor.utils import make_shape_tuple
from megengine.device import get_device_count
from megengine.module import LayerNorm


def test_where():
    maskv0 = np.array([[1, 0], [0, 1]], dtype=np.bool_)
    xv0 = np.array([[1, np.inf], [np.nan, 4]], dtype=np.float32)
    yv0 = np.array([[5, 6], [7, 8]], dtype=np.float32)

    maskv1 = np.array([[1, 0, 1], [1, 0, 0], [1, 1, 0]], dtype=np.bool_)
    xv1 = np.array([[1, np.inf, 2], [0, np.nan, 4], [1, 5, 7]], dtype=np.float32)
    yv1 = np.array([[5, 6, 9], [2, 7, 8], [2, 1, 9]], dtype=np.float32)

    cases = [
        {"input": [maskv0, xv0, yv0]},
        {"input": [maskv1, xv1, yv1]},
    ]
    opr_test(cases, F.where, ref_fn=np.where, test_trace=False)

    maskv2 = np.array([1, 1, 1], dtype=np.bool_)
    xv2 = np.array([1, 3, 2], dtype=np.float32)
    yv2 = np.array([5, 6, 9], dtype=np.float32)

    maskv3 = np.array([0, 0, 0], dtype=np.bool_)
    xv3 = np.array([1, 3, 2], dtype=np.float32)
    yv3 = np.array([5, 6, 9], dtype=np.float32)

    cases = [
        {"input": [maskv2, xv2, yv2]},
        {"input": [maskv3, xv3, yv3]},
    ]
    opr_test(cases, F.where, ref_fn=np.where, test_trace=False)


def test_dropout():
    from megengine.autodiff import GradManager
    from megengine.core._imperative_rt.ops import set_global_rng_seed

    def test_dropout_with_shape(shape, rate):
        data = tensor(np.ones(shape, dtype=np.float32))
        gm = GradManager().attach([data])
        with gm:
            out = F.nn.dropout(data, rate, training=True)
            gm.backward(out, tensor(np.ones(shape, dtype=np.float32)))
            assert not out.numpy().all()
            np.testing.assert_allclose(out.numpy(), data.grad.numpy(), 1e-7, 1e-7)

    def test_multiple_dropout(shape, rate):
        data = tensor(np.ones(shape, dtype=np.float32))
        gm = GradManager().attach([data])
        with gm:
            out1 = F.nn.dropout(data, rate, training=True)
            out2 = F.nn.dropout(out1, rate, training=True)
            out3 = F.nn.dropout(out2, rate, training=True)
            gm.backward(out3, tensor(np.ones(shape, dtype=np.float32)))
            np.testing.assert_allclose(out3.numpy(), data.grad.numpy(), 1e-7, 1e-7)

    def test_dropout_seed(shape, rate):
        data = tensor(np.random.randn(*shape), dtype="float32")
        set_global_rng_seed(111)
        out1 = F.nn.dropout(data, rate, training=True)
        out2 = F.nn.dropout(data, rate, training=True)
        assert not (out1.numpy() == out2.numpy()).all()

        set_global_rng_seed(111)
        out3 = F.nn.dropout(data, rate, training=True)
        assert (out1.numpy() == out3.numpy()).all()

        set_global_rng_seed(222)
        out4 = F.nn.dropout(data, rate, training=True)
        assert not (out1.numpy() == out4.numpy()).all()

    test_dropout_with_shape([13, 17, 63, 21], 0.4)
    test_dropout_with_shape([16, 32, 64], 0.3)
    test_multiple_dropout([1024], 0.2)
    test_dropout_seed([16, 32], 0.2)


def test_matinv():
    shape1 = (5, 5)
    shape2 = (3, 9, 9)
    data1 = np.random.random(shape1).astype("float32")
    data2 = np.random.random(shape2).astype("float32")
    # make matrix diagonally dominant for numerical stability
    data1 += (np.eye(shape1[0]) * shape1[0]).astype("float32")
    data2 += np.broadcast_to((np.eye(shape2[1]) * shape2[1]).astype("float32"), shape2)

    cases = [
        {"input": data1},
        {"input": data2},
    ]

    opr_test(
        cases,
        F.matinv,
        compare_fn=lambda x, y: np.testing.assert_allclose(x.numpy(), y, rtol=1e-4),
        ref_fn=np.linalg.inv,
    )


def test_matmul():
    shape1 = 3
    shape2 = 3
    shape3 = (3, 5)
    shape4 = (5, 6)
    data1 = np.random.random(shape1).astype("float32")
    data2 = np.random.random(shape2).astype("float32")
    data3 = np.random.random(shape3).astype("float32")
    data4 = np.random.random(shape4).astype("float32")

    cases = [
        {"input": [data1, data2]},
        {"input": [data2, data3]},
        {"input": [data3, data4]},
    ]
    opr_test(cases, F.matmul, ref_fn=np.matmul)

    batch_size = 10
    shape1 = (2,)
    shape2 = (batch_size, 2, 3)
    shape3 = (batch_size, 3, 4)
    shape4 = (batch_size, 10, 4, 2)
    shape5 = (batch_size, 10, 2, 4)
    data1 = np.random.random(shape1).astype("float32")
    data2 = np.random.random(shape2).astype("float32")
    data3 = np.random.random(shape3).astype("float32")
    data4 = np.random.random(shape4).astype("float32")
    data5 = np.random.random(shape5).astype("float32")

    cases = [
        {"input": [data1, data2]},
        {"input": [data2, data3]},
        {"input": [data3, data4]},
        {"input": [data4, data5]},
    ]
    opr_test(cases, F.matmul, ref_fn=np.matmul)

    opr_test(
        [{"input": [data1, data4]}],
        F.matmul,
        ref_fn=lambda x, y: np.matmul(x, y.transpose(0, 1, 3, 2)),
        transpose_b=True,
    )

    opr_test(
        [{"input": [data3, data2]}],
        F.matmul,
        ref_fn=lambda x, y: np.matmul(x.transpose(0, 2, 1), y.transpose(0, 2, 1)),
        transpose_a=True,
        transpose_b=True,
    )


@pytest.mark.parametrize(
    "shape_a, shape_b", [((0,), (0,)), ((10, 0), (0, 10)), ((3, 10, 0), (3, 0, 10)),],
)
@pytest.mark.parametrize("is_symbolic", [None, True, False])
def test_matmul_empty_tensor(shape_a, shape_b, is_symbolic):
    def func(a, b):
        return F.matmul(a, b)

    if is_symbolic is not None:
        func = jit.trace(symbolic=is_symbolic)(func)

    a = tensor(np.random.randn(*shape_a))
    b = tensor(np.random.randn(*shape_b))
    for _ in range(3):
        out = func(a, b)
        assert np.all(out.numpy() == 0)
        if is_symbolic is None:
            break


def test_interpolate():
    def linear_interpolate():
        inp = tensor(np.arange(1, 3, dtype=np.float32).reshape(1, 1, 2))

        out = F.vision.interpolate(inp, scale_factor=2.0, mode="linear")
        out2 = 
F.vision.interpolate(inp, 4, mode="linear")