#!/usr/bin/env python
"""
Matplotlib based photo ranking system using the Elo rating system.

Reference: http://en.wikipedia.org/wiki/Elo_rating_system

by Nick Hilton

This file is in the public domain.

"""

# Python
import argparse
import glob
import json
import os
import sys


# 3rd party
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import exifread


class Photo:

    LEFT = 0
    RIGHT = 1

    def __init__(self, filename, score = 1400.0, wins = 0, matches = 0):

        if not os.path.isfile(filename):
            raise ValueError("Could not find the file: %s" % filename)

        self._filename = filename
        self._score = score
        self._wins = wins
        self._matches = matches

        # read orientation with exifread

        with open(filename, 'rb') as fd:
            tags = exifread.process_file(fd)

        self._rotation = str(tags['Image Orientation'])


    def filename(self):
        return self._filename


    def matches(self):
        return self._matches


    def rotation(self):
        return self._rotation


    def score(self, s = None, is_winner = None):

        if s is None:
            return self._score

        assert is_winner is not None

        self._score = s

        self._matches += 1

        if is_winner:
            self._wins += 1


    def win_percentage(self):
        return 100.0 * float(self._wins) / float(self._matches)


    def __eq__(self, rhs):
        return self._filename == rhs._filename


    def to_dict(self):

        return {
            'filename' : self._filename,
            'score' : self._score,
            'matches' : self._matches,
            'wins' : self._wins,
        }


class Display(object):
    """
    Given two photos, displays them with Matplotlib and provides a graphical
    means of choosing the better photo.
    """


    def __init__(self, f1, f2, title = None, figsize = None):

        self._choice = None

        assert isinstance(f1, Photo)
        assert isinstance(f2, Photo)

        if figsize is None:
            figsize = [20,12]

        fig = plt.figure(figsize=figsize)

        h = 10

        ax11 = plt.subplot2grid((h,2), (0,0), rowspan = h - 1)
        ax12 = plt.subplot2grid((h,2), (0,1), rowspan = h - 1)

        ax21 = plt.subplot2grid((h,6), (h - 1, 1))
        ax22 = plt.subplot2grid((h,6), (h - 1, 4))

        kwargs = dict(s = 'Select', ha = 'center', va = 'center', fontsize=20)

        ax21.text(0.5, 0.5, **kwargs)
        ax22.text(0.5, 0.5, **kwargs)

        self._fig = fig
        self._ax_select_left = ax21
        self._ax_select_right = ax22

        fig.subplots_adjust(
            left = 0.02,
            bottom = 0.02,
            right = 0.98,
            top = 0.98,
            wspace = 0.05,
            hspace = 0,
        )

        left = self._read(f1)
        right = self._read(f2)

        ax11.imshow(left)
        ax12.imshow(right)

        for ax in [ax11, ax12, ax21, ax22]:
            ax.set_xticklabels([])
            ax.set_yticklabels([])

            ax.set_xticks([])
            ax.set_yticks([])

        self._attach_callbacks()

        if title:
            fig.suptitle(title, fontsize=20)

        plt.show()


    def _read(self, photo):

        data = mpimg.imread(photo.filename())

        r = photo.rotation()

#~        print "Rotation: ", r
#~        print "    data.shape = ", data.shape

        # rotate as necessary

        if r == 'Horizontal (normal)':
            pass

        elif r == 'Rotated 90 CW':

            data = np.rot90(data, 3)

        elif r == 'Rotated 90 CCW':

            data = np.rot90(data, 1)

        else:
            raise RuntimeError('Unhandled rotation "%s"' % r)

#~        print "    data.shape = ", data.shape

        return data


    def _on_click(self, event):

        if event.inaxes == self._ax_select_left:
            self._choice = Photo.LEFT
            plt.close(self._fig)

        elif event.inaxes == self._ax_select_right:
            self._choice = Photo.RIGHT
            plt.close(self._fig)


    def _attach_callbacks(self):
        self._fig.canvas.mpl_connect('button_press_event', self._on_click)


