text:
import QuantumOpticsBase, LinearAlgebra, OrdinaryDiffEq, QuadGK, DelayDiffEq

"""
    master_int(W0, tspan, Ham_0, Ham_I;
    \treltol=1.0e-12, abstol=1.0e-12, int_reltol=1.0e-8, int_abstol=0.0,
    \tfout=nothing, alg=DelayDiffEq.MethodOfSteps(DelayDiffEq.Vern6()))

Integrate Quantum Master equation

``\\frac{d}{d t} \\rho^{(I)}(t) = - \\frac{i}{\\hbar} [ \\hat{H}_I^{(I)}(t), \\rho^{(I)}(t_0) ]
-\\frac{1}{\\hbar^2} \\int_{t_0}^{t_1} \\text{d} \\tau \\: [ \\hat{H}_I^{(I)}(t), [ \\hat{H}_I^{(I)}(\\tau), \\rho^{(I)}(\\tau) ]] ``

``H = H_S + H_B + H_I = H_0 + H_I, \\quad \\hbar = 1. ``

# Arguments
* `W0`: Initial state vector (can be a bra or a ket) or initial propagator.
* `tspan`: Vector specifying the points of time for which output should be displayed.s
* `Ham_0`: System and bath Hamiltonian as Operator.
* `Ham_I`: Interaction Hamiltonian as Operator.
* `reltol`: Relative tolerance for DiffEqCallbacks solver and its inner states.
* `abstol`: Absolute tolerance for DiffEqCallbacks solver and its inner states.
* `int_reltol`: Relative tolerance for QuadGK solver and its inner states.
* `int_abstol`: Absolute tolerance for QuadGK solver and its inner states.
* `fout=nothing`: If given, this function `fout(t, rho)` is called every time
        an output should be displayed. ATTENTION: The state `rho` is neither
        normalized nor permanent! It is still in use by the ode solver and
        therefore must not be changed.
* `alg`: Algorithm with which DiffEqCallbacks will solve QME equation.
"""
function master_int(
    W0::T,
    tspan::Array,
    Ham_0::U,
    Ham_I::V;
    reltol::AbstractFloat = 1.0e-12,
    abstol::AbstractFloat = 1.0e-12,
    int_reltol::AbstractFloat = 1.0e-8,
    int_abstol::AbstractFloat = 0.0,
    alg::Any = DelayDiffEq.MethodOfSteps(DelayDiffEq.Vern6()),
    fout::Union{Function,Nothing} = nothing,
    kwargs...,
) where {B<:Basis,T<:Operator{B,B},U<:Operator{B,B},V<:Operator{B,B}}
    history_fun(p, t) = T(W0.basis_l, W0.basis_r, zeros(ComplexF64, size(W0.data)))
    # (du,u,h,p,t)
    tmp = copy(W0.data)
    dmaster_(t, rho::T, drho::T, history_fun, p) = dmaster_int(
        t,
        rho,
        drho,
        history_fun,
        tmp,
        p,
        W0,
        Ham_0,
        Ham_I,
        int_reltol,
        int_abstol,
    )
    tspan_ = convert(Vector{float(eltype(tspan))}, tspan)
    x0 = W0.data
    state = T(W0.basis_l, W0.basis_r, W0.data)
    dstate = T(W0.basis_l, W0.basis_r, W0.data)
    OpenQuantumSystems.integrate_delayed(
        tspan_,
        dmaster_,
        history_fun,
        x0,
        state,
        dstate,
        fout;
        reltol = reltol,
        abstol = abstol,
        alg = alg,
        kwargs...,
    )
end

function dmaster_int(
    t::AbstractFloat,
    rho::T,
    drho::T,
    history_fun,
    tmp::Array,
    p,
    W0,
    Ham_0::U,
    Ham_I::V,
    int_reltol::AbstractFloat,
    int_abstol::AbstractFloat,
) where {B<:Basis,T<:Operator{B,B},U<:Operator{B,B},V<:Operator{B,B}}
    Ham_II_t = getInteractionHamIPicture(Ham_0, Ham_I, t)
    QuantumOpticsBase.mul!(drho, Ham_II_t, W0, -eltype(rho)(im), zero(eltype(rho)))
    QuantumOpticsBase.mul!(drho, W0, Ham_II_t, eltype(rho)(im), one(eltype(rho)))

    kernel_integrated, err = QuadGK.quadgk(
        s -> kernel_int(t, s, tmp, history_fun, p, Ham_II_t.data, Ham_0, Ham_I),
        0,
        t,
        rtol = int_reltol,
        atol = int_abstol,
    )
    LinearAlgebra.mul!(
        drho.data,
        -one(eltype(rho)),
        kernel_integrated,
        one(eltype(rho)),
        one(eltype(rho)),
    )
    return drho
end

function kernel_int(t, s, tmp, h, p, Ham_II_t, Ham_0, Ham_I)
    Ham_II_s = getInteractionHamIPictureA(Ham_0, Ham_I, s)
    rho = h(p, s)

    if (typeof(rho) <: Operator)
        rho = rho.data
    end
    drho = deepcopy(rho)
    # commutator(Ham.data, commutator(Ham.data, tmp))
    QuantumOpticsBase.mul!(tmp, Ham_II_s, rho, eltype(rho)(1), zero(eltype(rho)))
    QuantumOpticsBase.mul!(tmp, rho, Ham_II_s, -eltype(rho)(1), one(eltype(rho)))

    QuantumOpticsBase.mul!(drho, Ham_II_t, tmp, eltype(rho)(1), zero(eltype(rho)))
    QuantumOpticsBase.mul!(drho, tmp, Ham_II_t, -eltype(rho)(1), one(eltype(rho)))
    return drho
end

"""
    master(W0, tspan, Ham;
    \treltol=1.0e-12, abstol=1.0e-12, int_reltol=1.0e-8, int_abstol=0.0,
    \tfout=nothing, alg=DelayDiffEq.MethodOfSteps(DelayDiffEq.Vern6()))

Integrate Quantum Master equation

``\\frac{d}{d t} \\rho(t) = - \\frac{i}{\\hbar} [ \\hat{H}, \\rho(t_0) ]
-\\frac{1}{\\hbar^2} \\int_{t_0}^{t_1} \\text{d} \\tau \\: [ \\hat{H}, [ \\hat{H}, \\rho(\\tau) ]]
,\\quad \\hbar = 1. ``

# Arguments
* `W0`: Initial state vector (can be a bra or a ket) or initial propagator.
* `tspan`: Vector specifying the points of time for which output should be displayed.s
* `Ham`: Arbitrary operator specifying the Hamiltonian.
* `reltol`: Relative tolerance for DiffEqCallbacks solver and its inner states.
* `abstol`: Absolute tolerance for DiffEqCallbacks solver and its inner states.
* `int_reltol`: Relative tolerance for QuadGK solver and its inner states.
* `int_abstol`: Absolute tolerance for QuadGK solver and its inner states.
* `fout=nothing`: If given, this function `fout(t, rho)` is called every time
        an output should be displayed. ATTENTION: The state `rho` is neither
        normalized nor permanent! It is still in use by the ode solver and
        therefore must not be changed.
* `alg`: Algorithm with which DiffEqCallbacks will solve QME equation.
"""
function master(
    W0::T,
    tspan,
    Ham::U;
    reltol::Float64 = 1.0e-12,
    abstol::Float64 = 1.0e-12,
    int_reltol::AbstractFloat = 1.0e-8,
    int_abstol::AbstractFloat = 0.0,
    alg::Any = DelayDiffEq.MethodOfSteps(DelayDiffEq.Vern6()),
    fout::Union{Function,Nothing} = nothing,
    kwargs...,
) where {B<:Basis,T<:Operator{B,B},U<:Operator{B,B}}
    history_fun(p, t) = T(W0.basis_l, W0.basis_r, zeros(ComplexF64, size(W0.data)))
    # (du,u,h,p,t)
    tmp = copy(W0.data)
    dmaster_(t, rho::T, drho::T, history_fun, p) =
        dmaster(t, rho, drho, history_fun, tmp, p, W0, Ham)
    tspan_ = convert(Vector{float(eltype(tspan))}, tspan)
    x0 = W0.data
    state = T(W0.basis_l, W0.basis_r, W0.data)
    dstate = T(W0.basis_l, W0.basis_r, W0.data)
    OpenQuantumSystems.integrate_delayed(
        tspan_,
        dmaster_,
        history_fun,
        x0,
        state,
        dstate,
        fout;
        reltol = reltol,
        abstol = abstol,
        alg = alg,
        kwargs...,
    )
end

function dmaster(
    t::AbstractFloat,
    rho::T,
    drho::T,
    history_fun,
    tmp::Array,
    p,
    W0,
    Ham::U,
) where {B<:Basis,T<:Operator{B,B},U<:Operator{B,B}}
    QuantumOpticsBase.mul!(drho, Ham, W0, -eltype(rho)(im), zero(eltype(rho)))
    QuantumOpticsBase.mul!(drho, W0, Ham, eltype(rho)(im), one(eltype(rho)))

    kernel_integrated, err =
        QuadGK.quadgk(s -> kernel(t, s, tmp, history_fun, p, Ham.data), 0, t, rtol = 1e-8)
    LinearAlgebra.mul!(
        drho.data,
        -one(eltype(rho)),
        kernel_integrated,
        one(eltype(rho)),
        one(eltype(rho)),
    )
    return drho
end

function kernel(t, s, tmp, h, p, Ham)
    rho = h(p, s)
    if (typeof(rho) <: Operator)
        rho = rho.data
    end
    drho = deepcopy(rho)
    # commutator(Ham.data, commutator(Ham.data, tmp))
    QuantumOpticsBase.mul!(tmp, Ham, rho, eltype(rho)(1), zero(eltype(rho)))
    QuantumOpticsBase.mul!(tmp, rho, Ham, -eltype(rho)(1), one(eltype(rho)))

    QuantumOpticsBase.mul!(drho, Ham, tmp, eltype(rho)(1), zero(eltype(rho)))
    QuantumOpticsBase.mul!(drho, tmp, Ham, -eltype(rho)(1), one(eltype(rho)))
    return drho
end


text:
"""
    unipartitemotifs()

The names of the motifs come from Stouffer et al. (2007) -- especially Fig. 1,
available online at
<http://rspb.royalsocietypublishing.org/content/274/1621/1931.figures-only>

The motifs are returned as a named tuple, with every motif identified by its
name in the original publication. The species are named :a. :b, and :c.
"""
function unipartitemotifs()

  # Everything is stored in a Dict, and the keys are symbols with the names of
  # the motifs.
  motifs = (
    S1 = UnipartiteNetwork([0 1 0; 0 0 1; 0 0 0].>0, [:a, :b, :c]),
    S2 = UnipartiteNetwork([0 1 1; 0 0 1; 0 0 0].>0, [:a, :b, :c]),
    S3 = UnipartiteNetwork([0 1 0; 0 0 1; 1 0 0].>0, [:a, :b, :c]),
    S4 = UnipartiteNetwork([0 0 1; 0 0 1; 0 0 0].>0, [:a, :b, :c]),
    S5 = UnipartiteNetwork([0 1 1; 0 0 0; 0 0 0].>0, [:a, :b, :c]),
    D1 = UnipartiteNetwork([0 1 1; 0 0 0; 1 1 0].>0, [:a, :b, :c]),
    D2 = UnipartiteNetwork([0 1 1; 0 0 1; 0 1 0].>0, [:a, :b, :c]),
    D3 = UnipartiteNetwork([0 0 1; 0 0 0; 1 1 0].>0, [:a, :b, :c]),
    D4 = UnipartiteNetwork([0 1 0; 0 0 1; 0 1 0].>0, [:a, :b, :c]),
    D5 = UnipartiteNetwork([0 1 0; 0 0 1; 1 1 0].>0, [:a, :b, :c]),
    D6 = UnipartiteNetwork([0 1 1; 1 0 1; 1 1 0].>0, [:a, :b, :c]),
    D7 = UnipartiteNetwork([0 1 1; 1 0 0; 1 1 0].>0, [:a, :b, :c]),
    D8 = UnipartiteNetwork([0 1 1; 1 0 0; 1 0 0].>0, [:a, :b, :c])
    )

  # Return
  return motifs
end

"""
Internal function

Returns all permutations of the adjacency matrix of a motif.
"""
function permute_motif(m::T) where {T<:UnipartiteNetwork}
    perm = Array{Bool,2}[]
    for s in permutations(1:richness(m))
        push!(perm, m.A[s,s])
    end
    return unique(perm)
end

"""
Internal function

Returns all permutations of the adjacency matrix of a motif.
"""
function permute_motif(m::T) where {T<:BipartiteNetwork}
    perm = Array{Bool,2}[]
    for ts in permutations(1:richness(m; dims=1)), bs in permutations(1:richness(m; dims=2))
        push!(perm, m.A[ts,bs])
    end
    return unique(perm)
end

function inner_find_motif(N::T1, m::T2) where {T1<:UnipartiteNetwork,T2<:UnipartiteNetwork}
    motif_permutations = permute_motif(m)
    matching_species = []
    for species_combination in combinations(species(N), richness(m))
        if N[species_combination].A ∈ motif_permutations
            push!(matching_species, (species_combination,))
        end
    end
    return matching_species
end

function inner_find_motif(N::T1, m::T2) where {T1<:BipartiteNetwork,T2<:BipartiteNetwork}
    motif_permutations = permute_motif(m)
    matching_species = []
    top_combinations = combinations(species(N; dims=1), richness(m; dims=1))
    bottom_combinations = combinations(species(N; dims=2), richness(m; dims=2))
    # Pre-allocate the species positions
    top_sp_pos = zeros(Int64, richness(m; dims=1))
    bot_sp_pos = zeros(Int64, richness(m; dims=2))
    # Positions of species (they don't change!)
    p1 = Dict(zip(species(N; dims=1), 1:richness(N; dims=1)))
    p2 = Dict(zip(species(N; dims=2), 1:richness(N; dims=2)))
    for top_species in top_combinations
        for i in 1:length(top_species)
            top_sp_pos[i] = p1[top_species[i]]
        end
        for bottom_species in bottom_combinations
            for i in 1:length(bottom_species)
                bot_sp_pos[i] = p2[bottom_species[i]]
            end
            if N.A[top_sp_pos, bot_sp_pos] ∈ motif_permutations
                push!(matching_species, (top_species, bottom_species))
            end
        end
    end
    return matching_species
end

function inner_find_motif(N::T1, m::T2) where {T1<:UnipartiteProbabilisticNetwork, T2<:UnipartiteNetwork}
    motif_permutations = permute_motif(m)
    all_combinations = []
    for species_combination in combinations(species(N), richness(m))
        isg = N[species_combination]
        motif_mean, motif_var = 0.0, 0.0
        for perm in motif_permutations
            imat = zeros(eltype(N.A), size(m))
            for i in eachindex(imat)
                imat[i] = (perm[i] ? 2.0*isg[i] : 1.0)-isg[i]
            end
            pmm, pmv = prod(imat), EcologicalNetworks.m_var(imat)
            motif_mean += pmm
            motif_var += pmv
        end
        push!(all_combinations, ((species_combination,),(motif_mean, motif_var)))
    end
    return all_combinations
