import numpy as np
def K0_Weiss(S, TC):
    TK=TC+273.15 #Temperatura de Celsius a Kelvin
    lnK0 = 9345.17/TK - 60.2409 + 23.3585 * np.log(TK/100) + S * (0.023517 - 0.00023656 * TK + 4.7036e-07 *
        TK * TK)
    K0=np.exp(lnK0)
    return K0


"""
Created on Mon Sep  3 11:51:00 2018

@author: fergomez
Esta funcion calcula K1 y K2 del sistema CO2-H20 segun
Roy et al., (1993), The dissociation constant of carbonic acid in sewater
at salinities of 5 to 45 and temperatures of 0 to 45 desgrees Celsius.
Marine Chemistry 44(2-4), 249-267.
Unidades
K1, K2=mol*kg-soln-1
Ejemplo de uso:
Input
K1, K2=K1_K2_Roy(35, 25)
print(K1, K2)
Output:
1.3921075396202872e-06 1.1887254858040348e-09
"""


def K1_K2_Roy(S, TC):
    TK=TC+273.15 #Temperatura de Celsius a Kelvin
    #K1 usando Roy et al., 1993
    tmp1 = 2.83655 - 2307.1266/TK - 1.5529413 * np.log(TK)
    tmp2 = -(0.207608410 + 4.0484/TK) * np.sqrt(S)
    tmp3 = 0.08468345 * S - 0.00654208 * S**1.5
    tmp4 = np.log(1 - 0.001005 * S)
    lnK1roy = tmp1 + tmp2 + tmp3 + tmp4
    K1 = np.exp(lnK1roy)
    #K2 usando Roy et al., 1993
    tmp1 = -9.226508 - 3351.6106/TK - 0.2005743 * np.log(TK)
    tmp2 = (-0.106901773 - 23.9722/TK) * np.sqrt(S)
    tmp3 = 0.1130822 * S - 0.00846934 * S**1.5 + np.log(1 -
        0.001005 * S)
    lnK2roy = tmp1 + tmp2 + tmp3
    K2 = np.exp(lnK2roy)
    return K1, K2

def K1_K2_SBY(S):
    pK1=6.1568-0.00352*S
    pK2=8.5503-0.0080*S
    K1=10**(-pK1)
    K2=10**(-pK2)
    return K1, K2
"""
Created on Mon Sep  3 11:53:53 2018

@author: fergomez
Esta funcion calcula K0 del sistema CO2-H20 segun recomienda Dickson et al 2007
Guide for best practices for ocean CO2 measurements. PICES SP 3, 191 pp.
Conforme a Mucci, A., 1983. The solubility of calcite and aragonite in seawater
at various salinities and temperatures and one atmosphere of total pressure.
American journal of Science, 283, 780-799.

Ejemplo de uso:
Input
Ksp_a=Kspa_Mucci(35, 25)
print(Ksp_a)
Output:
6.48175906801198e-07
"""

def Kspa_Mucci(S, TC):
    TK=TC+273.15 #Temperatura de Celsius a Kelvin
    tmp1 = -171.945 - 0.077993 * TK + 2903.293/TK + 71.595 * np.log10(TK)
    tmp2 = +(-0.068393 + 0.0017276 * TK + 88.135/TK) * np.sqrt(S)
    tmp3 = -0.10018 * S + 0.0059415 * S**1.5
    log10Kspa = tmp1 + tmp2 + tmp3
    Ksp_a=10**(log10Kspa)
    return Ksp_a


"""
Created on Mon Sep  3 11:54:42 2018

@author: fergomez

Esta funcion calcula K0 del sistema CO2-H20 segun recomienda Dickson et al 2007
Guide for best practices for ocean CO2 measurements. PICES SP 3, 191 pp.
Conforme a Mucci, A., 1983. The solubility of calcite and aragonite in seawater
at various salinities and temperatures and one atmosphere of total pressure.
American journal of Science, 283, 780-799.

Ejemplo de uso:
Input
Ksp_c=Kspc_Mucci(35, 25)
print(Ksp_c)
Output:
4.2723509278626e-07
"""

def Kspc_Mucci(S, TC):
    TK=TC+273.15 #Temperatura de Celsius a Kelvin
    tmp1 = -171.9065 - 0.077993 * TK + 2839.319/TK + 71.595 * np.log10(TK)
    tmp2 = +(-0.77712 + 0.0028426 * TK + 178.34/TK) * np.sqrt(S)
    tmp3 = -0.07711 * S + 0.0041249 * S**1.5
    log10Kspc = tmp1 + tmp2 + tmp3
    Ksp_c=10**(log10Kspc)
    return Ksp_c

# -*- coding: utf-8 -*-
"""
Created on Mon Sep  3 11:55:56 2018
Calcula el factor de correccion de la presion para las constantes de
equilibrio Ki usando Millero, F.J., (1995). Thermodynamics of the carbon
dioxide system in the oceans. Geochimica et Cosmochimica Acta 59:661-677.
@author: fergomez

Ejemplo de uso:
A presion atmosferica Kspa=4.27 e-07
Input (para aragonita):
fcorr=P_corr_Millero(-45.96, 0.5304, 0.0, -11.76e-03, 0.3692e-03, 0.0, 25, 300)
print(fcorr)
Output:
1.4787577790538065

entonces Kspa(presion atmosferica)*fcorr=Kspa_corregido
4.27 e-07*1.4787577790538065= 9.58 e-07
"""


#pagina 186 Andreas Hoffman PhD Thesis
def P_corr_Millero(a0, a1, a2, b0, b1, b2, TC, P):
    TK=TC+273.15 #Temperatura de Celsius a Kelvin
    R=83.131
    delta_V= a0+a1*TC+a2*TC*TC
    delta_k= (b0+b1*TC+b2*TC*TC)
    lnKp=-(delta_V/(R * TK))*P+0.5*(delta_k/(R*TK))*P*P
    fcorr=np.exp(lnKp)
    return fcorr

