#!/usr/bin/env python


'''

# do not reload if the file is similar


'''



import sys
import os
import string
import getopt
import types

from numpy import *
from pNbody import *
from pNbody import Movie
from pNbody.param import *
from pNbody.libutil import *
import copy


SAVEDPARAMS = 'saved_parameters.py'
global gparam



####################################################################################
def read_params(paramname):
####################################################################################
  '''
  read param from a parameter Nbody file
  '''
  
  gparams = Params(paramname,None)
          
  # create new params
  defparam = {}
  for param in gparams.params:
    param_name = param[0]
    param_valu = param[3]
    
    defparam[param_name]=param_valu
  
  # aditional parameters    
  defparam['exec'] = None
  defparam['macro']= None
  defparam['n1'] = None
  defparam['n2'] = None  
  defparam['select'] = None    
  defparam['ftype'] = None  
  
  defparam['subs'] = None
  defparam['fdir'] = None	
  defparam['tdir'] = None	# track directory
  defparam['pfile'] = None      
  
  return defparam





  


####################################################################################
def write_params(file,nh,nw,width,height,params):
####################################################################################

  fd = open(file,'w')
  fd.write('''
nh = %d  		# number of horizontal frame								     
nw = %d			# number of vertical frame	
# size of subfilms						     
width = %d							     
height = %d					     
# size of the film						     
numByte = width * nw						     
numLine = height * nh						     						     
# init parameters						     
param = initparams(nh,nw)\n'''%(nh,nw,width,height))

  
  # take params from the first file
  keys = params.keys() 
  keys.sort()
  
  params = params[keys[0]]  
  
  for i in param.keys():
    fd.write('\n')  
    for key in param[i]:
    
      name   = key
      tpe    = gparams.get_type(key)
      value  = param[i][key]
      string = write_ascii_value(value,tpe,name)
      
      if tpe == "String" and value!=None:
        line = '''param[%d]['%s'] = "%s"\n''' %(i,name,string)
      else:
        line = '''param[%d]['%s'] = %s\n'''   %(i,name,string)      
    
      fd.write(line)

  
  fd.close()
  
  
  
####################################################################################
def initparams(nh,nw):
####################################################################################  
  '''
  
  create a dictionary "param"    and full it with nh*nw empty dictionary
  create a dictionary "allparam" and full it with nh*nw empty dictionary
  
  allparam = all parameters (from file nbodyparams and from file .py)
  param    = parameters from file .py
  '''
    
  global gparams 
  global allparam
  
  
  param    = {}
  allparam = {}
  
  n = 1
  for i in range(nh*nw):
    param[n]    = {'tdir':None}
    param[n]    = {'pfile':None}
    allparam[n] = {}   
    n = n + 1			         
    
  return param


####################################################################################
def version():
####################################################################################
  print 'version 2.0'
  sys.exit(0)  

####################################################################################
def help_message():
####################################################################################
  print '''Usage : mkgmov [option] output files 
  Options: -h        -- this help message
  	   -p	     -- parameter file
	   -f	     -- g-parameter file
	   -s	     -- disable the softening of rsp 
	   -c	     -- enable auto cd for each image
	   -z        -- convert time in redshift 
	   -u        -- do not reload file
	   --info    -- give the optimal factor for each files in the list
           --help    -- this help message
           --version -- displays version
	   
    '''  
  sys.exit(0)

####################################################################################
def check_arguments(options,xarguments):
####################################################################################

  param     = None
  gparam    = None
  verbose   = None
  info 	    = None
  pio       = 'no'
  sofrsp    = 1
  arange    = 0
  redshift  = 0
  readonlyonce = 0

  for a in options[:]:
    if a[0] == '-h':
      help_message()
 
    if a[0] == '--help':
      help_message()
      
    if a[0] == '--version':
      version()     

    if a[0] == '--info':
      info = 1 

    if a[0] == '--pio':
      pio = "yes" 
      
    if a[0] == '-v':
      verbose = 1    
        
    if a[0] == '-s':
      sofrsp = 0      

    if a[0] == '-c':
      arange = 1  

    if a[0] == '-z':
      redshift = 1  

    if a[0] == '-u':
      readonlyonce = 1                    

    if a[0] == '-p':
      if a[1] == '':
        help_message()
      else:
        param = a[1]
        continue

    if a[0] == '-f':
      if a[1] == '':
        help_message()
      else:
        gparam = a[1]
        continue
  
  try:
    output = xarguments[0] 
  except:
    help_message()
      
  files =  xarguments[1:] 
  files.sort()
  
  return files,output,param,gparam,verbose,info,pio,sofrsp,arange,redshift,readonlyonce


    
