module GEOCLIM

using NetCDF
using UnPack
using MultiAssign
using Roots
using BasicInterpolators
using StaticArrays: SVector
using LinearAlgebra: â‹… #dot product
using CircularArrays

#------------------------------------------------------------------------------
# immutable physical constants

#ideal gas constant [J/K*mole]
const ğ��‘ = 8.31446262

#the Earth's mean radius [m]
const ğ��‘â‚‘ = 6.371e6

#seconds in a year
const ğ��²ğ��« = 31536000.0

#------------------------------------------------------------------------------
# types

include("Climatology.jl")
include("ClimatologyInterpolator.jl")

#------------------------------------------------------------------------------
# weathering functions producing weathering per unit area

export godderis, whak, mac

#==============================================================================
As implementated by
  * GoddÃ©ris, Y. et al. Onset and ending of the late Palaeozoic ice age triggered by tectonically paced rock weathering. Nature Geosci 10, 382â€“386 (2017).
  * Donnadieu, Y. et al. A GEOCLIM simulation of climatic and biogeochemical consequences of Pangea breakup: SIMULATION OF PANGEA BREAKUP. Geochem. Geophys. Geosyst. 7, n/a-n/a (2006).
and in other, similar studies.

Arguments
  r - runoff [m/s]
  T - temperature [K]
  k - calibration constant [mole/m^3]
  Eâ‚� - activation energy [J/mole]
  Tâ‚€ - reference temperature [K]
==============================================================================#
godderis(r, T, k, Eâ‚�, Tâ‚€) = k*r*exp((Eâ‚�/ğ��‘)*(1/Tâ‚€ - 1/T))

#==============================================================================
Originally named after
  * Walker, J. C. G., Hays, P. B. & Kasting, J. F. A negative feedback mechanism for the long-term stabilization of Earthâ€™s surface temperature. J. Geophys. Res. 86, 9776 (1981).
As implementated by
  * Abbot, D. S., Cowan, N. B. & Ciesla, F. J. Indication of Insensitivity of Planetary Weathering Behavior and Habitable Zone to Surface Land Fraction. ApJ 756, 178 (2012).
The important difference between this function and the godderis function is the
pCO2 dependence. The temperature dependence is also written differently.

Arguments
  r - runoff [m/s]
  T - temperature [K]
  pCO2 - carbon dioxide concentration [units must be same as pCO2â‚€]
  k - calibration constant [mole/m^3]
  Tâ‚‘ - scaling of temperature dependence [K]
  Tâ‚€ - reference temperature [K]
  pCO2â‚€ - reference carbon dioxide concentration [units must be same as pCO2]
==============================================================================#
whak(r, T, pCO2, k, Tâ‚‘, Tâ‚€, pCO2â‚€, Î²=0.2) = k*r*exp((T - Tâ‚€)/Tâ‚‘)*(pCO2/pCO2â‚€)^Î²

#==============================================================================
Originally named after
  * Maher, K. & Chamberlain, C. P. Hydrologic Regulation of Chemical Weathering and the Geologic Carbon Cycle. Science 343, 1502â€“1504 (2014).
As implementated by
  * Graham, R. J. & Pierrehumbert, R. Thermodynamic and Energetic Limits on Continental Silicate Weathering Strongly Impact the Climate and Habitability of Wet, Rocky Worlds. ApJ 896, 115 (2020).

Arguments
  r - runoff [m/s]
  T - temperature [K]
  pCO2 - carbon dioxide concentration [bar]
  Tâ‚‘ - scaling of temperature dependence [K]
  Tâ‚€ - reference temperature [K]
  pCO2â‚€ - reference carbon dioxide concentration [bar]
Paramters
  n - thermodynamic pCO2 dependence [-]
  Î› - thermodynamic coefficient for Ceq [-]
  L - flow path length [m] 
  Ï• - porosity [-]
  Ï� - mineral mass to fluid volume ratio [kg mâ�»Â³]
  kâ‚€ - reference rate constant [mol mâ�»Â² yrâ�»Â¹]
  ğ��€ - specific surface area (not weathering surface area) [mÂ²kgâ�»Â¹]
  X - reactive mineral conc. in fresh rock [-]
  tâ‚› - soil age [yr]
  m - mineral molar mass [kg/mol]
  Î¼ - scaling constant [-]
  Î² - pCO2 scaling [-]
==============================================================================#
function mac(r, T, pCO2, Tâ‚‘, Tâ‚€, pCO2â‚€;
             n=0.316,
             Î›=1.4e-3,
             L=1.0,
             Ï•=0.1,
             Ï�=12728.0,
             kâ‚€=8.7e-6,
             ğ��€=1e2,
             X=0.36,
             tâ‚›=1e5,
             m=0.27,
             Î¼=exp(2),
             Î²=0.2)
    #defined for convenience
    Î± = L*Ï•*Ï�*ğ��€*X*Î¼
    #equilibrium concentration
    Ceq = 1e3*Î›*(pCO2^n) #conversion from mol/liter to mol/m3
    #temperature dependence
    a = exp((T - Tâ‚€)/Tâ‚‘)
    #pCO2 dependence
    b = (pCO2/pCO2â‚€)^Î²
    #denominator
    d = 1/(kâ‚€*a*b) + m*ğ��€*tâ‚› + Î±/(r*ğ��²ğ��«*Ceq)
    #weathering per unit area 
    (Î±/d)/ğ��²ğ��«
