import tilelang.testing
from tilelang import carver
from tilelang.carver.roller import PrimFuncNode, OutputNode, Edge
from tilelang.carver.arch import auto_infer_current_arch
from tvm import te


def run_general_matmul_emit_configs(M, N, K, topk: int = 20):
    arch = auto_infer_current_arch()

    def gemm(M, N, K):
        A = te.placeholder((M, K), name="A", dtype="float16")
        B = te.placeholder((N, K), name="B", dtype="float16")

        # Describe the matrix multiplication in TE
        k = te.reduce_axis((0, K), name="k")

        C = te.compute((M, N), lambda i, j: te.sum(A[i, k].astype("float16") * B[j, k].astype("float16"), axis=[k]), name="C")

        return A, B, C

    arg1 = gemm(M, N, K)
    args = arg1

    func = te.create_prim_func(args)

    tensorized_func, tags = carver.utils.get_tensorized_func_and_tags(func, arch.target)
    print(tags)
    policy = carver.TensorCorePolicy.from_prim_func(func=tensorized_func, arch=arch, tags=tags, name="matmul_0")

    hints = policy.emit_config(topk=topk)

    for hint in hints:
        print(hint)

    assert len(hints) > 0, "Hints length is zero"

    prim_func_node = PrimFuncNode(tensorized_func, name="matmul_1")
    output_nodes = [OutputNode(prim_func_node)]
    policy = carver.TensorCorePolicy.from_output_nodes(output_nodes, arch=arch, tags=tags)

    hints = policy.emit_config(topk=10)

    for config in hints:
        print(config)

    assert len(hints) > 0, "Hints length is zero"


def test_general_matmul_emit_configs():
    run_general_matmul_emit_configs(128, 128, 128)


def run_general_matmul_matmul_emit_configs(M, N, K, topk: int = 20):
    arch = auto_infer_current_arch()

    def gemm(M, N, K):
        A = te.placeholder((M, K), name="A", dtype="float16")
        B = te.placeholder((N, K), name="B", dtype="float16")

        # Describe the matrix multiplication in TE
        k = te.reduce_axis((0, K), name="k")

        C = te.compute((M, N), lambda i, j: te.sum(A[i, k].astype("float16") * B[j, k].astype("float16"), axis=[k]), name="C")

        return A, B, C

    arg1 = gemm(M, N, K)
    args = arg1

    func = te.create_prim_func(args)

    tensorized_func, tags = carver.utils.get_tensorized_func_and_tags(func, arch.target)
    print(tags)

    node_0 = PrimFuncNode(tensorized_func, name="matmul_0")
    node_1 = PrimFuncNode(tensorized_func, name="matmul_1")

    edge = Edge(node_0, node_1, 0, 0)
    node_0._out_edges.append(edge)
    node_1.set_inputs(0, edge)

    output_nodes = [OutputNode(node_1)]
    policy = carver.TensorCorePolicy.from_output_nodes(output_nodes, arch=arch, tags=tags)

    hints = policy.emit_config(topk=topk)

    for config in hints:
        print(config)

    assert len(hints) > 0, "Hints length is zero"


def test_general_matmul_matmul_emit_configs():
    run_general_matmul_matmul_emit_configs(128, 128, 128)


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