#! /usr/bin/python 
# Last edited on 2009-12-22 15:46:47 by stolfi

import sys
import math
import copy

def quantize(fval) :
  # Converts a float sample in {[0 _ 1]} to an int in {0..255}
  if (fval <= 0) :
    val = 0
  elif (fval >= 1) :
    val = 255
  else : 
    val = 1 + math.floor(fval*254);
    if (val > 254) :
      val = 254
  return val
  # ----------------------------------------------------------------------

def write_ppm_image(name, pic, nx, ny) :
  wr = open(name, "w")
  wr.write("P3\n")
  wr.write("%d %d %d\n" % (nx, ny, 255))
  for iy in range(ny) :
    for ix in range(nx) :
      pix = pic[iy][ix]
      if ((pix == None) or (len(pix) != 3)) :
        sys.stderr.write("bogus color value in pixel (%d,%d) = %s\n" % (ix,iy,pix));
      for ic in range(3) :
        val = quantize(pix[ic])
        wr.write(" %d" % val )
      if (ix % 10 == 0) :
        wr.write("\n")
    wr.write("\n")
  wr.close()
  # ----------------------------------------------------------------------

def write_pgm_image(name, pic, nx, ny) :
  wr = open(name, "w")
  wr.write("P2\n")
  wr.write("%d %d %d\n" % (nx, ny, 255))
  for iy in range(ny) :
    for ix in range(nx) :
      val = quantize(pic[iy][ix])
      wr.write(" %d" % val )
      if (ix % 30 == 0) :
        wr.write("\n")
    wr.write("\n")
  wr.close()
  # ----------------------------------------------------------------------
      
def make_weave_pic(qx, qy, nb, threadpix) :
  # {qx} X size of image (squarelets)
  # {qy} Y size of image (squarelets)
  # {nb} size of a squarelet (pixels)
  # {threadpix(ru,rv,su,sv,tk,st)} thread-rendering procedure.
  nx = qx*nb;  # X size of complete image (pixels)
  ny = qy*nb;  # Y size of complete image (pixels)
  pic = make_array(nx, ny)
  sys.stderr.write('[')
  for iy in range(ny) :
    sys.stderr.write('.')
    for ix in range(nx) :
      # Paint top thread:
      pix = render_weave_layer(ix,iy,nb,nx,ny,1,threadpix);
      if (pix == None) :
        # Paint bottom thread:
        pix = render_weave_layer(ix,iy,nb,nx,ny,0,threadpix);
      if (pix == None) :
        # Paint background:
        pix = render_weave_layer(ix,iy,nb,nx,ny,-1,threadpix);
      assert pix != None
      pic[iy][ix] = pix;
  sys.stderr.write(']\n')
  # ----------------------------------------------------------------------

def render_weave_layer(ix,iy,nb,nx,ny,st,threadpix) :
  # {ix,iy} pixel indices in whole image.
  # {nb} size of a squarelet (pixels)
  # {nx,ny} dimensions of image (pixels).
  # {st} weave layer to evaluate (1 = top, 0 = bottom, -1 = background).
  # {threadpix(ru,rv,su,sv,tk,st)} thread-rendering procedure.
  # Returns {None} if pixel is not on the layer.
  if (st == -1) :
    return threadpix(0,0,1,1,-1,0)
  else :
    # Compute indices {kx,ky} of current squarelet:
    kx = ix // nb # X index of curr squarelet.
    ky = iy // nb # Y index of curr squarelet.
    # Compute squarelet indices {jx,jy} folded around mid axes:
    jx = min(kx, qx - 1 - kx)
    jy = min(ky, qy - 1 - ky)
    # Compute number {t} of tile band containing {jx,jy}:
    t = (jx + jy) // 3 # Index of curr tile band.
    # Compute which kind {tk} of thread is on top in this band:
    tk = t % 2 # 0 = X-thread, 1 = Y-thread.
    # Compute upper corner {bx,by} and pixel size {sx,sy} of tile:
    if (tk != st) :
      # Tile of the X-oriented thread:
      bx = nb*(3*t - ky)
      by = nb*ky
      sx = 3*nb;
      sy = nb
    else :
      # Tile of the Y-oriented thread:
      bx = nb*ky
      by = nb*(3*t - kx)
      sx = nb;
      sy = 3*nb;
    # Adjust tile size to account for mid-axis mirroring: 
    if (2*(bx + sx) >= nx) :
      sx = nx - 2*bx
    if (2*(by + sy) >= ny) :
      sy = ny - 2*by
    # Compute thread-oriented pixel coords {ru,rv} rel to tile, and tile size {su,sv}:
    if (tk != st) :
      # Tile of the X-oriented thread on top.
      ru = ix - bx;
      rv = iy - by;
      su = sx;
      sv = sy;
    else :
      # Tile of the X-oriented thread on top.
      ru = iy - by;
      rv = ix - bx;
      su = sy;
      sv = sx;
    # sys.stderr.write(" pixel (%4d, %4d) --> (%4d, %4d) t = %d\n" % (ix,iy,ru,rv,t))
    assert (ru >= 0) & (ru < su)
    assert (rv >= 0) & (rv < sv)
    return threadpix(ru,rv,su,sv,tk,st)
  # ----------------------------------------------------------------------