####################################################################################  
#
#  MAIN
#
#################################################################################### 

if mpi.mpi_IsMaster():

  try:
    options, xarguments = getopt.getopt(sys.argv[1:],'p:f:hvscfzu', ['info','pio','help','version'])
  except getopt.error:
    help_message()
    sys.exit(0)


  # check arguments
  files,output,param_file,gparam_file,verbose,info,pio,sofrsp,arange,redshift,readonlyonce = check_arguments(options,xarguments)

  nbody_parameter_file = gparam_file

  # verifie que output n'existe pas
  complete = 0
  if not info:
    if (os.path.exists(output)!=0):
    
      answer = raw_input('%s exists. Remove it (r) Continue it (c) Exit (q) ? ' %(output))    

      if len(answer)==0 or answer[0] == 'c':									
  	if not os.path.exists(SAVEDPARAMS):							
  	  print "Warning : %s does not exists."%(SAVEDPARAMS)					
  	  print "Using   : %s instead."%(SAVEDPARAMS)						
  	complete = 1
  	param_file = SAVEDPARAMS	
      elif answer[0] == 'r' :						       
  	os.remove(output)									
      else:
  	sys.exit()										
  												
       
      
  ###############################################################################################
  # a partir de snapshots ou de files, cree une liste de files
  ###############################################################################################

  ##############################################################
  # lit le fichier param s'il existe, sinon, valeurs par defaut

  if param_file == None:
    # nombre de sous film horiz et vertical
    nh = 1
    nw = 1
    # size of subfilms
    width = 128
    height = width/2 + width/4
    # info on subfilms
    param = initparams(nh,nw)
  elif os.path.exists(param_file):
    execfile(param_file) 
  else:
    print "Error : %s does not exists." %(param_file)
    sys.exit()  




  ##############################################################
  # compleete allparam with default values
  ##############################################################


  if nbody_parameter_file != None:
    paramname = nbody_parameter_file 
  else:
    paramname = PARAMETERFILE	

  # read the default parameters
  #paramname = PARAMETERFILE				  # ! ! ! not that good ! ! !
  gparams = Params(paramname,None)

  defparam = read_params(paramname)			  # ! ! ! read again	! ! !
  
  # compleete allparam
  n = 1
  for i in range(nh*nw):
    allparam[n] = copy.deepcopy(defparam)   
    n = n + 1				

  #############################################
  # read from pfile files if exists

  n = 1
  # loop over all sub-images
  for i in range(nh*nw):
    
    n = i+1
    
    if param[n]['pfile']!= None:  
  	  
      paramfile = param[n]['pfile']
    
      if not os.path.isfile(paramfile):
  	print "file %s does not exists..."%(paramfile)
  	sys.exit()
  	    
      allparam[n] = read_params(paramfile)
       
  ##############################################################
  # create params with allparam and add paramfiles
  ##############################################################

  params = {}
  listfiles = []

  basefiles = []
  for file in files:
    basefiles.append(os.path.basename(file))

  #############################################
  # read from trail files if exists

  n = 1
  # loop over all sub-images and verify that tdir exists
  for i in range(nh*nw):
    if not param[n].has_key('tdir'):    
      param[n]['tdir']= None  
    n = n + 1  

  n = 1
  # loop over all sub-images and verify that time exists
  for i in range(nh*nw):
    if not param[n].has_key('time'):    
      txt = "You must define the parameter time : ex param[1]['time'] = 'nb.time'"
      raise "parameter Error", txt
    n = n + 1  

  # loop over all sub-images
  n = 1
  for i in range(nh*nw):
          
    if param[n]['tdir']!= None: 		  # here, we use param (init with .py)
    
    
      directory = param[n]['tdir']	  
      if os.path.isdir(directory):
    	paramfiles = glob.glob('%s/*'%(directory))
    	paramfiles.sort()
    	
    
    	# loop over files
    	for paramfile in paramfiles:
    	
    	  snap_file = os.path.basename(paramfile)
    	  snap_file = os.path.splitext(snap_file)[0]
    					 
    	  try:  	  
    	    index = basefiles.index(snap_file) 
    	  except ValueError:					 
    	    print "file %s not in file list"%(snap_file)	 
    	    sys.exit()  					 
    								 
    	  listfiles.append(files[index]) 
    	  
    	# stop after first time here	  
    	break		  
      
    
  if len(listfiles)!=0:  
    files = listfiles
      
  # now, loop over files
  old_name = ''
  num = -1
  listfiles = []


  for file in files:

    # set name and redefine files (listfiles)
    name = os.path.basename(file)
    # add number to the file
    if name == old_name:
      num = num + 1
    else:
      num = 0  
  	    
    old_name = name
    name = "%s.%05d"%(name,num)
    file = "%s.%05d"%(file,num)
    listfiles.append(file)					      


    n=1
    # loop over images
    param_for_this_file = {}
    
    for i in range(nh*nw):  
      if param[n]['tdir']!= None:
  	directory = param[n]['tdir']
  	      
  	paramfile = os.path.join(directory,name)	   
  	
  	if not os.path.isfile(paramfile):
  	  print "file %s does not exists..."%(paramfile)
  	  sys.exit()
  	      
  	param_for_this_file[n] = read_params(paramfile)
  	      
      else:		  # if no params file for this image (use .py file)
  	param_for_this_file[n] = allparam[n]
  		      
  	
      # add parameter to the list
      params[file] = param_for_this_file				  
      n = n + 1


  files = listfiles  
    
  ##############################################################
  # compleete params with param
  ##############################################################


  # loop over files
  for file in files:
    n=1
    # loop over images
    for i in range(nh*nw):
      
      keys = param[n].keys()
      for key in keys:
  	params[file][n][key]=param[n][key]	  
      n = n + 1

  ##############################################################
  # open the movie file
  ##############################################################

  # size of the film						       
  numByte = width * nh  					       
  numLine = height * nw

  # ouverture du film
  if not info:
    f = Movie.Movie(output)
    
    if not complete:
      f.new(numByte,numLine)

    else:
      f.open('r+')



