# This function replaces the setup of NonLinearCoefficientsParameters and
# MomentTensorPotentialFittingParameters without the need of a loop.
fitting_parameters_list = scanOverNonLinearCoefficients(
    number_of_initial_guesses=30,
    basis_size=n_basis,
    outer_cutoff_radii=4.5*Ang,
    mtp_filename_suffix='MTP_fit.mtp',
    use_element_specific_coefficients=True,
    random_seed=42,
)

# Perform an initial MTP training using a list of fitting parameters.
mtp_training = MomentTensorPotentialTraining(
    filename='pre_training.hdf5',
    object_id='mtp_training',
    training_sets=initial_training_data,
    calculator=reference_calculator,
    calculate_stress=True,
    fitting_parameters_list=fitting_parameters_list,
    train_test_split=0.9,
    random_seed=13345,
    log_filename_prefix='pre_training',
)
mtp_training.update()

# Determine the best fit.
best_fit_index = mtp_training.rankFits(
    data_tags=None,
    weights=[[1, 1, 1], [1, 1, 1]],
    statistical_measure=R2Score
)[0][0]

# Get the parameters of the best fit which can be passed to the fitting_parameters keyword of
# ActiveLearningSimulation.
best_fitting_parameters = mtp_training.fittingParametersList()[best_fit_index]

# Set up the active ActiveLearningSimulation as usual.
active_learning = ActiveLearningSimulation(
    fitting_parameters=best_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,
)
