text:
module TestVariables

using MathOptInterface
using Test

const MOI = MathOptInterface
const MOIU = MOI.Utilities

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

function test_get_bounds_Float64()
    T = Float64
    model = MOIU.Model{T}()

    x = MOI.add_variable(model)
    y_v, y_c = MOI.add_constrained_variable(model, MOI.GreaterThan{T}(1.0))
    z_v, z_c = MOI.add_constrained_variable(model, MOI.LessThan{T}(-1.0))
    w_v, w_c = MOI.add_constrained_variable(model, MOI.Interval{T}(-2.0, 3.0))
    u_v, u_c =
        MOI.add_constrained_variable(model, MOI.Semicontinuous{T}(1.0, 3.5))
    v_v, v_c =
        MOI.add_constrained_variable(model, MOI.Semiinteger{T}(-1.0, 2.5))
    e_v, e_c = MOI.add_constrained_variable(model, MOI.EqualTo(1.35))

    @test -Inf == @inferred MOIU.get_bounds(model, T, x)[1]
    @test Inf == @inferred MOIU.get_bounds(model, T, x)[2]

    @test 1.0 == @inferred MOIU.get_bounds(model, T, y_v)[1]
    @test Inf == @inferred MOIU.get_bounds(model, T, y_v)[2]

    @test -Inf == @inferred MOIU.get_bounds(model, T, z_v)[1]
    @test -1.0 == @inferred MOIU.get_bounds(model, T, z_v)[2]

    @test -2.0 == @inferred MOIU.get_bounds(model, T, w_v)[1]
    @test 3.0 == @inferred MOIU.get_bounds(model, T, w_v)[2]

    @test 0.0 == @inferred MOIU.get_bounds(model, T, u_v)[1]
    @test 3.5 == @inferred MOIU.get_bounds(model, T, u_v)[2]

    @test -1.0 == @inferred MOIU.get_bounds(model, T, v_v)[1]
    @test 2.5 == @inferred MOIU.get_bounds(model, T, v_v)[2]

    @test 1.35 == @inferred MOIU.get_bounds(model, T, e_v)[1]
    @test 1.35 == @inferred MOIU.get_bounds(model, T, e_v)[2]

    MOI.add_constraint(model, y_v, MOI.LessThan{T}(3.6))

    @test 1.0 == @inferred MOIU.get_bounds(model, T, y_v)[1]
    @test 3.6 == @inferred MOIU.get_bounds(model, T, y_v)[2]
    return
end

function test_get_bounds_UInt128()
    T = UInt128
    model = MOIU.Model{T}()

    x = MOI.add_variable(model)
    y_v, y_c = MOI.add_constrained_variable(model, MOI.GreaterThan{T}(1))
    z_v, z_c = MOI.add_constrained_variable(model, MOI.LessThan{T}(1))
    w_v, w_c = MOI.add_constrained_variable(model, MOI.Interval{T}(2, 3))
    u_v, u_c = MOI.add_constrained_variable(model, MOI.Semicontinuous{T}(1, 3))
    v_v, v_c = MOI.add_constrained_variable(model, MOI.Semiinteger{T}(4, 7))

    @test typemin(T) == @inferred MOIU.get_bounds(model, T, x)[1]
    @test typemax(T) == @inferred MOIU.get_bounds(model, T, x)[2]

    @test 1 == @inferred MOIU.get_bounds(model, T, y_v)[1]
    @test typemax(T) == @inferred MOIU.get_bounds(model, T, y_v)[2]

    @test typemin(T) == @inferred MOIU.get_bounds(model, T, z_v)[1]
    @test 1 == @inferred MOIU.get_bounds(model, T, z_v)[2]

    @test 2 == @inferred MOIU.get_bounds(model, T, w_v)[1]
    @test 3 == @inferred MOIU.get_bounds(model, T, w_v)[2]

    @test 0 == @inferred MOIU.get_bounds(model, T, u_v)[1]
    @test 3 == @inferred MOIU.get_bounds(model, T, u_v)[2]

    @test 0 == @inferred MOIU.get_bounds(model, T, v_v)[1]
    @test 7 == @inferred MOIU.get_bounds(model, T, v_v)[2]

    MOI.add_constraint(model, y_v, MOI.LessThan{T}(5))

    @test 1 == @inferred MOIU.get_bounds(model, T, y_v)[1]
    @test 5 == @inferred MOIU.get_bounds(model, T, y_v)[2]
end

end  # module

TestVariables.runtests()


text:
## ---  Define a StatiCompiler- and LLVM-compatible dynamically-sized string type

    # Define the MallocString type, backed by a malloc'd heap of memory

    # Definition and constructors:
    """
    ```julia
    struct MallocString
        pointer::Ptr{UInt8}
        length::Int
    end
    ```
    A stringy object that contains `length` bytes (i.e., `UInt8`s, including
    the final null-termination `0x00`), at a location in memory specified by
    `pointer`.

    A `MallocString` should generally behave like a base Julia `String`, but is
    explicitly null-terminated, mutable, standalone-StaticCompiler-safe (does not
    require libjulia) and backed by `malloc`ed memory which is not tracked by
    the GC and should be `free`d when no longer in use.

    Can be constructed with the `m"..."` string macro.

    Unlike base Julia `String`s, slicing does not create a copy, but rather a view.
    You are responsible for ensuring that any such views are null-terminated if you
    wish to pass them to any functions (including most libc/system IO) that expect
    null-termination.

    ## Examples
    ```julia
    julia> s = m"Hello world!"
    m"Hello world!"

    julia> s[8:12] = c"there"; s
    m"Hello there!"

    julia> s[1:5]
    StringView: "Hello"

    julia> s[1:5] == "Hello"
    true

    julia> StaticString(s[1:5])
    c"Hello"

    julia> free(s)
    0
    ```
    """
    struct MallocString <: AbstractPointerString
        pointer::Ptr{UInt8}
        length::Int
    end

    """
    ```julia
    MallocString(undef, N)
    ```
    Construct an uninitialized `N`-byte (including null-termination!) `MallocString`.
    Here `undef` is the `UndefInitializer`.

    ## Examples
    ```julia
    julia> s = MallocString(undef, 10)
    m""

    julia> free(s)
    0
    ```
    """
    @inline function MallocString(::UndefInitializer, N::Int)
        s = MallocString(Ptr{UInt8}(malloc(N)), N)
        s[end] = 0x00
        return s
    end
    """
    ```julia
    MallocString(data::NTuple{N, UInt8})
    ```
    Construct a `MallocString` containing the `N` bytes specified by `data`.
    To yield a valid string, `data` must be null-terminated, i.e., end in `0x00`.

    ```julia
    MallocString(s::AbstractStaticString)
    ```
    Construct a `MallocString` containing the same data as the existing string `s`

    ## Examples
    ```julia
    julia> data = (0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x20, 0x77, 0x6f, 0x72, 0x6c, 0x64, 0x21, 0x00);

    julia> s = MallocString(data)
    m"Hello world!"

    julia> s2 = MallocString(s[1:5])
    m"Hello"

    julia> free(s)
    0

    julia> free(s2)
    0
    ```
    """
    @inline function MallocString(data::NTuple{N, UInt8}) where N
        s = MallocString(Ptr{UInt8}(malloc(N)), N)
        s[:] = data
        return s
    end
@inline function MallocString(s::AbstractStaticString)
    N = length(s) + 1 # Add room for null-termination
    c = MallocString(Ptr{UInt8}(malloc(N)), N)
    c[1:length(s)] = s
    c[end] = 0x00
    return c
end
    @inline MallocString(p::Ptr{UInt8}) = MallocString(p, strlen(p)+1)
    @inline MallocString(argv::Ptr{Ptr{UInt8}}, n::Integer) = MallocString(unsafe_load(argv, n))

    # String macro to create null-terminated `MallocStrings`s
    """
    ```julia
    @m_str -> MallocString
    ```
    Construct a `MallocString`, such as `m"Foo"`.

    A `MallocString` should generally behave like a base Julia `String`, but is
    explicitly null-terminated, mutable, standalone-StaticCompiler-safe (does not
    require libjulia) and is backed by `malloc`d memory which is not tracked by
    the GC and should be `free`d when no longer in use.

    ## Examples
    ```julia
    julia> s = m"Hello there!"
    m"Hello there!"

    julia> s == "Hello there!"
    true

    julia> free(s)
    0
    ```
    """
    macro m_str(s)
        n = _unsafe_unescape!(s)
        t = Expr(:tuple, codeunits(s)[1:n]..., 0x00)
        :(MallocString($t))
    end

    # Destructor:
    @inline free(s::MallocString) = free(s.pointer)

    # Fundamentals -- where overriding AbstractStaticString defaults
    @inline Base.length(s::MallocString) = s.length - 1     # For consistency with base
    @inline function Base.:(==)(a::Union{StaticString,MallocString}, b::Union{StaticString,MallocString})
        (N = length(a)) == length(b) || return false
        pa, pb = pointer(a), pointer(b)
        for n in 0:N
            unsafe_load(pa + n) == unsafe_load(pb + n) || return false
        end
        return true
    end

    # Custom replshow for interactive use (n.b. _NOT_ static-compilerable)
    function Base.show(io::IO, s::MallocString)
        Base.print(io, "m\"")
        Base.escape_string(io, Base.unsafe_string(pointer(s)))
        Base.print(io, "\"")
    end

    # Some of the AbstractString interface -- where overriding AbstractStaticString defaults
    @inline function Base.:*(a::MallocString, b::AbstractStaticString)  # Concatenation
        N = length(a) + length(b) + 1
        c = MallocString(undef, N)
        c[1:length(a)] = a
        c[length(a)+1:length(a)+length(b)] = b
        c[end] = 0x00 # Null-terminate
        return c
    end
    @inline function Base.:^(s::MallocString, n::Integer)       # Repetition
        l = length(s) # Excluding the null-termination
        N = n*l + 1
        c = MallocString(undef, N)
        for i=1:n
            c[(l*(i-1) + 1):(l*i)] = s
        end
        c[end] = 0x00 # Null-terminate
        return c
    end
    @inline function Base.copy(s::MallocString)
        new_s = MallocString(undef, sizeof(s))
        new_s[:] = s
        return new_s
    end

    # As Base.unsafe_string, but loading to a MallocString instead
    @inline function unsafe_mallocstring(p::Ptr{UInt8})
        len = strlen(p)+1
        s = MallocString(undef, len)
        for i=1:len
            s[i] = unsafe_load(p, i)
        end
        return s
    end

ext: jl


text:
import ACE: OneParticleBasis, AbstractState
import ACE.OrthPolys: TransformedPolys

import ForwardDiff
import NamedTupleTools
using NamedTupleTools: namedtuple

@doc raw"""
`struct Scal1pBasis <: OneParticleBasis`

One-particle basis of the form $P_n(x_i)$ for a general scalar, invariant
input `x`. This type basically just translates the `TransformedPolys` into a valid
one-particle basis.
"""
mutable struct Scal1pBasis{VSYM, VIDX, ISYM, T, TT, TJ} <: OneParticleBasis{T}
   P::TransformedPolys{T, TT, TJ}
   B_pool::VectorPool{T}
   dB_pool::VectorPool{T}
end

scal1pbasis(varsym::Symbol, idxsym::Symbol, args...; varidx = 1, kwargs...) =
            Scal1pBasis(varsym, varidx, idxsym,
                  ACE.OrthPolys.transformed_jacobi(args...; kwargs...))

Scal1pBasis(varsym::Symbol, idxsym::Symbol, P::TransformedPolys{T, TT, TJ}
            ) where {T, TT, TJ} = Scal1pBasis(varsym, 1, idxsym, P)

Scal1pBasis(varsym::Symbol, varidx::Integer, idxsym::Symbol, P::TransformedPolys{T, TT, TJ}
            ) where {T, TT, TJ} =
      Scal1pBasis{varsym, Int(varidx), idxsym, T, TT, TJ}(P)

Scal1pBasis{VSYM, VIDX, ISYM, T, TT, TJ}(P::TransformedPolys{T, TT, TJ}
            ) where {VSYM, VIDX, ISYM, T, TT, TJ} =
      Scal1pBasis{VSYM, VIDX, ISYM, T, TT, TJ}(P, VectorPool{T}(), VectorPool{T}())

