from visual import * #vector, mag, norm, rate
from operator import add

# Mike Sullivan

default_dt = 0.001
g = 9.8
G = 6.7e-11
gv = vector(0, -g, 0)

#### Graphics ###
# All graphics objects contain information about what they look like,
# but not about what they represent. What they represent is passed
# to update() in order to redraw the object

# A graphics object that does nothing
class NullPhysicsGraphics:
	def update(self, phys):
		pass

# Generic graphics for a physics object
class PhysicsObjectGraphics:
	def __init__(self, gfx=None, pArrow=None, doTrail=False):
		self.gfx = gfx
		self.pArrow = self.trail = None
		if (pArrow): self.pArrow = arrow(color=pArrow)
		if (doTrail): self.trail = curve(color=gfx.color)

	def update(self, phys):
		self.gfx.pos = phys.position
		if (self.pArrow):
			self.pArrow.pos = phys.position
			self.pArrow.axis = phys.momentum
		if (self.trail):
			self.trail.append(pos=phys.position)

# Graphics for a spring using just a line
class RopeGraphics:
	def __init__(self, curve):
		self.curve = curve
	def update(self, spring):
		self.curve.pos = [spring.m1.position, spring.m2.position]

# Graphics for a spring using a helix
class SpringGraphics:
	def __init__(self, helix):
		self.helix = helix
	def update(self, spring):
		self.helix.pos = spring.m1.position
		self.helix.axis = spring.m2.position - spring.m1.position

#### Physics ####

# An object that obeys the laws of physics
class PhysicsObject:
	def __init__(self, position=vector(0,0,0), momentum=vector(0,0,0), mass=1, gfx=NullPhysicsGraphics()):
		self.position = vector(position)
		self.momentum = vector(momentum)
		self.mass = mass
		self.gfx = gfx

	def apply_force(self, f, dt):
		self.momentum += f*dt

	def update_position(self, dt):
		if self.mass > 0:
			self.position += self.momentum/self.mass*dt

	def update(self):
		self.gfx.update(self)

# An object that is fixed in place
class FixedPhysicsObject(PhysicsObject):
	def apply_force(self, f, dt):
		pass
	def update_position(self, dt):
		pass

### Actors apply forces to physics objects ###

# An actor that applies a constant force on a set of objects (ex: gravity on earth)
class ConstantForce:
	def __init__(self, f, objs):
		self.f = f
		self.objs = objs
	def apply_forces(self, dt):
		for o in self.objs:
			o.apply_force(self.f, dt)
	def update(self):
		pass
			
# An actor that represents a spring connecting m1 and m2
class Spring:
	def __init__(self, k, length, m1, m2, gfx=NullPhysicsGraphics()):
		self.k = k
		self.length = length
		self.gfx = gfx
		self.m1 = m1
		self.m2 = m2

	# on m1 by m2
	def get_force(self, m1, m2):
		rv = self.m1.position - self.m2.position
		s = mag(rv) - self.length
		v = -self.k * s * norm(rv)
		return v

	# apply the spring 
	def apply_forces(self, dt):
		f = self.get_force(self.m1, self.m2)
		self.m1.apply_force(f, dt)
		self.m2.apply_force(-f, dt)

	def update(self):
		self.gfx.update(self)

# A spring with some damping
class DampedSpring(Spring):
	def __init__(self, kdamp, k, length, m1, m2, gfx=NullPhysicsGraphics()):
		Spring.__init__(self, k, length, m1, m2, gfx)
		self.kdamp = kdamp

	def apply_forces(self, dt):
		Spring.apply_forces(self, dt)
		v = -(self.m1.momentum/self.m1.mass - self.m2.momentum/self.m2.mass) * self.kdamp
		#print v
		self.m1.apply_force(v, dt)
		self.m2.apply_force(-v, dt)
		

# An actor that applies Newton's law of universal gravitation to a set of objects
class UniversalGravity:
	def __init__(self, objs):
		self.objs = objs

	# Calculate the gravitational force on m1 by m2
	def grav_force(self, m1, m2):
		rv = m1.position - m2.position
		return -G * m1.mass * m2.mass / mag2(rv) * norm(rv)

	# Solve (numerically) the n-body problem (sort of)
	#FIXME: implement this more efficiently
	def apply_forces(self, dt):
		for m1 in self.objs:
			for m2 in self.objs:
				if (m1 != m2):
					m1.apply_force(self.grav_force(m1, m2), dt)

	def update(self):
		pass

## Actors to be implemented ##
# Damping?
# Friction
# Coulomb's law
# Collision Detection
# Normal forces?

class Simulation:
	def __init__(self, actors, objs, dt=default_dt, scale=1, callback=None):
		self.actors = actors
		self.objs = objs
		self.dt = dt
		self.scale = scale
		self.callback = callback

	# Run one step of the simulation
	def step(self):
		for a in self.actors:
			a.apply_forces(self.dt)
		for obj in self.objs:
			obj.update_position(self.dt)
		for o in self.actors + self.objs:
			o.update()

	# Run the simulation!
	def simulate(self):
		i = 0
		while 1:
			rate(self.scale/self.dt)
			self.step()
			if self.callback:
				self.callback(self, i)
			i+=1
		

def debug_obj0(sim, i):
	print str(i) + " (" + str(i*sim.dt) + "s): " + str(sim.objs[0].momentum) + " - " + str(sim.objs[0].position)