"""Using filters, trajectories and python dynamics 
"""

from Dacapo import Dacapo
from ASE import Atom, ListOfAtoms
from ASE.Filters.Subset import Subset
from ASE.Dynamics.ConjugateGradient import ConjugateGradient
from ASE.Dynamics.MDMin import MDMin
from ASE.Dynamics.LineMinimizers.LM1 import LM1
from ASE.Trajectories import NetCDFTrajectory
import os

os.system('rm -f CO_in_a_box.nc CO_in_a_box.txt dynamics.nc dynamics.txt')

atoms = ListOfAtoms([Atom('C', (2, 2, 2),tag=1),
                     Atom('O', (3.1, 2, 2))],
                     cell=(4, 3.5, 3.5), periodic=1)

calc = Dacapo(planewavecutoff=300,   # in eV
              nbands=6,            # 1 extra empty bands
              out='CO_in_a_box.nc',txtout='CO_in_a_box.txt')

atoms.SetCalculator(calc)

energy = atoms.GetPotentialEnergy()
forces = atoms.GetCartesianForces()

print 'Initial energy:',energy
print 'Initial abs force:',max(max(forces))

calc.SetNetCDFFile('dynamics.nc')
calc.SetTxtFile('dynamics.txt')

# for efficiency keep dacapo running between steps
calc.StayAliveOn()

# Setup the filter constraining the fixed atoms
subset = Subset(atoms, mask=[atom.GetTag() != 1 for atom in atoms])

# Setup the minimizer
# use MDMin and not ConjugateGradient because of a 
# problem combining StayAliveOn and this minimizer, 
# dyn = ConjugateGradient(subset, fmax=0.05,lineMin=LM1(0.05))
# dyn.SetVerbosity(1)
dyn = MDMin(subset, fmax=0.2,dt=0.08)


# Create a trajectory for the minimization path
path = NetCDFTrajectory('CO-trajectory.nc', atoms)
# Put the initial state in the trajectory:
path.Update()

# attach the trajectory to the dynamics
dyn.Attach(path)

# Find the ground state
dyn.Converge()

energy = atoms.GetPotentialEnergy()
forces = atoms.GetCartesianForces()
print 'Final energy:', energy
print 'Final abs forces:',max(max(forces))