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

# cp = center, xs = scale, win = complex number window
cp = -0.75,0
xs = 1.5

# W = width, H = height
W,H = 100,100

mandel = 1

def init():
    global win
    ys = xs * float(H)/float(W)
    win = [cp[0]-xs, cp[0]+xs, cp[1]-ys, cp[1]+ys]     
    glClearColor(1,0,0,0)
    glViewport(0,0,W,H)
    glMatrixMode(GL_PROJECTION);
    glLoadIdentity()
    # place integer coordinates at the center of the pixel
    # pixels run from 0 to W-1 and H-1
    apply(gluOrtho2D, (-.5,W-.5,-.5,H-.5))
    glMatrixMode(GL_MODELVIEW);

def set_color(n):
    # go from dark blue through blueish green to yellow
    if n:
        v = n / 99.
        #v = n / 10.
        #if v > 1: v = 1
        glColor3f(v*v, v, 0.2)
    else:
        glColor3f(0,0,0)

def dwell(x,y,cx,cy):
    """dwell(x,y) -> n
    How many iterations of c <- c*c+(x,y) are required for the complex
    number c to break out from |c|<2, starting from (0,0) ? (0 for never)."""
    cr,ci = x,y
    for i in range(1,100):
        cr,ci = cr*cr - ci*ci + cx, 2*cr*ci + cy
        if cr*cr+ci*ci > 4.0:
            return i
    return 0

def display():
    import time
    t = time.time()
    glClear(GL_COLOR_BUFFER_BIT)
    glBegin(GL_POINTS)
    # xf, yf are steps corresponding to one pixel
    # in the space where the mandelbrot shape has been defined
    xf, yf = (win[1]-win[0]) / float(W), (win[3]-win[2]) / float(H)
    for y in range(H):
        yc = y * yf + win[2]
        for x in range(W):
            xc = x * xf + win[0]
            if mandel:
                set_color(dwell(xc,yc,xc,yc)) # mandelbrot
            else:
                set_color(dwell(xc,yc,-.5,.58)) # julia set
            #set_color(dwell(xc,yc,.32,.26)) # julia set
            glVertex2i(x,y)
    glEnd()
    glFlush()
    print time.time() - t

def reshape(w,h):
    global W,H
    W,H = w,h
    init()

def kb(key, x, y):
    global mandel
    if key == 'm':
        mandel = 1
    elif key == 'j':
        mandel = 0
    elif ord(key) == 27: # ESC
        sys.exit()
    glutPostRedisplay()        

def mouse(button,state,x,y):
    # Left mousebutton zooms in, right mousebutton zooms out
    if state==1:
        global cp,xs
        xx = x/float(W)
        yy = 1.0-(y/float(H))
        cp = (win[1]-win[0])*xx+win[0], (win[3]-win[2])*yy+win[2]    
        xs *= [0.5,1.,2.][button]
        init()
        glutPostRedisplay()

if __name__ == '__main__':

    glutInit([])
    glutInitDisplayMode( GLUT_SINGLE | GLUT_RGB )
    glutInitWindowSize( W,H )
    glutCreateWindow( 'Mandelbrot' )
    #glutCreateWindow('Julia c=(-.5, .58)')
    #glutCreateWindow('Julia c=(.32, .26)')
    glutReshapeFunc( reshape )
    glutMouseFunc( mouse )
    glutDisplayFunc( display )
    glutKeyboardFunc( kb )
    init()
    glutMainLoop()