Animation of the Koch Snowflake wiht polygonal segments

koch.py

# MIT License
#
# Copyright (c) 2021 Paolo Lammens
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import math
import queue
import threading
import time
import tkinter as tk
from colorsys import hsv_to_rgb
from turtle import RawTurtle, ScrolledCanvas, TurtleScreen, Vec2D
from typing import Callable, Tuple

import hdpitkinter as hdpitk
import numpy as np


SCALE: float = 100  # pixels per unit


root = hdpitk.HdpiTk()
root.geometry("1400x800")
root.title("Koch's Snowflake")

canvas = ScrolledCanvas(root)
canvas.pack(fill=tk.BOTH, expand=tk.YES)

screen = TurtleScreen(canvas)
turtle = RawTurtle(screen, visible=False)
turtle.speed(0)
turtle.pensize(2)


def _draw_koch_gcd(
    segment_func: Callable[[float], None],
    n: int,
    pos: Vec2D,
    heading: float,
    length: float,
    period: float = 1,
    hue_range: Tuple[float, float] = None,
):
    num_segments = 4 ** n
    segment_length = length / 3 ** n
    segment_period = period / num_segments

    if hue_range is None:
        hue_range, saturation, value = (0, 0), 0, 0
    else:
        saturation, value = 1, 0.9
    hues = np.linspace(*hue_range, endpoint=False, num=num_segments)

    screen.delay(delay=math.floor(segment_period * 1e3 / 5))
    turtle.penup()
    turtle.goto(pos * SCALE)  # noqa
    turtle.pendown()

    angles = [0, 60, -60, 0]
    for i, hue in zip(range(num_segments), hues):
        turtle.pencolor(*hsv_to_rgb(hue, saturation, value))

        angle = heading + sum(angles[(i >> (2 * j)) % 4] for j in range(n))
        turtle.setheading(angle)
        segment_func(segment_length)


def draw_koch(k: int, n: int, **kwargs):
    if k == 2:
        # classical snowflake, just draw segment (circle looks bad and grainy)
        _draw_koch_gcd(segment_func=lambda l: turtle.forward(l * SCALE), n=n, **kwargs)
    elif k is None or k > 2:

        def draw_segment(length):
            turtle.left(90)
            turtle.circle(radius=-length / 2 * SCALE, extent=180, steps=k and k // 2)

        _draw_koch_gcd(segment_func=draw_segment, n=n, **kwargs)
    else:
        assert False


def draw_koch_triangle(k: int, n: int, size=3, **kwargs):
    h = size * math.sqrt(3) / 2
    a = Vec2D(-size / 2, h / 3)
    b = Vec2D(size / 2, h / 3)
    c = Vec2D(0, -2 / 3 * h)

    hue_range = (0, 1)
    h0, h3 = hue_range
    h1 = h0 + 1 / 3 * (h3 - h0)
    h2 = h0 + 2 / 3 * (h3 - h0)

    draw_koch(k, n, pos=a, heading=0, length=size, hue_range=(h0, h1), **kwargs)
    draw_koch(k, n, pos=b, heading=-120, length=size, hue_range=(h1, h2), **kwargs)
    draw_koch(k, n, pos=c, heading=120, length=size, hue_range=(h2, h3), **kwargs)


class AnimationQueue(queue.Queue):
    def __init__(self):
        super().__init__(maxsize=1)

    def put_and_wait(self, animation: Callable[[], None]):
        self.put(animation)
        self.join()


def main(ks=(2, 4, 6, 8, None), max_n=4):
    legend_var = tk.StringVar(value="k = [num. sides of polygon]\nn = [iteration]")
    legend = tk.Label(
        canvas._canvas,
        textvariable=legend_var,
        font=("Arial", 18),
        anchor="w",
        justify=tk.LEFT,
    )
    legend.place(x=50, y=50)
    click_to_start = tk.Label(
        canvas._canvas,
        text="(click to start)",
        font=("Courier", 18),
        background="white",
    )
    click_to_start.place(relx=0.5, rely=0.5, anchor=tk.CENTER)

    animation_queue = AnimationQueue()

    def run_sequence():
        for overlay in False, True:
            for k in ks:
                animation_queue.put_and_wait(screen.clear)
                animation_queue.put_and_wait(turtle.home)

                for n in range(max_n + 1):
                    if not overlay:
                        animation_queue.put_and_wait(screen.clear)
                    animation_queue.put_and_wait(
                        lambda: legend_var.set(f"k = {k or '∞'}\nn = {n}")
                    )

                    animation_queue.put_and_wait(
                        lambda: draw_koch_triangle(k, n, period=0.5, size=5)
                    )

                    time.sleep(1.5)

                time.sleep(2)

        animation_queue.put(None)

    def run_turtle_worker():
        # have to poll to avoid freezing GUI while waiting in between animations
        try:
            if animation := animation_queue.get(block=False):
                animation()
                animation_queue.task_done()
                screen.ontimer(run_turtle_worker, t=12)
            else:
                return  # queue had None; we're done
        except queue.Empty:
            # schedule next poll
            screen.ontimer(run_turtle_worker, t=5)

    def start(x, y):
        click_to_start.place_forget()
        threading.Thread(target=run_sequence, daemon=True, name="Scheduler").start()
        run_turtle_worker()

    screen.onclick(start)
    root.mainloop()


main()