import phys,planets
from ClimateUtilities import *
import math

def ProbA():
    tag ='Workbook:thermo:TitanMethaneSatvp'
#Make a methane saturation vapor pressure function
#Note that this script computes a few additional things
#that weren't asked for in the problem, such as the
#molar concentration
    gas = phys.CH4
    psat = phys.satvps_function(gas.TriplePointT,gas.TriplePointP,gas.MolecularWeight,gas.L_vaporization)
    print 'psat = %e Pa at 95K and %e Pa at 120K'%(psat(95.),psat(120.))
    #Compute mixing ratios and concentrations:"
    pN2 = 1.5e5
    molarMix95 = psat(95.)/pN2
    molarMix120 = psat(120.)/pN2
    MCH4 = phys.CH4.MolecularWeight
    MN2 = phys.N2.MolecularWeight
    massMix95 = (MCH4/MN2)*molarMix95
    massMix120 = (MCH4/MN2)*molarMix120
    molarCon95 = psat(95.)/(pN2 + psat(95.))
    molarCon120 = psat(120.)/(pN2 + psat(120.))
    massSpCon95 = (psat(95.)*MCH4)/((pN2*MN2) + (psat(95.)*MCH4))
    massSpCon120 = (psat(120.)*MCH4)/((pN2*MN2) + (psat(120.)*MCH4))
    print "molar mixing ratios %f at 95K,%f at 120K"%(molarMix95,molarMix120)
    print "mass mixing ratios %f at 95K, %f at 120K"%(massMix95,massMix120)
    print "molar concentrations %f at 95K, %f at 120K"%(molarCon95,molarCon120)
    print "mass specific concentration %f at 95K %f at 120K"%(massSpCon95,massSpCon120)

def ProbB():
    tag = 'Workbook:thermo:SpringtimeEuropa'
    print 'Springtime for Europa-----------------------'
    T0 = phys.water.TriplePointT
    p0 = phys.water.TriplePointP
    M = phys.water.MolecularWeight
    L = phys.water.L_vaporization
    esat = phys.satvps_function(T0,p0,M,L)
    R = phys.water.R
    print "Surface pressure is %fmb"%(esat(280.)/100.)
    print "Scale height of atmosphere is %f km"%((R*280./1.3)/1000.)

    #Set up function for one-component moist adiabat
    def Tsat(p):
        return(1./(1./Tref - (R/L)*math.log(p/psat_ref)))

    #Set constants and plot
    Tref = 280.
    psat_ref = esat(280.)

    plist = [psat_ref - i*psat_ref/100 for i in range(1,100)]
    #Note the following trick using "list comprehension" to
    #generate the list of temperatures.
    Tlist = [Tsat(p) for p in plist]

    c = Curve()
    c.addCurve(plist)
    c.addCurve(Tlist)
    c.switchXY = c.reverseY=c.YlogAxis = True
    #Titles and axis labels
    c.PlotTitle = 'Springtime for Europa!'
    c.Xlabel = 'Pressure (Pa)'
    c.Ylabel = 'Temperature (K)'
    #Do the plot
    w = plot(c)
    w.save('SpringtimeEuropa')

def ProbC():
    tag = 'Workbook:thermo:ShomateCpEqn'
    #Define the Shomate function
    def cp(T):
        TT = T/1000.
        A,B,C,D,E = coeffs
        return A + B*TT + C*TT*TT + D*TT*TT*TT + E/(TT*TT)

    #Define coefficient arrays
    ShomateN2 = [931.857, 293.529 , -70.576, 5.688 , 1.587]
    ShomateCO2 = [568.122,1254.249,-765.713, 180.645, -3.105]
    ShomateNH3 = [1176.213, 2927.717, -904.470, 113.009, 11.128]

    #Now print out some values at 25C and compare with
    #the values in the table in the text (which are also
    #the values in the phys module

    coeffs = ShomateN2
    gas = phys.N2
    print gas.formula,gas.cp,cp(273.15+25.)
    coeffs = ShomateCO2
    gas = phys.CO2
    print gas.formula,gas.cp,cp(273.15+25.)
    coeffs = ShomateNH3
    gas = phys.NH3
    print gas.formula,gas.cp,cp(273.15+25.)
    #
    #**ToDo: Advanced tip-- show how to do this by
    #        looping over the coeffs

    #Now make a plot
    Tlist = [300. + i for i in range(700)] #Define temperature list
    c = Curve()
    c.addCurve(Tlist,'T','Temperature (K)')
    #
    #Now loop over coefficients, compute list of cp, and put in the curve
    #**ToDo: Find a more elegant index-free way to do this loop
    for i in range(3):
        if i==0:
            coeffs = ShomateN2
            name = phys.N2.formula
        elif i ==1:
            coeffs = ShomateCO2
            name = phys.CO2.formula
        else:
            coeffs = ShomateNH3
            name = phys.NH3.formula
        cpList = [cp(TT) for TT in Tlist]
        c.addCurve(cpList,name)
    #
    #Plot and axis titles
    c.PlotTitle = 'Cp from Shomate'
    c.Xlabel = 'Temperature (K)'
    c.Ylabel = 'Specific Heat (J/(kg K))'
    w = plot(c) #Do the plot
    w.save('ShomateCp')
    #
    #Now compute the energy change between 300K and 700K
    E = 0. #Initialize variables for accumulating sums
    #Try doing this for NH3 by changing the following line
    coeffs = ShomateN2
    dT = .1
    Tlist = [300. + dT*i for i in range(4000)]
    for T in Tlist:
        E = E + cp(T)*dT #This is implicitly multiplied by 1 kg
    Tmid = (Tlist[0]+Tlist[-1])/2.  # Index [-1] is shorthand for the last value
    print "Energy change = ",E,'J, vs ',cp(Tmid)*400.,'for midpoint'

