import numpy as np

from ase.build import bulk
from ase.calculators.lj import LennardJones
from ase.spacegroup.symmetrize import FixSymmetry, check_symmetry
from ase.optimize import BFGS
from ase.constraints import UnitCellFilter

# We setup a bcc Al cell - bcc is unstable with LJ potential
# so without constraint this would relax back to an fcc structure
atoms_prim = bulk('Al', 'bcc', a=2 / np.sqrt(3), cubic=True)

# Now we setup a 2x2x2 supercell, and break the symmetry slightly
atoms_init = atoms_prim * [2, 2, 2]
atoms_init.positions[0, 0] += 1.0e-7  # break symmetry by 1e-7

# We use an LJ calculator, and allow the cell and atomic positions to relax
atoms_unsym = atoms_init.copy()
atoms_unsym.calc = LennardJones()
ucf_unsym = UnitCellFilter(atoms_unsym)

dyn = BFGS(ucf_unsym)
print("Initial Energy", atoms_unsym.get_potential_energy())
dyn.run(fmax=0.001)
print("Final Energy", atoms_unsym.get_potential_energy())

# Now we repeat the optimisation with the symmetrization constraint in place
atoms_sym = atoms_init.copy()
atoms_sym.calc = LennardJones()
atoms_sym.set_constraint(FixSymmetry(atoms_sym))
ucf_sym = UnitCellFilter(atoms_sym)

dyn = BFGS(ucf_sym)
print("Initial Energy", atoms_sym.get_potential_energy())
dyn.run(fmax=0.001)
print("Final Energy", atoms_sym.get_potential_energy())

print("position difference", np.linalg.norm(atoms_unsym.get_positions() -
                                            atoms_sym.get_positions()))

# We print out the initial symmetry groups at two different precision levels
print("initial symmetry at precision 1e-6")
check_symmetry(atoms_init, 1.0e-6, verbose=True)
print("initial symmetry at precision 1e-8")
check_symmetry(atoms_init, 1.0e-8, verbose=True)

# Printing the final symmetries shows that
# the "unsym" case relaxes to a lower energy fcc structure
# with a change in spacegroup, while the "sym" case stays as bcc
print("unsym symmetry after relaxation")
d_unsym = check_symmetry(atoms_unsym, verbose=True)
print("sym symmetry after relaxation")
d_sym = check_symmetry(atoms_sym, verbose=True)
