def xtabs(data):
    xid, yid, vals, cats = data
    res = [0] * 5
    for i in range(n):
        for j in range(n):
            for k in range(len(res)):
                add = (xid[i] == yid[j]) * (cats[j] == k) * vals[i]
                res[k] += add
    return res

import sys

import secretflow as sf
import spu

# Check the version of your SecretFlow
print('The version of SecretFlow: {}'.format(sf.__version__))

# In case you have a running secretflow runtime already.
sf.shutdown()

sf.init(['alice', 'bob', 'carol', 'dave'], address='local')


aby3_config = sf.utils.testing.cluster_def(parties=['alice', 'bob', 'carol'])


spu_device = sf.SPU(aby3_config)


alice, dave = sf.PYU('alice'), sf.PYU('dave')


spu_io = sf.device.SPUIO(spu_device.conf, spu_device.world_size)


n = int(sys.argv[1])
data = [list(range(n))] * 3 + [[i % 5 for i in range(n)]]

pyu = alice(lambda: data)()

spu_data = pyu.to(spu_device)

def xtabs(data):
    xid, yid, vals, cats = data
    res = [0] * 5
    for i in range(n):
        for j in range(n):
            for k in range(len(res)):
                add = (xid[i] == yid[j]) * (cats[j] == k) * vals[i]
                res[k] += add
    return res

res = spu_device(xtabs)(spu_data)

print(sf.reveal(res))
print('expected', xtabs(data))
print('size', n)