import numba

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 = 100, 2000  # pixel density (= image width) and image height
n, r = 100000, 100000  # number of iterations and escape radius (r > 2)

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

S = np.zeros(n + 100, dtype=np.complex128)  # 100 iterations are chained
u, v = dc.Decimal(0), dc.Decimal(0)

for i in range(n + 100):
    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 * np.pi, num=d+1)
y = np.linspace(0, 2 * h / d * np.pi, num=h+1)

A, B = np.meshgrid(x, y, indexing="ij")  # transpose the matrices
C = (- 8.0) * np.exp((A + B * 1j) * 1j)

@numba.njit(parallel=True, fastmath=True)
def iteration_numba_bla(S, C):
    I, J = np.zeros(C.shape, dtype=np.int64), np.zeros(C.shape, dtype=np.complex128)
    E, Z, dZ = np.zeros_like(C), np.zeros_like(C), np.zeros_like(C)

    def iteration(S, dS, R, A, B, C):
        I, J = np.zeros(C.shape, dtype=np.int64), np.zeros(C.shape, dtype=np.complex128)
        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 step2(delta, index, epsilon, z, dz):
            index, epsilon = index + 1, (2 * S[index] + epsilon) * epsilon + delta
            z, dz = S[index] + epsilon, 2 * z * dz + 1
            index, epsilon = index + 1, (2 * S[index] + epsilon) * epsilon + delta
            z, dz = S[index] + epsilon, 2 * z * dz + 1
            return index, epsilon, z, dz

        def skip100(delta, index, e, z, dz):
            de = dz - dS[index]  # no catastrophic cancellation (don't try that with e)
            # for k in range(100):  # skip 100 iterations (using linear approximations)
            #     index, e, de = index + 1, 2 * S[index] * e + delta, 2 * S[index] * de
            index, e, de = index + 100, A[index] * e + B[index] * delta, A[index] * de
            z, dz = S[index] + e, dS[index] + de
            return index, e, z, dz

        def kernel(delta, index, epsilon, z, dz):
            i, j = 0, 0
            while i + j < n:
                if abs2(z) < abs2(r):
                    if abs2(epsilon) < abs2(1e-10 * R[index]):  # bla_impurity = 1e-10
                        index, epsilon, z, dz = skip100(delta, index, epsilon, z, dz)
                        j = j + 100
                    else:
                        if abs2(z) < abs2(epsilon):
                            index, epsilon = 0, z  # reset the reference orbit
                        index, epsilon, z, dz = step2(delta, index, epsilon, z, dz)
                        i = i + 2
                else:
                    break

            return complex(i + j, j), index, epsilon, z, dz

        for k in range(len(C)):
            J[k], I[k], E[k], Z[k], dZ[k] = kernel(C[k], I[k], E[k], Z[k], dZ[k])

        return J, I, E, Z, dZ

    A, B = np.ones(n, dtype=np.complex128), np.zeros(n, dtype=np.complex128)
    R, aS = np.full(n, 2, dtype=np.float64), np.where(np.abs(S) < 2, np.abs(S), 0)
    dS = np.zeros_like(S)

    for i in range(1, n + 100):  # derivation of the series (accuracy is not required)
        dS[i] = 2 * S[i - 1] * dS[i - 1] + 1

    for i in numba.prange(n):  # coefficients and radii for the bilinear approximation
        for j in range(100):
            A[i], B[i] = 2 * S[i + j] * A[i], 2 * S[i + j] * B[i] + 1
            R[i] = min(R[i], aS[i + j])  # validity radii and skip barriers (zeros)

    for i in numba.prange(len(C)):
        J[i], I[i], E[i], Z[i], dZ[i] = iteration(S, dS, R, A, B, C[i])

    return J, I, E, Z, dZ

J, I, E, Z, dZ = iteration_numba_bla(S, C)
D = np.zeros(C.shape)

skipped = J.imag.sum() / J.real.sum()
print("%.1f%% of all iterations were skipped." % (skipped * 100))

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 ** 0.015, cmap=plt.cm.gist_ncar, origin="lower")

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