# -*- coding: utf-8 -*-
"""
Created on Thu Mar  8 14:14:53 2018

@author: ip1016
"""
import numpy as np
import math as m
import matplotlib.pyplot as plt
import matplotlib.animation as animation


#r=np.array([0,0])
#v=np.array([0,0])
collision=[]

class Ball:
    def __init__(self,t,m,R,r=[0, 0],v=[0, 0],clr='r'):
        self.__r=np.array(r,dtype='float')
        self.m=m
        self.__R=float(R)
        self.__v=np.array(v,dtype = 'float')
        self.t=t
       # self.__dpos=np.array([r(0),r(1)],[v(0),v(1)])
        self.__patch=plt.Circle(self.__r,self.__R, ec=clr,fill=False,fc=clr)
        if len(self.__r)!=2:
            raise Exception("Vector parameter size")

    def pos(self):
        return self.__r

    def vel(self):
        return self.__v

    def move(self,dt):
        self.__r=self.__r+np.multiply(self.__v,dt)
        self.__patch.center=self.__r

    def rad(self):
        return self.__R
    def ty(self):
        return self.t
    def get_patch(self):
        return self.__patch

    def time_to_collision(self,other):
        a=np.dot(self.__v,self.__v)+np.dot(other.__v,other.__v)-2*np.dot(self.__v,other.__v)
        b=2*np.dot(self.__r,self.__v)-2*np.dot(self.__r,other.__v)-2*np.dot(self.__v,other.__r)+2*np.dot(other.__r,other.__v)
        if self.t==1 and other.t==1:
            c=np.dot(self.__r,self.__r)+np.dot(other.__r,other.__r)-2*np.dot(self.__r,other.__r)-(self.__R+other.__R)*(self.__R+other.__R)

            delta=b*b-4*a*c
            if delta<0:
                return 100000

            if delta==0:
                return 100000
            dt1=(-b+m.sqrt(delta))/(2*a)
            dt2=(-b-m.sqrt(delta))/(2*a)

            if dt1<0 and dt2<0:
                return 100000

            if dt1>0 and dt2>0:
                return min([dt1,dt2])

            if dt1==0 and dt2>0:
                return 0

            if dt2==0 and dt1>0:
                return 0

            if dt1==0 and dt2<0:
                return 100000

            if dt2==0 and dt1<0:
                return 100000



        if (self.t==-1 and other.t==1) or (self.t==1 and other.t==-1):
            c=np.dot(self.__r,self.__r)+np.dot(other.__r,other.__r)-2*np.dot(self.__r,other.__r)-(self.__R-other.__R)*(self.__R-other.__R)
            delta=b*b-4*a*c
            if delta<0:
                return 100000

            if delta==0:
                return 100000
            dt1=(-b+m.sqrt(delta))/(2*a)
            dt2=(-b-m.sqrt(delta))/(2*a)
            if dt1<0 and dt2>0:
                return dt2

            if dt1>0 and dt2<0:
                return dt1

            if dt1==0 or dt2==0:
                return 0





    def collide(self, other):
        other.__v=np.array(other.__v,dtype='float')

        self.__v=-(self.__v)
        other.__v=-(other.__v)
        #other.__v[1]=(np.linalg.norm(other.__v)*m.cos(theta2-phi)(other.m-self.m)+2*self.m*np.linalg.norm(self.__v)*m.cos(theta1-phi))*m.sin(phi)/(self.m+other.m)+np.linalg.norm(other.__v)*m.sin(theta2-phi)*m.cos(phi)
class Gas:
    def __init__(self,list=[Ball(1,1.0,1.0,[2,1], [0.2,0.3],'r'), Ball(1,1.0,1.0,[0,0], [-0.2,0.3],'r'),Ball(1,1.0,1.0,[0,5], [0.2,-0.3],'r'),Ball(1,1.0,1.0,[2,4], [-0.5,0.3],'r'),Ball(-1,500000.0,10.0,[0,0], [0,0],'g')]):
        """
        Initialise the Orbit class for 2 balls and some trivial text
        """
        #self.__ball1 = Ball(0,1.0,1.0,[2,1], [0.2,0.3],'r')
        #self.__ball2=Ball(0,1.0,1.0,[0,0], [-0.2,0.3],'r')
        #self.__ball3=Ball(0,1.0,1.0,[0,5], [0.2,-0.3],'r')
        #self.__ball4=Ball(0,1.0,1.0,[2,3], [-0.5,0.3],'r')
        #self.__container = Ball(1,500000.0,10.0,[0,0], [0,0],'g')
        #self.__list=[self.__ball1,self.__ball2,self.__ball3,self.__ball4,self.__container]
        self.list=[Ball(1,1.0,1.0,[2,1], [0.2,0.3],'r'), Ball(1,1.0,1.0,[0,0], [-0.2,0.3],'r'),Ball(1,1.0,1.0,[0,5], [0.2,-0.3],'r'),Ball(1,1.0,1.0,[2,-7], [-0.5,0.3],'r'),Ball(-1,500000.0,10.0,[0,0], [0,0],'g')]
        self.__text0 = None

    def init_figure(self):
        """
        Initialise the container diagram and add it to
        the plot.
        This method is called once by FuncAnimation with no arguments.
        Returns a list or tuple of the 'patches' to be animated.
        """
        # add the big circle to represent the container
        BigCirc = plt.Circle((0,0), 10, ec = 'b', fill = False, ls = 'solid')
        ax.add_artist(BigCirc)
        # initialise the text to be animated and add it to the plot
        self.__text0 = ax.text(-9.9,9,"f={:4d}".format(0,fontsize=12))
        patches = [self.__text0]

        # add the patches for the balls to the plot
        for b in self.list:
            pch = b.get_patch()
            ax.add_patch(pch)
            patches.append(pch)

        return patches

    def change_mom(self):
        change_mom=self.m*self.__v
        return change_mom




    def next_frame(self,framenumber):
        self.__text0.set_text("f={:4d}".format(framenumber))
        patches = [self.__text0]
        collision=[]
        for b in self.list:
            for i in range(len(self.list)):
                for j in range(i+1,len(self.list)):

                    collision.append(self.list[i].time_to_collision(self.list[j]))
                    #print ("a")
                    #print (abs(np.linalg.norm(self.list[i].pos()-self.list[j].pos())-(self.list[i].ty()*self.list[i].rad()+self.list[j].ty()*self.list[j].rad())))
                #print (self.list[i].time_to_collision(self.list[j]))
            #print(collision)
            # print("c")

            #print("j")


            if min(collision)==0:
                for i in range(len(self.list)):
                    for j in range(i+1,len(self.list)):
                        if abs(np.linalg.norm(self.list[i].pos()-self.list[j].pos())-(self.list[i].ty()*self.list[i].rad()+self.list[j].ty()*self.list[j].rad()))<0.1:
                            print ("d")
                            self.list[i].collide(self.list[j])
            else:
                #print(b.pos())
                b.move(min(collision))
                #print(b.pos())
                patches.append(b.get_patch())
                #print (collision)
            #print(b.vel())
           # print ("b")

        return patches


if __name__ == "__main__":

    fig=plt.figure()
    ax=plt.axes(xlim=(-10,10),ylim=(-10,10))
    ax.axes.set_aspect('equal')

    movie=Gas()

    anim=animation.FuncAnimation(fig,
                                 movie.next_frame,
                                 init_func = movie.init_figure,
                                 frames=1000,
                                 interval=20,
                                 blit=True)
    plt.show()