#!/usr/bin/env python


import pNbody
from pNbody import Movie
from pNbody import *
from pNbody.palette import *
from pNbody import cosmo

import string
from numpy import *
import Image
import ImageDraw
import ImageFont
import ImagePalette
import ImageTk
import ImageFont
import sys
import os
import string
import getopt
import math

from optparse import OptionParser


FONT_PATH = os.path.join(NBODYPATH,'fonts')
FFONT = os.path.join(FONT_PATH,'helvBO10.pil')


    

####################################################################################
def parse_options():
####################################################################################

  usage = "usage: %prog [options] file"
  parser = OptionParser(usage=usage)

  parser.add_option("-q","--quality",
                    action="store", 
		    dest="quality", 
		    type="int",
		    default=5,
                    help="quality 1,2,3,4,5")

  parser.add_option("-m","--mode",
                    action="store", 
		    dest="mode", 
		    type="string",
		    default=None,
                    help="image mode : L, P or RGB")

  parser.add_option("--rgb",
                    action="store_true", 
		    dest="rgb", 
		    default=False,
                    help="rgb mode")

  parser.add_option("-k","--keep",
                    action="store_true", 
		    dest="keep", 
		    default=False,
                    help="keep png images")
		    
  parser.add_option("--only_film",
                    action="store_true", 
		    dest="only_film", 
		    default=False,
                    help="make movie from dir dir, without reading a gmov film")
		    		    

  parser.add_option("--info",
                    action="store_true", 
		    dest="info", 
		    default=False,
                    help="give info on film")

  parser.add_option("--nofilm",
                    action="store_true", 
		    dest="nofilm", 
		    default=False,
                    help="do not create the film (extract image only)")

  parser.add_option("--with_time",
                    action="store_true", 
		    dest="with_time", 
		    default=False,
                    help="add time label")

  parser.add_option("--time_format",
                    action="store", 
		    dest="time_format", 
		    default=None,
                    help="time format")

  parser.add_option("--time_x",
                    action="store",
		    type='float', 
		    dest="time_x", 
		    default=None,
                    help="time relative position in x")
  
  parser.add_option("--time_y",
                    action="store",
		    type='float', 
		    dest="time_y", 
		    default=None,
                    help="time relative position in y")
		    
  parser.add_option("--text",
                    action="store", 
		    dest="text", 
		    type="string",
		    default=None,
                    help="text")

  parser.add_option("--text_x",
                    action="store",
		    type='float', 
		    dest="text_x", 
		    default=None,
                    help="text relative position in x")
  
  parser.add_option("--text_y",
                    action="store",
		    type='float', 
		    dest="text_y", 
		    default=None,
                    help="text relative position in y")


  parser.add_option("-o",
                    action="store", 
		    dest="outputfile", 
		    type="string",
		    default=None,
                    help="output file name")
		    
  parser.add_option("--fps",
                    action="store", 
		    dest="fps", 
		    type="int",
		    default=24,
                    help="frame per second")

  parser.add_option("--istart",
                    action="store", 
		    dest="istart", 
		    type="int",
		    default=0,
                    help="first number of frame")

  parser.add_option("-s","--step",
                    action="store", 
		    dest="step", 
		    type="int",
		    default=0,
                    help="number of step to leep")
		    
  parser.add_option("-i","--tinit",
                    action="store", 
		    dest="tinit", 
		    type="float",
		    default=0,
                    help="initial time")


  parser.add_option("-f","--tfinal",
                    action="store", 
		    dest="tfinal", 
		    type="float",
		    default=1e100,
                    help="final time")

  parser.add_option("-z","--zoom",
                    action="store", 
		    dest="zoom", 
		    type="float",
		    default=1,
                    help="zoom factor")


  parser.add_option("-p","--palette",
                    action="store", 
		    dest="palette", 
		    type="string",
		    default="ramp",
                    help="color palette")
		    
		    
  parser.add_option("-d","--dir",
                    action="store", 
		    dest="dir", 
		    type="string",
		    default="tmp",
                    help="directory for png files")	
		    
  parser.add_option("--text_color",
                    action="store", 
		    dest="text_color", 
		    type="string",
		    default=None,
                    help="text color")		    
		    	    
  parser.add_option("--codec",
                    action="store", 
		    dest="codec", 
		    type="string",
		    default=None,
                    help="msmpeg4v2, mpeg4, x264")		    

  parser.add_option("--format",
                    action="store", 
		    dest="format", 
		    type="string",
		    default=None,
                    help="film format")	


  parser.add_option("--ratio",
                    action="store", 
		    dest="ratio", 
		    type="float",
		    default=None,
                    help="film ratio")			    

  parser.add_option("--bitrate",
                    action="store", 
		    dest="bitrate", 
		    type="int",
		    default=None,
                    help="bitrate in k")	

  parser.add_option("--pass",
                    action="store", 
		    dest="passes", 
		    type="int",
		    default=1,
                    help="number of pass")

  (options, args) = parser.parse_args()

  if len(args) == 0:
    if not options.only_film:
      print "you must specify at least a filename"
      sys.exit(0)

  return args,options
  
        
