
import copy


################################################################################
class Units:
################################################################################
  '''
  Units
  '''
  def __init__(self,symbol,factor=1.,power=1,ulist=[]):
    '''
    '''
    self.symbol = symbol
    self.factor = float(factor)
    self.power  = int(power)
    self.ulist  = ulist
    
    self.GatherBaseUnits()

  def set_symbol(self,symbol):
    
    self.ulist  = [copy.deepcopy(self)]
    self.power  = 1
    self.factor = 1.0
    self.symbol = symbol
    self.GatherBaseUnits()  
    	
        
  def GatherBaseUnits(self):
    '''
    create the self.bases list (recursively)
    '''
    
    self.bases = []
    
    for unit in self.ulist:
      
      if len(unit.ulist)==0:		# a base unit
        self.bases.append(unit)   	# add the unit
      
      else:				# go deeper in the tree
        unit.GatherBaseUnits()
	self.bases = self.bases + unit.bases    

    self.ReduceBaseUnits()


  def ReduceBaseUnits(self):
  
    allunits = {}
    factorTot   = 1.0
    for base in self.bases:  
      factorTot = factorTot * base.factor
      if allunits.has_key(base.symbol):
        allunits[base.symbol] = allunits[base.symbol] +  base.power
      else:
        allunits[base.symbol] = base.power	

    self.bases = []
    flag = 1
    for key in allunits.keys():
      factor = 1.0
      if flag: 
        factor = factorTot
	flag = 0
      power  = allunits[key]
      symbol = key
      ulist  = [] 
      self.bases.append(Units(symbol,factor,power,ulist))	
    
    
    
    
  def str_base(self):
    
    string = ''
      
    if self.factor != 1.0:		         
      string = string + self.factor.__str__()
    					         
    if len(string) > 0: string = string + ' '    
    if self.symbol == None:		         
      string = string + '-'  		         
    else:				         
      string = string + self.symbol  	         
      					         
    if self.power != 1: 		         
      string = string + '^%d'%self.power
      
    return string  
       
          
  def __str__(self):  
    
    string = ''
    # (1) display using symbol 
    if self.factor != 0:
      string = string+self.str_base()
      
    str1 = "[%s]"%string  
             
    string = ''
    # (2) display using (base units)
    for base in self.bases:
      if len(string) > 0: string = string + ' '
      string = string + base.str_base()
    
    str2 = "[%s]"%string  

    if str2 != '[]':
      str1 = "%s\t%s"%(str1,str2)  
    
    return str1   
    
        
    
  def __mul__(self,y):  
    
    x = copy.deepcopy(self)
    y = copy.deepcopy(y)
    
    # multiply with a scalar
    if type(y)== int or type(y) == float :	
      x.factor = x.factor*y
      x.mul_list(x.ulist,y)
      x.GatherBaseUnits()
      return x  
    
    # multiply with a unit
    elif type(y)==type(x):
      factor = x.factor * y.factor
      
      if x.symbol == y.symbol:		# same units
        x.symbol = x.symbol
	x.power  = x.power + y.power
	x.ulist  = x.ulist + y.ulist
        x.GatherBaseUnits()
	return x		
	
      else:				# different unit
	symbol = None
	power  = 1
	factor = 1.0
	if x.ulist == []:			# case of a base unit
	  x.ulist = [copy.deepcopy(x)]
	if y.ulist == []:			# case of a base unit
	  y.ulist = [copy.deepcopy(y)]	  
	ulist = x.ulist + y.ulist
        return Units(symbol,factor,power,ulist)
    
    
  def __div__(self,y): 	
  
    x = copy.deepcopy(self)
    y = copy.deepcopy(y)
    
    # divide with a scalar
    if type(y) == int or type(y) == float:	
      x.factor = x.factor/y
      x.div_list(x.ulist,y)
      x.GatherBaseUnits()
      return x  
      
    # divide with a unit
    elif type(y)==type(x):
      factor = x.factor / y.factor
      
      if x.symbol == y.symbol:		# same units
        x.symbol = None
	x.power  = 0
	x.ulist  = []
        x.GatherBaseUnits()
	return x		
	
      else:				# different unit
	symbol = None
	power  = 1
	factor = 1.0
	if x.ulist == []:			# case of a base unit
	  x.ulist = [copy.deepcopy(x)]
	if y.ulist == []:			# case of a base unit
	  y.ulist = [copy.deepcopy(y)]	  
	y.pow_list(y.ulist,-1)
	y.GatherBaseUnits()
	ulist = x.ulist + y.ulist
        return Units(symbol,factor,power,ulist)

  def __rdiv__(self,y): 	
    x = copy.deepcopy(self)
    
    # divide with a scalar
    if type(y)==int or type(y)==float:	
      return y*x**-1  
    else:
      raise "TypeError","unsupported operand type(s) for / or div()"   
            
    
  def __pow__(self,y):
  
    x = copy.deepcopy(self)
    y = copy.deepcopy(y)
    
    # power with a scalar
    if type(y)==int or type(y)==float:    	
      x.factor = x.factor**y
      x.power  = x.power*y   
      x.pow_list(x.ulist,y)
      x.GatherBaseUnits()
      return x  
      
    else:
      raise "TypeError","unsupported operand type(s) for ** or pow()"         



  def mul_list(self,ulist,y):
   
    for u in ulist:
      if len(u.ulist)==0:		# a base unit
        u.factor = u.factor*y
      else:				# go deeper in the tree
        u.mul_list(u.ulist,y)
	u.GatherBaseUnits()
    
  def div_list(self,ulist,y):
   
    for u in ulist:
      if len(u.ulist)==0:		# a base unit
        u.factor = u.factor/y
      else:				# go deeper in the tree
        u.mul_list(u.ulist,y)
	u.GatherBaseUnits()

  def pow_list(self,ulist,y):
    
    for u in ulist:
      if len(u.ulist)==0:		# a base unit
        u.factor = u.factor**y
	u.power  = u.power*y      
      else:				# go deeper in the tree
        u.pow_list(u.ulist,y)
        u.GatherBaseUnits()
	
  __rmul__ = __mul__   
  	
    
  


    
