# RUN: %PYTHON %s

from mlir.dialects import arith, func, linalg
from mlir.dialects.linalg.opdsl.lang import *
from mlir.ir import *


def run(f):
    print("\nTEST:", f.__name__)
    f()
    return f


@run
def test_infer_contraction_dimensions_from_ops():
    with Context(), Location.unknown():
        module = Module.create()
        f32 = F32Type.get()
        with InsertionPoint(module.body):
            # === Static shapes ===
            m, n, k = 4, 4, 4
            a_type = RankedTensorType.get((m, k), f32)
            b_type = RankedTensorType.get((k, n), f32)
            c_type = RankedTensorType.get((m, n), f32)

            @func.FuncOp.from_py_func(a_type, b_type, c_type)
            def contraction_fn(a, b, c):
                zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.0), result=f32)
                filled = linalg.fill(zero, outs=[c])
                fill_op = filled.owner

                assert not linalg.isa_contraction_op(zero.operation)
                assert not linalg.isa_contraction_op(fill_op)
                assert linalg.infer_contraction_dimensions(fill_op) is None

                dim_m = AffineDimExpr.get(0)
                dim_n = AffineDimExpr.get(1)
                dim_k = AffineDimExpr.get(2)

                a_map = AffineMap.get(3, 0, [dim_m, dim_k])
                b_map = AffineMap.get(3, 0, [dim_k, dim_n])
                c_map = AffineMap.get(3, 0, [dim_m, dim_n])
                result = linalg.contract(
                    a,
                    b,
                    outs=(filled,),
                    indexing_maps=[a_map, b_map, c_map],
                )
                contraction_op = result.owner

                assert linalg.isa_contraction_op(contraction_op)
                dims = linalg.infer_contraction_dimensions(contraction_op)
                assert dims is not None

                # Expect m=[0], n=[1], k=[2] as per standard matmul.
                assert list(dims.m) == [0], f"Expected m=[0], got {list(dims.m)}"
                assert list(dims.n) == [1], f"Expected n=[1], got {list(dims.n)}"
                assert list(dims.k) == [2], f"Expected k=[2], got {list(dims.k)}"
                assert (
                    list(dims.batch) == []
                ), f"Expected batch=[], got {list(dims.batch)}"

            # === Dynamic shape case ===
            dyn = ShapedType.get_dynamic_size()
            a_dyn_type = RankedTensorType.get((4, dyn), f32)
            b_dyn_type = RankedTensorType.get((dyn, 4), f32)
            c_type = RankedTensorType.get((4, 4), f32)

            @func.FuncOp.from_py_func(a_dyn_type, b_dyn_type, c_type)
            def dynamic_contraction_fn(a, b, c):
                zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.0), result=f32)
                filled = linalg.fill(zero, outs=[c])
                dim_m = AffineDimExpr.get(0)
                dim_n = AffineDimExpr.get(1)
                dim_k = AffineDimExpr.get(2)

                a_map = AffineMap.get(3, 0, [dim_m, dim_k])
                b_map = AffineMap.get(3, 0, [dim_k, dim_n])
                c_map = AffineMap.get(3, 0, [dim_m, dim_n])

                result = linalg.contract(
                    a,
                    b,
                    outs=(filled,),
                    indexing_maps=[a_map, b_map, c_map],
                )
                contraction_op = result.owner

                assert linalg.isa_contraction_op(contraction_op)
                dims = linalg.infer_contraction_dimensions(contraction_op)
                assert dims is not None
                assert list(dims.m) == [0], f"Expected m=[0], got {list(dims.m)}"
                assert list(dims.n) == [1], f"Expected n=[1], got {list(dims.n)}"
                assert list(dims.k) == [2], f"Expected k=[2], got {list(dims.k)}"
                assert (
                    list(dims.batch) == []
                ), f"Expected batch=[], got {list(dims.batch)}"


