from scene import *
import numpy as np
import random

class MyScene(Scene):
    
    def setup(self):
        self.background_color='#00483a'
        self.width,self.height=self.size
        balls_name=['emj:Black_Circle',
                                'emj:Blue_Circle',
                                'emj:Red_Circle',
                                'emj:White_Circle',
                                'emj:Grinning',
                                'emj:Smiling_1',
                                'emj:Smiling_2',
                                'emj:Smiling_3',
                                'emj:Smiling_4',
                                'emj:Smiling_5',
                                'emj:Smiling_6']
        self.balls=[]
        for ball_name in balls_name:
            r=28
            px=random.randrange(r,self.width-r)
            py=random.randrange(r,self.height-r)
            vx=random.uniform(-10,10)
            vy=random.uniform(-10,10)
            ball=ExSpriteNode(ball_name,r,px,py,vx,vy)
            self.balls.append(ball)
            self.add_child(ball)
    
    def update(self):
        for ball in self.balls:
            ball.move(1.0)
        n=len(self.balls)
        for i1 in range(n-1):
            b1=self.balls[i1]
            for i2 in range(i1+1,n):
                b2=self.balls[i2]
                r=b1.radius+b2.radius
                d=self.get_distance(b1,b2)                    
                if round(d-r,2)<=0:
                    t=self.get_contact_time(b1,b2)
                    b1.position+=b1.velocity*t
                    b2.position+=b2.velocity*t
                    v1=self.get_new_velocity(b1,b2)
                    v2=self.get_new_velocity(b2,b1)
                    b1.velocity+=(-v1)+v2
                    b2.velocity+=v1+(-v2)
                    b1.position+=b1.velocity*(-t)
                    b2.position+=b2.velocity*(-t)
                                        
    def get_new_velocity(self,ball1,ball2):
        a,b=ball1.velocity
        X,Y=ball2.position-ball1.position
        k=(a*X+b*Y)/(X**2+Y**2)
        return np.array([k*X,k*Y])

    def get_distance(self,ball1,ball2):
        v=ball2.position-ball1.position
        r=ball1.radius+ball2.radius
        return np.linalg.norm(v)
        
    def get_contact_time(self,ball1,ball2):
        A,B=ball2.position-ball1.position
        C,D=ball2.velocity-ball1.velocity
        R=ball1.radius+ball2.radius
        a=C**2+D**2
        b=2*(A*C+B*D)
        c=A**2+B**2-R**2
        t=(-b-(b**2-4*a*c)**(1/2))/(2*a)
        return t
                            
class ExSpriteNode(SpriteNode):
        
    def __init__(self,texture,radius,px,py,vx,vy):
        super(ExSpriteNode,self).__init__(texture)
        self.radius=radius
        self.position=np.array([px,py])
        self.velocity=np.array([vx,vy])
        
    def move(self,decay):
        self.position+=self.velocity
        px,py=self.position
        vx,vy=self.velocity
        r=self.radius
        w=self.parent.width-r
        h=self.parent.height-r
        if px<r:
            px=r
            vx=-vx*decay            
        if px>w:
            px=w
            vx=-vx*decay
        if py<r:
            py=r
            vy=-vy*decay
        if py>h:
            py=h
            vy=-vy*decay
        self.velocity=np.array([vx,vy])
        self.position=np.array([px,py])        
            
run(MyScene())