text:
# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

module Exp

using LinearAlgebra
import ACE
using ACE: alloc_B, alloc_temp, evaluate!, order

_get_samples(basis, nargs, nsamples) =
   [  [ ACE.rand_vec(basis.J) for _=1:nargs ]
       for _=1:nsamples ]

function _asmlsq(basis, samples)
   A = zeros(ComplexF64, (length(samples), length(basis)))
   B = alloc_B(basis)
   tmp = alloc_temp(basis)
   Z = zeros(Int, length(samples[1]))
   for (irow, R) in enumerate(samples)
      evaluate!(B, tmp, basis, R, Z, 0)
      A[irow, :] .= B
   end
   return A
end

function determine_order(f, nargs, basis;
                         ntrain = 10_000,
                         ntest = 1_000)
   train = _get_samples(basis, nargs, ntrain)
   test = _get_samples(basis, nargs, ntest)

   # assemble the lsq system
   Y = f.(train)
   A = _asmlsq(basis, train)

   # now loop over all possible body-orders...
   rmse = zeros(order(basis))
   for N = 1:order(basis)
      # indices in the basis corresponding to order N
      IN = ACE.Utils.findall_basis_N(basis, N; verbose=false)
      # the relevant subset of the lsq system is
      AN = A[:, IN]
      # now solve the lsq system and record the relative RMSE
      θ = qr(AN) \ Y
      rmse[N] = norm(AN * θ - Y) / norm(Y)
   end
   return rmse
end

end


text:
module MathematicaParser

export parsemathematica, toexpr
export MxNode, Args, Fun, Inv, Neg, Pow, Prd, Sum

##
using ParserCombinator

################################################################################
##                                                                            ##
#                      Stucts and Utility Functions                            #
#                                                                              #
################################################################################

abstract type MxNode end

Base.:(==)(n1::MxNode, n2::MxNode) = n1.val == n2.val

struct Sum<:MxNode
    val
    function Sum(v::Vector)
        # Only create a Sum type if the input has 2+ inputs
        if length(v) == 1
            return v[1]
        else
            return new(v)
        end
    end
end

struct Prd<:MxNode
    val
    function Prd(v::Vector)
        # Only create a Prd type if the input has 2+ inputs
        if length(v) == 1
            return v[1]
        else
            return new(v)
        end
    end
end

# TODO: consider making a Div object with numerator and denominator
struct Inv<:MxNode val end
struct Neg<:MxNode val end
Neg(n::Number) = - one(n)

struct Pow<:MxNode
    base
    expn
end
function Pow(b, v::Vector)
    # Only create a Pow type if there is a nonempty second argument
    if length(v) == 0
        return b
    elseif length(v) == 1
        return Pow(b, v[1])
    else
        # If there is a "power tower" like 2^3^5
        return Pow(b, Pow(v[1], v[2:end]))
    end
end

Base.:(==)(p1::Pow, p2::Pow) = (p1.base == p2.base) && (p1.expn == p2.expn)

struct Args<:MxNode val end

struct Fun<:MxNode
    fun
    argsets::Vector{Args}
end

Base.:(==)(f1::Fun, f2::Fun) = (f1.fun == f2.fun) && (f1.argsets == f2.argsets)

# TODO: How to convert to other number types?
#   - Complex numbers & im
#   - Rationals
#   - Special numbers like pi
"""
    pnum(x)

Parses a string of a number into a number. If the string parses to an Int
without an InexactError, then it returns an Int type, otherwide returns
a Float type.
"""
function pnum(x)
    xfl = parse(Float64, x)
    try
        return convert(Int, xfl)  # catches inexact errors
    catch
        return xfl
    end
end

################################################################################
##                                                                            ##
#                                 Basic patterns                               #
#                                                                              #
################################################################################

# whitespace:
spc = "([\t ]+|[\r\n]+(#.*)?)"
wstar(x) = string(x, spc, "*")   # A pattern x with possible whitespace
wplus(x) = string(x, spc, "+")   # A pattern x with 1+ whitespaces
spc_star = ~Pattern(spc)[0:end]  # Matches and drops any number of spaces

# Bracketing
par_beg = E"(" + spc_star
par_end = E")" + spc_star
sqr_beg = E"[" + spc_star
sqr_end = E"]" + spc_star
crl_beg = E"{" + spc_star
crl_end = E"}" + spc_star

# Special characters
# Pi = p"Pi"
# parse_one("Pi", Pi)
# parse_one("Piedmont", Pi)

# A pattern for numbers
# evaluates the number with the function pnum
num = Pattern(wstar("-?(\\d*\\.?\\d+|\\d+\\.\\d*)([eE]-?\\d+)?")) > pnum

# A pattern for variable names
varpat = p"[A-Za-z][A-Za-z\d_]*(?![\w]*\[)" > Symbol
# varpat = p"[a-z][A-Za-z\d_]*(?![\w]*\[)" > Symbol
var = spc_star + varpat + spc_star

add_op = E"+" + spc_star
sub_op = E"-" + spc_star
mul_op = E"*" + spc_star
div_op = E"/" + spc_star

pow_op = E"^" + spc_star

################################################################################
##                                                                            ##
#                                 Expressions                                  #
#                                                                              #
# Patterns are nested in order of operator precedence
#   ⋅ Comparators/equals/etc
#   ⋅ Parseing numbers, parentheses, functions
#   ⋅ Negative numbers
#       ⋅  -1^3  →  Pow(-1, 3) or (-1)³
#       ⋅ 2-1^3  →  Sum(Any[2, Neg(Pow(1, 3))]) or 2 - (1³)
#   ⋅ Exponents
#   ⋅ Multiplication/division
#   ⋅ Addition/Subtraction
#                                                                              #
################################################################################

expr = Delayed();

# Pattern for arguments
#   ⋅ Starts with "[", includs one or more sub expression separated by a comma
#   ⋅ Creates an Args types from the input
#   ⋅ Ends with a "]"
args = sqr_beg + (expr + (P"," + spc_star + expr)[0:end] |> collect > Args) + sqr_end;

# A pattern for Functions
# Starts with a capital letter and is followed by 1 or more argument sets
fun = p"[A-Z][a-zA-Z\d]*" + args[1:end] |> f -> Fun(Symbol(f[1]), f[2:end]);

# A pattern for lists,
# Any number of expression separated by commas enclosed in curly braces
list = crl_beg + (expr + (P"," + spc_star + expr)[0:end] + crl_end |> collect);

# Anonymous Functions with derivative shorthand
# TODO: single letter functions including lower case
function deriv(f, n::Int, args)
    if n == 0
        Fun(Symbol(f), [args])
    else
        Fun(:Derivative, Args[Args(n), Args(Symbol(f)), args])
    end
end
anonfun = p"[a-zA-Z\d]+" + (e"'"[0:end] |> length) + args[1:end] > deriv

fun2 = fun | anonfun

# Pattern for a "value" which can be...
#   ⋅ a subexpression in parentheses,
#   ⋅ a function,
#   ⋅ a number, or
#   ⋅ a variable
# val = (par_beg + expr + par_end) | fun | num | var;
val = (par_beg + expr + par_end) | fun2 | num | var | list;

# TODO: perhaps have different behaviors for -(x+2) and -2
#   ⋅ -(x+2) could create a Neg type
#   ⋅ -2 could parse to a number
neg = Delayed();
neg.matcher = val | (E"-" + neg > Neg); # allow multiple (or no) negations (eg ---3)

pow = neg + (pow_op + neg)[0:end] |> x -> Pow(x[1], x[2:end]);

mul = mul_op + pow;
div = div_op + pow > Inv;
prd = pow + (mul | div)[0:end] |> Prd;

add = add_op + prd;
sub = sub_op + prd > Neg;

expr.matcher = prd + (add | sub)[0:end] |> Sum;

function parsemathematica(s::AbstractString)
    try
        parse_one(s, expr)
    catch
        @warn "Cannot parse expression: $s"
    end
end

##

toexpr(x::Number) = x
toexpr(x::Symbol) = x
toexpr(s::Symbol) = s
toexpr(a::Args) = toexpr.(a.val)
toexpr(av::Vector{Args}) = vcat([toexpr(a) for a in av]...) # Flattens arg sets

toexpr(f::Sum) = Expr(:call, :+, toexpr.(f.val)...)
toexpr(f::Prd) = Expr(:call, :*, toexpr.(f.val)...)
toexpr(f::Pow) = Expr(:call, :^, f.base, toexpr(f.expn))
toexpr(f::Inv) = Expr(:call, :/, 1, toexpr(f.val))
toexpr(f::Neg) = Expr(:call, :*, -1, toexpr(f.val))

function toexpr(f::Fun)
    # Apply the conversion between functions here
    Expr(:call, f.fun, toexpr(f.argsets)...)
end

##
# myshow(x::Number, depth=0) = x

# function myshow(x::MxNode, depth=0)
#     indent = "--"^depth
#     println(indent, typeof(x))
#     for f in fieldnames(typeof(x))
#         v = getfield(x,f)
#         # @info f, v
#         # print()
#         println("  "^(depth+1),myshow.(v, depth+1))
#     end
#     return nothing
# end

# foo = parsemathematica("1^2 + 3")[1]
# foo = parsemathematica("4*-6*(3 ^ 7 + 5)")[1]
# @show foo;
# myshow(foo);

end

ext: jl


text:
immutable RK4 <: TimeIntergral_onestep
end

immutable RK4_oneStage{ns}
  n :: Int64
  function RK4_oneStage()
    @assert eltype(ns) <: Int64
    @assert 0<ns
    @assert ns<=4
    new(ns)
  end
end
export RK4_oneStage

#= function TimeForward!( p :: ModelParams, t :: Float64, dt :: Float64) =#
#=   error("No function for abstract types") =#
#= end =#

#= function fTrend( data:: Tuple, p:: ModelParams) =#
#=   error("No function for abstract types") =#
#= end =#

function eval(tf :: RK4_oneStage{1}, p :: ModelParams, y :: Tuple, y_0 :: Tuple, dt :: Float64)
  k1 = fTrend(y, p)
  y1 = map((1:length(y)...)) do ii
    k1[ii]*dt/2 + y_0[ii]
  end
  (k1,y1,dt/2)
end
function eval(tf :: RK4_oneStage{2}, p :: ModelParams, y :: Tuple, y_0 :: Tuple, dt :: Float64)
  k1 = fTrend(y, p)
  y1 = map((1:length(y)...)) do ii
    k1[ii]*dt/2 + y_0[ii]
  end
  (k1,y1,0.0)
end
function eval(tf :: RK4_oneStage{3}, p :: ModelParams, y :: Tuple, y_0 :: Tuple, dt :: Float64)
  k1 = fTrend(y, p)
  y1 = map((1:length(y)...)) do ii
    k1[ii]*dt + y_0[ii]
  end
  (k1,y1,dt/2)
end
function eval(tf :: RK4_oneStage{4}, p :: ModelParams, y :: Tuple, y_0 :: Tuple, dt :: Float64)
  k1 = fTrend(y, p)
  y1 = y_0
  return (k1,y1,0.0)
end

function eval!(tf :: RK4, p :: ModelParams,
  data :: Tuple, trends :: Tuple,
  t :: Float64, dt :: Float64)

  (k1,y1,forwardDt1) = eval(RK4_oneStage{1}(), p , data, data, dt)
  TimeForward!(p, t, forwardDt1)
  t = t + forwardDt1

  (k2,y2,forwardDt2) = eval(RK4_oneStage{2}(), p , y1, data, dt)
  TimeForward!(p, t, forwardDt2)
  t = t + forwardDt2

  (k3,y3,forwardDt3) = eval(RK4_oneStage{3}(), p , y2, data, dt)
  TimeForward!(p, t, forwardDt3)
  t = t + forwardDt3

  (k4,y4,forwardDt4) = eval(RK4_oneStage{4}(), p , y3, data, dt)
  #= TimeForward!(p, t, forwardDt4) =#
  #= t = t + forwardDt4 =#

  yfinal = map((1:length(y1)...)) do ii
    data[ii] + (k1[ii]+2*k2[ii] + 2*k3[ii] +k4[ii])*dt/6
  end
  return (p, yfinal, t)