_varsym(basis::Scal1pBasis{VSYM}) where {VSYM} = VSYM
_varidx(basis::Scal1pBasis{VSYM, VIDX}) where {VSYM, VIDX} = VIDX
_idxsym(basis::Scal1pBasis{VSYM, VIDX, ISYM}) where {VSYM, VIDX, ISYM} = ISYM

_val(X::AbstractState, basis::Scal1pBasis) =
      getproperty(X, _varsym(basis))[_varidx(basis)]

_val(x::Number, basis::Scal1pBasis) = x

# ---------------------- Implementation of Scal1pBasis

Base.length(basis::Scal1pBasis) = length(basis.P)

get_spec(basis::Scal1pBasis, n::Integer) = NamedTuple{(_idxsym(basis),)}(n)

get_spec(basis::Scal1pBasis) = [ get_spec(basis, i) for i = 1:length(basis) ]

function set_spec!(basis::Scal1pBasis, spec::Vector)
   # we don't want to set anything, just check that its compatible with the spec
   # we produce on the fly
   old_spec = get_spec(basis)
   @assert old_spec == spec
   return basis
end

==(P1::Scal1pBasis, P2::Scal1pBasis) = (P1.P == P2.P)

write_dict(basis::Scal1pBasis{T}) where {T} = Dict(
      "__id__" => "ACE_Scal1pBasis",
          "P" => write_dict(basis.P) ,
          "varsym" => string(_varsym(basis)),
          "varidx" => _varidx(basis),
          "idxsym" => string(_idxsym(basis)) )

read_dict(::Val{:ACE_Scal1pBasis}, D::Dict) =
      Scal1pBasis(Symbol(D["varsym"]), Int(D["varidx"]), Symbol(D["idxsym"]),
                  read_dict(D["P"]))

valtype(basis::Scal1pBasis) = valtype(basis.P)

valtype(basis::Scal1pBasis, cfg::AbstractConfiguration) =
      valtype(basis, zero(eltype(cfg)))

valtype(basis::Scal1pBasis, X::AbstractState) =
      valtype(basis.P, _val(X, basis))

gradtype(basis::Scal1pBasis, X::AbstractState) =
      dstate_type(valtype(basis, X), X)

gradtype(basis::Scal1pBasis, cfg::AbstractConfiguration) =
      gradtype(basis, zero(eltype(cfg)))

argsyms(basis::Scal1pBasis) = ( _varsym(basis), )

symbols(basis::Scal1pBasis) = [ _idxsym(basis) ]

indexrange(basis::Scal1pBasis) = NamedTuple{(_idxsym(basis), )}((1:length(basis),))

_getidx(b, basis::Scal1pBasis) = b[_idxsym(basis) ]

isadmissible(b, basis::Scal1pBasis) = (1 <= _getidx(b, basis) <= length(basis))

degree(b, basis::Scal1pBasis) = _getidx(b, basis) - 1

get_index(basis::Scal1pBasis, b) = _getidx(b, basis)

# TODO: need better structure to support this ...
rand_radial(basis::Scal1pBasis) = rand_radial(basis.P)

# ---------------------------  Evaluation code
#

evaluate!(B, basis::Scal1pBasis, x::Number) =
      evaluate!(B, basis.P, x)

evaluate!(B, basis::Scal1pBasis, X::AbstractState) =
      evaluate!(B, basis.P, _val(X, basis))

"""
returns an `SVector{N}` of the form `x * e_I` where `e_I` is the Ith canonical basis vector.
"""
@generated function __e(::SVector{N}, ::Val{I}, x::T) where {N, I, T}
   code = "SA["
   for i = 1:N
      if i == I
         code *= "x,"
      else
         code *= "0,"
      end
   end
   code *= "]"
   quote
      $( Meta.parse(code) )
   end
end

__e(::Number, ::Any, x) = x

function _scal1pbasis_grad(TDX::Type, basis::Scal1pBasis, gval)
   gval_tdx = __e( getproperty(zero(TDX), _varsym(basis)),
                   Val(_varidx(basis)),
                   gval )
   return TDX( NamedTuple{(_varsym(basis),)}((gval_tdx,)) )
end

function evaluate_d!(dB, basis::Scal1pBasis, X::AbstractState)
   TDX = eltype(dB)
   x = _val(X, basis)
   dP = acquire_dB!(basis.P, x)
   evaluate_d!(dP, basis.P, x)
   for n = 1:length(basis)
      dB[n] = _scal1pbasis_grad(TDX, basis, dP[n])
   end
   release_dB!(basis.P, dP)
   return dB
end

function evaluate_ed!(B, dB, basis::Scal1pBasis, X::AbstractState)
   TDX = eltype(dB)
   x = _val(X, basis)
   dP = acquire_dB!(basis.P, x)
   evaluate!(B, basis.P, x)
   evaluate_d!(dP, basis.P, x)
   for n = 1:length(basis)
      dB[n] = _scal1pbasis_grad(TDX, basis, dP[n])
   end
   release_dB!(basis.P, dP)
   return B, dB
end

# this one we probably only need for training so can relax the efficiency a bit
function evaluate_dd(basis::Scal1pBasis, X::AbstractState)
   ddP = ForwardDiff.derivative(x -> evaluate_d(basis, _val(X, basis)))
   TDX = gradtype(basis, X)
   return _scal1pbasis_grad.(Ref(TDX), Ref(basis), ddP_n)
end

# -------------- AD codes

import ChainRules: rrule, ZeroTangent, NoTangent

function _rrule_evaluate(basis::Scal1pBasis, X::AbstractState,
                         w::AbstractVector{<: Number})
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   a = _rrule_evaluate(basis.P, x, real.(w))
   TDX = ACE.dstate_type(a, X)
   return TDX( NamedTuple{(_varsym(basis),)}( (a,) ) )
end

rrule(::typeof(evaluate), basis::Scal1pBasis, X::AbstractState) =
                  evaluate(basis, X),
                  w -> (NoTangent(), NoTangent(), _rrule_evaluate(basis, X, w))

function _rrule_evaluate_d(basis::Scal1pBasis, X::AbstractState,
                           w::AbstractVector)
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   w1 = [ _val(w, basis) for w in w ]
   a = _rrule_evaluate_d(basis.P, x, w1)
   TDX = ACE.dstate_type(a, X)
   return TDX( NamedTuple{(_varsym(basis),)}( (a,) ) )
end

function rrule(::typeof(evaluate_d), basis::Scal1pBasis, X::AbstractState)
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   dB_ = evaluate_d(basis.P, x)
   TDX = dstate_type(valtype(basis, X), X)
   dB = [ TDX( NamedTuple{(_varsym(basis),)}( (dx,) ) )  for dx in dB_ ]
   return dB,
          w -> (NoTangent(), NoTangent(), _rrule_evaluate_d(basis, X, w))
end


text:
#=
Copyright (c) 2015, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

module AstWalker

# See the documentation on the AstWalk function for a description of how to use this module.

import ..DebugMsg
DebugMsg.init()

using CompilerTools
using ..Helper

# Return this value to indicate to AstWalk that you didn't process the current node
# and that AstWalk should recursively process the given node.
abstract ASTWALK_DIDNT_MODIFY

immutable ASTWALK_RECURSE <: ASTWALK_DIDNT_MODIFY
end
immutable ASTWALK_RECURSE_DUPLICATE <: ASTWALK_DIDNT_MODIFY
end

# Return this node to indicate that AstWalk should remove the current node from
# the container in which it resides.
immutable ASTWALK_REMOVE
end

export AstWalk, ASTWALK_RECURSE, ASTWALK_REMOVE, ASTWALK_RECURSE_DUPLICATE

"""
AstWalk through a lambda expression.
Walk through each input parameters and the body of the lambda.
"""
# TODO - it seems we should walk some parts of meta as well.
function from_lambda(ast :: Array{Any,1}, depth, callback, cbdata :: ANY, top_level_number, read)
  assert(length(ast) == 3)
  param = ast[1]
  meta  = ast[2]
  body  = ast[3]

  @dprintln(3,"from_lambda pre-convert param = ", param, " typeof(param) = ", typeof(param))
  # AstWalk over each incoming parameter to the lambda.
  for i = 1:length(param)
    @dprintln(3,"from_lambda param[i] = ", param[i], " typeof(param[i]) = ", typeof(param[i]))
    param[i] = from_expr(param[i], depth, callback, cbdata, top_level_number, false, read)
  end
  @dprintln(3,"from_lambda post-convert param = ", param, " typeof(param) = ", typeof(param))

  @dprintln(3,"from_lambda pre-convert body = ", body, " typeof(body) = ", typeof(body))
  # AstWalk over the body of the function.
  body = from_expr(body, depth, callback, cbdata, top_level_number, false, read)
  @dprintln(3,"from_lambda post-convert body = ", body, " typeof(body) = ", typeof(body))
  # Make sure that what comes back is still a body type.
  if typeof(body) != Expr || body.head != :body
    @dprintln(0,"AstWalk from_lambda got a non-body returned from procesing body")
    @dprintln(0,body)
    throw(string("big problem"))
  end

  ast[1] = param
  ast[2] = meta
  ast[3] = body
  return ast
end

"""
AstWalk through a function body.
"""
function from_body(ast :: Array{Any,1}, depth, callback, cbdata :: ANY, top_level_number, read)
  len = length(ast)
  # We need to differentiate top-level statements from non-top level and so this function handles the
  # top-level ones.
  top_level = true

  body = Any[]

  # AstWalk over each statement in a lambda body.
  for i = 1:len
    @dprintln(2,"Processing top-level ast #",i," depth=",depth)

    @dprintln(3,"AstWalk from_exprs, ast[", i, "] = ", ast[i])
    new_expr = from_expr(ast[i], depth, callback, cbdata, i, top_level, read)
    @dprintln(3,"AstWalk from_exprs done, ast[", i, "] = ", new_expr)
    # If the return result doesn't indicate the statement should be removed then put the new version
    # of the statement (new_expr) into the new body.
    if new_expr != ASTWALK_REMOVE
      push!(body, new_expr)
    end
  end

  # Return the updated body.
  return body
end

"""
AstWalk through an array of expressions.
"""
function from_exprs(ast, depth, callback, cbdata :: ANY, top_level_number, read)
  len = length(ast)
  # We need to differentiate top-level statements from non-top level and so this function handles the
  # non-top-level ones.
  top_level = false

  body = Any[]

  for i = 1:len
    @dprintln(2,"Processing ast #",i," depth=",depth)
    @dprintln(3,"AstWalk from_exprs, ast[", i, "] = ", ast[i])
    new_expr = from_expr(ast[i], depth, callback, cbdata, i, top_level, read)
    @dprintln(3,"AstWalk from_exprs done, ast[", i, "] = ", new_expr)
    if new_expr != ASTWALK_REMOVE
      push!(body, new_expr)
    end
  end

  return body
end

"""
AstWalk through an assignment expression.
Recursively process the left and right hand sides with AstWalk.
"""
function from_assignment(ast :: Array{Any,1}, depth, callback, cbdata :: ANY, top_level_number, read)
  # Process an assignment statement by process first the left-hand side and then the right.
  @dprintln(3,"from_assignment, lhs = ", ast[1])
  ast[1] = from_expr(ast[1], depth, callback, cbdata, top_level_number, false, false)
  @dprintln(3,"from_assignment, rhs = ", ast[2])
  ast[2] = from_expr(ast[2], depth, callback, cbdata, top_level_number, false, read)
  return ast
end

"""
AstWalk through a call expression.
Recursively process the name of the function and each of its arguments.
"""
function from_call(ast :: Array{Any,1}, depth, callback, cbdata :: ANY, top_level_number, read)
  assert(length(ast) >= 1)
  # A call is a function followed by its arguments.  Extract these parts below.
  fun  = ast[1]
  args = ast[2:end]
  @dprintln(2,"from_call fun = ", fun, " typeof fun = ", typeof(fun))
  if length(args) > 0
    @dprintln(2,"first arg = ",args[1], " type = ", typeof(args[1]))
  end
  # Symbols don't need to be translated.
  # if typeof(fun) != Symbol
      # I suppose this "if" could be wrong.  If you wanted to replace all "x" functions with "y" then you'd need this wouldn't you?
  #    fun = from_expr(fun, depth, callback, cbdata, top_level_number, false, read)
  # end
  # Process the arguments to the function recursively.
  args = from_exprs(ast, depth+1, callback, cbdata, top_level_number, read)

  return args
