text:
const FT = Float64
import SciMLBase: step!
using OrdinaryDiffEq: ODEProblem, solve, SSPRK33, savevalues!
import ClimaCore.Utilities: PlusHalf
using DiffEqCallbacks

include("atmos_rhs.jl")
include("ocean_rhs.jl")
include("land_rhs.jl")
include("coupledbc.jl")

# model parameters that the coupler overwrites
const atm_T_ini = FT(270.0)
const MSLP = FT(1e5)
const grav = FT(9.8)
const R_d = FT(287.058)
const γ = FT(1.4)
const C_p = FT(R_d * γ / (γ - 1))
const C_v = FT(R_d / (γ - 1))
const R_m = R_d
cpl_parameters = (
    # atmos parameters
    atm_μ = FT(0.0001), # diffusion coefficient
    atm_T_top = FT(280.0), # fixed temperature at the top of the domain_atm
    atm_T_ini = atm_T_ini, # initial condition of at temperature (isothermal) [K]
    MSLP = MSLP, # mean sea level pressure
    grav = grav, # gravitational constant
    R_d = R_d, # R dry (gas constant / mol mass dry air)
    γ = γ, # heat capacity ratio
    C_p = C_p, # heat capacity at constant pressure
    C_v = C_v, # heat capacity at constant volume
    R_m = R_m, # moist R, assumed to be dry
    # land slab parameters
    lnd_h = FT(0.5), # depth of slab layer [m]
    lnd_ρ = FT(1500), # density [kg m^-3]
    lnd_c = FT(800), # specific heat [J K^-1 kg^-1]
    lnd_T_ini = FT(260.0), # initial condition of at temperature (isothermal) [K]
    # ocean slab parameters
    ocn_h = FT(0.5), # depth of slab layer [m]
    ocn_ρ = FT(1025), # density [kg m^-3]
    ocn_c = FT(3850), # specific heat [J K^-1 kg^-1]
    ocn_T_ini = FT(260.0), # initial condition of at temperature (isothermal) [K]
    # coupling parameters
    C_H = FT(0.0015),
)

@info "Init Models and Maps"

atm_Y_default, atm_bc, atm_domain = atm_init(
    xmin = -500,
    xmax = 500,
    zmin = 0,
    zmax = 1000,
    npoly = 4,
    helem = 20,
    velem = 20,
    bc = (ρθ = (bottom = CoupledFlux(), top = ZeroFlux()),),
)

ocn_Y_default, ocn_domain = ocn_init(xmin = -500, xmax = 0, helem = 10, npoly = 0)

lnd_Y_default, lnd_domain = lnd_init(xmin = 0, xmax = 500, helem = 10, npoly = 0)

atm_boundary = Spaces.level(atm_domain.hv_face_space, PlusHalf(0))

# Build remapping operators
atm_to_ocn = Operators.LinearRemap(ocn_domain, atm_boundary)
atm_to_lnd = Operators.LinearRemap(lnd_domain, atm_boundary)
ocn_to_atm = Operators.LinearRemap(atm_boundary, ocn_domain)
lnd_to_atm = Operators.LinearRemap(atm_boundary, lnd_domain)

function make_dss_func()
    function _dss!(x::Fields.Field)
        Spaces.weighted_dss!(x)
    end
    function _dss!(::Any)
        nothing
    end
    dss_func(Y, t, integrator) = foreach(_dss!, Fields._values(Y))
    return dss_func
end
dss_func = make_dss_func()
dss_callback = FunctionCallingCallback(dss_func, func_start = true)

function cpl_run()

    # timestepping
    t_start, t_end = (0.0, 5e2)
    cpl_Δt = 0.1
    saveat = 1e2
    atm_nsteps, ocn_nsteps, lnd_nsteps = (5, 1, 1)

    # initialize
    atm_T_sfc = Operators.remap(ocn_to_atm, ocn_Y_default.T_sfc) .+ Operators.remap(lnd_to_atm, lnd_Y_default.T_sfc) # masked arrays; regrid to atm grid
    atm_F_sfc = Fields.zeros(atm_boundary)

    atm_Y = Fields.FieldVector(Yc = atm_Y_default.Yc, ρw = atm_Y_default.ρw, F_sfc = atm_F_sfc)
    atm_prob = ODEProblem(
        atm_rhs!,
        atm_Y,
        (t_start, t_end),
        (cpl_p = cpl_parameters, T_sfc = atm_T_sfc, bc = atm_bc, domain = atm_domain),
    )
    atm_integ = init(
        atm_prob,
        SSPRK33(),
        dt = cpl_Δt / atm_nsteps,
        saveat = saveat,
        progress = true,
        progress_message = (dt, u, params, t) -> t,
        callback = dss_callback,
    )

    ocn_F_sfc = Fields.zeros(ocn_domain)

    ocn_Y = Fields.FieldVector(T_sfc = ocn_Y_default.T_sfc)
    ocn_prob = ODEProblem(ocn_rhs!, ocn_Y, (t_start, t_end), (cpl_parameters, F_sfc = ocn_F_sfc))
    ocn_integ = init(ocn_prob, SSPRK33(), dt = cpl_Δt / ocn_nsteps, saveat = saveat)

    lnd_F_sfc = Fields.zeros(lnd_domain)

    lnd_Y = Fields.FieldVector(T_sfc = lnd_Y_default.T_sfc)
    lnd_prob = ODEProblem(lnd_rhs!, lnd_Y, (t_start, t_end), (cpl_parameters, F_sfc = lnd_F_sfc))
    lnd_integ = init(lnd_prob, SSPRK33(), dt = cpl_Δt / lnd_nsteps, saveat = saveat)

    @info "Run model"
    # coupler stepping
    for t in (t_start:cpl_Δt:t_end)

        ## Atmos
        # pre: reset flux accumulator
        atm_F_sfc = atm_integ.u.F_sfc
        atm_F_sfc .= atm_F_sfc .* 0.0    # reset surface flux to be accumulated

        # run
        # NOTE: use (t - integ_atm.t) here instead of Δt_cpl to avoid accumulating roundoff error in our timestepping.
        step!(atm_integ, t - atm_integ.t, true)

        ## Ocean
        # pre: get accumulated flux from atmos
        ocn_F_sfc = ocn_integ.p.F_sfc
        Operators.remap!(ocn_F_sfc, atm_to_ocn, atm_integ.u.F_sfc ./ cpl_Δt)

        # run
        step!(ocn_integ, t - ocn_integ.t, true)
        # post: send ocean surface temp to atmos
        Operators.remap!(atm_integ.p.T_sfc, ocn_to_atm, ocn_integ.u.T_sfc)

        ## Land
        # pre: get accumulated flux from atmos
        lnd_F_sfc = lnd_integ.p.F_sfc
        Operators.remap!(lnd_F_sfc, atm_to_lnd, atm_integ.u.F_sfc ./ cpl_Δt)

        # run
        step!(lnd_integ, t - lnd_integ.t, true)
        # post: send ocean surface temp to atmos
        atm_integ.p.T_sfc .+= Operators.remap(lnd_to_atm, lnd_integ.u.T_sfc)
    end
    @info "Simulation Complete"
    return atm_integ, ocn_integ, lnd_integ
end

atm_integ, ocn_integ, lnd_integ = cpl_run()

sol = atm_integ.sol

# dirname = "sea_breeze_2d"
# path = joinpath(@__DIR__, "output", dirname)
# mkpath(path)

