text:
## Calculate survival function for all subjects on a grid given (shape,scale,delta) given class-membership
## size(Xe) = N*r
function surv_weibull(grid::Vector{T}, shape::T, scale::T, delta::Vector{T}, Xe::Vector{T}) where T <: Real
    exp.(-exp(dot(Xe,delta)) * weibull_cumhaz(grid, shape, scale))
end

function surv(time::Vector{T}, param::Param, data::Data) where T <: Real
    @unpack bl_shape, bl_scale, delta, cind = param

    N = size(data.Xe, 1)
    res = zeros(T, N, length(time))

    for i in 1:N
        g = cind[i]
        res[i,:] = surv_weibull(time, bl_shape[g], bl_scale[g], delta[g,:], data.Xe[i,:])
    end
    res
end

function surv_mixture(time::Vector{T}, param::Param, data::Data) where T <: Real
    @unpack G, xi, bl_shape, bl_scale, delta = param

    N = size(data.Xe, 1)
    res = zeros(T, N, length(time))

    cp = class_prob(xi, data.Xp)

    for i in 1:N
        for g in 1:G
            res[i,:] += cp[i,g].*surv_weibull(time, bl_shape[g], bl_scale[g], delta[g,:], data.Xe[i,:])
        end
    end
    res
end

function surv_spline(grid::Vector{T}, w::Vector{T}, delta::Vector{T}, Xe::Vector{T}, cumhaz_basis::Matrix{T}) where T <: Real
    exp.(-exp(dot(Xe,delta)) .* cumhaz_basis*exp.(w))
end

function surv(time::Vector{T}, param::ParamSpline, data::Data) where T <: Real
    @unpack w, delta, cind = param

    haz_basis, cumhaz_basis = eval_spline(data, time)

    N = size(data.Xe, 1)
    res = zeros(T, N, length(time))

    for i in 1:N
        g = cind[i]
        res[i,:] = surv_spline(time, w[g,:], delta[g,:], data.Xe[i,:], cumhaz_basis)
    end
    res
end

function surv_mixture(time::Vector{T}, param::ParamSpline, data::Data) where T <: Real
    @unpack G, xi, w, delta = param

    haz_basis, cumhaz_basis = eval_spline(data, time)

    N = size(data.Xe, 1)
    res = zeros(T, N, length(time))

    cp = class_prob(xi, data.Xp)

    for i in 1:N
        for g in 1:G
            res[i,:] += cp[i,g].*surv_spline(time, w[g,:], delta[g,:], data.Xe[i,:], cumhaz_basis)
        end
    end
    res
end

function km_survival(data::Data, grid=sort(unique(data.T))) where T <: AbstractParam
    R"library(survival)"
    df = DataFrame(time=data.T, status=data.E)
    @rput df
    @rput grid
    R"x <- survfit(Surv(time, status) ~ 1, data=df)"
    R"fs <- approxfun(x$time, x$surv, f=1, rule=2, method='constant')"
    R"fe <- approxfun(x$time, x$std.err, f=1, rule=2, method='constant')"
    R"surv <- fs(grid)"
    R"se <- fe(grid)"
    R"res <- data.frame(time=grid, surv=surv, ci_lo=pmax(0, surv - 1.96*se), ci_hi=pmin(1, surv + 1.96*se))"
    @rget res
    res
end

function marginal_survival(x::T, data::Data; grid=sort(unique(data.T))) where T <: Fit
    marginal_survival(merge(x.chains...), data; grid=grid)
end

function marginal_survival(x::T, data::Data; start=1.0, grid=sort(unique(data.T))) where T <: Chain
    ##df = by(conditional_survival(x, data, grid), :time) do z
    ##    (surv=mean(z.mean), ci_lo=mean(z.ci_lo), ci_hi=mean(z.ci_hi))
    ##end

    @assert minimum(grid) >= start

    M = length(x)
    mn = Vector{Float64}(undef, length(grid))
    ci_lo = Vector{Float64}(undef, length(grid))
    ci_hi = Vector{Float64}(undef, length(grid))
    S = Matrix{Float64}(undef, M, length(grid))

    z = get_samples(x)

    for t in 1:M
        S[t,:] = mean(surv_mixture(grid .- start, z[t], data), dims=1)[1,:]
    end

    for j in 1:length(grid)
        Sv = view(S, :, j)
        mn[j] = mean(Sv)
        ci_lo[j] = quantile(Sv, 0.025)
        ci_hi[j] = quantile(Sv, 0.975)
    end

    df = DataFrame(time=grid, surv=mn, ci_lo=ci_lo, ci_hi=ci_hi)

    km = km_survival(data, grid)
    df = vcat(df, km)
    df.method = repeat(["mc", "km"], inner=length(unique(grid)))
    df
end

function class_marginal_survival(x::Fit{T}, data::Data; start=1.0, grid=sort(unique(data.T))) where T <: AbstractParam
    ##df = by(conditional_survival(x, data, grid), :time) do z
    ##    (surv=mean(z.mean), ci_lo=mean(z.ci_lo), ci_hi=mean(z.ci_hi))
    ##end

    @assert minimum(grid) >= start

    G = x.theta[1].G

    M = length(x.theta)
    mn = Matrix{Float64}(undef, G, length(grid))
    ci_lo = Matrix{Float64}(undef, G, length(grid))
    ci_hi = Matrix{Float64}(undef, G, length(grid))
    S = ones(Float64, M, G, length(grid))

    for t in 1:M
        par = x.theta[t]
        for g in 1:G
            sel = par.cind .== g
            if sum(sel) > 0
                S[t,g,:] = mean(surv(grid .- start, par, slice_data(data, sel)), dims=1)[1,:]
            end
        end
    end

    for j in 1:length(grid)
        for g in 1:G
            Sv = view(S, :, g, j)
            mn[g,j] = mean(Sv)
            ci_lo[g,j] = quantile(Sv, 0.025)
            ci_hi[g,j] = quantile(Sv, 0.975)
        end
    end

    df = DataFrame(time=repeat(grid, inner=G), class=repeat(1:G, outer=length(grid)), surv=vec(mn), ci_lo=vec(ci_lo), ci_hi=vec(ci_hi))

    ##km = km_survival(data, grid)
    ##df = vcat(df, km)
    df.method = df.class ##repeat(["mc", "km"], inner=length(unique(grid)))
    df
end

function conditional_survival(x::Tf, Xe::Vector{Tx}, Xp::Vector{Tx}, data::Data; grid=sort(unique(data.T)), start=1.0) where {Tf <: Fit, Tx <: Real}
    conditional_survial(merge(x.chains...), Xe, Xp, data; grid=grid, start=start)
end

function conditional_survival(x::Tf, Xe::Vector{Tx}, Xp::Vector{Tx}, data::Data; grid=sort(unique(data.T)), start=1.0) where {Tf <: Chain, Tx <: Real}
    M = length(x.theta)

    @assert minimum(grid) >= start

    ##mn = SharedArray{Float64,1}((length(grid)))
    ##ci_lo = SharedArray{Float64,1}((length(grid)))
    ##ci_hi = SharedArray{Float64,1}((length(grid)))

    mn = Vector{Float64}(undef, length(grid))
    ci_lo = Vector{Float64}(undef, length(grid))
    ci_hi = Vector{Float64}(undef, length(grid))
    S = Matrix{Float64}(undef, M, length(grid))

    z = get_samples(x)

    for t in 1:M
        S[t,:] = surv_mixture(grid .- start, z[t], Xe, Xp, data)
    end

    for j in 1:length(grid)
        Sv = view(S, :, j)
        mn[j] = mean(Sv)
        ci_lo[j] = quantile(Sv, 0.025)
        ci_hi[j] = quantile(Sv, 0.975)
    end

    DataFrame(time=grid, mean=mn, ci_lo=ci_lo, ci_hi=ci_hi)
end


text:
#=
Popis
  Funkcia vypocita hodnotu eccentricity centrality pre kazdy vrchol.

Parametre
  g : datova struktura grafu :: GenericGraph
  dist_key : nazov atributu, v ktorom je ulozene ohodnotenie hrany :: String

Navratova hodnota
  eccentricity : slovnik s mapovanim (kluc=vrchol_grafu => hodnota=eccentricita) :: Dict{Int64, Float64}

Poznamky
  Implementacia tejto funkcie vyuziva schopnost poli, ktora sa nazyva list comprehension.
  Vdaka tejto vlastnosti je mozne vykonat zlozitejsie operacie, funkcie, podmienky priamo pri tvorbe pola.
=#

function centrality_eccentricity_list_comprehension(g::GenericGraph; dist_key::AbstractString="weight")

  return eccentricity = [ v.index => (z = 1.0/maximum(filter(x -> x != Inf, dijkstra_shortest_paths(g, AttributeEdgePropertyInspector{Float64}(dist_key), [v]).dists))) == Inf ? 0.0 : z  for v = vertices(g) ]

end

ext: jl

max_stars_repo_name: mpuk/GraphAnalysis

max_issues_repo_name: mpuk/GraphAnalysis

max_forks_repo_name: mpuk/GraphAnalysis


text:
module ZeroOne

	include("./ZeroOne/Sample.jl")
	include("./ZeroOne/Reinforcement.jl")
	using .Reinforcement

end


text:
if !@isdefined ExoplanetsSysSim
    using ExoplanetsSysSim
end
using Distributions
using Roots #if we want to use the 'find_zero()' function
using Optim

##### This code is a translation of 'MRpredict.R' at 'https://github.com/Bo-Ning/Predicting-exoplanet-mass-and-radius-relationship/blob/master/MR-predict/MRpredict.R'

struct MR_param_Ning2018
    Mass_min::Float64 #log10 scale
    Mass_max::Float64 #log10 scale
    Radius_min::Float64 #log10 scale
    Radius_max::Float64 #log10 scale
    deg::Int64
    weights_mle::Vector{Float64}
    indices_keep::Vector{Int64}
    indices_div::Vector{Int64}
    indices_rem::Vector{Int64}
end

function pdfnorm_beta(x::Float64, x_obs::Float64, x_sd::Float64, x_max::Float64, x_min::Float64, shape1::Int64, shape2::Int64 ; log::Bool=true)
    if log == true
        norm_dist = Normal(10^x, x_sd)
    else
        norm_dist = Normal(x, x_sd)
    end
    #beta_dist = Beta(shape1, shape2)
    return pdf.(norm_dist, x_obs) .* pdf.(Beta(shape1, shape2), (x-x_min)/(x_max-x_min)) / (x_max-x_min)
end

function fn_for_integrate(x_obs::Float64, x_sd::Float64, deg::Int, degree::Int, x_max::Float64, x_min::Float64 ; log::Bool=false, abs_tol::Float64=1e-10)
    shape1, shape2 = degree, deg-degree+1
    return quadgk(x -> pdfnorm_beta(x, x_obs, x_sd, x_max, x_min, shape1, shape2; log=log), x_min, x_max, abstol=abs_tol)[1]
end

function conditional_density(y::Float64, y_max::Float64, y_min::Float64, x_max::Float64, x_min::Float64, deg::Int64, w_hat::Vector{Float64}, indices_keep::Vector{Int64}, indices_div::Vector{Int64}, indices_rem::Vector{Int64}; y_sd::Union{Nothing,Float64}=nothing, qtl::Vector{Float64}=[0.16, 0.84])
    deg_vec = 1:deg

    #To compute conditional mean, variance, quantile, distribution:
    y_beta_indv = zeros(length(indices_rem))
    if y_sd == nothing
        map!(degree -> pdf.(Beta(degree, deg-degree+1), (y-y_min)/(y_max-y_min)), y_beta_indv, deg_vec[indices_rem])
        y_beta_indv /= (y_max-y_min)
    else
        map!(degree -> fn_for_integrate(y, y_sd, deg, degree, y_max, y_min), y_beta_indv, deg_vec[indices_rem])
    end
    denominator = sum(w_hat .* y_beta_indv)

    #Mean:
    mean = sum(w_hat::Vector{Float64} .* (((deg_vec::UnitRange{Int64})/(deg+1)*(x_max-x_min) .+ x_min)[indices_div] .* y_beta_indv::Vector{Float64})::Vector{Float64}) / denominator::Float64

    #Variance:
    var = sum(w_hat::Vector{Float64} .* (((deg_vec::UnitRange{Int64}) .* (deg .- deg_vec .+ 1) / ((deg+1)^2*(deg+2))*(x_max-x_min)^2)[indices_div] .* y_beta_indv::Vector{Float64})::Vector{Float64}) / denominator::Float64

    #Quantile:
    function pbeta_conditional_density(x::Float64)::Float64
        function mix_density(j::Float64)::Float64
            @fastmath @inbounds @views return sum(w_hat::Vector{Float64} .* (map(degree -> cdf.(Beta(degree, deg::Int64-degree+1), (j-x_min)/(x_max-x_min)), (deg_vec::UnitRange{Int64})[indices_div]) .* y_beta_indv::Vector{Float64})) / denominator::Float64
        end

        return mix_density(x)
    end

    function conditional_quantile(q::Float64)::Float64
        function g(x::Float64)::Float64
            return pbeta_conditional_density(x) - q
        end
        function g2(x::Float64)::Float64
            return (pbeta_conditional_density(x) - q)^2
        end

        #return find_zero(g, (x_min, x_max), FalsePosition())
        #return find_zero(g, (x_min, x_max))::Float64 ###'find_zero' is from 'using Roots' (is this a good function to use?); Eric suggested brent's method in the Optim package

        result = optimize(g2, x_min, x_max, Brent(), rel_tol=1e-3, abs_tol=1e-3) #Brent(); GoldenSection()
        if Optim.converged(result)
            return Optim.minimizer(result)::Float64
        else
            return find_zero(g, (x_min, x_max))::Float64
        end
    end

    quantile = map(conditional_quantile, qtl)

    return (mean, var, quantile, denominator, y_beta_indv)
