import torch
import matterhorn_pytorch.snn as snn
import matterhorn_pytorch.snn.functional as SF
import matterhorn_pytorch.data.functional as DF


if __name__ == "__main__":
    tensor_shape = (2, 3, 4, 5)
    indices = torch.tensor([[0, 1, 2, 3], [1, 2, 1, 0], [0, 1, 2, 3], [0, 0, 0, 0]], dtype = torch.long)
    print(indices)
    updated_tensor = DF.event_seq_to_spike_train(indices, tensor_shape, count = True)
    print(updated_tensor)
    event_seq = DF.spike_train_to_event_seq(updated_tensor)
    event_seq = DF.spike_train_to_event_seq(updated_tensor).long()
    print(event_seq)
    s = updated_tensor[tuple(event_seq[:, s] for s in range(event_seq.shape[1]))]
    print(s)