####################################################################################################################################
#
# GADGET CLASS
#
####################################################################################################################################    

class Nbody_gadget(NbodyDefault):

  
  def get_read_fcts(self):
    return [self.read_particles]
    
  def get_write_fcts(self):
    return [self.write_particles]    

  def get_mxntpe(self):
    return 6

  def get_default_spec_vars(self):
    '''
    return specific variables default values for the class
    '''
    return {'massarr'              :array([0,0,self.nbody,0,0,0]),
            'atime'                :0.,
            'redshift'             :0.,
	    'flag_sfr'             :0,
	    'flag_feedback'        :0,
	    'nall'                 :array([0,0,self.nbody,0,0,0]),
	    'flag_cooling'         :0,
	    'num_files'            :1,
	    'boxsize'              :0.,
	    'omega0'               :0.,
	    'omegalambda'          :0.,
	    'hubbleparam'          :0.,
	    'flag_age'             :0.,
	    'hubbleparam'          :0.,
	    'flag_metals'          :0.,
	    'nallhw'               :array([0,0,0,0,0,0]),
	    'flag_entr_ic'         :0,
	    'critical_energy_spec' :0.,
	    'empty'                :52*' '
	    }



  def get_massarr_and_nzero(self):
    """
    return massarr and nzero
    """
   
    if self.has_var('self.massarr') and self.has_var('self.nzero'):
      if self.massarr !=None and self.nzero!=None: 
        print "warning : get_massarr_and_nzero : here we use massarr and nzero",self.massarr,self.nzero
        return self.massarr,self.nzero
 
    massarr = zeros(len(self.npart),float)
    nzero = 0
    
    # for each particle type, see if masses are equal
    for i in range(len(self.npart)):
      first_elt = sum((arange(len(self.npart)) < i) * self.npart)
      last_elt  = first_elt + self.npart[i]

      if first_elt!=last_elt:
        c = (self.mass[first_elt]==self.mass[first_elt:last_elt]).astype(int)
        if sum(c)==len(c):
          massarr[i] = self.mass[first_elt]
        else:
          nzero = nzero+len(c)	
	
    return massarr.tolist(),nzero



    

  def read_particles(self,f):
    '''
    read gadget file
    '''
    ##########################################
    # read the header and send it to each proc
    ##########################################
    tpl = (24,48,float,float,int32,int32,24,int32,int32,float,float,float,float,int32,int32,24,int32,float,52)
    header = io.ReadBlock(f,tpl,byteorder=self.byteorder,pio=self.pio)
    npart,massarr,atime,redshift,flag_sfr,flag_feedback,nall,flag_cooling,num_files,boxsize,omega0,omegalambda,hubbleparam,flag_age,flag_metals,nallhw,flag_entr_ic,critical_energy_spec,empty = header 

    npart   = fromstring(npart,int32)
    massarr = fromstring(massarr,float)
    nall    = fromstring(nall,int32)
    nallhw  = fromstring(nallhw,int32)   

          	
    if sys.byteorder != self.byteorder:
      npart.byteswap(True)
      massarr.byteswap(True)
      nall.byteswap(True)
      nallhw.byteswap(True)
                  	     	
    ##########################################
    # computes nzero
    ##########################################	
    # count number of particles that have non constant masses and then
    # have masses storded further	     
            	       
    if self.pio == 'no':

      npart_tot    = npart	
      npart_all    = libutil.get_npart_all(npart,mpi.mpi_NTask())
      npart        = npart_all[mpi.mpi_ThisTask()]			# local
      npart_read   = npart_tot
      nbody_read   = sum(npart_read)
      
      npart_m_read = npart_tot * (massarr==0)
      
      ngas         = npart[0]
      ngas_read    = npart_tot[0]
      
      # compute nzero
      nzero=0							      
      mass = array([])  					      
      for i in range(len(npart_tot)):				      
        if massarr[i] == 0:					      
          nzero = nzero+npart_tot[i]				      
        else:	  
          mass = concatenate((mass,ones(npart[i])*massarr[i]))              
            
    else:
   
      npart_tot    = mpi.mpi_reduce(npart) 
      npart_all    = None						# each proc read for himself
      npart        = npart						# local
      npart_read   = None						# each proc read for himself
      nbody_read   = sum(npart)
      
      npart_m_read = None						# each proc read for himself
     
      ngas         = npart[0]
      ngas_read    = ngas
        

      # compute nzero
      nzero=0							      
      mass = array([])  					      
      for i in range(len(npart)):				      
        if massarr[i] == 0:					      
          nzero = nzero+npart[i]				      
        else:	  
          mass = concatenate((mass,ones(npart[i])*massarr[i]))     


    nbody	= sum(npart)
    nbody_tot	= sum(npart_tot)
    
    if npart_m_read !=None:
      if sum(npart_m_read) != nzero:
        raise "sum(npart_m) (%d) != nzero (%d)"%(sum(npart_m_read),nzero),npart_m_read

    
    		
    ##########################################
    # read and send particles attribute
    ##########################################        
    pos = io.ReadDataBlock(f,float32,shape=(nbody_read,3),byteorder=self.byteorder,pio=self.pio,npart=npart_read)
    vel = io.ReadDataBlock(f,float32,shape=(nbody_read,3),byteorder=self.byteorder,pio=self.pio,npart=npart_read)
    num = io.ReadDataBlock(f,int32  ,shape=(nbody_read,) ,byteorder=self.byteorder,pio=self.pio,npart=npart_read)
            
    ##########################################
    # read mass if needed
    ##########################################   
    if nzero != 0:   
      massnzero = io.ReadDataBlock(f,float32,shape=(nzero,),byteorder=self.byteorder,pio=self.pio,npart=npart_m_read)
      mass = concatenate((mass,massnzero))
                
      
      	
    # extentions
    u   = None
    rho = None
    rsp = None
    opt = None
    opt2= None
    erd = None
    dte = None
    pot = None
    
    if not io.end_of_file(f,MPI=MPI):
      u   = io.ReadDataBlock(f,float32,shape=(ngas_read,),byteorder=self.byteorder,pio=self.pio,npart=None)
      u   = concatenate((u,zeros(nbody-ngas).astype(float32)))
    
    if not io.end_of_file(f,MPI=MPI):
      rho = io.ReadDataBlock(f,float32,shape=(ngas_read,),byteorder=self.byteorder,pio=self.pio,npart=None)
      rho = concatenate((rho,zeros(nbody-ngas).astype(float32)))
    
    if not io.end_of_file(f,MPI=MPI): 
      rsp  = io.ReadDataBlock(f,float32,shape=(ngas_read,),byteorder=self.byteorder,pio=self.pio,npart=None)
      rsp  = concatenate((rsp,zeros(nbody-ngas).astype(float32)))
    
    # here it is the end of the minimal output
    if not io.end_of_file(f,MPI=MPI):
      opt  = io.ReadDataBlock(f,float32,shape=(ngas_read,),byteorder=self.byteorder,pio=self.pio,npart=None)
      if (len(opt)==ngas): 
        opt  = concatenate((opt,zeros(nbody-ngas).astype(float32)))   

    if not io.end_of_file(f,MPI=MPI):
      opt2 = io.ReadDataBlock(f,float32,shape=(ngas_read,),byteorder=self.byteorder,pio=self.pio,npart=None)
      if (len(opt2)==ngas): 
        opt2 = concatenate((opt2,zeros(nbody-ngas).astype(float32)))    

    if not io.end_of_file(f,MPI=MPI):
      erd  = io.ReadDataBlock(f,float32,shape=(ngas_read,),byteorder=self.byteorder,pio=self.pio,npart=None)
      if (len(erd)==ngas):
        erd  = concatenate((erd,zeros(nbody-ngas).astype(float32)))
    
    if not io.end_of_file(f,MPI=MPI):
      dte  = io.ReadDataBlock(f,float32,shape=(ngas_read,),byteorder=self.byteorder,pio=self.pio,npart=None)
      dte  = concatenate((dte,zeros(nbody-ngas).astype(float32)))
    
    if not io.end_of_file(f,MPI=MPI):
      pot  = io.ReadDataBlock(f,float32,shape=(ngas_read,),byteorder=self.byteorder,pio=self.pio,npart=None)
    	 

    # make global				      
    self.npart = npart      
    self.massarr = massarr
    		      
    self.atime = atime				      
    self.redshift = redshift			      
    self.flag_sfr = flag_sfr			      
    self.flag_feedback = flag_feedback		      
    self.nall = nall				      
    self.flag_cooling = flag_cooling		      
    self.num_files = num_files			      
    self.boxsize = boxsize			      
    self.omega0 = omega0			      
    self.omegalambda = omegalambda		      
    self.hubbleparam = hubbleparam		      
    self.flag_age = flag_age			      
    self.flag_metals = flag_metals		      
    self.nallhw = nallhw			      
    self.flag_entr_ic = flag_entr_ic		      
    self.critical_energy_spec = critical_energy_spec  
    self.empty = empty				      
    self.nbody =  nbody 			      
    						      
    self.pos = pos				      
    self.vel = vel				      
    self.mass = mass				      
    self.num = num				      
    
    self.tpe = array([],int32)
    for i in range(len(npart)):
      self.tpe = concatenate( (self.tpe,ones(npart[i])*i) )
    
    self.nzero = nzero				      
    self.u = u					      
    self.rho = rho	
    self.opt = opt
    self.opt2= opt2			      
    self.erd = erd				      
    self.dte = dte
    self.pot = pot				      
    
    self.rsp = rsp		
    
    
    if type(self.massarr) == ndarray:
      self.massarr = self.massarr.tolist()
    if type(self.nall) == ndarray:
      self.nall = self.nall.tolist()	  
    if type(self.nallhw) == ndarray:
      self.nallhw = self.nallhw.tolist()      
    
    
    
    
        

  def write_particles(self,f):
    '''
    specific format for particle file
    '''
    
    # here, we must let the user decide if we creates
    # the mass block or not, event if all particles have the same mass
    massarr,nzero = self.get_massarr_and_nzero()
        
    if self.pio == 'yes':
            
      '''
      here, we have to compute also 
      mass for each proc
      '''	    
	    
      npart = self.npart
      nall  = self.npart_tot
      num_files = mpi.NTask
      
      npart_write = None
      npart_m_write = None
      
    else:

      npart     = self.npart
      nall      = self.npart_tot
      npart_all = array(mpi.mpi_Allgather(npart))
      num_files = 1
      
      npart_write = self.npart
      npart_m_write = None
      
      # compute the global massarr and global nzero
      nzero_tot = mpi.mpi_sum(nzero)
      massarr_all = array(mpi.mpi_Allgather(massarr))
      massarr_tot = zeros(len(npart),float)
            
            
      for i in range(len(npart)):
        
	# keep only values where there are particles
	massarr_all_red = compress(npart_all[:,i]!=0,massarr_all[:,i])
		  
	if len(massarr_all_red)>0:
	  if (massarr_all_red == massarr_all_red).all():
	    massarr_tot[i] = massarr_all[0,i]
	  else: 		  # not equal
	    raise "this case is not implemented"
	    massarr_tot[i] = 0.0
	    nzero_tot = nzero_tot + sum(npart_write[:,i])
	
	
	  
      # now, re-compute nzero for the current node
      massarr = massarr_tot
      nzero = 0
      for i in range(len(npart)):
        if massarr[i] == 0:							       
          nzero = nzero+npart[i]   
              
	    
    # now that we have the right massarr and nzero, 
    # we can compute massnzero for each node
    nzero_all = zeros((mpi.NTask,len(self.npart))) 
    
    if nzero != 0:
      ni = 0
      massnzero=array([],float32)
      for i in range(len(self.npart)): 
    	if npart[i]!=0 and massarr[i]==0.:
    	  massnzero = concatenate((massnzero,self.mass[ni:ni+self.npart[i]]))
          nzero_all[mpi.ThisTask,i] = nzero_all[mpi.ThisTask,i] + npart[i] 
    	
	ni = ni + self.npart[i]

    nzero_all = mpi.mpi_reduce(nzero_all)
        	
	
    if self.pio == 'yes':
      pass
    else:
      npart = self.npart_tot
      nzero = mpi.mpi_reduce(nzero)	# to ensure that all nodes
                                        # will do -> write mass if needed

    # header																																							   
    if sys.byteorder == self.byteorder: 																																			   
      npart	    = array(npart,int32).tostring()	
      massarr	    = array(massarr,float).tostring()																																		   
    else:																																							   
      npart	    = array(npart,int32).byteswapped().tostring()																																   
      massarr	    = array(massarr,float).byteswapped().tostring()	    																															   
    
    atime	  = self.atime																																					   
    redshift	  = self.redshift																																				   
    flag_sfr	  = self.flag_sfr																																				   
    flag_feedback = self.flag_feedback																																				   
    
    if sys.byteorder == self.byteorder: 																																			   
      nall	    = array(nall,int32).tostring()																																		   
    else:																																							   
      nall	    = array(nall,int32).byteswapped().tostring()																																	   
    
    flag_cooling  = self.flag_cooling																																				   
    num_files	  = num_files																																					   
    boxsize	  = self.boxsize																																				   
    omega0	  = self.omega0 																																				   
    omegalambda   = self.omegalambda																																				   
    hubbleparam   = self.hubbleparam																																				   
    flag_age	  = self.flag_age																																				   
    flag_metals   = self.flag_metals																																				   
    
    if sys.byteorder == self.byteorder: 																																			   
      nallhw	    = array(self.nallhw,float).tostring()																																	   
    else:  																																							   
      nallhw	    = array(self.nallhw,float).byteswapped().tostring() 																															   
    
    flag_entr_ic	  = self.flag_entr_ic 																																			   
    critical_energy_spec  = self.critical_energy_spec	 																																	   
    empty		  = self.empty																																				   



    																																								   
    tpl = ((npart,24),(massarr,48),(atime,float),(redshift,float),(flag_sfr,int32),(flag_feedback,int32),(nall,24),(flag_cooling,int32),(num_files,int32),(boxsize,float),(omega0,float),(omegalambda,float),(hubbleparam,float),(flag_age,int32),(flag_metals,int32),(nallhw,24),(flag_entr_ic,int32),(critical_energy_spec,float),(empty,52))  
    io.WriteBlock(f,tpl,byteorder=self.byteorder)																																		   
        
    # positions 
    io.WriteArray(f,self.pos.astype(float32),byteorder=self.byteorder,pio=self.pio,npart=npart_write)																												   
    # velocities																																						   
    io.WriteArray(f,self.vel.astype(float32),byteorder=self.byteorder,pio=self.pio,npart=npart_write)																												   
    # id																																							   
    io.WriteArray(f,self.num.astype(int32),byteorder=self.byteorder,pio=self.pio,npart=npart_write)																													   
    
    # write mass if needed
    if nzero != 0:
      io.WriteArray(f,massnzero.astype(float32),byteorder=self.byteorder,pio=self.pio,npart=npart_m_write)

    # write extension																																						   
    if self.has_array('u'):	
      if self.u != None: 
        io.WriteArray(f,self.u[:self.npart[0]].astype(float32),byteorder=self.byteorder,pio=self.pio,npart=None)  																										   
    if self.has_array('rho'):
      if self.rho != None:
        io.WriteArray(f,self.rho[:self.npart[0]].astype(float32),byteorder=self.byteorder,pio=self.pio,npart=None) 																										   
    if self.has_array('rsp'):
      if self.rsp != None:
        io.WriteArray(f,self.rsp[:self.npart[0]].astype(float32),byteorder=self.byteorder,pio=self.pio,npart=None)			  																					     
    # this is the end of the minimal output
    if self.has_array('opt'):  
      if self.opt != None:															  																					     
        io.WriteArray(f,self.opt[:self.npart[0]].astype(float32),byteorder=self.byteorder,pio=self.pio,npart=None)			  																					     
    if self.has_array('opt2'):  
      if self.opt2 != None:															  																					     
        io.WriteArray(f,self.opt2[:self.npart[0]].astype(float32),byteorder=self.byteorder,pio=self.pio,npart=None)			  																					     



    if self.has_array('erd'):  
      if self.erd != None:															  																					     
        io.WriteArray(f,self.erd[:self.npart[0]].astype(float32),byteorder=self.byteorder,pio=self.pio,npart=None)			  																					     
    if self.has_array('dte'):
      if self.dte != None:															  																					     
        io.WriteArray(f,self.dte[:self.npart[0]].astype(float32),byteorder=self.byteorder,pio=self.pio,npart=None)			  																					     
    if self.has_array('pot'):
      if self.pot != None:															  																					     
        io.WriteArray(f,self.pot[:self.npart[0]].astype(float32),byteorder=self.byteorder,pio=self.pio,npart=None)			  																					     
    																																								   
    																																								   
	
	
        
        
  def spec_info(self):
    """
    Write spec info
    """	
    infolist = []
    infolist.append("")
    #infolist.append("nzero               : %s"%self.nzero)	
    #infolist.append("npart               : %s"%self.npart)		
    #infolist.append("massarr             : %s"%self.massarr)       
    infolist.append("atime               : %s"%self.atime)		
    infolist.append("redshift            : %s"%self.redshift)       
    infolist.append("flag_sfr            : %s"%self.flag_sfr)       
    infolist.append("flag_feedback       : %s"%self.flag_feedback)  
    infolist.append("nall                : %s"%self.nall)		
    infolist.append("flag_cooling        : %s"%self.flag_cooling)   
    infolist.append("num_files           : %s"%self.num_files)      
    infolist.append("boxsize             : %s"%self.boxsize)        
    infolist.append("omega0              : %s"%self.omega0)  	
    infolist.append("omegalambda         : %s"%self.omegalambda)    
    infolist.append("hubbleparam         : %s"%self.hubbleparam)   
    infolist.append("flag_age            : %s"%self.flag_age)
    infolist.append("flag_metals         : %s"%self.flag_metals)
    infolist.append("nallhw              : %s"%self.nallhw)
    infolist.append("flag_entr_ic        : %s"%self.flag_entr_ic) 
    infolist.append("critical_energy_spec: %s"%self.critical_energy_spec) 


    infolist.append("")
    if self.has_array('u'):
      infolist.append("len u               : %s"%len(self.u))
      infolist.append("u[0]                : %s"%self.u[0])
      infolist.append("u[-1]               : %s"%self.u[-1])
    if self.has_array('rho'):  
      infolist.append("len rho             : %s"%len(self.rho))
      infolist.append("rho[0]              : %s"%self.rho[0])
      infolist.append("rho[-1]             : %s"%self.rho[-1])  
    if self.has_array('rsp'):  
      infolist.append("len rsp             : %s"%len(self.rsp))
      infolist.append("rsp[0]              : %s"%self.rsp[0])
      infolist.append("rsp[-1]             : %s"%self.rsp[-1]) 
    if self.has_array('opt'):  
      infolist.append("len opt             : %s"%len(self.opt))
      infolist.append("opt[0]              : %s"%self.opt[0])
      infolist.append("opt[-1]             : %s"%self.opt[-1])     
    if self.has_array('opt2'):  
      infolist.append("len opt2            : %s"%len(self.opt2))
      infolist.append("opt2[0]             : %s"%self.opt2[0])
      infolist.append("opt2[-1]            : %s"%self.opt2[-1])     
    if self.has_array('erd'):  
      infolist.append("len erd             : %s"%len(self.erd))
      infolist.append("erd[0]              : %s"%self.erd[0])
      infolist.append("erd[-1]             : %s"%self.erd[-1]) 
    if self.has_array('dte'):  
      infolist.append("len dte             : %s"%len(self.dte))
      infolist.append("dte[0]              : %s"%self.dte[0])
      infolist.append("dte[-1]             : %s"%self.dte[-1]) 
      
    return infolist  
      


  #def select(self,tpe='gas'):
  def select(self,*arg,**kw):
    """ 
    Return an N-body object that contain only particles of a 
    certain type, defined by in gadget:

    gas		: 	gas particles
    halo	: 	halo particles
    disk	:	disk particles
    bulge	:	bulge particles
    stars	:	stars particles
    bndry	:	bndry particles
    
    sph         : 	gas with u > u_c
    sticky      : 	gas with u < u_c

    """   

    index = {'gas':0,'halo':1,'disk':2,'bulge':3,'stars':4,'bndry':5}

    # this allows to write nb.select(('gas','disk'))
    if len(arg)==1:
      if type(arg[0])==types.TupleType:
        arg = arg[0]

    tpes = arg    
           
    
    
    # create the selection vector
    c = zeros(self.nbody)
                 
    for tpe in tpes:
      if type(tpe) == types.StringType: 
      
      	if   (tpe=='sph'):
	  c = c+(self.u>self.critical_energy_spec)*(self.tpe==0)
	  
      	elif (tpe=='sticky'):
	  c = c+(self.u<self.critical_energy_spec)*(self.tpe==0)

      	elif (tpe=='wg'):
	  c = c+(self.u>0)*(self.tpe==0)
	  
      	elif (tpe=='cg'):
	  c = c+(self.u<0)*(self.tpe==0)


        elif not index.has_key(tpe):
          print "unknown type, do nothing %s"%(tpe)
        
	else:
          i = index[tpe]
          c = c+(self.tpe==i)
	  	  	
      elif type(tpe) == types.IntType:	
        
	c = c+(self.tpe==tpe)
	
	        
    return self.selectc(c)
      
    
    '''
    
    
    elif type(tpe) == types.StringType:
    
      if (tpe=='sph'):
        nb = self.select('gas')
        return nb.selectc((self.u>self.critical_energy_spec))
    
      if (tpe=='sticky'):
        nb = self.select('gas')
        return nb.selectc((nb.u<=nb.critical_energy_spec))
        
    

      if not index.has_key(tpe):
        print "unknown type %s"%(tpe)
        return self
    
      i = index[tpe]
      
    else:
      i = tpe
        
        
    if self.npart[i]==0:
      #print "no particle of type %s"%(tpe)
      return self.selectc(zeros(self.nbody))
    
    n1 = sum(self.npart[:i])
    n2 = n1 + self.npart[i]-1
    
    return self.sub(n1,n2)
    
    '''
     

  def subdis(self,mode='dd',val=None):
    """ 
    Equivalent of select
    """
    return self.select(mode)
   
  
 