end

#--------------------------------------
#each total weathering function has a similar form, can generalize

function totalweathering(ğ�’»w::F,
                         ğ�’¸::Climatology,
                         args::Vararg{Real,N};
                         kwargs...
                         ) where {F<:Function, N}
    #get climatology fields
    @unpack mask, r, T, A, f, n, m = ğ�’¸
    #repackage varargs
    X = ntuple(n->Float64(args[n]), N)
    #initialize the weathering sum
    Î£W = 0.0
    #sum weathering at all grid cells
    @inbounds for i âˆˆ 1:n, j âˆˆ 1:m
        if mask[i,j]
            #area of cell [m^2]
            Aáµ¢â±¼ = A[i,j]
            #land fraction of cell [-]
            fáµ¢â±¼ = f[i,j]
            #total weathering in cell
            Î£W += Aáµ¢â±¼*fáµ¢â±¼*ğ�’»w(r[i,j], T[i,j], X...; kwargs...)
        end
    end
    return Î£W
end

#--------------------------------------
#total weathering wrappers for each method to be applied to a Climatology

godderis(ğ�’¸::Climatology, args...) = totalweathering(godderis, ğ�’¸, args...)

whak(ğ�’¸::Climatology, args...) = totalweathering(whak, ğ�’¸, args...)

mac(ğ�’¸::Climatology, args...; kwargs...) = totalweathering(mac, ğ�’¸, args...; kwargs...)

#------------------------------------------------------------------------------

#====
This is a general function to perform root finding with a
ClimatologyInterpolator. The function will find the interpolation location x
where ğ�’»(x, â„�(x)), a function applied to a Climatology, equals the value y.
Typically the function ğ�’» would be a weathering function, but it could be
anything at all.
====#

export findequilibrium

function findequilibrium(â„�::ClimatologyInterpolator{I,ğ�’¯},
                         ğ�’»::F,
                         ğ�“�::Real;
                         tol::Real=1e-4,
                         maxevals::Int=1000
                         ) where {I,ğ�’¯,F}
    #===
    The function to zero is the difference between
    an operation on a Climatology (like a weathering
    estimate) and the desired value of that operation.
    The returned value is converted to the same numeric
    type as the interpolator.
    ===#
    â„±(ğ�“�) = convert(ğ�’¯, ğ�’»(ğ�“�, â„�(ğ�“�)) - ğ�“�)
    #the limits of the ClimatologyInterpolator's range
    ğ�“�â‚�, ğ�“�â‚‚ = â„�.x[1], â„�.x[end]
    #find the root with a bracketing method
    return find_zero(
        â„±,
        (ğ�“�â‚�,ğ�“�â‚‚),
        Roots.Brent(),
        atol=convert(ğ�’¯, tol),
        rtol=convert(ğ�’¯, tol),
        xatol=convert(ğ�’¯, tol),
        xrtol=convert(ğ�’¯, tol),
        maxevals=maxevals
    )
end

#------------------------------------------------------------------------------
# some other general functions

export landfraction
export meanlandlatitude, meanabslandlatitude

function readtopo(fn::String,
                  latname::String="lat", #variable name
                  toponame::String="topo") #variable name
    #read variables from file
    lat = ncread(fn, latname)
    topo = ncread(fn, toponame)
    #transpose if needed
    n = length(lat)
    if size(topo,2) == n
        return lat, collect(transpose(topo))
    else
        @assert size(topo,1) == n
        return lat, topo
    end
end

function landmean(X::AbstractMatrix{ğ�’¯},
                  lat::AbstractVector,
                  mask::BitMatrix,
                  cut::Real=Inf) where {ğ�’¯}
    n, m = size(X)
    @assert length(lat) == n
    num = zero(ğ�’¯)
    den = zero(ğ�’¯)
    @inbounds for i âˆˆ 1:n
        if -cut <= lat[i] <= cut
            #cell weight depends on latitude
            Î± = cos(lat[i]*Ï€/180)
            for j âˆˆ 1:m
                if mask[i,j]
                    num += Î±*X[i,j]
                    den += Î±
                end
            end
        end
    end
    return num/den
end

#computes land fraction of a topography file
#assumes latitude âˆˆ [-90, 90]Â°
#assumes land is where topo > 0
function landfraction(fn::String;
                      latname::String="lat",
                      toponame::String="topo",
                      cut::Real=Inf) #restrict to cells where -cut <= lat <= cut
    lat, topo = readtopo(fn, latname, toponame)
    landmean(topo .> 0, lat, trues(size(topo)), cut)
end

function meanlandlatitude(fn::String;
                          latname::String="lat",
                          toponame::String="topo",
                          cut::Real=Inf) #restrict to cells where -cut <= lat <= cut
    lat, topo = readtopo(fn, latname, toponame)
    latgrid = repeat(lat, 1, size(topo,2))
    landmean(latgrid, lat, topo .> 0, cut)
end

function meanabslandlatitude(fn::String;
                             latname::String="lat",
                             toponame::String="topo",
                             cut::Real=Inf) #restrict to cells where -cut <= lat <= cut
    lat, topo = readtopo(fn, latname, toponame)
    latgrid = repeat(abs.(lat), 1, size(topo,2))
    landmean(latgrid, lat, topo .> 0, cut)
end

end
