#! /usr/bin/env python

import pygtk
pygtk.require('2.0')
import gtk, gobject, cairo
import random
import math

class Screen(gtk.DrawingArea):

    # Draw in response to an expose event
    __gsignals__ = { "expose-event": "override" }

    # Handle the expose-event by drawing
    def do_expose_event(self, event):
	cr = self.window.cairo_create()

	cr.rectangle(event.area.x, event.area.y,
		event.area.width, event.area.height)
        cr.clip()

	self.draw(cr, *self.window.get_size())

    def draw(self, cr, width, height):
	cr.set_source_rgb(0.3, 0.3, 0.3)
	cr.rectangle(0, 0, width, height)
	cr.fill()
        cr.set_source_rgb(1.0, 1.0, 1.0)
        cr.rectangle(width/4,height/4,width/2,height/2);
	cr.stroke()

def run(Widget):
    window = gtk.Window()
    window.connect("delete-event", gtk.main_quit)
    widget = Widget()
    widget.show()
    window.add(widget)
    window.present()
    gtk.main()

class City:
  def __init__(self, x, y, vx, vy):
    self.x = x
    self.y = y
    self.vx = vx
    self.vy = vy
    self.t = 0
    self.friend = None
    self.history = []
  def move(self):
    if self.friend:
      # accellerate toward our friend
      dx = self.friend.x - self.x
      dy = self.friend.y - self.y
      self.vx += dx * 1.0e-3
      self.vy += dy * 1.0e-3
    # damp
    self.vx *= 0.939
    self.vy *= 0.939
    # euler integrate the position
    self.x += self.vx
    self.y += self.vy
    self.history.append((self.t, self.x, self.y, self.vx, self.vy))
    #print self.t, self.x, self.y, self.vx, self.vy
    self.t += 1
  def find_friend(self, cities):
    self.friend = random.choice(cities)
  def draw(self, cr):
    for (t, x, y, vx, vy) in self.history:
      if self.friend:
        cr.move_to(x, y)
        cr.line_to(self.friend.history[t][1], self.friend.history[t][2])
      else:
        cr.arc(x, y, 2.0, 0.0, 2.0*math.pi)
      cr.stroke()

class CityScape(Screen):
  def __init__(self, count=11):
    Screen.__init__(self)
    self.cities = []
    # generate a random collection of cities
    scale = 1000.0
    vscale = 100.0
    for i in xrange(count):
      x = scale*random.random() - 0.5*scale
      y = scale*random.random() - 0.5*scale
      vx = vscale*random.random() - 0.5*vscale
      vy = vscale*random.random() - 0.5*vscale
      city = City(x,y,vx,vy)
      self.cities.append(city)
    for city in self.cities:
      city.find_friend(self.cities)
    self.timer = gobject.timeout_add(100, self.update, self)
  def draw(self, cr, width, height):
    cr.set_source_rgb(1.0, 1.0, 1.0)
    cr.rectangle(0, 0, width, height)
    cr.fill()

    cr.set_source_rgba(0.0, 0.0, 0.0, 0.2)
    cr.translate(width/2.0, height/2.0)
    for city in self.cities:
      city.draw(cr)
  def update(self, window):
   for city in self.cities: city.move()
   self.window.invalidate_rect(window.allocation, False)
   return True

if __name__ == "__main__":
    run(CityScape)