iso7.py

# Nicolas Seriot
# 2021-09-24
# 
# https://gist.github.com/nst/1cfb01d0b78993f7ffe2df7c101b586c
# Thread: https://twitter.com/nst021/status/1437889678947110912
# Typical output: https://seriot.ch/visualization/iso7.gif

import cairo
import random
import numpy as np

import cProfile, pstats, io
from pstats import SortKey

import imageio
import os

MARGIN = 40
X_MAX, Y_MAX, Z_MAX = 36,24,16 # space size
DW, DH = 28, 14 # diamond size
DW2, DH2 = int(DW/2), int(DH/2)
TTL = 40
NB_SNAKES = 30
NB_IMAGES = 200

def draw_surface(c, points, color, vertices = []):
    
    c.save()

    c.set_source_rgb(*color)
    
    [c.line_to(*p) for p in points]

    c.fill()
    
    c.set_source_rgb(0,0,0)
    c.set_line_width(1)
    
    for pts in vertices:
        p1, p2 = pts
        c.move_to(*p1)
        c.line_to(*p2)
        c.stroke()
    
    c.restore()
        
def draw_world_box(c, front=False, back=False):

    """
          5
    2         6
        3 8       <- 3 front, 8 back
    1         7 
        4
    """

    o_x = DW2 * Y_MAX
    o_y = DH2

    p1 = (0, Y_MAX*DH2)
    p2 = (0, p1[1] + DH*Z_MAX)
    p3 = (Y_MAX*DW2, DH*Z_MAX)
    p4 = (p3[0], 0)
    p5 = (X_MAX*DW2, (X_MAX+Y_MAX)*DH2 + Z_MAX*DH)
    p6 = ((X_MAX + Y_MAX) * DW2, X_MAX*DH2 + Z_MAX*DH)
    p7 = (p6[0], X_MAX*DH2)
    p8 = (p5[0], (X_MAX+Y_MAX)*DH2)
    
    c.save()

    c.set_source_rgb(0,0,0)
    c.set_line_width(1)

    if back:

        [c.line_to(*p) for p in [p5, p8]]
        c.stroke()

        [c.line_to(*p) for p in [p1, p8, p7]]
        c.stroke()

    if front:

        [c.line_to(*p) for p in [p1, p2, p5, p6, p7, p4, p1]]
        c.stroke()
    
        [c.line_to(*p) for p in [p2, p3, p6]]
        c.stroke()
    
        [c.line_to(*p) for p in [p3, p4]]
        c.stroke()
    
    c.restore()

def draw_cube(c, x, y, z,
              nx = False, nx_ = False, ny = False, ny_ = False, nz = False, nz_ = False,
              nxy_=False, nxz = False, nx_z_ = False, nyz = False, ny_z_ = False,
              flat=False):
    
    x, y = (DW2 * (Y_MAX-1) + DW2*(x-y), DH2*(x+y))
    
    c.save()

    c.translate(x, y)

    """
      5
    2   6
      3      - DH
    1   7
      4
      | z_offset
    """
    
    z_offset = DH * z

    p1 = (0,   z_offset + DH2)
    p2 = (0,   z_offset + DH2 + DH)
    p3 = (DW2, z_offset + DH)
    p4 = (DW2, z_offset + 0)
    p5 = (DW2, z_offset + 2*DH)
    p6 = (DW,  z_offset + DH2 + DH)
    p7 = (DW,  z_offset + DH2)
    
    #COLOR_RIGHT = (0.5,0.5,0.5)
    #COLOR_TOP = (1,1,1)
    #COLOR_LEFT = (0,0,0)
    
    z_ratio = 0.4 + z*0.6/Z_MAX
    
    COLOR_TOP   = (1 * z_ratio,   0.5 * z_ratio, 0.5 * z_ratio)
    COLOR_LEFT  = (0.6 * z_ratio, 0,             0)
    COLOR_RIGHT = (1 * z_ratio,   0,             0)

    vertices = []
    
    if flat:
        vertices.append((p1,p3))
        vertices.append((p3,p7))
        vertices.append((p7,p4))
        vertices.append((p4,p1))
        draw_surface(c, [], (1,1,1), vertices)
    else:
    
        if not nx:#
            vertices.append((p5, p6))
            vertices.append((p6, p7))
        if not ny:
            vertices.append((p2, p5))
            vertices.append((p2, p1))
        if not nz_:
            vertices.append((p4, p7))
            vertices.append((p1, p4))
        if not nx_ and not ny_:
            vertices.append((p3, p4))
        if not nz and not ny_:
            vertices.append((p6, p3))
        if not nz and not nx_:
            vertices.append((p2, p3))
        if nxy_:
            vertices.append((p6, p7))
        if nx_z_:
            vertices.append((p1, p4))
        if nxz:
            vertices.append((p5, p6))
        if nyz:
            vertices.append((p2, p5))
        if ny_z_:
            vertices.append((p4, p7))
        
        draw_surface(c, [p2, p5, p6, p3], COLOR_TOP, vertices)
        draw_surface(c, [p1, p2, p3, p4], COLOR_LEFT, vertices)
        draw_surface(c, [p3, p6, p7, p4], COLOR_RIGHT, vertices)
    
    c.restore()

