# 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 numpy as np
import pytest

import tvm
import tvm.testing
from tvm.script import tirx as T
from tvm.testing import env


@T.prim_func(s_tir=True)
def get_valid_counts(
    data: T.handle,
    valid_count: T.handle,
    out: T.handle,
    out_indices: T.handle,
    score_threshold: T.float32,
    id_index: T.int32,
    score_index: T.int32,
) -> None:
    data_buf = T.match_buffer(data, (1, 2500, 6), "float32")
    valid_count_buf = T.match_buffer(valid_count, (1,), "int32")
    out_buf = T.match_buffer(out, (1, 2500, 6), "float32")
    out_indices_buf = T.match_buffer(out_indices, (1, 2500), "int32")

    with T.sblock("init"):
        vi = T.axis.S(1, 0)
        valid_count_buf[vi] = T.int32(0)
        for j in range(2500):
            with T.sblock("update"):
                vj = T.axis.S(2500, j)
                T.reads([data_buf[vi, vj, 6]])
                T.writes([valid_count_buf[vi], out_indices_buf[vi, vj], out_buf[vi, vj, 6]])
                if (data_buf[vi, vj, score_index] > score_threshold) and (
                    (id_index < 0) or (data_buf[vi, vj, id_index] >= T.float32(0))
                ):
                    for k in T.serial(0, 6):
                        out_buf[vi, valid_count_buf[vi], k] = data_buf[vi, vj, k]
                    out_indices_buf[vi, valid_count_buf[vi]] = vj
                    valid_count_buf[vi] = valid_count_buf[vi] + 1
                if vj >= valid_count_buf[vi]:
                    for k in T.serial(0, 6):
                        out_buf[vi, vj, k] = T.float32(-1)
                    out_indices_buf[vi, vj] = T.int32(-1)


def _check_get_valid_counts_with_numpy(f, dshape, score_threshold, id_index, score_index):
    dtype = "float32"
    ctx = tvm.cpu()
    batch_size, num_anchor, elem_length = dshape
    np_data = np.random.uniform(low=-2, high=2, size=dshape).astype(dtype)
    np_out1 = np.zeros(shape=(batch_size,), dtype="int32")
    np_out2 = np.zeros(shape=dshape).astype(dtype)
    np_out3 = np.zeros(shape=(batch_size, num_anchor), dtype="int32")
    for i in range(batch_size):
        np_out1[i] = 0
        inter_idx = 0
        for j in range(num_anchor):
            score = np_data[i, j, score_index]
            if score > score_threshold and (id_index < 0 or np_data[i, j, id_index] >= 0):
                for k in range(elem_length):
                    np_out2[i, inter_idx, k] = np_data[i, j, k]
                np_out1[i] += 1
                np_out3[i, inter_idx] = j
                inter_idx += 1
            if j >= np_out1[i]:
                for k in range(elem_length):
                    np_out2[i, j, k] = -1.0
                np_out3[i, j] = -1

    in_data = tvm.runtime.tensor(np_data, ctx)
    out1 = tvm.runtime.tensor(np_out1, ctx)
    out2 = tvm.runtime.tensor(np_out2, ctx)
    out3 = tvm.runtime.tensor(np_out3, ctx)
    f(in_data, out1, out2, out3, score_threshold, id_index, score_index)
    tvm.testing.assert_allclose(out1.numpy(), np_out1, rtol=1e-5)
    tvm.testing.assert_allclose(out2.numpy(), np_out2, rtol=1e-5)
    tvm.testing.assert_allclose(out3.numpy(), np_out3, rtol=1e-5)
    print("test get_valid_counts end")


def test_get_valid_counts_script_func():
    device = "llvm"
    # check lowering
    print(get_valid_counts.script())
    mod = tvm.ir.IRModule({"get_valid_counts": get_valid_counts})
    print(mod.script())
    # check building
    f = tvm.compile(mod["get_valid_counts"], target=device)
    _check_get_valid_counts_with_numpy(f, (1, 2500, 6), 0.0, 0, 1)


