text:
using Markdown

"return information about all matrices selected by pattern"
mdinfo(p::Pattern) = mdinfo(MATRIX_DB, p)
function mdinfo(db::MatrixDatabase, p::Pattern)
    check_symbols(p)
    mdbuffer = Markdown.MD([])
    md = mdbuffer
    for name in mdlist(p)
        try
            md = mdinfo(db.data[name])
        catch
            md = Markdown.parse("# $name\nno info available")
        finally
            append!(mdbuffer.content, md.content)
        end
    end
    mdbuffer
end

mdinfo(md::MatrixDescriptor) = mdinfo(md.data)

_mdheader(md::Markdown.MD, p, o) = isempty(md.content) ? (nothing, o) : _mdheader(md.content[1], md, o)
_mdheader(md::Markdown.Header, p, o) = (md, p)
_mdheader(md, p, o) = (nothing, o)

function mdinfo(data::GeneratedMatrixData)
    md = eval(Meta.parse("Docs.@doc $(data.func)", raise = false))
    # As md is cached internally, need to make copies
    mdh, md = _mdheader(md, nothing, md)
    if mdh != nothing
        mdh, md = Markdown.Header(copy(mdh.text)), copy(md)
        push!(mdh.text, " ($(data.name))")
        md.content[1] = mdh
    end
    md
end

_repl(a::AbstractString) = a
_repl(a::AbstractString, p::Pair, q::Pair...) = _repl(replace(a, p), q...)

function mdinfo(data::RemoteMatrixData)
    file = verify_loadinfo(data)
    txt = mmreadcomment(file)
    txt = _repl(txt, r"^%-+$"m => "---", r"^%%" => "###### ", r"%+" => "* ")
    md = Markdown.parse(txt)
    insert!(md.content, 1, Markdown.Header{1}([data.name]))
    md
end

"""
format output list to multi- column table form adapted to display width.
"""
function buildnametable(hdr, name::AbstractVector, maxrow::Integer=0)
    maxrow = maxrow <= 0 ? displaysize(stdout)[2] + maxrow : maxrow
    maxrow = clamp(maxrow, 20, 1000)
    hdr isa Vector || ( hdr = [hdr] )
    n = length(name)
    name = string.(name)
    cols = 1
    while sumwidth(name, cols+1) <= maxrow - 1 * cols && cols < n
        cols += 1
    end
    rows = (length(name) + cols - 1) ÷ cols
    n = length(name)
    while n < rows * cols
        push!(name, "")
        n += 1
    end
    name = reshape(name, rows, cols)
    name = vecvec(name)
    insert!(name, 1, [string.(hdr)..., tab("", cols-length(hdr))...])
    Markdown.Table(name, tab(:l, cols))
end

function sumwidth(name::AbstractVector, cols::Integer)
    n = length(name)
    rows = (n + cols - 1) ÷ cols
    wi = 0
    for k = 1:cols
        r = (k-1)*rows+1:1:min(k*rows,n)
        if !isempty(r)
            wi += maximum(length.(name[r]))
        end
    end
    wi
end

vecvec(tab::Matrix) = [ tab[i,:] for i in 1:size(tab,1)]
tab(a, cols::Integer) = [a for i in 1:cols]

function buildnametable1(db::MatrixDatabase, hdr, p::Pattern, maxrow::Integer=0)
    dali = mdata.(Ref(db), mdlist(db, p))
    lt(a, b) = a.id < b.id
    sort!(dali, lt=lt)
    item(data::MatrixData) = string(data.id, " ", data.name)
    MD(buildnametable(hdr, item.(dali), maxrow))
end

MD(a...) = Markdown.MD([a...])

listnames(p::Pattern) = listnames(MATRIX_DB, p)
function listnames(db::MatrixDatabase, p::Pattern)
    li = mdlist(db, p)
    MD(buildnametable("list($(length(li)))", li))
end

##############################
# display database information
##############################

"""
    overview([db])

return formatted overview about matrices and groups in the collection.
"""
function overview(db::MatrixDatabase)
    # Print information strings
    LMARG = -10
    hdr_mat = Markdown.Header{3}("Currently loaded Matrices")

    dli(p) = mdata.(Ref(db), mdlist(db, p))
    ldall(p) = [length(mdlist(p & isloaded)), length(mdlist(p))]

    bmat = buildnametable1(db, "builtin(#)", :builtin, LMARG)
    umat = buildnametable1(db, "user(#)", :user, LMARG)
    spdir = buildnametable(["Suite Sparse", "of"], ldall(sp(:)))
    mmdir = buildnametable(["MatrixMarket", "of"], ldall(mm(:)))

    hdr_groups = Markdown.Header{3}("Groups:")

    groups = buildnametable("Groups", listgroups(), LMARG)

    MD(([hdr_mat, bmat, umat, groups, spdir, mmdir]))
end

"""
    mdinfo()
Overview about matrices.
"""
mdinfo(db::MatrixDatabase=MATRIX_DB) = overview(db)


text:
###############################################################################
#
# Leaf boundary layer parameter set
#
###############################################################################
#= AbstractLeafBLParaSet type tree
---> LeafBLParaSetFixed
---> LeafBLParaSetGentine
=#

#=
"""
    AbstractLeafBLParaSet

Hierarchy of the `AbstractLeafBLParaSet`:
- [`LeafBLParaSetFixed`](@ref)
- [`LeafBLParaSetGentine`](@ref)
"""
abstract type AbstractLeafBLParaSet end

"""
    struct LeafBLParaSetFixed{FT<:AbstractFloat}

Leaf boundary layer with fixed resistance

# Fields
$(DocStringExtensions.FIELDS)
"""
struct LeafBLParaSetFixed{FT<:AbstractFloat} <: AbstractLeafBLParaSet
    "Resistance of the leaf boundary layer `[m² s mol⁻¹]`"
    ra::FT
end

"""
    struct LeafBLParaSetGentine

Gentine's boundary layer scheme that computes Monin Obhukow length and
    resistances across the leaf to Ca scale
"""
struct LeafBLParaSetGentine <: AbstractLeafBLParaSet end
=#

###############################################################################
#
# Grid information struct
# This struct should go to types.jl
#
###############################################################################
#=
"""
    struct SPACGrid{FT,lat,lom,ele}

Information for a soil-plant-air continuum in a grid
"""
Base.@kwdef mutable struct SPACGrid{FT}
    # TODO generalize this function?
    # Grid information
    "The latitutional size `[degree]`"
    dim_lat::FT = FT(0.25)
    "The longitutional size `[degree]`"
    dim_lon::FT = FT(0.25)
    "The latitude `[degree]`"
    lat    ::FT = FT(40.0)
    "The longitude `[degree]`"
    lon    ::FT = FT(-100.0)
    "The elevation `[m]`"
    ele    ::FT = FT(0.0)
    "Atmospheric pressure `[Pa]`"
    p_atm  ::FT = FT(101325.0)
    "Atmospheric O₂ partial pressure"
    p_O₂   ::FT = FT(21176.925)

    # Soil and Air profile from the soil to top of canopy
    "Total number of vertical layers above ground"
    n_above::Int = 20
    "Total number of vertical layers below ground"
    n_below::Int = 10
    "A list of upper z for the veritical layers above ground `[m]`"
    zl_above    ::Array{FT,1} = collect(FT,  1.0: 1.0:20.0)
    "A list of lower z for the veritical layers below ground `[m]`"
    al_below    ::Array{FT,1} = collect(FT, -0.2:-0.2:-2.0)
    "An array of CO₂ partial pressure along the vertical layers `[Pa]`"
    p_CO₂_array ::Array{FT,1} = zeros(FT, n_above) + 40
    "An array of H₂O partial pressure along the vertical layers `[Pa]`"
    p_H₂O_array ::Array{FT,1} = zeros(FT, n_above) + 1500
    "An array of soil water potential `[MPa]`"
    p_soil_array::Array{FT,1} = zeros(FT, n_above) + FT(-0.5)
    "An array of soil water content"
    swc_array   ::Array{FT,1} = zeros(FT, n_above) + FT(0.4)
    "An array of air temperature `[K]`"
    t_air_array ::Array{FT,1} = zeros(FT, n_above) + FT(298.15)
    "An array of soil temperature `[K]`"
    t_soil_array::Array{FT,1} = zeros(FT, n_above) + FT(298.15)

    # Plant information
    "An array of trees in the grid"
    tree_array::Array = []
    "An array of tree density in the grid"
    density_array::Array{FT,1} = []
end

# TODO add a function to map plants into a SPACGrid struct,
# automatically break the tree canopy into afew layers as defined in the SPACGrid
function generate_canopy_and_root()
    ;
end

"""
update_canopy_from_rt_module!(tree::Tree, canopy_rt::struct_canopy, canOpt_rt::struct_canopyOptProps, canRad_rt::struct_canopyRadiation)

Updates canopy information from the RT module, given
- `tree` A [`Tree`](@ref) type
- `canopy_rt` A [`struct_canopy`] type from the Plant module
- `cadOpt_rt` A [`struct_canopyOptProps`] type from the CanopyRT module
- `canRad_rt` A [`struct_canopyRadiation`] type from the CanopyRT module

This interface function is pending for leaf temperature...
"""
function update_canopy_from_rt_module!(tree::Tree, canopy_rt::Canopy4RT, canOpt_rt::AbstractCanopyOpti, canRad_rt::CanopyRadiation)
    # fraction of sunlit leaves in each layer
    fraction_sl = repeat(canopy_rt.lidf, outer=[ length(canopy_rt.lazitab) ]) / length(canopy_rt.lazitab)

    # update the PAR from canopyRT module to the Plant module
    if tree.n_canopy == canopy_rt.nlayers
        nlayers = canopy_rt.nlayers
        for pl_layer in 1:nlayers
            canopyi  = tree.canopy_list[pl_layer]
            rt_layer = nlayers + 1 - pl_layer

            # set the diffuse par to all the leaves
            canopyi.par_list .= canRad_rt.absPAR_shadeCab[rt_layer] * 1E6
            # add the direct PAR to sunlit leaves
            canopyi.par_list[1:end-1] .+= reshape(canRad_rt.absPAR_sunCab[:,:,rt_layer],(:,1))[:,1] * 1E6

            # calculate the fraction of sunlit and shaded leaves
            f_view = (canOpt_rt.Ps[rt_layer]+canOpt_rt.Ps[rt_layer+1]) / 2
            la_new = canopyi.la .* [f_view .* fraction_sl; 1-f_view]
            canopyi.f_view  = f_view
            canopyi.la_list = la_new

            # update leaf temperature
            canopyi.t_list[1:end-1] .= reshape(canRad_rt.T_sun3D[:,:,rt_layer],(:,1))[:,1]
            canopyi.t_list[end]      = canRad_rt.T_shade[rt_layer]

            # update leaf-to-air VPD
            canopyi.d_list = saturation_vapor_pressure.(canopyi.t_list) .- canopyi.p_H₂O
        end
    else
        println("Error: the canopy layer in the Tree differs from that in the struct_canopy!")
    end
end

"""
# compute surface energy balance of a leaf Rn-H-LE=0
# Arguments
- `met::meteo`: meteorological forcing.
- `Ts_t::Float32`: Leaf temperature forcing.
- `T::Number`: Leaf Temperature
"""
#function LeafEnergyWaterBalance(mod::AbstractPhotosynthesis, Tleaf, psileaf, Cc, met::meteo, l::LeafPhotosynthesisParams,   psi_s)
function LeafEnergyWaterBalance(mod::AbstractPhotosynthesis, Tleaf, psileaf, Cc, met::MeteoParams, l::LeafParams,   psi_s)
    # plugged Tleaf and psileaf outside of l as they are being iterated during fine time stepping - rest assumed constant
    # find solution to the surface energy LeafEnergyBalance, with a small thermal inertia: C dT/dt = Rn-H-LE and water balance of a leaf d.dH20/dt = Sap - T
    # first set parameter values that are assumed not to be changing with energy budget
    l.T            = Tleaf;
    l.psi_l        = psileaf;
    l.Cc           = Cc;

    if(Tleaf<200.0 || psileaf>0.0)
        println("Error in Tleaf and psi_leaf - unphysical values")
        @show Tleaf
        @show psileaf
        process.exit(10)
    end
    LeafPhotosynthesis!(mod, l, met); # compute H, LE in there as well
    LeafPhotosynthesis!(mod, l, met);
    # Leaf response time in seconds (15min)
    tau_gs = 15.0*60.0
    # TODO aerodynamic resistance should be double for the sensible heat flux at leaf level, but do we really care?
    l.Cleaf = l.LMA* ( (1.0-l.RWC)*physcon.Cdryleaf + l.RWC* CP_L )/(1.0-l.RWC); # leaf conductivity
    # println("Cleaf=",l.Cleaf)
    lv           =   latent_heat_vapor(Tleaf);
    l.Rn      =   (1-l.α)*met.S_down +  2 * (met.L_down - l.ε* K_STEFAN *Tleaf^4); # 2 is for two sides of the leaves
    #println("S_down=",met.S_down," , Ldown=",met.L_down," Tleaf=",Tleaf)
    #dRn_dTs = - 4*l.ε*physcon.σ*Tleaf^3;
    setkx!(l,psi_s, psileaf) ;# set hydraulic conductivity as a function of psis and psi_l
    l.kx = l.kx/1000
    #@show l.kx
    # TODO * viscosty or / viscosity?
    l.Sap     =   (psi_s - l.psi_l - ρ_H₂O * GRAVITY * l.height)*l.kx * relative_viscosity(Tleaf); # equal to int k(psi)dpsi + gravity term=-rho.g.mean(k)*height, includes also water viscosty;
    #println("Sap=",lv*l.Sap," W/m^2, kx=",l.kx,", rho.g.h=",physcon.ρw*physcon.grav*l.height)
    #@show flux.LE
    dT_dt        =   (l.Rn-l.H-l.LE)/l.Cleaf; # 2 times for up and down part of the leaves - TODO need to check this I am not sure I agree when integrated over the canopy
    dH2Ol_dt     =   (l.Sap-l.LE/lv)/l.Ctree*1000;
    dt           =   1.0; # one second time step for Cc - a bit arbitrary and not cruCcal but will be cahnged later for actual airspace
    #dCc_dt       =   0.0#(flux.An_diffusion - flux.An_biochemistry) / dt; #(l.Chloroplast_rel_volume*l.dleaf*l.LAI);
    dgs_dt       =   (l.gs_ss-l.gs)/tau_gs
    #println("Cc/(dCc/dt)=",1.0/(dCc_dt/l.Cc)," An_biochemistry=",flux.An_biochemistry," An_diffusion=",flux.An_diffusion)

    # print(flux.Cs, " ppm, ", l.VPD/1000.0, " (kPA), ", flux.An_biochemistry, " micromol/s/m2  "," An_diffusion=",flux.An_diffusion)
    #println("Sdown= " , met.S_down, "W/m2, Rn=",flux.Rn,"W/m2, SEB=",flux.Rn-flux.H-flux.LE,"W/m2, H= ",flux.H, "W /m2, LE= ",flux.LE, "W /m2, dT_dt=",dT_dt*3600," (K/hr), ra=",flux.ra, " (s/m) ")
    return dT_dt, dH2Ol_dt, dgs_dt
