#This script computes carbonate/bicarbonate/CO2 equilibrium
#
#Note that the atmospheric fraction computed in the routines
#below is the ratio of atmospheric carbon to oceanic carbon.
#Total carbon would be the atmospheric carbon time (1 + 1/f)
#

from ClimateUtilities import *
import math
import planets

cfact = 1./(1000./18.) #Convert mole/liter to mole fraction
#Remember, the expression for H would need to be converted, too,
#because the equilibrium constant for water dissociation,
#10**-14, is given in mole/liter

#Compute Arrhenius law coefficients
def Tcoeffs(coeff30,coeff0):
    results = []
    for i in range(len(coeff30)):
        c30 = coeff30[i]
        c0 = coeff0[i]
        Tstar = math.log(c30/c0)/(1./303.15 - 1./273.15)
        results.append(Tstar)
    return results

#Freshwater values, from Wikipedia or Faure.
#coeffFresh_25C_2m = [.43e-6,.561e-10,.447e-8]
#Check: Is value of Ksp really that low for fresh water?

#Data from Archer's calculator. [K1,K2,Ksp]
#Units based on concentrations in moles/liter
coeff35_30C_2m = [1.0474e-6,.85430e-9,.47525e-6]
coeff35_0C_2m  = [.87140e-6,.42772e-9,.47528e-6]
#The low salinity values may be out of range of fit
coeff5_30C_2m =  [.57356e-6,.333851e-9,.47525e-6]
coeff5_0C_2m =   [.3375e-6,.14690e-10,.47528e-6] #Note strong T dep at low salinity!
#High pressure, standard salinity
coeff35_30C_3000m = [1.355e-6,1.023e-9,.78332e-6]
coeff35_0C_3000m = [.8714e-6,.4277e-9,.82759e-6]

#Calculate temperature coefficients
c1,c2,csp = Tcoeffs(coeff35_30C_3000m,coeff35_0C_3000m)
K1,K2,Ksp = coeff35_0C_3000m





K1 = K1*cfact
K2 = K2*cfact
Ksp = Ksp*cfact*cfact

#Henry's law constants for CO2
kh0 = 1.6e8
ch = 2400.


#logKsp = math.log10(4.75e-7)
#These are the values at 0C
#From Higgins and Schrag. In moles/liter. (convert to mole fraction later)
#Values at 0C
#K1 = 6.33e-7
#K2 = 3.53e-10
#Kh = 15.89 #p in bar, concentration in mole/l
#KCO2 = 1./Kh
#Values at 30C#
#logK1 = math.log10(10.47e-7)
#logK2 = math.log10(8.77e-10)
#KCO2 = 2.52e-2 #From Higgins and Schrag; p in atm.


#Set the temperature
T = 275.

#Compute the temperature-scaled constants
kh = kh0*math.exp(-ch*(1./T - 1/298.15))
K1 = K1*math.exp(c1*(1./T-1./273.15))
K2 = K2*math.exp(c2*(1./T-1./273.15))
Ksp = Ksp*math.exp(csp*(1./T-1./273.15))


KCO2 = 1./kh
logK1 = math.log10(K1)
logK2 = math.log10(K2)
logKsp = math.log10(Ksp)

#pCO2 in Pascals
def carb(pH,pCO2):
    CO2 = pCO2*KCO2 #Henry's law constant
    
    logHCO3 = math.log10(CO2) + logK1 + (pH - math.log10(cfact))
    logCO3 = logHCO3 + (pH - math.log10(cfact)) + logK2
    return CO2, 10**logHCO3,10**logCO3


#Make graph of species fractions vs. pH
#Note that this graph is independent of pCO2,
#since all concentrations scale with pCO2
def PartitionGraph():
    pHlist = [1.+.1*i for i in range(130)]
    pCO2 = 1.e-4 #Pick any value. It doesn't matter
    CO2L =[]
    HCO3L=[]
    CO3L = []
    for pH in pHlist:
        CO2,HCO3,CO3 = carb(pH,pCO2)
        totCarb = CO2+HCO3+CO3
        CO2L.append(CO2/totCarb)
        HCO3L.append(HCO3/totCarb)
        CO3L.append(CO3/totCarb)
    c = Curve()
    c.YlogAxis = True
    c.addCurve(pHlist,'pH')
    c.addCurve(CO2L,'CO2')
    c.addCurve(HCO3L,'HCO3')
    c.addCurve(CO3L,'CO3')
    return c

#Make graph of ocean carbon storage vs. pH
#(What would this look like for a pure CO2 atmosphere?)
OceanMoles = 7.8e19
#AtmMoles = 1.8e17 #The relation of this to surface pressure depends on g
AEarth = 4.*math.pi*planets.Earth.a**2 #Surface area of Earth
ratEarth = AEarth*kh/(29.*9.8*OceanMoles)
def StorageGraph(rat):
    pHlist = [5.+.1*i for i in range(41)]
    fracList = []
    for pH in pHlist:
        fracList.append(AtmosFrac(rat,pH))
    c = Curve()
    c.YlogAxis = True
    c.addCurve(pHlist,'pH')
    c.addCurve(fracList,'atmFrac')
    return c

