# 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.
# ruff: noqa: E712, F401
import ctypes
import math

import numpy as np
import pytest

pytest.importorskip("scipy")

import scipy

import tvm
import tvm.testing
from tvm import te, tirx, topi
from tvm.script import tirx as T
from tvm.support import clang, utils


def test_nearbyint():
    m = te.var(
        "m",
    )
    A = te.placeholder((m,), name="A")
    A_rounded = te.compute((m,), lambda *i: tvm.tirx.nearbyint(A(*i)), name="A")

    # Convert to TIR and create schedule
    mod = te.create_prim_func([A, A_rounded])
    sch = tvm.s_tir.Schedule(mod)

    # Build from scheduled TIR
    func = tvm.compile(sch.mod, target="llvm")

    dev = tvm.cpu(0)
    n = 10
    a = tvm.runtime.tensor(np.random.uniform(high=100, size=n).astype(A.dtype.dtype), dev)
    a_rounded = tvm.runtime.tensor(np.random.uniform(size=n).astype(A_rounded.dtype.dtype), dev)
    func(a, a_rounded)
    # Note that numpys rint rounds to nearest integer with
    # ties to halfway is broken by rounding to even.
    # So that 1.5 and 2.5 will round 2.
    # This is the default rounding mode with libc as well.
    # However one can set a different rounding mode and in that
    # case numpy result might differ.
    tvm.testing.assert_allclose(a_rounded.numpy(), np.rint(a.numpy()))


def test_round_ties_to_even():
    """Test that tir.round uses ties-to-even (banker's rounding) semantics."""
    m = te.var("m")
    A = te.placeholder((m,), name="A")
    A_rounded = te.compute((m,), lambda *i: tvm.tirx.round(A(*i)), name="A")

    mod = te.create_prim_func([A, A_rounded])
    sch = tvm.s_tir.Schedule(mod)
    func = tvm.compile(sch.mod, target="llvm")

    dev = tvm.cpu(0)
    # Midpoint values where ties-to-even and ties-away differ
    test_values = np.array([0.5, 1.5, 2.5, 3.5, -0.5, -1.5, -2.5, -3.5], dtype="float32")
    expected = np.array([0.0, 2.0, 2.0, 4.0, 0.0, -2.0, -2.0, -4.0], dtype="float32")

    a = tvm.runtime.tensor(test_values, dev)
    a_rounded = tvm.runtime.tensor(np.zeros(len(test_values), dtype="float32"), dev)
    func(a, a_rounded)
    tvm.testing.assert_allclose(a_rounded.numpy(), expected)


def test_round_intrinsics_on_int():
    i = tvm.tirx.Var("i", "int32")
    for op in [tvm.tirx.round, tvm.tirx.trunc, tvm.tirx.ceil, tvm.tirx.floor, tvm.tirx.nearbyint]:
        assert op(tvm.tirx.const(10, "int32")).value == 10
        assert op(tvm.tirx.const(True, "bool")).value == True
        assert op(i).same_as(i)

    assert tvm.tirx.isnan(tvm.tirx.const(10, "int32")).value == False