####################################################################################  
#
#  MAIN
#
#################################################################################### 


files, options = parse_options()


if options.only_film:
  options.keep = True



if not options.only_film:


  fname = files[0]


  exec("options.text_color = %s"%(options.text_color))
  	  
  		  
  if options.rgb:
    mode="image_rgb"
    
    if options.text_color == None:
      options.text_color = (255,255,255) 
    
  else:
    mode="image"    

    if options.text_color == None:
      options.text_color = 255






  # initial and final time
  	  
  if (options.tfinal != 0):
    if (options.tfinal < options.tinit):
      print "number of final image must be larger than initial image !"
      sys.exit(0)

  #fname = xarguments[0]

  if (fname==" "):
    help_message()
    sys.exit(0)   
      
    
  # verifie que le fichier existe

  if (os.path.exists(fname)==0): 
    print "Error : the file ",fname," no not exist."
    sys.exit(0) 


  # output file
  if options.outputfile == None:
    options.outputfile = "%s.avi"%os.path.splitext(os.path.basename(fname))[0]
  else:
    # find output format
    ext = os.path.splitext(options.outputfile)[1]
    options.format = ext[1:]
    
      
    
  # ouvre le film  
  film = Movie.Movie(fname)
  film.open()  
    
  if options.info:
    film.info() 
    sys.exit(0)

  # create a palette object
  palette = Palette()
  pth =  os.path.join(PALETTEDIR,options.palette)
  palette.read(pth)


  if not os.path.exists(options.dir):
    os.mkdir(options.dir)
    
  ###################################################
  # taille de l'image de sortie  

  # largeur de l'image
  width =  int(film.numByte*options.zoom)

  #----------------				 
  # choix des font				 
  #----------------				 
  						 
  # determination de la taille  		 
  						 
  fsize = int(12*options.zoom - math.fmod(12*options.zoom,2))	 
  						 
  if (fsize > 24) :				 
    fsize = 24  				 
  if (fsize < 8) :				 
    fsize = 8					 
  						 
  if  (fsize < 10):				 
    size = '0'+`fsize`  			 
  else: 					 
    size = `fsize`				 
      
  fontbasename = os.path.basename(FFONT)  
  fontdirname  = os.path.dirname(FFONT)+'/'	
  fontname = fontdirname+fontbasename[:6]+size+ fontbasename[8:]

  font = ImageFont.load(fontname)    
  							
  xy = font.getsize('-')		
  	  
  if options.with_time and (options.time_x==None) and (options.time_y==None): 						
    dylab = int(2*xy[1])
  else :
    dylab = 0
  				
  											
  # hauteur de l'image  				
  height = int((film.numLine)*options.zoom)+dylab 
  shapezoom = (width,height)				

  ###########################
  # lecture de tout le film
  ###########################

  leap = 0
  i = -1 + options.istart
  while 1:
    #data = film.read_one()
    #time,image = read_one_with_time()
  	 
    i = i + 1  

  	    
    time,image = film.read_one_with_time(mode)  
      
    if time==None:
      break
    
    atime = film.current_time
    
      
    if atime>=options.tinit and atime<=options.tfinal:
    
      if leap==0:
      
  	# zoom
  	if options.zoom != 1.:
  	  image = image.transform(shapezoom,Image.AFFINE,(1./options.zoom,0.,0.,0.,1./options.zoom,0),Image.BICUBIC)  
  	if dylab!=0:	  
  	  image = image.transform(shapezoom,Image.AFFINE,(1,0.,0.,0.,1,-dylab),Image.BICUBIC)	     
  												      
  	# add time	  
  	if options.with_time:	
	    
  	  if   options.time_format==None:
  	    time = "%8.3f"%(atime)
  	  elif options.time_format=='cGyr':	   
  	    a = atime
  	    t = cosmo.Age_a(a)
  	    time = "t = %4.1f Gyr"%(-t)
  	  elif options.time_format=='caGyr':	    
  	    a = atime
  	    t = cosmo.Age_a(a)
  	    time = "a = %4.3f  t = %4.1f Gyr"%(a,-t)
  	  elif options.time_format=='cazGyr': 
  	    a = atime
  	    z = 1./a - 1 
  	    t = cosmo.Age_a(a)
  	    time = "a = %4.3f       z=%4.1f       t=%4.1f Gyr"%(a,z,-t)	    
  	  elif options.time_format=='czGyr': 
  	    a = atime
  	    z = 1./a - 1 
  	    t = cosmo.Age_a(a)
  	    time = "z=%4.1f    t=%4.1f Gyr"%(z,-t)
  	  elif options.time_format=='ca':	 
  	    a = atime
  	    time = "a = %4.3f"%(a)
  	  elif options.time_format=='cz':	 
  	    a = atime
  	    z = 1./a - 1 
  	    time = "z=%4.1f"%(z)    
  	  elif options.time_format=='Myr':	 
  	    a = atime
  	    time = "%04.0f Myr"%(a)   
	  else:
	    exec(options.time_format) 
	      			  
  	  xy = font.getsize(time)

	  # specific position
	  if (options.time_x!=None) and (options.time_y!=None):
  	    xref = width - (width- xy[0])/2						      
	    x = width*options.time_x  - xy[0]/2
	    y = height - height*options.time_y	- xy[1]/2
  	    poslab = (x,y)     
  	    draw = ImageDraw.Draw(image)  							      
  	    draw.text(poslab,time,fill=options.text_color,font=font)
	  
	  # default position outisde movie
	  else: 
  	    xref = width - (width- xy[0])/2						      
  	    x = xref - xy[0]	  
  	    y = int(dylab/4)
  	    poslab = (x,y)     
  	    draw = ImageDraw.Draw(image)  							      
  	    draw.rectangle((0,0,width,dylab),fill=1)	  
  	    draw.text(poslab,time,fill=options.text_color,font=font)
    
    
  	# add a text		  
  	if options.text != None:
  	  xy = font.getsize(options.text)	
	  
	  if (options.text_x!=None) and (options.text_y!=None): 
	    x = width*options.text_x  - xy[0]/2
	    y = height - height*options.text_y	- xy[1]/2	
	  else:						      
  	    x = 1*xy[1]
  	    y = height - xy[1] - 1*xy[1]  
	      	  
	  postext = (x,y)    
  	  draw = ImageDraw.Draw(image)  	  
  	  draw.text(postext,options.text,fill=options.text_color,font=font)	      
  								      
  	
	# add borders
	if options.ratio != None:
	  
	  # read the geometry
	  w = image.size[0]
	  h = image.size[1]
	  
	  f = options.ratio * float(h)/float(w)
	  dw =  int(0.5*w* (f - 1))	  
 	  image = image.transform((dw++dw+w,h),Image.AFFINE,(1,0.,-dw,0.,1,0),Image.BICUBIC)
	
												      
  	# include the palette
  	if not options.rgb:
  	  image.putpalette(palette.palette)	  
        
	if options.mode != None: 
          image = image.convert(options.mode)

    ####################################
    # en dessous, partie propre a mpeg 
    ####################################
  	  
  	# save it
  	bname = "%s/%08d"%(options.dir,i)
  	print bname, atime
  	image.save(bname+'.png') 
	
	if options.format == 'mp4':
	  os.system("convert -quality 100 %s %s "%(bname+'.png',bname+'.gif'))
      
      leap = fmod(leap+1,options.step)


    else:
      continue
  

