# 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, F401, F841

import numpy as np
import pytest

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


def test_buffer():
    m = tvm.tirx.Var("m", "int32")
    n = tvm.tirx.Var("n", "int32")
    l = tvm.tirx.Var("l", "int32")
    Ab = tvm.tirx.decl_buffer((m, n), "float32")
    Bb = tvm.tirx.decl_buffer((n, l), "float32")

    assert type(Ab) is tvm.ir.Var
    assert tvm.tirx.is_buffer_var(Ab)
    assert isinstance(Ab.ty, tvm.tirx.BufferType)
    assert Ab.ty.dtype == tvm.ir.PrimType("float32")
    assert tuple(Ab.ty.shape) == (m, n)
    assert not tvm.tirx.is_buffer_var(m)


def test_buffer_compatibility_alias_and_global_var_properties():
    scalar = tvm.ir.Var("scalar", tvm.ir.PrimType("int32"))
    buffer = tvm.tirx.decl_buffer((8,), "float32")

    assert tvm.tirx.Buffer is tvm.ir.Var
    assert isinstance(scalar, tvm.tirx.Buffer)
    assert not tvm.tirx.is_buffer_var(scalar)
    assert tvm.tirx.is_buffer_var(buffer)

    assert tuple(buffer.shape) == (8,)
    assert buffer.dtype == tvm.DataType("float32")
    assert buffer.data.args[0].same_as(buffer)

    for name in ("shape", "dtype", "data"):
        assert not hasattr(scalar, name)
        with pytest.raises(AttributeError, match="only available on a Var with BufferType"):
            getattr(scalar, name)


def test_buffer_data_is_typed_projection():
    buffer = tvm.tirx.decl_buffer((8,), "bool", scope="shared")

    assert buffer.ty.dtype == tvm.ir.PrimType("bool")
    assert tvm.tirx.buffer_data_pointer_type(buffer) == tvm.ir.PointerType(
        tvm.ir.PrimType("bool"), "shared"
    )
    assert buffer.data.op.name == "tirx.buffer_data"
    assert buffer.data.args[0].same_as(buffer)
    assert buffer.data.ty == tvm.tirx.buffer_data_pointer_type(buffer)


def test_buffer_pointer_type_derived_from_dtype_and_scope():
    data = tvm.ir.Var("storage", tvm.ir.PointerType(tvm.ir.PrimType("uint8"), "local"))
    buffer = tvm.tirx.decl_buffer((8,), "float16", data=data)

    assert buffer.ty.dtype == tvm.ir.PrimType("float16")
    assert buffer.ty.storage_scope == "local"
    assert buffer.data.ty == tvm.ir.PointerType(tvm.ir.PrimType("float16"), "local")


def test_decl_buffer_requires_physical_data_binding():
    buffer = tvm.tirx.decl_buffer((8,), "float32")
    data = tvm.tirx.Var("data", buffer.data.ty)

    with pytest.raises(TypeError, match="requires a physical data binding"):
        tvm.tirx.DeclBuffer(buffer)

    decl = tvm.tirx.DeclBuffer(buffer, data=data)
    assert decl.buffer.same_as(buffer)
    assert decl.data.same_as(data)


def test_buffer_access_ptr():
    m = tvm.tirx.Var("m", "int32")
    n = tvm.tirx.Var("n", "int32")
    Ab = tvm.tirx.decl_buffer((m, n), "float32", strides=[n + 1, 1])
    aptr = Ab.access_ptr("rw")
    assert isinstance(aptr.ty, tvm.ir.PointerType)
    assert aptr.ty.element_type == tvm.ir.PrimType("void")
    tvm.ir.assert_structural_equal(aptr.args[3], Ab.ty.strides[0] * m)
    assert aptr.args[0].ty == Ab.ty.dtype
    assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE
    typed_ptr = Ab.access_ptr("r", ptr_type="uint8")
    assert typed_ptr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8"))
    shared = tvm.tirx.decl_buffer((m, n), "float32", scope="shared")
    assert shared.access_ptr("r").ty == tvm.ir.PointerType(tvm.ir.PrimType("void"), "shared")
    assert shared.access_ptr("r", ptr_type="uint8").ty == tvm.ir.PointerType(
        tvm.ir.PrimType("uint8"), "shared"
    )
    aptr = Ab.access_ptr("w")
    assert aptr.args[4].value == BufferAccessKind.WRITE


def test_buffer_access_ptr_offset():
    m = tvm.tirx.Var("m", "int32")
    n = tvm.tirx.Var("n", "int32")
    Ab = tvm.tirx.decl_buffer((m, n), "float32")
    aptr = Ab.access_ptr("rw", offset=100)
    tvm.testing.assert_prim_expr_equal(aptr.args[2], 100)
    assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE
    v = tvm.tirx.Var("int32", "int32")
    aptr = Ab.access_ptr("rw", offset=100 + 100 + v)
    tvm.testing.assert_prim_expr_equal(aptr.args[2], 200 + v)
    assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE
    aptr = Ab.access_ptr("rw", offset=tvm.tirx.call_extern("int32", "test_call", 100 + 100 + v))
    tvm.testing.assert_prim_expr_equal(
        aptr.args[2], tvm.tirx.call_extern("int32", "test_call", 200 + v)
    )
    assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE


def test_buffer_access_ptr_extent():
    m = tvm.tirx.Var("m", "int32")
    n = tvm.tirx.Var("n", "int32")
    Ab = tvm.tirx.decl_buffer((m, n), "float32")
    aptr = Ab.access_ptr("rw")
    tvm.ir.assert_structural_equal(aptr.args[3], m * n)
    aptr = Ab.access_ptr("rw", offset=100)
    tvm.ir.assert_structural_equal(aptr.args[3], m * n - 100)
    Ab = tvm.tirx.decl_buffer((m, n), "float32", strides=[n + 1, 1])
    aptr = Ab.access_ptr("rw", offset=100)
    tvm.ir.assert_structural_equal(aptr.args[3], Ab.ty.strides[0] * m - 100)

    # Test extent from input params
    aptr = Ab.access_ptr("rw", extent=200)
    tvm.ir.assert_structural_equal(aptr.args[3], T.int32(200))
    aptr = Ab.access_ptr("rw", offset=100, extent=100)
    tvm.ir.assert_structural_equal(aptr.args[3], T.int32(100))


def test_buffer_vload():
    m = tvm.tirx.Var("m", "int32")
    n = tvm.tirx.Var("n", "int32")
    Ab = tvm.tirx.decl_buffer((m, n), "float32", elem_offset=100)
    load = Ab.vload([2, 3])
    tvm.ir.assert_structural_equal(load.indices, [T.int32(2), T.int32(3)])


def test_buffer_offset_of():
    m = tvm.tirx.Var("m", "int32")
    n = tvm.tirx.Var("n", "int32")
    Ab = tvm.tirx.decl_buffer((m, n), "float32", elem_offset=100)
    offset = Ab.offset_of([2, 3])
    tvm.ir.assert_structural_equal(offset, [n * 2 + 103])


def test_buffer_index_merge_mult_mod():
    m = tvm.tirx.Var("m", "int32")
    n = tvm.tirx.Var("n", "int32")
    s = tvm.tirx.Var("s", "int32")
    k0 = tvm.tirx.Var("k0", "int32")
    k1 = tvm.tirx.Var("k1", "int32")
    A = tvm.tirx.decl_buffer((m, n), "float32")
    A_stride = tvm.tirx.decl_buffer((m, n), "float32", strides=(s, 1))

    def assert_simplified_equal(index_simplified, index_direct):
        (
            tvm.ir.assert_structural_equal(index_simplified, index_direct),
            f"index_simplified={index_simplified}, index_direct={index_direct}",
        )

    idxd = tvm.tirx.indexdiv
    idxm = tvm.tirx.indexmod

    # Test Case1
    index_simplified = A_stride.offset_of(
        (idxd(idxm(k0, k1), s), idxm(idxm(k0, k1), s) + idxd(k0, k1) * k1)
    )
    index_direct = A_stride.offset_of((0, k0))
    assert_simplified_equal(index_simplified, index_direct)

    # Test Case2
    index_simplified = A.offset_of(
        (idxd(idxm(k0, idxd(k1, s)), n), idxm(idxm(k0, idxd(k1, s)), n) + idxm(k0, k1))
    )
    index_direct = A.offset_of((0, idxm(k0, idxd(k1, s)) + idxm(k0, k1)))
    assert_simplified_equal(index_simplified, index_direct)
    # Test Case3
    index_simplified = A.offset_of(
        (
            idxd((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n) + idxd(idxm(k0, idxd(k1, s)), n),
            idxm((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n) + idxm(idxm(k0, idxd(k1, s)), n),
        )
    )
    index_direct = A.offset_of((0, k0))
    assert_simplified_equal(index_simplified, index_direct)
    # Test Case4 (not able to simplify)
    index_simplified = A.offset_of(
        (idxd(idxm(k0, idxd(k1, s)), n), idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1))
    )
    index_direct = A.offset_of(
        (0, idxd(idxm(k0, idxd(k1, s)), n) * n + (idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1)))
    )
    assert_simplified_equal(index_simplified, index_direct)

    # Test Case5
    B = tvm.tirx.decl_buffer((1, 14, 14, 1024))
    i = tvm.tirx.Var("i", "int32")
    j = tvm.tirx.Var("j", "int32")
    k = tvm.tirx.Var("k", "int32")

    index_simplified1 = B.offset_of(
        (
            idxd(idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14),
            idxm(idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14),
            idxm(idxd((i * 50176 + j * 28672 + k), 1024), 14),
            idxm((i * 50176 + j * 28672 + k), 1024),
        )
    )
    index_simplified2 = B.offset_of(
        (
            idxd(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
            idxm(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
            idxm(i * 7 + idxd(k, 1024), 14),
            idxm(k, 1024),
        )
    )
    index_direct = B.offset_of((0, 0, 0, (i * 50176 + j * 28672 + k)))
    assert_simplified_equal(index_simplified1, index_direct)
    assert_simplified_equal(index_simplified2, index_direct)


def test_buffer_flatten():
    """A buffer should flatten to a 1-d shape"""
    buf = tvm.tirx.decl_buffer([16, 32])
    flat = buf.get_flattened_buffer()
    # A metadata-changing rewrite creates a fresh typed Var.  The physical
    # pointer is always derived from that Var instead of being stored as a
    # second buffer identity.
    assert not buf.same_as(flat)
    assert flat.data.args[0].same_as(flat)
    assert flat.data.op.name == "tirx.buffer_data"
    tvm.ir.assert_structural_equal(flat.ty.shape, [T.int32(16 * 32)])


def test_buffer_flatten_preserves_identity():
    """Flattening a 1-d buffer should return the original"""
    buf = tvm.tirx.decl_buffer([16])
    flat = buf.get_flattened_buffer()
    assert buf.same_as(flat)


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