# standard modules
import os,sys,string,types

# array module
from numpy import *

import pyfits

import mpi


#################################
def checkfile(name):
#################################
  '''
  check if file "name" exists
  '''
  						   
  if name == None:								   
    raise "CheckFileError"	   
  										   
  if not os.path.isfile(name):						   
    raise "CheckFileError" 
 
##################################
#def end_of_file(f):
##################################
#  '''
#  Return True if we have reached the end of the file f, False instead
#  '''
#  p1 = f.tell()
#  f.seek(0,2)						   
#  p2 = f.tell()
#  f.seek(p1)
#  
#  if p1 == p2:
#    return True
#  else:
#    return False  

#################################
def end_of_file(f,MPI=None):
#################################
  '''
  Return True if we have reached the end of the file f, False instead
  '''
  
  if mpi.ThisTask == 0:
  
    p1 = f.tell()
    f.seek(0,2) 	        		    
    p2 = f.tell()
    f.seek(p1)
  
    if p1 == p2:
      status = True
    else:
      status = False 
  
  else:
    status = None
      
  status = mpi.mpi_Bcast(0,status)     
      
  return status   

#####################################################
def write_array(file,vec):	
#####################################################     	 
  '''
  write an 1d array into ascii file
  '''
  f = open(file,'w')
  for i in range(len(vec)):
    f.write("%f\n"%vec[i])
  f.close()  
  
#####################################################
def read_ascii(file,columns,lines=None,dtype=float):	
#####################################################     	 
    """[X,Y,Z] = READ('FILE',[1,4,13],lines=[10,1000])	     	 
    Read columns 1,4 and 13 from 'FILE'  from line 10 to 1000	 
    into array X,Y and Z
    
    file is either fd or name file
    				     	 
    """ 
    try: import mmap					     	 
    except: pass					     	 

    if type(file) != types.FileType:
      f = open(file,'r+')
    else:
      f = file
       				     	 
    f.seek(0,2) 					     	 
    fsize = f.tell()					     	 
    f.seek(0,0) 					     	 
    start = 1						     	 
    end = 0						     	 
    numcols=len(columns)				     	 
    							     	 
    if lines is None:					     	 
	lines = [start,end]				     	 
    elif len(lines) != 2:				     	 
	raise "lines!=[start,end]"			     	 
    [start,end] = map(int,lines)			     	 
    start = start - 1					     	 

    nl = len(f.readlines())
    f.seek(0)							     	 
    if (end == 0) or (end > nl):		     	 
	end = nl				     	 
    numlines = end - start    
               				         							     	 
    block = 65536  # chunk to read, in bytes		     	 
    data = mmap.mmap(f.fileno(), fsize) 		     	 
    for i in range(start):				     	 
	junk = data.readline()				     	 
							     	 
    numcols =  len(columns)				     	 
    if dtype == int:					     	 
       a = zeros((numlines,numcols), 'l')	     	 
    else:						     	 
       a = zeros((numlines,numcols), 'd')	     	 

    i = 0						     	 
    line = data.readline()				     	 
    while line and ( i < numlines):			     	 
	d = array(map(dtype,string.split(line)))    
	a[i,:] = take(d,columns)		     	 
	line = data.readline()				     	 
	i = i + 1					     	 
    data.close()
    if type(file) != types.FileType:
      f.close()						     	 
    b = []						     	 
    for i in range(numcols):				     	 
	b.append(a[:,i])				     	 
    del a						     	 
    return  b	 
 
  
  
#################################
def readblock(f,data_type,shape=None,byteorder=sys.byteorder):
#################################
  '''
  data_type = int,float32,float
  or
  data_type = array
  
  shape	    = tuple
  '''  
  
  try:
    nb1 = fromstring(f.read(4),int32)
    if sys.byteorder != byteorder:
      nb1.byteswap(True)
    nb1 = nb1[0]  
    
  except IndexError:
    raise "ReadBlockError" 
      
  if type(data_type) == types.TupleType:
       
    data = []
    for tpe in data_type:
        
      if type(tpe) == int:
        val =  f.read(tpe)
      else:
        bytes = dtype(tpe).itemsize
	val = fromstring(f.read(bytes),tpe)
        if sys.byteorder != byteorder:
          val.byteswap(True)  
	 
	val = val[0]  
      
      data.append(val)
        
  else:
    data  = fromstring(f.read(nb1),data_type)   
    if sys.byteorder != byteorder:
      data.byteswap(True)  
      
  nb2 = fromstring(f.read(4),int32)
  if sys.byteorder != byteorder:
    nb2.byteswap(True)
  nb2 = nb2[0]  

  if nb1 != nb2:
    raise "ReadBlockError","nb1=%d nb2=%d"%(nb1,nb2)   
  
  # reshape if needed
  if shape != None:
    data.shape=shape

    
  return data


