#!/usr/bin/env python2
"""
Create collages from multiple image files.
"""
import sys, os, argparse, glob
try:
    from PIL import Image
except ImportError:
    print >>sys.stderr, "This program requires PIL (the Python Imaging Library) or Pillow."
    sys.exit(1)
Image.MAXBLOCK = 64 * 1024 * 1024

def jpeg_quality(x):
    x = int(x)
    if 1 <= x < 100: return x
    raise ValueError("JPEG quality value out of range")

def coord(x):
    w,h = (int(x.strip()) for x in x.lower().replace('x', ',').split(','))
    return (w,h)

def color(x):
    x = x.lower().lstrip('#')
    try:
        x = {
            "white":   "ffffff",
            "silver":  "c0c0c0",
            "gray":    "808080",
            "black":   "000000",
            "red":     "ff0000",
            "maroon":  "800000",
            "yellow":  "ffff00",
            "olive":   "808000",
            "lime":    "00ff00",
            "green":   "008000",
            "aqua":    "00ffff",
            "teal":    "008080",
            "blue":    "0000ff",
            "navy":    "000080",
            "fuchsia": "ff00ff",
            "purple":  "800080",
        }[x]
    except KeyError:
        pass
    if len(x) == 1: x = 6 * x
    if len(x) == 2: x = 3 * x
    if len(x) == 3: x = 2 * x[0] + 2 * x[1] + 2 * x[2]
    if len(x) != 6: raise ValueError("invalid color code '#%s'" % x)
    return tuple(int(x[n:n+2], 16) for n in (0,2,4))

def median(x):
    x = sorted(x)
    l = len(x)
    return x[(l - ((~l) & 1)) / 2]

def muldiv(a, b, c):
    return (a * b + c/2) / c


class CollageImage(object):
    def __init__(self, filename, crop_pos=50):
        try:
            xname, xcrop = filename.rsplit(':', 1)
            crop_pos = int(xcrop)
            filename = xname
        except ValueError:
            pass
        self.filename = filename
        self.img = Image.open(self.filename)
        self.size = self.img.size
        self.crop_pos = crop_pos

    def resize(self, target, axis, keep_aspect=True):
        current = self.size[axis]
        if current != target:
            other = self.size[1 - axis]
            if keep_aspect:
                other = muldiv(other, target, current)
            self.size = (other, target) if axis else (target, other)
        return self.size

    def process(self, draft=False):
        if draft:
            self.img.draft('RGB', self.size)
        w, h = min((self.img.size[0], muldiv(self.size[1], self.img.size[0], self.size[0])),
                   (muldiv(self.size[0], self.img.size[1], self.size[1]), self.img.size[1]))
        x0 = muldiv(self.img.size[0] - w, self.crop_pos, 100)
        y0 = muldiv(self.img.size[1] - h, self.crop_pos, 100)
        img = self.img.crop((x0, y0, x0 + w, y0 + h))
        if draft:
            return img.resize(self.size, Image.NEAREST).convert('RGB')
        elif img.size > self.size:
            return img.convert('RGB').resize(self.size, Image.ANTIALIAS)
        else:
            return img.convert('RGB').resize(self.size, Image.BICUBIC)

################################################################################

def SetGroupSize(group, target, axis, keep_aspect=False):
    for img in group:
        img.resize(target, axis, keep_aspect)

def DistributeGroup_Static(group, target, axis, border=0):
    n = len(group)
    remain = target - (n - 1) * border
    for img in group:
        this_size = muldiv(remain, 1, n)
        img.resize(this_size, axis, False)
        remain -= this_size
        n -= 1
    assert remain == 0

def DistributeGroup_Dynamic(group, target, axis, border=0):
    remain_new = target - (len(group) - 1) * border
    remain_old = sum(i.size[axis] for i in group)
    for img in group:
        this_old = img.size[axis]
        this_new = muldiv(this_old, remain_new, remain_old)
        img.resize(this_new, axis, False)
        remain_old -= this_old
        remain_new -= this_new
    assert remain_new == 0

