import timeit

import numpy as np

from ase.io import read, write  # noqa
from ase.build import bulk

from dynasor.trajectory.extxyz_trajectory_reader import ExtxyzTrajectoryReader  # noqa


fname = 'atoms.extxyz'
frames = 100
atoms = 2000


atoms = bulk('Al').repeat((atoms, 1, 1))

momenta = np.random.random(size=atoms.positions.shape)
atoms.set_momenta(momenta)
atoms.info['Time'] = 1.0
atoms.set_tags(np.ones(len(atoms)))

print(f'Number of atoms:  {len(atoms)}')
print(f'Number of frames: {frames}')
print(f'Number of bytes:  {frames * len(atoms) * 1.8 / 2**12 / 2**12:.3f} GB')

write(fname, [atoms]*frames)
timer = timeit.Timer("read(fname, index=':')", globals=globals())  # noqa
N, T = timer.autorange()
print(f'Time ASE: {T/N:<10.5f}', flush=True)
T0 = T/N
for n in [2, 3, 4, 5, 6, 8, 12]:
    timer = timeit.Timer('list(ExtxyzTrajectoryReader(fname, max_workers=n))', globals=globals())
    N, T = timer.autorange()
    print(f'Workers: {n}  Time: {T/N:<10.5f}  Speedup: {T0/(T/N):<10.5f}  '
          f'Efficiency: {T0/(T/N)/n:<10.5f}', flush=True)