"""Draw one step of the map, on the same ruled plane the maths pages use.

Squaring doubles the argument and squares the modulus, and adding c
translates the result. One figure showing all three, with the unit circle
drawn because squaring moves a point toward it or away from it.

The plane, addition, modulus, polar and multiplication diagrams that used
to live here moved to Mathematics, Fundamentals, where several articles can
use them. This script keeps only what is about the Mandelbrot map.

Writes squaring.svg beside this file.
"""
import numpy as np
import matplotlib
matplotlib.use("Agg")
matplotlib.rcParams["svg.fonttype"] = "path"
import matplotlib.pyplot as plt

TEAL, INDIGO, ROSE, AMBER = "#00706B", "#3A49A6", "#A93356", "#B36A1E"
INK, MUTED, GRID, FINE = "#0D1620", "#5A6874", "#D3DBE3", "#E1E7ED"


def imag_label(y):
    """Name a tick on the vertical axis, which counts in i rather than in units."""
    return "i" if y == 1 else "-i" if y == -1 else f"{y}i"


def grid(xlim, ylim, size, step=1.0, ticks=True, named=True):
    """One complex plane, ruled like graph paper and ready to draw on.

    Five fine squares to each ruled one, which is how squared paper is
    printed. The fine lines carry no meaning. They are there so a length
    can be read off by counting rather than estimated.
    """
    fig, ax = plt.subplots(figsize=size)
    lo_x, hi_x = int(np.floor(xlim[0])), int(np.ceil(xlim[1]))
    lo_y, hi_y = int(np.floor(ylim[0])), int(np.ceil(ylim[1]))
    fine = step / 5
    for x in np.arange(lo_x, hi_x + fine / 2, fine):
        ax.axvline(x, lw=0.5, color=FINE, zorder=0)
    for y in np.arange(lo_y, hi_y + fine / 2, fine):
        ax.axhline(y, lw=0.5, color=FINE, zorder=0)
    for x in np.arange(lo_x, hi_x + step / 2, step):
        ax.axvline(x, lw=0.8, color=GRID, zorder=0)
    for y in np.arange(lo_y, hi_y + step / 2, step):
        ax.axhline(y, lw=0.8, color=GRID, zorder=0)
    ax.axhline(0, lw=1.3, color=INK, zorder=1)
    ax.axvline(0, lw=1.3, color=INK, zorder=1)

    # Ticks are drawn by hand, because the two axes are labelled differently.
    for x in ([] if not ticks else range(lo_x, hi_x + 1)):
        if x == 0 or not xlim[0] < x < xlim[1]:
            continue
        ax.plot([x, x], [-0.09, 0.09], lw=1.1, color=INK, zorder=2)
        ax.annotate(f"{x}", (x, -0.17), ha="center", va="top", fontsize=10, color=INK)
    for y in ([] if not ticks else range(lo_y, hi_y + 1)):
        if y == 0 or not ylim[0] < y < ylim[1]:
            continue
        ax.plot([-0.07, 0.07], [y, y], lw=1.1, color=INK, zorder=2)
        ax.annotate(imag_label(y), (-0.17, y), ha="right", va="center",
                    fontsize=10, color=INK)
    if ticks:
        ax.annotate("0", (-0.17, -0.17), ha="right", va="top", fontsize=10, color=INK)
    # The first two figures are ordinary trigonometry, drawn on the same
    # squared paper but before the vertical axis counts in i.
    if named:
        ax.annotate("real axis", (xlim[1] - 0.1, 0.2), ha="right", va="bottom",
                    fontsize=11, color=MUTED)
        ax.annotate("imaginary axis", (0.2, ylim[1] - 0.1), ha="left", va="top",
                    fontsize=11, color=MUTED)

    ax.set_xlim(*xlim)
    ax.set_ylim(*ylim)
    ax.set_aspect("equal")
    ax.set_xticks([])
    ax.set_yticks([])
    for s in ax.spines.values():
        s.set_visible(False)
    return fig, ax


def arrow(ax, tip, colour, tail=0j, lw=1.6, style="-|>"):
    ax.annotate("", xy=(tip.real, tip.imag), xytext=(tail.real, tail.imag),
                arrowprops=dict(arrowstyle=style, lw=lw, color=colour,
                                shrinkA=0, shrinkB=0), zorder=3)


def arc(ax, ang, radius, colour, start=0.0):
    """An arc from one angle to another, for showing what a turn does."""
    t = np.linspace(start, ang, 240)
    ax.plot(radius * np.cos(t), radius * np.sin(t), lw=1.4, color=colour, zorder=2)
    mid = (start + ang) / 2
    return radius * np.cos(mid), radius * np.sin(mid)


# What each drawing colour becomes in the file. The page carries these
# drawings rather than linking them, so they inherit the reader's theme
# instead of answering the operating system. An img element could not:
# it is a separate document and cannot read --ink from the page.
PROPERTY = {TEAL: "--teal", INDIGO: "--indigo", ROSE: "--rose",
            AMBER: "--amber", INK: "--ink", MUTED: "--muted",
            GRID: "--line", FINE: "--sunk"}