end
export RK4


text:
using PlayGround

ws=PlayGround.run(10000);
PlayGround.clean()


text:
using LsqFit
using Optim
using DataFrames
import StatsBase.fitted

include("dynamics.jl")

struct SIRX <: Dynamics
    α::Float64
    β::Float64
    κ::Float64
    κ0::Float64
end

# default constructor
SIRX() = SIRX(1.0,1.0,0.0,0.0)

# another constructor
function getParams(κ::Float64,κ0::Float64,
    R0Free::Float64,TInfected::Float64,d::Dynamics)
    β = 1.0/TInfected
    α = R0Free*β
    return SIRX(α,β,κ,κ0)
end

# number of states
function nstates(d::SIRX)
    return 4
end

# names of states
function stateNames(d::SIRX)
    return ["S" "I" "R" "X"]
end

function R0Eff(d::SIRX)
    return d.α/(d.β+d.κ+d.κ0)
end

# change in a day
"""
change(s::Vector{Float64},d::SIRX)

Return the change in the state of the population in a day

s = state vector (S,I,R,X)
d = SIRX dynamics parameters
"""
function change(s::Vector{Float64},d::SIRX)
    # this is done purely for readability of the formula
    st = NamedTuple{(:S,:I,:R,:X)}(s)
    S = -( d.α * st.I + d.κ0) * st.S
    I = d.α*st.I*st.S - d.β*st.I - d.κ0*st.I - d.κ*st.I
    R = d.β*st.I + d.κ0*st.S
    X = (d.κ+d.κ0) * st.I
    return [S,I,R,X]
end

# initializer for states
function initialize(N::Float64,C0::Float64,IXRatio::Float64,d::SIRX)
    state = zeros(nstates(d))
    state[4] = C0/N # X0, page 1, Supplement
    state[2] = IXRatio*(C0/N) # I0, page 1, supplement
    state[1] = 1.0-state[2]-state[4] # S0
    return state
end

"""
caseModel(nt::Int64,N::Float64,C0::Float64,IXRatio::Float64,d::SIRX)

Returns vector of cases given model dynamics and population parameters.
"""
function caseModel(nt::Int64,N::Float64,C0::Float64,IXRatio::Float64,
    d::SIRX)
    s0 = initialize(N,C0,IXRatio,d)
    (d,ds) =evolve(N,s0,d,nt)
    return N.*d[4,:]
end

function estimatedStates(nt::Int64,N::Float64,C0::Float64,
    IXRatio::Float64,d::SIRX)
    s0 = initialize(N,C0,IXRatio,d)
    (d,ds) =evolve(N,s0,d,nt)
    s1 = DataFrame(N.*d',[:S,:I,:R,:X])
    return s1
end

function logCaseModelUnknown(nt::Int64,p::Vector{Float64},
    N::Float64,C0::Float64,R0Free::Float64,TInfected::Float64)
    d = getParams(p[1],p[2],R0Free,TInfected,SIRX())
    return log.(caseModel(nt,N,C0,p[3],d))
end

function sqrtCaseModelUnknown(nt::Int64,p::Vector{Float64},
    N::Float64,C0::Float64,R0Free::Float64,TInfected::Float64)
    d = getParams(p[1],p[2],R0Free,TInfected,SIRX())
    return sqrt.(caseModel(nt,N,C0,p[3],d))
end


text:
abstract CovarianceKernel

function gpcov(kernel::CovarianceKernel,
               v1::Vector{Float64},
               v2::Vector{Float64})

    n1 = size(v1, 1)
    n2 = size(v2, 1)
    cv = zeros(n1, n2)

    for i in 1:n1
        for j in 1:n2
            cv[i, j] = gpcov(kernel, v1[i], v2[j])
        end
    end

    cv
end

type SquaredExponential <: CovarianceKernel
    bandwidth::Float64
end

function gpcov(kernel::SquaredExponential, x1::Float64, x2::Float64)
    bw = kernel.bandwidth
    r = x1 - x2
    exp(- r ^ 2 / (2 * bw ^ 2))
end


text:
#=
    Contents: Functions for cleaning up connection pools.
=#

"""
Sets the connection c to unoccupied if target_pool_size is not exceeded, otherwise removes it from the pool.

Notes:
1. The call to pop! is left in both clauses of the if statement for readability.
   A bug would arise otherwise because the released connection would not be counted in the call to get_n_connections.
"""
function release!(cp::ConnectionPool, c)
    if get_n_connections(cp) <= get_target(cp)
	pop!(cp.occupied, c)
	push!(cp.unoccupied, c)
    else
	disconnect(c)
	pop!(cp.occupied, c)
    end
end

"""
Disconnects all connections and removes them from the pool and sets target_pool_size and peak_size to 0.

Notes:
1. Requires get_n_occupied(cp) == 0.
"""
function delete!(cp::ConnectionPool)
    assert(get_n_occupied(cp) == 0)
    while get_n_unoccupied(cp) > 0
	c = pop!(cp.unoccupied)
	disconnect(c)
    end
    set_target!(cp, 0)
    set_peak!(cp, 0)
end

# EOF


text:
## ONE FACTOR MODELS ##
mutable struct OneFactorShortRateTree{S <: ShortRateDynamics, T <: TrinomialTree} <: ShortRateTree
  tree::T
  dynamics::S
  tg::TimeGrid
  treeLattice::TreeLattice1D{OneFactorShortRateTree{S, T}}

  function OneFactorShortRateTree{S, T}(tree::T, dynamics::S, tg::TimeGrid) where {S, T}
    oneFactorTree = new{S, T}(tree, dynamics, tg)
    oneFactorTree.treeLattice = TreeLattice1D(tg, get_size(tree, 2), oneFactorTree)

    return oneFactorTree
  end
end

function rebuild_lattice!(lattice::OneFactorShortRateTree, tg::TimeGrid)
  println(lattice.tg)
  rebuild_tree!(lattice.tree, tg)
  lattice.tg = tg

  # update treeLattice
  lattice.treeLattice = TreeLattice1D(tg, get_size(lattice.tree, 2), lattice)
  return lattice
end

get_size(tr::OneFactorShortRateTree, i::Int) = get_size(tr.tree, i)

function discount(tr::OneFactorShortRateTree, i::Int, idx::Int)
  x = get_underlying(tr.tree, i, idx)
  # r = 1.0
  # try
  r = short_rate(tr.dynamics, tr.tg.times[i], x)
  # catch e
  #   println(e)
  #   println(i)
  #   println(tr.dynamics.fitting.values)
  #   error("DIE")
  # end
  return exp(-r * tr.tg.dt[i])
end

descendant(tr::OneFactorShortRateTree, i::Int, idx::Int, branch::Int) = descendant(tr.tree, i, idx, branch)
probability(tr::OneFactorShortRateTree, i::Int, idx::Int, branch::Int) = probability(tr.tree, i, idx, branch)

get_params(m::OneFactorModel) = Float64[get_a(m), get_sigma(m)]

struct RStarFinder{M <: ShortRateModel} <: Function
  model::M
  strike::Float64
  maturity::Float64
  valueTime::Float64
  fixedPayTimes::Vector{Float64}
  amounts::Vector{Float64}
end

function (rsf::RStarFinder)(x::Float64)
  _value = rsf.strike
  _B = discount_bond(rsf.model, rsf.maturity, rsf.valueTime, x)
  sz = length(rsf.fixedPayTimes)
  @simd for i = 1:sz
    @inbounds dbVal = discount_bond(rsf.model, rsf.maturity, rsf.fixedPayTimes[i], x) / _B
    @inbounds _value -= rsf.amounts[i] * dbVal
  end

  return _value
end

# function operator(rsf::RStarFinder)
#   function _inner(x::Float64)
#     _value = rsf.strike
#     _B = discount_bond(rsf.model, rsf.maturity, rsf.valueTime, x)
#     sz = length(rsf.fixedPayTimes)
#     for i = 1:sz
#       dbVal = discount_bond(rsf.model, rsf.maturity, rsf.fixedPayTimes[i], x) / _B
#       _value -= rsf.amounts[i] * dbVal
#     end
#
#     return _value
#   end
#
#   return _inner
# end

discount_bond(model::OneFactorModel, tNow::Float64, maturity::Float64, factors::Vector{Float64}) = discount_bond(model, tNow, maturity, factors[1])
discount_bond(model::OneFactorModel, tNow::Float64, maturity::Float64, _rate::Float64) = A(model, tNow, maturity) * exp(-B(model, tNow, maturity) * _rate)


text:
using RiskAdjustedLinearizations, SparseArrays, SparseDiffTools, Test
include(joinpath(dirname(@__FILE__), "..", "..", "examples", "rbc_cc", "rbc_cc.jl"))
include(joinpath(dirname(@__FILE__), "..", "..", "examples", "crw", "crw.jl"))

# Set up
n_strips = 0
m_rbc_cc = RBCCampbellCochraneHabits()
m_crw = CoeurdacierReyWinant()

# Test sparse Jacobians on RBC-CC
m_dense = rbc_cc(m_rbc_cc, n_strips)
z0 = copy(m_dense.z)
y0 = copy(m_dense.y)
Ψ0 = copy(m_dense.Ψ)

m = rbc_cc(m_rbc_cc, n_strips; sparse_jacobian = [:μ, :ξ, :𝒱])
@testset "Construct a RiskAdjustedLinearization that exploits sparsity in Jacobians (using RBC-CC)" begin
    @test isempty(m_dense.linearization.sparse_jac_caches)
    @test m.z ≈ z0
    @test m.y ≈ y0
    @test m.Ψ ≈ Ψ0
    for k in [:μz, :μy, :ξz, :ξy, :J𝒱]
        @test haskey(m.linearization.sparse_jac_caches, k)
        if k != :J𝒱
            @test issparse(m.linearization.sparse_jac_caches[k][:sparsity])
            @test isa(m.linearization.sparse_jac_caches[k][:colorvec], AbstractVector{Int})
        end
    end
    @test m.linearization.sparse_jac_caches[:J𝒱][:colorvec] == 1:2
    @test m.linearization.sparse_jac_caches[:J𝒱][:sparsity] == ones(size(m.Ψ))
end

@testset "Update a RiskAdjustedLinearization to exploit sparsity in Jacobians (using RBC-CC)" begin
    update_sparsity_pattern!(m_dense, :𝒱)
    for k in [:μz, :μy, :ξz, :ξy]
        @test !haskey(m_dense.linearization.sparse_jac_caches, k)
    end
    @test m_dense.linearization.sparse_jac_caches[:J𝒱][:colorvec] == 1:2
    @test m_dense.linearization.sparse_jac_caches[:J𝒱][:sparsity] == ones(size(m.Ψ))

    update_sparsity_pattern!(m_dense, [:μ, :ξ, :𝒱])
    for k in [:μz, :μy, :ξz, :ξy]
        @test haskey(m_dense.linearization.sparse_jac_caches, k)
        @test issparse(m_dense.linearization.sparse_jac_caches[k][:sparsity])
        @test isa(m_dense.linearization.sparse_jac_caches[k][:colorvec], AbstractVector{Int})
    end

    @test_broken update_sparsity_pattern!(m_dense, [:μ, :ξ, :𝒱]; sparsity_detection = true)
end

@testset "Calculate risk-adjusted linearization with sparse autodiff (using RBC-CC)" begin
    # Now provide the sparsity pattern and matrix coloring vector
    # to update the Jacobians of objects
    rbc_cc_out = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "reference", "rbccc_sss_iterative_output.jld2"), "r")
    m_dense = rbc_cc(m_rbc_cc, n_strips) # recompute to get dense Jacobians again
    update!(m_dense, vec(rbc_cc_out["z_rss"]), vec(rbc_cc_out["y_rss"]), rbc_cc_out["Psi_rss"])
    ztrue = copy(m_dense.z)
    ytrue = copy(m_dense.y)
    Ψtrue = copy(m_dense.Ψ)

    sparsity = Dict{Symbol, SparseMatrixCSC{Float64, Int64}}()
    colorvec = Dict{Symbol, Vector{Int64}}()
    sparsity[:μz] = sparse(m_dense[:Γ₁])
    sparsity[:μy] = sparse(m_dense[:Γ₂])
    sparsity[:ξz] = sparse(m_dense[:Γ₃])
    sparsity[:ξy] = sparse(m_dense[:Γ₄])
    sparsity[:J𝒱] = sparse(m_dense[:JV])
    for (k, v) in sparsity
        colorvec[k] = matrix_colors(v)
    end

    # Check updating dense Jacobians works
    update_sparsity_pattern!(m_dense, [:μ, :ξ, :𝒱])
    try # prone to weird non-deterministic behavior in nlsolve
        solve!(m_dense, ztrue * 1.005, ytrue * 1.005, Ψtrue * 1.005; algorithm = :relaxation,
               ftol = 1e-6, tol = 1e-5, verbose = :none)
        @test norm(steady_state_errors(m_dense), Inf) < 1e-4
    catch e
        println("Updating dense Jacobian with sparse Jacobian methods did not pass")
    end
    @test_broken update_sparsity_pattern!(m_dense, [:μ, :ξ, :𝒱]; sparsity_detection = true)

    # Check updating sparse Jacobians w/new patterns works
    update_sparsity_pattern!(m, :𝒱; sparsity = sparsity,
                             colorvec = colorvec)
    try # prone to weird non-deterministic behavior in nlsolve
        solve!(m, ztrue * 1.005, ytrue * 1.005, Ψtrue * 1.005; algorithm = :relaxation,
               ftol = 1e-6, tol = 1e-5, verbose = :none)
        @test norm(steady_state_errors(m), Inf) < 1e-4
    catch e
        println("Updating sparsity pattern of 𝒱 for an RAL w/sparse methods did not pass")
    end

    update_sparsity_pattern!(m, [:μ, :ξ, :𝒱]; sparsity = sparsity)
    update_sparsity_pattern!(m, [:μ, :ξ, :𝒱]; sparsity = sparsity,
                             colorvec = colorvec)
    try # prone to weird non-deterministic behavior in nlsolve
        solve!(m, ztrue * 1.005, ytrue * 1.005, Ψtrue * 1.005; algorithm = :relaxation,
               ftol = 1e-6, tol = 1e-5, verbose = :none)
        @test norm(steady_state_errors(m), Inf) < 1e-4
    catch e
        println("Updating sparsity pattern of μ, ξ, and 𝒱 for an RAL w/sparse methods did not pass")
    end

    close(rbc_cc_out)
    # caching appears to be failing somehow; the caches of μ, ξ, and 𝒱 are being set to NaN unexpectedly
    @test_broken solve!(m, ztrue * 1.005, ytrue * 1.005, Ψtrue * 1.005; algorithm = :homotopy, verbose = :none)
    #=
    @test m.z ≈ m_dense.z atol=1e-6
    @test m.y ≈ m_dense.y atol=1e-6
    @test m.Ψ ≈ m_dense.Ψ atol=1e-6
    =#
