import sys, operator as op, math
import PIL.Image as Image
import PIL.ImageChops as IChops
import PIL.ImageFilter as IFilter
import PIL.ImageDraw as IDraw
import PIL.ImageStat as IStat
import Numeric as Num

def floor(x):
    return int(math.floor(x))

def ceil(x):
    return int(math.ceil(x))

def arr2img(ar):
    """ Convert Numeric.array to PIL.Image.
    """
    return Image.fromstring('L', (len(ar), len(ar[0])), ar.astype('b').tostring())

def img2arr(im):
    """ Convert PIL.Image to Numeric.array.
    """
    return Num.reshape(Num.fromstring(im.tostring(), 'b'), im.size)

class Icon:
    """ An icon with some geometry and channel information pre-computed.
    """
    WHOLE = 'Entire icon'
    TOP = 'Top half'
    LEFT = 'Left half'
    RIGHT = 'Right half'
    BOTTOM = 'Bottom half'
    TOP_LEFT = 'Top-left half'
    TOP_RIGHT = 'Top-right half'
    BOTTOM_LEFT = 'Bottom-left half'
    BOTTOM_RIGHT = 'Bottom-right half'

    def __init__(self, filename):
        self.src = filename
        self.img = Image.open(self.src).convert('RGBA')

        self.W = self.img.size[0]
        self.H = self.img.size[1]

        self.R, self.G, self.B, self.A = map(img2arr, self.img.split())
    
class Crime:
    """ A type of crime with hints on how best to visually match its associated icon.
    """
    AND = 'intersection of colors'
    OR = 'union of colors'

    def __init__(self, icon, colors, choke, spread, combine=AND):
        self.icon = icon
        self.colors = colors
        self.choke = choke
        self.spread = spread
        self.combine = combine

def crimes():
    """ Each type of crime has an associated icon and visual match characteristics.
    """
    return {'Aggravated Assault': Crime(Icon('icons/aggravated_assault.png'),
                                        [(0xE5, 0x05, 0x08)], 3, 5),
            'Alcohol':            Crime(Icon('icons/alcohol.png'),
                                        [(0x6B, 0x06, 0x08)], 3, 7),
            'Arson':              Crime(Icon('icons/arson.png'),
                                        [(0xEF, 0x08, 0x0C)], 3, 7),
            'Burglary':           Crime(Icon('icons/burglary.png'),
                                        [(0x05, 0x03, 0x05), (0xFD, 0xFC, 0xEF)], 1, 5, Crime.AND),
            'Disturbing Peace':   Crime(Icon('icons/disturbing_the_peace.png'),
                                        [(0x04, 0x03, 0x05)], 3, 5),
            'Murder':             Crime(Icon('icons/murder.png'),
                                        [(0xE2, 0x0C, 0x0F), (0xB6, 0x20, 0x21)], 3, 5, Crime.OR),
            'Narcotics':          Crime(Icon('icons/narcotics.png'),
                                        [(0x10, 0x68, 0xE9), (0x45, 0x68, 0x7C)], 1, 5, Crime.OR),
            'Robbery':            Crime(Icon('icons/robbery.png'),
                                        [(0x05, 0x03, 0x07), (0xFD, 0xFB, 0xEF)], 1, 5, Crime.AND),
            'Prostitution':       Crime(Icon('icons/prostitution.png'),
                                        [(0xF7, 0x05, 0xC4), (0xD0, 0x1E, 0xAA)], 1, 5, Crime.OR),
            'Simple Assault':     Crime(Icon('icons/simple_assault.png'),
                                        [(0x05, 0x04, 0xE0)], 3, 5),
            'Theft':              Crime(Icon('icons/theft.png'),
                                        [(0x45, 0x9D, 0x1D)], 1, 5),
            'Vandalism':          Crime(Icon('icons/vandalism.png'),
                                        [(0x8D, 0x17, 0xBB)], 1, 7),
            'Vehicle Theft':      Crime(Icon('icons/vehicle_theft.png'),
                                        [(0xF2, 0x57, 0x0F), (0xFD, 0xF2, 0xCD)], 1, 7, Crime.AND)}

