import sys
import Image
from array import array
import os
import pickle

try:
    memo = pickle.load(open("memo"))
except:
    memo = {}

import atexit
atexit.register(lambda: pickle.dump(memo, open("memo", "w")))

def rgbpal(imdata):
    # For RGB imdata, return list of (r,g,b) triples and the palette
    li = array('B', imdata).tolist()
    rgbs = zip(li[0::3], li[1::3], li[2::3])
    palette = list(set(rgbs))
    return (rgbs, palette)

def getch(im, x, y):
    # return the RGB data for the 8x8 character at (x, y) in im
    # if the 8x8 RGB contains more than 4 colors, quantize it using
    # `scolorq <http://www.cs.berkeley.edu/~dcoetzee/downloads/scolorq/>`_.

    sub88 = im.crop((x, y, x + 8, y + 8))
    sub88d = sub88.tostring()
    if sub88d in memo:
        return Image.fromstring("RGB", (8, 8), memo[sub88d])
    else:
        (_, pal) = rgbpal(sub88d)
        if len(pal) > 4:
            open("rgb", "w").write(sub88d)
            os.system("scolorq/spatial_color_quant rgb 8 8 4 q.rgb")
            data = open("q.rgb").read()
        else:
            data = sub88d
        memo[sub88d] = data
        sub88p = Image.fromstring("RGB", (8, 8), data)
        return sub88p
    return tuple(tuple((int(0 != im.getpixel((x + j, y + i)))) for j in range(8)) for i in range(8))
    
def binary(imdata):
    # imdata is 8x8x3 RGB character, string of length 192
    # return the pixel data and palette data for it
    assert len(imdata) == (3 * 8 * 8)
    (rgbs, palette) = rgbpal(imdata)
    indices = [palette.index(c) for c in rgbs]
    indices_b = ""
    for i in range(0, len(indices), 4):
        if 0:
            c =   ((indices[i] << 6) +
                   (indices[i + 1] << 4) +
                   (indices[i + 2] << 2) +
                   (indices[i + 3]))
        else:
            lup = [ 0x00, 0x08, 0x80, 0x88 ]
            c = ((lup[indices[i]]) +
                 (lup[indices[i + 1]] >> 1) +
                 (lup[indices[i + 2]] >> 2) +
                 (lup[indices[i + 3]] >> 3));
        indices_b += (chr(c))
    palette = (palette + ([(0,0,0)] * 4))[:4]
    def rgb555(r, g, b):
        return ((r / 8) << 10) + ((g / 8) << 5) + (b / 8)
    ph = array('H', [rgb555(*p) for p in palette])
    ph.byteswap()
    palette_b = ph.tostring()
    return (indices_b, palette_b)

def encode(im):
    preview = Image.new("RGB", im.size)
    charset = {}
    picture = []
    for y in range(0, im.size[1], 8):
        print y, im.size[1]
        for x in range(0, im.size[0], 8):
            iglyph = getch(im, x, y)
            preview.paste(iglyph, (x, y))
            glyph = iglyph.tostring()
            if not glyph in charset:
                charset[glyph] = len(charset)
            picture.append(charset[glyph])
    preview.save("preview.png")
    picd = array('B', picture)
    cd = array('B', [0] * 16 * len(charset))
    pd = array('B', [0] * 8 * len(charset))
    for d,i in charset.items():
        for y in range(8):
            (char, pal) = binary(d)
            cd[16 * i:16 * (i+1)] = array('B', char)
            pd[8 * i:8 * (i+1)] = array('B', pal)
    return (picd, cd, pd)

def main(filename):
    im = Image.open(filename).convert("RGB")
    (picdata, chrdata, paldata) = encode(im)
    open(filename + ".pic", "w").write(picdata.tostring())
    open(filename + ".chr", "w").write(chrdata.tostring())
    open(filename + ".pal", "w").write(paldata.tostring())

if __name__ == "__main__":
    main(sys.argv[1])