import espressomd
import espressomd.io.writer

import os
import h5py
import numpy as np
import tempfile

# set particles outside the main box to get folded coordinates
system = espressomd.System(box_l=[6, 7, 8])
for i in range(12):
    p = system.part.add(pos=np.array(3 * [i - 4], dtype=float))

# write H5MD data
temp_directory = tempfile.TemporaryDirectory()
temp_file = os.path.join(temp_directory.name, 'test.h5')
units = espressomd.io.writer.h5md.UnitSystem(time='ps', mass='u', length='m', charge='e')
h5 = espressomd.io.writer.h5md.H5md(file_path=temp_file, unit_system=units)
h5.write()
h5.write()
h5.flush()
h5.close()

# read H5MD data
h5_params = h5.get_params()
py_file = h5py.File(temp_file, 'r')
py_box = py_file['particles/atoms/box/edges/value'][1]
py_pos = py_file['particles/atoms/position/value'][1]
py_img = py_file['particles/atoms/image/value'][1]
py_id = py_file['particles/atoms/id/value'][1]

# apply unfolding
pos_folded = py_pos
pos_unfolded = py_img * py_box + py_pos

# report
print('unfolded coordinates:')
print('  from ESPResSo')
print(system.part.all().pos)
print('  from H5MD')
print(pos_unfolded[np.argsort(py_id)])
print('folded coordinates:')
print('  from ESPResSo')
print(system.part.all().pos_folded)
print('  from H5MD')
print(pos_folded[np.argsort(py_id)])

# verify data integrity
tolerance = 1e-16
np.testing.assert_allclose(
    pos_folded[np.argsort(py_id)],
    system.part.all().pos_folded,
    atol=tolerance)
np.testing.assert_allclose(
    pos_unfolded[np.argsort(py_id)],
    system.part.all().pos,
    atol=tolerance)