
# Invivo AI challenge simulation environment
# Copyright 2009 Simon Funk
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import Model as m
from Box2D import *
import time

# See http://www.box2d.org/wiki/index.php?title=Hello.py
# and http://www.box2d.org/wiki/index.php?title=Python_Manual

class WorldSim(object):

	def __init__(self, model):
		self.model = model

		worldAABB=b2AABB()
		worldAABB.lowerBound.Set(model.bounds[0][0]-1., model.bounds[0][1]-1.)
		worldAABB.upperBound.Set(model.bounds[1][0]+1., model.bounds[1][1]+1.)

		gravity = b2Vec2(0, model.gravity)
		#doSleep = True 									# No need to actively sim resting bodies.
		doSleep = False 									# Hmm, sleeping objects may not respond to changing control forces.

		self.world = b2World(worldAABB, gravity, doSleep)

		# BUG workaround!  Must store local copy because SetContactListener does not add ref-count to param:
		self.contacter = MyContactListener(self)
		self.world.SetContactListener(self.contacter)

		#
		# Dicts mapping model_body to sim_body,
		# and model_joint to sim_joint:
		#
		self.bodymap  = {}
		self.jointmap = {}

		for mbody in model.bodies:

			bodyDef          = b2BodyDef()
			bodyDef.position = mbody.position
			bodyDef.angle    = mbody.angle
			bodyDef.userData = mbody

			sbody = self.world.CreateBody(bodyDef)

			for mshape in mbody.shapes:

				if isinstance(mshape, m.Poly):
					shapeDef = b2PolygonDef()
					shapeDef.setVertices(mshape.points)
				elif isinstance(mshape, m.Circle):
					shapeDef = b2CircleDef()
					shapeDef.radius = mshape.radius
					shapeDef.localPosition.Set(mshape.center[0], mshape.center[1])
				else:
					raise Exception("Unknown Shape class: %s"%(mshape.__class__.__name__,))

				if mbody.dynamic:
					shapeDef.density = mshape.density

				shapeDef.friction	= mshape.friction
				shapeDef.restitution = mshape.restitution

				shapeDef.filter.categoryBits = mshape.layers
				shapeDef.filter.maskBits	 = mshape.layers

				shapeDef.userData = mshape

				sshape = sbody.CreateShape(shapeDef)

				# Do we want to save sshape?

			sbody.SetMassFromShapes()
			sbody.linearVelocity  = mbody.velocity
			sbody.angularVelocity = mbody.rotating

			if mbody.ccd:
				sbody.isBullet = True

			self.bodymap[mbody] = sbody

		for mjoint in model.joints:

			if isinstance(mjoint, m.Pivot):
				jd = b2RevoluteJointDef()
				jd.Initialize(self.bodymap[mjoint.body1], self.bodymap[mjoint.body2], mjoint.position)
				jd.lowerAngle	 = mjoint.minAngle
				jd.upperAngle	 = mjoint.maxAngle
				jd.enableLimit	 = mjoint.limitAngle
				jd.enableMotor   = True	# Even if no torque applied, there may be friction.
				self.setTorque(jd, mjoint)
			else:
				raise Exception("Unknown Joint class: %s"%(mjoint.__class__.__name__,))

			self.jointmap[mjoint] = self.world.CreateJoint(jd).getAsType()

	@staticmethod
	def setTorque(spivot, mjoint):
		"""Updates the dynamic simulation pivot params from the current model pivot params.
		NOTE mjoint.speed must be current because this is used to implement friction and max speed.
		"""

		torque = mjoint.torque
		fric   = mjoint.friction
		speed  = mjoint.speed

		if torque > 0.:
			atorque = torque	# abs(torque)
			storque = 1.		# sign(torque)
		elif torque < 0.:
			atorque = -torque
			storque = -1.
		else:
			atorque = 0.
			storque = 0.

		#
		# We have to enumerate various cases because we are emulating friction
		#  with active torque and we have to take care not to add energy to
		#  the system when a reversal happens inside the timestep.
		#
		# Also, the underlying sim does not let us apply torque directly, but
		#  rather we have to factor it into maxMotorTorque and motorSpeed, the
		#  latter of which we set to an extreme we don't normally expect to reach.
		#
		if speed*storque >= 0.:	# Pushing in same direction as current motion (and against friction)
			if atorque > fric:		#     With enough force to overcome friction:
				spivot.maxMotorTorque = atorque	- fric	# Remaining torque (after friction)...
				spivot.motorSpeed     = storque * 100.	# serves to accelerate current motion
			else:					#     Without enough force to overcome friction:
				spivot.maxMotorTorque = fric - atorque	# Remaining friction...
				spivot.motorSpeed     = 0.				# ...serves to stop the motion (but no more)
		else:						# Pushing against current motion (and with friction)
			spivot.maxMotorTorque = fric + atorque	# Forces sum...
			spivot.motorSpeed     = storque * 100.	# ...in direction of torque
			# BUG: This case unfortunately adds energy during reversals because friction
			#  continues to contribute energy in the same direction after the reversal.
			# Is there a fix for this?

	def run(self, endtime=None):

		timeStep  = self.model.timestep # time betwen frames
		iterfac   = 1
		vel_iters = 10 * iterfac
		pos_iters = 8  * iterfac

		t0 = t = time.time()

		hasContacts = [mbody for mbody in self.bodymap if mbody.contacts is not None]

		while not self.model.stop and not (endtime and t-t0 >= endtime):

			for mbody in hasContacts:
				mbody.contacts = []

			self.world.Step(timeStep, vel_iters, pos_iters)

			for mbody, sbody in self.bodymap.iteritems():
				mbody.move(sbody.GetPosition().tuple(), sbody.GetAngle())
				mbody.velocity = sbody.linearVelocity.tuple()
				mbody.rotating = sbody.angularVelocity

			for mjoint, sjoint in self.jointmap.iteritems():
				mjoint.angle = sjoint.GetJointAngle()	# sjoint -> mjoint
				mjoint.speed = sjoint.GetJointSpeed()
				self.setTorque(sjoint, mjoint)			# mjoint -> sjoint

			t  += timeStep
			now = time.time()
			if now < t:
				time.sleep(t-now)

			self.model.changed(t - t0)

