import numpy as np
import matplotlib.pyplot as plt

import decimal as dc  # decimal floating point arithmetic with arbitrary precision
dc.getcontext().prec = 80  # set precision to 80 digits (about 256 bits)

d, h = 50, 1000  # pixel density (= image width) and image height
n, r = 80000, 100000.0  # number of iterations and escape radius (r > 2)

a = dc.Decimal("-1.256827152259138864846434197797294538253477389787308085590211144291")
b = dc.Decimal(".37933802890364143684096784819544060002129071484943239316486643285025")

S = np.zeros(n + 2, dtype=np.complex128)
u, v = dc.Decimal(0), dc.Decimal(0)

for i in range(n + 2):
    S[i] = float(u) + float(v) * 1j
    if u * u + v * v < r * r:
        u, v = u * u - v * v + a, 2 * u * v + b
    else:
        print("The reference sequence diverges within %s iterations." % i)
        break

x = np.linspace(0, 2, num=d+1, dtype=np.float64)
y = np.linspace(0, 2 * h / d, num=h+1, dtype=np.float64)

A, B = np.meshgrid(x * np.pi, y * np.pi)
C = (- 8.0) * np.exp((A + B * 1j) * 1j)

def iteration(S, C):
    I = np.zeros(C.shape, dtype=np.intp)
    E, Z, dZ = np.zeros_like(C), np.zeros_like(C), np.zeros_like(C)

    def abs2(z):
        return z.real * z.real + z.imag * z.imag

    def iterate2(C, I, E, Z, dZ):
        I, E = I + 1, (2 * S[I] + E) * E + C
        Z, dZ = S[I] + E, 2 * Z * dZ + 1
        I, E = I + 1, (2 * S[I] + E) * E + C
        Z, dZ = S[I] + E, 2 * Z * dZ + 1
        return I, E, Z, dZ

    for i in range(0, n, 2):
        M, R = abs2(Z) < abs2(r), abs2(Z) < abs2(E)
        I[R], E[R] = 0, Z[R]  # reset the reference orbit
        I[M], E[M], Z[M], dZ[M] = iterate2(C[M], I[M], E[M], Z[M], dZ[M])

    return I, E, Z, dZ

I, E, Z, dZ = iteration(S, C)
D = np.zeros(C.shape, dtype=np.float64)

fig = plt.figure(figsize=(12.8, 1.6))
fig.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95)

N = abs(Z) > 2  # exterior distance estimation
D[N] = np.log(abs(Z[N])) * abs(Z[N]) / abs(dZ[N])

ax1 = fig.add_subplot(1, 1, 1)
ax1.imshow(D.T ** 0.015, cmap=plt.cm.nipy_spectral, origin="lower")

fig.savefig("Mandelbrot_numpy_set_4.png", dpi=200)