end

"""
Entry point into the code to perform an AST walk.

This function will cause every node in the AST to be visited and a callback invoked on each one.
The callback can record information about nodes it sees through the opaque cbdata parameter.
The callback can also cause the current node to be replaced in the AST by returning a new AST
node.  If the callback doesn't wish to change the node then it returns ASTWALK_RECURSE which causes
AstWalk to recursively process the sub-trees under the current AST node.  If you want to modify the
current AST node and want the sub-trees of that AST node to be processed first then you manually have
to recursively call AstWalk on each one.  There are some cases where you don't want to recursively
process the sub-trees first and so this recursive process has to be left up to the user.

You generally pass a lambda expression as the first argument although any AST node is acceptable.
The third argument is an object that is opaque to AstWalk but that is passed to every callback.
You can use this object to collect data about the AST as it is walked or to hold information on
how to change the AST as you are walking over it.
The second argument is a callback function.  For each AST node, AstWalk will invoke this callback.
The signature of the callback must be (Any, Any, Int64, Bool, Bool).  The arguments to the callback
are as follows:
    1) The current node of the AST being walked.
    2) The callback data object that you originally passed as the first argument to AstWalk.
    3) Specifies the index of the body's statement that is currently being processed.
    4) True if the current AST node being walked is the root of a top-level statement, false if the AST node is a sub-tree of a top-level statement.
    5) True if the AST node is being read, false if it is being written.
"""
function AstWalk(ast :: ANY, callback, cbdata :: ANY)
  from_expr(ast, 1, callback, cbdata, 0, false, true)
end

#function from_expr(ast :: LambdaInfo, depth, callback, cbdata :: ANY, top_level_number, is_top_level, read)
    #LambdaVarInfo, body = CompilerTools.LambdaHandling.lambdaToLambdaVarInfo(ast)
    #new_body = from_expr(body, depth, callback, cbdata, top_level_number, is_top_level, read)
    #CompilerTools.LambdaHandling.LambdaVarInfoToLambda(LambdaVarInfo, new_body)
#end

"""
The main routine that switches on all the various AST node types.
The internal nodes of the AST are of type Expr with various different Expr.head field values such as :lambda, :body, :block, etc.
The leaf nodes of the AST all have different types.
There are some node types we don't currently recurse into.  Maybe this needs to be extended.
"""
function from_expr(ast :: ANY, depth, callback, cbdata :: ANY, top_level_number, is_top_level, read)
    @dprintln(2,"from_expr depth=",depth," ", " ", ast)

    # For each AST node, we first call the user-provided callback to see if they want to do something
    # with the node.
    ret = callback(ast, cbdata, top_level_number, is_top_level, read)
    @dprintln(2,"callback ret = ",ret)
    # If the return value of the callback isn't ASTWALK_RECURSE then the callback is replaced the current
    # AST node with "ret".
    if ret != ASTWALK_RECURSE && ret != ASTWALK_RECURSE_DUPLICATE
        return ret
    end

    # The user callback didn't replace the AST node so recursively process it.
    # We have a different from_expr_helper that is accurately typed for each possible AST node type.
    ast = from_expr_helper(ast, depth, callback, cbdata, top_level_number, is_top_level, read)
    if ret == ASTWALK_RECURSE_DUPLICATE
        ast = copy(ast)
    end

    @dprintln(3,"Before return for ", ast)
    return ast
end

# Handle recursively walking over an Expr type AST node.
function from_expr_helper(ast::Expr,
                          depth,
                          callback,
                          cbdata::ANY,
                          top_level_number,
                          is_top_level,
                          read)
    # If you have an assignment with an Expr as its left-hand side then you get here with "read = false"
    # but that doesn't  mean the whole Expr is written.  In fact, none of it is written so we set read
    # back to true and then restore in the incoming read value at the end.
    read_save = read
    read = true

    @dprint(2,"Expr ")
    head = ast.head
    args = ast.args
    typ  = ast.typ
    @dprintln(2,head, " ", args)
    if head == :lambda
        args = from_lambda(args, depth, callback, cbdata, top_level_number, read)
    elseif head == :body
        @dprintln(2,"Processing :body Expr in AstWalker.from_expr")
        args = from_body(args, depth+1, callback, cbdata, top_level_number, read)
        @dprintln(2,"Done processing :body Expr in AstWalker.from_expr")
    elseif head == :block
        args = from_exprs(args, depth+1, callback, cbdata, top_level_number, read)
    elseif head == :(.)
        args = from_exprs(args, depth+1, callback, cbdata, top_level_number, read)
    elseif head == :(=) || head == :(.=)
        args = from_assignment(args, depth, callback, cbdata, top_level_number, read)
    elseif head == :(::)
        assert(length(args) == 2)
        @dprintln(3, ":: args[1] = ", args[1])
        @dprintln(3, ":: args[2] = ", args[2])
        args[1] = from_expr(args[1], depth, callback, cbdata, top_level_number, false, read)
    elseif head == :return
        args = from_exprs(args, depth, callback, cbdata, top_level_number, read)
    elseif head == :invoke
        args = from_call(args, depth, callback, cbdata, top_level_number, read)
    elseif head == :call
        args = from_call(args, depth, callback, cbdata, top_level_number, read)
        # TODO: catch domain IR result here
    elseif head == :call1
        args = from_call(args, depth, callback, cbdata, top_level_number, read)
        # TODO?: tuple
    elseif head == :line
        # skip
    elseif head == :copy
        # turn array copy back to plain Julia call
        head = :call
        args = vcat(:copy, args)
    elseif head == :copyast
        @dprintln(2,"copyast type")
        # skip
    elseif head == :gotoifnot
        assert(length(args) == 2)
        args[1] = from_expr(args[1], depth, callback, cbdata, top_level_number, false, read)
    elseif head == :getindex
        args = from_exprs(args,depth, callback, cbdata, top_level_number, read)
    elseif head == :new
        args = from_exprs(args,depth, callback, cbdata, top_level_number, read)
    elseif head == :arraysize
        assert(length(args) == 2)
        args[1] = from_expr(args[1], depth, callback, cbdata, top_level_number, false, read)
        args[2] = from_expr(args[2], depth, callback, cbdata, top_level_number, false, read)
    elseif head == :alloc
        assert(length(args) == 2)
        args[2] = from_exprs(args[2], depth, callback, cbdata, top_level_number, read)
    elseif head == :boundscheck || head == :inbounds
        # skip
    elseif head == :type_goto
        assert(length(args) == 2)
        args[1] = from_expr(args[1], depth, callback, cbdata, top_level_number, false, read)
        args[2] = from_expr(args[2], depth, callback, cbdata, top_level_number, false, read)
    elseif head == :tuple
        for i = 1:length(args)
            args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
        end
    elseif head == :enter
        # skip
    elseif head == :leave
        # skip
    elseif head == :curly
        # skip
    elseif head == :the_exception
        # skip
    elseif head == :&
        # skip
    elseif head == :static_typeof
        for i = 1:length(args)
            args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
        end
    elseif head == :ccall
        for i = 1:length(args)
            args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
        end
    elseif head == :function
	@dprintln(3,"in function head")
	args[2] = from_expr(args[2], depth, callback, cbdata, top_level_number, false, read)
    elseif head == :vcat || head == :typed_vcat || head == :hcat || head == :typed_hcat
	@dprintln(3,"in vcat head")
	#skip
    elseif head == :ref
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :meta
	# ignore :meta for now. TODO: we might need to walk its args.
    elseif head == :comprehension || head == :vect || head == :generator
	# args are either Expr or Symbol
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :typed_comprehension
	# args are either Expr or Symbol
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :(->) || head == :(&&) || head == :(||)
	# args are either Expr or Symbol
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :(:)
	# args are either Expr or Symbol
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :const
	dump(ast,1000)
	# ignore :const for now.
    elseif head == :for
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head in Set([:(+=), :(/=), :(*=), :(-=)])
        args[1] = from_expr(args[1], depth, callback, cbdata, top_level_number, false, false)
        args[2] = from_expr(args[2], depth, callback, cbdata, top_level_number, false, read)
    elseif head == :if
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :comparison
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :while
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :let
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :local
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :quote
	for i = 1:length(args)
	    args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
	end
    elseif head == :simdloop
        # skip
    elseif head == :static_parameter
        # skip
    elseif head == :macrocall
        for i = 1:length(args)
            args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
        end
    elseif head == :(...) || head == :parameters || head == :kw
        args[1] = from_expr(args[1], depth, callback, cbdata, top_level_number, false, read)
    elseif head == Symbol("'")
        for i = 1:length(args)
            args[i] = from_expr(args[i], depth, callback, cbdata, top_level_number, false, read)
        end
    elseif head==:break
        # skip
    elseif head==:continue
        # skip
    else
        throw(string("CompilerTools.AstWalker.from_expr: unknown Expr head :", head, " in ", ast))
    end
    ast.head = head
    ast.args = args
    ast.typ  = typ

    read = read_save

    return ast
end

# The following are for non-Expr AST nodes are generally leaf nodes of the AST where no
# recursive processing is possible.
function from_expr_helper(ast::Union{RHSVar,TopNode,LHSVar,Symbol},
                          depth,
                          callback,
                          cbdata::ANY,
                          top_level_number,
                          is_top_level,
                          read)
    @dprintln(2, typeof(ast), " type")
    # Intentionally do nothing.
    return ast
end

function from_expr_helper(ast::Union{LineNumberNode,LabelNode,GotoNode,DataType,AbstractString,NewvarNode,Void,Function,Module},
                          depth,
                          callback,
                          cbdata::ANY,
                          top_level_number,
                          is_top_level,
                          read)
    # Intentionally do nothing.
    return ast
end

function from_expr_helper(ast::Tuple,
                          depth,
                          callback,
                          cbdata::ANY,
                          top_level_number,
                          is_top_level,
                          read)

    # N.B. This also handles the empty tuple correctly.

    new_tt = Expr(:tuple)
    for i = 1:length(ast)
        push!(new_tt.args, from_expr(ast[i], depth, callback, cbdata, top_level_number, false, read))
    end
    new_tt.typ = asttyp
    ast = eval(new_tt)

    return ast
end

function from_expr_helper(ast::QuoteNode, depth,
                          callback,
                          cbdata::ANY,
                          top_level_number,
                          is_top_level,
                          read)
    value = ast.value
    #TODO: fields: value
    @dprintln(2,"QuoteNode type ",typeof(value))

    return ast
end

"""
The catchall function to process other kinds of AST nodes.
"""
function from_expr_helper(ast::ANY,
                          depth,
                          callback,
                          cbdata::ANY,
                          top_level_number,
                          is_top_level,
                          read)
    asttyp = typeof(ast)

    if isdefined(:GetfieldNode) && asttyp == GetfieldNode  # GetfieldNode = value + name
        @dprintln(2,"GetfieldNode type ",typeof(ast.value), " ", ast)
    elseif isdefined(:GlobalRef) && asttyp == GlobalRef
        @dprintln(2,"GlobalRef type ",typeof(ast.mod), " ", ast)  # GlobalRef = mod + name
    elseif isbits(asttyp)
        #skip
    elseif is(asttyp, LambdaInfo)
        #skip
    else
        println(ast, " type = ", typeof(ast), " asttyp = ", asttyp)
        throw(string("from_expr: unknown AST (", typeof(ast), ",", ast, ")"))
    end

    return ast
end

end


text:
#==
Functions related to calculating physical parameters of the ionospheric plasma
(susceptibility) relevant to the propagation of electromagnetic waves.
==#

