
"""
Not for users. To create a movie use win.getMovieFrame() and then win.saveMovieFrames(filename)

makeAnimatedGIF is the main function here. Others (like mpeg codecs) may follow.

Many thanks to Ray Pascor (pascor at hotpop.com) for the public domain code on 
building an optimised gif palette (makeRGBhistogram, makePalette, rgb2palette are 
very heavily based on his code).
"""
from psychopy import log
import string, time
import Image, ImageChops
from GifImagePlugin import getheader, getdata #part of PIL
try: 
    import pymedia.video.vcodec as vcodec
    havePyMedia=True
except:
    havePyMedia=False

import numpy
    

# --------------------------------------------------------------------
# straightforward delta encoding

def makeAnimatedGIF(filename, images):
    """Convert list of image frames to a GIF animation file
    using simple delta coding"""

    frames = 0
    previous=None
    fp = open(filename, 'wb')
    if images[0].mode in ['RGB','RGBA']: 
        #first make an optimised palette
        optimPalette=makePalette(images, verbose=True)
        
    for n, im in enumerate(images):
        print 'converting frame %i of %i to GIF' %(n+1,len(images))
        if im.mode=='RGB': 
            im = rgb2palette(im, palette=optimPalette, verbose=False)
        if not previous:
            # global header
            for s in getheader(im) + getdata(im):
                fp.write(s)
        else:
            # delta frame
            delta = ImageChops.subtract_modulo(im, previous)
            bbox = delta.getbbox()
            # compress difference
            if bbox:
                for s in getdata(im.crop(bbox), offset = bbox[:2]):
                    fp.write(s)
            else:
                for s in getdata(im):
                    fp.write(s)
                    
        previous = im.copy()
        frames = frames + 1
    fp.write(";")
    fp.close()
    return frames


def RgbHistogram (images, verbose=False):
    """build a histogram of the colors in the image(s)
    with which we can build an optimized color palette"""
    hist = None

    #make a list if given only one image
    if type(images)== type([]):
        nImages = len(images)        
    else:
        nImages=1
        images= [images]
        
    # Form a histogram dictionary whose keys are the repr(color)
    if verbose:    print 'optimising palette ...'
    datalist=[]
    for imgRgb in images:
        datalist.extend(imgRgb.getdata())

    dicthist = {}
    xsize, ysize = imgRgb.size
    numcolors = 0
    for i in xrange (xsize * ysize * nImages):
        color = datalist [i]
        key = repr (color)

        if dicthist.has_key (key):  # color already exists
            dicthist [key] += 1     # increment the count
        else:                       # make a new key
            dicthist [key] = 1      # instantiate a new entry and init the count
            numcolors += 1

            if numcolors > 256:
                if verbose:    print '               ... too many colors'
                return None         # Error flag:  use PIL default color palette/dithering
    if verbose:    print '               ... OK'

    # reform the dictionary into a sorted histogram of the form: (count, (r, g, b))
    hist = []
    for key in dicthist.iterkeys():
        count = dicthist [key]
        color = eval (key)
        hist.append ( (count, color) )
    #end for

    hist.sort()
    hist.reverse()           # make largest counts first

    return hist

#end def RgbHistogram

def Getalphaindex (imgP, maskinv):

    # Find the least used color (palette entry, actually)
    # This will be the color to which transparency will be set when saving the file

    xsize, ysize = imgP.size
    hist = imgP.histogram (maskinv)     # get counts for all colors having non-active alphas

    indexleastused = 255                # arbitrary starting least used palette index
    leastcount = xsize * ysize          # max possible count
    for i in xrange (len (hist)):       # palette size
        if hist [i] < leastcount:
            leastcount = hist [i]       # the count
            indexleastused = i          # the palette index
        #end if

        if hist [i] == 0:    break      # first 0 entry: done
    #end if

    return (indexleastused, leastcount)

#end def Getalphaindex

