"""Draw the three structural figures: the loop, the topologies, the vector.

The first two are block diagrams and carry no calculation. The third is
computed: it runs the anchor construction and the two spherical
interpolations of src/dsp/MatrixMorph.hpp at N=6, K=6, mode 0, so the loop
drawn is the loop the late-field matrix actually travels.

Writes chain.svg, mixing.svg, loop.svg, topology.svg and vector.svg
beside this file.
"""
import sys

# figtheme sits beside this file; its bytecode would land in the page
# bundle and be published with it.
sys.dont_write_bytecode = True

import numpy as np
import matplotlib

matplotlib.use("Agg")
matplotlib.rcParams["svg.fonttype"] = "path"
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch, Circle

from figtheme import (TEAL, INDIGO, ROSE, AMBER, INK, MUTED, GRID, PAPER,
                      LINES, save)


def canvas(size, xlim, ylim):
    """A plain drawing area with no axes, in figure-independent units."""
    fig, ax = plt.subplots(figsize=size)
    ax.set_xlim(*xlim)
    ax.set_ylim(*ylim)
    ax.set_aspect("equal")
    ax.axis("off")
    return fig, ax


def box(ax, x, y, w, h, text, colour=INK, fontsize=8.5):
    """One labelled block. x and y give its centre."""
    ax.add_patch(FancyBboxPatch((x - w / 2, y - h / 2), w, h,
                                boxstyle="round,pad=0.04,rounding_size=0.12",
                                linewidth=1.1, edgecolor=colour,
                                facecolor=PAPER, zorder=2))
    ax.text(x, y, text, ha="center", va="center", color=colour,
            fontsize=fontsize, zorder=3, linespacing=1.35)


def arrow(ax, start, end, colour=MUTED, style="-|>", rad=0.0, lw=1.0):
    """An arrow that begins and ends exactly where it is told to.

    FancyArrowPatch defaults to shrinkA=2 and shrinkB=2, measured in points,
    so both ends are pulled in. On these figures that is about 0.07 of a
    unit, which is enough for a head to stop visibly short of the box it is
    meant to touch. Both are zero here and the coordinates say everything.
    """
    ax.add_patch(FancyArrowPatch(start, end, arrowstyle=style,
                                 mutation_scale=9, linewidth=lw,
                                 color=colour, zorder=1,
                                 shrinkA=0, shrinkB=0,
                                 connectionstyle=f"arc3,rad={rad}"))


def finish(fig, name):
    fig.tight_layout(pad=0.2)
    save(fig, name)


def loop():
    """The seven steps the late field runs per sample, in source order.

    The order is the one written in FDN::processImpl in src/dsp/FDN.cpp. The
    output tap comes second, before any damping, which is why the wet signal
    keeps more high end than the signal that stays in the loop.
    """
    fig, ax = canvas((6.8, 3.4), (0, 13.6), (0, 6.6))

    box(ax, 2.7, 5.2, 3.2, 0.9, "six delay lines\nread", INK)
    box(ax, 7.2, 5.2, 3.0, 0.9, "output taps C\n2 by 6", TEAL)
    box(ax, 11.4, 5.2, 3.0, 0.9, "wet output\nL and R", TEAL)
    box(ax, 2.7, 3.3, 3.2, 0.9, "three-band\ndamping", INK)
    box(ax, 6.8, 3.3, 3.0, 0.9, "cubic\nsaturation", INK)
    box(ax, 11.0, 3.3, 3.4, 0.9, "add room tone\n(pink noise)", AMBER)
    box(ax, 10.2, 1.4, 4.6, 0.9, "mixing matrix A(t)\nI minus 2 v v-transpose",
        ROSE)
    box(ax, 4.3, 1.4, 3.8, 0.9, "add input B\nwrite and advance", INDIGO)

    arrow(ax, (4.3, 5.2), (5.7, 5.2))
    arrow(ax, (8.7, 5.2), (9.9, 5.2), TEAL)
    arrow(ax, (2.7, 4.75), (2.7, 3.75))
    arrow(ax, (4.3, 3.3), (5.3, 3.3))
    arrow(ax, (8.3, 3.3), (9.3, 3.3))
    arrow(ax, (11.0, 2.85), (11.0, 1.85))
    arrow(ax, (7.9, 1.4), (6.2, 1.4), ROSE)

    # The write-back rail, up the left margin and into the delay lines. The
    # vertical leg sits at 0.35 rather than 0.6 so the arrow that finishes it
    # has 0.75 to travel and reads as a line rather than as a lone head.
    ax.plot([2.4, 0.35, 0.35], [1.4, 1.4, 5.2], color=INDIGO, lw=1.0, zorder=1)
    arrow(ax, (0.35, 5.2), (1.1, 5.2), INDIGO)

    ax.text(0.6, 6.2, "one sample of the late field", color=MUTED,
            fontsize=8.5, ha="left")
    finish(fig, "loop.svg")