end

function inner_find_motif(N::T1, m::T2) where {T1<:BipartiteProbabilisticNetwork, T2<:BipartiteNetwork}
    motif_permutations = permute_motif(m)
    all_combinations = []
    top_combinations = combinations(species(N; dims=1), richness(m; dims=1))
    bottom_combinations = combinations(species(N; dims=2), richness(m; dims=2))
    for top_species in top_combinations, bottom_species in bottom_combinations
        isg = N[top_species, bottom_species]
        motif_mean, motif_var = 0.0, 0.0
        for perm in motif_permutations
            imat = zeros(eltype(N.A), size(m))
            for i in eachindex(imat)
                imat[i] = (perm[i] ? 2.0*isg[i] : 1.0)-isg[i]
            end
            pmm, pmv = prod(imat), EcologicalNetworks.m_var(imat)
            motif_mean += pmm
            motif_var += pmv
        end
        push!(all_combinations, ((top_species, bottom_species),(motif_mean, motif_var)))
    end
    return all_combinations
end

"""
    find_motif(N::T1, m::T2) where {T1<:AbstractEcologicalNetwork, T2<:BinaryNetwork}

Returns an array of tuples, in which each tuple contains the species that are
part of the motif. The length of the array gives the number of times the motif
was found. For probabilistic networks, the tuple also contains the probability
of observing the species in the correct conformation for the motif, as well as
the variance.
"""
function find_motif(N::T1, m::T2) where {T1<:AbstractEcologicalNetwork, T2<:BinaryNetwork}
    M = copy(N)
    if typeof(N) <: AbstractUnipartiteNetwork
        M = nodiagonal(M)
    end
    if typeof(M) <: QuantitativeNetwork
        M = convert(BinaryNetwork, M)
    end
    return inner_find_motif(M, m)
end

"""
    expected_motif_count(s)

Get the expected number of motifs (and variance) from the output of `find_motif`
on a probabilistic network.
"""
function expected_motif_count(s)
    m = [x[2][1] for x in s]
    v = [x[2][2] for x in s]
    return (sum(m), sum(v))
end

ext: jl


text:
using Base.Meta: parse

#################
# Overload of ~ #
#################

macro VarName(ex::Union{Expr, Symbol})
  # Usage: @VarName x[1,2][1+5][45][3]
  #    return: (:x,[1,2],6,45,3)
  s = string(gensym())
  if isa(ex, Symbol)
    ex_str = string(ex)
    return :(Symbol($ex_str), (), Symbol($s))
  elseif ex.head == :ref
    _2 = ex
    _1 = ""
    while _2.head == :ref
      if length(_2.args) > 2
        _1 = "[" * foldl( (x,y)-> "$x, $y", map(string, _2.args[2:end])) * "], $_1"
      else
        _1 = "[" * string(_2.args[2]) * "], $_1"
      end
      _2   = _2.args[1]
      isa(_2, Symbol) && (_1 = ":($_2)" * ", ($_1), Symbol(\"$s\")"; break)
    end
    return esc(parse(_1))
  else
    error("VarName: Mis-formed variable name $(ex)!")
  end
end

"""
    var_name ~ Distribution()

Tilde notation `~` can be used to specifiy *a variable follows a distributions*.

If `var_name` is an un-defined variable or a container (e.g. Vector or Matrix), this variable will be treated as model parameter; otherwise if `var_name` is defined, this variable will be treated as data.
"""
macro ~(left, right)
  # Is multivariate a subtype of real, e.g. Vector, Matrix?
  if isa(left, Real)                  # value
    # Call observe
    esc(
      quote
        _lp += Turing.observe(
          sampler,
          $(right),   # Distribution
          $(left),    # Data point
          vi
        )
      end
    )
  else
    vsym = getvsym(left)
    vsym_str = string(vsym)
    # Require all data to be stored in data dictionary.
    if vsym in Turing._compiler_[:fargs]
      if ~(vsym in Turing._compiler_[:dvars])
        @info " Observe - `" * vsym_str * "` is an observation"
        push!(Turing._compiler_[:dvars], vsym)
      end
      esc(
        quote
          # Call observe
          _lp += Turing.observe(
            sampler,
            $(right),   # Distribution
            $(left),    # Data point
            vi
          )

        end
      )
    else
      if ~(vsym in Turing._compiler_[:pvars])
        msg = " Assume - `" * vsym_str * "` is a parameter"
        isdefined(Main, Symbol(vsym_str)) && (msg  *= " (ignoring `$(vsym_str)` found in global scope)")
        @info msg
        push!(Turing._compiler_[:pvars], vsym)
      end
      # The if statement is to deterimnet how to pass the prior.
      # It only supports pure symbol and Array(/Dict) now.
      #csym_str = string(gensym())
      if isa(left, Symbol)
        # Symbol
        sym, idcs, csym = @VarName(left)
        csym = Symbol(string(Turing._compiler_[:fname])*string(csym))
        syms = Symbol[csym, left]
        assume_ex = quote
          vn = Turing.VarName(vi, $syms, "")
          if isa($(right), Vector)
            $(left), __lp = Turing.assume(
              sampler,
              $(right),   # dist
              vn,         # VarName
              $(left),
              vi          # VarInfo
            )
            _lp += __lp
          else
            $(left), __lp = Turing.assume(
              sampler,
              $(right),   # dist
              vn,         # VarName
              vi          # VarInfo
            )
            _lp += __lp
          end
        end
      else
        assume_ex = quote
          sym, idcs, csym = @VarName $left
          csym_str = string(Turing._compiler_[:fname])*string(csym)
          indexing = isempty(idcs) ? "" : mapreduce(idx -> string(idx), *, idcs)
          vn = Turing.VarName(vi, Symbol(csym_str), sym, indexing)
          $(left), __lp = Turing.assume(
            sampler,
            $right,   # dist
            vn,       # VarName
            vi        # VarInfo
          )
          _lp += __lp
        end
      end
      esc(assume_ex)
    end
  end
end

#################
# Main Compiler #
#################

"""
    @model(name, fbody)

Wrapper for models.

Usage:

```julia
@model model() = begin
  # body
end
```

Example:

```julia
@model gauss() = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt.(s))
  1.5 ~ Normal(m, sqrt.(s))
  2.0 ~ Normal(m, sqrt.(s))
  return(s, m)
end
```
"""
macro model(fexpr)
   # Compiler design: sample(fname_compiletime(x,y), sampler)
   #   fname_compiletime(x=nothing,y=nothing; data=data,compiler=compiler) = begin
   #      ex = quote
   #          fname_runtime(vi::VarInfo,sampler::Sampler) = begin
   #              x=x,y=y
   #              # pour all variables in data dictionary, e.g.
   #              k = data[:k]
   #              # pour model definition `fbody`, e.g.
   #              x ~ Normal(0,1)
   #              k ~ Normal(x, 1)
   #          end
   #      end
   #      Main.eval(ex)
   #   end

  @debug fexpr
  fexpr = translate(fexpr)

  fname = fexpr.args[1].args[1]      # Get model name f
  fargs = fexpr.args[1].args[2:end]  # Get model parameters (x,y;z=..)
  fbody = fexpr.args[2]              # NOTE: nested args is used here because the orignal model expr is in a block
                                     # NOTE: the code above was `fbody = fexpr.args[2].args[end]`, but since Julia 0.7
                                     #       block doesn't need this nested trick to be fetched
  # Prepare for keyword arguments, e.g.
  #   f(x,y)
  #       ==> f(x,y;)
  #   f(x,y; c=1)
  #       ==> unchanged

  if (length(fargs) == 0 ||         # e.g. f()
          isa(fargs[1], Symbol) ||  # e.g. f(x,y)
          fargs[1].head == :kw)     # e.g. f(x,y=1)
    insert!(fargs, 1, Expr(:parameters))
  end

  @debug fname
  # @debug fargs
  @debug fbody

  # Remove positional arguments from inner function, e.g.
  #  f((x,y; c=1)
  #      ==> f(; c=1)
  #  f(x,y;)
  #      ==> f(;)
  # fargs_inner = deepcopy(fargs)[1:1]

  # Add keyword arguments, e.g.
  #  f(; c=1)
  #      ==> f(; c=1, :vi=VarInfo(), :sample=nothing)
  #  f(;)
  #      ==> f(; :vi=VarInfo(), :sample=nothing)
  # push!(fargs_inner[1].args, Expr(:kw, :vi, :(Turing.VarInfo())))
  # push!(fargs_inner[1].args, Expr(:kw, :sampler, :(nothing)))
  # @debug fargs_inner

  # Modify fbody, so that we always return VarInfo
  fbody_inner = deepcopy(fbody)

  return_ex = fbody.args[end] # get last statement of defined model
  if typeof(return_ex) == Symbol
    pop!(fbody_inner.args)
    # NOTE: code below is commented out to disable explict return
    # vstr = string(return_ex)
    # push!(fbody_inner.args, :(vn = Turing.VarName(:ret, Symbol($vstr*"_ret"), "", 1)))
    # NOTE: code above is commented out to disable explict return
  elseif return_ex.head == :return || return_ex.head == :tuple
    pop!(fbody_inner.args)
    # NOTE: code below is commented out to disable explict return
    # # Turn statement from return to tuple
    # if return_ex.head == :return && typeof(return_ex.args[1]) != Symbol && return_ex.args[1].head == :tuple
    #   return_ex = return_ex.args[1]
    # end
    #
    # # Replace :return or :tuple statement with corresponding operations on vi
    # for v = return_ex.args
    #   @assert typeof(v) == Symbol "Returned variable ($v) name must be a symbol."
    #   push!(fbody_inner.args, :(if sampler != nothing vi.pred[Symbol($(string(v)))] = Turing.realpart($v) end))
    # end
    # NOTE: code above is commented out to disable explict return
  end

  pushfirst!(fbody_inner.args, :(_lp = zero(Real)))
  push!(fbody_inner.args, :(vi.logp = _lp))
  push!(fbody_inner.args, Expr(:return, :vi)) # always return vi in the end of function body

  @debug fbody_inner

  fname_inner_str = "$(fname)_model"
  fname_inner = Symbol(fname_inner_str)
  fdefn_inner = Expr(:(=), fname_inner,
          Expr(:function, Expr(:call, fname_inner))) # fdefn = :( $fname() )
  # push!(fdefn_inner.args[2].args[1].args, fargs_inner...)   # set parameters (x,y;data..)

  push!(fdefn_inner.args[2].args[1].args, :(vi::Turing.VarInfo))
  push!(fdefn_inner.args[2].args[1].args, :(sampler::Union{Nothing,Turing.Sampler}))

  push!(fdefn_inner.args[2].args, deepcopy(fbody_inner))    # set function definition
  @debug fdefn_inner

  fdefn_inner_callback_1 = parse("$fname_inner_str(vi::Turing.VarInfo)=$fname_inner_str(vi,nothing)")
  fdefn_inner_callback_2 = parse("$fname_inner_str(sampler::Turing.Sampler)=$fname_inner_str(Turing.VarInfo(),nothing)")
  fdefn_inner_callback_3 = parse("$fname_inner_str()=$fname_inner_str(Turing.VarInfo(),nothing)")

  compiler = Dict(:fname => fname,
                  :fargs => fargs,
                  :fbody => fbody,
                  :dvars => Set{Symbol}(),  # data
                  :pvars => Set{Symbol}(),  # parameter
                  :fdefn_inner => fdefn_inner,
                  :fdefn_inner_callback_1 => fdefn_inner_callback_1,
                  :fdefn_inner_callback_2 => fdefn_inner_callback_2,
                  :fdefn_inner_callback_3 => fdefn_inner_callback_3)

  # Outer function defintion 1: f(x,y) ==> f(x,y;data=Dict())
  fargs_outer = deepcopy(fargs)
  # Add data argument to outer function
  push!(fargs_outer[1].args, Expr(:kw, :data, :(Dict{Symbol,Any}())))
  # Add data argument to outer function
  push!(fargs_outer[1].args, Expr(:kw, :compiler, compiler))
  for i = 2:length(fargs_outer)
    s = fargs_outer[i]
    if isa(s, Symbol)
      fargs_outer[i] = Expr(:kw, s, :nothing) # turn f(x;..) into f(x=nothing;..)
    end
  end

  fdefn_outer = Expr(:function, Expr(:call, fname, fargs_outer...),
                        Expr(:block, Expr(:return, fname_inner)))

  pushfirst!(fdefn_outer.args[2].args, :(Main.eval(fdefn_inner_callback_3)))
  pushfirst!(fdefn_outer.args[2].args, :(Main.eval(fdefn_inner_callback_2)))
  pushfirst!(fdefn_outer.args[2].args, :(Main.eval(fdefn_inner_callback_1)))
  pushfirst!(fdefn_outer.args[2].args, :(Main.eval(fdefn_inner)))
  pushfirst!(fdefn_outer.args[2].args, quote
      # Check fargs, data
      Turing.eval(:(_compiler_ = deepcopy($compiler)))
      fargs = Turing._compiler_[:fargs];

      # Copy the expr of function definition and callbacks
      fdefn_inner = Turing._compiler_[:fdefn_inner];
      fdefn_inner_callback_1 = Turing._compiler_[:fdefn_inner_callback_1];
      fdefn_inner_callback_2 = Turing._compiler_[:fdefn_inner_callback_2];
      fdefn_inner_callback_3 = Turing._compiler_[:fdefn_inner_callback_3];

      # Add gensym to function name
      fname_inner_with_gensym = gensym((fdefn_inner.args[2].args[1].args[1]));

      # Change the name of inner function definition to the one with gensym()
      fdefn_inner.args[2].args[1].args[1] = fname_inner_with_gensym
      fdefn_inner_callback_1.args[1].args[1] = fname_inner_with_gensym
      fdefn_inner_callback_1.args[2].args[2].args[1] = fname_inner_with_gensym
      fdefn_inner_callback_2.args[1].args[1] = fname_inner_with_gensym
      fdefn_inner_callback_2.args[2].args[2].args[1] = fname_inner_with_gensym
      fdefn_inner_callback_3.args[1].args[1] = fname_inner_with_gensym
      fdefn_inner_callback_3.args[2].args[2].args[1] = fname_inner_with_gensym

      # Copy data dictionary
      for k in keys(data)
        if fdefn_inner.args[2].args[2].args[1].head == :line
          # Preserve comments, useful for debuggers to
          # correctly locate source code oringin.
          insert!(fdefn_inner.args[2].args[2].args, 2, Expr(:(=), Symbol(k), data[k]))
        else
          insert!(fdefn_inner.args[2].args[2].args, 1, Expr(:(=), Symbol(k), data[k]))
        end
      end
      # @debug fdefn_inner
  end )

  for k in fargs
    if isa(k, Symbol)       # f(x,..)
      _k = k
    elseif k.head == :kw    # f(z=1,..)
      _k = k.args[1]
    else
      _k = nothing
    end
    if _k != nothing
      _k_str = string(_k)
      @debug _k_str, " = ", _k
      data_check_ex = quote
            if haskey(data, keytype(data)($_k_str))
              if nothing != $_k
                Turing.dwarn(0, " parameter "*$_k_str*" found twice, value in data dictionary will be used.")
              end
            else
              data[keytype(data)($_k_str)] = $_k
              data[keytype(data)($_k_str)] == nothing && @error("Data `"*$_k_str*"` is not provided.")
            end
          end
      pushfirst!(fdefn_outer.args[2].args, data_check_ex)
    end
  end
  pushfirst!(fdefn_outer.args[2].args, quote data = copy(data) end)

  @debug esc(fdefn_outer)
  esc(fdefn_outer)