# using JLD2
# save(joinpath(path, "last_sim.jld2"), "atm_integ", atm_integ, "ocn_integ", ocn_integ, "lnd_integ", lnd_integ)

# post-processing
# import Plots, ClimaCorePlots
# Plots.GRBackend()
# interp = 5

# anim = Plots.@animate for u in sol.u
#     Plots.contourf(u.Yc.ρθ ./ u.Yc.ρ)
# end
# Plots.mp4(anim, joinpath(path, "theta.mp4"), fps = 20)

#  If2c = Operators.InterpolateF2C()
#  anim = Plots.@animate for u in sol.u
#      Plots.contourf(If2c.(u.ρw) ./ u.Yc.ρ)
#  end
#  Plots.mp4(anim, joinpath(path, "vel_w.mp4"), fps = 20)

#  anim = Plots.@animate for u in sol.u
#      Plots.contourf(u.Yc.ρuₕ ./ u.Yc.ρ)
#  end
#  Plots.mp4(anim, joinpath(path, "vel_u.mp4"), fps = 20)


text:
function update_global_vars!(
    cache::Cache,
    inputs::SimulationInputs,
    x::AbstractArray{U},
) where {U <: Real}
    for (var, pointer) in get_global_vars_update_pointers(inputs)
        pointer == 0 && continue
        get_global_vars(cache, U)[var] = x[pointer]
    end
    return
end

ext: jl


text:
module TestGuards

using ConcurrentUtils
using Test

function check_guarding(Guard; guarding = guarding)
    ch1 = Channel{Nothing}()
    ch2 = Channel{Nothing}()
    ref = Threads.Atomic{Int}(0)
    g = Guard(ref)
    @sync begin
        Threads.@spawn begin
            guarding(g) do ref
                ref[] = 111
                put!(ch1, nothing)
            end
        end
        Threads.@spawn begin
            put!(ch2, nothing)
            guarding(g) do ref
                ref[] = 222
            end
        end
        take!(ch1)
        sleep(0.01)
        @test ref[] == 111
        take!(ch2)
    end
    @test ref[] == 222
end

test_default_guard_guarding() = check_guarding(Guard)
test_default_guard_guarding_read() = check_guarding(Guard; guarding = guarding_read)
test_read_write_guard_guarding() = check_guarding(ReadWriteGuard)

function test_guarding_read()
    ch1 = Channel{Nothing}()
    ch2 = Channel{Nothing}()
    ch3 = Channel{Nothing}()
    ref = Threads.Atomic{Int}(0)
    g = ReadWriteGuard(ref)
    @sync begin
        Threads.@spawn begin
            guarding_read(g) do _ref
                put!(ch1, nothing)
                put!(ch1, nothing)
            end
        end
        Threads.@spawn begin
            guarding_read(g) do _ref
                put!(ch2, nothing)
                put!(ch2, nothing)
            end
        end
        Threads.@spawn begin
            put!(ch3, nothing)
            guarding(g) do ref
                ref[] = 111
            end
        end

        # Check that read access can be obtained multiple times (otherwise deadlock):
        take!(ch1)
        take!(ch2)

        # Make sure the writer is "about to" get the guard:
        take!(ch3)

        # The writer should not be able to get the lock:
        sleep(0.1)
        @test ref[] == 0

        # Finish the readers.  Now the writer should be able to proceed (otherwise
        # deadlock):
        take!(ch1)
        take!(ch2)
    end
    @test ref[] == 111
end

end  # module


text:
#Additional considerations at the borehole wall

#--------------------------------------
# Activating vfr@vertical boundary
# separating Fluid-PoroElastic media
#--------------------------------------
function update_vr_vfr_1st_vertical(vr,trr,tpp,tzz,trz,vfr,pf,
    vr_org,vfr_org,
    Rho,Rhof,D1,D2,nr,nz,dr,dz,dt,LPML_z,
    Prz_T,Prz_B,
    ir_wall,Flag_AC,Flag_E)
    #Acoustic-Poroelastic vertical boundary
    #Calculating vrw and vr at the boundary

    j=ir_wall #wall
    #main region
    @inbounds Threads.@threads for k=LPML_z+1:nz-LPML_z

        #check if I need to activate vfr
        if(Flag_AC[k,j]==1 && Flag_AC[k,j+1]==0 && Flag_E[k,j+1]==0)#
             #I am accessing [k,j], but the definition of r_now changes with each component!
             r_now=(j-1)*dr+dr/2.0 #for vr (Mittet)
             vr_now=vr_org[k,j]
             trr_f1,trr_b1=trr[k,j+1],trr[k,j]
             trz_f1,trz_b1=trz[k,j],trz[k-1,j] #should be zero

             trr_av=0.5*(trr_f1+trr_b1) #this looks ok according to BC (pf=-trr, now trr_b1 is -pf)
             tpp_av=0.5*(tpp[k,j+1]+tpp[k,j]) #see trr.
#             tpp_av=tpp[k,j+1] #test. assuming tpp is discontinuous

             vfr_now=vfr_org[k,j]
             pf_f1,pf_b1=pf[k,j+1],pf[k,j] #

             rho_av=0.5*(Rho[k,j]+Rho[k,j+1]) #be carefull (normal)
             rhof_av=0.5*(Rhof[k,j]+Rhof[k,j+1]) #be carefull (normal)
             D1_av=0.5*(D1[k,j]+D1[k,j+1]) #be carefull (normal)
             D2_av=0.5*(D2[k,j]+D2[k,j+1]) #be carefull (normal)

             #Normal
             vr[k,j],vfr[k,j]=update_vr_1st_Por(vr_now,vfr_now,
                 trr_f1,trr_b1,trz_f1,trz_b1,trr_av,tpp_av,
                 pf_f1,pf_b1,
                 dz,dr,dt,r_now,rho_av,rhof_av,D1_av,D2_av,0)
             #

         end #if AC-PE BC
    end #k

#
    #Top PML
    @inbounds Threads.@threads for k=2:LPML_z
        #check if I need to activate vfr
        if(Flag_AC[k,j]==1 && Flag_AC[k,j+1]==0 && Flag_E[k,j+1]==0)#
              #I am accessing [k,j], but the definition of r_now changes with each component!
             r_now=(j-1)*dr+dr/2.0 #for vr (Mittet)
             vr_now=vr_org[k,j]
             trr_f1,trr_b1=trr[k,j+1],trr[k,j]
             trz_f1,trz_b1=trz[k,j],trz[k-1,j] #should be zero

             trr_av=0.5*(trr_f1+trr_b1) #this looks ok according to BC (pf=-trr, now trr_b1 is -pf)
             tpp_av=0.5*(tpp[k,j+1]+tpp[k,j]) #see trr.