@doc raw"""
    susceptibility(altitude, frequency, bfield, species; params=LMPParams())
    susceptibility(altitude, frequency, w::HomogeneousWaveguide; params=LMPParams())
    susceptibility(altitude, me::ModeEquation; params=LMPParams())

Compute the ionosphere susceptibility tensor `M` as a `SMatrix{3,3}` using
`species.numberdensity` and `species.collisionfrequency` at `altitude`.

Multiple species can be passed as an iterable. Use a `tuple` of `Species`, rather than a
`Vector`, for better performance.

If `params.earthcurvature == true`, `M` includes a first order correction for earth
curvature by means of a fictitious refractive index [Pappert1967].

The susceptibility matrix is calculated from the constitutive relations presented in
[Ratcliffe1959]. This includes the effect of earth's magnetic field vector and
collisional damping on electron motion.

The tensor is:
```math
M = -\frac{X}{U(U²-Y²)}
        \begin{pmatrix}
            U² - x²Y²    & -izUY - xyY² & iyUY - xzY² \\
            izUY - xyY²  & U² - y²Y²    & -ixUY - yzY² \\
            -iyUY - xzY² & ixUY - yzY²  & U² - z²Y²
        \end{pmatrix}
```
where ``X = ωₚ²/ω²``, ``Y = |ωₕ/ω|``, ``Z = ν/ω``, and ``U = 1 - iZ``. The earth curvature
correction subtracts ``2/Rₑ*(H - altitude)`` from the diagonal of ``M`` where ``H`` is
`params.curvatureheight`.

# References

[Pappert1967]: R. A. Pappert, E. E. Gossard, and I. J. Rothmuller, “A numerical
    investigation of classical approximations used in VLF propagation,” Radio Science,
    vol. 2, no. 4, pp. 387–400, Apr. 1967, doi: 10.1002/rds196724387.

[Ratcliffe1959]: J. A. Ratcliffe, "The magneto-ionic theory & its applications to the
    ionosphere," Cambridge University Press, 1959.
"""
function susceptibility(altitude, frequency, bfield, species; params=LMPParams())
    @unpack earthradius, earthcurvature, curvatureheight = params

    B, x, y, z = bfield.B, bfield.dcl, bfield.dcm, bfield.dcn
    ω = frequency.ω

    # Precompute constants (if multiple species)
    invω = inv(ω)
    invE0ω = invω/E0

    #== TODO:
    The zero type should be inferred instead of hard coded, but because we species N and nu
    are FunctionWrappers, we know the types will be ComplexF64.
    ==#
    U²D = zero(ComplexF64)
    Y²D = zero(ComplexF64)
    UYD = zero(ComplexF64)
    @inbounds for i in eachindex(species)
        X, Y, Z = _magnetoionicparameters(altitude, invω, invE0ω, bfield, species[i])

        U = 1 - 1im*Z
        U² = U^2
        Y² = Y^2

        D = -X/(U*(U² - Y²))

        U²D += U²*D
        Y²D += Y²*D
        UYD += U*Y*D
    end

    # Leverage partial symmetry of M to reduce computations
    izUYD = 1im*z*UYD
    xyY²D = x*y*Y²D
    iyUYD = 1im*y*UYD
    xzY²D = x*z*Y²D
    ixUYD = 1im*x*UYD
    yzY²D = y*z*Y²D

    # Elements of `M`
    M11 = U²D - x^2*Y²D
    M21 = izUYD - xyY²D
    M31 = -iyUYD - xzY²D
    M12 = -izUYD - xyY²D
    M22 = U²D - y^2*Y²D
    M32 = ixUYD - yzY²D
    M13 = iyUYD - xzY²D
    M23 = -ixUYD - yzY²D
    M33 = U²D - z^2*Y²D

    if earthcurvature
        curvaturecorrection = 2/earthradius*(curvatureheight - altitude)

        M11 -= curvaturecorrection
        M22 -= curvaturecorrection
        M33 -= curvaturecorrection
    end

    M = SMatrix{3,3}(M11, M21, M31, M12, M22, M32, M13, M23, M33)

    return M
end

susceptibility(altitude, me::ModeEquation; params=LMPParams()) =
    susceptibility(altitude, me.frequency, me.waveguide; params=params)

susceptibility(altitude, frequency, w::HomogeneousWaveguide; params=LMPParams()) =
    susceptibility(altitude, frequency, w.bfield, w.species; params=params)

"""
    susceptibilityspline(frequency, bfield, species; params=LMPParams())
    susceptibilityspline(frequency, w::HomogeneousWaveguide; params=LMPParams())
    susceptibilityspline(me::ModeEquation; params=LMPParams())

Construct a cubic interpolating spline of [`susceptibility`](@ref) and return the callable
`Interpolations` type.

`params.susceptibilitysplinestep` is the altitude step in meters used to construct the
spline between `BOTTOMHEIGHT` and `params.topheight`.
"""
function susceptibilityspline(frequency, bfield, species; params=LMPParams())
    zs = BOTTOMHEIGHT:params.susceptibilitysplinestep:params.topheight
    Ms = susceptibility.(zs, (frequency,), (bfield,), (species,))

    itp = CubicSplineInterpolation(zs, Ms)

    return itp
end

susceptibilityspline(me::ModeEquation; params=LMPParams()) =
    susceptibilityspline(me.frequency, me.waveguide; params=params)

susceptibilityspline(frequency, w::HomogeneousWaveguide; params=LMPParams()) =
    susceptibilityspline(frequency, w.bfield, w.species; params=params)


text:
module CirruParser
  export parse, pare

  include("Tree.jl")

  function parse(code, filename)
    buffer = Dict()

    state = Dict()
    state[:name] = :indent
    state[:x] = 1
    state[:y] = 1
    state[:level] = 1
    state[:indent] = 0
    state[:indented] = 0
    state[:nest] = 0
    state[:path] = filename

    res = parseRunner({}, buffer, state, code)
    res = map(Tree.resolveDollar, res)
    res = map(Tree.resolveComma, res)
    res
  end

  function shorten(xs)
    if typeof(xs) <: Array
      return map(shorten, xs)
    else
      return xs[:text]
    end
  end

  # eof

  function pare(code, filename)
    res = parse(code, filename)
    shorten(res)
  end

  function escape_eof(xs, buffer, state, code)
    error("EOF in escape state")
  end

  function string_eof(xs, buffer, state, code)
    error("EOF in string state")
  end

  function space_eof(xs, buffer, state, code)
    return xs
  end

  function token_eof(xs, buffer, state, code)
    buffer[:ex] = state[:x]
    buffer[:ey] = state[:y]
    xs = Tree.appendBuffer(xs, state[:level], buffer)
    buffer = Dict()
    xs
  end

  function indent_eof(xs, buffer, state, code)
    return xs
  end

  # escape

  function escape_newline(xs, buffer, state, code)
    error("new line whice escape")
  end

  function escape_n(xs, buffer, state, code)
    state[:x] += 1
    buffer[:text] *= "\n"
    state[:name] = :string
    parseRunner(xs, buffer, state, code[2:end])
  end

  function escape_t(xs, buffer, state, code)
    state[:x] += 1
    buffer.text *= "\t"
    state[:name] = :string
    parseRunner(xs, buffer, state, code[2:end])
  end

  function escape_else(xs, buffer, state, code)
    state[:x] += 1
    buffer[:text] *= string(code[1])
    state[:name] = :string
    parseRunner(xs, buffer, state, code[2:end])
  end

  # string

  function string_backslash(xs, buffer, state, code)
    state[:name] = :escape
    state[:x] += 1
    parseRunner(xs, buffer, state, code[2:end])
  end

  function string_newline(xs, buffer, state, code)
    error("newline in a string")
  end

  function string_quote(xs, buffer, state, code)
    state[:name] = :token
    state[:x] += 1
    parseRunner(xs, buffer, state, code[2:end])
  end

  function string_else(xs, buffer, state, code)
    state[:x] += 1
    buffer[:text] *= string(code[1])
    parseRunner(xs, buffer, state, code[2:end])
  end

  # space

  function space_space(xs, buffer, state, code)
    state[:x] += 1
    parseRunner(xs, buffer, state, code[2:end])
  end

  function space_newline(xs, buffer, state, code)
    if state[:nest] != 0
      error("incorrect nesting")
    end

    state[:name] = :indent
    state[:x] = 1
    state[:y] += 1
    state[:indented] = 0
    parseRunner(xs, buffer, state, code[2:end])
  end

  function space_open(xs, buffer, state, code)
    nesting = Tree.createNesting(1)
    xs = Tree.appendList(xs, state[:level], nesting)
    state[:nest] += 1
    state[:level] += 1
    state[:x] += 1
    parseRunner(xs, buffer, state, code[2:end])
  end

  function space_close(xs, buffer, state, code)
    state[:nest] -= 1
    state[:level] -= 1
    if state[:nest] < 0
      error("close at space")
    end
    state[:x] += 1
    parseRunner(xs, buffer, state, code[2:end])
  end

  function space_quote(xs, buffer, state, code)
    state[:name] = :string

    buffer = Dict()
    buffer[:text] = ""
    buffer[:x] = state[:x]
    buffer[:y] = state[:y]
    buffer[:path] = state[:path]

    state[:x] += 1
    parseRunner(xs, buffer, state, code[2:end])
  end

  function space_else(xs, buffer, state, code)
    state[:name] = :token

    buffer = Dict()
    buffer[:text] = string(code[1])
    buffer[:x] = state[:x]
    buffer[:y] = state[:y]
    buffer[:path] = state[:path]

    state[:x] += 1
    parseRunner(xs, buffer, state, code[2:end])
  end

  # token

  function token_space(xs, buffer, state, code)
    state[:name] = :space
    buffer[:ex] = state[:x]
    buffer[:ey] = state[:y]
    xs = Tree.appendBuffer(xs, state[:level], buffer)
    state[:x] += 1
    buffer = Dict()
    parseRunner(xs, buffer, state, code[2:end])
  end

  function token_newline(xs, buffer, state, code)
    state[:name] = :indent
    buffer[:ex] = state[:x]
    buffer[:ey] = state[:y]
    xs = Tree.appendBuffer(xs, state[:level], buffer)
    state[:indented] = 0
    state[:x] = 1
    state[:y] += 1
    buffer = Dict()
    parseRunner(xs, buffer, state, code[2:end])
  end

  function token_open(xs, buffer, state, code)
    error("open parenthesis in token")
  end

  function token_close(xs, buffer, state, code)
    state[:name] = :space
    buffer[:ex] = state[:x]
    buffer[:ey] = state[:y]
    xs = Tree.appendBuffer(xs, state[:level], buffer)
    buffer = Dict()
    parseRunner(xs, buffer, state, code)
  end

  function token_quote(xs, buffer, state, code)
    state[:name] = :string
    state[:x] += 1
    parseRunner(xs, buffer, state, code[2:end])
  end

  function token_else(xs, buffer, state, code)
    buffer[:text] *= string(code[1])
    state[:x] += 1
    parseRunner(xs, buffer, state, code[2:end])
  end

  # indent

  function indent_space(xs, buffer, state, code)
    state[:indented] += 1
    state[:x] += 1
    parseRunner(xs, buffer, state, code[2:end])
  end

  function indent_newline(xs, buffer, state, code)
    state[:x] = 1
    state[:y] += 1
    state[:indented] = 0
    parseRunner(xs, buffer, state, code[2:end])
  end

  function indent_close(xs, buffer, state, code)
    error("close parenthesis at indent")
  end

  function indent_else(xs, buffer, state, code)
    state[:name] = :space
    if (state[:indented] % 2) == 1
      error("odd indentation")
    end
    indented = state[:indented] / 2
    diff = indented - state[:indent]

    if diff <= 0
      nesting = Tree.createNesting(1)
      xs = Tree.appendList(xs, (state[:level] + diff - 1), nesting)
    elseif diff > 0
      nesting = Tree.createNesting(diff)
      xs = Tree.appendList(xs, state[:level], nesting)
    end

    state[:level] += diff
    state[:indent] = indented
    parseRunner(xs, buffer, state, code)
  end

  # parse

  function parseRunner(xs, buffer, state, code)
    args = {xs, buffer, state, code}
    if length(code) == 0
      eof = true
    else
      eof = false
      char = code[1]
    end

    if state[:name] == :escape
      if eof
        escape_eof(args...)
      elseif char == '\n'
        escape_newline(args...)
      elseif char == 'n'
        escape_n(args...)
      elseif char == 't'
        escape_t(args...)
      else
        escape_else(args...)
      end
    elseif state[:name] == :string
      if eof
        string_eof(args...)
      elseif char == '\\'
        string_backslash(args...)
      elseif char == '\n'
        string_newline(args...)
      elseif char == '"'
        string_quote(args...)
      else
        string_else(args...)
      end
    elseif state[:name] == :space
      if eof
        space_eof(args...)
      elseif char == ' '
        space_space(args...)
      elseif char == '\n'
        space_newline(args...)
      elseif char == '('
        space_open(args...)
      elseif char == ')'
        space_close(args...)
      elseif char =='"'
        space_quote(args...)
      else
        space_else(args...)
      end
    elseif state[:name] == :token
      if eof
        token_eof(args...)
      elseif char == ' '
        token_space(args...)
      elseif char == '\n'
        token_newline(args...)
      elseif char == '('
        token_open(args...)
      elseif char == ')'
        token_close(args...)
      elseif char == '"'
        token_quote(args...)
      else
        token_else(args...)
      end
    elseif state[:name] == :indent
      if eof
        indent_eof(args...)
      elseif char == ' '
        indent_space(args...)
      elseif char == '\n'
        indent_newline(args...)
      elseif char == ')'
        indent_close(args...)
      else
        indent_else(args...)
      end
    end
  end
