"""Draw one map, its orbits and its fixed points.

Two figures. The first is a cobweb, which reads an orbit off the graph of
the map without any arithmetic. The second is the same two orbits plotted
as sequences, which is the form the text works in.

Both use the map x -> x squared, whose fixed points are 0 and 1 and whose
behaviour changes at 1. It is the real-number half of the map that draws
the Mandelbrot set, so a reader meets the mechanism before the complex
arithmetic is added to it.

Writes cobweb.svg and orbits.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 square(x):
    return x * x


def orbit(x0, steps):
    """The sequence a starting point produces, as a list."""
    xs = [x0]
    for _ in range(steps):
        xs.append(square(xs[-1]))
    return xs


def axes(xlim, ylim, size, step=0.5):
    """Plain ruled paper with two labelled axes, for a graph rather than a plane."""
    fig, ax = plt.subplots(figsize=size)
    fine = step / 5
    for x in np.arange(np.floor(xlim[0]), xlim[1] + fine / 2, fine):
        ax.axvline(x, lw=0.5, color=FINE, zorder=0)
    for y in np.arange(np.floor(ylim[0]), ylim[1] + fine / 2, fine):
        ax.axhline(y, lw=0.5, color=FINE, zorder=0)
    for x in np.arange(np.floor(xlim[0]), xlim[1] + step / 2, step):
        ax.axvline(x, lw=0.8, color=GRID, zorder=0)
    for y in np.arange(np.floor(ylim[0]), ylim[1] + 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)
    ax.set_xlim(*xlim)
    ax.set_ylim(*ylim)
    ax.set_xticks([])
    ax.set_yticks([])
    for s in ax.spines.values():
        s.set_visible(False)
    return fig, ax


def cobweb(ax, x0, steps, colour, lw=1.5):
    """Draw the staircase that reads an orbit off the graph.

    Up to the curve, across to the diagonal, and repeat. Each horizontal
    move carries the output back to the input axis, which is what applying
    the map again means.
    """
    x, xs, ys = x0, [x0], [0.0]
    for _ in range(steps):
        y = square(x)
        xs += [x, y]
        ys += [y, y]
        x = y
    ax.plot(xs, ys, lw=lw, color=colour, zorder=3, solid_joinstyle="round")


# 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")


# The cobweb. The staircase is the orbit, read off the graph rather than
# computed, and the diagonal is where an output becomes the next input.
fig, ax = axes((-0.30, 1.90), (-0.30, 1.90), (6.8, 5.8))
t = np.linspace(0, 1.38, 400)
ax.plot(t, square(t), lw=2.2, color=TEAL, zorder=2)
ax.plot([0, 1.9], [0, 1.9], lw=1.3, ls=(0, (5, 4)), color=MUTED, zorder=2)
ax.annotate("y = x squared", (1.24, 1.66), ha="right", va="bottom",
            fontsize=11, color=TEAL)
ax.annotate("the diagonal", (0.98, 1.05), ha="right", va="bottom",
            fontsize=11, color=MUTED)

cobweb(ax, 0.80, 7, ROSE)
cobweb(ax, 1.05, 3, AMBER)
ax.annotate("", xy=(1.86, 1.86), xytext=(1.4775, 1.4775),
            arrowprops=dict(arrowstyle="-|>", lw=1.5, color=AMBER,
                            shrinkA=0, shrinkB=0), zorder=3)

for x, colour, lab in ((0.80, ROSE, "0.8"), (1.05, AMBER, "1.05")):
    ax.plot(x, 0, "o", ms=7, color=colour, zorder=4)
    ax.annotate(lab, (x, -0.09), ha="center", va="top", fontsize=11,
                color=colour)
ax.annotate("each staircase starts on the horizontal axis", (-0.28, -0.24),
            ha="left", va="center", fontsize=10, color=MUTED)

for x in (0.0, 1.0):
    ax.plot(x, x, "o", ms=8, color=INDIGO, zorder=5)
ax.annotate("fixed points, 0 and 1", (0.04, 1.44), ha="left", va="bottom",
            fontsize=11, color=INDIGO)

ax.annotate("from 0.8 the staircase falls to 0", (1.87, 0.42), ha="right",
            va="center", fontsize=10, color=ROSE)
ax.annotate("from 1.05 it climbs away", (1.87, 0.24), ha="right",
            va="center", fontsize=10, color=AMBER)
save(fig, "cobweb.svg")


# The same two orbits as sequences, which is how the text writes them.
fig, ax = axes((-0.95, 8.8), (-0.55, 2.20), (6.8, 4.6), step=0.5)
near, away = orbit(0.80, 8), orbit(1.12, 8)
for xs, colour in ((near, ROSE), (away, AMBER)):
    n = [i for i, v in enumerate(xs) if v <= 2.12]
    ax.plot(n, [xs[i] for i in n], lw=1.6, color=colour, zorder=3,
            marker="o", ms=6)

# The escaping orbit leaves the top of the frame between two steps, so the
# arrow stands in for every step after the last one drawn.
last = max(i for i, v in enumerate(away) if v <= 2.12)
ax.annotate("", xy=(last + 0.80, 2.17), xytext=(last, away[last]),
            arrowprops=dict(arrowstyle="-|>", lw=1.6, color=AMBER,
                            shrinkA=0, shrinkB=0), zorder=3)
ax.annotate("leaves every bound", (last + 0.90, 2.10), ha="left", va="top",
            fontsize=11, color=AMBER)
ax.annotate("settles on 0", (5.2, 0.13), ha="left", va="bottom",
            fontsize=11, color=ROSE)

ax.axhline(1.0, lw=1.2, ls=(0, (5, 4)), color=INDIGO, zorder=2)
ax.annotate("1, the fixed point between them", (8.7, 1.06), ha="right",
            va="bottom", fontsize=10, color=INDIGO)

for i in range(0, 9):
    ax.plot([i, i], [-0.06, 0.06], lw=1.0, color=INK, zorder=2)
    ax.annotate(f"{i}", (i, -0.13), ha="center", va="top", fontsize=10,
                color=INK)
for v in (1, 2):
    ax.plot([-0.07, 0.07], [v, v], lw=1.0, color=INK, zorder=2)
    ax.annotate(f"{v}", (-0.15, v), ha="right", va="center", fontsize=10,
                color=INK)
ax.annotate("step number", (4.0, -0.42), ha="center", va="center",
            fontsize=11, color=MUTED)
save(fig, "orbits.svg")