def save(fig, name):
    """Write the figure, then make it the page's own drawing.

    Three passes over what matplotlib produced. Every colour becomes a
    custom property with the literal kept as a fallback. Every id is given
    the file's stem, because several of these are inlined into one page and
    matplotlib numbers them from one in each. And six decimal places become
    two, which is 0.005 of a point, well under a pixel at any size the page
    uses. A number whose whole part is zero is left alone: those are the
    glyph scale factors, where rounding would resize the text.
    """
    import re

    fig.tight_layout()
    # transparent, so the page shows through. Without it matplotlib
    # paints an opaque white rectangle behind everything, and the
    # figure then reads as a light card whatever the theme does to
    # the ink drawn on it.
    fig.savefig(name, transparent=True)
    text = open(name, encoding="utf-8").read()

    stem = re.sub(r"[^A-Za-z0-9]+", "-", name.rsplit(".", 1)[0]).strip("-")
    text = re.sub(r'(\sid=")([^"]+)"',
                  lambda m: f'{m.group(1)}{stem}-{m.group(2)}"', text)
    text = re.sub(r'((?:xlink:)?href="#)([^"]+)"',
                  lambda m: f'{m.group(1)}{stem}-{m.group(2)}"', text)
    text = re.sub(r'(url\(#)([^)]+)\)',
                  lambda m: f'{m.group(1)}{stem}-{m.group(2)})', text)

    for literal, prop in PROPERTY.items():
        text = re.sub(re.escape(literal), f"var({prop}, {literal.lower()})",
                      text, flags=re.IGNORECASE)

    text = re.sub(r"(\d+)\.(\d{3,})",
                  lambda m: (m.group(0) if m.group(1) == "0"
                             else f"{float(m.group(0)):.2f}".rstrip("0").rstrip(".")),
                  text)

    with open(name, "w", encoding="utf-8") as f:
        f.write(text)
    print(f"{name} written")


# One step of the map. The unit circle matters here, because squaring
# moves a point toward it or away from it.
r0, th0 = 0.85, np.deg2rad(52)
zs = r0 * np.exp(1j * th0)
zs2 = zs ** 2
c = -0.45 + 0.30j
zn = zs2 + c
fig, ax = grid((-1.35, 1.35), (-1.25, 1.45), (6.8, 6.6), step=0.5, ticks=False)
# The half-unit ruling carries no labels, so mark the two that matter.
for v in (-1, 1):
    ax.plot([v, v], [-0.07, 0.07], lw=1.1, color=INK, zorder=2)
    ax.annotate(f"{v}", (v, -0.13), ha="center", va="top", fontsize=10, color=INK)
    ax.plot([-0.06, 0.06], [v, v], lw=1.1, color=INK, zorder=2)
    ax.annotate(imag_label(v), (-0.13, v), ha="right", va="center",
                fontsize=10, color=INK)
t = np.linspace(0, 2 * np.pi, 600)
ax.plot(np.cos(t), np.sin(t), lw=1.0, ls=(0, (4, 3)), color=MUTED, zorder=1)

arc(ax, th0, 0.30, TEAL)
arc(ax, 2 * th0, 0.46, INDIGO)
for v, col, lab in ((zs, TEAL, "z"), (zs2, INDIGO, "z squared")):
    ax.plot([0, v.real], [0, v.imag], lw=1.5, color=col, zorder=3)
    ax.plot(v.real, v.imag, "o", ms=6.5, color=col, zorder=4)
    ax.annotate(lab, (v.real, v.imag), textcoords="offset points",
                xytext=(9, 7), color=col, fontsize=11)
arrow(ax, zn, ROSE, tail=zs2, lw=1.5)
ax.plot(zn.real, zn.imag, "o", ms=6.5, color=ROSE, zorder=4)
ax.annotate("z squared plus c", (zn.real, zn.imag), textcoords="offset points",
            xytext=(-6, 12), color=ROSE, fontsize=11, ha="center")
mid = (zs2 + zn) / 2
ax.annotate("add c", (mid.real, mid.imag), textcoords="offset points",
            xytext=(-2, -16), color=ROSE, fontsize=10, ha="center")
ax.annotate("\u03b8", (0.30 * np.cos(th0 / 2), 0.30 * np.sin(th0 / 2)),
            textcoords="offset points", xytext=(7, -2), color=TEAL, fontsize=10)
ax.annotate("2\u03b8", (0.46 * np.cos(2 * th0), 0.46 * np.sin(2 * th0)),
            textcoords="offset points", xytext=(-4, 12), color=INDIGO,
            fontsize=10, ha="right")
ax.annotate("modulus 1", (np.cos(np.deg2rad(-42)), np.sin(np.deg2rad(-42))),
            textcoords="offset points", xytext=(8, -10), color=MUTED, fontsize=10)
save(fig, "squaring.svg")