def test_unary_intrin():
    test_funcs = [
        (tvm.tirx.exp, lambda x: np.exp(x)),
        (tvm.tirx.exp10, lambda x: np.power(10, x)),
        (tvm.tirx.log2, lambda x: np.log2(x)),
        (tvm.tirx.log10, lambda x: np.log10(x)),
        (tvm.tirx.sinh, lambda x: np.sinh(x)),
        (tvm.tirx.cosh, lambda x: np.cosh(x)),
        (tvm.tirx.log1p, lambda x: np.log1p(x)),
        (tvm.tirx.asin, lambda x: np.arcsin(x)),
        (tvm.tirx.acos, lambda x: np.arccos(x)),
        (tvm.tirx.atan, lambda x: np.arctan(x)),
        (tvm.tirx.asinh, lambda x: np.arcsinh(x)),
        (tvm.tirx.acosh, lambda x: np.arccosh(x)),
        (tvm.tirx.atanh, lambda x: np.arctanh(x)),
        (tvm.tirx.erf, lambda x: scipy.special.erf(x)),
    ]

    def run_test(tvm_intrin, np_func, atol=1e-5, rtol=1e-5):
        m = te.var(
            "m",
        )
        A = te.placeholder((m,), name="A")
        B = te.compute((m,), lambda *i: tvm_intrin(A(*i)), name="B")

        # Convert to TIR and create schedule
        mod = te.create_prim_func([A, B])
        sch = tvm.s_tir.Schedule(mod)

        # Build from scheduled TIR
        func = tvm.compile(sch.mod, target="llvm")

        dev = tvm.cpu(0)
        n = 10
        a = tvm.runtime.tensor(np.random.uniform(0.1, 0.5, size=n).astype(A.dtype.dtype), dev)
        b = tvm.runtime.tensor(np.random.uniform(size=n).astype(A.dtype.dtype), dev)
        func(a, b)
        tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=atol, rtol=rtol)

        # Out-of-bounds test for asin/acos
        name = tvm_intrin.__name__
        if name in ("asin", "acos"):
            # generate some values outside [-1, 1]
            n = 8
            out_np = np.concatenate(
                [
                    np.random.uniform(1.1, 2.0, size=n // 2),
                    np.random.uniform(-2.0, -1.1, size=n // 2),
                ]
            ).astype(A.dtype.dtype)
            a2 = tvm.runtime.tensor(out_np, dev)
            b2 = tvm.runtime.tensor(np.empty_like(out_np), dev)
            func(a2, b2)
            # all outputs should be NaN
            assert np.all(np.isnan(b2.numpy()))
        if name == "exp":
            n = 8
            out_np = np.random.randint(-20, 20, size=n).astype(A.dtype.dtype)
            a2 = tvm.runtime.tensor(out_np, dev)
            b2 = tvm.runtime.tensor(np.empty_like(out_np), dev)
            func(a2, b2)
            assert b2.numpy().dtype == np.float32
            # Verify correctness against NumPy exp
            expected = np.exp(out_np.astype(np.float32))
            tvm.testing.assert_allclose(b2.numpy(), expected, rtol=1e-5, atol=1e-5)

    for func in test_funcs:
        atol = rtol = 1e-3 if func[0].__name__ in ["asin", "acos", "atan"] else 1e-5
        run_test(*func, atol, rtol)


def test_asin_acos_boundary_values():
    """Test asin and acos with boundary values and threshold switching."""
    test_funcs = [
        (tvm.tirx.asin, lambda x: np.arcsin(x)),
        (tvm.tirx.acos, lambda x: np.arccos(x)),
    ]

    def run_test(tvm_intrin, np_func):
        m = te.var("m")
        A = te.placeholder((m,), name="A")
        B = te.compute((m,), lambda *i: tvm_intrin(A(*i)), name="B")

        mod = te.create_prim_func([A, B])
        sch = tvm.s_tir.Schedule(mod)
        func = tvm.compile(sch.mod, target="llvm")

        dev = tvm.cpu(0)

        # Test boundary values: ±1.0 (should use system library)
        boundary_values = np.array([1.0, -1.0], dtype=np.float32)
        a1 = tvm.runtime.tensor(boundary_values, dev)
        b1 = tvm.runtime.tensor(np.empty_like(boundary_values), dev)
        func(a1, b1)
        tvm.testing.assert_allclose(b1.numpy(), np_func(boundary_values), atol=1e-5, rtol=1e-5)

        # Test values at threshold: ±0.5 (should use system library)
        threshold_values = np.array([0.5, -0.5], dtype=np.float32)
        a2 = tvm.runtime.tensor(threshold_values, dev)
        b2 = tvm.runtime.tensor(np.empty_like(threshold_values), dev)
        func(a2, b2)
        tvm.testing.assert_allclose(b2.numpy(), np_func(threshold_values), atol=1e-4, rtol=1e-4)

        # Test values just below threshold: ±0.49 (should use Taylor series)
        below_threshold_values = np.array([0.49, -0.49, 0.3, -0.3, 0.0], dtype=np.float32)
        a3 = tvm.runtime.tensor(below_threshold_values, dev)
        b3 = tvm.runtime.tensor(np.empty_like(below_threshold_values), dev)
        func(a3, b3)
        tvm.testing.assert_allclose(
            b3.numpy(), np_func(below_threshold_values), atol=1e-3, rtol=1e-3
        )

        # Test out-of-domain values: should return NaN
        out_of_domain = np.array([1.1, -1.1, 2.0, -2.0], dtype=np.float32)
        a4 = tvm.runtime.tensor(out_of_domain, dev)
        b4 = tvm.runtime.tensor(np.empty_like(out_of_domain), dev)
        func(a4, b4)
        assert np.all(np.isnan(b4.numpy())), "Out-of-domain inputs should return NaN"

    for func in test_funcs:
        run_test(*func)


def test_binary_intrin():
    test_funcs = [
        (tvm.tirx.atan2, lambda x1, x2: np.arctan2(x1, x2)),
        (tvm.tirx.nextafter, lambda x1, x2: np.nextafter(x1, x2)),
        (tvm.tirx.copysign, lambda x1, x2: np.copysign(x1, x2)),
        (tvm.tirx.hypot, lambda x1, x2: np.hypot(x1, x2)),
    ]

    def run_test(tvm_intrin, np_func):
        m = te.var(
            "m",
        )
        A = te.placeholder((m,), name="A")
        B = te.placeholder((m,), name="B")
        C = te.compute((m,), lambda *i: tvm_intrin(A(*i), B(*i)), name="C")

        # Convert to TIR and create schedule
        mod = te.create_prim_func([A, B, C])
        sch = tvm.s_tir.Schedule(mod)

        # Build from scheduled TIR
        func = tvm.compile(sch.mod, target="llvm")

        dev = tvm.cpu(0)
        n = 10
        a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(A.dtype.dtype), dev)
        b = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(B.dtype.dtype), dev)
        c = tvm.runtime.tensor(np.random.uniform(size=n).astype(A.dtype.dtype), dev)
        func(a, b, c)
        tvm.testing.assert_allclose(c.numpy(), np_func(a.numpy(), b.numpy()), atol=1e-5, rtol=1e-5)

    for func in test_funcs:
        run_test(*func)


