# Set up lattice
lattice = Hexagonal(4.966693338812276*Angstrom, 5.4647751809821905*Angstrom)

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

# Define coordinates
fractional_coordinates = [[ 0.469217958328,  0.            ,  0.            ],
                          [ 0.            ,  0.469217958328,  0.666666666667],
                          [ 0.530782041672,  0.530782041672,  0.333333333333],
                          [ 0.409091887977,  0.269718793416,  0.115058361235],
                          [ 0.269718793416,  0.409091887977,  0.551608638765],
                          [ 0.730281206584,  0.139373094561,  0.781725361235],
                          [ 0.590908112023,  0.860626905439,  0.218274638765],
                          [ 0.860626905439,  0.590908112023,  0.448391361235],
                          [ 0.139373094561,  0.730281206584,  0.884941638765]]

# Set up reference configuration of quartz
reference_configuration = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates,
)

# Define the molecular dynamics method
initial_velocity = MaxwellBoltzmannDistribution(
    temperature=1500.0*Kelvin,
    remove_center_of_mass_momentum=True,
    random_seed=None,
    enforce_temperature=True,
)

method = Langevin(
    time_step=1*femtoSecond,
    reservoir_temperature=1500*Kelvin,
    friction=0.01*femtoSecond**-1,
    initial_velocity=initial_velocity,
    heating_rate=0*Kelvin/picoSecond,
)

# Define the molecular dynamics calculator
potentialSet = Tersoff_SiO_2007()
calculator = TremoloXCalculator(parameters=potentialSet)

# Define the MD snapshots parameters:
# Run 50000 MD steps and take 50 equally spaced samples for the final training set.
training_sets = MolecularDynamicsSnapshotsParameters(
    reference_configuration,
    calculator=calculator,
    sample_size=50,
    steps=50000,
    method=method,
    log_filename_prefix=None,
)

# Set up MTP training and run the training data generation and calculation.
moment_tensor_potential_training = MomentTensorPotentialTraining(
    filename='mtp_study.hdf5',
    object_id='training',
    training_sets=training_sets,
    calculator=LCAOCalculator(),
    calculate_stress=True,
)
moment_tensor_potential_training.update()