import simbios.simtk as simtk
from simbios.simtk import streambuf, Transform, adaptbx_ostream
import sys

# Reimplement PeriodicPdbWriter, since the C++ wrapper is not working
# But this one isn't working either...
class PeriodicPdbWriter3(simtk.PeriodicEventReporter):
    def __init__(self, system, pyfile, interval):
        super(PeriodicPdbWriter3, self).__init__(interval)
        self.system = system
        self.pyfile = pyfile
        self.buf = streambuf(pyfile)
        self.ostream = adaptbx_ostream(self.buf)
        self.modelNumber = 1

    def handleEvent(self, state):
        self.system.realize(state, simtk.Stage.Position)
        nextAtomSerialNumber = 1
        print >>self.pyfile, "MODEL     %4d" % self.modelNumber
        for c in range(self.system.getNumCompounds()):
            c = simtk.CompoundSystem.CompoundIndex(c)
            compound = self.system.getCompound(c) # same runtime error
            compound.writePdb(state, self.pyfile)
            pass
        print >>self.pyfile, "ENDMDL"
        self.pyfile.flush()
        self.modelNumber += 1

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

    forceField.loadAmber99Parameters()

    rna = simtk.RNA("AUG")
    rna.assignBiotypes()
    system.adoptCompound(rna)
    
    system.updDefaultSubsystem().addEventReporter(
           simtk.VTKEventReporter(system, 0.020))

    # os = adaptbx_ostream(sys.stdout)
    # pdb = simtk.PeriodicPdbWriter(system, os, 0.100)
    pdb = PeriodicPdbWriter3(system, sys.stdout, 0.200)
    system.updDefaultSubsystem().addEventReporter(pdb)

    # Maintain a constant temperature
    # This thermostat might be the trouble...
    # Maintain a constant temperature
    system.updDefaultSubsystem().addEventHandler(
        simtk.VelocityRescalingThermostat(system,  293.15, 0.1))


    # finalize multibody system
    system.modelCompounds() 

    # Instantiate simbody model
    system.realizeTopology()
    state = system.updDefaultState()

    # Relax the structure before dynamics run
    simtk.LocalEnergyMinimizer.minimizeEnergy(system, state, 0.5)

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

if __name__ == '__main__':
    main()