# 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.
"""Test code for reduce."""
import os
import sys

import numpy as np
import pytest

import tvm
import tvm.testing
import tvm.topi.testing

from tvm import te, topi, tir

in_shape, axis, keepdims, reduce_type, dtype = tvm.testing.parameters(
    ((32,), 0, False, "argmax", "float32"),
    ((128, 24, 128, 24), (1, 2, 3), True, "sum", "float32"),
    ((2, 3), None, True, "all", "bool"),
    ((128, 24 * 128 * 24), (1,), False, "max", "float32"),
    ((32, 128, 24), None, True, "sum", "float32"),
    ((32, 128, 24), None, True, "all", "bool"),
    ((128, 24, 128, 24), (0, 2), False, "min", "float32"),
    ((32, 128), 1, True, "argmax", "float32"),
    ((32, 24, 32, 24), 2, False, "argmin", "float32"),
    ((31, 21, 15), None, True, "argmax", "float32"),
    ((31, 21, 15), None, False, "sum", "float32"),
    ((128, 24, 128, 24), (1, 2, 3), True, "sum", "float64"),
    ((2, 3), None, True, "any", "bool"),
    ((32, 128, 24), None, True, "any", "bool"),
    ((1, 4, 7), 1, True, "any", "bool"),
    ((128, 24, 128, 24), 2, False, "any", "bool"),
    ((128, 24, 128, 24), 2, False, "sum", "bool"),
    ((128, 24, 128, 24), 0, True, "sum", "bool"),
)


@tvm.testing.fixture(cache_return_value=True)
def ref_data(in_shape, axis, keepdims, reduce_type, dtype):
    # Test
    if dtype == "bool":
        in_npy_map = in_npy = np.random.choice([True, False], size=in_shape)
    else:
        in_npy = np.random.uniform(-1, 1, size=in_shape).astype(dtype)
        in_npy_map = np.sqrt(np.exp(in_npy)).astype(dtype)

    if reduce_type == "sum":
        if dtype == "bool":
            out_npy = in_npy_map.sum(axis=axis, keepdims=keepdims, dtype="bool")
        else:
            out_npy = in_npy_map.sum(axis=axis, keepdims=keepdims)
    elif reduce_type == "all" and dtype == "bool":
        out_npy = in_npy_map.all(axis=axis, keepdims=keepdims)
    elif reduce_type == "any" and dtype == "bool":
        out_npy = in_npy_map.any(axis=axis, keepdims=keepdims)
    elif reduce_type == "max":
        out_npy = in_npy_map.max(axis=axis, keepdims=keepdims)
    elif reduce_type == "min":
        out_npy = in_npy_map.min(axis=axis, keepdims=keepdims)
    elif reduce_type == "argmax":
        out_npy = _my_npy_argmax(in_npy_map, axis=axis, keepdims=keepdims)
    elif reduce_type == "argmin":
        out_npy = _my_npy_argmin(in_npy_map, axis=axis, keepdims=keepdims)
    else:
        raise NotImplementedError

    return in_npy, in_npy_map, out_npy


def _my_npy_argmax(arr, axis, keepdims):
    if not keepdims:
        return arr.argmax(axis=axis)
    else:
        if axis is None:
            out_shape = [1 for _ in arr.shape]
        else:
            out_shape = list(arr.shape)
            out_shape[axis] = 1

        return arr.argmax(axis=axis).reshape(out_shape)


def _my_npy_argmin(arr, axis, keepdims):
    if not keepdims:
        return arr.argmin(axis=axis)
    else:
        if axis is None:
            out_shape = [1 for _ in arr.shape]
        else:
            out_shape = list(arr.shape)
            out_shape[axis] = 1
        return arr.argmin(axis=axis).reshape(out_shape)


