# Load a TrainingSet with bandgap property data saved under the additional property key 'bandgap'
training_set = nlread("path/to/property_dataset.hdf5", TrainingSet)[0]

# Define model parameters
# Fine-tune the property model from a pre-trained force field MACE foundation model
model_params = MACEModelParameters(
    foundation_model_path="path/to/mace-mp-0b3-medium.model",
)

# Define training parameters
training_params = TrainingParameters(
    experiment_name="bandgap_model",
    random_seed=42,
)

# Define dataset parameters for band gap property
# Specify that we are training a configuration-level property (one scalar value per configuration)
dataset_params = GeneralPropertyDatasetParameters(
    property_key="bandgap",
    task_type=MLParameterOptions.TASK_TYPE.GENERAL,  # Configuration-level property
    validation_fraction=0.1,
)

# Define the fitting parameters
fitting_params = GeneralPropertyFittingParameters(
    model_parameters=model_params,
    dataset_parameters=dataset_params,
    training_parameters=training_params,
)

# Create the MachineLearnedPropertyTrainer
trainer = MachineLearnedPropertyTrainer(
    fitting_parameters=fitting_params,
    training_sets=training_set,
    train_test_split=0.8,
)

# Train the model
trainer.train()

# Get the trained PropertyPredictor from the trainer
predictor = trainer.propertyPredictor()

# Alternatively, load a PropertyPredictor directly from the saved model file
# The trained model is automatically saved as 'bandgap_model_<random_seed>.qatkpt'
predictor = PropertyPredictor("bandgap_model_42.qatkpt")

# Use the predictor to predict properties for new configurations
new_configuration = ...  # Load or define your atomic configuration here
predicted_property = predictor.predict(new_configuration)
print("Predicted bandgap:", predicted_property)