end

function conditional_density(y::Float64, y_max::Float64, y_min::Float64, x_max::Float64, x_min::Float64, deg::Int64, w_hat::Vector{Float64}, indices_keep::Vector{Int64}; y_sd::Union{Nothing,Float64}=nothing, qtl::Vector{Float64}=[0.16, 0.84])
    deg_vec = 1:deg

    #To compute conditional mean, variance, quantile, distribution:
    y_beta_indv = zeros(deg)
    if y_sd == nothing
        map!(degree -> pdf.(Beta(degree, deg-degree+1), (y-y_min)/(y_max-y_min)), y_beta_indv, deg_vec)
        y_beta_indv /= (y_max-y_min)
    else
        map!(degree -> fn_for_integrate(y, y_sd, deg, degree, y_max, y_min), y_beta_indv, deg_vec)
    end
    y_beta_pdf = kron(ones(deg), y_beta_indv)[indices_keep]
    denominator = sum(w_hat .* y_beta_pdf)

    #Mean:
    mean = sum(w_hat::Vector{Float64} .* kron(deg_vec::UnitRange{Int64}/(deg+1)*(x_max-x_min)+x_min, y_beta_indv::Vector{Float64})[indices_keep]) / denominator::Float64

    #Variance:
    var = sum(w_hat::Vector{Float64} .* kron(deg_vec::UnitRange{Int64} .* (deg-deg_vec+1) / ((deg+1)^2*(deg+2))*(x_max-x_min)^2, y_beta_indv::Vector{Float64})[indices_keep]) / denominator::Float64

    #Quantile:
    function pbeta_conditional_density(x::Float64)::Float64
        function mix_density(j::Float64)::Float64
            @fastmath @inbounds @views return sum(w_hat::Vector{Float64} .* kron(map(degree -> cdf.(Beta(degree, deg-degree+1), (j-x_min)/(x_max-x_min)), deg_vec::UnitRange{Int64}), y_beta_indv::Vector{Float64})[indices_keep]) / denominator::Float64
            #@fastmath @inbounds @views return sum(w_hat::Vector{Float64} .* (map(degree -> cdf.(Beta(degree, deg::Int64-degree+1), (j-x_min)/(x_max-x_min)), (deg_vec::UnitRange{Int64})[indices_div]) .* y_beta_indv::Vector{Float64})) / denominator::Float64
        end

        return mix_density(x)
    end

    function conditional_quantile(q::Float64)::Float64
        function g(x::Float64)::Float64
            return pbeta_conditional_density(x) - q
        end
        function g2(x::Float64)::Float64
            return (pbeta_conditional_density(x) - q)^2
        end

        #return find_zero(g, (x_min, x_max), FalsePosition())
        #return find_zero(g, (x_min, x_max))::Float64 ###'find_zero' is from 'using Roots' (is this a good function to use?); Eric suggested brent's method in the Optim package

        result = optimize(g2, x_min, x_max, Brent(), rel_tol=1e-3, abs_tol=1e-3) #Brent(); GoldenSection()
        if Optim.converged(result)
            return Optim.minimizer(result)::Float64
        else
            return find_zero(g, (x_min, x_max))::Float64
        end
    end

    quantile = map(conditional_quantile, qtl)

    return (mean, var, quantile, denominator, y_beta_indv)
end

function predict_mass_given_radius(Radius::Float64, param::MR_param_Ning2018 ; R_sigma::Union{Nothing,Float64}=nothing, posterior_sample::Bool=false, qtl::Vector{Float64}=[0.16, 0.84])
    #Convert data to log scale:
    l_radius = log10.(Radius)

    #The 'posterior_sample == false' condition can deal with two cases:
    #Case I: if input data do not have measurement errors
    #Case II: if input data have measurement error
    if posterior_sample == false
        predicted_value = conditional_density(l_radius, param.Radius_max, param.Radius_min, param.Mass_max, param.Mass_min, param.deg, param.weights_mle, param.indices_keep, param.indices_div, param.indices_rem; y_sd=R_sigma, qtl=qtl)
        ###predicted_value = conditional_density(l_radius, param.Radius_max, param.Radius_min, param.Mass_max, param.Mass_min, param.deg, param.weights_mle, param.indices_keep; y_sd=R_sigma, qtl=qtl)
        predicted_mean = predicted_value[1]
        predicted_quantiles = predicted_value[3]
    elseif posterior_sample == true
        #Case III: if the input are posterior samples
        radius_sample = log10(Radius)
        k = length(radius_sample)
        denominator_sample = zeros(k)
        mean_sample = zeros(k)
        y_beta_indv_sample = zeros(k, param.deg)

        #Calculate mean:
        for i in 1:k
            results = cond_density_estimation(y=radius_sample[i], y_max=param.Radius_max, y_min=param.Radius_min, x_max=param.Mass_max, x_min=param.Mass_min, deg=param.deg, w_hat=param.weights_mle, qtl=quantile, only_output_mean=true) ###what function is this??? Is it part of some CRAN package...
            mean_sample[i] = (results[1]) ###why are there parentheses here? It makes no difference...
            denominator_sample[i] = results[2]
            y_beta_indv_sample[i,:] = results[3:57]
        end
        predicted_mean = mean(mean_sample)

        #Calculate 16% and 84% quantiles:
        #Mixture of the CDF of k conditional densities
        function pbeta_conditional_density(x)
            function mix_density(j)
                deg_vec = 1:param.deg
                x_indv_cdf = map(degree -> cdf.(Beta(degree, param.deg-degree+1), (j-x_min)/(x_max-x_min)), deg_vec)
                quantile_numerator = zeros(k)
                p_beta_sample = zeros(k)
                for ii in 1:k
                    quantile_numerator[ii] = sum(param.weights_mle .* kron(x_indv_cdf, y_beta_indv_sample[ii,:]))
                    p_beta_sample[ii] = quantile_numerator[ii]/denominator_sample[ii]
                end
                return p_beta = mean(p_beta_sample)
            end

            return map(mix_density, x)
        end

        function mixture_conditional_quantile(q, x_min, x_max)
            function g(x)
                return pbeta_conditional_density(x) - q
            end
            return find_zero(g, (x_min, x_max)) ###see note above
            #return Optim.minimum(optimize(g, x_min, x_max))
        end

        predicted_quantiles = map(q -> mixture_conditional_quantile(q, param.Mass_min, param.Mass_max), qtl)
    end

    #Return the output:
    return (predicted_mean, predicted_quantiles)::Tuple{Float64, Vector{Float64}}
end

function draw_planet_mass_from_radius_Ning2018(Radius::Float64, param::MR_param_Ning2018)
    #This function takes in a Radius (in solar radii) and draws a mass (returning in solar masses) probabilistically
    Radius_in_earths = Radius/ExoplanetsSysSim.earth_radius
    log_Radius_in_earths = log10.(Radius_in_earths)
    @assert param.Radius_min < log_Radius_in_earths < param.Radius_max
    q::Float64 = rand()
    l_mass = predict_mass_given_radius(Radius_in_earths, param; qtl=[q])[2][1]
    return (10^l_mass)*ExoplanetsSysSim.earth_mass
end

function generate_planet_mass_from_radius_Ning2018(Radius::Float64, sim_param::SimParam)
    global MR_param
    return draw_planet_mass_from_radius_Ning2018(Radius, MR_param)
end

function interpolate_planet_mass_from_radius_quantile_Ning2018_table(log_Radius::Float64, quantile::Float64)
    #If using "GridInterpolations" package (very slow for repeated use):
    #=
    global grid, gridData
    x = [log_Radius, quantile]
    return interpolate(grid, gridData, x)
    =#

    #If using "Interpolations" package (fastest):
    global scaled_itp
    return scaled_itp(log_Radius, quantile)
end

function generate_planet_mass_from_radius_Ning2018_table(Radius::Float64, sim_param::SimParam)
    Radius_in_earths = Radius/ExoplanetsSysSim.earth_radius
    log_Radius_in_earths = log10.(Radius_in_earths)
    q::Float64 = rand()
    log_Mass = interpolate_planet_mass_from_radius_quantile_Ning2018_table(log_Radius_in_earths, q)
    return (10^log_Mass)*ExoplanetsSysSim.earth_mass
end

##### To initialize the M-R model parameters:

weights_mle = CSV.read("weights.mle.csv")[:x]
degrees = 55

N_keep = 25
indices_keep = sortperm(weights_mle, rev=true)[1:N_keep]
indices_div = 1 .+ div.(indices_keep .- 1, degrees)
indices_rem = 1 .+ rem.(indices_keep .- 1, degrees)
weights = weights_mle[indices_keep]
MR_param = MR_param_Ning2018(-1., 3.809597, -0.302, 1.357509, degrees, weights, indices_keep, indices_div, indices_rem) #(-1., 3.809597, -0.3, 1.357509, 55, weights_mle)

##### Examples:

#=
#Observation without measurement errors:

Radius = 5. #original scale, not log scale
predict_result = predict_mass_given_radius(Radius, MR_param)
println(predict_result)

#Observation with a measurement error:

Radius = 5. #original scale, not log scale
R_sigma = 0.1
predict_result = predict_mass_given_radius(Radius, MR_param; R_sigma=0.1)
println(predict_result)

#If want to change 16% and 84% quantiles to 5% and 95% quantiles:

Radius = 5. #original scale, not log scale
predict_result = predict_mass_given_radius(Radius, MR_param; qtl=[0.05, 0.95])
println(predict_result)
=#

#Input are posterior samples: ###currently broken because the function 'cond_density_estimation' is undefined
#=
Radius_dist = Normal(5, 0.5) #original scale, not log scale
Radius_samples = rand(Radius_dist, 100)
predict_result = predict_mass_given_radius(Radius_samples, MR_param; R_sigma=0.1, posterior_sample=true)
println(predict_result)
=#

##### For plotting and timing the model:
#=
#Pkg.add("PyPlot")
using PyPlot

N_keep_array = [3025, 19, 18]
colors = ["blue", "green", "red"]
Radii = 10 .^(range(MR_param.Radius_min+0.01, stop=MR_param.Radius_max-0.01, length=1000))

fig, ax = subplots()
for (i,N_keep) in enumerate(N_keep_array)
    indices_keep = sortperm(weights_mle, rev=true)[1:N_keep]
    indices_div = 1 .+ div.(indices_keep .- 1, degrees)
    indices_rem = 1 .+ rem.(indices_keep .- 1, degrees)
    weights = weights_mle[indices_keep]
    MR_param = MR_param_Ning2018(-1., 3.809597, -0.302, 1.357509, degrees, weights, indices_keep, indices_div, indices_rem)

    @time Masses_predict = map(r -> predict_mass_given_radius(r, MR_param; qtl=[0.16, 0.84]), Radii)

    Masses_mean = 10 .^[x[1] for x in Masses_predict]
    Masses_lower = 10 .^[x[2][1] for x in Masses_predict]
    Masses_upper = 10 .^[x[2][2] for x in Masses_predict]

    loglog(Radii, Masses_mean, color=colors[i], linewidth=2.0, linestyle="-", label="M = "*string(N_keep))
    loglog(Radii, Masses_lower, color=colors[i], linewidth=1.0, linestyle="--")
    loglog(Radii, Masses_upper, color=colors[i], linewidth=1.0, linestyle="--")
end
xlabel(L"$R_p (R_\oplus)$")
ylabel(L"$M_p (M_\oplus)$")
legend(loc="upper left")
=#

##### If we want to generate and save a pre-computed table of masses on a radius vs. quantile grid:
#=
N_keep = 3025
indices_keep = sortperm(weights_mle, rev=true)[1:N_keep]
indices_div = 1 .+ div.(indices_keep .- 1, degrees)
indices_rem = 1 .+ rem.(indices_keep .- 1, degrees)
weights = weights_mle[indices_keep]
MR_param = MR_param_Ning2018(-1., 3.809597, -0.302, 1.357509, degrees, weights, indices_keep, indices_div, indices_rem)

N_radii = 1001
N_quantiles = 1001
log_Radii = collect(range(MR_param.Radius_min, stop=MR_param.Radius_max, length=N_radii))
Radii = 10 .^log_Radii
quantiles = collect(range(0., stop=1.0, length=N_quantiles))

file_name = "MRpredict_table_weights"*string(N_keep)*"_R"*string(N_radii)*"_Q"*string(N_quantiles)*".txt"
f = open(file_name, "w")
println(f, "# All masses are in log10; weights used = ", N_keep)
println(f, "# First uncommented line is the header with the column labels (first column contains the radii in log10 while the remaining columns correspond to the quantiles)")
writedlm(f, reshape(append!(["log_R"], "q".*string.(quantiles)), (1,:)), ", ")

log_Mass_table = zeros(N_radii, N_quantiles + 1)
for (i,r) in enumerate(Radii)
    log_Mass_quantiles = predict_mass_given_radius(r, MR_param; qtl=quantiles)[2]
    log_Mass_table[i,:] = append!([log_Radii[i]], log_Mass_quantiles)
    println(i, ", Radius = ", r)
end
writedlm(f, log_Mass_table, ", ")
close(f)
=#

##### To load a pre-computed table for interpolating the mass radius relation:

log_Mass_table = CSV.read("MRpredict_table_weights3025_R1001_Q1001.txt", header=3)

#One way to do the interpolation is to use "GridInterpolations" but this is very slow:
#=
#Pkg.add("GridInterpolations")
using GridInterpolations

N_quantiles = 1001
quantiles = collect(range(0., stop=1.0, length=N_quantiles))
grid = RectangleGrid(log_Mass_table[:,1], quantiles)
gridData = convert(Array, log_Mass_table[:,2:end])

#x = [0., 0.5] #an example data point to interpolate at
#interpolate(grid, gridData, x) #interpolate at the example data point
=#

#A faster way to do the interpolation is to first construct an interpolation object:

#Pkg.add("Interpolations") #NOTE: the README is somewhat out-dated so some of the examples/syntax do not actually work as the way they are shown
using Interpolations

N_radii = 1001
N_quantiles = 1001
log_Radii = range(MR_param.Radius_min, stop=MR_param.Radius_max, length=N_radii)
quantiles = range(0., stop=1.0, length=N_quantiles)

table_data = convert(Array, log_Mass_table[:,2:end])
itp = interpolate(table_data, BSpline(Cubic(Line(OnGrid())))) #interpolation object where the x and y axes are indices
scaled_itp = Interpolations.scale(itp, log_Radii, quantiles) #scaled interpolation object where the x and y axes are scaled to their physical units

#scaled_itp(0., 0.5) #calls the interpolation object to perform an interpolation at a given point

##### For timing the functions:

sim_param = setup_sim_param_model()
#Radii = ones(10000)*ExoplanetsSysSim.earth_radius
Radii = (10 .^(range(MR_param.Radius_min+0.01, stop=MR_param.Radius_max-0.01, length=10000)))*ExoplanetsSysSim.earth_radius
@time Masses = map(r -> generate_planet_mass_from_radius_Ning2018(r, sim_param), Radii)

@time Masses = map(r -> generate_planet_mass_from_radius_Ning2018_table(r, sim_param), Radii)


text:
function known_lastindex(::Type{T}) where {T}
    if known_offset1(T) === nothing || known_length(T) === nothing
        return nothing
    else
        return known_length(T) - known_offset1(T) + 1
    end
end
known_lastindex(@nospecialize x) = known_lastindex(typeof(x))

@inline static_lastindex(x) = Static.maybe_static(known_lastindex, lastindex, x)

function Base.first(x::AbstractVector, n::StaticInt)
    @boundscheck n < 0 && throw(ArgumentError("Number of elements must be nonnegative"))
    start = offset1(x)
    @inbounds x[start:min((start - one(start)) + n, static_lastindex(x))]
end

function Base.last(x::AbstractVector, n::StaticInt)
    @boundscheck n < 0 && throw(ArgumentError("Number of elements must be nonnegative"))
    stop = static_lastindex(x)
    @inbounds x[max(offset1(x), (stop + one(stop)) - n):stop]
end

"""
    ArrayInterface.to_indices(A, I::Tuple) -> Tuple

Converts the tuple of indexing arguments, `I`, into an appropriate form for indexing into `A`.
Typically, each index should be an `Int`, `StaticInt`, a collection with values of `Int`, or a collection with values of `CartesianIndex`
This is accomplished in three steps after the initial call to `to_indices`:

# Extended help

This implementation differs from that of `Base.to_indices` in the following ways:

*  `to_indices(A, I)` never results in recursive processing of `I` through
  `to_indices(A, axes(A), I)`. This is avoided through the use of an internal `@generated`
  method that aligns calls of `to_indices` and `to_index` based on the return values of
  `ndims_index`. This is beneficial because the compiler currently does not optimize away
  the increased time spent recursing through
    each additional argument that needs converting. For example:
    ```julia
    julia> x = rand(4,4,4,4,4,4,4,4,4,4);

    julia> inds1 = (1, 2, 1, 2, 1, 2, 1, 2, 1, 2);

    julia> inds2 = (1, CartesianIndex(1, 2), 1, CartesianIndex(1, 2), 1, CartesianIndex(1, 2), 1);

    julia> inds3 = (fill(true, 4, 4), 2, fill(true, 4, 4), 2, 1, fill(true, 4, 4), 1);

    julia> @btime Base.to_indices(\$x, \$inds2)
    1.105 μs (12 allocations: 672 bytes)
    (1, 1, 2, 1, 1, 2, 1, 1, 2, 1)

    julia> @btime ArrayInterface.to_indices(\$x, \$inds2)
    0.041 ns (0 allocations: 0 bytes)
    (1, 1, 2, 1, 1, 2, 1, 1, 2, 1)

    julia> @btime Base.to_indices(\$x, \$inds3);
    340.629 ns (14 allocations: 768 bytes)

    julia> @btime ArrayInterface.to_indices(\$x, \$inds3);
    11.614 ns (0 allocations: 0 bytes)

    ```
* Recursing through `to_indices(A, axes, I::Tuple{I1,Vararg{Any}})` is intended to provide
  context for processing `I1`. However, this doesn't tell use how many dimensions are
  consumed by what is in `Vararg{Any}`. Using `ndims_index` to directly align the axes of
  `A` with each value in `I` ensures that a `CartesiaIndex{3}` at the tail of `I` isn't
  incorrectly assumed to only consume one dimension.
* `Base.to_indices` may fail to infer the returned type. This is the case for `inds2` and
  `inds3` in the first bullet on Julia 1.6.4.
* Specializing by dispatch through method definitions like this:
  `to_indices(::ArrayType, ::Tuple{AxisType,Vararg{Any}}, ::Tuple{::IndexType,Vararg{Any}})`
  require an excessive number of hand written methods to avoid ambiguities. Furthermore, if
  `AxisType` is wrapping another axis that should have unique behavior, then unique parametric
  types need to also be explicitly defined.
* `to_index(axes(A, dim), index)` is called, as opposed to `Base.to_index(A, index)`. The
  `IndexStyle` of the resulting axis is used to allow indirect dispatch on nested axis types
  within `to_index`.
"""
to_indices(A, ::Tuple{}) = ()
@inline function to_indices(a::A, inds::I) where {A,I}
    _to_indices(a, inds, IndexStyle(A), static(ndims(A)), IndicesInfo(I))
end
@generated function _to_indices(a, inds, ::S, ::StaticInt{N}, ::IndicesInfo{NI,NS,IS}) where {S,N,NI,NS,IS}
    _to_indices_expr(S, N, NI, NS, IS)
end
function _to_indices_expr(S::DataType, N::Int, ni, ns, is)
    blk = Expr(:block, Expr(:meta, :inline))
    # check to see if we are dealing with linear indexing over a multidimensional array
    if length(ni) == 1 && ni[1] === 1
        push!(blk.args, :((to_index(LazyAxis{:}(a), getfield(inds, 1)),)))
    else
        indsexpr = Expr(:tuple)
        ndi = Int[]
        nds = Int[]
        isi = Bool[]
        # 1. unwrap AbstractCartesianIndex, CartesianIndices, Indices
        for i in 1:length(ns)
            ns_i = ns[i]
            if ns_i isa Tuple
                for j in 1:length(ns_i)
                    push!(ndi, 1)
                    push!(nds, ns_i[j])
                    push!(isi, false)
                    push!(indsexpr.args, :(getfield(getfield(getfield(inds, $i), 1), $j)))
                end
            else
                push!(indsexpr.args, :(getfield(inds, $i)))
                push!(ndi, ni[i])
                push!(nds, ns_i)
                push!(isi, is[i])
            end
        end

        # 2. find splat indices
        splat_position = 0
        remaining = N
        for i in eachindex(ndi, nds, isi)
            if isi[i] && splat_position == 0
                splat_position = i
            else
                remaining -= ndi[i]
            end
        end
        if splat_position !== 0
            for _ in 2:remaining
                insert!(ndi, splat_position, 1)
                insert!(nds, splat_position, 1)
                insert!(indsexpr.args, splat_position, indsexpr.args[splat_position])
            end
        end

        # 3. insert `to_index` calls
        dim = 0
        nndi = length(ndi)
        for i in 1:nndi
            ndi_i = ndi[i]
            if ndi_i == 1
                dim += 1
                indsexpr.args[i] = :(to_index($(_axis_expr(N, dim)), $(indsexpr.args[i])))
            else
                subaxs = Expr(:tuple)
                for _ in 1:ndi_i
                    dim += 1
                    push!(subaxs.args, _axis_expr(N, dim))
                end
                if i == nndi && S <: IndexLinear
                    indsexpr.args[i] = :(to_index(LinearIndices($(subaxs)), $(indsexpr.args[i])))
                else
                    indsexpr.args[i] = :(to_index(CartesianIndices($(subaxs)), $(indsexpr.args[i])))
                end
            end
        end
        push!(blk.args, Expr(:(=), :axs, :(lazy_axes(a))))
        push!(blk.args, :(flatten_tuples($(indsexpr))))
    end
    return blk
end

function _axis_expr(N::Int, d::Int)
    if d <= N
        :(getfield(axs, $d))
    else  # ndims(a)+ can only have indices 1:1
        :($(SOneTo(1)))
    end
end

@inline function flatten_tuples(inds::I) where {I}
    if @generated
        t = Expr(:tuple)
        for i in 1:fieldcount(I)
            p = fieldtype(I, i)
            if p <: Tuple
                for j in 1:fieldcount(p)
                    push!(t.args, :(@inbounds(getfield(getfield(inds, $i), $j))))
                end
            else
                push!(t.args, :(@inbounds(getfield(inds, $i))))
            end
        end
        Expr(:block, Expr(:meta, :inline), t)
    else
        out = ()
        for i in inds
            out = i isa Tuple ? (out..., i...) : (out..., i)
        end
        out
    end
end

"""
    ArrayInterface.to_index([::IndexStyle, ]axis, arg) -> index

Convert the argument `arg` that was originally passed to `ArrayInterface.getindex` for the
dimension corresponding to `axis` into a form for native indexing (`Int`, Vector{Int}, etc.).

`ArrayInterface.to_index` supports passing a function as an index. This function-index is
transformed into a proper index.

```julia
julia> using ArrayInterface, Static

julia> ArrayInterface.to_index(static(1):static(10), 5)
5

julia> ArrayInterface.to_index(static(1):static(10), <(5))
static(1):4

julia> ArrayInterface.to_index(static(1):static(10), <=(5))
static(1):5

julia> ArrayInterface.to_index(static(1):static(10), >(5))
6:static(10)

julia> ArrayInterface.to_index(static(1):static(10), >=(5))
5:static(10)

```

Use of a function-index helps ensure that indices are inbounds

```julia
julia> ArrayInterface.to_index(static(1):static(10), <(12))
static(1):10

julia> ArrayInterface.to_index(static(1):static(10), >(-1))
1:static(10)
```

New axis types with unique behavior should use an `IndexStyle` trait:
```julia
to_index(axis::MyAxisType, arg) = to_index(IndexStyle(axis), axis, arg)
to_index(::MyIndexStyle, axis, arg) = ...
```

"""
to_index(x, i::Slice) = i
to_index(x, ::Colon) = indices(x)
to_index(::LinearIndices{0,Tuple{}}, ::Colon) = Slice(static(1):static(1))
to_index(::CartesianIndices{0,Tuple{}}, ::Colon) = Slice(static(1):static(1))
# logical indexing
to_index(x, i::AbstractArray{Bool}) = LogicalIndex(i)
to_index(x::LinearIndices, i::AbstractArray{Bool}) = LogicalIndex{Int}(i)
# cartesian indexing
@inline to_index(x, i::CartesianIndices{0}) = i
@inline to_index(x, i::CartesianIndices) = getfield(i, :indices)
@inline to_index(x, i::CartesianIndex) = Tuple(i)
@inline to_index(x, i::NDIndex) = Tuple(i)
@inline to_index(x, i::AbstractArray{<:AbstractCartesianIndex}) = i
@inline function to_index(x, i::Base.Fix2{<:Union{typeof(<),typeof(isless)},<:Union{Base.BitInteger,StaticInt}})
    static_first(x):min(_sub1(canonicalize(i.x)), static_last(x))
end
@inline function to_index(x, i::Base.Fix2{typeof(<=),<:Union{Base.BitInteger,StaticInt}})
    static_first(x):min(canonicalize(i.x), static_last(x))
end
@inline function to_index(x, i::Base.Fix2{typeof(>=),<:Union{Base.BitInteger,StaticInt}})
    max(canonicalize(i.x), static_first(x)):static_last(x)
end
@inline function to_index(x, i::Base.Fix2{typeof(>),<:Union{Base.BitInteger,StaticInt}})
    max(_add1(canonicalize(i.x)), static_first(x)):static_last(x)
end
# integer indexing
to_index(x, i::AbstractArray{<:Integer}) = i
to_index(x, @nospecialize(i::StaticInt)) = i
to_index(x, i::Integer) = Int(i)
@inline to_index(x, i) = to_index(IndexStyle(x), x, i)
function to_index(S::IndexStyle, x, i)
    throw(ArgumentError(
        "invalid index: $S does not support indices of type $(typeof(i)) for instances of type $(typeof(x))."
    ))
end

"""
    unsafe_reconstruct(A, data; kwargs...)

Reconstruct `A` given the values in `data`. New methods using `unsafe_reconstruct`
should only dispatch on `A`.
"""
function unsafe_reconstruct(axis::OneTo, data; kwargs...)
    if axis === data
        return axis
    else
        return OneTo(data)
    end
end
function unsafe_reconstruct(axis::UnitRange, data; kwargs...)
    if axis === data
        return axis
    else
        return UnitRange(first(data), last(data))
    end
end
function unsafe_reconstruct(axis::OptionallyStaticUnitRange, data; kwargs...)
    if axis === data
        return axis
    else
        return OptionallyStaticUnitRange(static_first(data), static_last(data))
    end
end
function unsafe_reconstruct(A::AbstractUnitRange, data; kwargs...)
    return static_first(data):static_last(data)
end

"""
    to_axes(A, inds) -> Tuple

Construct new axes given the corresponding `inds` constructed after
`to_indices(A, args) -> inds`. This method iterates through each pair of axes and
indices calling [`to_axis`](@ref).
"""
@inline function to_axes(A, inds::Tuple)
    if ndims(A) === 1
        return (to_axis(axes(A, 1), first(inds)),)
    elseif Base.length(inds) === 1
        return (to_axis(eachindex(IndexLinear(), A), first(inds)),)
    else
        return to_axes(A, axes(A), inds)
    end
