import espressomd.checkpointing
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--perturbation', action='store', type=str, required=False,
                    default='continue', choices=['continue', 'load', 'recalc', 'resort'])
args = parser.parse_args()

def verbose_integration(nsteps, recalc_forces=False):
    p = system.part[0]
    for i in range(nsteps):
        system.integrator.run(1, recalc_forces=recalc_forces)
        print(f'rng_counter = {system.thermostat.get_state()[0]["counter"]}, {p.pos=} {p.f=}')

checkpoint = espressomd.checkpointing.Checkpoint(checkpoint_id='mycheckpoint')
if args.perturbation != 'load':
    system = espressomd.System(box_l=[1, 1, 1], time_step=0.01)
    system.cell_system.skin = 0.4
    system.thermostat.set_langevin(kT=1.0, gamma=1.0, seed=42)
    system.part.add(pos=[0, 0, 0])
    verbose_integration(2)
    checkpoint.register('system')
    checkpoint.save()
else:
    checkpoint.load()
    print('rng_counter = 1, unknown')
    print('rng_counter = 2, unknown')

if args.perturbation == 'recalc':
    verbose_integration(2, recalc_forces=True)
if args.perturbation == 'continue':
    verbose_integration(2)
if args.perturbation == 'load':
    verbose_integration(2)
if args.perturbation == 'resort':
    system.part.add(pos=[0, 0, 0], type=1).remove() # trigger cell system resort
    verbose_integration(2)