"""Draw a magnified piece of the edge, coloured in bands.

The centre is a point on the boundary near the seahorse valley. Colour
comes from the smooth escape count, which turns the integer step number
into a continuous value so the bands are even rather than stepped.
Writes zoom.png beside this file.
"""
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

CENTRE = complex(-0.743643887037151, 0.13182590420533)
HALF = 5.0e-4
W, H = 1800, 1350
NMAX = 3000
ESCAPE = 1e4                      # large, so the smooth count is accurate

re = np.linspace(CENTRE.real - HALF, CENTRE.real + HALF, W)
im = np.linspace(CENTRE.imag - HALF * H / W, CENTRE.imag + HALF * H / W, H)
c = re[None, :] + 1j * im[:, None]

z = np.zeros_like(c)
count = np.zeros(c.shape)
alive = np.ones(c.shape, dtype=bool)
for n in range(NMAX):
    z[alive] = z[alive] ** 2 + c[alive]
    just = alive & (np.abs(z) > ESCAPE)
    # n + 1 - log2(log|z|) removes the step, so the bands are continuous
    count[just] = n + 1 - np.log(np.log(np.abs(z[just]))) / np.log(2)
    alive &= ~just
count[alive] = np.nan                     # inside the set, painted flat

# deep blue, cyan, cream, amber, near-black. The classic banding, and it
# reads on a light page as well as a dark one.
cmap = LinearSegmentedColormap.from_list("bands", [
    (0.00, "#04102b"), (0.16, "#1c4fa1"), (0.42, "#63c7e8"),
    (0.60, "#fbf4dd"), (0.78, "#e0912f"), (0.92, "#6b2410"),
    (1.00, "#04102b")])
cmap.set_bad("#050608")

shown = np.sqrt(count - np.nanmin(count))          # even out the spread
shown = (shown % 12) / 12                          # cycle the palette

fig = plt.figure(figsize=(W / 200, H / 200), dpi=200)
ax = fig.add_axes([0, 0, 1, 1])
ax.imshow(shown, cmap=cmap, origin="lower", interpolation="bilinear", vmin=0, vmax=1)
ax.set_axis_off()
fig.savefig("zoom.png", dpi=200)

# The palette is bands, so 256 colours hold it with no visible loss and
# the file drops from 2.4MB to under one. Dithering keeps the wide
# gradient in the background from stepping.
from PIL import Image
img = Image.open("zoom.png").convert("RGB")
img.quantize(colors=256, method=Image.MEDIANCUT,
             dither=Image.FLOYDSTEINBERG).save("zoom.png", optimize=True)
print("zoom.png written")