################################################################################    
class UnitSystem:
################################################################################
  '''
  Units system
  '''
  def __init__(self,UnitSysName,UnitLst):
    '''
    UnitDic = {'lenght':UnitLenght,'mass':UnitMass,'time':UnitTime}
    '''
    
    UnitDic = {}
    for u in UnitLst:
      u.GatherBaseUnits() 
      
      if len(u.bases) > 1:
        raise "UnitSystemError","%s is not a pure Unit."%(u.symbol)  
      
      elif len(u.bases) == 1:
        UnitDic[u.bases[0].symbol] = u.bases[0] 
      
      elif len(u.bases) == 0:				# base unit
        UnitDic[u.symbol] = u
    
    dic_of_factors = {}
    dic_of_powers  = {}
    
    for UnitType in UnitDic.keys():
      u = UnitDic[UnitType]
      
      if u.bases ==[]:					# base unit
        dic_of_factors[UnitType] = 1/u.factor
	dic_of_powers[UnitType]  = u.power
      else:
        dic_of_factors[UnitType] = 1/u.bases[0].factor		
        dic_of_powers[UnitType]  = u.bases[0].power
    
    self.UnitDic = UnitDic
    self.dic_of_factors = dic_of_factors		# value of base unit in current units
    self.dic_of_powers  = dic_of_powers			# not usefull
        



  def get_UnitLength_in_cm(self):
    f = self.UnitDic['m'].factor
    c = PhysCte(f,Unit_m)		
    return c.into(cgs)
    

  def get_UnitMass_in_g(self):
    f = self.UnitDic['kg'].factor
    c = PhysCte(f,Unit_kg)		
    return c.into(cgs)


  def get_UnitTime_in_s(self):
    f = self.UnitDic['s'].factor
    c = PhysCte(f,Unit_s)		
    return c.into(cgs)


  def get_UnitVelocity_in_cm_per_s(self):
    return self.get_UnitLength_in_cm()/self.get_UnitTime_in_s()

  def get_UnitEnergy_in_cgs(self):
    return self.get_UnitMass_in_g()*self.get_UnitVelocity_in_cm_per_s()**2

  def get_UnitDensity_in_cgs(self):
    return self.get_UnitMass_in_g()/self.get_UnitLength_in_cm()**3









################################################################################
class PhysCte:
################################################################################
  '''
  Physical constant
  '''
  
  def __init__(self,value,Unit):
     #super(PhysCte,self).__init__(value)
     self.value = value
     self.Unit = Unit


  def factor_to_base(self):
  
     if self.Unit.bases == []:
       factor = self.Unit.factor
     else:
       factor = self.Unit.bases[0].factor
  
     return factor

     
  def __str__(self):
     return float.__str__(float(self.value)) + ' %s'%(self.Unit)
     
  
  def into(self,SystemName):			# here, we could return an object
       
    f = self.value
    
    if self.Unit.bases==[]:
      unit = self.Unit
      symbol = unit.symbol
      f = f * unit.factor * (SystemName.dic_of_factors[symbol])**unit.power
    else:
      for unit in self.Unit.bases:
        symbol = unit.symbol
        f = f * unit.factor * (SystemName.dic_of_factors[symbol])**unit.power
    
    return float(f)
    


  def __float__(self):
    return float(self.value)


  def __mul__(self,y):
    pass
    
  def __add__(self,y):
    pass
    
  def __sub__(self,y):
    pass
  
  def __div__(self,y):
    pass      
  
  def __pow__(self,y):
    pass     

################################################################################
# define some units
################################################################################

# base units
Unit_m	 = Units('m')
Unit_kg	 = Units('kg')
Unit_s	 = Units('s')
Unit_mol = Units('mol')
Unit_C   = Units('C')
Unit_K   = Units('K')


