#!/opt/local/bin/python2.7
import bz2
from collections import defaultdict
import cPickle
import imghdr
import os
import sys

import pHash

from image import Image

HAMMING_THRESHOLD = 10
CORRELATION_THRESHOLD = 70.0

# Find duplicates by looking over index
# write some description

def find_similar(total_done, index, image):
    index_length = len(index)
    image_digest = image.make_digest()
    
    # Search the index
    ctr = 0
    for index_image in index:
        ctr += 1
        percent_done = int((float(ctr) / index_length) * 100)
        sys.stdout.write('\r\tTotal: {:3d}% --- This image: {:3d}%'.format(total_done, percent_done))
        sys.stdout.flush()
        # Original might be in index, skip it
        if image.hash == index_image.hash and image.path == index_image.path:
            continue
        
        # Hamming distance hash compare
        hamming_distance = pHash.hamming_distance(image.hash, index_image.hash)
        
        # Cross correlation compare
        correlation = pHash.crosscorr(image_digest, index_image.make_digest())[1]

        # Determine if comparison criteria are within the thresholds
        if hamming_distance <= HAMMING_THRESHOLD or \
           correlation >= CORRELATION_THRESHOLD:
            sys.stdout.write('\r\t--- SIMILAR IMAGE FOUND ---              \n')
            print '\tHamming distance: {:d}'.format(hamming_distance)
            print '\tCross correlation: {:f}'.format(correlation)
            print '\tImage path: {:s}\n'.format(index_image.path)
    print


def main():

    # Make sure at least one command line argument was provided
    if len(sys.argv) < 3:
        print 'Usage: find_similar.py <image dir root path> <index path>'
        return -1;

    # Print out the path we're walking
    root_path = sys.argv[1]
    index_path = sys.argv[2]

    # Load image index
    index_file = bz2.BZ2File(index_path, 'r')
    image_index = cPickle.load(index_file)
    index_file.close()

    print 'Looking for similar images from from path: {:s}'.format(root_path)

    # Walk the root path and count all images
    # (is there a better way to do this?)
    image_counter = 0
    for root, dirs, files in os.walk(root_path):
        for file in files:
            image_path = os.path.join(root, file)
            # Check the make sure the file is a JPEG
            if not imghdr.what(image_path) == 'jpeg':
                continue
            image_counter += 1
    print 'Done counting, {:d} images total.'.format(image_counter)

    # Walk the root path
    ctr = 0
    for root, dirs, files in os.walk(root_path):
        for file in files:
            image_path = os.path.join(root, file)
            # Check the make sure the file is a JPEG
            if not imghdr.what(image_path) == 'jpeg':
                continue
            ctr += 1
            percent_done = int((float(ctr) / image_counter) * 100)
            print '\n\nLooking for similar -- {:s}'.format(image_path)
            image_obj = Image(image_path)
            # Find similar images by searching the index
            find_similar(percent_done, image_index, image_obj)

    print '\nDone.'

if __name__ == '__main__':
    sys.exit(main())