Koi pond

"""Koi seen from above in flat shapes, with ripple rings."""

import math
from typing import NamedTuple

from shapely import Point, Polygon
from shapely.geometry.base import BaseGeometry
from shapely.ops import unary_union

from walldye import (
    ACCENT,
    ACCENT_1,
    ACCENT_2,
    BG,
    BG_ALT,
    UI,
    UI_ALT,
    UI_HI,
    Canvas,
    Color,
    P,
    Vec,
    design,
    mix,
    polar,
)
from walldye.field import Noise
from walldye.geom import Affine, parts

type Patch = tuple[float, float, float, float]


class Koi(NamedTuple):
    head: Vec
    heading: float  # degrees the nose points, screen convention
    length: float  # nose to tail root
    turn: float  # radians the spine bends from nose to tail root
    phase: float  # of the swimming undulation
    half: float  # half-width at the fullest point
    body: Color
    fins: Color
    tail: Color
    patch: Color
    fade: float  # how strongly fins and tail tips show over the water, 0..1
    patches: tuple[Patch, ...]  # (spine from, spine to, radius and sideways drift in half-widths)


RIPPLES = (BG_ALT, mix(BG, BG_ALT, 0.7), mix(BG, BG_ALT, 0.5))  # inner to outer ring


def spine(
    head: Vec, heading: float, length: float, turn: float, phase: float, n: int = 60
) -> tuple[list[Vec], list[float]]:
    """`n + 1` points from nose to tail root, and the direction towards the tail at each in
    screen radians: a steady turn plus a swimming undulation that grows towards the tail."""
    pts, angs = [head], []
    back = math.radians(heading) + math.pi
    for i in range(n):
        t = i / (n - 1)
        a = back + turn * t + 0.28 * t * t * math.sin(phase + t * 2.2 * math.pi)
        angs.append(a)
        pts.append(polar(pts[-1], length / (n - 1), rad=a))
    angs.append(angs[-1])
    return pts, angs


def half_width(t: float, w: float) -> float:
    """Half-width at spine fraction `t`: round head, fullest behind the pectorals, slim peduncle."""
    if t < 0.32:
        return w * (0.62 + 0.38 * math.sin(math.pi / 2 * t / 0.32))
    return w * (1 - 0.8 * ((t - 0.32) / 0.68) ** 1.6)


def fin(root: Vec, a: float, size: float, width: float = 0.34) -> Polygon:
    """Teardrop fin from `root` along `a` (screen radians), broadening to a rounded trailing edge."""
    side = [
        (u * size, width * size * math.sin(math.pi * u) ** 0.6 * (0.45 + 0.55 * u))
        for u in (k / 40 for k in range(41))
    ]
    outline = side + [(x, -y) for x, y in side[::-1]]
    return Polygon(Affine.frame(root, rad=a).apply(outline))


def koi(k: Koi) -> tuple[BaseGeometry, BaseGeometry, BaseGeometry, list[Vec], list[float]]:
    """Body, paired fins and tail of one fish, plus its spine points and directions."""
    pts, angs = spine(k.head, k.heading, k.length, k.turn, k.phase, n=160)
    n = len(pts) - 1
    body = unary_union(
        [
            Point(p).buffer(half_width(i / n, k.half), quad_segs=16)
            for i, p in enumerate(pts)
            if i / n > 0.02
        ]
    )

    # two long slim lobes splayed ~35 degrees, rooted inside the peduncle so the fork doesn't pinch
    end, ea = pts[-1], angs[-1]
    root = polar(end, -0.07 * k.length, rad=ea)
    tail = unary_union([fin(root, ea + side * 0.3, 0.42 * k.length, 0.2) for side in (1, -1)])
    tail = tail.buffer(2, quad_segs=8).buffer(-2, quad_segs=8)

    # mirrored about the local spine direction: pectorals at 25%, pelvics at 55%
    fins: list[Polygon] = []
    for t, size, sweep in ((0.25, 0.22 * k.length, 0.95), (0.55, 0.12 * k.length, 0.75)):
        i = int(t * n)
        d = pts[i + 6] - pts[i - 6]
        a = math.atan2(d.y, d.x)
        for side in (1, -1):
            at = polar(pts[i], half_width(t, k.half) * 0.7, rad=a + side * math.pi / 2)
            fins.append(fin(at, a + side * sweep, size))
    return body, unary_union(fins), tail, pts, angs


