#!/usr/bin/env python3

# .SCRIPTDOC

descr = """

# Image processing utilities for the University of South Carolina Heterogeneous
# and Reconfigurable Computing (HeRC) research group. In particular, this tool
# is used for converting images to and from formats suitable for use with
# various embedded systems and FPGA platforms.

"""
descr = descr.replace("# ", " ")

# .DESCRIPTION

# Keep in mind that this script is not a tool for producing high-quality
# production images. It may introduce artifacting in some cases. Reading an
# image as a JPEG and saving it as a JPEG will not necessarily produce a
# bit-for-bit copy of that image. If you want high quality image manipulation
# of standardized formats, you should instead use ``convert`` from ImageMagic.

# Dependencies
# ------------
#
# To run this application, a Python 3 interpreter must be in ``$PATH`` and
# available via the name ``python3``. Additionally, the following Python 3
# packages must be installed:
#
# * ``PIL``
# * ``numpy``
# * ``skimage`` (only for canny edge detection)

# HeRC Header Format (H2F)
# ------------------------

# For the purpose of code-sharing, a number of in-house HeRC formats use the
# same header format, which is documented here. An H2F consists of 16 bytes
# of data, organized into the following fields. Unless otherwise noted, all
# integer fields are assumed to be unsigned, with large byte order.

# * Magic string [4 bytes] - the ASCII characters 'HeRC', hexadecimal
#   ``0x48655243``. This is used for validation.
# * Width in pixels [2 bytes] - unsigned int specifying image width.
# * Height in pixels [2 bytes] - unsigned int specifying image height.
# * File Format [1 byte] - see below for values.
# * Reserved [7 bytes] - reserved for future use, should be zeroed.

# The file format field is used to specify which of the HeRC image formats is
# in use. The following values are possible:

# =====   ==========
# Code    Format
# =====   ==========
# 0       ``de2rawh``
# 1       ``hif8``
# 2       ``hif1``
# =====   ==========

# Supported Image Formats
# -----------------------

# The following standard image formats are supported:
#
# * ``jpeg``
# * ``png``
# * ``ppm``
# * ``bmp`` (24-bit color bitmap)

# The remainder of this section documents non-standard image formats which
# are also supported.

# **NOTE**: To view a current listing of available image formats, as well as
# implementation status, please use the ``--list`` flag. The output from this
# option is generated directly from the ``herc-imtool`` image conversion
# matrix, and should always reflect what your copy of the tool is actually
# capable of.

# Altera DE2 RAW Image With Header (``de2rawh``)
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# This image format stores a contiguous sequence of 1-byte color channel
# values, organized into tuples comprising the RGB color channel of each pixel
# consecutively, such that each group of 24 bits specifies the color value
# for one single pixel. Note that this is only the case for the default format,
# and some variants (such as feature flag 1). Some feature flags (i.e. 2) use
# a diffrent format. See below for details.

# The ``de2rawh`` begins with a 16-byte H2F format header, and uses file format
# code 0.

# This image format is used on the Altera DE2-115 development board to load
# images for further processing. Note that this is not an Altera format, it is
# simply convenient for use with Altera's tooling. To that end, this image
# format may in fact be useful for other purposes, but it's primary purpose is
# to support the DE2 development board, in particular in the context of the
# University of South Carolina CSCE313 course.

# In particular, it is intended for this format to be convenient to utilize
# with C program, as it can be indexed directly as an array, being careful to
# start indexing at byte 16 (i.e. the first 16 bytes are ignored, so that the
# header is not treated as data).
#
# By convention, images in the ``de2rawh`` format use one of the following
# two file extensions: ``.de2rawh``, or ``.d2h``.

# **NOTE**: You can also have ``herc-imgtool`` pack your ``de2rawh`` file into
# a zip file suitable for use with Altera zipfs. This is accomplished by
# providing the ``--pack2zipfs`` flag.

# **NOTE**: The name of this image format is maintained for compatibility with
# herc-imgtool 0.0.1. This format would be more accurately dubbed HIF24.

# HeRC Image Format (HIF)
# ~~~~~~~~~~~~~~~~~~~~~~~

# HIF is actually a collection of different image formats. All HIF image
# types store pixel data contiguously in row-major format, with big
# byte-order.  HIF formats are suffixed with a number (i.e. ``hif8``), which is
# the number of bits per pixel. In this way, ``hif1`` refers to a
# 1-bit-per-pixel bitmap, ``hif8`` to a 1-byte-per-pixel greyscale image, and
# ``hif24`` to a 3-byte-per-pixel RGB image (identical to ``de2rawh``).

# All HIF types begin wtih a 16-byte H2F format header.

# The following HIF formats are supported:
#
# * ``hif1``
# * ``hif8``
# * ``hif24`` (``de2rawh``)