def ProbD():
    tag = 'Workbook:thermo:DryAdiabatTdepCp'
    ShomateCO2 = [568.122,1254.249,-765.713, 180.645, -3.105]
    def cp(T):
        TT = T/1000.
        A,B,C,D,E = ShomateCO2
        return A + B*TT + C*TT*TT + D*TT*TT*TT + E/(TT*TT)
    R = phys.CO2.R
    #Define the derivative function. In this case
    #we don't need the parameter argument, so we leave
    #it out (allowed by newer versions of ClimateUtilities)
    #We still need to include logp, even though it's not
    #used, though. Without this argument, the integrator
    #has no way of knowing that logT is the dependent
    #variable
    def dlogTdlogp(logp,logT):
        return R/cp(math.exp(logT))

    #Initial conditions. We start at the ground
    #and integrate upwards (i.e in the negative p direction)
    Tstart = 737.  #Venus ground temperature 
    pstart = 92.e5  #Venus ground pressure
    ptop =  .1e5 #Where to quit

    logTstart = math.log(Tstart)
    logpstart = math.log(pstart)

    #
    #Set the increment. Smaller is more accurate
    dlogp = -math.log(pstart/ptop)/100.
    #This is set so the calculation ends exactly at ptop
    #
    logp = logpstart
    logT = logTstart
    m = integrator(dlogTdlogp,logpstart,logTstart,dlogp)
    #
    #Loop to carry out the integration and save results
    pList = [math.exp(logp)]
    TList = [math.exp(logT)]
    while logp > math.log(ptop):
        logp,logT = m.next() #Compute next value
        pList.append(math.exp(logp))
        TList.append(math.exp(logT))

    print 'For Venus, the dry adiabatic temperature at %f bar is %f K'%(pList[-1]/1.e5,TList[-1])
    print 'For constant cp, temperature would be %f K'%(Tstart*(pList[-1]/pstart)**(R/phys.CO2.cp))

    #Make a plot of the temperature profile
    c2 = Curve()
    c2.addCurve(pList,'p')
    c2.addCurve(TList,'T')
    c2['Tconst'] = Tstart*(c2['p']/pstart)**(R/phys.CO2.cp) # constant cp case
    c2.YlogAxis = c2.reverseY = c2.switchXY = True
    c2.Xlabel = 'Pressure'
    c2.Ylabel = 'Temperature'
    w2 = plot(c2)
    w2.save('VenusTdepCp')

