#!/usr/pkg/bin/python
"""simple program to create Mandelbrot set

   this takes command-line arguments to create a PNG file"""

Copyright = """
    mandelgen -- generate image of Mandelbrot set
    Copyright (C) 2005  John Comeau <jc.jcomeau.com>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
    """
errormessage = "Not all needed libraries found, upgrade or check path: "
try:
 True # not defined in older Python releases
except:
 True, False = 1, 0
try:
 import sys, os, types, re, pwd
 sys.path.append(os.path.join(pwd.getpwuid(os.geteuid())[5], 'lib', 'python'))
 errormessage = errormessage + repr(sys.path)
 from com.jcomeau import gpl, jclicense
except:
 try:
  sys.stderr.write("%s\n" % errormessage)
 except:
  print errormessage
 raise

# get name this program was called as
myself = os.path.split(sys.argv[0])[1]
command = os.path.splitext(myself)[0]  # chop any suffix (extension)
# now get name we gave it when we wrote it
originalself = re.compile('[0-9A-Za-z]+').search(Copyright).group()

# globals and routines that should be in every program
# (yes, you could import them, but there are problems in that approach too)
import struct

def DebugPrint(*whatever):
 return False  # defined instead by pytest module, use that for debugging

def join(*args):
 "for pythons without str.join"
 string, array = args
 if type(array) == types.StringType: array = eval(array)
 if hasattr(str, 'join'):
  return string.join(array)
 else:
  joined = ''
  for index in range(0, len(array)):
   joined = joined + array[index]
   if index != (len(array) - 1):
    joined = joined + string
  return joined

def split(*args):
 "for pythons without str.split"
 string, string_to_split = args
 if not len(string): string = None
 if hasattr('str', 'split'):
  return string_to_split.split(string)
 else:
  return re.compile(re.escape(string)).split(string_to_split)

# other globals, specific to this program
current = {
 'x': -0.5,
 'y': 0.0,
 'width': 320,
 'height': 240,
 'ratio': 0,  # to be set in init()
 'zoom': 1.0,
 'output': sys.stdout,
 'colors': None,  # set in init()
 'maxiterations': 0,  # color values according to iterations
# if you make maxvalue 2, and only check z.real for it, you'll see what
# resemble lines of magnetic force around the Mandelbrot "lake".
 'maxvalue': 4,  # x^2 + y^2
 'png_maxval': 255,
 'crc_table': [], # set in init()
}

def mandelgen(*args):
 """main routine

    creates PPM image of Mandelbrot set at a particular location and zoom level
    each pixel is z -> z^2 + z0
 """
 DebugPrint('args', args)
 init(args)
 pngheader()
 DebugPrint('current', current['width'], current['height'], current['x'],
  current['y'], current['ratio'], current['zoom'])
 gen_by_pixel()

def gen_by_pixel(*args):
 showprogress, index = True, 0
 sys.stderr.write('Showing one dot per line processed\n')
 if DebugPrint('gen_by_pixel() starting'):
  current['maxiterations'] = 255
 while index < (current['height'] * current['width']):
  x, y = index % current['width'], int(index / current['width'])
  c, z, iterations = complex(widthvalue(x), heightvalue(y)), complex(0, 0), 0
  while iterations < current['maxiterations']:
   z = (z * z) + c
   try:
    if ((z.real ** 2) + (z.imag ** 2)) < current['maxvalue']:
     iterations += 1
    else:
     raise Exception, 'escaped'
   except:
    break
  pngpixel(iterations)
  if showprogress and index % current['width'] == 0:
   sys.stderr.write('.')
  index += 1
 if showprogress:
   sys.stderr.write('\n')

def complex_multiply(i, j):
 [a, b], [c, d] = i, j
 return [(a * c) - (b * d), (b * c) + (a * d)]

def complex_add(i, j):
 [a, b], [c, d] = i, j
 return [a + c, b + d]

def heightvalue(*args):
 heightspan = current['top'] - current['bottom']
 pixelheight = args[0]
 return current['top'] - ((heightspan * pixelheight) / current['height'])

def widthvalue(*args):
 widthspan = current['right'] - current['left']
 pixelwidth = args[0]
 return current['left'] + ((widthspan * pixelwidth) / current['width'])