# Altera DE2 RAW Image Without Header (``de2raw``)
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# This is exactly identical to de2rawh, but sans the 16 byte header. The
# resolution is assumed to be 320x240 exactly. Note that if the image is not
# already 320x240, it will be automatically resized as if you had passed the
# options ``--resize 320x240 --bg_fill 0,0,0``. For a non-default background
# color, you must specify -``-resize`` and ``--bg_fill`` manually.

# By convention, images in the ``de2raw`` format use one of the following two
# file extensions: ``.de2raw``, ``.dat``, or ``.d2r``. Note that the ``.dat``
# file extension is actually associated (see ``infer_format()``) with the
# ``de2raw`` format, meaning using ``.dat`` as a file extension for input or
# output implies that that file should be treated as a ``de2raw`` file.
# To override this behaviour, use ``--input_format``.

# **NOTE**: You can also have ``herc-imgtool`` pack your ``de2raw`` file into a
# zip file suitable for use with Altera zipfs. This is accomplished by
# providing the ``--pack2zipfs`` flag.


# ``herc-imgtool`` In-Memory Format (``hercimf``)
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# This format is simply the in-memory image format for this program. It's
# on-disk representation is simply a pickled copy of the data structure which
# holds the image in memory. This is generally not useful except for debugging
# this program. Note that this format is not guaranteed to be stable - it
# should not be used for archival.

# The data structure itself is a dict with the following keys:

# * ``width`` - integer number of horizontal pixels
# * ``height`` - integer number of vertical pixels
# * ``version`` - the version string for the herc-imgtool which generated
#   this image.
# * ``pixels`` - numpy 2-dimensional array of pixel color values in RGB color.
#   I believe this is technically a 3-dimensional array in-memory, since color
#   channels can be accessed via an additional level of subscription.
# * ``greyscale`` - boolean, if true, then the red channel is intended to be
#   treated as greyscale color values, and the G and B color channels are
#   to be ignored. The ``pixels`` array retains the same format.
# * ``pixmap`` - boolean, if true, ``greyscale`` is ignored, and the image is
#   assumed to be a 1-bit-per-pixel image. The red channel is used to store
#   the 1s or 0s (as 8 bit ints) while the G and B channels are ignored.

# Note that the hercimf on-disk format is bz2 compressed before writing, to
# save space.

# ``herc-imgtool`` Architecture Overview
# --------------------------------------

# ``herc-imgtool`` utilizes a 4-stage image processing pipeline, which is
# orchestrated by the ``main()`` function. The stages of the pipeline are:

# 1. **image load** - the input image is loaded into the in-memory format for
# further processing. In the case of ``hercimf`` image, this simply involves
# unpickling the object. For other formats, dedicated conversion functions are
# utilized.

# 2. **image manipulation** - any image manipulations specified by the user are
# executed sequentially in a pre-defined order.  For example ``--resize`` and
# ``--stretch``. In the future, more manipulations (i.e. channel masking) may
# be added. The order in which manipulations run is not user-configurable.

# 3. **image store** - the finalized image is written to disk in the desired
# format. For ``hercimf`` images, this simply involves pickling the object and
# writing it out. For other images, dedicated conversion functions are
# utilized.

# 4. **post processing** - in some cases, it may be desirable to post-process
# the generated output file using a script of some kind. For example, the
# ``--pack2zipfs`` flag causes a ``zip`` command to be run on the generated
# output file to pack a zip file suitable for use with Altera zipfs. At
# present, this is the only post-processing step available, but others may be
# added in the future. Post-processing steps are executed sequentially, in a
# non-user-configurable order.

# Of particular note, stages 1 and 3 utilize the somewhat novel *conversion
# matrix* this construct is a hashtable of hashtables. The outer key is a
# format (i.e. ``png``). The inner key is either ``load`` or ``store``.  The
# value associated with the inner key is always a function object (function
# pointer in C parlance). Every such function must conform to the same calling
# conventions, and must be pure with respect to the overall application state
# (which resides in ``main()``. The conversion matrix saves considerable time
# and complexity as compared to a case statement or a long ``if/elif/else``
# block. It is also useful, because individual formats may support read-only,
# write-only, or read-write access, all of which can be encoded trivially in
# the conversion matrix.

# Stage 2 is relatively uninteresting, and roughly consists of executing a
# number of pure functions against the working state of the image in memory.

# Stage 4 is equally uninteresting, and generally consists of conditionally
# executing shell commands, with appropriate glue.

# .SYNTAX

# Run this script with the ``--help`` parameter for usage information.

# .AUTHOR

# Charles Daniels


# .LICENSE

# Copyright 2018 Charles Daniels

#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:

#  1. Redistributions of source code must retain the above copyright notice,
#  this list of conditions and the following disclaimer.

