# --------------------------------------------------------------------------------------------------
# Create Methanol Molecule
# --------------------------------------------------------------------------------------------------
# Define elements
elements = [Carbon, Hydrogen, Hydrogen, Hydrogen, Oxygen, Hydrogen]

# Define coordinates
cartesian_coordinates = [[ 0.048333359398, -0.11839206808 , -0.068353692376],
                         [ 1.138333281204, -0.11839206808 , -0.068353692376],
                         [-0.31499994787 , -0.11839206808 , -1.096015473978],
                         [-0.31499994787 , -1.008373277446,  0.445477198425],
                         [-0.363333307269,  0.889981209366,  0.513830890801],
                         [-0.013333332377,  1.747302557837,  0.018856179479]]*Angstrom

# Set up configuration
methanol = MoleculeConfiguration(
    elements=elements,
    cartesian_coordinates=cartesian_coordinates
)

# Add tags
methanol.addTags('DREI_C_3',  [0])
methanol.addTags('DREI_H_',   [1, 2, 3])
methanol.addTags('DREI_H__A', [5])
methanol.addTags('DREI_O_3',  [4])

# Add bonds
bonds = [[0, 1],
         [0, 2],
         [0, 3],
         [0, 4],
         [4, 5]]
methanol.setBonds(bonds)

# --------------------------------------------------------------------------------------------------
# Calculate Atomic Charges
# --------------------------------------------------------------------------------------------------
# Create the charger and calculate the charges
charger = ReaxFFAtomicCharges()
methanol_charges = charger.charges(methanol)

# Print charges
print('METHANOL CHARGES')
print("----------------------------")
print(f'{"Number":<8s}{"Atom":<8s}{"Charge":>12s}')
print("----------------------------")
for i, q in enumerate(methanol_charges):
    element = methanol.elements()[i].symbol()
    charge = q.inUnitsOf(elementary_charge)
    print(f'{i+1:<8d}{element:<8s}{charge:>12.4f}')
print("----------------------------")

# --------------------------------------------------------------------------------------------------
# Create A Box Of Methanol Molecules
# --------------------------------------------------------------------------------------------------
# Define the volume in which the molecules should be packed
u1, u2, u3 = numpy.diag([40.0] * 3) * Angstrom
cell = UnitCell(u1, u2, u3)

# Set the size of the buffer region
buffer_size = 2.0 * Angstrom

# Set Packmol parameters
tolerance = 2.0 * Angstrom
maximum_loops = 30

# Run the Packmol script to generate the packed configuration
methanol_box, packing_successful, packmol_message = runPackmol(
    [(methanol, 900)],
    cell,
    tolerance,
    buffer_size,
    maximum_loops,
)

# --------------------------------------------------------------------------------------------------
# Add A Dreiding Calculator With The Calculated Charges
# --------------------------------------------------------------------------------------------------
# Repeat the charges for each molecule
methanol_box_charges = numpy.tile(
    methanol_charges.inUnitsOf(elementary_charge),
    900
) * elementary_charge

# Create the potential builder
potential_builder = DreidingPotentialBuilder()

# Create the calculator, setting the atomic charges
calculator = potential_builder.createCalculator(
    methanol_box,
    atomic_charges=methanol_box_charges
)

# Assign the calculator to the configuration
methanol_box.setCalculator(calculator)

# --------------------------------------------------------------------------------------------------
# Save The Configuration
# --------------------------------------------------------------------------------------------------
nlsave('Methanol_Box.hdf5', methanol_box)