def visibility_matrix(m):

    X,Y,Z = m.shape

    # nothing is visible except the three visible faces of the space

    v = np.full(m.shape, False)
    
    for x in range(X):
        for y in range(Y):
            v[x][y][Z-1] = True

    for x in range(X):
        for z in range(Z):
            v[x][0][z] = True

    for y in range(Y):
        for z in range(Z):
            v[0][y][z] = True

    # iterating from user's standpoint
    
    for x in range(X):
        for y in range(Y):
            for z in range(Z)[::-1]:

                # if not visible
                # no need to update visibility of "back" cube
                # continue to the next cube
                if not v[x][y][z]:
                    continue

                # if m is empty
                # "back" cube becomes visible
                if not m[x][y][z]:
                    if x < (X-1) and y < (Y-1) and z > 0:
                        v[x+1][y+1][z-1] = True
        
    return v

def offset_is_valid(m,x,y,z,a,b,c):

    return 0 <= (x+a) < X and 0 <= (y+b) < Y and 0 <= (z+c) < Z

def is_neighbour(ttls,oids,x,y,z,a,b,c):

    X,Y,Z = ttls.shape
    
    offset_is_valid = 0 <= (x+a) < X and 0 <= (y+b) < Y and 0 <= (z+c) < Z
    if not offset_is_valid:
        return False
    
    oid = oids[x,y,z]
    ttl = ttls[x,y,z]

    noid = oids[x+a,y+b,z+c]
    nttl = ttls[x+a,y+b,z+c]
    return noid == oid and abs(nttl-ttl) <= 1

def draw_model(ctx, ttls, oids):
    
    X,Y,Z = ttls.shape
    
    v = visibility_matrix(ttls)
    
    for x in range(X)[::-1]:
        for y in range(Y)[::-1]:
            for z in range(Z_MAX):
                
                if not v[x][y][z]:
                    continue

                oid = oids[x][y][z]
                if oid == 0:
                    continue
                
                ttl = ttls[x,y,z]
                
                nx  = is_neighbour(ttls,oids,x,y,z, 1, 0, 0)
                nx_ = is_neighbour(ttls,oids,x,y,z,-1, 0, 0)
                ny  = is_neighbour(ttls,oids,x,y,z, 0, 1, 0)
                ny_ = is_neighbour(ttls,oids,x,y,z, 0,-1, 0)
                nz  = is_neighbour(ttls,oids,x,y,z, 0, 0, 1)
                nz_ = is_neighbour(ttls,oids,x,y,z, 0, 0,-1)

                nxy_  = ttls[x+1][y-1][z] if (x+1) < X  and (y-1) >= 0 else False

                nxz   = ttls[x+1][y][z+1] if (x+1) < X  and (z+1) < Z  else False
                nx_z_ = ttls[x-1][y][z-1] if (x-1) >= 0 and (z-1) >= 0 else False

                nyz   = ttls[x][y+1][z+1] if (y+1) < Y  and (z+1) < Z  else False
                ny_z_ = ttls[x][y-1][z-1] if (y-1) >= 0 and (z-1) >= 0 else False
                
                draw_cube(ctx,x,y,z,nx,nx_,ny,ny_,nz,nz_,nxy_,nxz,nx_z_,nyz,ny_z_)

def images_paths():
    # ensure we can read images when working from another dir
    cwd = os.getcwd()
    dir_path = os.path.dirname(os.path.realpath(__file__))
    os.chdir(dir_path)
    files = [os.path.sep.join([dir_path, s])
             for s in os.listdir(dir_path)
             if s.endswith('.png')]
    files.sort()
    os.chdir(cwd)
    return files

def draw_gif(filename):

    images = []
    
    for s in images_paths():
        images.append(imageio.imread(s))
    
    imageio.mimsave(filename, images, format='GIF', duration=0.1)

