mkadvice-0.2

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from argparse import ArgumentParser
from collections import Counter
from cStringIO import StringIO
from Queue import Empty, Queue
from threading import Thread

from PIL import Image, ImageDraw, ImageFont
from progressbar import Bar, ProgressBar, Percentage, SimpleProgress
from requests import get as req_get

DEFAULT_OUT = 'output.jpg'
FONT = '/usr/share/fonts/ubuntu-font-family/Ubuntu-C.ttf'
WIDTH = 214  # Poster width
HIGHT = 317  # Poster hight

class Movie:
    def __init__(self, imdb_id, title, genre, poster_url=None, poster=None):
        self.imdb_id = imdb_id
        self.title = title
        self.genre = genre
        self.poster_url = poster_url
        self.poster = poster

    def dl_poster_url(self):
        payload = {'id': self.imdb_id, 'plot': 'none'}
        json = req_get('http://imdbapi.org/', params=payload).json()
        self.poster_url = json.get('poster')

    def dl_poster(self):
        if self.poster_url:
            self.poster = req_get(self.poster_url).content


def dl_links(q1, q2):
    """Downloads links, move object (movie) from queue1 to queue2."""
    while not q1.empty():
            movie = q1.get()
            q1.task_done()
            movie.dl_poster_url()
            q2.put(movie)
            q2.task_done()


def dl_posters(q2, q3, Empty, pbar):
    """Downloads posters, move object from queue2 to queue3."""
    while not q3.full():
        try:
            movie = q2.get()
        except Empty:
            pass
        else:
            movie.dl_poster()
            q3.put(movie)
            pbar.update(q3.qsize())
            q3.task_done()


def make_caption(title):
    """Returns image with caption."""
    fake_poster = Image.new('RGB', (WIDTH, HIGHT), 'lightgrey')
    caption = ImageDraw.Draw(fake_poster)
    caption_font = ImageFont.truetype(FONT, 24)
    caption.text((5, 10), title, fill='black', font=caption_font)
    return fake_poster


def main(input_file, output_file):
    # Read movies from file
    movies = [Movie(*(i.strip().capitalize() for i in line.split(',', 2)))
              for line in open(input_file, 'r')]

    # Get unique genres and maximum number of posters in column
    genres = Counter(mov.genre for mov in movies)
    unique_genres = genres.keys()
    unique_genres.sort()
    max_posters = max(genres.values())

    widgets = [Percentage(), Bar(left='[', right=']'), SimpleProgress()]
    pbar = ProgressBar(maxval=len(movies), widgets=widgets).start()

    q1, q2, q3 = (Queue(len(movies)) for q in range(3))
    for mov in movies:
        q1.put(mov)

    # Create two thread pools
    threads_num = int(len(movies) / 2)  # Should find optimal value
    for thread in range(threads_num):
        link_task = Thread(target=dl_links, args=(q1, q2))
        link_task.setDaemon(True)
        link_task.start()
        poster_task = Thread(target=dl_posters, args=(q2, q3, Empty, pbar))
        poster_task.setDaemon(True)
        poster_task.start()

    while not q3.full(): pass  # Wait until all downloads are completed

    pbar.finish()

    # Create canvas, add genres and posters to canvas
    canv_width = 20 + (WIDTH+20) * len(unique_genres)
    canv_hight = 100 + HIGHT*max_posters + 20
    canvas = Image.new('RGB', (canv_width, canv_hight), 'black')
    draw = ImageDraw.Draw(canvas)
    genre_font = ImageFont.truetype(FONT, 42)
    shift = 0
    for genre in unique_genres:
        draw.text((20 + shift, 20), genre, fill='white', font=genre_font)
        shift += (WIDTH+20)

    p_counter = Counter()  # Current number of posters in each column (genre)
    y_counter = Counter()  # Current y value in each column
    x, y = (0, 0)
    for mov in movies:
        if mov.poster:
            poster = Image.open(StringIO(mov.poster))
            if poster.size != (WIDTH, HIGHT):
                poster = poster.resize((WIDTH, HIGHT))
        else:
            poster = make_caption(mov.title)

        # Set the (x, y) in appliance with genre
        for genre in unique_genres:
            if genre == mov.genre:
                x = 20 + (WIDTH+20) * unique_genres.index(genre)
                y = 100 + y_counter[genre]
                p_counter[genre] += 1
                y_counter[genre] = HIGHT * p_counter[genre]

        canvas.paste(poster, (x, y))
    canvas.save(output_file)


if __name__ == '__main__':
    parser = ArgumentParser(description='Downloads posters and makes "advice"')
    parser.add_argument('input_file',
                        metavar='FILE',
                        help='Text file containing movies (id, title, genre)')
    parser.add_argument('-o',
                        dest='output_file',
                        metavar='OUTPUT',
                        default=DEFAULT_OUT,
                        help='Name of out image (default: '+ DEFAULT_OUT +')')
    args = parser.parse_args()
    main(args.input_file, args.output_file)