#!/usr/bin/python

# Standard imports
import sys
from time import time

from MMTK import *
from MMTK.ForceFields import Amber94ForceField
from MMTK.ForceFields import HarmonicForceField
from MMTK.Proteins import Protein
from MMTK.PDB import PDBConfiguration

from MMTK.Trajectory import Trajectory, TrajectoryOutput, StandardLogOutput
from MMTK.Dynamics import VelocityVerletIntegrator


def main():
	# Get the protein filename that we are dealing with
	filename = 'insulin.pdb'
	if len(sys.argv) > 1:	
		filename = sys.argv[1]
	

	# Create universe and read in the protein
	universe = InfiniteUniverse(Amber94ForceField(None, {'method': 'multipole'}))
	#universe = InfiniteUniverse(Amber94ForceField())
	conf = PDBConfiguration(filename)
	chains = conf.createPeptideChains()
	universe.protein = Protein(chains)

	# Break the protein up into pieces for evaluation
	atoms = universe.protein.atoms
	s1 = Collection(atoms[:len(atoms)/2])
	s2 = Collection(atoms[len(atoms)/2:])
	
	# Compute energies for the universe as a whole and for sets separately
	num_times = 20
	
	e_total = universe.energy()
	t1b = time()
	for i in range(num_times):
		e_total = universe.energy()
	t1e = time()


	e_s1 = universe.energy(s1)
	t2b = time()
	for i in range(num_times):
		e_s1 = universe.energy(s1)
	t2e = time()
	
		
	e_s2 = universe.energy(s2) 
	t3b = time()
	for i in range(num_times):
		e_s2 = universe.energy(s2) 
	t3e = time()


	e_s1s2 = universe.energy(s1,s2)
	t4b = time()
	for i in range(num_times):
		e_s1s2 = universe.energy(s1,s2)
	t4e = time()


	print "%5.2f + %5.2f + %5.2f = %5.2f" % (e_s1, e_s2, e_s1s2, e_total)
	print "Time as a whole", t1e - t1b
	print "Time for the first set", t2e - t2b
	print "Time for the second set", t3e - t3b
	print "Time for between the sets", t4e - t4b
	

if __name__ == '__main__':
	main()