import simbios.simtk as simtk
import sys

def main():
    # molecule-specialized simbody System
    system = simtk.CompoundSystem()
    matter = simtk.SimbodyMatterSubsystem(system)
    decorations = simtk.DecorationSubsystem(system)
 
    # molecular force field
    forceField = simtk.DuMMForceFieldSubsystem(system) 

    # GeneralForceSubsystem forces(system)
    # VanderWallSphere boundary(forces, forceField, Vec3(0,0,0), 0.50, 0.2, 0.001)
    
    # Define an atom class for argon
    forceField.defineAtomClass( 
        forceField.getNextUnusedAtomClassIndex(), 
        "argon", 
        18, 
        0, 
        0.188, 
        0.100
        )
    forceField.defineChargedAtomType(
        forceField.getNextUnusedChargedAtomTypeIndex(), 
        "argon", 
        forceField.getAtomClassIndex("argon"), 
        0.0 # neutral charge
        )

    if not simtk.Biotype.exists("argon", "argon"):
        simtk.Biotype.defineBiotype(simtk.Element.Argon(), 0, "argon", "argon")

    forceField.setBiotypeChargedAtomType( forceField.getChargedAtomTypeIndex("argon"), simtk.Biotype.get("argon", "argon").getIndex() )
    forceField.setGbsaGlobalScaleFactor(0)

    argonAtom1 = simtk.Argon()
    argonAtom2 = simtk.Argon()

    # place first argon atom, units are nanometers
    system.adoptCompound(argonAtom1, simtk.Vec3(-0.3, 0, 0)) 

    # place second argon atom, units are nanometers
    system.adoptCompound(argonAtom2, simtk.Vec3( 0.3, 0, 0)) 

    system.defaultSubsystem.addEventReporter(
            simtk.VTKEventReporter(system, 0.200))
    
    system.modelCompounds() # finalize multibody system

    # Add larger space filling sphere for each argon atom
    for argon in argonAtom1, argonAtom2:
	atomIx = simtk.Compound.AtomIndex(0)
        decorations.addBodyFixedDecoration(
            argon.getAtomMobilizedBodyIndex(atomIx),
            argon.getAtomLocationInMobilizedBodyFrame(atomIx),
            simtk.DecorativeSphere(0.188).setColor((0.5,1,0.5)).setOpacity(1.0).setResolution(3))

    state = system.realizeTopology() 

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

if __name__ == '__main__':
    main()