#!/venv/bin/python3

# pipdither
# Error Diffusing image dithering.
# Now with serpentine scanning.

# See http://www.efg2.com/Lab/Library/ImageProcessing/DHALF.TXT

# http://www.python.org/doc/2.4.4/lib/module-bisect.html
from bisect import bisect_left


import png


def dither(
    out,
    input,
    bitdepth=1,
    linear=False,
    defaultgamma=1.0,
    targetgamma=None,
    cutoff=0.5,  # see :cutoff:default
):
    """Dither the input PNG `inp` into an image with a smaller bit depth
    and write the result image onto `out`.  `bitdepth` specifies the bit
    depth of the new image.

    Normally the source image gamma is honoured (the image is
    converted into a linear light space before being dithered), but
    if the `linear` argument is true then the image is treated as
    being linear already: no gamma conversion is done (this is
    quicker, and if you don't care much about accuracy, it won't
    matter much).

    Images with no gamma indication (no ``gAMA`` chunk) are normally
    treated as linear (gamma = 1.0), but often it can be better
    to assume a different gamma value: For example continuous tone
    photographs intended for presentation on the web often carry
    an implicit assumption of being encoded with a gamma of about
    0.45 (because that's what you get if you just "blat the pixels"
    onto a PC framebuffer), so ``defaultgamma=0.45`` might be a
    good idea.  `defaultgamma` does not override a gamma value
    specified in the file itself: It is only used when the file
    does not specify a gamma.

    If you (pointlessly) specify both `linear` and `defaultgamma`,
    `linear` wins.

    The gamma of the output image is, by default, the same as the input
    image.  The `targetgamma` argument can be used to specify a
    different gamma for the output image.  This effectively recodes the
    image to a different gamma, dithering as we go.  The gamma specified
    is the exponent used to encode the output file (and appears in the
    output PNG's ``gAMA`` chunk); it is usually less than 1.

    """

    # Encoding is what happened when the PNG was made (and also what
    # happens when we output the PNG).  Decoding is what we do to the
    # source PNG in order to process it.

    # The dithering algorithm is not completely general; it
    # can only do bit depth reduction, not arbitrary palette changes.
    import operator

    maxval = 2 ** bitdepth - 1
    r = png.Reader(file=input)

    _, _, pixels, info = r.asDirect()
    planes = info["planes"]
    # :todo: make an Exception
    assert planes == 1
    width = info["size"][0]
    sourcemaxval = 2 ** info["bitdepth"] - 1

    if linear:
        gamma = 1
    else:
        gamma = info.get("gamma") or defaultgamma

    # Calculate an effective gamma for input and output;
    # then build tables using those.

    # `gamma` (whether it was obtained from the input file or an
    # assumed value) is the encoding gamma.
    # We need the decoding gamma, which is the reciprocal.
    decode = 1.0 / gamma

    # `targetdecode` is the assumed gamma that is going to be used
    # to decoding the target PNG.
    # Note that even though we will _encode_ the target PNG we
    # still need the decoding gamma, because
    # the table we use maps from PNG pixel value to linear light level.
    if targetgamma is None:
        targetdecode = decode
    else:
        targetdecode = 1.0 / targetgamma

    incode = build_decode_table(sourcemaxval, decode)

    # For encoding, we still build a decode table, because we
    # use it inverted (searching with bisect).
    outcode = build_decode_table(maxval, targetdecode)

    # The table used for choosing output codes.  These values represent
    # the cutoff points between two adjacent output codes.
    # The cutoff parameter can be varied between 0 and 1 to
    # preferentially choose lighter (when cutoff > 0.5) or
    # darker (when cutoff < 0.5) values.
    # :cutoff:default: The default for this used to be 0.75, but
    # testing by drj on 2021-07-30 showed that this produces
    # banding when dithering left-to-right gradients;
    # test with:
    # priforgepng grl | priditherpng | kitty icat
    choosecode = list(zip(outcode[1:], outcode))
    p = cutoff
    choosecode = [x[0] * p + x[1] * (1.0 - p) for x in choosecode]

    rows = repeat_header(pixels)
    dithered_rows = run_dither(incode, choosecode, outcode, width, rows)
    dithered_rows = remove_header(dithered_rows)

    info["bitdepth"] = bitdepth
    info["gamma"] = 1.0 / targetdecode
    w = png.Writer(**info)
    w.write(out, dithered_rows)


def build_decode_table(maxval, gamma):
    """Build a lookup table for decoding;
    table converts from pixel values to linear space.
    """

    assert maxval == int(maxval)
    assert maxval > 0

    f = 1.0 / maxval
    table = [f * v for v in range(maxval + 1)]
    if gamma != 1.0:
        table = [v ** gamma for v in table]
    return table


def run_dither(incode, choosecode, outcode, width, rows):
    """
    Run an serpentine dither.
    Using the incode and choosecode tables.
    """

    # Errors diffused downwards (into next row)
    ed = [0.0] * width
    flipped = False
    for row in rows:
        # Convert to linear...
        row = [incode[v] for v in row]
        # Add errors...
        row = [e + v for e, v in zip(ed, row)]

        if flipped:
            row = row[::-1]
        targetrow = [0] * width

        for i, v in enumerate(row):
            # `it` will be the index of the chosen target colour;
            it = bisect_left(choosecode, v)
            targetrow[i] = it
            t = outcode[it]
            # err is the error that needs distributing.
            err = v - t

            # Sierra "Filter Lite" distributes          * 2
            # as per this diagram.                    1 1
            ef = err * 0.5
            # :todo: consider making rows one wider at each end and
            # removing "if"s
            if i + 1 < width:
                row[i + 1] += ef
            ef *= 0.5
            ed[i] = ef
            if i:
                ed[i - 1] += ef

        if flipped:
            ed = ed[::-1]
            targetrow = targetrow[::-1]
        yield targetrow
        flipped = not flipped


WARMUP_ROWS = 32


def repeat_header(rows):
    """Repeat the first row, to "warm up" the error register."""
    for row in rows:
        yield row
        for _ in range(WARMUP_ROWS):
            yield row
        break
    yield from rows


def remove_header(rows):
    """Remove the same number of rows that repeat_header added."""

    for _ in range(WARMUP_ROWS):
        next(rows)
    yield from rows


def main(argv=None):
    import sys

    # https://docs.python.org/3.5/library/argparse.html
    import argparse

    parser = argparse.ArgumentParser()

    if argv is None:
        argv = sys.argv

    progname, *args = argv

    parser.add_argument("--bitdepth", type=int, default=1, help="bitdepth of output")
    parser.add_argument(
        "--cutoff",
        type=float,
        default=0.5,
        help="cutoff to select adjacent output values",
    )
    parser.add_argument(
        "--defaultgamma",
        type=float,
        default=1.0,
        help="gamma value to use when no gamma in input",
    )
    parser.add_argument("--linear", action="store_true", help="force linear input")
    parser.add_argument(
        "--targetgamma",
        type=float,
        help="gamma to use in output (target), defaults to input gamma",
    )
    parser.add_argument(
        "input", nargs="?", default="-", type=png.cli_open, metavar="PNG"
    )

    ns = parser.parse_args(args)

    return dither(png.binary_stdout(), **vars(ns))


if __name__ == "__main__":
    main()