#!/usr/bin/env python3

# Copyright (C) 2012 Vladimir Niksic
# Copyright (C) 2012 Hrvoje Niksic
#
# This work is licensed under the Creative Commons Attribution 3.0
# Unported License. To view a copy of this license, visit
# http://creativecommons.org/licenses/by/3.0/ or send a letter to
# Creative Commons, 444 Castro Street, Suite 900, Mountain View,
# California, 94041, USA.

import optparse
import os, sys, time
import re, itertools
import random, subprocess, shutil, tempfile, signal
import fractions
import PIL.Image, PIL.ImageOps, PIL.ExifTags
import multiprocessing.pool
from tinytag import TinyTag

FRAMES_PER_SEC = 25
DEFAULT_OUTPUT_FILE = 'out.mp4'
DEFAULT_GEOMETRY = (1920, 1080)

def thread_count():
    return os.sysconf('SC_NPROCESSORS_ONLN')

def get_tagid(name):
    for tagid, tagname in PIL.ExifTags.TAGS.items():
        if tagname == name:
            return tagid

def get_tagval(img, tagname):
    if not hasattr(img, '_getexif'):
        return
    exif = img._getexif()
    if exif is not None:
        return exif.get(get_tagid(tagname))

transposition_map = {
    2: (PIL.Image.FLIP_LEFT_RIGHT,),
    3: (PIL.Image.ROTATE_180,),
    4: (PIL.Image.FLIP_TOP_BOTTOM,),
    5: (PIL.Image.FLIP_TOP_BOTTOM, PIL.Image.ROTATE_270),
    6: (PIL.Image.ROTATE_270,),
    7: (PIL.Image.FLIP_LEFT_RIGHT, PIL.Image.ROTATE_270),
    8: (PIL.Image.ROTATE_90,),
}

def apply_exif_orientation(img):
    orient = get_tagval(img, 'Orientation')
    for method in transposition_map.get(orient, ()):
        img = img.transpose(method)
    return img

def scaled_size(size, geom):
    aspect = fractions.Fraction(*size)
    wanted_aspect = fractions.Fraction(*geom)
    w, h = geom
    if aspect < wanted_aspect:
        w = int(h * aspect)     # black bars left and right
    else:
        h = int(w / aspect)     # black bars up and down
    return w, h

def letterbox(img, geom):
    """Fit the image into a letterbox of size GEOM."""
    fitsize = scaled_size(img.size, geom)
    if fitsize != img.size:
        img = img.resize(fitsize, PIL.Image.ANTIALIAS)
        if img.size != geom:
            full = PIL.Image.new(img.mode, geom)
            w, h = geom
            scaled_w, scaled_h = fitsize
            full.paste(img, ((w - scaled_w) // 2, (h - scaled_h) // 2))
            img = full
    return img

def transform_single(src, target, geom):
    if geom is None:
        os.symlink(os.path.abspath(src), target)
        return
    try:
        img = PIL.Image.open(src)
    except IOError:
        print('warning: cannot open %s' % src)
        return
    img = apply_exif_orientation(img)
    img = letterbox(img, geom)
    img.save(target, 'JPEG', quality=90)

def report(s):
    sys.stdout.write(s)
    sys.stdout.flush()

def task_start(src, target, geom, taskid):
    if geom:
        report(' %d' % taskid)
    transform_single(src, target, geom)

def wait_results(results):
    while results:
        done = frozenset(r for r in results if r.ready())
        results -= done
        for r in done:
            r.get()         # propagate exception if any
        time.sleep(min(.5, .1 * len(results)))

def transform_pics(pics, geom, picdir):
    pool = multiprocessing.pool.Pool(thread_count())
    if geom:
        report('scaling %d images in %d threads...'
               % (len(pics), thread_count()))
    outfiles = []
    results = set()
    for cnt, pic in enumerate(pics, 1):
        newfile = os.path.join(picdir, 'pic-%04d.jpg' % cnt)
        results.add(
            pool.apply_async(task_start, (pic, newfile, geom, cnt)))
        outfiles.append(newfile)
    try:
        wait_results(results)
    except:
        pool.terminate()
        raise
    pool.close()
    pool.join()
    if geom:
        report('\n')
    return outfiles

def create_frames(pics, geom, frames, picdir):
    processed = transform_pics(pics, geom, picdir)
    # non-existent processed pictures occur when a source image cannot
    # be read
    processed[:] = [p for p in processed if os.path.exists(p)]
    report('generating %d frames out of %d images...'
           % (frames, len(processed)))
    framesperpic = fractions.Fraction(frames, len(processed))
    startframe = 0
    for pic in processed:
        for frame in range(int(startframe),
                            int(startframe + framesperpic)):
            os.symlink(pic, os.path.join(picdir, 'frame%d.jpg' % frame))
        startframe += framesperpic
    report(' done\n')

def encode_video(mp3file, outfile, picdir):
    subprocess.call(["ffmpeg",
                     "-r", str(FRAMES_PER_SEC),
                     "-i", os.path.join(picdir, 'frame%d.jpg'),
                     "-i", mp3file,
                     #"-loglevel", "warning",
                     "-vcodec", "libx264",
                     "-acodec", "copy",
                     "-threads", str(thread_count()),
                     outfile])

def mp3_length(mp3file):
    tag = TinyTag.get(mp3file)
    return int(tag.duration)

def create_video(pics, mp3, geom, outfile):
    frames = mp3_length(mp3) * FRAMES_PER_SEC
    if len(pics) > frames:
        raise RuntimeError('The number of pictures must not exceed '
                           'the target number of frames')
    picdir = tempfile.mkdtemp(prefix='img2video')
    try:
        create_frames(pics, geom, frames, picdir)
        encode_video(mp3, outfile, picdir)
    finally:
        print('deleting', picdir)
        try:
            shutil.rmtree(picdir)
        except OSError:
            pass

def determine_geometry(pics):
    sizes = [PIL.Image.open(pic).size for pic in pics]
    widths, heights = zip(*sizes)
    return max(widths), max(heights)

GEOM_RE = re.compile(r'(\d+)x(\d+)$')

def parse_geometry(s):
    match = GEOM_RE.match(s)
    if not match:
        raise ValueError("invalid geometry %r" % s)
    return (int(match.group(1)), int(match.group(2)))

def main():
    parser = optparse.OptionParser(usage="%prog [-s] [o OUTFILE] MP3 IMAGE...")
    parser.add_option('-g', '--geom', dest='geom',
                      metavar='wxh|auto|none',
                      help='desired image geometry, defaults to 1920x1080')
    parser.add_option('-s', '--shuffle', dest='shuffle',
                      default=False, action='store_true',
                      help='shuffle images before encoding')
    parser.add_option('-o', '--output', dest='output',
                      help='output file name, defaults to out.mp4')
    options, args = parser.parse_args()
    if len(args) < 2:
        parser.print_usage(sys.stderr)
        sys.exit(2)
    mp3 = args[0]
    pics = args[1:]
    if options.shuffle:
        random.shuffle(pics)
    if options.geom is None:
        geom = DEFAULT_GEOMETRY
    elif options.geom == 'auto':
        geom = determine_geometry(pics)
    elif options.geom == 'none':
        geom = None
    else:
        geom = parse_geometry(options.geom)
    print('using geometry', geom)
    create_video(pics, mp3, geom, options.output or DEFAULT_OUTPUT_FILE)

if __name__ == '__main__':
    main()
