'''PhotomosaicsWithPython.py
This program reads in one or more tile-database images and one 
subject image. Then it constructs a photomosaic version of the 
subject image using the tile images.
'''

TILE_WIDTH = 128
TILE_HEIGHT = 128

# The following list of filenames can be arbitrarily long.
# Each file should name an image that holds one or more 128x128 tile images.
DBImage_filenames = ['FlowerDBImage.JPG']
prefix = '' # Change if your database images are not in current folder.
TILES_PACKAGED_IN = []  # Window num of database image containing the tile.
TILES_COLUMN_INDEX = [] # Horiz. pos. of tile in its database image (e.g., 0, 1, 2, 3)
TILES_ROW_INDEX = []    # Vert. pos. of tile in its database image.
TILES_AVG_COLOR = []    # RGB value representing the average color of the tile.
TILES_4X4_REDUCTION = []# A list of 16 RGB triples, used in matching.

# NTILES tells how many of the tiles read in should be used.
NTILES=92 # Any additional tiles will be ignored.

def load_image_database():
  for dbimage in DBImage_filenames:
    wn = pmOpenImage(0, prefix + dbimage)
    w = pmGetImageWidth(wn)
    h = pmGetImageHeight(wn)
    ncols = w / TILE_WIDTH
    nrows = h / TILE_HEIGHT
   
    count = 0
    print 'Image '+dbimage+ ' has width '+str(w)+\
          ' and height '+str(h)+' and seems to have '+\
          str(ncols*nrows)+' tiles.'
    for j in range(ncols):
      for i in range(nrows):
        # Processing a tile:
        TILES_PACKAGED_IN.append(wn)
        TILES_COLUMN_INDEX.append(j)
        TILES_ROW_INDEX.append(i)
        four_by_four = reduce_tile(wn,i,j)
        TILES_4X4_REDUCTION.append(four_by_four)
        avg = get_RGB_avg(four_by_four)
        TILES_AVG_COLOR.append(avg)
        count += 1
        if count==NTILES: return

def reduce_tile(wn,i,j):
  reduction = []
  horiz_factor = TILE_WIDTH / 4
  vert_factor = TILE_HEIGHT / 4
  xstart = j*TILE_WIDTH
  ystart = i*TILE_HEIGHT
  averaging_factor = horiz_factor*vert_factor
  for mx in range(4):
    for my in range(4):
       r=0; g=0; b=0
       for x in range(horiz_factor):
         for y in range(vert_factor):
           rgb = pmGetPixel(wn,xstart + mx*horiz_factor + x,
                               ystart + my*vert_factor  + y)
           r+=rgb[0]; g+=rgb[1]; b+=rgb[2]
       avg_r = r/averaging_factor
       avg_g = g/averaging_factor
       avg_b = b/averaging_factor
       reduction.append((avg_r, avg_g, avg_b))
  return reduction

def get_RGB_avg(rgb_list):
  r = 0; g = 0; b = 0
  the_len = len(rgb_list)
  for (dr, dg, db) in rgb_list:
    r+=dr; g+=dg; b+=db
  print 'Average RGB for this tile or patch: '+\
        str((r/the_len, g/the_len, b/the_len))
  return (r/the_len, g/the_len, b/the_len)

def print_database_info():
  print 'TILES_PACKAGED_IN = ' + str(TILES_PACKAGED_IN)
  print 'TILES_AVG_COLOR   = ' + str(TILES_AVG_COLOR)

load_image_database()
print_database_info()

''' We assume that the database of tile images is now loaded.
Now we load in the original image and compute a slightly reduced
version that we call the 'subject' image. This is the main image
to be rendered using the tiles.
'''
path = 'AnnaAtConservatory2.jpg'
ORIG = pmOpenImage(0, path)
SUBJ_WIDTH  = 256
SUBJ_HEIGHT = 256
subj = pmNewImage(0, 'Subject image to be rendered',\
                  SUBJ_WIDTH, SUBJ_HEIGHT, 0,0,0)
pmSetFormula('S1(x*w1/w,y*h1/h)')
pmSetSource1(ORIG); pmSetDestination(subj); pmCompute()

''' And now we create a blank image to hold the photomosaic.
'''
TILE_H_COVERAGE = 8 # Horiz. size of a subject image zone.
TILE_V_COVERAGE = 8 # Vert.   "
EFFECTIVE_TILE_WIDTH = 32 # Not more than width of tiles in database
EFFECTIVE_TILE_HEIGHT= 32
PHOTOM_WIDTH  = SUBJ_WIDTH * EFFECTIVE_TILE_WIDTH / TILE_H_COVERAGE
PHOTOM_HEIGHT = SUBJ_HEIGHT* EFFECTIVE_TILE_HEIGHT/ TILE_V_COVERAGE

NROWS_OF_TILES = PHOTOM_HEIGHT / EFFECTIVE_TILE_HEIGHT
NCOLS_OF_TILES = PHOTOM_WIDTH / EFFECTIVE_TILE_WIDTH
print 'There will be ' + str(NROWS_OF_TILES)+' rows of tiles'
print 'There will be ' + str(NCOLS_OF_TILES)+' columns of tiles'