def test_reduce_map(target, dev, ref_data, in_shape, axis, keepdims, reduce_type, dtype):
    target = tvm.target.Target(target)
    if target.kind.name == "vulkan" and reduce_type in ["sum", "any", "all"]:
        pytest.xfail(f"Vulkan backend has known errors on {reduce_type}")

    in_npy, in_npy_map, out_npy = ref_data

    # Build the logic and compile the function
    A = te.placeholder(shape=in_shape, name="A", dtype=dtype)
    A1 = topi.sqrt(topi.exp(A))
    out_dtype = dtype
    if reduce_type == "sum":
        if dtype == "bool":
            B = topi.sum(A, axis=axis, keepdims=keepdims)
        else:
            B = topi.sum(A1, axis=axis, keepdims=keepdims)
    elif reduce_type == "all":
        B = topi.all(A, axis=axis, keepdims=keepdims)
    elif reduce_type == "any":
        B = topi.any(A, axis=axis, keepdims=keepdims)
    elif reduce_type == "max":
        B = topi.max(A1, axis=axis, keepdims=keepdims)
    elif reduce_type == "min":
        B = topi.min(A1, axis=axis, keepdims=keepdims)
    elif reduce_type == "argmax":
        B = topi.argmax(A1, axis=axis, keepdims=keepdims)
        out_dtype = "int32"
    elif reduce_type == "argmin":
        B = topi.argmin(A1, axis=axis, keepdims=keepdims)
        out_dtype = "int32"
    else:
        raise NotImplementedError

    with tvm.target.Target(target):
        s = tvm.topi.testing.get_reduce_schedule(target)(B)

    foo = tvm.build(s, [A, B], target, name=reduce_type)

    data_tvm = tvm.nd.array(in_npy, device=dev)
    out_tvm = tvm.nd.empty(shape=out_npy.shape, device=dev, dtype=out_dtype)
    foo(data_tvm, out_tvm)

    if reduce_type == "argmax" or reduce_type == "argmin":
        out_tvm_indices = out_tvm.numpy()
        if keepdims:
            out_tvm_indices = np.take(out_tvm_indices, indices=0, axis=axis)
        if axis is None:
            out_tvm_val = in_npy_map.ravel()[out_tvm_indices]
        else:
            other_indices = tuple(np.indices(in_shape[0:axis] + in_shape[(axis + 1) :]))
            sel_indices = other_indices[0:axis] + (out_tvm_indices,) + other_indices[axis:]
            out_tvm_val = in_npy_map[sel_indices]
        if reduce_type == "argmax":
            tvm.testing.assert_allclose(out_tvm_val, in_npy_map.max(axis=axis), 1e-3, 1e-3)
        elif reduce_type == "argmin":
            tvm.testing.assert_allclose(out_tvm_val, in_npy_map.min(axis=axis), 1e-3, 1e-3)
    else:
        tvm.testing.assert_allclose(out_tvm.numpy(), out_npy, 1e-3, 1e-3)


def test_complex_reduce(target, dev):
    in_shape = (2, 3)
    dtype = "float32"
    axis = 0
    keepdims = False
    A = te.placeholder(shape=in_shape, name="A", dtype=dtype)
    B = topi.sum(A, axis=axis, keepdims=keepdims)
    C = topi.add(B, B)
    D = topi.multiply(B, B)
    E = topi.add(C, D)

    with tvm.target.Target(target):
        s = tvm.topi.testing.get_reduce_schedule(target)(E)
    foo = tvm.build(s, [A, E], target, name="sum")

    in_npy = np.random.uniform(-1, 1, size=in_shape).astype(dtype)
    sum_npy = in_npy.sum(axis=axis, keepdims=keepdims)
    out_npy = sum_npy * 2 + sum_npy * sum_npy

    data_tvm = tvm.nd.array(in_npy, device=dev)
    out_tvm = tvm.nd.empty(shape=out_npy.shape, device=dev, dtype=dtype)
    foo(data_tvm, out_tvm)
    tvm.testing.assert_allclose(out_tvm.numpy(), out_npy, 1e-3, 1e-3)


n = tir.Var("n", "int32")
m = tir.Var("m", "int32")
true_value_map = {n: 3, m: 5}

data_shape, target_shape = tvm.testing.parameters(
    ((2, 3), (3,)),
    ((2, 3, 4), (2, 1, 4)),
    ((2, 3, 4, 5), (3, 1, 5)),
    ((2, n, 4, m), (n, 1, m)),
)


def _my_npy_collapse_sum(data, target_shape):
    reduce_axes = []
    i = data.ndim - 1
    j = len(target_shape) - 1
    while i >= 0:
        if j < 0:
            reduce_axes.append(i)
        elif target_shape[j] == 1 and data.shape[i] > 1:
            reduce_axes.append(i)
        i -= 1
        j -= 1
    return np.sum(data, tuple(reduce_axes)).reshape(target_shape)


def test_collapse_sum(data_shape, target_shape):
    A = te.placeholder(data_shape, name="A")
    B = topi.collapse_sum(A, target_shape)
    s = te.create_schedule([B.op])

    data_shape_const = [int(s) if s not in true_value_map else true_value_map[s] for s in A.shape]
    target_shape_const = [
        int(s) if s not in true_value_map else true_value_map[s] for s in target_shape
    ]
    a_np = np.random.uniform(size=data_shape_const).astype(A.dtype)
    b_np = _my_npy_collapse_sum(a_np, target_shape_const)
    dev = tvm.cpu(0)
    a = tvm.nd.array(a_np, dev)
    B_shape_const = [int(s) if s not in true_value_map else true_value_map[s] for s in B.shape]
    b = tvm.nd.array(np.zeros(B_shape_const, dtype=B.dtype), dev)
    # Building with the CSE pass disabled
    with tvm.transform.PassContext(opt_level=3, disabled_pass=["tir.CommonSubexprElimTIR"]):
        foo = tvm.build(s, [A, B], "llvm", name="collapse_sum")
    foo(a, b)
    tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-5)


if __name__ == "__main__":
    tvm.testing.main()