# RUN: env SUPPORT_LIB=%mlir_c_runner_utils \
# RUN:   %PYTHON %s | FileCheck %s

import ctypes
import os
import sys
import tempfile

from mlir import ir
from mlir import runtime as rt
from mlir.dialects import builtin
from mlir.dialects import sparse_tensor as st

_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
sys.path.append(_SCRIPT_PATH)
from tools import sparsifier


def boilerplate(attr: st.EncodingAttr):
    """Returns boilerplate main method."""
    return f"""
func.func @main(%p : !llvm.ptr) -> () attributes {{ llvm.emit_c_interface }} {{
  %d = arith.constant sparse<[[0, 0], [1, 1], [0, 9], [9, 0], [4, 4]],
                             [1.0, 2.0, 3.0, 4.0, 5.0]> : tensor<10x10xf64>
  %a = sparse_tensor.convert %d : tensor<10x10xf64> to tensor<10x10xf64, {attr}>
  sparse_tensor.out %a, %p : tensor<10x10xf64, {attr}>, !llvm.ptr
  return
}}
"""


def expected(id_map):
    """Returns expected contents of output.

    +-----+-----+-----+-----+-----+
    | 1 0 | . . | . . | . . | 0 3 |
    | 0 2 | . . | . . | . . | 0 0 |
    +-----+-----+-----+-----+-----+
    | . . | . . | . . | . . | . . |
    | . . | . . | . . | . . | . . |
    +-----+-----+-----+-----+-----+
    | . . | . . | 5 0 | . . | . . |
    | . . | . . | 0 0 | . . | . . |
    +-----+-----+-----+-----+-----+
    | . . | . . | . . | . . | . . |
    | . . | . . | . . | . . | . . |
    +-----+-----+-----+-----+-----+
    | 0 0 | . . | . . | . . | . . |
    | 4 0 | . . | . . | . . | . . |
    +-----+-----+-----+-----+-----+

    Output appears as dimension coordinates but lexicographically
    sorted by level coordinates. For BSR, the blocks are filled.
    """
    if id_map is 0:
        return f"""# extended FROSTT format
2 5
10 10
1 1 1
1 10 3
2 2 2
5 5 5
10 1 4
"""
    if id_map is 1:
        return f"""# extended FROSTT format
2 5
10 10
1 1 1
10 1 4
2 2 2
5 5 5
1 10 3
"""
    if id_map is 2:
        return f"""# extended FROSTT format
2 16
10 10
1 1 1
1 2 0
2 1 0
2 2 2
1 9 0
1 10 3
2 9 0
2 10 0
5 5 5
5 6 0
6 5 0
6 6 0
9 1 0
9 2 0
10 1 4
10 2 0
"""
    raise AssertionError("unexpected id_map")


def build_compile_and_run_output(attr: st.EncodingAttr, compiler, expected):
    # Build and Compile.
    module = ir.Module.parse(boilerplate(attr))
    engine = compiler.compile_and_jit(module)
    # Invoke the kernel and compare output.
    with tempfile.TemporaryDirectory() as test_dir:
        out = os.path.join(test_dir, "out.tns")
        buf = out.encode("utf-8")
        mem_a = ctypes.pointer(ctypes.pointer(ctypes.create_string_buffer(buf)))
        engine.invoke("main", mem_a)
        actual = open(out).read()
        if actual != expected:
            quit("FAILURE")


def main():
    support_lib = os.getenv("SUPPORT_LIB")
    assert support_lib is not None, "SUPPORT_LIB is undefined"
    if not os.path.exists(support_lib):
        raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)

    # CHECK-LABEL: TEST: test_output
    print("\nTEST: test_output")
    count = 0
    with ir.Context() as ctx, ir.Location.unknown():
        # Loop over various sparse types (COO, CSR, DCSR, CSC, DCSC) with
        # regular and loose compression and various metadata bitwidths.
        # For these simple orderings, dim2lvl and lvl2dim are the same.
        builder = st.EncodingAttr.build_level_type
        fmt = st.LevelFormat
        prop = st.LevelProperty
        levels = [
            [builder(fmt.compressed, [prop.non_unique]), builder(fmt.singleton)],
            [
                builder(fmt.compressed, [prop.non_unique]),
                builder(fmt.singleton, [prop.soa]),
            ],
            [builder(fmt.dense), builder(fmt.compressed)],
            [builder(fmt.dense), builder(fmt.loose_compressed)],
            [builder(fmt.compressed), builder(fmt.compressed)],
        ]
        orderings = [
            (ir.AffineMap.get_permutation([0, 1]), 0),
            (ir.AffineMap.get_permutation([1, 0]), 1),
        ]
        bitwidths = [8, 64]
        compiler = sparsifier.Sparsifier(
            extras="", options="", opt_level=2, shared_libs=[support_lib]
        )
        for level in levels:
            for ordering, id_map in orderings:
                for bwidth in bitwidths:
                    attr = st.EncodingAttr.get(
                        level, ordering, ordering, bwidth, bwidth
                    )
                    build_compile_and_run_output(attr, compiler, expected(id_map))
                    count = count + 1

        # Now do the same for BSR.
        level = [
            builder(fmt.dense),
            builder(fmt.compressed),
            builder(fmt.dense),
            builder(fmt.dense),
        ]
        d0 = ir.AffineDimExpr.get(0)
        d1 = ir.AffineDimExpr.get(1)
        c2 = ir.AffineConstantExpr.get(2)
        dim2lvl = ir.AffineMap.get(
            2,
            0,
            [
                ir.AffineExpr.get_floor_div(d0, c2),
                ir.AffineExpr.get_floor_div(d1, c2),
                ir.AffineExpr.get_mod(d0, c2),
                ir.AffineExpr.get_mod(d1, c2),
            ],
        )
        l0 = ir.AffineDimExpr.get(0)
        l1 = ir.AffineDimExpr.get(1)
        l2 = ir.AffineDimExpr.get(2)
        l3 = ir.AffineDimExpr.get(3)
        lvl2dim = ir.AffineMap.get(4, 0, [2 * l0 + l2, 2 * l1 + l3])
        attr = st.EncodingAttr.get(level, dim2lvl, lvl2dim, 0, 0)
        build_compile_and_run_output(attr, compiler, expected(2))
        count = count + 1

    # CHECK: Passed 21 tests
    print("Passed", count, "tests")


if __name__ == "__main__":
    main()