# 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: E741
import math
import random

import numpy as np
import pytest

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


@pytest.mark.parametrize(
    "dtype, literals",
    [
        ["int8", [-128, 0, 127]],
        ["uint8", [0, 255]],
        ["int32", [-2147483648, 2147483647]],
        ["uint32", [0, 4294967295]],
        ["int64", [-9223372036854775808, 9223372036854775807]],
        ["uint64", [0, 9223372036854775807]],
    ],
)
def test_tir_make_intimm(dtype, literals):
    for l in literals:
        imm = tirx.const(l, dtype)
        assert imm.value == l, imm


@pytest.mark.parametrize(
    "dtype, literals",
    [
        ["int8", [-129, 128]],
        ["uint8", [-1, 256]],
        ["int32", [-2147483650, 2147483648]],
        ["uint32", [-1, 4294967296]],
        ["uint64", [-1, 18446744073709551616]],
    ],
)
def test_tir_invalid_intimm(dtype, literals):
    for l in literals:
        # Out-of-range positive literals raise a builtin ValueError from
        # the IntImm range check; negative-into-unsigned raises an
        # InternalError ("cannot make uint from negative value") which is a
        # RuntimeError subclass. Accept either.
        with pytest.raises((RuntimeError, ValueError)):
            tirx.const(l, dtype)


@pytest.mark.parametrize(
    "dtype, literals",
    [
        [
            "uint64",
            {
                9223372036854775807: 9223372036854775807,
                18446744073709551615: 18446744073709551615,
            },
        ],
    ],
)
def test_tir_large_py_int_literals(dtype, literals):
    """
    For large uint value, use LargeUIntImm intrin,
    """
    for l in literals:
        x = tirx.const(l, dtype)
        if isinstance(x, tirx.IntImm | tirx.FloatImm):
            assert x.value == literals[l]
        else:
            # LargeUIntImm(low32, hi32)
            assert (int(x.args[1]) << 32) + int(x.args[0]) == literals[l]


