#+-------------------------------------------------------------+
#| Script to set-up a vacancy defect state in a bulk           |
#| crystal,relax it and calculate the total energy.            |
#+-------------------------------------------------------------+

#User-specified parameters:

#+-------------------------------------------------------------+
#Input-output files parameters                                 |
#+-------------------------------------------------------------+

#Name of file containing the initial configuration. 
#The file should include the bulk configuration of 
#the supercell and the calculator to be used
filename = 'big-bulk-GaAs.hdf5'

#Name of the file the results will be saved to
saved_filename = 'vac-0-GaAs.hdf5'

#+-------------------------------------------------------------+
#Defect state parameters                                       |
#+-------------------------------------------------------------+

#Index number of atom to be replaced with the defect (integer)
defect_index = 10

#Charge of the defect (integer or float, with sign)
defect_charge = 0

#+-------------------------------------------------------------+
#Geometry optimization parameters                              |
#+-------------------------------------------------------------+

#Should a structural relaxation be performed
relax_structure = True

#Convergence threshold for the atomic forces (in eV/Angstrom)
max_forces = 0.05

#Maximum number of relaxation steps
max_steps = 200

#+-------------------------------------------------------------+
#                                                              | 
#     Nothing to be changed below this line!!                  |
#                                                              |
#+-------------------------------------------------------------+

# Using the last configuration ensures that it is the optimized one if it exists
bulk_configuration = nlread(filename, BulkConfiguration)[-1]

# Checking if the the lattice constant is sensible by measuring the stress and comparing to
# max_stress with a default value of 0.05 (default value for OptimizeGeometry)
stress = Stress(bulk_configuration)
nlprint(stress)

max_stress = 0.05

if numpy.max(numpy.abs(stress.evaluate().inUnitsOf(eV/Ang**3))) > max_stress:
    print('Warning! - The bulk configuration seems to have a wrong lattice parameter.')
    print('The maximum component of the stress tensor is larger than ', max_stress, ' eV/Ang**3.')

# Extracting various objects for possible modification and use when reconstructing the BulkConfiguration
calculator = bulk_configuration.calculator()
elements = bulk_configuration.elements()
fractional_coordinates = bulk_configuration.fractionalCoordinates()
bravais_lattice = bulk_configuration.bravaisLattice()

# Replace atom at vacancy by a ghost atom
ghost_atoms = [defect_index]

# Reconstruct BulkConfiguration
bulk_configuration = BulkConfiguration(
bravais_lattice = bravais_lattice,
elements = elements,
fractional_coordinates = fractional_coordinates,
ghost_atoms = ghost_atoms,
)

# Update the calculator with the charge of the vacancy
bulk_configuration.setCalculator(
      calculator(charge = defect_charge))

bulk_configuration.update()

# Save the new BulkConfiguration with the user-specified filename
nlsave(saved_filename, bulk_configuration)

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------

# Optimize geometry if selected by the user
if relax_structure is True:

    traj_filename=saved_filename[:-3]+'-traj.hdf5'

    constraints=[defect_index]

    bulk_configuration = OptimizeGeometry(
        bulk_configuration,
        max_forces=max_forces*eV/Ang,
        max_steps=max_steps,
        max_step_length=0.2*Ang,
        constraints=constraints,
        trajectory_filename=traj_filename,
        disable_stress=True,
        optimizer_method=LBFGS(),
        )

    nlsave(saved_filename, bulk_configuration)
    nlprint(bulk_configuration)

# -------------------------------------------------------------
# Total Energy
# -------------------------------------------------------------
total_energy = TotalEnergy(bulk_configuration)
nlsave(saved_filename, total_energy)
nlprint(total_energy)