from math import cos, sin, pi

from ASE import ListOfAtoms, Atom
from ASE.Visualization.RasMol import RasMol
from ASE.Dynamics.ConjugateGradient import ConjugateGradient
from ASE.Trajectories import NetCDFTrajectory
from Dacapo import Dacapo


d = 2.0
r = 0.9575
t = pi / 180 * 104.51
dimer = ListOfAtoms([Atom('O', (0, 0, 0)),
                     Atom('H', (-r * cos(t / 2), r * sin(t / 2), 0)),
                     Atom('H', (-r * cos(t / 2), -r * sin(t / 2), 0)),
                     Atom('O', (d + r, 0, 0)),
                     Atom('H', (d, 0, 0)),
                     Atom('H', (d + r - r * cos(t), 0, r * sin(t)))],
                    cell=(9, 6, 7))
dimer.SetCartesianPositions(dimer.GetCartesianPositions() + (4, 3, 3.5))
plot = RasMol(dimer, (2, 2, 2))
calc = Dacapo(nbands=12,
              planewavecutoff=340,
              densitycutoff=500,
              out='2H2O.nc')
calc.StayAliveOff()
dimer.SetCalculator(calc)
traj = NetCDFTrajectory('2H2O.traj', dimer)
cg = ConjugateGradient(dimer, fmax=0.05)

cg.Attach(traj)
cg.Attach(plot)
cg.Converge()