#!/usr/bin/env python
from math import sqrt
from numpy import array, zeros
from matplotlib.pyplot import subplot, imread, imshow, show, plot
from scipy.ndimage import correlate
from scipy.cluster.vq import kmeans

from intersect import col2bin, domainIterator, colHist, histogramIntersect

def find_peaks(b, threshold, d):
    """Find the location of the peak values of a back projection histogram."""
    # Find all pixels with a value higher than a threshold of the maximum
    # value. Collect K-means estimators on-the-fly.
    threshold *= b.max()
    means = []
    use = []
    d_sq = d ** 2

    for p in domainIterator(b):
        if b[p] >= threshold:
            use.append(p)
            found = False

            for mean in means:
                if (mean[0] - p[0]) ** 2 + (mean[1] - p[1]) ** 2 < d_sq:
                    found = True
                    break

            if not found:
                means.append(p)

    # Use K-means to identify possible matches
    m = kmeans(array(use), array(means))[0]

    return [m[i].tolist() for i in xrange(m.shape[0])]

def convolution(image, radius):
    """Calculate the convolution of an image using a specified radius."""
    # Loop to the square that surrounds the circle, and check if the pixel
    # is inside the circle
    r_sq = radius ** 2
    mask = zeros((2 * radius + 1, 2 * radius + 1), dtype=int)

    for x, y in domainIterator(mask):
        if (x - radius) ** 2 + (y - radius) ** 2 < r_sq:
            mask[x, y] = 1

    return correlate(image, mask, mode='nearest')

def hbp(image, environment, bins, model, radius, mask=None):
    """Create the histogram back projection of two images."""
    # Create image histograms
    print 'Creating histograms...'
    M = colHist(image, bins, model, mask)
    I = colHist(environment, bins, model)

    # Create ratio histogram
    R = zeros(bins)

    for c in domainIterator(R, 3):
        if I[c] != 0:
            R[c] = float(M[c]) / I[c]

    # Apply color model
    use = environment.astype(float)

    if model == 'hsv':
        for p in domainIterator(image):
            use[p] = rgb_to_hsv(*use[p].tolist())

    use *= map(lambda x: x - 1, bins)

    # Create back projection
    print 'Creating back projection...'
    b = zeros(environment.shape[:2])

    for p in domainIterator(b):
        b[p] = min(R[col2bin(use[p])], 1.)

    # Create convolution to concentrate clusters into peak values
    print 'Creating convolution...'
    return convolution(b, radius)

def exclude_color(color, image):
    """Create an image mask that excludes a specific color."""
    mask = zeros(image.shape[:2], dtype=bool)
    color = array(color)

    for p in domainIterator(image):
        if (image[p] != color).any():
            mask[p] = True

    return mask

if __name__ == '__main__':
    print 'Reading images...'
    waldo = imread('waldo.tiff') / 255.
    env = imread('waldo_env.tiff') / 255.

    # Use a mask to ignore the white background in the Waldo image
    mask = exclude_color([255] * 3, waldo)

    # Execute the back projection algorithm
    b = hbp(waldo, env, [32] * 3, 'rgb', 10, mask)

    w, h = waldo.shape[:2]

    def plt(peaks):
        """Draw a rectangle around a list of center pixels."""
        for x, y in peaks:
            l = x - w / 2
            r = x + w / 2
            b = y - h / 2
            t = y + h / 2
            plot([t, t, b, b, t], [l, r, r, l, l], 'r-')

    # Experimental: use K-means to locate clusters, and draw a rectangle
    # with the size of the model image around each cluster
    print 'Locating peaks...'
    peaks = find_peaks(b, .25, sqrt(w ** 2 + h ** 2))
    print 'Done'

    subplot(121)
    plt(peaks)
    imshow(env, origin='lower')
    subplot(122)
    plt(peaks)
    imshow(b, origin='lower')
    show()