def Layout_FixedSize_Static(imgs, major_axis, minor_axis, outsize, border=0):
    n = len(imgs)
    remain = outsize[major_axis] - (n - 1) * border
    for group in imgs:
        this_size = muldiv(remain, 1, n)
        SetGroupSize(group, this_size, major_axis)
        DistributeGroup_Static(group, outsize[minor_axis], minor_axis, border)
        remain -= this_size
        n -= 1
    assert remain == 0

def Layout_FixedSize_Dynamic(imgs, major_axis, minor_axis, outsize, border=0):
    group_major_sizes = [median(i.size[major_axis] for i in group) for group in imgs]
    remain_new = outsize[major_axis] - (len(imgs) - 1) * border
    remain_old = sum(group_major_sizes)
    for group, group_major_size in zip(imgs, group_major_sizes):
        this_new = muldiv(group_major_size, remain_new, remain_old)
        SetGroupSize(group, this_new, major_axis, True)
        DistributeGroup_Dynamic(group, outsize[minor_axis], minor_axis, border)
        remain_old -= group_major_size
        remain_new -= this_new
    assert remain_new == 0

def Layout_VariableSize_Static(imgs, major_axis, minor_axis, border=0):
    # estimate the group major and and total minor sizes
    group_major_size = median(median(i.size[major_axis] for i in group) for group in imgs)
    total_minor_size = median(sum(muldiv(i.size[minor_axis], group_major_size, i.size[major_axis]) for i in group) + (len(group) - 1) * border for group in imgs)
    # assign these sizes to all groups
    for group in imgs:
        SetGroupSize(group, group_major_size, major_axis)
        DistributeGroup_Static(group, total_minor_size, minor_axis, border)
    # compute total size
    outsize = 2 * [total_minor_size]
    outsize[major_axis] = (group_major_size + border) * len(imgs) - border
    return tuple(outsize)

def Layout_VariableSize_Dynamic(imgs, major_axis, minor_axis, border=0):
    # first, turn all groups into rectangles, preserving aspect ratios
    group_minor_sizes = []
    for group in imgs:
        group_major_size = median(i.size[major_axis] for i in group)
        group_minor_sizes.append(sum(i.resize(group_major_size, major_axis, True)[minor_axis] for i in group) + (len(group) - 1) * border)
    # then, scale all groups to a common minor size
    total_minor_size = median(group_minor_sizes)
    total_major_size = 0
    for group, group_minor_size in zip(imgs, group_minor_sizes):
        # first, scale the major size accordingly
        group_major_size = muldiv(group[0].size[major_axis], total_minor_size, group_minor_size)
        SetGroupSize(group, group_major_size, major_axis)
        total_major_size += group_major_size
        # then, distribute the items along the group
        DistributeGroup_Dynamic(group, total_minor_size, minor_axis, border)
    # compute total size
    outsize = 2 * [total_minor_size]
    outsize[major_axis] = total_major_size + (len(imgs) - 1) * border
    return tuple(outsize)