end
=#

#=
"""
  CanopyPhotosynthesis!(mo::AbstractPhotosynthesis, leaf::LeafPhotosynthesisParams, met::meteo, canRad)

Compute Photosynthesis for all leaves in a canopy, given:
- `mod`    One [`PhotoMods`](@ref) model setup structure
- `leaf`   One [`LeafPhotosynthesisParams`](@ref) structure
- `met`    One [`lmeteo`](@ref) structure
- `canRad` One [`struct_canopyRadiation`](@ref) structure
"""
function CanopyPhotosynthesis!(mo::AbstractPhotosynthesis, leaf::LeafPhotosynthesisParams, met::meteo, canRad)
  ISun   = CartesianIndices(canRad.absPAR_sunCab)
  IShade = CartesianIndices(canRad.absPAR_shadeCab)
  for i in IShade
    leaf.APAR = 1e6*canRad.absPAR_shadeCab[i]
    leaf.T    = canRad.T_shade[i]
    LeafPhotosynthesis!(mo, leaf, met)
    canRad.GPP_shade[i] = leaf.An;
    canRad.ϕ_shade[i]   = leaf.ϕs;
    canRad.gs_shade[i]  = leaf.gs;
    canRad.H_shade[i]   = leaf.H;
    canRad.LE_shade[i]  = leaf.LE;
    canRad.Cc_shade[i]  = leaf.Cc;
    canRad.NPQ_shade[i] = leaf.NPQ;
  end
  for i in ISun
    leaf.APAR = 1e6*canRad.absPAR_sunCab[i]
    leaf.T    = canRad.T_sun[i[3]]
    LeafPhotosynthesis!(mo, leaf, met)
    canRad.GPP_sun[i] = leaf.An;
    canRad.ϕ_sun[i]   = leaf.ϕs;
    canRad.gs_sun[i]  = leaf.gs;
    canRad.H_sun[i]   = leaf.H;
    canRad.LE_sun[i]  = leaf.LE;
    canRad.Cc_sun[i]  = leaf.Cc;
    canRad.NPQ_sun[i] = leaf.NPQ;
  end
end

=#

###############################################################################
#
# Update leaf boundary layer conductance
#
###############################################################################

#=
function boundary_layer_resistance!(bl::LeafBLParaSetGentine, leaf::LeafParams,  envir::MeteoParams)
    # based on Monin-Obukhov Similiarity theory -> to be changed for LES
    # compute Obukhov length
    # iterate a few times

    # TODO ideally should not have any information about the leaves as it is the turbuence above canopy in log profile
    # first update roughness (if phenology is changing)
    # need to make this abstract as well.
    setRoughness!(leaf)

    if(leaf.height>envir.zscreen)
        println("Vegetation height is higher than screen level height")
        #process.exit(20)
    end

    rmin  = 10.0;
    Lold  = -1e6;
    raw_full = -999.0;
    H     = -999.0;
    LE    = -999.0;
    ustar = -999.0;

    # TODO change this later - 0 leaf boundary layer resistance - I prefer this as it is somewhat included in z0
    leaf.Cd     = Inf;

    # need to compute the evolving buoyancy flux
    Tv       = envir.t_air*(1.0+0.61*envir.e_air);
    ra_leaf  = 1/leaf.Cd / sqrt(envir.U/leaf.dleaf);
    #println("ra_leaf=",ra_leaf)
    ΔT   =   leaf.T - envir.t_air;
    VPD      =   max(leaf.esat-envir.e_air,1.0)
    ρd       =   envir.P_air/(R_D * envir.t_air)

    # TODO use ThermoDynamics to compute latent heat flux
    lv       =   latent_heat_vapor(leaf.T);
    L        =   envir.L; # initial Obukhov length
    counter = 1;
    #@show abs(1.0-Lold/L)
    while (counter<20 && abs(1.0-Lold/L)>1e-4) # 1% error
      #println("L=",L," ,Lold=",Lold)
      ra_m     =   max(1.0/(VON_KARMAN_CONST^2*envir.U) * ( log((envir.zscreen - leaf.d)/leaf.z0m) - ψ_m((envir.zscreen - leaf.d)/L,envir.stab_type_stable) + ψ_m(leaf.z0m/L,envir.stab_type_stable) ) * ( log((envir.zscreen - leaf.d)/leaf.z0m) - ψ_m((envir.zscreen - leaf.d)/L,envir.stab_type_stable) + ψ_h(leaf.z0h/L,envir.stab_type_stable) ), rmin) ;# momentum aerodynamic resistance
      ra_w     =   max(1.0/(VON_KARMAN_CONST^2*envir.U) * ( log((envir.zscreen - leaf.d)/leaf.z0m) - ψ_m((envir.zscreen - leaf.d)/L,envir.stab_type_stable) + ψ_m(leaf.z0m/L,envir.stab_type_stable) ) * ( log((envir.zscreen - leaf.d)/leaf.z0h) - ψ_h((envir.zscreen - leaf.d)/L,envir.stab_type_stable) + ψ_h(leaf.z0h/L,envir.stab_type_stable) ), rmin) ;# water aerodynamic resistance
      #println("ra_m=",ra_m)
      #@show ra_leaf
      #@show ra_w
      ram_full =   ra_leaf + ra_m;
      raw_full =   ra_leaf + ra_w;
      H        =   ρd * CP_D * ΔT / raw_full;
      rs_s_m   =   envir.g_m_s_to_mol_m2_s/leaf.gs;
      LE       =   WATER_AIR_MRATIO / envir.P_air * ρd * lv * VPD / (rs_s_m + raw_full)
      #@show LE
      ustar    =   sqrt(envir.U/ram_full);
      Hv_s     =   H + 0.61 * CP_D / lv * envir.t_air * LE
      Lold     =   L;
      L        =   - ustar^3*Tv/(GRAVITY * VON_KARMAN_CONST * Hv_s); # update Obukhov length
      #println("L=",L, " ra=",ra_w," (s/m), H=", H, " (W/m2), counter=", counter)
      counter = counter+1
    end

    # TODO Is the envir more a "global" variable set for the tree? Or a "local" variable for each layer/leaf?
    # If for the whole plant or layer, then envir should be updated only once?
    # save these values in leaf and flux structures
    envir.L   = L
    leaf.H  = H
    leaf.LE = LE
    envir.ustar = ustar
    envir.ra = raw_full
end

"""
ψ_h(ζ,stab_type_stable)

momentum correction function
"""
function ψ_h(ζ,stab_type_stable)
if(ζ>=0.0) # stable conditions
    if(stab_type_stable==1) # Grachev default value - Grachev et al SHEBA 2007
        # issue when x ~ 0
        x  = (1.0+ζ)^(1.0/3.0);
        am = 5.0;
        bm = am/6.5;
        ah = bh = 5.0;
        ch = 3.0;
        Bm = ((1.0-bm)/bm)^(1.0/3.0);
        Bh = sqrt(5);
        ψ = -bh/2*log(1+ch*ζ+ζ*ζ) + (-ah/Bh + bh*ch/(2.0*Bh))*(log((2.0*ζ+ch-Bh)/(2.0*ζ+ch+Bh)) - log((ch-Bh)/(ch+Bh)) );
    elseif(stab_type_stable==2) # Webb correction as in NoahMP
        alpha = 5.0;
        #if(ζ>=0.0 && ζ<=1)
        #  ψ = -alpha * ζ;
        #else
        #  ψ = -alpha;
        #end
        # Pierre - I didn't like this so I changed it to be C1
        ψ = -alpha*tanh(ζ);
    elseif(stab_type_stable==3) # Beljaars-Holtslag, 1991, Journal of Applied Meteorology
        # Constants
        a = 1.0;
        b = 0.667;
        c = 5;
        d = 0.35;
        # Stability correction for momentum
        ψ = (- (1. + (2. / 3.) * a * ζ) ^ (3. / 2.)
                  - b * (ζ - (c / d)) * exp(-d * ζ)
                  - (b * c) / d + 1);
    elseif(stab_type_stable==4) #Cheng and Brutsaert (2005) as described by Jimenez et al. (2012)
        a = 6.1;
        b = 2.5;
        c = 5.3;
        d = 1.1;
        ψ =  -c * log(ζ + (1 + ζ^d) ^ (1. / d));
    end
else # (ζ<0) unstable
    x   = (1.0-16.0*ζ)^0.25;
    ψ = 2.0*log((1.0+x*x)/2.0);
end

return ψ
end

"""
ψ_m(ζ,stab_type_stable)

Momentum correction function
"""
function ψ_m(ζ,stab_type_stable)
FT = eltype(ζ)

if(ζ>=0) # stable conditions
    if(stab_type_stable==1) # Grachev default value - Grachev et al SHEBA 2007
        x  = (1+ζ)^FT(1/3);
        am = FT(5.0);
        bm = am/FT(6.5);
        ah = bh = FT(5.0);
        ch = FT(3.0);
        Bm = ((FT(1.0)-bm)/bm)^FT(1/3);
        Bh = sqrt(5);
        ψ = -FT(3)*am/bm*(x-1) + am*Bm/(2bm)*( FT(2)*log((x+Bm)/(1+Bm))
                                                 - log( (x*x-x*Bm+Bm*Bm)/(1-Bm+Bm*Bm) )
                                                 +2*sqrt(FT(3.0))*( atan((2x-Bm)/(sqrt(FT(3.0))*Bm)) -  atan((2-Bm)/(sqrt(FT(3.0))*Bm)))  );
    elseif(stab_type_stable==2) # Webb correction as in NoahMP, quite numerically unstable at night - lots of oscillations
        alpha = FT(5.0);
        """ should be the correct one but not continuously differntialble
        if(ζ>=0.0 && ζ<=1)
          ψ = -alpha * ζ;
        else
          ψ = -alpha; # should be z less anyways but here just made to create a plateau in the value
        end
        """
        # Pierre's new one:
        ψ = -alpha*tanh(ζ);
    elseif(stab_type_stable==3) # Beljaars-Holtslag, 1991, Journal of Applied Meteorology
        # Constants
        a = FT(1.0);
        b = FT(0.667);
        c = FT(5);
        d = FT(0.35);
        # Stability correction for momentum
        ψ = -(a * ζ
                 + b * (ζ - c / d) * exp(-d * ζ)
                 + b * c / d);
    elseif(stab_type_stable==4) #Cheng and Brutsaert (2005) as described by Jimenez et al. (2012)
        a = FT(6.1);
        b = FT(2.5);
        c = FT(5.3);
        d = FT(1.1);
        ψ = -a * log(ζ + (1 + ζ^b) ^ (1 / b));
    end

else #  (ζ<0) unstable
    x   =   (1-16ζ)^FT(0.25);
    ψ   =   FT(2.0)*log((1+x)/2) + log((1+x*x)/2) - 2*atan(x) + FT(π)/2;
end

return ψ
end

"""
setRoughness!(leaf::LeafPhotosynthesisParams)

Computes leaf roughness lengths, given
- `leaf` One [`LeafPhotosynthesisParams`](@ref) structure
"""
function setRoughness!(leaf::LeafParams)
    # adjust roughness coefficients
    leaf.z0m = 0.1*leaf.height;    # tree roughness (m)
    leaf.z0h = 0.1*leaf.height;    # tree roughness (m) - TODO should be changed later
    leaf.d   = 2/3*leaf.height;    # tree displacement height (m)
end
=#

ext: jl

max_stars_repo_name: bbuman/Land

max_issues_repo_name: bbuman/Land

max_forks_repo_name: bbuman/Land


text:
using LiiBRA
using Test

@testset "DRA_test.jl" begin include("DRA_test.jl") end

max_stars_repo_name: BradyPlanden/LiBRA

max_forks_repo_name: BradyPlanden/LiBRA


text:
### A Pluto.jl notebook ###
# v0.12.6

using Markdown
using InteractiveUtils

# ╔═╡ de7a5e8a-1ee4-11eb-2bbf-0f0eb3e489ec
using Pkg, DrWatson

# ╔═╡ e4eeb13a-1ee4-11eb-0113-91f9a9c3b659
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Combinatorics, Plots, LaTeXStrings
	Random.seed!(123)
end

# ╔═╡ ac3ceabe-1ee4-11eb-3935-313e96aafd62
md"## Listing 2.5"

# ╔═╡ d212f38c-1ee4-11eb-360a-b7d524890c8f
function isUpperLattice(v)
    for i in 1:Int(length(v)/2)
        sum(v[1:2*i-1]) >= i ? continue : return false
    end
    return true
