setVerbosity(MinimalLog)

#----------------------------------------
# Exchange-Correlation
#----------------------------------------
exchange_correlation = SGGA.PBE

k_point_sampling = KpointDensity(
    density_a=7.0*Angstrom,
    density_b=7.0*Angstrom,
    density_c=7.0*Angstrom
)

numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=k_point_sampling
)

#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = [
    BasisGGAPAW.Molybdenum_Accurate,
    BasisGGAPAW.Ruthenium_Accurate,
    BasisGGAPAW.Copper_Accurate,
    BasisGGAPAW.Cobalt_Accurate,
    ]

iteration_control_parameters = IterationControlParameters(
    tolerance=5e-05,
    max_steps=400,
    damping_factor=0.2,
    number_of_history_steps=15,
    non_convergence_behavior=StopCalculation(),
    enable_scf_stop_file=False
)

density_matrix_method = DiagonalizationSolver(
    optimize_for_speed_over_memory=True
)

algorithm_parameters = AlgorithmParameters(
    density_matrix_method=density_matrix_method,
    reciprocal_space_paw_quantities=True,
    scf_restart_step_length=0.3*Angstrom
)

calculator = LCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    iteration_control_parameters=iteration_control_parameters,
    checkpoint_handler=NoCheckpointHandler,
    algorithm_parameters=algorithm_parameters
)

# Read the initial training dataset for the active learning simulation.
dataset = nlread('Initial_training_data.hdf5', MomentTensorPotentialTraining)[-1]
initial_training_data = [TrainingSet(dataset, recalculate_training_data=False)]

# Define the MTP fitting parameters.
nl_parameters = NonLinearCoefficientsParameters(
    perform_optimization=False,
    initial_coefficients=Random,
)

fitting_parameters = MomentTensorPotentialFittingParameters(
    basis_size=800,
    outer_cutoff_radii=4.5*Ang,
    mtp_filename='MTP_CSP_AL.mtp',
    non_linear_coefficients_parameters=nl_parameters,
    forces_cap=500.0*eV/Ang,
    use_element_specific_coefficients=True,
)

# Set up the active learning simulation.
al = ActiveLearningSimulation(
    fitting_parameters=fitting_parameters,
    initial_training_data=initial_training_data,
    mtp_study_filename='MTP_training_CSP_AL.hdf5',
    mtp_study_object_id='training',
    reference_calculator=calculator,
    candidate_threshold=1.0,
    retrain_threshold=3.0,
    check_interval=1,
    max_forces_check=10.0*eV/Ang,
    minimum_bond_length_percent=0.2,
    candidate_trajectory_filename='MTP_training_CSP_AL_candidates.hdf5',
    restart_simulation=False,
    extrapolation_selection_parameters=ExtrapolationSelectionParameters(
        extrapolation_grade_algorithm=QueryByCommittee,
        committee_size=6,
        descriptor_cutoff=0.2,
        descriptor_basis_size=100,
    ),
)

# Read the initial population from file.
initial_population = nlread('Initial_population.hdf5', Population)[0]

# Run the crystal structure predicition with active learning.
global_optimization = CrystalStructurePrediction(
    initial_population,
    number_of_generations=10,
    selection_pressure=2.0,
    number_of_elites=10,
    number_to_promote=4,
    heredity_probability=50.0,
    permutation_probability=15.0,
    mutation_probability=35.0,
    sigma_lattice=0.7,
    max_forces=0.01*eV/Angstrom,
    max_stress=0.1*GPa,
    max_steps=1000,
    max_step_length=0.2*Angstrom,
    external_pressure=0*GPa,
    active_learning=al,
    log_filename_prefix='csp_al',
)