end

###################
# Helper function #
###################

function insdelim(c, deli=",")
  reduce((e, res) -> append!(e, [res, deli]), c; init = [])[1:end-1]
end

getvsym(s::Symbol) = s
getvsym(expr::Expr) = begin
  @assert expr.head == :ref "expr needs to be an indexing expression, e.g. :(x[1])"
  curr = expr
  while isa(curr, Expr) && curr.head == :ref
    curr = curr.args[1]
  end
  curr
end

translate!(ex::Any) = ex
translate!(ex::Expr) = begin
  if (ex.head === :call && ex.args[1] === :(~))
    ex.head = :macrocall; ex.args[1] = Symbol("@~")
    insert!(ex.args, 2, LineNumberNode(-1)) # NOTE: a `LineNumberNode` object is required
                                            #       at the second args for macro call in 0.7
  else
    map(translate!, ex.args)
  end
  ex
end
translate(ex::Expr) = translate!(deepcopy(ex))


text:
# SECTION 2: fertility experiments plotting
# ---------------------------------------
# 2.1 Fertility experiments and UE-BMI
   # 2.1.1 a temp macro for fast plotting
   macro fastexpr( Ylab::String, Loc::String )
      return :(
         grid(true); xlabel("Year"); ylabel($Ylab);
         legend(( "Benchmark", "High Fertility", "Low Fertility" ), loc = $Loc );
      )
   end
   # 2.1.2 do plotting
   # NOTE: it is a 4-subfigure figure
   figure(figsize = (13,8))
      subplot(2,2,1)  # gap / pool expenditure
         plot( DatPkg_base.Dt[:Year][idx_plot], 100 .* ( DatPkg_base.Dt[:LI] ./ DatPkg_base.Dt[:AggPoolExp] )[idx_plot]  );
         plot( DatPkg_high.Dt[:Year][idx_plot], 100 .* ( DatPkg_high.Dt[:LI] ./ DatPkg_high.Dt[:AggPoolExp] )[idx_plot], "-.b" );
         plot( DatPkg_low.Dt[:Year][idx_plot], 100 .* ( DatPkg_low.Dt[:LI] ./ DatPkg_low.Dt[:AggPoolExp] )[idx_plot], "--r" );
         @fastexpr( "Pool Gap / Pool Expenditures (%)", "upper left" )
      subplot(2,2,2)  # gap / pool income
         plot( DatPkg_base.Dt[:Year][idx_plot], 100 .* ( DatPkg_base.Dt[:LI] ./ DatPkg_base.Dt[:AggPoolIn] )[idx_plot]  );
         plot( DatPkg_high.Dt[:Year][idx_plot], 100 .* ( DatPkg_high.Dt[:LI] ./ DatPkg_high.Dt[:AggPoolIn] )[idx_plot], "-.b" );
         plot( DatPkg_low.Dt[:Year][idx_plot], 100 .* ( DatPkg_low.Dt[:LI] ./ DatPkg_low.Dt[:AggPoolIn] )[idx_plot], "--r" );
         @fastexpr( "Pool Gap / Pool Incomes (%)", "upper left" )
      subplot(2,2,3)  # gap / GDP
         plot( DatPkg_base.Dt[:Year][idx_plot], 100 .* ( DatPkg_base.Dt[:LI] ./ DatPkg_base.Dt[:Y] )[idx_plot]  );
         plot( DatPkg_high.Dt[:Year][idx_plot], 100 .* ( DatPkg_high.Dt[:LI] ./ DatPkg_high.Dt[:Y] )[idx_plot], "-.b" );
         plot( DatPkg_low.Dt[:Year][idx_plot], 100 .* ( DatPkg_low.Dt[:LI] ./ DatPkg_low.Dt[:Y] )[idx_plot], "--r" );
         @fastexpr( "Pool Gap / GDP (%)", "upper left" )
      subplot(2,2,4)  # gap / tax revenues
         plot( DatPkg_base.Dt[:Year][idx_plot], 100 .* ( DatPkg_base.Dt[:LI] ./ (DatPkg_base.Dt[:TRw] .+ DatPkg_base.Dt[:TRc]) )[idx_plot]  );
         plot( DatPkg_high.Dt[:Year][idx_plot], 100 .* ( DatPkg_high.Dt[:LI] ./ (DatPkg_high.Dt[:TRw] .+ DatPkg_high.Dt[:TRc]) )[idx_plot], "-.b" );
         plot( DatPkg_low.Dt[:Year][idx_plot], 100 .* ( DatPkg_low.Dt[:LI] ./ (DatPkg_low.Dt[:TRw] .+ DatPkg_low.Dt[:TRc]) )[idx_plot], "--r" );
         @fastexpr( "Pool Gap / Tax Revenues (%)", "upper left" )
      tight_layout()
   # 2.1.3 save figure
   savefig( "./output/FertilityShockProfile.pdf", format = "pdf" )

   # 2.2 Summary table for benchmark & fertility experiments
      # 2.2.1 a function to summarize a DatPkg_, where x is the index to slice the transition path
      afunc_SummaryDatPkg( d::NamedTuple, year::Int ) = begin
         local v::Vector{Any} = []
         # convert year -> index
         x = year - env.START_YEAR + 1

         # 0. population size
         push!(v, d.Dt[:N][x] )
         # 1. gdp per capita
         push!(v, d.Dt[:Y][x] / d.Dt[:N][x] )
         # 2. capital per capita
         push!(v, d.Dt[:K][x] / d.Dt[:N][x] )
         # 3. consumption per capita
         push!(v, d.Dt[:C][x] / d.Dt[:N][x] )
         # 4. health expenditure per capita
         push!(v, sum(d.Dst[:m][x,:] .* d.Ps[:N][x,:]) / d.Dt[:N][x] )
         # 5. outpatient expenditure per capita
         push!(v, sum(d.Dst[:MA][x,:] .* d.Ps[:N][x,:]) / d.Dt[:N][x] )
         # 6. inpatient expenditure per capita
         push!(v, sum(d.Dst[:MB][x,:] .* d.Ps[:N][x,:]) / d.Dt[:N][x] )
         # 7. investment per capita
         push!(v, d.Dt[:I][x] / d.Dt[:N][x] )
         # 8. labor supply per capita (endowment = 1)
         push!(v, d.Dt[:L][x] / d.Dt[:N][x] )
         # 9. average wage level (wage rate)
         push!(v, d.Dt[:w̄][x] )
         # 10. interest rate (net investment returns) (%)
         push!(v, 100 * d.Dt[:r][x] )
         # 11. income tax revenues
         push!(v, d.Dt[:TRw][x] )
         # 12. consumption tax revenues
         push!(v, d.Dt[:TRc][x] )
         # 13. gov purchase / GDP (%)
         push!(v, 100 * d.Dt[:G][x] / d.Dt[:Y][x] )
         # 14. UEBMI pooling gap / GDP (%)
         push!(v, 100 * d.Dt[:LI][x] / d.Dt[:Y][x] )

         return v::Vector{Any}
      end # func
      # 2.2.2 prepare a column of variable names
      tmpvarlist = [
         "Population size", "Output per capita", "Capital per capita",
         "Consumption per capita", "Health expenditure per capita", "Outpatient expenditure per capita",
         "Inpatient expenditure per capita", "Investment per capita",
         "Labor supply per capita (endowment=1)", "Average wage rate",
         "Interest rate (%)", "Income tax revenues", "Consumption tax revenues",
         "Government purchase / Output (%)", "UE-BMI's pooling gap / Output (%)"
      ]
      # 2.2.3 summary benchmark in different years
      df_SummaryFertility = DataFrames.DataFrame(
         :Variables => tmpvarlist,
         :Year2010 => afunc_SummaryDatPkg( DatPkg_base, 2010 ),
         :Year2020 => afunc_SummaryDatPkg( DatPkg_base, 2020 ),
         :Year2050 => afunc_SummaryDatPkg( DatPkg_base, 2050 ),
         :Year2110 => afunc_SummaryDatPkg( DatPkg_base, 2110 ),
         :FinalSS  => afunc_SummaryDatPkg( DatPkg_base, 2344 ),
      )
      # 2.2.4 output to a csv file
      CSV.write( "$(pwd())/output/SummaryBench.csv" , df_SummaryFertility )

#


text:
struct TrimConditions{T, AC<:Aircraft, AT<:Atmosphere, G<:Gravity}
    tas::T  # m/s
    ψ::T  # rad
    position::SVector{3, T}  # [x, y, z] earth (m)
    ψ_dot::T  # rad/s
    γ::T  # rad
    aircraft::AC
    atmosphere::AT
    gravity::G
end

struct TrimSolution{T, N}
    α::T
    β::T
    controls::SVector{N, T}
    θ::T
    ϕ::T
    p::T
    q::T
    r::T
end

function trim(
    dss_guess, controls_guess, aircraft, atmosphere, gravity, γ=0.0, ψ_dot=0.0;
    show_trace=false,
    ftol=1e-16,
    iterations=5000
)

    # TRIMMING SOLUTION
    sol_gues = [
        get_α(dss_guess),  # alpha (rad)
        get_β(dss_guess),  # beta (rad)
        controls_guess..., # thtl  (0-1), el (deg), ail (deg), rdr (deg)
    ]

    # CONSTS
    trim_conditions = TrimConditions(
        get_tas(dss_guess),  # TAS (m/s)
        get_euler_angles(dss_guess)[1],  # psi (rad)
        SVector{3}(get_earth_position(dss_guess)),  # north, east, down (m)
        ψ_dot,  # ψ_dot (rad/s)
        γ,  # γ (rad)
        aircraft,
        atmosphere,
        gravity,
    )

    ds_type = typeof(dss_guess)

    f_opt(sol) = trim_cost_function(ds_type, sol, trim_conditions; full_output=false)

    result = nlsolve(
        f_opt, sol_gues;
        ftol=ftol, show_trace=show_trace, iterations=iterations
    )

     if show_trace
            println(result)
     end

    sol = result.zero
    dssd, controls, outputs, cost = trim_cost_function(
        ds_type, sol, trim_conditions; full_output=true
    )

    controls = ConstantInput.(controls)

    return dssd, controls, outputs, cost
end

function trim_cost_function(ds::Type{T}, sol, trim_conditions; full_output=false) where {T<:DSState}

    tas = trim_conditions.tas
    ψ = trim_conditions.ψ
    x, y, z = trim_conditions.position
    ψ_dot = trim_conditions.ψ_dot
    γ = trim_conditions.γ
    aircraft = trim_conditions.aircraft
    atmosphere = trim_conditions.atmosphere
    gravity = trim_conditions.gravity

    gd = get_gravity_accel(gravity)

    ϕ, θ, p, q, r = apply_trimmer_constrains(
        tas,  # m/s
        sol[1],  # α (rad)
        sol[2],  # β (rad)
        γ,  # rad
        ψ_dot,  # rad/s
        gd,  # m/s²
    )

    trim_solution = TrimSolution(
        sol[1],  # α (rad)
        sol[2],  # β (rad)
        SVector{4}(sol[3:6]),  # controls  -> thtl  (0-1), el (deg), ail (deg), rdr (deg)
        θ,  # rad
        ϕ,  # rad
        p,  # rad/s
        q,  # rad/s
        r,  # rad/s
    )
    # Construct state vector
    dss = ds(trim_conditions, trim_solution, aircraft)
    dssd, outputs = f(time, dss, trim_solution.controls, aircraft, atmosphere, gravity)

    cost = get_trimmer_cost(dssd)

    if full_output
        return dssd, trim_solution.controls, outputs, cost
    else
        return cost
    end
end

function apply_trimmer_constrains(tas, α, β, γ, ψ_dot, gd)
    # Coordinated turn bank --> phi
    ϕ = coordinated_turn_bank(ψ_dot, α, β, tas, γ, gd)
    # Climb -> theta
    θ = rate_of_climb_constrain_no_wind(γ, α, β, ϕ)
    # Angular kinemtic -> p, q, r
    p, q, r = ψθϕ_dot_2_pqr(ψ_dot, 0, 0, θ, ϕ)

    return ϕ, θ, p, q, r
end

function SixDOFAeroEuler{T}(tc::TrimConditions, ts::TrimSolution, ac::Aircraft) where {T}
    x = [
        tc.tas,
        ts.α,
        ts.β,
        ts.ϕ,
        ts.θ,
        tc.ψ,
        ts.p,
        ts.q,
        ts.r,
        tc.position...,
        tgear(ac, ts.controls[1])
    ]
    return SixDOFAeroEuler(x)
end

function SixDOFBodyEuler{T}(tc::TrimConditions, ts::TrimSolution, ac::Aircraft) where {T}
    x = [
        wind2body(tc.tas, 0, 0, ts.α, ts.β)...,
        ts.ϕ,
        ts.θ,
        tc.ψ,
        ts.p,
        ts.q,
        ts.r,
        tc.position...,
        tgear(ac, ts.controls[1])
    ]
    return SixDOFBodyEuler(x)
end


text:
"""
    SteppedLayeredModel <: SeisModel1D

A `SteppedLayeredModel` contains `n` layers each with constant properties.
"""
struct SteppedLayeredModel <: SeisModel1D
    "Earth radius in km"
    a :: Float64
    "Number of layers"
    n :: Int
    "Radii in km of top of each layer"
    r :: Vector{Float64}
    "Physical parameters"
    vp :: Vector{Float64}
    vs :: Vector{Float64}
    density :: Vector{Float64}
    "Optional anisotropic parameters"
    aniso :: Bool
    vph :: Vector{Float64}
    vpv :: Vector{Float64}
    vsh :: Vector{Float64}
    vsv :: Vector{Float64}
    eta :: Vector{Float64}
    "Optional attenuation parameters"
    attenuation :: Bool
    Qμ :: Vector{Float64}
    Qκ :: Vector{Float64}
end