class MyContactListener(b2ContactListener):

	def __init__(self, world):
		super(MyContactListener, self).__init__() 

		self.world = world	# WorldSim object

	if False:
		def Add(self, cr):
			if cr.shape1.GetBody().isBullet or cr.shape2.GetBody().isBullet:
				print "Add %s"%(cr.position,)

		def Persist(self, cr):
			if cr.shape1.GetBody().isBullet or cr.shape2.GetBody().isBullet:
				print "Persist %s"%(cr.position,)

		def Remove(self, cr):
			if cr.shape1.GetBody().isBullet or cr.shape2.GetBody().isBullet:
				print "Remove %s"%(cr.position,)

	def Result(self, cr):
		sb1 = cr.shape1.GetBody()	# Sim bodies
		sb2 = cr.shape2.GetBody()
		mb1 = sb1.userData			# Model bodies
		mb2 = sb2.userData
		if mb1 and mb2:
			if mb1.contacts is not None:
				mb1.contacts.append(m.Contact(mb2, cr.normalImpulse, cr.position.tuple()))
			if mb2.contacts is not None:
				mb2.contacts.append(m.Contact(mb1, cr.normalImpulse, cr.position.tuple()))


# b2ContactResult(
#       id             = b2ContactID(
#		   features = b2ContactID_features(
#			   flip           = 0,
#			   incidentEdge   = 0,
#			   incidentVertex = 255,
#			   referenceEdge  = 0),
#		   key      = 16711680),
#       normal         = b2Vec2(0.071247,0.997459),
#       normalImpulse  = 0.0801500678062,
#       position       = b2Vec2(6.53103,0.0280052),
#       shape1         = b2Shape(
#		   density          = 0.0,
#		   filter           = b2FilterData(
#			   categoryBits = 65535,
#			   groupIndex   = 0,
#			   maskBits     = 65535),
#		   friction         = 0.20000000298,
#		   restitution      = 0.0,
#		   userData         = None,
#		   GetBody()        = b2Body(
#			   allowSleep        = True,
#			   angle             = 0.0,
#			   angularDamping    = 0.0,
#			   angularVelocity   = 0.0,
#			   fixedRotation     = False,
#			   isBullet          = False,
#			   isSleeping        = False,
#			   linearDamping     = 0.0,
#			   linearVelocity    = b2Vec2(0,0),
#			   massData          = b2MassData(
#				   I      = 0.0,
#				   center = b2Vec2(0,0),
#				   mass   = 0.0),
#			   position          = b2Vec2(0,0),
#			   userData          = None,
#			   GetInertia()      = 0.0,
#			   GetLocalCenter()  = b2Vec2(0,0),
#			   GetMass()         = 0.0,
#			   GetWorldCenter()  = b2Vec2(0,0),
#			   GetXForm()        = b2XForm(
#				   R        = b2Mat22(
#					   col1       = b2Vec2(1,0),
#					   col2       = b2Vec2(-0,1),
#					   GetAngle() = 0.0),
#				   position = b2Vec2(0,0)),
#			   IsBullet()        = False,
#			   IsDynamic()       = False,
#			   IsFrozen()        = False,
#			   IsFixedRotation() = False,
#			   IsSleeping()      = False,
#			   IsStatic()        = True),
#		   GetFilterData()  = b2FilterData(
#		   categoryBits = 65535,
#		   groupIndex   = 0,
#		   maskBits     = 65535),
#		   GetSweepRadius() = 5.87870931625,
#		   GetType()        = 1,
#		   IsSensor()       = False),
#       shape2         = b2Shape(
#		   density          = 0.20000000298,
#		   filter           = b2FilterData(
#			   categoryBits = 2,
#			   groupIndex   = 0,
#			   maskBits     = 2),
#		   friction         = 0.20000000298,
#		   restitution      = 0.0,
#		   userData         = None,
#		   GetBody()        = b2Body(
#			   allowSleep        = True,
#			   angle             = 4.00141429901,
#			   angularDamping    = 0.0,
#			   angularVelocity   = -0.0157184563577,
#			   fixedRotation     = False,
#			   isBullet          = False,
#			   isSleeping        = False,
#			   linearDamping     = 0.0,
#			   linearVelocity    = b2Vec2(0.0071337,0.000461275),
#			   massData          = b2MassData(
#				   I      = 0.00247265840881,
#				   center = b2Vec2(0,-2.48175e-10),
#				   mass   = 0.0231415946037),
#			   position          = b2Vec2(6.86263,0.456213),
#			   userData          = None,
#			   GetInertia()      = 0.00247265840881,
#			   GetLocalCenter()  = b2Vec2(0,-2.48175e-10),
#			   GetMass()         = 0.0231415946037,
#			   GetWorldCenter()  = b2Vec2(6.86263,0.456213),
#			   GetXForm()        = b2XForm(
#				   R        = b2Mat22(
#					   col1       = b2Vec2(-0.652573,-0.757726),
#					   col2       = b2Vec2(0.757726,-0.652573),
#					   GetAngle() = -2.2817709446),
#				   position = b2Vec2(6.86263,0.456213)),
#			   IsBullet()        = False,
#			   IsDynamic()       = True,
#			   IsFrozen()        = False,
#			   IsFixedRotation() = False,
#			   IsSleeping()      = False,
#			   IsStatic()        = False),
#		   GetFilterData()  = b2FilterData(
#		   categoryBits = 2,
#		   groupIndex   = 0,
#		   maskBits     = 2),
#		   GetSweepRadius() = 0.509999990463,
#		   GetType()        = 0,
#		   IsSensor()       = False),
#       tangentImpulse = 0.00880656857044)

