# XXX Verify that this works on alumne computers!!
# Maybe this is more advanced functionality than we want at the moment?
import os
from ase import Atoms
from ase.optimize import BFGS
from ase.calculators.aims import Aims
from ase.calculators.socketio import SocketIOCalculator
from ase.utils import workdir


os.environ['ASE_AIMS_COMMAND'] = 'aims.x'  # XXX other command
os.environ['AIMS_SPECIES_DIR'] = '/home/alumne/software/FHIaims/species_defaults/light'

atoms = Atoms('HOH',
              positions=[[0, 0, -1], [0, 1, 0], [0, 0, 1]])
opt = BFGS(atoms, trajectory='opt-aims-socketio.traj')

aims = Aims(compute_forces=True,
            use_pimd_wrapper=('localhost', 31415))

with SocketIOCalculator(aims, port=31415) as calc:
    atoms.calc = calc:
    with workdir('aims-output-socketio', mkdir=True):
        opt.run(fmax=0.05)