def AtmosFrac(rat,pH):
    H = cfact*10.**(-pH) #Convert to mole fraction
    return rat/(1.+K1/H + K1*K2/H**2)

#**ToDo: Incorporate temperature dependence of water
#dissociation constant here. (deviates from -14).
def ChargeBalance(pH,pCO2,pAlk = None):
    CO2,HCO3,CO3 = carb(pH,pCO2)
    #Determine pAlk in saturation instead of as a parameter
    if pAlk == None:
        pAlk = -(logKsp-math.log10(CO3))
    balance =  2.*10.**-pAlk + cfact*10.**-pH - (HCO3 + 2.*CO3 + cfact*10.**(-14+pH))
    return balance



#Calculate pH by iterating. Case for carbonate in saturation
def f(pH,pCO2):
    return ChargeBalance(pH,pCO2)
m = newtSolve(f)
def pH(pCO2):
    m.setParams(pCO2)
    return m(7.)

#Calculate pH by iterating. Case for fixed alkalinity
def f1(pH,params):
    pCO2 = params[0]
    pAlk = params[1]
    return ChargeBalance(pH,pCO2,pAlk)
m1 = newtSolve(f1)
def phFixedAlk(pCO2,pAlk):
    m1.setParams([pCO2,pAlk])
    return m1(7.)

#Computes log10  alkalinity for a specified pH and pCO2.
#Useful for global-warming type calculations where we keep
#alkalinity fixed.
def pAlk(pCO2,pH):
    CO2,HCO3,CO3 = carb(pH,pCO2)
    return -math.log10(.5*(-cfact*10.**-pH + (HCO3 + 2.*CO3 + cfact*10.**(-14+pH))))
#The following total carbon calcs are for Earth parameters. Just edit
#the AtmosCarb= line, to do it for a different planet.
#Note that these functions return the number of Moles of total
#carbon. Multiply by 12 to get Kg. (This was done in the discussion
#of air fraction vs. cumulative carbon in the text)
#
def TotCarbFixedAlk(pCO2,pAlk):
    AtmosCarb = AEarth*pCO2/(29.*9.8)
    fa = AtmosFrac(ratEarth,phFixedAlk(pCO2,pAlk))
    return AtmosCarb*(1. + 1/fa)
def TotCarbBuffered(pCO2):
    AtmosCarb = AEarth*pCO2/(29.*9.8)
    fa = AtmosFrac(ratEarth,pH(pCO2))
    return AtmosCarb*(1. + 1/fa)



#Atmos fraction graph vs pCO2 for saturated carbonate
#In this case, since we typically want to go to high CO2
#where the CO2 changes the mean molecular weight of the
#atmosphere, we compute ratEarth taking this into account
#pAir0 is the partial pressure the atmosphere would have
#in the absence of CO2 -- NOT the partial pressure of air
#in the mixture!
def StorageGraphBuffered(pAir0):
    #The pressure of CO2 in isolation.
    #We will compute the pCO2 from this
    pCO2_0list = [10.**(.1*i) for i in range(51)]
    pCO2list = []
    fracList = []
    pHList = []
    for pCO2_0 in pCO2_0list:
        #Determine pCO2 (partial pressure in the mixture)
        Ma = 29.
        MCO2 = 44.
        pCO2 = pCO2_0*(1.+(pAir0/pCO2_0))/(1.+(pAir0/pCO2_0)*(MCO2/Ma))
        molefrac = pCO2/(pAir0+pCO2_0)
        Mbar = molefrac*MCO2 + (1.-molefrac)*Ma
        #Edit OceanMoles, g and AEarth for other planets. 
        g = 9.8 #Acceleration of gravity
        rat = AEarth*kh/(Mbar*g*OceanMoles)
        #
        pCO2list.append(pCO2)
        fracList.append(AtmosFrac(rat,pH(pCO2)))
        pHList.append(pH(pCO2))
    #
    c = Curve()
    c.YlogAxis = True
    c.XlogAxis = True
    c.addCurve(pCO2list,'pCO2')
    c.addCurve(fracList,'atmFrac')
    #
    c1 = Curve()
    c1.XlogAxis = True
    c1.addCurve(pCO2list,'pCO2')
    c1.addCurve(pHList,'pH')
    #
    c2 = Curve()
    c2.XlogAxis = True
    c2.YlogAxis = True
    c2.addCurve(pCO2list,'pCO2')
    c2.addCurve(pCO2_0list,'pCO2Alone')
    return c,c1,c2

#Atmos fraction graph vs pCO2 for fixed alkalinity
def StorageGraphFixedAlk(rat,pAlk):
    pCO2list = [10.**(.1*i) for i in range(10,41)]
    fracList = []
    for pCO2 in pCO2list:
        fracList.append(AtmosFrac(rat,phFixedAlk(pCO2,pAlk)))
    c = Curve()
    c.YlogAxis = True
    c.XlogAxis = True
    c.addCurve(pCO2list,'pCO2')
    c.addCurve(fracList,'atmFrac')
    return c