#Llamamos las librerias que utilizaremos
import numpy as np
#import matplotlib.pyplot as plt
import math

from numpy import exp,arange
from pylab import meshgrid,cm,imshow,contour, clabel,colorbar,axis,title,show, contourf

#from constantes import K0_Weiss, K1_K2_Roy, Kspc_Mucci, Kspa_Mucci, P_corr_Millero
import matplotlib.pyplot as plt

#============================================================================
#Inicio de seccion de INPUT==================================================
#============================================================================

S=6.88 #ppt
TC=6.42 #Celsius
P=1.0 #atmosferas



K0=K0_Weiss(S, TC)
K1, K2=K1_K2_Roy(S, TC)
Ksp_c=Kspc_Mucci(S, TC)
Ksp_a=Kspa_Mucci(S, TC)

#============================================================================
#Fin de seccion de INPUT======================================================
#============================================================================

#=============================================================================
#Estimacion de Calcio de Tyrrell et al., 2008 y valores de ejemplo agua de mar
#============================================================================
#Ca=0.0103 #valores agua de mar
Ca=(0.331*S+0.392)*1e-03 #Mar Baltico
#Ca=(0.375*S+0.0368)*1e-03 # Bothnian Bay (parte norte del mar Baltico)
#=============================================================================

#def co2_t(pco2, K0):
#    co2=pco2*K0
#    return co2
#=============================================================================
#vco2_t = np.vectorize(co2_t)
#Pasamos de pco2 a co2
#co2=co2_t(pco2, K0)


#=============================================================================
#=============================================================================
#Calculo del factor de correccion por presion
#=============================================================================
fcorr_K1=P_corr_Millero(-25.50, 0.1271, 0.0, -3.08e-03, 0.0877e-03, 0.0, TC, P)
fcorr_K2=P_corr_Millero(-15.82, -0.219, 0.0, 1.136e-03, -0.1475e-03, 0.0, TC, P)
fcorr_Kspc=P_corr_Millero(-48.76, 0.5304, 0.0, -11.76e-03, 0.3692e-03, 0.0, TC, P)
fcorr_Kspa=P_corr_Millero(-45.96, 0.5304, 0.0, -11.76e-03, 0.3692e-03, 0.0, TC, P)
#=============================================================================

#=============================================================================
#Correccion de ctes de equilibrio por presion
#=============================================================================
K1=K1*fcorr_K1
K2=K2*fcorr_K2
Ksp_c=Ksp_c*fcorr_Kspc
Ksp_a=Ksp_a*fcorr_Kspa

#=============================================================================
#Funcion que realiza calculos del sistema carbonato y omega
#=============================================================================

def carbEq(co2, dic, Ca):
    #-----------------------------------------------
    # Resolvemos para obtener H+ (cf. a Zeebe and Wolf-Gladrow, 2000)
    a1=co2-dic
    a2=K1*co2
    a3=K1*K2*co2

    p= [a1, a2, a3]
    r = np.roots(p)
    h= max(np.real(r)) # Esto es para seleccionar la raiz real mas grande
    #
    # Calculamos HCO3, CO3 and CO2aq, usando DIC, AlK y H+
    hco3=dic/(1+h/K1+K2/h)
    co3=dic/(1+h/K2+h*h/(K1*K2))
    #co2=dic/(1+K1/h+K1*K2/(h*h))
    fco2=co2 / K0
    pH=-math.log10(h)
    alk=2*co3+hco3

    #Saturacion de Calcita y Aragonita
    omega_ar=Ca*co3/Ksp_a
    omega_cal=Ca*co3/Ksp_c

    return fco2, pH, co2, hco3, co3, h, dic, alk, omega_ar, omega_cal

vcarbEq = np.vectorize(carbEq)
#=============================================================================
#Uso de la funcion - ejemplo
#=============================================================================
#co2=co2*1000000
#dic=dic*1000000

pco2=np.linspace(0.000090, 0.001100, 100)
co2=pco2*K0
dic=np.linspace(0.001526, 0.002100, 100) #moles (para pasar de micromoles a moles, 1587*e-06)

#pco2=pco2*1000
#dic=dic*1000
co2, dic=meshgrid(co2*1000000, dic*1000000)

fco2, pH, co2, hco3, co3, h, dic, alk, omega_ar, omega_cal=vcarbEq(co2=co2, dic=dic, Ca=Ca)


#Ejemplo con valores de referencia de uso para agua de mar normal
#fco2, pH, co2, hco3, co3, h, dic, alk, omega_ar, omega_cal=carbEq(co2=0.00001032997, dic=0.002108, Ca=Ca)
#para este usar calcio=0.0103
#=============================================================================
#Grafico de los resultados
#=============================================================================
print(co2)
print('====================================')
print(dic)
print('====================================')
print(pH)
#print(omega_cal)
#fig, ax = plt.subplots()
im = imshow(pH, cmap=cm.RdBu) # dibujo la funcion
plt.xlabel('CO2')
plt.ylabel('DIC')
plt.ylim(0, 100)

plt.scatter(40, 40, color='k')
#ax.set_xlim(0.00000525, 0.000065)
#ax.set_ylim(0.001, 0.0021)
#plt.xlim(min(co2), max(co2))
# agrego lineas de contorno y rotulos
cset = contour(pH, arange(6.0,9.0,0.2),linewidths=2, cmap=cm.Set2)
clabel(cset,inline=True,fmt='%1.1f',fontsize=10)
colorbar(im) # agrego la barra de colores al costado