import numpy as np

from cleanlab.internal.util import value_counts

np.random.seed(0)
%load_ext memory_profiler

# Simulate a batch_size of 10_000 and a total of 250_000 classes
x = np.random.randint(250_000, size=10_000)

# Simulate a large dataset with most classes present.
y = np.random.randint(249_950, size=500_000_000)

%timeit value_counts(x, num_classes=250_000)
%memit value_counts(x, num_classes=250_000)
# 48.1 ms ± 1.24 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
# peak memory: 4103.93 MiB, increment: 4.66 MiB

%%timeit 
%memit value_counts(y, num_classes=250_000)
# peak memory: 8465.91 MiB, increment: 4385.38 MiB
# peak memory: 8472.33 MiB, increment: 4391.65 MiB
# peak memory: 8372.47 MiB, increment: 4291.54 MiB
# peak memory: 8570.36 MiB, increment: 4489.45 MiB
# peak memory: 8372.46 MiB, increment: 4291.54 MiB
# peak memory: 8582.37 MiB, increment: 4501.44 MiB
# peak memory: 8692.37 MiB, increment: 4611.44 MiB
# peak memory: 8760.37 MiB, increment: 4679.44 MiB
# 1min 5s ± 971 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit value_counts(x, num_classes=250_000)
%memit value_counts(x, num_classes=250_000)
# 16.9 ms ± 402 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# peak memory: 4093.53 MiB, increment: 0.00 MiB

%%timeit 
%memit value_counts(y, num_classes=250_000)
# peak memory: 8777.02 MiB, increment: 4712.53 MiB
# peak memory: 8587.01 MiB, increment: 4524.52 MiB
# peak memory: 8433.01 MiB, increment: 4370.52 MiB
# peak memory: 8355.69 MiB, increment: 4293.21 MiB
# peak memory: 8404.86 MiB, increment: 4342.38 MiB
# peak memory: 8789.01 MiB, increment: 4726.52 MiB
# peak memory: 8769.01 MiB, increment: 4706.52 MiB
# peak memory: 8355.70 MiB, increment: 4293.21 MiB
# 28.7 s ± 322 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)