text:
# diversity_measures.jl
# =====================
#
# Compute measures of genetic diversity with BioJulia data types.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/GeneticVariation.jl/blob/master/LICENSE.md

function pdist(d::Matrix{Tuple{Int,Int}})
    o = similar(d, Float64)
    @inbounds for i in eachindex(d)
        o[i] = d[i][1] / d[i][2]
    end
    return o
end

"""
    NL79(m::M, f::V) where {M<:AbstractMatrix{Float64},V<:AbstractVector{Float64}}

Compute nucleotide diversity using a matrix of the number of mutations
between sequence pairs, and a vector of the frequencies of each sequence
in the population.
"""
function NL79(m::AbstractMatrix{Float64}, f::AbstractVector{Float64})
    π = 0.0
    @inbounds for i = 1:lastindex(f), j = (i + 1):lastindex(f)
        π += m[i, j] * f[i] * f[j]
    end
    return 2 * π
end

"""
    NL79(sequences)

Compute nucleotide diversity, as described by Nei and Li in 1979.

This measure is defined as the average number of nucleotide differences per site
between two DNA sequences in all possible pairs in the sample population, and is
often denoted by the greek letter pi.

`Sequences` should be any iterable that yields biosequence types.

# Examples

```jldoctest
julia> testSeqs = [dna"AAAACTTTTACCCCCGGGGG",
                   dna"AAAACTTTTACCCCCGGGGG",
                   dna"AAAACTTTTACCCCCGGGGG",
                   dna"AAAACTTTTACCCCCGGGGG",
                   dna"AAAAATTTTACCCCCGTGGG",
                   dna"AAAAATTTTACCCCCGTGGG",
                   dna"AAAACTTTTTCCCCCGTAGG",
                   dna"AAAACTTTTTCCCCCGTAGG",
                   dna"AAAAATTTTTCCCCCGGAGG",
                   dna"AAAAATTTTTCCCCCGGAGG"]
10-element Array{BioSequences.BioSequence{BioSequences.DNAAlphabet{4}},1}:
 AAAACTTTTACCCCCGGGGG
 AAAACTTTTACCCCCGGGGG
 AAAACTTTTACCCCCGGGGG
 AAAACTTTTACCCCCGGGGG
 AAAAATTTTACCCCCGTGGG
 AAAAATTTTACCCCCGTGGG
 AAAACTTTTTCCCCCGTAGG
 AAAACTTTTTCCCCCGTAGG
 AAAAATTTTTCCCCCGGAGG
 AAAAATTTTTCCCCCGGAGG

 julia> NL79(testSeqs)
 0.096

```
"""
function NL79(sequences)
    frequencies = gene_frequencies(sequences)
    unique_sequences = collect(keys(frequencies))
    n = length(unique_sequences)
    if n < 2
        return 0.0
    else
        mutations = pdist(BioSequences.count_pairwise(Mutated, unique_sequences...))
        return NL79(mutations, collect(values(frequencies)))
    end
end

"""
    avg_mut(sequences)

The average number of mutations found in (n choose 2) pairwise comparisons of
sequences (i, j) in a sample of sequences.

`sequences` should be any indexable container of DNA sequence types.
"""
function avg_mut(sequences)
    s = length(sequences)
    @assert s ≥ 2 "At least 2 sequences are required."
    nmut = 0
    n = 0
    @inbounds for i in eachindex(sequences)
        si = sequences[i]
        for j in (i + 1):lastindex(sequences)
            nmut += count(Mutated, si, sequences[j])[1]
            n += 1
        end
    end
    return nmut / div(s * (s - 1), 2)
end


text:
# create a new "build_series_args" which converts all inputs into xs = Any[xitems], ys = Any[yitems].
# Special handling for: no args, xmin/xmax, parametric, dataframes
# Then once inputs have been converted, build the series args, map functions, etc.
# This should cut down on boilerplate code and allow more focused dispatch on type
# note: returns meta information... mainly for use with automatic labeling from DataFrames for now

typealias FuncOrFuncs @compat(Union{Function, AVec{Function}})

all3D(d::KW) = trueOrAllTrue(lt -> lt in (:contour, :heatmap, :surface, :wireframe, :contour3d), get(d, :linetype, :none))

# missing
convertToAnyVector(v::@compat(Void), d::KW) = Any[nothing], nothing

# fixed number of blank series
convertToAnyVector(n::Integer, d::KW) = Any[zeros(0) for i in 1:n], nothing

# numeric vector
convertToAnyVector{T<:Number}(v::AVec{T}, d::KW) = Any[v], nothing

# string vector
convertToAnyVector{T<:@compat(AbstractString)}(v::AVec{T}, d::KW) = Any[v], nothing

# numeric matrix
function convertToAnyVector{T<:Number}(v::AMat{T}, d::KW)
    if all3D(d)
        Any[Surface(v)]
    else
        Any[v[:,i] for i in 1:size(v,2)]
    end, nothing
end

# other matrix... vector of columns
function convertToAnyVector(m::AMat, d::KW)
    Any[begin
        v = vec(m[:,i])
        length(v) == 1 ? v[1] : v
    end for i=1:size(m,2)], nothing
end

# function
convertToAnyVector(f::Function, d::KW) = Any[f], nothing

# surface
convertToAnyVector(s::Surface, d::KW) = Any[s], nothing

# vector of OHLC
convertToAnyVector(v::AVec{OHLC}, d::KW) = Any[v], nothing

# dates
convertToAnyVector{D<:Union{Date,DateTime}}(dts::AVec{D}, d::KW) = Any[dts], nothing

# list of things (maybe other vectors, functions, or something else)
function convertToAnyVector(v::AVec, d::KW)
    if all(x -> typeof(x) <: Number, v)
        # all real numbers wrap the whole vector as one item
        Any[convert(Vector{Float64}, v)], nothing
    else
        # something else... treat each element as an item
        vcat(Any[convertToAnyVector(vi, d)[1] for vi in v]...), nothing
        # Any[vi for vi in v], nothing
    end
end

# --------------------------------------------------------------------

# TODO: can we avoid the copy here?  one error that crops up is that mapping functions over the same array
#       result in that array being shared.  push!, etc will add too many items to that array

compute_x(x::Void, y::Void, z)      = 1:size(z,1)
compute_x(x::Void, y, z)            = 1:size(y,1)
compute_x(x::Function, y, z)        = map(x, y)
compute_x(x, y, z)                  = copy(x)

# compute_y(x::Void, y::Function, z)  = error()
compute_y(x::Void, y::Void, z)      = 1:size(z,2)
compute_y(x, y::Function, z)        = map(y, x)
compute_y(x, y, z)                  = copy(y)

compute_z(x, y, z::Function)        = map(z, x, y)
compute_z(x, y, z::AbstractMatrix)  = Surface(z)
compute_z(x, y, z::Void)            = nothing
compute_z(x, y, z)                  = copy(z)

@noinline function compute_xyz(x, y, z)
    x = compute_x(x,y,z)
    y = compute_y(x,y,z)
    z = compute_z(x,y,z)
    x, y, z
end

# not allowed
compute_xyz(x::Void, y::FuncOrFuncs, z)       = error("If you want to plot the function `$y`, you need to define the x values!")
compute_xyz(x::Void, y::Void, z::FuncOrFuncs) = error("If you want to plot the function `$z`, you need to define x and y values!")
compute_xyz(x::Void, y::Void, z::Void)        = error("x/y/z are all nothing!")

# --------------------------------------------------------------------

# create n=max(mx,my) series arguments. the shorter list is cycled through
# note: everything should flow through this
function build_series_args(plt::AbstractPlot, kw::KW) #, idxfilter)
    x, y, z = map(sym -> pop!(kw, sym, nothing), (:x, :y, :z))
    if nothing == x == y == z
        return [], nothing, nothing
    end

    xs, xmeta = convertToAnyVector(x, kw)
    ys, ymeta = convertToAnyVector(y, kw)
    zs, zmeta = convertToAnyVector(z, kw)

    fr = pop!(kw, :fillrange, nothing)
    fillranges, _ = if typeof(fr) <: Number
        ([fr],nothing)
    else
        convertToAnyVector(fr, kw)
    end

    mx = length(xs)
    my = length(ys)
    mz = length(zs)
    ret = Any[]
    for i in 1:max(mx, my, mz)

        # try to set labels using ymeta
        d = copy(kw)
        if !haskey(d, :label) && ymeta != nothing
            if isa(ymeta, Symbol)
                d[:label] = string(ymeta)
            elseif isa(ymeta, AVec{Symbol})
                d[:label] = string(ymeta[mod1(i,length(ymeta))])
            end
        end

        # build the series arg dict
        numUncounted = pop!(d, :numUncounted, 0)
        commandIndex = i + numUncounted
        n = plt.n + i

        dumpdict(d, "before getSeriesArgs")
        d = getSeriesArgs(plt.backend, getplotargs(plt, n), d, commandIndex, convertSeriesIndex(plt, n), n)
        dumpdict(d, "after getSeriesArgs")

        d[:x], d[:y], d[:z] = compute_xyz(xs[mod1(i,mx)], ys[mod1(i,my)], zs[mod1(i,mz)])
        lt = d[:linetype]

        # for linetype `line`, need to sort by x values
        if lt == :line
            # order by x
            indices = sortperm(d[:x])
            d[:x] = d[:x][indices]
            d[:y] = d[:y][indices]
            d[:linetype] = :path
        end

        # special handling for missing x in box plot... all the same category
        if lt == :box && xs[mod1(i,mx)] == nothing
            d[:x] = ones(Int, length(d[:y]))
        end

        # map functions to vectors
        if isa(d[:marker_z], Function)
            d[:marker_z] = map(d[:marker_z], d[:x])
        end

        # @show fillranges
        d[:fillrange] = fillranges[mod1(i,length(fillranges))]
        if isa(d[:fillrange], Function)
            d[:fillrange] = map(d[:fillrange], d[:x])
        end

        # handle error bars
        for esym in (:xerror, :yerror)
            if get(d, esym, nothing) != nothing
                # we make a copy of the KW and apply an errorbar recipe
                append!(ret, apply_series_recipe(copy(d), Val{esym}))
            end
        end

        # handle ribbons
        if get(d, :ribbon, nothing) != nothing
            rib = d[:ribbon]
            d[:fillrange] = (d[:y] - rib, d[:y] + rib)
        end

        # handle quiver plots
        if lt == :quiver
            d[:linetype] = lt = :path
            d[:linewidth] = 0
        end
        if get(d, :quiver, nothing) != nothing
            append!(ret, apply_series_recipe(copy(d), Val{:quiver}))
        end

        # now that we've processed a given series... optionally split into
        # multiple dicts through a recipe (for example, a box plot is split into component
        # parts... polygons, lines, and scatters)
        # note: we pass in a Val type (i.e. Val{:box}) so that we can dispatch on the linetype
        kwlist = apply_series_recipe(d, Val{lt})
        append!(ret, kwlist)

        # # add it to our series list
        # push!(ret, d)
    end

    ret, xmeta, ymeta
end

# --------------------------------------------------------------------
# process_inputs
# --------------------------------------------------------------------

# These methods take a plot and the keyword arguments, and processes the input
# arguments (x/y/z, group, etc), populating the KW dict with appropriate values.

# --------------------------------------------------------------------
# 0 arguments
# --------------------------------------------------------------------

# don't do anything
function process_inputs(plt::AbstractPlot, d::KW)
end

# --------------------------------------------------------------------
# 1 argument
# --------------------------------------------------------------------

function process_inputs(plt::AbstractPlot, d::KW, n::Integer)
    # d[:x], d[:y], d[:z] = zeros(0), zeros(0), zeros(0)
    d[:x] = d[:y] = d[:z] = n
end

# no special handling... assume x and z are nothing
function process_inputs(plt::AbstractPlot, d::KW, y)
    d[:y] = y
end

# matrix... is it z or y?
function process_inputs{T<:Number}(plt::AbstractPlot, d::KW, mat::AMat{T})
    if all3D(d)
        n,m = size(mat)
        d[:x], d[:y], d[:z] = 1:n, 1:m, mat
    else
        d[:y] = mat
    end
end

# images - grays
function process_inputs{T<:Gray}(plt::AbstractPlot, d::KW, mat::AMat{T})
    d[:linetype] = :image
    n,m = size(mat)
    d[:x], d[:y], d[:z] = 1:n, 1:m, Surface(mat)
    # handle images... when not supported natively, do a hack to use heatmap machinery
    if !nativeImagesSupported()
        d[:linetype] = :heatmap
        d[:yflip] = true
        d[:z] = Surface(convert(Matrix{Float64}, mat.surf))
        d[:fillcolor] = ColorGradient([:black, :white])
    end
end

# images - colors
function process_inputs{T<:Colorant}(plt::AbstractPlot, d::KW, mat::AMat{T})
    d[:linetype] = :image
    n,m = size(mat)
    d[:x], d[:y], d[:z] = 1:n, 1:m, Surface(mat)
    # handle images... when not supported natively, do a hack to use heatmap machinery
    if !nativeImagesSupported()
        d[:yflip] = true
        imageHack(d)
    end
end

# plotting arbitrary shapes/polygons
function process_inputs(plt::AbstractPlot, d::KW, shape::Shape)
    d[:x], d[:y] = shape_coords(shape)
    d[:linetype] = :shape
end
function process_inputs(plt::AbstractPlot, d::KW, shapes::AVec{Shape})
    d[:x], d[:y] = shape_coords(shapes)
    d[:linetype] = :shape