def chain():
    """The whole signal path, which the loop figure only shows the end of.

    Read off Reverb::process in src/dsp/Reverb.cpp, including three things
    worth getting right. The dry copy is taken at line 442, before the low
    cut on 443, so it really is untouched. Width acts on the wet sum and
    not on the dry. And the early tap is the plain early output e0, while
    the late field is fed e0 crossfaded with the diffused copy, which is
    what the Diffusion control does.

    Two rows, because six stages in one leave no room for a line between
    them, and a stage joined by nothing but an arrowhead reads as a list.
    """
    fig, ax = canvas((7.0, 3.0), (0, 16.6), (0.55, 7.65))

    top, bot, lane = 6.3, 1.9, 4.0
    plain, field = 1.7, 1.9
    xs = [1.3, 3.9, 6.5, 9.1, 11.7, 14.3]
    box(ax, xs[0], top, plain, 1.0, "in", INK)
    box(ax, xs[1], top, plain, 1.0, "low cut", INK)
    box(ax, xs[2], top, plain, 1.0, "pre-delay", INK)
    box(ax, xs[3], top, field, 1.0, "early field\n3 lines", TEAL, 8.0)
    box(ax, xs[4], top, field, 1.0, "allpass\n8 sections", ROSE, 8.0)
    box(ax, xs[5], top, field, 1.0, "late field\n6 lines", INDIGO, 8.0)
    half = [plain, plain, plain, field, field, field]
    for i in range(5):
        arrow(ax, (xs[i] + half[i] / 2 + 0.05, top),
                  (xs[i + 1] - half[i + 1] / 2 - 0.05, top))

    box(ax, 3.2, bot, plain, 1.0, "mix", INK)
    box(ax, 7.2, bot, 2.4, 1.0, "output gain", INK)
    box(ax, 11.2, bot, plain, 1.0, "out", INK)
    arrow(ax, (4.1, bot), (5.95, bot))
    arrow(ax, (8.45, bot), (10.3, bot))

    # the wet return: the late field drops into a lane, the early tap joins
    # it, the sum passes through Width, and the result enters the mix
    wide = 1.6
    wl, wr = 7.5 - wide / 2, 7.5 + wide / 2
    box(ax, 7.5, lane, wide, 0.85, "width", INDIGO, 8.0)
    # The late field drops into the lane and the early tap joins it at the
    # junction, so the run from there into Width is the arrow itself rather
    # than a short cap laid over the end of a line.
    ax.plot([xs[5], xs[5], xs[3]], [top - 0.5, lane, lane],
            color=INDIGO, lw=1.1, zorder=1)
    arrow(ax, (xs[3], lane), (wr, lane), INDIGO, lw=1.1)
    ax.plot([wl, 3.2, 3.2], [lane, lane, bot + 0.95],
            color=INDIGO, lw=1.1, zorder=1)
    arrow(ax, (3.2, bot + 0.95), (3.2, bot + 0.5), INDIGO, lw=1.1)
    ax.plot([xs[3], xs[3]], [top - 0.5, lane], color=TEAL, lw=1.1, zorder=1)
    ax.add_patch(Circle((xs[3], lane), 0.09, facecolor=INDIGO,
                        edgecolor="none", zorder=3))
    ax.text(xs[3] + 0.22, top - 1.35, "under Early Level", color=TEAL,
            fontsize=8)
    ax.text(9.6, lane + 0.28, "wet: the early tap added to the late field",
            color=INDIGO, fontsize=8)

    # the dry signal, taken before the low cut and mixed back at the end
    ax.plot([xs[0], xs[0], 3.2], [top - 0.5, 1.0, 1.0],
            color=MUTED, lw=1.1, zorder=1)
    arrow(ax, (3.2, 1.0), (3.2, bot - 0.5), MUTED, lw=1.1)
    ax.text(4.0, 0.78, "dry, taken before the low cut and never filtered",
            color=MUTED, fontsize=8)

    # Diffusion is a crossfade between two ways into the late field
    ax.plot([10.4, 10.4, 13.0], [top + 0.5, top + 1.0, top + 1.0],
            color=ROSE, lw=1.1, ls=(0, (3, 2)), zorder=1)
    arrow(ax, (13.0, top + 1.0), (13.0, top + 0.5), ROSE, lw=1.1)
    ax.text(11.7, top + 1.14,
            "Diffusion sets how much of each way in", color=ROSE,
            fontsize=8, ha="center")
    save(fig, "chain.svg")


