# Metaballs by Gynvael Coldwind of Vexillium
# http://gynvael.coldwind.pl
# http://vexillium.org
# mailto: gynvael@vexillium.org
#
# Note:
#  This code was created for educational purposes, so it was supposed
#  to be as simple as possible, hence no real optimisation was used,
#  which leads to the fact that this code is ultra slow.
#  This version should work under both Windows and Linux platforms,
#  however it was only tested on Windows.
#
#  It requires pygame to run.
#
# License (BSD):
#  Copyright (c) 2008, Gynvael Coldwind of Vexillium
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#      * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#      * Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#      * Neither the name of the Gynvael Coldwind nor Vexillium nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY Gynvael Coldwind ''AS IS'' AND ANY
#  EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL Gynvael Coldwind BE LIABLE FOR ANY
#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import pygame, sys,os, time
from pygame.locals import * 
from math import sqrt
 
WINDOW_WIDTH  = 640
WINDOW_HEIGHT = 480
THRESHOLD     = 0.001

def Events(evs): 
  for ev in evs: 
    if ev.type == QUIT: 
      sys.exit(0) 
 
# Add a metaball
def AddMetaball(metaballs, x, y, vx, vy, r, g, b):
  metaballs.append({'x':x, 'y':y, 'vx':vx, 'vy':vy, 'r':r, 'g':g, 'b':b })

# Draw scene
def Scene(metaballs, canvas):
  # Calc metaballs
  for j in range(0, WINDOW_HEIGHT):
    for i in range(0, WINDOW_WIDTH):

      # Setup 
      r = 0.0
      g = 0.0
      b = 0.0
      power = 0.0

      # Calc power
      for m in metaballs:

        # Calc inverted dst
        dx = i - m['x']
        dy = j - m['y']
        if dx == 0 and dy == 0:
          idst = 1.0
        else:
          idst = 1.0 / (dx * dx + dy * dy)

        # Add stuff
        r = r + idst * m['r']
        g = g + idst * m['g']
        b = b + idst * m['b']
        power = power + idst

      # Normalize rgb
      len = sqrt(r*r + g*g + b*b)
      r = r / len
      g = g / len
      b = b / len

      # Check threshold
      colorpower = 128
      if power >= THRESHOLD:
        colorpower = 255

      # Recalc rgb
      r = (int)(r * colorpower)
      g = (int)(g * colorpower)
      b = (int)(b * colorpower)

      # Set color
      cl = pygame.Color(r,g,b,255)
      canvas.set_at((i, j), cl)

  # Flip
  pygame.display.flip()


# Calc metaball movements
def Calcs(metaballs):
  # For each metaball
  for m in metaballs:

    # Move the metaball
    m['x'] = m['x'] + m['vx']
    m['y'] = m['y'] + m['vy']

    # Check!
    if m['x'] >= WINDOW_WIDTH:
      m['x'] = WINDOW_WIDTH
      m['vx'] = -m['vx']

    if m['x'] < 0:
      m['x'] = 0
      m['vx'] = -m['vx']

    if m['y'] >= WINDOW_HEIGHT:
      m['y'] = WINDOW_HEIGHT
      m['vy'] = -m['vy']

    if m['y'] >= WINDOW_HEIGHT:
      m['y'] = WINDOW_HEIGHT
      m['vy'] = -m['vy']    

  
# Init pygame 
pygame.init()
 
# Crate a window
window = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT)) 
pygame.display.set_caption('Metaballs Python Psyco test') 

# Get canvas
canvas = pygame.display.get_surface() 

#screen.blit(monkey_surface, (0,0)) 

# Get metaballs
metaballs = []
AddMetaball(metaballs, 10.0, 20.0, 1.0, 2.0, 1.0, 0.0, 0.0)
AddMetaball(metaballs, 30.0, 50.0, 2.0, 1.0, 0.0, 0.0, 1.0)

# Init timer
lasttime = time.clock()

while True: 
   Events(pygame.event.get())
   Scene(metaballs, canvas)
   Calcs(metaballs)

   # Calc
   now = time.clock()
   elapsed = now - lasttime
   lasttime = now

   print "%f SPF" % elapsed