#  2. Redistributions in binary form must reproduce the above copyright notice,
#  this list of conditions and the following disclaimer in the documentation
#  and/or other materials provided with the distribution.

#  3. Neither the name of the copyright holder nor the names of its
#  contributors may be used to endorse or promote products derived from this
#  software without specific prior written permission.

#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.

# .CHANGELOG
#
# * 0.0.1
#
#   + 2018-02-12 - Charles Daniels
#
#       - Initial version with the following features:
#
#           * Letterbox and stretch resizing, with arbitrary BG fill color.
#
#           * Read and write support for jpeg, png, bmp, ppm, de2raw, de2rawh,
#             and hercimf format image files.
#
#           * Automatic packing of Altera zipfs-compatible zip files (via
#             ``--pack2zipfs``).
#
# * 0.0.2
#
#   + 2018-02-19 - Charles Daniels
#
#       - Added --canny
#       - Added --canny_sigma
#       - Added --greyscale option
#       - Added --threshold option
#
#   + 2018-02-20 - Charles Daniels
#
#       - Specify H2F and HIF formats.
#       - Added --display.
#       - Implemented HIF8 support.
#       - Added methods to handle generating and parsing H2F headers.
#
#   + 2018-02-22 - Charles Daniels
#
#       - Added HIF1 support.
#
# * 0.0.3


# .ENDOC


herc_imgtool_version = "0.0.3-git"

import os
import sys
import argparse
import numpy
import pickle
from PIL import Image
import bz2
import subprocess

def load_hercimf(path):
    im = None
    with open(path, 'rb') as f:
        im = f.read()
    im = bz2.decompress(im)
    im = pickle.loads(im)
    # saity check that the values we expect are here
    for k in ["width", "height", "version", "pixels"]:
        assert k in im
    return im

def store_hercimf(im, path):
    data = pickle.dumps(im)
    data = bz2.compress(data)
    with open(path, 'wb') as f:
        f.write(data)

def load_from_PIL(path):
    pil_image = Image.open(path)

    # ensure the image is in the RGB colorspace
    pil_image = pil_image.convert("RGB")

    # convert pixel data to a matrix
    pixels = numpy.array(pil_image)
    width, height = pil_image.size

    # populate hercimf data structure
    im = { "width": width, "height": height, "version": herc_imgtool_version,
            "pixels": pixels, "greyscale": False, "pixmap": False}
    return im

def store_with_PIL(im, path, fmt, **kwargs):
    pil_image = None
    if im["pixmap"]:
        im = pixmap2greyscale(im)

    if im["greyscale"]:
        pil_image = Image.fromarray(extract_greyscale(im))
        pil_image = pil_image.convert("L")
    else:
        pil_image = Image.fromarray(im["pixels"])
        pil_image = pil_image.convert("RGB")

    pil_image = pil_image.convert("RGB")
    pil_image.save(path, format=fmt, **kwargs)

def store_jpeg(im, path):
    store_with_PIL(im, path, "jpeg", quality=90)

def store_png(im, path):
    store_with_PIL(im, path, "png")

def store_bmp(im, path):
    store_with_PIL(im, path, "bmp")

def store_ppm(im, path):
    store_with_PIL(im, path, "ppm")

def load_de2rawh(path):
    return load_de2raw(path, useheader=True)

def parse_h2f(data):
    """parse_h2f

    Extract h2f header data as a dict from raw bytes.

    :param data:
    """

    header_data = {}

    header = data[0:16]
    magic = header[0:4]
    width = header[4:6]
    height = header[6:8]
    format_code = header[9:10]

    # cast back to int
    width = int.from_bytes(width, byteorder="big")
    height = int.from_bytes(height, byteorder="big")
    format_code = int.from_bytes(format_code, byteorder="big")

    # validate the magic string
    if magic != b'HeRC':
        sys.stderr.write("ERROR: incorrect magic! got '{}' expected '{}'\n"
                         .format(magic, b'HeRC'))
        raise ValueError

    header_data["magic"] = magic
    header_data["width"] = width
    header_data["height"] = height
    header_data["format_code"] = format_code

    return header_data

def make_h2f(width, height, format_code):
    header = bytearray(bytes(16))
    header[0:4] = b'HeRC'
    header[4:6] = int(width).to_bytes(2, 'big')
    header[6:8] = int(height).to_bytes(2, 'big')
    header[9:10] = int(format_code).to_bytes(1, 'big')

    assert(len(header) == 16)

    return bytes(header)