def find_maxval():
 init()
 maxrgb = map(max, current['colors'])
 maxval = max(maxrgb)
 DebugPrint('maxval', maxval)
 return maxval

def init(*args):
 if current['colors'] == None:
  if len(args) == 1: args = args[0]
  DebugPrint('init() args', args)
  if len(args) >= 2:
   current['width'], current['height'] = int(args[0]), int(args[1])
  current['ratio'] = 640.0 / (current['width'] * 100)
  if len(args) >= 4:
   current['x'], current['y'] = float(args[2]), float(args[3])
  if len(args) >= 5:
   current['zoom'] = float(args[4])
  colorgenerator = rgbvalues
  if len(args) >= 6:
   colorgenerator = eval(args[5])
  if len(args) >= 7:
   try:
    import psyco
    psyco.full()
    sys.stderr.write('%s\n' % 'Psyco loaded for faster execution')
   except:
    sys.stderr.write('%s\n' % 'Psyco could not be loaded; expect slowdown')
  current['colors'] = colorgenerator()
  current['png_maxval'] = find_maxval()
  # maximum iterations must be one less than the number of colors; if there
  # are just two colors, black and white, one iteration is sufficient
  current['maxiterations'] = len(current['colors']) - 1
  DebugPrint('maxiterations', current['maxiterations'])
  current['ratio'] = float(current['ratio']) / current['zoom']
  current['bottom'] = current['y'] - \
   ((current['height'] * current['ratio']) / 2)
  current['top'] = current['bottom'] + (current['height'] * current['ratio'])
  current['left'] = current['x'] - ((current['width'] * current['ratio']) / 2)
  current['right'] = current['left'] + (current['width'] * current['ratio'])
  current['pixels'] = colormap()
  current['crc_table'] = crc_table()

def binary(binary_string):
 number = 0
 for character in binary_string:
  number <<= 1
  number |= (character == "1")
 return number

def rgb(*args):
 """simple 3-pass coloring scheme to match WikiHow article"""
 return [(255, 0, 0), (0, 255, 0), (0, 0, 255)]

def rgb2black(*args):
 """alternate r, g, b darker and darker to 0, 0, 0"""
 values = []
 colors = [[(color, 0, 0), (0, color, 0),
  (0, 0, color)] for color in range(255, -1, -1)]
 for triplet in colors:
  values += triplet
 return values

def rgbvalues(*args):
 """define start color as cyan: r=0 g>=b

    decrement for each iteration such that 0 is black
    to match the colorForth program, g must be 6 bits and b 5, with the low
    bits of each cleared, and two subtracted each time
    also, found out that the g and b values must be shifted into the
    high-order bits of their respective bytes, to match colorForth results
 """
 cyan = binary("11111011110")
 divisor = binary("100000")
 values = [(0, (color / divisor) << 2,
  (color % divisor) << 3) for color in range(cyan, -2, -2)]
 #DebugPrint('%d values' % len(values))
 return values

def bluevalues(*args):
 """use blue values from 1023 to 0"""
 values = [(0, 0, color) for color in range(1023, -1, -1)]
 return values

def shadesofblue(*args):
 """progress from white to cyan to blue to black"""
 values = \
  [(red, 255, 255) for red in range(255, -1, -1)] + \
  [(0, green, 255) for green in range(254, -1, -1)] + \
  [(0, 0, blue) for blue in range(254, -1, -1)]
 return values

def shadesofblue2(*args):
 """progress from cyan to blue to black"""
 values = \
  [(0, green, 511) for green in range(511, 0, -1)] + \
  [(0, 0, blue) for blue in range(510, -1, -1)]
 return values

def inverse_red(*args):
 """go from dark to light until the black of the Mandelbrot lake"""
 values = \
  [(1023, color, color) for color in range(1, 1024)] + [(0, 0, 0)]
 return values

def inverse_blue(*args):
 """go from dark to light until the black of the Mandelbrot lake"""
 values = \
  [(color, color, 1023) for color in range(0, 1024)] + [(0, 0, 0)]
 return values

def fnbrot(*args):
 """Fn's alternating color scheme

    e.g. yellow rg0 fading into magenta r0b"""
 values = []
 for step in range(252, -8, -8):
  for color in range(step, -4, -4):
   values += [(step, step - color, color)]
 return values