end


text:
"""
FLOYDWARSHALL
-------------
    compute all shortest paths using the Floyd-Warshall algorithm.

    (D,P) = floydwarshall(A) returns the shortest distance matrix between all pairs
    of nodes in the graph A in matrix D.  If A has a negative weight cycle, then this
    algorithm will throw an error. P is the matrix of predecessors.

Functions
---------
- (D,P) = floydwarshall(A::MatrixNetwork)
- (D,P) = floydwarshall{T}(A::SparseMatrixCSC{T,Int64})

Example
-------
~~~
A = load_matrix_network("all_shortest_paths_example")
(D,P) = floydwarshall(A)
~~~
"""
:floydwarshall

## setup functions:

function floydwarshall_phase1(A::MatrixNetwork)
    (nzi,nzj,nzv) = csr_to_sparse(A.rp,A.ci,A.vals,A.n)
    return (nzi,nzj,nzv,A.n)
end

function floydwarshall_phase1{T}(A::SparseMatrixCSC{T,Int64})
    (ri,ci,ai) = findnz(A)
    return (ri,ci,ai,A.n)
end

function floydwarshall_phase2{T}(ri::Vector{Int64},ci::Vector{Int64},ai::Vector{T},n::Int64)

    nz = length(ai)
    D = Inf*ones(Int64,n,n)

    #TODO: check: always compute P or give the option of just computing D?

    P = zeros(Int64,n,n)
    # initialize the distance and predecessor matrix
    for ei = 1:nz
        i = ri[ei]
        j = ci[ei]
        v = ai[ei]
        if v < D[i,j]
            D[i,j] = v
            P[i,j] = i
        end
    end

    ids = sub2ind((n,n),1:n,1:n)
    D[ids] = 0 # set diagonal to 0

    for k=1:n
        for i=1:n
            for j=1:n
                if D[i,k]+D[k,j] < D[i,j]
                    D[i,j] = D[i,k]+D[k,j]
                    P[i,j] = P[k,j]
                end
            end
        end
    end

    if any(diag(D).<0)
        warn("floydwarshall:negativeCycle","negative weight cycle detected")
    end

    return (D,P)
end

## floyd warshall
function floydwarshall(A::MatrixNetwork)
    (nzi,nzj,nzv,n) = floydwarshall_phase1(A)
    (D,P) = floydwarshall_phase2(nzi,nzj,nzv,n)
    return (D,P)
end

function floydwarshall{T}(A::SparseMatrixCSC{T,Int64})
    (nzi,nzj,nzv,n) = floydwarshall_phase1(A)
    (D,P) = floydwarshall_phase2(nzi,nzj,nzv,n)
    return (D,P)
end


text:
# # Numerical Aspects
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/numerical.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/numerical.ipynb)

#md # !!! note "Note"
#md #     This section is still under development.

using RHEOS
#-

# ## Constant vs. Variable Sample Rate
# #### Weight indices rather than using a variable sample rate
# There may be situations when it is appropriate to focus the model fitting process on particular regions of the data. This may occur, for example, when the small time scale loading region is of particular interest. To achieve this, we could send data to fit that is sampled variably, with a greater sample density near the region of interest. However, the variable sampling rate means that the method RHEOS uses to compute the model predictions is slower and less accurate.

# A more optimal solution is to send constant sample rate data to the fitting function, but ask RHEOS to modify the final cost function to focus on specific points. In this way, RHEOS will be able to benefit from the much faster and more accurate constant sample rate model simulation methods.

# In concrete terms, consider the following array of time samples:

t = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5];

# We can send another array to the fitting process, `weightings`, that tells RHEOS to weight the first three time samples as doubly important, and the subsequent samples as half as important. The array we send would look like this:

weightings = [1, 1, 2, 2, 3, 3, 4, 6];

# In words, the elements in the time region we are particularly concerned about (0 ≤ t ≤ 0.2) are considered twice in the final fitting cost function, whereas we only consider every other element for times after 0.2 seconds.

# For practical situations, when working with data embedded in [`RheoTimeData`](@ref) structs, RHEOS has a convenience function [`indexweight`](@ref) which makes it easy to generate a weightings array that can then be passed to either [`modelfit`](@ref) or [`modelstepfit`](@ref). The full description of [`indexweight`](@ref) can be found in the [API](@ref) section. Below, we include a typical example of its use, based on the same data used in [Fitting and Predicting - Time Data](@ref).

## As in original example, load in the data
data = importcsv("assets/data_time.csv", t_col = 1, ϵ_col = 2, σ_col = 3)
rheotimedatatype(data)
#-

# Now instead of going straight to fitting, we will construct a weightings array using [`indexweight`](@ref).

## construct weighted array
weightings = indexweight(data; elperiods = [2, 3, -2], time_boundaries = [0.0, 0.5, 1.0, data.t[end]])

## fit model with weighted array
modelfit(data, Maxwell, strain_imposed, weights = weightings)

# We note that the results are slightly different from the original example in [Fitting and Predicting - Time Data](@ref) due to the element weighting provided to the fitting function.


text:
# syntax: proto3
using ProtoBuf
import ProtoBuf.meta
import ProtoBuf.google.protobuf

struct __enum_FieldBehavior <: ProtoEnum
    FIELD_BEHAVIOR_UNSPECIFIED::Int32
    OPTIONAL::Int32
    REQUIRED::Int32
    OUTPUT_ONLY::Int32
    INPUT_ONLY::Int32
    IMMUTABLE::Int32
    __enum_FieldBehavior() = new(0,1,2,3,4,5)
end #struct __enum_FieldBehavior
const FieldBehavior = __enum_FieldBehavior()

export FieldBehavior

max_stars_repo_name: sivakon/grpc-julia-gcp-example

max_issues_repo_name: sivakon/grpc-julia-gcp-example

max_forks_repo_name: sivakon/grpc-julia-gcp-example


text:
function train_layer_1_rbf_new(neurons,vars,obs,n_per_part,inds_all,x,y,cc)

    # a_all=zeros(neurons,1000vars)
    a_all=Vector{Vector{Float64}}()
    ooops=0
    for i=1:neurons
        println(Dates.format(now(), "HH:MM:SS")," Calculating weights w for ",i," neuron, of ",neurons," Total.")
        if n_per_part[i+1]-n_per_part[i]!=0
            try
                x1=copy(x[inds_all[n_per_part[i]+1:n_per_part[i+1]],:])
                obs1=size(x1,1)
                phi1 = calc_phi(x1,x1,obs1,obs1,cc,vars)
                atmp=phi1\y[inds_all[n_per_part[i]+1:n_per_part[i+1]]]
                if sum(isnan.(atmp))==0 && sum(isinf.(atmp))==0 push!(a_all,atmp) else push!(a_all,zeros(obs1)); ooops+=1 end
            catch ex
                ooops+=1; println(i," ",ex)
            end
        end
    end
    layer1=zeros(obs,neurons)
    for i=1:neurons
        println(Dates.format(now(), "HH:MM:SS")," Calculating output of neuron ",i," , of ",neurons," Total.")
        x1=copy(x[inds_all[n_per_part[i]+1:n_per_part[i+1]],:])
        obs1=size(x1,1)
        phi1 = calc_phi(x1,x,obs1,obs,cc,vars)
        layer1[:,i].=phi1*a_all[i]

        for ii=1:nn
            for jj=1:n
            end
        end
    end
    println(Dates.format(now(), "HH:MM:SS")," Solving system (",obs," by ",neurons+1,") for weights v in the output Layer")
    a_layer1=[layer1 ones(obs)]\y
    println(Dates.format(now(), "HH:MM:SS")," All weights computed.")

    return a_all,a_layer1,layer1

end


text:
using Documenter, PortHamiltonianBenchmarkSystems

push!(LOAD_PATH, "../src/")
makedocs(
 sitename="PortHamiltonianBenchmarkSystems",
 pages = [
          "Home" => "index.md",
          "Benchmark Systems" => [
                                  "GugercinMSDChain.md", "PoroModel.md", "RclCircuits.md"
            ]
         ]
)
deploydocs(
  repo = "github.com/Algopaul/PortHamiltonianBenchmarkSystems.git",
  versions = nothing
)


text:
############# Network with [W]⁺, [W]⁻ and layers that store their index ########

struct LayerNegPosIdx{F<:ActivationFunction, N<:Number}
    weights::Matrix{N}
    bias::Vector{N}
    W_neg::Matrix{N}
    W_pos::Matrix{N}
    activation::F
    index::Int64
end

LayerNegPosIdx(L::Layer, idx::Int64) = LayerNegPosIdx(L.weights, L.bias, min.(L.weights, 0), max.(L.weights, 0),
                                                           L.activation, idx)
n_nodes(L::LayerNegPosIdx) = length(L.bias)
affine_map(L::LayerNegPosIdx, x)  = L.weights*x .+ L.bias

struct NetworkNegPosIdx <: AbstractNetwork
    layers::Vector{LayerNegPosIdx}
end

NetworkNegPosIdx(net::Network) = NetworkNegPosIdx([LayerNegPosIdx(L, i) for (i, L) in enumerate(net.layers)])

function compute_output(nnet::NetworkNegPosIdx, input)
    curr_value = input
    for layer in nnet.layers
        curr_value = layer.activation(affine_map(layer, curr_value))
    end
    return curr_value
end


text:
using DecisionMakingEnvironments
import DecisionMakingEnvironments: create_simple_chain, create_simple_chain_finitetime, create_finitetime_cartpole
using Test
using Distributions

# create_finitetime_cartpole
@testset "Simple Chain Tests" begin
    prob = create_simple_chain(3, stochastic=false)
    @test length(prob.A) == 2
    @test eltype(prob.A) <: Int
    @test length(prob.S) == 3
    @test eltype(prob.S) <: Int
    @test prob.d0() isa Int
    @test prob.d0() == 1
    @test prob.p(prob.d0(),2) isa Int
    @test prob.p(1,2) == 2
    @test prob.p(3,2) == 1
    @test prob.r(1,2,1) isa Float64
    @test prob.r(1,2,1) == -1.0
    @test prob.γ(prob.d0(),2,1) isa Float64
    @test prob.γ(3,1,2) == 1.0
    @test prob.γ(3,2,1) == 0.0
    meta = meta_information(prob)
    @test meta[:minreturn] isa Float64
    @test meta[:minreward] == -1
    @test meta[:maxreturn] == -3
    @test meta[:minhorizon] == 3
    @test meta[:maxhorizon] == Inf
    prob = create_simple_chain(3, stochastic=true)
    @test prob.p(1,2) isa Int

    prob = create_simple_chain_finitetime(3, stochastic=false, droptime=false)
    maxT = 3*20
    @test length(prob.A) == 2
    @test eltype(prob.A) <: Int
    @test length(prob.S) == 2
    @test length(prob.S[1]) == maxT
    @test length(prob.S[2]) == 3
    @test eltype(prob.S[1]) <: Int
    @test eltype(prob.S[2]) <: Int
    s0 = prob.d0()
    @test s0 isa Tuple{Int,Int}
    @test s0[1] isa Int
    @test s0[2] isa Int
    @test s0[1] == 1
    @test s0[2] == 1
    @test prob.p(prob.d0(),2) isa Tuple{Int,Int}
    @test prob.p(s0,2) == (2,2)
    @test prob.p((3,3),2) == (1,1)
    @test prob.p((maxT-1,1),2) == (maxT,2)
    @test prob.p((maxT,1),2) == (1,1)
    @test prob.r(s0,2,(2,2)) isa Float64
    @test prob.r(s0,2,(2,2)) == -1.0
    @test prob.γ((maxT-1,1),2,(maxT,1)) isa Float64
    @test prob.γ((maxT-1,1),2,(maxT,1)) == 1.0
    @test prob.γ((maxT,1),2,(1,1)) == 0.0
    @test prob.γ((1,3),2,(2,1)) == 0.0
    @test prob.γ((1,2),2,(2,3)) == 1.0
    meta = meta_information(prob)
    @test meta[:minreturn] isa Float64
    @test meta[:minreturn] == -maxT
    @test meta[:minreward] == -1
    @test meta[:maxreward] == -1
    @test meta[:maxreturn] == -3
    @test meta[:minhorizon] == 3
    @test meta[:maxhorizon] == maxT

    prob2 = create_simple_chain_finitetime(3, stochastic=false, droptime=true)
    @test prob2.S == prob.S
    @test prob2.A == prob.A
    @test prob2.X == prob2.S[2]
    x = prob2.obs(s0)
    @test x == 1
    @test x isa Int
