# -*- coding: utf-8 -*-
"""
    Demo 4: Collision Demo
   
    Simple multi-ball collision demo
"""
#
# IMPORTS
#
# StdLib Imports
import sys
from os.path import dirname, join as pathjoin
import time
from datetime import date
import cProfile, pstats
import random
import math

import pygame
from pygame.color import THECOLORS as RGB

from vec2d import vec2d


#
# Globals / Constants
#
class GameOver(Exception): pass

#
# API
#
def rint(number):
    return int(round(number))

def vec_as_degrees(vec):
    return math.atan2(vec.y, vec.x) * 180 / math.pi

def cartesian_to_screen((screen_w, screen_h), x_or_tuple, y=None):
    if y == None:
        x, y = x_or_tuple[0], x_or_tuple[1]
    else:
        x, y = x_or_tuple, y
    cx = rint(screen_w/2.0 + x)
    cy = rint(screen_h/2.0 - y)
    return cx, cy

class Layer(object):
    def __init__(self, image, pos):
        self.image = image
        self.pos = vec2d(pos)    

class Graphics:
    canvas = None
    layers = []
    width = 0
    height = 0
   
    @staticmethod
    def init(screen_width, screen_height):
        Graphics.width = screen_width
        Graphics.height = screen_height
        pygame.init()
        Graphics.canvas = pygame.display.set_mode((screen_width, screen_height))
   
    @staticmethod
    def add_layer_to_canvas(layer):
        Graphics.layers.append(layer)
   
    @staticmethod
    def render():
        for layer in Graphics.layers:
            screen_pos = cartesian_to_screen((Graphics.width, Graphics.height),
                layer.pos)
            Graphics.canvas.blit(layer.image, screen_pos)
        pygame.display.flip()
        Graphics.clear_canvas()
   
    @staticmethod
    def clear_canvas():
        Graphics.layers = []
       
    @staticmethod
    def draw_circle(circle):
        r = rint(circle.radius)
        dim = r * 2
       
        # draw transparent surface
        base = Graphics.draw_transparent_surface(dim, dim)
       
        # draw circle
        pygame.draw.circle(base, circle.color, (r, r), r, 0)
       
        # adjust pos
        pt = circle.pos + vec2d(-r, r)
       
        return Layer(base, pt)
       
    @staticmethod
    def draw_token(token):
        r = rint(token.radius)
        dim = r * 2
        transparency = 50
       
        image = Graphics.draw_transparent_surface(dim, dim)
        alpha_mask = Graphics.draw_alpha_surface(dim, dim, transparency)
        pygame.draw.circle(alpha_mask, token.color, (r,r), r)

        image = Graphics.draw_image_on_surface(image, alpha_mask, (0,0))
        pt = token.pos + vec2d(-r, r)
       
        return Layer(image, pt)
       
    @staticmethod
    def draw_transparent_surface(width, height):
        flags = pygame.SRCALPHA
        depth = 32
        return pygame.Surface((width, height), flags, depth)
   
    @staticmethod
    def draw_alpha_surface(width, height, transparency):
        surface = pygame.Surface((width, height))
        surface.set_colorkey(RGB['black'])
        surface.set_alpha(rint(255 * (100-transparency) / 100.0))
        return surface
   
    @staticmethod
    def draw_image_on_surface(image, surface, image_pos=(0,0)):
        left = rint(surface.get_width() / 2.0)
        top = rint(surface.get_height() / 2.0)
        x = left + image_pos[0] - rint(image.get_width() / 2.0)
        y = top - image_pos[1] - rint(image.get_height() / 2.0)
       
        surface.blit(image, (x,y))
        return surface
       
       
