#!/usr/bin/env python3
from argparse import ArgumentParser
from functools import wraps

import cv2
import numpy as np


def image_required(fun):
    @wraps(fun)
    def wrapper(self, *args, **kwargs):
        if self.image is None:
            raise RuntimeError('Property {}.{} requires an image!'
                               .format(self.__class__.__name__, fun.__name__))
        return fun(self, *args, **kwargs)
    return wrapper


class imgrate:
    def __init__(self, imgfile=None):
        self.image = None
        if imgfile:
            self.load_image(imgfile)

    def load_image(self, imgfile):
        image = cv2.imread(imgfile)
        self.image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    @property
    @image_required
    def quality(self):
        return self.laplacian * self.fdibm

    @property
    @image_required
    def laplacian(self):
        return cv2.Laplacian(self.image, cv2.CV_64F).var()

    @property
    @image_required
    def fdibm(self):
        F = np.fft.fft2(self.image)
        Fcentered = np.fft.fftshift(F)
        Fabs = np.abs(Fcentered)
        Fmax = Fabs.max()
        Th = np.count_nonzero(Fabs > Fmax/1000)
        return Th / (F.shape[0] * F.shape[1])


if __name__ == '__main__':
    from os import remove
    from sys import exit
    from signal import signal, SIGINT
    from multiprocessing import Pool, cpu_count

    signal(SIGINT, lambda a, b: exit('\nExiting!'))


    def parse_arguments():
        ap = ArgumentParser(description='Find the best quality one among similar images.'
                                        'Note: image dimensions should match!')
        ap.add_argument('images', type=str, nargs='+', help='')
        ap.add_argument('-d', '--delete', action='store_true', help='Delete all but the best quality image.')
        ap.add_argument('-q', '--quiet', action='store_true', help='Print quality measurements only.')
        ap.add_argument('-m', '--method', type=str, default='quality', help='Select quality measurement to use.')
        return ap.parse_args()


    def run(args):
        if not hasattr(imgrate, args.method):
            exit('Invalid --method option!')

        ratings = calculate_all_ratings(args.images, args.method)
        maximum = max(ratings, key=ratings.get)

        fancy_print_ratings(ratings, args.quiet, maximum)

        if args.delete:
            [remove(image) for image in ratings if image != maximum]


    def calculate_all_ratings(images, method):
        ratings = []
        with Pool(processes=cpu_count()) as pool:
            for image in images:
                ratings.append((image, pool.apply_async(calculate_rating, (image, method))))

            return {image: promise.get() for image, promise in ratings}


    def calculate_rating(image, method):
        return getattr(imgrate(image), method)


    def fancy_print_ratings(ratings, quiet, maximum):
        for rating in ratings.keys():
            if quiet:
                print(ratings[rating])
            else:
                maxmark = '*' if rating == maximum and len(ratings) > 1 else ' '
                print('{}imgrate("{}") = {}'.format(maxmark, rating, ratings[rating]))


    run(parse_arguments())