from numpy import *
from ctes import *
import units
import thermodyn
import io

####################################################################################################################################
#
# THERMODYNAMIC RELATIONS
#
####################################################################################################################################    

# deflauts parameters in cgs
defaultpars = {"k":BOLTZMANN,"mh":PROTONMASS,"mu":2,"gamma":5/3.}


###################
def Prt(rho,T,pars=defaultpars):
###################
  '''
  P(rho,T)
  '''
  
  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
  
  mumh = mu*mh
  
  return k*T/(mumh/rho) 


###################
def Trp(rho,P,pars=defaultpars):
###################
  '''
  T(rho,P)
  '''
  
  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
  
  mumh = mu*mh
  
  return (mumh/rho)/k * ( P )



###################
def Art(rho,T,pars=defaultpars):
###################
  '''
  A(rho,T)
  '''

  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
  
  mumh = mu*mh

  return k/mumh * rho**(1.-gamma) * T
  
###################
def Tra(rho,A,pars=defaultpars):
###################
  '''
  T(rho,A)
  '''

  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
  
  mumh = mu*mh

  return mumh/k * rho**(gamma-1.) * A


###################
def Urt(rho,T,pars=defaultpars):
###################
  '''
  U(rho,T)
  '''
  
  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
  
  mumh = mu*mh
  
  return 1./(gamma-1.)* k*T/mumh


###################
def Tru(rho,U,pars=defaultpars):
###################
  '''
  T(rho,U)
  '''  
  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
  
  mumh = mu*mh
  return (gamma-1.)* mumh/k * U
  
    
   
  
###################
def Pra(rho,A,pars=defaultpars):
###################
  '''
  P(rho,A)
  '''
  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
  
  mumh = mu*mh
    
  return rho**gamma * A 


###################
def Arp(rho,P,pars=defaultpars):
###################
  '''
  A(rho,P)
  '''
  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
  
  mumh = mu*mh
  
  return rho**-gamma * P



###################
def Pru(rho,U,pars=defaultpars):
###################
  '''
  P(rho,U)
  '''
  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
  
  mumh = mu*mh
  
  return (gamma-1.) * rho * U 


###################
def Urp(rho,P,pars=defaultpars):
###################
  '''
  U(rho,P)
  '''
  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
    
  mumh = mu*mh
  
  return  1./(gamma-1.) * (1/rho) *  P


###################
def Ura(rho,A,pars=defaultpars):
###################
  '''
  U(rho,A)
  '''
  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
  
  mumh = mu*mh
  
  return 1./(gamma-1.) * rho**(gamma-1.) * A


###################
def Aru(rho,U,pars=defaultpars):
###################
  '''
  A(rho,U)
  '''
  k     = pars['k']
  mh    = pars['mh']
  mu    = pars['mu']
  gamma = pars['gamma']
    
  mumh = mu*mh
  
  return (gamma-1.) * rho**(1.-gamma) * U
  
  
  
  
###################
def MeanWeight(Xi,ionized=0):
###################

  if  ionized: 
    return 4./(8.-5.*(1.-Xi))
  else:
    return 4./(1.+3.*Xi)


###################
def ElectronDensity(rho,pars):
###################
  '''
  Electron density for a mixture of H + He
  '''

  Xi		= pars['Xi']
  mh    	= pars['mh']
  ionisation	= pars['ionisation']

  if  ionisation: 
    return rho.astype(float)/mh*(Xi+(1-Xi)/2.)
  else:
    return rho.astype(float)*0