#################################
def ReadBlock(f,data_type,shape=None,byteorder=sys.byteorder,pio='no'):
#################################
  '''
  data_type = int,float32,float
  or
  data_type = array
  
  shape	    = tuple
  
  pio   : parallel io, 'yes' or 'no'
          if 'yes', each proc read each file
	  if 'no',  proc 0 read and send to each other
  
  '''
  
  if mpi.NTask==1:
    data = readblock(f,data_type=data_type,shape=shape,byteorder=byteorder) 
    return data
    
  if pio == 'yes':
    data = readblock(f,data_type=data_type,shape=shape,byteorder=byteorder) 
    return data
    
  else:
    data = mpi.mpi_ReadAndSendBlock(f,data_type=data_type,shape=shape,byteorder=byteorder) 
    return data
  
  
    
#################################
def ReadArray(f,data_type,shape=None,byteorder=sys.byteorder,pio='no',nlocal=None):
#################################
  '''
  data_type = int,float32,float
  or
  data_type = array
  
  shape	    = tuple
    
  '''
  
  if mpi.NTask==1:
    data = readblock(f,data_type=data_type,shape=shape,byteorder=byteorder) 
    return data
    
  if pio == 'yes':
    data = readblock(f,data_type=data_type,shape=shape,byteorder=byteorder) 
    return data
    
  else:
    data = mpi.mpi_OldReadAndSendArray(f,data_type,shape=shape,byteorder=byteorder,nlocal=nlocal)


  return data

#################################
def ReadDataBlock(f,data_type,shape=None,byteorder=sys.byteorder,pio='no',npart=None):
#################################
  '''
  
  Read a block containg data.
  If NTask = 1 or  pio = 'yes', the block is read normally.
  If NTask > 1 and pio = 'no',  the master reads the block and send the data to the slaves.
  
 
  In the second case : 
  
  a) the master send N/Ntask element to each task.
  b) if the var npart is present, he send Np/Ntask to each task, for each Np of npart. 
  
 
  data_type = array
  
  shape	    = tuple
    
  '''
  
  if mpi.NTask==1 or pio == 'yes':
    data = readblock(f,data_type=data_type,shape=shape,byteorder=byteorder) 
    return data
        
  else:
    data = mpi.mpi_ReadAndSendArray(f,data_type,shape=shape,byteorder=byteorder,npart=npart)

  return data

    
#################################
def writeblock(f,data,byteorder=sys.byteorder):
#################################
  '''
  data = array
  or
  data = ((x,float32),(y,int),(z,float32),(label,40))
  
  shape	    = tuple
  '''
  
  if type(data) == types.TupleType: 
    
    # first, compute nbytes
    nbytes = 0
    for dat in data:
      if type(dat[0])==types.StringType:
        nbytes = nbytes + dat[1]
      else:
        nbytes = nbytes + array([dat[0]],dat[1]).type().bytes* array([dat[0]],dat[1]).size()
      
    nbytes = array([nbytes],int32)
    
    # write block
    if sys.byteorder != byteorder:
      nbytes.byteswap(True)
    
    f.write(nbytes.tostring())
    for dat in data:
      if type(dat[0])==types.StringType:
        f.write(string.ljust(dat[0],dat[1])[:dat[1]])
      else:
        ar = array([dat[0]],dat[1])
        if sys.byteorder != byteorder:
	  ar.byteswap(True)
        f.write(ar.tostring())
	
    f.write(nbytes.tostring())
  
  else:
    # write block
    #nbytes = array([data.type().bytes*data.size()],int)
    nbytes = array([data.nbytes],int32)
   
    if sys.byteorder != byteorder:
      nbytes.byteswap(True)
      data.byteswap(True)  
    
    f.write(nbytes.tostring())  
    f.write(data.tostring())
    f.write(nbytes.tostring()) 
  
  
  