def draw_png(m, oids, filename, draw_grid=False, draw_box=False):

    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32,
                                 (X_MAX+Y_MAX)*DW2 + MARGIN*2,
                                 (X_MAX+Y_MAX)*DH2 + Z_MAX*DH + MARGIN*2)
    c = cairo.Context(surface)
    
    cm = cairo.Matrix(yy=-1, y0=surface.get_height())
    c.transform(cm)
    
    c.translate(MARGIN, MARGIN)
    
    c.set_antialias(cairo.ANTIALIAS_NONE)
    
    #background
    c.set_source_rgb(1,1,1)
    c.paint()
    
    #pen
    c.set_source_rgb(0,0,0)
    c.set_line_width(1)

    if draw_grid:
        for x in range(X_MAX):
            for y in range(Y_MAX):        
                draw_cube(c, x, y, 0, flat=True)
    
    if draw_box:
        draw_world_box(c, back=True)

    draw_model(c, m, oids)

    if draw_box:
        draw_world_box(c, front=True)
    
    surface.write_to_png(filename)

def can_move_to(m, x, y, z):

    X,Y,Z = m.shape

    if not x in range(X):
        return False

    if not y in range(Y):
        return False

    if not z in range(Z):
        return False
    
    if m[x][y][z] > 0:
        return False

    return True    

def next_head(ttls, head, direction):

    x,y,z = head
    a,b,c = direction
    
    must_turn = can_move_to(ttls, x+a, y+b, z+c) == False

    must_keep_direction = random.uniform(0, 1) < 0.7

    if not must_turn and must_keep_direction:
        return (x+a, y+b, z+c), direction
        
    directions = [(-1,0,0), (1,0,0),
                  (0,-1,0), (0,1,0),
                  (0,0,-1), (0,0,1)]
               
    possible_directions = []

    for a,b,c in directions:
    
        x_ = x+a
        y_ = y+b
        z_ = z+c
        
        if x_ in range(X_MAX) \
            and y_ in range(Y_MAX) \
            and z_ in range(Z_MAX) \
            and ttls[x_][y_][z_] == 0:
            possible_directions.append((a,b,c))
        
    if len(possible_directions) == 0:
        print("** can't move")
        return head, direction
        
    a,b,c = random.choice(possible_directions)
    new_head = (x+a,y+b,z+c)
    
    return new_head, (a,b,c)

def main():

    #random.seed(0)

    ttls = np.full((X_MAX, Y_MAX, Z_MAX), 0) # TTLs
    oids = np.full((X_MAX, Y_MAX, Z_MAX), 0) # object ids

    """
    oids[0]      [0]      [0]       = 1
    oids[X_MAX-1][0]      [0]       = 1
    oids[X_MAX-1][Y_MAX-1][0]       = 1
    oids[X_MAX-1][Y_MAX-1][Z_MAX-1] = 1
    oids[X_MAX-1][0]      [Z_MAX-1] = 1
    oids[0]      [Y_MAX-1][0]       = 1
    oids[0]      [Y_MAX-1][Z_MAX-1] = 1
    oids[0]      [0]      [Z_MAX-1] = 1
    
    draw_png(ttls, oids, "x.png", draw_grid=True, draw_box=True)
    
    import sys
    sys.exit(0)
    
    """
    
    directions = [(-1,0,0), (1,0,0),
                  (0,-1,0), (0,1,0),
                  (0,0,-1), (0,0,1)]
    
    snakes = []
    for _ in range(NB_SNAKES):
        d = random.choice(directions)
        x = random.randrange(X_MAX)
        y = random.randrange(Y_MAX)
        z = random.randrange(Z_MAX)
        snake = {"h":(x,y,z), "d":d}
        snakes.append(snake)
    
    for oid,snake in enumerate(snakes):
        x,y,z = snake["h"]
        ttls[x][y][z] = TTL # time to live
        oids[x][y][z] = oid+1 # 0 means no object
    
    draw_png(ttls, oids, "000.png", draw_grid=True, draw_box=True)    
    
#    pr = cProfile.Profile()
#    pr.enable()

    for i in range(1, NB_IMAGES):
    
        for oid, s in enumerate(snakes):
            h = s["h"]
            d = s["d"]
            
            h_, d_ = next_head(ttls, h, d)
            s["h"] = h_
            s["d"] = d_
            x,y,z = h_
            ttls[x][y][z] = TTL
            oids[x][y][z] = oid+1

        for index, ttl in np.ndenumerate(ttls):
            new_ttl = max(0, ttl-1)
            ttls[index] = new_ttl
            if new_ttl == 0:
                oids[index] = 0
            
        filename = f'{i:03}.png'
        draw_png(ttls, oids, filename, draw_grid=True, draw_box=True)
        print(filename)

#    pr.disable()
#    s = io.StringIO()
#    sortby = SortKey.CUMULATIVE
#    ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
#    ps.print_stats()
#    print(s.getvalue())

    draw_gif("iso7.gif")
    
if __name__ == "__main__":
    
    main()