# -------------------------------------------------------------
# Bulk configuration
# -------------------------------------------------------------

# Set up lattice
lattice = Hexagonal(3.1604*Angstrom, 25.0*Angstrom)

# Define elements
elements = [Molybdenum, Sulfur, Sulfur]

# Define coordinates
fractional_coordinates = [[ 0.333333333333,  0.666666666667,  0.5           ],
                          [ 0.666666666667,  0.333333333333,  0.4365578     ],
                          [ 0.666666666667,  0.333333333333,  0.5634422     ]]

# Set up configuration
bulk_configuration = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates
    )

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------

potentialSet = StillingerWeber_MoS_2013()
calculator = TremoloXCalculator(parameters=potentialSet)

bulk_configuration.setCalculator(calculator)
nlprint(bulk_configuration)
bulk_configuration.update()
nlsave('mos2_relax.nc', bulk_configuration)

constraints = [FixStrain(x=False, y=False, z=True)]
bulk_configuration = OptimizeGeometry(
        bulk_configuration,
        max_forces=0.001*eV/Ang,
        max_stress=0.001*eV/Ang**3,
        max_steps=200,
        max_step_length=0.2*Ang,
        constraints=constraints,
        trajectory_filename=None,
        optimizer_method=LBFGS(),
        )
nlsave('mos2_relax.nc', bulk_configuration)
nlprint(bulk_configuration)