end
function process_inputs(plt::AbstractPlot, d::KW, shapes::AMat{Shape})
    x, y = [], []
    for j in 1:size(shapes, 2)
        tmpx, tmpy = shape_coords(vec(shapes[:,j]))
        push!(x, tmpx)
        push!(y, tmpy)
    end
    d[:x], d[:y] = x, y
    d[:linetype] = :shape
end

# function without range... use the current range of the x-axis
function process_inputs(plt::AbstractPlot, d::KW, f::FuncOrFuncs)
    process_inputs(plt, d, f, xmin(plt), xmax(plt))
end

# --------------------------------------------------------------------
# 2 arguments
# --------------------------------------------------------------------

function process_inputs(plt::AbstractPlot, d::KW, x, y)
    d[:x], d[:y] = x, y
end

# if functions come first, just swap the order (not to be confused with parametric functions...
# as there would be more than one function passed in)
function process_inputs(plt::AbstractPlot, d::KW, f::FuncOrFuncs, x)
    @assert !(typeof(x) <: FuncOrFuncs)  # otherwise we'd hit infinite recursion here
    process_inputs(plt, d, x, f)
end

# --------------------------------------------------------------------
# 3 arguments
# --------------------------------------------------------------------

# no special handling... just pass them through
function process_inputs(plt::AbstractPlot, d::KW, x, y, z)
    d[:x], d[:y], d[:z] = x, y, z
end

# 3d line or scatter
function process_inputs(plt::AbstractPlot, d::KW, x::AVec, y::AVec, zvec::AVec)
    # default to path3d if we haven't set a 3d linetype
    lt = get(d, :linetype, :none)
    if lt == :scatter
        d[:linetype] = :scatter3d
    elseif !(lt in _3dTypes)
        d[:linetype] = :path3d
    end
    d[:x], d[:y], d[:z] = x, y, zvec
end

# surface-like... function
function process_inputs{TX,TY}(plt::AbstractPlot, d::KW, x::AVec{TX}, y::AVec{TY}, zf::Function)
    x = TX <: Number ? sort(x) : x
    y = TY <: Number ? sort(y) : y
    # x, y = sort(x), sort(y)
    d[:z] = Surface(zf, x, y)  # TODO: replace with SurfaceFunction when supported
    d[:x], d[:y] = x, y
end

# surface-like... matrix grid
function process_inputs{TX,TY,TZ}(plt::AbstractPlot, d::KW, x::AVec{TX}, y::AVec{TY}, zmat::AMat{TZ})
    @assert size(zmat) == (length(x), length(y))
    if TX <: Number && !issorted(x)
        idx = sortperm(x)
        x, zmat = x[idx], zmat[idx, :]
    end
    if TY <: Number && !issorted(y)
        idx = sortperm(y)
        y, zmat = y[idx], zmat[:, idx]
    end
    d[:x], d[:y], d[:z] = x, y, Surface{Matrix{TZ}}(zmat)
    if !like_surface(get(d, :linetype, :none))
        d[:linetype] = :contour
    end
end

# surfaces-like... general x, y grid
function process_inputs{T<:Number}(plt::AbstractPlot, d::KW, x::AMat{T}, y::AMat{T}, zmat::AMat{T})
    @assert size(zmat) == size(x) == size(y)
    # d[:x], d[:y], d[:z] = Any[x], Any[y], Surface{Matrix{Float64}}(zmat)
    d[:x], d[:y], d[:z] = map(Surface{Matrix{Float64}}, (x, y, zmat))
    if !like_surface(get(d, :linetype, :none))
        d[:linetype] = :contour
    end
end

# --------------------------------------------------------------------
# Parametric functions
# --------------------------------------------------------------------

# special handling... xmin/xmax with function(s)
function process_inputs(plt::AbstractPlot, d::KW, f::FuncOrFuncs, xmin::Number, xmax::Number)
    width = get(plt.plotargs, :size, (100,))[1]
    x = linspace(xmin, xmax, width)
    process_inputs(plt, d, x, f)
end

# special handling... xmin/xmax with parametric function(s)
process_inputs{T<:Number}(plt::AbstractPlot, d::KW, fx::FuncOrFuncs, fy::FuncOrFuncs, u::AVec{T}) = process_inputs(plt, d, mapFuncOrFuncs(fx, u), mapFuncOrFuncs(fy, u))
process_inputs{T<:Number}(plt::AbstractPlot, d::KW, u::AVec{T}, fx::FuncOrFuncs, fy::FuncOrFuncs) = process_inputs(plt, d, mapFuncOrFuncs(fx, u), mapFuncOrFuncs(fy, u))
process_inputs(plt::AbstractPlot, d::KW, fx::FuncOrFuncs, fy::FuncOrFuncs, umin::Number, umax::Number, numPoints::Int = 1000) = process_inputs(plt, d, fx, fy, linspace(umin, umax, numPoints))

# special handling... 3D parametric function(s)
process_inputs{T<:Number}(plt::AbstractPlot, d::KW, fx::FuncOrFuncs, fy::FuncOrFuncs, fz::FuncOrFuncs, u::AVec{T}) = process_inputs(plt, d, mapFuncOrFuncs(fx, u), mapFuncOrFuncs(fy, u), mapFuncOrFuncs(fz, u))
process_inputs{T<:Number}(plt::AbstractPlot, d::KW, u::AVec{T}, fx::FuncOrFuncs, fy::FuncOrFuncs, fz::FuncOrFuncs) = process_inputs(plt, d, mapFuncOrFuncs(fx, u), mapFuncOrFuncs(fy, u), mapFuncOrFuncs(fz, u))
process_inputs(plt::AbstractPlot, d::KW, fx::FuncOrFuncs, fy::FuncOrFuncs, fz::FuncOrFuncs, umin::Number, umax::Number, numPoints::Int = 1000) = process_inputs(plt, d, fx, fy, fz, linspace(umin, umax, numPoints))

# --------------------------------------------------------------------
# Lists of tuples and FixedSizeArrays
# --------------------------------------------------------------------

# if we get an unhandled tuple, just splat it in
function process_inputs(plt::AbstractPlot, d::KW, tup::Tuple)
    process_inputs(plt, d, tup...)
end

# (x,y) tuples
function process_inputs{R1<:Number,R2<:Number}(plt::AbstractPlot, d::KW, xy::AVec{Tuple{R1,R2}})
    process_inputs(plt, d, unzip(xy)...)
end
function process_inputs{R1<:Number,R2<:Number}(plt::AbstractPlot, d::KW, xy::Tuple{R1,R2})
    process_inputs(plt, d, [xy[1]], [xy[2]])
end

# (x,y,z) tuples
function process_inputs{R1<:Number,R2<:Number,R3<:Number}(plt::AbstractPlot, d::KW, xyz::AVec{Tuple{R1,R2,R3}})
    process_inputs(plt, d, unzip(xyz)...)
end
function process_inputs{R1<:Number,R2<:Number,R3<:Number}(plt::AbstractPlot, d::KW, xyz::Tuple{R1,R2,R3})
    process_inputs(plt, d, [xyz[1]], [xyz[2]], [xyz[3]])
end

# 2D FixedSizeArrays
function process_inputs{T<:Number}(plt::AbstractPlot, d::KW, xy::AVec{FixedSizeArrays.Vec{2,T}})
    process_inputs(plt, d, unzip(xy)...)
end
function process_inputs{T<:Number}(plt::AbstractPlot, d::KW, xy::FixedSizeArrays.Vec{2,T})
    process_inputs(plt, d, [xy[1]], [xy[2]])
end

# 3D FixedSizeArrays
function process_inputs{T<:Number}(plt::AbstractPlot, d::KW, xyz::AVec{FixedSizeArrays.Vec{3,T}})
    process_inputs(plt, d, unzip(xyz)...)
end
function process_inputs{T<:Number}(plt::AbstractPlot, d::KW, xyz::FixedSizeArrays.Vec{3,T})
    process_inputs(plt, d, [xyz[1]], [xyz[2]], [xyz[3]])
end

# --------------------------------------------------------------------
# handle grouping
# --------------------------------------------------------------------

# function process_inputs(plt::AbstractPlot, d::KW, groupby::GroupBy, args...)
#     ret = Any[]
#     error("unfinished after series reorg")
#     for (i,glab) in enumerate(groupby.groupLabels)
#         # TODO: don't automatically overwrite labels
#         kwlist, xmeta, ymeta = process_inputs(plt, d, args...,
#                                             idxfilter = groupby.groupIds[i],
#                                             label = string(glab),
#                                             numUncounted = length(ret))  # we count the idx from plt.n + numUncounted + i
#         append!(ret, kwlist)
#     end
#     ret, nothing, nothing # TODO: handle passing meta through
# end

# --------------------------------------------------------------------
# For DataFrame support.  Imports DataFrames and defines the necessary methods which support them.
# --------------------------------------------------------------------

function setup_dataframes()
    @require DataFrames begin
    # @eval begin
        # import DataFrames

        DFS = Union{Symbol, AbstractArray{Symbol}}

        function handle_dfs(df::DataFrames.AbstractDataFrame, d::KW, letter, dfs::DFS)
            if isa(dfs, Symbol)
                get!(d, symbol(letter * "label"), string(dfs))
                collect(df[dfs])
            else
                get!(d, :label, reshape(dfs, 1, length(dfs)))
                Any[collect(df[s]) for s in dfs]
            end
        end

        function handle_group(df::DataFrames.AbstractDataFrame, d::KW)
            if haskey(d, :group)
                g = d[:group]
                if isa(g, Symbol)
                    d[:group] = collect(df[g])
                end
            end
        end

        @recipe function plot(df::DataFrames.AbstractDataFrame, sy::DFS)
            handle_group(df, d)
            handle_dfs(df, d, "y", sy)
        end

        @recipe function plot(df::DataFrames.AbstractDataFrame, sx::DFS, sy::DFS)
            handle_group(df, d)
            x = handle_dfs(df, d, "x", sx)
            y = handle_dfs(df, d, "y", sy)
            x, y
        end

        @recipe function plot(df::DataFrames.AbstractDataFrame, sx::DFS, sy::DFS, sz::DFS)
            handle_group(df, d)
            x = handle_dfs(df, d, "x", sx)
            y = handle_dfs(df, d, "y", sy)
            z = handle_dfs(df, d, "z", sz)
            x, y, z
        end

        # get_data(df::DataFrames.AbstractDataFrame, arg::Symbol) = df[arg]
        # get_data(df::DataFrames.AbstractDataFrame, arg) = arg
        #
        # function process_inputs(plt::AbstractPlot, d::KW, df::DataFrames.AbstractDataFrame, args...)
        #     # d[:dataframe] = df
        #     process_inputs(plt, d, map(arg -> get_data(df, arg), args)...)
        # end
        #
        # # expecting the column name of a dataframe that was passed in... anything else should error
        # function extractGroupArgs(s::Symbol, df::DataFrames.AbstractDataFrame, args...)
        #     if haskey(df, s)
        #         return extractGroupArgs(df[s])
        #     else
        #         error("Got a symbol, and expected that to be a key in d[:dataframe]. s=$s d=$d")
        #     end
        # end

        # function getDataFrameFromKW(d::KW)
        #     get(d, :dataframe) do
        #         error("Missing dataframe argument!")
        #     end
        # end

        # # the conversion functions for when we pass symbols or vectors of symbols to reference dataframes
        # convertToAnyVector(s::Symbol, d::KW) = Any[getDataFrameFromKW(d)[s]], s
        # convertToAnyVector(v::AVec{Symbol}, d::KW) = (df = getDataFrameFromKW(d); Any[df[s] for s in v]), v

    end
end

ext: jl


text:
using MagmaThermoKinematics.Diffusion2D
using ParallelStencil
using ParallelStencil.FiniteDifferences2D
using Plots
using LinearAlgebra
using SpecialFunctions
using Test

# Initialize for multiple threads (GPU is not tested here)
@init_parallel_stencil(Threads, Float64, 2);    # initialize parallel stencil in 2D

const CreatePlots = false      # easy way to deactivate plotting throughout

function Diffusion_SteadyState2D(Setup="Constant_Zdirection")
# steady state diffusion in x and z-direction for constant and variable K

# Model parameters
if      Setup=="Constant_Zdirection" || Setup=="VariableK_Zdirection"
    W,H                 =   3, 30;                              # Width, Height in km
    k_rock1             =   3;
    k_rock2             =   1.5;
    Nx, Nz              =   10, 100;                            # resolution
    Tbot                =   600;
    H1                  =   20;
    GeoT                =   Tbot/H;
elseif  Setup=="Constant_Xdirection" || Setup=="VariableK_Xdirection"
    W,H                 =   40, 3;                              # Width, Height in km
    k_rock1             =   3;
    k_rock2             =   1.5;
    Nx, Nz              =   100, 10;                            # resolution
    Tbot                =   600;
    W1                  =   20;
    GeoT                =   Tbot/W;
else
    error("unknown setup")
end

SecYear                 =   365.25*24*3600
ρ                       =   2800;                               # Density
cp                      =   1050;                               # Heat capacity
L                       =   350e3;                              # Latent heat J/kg/K
dx                      =   W/(Nx-1)*1e3; dz = H*1e3/(Nz-1);    # grid size [m]
κ                       =   k_rock1./(ρ*cp);                     # thermal diffusivity
dt                      =   min(dx^2, dz^2)./κ/10;             # stable timestep (required for explicit FD)