def test_ldexp():
    m = te.var(
        "m",
    )
    A = te.placeholder((m,), name="A")
    B = te.placeholder((m,), name="B", dtype="int32")
    C = te.compute((m,), lambda *i: tvm.tirx.ldexp(A(*i), B(*i)), name="C")

    # Convert to TIR and create schedule
    mod = te.create_prim_func([A, B, C])
    sch = tvm.s_tir.Schedule(mod)

    # Build from scheduled TIR
    func = tvm.compile(sch.mod, target="llvm")

    dev = tvm.cpu(0)
    n = 10
    a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(A.dtype.dtype), dev)
    b = tvm.runtime.tensor(np.random.randint(0, 5, size=n).astype(B.dtype.dtype), dev)
    c = tvm.runtime.tensor(np.random.uniform(size=n).astype(A.dtype.dtype), dev)
    func(a, b, c)
    tvm.testing.assert_allclose(c.numpy(), np.ldexp(a.numpy(), b.numpy()), atol=1e-5, rtol=1e-5)


dtype = tvm.testing.parameter("int32", "int64")


@pytest.mark.parametrize(
    "target",
    ["llvm", pytest.param({"kind": "vulkan", "from_device": 0}, marks=pytest.mark.gpu)],
)
def test_clz(target, dtype):
    if not tvm.testing.device_enabled(target):
        pytest.skip(f"{target} not enabled")
    target = tvm.target.Target(target)
    if (
        target.kind.name == "vulkan"
        and dtype == "int64"
        and not target.attrs.get("supports_int64", False)
    ):
        pytest.xfail("Vulkan target does not support Int64 types")

    def clz_np(x, dtype):
        ceil_log2 = np.ceil(np.log2(x)).astype(dtype)
        bits = int(dtype[-2:])
        clz = bits - ceil_log2
        clz[np.bitwise_and(x, x - 1) == 0] -= 1
        return clz

    m = te.var("m")
    A = te.placeholder((m,), name="A", dtype=dtype)
    B = te.compute((m,), lambda *i: tvm.tirx.clz(A(*i)), name="B")

    # Convert to TIR and create schedule
    mod = te.create_prim_func([A, B])
    sch = tvm.s_tir.Schedule(mod)

    # Apply scheduling primitives if target is Vulkan
    if target.kind.name == "vulkan":
        block = sch.get_sblock("B")
        loop = sch.get_loops(block)[0]
        bx, tx = sch.split(loop, factors=[None, 64])
        sch.bind(bx, "blockIdx.x")
        sch.bind(tx, "threadIdx.x")

    # Build from scheduled TIR
    func = tvm.compile(sch.mod, target=target)

    def run_and_check():
        dev = tvm.device(target.kind.name)
        n = 10
        highs = [10, 100, 1000, 10000, 100000, 1000000]

        if dtype == "int64":
            highs.append((1 << 63) - 1)

        for high in highs:
            a_np = np.random.randint(1, high=high, size=(n,), dtype=dtype)
            a = tvm.runtime.tensor(a_np, dev)
            b = tvm.runtime.tensor(np.zeros((n,)).astype("int32"), dev)
            func(a, b)
            ref = clz_np(a_np, dtype)
            np.testing.assert_equal(b.numpy(), ref)

    if target.kind.name == "llvm":
        run_and_check()
    else:
        tvm.testing.run_with_gpu_lock(run_and_check)