def load_hif8(path):
    # unpack data into memory for processing
    sys.stderr.write("INFO: Loading hif8 image from '{}'...\n".format(path))

    data = None
    with open(path, 'rb') as f:
        data = f.read()

    if data is None:
        sys.stderr.write("ERROR: no data read.\n")
        exit(1)

    greyscale = True
    pixmap = False

    raw_pixels = data[16:]
    header = parse_h2f(data)
    width = header["width"]
    height = header["height"]

    # populate pixels array
    pixels = numpy.empty(shape = (height, width, 3), dtype=numpy.uint8)
    length = width * height
    col = 0
    row = 0

    if len(raw_pixels) < length:
        sys.stderr.write("WARN: expected at least length {} but got {}\n"
                         .format(length, len(raw_pixels)))

    errcount = 0
    for i in range(length):
        try:
            pixels[row][col][0] = raw_pixels[i]
            pixels[row][col][1] = 0
            pixels[row][col][2] = 0
        except Exception:
            pixels[row][col][0] = 0
            pixels[row][col][1] = 0
            pixels[row][col][2] = 0
            errcount += 1

        # update subscripts
        col += 1
        if col >= width:
            col = 0
            row += 1

    if errcount > 0:
        sys.stderr.write("WARN: {} pixels were missing, ".format(errcount) +
                         "corrupted.\n")


    return { "width": width, "height": height, "version": herc_imgtool_version,
            "pixels": pixels, "greyscale" : greyscale, "pixmap": pixmap}

def load_de2raw(path, useheader=False):
    # unpack data into memory for processing
    sys.stderr.write("INFO: Loading de2raw image from '{}'...\n".format(path))

    data = None
    with open(path, 'rb') as f:
        data = f.read()

    if data is None:
        sys.stderr.write("ERROR: no data read.\n")
        exit(1)

    # handle useheader
    header = None
    raw_pixels = None
    magic = None
    width = None
    height = None

    # these are always False for HIF24
    greyscale = False
    pixmap = False

    if useheader:
        header = parse_h2f(data)
        raw_pixels = data[16:]
        width = header["width"]
        height = header["height"]
        format_code = header["format_code"]

    else:
        sys.stderr.write("INFO: de2raw decoder is running in headerless " +
                         "mode, header data will be faked\n")
        raw_pixels = data
        magic = b'HeRC'
        height = 240
        width = 320
        format_code = 0

    if format_code != 0:
        sys.stderr.write("INFO: bad format code '{}' for HIF24.\n"
                         .format(format_code))
        exit(1)

    # populate pixels array
    pixels = numpy.empty(shape = (height, width, 3), dtype=numpy.uint8)
    pixel_counter = 0
    length = width * height * 3
    col = 0
    row = 0

    if len(raw_pixels) < length:
        sys.stderr.write("WARN: expected at least length {} but got {}\n"
                         .format(length, len(raw_pixels)))

    errcount = 0
    for i in range(length):
        try:
            pixels[row][col][pixel_counter] = raw_pixels[i]
        except Exception:
            pixels[row][col][pixel_counter] = 0
            errcount += 1

        # update subscripts
        pixel_counter += 1
        if pixel_counter > 2:
            pixel_counter = 0
            col += 1
            if col >= width:
                col = 0
                row += 1

    if errcount > 0:
        sys.stderr.write("WARN: {} channels were missing, ".format(errcount) +
                         "corrupted.\n")


    return { "width": width, "height": height, "version": herc_imgtool_version,
            "pixels": pixels, "greyscale" : greyscale, "pixmap": pixmap}

def store_de2rawh(im, path):
        return store_de2raw(im, path, useheader=True)

def store_de2raw(im, path, useheader=False):
    format_code = 0
    with open(path, 'wb') as f:
        if useheader:
            # write header, if enabled
            f.write(make_h2f(im["width"], im["height"], format_code))
        else:
            # if we are running sans header, assert that the resolution is
            # 320x240

            if (im["height"] != 240) or (im["width"] != 320):
                sys.stderr.write("WARN: resolution '{}x{}' will be implicitly"
                                 .format(im["width"], im["height"]) +
                                 " resized to 320x240.\n")
                im = resize_image(im, "320x240")

        # write data
        for row in range(len(im["pixels"])):
            for col in range(len(im["pixels"][row])):
                for channel in range(len(im["pixels"][row][col])):
                    elem = int(im["pixels"][row][col][channel])
                    f.write(elem.to_bytes(1, 'big'))

def store_hif8(im, path):
    if not im["greyscale"]:
        im = to_greyscale(im)

    format_code = 1
    with open(path, 'wb') as f:
        f.write(make_h2f(im["width"], im["height"], format_code))

        # write data
        for row in range(len(im["pixels"])):
            for col in range(len(im["pixels"][row])):
                elem = int(im["pixels"][row][col][0])
                f.write(elem.to_bytes(1, 'big'))