####################################
# conversion into gif if needed
####################################  
'''
if options.format == 'mp4' and  options.only_film:
  
  files = glob.glob(os.path.join(options.dir,'*.png'))
  
  for file in files:
    bname = os.path.splitext(file)[0]
    cmd = "convert -quality 100 %s %s "%(bname+'.png',bname+'.gif')
    print cmd
    os.system(cmd)  
'''  
  
####################################
# creation du film mpeg
####################################

'''
#mencoder mf://tmp/*.png -mf type=png:fps=25  -ovc lavc -lavcopts vcodec=mpeg4 -oac copy -o film.mpg

#tres basse qualite vbitrate=800
mencoder mf://tmp/*.png -mf type=png:fps=25  -ovc lavc -lavcopts vcodec=mpeg4:vbitrate=256 -oac copy -o film.mpg    


#basse qualite vbitrate=800
mencoder mf://tmp/*.png -mf type=png:fps=25  -ovc lavc -lavcopts vcodec=mpeg4 -oac copy -o film.mpg	      
  
#basse qualite
mencoder mf://tmp/*.png -mf type=png:fps=25  -ovc lavc -lavcopts vcodec=mpeg4:vbitrate=1600 -oac copy -o film.mpg 
   
tres bonne qualite
mencoder mf://tmp/*.png -mf type=png:fps=25  -ovc lavc -lavcopts vcodec=mpeg4:vbitrate=2400:mbd=2:mv0:trell:v4mv:cbp:last_pred=3:predia=2:dia=2:vmax_b_frames=2:vb_strategy=1:precmp=2:cmp=2:subcmp=2:preme=2:qns=2 -oac copy -o film.mpg

super qualite
mencoder mf://tmp/*.png -mf type=png:fps=25  -ovc lavc -lavcopts vcodec=mpeg4:vbitrate=8000:vrc_buf_size=1835:vrc_maxrate=9800:keyint=15:trell:mbd=2:precmp=2:subcmp=2:cmp=2:dia=-10:predia=-10:cbp:mv0:vqmin=1:lmin=1:dc=10:vstrict=0 -oac copy -o film.mpg


sans compresseion
mencoder mf://tmp/*.png -mf type=png:fps=25  -ovc copy -oac copy -o film.avi  
'''