def color_field(ix,iy,sx,sy,tk,st) :
  # {ix,iy} pixel indices relative to a basic weave tile.
  # {sx,sy} dimensions of tile (pixels).
  # {tk} which kind of thread is to be painted.
  # {st} layer where the thread is (0 = bottom, 1 = top).
  # Assumes that the thread runs in the X direction.
  # If {tk} is -1, returns the background color.
  # Otherwise returns {None} if {tk} is not on the thread. 
  if (tk == -1) :
    # Return background color:
    return [ 0.1, 0.1, 0.1]
  on, hz = weave_params(ix,iy,sx,sy,st)
  if (not on) :
    # Pixel is not on layer:
    return None
  else :
    # Compute color from rel height {hz} and tread type {tk}.
    # Compute shading factor {fc}:
    fc = 0.7 + 0.3*hz
    # Return proper color:
    return [ 
      (tk*0.9000 + (1-tk)*0.1000)*fc +  0.200*(1-fc),
      (tk*0.8500 + (1-tk)*0.1500)*fc +  0.200*(1-fc),
      (tk*0.7500 + (1-tk)*0.2500)*fc +  0.200*(1-fc)
    ]
  # ----------------------------------------------------------------------

def height_field(ix,iy,sx,sy,tk,st) :
  # {ix,iy} pixel indices relative to a basic tile.
  # {sx,sy} dimensions of tile (pixels).
  # {tk} which kind of thread is to be painted.
  # {st} layer where the thread is (0 = bottom, 1 = top).
  # Assumes that the thread runs in the X direction.
  if (tk == -1) :
    # Return background height:
    return 0.0
  on, hz = weave_params(ix,iy,sx,sy,st)
  if (not on) :
    # Pixel is not on layer:
    return None
  else :
    # Pixel is on a thread.
    # Compute actual height in {[0_1]}:
    return 0.500 + 0.450 * hz
  # ----------------------------------------------------------------------

def weave_params(ix,iy,sx,sy,st) :
  # {ix,iy} pixel indices relative to a basic tile.
  # {sx,sy} dimensions of tile (pixels).
  # {st} layer where the thread is (0 = bottom, 1 = top).
  # Assumes that the thread runs in the X direction.
  # Returns pair {(on,hz)} where
  # {on} is a boolean that tells whether the pixel is over the thread.
  # {hz} rel height of thread's surface ([-1.0 _ +1.0]), or 0 if not {on}.
  mb = sy // 10 # Margin pixels.
  if ((iy < mb) or (iy >= sy - mb)) :
    # Pixel is not over thread:
    on = 0;
    hz = 0;
  else :
    # Pixel is over thread:
    on = 1;
    # Thread's surface height relative to thread axis:
    gw = nb - 2*mb;
    gz0 = (iy - mb + 0.5)/gw;
    gz1 = (nb - mb - 0.5 - iy)/gw;
    gz = math.sqrt(4*gz0*gz1)
    gz = 1 - (1 - gz)*(1 - gz); # Flatten it.
    # Negate axis height {gz} if thread is in bottom layer:
    if (st != 1) :
      gz = - gz;
    # X factor -- axis height:
    fw = 3*nb;
    fz0 = (ix + 0.5)/fw;
    fz1 = (3*nb - 0.5 - ix)/fw;
    fz = math.sqrt(4*fz0*fz1)
    fz = 1 - (1 - fz)*(1 - fz);
    # Height is sum of axis and surface:
    hz = gz + fz;
  return (tk, hz)
  # ----------------------------------------------------------------------

def make_array(nx,ny) :
  sys.stderr.write("allocating array\n");
  a = ny*[None];
  for iy in range(ny) :
    a[iy] = copy.copy(nx*[None])
  return a;

# MAIN

# Get arguments:
name = sys.argv[1]; # Name of output files sans extension.
qx = int(sys.argv[2]); # Number of squarelets in X.
qy = int(sys.argv[3]); # Number of squarelets in Y.
nb = int(sys.argv[4]); # Side of each squarelet (pixels).

nx = qx*nb
ny = qy*nb

sys.stderr.write('creating color image\n')
pic_color = make_weave_pic(qx, qy, nb, color_field)
write_ppm_image(name + ".ppm", pic_color, nx, ny)

sys.stderr.write('creating height map\n')
pic_height = make_weave_pic(qx, qy, nb, height_field)
write_pgm_image(name + ".pgm", pic_height, nx, ny)