import sys
import math

from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *


class Complex:
    def __init__(self, r = 0, i = 0):
        self.re = r
        self.im = i

class Extent:
    def __init__(self, l = 0, r = 0, b = 0, t = 0):
        self.l = l
        self.r = r
        self.b = b
        self.t = t


world = Extent(-1, 1, -1, 1)
c = Complex(0.109, 0.603)
width = 200
height = 200


def Julia(p, c):
    maxIterations = 256
    for its in range(maxIterations):
        (p.re, p.im) = (p.re*p.re - p.im*p.im + c.re,
                        2*p.re*p.im + c.im)
        rsqr = p.re*p.re + p.im*p.im
        if rsqr > 4.0:
            break
    return (its, math.sqrt(rsqr))

def Mandelbrot(c):
    maxIterations = 256
    p = Complex(0,0)
    for its in range(maxIterations):
        (p.re, p.im) = (p.re*p.re - p.im*p.im + c.re,
                        2*p.re*p.im + c.im)
        rsqr = p.re*p.re + p.im*p.im
        if rsqr > 4.0:
            break
    return (its, math.sqrt(rsqr))

def display():
    initialize()    
    delta = (world.r - world.l)/float(width)
    p = Complex()
    for j in xrange(height):
        y = world.b + j*delta
        for i in xrange(width):
            x = world.l + i*delta
            (p.re, p.im) = (x, y)
            (its, R) = Julia(p, c)
            #(its, R) = Mandelbrot(p)
            if its == 255:
                glColor3d(0,0,0)
            else:
                r = R/float(3)
                g = its/float(128)
                b = R/float(its+1)
                glColor3d(r,g,b)
            glBegin(GL_POLYGON)
            glVertex2d(x, y)
            glVertex2d(x, y+delta)
            glVertex2d(x+delta, y+delta)
            glVertex2d(x+delta, y)
            glEnd()
        glFlush()


def xScreenToWorld(scrX):
    return (world.r - world.l) * scrX / float(width) + world.l


def yScreenToWorld(scrY):
    return (world.t - world.b) * (1 - scrY / float(height)) + world.b


def onMouseButton(button, state, mx, my):
    x = xScreenToWorld(mx)
    y = yScreenToWorld(my)
    dx = world.r - world.l
    dy = world.t - world.b
    if (button == GLUT_LEFT_BUTTON) & (state == GLUT_DOWN):
        world.l = x - dx/4.0
        world.r = x + dx/4.0
        world.b = y - dy/4.0
        world.t = y + dy/4.0
        display()
    elif (button == GLUT_RIGHT_BUTTON) & (state == GLUT_DOWN):
        world.l = x - dx
        world.r = x + dx
        world.b = y - dy
        world.t = y + dy
        display()



def onKeyPress(key, x, y):
    global c
    print key
    if (key == 'r') | (key == 'R'): # reset view
        (c.re, c.im) = 0.109, 0.603
        world.l = -1
        world.r = 1
        world.b = -1
        world.t = 1
        display()
    elif (key == 'c') | (key == 'C'): # set c=(0,0)
        (c.re, c.im) = (0, 0)
        world.l = -1
        world.r = 1
        world.b = -1
        world.t = 1
        display()


def onResize(w, h):
    global width,height
    (width, height) = (w, h)
    glViewport(0, 0, w, h)
    cx = 0.5*(world.r + world.l)
    dy = world.t - world.b;
    world.l = cx - 0.5*dy * w/h
    world.r = cx + 0.5*dy * w/h
    
def initialize():
    global world
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    gluOrtho2D(world.l, world.r, world.b, world.t)
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()
    
    glColor3d(0,0,1)
    glBegin(GL_POLYGON)
    glVertex2d(world.l, world.b)
    glVertex2d(world.r, world.b)
    glVertex2d(world.r, world.t)
    glVertex2d(world.l, world.t)
    glEnd()



def main():
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_SINGLE | GLUT_RGB)
    glutInitWindowSize(width, height)
    glutInitWindowPosition(100, 100)
    glutCreateWindow("Julia Set (a.k.a. Whoville)")
    glutDisplayFunc(display)
    glutMouseFunc(onMouseButton)
    glutKeyboardFunc(onKeyPress)
    glutReshapeFunc(onResize)
    initialize()
    glutMainLoop()


if __name__ == "__main__":
    main()