def makePalette(images, verbose=False):
    palette = []                        # will be [0, 0, 0, ... 255, 255, 255]
    for i in xrange (256):
        palette.append (i); palette.append (i); palette.append (i) 
    #end for
    
    hist = RgbHistogram (images, verbose=verbose)
    #check for alpha in one of the images
    if type(images)==type([]):
        img = images[0]
    else: img = images
    
    if hist == None: # colors > 256:  use PIL dithered image & palette
        palette=None
    else:
        # Make two lists of the colors.
        colors = []
        colorsAndIndices = []
        for i in xrange (len (hist)):
            if img.mode=='RGBA':
                r, g, b, a = hist [i][1]            # pick off the color tuple
            else:
                r, g, b = hist [i][1]               # pick off the color tuple
            #end if    
            palette [i*3 + 0] = r
            palette [i*3 + 1] = g
            palette [i*3 + 2] = b
    return palette
    
def rgb2palette (imgRgb, palette=None, verbose=False):     # image could be a "RGBA"
    """
    Converts an RGB image to a palettised version (for saving as gif).
    """
    #verbose = False

    imgP = None
    datalist = None

    size = imgRgb.size
    xsize = size [0]
    ysize = size [1]

    hasalpha = False
    if imgRgb.mode == 'RGBA':
        hasalpha = True                            # for post=processing to create transparency
        if verbose:    print 'Rgb2p:  Input image is RGBA'

        # Create a mask and its inverse
        source = imgRgb.split()
        R, G, B, A = 0, 1, 2, 3         # band indices
        mask    = source [A].point (lambda i: i <  255 and 255)     # = True on active alphas
        maskinv = source [A].point (lambda i: i == 255        )     # = True on inactive alphas
    #end if

    # find the most popular colors, limiting the max number to 256
    # any "excess" colors with be transformed later to the closest palette match
    if palette==None:
        palette=makePalette(imgRgb)
    
    numPalette= numpy.reshape(numpy.asarray(palette), [256,3])
    listPalette = numPalette.tolist()
    imgP = Image.new ('P', size)            # Create a brand new paletted image
    imgP.putpalette (palette)               # Install the palette

    # Rewrite the entire image using new palette's indices.
    if verbose:    print 'Defining the new image using the newly created palette ...'

    # Each pixel gets a palette color index
    if datalist == None:
        datalist = list (imgRgb.getdata())      # xsize*ysize list of color 3-tuples
    #end if

    pxlctr = 0
    colctr = 0
    for yord in xrange (ysize):
        for xord in xrange (xsize):
            pxlcolor = list(datalist [yord*xsize + xord])     # org image color tuple
            
            if pxlcolor in listPalette:
                index = listPalette.index(list(pxlcolor))
            else: index=0
            #paletteindex = colorsAndIndices [index] [1]     # a simple lookup

            imgP.putpixel ((xord, yord), index)
            
    if hasalpha:
        indexleastused, leastcount = Getalphaindex (imgP, maskinv)
        if verbose:
            print
            print 'Low color-count image:   least used color index, leastcount =', indexleastused, leastcount
        #end if

        return (imgP, indexleastused, mask)
    else:
        return (imgP)
    #end if

#end def Rgb2p

def makeMPEG(filename, images, codec='mpeg1video', codecParams = None, verbose=False):
    if not havePyMedia:
        log.error('pymedia (www.pymedia.org) needed to make mpeg movies')
        return 0
    
    fw= open( filename, 'wb' )
    t= time.time()	
    
    #set bitrate
    if codec== 'mpeg1video':
        bitrate= 2700000
    else:
        bitrate= 9800000
        
    #set other params (or receive params dictionary)
    if codecParams == None:
        codecParams= { \
            'type': 0,
            'gop_size': 12,
            'frame_rate_base': 125,
            'max_b_frames': 0,
            'width': images[0].size[0],
            'height': images[0].size[1],
            'frame_rate': 3125,
            'deinterlace': 0,
            'bitrate': bitrate,
            'id': vcodec.getCodecID( codec )
            }
        log.info('Setting codec to ' + str(codecParams))
    encoder= vcodec.Encoder( codecParams )    
    
    for im in images:	    
        # Create VFrame
        imStr = im.tostring()
        bmpFrame= vcodec.VFrame( vcodec.formats.PIX_FMT_RGB24, im.size, (imStr,None,None))
        yuvFrame= bmpFrame.convert( vcodec.formats.PIX_FMT_YUV420P, im.size )
        d = encoder.encode( yuvFrame )
        try:
            fw.write( d.data )#this is what works!
        except:
            fw.write( d )#this is what pymedia demo recommends
    else:
        log.info('%d frames written in %.2f secs' % ( len(images), time.time()- t))
        i= 0
    fw.close()
    