"""Draw the trigonometry a complex number needs, and its polar form.

Five figures on one ruled plane, so the paper a reader learns in the first
is the paper the last argues on. One unit per square, five fine squares to
each ruled one, the way squared paper is printed.

The 3-4-5 triangle carries the first three figures. It is the smallest
right-angled triangle whose sides are all whole numbers, so a reader can
check every length by counting squares rather than trusting the drawing.

Writes triangle.svg, circle.svg, modulus.svg, polar.svg and
multiplication.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")



# A right-angled triangle, which is where sine and cosine are defined.
# Sides 3, 4 and 5, so every length can be checked by counting squares.
fig, ax = grid((-0.7, 6.0), (-1.9, 3.9), (7.0, 5.2), ticks=False,
            named=False)
ax.plot([0, 4], [0, 0], lw=2.6, color=AMBER, zorder=2)
ax.plot([4, 4], [0, 3], lw=2.6, color=INDIGO, zorder=2)
ax.plot([0, 4], [0, 3], lw=2.2, color=ROSE, zorder=3)

k = 0.22
ax.plot([4 - k, 4 - k, 4], [0, k, k], lw=1.1, color=MUTED, zorder=3)
mx, my = arc(ax, np.arctan2(3, 4), 0.85, TEAL)
ax.annotate("\u03b8", (mx + 0.10, my + 0.10), ha="left", va="bottom",
            fontsize=12, color=TEAL)

ax.annotate("adjacent, 4", (2.0, -0.16), ha="center", va="top", fontsize=11,
            color=AMBER)
ax.annotate("opposite, 3", (4.16, 1.5), ha="left", va="center", fontsize=11,
            color=INDIGO)
ax.annotate("hypotenuse, 5", (1.7, 1.5), ha="right", va="bottom", fontsize=12,
            color=ROSE)
ax.annotate("cos \u03b8 = adjacent / hypotenuse = 4 / 5 = 0.8",
            (0.15, -1.32), ha="left", va="center", fontsize=10, color=MUTED)
ax.annotate("sin \u03b8 = opposite / hypotenuse = 3 / 5 = 0.6",
            (0.15, -1.66), ha="left", va="center", fontsize=10, color=MUTED)
save(fig, "triangle.svg")


# The unit circle, where the two ratios become two coordinates. The arc is
# drawn heavily because its length is the angle itself, measured in radians.
th = np.deg2rad(50)
c, s = np.cos(th), np.sin(th)
fig, ax = grid((-1.8, 2.1), (-1.7, 1.9), (6.8, 5.6), step=0.5,
            ticks=False, named=False)
t = np.linspace(0, 2 * np.pi, 400)
ax.plot(np.cos(t), np.sin(t), lw=1.4, color=MUTED, zorder=2)

# Both ratios are drawn on the axes they measure, with the point dropped
# onto each by a dashed line. That keeps the right of the circle clear for
# the arc, which is the one thing this figure exists to show.
ax.plot([c, c], [s, 0], lw=1.0, ls=(0, (4, 3)), color=MUTED, zorder=2)
ax.plot([c, 0], [s, s], lw=1.0, ls=(0, (4, 3)), color=MUTED, zorder=2)
ax.plot([0, c], [0, 0], lw=2.8, color=AMBER, zorder=3)
ax.plot([0, 0], [0, s], lw=2.8, color=INDIGO, zorder=3)
arrow(ax, complex(c, s), ROSE, lw=1.8)
ax.plot(c, s, "o", ms=8, color=ROSE, zorder=4)

t = np.linspace(0, th, 200)
ax.plot(np.cos(t), np.sin(t), lw=3.6, color=TEAL, zorder=3)
ax.annotate("the arc is \u03b8", (1.14 * np.cos(np.deg2rad(47)),
            1.14 * np.sin(np.deg2rad(47))), ha="left", va="bottom",
            fontsize=11, color=TEAL)
ax.annotate("cos \u03b8", (c / 2, -0.14), ha="center", va="top",
            fontsize=11, color=AMBER)
ax.annotate("sin \u03b8", (-0.12, s / 2), ha="right", va="center",
            fontsize=11, color=INDIGO)

# The label for the radius goes at right angles to it, clear of the line.
q = np.deg2rad(140)
ax.annotate("1", (0.55 * c + 0.11 * np.cos(q), 0.55 * s + 0.11 * np.sin(q)),
            ha="center", va="center", fontsize=12, color=ROSE)

# Three turns, named exactly. The angle in radians is a length on this
# circle, so a whole turn is its circumference and that is 2 pi.
for x, y, lab, ha, va, dx, dy in (
        (1, 0, "0", "left", "top", 0.06, -0.08),
        (0, 1, "a quarter turn, \u03c0/2", "left", "bottom", 0.10, 0.06),
        (-1, 0, "a half turn, \u03c0", "right", "top", -0.06, -0.10)):
    ax.plot(x, y, "o", ms=5, color=MUTED, zorder=4)
    ax.annotate(lab, (x + dx, y + dy), ha=ha, va=va, fontsize=10, color=MUTED)

ax.annotate("a whole turn is 2\u03c0, the circumference of the circle",
            (-1.75, -1.42), ha="left", va="center", fontsize=10, color=MUTED)
save(fig, "circle.svg")

# The modulus, as the hypotenuse of the triangle the coordinates make.
# 3 + 4i is the example the page works, and its sides land on the ruling.
z = 3 + 4j
fig, ax = grid((-1.4, 4.8), (-1.2, 5.2), (6.8, 5.8))
ax.plot([0, z.real], [0, 0], lw=2.6, color=AMBER, zorder=2)
ax.plot([z.real, z.real], [0, z.imag], lw=2.6, color=INDIGO, zorder=2)
ax.plot([0, z.real], [0, z.imag], lw=2.2, color=ROSE, zorder=3)

# The right angle, marked the way it is on paper.
k = 0.22
ax.plot([z.real - k, z.real - k, z.real], [0, k, k], lw=1.1, color=MUTED, zorder=3)

ax.plot(z.real, z.imag, "o", ms=8, color=ROSE, zorder=4)
ax.annotate("3 + 4i", (z.real + 0.18, z.imag), ha="left", va="center",
            fontsize=13, color=ROSE)
ax.annotate("a = 3", (z.real / 2, 0.16), ha="center", va="bottom",
            fontsize=11, color=AMBER)
ax.annotate("b = 4", (z.real + 0.16, z.imag / 2), ha="left", va="center",
            fontsize=11, color=INDIGO)
ax.annotate("modulus 5", (z.real / 2 - 0.5, z.imag / 2 + 0.3), ha="right",
            va="bottom", fontsize=12, color=ROSE)
ax.annotate("3 squared is 9, and 4 squared is 16", (0.35, -0.72), ha="left",
            va="center", fontsize=10, color=MUTED)
ax.annotate("9 and 16 make 25, and 5 squared is 25", (0.35, -1.02), ha="left",
            va="center", fontsize=10, color=MUTED)
save(fig, "modulus.svg")


# Polar form: the same point named by a distance and a turn.
z = 3 + 2j
r, th = abs(z), np.angle(z)
fig, ax = grid((-1.6, 4.6), (-1.4, 3.4), (7.0, 5.2))
ax.plot([z.real, z.real], [0, z.imag], lw=1.2, ls=(0, (4, 3)), color=MUTED, zorder=2)
ax.plot([0, z.real], [0, 0], lw=2.2, color=AMBER, zorder=2)
arrow(ax, z, ROSE, lw=1.8)
ax.plot(z.real, z.imag, "o", ms=8, color=ROSE, zorder=4)
mx, my = arc(ax, th, 1.15, TEAL)
ax.annotate("theta, about 33.7 degrees", (mx + 0.12, my + 0.06), ha="left",
            va="bottom", fontsize=11, color=TEAL)
ax.annotate("r = 3.61", (z.real / 2 - 0.35, z.imag / 2 + 0.28), ha="right",
            va="bottom", fontsize=12, color=ROSE)
ax.annotate("a = 3", (z.real / 2, 0.14), ha="center", va="bottom", fontsize=11, color=AMBER)
ax.annotate("b = 2", (z.real + 0.14, z.imag / 2), ha="left", va="center",
            fontsize=11, color=MUTED)
ax.annotate("3 + 2i", (z.real + 0.18, z.imag + 0.16), ha="left", va="bottom",
            fontsize=13, color=ROSE)
ax.annotate("two coordinates, or one distance and one turn",
            (-1.5, -1.1), ha="left", va="center", fontsize=10, color=MUTED)
save(fig, "polar.svg")


# Multiplication, as moduli multiplied and arguments added.
w1, w2 = 1 + 1j, 2j
prod = w1 * w2
fig, ax = grid((-3.4, 3.0), (-1.0, 3.8), (7.0, 5.6))
for z, col, lab, off in ((w1, TEAL, "1 + i", (0.18, -0.12)),
                         (w2, INDIGO, "2i", (0.20, 0.0)),
                         (prod, ROSE, "-2 + 2i", (-0.20, 0.18))):
    arrow(ax, z, col, lw=1.8)
    ax.plot(z.real, z.imag, "o", ms=7.5, color=col, zorder=4)
    ha = "right" if off[0] < 0 else "left"
    ax.annotate(lab, (z.real + off[0], z.imag + off[1]), ha=ha, va="center",
                fontsize=12, color=col)

# Each arc starts where the previous one ended, so the addition of the
# two turns is the picture rather than a caption.
mx, my = arc(ax, np.angle(w1), 0.70, TEAL)
ax.annotate("45", (mx + 0.10, my), ha="left", va="center", fontsize=10, color=TEAL)
arc(ax, np.angle(prod), 1.05, INDIGO, start=np.angle(w1))
t = np.deg2rad(68)
ax.annotate("a further 90", (1.05 * np.cos(t) + 0.12, 1.05 * np.sin(t)),
            ha="left", va="center", fontsize=10, color=INDIGO)
arc(ax, np.angle(prod), 1.45, ROSE)
t = np.deg2rad(128)
ax.annotate("135 in all", (1.45 * np.cos(t) - 0.12, 1.45 * np.sin(t) + 0.06),
            ha="right", va="bottom", fontsize=10, color=ROSE)

ax.annotate("moduli multiply: 1.41 times 2 is 2.83", (-3.3, 3.55), ha="left",
            va="center", fontsize=10, color=MUTED)
ax.annotate("arguments add: 45 and 90 make 135", (-3.3, 3.25), ha="left",
            va="center", fontsize=10, color=MUTED)
save(fig, "multiplication.svg")