#################################
def WriteBlock(f,data,byteorder=sys.byteorder):
#################################
  '''
  data = ((x,float32),(y,int),(z,float32),(label,40))
  
  shape	    = tuple
  '''
  
  if f!=None:
  
    if type(data) == types.TupleType: 
      
      # first, compute nbytes
      nbytes = 0
      for dat in data:
    	if type(dat[0])==types.StringType:
    	  nbytes = nbytes + dat[1]
    	elif type(dat[0]) == string_:
	  nbytes = nbytes + dat[1]
	else:
	  #nbytes = nbytes + array([dat[0]],dat[1]).type().bytes* array([dat[0]],dat[1]).size()
	  nbytes = nbytes + array([dat[0]],dat[1]).nbytes
    	
      nbytes = array([nbytes],int32)
      
      # write block
      if sys.byteorder != byteorder:
    	nbytes.byteswap(True)
      
      f.write(nbytes.tostring())
      for dat in data:
    	if type(dat[0])==types.StringType:
    	  f.write(string.ljust(dat[0],dat[1])[:dat[1]])
    	elif type(dat[0]) == string_:
	  f.write(string.ljust(dat[0],dat[1])[:dat[1]])
    	else:
    	  ar = array([dat[0]],dat[1])
    	  if sys.byteorder != byteorder:
    	    ar.byteswap(True)
    	  f.write(ar.tostring())
    	  
      f.write(nbytes.tostring())
  


#################################
def WriteArray(f,data,byteorder=sys.byteorder,pio='no',npart=None):
#################################
  '''
  data = array
  
  shape	    = tuple
  '''
  
  
  if mpi.NTask==1 or pio == 'yes':
    writeblock(f,data,byteorder=byteorder)
        
  else:
    mpi.mpi_GatherAndWriteArray(f,data,byteorder=byteorder,npart=npart)


#################################
def WriteDataBlock(f,data,byteorder=sys.byteorder,pio='no',npart=None):
#################################
  '''

  Write a block containg data.
  If NTask = 1 or  pio = 'yes', the block is written normally.
  If NTask > 1 and pio = 'no',  the master get the block from the slaves and write it.
  
  In the second case : 
  
  a) the master get N/Ntask element from each task.
  b) if the var npart is present, he get Np/Ntask from each task, for each Np of npart. 


  data = array
  
  shape	    = tuple
  '''
  
  if mpi.NTask==1 or pio == 'yes':
    writeblock(f,data,byteorder=byteorder)
    
  else:
    mpi.mpi_GatherAndWriteArray(f,data,byteorder=byteorder,npart=npart)


#####################################################
def ReadFits(filename) :
#####################################################

  # read image
  fitsimg = pyfits.open(filename)
  data = fitsimg[0].data
  return data
  

#####################################################
def WriteFits(data, filename, extraHeader = None) :
#####################################################

  # image creation							        		   
  fitsimg = pyfits.HDUList()						        		   


  # add data
  hdu = pyfits.PrimaryHDU()						        		   
  hdu.data = data							        		   
  fitsimg.append(hdu)


  # add keys	
  keys = []  
  if extraHeader != None:
    #keys.append(('INSTRUME','st4 SBIG ccd camera','Instrument name'))
    #keys.append(('LOCATION',"175 OFXB St-Luc (VS)",'Location'))
    keys = extraHeader

  									       
  hdr = fitsimg[0].header						        	    
  for key in keys:							        	    
    hdr.update(key[0],key[1],comment=key[2])				        	    


  fitsimg.writeto(filename) 



