import png # purepng
import math, cmath
import sys

# Adapted from https://github.com/cbertram/mandelbrot/blob/master/mandelbrot.py, which is copyright Christian Betram under MIT license.
# Uses purepng pure-Python library ==> runs in PyPy, astronomically faster than in CPython.

preset = 1
if preset == 1:
  name = "MandelbrotSetVariant.png"
  width, height = 2000, 2000
  centerX, xSize = -0.5, 3
  centerY, ySize = 0, 3
  maxIterCount = 256 * 2
  z0 = 0
  
baseX = centerX - xSize / 2.0
baseY = centerY + ySize / 2.0

def setUpColors1(maxIterCount):
    # Use gray and a single hue
    red, green, blue = [0] * maxIterCount, [0] * maxIterCount, [0] * maxIterCount
    for iterCount in xrange(maxIterCount):
        a = (maxIterCount - iterCount) % 32
        b = (maxIterCount - iterCount) % 64
        if a == b:#b < a
            a = 32 - a
        baseColor = a * 7 + 31
        redFactor = a / 32.0
        greenFactor = (0.5 + redFactor * 0.5)
        red[iterCount] = int(baseColor * redFactor)
        green[iterCount] = int(baseColor * greenFactor)
        blue[iterCount] = baseColor
    return red, green, blue

def setUpColors2(maxIterCount):
    # Cycle various hues to better see iteration count escape boundaries
    red, green, blue = [0] * maxIterCount, [0] * maxIterCount, [0] * maxIterCount
    for iterCount in xrange(maxIterCount):
        m = iterCount
        baseColor = (maxIterCount - m) % 256
        if m % 5 == 0:
            red[iterCount] = int(baseColor * 0.75)
            green[iterCount] = baseColor
            blue[iterCount] = baseColor
        elif m % 5 == 1:
            red[iterCount] = int(baseColor * 0.75)
            green[iterCount] = int(baseColor * 0.75)
            blue[iterCount] = baseColor
        elif m % 5 == 2:
            red[iterCount] = baseColor
            green[iterCount] = baseColor
            blue[iterCount] = int(baseColor * 0.75)
        elif m % 5 == 3:
            red[iterCount] = baseColor
            green[iterCount] = int(baseColor * 0.75)
            blue[iterCount] = int(baseColor * 0.75)
        elif m % 5 == 4:
            red[iterCount] = int(baseColor * 0.75)
            green[iterCount] = baseColor
            blue[iterCount] = int(baseColor * 0.75)
    return red, green, blue

def setUpColors3(maxIterCount):
    # Use a single hue and very slowly decreasing color brightness
    red, green, blue = [0] * maxIterCount, [0] * maxIterCount, [0] * maxIterCount
    for iterCount in xrange(maxIterCount):
        a = (maxIterCount - iterCount) % 256
        b = (maxIterCount - iterCount) % 512
        if a == b:
            a = 255 - a
        baseColor = a
        redFactor = a / 256.0
        greenFactor = (0.5 + redFactor * 0.5)
        red[iterCount] = int(baseColor * redFactor)
        green[iterCount] = int(baseColor * greenFactor)
        blue[iterCount] = baseColor
    return red, green, blue

def iterate(z0, c, maxIterCount):
    z = z0
    iterCount = 0
    while iterCount < maxIterCount:
        z = z ** 2 + c
        #z = z ** 1 + c
        #z = z ** 2.1 + c
        if z.real**2 + z.imag**2 > 4:
            return iterCount
        iterCount += 1
    return None

def main():
    red, green, blue = setUpColors1(maxIterCount)
    image = []
    for h in xrange(height):
        if h % 10 == 0:
            sys.stdout.write(str(100 * h / height) + "% complete \r")
            sys.stdout.flush()
        row = [0] * width * 3
        for w in xrange(width):
            x = baseX + w * float(xSize) / width
            y = baseY - h * float(ySize) / height
            m = iterate(z0, complex(x, y), maxIterCount)
            if m is None:
                row[3 * w] = 0
                row[3 * w + 1] = 0
                row[3 * w + 2] = 0
            else:
                row[3 * w] = red[m]
                row[3 * w + 1] = green[m]
                row[3 * w + 2] = blue[m]
        image.append(row)
    print ""
    with open(name, "wb") as file1:
        w = png.Writer(width, height, greyscale=False)
        w.write(file1, image)

if __name__ == "__main__":
    main()