"""
    SteppedLayeredModel(; kwargs...) -> m

Construct a `SteppedLayeredModel` by specifying keyword arguments.  The length of
arrays defines the number of layers and the presence of arguments determines
whether the model has density, anisotropy and attenuation.

Note that radii are given in km, velocities in km/s, densities in g/cm³
and frequency in Hz.

### Keyword arguments
#### Radii
- `r`: The upper radius of a set of constant layers in km.  The uppermost
  layer radius is taken to be the planetary radius.  This is the only
  mandatory argument.
#### Isotropic properties
These must be given each as a `Vector` of values for each layer.
Each parameter (`vp`, `density`, `Qμ`, etc.) must have the correct number
of layers.
- `vp`: P-wave velocity
- `vs`: S-wave velocity
- `density`: Material density
#### Anisotropic parameters
If the model is anisotropic and the radial anisotropy parameters are all
provided, then the Voigt average velocities are constructed automatically
unless `vp` and/or `vs` are also given separately.
- `vpv`: Vertical P-wave velocity
- `vph`: Horizontal P-wave velocity
- `vsv`: Vertical S-wave velocity
- `vsh`: Horizontally-polarised S-wave velocity
- `eta`: The ratio C₁₃/(Vph - 2Vsv), where cᵢⱼ is the Voigt elasticity
  matrix, the 1-direction is horizontal and the 3-direction is radial
#### Attenuation parameters
ASCII versions (`Qmu`, `Qkappa`) of these exist for ease of use as well as
the Unicode versions (`Qμ`, `Qκ`).  If both are supplied, an error is thrown.
- `Qμ` or `Qmu`: Shear modulus quality factor
- `Qκ` or `Qkappa`: Bulk modulus quality factor
"""
function SteppedLayeredModel(; r, vp=nothing, vs=nothing, density=nothing,
                             vph=nothing, vpv=nothing, vsh=nothing, vsv=nothing,
                             eta=nothing, Qμ=nothing, Qmu=nothing,
                             Qκ=nothing, Qkappa=nothing, fref=nothing)
    n = length(r)
    all(x->x>0, diff(r)) || throw(ArgumentError("radii must increase monotonically"))
    R = last(r)
    aniso = attenuation = false
    # Anisotropy
    if all(!isnothing, (vph, vpv, vsh, vsv, eta))
        aniso = true
        # Set isotropic velocities automatically if not provided
        if any(isnothing, (vp, vs))
            voigt_vels = voigt_velocities.(vpv, vsv, vph, vsh, eta)
            vp === nothing && (vp = first.(voigt_vels))
            vs === nothing && (vs = last.(voigt_vels))
        end
    end
    # Attenuation
    # Default to ASCII versions if given
    Qmu !== nothing && (Qμ = Qmu)
    Qkappa !== nothing && (Qκ = Qkappa)
    if all(!isnothing, (Qμ, Qκ))
        attenuation = true
    end
    # Convert to correct shape, and check for length
    vp = _set_coeffs_size(SteppedLayeredModel, n, vp)
    vs = _set_coeffs_size(SteppedLayeredModel, n, vs)
    density = _set_coeffs_size(SteppedLayeredModel, n, density)
    vph = _set_coeffs_size(SteppedLayeredModel, n, vph)
    vsh = _set_coeffs_size(SteppedLayeredModel, n, vsh)
    vpv = _set_coeffs_size(SteppedLayeredModel, n, vpv)
    vsv = _set_coeffs_size(SteppedLayeredModel, n, vsv)
    eta = _set_coeffs_size(SteppedLayeredModel, n, eta)
    Qμ = _set_coeffs_size(SteppedLayeredModel, n, Qμ)
    Qκ = _set_coeffs_size(SteppedLayeredModel, n, Qκ)
    SteppedLayeredModel(R, n, r, vp, vs, density, aniso, vph, vpv,
        vsh, vsv, eta, attenuation, Qμ, Qκ)
end

_set_coeffs_size(::Type{SteppedLayeredModel}, n, v::Nothing) = []
function _set_coeffs_size(::Type{SteppedLayeredModel}, n, v::Vector)
    length(v) == n || throw(ArgumentError("number of value must be $n"))
    v
end

for sym in model_variables_SeisModel1D
    @eval function ($sym)(m::SteppedLayeredModel, r::Real; depth::Bool=false)
        length(m.$sym) > 0 ||
            throw(ArgumentError("'$split(string($sym), ".")[end]' not defined for model"))
        depth && (r = radius(m, r))
        m.$(sym)[findlayer(m, r)]
    end
end

function evaluate(m::SteppedLayeredModel, field::Symbol, r; depth::Bool=false)
    length(getfield(m, field)) > 0 ||
        throw(ArgumentError("$field' not defined for model"))
    depth && (r = radius(m, r))
    getfield(m, field)[findlayer(m, r)]
end

function mass(m::SteppedLayeredModel, r; depth::Bool=false)
    hasdensity(m) || throw(ArgumentError("model does not contain density"))
    depth && (r = radius(m, r))
    l = findlayer(m, r)
    r *= 1.e3
    M = 0.
    for i in 1:l
        R0 = i == 1 ? 0. : m.r[i-1]*1.e3
        R = i < l ? m.r[i]*1.e3 : r
        M += 1.e3*m.density[i]*(R^3 - R0^3)
    end
    4/3*π*M
end


text:
@testset "sobolindices" begin
    x = RandomVariable.(Uniform(-π, π), [:x1, :x2, :x3])

    ishigami = Model(
        df -> sin.(df.x1) + 7 .* sin.(df.x2) .^ 2 + 0.1 .* df.x3 .^ 4 .* sin.(df.x1), :f
    )

    n = 1000

    firstorder_analytical = [0.3138, 0.4424, 0.00]
    totaleffect_analytical = [0.5574, 0.4424, 0.2436]

    @testset "Monte Carlo" begin
        Random.seed!(8128)

        si = sobolindices([ishigami], x, :f, MonteCarlo(n))

        Random.seed!()

        @test all(isapprox(si.FirstOrder, firstorder_analytical; rtol=0.1))
        @test all(isapprox(si.TotalEffect, totaleffect_analytical; rtol=0.1))
    end

    @testset "Sobol" begin
        Random.seed!(8128)

        si = sobolindices([ishigami], x, :f, SobolSampling(n))

        Random.seed!()

        @test all(isapprox(si.FirstOrder, firstorder_analytical; rtol=0.1))
        @test all(isapprox(si.TotalEffect, totaleffect_analytical; rtol=0.1))
    end

    @testset "Halton" begin
        Random.seed!(8128)

        si = sobolindices([ishigami], x, :f, MonteCarlo(n))

        Random.seed!()

        @test all(isapprox(si.FirstOrder, firstorder_analytical; rtol=0.1))
        @test all(isapprox(si.TotalEffect, totaleffect_analytical; rtol=0.1))
    end

    @testset "Latin Hypercube" begin
        Random.seed!(8128)

        si = sobolindices([ishigami], x, :f, MonteCarlo(n))

        Random.seed!()

        @test all(isapprox(si.FirstOrder, firstorder_analytical; rtol=0.1))
        @test all(isapprox(si.TotalEffect, totaleffect_analytical; rtol=0.1))
    end
end

max_stars_repo_name: FriesischScott/julia-uq

max_issues_repo_name: FriesischScott/julia-uq

max_forks_repo_name: FriesischScott/julia-uq


text:
using OrdinaryDiffEq
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_2Dlinear
prob = deepcopy(prob_ode_2Dlinear)
prob.tspan = (0.0,100.0)
## Solve and plot
sol =solve(prob,DP5();dt=1//2^(22),qmax=1.000005,progress=true,
          maxiters=Int(1e7),progress_steps=Int(1e3),progressbar_name="Test")

sol =solve(prob,DP5();dt=1//2^(22),qmax=1.000005,progress=false,
          maxiters=Int(1e7),progressbar_name="Test")


text:
module Handlers

function val(v)
    @static if VERSION < v"0.7.0-DEV.1395"
        Val{v}()
    else
        Val(v)
    end
end

export handle, Handler, HandlerFunction, Router, register!

using HTTP

"""
handle(handler::Handler, request, response) => Response

Function used to dispatch to a Handler. Called from the core HTTP.serve method with the
initial Handler passed to `HTTP.serve(handler=handler)`.
"""
function handle end
handle(handler, req, resp, vals...) = handle(handler, req, resp)

"""
Abstract type representing an object that knows how to "handle" a server request.

Types of handlers include `HandlerFunction` (a julia function of the form `f(request, response`) and
`Router` (which pattern matches request url paths to other specific `Handler` types).
"""
abstract type Handler end

"""
HandlerFunction(f::Function)

A Function-wrapper type that is a subtype of `Handler`. Takes a single Function as an argument.
The provided argument should be of the form `f(request, response) => Response`, i.e. it accepts
both a `Request` and `Response` and returns a `Response`.
"""
struct HandlerFunction{F <: Function} <: Handler
    func::F # func(req, resp)
end

handle(h::HandlerFunction, req, resp) = h.func(req, resp)

"A default 404 Handler"
const FourOhFour = HandlerFunction((req, resp) -> Response(404))

"""
Router(h::Handler)
Router(f::Function)
Router()

An `HTTP.Handler` type that supports mapping request url paths to other `HTTP.Handler` types.
Can accept a default `Handler` or `Function` that will be used in case no other handlers match; by
default, a 404 response handler is used.
Paths can be mapped to a handler via `HTTP.register!(r::Router, path, handler)`, see `?HTTP.register!` for more details.
"""
struct Router <: Handler
    segments::Dict{String, Val}
    sym::Symbol
    func::Function
    function Router(ff::Union{Handler, Function, Void}=nothing)
        sym = gensym()
        if ff == nothing
            f = @eval $sym(args...) = FourOhFour
        else
            f = ff isa Function ? HandlerFunction(ff) : ff
        end
        r = new(Dict{String, Val}(), sym, f)
        return r
    end
end

const SCHEMES = Dict{String, Val}("http" => val(:http), "https" => val(:https))
const METHODS = Dict{String, Val}()
for m in instances(HTTP.Method)
    METHODS[string(m)] = val(Symbol(m))
end
const EMPTYVAL = val(())

"""
HTTP.register!(r::Router, url, handler)
HTTP.register!(r::Router, m::Union{Method, String}, url, handler)

Function to map request urls matching `url` and an optional method `m` to another `handler::HTTP.Handler`.
URLs are registered one at a time, and multiple urls can map to the same handler.
Methods can be passed as a string `"GET"` or enum object directly `HTTP.GET`.
The URL can be passed as a String or `HTTP.URI` object directly. Requests can be routed based on: method, scheme,
hostname, or path.
The following examples show how various urls will direct how a request is routed by a server:

- `"http://*"`: match all HTTP requests, regardless of path
- `"https://*"`: match all HTTPS requests, regardless of path
- `"google"`: regardless of scheme, match requests to the hostname "google"
- `"google/gmail"`: match requests to hostname "google", and path starting with "gmail"
- `"/gmail"`: regardless of scheme or host, match any request with a path starting with "gmail"
- `"/gmail/userId/*/inbox`: match any request matching the path pattern, "*" is used as a wildcard that matches any value between the two "/"
"""
register!(r::Router, url, handler) = register!(r, "", url, handler)
register!(r::Router, m::Method, url, handler) = register!(r, string(m), url, handler)

function register!(r::Router, method::String, url, handler)
    m = isempty(method) ? Any : typeof(METHODS[method])
    # get scheme, host, split path into strings & vals
    uri = url isa String ? HTTP.URI(url) : url
    s = HTTP.scheme(uri)
    sch = HTTP.hasscheme(uri) ? typeof(get!(SCHEMES, s, val(s))) : Any
    h = HTTP.hashostname(uri) ? Val{Symbol(HTTP.hostname(uri))} : Any
    hand = handler isa Function ? HandleFunction(handler) : handler
    register!(r, m, sch, h, HTTP.path(uri), hand)
end

function splitsegments(r::Router, h::Handler, segments)
    vals = Expr[]
    for s in segments
        if s == "*" #TODO: or variable, keep track of variable types and store in handler
            T = Any
        else
            v = val(Symbol(s))
            r.segments[s] = v
            T = typeof(v)
        end
        push!(vals, Expr(:(::), T))
    end
    return vals
end

function register!(r::Router, method::DataType, scheme, host, path, handler)
    # save string => Val mappings in r.segments
    segments = map(String, split(path, '/'; keep=false))
    vals = splitsegments(r, handler, segments)
    # return a method to get dispatched to
    #TODO: detect whether defining this method will create ambiguity?
    @eval $(r.sym)(::$method, ::$scheme, ::$host, $(vals...), args...) = $handler
    return
end

function handle(r::Router, req, resp)
    # get the url/path of the request
    m = val(HTTP.method(req))
    uri = HTTP.uri(req)
    # get scheme, host, split path into strings and get Vals
    s = get(SCHEMES, HTTP.scheme(uri), EMPTYVAL)
    h = val(Symbol(HTTP.hostname(uri)))
    p = HTTP.path(uri)
    segments = split(p, '/'; keep=false)
    # dispatch to the most specific handler, given the path
    vals = (get(r.segments, s, EMPTYVAL) for s in segments)
    handler = r.func(m, s, h, vals...)
    # pass the request & response to the handler and return
    return handle(handler, req, resp, vals...)
end

end # module


text:
using Distributed
using TensorOperations, KrylovKit, LinearAlgebra
using Combinatorics
using Base.Threads
using MPOmodule
using finiteMPS
using JLD

############################################################################
#
#	TDVP functions
#
############################################################################

function simple_taylor(Heff, x::AbstractArray{<:Any}, dτ::Number, order::Int)
	res = x
	tmp = x
	C = 1.0
	for i in 1:order
		tmp = Heff(tmp)
		C/= i
		res .+= C*tmp*(dτ^i)
	end

	return res

end

function simple_taylor_Heff!(MPO, Lenv, Renv, x::AbstractArray{<:Any}, dτ::Number, order::Int)
	res = x
	tmp = x
	C = 1.0
	for i in 1:order
		tmp = Heff(tmp)
		C/= i
		res .+= C*tmp*(dτ^i)
	end

	return res

end