end

# Test sparse Jacobians on CRW
m_dense = crw(m_crw)
z0 = copy(m_dense.z)
y0 = copy(m_dense.y)
Ψ0 = copy(m_dense.Ψ)

m = crw(m_crw; Ψ = zero(Ψ0), sparse_jacobian = [:μ, :ξ, :𝒱])
m_dense.Ψ .= 0.
@testset "Construct a RiskAdjustedLinearization that exploits sparsity in Jacobians (using CRW)" begin
    @test isempty(m_dense.linearization.sparse_jac_caches)
    @test m.z ≈ z0
    @test m.y ≈ y0
    for k in [:μz, :μy, :ξz, :ξy, :J𝒱]
        @test haskey(m.linearization.sparse_jac_caches, k)
        if k != :J𝒱
            @test issparse(m.linearization.sparse_jac_caches[k][:sparsity])
            @test isa(m.linearization.sparse_jac_caches[k][:colorvec], AbstractVector{Int})
        end
    end
    @test m.linearization.sparse_jac_caches[:J𝒱][:colorvec] == 1:3
    @test m.linearization.sparse_jac_caches[:J𝒱][:sparsity] == ones(size(m.Ψ))
end

@testset "Update a RiskAdjustedLinearization to exploit sparsity in Jacobians (using CRW)" begin
    update_sparsity_pattern!(m_dense, :𝒱)
    for k in [:μz, :μy, :ξz, :ξy]
        @test !haskey(m_dense.linearization.sparse_jac_caches, k)
    end
    @test m_dense.linearization.sparse_jac_caches[:J𝒱][:colorvec] == 1:3
    @test m_dense.linearization.sparse_jac_caches[:J𝒱][:sparsity] == ones(size(m.Ψ))

    update_sparsity_pattern!(m_dense, [:μ, :ξ, :𝒱])
    for k in [:μz, :μy]
        @test haskey(m_dense.linearization.sparse_jac_caches, k)
        @test issparse(m_dense.linearization.sparse_jac_caches[k][:sparsity])
        @test isa(m_dense.linearization.sparse_jac_caches[k][:colorvec], AbstractVector{Int})
    end
    @test_broken update_sparsity_pattern!(m_dense, [:μ, :ξ, :𝒱]; sparsity_detection = true)
end

@testset "Calculate risk-adjusted linearization with sparse autodiff (using CRW)" begin
    # Now provide the sparsity pattern and matrix coloring vector
    # to update the Jacobians of objects
    m_dense = crw(m_crw) # recompute to get dense Jacobians again
    crw_out = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "reference/crw_sss.jld2"), "r")
    update!(m_dense, vec(crw_out["z_rss"]), vec(crw_out["y_rss"]), copy(crw_out["Psi_rss"]))
    ztrue = copy(m_dense.z)
    ytrue = copy(m_dense.y)
    Ψtrue = copy(m_dense.Ψ)

    sparsity = Dict{Symbol, SparseMatrixCSC{Float64, Int64}}()
    colorvec = Dict{Symbol, Vector{Int64}}()
    sparsity[:μz] = sparse(m_dense[:Γ₁])
    sparsity[:μy] = sparse(m_dense[:Γ₂])
    sparsity[:ξz] = sparse(ones(size(m_dense[:Γ₃])))
    sparsity[:ξy] = sparse(m_dense[:Γ₄])
    sparsity[:J𝒱] = sparse(m_dense[:JV])
    for (k, v) in sparsity
        if k == :ξz
            colorvec[k] = 1:3
        elseif k ==  :J𝒱
            colorvec[k] = ones(Int64, 3)
        else
            colorvec[k] = matrix_colors(v)
        end
    end

    # Check updating dense Jacobians works
    update_sparsity_pattern!(m_dense, [:μ, :ξ])
    try
        solve!(m_dense, ztrue, ytrue, Ψtrue; algorithm = :relaxation, ftol = 5e-4, tol = 1e-3, verbose = :none)
        @test norm(steady_state_errors(m_dense), Inf) < 1e-3
    catch e
        println("Updating dense Jacobian with sparse Jacobian methods did not pass")
    end

    # Check updating sparse Jacobians w/new patterns works
    update_sparsity_pattern!(m, :𝒱; sparsity = sparsity,
                             colorvec = colorvec)
    try
        solve!(m, ztrue, ytrue, Ψtrue; algorithm = :relaxation, ftol = 5e-4, tol = 1e-3, verbose = :none)
        @test norm(steady_state_errors(m), Inf) < 1e-3
    catch e
        println("Updating sparsity pattern of 𝒱 for an RAL w/sparse methods did not pass")
    end

    update_sparsity_pattern!(m, [:μ, :ξ]; sparsity = sparsity) # don't do for 𝒱 b/c its Jacobian is all zeros
    update_sparsity_pattern!(m, [:μ, :ξ, :𝒱]; sparsity = sparsity,
                             colorvec = colorvec)
    try
        solve!(m, ztrue, ytrue, Ψtrue; algorithm = :relaxation, ftol = 5e-4, tol = 1e-3, verbose = :none)
        @test norm(steady_state_errors(m), Inf) < 1e-3
    catch e
        println("Updating sparsity pattern of μ, ξ, and 𝒱 for an RAL w/sparse methods did not pass")
    end
    @test_broken update_sparsity_pattern!(m_dense, [:μ, :ξ, :𝒱]; sparsity_detection = true)

    @test_broken solve!(m, ztrue, ytrue, Ψtrue; algorithm = :homotopy, ftol = 5e-4, tol = 1e-3, verbose = :none)
end


text:
# Module-related types and auxiliary functions

import Base: unsafe_convert

export
    CuModule, CuModuleFile, unload

include("module/jit.jl")

typealias CuModule_t Ptr{Void}

immutable CuModule
    handle::CuModule_t

    """
    Create a CUDA module from a string containing PTX code.

    If the Julia debug level is 2 or higher (or, on 0.5, if CUDAdrv is loaded in DEBUG
    mode), line number and debug information will be requested when loading the PTX code.
    """
    function CuModule(data)
        handle_ref = Ref{CuModule_t}()

        options = Dict{CUjit_option,Any}()
        options[ERROR_LOG_BUFFER] = Array(UInt8, 1024*1024)
        @static if (VERSION >= v"0.6.0-dev.779" && Base.JLOptions().debug_level >= 2) ||
                   DEBUG
            options[GENERATE_LINE_INFO] = true
            options[GENERATE_DEBUG_INFO] = true
        end
        @static if DEBUG
            options[INFO_LOG_BUFFER] = Array(UInt8, 1024*1024)
            options[LOG_VERBOSE] = true
        end
        optionKeys, optionVals = encode(options)

        try
            @apicall(:cuModuleLoadDataEx,
                     (Ptr{CuModule_t}, Ptr{Cchar}, Cuint, Ptr{CUjit_option}, Ptr{Ptr{Void}}),
                     handle_ref, data, length(optionKeys), optionKeys, optionVals)
        catch err
            (err == ERROR_NO_BINARY_FOR_GPU || err == ERROR_INVALID_IMAGE) || rethrow(err)
            options = decode(optionKeys, optionVals)
            rethrow(CuError(err.code, options[ERROR_LOG_BUFFER]))
        end

        @static if DEBUG
            options = decode(optionKeys, optionVals)
            if isempty(options[INFO_LOG_BUFFER])
                debug("JIT info log is empty")
            else
                debug("JIT info log: ", repr_indented(options[INFO_LOG_BUFFER]; abbrev=false))
            end
        end

        new(handle_ref[])
    end
end

unsafe_convert(::Type{CuModule_t}, mod::CuModule) = mod.handle

"""
Unload a CUDA module.
"""
function unload(mod::CuModule)
    @apicall(:cuModuleUnload, (CuModule_t,), mod)
end

"""
Create a CUDA module from a file containing PTX code.

Note that for improved error reporting, this does not rely on the corresponding CUDA driver
call, but opens and reads the file from within Julia instead.
"""
CuModuleFile(path) = CuModule(open(readstring, path))

# do syntax, f(module)
function CuModuleFile(f::Function, path::AbstractString)
    mod = CuModuleFile(path)
    local ret
    try
        ret = f(mod)
    finally
        unload(mod)
    end
    ret
end

include("module/linker.jl")
include("module/function.jl")
include("module/global.jl")


text:
module OpenVX
using CEnum

include("libopenvx/VX.jl")

using .VX
export VX
end # module


text:
println(9223372036854775807)
println(-9223372036854775808)
println(9223372036854775807 + 1)
println(-9223372036854775808 - 1)
println(2 ^ 63)
println(2 ^ 63 - 1)
println(- 2 ^ 63)
println(1729 ^ 42)

max_stars_repo_name: gabriel-doriath-dohler/pjulia

max_issues_repo_name: gabriel-doriath-dohler/pjulia

max_forks_repo_name: gabriel-doriath-dohler/pjulia


text:
using Test
using CGE

for e in ["-g","test.edgelist","-c","test1col.ecg","-e","test_n2v.embedding","-l","20","-f","1","-m","rss"]
    push!(ARGS,e)
end
edges, weights, vweights, comm, clusters, embed, verbose, land, forced, method = parseargs()

empty!(ARGS)
for e in ["-g","test.edgelist","-c","test2col.ecg","-e","test_ordered.embedding","-l","20","-f","1","-m","rss"]
    push!(ARGS,e)
end
edges, weights, vweights, comm, clusters, embed, verbose, land, forced, method = parseargs()

