#InspectDR:
#-------------------------------------------------------------------------------
#=
TAGS:
	#WANTCONST, HIDEWARN_0.7
=#

const _devmode = false
if _devmode
@warn("Development mode - precompile deactivated.")
end
__precompile__(!_devmode)

module InspectDR

using Colors
using Graphics
using NumericIO
using Pkg
import Printf: @sprintf
import Cairo
import Cairo: CairoContext

const GtkAvailable = true
import Gtk
const _Gtk = Gtk.ShortNames
#=
REMINDER:
Conditional inclusion causes issues with __precompile__() (source: tkelman).
Not worth risking precompile issues just for current troubles with JuliaBox/Gtk.

try
	import Gtk
	eval(:(const _Gtk = Gtk.ShortNames))
	GtkAvailable = true
catch
	@warn("InspectDR: Error loading Gtk.  GUI-based features are unavailable.")
end
=#

#==Constants: Initial default values
===============================================================================#
const DTPPOINTS_PER_INCH = 72 #Typography (desktop publishing) "points per inch"

#Default font:
const DEFAULT_FONTNAME = (@static Sys.iswindows() ? "Cambria" : "Serif")
#Cairo "built-in": Serif, Sans, Serif, Fantasy, Monospace
#NOTE: "Serif" does not work well in Windows.

#==Aliases
===============================================================================#
NullOr{T} = Union{Nothing, T}

include("codegen.jl")
include("styles.jl")
include("math.jl")
include("math_graphics.jl")
include("math_coordinates.jl")
include("numericfmt.jl")
include("events.jl")
include("datasets.jl")
include("grids.jl")
include("base.jl")
include("stylesheets.jl")
include("defaults.jl")
include("graph2d.jl")
include("glyphs.jl")
include("cairo_ext.jl")
include("cairo_base.jl")
include("cairo_annotation.jl")
include("cairo_smithpolar.jl")
include("cairo_legends.jl")
include("cairo_top.jl")
include("cairo_io.jl")
include("templates.jl")
include("compat.jl")

if GtkAvailable
include("gtk_base.jl")
include("gtk_events.jl")
include("gtk_zoom.jl")
include("gtk_markers.jl")
include("gtk_traces.jl")
include("gtk_input.jl")
include("gtk_top.jl")
end

include("show.jl")

#==Comments
================================================================================
-SVG MIME output (using show) does not show up properly in Notebook.  There
 appears to be an issue with determining the extents of the image (bounding box).
-Data area referred to as "graph" area.  This might not be correct.

DReal vs Float64:
	Float64: Used in API dealing with Gtk/Cairo coordinates.
	DReal: Used in API dealing with plot data itself
		(Attempt to reduce code bloat/different code paths by using concrete
		data type in internal data structures/algorithms... despite potential
		inefficiencies).

Transformations between coordinate systems:
         nlxfrm      lxfrm
	data/ -----> axis -----> device
	world                    /view

	where:
		lxfrm is a linear transformation
		nlxfrm is potentially a nonlinear transformation

	- Data coordinates hold whichever units is associated with the data (D).
	- Axis units are transformations of D (ex: D, log(D), dB20(D), ...)
	- Device units can typically be thought of in pixels.

Auxiliary transformations:
	     nlxfrm
	axis -----> readable
                (read)
	Readable coordinates are basically axis coordinates in a readable form.
	Ex (for log10-X scale): X{axis}=2.5 --> X{read}=10^2.5

PExtents2D vs BoundingBox:
	PExtents2D: Extents of data (data coordinate system).
	BoundingBox: Extents of graphical element ("device" coordinate system).

About extents:
	Extents are typically stored using user-level coordinates (USER).
	When drawing the plot, extents are manipulated in a transformed
	coordinate system (XFRM).  For example, XFRM coordinates of a log-x
	plot are given by XFRM=log10(USER).

	ext_data::PExtents2D #Maximum extents of data
	ext_full::PExtents2D #Defines "full" zoom when combined with ext_data
		#Set fields to NaN to keep ext_data values
	ext::PExtents2D #Current/active extents

	setextents(::Plot2D, ::PExtents2D): Also invalidates rendered plot.
	getextents(...): Complementary accessor.

About bounds:
	plotbounds(...): Bounding box of entire plot.
	databounds(...): Bounding box of entire data area (multiple graph strips).
	graphbounds(...): Bounding box of individual graph strip.
=#

#==Exported interface
===============================================================================#
add = _add #Danger: high risk of collision (common name)
export add, line, glyph
export vmarker, hmarker, atext

#==Unexported interface
================================================================================
Supported plot types:
	InspectDR.Multiplot() #Figure supporting multiple plots
	InspectDR.Plot2D <: Plot #2D plot object

Convenience functions:
	write_png(path, ::Plot)
	write_svg(path, ::Plot)
	write_eps(path, ::Plot)
	write_pdf(path, ::Plot)

Other:
	refresh(::PlotWidget)
	refresh(::GtkPlot)
	clearsubplots(::GtkPlot)
=#

#Create new plot window (Is this a good idea?):
#	figure(args...; kwargs...) = GtkPlot(args...; kwargs...)

#Accessors:
if GtkAvailable
	Plot(pwidget::PlotWidget) = pwidget.src
	Plot(gplot::GtkPlot, i::Int) = gplot.subplots[i]
end

#==Already exported
================================================================================
#Displaying a plot object:
	Base.display(d::InspectDR.GtkDisplay, p::Plot) #Gtk mutltiplot window

#Writing plot to IO stream:
	Base.show(::IO, ::MIME, ::Plot2D)
for the following MIMEs:
		MIME"image/png"
		MIME"image/svg+xml"
		MIME"image/eps"
		MIME"application/pdf"
=#

function __init__()
	global defaults
	_initialize(defaults)

	if GtkAvailable
		initialize_cursors()
		_setdefaults(keybindings)
	end

	checkcompat_plots()
end

#include("precompile.jl")

end #module


using Distributed
addprocs(2)

@everywhere begin
using MLJ
using MLJBase
using Test
using Random
end

@testset "ensembles" begin
    @test include("ensembles.jl")
end

@testset "matching models to data" begin
    @test include("model_matching.jl")
end

@testset "scitypes" begin
    @test include("scitypes.jl")
end


using Test, Pkg
@testset "Land" begin
    include("test_heat_parameterizations.jl")
    include("test_water_parameterizations.jl")
    include("prescribed_twice.jl")
end


# These do most of the work in the package.
# They are all @generated or recusive functions for performance.

const UnionAllTupleOrVector = Union{Vector{UnionAll},Tuple{UnionAll,Vararg}}

"""
Sort dimensions into the order they take in the array.

Missing dimensions are replaced with `nothing`
"""
@inline Base.permutedims(tosort::AbDimTuple, perm::Union{Vector{<:Integer},Tuple{<:Integer,Vararg}}) =
    map(p -> tosort[p], Tuple(perm))
@inline Base.permutedims(tosort::Union{Vector{<:Integer},Tuple{<:Integer,Vararg}}, perm::AbDimTuple) =
    tosort

@inline Base.permutedims(tosort::AbDimTuple, order::UnionAllTupleOrVector ) =
    permutedims(tosort, Tuple(map(u -> u(), order)))
@inline Base.permutedims(tosort::UnionAllTupleOrVector, order::AbDimTuple) =
    permutedims(Tuple(map(u -> u(), tosort)), order)
@inline Base.permutedims(tosort::AbDimTuple, order::AbDimVector) =
    permutedims(tosort, Tuple(order))
@inline Base.permutedims(tosort::AbDimVector, order::AbDimTuple) =
    permutedims(Tuple(tosort), order)

@generated Base.permutedims(tosort::AbDimTuple, order::AbDimTuple) =
    permutedims_inner(tosort, order)
permutedims_inner(tosort::Type, order::Type) = begin
    indexexps = []
    for dim in order.parameters
        index = findfirst(d -> basetype(d) <: basetype(dim), tosort.parameters)
        if index == nothing
            push!(indexexps, :(nothing))
        else
            push!(indexexps, :(tosort[$index]))
        end
    end
    Expr(:tuple, indexexps...)
end

"""
Convert a tuple of AbstractDimension to indices, ranges or Colon.
"""
@inline dims2indices(a, lookup, emptyval=Colon()) =
    dims2indices(dims(a), lookup, emptyval)
@inline dims2indices(dims::Tuple, lookup, emptyval=Colon()) =
    dims2indices(dims, (lookup,), emptyval)
@inline dims2indices(dims::Tuple, lookup::Tuple, emptyval=Colon()) =
    _dims2indices(dims, permutedims(lookup, dims), emptyval)

@inline _dims2indices(dims::Tuple, lookup::Tuple, emptyval) =
    (dims2indices(dims[1], lookup[1], emptyval),
     _dims2indices(tail(dims), tail(lookup), emptyval)...)
@inline _dims2indices(dims::Tuple, lookup::Tuple{}, emptyval) =
    (emptyval, _dims2indices(tail(dims), (), emptyval)...)
@inline _dims2indices(dims::Tuple{}, lookup::Tuple{}, emptyval) = ()

@inline dims2indices(dim::AbDim, lookup::AbDim, emptyval) = val(lookup)
@inline dims2indices(dim::AbDim, lookup::Type{<:AbDim}, emptyval) = Colon()
@inline dims2indices(dim::AbDim, lookup::Nothing, emptyval) = emptyval

"""
Slice the dimensions to match the axis values of the new array

All methods returns a tuple conatining two tuples: the new dimensions,
and the reference dimensions. The ref dimensions are no longer used in
the new struct but are useful to give context to plots.

Called at the array level the returned tuple will also include the
previous reference dims attached to the array.
"""
@inline slicedims(a, dims::AbDimTuple) =
    slicedims(a, dims2indices(a, dims))
@inline slicedims(a, I::Tuple) = slicedims(dims(a), refdims(a), I)
@inline slicedims(dims::Tuple, refdims::Tuple, I::Tuple) = begin
    newdims, newrefdims = slicedims(dims, I)
    newdims, (refdims..., newrefdims...)
end
@inline slicedims(dims::Tuple{}, I::Tuple) = ((), ())
@inline slicedims(dims::AbDimTuple, I::Tuple) = begin
    d = slicedims(dims[1], I[1])
    ds = slicedims(tail(dims), tail(I))
    (d[1]..., ds[1]...), (d[2]..., ds[2]...)
end
@inline slicedims(dims::Tuple{}, I::Tuple{}) = ((), ())
@inline slicedims(d::AbDim, i::Colon) = ((rebuild(d, val(d)),), ())
@inline slicedims(d::AbDim, i::Number) = ((), (rebuild(d, val(d)[i]),))
@inline slicedims(d::AbDim{<:AbstractArray}, i::AbstractArray) =
    ((rebuild(d, val(d)[i]),), ())
@inline slicedims(d::AbDim{<:LinRange}, i::UnitRange) = begin
    range = val(d)
    start, stop, len = range[first(i)], range[last(i)], length(i) ÷ step(i)
    ((rebuild(d, LinRange(start, stop, len)),), ())
end

"""
Get the number of an AbstractDimension as ordered in the array
"""
@inline dimnum(a, lookup) = dimnum(typeof(dims(a)), lookup)
@inline dimnum(dimtypes::Type, lookup::AbstractArray) = dimnum(dimtypes, (lookup...,))
@inline dimnum(dimtypes::Type, lookup::Number) = lookup
@inline dimnum(dimtypes::Type, lookup::Tuple) =
    (dimnum(dimtypes, lookup[1]), dimnum(dimtypes, tail(lookup))...,)
@inline dimnum(dimtypes::Type, lookup::Tuple{}) = ()
@inline dimnum(dimtypes::Type, lookup::AbDim) = dimnum(dimtypes, typeof(lookup))
@generated dimnum(dimtypes::Type{DTS}, lookup::Type{D}) where {DTS,D} = begin
    index = findfirst(dt -> D <: basetype(dt), DTS.parameters)
    if index == nothing
        :(throw(ArgumentError("No $lookup in $dimtypes")))
    else
        :($index)
    end
end

"""
Format the dimension to match internal standards.

Mostily this means converting tuples and UnitRanges to LinRange,
which is easier to handle.

Errors are thrown if dims don't match the array dims or size.
"""
@inline formatdims(a::AbstractArray{T,N}, dims::Tuple) where {T,N} = begin
    dimlen = length(dims)
    dimlen == N || throw(ArgumentError("dims ($dimlen) don't match array dimensions $(N)"))
    formatdims(size(a), dims)