# Array initializations (1 - main arrays on which we can initialize properties)
T                       =   @zeros(Nx,   Nz);
K                       =   @ones(Nx,    Nz)*k_rock1;
Rho                     =   @ones(Nx,    Nz)*ρ;
Cp                      =   @ones(Nx,    Nz)*cp;
dPhi_dt                 =   @zeros(Nx,   Nz);

# Work array initialization
Tnew, qx,qz, Kx, Kz     =   @zeros(Nx,   Nz), @zeros(Nx-1, Nz),     @zeros(Nx,   Nz-1), @zeros(Nx-1, Nz), @zeros(Nx,   Nz-1)    # thermal solver
X,Xc,Z                  =   @zeros(Nx,   Nz), @zeros(Nx-1, Nz-1),   @zeros(Nx,   Nz)    # 2D gridpoints

# Set up model geometry & initial T structure
x,z                     =   0:dx:W*1e3, -H*1e3:dz:(-H*1e3+(Nz-1)*dz);
X,Z                     =   ones(Nz)' .* x, z' .* ones(Nx);                             # 2D coordinate grids
Xc                      =   (X[2:Nx,:] + X[1:Nx-1,:])/2.0;
Grid, Spacing           =   (X,Z), (dx,dz);
T                       .=   Tbot;                                             # initial (linear) temperature profile

if      Setup=="Constant_Zdirection" || Setup=="VariableK_Zdirection"
    # bottom BC
    T[1,:]              .=   0;
else
    T[:,1]              .=   0;
end

if Setup=="VariableK_Zdirection"
    # variable k
    K[Z.<-H1*1e3] .= k_rock2;

elseif Setup=="VariableK_Xdirection"
    # variable k
    K[X.< W1*1e3] .= k_rock2;
end

time,time_kyrs, dike_inj = 0.0, 0.0, 0.0;
err = 100;
it = 0;
while (err>1e-10) & (it<1e6)

    it += 1
    # Perform a diffusion step
    @parallel diffusion2D_step!(Tnew, T, qx, qz, K, Kx, Kz, Rho, Cp, dt, dx, dz,  L, dPhi_dt);
    if Setup=="Constant_Zdirection" || Setup=="VariableK_Zdirection"
         # diffusion in z-direction
        @parallel (1:size(T,2)) bc2D_x!(Tnew);                                                      # set lateral boundary conditions (flux-free)
        Tnew[:,1] .= Tbot; Tnew[:,end] .= 0.0;                                                    # bottom & top temperature (constant)

    else
        # diffusion in x-direction
        @parallel (1:size(T,1)) bc2D_z!(Tnew);                                                      # set lateral boundary conditions (flux-free)
        Tnew[1,:] .= 0; Tnew[end,:] .= Tbot;                                                    # bottom & top temperature (constant)
    end

    err        =   maximum(abs.(Tnew-T));

    # Tracers         =   UpdateTracers(Tracers, Grid, Spacing, Tnew, Phi);                       # Update info on tracers
    T, Tnew         =   Tnew, T;                                                                # Update temperature
    time, time_kyrs =   time + dt, time/SecYear/1e3;                                            # Keep track of evolved time

    if mod(it,10000)==0  # Visualisation
    #    println(" Timestep $it = $(round(time/SecYear)/1e3) kyrs")
    end

end

x_km, z_km  =   x./1e3, z./1e3;

# compute analytical solution
if  Setup=="Constant_Zdirection"
    Tanal = -z_km.*GeoT;
    Tnum  = T[1,:];
    fname = "Diffusion_2D_SS_constantK_Z";

elseif  Setup=="VariableK_Zdirection"
    # 1D steady steate analytical solution for variable K is given by the folliwing balance equations
    #   k1*dT/dz|_1                     =   k2*dT/dz|_2     (heat flux)
    #   dT/dz|_1 * H1 + dT/dz|_2 * H2 =   Tbot            (assuming Ttop=0)
    #   H1 + H2                         =   H               (total thickness)
    #
    # substitute eq. 1 into eq 2 to eliminate dT/dz|_2
    #  dT/dz|_1 * H_1 + dT/dz|_1 * k1/k2*H2 =   Tbot
    #  dT/dz|_1 = Tbot/(H1 + k1/k2*H2)

    H2      =   H - H1;
    dTdz1   =   Tbot/(H1 + k_rock1/k_rock2*H2)
    dTdz2   =   k_rock1/k_rock2.*dTdz1

    Tanal   =   zeros(size(z_km));              Tanal2  =   copy(Tanal);
    Tanal   .=   -z_km[:].*dTdz1;
    Tanal2  .=   -z_km.*dTdz2 .- dTdz1*H1;
    Tanal[z_km .< -H1 ] .= Tanal2[z_km .< -H1 ]

    Tnum    =    T[1,:];
    fname = "Diffusion_2D_SS_variableK_Z";

elseif  Setup=="Constant_Xdirection"
    Tanal = x_km.*GeoT;
    Tnum  = T[:,1];
    fname = "Diffusion_2D_SS_constantK_X";

elseif  Setup=="VariableK_Xdirection"
    W2      =   W - W1;
    dTdx1   =   Tbot/(W1 + k_rock1/k_rock2*W2)
    dTdx2   =   k_rock1/k_rock2.*dTdx1

    Tanal   =   zeros(size(x_km));              Tanal2  =   copy(Tanal);
    Tanal  .=   x_km.*dTdx2 ;
    Tanal2  .=   x_km[:].*dTdx1 .+ dTdx1*W2;
    Tanal[x_km .> W1 ] .= Tanal2[x_km .> W1 ]

    Tnum    =    T[:,1];
    fname = "Diffusion_2D_SS_variableK_X";
end

if      Setup=="Constant_Zdirection" || Setup=="VariableK_Zdirection"
    if CreatePlots
        # create plot
        plot(Tanal,z_km, label = "Analytics");
        plot!(Tnum,z_km,ylabel="Depth [km]",xlabel="Temperature [C]", label = "Numerics",  marker = 2,   linewidth = 0);
    end
    error = norm(T[1,:] .- Tanal,2);
else
    if CreatePlots
        # create plot
        plot(x_km, Tanal,label = "Analytics");
        plot!(x_km,Tnum, xlabel="Width [km]",ylabel="Temperature [C]", label = "Numerics",  marker = 2,   linewidth = 0);
    end
    error = norm(T[:,1] .- Tanal,2);
end
if CreatePlots
    png(fname)
end

return error;        # return error

end # end of steady state diffusion test

function Diffusion_Halfspace2D()
    # Halfspace cooling example

    # Model parameters
    W,H                     =   3, 300;                             # Width, Height in km
    k_rock1                 =   3;
    Nx, Nz                  =   10, 100;                            # resolution
    Tbot                    =   1200;
    SecYear                 =   365.25*24*3600
    CoolingAge              =   30e6*SecYear;                       # thermal cooling age
    ρ                       =   2800;                               # Density
    cp                      =   1050;                               # Heat capacity
    L                       =   350e3;                              # Latent heat J/kg/K
    dx                      =   W/(Nx-1)*1e3; dz = H*1e3/(Nz-1);    # grid size [m]
    κ                       =   k_rock1./(ρ*cp);                     # thermal diffusivity
    dt                      =   min(dx^2, dz^2)./κ/10;             # stable timestep (required for explicit FD)

    numTime                 =   ceil(CoolingAge/dt);
    dt                      =   CoolingAge/numTime;
    nt                      =   Int(numTime);

    # Array initializations (1 - main arrays on which we can initialize properties)
    T                       =   @ones(Nx,    Nz)*Tbot;
    K                       =   @ones(Nx,    Nz)*k_rock1;
    Rho                     =   @ones(Nx,    Nz)*ρ;
    Cp                      =   @ones(Nx,    Nz)*cp;
    dPhi_dt                 =   @zeros(Nx,   Nz);

    # Work array initialization
    Tnew, qx,qz, Kx, Kz     =   @zeros(Nx,   Nz), @zeros(Nx-1, Nz),     @zeros(Nx,   Nz-1), @zeros(Nx-1, Nz), @zeros(Nx,   Nz-1)    # thermal solver
    X,Xc,Z                  =   @zeros(Nx,   Nz), @zeros(Nx-1, Nz-1),   @zeros(Nx,   Nz)    # 2D gridpoints

    # Set up model geometry & initial T structure
    x,z                     =   0:dx:W*1e3, -H*1e3:dz:(-H*1e3+(Nz-1)*dz);
    X,Z                     =   ones(Nz)' .* x, z' .* ones(Nx);                             # 2D coordinate grids
    Xc                      =   (X[2:Nx,:] + X[1:Nx-1,:])/2.0;
    Grid, Spacing           =   (X,Z), (dx,dz);
    T                       .=   Tbot;                                             # initial (linear) temperature profile

    T[:,end]               .=   0;     # top BC
    Tnew                   .= T;

    time,time_kyrs, dike_inj = 0.0, 0.0, 0.0;
    err = 100;
    it = 0;
    for it=1:nt

        # Perform a diffusion step
        @parallel diffusion2D_step!(Tnew, T, qx, qz, K, Kx, Kz, Rho, Cp, dt, dx, dz,  L, dPhi_dt);
        # diffusion in z-direction
        @parallel (1:size(T,2)) bc2D_x!(Tnew);                                                      # set lateral boundary conditions (flux-free)
        Tnew[:,1] .= Tbot; Tnew[:,end] .= 0.0;                                                    # bottom & top temperature (constant)

        # Tracers         =   UpdateTracers(Tracers, Grid, Spacing, Tnew, Phi);                       # Update info on tracers
        T, Tnew         =   Tnew, T;                                                                # Update temperature
        time, time_kyrs =   time + dt, time/SecYear/1e3;                                            # Keep track of evolved time

        if mod(it,1000)==0  # print progress
        #    println(" Timestep $it = $(round(time/SecYear)/1e3) kyrs")
        end

    end

    x_km, z_km  =   x./1e3, z./1e3;

    # compute analytical solution
    Tanal = -z_km;

    Tanal      =   (0-Tbot).*erfc.((abs.(z_km).*1e3)./(2*sqrt(κ*CoolingAge))) .+ Tbot;

    Tnum  = T[1,:];
    fname = "Diffusion_2D_Halfspace";

    if CreatePlots
        # create plot
        plot(Tanal,z_km, label = "Analytics");
        plot!(Tnum,z_km,ylabel="Depth [km]",xlabel="Temperature [C]", label = "Numerics",  marker = 2,   linewidth = 0);

        png(fname)
    end

    error = norm(T[1,:] .- Tanal,2);
    return error;        # return error

end # end of halfspace cooling test