def candidate(img, x, y, icon, part=Icon.WHOLE):
    """ Determine whether an icon exists at a given x, y location in an image.
    """
    box = ((x - ceil(icon.W/2.)), (y - ceil(icon.H/2.)), (x + floor(icon.W/2.)), (y + floor(icon.H/2.)))
    
    # area of interest from the complete image
    area = img.crop(box)
    areaR, areaG, areaB = map(img2arr, area.split())

    # alpha mask of meaningful icon pixels
    alpha = icon.A.copy()
    
    if part == Icon.TOP:
        # interested only in the top - bottom of alpha channel goes to black
        alpha[:,floor(area.size[1]*.35):] = 0

    elif part == Icon.LEFT:
        # interested only in the left - right of alpha channel goes to black
        alpha[floor(area.size[0]*.35):,:] = 0

    elif part == Icon.RIGHT:
        # interested only in the right - left of alpha channel goes to black
        alpha[:ceil(area.size[0]*.65),:] = 0

    elif part == Icon.BOTTOM:
        # interested only in the bottom - top of alpha channel goes to black
        alpha[:,:ceil(area.size[1]*.65)] = 0

    # determine how different the image is from the icon
    diffR = Num.absolute(icon.R.astype('l') - areaR.astype('l'))
    diffG = Num.absolute(icon.G.astype('l') - areaG.astype('l'))
    diffB = Num.absolute(icon.B.astype('l') - areaB.astype('l'))
    diffSum = (diffR + diffG + diffB)
    diffMask = alpha * 255 * 3
    diffFinal = Num.minimum(diffSum, diffMask)
    
    distance = sum(sum(diffFinal))
    threshold = 2.0 * 32 * sum(sum(alpha / 255))
    match = distance / threshold
    
    if match < 1.0:
        # close enough
        print '  Found:', part, box, distance, match
        return {'box': box, 'area': area, 'distance': distance, 'icon': icon, 'match': match, 'part': part}
        
    elif match < 2.0 and part == Icon.WHOLE:
        # close enough to check for a partially-occluded icon

        top = candidate(img, x, y, icon, Icon.TOP)
        if top:
            return top
        
        left = candidate(img, x, y, icon, Icon.LEFT)
        if left:
            return left
        
        right = candidate(img, x, y, icon, Icon.RIGHT)
        if right:
            return right
        
        bottom = candidate(img, x, y, icon, Icon.BOTTOM)
        if bottom:
            return bottom
        
    return None
        
def candidates(img, mask, crime, box=None):
    """ Search for possible crime icon instances in an image, with a mask of possible locations.
    """
    matches = []

    if box is None:
        # start at the whole box
        box = (0, 0, img.size[0], img.size[1])

    width  = box[2] - box[0]
    height = box[3] - box[1]

    hits = mask.crop(box).histogram()[-1]
    
    if hits > 0:
        # there exists a candidate somewhere in the current box

        if (width == 1) or (height == 1):
            # the box is very small, or at least very thin - check each pixel for a match
            for x in range(box[0], box[2]):
                for y in range(box[1], box[3]):
                    if mask.getpixel((x, y)):
                        match = candidate(img, x, y, crime.icon)
                        if match is not None:
                            matches.append(match)

        else:
            # break the box down into quadrants, and recursively search each one
            boxTL = (box[0],         box[1],          box[0]+width/2, box[1]+height/2)
            boxTR = (box[0]+width/2, box[1],          box[2],         box[1]+height/2)
            boxBL = (box[0],         box[1]+height/2, box[0]+width/2, box[3])
            boxBR = (box[0]+width/2, box[1]+height/2, box[2],         box[3])
            
            matches += candidates(img, mask, crime, boxTL)
            matches += candidates(img, mask, crime, boxTR)
            matches += candidates(img, mask, crime, boxBL)
            matches += candidates(img, mask, crime, boxBR)

    return matches

if __name__ == '__main__':
    img = Image.open(sys.argv[-1])
    
    # separate channels of original image, as arrays
    imgR, imgG, imgB = [img2arr(chan) for chan in img.split()]
    
    # find icons for each type of crime, separately
    for name, crime in crimes().items():
        print 'Isolating', name, 'pixels...'
        
        # start with an empty list for mask images
        masks = []
        
        for R, G, B in crime.colors:
            # calculate difference between image and characteristic colors for crime icon
            diffR = Num.absolute(imgR - (Num.ones(img.size) * R))
            diffG = Num.absolute(imgG - (Num.ones(img.size) * G))
            diffB = Num.absolute(imgB - (Num.ones(img.size) * B))
            diffD = diffR + diffG + diffB

            mask = Num.less(diffD, Num.ones(img.size) * 48) * 255
            masks.append(arr2img(mask))
            
        if crime.combine == Crime.OR:
            # combine mask channels into one image
            mask = reduce(IChops.lighter, masks)
        
            # choke out stray pixels, inflate others
            mask = mask.filter(IFilter.MinFilter(crime.choke))
            mask = mask.filter(IFilter.MaxFilter(crime.spread))
            
        else:
            # choke out stray pixels, inflate others
            masks = [mask.filter(IFilter.MinFilter(crime.choke)) for mask in masks]
            masks = [mask.filter(IFilter.MaxFilter(crime.spread)) for mask in masks]
            
            # combine mask channels into one image
            mask = reduce(IChops.darker, masks)
        
        print '  Considering', mask.histogram()[-1], 'of', (mask.size[0] * mask.size[1]), 'total pixels...'

        #mask.show()
        
        # look for possible matches
        for match in candidates(img, mask, crime):
            box = (match['box'][0] - 1, match['box'][1] - 1, match['box'][2], match['box'][3])
            
            if match['part'] == Icon.WHOLE:
                IDraw.Draw(img).rectangle(box, outline=(0x33, 0x33, 0x66), fill=None)

            else:
                IDraw.Draw(img).rectangle(box, outline=(0xCC, 0x99, 0x99), fill=None)
            
    img.show()