2d Collisions

import asyncio
import random
import numpy as np
import pygame as pg
SimulationUpdates, SimElapesdTime, MaxSimulationSteps = 0, 0, 0
OLDCollidingPairs = []
pg.init()
clock = pg.time.Clock()
WIDTH, HEIGHT = 1400 , 800
screen = pg.display.set_mode((WIDTH, HEIGHT))

def Text(text,x,y,size):
  font = pg.font.Font(None, size)
  text = font.render(text, True, (255, 255, 255))
  screen.blit(text, (x, y))


def circles_overlap(x1,y1,size1,x2,y2,size2):
  dist = (x1-x2) * (x1-x2) + (y1-y2) * (y1-y2)
  if dist <= (size1+size2) * (size1+size2):
    return True
  else:
    return False

class LineSegment:
  instances=[]
  def __init__(self,sx,sy,ex,ey,rad):
    self.sx = sx
    self.sy = sy
    self.ex = ex
    self.ey = ey
    self.radius = rad



class Ball:
  instances=[]
  def __init__(self, x, y, size):
    self.x = x
    self.y = y
    self.vx = 0
    self.vy = 0
    self.ax = 0
    self.ay = 0
    self.size = size
    col=random.randint(100,200)
    self.colour = (20,col,col)
    self.mass = size * 10
    self.SimTimeRemaining = 0
    self.oldx = 0
    self.oldy = 0
    col=random.randint(1,3)
    if col==1:
      self.colour=(255,0,0)
    elif col==2:
      self.colour=(0,255,0)
    else:
      self.colour=(0,0,255)
  def Collide(self):
    collided=False
    for ball in Ball.instances:
      if ball != self:
        if circles_overlap(self.x,self.y,self.size, ball.x,ball.y,ball.size):
          dist = ((self.x-ball.x)*(self.x-ball.x) + (self.y-ball.y)*(self.y-ball.y)) ** 0.5
          overlap = 0.5 * (dist - self.size - ball.size)

          #displace self
          self.x -= overlap * (self.x - ball.x) / dist
          self.y -= overlap * (self.y - ball.y) / dist

          #displace target
          ball.x += overlap * (self.x - ball.x) / dist
          ball.y += overlap * (self.y - ball.y) / dist
          
          return Ball.instances.index(self), Ball.instances.index(ball)
          collided=True

    
    #if collided:#time displacement
    try:
      IntendedSpeed = (self.vx*self.vx + self.vy*self.vy)**0.5
      IntendedDistance = IntendedSpeed * self.SimTimeRemaining
      ActualDistance = ((self.x-self.oldx)*(self.x-self.oldx) + (self.y-self.oldy)*(self.y-self.oldy))**0.5
      ActualTime = ActualDistance / IntendedSpeed

      self.SimTimeRemaining = self.SimTimeRemaining - ActualTime
    except:
      pass

  def Move(self):
    if self.SimTimeRemaining > 0:

      self.oldx = self.x
      self.oldy = self.y

      #drag to simulate rolling friction
      self.ax = -self.vx * 0.08
      self.ay = -self.vy * 0.08 + gravity #gravity

      #update ball physics
      self.vx += self.ax * self.SimTimeRemaining / TargetFPS
      self.vy += self.ay * self.SimTimeRemaining / TargetFPS
      self.x  += self.vx * self.SimTimeRemaining / TargetFPS
      self.y  += self.vy * self.SimTimeRemaining / TargetFPS

      #move to other side of screen
      if self.x <0:
        self.x += WIDTH
      if self.x >= WIDTH:
        self.x -= WIDTH
      if self.y <0:
        self.y += HEIGHT
      if self.y >= HEIGHT:
        self.y -= HEIGHT

      if self.vx*self.vx + self.vy*self.vy < 0.0001:
        self.vx = 0
        self.vy = 0
      return self.Collide()

def MoveAll():
  global OLDCollidingPairs
  CollNoise = False
  BopNoise = False
  for i in range(SimulationUpdates):
    
    for ball in Ball.instances:
      ball.SimTimeRemaining = SimElapesdTime

    for j in range(MaxSimulationSteps):
      CollidingPairs=[]
      FakeBalls=[]