PHOTOM = pmNewImage(0, 'Photomosaic', PHOTOM_WIDTH, PHOTOM_HEIGHT, 0,0,0)

# Returns the sum of the squares of the differences of 
# two sets of RGB values:
def colorDistance(c1, c2):
 dr = c1[0]-c2[0]
 dg = c1[1]-c2[1]
 db = c1[2]-c2[2]
 return dr*dr + dg*dg + db*db

def reduced_patch_distance(p1, p2):
  color_distances = map(colorDistance, p1, p2)
  return reduce(lambda x,y: x+y, color_distances)

def computePhotomosaic():
  # Loop through all the "effective tile" locations
  for i in range(NROWS_OF_TILES):
    for j in range(NCOLS_OF_TILES):
      # Compute a reduced (4 by 4) version of the corresponding 
      # patch of the subject image:
      #print 'Preparing to render patch ('+str(i)+','+str(j)+')'
      patch_reduction = reduce_patch(i, j)
      avg = get_RGB_avg(patch_reduction)
      index_of_best_tile = choose_best_tile(patch_reduction, avg)
      render_patch(i,j,index_of_best_tile, avg)

def choose_best_tile(patch_reduction, avgRGB):
  '''This computes the sum of the pixel distances between the
  adjusted representative values from each tile and the reduced 
  patch pixels passed in. Selects the tile with the minimum distance.  
  '''
  best = -1
  best_dist = 100000
  for i in range(NTILES):
    this_dist = reduced_patch_distance(patch_reduction,\
                                       TILES_4X4_REDUCTION[i])
    if this_dist < best_dist:
      best = i; best_dist = this_dist
  #print "best_dist = "+str(best_dist)+\
  # "; best tile index = "+str(best)
  return best

def compute_RGB_scale_factors(color1, color2):
  scale_red = (color2[0]+0.0)/(color1[0]+1.0)
  scale_green = (color2[1]+0.0)/(color1[1]+1.0)
  scale_blue = (color2[2]+0.0)/(color1[2]+1.0)
  return (scale_red, scale_green, scale_blue)

def render_patch(i,j,tile_idx,patch_avg):
  tile_win = TILES_PACKAGED_IN[tile_idx]
  tile_row = TILES_ROW_INDEX[tile_idx]
  tile_col = TILES_COLUMN_INDEX[tile_idx]
  tile_rgb = TILES_AVG_COLOR[tile_idx]
  (scale_red, scale_green, scale_blue) =\
    compute_RGB_scale_factors(tile_rgb, patch_avg)
  #print "Scale factors for R,G, and B are: "+\
  #  str(scale_red)+','+str(scale_green)+','+str(scale_blue)
  x_mos_start = j*EFFECTIVE_TILE_WIDTH
  y_mos_start = i*EFFECTIVE_TILE_HEIGHT
  x_tile_start = tile_col*TILE_WIDTH
  y_tile_start = tile_row*TILE_HEIGHT
  x_tile = x_tile_start
  x_mos = x_mos_start
  for dx_mos in range(EFFECTIVE_TILE_WIDTH):
    y_tile = y_tile_start
    y_mos = y_mos_start
    for dy_mos in range(EFFECTIVE_TILE_HEIGHT):
      tile_pixel = pmGetPixel(tile_win, int(x_tile), int(y_tile))
      new_r = scale_red * tile_pixel[0]
      new_g = scale_green*tile_pixel[1]
      new_b = scale_blue* tile_pixel[2]
      pmSetPixel(PHOTOM, int(x_mos), int(y_mos),\
                 int(new_r), int(new_g), int(new_b))
      y_mos += 1
      y_tile += (TILE_HEIGHT + 0.0) / EFFECTIVE_TILE_HEIGHT
    x_mos += 1
    x_tile += (TILE_WIDTH + 0.0) / EFFECTIVE_TILE_WIDTH

def reduce_patch(i,j):
  reduction = []
  horiz_factor = (SUBJ_WIDTH/NCOLS_OF_TILES)  / 4
  vert_factor  = (SUBJ_HEIGHT/NROWS_OF_TILES) / 4
  xstart = j*(SUBJ_WIDTH/NCOLS_OF_TILES)
  ystart = i*(SUBJ_HEIGHT/NROWS_OF_TILES)
  averaging_factor = horiz_factor*vert_factor
  for mx in range(4):
    for my in range(4):
       r=0; g=0; b=0
       for x in range(horiz_factor):
         for y in range(vert_factor):
           rgb = pmGetPixel(subj,xstart + mx*horiz_factor + x,
                            ystart + my*vert_factor + y)
           r+=rgb[0]; g+=rgb[1]; b+=rgb[2]
       avg_r = r/averaging_factor
       avg_g = g/averaging_factor
       avg_b = b/averaging_factor
       reduction.append((avg_r, avg_g, avg_b))
  return reduction

computePhotomosaic()