#---------------------Description--------------------------------------
#
#      Script to compute molar concentration at the cold trap
#
#
#-----------------------------------------------------------------------
#ToDo:

#

#Data on section of text which this script is associated with
Chapter = '8'
Figure = ''
Table = ''
#

from ClimateUtilities import *

import planets
import phys
import math

#Set the atmospheric composition
condensible = phys.water
noncondensible = phys.CO2


#Set up the needed thermodynamic functions
m = phys.MoistAdiabat(condensible,noncondensible)
m.ptop = .01
es = phys.satvps_function(condensible)

#Newton solver to get the lifted condensation level
def f(p,params):
    eta = params[0]
    psTot = params[1]
    Ts = params[2]
    Rcp = params[3]
    psCon = psTot*eta
    psNonCon = psTot*(1.-eta)
    Tdry = Ts*(p/psTot)**Rcp
    return eta*p - es(Tdry)
mLift= newtSolve(f)

#psNonCon is the surface partial pressure of the noncondensible
#gas.  Tcold is the temperature of the cold trap (the tropopause
#temperature).  Ts is the surface temperature.
#rhSurf is the surface relative humidity of the condensible (as a fraction)
#
def etaColdTrap(psNonCon,rhSurf,Ts=540.,Tcold=200.):
    psCon = rhSurf*es(Ts)
    eta = psCon/(psCon+psNonCon) #(constant) molar concentration below
                                 # the condensing level
    #Compute R/cp for the unsaturated mixture in the lower atmosphere
    Mc = condensible.MolecularWeight
    Mnc = noncondensible.MolecularWeight
    massCon = eta*Mc/((1.-eta)*Mnc + eta*Mc)
    cp = massCon*condensible.cp + (1.-massCon)*noncondensible.cp
    M = 1./(massCon/Mc + (1.-massCon)/Mnc)
    R = phys.Rstar/M
    Rcp = R/cp
    #Compute the lifted condensation level (pressure)
    #Check the computation of the lifted condensation level
    mLift.setParams([eta,psCon+psNonCon,Ts,Rcp])
    pcond = mLift(psCon + psNonCon)
    Tcond = Ts*(pcond/(psCon+psNonCon))**Rcp
    #print pcond/1.e5,(psCon+psNonCon)/1.e5
    #
    #Compute moist adiabat starting from condensation level
    #Remember: pressure argument is the noncondensible partial pressure
    p,Tad,MolarCon,MassCon = m(pcond*(1.-eta),Tcond)
    #Find index of cold trap
    i = Numeric.searchsorted(-Tad,-Tcold)
    #Compute the mass of the condensible
    #First, mass of the noncondensing part below the condensation level
    #(Mass expressed as partial pressure of pure condensible)
    pLower = (psCon/(psCon+psNonCon))*(psNonCon+psCon-pcond)
    #Now sum up mass in the saturated layer (up to the cold trap
    pUpper = 0.
    for j in range(i-1):
        pUpper += .5*(MolarCon[j]+MolarCon[j+1])*(p[j]-p[j+1])
    #print "Condensible pressure:",pUpper/1.e5,pLower/1.e5,(pUpper+pLower)/1.e5
    #(Ignore mass above the cold trap)
    #**ToDo:  put in check to see if Tcold is colder than min
    #return TsSat,pcond,p[i]/1.e5,MolarCon[i]
    #print "pcond =",pcond/1.e5,p[0]/1.e5,"pUpper =",pUpper/1.e5,"pLower =",pLower/1.e5
    return MolarCon[i],(pUpper+pLower)/1.e5