#static collisions
      for ball in Ball.instances:
        isfake = False
        for i in range(len(FakeBalls)):
          if Ball.instances.index(FakeBalls[i]) == Ball.instances.index(ball):
            isfake = True
        if not isfake:
          

          coll = ball.Move()
          if coll!=None:
            CollidingPairs.append(coll)

          for edge in LineSegment.instances:
            LineX1 = edge.ex - edge.sx
            LineY1 = edge.ey - edge.sy

            LineX2 = ball.x - edge.sx
            LineY2 = ball.y - edge.sy

            EdgeLength = LineX1*LineX1 + LineY1*LineY1

            t = max(0, min(EdgeLength, (LineX1*LineX2 + LineY1*LineY2))) / EdgeLength

            ClosestPointX = edge.sx + t * LineX1
            ClosestPointY = edge.sy + t * LineY1

            Distance = ((ball.x- ClosestPointX)*(ball.x- ClosestPointX) + (ball.y - ClosestPointY)*(ball.y - ClosestPointY))**0.5

            if Distance <= ball.size + edge.radius:
              #static collision occoured

              Ball.instances.append(Ball(ClosestPointX,ClosestPointY,edge.radius))
              fakeball = Ball.instances[-1]

              fakeball.mass = ball.mass * 1
              fakeball.vx = -ball.vx
              fakeball.vy = -ball.vy

              #fakeball.mass = ball.mass * 1
              #fakeball.vx = -ball.vx
              #fakeball.vy = -ball.vy

              FakeBalls.append(fakeball)
              coll = Ball.instances.index(ball), Ball.instances.index(fakeball)
              CollidingPairs.append(coll)

              overlap = 1 * (Distance - ball.size - fakeball.size)

              #displace target
              ball.x -= overlap * (ball.x - fakeball.x) / Distance
              ball.y -= overlap * (ball.y - fakeball.y) / Distance


#dynamic collisions
      for i in range(len(CollidingPairs)):
        b1 = Ball.instances[CollidingPairs[i][0]]
        b2 = Ball.instances[CollidingPairs[i][1]]

        #distance
        dist = ((b1.x-b2.x)*(b1.x-b2.x) + (b1.y-b2.y)*(b1.y-b2.y)) ** 0.5

        #normal
        nx = (b2.x - b1.x) / dist
        ny = (b2.y - b1.y) / dist

        #tangent
        tx = -ny
        ty = nx

        #dot product tangent
        dpTan1 = b1.vx * tx + b1.vy * ty
        dpTan2 = b2.vx * tx + b2.vy * ty

        #dot product normal
        dpNorm1 = b1.vx * nx + b1.vy * ny
        dpNorm2 = b2.vx * nx + b2.vy * ny

        #conservation of momentum in 1D
        m1 = (dpNorm1 * (b1.mass - b2.mass) + 2 * b2.mass * dpNorm2) / (b1.mass + b2.mass)
        m2 = (dpNorm2 * (b2.mass - b1.mass) + 2 * b1.mass * dpNorm1) / (b1.mass + b2.mass)

        #update velocitys
        b1.vx = tx * dpTan1 + nx * m1
        b1.vy = ty * dpTan1 + ny * m1

        b2.vx = tx * dpTan2 + nx * m2
        b2.vy = ty * dpTan2 + ny * m2

        IntendedSpeed1 = (b1.vx*b1.vx + b1.vy*b1.vy)**0.5
        IntendedSpeed2 = (b2.vx*b2.vx + b2.vy*b2.vy)**0.5

      for fakeball in FakeBalls:
        Ball.instances.remove(fakeball)
      CollidingPairs=[]
      FakeBalls=[]


