#!python
# -*- mode: python; Encoding: utf-8; coding: utf-8 -*-
# Last updated: <2022/07/13 08:15:46 +0900>
"""
Yliluoma's ordered dithering algorithm 1,tri-tone

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 * 1000.0)
    luma2 = (r2 * 299 + g2 * 587 + b2 * 114) / (255 * 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_plan_tritone(dt):
    col = dt["col"]
    pal = dt["pal"]
    d_range = dt["d_range"]

    r, g, b = col
    r_colors = [0, 0, 0, 0]
    r_ratio = 0.5
    least_penalty = 1e99

    len_pal = len(pal)
    for index1 in range(len_pal):
        for index2 in range(index1, len_pal, 1):
            # Determine the two component colors
            color1 = pal[index1]
            color2 = pal[index2]
            r1, g1, b1 = color1
            r2, g2, b2 = color2
            ratio = d_range / 2
            if color1 != color2:
                # Determine the ratio of mixing for each channel.
                # solve(r1 + ratio*(r2-r1)/64 = r, ratio)
                # Take a weighed average of these three ratios according to the
                # perceived luminosity of each channel (according to CCIR 601).

                if r2 != r1:
                    rr = (299 * d_range * int(r - r1) / int(r2 - r1))
                    rd = 299
                else:
                    rr = 0
                    rd = 0

                if g2 != g1:
                    gg = (587 * d_range * int(g - g1) / int(g2 - g1))
                    gd = 587
                else:
                    gg = 0
                    gd = 0

                if b2 != b1:
                    bb = (114 * d_range * int(b - b1) / int(b2 - b1))
                    bd = 114
                else:
                    bb = 0
                    bd = 0

                ratio = (rr + gg + bb) / (rd + gd + bd)
                if ratio < 0:
                    ratio = 0
                elif ratio > (d_range - 1):
                    ratio = d_range - 1

            # Determine what mixing them in this proportion will produce
            r0 = r1 + ratio * (r2 - r1) / d_range
            g0 = g1 + ratio * (g2 - g1) / d_range
            b0 = b1 + ratio * (b2 - b1) / d_range

            penalty = color_compare_ccir601(r, g, b, r0, g0, b0) \
                + color_compare_ccir601(r1, g1, b1, r2, g2, b2) * 0.1 \
                * (abs((ratio / d_range) - 0.5) + 0.5)

            if penalty < least_penalty:
                least_penalty = penalty
                r_colors[0] = index1
                r_colors[1] = index2
                r_ratio = ratio / d_range

            if index1 != index2:
                for index3 in range(len_pal):
                    if index3 == index2 or index3 == index1:
                        continue

                    # 50% index3, 25% index2, 25% index1
                    color3 = pal[index3]
                    r3, g3, b3 = color3
                    r0 = (r1 + r2 + r3 * 2.0) / 4.0
                    g0 = (g1 + g2 + g3 * 2.0) / 4.0
                    b0 = (b1 + b2 + b3 * 2.0) / 4.0

                    cc0 = color_compare_ccir601(r, g, b, r0, g0, b0)
                    cc1 = color_compare_ccir601(r1, g1, b1, r2, g2, b2)

                    cc2 = color_compare_ccir601((r1 + r2) / 2, (g1 + g2) / 2, (b1 + b2) / 2,
                                                r3, g3, b3)

                    # original
                    # cc2 = color_compare_ccir601((r1 + g1) / 2, (g1 + g2) / 2, (b1 + b2) / 2,
                    #                             r3, g3, b3)

                    penalty = cc0 + cc1 * 0.025 + cc2 * 0.025
                    if penalty < least_penalty:
                        least_penalty = penalty
                        r_colors[0] = index3  # (0,0) index3 occurs twice
                        r_colors[1] = index1  # (0,1)
                        r_colors[2] = index2  # (1,0)
                        r_colors[3] = index3  # (1,1)
                        r_ratio = 4.0

    x = dt["x"]
    y = dt["y"]
    if r_ratio == 4.0:
        # Tri-tone or quad-tone dithering
        i = (y % 2) * 2 + (x % 2)
        return (x, y, pal[r_colors[i]])

    i = 1 if (dt["threshold"] / d_range) < ratio else 0
    return (x, y, pal[r_colors[i]])


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

    # algorithm 1, tri-tone
    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,
                        "d_range": d_range, "pal": pal,
                        "threshold": dmap[y % dh][x % dw],
                    }
                    futures.append(executor.submit(devise_best_mixing_plan_tritone, 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,
                    "d_range": d_range, "pal": pal,
                    "threshold": dmap[y % dh][x % dw],
                }
                _x, _y, _col = devise_best_mixing_plan_tritone(dt)
                dst[_x, _y] = _col

    return im


def main():
    parser = argparse.ArgumentParser(description="Yliluoma ordered dithering 1 tri-tone")
    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
        p = Palette(args.palette)
        palette = p.palette
    else:
        # default palette
        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()