else:
  
  info         = None    
  pio          = None  
  arange       = None    
  redshift     = None   
  readonlyonce = None
  
  nh = None
  nw = None
  width = None
  height = None

  files = None
  params = None
  
#############################################################
# gather all parameters to the slaves
#############################################################

# gather options

info        = mpi.mpi_Bcast(0,info)
pio         = mpi.mpi_Bcast(0,pio)
arange      = mpi.mpi_Bcast(0,arange)
redshift    = mpi.mpi_Bcast(0,redshift)
readonlyonce= mpi.mpi_Bcast(0,readonlyonce)

nh = mpi.mpi_Bcast(0,nh)
nw = mpi.mpi_Bcast(0,nw)
width = mpi.mpi_Bcast(0,width)
height = mpi.mpi_Bcast(0,height)

files  = mpi.mpi_Bcast(0,files)
params = mpi.mpi_Bcast(0,params)


  																  	    
###################################################################################
###################################################################################
##
## main loop over all files
##
###################################################################################
###################################################################################




mn_opt = 0.
mx_opt = 0.
cd_opt = 0.

cflag = 1 

last_file = None

# loop over all files
for file in files:
  
  # check if the frame already exists
  if mpi.mpi_IsMaster():
    if complete and cflag:
      data = f.read_one()
    	  
      if f.current_time < f.stoptime:
    	continue   
      elif f.current_time == f.stoptime:
    	cflag = 0
    	continue
        


  # loop over all images
  data = []
  for fnum in range(1,nh*nw+1):
    
    
    ###############################			       
    ## take parameters from params 			       
    ###############################			       
  
    obs  =		params[file][fnum]['obs']	       
    if type(obs)==types.ListType:			       
      obs = read_ascii_value(obs,'ArrayObs','obs')	
          							       
    if params[file][fnum]['x0'] == None:		       
      x0 = None 					       
    else:						       
      x0 =		array(params[file][fnum]['x0'],float)  
    if params[file][fnum]['xp'] == None:		       
      xp = None 					       
    else:						       
      xp =		array(params[file][fnum]['xp'],float)  
    
    alpha  =		params[file][fnum]['alpha']	       
    view =		params[file][fnum]['view']	       
    r_obs =		params[file][fnum]['r_obs']	       
    clip =		params[file][fnum]['clip']	       
    cut =		params[file][fnum]['cut']	       
    eye =		params[file][fnum]['eye']	       
    dist_eye =  	params[file][fnum]['dist_eye']	       
    foc =		params[file][fnum]['foc']	       
    persp =		params[file][fnum]['persp']	       
    
    shape =		(width,height)			       
    #center =		params[file][fnum]['center']	       
    size =		params[file][fnum]['size']	       
    frsp =		params[file][fnum]['frsp']	       
    space =		params[file][fnum]['space']	       
    mode =		params[file][fnum]['mode']	       
    filter_name =	params[file][fnum]['filter_name']      
    filter_opts =	params[file][fnum]['filter_opts']      
    scale =		params[file][fnum]['scale']	       
    cd =		params[file][fnum]['cd']	       
    mn =		params[file][fnum]['mn']	       
    mx =		params[file][fnum]['mx']	       
    l_n =		params[file][fnum]['l_n']	       
    l_min =		params[file][fnum]['l_min']	       
    l_max =		params[file][fnum]['l_max']	       
    l_kx =		params[file][fnum]['l_kx']	       
    l_ky =		params[file][fnum]['l_ky']	       
    l_color =		params[file][fnum]['l_color']	       
    l_crush =		params[file][fnum]['l_crush']	       
    b_weight =  	params[file][fnum]['b_weight']	       
    b_xopts =		params[file][fnum]['b_xopts']	       
    b_yopts =		params[file][fnum]['b_yopts']	       
    b_color =		params[file][fnum]['b_color']	       
    
    # other params					       
    n1 =		params[file][fnum]['n1']	       
    n2 =		params[file][fnum]['n2'] 	       
    select =		params[file][fnum]['select']	       
    exec_param =	params[file][fnum]['exec']	       
    macro =		params[file][fnum]['macro']	       
    ftype =		params[file][fnum]['ftype']	
    time  =             params[file][fnum]['time']    
    
    # multi images					       
    subs =		params[file][fnum]['subs']	       
    
    # dir parameter
    fdir = 		params[file][fnum]['fdir']
                
    
    # set min and max			       
    if mn==0. and mx==0. and cd==0. and file!=files[0]:        
      mn = params[files[0]][fnum]['mn'] 		       
      mx = params[files[0]][fnum]['mx'] 		       
      cd = params[files[0]][fnum]['cd']     
    
    
    #
    # check if the file need to be read
    #
    must_be_read  = True
    if readonlyonce:  
      if last_file == os.path.splitext(file)[0]:
        must_be_read = False

    last_file = os.path.splitext(file)[0]
    
    
    
    if must_be_read:
      # open the model (remove the particle if there is one !!! not good !!!)
      fl = os.path.splitext(file)[0]
      if pio == 'yes':
        fl = os.path.splitext(fl)[0]
      
      if mpi.mpi_IsMaster():
        print "reading ",fl
      nb = Nbody(fl,ftype=ftype,pio=pio)
      
    
    
    ###############################			       
    ## take parameters from dir 			       
    ###############################	
    '''
    not implemented now
    '''

    if redshift:
      nb.tnow = 1/nb.tnow -1 

    if fdir != None:
      tnow = nb.tnow
           
    if arange == 1 :
      mn = 0.
      mx = 0.
      cd = 0.
        
    if view=='dk':
      mat = zeros(width*height)
      mat = mat.astype(int8)
      mat = mat.tostring()
      data.append(mat)
    else:

  
      if must_be_read:
      
        
      	############
      	# set time
      	############
	exec("nb.tnow = %s"%time)   
	      
            
      	############
      	# exec
      	############
      	if exec_param != None:
      	  exec(exec_param) 
	  
      	############	  
      	# macro
      	############
      	if macro != None:
      	  execfile(macro) 
 
      	############	  
      	# sub (n1,n2)
      	############
        if n1 != None and n2 != None:			  		     
      	  nb = nb.sub(n1,n2)  
			  

      	############	  
      	# select
      	############
        if select != None:				  	     
      	  nb = nb.select(select)
 
        
      # compute map1
      mat,matint,mn_opts,mx_opts,cd_opts = nb.Map(obs=obs,
        		   			   x0=x0,			   
        		   			   xp=xp,			   
        		   			   alpha=alpha,  		   
        		   			   view=view,			   
        		   			   r_obs=r_obs, 		   
        		   			   eye=eye,			   
        		   			   dist_eye=dist_eye,		   
        		   			   foc=foc,			   
        		   			   mode=mode,			   
        		   			   space=space, 		   
        		   			   persp=persp, 		   
        		   			   clip=clip,			   
        		   			   size=size,			   
        		   			   cut=cut,			   
        		   			   frsp=frsp,			   
        		   			   shape=shape, 		   
        		   			   filter_name=filter_name,	   
        		   			   filter_opts=filter_opts,	   
        		   			   scale=scale, 		   
        		   			   cd=cd,			   
        		   			   mn=mn,			   
        		   			   mx=mx,			   
        		   			   l_color=l_color,		   
        		   			   l_n=l_n,			   
        		   			   l_min=l_min, 		   
        		   			   l_max=l_max, 		   
        		   			   l_kx=l_kx,			   
        		   			   l_ky=l_ky,			   
        		   			   l_crush=l_crush,		   
        		   			   b_weight=b_weight,		   
        		   			   b_xopts=b_xopts,		   
        		   			   b_yopts=b_yopts,		   
        		   			   b_color=b_color,		   
			   			   subs=subs)	

      			   
      # !!!!!!!!!!!!!!!!!!!!
      # il faut recuperer les valeurs optimales
      # !!!!!!!!!!!!!!!!!!!!
      mn_opt = mn_opts[0]		# !!! source de probleme, il faut tout garder...
      mx_opt = mx_opts[0]
      cd_opt = cd_opts[0]
     
     
      # add the image to the matrices
      matint = transpose(matint.astype(int8)) 
      data.append(matint.tostring())
      
      if mpi.mpi_IsMaster():  
        print "%8.3f %002d min=%10.3e max=%10.3e cd=%10.3e"%(nb.tnow,fnum,mn_opt,mx_opt,cd_opt)
      

    if file==files[0] and mn==0. and mx==0. and cd==0.:
      params[file][fnum]['cd'] = cd_opt 	   
      params[file][fnum]['mn'] = mn_opt  	   
      params[file][fnum]['mx'] = mx_opt


  # save used parameters   
  if file==files[0]:
    if mpi.mpi_IsMaster():
      write_params(SAVEDPARAMS,nh,nw,width,height,params)				   
   
  
  datac = ''
  for i in range(nw):	# loop over the lines
    # append the nw pictures and sum to the one of the previous line	
    datac = datac + Movie.append_h(width,height,data[i*nh:(i+1)*nh])	  
    


  if not info:
    if mpi.mpi_IsMaster():
      f.write_pic(nb.tnow,datac)

  
if not info:
  if mpi.mpi_IsMaster():
    f.close()
