# 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.
# pylint: disable=invalid-name, missing-docstring
# ruff: noqa: F401, F841
"""Unittests for tvm.script.ir_builder.tirx"""

import numpy as np
import pytest

import tvm
import tvm.runtime
import tvm.testing
from tvm import tirx
from tvm.ir.base import assert_structural_equal
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import tirx as T


def test_ir_builder_tir_primfunc_base():
    with IRBuilder() as ib:
        with T.prim_func(s_tir=True):
            T.evaluate(0)

    # the prim_func generated by IRBuilder
    prim_func_actual = ib.get()

    # the expected prim_func
    prim_func_expected = tirx.PrimFunc(
        params=[],
        body=tirx.Evaluate(0),
        ret_type=None,
        buffer_map=None,
        attrs=tvm.ir.make_node("ir.DictAttrs", s_tir=True),
    )

    # Check if the generated ir is expected
    assert_structural_equal(prim_func_actual, prim_func_expected, map_free_vars=True)


def test_ir_builder_tir_primfunc_complete():
    with IRBuilder() as ib:
        with T.prim_func(s_tir=True):
            T.arg("a", T.handle())
            T.arg("b", T.int64())
            T.arg("c", T.Buffer((128, 128), "float32"))
            d = T.arg("d", T.handle())
            e = T.arg("e", T.Buffer((1024,), "int8"))
            T.func_attr({"key": "value"})
            T.func_ret(tvm.ir.PrimType("int64"))
            buffer_d = T.match_buffer(d, (64, 64), "int64")
            T.evaluate(0)

    # the prim_func generated by IRBuilder
    prim_func_actual = ib.get()

    # the expected prim_func
    c_handle, c_buffer = (
        tirx.Var("c_handle", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
        tirx.decl_buffer((128, 128), "float32", name="c", layout=None),
    )
    d_handle, d_buffer = (
        tirx.Var("d", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
        tirx.decl_buffer((64, 64), "int64", name="d", layout=None),
    )
    e_handle, e_buffer = (
        tirx.Var("e_handle", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
        tirx.decl_buffer((1024,), "int8", name="e", layout=None),
    )
    prim_func_expected = tirx.PrimFunc(
        params=[
            tirx.Var("a", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
            tirx.Var("b", "int64"),
            c_handle,
            d_handle,
            e_handle,
        ],
        body=tirx.Evaluate(0),
        ret_type=tvm.ir.PrimType("int64"),
        buffer_map={c_handle: c_buffer, d_handle: d_buffer, e_handle: e_buffer},
        attrs=tvm.ir.make_node("ir.DictAttrs", key="value", s_tir=True),
    )

    # Check if the generated ir is expected
    assert_structural_equal(prim_func_actual, prim_func_expected, map_free_vars=True)


def test_ir_builder_tir_block_base():
    with IRBuilder() as ib:
        with T.sblock("block"):
            T.evaluate(0)

    # the block generated by IRBuilder
    block_realize_actual = ib.get()

    # the expected block
    block_expected = tirx.SBlock(
        iter_vars=[],
        reads=[],
        writes=[],
        name_hint="block",
        body=tirx.Evaluate(0),
        alloc_buffers=None,
        match_buffers=None,
        annotations={"tirx.script_parsing_detect_access": tirx.IntImm("int64", 3)},
    )
    block_realize_expected = tirx.SBlockRealize(
        iter_values=[],
        predicate=True,
        block=block_expected,
    )

    # Check if the generated ir is expected
    assert_structural_equal(block_realize_actual, block_realize_expected, map_free_vars=True)


def test_ir_builder_tir_block_complete():
    with IRBuilder() as ib:
        a = T.int64()
        b = T.Buffer((128, 128), "float32")
        c = T.Buffer((128, 128), "float32")
        d = T.int32()
        e = T.Buffer((128, 128), "float32")
        f = T.int32()
        with T.sblock("block"):
            T.where(a > 1)
            T.reads(b[0:16, 0:16])
            T.writes(c[d:128, d:128])
            T.sblock_attr({"key": "value"})
            T.sblock_alloc_buffer((128, 128), "float32")
            T.match_buffer(e[0:32, 0:32], (32, 32), "float32")
            T.axis.spatial(128, f)
            T.evaluate(0)

    # the block generated by IRBuilder
    block_realize_actual = ib.get()

    # the expected block
    var_a = tirx.Var("a", "int64")
    buffer_b = tirx.decl_buffer((128, 128), "float32", name="b")
    buffer_c = tirx.decl_buffer((128, 128), "float32", name="c")
    var_d = tirx.Var("d", "int32")
    buffer_e = tirx.decl_buffer((128, 128), "float32", name="c")
    var_f = tirx.Var("f", "int32")
    block_expected = tirx.SBlock(
        iter_vars=[tirx.IterVar((0, 128), tirx.Var("", "int32"), iter_type=tirx.IterVar.DataPar)],
        reads=[buffer_b[0:16, 0:16]],
        writes=[buffer_c[var_d:128, var_d:128]],
        name_hint="block",
        body=tirx.Evaluate(0),
        alloc_buffers=[tirx.decl_buffer((128, 128), "float32")],
        match_buffers=[
            tirx.MatchBufferRegion(tirx.decl_buffer((32, 32), "float32"), buffer_e[0:32, 0:32])
        ],
        annotations={"key": "value"},
    )
    block_realize_expected = tirx.SBlockRealize(
        iter_values=[var_f],
        predicate=var_a > 1,
        block=block_expected,
    )

    # Check if the generated ir is expected
    assert_structural_equal(block_realize_actual, block_realize_expected, map_free_vars=True)


def test_ir_builder_tir_axis():
    with IRBuilder() as ib:
        a = T.int32()
        b = T.int32()
        c = T.int32()
        d = T.int32()
        with T.sblock("block"):
            T.axis.spatial(8, a)
            T.axis.reduce(16, b)
            T.axis.scan(32, c)
            T.axis.opaque(64, d)
            T.evaluate(0)

    # the block generated by IRBuilder
    block_realize_actual = ib.get()

    # the expected block
    var_a = tirx.Var("a", "int32")
    var_b = tirx.Var("b", "int32")
    var_c = tirx.Var("c", "int32")
    var_d = tirx.Var("d", "int32")
    block_expected = tirx.SBlock(
        iter_vars=[
            tirx.IterVar((0, 8), tirx.Var("", "int32"), iter_type=tirx.IterVar.DataPar),
            tirx.IterVar((0, 16), tirx.Var("", "int32"), iter_type=tirx.IterVar.CommReduce),
            tirx.IterVar((0, 32), tirx.Var("", "int32"), iter_type=tirx.IterVar.Ordered),
            tirx.IterVar((0, 64), tirx.Var("", "int32"), iter_type=tirx.IterVar.Opaque),
        ],
        reads=[],
        writes=[],
        name_hint="block",
        body=tirx.Evaluate(0),
        annotations={"tirx.script_parsing_detect_access": tirx.IntImm("int64", 3)},
    )
    block_realize_expected = tirx.SBlockRealize(
        iter_values=[var_a, var_b, var_c, var_d],
        predicate=True,
        block=block_expected,
    )

    # Check if the generated ir is expected
    assert_structural_equal(block_realize_actual, block_realize_expected, map_free_vars=True)


def test_ir_builder_tir_for():
    with IRBuilder() as ib:
        with T.serial(128) as a:
            with T.parallel(64) as b:
                with T.vectorized(32) as c:
                    with T.unroll(16) as d:
                        with T.thread_binding(8, thread="threadIdx.x") as e:
                            T.evaluate(0)

    # the for generated by IRBuilder
    for_actual = ib.get()

    # the expected for
    thread_binding_expected = tirx.For(
        loop_var=tirx.Var("", "int32"),
        min=0,
        extent=8,
        kind=tirx.ForKind.THREAD_BINDING,
        body=tirx.Evaluate(0),
        thread_binding=tirx.IterVar(
            None, tirx.Var("", "int32"), tirx.IterVar.ThreadIndex, "threadIdx.x"
        ),
    )
    unroll_expected = tirx.For(
        loop_var=tirx.Var("", "int32"),
        min=0,
        extent=16,
        kind=tirx.ForKind.UNROLLED,
        body=thread_binding_expected,
    )
    vectorized_expected = tirx.For(
        loop_var=tirx.Var("", "int32"),
        min=0,
        extent=32,
        kind=tirx.ForKind.VECTORIZED,
        body=unroll_expected,
    )
    parallel_expected = tirx.For(
        loop_var=tirx.Var("", "int32"),
        min=0,
        extent=64,
        kind=tirx.ForKind.PARALLEL,
        body=vectorized_expected,
    )
    for_expected = tirx.For(
        loop_var=tirx.Var("", "int32"),
        min=0,
        extent=128,
        kind=tirx.ForKind.SERIAL,
        body=parallel_expected,
    )

    # Check if the generated ir is expected
    assert_structural_equal(for_actual, for_expected, map_free_vars=True)


def test_ir_builder_tir_for_uint():
    with IRBuilder() as ib:
        with T.serial(tirx.const(128, "uint32")) as a:
            T.evaluate(0)

    # the for generated by IRBuilder
    for_actual = ib.get()

    for_expected = tirx.For(
        loop_var=tirx.Var("", "uint32"),
        min=tirx.const(0, "uint32"),
        extent=tirx.const(128, "uint32"),
        kind=tirx.ForKind.SERIAL,
        body=tirx.Evaluate(0),
    )

    # Check if the generated ir is expected
    assert_structural_equal(for_actual, for_expected, map_free_vars=True)


def test_ir_builder_tir_assert():
    with IRBuilder() as ib:
        with T.Assert(T.int32() == 0, message="a is 0"):
            T.evaluate(0)
    # the assert generated by IRBuilder
    assert_actual = ib.get()

    # AssertStmt is a leaf. The frame emits the assert and then the body stmts as siblings.
    assert_expected = tirx.SeqStmt(
        [
            tirx.AssertStmt(
                T.int32() == 0,
                tirx.StringImm("RuntimeError"),
                [tirx.StringImm("a is 0")],
            ),
            tirx.Evaluate(0),
        ]
    )

    # Check if the generated ir is expected
    assert_structural_equal(assert_actual, assert_expected, map_free_vars=True)


def test_ir_builder_tir_bind():
    # Test that T.bind emits a flat Bind statement and returns the Var.
    with IRBuilder() as ib:
        v = T.bind(tirx.IntImm("int32", 2))
    # the let binding generated by IRBuilder
    let_actual = ib.get()

    # Bind is now flat (no body), so a single Bind stmt is emitted.
    let_expected = tirx.Bind(T.int32(), tirx.IntImm("int32", 2))

    # Check if the generated ir is expected
    assert_structural_equal(let_actual, let_expected, map_free_vars=True)
    # Check that the returned value is a Var
    assert isinstance(v, tirx.Var)


def test_ir_builder_tir_thread():
    with IRBuilder() as ib:
        with T.prim_func(s_tir=True):
            brow = T.env_thread("blockIdx.y")
            with T.launch_thread(brow, 1):
                T.evaluate(0)

    # the prim_func generated by IRBuilder
    ir_actual = ib.get()

    # the expected prim_func
    iter_var = tirx.IterVar((0, 1), "v", iter_type=1, thread_tag="blockIdx.y")
    attr_stmt = tirx.AttrStmt(iter_var, "thread_extent", 1, tirx.Evaluate(0))
    func = tirx.PrimFunc([], attr_stmt).with_attr("s_tir", True)

    # Check if the generated ir is expected
    assert_structural_equal(ir_actual, func, map_free_vars=True)


def test_ir_builder_tir_allocate():
    with IRBuilder() as ib:
        with T.prim_func(s_tir=True):
            T.func_name("test")
            buf = T.alloc_buffer([10], "float32", scope="local")
            T.evaluate(1)

    # the allocate generated by IRBuilder
    ir_actual = ib.get()
    body = ir_actual.body

    # AllocBuffer is flat: body should be a SeqStmt with [AllocBuffer, Evaluate(1)]
    assert isinstance(body, tirx.SeqStmt), f"Expected SeqStmt but got {type(body)}"
    assert len(body) == 2
    assert isinstance(body[0], tirx.AllocBuffer)
    assert isinstance(body[1], tirx.Evaluate)
    assert body[1].value.value == 1


def test_ir_builder_tir_while():
    with IRBuilder() as ib:
        with T.While(T.int32() > 0):
            T.evaluate(0)

    # the while generated by IRBuilder
    ir_actual = ib.get()

    # the expected while
    ir_expected = tirx.While(tirx.Var("x", "int32") > 0, tirx.Evaluate(0))

    # Check if the generated ir is expected
    assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)


def test_ir_builder_tir_if_then_else():
    with IRBuilder() as ib:
        with T.If(T.int32() < 12):
            with T.Then():
                T.evaluate(T.int32(0))
            with T.Else():
                T.evaluate(T.int32(1))

    # the if_then_else generated by IRBuilder
    ir_actual = ib.get()

    # the expected if_then_else
    ir_expected = tirx.IfThenElse(
        tirx.Var("c", "int32") < 12,
        tirx.Evaluate(tirx.IntImm("int32", 0)),
        tirx.Evaluate(tirx.IntImm("int32", 1)),
    )

    # Check if the generated ir is expected
    assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)


def test_ir_builder_tir_buffer_store():
    buffer_a = T.Buffer((10, 10), "float32")
    i = T.int32()
    with IRBuilder() as ib:
        T.buffer_store(buffer_a, 0.1, [0, i])

    # the buffer store generated by IRBuilder
    ir_actual = ib.get()

    # the expected buffer store
    ir_expected = tirx.BufferStore(buffer_a, 0.1, [0, i])

    # Check if the generated ir is expected
    assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)


def test_ir_builder_tir_buffer_store_scalable_vec():
    buffer_a = T.Buffer((30,), "float32")
    value = T.broadcast(0.11, 4 * tvm.tirx.vscale())
    index = T.ramp(0, 1, 4 * tvm.tirx.vscale())

    with IRBuilder() as ib:
        T.buffer_store(buffer_a, value, [index])

    # the buffer store generated by IRBuilder
    ir_actual = ib.get()

    # the expected buffer store
    ir_expected = tirx.BufferStore(buffer_a, value, [index])

    # Check if the generated ir is expected
    assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)


