import math

WIDTH, HEIGHT = 1920*2//3, 1080*2//3

fps = 60
c = 400.0
GM = 50*c**2
n_balls = 20000//1
too_close_factor = 1.00001
line_width = 100.0/n_balls #1.5
steps_per_frame = 1
refresh_dt = 0.05

save_frames = True
filepath = 'frames\\'

hide_circles = True
hide_absorbed = True
absorbed_color = color(25,0,55,1000000/n_balls)
light_color = color(255,255,0,1000000/n_balls) #color(255,230,0,10)
fill_color = color(0) #color(0) or light_color
background_color = color(100)

p0 = (WIDTH/2.0, HEIGHT/2.0)
x_sep = 0.0

rs = 2*GM/c**2

def norm(v):
    return sqrt(v[0]**2 + v[1]**2)
def add(v1, v2):
    return (v1[0]+v2[0], v1[1]+v2[1])
def subtract(v1, v2):
    return (v1[0]-v2[0], v1[1]-v2[1])
def distance(v1, v2):
    return norm(subtract(v1, v2))
def inner(v1, v2):
    return v1[0]*v2[0] + v1[1]*v2[1]
def product(v, a):
    return (a*v[0], a*v[1])
def proj(v1, v2):
    return product(v2, inner(v1,v2)/(norm(v2)**2))
def unitary(v):
    return product(v, 1.0/norm(v))
def angle(v1, v2):
    return math.atan2(v1[0]*v2[1]-v1[1]*v2[0],
                  v1[0]*v2[0]+v1[1]*v2[1])
def rotate(v, theta):
    s, c = math.sin(theta), math.cos(theta)
    x, y = v[0], v[1]
    return (c*x-s*y, s*x+c*y)
    
def sign(x):
    if x<0:
        return -1
    if x>0:
        return 1
    return 0

def vround(v):
    return (int(round(v[0])), int(round(v[1])))

def wrapToPi(x):
    if x > math.pi:
        return x-2*math.pi
    if x < -math.pi:
        return x+2*math.pi
    return x

class Ball:
    def __init__(self, p, v, radius):
        self.p = p
        self.v = v
        self.phi_step = 0.001
        self.vu = self.get_vu(0.001)
        
        self.radius = max(radius, 0.0)
        
        self.path = [self.p, self.p]

        self.fill_color = fill_color
        self.line_color = light_color

        self.inactive = False
        self.absorbed = False
        
        self.ellapsed_time = 0

    def draw_circle(self):
        if hide_circles or self.inactive:
            return
        fill(self.fill_color)
        ellipse(self.p[0], self.p[1], 2*self.radius, 2*self.radius)
        
    def get_phi(self, p):
        return angle((1.0,0), subtract(p, p0))
        
    def get_vu(self, dt):
        dp = product(self.v, dt)
        p_next = add(self.p, dp)
        
        du = 1/distance(p_next, p0) - 1/distance(self.p, p0)
        dphi = self.get_phi(p_next) - self.get_phi(self.p)
        return du/dphi
        
    def update(self, dt):
        if self.inactive:
            return
        
        self.ellapsed_time += dt
        
        direction = sign(angle(subtract(self.p, p0), self.v))
        
        dphi = direction*self.phi_step
        u = 1/distance(self.p, p0)
        phi = self.get_phi(self.p)
        u_next = u + self.vu*dphi
        phi_next = phi + dphi
        p_next = (cos(phi_next)/u_next, sin(phi_next)/u_next)
        p_next = add(p_next, p0)
        
        self.v = subtract(p_next, self.p)
        dp = self.get_dp(dt)
        p_prev = self.p
        self.p = add(self.p, dp)
        self.v = product(dp, 1.0/dt)
        
        dphi = wrapToPi(self.get_phi(self.p) - self.get_phi(p_prev))
        dvu = (1.5*rs*u**2-u)*dphi
        self.vu += dvu

        if (self.ellapsed_time + dt)//refresh_dt > self.ellapsed_time//refresh_dt:
            self.path.append(self.p)
        else:
            self.path[-1] = self.p
        
    def get_dp(self, dt):
        # v only gives the local orbit direction
        r_vec = subtract(self.p, p0)
        x, y = r_vec
        r = norm(r_vec)
        phi = self.get_phi(self.p)
        
        k = self.v[1]/self.v[0]
        b = 1-rs/r
        
        #https://bit.ly/2WaCIcB
        k1 = sqrt(  -1/( 2*(b-1)*k*x*y - (b*x**2+y**2)*k**2 - b*y**2 - x**2 )  )
        
        dx = sign(self.v[0])*abs(b*c*dt*r*k1)
        dy = sign(self.v[1])*abs(k*dx)
        
        return (dx, dy)
        
    def is_too_close(self):
        if norm(subtract(self.p, p0)) < too_close_factor*(2*GM/c**2):
            return True
        
    def set_too_close(self):
        self.inactive = True
        self.absorbed = True
        self.line_color = absorbed_color

    def is_out_of_bounds(self):
        dx = 4*abs(self.v[0])*refresh_dt
        dy = 4*abs(self.v[1])*refresh_dt
        if (self.p[0]<-dx       and self.v[0]<-1) or \
           (self.p[0]>WIDTH+dx  and self.v[0]> 1) or \
           (self.p[1]<-dy       and self.v[1]<-1) or \
           (self.p[1]>HEIGHT+dy and self.v[1]> 1):
            return True
        return False
    
    def set_out_of_bounds(self):
        self.inactive = True