end

@testset "Simple Chain Interface Tests" begin
    prob = create_simple_chain(3, stochastic=false)
    @test prob.d0() isa Int
    @test prob.d0() == 1
    @test prob.p(prob.d0(),2) isa Int
    @test prob.p(1,2) == 2
    @test prob.p(3,2) == 1
    @test prob.r(1,2,1) isa Float64
    @test prob.r(1,2,1) == -1.0
    @test prob.γ(prob.d0(),2,1) isa Float64
    @test prob.γ(3,1,2) == 1.0
    @test prob.γ(3,2,1) == 0.0
    meta = meta_information(prob)
    @test meta[:minreturn] isa Float64
    @test meta[:minreward] == -1
    @test meta[:maxreturn] == -3
    @test meta[:minhorizon] == 3
    @test meta[:maxhorizon] == Inf
    prob = create_simple_chain(3, stochastic=true)
    @test prob.p(1,2) isa Int

    prob = create_simple_chain_finitetime(3, stochastic=false, droptime=false)
    maxT = 3*20
    @test length(prob.A) == 2
    @test eltype(prob.A) <: Int
    @test length(prob.S) == 2
    @test length(prob.S[1]) == maxT
    @test length(prob.S[2]) == 3
    @test eltype(prob.S[1]) <: Int
    @test eltype(prob.S[2]) <: Int
    s0 = prob.d0()
    @test s0 isa Tuple{Int,Int}
    @test s0[1] isa Int
    @test s0[2] isa Int
    @test s0[1] == 1
    @test s0[2] == 1
    @test prob.p(prob.d0(),2) isa Tuple{Int,Int}
    @test prob.p(s0,2) == (2,2)
    @test prob.p((3,3),2) == (1,1)
    @test prob.p((maxT-1,1),2) == (maxT,2)
    @test prob.p((maxT,1),2) == (1,1)
    @test prob.r(s0,2,(2,2)) isa Float64
    @test prob.r(s0,2,(2,2)) == -1.0
    @test prob.γ((maxT-1,1),2,(maxT,1)) isa Float64
    @test prob.γ((maxT-1,1),2,(maxT,1)) == 1.0
    @test prob.γ((maxT,1),2,(1,1)) == 0.0
    @test prob.γ((1,3),2,(2,1)) == 0.0
    @test prob.γ((1,2),2,(2,3)) == 1.0
    meta = meta_information(prob)
    @test meta[:minreturn] isa Float64
    @test meta[:minreturn] == -maxT
    @test meta[:minreward] == -1
    @test meta[:maxreward] == -1
    @test meta[:maxreturn] == -3
    @test meta[:minhorizon] == 3
    @test meta[:maxhorizon] == maxT

    prob2 = create_simple_chain_finitetime(3, stochastic=false, droptime=true)
    @test prob2.S == prob.S
    @test prob2.A == prob.A
    @test prob2.X == prob2.S[2]
    x = prob2.obs(s0)
    @test x == 1
    @test x isa Int
end

@testset "Cart Pole Interface Tests" begin
    tMax = 20.0
    dt = 0.02
    prob = create_finitetime_cartpole(randomize=false, tMax=tMax, dt=dt, Atype=:Discrete, droptime=true)
    s0,x0 = prob.d0()
    @test s0 isa Tuple{Float64,Array{Float64,1}}
    @test all(x0 .== zeros(4))
    s,x,r,γ = sample(prob, s0, 1)
    @test s isa Tuple{Float64,Array{Float64,1}}
    @test x isa Array{Float64,1}
    @test r == 1.0
    @test γ isa Float64
    @test γ == 1.0

    s,x,r,γ = sample(prob, (20.0, zeros(4)), 2)
    @test γ == 0.0
    meta = meta_information(prob)
    @test meta[:minreturn] isa Float64
    @test meta[:minreturn] == 9.0
    @test meta[:minreward] == 1.0
    @test meta[:maxreturn] == ceil(tMax / dt)
    @test meta[:minhorizon] == 9
    @test meta[:maxhorizon] == ceil(Int, tMax / dt)
    @test meta[:discounted] == false

    π = s -> Categorical([1.0,0.0])
    τ = Trajectory(prob)
    sample_trajectory!(prob, τ, π)
    @test sum(τ.rewards) == meta[:minreturn]
    @test length(τ.rewards) == 9
    @test length(τ.states) == 9
    @test eltype(τ.states) == Array{Float64,1}
    @test eltype(τ.actions) == Int
    @test length(τ.actions) == 9
    @test length(τ.blogps) == 9
    @test eltype(τ.blogps) == Float64

    J = sample_objective(prob, π)
    @test J == 9.0
    @test J isa Float64
    @test J ≥ meta[:minreturn]
    @test J ≤ meta[:maxreturn]

end


text:
n = rand(1:7)
A = matrixdepot("invol", n)
# A*A \approx eye(n)
@test_approx_eq_eps A*A eye(n,n) 1e-5
println("'invol' passed test...")


text:
using DimensionalData, Test, Unitful, Adapt

struct CustomArray{T,N} <: AbstractArray{T,N}
    arr::Array
end

CustomArray(x::Array{T,N}) where {T,N} = CustomArray{T,N}(x)
Adapt.adapt_storage(::Type{<:CustomArray}, xs::Array) = CustomArray(xs)

Base.size(x::CustomArray, y...) = size(x.arr, y...)
Base.getindex(x::CustomArray, y...) = getindex(x.arr, y...)
Base.count(x::CustomArray) = count(x.arr)

@testset "Metadata" begin
    @test adapt(CustomArray, Metadata(:a=>"1", :b=>"2")) == NoMetadata()
end

@testset "Dimension" begin
    d = X([1:10...]; metadata=Metadata(:a=>"1", :b=>"2"))
    d1 = adapt(CustomArray, d)
    @test val(d1) isa CustomArray
    @test val(d1).arr == [1:10...]
    @test metadata(d1) == NoMetadata()
end

@testset "DimArray" begin
    A = rand(4, 5)
    da = DimArray(A, (X, Y))
    da1 = adapt(CustomArray, da)
    @test parent(da1) isa CustomArray
    @test parent(da1).arr == A
    @test metadata(da1) == NoMetadata()
end

@testset "DimStack" begin
    A = rand(4, 5)
    B = rand(4, 5)
    ds = DimStack((a=A, b=B), (X, Y))
    ds1 = adapt(CustomArray, ds)
    @test parent(ds1[:a]) isa CustomArray
    @test parent(ds1[:a]).arr == A
    @test parent(ds1[:b]) isa CustomArray
    @test parent(ds1[:b]).arr == B
    @test metadata(ds) == NoMetadata()
end


text:
const δθ = π / 18
const π_third = π / 3
const π_fourth = π / 4
const π_sixth = π / 6

"""
    η_B(θ_B::Real, θ_crit::Real)

Calculate B angle dependent efficiency component following Pais et. al. 2018, eq. 7, https://arxiv.org/pdf/1805.00128.pdf
"""
@inline function η_B(θ_B::Real, θ_crit::Real)
    (tanh((θ_crit - θ_B) / δθ) + 1) / 2
end

"""
    ηB_acc_p(θ_B::Real)

Magnetic field geometry dependent efficiency for protons at initial acceleration.
"""
function ηB_acc_p(θ_B::Real)
    η_B(θ_B, π_fourth)
end

"""
    ηB_reacc_p(θ_B::Real)

Magnetic field geometry dependent efficiency for protons at reacceleration.
"""
function ηB_reacc_p(θ_B::Real)
    η_B(θ_B, π_third)
end

"""
    ηB_acc_e(θ_B::Real)

Magnetic field geometry dependent efficiency for electrons at initial acceleration.
"""
function ηB_acc_e(θ_B::Real)
    1 - η_B(θ_B, π_fourth)
end

"""
    ηB_reacc_e(θ_B::Real)

Magnetic field geometry dependent efficiency for electrons at reacceleration.
"""
function ηB_reacc_e(θ_B::Real)
    1 - η_B(θ_B, π_sixth)
end


text:
using MPI
using ClimateMachine
using Logging
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.DGMethods
using ClimateMachine.BalanceLaws: update_auxiliary_state!
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear
using ClimateMachine.MPIStateArrays
using ClimateMachine.ODESolvers
using ClimateMachine.Atmos: SphericalOrientation, latitude, longitude
using ClimateMachine.Orientations
using LinearAlgebra
using Printf
using Dates
using ClimateMachine.GenericCallbacks:
    EveryXWallTimeSeconds, EveryXSimulationSteps
using ClimateMachine.VTK: writevtk, writepvtu
import ClimateMachine.BalanceLaws: boundary_state!
using CLIMAParameters.Planet: planet_radius

if !@isdefined integration_testing
    const integration_testing = parse(
        Bool,
        lowercase(get(ENV, "JULIA_CLIMA_INTEGRATION_TESTING", "false")),
    )
end

const output = parse(Bool, lowercase(get(ENV, "JULIA_CLIMA_OUTPUT", "false")))

include("advection_diffusion_model.jl")

# This is a setup similar to the one presented in [Williamson1992](@cite)
struct SolidBodyRotation <: AdvectionDiffusionProblem end
function init_velocity_diffusion!(
    ::SolidBodyRotation,
    aux::Vars,
    geom::LocalGeometry,
)
    FT = eltype(aux)
    λ = longitude(SphericalOrientation(), aux)
    φ = latitude(SphericalOrientation(), aux)
    r = norm(geom.coord)

    uλ = 2 * FT(π) * cos(φ) * r
    uφ = 0
    aux.advection.u = SVector(
        -uλ * sin(λ) - uφ * cos(λ) * sin(φ),
        +uλ * cos(λ) - uφ * sin(λ) * sin(φ),
        +uφ * cos(φ),
    )
end
function initial_condition!(::SolidBodyRotation, state, aux, localgeo, t)
    λ = longitude(SphericalOrientation(), aux)
    φ = latitude(SphericalOrientation(), aux)
    state.ρ = exp(-((3λ)^2 + (3φ)^2))
end
finaltime(::SolidBodyRotation) = 1
u_scale(::SolidBodyRotation) = 2π

"""
This is the Divergent flow with smooth initial condition test case, the Case 4 in

@article{nair2010class,
  title={A class of deformational flow test cases for linear transport problems on the sphere},
  author={Nair, Ramachandran D and Lauritzen, Peter H},
  journal={Journal of Computational Physics},
  volume={229},
  number={23},
  pages={8868--8887},
  year={2010},
  publisher={Elsevier}
}

"""
struct ReversingDeformationalFlow <: AdvectionDiffusionProblem end
init_velocity_diffusion!(
    ::ReversingDeformationalFlow,
    aux::Vars,
    geom::LocalGeometry,
) = nothing
function initial_condition!(::ReversingDeformationalFlow, state, aux, coord, t)
    x, y, z = aux.coord
    r = norm(aux.coord)
    h_max = 0.95
    b = 5
    state.ρ = 0
    for (λ, φ) in ((5π / 6, 0), (7π / 6, 0))
        xi = r * cos(φ) * cos(λ)
        yi = r * cos(φ) * sin(λ)
        zi = r * sin(φ)
        state.ρ +=
            h_max * exp(-b * ((x - xi)^2 + (y - yi)^2 + (z - zi)^2) / r^2)
    end