#             tpp_av=tpp[k,j+1] #test. assuming tpp is discontinuous

             vfr_now=vfr_org[k,j]
             pf_f1,pf_b1=pf[k,j+1],pf[k,j] #

             rho_av=0.5*(Rho[k,j]+Rho[k,j+1]) #be carefull (normal)
             rhof_av=0.5*(Rhof[k,j]+Rhof[k,j+1]) #be carefull (normal)
             D1_av=0.5*(D1[k,j]+D1[k,j+1]) #be carefull (normal)
             D2_av=0.5*(D2[k,j]+D2[k,j+1]) #be carefull (normal)

             vfr_old=vfr_now

             #Normal
             vr_dummy,vfr[k,j]=update_vr_1st_Por(vr_now,vfr_now,
                    trr_f1,trr_b1,trz_f1,trz_b1,trr_av,tpp_av,
                    pf_f1,pf_b1,
                    dz,dr,dt,r_now,rho_av,rhof_av,D1_av,D2_av,0)

              #PML addition (vfr)
              vfr[k,j]=vfr[k,j]+
                     (D2_av/2+(D1_av-rhof_av^2/rho_av)/dt)^(-1)*(
                                       -rhof_av/rho_av*Prz_T[LPML_z-k+1,j]
                                       )

               #PML (vr) calculate elastic one then vfr addition
               vr[k,j]=update_vr(vr_now,trr_f1,trr_b1,trz_f1,trz_b1,trr_av,tpp_av,
                      0,dz,dr,dt,r_now,rho_av,0)
               vr[k,j]=vr[k,j]+dt/rho_av*Prz_T[LPML_z-k+1,j] #elastic+PML
               dert_vfr=(vfr[k,j]-vfr_old)/dt
               vr[k,j]=vr[k,j]-dt*rhof_av/rho_av*dert_vfr #poroelastic PML
             #
         end #if AC-PE BC
    end #k

#
    #Bottom PML
    @inbounds Threads.@threads for k=nz-LPML_z+1:nz-1

        #check if I need to activate vfr
        if(Flag_AC[k,j]==1 && Flag_AC[k,j+1]==0 && Flag_E[k,j+1]==0)#
              #I am always accessing [k,j], but the definition of r_now changes with components!
             r_now=(j-1)*dr+dr/2.0 #for vr (Mittet)
             vr_now=vr_org[k,j]
             trr_f1,trr_b1=trr[k,j+1],trr[k,j]
             trz_f1,trz_b1=trz[k,j],trz[k-1,j] #should be zero

             trr_av=0.5*(trr_f1+trr_b1) #this looks ok according to BC (pf=-trr, now trr_b1 is -pf)
             tpp_av=0.5*(tpp[k,j+1]+tpp[k,j]) #see trr.
#             tpp_av=tpp[k,j+1] #test. assuming tpp is discontinuous

             vfr_now=vfr_org[k,j]
             pf_f1,pf_b1=pf[k,j+1],pf[k,j] #

             rho_av=0.5*(Rho[k,j]+Rho[k,j+1]) #be carefull (normal)
             rhof_av=0.5*(Rhof[k,j]+Rhof[k,j+1]) #be carefull (normal)
             D1_av=0.5*(D1[k,j]+D1[k,j+1]) #be carefull (normal)
             D2_av=0.5*(D2[k,j]+D2[k,j+1]) #be carefull (normal)

             vfr_old=vfr_now

             #Normal
             vr_dummy,vfr[k,j]=update_vr_1st_Por(vr_now,vfr_now,
                    trr_f1,trr_b1,trz_f1,trz_b1,trr_av,tpp_av,
                    pf_f1,pf_b1,
                    dz,dr,dt,r_now,rho_av,rhof_av,D1_av,D2_av,0)

              #PML addition (vfr)
              vfr[k,j]=vfr[k,j]+
                     (D2_av/2+(D1_av-rhof_av^2/rho_av)/dt)^(-1)*(
                                       -rhof_av/rho_av*Prz_B[k-nz+LPML_z,j]
                                       )

               #PML (vr) calculate elastic one then vfr addition
               vr[k,j]=update_vr(vr_now,trr_f1,trr_b1,trz_f1,trz_b1,trr_av,tpp_av,
                      0,dz,dr,dt,r_now,rho_av,0)
               vr[k,j]=vr[k,j]+dt/rho_av*Prz_B[k-nz+LPML_z,j] #elastic PML
               dert_vfr=(vfr[k,j]-vfr_old)/dt
               vr[k,j]=vr[k,j]-dt*rhof_av/rho_av*dert_vfr #poroelastic PML
             #
         end #if AC-PE BC
    end #k

end


text:
# LinearMixedModel
fit(::Type{MixedModel}, f::FormulaTerm, tbl;
    wts = [],
    contrasts = Dict{Symbol,Any}(),
    verbose::Bool = false,
    REML::Bool = false) =
    fit(LinearMixedModel,
        f, tbl, wts = wts, contrasts = contrasts,
        verbose = verbose, REML = REML)
# GeneralizedLinearMixedModel
fit(::Type{MixedModel}, f::FormulaTerm, tbl,
    d::Distribution, l::Link = canonicallink(d);
    wts = [],
    contrasts = Dict{Symbol,Any}(),
    offset = [],
    verbose::Bool = false,
    REML::Bool = false,
    fast::Bool = false,
    nAGQ::Integer = 1) =
    fit(GeneralizedLinearMixedModel,
        f, tbl, d, l, wts = wts, contrasts = contrasts, offset = offset,
        verbose = verbose, fast = fast, nAGQ = nAGQ)
# LinearMixedModel
fit(::Type{MixedModel}, f::FormulaTerm, tbl,
    d::Normal, l::IdentityLink;
    wts = [],
    contrasts = Dict{Symbol,Any}(),
    verbose::Bool = false,
    REML::Bool = false,
    offset = [],
    fast::Bool = false,
    nAGQ::Integer = 1) =
    fit(LinearMixedModel, f, tbl,
        wts = wts, contrasts = contrasts,
        verbose = verbose, REML = REML)


text:
module Bitmanip

using FloatProps

export visfinite, visinf, visnan, visnormal, vissubnormal, vsignbit
export vfabs, vflipsign, vcopysign, vneg
export vilogb, vldexp

@inline function visfinite{T}(x::T)
    e = asInt(T,x) & exponent_mask(T)
    return e != exponent_mask(T)
end
@inline function visinf{T}(x::T)
    em = asInt(T,x) & (exponent_mask(T) | mantissa_mask(T))
    return em == exponent_mask(T)
end
@inline function visnan{T}(x::T)
    e = asInt(T,x) & exponent_mask(T)
    m = asInt(T,x) & mantissa_mask(T)
    return e == exponent_mask(T) && m != mkInt(T,0)
end
@inline function visnormal{T}(x::T)
    e = asInt(T,x) & exponent_mask(T)
    return e != exponent_mask(T) && e != mkInt(T,0)
end
@inline function vissubnormal{T}(x::T)
    e = asInt(T,x) & exponent_mask(T)
    m = asInt(T,x) & mantissa_mask(T)
    return e == mkInt(T,0) && m != mkInt(T,0)
end
@inline function vsignbit{T}(x::T)
    s = asInt(T,x) & signbit_mask(T)
    return s != mkInt(T,0)
end

@inline function vfabs{T}(x::T)
    return asFloat(T, asInt(T,x) & ~signbit_mask(T))
end
@inline function vflipsign{T}(x::T, y::T)
    ix = asInt(T,x)
    iysgn = asInt(T,y) & signbit_mask(T)
    return asFloat(T, ix $ iysgn)
end
@inline function vcopysign{T}(x::T, y::T)
    ixmag = asInt(T,x) & ~signbit_mask(T)
    iysgn = asInt(T,y) & signbit_mask(T)
    return asFloat(T, ixmag | iysgn)
end
@inline function vneg{T}(x::T)
    return asFloat(T, asInt(T,x) $ signbit_mask(T))
end

@inline function vilogb{T}(x::T)
    e = asInt(T,x) & exponent_mask(T) >> mantissa_bits(T) - exponent_offset(T)
    return e