balls = []
M = 2.0
radius = M*HEIGHT/(2.0*n_balls)
for ky in range(n_balls/2):
    x = WIDTH + radius + x_sep
    y = radius + ky*2*radius
    balls += [
        Ball((x, HEIGHT/2.0+y), (-c, 0), radius),
        Ball((x, HEIGHT/2.0-y), (-c, 0), radius)
    ]
ordered_balls = list(range(n_balls))
for k in range(n_balls/2):
    ordered_balls[n_balls/2+k] = balls[2*k]
    ordered_balls[n_balls/2-k-1] = balls[2*k+1]
#balls.append(Ball((0, 1.51*2*GM/c**2), (-c, 0), radius))
#balls.append(Ball((0, 1.49*2*GM/c**2), (-c, 0), radius))

def draw_wave_front(inactive_balls):
    active_ratio = (n_balls-inactive_balls)/(1.0*n_balls)
    cuttoff = 0.1
    if active_ratio < cuttoff:
        return
    
    stroke(color(255,255,0,255.0*(active_ratio-cuttoff)))
    strokeWeight(1.5)
    noFill()
    
    beginShape()
    started = False
    for k in range(n_balls/2+1):
        ball = ordered_balls[n_balls/2-k]
        p = ball.p
        if ball.absorbed:
            continue
        if distance(p, p0) < rs + 0.5:
            continue
        if not started:
            started = True
            x, y = p
            curveVertex(x, y)
            curveVertex(x, y)
            continue
        if distance((x,y), p) < 10:
            continue
        x, y = p
        curveVertex(x, y)
    x, y = p
    curveVertex(x, y)
    curveVertex(x, y)
    endShape()
    
    beginShape()
    started = False
    for k in range(1,n_balls/2+1):
        ball = ordered_balls[k-n_balls/2-1]
        p = ball.p
        if ball.absorbed:
            continue
        if distance(p, p0) < rs + 0.5:
            continue
        if not started:
            started = True
            x, y = p
            curveVertex(x, y)
            curveVertex(x, y)
            continue
        if distance((x,y), p) < 10:
            continue
        x, y = p
        curveVertex(x, y)
    x, y = p
    curveVertex(x, y)
    curveVertex(x, y)
    endShape()
        
    stroke(0)
    strokeWeight(1)

def draw_solid_path():
    def draw_solid_path_for(ball1, ball2):
        if ball1.absorbed or ball2.absorbed:
            return
        
        def get_xy(ball, k2):
            if k2 < len(ball.path):
                return ball.path[k2]
            return ball.path[-1]
        
        d1 = distance(ball1.p, p0)
        d2 = distance(ball2.p, p0)
        closeness = 1.0
        cuttoff = 1.5
        if d1 < cuttoff*rs or d2 < cuttoff*rs:
            x = min(d1, d2)/rs
            closeness = sqrt((x-1.0)/(cuttoff-1.0))

        noStroke()
        for k2 in range(min(len(ball1.path), len(ball2.path)) -1):
            x1, y1 = get_xy(ball1, k2)
            x2, y2 = get_xy(ball1, k2+1)
            x3, y3 = get_xy(ball2, k2+1)
            x4, y4 = get_xy(ball2, k2)
            d = distance((x2, y2), (x3, y3))
            a = closeness*(255.0*ball1.radius)/d
            a = max(0,min(a,255))
            fill(color(255,255,0,a))
            quad(x1, y1, x2, y2, x3, y3, x4, y4)
        stroke(0)

    for k in range(n_balls/2):
        ball1 = ordered_balls[n_balls/2-k]
        ball2 = ordered_balls[n_balls/2-k-1]
        draw_solid_path_for(ball1, ball2)
    for k in range(1,n_balls/2):
        ball1 = ordered_balls[k-n_balls/2-1]
        ball2 = ordered_balls[k-n_balls/2]
        draw_solid_path_for(ball1, ball2)

