#!/usr/bin/python3
# Recolor PNGs
#
# (c) 2022 Sacha Chua (sacha@sachachua.com) - MIT License
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:

# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np
import os
import csv
import argparse
from PIL import Image

DARK_GRAY = 0x9d
GRAY = 0xc9
HEADER_GRAY = 0xca
WHITE = 0xfe

color_dict = {}

def color_to_tuple(color_dict, s):
    if s in color_dict:
        s = color_dict[s]
    s = s.lstrip('#')
    if (s == '.'):
        return (None, None, None)
    elif (len(s) == 2):
        return (int(s, 16), int(s, 16), int(s, 16))
    else:
        return tuple(int(s[i:i + 2], 16) for i in (0, 2, 4))

def load_color_dict(filename):
    dict = {}
    with open(os.path.expanduser(filename), newline='') as csvfile:
        reader = csv.reader(csvfile, delimiter=',', quotechar='"')
        for row in reader:
            dict[row[0]] = row[1]
    return dict

def remove_grid(input):
    if isinstance(input, str):
        im = Image.open(input).convert('RGB')
    else:
        im = input
    data = np.array(im)
    freq = get_colors_by_freq(input)
    print(freq)
    return Image.fromarray(data)

def map_colors(input, color_map):
    if isinstance(input, str):
        im = Image.open(input).convert('RGB')
    else:
        im = input
    data = np.array(im)
    red, green, blue = data[:, :, 0], data[:, :, 1], data[:, :, 2]
    for from_c, to_c in color_map.items():
        from_r, from_g, from_b = color_to_tuple(color_dict, from_c)
        to_r, to_g, to_b = color_to_tuple(color_dict, to_c)
        mask = (red == from_r) & (green == from_g) & (blue == from_b)
        data[:, :, :3][mask] = [to_r, to_g, to_b]
    return Image.fromarray(data)

def set_colors_by_freq(input, color_list):
    if isinstance(input, str):
        im = Image.open(input).convert('RGB')
    else:
        im = input
    data = np.array(im)
    red, green, blue = data[:, :, 0], data[:, :, 1], data[:, :, 2]
    sorted_colors = get_colors_by_freq(input)
    freq = iter(color_list.split(','))
    for i, f in enumerate(freq):
        if f != '.':
            to_r, to_g, to_b = color_to_tuple(color_dict, f)
            by_freq = sorted_colors[i][1]
            if isinstance(by_freq, np.uint8):
                mask = (red == by_freq) & (green == by_freq) & (blue == by_freq)
            else:
                mask = (red == by_freq[0]) & (green == by_freq[1]) & (blue == by_freq[2])
            data[:, :, :3][mask] = [to_r, to_b, to_g]
    return Image.fromarray(data)

def color_string_to_map(s):
    color_map = {}
    colors = iter(args.colors.split(','))
    for from_c in colors:
        to_c = next(colors)
        color_map[from_c] = to_c
    return color_map

def get_colors_by_freq(input):
    if isinstance(input, str):
        im = Image.open(input).convert('RGB')
    else:
        im = input
    colors = im.getcolors(im.size[0] * im.size[1])
    return sorted(colors, key=lambda x: x[0], reverse=True)


def print_colors(input):
    sorted_colors = get_colors_by_freq(input)
    for x in sorted_colors:
        if x[0] > 10:
            if isinstance(x[1], np.uint8):
                print('%02x %d' % (x[1], x[0]))
            else:
                print(''.join(['%02x' % c for c in x[1]]) + ' %d' % x[0])


def process_file(input):
    print(input)
    if args.preview:
        output = None
    else:
        output = args.output if args.output else input
        if os.path.isdir(output):
            output = os.path.join(output, os.path.basename(input))
    im = Image.open(input).convert('RGB')
    if args.colors:
        im = map_colors(im, color_string_to_map(args.colors))
    elif args.freq:
        im = set_colors_by_freq(im, args.freq)
    else:
        print_colors(im)
        exit(0)
    if args.preview:
        im.thumbnail((700, 700))
        im.show()
    elif output:
        im.save(output)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Recolor a PNG.',
        formatter_class=argparse.RawTextHelpFormatter,
        epilog="If neither --colors nor --freq are specified, "
        + "display the most frequent colours in the image.")
    parser.add_argument('--colors', help="""Comma-separated list of RGB hex values in the form of old,new,old,new
    Examples:
    9d,ffaaaa,c9,ffd2d2 - reddish
    c9,ffea96 - yellow highlighter
    c9,d2d2ff - light blue
    """)
    parser.add_argument('--freq', help="Color replacements in order of descending frequency (ex: .,ffea96). .: use original color")
    parser.add_argument('--csv', help="CSV of color names to use in the form of colorname,hex")
    parser.add_argument('--preview', help="Preview only", action='store_const', const=True)
    parser.add_argument('input', nargs="+", help="Input file")
    parser.add_argument('--output', help="Output file. If not specified, overwrite input file.")

    args = parser.parse_args()
    color_dict = load_color_dict(args.csv) if args.csv else {}
    for input in args.input:
        process_file(os.path.join(os.getcwd(), input))