def twotoneblue(*args):
 """cyan to black, blue to black"""
 values = \
  [(0, color, color) for color in range(511, -1, -1)] + \
  [(0, 0, blue) for blue in range(511, -1, -1)]
 return values

def maxelbrot(*args):
 """switch between different color patterns from white to black

    it looks Peter Max-ish, hence the name"""
 values = [[
   (color, color, color), # white
   (color, color, 0), # yellow
   (color, 0, color), # magenta
   (0, color, color), # cyan
   (0, 0, color), # blue
   (color, 0, 0), # red
   (0, color, 0), # green
   ] for color in range(255, 0, -2)]
 colors = []
 for triplet in values:
  colors.extend(triplet)
 colors += [(0, 0, 0)]
 return colors[-1023:]

def colormap(*args):
 init()
 if DebugPrint('testing colormap'):
  current['maxiterations'] = 255
  pass
 if (current['maxiterations'] + 1) < len(current['colors']):
  divisor = len(current['colors']) / (current['maxiterations'] + 1)
  if len(current['colors']) % (current['maxiterations'] + 1):
   divisor += 1
  DebugPrint('divisor', divisor)
  colormap = filter(lambda s: current['colors'].index(s) % divisor == 0,
   current['colors'])
  while len(colormap) < (current['maxiterations'] + 1):
   colormap.append(current['colors'][-1])
 else:
  colormap = current['colors']
 DebugPrint('len(colormap)', len(colormap))
 return colormap

def pngheader():
 current['output'].write('\x89PNG\r\n\x1a\n')
 bit_depth = 8 + (8 * (current['png_maxval'] > 255))
 color_type = 2  # each pixel is RGB triple
 compression_method = 0
 filter_method = 0
 interlace_method = 0
 ihdr = struct.pack('>L', current['width']) + \
  struct.pack('>L', current['height']) + \
  chr(bit_depth) + chr(color_type) + chr(compression_method) + \
  chr(filter_method) + chr(interlace_method);
 ihdr = struct.pack('>L', len(ihdr)) + 'IHDR' + ihdr + \
  struct.pack('>L', update_crc(ihdr))
 current['output'].write(ihdr)

def pngpixel(iterations):
 r, g, b = current['pixels'][iterations]
 if current['png_maxval'] < 256:
  current['output'].write('%c%c%c' % (r, g, b))
 else:
  current['output'].write('%c%c%c%c%c%c' % (r >> 8, r % 256,
   g >> 8, g % 256, b >> 8, b % 256))

def colorlength(*args):
 current['output'] = open('/dev/null', 'w')
 colortest(*args)
 print len(current['colors'])
 
def colortest(*args):
 colorgenerator, height = None, 10
 if len(args):
  colors = eval(args[0])()
 else:
  sys.stderr.write('colortest must be called with a function name\n')
  exit(1)
 init(len(colors), height, 0, 0, 1, args[0])
 pngheader()
 DebugPrint(current['pixels'])
 for row in range(height):
  for column in range(len(colors)):
   pngpixel(column)

def crc_table():
 """table of CRCs of all 8-bit messages

    from PNG-1.1 spec
 """
 table = []
 for value in range(256):
  crc = value
  for bit in range(8):
   if crc & 1:
    crc = 0xedb88320L ^ (crc >> 1)
   else:
    crc = crc >> 1
  table.append(crc)
 return table
   
def update_crc(*args):
 """update a running CRC with the bytes buf[0..len-1]

    the CRC should be initialized to all 1's, and the transmitted 
    value is the 1's complement of the final running CRC
    (from PNG-1.1 spec)
 """
 init()
 mask = 0L
 if len(args) < 2:
  args = (0xffffffffL, args[0])
  mask = 0xffffffffL
 crc, bytes = args
 for byte in bytes:
  crc = current['crc_table'][(crc ^ ord(byte)) & 0xff] ^ (crc >> 8)
 return crc ^ mask

if __name__ == '__main__':
 # if this program was imported by another, the above test will fail,
 # and this following code won't be used...
 function = command; args = sys.argv[1:]  # default case
 if command == originalself:
  try:
   if len(args) and eval('type(%s) == types.FunctionType' % args[0]):
    function = sys.argv[1]; args = sys.argv[2:]
  except: pass
 print eval('%s%s' % (function, repr(tuple(args)))) or ''
else:
 # if you want something to be done on import, do it here; otherwise pass
 pass