#/usr/bin/env python

import pygame
from pygame.locals import *
from math import *

def dot(A, B, C):
    AB = [B[0]-A[0],B[1]-A[1]]
    BC = [C[0]-B[0],C[1]-B[1]]
    return AB[0] * BC[0] + AB[1] * BC[1]

def cross(A, B, C):    
    AB = [B[0]-A[0],B[1]-A[1]]
    AC = [C[0]-A[0],C[1]-A[1]]    
    return AB[0] * AC[1] - AB[1] * AC[0]

def distance(A, B):
    d1 = A[0] - B[0]
    d2 = A[1] - B[1]
    return sqrt(d1*d1+d2*d2)

def pointOnLine(P,L):
    x1=L[0][0]   
    y1=L[0][1]
    x2=L[1][0]   
    y2=L[1][1]
    A1 = y2-y1
    B1 = x1-x2
    C1 = A1*x1+B1*y1
    
    x1 = P[0]
    y1 = P[1]
    A2 = -B1
    B2 = A1
    C2 = A2*x1+B2*y1
    
    det = A1*B2 - A2*B1
    x = (B2*C1 - B1*C2)/det
    y = (A1*C2 - A2*C1)/det      
    return (x,y)    

def linePointDist(A, B, C):    
    dot1 = dot(A,B,C)
    if(dot1 > 0):
        # End point B is closest
        return (distance(B,C),tuple(B))
    dot2 = dot(B,A,C)
    if(dot2 > 0): 
        # End point A is closest
        return (distance(A,C),tuple(A))
    
    # Point ON the line is closest
    dist = cross(A,B,C) / distance(A,B)
        
    return (abs(dist),pointOnLine(C,[[A[0],A[1]],[B[0],B[1]]]))

class Joint:
    def __init__(self,pos):
        self.pos = pos
        self.rect = pygame.Rect(0,0,10,10)
        self.rect.center = pos           

    def setPos(self,pos):
        self.pos = pos
        self.rect.center = pos           

class JointedLine:
    
    def __init__(self):
        self.joints=[]
        self.rect = None
        
    def updateRect(self):
        self.rect = None
        if not len(self.joints):
            return
                    
        rects = []
        for j in self.joints:
            if not self.rect:
                self.rect = j.rect
            else:
                rects.append(j.rect)
 
        self.rect = self.rect.unionall(rects)
    
    def insertJoint(self,pos):
        if len(self.joints) < 2:
            return None
        if not self.rect:
            return None
        if not self.rect.collidepoint(pos[0], pos[1]):
            return None

        p1 = None
        p2 = None
        for j in self.joints:
            if p2 == None:
                p2 = j.pos
                continue
            p1 = p2
            p2 = j.pos
            l = [p1,p2]
            d,p = linePointDist(p1,p2,pos)
            if d <= 5:
                nj = Joint(p)
                self.joints.insert(self.joints.index(j),nj)
                self.updateRect()
                return nj
                    
        return None
        
    def appendJoint(self,pos):
        j = Joint(pos)
        self.joints.append(j)
        self.updateRect()
        return j
        
    def deleteJoint(self,pos):
        j = self.getJoint(pos)
        if j:
            self.joints.remove(j)
            self.updateRect()
    
    def getJoint(self,pos):
        if not self.rect:
            return None
        if self.rect.collidepoint(pos[0], pos[1]):
            for p in self.joints:
                if p.rect.collidepoint(pos[0], pos[1]):
                    return p
        return None
         
    def draw(self,surface,color,drawjoints=False):

        if len(self.joints) >= 2:
            p = []
            for j in self.joints:
                p.append(j.pos)
            pygame.draw.aalines(surface, color, False, p,True)
    
        if drawjoints:
            for p in self.joints:
                pygame.draw.rect(surface, color, p.rect)                
 
def main():

    draw_joints = True

    # init pygame
    pygame.init()    

    # setup the display
    screen = pygame.display.set_mode((800, 600),1)

    fnt = pygame.font.Font(pygame.font.get_default_font(),18)
    
    back = pygame.Surface((800,600),1)
    back= back.convert()
    back.fill((200,200,200))

    txt = """Left click on an empty space to append a new joint.
Left click on a line to insert a new joint on that line.
Hold down left mouse button to drag joints.
Right click on a joint to delete it.
Press SPACE to toggle visual joints."""
    
    x = 10
    y = 10
    for t in txt.splitlines():
        t = fnt.render(t, True, (255,255,255))
        back.blit(t,(x,y))
        y+=20;

    clock = pygame.time.Clock()
    
    l = JointedLine()
    sel_joint = None

    while 1:
        clock.tick(30)        

        for event in pygame.event.get():
            if event.type == QUIT:
                return
            elif event.type == KEYDOWN:
                if event.key == K_ESCAPE:
                    return
                if event.key == K_SPACE:
                    draw_joints = not draw_joints
            elif event.type == MOUSEBUTTONDOWN:
                if event.button==1:
                    sel_joint = l.getJoint(event.pos)
                    if not sel_joint:
                        sel_joint = l.insertJoint(event.pos)
                    if not sel_joint:
                        sel_joint = l.appendJoint(event.pos) 
                elif event.button==3:
                    l.deleteJoint(event.pos) 
            elif event.type == MOUSEBUTTONUP:
                sel_joint = None
                l.updateRect()
            elif event.type == MOUSEMOTION:
                if sel_joint:
                    sel_joint.setPos(event.pos)

        # clear screen
        screen.blit(back,(0,0))
        l.draw(screen,(0,0,0),draw_joints)
        
        pygame.display.flip()
             
 
#this calls the 'main' function when this script is executed
if __name__ == '__main__': main()