def ProbE():
    tag = 'Workbook:thermo:EarlyEarthSteamAtm'
    #Exponential saturation v.p. based on triple point
    esTP = phys.satvps_function(phys.water)
    
    #Exponential saturation v.p. based on boiling point
    #  In this case we can't just use the gas object, but
    #  have to pass the data explicitly in the argument list
    M = phys.water.MolecularWeight
    L = phys.water.L_vaporization
    esBP = phys.satvps_function(373.15,1.e5,M,L)

    #Antoine equation saturation v.p.
    def esAntoine(T):
        A,B,C = (3.55959,643.748,-198.043)
        return 1.e5*10.**(A-B/(T+C)) #Convert to Pascals
    #
    #Now print out vapor pressures and molar concentrations
    #ToDo: Pretty up alignment of table header
    print 'T','esTP','conTP','esBP','conBP','esAntoine','conAntoine'
    for T in [373.15,473.15]:
        conTP = esTP(T)/(esTP(T) + 1.e5)
        conBP = esBP(T)/(esBP(T)+1.e5)
        conAntoine = esAntoine(T)/(esAntoine(T)+1.e5)
        #The things like '%.2e'% convert the numbers to a nice output format
        print T,'%.2e'%esTP(T),'%.2f'%conTP,'%.2e'% esBP(T),'%.2f'%conBP,'%.2e'%esAntoine(T),'%.2f'%conAntoine

    #Check of the Antoine equation at critical point
    print
    print 'At critical point T, Antoine gives',esAntoine(phys.water.CriticalPointT)/1.e5,'bar'
    print 'vs correct answer',phys.water.CriticalPointP/1.e5,'bar'
    Mocean = 1.4e21
    Aearth = 4.*math.pi*planets.Earth.a**2
    Mvapor = Aearth*esAntoine(473.15)/planets.Earth.g
    print
    print "Based on Antoine equation:"
    print "     At 200C",Mvapor/Mocean, "of total water is in the atmosphere"
    #
    #Now find when atmospheric water exceeds ocean water. Do this
    #by just printing out a list of values. (A smarter way would be
    #to use Newton's method)
    print
    print 'Water left in ocean vs. T, for t.p, b.p. and Antoine vapor pressure'
    for T in [400. + 5.*i for i in range(100)]:
        Mvapor1 = Aearth*esTP(T)/planets.Earth.g
        Mvapor2 = Aearth*esBP(T)/planets.Earth.g
        Mvapor3 = Aearth*esAntoine(T)/planets.Earth.g
        print T,Mocean-Mvapor1,Mocean-Mvapor2,Mocean-Mvapor3
    
def ProbF():
    tag = 'Workbook:thermo:LfromClausius'
    print '---Latent Heat from Clausius Clapeyron---'
    #Define the latent heat function
    #The gas constant R is a global
    def Lcomp(T,e,dT):
        deriv = (math.log(e(T+dT)) - math.log(e(T-dT)))/(2.*dT)
        return R*T**2 * deriv

    #Now define a test saturation vapor pressure, define
    #the other globals, and try out L(T)
    gas = phys.water
    e1 = phys.satvps_function(gas.TriplePointT,gas.TriplePointP,gas.MolecularWeight,gas.L_vaporization)
    R = gas.R

    #This loop checks the computed value against the correct value for various dT
    for dT in [1.e-30,1.e-6,1.e-3,1.,10.,100.]:
        print dT,Lcomp(300.,e1,dT),gas.L_vaporization


    #Now do it for the empirical function, satvpg.

    #First try some checks to see the sensitivity to dT
    print 'Convergence checks for satvpg(T)'
    for dT in [100.,10.,1.,1.e-3,1.e-6,1.e-9]:
        print dT,Lcomp(300.,phys.satvpg,dT)

    #Looks like dT = 1.e-6 is OK
    #
    #Make a table for plotting
    TT = [200.+ 2.*i for i in range(100)]
    Lsatvpg = [Lcomp(T,phys.satvpg,1.e-6) for T in TT]
    Lsatvpw = [Lcomp(T,phys.satvpw,1.e-6) for T in TT]
    Lsatvpi = [Lcomp(T,phys.satvpi,1.e-6) for T in TT]
    c1 = Curve()
    c1.addCurve(TT,'T')
    c1.addCurve(Lsatvpg,'Lsatvpg')
    #Comment out the next two statments
    #if you want to see a plot for just satvpg,
    #without the distraction of the water and ice values
    #that I added for comparison
    c1.addCurve(Lsatvpw,'Lsatvpw')
    c1.addCurve(Lsatvpi,'Lsatvpi')
    #
    c1.Xlabel = 'Temperature'
    c1.Ylabel = 'Computed Latent Heat'
    c1.PlotTitle = 'L computed from empirical functions'
    w = plot(c1)
    w.save('LfromClaus')
    #
    #On the graph, you can hardly see the computed L for
    #phys.satvpi is varying at all at low values.  Try executing
    #Lcomp(T,phys.satvpi,1.e-6) for T = 100, 200 and 273. to
    #see how much the latent heat is actually varying