function evo_sweep_TDVP(MPSvec::Vector{<:Any}, Cpre::AbstractArray{<:Number}, MPOvec::Vector{MPO},
			RBlocks::Vector{<:Any}, dτ::Union{Float64,ComplexF64})

	order = 3
	taylor_t = []
	env_t = []

	#
	#	foward sweep
	#
	Lenv = [Array{ComplexF64,2}(I,1,1)]
	LBlocks = [Lenv]
	for site in 1:length(MPSvec)-1

		#
		# evolve A_site forward in time
		#
		localMPO = MPOvec[site]
		Renv = RBlocks[end-(site-1)]
		Lenv = LBlocks[site]
		Heff(x) = applyHeff(x, localMPO, Lenv, Renv)

		AR_site = MPSvec[site]
		@tensor MPSvec[site][α, d, β] := Cpre[α, γ]*AR_site[γ, d, β]

		ACnew = simple_taylor(Heff, MPSvec[site], -0.5*dτ, order)
		#ACnew, info = exponentiate(x->Heff(x), -0.5*dτ, MPSvec[site]; ishermitian = false, tol = 1e-14)

		#
		# split A_new_site and left can it
		#
		ALnew, S, Vdagger = leftCanSite(ACnew; optSVD = true)
		S = S./sign(S[1,1])
		S = S./(sqrt(sum(S.^2))) #normalize S
		C = S*Vdagger
		MPSvec[site] = ALnew

		#
		# build up Lenv
		#
		Lenv = applyTM_MPO([MPSvec[site]], [localMPO], LBlocks[end]; left=true)
		push!(LBlocks, Lenv) # (int,1,....,N-1)

		#
		# evolve C back in time
		#
		Hceff(x) = applyHCeff(x, Lenv, Renv)

		Cpre = simple_taylor(Hceff, C, 0.5*dτ, order)
		#Cpre, info = exponentiate(x->Hceff(x), 0.5*dτ, C; ishermitian = false, tol = 1e-14)

	end
	#
	#	evolve site N
	#
	Lenv=LBlocks[end]
	Renv=RBlocks[1]
	localMPO = MPOvec[end]
	Heff(x) = applyHeff(x, localMPO, Lenv, Renv)
	AR_site = MPSvec[end]
	@tensor v_0[α, d, β] := Cpre[α, γ]*AR_site[γ, d, β]

	ACnew = simple_taylor(Heff, v_0, -1.0*dτ, order)

	MPSvec[end] = ACnew

	RBlocks = [Renv]

	#
	#	evolve back
	#
	Cvec = []
	for site in length(MPSvec):-1:2

		AR, S, U = rightCanSite(MPSvec[site];optSVD=true)
		MPSvec[site] = AR
		S = S./sign(S[1,1])
		snorm = sqrt(sum(S.^2))
		S = S./snorm
		C = U*S
		push!(Cvec, C)

		#
		# add RBlock
		#
		Renv = applyTM_MPO([MPSvec[site]],[MPOvec[site]], RBlocks[end]; left = false)
		push!(RBlocks, Renv)#[int,N,.....,2]

		#
		#	evolve C backwards
		#
		Lenv = LBlocks[site]
		Hceff(x) = applyHCeff(x, Lenv, Renv)
		#Cnext, info = exponentiate(Hceff, 0.5*dτ, C; ishermitian = false, tol = 10e-16)
		Cnext = simple_taylor(Hceff, C, 0.5*dτ, order)

		#
		#	evolve AC forward
		#
		@tensor v_0[α, d, β] := MPSvec[site-1][α, d, γ]*Cnext[γ, β]
		Lenv = LBlocks[site-1]
		Renv = RBlocks[end]
		localMPO = MPOvec[site-1]
		Heff(x) = applyHeff(x, localMPO, Lenv, Renv)

		#ACnew, info = exponentiate(Heff, -0.5*dτ, v_0; ishermitian = false, tol=10e-16, verbosity = 0, maxiter = 20)
		ACnew = simple_taylor(Heff, v_0, -0.5*dτ, order)

		MPSvec[site-1] = ACnew

	end

	AR, S, U = rightCanSite(MPSvec[1]; optSVD = true)
	MPSvec[1] = AR

	S = S./sign(S[1,1])
	snorm = sqrt(sum(S.^2))
	S = S./snorm
	Cpre = U*S

	return Cpre, RBlocks, Cvec[end:-1:1]
end

function get1RDM(MPSvec::Vector{<:Any})
	γ = zeros(ComplexF64, length(MPSvec), length(MPSvec))
	for i in 1:length(MPSvec)
	 	for j in i:length(MPSvec)

	 		mpo = CMPO(i,j, length(MPSvec))
			c = getExpValue(MPSvec, mpo)
			γ[i,j] = c

			γ[j,i] = conj(c)

		end
	end

	return γ
end

############################################################################
#
#	MPS_SEE struct
#
############################################################################

struct MPS_SSE
	init_MPS::Vector{Array{ComplexF64,3}}
	H_MPO::Vector{MPO}
	Heff::Vector{MPO}
	Decay_MPO::Vector{Vector{MPO}}
	Jump_MPO::Vector{Vector{MPO}}
	n_rel::Float64 # expected number of jumps per time step
	evoT::Tuple{Int, Int}
	maxDim::Int
	measureDict::Dict
end

function saveMPS_SSE(name, A::MPS_SSE)

	save(name, "MPS_SSE", A)

end

function loadMPS_SSE(name)

	return load(name)["MPS_SSE"]

end

function get_dτ(A::MPS_SSE, Decay_sum::Vector{MPO})
	Γ_act = opnorm(creatOp_from_MPO(Decay_sum)[:,1,1,:])   # how do I change this for large systems
	dτ = A.n_rel/Γ_act
	return dτ
end

function get_Heff(A::MPS_SSE; depara_opt = false)
	H_MPO = A.H_MPO
	Decay_MPO = A.Decay_MPO

	Heff = H_MPO
	for decay_mpo in Decay_MPO

		Heff = Heff + (-0.5im*decay_mpo)
	end

	if depara_opt == true
		# deparallize mpo
		# add further MPO comperssion methods
		println("not implemented")
	end

	return Heff
end

function get_Decay_MPO(A::MPS_SSE)
	Decay_Sum_MPO = 0
	for (i,decay) in enumerate(A.Decay_MPO)
		Decay_Sum_MPO = i == 1 ? decay : Decay_Sum_MPO + decay

	end

	return Decay_Sum_MPO
end

function get_measurements(MPSvec::Vector{<:Any}, C::Vector{<:Any}, T::Float64, measureDict::Dict)
	N = size(MPSvec)[1]

	# ensure normalization
	leftCanMPS(MPSvec)
	Cpre, __, Cvec = rightCanMPS(MPSvec)
	measure_results = []
	push!(measure_results, T) # time
	Cvec = Cvec[end:-1:1]
	push!(Cvec, Cpre)
	Cvec = Cvec[end:-1:1]

	for x in measureDict
		#@show "measure: $(x[1])"

		if typeof(x[2]) == Vector{MPO}

			Renv = [Array{ComplexF64,2}(I,1,1)]
			measure_x= applyTM_MPO(MPSvec[1:end], x[2][1:end], Renv; left = false)[1,1][1]
			#@show x[1]
			#@show measure_x

		# eventuall add routine for single site measurement
		else
			measure_x = singleSiteExpValue(MPSvec, Cvec[x[2][1]], x[2][2],x[2][1]; can_left=false)
		end

		push!(measure_results, measure_x)
	end

	#entanglement
	for i in 0:Int(round(N/2, RoundDown))-1

		#@show Cvec[Int(round(N/2, RoundDown))-i]
		λ = diag(Cvec[Int(round(N/2, RoundDown))-i])
		#push!(measure_results, diag(Cvec[Int(round(N/2, RoundDown))]))
		#@show -2.0*sum((λ.^2).*log.(λ))
		push!(measure_results , -2.0*sum((λ.^2).*log.(λ)))

	end

	return measure_results
end

function evo_TDVP_SSE(MPS_SSE::MPS_SSE; MPOvec::Union{Nothing, Vector{MPO}} = nothing, dτ::Union{Number,Nothing} = nothing, name = nothing, saveState = false)

	#
	#	initialize result txt
	#
	if name != nothing
		dir_name = name[1]
		id = name[2]
		name_measurement_files = string(dir_name,"/res_id$(id).txt")

		test =[x[1] for  x in MPS_SSE.measureDict]
		measurements = Array{Any,2}(undef, 1, length(test))
		measurements[1,:] = test[:]
		writedlm(name_measurement_files, measurements)
	end

	# initlize run
	MPSvec = MPS_SSE.init_MPS
	MPOvec = MPOvec == nothing ? get_Heff(MPS_SSE) : MPOvec
	maxDim = MPS_SSE.maxDim

	MPOvec_decay = deparaMPOChain!(get_Decay_MPO(MPS_SSE))
	dτ = dτ == nothing ?  get_dτ(MPS_SSE, MPOvec_decay) : dτ
	(Tinit, Tmax) = MPS_SSE.evoT

	#right can. initial state
	Cpre, R, Cvec = rightCanMPS(MPSvec)

	#
	# build up right environment blocks
	#
	Renv = [Array{ComplexF64,2}(I,1,1)]
	RBlocks = [Renv]

	for site = 1:length(MPSvec)-1
		Renv = applyTM_MPO([MPSvec[end-(site-1)]],[MPOvec[end-(site-1)]], RBlocks[site]; left = false)
		push!(RBlocks, Renv)#[init,N,.....,2]
	end

	@show Tmax
	T = Tinit
	n_jumps = 0
	dτ = dτ
	t_measure = 20*dτ

	measure_list = []
	#@showprogress for iter in 0:dτ:(Tmax+dτ) #while T < Tmax

	runtime = @elapsed while T < Tmax

		#@show T
		T += dτ; t_measure += dτ

		#
		#	TDVP - evo
		#
		preMPS = deepcopy(MPSvec)
		Cpre, RBlocks, Cvec = evo_sweep_TDVP(MPSvec, Cpre, MPOvec, RBlocks, im*dτ)
		#
		#	stochastic jump
		#
		number_channels = length(MPS_SSE.Jump_MPO)

		#
		#	get decay prob.
		#
		#@time begin
		R = [Array{ComplexF64,2}(I,1,1)]
		decay_prob = applyTM_MPO(MPSvec, MPOvec_decay, R;
					 left = false)[1,1][1]*dτ

		if abs(decay_prob) > 1 @error "decay prob > 1" end

		# get random number
		rand_jump = rand(1)[1]

		# perform jump
		if rand_jump < abs(decay_prob)

			# rand_jump is uniformly distributed between
			# [0,decay_prob) we can use it to select channel
			rand_channel = rand_jump
			sum_d_channel = 0
			nop = 1
			Jump_MPO = nothing

			for n_channel in 1:number_channels
				D_MPO = MPS_SSE.Decay_MPO[n_channel]

				R = [Array{ComplexF64,2}(I,1,1)]
				new_decay_contribution = applyTM_MPO(MPSvec, D_MPO, R; left = false)[1,1][1]
				sum_d_channel += new_decay_contribution*dτ
				#@show sum_d_channel, rand_channel, nop
				# check in which channel rand_channel lies
				if abs(sum_d_channel) > rand_channel

 					Jump_MPO =  MPS_SSE.Jump_MPO[nop]
					#println("*********jump occured at channel $(nop) ****************")
					break
				else
					nop += 1
				end

				#@show sum_d_channel, rand_channel
				@assert nop <= size(MPS_SSE.Decay_MPO)[1] "prob. didn't add up"
			end
			#
			#	variational application
			#
			#measure_results = get_measurements(deepcopy(MPSvec), pushfirst!(Cvec,Cpre), T, MPS_SSE.measureDict)

			#MPSvec = iter_applyMPO(MPSvec, Jump_MPO, maxDim; Niter = 1)
			MPSvec = iter_applyMPO(preMPS, Jump_MPO, maxDim; Niter = 1)

			#
			#	variational truncation
			#
			#MPSvec = iter_trunc(MPSvec, 140; Niter = 1)

			leftCanMPS(MPSvec)
			Cpre, R, Cvec = rightCanMPS(MPSvec)
			n_jumps += 1
			#@show n_jumps

			#
			#	svd trunc
			#

			#=
			Renv = [Array{ComplexF64,2}(I,1,1)]
			@show applyTM_MPO(newMPSvec, measureDict["Np"], Renv; left = false)

			# show particle number in Schmidt vector
			s = mixedCanMPS(newMPSvec, Int(round(N/2, RoundDown)))
			#println("schmidt coefficients: $(s)")

			Renv = [Array{ComplexF64,2}(I,1,1)]
			np_half = [measureDict["Np"][1] ,measureDict["Np"][Int(round(N/2, RoundDown)+2):end]...]
			Rblock_total_N = applyTM_MPO(newMPSvec[Int(round(N/2, RoundDown)+1):end], np_half, Renv; left = false)
			#Rblock_total_N = sum(Rblock_total_N)
			N_λ_right = [Rblock_total_N[1][i,i] for i in 1:size(Rblock_total_N[1])[1]]
			#println("particle number right: $(N_λ_right)")

			Renv = [Array{ComplexF64,2}(I,1,1)]
			np_half = [measureDict["Np"][1:Int(round(N/2, RoundDown)-1)]..., measureDict["Np"][end]]
			Lblock_total_N = applyTM_MPO(newMPSvec[1:Int(round(N/2, RoundDown))], np_half, Renv; left = true)
			#Lblock_total_N = sum(Lblock_total_N)
			N_λ_left = [Lblock_total_N[1][i,i] for i in 1:size(Lblock_total_N[1])[1]]
			#println("particle number left: $(N_λ_left)")

			Np_s = Dict()
			for (i,x) in enumerate(s)
				Np_s[x] = (N_λ_left[i], N_λ_right[i])
			end

			for x in sort!([x for x in Np_s], by = x->x[1])
				println(x)
			end

			@tensor res[a,b,c] := newMPSvec[Int(round(N/2, RoundDown))][a,b, i]*diagm(0=>s)[i,c]
			newMPSvec[Int(round(N/2, RoundDown))] = res
			leftCanMPS(newMPSvec)
			rightCanMPS(newMPSvec)
			=#

			#
			#	you have to build a new right block
			#	bcs. MPS has changed from jump
			#
			Renv = [Array{ComplexF64,2}(I,1,1)]
			RBlocks = [Renv]
			for site = 1:length(MPSvec)-1
				Renv = applyTM_MPO([MPSvec[end-(site-1)]],[MPOvec[end-(site-1)]], RBlocks[site]; left = false)
				push!(RBlocks, Renv)#[init,N,.....,2]
			end

		end

		#
		# measure stuff MPS comes in right can. form
		#

		if t_measure >= 20*dτ
			@show T
			T = Float64(T)
			measure_results = get_measurements(deepcopy(MPSvec), pushfirst!(Cvec,Cpre), T, MPS_SSE.measureDict)
			t_measure = 0.0
			push!(measure_results, n_jumps)
			push!(measure_results, abs(decay_prob))
			push!(measure_list, measure_results)

			#
			#	save intermidiate state
			#
			if name != nothing
				dir_name = name[1]
				id = name[2]
				io = open(name_measurement_files, "a")
				write_res = Array{Any,2}(undef, 1, length(measure_results))
				write_res[1,:] = measure_results[:]
				writedlm(io, write_res)
				close(io)
				if saveState == true
					name_inter_state = string(dir_name,"/inter_states/inter_state_id$(id).jld")
					name_1RDM = string(dir_name,"/inter_states/1RDM/1RDM_id$(id)_T$(T).jld")
					C = get1RDM(MPSvec)
					save(name_1RDM, "1RDM", C, "T", T)
					save(name_inter_state, "time", T, "MPSvec", MPSvec)
				end
			end
		end

	end

	rightCanMPS(MPSvec)
	name_inter_state = string(dir_name,"/inter_states/inter_state_id$(id).jld")
	save(name_inter_state, "time", T, "MPSvec", MPSvec)
	return runtime, measure_list#, getCoef(MPSvec)
