import torch
import torch.nn as nn
import numpy as np
import mmcv
from mmcv.ops import SparseConv3d, SparseConvTensor
import time
import sys

class SparseConvModel(nn.Module):
    def __init__(self):
        super(SparseConvModel, self).__init__()
        self.conv1 = SparseConv3d(
            in_channels=4,
            out_channels=64,
            kernel_size=3,
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x.features)
        return x

model = SparseConvModel().eval().cuda().half()
print(model)

with torch.no_grad():
    features = np.random.rand(100, 4).astype(np.float32)
    coordinates = np.random.randint(0, 128, size=(100, 4))
    coordinates[:, 0] = 0

    features = torch.tensor(features).cuda().half()
    coordinates = torch.tensor(coordinates, dtype=torch.int32).cuda()

    torch.cuda.synchronize()
    spatial_shape = (128, 128, 128)

    sparse_tensor = SparseConvTensor(
        features=features,
        indices=coordinates,
        spatial_shape=spatial_shape,
        batch_size=1
    )
    while(1):
        start_time = time.perf_counter()
        output = model(sparse_tensor)
        end_time = time.perf_counter()
        print("consum: {:.2f}ms".format((end_time - start_time) * 1000))
        print(output[0,0:10])
        time.sleep(1)

print(output.shape)
print(output.dtype)