import matplotlib
matplotlib.use('Agg')

from mpi4py import MPI
import numpy as np
import pylab

INITIAL_Z = 0.4
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

min_x = float(input("Enter min_x:\n"))
min_y = float(input("Enter max_x:\n"))
max_x = float(input("Enter min_y:\n"))
max_y = float(input("Enter max_y:\n"))
x = int(input("Enter x (width):\n"))
y = int(input("Enter y (height):\n"))
iters = int(input("Enter number of iterations (0, 255):\n"))
print('Please wait...')


def iterations(real, imag, iters):
    z = INITIAL_Z
    r = complex(real, imag)
    for i in range(iters):
        z = z * r * (1 - z)
        if z.real * z.real + z.imag * z.imag > 3:
            return 255 * i // iters
    return 255


def create_fractal(min_x, max_x, min_y, max_y, image, iters):
    height = image.shape[0]
    width = image.shape[1]
    pixel_size_x = (max_x - min_x) / width
    pixel_size_y = (max_y - min_y) / height
    for x in range(width):
        real = min_x + x * pixel_size_x
        for y in range(height):
            imag = min_y + y * pixel_size_y
            color = iterations(real, imag, iters)
            image[y, x] = color
    return image


image = np.zeros((y, x))
create_fractal(min_x, min_y, max_x, max_y, image, iters)
pylab.imshow(image, aspect='auto', extent=[max_x, max_y, min_y, min_x])
pylab.gca().invert_yaxis()
pylab.gray()
pylab.savefig('result.png')
print('Done.')