setVerbosity(MinimalLog)

# -------------------------------------------------------------
# Calculator: DFT Reference calculator.
# -------------------------------------------------------------
k_point_sampling = KpointDensity(
    density_a=7.0*Angstrom,
    )
numerical_accuracy_parameters = NumericalAccuracyParameters(
    density_mesh_cutoff=125.0*Hartree,
    k_point_sampling=k_point_sampling,
    )

iteration_control_parameters = IterationControlParameters(
    tolerance=5e-05,
    damping_factor=0.3,
    number_of_history_steps=12,
    )

reference_calculator = LCAOCalculator(
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    iteration_control_parameters=iteration_control_parameters,
    )

# Load the training data. This can be one or several TrainingSet, MomentTensorPotentialTraining or
# Trajectory objects, which contain energy, forces, stress calculated with the reference
# calculator.
initial_training_data = [nlread('HfO2_crystal_training.hdf5', TrainingSet)[0]]

# Use the predefined small MTP basis.
mtp_basis = PredefinedBasisSmall

# Optimize the non-linear coefficients on the energy only.
nl_parameters = NonLinearCoefficientsParameters(
    perform_optimization=True,
    energy_only=True,
)

fitting_parameters = MomentTensorPotentialFittingParameters(
    basis_size=mtp_basis,
    outer_cutoff_radii=4.5*Ang,
    mtp_filename='HfO2_active_learning.mtp',
    non_linear_coefficients_parameters=nl_parameters,
    use_element_specific_coefficients=True,
)

active_learning = ActiveLearningSimulation(
    fitting_parameters=fitting_parameters,
    initial_training_data=initial_training_data,
    mtp_study_filename='HfO2_mtp_study',
    mtp_study_object_id='HfO2',
    reference_calculator=reference_calculator,
    candidate_threshold=1.0,
    retrain_threshold=3.0,
    check_interval=20,
    max_forces_check=10.0*eV/Ang,
    use_stress=True,
    candidate_trajectory_filename='HfO2_am_active_learning_candidates.hdf5',
    restart_simulation=True,
    extrapolation_selection_parameters=ExtrapolationSelectionParameters(
        extrapolation_grade_algorithm=QueryByCommitteeForces,
        descriptor_cutoff=0.1,
    ),
)

# Set up a high-temperature MD at 3000 K.
initial_velocity = MaxwellBoltzmannDistribution(
    temperature=3000.0*Kelvin,
)

method = Langevin(
    time_step=1*femtoSecond,
    reservoir_temperature=3000*Kelvin,
    friction=0.01*femtoSecond**-1,
    initial_velocity=initial_velocity,
)

constraints = [FixCenterOfMass()]

# Use 8 different initial configurations which are loaded from file and collected in a list. The
# trajectory filenames are also set up here, so that each MD trajectory is written to a different
# file.
number_of_initial_configurations = 8
initial_configuration_list = []
trajectory_filename_list = []
for i in range(number_of_initial_configurations):
    configuration = nlread(f'initial_configuration_{i}.hdf5', BulkConfiguration)[0]
    initial_configuration_list.append(configuration)
    trajectory_filename_list.append(f'HfO2_am_active_learning_3000K_{i}.hdf5')

# Run the MD simulation through the active learning object.
md_trajectory = active_learning.runMolecularDynamics(
    initial_configuration_list,
    constraints=constraints,
    trajectory_filename=trajectory_filename_list,
    steps=100000,
    log_interval=1000,
    method=method,
    domain_decomposition_pattern=[1, 1, 1],
)

# Extract the additional training data that has been added during active learning, as TrainingSet
# object.
additional_training_data = active_learning.additionalTrainingSet()