# Tk_elastic_ball_collisions.py

# Tk_elastic_ball_collisions.py
# to fix: will eventually gets wonky overlapped despite the ghosting feature

from math import sqrt, sin, cos
import tkinter as tk
import time

def _create_circle(self, x, y, r, **kwargs):
    return self.create_oval(x-r, y-r, x+r, y+r, **kwargs)
tk.Canvas.create_circle = _create_circle


WINDOW_WIDTH = 200
WINDOW_HEIGHT = 200

# The coefficient of restitution
# Set to 1 for perfectly elastic collisions
e = 1


class Ball:
    def __init__(self, mass:float, r:float, x:float, y:float,
                 vx:float=0, vy:float=0, ghost:int=0, **kwargs):
        """
        This is a class that defines what a ball is and how it interacts
        with other balls.

        Arguments:
            mass:float    # The mass of the ball
                          ------------------------------------------------------
            r:float       # The radius of the ball, must be >0
                          ------------------------------------------------------
            x:float       # The x position of the ball
                          #   must be >0 and <WINDOW_WIDTH
            y:float       # The y position of the ball
                          #   must be >0 and <WINDOW_HEIGHT
                          ------------------------------------------------------
            vx:float      # The x velocity of the ball
            vy:float      # The y velocity of the ball
                          ------------------------------------------------------
            **kwargs      # All of the args to be passed in to `create_circle`
        """
        self.m = mass
        self.r = r
        self.x = x
        self.y = y
        self.vx = vx
        self.vy = vy
        self.ghost = ghost
        self.kwargs = kwargs

    def display(self, canvas:tk.Canvas) -> int:
        """
        Displays the ball on the screen and returns the canvas_id (which
        is a normal python int).
        """
        canvas_id = canvas.create_circle(self.x, self.y, self.r, **self.kwargs)
        return canvas_id

    def move(self, all_balls:list, dt:float) -> (float, float):
        """
        This moves the ball according to `self.vx` and `self.vy`.
        It also checks for collisions with other balls.

        Arguments:
            all_balls:list   # A list of all balls that are in the
                             # simulation. It can include this ball.
                             ---------------------------------------------------
            dt:float         # delta time - used to move the balls

        Returns:
            dx:float         # How much the ball has moved in the x direction
            dy:float         # How much the ball has moved in the y direction

        Note: This function isn't optimised in any way. If you optimise
        it, it should run faster.
        """

        # Check if there are any wall collisions:
        if self.x - self.r < 0:
            self.vx = self.vx * -1
            self.x += 0.01
        elif self.x + self.r > WINDOW_WIDTH:
            self.vx = self.vx * -1
            self.x -= 0.01
        if self.y - self.r < 0:
            self.vy = self.vy * -1
            self.y += 0.01
        elif self.y + self.r > WINDOW_HEIGHT:
            self.vy = self.vy * -1
            self.y -= 0.01

        # Check if there are any ball collisions:
        for other, _ in all_balls:
            # Skip is `ball` is the same as this ball
            if id(other) == id(self):
                continue
            # Check if there is a collision:
            distance_squared = (other.x - self.x)**2 + (other.y - self.y)**2
            if distance_squared <= (other.r + self.r)**2:
                # Now the fun part - calulating the resultant velocity.
                if not self.ghost:
                    self.ghost = 5

                    # First I will find the normal vector of the balls' radii.
                    # That is just the unit vector in the direction of the
                    #  balls' radii
                    ball_radii_vector_x = other.x - self.x
                    ball_radii_vector_y = other.y - self.y
                    abs_ball_radii_vector = sqrt(ball_radii_vector_x**2 +\
                                                 ball_radii_vector_y**2)
                    nx = ball_radii_vector_x / abs_ball_radii_vector        **2*2
                    ny = ball_radii_vector_y / abs_ball_radii_vector        **2*2

                    # Now I will calculate the tangent
                    tx = -ny
                    ty = nx

                    """ Only for debug
                    print("n =", (nx, ny), "\t t =", (tx, ty))
                    print("u1 =", (self.vx, self.vy),
                          "\t u2 =", (other.vx, other.vy))
                    #"""

                    # Now I will split the balls' velocity vectors to the sum
                    #  of 2 vectors parallel to n and t
                    # self_velocity  = λ*n + μ*t
                    # other_velocity = a*n + b*t
                    λ = (self.vx*ty - self.vy*tx) / (nx*ty - ny*tx)
                    # Sometimes `tx` can be 0 if so we are going to use `ty` instead
                    try:
                        μ = (self.vx - λ*nx) / tx
                    except ZeroDivisionError:
                        μ = (self.vy - λ*ny) / ty

                    """ Only for debug
                    print("λ =", λ, "\t μ =", μ)
                    #"""

                    a = (other.vx*ty - other.vy*tx) / (nx*ty - ny*tx)
                    # Sometimes `tx` can be 0 if so we are going to use `ty` instead
                    try:
                        b = (other.vx - a*nx) / tx
                    except ZeroDivisionError:
                        b = (other.vy - a*ny) / ty

                    """ Only for debug
                    print("a =", a, "\t b =", b)
                    #"""

                    self_u = λ
                    other_u = a

                    sum_mass_u = self.m*self_u + other.m*other_u
                    sum_masses = self.m + other.m

                    # Inelastic_collision
                    self_v = (e*other.m*(other_u-self_u) + sum_mass_u)/sum_masses
                    other_v = (e*self.m*(self_u-other_u) + sum_mass_u)/sum_masses

                    self.vx = self_v*nx + μ*tx
                    self.vy = self_v*ny + μ*ty
                    other.vx = other_v*nx + b*tx
                    other.vy = other_v*ny + b*ty

                    print("v1 =", (self.vx, self.vy),
                          "\t v2 =", (other.vx, other.vy))

            else:
                self.ghost = max(0, self.ghost - 1)

        # Move the ball
        dx = self.vx * dt
        dy = self.vy * dt
        self.x += dx
        self.y += dy
        return dx, dy


class Simulator:
    def __init__(self):
        self.balls = [] # Contains tuples of (<Ball>, <canvas_id>)
        self.root = tk.Tk()
        self.root.resizable(False, False)
        self.canvas = tk.Canvas(self.root, width=WINDOW_WIDTH,
                                height=WINDOW_HEIGHT)
        self.canvas.pack()

    def step(self, dt:float) -> None:
        """
        Steps the simulation as id `dt` seconds have passed
        """
        for ball, canvas_id in self.balls:
            dx, dy = ball.move(self.balls, dt)
            self.canvas.move(canvas_id, dx, dy)

    def run(self, dt:float, total_time:float) -> None:
        """
        Note: This function is just for proof of concept which disregards
        the slowing simulation as a whole.
        """
        while 1:
            self.step(dt)
            self.root.update()
            time.sleep(dt)

    def add_ball(self, *args, **kwargs) -> None:
        """
        Adds a ball by passing all of the args and keyword args to `Ball`.
        It also displays the ball and appends it to the list of balls.
        """
        ball = Ball(*args, **kwargs)
        canvas_id = ball.display(self.canvas)
        self.balls.append((ball, canvas_id))


app = Simulator()
app.add_ball(mass=1, r=20, x=20, y=100, vx=240, vy=0, fill="red")
app.add_ball(mass=1, r=20, x=100, y=100, vx=0, vy=0, fill="blue")
app.add_ball(mass=1, r=20, x=180, y=98, vx=0, vy=0, fill="yellow")
app.run(0.01, 10000)