"""Draw the set, with the two components whose shapes are known exactly.

Escape time on a grid, then the main cardioid and the period-2 disc
plotted from their formulas over the top. Writes mandelbrot.png beside
this file.
"""
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

RE = (-2.2, 0.8)
IM = (-1.3, 1.3)
W, H = 1600, 1387
NMAX = 400

re = np.linspace(*RE, W)
im = np.linspace(*IM, H)
c = re[None, :] + 1j * im[:, None]

z = np.zeros_like(c)
escaped_at = np.full(c.shape, NMAX, dtype=np.int32)
alive = np.ones(c.shape, dtype=bool)
for n in range(NMAX):
    z[alive] = z[alive] ** 2 + c[alive]
    just = alive & (np.abs(z) > 2.0)
    escaped_at[just] = n
    alive &= ~just

fig, ax = plt.subplots(figsize=(8.0, 6.94), dpi=200)
ax.imshow(np.log1p(escaped_at), extent=[*RE, *IM], origin="lower",
          cmap="bone", interpolation="bilinear")

# The main cardioid, c = mu/2 - mu^2/4 with mu on the unit circle.
t = np.linspace(0, 2 * np.pi, 2000)
mu = np.exp(1j * t)
card = mu / 2 - mu ** 2 / 4
ax.plot(card.real, card.imag, lw=1.1, color="#B36A1E")

# The period-2 component, the disc of radius 1/4 centred at -1.
ax.plot(-1 + 0.25 * np.cos(t), 0.25 * np.sin(t), lw=1.1, color="#00706B")

ax.set_xlabel("Re(c)")
ax.set_ylabel("Im(c)")
ax.set_xticks([-2, -1.5, -1, -0.5, 0, 0.5])
ax.set_yticks([-1, -0.5, 0, 0.5, 1])
for s in ax.spines.values():
    s.set_linewidth(0.6)
fig.tight_layout()
fig.savefig("mandelbrot.png", dpi=200)
print("mandelbrot.png written")