#####################################################
def old_WriteFits(data, filename, extraHeader = None,LittleEndian=1,bitpix=None) :
#####################################################
    
    if bitpix==None:   
      bitpix = data.dtype.itemsize*8
      
      if data.dtype == float64 or data.dtype == float32:
        bitpix = -bitpix


    shape = list(data.shape)
    shape.reverse()
    naxis = len(shape)
    header = [ 'SIMPLE  = T', 'BITPIX  = %d' % bitpix, 'NAXIS   = %d' % naxis ]
    for i in range(naxis) :
        keyword = 'NAXIS%d' % (i + 1)
        keyword = string.ljust(keyword, 8)
        header.append('%s= %d' % (keyword, shape[i]))
    if extraHeader != None :
        for rec in extraHeader :
            try :
                key = string.split(rec)[0]
            except IndexError :
                pass
            else :
                if key != 'SIMPLE' and \
                   key != 'BITPIX' and \
                   key[:5] != 'NAXIS' and \
                   key != 'END' :
                    header.append(rec)
    header.append('END')
    header = map(lambda x: string.ljust(x,80)[:80], header)
    header = string.join(header,'')
    numBlock = (len(header) + 2880 - 1) / 2880
    header = string.ljust(string.join(header,''), numBlock*2880)
    file = open(filename, 'w')
    file.write(header)
    if LittleEndian : data = data.byteswap()
    data = data.tostring()
    file.write(data)
    numBlock = (len(data) + 2880 - 1) / 2880
    padding = ' ' * (numBlock*2880 - len(data))
    file.write(padding)


#################################
def read_anul(file):
#################################
  '''																			   
  Read a file of type 'anul' created with the program 'coupe_sph5'	
  this file represents the ring decomposition of a galaxy														   

  file : name of the file	
  				
  The output consists of 
  
  R 		: radius of the ring
  ap,pp		: alpha and phi determined by the positions
  av,pv		: alpha and phi determined by the velocities
  xp,yp,zp	: normal vector of the ring determined by the positions   
  xv,yv,zv	: normal vector of the ring determined by the velocities   
  lx,ly,lz	: normal vector to the inner regions of the galaxy, that defines
  		  the referential frame
  																			   															   
  '''	 
  
  f = open(file)
  lines = f.readlines()
  f.close()
  
  # R alpha_p phi_p alpha_v phi_v lx_p ly_p lz_p lx_v ly_v lz_v lx, ly ,lz, -, -, -, n
  R  = array([],float)
  ap = array([],float)
  pp = array([],float)
  av = array([],float)
  pv = array([],float)
  xp = array([],float)
  yp = array([],float)
  zp = array([],float)
  xv = array([],float)
  yv = array([],float)
  zv = array([],float)
  n = array([],int)
  
  for line in lines[1:]:
    if line != "":
      R  = concatenate((R ,[float(line.split()[0])]))
      ap = concatenate((ap,[float(line.split()[1])]))
      pp = concatenate((pp,[float(line.split()[2])]))
      av = concatenate((av,[float(line.split()[3])]))
      pv = concatenate((pv,[float(line.split()[4])]))
      xp = concatenate((xp,[float(line.split()[5])]))      
      yp = concatenate((yp,[float(line.split()[6])]))
      zp = concatenate((zp,[float(line.split()[7])]))
      xv = concatenate((xv,[float(line.split()[8])]))      
      yv = concatenate((yv,[float(line.split()[9])]))
      zv = concatenate((zv,[float(line.split()[10])]))
      lx = float(line.split()[11])
      ly = float(line.split()[12])
      lz = float(line.split()[13])
      n = concatenate((n,[float(line.split()[17])]))
      
  return R,ap,pp,av,pv,xp,yp,zp,xv,yv,zv,lx,ly,lz,n   


#################################
def read_cooling(file):
#################################
  '''
  Read cooling file
  '''

  f = open(file,'r')
  f.readline()
  f.readline()
  lines = f.readlines()
  f.close()
  
  lines = map(string.strip,lines)
  elts = map(string.split,lines)
  
  logT	= array(map(lambda x:float(x[0]),elts))
  logL0	= array(map(lambda x:float(x[1]),elts))
  logL1	= array(map(lambda x:float(x[2]),elts))
  logL2	= array(map(lambda x:float(x[3]),elts))
  logL3	= array(map(lambda x:float(x[4]),elts))
  logL4	= array(map(lambda x:float(x[5]),elts))
  logL5	= array(map(lambda x:float(x[6]),elts))
  logL6	= array(map(lambda x:float(x[7]),elts))          
      
  return logT,logL0,logL1,logL2,logL3,logL4,logL5,logL6