def test_tir_intimm_overflow():
    assert int(tirx.const(255, "uint8") + tirx.const(1, "uint8")) == 0
    assert int(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32")) == -(2**31)
    assert int(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32")) == 0
    assert int(tirx.const(2**63 - 1, "int64") + tirx.const(1, "int64")) == -(2**63)
    assert int(tirx.const(2**32, "uint64") * tirx.const(2**32, "uint64")) == 0
    # customized int types
    assert int(tirx.const(7, "int4") + tirx.const(1, "int4")) == -8
    assert int(tirx.const(2**39 - 1, "int40") + tirx.const(1, "int40")) == -(2**39)


def compare_float_value(value, expect, msg):
    if math.isfinite(value):
        assert np.abs(value - expect) < 1e-5, f"{value} vs {expect}, {msg}"
    elif math.isnan(value):
        assert math.isnan(expect), f"{value} vs {expect}, {msg}"
    elif math.isinf(value):
        assert math.isinf(expect), f"{value} vs {expect}, {msg}"


@pytest.mark.parametrize(
    "dtype, literals",
    [
        ["float16", [-65504.0, 3.14, 65504.0, np.inf, np.nan]],
        ["bfloat16", [-3.38953139e38, 3.38953139e38, 3.14]],
        ["float32", [np.finfo("float32").min, 3.14, np.finfo("float32").max, np.inf, np.nan]],
        ["float64", [np.finfo("float64").min, 3.14, np.finfo("float64").max, np.inf, np.nan]],
    ],
)
def test_tir_make_floatimm(dtype, literals):
    for l in literals:
        imm = tirx.const(l, dtype)
        compare_float_value(imm.value, l, "imm value should match feed value")


@pytest.mark.parametrize(
    "dtype, literals",
    [
        ["float16", [-65505.0, 65505.0]],
        ["float32", [-3.402e39, 3.402e39]],
    ],
)
def test_tir_invalid_floatimm(dtype, literals):
    """Currently only fp16 and fp32 have range check."""
    for l in literals:
        # FloatImm out-of-range raises a builtin ValueError.
        with pytest.raises(ValueError):
            tirx.const(l, dtype)


@pytest.mark.parametrize("dtype", ["float16", "float32", "float64"])
@pytest.mark.parametrize("literal", [3.14, np.nan, np.inf])
def test_tir_special_floatimms(dtype, literal):
    x = tirx.const(literal, dtype)
    compare_float_value(x.value, literal, "imm value should match feed value")


@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_tir_too_large_literal_f64():
    # Behavior check: if literal f64 value is out of dtype range, the
    # object is still constructed, and eval to infinity.
    @T.prim_func(s_tir=True)
    def imm_overflow_fp64() -> T.float64:
        T.evaluate(T.ret(T.float64(1.7976e309), dtype="float64"))

    f = tvm.compile(imm_overflow_fp64, target="llvm")
    assert math.isinf(f())


@pytest.mark.parametrize(
    "literal, expect_dtype",
    [
        (256, "int32"),
        (2147483647, "int32"),
        (-2147483648, "int32"),
        (2147483648, "int64"),
        (-2147483649, "int64"),
        (3.14159, "float32"),
        (np.finfo("float32").min, "float32"),
        (np.finfo("float32").max, "float32"),
        (-3.402e39, "float64"),
        (3.402e39, "float64"),
    ],
)
def test_tir_const_auto_dtype(literal, expect_dtype):
    x = tirx.const(literal, dtype=None)
    assert x.ty.dtype == expect_dtype
    assert x.value == literal


def check_tir_const_fold(
    dtype, foldf, calcf, x_range=None, y_range=None, expect=None, skip_overflow=False
):
    """Helper to check constant folding behavior

    Parameters
    ----------
    dtype: str
        Datatype of constants

    foldf: (x, y) -> z
        Folding function to call

    calcf: (x, y) -> z
        Compiled calculation function to call

    x_range: Union[int, float, tuple]
        Single value or value range [min, max]

    y_range: Union[int, float, tuple]
        Single value or value range [min, max]

    expect: Union[int, float]
        Expected calculation result

    skip_overflow: bool
        Skip assertion if the overflow happens
    """
    seed = random.randint(0, 2147483648)
    np.random.seed(seed)
    ninfo = np.finfo(dtype) if dtype.startswith("float") else np.iinfo(dtype)

    if x_range is None:
        x_range = (ninfo.min, ninfo.max)
    if isinstance(x_range, int | float):
        x = x_range
    elif dtype.startswith("int") or dtype.startswith("uint"):
        x = np.random.randint(x_range[0], x_range[1] + 1, dtype=dtype)
    else:
        x = np.random.uniform(x_range[0], x_range[1])

    if y_range is None:
        y_range = (ninfo.min, ninfo.max)
    if isinstance(y_range, int | float):
        y = y_range
    elif dtype.startswith("int") or dtype.startswith("uint"):
        y = np.random.randint(y_range[0], y_range[1] + 1, dtype=dtype)
    else:
        y = np.random.uniform(y_range[0], y_range[1])

    if skip_overflow:
        py_res = foldf(x, y)
        if isinstance(py_res, tirx.IntImm | tirx.FloatImm):
            py_res = py_res.value
        if not (ninfo.min <= py_res <= ninfo.max):
            # If the result overflow, certain arithmetics is non-defined
            # thus we intentionally do not make the test failed.
            return

    fold_res = foldf(tirx.const(x, dtype), tirx.const(y, dtype))
    calc_res = calcf(x, y)

    flaky_msg = (
        f"{dtype} ({x}, {y}, {expect}) const folding check failed.\n"
        + "This test is intentionally non-deterministic, "
        + f"if it fails please report it in GitHub issue together with this seed {seed}\n"
    )
    if dtype.startswith("float"):
        compare_float_value(calc_res, fold_res.value, flaky_msg)
        if expect:
            compare_float_value(expect, calc_res, flaky_msg)
    else:
        assert calc_res == fold_res.value, flaky_msg
        if expect:
            assert expect == calc_res, flaky_msg


@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_tir_floatimm_const_fold():
    """Behavior check: folding fp32 match platform f32 arithmetic"""

    @T.prim_func(s_tir=True)
    def float_imm_multiply(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
        z[()] = x * y

    @T.prim_func(s_tir=True)
    def float_imm_add(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
        z[()] = x + y

    @T.prim_func(s_tir=True)
    def float_imm_sub(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
        z[()] = x - y

    @T.prim_func(s_tir=True)
    def float_imm_div(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
        z[()] = x / y

    def __wrap_build(f):
        lib = tvm.compile(f, target="llvm")
        z = tvm.runtime.tensor(np.zeros([]).astype("float32"))

        def _func(x, y):
            lib(x, y, z)
            return z.numpy()

        return _func

    fmul = __wrap_build(float_imm_multiply)
    fadd = __wrap_build(float_imm_add)
    fsub = __wrap_build(float_imm_sub)
    fdiv = __wrap_build(float_imm_div)

    # overflow
    check_tir_const_fold("float32", lambda x, y: x * y, fmul, 3.0e30, 3.0e30, np.inf)
    check_tir_const_fold("float32", lambda x, y: x * y, fmul, 3.0e30, -3.0e30, -np.inf)
    check_tir_const_fold("float32", lambda x, y: x / y, fdiv, 3.0e30, 3.0e-30, np.inf)

    # divide by zero
    with pytest.raises(RuntimeError):
        check_tir_const_fold("float32", lambda x, y: x / y, fdiv, 1.0, 0.0)

    # nan and inf
    check_tir_const_fold("float32", lambda x, y: x + y, fadd, 1.0, np.nan, np.nan)
    check_tir_const_fold("float32", lambda x, y: x + y, fadd, 1.0, np.inf, np.inf)
    check_tir_const_fold("float32", lambda x, y: x + y, fadd, 1.0, -np.inf, -np.inf)

    # randomized check
    check_tir_const_fold("float32", lambda x, y: x * y, fmul)
    check_tir_const_fold("float32", lambda x, y: x + y, fadd)
    check_tir_const_fold("float32", lambda x, y: x - y, fsub)
    check_tir_const_fold(
        "float32", lambda x, y: x / y, fdiv, y_range=(0.01, np.finfo("float32").max)
    )


@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_tir_int8_const_fold():
    """Behavior check: folding i8 operation match platform i8 arithmetic"""

    @T.prim_func(s_tir=True)
    def imm_multiply(x: T.int8, y: T.int8) -> T.int8:
        T.evaluate(T.ret(x * y, dtype="int8"))

    @T.prim_func(s_tir=True)
    def imm_add(x: T.int8, y: T.int8) -> T.int8:
        T.evaluate(T.ret(x + y, dtype="int8"))

    @T.prim_func(s_tir=True)
    def imm_sub(x: T.int8, y: T.int8) -> T.int8:
        T.evaluate(T.ret(x - y, dtype="int8"))

    @T.prim_func(s_tir=True)
    def imm_truncdiv(x: T.int8, y: T.int8) -> T.int8:
        T.evaluate(T.ret(T.truncdiv(x, y), dtype="int8"))

    @T.prim_func(s_tir=True)
    def imm_floordiv(x: T.int8, y: T.int8) -> T.int8:
        T.evaluate(T.ret(T.floordiv(x, y), dtype="int8"))

    fmul = tvm.compile(imm_multiply, target="llvm")
    fadd = tvm.compile(imm_add, target="llvm")
    fsub = tvm.compile(imm_sub, target="llvm")
    ffloordiv = tvm.compile(imm_floordiv, target="llvm")
    ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")

    # overflow
    check_tir_const_fold("int8", lambda x, y: x + y, fadd, 127, 1, -128)
    check_tir_const_fold("int8", lambda x, y: x * y, fmul, 127, 127, 1)

    # divide by zero
    with pytest.raises(RuntimeError):
        check_tir_const_fold("int8", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
    with pytest.raises(RuntimeError):
        check_tir_const_fold("int8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)

    # i8 mod folding is not implemented
    assert not isinstance(tirx.floormod(tirx.const(7, "int8"), tirx.const(3, "int8")), tirx.IntImm)
    assert not isinstance(tirx.truncmod(tirx.const(7, "int8"), tirx.const(3, "int8")), tirx.IntImm)

    # randomized check
    check_tir_const_fold("int8", lambda x, y: x * y, fmul)
    check_tir_const_fold("int8", lambda x, y: x + y, fadd)
    check_tir_const_fold("int8", lambda x, y: x - y, fsub)
    check_tir_const_fold(
        "int8", lambda x, y: tirx.floordiv(x, y), ffloordiv, y_range=(1, np.iinfo("int8").max)
    )
    check_tir_const_fold(
        "int8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, y_range=(1, np.iinfo("int8").max)
    )


@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_tir_uint8_const_fold():
    """Behavior check: folding u8 operation match platform u8 arithmetic"""

    @T.prim_func(s_tir=True)
    def imm_multiply(x: T.uint8, y: T.uint8) -> T.uint8:
        T.evaluate(T.ret(x * y, dtype="uint8"))

    @T.prim_func(s_tir=True)
    def imm_add(x: T.uint8, y: T.uint8) -> T.uint8:
        T.evaluate(T.ret(x + y, dtype="uint8"))

    @T.prim_func(s_tir=True)
    def imm_sub(x: T.uint8, y: T.uint8) -> T.uint8:
        T.evaluate(T.ret(x - y, dtype="uint8"))

    @T.prim_func(s_tir=True)
    def imm_truncdiv(x: T.uint8, y: T.uint8) -> T.uint8:
        T.evaluate(T.ret(T.truncdiv(x, y), dtype="uint8"))

    @T.prim_func(s_tir=True)
    def imm_floordiv(x: T.uint8, y: T.uint8) -> T.uint8:
        T.evaluate(T.ret(T.floordiv(x, y), dtype="uint8"))

    fmul = tvm.compile(imm_multiply, target="llvm")
    fadd = tvm.compile(imm_add, target="llvm")
    fsub = tvm.compile(imm_sub, target="llvm")
    ffloordiv = tvm.compile(imm_floordiv, target="llvm")
    ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")

    # overflow
    check_tir_const_fold("uint8", lambda x, y: x + y, fadd, 255, 1, 0)

    # zero sub
    with pytest.raises(RuntimeError):
        check_tir_const_fold("uint8", lambda x, y: x - y, fsub, 0, 10)

    # divide by zero
    with pytest.raises(RuntimeError):
        check_tir_const_fold("uint8", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
    with pytest.raises(RuntimeError):
        check_tir_const_fold("uint8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)

    # u8 floormod folding is overflow-free and implemented
    folded_floormod = tirx.floormod(tirx.const(7, "uint8"), tirx.const(3, "uint8"))
    assert isinstance(folded_floormod, tirx.IntImm)
    assert int(folded_floormod) == 1

    # u8 truncmod folding is not implemented
    assert not isinstance(
        tirx.truncmod(tirx.const(7, "uint8"), tirx.const(3, "uint8")), tirx.IntImm
    )

    # randomized check
    check_tir_const_fold("uint8", lambda x, y: x * y, fmul)
    check_tir_const_fold("uint8", lambda x, y: x + y, fadd)
    check_tir_const_fold("uint8", lambda x, y: x - y, fsub)
    check_tir_const_fold(
        "uint8", lambda x, y: tirx.floordiv(x, y), ffloordiv, y_range=(1, np.iinfo("uint8").max)
    )
    check_tir_const_fold(
        "uint8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, y_range=(1, np.iinfo("uint8").max)
    )


@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_tir_int32_const_fold():
    """Behavior check: folding i32 operation match platform i32 arithmetic"""

    @T.prim_func(s_tir=True)
    def imm_multiply(x: T.int32, y: T.int32) -> T.int32:
        T.evaluate(T.ret(x * y, dtype="int32"))

    @T.prim_func(s_tir=True)
    def imm_add(x: T.int32, y: T.int32) -> T.int32:
        T.evaluate(T.ret(x + y, dtype="int32"))

    @T.prim_func(s_tir=True)
    def imm_sub(x: T.int32, y: T.int32) -> T.int32:
        T.evaluate(T.ret(x - y, dtype="int32"))

    @T.prim_func(s_tir=True)
    def imm_truncdiv(x: T.int32, y: T.int32) -> T.int32:
        T.evaluate(T.ret(T.truncdiv(x, y), dtype="int32"))

    @T.prim_func(s_tir=True)
    def imm_truncmod(x: T.int32, y: T.int32) -> T.int32:
        T.evaluate(T.ret(T.truncmod(x, y), dtype="int32"))

    @T.prim_func(s_tir=True)
    def imm_floordiv(x: T.int32, y: T.int32) -> T.int32:
        T.evaluate(T.ret(T.floordiv(x, y), dtype="int32"))

    @T.prim_func(s_tir=True)
    def imm_floormod(x: T.int32, y: T.int32) -> T.int32:
        T.evaluate(T.ret(T.floormod(x, y), dtype="int32"))

    fmul = tvm.compile(imm_multiply, target="llvm")
    fadd = tvm.compile(imm_add, target="llvm")
    fsub = tvm.compile(imm_sub, target="llvm")
    ffloordiv = tvm.compile(imm_floordiv, target="llvm")
    ffloormod = tvm.compile(imm_floormod, target="llvm")
    ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
    ftruncmod = tvm.compile(imm_truncmod, target="llvm")

    # i32 overflow is not specified, only check for range
    assert -(2**31) <= int(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32")) < 2**31
    assert -(2**31) <= int(tirx.const(-(2**31), "int32") - tirx.const(1, "int32")) < 2**31

    # divide by zero
    with pytest.raises(RuntimeError):
        check_tir_const_fold("int32", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
    with pytest.raises(RuntimeError):
        check_tir_const_fold("int32", lambda x, y: tirx.floormod(x, y), ffloormod, 1, 0)
    with pytest.raises(RuntimeError):
        check_tir_const_fold("int32", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)
    with pytest.raises(RuntimeError):
        check_tir_const_fold("int32", lambda x, y: tirx.truncmod(x, y), ftruncmod, 1, 0)

    # randomized check
    check_tir_const_fold("int32", lambda x, y: x * y, fmul, skip_overflow=True)
    check_tir_const_fold("int32", lambda x, y: x + y, fadd, skip_overflow=True)
    check_tir_const_fold("int32", lambda x, y: x - y, fsub, skip_overflow=True)
    check_tir_const_fold(
        "int32",
        lambda x, y: tirx.floordiv(x, y),
        ffloordiv,
        y_range=(1, np.iinfo("int32").max),
        skip_overflow=True,
    )
    check_tir_const_fold(
        "int32",
        lambda x, y: tirx.truncdiv(x, y),
        ftruncdiv,
        y_range=(1, np.iinfo("int32").max),
        skip_overflow=True,
    )
    check_tir_const_fold(
        "int32",
        lambda x, y: tirx.floormod(x, y),
        ffloormod,
        y_range=(1, np.iinfo("int32").max),
        skip_overflow=False,
    )
    check_tir_const_fold(
        "int32",
        lambda x, y: tirx.truncmod(x, y),
        ftruncmod,
        y_range=(1, np.iinfo("int32").max),
        skip_overflow=False,
    )


@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_tir_uint32_const_fold():
    """Behavior check: folding u32 operation match platform u32 arithmetic"""

    @T.prim_func(s_tir=True)
    def imm_multiply(x: T.uint32, y: T.uint32) -> T.uint32:
        T.evaluate(T.ret(x * y, dtype="uint32"))

    @T.prim_func(s_tir=True)
    def imm_add(x: T.uint32, y: T.uint32) -> T.uint32:
        T.evaluate(T.ret(x + y, dtype="uint32"))

    @T.prim_func(s_tir=True)
    def imm_sub(x: T.uint32, y: T.uint32) -> T.uint32:
        T.evaluate(T.ret(x - y, dtype="uint32"))

    @T.prim_func(s_tir=True)
    def imm_truncdiv(x: T.uint32, y: T.uint32) -> T.uint32:
        T.evaluate(T.ret(T.truncdiv(x, y), dtype="uint32"))

    @T.prim_func(s_tir=True)
    def imm_floordiv(x: T.uint32, y: T.uint32) -> T.uint32:
        T.evaluate(T.ret(T.floordiv(x, y), dtype="uint32"))

    fmul = tvm.compile(imm_multiply, target="llvm")
    fadd = tvm.compile(imm_add, target="llvm")
    fsub = tvm.compile(imm_sub, target="llvm")
    ffloordiv = tvm.compile(imm_floordiv, target="llvm")
    ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")

    # u32 overflow is not specified, only check for range
    assert 0 <= int(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32")) < 2**32

    # divide by zero
    with pytest.raises(RuntimeError):
        check_tir_const_fold("uint32", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
    with pytest.raises(RuntimeError):
        check_tir_const_fold("uint32", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)

    # u32 floormod folding is overflow-free and implemented
    folded_floormod = tirx.floormod(tirx.const(7, "uint32"), tirx.const(3, "uint32"))
    assert isinstance(folded_floormod, tirx.IntImm)
    assert int(folded_floormod) == 1

    # u32 truncmod folding is not implemented
    assert not isinstance(
        tirx.truncmod(tirx.const(7, "uint32"), tirx.const(3, "uint32")), tirx.IntImm
    )

    # randomized check
    check_tir_const_fold("uint32", lambda x, y: x * y, fmul, skip_overflow=True)
    check_tir_const_fold("uint32", lambda x, y: x + y, fadd, skip_overflow=True)
    check_tir_const_fold("uint32", lambda x, y: x - y, fsub, skip_overflow=True)
    check_tir_const_fold(
        "uint32",
        lambda x, y: tirx.floordiv(x, y),
        ffloordiv,
        y_range=(1, np.iinfo("uint32").max),
        skip_overflow=False,
    )
    check_tir_const_fold(
        "uint32",
        lambda x, y: tirx.truncdiv(x, y),
        ftruncdiv,
        y_range=(1, np.iinfo("uint32").max),
        skip_overflow=False,
    )


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