def store_hif1(im, path):
    if not im["pixmap"]:
        im = to_pixmap(im)

    pixmap = extract_pixmap(im)

    format_code = 2
    with open(path, 'wb') as f:
        f.write(make_h2f(im["width"], im["height"], format_code))

        # write data
        bit_index = 0
        current_byte = 0
        for row in range(len(im["pixels"])):
            for col in range(len(im["pixels"][row])):
                elem = pixmap[row][col]
                if elem:
                    # this makes the bit-order what we would expect, so the
                    # leftmost bit comes earlier in the image.
                    current_byte += (1 << (7 - bit_index))

                bit_index += 1
                if bit_index > 7:
                    f.write(current_byte.to_bytes(1, 'big'))
                    current_byte = 0
                    bit_index = 0

        # pad out the end of the last byte with zeros
        while (bit_index < 8) and (bit_index != 0):
            current_byte += (1 << (bit_index))
            bit_index += 1
        f.write(current_byte.to_bytes(1, 'big'))

def load_hif1(path):
    # unpack data into memory for processing
    sys.stderr.write("INFO: Loading hif1 image from '{}'...\n".format(path))

    data = None
    with open(path, 'rb') as f:
        data = f.read()

    if data is None:
        sys.stderr.write("ERROR: no data read.\n")
        exit(1)

    greyscale = True
    pixmap = True

    raw_pixels = bytearray(data[16:])
    header = parse_h2f(data)
    width = header["width"]
    height = header["height"]

    if header["format_code"] != 2:
        sys.stderr.write("ERROR: bad format code for HIF1: '{}'\n"
                         .format(header["format_code"]))

    # populate pixels array
    pixels = numpy.empty(shape = (height, width, 3), dtype=numpy.uint8)
    length = (width * height)
    raw_length = len(raw_pixels)
    col = 0
    row = 0

    if len(raw_pixels) * 8 < length:
        sys.stderr.write("WARN: expected at least length {} but got {}\n"
                         .format(length, len(raw_pixels)))

    errcount = 0
    bit_index = 0
    for i in range(0, raw_length):
        the_val = bytes(1)
        try:
            the_val = raw_pixels[i]
        except Exception:
            #  the_val = bytes(1)
            the_val = 0
            errcount += 1

        # we convert the value into a binary string, because it makes it easy
        # to query out certain bits by their subscripts.
        assert(the_val < 256)
        the_val = str(bin(the_val))[2:]
        the_val = '0'*(8 - len(the_val)) + the_val
        assert(len(the_val) == 8)

        pixel = 0
        for bit_index in range(0, 8):
            pixel = int(the_val[bit_index], 2)

            assert row < height
            assert col < width

            pixels[row][col][0] = pixel
            pixels[row][col][1] = 0
            pixels[row][col][2] = 0

            # update subscripts
            col += 1
            if col >= width:
                col = 0
                row += 1
                if row >= height:
                    # we are just reading padding now
                    break
        if row >= height:
            # break out of the outer loop also
            break


    if errcount > 0:
        sys.stderr.write("WARN: {} pixels were missing, ".format(errcount) +
                         "corrupted.\n")


    return { "width": width, "height": height, "version": herc_imgtool_version,
            "pixels": pixels, "greyscale" : greyscale, "pixmap": pixmap}

def parse_size(size):
    try:
        size = [int(x) for x in size.split('x')]
        width = size[0]
        height = size[1]
        return (width, height)
    except Exception:
        sys.stderr.write("ERROR: Failed to parse WxH size '{}'\n".format(size))
        raise ValueError


