#!/usr/bin/python3
# lid: lo-fi image dithering
# written by sloum
# licence: public domain - use it, improve it, share it!
#########

from PIL import Image
import sys
import os, os.path

# Default dithering: 4 level ordered dither
def ordered_dither_4(r):
    im = Image.open(r["path"]).convert("L")
    width, height = im.size
    new = create_image(width, height)
    newpixels = new.load()

    dots = [
        [64, 128],
        [192, 0]
    ]

    total_px_count = width * height
    counter = 1
    last_percent_val = 0
    last_hash_val = 0
    print("Dither method: ordered (4 levels)\033[?25l")
    for row in range(0, height):
        for col in range(0, width):
            dotrow = 1 if row % 2 else 0
            dotcol = 1 if col % 2 else 0
            impx = get_pixel(im, col, row)
            newpixels[col, row] = int(impx > dots[dotrow][dotcol])
            counter += 1
            new_percent_val = round(counter / total_px_count * 100)
            new_hash_val = round(counter / total_px_count * 10)
            if new_percent_val != last_percent_val or new_hash_val != last_hash_val:
                print("\r{:>3}% |{:<10}|".format(new_percent_val, "#" * new_hash_val), end="")
                last_percent_val = new_percent_val
                last_hash_val = new_hash_val

    new.save("{}.{}".format(r["out"], r["-f"]), r["-f"].upper(), optimize=True, quality=r["-q"])
    print("\033[20C\nFinished! ~ {}/{}.{}\033[?25h".format(os.getcwd(), r["out"], r["-f"].lower()))

# 9 level ordered dither
def ordered_dither_9(r):
    im = Image.open(r["path"]).convert("L")
    width, height = im.size
    new = create_image(width, height)
    newpixels = new.load()

    dots = [
        [0, 196, 84],
        [168, 140, 56],
        [112, 28, 224]
    ]

    total_px_count = width * height
    counter = 1
    last_percent_val = 0
    last_hash_val = 0
    print("Dither method: ordered (9 levels)\033[?25l")
    for row in range(0, height):
        for col in range(0, width):
            if not row % 3:
                dotrow = 2
            elif not row % 2:
                dotrow = 1
            else:
                dotrow = 0

            if not col % 3:
                dotcol = 2
            elif not col % 2:
                dotcol = 1
            else:
                dotcol = 0
            impx = get_pixel(im, col, row)
            newpixels[col, row] = int(impx > dots[dotrow][dotcol])
            counter += 1
            new_percent_val = round(counter / total_px_count * 100)
            new_hash_val = round(counter / total_px_count * 10)
            if new_percent_val != last_percent_val or new_hash_val != last_hash_val:
                print("\r{:>3}% |{:<10}|".format(new_percent_val, "#" * new_hash_val), end="")
                last_percent_val = new_percent_val
                last_hash_val = new_hash_val

    new.save("{}.{}".format(r["out"], r["-f"]), r["-f"].upper(), optimize=True, quality=r["-q"])
    print("\033[20C\nFinished! ~ {}/{}.{}\033[?25h".format(os.getcwd(), r["out"], r["-f"].lower()))

# Alternate dither mode
def error_dither(r):
    im = Image.open(r["path"]).convert("L")
    width, height = im.size
    new = create_image(width, height)
    newpixels = new.load()
    i = im.load()

    total_px_count = width * height
    counter = 1
    last_percent_val = 0
    last_hash_val = 0
    print("Dither method: error diffusion\033[?25l")
    for row in range(0, height):
        for col in range(0, width):
            newpixels[col, row] = update_error(im, i, row, col)
            counter += 1
            new_percent_val = round(counter / total_px_count * 100)
            new_hash_val = round(counter / total_px_count * 10)
            if new_percent_val != last_percent_val or new_hash_val != last_hash_val:
                print("\r{:>3}% |{:<10}|".format(new_percent_val, "#" * new_hash_val), end="")
                last_percent_val = new_percent_val
                last_hash_val = new_hash_val

    new.save("{}.{}".format(r["out"], r["-f"]), r["-f"].upper(), optimize=True, quality=r["-q"])
    print("\033[20C\nFinished! ~ {}/{}.{}\033[?25h".format(os.getcwd(), r["out"], r["-f"].lower()))