end
has_variable_coefficients(::ReversingDeformationalFlow) = true
function update_velocity_diffusion!(
    ::ReversingDeformationalFlow,
    ::AdvectionDiffusion,
    state::Vars,
    aux::Vars,
    t::Real,
)
    FT = eltype(aux)
    λ = longitude(SphericalOrientation(), aux)
    φ = latitude(SphericalOrientation(), aux)
    r = norm(aux.coord)
    T = FT(5)
    λp = λ - FT(2π) * t / T
    uλ =
        10 * r / T * sin(λp)^2 * sin(2φ) * cos(FT(π) * t / T) +
        FT(2π) * r / T * cos(φ)
    uφ = 10 * r / T * sin(2λp) * cos(φ) * cos(FT(π) * t / T)
    aux.advection.u = SVector(
        -uλ * sin(λ) - uφ * cos(λ) * sin(φ),
        +uλ * cos(λ) - uφ * sin(λ) * sin(φ),
        +uφ * cos(φ),
    )
end
u_scale(::ReversingDeformationalFlow) = 2.9
finaltime(::ReversingDeformationalFlow) = 5

function advective_courant(
    m::AdvectionDiffusion,
    state::Vars,
    aux::Vars,
    diffusive::Vars,
    Δx,
    Δt,
    direction,
)
    return Δt * norm(aux.advection.u) / Δx
end

struct NoFlowBC end
function boundary_state!(
    ::RusanovNumericalFlux,
    ::NoFlowBC,
    ::AdvectionDiffusion,
    stateP::Vars,
    auxP::Vars,
    nM,
    stateM::Vars,
    auxM::Vars,
    t,
    _...,
)
    auxP.advection.u = -auxM.advection.u
end

function do_output(mpicomm, vtkdir, vtkstep, dgfvm, Q, Qe, model, testname)
    ## Name of the file that this MPI rank will write
    filename = @sprintf(
        "%s/%s_mpirank%04d_step%04d",
        vtkdir,
        testname,
        MPI.Comm_rank(mpicomm),
        vtkstep
    )

    statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))
    exactnames = statenames .* "_exact"

    writevtk(filename, Q, dgfvm, statenames, Qe, exactnames)

    ## Generate the pvtu file for these vtk files
    if MPI.Comm_rank(mpicomm) == 0
        ## Name of the pvtu file
        pvtuprefix = @sprintf("%s/%s_step%04d", vtkdir, testname, vtkstep)

        ## Name of each of the ranks vtk files
        prefixes = ntuple(MPI.Comm_size(mpicomm)) do i
            @sprintf("%s_mpirank%04d_step%04d", testname, i - 1, vtkstep)
        end

        writepvtu(
            pvtuprefix,
            prefixes,
            (statenames..., exactnames...),
            eltype(Q),
        )

        @info "Done writing VTK: $pvtuprefix"
    end
end

function test_run(
    mpicomm,
    ArrayType,
    vert_range,
    topl,
    problem,
    explicit_method,
    fvmethod,
    cfl,
    N,
    timeend,
    FT,
    vtkdir,
    outputtime,
)
    grid = DiscontinuousSpectralElementGrid(
        topl,
        FloatType = FT,
        DeviceArray = ArrayType,
        polynomialorder = (N, 0),
        meshwarp = equiangular_cubed_sphere_warp,
    )

    dx = min_node_distance(grid, HorizontalDirection())
    dt = FT(cfl * dx / (vert_range[2] * u_scale(problem())) / N)
    dt = outputtime / ceil(Int64, outputtime / dt)

    bcs = (NoFlowBC(),)
    model = AdvectionDiffusion{3}(problem(), bcs, diffusion = false)
    dgfvm = DGFVModel(
        model,
        grid,
        fvmethod,
        RusanovNumericalFlux(),
        CentralNumericalFluxSecondOrder(),
        CentralNumericalFluxGradient(),
    )

    Q = init_ode_state(dgfvm, FT(0))

    odesolver = explicit_method(dgfvm, Q; dt = dt, t0 = 0)

    eng0 = norm(Q)
    @info @sprintf """Starting
    problem           = %s
    ArrayType         = %s
    FV Reconstruction = %s
    method            = %s
    time step         = %.16e
    norm(Q₀)          = %.16e""" problem ArrayType fvmethod explicit_method dt eng0

    # Set up the information callback
    starttime = Ref(now())
    cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)
        if s
            starttime[] = now()
        else
            energy = norm(Q)
            @info @sprintf(
                """Update
                simtime = %.16e
                runtime = %s
                norm(Q) = %.16e""",
                gettime(odesolver),
                Dates.format(
                    convert(Dates.DateTime, Dates.now() - starttime[]),
                    Dates.dateformat"HH:MM:SS",
                ),
                energy
            )
        end
    end
    cbcfl = EveryXSimulationSteps(10) do
        dt = ODESolvers.getdt(odesolver)
        cfl = DGMethods.courant(
            advective_courant,
            dgfvm,
            model,
            Q,
            dt,
            HorizontalDirection(),
        )
        @info @sprintf(
            """Courant number
            simtime = %.16e
            courant = %.16e""",
            gettime(odesolver),
            cfl
        )
    end
    callbacks = (cbinfo,)
    if ~isnothing(vtkdir)
        # Create vtk dir
        mkpath(vtkdir)

        vtkstep = 0
        # Output initial step
        do_output(
            mpicomm,
            vtkdir,
            vtkstep,
            dgfvm,
            Q,
            Q,
            model,
            "fvm_advection_sphere",
        )

        # Setup the output callback
        cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do
            vtkstep += 1
            Qe = init_ode_state(dgfvm, gettime(odesolver))
            do_output(
                mpicomm,
                vtkdir,
                vtkstep,
                dgfvm,
                Q,
                Qe,
                model,
                "fvm_advection_sphere",
            )
        end
        callbacks = (callbacks..., cbvtk)
    end

    solve!(Q, odesolver; timeend = timeend, callbacks = callbacks)

    # Print some end of the simulation information
    engf = norm(Q)
    Qe = init_ode_state(dgfvm, FT(timeend))

    engfe = norm(Qe)
    errf = euclidean_distance(Q, Qe)
    Δmass = abs(weightedsum(Q) - weightedsum(Qe)) / weightedsum(Qe)
    @info @sprintf """Finished
    Δmass                   = %.16e
    norm(Q)                 = %.16e
    norm(Q) / norm(Q₀)      = %.16e
    norm(Q) - norm(Q₀)      = %.16e
    norm(Q - Qe)            = %.16e
    norm(Q - Qe) / norm(Qe) = %.16e
    """ Δmass engf engf / eng0 engf - eng0 errf errf / engfe
    return errf, Δmass
end

using Test
let
    ClimateMachine.init()
    ArrayType = ClimateMachine.array_type()

    mpicomm = MPI.COMM_WORLD

    polynomialorder = 4
    base_num_elem = 3

    max_cfl = Dict(LSRK144NiegemannDiehlBusch => 5.0)

    expected_result = Dict()

    expected_result[SolidBodyRotation, 1, FVConstant] = 1.6249678501611961e+07
    expected_result[SolidBodyRotation, 2, FVConstant] = 7.2020207047738554e+05
    expected_result[SolidBodyRotation, 3, FVConstant] = 5.2452627634607365e+04
    expected_result[SolidBodyRotation, 4, FVConstant] = 3.1132403618328990e+03

    expected_result[SolidBodyRotation, 1, FVLinear] = 1.6649963041466445e+07
    expected_result[SolidBodyRotation, 2, FVLinear] = 7.2518593691652094e+05
    expected_result[SolidBodyRotation, 3, FVLinear] = 5.2473203187596591e+04
    expected_result[SolidBodyRotation, 4, FVLinear] = 3.1133127473480104e+03

    expected_result[ReversingDeformationalFlow, 1, FVConstant] =
        2.0097347028034222e+08
    expected_result[ReversingDeformationalFlow, 2, FVConstant] =
        7.0092390520693690e+07
    expected_result[ReversingDeformationalFlow, 3, FVConstant] =
        7.7527847763998993e+06
    expected_result[ReversingDeformationalFlow, 4, FVConstant] =
        1.4209138735343830e+05

    expected_result[ReversingDeformationalFlow, 1, FVLinear] =
        2.0156864805489743e+08
    expected_result[ReversingDeformationalFlow, 2, FVLinear] =
        7.0092714938572392e+07
    expected_result[ReversingDeformationalFlow, 3, FVLinear] =
        7.7527849036761308e+06
    expected_result[ReversingDeformationalFlow, 4, FVLinear] =
        1.4209138776027536e+05

    numlevels =
        integration_testing || ClimateMachine.Settings.integration_testing ? 4 :
        1
    explicit_method = LSRK144NiegemannDiehlBusch
    @testset "$(@__FILE__)" begin
        for FT in (Float64,)
            for problem in (SolidBodyRotation, ReversingDeformationalFlow)
                for fvmethod in (FVConstant, FVLinear)
                    cfl = max_cfl[explicit_method]
                    result = zeros(FT, numlevels)
                    for l in 1:numlevels
                        numelems_horizontal = 2^(l - 1) * base_num_elem
                        numelems_vertical = 3

                        _planet_radius = FT(planet_radius(param_set))
                        domain_height = FT(10e3)
                        vert_range =
                            (_planet_radius, _planet_radius + domain_height)

                        topl = StackedCubedSphereTopology(
                            mpicomm,
                            numelems_horizontal,
                            range(
                                vert_range[1],
                                stop = vert_range[2],
                                length = numelems_vertical + 1,
                            ),
                        )

                        timeend = finaltime(problem())
                        outputtime = timeend

                        @info (ArrayType, FT)
                        vtkdir =
                            output ?
                            "vtk_fvm_advection_sphere" *
                            "_$problem" *
                            "_$explicit_method" *
                            "_poly$(polynomialorder)" *
                            "_$(ArrayType)_$(FT)" *
                            "_level$(l)" :
                            nothing

                        result[l], Δmass = test_run(
                            mpicomm,
                            ArrayType,
                            vert_range,
                            topl,
                            problem,
                            explicit_method,
                            fvmethod(),
                            cfl,
                            polynomialorder,
                            timeend,
                            FT,
                            vtkdir,
                            outputtime,
                        )
                        @test result[l] ≈
                              FT(expected_result[problem, l, fvmethod])
                        @test Δmass <= FT(1e-12)
                    end
                    @info begin
                        msg = ""
                        for l in 1:(numlevels - 1)
                            rate = log2(result[l]) - log2(result[l + 1])
                            msg *= @sprintf(
                                "\n  rate for level %d = %e\n",
                                l,
                                rate
                            )
                        end
                        msg
                    end
                end
            end
        end
    end
end


text:
module Expronicon

using MLStyle
using MLStyle.MatchImpl
using MLStyle.AbstractPatterns

export
    # types
    NoDefault, JLExpr, JLFor, JLIfElse,
    JLFunction, JLField, JLKwField, JLStruct, JLKwStruct,
    # analysis
    @expr, @test_expr, compare_expr, compare_vars,
    AnalysisError, is_function, is_kw_function, is_struct,
    is_ifelse, is_for, is_field, is_field_default,
    split_function, split_function_head, split_struct,
    split_struct_name, split_ifelse, annotations,
    uninferrable_typevars, has_symbol,
    is_literal, is_gensym,
    alias_gensym,
    has_kwfn_constructor,
    has_plain_constructor,
    # transformations
    no_default, prettify, rm_lineinfo, flatten_blocks, name_only,
    rm_annotations, rm_single_block, rm_nothing,
    replace_symbol, subtitute, eval_interp, eval_literal,
    # codegen
    codegen_ast,
    codegen_ast_kwfn,
    codegen_ast_kwfn_plain,
    codegen_ast_kwfn_infer,
    codegen_ast_struct,
    codegen_ast_struct_head,
    codegen_ast_struct_body,
    construct_method_plain,
    construct_method_infer,
    struct_name_plain,
    struct_name_without_inferable,
    # x functions
    xtuple,
    xnamedtuple,
    xcall,
    xpush,
    xgetindex,
    xfirst,
    xlast,
    xprint,
    xprintln,
    xmap,
    xmapreduce,
    xiterate,
    # printings
    print_expr, sprint_expr