def ProbG():
    tag = 'Workbook:thermo:MoistAdiabatSlope'
    pa = 1.e5
    Ra = phys.air.R
    L = phys.water.L_vaporization
    es = phys.satvps_function(phys.water.TriplePointT,phys.water.TriplePointP,phys.water.MolecularWeight,L)
    def slope(T):
        qsat = (phys.water.MolecularWeight/phys.air.MolecularWeight)*es(T)/pa
        A = (L/(Ra*T))*qsat
        B1 = phys.water.cp/phys.air.cp
        B2 = L/(phys.water.R*T)
        B3 = L/(phys.air.cp*T)
        B = (B1 + (B2-1.)*B3)
        Num = 1. + A
        Den = 1. + B*qsat
        drySlope = Ra/phys.air.cp
        print "T = %f and qsat = %f"%(T,qsat)
        print "Numerator (1+A), A=",A
        print "Denominator 1+ (B1 +(B2-1)*B3)*qsat"
        print "B1,B2,B3 = ",B1,B2,B3
        print "Denominator = ",1.+B*qsat,'with B=',B
        print "Slope = %f,dry adiabat slope = %f"%(drySlope*Num/Den,drySlope)

    print "Moist Adiabat Slope Results"
    slope(250.)
    slope(300.)
    
def ProbH():
    tag = 'Workbook:thermo:MoistAdiabatTitan'
    m = phys.MoistAdiabat(phys.CH4,phys.N2)
    #This function uses the MoistAdiabat object to
    #make the plots of the same type as output by
    #the chapter script MoistAdiabat.py. The argument
    #is the surface temperature you want; it assumes
    #an N2 surface partial pressure of 1.5 bars
    Rc = phys.CH4.R
    L = phys.CH4.L_vaporization
    Tr = phys.CH4.TriplePointT
    psat_ref = phys.CH4.TriplePointP
    Ra = phys.N2.R
    cpa = phys.N2.cp
    def Tsat(p):
        return(1./(1./Tr - (Rc/L)*Numeric.log(p/psat_ref)))
    def Tdry(Ts,p):
        return Ts*(p/p[0])**(Ra/cpa) #Note p is a Numeric array
    def results(Ts):
        p,T,qMolar,qMass = m(1.e5,Ts)
        #Temperature plot
        c1 = Curve()
        c1.addCurve(p,'p')
        c1.addCurve(T,'MoistAdiabat')
        c1.addCurve([Tsat(pp) for pp in p],'Tsat','One Component Saturated')
        c1.addCurve(Tdry(Ts,p),'Tdry','Dry adiabat')
        c1.switchXY = c1.YlogAxis = c1.reverseY = True
        c1.PlotTitle = 'Titan Adiabat for Ts = %4.0f K'%Ts
        #Molar concentration plot
        c2 = Curve()
        c2.addCurve(p,'p')
        c2.addCurve(qMolar,'qMolar','Molar Concentration')
        c2.switchXY = c2.YlogAxis = c2.reverseY = c2.XlogAxis = True
        c2.PlotTitle = 'Titan CH4 molar concentration for Ts = %4.0f K'%Ts
        return c1,c2
    for Ts in [80.,95.,120.]:
        c1,c2 = results(Ts)
        w1 =plot(c1)
        w2 =plot(c2)
        #Save the plots here, if you want them.
        #The formatting character %.0f makes sure there
        #are no spaces or unwanted decimal places in the filename
        #w1.save('TitanT%.0fK'%Ts)
        #w2.save('Titanq%.0fK'%Ts)
    #
    #**ToDo: Illustrate how to use the interpolation option to put
    #  all the temperature profiles on the same graph.
    #  Also do same for concentration profiles
    #
    #**ToDo: Do a comparison with Huyghens data
        

def ProbI():
    tag = 'Workbook:thermo:MoistAdiabatAirWaterMass'
    m = phys.MoistAdiabat() #Default is air/water
    #
    #Write a simple integration loop to find water mass
    #(per unit area)for any given surface air partial pressure and
    #and any given surface temperature.
    #The air mass is the difference between total mass and
    #water mass. The total mass is psurf/g, by the hydrostatic
    #relation
    def Masses(pa,Ts):
        p,T,qMolar,qMass = m(pa,Ts)
        WaterMass = 0.
        n = len(p)
        for i in range(n-1):
            dp = -(p[i+1]-p[i])
            qbar = .5*(qMass[i+1]+qMass[i])
            WaterMass = WaterMass + qbar*dp/planets.Earth.g
        return WaterMass,p[0]/planets.Earth.g - WaterMass
    #
    #Print out a table
    print 'Water and air mass paths, in kg/m**2'
    print 'Ts  WaterMass  pw(Ts)/g   AirMass   pa/g'
    psatw = phys.satvps_function(phys.water)
    for T in [270.+5.*i for i in range(27)]:
        Mw1 = psatw(T)/planets.Earth.g
        Ma1 = 1.e5/planets.Earth.g
        WaterMass,AirMass = Masses(1.e5,T)
        print T,'%0.2f'%WaterMass,'%0.2f'%Mw1,'%0.2f'%AirMass,'%0.2f'%Ma1