def mixing():
    """What the mixing matrix does to six numbers, at both ends of Diffusion.

    The article proves that A preserves length and never shows the table or
    a multiplication. This does both, with an input of one line carrying
    everything, so the spread is visible and the total is checkable: at
    Diffusion 1 the output is 2/3 and five of -1/3, whose squares sum to
    4/9 + 5/9, which is 1 again.

    Cells are filled at an opacity set by the size of the entry, and the
    colour carries the sign. Matplotlib writes the opacity separately from
    the fill, so the two hexes still become the theme's own properties.
    """
    fig, ax = canvas((6.9, 2.9), (0, 10.4), (0, 4.35))
    cell = 0.42

    def grid(ox, oy, values, cols, fontsize=6.5):
        rows = len(values) // cols
        for r in range(rows):
            for c in range(cols):
                v = values[r * cols + c]
                x, y = ox + c * cell, oy - r * cell
                if v:
                    ax.add_patch(FancyBboxPatch(
                        (x, y - cell), cell, cell,
                        boxstyle="square,pad=0", linewidth=0.6,
                        edgecolor=GRID,
                        facecolor=TEAL if v > 0 else ROSE,
                        alpha=min(1.0, abs(v)), zorder=2))
                else:
                    ax.add_patch(FancyBboxPatch(
                        (x, y - cell), cell, cell,
                        boxstyle="square,pad=0", linewidth=0.6,
                        edgecolor=GRID, facecolor="none", zorder=2))
                if v:
                    # A cell is filled at an opacity set by the size of its
                    # entry, so the strong ones need the light text and the
                    # faint ones the dark. PAPER is the theme's --surface,
                    # which is white in the light theme and dark in the dark
                    # one, so the contrast holds either way round.
                    ink = PAPER if abs(v) >= 0.5 else INK
                    ax.text(x + cell / 2, y - cell / 2, label(v), ha="center",
                            va="center", fontsize=fontsize, color=ink, zorder=4)

    def label(v):
        if v == int(v):
            return f"{int(v)}"
        return "2/3" if abs(v - 2 / 3) < 1e-9 else ("-1/3" if v < 0 else "1/3")

    def panel(ox, title, diag, off, note):
        # diag is the whole diagonal, not one value repeated. At Diffusion 0
        # it reads (-1, 1, 1, 1, 1, 1), because A = I - 2vv' with v on the
        # first axis flips that axis and leaves the other five alone. Drawing
        # -1 six times would be -I, which is a different matrix.
        top = 3.35
        a = [diag[r] if r == c else off for r in range(6) for c in range(6)]
        x_in = [1, 0, 0, 0, 0, 0]
        y_out = [diag[0], off, off, off, off, off]
        ax.text(ox + 3 * cell, top + 0.42, title, color=MUTED, fontsize=8.5,
                ha="center")
        grid(ox, top, a, 6)
        ax.text(ox + 6 * cell + 0.22, top - 3 * cell, "\u00d7", color=MUTED,
                fontsize=10, ha="center", va="center")
        grid(ox + 6 * cell + 0.44, top, x_in, 1)
        ax.text(ox + 7 * cell + 0.66, top - 3 * cell, "=", color=MUTED,
                fontsize=10, ha="center", va="center")
        grid(ox + 7 * cell + 0.88, top, y_out, 1)
        ax.text(ox + 3 * cell, top - 6 * cell - 0.28, note, color=MUTED,
                fontsize=8, ha="center")

    panel(0.35, "Diffusion 1: every line into every other",
          [2 / 3] * 6, -1 / 3, "one line in, all six out, length still 1")
    panel(5.55, "Diffusion 0: nothing mixes",
          [-1] + [1] * 5, 0, "one line in, the same one out, length still 1")
    ax.plot([5.2, 5.2], [0.35, 4.0], color=GRID, lw=1.0)
    save(fig, "mixing.svg")


