# 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 pytest

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


def test_simplify_reshape_flattened_index():
    ana = tvm.arith.Analyzer()

    i0 = tirx.Var("i0", "int64")
    i1 = tirx.Var("i1", "int64")
    ana.bind(i0, tvm.ir.Range(0, 8))
    ana.bind(i1, tvm.ir.Range(0, 3))

    i_flattened = i0 * 3 + i1
    tvm.ir.assert_structural_equal(
        ana.simplify((i_flattened) // 12 * 12 + (i_flattened) % 12 // 4 * 4 + (i_flattened) % 4),
        i_flattened,
    )


dtype = tvm.testing.parameter(
    "uint8",
    "uint16",
    "uint32",
    "uint64",
    "int8",
    "int16",
    "int32",
    "int64",
    "float16",
    "float32",
    "float64",
)


def test_can_prove_self_identity(dtype):
    ana = tvm.arith.Analyzer()

    n = tirx.Var("n", dtype)
    assert ana.can_prove(n == n)


def test_can_prove_self_equal_to_self(dtype):
    ana = tvm.arith.Analyzer()

    n = tirx.Var("n", dtype)
    assert ana.can_prove_equal(n, n)


def test_simplify_symbolic_comparison():
    ana = tvm.arith.Analyzer()

    i0 = tirx.Var("i0", "int64")
    i1 = tirx.Var("i1", "int64")
    n, m = tvm.tirx.Var("n", "int64"), tvm.tirx.Var("m", "int64")
    outer = (n + 31) // 32
    PS = tvm.arith.ProofStrength

    non_negative = tvm.arith.ConstIntBound(0, tvm.arith.ConstIntBound.POS_INF)
    ana.update(n, non_negative)
    ana.update(m, non_negative)
    ana.bind(i0, tvm.ir.Range(0, outer))
    ana.bind(i1, tvm.ir.Range(0, 32))
    assert not ana.can_prove(i0 * 32 + i1 < (n + 31) // 32 * 32, PS.DEFAULT)
    assert ana.can_prove(i0 * 32 + i1 < (n + 31) // 32 * 32, PS.SYMBOLIC_BOUND)
    assert ana.can_prove(i0 * 32 + i1 < (n + 31) // 32 * 32 + m, PS.SYMBOLIC_BOUND)
    assert ana.can_prove(i0 * 32 + i1 + 1 <= (n + 31) // 32 * 32, PS.SYMBOLIC_BOUND)
    assert ana.can_prove((n + 31) // 32 * 32 >= i0 * 32 + i1 + 1, PS.SYMBOLIC_BOUND)
    assert ana.can_prove((n + 31) // 32 * 32 >= i0 * 32 + i1, PS.SYMBOLIC_BOUND)


# These tests exercised arith::CanProve's substitution-based proof loop for
# vscale-bearing expressions (iterating over known vscale values for a VLA target).
# That loop has been removed -- arith no longer attempts target-dependent proofs
# about scalable-vector lengths. The LOG(WARNING) for non-VLA targets is also gone.
@pytest.mark.xfail(reason="arith no longer proves vscale-bearing inequalities via substitution")
@pytest.mark.parametrize(
    "expression",
    [
        T.vscale() * 32 < T.vscale() * 64,
        T.vscale() * 2 * (T.vscale() * 2) >= T.vscale() * 4,
        (T.vscale() * 4 + 114) // (T.vscale() * 4) * (T.vscale() * 4) >= 115,
        64 % T.vscale() <= T.vscale(),
    ],
)
def test_simplify_vscale_comparison_with_sve_target(expression):
    ana = tvm.arith.Analyzer()

    with tvm.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}):
        assert ana.can_prove(expression)


@pytest.mark.xfail(
    reason="arith no longer emits a LOG(WARNING) for vscale proofs on non-VLA targets"
)
def test_simplify_vscale_comparison_without_sve_target(capfd):
    ana = tvm.arith.Analyzer()
    vs = tvm.tirx.vscale()

    with pytest.raises(AssertionError):
        with tvm.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu"}):
            assert ana.can_prove(vs * 32 < vs * 64)

    warning_prefix = (
        "Warning: The expression contains scalable values. An attempt to prove by substituting "
        "with known values of vscale was not performed. This proof currently only supports "
        "VLA targets, but the target was "
    )
    capture = capfd.readouterr().err
    assert warning_prefix in capture
    assert '"kind":"llvm"' in capture
    assert '"mtriple":"aarch64-linux-gnu"' in capture


def test_regression_simplify_inf_recursion():
    ana = tvm.arith.Analyzer()
    cond = tirx.Var("cond", "int32")

    res = (tvm.tirx.NE(cond, 0).astype("int8") - tvm.tirx.NE(cond, 0).astype("int8")).astype(
        "int32"
    ) == 0
    # regression in a previous case
    # try compare and int set recursive call can cause infinite loop
    ana.rewrite_simplify(res)


def test_bind_allow_override():
    ana = tvm.arith.Analyzer()
    x = tirx.Var("x", "int64")

    ana.bind(x, tvm.ir.Range(0, 10))
    ana.bind(x, tvm.ir.Range(0, 5), allow_override=True)
    assert ana.can_prove(x < 5)

    with pytest.raises(RuntimeError, match="Trying to update var 'x' with a different const bound"):
        ana.bind(x, tvm.ir.Range(0, 3))


def test_simplify_floor_mod_with_linear_offset():
    """
    Test that the floor_mod is simplified correctly when the offset is linear.
    """
    ana = tvm.arith.Analyzer()
    past_decoder_sequence_length = tirx.Var("past_decoder_sequence_length", "int64")
    expr1 = (past_decoder_sequence_length + 1) * 64
    divisor1 = (past_decoder_sequence_length + 1) * 32
    assert ana.can_prove_equal(tvm.tirx.floormod(expr1, divisor1), 0)
    divisor2 = 32 * (past_decoder_sequence_length + 1)
    assert ana.can_prove_equal(tvm.tirx.floormod(expr1, divisor2), 0)


def test_simplify_uint_floormod_const_scale_divisible():
    """uint32 floormod(x * c1, c2) -> 0 when c1 % c2 == 0 (overflow-free)."""
    ana = tvm.arith.Analyzer()
    q = tirx.Var("q_stage_idx", "uint32")
    expr = q * tirx.Cast("uint32", 128)
    mod = expr % tirx.const(4, "uint32")
    assert ana.can_prove_equal(mod, tirx.const(0, "uint32"))
    tvm.ir.assert_structural_equal(ana.rewrite_simplify(mod), tirx.const(0, "uint32"))


def test_simplify_float_division():
    # Test for the discussion:
    # https://discuss.tvm.apache.org/t/discuss-is-constant-division-to-multiplication-rewrite-in-tvm-necessary/18615
    ana = tvm.arith.Analyzer()
    x = tirx.Var("x", "float32")
    ry = x / 27
    # in old version, the division will be rewritten into x * T.float32(1 / 27)
    sy = ana.rewrite_simplify(ry)
    tvm.ir.assert_structural_equal(ry, sy)


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