end


text:
# Adapted from the Yggdrasil build script
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

Pkg.setprotocol!(domain = "github.com", protocol = "ssh")

function getversion(headerfile)
    a = b = c = 0
    for l in readlines(headerfile)
        if startswith(l, "#define JLCXX_VERSION_MAJOR")
            a = parse(Int, split(l)[end])
        end
        if startswith(l, "#define JLCXX_VERSION_MINOR")
            b = parse(Int, split(l)[end])
        end
        if startswith(l, "#define JLCXX_VERSION_PATCH")
            c = parse(Int, split(l)[end])
            break
        end
    end
    return VersionNumber(a,b,c)
end

basepath = dirname(@__DIR__)

name = "libcxxwrap_julia"
version = getversion(joinpath(basepath, "include", "jlcxx", "jlcxx_config.hpp"))

julia_versions = [v"1.6.0", v"1.7.0", v"1.8.0"]

# Collection of sources required to complete build
sources = [
    DirectorySource(basepath; target="libcxxwrap-julia")
]

# Bash recipe for building across all platforms
script = raw"""
mkdir build
cd build

cmake \
    -DJulia_PREFIX=$prefix \
    -DCMAKE_INSTALL_PREFIX=$prefix \
    -DCMAKE_FIND_ROOT_PATH=$prefix \
    -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} \
    -DCMAKE_BUILD_TYPE=Release \
    ../libcxxwrap-julia
VERBOSE=ON cmake --build . --config Release --target install -- -j${nproc}
install_license $WORKSPACE/srcdir/libcxxwrap-julia*/LICENSE.md
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line

function libjulia_platforms(julia_version)
  platforms = supported_platforms(; experimental=julia_version ≥ v"1.7")

  filter!(p -> libc(p) != "musl" && arch(p) != "i686"  && !contains(arch(p), "arm") && arch(p) != "aarch64" && arch(p) != "powerpc64le" && !Sys.isfreebsd(p), platforms)

  for p in platforms
      p["julia_version"] = string(julia_version)
  end

  platforms = expand_cxxstring_abis(platforms)

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

  return platforms
end

platforms = vcat(libjulia_platforms.(julia_versions)...)

# The products that we will ensure are always built
products = [
    LibraryProduct("libcxxwrap_julia", :libcxxwrap_julia; dlopen_flags=[:RTLD_GLOBAL]),
    LibraryProduct("libcxxwrap_julia_stl", :libcxxwrap_julia_stl; dlopen_flags=[:RTLD_GLOBAL]),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    BuildDependency("libjulia_jll"),
]

deployingargs = deepcopy(ARGS)
deploytarget = ENV["GITHUB_REF_NAME"] == "main" ? "barche/libcxxwrap_julia_jll.jl" : "local"
push!(deployingargs, "--deploy=$deploytarget")

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(deployingargs, name, version, sources, script, platforms, products, dependencies;
    preferred_gcc_version = v"9", julia_compat = "1.6")

if ENV["GITHUB_REF_NAME"] == "main"
    Pkg.Registry.add(RegistrySpec(url = "https://github.com/barche/CxxWrapTestRegistry.git"))
    Pkg.develop(PackageSpec(url="https://github.com/barche/libcxxwrap_julia_jll.jl.git"))
    import libcxxwrap_julia_jll
    using LocalRegistry
    register(libcxxwrap_julia_jll)
end

max_stars_repo_name: eloop/libcxxwrap-julia

max_issues_repo_name: eloop/libcxxwrap-julia

max_forks_repo_name: eloop/libcxxwrap-julia


text:
function ccf_total(order_list_timeseries::AbstractChunkListTimeseries, line_list_df::DataFrame, pipeline::PipelinePlan; recalc::Bool = false,
                  output_fn_suffix::String = "",
                  range_no_mask_change::Real=RvSpectMLBase.max_bc, ccf_mid_velocity::Real=0.0,
                  v_step::Real=250.0, v_max=RvSpectMLBase.max_bc, mask_scale_factor::Real=1,
                  ccf_var_scale::Real=1.0, mask_type::Symbol = :tophat, Δfwhm::AbstractVector{T} = zeros(0),
                  calc_ccf_var::Bool = false, calc_ccf_covar::Bool = false, allow_nans::Bool = false, verbose::Bool = false ) where { T<:Real }
    if need_to(pipeline,:ccf_total) || recalc
      if verbose println("# Computing CCF.")  end
      @assert !need_to(pipeline,:extract_orders)
      @assert !need_to(pipeline,:clean_line_list_tellurics)
      if mask_type == :tophat
          mask_shape = CCFs.TopHatCCFMask(order_list_timeseries.inst, scale_factor=mask_scale_factor)
      elseif mask_type == :gaussian
          mask_shape = CCFs.GaussianCCFMask(order_list_timeseries.inst, σ_scale_factor=mask_scale_factor)
      elseif mask_type == :supergaussian
            mask_shape = CCFs.SuperGaussianCCFMask(order_list_timeseries.inst, σ_scale_factor=mask_scale_factor)
      elseif mask_type == :halfcos
          mask_shape = CCFs.CosCCFMask(order_list_timeseries.inst, scale_factor=mask_scale_factor)
      else
        @error("Requested mask shape (" * string(mask_type) * " not avaliable.")
      end

      line_list = hasproperty(line_list_df, :order) ? CCFs.BasicLineList2D(line_list_df.lambda, line_list_df.weight, line_list_df.order) :
                  CCFs.BasicLineList(line_list_df.lambda, line_list_df.weight )
      ccf_plan = CCFs.BasicCCFPlan(mask_shape = mask_shape, line_list=line_list, midpoint=ccf_mid_velocity, range_no_mask_change=range_no_mask_change, step=v_step, max=v_max, allow_nans=allow_nans)
      v_grid = calc_ccf_v_grid(ccf_plan)

      ccfs = calc_ccf_chunklist_timeseries(order_list_timeseries, ccf_plan, ccf_var_scale=ccf_var_scale, calc_ccf_var=calc_ccf_var, calc_ccf_covar=calc_ccf_covar, Δfwhm=Δfwhm, verbose=verbose)
      if save_data(pipeline, :ccf_total)
         CSV.write(joinpath(output_dir,target_subdir * "_ccfs" * output_fn_suffix * ".csv"),Tables.table(ccfs',header=Symbol.(v_grid)))
         #CSV.write(joinpath(output_dir,target_subdir * "_ccfs_expr.csv"),Tables.table(ccfs_expr',header=Symbol.(v_grid)))
      end
      set_cache!(pipeline, :ccf_total, (ccfs=ccfs, v_grid=v_grid) )
      dont_need_to!(pipeline,:ccf_total)
    end

    if has_cache(pipeline,:ccf_total) return read_cache(pipeline,:ccf_total)
    else   @error("Invalid pipeline state.")          end
end

#= Tried to make this for one chunk_list instead of a timeseries, but didn't get it working
function ccf_total(order_list::AbstractChunkList, line_list_df::DataFrame, inst::AbstractInstrument; recalc::Bool = false,
                  output_fn_suffix::String = "", range_no_mask_change::Real=RvSpectMLBase.max_bc, ccf_mid_velocity::Real=0.0, v_step::Real=250.0,
                  mask_scale_factor::Real=1, mask_type::Symbol = :tophat,
                  use_old::Bool = false, calc_ccf_var::Bool = false,  verbose::Bool = false )
    #if need_to(pipeline,:ccf_total) || recalc
      if verbose println("# Computing CCF.")  end
      #@assert !need_to(pipeline,:extract_orders)
      #@assert !need_to(pipeline,:clean_line_list_tellurics)
      if mask_type == :tophat
          mask_shape = CCFs.TopHatCCFMask(inst, scale_factor=mask_scale_factor)
      elseif mask_type == :gaussian
          mask_shape = CCFs.GaussianCCFMask(inst, σ_scale_factor=mask_scale_factor)
      elseif mask_type == :supergaussian
            mask_shape = CCFs.SuperGaussianCCFMask(inst, σ_scale_factor=mask_scale_factor)
      elseif mask_type == :halfcos
          mask_shape = CCFs.CosCCFMask(inst, scale_factor=mask_scale_factor)
      else
        @error("Requested mask shape (" * string(mask_type) * " not avaliable.")
      end

      line_list = CCFs.BasicLineList(line_list_df.lambda, line_list_df.weight)
      ccf_plan = CCFs.BasicCCFPlan(mask_shape = mask_shape, line_list=line_list, midpoint=ccf_mid_velocity, range_no_mask_change=range_no_mask_change, step=v_step, max=RvSpectMLBase.max_bc)
      v_grid = calc_ccf_v_grid(ccf_plan)
      @time ccfs = calc_ccf_chunklist(order_list, ccf_plan)
      #=
      if use_old
        @time ccf = calc_ccf_chunklist(order_list, ccf_plan)
      else
        @time ccf = calc_ccf_chunklist(order_list, ccf_plan, calc_ccf_var=calc_ccf_var)
      end
      =#
      #=
      if save_data(pipeline, :ccf_total)
         CSV.write(joinpath(output_dir,target_subdir * "_ccfs" * output_fn_suffix * ".csv"),Tables.table(ccfs',header=Symbol.(v_grid)))
         #CSV.write(joinpath(output_dir,target_subdir * "_ccfs_expr.csv"),Tables.table(ccfs_expr',header=Symbol.(v_grid)))
      end
      set_cache!(pipeline, :ccf_total, (ccfs=ccfs, v_grid=v_grid) )
      dont_need_to!(pipeline,:ccf_total)
      =#
    #end # if pipeline_plan

    #if has_cache(pipeline,:ccf_total) return read_cache(pipeline,:ccf_total)
    #else   @error("Invalid pipeline state.")          end
    return ccf
end
=#


text:
module BlobTracking
using Statistics, LinearAlgebra
using Images, ImageFiltering, ImageDraw, VideoIO
using LowLevelParticleFilters, Hungarian, StaticArrays, Distributions, Distances, Interact, NearestNeighbors
using JuliennedArrays # For faster median etc.

export BlobTracker, Blob, Recorder, track_blobs, showblobs, tune_sizes, FrameBuffer, MedianBackground, DiffBackground, background, foreground, update!, TrackingResult, Measurement, location, threshold, invthreshold, OOB, lifetime, Trace, trace, tracem, allblobs, draw!, to_static, get_coordinates

export AbstractCorrespondence, HungarianCorrespondence, NearestNeighborCorrespondence, MCCorrespondence

"""
    Workspace{T1, T2}

Contains buffer arrays that can be resued to minimize allocations
"""
struct Workspace{T1,T2}
    storage::T1
    blob_storage::T2
end

"""
    Workspace(img::AbstractMatrix, n::Int)

Provide example image and length of `bt.sizes`
"""
function Workspace(img::AbstractMatrix, n::Int)
    storage = Gray.(img)
    blob_storage = Array{Float64}(undef, size(img)..., n)
    Workspace(storage,blob_storage)
end

struct Measurement
    coordinates
    assi
end

const OOB = CartesianIndex(0,0)
@inline to_static(a::Number) = a
@inline to_static(a::AbstractMatrix) = SMatrix{size(a)...}(a)
@inline to_static(a::AbstractVector) = SVector{length(a)}(a)
@inline to_static(a::CartesianIndex{2}) = SVector{2}(a.I)

include("correspondence_types.jl")
include("framebuffer.jl")
include("background.jl")
include("blob.jl")
include("trackingresult.jl")
include("correspondence.jl")
include("display.jl")

Workspace(img::AbstractMatrix, bt::BlobTracker) = Workspace(img, length(bt.sizes))

function Base.iterate(vid::VideoIO.VideoReader)
    state = read(vid)
    (Gray.(state)),state
end
function Base.iterate(vid::VideoIO.VideoReader, state)
    eof(vid) && return nothing
    read!(vid,state)
    (Gray.(state)),state
end

function Base.filter!(result::TrackingResult, bt::BlobTracker, m::Measurement, step=1)
    for (bi, ass) in enumerate(m.assi)
        blob = result.blobs[bi]
        if ass == 0 || too_far(bt.correspondence, blob,m.coordinates[ass]) # penalize not found
            blob.counter += step
            m.assi[bi] = 0
        else
            blob.counter = max(0, blob.counter-step) # decrement counter if measurement found
        end
    end
    m
end

function Base.filter!(result::TrackingResult, bt::BlobTracker, measurements::Vector{Measurement})
    N = length(measurements)
    map(measurements) do m
        filter!(result, bt, m, 1/N)
    end
end

"""
    predict!(blobs)

Advance all blobs by predicting their next state
"""
function LowLevelParticleFilters.predict!(blobs)
    foreach(blob->predict!(blob.kf,0), blobs)
end

"""
    correct!(blobs, measurement::Measurement)

Correct the state of the blobs by incorporating the measurement in the Kalman filter
"""
function LowLevelParticleFilters.correct!(blobs::Vector, measurement::Measurement)
    for (bi, ass) in enumerate(measurement.assi)
        if ass != 0
            ll = correct!(blobs[bi].kf,SVector(measurement.coordinates[ass].I))
            push!(blobs[bi].tracem, measurement.coordinates[ass])
        else
            push!(blobs[bi].tracem, OOB)
        end
        push!(blobs[bi].trace, location(blobs[bi]))
    end
end

LowLevelParticleFilters.predict!(result::TrackingResult) = predict!(result.blobs)
LowLevelParticleFilters.correct!(result::TrackingResult, measurement) = correct!(result.blobs, measurement)

function LowLevelParticleFilters.correct!(blobs::Vector, measurements::Vector{Measurement})
    N = length(measurements)
    R2 = copy(blobs[1].kf.R2)
    for measurement in measurements
        for (bi, ass) in enumerate(measurement.assi)
            blob = blobs[bi]
            if ass != 0
                blob.kf.R2 .*= N # QUESTION N or N^2?
                ll = correct!(blob.kf,SVector(measurement.coordinates[ass].I))
                blob.kf.R2 .= R2
                push!(blob.tracem, measurement.coordinates[ass])
            else
                push!(blob.tracem, OOB)
            end
            push!(blob.trace, location(blob))
        end
    end
 end

"""
    update!(ws::Workspace, bt::BlobTracker, coords_or_img, result::TrackingResult)

Perform one iteration of predict and correct