def update_error(im, i, r, c):
    w, h = im.size
    current = get_pixel(im, c, r)
    if current > 128:
        res = 1
        diff = -(255 - current)
    else:
        res = 0
        diff = abs(0 - current)
    mov = [[0, 1, 0.4375],[1, 1, 0.0625],[1, 0, 0.3125],[1, -1, 0.1875]]
    for x in mov:
        if r + x[0] >= h or c + x[1] >= w or c + x[1] <= 0:
            continue
        p = get_pixel(im, c + x[1], r + x[0])
        p = round(diff * x[2] + p)
        if p < 0:
            p = 0
        elif p > 255:
            p = 255
        i[c + x[1], r + x[0]] = p
    return res



def print_spinner(n):
    print("\r|{:<10}|".format("#" * n), end="")

# Create a new image with the given size
def create_image(i, j):
    image = Image.new("1", (i, j))
    return image


# Get the pixel from the given image
def get_pixel(image, i, j):
    width, height = image.size
    if i > width or j > height:
        return None

    # Get Pixel
    pixel = image.getpixel((i, j))
    return pixel


def display_help():
    helptext = """
    lid - lo-fi image dithering

    syntax:
    lid [option]... [infile] [outfile]

    options:
    -f          output format. defaults to: png (recommended)
    -m          dither mode. defaults to: o4. other options: e, o9
    -q          image quality. defaults to: 90. only has effect on png/jpg

    arguments:
    infile      path to a valid image file on your system
    outfile     file name to output, do not include file extension

    example usage:
    lid -f jpeg -q 75 -m e ~/my_picture.png my_dithered_picture

    notes:
    it is highly recommended that users output their files as PNG. the file
    size will be significantly smaller than jpeg and the quality in most
    cases will be about the same.

    the quality flag mostly seems to effect jpegs and does not do a whole
    heap to PNG files, but play around with it as mileage may vary.
    """
    print(helptext)


# Get and validate the arguments
# returns dict or None
def parse_input():
    args = sys.argv[1:]
    argmap = {
        "-f": "png",       # output format
        "-q": 25,          # optimization quality
        "-m": "o4",         # mode, defaults to ordered
        "path": None,      # path to file
        "out": "lid_file" # output filename
    }
    if len(args) < 1:
        print("lid error: expected image filepath, received None")
        return None

    if "help" in args or "-h" in args or "--help" in args:
        display_help()
        sys.exit(1)

    item = None
    for x in args:
        if x in argmap:
            item = x
            continue
        elif item:
            argmap[item] = x
            item = None
        elif len(x) > 0 and x[0] == "-":
            print("lid error: unknown flag {}".format(x))
            return None
        else:
            if argmap["path"]:
                argmap["out"] = x
            else:
                argmap["path"] = x

    if not argmap["-f"].lower() in ["jpeg", "jpg", "png", "gif"]:
        print("lid error: invalid format requested with -f flag. Options: jpeg, png, gif. Defaults to png")
        return None
    elif argmap["-f"] == "jpg":
        argmap["-f"] = "jpeg"

    if not argmap["-m"] in ["o4", "e", "o9"]:
        print("lid error: invalid dither mode provided for -m. Options: o4, o9, e. Default: o4")
        return None

    try:
        argmap["-q"] = int(argmap["-q"])
        if argmap["-q"] < 1 or argmap["-q"] > 100:
            print("lid error: -q value must be between 1 and 100")
            return None
    except ValueError:
        print("lid error: -q value must be an integer")
        return None

    if not argmap["path"] or not os.path.isfile(argmap["path"]):
        print("lid error: invalid filepath")
        return None

    return argmap

if __name__ == "__main__":
    request = parse_input()
    if request:
        if request["-m"] == "o4":
            ordered_dither_4(request)
        elif request["-m"] == "o9":
            ordered_dither_9(request)
        elif request["-m"] == "e":
            error_dither(request)
        else:
            print("Unknown mode flag")
            sys.exit(2)
        sys.exit(0)
    sys.exit(1)