#############################
def Lambda(rho,u,localsystem,thermopars,coolingfile):
#############################
  '''
  This corresponds to Lambda normalized
  
  Ln = L / nh 2
  
  nh = (xi*rho/mh)
  '''

  UnitLength_in_cm = PhysCte(1,localsystem.UnitDic['m']).into(units.cgs)
  UnitTime_in_s    = PhysCte(1,localsystem.UnitDic['s']).into(units.cgs)
  UnitMass_in_g    = PhysCte(1,localsystem.UnitDic['kg']).into(units.cgs)
 
  UnitVelocity_in_cm_per_s = UnitLength_in_cm/UnitTime_in_s
  UnitEnergy_in_cgs = UnitMass_in_g * pow(UnitLength_in_cm, 2) / pow(UnitTime_in_s, 2)
  
  metalicity       = thermopars['metalicity']
  hubbleparam      = thermopars['hubbleparam']
    
  # compute cooling time
  logT,logL0,logL1,logL2,logL3,logL4,logL5,logL6 = io.read_cooling(coolingfile)
  
  if  metalicity==0:
    logL = logL0
  elif metalicity==1:  
    logL = logL1
  elif metalicity==2: 
    logL = logL2
  elif metalicity==3: 
    logL = logL3
  elif metalicity==4: 
    logL = logL4
  elif metalicity==5: 
    logL = logL5
  elif metalicity==6: 
    logL = logL6 
    
  # compute gas temp  
  logTm = log10(thermodyn.Tru(rho,u,thermopars))  
  c   = ((logTm>=4)*(logTm<8.5))
  u   = where(c,u,1)
  rho = where(c,rho,1)
  # recompute gas temp
  logTm = log10(thermodyn.Tru(rho,u,thermopars))
  
  # get the right L for a given mT
  logLm = take(logL,searchsorted(logT,logTm))
  Lm = 10**logLm
  
  # transform in user units
  Lm =  Lm / UnitEnergy_in_cgs /pow(UnitLength_in_cm,3) *  UnitTime_in_s
  
  L = Lm * hubbleparam
    
    
  return L


#############################
def CoolingTime(rho,u,localsystem,thermopars,coolingfile):
#############################

  UnitLength_in_cm = PhysCte(1,localsystem.UnitDic['m']).into(units.cgs)
  UnitTime_in_s    = PhysCte(1,localsystem.UnitDic['s']).into(units.cgs)
  UnitMass_in_g    = PhysCte(1,localsystem.UnitDic['kg']).into(units.cgs)
 
  UnitVelocity_in_cm_per_s = UnitLength_in_cm/UnitTime_in_s
  UnitEnergy_in_cgs = UnitMass_in_g * pow(UnitLength_in_cm, 2) / pow(UnitTime_in_s, 2)
  
  ProtonMass       = thermopars['mh']
  Xi               = thermopars['Xi']
  metalicity       = thermopars['metalicity']
  hubbleparam      = thermopars['hubbleparam']
    
  # compute cooling time
  logT,logL0,logL1,logL2,logL3,logL4,logL5,logL6 = io.read_cooling(coolingfile)
  
  if  metalicity==0:
    logL = logL0
  elif metalicity==1:  
    logL = logL1
  elif metalicity==2: 
    logL = logL2
  elif metalicity==3: 
    logL = logL3
  elif metalicity==4: 
    logL = logL4
  elif metalicity==5: 
    logL = logL5
  elif metalicity==6: 
    logL = logL6 
    
  # compute gas temp  
  logTm = log10(thermodyn.Tru(rho,u,thermopars))  
  c   = ((logTm>=4)*(logTm<8.5))
  u   = where(c,u,1)
  rho = where(c,rho,1)
  # recompute gas temp
  logTm = log10(thermodyn.Tru(rho,u,thermopars))
  
  # get the right L for a given mT
  logLm = take(logL,searchsorted(logT,logTm))
  Lm = 10**logLm
  
  # transform in user units
  Lm =  Lm / UnitEnergy_in_cgs /pow(UnitLength_in_cm,3) *  UnitTime_in_s
  
  L = Lm * hubbleparam
    
  tc = (ProtonMass)**2/(L*Xi**2) * u/rho
  tc = where(c,tc,0)
    
  return tc,c