################################################################################

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        usage="%(prog)s [OPTIONS...] -o <OUTPUT> <IMAGE_A1> <IMAGE_A2> [...] <IMAGE_An> , <IMAGE_B1> [...]",
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("infiles", metavar="IMAGE_xx[:PERCENT]", nargs='+',
                        help="input files; use ',', ';' or '/' to separate rows or columns; append ':0' to ':100' to specify crop position if cropping is to be done")
    g = parser.add_argument_group("output options")
    g.add_argument("-o", "--output", metavar="FILE",
                   help="output image file name")
    g.add_argument("-p", "--preview", action="store_true",
                   help="show output preview")
    g.add_argument("-s", "--size", metavar="WxH", type=coord,
                   help="set output image size in pixels [default: autodetect]")
    g.add_argument("-q", "--quality", metavar="N", type=jpeg_quality, default=95,
                   help="JPEG output quality value [1-99, default: %(default)s]")
    g.add_argument("-l", "--low-quality", "--draft", action="store_true",
                   help="use faster, but (much) lower-quality scaling")
    g.add_argument("-n", "--dry-run", action="store_true",
                   help="don't generate any images, just show the resulting geometry")
    g = parser.add_argument_group("layout options")
    g.add_argument("-r", "--rows", "--row-major", "--horizontal", action="store_const", dest="vertical", const=False, default=False,
                   help="specify inputs in row-major order (normal English reading order: left to right, then top to bottom) [default]")
    g.add_argument("-c", "--cols", "--column-major", "--vertical", action="store_const", dest="vertical", const=True,
                   help="specify inputs in column-major order (top to bottom, then left to right)")
    g.add_argument("-t", "--autotile", metavar="N", type=int, default=0,
                   help="automatically spread images across rows (in -r mode) or columns (in -c mode) of N images each")
    g.add_argument("-d", "--dynamic", action="store_true",
                   help="allow each row/column/image to have a different size [default: same size for all images, crop if necessary]")
    g = parser.add_argument_group("border options")
    g.add_argument("-i", "--inner-border", metavar="PX", type=int, default=0,
                   help="set border width between images")
    g.add_argument("-j", "--outer-border", metavar="PX", type=int, default=0,
                   help="set border width around the collage")
    g.add_argument("-b", "--border-color", metavar="COLOR", type=color, default=(0,0,0),
                   help="set border color [HTML color name or hex code; default: black]")
    args = parser.parse_args()
    if not(args.output or args.preview or args.dry_run):
        parser.error("at least one of -n, -o or -p must be specified")

    major_name = "column" if args.vertical else "row"
    minor_name = "row" if args.vertical else "column"
    minor_axis = int(args.vertical)
    major_axis = 1 - minor_axis
    
    # import source images
    imgs = [[]]
    for arg in args.infiles:
        if any((c.isalnum() or (c in "*?")) for c in arg):
            # at least one alphanumeric character or wildcard -> assume image file name(s)
            if ("*" in arg) or ("?" in arg):
                files = glob.glob(arg)
            else:
                files = []
            if not files:
                files = [arg]
            for f in files:
                try:
                    if (args.autotile > 0) and (len(imgs[-1]) >= args.autotile):
                        imgs.append([])
                    imgs[-1].append(CollageImage(f))
                except EnvironmentError, e:
                    print >>sys.stderr, "ERROR: can't read input image '%s' - %s" % (f, e)
                    sys.exit(1)
        else:
            # row/column separator found
            imgs.append([])
    l = map(len, imgs)
    lmin = min(l)
    lmax = max(l)
    if not lmin:
        parser.error("empty input %s(s) found" % major_name)
    print "%d %s(s) x %s %s(s)," % (len(imgs), major_name, ("%d-%d" % (lmin, lmax)) if (lmin != lmax) else lmin, minor_name),

    # compute layout
    if args.size:
        outsize = tuple(x - 2 * args.outer_border for x in args.size)
        if args.dynamic:
            Layout_FixedSize_Dynamic(imgs, major_axis, minor_axis, outsize, args.inner_border)
        else:
            Layout_FixedSize_Static(imgs, major_axis, minor_axis, outsize, args.inner_border)
    elif args.dynamic:
        outsize = Layout_VariableSize_Dynamic(imgs, major_axis, minor_axis, args.inner_border)
    else:
        outsize = Layout_VariableSize_Static(imgs, major_axis, minor_axis, args.inner_border)
    outsize = tuple(x + 2 * args.outer_border for x in outsize)
    print "output size: %dx%d" % outsize

    # create image
    if not args.dry_run:
        result = Image.new('RGB', outsize, args.border_color)
    pos = 2 * [args.outer_border]
    for group in imgs:
        pos[minor_axis] = args.outer_border
        major_size = 0
        for img in group:
            print "%s: %dx%d -> %dx%d+%d+%d" % (img.filename, img.img.size[0], img.img.size[1], img.size[0], img.size[1], pos[0], pos[1])
            if not args.dry_run:
                try:
                    result.paste(img.process(args.low_quality), tuple(pos))
                except EnvironmentError, e:
                    print >>sys.stderr, "ERROR: can't process input image '%s' - %s" % (img.filename, e)
                    sys.exit(1)
            major_size = max(major_size, img.size[major_axis])
            pos[minor_axis] += img.size[minor_axis] + args.inner_border
        pos[major_axis] += major_size + args.inner_border

    # export image
    if not args.dry_run:
        if args.preview:
            result.show()
        if args.output:
            if args.output.lower().endswith((".jpg", ".jpeg")):
                result.save(args.output, quality=args.quality)
            else:
                result.save(args.output)