! Example of Ulam spiral.
! Author: Francois Tonneau

! Plot idea from Alain Matthes:
! https://texample.net/ulam-spiral/

size 18.4 18.4

xpos = pagewidth()/2
ypos = pageheight()/2

unit = 2

numstart = 41
numstop  = 121

length = 1

declare sub run length direction$
declare sub paint number xpos ypos
declare sub primality number

declare sub lighting x y
declare sub blueish z
declare sub pinkish z

number = numstart
paint number xpos ypos

! ------------------------------
! Draw the spiral by successive turns.
! ------------------------------

while number <= 96
    ! Do one full turn.
    run length up
    run length right
    length = length + 1
    run length down
    run length left
    length = length + 1
next
! Add last ladder.
run length-1 up

sub run length direction$
    local dx, dy
    if direction$ = "up" then
        dx = 0; dy = unit
    else if direction$ = "right" then
        dx = unit; dy = 0
    else if direction$ = "down" then
        dx = 0; dy = -unit
    else if direction$ = "left" then
        dx = -unit; dy = 0
    end if
    local k
    for k = 1 to length
        number = number + 1
        xpos = xpos + dx
        ypos = ypos + dy
        paint number xpos ypos
    next k
end sub

sub paint number xpos ypos
    amove xpos ypos
    set color black lwidth 0.06
    box unit unit justify cc fill "#282820"
    gsave
    begin clip
        begin path clip
            circle (unit/2)-0.05        ! -0.05: => avoid biting on box borders
        end path
        rmove -unit/2 -unit/2
        if primality(number) = 1 then
            colormap "lighting(x,y)" -1 1 -1 1 100 100 unit unit palette blueish
        else
            colormap "lighting(x,y)" -1 1 -1 1 100 100 unit unit palette pinkish
        end if
    end clip
    grestore
    amove xpos-0.08 ypos+0.08           ! +-0.08: => improve apparent centering
    set just cc hei 0.55 color black
    write number
end sub

! ------------------------------
! Determine whether a number is prime
! ------------------------------

sub divides number factor
    local quotient = number/factor
    if int(quotient) = quotient then
        return 1
    else
        return 0
    end if
end sub

sub primality number
    ! This simple test works only for integers <= 121!
    if divides(number, 2)  = 1 then return 0
    if divides(number, 3)  = 1 then return 0
    if divides(number, 5)  = 1 then return 0
    if divides(number, 7)  = 1 then return 0
    if divides(number, 11) = 1 then return 0
    return 1
end sub

! ------------------------------
! Handle ball coloring
! ------------------------------

sub lighting x y
    local gloss = 8.5       ! gloss exponent
    local radius = 2.0      ! radius of a hypothetical x, y, z ball
    local rsquare = radius * radius
    if x * x + y * y > rsquare then return 0
    local z = sqrt(rsquare - x * x - y * y)
    local x_hot = -0.30; local y_hot =  0.30      ! position of brightest spot
    local z_hot = sqrt(rsquare - x_hot * x_hot - y_hot * y_hot)
    local dot_product  = x * x_hot + y * y_hot + z * z_hot
    local cosine = dot_product / rsquare
    if cosine < 0 then cosine = 0
    return cosine^gloss
end sub

sub tone r_base g_base b_base z
    local r_black =   0; local g_black =   0;  local b_black =   0
    local r_white = 235; local g_white = 235;  local b_white = 235
    local w, r, g, h
    if z < 0.50 then
        w = z / 0.50
        r = (1 - w) * r_black + w * r_base
        g = (1 - w) * g_black + w * g_base
        b = (1 - w) * b_black + w * b_base
    else
        w = (z - 0.50) / 0.50
        r = (1 - w) * r_base + w * r_white
        g = (1 - w) * g_base + w * g_white
        b = (1 - w) * b_base + w * b_white
    end if
    return rgb255(r, g, b)
end sub

sub blueish z
    return tone(90, 90, 180, z)
end sub

sub pinkish z
    return tone(180, 100, 100, z)
end sub