end

@inline function vldexp{T}(x::T, i::Integer)
    i::intType(T)
    # Use direct integer manipulation
    # Extract integer as lowest mantissa bits (highest bits still
    # contain offset, exponent, and sign)
    ix = asInt(T,x)
    # Construct scale factor by setting exponent (this shifts out the
    # highest bits)
    scale = asFloat(T, ix << mantissa_bits(T))
    return x * scale
end
@inline function vldexp{T}(x::T, y::T)
    # Use direct integer manipulation
    # Add a large number to shift the integer bits into the rightmost
    # bits of the mantissa. Also already add the exponent offset that
    # we need below.
    offset = mkFloat(T, mkInt(T,1) << mantissa_bits(T) + exponent_offset(T))
    y = y + offset
    # Extract integer as lowest mantissa bits (highest bits still
    # contain offset, exponent, and sign)
    iy = asInt(T,y)
    # Construct scale factor by setting exponent. This shifts out the
    # highest bits, and shifts the lowest mantissa bits into the
    # exponent. We already added the exponent offset above.
    scale = asFloat(T, iy << mantissa_bits(T))
    r = x * scale
    return r
end

end


text:
### analyzeInput! ###
# DESCRIPTION: analyzes input values of the simulation to assure there is no mistake
# INPUT: r, v, m, g, t_step, t_max, sld
# @param r: NxD matrix where N is the number of bodies and D the number of dimensions,
#           it contains the initial positions of the bodies
# @param v: NxD matrix where N is the number of bodies and D the number of dimensions,
#           it contains the initial velocities of the bodies
# @param m: array of N elements, each containing the mass of the particle
# @param g: gravitational constant
# @param t_step: time-step of the simulation
# @param t_max: specifies when stops the simulation
# @param sld: simulation logic dictionary, a dictionary where all posible options of the
#             simulation can be specified
#    @key bool_energy: if true, energy of the hamiltonian will be calculated every output step
#    @key bool_print_kep: information about the kepler-flow will be displayed in the console
#    @key ignore_H0: the lineal movement of the whole system will be ignored if this boolean is true.
#    @key out_step: the output will be calculated on every 'out_step'-th step
# @return exit: if there are some errors, it's true; otherwise, false

function analyzeInput!(r, v, m, g, t_step, t_max, sld)

    exit = false

    n = size(r)[1]
    if n != size(v)[1] || n != size(m)[1]
        println("Error: Different amount of particles in input data.")
        exit = true
    end

    if size(r)[2] != 3 || size(v)[2] != 3
        println("Error: Initial positions and velocities must have 3 coordenates.")
        exit = true
    elseif length(size(m)) != 1
        println("Error: Each particle must have one value as its mass.")
        exit = true
    end

    if isnan(g) || g == 0.0
        g = 1.0
        println("Warning: 1 is assigned to the gravitational constant by default.")
    end

    if t_step == 0.0 || isnan(t_step)
        println("Error: t_step can't be 0.")
        exit = true
    elseif t_step < 0.0 && t_max > t_step
        t_max = min(t_step, -t_max)
        println("Warning: t_step and t_max must be both positive or negative.")
    elseif t_step > 0.0 && t_max < t_step
        t_max = min(t_step, -t_max)
        println("Warning: t_step and t_max must be both positive or negative.")
    end

    # SIMULATION LOGIC DICTIONARY
    if isa(sld, Dict)

        if haskey(sld, "bool_energy")
            if !isa(sld["bool_energy"], Bool)
                sld["bool_energy"] = false
                println("Warning: bool_energy must be boolean. Default: false.")
            end
        else
            sld["bool_energy"] = false # DEFAULT false
        end

        if haskey(sld, "bool_print_kep")
            if !isa(sld["bool_print_kep"], Bool)
                sld["bool_print_kep"] = false
                println("Warning: bool_print_kep must be boolean. Default: false.")
            end
        else
            sld["bool_print_kep"] = false # DEFAULT false
        end

        if haskey(sld, "ignore_H0")
            if !isa(sld["ignore_H0"], Bool)
                sld["ignore_H0"] = false
                println("Warning: ignore_H0 wasn't well defined. Default: false.")
            end
        else
            sld["ignore_H0"] = false
            println("Warning: ignore_H0 wasn't defined. Default: false.")
        end

        if haskey(sld, "out_step")
            if abs(t_max / sld["out_step"]) < abs(t_step)
                sld["out_step"] = 1
                println("Warning: out_step was out of bounds. All steps will have an output.")
            end
        else
            sld["out_step"] = 1
            println("Warning: out_step was not defined. All steps will have an output.")
        end

    else
        println("Error: sld must be a dictionary.")
        exit = true
    end

    return exit
end

"""
    saveOutput(file::String, r = NaN, v = NaN, e = NaN, time_step=1.0, delimeter::String =" ", digits::Int64 =64)

Saves calculated positions, velocities and/or energies in a .txt file.

# Args

* `file`: the name of the output file (.txt format).
* `r`: OxNxD matrix, where O is the number of outputs, N the number of bodies and D the number of dimensions. It contains the positions of each body on each output time.
* `v`: OxNxD matrix (optional), it contains the velocities of each body on each output time.
* `e`: array of O elements (optional), it contains the energies of the system on each output time.
* `time_step`: specifies the time between two output steps, by default 1.
* `delimeter`: specifies what will be written between values.
* `digits`: specifies, at most, how many decimal digits will be written
"""
function saveOutput(file::String, r = NaN, v = NaN, e = NaN, time_step=1.0, delimeter::String =" ", digits::Int64 =64)

    x = size(r)[1]
    y = size(r)[2]
    z = size(r)[3]

    if !((size(v) == () || x == size(v)[1]) && (x == length(e)) || size(e) == ())
        println("Error: first dimensions of r, v and e don't match.")
        return
    elseif !(size(v) == () || (y == size(v)[2] && z == size(v)[3] && z == 3))
        println("Error: r and v's dimensions don't match.")
        return
    end

    time = 0.0

    f = 0
    if file != ""
        if file[(end-3):end] == ".txt"
            f = open(file, "w")
        else
            println("Warning: filename isn't a txt file. 'output.txt' will be its name.")
            f = open("output.txt")
        end
    else
        return
    end

    for m in 1:x
        if size(r) != ()
            for n in 1:y
                write(f, string("$(round(r[m,n,1];digits=digits))", delimeter,"$(round(r[m,n,2];digits=digits))", delimeter,"$(round(r[m,n,3];digits=digits))",delimeter))
            end
            write(f, "\n")
        end

        if size(v) != ()
            for n in 1:y
                write(f, string("$(round(v[m,n,1];digits=digits))", delimeter,"$(round(v[m,n,2];digits=digits))", delimeter,"$(round(v[m,n,3];digits=digits))",delimeter))
            end
            write(f, "\n")
        end

        if size(e) != ()
            time += time_step
            write(f, "$time $(e[m])\n")
        end
    end

    close(f)

end