def test_ir_builder_tir_buffer_store_predicate():
    buffer_a = T.Buffer((30,), "float32")
    value = T.broadcast(0.11, T.vscale() * 4)
    index = T.ramp(0, 1, T.vscale() * 4)
    predicate = T.broadcast(T.bool(True), T.vscale() * 4)

    with IRBuilder() as ib:
        T.buffer_store(buffer_a, value, [index], predicate)

    ir_actual = ib.get()
    ir_expected = tirx.BufferStore(buffer_a, value, [index], predicate)
    assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)


def test_ir_builder_tir_evaluate():
    with IRBuilder() as ib:
        T.evaluate(0)
    # the evaluate generated by IRBuilder
    eval_actual = ib.get()

    # the expected evaluate
    eval_expected = tirx.Evaluate(0)

    # Check if the generated ir is expected
    assert_structural_equal(eval_actual, eval_expected, map_free_vars=True)


def test_ir_builder_tir_decl_buffer():
    with IRBuilder() as ib:
        with T.prim_func(s_tir=True):
            T.func_name("test")
            buf = T.decl_buffer([128, 128], "float32")
            T.evaluate(1)

    # the decl_buffer generated by IRBuilder
    ir_actual = ib.get()
    body = ir_actual.body

    # decl_buffer without data emits AllocBuffer (flat): body should be SeqStmt
    assert isinstance(body, tirx.SeqStmt), f"Expected SeqStmt but got {type(body)}"
    assert len(body) == 2
    assert isinstance(body[0], tirx.AllocBuffer)
    assert isinstance(body[1], tirx.Evaluate)
    assert body[1].value.value == 1


def test_ir_builder_tir_inline():
    with IRBuilder() as ib:
        m, n = T.meta_var(1), T.meta_var(2)
        a, b = T.meta_var([3, 4])
        T.evaluate(m.value + n.value + a.value + b.value)
    # the evaluate generated by IRBuilder
    eval_actual = ib.get()

    # the expected evaluate
    eval_expected = tirx.Evaluate(10)

    # Check if the generated ir is expected
    assert_structural_equal(eval_actual, eval_expected, map_free_vars=True)


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