empty!(ARGS)
for e in ["-g","test_weights.edgelist","-c","test2col.ecg","-e","test_unordered.embedding","-l","20","-f","1","-m","rss"]
    push!(ARGS,e)
end
edges, weights, vweights, comm, clusters, embed, verbose, land, forced, method = parseargs()

@testset "parsing" begin

@test typeof(edges) == Array{Int,2}
@test minimum(edges) == 1

@test typeof(weights) == Array{Float64,1}
@test typeof(vweights) == Array{Float64,1}

@test typeof(comm) == Array{Int,2}
@test size(comm,2) == 1
@test minimum(comm) == 1

@test typeof(clusters) == Vector{Vector{Int}}
@test typeof(embed) == Array{Float64,2}

@test typeof(verbose) == Bool

@test typeof(land) == Int
@test typeof(forced) == Int

end

distances, lembed, lcomm, ledges, lweights  = landmarks(edges, weights, vweights, clusters, comm,
embed, verbose, land, forced, method, false);

@testset "landmarks-rss" begin

    @test typeof(ledges) == Array{Int,2}
    @test minimum(ledges) == 1
    @test typeof(lweights) == Array{Float64,1}
    @test typeof(lcomm) == Array{Int,2}
    @test size(lcomm,2) == 1
    @test typeof(distances) == Array{Float64,1}
end

distances, lembed, lcomm, ledges, lweights  = landmarks(edges, weights, vweights, clusters, comm,
embed, verbose, land, forced, CGE.split_cluster_rss2,false);

@testset "landmarks-rss2" begin

@test typeof(ledges) == Array{Int,2}
@test minimum(ledges) == 1
@test typeof(lweights) == Array{Float64,1}
@test typeof(lcomm) == Array{Int,2}
@test size(lcomm,2) == 1
@test typeof(distances) == Array{Float64,1}
end

distances, lembed, lcomm, ledges, lweights  = landmarks(edges, weights, vweights, clusters, comm,
embed, verbose, land, forced, CGE.split_cluster_size,false);

@testset "landmarks-size" begin

@test typeof(ledges) == Array{Int,2}
@test minimum(ledges) == 1
@test typeof(lweights) == Array{Float64,1}
@test typeof(lcomm) == Array{Int,2}
@test size(lcomm,2) == 1
@test typeof(distances) == Array{Float64,1}
end

distances, lembed, lcomm, ledges, lweights, lvweights, v_to_l  = landmarks(edges, weights, vweights, clusters, comm,
embed, verbose, land, forced, CGE.split_cluster_diameter,false);

@testset "landmarks-diameter" begin

@test typeof(ledges) == Array{Int,2}
@test minimum(ledges) == 1
@test typeof(lweights) == Array{Float64,1}
@test typeof(lcomm) == Array{Int,2}
@test size(lcomm,2) == 1
@test typeof(distances) == Array{Float64,1}
end

results = wGCL(ledges, lweights, lcomm, lembed, distances, lvweights,
    Float64[], Int64[], Matrix{Int}(undef,0,0), Float64[],Matrix{Float64}(undef,0,0),false, 42, 10000,verbose);

@testset "wgcl" begin

@test typeof(results) == Array{Float64,1}
@test results[1] <=10.0

end

@testset "clustering" begin

louvain_clust("test.edgelist")
@test isfile("test.edgelist.ecg")
isfile("test.edgelist.ecg") && rm("test.edgelist.ecg")

louvain_clust("testw.edgelist",edges, weights)
@test isfile("testw.edgelist.ecg")
isfile("testw.edgelist.ecg") && rm("testw.edgelist.ecg")
end


text:
mutable struct Source <: Data.Source
    schema::Data.Schema
    featurelayer::AbstractFeatureLayer
    feature::ArchGDAL.Feature
    ngeom::Int
end

function Source(layer::AbstractFeatureLayer)
    featuredefn = layerdefn(layer)
    ngeometries = ngeom(featuredefn)
    nfld = nfield(featuredefn)
    header = [
        ["geometry$(i-1)" for i in 1:ngeometries];
        [getname(getfielddefn(featuredefn,i-1)) for i in 1:nfld]
    ]
    types = [
        [IGeometry for i in 1:ngeometries];
        [_FIELDTYPE[gettype(getfielddefn(featuredefn,i-1))] for i in 1:nfld]
    ]
    ArchGDAL.Source(
        Data.Schema(types, header, nfeature(layer)),
        layer,
        unsafe_nextfeature(layer),
        ngeometries
    )
end
Data.schema(source::ArchGDAL.Source) = source.schema
Data.isdone(s::ArchGDAL.Source, row, col) = s.feature.ptr == C_NULL
Data.streamtype(::Type{ArchGDAL.Source}, ::Type{Data.Field}) = true
Data.accesspattern(source::ArchGDAL.Source) = Data.Sequential
Data.reset!(source::ArchGDAL.Source) = resetreading!(source.featurelayer)

function Data.streamfrom(
        source::ArchGDAL.Source,
        ::Type{Data.Field},
        ::Type{T},
        row,
        col
    ) where T
    val = if col <= source.ngeom
        @assert T <: IGeometry
        getgeom(source.feature, col-1)
    else
        T(getfield(source.feature, col - source.ngeom - 1))
    end
    if col == source.schema.cols
        destroy(source.feature)
        source.feature.ptr = GDAL.ogr_l_getnextfeature(source.featurelayer.ptr)
        if row == source.schema.rows
            @assert source.feature.ptr == C_NULL
            resetreading!(source.featurelayer)
        end
    end
    val
end


text:
using WhereTheWaterFlows

function upstream_catchments(dem::AbstractArray, pourpt::CartesianIndex{2}; snap_pourpt::Bool=true,search_radius=10)

    area, slen, dir, nout, nin, pits, c, bnds = waterflows(dem)

    xidx,yidx = Tuple(pourpt)

    if snap_pourpt
        window = area[yidx-search_radius:yidx+search_radius,xidx-search_radius:xidx+search_radius]
        val, outletindex = findmax(window)
        outletindex = (outletindex - CartesianIndex(search_radius+1,search_radius+1)) + CartesianIndex(yidx,xidx)
    else
        outletindex = CartesianIndex(yidx,xidx)
    end

    catchments = catchment(dir,outletindex)

    return catchments

end

function flow_sequence(fdir::Array{Int8,2}, start::CartesianIndex)

    c = falses(size(fdir))
    idxes= CartesianIndex{2}[]

    _flowseq!(c,idxes,fdir,start,1)

    return idxes
end

function _flowseq!(c,idxes, dir, ij, iter)
    push!(idxes,ij)
    c[ij] = true
    # proc upstream points
    for IJ in WhereTheWaterFlows.iterate_D9(ij, c)
        ij==IJ && continue
        if WhereTheWaterFlows.flowsinto(IJ, dir[IJ], ij)
            _flowseq!(c, idxes, dir, IJ, iter+1)
        end
    end
end


text:
export excess_solution_check, excess_solution_check!, OverdeterminedTracker
"""
    excess_solution_check!(path_result::PathResult,
                           F::RandomizedSystem,
                           newton_cache = NewtonCache(F.system))

Assigns to the [`PathResult`](@ref) `path_result` the `return_code` `:excess_solution` if
the `path_result` is a solution of the randomized system `F` but not of
the polynomial system underlying `F`.
This is performed by using Newton's method for non-singular solutions and comparing the
residuals of the solutions for singular solutions.
"""
function excess_solution_check!(
    path_result::PathResult,
    F::RandomizedSystem,
    newton_cache::NewtonCache = NewtonCache(F.system),
)
    is_success(path_result) || return path_result
    if is_nonsingular(path_result)
        acc = max(accuracy(path_result), eps())
        max_abs_norm_first_update = 100 * sqrt(acc)
        res = newton(
            F.system,
            solution(path_result),
            nothing,
            InfNorm(),
            newton_cache;
            extended_precision = true,
            atol = 1e3 * acc,
            rtol = 1e3 * acc,
            max_abs_norm_first_update = max_abs_norm_first_update,
        )
        if !is_success(res)
            path_result.return_code = :excess_solution
            path_result.accuracy = res.accuracy
            path_result.residual = res.residual
        end
    else
        # for singular solutions compare residuals due to lack of something better right now
        evaluate!(newton_cache.r, F.system, solution(path_result))
        system_residual = LA.norm(newton_cache.r, InfNorm())
        if system_residual > 100 * residual(path_result)
            path_result.return_code = :excess_solution
            path_result.residual = system_residual
        end
    end

    path_result
end

struct ExcessSolutionCheck{S<:RandomizedSystem,M} <: Function
    system::S
    newton_cache::NewtonCache{M}
end
(check::ExcessSolutionCheck)(R::PathResult) =
    excess_solution_check!(R, check.system, check.newton_cache)

"""
    excess_solution_check(F::RandomizedSystem)

Returns a function `λ(::PathResult)` which performs the excess solution check.
The call `excess_solution_check(F)(path_result)` is identical to
`excess_solution_check!(F, path_result)`. See also [`excess_solution_check!`](@ref).
"""
excess_solution_check(F::RandomizedSystem) = ExcessSolutionCheck(F, NewtonCache(F.system))

"""
    OverdeterminedTracker(tracker::AbstractPathTracker, F::RandomizedSystem)

Wraps the given [`AbstractPathTracker`](@ref) `tracker` to apply
[`excess_solution_check`](@ref) for the given randomized system `F` on each path result.
"""
struct OverdeterminedTracker{T<:AbstractPathTracker,E<:ExcessSolutionCheck} <:
       AbstractPathTracker
    tracker::T
    excess_solution_check::E
end
OverdeterminedTracker(T::AbstractPathTracker, F::RandomizedSystem) =
    OverdeterminedTracker(T, excess_solution_check(F))

track(T::OverdeterminedTracker, x; kwargs...) =
    T.excess_solution_check(track(T.tracker, x; kwargs...))


text:
abstract type AbstractDisplayProperties end

Base.@kwdef mutable struct FileDialogDisplayProperties <: AbstractDisplayProperties
    caption::String
    action::String
    position::Float64 = Float64(32)
    width::Float64 = Float64(320)
    heigh::Float64 = Float64(200)
end

dump(Base.@kwdef mutable struct FileDialogDisplayProperties <: AbstractDisplayProperties
    caption::String
    action::String
    position::Float64 = Float64(32)
    width::Float64 = Float64(320)
    heigh::Float64 = Float64(200)
end)

@macroexpand Base.@kwdef mutable struct FileDialogDisplayProperties <: AbstractDisplayProperties
    caption::String
    action::String
    position::Float64 = Float64(32)
    width::Float64 = Float64(320)
    heigh::Float64 = Float64(200)
end

@macroexpand @mytest Base.@kwdef mutable struct FileDialogDisplayProperties <: AbstractDisplayProperties
    caption::String
    action::String
    position::Float64 = Float64(32)
    width::Float64 = Float64(320)
    heigh::Float64 = Float64(200)
end

macro mytest(expr)
    expr = macroexpand(__module__, expr) # to expand @static
    kwdefs = dump(expr)

    # expr isa Expr && expr.head == :struct || error("Invalid usage of @kwdef")
    # T = expr.args[2]
    # if T isa Expr && T.head == :<:
    #     T = T.args[1]
    # end
    #
    # params_ex = Expr(:parameters)
    # call_args = Any[]
    #
    # kwdefs = nothing

    # _kwdef!(expr.args[3], params_ex.args, call_args)
    # # Only define a constructor if the type has fields, otherwise we'll get a stack
    # # overflow on construction
    # if !isempty(params_ex.args)
    #     if T isa Symbol
    #         kwdefs = :(($(esc(T)))($params_ex) = ($(esc(T)))($(call_args...)))
    #     elseif T isa Expr && T.head == :curly
    #         # if T == S{A<:AA,B<:BB}, define two methods
    #         #   S(...) = ...
    #         #   S{A,B}(...) where {A<:AA,B<:BB} = ...
    #         S = T.args[1]
    #         P = T.args[2:end]
    #         Q = [U isa Expr && U.head == :<: ? U.args[1] : U for U in P]
    #         SQ = :($S{$(Q...)})
    #         kwdefs = quote
    #             ($(esc(S)))($params_ex) =($(esc(S)))($(call_args...))
    #             ($(esc(SQ)))($params_ex) where {$(esc.(P)...)} =
    #                 ($(esc(SQ)))($(call_args...))
    #         end
    #     else
    #         error("Invalid usage of @kwdef")
    #     end
    # else
    #     kwdefs = nothing
    # end
    quote
        Base.@__doc__($(esc(expr)))
        $kwdefs
    end