"""
    showOutput(r, v, e = NaN, names = "", time_step = 1.0)

Shows on the console the outputs that have been calculated.

# Args

* `r`: OxNxD matrix, where O is the number of outputs, N the number of bodies and D the number of dimensions, it contains the positions of each body on each output time.
* `v`: OxNxD matrix,it contains the velocities of each body on each output time.
* `e`: array of O elements (optional), it contains the energies of the system on each output time.
* `names`: array of N strings (optional), it contains the names of each body.
* `time_step`: specifies the time between two output steps, by default 1.0.
"""
function showOutput(r, v, e = NaN, names = "", time_step = 1.0)
    if !(size(r)[1] == size(v)[1] && (size(e) == () || size(r)[1] == length(e)))
        println("Error: first dimensions of r, v and e don't match.")
        return
    elseif size(r)[2] != size(v)[2] || size(r)[3] != size(v)[3] || size(r)[3] != 3
        println("Error: r and v's dimensions don't match.")
        return
    end

    total_dt = 0
    if names == ""
        names = string.(collect(1:n))
    end
    for i in 1:size(r)[1]
        println("--- Iteration $(i-1) - Time $total_dt ---")
        if size(e) != ()
            println("Total energy: $(e[i])")
        end
        println("NAME: X, Y, Z; Vx, Vy, Vz")
        for j in 1:size(r)[2]
            println("""$(names[j]): $(r[i,j,1]), $(r[i,j,2]), $(r[i,j,3]); $(v[i,j,1]), $(v[i,j,2]), $(v[i,j,3])""")
        end
        println("")
        total_dt += time_step
    end
end

max_stars_repo_name: salanueva/KeplerFlow

max_issues_repo_name: salanueva/KeplerFlow

max_forks_repo_name: salanueva/KeplerFlow


text:
function floorsum(n, m, a, b)
    ans = 0
    if a >= m
        ans += (n-1) * n * (a÷m) ÷ 2
        a %= m
    end
    if b >= m
        ans += n * (b÷m);
        b %= m
    end

    ymax = (a * n + b) ÷ m
    xmax = (ymax * m - b)
    ymax == 0 && return ans
    ans += (n - (xmax + a - 1) ÷a ) * ymax
    ans += floorsum(ymax, a, m, (a - xmax % a) % a)
    return ans
end


text:
"""
    Embed(size::Int, vocab_size::Int)

The Embedding Layer, `size` is the hidden size. `vocab_size` is the number of the vocabulary. Just a wrapper for embedding matrix.
"""
struct Embed{F ,W <: AbstractArray{F}} <: AbstractEmbed{F}
    scale::F
    embedding::W
end

@treelike Embed

Base.size(e::Embed, s...) = size(e.embedding, s...)

Embed(size::Int, vocab_size::Int; scale = one(Float32)) = Embed(Float32(scale), param(randn(Float32, size, vocab_size)))

function (e::Embed)(x::AbstractArray{Int})
    if isone(e.scale)
        gather(e.embedding, x)
    else
        e(x, e.scale)
    end
end

(e::Embed{F})(x, scale) where {F} = gather(e.embedding, x) .* convert(F, scale)

Base.show(io::IO, e::Embed) = print(io, "Embed($(size(e.embedding, 1)))")


text:
using FTPServer
using FTPClient
using Test

@testset "FTPServer.jl" begin
    FTPServer.init()

    @testset "no-ssl" begin
        FTPServer.serve() do server
            opts = (
                :hostname => FTPServer.hostname(server),
                :port => FTPServer.port(server),
                :username => FTPServer.username(server),
                :password => FTPServer.password(server),
            )

            options = RequestOptions(; opts..., ssl=false)
            ctxt, resp = ftp_connect(options)
            @test resp.code == 226
        end
    end
    @testset "ssl - $mode" for mode in (:explicit, :implicit)
        FTPServer.serve(; security=mode) do server
            opts = (
                :hostname => FTPServer.hostname(server),
                :port => FTPServer.port(server),
                :username => FTPServer.username(server),
                :password => FTPServer.password(server),
                :ssl => true,
                :implicit => mode === :implicit,
                :verify_peer => false,
            )

            options = RequestOptions(; opts...)
            # Test implicit/exlicit ftp ssl scheme is set correctly
            @test options.uri.scheme == (mode === :implicit ? "ftps" : "ftpes")
            ctxt, resp = ftp_connect(options)
            @test resp.code == 226
        end
    end

    FTPServer.cleanup()

    @test !isfile(FTPServer.CERT)
    @test !isfile(FTPServer.KEY)
    @test !isdir(FTPServer.HOMEDIR)
end


text:
# Multinomial Stein Distribution
#
# Represents a Multinomial distribution with parameters ntrials, probs

using Distributions: Multinomial
import Distributions

type SteinMultinomial <: SteinDistribution
    # Number of trials
    ntrials::Int64
    # Probability vector, sums to 1
    probs::Array{Float64}
    # Stein factors, length is 1 less than that of probs
    c1::Float64
    c2::Float64
    c3::Float64
end

# Constructor with (c1,c2,c3) = (1,1,1)
SteinMultinomial(ntrials::Int64, probs::Array{Float64}) = SteinMultinomial(
    ntrials,
    probs,
    1.0,
    1.0,
    1.0
)

# Constructor with (c1,c2,c3) = (1,1,1), all categories equiprobable
SteinMultinomial(ntrials::Int64, ncategories::Int64) = SteinMultinomial(
    ntrials,
    [1.0/ncategories for i=1:ncategories],
    1.0,
    1.0,
    1.0
)

# Draw n independent samples from distribution; last column is dropped
function rand(d::SteinMultinomial, n::Int64)
    Distributions.rand(Multinomial(d.ntrials,d.probs),n)'[:,1:end-1]
end

# Joint distribution evaluated at a point y
function jointdistribution(d::SteinMultinomial, y::Array{Float64})
    Distributions.pdf(Multinomial(d.ntrials,d.probs),[y,d.ntrials-sum(y)])
end

# Conditional distribution of jth coordinate of y, given all other coordinates,
# represented as a probability vector
function conddistribution(d::SteinMultinomial, y::Array{Float64,1}, j::Int64)
    conddist = [jointdistribution(d, [y[1:(j-1)], yj, y[(j+1):end]])
                for yj in 0:d.ntrials]
    conddist / sum(conddist)
end

function supportlowerbound(d::SteinMultinomial, j::Int64)
    0
end

function supportupperbound(d::SteinMultinomial, j::Int64)
    d.ntrials
end

function numdimensions(d::SteinMultinomial)
    length(d.probs) - 1
end

# Cumulative distribution function (only valid when X is univariate)


text:
using StatsBase, Statistics

"""
   digitize(data, cal0, cal1, mode=:equal)

  `data`  --  iterable holding floating point values to be digitized

  `cal0`  --  vectors of floating point values corresponding to samples of
  `cal1`      voltages from measurements of the 0 and 1 states

  `mode`  --  mode of operation for thresholding. Supported options are `:equal` and `:max`

  Translate floating point values in `data` to binary outcomes based on
calibration samples for 0 and 1 outcomes. It is assumed that the number of samples
in `cal0` and `cal1` is large enough that the lower and upper quartile estimate are
accurate. The output is a `(ddata, pe0, pe1)`, where

  `ddata` -- digitized version of ddata (0s and 1s as `Int`s)

  `pe0`   -- probability of error when preparing a 0

  `pe1`   -- probability of error when preparing a 1

  `thr`   -- chosen threshold

  The `:equal` mode of operation tries to equalize the conditional error probabilities for the
two possible outcomes, and it is generally less sensitive to statistical fluctuations. The
`:max` mode of operation tries to minimize the sum of error probabilities, but is is much more
sensitive to statistical fluctuations, and is therefore not recommended unless the ensemble
size for the calibrations is very large.

"""
function digitize(data, cal0::Vector{Float64}, cal1::Vector{Float64}; mode = :equal, thr = NaN)
    minA, maxA = extrema([quantile(cal0,[.25,.75]); quantile(cal1,[.25,.75])])
    Arange = range(minA,stop=maxA,length=1000)
    if isnan(thr)
      thr = if mode==:max
              Arange[findmax(abs.(ecdf(cal0)(Arange)-ecdf(cal1)(Arange)))[2]]
            elseif mode==:equal
              Arange[findmin(abs.(ecdf(cal0)(Arange).-1+ecdf(cal1)(Arange)))[2]]
            else
              error("Unsupported digitization mode")
            end
    end

    ddata = data .> thr
    dcal0 = cal0 .> thr
    dcal1 = cal1 .> thr

    if sum(dcal0)/length(cal0) > 1/2
        ddata = .~ddata
        dcal0 = .~dcal0
        dcal1 = .~dcal1
    end

    return round.(Int,ddata), sum(dcal0)/length(cal0), 1.0-sum(dcal1)/length(cal1), thr
