#!python
# -*- mode: python; Encoding: utf-8; coding: utf-8 -*-
# Last updated: <2022/07/13 08:39:11 +0900>
"""
Yliluoma's ordered dithering algorithm 4 (adobe like)

Arbitrary-palette positional dithering algorithm
https://bisqwit.iki.fi/story/howto/dither/jy/

Usage:
    py ThisScript.py -i INPUT.png -o OUTPUT.png [-d num] [-p PALETTE] [--mp]

-d num, --dither num : 2 or 4 or 8 (Dither 2x2, 4x4, 8x8)
-p PALETTE, --palette PALETTE : Palette file (.png or .gpl)
--mp : enable multi process

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

import os
import sys
from PIL import Image
from tqdm import tqdm
import argparse
import re
import concurrent.futures as cf


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

pal = [
    # (r, g, b)
    (8, 0, 0), (32, 26, 11), (67, 40, 23), (73, 41, 16),
    (35, 67, 9), (93, 79, 30), (156, 107, 32), (169, 34, 15),
    (43, 52, 124), (43, 116, 9), (208, 202, 64), (232, 160, 119),
    (106, 148, 171), (213, 196, 179), (252, 231, 110), (252, 250, 226)
]


class Palette:
    """ Palette class. """

    def __init__(self, filename):
        _, ext = os.path.splitext(filename)
        if ext == ".gpl":
            # GIMP palette file
            self.get_palette_from_gpl(filename)
        elif ext == ".png" or ext == ".gif":
            self.get_palette_from_image(filename)
        else:
            print("Error: Unsupported file = %s" % filename)
            sys.exit()

    def get_palette_from_gpl(self, filename):
        """ Get palette data from GIMP palette file (.gpl) """
        with open(filename) as f:
            data = f.read()
        lst = data.splitlines()

        self.colors = []
        self.palette = []
        for s in lst:
            if re.match(r"^\s*#", s):
                continue
            if re.match("^GIMP Palette", s):
                continue
            m = re.match(r"^Name:\s*(.+)$", s)
            if m:
                # print("Palette name: %s" % m.groups()[0])
                continue
            if re.match(r"^Columns", s):
                continue
            m = re.match(r"^\s*(\d+)\s+(\d+)\s+(\d+)", s)
            if m:
                r, g, b = m.groups()
                r = int(r)
                g = int(g)
                b = int(b)
                self.palette.append((r, g, b))

    def get_palette_from_image(self, filename):
        """ Get palette data from image. """
        im = Image.open(filename)
        w, h = im.size

        # print("# im.mode = %s" % im.mode)
        # print("%d x %d" % im.size)

        self.colors = []
        self.palette = []

        if im.mode == "P":
            # index color image
            cols = im.getcolors()
            p = im.getpalette()

            if p is None:
                print("Error: Not found palette.")
                sys.exit()

            for cnt, i in cols:
                r, g, b = p[i * 3], p[i * 3 + 1], p[i * 3 + 2]
                self.palette.append((r, g, b))
                self.colors.append((cnt, (r, g, b)))

        elif im.mode == "RGB":
            # RGB image
            rgb_count = {}
            src = im.load()
            for y in range(h):
                for x in range(w):
                    col = src[x, y]
                    if col in rgb_count:
                        rgb_count[col] += 1
                    else:
                        rgb_count[col] = 1
            self.palette = list(rgb_count.keys())

            for c in self.palette:
                cnt = rgb_count[c]
                self.colors.append((cnt, c))

        else:
            print("Error: Unsupported image mode = %s" % im.mode)
            sys.exit()

    def count(self):
        return sorted(self.colors, key=lambda x: x[0], reverse=True)


def color_compare_ccir601(r1, g1, b1, r2, g2, b2):
    """
    Compare the difference of two RGB values.
    weigh by CCIR 601 luminosity.
    """
    luma1 = (r1 * 299 + g1 * 587 + b1 * 114) / (255.0 * 1000.0)
    luma2 = (r2 * 299 + g2 * 587 + b2 * 114) / (255.0 * 1000.0)
    lumad = luma1 - luma2
    r = (r1 - r2) / 255.0
    g = (g1 - g2) / 255.0
    b = (b1 - b2) / 255.0
    return (r * r * 0.299 + g * g * 0.587 + b * b * 0.114) * 0.75 + lumad * lumad


def devise_best_mixing_plan4(dt):
    srccol = dt["col"]
    luma = dt["luma"]
    d_range = dt["d_range"]
    pal = dt["pal"]
    map_value = dt["mapvalue"]

    n_colors = len(pal)
    r_colors = [0] * d_range
    src = list(srccol)
    lx = 0.09   # Error multiplier
    e = [0, 0, 0]   # Error accumulator

    for c in range(d_range):
        # Current temporary value
        t = [
            int(src[0] + e[0] * lx),
            int(src[1] + e[1] * lx),
            int(src[2] + e[2] * lx)
        ]

        # Clamp it in the allowed RGB range
        if t[0] < 0:
            t[0] = 0
        elif t[0] > 255:
            t[0] = 255
        if t[1] < 0:
            t[1] = 0
        elif t[1] > 255:
            t[1] = 255
        if t[2] < 0:
            t[2] = 0
        elif t[2] > 255:
            t[2] = 255

        # Find the closest color from the palette
        least_penalty = 1e99
        chosen = c % n_colors

        for index in range(n_colors):
            pc = list(pal[index])
            penalty = color_compare_ccir601(pc[0], pc[1], pc[2], t[0], t[1], t[2])
            if penalty < least_penalty:
                least_penalty = penalty
                chosen = index

        # Add it to candidates and update the error
        r_colors[c] = chosen
        pc = list(pal[chosen])
        e[0] += src[0] - pc[0]
        e[1] += src[1] - pc[1]
        e[2] += src[2] - pc[2]

    # Sort the colors according to luminance
    # std::sort(result.colors, result.colors + 64, PaletteCompareLuma);
    cols = sorted(r_colors, key=lambda x: luma[x])

    return (dt["x"], dt["y"], pal[cols[map_value]])


def convert_dither(srcim, dither, pal, mp):
    w, h = srcim.size
    im = Image.new("RGB", (w, h))
    src = srcim.load()
    dst = im.load()

    dmap = dithermaps[dither]
    dh = len(dmap)
    dw = len(dmap[0])
    d_range = dw * dh  # dither level range

    # Luminance for each palette entry,
    # to be initialized as soon as the program begins
    luma = [(r * 299 + g * 587 + b * 114) for r, g, b in pal]

    # adobe like pattern dither
    if mp:
        # multi process
        for y in tqdm(range(h), ascii=True):
            with cf.ProcessPoolExecutor(max_workers=None) as executor:
                futures = []
                for x in range(w):
                    dt = {
                        "col": src[x, y], "x": x, "y": y,
                        "luma": luma, "d_range": d_range, "pal": pal,
                        "mapvalue": dmap[y % dh][x % dw],
                    }
                    futures.append(executor.submit(devise_best_mixing_plan4, dt))

                for future in cf.as_completed(futures):
                    _x, _y, _col = future.result()
                    dst[_x, _y] = _col
    else:
        # single process
        for y in tqdm(range(h), ascii=True):
            for x in range(w):
                dt = {
                    "col": src[x, y], "x": x, "y": y,
                    "luma": luma, "d_range": d_range, "pal": pal,
                    "mapvalue": dmap[y % dh][x % dw],
                }
                _x, _y, _col = devise_best_mixing_plan4(dt)
                dst[_x, _y] = _col

    return im


def main():
    parser = argparse.ArgumentParser(description="Yliluoma ordered dithering 4")
    parser.add_argument("-i", "--input", required=True, help="Input png filename")
    parser.add_argument("-o", "--output", required=True, help="Output png filename")
    parser.add_argument("-p", "--palette", help="Palette file (.png or .gpl)")
    parser.add_argument("-d", "--dither", type=int, default=8,
                        help="Dither type 2,4,8 (2x2,4x4,8x8). default: 8")
    parser.add_argument("--mp", action="store_true", help="Enable multi process")
    args = parser.parse_args()

    if args.dither not in [2, 4, 8]:
        print("Error: Unknown dither = %d" % args.dither)
        sys.exit()

    if args.palette is not None:
        # load palette file
        # print("Palette file : %s" % args.palette)
        p = Palette(args.palette)
        palette = p.palette
    else:
        # print("Palette file is None")
        palette = pal

    srcim = Image.open(args.input)
    im = convert_dither(srcim=srcim, dither=args.dither,
                        pal=palette, mp=args.mp)
    im.save(args.output)


if __name__ == '__main__':
    main()