end


text:
ContextManagers.maybeenter(::Any) = nothing

function _enter(source)
    context = ContextManagers.maybeenter(source)
    context === nothing && notasource(source)
    return something(context)
end

function notasource(@nospecialize(source))::Union{}
    error("Not a context source: ", source)
end

function ContextManagers.exit(x, @nospecialize(err))
    ContextManagers.exit(x)
    return nothing
end

ContextManagers.exit(ctx) = close(ctx)

Base.getindex(err::ContextManagers.NonException) = err.value

function _exit(context, @nospecialize(err))::Bool
    if !(err isa Exception)
        err = ContextManagers.NonException(err)
    end
    y = ContextManagers.exit(context, err)::Union{Handled,Nothing}
    # TODO: better error message
    return y isa Handled
end

macro with(doblock::Expr, bindings...)
    unsupported() = error("unsupported syntax")
    Meta.isexpr(doblock, :(->), 2) && doblock.args[1] == Expr(:tuple) || unsupported()
    body = doblock.args[2]
    ex = foldr(bindings; init = body) do b, ex
        @gensym err ans handled context
        if b isa Symbol
            value = resource = b
        elseif Meta.isexpr(b, :kw, 2)
            value, resource = b.args
            if value === :_
                @gensym value
            end
        else
            @gensym value
            resource = b
        end
        quote
            # Using `let` so that it works with `value === resource`
            let $context = $_enter($resource),
                $value = $ContextManagers.value($context),
                $ans = nothing,
                $handled = false

                try
                    $ans = $ex
                catch $err
                    $handled = $_exit($context, $err)
                    $handled || rethrow()
                end
                $handled || $ContextManagers.exit($context, nothing)
                $ans
            end
        end
    end
    return esc(ex)
end

ContextManagers.with(f::F, resource1, resources...) where {F} =
    _with(f, resource1, resources...)

_with(f) = f()
function _with(f::F, resource1, resources...) where {F}
    ans = nothing
    handled = false
    context = _enter(resource1)
    try
        x = ContextManagers.value(context)
        ans = _with(resources...) do args...
            f(x, args...)
        end
    catch err
        handled = _exit(context, err)
        handled || rethrow()
    end
    handled || _exit(context, nothing)
    return ans
end

ContextManagers.value(x) = x

ContextManagers.closing(x) = ContextManagers.onexit(close, x)

"""
    ContextManagers.onexit(cleanup, x)
    ContextManagers.onexit(cleanup)

Create a context manager that runs `cleanup(x)` upon exit.

# Example

```jldoctest
julia> using ContextManagers: @with, onexit

julia> @with(
           onexit(111) do x
               @show x
           end,
       ) do
       end;
x = 111
```
"""
ContextManagers.onexit

struct OnExit{C,T}
    close::C
    value::T
end

ContextManagers.onexit(close::F, value) where {F} = OnExit(close, value)
ContextManagers.maybeenter(c::OnExit) = c
ContextManagers.value(c::OnExit) = c.value
ContextManagers.exit(c::OnExit) = c.close(c.value)

ContextManagers.onexit(close::F) where {F} = OnExit(_ -> close(), nothing)

"""
    ContextManagers.onfail(cleanup, x)
    ContextManagers.onfail(cleanup)

Create a context manager that runs `cleanup(x)` upon unsuccessful exit.

# Example

```jldoctest
julia> using ContextManagers: @with, IgnoreError, onfail

julia> @with(
           onfail(111) do x
               @show x
           end,
       ) do
       end;  # prints nothing

julia> @with(
           IgnoreError(),
           onfail(111) do x
               @show x
           end,
       ) do
           error("error")
       end;
x = 111
```
"""
ContextManagers.onfail

struct OnFail{C,T}
    close::C
    value::T
end

ContextManagers.onfail(close::F, value) where {F} = OnFail(close, value)
ContextManagers.maybeenter(c::OnFail) = c
ContextManagers.value(c::OnFail) = c.value
function ContextManagers.exit(c::OnFail, err)
    if err !== nothing
        c.close(c.value)
    end
    return nothing
end

ContextManagers.onfail(close::F) where {F} = OnFail(_ -> close(), nothing)


text:
#=

Calculation of atmoshpere heat transport is not simple. Detail must be careful.
The standard version I found is here:
    http://www.atmos.albany.edu/facstaff/brose/classes/ATM623_Spring2015/Notes/Lectures/Lecture13%20--%20Heat%20transport.html

In particular, snow flux is not part of latent heat flux so must be added separately.

=#

using NCDatasets
using Formatting
using ArgParse
using Statistics
using JSON

include("MapTransform.jl")
include("constants.jl")
include("CESMReader.jl")

using .CESMReader
using .MapTransform

function mreplace(x)
    return replace(x, missing=>NaN)
end

function parse_commandline()

    s = ArgParseSettings()
    @add_arg_table s begin

        "--data-file-prefix"
            help = "Data filename prefix including folder and path until the timestamp. File extension `nc` is assumed."
            arg_type = String
            required = true

        "--data-file-timestamp-form"
            help = "Data filename timestamp form. Either `YEAR` or `YEAR_MONTH`."
            arg_type = String
            required = true

        "--output-file"
            help = "Output file."
            arg_type = String
            required = true

        "--domain-file"
            help = "Ocn domain file."
            arg_type = String
            required = true

        "--beg-year"
            help = "Year of begin."
            arg_type = Int64
            required = true

        "--end-year"
            help = "Year of end."
            arg_type = Int64
            required = true

    end

    return parse_args(ARGS, s)
end

parsed = parse_commandline()
print(json(parsed, 4))

Dataset(parsed["domain-file"], "r") do ds
    global area = mreplace(ds["area"][:])
    global mask = mreplace(ds["mask"][:])
    global lat  = mreplace(ds["yc"][:])

#    mask .= 1
    area .*= 4π * Re^2 / sum(area)
end

lat_bnd = collect(Float64, -90:1:90)
r = MapTransform.Relation(
    lat = lat,
    area = area,
    mask = mask,
    lat_bnd = lat_bnd,
)

_proxy = area * 0 .+ 1.0
sum_valid_area = MapTransform.∫∂a(r, _proxy)[end]

println("Sum of valid area: ", sum_valid_area, "; ratio: ", sum_valid_area / sum(area))

let

    if parsed["data-file-timestamp-form"] == "YEAR"
        filename_format = format("{:s}{{:04d}}.nc", joinpath(parsed["data-file-prefix"]))
        form = :YEAR
    elseif parsed["data-file-timestamp-form"] == "YEAR_MONTH"
        filename_format = format("{:s}{{:04d}}-{{:02d}}.nc", joinpath(parsed["data-file-prefix"]))
        form = :YEAR_MONTH
    end

    fh = FileHandler(filename_format=filename_format, form=form)

    beg_t = (parsed["beg-year"] - 1) * 12 + 1
    end_t = (parsed["end-year"] - 1) * 12 + 12

    _ADVT, _WKRSTT = getData(fh, ["ADVT", "WKRSTT"], (parsed["beg-year"], parsed["end-year"]), (:, :))

    Nt = end_t - beg_t + 1

    global ADVT      = zeros(Float64, length(r.lat_bnd)-1, Nt)
    global WKRSTT     = zeros(Float64, length(r.lat_bnd)-1, Nt)
    global WKRSTT_avg = zeros(Float64, Nt)
    global OHT_ADVT  = zeros(Float64, length(r.lat_bnd), Nt)
    global OHT_WKRSTT = zeros(Float64, length(r.lat_bnd), Nt)
    global OHT       = zeros(Float64, length(r.lat_bnd), Nt)

#    @. _WKRSTT += _Q_LOST / ρc

    for t = 1:Nt
        vw_ADVT = view(_ADVT, :, :, t) * ρc
        ADVT[:, t]     = MapTransform.transform(r, vw_ADVT)
        OHT_ADVT[:, t] = - MapTransform.∫∂a(r, vw_ADVT)

        vw_WKRSTT = view(_WKRSTT, :, :, t) * ρc
        WKRSTT[:, t]     = MapTransform.transform(r, vw_WKRSTT)
        WKRSTT_avg[t]    = MapTransform.mean(r, vw_WKRSTT)
        OHT_WKRSTT[:, t] = - MapTransform.∫∂a(r, (vw_WKRSTT .- WKRSTT_avg[t]))
    end

    global SHF = ADVT + WKRSTT
    global OHT = OHT_ADVT + OHT_WKRSTT

    global years = Int(Nt/12)
    global OHT_AM        = zeros(Float64, length(r.lat_bnd), years)
    global OHT_ADVT_AM   = zeros(Float64, length(r.lat_bnd), years)
    global OHT_WKRSTT_AM  = zeros(Float64, length(r.lat_bnd), years)
    for y = 1:years
        OHT_AM[:, y]       = mean(OHT[:, ((y-1)*12+1):(y*12)], dims=2)[:, 1]
        OHT_ADVT_AM[:, y]  = mean(OHT_ADVT[:, ((y-1)*12+1):(y*12)], dims=2)[:, 1]
        OHT_WKRSTT_AM[:, y] = mean(OHT_WKRSTT[:, ((y-1)*12+1):(y*12)], dims=2)[:, 1]
    end

end

Dataset(parsed["output-file"], "c") do ds

    defDim(ds, "time", Inf)
    defDim(ds, "year", years)
    defDim(ds, "N1", 1)
    defDim(ds, "lat_bnd", length(r.lat_bnd))
    defDim(ds, "lat",     length(r.lat_bnd)-1)

    ds.attrib["ocean_area"] = sum_valid_area

    for (varname, vardata, vardim, attrib) in [
        ("ADVT",    ADVT,        ("lat", "time"), Dict()),
        ("WKRSTT",   WKRSTT,       ("lat", "time"), Dict()),
        ("WKRSTT_avg",   WKRSTT_avg,       ("time",), Dict()),
        ("SHF",     SHF,         ("lat", "time"), Dict()),
        ("OHT_ADVT",  OHT_ADVT,  ("lat_bnd", "time"), Dict()),
        ("OHT_WKRSTT", OHT_WKRSTT, ("lat_bnd", "time"), Dict()),
        ("OHT",       OHT,       ("lat_bnd", "time"), Dict()),

        ("ADVT_MEAN",  mean(ADVT, dims=2)[:, 1],      ("lat",), Dict()),
        ("WKRSTT_MEAN", mean(WKRSTT, dims=2)[:, 1],     ("lat",), Dict()),
        ("SHF_MEAN",   mean(SHF, dims=2)[:, 1],       ("lat",), Dict()),
        ("OHT_ADVT_MEAN",   mean(OHT_ADVT, dims=2)[:, 1],  ("lat_bnd",), Dict()),
        ("OHT_WKRSTT_MEAN",  mean(OHT_WKRSTT, dims=2)[:, 1], ("lat_bnd",), Dict()),
        ("OHT_MEAN",        mean(OHT, dims=2)[:, 1],       ("lat_bnd",), Dict()),

        ("lat_bnd",          r.lat_bnd,                      ("lat_bnd",), Dict()),
        ("area",             r.∂a,                           ("lat",), Dict()),
    ]

        println("Doing var: ", varname)

        var = defVar(ds, varname, Float64, vardim)
        var.attrib["_FillValue"] = 1e20

        var = ds[varname]

        for (k, v) in attrib
            var.attrib[k] = v
        end

        rng = []
        for i in 1:length(vardim)-1
            push!(rng, Colon())
        end
        push!(rng, 1:size(vardata)[end])
        var[rng...] = vardata

    end

end

max_stars_repo_name: meteorologytoday/CESM-diagnostic

max_issues_repo_name: meteorologytoday/CESM-diagnostic