end
# drop this dimension
to_axes(A, a::Tuple, i::Tuple{<:CanonicalInt,Vararg{Any}}) = to_axes(A, _maybe_tail(a), tail(i))
to_axes(A, a::Tuple, i::Tuple{I,Vararg{Any}}) where {I} = _to_axes(StaticInt(ndims_index(I)), A, a, i)
function _to_axes(::StaticInt{1}, A, axs::Tuple, inds::Tuple)
    return (to_axis(_maybe_first(axs), first(inds)), to_axes(A, _maybe_tail(axs), tail(inds))...)
end
@propagate_inbounds function _to_axes(::StaticInt{N}, A, axs::Tuple, inds::Tuple) where {N}
    axes_front, axes_tail = Base.IteratorsMD.split(axs, Val(N))
    if IndexStyle(A) === IndexLinear()
        axis = to_axis(LinearIndices(axes_front), getfield(inds, 1))
    else
        axis = to_axis(CartesianIndices(axes_front), getfield(inds, 1))
    end
    return (axis, to_axes(A, axes_tail, tail(inds))...)
end
to_axes(A, ::Tuple{Ax,Vararg{Any}}, ::Tuple{}) where {Ax} = ()
to_axes(A, ::Tuple{}, ::Tuple{}) = ()

_maybe_first(::Tuple{}) = static(1):static(1)
_maybe_first(t::Tuple) = first(t)
_maybe_tail(::Tuple{}) = ()
_maybe_tail(t::Tuple) = tail(t)

"""
    to_axis(old_axis, index) -> new_axis

Construct an `new_axis` for a newly constructed array that corresponds to the
previously executed `to_index(old_axis, arg) -> index`. `to_axis` assumes that
`index` has already been confirmed to be in bounds. The underlying indices of
`new_axis` begins at one and extends the length of `index` (i.e., one-based indexing).
"""
@inline function to_axis(axis, inds)
    if !can_change_size(axis) &&
       (known_length(inds) !== nothing && known_length(axis) === known_length(inds))
        return axis
    else
        return to_axis(IndexStyle(axis), axis, inds)
    end
end

# don't need to check size b/c slice means it's the entire axis
@inline function to_axis(axis, inds::Slice)
    if can_change_size(axis)
        return copy(axis)
    else
        return axis
    end
end
to_axis(S::IndexLinear, axis, inds) = StaticInt(1):length(inds)

"""
    ArrayInterface.getindex(A, args...)

Retrieve the value(s) stored at the given key or index within a collection. Creating
another instance of `ArrayInterface.getindex` should only be done by overloading `A`.
Changing indexing based on a given argument from `args` should be done through,
[`to_index`](@ref), or [`to_axis`](@ref).
"""
function getindex(A, args...)
    inds = to_indices(A, args)
    @boundscheck checkbounds(A, inds...)
    unsafe_getindex(A, inds...)
end
@propagate_inbounds function getindex(A; kwargs...)
    inds = to_indices(A, find_all_dimnames(dimnames(A), static(keys(kwargs)), Tuple(values(kwargs)), :))
    @boundscheck checkbounds(A, inds...)
    unsafe_getindex(A, inds...)
end
@propagate_inbounds getindex(x::Tuple, i::Int) = getfield(x, i)
@propagate_inbounds getindex(x::Tuple, ::StaticInt{i}) where {i} = getfield(x, i)

## unsafe_getindex ##
function unsafe_getindex(a::A) where {A}
    is_forwarding_wrapper(A) || throw(MethodError(unsafe_getindex, (A,)))
    unsafe_getindex(parent(a))
end

# TODO Need to manage index transformations between nested layers of arrays
function unsafe_getindex(a::A, i::CanonicalInt) where {A}
    if IndexStyle(A) === IndexLinear()
        is_forwarding_wrapper(A) || throw(MethodError(unsafe_getindex, (A, i)))
        return unsafe_getindex(parent(a), i)
    else
        return unsafe_getindex(a, _to_cartesian(a, i)...)
    end
end
function unsafe_getindex(a::A, i::CanonicalInt, ii::Vararg{CanonicalInt}) where {A}
    if IndexStyle(A) === IndexLinear()
        return unsafe_getindex(a, _to_linear(a, (i, ii...)))
    else
        is_forwarding_wrapper(A) || throw(MethodError(unsafe_getindex, (A, i)))
        return unsafe_getindex(parent(a), i, ii...)
    end
end

unsafe_getindex(a, i::Vararg{Any}) = unsafe_get_collection(a, i)

unsafe_getindex(A::Array) = Base.arrayref(false, A, 1)
unsafe_getindex(A::Array, i::CanonicalInt) = Base.arrayref(false, A, Int(i))
@inline function unsafe_getindex(A::Array, i::CanonicalInt, ii::Vararg{CanonicalInt})
    unsafe_getindex(A, _to_linear(A, (i, ii...)))
end

unsafe_getindex(A::LinearIndices, i::CanonicalInt) = Int(i)
unsafe_getindex(A::CartesianIndices{N}, ii::Vararg{CanonicalInt,N}) where {N} = CartesianIndex(ii...)
unsafe_getindex(A::CartesianIndices, ii::Vararg{CanonicalInt}) =
    unsafe_getindex(A, Base.front(ii)...)
unsafe_getindex(A::CartesianIndices, i::CanonicalInt) = @inbounds(A[i])

unsafe_getindex(A::ReshapedArray, i::CanonicalInt) = @inbounds(parent(A)[i])
function unsafe_getindex(A::ReshapedArray, i::CanonicalInt, ii::Vararg{CanonicalInt})
    @inbounds(parent(A)[_to_linear(A, (i, ii...))])
end

unsafe_getindex(A::SubArray, i::CanonicalInt) = @inbounds(A[i])
unsafe_getindex(A::SubArray, i::CanonicalInt, ii::Vararg{CanonicalInt}) = @inbounds(A[i, ii...])

# This is based on Base._unsafe_getindex from https://github.com/JuliaLang/julia/blob/c5ede45829bf8eb09f2145bfd6f089459d77b2b1/base/multidimensional.jl#L755.
#=
    unsafe_get_collection(A, inds)

Returns a collection of `A` given `inds`. `inds` is assumed to have been bounds-checked.
=#
function unsafe_get_collection(A, inds)
    axs = to_axes(A, inds)
    dest = similar(A, axs)
    if map(length, axes(dest)) == map(length, axs)
        Base._unsafe_getindex!(dest, A, inds...)
    else
        Base.throw_checksize_error(dest, axs)
    end
    return dest
end
_ints2range(x::CanonicalInt) = x:x
_ints2range(x::AbstractRange) = x
# apply _ints2range to front N elements
_ints2range_front(::Val{N}, ind, inds...) where {N} =
    (_ints2range(ind), _ints2range_front(Val(N - 1), inds...)...)
_ints2range_front(::Val{0}, ind, inds...) = ()
_ints2range_front(::Val{0}) = ()
# get output shape with given indices
_output_shape(::CanonicalInt, inds...) = _output_shape(inds...)
_output_shape(ind::AbstractRange, inds...) = (Base.length(ind), _output_shape(inds...)...)
_output_shape(::CanonicalInt) = ()
_output_shape(x::AbstractRange) = (Base.length(x),)
@inline function unsafe_get_collection(A::CartesianIndices{N}, inds) where {N}
    if (Base.length(inds) === 1 && N > 1) || stride_preserving_index(typeof(inds)) === False()
        return Base._getindex(IndexStyle(A), A, inds...)
    else
        return reshape(
            CartesianIndices(_ints2range_front(Val(N), inds...)),
            _output_shape(inds...)
        )
    end
end
_known_first_isone(ind) = known_first(ind) !== nothing && isone(known_first(ind))
@inline function unsafe_get_collection(A::LinearIndices{N}, inds) where {N}
    if Base.length(inds) === 1 && ndims_index(typeof(first(inds))) === 1
        return @inbounds(eachindex(A)[first(inds)])
    elseif stride_preserving_index(typeof(inds)) === True() &&
            reduce_tup(&, map(_known_first_isone, inds))
        # create a LinearIndices when first(ind) != 1 is imposable
        return reshape(
            LinearIndices(_ints2range_front(Val(N), inds...)),
            _output_shape(inds...)
        )
    else
        return Base._getindex(IndexStyle(A), A, inds...)
    end
end

"""
    ArrayInterface.setindex!(A, args...)

Store the given values at the given key or index within a collection.
"""
@propagate_inbounds function setindex!(A, val, args...)
    can_setindex(A) || error("Instance of type $(typeof(A)) are not mutable and cannot change elements after construction.")
    inds = to_indices(A, args)
    @boundscheck checkbounds(A, inds...)
    unsafe_setindex!(A, val, inds...)
end
@propagate_inbounds function setindex!(A, val; kwargs...)
    can_setindex(A) || error("Instance of type $(typeof(A)) are not mutable and cannot change elements after construction.")
    inds = to_indices(A, find_all_dimnames(dimnames(A), static(keys(kwargs)), Tuple(values(kwargs)), :))
    @boundscheck checkbounds(A, inds...)
    unsafe_setindex!(A, val, inds...)
end

## unsafe_setindex! ##
function unsafe_setindex!(a::A, v) where {A}
    is_forwarding_wrapper(A) || throw(MethodError(unsafe_setindex!, (A, v)))
    return unsafe_setindex!(parent(a), v)
end
# TODO Need to manage index transformations between nested layers of arrays
function unsafe_setindex!(a::A, v, i::CanonicalInt) where {A}
    if IndexStyle(A) === IndexLinear()
        is_forwarding_wrapper(A) || throw(MethodError(unsafe_setindex!, (A, v, i)))
        return unsafe_setindex!(parent(a), v, i)
    else
        return unsafe_setindex!(a, v, _to_cartesian(a, i)...)
    end
end
function unsafe_setindex!(a::A, v, i::CanonicalInt, ii::Vararg{CanonicalInt}) where {A}
    if IndexStyle(A) === IndexLinear()
        return unsafe_setindex!(a, v, _to_linear(a, (i, ii...)))
    else
        is_forwarding_wrapper(A) || throw(MethodError(unsafe_setindex!, (A, v, i, ii...)))
        return unsafe_setindex!(parent(a), v, i, ii...)
    end
end

function unsafe_setindex!(A::Array{T}, v) where {T}
    Base.arrayset(false, A, convert(T, v)::T, 1)
end
function unsafe_setindex!(A::Array{T}, v, i::CanonicalInt) where {T}
    return Base.arrayset(false, A, convert(T, v)::T, Int(i))
end

unsafe_setindex!(a, v, i::Vararg{Any}) = unsafe_set_collection!(a, v, i)

# This is based on Base._unsafe_setindex!.
#=
    unsafe_set_collection!(A, val, inds)

Sets `inds` of `A` to `val`. `inds` is assumed to have been bounds-checked.
=#
unsafe_set_collection!(A, v, i) = Base._unsafe_setindex!(IndexStyle(A), A, v, i...)


text:
mutable struct SimpleNetwork{T<:Integer} <: AbstractNetwork{T}
    # structure
    ## forward star
    fadjlist::Vector{Vector{T}}
    fpoints::Vector{T}

    ## backward star
    badjlist::Vector{Vector{T}}
    bpoints::Vector{T}
    ## need to make sure both of the above representations are consistent

    # data
    ## zones
    zones::Vector{<:AbstractZone{T}}

    ## links
    links::Vector{<:AbstractLink{T}}
end

function SimpleNetwork(nnodes::T; zonetype=SimpleZone{T}, linktype=SimpleLink{T}) where T<:Integer
    fadjlist = [Vector{T}() for _ in one(T):nnodes]
    fpoints = ones(T, nnodes+1)

    badjlist = [Vector{T}() for _ in one(T):nnodes]
    bpoints = ones(T, nnodes+1)

    zones = Vector{zonetype}()

    links = Vector{linktype}()

    SimpleNetwork{T}(fadjlist, fpoints, badjlist, bpoints, zones, links)
end

function SimpleNetwork(nnodes::T, links::Vector{<:AbstractLink{T}}, zones::Vector{<:AbstractZone{T}}) where {T<:Integer}
    net = SimpleNetwork(nnodes)

    for link in links
        add_link!(net, link)
    end

    for zone in zones
        add_zone!(net, zone)
    end
    net
end

function SimpleNetwork(nnodes::T, linkdata, zonedata, args...; linktype::Type{<:AbstractLink{T}}=SimpleLink{T}, zonetype::Type{<:AbstractZone{T}}=SimpleZone{T}, kwargs...) where {T<:Integer}
    links = [linktype(row, args...; kwargs...) for row in eachrow(linkdata)]
    zones = [zonetype(row, args...; kwargs...) for row in eachrow(zonedata)]
    SimpleNetwork(nnodes, links, zones)
end

function add_zone!(net::SimpleNetwork{T}, z::AbstractZone{T}) where T<:Integer
    # check if z's index already taken
    (id(z) ∉ id.(net.zones)) || ((@warn "Zone ID $(id(z)) already taken! New zone not added."); return false)
    (id(z) <= numnodes(net)) || ((@warn "Zone ID $(id(z)) greater than the number of nodes! New zone not added."); return false)
    push!(net.zones, z)
    return true
end

function rem_zone!(net::SimpleNetwork{T}, i::T) where T<:Integer
    for (zi,z) in enumerate(net.zones)
        if i==id(z)
            deleteat!(net.zones, zi)
            return true
        end
    end
    return false
end

function add_link!(net::SimpleNetwork{T}, l::AbstractLink{T}) where T<:Integer
    # check if l's index already taken
    (idx(net, l) in idx.(Ref(net), net.links)) && return false
    # update star representations
    i, j = upn(l), dwn(l)

    @inbounds list = net.badjlist[j]
    index = searchsortedfirst(list, i)
    insert!(list, index, i)
    net.bpoints[(j+1):end] .+= 1

    @inbounds list = net.fadjlist[i]
    index = searchsortedfirst(list, j)
    insert!(list, index, j)
    net.fpoints[(i+1):end] .+= 1

    # add to list
    insert!(net.links, net.fpoints[i]-1+index, l)

    return true