@tvm.script.ir_module
class Module:
    @T.prim_func(s_tir=True)
    def test_tir_fma(A: T.handle, B: T.handle, C: T.handle, d: T.handle) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "test_fma", "tirx.noalias": True})
        n = T.int32()
        stride = T.int32()
        stride_1 = T.int32()
        stride_2 = T.int32()
        stride_3 = T.int32()
        A_1 = T.match_buffer(
            A,
            [n],
            strides=[stride],
            elem_offset=0,
            align=64,
            offset_factor=1,
            buffer_type="auto",
        )
        B_1 = T.match_buffer(
            B,
            [n],
            strides=[stride_1],
            elem_offset=0,
            align=64,
            offset_factor=1,
            buffer_type="auto",
        )
        C_1 = T.match_buffer(
            C,
            [n],
            strides=[stride_2],
            elem_offset=0,
            align=64,
            offset_factor=1,
            buffer_type="auto",
        )
        d_1 = T.match_buffer(
            d,
            [n],
            strides=[stride_3],
            elem_offset=0,
            align=64,
            offset_factor=1,
            buffer_type="auto",
        )
        # body
        for i in T.serial(0, n):
            d_1[(i * stride_3)] = (A_1[(i * stride)] * B_1[(i * stride_1)]) + C_1[(i * stride_2)]


def test_fma():
    opt = tvm.transform.Sequential(
        [
            tvm.tirx.transform.Apply(lambda f: f.with_attr("target", tvm.target.Target("llvm"))),
            tvm.tirx.transform.LowerIntrin(),
        ]
    )
    mod = opt(Module)
    assert mod["test_tir_fma"].body.body.value.op.name == "tirx.call_llvm_pure_intrin"


if __name__ == "__main__":
    test_nearbyint()
    test_unary_intrin()
    test_round_intrinsics_on_int()
    test_asin_acos_boundary_values()
    test_binary_intrin()
    test_ldexp()
    test_clz()
    test_fma()