include("patches.jl")
include("types.jl")
include("transform.jl")
include("analysis.jl")
include("codegen.jl")
include("printing.jl")
include("match.jl")
include("expand.jl")

end


text:
using LabMaze
using Test

@testset "Connections" begin
    cell1 = Cell(Rect, ; row=3, col=4)
    cell2 = Cell(Rect, 3, 5)
    connect!(cell1, cell2)
    @test is_connected(cell1, cell2)
    @test length(cell1.connections) == 1
    @test length(cell2.connections) == 1
    disconnect!(cell1, cell2)
    @test !is_connected(cell1, cell2)
end

@testset "Triangle" begin
    nrows = 3
    ncols = 4
    grid = TriaGrid(nrows, ncols)
    @test grid.nrows == nrows
    @test grid.ncols == ncols
    periodic_grid = PeriodicTriaGrid(nrows, ncols)
    # periodic triagonal grids have to have even ncols
    @test_throws AssertionError PeriodicTriaGrid(nrows, 3)
    @test grid.nrows == nrows
    @test grid.ncols == ncols
end

@testset "Rectangle" begin
    nrows = 3
    ncols = 4
    grid = RectGrid(nrows, ncols)
    @test grid.nrows == nrows
    @test grid.ncols == ncols
    periodic_grid = PeriodicRectGrid(nrows, ncols)
    @test grid.nrows == nrows
    @test grid.ncols == ncols
end

@testset "Hexagon" begin
    nrows = 3
    ncols = 4
    grid = HexaGrid(nrows, ncols)
    @test grid.nrows == nrows
    @test grid.ncols == ncols
    periodic_grid = PeriodicHexaGrid(nrows, ncols)
    # periodic hexagonal grids have to have even ncols
    @test_throws AssertionError PeriodicHexaGrid(nrows, 3)
    @test grid.nrows == nrows
    @test grid.ncols == ncols
end

@testset "Deadends" begin
    nrows = 3
    ncols = 4
    grid = HexaGrid(nrows, ncols)
    # there should be no deadends in the beginning
    @test length(get_deadends(grid)) == 0
    recursive_backtracker!(grid)
    n_deadends = length(get_deadends(grid))
    remove_deadends!(grid, 0.8)
    new_n_deadends = length(get_deadends(grid))
    @assert new_n_deadends <= n_deadends
end


text:
@testset "Hooks for Core integration" begin
    JuliaSyntax.enable_in_core!()

    @test Meta.parse("x + 1") == :(x + 1)
    @test Meta.parse("x + 1", 1) == (:(x + 1), 6)

    # Test that parsing statements incrementally works
    @test Meta.parse("x + 1\n(y)", 1) == (:(x + 1), 6)
    @test Meta.parse("x + 1\n(y)", 6) == (:y, 10)

    # Check that Meta.parse throws the JuliaSyntax.ParseError rather than
    # Meta.ParseError when Core integration is enabled.
    @test_throws JuliaSyntax.ParseError Meta.parse("[x")

    JuliaSyntax.disable_in_core!()
end


text:
mutable struct Parameters{I, F}
      H::F            # meters               Physical
     Ly::F            # meters
     Lϕ::F            # radians
   ϕmin::F            # radials
   ϕmid::F            # radians
   ϕmax::F            # radials
      a::F            # meters

   TwoΩ::F            # 1/second             Coriolis
     f₀::F            # 1/second

      g::F            # meters/second^2      Buoyancy

      N::I            # Integer              Numerical

     Lj::F            # meters               Jet (specify profile)
     Uj::F            # meters/second

     dk::F            # 1/meters             Wavenumbers
   kmin::F            # 1/metres
   kmax::F            # 1/meters

     Fr::F            # Nondimensional       Nondimensional
     Ro::F            # Nondimensional
end

"""

Cartesian version of parameters function.

"""
function Parameters(geometry::Cartesian;
     H = 1,
    Ly = 20,
    Lϕ = nothing,
  ϕmin = nothing,
  ϕmid = nothing,
  ϕmax = nothing,
     a = nothing,

  TwoΩ = nothing,
    f₀ = 1,

     g = 9.81,

     N = 100,

    Lj = 1,
    Uj = 0.1,

    dk = 1e-1,
  kmin = dk,
  kmax = 2,

    Fr = 0.0,
    Ro = 0.0
    )
    return Parameters(H, Ly, f₀, g, N, Lj, Uj, N, dk, kmin, kmax, Fr, Ro)
end

"""

Spherical version of parameters function.

"""
function Parameters(geometry::Spherical;
     H = 1,
    Ly = nothing,
    Lϕ =  π/6,
  ϕmin =  π/8,
  ϕmid =  π/4,
  ϕmax = 3π/8,
     a = 1,

  TwoΩ = 1,
    f₀ = nothing,

     g = 9.81,

     N = 100,

    Lj = 1,
    Uj = 0.1,

    dk = 1e-1,
  kmin = dk,
  kmax = 2,

    Fr = 0.0,
    Ro = 0.0
  )
  return Parameters(H, Lϕ, ϕmin, ϕmid, ϕmax, a, TwoΩ, g, N, Lj, Uj, dk, kmin, kmax, Fr, Ro)
end


text:
export GenericCircuit, Circuit

function evaluate end

export PrimitiveCircuit
# Primitive Routines
struct PrimitiveCircuit{name} end

# avoid splatting
evaluate(c::PrimitiveCircuit) = c()
evaluate(c::PrimitiveCircuit, x) = c(x)
evaluate(c::PrimitiveCircuit, x, y) = c(x, y)
evaluate(c::PrimitiveCircuit, x, y, xs...) = c(x, y, xs...)

function Base.show(io::IO, x::PrimitiveCircuit{name}) where {name}
    print(io, name, " (primitive circuit)")
end

"""
    GenericCircuit{name}

Generic quantum circuit is the quantum counterpart of generic function.
"""
struct GenericCircuit{name} end

function Base.show(io::IO, x::GenericCircuit{name}) where {name}
    print(io, name, " (generic circuit with ", length(methods(x).ms), " methods)")
end

struct Circuit{name,F<:Function,Free<:Tuple}
    fn::Function
    free::Free
end

Circuit{name}(fn::Function, free::Tuple) where {name} =
    Circuit{name,typeof(fn),typeof(free)}(fn, free)
Circuit{name}(fn::Function) where {name} = Circuit{name}(fn, ())

function Base.show(io::IO, x::Circuit{name}) where {name}
    print(io, name, " (quantum circuit)")
end

# syntax sugar
(circ::Circuit)(r::AbstractRegister) = circ(r, 1:nactive(r))
(circ::Circuit)(locs) = r -> (circ(r, locs); r)
(circ::Circuit)(locs, ctrl_locs) = r -> (circ(r, locs, ctrl_locs); r)

function Base.:(|>)(r::AbstractRegister, circ::Circuit)
    circ(r)
    return r
end

# we only convert to Locations right before we call the stubs
(circ::Circuit)(register::AbstractRegister, locs) = circ.fn(circ, register, Locations(locs))
(circ::Circuit)(register::AbstractRegister, locs, ctrl_locs) =
    circ.fn(circ, register, Locations(locs), CtrlLocations(ctrl_locs))

# allow primitive circuits to be called directly
(c::PrimitiveCircuit)(r::AbstractRegister, locs) = c()(r, locs)
(c::PrimitiveCircuit)(r::AbstractRegister) = c()(r)
(c::PrimitiveCircuit)(locs::Locations) = r -> c(r, locs)
(c::PrimitiveCircuit)(locs::Locations, ctrl_locs::CtrlLocations) = r -> c(r, locs, ctrl_locs)


text:
# create the Dict for datas, but without any calculations
function data_djEs_sketch(at::Atoms, h::Float64, cutoffs::Array{Float64,1})

    # X  = positions(at) |> mat;
    atd = deepcopy(at);
    # neighbour lists for different cutoffs
    rcut1, rcut2, rcut3 = cutoffs[:];
    nlist_d1 = neighbourlist(at, rcut1);
    nlist_d2 = neighbourlist(at, rcut2);
    nlist_d3 = neighbourlist(at, rcut3);

    # calculators
    # eF = 10.89102; for Carbon and 5.521241 for Si
    # tbm = NRLTB.NRLTBModel(:C or Si,
    #                ZeroTGrand(eF = eF),
    #                bzquad = TB.GammaPoint(),
    #                cutoff = :original );

    # data
    D = Dict{String, Any}();
    D["species"] = JuLIP.Chemistry.chemical_symbol(at.Z[1]);
    D["at"] = Dict(atd);
    D["Info"] = "djEs type training set";
    D["cutoff_d"] = rcut1
    D["cutoff_d2"] = rcut2;
    D["cutoff_d3"] = rcut3;
    data = Dict[]
    D["data"] = data;

    # 0. compute Es on reference
    # println("on reference: calculate Es")
    l0 = 1;
    ### Es = site_energy(tbm, atd, l0);
    # store Es
    dat = Dict{String, Any}();
    dat["Info"] = "Es of the origin site";
    ### dat["Es"] = Es;
    dat["datatype"] = "Es"
    push!(data, dat);

    # 1. compute dEs on reference and only store the derivatives
    # println("calculating dEs_", l0)
    ### dEs_ref = site_energy_d(tbm, atd, l0) |> mat
    # store dEs
    dat = Dict{String, Any}();
    dat["Info"] = "dEs of the origin site";
### dat["dEs"] = dEs_ref;
    dat["datatype"] = "dEs"
    push!(data, dat);

    # 2. compute FD approximations for dE2 and only store the FD within rcut_d2
    #    d2Esh = 1/2h * ( E_{ℓ,n}(y+h⋅e0) - E_{ℓ,n}(y-h⋅e0) )
    index_rcut2_0 = [];
    for (n, neigs, R) in sites(nlist_d2)
        if n == l0
            index_rcut2_0 = neigs;
        end
    end
    N₂ = length(index_rcut2_0);
    d2Esh = zeros(3, N₂, 3, length(atd));
    dEs_pert_0p = zeros(3, N₂, 3, length(atd));
    dEs_pert_0n = zeros(3, N₂, 3, length(atd));
    for i = 1 : 3
        # println("perturb the origin atom in direction ", i, "+h :")
        # X[i,l0] += h;
        # atd = deepcopy(at);
        # set_positions!(atd, X);
        for j = 1 : N₂
            ℓ = index_rcut2_0[j];
            # println("calculating dE_", ℓ)
###         dEs_pert_0p[i, j, :, :] = site_energy_d(tbm, atd, ℓ) |> mat;
        end
        # println("perturb the origin atom in direction ", i, "-h :")
        # X[i,l0] -= 2*h;
        # atd = deepcopy(at);
        # set_positions!(atd, X);
        for j = 1 : N₂
            ℓ = index_rcut2_0[j];
            # println("calculating dE_", ℓ)
###         dEs_pert_0n[i, j, :, :] = site_energy_d(tbm, atd, ℓ) |> mat;
        end
        # X[i,l0] += h;
        # compute the central finite differences
        for j = 1 : N₂
            ℓ = index_rcut2_0[j];
            index_rcut2_ℓ = [];
            # # TODO: make sure we really don't need this!
            # for (k, neigs, r, R) in sites(nlist_d2)
            #     if k == ℓ
            #         index_rcut2_ℓ = neigs;
            #     end
            # end
###         d2Esh[i, j, :, :] = 1.0/(2.0*h) *
                    # ( dEs_pert_0p[i,j,:,:] - dEs_pert_0n[i,j,:,:] );
            # store d2Esh
            dat = Dict{String, Any}();
            dat["Info"] = "perturb the origin atom in direction i
                and store d2Esh = 1/2h * ( E_{ℓ,n}(y+h*e₀) - E_{ℓ,n}(y-h*e₀) )";
            dat["h"] = h;
            dat["i"] = i;
            dat["l"] = ℓ;
            # dat["n"] = index_rcut2_ℓ;
###         dat["d2Esh"] = d2Esh[i, j, :, :];
            dat["datatype"] = "d2Esh";
            push!(data, dat);
        end
    end

    return D
end