def markings(nz: Noise, key: int, pts: list[Vec], k: Koi) -> BaseGeometry:
    """Kohaku-style patches: irregular blobs over the back, each drifting a little off the spine."""
    n = len(pts) - 1
    out = []
    for s0, s1, size, drift in k.patches:
        blobs = []
        for i in range(int(s0 * n), int(s1 * n) + 1):
            t = i / n
            w = half_width(t, k.half)
            swell = max(0.0, math.sin(math.pi * (t - s0) / (s1 - s0))) ** 0.3
            rad = w * size * swell * (1 + 0.4 * nz.fbm(t * 14, key, 3))
            d = pts[min(i + 1, n)] - pts[i - 1]
            off = (drift + 0.35 * nz.fbm(t * 11, key + 3.7, 2)) * w
            if rad > 1:
                at = polar(pts[i], off, rad=math.atan2(d.y, d.x) + math.pi / 2)
                blobs.append(Point(at).buffer(rad, quad_segs=12))
        out.append(unary_union(blobs))
    return unary_union(out)


@design(aspects="any")
def draw(s: Canvas) -> None:
    # landscape: the pair right of center, its outer ripple kept ~120 from the right edge;
    # portrait: the pair low and to the right, under a clock
    if s.landscape:
        lit, stray = Vec(min(0.615 * s.w, s.w - 740), 430), Vec(150, 700)
    else:
        lit, stray = s.frac(0.3, 0.42), s.frac(0.14, 0.82)
    fish = [
        Koi(
            lit, -160, 300, 0.9, 0.6, 36, ACCENT, ACCENT_1, ACCENT_2, ACCENT_1, 0.55,
            ((0.04, 0.15, 0.6, 0.2), (0.24, 0.52, 0.75, -0.2), (0.6, 0.72, 0.55, 0.3)),
        ),
        Koi(
            lit + (320, 170), -150, 270, 0.75, 2.2, 32, UI_ALT, UI, UI, UI_HI, 0.8,
            ((0.08, 0.2, 0.55, -0.15), (0.32, 0.56, 0.7, 0.2)),
        ),
        # a third, barely there, nosing in from the left edge
        Koi(
            stray, -12, 190, -0.35, 1.1, 22, BG_ALT, BG_ALT, BG_ALT, UI, 0.8,
            ((0.3, 0.5, 0.55, 0.0),),
        ),
    ]  # fmt: skip

    # three rings spreading from where each tail of the pair last flicked, the outermost faintest
    with s.buckets(RIPPLES, "stroke", stroke_width=1.4) as rings:
        for k in fish[:2]:
            pts, _ = spine(k.head, k.heading, k.length * 1.35, k.turn * 1.35, k.phase)
            for j in range(len(RIPPLES)):
                rings[j].circle(pts[-1], 0.12 * k.length * (j + 1) ** 1.3)

    for i, k in enumerate(fish):
        body, fins, tail, pts, angs = koi(k)
        # the tail thins towards its tips, sinking a little into the water
        end = pts[-1]
        tip = mix(BG, k.tail, 0.7 * k.fade)
        fade = s.linear_gradient(
            [(0, k.tail), (0.25, k.tail), (1, tip)], end, polar(end, 0.4 * k.length, rad=angs[-1])
        )
        s.fill(P().shape(fins), mix(BG, k.fins, k.fade))
        s.fill(P().shape(body), k.body)
        # tail over the body so its root hides the rounded end of the peduncle
        s.fill(P().shape(tail), fade)
        mk = markings(s.noise(5 + i), 5 + i, pts, k)
        mk = mk.buffer(-9).buffer(13).buffer(-4).intersection(body.buffer(-5))
        patches = P()
        for g in parts(mk):
            if isinstance(g, Polygon):
                patches.shape(g)
        s.fill(patches, k.patch)