def DrawAll(mousemoving2,mouse_x,mouse_y):
  screen.fill((0,0,0))
  for ball in Ball.instances:
    pg.draw.circle(screen,ball.colour,(ball.x,ball.y),ball.size)

  for line in LineSegment.instances:
    pg.draw.circle(screen,(200,255,200),(line.sx,line.sy),line.radius,5)
    pg.draw.circle(screen,(200,255,200),(line.ex,line.ey),line.radius,5)
    
    nx = -(line.ey - line.sy)
    ny = (line.ex - line.sx)
    d = (nx*nx + ny*ny)**0.5
    nx /= d
    ny /= d
    rad = line.radius-5/2
    pg.draw.line(screen,(200,255,200),(line.sx + nx*rad, line.sy + ny*rad), (line.ex + nx*rad, line.ey + ny*rad),5)
    pg.draw.line(screen,(200,255,200),(line.sx - nx*rad, line.sy - ny*rad), (line.ex - nx*rad, line.ey - ny*rad),5)

  if mousemoving2!=-1:
    pg.draw.line(screen,(0,0,255),(int(Ball.instances[mousemoving2].x), int(Ball.instances[mousemoving2].y)),(mouse_x,mouse_y),5)

async def Main():
  global SimulationUpdates, SimElapesdTime, MaxSimulationSteps
  mousemoving0=-1
  mousemoving2=-1

  while True:
    SimulationUpdates = 4
    MaxSimulationSteps = 10
    MaxSimulationSteps = 5
    SimElapesdTime = clock.get_fps() / SimulationUpdates / MaxSimulationSteps

    for event in pg.event.get():
      if event.type == pg.QUIT:
        exit()
    
    key = pg.key.get_pressed()
    mouse_x, mouse_y = pg.mouse.get_pos()
    mouse_pressed = pg.mouse.get_pressed()

    #left click
    if mouse_pressed[0]:
      if mousemoving0==-1:
        for ball in Ball.instances:
          if circles_overlap(ball.x,ball.y,ball.size,mouse_x,mouse_y,0):
            mousemoving0=Ball.instances.index(ball)
            movingtype="Ball"
        for line in LineSegment.instances:
          if circles_overlap(line.sx,line.sy,line.radius,mouse_x,mouse_y,0):
            mousemoving0=line.instances.index(line)
            movingtype="SLine"
          if circles_overlap(line.ex,line.ey,line.radius,mouse_x,mouse_y,0):
            mousemoving0=line.instances.index(line)
            movingtype="ELine"
    else:
      mousemoving0=-1

    if mousemoving0!=-1:
      if movingtype == "Ball":
        Ball.instances[mousemoving0].x = mouse_x
        Ball.instances[mousemoving0].y = mouse_y
      elif movingtype == "SLine":
        LineSegment.instances[mousemoving0].sx = mouse_x
        LineSegment.instances[mousemoving0].sy = mouse_y
      elif movingtype == "ELine":
        LineSegment.instances[mousemoving0].ex = mouse_x
        LineSegment.instances[mousemoving0].ey = mouse_y

    #right click
    if mouse_pressed[2]:
      if mousemoving2==-1:
        for ball in Ball.instances:
          if circles_overlap(ball.x,ball.y,ball.size,mouse_x,mouse_y,0):
            mousemoving2=Ball.instances.index(ball)
    else:
      if mousemoving2!=-1:
        Ball.instances[mousemoving2].vx = 0.2 * (Ball.instances[mousemoving2].x - mouse_x)
        Ball.instances[mousemoving2].vy = 0.2 * (Ball.instances[mousemoving2].y - mouse_y)
      mousemoving2=-1

    #move everything
    MoveAll()

    #draw everything
    DrawAll(mousemoving2,mouse_x,mouse_y)

    pg.display.set_caption(f"2d Collisions - left click+drag circles to move them - right click+drag to add velocity to ball")
    pg.display.flip()
    clock.tick(TargetFPS)
    await asyncio.sleep(0)
    #for line in LineSegment.instances:
    # print(f"LineSegment({line.sx},{line.sy},{line.ex},{line.ey},20)")


TargetFPS=60
gravity = 1
LineSegment.instances.append(LineSegment(873,706,598,640,20))
LineSegment.instances.append(LineSegment(5,80,332,205,20))
LineSegment.instances.append(LineSegment(989,324,1371,152,20))
LineSegment.instances.append(LineSegment(340,402,609,417,20))
for i in range(50):
  Ball.instances.append(Ball(random.randint(50,WIDTH-50),random.randint(50,HEIGHT-50),random.randint(20,40)))

asyncio.run(Main())