#!/usr/bin/env python
"""mungeMPO is a simple utility to extract and process 3D image pairs from 
   .MPO format image files

    This simple python project can easily be imported as a module as part
    of a larger image management
    system

"""

import StringIO
import argparse

import numpy
from PIL import Image


# anaglyphic methods as described here:
# http://www.3dtv.at/Knowhow/AnaglyphComparison_en.aspx
# with thanks to Peter Wimmer

_MAGIC = [0.299, 0.587, 0.114]              # special magic matrix
_ZERO = [0, 0, 0]                           # zero matrix
_IDENT = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]  # identity matrix
REAL = ([_MAGIC, _ZERO, _ZERO], [_ZERO, _ZERO, _MAGIC])
GRAY = ([_MAGIC, _ZERO, _ZERO], [_ZERO, _MAGIC, _MAGIC])
COLOUR = ([_IDENT[0], _ZERO, _ZERO], [_ZERO, _IDENT[1], _IDENT[2]])
HALFCOLOUR = ([_MAGIC, _ZERO, _ZERO], [_ZERO, _IDENT[1], _IDENT[2]])
OPTIMISED = ([[0, 0.7, 0.3], _ZERO, _ZERO], [_ZERO, _IDENT[1], _IDENT[2]])
MODE_DICT = ({'real' : REAL, 'gray':GRAY, 'colour':COLOUR, 
        'halfcolour':HALFCOLOUR, 'optimised':OPTIMISED}) 
METHODS = [REAL, GRAY, COLOUR, HALFCOLOUR, OPTIMISED]

def image_to_array(image):
    """convert PIL Image object to numpy array"""
    ist = image.tostring()
    dim = len(image.getbands())
    return numpy.fromstring(ist, numpy.uint8).reshape(len(ist)/dim, dim)

def array_to_image(image_mode, image_size, src_array):
    """Convert numpy array to PIL Image object"""
    return Image.fromstring(image_mode, image_size, 
           src_array.reshape(len(src_array) * 
           len(image_mode), 1).astype(numpy.uint8).tostring())
   
        
def get_mpo_data(filename):
    """loads data from file, returns a tuple containing JPG data as strings"""
    
    infile = open(filename, 'rb')
    imagedata = infile.read() 
    #read in whole file
    offset = imagedata.find('\xFF\xD8\xFF\xE1', 4)
    #left image is first
    left = imagedata[:offset - 4]
    right = imagedata[offset:]
    return (left, right)

def get_image_pair(filename):
    """loads data from MPO file, returns a pair of PIL Image objects"""
    left, right = get_mpo_data(filename)
    image_buffer = StringIO.StringIO(left)
    left_image = Image.open(image_buffer)
    image_buffer = StringIO.StringIO(right)
    right_image = Image.open(image_buffer)
    return (left_image, right_image)

def get_anaglyph(filename, method=OPTIMISED):
    """generates an anaglyph from MPO file and returns it as a PIL image
       modes are: true_anaglyph
              gray_anaglyph
              colour_anaglyph
              half_colour_anaglyph
              optimised_anaglyph
       """     

    matrix1, matrix2 = [numpy.array(m).transpose() for m in method]
    image1, image2 = get_image_pair(filename)
    im1, im2 = image_to_array(image1), image_to_array(image2)
    composite = numpy.dot(im1, matrix1) + numpy.dot(im2, matrix2)
    result = array_to_image(image1.mode, image1.size, composite)
    return result
        
def save_anaglyph(filename, outfile, method=OPTIMISED):
    """ Save composite anaglyph of requested mode to given filename"""
    image = get_anaglyph(filename, method)
    image.save(outfile)
    
def save_pair(filename, outfile1, outfile2):
    """Save stero pair of images to given filenames"""
    image1, image2 = get_image_pair(filename)
    image1.save(outfile1)
    image2.save(outfile2)

def _frename(filename, suffix):
    """Rename file """
    #TO_DO: Make this safer/better
    return filename[:-4] + suffix
    
def main():
    """run in command line mode - parse options and process files"""
    # set up argparse-based commandline options parser
    m_text = """ The anaglyphic mode to use. where mode is one of:
        [optimised]: Default mode- best option for images with not much red in.
        [real]: dark, colourised anaglyph
        [gray]: black and white (but in red and cyan!)
        [colour]: straightforward colour - bright colours clash though
        [halfcolour]: toned down version.
        [split]: seperates the image into left and right files"""
        
    d_text = """mungempo is a simple commandline tool (and Python helper 
        library) to access stereo images stored inside .MPO file formats often 
        used by 3D cameras such as the Fuji W3, and for the stills mode of various
        3d camcorders.Using the commandline options you can generate separate 
        images or a combined anaglyphic image for viewing with coloured 
        3D glasses."""
        
    s_text = """Suffix to add to generated outfile(s). Default is _3D.jpg.
        When using split mode, suffix is autogenerated as _L and _R. """
        
    parser = argparse.ArgumentParser(description=d_text)
    parser.add_argument('-m', '--mode', nargs='?', default='optimised', 
                         choices=('optimised','real','gray','colour', 
                         'halfcolour','split'), help=m_text)
    parser.add_argument('filename', nargs='+', 
                         help=".MPO file(s) to process")
    parser.add_argument('-s', '--suffix', default='_3D.jpg', 
                         help=s_text)
    args = parser.parse_args()
    
    #execute actions depending on provided switchs/filenames
    if (args.mode != 'split'): 
        modematrix = MODE_DICT[args.mode]
        for name in args.filename:
            print "Processing " + name
            save_anaglyph(name, _frename(name, args.suffix), modematrix)
            
    else:
        for name in args.filename:
            out1 = _frename(name,'_L.jpg')
            out2 = _frename(name,'_R.jpg')
            save_pair(name, out1, out2)
            print "Processed " + name
    print "Operation complete!"
    
if __name__ == "__main__":
    main()
    
 
  
    