class Physics:
   
    #
    # Math
    #
    @staticmethod
    def distance_squared((x1,y1), (x2,y2)):
        return (y2-y1) * (y2-y1) + (x2-x1) * (x2-x1)
   
    @staticmethod
    def min_dist_point_to_segment(point, segment):
        """Segment is a tuple of points. Return minimum distance between point
        and segment."""
        endpt1, endpt2 = vec2d(segment[0]), vec2d(segment[1])
        point = vec2d(point)
        vpt = endpt1 - point
        vseg = endpt2 - endpt1
        t = (vpt).dot(vseg) / vseg.get_length_sqrd()
        return endpt1 + t * (endpt1 - endpt2)
   
    #
    # Circle vs circle
    #
    @staticmethod
    def circles_are_colliding(c1, c2):
        return Physics.circles_overlap((c1.pos.x, c1.pos.y, c1.radius),
            (c2.pos.x, c2.pos.y, c2.radius))
   
    @staticmethod
    def circles_overlap((c1_x, c1_y, c1_r), (c2_x, c2_y, c2_r)):
        """Tests whether circles overlap. Arguments are (x, y, radius) for each
        circle. Returns True/False. Optimized per suggestion
        here: http://stackoverflow.com/questions/780169"""
        dist_sq = Physics.distance_squared((c1_x, c1_y), (c2_x, c2_y))
        len_sq = (c1_r + c2_r) * (c1_r + c2_r)
        return dist_sq <= len_sq
   
    @staticmethod
    def circle_segment_overlap((cx, cy, cr), (pt1, pt2)):
        """Tests whether circle overlaps line segment."""
        min_vec = Physics.min_dist_point_to_segment((cx, cy), (pt1, pt2))
        return min_vec.get_dist_sqrd((cx, cy)) <= cr * cr
   
    @staticmethod
    def collide_circle_with_circle(
        (c1_x, c1_y, c1_vx, c1_vy, c1_r, c1_m, c1_e),
        (c2_x, c2_y, c2_vx, c2_vy, c2_r, c2_m, c2_e)):
        """Resolve collision between 2 circles with given parameters. Returns a
        tuple of tuples representing new positions and velocities for each
        circle."""
        # translate colliding circles
        pos1, pos2 = Physics.translate_colliding_circles(
            (c1_x, c1_y, c1_r, c1_m),
            (c2_x, c2_y, c2_r, c2_m))
       
        # reflect colliding circles
        vel1, vel2 = Physics.reflect_colliding_circles(
            (c1_x, c1_y, c1_vx, c1_vy, c1_r, c1_m, c1_e),
            (c2_x, c2_y, c2_vx, c2_vy, c2_r, c2_m, c2_e))
       
        # return collisions for circle1, circle2
        return (pos1, vel1), (pos2, vel2)
   
    @staticmethod
    def translate_colliding_circles((c1_x, c1_y, c1_r, c1_m),
        (c2_x, c2_y, c2_r, c2_m)):
        """returns new circle positions as vectors"""
        # inverse masses
        im1 = 1.0 / c1_m
        im2 = 1.0 / c2_m
       
        # get the minimum translation distance to push overlapping circles apart
        mtd = Physics.mtd((c1_x, c1_y, c1_r), (c2_x, c2_y, c2_r))
       
        # push-pull apart based on mass
        new_c1_pos = vec2d(c1_x, c1_y) + (mtd * (im1 / (im1 + im2)))
        new_c2_pos = vec2d(c2_x, c2_y) - (mtd * (im2 / (im1 + im2)))
       
        return new_c1_pos, new_c2_pos
   
    @staticmethod
    def reflect_colliding_circles(
        (c1_x, c1_y, c1_vx, c1_vy, c1_r, c1_m, c1_e),
        (c2_x, c2_y, c2_vx, c2_vy, c2_r, c2_m, c2_e)):
        # inverse masses, mtd, restitution
        im1 = 1.0 / c1_m
        im2 = 1.0 / c2_m
        mtd = Physics.mtd((c1_x, c1_y, c1_r), (c2_x, c2_y, c2_r))
        restitution = c1_e * c2_e
       
        # impact speed
        v = vec2d(c1_vx, c1_vy) - vec2d(c2_vx, c2_vy)
        vn = v.dot(mtd.normalized())
       
        # circle moving away from each other already -- return
        # original velocities
        if vn > 0.0:
            return vec2d(c1_vx, c1_vy), vec2d(c2_vx, c2_vy)
       
        # collision impulse
        i = (-1.0 * (1.0 + restitution) * vn) / (im1 + im2)
        impulse = mtd.normalized() * i
       
        # change in momentun
        new_c1_v = vec2d(c1_vx, c1_vy) + (impulse * im1)
        new_c2_v = vec2d(c2_vx, c2_vy) - (impulse * im2)
       
        return new_c1_v, new_c2_v
   
    @staticmethod
    def mtd((c1_x, c1_y, c1_r), (c2_x, c2_y, c2_r)):
        """source: http://stackoverflow.com/q/345838/1093087"""
        delta = vec2d(c1_x, c1_y) - vec2d(c2_x, c2_y)
        d = delta.length
        mtd = delta * (c1_r + c2_r - d) / d
        return mtd
   
    #
    # Circle in circle
    #
    @staticmethod
    def circle_contains_circle(outer_x, outer_y, outer_r, inner_x, inner_y,
        inner_r):
        distension = Physics.circle_distends_circle(outer_x, outer_y, outer_r,
            inner_x, inner_y, inner_r)
        return distension <= 0
   
    @staticmethod
    def circle_distends_circle(outer_x, outer_y, outer_r, inner_x, inner_y,
        inner_r):
        return (vec2d(outer_x, outer_y).get_distance(
            vec2d(inner_x, inner_y)) + inner_r) - outer_r
   
    @staticmethod
    def resolve_circle_in_circle_collision(
        (inner_x, inner_y, inner_r, inner_vx, inner_vy, inner_e),
        (outer_x, outer_y, outer_r, outer_e)):
        """Resolve collision between a circle crossing the perimeter of a larger
        containing circle. c1 is the inner circle, c2 the outer containing."""
        new_inner_circle_position = Physics.translate_circle_distending_circle(
            (inner_x, inner_y, inner_r, inner_vx, inner_vy),
            (outer_x, outer_y, outer_r))
   
        (impact_x, impact_y) = new_inner_circle_position
   
        new_inner_circle_velocity = Physics.reflect_circle_in_circle(
            (impact_x, impact_y, inner_vx, inner_vy, inner_e),
            (outer_x, outer_y, outer_e))
   
        return new_inner_circle_position, new_inner_circle_velocity
   
    @staticmethod
    def translate_circle_distending_circle(
        (inner_x, inner_y, inner_r, inner_vx, inner_vy),
        (outer_x, outer_y, outer_r)):
        """returns new inner circle position as vector"""        
        if Physics.circle_contains_circle(outer_x, outer_y, outer_r,
            inner_x, inner_y, inner_r):
            return vec2d(inner_x, inner_y)
       
        # vectors
        inner_pos = vec2d(inner_x, inner_y)
        inner_v = vec2d(inner_vx, inner_vy)
   
        # get distension vector
        inner_pt = vec2d(inner_x, inner_y)
        distension = Physics.circle_distends_circle(outer_x, outer_y, outer_r,
            inner_x, inner_y, inner_r)
        dv = inner_pt.normalized() * distension
   
        # project dv onto -v : |c| = a.b / |b| (a onto b)
        retract_dist = Physics.project_v1_on_v2(dv, -inner_v)
   
        # subtract velocity projection from current c.pt
        translation = inner_v.normalized() * retract_dist
       
        return inner_pos + translation
   
    @staticmethod
    def reflect_circle_in_circle((inner_x, inner_y, inner_vx, inner_vy, inner_e),
        (outer_x, outer_y, outer_e)):
        """Returns new velocity for inner circle as vector"""
        n = (vec2d(inner_x, inner_y) - vec2d(outer_x, outer_y)).normalized()
        v = vec2d(inner_vx, inner_vy)
        restitution = inner_e * outer_e
   
        # Solve Reflection
        # based on: http://stackoverflow.com/questions/573084/bounce-angle
        u = n * (v.dot(n))
        w = v - u
        v_ = w - u
        reflection = (v_ - v) * restitution
       
        new_inner_circle_velocity = v + reflection
        return new_inner_circle_velocity
   
    @staticmethod
    def project_v1_on_v2(vec1, vec2):
        scalar = vec1.dot(vec2) / vec2.length
        return scalar
   
       
