#!/bin/python

BASEDIR = "/home/tastyfish/git/drummy_utils/img_averager/"
CHUNKSIZE = 256
WIDTH = 640
HEIGHT = 480

def nToXY(n):
  return (int(n % WIDTH),int(n / WIDTH))

def rgbTo332(rgb):
  return ((int(rgb[0]) >> 5) << 5) | ((int(rgb[1]) >> 5) << 2) | (int(rgb[2]) >> 6)

def rgbFrom332(n):
  r = [0,0,0]

  v = (int(n) >> 5) & 0x07;
  r[0] = int(v * 36 if v != 7 else 255)

  v = (n >> 2) & 0x07;
  r[1] = int(v * 36 if v != 7 else 255)
 
  v = n & 0x03;
  r[2] = int(v * 72 if v != 3 else 255)

  return tuple(r)

import glob
import os
import sys
from PIL import Image

os.chdir(BASEDIR)

imgMax = Image.new("RGB",(WIDTH,HEIGHT),color="white")
imgMaxPixels = imgMax.load()

imgAvg = Image.new("RGB",(WIDTH,HEIGHT),color="white")
imgAvgPixels = imgAvg.load()

hist = [[0 for j in range(256)] for i in range(CHUNKSIZE)]
sums = [[0,0,0] for i in range(CHUNKSIZE)]

for i in range(WIDTH * HEIGHT):
  if i % CHUNKSIZE == 0:
    print("chunk " + str(int(i / CHUNKSIZE)) + "/" + str(int((WIDTH * HEIGHT) / CHUNKSIZE)))

    for j in range(CHUNKSIZE):
      hist[j] = [0 for k in range(256)]
      sums[j] = [0,0,0]

    numImages = 0

    for ff in glob.glob("*.jpg") + glob.glob("*.png"):
      # print("  " + ff)
      numImages += 1
      ffi = Image.open(ff).convert("RGB")
      ffip = ffi.load()

      for k in range(CHUNKSIZE):
        try:
          p = ffip[nToXY(i + k)]
          sums[k][0] += p[0]
          sums[k][1] += p[1]
          sums[k][2] += p[2]
          hist[k][rgbTo332(p)] += 1
        except Exception as e:
          print("A " + str(e))
          pass

    for j in range(CHUNKSIZE):
      sums[j][0] = int(sums[j][0] / numImages)
      sums[j][1] = int(sums[j][1] / numImages)
      sums[j][2] = int(sums[j][2] / numImages)

  histMaxI = 0

  for k in range(256):
    if hist[i % CHUNKSIZE][k] > hist[i % CHUNKSIZE][histMaxI]:
      histMaxI = k

  try:
    imgMaxPixels[nToXY(int(i))] = rgbFrom332(histMaxI)
  except Exception as e:
    print("B " + str(e))

  try:
    imgAvgPixels[nToXY(int(i))] = tuple(sums[i % CHUNKSIZE])
  except Exception as e:
    print("C " + str(e))

imgMax.save("img_max.png")
imgAvg.save("img_avg.png")