end

function rem_link!(net::SimpleNetwork{T}, idx::U) where {T<:Integer, U<:Integer}
    try
        l = net.links[idx]
        i, j = upn(l), dwn(l)
        deleteat!(net, idx)
        deletesortedfirst!(net.fadjlist[i], j)
        net.fpoints[(i+1):end] .-= 1
        deletesortedfirst!(net.badjlist[j], i)
        net.bpoints[(j+1):end] .-= 1
        return true
    catch
        return false
    end
end

eltype(x::SimpleNetwork{T}) where T = T

has_link(net::SimpleNetwork{T}, i::U, j::U) where {T<:Integer, U<:Integer} = (j in net.fadjlist[i])

idx(net::SimpleNetwork{T}, i::U, j::U) where {T<:Integer, U<:Integer} = has_link(net, i, j) ? net.fpoints[i]-1+searchsortedfirst(net.fadjlist[i],j) : -1
idx(net::SimpleNetwork{T}, l::AbstractLink{U}) where {T<:Integer, U<:Integer} = idx(net, upn(l), dwn(l))
idx(net::SimpleNetwork{T}, t::Tuple{U,U}) where {T<:Integer, U<:Integer} = idx(net, t...)

numnodes(net::SimpleNetwork) = length(net.fpoints)-1
numzones(net::SimpleNetwork) = length(net.zones)
numlinks(net::SimpleNetwork) = length(net.links)
numsources(net::SimpleNetwork) = sum(issource.(net.zones))
numsinks(net::SimpleNetwork) = sum(issink.(net.zones))

sources(net::SimpleNetwork) = id.(filter(issource, net.zones))
sinks(net::SimpleNetwork) = id.(filter(issink, net.zones))
zones(net::SimpleNetwork) = net.zones
links(net::SimpleNetwork) = net.links

outneighbors(net::SimpleNetwork{T}, i::U) where {T<:Integer, U<:Integer} = net.fadjlist[i]
inneighbors(net::SimpleNetwork{T}, j::U) where {T<:Integer, U<:Integer} = net.badjlist[j]

outdegree(net::SimpleNetwork{T}, i::U) where {T<:Integer, U<:Integer} = length(net.fadjlist[i])
indegree(net::SimpleNetwork{T}, j::U) where {T<:Integer, U<:Integer} = length(net.badjlist[j])

fstar(net::SimpleNetwork{<:Integer}) = (adjlist = net.fadjlist, points = net.fpoints)
bstar(net::SimpleNetwork{<:Integer}) = (adjlist = net.badjlist, points = net.bpoints)

issource(net::SimpleNetwork{T}, i::U) where {T<:Integer, U<:Integer} = issource(net.zones[i])
issink(net::SimpleNetwork{T}, i::U) where {T<:Integer, U<:Integer} = issink(net.zones[i])
throughflowallowed(net::SimpleNetwork{T}, i::U) where {T<:Integer, U<:Integer} = throughflowallowed(net.zones[i])
props(net::SimpleNetwork{T}, i::U) where {T<:Integer, U<:Integer} = props(net.zones[i])


text:
X = Float64[2 5; 6 4; 5 3; 2 2; 1 4; 5 4; 3 3; 2 3; 2 4; 8 2; 9 2; 10 2; 11 2; 10 3; 9 1]
d2(x⃗, y⃗) = sqrt(sum((x⃗ - y⃗).^2))
d2²(x⃗, y⃗) = sum((x⃗ - y⃗).^2)
N = size(X)[1]

# prototypes = X[1:4, :]
prototypes = X[12:15, :]
K = size(prototypes)[1]

function compute_distances(X, prototypes)
    N = size(X)[1]
    K = size(prototypes)[1]

    distances = Matrix{Float64}(undef, N, K)
    for k = 1:K
        for n = 1:N
            distances[n, k] = d2(X[n, :], prototypes[k, :])
        end
    end
    distances
end

for t = 1:10
    previous_prototypes = copy(prototypes)

    distances = compute_distances(X, prototypes)

    cluster_assignments = map(argmin, eachrow(distances))
    println("cluster t=$t @ ", cluster_assignments)
    for k = 1:K
        indices = findall(x->x==k, cluster_assignments)
        prototypes[k, :] = sum(X[indices, :], dims=1) / length(indices)
    end
    if all(prototypes == previous_prototypes)
        break
    end
end

distances = compute_distances(X, prototypes)
cluster_assignments = map(argmin, eachrow(distances))
println("cluster @ ", cluster_assignments)

J = 0.0
for k = 1:K
    indices = findall(x->x==k, cluster_assignments)
    for i = 1:length(indices)
        J += d2²(X[indices[i], :], prototypes[k, :])
    end
end
println("J = ", J)


text:
__precompile__()

module CardinalDicts

export AbstractCardinalDictionary,
       CardinalDict, CardinalPairDict,
       keymax, keysmax, clearindex!, isfull

# the order of the params is reversed so constructors work more cleanly
abstract type AbstractCardinalDictionary{V,K} <: Associative{K,V} end

include("support/pairing.jl")
include("support/typedindexing.jl")

include("types/CardinalDict.jl")
include("types/CD_indexed_api.jl")
include("types/CD_dict_api.jl")
include("types/CD_string_show.jl")

include("types/CardinalPairDict.jl")
include("types/CPD_indexed_api.jl")
include("types/CPD_dict_api.jl")
include("types/CPD_string_show.jl")

end # module


text:
################################################################################
# Air.jl
#
# The Air library.
# Functional collections and tools for Julia.
#
# @author Noah C. Benson
#
# MIT License
# Copyright (c) 2020-2021 Noah C. Benson

module Air

include("util.jl")
include("api.jl")
include("ptree.jl")
include("pset.jl")
include("pdict.jl")
include("lazydict.jl")
include("parray.jl")
include("pheap.jl")
include("pwdict.jl")
include("pwset.jl")
include("vars.jl")

include("TX.jl")

"""
    @p{k1 => v1, k2 => v2, ...}
    @p[x1, x2, ...]
    @p[x1 x2...]
    @p(x1, x2, ...)

Yields a persistent data structure, depending on how the macro is called. This
is a shorthand for calling the various constructors directly. The following
expressions are equivalent:
* `@p{k1 => v1, k2 => v2, ...}` and `PDict(k1 => v1, k2 => v2, ...)`
* `@p[x1, x2, ...]` and `PVector([x1, x2, ...])`
* `@p[x1 x2 ...]` and `PArray([x1 x2 ...])`
* `@p(x1, x2, ...)` and `PSet([x1, x2, ...]`

See also: [`PDict`](@ref), [`PVector`](@ref), [`PArray`](@ref), [`PSet`](@ref).

# Examples

```@meta
DocTestSetup = quote
    using Air
end
```

```jldoctest; filter=r"PDict{Symbol, ?Int64} with 3 entries:"
julia> @p{:a => 1, :b => 2, :c => 3}
PDict{Symbol,Int64} with 3 entries:
  :c => 3
  :a => 1
  :b => 2
```

```jldoctest; filter=r"4-element (PArray{Int64, ?1}|PVector{Int64}):"
julia> @p[1, 2, 3, 4]
4-element PArray{Int64,1}:
 1
 2
 3
 4
```

```jldoctest; filter=r"2×2 (PArray{Symbol, ?2}|PMatrix{Symbol}):"
julia> @p[:q2 :q1; :q3 :q4]
2×2 PArray{Symbol,2}:
 :q2  :q1
 :q3  :q4
```

```jldoctest
julia> @p(:a, :b, :a, :c)
PSet{Symbol} with 3 elements:
  :c
  :a
  :b
```
"""
macro p(exprs...)
    return esc(:(Air.PSet([$(exprs...)])))
end
macro p(expr::Expr)
    if expr.head == :braces
        return esc(:(Air.PDict($(expr.args...))))
    elseif expr.head == :vect
        return esc(:(Air.PVector([$(expr.args...)])))
    elseif expr.head == :vcat
        return esc(:(Air.PArray($expr)))
    elseif expr.head == :hcat
        return esc(:(Air.PArray($expr)))
    elseif expr.head == :tuple
        return esc(:(Air.PSet($expr)))
    else
        return esc(:(Air.PSet([$expr])))
    end
end

export Var, @var, Volatile, Actor, Source, tx, @tx, ReentrantRef,
    TransactionalRef, AbstractSourceKernel, getfilter, getfinalize,
    setfilter!, setfinalize!, send, geterror, receive, reset, @p

end


text:
module TestReducedim
using DataArrays, Base.Test, Compat

# Test for fast unit stride BitArray functions
function test_any()
    bits = bitrand(64*5)
    for i = 1:length(bits), j = i:length(bits)
        v = false
        for k = i:j
            v |= bits[k]
        end
        Base.Test.@test DataArrays._any(bits, i, j) == v
    end
end
test_any()

function test_count()
    bits = bitrand(64*5)
    for i = 1:length(bits), j = i:length(bits)
        v = 0
        for k = i:j
            v += bits[k]
        end
        Base.Test.@test DataArrays._count(bits, i, j) == v
    end
end
test_count()

# mapslices from Base, hacked to work for these cases
function safe_mapslices{T}(f::Function, A::AbstractArray{T}, region, skipna)
    dims = intersect(region, 1:ndims(A))
    if isempty(dims)
        if skipna
            naval = f(T[], 1)
            A = copy(A)
            A[isna(A)] = isempty(naval) ? NA : naval[1]
        end
        return A
    end

    if isempty(dims)
        return map(f,A)
    end

    dimsA = collect(size(A))
    ndimsA = ndims(A)
    alldims = collect(1:ndimsA)

    otherdims = setdiff(alldims, dims)

    idx = Vector{Any}(ndimsA)
    fill!(idx, 1)
    Asliceshape = tuple(dimsA[dims]...)
    itershape   = tuple(dimsA[otherdims]...)
    for d in dims
        idx[d] = 1:size(A,d)
    end

    r1 = f(copy(reshape(A[idx...], Asliceshape)); skipna=skipna)

    # determine result size and allocate
    Rsize = copy(dimsA)
    # TODO: maybe support removing dimensions
    if !isa(r1, AbstractArray) || ndims(r1) == 0
        r2 = similar(A, T, 1)
        r2[1] = r1
        r1 = r2
    end
    Rsize[dims] = [size(r1)...; ones(Int,max(0,length(dims)-ndims(r1)))]
    R = similar(r1, tuple(Rsize...))

    ridx = Vector{Any}(ndims(R))
    fill!(ridx, 1)
    for d in dims
        ridx[d] = 1:size(R,d)
    end

    R[ridx...] = r1

    first = true

    for idxs = CartesianRange(itershape)
        if first
            first = false
        else
            ia = [idxs.I...]
            idx[otherdims] = ia
            ridx[otherdims] = ia
            try
                R[ridx...] = f(copy(reshape(A[idx...], Asliceshape)); skipna=skipna)
            catch e
                if (isa(e, ErrorException) && e.msg == "Reducing over an empty array is not allowed.") || (isa(e, ArgumentError) && e.msg == "reducing over an empty collection is not allowed")

                    R[ridx...] = NA
                else
                    println(typeof(e))
                    rethrow(e)
                end
            end
        end
    end

    return R
end

macro test_da_approx_eq(da1, da2)
    quote
        v1 = $da1
        v2 = $da2
        na = isna(v1)
        @test na == isna(v2)
        defined = !na
        if any(defined)
            @test_approx_eq v1[defined] v2[defined]
        end
    end
end

myvarzm(x; skipna::Bool=false) = var(x; mean=0, skipna=skipna)
myvar1m(x; skipna::Bool=false) = var(x; mean=1, skipna=skipna)

for Areduc in (DataArray(rand(3, 4, 5, 6)),
               DataArray(rand(3, 4, 5, 6), rand(3, 4, 5, 6) .< 0.2))
    for skipna = (false, true)
        for region in Any[
            1, 2, 3, 4, 5, (1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4),
            (1, 2, 3), (1, 3, 4), (2, 3, 4), (1, 2, 3, 4)]
            # println("region = $region, skipna = $skipna")

            outputs = Any[DataArray(fill(NaN, Base.reduced_dims(size(Areduc), region)))]
            has_na = anyna(Areduc)
            if has_na && !skipna
                # Should throw an error reducing to non-DataArray
                @test_throws NAException sum!(outputs[1].data, Areduc; skipna=skipna)
            else
                # Should be able to reduce to non-DataArray
                push!(outputs, outputs[1].data)
            end

            for r in outputs
                @test_da_approx_eq sum!(r, Areduc; skipna=skipna) safe_mapslices(sum, Areduc, region, skipna)
                @test_da_approx_eq prod!(r, Areduc; skipna=skipna) safe_mapslices(prod, Areduc, region, skipna)
                if !has_na
                    @test_da_approx_eq maximum!(r, Areduc; skipna=skipna) safe_mapslices(maximum, Areduc, region, skipna)
                    @test_da_approx_eq minimum!(r, Areduc; skipna=skipna) safe_mapslices(minimum, Areduc, region, skipna)
                end
                @test_da_approx_eq Base.sumabs!(r, Areduc; skipna=skipna) safe_mapslices(sum, abs(Areduc), region, skipna)
                @test_da_approx_eq Base.sumabs2!(r, Areduc; skipna=skipna) safe_mapslices(sum, abs2(Areduc), region, skipna)
                @test_da_approx_eq mean!(r, Areduc; skipna=skipna) safe_mapslices(mean, Areduc, region, skipna)
            end

            @test_da_approx_eq sum(Areduc, region; skipna=skipna) safe_mapslices(sum, Areduc, region, skipna)
            @test_da_approx_eq prod(Areduc, region; skipna=skipna) safe_mapslices(prod, Areduc, region, skipna)
            @test_da_approx_eq maximum(Areduc, region; skipna=skipna) safe_mapslices(maximum, Areduc, region, skipna)
            @test_da_approx_eq minimum(Areduc, region; skipna=skipna) safe_mapslices(minimum, Areduc, region, skipna)
            @test_da_approx_eq Base.sumabs(Areduc, region; skipna=skipna) safe_mapslices(sum, abs(Areduc), region, skipna)
            @test_da_approx_eq Base.sumabs2(Areduc, region; skipna=skipna) safe_mapslices(sum, abs2(Areduc), region, skipna)
            @test_da_approx_eq mean(Areduc, region; skipna=skipna) safe_mapslices(mean, Areduc, region, skipna)

            if region != 5
                @test_da_approx_eq var(Areduc, region; skipna=skipna) safe_mapslices(var, Areduc, region, skipna)
                @test_da_approx_eq var(Areduc, region; mean=0, skipna=skipna) safe_mapslices(myvarzm, Areduc, region, skipna)
                for r in outputs
                    @test_da_approx_eq var(Areduc, region; mean=fill!(r, 1), skipna=skipna) safe_mapslices(myvar1m, Areduc, region, skipna)
                end
            end
        end
    end