end

"""
    digitize_2D(data_A, data_B, cal00, cal01, cal10, cal11; n)

Digitize 2-qubit readout voltages data_A, data_B using the calibrations cal00, cal01, cal10, cal11.
Voltages are binned in a n x n grid
"""
function digitize_2D(data_A, data_B, cal00, cal01, cal10, cal11; n=51)
  k00 = kde((cal00[1],cal00[2]))
  k01 = kde((cal01[1],cal01[2]))
  k10 = kde((cal10[1],cal10[2]))
  k11 = kde((cal11[1],cal11[2]))
  xpoints, ypoints, dig_map = map_2D(k00,k01,k10,k11, n)
  # assign each coordinate (data_A, data_B) to a 2-bit value using dig_map
  dig_data = [dig_map[max(1,searchsortedlast(xpoints, d_A)), max(1, searchsortedlast(ypoints, d_B))] for (d_A, d_B) in zip(data_A, data_B)]
  return dig_data
end

"""
    map_2D(k00, k01, k10, k11, n)

Digitize readout voltages in a n x n grid, using the estimated
distributions for the computational states k00, k01, 01, k11
"""
function map_2D(k00, k01, k10, k11, n; thresholds=(NaN, NaN))
  xmin = minimum([k00.x; k01.x; k10.x; k11.x])
  xmax = maximum([k00.x; k01.x; k10.x; k11.x])
  ymin = minimum([k00.y; k01.y; k10.y; k11.y])
  ymax = maximum([k00.y; k01.y; k10.y; k11.y])
  xpoints = range(xmin,stop=xmax,length=n)
  ypoints = range(ymin,stop=ymax,length=n)
  assignmat = zeros(n,n)
  distrs = [k00 k01 k10 k11]
  for yi in 1:n
    y_ind = zeros(Int,4)
    for k=1:4
      y_ind[k] = max(1,searchsortedlast(distrs[k].y, ypoints[yi]))
    end
    for xi in 1:n
      x_ind = zeros(Int,4)
      for m=1:4
        x_ind[m] = max(1,searchsortedlast(distrs[m].x, xpoints[xi]))
      end
      if any(isnan, thresholds)
        highest_kde_ind =  argmax([distrs[1].density[x_ind[1],y_ind[1]] distrs[2].density[x_ind[2],y_ind[2]] distrs[3].density[x_ind[3],y_ind[3]] distrs[4].density[x_ind[4],y_ind[4]]])
        if distrs[highest_kde_ind].density[x_ind[highest_kde_ind],y_ind[highest_kde_ind]]>0
          assignmat[xi, yi] = highest_kde_ind
        end
      else
        assign_a1 = mean(k00.density) > mean(k01.density) ? Int(xpoints[xi] < thresholds[1]) : Int(xpoints[xi] > thresholds[1])
        assign_a2 = mean(k00.density) > mean(k10.density) ? Int(ypoints[yi] < thresholds[2]) : Int(ypoints[yi] > thresholds[2])
        assignmat[xi, yi] = assign_a1 + 2*assign_a2 + 1
      end
    end
  end
  return xpoints, ypoints, assignmat
end

"""
  smooth_2D_map(digmat; dig_thr)

Smooth a 2D classification map by doing a majority vote over the nearest neighbors.
dig_thr is the minimum number of equal nearest neighbors to flip a state
Example: point digmat[i,j] is surrounded by 2 points equal to 1, 3 equal to 2, 3 equal to 3, 1 equal to 4.
If dig_thr <=3, digmat[i,j] will be set to a value chosen randomly between 2 and 3. dig_thr
is considered to be 1 if digmat[i,j] = 0 (unassigned)
"""
function smooth_2D_map(digmat; dig_thr = 8)
  roundvec = zeros(Int, 8) #8 nearest neighbors
  nn_count = zeros(Int, 4) #counts/state
  for n in 2:size(digmat,2)-1
    for m in 2:size(digmat,1)-1
      roundvec = [digmat[m-1,n] digmat[m+1,n] digmat[m, n-1] digmat[m, n+1] digmat[m-1, n+1] digmat[m-1, n-1] digmat[m+1, n+1] digmat[m+1, n-1]]
      for k = 1:4
        nn_count[k] = count(x -> x == k, roundvec)
      end
      highest_nn = maximum(nn_count)
      indmaxs = findall(x -> x == highest_nn, nn_count) #find all the dominant colors
      indmax = indmaxs[rand(1:end)] #break ties randomly
      if digmat[m,n] == 0 || maximum(nn_count) >= dig_thr #flip states if surrounded
        #by dig_thr points of the same color, or if the state is unassigned
        digmat[m,n] = indmax
      end
    end
  end
  return digmat
end

"""
    get_fidelities_2D(distrs, xpoints, ypoints, digmat)

Return readout assignment probabilities in a 2-qubit space
distrs: distributions for the 4 computational states
digmat: digitization map with grid [xpoints; ypoints]
"""
function get_fidelities_2D(distrs, xpoints, ypoints, digmat)
  probmat = zeros(4,4)
  for k = 1:4
    #ensure that the distributions are normalized
    distrs[k].density/=sum(distrs[k].density)
    for (n, yi) in enumerate(distrs[k].y)
      #search in digmat
      y_ind = max(1,searchsortedlast(ypoints, yi))
      for (m, xi) in enumerate(distrs[k].x)
        x_ind = max(1,searchsortedlast(xpoints, xi))
        for d=1:4
          probmat[k,d] += distrs[k].density[m,n]*(digmat[x_ind,y_ind] == d)
        end
      end
    end
  end
  println("Assign fid. P_00 = $(probmat[1,1])")
  println("Assign fid. P_01 = $(probmat[2,2])")
  println("Assign fid. P_10 = $(probmat[3,3])")
  println("Assign fid. P_11 = $(probmat[4,4])")
  p_0x = sum(probmat[1:2,1:2])/2 #assignment prob. for A1 independent of A2
  p_1x = sum(probmat[3:4,3:4])/2
  p_x0 = (probmat[1,1]+probmat[1,3]+probmat[3,1]+probmat[3,3])/2
  p_x1 = (probmat[2,2]+probmat[2,4]+probmat[4,2]+probmat[4,4])/2
  println("Single-qubit fid. A1 = $((p_0x+p_1x)/2)")
  println("Single-qubit fid. A2 = $((p_x0+p_x1)/2)")
  return probmat