@T.prim_func(s_tir=True)
def alloc_zero_dim_buffer(a: T.handle, b: T.handle) -> None:
    A = T.match_buffer(a, [], dtype="float32")
    B = T.match_buffer(b, [], dtype="float32")
    # body
    # tirx.with block("root")
    C = T.sblock_alloc_buffer([], dtype="float32")
    A[()] = T.float32(2)
    C[()] = A[()] + B[()]
    B[()] = C[()]


@T.prim_func(s_tir=True)
def alloc_zero_dim_buffer_block(a: T.handle, b: T.handle) -> None:
    A = T.match_buffer(a, (), "float32")
    B = T.match_buffer(b, (), "float32")
    with T.sblock("root"):
        T.reads([])
        T.writes([])
        C = T.sblock_alloc_buffer((), "float32")
        A[()] = T.float32(2)
        C[()] = A[()] + B[()]
        B[()] = C[()]


def _check_alloc_zero_dim_buffer(f):
    dtype = "float32"
    ctx = tvm.cpu()

    np_data = np.zeros(shape=()).astype(dtype)
    np_out = np.zeros(shape=()).astype(dtype)
    tvm_data = tvm.runtime.tensor(np_data, ctx)
    tvm_out = tvm.runtime.tensor(np_out, ctx)

    # np func exection
    np_inter = np.array(1)
    np_data[()] = 2.0
    np_inter[()] = np_data[()] + np_out[()]
    np_out[()] = np_inter[()]

    # tvm func execution
    f(tvm_data, tvm_out)
    tvm.testing.assert_allclose(tvm_out.numpy(), np_out, rtol=1e-5)


def test_alloc_zero_dim_buffer_round_trip():
    func = alloc_zero_dim_buffer
    func_with_block = alloc_zero_dim_buffer_block
    rt_func = tvm.script.from_source(func.script())
    rt_func_with_block = tvm.script.from_source(func_with_block.script())
    rt_mod = tvm.compile(rt_func, "llvm")
    rt_mod_with_block = tvm.compile(rt_func_with_block, "llvm")
    tvm.ir.assert_structural_equal(
        func.with_attr("global_symbol", "main"), func_with_block.with_attr("global_symbol", "main")
    )
    tvm.ir.assert_structural_equal(
        rt_func.with_attr("global_symbol", "main"),
        rt_func_with_block.with_attr("global_symbol", "main"),
    )
    _check_alloc_zero_dim_buffer(rt_mod)
    _check_alloc_zero_dim_buffer(rt_mod_with_block)


@T.prim_func(s_tir=True)
def ceildiv_test(A: T.Buffer(16, "int32")):
    for i in range(16):
        A[i] = T.ceildiv(A[i], 4)


@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_ceildiv():
    f = tvm.compile(ceildiv_test, "llvm")
    a = tvm.runtime.tensor(np.arange(16).astype("int32"))
    f(a)
    ref = (np.arange(16) + 3) // 4
    tvm.testing.assert_allclose(a.numpy(), ref)


