using AIBECS
using PyPlot, PyCall
using LinearAlgebra
using GR, Interact

wet3D, grd, T = Primeau_2x2x2.load();

wet3D

grd

grd.depth_3D

T

Matrix(T)

G(x,p) = Λ(x,p) - x / p.τ 

function Λ(x,p)
    λ, h, Ratm = p.λ, p.h, p.Ratm
    return @. λ / h * (Ratm - x) * (z == z₀)
end

iwet = findall(wet3D)
z = grd.depth_3D[iwet] 

z₀ = z[1]

z .== z₀

t = empty_parameter_table()
add_parameter!(t,    :τ, 5730u"yr" / log(2)) # radioactive decay e-folding timescale
add_parameter!(t,    :λ,   50u"m" / 10u"yr") # piston velocity
add_parameter!(t,    :h,      grd.δdepth[1]) # top layer height
add_parameter!(t, :Ratm,      1.0u"mol/m^3") # atmospheric concentration
t 

initialize_Parameters_type(t, "C14_parameters")
p = C14_parameters()   

F, ∇ₓF = state_function_and_Jacobian(p -> T, G) ; 

x = 0.5p.Ratm * ones(5)
F(x,p)

Matrix(∇ₓF(x,p))

function time_steps(x₀, Δt, n, F, ∇ₓF)
    x_hist = [x₀]
    δt = Δt / n
    for i in 1:n
        push!(x_hist, AIBECS.crank_nicolson_step(last(x_hist), p, δt, F, ∇ₓF))
    end
    return reduce(hcat, x_hist), 0:δt:Δt
end

Δt = 5000u"yr" |> u"s" |> ustrip
x₀ = p.Ratm * ones(5)             
x_hist, t_hist = time_steps(x₀, Δt, 1000, F, ∇ₓF)  