def resize_image(im, size, pad_color=(0, 0, 0), stretch=False):
    """resize_image

    Resize an image to the specified size. Size should be a string in WxH
    format, and is parsed by parse_size().


    :param im: hercimf format image
    :param size: size in WxH format, as a string
    :pad_color: (R, G, B) tupe specifying bg fill color, if one is needed.
    :stretch: if asserted, aspect ratio is not preserved and no letterboxing
    occurs. Otherwise, aspect ratio is preserved, and image is letterboxed
    to required dimensions using the specified fill color.
    """

    width, height = parse_size(size)

    # convert to PIL image for resizing
    pil_image = Image.fromarray(im["pixels"])
    pil_image.thumbnail((width, height), Image.ANTIALIAS)
    pil_image = pil_image.convert("RGB")

    pixels = None

    if stretch:
        # aspect ratio is not preserved
         pixels = numpy.array(pil_image)
    else:
        # PIL will not automatically pad (letterbox) for us, so all we have is
        # a guaranteed upper bound on the output image size.
        bg = Image.new('RGB', (width, height), pad_color);
        offset = ((bg.width - pil_image.width) // 2,
                  (bg.height - pil_image.height) // 2)
        bg.paste(pil_image, offset)

        # generate pixels
        pixels = numpy.array(bg)

    return { "width": width, "height": height, "version": herc_imgtool_version,
            "pixels": pixels, "greyscale":im["greyscale"],
            "pixmap": im["pixmap"]}

def to_greyscale(im):
    pil_image = Image.fromarray(im["pixels"])
    pil_image = pil_image.convert('L')
    greyscale_pixels =  numpy.array(pil_image)
    for x in range(im["width"]):
        for y in range(im["height"]):
            im["pixels"][y][x][0] = greyscale_pixels[y][x]
            im["pixels"][y][x][1] = 0
            im["pixels"][y][x][2] = 0

    im["greyscale"] = True
    return im

def threshold(im, threshold_val="1,1,1"):
    """threshold

    Convert the image to 1-bit-per-pixel, stored in the red channel. All pixels
    which exceed the given threshold are passed through, and the rest are
    zeroed.  If the image is greyscale, then only the R channel of the
    threshold is used.

    :param im:
    :param threshold_val: RGB color tuple - if all of the input color channels
    for
    a pixel are >= the channel in this tuple, it is passed through.
    """

    if im["pixmap"]:
        # It is meaningless to threshold a pixmap.
        return im

    threshold_val = parse_RGB(threshold_val)
    tr, tg, tb = threshold_val
    for x in range(im["width"]):
        for y in range(im["height"]):
            r, g, b = im["pixels"][y][x]
            if im["greyscale"]:
                # for greyscale images, we only want to consider the red
                # channel.
                g, b = (255, 255)
            if (r >= tr) and (g >= tg) and (b >= tb):
                im["pixels"][y][x][0] = 0;
                im["pixels"][y][x][1] = 0;
                im["pixels"][y][x][2] = 0;

    return im

def to_pixmap(im):
    """to_pixmap

    Convers the image to a boolean array. All pixel values which are 0 are
    output as 0, and all others are 1. The pixmap is stored in the R channel
    and the G and B channels are zeroed.

    :param im:
    """
    for x in range(im["width"]):
        for y in range(im["height"]):
            r, g, b = im["pixels"][y][x]
            if (r + g + b) > 0:
                im["pixels"][y][x] = [1,0,0]
            else:
                im["pixels"][y][x] = [0,0,0]

    im["pixmap"] = True
    return im

def pixmap2greyscale(im):
    """pixmap2greyscale

    Converts an image to greyscale with 1 being 255 and 0 being 0.

    :param im:
    """

    if not im["pixmap"]:
        im = to_pixmap(im)
    for x in range(im["width"]):
        for y in range(im["height"]):
            r, g, b = im["pixels"][y][x]
            if (r + g + b) > 0:
                im["pixels"][y][x] = [255,0,0]
            else:
                im["pixels"][y][x] = [0,0,0]

    im["pixmap"] = False
    im["greyscale"] = True

    return im

def extract_pixmap(im):
    """extract_pixmap

    Extract and return 2-d array of bools representing image pixmap. If the
    image is not a pixmap already, then it will be silently converted to one
    with the threshold ``127,127,127``.

    :param im:
    """

    if not im["pixmap"]:
        im = to_pixmap(im)
    pixmap = numpy.zeros(shape=(im["height"], im["width"]), dtype=bool)
    for col in range(im["width"]):
        for row in range(im["height"]):
            r, g, b = im["pixels"][row][col]
            pixmap[row][col] = (r > 0)

    return pixmap



def extract_greyscale(im):
    """extract_greyscale

    Return 2-d matrix of greyscale pixel values with the G and B channels
    removed explicitly. If the image is not greyscale already, it is converted
    to such.

    :param im:
    """

    if not im["greyscale"]:
        im = to_greyscale(im)

    greyscale_pixels = numpy.zeros(shape=(im["height"], im["width"]))
    for x in range(im["width"]):
        for y in range(im["height"]):
            greyscale_pixels[y][x] = im["pixels"][y][x][0]

    return greyscale_pixels


def canny_filter(im, sigma):
    try:
        import skimage.feature

        greymap = extract_greyscale(im)
        bitmap = skimage.feature.canny(greymap, sigma)
        for x in range(len(bitmap)):
            for y in range(len(bitmap[x])):
                val = bitmap[x][y]
                im['pixels'][x][y] = [val, 0, 0]

        im["pixmap"] = True
        return im

    except ImportError:
        sys.stderr.write("ERROR: could not import skimage.feature, is " +
                         "scikit-image installed?")
        exit(1)

def infer_format(s):
    """infer_format

    Given some string, guess the image format if possible. For example,
    given s = "foo.jpeg", this function would return "jpeg".

    :param s:
    """

    if s is None:
        return "dummy"

    s = s.lower().strip()
    if s.endswith("jpeg"):
        return "jpeg"
    elif s.endswith("jpg"):
        return "jpeg"
    elif s.endswith("png"):
        return "png"
    elif s.endswith("bmp"):
        return "bmp"
    elif s.endswith("hercimf"):
        return "hercimf"
    elif s.endswith("de2raw"):
        return "de2raw"
    elif s.endswith("d2r"):
        return "de2raw"
    elif s.endswith("de2rawh"):
        return "de2rawh"
    elif s.endswith("d2h"):
        return "de2rawh"
    elif s.endswith("ppm"):
        return "ppm"
    elif s.endswith("dat"):
        return "de2raw"
    elif s.endswith("hif24"):
        return "de2rawh"
    elif s.endswith("hif8"):
        return "hif8"
    elif s.endswith("hif1"):
        return "hif1"
    else:
        return None

def parse_RGB(s):
    """parse_RGB

    Parse a string of the form R,G,B into a color tuple of ints.

    :param s:
    """
    try:
        R, G, B = [numpy.uint8(x) for x in s.split(",")]
    except Exception:
        sys.stderr.write("ERROR: invalid RGB color tuple '{}'.\n".format(s))
        raise ValueError

    return (R, G, B)

def display_image(im):

    try:
        import matplotlib.pyplot

        if im["pixmap"]:
            matplotlib.pyplot.imshow(extract_pixmap(im), cmap="gray")
        elif im["greyscale"]:
            matplotlib.pyplot.imshow(extract_greyscale(im), cmap="gray")
        else:
            matplotlib.pyplot.imshow(im["pixels"])

        matplotlib.pyplot.show()

    except ImportError:
        sys.stderr.write("WARNING: cannot display image, missing matplotlib.\n")

def do_nothing(**kwargs):
    pass

# The conversion matrix associates an image format with a dict which has the
# keys load and store. The values of these keys should be functions.  The load
# function should accept a path on disk, and return the image in herc-imgtool's
# in-memory format. The store function accepts as it's only arguments an image
# in herc-imgtool's in-memory format and a path on disk, to which the image is
# written.

conversion_matrix = {}
conversion_matrix["jpg"] = {"load": load_from_PIL, "store": store_jpeg}
conversion_matrix["jpeg"] = conversion_matrix["jpg"]
conversion_matrix["hercimf"] = {"load": load_hercimf, "store": store_hercimf}
conversion_matrix["png"] = {"load": load_from_PIL, "store": store_png}
conversion_matrix["bmp"] = {"load": load_from_PIL, "store": store_bmp}
conversion_matrix["ppm"] = {"load": load_from_PIL, "store": store_ppm}
conversion_matrix["de2raw"] = {"load": load_de2raw, "store": store_de2raw}
conversion_matrix["de2rawh"] = {"load": load_de2rawh, "store": store_de2rawh}
conversion_matrix["hif8"] = {"load": load_hif8, "store": store_hif8}
conversion_matrix["hif1"] = {"load": load_hif1, "store": store_hif1}
conversion_matrix["dummy"] = {"load": do_nothing, "store": do_nothing}

def main():

    parser = argparse.ArgumentParser(description=descr)

    parser.add_argument("--version", "-v", action="version",
                        version="%(prog)s " + herc_imgtool_version)

    parser.add_argument("--input", "-i", default=None,
                        help="Input image to process.")

    parser.add_argument("--output", "-o", default=None,
                        help="Output path to write result to.")

    parser.add_argument("--resize", "-r", default=None,
                        help="Re-scale the image to the specified size. Use" +
                        " the format WxH. Aspect ratio is preserved, and " +
                        "the image is letterboxed using the bg_fill color." +
                        " (default: do not resize)")

    parser.add_argument("--stretch", "-s", default=None,
                        help="Re-scape the image as with --resize, but " +
                        "stretch the image to fit the desired resolution, " +
                        "rather than preserving aspect ratio. This option " +
                        "has no effect when combined with --resize." +
                        " (default: do not stretch)")

    parser.add_argument("--bg_fill", "-g", default="0,0,0",
                        help="RGB color tuple of the form R,B,G. This color " +
                        "is used for any padding or background which needs " +
                        "to be generated. (default: 0,0,0)")

    parser.add_argument("--input_format", "-f", default="auto",
                        help="Specify format for input image. " +
                        "(default: auto)")

    parser.add_argument("--output_format", "-u", default="auto",
                        help="Specify format for output image. " +
                        "(default: auto)")

    parser.add_argument("--list_formats", "-l", action="store_true",
                        default=False, help="Print supported image formats" +
                        " then exit.")

    parser.add_argument("--pack2zipfs", "-z", action="store_true",
                        default=False, help="Assert to pack the output " +
                        "file into an uncompressd zip file suitable for use " +
                        "with Altera zipfs.")

    parser.add_argument("--canny", "-c", action="store_true", default=False,
                        help="Apply a canny image filter, requires " +
                        "scikit-image to be installed. Use --canny_sigma" +
                        " to modify the value of sigma. Implies " +
                        "--greyscale. This runs after --threshold.")

    parser.add_argument("--canny_sigma", "-S", type=float, default=1.0,
                        help="Specify value of sigma for canny image filter." +
                        " (default: 1.0)")

    parser.add_argument("--greyscale", "-G", default=False, action="store_true",
                        help="Convert image to greyscale and store the " +
                        "result in the red channel only.")

    parser.add_argument("--threshold", "-T", default=None,
                        help="Drop all pixels whose RGB color values are " +
                        "less than the color values specified here. This " +
                        "parameter takes an RGB color tuple of the form " +
                        "R,G,B, for example 1,2,3. Runs after --greyscale.")

    parser.add_argument("--display", "-D", default=False, action="store_true",
                        help="As an image processing step, display the " +
                        "working state of the image after all other " +
                        "image manipulation steps have finished.")


    args = parser.parse_args()

    # list formats and exit
    if args.list_formats:
        for fmt in conversion_matrix:
            c = conversion_matrix[fmt]
            rw = "NOT IMPLEMENTED"
            if (c["load"] is None) and (c["store"] is not None):
                rw = "WRITE ONLY"
            elif (c["load"] is not None) and (c["store"] is None):
                rw = "READ ONLY"
            elif (c["load"] is not None) and (c["store"] is not None):
                rw = "READ + WRITE"
            print("{:20} {}".format(fmt, rw))

        exit(0)

    if ((not args.output) or (not args.input)) and not args.display:
        sys.stderr.write("ERROR: --input and --output must be specified.\n")
        exit(1)

    # parse the BG fill color
    args.bg_fill = parse_RGB(args.bg_fill)

    # infer input and output format if required
    if args.input_format == "auto":
        args.input_format = infer_format(args.input)
        if args.input_format is None:
            sys.stderr.write("ERROR: could not infer format from '{}'\n"
                             .format(args.input))
            exit(1)
        else:
            sys.stderr.write("INFO: inferred format '{}' from '{}'\n"
                             .format(args.input_format, args.input))

    if args.output_format == "auto":
        args.output_format = infer_format(args.output)
        if args.output_format is None:
            sys.stderr.write("ERROR: could not infer format from '{}'\n"
                             .format(args.output))
            exit(1)
        else:
            sys.stderr.write("INFO: inferred format '{}' from '{}'\n"
                             .format(args.output_format, args.output))

    # make sure input and output were specified
    if args.input is None:
        sys.systderr.write("ERROR: no input file specified\n")
        exit(1)
    if (args.output is None) and not args.display:
        sys.stderr.write("ERROR: no output file specified\n")
        exit(1)


    # select the appropriate load and store functions from the conversion
    # matrix
    load_func = conversion_matrix[args.input_format]["load"]
    store_func = conversion_matrix[args.output_format]["store"]
    if load_func is None:
        sys.stderr.write("ERROR: no load function for format '{}'\n"
                         .format(args.input_format))
        exit(1)
    if store_func is None:
        sys.stderr.write("ERROR: no store function for format '{}'\n"
                         .format(args.output_format))
        exit(1)

    sys.stderr.write("INFO: loading image...\n")
    working_image = \
            conversion_matrix[args.input_format]["load"](path=args.input)

    sys.stderr.write("INFO: processing image...\n")
    # perform the resize operation, if requested
    if args.resize is not None:
        sys.stderr.write("INFO: resizing image to '{}'...\n"
                         .format(args.resize))
        working_image = resize_image(working_image,
                                     args.resize,
                                     args.bg_fill,
                                     stretch=False)
    if args.stretch is not None:
        # because this executes after letterbox resize, this has no effect
        # if the letterbox resize option is selected.
        sys.stderr.write("INFO: resizing image to '{}'...\n"
                         .format(args.resize))
        working_image = resize_image(working_image,
                                     args.resize,
                                     args.bg_fill,
                                     stretch=True)

    if args.greyscale:
        sys.stderr.write("INFO: converting image to greyscale... \n")
        working_image = to_greyscale(working_image)

    if args.threshold is not None:
        sys.stderr.write("INFO: thresholding image... \n")
        working_image = threshold(working_image, args.threshold)

    if args.canny:
        sys.stderr.write("INFO: performing canny edge detection... \n")
        working_image = canny_filter(working_image, args.canny_sigma)

    if args.display:
        sys.stderr.write("INFO: displaying image... \n")
        display_image(working_image)

    sys.stderr.write("INFO: saving output... \n")
    conversion_matrix[args.output_format]["store"](im=working_image,
                                                   path=args.output)

    if args.pack2zipfs:
        sys.stderr.write("INFO: packing zipfs file... \n")
        zip_name = args.output + ".zip"
        if '.' in args.output:
            zip_name = '.'.join(args.output.split('.')[:-1]) + ".zip"
        subprocess.check_output(["zip", "-0", zip_name, args.output])
        sys.stderr.write("INFO: zip file packed to '{}'\n".format(zip_name))



if __name__ == "__main__":
    main()