end

# ╔═╡ 6263427a-1ee5-11eb-12d1-1d576fa18461
begin
	n, N = 5, 10^5

	omega = unique(permutations([zeros(Int,n);ones(Int,n)]))
	A = omega[isUpperLattice.(omega)]
	pA_modelI = length(A)/length(omega)
end

# ╔═╡ 62637f3a-1ee5-11eb-34e2-77e37f188fe5
function randomWalkPath(n)
    x, y = 0, 0
    path = []
    while x<n && y<n
       if rand()<0.5
            x += 1
            push!(path,0)
        else
            y += 1
            push!(path,1)
        end
    end
    append!(path, x<n ? zeros(Int64,n-x) : ones(Int64,n-y))
    return path
end

# ╔═╡ 6263f402-1ee5-11eb-20e1-5d1ebc78f310
begin
	pA_modelIIest = sum([isUpperLattice(randomWalkPath(n)) for _ in 1:N])/N
	("Model I: ",pA_modelI, "\t Model II: ", pA_modelIIest)
end

# ╔═╡ 6276d39e-1ee5-11eb-23cc-2fe5377e6ce9
function plotPath(v,l,c)
    x,y = 0,0
    graphX, graphY = [x], [y]
    for i in v
        if i == 0
            x += 1
        else
            y += 1
        end
        push!(graphX,x), push!(graphY,y)
    end
    plot!(graphX, graphY,
            la=0.8, lw=2, label=l, c=c, ratio=:equal, legend=:topleft,
            xlims=(0,n), ylims=(0,n),
            xlabel=L"East\rightarrow", ylabel=L"North\rightarrow")
end

# ╔═╡ 62781fb2-1ee5-11eb-0ba1-035fa45106f5
begin
	plot()
	plotPath(rand(A), "Upper lattice path", :blue)
	plotPath(rand(setdiff(omega,A)), "Non-upper lattice path", :red)
	plot!([0, n], [0,n], ls=:dash, c=:black, label="")
end

# ╔═╡ 6285dea2-1ee5-11eb-2b6a-f9906c0d3351
md"## End of listing 2.5"

# ╔═╡ Cell order:
# ╟─ac3ceabe-1ee4-11eb-3935-313e96aafd62
# ╠═de7a5e8a-1ee4-11eb-2bbf-0f0eb3e489ec
# ╠═e4eeb13a-1ee4-11eb-0113-91f9a9c3b659
# ╠═d212f38c-1ee4-11eb-360a-b7d524890c8f
# ╠═6263427a-1ee5-11eb-12d1-1d576fa18461
# ╠═62637f3a-1ee5-11eb-34e2-77e37f188fe5
# ╠═6263f402-1ee5-11eb-20e1-5d1ebc78f310
# ╠═6276d39e-1ee5-11eb-23cc-2fe5377e6ce9
# ╠═62781fb2-1ee5-11eb-0ba1-035fa45106f5
# ╟─6285dea2-1ee5-11eb-2b6a-f9906c0d3351


text:
using Ipopt, JuMP, NLPModels, MathProgBase

# pass an ADNLPModel to IPOPT
nlp = ADNLPModel(x->dot(x,x), ones(2), lvar=-ones(2), uvar=ones(2),
                 c=x->[sum(x)-1], lcon=[0.0], ucon=[0.0])
show(nlp.meta)
print(nlp.meta)
model = NLPtoMPB(nlp, IpoptSolver())
@assert isa(model, Ipopt.IpoptMathProgModel)
MathProgBase.optimize!(model)
@assert MathProgBase.getobjval(model) ≈ 0.5

# Calling functions to incrase coverage
d = NLPModelEvaluator(nlp)
Jv = zeros(1)
MathProgBase.eval_jac_prod(d, Jv, zeros(2), ones(2))
Jtv = zeros(2)
MathProgBase.eval_jac_prod_t(d, Jv, zeros(2), ones(1))
Hv = zeros(2)
MathProgBase.eval_hesslag_prod(d, Hv, zeros(2), ones(2), 1.0, ones(1))
MathProgBase.isobjlinear(d)
MathProgBase.isobjquadratic(d)
MathProgBase.isconstrlinear(d, 1)

# Testing NLPtoMPB of a MathProgNLPModel
nlp = MathProgNLPModel(hs6())
model = NLPtoMPB(nlp, IpoptSolver())
MathProgBase.optimize!(model)

# Testing NLPtoMPB of a MathProgNLSModel
nls = mgh35()
model = NLPtoMPB(nls, IpoptSolver())
MathProgBase.optimize!(model)


text:
function cudnnCreate()
    handle = Ref{cudnnHandle_t}()
    cudnnCreate(handle)
    return handle[]
end

function cudnnGetProperty(property::CUDAapi.libraryPropertyType)
  value_ref = Ref{Cint}()
  cudnnGetProperty(property, value_ref)
  value_ref[]
end

version() = VersionNumber(cudnnGetProperty(CUDAapi.MAJOR_VERSION),
                          cudnnGetProperty(CUDAapi.MINOR_VERSION),
                          cudnnGetProperty(CUDAapi.PATCH_LEVEL))


text:
function createBroker(number=1)
    initstring = "-f $number --name=mainbroker --loglevel=0"
    @test_throws h.HELICSErrorInvalidArgument broker = h.helicsCreateBroker("mq", "", initstring)
    broker = h.helicsCreateBroker("zmq", "", initstring)
    @test broker isa h.Broker
    @test h.helicsBrokerIsConnected(broker) == true
    return broker
end

function setupFederateInfo(name="A Core", number=1, deltat=0.01)
    fedinitstring = "--broker=mainbroker --federates=$number"

    # Create Federate Info object that describes the federate properties
    fedinfo = h.helicsCreateFederateInfo()
    @test fedinfo isa h.FederateInfo

    # # Set Federate name
    h.helicsFederateInfoSetCoreName(fedinfo, "Test$name")

    # # Set core type from string
    h.helicsFederateInfoSetCoreTypeFromString(fedinfo, "zmq")

    # # Federate init string
    h.helicsFederateInfoSetCoreInitString(fedinfo, fedinitstring)

    # # Set the message interval (timedelta) for federate. Note th#
    # # HELICS minimum message time interval is 1 ns and by default
    # # it uses a time delta of 1 second. What is provided to the
    # # setTimedelta routine is a multiplier for the default timedelta.

    # # Set one second message interval
    h.helicsFederateInfoSetTimeProperty(fedinfo, h.HELICS_PROPERTY_TIME_DELTA, deltat)
    h.helicsFederateInfoSetIntegerProperty(fedinfo, h.HELICS_PROPERTY_INT_LOG_LEVEL, -1)
    return fedinfo
end

function createValueFederate(federates=1, name="A Federate", deltat=0.01)
    fedinfo= setupFederateInfo(name, federates, deltat)
    vFed = h.helicsCreateValueFederate("Test$name", fedinfo)
    @test vFed isa h.ValueFederate
    return vFed, fedinfo
end

function createMessageFederate(federates=1, name="A Federate", deltat=0.01)
    fedinfo = setupFederateInfo(name, federates, deltat)
    mFed = h.helicsCreateMessageFederate("Test$name", fedinfo)
    @test mFed isa h.MessageFederate
    return mFed, fedinfo
end

function destroyBroker(broker)
    h.helicsBrokerDisconnect(broker)
    h.helicsCloseLibrary()
end

function destroyFederate(fed, fedinfo, broker=nothing)
    h.helicsFederateFinalize(fed)
    state = h.helicsFederateGetState(fed)
    if broker !== nothing
        while (h.helicsBrokerIsConnected(broker))
            sleep(1)
        end
    end
    h.helicsFederateInfoFree(fedinfo)
    h.helicsFederateFree(fed)
    if broker !== nothing
        destroyBroker(broker)
    end
end

destroyValueFederate(args...) = destroyFederate(args...)
destroyMessageFederate(args...) = destroyFederate(args...)


text:
# opening and closing a file
f = open("file.out","w")
try
    write(f, "hello world")
finally
    close(f)
end

max_stars_repo_name: elucian/julia

max_issues_repo_name: elucian/julia

max_forks_repo_name: elucian/julia


text:
ϵ = 1e-13
safedivide(a, b, v=1.) = (a+(ϵ*v))/(b+ϵ)
decimal(a) = a - floor(a)
argfilter(fn, x) = [i for (i,x) in enumerate(x) if fn(x)]

function getindex(arr::AbstractArray, f::AbstractFloat, rest...)
    if f < 0.
        f = 0.
    end
    i = Int(ceil(f + ϵ))
    a = (i > 1) ? arr[i-1, rest...] : zeros(size(arr)[2:end])
    b = arr[i, rest...]
    return a .+ (b .- a) .* decimal(f)
end

function squeeze(a; dims=nothing)
    if dims === nothing
        return a[((i == 1) ? 1 : Colon() for i in size(a))...]
    else
        return a[((i ∈ dims) ? 1 : Colon() for i in 1:ndims(a))...]
    end
end

function expand(a; dims)
    dims = sort([dims...])
    i = [size(a)...]
    for d in dims
        insert!(i, d, 1)
    end
    reshape(a, i...)
end

squeezesum(a::AbstractArray; dims) = squeeze(sum(a, dims=dims), dims=dims)

function mucsum(a::AbstractArray; dims::Integer=1)
    b = copy(a)
    colonvec(i) = (tmp = Array{Any}(fill(:, ndims(a))); tmp[dims] = i; tmp)
    for i in size(a)[dims]:-1:2
        b[colonvec(i)...] -= b[colonvec(i-1)...]
    end
    b
end

function mucsum(v::AbstractVector, rng::ContinuousRange)
    i, j = firstidx(rng), lastidx(rng)
    rv = mucsum(v[i:j] .- ((i == 1) ? 0. : v[i-1]))

    if i == j
        rv .*= size(rng)
    elseif i < j
        rv[1] *= (1 - decimal(start(rng)))
        rv[end] *= decimal(stop(rng))
    end
    return rv
end

function cvalarg(cvec::AbstractVector, cval, starttrac, maxt)
    starttrac = (starttrac < 0) ? 0. : starttrac
    t = Int(ceil(starttrac + ϵ))
    while t <= maxt
        if cvec[t] > cval
            tmp = (t > 1) ? cvec[t-1] : 0.
            return (t-1) + (cval - tmp)/(cvec[t] - tmp)
        else
            t += 1
        end
    end
    return maxt
end

function rcvalarg(cvec::AbstractVector, cval, endtrac, mint)
    t = Int(floor(endtrac))
    while (t >= mint)
        if (cvec[t] < cval)# && !(cvec[t] ≈ cval)
            return (t+1) - (cvec[t+1] - cval)/(cvec[t+1] - cvec[t])
        else
            t -= 1
        end
    end
    return mint
end

function valarg(r::AbstractArray, v, starti=1, maxi=size(r)[1])
    starti = (starti < 1) ? 1 : starti
    i = starti
    while i <= maxi
        totr = sum(r[i,..])
        if v < totr
            return (i-1) + (v/totr)
        else
            v -= totr
        end
        i += 1
    end

    return maxi
end

argcut(r::AbstractArray, j, i) = squeezesum(r[i:j,..], dims=1)

function projpx(u::Array) # projection onto a probability simplex
    v = sort(u, rev=true)

    vsum, lv = 0., 0.
    for j in 1:length(u)
        if v[j] + (1.0/j)*(1 - vsum - v[j]) > 0.
            vsum += v[j]
        else
            lv = (1.0/(j-1))*(1 - vsum) # lagrange variable
            break
        end
    end

    return max.(u .+ lv, Ref(0.))
end


text:
"""
$(TYPEDEF)

Partial prior belief (absolute data) on any variable, given `<:SamplableBelief` and which dimensions of the intended variable.
"""
struct PartialPrior{T <: SamplableBelief,P <: Tuple} <: AbstractPrior
  Z::T
  partial::P
end
getSample(cf::CalcFactor{<:PartialPrior}, N::Int=1) = (reshape(rand(cf.factor.Z,N),:,N), )

"""
$(TYPEDEF)

Serialization type for `PartialPrior`.
"""
mutable struct PackedPartialPrior <: PackedInferenceType
  Z::String
  partials::Vector{Int}
end

function convert(::Type{PackedPartialPrior}, d::PartialPrior)
  PackedPartialPrior(convert(PackedSamplableBelief, d.Z), [d.partial...;])
end
function convert(::Type{PartialPrior}, d::PackedPartialPrior)
  PartialPrior(convert(SamplableBelief, d.Z),(d.partials...,))
end


text:
using JuMP
using GLPK

model_auto = Model(GLPK.Optimizer)
@variable(model_auto, 0 <= x <= 1)
@variable(model_auto, 0 <= y <= 1)
@constraint(model_auto, x + y <= 1)
@objective(model_auto, Max, x + 2y)
optimize!(model_auto)
objective_value(model_auto)

model_auto_no = Model()
@variable(model_auto_no, 0 <= x <= 1)
@variable(model_auto_no, 0 <= y <= 1)
@constraint(model_auto_no, x + y <= 1)
@objective(model_auto_no, Max, x + 2y)
set_optimizer(model_auto_no, GLPK.Optimizer)
optimize!(model_auto_no)
objective_value(model_auto_no)

model_manual = Model(GLPK.Optimizer, caching_mode = MOIU.MANUAL)
@variable(model_manual, 0 <= x <= 1)
@variable(model_manual, 0 <= y <= 1)
@constraint(model_manual, x + y <= 1)
@objective(model_manual, Max, x + 2y)
MOIU.attach_optimizer(model_manual)
optimize!(model_manual)
objective_value(model_manual)

model_direct = direct_model(GLPK.Optimizer())
@variable(model_direct, 0 <= x <= 1)
@variable(model_direct, 0 <= y <= 1)
@constraint(model_direct, x + y <= 1)
@objective(model_direct, Max, x + 2y)
optimize!(model_direct)
objective_value(model_direct)

