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

# Set up lattice
vector_a = [5.4306, 0.0, 0.0]*Angstrom
vector_b = [0.0, 5.4306, 0.0]*Angstrom
vector_c = [0.0, 0.0, 5.4306]*Angstrom
lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
elements = [Silicon, Silicon, Silicon, Silicon, Silicon, Silicon, Silicon,
            Silicon]

# Define coordinates
fractional_coordinates = [[0.0 ,  0.0 ,  0.0 ],
                          [0.25,  0.25,  0.25],
                          [0.5 ,  0.5 ,  0.0 ],
                          [0.75,  0.75,  0.25],
                          [0.5 ,  0.0 ,  0.5 ],
                          [0.75,  0.25,  0.75],
                          [0.0 ,  0.5 ,  0.5 ],
                          [0.25,  0.75,  0.75]]

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

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

potentialSet = StillingerWeber_Si_1985()
calculator = TremoloXCalculator(parameters=potentialSet)

bulk_configuration.setCalculator(calculator)
bulk_configuration.update()

# -------------------------------------------------------------
# Molecular Dynamics
# -------------------------------------------------------------

initial_velocity = MaxwellBoltzmannDistribution(
    temperature=300.0*Kelvin,
    remove_center_of_mass_momentum=True
)

# Define a hook function that prints the instantaneous cell volume and the pressure at each MD step
def print_cell_and_stress(step, time, configuration, forces, stress):

    # Get the observables from the configuration
    volume = configuration.bravaisLattice().unitCellVolume().inUnitsOf(Ang**3)

    # Calculate the hydrostatic pressure from the stress tensor
    pressure = -(stress[0, 0] + stress[1, 1] + stress [2, 2]) / 3.0

    # Print the output to the log file
    print("| MD step %i:  Pressure : %f bar    Volume:  %f Angstrom**3" % (step, pressure, volume))

# Run the MD simulation
md_trajectory = MolecularDynamics(
    bulk_configuration,
    constraints=[],
    trajectory_filename='trajectory.nc',
    steps=200,
    post_step_hook=print_cell_and_stress,
    log_interval=50,
    method=NPTMartynaTobiasKlein()
)