# RUN: %PYTHON %s | FileCheck %s

import torch
import numpy as np

from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run

from mpact.models.kernels import SqSum

net = SqSum()

# Construct adjacency matrix.
V = 8
edges = np.array([[0, 1], [0, 4], [1, 4], [3, 4], [4, 3]], dtype=np.int32)
E = edges.shape[0]
adj_mat = torch.sparse_coo_tensor(edges.T, torch.ones(E), (V, V), dtype=torch.int64)

#
# CHECK: pytorch
# CHECK:   tensor(5)
# CHECK: mpact
# CHECK:   5

# Run it with PyTorch.
print("pytorch")
res = net(adj_mat)
print(res)

# Run it with MPACT.
print("mpact")
# TODO: make this work, expose `sparse-emit-strategy=sparse-iterator` to
# mini-pipeline.
# res = mpact_jit(net, adj_mat, use_sp_it=True)
res = mpact_jit(net, adj_mat)
print(res)