using Cbc

model = Model(optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0));

model = Model(optimizer_with_attributes(Cbc.Optimizer, "max_iters" => 10000));

model = Model(optimizer_with_attributes(Cbc.Optimizer, "seconds" => 5));

termination_status(model_auto)

display(typeof(MOI.OPTIMAL))

primal_status(model_auto)

dual_status(model_auto)

display(typeof(MOI.FEASIBLE_POINT))

@show value(x)
@show value(y)
@show objective_value(model_auto)


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule WhiteboxTools_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("WhiteboxTools")
JLLWrappers.@generate_main_file("WhiteboxTools", UUID("8bca1bde-e0b5-5803-b535-ac666483fcca"))
end  # module WhiteboxTools_jll


text:
condtable = DataFrame(Ori=repeat([0,90,180,270],2),SpatialFreq=repeat([2,4],inner=4))
ctc = repeat(condtable,8)

@test_throws ErrorException flin(DataFrame(i=1:10))
@test_throws ErrorException flin(DataFrame(n=1:10))
@test_throws ErrorException flin(DataFrame(i=1:10,n=1:10))
@test_throws ErrorException condin(DataFrame(i=1:10))
@test_throws ErrorException condin(DataFrame(n=1:10))
@test_throws ErrorException condin(DataFrame(i=1:10,n=1:10))

cond = condin(ctc)
@test all(cond.n.==8)

fl = flin(ctc)
@test all(fl[:Ori].n.==16)
@test all(fl[:SpatialFreq].n.==32)

[condfactor(r) for r in eachrow(cond)]
@test condfactor(cond) == collect(keys(fl))

finalfactor(cond)
condstring(cond)

condresponse(rand(nrow(ctc)),cond.i)
df = condresponse(rand(nrow(ctc)),cond)
condresponse(Dict(1=>rand(nrow(ctc))),cond)
condresponse(Dict(2=>rand(nrow(ctc))),ctc,[:Ori])

factorresponse(df)

findcond(cond,Ori=0,SpatialFreq=2)
@test isresponsive(randn(nrow(ctc)),randn(nrow(ctc)).+1)
@test isresponsive(randn(nrow(ctc)),randn(nrow(ctc)).+1,cond.i)
ismodulative([DataFrame(Y=randn(nrow(ctc))) ctc])


text:
# Random phase approxiamtion for electron gas

using QuantumStatistics

"""
3D Imaginary-time effective interaction derived from random phase approximation.
Return ``dW_0(q, τ)/v_q`` where ``v_q`` is the bare Coulomb interaction and ``dW_0`` is the dynamic part of the effective interaction.

The total effective interaction can be recoverd using,
```math
W_0(q, τ) = v_q δ(τ) + dW_0(q, τ).
```

The dynamic contribution is the fourier transform of,
```math
dW_0(q, iω_n)=v_q^2 Π(q, iω_n)/(1-v_q Π(q, iω_n))
```
Note that this dynamic contribution ``dW_0'' diverges at small q. For this reason, this function returns ``dW_0/v_q``
"""
function dWRPA(vqinv, qgrid, τgrid, kF, β, spin, mass)
    @assert all(qgrid .!= 0.0)
    EF = kF^2 / (2mass)
    dlr = DLR.DLRGrid(:corr, 10EF, β, 1e-10) # effective interaction is a correlation function of the form <O(τ)O(0)>
    Nq, Nτ = length(qgrid), length(τgrid)
    Π = zeros(Complex{Float64}, (Nq, dlr.size)) # Matsubara grid is the optimized sparse DLR grid
    dW0norm = similar(Π)
    for (ni, n) in enumerate(dlr.n)
        for (qi, q) in enumerate(qgrid)
            Π[qi, ni] = TwoPoint.LindhardΩnFiniteTemperature(3, q, n, kF, β, mass, spin)[1]
        end
        dW0norm[:, ni] = @. Π[:, ni] / (vqinv - Π[:, ni])
        # println("ω_n=2π/β*$(n), Π(q=0, n=0)=$(Π[1, ni])")
        # println("$ni  $(dW0norm[2, ni])")
    end
    dW0norm = DLR.matfreq2tau(:corr, dW0norm, dlr, τgrid, axis=2) # dW0/vq in imaginary-time representation, real-valued but in complex format

    # println(dW0norm[1, :])
    # println(DLR.matfreq2tau(:corr, dW0norm[1, :], dlr, τgrid, axis=1))
    # println(DLR.matfreq2tau(:corr, dW0norm[2, :], dlr, τgrid, axis=1))
    # coeff = DLR.matfreq2dlr(:corr, dW0norm[1, :], dlr)
    # fitted = DLR.dlr2matfreq(:corr, coeff, dlr, dlr.n)
    # for (ni, n) in enumerate(dlr.n)
    #     println(dW0norm[1, ni], " vs ", fitted[ni])
    # end
    return real.(dW0norm)
end

if abspath(PROGRAM_FILE) == @__FILE__
    using Gaston
    include("parameter.jl")

    qgrid = Grid.boseK(kF, 3kF, 0.2kF, 32)
    τgrid = Grid.tau(β, EF / 20, 128)
    # println("qGrid: ", qgrid.grid)
    println("τGrid: ", τgrid.grid)
    vqinv = [(q^2 + mass2) / (4π * e0^2) for q in qgrid.grid] # instantaneous interaction (Coulomb interaction)

    dW0norm = dWRPA(vqinv, qgrid.grid, τgrid.grid, kF, β, spin, me)
    for (qi, q) in enumerate(qgrid.grid)
        println("$q   $(dW0norm[qi, 1])")
    end

    display(plot(qgrid.grid ./ kF, dW0norm[:, 1]))
    sleep(100)
end


text:
using Base.Threads, SortingAlgorithms

import SortingAlgorithms: StringRadixSort, StringRadixSortAlg, uint_mapping, load_bits
import StatsBase: BaseRadixSortSafeTypes
import Base: Forward, ForwardOrdering, ReverseOrdering, sort!, Reverse

# winner from benchmarks/which_is_the_fastest_UInt_histogram.jl
function uint_hist(bits::Vector{T}) where T <: Unsigned
    iter = sizeof(T)÷2
    hist = zeros(UInt32, 65536, iter, nthreads())
    @threads for j = 1:length(bits)
        for i = 0:iter-1
            @inbounds hist[1+Int((bits[j] >> (i << 4)) & 0xffff), i+1, threadid()] += 1
        end
    end
    @threads for j in 1:iter
        for i = 2:nthreads()
            @inbounds hist[:, j, 1] .+= hist[:, j, i]
        end
    end
    hist[:,:,1]
end

# sort it by using sorttwo! on the pointer
function sort_1fasterhistogram!(svec::Vector{String}, lo::Int, hi::Int, ::StringRadixSortAlg, o::O) where O <: Union{ForwardOrdering, ReverseOrdering}
    if lo >= hi;  return svec;  end
    # the length subarray to sort
    l = hi - lo + 1

    # find the maximum string length
    lens = maximum(sizeof, svec)
    skipbytes = lens
    # ptrs = pointer.(svec)
    if lens > 0
        while lens > 4
            skipbytes = max(0, skipbytes - 8)
            bits64 = zeros(UInt64, l)
            if o == Reverse
                bits64[lo:hi] .= .~load_bits.(UInt64, @view(svec[lo:hi]), skipbytes)
                SortingLab.sorttwo!(bits64, svec, lo, hi)
            else
                bits64[lo:hi] .= load_bits.(UInt64, @view(svec[lo:hi]), skipbytes)
                SortingLab.sorttwo!(bits64, svec, lo, hi)
            end
            lens -= 8
        end
        if lens > 0
            skipbytes = max(0, skipbytes - 4)
            bits32 = zeros(UInt32, l)
            if o == Reverse
                bits32[lo:hi] .= .~load_bits.(UInt32, @view(svec[lo:hi]), skipbytes)
                SortingLab.sorttwo!(bits32, svec, lo, hi)
            else
                bits32[lo:hi] .= load_bits.(UInt32, @view(svec[lo:hi]), skipbytes)
                SortingLab.sorttwo!(bits32, svec, lo, hi)
            end
            lens -= 4
        end
    end
    # unsafe_pointer_to_objref.(ptrs - 8)
    svec
end

"""
    sorttwo!(vs, index)

Sort both the `vs` and reorder `index` at the same. This allows for faster sortperm
for radix sort.
"""
function sorttwo!(vs::Vector{T}, index, lo::Int = 1, hi::Int=length(vs), RADIX_SIZE = 16, RADIX_MASK = 0xffff) where T <:BaseRadixSortSafeTypes
    # Input checking
    if lo >= hi;  return (vs, index);  end

    o = Forward

    # Init
    iters = ceil(Integer, sizeof(T)*8/RADIX_SIZE)
    # number of buckets in the counting step
    nbuckets = 2^RADIX_SIZE

    # Histogram for each element, radix
    bin = uint_hist(vs)

    # bin = zeros(UInt32, nbuckets, iters)
    # if lo > 1;  bin[1,:] = lo-1;  end

    # Sort!
    swaps = 0
    len = hi-lo+1

    index1 = similar(index)
    ts=similar(vs)
    for j = 1:iters
        # Unroll first data iteration, check for degenerate case
        v = uint_mapping(o, vs[hi])
        idx = Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK) + 1

        # are all values the same at this radix?
        if bin[idx,j] == len;  continue;  end

        # cbin = cumsum(bin[:,j])
        # tries to achieve the above one-liner with more efficiency
        cbin = zeros(UInt32, nbuckets)
        cbin[1] = bin[1,j]
        for i in 2:nbuckets
            cbin[i] = cbin[i-1] + bin[i,j]
        end

        ci = cbin[idx]
        ts[ci] = vs[hi]
        index1[ci] = index[hi]
        cbin[idx] -= 1

        # Finish the loop...
        @inbounds for i in hi-1:-1:lo
            v = uint_mapping(o, vs[i])
            idx = Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK) + 1
            ci = cbin[idx]
            ts[ci] = vs[i]
            index1[ci] = index[i]
            cbin[idx] -= 1
        end
        vs,ts = ts,vs
        index, index1 = index1, index
        swaps += 1
    end

    if isodd(swaps)
        vs,ts = ts,vs
        index, index1 = index1, index
        for i = lo:hi
            @inbounds vs[i] = ts[i]
            @inbounds index[i] = index1[i]
        end
    end
    (vs, index)
end

using SortingLab
using Base.Test
using SortingAlgorithms
import SortingAlgorithms: load_bits

# write your own tests here
# @test 1 == 2
N = 100_000_000
K = 100
samplespace = "id".*dec.(1:N÷K,10);
srand(1)
svec = rand(samplespace, N);

using BenchmarkTools

# overall sorting is faster too
function ghi(svec)
    csvec = copy(svec)
    gc()
    @time a = @elapsed sort_1fasterhistogram!(csvec, 1, length(csvec), StringRadixSort, Base.Forward)

    csvec = copy(svec)
    gc()
    @time b = @elapsed SortingAlgorithms.sort!(csvec, 1, length(csvec), StringRadixSort, Base.Forward)
    a,b
end
@time [ghi(svec) for i=repeat([4], outer=5)]


