from ASE import ListOfAtoms, Atom
from ASE.Dynamics.MDMin import MDMin
from ASE.Dynamics.ConjugateGradient import ConjugateGradient
from ASE.Trajectories import WriteNetCDFTrajectory
from ASE.Filters.Subset import Subset
from ASE.Filters.NudgedElasticBand import NudgedElasticBand
from math import sqrt

import Asap


a = 4.0614
b = a / sqrt(2)
h = b / 2
initial = ListOfAtoms([Atom('Al', (0, 0, 0)),
                       Atom('Al', (a / 2, b / 2, -h))],
                      periodic=(1, 1, 0),
                      cell=(a, b, -2 * h))
initial = initial.Repeat((4, 5, 2))
initial.append(Atom('Al', (a / 2, b / 2, h)))
initial = Asap.ListOfAtoms(initial)

final = initial.Copy()
final[-1].SetCartesianPosition((a / 2, 3 * b / 2, h))

# Construct a list of images:
images = [initial]
for i in range(??):
    images.append(initial.Copy())
images.append(final)

# Let all images use an Asap-EMT calculator:
for image in images:
    image.SetCalculator(Asap.EMT())

# Make a mask of zeros and ones that select the dynamic atoms (the
# three topmost layers):
mask = [atom.GetCartesianPosition()[2] > -2.5 * h for atom in initial]

# The subsets contain only the three topmost layers - the bottom layer
# is filtered out (fixed):
image_subsets = [Subset(image, mask=mask) for image in images]

# Relax the initial state:
cg = ConjugateGradient(image_subsets[0], fmax=0.02)
cg.Converge()
    
# Relax the final state:
cg = ConjugateGradient(image_subsets[-1], fmax=0.02)
cg.Converge()

# Create a Nudged Elastic Band:
neb = NudgedElasticBand(image_subsets)

# Mak a starting guess for the minimum energy path (a straight line
# from the initial to the final state):
neb.SetInterpolatedPositions()

# Use MDMin to relax the path:
minimizer = MDMin(neb, dt=0.08)

# Relax the path and print out the energy barrier:
for i in range(5):
    minimizer.Run(10)
    energies = neb.GetListOfEnergies()
    Ea = max(energies) - energies[0]
    print i, Ea

# Write the path to a trajectory:
WriteNetCDFTrajectory(images, 'jump1.traj')