#!/usr/bin/env python
from numpy import zeros
from matplotlib.pyplot import imread
from colorsys import rgb_to_hsv

def col2bin(color):
    """Get the histogram bin coordinates of a color."""
    return tuple(color.astype(int))

def domainIterator(image, dim=2):
    """Pixel iterator for arrays with 2 or 3 dimensions."""
    if dim == 2:
        for x in xrange(image.shape[0]):
            for y in xrange(image.shape[1]):
                yield x, y
    elif dim == 3:
        for x in xrange(image.shape[0]):
            for y in xrange(image.shape[1]):
                for z in xrange(image.shape[2]):
                    yield x, y, z

def colHist(image, bins, model, mask=None):
    """Create the color histogram of an image."""
    h = zeros(bins, dtype=int)
    use = image.astype(float)

    if model == 'hsv':
        for p in domainIterator(image):
            use[p] = rgb_to_hsv(*use[p].tolist())
    elif model != 'rgb':
        raise ValueError('Color model "%s" is not supported.' % model)

    # Multiply with the number of bins to normalize bin coordinates
    use *= map(lambda x: x - 1, bins)

    for p in domainIterator(image):
        # Optionally use a mask over the image
        if mask is None or mask[p]:
            h[col2bin(use[p])] += 1

    return h

def histogramIntersect(h1, h2):
    """Calculate the intersection between two color histograms."""
    if h1.shape != h2.shape:
        raise ValueError('Histograms h1 and h2 are not aligned.')

    match = 0

    # Add the minimum of each bin to the result
    for c in domainIterator(h1, 3):
        match += min(h1[c], h2[c])

    # Normalize by dividing by the number of pixels
    return float(match) / h2.sum()

if __name__ == '__main__':
    from sys import argv

    count = 20
    model = 'rgb'
    bins = [64] * 3

    # Usage: python intersect.py [ count [ model ] ]
    if len(argv) > 1:
        count = int(argv[1])

        if len(argv) > 2:
            model = argv[2]

    # Create table of intersections of images in the database
    hist = [colHist(imread('database/%d.jpg' % i) / 255., \
            bins, model) for i in xrange(1, count + 1)]

    table = zeros((len(hist), len(hist)), dtype=int)

    for i, im in enumerate(hist):
        table[i, i] = 100

        for j in xrange(i + 1, len(hist)):
            table[i, j] = table[j, i] \
                    = round(histogramIntersect(hist[i], hist[j]) * 100)

    print table