#if options.format == 'mp4':
#  
#  #ffmpeg -r 24 -b 16000k -i %08d.gif  qq.mp4
#  
#  #cmd = """ffmpeg  -i %s/%%08d.gif  -r %d  -b 16000k  %s"""%(options.dir,options.fps,options.outputfile)
#  cmd = """ffmpeg  -i %s/%%08d.png  -r %d  -b 16000k  %s"""%(options.dir,options.fps,options.outputfile)
  


# read one file and find width height
files = glob.glob(os.path.join(options.dir,'*.png'))
img = Image.open(files[0])
width = img.size[0]
height= img.size[1] 
print "image size = %d x %d"%(width,height)

  
oac =  "-oac copy"

if   options.codec=='msmpeg4v2':

  ovco = "-ovc lavc -lavcopts"

  if options.bitrate!=None:
    vbitrate = options.bitrate
  elif options.quality == 1:
    vbitrate = 256
  elif options.quality == 2:
    vbitrate = 1600
  elif options.quality == 3:
    vbitrate = 2400
  elif options.quality == 4:
    vbitrate = 8000

  ovc =  "vcodec=%s:vbitrate=%d"%(options.codec,vbitrate)
  


elif options.codec=='x264':
 
  # http://www.mplayerhq.hu/DOCS/HTML/en/menc-feat-x264.html

  ovco = "-ovc x264 -x264encopts"

  if options.bitrate!=None:
    vbitrate = options.bitrate
  else:
    vbitrate = int(50 * options.fps * width * height / 256. / 1024.)
  
  print "vbitrate=%d"%(vbitrate)
  
  if options.quality == 1:
    ovc = "subq=4:bframes=2:b_pyramid:weight_b:bitrate=%d"%vbitrate
  if options.quality == 2:
    ovc = "subq=5:8x8dct:frameref=2:bframes=3:b_pyramid:weight_b:bitrate=%d"%vbitrate
  else:
    ovc = "subq=6:partitions=all:8x8dct:me=umh:frameref=5:bframes=3:b_pyramid:weight_b:bitrate=%d"%vbitrate


else:
  ovc =  "-ovc copy"


cmds = []
if options.codec=='x264':
  for p in xrange(options.passes):
    tmp = "%s pass=%s:%s"%(ovco,p+1,ovc)
    cmd = """mencoder mf://%s/*.png  -mf type=png:fps=%d %s %s -o %s"""%(options.dir,options.fps,tmp,oac,options.outputfile)
    cmds.append(cmd)  
else:
  ovc = "%s %s"%(ovco,ovc)
  cmd = """mencoder mf://%s/*.png  -mf type=png:fps=%d %s %s -o %s"""%(options.dir,options.fps,ovc,oac,options.outputfile)
  cmds.append(cmd)







if not options.nofilm:

  for cmd in cmds:
    print "############################"
    print cmd
    print "############################"
    os.system(cmd)

else:
  options.keep = True


if not options.keep:  
  os.system('rm %s/*'%(options.dir)) 
  os.removedirs('%s'%(options.dir))


