from ase.build import bulk, molecule
from ase.constraints import FixAtoms, FixBondLength
from ase.io import read
from ase.io.trajectory import Trajectory, get_header_data


def test_trajectory_heterogeneous():
    a0 = molecule('H2O')
    a1 = a0.copy()
    a1.rattle(stdev=0.5)
    a2 = a0.copy()
    a2.set_masses()
    a2.center(vacuum=2.0)
    a2.rattle(stdev=0.2)
    a3 = molecule('CH3CH2OH')
    a4 = bulk('Au').repeat((2, 2, 2))
    a5 = bulk('Cu').repeat((2, 2, 3))

    # Add constraints to some of the images:
    images = [a0, a1, a2, a3, a4, a5]
    for i, img in enumerate(images[3:]):
        img.set_constraint(FixAtoms(indices=range(i + 3)))
        if i == 2:
            img.constraints.append(FixBondLength(5, 6))

    with Trajectory('out.traj', 'w') as traj:
        for i, img in enumerate(images):
            traj.write(img)
            print(i, traj.multiple_headers)
            assert traj.multiple_headers == (i >= 2)

    with Trajectory('out.traj') as rtraj:
        newimages = list(rtraj)

    assert len(images) == len(newimages)
    for i in range(len(images)):
        assert images[i] == newimages[i], i
        h1 = get_header_data(images[i])
        h2 = get_header_data(newimages[i])
        print(i, images[i])
        print(h1)
        print(h2)
        print()
        # assert headers_equal(h1, h2)

    # Test append mode:
    with Trajectory('out.traj', 'a') as atraj:
        atraj.write(molecule('H2'))
        atraj.write(molecule('H2'))
    read('out.traj', index=':')