class Circle(object):
    pos = vec2d(0,0)
    radius = 1
    color = RGB['white']
   
    def render_as_layer(self):
        r = rint(self.radius)
        dim = r * 2
       
        # draw transparent surface
        flags = pygame.SRCALPHA
        depth = 32
        base = pygame.Surface((dim, dim), flags, depth)
        base.set_alpha(128)
       
        # draw circle
        pygame.draw.circle(base, self.color, (r, r), r, 0)
       
        # adjust pos
        pt = self.pos + vec2d(-r, r)
       
        return Layer(base, pt)
   

class Court(Circle):
    radius = 100
    elasticity = 1.0
   
   
class Token(Circle):
   
    pos = vec2d(0,0)
    vel = vec2d(0,0)
    radius = 1
    mass = 10
    elasticity = 1.0
    color = RGB['blue']
    lost_vel = 0
   
    def __init__(self):
        self.lost_vel = 0
       
    def update(self):
        #self.redirect_slightly()
        #self.limit_speed()
        self.pos += self.vel
       
    def redirect_slightly(self):
        # change direction by +/- n degrees
        # ref: http://stackoverflow.com/q/4780119/1093087
        max_rotation = 4
        rotation = random.randint(-max_rotation, max_rotation)
       
        # new direction to vector
        theta = math.radians(rotation)
        ct = math.cos(theta);
        st = math.sin(theta);
        vx = self.vel.x * ct - self.vel.y * st;
        vy = self.vel.x * st + self.vel.y * ct;
        self.vel = vec2d(vx, vy)
       
        return self
       
    def limit_speed(self):
        speed = self.vel.length
        if speed > self.max_speed:
            print "slowing speed from %s to %s" % (speed, self.max_speed)
            self.vel.length = self.max_speed
            self.lost_vel += (speed - self.max_speed)
        return self

    @property
    def max_speed(self):
        return rint(self.radius / 1.5) - 1
   
    def __repr__(self):
        return "<%s Token %s pt=%s v=%s>" % (
            self.color, id(self), self.pos, self.vel)


