"""Draw a matrix as a bracketed grid, and the three operations on it.

Three figures. The first names the parts, the second adds two matrices
position by position, and the third multiplies a row into a column. The
third is the one that needs a picture, because it is the only rule here
that does not work position by position.

The numbers are the textbook pair, and the script checks its own arithmetic
against NumPy before drawing it, so a wrong entry cannot reach the page.

Writes shape.svg, sum.svg and product.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"

CW, CH = 0.98, 0.72          # one cell, in drawing units


def cell(x0, y0, i, j):
    """Centre of the cell in row i and column j, counting from zero."""
    return x0 + (j + 0.5) * CW, y0 - (i + 0.5) * CH


def matrix(ax, x0, y0, entries, colour=INK, size=13, tint=None, ink=None):
    """Draw a bracketed grid with its top-left corner at (x0, y0).

    tint takes {(i, j): colour} and fills those cells, which is how a row,
    a column and the entry they produce are picked out of three grids. The
    text in a filled cell takes the same colour unless ink says otherwise,
    which is what a whole band wants: the band is tinted, and only the one
    entry it is about is recoloured.
    """
    rows, cols = len(entries), len(entries[0])
    w, h = cols * CW, rows * CH

    for (i, j), fill in (tint or {}).items():
        cx, cy = cell(x0, y0, i, j)
        ax.add_patch(plt.Rectangle((cx - CW / 2, cy - CH / 2), CW, CH,
                                   facecolor=fill, alpha=0.16, lw=0, zorder=1))
    for i in range(rows):
        for j in range(cols):
            cx, cy = cell(x0, y0, i, j)
            at = (ink.get((i, j)) if ink is not None
                  else (tint or {}).get((i, j)))
            ax.annotate(entries[i][j], (cx, cy), ha="center", va="center",
                        fontsize=size, color=at or colour, zorder=3)

    # Square brackets, drawn as three strokes each, the way they are printed.
    ear = 0.16
    for side in (0, 1):
        bx = x0 - 0.14 if side == 0 else x0 + w + 0.14
        dx = ear if side == 0 else -ear
        ax.plot([bx, bx], [y0 + 0.10, y0 - h - 0.10], lw=1.6, color=MUTED,
                zorder=2, solid_capstyle="round")
        for yy in (y0 + 0.10, y0 - h - 0.10):
            ax.plot([bx, bx + dx], [yy, yy], lw=1.6, color=MUTED, zorder=2,
                    solid_capstyle="round")
    return w, h


def plane(xlim, ylim, size=None, ax=None, step=1.0):
    """Ruled paper with two axes, for drawing a vector as an arrow."""
    if ax is None:
        _, 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_aspect("equal")
    ax.set_xticks([])
    ax.set_yticks([])
    for sp in ax.spines.values():
        sp.set_visible(False)
    return ax


def vec(ax, tip, colour, tail=(0.0, 0.0), lw=2.0, style="-|>"):
    ax.annotate("", xy=tuple(tip), xytext=tuple(tail),
                arrowprops=dict(arrowstyle=style, lw=lw, color=colour,
                                shrinkA=0, shrinkB=0), zorder=4)


def board(size, xlim, ylim):
    fig, ax = plt.subplots(figsize=size)
    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, tail, tip, colour, lw=1.5, rad=0.0):
    ax.annotate("", xy=tip, xytext=tail,
                arrowprops=dict(arrowstyle="-|>", lw=lw, color=colour,
                                shrinkA=0, shrinkB=0,
                                connectionstyle=f"arc3,rad={rad}"), zorder=4)


# 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

    # bbox_inches trims the canvas to what was actually drawn. These
    # figures set their limits by hand to leave room for the labels, and
    # without the crop that room is shipped as blank margin on the page.
    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,
                bbox_inches="tight", pad_inches=0.04)
    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 parts of a matrix, named. Rows across, columns down, and one entry
# picked out to show which of the two subscripts counts which.
SYM = [[r"$a_{11}$", r"$a_{12}$", r"$a_{13}$", r"$a_{14}$"],
       [r"$a_{21}$", r"$a_{22}$", r"$a_{23}$", r"$a_{24}$"],
       [r"$a_{31}$", r"$a_{32}$", r"$a_{33}$", r"$a_{34}$"]]

fig, ax = board((5.6, 1.9), (-1.15, 8.30), (-1.95, 1.05))
band = {(1, j): TEAL for j in range(4)}
band.update({(i, 2): INDIGO for i in range(3)})
band[(1, 2)] = ROSE
w, h = matrix(ax, 0, 0.55, SYM, tint=band, ink={(1, 2): ROSE})

ax.annotate("row 2", (-0.42, cell(0, 0.55, 1, 0)[1]), ha="right", va="center",
            fontsize=11, color=TEAL)
ax.annotate("column 3", (cell(0, 0.55, 0, 2)[0], 0.72), ha="center",
            va="bottom", fontsize=11, color=INDIGO)

ax.annotate("3 rows and 4 columns,", (4.45, -0.30), ha="left", va="center",
            fontsize=11, color=INK)
ax.annotate("so this matrix is 3 by 4", (4.45, -0.72), ha="left", va="center",
            fontsize=11, color=INK)
save(fig, "shape.svg")


# Addition, which works one position at a time. Two matrices of the same
# shape only, and the shape of the answer is that shape again.
A = np.array([[2, 5, 1], [0, 3, 4]])
B = np.array([[7, 1, 2], [6, 0, 5]])
S = A + B
txt = lambda M: [[f"{v}" for v in row] for row in M]

fig, ax = board((7.4, 2.0), (-0.30, 12.20), (-2.10, 0.60))
x = 0.0
for M, tint in ((A, {(0, 1): TEAL}), (B, {(0, 1): INDIGO}), (S, {(0, 1): ROSE})):
    w, h = matrix(ax, x, 0.35, txt(M), tint=tint)
    x += w + 1.55
mid = 0.35 - h / 2
ax.annotate("+", (3.71, mid), ha="center", va="center", fontsize=17, color=MUTED)
ax.annotate("=", (8.03, mid), ha="center", va="center", fontsize=17, color=MUTED)

ax.annotate("5 + 1 = 6, and every other position is added the same way",
            (-0.7, -1.55), ha="left", va="center", fontsize=11, color=ROSE)
ax.annotate("subtraction works the same way, one position at a time",
            (-0.7, -1.92), ha="left", va="center", fontsize=10, color=MUTED)
save(fig, "sum.svg")


# Multiplication, the one rule that does not work position by position.
# One row of the left matrix is spent against one column of the right.
P = np.array([[1, 2, 3], [4, 5, 6]])
Q = np.array([[7, 8], [9, 10], [11, 12]])
R = P @ Q
assert R.tolist() == [[58, 64], [139, 154]], R
assert R[1, 0] == 4 * 7 + 5 * 9 + 6 * 11 == 139

fig, ax = board((7.4, 3.4), (-0.95, 10.75), (-3.40, 1.15))
xa, ya = 0.0, 0.60
wa, ha_ = matrix(ax, xa, ya, txt(P),
                 tint={(1, 0): TEAL, (1, 1): TEAL, (1, 2): TEAL})
xb = xa + wa + 1.45
wb, hb = matrix(ax, xb, ya + 0.36, txt(Q),
                tint={(0, 0): INDIGO, (1, 0): INDIGO, (2, 0): INDIGO})
xc = xb + wb + 1.45
wc, hc = matrix(ax, xc, ya, txt(R), tint={(1, 0): ROSE})

ax.annotate("×", (xa + wa + 0.72, ya - ha_ / 2), ha="center", va="center",
            fontsize=16, color=MUTED)
ax.annotate("=", (xb + wb + 0.72, ya - ha_ / 2), ha="center", va="center",
            fontsize=16, color=MUTED)

# The row and the column both point at the entry they make between them.
tip = cell(xc, ya, 1, 0)
arrow(ax, (xa + wa + 0.20, cell(xa, ya, 1, 2)[1]),
      (tip[0] - CW / 2 - 0.10, tip[1] - 0.10), TEAL, rad=0.32)
arrow(ax, (cell(xb, ya + 0.36, 2, 0)[0], ya + 0.36 - hb - 0.22),
      (tip[0] - 0.10, tip[1] - CH / 2 - 0.10), INDIGO, rad=0.22)

ax.annotate("row 2", (xa - 0.30, cell(xa, ya, 1, 0)[1]), ha="right",
            va="center", fontsize=11, color=TEAL)
ax.annotate("column 1", (cell(xb, ya + 0.36, 0, 0)[0], ya + 0.52), ha="center",
            va="bottom", fontsize=11, color=INDIGO)

ax.annotate(r"$4 \times 7 \; + \; 5 \times 9 \; + \; 6 \times 11 \;=\; 139$",
            (xa - 0.8, -2.35), ha="left", va="center", fontsize=13, color=ROSE)
ax.annotate("each pair is one step along the row and one step down the column",
            (xa - 0.8, -2.85), ha="left", va="center", fontsize=10, color=MUTED)
ax.annotate("2 by 3 times 3 by 2 gives 2 by 2, and the two 3s must agree",
            (xa - 0.8, -3.22), ha="left", va="center", fontsize=10, color=MUTED)
save(fig, "product.svg")


# A vector as an arrow, its length by Pythagoras, and the unit vector along
# the same direction. The 3-4-5 triangle again, lying down this time: it
# says the same thing in half the height, and a figure is as tall as the
# page has to make room for.
v = np.array([4.0, 3.0])
n = float(np.hypot(*v))
u = v / n
assert n == 5.0 and u.tolist() == [0.8, 0.6]

ax = plane((-0.6, 9.4), (-0.75, 3.65), size=(5.6, 2.45))
ax.plot([0, v[0]], [0, 0], lw=2.4, color=AMBER, zorder=2)
ax.plot([v[0], v[0]], [0, v[1]], lw=2.4, color=INDIGO, zorder=2)
t = np.linspace(0, np.pi / 2, 200)
ax.plot(np.cos(t), np.sin(t), lw=1.1, ls=(0, (4, 3)), color=MUTED, zorder=1)
vec(ax, v, ROSE, lw=2.2)
vec(ax, u, TEAL, lw=2.6)
ax.plot(*v, "o", ms=8, color=ROSE, zorder=5)
ax.plot(*u, "o", ms=7, color=TEAL, zorder=5)

ax.annotate("v = (4, 3)", (v[0] + 0.16, v[1]), ha="left", va="center",
            fontsize=12, color=ROSE)
ax.annotate("length 5", (2.05, 1.72), ha="right", va="bottom", fontsize=12,
            color=ROSE)
ax.annotate("4", (2.0, -0.16), ha="center", va="top", fontsize=11, color=AMBER)
ax.annotate("3", (v[0] + 0.14, 1.5), ha="left", va="center", fontsize=11,
            color=INDIGO)
ax.annotate("(0.8, 0.6)", (1.02, 0.36), ha="left", va="top",
            fontsize=11, color=TEAL)
ax.annotate("16 and 9 make 25, and 5 squared is 25", (5.55, 2.35), ha="left",
            va="center", fontsize=10, color=MUTED)
ax.annotate("dividing v by 5 sets its length to 1", (5.55, 1.95), ha="left",
            va="center", fontsize=10, color=TEAL)
save(plt.gcf(), "vector.svg")


# The dot product, twice. Once between two vectors at an angle, and once
# between two at a right angle, where it comes out zero.
a, b = np.array([4.0, 1.0]), np.array([1.0, 3.0])
c = np.array([-1.0, 4.0])
assert float(a @ b) == 7.0 and float(a @ c) == 0.0
ang = np.degrees(np.arccos(a @ b / (np.hypot(*a) * np.hypot(*b))))
assert abs(ang - 57.5288) < 1e-3, ang

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(7.4, 3.4))
for ax, (p, q, pc, qc, pl, ql) in zip(
        (ax1, ax2), (((a, b, TEAL, INDIGO, "(4, 1)", "(1, 3)")),
                     ((a, c, TEAL, ROSE, "(4, 1)", "(-1, 4)")))):
    plane((-1.75, 5.10), (-1.65, 4.55), ax=ax)
    vec(ax, p, pc)
    vec(ax, q, qc)
    ax.annotate(pl, (p[0] + 0.16, p[1] - 0.10), ha="left", va="top",
                fontsize=11, color=pc)
    ax.annotate(ql, (q[0] + 0.16, q[1] + 0.10), ha="left", va="bottom",
                fontsize=11, color=qc)

t = np.linspace(np.arctan2(*a[::-1]), np.arctan2(*b[::-1]), 200)
ax1.plot(1.5 * np.cos(t), 1.5 * np.sin(t), lw=1.4, color=MUTED, zorder=3)
ax1.annotate("57.5 degrees", (1.70, 1.36), ha="left", va="center", fontsize=10,
             color=MUTED)
ax1.annotate("4(1) + 1(3) = 7", (0.15, -1.35), ha="left", va="center",
             fontsize=11, color=INK)

# The right angle, marked where the two directions meet.
k = 0.42
e1, e2 = a / np.hypot(*a), c / np.hypot(*c)
corner = k * (e1 + e2)
ax2.plot([k * e1[0], corner[0], k * e2[0]], [k * e1[1], corner[1], k * e2[1]],
         lw=1.2, color=MUTED, zorder=3)
ax2.annotate("4(-1) + 1(4) = 0", (0.15, -1.35), ha="left", va="center",
             fontsize=11, color=ROSE)
save(fig, "dot.svg")


# Eigenvectors. The matrix sends the circle of every unit direction to an
# oval, and the two directions that are only stretched are the oval's axes.
A = np.array([[2.0, 1.0], [1.0, 2.0]])
# eigvalsh, not eig: A is symmetric, and the general routine can hand back
# a complex dtype with zero imaginary parts, which is awkward to compare.
assert sorted(np.round(np.linalg.eigvalsh(A), 9).tolist()) == [1.0, 3.0]

ax = plane((-2.70, 5.90), (-2.60, 2.55), size=(5.6, 3.5))
t = np.linspace(0, 2 * np.pi, 500)
ring = np.vstack([np.cos(t), np.sin(t)])
oval = A @ ring
ax.plot(ring[0], ring[1], lw=1.3, ls=(0, (4, 3)), color=MUTED, zorder=2)
ax.plot(oval[0], oval[1], lw=1.5, color=INDIGO, alpha=0.75, zorder=2)

# One ordinary direction, to show that most are turned on the way.
d = np.array([np.cos(np.radians(18)), np.sin(np.radians(18))])
Ad = A @ d
vec(ax, d, MUTED, lw=1.5)
vec(ax, Ad, INDIGO, lw=1.5)
ax.plot([d[0], Ad[0]], [d[1], Ad[1]], lw=1.0, ls=(0, (2, 3)), color=MUTED,
        zorder=3)
ax.annotate("turned as well as stretched", (Ad[0] + 0.18, Ad[1] - 0.10),
            ha="left", va="top", fontsize=11, color=INDIGO)

for e, lam, colour, lab in (
        (np.array([1.0, 1.0]) / np.sqrt(2), 3.0, ROSE, "stretched by 3"),
        (np.array([1.0, -1.0]) / np.sqrt(2), 1.0, TEAL, "left as it was")):
    assert np.allclose(A @ e, lam * e)
    vec(ax, lam * e, colour, lw=2.6)
    ax.plot(*e, "o", ms=7, color=colour, zorder=5)
    ax.plot(*(lam * e), "o", ms=8, color=colour, zorder=5)
ax.annotate("stretched by 3", (2.30, 2.14), ha="left", va="center",
            fontsize=11, color=ROSE)
ax.annotate("left as it was", (0.86, -0.80), ha="left", va="center",
            fontsize=11, color=TEAL)

ax.annotate("the dashed circle is every", (2.55, -1.32), ha="left",
            va="center", fontsize=10, color=MUTED)
ax.annotate("direction of length 1", (2.55, -1.68), ha="left", va="center",
            fontsize=10, color=MUTED)
ax.annotate("the oval is where the matrix", (2.55, -2.14), ha="left",
            va="center", fontsize=10, color=INDIGO)
ax.annotate("sends those directions", (2.55, -2.50), ha="left", va="center",
            fontsize=10, color=INDIGO)
save(plt.gcf(), "eigen.svg")


# Rotation, in two dimensions and then in three. The right panel is drawn
# axonometrically: each axis is given a direction on the page and a point
# is the sum of its three coordinates along them. It is a drawing of three
# dimensions, not a projection anything can be measured off.
def turn(v, deg):
    """The two by two rotation matrix, applied to v."""
    a = np.radians(deg)
    R = np.array([[np.cos(a), -np.sin(a)], [np.sin(a), np.cos(a)]])
    return R @ v


fig, (axA, axB) = plt.subplots(1, 2, figsize=(6.6, 2.7))

# Two dimensions: one matrix, one angle, and the circle it moves along.
plane((-1.65, 1.85), (-0.45, 1.75), ax=axA, step=0.5)
t = np.linspace(0, np.pi, 300)
axA.plot(np.cos(t), np.sin(t), lw=1.1, ls=(0, (4, 3)), color=MUTED, zorder=1)
a0, a1 = 22.0, 78.0
v0 = np.array([np.cos(np.radians(a0)), np.sin(np.radians(a0))])
v1 = turn(v0, a1 - a0)
assert np.allclose(np.hypot(*v1), 1.0)
vec(axA, v0, TEAL, lw=2.4)
vec(axA, v1, ROSE, lw=2.4)
t = np.linspace(np.radians(a0), np.radians(a1), 160)
axA.plot(0.52 * np.cos(t), 0.52 * np.sin(t), lw=1.5, color=INDIGO, zorder=3)
mid_a = np.radians((a0 + a1) / 2)
axA.annotate("θ", (0.62 * np.cos(mid_a), 0.62 * np.sin(mid_a)),
             ha="left", va="bottom", fontsize=12, color=INDIGO)
axA.annotate("v", (v0[0] + 0.10, v0[1] - 0.04), ha="left", va="top",
             fontsize=12, color=TEAL)
axA.annotate("Rv", (v1[0] - 0.06, v1[1] + 0.08), ha="right", va="bottom",
             fontsize=12, color=ROSE)
axA.annotate("length unchanged", (-1.60, -0.30), ha="left",
             va="center", fontsize=10, color=MUTED)

# Three dimensions: the same turn, about the axis that stays still.
EX, EY, EZ = np.array([1.0, -0.30]), np.array([0.56, 0.42]), np.array([0.0, 1.0])


def to2d(x, y, z=0.0):
    return x * EX + y * EY + z * EZ


axB.set_aspect("equal")
axB.set_xticks([])
axB.set_yticks([])
for sp in axB.spines.values():
    sp.set_visible(False)
for end, lab, off in ((to2d(2.5, 0), "x", (0.10, -0.10)),
                      (to2d(0, 2.5), "y", (0.10, 0.04)),
                      (to2d(0, 0, 2.0), "z", (0.06, 0.10))):
    axB.annotate("", xy=tuple(end), xytext=(0, 0),
                 arrowprops=dict(arrowstyle="-|>", lw=1.3, color=MUTED,
                                 shrinkA=0, shrinkB=0), zorder=2)
    axB.annotate(lab, (end[0] + off[0], end[1] + off[1]), ha="left",
                 va="center", fontsize=11, color=MUTED)

r3 = 1.7
p0 = to2d(r3 * np.cos(np.radians(a0)), r3 * np.sin(np.radians(a0)))
p1 = to2d(r3 * np.cos(np.radians(a1)), r3 * np.sin(np.radians(a1)))
t = np.linspace(np.radians(a0), np.radians(a1), 160)
arc = np.array([to2d(0.95 * np.cos(u), 0.95 * np.sin(u)) for u in t])
axB.plot(arc[:, 0], arc[:, 1], lw=1.5, color=INDIGO, zorder=3)
axB.annotate("", xy=tuple(p0), xytext=(0, 0),
             arrowprops=dict(arrowstyle="-|>", lw=2.4, color=TEAL,
                             shrinkA=0, shrinkB=0), zorder=4)
axB.annotate("", xy=tuple(p1), xytext=(0, 0),
             arrowprops=dict(arrowstyle="-|>", lw=2.4, color=ROSE,
                             shrinkA=0, shrinkB=0), zorder=4)
axB.annotate("v", (p0[0] + 0.10, p0[1] - 0.06), ha="left", va="top",
             fontsize=12, color=TEAL)
axB.annotate("Rv", (p1[0] + 0.10, p1[1] + 0.04), ha="left", va="bottom",
             fontsize=12, color=ROSE)
mid3 = to2d(1.05 * np.cos(mid_a), 1.05 * np.sin(mid_a))
axB.annotate("θ", (mid3[0] + 0.06, mid3[1] - 0.02), ha="left",
             va="center", fontsize=12, color=INDIGO)
axB.annotate("turning about z leaves z alone", (-1.35, -1.05), ha="left",
             va="center", fontsize=10, color=MUTED)
axB.set_xlim(-1.4, 3.1)
axB.set_ylim(-1.25, 2.25)
save(fig, "rotation.svg")
