#!python
# -*- mode: python; Encoding: utf-8; coding: utf-8 -*-
# Last updated: <2022/07/13 08:35:19 +0900>
"""
Yliluoma's ordered dithering algorithm 3

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] [-c] [--mp]

-d num, --dither num : 2 or 4 or 8 (Dither 2x2, 4x4, 8x8)
-p PALETTE, --palette PALETTE : Palette file (.png or .gpl)
-c, --ciede2000 : enable CIEDE2000
--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 math
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)
]

# CIE C illuminant
illum = [
    0.488718, 0.176204, 0.000000,
    0.310680, 0.812985, 0.0102048,
    0.200602, 0.0108109, 0.989795
]

gamma = 2.2  # Gamma correction we use.


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 gamma_correct(v):
    return pow(v, gamma)


def gamma_uncorrect(v):
    return pow(v, 1.0 / gamma)


class LabItem:
    """ CIE L*a*b* color value with C and h added. """

    def __init__(self, r, g, b, divide):
        self.set(r / divide, g / divide, b / divide)

    def set(self, r, g, b):
        xx = illum[0] * r + illum[3] * g + illum[6] * b
        yy = illum[1] * r + illum[4] * g + illum[7] * b
        zz = illum[2] * r + illum[5] * g + illum[8] * b
        x = xx / (illum[0] + illum[1] + illum[2])
        y = yy / (illum[3] + illum[4] + illum[5])
        z = zz / (illum[6] + illum[7] + illum[8])

        threshold1 = (6 * 6 * 6.0) / (29 * 29 * 29.0)
        threshold2 = (29 * 29.0) / (6 * 6 * 3.0)

        x1 = pow(x, 1.0 / 3.0) if x > threshold1 else ((threshold2 * x) + (4 / 29.0))
        y1 = pow(y, 1.0 / 3.0) if y > threshold1 else ((threshold2 * y) + (4 / 29.0))
        z1 = pow(z, 1.0 / 3.0) if z > threshold1 else ((threshold2 * z) + (4 / 29.0))

        self.lum = (29 * 4) * y1 - (4 * 4)
        self.a = 500 * (x1 - y1)
        self.b = 200 * (y1 - z1)

        ct = self.a * self.a + self.b + self.b
        # ct = self.a * self.a + self.b * self.b
        if ct < 0:
            ct = 0
        self.c = math.sqrt(ct)
        self.h = math.atan2(self.b, self.a)


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 color_compare_ciede2000(lab1, lab2):
    """
    From the paper "The CIEDE2000 Color-Difference Formula: Implementation Notes,
    Supplementary Test Data, and Mathematical Observations",
    by Gaurav Sharma, Wencheng Wu and Edul N. Dalal,
    Color Res. Appl., vol. 30, no. 1, pp. 21-30, Feb. 2005.
    Return the CIEDE2000 Delta E color difference measure squared, for two Lab values
    """

    # Compute Cromanance and Hue angles
    cab = 0.5 * (lab1.c + lab2.c)
    cab7 = pow(cab, 7.0)
    g = 0.5 * (1.0 - math.sqrt(cab7 / (cab7 + 6103515625.0)))
    a1 = (1.0 + g) * lab1.a
    a2 = (1.0 + g) * lab2.a
    c1 = math.sqrt(a1 * a1 + lab1.b * lab1.b)
    c2 = math.sqrt(a2 * a2 + lab2.b * lab2.b)

    if c1 < 1e-9:
        h1 = 0.0
    else:
        h1 = math.degrees(math.atan2(lab1.b, a1))
        if h1 < 0.0:
            h1 += 360.0

    if c2 < 1e-9:
        h2 = 0.0
    else:
        h2 = math.degrees(math.atan2(lab2.b, a2))
        if h2 < 0.0:
            h2 += 360.0

    # Compute delta L, C and H
    dl = lab2.lum - lab1.lum
    dc = c2 - c1

    if c1 < 1e-9 or c2 < 1e-9:
        dhh = 0.0
    else:
        dhh = h2 - h1
        if dhh > 180.0:
            dhh -= 360.0
        elif dhh < -180.0:
            dhh += 360.0
    dh = 2.0 * math.sqrt(c1 * c2) * math.sin(math.radians(0.5 * dhh))

    lum = 0.5 * (lab1.lum + lab2.lum)
    c = 0.5 * (c1 + c2)
    if c1 < 1e-9 or c2 < 1e-9:
        h = h1 + h2
    else:
        h = h1 + h2
        if abs(h1 - h2) > 180.0:
            if h < 360.0:
                h += 360.0
            elif h >= 360.0:
                h -= 360.0
        h *= 0.5

    t = 1.0 \
        - 0.17 * math.cos(math.radians(h - 30.0)) \
        + 0.24 * math.cos(math.radians(2.0 * h)) \
        + 0.32 * math.cos(math.radians(3.0 * h + 6.0)) \
        - 0.2 * math.cos(math.radians(4.0 * h - 63.0))

    hh = (h - 275.0) / 25.0
    ddeg = 30.0 * math.exp(-hh * hh)
    c7 = pow(c, 7.0)
    rc = 2.0 * math.sqrt(c7 / (c7 + 6103515625.0))
    l50sq = (lum - 50.0) * (lum - 50.0)
    sl = 1.0 + (0.015 * l50sq) / math.sqrt(20.0 + l50sq)
    sc = 1.0 + 0.045 * c
    sh = 1.0 + 0.015 * c * t
    rt = -math.sin(math.radians(2 * ddeg)) * rc
    dlsq = dl / sl
    dcsq = dc / sc
    dhsq = dh / sh

    return dlsq * dlsq + dcsq * dcsq + dhsq * dhsq + rt * dcsq * dhsq


def devise_best_mixing_plan3(dt):
    src = dt["col"]
    n_colors = dt["n_colors"]
    limit = dt["limit"]
    luma = dt["luma"]
    pal_g = dt["pal_g"]
    # meta = dt["meta"]
    ciede2000 = dt["ciede2000"]
    pal = dt["pal"]

    r, g, b = src

    # Input color in RGB
    input_rgb = [r, g, b]

    # Input color in CIE L*a*b*
    inputlab = LabItem(r, g, b, 255.0)

    solution = {}

    # The penalty of our currently "best" solution.
    current_penalty = -1

    # First, find the closest color to the input color. It is our seed.
    if True:
        chosen = 0
        for index in range(n_colors):
            cr, cg, cb = pal[index]
            if ciede2000:
                test_lab = LabItem(cr, cg, cb, 255.0)
                penalty = color_compare_ciede2000(inputlab, test_lab)
            else:
                penalty = color_compare_ccir601(
                    input_rgb[0], input_rgb[1], input_rgb[2],
                    cr, cg, cb
                )
            if penalty < current_penalty or current_penalty < 0:
                current_penalty = penalty
                chosen = index
        solution[chosen] = limit

    dbllimit = 1.0 / limit
    while current_penalty != 0.0:
        # Find out if there is a region in Solution that
        # can be split in two for benefit.
        best_penalty = current_penalty
        best_splitfrom = 0xffffffff
        best_split_to = [0, 0]

        for split_color, split_count in solution.items():

            # if split_count <= 1:
            #     continue

            # Tally the other colors
            sum = [0, 0, 0]
            for col, cnt in solution.items():
                if col == split_color:
                    continue

                sum[0] += pal_g[col][0] * cnt * dbllimit
                sum[1] += pal_g[col][1] * cnt * dbllimit
                sum[2] += pal_g[col][2] * cnt * dbllimit

            portion1 = (split_count / 2.0) * dbllimit
            portion2 = (split_count - split_count / 2.0) * dbllimit

            for a in range(n_colors):
                # if(a != split_color && Solution.find(a) != Solution.end()) continue;

                firstb = 0
                if portion1 == portion2:
                    firstb = a + 1

                for b in range(firstb, n_colors, 1):
                    if a == b:
                        continue

                    # if(b != split_color && Solution.find(b) != Solution.end()) continue;

                    lumadiff = luma[a] - luma[b]
                    if lumadiff < 0:
                        lumadiff = -lumadiff
                    if lumadiff > 80000:
                        continue

                    test = [
                        gamma_uncorrect(sum[0] + pal_g[a][0] * portion1 + pal_g[b][0] * portion2),
                        gamma_uncorrect(sum[1] + pal_g[a][1] * portion1 + pal_g[b][1] * portion2),
                        gamma_uncorrect(sum[2] + pal_g[a][2] * portion1 + pal_g[b][2] * portion2)
                    ]

                    # Figure out if this split is better than what we had

                    if ciede2000:
                        test_lab = LabItem(test[0], test[1], test[2], 1)
                        penalty = color_compare_ciede2000(inputlab, test_lab)
                    else:
                        penalty = color_compare_ccir601(
                            input_rgb[0], input_rgb[1], input_rgb[2],
                            test[0] * 255, test[1] * 255, test[2] * 255
                        )

                    if penalty < best_penalty:
                        best_penalty = penalty
                        best_splitfrom = split_color
                        best_split_to[0] = a
                        best_split_to[1] = b

                    if portion2 == 0:
                        break

        if best_penalty == current_penalty:
            break  # No better solution was found.

        split_count = solution[best_splitfrom]
        split1 = split_count / 2.0
        split2 = split_count - split1

        del solution[best_splitfrom]

        if split1 > 0:
            if best_split_to[0] in solution.keys():
                solution[best_split_to[0]] += split1
            else:
                solution[best_split_to[0]] = split1
        if split2 > 0:
            if best_split_to[1] in solution.keys():
                solution[best_split_to[1]] += split2
            else:
                solution[best_split_to[1]] = split2

        current_penalty = best_penalty

    # Sequence the solution.
    r_colors = []
    for col, cnt in solution.items():
        size = len(r_colors) + cnt
        while len(r_colors) < size:
            r_colors.append(col)

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

    map_value = int(dt["threshold"] * len(cols) // dt["d_range"])
    r_col = pal[cols[map_value]]

    return (dt["x"], dt["y"], r_col)


def convert_dither(srcim, dither, ciede2000, 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

    # get palette length
    n_colors = len(pal)

    # 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]

    # algorithm 3
    pal_g = []
    meta = []
    for r, g, b in pal:
        gr = gamma_correct(r / 255.0)
        gg = gamma_correct(g / 255.0)
        gb = gamma_correct(b / 255.0)
        pal_g.append([gr, gg, gb])
        meta.append(LabItem(r, g, b, 255.0))

    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, "pal_g": pal_g, "meta": meta, "pal": pal,
                        "n_colors": n_colors, "limit": n_colors,
                        "d_range": d_range,
                        "ciede2000": ciede2000,
                        "threshold": dmap[y % dh][x % dw],
                    }
                    futures.append(executor.submit(devise_best_mixing_plan3, 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, "pal_g": pal_g, "meta": meta, "pal": pal,
                    "n_colors": n_colors, "limit": n_colors,
                    "d_range": d_range,
                    "ciede2000": ciede2000,
                    "threshold": dmap[y % dh][x % dw],
                }
                x, y, col = devise_best_mixing_plan3(dt)
                dst[x, y] = col

    return im


def main():
    parser = argparse.ArgumentParser(description="Yliluoma ordered dithering 3")
    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("-c", "--ciede2000", action="store_true", help="Enable CIEDE2000")
    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,
                        ciede2000=args.ciede2000, pal=palette,
                        mp=args.mp)
    im.save(args.output)


if __name__ == '__main__':
    main()