function Diffusion_Gaussian2D(Setup="2D")
    # Gaussian diffusion test in 2D

    # Model parameters
    W,H                     =   300, 300;                           # Width, Height in km
    k_rock1                 =   3;
    Nx, Nz                  =   100, 100;                           # resolution
    Tbot                    =   0;
    SecYear                 =   365.25*24*3600
    σ                       =   15e3;                               # halfwidth of gaussian
    Tmax                    =   1000;                                # max. of gaussian
    TotalTime               =   3e6*SecYear;                        # thermal cooling age
    ρ                       =   2800;                               # Density
    cp                      =   1050;                               # Heat capacity
    L                       =   350e3;                              # Latent heat J/kg/K
    dx                      =   W/(Nx-1)*1e3; dz = H*1e3/(Nz-1);    # grid size [m]
    κ                       =   k_rock1./(ρ*cp);                    # thermal diffusivity
    dt                      =   min(dx^2, dz^2)./κ/10;              # stable timestep (required for explicit FD)

    numTime                 =   ceil(TotalTime/dt);
    dt                      =   TotalTime/numTime/100;
    nt                      =   Int(numTime);

    # Array initializations (1 - main arrays on which we can initialize properties)
    T                       =   @ones(Nx,    Nz)*Tbot;
    K                       =   @ones(Nx,    Nz)*k_rock1;
    Rho                     =   @ones(Nx,    Nz)*ρ;
    Cp                      =   @ones(Nx,    Nz)*cp;
    dPhi_dt                 =   @zeros(Nx,   Nz);

    # Work array initialization
    Tnew, qx,qz, Kx, Kz     =   @zeros(Nx,   Nz), @zeros(Nx-1, Nz),     @zeros(Nx,   Nz-1), @zeros(Nx-1, Nz), @zeros(Nx,   Nz-1)    # thermal solver
    X,Xc,Z                  =   @zeros(Nx,   Nz), @zeros(Nx-1, Nz-1),   @zeros(Nx,   Nz)    # 2D gridpoints

    # Set up model geometry & initial T structure
    x,z                     =   -W/2*1e3:dx:W/2*1e3, -H/2*1e3:dz:(-H/2*1e3+(Nz-1)*dz);
    X,Z                     =   ones(Nz)' .* x, z' .* ones(Nx);                             # 2D coordinate grids
    Xc                      =   (X[2:Nx,:] + X[1:Nx-1,:])/2.0;
    Grid, Spacing           =   (X,Z), (dx,dz);

    if Setup=="2D"
        T                  .=  Tmax.*exp.(  -(X.^2 .+ Z.^2)./(σ^2));                     # initial gaussian profile
    elseif Setup=="Axisymmetric"
        T                  .=  Tmax.*exp.(  -(X.^2 .+ Z.^2)./(σ^2));                     # initial gaussian profile
    else
        error("Unknown setup")
    end

    Tnew                   .=   T;

    #ENV["GKSwstype"]="nul"; if isdir("viz2D_out")==false mkdir("viz2D_out") end; loadpath = "./viz2D_out/"; anim = Animation(loadpath,String[])
    #println("Animation directory: $(anim.dir)")

    time,time_kyrs          = 0.0, 0.0;
    err = 100;
    it = 0;
    nt = 500
    for it=1:nt

        # Perform a diffusion step
        if Setup=="2D"
            @parallel diffusion2D_step!(Tnew, T, qx, qz, K, Kx, Kz, Rho, Cp, dt, dx, dz,  L, dPhi_dt);
        elseif Setup=="Axisymmetric"
            @parallel diffusion2D_AxiSymm_step!(Tnew, T, X, Xc, qx, qz, K, Kx, Kz, Rho, Cp, dt, dx, dz,  L, dPhi_dt);
        end

        # diffusion in z-direction
        @parallel (1:size(T,2)) bc2D_x!(Tnew);                                                      # set lateral boundary conditions (flux-free)
        Tnew[:,1] .= Tbot; Tnew[:,end] .= 0.0;                                                    # bottom & top temperature (constant)

        T, Tnew         =   Tnew, T;                                                                # Update temperature
        time, time_kyrs =   time + dt, time/SecYear/1e3;                                            # Keep track of evolved time

        if mod(it,50)==0  # print progress
        #    println(" Timestep $it = $(round(time/SecYear)/1e3) kyrs")

        #    x_km, z_km  =   x./1e3, z./1e3;
        #    p1          =   heatmap(x_km, z_km, T',         aspect_ratio=1, xlims=(x_km[1],x_km[end]), ylims=(z_km[1],z_km[end]),   c=:inferno, title="Temperature, $(round(time_kyrs, digits=2)) kyrs",  dpi=150)
        #    plot(p1); frame(anim)
        end

    end

    x_km, z_km  =   x./1e3, z./1e3;

    # compute analytical solution

    if Setup=="2D"
        Tanal  =  Tmax./(1 + 4*time*κ/σ^2)^(2/2).*exp.(  -(X.^2 .+ Z.^2)./(σ^2 + 4*time*κ));                     # initial gaussian profile
        fname = "Diffusion_2D_Gaussian";

    elseif Setup=="Axisymmetric"
        # Axisymmetric is like 3D, where Y=X. Hence we can use the 3D solution, which is
        Tanal  =  Tmax./( (1 + 4*time*κ/σ^2)^(3/2) ).*exp.(  -(X.^2 .+ Z.^2)./(σ^2 + 4*time*κ));                     # initial gaussian profile
        fname = "Diffusion_Axisymmetric_Gaussian";
    end
    Terror =  T - Tanal

    if CreatePlots
        # create plot
        p1          =   heatmap(x_km, z_km, Terror',         aspect_ratio=1, xlims=(x_km[1],x_km[end]), ylims=(z_km[1],z_km[end]),   c=:inferno, title="T error 2D $(round(time_kyrs/1e3, digits=2)) Myrs",  dpi=150)
        plot(p1);
        png(fname)
    end

    error = norm(Terror[:],2);

    return error;        # return error

end # end of gaussian diffusion test

# Create a range of 2D diffusion tests which calls the routines above
@testset "2D steady state diffusion" begin
    @test Diffusion_SteadyState2D("VariableK_Zdirection") ≈ 8.749  atol=1e-3;
    @test Diffusion_SteadyState2D("Constant_Zdirection")  ≈ 1e-5   atol=1e-4;
    @test Diffusion_SteadyState2D("VariableK_Xdirection") ≈ 2.041  atol=1e-3;
    @test Diffusion_SteadyState2D("Constant_Xdirection")  ≈ 1e-4   atol=1e-4;
end;

@testset "2D halfspace cooling" begin
    @test  Diffusion_Halfspace2D() ≈ 0.6247443434361517 atol=1e-5;
end;
@testset "2D Gaussian diffusion" begin
    @test Diffusion_Gaussian2D("2D")           ≈  5.688521713402446 atol=1e-3;
    @test Diffusion_Gaussian2D("Axisymmetric") ≈ 11.361627247736031 atol=1e-3;
end;


text:
using DiffEqFlux, OrdinaryDiffEq, Flux, Optim, Plots
using DiffEqSensitivity
using Zygote
using ForwardDiff
using LinearAlgebra
using Random
using Statistics
using ProgressBars, Printf
using Flux.Optimise: update!
using Flux.Losses: mae, mse
using BSON: @save, @load

is_restart = false
n_epoch = 25000;
ntotal = 20
n_plot = 50;
grad_max = 1.e2;
batch_size = ntotal;

opt = ADAMW(0.005, (0.9, 0.999), 1.f-6);

lb = 1.e-6;
ub = 1.e5;
ode_solver = AutoTsit5(Rosenbrock23(autodiff=false));

k1 = .35e0
k2 = .266e2
k3 = .123e5
k4 = .86e-3
k5 = .82e-3
k6 = .15e5
k7 = .13e-3
k8 = .24e5
k9 = .165e5
k10 = .9e4
k11 = .22e-1
k12 = .12e5
k13 = .188e1
k14 = .163e5
k15 = .48e7
k16 = .35e-3
k17 = .175e-1
k18 = .1e9
k19 = .444e12
k20 = .124e4
k21 = .21e1
k22 = .578e1
k23 = .474e-1
k24 = .178e4
k25 = .312e1

function pollu!(dy, y, p, t)
    r1  = k1 * y[1]
    r2  = k2 * y[2] * y[4]
    r3  = k3 * y[5] * y[2]
    r4  = k4 * y[7]
    r5  = k5 * y[7]
    r6  = k6 * y[7] * y[6]
    r7  = k7 * y[9]
    r8  = k8 * y[9] * y[6]
    r9  = k9 * y[11] * y[2]
    r10 = k10 * y[11] * y[1]
    r11 = k11 * y[13]
    r12 = k12 * y[10] * y[2]
    r13 = k13 * y[14]
    r14 = k14 * y[1] * y[6]
    r15 = k15 * y[3]
    r16 = k16 * y[4]
    r17 = k17 * y[4]
    r18 = k18 * y[16]
    r19 = k19 * y[16]
    r20 = k20 * y[17] * y[6]
    r21 = k21 * y[19]
    r22 = k22 * y[19]
    r23 = k23 * y[1] * y[4]
    r24 = k24 * y[19] * y[1]
    r25 = k25 * y[20]

    dy[1]  = -r1 - r10 - r14 - r23 - r24 + r2 + r3 + r9 + r11 + r12 + r22 + r25
    dy[2]  = -r2 - r3 - r9 - r12 + r1 + r21
    dy[3]  = -r15 + r1 + r17 + r19 + r22
    dy[4]  = -r2 - r16 - r17 - r23 + r15
    dy[5]  = -r3 + r4 + r4 + r6 + r7 + r13 + r20
    dy[6]  = -r6 - r8 - r14 - r20 + r3 + r18 + r18
    dy[7]  = -r4 - r5 - r6 + r13
    dy[8]  = r4 + r5 + r6 + r7
    dy[9]  = -r7 - r8
    dy[10] = -r12 + r7 + r9
    dy[11] = -r9 - r10 + r8 + r11
    dy[12] = r9
    dy[13] = -r11 + r10
    dy[14] = -r13 + r12
    dy[15] = r14
    dy[16] = -r18 - r19 + r16
    dy[17] = -r20
    dy[18] = r20
    dy[19] = -r21 - r22 - r24 + r23 + r25
    dy[20] = -r25 + r24
end

u0 = zeros(20);
u0[2]  = 0.2;
u0[4]  = 0.04;
u0[7]  = 0.1;
u0[8]  = 0.3;
u0[9]  = 0.01;
u0[17] = 0.007;

tspan = (0.0, 60.0);
t_end = tspan[2];
k = [0.04, 3e7, 1e4];
tsteps = range(0, t_end, length=ntotal);
prob_rober = ODEProblem(pollu!, u0, tspan, k);
sol_rober = solve(prob_rober, ode_solver, saveat=tsteps, atol=1.f-6, rtol=1e-12);
normdata = Array(sol_rober)

i_slow = 1:20
nslow = length(i_slow)
yscale = maximum(normdata, dims=2)[i_slow] - minimum(normdata, dims=2)[i_slow]

nslow = length(i_slow);
node = 10;
dudt2 = Chain(x -> x,
              Dense(nslow, node, gelu),
              Dense(node, node, gelu),
              Dense(node, node, gelu),
              Dense(node, nslow))

p, re = Flux.destructure(dudt2);
rep = re(p)

function dudt!(du, u, p, t)
    du .= rep(u) .* yscale / t_end
end

prob = ODEProblem(dudt!, u0[i_slow], tspan)
sense = BacksolveAdjoint(checkpointing=true; autojacvec=ZygoteVJP());

function predict_n_ode(p, sample)
    global rep = re(p)
    _prob = remake(prob, p=p, tspan=[0, tsteps[sample]])
    pred = Array(solve(_prob, ode_solver, saveat=tsteps[1:sample], atol=lb, sensalg=sense))
end
pred = predict_n_ode(p, ntotal)

function loss_n_ode(p, sample=ntotal)
    pred = predict_n_ode(p, sample)
    loss = mae(pred ./ yscale, normdata[i_slow, 1:size(pred)[2]] ./ yscale)
    return loss
end
loss_n_ode(p, ntotal)

list_loss = []
list_grad = []
iter = 1
cb = function (p, loss_mean, g_norm)
    global list_loss, list_grad, iter
    push!(list_loss, loss_mean)
    push!(list_grad, g_norm)

    if iter % n_plot == 0
        pred = predict_n_ode(p, ntotal)

        list_plt = []
        for i in 1:nslow
            j = i_slow[i]
            plt = scatter(tsteps[:], normdata[j,:], xscale=:identity, label="data")
            plot!(plt, tsteps[:], pred[i,:], lw=2, xscale=:identity, label="pred", framestyle=:box)
            if i == 1
                plot!(plt, legend=:best)
            else
                plot!(plt, legend=false)
            end
            push!(list_plt, plt)
        end
        plt_all = plot(list_plt..., size=(1500, 1500))
        png(plt_all, "figs/pred.png")

        plt_loss = plot(list_loss, xscale=:identity, yscale=:log10, label="loss", legend=:topright)
        plt_grad = plot(list_grad, xscale=:identity, yscale=:log10, label="grad_norm", legend=:bottomright)
        xlabel!(plt_loss, "Epoch")
        xlabel!(plt_grad, "Epoch")
        ylabel!(plt_loss, "Loss")
        ylabel!(plt_grad, "Gradient Norm")
        plt_all = plot([plt_loss, plt_grad]..., framestyle=:box, size=(1000, 400))
        plot!(plt_all, xtickfontsize=10, ytickfontsize=10, xguidefontsize=12, yguidefontsize=12)
        png(plt_all, "figs/loss_grad")

        @save "./checkpoint/mymodel.bson" p opt list_loss list_grad iter
    end
    iter += 1
    return false
end

if is_restart
    @load "./checkpoint/mymodel.bson" p opt list_loss list_grad iter
    iter += 1
end

epochs = ProgressBar(iter:n_epoch);
for epoch in epochs
    global p
    sample = rand(batch_size:ntotal)

    loss = loss_n_ode(p, sample)
    grad = ForwardDiff.gradient(x -> loss_n_ode(x, sample), p)
    grad_norm = norm(grad, 2)
    grad = grad ./ grad_norm .* grad_max

    update!(opt, p, grad)

    set_description(epochs, string(@sprintf("Loss: %.4e grad: %.2e", loss, grad_norm)))
    cb(p, loss, grad_norm)
end

max_stars_repo_name: jiweiqi/StiffNeuralODE

max_issues_repo_name: jiweiqi/StiffNeuralODE

max_forks_repo_name: jiweiqi/StiffNeuralODE


text:
# TODO: add more target distributions and make them iteratable

# Dimension of testing distribution
const D = 5
# Tolerance ratio
const TRATIO = Int == Int64 ? 1 : 2
# Deterministic tolerance
const DETATOL = 1e-3 * D * TRATIO
# Random tolerance
const RNDATOL = 5e-2 * D * TRATIO * 2

# Hand-coded multivaite Gaussain

const gaussian_m = zeros(D)
const gaussian_s = ones(D)

function ℓπ(m, s, x::AbstractVecOrMat{T}) where {T}
    diff = x .- m
    v = s.^2
    return -(log(2 * T(pi)) .+ log.(v) .+ diff .* diff ./ v) / 2
end

function ℓπ(θ::AbstractVector)
    return sum(ℓπ(gaussian_m, gaussian_s, θ))
end

function ℓπ(θ::AbstractMatrix)
    return dropdims(sum(ℓπ(gaussian_m, gaussian_s, θ); dims=1); dims=1)
end

function ∂ℓπ∂θ(m, s, x::AbstractVecOrMat{T}) where {T}
    diff = x .- m
    v = s.^2
    v = -(log(2 * T(pi)) .+ log.(v) .+ diff .* diff ./ v) / 2
    g = -diff
    return v, g
end

function ∂ℓπ∂θ(θ::AbstractVector)
    v, g = ∂ℓπ∂θ(gaussian_m, gaussian_s, θ)
    return sum(v), g
end

function ∂ℓπ∂θ(θ::AbstractMatrix)
    v, g = ∂ℓπ∂θ(gaussian_m, gaussian_s, θ)
    return dropdims(sum(v; dims=1); dims=1), g
end

# For the Turing model
# @model gdemo() = 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

using Distributions: logpdf, InverseGamma, Normal
using Bijectors: invlink, logpdf_with_trans

function invlink_gdemo(θ)
    s = invlink(InverseGamma(2, 3), θ[1])
    m = θ[2]
    return [s, m]
end

function ℓπ_gdemo(θ)
    s, m = invlink_gdemo(θ)
    logprior = logpdf_with_trans(InverseGamma(2, 3), s, true) + logpdf(Normal(0, sqrt(s)), m)
    loglikelihood = logpdf(Normal(m, sqrt(s)), 1.5) + logpdf(Normal(m, sqrt(s)), 2.0)
    return logprior + loglikelihood
end

using Distributions: MvNormal
import Turing

Turing.@model function mvntest(θ=missing, x=missing)
    θ ~ MvNormal(zeros(D), 2)
    x ~ Normal(sum(θ), 1)
    return θ, x
end

function get_primitives(x, modelgen)
    spl_prior = Turing.SampleFromPrior()
    function ℓπ(θ)
        vi = Turing.VarInfo(model)
        vi[spl_prior] = θ
        model(vi, spl_prior)
        Turing.getlogp(vi)
    end
    adbackend = Turing.Core.ForwardDiffAD{40}
    alg_ad = Turing.HMC{adbackend}(0.1, 1)
    model = modelgen(missing, x)
    vi = Turing.VarInfo(model)
    spl = Turing.Sampler(alg_ad, model)
    Turing.Core.link!(vi, spl)
    ∂ℓπ∂θ = θ -> Turing.Core.gradient_logp(adbackend(), θ, vi, model, spl)
    θ₀ = Turing.VarInfo(model)[Turing.SampleFromPrior()]
    return ℓπ, ∂ℓπ∂θ, θ₀
end

function rand_θ_given(x, modelgen, metric, κ; n_samples=20)
    ℓπ, ∂ℓπ∂θ, θ₀ = get_primitives(x, modelgen)
    h = Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)
    samples, stats = sample(h, κ, θ₀, n_samples; verbose=false, progress=false)
    s = samples[end]
    return length(s) == 1 ? s[1] : s