def topology():
    """Thirty-six delay lines in series against six in one loop.

    Left: the Airwindows late field, six diffusion stages of six lines each,
    the output of the sixth feeding back to the first. Right: the SD-Reverb
    late field, one stage of six lines.

    The six stages are spaced 1.5 apart rather than 1.0, which is what makes
    the connector between them 0.7 long instead of 0.26. At 0.26 the
    arrowhead was the whole arrow and the stages read as a row of separate
    things rather than as a path. The width for it comes from the right
    panel, where a box sized for three lines held two.
    """
    fig, ax = canvas((6.8, 2.31), (0, 15.0), (1.75, 6.85))

    def bank(x, colour):
        """Six delay lines drawn as one column of dots inside a frame."""
        # The frame has to hold six dots at 0.46 apart, so its span is
        # 5 * 0.46 plus a radius at each end plus the clearance. Drawn from
        # the dots rather than typed: the first version ran 3.3 to 5.9 and
        # the sixth dot, centred on 3.30, hung 0.11 below the line.
        top, step, r, gap = 5.6, 0.46, 0.11, 0.19
        low = top - 5 * step
        ax.add_patch(FancyBboxPatch((x - 0.32, low - r - gap), 0.64,
                                    (top - low) + 2 * (r + gap),
                                    boxstyle="round,pad=0.03,rounding_size=0.1",
                                    linewidth=1.0, edgecolor=colour,
                                    facecolor=PAPER, zorder=2))
        for i in range(6):
            ax.add_patch(Circle((x, top - i * step), r, facecolor=colour,
                                edgecolor="none", zorder=3))

    # The first stage sits well clear of the left edge so the return has a
    # run to travel along. An arrow with 0.2 to cover is an arrowhead and
    # nothing else, which is what the earlier version drew. The stages are
    # 1.4 apart rather than 1.5 to pay for it, and 0.6 between them is still
    # a line the eye follows.
    left = [1.7 + s * 1.4 for s in range(6)]
    ax.text(4.9, 6.5, "six stages, thirty-six lines", color=MUTED,
            fontsize=9, ha="center")
    for i, x in enumerate(left):
        bank(x, INDIGO)
        if i < 5:
            arrow(ax, (x + 0.40, 4.6), (x + 1.00, 4.6), MUTED, lw=0.9)
    ax.plot([left[-1] + 0.40, 9.5, 9.5, 0.35, 0.35],
            [4.6, 4.6, 2.6, 2.6, 4.6], color=MUTED, lw=0.9, zorder=1)
    arrow(ax, (0.35, 4.6), (1.35, 4.6), MUTED, lw=0.9)
    ax.text(4.9, 2.1, "the same matrix, applied six times",
            color=MUTED, fontsize=8.5, ha="center")

    ax.plot([10.0, 10.0], [1.9, 6.7], color=GRID, lw=1.0)

    ax.text(12.4, 6.5, "one stage, six lines", color=MUTED,
            fontsize=9, ha="center")
    bank(11.5, ROSE)
    box(ax, 13.5, 4.6, 2.0, 0.9, "A(t), rebuilt\nonce per block", ROSE, 8.0)
    arrow(ax, (11.85, 4.6), (12.45, 4.6), MUTED, lw=0.9)
    ax.plot([13.5, 13.5, 10.35, 10.35], [4.15, 2.6, 2.6, 4.6],
            color=MUTED, lw=0.9, zorder=1)
    arrow(ax, (10.35, 4.6), (11.15, 4.6), MUTED, lw=0.9)
    ax.text(12.4, 2.1, "one matrix, moving", color=MUTED,
            fontsize=8.5, ha="center")

    finish(fig, "topology.svg")


# ── the trajectory, computed exactly as src/dsp/MatrixMorph.hpp computes it ──

def normalise(v):
    return v / np.linalg.norm(v)


def make_vector(balance, n=6):
    """The base reflection vector. balance 1 gives the even one."""
    bal = 1.0 / np.sqrt(n)
    v = np.full(n, balance * bal)
    v[0] = 1.0 + balance * (bal - 1.0)
    return normalise(v)


