import numpy as np
import matplotlib.pyplot as plt

d, h = 800, 600  # pixel density (= image width) and image height
n, r = 200, 500  # number of iterations and escape radius (r > 2)

direction, height = 45.0, 1.5  # direction and height of the light
density, intensity = 4.0, 0.5  # density and intensity of the stripes

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

A, B = np.meshgrid(x - 1, y - h / d)
C = (1.5 + 1.0j) * (A + B * 1j) - 0.5

def iteration(C):
    S, T = np.zeros(C.shape), np.zeros(C.shape)
    Z, dZ, ddZ = np.zeros_like(C), np.zeros_like(C), np.zeros_like(C)

    def iterate(C, S, T, Z, dZ, ddZ):
        S, T = S + np.sin(density * np.angle(Z)), T + 1
        Z, dZ, ddZ = Z * Z + C, 2 * Z * dZ + 1, 2 * (dZ * dZ + Z * ddZ)
        return S, T, Z, dZ, ddZ

    for i in range(n):
        M = abs(Z) < r
        S[M], T[M], Z[M], dZ[M], ddZ[M] = iterate(C[M], S[M], T[M], Z[M], dZ[M], ddZ[M])

    return S, T, Z, dZ, ddZ

S, T, Z, dZ, ddZ = iteration(C)
D = np.zeros(C.shape)

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

N = abs(Z) >= r  # basic normal map effect and stripe average coloring (potential function)
P, Q = S[N] / T[N], (S[N] + np.sin(density * np.angle(Z[N]))) / (T[N] + 1)
U, V = Z[N] / dZ[N], 1 + (np.log2(np.log(abs(Z[N])) / np.log(r)) * (P - Q) + Q) * intensity
U, v = U / abs(U), np.exp(direction / 180 * np.pi * 1j)  # unit normal vectors and light vector
D[N] = np.maximum((U.real * v.real + U.imag * v.imag + V * height) / (1 + height), 0)

ax1 = fig.add_subplot(1, 2, 1)
ax1.imshow(D ** 1.0, cmap=plt.cm.bone, origin="lower")

N = abs(Z) > 2  # advanced normal map effect using higher derivatives (distance estimation)
H, K, L = abs(Z[N] / dZ[N]), abs(dZ[N] / ddZ[N]), np.log(abs(Z[N]))
U = Z[N] / dZ[N] - (H / K) ** 2 * L / (1 + L) * dZ[N] / ddZ[N]
U, v = U / abs(U), np.exp(direction / 180 * np.pi * 1j)  # unit normal vectors and light vector
D[N] = np.maximum((U.real * v.real + U.imag * v.imag + height) / (1 + height), 0)

ax2 = fig.add_subplot(1, 2, 2)
ax2.imshow(D ** 1.0, cmap=plt.cm.afmhot, origin="lower")

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