import numpy as np
import matplotlib.pyplot as plt
import timeit

setup = "import numpy as np; idxlen = {}; rng = np.random.default_rng(); matrix = -np.ones([100,idxlen],dtype=int)"

test1 = "rand_idxs = np.arange(idxlen); rng.shuffle(rand_idxs); matrix[0][:] = rand_idxs"
test2 = "matrix[0][:] = rng.choice(idxlen, idxlen, replace=False)"
test3 = "matrix[0][:] = rng.permutation(idxlen)"
idx_lengths = np.logspace(1, 5, num=10, dtype=int)
results1 = []
results2 = []
results3 = []

for idxlen in idx_lengths:
    time1 = timeit.timeit(test1.format(idxlen), setup.format(idxlen), number=10000)
    time2 = timeit.timeit(test2.format(idxlen), setup.format(idxlen), number=10000)
    time3 = timeit.timeit(test3.format(idxlen), setup.format(idxlen), number=10000)
    results1.append(time1)
    results2.append(time2)
    results3.append(time3)

plt.figure(figsize=(10, 6))
plt.plot(idx_lengths, results1, label='rng.shuffle')
plt.plot(idx_lengths, results2, label='rng.choice')
plt.plot(idx_lengths, results3, label='rng.permutation')
plt.xscale('log')
plt.yscale('linear')

plt.title('Shuffle Timings for Different Methods')
plt.xlabel('Index Length (idxlen)')
plt.ylabel('Time (seconds)')
plt.legend()
plt.show()