'''MandelbrotSet.py
This program computes a fractal image based on the Mandelbrot set.
The program facilitates changing the resolution of the image (and thus
the time required to compute it) and the region of the complex plane
that is being considered.

If MONOCHROME is set to True, then the image is grayscale, and the
mapping is chosen to make the image look good in that mode.
If not, then if PSEUDOCOLOR is set to True, then a color image is 
produced that achieves a balance between
having different colors and maintaining some understandable relationship
between (most of) the colors and the "count" of iterations used to determine
that status of a pixel.  The scheme here is the following:
Divergence: Dark blue
Low count values: Medium green.
Count values approaching 255: red.
Count values above 255: red with amounts of green added that cycle from
nothing to max every 128 values, thereby creating yellows and whites.
High values of count (e.g., 1000): more and more blue, so that some pixels
can appear white hot.
If PSEUDOCOLOR is set of False, then the "pure" count values are stored into
the image being produced, encoded in the blue and green components.
The procedure makeColor can then be used to create a color image that is
an alternative to what you get with PSEUDOCOLOR turned on.

The code below contains commented-out versions of various useful settings.
'''

#W=192; H=128 # For the initial range, showing the whole Mandelbrot set
#W=64; H=64
#W=128; H=128
#W=256; H=256
#W=512; H=512
W=1024; H=1024 # On some parts of the complex plane, this will take a very long time.
#REAL_RANGE = (-2, 1); IMAGINARY_RANGE = (-1, 1) # The initial range, showing the whole Mandelbrot set
#REAL_RANGE = (0, 0.2); IMAGINARY_RANGE = (0.5, 0.7)
#REAL_RANGE = (0.1, 0.15); IMAGINARY_RANGE = (0.6, 0.65) # One nice detail area.

# The following region contains a tree-like detail:
REAL_RANGE = [0.1208282470703125, 0.1283050537109375]
IMAGINARY_RANGE = [0.6367431640625, 0.644219970703125]

MONOCHROME=False
PSEUDOCOLOR=False

WIN = pmNewComputedImage("Mandelbrot Set within a in "+\
        str(REAL_RANGE)+" and b in "+\
        str(IMAGINARY_RANGE),\
      W,H,"rgb(255, 230, 230)")

def createMandelbrotImage():
  global WIN
  ncols = pmGetImageWidth(WIN)
  nrows = pmGetImageHeight(WIN)
  rectWidth = float(REAL_RANGE[1]-REAL_RANGE[0])
  rectHeight = float(IMAGINARY_RANGE[1]-IMAGINARY_RANGE[0])
  for j in range(ncols):
    for i in range(nrows):
      real = REAL_RANGE[0]+j*rectWidth/ncols
      imag = IMAGINARY_RANGE[0]+i*rectHeight/nrows
      c = complex(real, imag)
      count = iterate(c)
      (r,g,b) = getColor(count)
      pmSetPixel(WIN, j, i, r, g, b)

NMAX = 2000
THRESH = 2000
def iterate(c):
  global NMAX, THRESH
  n = 0
  z = complex(0,0)
  while n<NMAX:
    z = z*z + c
    if abs(z)>THRESH: return n
    n += 1
    #if n<10: report(n,z)
  return -1

RCOUNT = 0
def report(n,z):
  global RCOUNT
  if RCOUNT < 200:
    print "n="+str(n)+"; z="+str(z)
    RCOUNT += 1

def getColor(count):
  global NMAX, MONOCHROME, PSEUDOCOLOR
  if count==-1:
    if MONOCHROME: return (0,0,0)
    else: return (0, 0, 96)
  else:
    if MONOCHROME:
      b = int(count/6) #% 256
      b = 10*b; r = b; g = b # Special for monochrome version.
    else:
      if PSEUDOCOLOR:
        r = count
        if count>255: g = (count*2) % 256
        else: g = max(0,192-count*2)
        b = int(count/8)
      else: # Preserve full count info in output image...
        b = count % 256
        g = (count / 256) % 256
        r = (count / 65536) % 256

    return (r,g,b)

def findNewRectangle(x, y, width, height):
  ''' Given the current ranges and these arguments,
      compute a new pair of ranges that makes it easy to
      zoom in on the selected region of the complex plane.
      The center of the new region will be (x,y) and the
      width, in pixels of the current image, will be used
      to compute the intervals.'''
  global W, H, REAL_RANGE, IMAGINARY_RANGE
  (rstart,rend)=REAL_RANGE
  (istart,iend)=IMAGINARY_RANGE
  rextent = rend-rstart
  iextent = iend-istart
  new_x_center = rstart+(float(x)/W)*rextent
  new_y_center = istart+(float(y)/H)*iextent
  new_half_width = rextent*width/(2*W)
  new_real_range = [new_x_center - new_half_width, new_x_center + new_half_width]
  new_half_height = iextent*height/(2*H)
  new_imaginary_range = [new_y_center - new_half_height, new_y_center + new_half_height]
  print "Use REAL_RANGE = "+str(new_real_range)
  print "and IMAGINARY_RANGE = "+str(new_imaginary_range)

def makeColor():
  '''Produce a colored version of the image from the pure counts image.'''
  global WIN
  newWin = pmCloneImage(WIN)
  pmSetSource1(WIN)
  pmSetDestination(newWin)
  # The mapping here does well on the tree-like detail, making the branches
  # brown and putting green around them.
  pmSetFormula("RGB(Blue1(x,y)*2, (4*Blue1(x,y)) mod 256 + Green1(x,y)*100,50)")
  pmCompute()



createMandelbrotImage()
if not MONOCHROME and not PSEUDCOLOR: makeColor()