end

# Test NA-skipping behavior for maximum
a = @data([NA NA; 3 4])
@test isequal(maximum(a, 1; skipna=true), [3 4])
@test isequal(maximum!(zeros(1, 2), a; skipna=true), [3 4])

# Maximum should give an NA in the output if all values along dimension are NA
@test isequal(maximum(a, 2; skipna=true), @data([NA 4])')
# Maximum should refuse to reduce to a non-DataArray
@test_throws NAException maximum!(zeros(2, 1), a; skipna=true)
end


text:
module LargeCats
export largecat

function similar_ntuple(f, t::NTuple{N,<:Any}) where {N}
    ntuple(f, N)
end

function ifelse_dims(f_in_dims, f_else, t1::Tuple, t2::Tuple, dims)
    map(similar_ntuple(identity, t1), t1, t2) do i, x1, x2
        if i in dims
            f_in_dims(x1, x2)
        else
            f_else(x1, x2)
        end::Int
    end
end

function ifelse_dims(f_in_dims, f_else, c1::CartesianIndex, c2::CartesianIndex, dims)
    t = ifelse_dims(f_in_dims, f_else, Tuple(c1), Tuple(c2), dims)
    CartesianIndex(t)
end

firstarg(x1, x2) = x1
secondarg(x1, x2) = x2
samearg(x1, x2) = (@assert x1 === x2; x1)
@noinline function largecat!(out::AbstractArray{N}, itr; dims) where {N}
    Base.require_one_based_indexing(out)
    I_start = first(CartesianIndices(out))
    I_outsize = CartesianIndex(size(out))
    I_stop = ifelse_dims((i1, i2) -> i1 - 1, secondarg, I_start, I_outsize, dims)
    for x in itr
        Base.require_one_based_indexing(x)
        I_size = CartesianIndex(size(x))
        I_stop = ifelse_dims(+, firstarg, I_stop, I_size, dims)
        # @show I_size
        # @show I_start
        # @show I_stop
        if length(x) < 100
            out[I_start:I_stop] .= x # this one is faster for lots of very small arrays
        else
            copyto!(out, I_start:I_stop, x, CartesianIndices(x))
        end
        I_start = ifelse_dims(+, firstarg, I_start, I_size, dims)
    end
    out
end

_allocate_largecat_output(T::Type{<:Number}, outsize) = zeros(T, outsize)
_allocate_largecat_output(::Type{T}, outsize) where {T} = Array{T}(undef, outsize)

function allocate_largecat_output(itr, dims)
    T = largecat_eltype(itr)
    outsize = largecat_outsize(itr, dims = dims)
    N = length(outsize)
    _allocate_largecat_output(T, outsize)
end

function largecat(itr; dims)
    out = allocate_largecat_output(itr, dims)
    largecat!(out, itr, dims = dims)
end

function largecat_outsize(itr; dims)
    dims1 = size(first(itr))
    out_size = map(similar_ntuple(identity, dims1), dims1) do i, d
        if i in dims
            0
        else
            d
        end
    end
    for x in itr
        Base.require_one_based_indexing(x)
        out_size = ifelse_dims(+, samearg, out_size, size(x), dims)
    end
    return out_size
end

function largecat_eltype(itr)
    T = eltype(eltype(itr))
    if Base.isconcretetype(T)
        return T
    else
        compute_largecat_eltype(itr)
    end
end

function compute_largecat_eltype(itr)
    mapreduce(eltype, promote_type, itr)
end

end #module


text:
module ImmersedBoundaries

using Adapt

using Oceananigans.Grids
using Oceananigans.Operators
using Oceananigans.Fields
using Oceananigans.Utils
using Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, time_discretization

using Oceananigans.TurbulenceClosures:
    viscous_flux_ux,
    viscous_flux_uy,
    viscous_flux_uz,
    viscous_flux_vx,
    viscous_flux_vy,
    viscous_flux_vz,
    viscous_flux_wx,
    viscous_flux_wy,
    viscous_flux_wz,
    diffusive_flux_x,
    diffusive_flux_y,
    diffusive_flux_z

using Oceananigans.Advection:
    advective_momentum_flux_Uu,
    advective_momentum_flux_Uv,
    advective_momentum_flux_Uw,
    advective_momentum_flux_Vu,
    advective_momentum_flux_Vv,
    advective_momentum_flux_Vw,
    advective_momentum_flux_Wu,
    advective_momentum_flux_Wv,
    advective_momentum_flux_Ww,
    advective_tracer_flux_x,
    advective_tracer_flux_y,
    advective_tracer_flux_z

import Oceananigans.Utils: cell_advection_timescale
import Oceananigans.Grids: with_halo
import Oceananigans.Coriolis: φᶠᶠᵃ
import Oceananigans.Grids: with_halo, xnode, ynode, znode, all_x_nodes, all_y_nodes, all_z_nodes

import Oceananigans.Advection:
    _advective_momentum_flux_Uu,
    _advective_momentum_flux_Uv,
    _advective_momentum_flux_Uw,
    _advective_momentum_flux_Vu,
    _advective_momentum_flux_Vv,
    _advective_momentum_flux_Vw,
    _advective_momentum_flux_Wu,
    _advective_momentum_flux_Wv,
    _advective_momentum_flux_Ww,
    _advective_tracer_flux_x,
    _advective_tracer_flux_y,
    _advective_tracer_flux_z

import Oceananigans.TurbulenceClosures:
    _viscous_flux_ux,
    _viscous_flux_uy,
    _viscous_flux_uz,
    _viscous_flux_vx,
    _viscous_flux_vy,
    _viscous_flux_vz,
    _viscous_flux_wx,
    _viscous_flux_wy,
    _viscous_flux_wz,
    _diffusive_flux_x,
    _diffusive_flux_y,
    _diffusive_flux_z,
    κᶠᶜᶜ,
    κᶜᶠᶜ,
    κᶜᶜᶠ,
    νᶜᶜᶜ,
    νᶠᶠᶜ,
    νᶜᶠᶠ,
    νᶠᶜᶠ

export AbstractImmersedBoundary

"""
AbstractImmersedBoundary

Abstract supertype for immersed boundary grids.
"""
abstract type AbstractImmersedBoundary end

struct ImmersedBoundaryGrid{FT, TX, TY, TZ, G, I} <: AbstractGrid{FT, TX, TY, TZ}
    grid :: G
    immersed_boundary :: I

    function ImmersedBoundaryGrid{TX, TY, TZ}(grid::G, ib::I) where {TX, TY, TZ, G <: AbstractUnderlyingGrid, I}
        FT = eltype(grid)
        return new{FT, TX, TY, TZ, G, I}(grid, ib)
    end
end

function ImmersedBoundaryGrid(grid, ib)
    @warn "ImmersedBoundaryGrid is unvalidated and may produce incorrect results. " *
              "Don't hesitate to help validate ImmersedBoundaryGrid by reporting any bugs " *
              "or unexpected behavior to https://github.com/CliMA/Oceananigans.jl/issues"

    TX, TY, TZ = topology(grid)
    return ImmersedBoundaryGrid{TX, TY, TZ}(grid, ib)
end

const IBG = ImmersedBoundaryGrid

@inline Base.getproperty(ibg::IBG, property::Symbol) = get_ibg_property(ibg, Val(property))
@inline get_ibg_property(ibg::IBG, ::Val{property}) where property = getfield(getfield(ibg, :grid), property)
@inline get_ibg_property(ibg::IBG, ::Val{:immersed_boundary}) = getfield(ibg, :immersed_boundary)
@inline get_ibg_property(ibg::IBG, ::Val{:grid}) = getfield(ibg, :grid)

Adapt.adapt_structure(to, ibg::IBG{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} =
    ImmersedBoundaryGrid{TX, TY, TZ}(adapt(to, ibg.grid), adapt(to, ibg.immersed_boundary))

with_halo(halo, ibg::ImmersedBoundaryGrid) = ImmersedBoundaryGrid(with_halo(halo, ibg.grid), ibg.immersed_boundary)

@inline cell_advection_timescale(u, v, w, ibg::ImmersedBoundaryGrid) = cell_advection_timescale(u, v, w, ibg.grid)
@inline φᶠᶠᵃ(i, j, k, ibg::ImmersedBoundaryGrid) = φᶠᶠᵃ(i, j, k, ibg.grid)

@inline xnode(LX, LY, LZ, i, j, k, ibg::ImmersedBoundaryGrid) = xnode(LX, LY, LZ, i, j, k, ibg.grid)
@inline ynode(LX, LY, LZ, i, j, k, ibg::ImmersedBoundaryGrid) = ynode(LX, LY, LZ, i, j, k, ibg.grid)
@inline znode(LX, LY, LZ, i, j, k, ibg::ImmersedBoundaryGrid) = znode(LX, LY, LZ, i, j, k, ibg.grid)

all_x_nodes(loc, ibg::ImmersedBoundaryGrid) = all_x_nodes(loc, ibg.grid)
all_y_nodes(loc, ibg::ImmersedBoundaryGrid) = all_y_nodes(loc, ibg.grid)
all_z_nodes(loc, ibg::ImmersedBoundaryGrid) = all_z_nodes(loc, ibg.grid)

include("immersed_grid_metrics.jl")
include("grid_fitted_immersed_boundary.jl")
include("mask_immersed_field.jl")

#####
##### Diffusivities (for VerticallyImplicitTimeDiscretization)
#####

for (locate_coeff, loc) in ((:κᶠᶜᶜ, (f, c, c)),
                            (:κᶜᶠᶜ, (c, f, c)),
                            (:κᶜᶜᶠ, (c, c, f)),
                            (:νᶜᶜᶜ, (c, c, c)),
                            (:νᶠᶠᶜ, (f, f, c)),
                            (:νᶠᶜᶠ, (f, c, f)),
                            (:νᶜᶠᶠ, (c, f, f)))

    @eval begin
        @inline $locate_coeff(i, j, k, ibg::IBG{FT}, coeff) where FT =
            ifelse(solid_cell(loc..., i, j, k, ibg), $locate_coeff(i, j, k, ibg.grid, coeff), zero(FT))
    end
end

end # module


text:
#=
     canonical(n, d)    n, d   two integers
     canonical(n__d)    n//d   one rational

     Map an ordered pair of integers in their given ratio
     to an ordered pair of integers, preserving their ratio
     while proportional reduction finds the ordered pair of
     integers with the given ratio and of least magnitude.
=#

#=
    System Rationals are maintained in lowest terms, no additional work needed
    FastRational{T,IsReduced} is already in lowest terms, no additional work needed
    FastRational{T,MayReduce} may or may not be in lowest terms, must apply work
=#

@inline canonical(q::Rational{T}) where{T} = q
@inline canonical(q::FastRational{T,IsReduced}) where{T} = q
@inline canonical(q::FastRational{T,MayReduce}) where{T} = canonical(q.num, q.den)

"""
    canonical(num::T, den::T) where T<:Union{Signed, Unsigned}

presumes
  num, den are "normative": den != 0, both are finite
  num, den share an integer type, Signed or Unsigned

provides
  num, den where either
      1 == |num| == |den|
  or
           |num| != |den|
      and
           either both are odd or
           one is odd, one even

assures
  num, den have no common factors
  den is strictly positive
  num is negative onlyif num/den < 0

invariant
  The integer type given num, den is carried in this
  functional action and is carried out unchanged.
""" canonical

function canonical(num::T, den::T) where {T<:Unsigned}
    num, den = canonical_valued(num, den)
    return num, den
end

function canonical(num::T, den::T) where {T<:Signed}
    num, den = canonical_signed(num, den)
    num, den = canonical_valued(num, den)
    return num, den
end

@inline function canonical_signed(num::T, den::T) where {T<:Signed}
    return flipsign(num, den), abs(den)
end

@inline function canonical_valued(num::T, den::T) where {T<:Signed}
    gcdval = gcd(num, den)
    gcdval === one(T) && return num, den
    num = div(num, gcdval)
    den = div(den, gcdval)
    return num ÷ den
end


text:
# Copyright 2017 Eric S. Tellez
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

using KernelMethods
using KernelMethods: precision

using SimilaritySearch
# import KernelMethods.Nets: KlusterClassifier
using Test

@testset "Scores" begin
    @test accuracy([1,1,1,1,1], [1,1,1,1,1]) == 1.0
    @test accuracy([1,1,1,1,1], [0,0,0,0,0]) == 0.0
    @test accuracy([1,1,1,1,0], [0,1,1,1,1]) == 0.6
    @test precision_recall([0,1,1,1,0,1], [0,1,1,1,1,1]) == (precision=0.8333333333333334, recall=0.8333333333333334,
        per_class=Dict(0 => (precision=1.0, recall=0.5, population=2), 1 => (precision=0.8, recall=1.0, population=4)))
    @test precision([0,1,1,1,0,1], [0,1,1,1,1,1]) == 0.9
    @test recall([0,1,1,1,0,1], [0,1,1,1,1,1]) == 0.75
    @test precision([0,1,1,1,0,1], [0,1,1,1,1,1], weight=:weighted) == (1.0 * 2/6 + 0.8 * 4/6) / 2
    @test recall([0,1,1,1,0,1], [0,1,1,1,1,1], weight=:weighted) == (0.5 * 2/6 + 1.0 * 4/6) / 2
    @test f1([0,1,1,1,0,1], [0,1,1,1,1,1], weight=:macro) ≈ (2 * 0.5 / 1.5 + 2 * 0.8 / 1.8) / 2
    #@show f1([0,1,1,1,0,1], [0,1,1,1,1,1], weight=:weighted) # ≈ (2/6 * 2 * 0.5 / 1.5 + 4 / 6 * 2 * 0.8 / 1.8) / 2
end

@testset "CrossValidation" begin
    data = collect(1:100)
    function f(train_X, train_y, test_X, test_y)
        @test train_X == train_y
        @test test_X == test_y
        @test length(train_X ∩ test_X) == 0
        @test length(train_X ∪ test_X) >= 99
        1
    end
    @test montecarlo(f, data, data, runs=10) |> sum == 10
    @test kfolds(f, data, data, folds=10, shuffle=true) |> sum == 10
end

include("oneclass.jl")
include("loaddata.jl")

@testset "KNN" begin
    X, y = loadiris()
    nnc = NearNeighborClassifier(l2_distance, X, y)
    @test optimize!(nnc, accuracy, runs=5, trainratio=0.2, testratio=0.2)[1][1] > 0.8
    @test optimize!(nnc, accuracy, runs=5, trainratio=0.3, testratio=0.3)[1][1] > 0.8
    @test optimize!(nnc, accuracy, runs=5, trainratio=0.7, testratio=0.3)[1][1] > 0.8

    @test optimize!(nnc, accuracy, folds=2)[1][1] > 0.8
    @test optimize!(nnc, accuracy, folds=3)[1][1] > 0.8
    @test optimize!(nnc, accuracy, folds=5)[1][1] > 0.85
    @test optimize!(nnc, accuracy, folds=10)[1][1] > 0.85
    @show optimize!(nnc, accuracy, folds=5)
    @test sum([maximum(x) for x in predict_proba(nnc, X, smoothing=0)])/ length(X) > 0.8 ## close to have all ones, just in case
    @test sum([maximum(x) for x in predict_proba(nnc, X, smoothing=0.01)])/ length(X) > 0.8 ## close to have all ones, just in case
end

@testset "NB" begin
    X, y = loadiris()
    nbc = NaiveBayesClassifier(X, y)
    @test optimize!(nbc, X, y, accuracy, runs=5, trainratio=0.2, testratio=0.2)[1][1] > 0.8
    @test optimize!(nbc, X, y, accuracy, runs=5, trainratio=0.3, testratio=0.3)[1][1] > 0.8
    @test optimize!(nbc, X, y, accuracy, runs=5, trainratio=0.5, testratio=0.5)[1][1] > 0.8
    @test optimize!(nbc, X, y, accuracy, runs=5, trainratio=0.7, testratio=0.3)[1][1] > 0.8
    @show optimize!(nbc, X, y, accuracy, runs=5, trainratio=0.7, testratio=0.3)
end

include("knnreg.jl")
include("kmap.jl")

include("kclass.jl")


text:
"""
This file stores the major front-end functions of ClimateSatellite.jl that are important for
raw data downloading and extraction.

Current functionalities stored in this file include:
- Downloading of satellite data (general-purpose function that calls the specific backend)
- Saving of satellite data into NetCDF files
- Extraction of downloaded satellite data for the following
    - Domain-name specific (previously defined in ClimateEasy)
    - Point location (nearest-neighbour)
    - Within a given boundary found in a specific domain

"""

function clisatdownload(
    productID::AbstractString, date::TimeType;
    email::AbstractString="", path::AbstractString="",
    regions::Array{<:AbstractString,1}=["GLB"],
    overwrite::Bool=false
)

    if path == ""; dataroot = clisatroot(productID);
    else;          dataroot = clisatroot(productID,path);
    end

    info = Dict{Any,Any}("root"=>dataroot,"email"=>replace(email,"@"=>"%40"));
    clisatinfo!(info,productID);

    if info["source"] == "PMM"; checkemail(email)
        if     isprod(info,"gpm");  gpmdwn(regions,date,info,overwrite=overwrite);
        elseif isprod(info,"3b42"); trmmdwn(regions,date,info,overwrite=overwrite);
        end
    elseif info["source"] == "MIMIC"; mimicdwn(regions,date,info,overwrite=overwrite);
    elseif info["source"] == "RSS"; checkemail(email)
        if     isprod(info,"trmm"); rtmidwn(regions,date,info,email,overwrite=overwrite);
        elseif isprod(info,"gpm");  rgmidwn(regions,date,info,email,overwrite=overwrite);
        elseif isprod(info,"smif"); rsmidwn(regions,date,info,email,overwrite=overwrite);
        elseif isprod(info,"wind"); rwnddwn(regions,date,info,email,overwrite=overwrite);
        elseif isprod(info,"amsr"); rmsrdwn(regions,date,info,email,overwrite=overwrite);
        end
    end

end

function clisatrawsave(
    data::Array{<:Real}, grid::Vector{Any},
    region::AbstractString, info::Dict, date::TimeType
)

    @info "$(Dates.now()) - Saving $(info["source"]) $(info["product"]) data for the $(gregionfullname(region)) region..."

    fol = clisatrawfol(info,date,region);
    fnc = joinpath(fol,clisatrawname(info,date,region)); lon,lat,t,tunit = grid;
    nlon,nlat,nt = size(data); nvar = length(info["variable"]);
    yrstr = @sprintf("%04d",Dates.year(date)); mostr = Dates.month(date);

    if nlon != length(lon)
        error("$(Dates.now()) - nlon is $(nlon) but lon contains $(length(lon)) elements")
    end
    if nlat != length(lat)
        error("$(Dates.now()) - nlat is $(nlat) but lat contains $(length(lat)) elements")
    end

    var_var = Vector{AbstractString}(undef,nvar)
    att_var = map(x->Dict(),1:nvar)

    for ii = 1 : nvar
        var_var[ii]                  = info["varID"][ii];
        att_var[ii]["units"]         = info["units"][ii];
        att_var[ii]["standard_name"] = info["standard"][ii];
        att_var[ii]["long_name"]     = info["variable"][ii];
        att_var[ii]["scale_factor"]  = info["scale"][ii];
        att_var[ii]["add_offset"]    = info["offset"][ii];
        att_var[ii]["missing_value"] = -32768;
    end

    att_lon = Dict("units"=>"degrees_east","long_name"=>"longitude");
    att_lat = Dict("units"=>"degrees_north","long_name"=>"latitude");
    att_t   = Dict("calendar"=>"gregorian","long_name"=>"time",
                   "units"=>"$(tunit) since $(yrstr)-$(mostr)-1 0:0:0");

    if isfile(fnc)
        @info "$(Dates.now()) - Stale NetCDF file $(fnc) detected.  Overwriting ..."
        rm(fnc);
    end

    ## Write data here

    @debug "$(Dates.now()) - Creating $(info["source"]) $(info["product"]) netCDF file $(fnc) ..."

    ds = Dataset(fnc,"c");
    ds.dim["longitude"] = nlon; ds.dim["latitude"] = nlat; ds.dim["time"] = nt;

    @debug "$(Dates.now()) - Saving $(info["source"]) $(info["product"]) data to netCDF file $(fnc) ..."

    if nvar != 1
        for ii = 1 : nvar
            v = defVar(ds,var_var[ii],Int16,("longitude","latitude","time"),
                       attrib=att_var[ii]);
            v.var[:] = data[:,:,:,ii];
        end
    else;
        v = defVar(ds,var_var[1],Int16,("longitude","latitude","time"),attrib=att_var[1]);
        v.var[:] = data;
    end

    defVar(ds,"longitude",lon,("longitude",),attrib=att_lon)
    defVar(ds,"latitude",lat,("latitude",),attrib=att_lat)
    defVar(ds,"time",t,("time",),attrib=att_t)

    close(ds);

    ## Write data end

    @info "$(Dates.now()) - $(info["source"]) $(info["product"]) data for the $(gregionfullname(region)) region has been saved into file $(fnc)"

end

# Root Functions
function clisatrawregion(
    productID::AbstractString, varname::AbstractString,
    start::TimeType, finish::TimeType;
    path::AbstractString="",
    region::AbstractString="GLB",
    unpack::Bool=true
)

    if path == ""; dataroot = clisatroot(productID);
    else;          dataroot = clisatroot(productID,path);
    end
    isgeoregion(region);

    info = Dict{Any,Any}("root"=>dataroot); clisatinfo!(info,productID);

    @info "$(Dates.now()) - Extracting $(info["source"]) $(info["product"]) data for the entire $(gregionfullname(region)) region ..."

    if !isdir(clisatregfol(info,region))
        error("$(Dates.now()) - No data has been downloaded from $(info["source"]) $(info["product"]) in the $(gregionfullname(region))")
    end

    clisatvarinfo!(info,productID,varname=varname);
    dvec,dys,dyf,ndy = clisatextractdate(start,finish);
    nt = info["dayfreq"]; ndates = length(dvec);

    lon,lat = clisatlonlat(info); rlon,rlat,rinfo = gregiongridvec(region,lon,lat);
    nlon = length(rlon); nlat = length(rlat);
    data = zeros(Int16,nlon,nlat,ndy*nt);

    for ii = 1 : ndates; dateii = dvec[ii];

        fol = clisatrawfol(info,dateii,region);
        if !isdir(fol)
            error("$(Dates.now()) - There is no data for $(info["source"]) $(info["product"]) for $(yrmo2dir(dateii)).")
        end

        fnc = joinpath(fol,clisatrawname(info,dateii,region));
        ds  = Dataset(fnc,"r"); vds = ds[varname];

        if     ii == 1 && ii != ndates
            moday = daysinmonth(dateii);
            ibeg = (dys-1)*nt + 1; iend = (moday+1-dys)*nt;
            data[:,:,1:iend] = vds.var[:,:,ibeg:end];
        elseif ii != 1 && ii == ndates
            ibeg = iend+1; iend = dyf*nt;
            data[:,:,ibeg:end] = vds.var[:,:,1:iend];
        elseif ii == 1 && ii == ndates
            ibeg = (dys-1)*nt + 1; iend = dyf*nt;
            data = vds.var[:,:,ibeg:iend];
        else
            moday = daysinmonth(dateii);
            ibeg = iend+1; iend = ibeg-1 + moday*nt;
            data[:,:,ibeg:iend] = vds.var[:];
        end

        close(ds);

    end

    @info "$(Dates.now()) - $(info["source"]) $(info["product"]) data for the entire $(gregionfullname(region)) region has been extracted."

    if !unpack; return datavec,info,[rlon,rlat]
    else; offset = info["offset"]; scale = info["scale"];
          return datavec.*scale.+offset,info,[rlon,rlat]
    end

end

function clisatrawpoint(
    productID::AbstractString, varname::AbstractString,
    start::TimeType, finish::TimeType;
    coord::Array{<:Real,1},
    path::AbstractString="",
    region::AbstractString="GLB",
    unpack::Bool=true
)

    if path == ""; dataroot = clisatroot(productID);
    else;          dataroot = clisatroot(productID,path);
    end
    isgeoregion(region);

    if length(coord) != 2
        error("$(Dates.now()) - Coordinate vector must be in the form [lon,lat]")
    end

    info = Dict{Any,Any}("root"=>dataroot); clisatinfo!(info,productID);

    @info "$(Dates.now()) - Extracting $(info["source"]) $(info["product"]) data at coordinates $(coord) ..."

    if !isdir(clisatregfol(info,region))
        error("$(Dates.now()) - No data has been downloaded from $(info["source"]) $(info["product"]) in the $(gregionfullname(region))")
    end

    clisatvarinfo!(info,productID,varname=varname);
    dvec,dys,dyf,ndy = clisatextractdate(start,finish);
    nt = info["dayfreq"]; ndates = length(dvec);

    lon,lat = clisatlonlat(info); rlon,rlat,rinfo = gregiongridvec(region,lon,lat);
    plon,plat = coord; ilon,ilat = regionpoint(plon,plat,rlon,rlat);
    data = zeros(Int16,ndy*nt);

    for ii = 1 : ndates; dateii = dvec[ii];

        fol = clisatrawfol(info,dateii,region);
        if !isdir(fol)
            error("$(Dates.now()) - There is no data for $(info["source"]) $(info["product"]) for $(yrmo2dir(dateii)).")
        end

        fnc = joinpath(fol,clisatrawname(info,dateii,region));
        ds  = Dataset(fnc,"r"); vds = ds[varname];

        if     ii == 1 && ii != ndates
            moday = daysinmonth(dateii);
            ibeg = (dys-1)*nt + 1; iend = (moday+1-dys)*nt;
            data[1:iend] = vds.var[ilon,ilat,ibeg:end];
        elseif ii != 1 && ii == ndates
            ibeg = iend+1; iend = dyf*nt;
            data[ibeg:end] = vds.var[ilon,ilat,1:iend];
        elseif ii == 1 && ii == ndates
            ibeg = (dys-1)*nt + 1; iend = dyf*nt;
            data = vds.var[ilon,ilat,ibeg:iend];
        else
            moday = daysinmonth(dateii);
            ibeg = iend+1; iend = ibeg-1 + moday*nt;
            data[ibeg:iend] = vds.var[ilon,ilat,:];
        end

        close(ds);

    end

    @info "$(Dates.now()) - $(info["source"]) $(info["product"]) data for the coordinates $(coord) has been extracted."

    if !unpack; return datavec,info
    else; offset = info["offset"]; scale = info["scale"];
          return datavec.*scale.+offset,info
    end

end

function clisatrawgrid(
    productID::AbstractString, varname::AbstractString,
    start::TimeType, finish::TimeType;
    grid::Array{<:Real,1},
    path::AbstractString="",
    region::AbstractString="GLB",
    unpack::Bool=true
)

    if path == ""; dataroot = clisatroot(productID);
    else;          dataroot = clisatroot(productID,path);
    end
    isgeoregion(region);

    if length(grid) != 4
        error("$(Dates.now()) - Grid vector must be in the form [N,S,E,W]")
    end

    info = Dict{Any,Any}("root"=>dataroot); clisatinfo!(info,productID);

    @info "$(Dates.now()) - Extracting $(info["source"]) $(info["product"]) data for the [N,S,E,W] bounds $(grid) ..."

    if !isdir(clisatregfol(info,region))
        error("$(Dates.now()) - No data has been downloaded from $(info["source"]) $(info["product"]) in the $(gregionfullname(region))")
    end

    clisatvarinfo!(info,productID,varname=varname);
    dvec,dys,dyf,ndy = clisatextractdate(start,finish);
    nt = info["dayfreq"]; ndates = length(dvec);

    lon,lat = clisatlonlat(info); isgridinregion(grid,region);
    rlon,rlat,rinfo = gregiongridvec(region,lon,lat);
    glon,glat,ginfo = regiongridvec(grid,rlon,rlat); iWE,iNS = ginfo["IDvec"];
    nlon = length(glon); nlat = length(glat);
    data = zeros(Int16,nlon,nlat,ndy*nt);

    for ii = 1 : ndates; dateii = dvec[ii];

        fol = clisatrawfol(info,dateii,region);
        if !isdir(fol)
            error("$(Dates.now()) - There is no data for $(info["source"]) $(info["product"]) for $(yrmo2dir(dateii)).")
        end

        fnc = joinpath(fol,clisatrawname(info,dateii,region));
        ds  = Dataset(fnc,"r"); vds = ds[varname];

        if     ii == 1 && ii != ndates
            moday = daysinmonth(dateii);
            ibeg = (dys-1)*nt + 1; iend = (moday+1-dys)*nt;
            data[:,:,1:iend] = vds.var[iWE,iNS,ibeg:end];
        elseif ii != 1 && ii == ndates
            ibeg = iend+1; iend = dyf*nt;
            data[:,:,ibeg:end] = vds.var[iWE,iNS,1:iend];
        elseif ii == 1 && ii == ndates
            ibeg = (dys-1)*nt + 1; iend = dyf*nt;
            data = vds.var[iWE,iNS,ibeg:iend];
        else
            moday = daysinmonth(dateii);
            ibeg = iend+1; iend = ibeg-1 + moday*nt;
            data[:,:,ibeg:iend] = vds.var[iWE,iNS,:];
        end

        close(ds);

    end

    @info "$(Dates.now()) - $(info["source"]) $(info["product"]) data within the [N,S,E,W] bounds $(grid) has been extracted."

    if !unpack; return datavec,info,[glon,glat]
    else; offset = info["offset"]; scale = info["scale"];
          return datavec.*scale.+offset,info,[glon,glat]
    end

end


text:
"""
Intelligent driver model
"""

function IDM(;Vego::Float64, Vfront::Float64, Vref::Float64, Snet::Float64, T::Float64=0.1, Amax::Float64, Bdec::Float64, Smin::Float64=0.1, δ::Float64=3.0)
    """
        Vego:   Velocity of ego car;
        Vfront: Velocity of front car;
        Vref:   the velocity the vehicle would drive at in free traffic;
        Snet:   Distance between ego and front car;
        Smin:   a minimum desired net distance;
        T:      the minimum possible time to the vehicle in front, or Safe time headway;
        Amax:   the maximum vehicle acceleration of ego car;
        Bdec:   comfortable braking deceleration, must be positive;
        δ:      just a exponent, usually set to 4.
    """
    ΔV = Vego - Vfront
    Sstar = Smin + Vego*T + Vego*ΔV/(2*sqrt(Amax*Bdec))
    Aego = Amax*(1 - (Vego/Vref)^δ - (Sstar/Snet)^2)
end

#= # Test
@show IDM(Vego=4.0, Vfront=0.0, Vref=4.0, Snet=4.0, Smin=0.1, T=0.1, Amax=3.0, Bdec=1.0)
=#


text:
abstract type AbstractHook end

mutable struct Hook <: AbstractHook
    layer_num::Int
    logg::Array{Real,1}
    count::Int
    function Hook(n,logg;count=0)
        new(n,logg,count)
    end
end

function logg!(x,h::Hook)
    h.count += 1
    h.logg[h.count] = mean(x)
    x
end

function hook_layer(x;h::Hook)
    return logg!(x,h)
end

#todo https://fluxml.ai/Zygote.jl/latest/adjoints/
@adjoint hook_layer(x; kwargs...) = hook_layer(x; kwargs...), b -> (b,)

function make_hooks(num_hooks, n)
    return Dict(i => Hook(i,zeros(n)) for i in num_hooks)
end

function make_new_layers(model, dic_hooks)
    t = ()
    for (i,l) in enumerate(model.layers)
        if haskey(dic_hooks,i)
            t = tuple(t..., l, x -> hook_layer(x,h=dic_hooks[i]))
        else
            t = tuple(t..., l)
        end
    end
    return t
end

function get_model_hooks(model,num_hooks;n=10)
    dic_hooks = make_hooks(num_hooks,n)
    n_layers = make_new_layers(model,dic_hooks)
    return (Chain(n_layers...), dic_hooks)
end


text:
@testset "extrema_enclosure" begin
    @testset "test problems" begin
        demo_problems = demo_problem.(1:length(demo_problems_definitions))

        for (f, (a, b), min, max, absmin, absmax) in demo_problems
            for abs_value in (false, true)
                # Direct evaluation
                min1, max1 = ArbExtras.extrema_enclosure(f, a, b, degree = -1; abs_value)
                min2 = ArbExtras.minimum_enclosure(f, a, b, degree = -1; abs_value)
                max2 = ArbExtras.maximum_enclosure(f, a, b, degree = -1; abs_value)

                # Taylor evaluation
                min3, max3 = ArbExtras.extrema_enclosure(f, a, b; abs_value)
                min4 = ArbExtras.minimum_enclosure(f, a, b; abs_value)
                max4 = ArbExtras.maximum_enclosure(f, a, b; abs_value)

                if abs_value
                    @test Arblib.overlaps(min1, absmin)
                    @test Arblib.overlaps(min2, absmin)
                    @test Arblib.overlaps(min3, absmin)
                    @test Arblib.overlaps(min4, absmin)

                    @test Arblib.overlaps(max1, absmax)
                    @test Arblib.overlaps(max2, absmax)
                    @test Arblib.overlaps(max3, absmax)
                    @test Arblib.overlaps(max4, absmax)
                else
                    @test Arblib.overlaps(min1, min)
                    @test Arblib.overlaps(min2, min)
                    @test Arblib.overlaps(min3, min)
                    @test Arblib.overlaps(min4, min)

                    @test Arblib.overlaps(max1, max)
                    @test Arblib.overlaps(max2, max)
                    @test Arblib.overlaps(max3, max)
                    @test Arblib.overlaps(max4, max)
                end

                # Check that the enclosure satisfies the relative error
                if abs_value && Arblib.contains_zero(min3)
                    # This case is not handled very well at the moment
                    # so it might not return exactly zero for the
                    # minimum, therefore we have to check the absolute
                    # tolerance instead of the relative one.
                    @test Arblib.radius(Arb, min3) <= sqrt(eps(Arb))
                else
                    @test Arblib.radius(Arb, min3) <= sqrt(eps(Arb)) * abs(min3)
                end
                @test Arblib.radius(Arb, max3) <= sqrt(eps(Arb)) * abs(max3)
                if abs_value && Arblib.contains_zero(min4)
                    # This case is not handled very well at the moment
                    # so it might not return exactly zero for the
                    # minimum, therefore we have to check the absolute
                    # tolerance instead of the relative one.
                    @test Arblib.radius(Arb, min4) <= sqrt(eps(Arb))
                else
                    @test Arblib.radius(Arb, min4) <= sqrt(eps(Arb)) * abs(min4)
                end
                @test Arblib.radius(Arb, max4) <= sqrt(eps(Arb)) * abs(max4)
            end
        end
    end

    @testset "special cases" begin
        # Thin interval
        @test ArbExtras.extrema_enclosure(x -> -2x, Arf(2), Arf(2)) == (Arb(-4), Arb(-4))
        @test ArbExtras.minimum_enclosure(x -> -2x, Arf(2), Arf(2)) == Arb(-4)
        @test ArbExtras.maximum_enclosure(x -> -2x, Arf(2), Arf(2)) == Arb(-4)
        @test ArbExtras.extrema_enclosure(x -> -2x, Arf(2), Arf(2), abs_value = true) ==
              (Arb(4), Arb(4))
        @test ArbExtras.minimum_enclosure(x -> -2x, Arf(2), Arf(2), abs_value = true) ==
              Arb(4)
        @test ArbExtras.maximum_enclosure(x -> -2x, Arf(2), Arf(2), abs_value = true) ==
              Arb(4)

        # Non-ordered endpoints
        @test_throws ArgumentError ArbExtras.extrema_enclosure(identity, Arf(2), Arf(1))
        @test_throws ArgumentError ArbExtras.minimum_enclosure(identity, Arf(2), Arf(1))
        @test_throws ArgumentError ArbExtras.maximum_enclosure(identity, Arf(2), Arf(1))

        # Infinite endpoints
        @test_throws ArgumentError ArbExtras.extrema_enclosure(identity, Arf(2), Arf(Inf))
        @test_throws ArgumentError ArbExtras.minimum_enclosure(identity, Arf(2), Arf(Inf))
        @test_throws ArgumentError ArbExtras.maximum_enclosure(identity, Arf(2), Arf(Inf))

        # Threading enabled
        @test Arblib.isequal(
            ArbExtras.extrema_enclosure(
                x -> -2x,
                Arf(2),
                Arf(3),
                degree = -1,
                threaded = true,
            ),
            ArbExtras.extrema_enclosure(x -> -2x, Arf(2), Arf(3), degree = -1),
        )
        @test Arblib.isequal(
            ArbExtras.extrema_enclosure(x -> -2x, Arf(2), Arf(3), threaded = true),
            ArbExtras.extrema_enclosure(x -> -2x, Arf(2), Arf(3)),
        )
        @test Arblib.isequal(
            ArbExtras.minimum_enclosure(
                x -> -2x,
                Arf(2),
                Arf(3),
                degree = -1,
                threaded = true,
            ),
            ArbExtras.minimum_enclosure(x -> -2x, Arf(2), Arf(3), degree = -1),
        )
        @test Arblib.isequal(
            ArbExtras.minimum_enclosure(x -> -2x, Arf(2), Arf(3), threaded = true),
            ArbExtras.minimum_enclosure(x -> -2x, Arf(2), Arf(3)),
        )
        @test Arblib.isequal(
            ArbExtras.maximum_enclosure(
                x -> -2x,
                Arf(2),
                Arf(3),
                degree = -1,
                threaded = true,
            ),
            ArbExtras.maximum_enclosure(x -> -2x, Arf(2), Arf(3), degree = -1),
        )
        @test Arblib.isequal(
            ArbExtras.maximum_enclosure(x -> -2x, Arf(2), Arf(3), threaded = true),
            ArbExtras.maximum_enclosure(x -> -2x, Arf(2), Arf(3)),
        )

        # Verbose = true
        @test_logs (:info, "reached maximum depth 20") match_mode = :any ArbExtras.extrema_enclosure(
            inv,
            Arf(0),
            Arf(1),
            verbose = true,
        )
        @test_logs (:info, "reached maximum depth 20") match_mode = :any ArbExtras.minimum_enclosure(
            inv,
            Arf(0),
            Arf(1),
            verbose = true,
        )
        @test_logs (:info, "reached maximum depth 20") match_mode = :any ArbExtras.maximum_enclosure(
            inv,
            Arf(0),
            Arf(1),
            verbose = true,
        )

        @test_logs (:info, "reached maximum number of evaluations 117 >= 100") match_mode =
            :any ArbExtras.extrema_enclosure(
            sin ∘ exp,
            Arf(0),
            Arf(10),
            maxevals = 100,
            verbose = true,
        )
        @test_logs (:info, "reached maximum number of evaluations 107 >= 100") match_mode =
            :any ArbExtras.minimum_enclosure(
            sin ∘ exp,
            Arf(0),
            Arf(10),
            maxevals = 100,
            verbose = true,
        )
        @test_logs (:info, "reached maximum number of evaluations 111 >= 100") match_mode =
            :any ArbExtras.maximum_enclosure(
            sin ∘ exp,
            Arf(0),
            Arf(10),
            maxevals = 100,
            verbose = true,
        )
    end
end


text:
pxyequal(a,b)=(isa(a,Vector) ? isequal(a,b) : (isequal(a[1],b[1]) && xyequal(a,b)))
xyequal(a,b)=myapprox(sortcols(vcat(hcat(a[2]...),hcat(a[3]...))), sortcols(vcat(hcat(b[2]...),hcat(b[3]...))))
function myapprox(x, y;
                  maxeps::Real = max(eps(eltype(x)), eps(eltype(y))),
                  rtol::Real=maxeps^(1/3), atol::Real=maxeps^(1/2))
    size(x) == size(y) || (warn("myapprox: $(size(x))!=$(size(y))"); return false)
    d = abs(x-y)
    s = abs(x)+abs(y)
    all(d .< (atol + rtol * s))
end

:ok