try:

    @T.prim_func(s_tir=True)
    def slice_op_test(
        A: T.Buffer((10,), "float32"), B: T.Buffer((10,), "float32"), C: T.Buffer((10,), "uint32")
    ):
        B[0:5] = A[0:5] + B[0:5]
        B[0:5] = A[0:5] - B[0:5]
        B[0:5] = A[0:5] * B[0:5]
        B[0:5] = A[0:5] / B[0:5]
        C[0:5] = C[0:5] % T.broadcast(T.uint32(5), 5)
        B[0:5] = -B[0:5]
        C[0:5] = C[0:5] >> 4
        C[0:5] = C[0:5] << 4
        C[0:5] = C[0:5] << C[0:5]
        C[0:5] = C[0:5] >> C[0:5]
        T.evaluate(A[0:5] > B[0:5])
        T.evaluate(A[0:5] > 5)
        T.evaluate(A[0:5] >= B[0:5])
        T.evaluate(A[0:5] >= 5)
        T.evaluate(A[0:5] < B[0:5])
        T.evaluate(A[0:5] < 5)
        T.evaluate(A[0:5] <= B[0:5])
        T.evaluate(A[0:5] <= 5)
        T.evaluate(A[0:5] == B[0:5])
        T.evaluate(A[0:5] == 5)
        T.evaluate(A[0:5] != B[0:5])
        T.evaluate(A[0:5] != 5)
        T.evaluate((A[0:5] > 0) and (B[0:5] > 0))
        T.evaluate((A[0:5] > 0) or (B[0:5] > 0))
        T.evaluate((A[0:5] < 0) and (1 > 0))
        T.evaluate((A[0:5] > 0) or (1 > 0))

    @T.prim_func(s_tir=True)
    def slice_op_test_ref(
        A: T.Buffer((10,), "float32"), B: T.Buffer((10,), "float32"), C: T.Buffer((10,), "uint32")
    ):
        B[0:5] = A[0:5] + B[0:5]
        B[0:5] = A[0:5] - B[0:5]
        B[0:5] = A[0:5] * B[0:5]
        B[0:5] = A[0:5] / B[0:5]
        C[0:5] = C[0:5] % T.Broadcast(T.uint32(5), 5)
        B[0:5] = B[0:5] * T.Broadcast(T.float32(-1), 5)
        C[0:5] = T.shift_right(C[0:5], T.Broadcast(T.uint32(4), 5))
        C[0:5] = T.shift_left(C[0:5], T.Broadcast(T.uint32(4), 5))
        C[0:5] = T.shift_left(C[0:5], C[0:5])
        C[0:5] = T.shift_right(C[0:5], C[0:5])
        T.evaluate(A[0:5] > B[0:5])
        T.evaluate(A[0:5] > T.Broadcast(T.float32(5), 5))
        T.evaluate(A[0:5] >= B[0:5])
        T.evaluate(A[0:5] >= T.Broadcast(T.float32(5), 5))
        T.evaluate(A[0:5] < B[0:5])
        T.evaluate(A[0:5] < T.Broadcast(T.float32(5), 5))
        T.evaluate(A[0:5] <= B[0:5])
        T.evaluate(A[0:5] <= T.Broadcast(T.float32(5), 5))
        T.evaluate(A[0:5] == B[0:5])
        T.evaluate(A[0:5] == T.Broadcast(T.float32(5), 5))
        T.evaluate(A[0:5] != B[0:5])
        T.evaluate(A[0:5] != T.Broadcast(T.float32(5), 5))
        T.bitwise_and(A[0:5] > T.Broadcast(T.float32(0), 5), B[0:5] > T.Broadcast(T.float32(0), 5))
        T.bitwise_or(A[0:5] > T.Broadcast(T.float32(0), 5), B[0:5] > T.Broadcast(T.float32(0), 5))
        T.bitwise_and(A[0:5] < T.Broadcast(T.float32(0), 5), T.Broadcast(T.bool(1), 5))
        T.bitwise_or(A[0:5] > T.Broadcast(T.float32(0), 5), T.Broadcast(T.bool(1), 5))
except tvm.error.DiagnosticError:
    slice_op_test = None
    slice_op_test_ref = None


def test_slice_op():
    if slice_op_test is None:
        pytest.skip("slice arithmetic on BufferRegion is not defined")
    tvm.ir.assert_structural_equal(
        slice_op_test.with_attr("global_symbol", "main"),
        slice_op_test_ref.with_attr("global_symbol", "main"),
    )


if __name__ == "__main__":
    test_get_valid_counts_script_func()
    test_alloc_zero_dim_buffer_round_trip()
    test_slice_op()