from __future__ import division
import sys, random, time, math
import pygame
from pymunk.vec2d import Vec2d

class Box(object):
    def __init__(self, x, y, w, h):
        self.x, self.y, self.w, self.h = x, y, w, h

    def center(self):
        return Vec2d(self.x + self.w / 2, self.y + self.h / 2)

class Anchor(object):
    """ignores updates to position"""
    def __init__(self, x, y):
        self._pos = x, y
        self.w = self.h = .00001

    x = property(lambda self: self._pos[0], lambda self, v: None)
    y = property(lambda self: self._pos[1], lambda self, v: None)
    def center(self):
        return Vec2d(self._pos)

class Spring(object):
    def __init__(self, b1, b2, rest=1.0):
        self.b1, self.b2 = b1, b2
        self.rest = rest

    def length(self):
        b1 = self.b1
        b2 = self.b2
        return math.hypot((b1.x + b1.w/2) - (b2.x + b2.w/2),
                          (b1.y + b1.h/2) - (b2.y + b2.h/2))

def absMin(a, absB):
    if a > 0 and a > absB:
        return absB
    if a < 0 and a < -absB:
        return -absB
    return a


class World(object):
    def __init__(self):
        self.steps = 0
        self.boxes = []
        self.horizSprings = []
        self.vertSprings = []

        self.horizSpringRest = 1.0
        self.vertGap = .2

    def measure(self):
        """note current bounding box of all boxes"""
        self.xmin = min(b.x for b in self.boxes)
        self.xmax = max(b.x + b.w for b in self.boxes)
        self.ymin = min(b.y for b in self.boxes)
        self.ymax = max(b.y + b.h for b in self.boxes)


    def step(self, dt):
        force = {} # box : y force
        startTime = time.time()
        for b in self.boxes:
            force[b] = 0

        self.stepHorizSprings(force)
        
        overlapCount = self.stepVertSprings(force)
            
        totalMove = 0
        for b in self.boxes:
            dy = force[b] * dt
            b.y += dy
            totalMove += abs(dy)
        self.steps += 1
        print "step %s: %s vertSprings overlapping; avg move per box %s; elapsed %.04f" % (
            self.steps,
            overlapCount,
            totalMove / len(self.boxes),
            time.time() - startTime)

    def stepHorizSprings(self, force):

        for s in self.horizSprings:
            stretch = (s.length() - s.rest) / s.rest
            f = (s.b2.y + s.b2.h/2 - s.b1.y + s.b1.h/2) * stretch * .004
            # shorten horiz springs
            force[s.b2] -= f
            force[s.b1] += f
        

    def stepVertSprings(self, force):
        overlapCount = 0
        for s in self.vertSprings:

            upper, lower = s.b1, s.b2
            if lower.y < upper.y:
                upper, lower = lower, upper
            
            minSpace = self.vertGap
            curSpace = lower.y - (upper.y + upper.h)
            if curSpace < minSpace:
                # overlapping!
                f = (minSpace - curSpace) * 1200
                overlapCount += 1
                force[upper] = force[lower] = 0
            else:
                # reduce gap
                f = - (curSpace - minSpace) * 250


            if abs(f) > 2:
                force[upper] += - f
                force[lower] += f
        return overlapCount


def pygameDrawWorld(screen, world, withSprings=True):

    xmin, xmax = world.xmin, world.xmax
    ymin, ymax = world.ymin, world.ymax

    xscl = screen.get_width() / (xmax - xmin)
    yscl = screen.get_height() / (ymax - ymin)

    for b in world.boxes:
        pygame.draw.rect(screen, (150, 0, 0),
                         ((b.x - xmin) * xscl,
                          (b.y - ymin) * yscl,
                          b.w * xscl - 2,
                          b.h * yscl - 2))

    if withSprings:
        def trans(v):
            return (v - (xmin, ymin)) * (xscl, yscl)

        for s in world.horizSprings:
            pygame.draw.line(screen, (150, 150, 0),
                             trans(s.b1.center()),
                             trans(s.b2.center()))

        for s in world.vertSprings:
            pygame.draw.line(screen, (255, 0, 255),
                             trans(s.b1.center()),
                             trans(s.b2.center()))

 
def main():
    pygame.init()
    screen = pygame.display.set_mode((600, 600))
    clock = pygame.time.Clock()
    running = True
    
    w = World()
    anc1 = Anchor(.5, 3)
    anc2 = Anchor(6.5, 3)
    w.boxes.extend(
        [Box(1, 1, 1, 3),
         Box(2, 1, 1, 2.5),
         Box(3, 1, 1, 3),
         Box(4, 1, 1, 4),
         Box(5, 1, 1, 3),

         Box(1, 3, 1, 1.5),
         Box(2, 4, 1, 2),
         Box(3, 4, 1, 1.7),
         anc1, anc2,
             ])

    b = w.boxes
    w.horizSprings.extend(
        [Spring(b[0], b[1]),
         Spring(b[1], b[2]),
         Spring(b[2], b[3]),
         Spring(b[3], b[4]),
         
         Spring(b[5], b[6]),
         Spring(b[6], b[7]),
         
         Spring(anc1, b[0]),
         Spring(b[4], anc2),
         ])

    w.vertSprings.extend([Spring(b[0], b[5]),
                          Spring(b[1], b[6]),
                          Spring(b[2], b[7]),
                          ])
    w.measure()

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
                running = False
        
        screen.fill((255, 255, 255))
        
        dt = 1/50.0

        pygameDrawWorld(screen, w)
        w.step(dt)
        
        pygame.display.flip()
        clock.tick(1/dt)
        
if __name__ == '__main__':
    sys.exit(main())