text:
##################################################################
function multigp(x, y, σ2, idx; kernel=OU, maxiter = 10, maxrandom = 100, seed = 1, show_trace = true)
##################################################################

    #-----------------------------------------------------------------
    # Preliminary stuff, check dimensions, set random seed...
    #-----------------------------------------------------------------

    # check input parameters: there should be as many inputs x as outputs y as
    # many squared errors σ2 as many labels in idx
    @assert(length(x) == length(y) == length(σ2) == length(idx))

    # make sure labels run continuously from 1 to maximum number of labels
    @assert(all(unique(idx) .== collect(1:length(unique(idx)))))

    # instantiate random number generator
    rg = MersenneTwister(seed)

    # How much pertubation for numerical stability
    JITTER = 1e-8

    #-----------------------------------------------------------------
    # Split into seperate datasets according to identifier in `idx`
    #-----------------------------------------------------------------

    X =  [x[findall(idx .== i)] for i in unique(idx)]

    Y =  [y[findall(idx .== i)] for i in unique(idx)]

    Σ =  [Diagonal(σ2[findall(idx .== i)]) for i in unique(idx)]

    D = length(Y)

    #-----------------------------------------------------------------
    # pre-calculate distances and pre-allocate kernel matrices
    #-----------------------------------------------------------------

    SqDist = [pairwise(SqEuclidean(), reshape(x, 1, length(x))) for x in X]

    K = [zeros(length(x), length(x)) for x in X]

    #-----------------------------------------------------------------
    # Various tests for checking dimensions
    #-----------------------------------------------------------------

    for d in 1:D
        @assert(size(SqDist[d], 1) == size(SqDist[d], 2) == length(X[d]))
    end

    @assert(all(map(length, X) .== map(length, Y) .== map(s->size(s, 1), Σ) .== map(s->size(s, 2), Σ)))

    @assert(D == length(SqDist) == length(X) == length(Y) == length(Σ) == length(K))

    #-----------------------------------------------------------------
    # Report
    #-----------------------------------------------------------------

    @printf("Running multigp\n")

    @printf("\t There is a total of %d observations\n", length(x))

    @printf("\t There are a %d dimensions\n", D)

    #----------------------------------------------------
    # Auxiliary for sorting out parameters for GP
    #----------------------------------------------------

    function unpack(hyper)
        local logα = hyper[1:D]
        local logs = hyper[D+1]
        return logα, logs
    end

    #----------------------------------------------------
    # Define objective as marginal log-likelihood
    #----------------------------------------------------

    function objective(hyper)

        local logα, logs = unpack(hyper)

        for i in 1:length(K)
            calculatekernelmatrix!(K[i], SqDist[i], kernel, [logα[i]; logs])
            K[i] += Σ[i] + JITTER*I
        end

        f = zero(eltype(hyper))

        try

            # below is Eq. 2.30 in RW for D dimensions
            for i in 1:D
                C = cholesky(K[i]).L
                #f += -0.5*(Y[i]'*(K[i]\Y[i])) - 0.5*logdet(K[i])
                f += -0.5*sum(abs2.(C\Y[i])) - 0.5*2*sum(log.(diag(C)))
            end

        catch err

            if isa(err, PosDefException)

                @warn("Cholesky failed, matrix not posdef. Return -Inf.")

                return Inf

            else

                throw(err)

            end

        end

        -1.0 * f # we are minimising

    end

    #--------------------------------------------------------
    # Initialise GP hyperparameters with random search
    # and few iterations
    #--------------------------------------------------------

    initsolutions       = [randn(rg, 2*D+1)*7.5 for i=1:maxrandom]

    initoptimise(x)     = Optim.optimize(objective, x, NelderMead(),
                                Optim.Options(iterations=200)).minimum

    initialfitness      = @showprogress "\t init random search " map(initoptimise, initsolutions)

    bestinitialsolution = initsolutions[argmin(initialfitness)]

    #----------------------------------------------------
    # Optimise hyperparameters
    #----------------------------------------------------

    result = Optim.optimize(objective, bestinitialsolution, NelderMead(),
        Optim.Options(show_every=100, show_trace=show_trace, iterations=maxiter))

    # instantiate learnt matrix and observed variance parameter
    logα, logs = unpack(result.minimizer)

    for i in 1:length(K)

        calculatekernelmatrix!(K[i], SqDist[i], kernel, [logα[i]; logs])

        K[i] += Σ[i] + JITTER*I

    end

    #----------------------------------------------------
    # return prediction function
    #----------------------------------------------------

    function predictTest(xtest, d)

        # dimensions: N × Ntest
        k = calculatekernelmatrix(X[d], xtest, kernel,  [logα[d]; logs])

        # Ntest × 1
        c = calculatekernelmatrix(xtest, xtest, kernel, [logα[d]; logs])

        # predictive covariance and mean
        Σpred = c-k'*(K[d]\k)  + I*JITTER

        Σpred = 0.5*(Σpred+Σpred')

        # D × Ntest
        μpred = k'*(K[d]\Y[d])

        return μpred, Σpred

    end

    return logs, predictTest

end


text:
load("cuda.jl")
load("cufft.jl")

# The CuMatrix class
type CuMatrix
    T::Type
    ptr::Ptr
    dims::(Integer,Integer)

    # Construct Matrix on device
    function CuMatrix(T::Type, dims::(Integer, Integer))
        if T != Float32   && T != Float64 &&
           T != Complex64 && T != Complex128
            error("No integer or boolean support yet")
        end
        ptr = cuda_malloc(T, dims)
        Res = new(T, ptr, dims)
        finalizer(Res, CuFree)
        Res
    end

    # Copy Matrix from host
    function CuMatrix{T}(in::Matrix{T})
        dims = size(in)
        ptr = cuda_malloc(T, dims)
        mem_device(ptr, in)
        Res = new(T, ptr, dims)
        finalizer(Res, CuFree)
        Res
    end

    # Constructor with existing device pointer
    function CuMatrix{T}(ptr::Ptr{T}, dims::(Integer, Integer))
        Res = new(T, ptr, dims)
        finalizer(Res, CuFree)
        Res
    end

    # Aliasing
    CuMatrix(T::Type, rows::Integer, cols::Integer) = CuMatrix(T, (rows, cols))
    CuMatrix(T::Type, rows::Integer) = CuMatrix(T, (rows, 1))

    # Default type: Single precision (for compatibility with older cards)
    CuMatrix(rows::Integer, cols::Integer) = CuMatrix(Float32, rows, cols)
    CuMatrix(rows::Integer) = CuMatrix(Float32, rows, 1)
end

# Get matrix from device to host
function Array(in::CuMatrix)
    out = Array(eltype(in), in.dims[1], in.dims[2])
    mem_host(out, in.ptr)
    return out
end

# Perform a deep copy
function copy(in::CuMatrix)
    ptr = cuda_malloc(eltype(in), in.dims)
    bytes::Int32 = numel(in) * sizeof(eltype(in))
    mem_copy(ptr, in.ptr, bytes)
    CuMatrix(ptr, in.dims)
end

# Display function
function print(in::CuMatrix)
    print("On GPU\n", Array(in), "\n")
end

# Display function
function show(in::CuMatrix)
    print(in)
end

# Return number of elements
numel(A::CuMatrix) = A.dims[1]*A.dims[2]
numel(T::Type, A::CuMatrix) = convert(T, numel(A))

# Return type of elements
eltype(A::CuMatrix) = A.T

# Freeing memory
function CuFree(in::CuMatrix)
    cuda_free(in.ptr)
end

# Random Number Generators
function curand(T::Type, rows::Integer, cols::Integer)
    out = CuMatrix(T, rows, cols)
    cuda_rand(out.ptr, rows * cols)
    return out
end

function curandn(T::Type, rows::Integer, cols::Integer)
    # curand only supports randn for even number of elements
    count = rows * cols
    count += count & 1 # Increase the malloc size
    ptr = cuda_malloc(T, count)
    cuda_randn(ptr, count)
    # But only use the required number of elements
    CuMatrix(T, ptr, (rows, cols))
end

curand(rows::Integer, cols::Integer) = curand(Float32, rows, cols)
curandn(rows::Integer, cols::Integer) = curandn(Float32, rows, cols)

# BLAS Functions
function (*)(A::CuMatrix, B::CuMatrix)
    if (A.dims[2] != B.dims[1])
        error("Inner dimension mismatch in Matrix multiply")
    end
    if (eltype(A) != eltype(B))
        error("Precision mismatch in Matrix multiply")
    end
    C = CuMatrix(eltype(A), (A.dims[1], B.dims[2]))

    m = convert(Int32, C.dims[1])
    n = convert(Int32, C.dims[2])
    k = convert(Int32, B.dims[1])

    cuda_gemm('N', 'N', m, n, k,
            one(A.T), A.ptr, m, B.ptr, k,
            zero(A.T), C.ptr, m)
    return C
end

amax(A::CuMatrix) = cuda_amax(numel(Int32, A), A.ptr)
amin(A::CuMatrix) = cuda_amin(numel(Int32, A), A.ptr)
asum(A::CuMatrix) = cuda_asum(numel(Int32, A), A.ptr)

function (*)(A::CuMatrix, alpha)
    n = numel(Int32, A)
    B = copy(A)
    cuda_scal(n, B.ptr, alpha)
    return B
end

(*)(alpha, A::CuMatrix) = (*)(A::CuMatrix, alpha)

function dot(A::CuMatrix, B::CuMatrix)
    if eltype(A) != eltype(B)
        error("Precision mismatch in Dot product")
    end

    n = numel(Int32, A)
    m = numel(Int32, A)

    if m != n
        error("Size mismatch in Dot product")
    end

    cuda_dot(n, A.ptr, B.ptr)
end

# Euclidean norm
nrm2(A::CuMatrix) = cuda_nrm2(numel(Int32, A), A.ptr)

# Thrust functions

# lapack functions
function norm(A::CuMatrix, p)
    if (p == 2)
        nrm2(A::CuMatrix)
    else
        error("norm not supported for p == ", p)
    end
end

# CUFFT functions
cufftExec(plan::Uint32, idata::Ptr{Complex64}, odata::Ptr{Complex64}, direction::Int32) = cufftExecC2C(plan, idata, odata, direction)
cufftExec(plan::Uint32, idata::Ptr{Float32}, odata::Ptr{Complex64}) = cufftExecR2C(plan, idata, odata)
cufftExec(plan::Uint32, idata::Ptr{Complex64}, odata::Ptr{Float32}) = cufftExecC2R(plan, idata, odata)
cufftExec(plan::Uint32, idata::Ptr{Complex128}, odata::Ptr{Complex128}, direction::Int32) = cufftExecZ2Z(plan, idata, odata, direction)
cufftExec(plan::Uint32, idata::Ptr{Float64}, odata::Ptr{Complex128}) = cufftExecD2Z(plan, idata, odata)
cufftExec(plan::Uint32, idata::Ptr{Complex128}, odata::Ptr{Float64}) = cufftExecZ2D(plan, idata, odata)
cufftExec(plan::Uint32, idata::CuMatrix, odata::CuMatrix, direction::Int32) = cufftExec(plan, idata.ptr, odata.ptr, direction)
cufftExec(plan::Uint32, idata::CuMatrix, odata::CuMatrix) = cufftExec(plan, idata.ptr, odata.ptr)

max_stars_repo_name: stefan-k/julia-CuMatrix

max_issues_repo_name: stefan-k/julia-CuMatrix

max_forks_repo_name: stefan-k/julia-CuMatrix


text:
"""
$(TYPEDEF)

A simple feedback linearizable unicycle model.

# Fields

$(TYPEDFIELDS)
"""
struct Unicycle4D{ΔT} <: ControlSystem{ΔT, 4, 2} end

function dx(cs::Unicycle4D, x::SVector{4}, u::SVector{2}, t::AbstractFloat)
    # position: x
    dx1 = x[4] * cos(x[3])
    # position: y
    dx2 = x[4] * sin(x[3])
    # orientation
    dx3 = u[1]
    # velocity
    dx4 = u[2]

    return @SVector[dx1, dx2, dx3, dx4]
end

function linearize_discrete(cs::Unicycle4D, x::SVector{4}, u::SVector{2},
                            t::AbstractFloat)
    ΔT = samplingtime(cs)
    cθ = cos(x[3]) * ΔT
    sθ = sin(x[3]) * ΔT

    A = @SMatrix [1 0 -x[4]*sθ cθ;
                  0 1 x[4]*cθ  sθ;
                  0 0 1        0;
                  0 0 0        1]

    B = @SMatrix [0  0;
                  0  0;
                  ΔT 0;
                  0  ΔT];
    return LinearSystem{ΔT}(A, B)
end

xyindex(cs::Unicycle4D) = SVector(1, 2)

"------------------- Implement Feedback Linearization Interface -------------------"

n_linstates(cs::Unicycle4D) = 4
ξxyindex(cs::Unicycle4D) = SVector(1, 3)

# TODO: Maybe use StaticArrays.FieldVector here to have a type safe desctiction
# between ξ and x coordinates.

# implement holy trait
LinearizationStyle(::Unicycle4D) = FeedbackLinearization()

@inline function feedbacklin(cs::Unicycle4D)
    ΔT = samplingtime(cs)
    # ξ = (px, pẋ, py, pẏ)
    A = @SMatrix [1. ΔT 0. 0.;
                  0. 1. 0. 0.;
                  0. 0. 1. ΔT;
                  0. 0. 0. 1.;]
    B = @SMatrix [0. 0.;
                  ΔT 0.;
                  0. 0.;
                  0. ΔT];

    dyn = LinearSystem{ΔT}(A, B)
    return LTISystem(dyn, ξxyindex(cs))
end

function x_from(cs::Unicycle4D, ξ::SVector{4})
    # px, py, θ, v
    return @SVector [ξ[1], ξ[3], atan(ξ[4], ξ[2]), sqrt(ξ[2]^2 + ξ[4]^2)]
end

function ξ_from(cs::Unicycle4D, x::SVector{4})
    # px, pẋ, py, pẏ
    return @SVector [x[1], cos(x[3])*x[4], x[2], sin(x[3])*x[4]]
end

function λ_issingular(cs::Unicycle4D, ξ::SVector{4})
    vx = ξ[2]; vy = ξ[4]
    ϵ = 0.01
    # cant invert for velocities near 0
    return isnan(vx) || isnan(vy) || (abs(vx) < ϵ && abs(vy) < ϵ)
end

function inverse_decoupling_matrix(cs::Unicycle4D, x::SVector{4})
    cθ = cos(x[3]); sθ = sin(x[3]); v = x[4]
    return @SMatrix[cθ    sθ;
                   -sθ/v -cθ/v]
end

@inline decoupling_drift_term(cs::Unicycle4D, x::SVector{4}) = @SVector(zeros(n_controls(cs)))


text:
# Hello world in Julia

println("Hello, World!")


text:
module SparseMatrixAssemblers

using Test
using Gridap.Arrays
using Gridap.TensorValues
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.Integration
using Gridap.Fields
using Gridap.Algebra
using SparseArrays
using Gridap.FESpaces
using Gridap.CellData

domain =(0,1,0,1)
partition = (2,2)
model = CartesianDiscreteModel(domain,partition)

order = 1
grid_topology = get_grid_topology(model)
polytopes = get_polytopes(grid_topology)
reffes = [LagrangianRefFE(Float64,p,order) for p in polytopes]

dirichlet_tags = [1,2,3,4,6,5]
V = GradConformingFESpace(reffes,model,dirichlet_tags)

U = TrialFESpace(V)

b(x) = x[2]

v = get_cell_basis(V)
u = get_cell_basis(U)

trian = get_triangulation(model)
degree = 2
quad = CellQuadrature(trian,degree)

btrian = BoundaryTriangulation(model)
bquad = CellQuadrature(btrian,degree)

bu = restrict(u,btrian)
bv = restrict(v,btrian)

cellmat = integrate(∇(v)⊙∇(u),trian,quad)
cellvec = integrate(v⊙b,trian,quad)
cellmatvec = pair_arrays(cellmat,cellvec)
cellids = collect(1:num_cells(trian))

bcellmat = integrate(bv*bu,btrian,bquad)
bcellvec = integrate(bv*3,btrian,bquad)
bcellmatvec = pair_arrays(bcellmat,bcellvec)
bcellids = get_face_to_cell(btrian)

term_to_cellmat = [cellmat, bcellmat]
term_to_cellvec = [cellvec, bcellvec]
term_to_cellids = [cellids, bcellids]
term_to_cellmatvec = [ cellmatvec, bcellmatvec ]

mtypes = [
  SparseMatrixCSC,
  SparseMatrixCSR,
  SparseMatrixCSR{0},
  SparseMatrixCSR{1},
  SparseMatrixCSR{0,Float64, Int},
  SparseMatrixCSR{1,Float64, Int},
  SymSparseMatrixCSR,
  SymSparseMatrixCSR{0},
  SymSparseMatrixCSR{1},
  SymSparseMatrixCSR{0, Float64, Int},
  SymSparseMatrixCSR{1, Float64, Int}]

for T in mtypes

  matvecdata = ( term_to_cellmatvec , term_to_cellids, term_to_cellids)
  matdata = (term_to_cellmat,term_to_cellids,term_to_cellids)
  vecdata = (term_to_cellvec,term_to_cellids)
  data = (matvecdata,matdata,vecdata)

  assem = SparseMatrixAssembler(T,Vector{Float64},U,V)
  test_sparse_matrix_assembler(assem,matdata,vecdata,data)

  matdata = ([cellmat],[cellids],[cellids])
  vecdata = ([cellvec],[cellids])

  mat = assemble_matrix(assem,matdata)
  vec = assemble_vector(assem,vecdata)

  x = mat \ vec

  assemble_matrix!(mat,assem,matdata)
  assemble_vector!(vec,assem,vecdata)

  x2 = mat \ vec
  @test x ≈ x2

  @test vec ≈ [0.0625, 0.125, 0.0625]
  @test mat[1, 1]  ≈  1.333333333333333
  @test mat[2, 1]  ≈ -0.33333333333333
  @test mat[1, 2]  ≈ -0.33333333333333
  @test mat[2, 2]  ≈ 2.666666666666666
  @test mat[3, 2]  ≈ -0.33333333333333
  @test mat[2, 3]  ≈ -0.33333333333333
  @test mat[3, 3]  ≈ 1.333333333333333

  data = (([cellmatvec],[cellids],[cellids]),([],[],[]),([],[]))
  mat, vec = allocate_matrix_and_vector(assem,data)
  assemble_matrix_and_vector!(mat,vec,assem,data)
  assemble_matrix_and_vector!(mat,vec,assem,data)

  @test vec ≈ [0.0625, 0.125, 0.0625]
  @test mat[1, 1]  ≈  1.333333333333333
  @test mat[2, 1]  ≈ -0.33333333333333

  x3 = mat \ vec
  @test x ≈ x3

  mat, vec = assemble_matrix_and_vector(assem,data)

  x4 = mat \ vec
  @test x ≈ x4

  @test vec ≈ [0.0625, 0.125, 0.0625]
  @test mat[1, 1]  ≈  1.333333333333333
  @test mat[2, 1]  ≈ -0.33333333333333

end

strian = SkeletonTriangulation(model)
squad = CellQuadrature(strian,degree)

su = restrict(u,strian)
sv = restrict(v,strian)

scellmat = integrate(jump(sv)*su.⁻,strian,squad)
scellvec = integrate(mean(sv*3),strian,squad)
scellmatvec = pair_arrays(scellmat,scellvec)
scellids = get_cell_id(strian)
zh = zero(V)
scellvals = get_cell_values(zh,scellids)
scellmatvec = attach_dirichlet(scellmatvec,scellvals)

matvecdata = ([scellmatvec],[scellids],[scellids])
matdata = ([scellmat],[scellids],[scellids])
vecdata = ([scellvec],[scellids])
data = (matvecdata,matdata,vecdata)

assem = SparseMatrixAssembler(U,V)
test_sparse_matrix_assembler(assem,matdata,vecdata,data)

A = assemble_matrix(assem,matdata)
@test A == zeros(num_free_dofs(V),num_free_dofs(U))

# Allowing a non-blocked matrix in a context where you would expect a blocked one.
scellmat = map(i->zeros(size(i)),scellmat)
matdata = ([scellmat],[scellids],[scellids])
A = assemble_matrix(assem,matdata)
@test A == zeros(num_free_dofs(V),num_free_dofs(U))

end # module


text:
using SpecialFunctions
using LinearAlgebra

function bessels_ref!(j, j′, y, y′, x)
    for i = 1:length(j)
        n = i - 1
        ν = (2n + 1)/2
        j[i] = besselj(ν, x)
        y[i] = bessely(ν, x)
    end

    ω = √(π/2x)
    lmul!(ω, j)
    lmul!(ω, y)

    j′[1] = -y[1]-j[1]/x
    y′[1] = j[1]-y[1]/x

    for i = 2:length(j)
        n = i - 1
        ν = (n + 1)/x
        j′[i] = j[i-1] - ν*j[i]
        y′[i] = y[i-1] - ν*y[i]
    end
end


text:
using OrdinaryDiffEq
using DiffEqCallbacks
using Surrogates
using Conda
using PyCall
using Random
using Plots
using BenchmarkTools;

env = Conda.ROOTENV
Conda.pip_interop(true, env)
Conda.pip("install", "mogp-emulator");

random = pyimport("random")
np = pyimport("numpy")
mogp = pyimport("mogp_emulator");

Random.seed!(123)
random.seed(123)
np.random.seed(123);

function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,γ) = p
    @inbounds begin
        du[1] = -β*S*I
        du[2] = β*S*I - γ*I
        du[3] = γ*I
    end
    nothing
