from simbios.simtk import *
import sys

def main():
    system = CompoundSystem()
    matter = SimbodyMatterSubsystem(system)
    decorations = DecorationSubsystem(system)
    forceField = DuMMForceFieldSubsystem(system) 

    # Atom classes are available, but not charged atom types for ethane
    # in standard Amber force field
    forceField.loadAmber99Parameters()

    if not Biotype.exists("ethane", "C"):
        Biotype.defineBiotype(Element.Carbon(), 4, "ethane", "C")
    if not Biotype.exists("ethane", "H"):
        Biotype.defineBiotype(Element.Hydrogen(), 1, "ethane", "H")

    forceField.defineChargedAtomType(
        forceField.getNextUnusedChargedAtomTypeIndex(), 
        "ethane C", 
        forceField.getAtomClassIndex("CT"), 
        -0.060 # made up
        )
    forceField.setBiotypeChargedAtomType(
        forceField.getChargedAtomTypeIndex("ethane C"),
        Biotype.get("ethane", "C").getIndex() )

    forceField.defineChargedAtomType(
        forceField.getNextUnusedChargedAtomTypeIndex(), 
        "ethane H", 
        forceField.getAtomClassIndex("HC"), 
        0.020 # made up, use net neutral charge
        )
    forceField.setBiotypeChargedAtomType( 
        forceField.getChargedAtomTypeIndex("ethane H"),
        Biotype.get("ethane", "H").getIndex() )

    ethane1 = Ethane()
    ethane2 = Ethane()
    ethane1.writeDefaultPdb(sys.stdout)

    # place first ethane, units are nanometers
    # skew it a little to break strict symmetry
    system.adoptCompound(ethane1, Transform(Vec3(-0.5, 0, 0)) * Transform(Rotation(0.1, YAxis)) ) 

    # place second ethane, units are nanometers
    system.adoptCompound(ethane2, Vec3( 0.5, 0, 0)) 

    system.updDefaultSubsystem().addEventReporter(VTKEventReporter(system,
        0.050))

    system.modelCompounds() # finalize multibody system

    state = system.realizeTopology()

    # Simulate it.
    integ = VerletIntegrator(system)
    ts = TimeStepper(system, integ)
    ts.initialize(state)
    ts.stepTo(20.0)

if __name__ == '__main__':
    main()