
from types import *
import math,random

dt=0.001
CRASH_DIST=0.1

class Crashed:
    pass

class Vector:
    def __init__(self, *args):
        if len(args)==2:
            if isinstance(args[0],Vector) and isinstance(args[1],Vector):
                from_,to=args
                self.x=to.x-from_.x
                self.y=to.y-from_.y
            elif type(args[0]) in [IntType, FloatType] and \
                 type(args[1]) in [IntType, FloatType]:
                self.x=float(args[0])
                self.y=float(args[1])
            else:
                assert 0
        elif len(args)==0:
            self.x=0.0
            self.y=0.0
        elif len(args)==1:
            assert isinstance(args[0],Vector)
            self.x=args[0].x
            self.y=args[0].y
        else:
            assert 0

    def setLength(self,value):
        self.normalize()
        self.scale(value)
        #assert self.length()==value, "%f %f"%(self.length(),value)

    def scale(self, f):
        self.x*=f
        self.y*=f

    def length(self):
        return math.sqrt(self.x**2 + self.y**2)

    def normalize(self):
        self.scale(1.0/self.length())
        #assert self.length()==1.0,"%f %s"%(self.length(),self)

    def add(self,other):
        self.x+=other.x
        self.y+=other.y

    def __repr__(self):
        return "{%f,%f}"%(self.x,self.y)
        
class Object:
    def __init__(self, mass, pos, vel, fix):
        self.mass=mass
        self.pos=pos
        self.vel=vel
        self.fix=fix

    def __repr__(self):
        return "[mass=%f, pos=%s, vel=%s]"%(self.mass,self.pos,self.vel)
        
class World:
    def __init__(self):
        self.objects=[]
        self.steps=0

    def add(self,object):
        self.objects.append(object)
        
    def step(self):
        """
Updates the world.
        """
        for obj in self.objects:
            total_force=Vector()
            for other in self.objects:
                if obj==other: continue
                dist=Vector(obj.pos,other.pos).length()
                if dist < CRASH_DIST:
                    raise Crashed()
                force=Vector(obj.pos,other.pos)
                force.setLength((obj.mass * other.mass) /
                       dist**2)
                total_force.add(force)
                
            acceleration=Vector(total_force)
            acceleration.scale(1.0/obj.mass)
            obj.vel.add(acceleration)
        for obj in self.objects:
            if not obj.fix:
                tmp=Vector(obj.vel)
                tmp.scale(dt)
                obj.pos.add(tmp)
        self.steps+=1

    def dump(self):
        if 0:
            print self.objects[1]
        else:
            for obj in self.objects:
                print "  ",obj
            print

def get_presets():
    """
Returns a list of descriptions of the available presets.
    """
    return [
        "3 fix, 1 moving",
        "oscillating 1",
        "oscillating 2",
        "different masses",
        "random",
        "initial velocity 1",
        "initial velocity 2",
        "2 fix, 2 moving",
        "4 fix, 1 moving",
        "0 fix, 2 moving, oscillating",
        "one standing",
        "sun & earth",
        "earth & moon",
        "two planets"
        ]

def init_world(world,preset):
    """
Populate the world with objects according to the preset.
    """
    if preset==0:
        world.add(Object(1.0,Vector(-1,1),Vector(),1))
        world.add(Object(2.0,Vector(+1,1),Vector(),1))
        world.add(Object(1.0,Vector(-3,-1),Vector(),1))
        world.add(Object(1.0,Vector(0,-1),Vector(),0))
    elif preset==1:
        world.add(Object(1.0,Vector(-1,1),Vector(),1))
        world.add(Object(1.0,Vector(+1,1),Vector(),1))
        world.add(Object(2.0,Vector(0,-1),Vector(),0))
    elif preset==2:
        world.add(Object(1.0,Vector(-1,3),Vector(),1))
        world.add(Object(1.0,Vector(+1,3),Vector(),1))
        world.add(Object(1.0,Vector(-1,-3),Vector(),1))
        world.add(Object(1.0,Vector(+1,-3),Vector(),1))

        world.add(Object(2.0,Vector(0,-1),Vector(0,20),0))
    elif preset==3:
        world.add(Object(1.0,Vector(-1,1),Vector(),1))
        world.add(Object(2.0,Vector(+1,1),Vector(),1))
        world.add(Object(1.0,Vector(0,-1),Vector(),0))
    elif preset==4:
        world.add(Object(1.0,Vector(0,-3.5),Vector(),0))
        for i in range(3):
            world.add(Object(1.0,Vector(random.random()*7-3.5,
                                        random.random()*7-3.5),Vector(),1))
    elif preset==5:
        world.add(Object(1.0,Vector(0,-3.5),Vector(10,10),0))
        world.add(Object(1.0,Vector(-0.5,0.5),Vector(),1))
        world.add(Object(1.0,Vector(-2.8,1.8),Vector(),1))
        world.add(Object(1.0,Vector(-3,2),Vector(),1))
    elif preset==6:
        world.add(Object(1.0,Vector(0,-3.5),Vector(10,0),0))
        world.add(Object(1.0,Vector(-2,1),Vector(),1))        
        world.add(Object(1.0,Vector(2,1),Vector(),1))
    elif preset==7:
        world.add(Object(1.0,Vector(-1,1),Vector(),1))
        world.add(Object(1.0,Vector(+1,1),Vector(),1))
        world.add(Object(2.0,Vector(0,-1),Vector(-10,0),0))
        world.add(Object(1.0,Vector(0,2),Vector(10,0),0))
    elif preset==8:
        world.add(Object(0.1,Vector(+1,+1),Vector(),1))
        world.add(Object(1.0,Vector(+1,-1),Vector(),1))
        world.add(Object(1.0,Vector(-1,+1),Vector(),1))
        world.add(Object(1.0,Vector(-1,-1),Vector(),1))
        
        world.add(Object(0.5,Vector(0,-2),Vector(30,-30),0))
    elif preset==9:
        world.add(Object(1.0,Vector(-1,+1),Vector(10,0),0))
        world.add(Object(1.0,Vector(+1,-1),Vector(-10,0),0))
    elif preset==10:
        world.add(Object(1.0,Vector(-1,+1),Vector(10,1),0))
        world.add(Object(1.0,Vector(+1,-1),Vector(0,0),0))
    elif preset==11:
        world.add(Object(1.0,Vector(+0,+1),Vector(40,0),0))
        world.add(Object(1.0,Vector(+0,+0),Vector(0,0),1))
    elif preset==12:
        world.add(Object(1.0,Vector(+0,+1),Vector(32,0),0))
        world.add(Object(1.0,Vector(+0,+0),Vector(0,0),0))
    elif preset==13:
        world.add(Object(0.00001,Vector(+0,+1),Vector(60,0),0))
        world.add(Object(0.00001,Vector(+0,+3),Vector(30,0),0))
        world.add(Object(3.0,Vector(+0,+0),Vector(0,0),0))

def test():
    w=World()
    a=Object(0.1,Vector(-1,1),Vector(),1)
    b=Object(0.1,Vector(+1,1),Vector(),1)
    c=Object(0.1,Vector(0,-1),Vector(),1)
    w.add(a)
    w.add(b)
    w.add(c)
    for i in range(200):
        w.dump()
        w.step()
        
if __name__=="__main__":
    test()

