#!/usr/bin/env python

"""
pipcolours - extract all colours present in source image.
Produces a PNG that has each colour exactly once.

Also performs grey and alpha reduction when possible:
if all colour pixels are grey, output PNG is grey;
if all pixels are opaque, output PNG has no alpha channel.
"""

import sys


def colours(out, inp):
    import itertools
    import png

    r = png.Reader(file=inp)
    _,_,pixels,info = r.asDirect()
    planes = info['planes']
    bitdepth = info['bitdepth']

    col = set()
    for row in pixels:
        col = col.union(png.group(row, planes))

    col, planes = channel_reduce(col, planes, bitdepth)

    col = sorted(col)
    col = list(itertools.chain(*col))
    width = len(col) // planes
    greyscale = planes in (1,2)
    alpha = planes in (2,4)
    w = png.Writer(width, 1,
        bitdepth=bitdepth, greyscale=greyscale, alpha=alpha)
    w.write(out, [col])


def channel_reduce(col, planes, bitdepth):
    """Attempt to reduce the number of channels in the set of colours."""
    col, planes = reduce_grey(col, planes)
    col, planes = reduce_alpha(col, planes, bitdepth)
    return col, planes


def reduce_grey(col, planes):
    """
    Reduce a colour image to grey if
    all intensities match in all pixels.
    """
    if planes >= 3:
        def isgrey(c):
            return c[0] == c[1] == c[2]
        if min(map(isgrey, col)) == True:
            # Every colour is grey.
            col = set(map(lambda x: x[0::3], col))
            planes -= 2
    return col, planes


def reduce_alpha(col, planes, bitdepth):
    """
    Remove alpha channel if all pixels are fully opaque.
    """
    maxval = 2 ** bitdepth - 1
    if planes in (2, 4):
        def isopaque(c):
            return c[-1] == maxval
        if min(map(isopaque, col)) == True:
            # Every pixel is opaque
            col = set(map(lambda x: x[:-1], col))
            planes -= 1
    return col, planes


def ensure_binary_stdin():
    """
    Ensure sys.stdin returns bytes.
    """

    try:
        # Ensure sys.stdin returns bytes
        sys.stdin = sys.stdin.buffer
    except AttributeError:
        # Probably Python 2, where bytes are strings.
        pass


def ensure_binary_stdout():
    """
    Ensure sys.stdout accepts bytes.
    """

    # First there is a Python3 issue.
    try:
        sys.stdout = sys.stdout.buffer
    except AttributeError:
        # Probably Python 2, where bytes are strings.
        pass

    # On Windows the C runtime file orientation needs changing.
    if sys.platform == "win32":
        import msvcrt
        import os
        msvcrt.setmode(sys.stdout.fileno(), os.O_BINARY)


def main(argv=None):
    if argv is None:
        argv = sys.argv

    argv = argv[1:]
    if len(argv) > 0:
        f = open(argv[0], 'rb')
    else:
        ensure_binary_stdin()
        f = sys.stdin
    ensure_binary_stdout()
    return colours(sys.stdout, f)


if __name__ == '__main__':
    main()
