from simbios.simtk import *

# Propane is a three carbon linear alkane
# C(3)H(8), or CH3-CH2-CH3
class Propane(Molecule):
    # constructor
    def __init__(self):
        super(Propane, self).__init__()
        self.setCompoundName("Propane")

        # First atom
        self.setBaseAtom( AliphaticCarbon("C1") )
        self.setAtomBiotype("C1", "propane", "C1_or_C3")
        self.convertInboardBondCenterToOutboard() # this is the root of the parent compound

        # Second atom
        self.bondAtom( AliphaticCarbon("C2"), "C1/bond1" )
        self.setAtomBiotype("C2", "propane", "C2")

        # Third atom
        self.bondAtom( AliphaticCarbon("C3"), "C2/bond2" )
        self.setAtomBiotype("C3", "propane", "C1_or_C3")

        # First methyl hydrogens
        self.bondAtom( AliphaticHydrogen("H11"), "C1/bond2" )
        self.bondAtom( AliphaticHydrogen("H12"), "C1/bond3" )
        self.bondAtom( AliphaticHydrogen("H13"), "C1/bond4" )
        self.setAtomBiotype("H11", "propane", "H1_or_H3")
        self.setAtomBiotype("H12", "propane", "H1_or_H3")
        self.setAtomBiotype("H13", "propane", "H1_or_H3")

        # Second methylene hydrogens
        self.bondAtom( AliphaticHydrogen("H21"), "C2/bond3" )
        self.bondAtom( AliphaticHydrogen("H22"), "C2/bond4" )
        self.setAtomBiotype("H21", "propane", "H2")
        self.setAtomBiotype("H22", "propane", "H2")

        # Third methyl hydrogens
        self.bondAtom( AliphaticHydrogen("H31"), "C3/bond2" )
        self.bondAtom( AliphaticHydrogen("H32"), "C3/bond3" )
        self.bondAtom( AliphaticHydrogen("H33"), "C3/bond4" )
        self.setAtomBiotype("H31", "propane", "H1_or_H3")
        self.setAtomBiotype("H32", "propane", "H1_or_H3")
        self.setAtomBiotype("H33", "propane", "H1_or_H3")
    

    # create charged atom types
    # ensure that charges sum to zero, unless molecule has a formal charge
    # Must be called AFTER first propane is declared, so Biotypes and atom classes will be defined
    def setAmberLikeParameters(forceField):
    
        forceField.defineChargedAtomType(
            forceField.getNextUnusedChargedAtomTypeIndex(),
            "propane C1_or_C3", 
            forceField.getAtomClassIndex("CT"), # amber tetrahedral carbon
            -0.060 # made up
            )
        forceField.setBiotypeChargedAtomType( 
            forceField.getChargedAtomTypeIndex("propane C1_or_C3"),
            Biotype.get("propane", "C1_or_C3").getIndex() )

        forceField.defineChargedAtomType(
            forceField.getNextUnusedChargedAtomTypeIndex(),
            "propane C2", 
            forceField.getAtomClassIndex("CT"), # amber tetrahedral carbon
            -0.040 # made up
            )
        forceField.setBiotypeChargedAtomType( 
            forceField.getChargedAtomTypeIndex("propane C2"),
            Biotype.get("propane", "C2").getIndex() )

        forceField.defineChargedAtomType(
            forceField.getNextUnusedChargedAtomTypeIndex(),
            "propane H1_or_H3", 
            forceField.getAtomClassIndex("HC"), # amber tetrahedral carbon
            0.020 # made up, use net neutral charge
            )
        forceField.setBiotypeChargedAtomType( 
            forceField.getChargedAtomTypeIndex("propane H1_or_H3"),
            Biotype.get("propane", "H1_or_H3").getIndex() )

        forceField.defineChargedAtomType(
            forceField.getNextUnusedChargedAtomTypeIndex(),
            "propane H2", 
            forceField.getAtomClassIndex("HC"), # amber tetrahedral carbon
            0.020 # made up, use net neutral charge
            )
        forceField.setBiotypeChargedAtomType( 
            forceField.getChargedAtomTypeIndex("propane H2"),
            Biotype.get("propane", "H2").getIndex() )

    setAmberLikeParameters = staticmethod(setAmberLikeParameters)

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

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

    propane1 = Propane()
    propane2 = Propane()

    Propane.setAmberLikeParameters(forceField)

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

    # place second propane, units are nanometers
    system.adoptCompound(propane2, Vec3( 0.5, 0, 0)) 

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

    system.modelCompounds() # finalize multibody system

    state = system.realizeTopology()

    integ = VerletIntegrator(system)
    ts = TimeStepper(system, integ)
    ts.initialize(state)
    ts.stepTo(30.0)

if __name__ == "__main__":
    main()