img_shadow = None
img_photon_sphere = None
img_event_horizon = None
def setup():
    global img_shadow, img_photon_sphere, img_event_horizon
    size(WIDTH, HEIGHT)
    frameRate(fps)
    smooth(8)
    img_shadow = loadImage('Shadow.png')
    img_photon_sphere = loadImage('Photon Sphere.png')
    img_event_horizon = loadImage('Event Horizon.png')
    s = 0.3
    img_shadow.resize(int(round(s*img_shadow.width)), 0)
    img_photon_sphere.resize(int(round(s*img_photon_sphere.width)), 0)
    img_event_horizon.resize(int(round(s*img_event_horizon.width)), 0)

frame = 1
post_frame = 0
finished = False
def draw():
    global frame, post_frame, finished
    if finished:
        return
    background(background_color)
    
    dt = 1.0/fps

    inactive_balls = 0
    for ball in balls:
        if ball.inactive:
            inactive_balls += 1
            continue
        if ball.is_out_of_bounds():
            ball.set_out_of_bounds()
            continue
        if ball.is_too_close():
            ball.set_too_close()
            continue

    if balls[0].p[0] > WIDTH + balls[0].radius:
        print balls[0].p[0] - WIDTH
        fill(color(0, 0, 0))
        ellipse(p0[0], p0[1], 2*rs, 2*rs)
        
        for ball in balls:
            for k in range(steps_per_frame):
                ball.update(dt/steps_per_frame)
        return

    for ball in balls:
        ball.draw_circle()
        
    draw_wave_front(inactive_balls)
    draw_solid_path()
      
    fill(color(0, 0, 0))
    ellipse(p0[0], p0[1], 2*rs, 2*rs)
    
    if inactive_balls > 0.96*n_balls:
        post_frame += 1
    if True:
        a = 255  # post_frame*255.0/30
        
        M = sqrt(27)/2
        noFill()
        strokeWeight(2.0)
        stroke(color(0, 0, 100, a))
        line(p0[0],p0[1]-M*rs, p0[0]+WIDTH-p0[0], p0[1]-M*rs)
        line(p0[0],p0[1]-M*rs+2*M*rs, p0[0]+WIDTH-p0[0], p0[1]-M*rs+2*M*rs)
        ellipse(p0[0], p0[1], M*2*rs, M*2*rs)
        stroke(color(255,255,0,a))
        ellipse(p0[0], p0[1], 1.5*2*rs, 1.5*2*rs)
        strokeWeight(1.0)
        
        font = createFont('Georgia', 32)
        textFont(font)
        
        fill(color(0, 0, 100, a))
        s = 'Shadow:'
        text(s, p0[0], p0[1]-M*rs-2-22)
        tint(255, a)
        image(img_shadow, p0[0]+textWidth(s)+10, p0[1]-M*rs-2-74)
        
        fill(color(255,255,0,a))
        s = 'Photon Sphere:'
        buff = 0.1
        text(s, p0[0], p0[1]-(1.5 + buff)*rs-2-5)
        tint(255, a)
        image(img_photon_sphere, p0[0]+textWidth(s)+10, p0[1]-(1.5 + buff)*rs-2-26)
        
        fill(color(0,0,0,a))
        s = 'Event Horizon:'
        text(s, p0[0], p0[1]-rs-2)
        tint(255, a)
        image(img_event_horizon, p0[0]+textWidth(s)+10, p0[1]-rs-2-67+22)
        
        fill(color(255,255,255,a))
        stroke(color(255,255,255,a))
        text('Singularity', p0[0],p0[1]-5)
        ellipse(p0[0], p0[1], 3, 3)
        
        stroke(0)

    for ball in balls:
        for k in range(steps_per_frame):
            ball.update(dt/steps_per_frame)
    
    if save_frames and post_frame<180:
        saveFrame(filepath + str(frame).zfill(4) + '.png')
    if post_frame<180:
        print frame
    if post_frame==180:
        print 'Done!'
        finished = True
    
    frame += 1