Double slit

"""Young's double slit seen side on, with a comb of lines for the fringes on the screen."""

import itertools
import math
from collections.abc import Sequence

import numpy as np
from numpy.typing import NDArray
from shapely.geometry import Point, Polygon
from shapely.ops import unary_union

from walldye import (
    ACCENT,
    ACCENT_1,
    ACCENT_4,
    BG,
    BG_ALT,
    UI,
    UI_ALT,
    UI_HI,
    Buckets,
    Canvas,
    P,
    by_regime,
    design,
    mix,
    smoothstep,
)
from walldye.geom import Affine, parts

WAVE = 24  # wavelength: the spacing of the crest arcs
SLIT = 104  # slit width
SEP = 4 * SLIT  # slit pitch; every fourth order falls on an envelope zero and goes missing
DROP = 820  # barrier to screen
SPAN = 426  # half the screen's length: it ends on the envelope's second zero
RINGS = int(math.hypot(DROP, SPAN + SEP) / WAVE)  # crests from each slit that reach the screen
HALF = 4  # half the width of a crest band where two crests overlap
ARCS = (BG_ALT, mix(BG_ALT, UI, 0.5))
# a step lower on paper, where UI_HI beads would outweigh the comb
BEADS = (
    by_regime(BG_ALT, mix(BG, BG_ALT, 0.6)),
    by_regime(UI, BG_ALT),
    by_regime(UI_ALT, UI),
    by_regime(UI_HI, UI_ALT),
)
COMB = (UI, ACCENT_4, ACCENT_1, ACCENT)
CUTS = (0.3, 0.6, 0.85)  # COMB thresholds on the square root of the intensity
PITCH, PEAK = 4, 190  # comb line pitch; the central fringe's length


def envelope(sin: NDArray[np.float64]) -> NDArray[np.float64]:
    """Single-slit intensity for the sine of the angle off the axis, 1 on it."""
    return np.sinc(SLIT * sin / WAVE) ** 2


def intensity(v: NDArray[np.float64]) -> NDArray[np.float64]:
    """Two-slit intensity at height `v` on the screen: the exact path difference sets the
    fringes, the single-slit envelope their brightness."""
    diff = np.hypot(DROP, v + SEP / 2) - np.hypot(DROP, v - SEP / 2)
    return np.cos(math.pi * diff / WAVE) ** 2 * envelope(v / np.hypot(DROP, v))


def graded(b: Buckets, pts: NDArray[np.floating], tones: Sequence[int | None]) -> None:
    """Draw a polyline as runs of equal tone index, each reaching the next run's first
    vertex; single-vertex runs and None runs are left out."""
    i = 0
    for tone, run in itertools.groupby(tones):
        n = len(list(run))
        if tone is not None and n > 1:
            b[tone].poly(pts[i : i + n + 1])
        i += n


@design(aspects="any")
def draw(s: Canvas) -> None:
    # Local frame: u runs from the barrier toward the screen, v along the screen; landscape
    # screens send the light rightward, portrait ones downward. `o` is the screen's center.
    o = s.pick(landscape=(0.66, 0.5), portrait=(0.5, 0.62))
    if s.landscape:
        m = Affine(1, 0, 0, 1, o.x - DROP, o.y)
    else:
        m = Affine(0, 1, 1, 0, o.x, o.y - DROP)
    slits = (-SEP / 2, SEP / 2)
    v0 = SEP / 2 + SLIT / 2 + 30
    lit = Polygon([(0, -v0), (DROP, -SPAN), (DROP, SPAN), (0, v0)])

    def inside(pts: NDArray[np.float64]) -> NDArray[np.float64]:
        """How far into the lit trapezoid each point sits, eased to 1 over 70 units."""
        u, v = pts[:, 0], np.abs(pts[:, 1])
        edge = v0 + (SPAN - v0) * u / DROP
        slope = math.hypot(1, (SPAN - v0) / DROP)
        ok = (u >= 0) & (u <= DROP)
        return np.where(ok, smoothstep(0, 70, (edge - v) / slope), 0.0)

    with s.group(transform=m):
        # plane wave arriving at the barrier, its crests fading at the ends and further back
        vs = np.linspace(-v0 - 40, v0 + 40, 60)
        with s.buckets(ARCS, "stroke", stroke_width=1.2) as waves:
            for n in range(1, 7):
                pts = np.column_stack([np.full_like(vs, -n * WAVE), vs])
                f = smoothstep(v0 + 40, v0 - 80, np.abs(vs)) * (1.15 - 0.15 * n)
                graded(waves, pts, [None if t < 0.2 else int(t > 0.6) for t in f])

        # crest arcs from each slit, clipped to the lit region
        with s.buckets(ARCS, "stroke", stroke_width=1.2, stroke_linecap="round") as arcs:
            for sv in slits:
                for n in range(1, RINGS + 1):
                    r = n * WAVE
                    a = np.linspace(-math.pi / 2, math.pi / 2, 24 + 4 * n)
                    pts = np.column_stack([r * np.cos(a), sv + r * np.sin(a)])
                    f = inside(pts)
                    graded(arcs, pts, [None if t < 0.05 else int(t > 0.5) for t in f])

        # where crests from both slits overlap: beads along the bright directions
        crests = []
        for sv in slits:
            o2 = Point(0, sv)
            rings = [
                o2.buffer(n * WAVE + HALF, 64).difference(o2.buffer(n * WAVE - HALF, 64))
                for n in range(1, RINGS + 1)
            ]
            crests.append(unary_union(rings).intersection(lit))
        beads = [g for g in parts(crests[0].intersection(crests[1])) if g.area > 2]
        c = np.array([g.centroid.coords[0] for g in beads])
        # far-field brightness in the bead's direction as seen from between the slits
        glow = envelope(c[:, 1] / np.hypot(c[:, 0], c[:, 1])) ** 0.5 * inside(c)
        with s.buckets(BEADS, "fill") as fills:
            for g, t in zip(beads, glow):
                if t > 0.08:
                    fills[min(3, int(t * 4))].shape(g)

        # the barrier, cut by two slits
        bar = P()
        edges = [-v0 - 60, *(sv + d for sv in slits for d in (-SLIT / 2, SLIT / 2)), v0 + 60]
        for a0, a1 in zip(edges[::2], edges[1::2]):
            bar.rect(-4, a0, 8, a1 - a0)
        s.fill(bar, UI_ALT)

        # the screen, and the fringes it records as a comb of lines whose length and tone
        # follow the intensity, under the dashed single-slit envelope
        s.fill(P().rect(DROP - 1.5, -SPAN - 10, 3, 2 * SPAN + 20), UI_ALT)
        n = int(SPAN // PITCH)
        v = PITCH * np.arange(-n, n + 1, dtype=float)
        i = intensity(v)
        with s.buckets(COMB, "stroke", stroke_width=PITCH / 2) as comb:
            for vv, ii in zip(v, i):
                ln = PEAK * ii
                if ln >= 2:
                    comb[int(np.searchsorted(CUTS, ii**0.5))].M(DROP + 6, vv).H(DROP + 6 + ln)
        v = np.linspace(-SPAN, SPAN, 900)
        env = envelope(v / np.hypot(DROP, v))
        s.stroke(P().poly(np.column_stack([DROP + 6 + PEAK * env, v])), UI, 1.2, dash=(4, 5))