# Define the propylene monomer
elements = [Carbon, Hydrogen, Hydrogen, Hydrogen, Carbon, Hydrogen, Hydrogen,
            Carbon, Hydrogen, Hydrogen, Hydrogen]
cartesian_coordinates = [[ 0.000000111747,  0.00000161924 , -0.000056562984],
                         [-0.630526888897,  0.89112270661 ,  0.025980951533],
                         [ 0.630920857538,  0.026416367051, -0.890816508519],
                         [-0.630272811723, -0.891281341673, -0.026694611157],
                         [ 0.872716358782, -0.036383769323,  1.235342747223],
                         [ 0.241761994012, -0.062941276536,  2.12614593826 ],
                         [ 1.502988846798,  0.854823441738,  1.262143739191],
                         [ 1.746330901043, -1.271157968252,  1.199296815275],
                         [ 2.797055327872, -0.974121818704,  1.208611065519],
                         [ 1.535794456846, -1.891496275549,  2.072978259283],
                         [ 1.536636911043, -1.839009024623,  0.29056935134 ]]*Angstrom
monomer = MoleculeConfiguration(
    elements=elements,
    cartesian_coordinates=cartesian_coordinates
)
monomer.findBonds()
monomer.addTags('HEAD_CONNECT',   [0])
monomer.addTags('HEAD_GROUP_A',   [3])
monomer.addTags('TAIL_CONNECT',   [4])
monomer.addTags('TAIL_GROUP_A',   [5])
monomer.addTags('TAIL_GROUP_B',   [6])

# Define the polymer sequence
polymer_sequence = PolymerSequence(
    monomer_configurations=[monomer],
    number_of_monomers=33,
    number_of_chains=33,
    tactic_ratio=0.5,
    head_tail_flip_probability=0.0,
)

# Define the forcefield
builder = OPLSPotentialBuilder(include_electrostatic=False)

# Build the initial configuration
polymer_builder = PolymerMonteCarloBuilder(
    polymer_sequence=[polymer_sequence],
    density=0.90*gram/cm**3,
    potential_builder=builder,
)
initial_configuration = polymer_builder.buildPolymerConfiguration()

# Set a calculator without electrostatics on the configuration
calculator = SoftMatterCalculator(builder)
initial_configuration.setCalculator(calculator)

# Perform pre-equilibration of the configuration
equilibration_method = ForceCappedEquilibration()
pre_configuration = equilibration_method.runEquilibration(initial_configuration)

# Add a new calculator with electrostatics
builder = OPLSPotentialBuilder()
calculator = SoftMatterCalculator(builder)
pre_configuration.setCalculator(calculator)

# Define the equilibration simulation
simulation = SoftMatterDynamicsSimulation(
    pre_configuration,
    NPTNoseHooverMonteCarlo(),
)
simulation.setLogging(10000)
simulation.setTrajectory(5000, f"equilibrated_pp.hdf5")

# Set the equilibration profiles
profiles = polymerEquilibrationProfiles()

# Run the equilibration
simulation.simulate(1560000, profiles=profiles)
final_configuration = simulation.currentConfiguration()