#!/usr/bin/env python
from gi.repository import Gtk
from gi.repository import GLib
import cairo
import math
import cmath
import random
from numpy import poly1d

##
## This is really inefficient quick and dirty
##

# This is ugly, scale changes what points we take, drawscale changes what
# we plot, mx and my move things (but don't move axes)
scale = 1.0;
drawscale = 0.5;
drawmx = 0.0;
drawmy = 0.0;

##
## The polynomial is set below in the compute_few function
##

#scale = 0.2;
#drawscale = 0.15;
#scale = 1.0;
#drawscale = 4.0;
#drawmx = 0.6;
#drawmy = 0.0;


x = [];
y = [];
col1 = [];
col2 = [];

zx = [];
zy = [];
zcol1 = [];
zcol2 = [];

stop_the_drawing = False

def OnDraw(w, cr):
  cr.set_line_width(1)
  cr.set_source_rgb(1, 1, 1)
  cr.rectangle (0, 0, 1200, 600)
  cr.fill ()

  cr.set_source_rgb(0, 0, 0)

  cr.move_to(600,0)
  cr.line_to(600,600)
  cr.stroke()

  cr.set_source_rgb(0.5, 0.5, 0.5)

  cr.move_to(0,300)
  cr.line_to(600,300)
  cr.move_to(300,0)
  cr.line_to(300,600)
  cr.stroke()

  cr.move_to(600,300)
  cr.line_to(1200,300)
  cr.move_to(900,0)
  cr.line_to(900,600)
  cr.stroke()

  for i in range(len(x)):
    cr.set_source_rgba(0.9*(1.0-col1[i]), 0.9*col1[i], 0.9*col2[i], 0.2)
    cr.arc(x[i],y[i],1, 0, 2*math.pi)
    cr.fill()

  for i in range(len(zx)):
    cr.set_source_rgba(0.9*(1.0-zcol1[i]), 0.9*zcol1[i], 0.9*zcol2[i], 1)
    cr.arc(zx[i],zy[i],1, 0, 2*math.pi)
    cr.fill()

def stop_drawing(w):
  global stop_the_drawing
  stop_the_drawing = True

w = Gtk.Window()
b = Gtk.VBox()
w.add(b)
a = Gtk.DrawingArea()
a.set_size_request(1200, 600)
b.add(a)

b2 = Gtk.HBox()
b.add(b2)


but = Gtk.Button("Stop")
but.connect('clicked', stop_drawing)
b2.add(but)

but = Gtk.Button("Quit")
but.connect('clicked', Gtk.main_quit)
b2.add(but)

def compute_few():
  for i in range(5000):
    if i & 0x11:
      while Gtk.events_pending(): Gtk.main_iteration()
    if stop_the_drawing:
      break
    z = random.uniform(-1.0/scale,1.0/scale)+1j*random.uniform(-1.0/scale,1.0/scale)
    #z = cmath.rect(random.uniform(0.0,1.0/scale),random.uniform(-math.pi,math.pi))
    # This is the polynomial
    #P = poly1d([1,z,0.2], variable = 't')
    P = poly1d([1,z,z], variable = 't')
    #P = poly1d([1,z**2,z/3], variable = 't')
    #P = poly1d([1,0,0,-z,z**3], variable = 't')
    #P = poly1d([1,-z], variable = 't')

    zcol1.append((z.real*scale + 1.0)/2.0)
    zcol2.append((z.imag*scale + 1.0)/2.0)
    zx.append(drawscale*(z.real+drawmx)*300+300)
    zy.append(-drawscale*(z.imag+drawmy)*300+300)
    for s in P.r:
      s = complex(s)
      col1.append((z.real*scale + 1.0)/2.0)
      col2.append((z.imag*scale + 1.0)/2.0)
      x.append(drawscale*(s.real+drawmx)*300+900)
      y.append(-drawscale*(s.imag+drawmy)*300+300)
  a.queue_draw()
  if len(x) > 1000000 or stop_the_drawing:
    return False
  else:
    return True
	    
w.connect('destroy', Gtk.main_quit)
a.connect('draw', OnDraw)

w.show_all()

GLib.idle_add(compute_few)

Gtk.main()