#Arguments:
- `ws`: a Workspace object
- `bt`: the blob tracker
- `coords_or_img`: vector of coordinates or an image
- `result`: a `TrackingResult`
"""
function LowLevelParticleFilters.update!(ws, bt::BlobTracker, coords_or_img, result::TrackingResult)
    blobs = result.blobs
    measurement = Measurement(ws, bt, coords_or_img, result)
    predict!(result)
    filter!(result, bt, measurement)
    correct!(result, measurement)
    spawn_blobs!(result, bt, measurement)
    kill_blobs!(result, bt)
    measurement
end

function spawn_blobs!(result::TrackingResult, bt::BlobTracker, measurement)
    newcoordinds = setdiff(1:length(measurement.coordinates), measurement.assi)
    newblobs = Blob.(bt.params, measurement.coordinates[newcoordinds])
    append!(result.blobs, newblobs)
end

function spawn_blobs!(result::TrackingResult, bt::BlobTracker, measurements::Vector{Measurement})
    for measurement in measurements
        spawn_blobs!(result, bt, measurement)
    end
end

"""
    kill_blobs!(result::TrackingResult, bt::BlobTracker)

Kill all blobs that have not seen a measurement for the duration `bt.kill_counter_th`
"""
function kill_blobs!(result::TrackingResult, bt::BlobTracker)
    blobs,dead = result.blobs, result.dead
    bi = 1
    while bi <= length(blobs)
        if blobs[bi].counter > bt.kill_counter_th
            push!(dead, blobs[bi])
            deleteat!(blobs, bi)
        else
            bi += 1
        end
    end
end

"""
    coordinates = measure(ws::Workspace, bt::BlobTracker, img)

Detect blobs in the image and return the coordinates

#Arguments:
- `ws`: Workspace
- `bt`: BlobTracker
- `img`: image
"""
function measure(ws::Workspace, bt::BlobTracker, img)
    prepare_image!(ws,bt,img)
    coordinates = detect_blobs!(ws, bt, img)
    bt.mask === nothing || (coordinates = filter!(c->bt.mask[c] != 0, coordinates))
    coordinates
end

function Measurement(ws, bt::BlobTracker, img::AbstractMatrix, result)
    coordinates = measure(ws, bt, img)
    if isempty(result.blobs) && bt.correspondence isa MCCorrespondence
        measurement = assign(bt.correspondence.inner, result.blobs, coordinates)
    else
        measurement = assign(bt.correspondence, result.blobs, coordinates)
    end
    log_measurement(result, measurement)
    measurement
end

function Measurement(_, bt::BlobTracker, coordinates::Trace, result)
    if isempty(result.blobs) && bt.correspondence isa MCCorrespondence
        measurement = assign(bt.correspondence.inner, result.blobs, coordinates)
    else
        measurement = assign(bt.correspondence, result.blobs, coordinates)
    end
    log_measurement(result, measurement)
    measurement
end

"""
    track_blobs(bt::BlobTracker, vid; display=false, recorder=nothing, threads=Threads.nthreads() > 1, ignoreempty=false)

Main entry point to tracking blobs

#Arguments:
- `bt`: a BlobTracker
- `vid`: Some iterable thing that iterates images
- `display = Base.display`: function to display images live. Displaying live slows things down a bit. Use `display=nothing` to not display anything. Consider also `c = imshow(img);
displayfun = img -> imshow!(c["gui"]["canvas"],img); `.
- `recorder`: an optional `Recorder` that can record each frame to a video on disk. Recording things does not slow things down much and also does not affect memory usage much.
- `threads`: Use threaded processing of frames? Only useful if Julia is started with multiple threads.
- `ignoreempty=false`: wether or not to ignore display and recording of frames which contains no blobs and no measurements.
"""
function track_blobs(bt::BlobTracker, vid; display=nothing, recorder=nothing, threads=Threads.nthreads()>1, ignoreempty=false)
    result = TrackingResult()
    buffer = threads ? coordinate_iterator(bt, vid) : vid
    img,buffer = Iterators.peel(buffer)
    ws = Workspace(img, length(bt.sizes))
    img, coord_or_img = img isa Tuple ? img : (img,img)
    measurement = Measurement(ws, bt, coord_or_img, result)
    spawn_blobs!(result, bt, measurement)
    showblobs(RGB.(Gray.(img)), result, measurement, recorder = recorder, display=display)

    try
        for (ind,img) in enumerate(buffer)
            println("Frame $ind")
            img, coord_or_img = img isa Tuple ? img : (img,img)
            measurement = update!(ws, bt, coord_or_img , result)
            showblobs(RGB.(Gray.(img)),result,measurement, rad=6, recorder=recorder, display=display, ignoreempty=ignoreempty)
        end
    finally
        finalize(recorder)
    end
    result
end

function track_blobs(bt::BlobTracker, coords::Vector{Trace})
    result = TrackingResult()
    measurement = Measurement(nothing, bt, coords[1], result)
    spawn_blobs!(result, bt, measurement)
    for coord in coords[2:end]
        measurement = update!(nothing, bt, coord , result)
    end
    result
end

function get_coordinates(bt::BlobTracker, vid; threads=Threads.nthreads()>1)
    coords = Trace[]
    if threads
        for (img, coord) in coordinate_iterator(bt, vid)
            push!(coords, coord)
        end
    else
        img,vid = Iterators.peel(vid)
        ws = Workspace(copy(img), length(bt.sizes))
        coord = measure(ws, bt, img)
        push!(coords, coord)
        for img in vid
            coord = measure(ws, bt, img)
            push!(coords, coord)
        end
    end
    coords
end

function coordinate_iterator(bt, vid)
    img1 = iterate(vid)
    img1 === nothing && return
    img1,state = img1

    ws1 = Workspace(copy(img1), length(bt.sizes))
    ws2 = Workspace(copy(img1), length(bt.sizes))
    m1 = measure(ws1, bt, img1)

    Channel{Tuple{typeof(img1), Trace}}(3, spawn=false) do ch
        put!(ch,(img1, m1))
        while true
            img1 = iterate(vid, state)
            img1 === nothing && return
            img1,state = img1
            m1 = Threads.@spawn measure(ws1, bt, img1)
            img2 = iterate(vid, state)
            if img2 !== nothing
                img2,state = img2
                m2 = Threads.@spawn measure(ws2, bt, img2)
            end
            put!(ch,(img1, fetch(m1)))
            if img2 === nothing
                return
            else
                put!(ch,(img2, fetch(m2)))
            end
        end
    end
end

end # module


text:
using ReduxImGui
using ReduxImGui.Redux
using ReduxImGui.CImGui

function naive_color_edits(store::AbstractStore, get_state=Redux.get_state)
    if ReduxImGui.ColorEdit(store, s->get_state(s).color_edits["basic_color_edit_no_alpha"])
        @info "This triggers $(@__FILE__):$(@__LINE__)."
    end
    CImGui.SameLine()
    ShowHelpMarker("Click on the colored square to open a color picker.\nClick and hold to use drag and drop.\nRight-click on the colored square to show options.\nCTRL+click on individual component to input value.\n")
    if ReduxImGui.ColorEdit(store, s->get_state(s).color_edits["basic_color_edit_with_alpha"])
        @info "This triggers $(@__FILE__):$(@__LINE__)."
    end
end


text:
module ISICArchive

import HTTP
import HTTP: get
import Base: convert, show

using JSON
using Parameters
using Images
using ImageMagick
using Dates

export
    DatasetListRequest,
    DatasetMetadataRequest,

    ImageListRequest,
    ImageMetadataRequest,
    ImageDownloadRequest,
    ImageThumbnailRequest

include("common.jl")
include("dataset.jl")
include("image.jl")

end # module


text:
using MLFlowClient
using Documenter

DocMeta.setdocmeta!(MLFlowClient, :DocTestSetup, :(using MLFlowClient); recursive=true)

makedocs(;
    modules=[MLFlowClient],
    authors="@deyandyankov and contributors",
    repo="https://github.com/JuliaAI.jl/blob/{commit}{path}#{line}",
    sitename="MLFlowClient.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://juliaai.github.io/MLFlowClient.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
        "Tutorial" => "tutorial.md",
        "Reference" => "reference.md"
    ],
)

deploydocs(;
    repo="github.com/JuliaAI/MLFlowClient.jl",
    devbranch="main",
)

max_stars_repo_name: JuliaAI/MLFlowClient

max_issues_repo_name: JuliaAI/MLFlowClient

max_forks_repo_name: JuliaAI/MLFlowClient


text:
"""
gradient(
    θ::AbstractVector{<:Real},
    vi::VarInfo,
    model::Function,
    sampler::Union{Nothing, Sampler}=nothing,
)

Computes the gradient of the log joint of `θ` for the model specified by
`(vi, sampler, model)` using whichever automatic differentation tool is currently active.
"""
function gradient(
    θ::AbstractVector{<:Real},
    vi::VarInfo,
    model::Function,
    sampler::Union{Nothing, Sampler}=nothing,
    backend::Symbol=ADBACKEND[],
)
    @assert backend ∈ (:forward_diff, :reverse_diff)
    if backend == :forward_diff
        return gradient_forward(θ, vi, model, sampler)
    else
        return gradient_reverse(θ, vi, model, sampler)
    end
end

"""
gradient_forward(
    θ::AbstractVector{<:Real},
    vi::VarInfo,
    model::Function,
    spl::Union{Nothing, Sampler}=nothing,
    chunk_size::Int=CHUNKSIZE[],
)

Computes the gradient of the log joint of `θ` for the model specified by `(vi, spl, model)`
using forwards-mode AD from ForwardDiff.jl.
"""
function gradient_forward(
    θ::AbstractVector{<:Real},
    vi::VarInfo,
    model::Function,
    sampler::Union{Nothing, Sampler}=nothing,
    chunk_size::Int=CHUNKSIZE[],
)
    # Record old parameters.
    vals_old, logp_old = copy(vi.vals), copy(vi.logp)

    # Define function to compute log joint.
    function f(θ)
        vi[sampler] = θ
        return -runmodel!(model, vi, sampler).logp
    end

    # Set chunk size and do ForwardMode.
    chunk = ForwardDiff.Chunk(min(length(θ), chunk_size))
    config = ForwardDiff.GradientConfig(f, θ, chunk)
    ∂l∂θ = ForwardDiff.gradient!(similar(θ), f, θ, config)
    l = vi.logp.value

    # Replace old parameters to ensure this function doesn't mutate `vi`.
    vi.vals, vi.logp = vals_old, logp_old

    # Strip tracking info from θ to avoid mutating it.
    θ .= ForwardDiff.value.(θ)

    return l, ∂l∂θ
end

"""
gradient_reverse(
    θ::AbstractVector{<:Real},
    vi::VarInfo,
    model::Function,
    sampler::Union{Nothing, Sampler}=nothing,
)