#
# Simulation Class
#
class Simulation(object):
   
    screen_width = 300
    screen_height = 300
    seconds = 60.0
    fps = 30
   
    def __init__(self):
        Graphics.init(Simulation.screen_width, Simulation.screen_height)
        self.frame = 0
        self.start_at = time.time()
        self.max_frame = Simulation.fps * Simulation.seconds
        self.clock = pygame.time.Clock()
       
        self.court = Court()
        self.court.radius = 150
        self.court.color = RGB['gray']
        self.vlog = []
       
        ri = lambda min_, max_: random.randint(min_, max_)
       
        # tokens
        colors = ('red', 'blue', 'white')
        tokens_per_color = 5
        self.tokens = []
        for color in colors:
            for n in range(tokens_per_color):
                token = Token()
                token.radius = 16
                token.color = RGB[color]
                token.vel = vec2d(ri(-100,100), ri(-100,100))
                token.vel.length = token.max_speed * ri(25,50) / 100.0
                token.pos = vec2d(ri(-100,100), ri(-100,100))

                self.tokens.append(token)
   
    def update(self):
        # tick clock
        self.frame += 1
        self.clock.tick(Simulation.fps)
       
        # handle events
        self.handle_events()

        # update token
        self.update_tokens()
       
        # collisions
        self.resolve_collisions()
        self.enforce_court_boundaries()
       
        # game over?
        if self.frame >= self.max_frame:
            raise GameOver("time is up")
           
   
    def paint(self):
        Graphics.add_layer_to_canvas(self.court.render_as_layer())
       
        for token in self.tokens:
            Graphics.add_layer_to_canvas(Graphics.draw_token(token))
       
        pygame.display.set_caption("Frame: %d / FPS: %.4f" % (
            self.frame, self.frame / self.duration))
       
        Graphics.render()
   
    def run(self):
        while True:
            self.update()
            self.paint()
           
    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                raise GameOver("user requested game over")
           
    def update_tokens(self):
        for token in self.tokens:
            token.update()
           
    def resolve_collisions(self):
        collisions = 0
       
        for i, token in enumerate(self.tokens):
            for other_token in self.tokens[i+1:]:
                if token.color == other_token.color:
                    continue
               
                if Physics.circles_are_colliding(token, other_token):
                    self.collide_tokens(token, other_token)
                    collisions += 1
       
        return collisions
           
    def tokens_are_colliding(self, token1, token2):
        return Physics.circles_are_colliding(token1, token2)
        return Physics.circles_overlap(
            (token1.pos.x, token1.pos.y, token1.radius),
            (token2.pos.x, token2.pos.y, token2.radius)
        )
   
    def collide_tokens(self, token1, token2):
        t1_r, t2_r = Physics.collide_circle_with_circle(
            (token1.pos.x, token1.pos.y, token1.vel.x, token1.vel.y,
             token1.radius, token1.mass, token1.elasticity),
            (token2.pos.x, token2.pos.y, token2.vel.x, token2.vel.y,
             token2.radius, token2.mass, token2.elasticity))
       
        # apply collision resolution
        pos, vel = t1_r
        token1.pos = vec2d(pos)
        token1.vel = vec2d(vel)
       
        # apply collision resolution
        pos, vel = t2_r
        token2.pos = vec2d(pos)
        token2.vel = vec2d(vel)
           
    def enforce_court_boundaries(self):
        bounces = 0
       
        for token in self.tokens:
            if not Physics.circle_contains_circle(self.court.pos.x,
                self.court.pos.y, self.court.radius, token.pos.x,
                token.pos.y, token.radius):
                pos, vel = Physics.resolve_circle_in_circle_collision(
                    (token.pos.x, token.pos.y, token.radius,
                     token.vel.x, token.vel.y, token.elasticity),
                    (self.court.pos.x, self.court.pos.y, self.court.radius,
                     self.court.elasticity))
               
                # translate and reflect
                token.pos = vec2d(pos)
                token.vel = vec2d(vel)
               
                bounces += 1
               
        return bounces
           
    @property
    def duration(self):
        return time.time() - self.start_at
           
 

#
# Main
#
def main():
    try:
        sim = Simulation()
        sim.run()
    except GameOver, e:
        print "[%s] %s frames complete in %.4fs" % (e, sim.frame, sim.duration)
        time.sleep(2)
        pygame.quit()
        sys.exit(0)

def profile_main():
    command = "main()"
    fname = 'main-%s.cprof' % (str(date.today()).replace('-',''))
    fpath = pathjoin(dirname(__file__), fname)

    t0 = time.time()
    cProfile.run(command, filename=fpath)
    t1 = time.time() - t0
    print "\nprofiling complete is %.4fs\n" % (t1)

    p = pstats.Stats(fpath)
    p.sort_stats('time', 'cumulative').print_stats(100)

if __name__ == "__main__":
    main_mod = sys.modules[globals()['__name__']]
    if "profile" in sys.argv:
        profile_main()
    else:
        main()