#!/usr/bin/python
# Code originally from the site:
# http://0pointer.de/blog/projects/mandelbrot.html
# Modified by D. Vrajitoru

import Image, ImageDraw, math, colorsys
from bpnn import *

dimensions = (600, 600)
scale = 1.0/(dimensions[0]/3)
#center = (2.2, 1.5)       # Use this for Mandelbrot set
center = (1.5, 1.5)       # Use this for Julia set
iterate_max = 100
colors_max = 50

img = Image.new("RGB", dimensions)
d = ImageDraw.Draw(img)

# Calculate a tolerable palette
palette = [0] * colors_max

def original_palette():
    for i in xrange(colors_max):
        f = 1-abs((float(i)/colors_max-1)**15)
        r, g, b = colorsys.hsv_to_rgb(.66+f/3, 1-f/2, f)
        palette[i] = (int(r*255), int(g*255), int(b*255))

#Create a palette with a neural network that we already trained
def nn_palette(neural_net):
    for i in xrange(colors_max):
        f = 1-abs((float(i)/colors_max-1)**15)
        rgb = neural_net.update([f])
        palette[i] = (int(rgb[0]*255), int(rgb[1]*255), int(rgb[2]*255))

# Calculate the mandelbrot sequence for the point c with start value z
def iterate_mandelbrot(c, z = 0):
    for n in xrange(iterate_max + 1):
        z = z*z +c
        if abs(z) > 2:
            return n
    return None

# Draw our image
def draw_image():
    for y in xrange(dimensions[1]):
        for x in xrange(dimensions[0]):
            c = complex(x * scale - center[0], y * scale - center[1])

            #n = iterate_mandelbrot(c)            # Use this for Mandelbrot set
            n = iterate_mandelbrot(complex(0.3, 0.6), c)  # Use this for Julia set

            if n is None:
                v = 1
            else:
                v = n/100.0
                
            d.point((x, y), fill = palette[int(v * (colors_max-1))])

if __name__ == '__main__':
    # one input: f, 3 hidden, and 3 output: r, g, b
    nnet = NN(1, 3, 3)
    colors=[[[0], [0, 0, 1]],
            #[[0.75], [0, 1, 0]],
            [[1], [1, 0, 0]]]
    nnet.train(colors)
    nn_palette(nnet)
    draw_image()
    
    del d
    img.save("result.png")