Computes the gradient of the log joint of `θ` for the model specified by
`(vi, sampler, model)` using reverse-mode AD from Flux.jl.
"""
function gradient_reverse(
    θ::AbstractVector{<:Real},
    vi::Turing.VarInfo,
    model::Function,
    sampler::Union{Nothing, Sampler}=nothing,
)
    vals_old, logp_old = copy(vi.vals), copy(vi.logp)

    # Specify objective function.
    function f(θ)
        vi[sampler] = θ
        return -runmodel!(model, vi, sampler).logp
    end

    # Compute forward and reverse passes.
    l_tracked, ȳ = Tracker.forward(f, θ)
    l, ∂l∂θ = Tracker.data(l_tracked), Tracker.data(ȳ(1)[1])

    # Remove tracking info from variables in model (because mutable state).
    vi.vals, vi.logp = vals_old, logp_old

    # Strip tracking info from θ to avoid mutating it.
    θ .= Tracker.data.(θ)

    # Return non-tracked gradient value
    return l, ∂l∂θ
end

import Base: <=
<=(a::Tracker.TrackedReal, b::Tracker.TrackedReal) = a.data <= b.data

function verifygrad(grad::AbstractVector{<:Real})
    if any(isnan.(grad)) || any(isinf.(grad))
        @warn("Numerical error has been found in gradients.")
        @warn("grad = $(grad)")
        return false
    else
        return true
    end
end


text:
using Revise
using FwiFlow
using PyCall
using LinearAlgebra
using DelimitedFiles
using PyPlot
using MAT
matplotlib.use("agg")
np = pyimport("numpy")

# mode
# 0 -- generate data
# 1 -- running inverse modeling
mode = 1
# mode
# 0 -- small data
# 1 -- large data
datamode = 0
sparsity = 0.01

# mode = parse(Int64, ARGS[1])
# datamode = parse(Int64, ARGS[2])
# sparsity = parse(Float64, ARGS[3])

FLDR = "datamode$(datamode)sparsity$sparsity"
if !isdir(FLDR)
    mkdir(FLDR)
end

# data structure for flow simulation
const K_CONST =  9.869232667160130e-16 * 86400 * 1e3
const ALPHA = 1.0
mutable struct Ctx
  m; n; h; NT; Δt; Z; X; ρw; ρo;
  μw; μo; K; g; ϕ; qw; qo; sw0
end

function tfCtxGen(m,n,h,NT,Δt,Z,X,ρw,ρo,μw,μo,K,g,ϕ,qw,qo,sw0,ifTrue)
  tf_h = constant(h)
  # tf_NT = constant(NT)
  tf_Δt = constant(Δt)
  tf_Z = constant(Z)
  tf_X= constant(X)
  tf_ρw = constant(ρw)
  tf_ρo = constant(ρo)
  tf_μw = constant(μw)
  tf_μo = constant(μo)
  # tf_K = isa(K,Array) ? Variable(K) : K
  if ifTrue
    tf_K = constant(K)
  else
    tf_K = Variable(K)
  end
  tf_g = constant(g)
  # tf_ϕ = Variable(ϕ)
  tf_ϕ = constant(ϕ)
  tf_qw = constant(qw)
  tf_qo = constant(qo)
  tf_sw0 = constant(sw0)
  return Ctx(m,n,tf_h,NT,tf_Δt,tf_Z,tf_X,tf_ρw,tf_ρo,tf_μw,tf_μo,tf_K,tf_g,tf_ϕ,tf_qw,tf_qo,tf_sw0)
end

if mode == 0
    global Krw, Kro
    # LET-type
    Lw = 1.8
    Ew = 2.1
    Tw = 2.3
    Krwo = 0.6
    function Krw(Sw)
        return Krwo*Sw^Lw/(Sw^Lw + Ew*(1-Sw)^Tw)
    end

    function Kro(So)
        return So^Lw / (So^Lw + Ew*(1-So)^Tw)
    end
elseif mode == 1
    global Krw, Kro, θ1, θ2
    θ1 = Variable(ae_init([1,20,20,20,1]))
    θ2 = Variable(ae_init([1,20,20,20,1]))
    function Krw(Sw)
        Sw_ = tf.reshape(Sw, (-1,1))
        y = ae(Sw_, [20,20,20,1],θ1)
        tf.reshape((tanh(y)+1)/2, Sw.shape)
    end
    function Kro(So)
        So_ = tf.reshape(So, (-1,1))
        y = ae(So_, [20,20,20,1],θ2)
        tf.reshape(1-(tanh(y)+1)/2, So.shape)
    end
end

function krw_and_kro()
    sw = LinRange(0,1,100)|>collect|>constant
    Krw(sw), Kro(1-sw)
end

# IMPES for flow simulation
function ave_normal(quantity, m, n)
    aa = sum(quantity)
    return aa/(m*n)
end

function onestep(sw, p, m, n, h, Δt, Z, ρw, ρo, μw, μo, K, g, ϕ, qw, qo)
    # step 1: update p
    λw = Krw(sw)/μw
    λo = Kro(1-sw)/μo
    # λw = sw.*sw/μw
    # λo = (1-sw).*(1-sw)/μo
    λ = λw + λo
    q = qw + qo + λw/(λo+1e-16).*qo
    # q = qw + qo
    potential_c = (ρw - ρo)*g .* Z

    # Step 1: implicit potential
    Θ = upwlap_op(K * K_CONST, λo, potential_c, h, constant(0.0))

    load_normal = (Θ+q/ALPHA) - ave_normal(Θ+q/ALPHA, m, n)

    # p = poisson_op(λ.*K* K_CONST, load_normal, h, constant(0.0), constant(1))
    p = upwps_op(K * K_CONST, λ, load_normal, p, h, constant(0.0), constant(0)) # potential p = pw - ρw*g*h

    # step 2: implicit transport
    sw = sat_op(sw, p, K * K_CONST, ϕ, qw, qo, μw, μo, sw, Δt, h)
    return sw, p
end

function imseq(tf_ctx)
    ta_sw, ta_p = TensorArray(NT+1), TensorArray(NT+1)
    ta_sw = write(ta_sw, 1, tf_ctx.sw0)
    ta_p = write(ta_p, 1, constant(zeros(tf_ctx.m, tf_ctx.n)))
    i = constant(1, dtype=Int32)
    function condition(i, tas...)
        i <= tf_ctx.NT
    end
    function body(i, tas...)
        ta_sw, ta_p = tas
        sw, p = onestep(read(ta_sw, i), read(ta_p, i), tf_ctx.m, tf_ctx.n, tf_ctx.h, tf_ctx.Δt, tf_ctx.Z, tf_ctx.ρw, tf_ctx.ρo, tf_ctx.μw, tf_ctx.μo, tf_ctx.K, tf_ctx.g, tf_ctx.ϕ, tf_ctx.qw[i], tf_ctx.qo[i])
        ta_sw = write(ta_sw, i+1, sw)
        ta_p = write(ta_p, i+1, p)
        i+1, ta_sw, ta_p
    end

    _, ta_sw, ta_p = while_loop(condition, body, [i, ta_sw, ta_p])
    out_sw, out_p = stack(ta_sw), stack(ta_p)
end

# visualization functions
function plot_saturation_series(S)

    z_inj = (Int(round(0.6*m))-1)*h + h/2.0
    x_inj = (Int(round(0.1*n))-1)*h + h/2.0
    z_prod = (Int(round(0.6*m))-1)*h + h/2.0
    x_prod = (Int(round(0.9*n))-1)*h + h/2.0
    fig2,axs = subplots(3,3, figsize=[30,15], sharex=true, sharey=true)
    ims = Array{Any}(undef, 9)
    for iPrj = 1:3
        for jPrj = 1:3
            @info iPrj, jPrj
            ims[(iPrj-1)*3+jPrj] = axs[iPrj,jPrj].imshow(S[survey_indices[(iPrj-1)*3+jPrj], :, :], extent=[0,n*h,m*h,0], vmin=0.0, vmax=0.6);
            axs[iPrj,jPrj].title.set_text("Snapshot $((iPrj-1)*3+jPrj)")
            if jPrj == 1 || jPrj == 1
                axs[iPrj,jPrj].set_ylabel("Depth (m)")
            end
            if iPrj == 3 || iPrj == 3
                axs[iPrj,jPrj].set_xlabel("Distance (m)")
            end
            # if iPrj ==2 && jPrj == 3
            # cb = fig2.colorbar(ims[(iPrj-1)*3+jPrj], ax=axs[iPrj,jPrj])
            # cb.set_label("Saturation")
            axs[iPrj,jPrj].scatter(x_inj, z_inj, c="r", marker=">", s=128)
            axs[iPrj,jPrj].scatter(x_prod, z_prod, c="r", marker="<", s=128)
        end
    end
    # fig2.subplots_adjust(wspace=0.04, hspace=0.042)
    fig2.subplots_adjust(wspace=0.02, hspace=0.18)
    cbar_ax = fig2.add_axes([0.91, 0.08, 0.01, 0.82])
    cb2 = fig2.colorbar(ims[1], cax=cbar_ax)
    cb2.set_label("Saturation")
    # savefig("figures_summary/Saturation_evo_patchy_init.pdf",bbox_inches="tight",pad_inches = 0);

end

function plot_saturation(S)

    z_inj = (Int(round(0.6*m))-1)*h + h/2.0
    x_inj = (Int(round(0.1*n))-1)*h + h/2.0
    z_prod = (Int(round(0.6*m))-1)*h + h/2.0
    x_prod = (Int(round(0.9*n))-1)*h + h/2.0
    imshow(S[end, :, :], extent=[0,n*h,m*h,0], vmin=0.0, vmax=0.6)
    ylabel("Depth (m)")
    xlabel("Distance (m)")
    colorbar()
    scatter(x_inj, z_inj, c="r", marker=">", s=128)
    scatter(x_prod, z_prod, c="r", marker="<", s=128)

end

function plot_kr(krw, kro, w=missing, o=missing)
    close("all")
    plot(LinRange(0,1,100), krw, "b-", label="\$K_{rw}\$")
    plot(LinRange(0,1,100), kro, "r-", label="\$K_{ro}\$")
    if !ismissing(w)
        plot(LinRange(0,1,100), w, "g--", label="True \$K_{rw}\$")
        plot(LinRange(0,1,100), o, "c--", label="True \$K_{ro}\$")
    end
    xlabel("\$S_w\$")
    ylabel("\$K_r\$")
    legend()
end

# parameters for flow simulation
const SRC_CONST = 86400.0 #
const GRAV_CONST = 9.8    # gravity constant

if datamode==0
    # Hyperparameter for flow simulation
    global m = 15
    global n = 30
    global h = 30.0 # meter
    global NT  = 50
    global dt_survey = 5
    global Δt = 20.0 # day
else
    global m = 45
    global n = 90
    global h = 10.0 # meter
    global NT  = 100
    global dt_survey = 10
    global Δt = 10.0 # day
end

z = (1:m)*h|>collect
x = (1:n)*h|>collect
X, Z = np.meshgrid(x, z)

ρw = 501.9
ρo = 1053.0
μw = 0.1
μo = 1.0

K_init = 20.0 .* ones(m,n) # initial guess of permeability

g = GRAV_CONST
ϕ = 0.25 .* ones(m,n)
qw = zeros(NT, m, n)
qw[:,Int(round(0.6*m)),Int(round(0.1*n))] .= 0.005 * (1/h^2)/10.0 * SRC_CONST
qo = zeros(NT, m, n)
qo[:,Int(round(0.6*m)),Int(round(0.9*n))] .= -0.005 * (1/h^2)/10.0 * SRC_CONST
sw0 = zeros(m, n)
survey_indices = collect(1:dt_survey:NT+1) # 10 stages
n_survey = length(survey_indices)

K = 20.0 .* ones(m,n) # millidarcy
K[Int(round(0.52*m)):Int(round(0.67*m)),:] .= 120.0
tfCtxTrue = tfCtxGen(m,n,h,NT,Δt,Z,X,ρw,ρo,μw,μo,K,g,ϕ,qw,qo, sw0, true)

out_sw_true, out_p_true = imseq(tfCtxTrue)
krw, kro = krw_and_kro()

using Random; Random.seed!(233)
obs_ids = rand(1:m*n, Int(round(m*n*sparsity)))
obs = tf.reshape(out_sw_true[1:NT+1], (NT+1,-1))[:, obs_ids]

error()
# executing the simulation
if mode == 0
    # generate data
    sess = Session(); init(sess)
    S,obs_ = run(sess, [out_sw_true, obs])
    krw_, kro_ = run(sess, [krw, kro])
    matwrite("$FLDR/Data.mat", Dict("S"=>S, "krw"=>krw_, "kro"=>kro_, "obs"=>obs_))
    close("all");plot_kr(krw_, kro_); savefig("$FLDR/krwo.png")
    close("all");plot_saturation(S);
    X, Y = np.meshgrid(collect((0:n-1)*h), collect((0:m-1)*h))
    x = X[:][obs_ids] .+h/2; y = Y[:][obs_ids].+h/2
    scatter(x, y, marker = "*", s=80, color="magenta")
    savefig("$FLDR/sat.png")
    # savefig("$FLDR/sat.pdf")

    # plot_saturation_series(S)

    # savefig("$FLDR/satnn.pdf")
    # savefig("$FLDR/satall.pdf")

else
    dat = Dict{String, Any}("loss" => Float64[], "errw"=>Float64[], "erro"=>Float64[])
    summary = (vs, i, loss_)->begin
        global dat
        S, krw_, kro_, θ1, θ2 = vs
        errw = norm(krw_-wref)
        erro = norm(kro_-oref)
        # dat["S$i"] = S
        dat["w$i"] = krw_
        dat["o$i"] = kro_
        dat["loss"] = [dat["loss"]; loss_]
        dat["errw"] = [dat["errw"]; errw]
        dat["erro"] = [dat["erro"]; erro]
        dat["theta1"] = θ1
        dat["theta2"] = θ2
        if mod(i, 10)==1
            close("all");plot_kr(krw_, kro_, wref, oref); savefig("$FLDR/krwo$i.png")
            close("all");plot_saturation(S); savefig("$FLDR/sat$i.png")
            close("all");
            semilogy(dat["errw"], label="\$k_{r1}\$");
            semilogy(dat["erro"], label="\$k_{r2}\$"); legend()
            xlabel("Iteration"); ylabel("Error"); savefig("$FLDR/err.png")
            close("all");semilogy(dat["loss"]); xlabel("Iteration"); ylabel("Loss"); savefig("$FLDR/loss.png")
        end
        @show i, loss_, errw, erro
        matwrite("$FLDR/invData.mat", dat)
    end
    d = matread("$FLDR/Data.mat")
    Sref, wref, oref, obsref = d["S"], d["krw"][:], d["kro"][:], d["obs"]
    # loss = sum((out_sw_true - Sref)^2)
    loss = sum((obsref-obs)^2)
    sess = Session(); init(sess)
    BFGS!(sess, loss, 500;vars = [out_sw_true, krw, kro, θ1, θ2], callback = summary)
end

# dd = matread("$FLDR/invData.mat")
# run(sess, [assign(θ1,dd["theta1"]),assign(θ2,dd["theta2"])] )


text:
using Draftcheck

check(joinpath(dirname(@__FILE__), "..", "example", "test.tex"), joinpath(dirname(@__FILE__), "..", "example", "example.jl"))


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

JLLWrappers.@generate_main_file_header("HiGHSstatic")
JLLWrappers.@generate_main_file("HiGHSstatic", UUID("30d44bb2-ff8b-57e8-b516-b10f428bad0e"))
end  # module HiGHSstatic_jll


text:
using Assets, Nasdaq

using Test

response = Nasdaq.load(symbols=[:MSFT,:GOOG,:AAPL], force=true)

# Check that some basic stocks listed on Nasdaq have been loaded correctly
const stocks = ((:MSFT, "Microsoft Corp"),
                (:GOOG, "Alphabet Cl C"),
                (:AAPL, "Apple Inc"))

@testset "Basic currencies" begin
    for (sym, nam) in stocks
        @test response[sym] == nam
    end
end


text:
# TODO: GitHub's list users only uses `since`.

export @paginate

"""
    @paginate fun(args...; kwargs...) page=1 per_page=100 -> Paginator

Create an iterator that paginates the results of repeatedly calling `fun(args...; kwargs...)`.
`fun` must take a [`Forge`](@ref) as its first argument and return a `Result{Vector{T}}`.

## Keywords
- `page::Int=1`: Starting page.
- `per_page::Int=100`: Number of entries per page.
"""
macro paginate(ex::Expr, kws::Expr...)
    (ex.args[2] isa Expr && ex.args[2].head === :parameters) ||
        insert!(ex.args, 2, Expr(:parameters))
    append!(ex.args[2].args, map(kw -> Expr(:kw, kw.args...), kws))
    insert!(ex.args, 3, ex.args[1])
    isempty(ex.args[2].args) && deleteat!(ex.args, 2)
    ex.args[1] = :(GitForge.paginate)
    esc(ex)
end

mutable struct Paginator{T}
    f::Function
    page::Int
    next::Dict
    last::Int
    r::Result{Vector{T}}

    function Paginator{T}(
        f::Function, args::Tuple, kwargs::Iterators.Pairs,
        page::Int, per_page::Int,
    ) where T
        pager = Dict("page" => page, "per_page" => per_page)
        query = merge(get(kwargs, :query, Dict()), pager)
        return new{T}(
            q::Dict -> f(args...; kwargs..., query=merge(query, q)),
            page, Dict(), typemax(Int),
        )
    end
end

function paginate(
    fun::Function, f::Forge, args...;
    page::Int=1, per_page::Int=100, kwargs...,
)
    V = into(f, fun)
    V <: Vector || throw(ArgumentError("Function must return Result{Vector{T}}"))
    return Paginator{eltype(V)}(fun, (f, args...), kwargs, page, per_page)
end

function Base.iterate(p::Paginator, state=nothing)
    # If we have no results yet, fetch one.
    isdefined(p, :r) || (p.r = p.f(p.next))
    exception(p.r) === nothing || return nothing

    # The state is nothing when we are starting a new page.
    it = state === nothing ? iterate(value(p.r)) : iterate(value(p.r), state)
    if it === nothing
        # This page is finished, so either fetch a new page or quit.
        p.page > p.last && return nothing
        nextpage!(p) || return nothing
        return iterate(p)
    else
        # More results in the page.
        return it
    end
end

function nextpage!(p::Paginator)
    p.r = p.f(p.next)
    exception(p.r) === nothing || return false
    p.page += 1
    next, last = parserels(HTTP.header(response(p.r), "Link"))
    next === nothing || (p.next = next)
    last === nothing || (p.last = last)
    return true
end

function unescape(url::AStr)
    ms = eachmatch(r"[\?&](.*?)=(.*?)(?:&|$)", url; overlap=true)
    return Dict(m[1] => m[2] for m in ms)
end

function parserels(link::AStr)
    nextm = match(r"(<.*?)>; rel=\"next\"", link)
    lastm = match(r"(<.*?)>; rel=\"last\"", link)
    nextd = nextm === nothing ? nothing : unescape(nextm[1])
    lastd = lastm === nothing ? nothing : unescape(lastm[1])
    last = if lastd == nothing
        nothing
    else
        l = get(lastd, "page", nothing)
        l === nothing ? nothing : parse(Int, l)
    end
    return nextd, last
end