from asap3 import *
from ase.lattice.cubic import FaceCenteredCubic
from asap3.testtools import ReportTest
from asap3.md.velocitydistribution import *
from asap3.analysis import CoordinationNumbers, PTM
from ase.parallel import world
from ase.visualize import view

debug = 1
if debug == 1:
    DebugOutput("safe_nblist%d.log", nomaster=True)
elif debug == 2:
    time.sleep(world.rank)
    print("PID:", os.getpid())
    time.sleep(20)

# Important: always at least two cores along x direction
ismaster = world.rank == 0
isparallel = world.size != 1
if world.size == 1:
    cpulayout = None
elif world.size == 2:
    cpulayout = [2,1,1]
elif world.size == 3:
    cpulayout = [3,1,1]
elif world.size == 4:
    cpulayout = [2,1,2]


def make_atoms():
    """Make an atomic system.

    Some vacuum is added, so that shearing it changes the number of
    atoms on the different MPI tasks.
    """
    if ismaster:
        atoms = FaceCenteredCubic('Cu', size=(20,20,20))
        atoms.pbc = False
        atoms.center(vacuum=5.0, axis=0)
        atoms.pbc = True
    else:
        atoms = None
    if world.size > 1:
        atoms = MakeParallelAtoms(atoms, cpulayout)
    atoms.calc = EMT()
    atoms.get_forces()
    return atoms

def deform_atoms(atoms):
    """Shear the atoms."""
    r = atoms.get_positions()
    df =  0.1 * r[:,1]
    r[:,0] += df
    atoms.set_positions(r)
    
atoms = make_atoms()
e0 = atoms.get_potential_energy()
deform_atoms(atoms)
e1 = atoms.get_potential_energy()
if world.rank == 0:
    print("\n\nEnergy of atoms changed from {} to {}".format(e0, e1), file=sys.stderr)

atoms = make_atoms()
print("A: {} atoms on cpu {}".format(len(atoms), world.rank), file=sys.stderr)
deform_atoms(atoms)
#atoms.get_potential_energy()    # Workaround
CoordinationNumbers(atoms)
atoms.get_forces()
e2 = atoms.get_potential_energy()
print("B: {} atoms on cpu {}".format(len(atoms), world.rank), file=sys.stderr)
ReportTest("Energy after CoordinationNumbers()", e2, e1, 0.001)


atoms = make_atoms()
deform_atoms(atoms)
#atoms.get_potential_energy()    # Workaround
nbl = atoms.calc.get_neighborlist()
nbl.check_and_update(atoms)
atoms.get_forces()
e2 = atoms.get_potential_energy()
ReportTest("Energy after borrowed neighborlist update", e2, e1, 0.001)

atoms = make_atoms()
deform_atoms(atoms)
#atoms.get_potential_energy()    # Workaround
PTM(atoms, cutoff=4.7)
atoms.get_forces()
e2 = atoms.get_potential_energy()
ReportTest("Energy after PTM()", e2, e1, 0.001)


# The following are not expected to be affected by issue #43.

for i, nbtype in enumerate((FullNeighborList, NeighborCellLocator)):
    atoms = make_atoms()
    nbl = nbtype(4.8, atoms)
    deform_atoms(atoms)
    #atoms.get_potential_energy()    # Workaround
    nbl.check_and_update(atoms)
    atoms.get_forces()
    e2 = atoms.get_potential_energy()
    ReportTest("Energy neighbor list type {}".format(i), e2, e1, 0.001)

ReportTest.Summary()