end;

tmax = 10000.0
tspan = (0.0,tmax)
δt = 1.0;

n_train = 50 # Number of training samples
n_test = 1000; # Number of test samples

# Parameters are β, γ
lb = [0.00005, 0.1]
ub = [0.001, 1.0];

N = 1000.0
u0 = [990.0,10.0,0.0]
p = [0.0005,0.25]
prob_ode = ODEProblem(sir_ode!,u0,tspan,p);

sampler = LatinHypercubeSample();

θ = Surrogates.sample(n_train,lb,ub,sampler);

logit = (x) -> log(x/(1-x))
invlogit = (x) -> exp(x)/(exp(x)+1.0)
cb_ss = TerminateSteadyState()
logit_final_size = function(z)
  prob = remake(prob_ode;p=z)
  sol = solve(prob, ROS34PW3(),callback=cb_ss)
  fsp = sol[end][3]/N
  logit(fsp)
end;

lfs = logit_final_size.(θ);

gp = mogp.GaussianProcess(θ, lfs, nugget="fit");

gp = mogp.fit_GP_MAP(gp, n_tries=100);

lfs_train_pred = gp.predict(θ);

scatter(invlogit.(lfs),
        invlogit.(lfs_train_pred["mean"]),
        xlabel = "Model final size",
        ylabel = "Surrogate final size",
        legend = false,
        title = "Training set")

θ_test = sample(n_test,lb,ub,sampler)
lfs_test = logit_final_size.(θ_test)
lfs_test_pred = gp.predict(θ_test);

scatter(invlogit.(lfs_test),
        invlogit.(lfs_test_pred["mean"]),
        xlabel = "Model final size",
        ylabel = "Surrogate final size",
        legend = false,
        title = "Test set")

β_grid = collect(lb[1]:0.00001:ub[1])
θ_eval = [[βᵢ,0.25] for βᵢ in β_grid]
lfs_eval = gp.predict(θ_eval)
fs_eval = invlogit.(lfs_eval["mean"])
fs_eval_uc = invlogit.(lfs_eval["mean"] .+ 1.96 .* sqrt.(lfs_eval["unc"]))
fs_eval_lc = invlogit.(lfs_eval["mean"] .- 1.96 .* sqrt.(lfs_eval["unc"]))
plot(β_grid,
     fs_eval,
     xlabel = "Infectivity parameter, β",
     ylabel = "Final size",
     label = "Model")
plot!(β_grid,
      invlogit.(logit_final_size.(θ_eval)),
      ribbon = (fs_eval .- fs_eval_lc, fs_eval_uc - fs_eval),
      label = "Surrogate",
      legend = :right)

γ_grid = collect(lb[2]:0.001:ub[2])
θ_eval = [[0.001,γᵢ] for γᵢ in γ_grid]
lfs_eval = gp.predict(θ_eval)
fs_eval = invlogit.(lfs_eval["mean"])
fs_eval_uc = invlogit.(lfs_eval["mean"] .+ 1.96 .* sqrt.(lfs_eval["unc"]))
fs_eval_lc = invlogit.(lfs_eval["mean"] .- 1.96 .* sqrt.(lfs_eval["unc"]))
plot(γ_grid,
     fs_eval,
     xlabel = "Recovery rate, γ",
     ylabel = "Final size",
     label = "Model")
plot!(γ_grid,
      invlogit.(logit_final_size.(θ_eval)),
      ribbon = (fs_eval .- fs_eval_lc, fs_eval_uc - fs_eval),
      label = "Surrogate")

obs = logit_final_size(p)
invlogit(obs)

hm = mogp.HistoryMatching(gp=gp,
                          obs=obs,
                          coords=np.array(θ_test),
                          threshold=3.0);

nroy_points = hm.get_NROY() .+ 1
length(nroy_points),n_test

x = [θᵢ[1] for θᵢ in θ_test]
y = [θᵢ[2] for θᵢ in θ_test]
l = @layout [a b]
pl1 = histogram(x[nroy_points],legend=false,xlim=(lb[1],ub[1]),bins=lb[1]:0.00005:ub[1],title="NROY values for β")
vline!(pl1,[p[1]])
pl2 = histogram(y[nroy_points],legend=false,xlim=(lb[2],ub[2]),bins=lb[2]:0.05:ub[2],title="NROY values for γ")
vline!(pl2,[p[2]])
plot(pl1, pl2, layout = l)

@benchmark logit_final_size(p)

@benchmark gp.predict(p)


text:
include("roblib.jl")

inverse(side::HorizonSide)=HorizonSide(mod(Int(side)+2,4))
next(side::HorizonSide)=HorizonSide(mod(Int(side)+3,4))

function mark_cross!(r::Robot)
    for side in (Nord,Ost,Sud,West)
        putmarkers!(r,side)
        if isborder(r,side)==true && isborder(r,next(side))==false
            obhod(r,side)
            putmarker!(r)
            putmarkers!(r,side)
        end
        side_inv=inverse(side)
        move_by_marker!(r,side_inv)
    end
putmarker!(r)
end

function putmarkers!(r::Robot,side::HorizonSide)
    while isborder(r,side)==false
        move!(r,side)
        putmarker!(r)
    end
end

function move_by_marker!(r::Robot,side::HorizonSide)
    while ismarker(r)==true
        if isborder(r,side)==false
            move!(r,side)
        elseif isborder(r,side)==true && isborder(r,next(side))==false
            obhod(r,side)
        end
    end
end

mark_cross!(r)


text:
function same_geometry(a::AbstractVTKUnstructuredData, b::AbstractVTKUnstructuredData)
    dim(a) != dim(b) && return false
    num_of_points(a) != num_of_points(b) && return false
    num_of_cells(a) != num_of_cells(b) && return false

    _num_of_points = num_of_points(a)
    b_image_point_inds = zeros(Int, _num_of_points)
    for i in 1:_num_of_points
        for j in 1:_num_of_points
            @views if b.point_coords[:,i] == a.point_coords[:,j]
                b_image_point_inds[i] = j
                break
            end
        end
        if b_image_point_inds[i] == 0
            return false
        end
    end

    function map_cell_connectivity(b_cell_connectivity)
        b_cell_connectivity_mapped = zeros(b_cell_connectivity)
        for i in 1:length(b_cell_connectivity)
            b_cell_connectivity_mapped[i] = b_image_point_inds[b_cell_connectivity[i]]
        end
        b_cell_connectivity_mapped
    end

    _num_of_cells = num_of_cells(a)
    b_cell_connectivity_mapped = map(map_cell_connectivity, b.cell_connectivity)
    b_image_cell_inds = zeros(Int, _num_of_cells)
    for i in 1:_num_of_cells
        for j in 1:_num_of_cells
            if similar_cells(b_cell_connectivity_mapped[i], cell_type(b,i), a.cell_connectivity[j], cell_type(a,j))
                b_image_cell_inds[i] = j
                break
            end
        end
        if b_image_cell_inds[i] == 0
            return false
        end
    end
    return true
end
function same_geometry(a::VTKStructuredData, b::VTKStructuredData)
    return a.point_coords == b.point_coords
end
function same_geometry(a::VTKRectilinearData, b::VTKRectilinearData)
    return extents(a) == extents(b) && all(a.point_coords .== b.point_coords)
end
function same_geometry(a::VTKUniformRectilinearData, b::VTKUniformRectilinearData)
    return a.spacing == b.spacing && extents(a) == extents(b)
end
function same_geometry(a::AbstractStaticVTKData, b::AbstractStaticVTKData)
    return same_geometry(promote(a,b)...)
end
function same_ordered_geometry(a::AbstractVTKUnstructuredData, b::AbstractVTKUnstructuredData)
    dim(a) != dim(b) && return false
    num_of_points(a) != num_of_points(b) && return false
    num_of_cells(a) != num_of_cells(b) && return false

    _num_of_points = num_of_points(a)
    @views for i in 1:_num_of_points
        b.point_coords[:,i] == a.point_coords[:,i] || return false
    end

    _num_of_cells = num_of_cells(a)
    for i in 1:_num_of_cells
        similar_cells(a.cell_connectivity[i], cell_type(a,i), b.cell_connectivity[i], cell_type(b,i)) || return false
    end
    return true
end
function same_ordered_geometry(a::AbstractVTKMultiblockData, b::AbstractVTKMultiblockData)
    length(a) == length(b) || return false
    for i in 1:length(a)
        typeof(a[i]) == typeof(b[i]) && (typeof(a[i]) <: AbstractVTKMultiblockData || typeof(a[i]) <: AbstractVTKUnstructuredData ? same_ordered_geometry(a[i], b[i]) : a[i] == b[i]) || return false
    end
    return true