@run
def test_infer_convolution_dimensions_from_ops():
    with Context(), Location.unknown():
        module = Module.create()
        f32 = F32Type.get()

        with InsertionPoint(module.body):
            # === Static shapes ===
            batch, h, w, c_in, kh, kw, c_out = 1, 8, 8, 4, 3, 3, 16
            input_type = RankedTensorType.get((batch, h, w, c_in), f32)
            filter_type = RankedTensorType.get((kh, kw, c_in, c_out), f32)
            output_type = RankedTensorType.get(
                (batch, h - kh + 1, w - kw + 1, c_out), f32
            )

            @func.FuncOp.from_py_func(input_type, filter_type, output_type)
            def conv_fn(input, filter, output):
                zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.0), result=f32)
                filled = linalg.fill(zero, outs=[output])
                fill_op = filled.owner

                assert not linalg.isa_convolution_op(fill_op)
                assert linalg.infer_convolution_dimensions(fill_op) is None

                result = linalg.conv_2d_nhwc_hwcf(input, filter, outs=[filled])
                conv_op = result.owner

                assert linalg.isa_convolution_op(conv_op)
                dims = linalg.infer_convolution_dimensions(conv_op)
                assert dims is not None
                assert list(dims.batch) == [0]
                assert list(dims.output_image) == [1, 2]
                assert list(dims.output_channel) == [3]
                assert list(dims.filter_loop) == [4, 5]
                assert list(dims.input_channel) == [6]
                assert list(dims.depth) == []
                assert list(dims.strides) == [1, 1]
                assert list(dims.dilations) == [1, 1]

            # === Dynamic shapes ===
            dyn = ShapedType.get_dynamic_size()
            dyn_input_type = RankedTensorType.get((batch, dyn, dyn, c_in), f32)
            dyn_output_type = RankedTensorType.get((batch, dyn, dyn, c_out), f32)

            @func.FuncOp.from_py_func(dyn_input_type, filter_type, dyn_output_type)
            def dyn_conv_fn(input, filter, output):
                zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.0), result=f32)
                filled = linalg.fill(zero, outs=[output])
                result = linalg.conv_2d_nhwc_hwcf(input, filter, outs=[filled])
                conv_op = result.owner

                assert linalg.isa_convolution_op(conv_op)
                dims = linalg.infer_convolution_dimensions(conv_op)
                assert dims is not None
                assert list(dims.batch) == [0]
                assert list(dims.output_image) == [1, 2]
                assert list(dims.output_channel) == [3]
                assert list(dims.filter_loop) == [4, 5]
                assert list(dims.input_channel) == [6]
                assert list(dims.depth) == []
                assert list(dims.strides) == [1, 1]
                assert list(dims.dilations) == [1, 1]


@run
def test_get_indexing_maps_attr():
    with Context(), Location.unknown():
        module = Module.create()
        f32 = F32Type.get()
        with InsertionPoint(module.body):
            a_type = RankedTensorType.get((4, 8), f32)
            b_type = RankedTensorType.get((8, 16), f32)
            c_type = RankedTensorType.get((4, 16), f32)

            dim_m = AffineDimExpr.get(0)
            dim_n = AffineDimExpr.get(1)
            dim_k = AffineDimExpr.get(2)

            a_map = AffineMap.get(3, 0, [dim_m, dim_k])
            b_map = AffineMap.get(3, 0, [dim_k, dim_n])
            c_map = AffineMap.get(3, 0, [dim_m, dim_n])

            @func.FuncOp.from_py_func(a_type, b_type, c_type)
            def matmul_func(a, b, c):
                zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.0), result=f32)
                assert not linalg.get_indexing_maps(
                    zero.operation
                ), "Expected no indexing_maps on non-linalg op"

                init = linalg.fill(zero, outs=[c])
                fill_op = init.owner
                fill_maps = linalg.get_indexing_maps(fill_op)
                assert fill_maps is not None
                assert len(fill_maps) == 2

                # The fill op should have maps like (d0, d1) -> () and (d0, d1).
                fill_input_map = fill_maps[0].value
                fill_output_map = fill_maps[1].value
                assert fill_input_map == AffineMap.get(2, 0, [])
                assert fill_output_map == AffineMap.get(2, 0, [dim_m, dim_n])

                result = linalg.matmul(a, b, outs=(init,))
                matmul_op = result.owner
                matmul_maps = linalg.get_indexing_maps(matmul_op)
                assert matmul_maps is not None
                assert len(matmul_maps) == 3

                maps = [map_attr.value for map_attr in matmul_maps]
                assert maps[0] == a_map
                assert maps[1] == b_map
                assert maps[2] == c_map


@run
def test_infer_contraction_dimensions_from_maps():
    with Context(), Location.unknown():
        module = Module.create()
        with InsertionPoint(module.body):
            # === Test valid contraction (matmul) ===
            dim_m = AffineDimExpr.get(0)
            dim_n = AffineDimExpr.get(1)
            dim_k = AffineDimExpr.get(2)
            a_map = AffineMap.get(3, 0, [dim_m, dim_k])
            b_map = AffineMap.get(3, 0, [dim_k, dim_n])
            c_map = AffineMap.get(3, 0, [dim_m, dim_n])

            dims = linalg.infer_contraction_dimensions_from_maps([a_map, b_map, c_map])
            assert dims is not None

            # Expect m=[0], n=[1], k=[2] as per standard matmul.
            assert list(dims.m) == [0], f"Expected m=[0], got {list(dims.m)}"
            assert list(dims.n) == [1], f"Expected n=[1], got {list(dims.n)}"
            assert list(dims.k) == [2], f"Expected k=[2], got {list(dims.k)}"
            assert list(dims.batch) == [], f"Expected batch=[], got {list(dims.batch)}"

            # === Test invalid input (wrong number of maps) ===
            invalid_dims = linalg.infer_contraction_dimensions_from_maps([a_map, b_map])
            assert invalid_dims is None

            # === Test element-wise operation ===
            dim_i = AffineDimExpr.get(0)
            dim_j = AffineDimExpr.get(1)
            elementwise_map = AffineMap.get(2, 0, [dim_i, dim_j])
            elementwise_dims = linalg.infer_contraction_dimensions_from_maps(
                [elementwise_map, elementwise_map, elementwise_map]
            )
            assert elementwise_dims is not None
            assert len(elementwise_dims.m) == 0
            assert len(elementwise_dims.n) == 0
            assert len(elementwise_dims.k) == 0
            assert list(elementwise_dims.batch) == [0, 1]