class EloTable:


    def __init__(self, max_increase = 32.0):
        self._K = max_increase
        self._photos = {}
        self._shuffled_keys = []


    def add_photo(self, photo):

        filename = photo.filename()

        if filename not in self._photos:
            self._photos[filename] = photo


    def get_ranked_list(self):

        # Convert the dictionary into a list and then sort by score.

        ranked_list = self._photos.values()

        ranked_list = sorted(
            ranked_list,
            key = lambda record : record.score(),
            reverse = True)

        return ranked_list


    def rank_photos(self, n_iterations, figsize):
        """
        Displays two photos using the command "gnome-open".  Then asks which
        photo is better.
        """

        n_photos = len(self._photos)

        keys = self._photos.keys()

        for i in xrange(n_iterations):

            np.random.shuffle(keys)

            n_matchups = n_photos / 2

            for j in xrange(0, n_photos - 1, 2):

                match_up = j / 2

                title = 'Round %d / %d, Match Up %d / %d' % (
                    i + 1, n_iterations,
                    match_up + 1,
                    n_matchups)

                photo_a = self._photos[keys[j]]
                photo_b = self._photos[keys[j+1]]

                d = Display(photo_a, photo_b, title, figsize)

                if d == Photo.LEFT:
                    self.__score_result(photo_a, photo_b)
                else:
                    self.__score_result(photo_b, photo_a)


    def __score_result(self, winning_photo, loosing_photo):

        # Current ratings
        R_a = winning_photo.score()
        R_b = loosing_photo.score()

        # Expectation

        E_a = 1.0 / (1.0 + 10.0 ** ((R_a - R_b) / 400.0))

        E_b = 1.0 / (1.0 + 10.0 ** ((R_b - R_a) / 400.0))

        # New ratings
        R_a = R_a + self._K * (1.0 - E_a)
        R_b = R_b + self._K * (0.0 - E_b)

        winning_photo.score(R_a, True)
        loosing_photo.score(R_b, False)


    def to_dict(self):

        rl = self.get_ranked_list()

        rl = [x.to_dict() for x in rl]

        return {'photos' : rl}



def main():

    description = """\
Uses the Elo ranking algorithm to sort your images by rank.  The program reads
the comand line for images to present to you in random order, then you select
the better photo.  After N iteration the resulting rankings are displayed.
"""
    parser = argparse.ArgumentParser(description = description)

    parser.add_argument(
        "-r",
        "--n-rounds",
        type = int,
        default = 3,
        help = "Specifies the number of rounds to pass through the photo set"
    )

    parser.add_argument(
        "-f",
        "--figsize",
        nargs = 2,
        type = int,
        default = [20, 12],
        help = "Specifies width and height of the Matplotlib figsize (20, 12)"
    )

    parser.add_argument(
        "photo_dir",
        help = "The photo directory to scan for .jpg images"
    )

    args = parser.parse_args()

    assert os.path.isdir(args.photo_dir)

    os.chdir(args.photo_dir)

    ranking_table_json = 'ranking_table.json'
    ranked_txt         = 'ranked.txt'

    # Create the ranking table and add photos to it.

    table = EloTable()

    #--------------------------------------------------------------------------
    # Read in table .json if present

    if os.path.isfile(ranking_table_json):
        with open(ranking_table_json, 'r') as fd:
            d = json.load(fd)

        # create photos and add to table

        for p in d['photos']:

            photo = Photo(**p)

            table.add_photo(photo)

    #--------------------------------------------------------------------------
    # glob for files, to include newly added files

    filelist = glob.glob('*.jpg')

    photos = [Photo(x) for x in filelist]

    for p in photos:
        table.add_photo(p)

    #--------------------------------------------------------------------------
    # Rank the photos!

    table.rank_photos(args.n_rounds, args.figsize)

    #--------------------------------------------------------------------------
    # save the table

    with open(ranking_table_json, 'w') as fd:

        d = table.to_dict()

        jstr = json.dumps(d, indent = 4, separators=(',', ' : '))

        fd.write(jstr)

    #--------------------------------------------------------------------------
    # dump ranked list to disk

    with open(ranked_txt, 'w') as fd:

        ranked_list = table.get_ranked_list()

        heading_fmt = "%4d    %4.0f    %7d    %7.2f    %s\n"

        heading = "Rank    Score    Matches    Win %    Filename\n"

        fd.write(heading)

        for i, photo in enumerate(ranked_list):

            line = heading_fmt %(
                i + 1,
                photo.score(),
                photo.matches(),
                photo.win_percentage(),
                photo.filename())

            fd.write(line)

    #--------------------------------------------------------------------------
    # dump ranked list to screen

    print "Final Ranking:"

    with open(ranked_txt, 'r') as fd:
        text = fd.read()

    print text



if __name__ == "__main__": main()