end

function same_geometry_shape(a::VTKStructuredData, b::VTKStructuredData)
    return size(a.point_coords) == size(b.point_coords)
end
function same_geometry_shape(a::VTKRectilinearData, b::VTKRectilinearData)
    return extents(a) == extents(b)
end
function same_geometry_shape(a::VTKUniformRectilinearData, b::VTKUniformRectilinearData)
    return extents(a) == extents(b)
end
function same_geometry_shape(a::AbstractVTKUnstructuredData, b::AbstractVTKUnstructuredData)
    dim(a) != dim(b) && return false
    num_of_points(a) != num_of_points(b) && return false
    num_of_cells(a) != num_of_cells(b) && return false

    _num_of_cells = num_of_cells(a)
    for i in 1:_num_of_cells
        similar_cells(a.cell_connectivity[i], cell_type(a,i), b.cell_connectivity[i], cell_type(b,i)) || return false
    end
    return true
end
same_geometry_shape(a::AbstractStaticVTKData, b::AbstractStaticVTKData) = false

function same_ordered_geometry_shape(a::AbstractVTKMultiblockData, b::AbstractVTKMultiblockData)
    length(a) == length(b) || return false
    for i in 1:length(a)
        typeof(a[i]) == typeof(b[i]) || return false
        if typeof(a[i]) <: AbstractVTKMultiblockData
            same_ordered_geometry_shape(a[i], b[i]) || return false
        else
            same_geometry_shape(a[i], b[i]) || return false
        end
    end
    return true
end

function same_data_shape(a::T, b::T) where {T<:AbstractVTKSimpleData}
    a_keys = keys(a.point_data)
    b_keys = keys(b.point_data)
    for i in a_keys
        in(i, b_keys) && size(a.point_data[i]) == size(b.point_data[i]) || return false
    end
    a_keys = keys(a.cell_data)
    b_keys = keys(b.cell_data)
    for i in a_keys
        in(i, b_keys) && size(a.cell_data[i]) == size(b.cell_data[i]) || return false
    end
    return true
end
same_data_shape(a::AbstractVTKStructuredData, b::AbstractVTKUnstructuredData) = false
same_data_shape(a::AbstractVTKMultiblockData, b::AbstractVTKSimpleData) = false

function same_data_shape(a::T, b::T) where {T<:AbstractVTKMultiblockData}
    length(a) == length(b)
    for i in 1:length(a)
        same_data_shape(a[i], b[i]) || return false
    end
    return true
end


text:
using PyPlot
using Gen

import Random
Random.seed!(1)
include("poisson_process.jl")
include("load_data.jl")

# Example from Section 4 of Reversible jump Markov chain Monte Carlo
# computation and Bayesian model determination

#########
# model #
#########

const K = :k
const EVENTS = :events
const UNSORTED_CHANGEPT = :unsorted_changept
const UNSORTED_RATE = :unsorted_rate

function sort_change_pts(unsorted_change_pts::Vector)
    k = length(unsorted_change_pts)
    sorted_to_unsorted = sortperm(unsorted_change_pts)
    unsorted_to_sorted = Vector{Int}(undef, k)
    for i in 1:k
        unsorted_to_sorted[sorted_to_unsorted[i]] = i
    end
    sorted_change_pts = unsorted_change_pts[sorted_to_unsorted]
    return (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted)
end

function test_sort()
    unsorted_change_pts = [0.3, 0.1, 0.2, 0.0, -1.0]
    (change_pts, unsorted_to_sorted, sorted_to_unsorted) = sort_change_pts(unsorted_change_pts)
    @assert change_pts == [-1.0, 0.0, 0.1, 0.2, 0.3]
    @assert sorted_to_unsorted == [5, 4, 2, 3, 1]
    @assert unsorted_to_sorted == [5, 3, 4, 2, 1]
end

test_sort()

@gen function model(T::Float64)

    # prior on number of change points
    k = @trace(poisson(3.), K)

    # prior distribution on change points
    (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted) = sort_change_pts(
        Float64[({(UNSORTED_CHANGEPT, i)} ~ uniform(0, T)) for i in 1:k])

    # k + 1 rate values
    # unsorted_rates[i] for i > 0 is the rate of the segment immediately before
    # the i'th unsorted changepoint. unsorted_rates[k+1] is the rate of the
    # last segment
    alpha = 1.0; beta = 200.0
    unsorted_rates = Float64[@trace(Gen.gamma(alpha, 1. / beta), (UNSORTED_RATE, i)) for i=1:k+1]
    sorted_rates = Vector{Float64}(undef, k+1)
    sorted_rates[1:k] = unsorted_rates[sorted_to_unsorted]
    sorted_rates[k+1] = unsorted_rates[k+1]

    # poisson process
    bounds = vcat([0.], sorted_change_pts, [T])
    @trace(piecewise_poisson_process(bounds, sorted_rates), EVENTS)

    return (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted, sorted_rates)
end

get_t_min(trace) = 0.0
get_t_max(trace) = get_args(trace)[1]
get_num_change_pts(trace) = trace[:k]
get_sorted_change_pts(trace) = get_retval(trace)[1]
get_sorted_rates(trace) = get_retval(trace)[4]

function render(trace; ymax=0.02)
    k = get_num_change_pts(trace)
    bounds = vcat([get_t_min(trace)], get_sorted_change_pts(trace), [get_t_max(trace)])
    rates = get_sorted_rates(trace)
    for i=1:length(rates)
        lower = bounds[i]
        upper = bounds[i+1]
        rate = rates[i]
        plot([lower, upper], [rate, rate], color="black", linewidth=2)
    end
    points = trace[EVENTS]
    scatter(points, -rand(length(points)) * (ymax/5.), color="black", s=5)
    ax = gca()
    xlim = [get_t_min(trace), get_t_max(trace)]
    plot(xlim, [0., 0.], "--")
    ax[:set_xlim](xlim)
    ax[:set_ylim](-ymax/5., ymax)
end

function show_prior_samples()
    figure(figsize=(16,16))
    T = 40000.
    for i=1:16
        println("simulating $i")
        subplot(4, 4, i)
        (trace, ) = generate(model, (T,), EmptyChoiceMap())
        render(trace; ymax=0.015)
    end
    tight_layout(pad=0)
    savefig("prior_samples.pdf")
end

#############
# rate move #
#############

@gen function rate_proposal(trace)

    # pick a random segment whose rate to change
    i = @trace(uniform_discrete(1, trace[K]+1), :i)

    # propose new value for the rate
    #cur_rate = trace[(UNSORTED_RATE, i)]
    #I@trace(uniform_continuous(cur_rate/2., cur_rate*2.), :new_rate)
    new_rate_scaled ~ uniform(0, 1)

    nothing
end

# it is an involution because it:
# - maintains i = fwd_choices[:i] constant
# - swaps choices[(RATE, i)] with fwd_choices[:new_rate]

@transform rate_involution (model_in, aux_in) to (model_out, aux_out) begin
    i = @read(aux_in[:i], :discrete)
    @write(aux_out[:i], i, :discrete)
    new_rate_scaled = @read(aux_in[:new_rate_scaled], :continuous)
    cur_rate = @read(model_in[(UNSORTED_RATE, i)], :continuous)
    lower_bound = cur_rate / 2.0
    upper_bound = cur_rate * 2.0
    new_rate = lower_bound + new_rate_scaled * (upper_bound - lower_bound)
    @write(model_out[(UNSORTED_RATE, i)], new_rate, :continuous)
    prev_rate_scaled = (cur_rate - (new_rate / 2.0)) / (new_rate * 2.0 - new_rate / 2.0)
    @write(aux_out[:new_rate_scaled], prev_rate_scaled, :continuous)
end

is_involution!(rate_involution)

rate_move(trace) = metropolis_hastings(trace, rate_proposal, (), rate_involution, check=false)

#################
# position move #
#################

function get_sorted_idx(trace, unsorted_idx::Int)
    (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted, _) = get_retval(trace)
    return unsorted_to_sorted[unsorted_idx]
end

function get_lower_bound(trace, unsorted_idx::Int)
    if get_sorted_idx(trace, unsorted_idx) == 1
        return get_t_min(trace)
    else
        return sorted_change_pts[sorted_idx - 1]
    end
end

function get_upper_bound(trace, unsorted_idx::Int)
    if get_sorted_idx(trace, unsorted_idx) == get_num_change_pts(trace)
        return get_t_max(trace)
    else
        return sorted_change_pts[sorted_idx + 1]
    end
end

# TODO, note: the density is not nonzero everywhere with respect to the reference measure...
# but it could easily be made so by sampling from uniform(0, 1) the location between lower and upper bound...

@gen function position_proposal(trace)
    k = trace[K]
    @assert k > 0

    # pick a random unsorted changepoint to change
    i = @trace(uniform_discrete(1, k), :i)

    sorted_idx = get_sorted_idx(trace, i)
    sorted_to_unsorted = get_retval(trace)[3]
    #lower = get_lower_bound(trace, i)
    #upper = get_upper_bound(trace, i)
    @trace(uniform_continuous(0, 1), :new_changept_scaled)
    #@trace(uniform_continuous(lower, upper), :new_changept)

    return (sorted_idx, sorted_to_unsorted)
end

# it is an involution because it:
# - maintains i = fwd_choices[:i] constant
# - swaps choices[(CHANGEPT, i)] with fwd_choices[:new_changept]

@transform position_involution (model_in, aux_in) to (model_out, aux_out) begin
    (sorted_idx, sorted_to_unsorted) = @read(aux_in[], :discrete)
    k = @read(model_in[:k], :discrete)
    i = @read(aux_in[:i], :discrete)
    @write(aux_out[:i], i, :discrete)
    new_changept_scaled = @read(aux_in[:new_changept_scaled], :continuous)
    lower_bound = sorted_idx == 1 ? 0.0 : @read(model_in[(UNSORTED_CHANGEPT, sorted_to_unsorted[sorted_idx-1])], :continuous)
    upper_bound = sorted_idx == k ? T : @read(model_in[(UNSORTED_CHANGEPT, sorted_to_unsorted[sorted_idx+1])], :continuous)
    new_changept = lower_bound + new_changept_scaled * (upper_bound - lower_bound)
    @write(model_out[(UNSORTED_CHANGEPT, i)], new_changept, :continuous)

    prev_changept = @read(model_in[(UNSORTED_CHANGEPT, i)], :continuous)
    prev_changept_scaled = (prev_changept - lower_bound) / (upper_bound - lower_bound)
    @write(aux_out[:new_changept_scaled], prev_changept_scaled, :continuous)
end

is_involution!(position_involution)

position_move(trace) = metropolis_hastings(trace, position_proposal, (), position_involution, check=false)

######################
# birth / death move #
######################

const CHOSEN = :chosen
const IS_BIRTH = :is_birth
const NEW_CHANGEPT = :new_changept
const U = :u

@gen function birth_death_proposal(trace)
    #T = get_t_max(trace)
    k = trace[K]
    # if k = 0, then always do a birth move
    # if k > 0, then randomly choose a birth or death move
    isbirth = (k == 0) ? true : @trace(bernoulli(0.5), IS_BIRTH)
    if isbirth
        {NEW_CHANGEPT} ~ uniform_continuous(get_t_min(trace), get_t_max(trace))
        {U} ~ uniform_continuous(0.0, 1.0)
    end
    (_, unsorted_to_sorted, sorted_to_unsorted, _) = get_retval(trace)
    return (unsorted_to_sorted, sorted_to_unsorted) # NOTE: based on testing predicates on real values
end

function new_rates(cur_rate, u, cur_cp, prev_cp, next_cp)
    d_prev = cur_cp - prev_cp
    d_next = next_cp - cur_cp
    @assert d_prev > 0
    @assert d_next > 0
    d_total = d_prev + d_next
    log_cur_rate = log(cur_rate)
    log_ratio = log(1 - u) - log(u)
    prev_rate = exp(log_cur_rate - (d_next / d_total) * log_ratio)
    next_rate = exp(log_cur_rate + (d_prev / d_total) * log_ratio)
    @assert prev_rate > 0.
    @assert next_rate > 0.
    (prev_rate, next_rate)
end

function new_rates_inverse(prev_rate, next_rate, cur_cp, prev_cp, next_cp)
    d_prev = cur_cp - prev_cp
    d_next = next_cp - cur_cp
    @assert d_prev > 0
    @assert d_next > 0
    d_total = d_prev + d_next
    log_prev_rate = log(prev_rate)
    log_next_rate = log(next_rate)
    cur_rate = exp((d_prev / d_total) * log_prev_rate + (d_next / d_total) * log_next_rate)
    u = prev_rate / (prev_rate + next_rate)
    @assert cur_rate > 0.
    (cur_rate, u)
end

# it is an involution because:
# - it switches back and forth between birth move and death move
# - it maintains fwd_choices[CHOSEN] constant (applying it twice will first insert a new
#   changepoint and then remove that same changepoint)
# - new_rates, curried on cp_new, cp_prev, and cp_next, is the inverse of new_rates_inverse.

@transform birth_death_involution (model_in, aux_in) to (model_out, aux_out) begin
    T = get_args(model_in)[1]
    (unsorted_to_sorted, sorted_to_unsorted) = @read(aux_in[], :discrete)

    # current number of changepoints
    k = @read(model_in[K], :discrete)

    # if k == 0, then we can only do a birth move
    isbirth = (k == 0) || @read(aux_in[IS_BIRTH], :discrete)

    # if we are a birth move, the inverse is a death move
    if k > 1 || isbirth
        @write(aux_out[IS_BIRTH], !isbirth, :discrete)
    end

    if isbirth
        @tcall(birth(k, unsorted_to_sorted, sorted_to_unsorted))
    else
        @tcall(death(k, unsorted_to_sorted, sorted_to_unsorted))
    end