end
@inline formatdims(size::Tuple, dims::Tuple) =
    (formatdims(size[1], dims[1]), formatdims(tail(size), tail(dims))...,)
@inline formatdims(size::Tuple{}, dims::Tuple{}) = ()
@inline formatdims(len::Integer, dim::AbDim{<:AbstractArray}) =
    if length(val(dim)) == len
        dim
    else
        throw(ArgumentError("length of $dim $(length(val(dim))) does not match
                             size of array dimension $len"))
    end
@inline formatdims(len::Integer, dim::AbDim{<:Union{UnitRange,NTuple{2}}}) = linrange(dim, len)

linrange(dim, len) = begin
    range = val(dim)
    start, stop = first(range), last(range)
    rebuild(dim, LinRange(start, stop, len))
end

"""
Replace the specified dimensions with an index of 1 to match
a new array size where the dimension has been reduced to a length
of 1, but the number of dimensions has not changed.

Used in mean, reduce, etc.
"""
@inline reducedims(dim) = reducedims((dim,))
@inline reducedims(dims::Tuple) = map(d -> basetype(d)(1), dims)

"""
Get the dimension(s) matching the type(s) of the lookup dimension.
"""
@inline dims(a::AbstractArray, lookup) = dims(dims(a), lookup)
@inline dims(ds::AbDimTuple, lookup::Integer) = ds[lookup]
@inline dims(ds::AbDimTuple, lookup::Tuple) =
    (dims(ds, lookup[1]), dims(ds, tail(lookup))...)
@inline dims(ds::AbDimTuple, lookup::Tuple{}) = ()
@inline dims(ds::AbDimTuple, lookup) = dims(ds, basetype(lookup))
@generated dims(ds::DT, lookup::Type{L}) where {DT<:AbDimTuple,L} = begin
    index = findfirst(dt -> dt <: L, DT.parameters)
    if index == nothing
        :(throw(ArgumentError("No $lookup in $dims")))
    else
        :(ds[$index])
    end
end


function maxpool1d(ip, kernel, strides=1, pads=0)
    l = size(ip)[2]
    ip1 = ip[:, 1]
    temp = []
    for x=1:length(ip1)-kernel+1
        push!(temp, findmax(ip1[x:x+kernel])[1])
    end
    return temp
end


module FullyThreadedTree

    using Base.Cartesian

    abstract type AbstractTree{N} end
    abstract type AbstractFace{N} end

    include("face.jl")

    export Face,
           cells,
           at_boundary,
           at_refinement,
           regular,
           active,
           level,
           collect_faces

    include("tree.jl")

    export Tree,
           initialize_tree,
           refine!,
           coarsen!,
           active,
           parent_of_active,
           initialized,
           level,
           faces,
           siblings,
           root,
           collect_cells

    include("interface.jl")

    export iterate,
           show,
           length,
           cells,
           all_faces

    include("tools.jl")

    export volume,
           centroid,
           integrate,
           cell_distance,
           cell_volume,
           area,
           volume,
           levels

    include("plotting.jl")

    export plot,
           plot!

end # module


module SegmentTrees

export
    SegmentTree,
    SegmentMap,
    AbstractSegment,
    Segment,
    SegmentValue,
    depth,
    hasintersection,
    from,
    value

using Base: notnothing

include("slice.jl")
include("creation.jl")
include("iteration.jl")
include("insertion.jl")
include("deletion.jl")
include("merging.jl")
include("search.jl")
include("intersection.jl")
include("map.jl")

# Diagnostics
# -----------

# Dumb tree printing, useful only for debuging.
#function showtree(io::IO, t::SegmentBTree)
    #showtree(io, t.root, 0)
#end

#function showtree(t::SegmentBTree)
    #showtree(STDOUT, t)
#end

#function showtree(io::IO, t::InternalNode, indent::Int)
    #for (i, child) in enumerate(t.children)
        #showtree(io, child, indent+1)
        #if i <= length(t.keys)
            #print(io, repeat("  ", indent), t.keys[i], ":\n")
        #end
    #end
#end

#function showtree(io::IO, t::LeafNode, indent::Int)
    #for (k, v) in zip(t.keys, t.values)
        #print(io, repeat("  ", indent), k, ": ", v, "\n")
    #end
#end

end  # module SegmentTrees


module Compress

import Pkg.TOML
import Pkg.Types: VersionSpec, compress_versions

function load_versions(path::String)
    versions_file = joinpath(dirname(path), "Versions.toml")
    versions_dict = TOML.parsefile(versions_file)
    sort!([VersionNumber(v) for v in keys(versions_dict)])
end

function load(path::String,
    versions::Vector{VersionNumber} = load_versions(path))
    compressed = TOML.parsefile(path)
    uncompressed = Dict{VersionNumber,Dict{Any,Any}}()
    for (vers, data) in compressed
        vs = VersionSpec(vers)
        for v in versions
            v in vs || continue
            merge!(get!(uncompressed, v, Dict()), deepcopy(data))
        end
    end
    return uncompressed
end

function compress(path::String, uncompressed::Dict,
    versions::Vector{VersionNumber} = load_versions(path))
    inverted = Dict()
    for (ver, data) in uncompressed, (key, val) in data
        val isa TOML.TYPE || (val = string(val))
        push!(get!(inverted, key => val, VersionNumber[]), ver)
    end
    compressed = Dict()
    for ((k, v), vers) in inverted
        for r in compress_versions(versions, sort!(vers)).ranges
            get!(compressed, string(r), Dict{String,Any}())[k] = v
        end
    end
    return compressed
end

function save(path::String, uncompressed::Dict,
    versions::Vector{VersionNumber} = load_versions(path))
    compressed = compress(path, uncompressed)
    open(path, write=true) do io
        TOML.print(io, compressed, sorted=true)
    end
end

end # module


"""
    UncertainValueDataset

A dataset of uncertain values.

## Fields

- **`values::AbstractVector{<:AbstractUncertainValue}`**: The uncertain values. Each value is
    represented by an `AbstractUncertainValue`.
"""
struct UncertainValueDataset <: AbstractUncertainValueDataset
    values::AbstractVector{<:AbstractUncertainValue}
end

"""
    ConstrainedUncertainValueDataset

Generic constrained dataset containing uncertain values.

## Fields
- **`values::AbstractVector{<:AbstractUncertainValue}`**: The uncertain values.
"""
struct ConstrainedUncertainValueDataset <: AbstractUncertainValueDataset
    values::AbstractVector{<:AbstractUncertainValue}
end

function UncertainValueDataset(x::AbstractArray{T, 1}) where T
    UncertainValueDataset(CertainValue.(x))
end

export
UncertainValueDataset,
ConstrainedUncertainValueDataset


"""
    set_geometry(cascade, ::Type{T}; kwargs...) where T<:Geometry
Given a cascade and geometry, return a cascade with an updated type.
"""
function set_geometry(cascade, ::Type{Violin}, args...; kwargs...)
    return _set_geometry(cascade, Violin, Poly, Coloring, args...; style=:fill, kwargs...)
end

function set_geometry(cascade, ::Type{Vertical}, args...; kwargs...)
    return _set_geometry(cascade, Vertical, Box, Blending, args...;
        style=:fill,
        subdivide=true,
        space=true,
        kwargs...,
    )
end

function set_geometry(cascade, ::Type{Horizontal}, args...; usegradient=missing, kwargs...)
    # !!!! want to show gradient if there's only one cascade element. Is this working??
    usegradient = coalesce(usegradient, length(cascade)==1)

    return _set_geometry(cascade, Horizontal, Box, usegradient ? Blending : Coloring, args...;
        alpha=length(cascade),
        style=:fill,
        subdivide=false,
        space=true,
        kwargs...,
    )
end

function set_geometry(cascade, ::Type{Parallel}, args...; slope::Bool=true, kwargs...)
    return _set_geometry(cascade, Parallel, Line, Coloring, args...;
        quantile = convert(Float64, slope),
        alpha = length(cascade),
        factor = 0.5,
        style = :stroke,
        subdivide = false,
        space = false,
        kwargs...,
    )
end

"""
    _set_geometry(cascade::Cascade{Data}, Geometry, Shape, Color; kwargs...)
"""
function _set_geometry(cascade::Cascade{Data}, Geometry, Shape, Color, args...;
    colorcycle::Bool = false,
    kwargs...,
)
    pos = scale_for(cascade, Geometry, args...; kwargs...)
    sgn = sign(cascade)

    label = get_label.(collect_data(cascade))
    shape = vectorize.(_define_from(Shape, Color, pos, sgn; kwargs...))

    data = Geometry.(label, shape, length(cascade))

    return Cascade(
        start = set_hue!(first(data), "black"),
        stop = set_hue!(last(data), "black"),
        steps = colorcycle ? set_hue!.(data[2:end-1], cascade.permutation) : data[2:end-1],
        permutation = cascade.permutation,
        correlation = cascade.correlation,
        ispermuted = cascade.ispermuted,
        iscorrelated = cascade.iscorrelated,
    )
end

"""
    set_hue!(x, h)
"""
function set_hue!(x::Coloring, h; kwargs...)
    x.hue = define_from(Luxor.RGB, h; kwargs...)
    return x
end

set_hue!(x::Handle, h; kwargs...) = begin set_hue!(x.shape, h; kwargs...); return x end
set_hue!(x::Blending, h) = begin x.hue = _define_gradient(h); return x end
set_hue!(x::T, h) where T<:Geometry = begin set_hue!(x.shape, h); return x end
set_hue!(x::T, h) where T<:Shape = begin set_hue!(x.color, h); return x end
set_hue!(x::Vector{T}, h) where T<:Shape = begin set_hue!.(x, h); return x end

"""
    set_position!(handle::Handle)
This method calculates the position of a Handle's shape and label such that the shape is
centered at (x0,y0) and fits within the bounds of `(2*dx, 2*dy)`. The label is aligned
middle/left at `(x0+dx+space, y0)`
"""
function set_position!(handle::Handle;
    x0 = WIDTH-100,
    y0 = 3*SEP,
    dx = 2*SEP,
    dy = SEP,
    space = SEP,
    idx = 1,
    kwargs...,
)
    y0 = y0 + (2*dy+space)*(idx-1)

    setproperty!(handle.shape, :position, (
        Luxor.Point(x0-dx, y0+dy),
        Luxor.Point(x0+dx, y0-dy),
    ))

    setproperty!(handle.label, :position, Luxor.Point(x0+dx+space, y0))

    return handle
end

"""
    get_shape(cascade; kwargs...)
This function returns the shape used in `cascade`, with its position updated, as defined by
keyword arguments.
"""
get_shape(cascade::Cascade{T}) where T<:Geometry = copy(first(cascade.start.shape))
get_shape(cascade::Cascade{Violin}) = copy(cascade.start.shape)


"""
smoother of solution

ref:
- Dechter, R., Meiri, I., and Pearl, J. (1991).
  Temporal constraint networks.
  Artificial Intelligence (AIJ).
- Honing, W. et al. (2016).
  Multi-agent path finding with kinematic constraints.
  In ICAPS.
"""

mutable struct Action{State<:AbsState}
    id::String
    from::Node{State}
    to::Node{State}
    t::Int64
    agent::Int64
    predecessors::Vector{Tuple{Int64,String}}
    successors::Vector{Tuple{Int64,String}}
end

function get_action_id(
    from::Node{State},
    to::Node{State},
    t::Int64,
)::String where {State<:AbsState}
    return string(from.id) * "_" * string(to.id) * "_" * string(t)
end

"""create temporal plan graph temporal plan graph"""
function get_temporal_plan_graph(
    solution::Vector{Vector{Node{State}}},
    collide::Function,
    connect::Function;
    skip_connection::Bool = true,
)::Vector{Vector{Action{State}}} where {State<:AbsState}

    N = length(solution[1])
    T = length(solution)

    # temporal plan graph
    TPG = [Vector{Action{State}}() for i = 1:N]

    # type-1 dependency
    for i = 1:N
        v_last = solution[1][i]
        for (t, Q) in enumerate(solution[1:end])
            v_current = Q[i]

            if v_current.id != v_last.id
                action = Action(
                    get_action_id(v_last, v_current, t),
                    v_last,
                    v_current,
                    t,
                    i,
                    Vector{Tuple{Int64,String}}(),
                    Vector{Tuple{Int64,String}}(),
                )

                # add type-1 dependency
                if length(TPG[i]) > 0
                    action_pre = TPG[i][end]
                    push!(action.predecessors, (i, action_pre.id))
                end

                # insert
                push!(TPG[i], action)

                # update
                v_last = v_current
            end
        end
    end

    # type-2 dependency
    for i = 1:N
        for action_self in TPG[i]
            for j = 1:N
                if j == i
                    continue
                end

                # exclude ealier actions
                for action_other in filter(a -> a.t > action_self.t, TPG[j])

                    # create type-2 dependency
                    if collide(
                        action_self.from.q,
                        action_self.to.q,
                        action_other.from.q,
                        action_other.to.q,
                        i,
                        j,
                    )
                        push!(
                            action_other.predecessors,
                            (action_self.agent, action_self.id),
                        )
                        break
                    end
                end
            end
        end
    end

    # remove redundant dependencies
    for i = 1:N
        for action in TPG[i]
            latest_actions = Dict{Int64,Action}()
            for (j, id) in action.predecessors
                action_other = TPG[j][findfirst(a -> a.id == id, TPG[j])]
                if !haskey(latest_actions, j) || latest_actions[j].t < action_other.t
                    latest_actions[j] = action_other
                end
            end
            action.predecessors = map(a -> (a.agent, a.id), values(latest_actions))

            # update successors
            for (j, id) in action.predecessors
                action_other = TPG[j][findfirst(a -> a.id == id, TPG[j])]
                push!(action_other.successors, (i, action.id))
            end
        end
    end

    skip_connection && try_skip_connection!(TPG, collide, connect)
    return TPG
end

"""remove redundant actions in temporal plan graph"""
function try_skip_connection!(
    TPG::Vector{Vector{Action{State}}},
    collide::Function,
    connect::Function,
) where {State<:AbsState}

    N = length(TPG)
    get_action = (j, id) -> TPG[j][findfirst(a -> a.id == id, TPG[j])]

    for i = 1:N
        k = 1
        while k < length(TPG[i])
            k += 1
            a1 = TPG[i][k-1]
            a2 = TPG[i][k]

            # check connection
            !connect(a1.from.q, a2.to.q, i) && continue

            # check dependencies
            any([j != i for (j, _) in a2.predecessors]) && continue
            any([j != i for (j, _) in a1.successors]) && continue

            # new action candidate
            a3 = Action(
                get_action_id(a1.from, a2.to, a1.t),
                a1.from,
                a2.to,
                a1.t,
                a1.agent,
                a1.predecessors,
                a2.successors,
            )

            # get causalities
            causal_actions = get_ancestors(a1, TPG)
            union!(causal_actions, get_descendants(a2, TPG))
            causal_actions = filter(val -> val[1] != i, causal_actions)

            # check collisions
            conflicted = false
            for j = 1:N
                if j == i
                    continue
                end
                for a4 in filter(a -> !((j, a.id) in causal_actions), TPG[j])
                    if collide(a3.from.q, a3.to.q, a4.from.q, a4.to.q, i, j)
                        conflicted = true
                        break
                    end
                end
                if conflicted
                    continue
                end
                # check last location
                if collide(a3.from.q, a3.to.q, TPG[j][end].to.q, TPG[j][end].to.q, i, j)
                    conflicted = true
                    break
                end
            end
            if conflicted
                continue
            end

            # update action orders
            deleteat!(TPG[i], k-1:k)
            insert!(TPG[i], k - 1, a3)

            # update dependencies
            for (j, id) in a3.predecessors
                a4 = get_action(j, id)
                m = findfirst(tpl -> tpl[1] == i && tpl[2] == a1.id, a4.successors)
                deleteat!(a4.successors, m)
                push!(a4.successors, (i, a3.id))
            end
            for (j, id) in a3.successors
                a4 = get_action(j, id)
                m = findfirst(tpl -> tpl[1] == i && tpl[2] == a2.id, a4.predecessors)
                deleteat!(a4.predecessors, m)
                push!(a4.predecessors, (i, a3.id))
            end
            k -= 1
        end
    end
end

"""obtain actions with causal dependencies"""
function get_causal_actions(
    action::Action{State},
    TPG::Vector{Vector{Action{State}}},
    for_ancestors::Bool,
) where {State<:AbsState}

    N = length(TPG)
    tables = [Dict{String,Set{Tuple{Int64,String}}}() for i = 1:N]
    function f(i, id)
        action = TPG[i][findfirst(a -> a.id == id, TPG[i])]
        if haskey(tables[i], id)
            return tables[i][id]
        end
        causal_actions = Set{Tuple{Int64,String}}()
        for (j, id_j) in ((for_ancestors) ? action.predecessors : action.successors)
            push!(causal_actions, (j, id_j))
            foreach(val -> push!(causal_actions, val), f(j, id_j))
        end
        tables[i][id] = causal_actions
        return causal_actions
    end
    return f(action.agent, action.id)
end

function get_ancestors(
    action::Action{State},
    TPG::Vector{Vector{Action{State}}},
) where {State<:AbsState}
    return get_causal_actions(action, TPG, true)
end

function get_descendants(
    action::Action{State},
    TPG::Vector{Vector{Action{State}}},
) where {State<:AbsState}
    return get_causal_actions(action, TPG, false)
end

"""sampling one solution from temporal plan graph"""
function get_greedy_solution(
    TPG::Vector{Vector{Action{State}}},
    config_goal::Vector{State},
)::Vector{Vector{Node{State}}} where {State<:AbsState}

    N = length(TPG)
    indexes = [1 for i = 1:N]  # internal clock of agents
    solution = [[
        (length(arr) > 0 ? arr[1].from : Node(config_goal[i], 0, Vector{Int64}())) for
        (i, arr) in enumerate(TPG)
    ]]

    while !all([k > length(TPG[i]) for (i, k) in enumerate(indexes)])
        config = []
        indexes_last = copy(indexes)
        for i = 1:N
            if length(TPG[i]) > 0
                action = TPG[i][min(indexes[i], length(TPG[i]))]
                if all([
                    findfirst(a -> a.id == id, TPG[j]) < indexes_last[j] for
                    (j, id) in action.predecessors
                ])
                    # increment
                    indexes[i] = min(indexes[i] + 1, length(TPG[i]) + 1)
                    push!(config, action.to)
                else
                    # stay
                    push!(config, action.from)
                end
            else
                push!(config, Node(config_goal[i], 0, Vector{Int64}()))
            end
        end
        push!(solution, config)
    end
    return solution
end

"""
    get_solution_cost(
        solution::Union{Nothing,Vector{Vector{Node{State}}}},
    )::Union{Nothing,Dict{Symbol,Float64}} where {State<:AbsState}

compute costs (sum_of_costs & makespan) of solution
"""
function get_solution_cost(
    solution::Union{Nothing,Vector{Vector{Node{State}}}},
)::Union{Nothing,Dict{Symbol,Float64}} where {State<:AbsState}

    isnothing(solution) && return nothing
    N = length(solution[1])
    T = length(solution)

    # compute last timesteps
    last_timesteps = fill(T, N)
    for i = 1:N
        for t in reverse(collect(1:T-1))
            solution[t][i] != solution[t+1][i] && break
            last_timesteps[i] = t
        end
    end

    makespan = 0
    sum_of_cost = 0
    for t = 2:T
        c = maximum(i -> dist(solution[t-1][i], solution[t][i]), 1:N)
        sum_of_cost += c * length(filter(e -> t <= e, last_timesteps))
        makespan += c
    end

    return Dict(:sum_of_cost => sum_of_cost, :makespan => makespan)
end

"""
    get_tpg_cost(
        TPG::Union{Nothing,Vector{Vector{Action{State}}}},
    )::Union{Nothing,Dict{Symbol,Float64}} where {State<:AbsState}

compute costs (sum_of_costs & makespan) of temporal plan graph (by dynamic programming)
"""
function get_tpg_cost(
    TPG::Union{Nothing,Vector{Vector{Action{State}}}},
)::Union{Nothing,Dict{Symbol,Float64}} where {State<:AbsState}

    isnothing(TPG) && return nothing

    N = length(TPG)
    cost_tables = [Dict{String,Float64}() for i = 1:N]

    function f(i, id)
        action = TPG[i][findfirst(a -> a.id == id, TPG[i])]
        c = dist(action.from.q, action.to.q)
        if !isempty(action.predecessors)
            c += maximum([
                (haskey(cost_tables[j], id_j) ? cost_tables[j][id_j] : f(j, id_j)) for
                (j, id_j) in action.predecessors
            ])
        end
        cost_tables[i][id] = c
        return c
    end

    arr = [(length(TPG[i]) > 0 ? f(i, TPG[i][end].id) : 0) for i = 1:N]
    return Dict(:sum_of_cost => sum(arr), :makespan => maximum(arr))
end

"""
    smoothing(
        solution::Vector{Vector{Node{State}}},
        connect::Function,
        collide::Function;
        VERBOSE::Int64 = 0,
    )::Tuple{
        Vector{Vector{Action{State}}},  # temporal plan graph
        Vector{Vector{Node{State}}},    # solution
        Dict{Symbol,Float64},           # final cost
    } where {State<:AbsState}

smoothing by temporal plan graph
"""
function smoothing(
    solution::Vector{Vector{Node{State}}},
    connect::Function,
    collide::Function;
    VERBOSE::Int64 = 0,
)::Tuple{
    Vector{Vector{Action{State}}},  # temporal plan graph
    Vector{Vector{Node{State}}},  # solution
    Dict{Symbol,Float64},  # final cost
} where {State<:AbsState}

    isnothing(solution) && return nothing

    solution_tmp = solution
    config_goal = map(v -> v.q, solution[end])
    cost_last = nothing
    sum_of_cost_last = Inf
    while true
        # 1. create temporal plan graph
        TPG = get_temporal_plan_graph(solution_tmp, collide, connect)
        # 2. sampling from temporal plan graph
        solution_tmp = get_greedy_solution(TPG, config_goal)
        cost = get_tpg_cost(TPG)
        if sum_of_cost_last >= cost[:sum_of_cost]
            return (TPG, solution_tmp, cost)
        else
            # 3. update solution
            VERBOSE > 0 && @info @sprintf("cost is updated: %f -> %f", cost_last, cost)
            cost_last = cost
            sum_of_cost_last = get(cost, :sum_of_cost)
        end
    end
end

function smoothing(solution::Nothing, args...; kwargs...)::Nothing
    nothing
end


function insert_zero(x, loc)
    return x + x & ~(1 << (loc - 1) - 1)
end


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Functions to render the numbers in text backend.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

function _render_number_text(
    number::Rational;
    compact::Bool = true
)
    if compact
        aux = abs(number.num)
        num_str = ""

        while aux ≥ 1
            i = aux % 10
            num_str = _EXPONENTS[i + 1] * num_str
            aux = floor(aux / 10)
        end

        aux = abs(number.den)
        den_str = ""

        while aux ≥ 1
            i = aux % 10
            den_str = _SUBNUMBERS[i + 1] * den_str
            aux = floor(aux / 10)
        end

        number_str = (number.num ≥ 0 ? "" : "- ") * num_str * "/" * den_str

        return number_str
    else
        num_str = sprint(print, abs(number.num))
        den_str = sprint(print, abs(number.den))

        num_width = textwidth(num_str)
        den_width = textwidth(den_str)

        fraction_width = max(num_width, den_width)

        # Align the numbers in the center.
        Δ = div(fraction_width - num_width, 2)
        num_str = " "^Δ * num_str

        Δ = div(fraction_width - den_width, 2)
        den_str = " "^Δ * den_str

        # Check if the fraction is negative. In this case, we must change the
        # alignemnt.
        is_neg = number.num < 0

        num_str  = (!is_neg ? "" : "  ") * num_str * "\n"
        num_str *= (!is_neg ? "" : "- ") * "—"^fraction_width * "\n"
        num_str *= (!is_neg ? "" : "  ") * den_str

        return num_str
    end
end

function _render_number_text(number::Irrational; kwargs...)
    return string(number)
end

function _render_number_text(
    number::Number;
    always_print_base::Bool = false,
    multiplication_sign::Char = '×',
    significand_format::String = "%g",
    show_base::Bool = true,
    show_significand::Bool = true,
    new_decimal_base::Union{Nothing, Number} = nothing
)
    # Get the significand and the base.
    if new_decimal_base !== nothing
        significand, base = _get_significand_and_base(number, new_decimal_base)
    else
        significand, base = _get_significand_and_base(number)
    end

    # Significand
    # ==========================================================================

    significand_str = ""

    if show_significand
        fmt = Printf.Format(significand_format)
        significand_str = Printf.format(fmt, significand)
    end

    # Base
    # ==========================================================================

    # If `base` is 0, then only show it if the user wants.
    base_str = ""

    if show_base && ((base != 0) || always_print_base)
        # Create the string representation.
        aux = abs(base)
        exponent_str = ""

        while aux ≥ 1
            i = aux % 10
            exponent_str = _EXPONENTS[i + 1] * exponent_str
            aux = floor(aux / 10)
        end

        # If `exponent_str` is empty, then the base is 0.
        if always_print_base
            exponent_str = "⁰"
        end

        if base < 0
            exponent_str = "⁻" * exponent_str
        end

        if show_significand
            base_str *= " " * multiplication_sign * " "
        end

        base_str *= "10" * exponent_str
    end

    # Output
    # ==========================================================================

    number_str = significand_str * base_str

    return number_str
end


# reso.jl

#include("misc.jl")
#include("primitives.jl")
#include("types.jl")
#include("unify.jl")

# input form
# "[x,y].[+P(x,f(x),-Q(x)]"

# internal form
# CForm = ([:x,:y], [:(+(P(x,f(x)))), :(-(Q(x)))])
# CForm2 = (:C12, [:x,:y], [:(+(P(x,f(x)))), :(-(Q(x)))])

## result
# :NAP is Not APplicable

Clause = Array
Literal = Expr

EmptyClause = []

CForm = Tuple{Array,Array}
EmptyCForm = ([],EmptyClause)

## resolution of 2 clauses
function genvars(vars::Vlist)::Vlist
 map(v->gensym(v),vars)
end

function rename(vars, C::Clause, nvar)::Clause
 if isempty(vars); return C end
 apply(vars, C, nvar)
end

"""
resolution clash the i1'th of c1 with i2'th of c2
"""
function resolution(vars::Vlist, c1::Clause, c2::Clause, i1::Int, i2::Int)
# clause c1,c2 is an array.

 if length(c1) < i1; println("i1 over the c1's length");return :NAP end
 if length(c2) < i2; println("i2 over the c2's length");return :NAP end
 if i1 <= 0; println("i1 should be >=0");return :NAP end
 if i2 <= 0; println("i2 should be >=0");return :NAP end

 lit1 = c1[i1]
 lit2 = c2[i2]

 if lit1.args[1] == lit2.args[1]; return :NAP end
 lit1.args[1] = lit2.args[1]  # temporarily let signs same

 try
  sigma = unify(vars, lit1, lit2)
  c1=vcat(c1[1:(i1-1)],c1[(i1+1):end],c2[1:(i2-1)],c2[(i2+1):end])
  r1,s1 = apply(vars,c1,sigma), sigma
  return r1,s1
 catch e
  if isa(e, Loop)
   return :FAIL
  end
 end
 return :NAP
end

function resolution(var1::Vlist, c1::Clause, i1::Int, var2::Vlist, c2::Clause, i2::Int)
 vs1 = genvars(var1)
 vs2 = genvars(var2)
 var = vcat(vs1,vs2)
 rc1=rename(var1, c1, vs1)
 rc2=rename(var2, c2, vs2)
 rv = resolution(var, rc1, rc2, i1, i2)
 if typeof(rv) == Symbol;return rv end
 r1,sigma = rv
 return (var, r1), sigma
end

function resolution(c1::CForm, i1::Int, c2::CForm,i2::Int)
 cc1=deepcopy(c1)
 cc2=deepcopy(c2)
 rv =resolution(cc1[1],cc1[2],i1, cc2[1],cc2[2],i2)
 if typeof(rv) == Symbol; return rv end
 cf,sigma=rv
 return renamereadable(cf,sigma)
end

"""
reduction
"""
function reduction(vars::Vlist, c1::Clause, i1::Int)
 if length(c1) < i1; println("i1 over the c1's length");return :NAP end
 if i1 <= 0; println("i1 should be >=0");return :NAP end
 lit1 = c1[i1]
 atm1 = lit1.args[2]

 for i2 in 1:length(c1)
  if i2==i1; continue end
   lit2 = c1[i2]
   atm2 = lit2.args[2]
   if atm1.args[1] != atm2.args[1]; continue end
   if lit1.args[1] != lit2.args[1]; continue end

   try
    sigma = unify(vars, lit1, lit2)
    c1=vcat(c1[1:(i1-1)],c1[(i1+1):end])
    r1=apply(vars,c1,sigma)
    return renamereadable((vars, r1), sigma)
   catch e
    continue
   end
 end
 return ((vars,c1),[])
end

function reduction(c1::CForm, i1::Int)
 reduction(c1[1], c1[2], i1)
end

"""
satisfiable check
"""
function satisfiable(vars::Vlist, c1::Clause)
 for i1 in 1:length(c1)
  lit1 = c1[i1]
  atm1 = lit1.args[2]

  for i2 in i1+1:length(c1)
   lit2 = c1[i2]
   atm2 = lit2.args[2]
   if atm1.args[1] != atm2.args[1]; continue end
   if lit1.args[1] == lit2.args[1]; continue end
   lit2.args[1] = lit1.args[1] # temporaly let it same

   try
    sigma = unify(vars, lit1, lit2)
    return true
   catch e
    continue
   end
  end
 end
 return false

end

function satisfiable(c1::CForm)
 satisfiable(c1[1], c1[2])
end

### variable lists

numberingvar(vs::Array{Symbol},i::Int) = map(v->Symbol(v,i),vs)

function readablevars(n::Int)
 rv = [:u,:v,:w,:x,:y,:z]
 vars=[]
 for i in 1:ceil(Int, n/length(rv))
   append!(vars, numberingvar(rv, i))
 end
 vars[1:n]::Vlist
end

function containvar(vars::Vlist, form::Expr)
  cvars=[]::Vlist
  for arg in form.args
   if typeof(arg) == Symbol
    if arg in vars
      push!(cvars, arg)
    end
   else
    pvars =  containvar(vars, arg)
    append!(cvars, pvars)
   end
  end
  return union(cvars,cvars)
end

function containvar(vars::Vlist, cls::Clause)
  cvars=[]::Vlist
  for lit in cls
   pvars =  containvar(vars, lit)
   append!(cvars, pvars)
  end
  return union(cvars,cvars)
end

function fitvars(vars::Vlist, cls::Clause, sub::Tlist)
 fvars = containvar(vars, cls)
 fsub = reducesub(vars, fvars, sub)
 return((fvars, cls), fsub)
end

function renamereadable(clause::CForm, sub::Tlist)
 vars = clause[1]
 cls  = clause[2]
 rvars = readablevars(length(vars))
 rcls = apply(vars, cls, rvars)
 rsub = apply(vars, sub, rvars)
 fcls, fsub = fitvars(rvars, rcls, rsub)
 return fcls,fsub
end

function equalclause(clause1::CForm, clause2::CForm)
 rcls1,rsub1=renamereadable(clause1, [])
 rcls2,rsub2=renamereadable(clause2, [])

 var1=rcls1[1]
 cls1=rcls1[2]
 var2=rcls2[1]
 cls2=rcls2[2]

 if length(var1)!=length(var2);return false end
 if length(cls1)!=length(cls2);return false end

 scls2 = apply(var2, cls2, var1)
 if cls1 == scls2; return true end
 return false
end

function reducesub(vars::Vlist, fvars::Vlist, sub::Tlist)
 fsub  = []
 if isempty(sub); return [] end
 for i in 1:length(vars)
  if vars[i] in fvars
    push!(fsub, sub[i])
  end
 end
 return fsub
end


module DiscreteWrapper

using JavaCall

const AbstractIntegerDistribution = @jimport org.hipparchus.distribution.discrete.AbstractIntegerDistribution
const BinomialDistribution = @jimport org.hipparchus.distribution.discrete.BinomialDistribution
const Class = @jimport java.lang.Class
const EnumeratedIntegerDistribution = @jimport org.hipparchus.distribution.discrete.EnumeratedIntegerDistribution
const GeometricDistribution = @jimport org.hipparchus.distribution.discrete.GeometricDistribution
const HypergeometricDistribution = @jimport org.hipparchus.distribution.discrete.HypergeometricDistribution
const IntegerDistribution = @jimport org.hipparchus.distribution.IntegerDistribution
const JString = @jimport java.lang.String
const List = @jimport java.util.List
const Object = @jimport java.lang.Object
const PascalDistribution = @jimport org.hipparchus.distribution.discrete.PascalDistribution
const PoissonDistribution = @jimport org.hipparchus.distribution.discrete.PoissonDistribution
const UniformIntegerDistribution = @jimport org.hipparchus.distribution.discrete.UniformIntegerDistribution
const ZipfDistribution = @jimport org.hipparchus.distribution.discrete.ZipfDistribution

export AbstractIntegerDistribution
export BinomialDistribution
export EnumeratedIntegerDistribution
export GeometricDistribution
export HypergeometricDistribution
export PascalDistribution
export PoissonDistribution
export SaddlePointExpansion
export UniformIntegerDistribution
export ZipfDistribution
export cumulative_probability
export equals
export get_class
export get_exponent
export get_mean
export get_number_of_elements
export get_number_of_successes
export get_number_of_trials
export get_numerical_mean
export get_numerical_variance
export get_pmf
export get_population_size
export get_probability_of_success
export get_sample_size
export get_support_lower_bound
export get_support_upper_bound
export hash_code
export inverse_cumulative_probability
export is_support_connected
export log_probability
export normal_approximate_probability
export notify
export notify_all
export probability
export to_string
export upper_cumulative_probability
export wait

include("abstract_integer_distribution.jl")
include("binomial_distribution.jl")
include("enumerated_integer_distribution.jl")
include("geometric_distribution.jl")
include("hypergeometric_distribution.jl")
include("pascal_distribution.jl")
include("poisson_distribution.jl")
include("saddle_point_expansion.jl")
include("uniform_integer_distribution.jl")
include("zipf_distribution.jl")

end


# This file was generated, do not modify it. # hide
using StatsBase #hide

@show describe(D.localrank)


export paraGen, surveyGen, padding, velocity_to_moduli
"""
	paraGen(nz::Int64, nx::Int64, dz::Real, dx::Real, nSteps::Int64, dt::Real,
	f0::Real, nPml::Int64, nPad::Int64, para_fname::String, survey_fname::String, data_dir_name::String;
	if_win::Bool=false, filter_para=nothing, if_src_update::Bool=false, scratch_dir_name::String="")

Generates a parameter file consumed by [`fwi_op`](@ref) and [`fwi_op_ops`](@ref)

- `nSteps` : Number of time steps in the simulation
- `dt` : Time step size
- `f0` : Source reference frequency for CPML (usually chosen a th dominant frequency)
- `nPoints_pml` : Number of points in CPML boundary condition
- `nPad` : Padding width in the z direction, in order to make the GPU memory coalesced
- `if_win` : Whether to apply window function to data gathers
- `filter` : Whether apply band-pass filters to data
- `if_src_update` : Whether update source signatures
- `survey_fname` : The name of the survey file
- `data_dir_name` : Locations for storing observation data (generated by [`fwi_op_ops`](@ref), in the form of `.bin`)
- `scratch_dir_name` : Temporary data location. If present, will output intermediate files for QC.
"""
function paraGen(nz::Int64, nx::Int64, dz::Real, dx::Real, nSteps::Int64, dt::Real,
		f0::Real, nPml::Int64, nPad::Int64, para_fname::String, survey_fname::String, data_dir_name::String;
		if_win::Bool=false, filter_para=nothing, if_src_update::Bool=false, scratch_dir_name::String="")
	para = OrderedDict()
	para["nz"] = nz
	para["nx"] = nx
	para["dz"] = dz
	para["dx"] = dx
	para["nSteps"] = nSteps
	para["dt"] = dt
	para["f0"] = f0
	para["nPoints_pml"] = nPml
	para["nPad"] = nPad
	if if_win != false
		para["if_win"] = true
	end
	if filter_para != nothing
		para["filter"] = filter_para
	end
	if if_src_update != false
		para["if_src_update"] = true
	end
	para["survey_fname"] = survey_fname
	para["data_dir_name"] = data_dir_name
	if !isdir(data_dir_name)
		mkdir(data_dir_name)
	end
	# if nStepsWrap != nothing
	#   para["nStepsWrap"] = nStepsWrap
	# end

	if(scratch_dir_name != "")
			para["scratch_dir_name"] = scratch_dir_name
			if !isdir(scratch_dir_name)
				mkdir(scratch_dir_name)
			end
	end
	para_string = JSON.json(para)

	open(para_fname,"w") do f
		write(f, para_string)
	end
end

# all shots share the same number of receivers
@doc raw"""
	surveyGen(z_src::Array{T}, x_src::Array{T},
	z_rec::Array{T}, x_rec::Array{T}, survey_fname::String; Windows=nothing, Weights=nothing) where T<:Integer

Generates the survey parameter file.

- `z_src` : $z$ coordinates of sources
- `x_src` : $x$ coordinates of sources
- `z_rec` : $z$ coordinates of receivers
- `x_rec` : $x$ coordinates of receiverss
- `survey_fname` : The name of the survey file
- `Windows` :
- `Weights` :
"""
function surveyGen(z_src::Array{T}, x_src::Array{T},
	z_rec::Array{T}, x_rec::Array{T}, survey_fname::String; Windows=nothing, Weights=nothing) where T<:Integer
	nsrc = length(x_src)
	nrec = length(x_rec)
	survey = OrderedDict()
	survey["nShots"] = nsrc
	for i = 1:nsrc
		shot = OrderedDict()
		shot["z_src"] = z_src[i]
		shot["x_src"] = x_src[i]
		shot["nrec"] = nrec
		shot["z_rec"] = z_rec
		shot["x_rec"] = x_rec
		if Windows != nothing
			shot["win_start"] = Windows["shot$(i-1)"][:start]
			shot["win_end"] = Windows["shot$(i-1)"][:end]
		end
		if Weights != nothing
			# shot["weights"] = Int64.(Weights["shot$(i-1)"][:weights])
			shot["weights"] = Weights["shot$(i-1)"][:weights]
		end
		survey["shot$(i-1)"] = shot
	end

	survey_string = JSON.json(survey)
	open(survey_fname,"w") do f
		write(f, survey_string)
	end

end

"""
		sourceGene(f::Float64, nStep::Integer, delta_t::Float64)

Generates Ricker wavelets.
"""
function sourceGene(f::Float64, nStep::Integer, delta_t::Float64)
#  Ricker wavelet generation and integration for source
#  Dongzhuo Li @ Stanford
#  May, 2015

	e = pi*pi*f*f;
	t_delay = 1.2/f;
	source = Matrix{Float64}(undef, 1, nStep)
	for it = 1:nStep
			source[it] = (1-2*e*(delta_t*(it-1)-t_delay)^2)*exp(-e*(delta_t*(it-1)-t_delay)^2);
	end

	for it = 2:nStep
			source[it] = source[it] + source[it-1];
	end
	source = source * delta_t;
end

"""

get `vs` upper and lower bounds from log point cloud
- 1st row of Bounds: `vp ref line`
- 2nd row of Bounds: `vs high ref line`
- 3rd row of Bounds: `vs low ref line`
"""
function cs_bounds_cloud(cpImg, Bounds)
	cs_high_itp = Spline1D(Bounds[1,:], Bounds[2,:]; k=1)
	cs_low_itp = Spline1D(Bounds[1,:], Bounds[3,:]; k=1)
	csHigh = zeros(size(cpImg))
	csLow = zeros(size(cpImg))
	for i = 1:size(cpImg, 1)
		for j = 1:size(cpImg, 2)
			csHigh[i,j] = cs_high_itp(cpImg[i,j])
			csLow[i,j] = cs_low_itp(cpImg[i,j])
		end
	end
	return csHigh, csLow
end

"""
	klauderWave(fmin, fmax, t_sweep, nStepTotal, nStepDelay, delta_t)

Generates Klauder wavelet.
"""
function klauderWave(fmin, fmax, t_sweep, nStepTotal, nStepDelay, delta_t)

	nStep = nStepTotal - nStepDelay
	source = Matrix{Float64}(undef, 1, nStep+nStep-1)
	source_half = Matrix{Float64}(undef, 1, nStep-1)
	K = (fmax - fmin) / t_sweep
	f0 = (fmin + fmax) / 2.0
	t_axis = delta_t:delta_t:(nStep-1)*delta_t
	source_half = sin.(pi * K .* t_axis .* (t_sweep .- t_axis)) .* cos.(2.0 * pi * f0 .* t_axis) ./ (pi*K.*t_axis*t_sweep)
	for i = 1:nStep-1
		source[i] = source_half[end-i+1]
	end
	for i = nStep+1:2*nStep-1
		source[i] = source_half[i-nStep]
	end
	source[nStep] = 1.0
	source_crop = source[:,nStep-nStepDelay:end]
	return source_crop
end

@doc raw"""
	padding(cp, cs, den, nz_orig, nx_orig, nz, nx, nPml, nPad)

Adds PML boundaries to `cp`, `cs` and `den`.
The original `nz_orig x nx_orig` grid is resampled to `nz x nx`.

# Note
`nPad` is used to make the number of nodes in the $z$ direction a multiple of 32 (for coalesced memory access).
"""
function padding(cp, cs, den, nz_orig, nx_orig, nz, nx, nPml, nPad)
	tran_cp = tf.reshape(cp, (1, nz_orig, nx_orig, 1))
	tran_cs = tf.reshape(cs, (1, nz_orig, nx_orig, 1))
	tran_den = tf.reshape(den, (1, nz_orig, nx_orig, 1))

	tran_cp = squeeze(tf.image.resize_bilinear(tran_cp, (nz, nx)))
	tran_cs = squeeze(tf.image.resize_bilinear(tran_cs, (nz, nx)))
	tran_den = squeeze(tf.image.resize_bilinear(tran_den, (nz, nx)))

	# cp_pad = tf.pad(tran_cp, [nPml (nPml+nPad); nPml nPml], constant_values=5500.0)
	# cs_pad = tf.pad(tran_cs, [nPml (nPml+nPad); nPml nPml], constant_values=0.0)
	# den_pad = tf.pad(tran_den, [nPml (nPml+nPad); nPml nPml], constant_values=2500.0)
	cp_pad = tf.pad(tran_cp, [nPml (nPml+nPad); nPml nPml], "SYMMETRIC")
	cs_pad = tf.pad(tran_cs, [nPml (nPml+nPad); nPml nPml], "SYMMETRIC")
	den_pad = tf.pad(tran_den, [nPml (nPml+nPad); nPml nPml], "SYMMETRIC")

	cp_pad = cast(cp_pad, Float64)
	cs_pad = cast(cs_pad, Float64)
	den_pad = cast(den_pad, Float64)
	return cp_pad, cs_pad, den_pad
end

"""
	velocity_to_moduli(cp::Union{Array{Float64, 2}, PyObject},
	cs::Union{Array{Float64, 2}, PyObject},
	den::Union{Array{Float64, 2}, PyObject})

Computes the Lamé parameters using velocities.
"""
function velocity_to_moduli(cp::Union{Array{Float64, 2}, PyObject},
	cs::Union{Array{Float64, 2}, PyObject},
	den::Union{Array{Float64, 2}, PyObject})
    lambda = (cp.*cp - 2.0 * cs.*cs) .* den / 1e6
    mu = cs.*cs .* den / 1e6
    return lambda, mu
end


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct TGTypeII <: TurbineGov
        R::Float64
        T1::Float64
        T2::Float64
        τ_min::Float64
        τ_max::Float64
        P_ref::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int64
        internal::InfrastructureSystemsInternal
    end

Parameters of a Turbine Governor Type II.

# Arguments
- `R::Float64`: Droop parameter, validation range: (0, nothing)
- `T1::Float64`: Transient gain time constant, validation range: (0, nothing)
- `T2::Float64`: Power fraction time constant, validation range: (0, nothing)
- `τ_min::Float64`: Min Power into the Governor, validation range: (0, nothing)
- `τ_max::Float64`: Max Power into the Governor, validation range: (0, nothing)
- `P_ref::Float64`: Reference Power Set-point, validation range: (0, nothing)
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`
- `n_states::Int64`
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct TGTypeII <: TurbineGov
    "Droop parameter"
    R::Float64
    "Transient gain time constant"
    T1::Float64
    "Power fraction time constant"
    T2::Float64
    "Min Power into the Governor"
    τ_min::Float64
    "Max Power into the Governor"
    τ_max::Float64
    "Reference Power Set-point"
    P_ref::Float64
    ext::Dict{String, Any}
    states::Vector{Symbol}
    n_states::Int64
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function TGTypeII(R, T1, T2, τ_min, τ_max, P_ref=1.0, ext=Dict{String, Any}(), )
    TGTypeII(R, T1, T2, τ_min, τ_max, P_ref, ext, [:xg], 1, InfrastructureSystemsInternal(), )
end

function TGTypeII(; R, T1, T2, τ_min, τ_max, P_ref=1.0, ext=Dict{String, Any}(), )
    TGTypeII(R, T1, T2, τ_min, τ_max, P_ref, ext, )
end

# Constructor for demo purposes; non-functional.
function TGTypeII(::Nothing)
    TGTypeII(;
        R=0,
        T1=0,
        T2=0,
        τ_min=0,
        τ_max=0,
        P_ref=0,
        ext=Dict{String, Any}(),
    )
end

"""Get TGTypeII R."""
get_R(value::TGTypeII) = value.R
"""Get TGTypeII T1."""
get_T1(value::TGTypeII) = value.T1
"""Get TGTypeII T2."""
get_T2(value::TGTypeII) = value.T2
"""Get TGTypeII τ_min."""
get_τ_min(value::TGTypeII) = value.τ_min
"""Get TGTypeII τ_max."""
get_τ_max(value::TGTypeII) = value.τ_max
"""Get TGTypeII P_ref."""
get_P_ref(value::TGTypeII) = value.P_ref
"""Get TGTypeII ext."""
get_ext(value::TGTypeII) = value.ext
"""Get TGTypeII states."""
get_states(value::TGTypeII) = value.states
"""Get TGTypeII n_states."""
get_n_states(value::TGTypeII) = value.n_states
"""Get TGTypeII internal."""
get_internal(value::TGTypeII) = value.internal

"""Set TGTypeII R."""
set_R!(value::TGTypeII, val::Float64) = value.R = val
"""Set TGTypeII T1."""
set_T1!(value::TGTypeII, val::Float64) = value.T1 = val
"""Set TGTypeII T2."""
set_T2!(value::TGTypeII, val::Float64) = value.T2 = val
"""Set TGTypeII τ_min."""
set_τ_min!(value::TGTypeII, val::Float64) = value.τ_min = val
"""Set TGTypeII τ_max."""
set_τ_max!(value::TGTypeII, val::Float64) = value.τ_max = val
"""Set TGTypeII P_ref."""
set_P_ref!(value::TGTypeII, val::Float64) = value.P_ref = val
"""Set TGTypeII ext."""
set_ext!(value::TGTypeII, val::Dict{String, Any}) = value.ext = val
"""Set TGTypeII states."""
set_states!(value::TGTypeII, val::Vector{Symbol}) = value.states = val
"""Set TGTypeII n_states."""
set_n_states!(value::TGTypeII, val::Int64) = value.n_states = val
"""Set TGTypeII internal."""
set_internal!(value::TGTypeII, val::InfrastructureSystemsInternal) = value.internal = val


#!/usr/bin/env julia

# Given some model parameters, synthesize and plot the channel maps and
# integrated spectrum

using ArgParse

s = ArgParseSettings()
@add_arg_table s begin
    "--norad"
    help = "Use the image already here."
    action = :store_true
    "--config"
    help = "a YAML configuration file"
    default = "config.yaml"
end

parsed_args = parse_args(ARGS, s)

import YAML
config = YAML.load(open(parsed_args["config"]))

using DiskJockey.constants
using DiskJockey.image
using DiskJockey.hmodel
# using constants
# using image
# using model
using HDF5

import PyPlot.plt
using LaTeXStrings
import Images

species = config["species"]
transition = config["transition"]
lam0 = lam0s[species*transition]

# First contour is at 3 sigma, and then contours go up (or down) in multiples of spacing
function get_levels(rms::Float64, vmax::Float64, spacing=3)
    levels = Float64[]

    val = 3 * rms
    while (val < vmax)
        append!(levels, [val])
        val += rms * spacing
    end

    return levels
end

# function plot_beam(ax, BMAJ, BMIN, xy=(1,-1))
#     BMAJ = 3600. * header["BMAJ"] # [arcsec]
#     BMIN = 3600. * header["BMIN"] # [arcsec]
#     BPA =  header["BPA"] # degrees East of North
#     # from matplotlib.patches import Ellipse
#     ax[:add_artist](PyPlot.matplotlib[:patches][:Ellipse](xy=xy, width=BMIN, height=BMAJ, angle=BPA, facecolor="0.8", linewidth=0.2))
# end

# Plot the raw channel maps directly from RADMC
function plot_chmaps(img::image.RawImage)

    (im_ny, im_nx) = size(img.data)[1:2] # y and x dimensions of the image

    # CO 2-1 rest frame
    # lam0 = cc/230.538e9 * 1e4 # [microns]
    # nlam = length(img.lams)

    # convert wavelengths to velocities
    # vels = c_kms * (img.lams .- lam0)/lam0

    fig, ax = plt.subplots(nrows=2, ncols=12, figsize=(12, 2.8))

    ext = (1, im_nx, 1, im_ny) # Python array convention for Matplotlib

    for row=1:2
        for col=1:12
            iframe = col + (row - 1) * 12

            if iframe > nlam
                # Stop if we run out of channels
                break
            end

            frame = img.data[:,:,iframe]
            frame += 1e-99 #Add a tiny bit so that we don't have log10(0)
            max = maximum(log10(frame))
            ax[row, col][:imshow](log10(frame), vmin=max - 6, vmax=max, interpolation="none", origin="lower", cmap=plt.get_cmap("PuBu"), extent=ext)
            levels = linspace(max - 0.8, max, 5)
            ax[row, col][:contour](log10(frame), origin="lower", colors="k", levels=levels, linestyles="solid", linewidths=0.2, extent=ext)

            if col != 1 || row != 2
                ax[row, col][:xaxis][:set_ticklabels]([])
                ax[row, col][:yaxis][:set_ticklabels]([])
            else
                ax[row, col][:set_xlabel](L"$X$")
                ax[row, col][:set_ylabel](L"$Y$")
            end

            ax[row, col][:annotate](@sprintf("%.1f", vels[iframe]), (0.1, 0.8), xycoords="axes fraction", size=8)

        end
    end

    fig[:subplots_adjust](hspace=0.01, wspace=0.05, top=0.9, bottom=0.1, left=0.05, right=0.95)

    plt.savefig("channel_maps_raw_image.png")

end

# Plot the channel maps using sky convention
function plot_chmaps(img::image.SkyImage, fname="channel_maps_sky.png")

    (im_ny, im_nx) = size(img.data)[1:2] # y and x dimensions of the image

    # Image needs to be flipped along RA dimension
    ext = (img.ra[end], img.ra[1], img.dec[1], img.dec[end])

    # Figure out how many plots we'll have.
    ncols = 8
    nrows = iceil(nlam/ncols)

    fig, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=(12, 1.5 * nrows))

    for row=1:nrows
        for col=1:ncols
            iframe = col + (row - 1) * ncols

            if col != 1 || row != nrows
                ax[row, col][:xaxis][:set_ticklabels]([])
                ax[row, col][:yaxis][:set_ticklabels]([])
            else
                ax[row, col][:set_xlabel](L"$\Delta \alpha$ ('')")
                ax[row, col][:set_ylabel](L"$\Delta \delta$ ('')")
            end

            if iframe > nlam
                # Stop if we run out of channels
                # Plot a blank square
                ax[row, col][:imshow](zeros((im_ny, im_nx)), cmap=plt.get_cmap("PuBu"), vmin=0, vmax=20, extent=ext, origin="lower")

            else
                #Flip the frame for Sky convention
                frame = fliplr(img.data[:,:,iframe])
                frame += 1e-99 #Add a tiny bit so that we don't have log10(0)
                lframe = log10(frame)
                max = maximum(lframe)
                ix,iy = ind2sub(size(lframe), indmax(lframe))

                # ax[row, col][:imshow](lframe, extent=ext, interpolation="none", vmin=max - 6, vmax=max, origin="lower", cmap=plt.get_cmap("PuBu"))
                ax[row, col][:imshow](frame, extent=ext, interpolation="none", origin="lower", cmap=plt.get_cmap("PuBu"), norm=norm)

                # println("frame ", iframe)
                # println("levels ", levels)
                # println("extrema ", extrema(frame))
                if length(levels) > 0
                    ax[row, col][:contour](frame, origin="lower", colors="k", levels=levels, extent=ext, linestyles="solid", linewidths=0.2)
                end

                ax[row, col][:annotate](@sprintf("%.1f", vels[iframe]), (0.1, 0.8), xycoords="axes fraction", size=8)
            end

        end
    end

    fig[:subplots_adjust](hspace=0.06, wspace=0.01, top=0.9, bottom=0.1, left=0.05, right=0.95)

    plt.savefig(fname)

end

# Plot the raw array for the Sky image
function plot_chmaps_data(img::image.SkyImage)

    (im_ny, im_nx) = size(img.data)[1:2] # y and x dimensions of the image

    ll = sin(img.ra .* arcsec)
    mm = sin(img.dec .* arcsec)
    ext = (ll[1], ll[end], mm[1], mm[end])

    nlam = length(img.lams)

    fig, ax = plt.subplots(nrows=2, ncols=12, figsize=(12, 2.8))

    for row=1:2
        for col=1:12
            iframe = col + (row - 1) * 12

            if iframe > nlam
                # Stop if we run out of channels
                break
            end

            frame = img.data[:,:,iframe]
            frame += 1e-99 #Add a tiny bit so that we don't have log10(0)
            max = maximum(log10(frame))
            ax[row, col][:imshow](log10(frame), extent=ext, vmin=max - 6, vmax=max, interpolation="none", origin="lower", cmap=plt.get_cmap("PuBu"))
            levels = linspace(max - 0.8, max, 5)
            ax[row, col][:contour](log10(frame), origin="lower", colors="k", levels=levels, extent=ext, linestyles="solid", linewidths=0.2)

            if col != 1 || row != 2
                ax[row, col][:xaxis][:set_ticklabels]([])
                ax[row, col][:yaxis][:set_ticklabels]([])
            else
                ax[row, col][:set_xlabel](L"$ll$")
                ax[row, col][:set_ylabel](L"$mm$")
            end

            ax[row, col][:annotate](@sprintf("%.1f", vels[iframe]), (0.1, 0.8), xycoords="axes fraction", size=8)

        end
    end

    fig[:subplots_adjust](hspace=0.01, wspace=0.05, top=0.9, bottom=0.1, left=0.05, right=0.95)

    plt.savefig("channel_maps_sky_raw.png")

end

# Plot the spatially-integrated spectrum
function plot_spectrum(img::image.SkyImage)

    fig = plt.figure()
    ax = fig[:add_subplot](111)

    spec = imToSpec(img)

    ax[:plot](vels, spec[:,2], ls="steps-mid")
    # ax[:plot](vels, reverse(spec[:,2]), ls="steps-mid")
    ax[:set_ylabel](L"$f_\nu$ [Jy]")
    ax[:set_xlabel](L"$v$ [km/s]")

    fig[:subplots_adjust](left=0.15, bottom=0.15, right=0.85)

    plt.savefig("spectrum.png")
end

pp = config["parameters"]
params = ["M_star", "r_c", "T_10m", "q_m", "T_10a", "q_a", "T_freeze", "X_freeze", "gamma", "h", "delta", "logM_gas", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]
nparam = length(params)
starting_param = Array(Float64, nparam)

for i=1:nparam
    starting_param[i] = pp[params[i]][1]
end

# Convert logM_gas to M_gas
starting_param[12] = 10^starting_param[12]

pars = Parameters(starting_param...)
#
# vel = pars.vel # [km/s]
# # RADMC conventions for inclination and PA
incl = pars.incl # [deg]
PA = pars.PA # [deg] Position angle runs counter clockwise
npix = config["npix"] # number of pixels
#
# # Doppler shift the dataset wavelength to rest-frame wavelength
# beta = vel/c_kms # relativistic Doppler formula
# shift_lams =  lams .* sqrt((1. - beta) / (1. + beta)) # [microns]
#
# grd = config["grid"]
# grid = Grid(grd["nr"], grd["ntheta"], grd["r_in"], grd["r_out"], true)
#
# # Comment out afterward
# # nchan = 251
# # global const vels = linspace(-10., 10, nchan) # [km/s]
# # # CO 2-1 rest frame
# # lam0 = cc/230.538e9 * 1e4 # [microns]
# # # convert velocities to wavelengths
# # shift_lams = lam0 * (vels/c_kms + 1)

if !parsed_args["norad"]
    run(`radmc3d image incl $incl posang $PA npix $npix loadlambda`)
end

im = imread()

# plot_chmaps(im)

skim = imToSky(im, pars.dpc)

# Do the velocity conversion here
global nlam = length(skim.lams)
# convert wavelengths to velocities
global vels = c_kms * (skim.lams .- lam0)/lam0

# get the colorscale corresponding to the velocities
vmin, vmax = extrema(skim.data)
vvmax = maxabs(skim.data)

ldata = log10(skim.data + 1e-99)
vlmax = maximum(ldata)

# if in(config, "beam")
beam = config["beam"]
rms = beam["rms"] # Jy/beam
BMAJ = beam["BMAJ"]/2 # semi-major axis [arcsec]
BMIN = beam["BMIN"]/2 # semi-minor axis [arcsec]
BAVG = (BMAJ + BMIN)/2
BPA = beam["BPA"] # position angle East of North [degrees]

println("Beam sigma ", BAVG, " [arcsec]")

arcsec_ster = (4.25e10)
# Convert beam from arcsec^2 to Steradians
rms = rms/(pi * BMAJ * BMIN) * arcsec_ster
# println("RMS: ", rms, " vmax: ", vmax)
global levels = get_levels(rms, vmax)

norm = PyPlot.matplotlib[:colors][:Normalize](vlmax - 8, vlmax)

# println("Plotting hires maps")
# plot_chmaps(skim)

println("bluring maps")
sk_blur = blur(skim, [BAVG, BAVG])

# Now redo all this for the blurred
vmin, vmax = extrema(sk_blur.data)
global levels = get_levels(rms, vmax)
norm = PyPlot.matplotlib[:colors][:Normalize](0, vmax)

println("Plotting blured maps")
plot_chmaps(sk_blur, "channel_maps_blur.png")

plot_spectrum(skim)


#=##############################################################################################
Filename: solve_wing_viscous_drag.jl
Author: Ryan Anderson
Contact: rymanderson@gmail.com
README: define an `Action` object to solve a CCBlade rotor
=###############################################################################################

"""
    solve_wing_viscous_drag <: Action

Inputs:

* `aircraft::Aircraft` : `Aircraft` system object
* `parameters<:Parameters` `Parameters` struct
* `freestream::Freestream` : `Freestream` object
* `environment::Environment` `Environment` object
* `step_range::AbstractArray` : array of times for which the simulation is run
* `stepi::Int` : index of the current step
* `step_symbol::String` : defines the step, e.g. `alpha` or `time`

`parameters <: Parameters` requires the following elements:

* `cls::Vector{Array{Float64,2}}` : each element is an array of size (nspanwisepanels, nsteps) containing local lift coefficients at each lifting line section, corresponding to each lifting surface
* `wing_strip_ys::Vector{Vector{Float64}}` : [i][j]th element is the jth dimensional spanwise coordinate coinciding with control points of the lifting line of the ith lifting surface
* `wing_strip_polars Vector{Vector{<:CCBlade.AFType}}` : [i][j]th element is a `<: CCBlade.AFType` object corresponding to the jth lifting line section of the ith lifting surface
* `cds_viscous::Vector{Array{Float64,2}}` : each element is an array of size (nspanwisepanels, nsteps) containing viscous-corrected local drag coefficients at each lifting line section, corresponding to each lifting surface
* `CDs_viscous::Array{Float64,1}` : ith element is the drag coefficient of the aircraft at the ith step

"""
function solve_wing_viscous_drag(aircraft, parameters, freestream, environment, step_range, stepi, step_symbol)
    # extract parameters
    cls = parameters.cls
    cds = parameters.cds
    wing_strip_ys = parameters.wing_strip_ys
    wing_strip_polars = parameters.wing_strip_polars
    # apply viscous correction
    nwings = length(aircraft.wing_system.system.surfaces)
    for iwing in 1:nwings

    end
    # update cds of aircraft.wing_system.system

    # extract forces and moments
    CF_viscous, CM_viscous = VL.body_forces(aircraft.wing_system.system; frame=VL.Wind())
    # store to `parameters`
    for iwing in 1:nwings
        parameters.cds[iwing][:,stepi] = cfs[iwing][1,:]
        parameters.cmxs[iwing][:,stepi] = cms[iwing][1,:]
        parameters.cmys[iwing][:,stepi] = cms[iwing][2,:]
        parameters.cmzs[iwing][:,stepi] = cms[iwing][3,:]
    end
    parameters.cds_viscous[stepi] .= cds_viscous
    parameters.CDs_viscous[stepi] = CF_viscous[1]

    return false
end

"""
solve_wing_viscous_drag(system, step_range)

Method returns initialized elements required for the `parameters <: Parameters` struct during simulation.

Inputs:

* `aircraft::Aircraft` : aircraft system to be simulated
* `step_range::AbstractArray` : defines each step of the simulation

Outputs:

* `CLs::Array{Float64,1}` : ith element is the lift coefficient of the aircraft at the ith step
* `CDs::Array{Float64,1}` : ith element is the drag coefficient of the aircraft at the ith step
* `CYs::Array{Float64,1}` : ith element is the side force coefficient of the aircraft at the ith step

"""
function solve_wing_viscous_drag(aircraft, step_range)

    CLs = zeros(length(step_range))
    CDs = zeros(length(step_range))
    CYs = zeros(length(step_range))

    return CLs, CDs, CYs
end


#*****************************************************************************#
# General initialization functions for all integrators                        #
#*****************************************************************************#

"Print error for integrators not implemented, yet."
function Integrator(equation::Equation, tableau::AbstractTableau, Δt)
    error("No integrator found for equation ", equation, " and tableau ", tableau)
end

#*****************************************************************************#
# Initialization functions for deterministic integrators                      #
#*****************************************************************************#

"Create integrator for explicit Runge-Kutta tableau."
function Integrator(equation::ODE, tableau::TableauERK, Δt)
    IntegratorERK(equation, tableau, Δt)
end

"Create integrator for diagonally implicit Runge-Kutta tableau."
function Integrator(equation::ODE, tableau::TableauDIRK, Δt)
    IntegratorDIRK(equation, tableau, Δt)
end

"Create integrator for fully implicit Runge-Kutta tableau."
function Integrator(equation::ODE, tableau::TableauFIRK, Δt)
    IntegratorFIRK(equation, tableau, Δt)
end

"Create integrator for explicit partitioned Runge-Kutta tableau."
function Integrator(equation::PODE, tableau::TableauEPRK, Δt)
    IntegratorEPRK(equation, tableau, Δt)
end

"Create integrator for implicit partitioned Runge-Kutta tableau."
function Integrator(equation::PODE, tableau::TableauIPRK, Δt)
    IntegratorIPRK(equation, tableau, Δt)
end

"Create integrator for variational partitioned Runge-Kutta tableau."
function Integrator(equation::IODE, tableau::TableauVPRK, Δt)
    IntegratorVPRK(equation, tableau, Δt)
end

"Create integrator for formal Lagrangian Runge-Kutta tableau."
function Integrator(equation::VODE, tableau::TableauFIRK, Δt)
    IntegratorFLRK(equation, tableau, Δt)
end

"Create integrator for Projected Gauss-Legendre Runge-Kutta tableau."
function Integrator(equation::IODE, tableau::CoefficientsPGLRK, Δt)
    IntegratorPGLRK(equation, tableau, Δt)
end

"Create integrator for variational partitioned additive Runge-Kutta tableau."
function Integrator(equation::IDAE, tableau::TableauVPARK, Δt)
    IntegratorVPARK(equation, tableau, Δt)
end

"Create integrator for variational special partitioned additive Runge-Kutta tableau."
function Integrator(equation::IDAE, tableau::TableauVSPARK, Δt)
    IntegratorVSPARK(equation, tableau, Δt)
end

"Create integrator for variational special partitioned additive Runge-Kutta tableau with projection on primary constraint."
function Integrator(equation::IDAE, tableau::TableauVSPARKprimary, Δt)
    IntegratorVSPARKprimary(equation, tableau, Δt)
end

"Create integrator for variational special partitioned additive Runge-Kutta tableau with projection on secondary constraint."
function Integrator(equation::VDAE, tableau::TableauVSPARKsecondary, Δt)
    IntegratorVSPARKsecondary(equation, tableau, Δt)
end

"Create integrator for Hamiltonian partitioned additive Runge-Kutta tableau."
function Integrator(equation::PDAE, tableau::TableauHPARK, Δt)
    IntegratorHPARK(equation, tableau, Δt)
end

"Create integrator for Hamiltonian special partitioned additive Runge-Kutta tableau."
function Integrator(equation::PDAE, tableau::TableauHSPARK, Δt)
    IntegratorHSPARK(equation, tableau, Δt)
end

"Create integrator for Hamiltonian special partitioned additive Runge-Kutta tableau with projection on primary constraint."
function Integrator(equation::PDAE, tableau::TableauHSPARKprimary, Δt)
    IntegratorHSPARKprimary(equation, tableau, Δt)
end

"Create integrator for splitting tableau."
function Integrator(equation::SODE, tableau::AbstractTableauSplitting, Δt)
    IntegratorSplitting(equation, tableau, Δt)
end

#*****************************************************************************#
# Initialization functions for stochastic integrators                         #
#*****************************************************************************#

"Create integrator for stochastic explicit Runge-Kutta tableau."
function Integrator(equation::SDE, tableau::TableauSERK, Δt)
    IntegratorSERK(equation, tableau, Δt)
end

"Create integrator for weak explicit Runge-Kutta tableau."
function Integrator(equation::SDE, tableau::TableauWERK, Δt)
    IntegratorWERK(equation, tableau, Δt)
end

"Create integrator for stochastic fully implicit Runge-Kutta tableau."
function Integrator(equation::SDE, tableau::TableauSIRK, Δt; K::Int=0)
    IntegratorSIRK(equation, tableau, Δt, K=K)
end

"Create integrator for stochastic fully implicit partitioned Runge-Kutta tableau."
function Integrator(equation::PSDE, tableau::TableauSIPRK, Δt; K::Int=0)
    IntegratorSIPRK(equation, tableau, Δt, K=K)
end

"Create integrator for stochastic fully implicit split partitioned Runge-Kutta tableau."
function Integrator(equation::SPSDE, tableau::TableauSISPRK, Δt; K::Int=0)
    IntegratorSISPRK(equation, tableau, Δt, K=K)
end

"Create integrator for weak fully implicit Runge-Kutta tableau."
function Integrator(equation::SDE, tableau::TableauWIRK, Δt)
    IntegratorWIRK(equation, tableau, Δt)
end

#*****************************************************************************#
# General integration functions for all integrators                           #
#*****************************************************************************#

"Apply integrator for ntime time steps and return solution."
function integrate(integrator::Integrator, ntime::Int, nsave::Int=DEFAULT_NSAVE)
    solution = Solution(equation(integrator), timestep(integrator), ntime, nsave)
    integrate!(integrator, solution)
    return solution
end

"Integrate given equation with given tableau for ntime time steps and return solution."
function integrate(equation::Equation, tableau::AbstractTableau, Δt, ntime, nsave=DEFAULT_NSAVE)
    return integrate(Integrator(equation, tableau, Δt), ntime, nsave)
end

"Integrate ODE specified by vector field and initial condition with given tableau for ntime time steps and return solution."
function integrate(f::Function, x₀::Vector, tableau::AbstractTableau, Δt, ntime, nsave=DEFAULT_NSAVE; t₀=0)
    return integrate(ODE(f, t₀, x₀), tableau, Δt, ntime, nsave)
end

"Integrate PODE specified by two vector fields and initial conditions with given tableau for ntime time steps and return solution."
function integrate(v::Function, f::Function, q₀::Vector, p₀::Vector, tableau::AbstractTableau, Δt, ntime, nsave=DEFAULT_NSAVE; t₀=0)
    return integrate(PODE(v, f, t₀, q₀, p₀), tableau, Δt, ntime, nsave)
end

#*****************************************************************************#
# Integration functions for deterministic integrators                         #
#*****************************************************************************#

"Integrate equation for all initial conditions."
function integrate!(int::DeterministicIntegrator, sol::Solution)
    integrate!(int, sol, 1, sol.ni)
end

"Integrate ODE for initial conditions m with m₁ ≤ m ≤ m₂."
function integrate!(int::DeterministicIntegrator, sol::Solution, m1, m2)
    # integrate initial conditions m with m₁ ≤ m ≤ m₂ for all time steps
    integrate!(int, sol, m1, m2, 1, sol.ntime)
end

"Integrate ODE for initial conditions m with m₁ ≤ m ≤ m₂ for time steps n with n₁ ≤ n ≤ n₂."
function integrate!(int::DeterministicIntegrator{DT,TT}, sol::Solution{DT,TT,N}, m1::Int, m2::Int, n1::Int, n2::Int) where {DT,TT,N}
    @assert m1 ≥ 1
    @assert m2 ≥ m1
    @assert m2 ≤ sol.ni

    @assert n1 ≥ 1
    @assert n2 ≥ n1
    @assert n2 ≤ sol.ntime

    asol = AtomicSolution(equation(int))

    # loop over initial conditions showing progress bar
    for m in m1:m2
        # get cache from solution
        get_initial_conditions!(sol, asol, m, n1)
        initialize!(int, asol)

        # loop over time steps
        for n in n1:n2
            # try
                # integrate one initial condition for one time step
                integrate_step!(int, asol)

                # take care of periodic solutions
                cut_periodic_solution!(asol, periodicity(equation(int)))

                # copy solution from cache to solution
                set_solution!(sol, asol, n, m)
            # catch ex
            #     tstr = " in time step " * string(n)
            #
            #     if m1 ≠ m2
            #         tstr *= " for initial condition " * string(m)
            #     end
            #
            #     tstr *= "."
            #
            #     if isa(ex, DomainError)
            #         @warn("Domain error" * tstr)
            #     elseif isa(ex, ErrorException)
            #         @warn("Simulation exited early" * tstr)
            #         @warn(ex.msg)
            #     else
            #         @warn(string(typeof(ex)) * tstr)
            #         throw(ex)
            #     end
            # end
        end
    end
end

#*****************************************************************************#
# Integration functions for stochastic integrators                            #
#*****************************************************************************#

"Integrate SDE for all sample paths and initial conditions."
function integrate!(int::StochasticIntegrator, sol::StochasticSolution)
    integrate!(int, sol, 1, sol.ns, 1, sol.ni)
end

"Integrate SDE for the sample paths k with k₁ ≤ k ≤ k₂ and initial conditions m with m₁ ≤ m ≤ m₂."
function integrate!(int::StochasticIntegrator, sol::StochasticSolution, k1::Int, k2::Int, m1::Int, m2::Int)
    # initialize integrator for initial conditions m with m₁ ≤ m ≤ m₂ and time step 0
    initialize!(int, sol, k1, k2, m1, m2)

    # integrate initial conditions m with m₁ ≤ m ≤ m₂ for all time steps
    integrate!(int, sol, k1, k2, m1, m2, 1, sol.ntime)
end

"Integrate SDE for the sample paths k with k₁ ≤ k ≤ k₂, the initial conditions m with m₁ ≤ m ≤ m₂, and the time steps n with n₁ ≤ n ≤ n₂."
function integrate!(int::StochasticIntegrator{DT,TT}, sol::StochasticSolution{DT,TT,N}, k1::Int, k2::Int, m1::Int, m2::Int, n1::Int, n2::Int) where {DT,TT,N}
    @assert k1 ≥ 1
    @assert k2 ≥ k1
    @assert k2 ≤ sol.ns

    @assert m1 ≥ 1
    @assert m2 ≥ m1
    @assert m2 ≤ sol.ni

    @assert n1 ≥ 1
    @assert n2 ≥ n1
    @assert n2 ≤ sol.ntime

    # loop over initial conditions
    for m in m1:m2
        # loop over sample paths
        for k in k1:k2
            # loop over time steps
            for n in n1:n2
                # try
                    # integrate one initial condition for one time step
                    integrate_step!(int, sol, k, m, n)
                # catch ex
                #     tstr = " in time step " * string(n)
                #
                #     if m1 ≠ m2
                #         tstr *= " for initial condition " * string(m)
                #     end
                #
                #     tstr *= "."
                #
                #     if isa(ex, DomainError)
                #         @warn("Domain error" * tstr)
                #     elseif isa(ex, ErrorException)
                #         @warn("Simulation exited early" * tstr)
                #         @warn(ex.msg)
                #         throw(ex)
                #     else
                #         @warn(string(typeof(ex)) * tstr)
                #         throw(ex)
                #     end
                # end
            end
        end
    end
end

"Initialize stochastic integrator for the sample paths k with k₁ ≤ k ≤ k₂, initial conditions m with m₁ ≤ m ≤ m₂ and time step 0."
function initialize!(int::StochasticIntegrator, sol::StochasticSolution, k1::Int, k2::Int, m1::Int, m2::Int)
    for m in m1:m2
        for k in k1:k2
            # initialize integrator for the sample path k, initial condition m and time step 0
            initialize!(int, sol, k, m)
        end
    end
end

# TODO Add solver status information to all integrators (if requested).


function menu(savelocation::String = "."; compile::Bool = false)
    println('='^80)
    println("Policy Evaluation with Multiple Instrumental Variables")
    println("Mogstad, Torgovitsky, and Walters (2021)")
    println('='^80)
    println("What do you want to run?")
    println("\t 1. Illustration of mutual consistency (Figure 2)")
    println("\t 2. Numerical illustration for ATT (Figure 4)")
    println("\t 3. Numerical illustration for LATE(+20) (Figure 5)")
    println("\t 4. PRTE misspecification example (Figure 6)")
    println("\t 5. Everything")
    print("Enter choice: ")
    choice = readline()
    choice = parse(Int64, choice)

    if choice == 1
        savedir, _ = setup(savelocation, stub = "illustrate-mc")
        run_illustrate_mc(savedir, compile)
    elseif choice == 2
        savedir, _ = setup(savelocation, stub = "simulation-att")
        run_simulation_att(savedir, compile)
    elseif choice == 3
        savedir, _ = setup(savelocation, stub = "simulation-prte")
        run_simulation_prte(savedir, compile)
    elseif choice == 4
        savedir, _ = setup(savelocation, stub = "prte-misspecification")
        run_prte_misspecification(savedir, compile)
    elseif choice == 5
        savedir, _ = setup(savelocation, stub = "everything")
        results_mc = run_illustrate_mc(savedir, compile)
        results_att = run_simulation_att(savedir, compile)
        results_prte = run_simulation_prte(savedir, compile)
        results_prte_misspecification =
            run_prte_misspecification(savedir, compile)
        return Dict(:results_mc => results_mc,
                    :results_att => results_att,
                    :results_prte => results_prte)
    end
end
export menu

function run_illustrate_mc(savedir::String, compile::Bool = false)
    results = illustrate_mc()
    fn_results = joinpath(savedir, "illustrate-mc.csv")
    CSV.write(fn_results, results)
    if compile
        compile_latex(joinpath(savedir, "FigureMTRConsistent.tex"))
        compile_latex(joinpath(savedir, "FigureMTRInconsistent.tex"))
    end
    return results
end
export run_illustrate_mc

function run_simulation_att(savedir::String, compile::Bool = false)
    results = simulation_att()
    fn_results = joinpath(savedir, "simulation-att.csv")
    CSV.write(fn_results, results)
    if compile
        compile_latex(joinpath(savedir, "FigureATT.tex"))
    end
    return results
end
export run_simulation_att

function run_simulation_prte(savedir::String, compile::Bool = false)
    results = simulation_prte()
    fn_results = joinpath(savedir, "simulation-prte.csv")
    CSV.write(fn_results, results)
    if compile
        compile_latex(joinpath(savedir, "FigurePRTE.tex"))
    end
    return results
end
export run_simulation_prte

function run_prte_misspecification(savedir::String, compile::Bool = false)
    results = prte_misspecification()
    fn_results = joinpath(savedir, "prte-misspecification.csv")
    CSV.write(fn_results, results)
    if compile
        compile_latex(joinpath(savedir, "FigurePRTEMisspecification.tex"))
    end
    return results
end
export run_simulation_prte

function compile_latex(fn::String)
    oldwd = pwd()
    try
        cd(dirname(fn))
        cstr = `pdflatex -halt-on-error $(basename(fn)) "|" grep -a3 ^!`
        @suppress begin
            run(cstr)
            run(`latexmk -c`)
        end
        cd(oldwd)
    catch err
        cd(oldwd)
    end
end


function validate_libsize(libsize, source, dim, τ, η, replace)
    n_available_pts = length(source) - dim*τ - abs(η)
    if libsize > n_available_pts
        if replace
            #@warn "libsize = $libsize > n_available_pts = $n_available_pts. Sampling with replacement."
        else
            throw(DomainError(libsize, "libsize = $libsize > n_available_pts = $n_available_pts. Reduce `libsize` (preferably to something much smaller than `n_available_pts = $n_available_pts`) or enable sampling with replacement."))
        end
    end
end

function validate_embedding!(embedding, jitter)
    maxval = abs.(maximum(embedding))
    unique_pts = unique(embedding, dims = 2)
    if size(unique_pts, 2) < size(embedding, 2)
        #@warn "Not all embedding points are unique. Jittering coordinates by `rand(Uniform(-$jitter*maximum(embedding), $jitter*maximum(embedding)))`"

        for i in 1:length(embedding)
            embedding[i] += rand(Uniform(-jitter*maxval, jitter*maxval))
        end
    end
end

function validate_theiler_window!(theiler_window, points_available)
    if theiler_window < 0
        throw(DomainError(theiler_window, "`theiler_window=$theiler_window`. Must be ≧ 0."))
    end
    if theiler_window >= ceil(Int, 0.5*points_available)
        throw(DomainError(theiler_window, "`theiler_window=$theiler_window >= ceil(Int, 0.5*(length(target) - dim*τ))=$points_available`. Please reduce `theiler_window`."))
    end
end

function validate_embedding_params(dim, τ, points_available, theiler_window)
    if dim == 1
        @warn "`dim=$dim`, but must be at least 2 to construct an embedding."
    end
    if dim <= 0
        throw(DomainError(dim, "`dim=$dim` must be at least 2 to construct an embedding."))
    end
    if τ == 0
        @warn "`τ=$τ` does not produce a delay embedding!"
    end
    if τ < 0
        @warn "`τ=$τ` is negative. The cross mapping algorithm is implemented assuming τ is positive."
    end
    if dim*τ >= ceil(Int, 0.5*points_available)
        throw(DomainError(theiler_window, "`theiler_window=$theiler_window >= ceil(Int, 0.5*(length(target) - dim*τ))=$points_available`. Please reduce `theiler_window`."))
    end
end

function validate_uncertainty_measure(uncertainty_measure)
    if uncertainty_measure ∉ [:quantile, :std, :none]
        throw(DomainError(uncertainty_measure, "uncertainty_measure = $uncertainty_measure is invalid. Must be either :quantile, :std or :none."))
    end
end

function validate_average_measure(average_measure)
    if average_measure ∉ [:mean, :median, :none]
        throw(DomainError(average_measure, "average_measure = $average_measure is invalid. Must be either :median, :mean or :none."))
    end
end

function validate_output_selection(average_measure, uncertainty_measure, summarise)
    if ((average_measure, uncertainty_measure) == (:none, :none)) && summarise
        throw(ErrorException("Setting both average_measure and uncertainty_measure to :none while summarise = true will not produce output."))
    end
end

export
validate_libsize,
validate_embedding!,
validate_embedding_params,
validate_theiler_window!,
validate_uncertainty_measure,
validate_average_measure,
validate_output_selection


Base.size(x::Atom{T}) where {T} = x.sizes
Base.size(x::Atom{Cylinder}) = x.sizes[:h]

import Base.eltype
eltype(laser::Laser{PlaneWave3D}) = "PlaneWave3D"
eltype(laser::Laser{Gaussian3D}) = "Gaussian3D"

eltype(problem::LinearOptics{Scalar}) = "Scalar"
eltype(problem::LinearOptics{Vectorial}) = "Vectorial"

eltype(problem::NonLinearOptics{MeanField}) = "MeanField"

function Base.show(io::IO, atoms::Atom{Cube})
    geometry_text = "Atoms on a $( highlight("Cube", :yellow) )"
    N_text = " with N=$(highlight(atoms.N, :yellow)) and "
    size_text = "side kL=$(highlight(make_short(atoms.sizes), :yellow))"
    return printstyled(io, geometry_text * N_text * size_text)
end
function Base.show(io::IO, atoms::Atom{Sphere})
    geometry_text = "Atoms on a $( highlight("Sphere", :yellow) )"
    N_text = " with N=$(highlight(atoms.N, :yellow)) and "
    size_text = "side kR=$(highlight(make_short(atoms.sizes), :yellow))"
    return printstyled(io, geometry_text * N_text * size_text)
end
function Base.show(io::IO, atoms::Atom{Cylinder})
    geometry_text = "Atoms on a $( highlight("Cylinder", :yellow) )"
    N_text = " with N=$(highlight(atoms.N, :yellow)) and "
    size_text = "[R=$(highlight(make_short(atoms.sizes[:R]), :yellow)), h=$(highlight(make_short(atoms.sizes[:h]), :yellow))]"
    return printstyled(io, geometry_text * N_text * size_text)
end

function Base.show(io::IO, laser::Laser{PlaneWave3D})
    l_t = "$( highlight("PlaneWave 3D", :yellow) ) "
    dr_t = " laser at direction=$(highlight(laser.pump.direction, :yellow)), "
    s_t = " s=$(highlight(make_short(laser.s), :yellow)) and "
    d_t = "Δ=$(highlight(make_short(laser.Δ), :yellow))"
    return printstyled(io, l_t * dr_t * s_t * d_t)
end

function Base.show(io::IO, laser::Laser{Gaussian3D})
    l_t = "$( highlight("Gaussian 3D", :yellow) )"
    w_t = " laser with waist w₀=$(highlight(laser.pump.w₀, :yellow)), "
    s_t = "s=$(highlight(make_short(laser.s), :yellow)) and "
    d_t = "Δ=$(highlight(make_short(laser.Δ), :yellow))"
    return printstyled(io, l_t * w_t * s_t * d_t)
end

function Base.show(io::IO, problem::LinearOptics{T}) where {T<:Linear}
    s_t = "$( highlight("LinearOptics", :blue) ): $( highlight(eltype(problem), :yellow) ) problem "
    N_t = "with N=$(highlight(problem.atoms.N, :yellow)) atoms, "
    l_t = "and $(highlight(eltype(problem.laser), :yellow)) laser"
    return printstyled(io, s_t * N_t * l_t)
end
function Base.show(io::IO, problem::NonLinearOptics{T}) where {T<:NonLinear}
    s_t = "$( highlight("NonLinearOptics", :blue) ): $( highlight(eltype(problem), :yellow) ) problem "
    N_t = "with N=$(highlight(problem.atoms.N, :yellow)) atoms, "
    l_t = "and $(highlight(eltype(problem.laser), :yellow)) laser"
    return printstyled(io, s_t * N_t * l_t)
end

function highlight(s, colour)
    io = IOBuffer()
    printstyled(IOContext(io, :color => true), s, color = colour)
    return io |> take! |> String
end
function make_short(x)
    if (abs(x) ≈ 0)
        return x
    elseif (abs(x) ≥ 1e3) || (abs(x) ≤ 1e-3)
        return @sprintf("%.2E", x)
    else
        return round(x; digits = 2)
    end
end


#= Atom Operations
Functions to extract information from OxygenAtoms.

=#

function valence_electrons(a::OxygenAtom)::Int
    sum(map(e -> e[end], a.orbitals.orbitals))
end

function neighbors(mol::OxygenMol, idx::Int)::Array{OxygenAtom,1}
    atom_neighbors = mol.adj[idx]
    [mol.atoms[i] for (i, _) in atom_neighbors]
end