end

"""
    single_shot_fidelities_2D(filenum, datestr; ch_a1, ch_a2, grid_dim, nn_quorum, showPlot)

Wrapper to load and digitize 2-q single-shot data.
Data format: single-shot alternating between 00, 01, 10, 11 in channels {ch_a1, ch_a2}
grid_dim: number of bins / qubit
nn_quorum: number of votes required for the nearest neighbors to flip a value in the assignment map
indep_thrs: if true, the thresholds are chosen independently of the other qubit
"""
function single_shot_fidelities_2D(filenum, datestr; ch_a1 = 1, ch_a2 = 2, grid_dim=51, nn_quorum = 6, showPlot = false, indep_thrs = false)
  data=load_data(datapath, filenum, datestr)[2];
  a1 = real(data[ch_a1]["data"])
  a2 = real(data[ch_a2]["data"])

  if indep_thrs
    #pick thresholds from single-excitation distributions
    thresholds = (digitize([], a1[1:4:end], a1[2:4:end])[4], digitize([], a2[1:4:end], a2[3:4:end])[4])
  else
    #2D thresholds will be computed next
    thresholds = (NaN, NaN)
  end
  cal00 = (a1[1:4:end],a2[1:4:end])
  cal01 = (a1[2:4:end],a2[2:4:end])
  cal10 = (a1[3:4:end],a2[3:4:end])
  cal11 = (a1[4:4:end],a2[4:4:end])
  k00 = kde((cal00[1],cal00[2]))
  k01 = kde((cal01[1],cal01[2]))
  k10 = kde((cal10[1],cal10[2]))
  k11 = kde((cal11[1],cal11[2]))
  xpts,ypts,assign_mat = map_2D(k00, k01, k10, k11, grid_dim; thresholds = thresholds)
  smoothed_mat = smooth_2D_map(assign_mat; dig_thr = nn_quorum)
  fidelity_mat = get_fidelities_2D([k00 k01 k10 k11], xpts, ypts, smoothed_mat)
  if showPlot
    fig = figure("pyplot_surfaceplot",figsize=(5,4))
    xpoints = repmat(xpts,1,length(ypts))
    ypoints = repmat(ypts',length(xpts),1)
    ax = fig[:add_subplot](1,1,1)
    pcolormesh(xpts, ypts, smoothed_mat',cmap = PyPlot.get_cmap("Accent", 5),vmin=0,vmax=4)
    colorbar(ticks=0:4)
    xlabel(L"$A_1$ (V)")
    ylabel(L"$A_2$ (V)")
  end
  return smoothed_mat, fidelity_mat
end

"""
Load multi-qubit calibration data (e.g., measurement results for computational states)
"""
function load_ss(data, qubits; group = "shots", ref_group = "main", quad = :real)
    shots = Dict()
    for (ct, qubit) in enumerate(reverse(qubits)) #from LSB to MSB
      shots[qubit] = Dict()
      metadata_key = []
      for k in keys(data[qubit*"-"*ref_group])
        if contains(k, "metadata")
          push!(metadata_key, k)
        end
      end
      if length(metadata_key) != 1
        error("Invalid metadata")
      end
      data_out = []
      for m = 0:2^length(qubits)-1
        label = bin(m, length(qubits))
        ind = findall(x -> x==parse(UInt8, label, 2), data[qubit*"-"*group][metadata_key[1]])
        shots[qubit][label] = eval(quad).(data[qubit*"-"*group]["Data"][ind])
      end
    end
    return shots
end

"""
Digitize an N-qubit measurement with independent single-qubit thresholds. Calculate assignment probabilities. For N=2, you
can use single_shot_fidelities_2D and related functions instead, to optimize assingment probabilities in 2D.
shots: single-shots (loaded by load_ss)
qubits: list/tuple of qubit names
thresholds: list of digitization threshold (determined by digitize or get_fidelity)
Returns a dictionary (with computational states as keys and counts as values) and the corresponding matrix with
assignment probabilities
"""
function bin_N(shots, qubits, thresholds)
    nshots = length(shots[qubits[1]][bin(0,length(qubits))])
    nstates = 2^length(qubits)
    keys = []
    for m = 0:nstates-1
        label = bin(m, length(qubits))
        keys = push!(keys, label)
    end
    assignmat = Dict()
    for key in keys
        assignmat[key] = zeros(nshots)
        for n = 1:nshots
            for q = 1:length(qubits)
                assignmat[key][n] += (2^(length(qubits)-q))*(shots[qubits[q]][key][n] < thresholds[q])
            end
        end
    end
    probmat = zeros(nstates,nstates);
    for j in 1:nstates
        for k in 1:nstates
            probmat[j, k] = countmap(assignmat[bin(j-1,3)])[k-1]/nshots
        end
    end
    return assignmat, probmat
end

"""
Get single-qubit assingment probabilities as a function of the other 2-qubit states. Useful to assess readout
crosstalk.
"""
#TODO: generalize for N>3 qubits
function get_1q_fid_3q(probmat)
  inds= 1:4
  inds2 = [1,2,5,6]
  inds3 = 1:2:7

  pq1_0 = []; pq1_1 = []; pq2_0 = []; pq2_1 = []; pq3_0 = []; pq3_1 = []

  for k in 1:4
      push!(pq1_0, sum(probmat[k,inds]/nshots))
      push!(pq1_1, sum(probmat[k+4,inds+4]/nshots))
      push!(pq2_0, sum(probmat[inds2[k],inds2]/nshots))
      push!(pq2_1, sum(probmat[inds2[k]+2,inds2+2]/nshots))
      push!(pq3_0, sum(probmat[inds3[k], inds3]/nshots))
      push!(pq3_1, sum(probmat[inds3[k]+1, inds3+1]/nshots))
  end
  return [pq1_0, pq1_1, pq2_0, pq2_1, pq3_0, pq3_1]
end


text:
"""
    @add_kwonly function_definition

Define keyword-only version of the `function_definition`.

    @add_kwonly function f(x; y=1)
        ...
    end

expands to:

    function f(x; y=1)
        ...
    end
    function f(; x = error("No argument x"), y=1)
        ...
    end
"""
macro add_kwonly(ex)
  esc(add_kwonly(ex))
end

add_kwonly(ex::Expr) = add_kwonly(Val{ex.head}, ex)

function add_kwonly(::Type{<: Val}, ex)
  error("add_only does not work with expression $(ex.head)")
end

function add_kwonly(::Union{Type{Val{:function}},
                            Type{Val{:(=)}}}, ex::Expr)
  body = ex.args[2:end]  # function body
  default_call = ex.args[1]  # e.g., :(f(a, b=2; c=3))
  kwonly_call = add_kwonly(default_call)
  if kwonly_call === nothing
    return ex
  end

  return quote
    begin
      $ex
      $(Expr(ex.head, kwonly_call, body...))
    end
  end
end

function add_kwonly(::Type{Val{:where}}, ex::Expr)
  default_call = ex.args[1]
  rest = ex.args[2:end]
  kwonly_call = add_kwonly(default_call)
  if kwonly_call === nothing
    return nothing
  end
  return Expr(:where, kwonly_call, rest...)
end

function add_kwonly(::Type{Val{:call}}, default_call::Expr)
  # default_call is, e.g., :(f(a, b=2; c=3))
  funcname = default_call.args[1]  # e.g., :f
  required = []  # required positional arguments; e.g., [:a]
  optional = []  # optional positional arguments; e.g., [:(b=2)]
  default_kwargs = []
  for arg in default_call.args[2:end]
    if isa(arg, Symbol)
      push!(required, arg)
    elseif arg.head == :(::)
      push!(required, arg)
    elseif arg.head == :kw
      push!(optional, arg)
    elseif arg.head == :parameters
      @assert default_kwargs == []  # can I have :parameters twice?
      default_kwargs = arg.args
    else
      error("Not expecting to see: $arg")
    end
  end
  if isempty(required) && isempty(optional)
    # If the function is already keyword-only, do nothing:
    return nothing
  end
  if isempty(required)
    # It's not clear what should be done.  Let's not support it at
    # the moment:
    error("At least one positional mandatory argument is required.")
  end

  kwonly_kwargs = Expr(:parameters, [
    Expr(:kw, pa, :(error($("No argument $pa"))))
    for pa in required
  ]..., optional..., default_kwargs...)
  kwonly_call = Expr(:call, funcname, kwonly_kwargs)
  # e.g., :(f(; a=error(...), b=error(...), c=1, d=2))

  return kwonly_call
end

function num_types_in_tuple(sig)
  length(sig.parameters)
end

function num_types_in_tuple(sig::UnionAll)
  length(Base.unwrap_unionall(sig).parameters)
end

function numargs(f)
  typ = Tuple{Any, Val{:analytic}, Vararg}
  typ2 = Tuple{Any, Type{Val{:analytic}}, Vararg} # This one is required for overloaded types
  typ3 = Tuple{Any, Val{:jac}, Vararg}
  typ4 = Tuple{Any, Type{Val{:jac}}, Vararg} # This one is required for overloaded types
  typ5 = Tuple{Any, Val{:tgrad}, Vararg}
  typ6 = Tuple{Any, Type{Val{:tgrad}}, Vararg} # This one is required for overloaded types
  numparam = maximum([(m.sig<:typ || m.sig<:typ2 || m.sig<:typ3 || m.sig<:typ4 || m.sig<:typ5 || m.sig<:typ6) ? 0 : num_types_in_tuple(m.sig) for m in methods(f)])
  return (numparam-1) #-1 in v0.5 since it adds f as the first parameter
end

function isinplace(f,inplace_param_number)
  numargs(f)>=inplace_param_number
end

### Default Linsolve

# Try to be as smart as possible
# lu! if Matrix
# lu if sparse
# gmres if operator

mutable struct DefaultLinSolve
  A
  iterable
end
DefaultLinSolve() = DefaultLinSolve(nothing, nothing)

function (p::DefaultLinSolve)(x,A,b,update_matrix=false;tol=nothing, kwargs...)
  if p.iterable isa Vector && eltype(p.iterable) <: LinearAlgebra.BlasInt # `iterable` here is the pivoting vector
    F = LU{eltype(A)}(A, p.iterable, zero(LinearAlgebra.BlasInt))
    ldiv!(x, F, b)
    return nothing
  end
  if update_matrix
    if typeof(A) <: Matrix
      blasvendor = BLAS.vendor()
      # if the user doesn't use OpenBLAS, we assume that is a better BLAS
      # implementation like MKL
      #
      # RecursiveFactorization seems to be consistantly winning below 100
      # https://discourse.julialang.org/t/ann-recursivefactorization-jl/39213
      if ArrayInterface.can_setindex(x) && (size(A,1) <= 100 || ((blasvendor === :openblas || blasvendor === :openblas64) && size(A,1) <= 500))
        p.A = RecursiveFactorization.lu!(A)
      else
        p.A = lu!(A)
      end
    elseif typeof(A) <: Tridiagonal
      p.A = lu!(A)
    elseif typeof(A) <: Union{SymTridiagonal}
      p.A = ldlt!(A)
    elseif typeof(A) <: Union{Symmetric,Hermitian}
      p.A = bunchkaufman!(A)
    elseif typeof(A) <: SparseMatrixCSC
      p.A = lu(A)
    elseif ArrayInterface.isstructured(A)
      p.A = factorize(A)
    elseif !(typeof(A) <: AbstractDiffEqOperator)
      # Most likely QR is the one that is overloaded
      # Works on things like CuArrays
      p.A = qr(A)
    end
  end

  if typeof(A) <: Union{Matrix,SymTridiagonal,Tridiagonal,Symmetric,Hermitian} # No 2-arg form for SparseArrays!
    x .= b
    ldiv!(p.A,x)
  # Missing a little bit of efficiency in a rare case
  #elseif typeof(A) <: DiffEqArrayOperator
  #  ldiv!(x,p.A,b)
  elseif ArrayInterface.isstructured(A) || A isa SparseMatrixCSC
    ldiv!(x,p.A,b)
  elseif typeof(A) <: AbstractDiffEqOperator
    # No good starting guess, so guess zero
    if p.iterable === nothing
      p.iterable = IterativeSolvers.gmres_iterable!(x,A,b;initially_zero=true,restart=5,maxiter=5,tol=1e-16,kwargs...)
      p.iterable.reltol = tol
    end
    x .= false
    iter = p.iterable
    purge_history!(iter, x, b)

    for residual in iter
    end
  else
    ldiv!(x,p.A,b)
  end
  return nothing
end

function (p::DefaultLinSolve)(::Type{Val{:init}},f,u0_prototype)
  DefaultLinSolve()
end

const DEFAULT_LINSOLVE = DefaultLinSolve()

@inline UNITLESS_ABS2(x) = real(abs2(x))
@inline DEFAULT_NORM(u::Union{AbstractFloat,Complex}) = @fastmath abs(u)
@inline DEFAULT_NORM(u::Array{T}) where T<:Union{AbstractFloat,Complex} =
                                         sqrt(real(sum(abs2,u)) / length(u))
@inline DEFAULT_NORM(u::StaticArray{T}) where T<:Union{AbstractFloat,Complex} =
                                            sqrt(real(sum(abs2,u)) / length(u))
@inline DEFAULT_NORM(u::RecursiveArrayTools.AbstractVectorOfArray) =
                                        sum(sqrt(real(sum(UNITLESS_ABS2,_u)) / length(_u)) for _u in u.u)
@inline DEFAULT_NORM(u::AbstractArray) = sqrt(real(sum(UNITLESS_ABS2,u)) / length(u))
@inline DEFAULT_NORM(u) = norm(u)

"""
  prevfloat_tdir(x, x0, x1)

Move `x` one floating point towards x0.
"""
function prevfloat_tdir(x, x0, x1)
  x1 > x0 ? prevfloat(x) : nextfloat(x)
end

function nextfloat_tdir(x, x0, x1)
  x1 > x0 ? nextfloat(x) : prevfloat(x)
end

function max_tdir(a, b, x0, x1)
  x1 > x0 ? max(a, b) : min(a, b)
end

alg_autodiff(alg::AbstractNewtonAlgorithm{CS,AD}) where {CS,AD} = AD
alg_autodiff(alg) = false

"""
  value_derivative(f, x)

Compute `f(x), d/dx f(x)` in the most efficient way.
"""
function value_derivative(f::F, x::R) where {F,R}
    T = typeof(ForwardDiff.Tag(f, R))
    out = f(ForwardDiff.Dual{T}(x, one(x)))
    ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out)
end

value(x) = x
value(x::Dual) = ForwardDiff.value(x)
value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x)