C14age_hist = log.(p.Ratm ./ x_hist) * (p.τ * u"s" |> u"yr" |> ustrip)
PyPlot.figure(figsize=(8,4))
PyPlot.plot(t_hist .* 1u"s" .|> u"yr" .|> ustrip, C14age_hist')
PyPlot.xlabel("simulation time (years)")
PyPlot.ylabel("¹⁴C age (years)")
PyPlot.legend("box " .* string.(findall(vec(wet3D))))
PyPlot.title("Simulation of the evolution of ¹⁴C age with Crank-Nicolson time steps")

prob = SteadyStateProblem(F, ∇ₓF, x, p)
s = solve(prob, CTKAlg()).u

log.(p.Ratm ./ s) * (p.τ * u"s" |> u"yr")

Δt = 40000u"yr" |> u"s" |> ustrip
x_hist, t_hist = time_steps(x₀, Δt, 4000, F, ∇ₓF)
PyPlot.figure(figsize=(7,3))
PyPlot.semilogy(t_hist .* 1u"s" .|> u"yr" .|> ustrip, [norm(F(s,p)) for i in 1:size(x_hist,2)], label="steady-state")
PyPlot.semilogy(t_hist .* 1u"s" .|> u"yr" .|> ustrip, [norm(F(x_hist[:,i],p)) for i in 1:size(x_hist,2)], label="time-stepping")
PyPlot.xlabel("simulation time (years)"); PyPlot.ylabel("|F(x,p)| (mol m⁻³ s⁻¹)");
PyPlot.legend(); PyPlot.title("Stability of ¹⁴C age with simulation time")

wet3D, grd, T = OCIM1.load() ;

F, ∇ₓF = state_function_and_Jacobian(p -> T, G) ;

iwet = findall(wet3D)
x = p.Ratm * ones(length(iwet))

p.h = grd.δdepth[1] |> upreferred |> ustrip
z = grd.depth_3D[iwet]
z₀ = z[1]

prob = SteadyStateProblem(F, ∇ₓF, x, p)
s = solve(prob, CTKAlg()).u 
C14age = log.(p.Ratm ./ s) * p.τ * u"s" .|> u"yr"

lon, lat = grd.lon |> ustrip, grd.lat |> ustrip
function horizontal_slice(x, levels, title)
    x_3D = fill(NaN, size(grd))
    x_3D[iwet] .= x .|> ustrip
    GR.figure(figsize=(10,5))
    @manipulate for iz in 1:size(grd)[3]
        GR.xlim([0,360])
        GR.title(string(title, " at $(AIBECS.round(grd.depth[iz])) depth"))
        GR.contourf(lon, lat, x_3D[:,:,iz]', levels=levels)
    end
end  
horizontal_slice(C14age, 0:100:2400, "14C age [yr] using the OCIM1 circulation")

function zonal_slice(x, levels, title)
    x_3D = fill(NaN, size(grd))
    x_3D[iwet] .= x .|> ustrip
    GR.figure(figsize=(10,5))
    @manipulate for longitude in (grd.lon .|> ustrip)
        GR.title(string(title, " at $(round(longitude))°"))
        ilon = findfirst(ustrip.(grd.lon) .== longitude)
        GR.contourf(lat, reverse(-grd.depth .|> ustrip), reverse(x_3D[:,ilon,:], dims=2), levels=levels)
    end
end  
zonal_slice(C14age, 0:100:2400, "14C age [yr] using the OCIM1 circulation")

T_DIP(p) = T

T_POP(p) = buildPFD(grd, wet3D, sinking_speed = w(p))

w(p) = p.w₀ .+ p.w′ * (z .|> ustrip)

relu(x) = (x .≥ 0) .* x
zₑ = 80u"m"    # depth of the euphotic zone
function uptake(DIP, p)
    τ, k = p.τ, p.k
    DIP⁺ = relu(DIP)
    return 1/τ * DIP⁺.^2 ./ (DIP⁺ .+ k) .* (z .≤ zₑ)
end

remineralization(POP, p) = p.κ * POP

geores(x, p) = (p.xgeo .- x) / p.τgeo

G_DIP(DIP, POP, p) = -uptake(DIP, p) + remineralization(POP, p) + geores(DIP, p)
G_POP(DIP, POP, p) =  uptake(DIP, p) - remineralization(POP, p)

t = empty_parameter_table()    # empty table of parameters
add_parameter!(t, :xgeo, 2.12u"mmol/m^3")
add_parameter!(t, :τgeo, 1.0u"Myr")
add_parameter!(t, :k, 6.62u"μmol/m^3")
add_parameter!(t, :w₀, 0.64u"m/d")
add_parameter!(t, :w′, 0.13u"1/d")
add_parameter!(t, :κ, 0.19u"1/d")
add_parameter!(t, :τ, 236.52u"d")
initialize_Parameters_type(t, "Pcycle_Parameters")   # Generate the parameter type
p = Pcycle_Parameters()

nb = length(iwet); x = ones(2nb)
F, ∇ₓF = state_function_and_Jacobian((T_DIP,T_POP), (G_DIP,G_POP), nb) ;

prob = SteadyStateProblem(F, ∇ₓF, x, p)
s = solve(prob, CTKAlg(), preprint=" ").u

DIP, POP = state_to_tracers(s, nb, 2)
DIP_unit = u"mmol/m^3"
horizontal_slice(DIP * u"mol/m^3" .|> DIP_unit, 0:0.3:3, "P-cycle model: DIP [mmol m^{-3}]")

zonal_slice(POP .|> relu .|> log10, -10:1:10, "P-cycle model: POP [log\\(mol m^{-3}\\)]")

ccrs = pyimport("cartopy.crs")
cfeature = pyimport("cartopy.feature")
lon_cyc = [lon; 360+lon[1]]
DIP_3D = fill(NaN, size(grd)); DIP_3D[iwet] = DIP * u"mol/m^3" .|> u"mmol/m^3" .|> ustrip
DIP_3D_cyc = cat(DIP_3D, DIP_3D[:,[1],:], dims=2)
f1 = PyPlot.figure(figsize=(12,5))
@manipulate for iz in 1:24
    withfig(f1, clear=true) do
        ax = PyPlot.subplot(projection = ccrs.EqualEarth(central_longitude=-155.0))
        ax.add_feature(cfeature.COASTLINE, edgecolor="#000000") # black coast lines
        ax.add_feature(cfeature.LAND, facecolor="#CCCCCC")      # gray land
        plt = PyPlot.contourf(lon_cyc, lat, DIP_3D_cyc[:,:,iz], levels=0:0.1:3.5, 
            transform=ccrs.PlateCarree(), zorder=-1, extend="both")
        plt2 = PyPlot.contour(plt, lon_cyc, lat, DIP_3D_cyc[:,:,iz], levels=plt.levels[1:5:end], 
            transform=ccrs.PlateCarree(), zorder=0, extend="both", colors="k")
        ax.clabel(plt2, fmt="%2.1f", colors="k", fontsize=14)
        cbar = PyPlot.colorbar(plt, orientation="vertical", extend="both", ticks=plt2.levels)
        cbar.add_lines(plt2)
        cbar.ax.tick_params(labelsize=14)
        cbar.set_label(label="mmol / m³", fontsize=16)
        PyPlot.title("Publication-quality DIP with Cartopy; depth = $(string(round(typeof(1u"m"),grd.depth[iz])))", fontsize=20)
    end
end