end

@transform birth(k::Int, unsorted_to_sorted, sorted_to_unsorted) (model_in, aux_in) to (model_out, aux_out) begin
    k = @read(model_in[K], :discrete)
    @write(model_out[K], k+1, :discrete)
    cp_new = @read(aux_in[NEW_CHANGEPT], :continuous)

    # set new changepoint
    @copy(aux_in[NEW_CHANGEPT], model_out[(UNSORTED_CHANGEPT, k+1)])

    # re-sort the changepoints, with the new one added
    unsorted_change_pts = [@read(model_in[(UNSORTED_CHANGEPT, i)], :continuous) for i in 1:k]
    push!(unsorted_change_pts, cp_new)
    (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted) = sort_change_pts(unsorted_change_pts)
    sorted_idx = unsorted_to_sorted[k+1]

    # the unsorted index of the next rate choice
    prev_next_unsorted_rate_idx = (sorted_idx == k+1) ? k+1 : sorted_to_unsorted[sorted_idx+1]
    new_next_unsorted_rate_idx = (sorted_idx == k+1) ? k+2 : sorted_to_unsorted[sorted_idx+1]

    # the previous and next changepoints
    cp_prev = (sorted_idx == 1) ? 0.0 : sorted_change_pts[sorted_idx-1]
    cp_next = (sorted_idx == k+1) ? T : sorted_change_pts[sorted_idx+1]

    # compute new rates
    h_cur = @read(model_in[(UNSORTED_RATE, prev_next_unsorted_rate_idx)], :continuous)
    u = @read(aux_in[U], :continuous)
    (h_prev, h_next) = new_rates(h_cur, u, cp_new, cp_prev, cp_next)

    # set new rates
    @write(model_out[(UNSORTED_RATE, k+1)], h_prev, :continuous)
    @write(model_out[(UNSORTED_RATE, new_next_unsorted_rate_idx)], h_next, :continuous)

    # shift the final rate up, only if it is not where we inserted it...
    if sorted_idx != k+1
        @copy(model_in[(UNSORTED_RATE, k+1)], model_out[(UNSORTED_RATE, k+2)])
    end
end

# TODO rename the last rate to 'last_rate' so we don't need to shift it around...

@transform death(k::Int, unsorted_to_sorted, sorted_to_unsorted) (model_in, aux_in) to (model_out, aux_out) begin
    k = @read(model_in[K], :discrete)
    @write(model_out[K], k-1, :discrete)
    cp_deleted = @read(model_in[(UNSORTED_CHANGEPT, k)], :continuous) # NOTE: will not be copied...

    # set proposed changepoint
    @copy(model_in[(UNSORTED_CHANGEPT, k)], aux_out[NEW_CHANGEPT]) # but gets copied here..

    # re-sort the changepoints (we could use the result of the sorting from the model trace, if we wanted)
    unsorted_change_pts = [@read(model_in[(UNSORTED_CHANGEPT, i)], :continuous) for i in 1:k] # all get copied..
    (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted) = sort_change_pts(unsorted_change_pts)
    sorted_idx = unsorted_to_sorted[k]

    # the unsorted index of the next rate choice
    prev_next_unsorted_rate_idx = (sorted_idx == k) ? k+1 : sorted_to_unsorted[sorted_idx+1]
    new_next_unsorted_rate_idx = (sorted_idx == k) ? k : sorted_to_unsorted[sorted_idx+1]
    @assert prev_next_unsorted_rate_idx != k

    # the previous and next changepoints
    cp_prev = (sorted_idx == 1) ? 0.0 : sorted_change_pts[sorted_idx-1]
    cp_next = (sorted_idx == k) ? T : sorted_change_pts[sorted_idx+1]

    # compute cur rate and u
    h_prev = @read(model_in[(UNSORTED_RATE, k)], :continuous)
    h_next = @read(model_in[(UNSORTED_RATE, prev_next_unsorted_rate_idx)], :continuous)
    (h_cur, u) = new_rates_inverse(h_prev, h_next, cp_deleted, cp_prev, cp_next)
    @write(aux_out[U], u, :continuous)

    # set cur rate
    @write(model_out[(UNSORTED_RATE, new_next_unsorted_rate_idx)], h_cur, :continuous)

    # shift the final rate, only if we are not deleting the last one..
    if sorted_idx != k
        @copy(model_in[(UNSORTED_RATE, k+1)], model_out[(UNSORTED_RATE, k)])
    end
end

is_involution!(birth_death_involution)

@gen function permutation_proposal(trace)
    k = get_num_change_pts(trace)
    to_swap ~ uniform_discrete(1, k-1)
end

@transform permutation_involution (model_in, aux_in) to (model_out, aux_out) begin
    k = @read(model_in[:k], :discrete)
    to_swap_with_k = @read(aux_in[:to_swap], :discrete)
    @copy(aux_in[:to_swap], aux_out[:to_swap])
    @copy(model_in[(UNSORTED_RATE, k)], model_out[(UNSORTED_RATE, to_swap_with_k)])
    @copy(model_in[(UNSORTED_RATE, to_swap_with_k)], model_out[(UNSORTED_RATE, k)])
    @copy(model_in[(UNSORTED_CHANGEPT, k)], model_out[(UNSORTED_CHANGEPT, to_swap_with_k)])
    @copy(model_in[(UNSORTED_CHANGEPT, to_swap_with_k)], model_out[(UNSORTED_CHANGEPT, k)])
end

is_involution!(permutation_involution)

function birth_death_move(trace)
    if trace[:k] > 1
        trace, acc = metropolis_hastings(trace, permutation_proposal, (), permutation_involution, check=false)
        @assert acc
    end
    return metropolis_hastings(trace, birth_death_proposal, (), birth_death_involution, check=false)
end

function mcmc_step(trace)
    (trace, _) = rate_move(trace)
    if trace[K] > 0
        (trace, _) = position_move(trace)
    end
    (trace, _) = birth_death_move(trace)
    trace
end

function simple_mcmc_step(trace)
    (trace, _) = rate_move(trace)
    if trace[K] > 0
        (trace, _) = position_move(trace)
    end
    (trace, _) = metropolis_hastings(trace, k_selection)
    trace
end

function do_mcmc(T, num_steps::Int)
    (trace, _) = generate(model, (T,), observations)
    for iter=1:num_steps
        k = trace[K]
        if iter % 1000 == 0
            println("iter $iter of $num_steps, k: $k")
        end
        trace = mcmc_step(trace)
    end
    trace
end

const k_selection = select(K)

function do_simple_mcmc(T, num_steps::Int)
    (trace, _) = generate(model, (T,), observations)
    for iter=1:num_steps
        k = trace[K]
        if iter % 1000 == 0
            println("iter $iter of $num_steps, k: $k")
        end
        trace = simple_mcmc_step(trace)
    end
    trace
end

########################
# inference experiment #
########################

import Random
Random.seed!(1)

const points = load_data_set()
const T = maximum(points)
const observations = choicemap()
observations[EVENTS] = points

function show_posterior_samples()
    figure(figsize=(16,16))
    for i=1:16
        println("replicate $i")
        subplot(4, 4, i)
        trace = do_mcmc(T, 200)
        render(trace; ymax=0.015)
    end
    tight_layout(pad=0)
    savefig("posterior_samples.pdf")

    figure(figsize=(16,16))
    for i=1:16
        println("replicate $i")
        subplot(4, 4, i)
        trace = do_simple_mcmc(T, 200)
        render(trace; ymax=0.015)
    end
    tight_layout(pad=0)
    savefig("posterior_samples_simple.pdf")
end

function get_rate_vector(trace, test_points)
    k = trace[K]
    cps = get_sorted_change_pts(trace)
    hs = get_sorted_rates(trace)
    rate = Vector{Float64}()
    cur_h_idx = 1
    cur_h = hs[cur_h_idx]
    next_cp_idx = 1
    upper = (next_cp_idx == k + 1) ? T : cps[next_cp_idx]
    for x in test_points
        while x > upper
            next_cp_idx += 1
            upper = (next_cp_idx == k + 1) ? T : cps[next_cp_idx]
            cur_h_idx += 1
            cur_h = hs[cur_h_idx]
        end
        push!(rate, cur_h)
    end
    rate
end

# compute posterior mean rate curve

function plot_posterior_mean_rate()
    test_points = collect(1.0:10.0:T)
    rates = Vector{Vector{Float64}}()
    num_samples = 0
    num_steps = 8000
    for reps=1:20
        (trace, _) = generate(model, (T,), observations)
        for iter=1:num_steps
            if iter % 1000 == 0
                println("iter $iter of $num_steps, k: $(trace[K])")
            end
            trace = mcmc_step(trace)
            if iter > 4000
                num_samples += 1
                rate_vector = get_rate_vector(trace, test_points)
                @assert length(rate_vector) == length(test_points)
                push!(rates, rate_vector)
            end
        end
    end
    posterior_mean_rate = zeros(length(test_points))
    for rate in rates
        posterior_mean_rate += rate / Float64(num_samples)
    end
    ymax = 0.010
    figure()
    plot(test_points, posterior_mean_rate, color="black")
    scatter(points, -rand(length(points)) * (ymax/6.), color="black", s=5)
    ax = gca()
    xlim = [0., T]
    plot(xlim, [0., 0.], "--")
    ax[:set_xlim](xlim)
    ax[:set_ylim](-ymax/5., ymax)
    savefig("posterior_mean_rate.pdf")
end

function plot_trace_plot()
    figure(figsize=(8, 4))

    # reversible jump
    (trace, _) = generate(model, (T,), observations)
    rate1 = Float64[]
    num_clusters_vec = Int[]
    burn_in = 0
    for iter=1:burn_in + 1000
        trace = mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, trace[K])
        end
    end
    subplot(2, 1, 1)
    plot(num_clusters_vec, "b")

    # simple MCMC
    (trace, _) = generate(model, (T,), observations)
    rate1 = Float64[]
    num_clusters_vec = Int[]
    burn_in = 0
    for iter=1:burn_in + 1000
        trace = simple_mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, trace[K])
        end
    end
    subplot(2, 1, 2)
    plot(num_clusters_vec, "b")

    savefig("trace_plot.pdf")
end

println("showing prior samples...")
show_prior_samples()

println("showing posterior samples...")
show_posterior_samples()

println("estimating posterior mean rate...")
plot_posterior_mean_rate()

println("making trace plot...")
plot_trace_plot()


text:
module CorrelationsController
using DataFrames, StatsBase, Statistics, Query, Combinatorics
using ..Stats
using JSON

export call

"""
call(message)

Computes two column correlations of pearson and spearman from records.

# Arguments
- `message::Dict`: the arguments dict
  - `"records"`: the array of array of rows

# Return
the json string of result object

# Examples
```jldoctest
julia> CorrelationsController.call(Dict("records" => [[1.0, 2.0, 0.0], [2.0, 1.0, 2.0], [3.0, 0.0, 9.0]]))
"{\"spearman_correlations\":[-1.0,1.0,-1.0],\"correlations\":[-1.0,0.9522165814091075,-0.9522165814091075],\"combs\":[[1,2],[1,3],[2,3]]}"
```
"""
function call(message)
  _records = hcat(message["records"]...)

  combs = collect(combinations(1:size(_records, 1), 2))
  cors = []
  corsspearman = []
  for i = combs
    cors = vcat(cors, Statistics.cor(_records[i[1], :], _records[i[2], :]))
    corsspearman = vcat(corsspearman, StatsBase.corspearman(_records[i[1], :], _records[i[2], :]))
  end

  Dict(
        :combs => combs,
        :correlations => cors,
        :spearman_correlations => corsspearman
       ) |> JSON.json
end
end


text:
using FileIO
using GaussianMixtures
using Distributed
using DelimitedFiles
using RDatasets
using Test
using ScikitLearnBase
using Random

import Base.isapprox
isapprox(a::Tuple, b::Tuple) = all(Bool[isapprox(x,y) for (x,y) in zip(a,b)])

GM = GaussianMixtures  # alias

setmem(0.1)
include("data.jl")
include("bayes.jl")
include("train.jl")
include("scikitlearn.jl")


text:
# _less_than_param
concept_less_than_param(x; param) = x[1] ≤ param
error_less_than_param(x; param, dom_size = 0) = max(0.0, x[1] - param) / dom_size

"""
    less_than_param
Constraint ensuring that the value of `x` is less than a given parameter `param`.
"""
const less_than_param = Constraint(
    args_length = 1,
    concept = concept_less_than_param,
    error = make_error(:less_than_param),
    param = 1,
)


text:
using JuMP, FactCheck, AmplNLWriter

if !isdefined(:solver); solver = IpoptNLSolver(); end
# Note min and max not implemented in Couenne

## Solve test problem with simple min functions
 #
 #  max   min( x^2, x )
 #  s.t.  -0.5 <= x <= 0.5
 #
 #  The optimal objective value is 0.25.
 #      x = 0.5
##
context("example: maxmin") do
    m = Model(solver=solver)
    @variable(m, -0.5 <= x <= 0.5, start = 0.25)
    @NLobjective(m, Max, min(x^2, 0.3, x))
    @fact solve(m) --> :Optimal
    @fact getobjectivevalue(m) --> roughly(0.25, 1e-2)
    @fact getvalue(x) --> roughly(0.5, 1e-2)
end

## Solve test problem with simple max functions
 #
 #  min   max( x^2, x )
 #  s.t.  -0.5 <= x <= 0.5
 #
 #  The optimal objective value is 0.
 #      x = 0.
##
context("example: minmax") do
    m = Model(solver=solver)
    @variable(m, -1 <= x <= 1, start=-1)
    @NLobjective(m, Min, max(x^2, x, -1))
    @fact solve(m) --> :Optimal
    @fact getobjectivevalue(m) --> roughly(0, 1e-2)
    @fact getvalue(x) --> roughly(0, 1e-2)
end