# other lenght units
Unit_cm = 0.01*Unit_m
Unit_cm.set_symbol('cm') 
Unit_km = 1e3*Unit_m
Unit_km.set_symbol('km')
Unit_Mm = 1e6*Unit_m
Unit_Mm.set_symbol('Mm')
Unit_Gm = 1e9*Unit_m
Unit_Gm.set_symbol('Gm')
Unit_kpc= 1000*3.085e18*Unit_cm
Unit_kpc.set_symbol('kpc')
Unit_pc= 3.085e18*Unit_cm
Unit_pc.set_symbol('pc')
Unit_ua= 1.495978e11*Unit_m
Unit_ua.set_symbol('ua')

# other mass units
Unit_g  = 0.001*Unit_kg
Unit_g.set_symbol('g')
Unit_Ms = 1.9891e33*Unit_g
Unit_Ms.set_symbol('Ms')
Unit_Mg = 2.23e11*Unit_Ms
Unit_Mg.set_symbol('Mg')
Unit_Mt = 5.9742e24*Unit_kg
Unit_Mt.set_symbol('Mt')
Unit_Mj = 317.893*Unit_Mt
Unit_Mj.set_symbol('Mj')

# other time units
Unit_h  = 3600*Unit_s 
Unit_h.set_symbol('h')
Unit_yr = 31536000*Unit_s
Unit_yr.set_symbol('yr')
Unit_kyr= 1e3*Unit_yr
Unit_kyr.set_symbol('kyr')
Unit_Myr= 1e6*Unit_yr
Unit_Myr.set_symbol('Myr')
Unit_dy= 86400*Unit_s
Unit_dy.set_symbol('days')
Unit_hr= 3600*Unit_s
Unit_hr.set_symbol('hr')

# other speed units
Unit_kmh= Unit_km/Unit_h
Unit_kmh.set_symbol('kmh')


# other units
Unit_V  = Unit_m/Unit_s
Unit_V2 = Unit_V**2
Unit_J  = Unit_kg *  Unit_V2
Unit_J.set_symbol('J') 
Unit_erg  = Unit_g * (Unit_cm/Unit_s)**2
Unit_erg.set_symbol('erg') 
Unit_N  = Unit_kg *   Unit_m/(Unit_s)**2
Unit_N.set_symbol('N') 
Unit_Pa = Unit_N/Unit_m**2
Unit_Pa.set_symbol('Pa') 
Unit_G  = Unit_N*Unit_m**2/Unit_kg**2
Unit_d  = Unit_kg/Unit_m**3

################################################################################
# define some common unit systems
################################################################################

mks = UnitSystem('mks',[Unit_m,  Unit_kg, Unit_s  , Unit_K, Unit_mol, Unit_C])
cgs = UnitSystem('cgs',[Unit_cm, Unit_g,  Unit_s  , Unit_K, Unit_mol, Unit_C])
gal = UnitSystem('gal',[Unit_kpc,Unit_Mg, Unit_Myr, Unit_K, Unit_mol, Unit_C])



'''
constantes are now defined in ctes.py

##################################################
# define some const. in mks
##################################################

slight 		= PhysCte(2.99792458e8,Unit_m/Unit_s)
G      		= PhysCte(6.6732e-11,Unit_G)
q_electron 	= PhysCte(1.6022e-19,Unit_C)
planck		= PhysCte(6.6262e-34,Unit_J*Unit_s)
boltzmann       = PhysCte(1.3807e-23,Unit_J/Unit_K)
m_electron	= PhysCte(9.1095e-31,Unit_kg)
m_proton	= PhysCte(1.6726e-27,Unit_kg)
m_neutron	= PhysCte(1.6750e-27,Unit_kg)
Na		= PhysCte(6.0220e+23,Unit_mol)
R		= PhysCte(8.3144e+00,Unit_J/Unit_mol)
Av		= PhysCte(6.828e-50 ,Unit_Pa*Unit_m**6)
Bv		= PhysCte(4.419e-29,Unit_m**3)

'''



#################################
def Set_SystemUnits_From_Params(params):
#################################
  '''
  return a system of units from given parameters
  
  params is a dictionary that must constains at least
  
  params['UnitVelocity_in_cm_per_s']
  params['UnitMass_in_g']
  params['UnitLength_in_cm']
  
  ''' 
  UnitVelocity_in_cm_per_s = params['UnitVelocity_in_cm_per_s']
  UnitMass_in_g 	   = params['UnitMass_in_g']
  UnitLength_in_cm	   = params['UnitLength_in_cm']
  UnitTime_in_s 	   = UnitLength_in_cm / UnitVelocity_in_cm_per_s

  # now from the params, define a system of units

  Unit_lenght	  = Unit_cm * UnitLength_in_cm
  Unit_mass	  = Unit_g  * UnitMass_in_g
  Unit_time	  = Unit_s  * UnitTime_in_s

  localsystem = UnitSystem('local',[Unit_lenght,  Unit_mass, Unit_time  , Unit_K, Unit_mol, Unit_C])
  
  return localsystem


