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

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

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

# Define coordinates
fractional_coordinates = [[ 0.333333333333,  0.666666666667,  0.25          ],
                          [ 0.666666666667,  0.333333333333,  0.75          ],
                          [ 0.333333333333,  0.666666666667,  0.621         ],
                          [ 0.666666666667,  0.333333333333,  0.121         ],
                          [ 0.333333333333,  0.666666666667,  0.879         ],
                          [ 0.666666666667,  0.333333333333,  0.379         ]]

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

# Add tags
bulk_configuration.addTags('layer1', [0, 3, 5])
bulk_configuration.addTags('layer2', [1, 2, 4])

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

sw_layer1 = StillingerWeber_MoS_2013(tags='layer1')
sw_layer2 = StillingerWeber_MoS_2013(tags='layer2')

# Define a new potential for the interlayer interaction.
lj_interlayer_potential = TremoloXPotentialSet(name="InterLayerPotential")

# Add particle type definitions for both types.
lj_interlayer_potential.addParticleType(ParticleType.fromElement(Molybdenum))
lj_interlayer_potential.addParticleType(ParticleType.fromElement(Sulfur, sigma=3.13*Angstrom, epsilon=0.00693*eV))

# Add Lennard-Jones potentials between the sulfur atoms of different layers.
lj_interlayer_potential.addPotential(LennardJonesPotential('S', 'S', r_cut=10.0 * Angstrom))
lj_interlayer_potential.setTags(['layer1', 'layer2'])

# Combine all 3 potential sets in a single calculator.
calculator = TremoloXCalculator(parameters=[sw_layer1, sw_layer2, lj_interlayer_potential])

bulk_configuration.setCalculator(calculator)
bulk_configuration.update()

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,
        trajectory_filename=None,
        optimizer_method=LBFGS(),
        )
nlprint(bulk_configuration)