# Example: Creating a TrainingSet with additional properties using Table

# Assume we have a list of configurations
configurations = [...]  # List of BulkConfiguration or MoleculeConfiguration instances

# Configuration-level property (one scalar value per configuration)
bandgaps = [1.0*eV, 1.1*eV, 1.3*eV]  # One value per configuration

# Atom-wise properties (one array per configuration, with one value per atom)
# Each configuration may have a different number of atoms
atomic_charges = [
    [0.1, -0.1, 0.2, -0.2]*e,  # 4 atoms in first configuration
    [0.15, -0.15]*e,  # 2 atoms in second configuration
    [0.1, -0.1, 0.1, -0.1]*e,  # 4 atoms in third configuration
]

# Create a Table with configurations and properties
table = Table()

# Add configuration column
table.addInstanceColumn(
    key="configuration",
    title="Configuration",
    types=(BulkConfiguration, MoleculeConfiguration),
)

# Add configuration-level scalar property
table.addQuantityColumn(key="bandgap", title="Band Gap", unit=eV)

# Add atom-wise property column
table.addInstanceColumn(
    key="atomic_charges",
    title="Atomic Charges",
    types=(PhysicalQuantity,),
)

# Populate the table
for configuration, bandgap, charges in zip(configurations, bandgaps, atomic_charges):
    table.append(
        configuration=configuration,
        bandgap=bandgap,
        atomic_charges=charges,
    )

# Create TrainingSet from Table - additional properties are automatically included
training_set = TrainingSet(configurations=table)

# Access additional properties
print("Available properties:", training_set.additionalPropertyNames())

# Retrieve all values for a property
all_bandgaps = training_set.additionalProperty("bandgap")
print(f"Band gaps: {all_bandgaps}")

# Retrieve property for a specific configuration
bandgap_0 = training_set.additionalProperty("bandgap", index=0)
print(f"Band gap of first configuration: {bandgap_0}")

# Retrieve atom-wise property
charges_0 = training_set.additionalProperty("atomic_charges", index=0)
print(f"Atomic charges of first configuration: {charges_0}")