def anchors(n=6, k=6, mode=0):
    """The k perturbation directions, zero-mean, one per loop corner."""
    mode_phase = 0.7 * mode
    out = []
    for j in range(k):
        a = np.array([np.sin(2.0 * np.pi * (j + 1) * (i + 1) / n
                             + mode_phase + 0.31 * j) for i in range(n)])
        out.append(a - a.mean())
    return out


def slerp(a, b, t):
    """Interpolate along the great circle joining two unit vectors."""
    dot = float(np.clip(np.dot(a, b), -1.0, 1.0))
    theta = np.arccos(dot)
    if theta < 1e-5:
        return normalise(a + t * (b - a))
    s = np.sin(theta)
    return (np.sin((1.0 - t) * theta) / s) * a + (np.sin(t * theta) / s) * b


def trajectory(depth, spread=0.35, n=6, k=6, steps=400):
    """The loop, and the depth-blended point actually handed to the matrix."""
    r = make_vector(1.0, n)
    corners = [normalise(r + spread * a) for a in anchors(n=n, k=k)]
    full, blended = [], []
    for step in range(steps + 1):
        phase = step / steps
        fp = phase * k
        i = int(fp) % k
        point = slerp(corners[i], corners[(i + 1) % k], fp - np.floor(fp))
        full.append(point)
        blended.append(slerp(r, point, depth))
    return r, corners, np.array(full), np.array(blended)


def vector():
    """What the six numbers do over one turn of the loop.

    The upper panel plots every component of v against loop phase at Morph
    Depth 1. The lower panel plots the angle between v and r, which is the
    one quantity a projection of a five-dimensional sphere would keep exact.
    """
    r, corners, full, blended = trajectory(depth=0.4, steps=300)
    phase = np.linspace(0.0, 1.0, full.shape[0])

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6.4, 4.2),
                                   height_ratios=[3, 2], sharex=True)
    for ax in (ax1, ax2):
        ax.grid(True, color=GRID, lw=0.6, zorder=0)
        ax.set_axisbelow(True)
        for side in ("top", "right"):
            ax.spines[side].set_visible(False)
        for side in ("left", "bottom"):
            ax.spines[side].set_color(MUTED)
        ax.tick_params(which="both", colors=MUTED, labelsize=8)

    # The loop passes through its own base once per turn. With K = N the
    # last anchor is constant across i, so the mean subtraction zeroes it
    # and corner K-1 is r itself. Marking the phase says where to look.
    base_phase = (len(anchors()) - 1) / len(anchors())

    for i, colour in enumerate(LINES):
        ax1.plot(phase, full[:, i], color=colour, lw=1.2)
    ax1.axhline(float(r[0]), color=MUTED, lw=0.9, ls=(0, (4, 3)))
    ax1.set_ylim(0.0, 0.74)
    ax1.text(0.015, 0.685, "r, every component 0.408",
             color=MUTED, fontsize=8)
    ax1.set_ylabel("component of v", color=INK, fontsize=9)

    for points, colour, label in ((full, ROSE, "Morph Depth 1"),
                                  (blended, INDIGO, "Morph Depth 0.4")):
        angle = np.degrees(np.arccos(np.clip(points @ r, -1.0, 1.0)))
        ax2.plot(phase, angle, color=colour, lw=1.4, label=label)
    ax2.set_ylabel("angle from r\n(degrees)", color=INK, fontsize=9)
    ax2.set_xlabel("position around the loop", color=INK, fontsize=9)
    ax2.set_xlim(0.0, 1.0)
    ax2.set_ylim(0.0, 36.0)
    for ax, top in ((ax1, 0.74), (ax2, 36.0)):
        ax.axvline(base_phase, color=MUTED, lw=0.9, ls=(0, (2, 3)), zorder=1)
    ax1.text(base_phase - 0.008, 0.685, "v = r", color=MUTED, fontsize=8,
             ha="right")
    leg = ax2.legend(frameon=False, fontsize=8.5, loc="lower center", ncols=2)
    for text in leg.get_texts():
        text.set_color(INK)

    fig.tight_layout(pad=0.4)
    save(fig, "vector.svg")


if __name__ == "__main__":
    chain()
    mixing()
    loop()
    topology()
    vector()