end

# Test function
geweke_g(θ, x) = cat(θ, x; dims=1)


text:
A = zeros(UInt8, 10, 10, 10)
for i in 1:8
    A[i, 5, i:i+2] = 0x01
end
A[9:10, 5, 8:10] = 0x01

blob_series = form_blobs(A)
@test length(blob_series) == 10
for blobs in blob_series
    @test length(blobs) == 1
    blob = blobs[1]
    @test blob.x.p == 5
    @test blob.ratio ≈ 3.0
    @test blob.area == 3
    @test blob.source == 1
    @test blob.θ ≈ 0.0
    for shape in [blob.shape1, blob.shape2]
        @test shape[2:end, :] == [0.0 0.0; 0.0 0.0; 0.0 0.0]
        @test abs([0.0 , 1.0]' * shape[1, :]) == 1.0
    end
end


text:
@enum Shape CIRCLE RECTANGLE ROUNDED_RECTANGLE DISTANCEFIELD TRIANGLE
@enum CubeSides TOP BOTTOM FRONT BACK RIGHT LEFT

struct Grid{N, T <: AbstractRange}
    dims::NTuple{N, T}
end
Base.ndims(::Grid{N,T}) where {N,T} = N

Grid(ranges::AbstractRange...) = Grid(ranges)
function Grid(a::Array{T, N}) where {N, T}
    s = Vec{N, Float32}(size(a))
    smax = maximum(s)
    s = s./smax
    Grid(ntuple(Val{N}) do i
        range(0, stop=s[i], length=size(a, i))
    end)
end

Grid(a::AbstractArray, ranges...) = Grid(a, ranges)

"""
This constructor constructs a grid from ranges given as a tuple.
Due to the approach, the tuple `ranges` can consist of NTuple(2, T)
and all kind of range types. The constructor will make sure that all ranges match
the size of the dimension of the array `a`.
"""
function Grid(a::AbstractArray{T, N}, ranges::Tuple) where {T, N}
    length(ranges) =! N && throw(ArgumentError(
        "You need to supply a range for every dimension of the array. Given: $ranges
        given Array: $(typeof(a))"
    ))
    Grid(ntuple(Val(N)) do i
        range(first(ranges[i]), stop=last(ranges[i]), length=size(a, i))
    end)
end

Base.length(p::Grid) = prod(size(p))
Base.size(p::Grid) = map(length, p.dims)
function Base.getindex(p::Grid{N,T}, i) where {N,T}
    inds = ind2sub(size(p), i)
    Point{N, eltype(T)}(ntuple(Val(N)) do i
        p.dims[i][inds[i]]
    end)
end

Base.iterate(g::Grid, i = 1) = i <= length(g) ? (g[i], i + 1) : nothing

GLAbstraction.isa_gl_struct(x::Grid) = true
GLAbstraction.toglsltype_string(t::Grid{N,T}) where {N,T} = "uniform Grid$(N)D"
function GLAbstraction.gl_convert_struct(g::Grid{N, T}, uniform_name::Symbol) where {N,T}
    return Dict{Symbol, Any}(
        Symbol("$uniform_name.start") => Vec{N, Float32}(minimum.(g.dims)),
        Symbol("$uniform_name.stop") => Vec{N, Float32}(maximum.(g.dims)),
        Symbol("$uniform_name.lendiv") => Vec{N, Cint}(length.(g.dims) .- 1),
        Symbol("$uniform_name.dims") => Vec{N, Cint}(map(length, g.dims))
    )
end
function GLAbstraction.gl_convert_struct(g::Grid{1, T}, uniform_name::Symbol) where T
    x = g.dims[1]
    return Dict{Symbol, Any}(
        Symbol("$uniform_name.start") => Float32(minimum(x)),
        Symbol("$uniform_name.stop") => Float32(maximum(x)),
        Symbol("$uniform_name.lendiv") => Cint(length(x) - 1),
        Symbol("$uniform_name.dims") => Cint(length(x))
    )
end
import Base: getindex, length, iterate, ndims, setindex!, eltype

to_cpu_mem(x) = x
to_cpu_mem(x::GPUArray) = gpu_data(x)

const ScaleTypes = Union{Vector, Vec, AbstractFloat, Nothing, Grid}
const PositionTypes = Union{Vector, Point, AbstractFloat, Nothing, Grid}

mutable struct ScalarRepeat{T}
    scalar::T
end
ndims(::ScalarRepeat) = 1
getindex(s::ScalarRepeat, i...) = s.scalar
#should setindex! really be allowed? It will set the index for the whole row...
setindex!(s::ScalarRepeat{T}, value, i...) where {T} = (s.scalar = T(value))
eltype(::ScalarRepeat{T}) where {T} = T

iterate(sr::ScalarRepeat, i = 1) = (sr.scalar, i + 1)

struct Instances{P,T,S,R}
    primitive::P
    translation::T
    scale::S
    rotation::R
end

function _Instances(position,px,py,pz, scale,sx,sy,sz, rotation, primitive)
    args = (position,px,py,pz, scale,sx,sy,sz, rotation, primitive)
    args = map(to_cpu_mem, args)
    p = const_lift(ArrayOrStructOfArray, Point3f0, args[1:4]...)
    s = const_lift(ArrayOrStructOfArray, Vec3f0, args[5:8]...)
    r = const_lift(ArrayOrStructOfArray, Vec4f0, args[9])
    const_lift(Instances, args[10], p, s, r)
end
function _Instances(position, scale, rotation, primitive)
    p = const_lift(ArrayOrStructOfArray, Point3f0, position)
    s = const_lift(ArrayOrStructOfArray, Vec3f0, scale)
    r = const_lift(ArrayOrStructOfArray, Vec4f0, rotation)
    const_lift(Instances, primitive, p, s, r)
end

struct GridZRepeat{G, T, N} <: AbstractArray{Point{3, T}, N}
    grid::G
    z::Array{T, N}
end
Base.size(g::GridZRepeat) = size(g.z)
Base.size(g::GridZRepeat, i) = size(g.z, i)
Base.IndexStyle(::Type{<:GridZRepeat}) = Base.IndexLinear()

function Base.getindex(g::GridZRepeat{G, T}, i) where {G,T}
    pxy = g.grid[i]
    Point{3, T}(pxy[1], pxy[2], g.z[i])
end

function ArrayOrStructOfArray(::Type{T}, array::Nothing, a, elements...) where T
    StructOfArrays(T, a, elements...)
end
function ArrayOrStructOfArray(::Type{T}, array::StaticVector, a, elements...) where T
    StructOfArrays(T, a, elements...)
end
function ArrayOrStructOfArray(::Type{T}, scalar::StaticVector, a::Nothing, elements::Nothing...) where T
    ScalarRepeat(transform_convert(T, scalar))
end
function ArrayOrStructOfArray(::Type{T1}, array::Array{T2}, a::Nothing, elements::Nothing...) where {T1,T2}
    array
end
function ArrayOrStructOfArray(::Type{T1}, grid::Grid, x::Nothing, y::Nothing, z::Array) where T1<:Point
    GridZRepeat(grid, z)
end
function ArrayOrStructOfArray(::Type{T1}, array::Grid, a::Nothing, elements::Nothing...) where T1<:Point
    array
end
function ArrayOrStructOfArray(::Type{T}, scalar::T) where T
    ScalarRepeat(scalar)
end
function ArrayOrStructOfArray(::Type{T}, array::Array) where T
    array
end

struct TransformationIterator{T,S,R}
    translation::T
    scale::S
    rotation::R
end
function TransformationIterator(instances::Instances)
    TransformationIterator(
        instances.translation,
        instances.scale,
        instances.rotation
    )
end
function iterate(t::TransformationIterator, state = 1)
    states = (
        iterate(t.translation, state),
        iterate(t.scale, state),
        iterate(t.rotation, state)
    )
    all(x-> x !== nothing, states) || return nothing
    _translation, _scale, _rotation = first.(states)
    translation = Point3f0(transform_convert(Vec3f0, _translation))
    scale = Vec3f0(transform_convert(Point3f0, _scale))
    rotation = to_rotation_mat(_rotation)
    (translation, scale, rotation), last.(states)
end

import GeometryTypes: transform_convert

# For quaternions
function to_rotation_mat(x::Vec{4, T}) where T
    Mat4{T}(AbstractPlotting.Quaternionf0(x[1], x[2], x[3], x[4]))
end
# For relative rotations of a vector
function to_rotation_mat(x::StaticVector{2, T}) where T
    to_rotation_mat(Vec3f0(x[1], x[2], 0))
end
function to_rotation_mat(x::StaticVector{3, T}) where T
    rotation = Vec3f0(transform_convert(Point3f0, x))
    v, u = normalize(rotation), Vec3f0(0, 0, 1)
    # Unfortunately, we have to check for when u == -v, as u + v
    # in this case will be (0, 0, 0), which cannot be normalized.
    q = if (u == -v)
        # 180 degree rotation around any orthogonal vector
        other = (abs(dot(u, Vec{3, T}(1,0,0))) < 1.0) ? Vec{3, T}(1,0,0) : Vec{3, T}(0,1,0)
        AbstractPlotting.qrotation(normalize(cross(u, other)), T(180))
    else
        half = normalize(u+v)
        vc = cross(u, half)
        AbstractPlotting.Quaternionf0(dot(u, half), vc[1], vc[2], vc[3])
    end
    Mat4{T}(q)
end

struct Intensity{T <: AbstractFloat} <: FieldVector{1, T}
    i::T
end
@inline (I::Type{Intensity{T}})(i::Tuple) where {T <: AbstractFloat} = I(i...)
@inline (I::Type{Intensity{T}})(i::Intensity) where {T <: AbstractFloat} = I(i.i)
Intensity{T}(x::Color{Tc, 1}) where {T <: AbstractFloat, Tc} = Intensity{T}(gray(x))

const GLIntensity = Intensity{Float32}
export Intensity, GLIntensity

struct GLVisualizeShader <: AbstractLazyShader
    paths::Tuple
    kw_args::Dict{Symbol, Any}
    function GLVisualizeShader(paths::String...; view = Dict{String, String}(), kw_args...)
        # TODO properly check what extensions are available
        @static if !Sys.isapple()
            view["GLSL_EXTENSIONS"] = "#extension GL_ARB_conservative_depth: enable"
            view["SUPPORTED_EXTENSIONS"] = "#define DETPH_LAYOUT"
        end
        args = Dict{Symbol, Any}(kw_args)
        args[:view] = view
        args[:fragdatalocation] = [(0, "fragment_color"), (1, "fragment_groupid")]
        new(map(x-> assetpath("shader", x), paths), args)
    end
end


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

mutable struct Int64Parameter <: ProtoType
    __protobuf_jl_internal_meta::ProtoMeta
    __protobuf_jl_internal_values::Dict{Symbol,Any}
    __protobuf_jl_internal_defaultset::Set{Symbol}

    function Int64Parameter(; kwargs...)
        obj = new(meta(Int64Parameter), Dict{Symbol,Any}(), Set{Symbol}())
        values = obj.__protobuf_jl_internal_values
        symdict = obj.__protobuf_jl_internal_meta.symdict
        for nv in kwargs
            fldname, fldval = nv
            fldtype = symdict[fldname].jtyp
            (fldname in keys(symdict)) || error(string(typeof(obj), " has no field with name ", fldname))
            if fldval !== nothing
                values[fldname] = isa(fldval, fldtype) ? fldval : convert(fldtype, fldval)
            end
        end
        obj
    end
end # mutable struct Int64Parameter
const __meta_Int64Parameter = Ref{ProtoMeta}()
function meta(::Type{Int64Parameter})
    ProtoBuf.metalock() do
        if !isassigned(__meta_Int64Parameter)
            __meta_Int64Parameter[] = target = ProtoMeta(Int64Parameter)
            fnum = Int[1,10,11]
            allflds = Pair{Symbol,Union{Type,String}}[:defaultValue => Int64, :range => Int64Range, :set => Int64Set]
            oneofs = Int[0,1,1]
            oneof_names = Symbol[Symbol("AllowedValues")]
            meta(target, Int64Parameter, allflds, ProtoBuf.DEF_REQ, fnum, ProtoBuf.DEF_VAL, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, oneofs, oneof_names)
        end
        __meta_Int64Parameter[]
    end
end
function Base.getproperty(obj::Int64Parameter, name::Symbol)
    if name === :defaultValue
        return (obj.__protobuf_jl_internal_values[name])::Int64
    elseif name === :range
        return (obj.__protobuf_jl_internal_values[name])::Int64Range
    elseif name === :set
        return (obj.__protobuf_jl_internal_values[name])::Int64Set
    else
        getfield(obj, name)
    end
end

mutable struct DoubleParameter <: ProtoType
    __protobuf_jl_internal_meta::ProtoMeta
    __protobuf_jl_internal_values::Dict{Symbol,Any}
    __protobuf_jl_internal_defaultset::Set{Symbol}

    function DoubleParameter(; kwargs...)
        obj = new(meta(DoubleParameter), Dict{Symbol,Any}(), Set{Symbol}())
        values = obj.__protobuf_jl_internal_values
        symdict = obj.__protobuf_jl_internal_meta.symdict
        for nv in kwargs
            fldname, fldval = nv
            fldtype = symdict[fldname].jtyp
            (fldname in keys(symdict)) || error(string(typeof(obj), " has no field with name ", fldname))
            if fldval !== nothing
                values[fldname] = isa(fldval, fldtype) ? fldval : convert(fldtype, fldval)
            end
        end
        obj
    end
end # mutable struct DoubleParameter
const __meta_DoubleParameter = Ref{ProtoMeta}()
function meta(::Type{DoubleParameter})
    ProtoBuf.metalock() do
        if !isassigned(__meta_DoubleParameter)
            __meta_DoubleParameter[] = target = ProtoMeta(DoubleParameter)
            fnum = Int[1,10]
            allflds = Pair{Symbol,Union{Type,String}}[:defaultValue => Float64, :range => DoubleRange]
            oneofs = Int[0,1]
            oneof_names = Symbol[Symbol("AllowedValues")]
            meta(target, DoubleParameter, allflds, ProtoBuf.DEF_REQ, fnum, ProtoBuf.DEF_VAL, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, oneofs, oneof_names)
        end
        __meta_DoubleParameter[]
    end
end
function Base.getproperty(obj::DoubleParameter, name::Symbol)
    if name === :defaultValue
        return (obj.__protobuf_jl_internal_values[name])::Float64
    elseif name === :range
        return (obj.__protobuf_jl_internal_values[name])::DoubleRange
    else
        getfield(obj, name)
    end
end

mutable struct StringParameter <: ProtoType
    __protobuf_jl_internal_meta::ProtoMeta
    __protobuf_jl_internal_values::Dict{Symbol,Any}
    __protobuf_jl_internal_defaultset::Set{Symbol}

    function StringParameter(; kwargs...)
        obj = new(meta(StringParameter), Dict{Symbol,Any}(), Set{Symbol}())
        values = obj.__protobuf_jl_internal_values
        symdict = obj.__protobuf_jl_internal_meta.symdict
        for nv in kwargs
            fldname, fldval = nv
            fldtype = symdict[fldname].jtyp
            (fldname in keys(symdict)) || error(string(typeof(obj), " has no field with name ", fldname))
            if fldval !== nothing
                values[fldname] = isa(fldval, fldtype) ? fldval : convert(fldtype, fldval)
            end
        end
        obj
    end
end # mutable struct StringParameter
const __meta_StringParameter = Ref{ProtoMeta}()
function meta(::Type{StringParameter})
    ProtoBuf.metalock() do
        if !isassigned(__meta_StringParameter)
            __meta_StringParameter[] = target = ProtoMeta(StringParameter)
            allflds = Pair{Symbol,Union{Type,String}}[:defaultValue => AbstractString]
            meta(target, StringParameter, allflds, ProtoBuf.DEF_REQ, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES)
        end
        __meta_StringParameter[]
    end
end
function Base.getproperty(obj::StringParameter, name::Symbol)
    if name === :defaultValue
        return (obj.__protobuf_jl_internal_values[name])::AbstractString
    else
        getfield(obj, name)
    end
end

mutable struct BoolParameter <: ProtoType
    __protobuf_jl_internal_meta::ProtoMeta
    __protobuf_jl_internal_values::Dict{Symbol,Any}
    __protobuf_jl_internal_defaultset::Set{Symbol}

    function BoolParameter(; kwargs...)
        obj = new(meta(BoolParameter), Dict{Symbol,Any}(), Set{Symbol}())
        values = obj.__protobuf_jl_internal_values
        symdict = obj.__protobuf_jl_internal_meta.symdict
        for nv in kwargs
            fldname, fldval = nv
            fldtype = symdict[fldname].jtyp
            (fldname in keys(symdict)) || error(string(typeof(obj), " has no field with name ", fldname))
            if fldval !== nothing
                values[fldname] = isa(fldval, fldtype) ? fldval : convert(fldtype, fldval)
            end
        end
        obj
    end
end # mutable struct BoolParameter
const __meta_BoolParameter = Ref{ProtoMeta}()
function meta(::Type{BoolParameter})
    ProtoBuf.metalock() do
        if !isassigned(__meta_BoolParameter)
            __meta_BoolParameter[] = target = ProtoMeta(BoolParameter)
            allflds = Pair{Symbol,Union{Type,String}}[:defaultValue => Bool]
            meta(target, BoolParameter, allflds, ProtoBuf.DEF_REQ, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES)
        end
        __meta_BoolParameter[]
    end
end
function Base.getproperty(obj::BoolParameter, name::Symbol)
    if name === :defaultValue
        return (obj.__protobuf_jl_internal_values[name])::Bool
    else
        getfield(obj, name)
    end
end

export Int64Parameter, DoubleParameter, StringParameter, BoolParameter


text:
"""
    VTKCellTypes

Module defining cell types for unstructured datasets.

Definitions are adapted from the [`VTK source
code`](https://raw.githubusercontent.com/Kitware/VTK/master/Common/DataModel/vtkCellType.h).
"""
module VTKCellTypes

export VTKCellType, nodes

struct VTKCellType
    vtk_name::String
    vtk_id::UInt8
    nodes::Int
end

"""
    nodes(c::VTKCellTypes)

Returns the number of nodes (or grid points) required by the cell type.

For instance, this returns 3 for `VTK_TRIANGLE`.

For cell types that can take any number of nodes, such as `VTK_POLY_LINE`,
this returns -1.
"""
nodes(c::VTKCellType) = c.nodes

const INVALID_CELL_TYPE = VTKCellType("INVALID", typemax(UInt8), -1)

# Array of known cell types
const VTK_CELL_TYPES = fill(INVALID_CELL_TYPE, typemax(UInt8))

# Construction of VTKCellType objects from the integer vtk id
function VTKCellType(vtkid::Integer)
    n = vtkid + 1  # ids start at zero
    cell = VTK_CELL_TYPES[n]
    if cell === INVALID_CELL_TYPE
        throw(ArgumentError("unknown vtkid = $vtkid"))
    end
    cell
end

# Define a VTK cell type
macro add(ctype, id, nodes=-1)
    name = string(ctype)
    quote
        const $(esc(ctype)) = VTKCellType($name, $id, $nodes)
        VTK_CELL_TYPES[$id + 1] = $(esc(ctype))
    end
end

# Linear cells
@add VTK_EMPTY_CELL 0 0
@add VTK_VERTEX 1  1
@add VTK_POLY_VERTEX 2 -1
@add VTK_LINE 3  2
@add VTK_POLY_LINE 4 -1
@add VTK_TRIANGLE 5  3
@add VTK_TRIANGLE_STRIP 6 -1
@add VTK_POLYGON 7 -1
@add VTK_PIXEL 8  4
@add VTK_QUAD 9  4
@add VTK_TETRA 10 4
@add VTK_VOXEL 11 8
@add VTK_HEXAHEDRON 12 8
@add VTK_WEDGE 13 6
@add VTK_PYRAMID 14 5
@add VTK_PENTAGONAL_PRISM 15 10
@add VTK_HEXAGONAL_PRISM 16 12

# Quadratic, isoparametric cells
@add VTK_QUADRATIC_EDGE 21 3
@add VTK_QUADRATIC_TRIANGLE 22 6
@add VTK_QUADRATIC_QUAD 23 8
@add VTK_QUADRATIC_POLYGON 36 -1
@add VTK_QUADRATIC_TETRA 24 10
@add VTK_QUADRATIC_HEXAHEDRON 25 20
@add VTK_QUADRATIC_WEDGE 26 15
@add VTK_QUADRATIC_PYRAMID 27 13
@add VTK_BIQUADRATIC_QUAD 28 9
@add VTK_TRIQUADRATIC_HEXAHEDRON 29 27
@add VTK_QUADRATIC_LINEAR_QUAD 30 6
@add VTK_QUADRATIC_LINEAR_WEDGE 31 12
@add VTK_BIQUADRATIC_QUADRATIC_WEDGE 32 18
@add VTK_BIQUADRATIC_QUADRATIC_HEXAHEDRON 33 24
@add VTK_BIQUADRATIC_TRIANGLE 34 7

# Cubic, isoparametric cell
@add VTK_CUBIC_LINE 35 4

# Special class of cells formed by convex group of points
@add VTK_CONVEX_POINT_SET 41

# Polyhedron cell (consisting of polygonal faces)
@add VTK_POLYHEDRON 42

# Higher order cells in parametric form
@add VTK_PARAMETRIC_CURVE 51
@add VTK_PARAMETRIC_SURFACE 52
@add VTK_PARAMETRIC_TRI_SURFACE 53
@add VTK_PARAMETRIC_QUAD_SURFACE 54
@add VTK_PARAMETRIC_TETRA_REGION 55
@add VTK_PARAMETRIC_HEX_REGION 56

# Higher order cells
@add VTK_HIGHER_ORDER_EDGE 60
@add VTK_HIGHER_ORDER_TRIANGLE 61
@add VTK_HIGHER_ORDER_QUAD 62
@add VTK_HIGHER_ORDER_POLYGON 63
@add VTK_HIGHER_ORDER_TETRAHEDRON 64
@add VTK_HIGHER_ORDER_WEDGE 65
@add VTK_HIGHER_ORDER_PYRAMID 66
@add VTK_HIGHER_ORDER_HEXAHEDRON 67

# Arbitrary order Lagrange elements (formulated separated from generic higher order cells)
@add VTK_LAGRANGE_CURVE 68
@add VTK_LAGRANGE_TRIANGLE 69
@add VTK_LAGRANGE_QUADRILATERAL 70
@add VTK_LAGRANGE_TETRAHEDRON 71
@add VTK_LAGRANGE_HEXAHEDRON 72
@add VTK_LAGRANGE_WEDGE 73
@add VTK_LAGRANGE_PYRAMID 74

# Arbitrary order Bezier elements (formulated separated from generic higher order cells)
@add VTK_BEZIER_CURVE 75
@add VTK_BEZIER_TRIANGLE 76
@add VTK_BEZIER_QUADRILATERAL 77
@add VTK_BEZIER_TETRAHEDRON 78
@add VTK_BEZIER_HEXAHEDRON 79
@add VTK_BEZIER_WEDGE 80
@add VTK_BEZIER_PYRAMID 81

end # module


text:
function Base.show(io::IO, C::ClassField_pp{S, T}) where {S, T}
  println(IOContext(io, :compact => true), "Cyclic class field of degree $(degree(C)) defined modulo $(defining_modulus(C))")
  if isdefined(C, :a)
    println(io, "Kummer generator ", C.a)
  end
  if isdefined(C, :K)
    println(io, "Kummer extension: ", C.K)
  end
  if isdefined(C, :A)
    println(io, "Defining  polynomial ", C.A.pol)
  end
end

function Base.show(io::IO, CF::ClassField)
  @show_name(io, CF)
  @show_special(io, CF)
  print(IOContext(io, :compact => true), "Class field defined mod ",
                   defining_modulus(CF), " of structure ",
                   codomain(CF.quotientmap))
end

###############################################################################
#
#  Modulus function
#
###############################################################################

@doc Markdown.doc"""
    defining_modulus(CF::ClassField)

The modulus, i.e. an ideal of the set of real places, used to create the
class field.
"""
function defining_modulus(CF::ClassField)
  return _modulus(CF.rayclassgroupmap)
end

function defining_modulus(CF::ClassField_pp)
  return _modulus(CF.rayclassgroupmap)
end

function _modulus(mq::MapRayClassGrp)
  return mq.defining_modulus::Tuple{NfOrdIdl, Vector{InfPlc}}
end

function _modulus(mq::MapClassGrp)
  return (ideal(order(codomain(mq)), 1), InfPlc[])::Tuple{NfOrdIdl, Vector{InfPlc}}
end

###############################################################################
#
#  Base ring and Base field
#
###############################################################################

@doc Markdown.doc"""
    base_ring(A::ClassField)

The maximal order of the field that $A$ is defined over.
"""
function base_ring(A::ClassField)
  return order(codomain(A.rayclassgroupmap))
end

@doc Markdown.doc"""
    base_field(A::ClassField)

The number field that $A$ is defined over.
"""
function base_field(A::ClassField)
  return number_field(base_ring(A))
end

function base_ring(A::ClassField_pp)
  return order(codomain(A.rayclassgroupmap))
end

function base_field(A::ClassField_pp)
  return number_field(base_ring(A))
end

###############################################################################
#
#  Degree
#
###############################################################################

@doc Markdown.doc"""
    degree(A::ClassField)

The degree of $A$ over its base field, i.e. the size of the defining ideal group.
"""
function degree(A::ClassField)
  if A.degree == -1
    A.degree = Int(order(codomain(A.quotientmap)))
  end
  return A.degree
end

function degree(A::ClassField_pp{S, T}) where {S, T}
  if A.degree == -1
    A.degree = Int(order(codomain(A.quotientmap)))
  end
  return A.degree
end

###############################################################################
#
#  Exponent
#
###############################################################################

@doc Markdown.doc"""
    exponent(A::ClassField)

The exponent of $A$ over its base field, i.e. the exponent of the Galois
group of the extension.
"""
function exponent(A::ClassField{S, T}) where {S, T}
  return Int(exponent(codomain(A.quotientmap)))
end

function exponent(A::ClassField_pp{S, T}) where {S, T}
  return Int(exponent(codomain(A.quotientmap)))
end

###############################################################################
#
#  Compositum
#
###############################################################################

@doc Markdown.doc"""
    compositum(a::ClassField, b::ClassField) -> ClassField

The compositum of $a$ and $b$ as a (formal) class field.
"""
function compositum(a::ClassField, b::ClassField)
  @assert base_ring(a) == base_ring(b)
  c = lcm(defining_modulus(a)[1], defining_modulus(b)[1])
  d = lcm(degree(a), degree(b))
  c_inf = union(defining_modulus(a)[2], defining_modulus(b)[2])
  r, mr = ray_class_group(c, c_inf, n_quo = Int(d))
  C = ray_class_field(mr)
  @assert domain(C.rayclassgroupmap) == r
  h = norm_group_map(C, [a,b])
  U = intersect(kernel(h[1])[1], kernel(h[2])[1])
  q, mq = quo(codomain(C.quotientmap), U)
  return ray_class_field(mr, GrpAbFinGenMap(C.quotientmap * mq))
end

@doc Markdown.doc"""
    *(A::ClassField, B::ClassField) -> ClassField

The compositum of $a$ and $b$ as a (formal) class field.
"""
*(a::ClassField, b::ClassField) = compositum(a, b)

###############################################################################
#
#  Intersection
#
###############################################################################

@doc Markdown.doc"""
    intersect(a::ClassField, b::ClassField) -> ClassField

The intersection of $a$ and $b$ as a class field.
"""
function Base.intersect(a::ClassField, b::ClassField)
  @assert base_ring(a) == base_ring(b)
  c = lcm(defining_modulus(a)[1], defining_modulus(b)[1])
  c_inf = union(defining_modulus(a)[2], defining_modulus(b)[2])
  d = lcm(degree(a), degree(b))

  r, mr = ray_class_group(c, c_inf, n_quo = Int(d))
  C = ray_class_field(mr)
  h = norm_group_map(C, [a,b])
  U = kernel(h[1])[1] + kernel(h[2])[1]
  q, mq = quo(codomain(C.quotientmap), U)
  return ray_class_field(mr, GrpAbFinGenMap(C.quotientmap * mq))
end

###############################################################################
#
#  Subfield
#
###############################################################################

@doc Markdown.doc"""
    issubfield(a::ClassField, b::ClassField) -> Bool

Determines if $a$ is a subfield of $b$.
"""
function issubfield(a::ClassField, b::ClassField)
  @assert base_ring(a) == base_ring(b)
  c = lcm(defining_modulus(a)[1], defining_modulus(b)[1])
  c_inf = union(defining_modulus(a)[2], defining_modulus(b)[2])
  d = lcm(degree(a), degree(b))

  r, mr = ray_class_group(c, c_inf, n_quo = Int(d))
  C = ray_class_field(mr)
  h = norm_group_map(C, [a,b])
  return issubset(kernel(h[2])[1], kernel(h[1])[1])
end

###############################################################################
#
#  Equality
#
###############################################################################

@doc Markdown.doc"""
    ==(a::ClassField, b::ClassField)

Tests if $a$ and $b$ are equal.
"""
function ==(a::ClassField, b::ClassField)
  @assert base_ring(a) == base_ring(b)
  mq1 = a.quotientmap
  mq2 = b.quotientmap
  if !isisomorphic(codomain(mq1), codomain(mq2))
    return false
  end
  expo = Int(exponent(codomain(mq1)))
  c = lcm(defining_modulus(a)[1], defining_modulus(b)[1])
  c_inf = union(defining_modulus(a)[2], defining_modulus(b)[2])
  r, mr = ray_class_group(c, c_inf, n_quo = expo)
  C = ray_class_field(mr)
  @assert defining_modulus(C) == (c, c_inf)
  h = norm_group_map(C, [a,b])
  return iseq(kernel(h[2])[1], kernel(h[1])[1])
end

###############################################################################
#
#  IsCyclic
#
###############################################################################

@doc Markdown.doc"""
    iscyclic(C::ClassField)

Tests if the (relative) automorphism group of $C$ is cyclic (by checking
the defining ideal group).
"""
function iscyclic(C::ClassField)
  mp = C.quotientmap
  return iscyclic(codomain(mp))
end

###############################################################################
#
#  Maximal p-subfield
#
###############################################################################

@doc Markdown.doc"""
    maximal_p_subfield(C::ClassField, p::Int)

Returns the class field corresponding to the maximal subextension of
prime power degree.
"""
function maximal_p_subfield(C::ClassField, p::Int)
  mg = C.quotientmap
  v = valuation(degree(C), p)
  q, mq = quo(codomain(mg), p^v, false)
  s, ms = snf(q)
  mg1 = mg*mq*inv(ms)
  return ray_class_field(C.rayclassgroupmap, mg1)
end

@doc Markdown.doc"""
    islocal_norm(r::ClassField, a::NfAbsOrdElem, p::NfAbsOrdIdl) -> Bool

Tests if $a$ is a local norm at $p$ in the extension implictly given by $r$.
Currently the conductor cannot have infinite places.
"""
function islocal_norm(r::ClassField, a::NfAbsOrdElem, p::NfAbsOrdIdl)
  m0, minf = conductor(r)
  if length(minf) > 0
    error("not implemented yet")
  end
  m0 = defining_modulus(r)[1] #need the maps...
  @assert isprime(p)
  v1 = valuation(a, p)
  v2 = valuation(m0, p)
  n0 = divexact(m0, p^v2)
  o0 = p^(v1 + v2)
  y = crt(order(p)(1), n0, a, o0)
  y = y*order(p)
  y = divexact(y, p^v1)

  return isone(r.quotientmap(preimage(r.rayclassgroupmap, y)))
end

################################################################################
#
#  Some applications
#
################################################################################

@doc Markdown.doc"""
    islocal_norm(r::ClassField, a::NfAbsOrdElem) -> Bool

Tests if $a$ is a local norm at all finite places in the extension implictly given by $r$.
"""
function islocal_norm(r::ClassField, a::NfAbsOrdElem)
  K = base_field(r)
  m0, minf = conductor(r)
  if !ispositive(a, minf)
    return false
  end
  fl = factor(m0*a)
  return all(x -> islocal_norm(r, a, x), keys(fl))
end

@doc Markdown.doc"""
    prime_decomposition_type(C::ClassField, p::NfAbsOrdIdl) -> (Int, Int, Int)

For a prime $p$ in the base ring of $r$, determine the splitting type of $p$
in $r$. ie. the tuple $(e, f, g)$ giving the ramification degree, the inertia
and the number of primes above $p$.
"""
function prime_decomposition_type(C::T, p::NfAbsOrdIdl) where T <: Union{ClassField, ClassField_pp}
  @hassert :ClassField 1 isprime(p)
  mR = C.rayclassgroupmap
  m0 = defining_modulus(C)[1]
  R = domain(mR)

  v = valuation(m0, p)
  if v == 0
    f = Int(order(C.quotientmap(mR\p)))
    return (1, f, divexact(degree(C), f))
  end
  r, mr = ray_class_group(divexact(m0, p^v), defining_modulus(C)[2], n_quo = Int(exponent(R)))
  lp, sR = find_gens(mR, coprime_to = minimum(m0))
  h = hom(sR, [preimage(mr, p) for p = lp])
  k, mk = kernel(GrpAbFinGenMap(C.quotientmap), false)
  q, mq = cokernel(mk*h, false)
  f = Int(order(mq(preimage(mr, p))))
  e = Int(divexact(degree(C), Int(order(q))))
  return (e, f, Int(divexact(order(q), f)))
end

###############################################################################
#
#  Ray Class Field interface
#
###############################################################################

@doc Markdown.doc"""
    ray_class_field(m::MapClassGrp) -> ClassField
    ray_class_field(m::MapRayClassGrp) -> ClassField

Creates the (formal) abelian extension defined by the map $m: A \to I$
where $I$ is the set of ideals coprime to the modulus defining $m$ and $A$
is a quotient of the ray class group (or class group). The map $m$
must be the map returned from a call to {class_group} or {ray_class_group}.
"""
function ray_class_field(m::Union{MapClassGrp, MapRayClassGrp})
  return ray_class_field(m, id_hom(domain(m)))
end

@doc Markdown.doc"""
    ray_class_field(m::Union{MapClassGrp, MapRayClassGrp}, quomap::GrpAbFinGenMap) -> ClassField

For $m$ a map computed by either {ray_class_group} or {class_group} and
$q$ a canonical projection (quotient map) as returned by {quo} for q
quotient of the domain of $m$ and a subgroup of $m$, create the
(formal) abelian extension where the (relative) automorphism group
is canonically isomorphic to the codomain of $q$.
"""
function ray_class_field(m::S, quomap::T) where {S <: Union{MapClassGrp, MapRayClassGrp}, T}
  domain(quomap) == domain(m) || error("1st map must be a (ray)class group map, and the 2nd must be a projection of the domain of the 1st")
  CF = ClassField{S, T}()
  CF.rayclassgroupmap = m
  D = codomain(quomap)
  S1, mS1 = snf(D)
  iS1 = inv(mS1)#GrpAbFinGenMap(D, S1, mS1.imap, mS1.map)
  CF.quotientmap = Hecke.compose(quomap, iS1)
  return CF
end

@doc Markdown.doc"""
    hilbert_class_field(k::AnticNumberField) -> ClassField

The Hilbert class field of $k$ as a formal (ray-) class field.
"""
function hilbert_class_field(k::AnticNumberField)
  return ray_class_field(class_group(k)[2])
end

@doc Markdown.doc"""
    ray_class_field(I::NfAbsOrdIdl; n_quo = 0) -> ClassField

The ray class field modulo $I$. If `n_quo` is given, then the largest
subfield of exponent $n$ is computed.
"""
function ray_class_field(I::NfAbsOrdIdl; n_quo = -1)
  return ray_class_field(ray_class_group(I, n_quo = n_quo)[2])
end

@doc Markdown.doc"""
    ray_class_field(I::NfAbsOrdIdl, inf::Array{InfPlc, 1}; n_quo = 0) -> ClassField

The ray class field modulo $I$ and the infinite places given. If `n_quo` is given, then the largest
subfield of exponent $n$ is computed.
"""
function ray_class_field(I::NfAbsOrdIdl, inf::Array{InfPlc, 1}; n_quo = -1)
  return ray_class_field(ray_class_group(I, inf, n_quo = n_quo)[2])
end


text:
struct Table{names, T <: Tuple{Vararg{AbstractVector}}, I <: Tuple{Vararg{Index}}} <: Relation{names}
    data::T
    indexes::I
end

@inline Table(in::Pair{<:Label}...) = Table{getnames(in)}(getvalues(in))
@inline Table{names}(in::Tuple) where {names} = Table{names, typeof(in)}(in)

@inline start(t::Table) = start(first(t.data))
@inline next(t::Table, i) = next(first(t.data), i)
@inline done(t::Table, i) = done(first(t.data), i)

@inline length(t::Table) = length(first(t.data))
@inline endof(t::Table) = endof(first(t.data))

@inline first(t::Table) = @inbounds return t[1]
@inline last(t::Table) = @inbounds return t[end]

@generated function getindex(t::Table{names}, i::Integer) where {names}
    expr = [:(t.data[$j][i]) for j = 1:length(names)]
    return quote
        @_propagate_inbounds_meta
        Row{names}(tuple($(expr...)))
    end
end

@generated function getindex(t::Table{names}, i) where {names}
    expr = [:(t.data[$j][i]) for j = 1:length(names)]
    return quote
        @_propagate_inbounds_meta
        Table{names}(tuple($(expr...)))
    end
end

function show(io::IO, table::Table)
    n = length(table)
    println(io, "Table with $n rows")
    for i = 1:min(n, 5)
        print(io, " ", table[i])
        if i != min(n, 5)
            print(io, "\n")
        end
    end
    if n > 5
        print(io, "\n ...")
    end
end


text:
# Autogenerated wrapper script for ADIOS2_jll for powerpc64le-linux-gnu-libgfortran4-cxx11
export bpls, libadios2_c, libadios2_c_mpi, libadios2_core, libadios2_core_mpi, libadios2_cxx11, libadios2_cxx11_mpi

using Blosc_jll
using Bzip2_jll
using CompilerSupportLibraries_jll
using MPICH_jll
using ZeroMQ_jll
using libpng_jll
using zfp_jll
JLLWrappers.@generate_wrapper_header("ADIOS2")
JLLWrappers.@declare_library_product(libadios2_c, "libadios2_c.so.2")
JLLWrappers.@declare_library_product(libadios2_c_mpi, "libadios2_c_mpi.so.2")
JLLWrappers.@declare_library_product(libadios2_core, "libadios2_core.so.2")
JLLWrappers.@declare_library_product(libadios2_core_mpi, "libadios2_core_mpi.so.2")
JLLWrappers.@declare_library_product(libadios2_cxx11, "libadios2_cxx11.so.2")
JLLWrappers.@declare_library_product(libadios2_cxx11_mpi, "libadios2_cxx11_mpi.so.2")
JLLWrappers.@declare_executable_product(bpls)
function __init__()
    JLLWrappers.@generate_init_header(Blosc_jll, Bzip2_jll, CompilerSupportLibraries_jll, MPICH_jll, ZeroMQ_jll, libpng_jll, zfp_jll)
    JLLWrappers.@init_library_product(
        libadios2_c,
        "lib/libadios2_c.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libadios2_c_mpi,
        "lib/libadios2_c_mpi.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libadios2_core,
        "lib/libadios2_core.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libadios2_core_mpi,
        "lib/libadios2_core_mpi.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libadios2_cxx11,
        "lib/libadios2_cxx11.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libadios2_cxx11_mpi,
        "lib/libadios2_cxx11_mpi.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_executable_product(
        bpls,
        "bin/bpls",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using PhotometricFilters
using Test

@testset "PhotometricFilters.jl" begin
    # Write your tests here.
end


text:
# This file holds the definition of the Retirement type. This type is used to
#   describe the particulars of retirement due to age and/or tenure.

export Retirement
"""
The `Retirement` type defines a retirement scheme for mandatory retirement of personnel members based on age or tenure. This type is used only to define the default retirement scheme, and will generally be overruled by OUT transitions attached to specific nodes.

The type contains the following fields:
* `maxCareerLength::Float64`: the maximum length of a personnel member's career. A value of 0 means this criterion is ignored. Default = 0.0
* `retirementAge::Float64`: the mandatory retirement age. A value of 0 means this criterion is ignored. Default = 0.0
* 'isEither::Bool': a flag indicating whether either of the criteria (age or tenure) must be satisfied for retirement, or both. Default = `true`
* `freq::Float64`: the length of the time interval between two retirement checks. Default = 1.0
* `offset::Float64`: the offset of the retirement schedule with respect to the start of the simulation. Default = 0.0

Constructor:
```
Retirement()
```
This constructor creates a `Retirement` object with the above mentioned default values for the parameters.
"""
mutable struct Retirement

    maxCareerLength::Float64
    retirementAge::Float64
    isEither::Bool
    freq::Float64
    offset::Float64

    function Retirement()::Retirement

        newRet = new()
        newRet.maxCareerLength = 0.0
        newRet.retirementAge = 0.0
        newRet.isEither = true
        newRet.freq = 1.0
        newRet.offset = 0.0
        return newRet

    end  # Retirement()

end  # mutable struct Retirement


text:
const current_called =
    fill(Trace[]) # the current vector onto which to tack the next function calls

# TODO: consider @memoize Dict{Tuple{Expr}, Any}
# But I need to make sure that Revise doesn't modify the Expr
splitdef_memo(def::Expr) = MacroTools.splitdef(def)

is_traceable(def) =
    (is_generated_function_definition(def) ?
     is_traceable(generated2normal(def)) :
     (is_function_definition(def) && (di=splitdef_memo(def); !is_call_definition(di)) &&
      !is_fancy_constructor_definition(di)))

""" `@traceable_loose(expr)` is like `traceable(expr)`, but doesn't error on non-function
definitions (it's a helper for `@traceable begin ... end`) """
macro traceable_loose(expr)
    is_traceable(expr) ? esc(:($TraceCalls.@traceable $expr)) : esc(expr)
end

global gensym_counter = 0  # necessary because of JuliaLang#23809
my_gensym() = Symbol(:__tc_gensym_, (global gensym_counter+=1;))

""" Turns `::Int=5` (and `_`) into `some_gensym::Int=5` """
function handle_missing_arg(arg)
    arg_name, arg_type, is_splat, default = splitarg(arg)
    combinearg((arg_name === nothing || arg_name == :_) ? my_gensym() : arg_name,
               arg_type, is_splat, default)
end

""" Overload `TraceCalls.store(x)` for specific types to change how `TraceCalls` stores
its arguments. For example, `TraceCalls.store(x::Vector) = copy(x)`. See also ?REPR. """
store(x) = x

arg_name(arg) = splitarg(arg)[1]
is_splat(arg) = splitarg(arg)[3]
arg_name_splat(arg) = is_splat(arg) ? Expr(:..., arg_name(arg)) : arg_name(arg)
function typed_arg(arg)
    name, arg_type = splitarg(arg)
    if is_splat(arg) arg_type = Any end
    return :($name::$arg_type)
end

function tracing_fun(splitdef_di::Dict, do_body_name::Symbol)::Dict
    # Returns a function definition that stores its arguments as a Trace and passes
    # them to do_body_name
    di = splitdef_di
    di[:args] = map(handle_missing_arg, di[:args])
    di[:kwargs] = map(handle_missing_arg, di[:kwargs])
    fname = di[:name]

    updated_fun_di = copy(di)
    passed_args = map(arg_name_splat, di[:args])
    passed_kwargs = [is_splat(kwa) ?
                     arg_name_splat(kwa) :
                     :(($(Expr(:quote, arg_name(kwa))), $(arg_name(kwa))))
                     for kwa in di[:kwargs]]
    @gensym prev_called trace_args trace_kwargs e res called
    updated_fun_di[:body] = quote
        local $res
        $prev_called = $TraceCalls.current_called[]
        $trace_args = map($TraceCalls.store, ($(passed_args...),))
        # We don't call store on kwargs. Does any function ever mutate a kwarg? Seems
        # unlikely.
        $trace_kwargs = ($(passed_kwargs...),)
        $TraceCalls.current_called[] = $called = $TraceCalls.Trace[]
        try
            $res = $do_body_name($(map(arg_name_splat, di[:args])...);
                                 $(passed_kwargs...))
        catch $e
            push!($prev_called,
                  $TraceCalls.Trace($fname, $trace_args, $trace_kwargs,
                                    $called, $e))
            rethrow()
        finally
            $TraceCalls.current_called[] = $prev_called
        end
        push!($prev_called,
              $TraceCalls.Trace($fname, $trace_args, $trace_kwargs,
                                $called, $TraceCalls.store($res)))
        return $res
    end
    updated_fun_di
end

function do_body_fun(di)
    fname = di[:name]
    do_body_di = copy(di)
    do_body_di[:name] = fname = isa(fname, Expr) ? gensym() : gensym(fname)
    combinedef(do_body_di), fname
end

function tracing_code_function(fdef::Expr)
    di = splitdef_memo(fdef)
    do_body_fdef, do_body = do_body_fun(di)
    quote
        @inline $do_body_fdef
        $(combinedef(tracing_fun(di, do_body)))
    end
end

function tracing_code_generated(fdef::Expr)
    di = splitdef_memo(generated2normal(fdef))
    do_body_fdef, do_body = do_body_fun(di)
    tfun_di = tracing_fun(di, do_body)
    tfun_di[:body] = Expr(:quote, tfun_di[:body])
    quote
        $(normal2generated(do_body_fdef))
        $(normal2generated(combinedef(tfun_di)))
    end
end

#"""  Takes a function definition, and returns a tracing version of it. """
tracing_code(fdef::Expr) =
    (is_generated_function_definition(fdef) ?
     tracing_code_generated(fdef) :
     tracing_code_function(fdef))

""" `@traceable function foo(...) ... end` makes that function definition traceable
with the `@trace` macro (eg. as `@trace () foo(...)`). """
macro traceable(fdef::Expr)
    @assert !is_call_definition(fdef) "@traceable cannot handle callable object definitions"

    if fdef.head == :block
        # TODO: just process it right away... It will be faster
        return esc(quote
            $([:($TraceCalls.@traceable_loose $expr) for expr in fdef.args]...)
        end)
    end

    file = (@__FILE__) == "" ? nothing : (@__FILE__)
    mod = current_module()
    di = splitdef_memo(fdef)
    signature = (module_name(mod),         # loose approximation of the real signature
                 :($(di[:name])($([splitarg(arg)[2] for arg in di[:args]]...))))

    traceable_definitions[signature] =
          RevertibleCodeUpdate(tracing_code,
                               CodeUpdate(ModDict(mod=>OrderedSet([MakeRelocatableExpr(fdef)]))))

    return esc(fdef)
end

#@memoize Dict{Tuple{Expr}, Any} strip_docstring_memo(expr0) = strip_docstring(expr0)
# There might be an issue with the memo if we decide to macroexpand the code, since the
# macroexpansion depends on the environment. It's probably a negligible issue.
@memoize Dict{Tuple{Expr, Bool}, Any} function traceable_update_handle_expr(expr0::Expr, trace_generated_functions)
    # maybe strip_docstring is not necessary anymore - September '17
    expr = strip_docstring(expr0)
    (!is_traceable(expr) || (!trace_generated_functions &&
                             is_generated_function_definition(expr))
     ? nothing : tracing_code(expr))
end
traceable_update_handle_expr(::Any, ::Any) = nothing
clear_handle_expr_memo!() =
    empty!(eval(TraceCalls, Symbol("##traceable_update_handle_expr_memoized_cache")))

""" `traceable_update(obj)::RevertibleCodeUpdate` produces the traceable code for the
given object. """
function traceable_update end

custom_when_missing(x) = warn(x)
custom_when_missing(fail::UpdateInteractiveFailure) =
    warn("Cannot trace $(fail.fn); use `@traceable` to trace methods defined interactively. See ?@traceable and the user manual for details.")
traceable_update(obj::Union{Module, String}, trace_generated_functions) =
    update_code_revertible(x->traceable_update_handle_expr(x, trace_generated_functions),
                           obj)
traceable_update(f::Union{Function, Type}, trace_generated_functions) =
    update_code_revertible(x->traceable_update_handle_expr(x, trace_generated_functions),
                           f, when_missing=custom_when_missing)

traceable_update(tup::Tuple, trace_generated_functions) =
    merge(map(x->traceable_update(x, trace_generated_functions), tup)...)
traceable_update(tup::Tuple{}, ::Bool) = EmptyRevertibleCodeUpdate()

""" The `RevertibleCodeUpdate` for the code from the `@traceable` macros. """
traceable_macro_update() = merge(EmptyRevertibleCodeUpdate(),
                                 values(traceable_definitions)...)

""" `@trace (foo, NoTraceable()) ...` will only trace `foo`; not the `@traceable`
functions. """
struct NoTraceable end
traceable_update(::NoTraceable, ::Any) = EmptyRevertibleCodeUpdate()
""" `traceable_update_code(x)` returns the code of the update. For debugging. """
traceable_update_code(x, trace_generated_functions) =
    map(to_expr, collect(reduce((s1,s2)->OrderedSet([collect(s1); collect(s2)]),
                                collect(values(traceable_update(x, trace_generated_functions).apply.md)))))
traceable_update_revert_code(x) =
    map(to_expr, collect(reduce(merge, collect(values(traceable_update(x, trace_generated_functions).revert.md)))))

function with_tracing_definitions(body::Function, obj, trace_generated_functions)
    if obj isa Tuple && NoTraceable() in obj
        upd = traceable_update(obj, trace_generated_functions)
    else
        upd = merge(traceable_update(obj, trace_generated_functions),
                    traceable_macro_update())
    end
    upd() do
        body()
    end
end

""" `get_error_or_value(f::Function)` runs `f()` in a try and returns either its outcome
or the caught exception. """
function get_error_or_value(f)
    try
        f()
    catch e
        e
    end
end

struct Root
    fun::Function
end
Base.show(io::IO, root::Root) = write(io, "Root$(root.fun)")
(r::Root)() = r.fun()

is_root(::Root) = true
is_root(::Any) = false
is_root(tr::Trace) = is_root(tr.func)

""" `recording_trace(fun::Function)` sets up a fresh Trace recording state, then executes
`fun` and returns the resulting Trace object. """
function recording_trace(fun::Function)
    current_called[] = called = Trace[]
    res = get_error_or_value(fun)
    top_trace = Trace(Root(fun), (), (), called, res)
    current_called[] = Trace[] # don't hang on to that memory unnecessarily
    top_trace
end
recording_trace(tr::Trace) = recording_trace(()->tr())  # could be handled more gracefully

function Base.trace(body::Union{Function, Trace}, to_trace=();
                    generated_functions::Bool=true)
    with_tracing_definitions(to_trace, generated_functions) do
        # To debug, just use `with_tracing_definitions()` interactively
        recording_trace(body)
    end
end

macro trace(expr)
    :($TraceCalls.trace(()->$(esc(expr))))
end

macro trace(to_trace, expr)
    :($TraceCalls.trace(()->$(esc(expr)), $(esc(to_trace))))
end