from xpstructures.raggedarray import RaggedArray
from xpstructures.hashtable import Counter 
from npstructures.raggedarray import RaggedArray as Orig_RaggedArray
from npstructures.hashtable import Counter as Orig_Counter

import numpy as np
import cupy as cp

import time

mempool = cp.get_default_memory_pool()

num_kmers = 75000000

kmer_matrix = np.load("testkmers.npy")[:num_kmers]
print(f"kmer_matrix.shape={kmer_matrix.shape}")

unique_kmer_matrix = np.unique(kmer_matrix)

counter = Orig_Counter(unique_kmer_matrix, mod=100_000_033)
t = time.time_ns()
counter.count(kmer_matrix)
t = time.time_ns() - t
print("cpu count time:", t / 1e6, "ms\n")
del counter

counter = Counter(unique_kmer_matrix, mod=100_000_033)
counter.to_cuda()

kmer_matrix = cp.asarray(kmer_matrix)
print(f"kmer_matrix moved to CUDA. Allocated CUDA bytes: {mempool.used_bytes()}")

t = time.time_ns()
counter.count(cp.asarray(kmer_matrix))
t = time.time_ns() - t
print("cuda count time:", t / 1e6, "ms")

print(f"Allocated CUDA bytes: {mempool.used_bytes()}")
del counter
print(f"CUDA counter deleted. Allocated CUDA bytes: {mempool.used_bytes()}")
del kmer_matrix
print(f"CUDA kmer_matrix deleted. Allocated CUDA bytes: {mempool.used_bytes()}")
