#!python
# -*- mode: python; Encoding: utf-8; coding: utf-8 -*-
# Last updated: <2022/06/20 19:54:00 +0900>
"""
Ordered dithering. N colos.

Windows10 x64 21H2 + Python 3.9.13 64bit + Pillow 9.1.1
"""

from PIL import Image
import sys

odtbl2x2 = [
    [1, 3],
    [4, 2]
]

odtbl3x3 = [
    [3, 7, 4],
    [6, 1, 9],
    [2, 8, 5]
]

# odtbl3x3 = [
#     [7, 9, 5],
#     [2, 1, 4],
#     [6, 3, 8]
# ]

odtbl4x4 = [
    [1, 9, 3, 11],
    [13, 5, 15, 7],
    [4, 12, 2, 10],
    [16, 8, 14, 6]
]

odtbl8x8 = [
    [1, 33, 9, 41, 3, 35, 11, 43],
    [49, 17, 57, 25, 51, 19, 59, 27],
    [13, 45, 5, 37, 15, 47, 7, 39],
    [61, 29, 53, 21, 63, 31, 55, 23],
    [4, 36, 12, 44, 2, 34, 10, 42],
    [52, 20, 60, 28, 50, 18, 58, 26],
    [16, 48, 8, 40, 14, 46, 6, 38],
    [64, 32, 56, 24, 62, 30, 54, 22]
]

odtbls = {
    "2x2": [5, odtbl2x2],
    "3x3": [10, odtbl3x3],
    "4x4": [17, odtbl4x4],
    "8x8": [65, odtbl8x8]
}


def get_dither_table(odtype, vrange, ofs):
    d, tbl = odtbls[odtype]
    if d >= vrange:
        print("Error: dither table range >= vrange")
        sys.exit()

    w = len(tbl[0])
    h = len(tbl)
    odtbl = [[0 for i in range(w)] for j in range(h)]

    for y in range(h):
        for x in range(w):
            odtbl[y][x] = tbl[y][x] * vrange / d + ofs

    return w, h, odtbl


def get_dither_image(odtype, im, level):
    width, height = im.size
    oim = Image.new("L", im.size)

    tw, th, _ = get_dither_table(odtype, 256, 0)

    ods = []
    for i in range(level):
        v0 = (256 * i) // level
        v1 = (256 * (i + 1)) // level
        vrange = v1 - v0
        _, _, od = get_dither_table(odtype, vrange, v0)
        ods.append([v0, v1, od])

    src_pixel = im.load()
    dst_pixel = oim.load()

    for y in range(height):
        dy = y % th
        for x in range(width):
            dx = x % tw
            v = src_pixel[x, y]
            for v0, v1, od in ods:
                if v0 <= v and v <= v1:
                    if v < od[dy][dx]:
                        dst_pixel[x, y] = v0
                    else:
                        dst_pixel[x, y] = v1

    return oim


def main():
    im = Image.open("mandrill.png")
    r, g, b = im.split()

    # odtype = "2x2"
    odtype = "4x4"
    level = 4
    
    rim = get_dither_image(odtype, r, level).convert("L")
    gim = get_dither_image(odtype, g, level).convert("L")
    bim = get_dither_image(odtype, b, level).convert("L")

    oim = Image.merge("RGB", (rim, gim, bim))
    oim.save("output_ncol_%s_lvl%d.png" % (odtype, level))


if __name__ == '__main__':
    main()