import hidet
from hidet.graph.operator import Operator
from hidet.ir.compute import TensorInput, compute, reduce
from hidet.ir.task import Task, task_compiled_func_type
from hidet.runtime import CompiledFunction
from hidet.graph.ops.definitions.utils import input_like


class NaiveMatmulTask(Task):
    def __init__(self, a: TensorInput, b: TensorInput):
        if a.ndim != 2 or b.ndim != 2:
            raise ValueError('Only support matrix multiplication')
        c = compute(
            name='c',
            shape=[a.shape[0], b.shape[1]],
            fcompute=lambda i, j: reduce(
                shape=[a.shape[1]],
                fcompute=lambda k: a[i, k] * b[k, j],
                reduce_type='sum'
            )
        )
        super().__init__(
            name='naive_matmul', inputs=[a, b], outputs=[c],
            arguments=[
                a,
                b,
                c,
                a.shape[0],
                b.shape[1],
                a.shape[1]
            ]
        )

    def build(self, target: str) -> CompiledFunction:
        from hidet.backend import load_lib_func
        return load_lib_func(
            lib_path='./naive_matmul.so',
            func_name='naive_matmul',
            func_type=task_compiled_func_type(self)
        )


class NaiveMatmulOp(Operator):
    def __init__(self, a: hidet.Tensor, b: hidet.Tensor):
        super().__init__(
            inputs=[a, b],
            task=NaiveMatmulTask(input_like(a, 'a'), input_like(b, 'b'))
        )


def naive_matmul(a, b):
    return NaiveMatmulOp(a, b).get_output(0)


def main():
    a = hidet.ones([3, 4]).cuda()
    b = hidet.ones([4, 5]).cuda()
    c = naive_matmul(a, b)
    hidet.cuda.synchronize()
    print(a)
    print(b)
    print(c)


if __name__ == '__main__':
    main()