max_forks_repo_name: meteorologytoday/CESM-diagnostic


text:
include("../coin-or-common.jl")

name = "Bonmin"
version = Bonmin_version

sources = [
    GitSource("https://github.com/coin-or/Bonmin.git", Bonmin_gitsha),
    DirectorySource("./bundled")
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/Bonmin

if [[ ${target} == *mingw* ]]; then
    sed -i s/dllimport/dllexport/ "${includedir}/coin-or/IpoptConfig.h"
    atomic_patch -p1 ${WORKSPACE}/srcdir/patches/fix_dllexport.patch
fi

# Remove misleading libtool files
rm -f ${libdir}/*.la
update_configure_scripts

# old and custom autoconf
sed -i s/elf64ppc/elf64lppc/ configure

export CPPFLAGS="${CPPFLAGS} -I${includedir} -I${includedir}/coin -I${includedir}/coin-or"
export CXXFLAGS="${CXXFLAGS} -std=c++11"

if [[ ${target} == *mingw* ]]; then
    export LDFLAGS="-L${libdir}"
    export LT_LDFLAGS="-no-undefined"
fi

./configure \
    --prefix=${prefix} \
    --build=${MACHTYPE} \
    --host=${target} \
    --disable-pkg-config \
    --disable-debug \
    --enable-shared \
    lt_cv_deplibs_check_method=pass_all \
    --with-asl-lib="-lasl" \
    --with-coindepend-lib="-lCbc -lCgl -lOsiClp -lClp -lOsi -lCoinUtils -lipoptamplinterface -lipopt"

make -j${nproc}
make install

# Prevent IpoptConfig.h from being included in the tarball after editing it above.
rm "${includedir}/coin-or/IpoptConfig.h"
"""

filter!(x -> cxxstring_abi(x) != "cxx03", platforms)

# The products that we will ensure are always built
products = [
    LibraryProduct("libbonmin", :libbonmin),
    LibraryProduct("libbonminampl", :libbonminampl),
    ExecutableProduct("bonmin", :amplexe),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency("ASL_jll", ASL_version),
    Dependency("Cbc_jll", Cbc_version, compat="=$(Cbc_version)"),
    Dependency("Cgl_jll", Cgl_version),
    Dependency("Clp_jll", Clp_version),
    Dependency("Osi_jll", Osi_version),
    Dependency("CoinUtils_jll", CoinUtils_version),
    Dependency("Ipopt_jll", Ipopt_version, compat="=$(Ipopt_version)"),
    Dependency("CompilerSupportLibraries_jll"),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)

max_stars_repo_name: waralex/Yggdrasil

max_issues_repo_name: waralex/Yggdrasil

max_forks_repo_name: waralex/Yggdrasil


text:
module DataFrameMacros

using Base: ident_cmp
using DataFrames: transform, transform!, select, select!, combine, subset, subset!, ByRow, passmissing, groupby, AsTable, DataFrame

export @transform, @transform!, @select, @select!, @combine, @subset, @subset!, @groupby, @sort, @sort!, @unique

funcsymbols = :transform, :transform!, :select, :select!, :combine, :subset, :subset!, :unique

for f in funcsymbols
    @eval begin
        """
            @$($f)(df, args...; kwargs...)

        The `@$($f)` macro builds a `DataFrames.$($f)` call. Each expression in `args` is converted to a `src => function => sink` construct that conforms to the transformation mini-language of DataFrames.

        Keyword arguments `kwargs` are passed down to `$($f)` but have to be separated from the positional arguments by a semicolon `;`.

        The transformation logic for all DataFrameMacros macros is explained in the `DataFrameMacros` module docstring, accessible via `?DataFrameMacros`.
        """
        macro $f(exprs...)
            macrohelper($f, exprs...)
        end
    end
end

defaultbyrow(::typeof(transform)) = true
defaultbyrow(::typeof(select)) = true
defaultbyrow(::typeof(subset)) = true
defaultbyrow(::typeof(transform!)) = true
defaultbyrow(::typeof(select!)) = true
defaultbyrow(::typeof(subset!)) = true
defaultbyrow(::typeof(combine)) = false
defaultbyrow(::typeof(unique)) = true

macro groupby(exprs...)
    f = select
    df, func, kw_exprs = df_funcs_kwexprs(f, exprs...)

    select_kwexprs = [Expr(:kw, :copycols, false)]
    select_call = build_call(f, df, func, select_kwexprs)
    quote
        temp = $select_call
        df_copy = copy($(esc(df)))
        for name in names(temp)
            df_copy[!, name] = temp[!, name]
        end
        groupby(df_copy, names(temp); $(kw_exprs...))
    end
end

macro sort(exprs...)
    sorthelper(exprs...; mutate = false)
end

macro sort!(exprs...)
    sorthelper(exprs...; mutate = true)
end

function sorthelper(exprs...; mutate)
    f = select
    df, func, kw_exprs = df_funcs_kwexprs(f, exprs...)

    select_kwexprs = [Expr(:kw, :copycols, false)]
    select_call = build_call(f, df, func, select_kwexprs)
    if mutate
        quote
            temp = $select_call
            sp = sortperm(temp; $(kw_exprs...))
            $(esc(df)) = $(esc(df))[sp, :]
        end
    else
        quote
            temp = $select_call

            sp = sortperm(temp; $(kw_exprs...))

            df_copy = copy($(esc(df)))
            df_copy[sp, :]
        end
    end
end

function macrohelper(f, exprs...)
    df, converted, kw_exprs = df_funcs_kwexprs(f, exprs...)
    build_call(f, df, converted, kw_exprs)
end

function build_call(f, df, converted, kw_exprs)
    :($f($(esc(df)), $(converted...); $(map(esc, kw_exprs)...)))
end

function df_funcs_kwexprs(f, exprs...)
    if length(exprs) >= 1 && exprs[1] isa Expr && exprs[1].head == :parameters
        kw_exprs = exprs[1].args
        df = exprs[2]
        source_func_sink_exprs = exprs[3:end]
    else
        df = exprs[1]
        source_func_sink_exprs = exprs[2:end]
        kw_exprs = []
    end

    if length(source_func_sink_exprs) == 1 && is_block_expression(source_func_sink_exprs[1])
        source_func_sink_exprs = extract_source_func_sink_exprs_from_block(
            source_func_sink_exprs[1])
    end

    converted = map(e -> convert_source_funk_sink_expr(f, e, df), source_func_sink_exprs)
    df, converted, kw_exprs
end

is_block_expression(x) = false
is_block_expression(e::Expr) = e.head == :block

function extract_source_func_sink_exprs_from_block(block::Expr)
    filter(x -> !(x isa LineNumberNode), block.args)
end

convert_source_funk_sink_expr(f, x, df) = x

function convert_source_funk_sink_expr(f, e::Expr, df)
    target, formula = split_formula(e)
    flags, formula = extract_macro_flags(formula)

    if 't' in flags
        if target !== nothing
            error("There should be no target expression when the @t flag is used. The implicit target is `AsTable`. Target received was $target")
        end
        target = :(AsTable)
        formula = convert_automatic_astable_formula(formula)
    end

    columns = gather_columns(formula)
    func, columns = make_function_expr(formula, columns)
    clean_columns = map(c -> clean_column(c, df), columns)

    byrow = (defaultbyrow(f) && !('c' in flags)) ||
        (!defaultbyrow(f) && ('r' in flags))

    pass_missing = 'm' in flags

    if pass_missing
        func = :(passmissing($func))
    end

    func = byrow ? :(ByRow($func)) : :($func)

    trans_expr = if target === nothing
        :([$(clean_columns...)] => $(esc(func)))
    else
        :([$(clean_columns...)] => $(esc(func)) => $(esc(target)))
    end
end

function split_formula(e::Expr)
    if e.head != :(=)
        target = nothing
        formula = e
    else
        target = e.args[1]
        formula = e.args[2]
    end
    target, formula
end

function gather_columns(x; unique = true)
    columns = []
    gather_columns!(columns, x; unique = unique)
    columns
end

function gather_columns!(columns, x; unique)
    if is_column_expr(x)
        if !unique || x ∉ columns
            push!(columns, x)
        end
    elseif x isa Expr && !(is_escaped_symbol(x)) && !(is_dot_quotenode_expr(x))
        # we have to exclude quotenodes in dot syntax such as the b in `:a.b`
        args_to_scan = if x.head == :. && length(x.args) == 2 && x.args[2] isa QuoteNode
            x.args[1:1]
        else
            x.args
        end
        foreach(args_to_scan) do arg
            gather_columns!(columns, arg; unique = unique)
        end
    end
end

flagchars = "crmt"

extract_macro_flags(x) = "", x

function extract_macro_flags(e::Expr)
    if e.head == :macrocall && all(occursin(flagchars), string(e.args[1])[2:end]) &&
            length(e.args) == 3
        string(e.args[1])[2:end], e.args[3]
    else
        "", e
    end
end

is_dot_quotenode_expr(x) = false
# matches Module.[Submodule.Subsubmodule...].value
function is_dot_quotenode_expr(e::Expr)
    e.head == :(.) &&
        length(e.args) == 2 &&
        (e.args[1] isa Symbol || is_dot_quotenode_expr(e.args[1])) &&
        e.args[2] isa QuoteNode
end

is_column_expr(q::QuoteNode) = true
is_column_expr(x) = false
function is_column_expr(e::Expr)
    e.head == :$ && !is_escaped_symbol(e)
end

function make_function_expr(formula, columns)
    is_simple, symbol = is_simple_function_call(formula, columns)
    if is_simple
        # :a + :a returns only +(:a) if we don't set unique false
        return symbol, gather_columns(formula; unique = false)
    end

    newsyms = map(x -> gensym(), columns)
    replaced_formula = postwalk(formula) do e
        # check first if this is an escaped symbol
        # and if yes return it unwrapped
        if is_escaped_symbol(e)
            return e.args[1]
        end

        # check if this expression matches one of the column expressions
        # and replace it with a newsym if it matches
        i = findfirst(c -> c == e, columns)
        if i === nothing
            e
        else
            newsyms[i]
        end
    end
    expr = quote
        ($(newsyms...),) -> $replaced_formula
    end
    expr, columns
end

clean_column(x::QuoteNode, df) = x
clean_column(x, df) = :(symbolarg($x, $df))
function clean_column(e::Expr, df)
    e = if e.head == :$
        e.args[1]
    else
        e
    end
    if e isa String
        QuoteNode(Symbol(e))
    else
        :(symbolarg($(esc(e)), $(esc(df))))
    end
end

symbolarg(x::Int, df) = Symbol(names(df)[x])
symbolarg(sym::Symbol, df) = sym
symbolarg(s::String, df) = Symbol(s)

is_escaped_symbol(e::Expr) = e.head == :$ && length(e.args) == 1 && e.args[1] isa QuoteNode
is_escaped_symbol(x) = false

is_simple_function_call(x, columns) = false, nothing
function is_simple_function_call(expr::Expr, columns)
    is_call = expr.head == :call
    no_columns_in_funcpart = isempty(gather_columns(expr.args[1]))
    only_columns_in_argpart = length(expr.args) >= 2 &&
        all(x -> x in columns, expr.args[2:end])

    is_simple = is_call && no_columns_in_funcpart && only_columns_in_argpart

    is_simple, expr.args[1]
end

# from macrotools
walk(x, inner, outer) = outer(x)
walk(x::Expr, inner, outer) = outer(Expr(x.head, map(inner, x.args)...))
postwalk(f, x) = walk(x, x -> postwalk(f, x), f)

convert_automatic_astable_formula(x) = x
function convert_automatic_astable_formula(e::Expr)
    e_replaced, assigned_symbols, gensyms = replace_assigned_symbols(e)
    Expr(:block,
        e_replaced,
        Expr(:tuple,
            map(assigned_symbols, gensyms) do a, g
                Expr(:(=), a, g)
            end...
        )
    )
end

function replace_assigned_symbols(e)
    symbols = Symbol[]
    gensyms = Symbol[]

    function gensym_for_sym(sym)
        i = findfirst(==(sym), symbols)

        if i === nothing
            push!(symbols, sym)
            gs = gensym()
            push!(gensyms, gs)
        else
            gs = gensyms[i]
        end
        gs
    end

    new_ex = postwalk(e) do x
        # normal case :symbol = expression
        if x isa Expr && x.head == :(=) && x.args[1] isa QuoteNode
            sym = x.args[1].value
            gs = gensym_for_sym(sym)
            Expr(:(=), gs, x.args[2:end]...)
        # tuple destructuring case x, y = expr1, expr2 where x or y (or others) are :symbols
        elseif x isa Expr && x.head == :(=) && x.args[1] isa Expr && x.args[1].head == :tuple
            for (i, a) in enumerate(x.args[1].args)
                if a isa QuoteNode
                    sym = a.value
                    x.args[1].args[i] = gensym_for_sym(sym)
                end
            end
            x
        else
            x
        end
    end
    new_ex, symbols, gensyms
end

@doc """
DataFrameMacros offers macros which transform expressions for DataFrames functions that use the `source => function => sink` mini-language.
The supported functions are `@transform`/`@transform!`, `@select/@select!`, `@groupby`, `@combine`, `@subset`/`@subset!`, `@sort`/`@sort!` and `@unique`.

All macros have signatures of the form:
```julia
@macro(df, args...; kwargs...)
```

Each positional argument in `args` is converted to a `source => function => sink` expression for the transformation mini-language of DataFrames.
By default, all macros execute the given function **by-row**, only `@combine` executes **by-column**.

For example, the following pairs of expressions are equivalent:

```julia
transform(df, :x => ByRow(x -> x + 1) => :y)
@transform(df, :y = :x + 1)

sort(df, :x => ByRow(x -> x ^ 2))
@sort(df, :x ^ 2)

combine(df, :x => (x -> sum(x) / 5) => :result)
@combine(df, :result = sum(:x) / 5)
```

## Column references

Each positional argument must be of the form `[sink =] some_expression`.
Columns can be referenced within `sink` or `some_expression` using a `Symbol`, a `String`, or an `Int`.
Any column identifier that is not a `Symbol` must be prefaced with the interpolation symbol `\$`.
The `\$` interpolation symbol also allows to use variables or expressions that evaluate to column identifiers.

The five expressions in the following code block are equivalent.

```julia
using DataFrames
using DataFrameMacros

df = DataFrame(x = 1:3)

@transform(df, :y = :x + 1)
@transform(df, :y = \$"x" + 1)
@transform(df, :y = \$1 + 1)
col = :x
@transform(df, :y = \$col + 1)
cols = [:x, :y, :z]
@transform(df, :y = \$(cols[1]) + 1)
```

## Passing multiple expressions

Multiple expressions can be passed as multiple positional arguments, or alternatively as separate lines in a `begin end` block. You can use parentheses, or omit them. The following expressions are equivalent:

```julia
@transform(df, :y = :x + 1, :z = :x * 2)
@transform df :y = :x + 1 :z = :x * 2
@transform df begin
    :y = :x + 1
    :z = :x * 2
end
@transform(df, begin
    :y = :x + 1
    :z = :x * 2
end)
```

## Flag macros

You can modify the behavior of all macros using flag macros, which are not real macros but only signal changed behavior for a positional argument to the outer macro.

Each flag is specified with a single character, and you can combine these characters as well.
The supported flags are:

| character | meaning |
|:--|:--|
| r | Switch to **by-row** processing. |
| c | Switch to **by-column** processing. |
| m | Wrap the function expression in `passmissing`. |
| t | Collect all `:symbol = expression` expressions into a `NamedTuple` where `(; symbol = expression, ...)` and set the sink to `AsTable`. |

### Example `@c`

To compute a centered column with `@transform`, you need access to the whole column at once and signal this with the `@c` flag.

```julia
using Statistics
using DataFrames
using DataFrameMacros

julia> df = DataFrame(x = 1:3)
3×1 DataFrame
 Row │ x
     │ Int64
─────┼───────
   1 │     1
   2 │     2
   3 │     3

julia> @transform(df, :x_centered = @c :x .- mean(:x))
3×2 DataFrame
 Row │ x      x_centered
     │ Int64  Float64
─────┼───────────────────
   1 │     1        -1.0
   2 │     2         0.0
   3 │     3         1.0
```

### Example `@m`

Many functions need to be wrapped in `passmissing` to correctly return `missing` if any input is `missing`.
This can be achieved with the `@m` flag macro.

```julia
julia> df = DataFrame(name = ["alice", "bob", missing])
3×1 DataFrame
 Row │ name
     │ String?
─────┼─────────
   1 │ alice
   2 │ bob
   3 │ missing

julia> @transform(df, :name_upper = @m uppercasefirst(:name))
3×2 DataFrame
 Row │ name     name_upper
     │ String?  String?
─────┼─────────────────────
   1 │ alice    Alice
   2 │ bob      Bob
   3 │ missing  missing
```

### Example `@t`

In DataFrames, you can return a `NamedTuple` from a function and then automatically expand it into separate columns by using `AsTable` as the sink value. To simplify this process, you can use the `@t` flag macro, which collects all statements of the form `:symbol = expression` in the function body, collects them into a `NamedTuple`, and sets the sink argument to `AsTable`.

```julia
julia> df = DataFrame(name = ["Alice Smith", "Bob Miller"])
2×1 DataFrame
 Row │ name
     │ String
─────┼─────────────
   1 │ Alice Smith
   2 │ Bob Miller

julia> @transform(df, @t begin
           s = split(:name)
           :first_name = s[1]
           :last_name = s[2]
       end)
2×3 DataFrame
 Row │ name         first_name  last_name
     │ String       SubString…  SubString…
─────┼─────────────────────────────────────
   1 │ Alice Smith  Alice       Smith
   2 │ Bob Miller   Bob         Miller
```

The `@t` flag also works with tuple destructuring syntax, so the previous example can be shortened to:

```julia
@transform(df, @t :first_name, :last_name = split(:name))
```

""" DataFrameMacros

const _titanic = include("titanic.jl")
titanic() = deepcopy(_titanic)

end


text:
export PointRobot2D

# ---------- Point Robot ----------

type PointRobot2D{S<:Shape2D} <: SweptCollisionChecker
    obstacles::S      # pretty much always Compound2D
    count::Int
end
PointRobot2D(obstacles) = PointRobot2D(obstacles, 0)
changeprecision{T<:AbstractFloat}(::Type{T}, CC::PointRobot2D) = PointRobot2D(changeprecision(T, CC.obstacles))

is_free_state(v::AbstractVector, CC::PointRobot2D) = !colliding(v, CC.obstacles)
is_free_motion(v::AbstractVector, w::AbstractVector, CC::PointRobot2D) = (CC.count += 1;
                                                                          !colliding(Line(v,w), CC.obstacles))
function is_free_path(path::Path, CC::PointRobot2D)
    for i in 1:length(path)-1
        !is_free_motion(path[i], path[i+1], CC) && return false
    end
    true
end
inflate(CC::PointRobot2D, ε; roundcorners = true) =
    ε > 0 ? PointRobot2D(inflate(CC.obstacles, ε, roundcorners = roundcorners)) : CC
addobstacle(CC::PointRobot2D, o::Shape2D) = PointRobot2D(Compound2D(CC.obstacles, o))
addblocker(CC::PointRobot2D, p::AbstractVector, r) = addobstacle(CC, Circle(p, r))
closest(p::AbstractVector, CC::PointRobot2D, W::AbstractMatrix) = closest(p, CC.obstacles, W)
closeR(p::AbstractVector, CC::PointRobot2D, W::AbstractMatrix, r2) = closeR(p, CC.obstacles, W, r2)

function plot(CC::PointRobot2D, lo = zeros(2), hi = ones(2); kwargs...)
    plot(CC.obstacles, color = "red", edgecolor = "none",
         xmin = lo[1], xmax = hi[1], ymin = lo[2], ymax = hi[2]; kwargs...)
end
# composable(CC::PointRobot2D, opts = (stroke("black"),)) = composable(CC.obstacles, opts)


text:
"""
    space(::SiteType"Fermion";
          conserve_qns=false,
          conserve_nf=conserve_qns,
          conserve_nfparity=conserve_qns,
          qnname_nf = "Nf",
          qnname_nfparity = "NfParity",
          qnname_sz = "Sz",
          conserve_sz = false)

Create the Hilbert space for a site of type "Fermion".

Optionally specify the conserved symmetries and their quantum number labels.
"""
function space(
  ::SiteType"Fermion";
  conserve_qns=false,
  conserve_nf=conserve_qns,
  conserve_nfparity=conserve_qns,
  qnname_nf="Nf",
  qnname_nfparity="NfParity",
  qnname_sz="Sz",
  conserve_sz=false,
  # Deprecated
  conserve_parity=nothing,
)
  if !isnothing(conserve_parity)
    conserve_nfparity = conserve_parity
  end
  if conserve_sz == true
    conserve_sz = "Up"
  end
  if conserve_nf && conserve_sz == "Up"
    zer = QN((qnname_nf, 0, -1), (qnname_sz, 0)) => 1
    one = QN((qnname_nf, 1, -1), (qnname_sz, 1)) => 1
    return [zer, one]
  elseif conserve_nf && conserve_sz == "Dn"
    zer = QN((qnname_nf, 0, -1), (qnname_sz, 0)) => 1
    one = QN((qnname_nf, 1, -1), (qnname_sz, -1)) => 1
    return [zer, one]
  elseif conserve_nfparity && conserve_sz == "Up"
    zer = QN((qnname_nfparity, 0, -2), (qnname_sz, 0)) => 1
    one = QN((qnname_nfparity, 1, -2), (qnname_sz, 1)) => 1
    return [zer, one]
  elseif conserve_nfparity && conserve_sz == "Dn"
    zer = QN((qnname_nfparity, 0, -2), (qnname_sz, 0)) => 1
    one = QN((qnname_nfparity, 1, -2), (qnname_sz, -1)) => 1
    return [zer, one]
  elseif conserve_nf
    zer = QN(qnname_nf, 0, -1) => 1
    one = QN(qnname_nf, 1, -1) => 1
    return [zer, one]
  elseif conserve_nfparity
    zer = QN(qnname_nfparity, 0, -2) => 1
    one = QN(qnname_nfparity, 1, -2) => 1
    return [zer, one]
  end
  return 2
end

val(::ValName"Emp", ::SiteType"Fermion") = 1
val(::ValName"Occ", ::SiteType"Fermion") = 2
val(::ValName"0", st::SiteType"Fermion") = val(ValName("Emp"), st)
val(::ValName"1", st::SiteType"Fermion") = val(ValName("Occ"), st)

state(::StateName"Emp", ::SiteType"Fermion") = [1.0 0.0]
state(::StateName"Occ", ::SiteType"Fermion") = [0.0 1.0]
state(::StateName"0", st::SiteType"Fermion") = state(StateName("Emp"), st)
state(::StateName"1", st::SiteType"Fermion") = state(StateName("Occ"), st)

function op!(Op::ITensor, ::OpName"N", ::SiteType"Fermion", s::Index)
  return Op[s' => 2, s => 2] = 1.0
end

function op!(Op::ITensor, ::OpName"C", ::SiteType"Fermion", s::Index)
  return Op[s' => 1, s => 2] = 1.0
end

function op!(Op::ITensor, ::OpName"Cdag", ::SiteType"Fermion", s::Index)
  return Op[s' => 2, s => 1] = 1.0
end

function op!(Op::ITensor, ::OpName"F", ::SiteType"Fermion", s::Index)
  Op[s' => 1, s => 1] = +1.0
  return Op[s' => 2, s => 2] = -1.0
end

has_fermion_string(::OpName"C", ::SiteType"Fermion") = true
has_fermion_string(::OpName"Cdag", ::SiteType"Fermion") = true


text:
module CellDataTests

using Test

@testset "CellFields" begin include("CellFieldsTests.jl") end

@testset "CellQuadratures" begin include("CellQuadraturesTests.jl") end

@testset "DomainContributions" begin include("DomainContributionsTests.jl") end

@testset "DiracDeltas" begin include("DiracDeltasTests.jl") end

@testset "CellStates" begin include("CellStatesTests.jl") end

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

@testset "AttachDirichlet" begin include("AttachDirichletTests.jl") end

@testset "AttachConstraints" begin include("AttachConstraintsTests.jl") end

end # module


text:
using Test
using ModelingToolkit, OrdinaryDiffEq

@parameters t
@connector function Pin(;name)
    sts = @variables v(t)=1.0 i(t)=1.0
    ODESystem(Equation[], t, sts, []; name=name)
end

function ModelingToolkit.connect(::Type{Pin}, ps...)
    eqs = [
           0 ~ sum(p->p.i, ps) # KCL
          ]
    # KVL
    for i in 1:length(ps)-1
        push!(eqs, ps[i].v ~ ps[i+1].v)
    end

    return eqs
end

function Ground(;name)
    @named g = Pin()
    eqs = [g.v ~ 0]
    compose(ODESystem(eqs, t, [], []; name=name), g)
end

function OnePort(;name)
    @named p = Pin()
    @named n = Pin()
    sts = @variables v(t)=1.0 i(t)=1.0
    eqs = [
           v ~ p.v - n.v
           0 ~ p.i + n.i
           i ~ p.i
          ]
    compose(ODESystem(eqs, t, sts, []; name=name), p, n)
end

function Resistor(;name, R = 1.0)
    @named oneport = OnePort()
    @unpack v, i = oneport
    ps = @parameters R=R
    eqs = [
           v ~ i * R
          ]
    extend(ODESystem(eqs, t, [], ps; name=name), oneport)
end

function Capacitor(;name, C = 1.0)
    @named oneport = OnePort()
    @unpack v, i = oneport
    ps = @parameters C=C
    D = Differential(t)
    eqs = [
           D(v) ~ i / C
          ]
    extend(ODESystem(eqs, t, [], ps; name=name), oneport)
end

function ConstantVoltage(;name, V = 1.0)
    @named oneport = OnePort()
    @unpack v = oneport
    ps = @parameters V=V
    eqs = [
           V ~ v
          ]
    extend(ODESystem(eqs, t, [], ps; name=name), oneport)
end

function Inductor(; name, L = 1.0)
    @named oneport = OnePort()
    @unpack v, i = oneport
    ps = @parameters L=L
    D = Differential(t)
    eqs = [
           D(i) ~ v / L
          ]
    extend(ODESystem(eqs, t, [], ps; name=name), oneport)
end


text:
f_qbftest(t) = 1/(2+cos(2*pi*t))^2

function test_qbf()
    N = 32
    for degree in 1:4
        basis = BSplineTranslatesBasis(N, degree, 0, 1)
        coef_direct = [DomainIntegrals.integral(t -> f_qbftest(t)*BasisFunctions.unsafe_eval_element(basis, i, t), FrameFun.support(basis, i)) for i in 1:length(basis)]
        for oversamplingfactor in 1:5
            A = IntegralEquations2D.qbf_operator(basis, degree, oversamplingfactor)
            coef = A * f_qbftest
            @test norm(coef-coef_direct) < 1e-4
        end
    end

    # Do a BigFloat test
    N = 500
    degree = 11
    oversamplingfactor = 6
    basis = BSplineTranslatesBasis(N, degree, big(0.), big(1.))
    A = IntegralEquations2D.qbf_operator(basis, degree, oversamplingfactor)
    coef = A * f_qbftest
    # c10 = integral(t -> f_qbftest(t)*BasisFunctions.unsafe_eval_element(basis, 10, t), FrameFun.support(basis, 10))
    c10 = BigFloat("0.004990536452129881961287340961088526093027349985606")
    @test abs(coef[10]-c10) < 1e-45
end


text:
using FractionalSystems, Plots

function Danca(du, u, p, t)
    du[1] = -u[1]+u[2]
    du[2] = -u[3]*sign(u[1])+u[4]
    du[3] = abs(u[1])-1
    du[4] = -0.5*u[2]
end

LE, tspan=FOLyapunov(Danca, 0.98, 0, 0.02, 300, [0.1; 0.1; 0.1; 0.1], 0.005, 1000)

plot(tspan, LE[1, :])
plot!(tspan, LE[2, :])
plot!(tspan, LE[3, :])
plot!(tspan, LE[4, :])


text:
struct TenseLayer{M,B,F} <: SugarLayer
    weight::M
    weight_hooks::Vector{GradientHook}
    bias::B
    bias_hooks::Vector{GradientHook}
    σ::F
end

function TenseLayer(W::M, bias = true, σ::F = identity) where {M<:AbstractMatrix, F}
    b = Matrix(Flux.create_bias( W, bias, size( W, 2 ) )')
    TenseLayer( W, Vector{GradientHook}(undef,0), b, Vector{GradientHook}(undef,0), σ )
end

function TenseLayer(    in::Integer, out::Integer, σ = identity;
                        initW = nothing, initb = nothing,
                        init = Flux.glorot_uniform, bias = true )
    W = if initW !== nothing
        Base.depwarn("keyword initW is deprecated, please use init (which similarly accepts a funtion like randn)", :Dense)
        Flux.initW(in, out)
    else
        init(in, out)
    end
    b = if bias === true && initb !== nothing
        Base.depwarn("keyword initb is deprecated, please simply supply the bias vector, bias=initb(out)", :Dense)
        Flux.initb(out)
    else
        bias
    end
    return TenseLayer( W, b, σ )
end
Flux.@functor TenseLayer

function ( a::TenseLayer )( x::AbstractVecOrMat )
    W, b, σ = a.weight, a.bias, a.σ
    W_h = composer( a.weight_hooks )( W )
    b_h = composer( a.bias_hooks )( b )
    return σ.( x * W_h .+ b_h )
end

function Base.show(io::IO, l::TenseLayer)
    print(io, "TenseLayer(", size(l.weight, 2), ", ", size(l.weight, 1))
    l.σ == identity || print(io, ", ", l.σ)
    l.bias == Flux.Zeros() && print(io, ", bias=false")
    print(io, ")")
end


text:
using Oceananigans: velocity_div!, compute_w_from_continuity!

function time_stepping_works(arch, ft)
    Nx, Ny, Nz = 16, 16, 16
    Lx, Ly, Lz = 1, 2, 3
    Δt = 1

    model = Model(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz), arch=arch, float_type=ft)
    time_step!(model, 1, Δt)

    # Just testing that no errors/crashes happen when time stepping.
    return true
end

function run_first_AB2_time_step_tests(arch, ft)
    Nx, Ny, Nz = 16, 16, 16
    Lx, Ly, Lz = 1, 2, 3
    Δt = 1

    model = Model(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz), arch=arch, float_type=ft)

    add_ones(args...) = 1.0
    model.forcing = Forcing(nothing, nothing, nothing, add_ones, nothing)

    time_step!(model, 1, Δt)

    # Test that GT = 1 after first time step and that AB2 actually reduced to forward Euler.
    @test all(model.G.Gu.data .≈ 0)
    @test all(model.G.Gv.data .≈ 0)
    @test all(model.G.Gw.data .≈ 0)
    @test all(model.G.GT.data .≈ 1.0)
    @test all(model.G.GS.data .≈ 0)

    return nothing
end

"""
    This test ensures that when we compute w from the continuity equation that the full velocity field
    is divergence-free.
"""
function test_compute_w_from_continuity(arch, ft)
    Nx, Ny, Nz = 16, 16, 16
    Lx, Ly, Lz = 16, 16, 16

    grid = RegularCartesianGrid(ft, (Nx, Ny, Nz), (Lx, Ly, Lz))

    u = FaceFieldX(arch, grid)
    v = FaceFieldY(arch, grid)
    w = FaceFieldZ(arch, grid)
    div_u = CellField(arch, grid)

    u.data .= rand(Nx, Ny, Nz)
    v.data .= rand(Nx, Ny, Nz)

    compute_w_from_continuity!(grid, u.data, v.data, w.data)

    velocity_div!(grid, u.data, v.data, w.data, div_u.data)

    # Set div_u to zero at the bottom because the initial velocity field is not divergence-free
    # so we end up some divergence at the bottom if we don't do this.
    div_u.data[:, :, end] .= zero(ft)

    min_div = minimum(div_u)
    max_div = minimum(div_u)
    sum_div = sum(div_u)
    abs_sum_div = sum(abs.(div_u))
    @info "Velocity divergence after recomputing w ($arch, $ft): min=$min_div, max=$max_div, sum=$sum_div, abs_sum=$abs_sum_div"

    all(isapprox.(div_u.data, 0; atol=2*eps(ft)))
end

"""
    This tests to make sure that the velocity field remains incompressible (or divergence-free) as the model is time
    stepped. It just initializes a cube shaped hot bubble perturbation in the center of the 3D domain to induce a
    velocity field.
"""
function incompressible_in_time(arch, ft, Nt)
    Nx, Ny, Nz = 32, 32, 32
    Lx, Ly, Lz = 10, 10, 10

    model = Model(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz), arch=arch, float_type=ft)

    grid = model.grid
    u, v, w = model.velocities.u, model.velocities.v, model.velocities.w

    div_u = CellField(arch, model.grid)

    # Just add a temperature perturbation so we get some velocity field.
    @. model.tracers.T.data[8:24, 8:24, 8:24] += 0.01

    time_step!(model, Nt, 0.05)

    velocity_div!(grid, u, v, w, div_u)

    min_div = minimum(div_u)
    max_div = minimum(div_u)
    sum_div = sum(div_u)
    abs_sum_div = sum(abs.(div_u))
    @info "Velocity divergence after $Nt time steps ($arch, $ft): min=$min_div, max=$max_div, sum=$sum_div, abs_sum=$abs_sum_div"

    # We are comparing with 0 so we use absolute tolerances. They are a bit larger than eps(Float64) and eps(Float32)
    # because we are summing over the absolute value of many machine epsilons. A better atol value may be
    # Nx*Ny*Nz*eps(ft) but it's much higher than the observed abs_sum_div.
    if ft == Float64
        return isapprox(abs_sum_div, 0; atol=5e-16)
    elseif ft == Float32
        return isapprox(abs_sum_div, 0; atol=1e-7)
    end
end


text:
using SententialDecisionDiagrams
const SDD = SententialDecisionDiagrams

# set up vtree and manager
vtree = SDD.read_vtree("$(@__DIR__)/input/big-swap.vtree")
manager = SDD.sdd_manager(vtree)

println("reading sdd from file ...")
α = SDD.read_sdd("$(@__DIR__)/input/big-swap.sdd", manager)
println("sdd size = $(SDD.size(α))")

# to perform a swap, we need the manager's vtree
manager_vtree = SDD.vtree(manager)

# ref alpha (no dead nodes when swapping)
SDD.ref(α)
#
# # using size of sdd normalized for manager_vtree as baseline for limit
SDD.init_vtree_size_limit(manager_vtree, manager)
#
limit = 2.0
SDD.set_vtree_operation_size_limit(limit, manager)

println("modifying vtree (swap node 7) (limit growth by $(limit)x) ... ")
succeeded = SDD.swap(manager_vtree, manager, 1)
if succeeded == 1; println("succeeded!") else println("did not succeed!") end
println("sdd size = $(SDD.size(α))")

println("modifying vtree (swap node 7) (no limit) ... ")
succeeded = SDD.swap(manager_vtree, manager, 0)
if succeeded == 1; println("succeeded!") else println("did not succeed!") end
println("sdd size = $(SDD.size(α))")

println("updating baseline of size limit ...")
SDD.update_vtree_size_limit(manager)

left_vtree = SDD.left(manager_vtree)
limit = 1.2
SDD.set_vtree_operation_size_limit(limit, manager)
println("modifying vtree (swap node 5) (limit growth by $(limit)x) ... ")
succeeded = SDD.swap(left_vtree, manager, 1)
if succeeded == 1; println("succeeded!") else println("did not succeed!") end
println("sdd size = $(SDD.size(α))")

limit = 1.3
SDD.set_vtree_operation_size_limit(limit, manager)
println("modifying vtree (swap node 5) (limit growth by $(limit)x) ... ")
succeeded = SDD.swap(left_vtree, manager, 1)
if succeeded == 1; println("succeeded!") else println("did not succeed!") end
println("sdd size = $(SDD.size(α))")

# deref alpha, since ref's are no longer needed
SDD.deref(α)