function set_velocity_tracer_fields(arch, grid, fieldname, value, answer)
    model = IncompressibleModel(architecture=arch, float_type=eltype(grid), grid=grid)
    kwarg = Dict(fieldname=>value)
    set!(model; kwarg...)

    if fieldname ∈ propertynames(model.velocities)
        ϕ = getproperty(model.velocities, fieldname)
    else
        ϕ = getproperty(model.tracers, fieldname)
    end

    return interior(ϕ) ≈ answer
end

function initial_conditions_correctly_set(arch, FT)
    model = IncompressibleModel(grid=RegularCartesianGrid(FT, size=(16, 16, 8), extent=(1, 2, 3)),
                                architecture=arch, float_type=FT)

    # Set initial condition to some basic function we can easily check for.
    # We offset the functions by an integer so that we don't end up comparing
    # zero values with other zero values. I was too lazy to pick clever functions.
    u₀(x, y, z) = 1 + x + y + z
    v₀(x, y, z) = 2 + sin(x * y * z)
    w₀(x, y, z) = 3 + y * z
    T₀(x, y, z) = 4 + tanh(x + y - z)
    S₀(x, y, z) = 5

    set!(model, u=u₀, v=v₀, w=w₀, T=T₀, S=S₀)

    Nx, Ny, Nz = model.grid.Nx, model.grid.Ny, model.grid.Nz
    xC, yC, zC = nodes(model.tracers.T)
    xF, yF, zF = nodes((Face, Face, Face), model.grid)
    u, v, w = model.velocities.u.data, model.velocities.v.data, model.velocities.w.data
    T, S = model.tracers.T.data, model.tracers.S.data

    all_values_match = true
    for i in 1:Nx, j in 1:Ny, k in 1:Nz
        values_match = ( u[i, j, k] ≈ 1 + xF[i] + yC[j] + zC[k]       &&
                         v[i, j, k] ≈ 2 + sin(xC[i] * yF[j] * zC[k])  &&
                         w[i, j, k] ≈ 3 + yC[j] * zF[k]               &&
                         T[i, j, k] ≈ 4 + tanh(xC[i] + yC[j] - zC[k]) &&
                         S[i, j, k] ≈ 5)
        all_values_match = all_values_match & values_match
    end

    return all_values_match
end

@testset "Models" begin
    @info "Testing models..."

    @testset "Model constructor errors" begin
        grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1))
        @test_throws TypeError IncompressibleModel(architecture=CPU, grid=grid)
        @test_throws TypeError IncompressibleModel(architecture=GPU, grid=grid)
        @test_throws TypeError IncompressibleModel(grid=grid, boundary_conditions=1)
        @test_throws TypeError IncompressibleModel(grid=grid, forcing=2)
        @test_throws TypeError IncompressibleModel(grid=grid, background_fields=3)
    end

    topos = ((Periodic, Periodic, Periodic),
             (Periodic, Periodic,  Bounded),
             (Periodic,  Bounded,  Bounded),
             (Bounded,   Bounded,  Bounded))

    for topo in topos
        @testset "$topo model construction" begin
            @info "  Testing $topo model construction..."
            for arch in archs, FT in float_types
		        arch isa GPU && topo == (Bounded, Bounded, Bounded) && continue

                grid = RegularCartesianGrid(FT, topology=topo, size=(16, 16, 2), extent=(1, 2, 3))
                model = IncompressibleModel(grid=grid, architecture=arch, float_type=FT)

                # Just testing that the model was constructed with no errors/crashes.
                @test model isa IncompressibleModel

                # Test that the grid didn't get mangled
                @test grid === model.grid
            end
        end
    end

    @testset "Adjustment of halos in IncompressibleModel constructor" begin
        @info "  Testing adjustment of halos in IncompressibleModel constructor..."

        default_grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 2, 3))
        funny_grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 2, 3), halo=(1, 3, 4))

        # Model ensures that halos are at least of size 1
        model = IncompressibleModel(grid=default_grid)
        @test model.grid.Hx == 1 && model.grid.Hy == 1 && model.grid.Hz == 1

        model = IncompressibleModel(grid=funny_grid)
        @test model.grid.Hx == 1 && model.grid.Hy == 3 && model.grid.Hz == 4

        # Model ensures that halos are at least of size 2
        for scheme in (CenteredFourthOrder(), UpwindBiasedThirdOrder())
            model = IncompressibleModel(advection=scheme, grid=default_grid)
            @test model.grid.Hx == 2 && model.grid.Hy == 2 && model.grid.Hz == 2

            model = IncompressibleModel(advection=scheme, grid=funny_grid)
            @test model.grid.Hx == 2 && model.grid.Hy == 3 && model.grid.Hz == 4
        end

        # Model ensures that halos are at least of size 3
        for scheme in (WENO5(), UpwindBiasedFifthOrder())
            model = IncompressibleModel(advection=scheme, grid=default_grid)
            @test model.grid.Hx == 3 && model.grid.Hy == 3 && model.grid.Hz == 3

            model = IncompressibleModel(advection=scheme, grid=funny_grid)
            @test model.grid.Hx == 3 && model.grid.Hy == 3 && model.grid.Hz == 4
        end

        # Model ensures that halos are at least of size 2 with AnisotropicBiharmonicDiffusivity
        model = IncompressibleModel(closure=AnisotropicBiharmonicDiffusivity(), grid=default_grid)
        @test model.grid.Hx == 2 && model.grid.Hy == 2 && model.grid.Hz == 2

        model = IncompressibleModel(closure=AnisotropicBiharmonicDiffusivity(), grid=funny_grid)
        @test model.grid.Hx == 2 && model.grid.Hy == 3 && model.grid.Hz == 4
    end

    @testset "Model construction with single tracer and nothing tracer" begin
        @info "  Testing model construction with single tracer and nothing tracer..."
        for arch in archs
            grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 2, 3))

            model = IncompressibleModel(grid=grid, architecture=arch, tracers=:c, buoyancy=nothing)
            @test model isa IncompressibleModel

            model = IncompressibleModel(grid=grid, architecture=arch, tracers=nothing, buoyancy=nothing)
            @test model isa IncompressibleModel
        end
    end

    @testset "Non-dimensional model" begin
        @info "  Testing non-dimensional model construction..."
        for arch in archs, FT in float_types
            grid = RegularCartesianGrid(FT, size=(16, 16, 2), extent=(3, 2, 1))
            model = NonDimensionalModel(architecture=arch, float_type=FT, grid=grid, Re=1, Pr=1, Ro=Inf)

            # Just testing that a NonDimensionalModel was constructed with no errors/crashes.
            @test model isa IncompressibleModel
        end
    end

    @testset "Setting model fields" begin
        @info "  Testing setting model fields..."
        for arch in archs, FT in float_types
            N = (4, 4, 4)
            L = (2π, 3π, 5π)

            grid = RegularCartesianGrid(FT, size=N, extent=L)
            x, y, z = nodes((Face, Cell, Cell), grid, reshape=true)

            u₀(x, y, z) = x * y^2 * z^3
            u_answer = @. x * y^2 * z^3

            T₀ = rand(size(grid)...)
            T_answer = deepcopy(T₀)

            @test set_velocity_tracer_fields(arch, grid, :u, u₀, u_answer)
            @test set_velocity_tracer_fields(arch, grid, :T, T₀, T_answer)
            @test initial_conditions_correctly_set(arch, FT)
        end
    end
end


# Train an normalizing flow for seismic denosing
# Authors: Ali Siahkoohi, alisk@gatech.edu
#          Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
#          Philipp Witte, pwitte3@gatech.edu
# Date: September 2020
"""
Example:
julia train_hint_denoising.jl --batchsize 256 --max_epoch 5
"""

using DrWatson
@quickactivate :FastApproximateInference

using InvertibleNetworks
using JLD
using JLD2
using Random
using Statistics
using ArgParse
using BSON
using Printf
using Flux: gpu, cpu
using Flux.Data: DataLoader
using Flux.Optimise: Optimiser, ExpDecay, update!, ADAM

# Random seed
Random.seed!(19)

s = ArgParseSettings()
@add_arg_table! s begin
    "--max_epoch"
        help = "Maximum # of epochs"
        arg_type = Int
        default = 50
    "--lr"
        help = "starting learning rate"
        arg_type = Float32
        default = 1f-3
    "--lr_step"
        help = "lr scheduler applied at evey lr_step epoch"
        arg_type = Int
        default = 1
    "--batchsize"
        help = "batch size"
        arg_type = Int
        default = 4
    "--n_hidden"
        help = "# hidden channels"
        arg_type = Int
        default = 64
    "--depth"
        help = "depth of the network"
        arg_type = Int
        default = 12
    "--sigma"
        help = "noise std"
        arg_type = Float32
        default = 1f-1
    "--sim_name"
        help = "simulation name"
        arg_type = String
        default = "seismic-denoising-256x256-0.3CS"
end
parsed_args = parse_args(s)

max_epoch = parsed_args["max_epoch"]
lr = parsed_args["lr"]
lr_step = parsed_args["lr_step"]
batchsize = parsed_args["batchsize"]
n_hidden = parsed_args["n_hidden"]
depth = parsed_args["depth"]
sigma = parsed_args["sigma"]
sim_name = parsed_args["sim_name"]

# Define raw data directory
mkpath(datadir("training-data"))
data_path = datadir("training-data", "seismic_samples_256_by_256_num_10k.jld")
label_path = datadir("training-data", "seismic_samples_256_by_256_num_10k_labels.jld")

# Download the dataset into the data directory if it does not exist
if isfile(data_path) == false
    run(`wget https://www.dropbox.com/s/vapl62yhh8fxgwy/'
        'seismic_samples_256_by_256_num_10k.jld -q -O $data_path`)
end
if isfile(label_path) == false
    run(`wget https://www.dropbox.com/s/blxkh6tdszudlcq/'
        'seismic_samples_256_by_256_num_10k_labels.jld -q -O $label_path`)
end

# Load seismic images and create training and testing data
JLD.jldopen(data_path, "r") do file
    X_orig = read(file, "X")

    labels = JLD.jldopen(label_path, "r")["labels"][:]
    idx = findall(x -> x == 1, labels)

    X_orig = X_orig[:, :, :, idx]
    global nx, ny, nc, nsamples = size(X_orig)

    # Whiten data
    X_orig = whiten_input(X_orig)

    # Forward operator
    global A = RombergOP(nx, ny, nc, 2f0)

    # Create observed data
    Y_orig = A*X_orig + sigma*randn(Float32, size(X_orig))

    # Split in training/testing
    global ntrain = Int(floor((nsamples*.9)))
    global train_idx = randperm(nsamples)[1:ntrain]

    # Save data for testing phase
    save_dict = @strdict train_idx Y_orig A
    @tagsave(
        datadir("testing-data", "noisy_seismic_samples_256_by_256_num_10k.jld2"),
        save_dict
    )

    # Dimensions after wavelet squeeze to increase no. of channels
    global nx = Int(nx/2)
    global ny = Int(ny/2)
    global n_in = Int(nc*4)
end

# Create network
CH = NetworkConditionalHINT(
    n_in, n_hidden, depth, k1=3, k2=3, p1=1, p2=1
) |> gpu

# Training
# Batch extractor
num_batches = cld(ntrain, batchsize)

X_train = JLD.jldopen(data_path, "r")["X"]
labels = JLD.jldopen(label_path, "r")["labels"][:]
idx = findall(x -> x == 1, labels)
X_train = X_train[:, :, :, idx]

Y_train = JLD2.jldopen(
    datadir("testing-data", "noisy_seismic_samples_256_by_256_num_10k.jld2"),
    "r"
)["Y_orig"]

train_loader = DataLoader(train_idx, batchsize=batchsize, shuffle=true)

# Optimizer
opt = Optimiser(ExpDecay(lr, .9f0, num_batches*lr_step, 0f0), ADAM(lr))

# Training log keeper
fval = zeros(Float32, num_batches*max_epoch)

for epoch=1:max_epoch
    for (itr, idx) in enumerate(train_loader)

        # Apply wavelet squeeze (change dimensions to -> n1/2 x n2/2 x nc*4 x ntrain)
        X = wavelet_squeeze(whiten_input(X_train[:, :, :, idx]))
        Y = wavelet_squeeze(Y_train[:, :, :, idx])

        X = X |> gpu
        Y = Y |> gpu

        fval[(epoch-1)*num_batches + itr] = loss_supervised(CH, X, Y)[1]

        @printf(
            "Epoch: [%d/%d] | Itr: [%d/%d] | Average negative log-likelihood: %4.2f\r",
            epoch, max_epoch, itr, num_batches, fval[(epoch-1)*num_batches + itr]
        )

        # Update params
        for p in get_params(CH)
            update!(opt, p.data, p.grad)
        end
        clear_grad!(CH)
    end

    # Saving parameters and logs
    Params = get_params(CH) |> cpu
    save_dict = @strdict epoch max_epoch lr lr_step batchsize n_hidden depth sigma sim_name A Params fval
    @tagsave(
        datadir(sim_name, savename(save_dict, "jld2"; digits=6)),
        save_dict;
        safe=true
    )

end


using DICOM, Lucas
using DocOpt

args = docopt("""showdicom.jl

Usage:
  dicom.jl show [-v] [--tag=<tag>] (--dir=<dir> | <files>...)
  dicom.jl (dump | debug) [-v] (--dir=<dir> | <files>...)
  dicom.jl -h | --help

Options:
  -h --help     Show this screen.
  -v --verbose  Show more verbose output.
  --tag=<tag>   Shows only one specific tag (tag syntax is 0xabcd,0x1234).
  --dir=<dir>   Use all files of `dir`, which is either `train` or `test`.
""")

if args["--tag"] != nothing
    args["--tag"] = tuple([parse(UInt16, t) for t in split(args["--tag"], ',')]...)
end

if args["--dir"] != nothing
    println("Collecting files...")
    args["<files>"] = UTF8String[]
    for d in readdir(args["--dir"])
        patientdir = joinpath(args["--dir"], d, "study")
        for s in readdir(patientdir)
            for f in readdir(joinpath(patientdir, s))
                if splitext(f)[2] == ".dcm"
                    push!(args["<files>"], joinpath(patientdir, s, f))
                end
            end
        end
    end
end

for name in args["<files>"]
    if args["--verbose"]
        println("FILE: ", name)
        println("=====")
    end
    dcm = dcm_parse(open(name));

    if args["show"]
        for el in dcm
            if args["--tag"] == nothing || el.tag == args["--tag"]
                print(el.tag, ": ")
                display(el)
                println()
            end
        end
    elseif args["dump"]
        # Dumps the image into a csv file.
        if (el = DICOM.lookup(dcm, (0x7fe0,0x0010))) != false
            writecsv("$name.csv", el.data[1][:,:,1])
            # TODO: Possibly also dump these. at least look into them!
            tags = Dict(
                (0x0008,0x0008) => "imtype",  # Image Type
                (0x0008,0x0060) => "modality",  # Modality: ASCIIString["MR"]
                (0x0008,0x0070) => "manuf",  # Manufacturer: ASCIIString["SIEMENS"]
                (0x0008,0x103e) => "series",  # Series Description: ASCIIString["2ch"]
                (0x0008,0x1090) => "model",  # Manufacturer's Model Name: ASCIIString["Aera"]
                (0x0018,0x0022) => "opts",  # Scan Options: ASCIIString["CT"]
                (0x0018,0x0023) => "aqtype",  # MR Acquisition Type: ASCIIString["2D"]
                (0x0018,0x0050) => "slice_thick",  # Slice Thickness: [8.0]
                (0x0018,0x1314) => "flip_angle",  # Flip Angle: [56.0]
                (0x0018,0x5100) => "pos",  # Patient Position: ASCIIString["HFS"]
                (0x0020,0x0032) => "imgpos",  # Image Position (Patient): [-87.771005468217,93.799321675558,195.08339402037]
                (0x0020,0x0037) => "imgori",  # Image Orientation (Patient): [0.8127451738409,-0.5549969637339,0.1772671787004,0.2130993502411,3.956532e-8,-0.9770305353093]
                (0x0020,0x1041) => "sliceloc",  # Slice Location: [53.505870937259]
                (0x0028,0x0004) => "photometric",  # Photometric Interpretation: ASCIIString["MONOCHROME2"]
                (0x0028,0x0030) => "pixelsp",  # Pixel Spacing: [1.40625,1.40625]
                (0x0028,0x1050) => "win_cent",  # Window Center: [344.0]
                (0x0028,0x1051) => "win_width",  # Window Width: [848.0]
                (0x0028,0x1055) => "win_expl",  # Window Center & Width Explanation: ASCIIString["Algo1"]
                (0x7fe0,0x0010) => "data",  # Variable Pixel Data: 256x192x1 Array{UInt16,3}:
            )
        end
    elseif args["debug"]
        # TODO: Write whatever code you're curious about right now here!

        # Show image size
        #if (el = DICOM.lookup(dcm, (0x7fe0,0x0010))) != false
        #    println(size(el.data[1]))
        #end

        # Check if "max pixel" == max(pixels) or not
        # Yes, it is.
        #maxpix = DICOM.lookup(dcm, (0x0028,0x0107)).data[1]
        #imgdata = DICOM.lookup(dcm, (0x7fe0,0x0010)).data[1]
        #println(maximum(imgdata) == maxpix)
    end
end


# Decision #
# ========================== #
@enum DecisionState NotTaken Taken Known

mutable struct Decision{T}
    state::DecisionState
    value::T
end
state(decision::Decision) = decision.state
decision_value(decision::Decision) = state(decision) == NotTaken ? NaN : decision.value

function Decision(value, ::Type{T}) where T
    return Decision(Taken, T(value))
end

function KnownDecision(value, ::Type{T}) where T
    return Decision(Known, T(value))
end

function Decision(info::JuMP.VariableInfo, ::Type{T}) where T
    if info.has_fix
        return Decision(Taken, T(info.fixed_value))
    end
    return Decision(NotTaken, T(NaN))
end

function KnownDecision(info::JuMP.VariableInfo, ::Type{T}) where T
    value = if info.has_fix
        # Value fixed at construction
        T(info.fixed_value)
    else
        # Value is known, but not yet set. Pick a feasible value for now
        value = if info.has_lb
            T(info.lower_bound)
        elseif info.has_ub
            T(info.upper_bound)
        else
            zero(T)
        end
    end
    return Decision(Known, value)
end

Base.copy(decision::Decision) = Decision(decision.state, decision.value)

# Decisions #
# ========================== #
struct IgnoreDecisions end

struct Decisions
    decisions::OrderedDict{MOI.VariableIndex, Decision{Float64}}

    function Decisions()
        return new(OrderedDict{MOI.VariableIndex, Decision{Float64}}())
    end
end

has_decision(decisions::Decisions, index::MOI.VariableIndex) = haskey(decisions.decisions, index)
decision(decisions::Decisions, index::MOI.VariableIndex) = decisions.decisions[index]
decision_value(decisions::Decisions, index::MOI.VariableIndex) = decision_value(decision(decisions, index))

function set_decision!(decisions::Decisions, index::MOI.VariableIndex, decision::Decision)
    decisions.decisions[index] = decision
    return nothing
end

function remove_decision!(::IgnoreDecisions, ::MOI.VariableIndex)
    return nothing
end

function remove_decision!(decisions::Decisions, index::MOI.VariableIndex)
    if !haskey(decisions.decisions, index)
        return nothing
    end
    delete!(decisions.decisions, index)
    return nothing
end

function clear!(decisions::Decisions)
    empty!(decisions.decisions)
    return nothing
end

function all_decisions(decisions::Decisions)
    return filter(decisions.decisions) do (index, decision)
        state(decision) != Known
    end |> keys |> collect
end

function all_known_decisions(decisions::Decisions)
    return filter(decisions.decisions) do (index, decision)
        state(decision) == Known
    end |> keys |> collect
end

function num_decisions(decisions::Decisions)
    return count(decisions.decisions) do (index, decision)
        return state(decision) != Known
    end
end

function num_known_decisions(decisions::Decisions)
    return count(decisions.decisions) do (index, decision)
        state(decision) == Known
    end
end

function untake_decisions!(decisions::Decisions)
    need_update = false
    # Update states and values
    for decision in all_decisions(decisions)
        if decision.state != NotTaken
            decision.state = NotTaken
            need_update |= true
        end
    end
    return need_update
end

function take_decisions!(decisions::Decisions, x::AbstractVector)
    # Check decision length
    num_decisions(decisions) == length(x) || error("Given decision of length $(length(x)) not compatible with number of defined decision variables $(num_decisions(decisions)).")
    # Update states and values (assume x given in sorted order)
    for (decision, val) in zip(all_decisions(decisions), x)
        decision.state = Taken
        decision.value = val
    end
    return nothing
end

function update_known_decisions!(decisions::Decisions, x::AbstractVector)
    # Check decision length
    num_known_decisions(decisions) == length(x) || error("Given decision of length $(length(x)) not compatible with number of defined known decision variables $(num_known_decisions(decisions)).")
    # Update values (assume x given in sorted order)
    for (key, val) in zip(all_known_decisions(decisions), x)
        decisions.decisions[key].value = val
    end
    return nothing
end

is_decision_type(::DataType) = false

include("variable_interface.jl")
include("expressions/expressions.jl")
include("functions/functions.jl")
include("sets.jl")
include("modifications.jl")
include("bridges/bridges.jl")
include("macros.jl")
include("updates.jl")
include("moi_overrides.jl")


"""
## Extract epoch data from SSR

#### Arguments
* `a`: A SSR object
* `valid_triggers`: Trigger numbers that are considered valid ([1,2])
* `remove_first`: Remove the first n triggers (0)
* `remove_last`: Remove the last n triggers (0)

#### Example

```julia
epochs = extract_epochs(SSR, valid_triggers=[1,2])
```
"""
function extract_epochs(
    a::SSR;
    valid_triggers::Union{AbstractArray,Int} = [1, 2],
    remove_first::Int = 0,
    remove_last::Int = 0,
    kwargs...,
)

    merge!(
        a.processing,
        Dict(
            "epochs" => extract_epochs(
                a.data,
                a.triggers,
                valid_triggers,
                remove_first,
                remove_last,
            ),
        ),
    )

    return a
end

function epoch_rejection(a::SSR; retain_percentage::Number = 0.95, kwargs...)

    a.processing["epochs"] =
        epoch_rejection(a.processing["epochs"], retain_percentage; kwargs...)

    return a
end

function create_sweeps(a::SSR; epochsPerSweep::Int = 64, kwargs...)

    if epochsPerSweep > size(a.processing["epochs"], 2)
        error("Sweep length is longer than number of epochs will allow")
    end

    merge!(
        a.processing,
        Dict("sweeps" => create_sweeps(a.processing["epochs"], epochsPerSweep)),
    )

    return a
end

#######################################
#
# Add triggers for more epochs
#
#######################################

function add_triggers(a::SSR; kwargs...)

    @debug("Adding triggers to reduce SSR. Using SSR modulation frequency")

    add_triggers(a, modulationrate(a); kwargs...)
end

function add_triggers(a::SSR, mod_freq::Number; kwargs...)

    @debug("Adding triggers to reduce SSR. Using $(mod_freq)Hz")

    epochIndex =
        DataFrame(Code = a.triggers["Code"], Index = [Int(i) for i in a.triggers["Index"]])
    epochIndex[:Code] = epochIndex[:Code] .- 252

    add_triggers(a, mod_freq, epochIndex; kwargs...)
end

function add_triggers(
    a::SSR,
    mod_freq::Number,
    epochIndex;
    cycle_per_epoch::Int = 1,
    kwargs...,
)

    @info(
        "Adding triggers to reduce SSR. Reducing $(mod_freq)Hz to $cycle_per_epoch cycle(s)."
    )

    # Existing epochs
    existing_epoch_length = median(diff(epochIndex[:Index]))     # samples
    existing_epoch_length_s = existing_epoch_length / samplingrate(a)
    @debug("Existing epoch length: $(existing_epoch_length_s)s")

    # New epochs
    new_epoch_length_s = cycle_per_epoch / mod_freq
    new_epochs_num = round.(existing_epoch_length_s / new_epoch_length_s) - 2
    new_epoch_times = collect(1:new_epochs_num) * new_epoch_length_s
    new_epoch_indx = [0; round.(new_epoch_times * samplingrate(a))]
    @debug("New epoch length = $new_epoch_length_s")
    @debug("New # epochs     = $new_epochs_num")

    # Place new epoch indices
    @debug("Was $(length(epochIndex[:Index])) indices")
    new_indx = epochIndex[:Index][1:end-1] .+ new_epoch_indx'
    new_indx = reshape(new_indx', length(new_indx), 1)[1:end-1]
    @debug("Now $(length(new_indx)) indices")

    # Place in dict
    new_code = round.(Int, ones(1, length(new_indx))) .+ 252
    a.triggers = Dict(
        "Index" => vec((new_indx)'),
        "Code" => vec(new_code),
        "Duration" => ones(length(new_code), 1)',
    )
    #TODO Possible the trigger duration of one is not long enough

    return a
end

#######################################
#
# Rejection channels
#
#######################################

function channel_rejection(
    a::SSR;
    threshold_abs::Number = 1000,
    threshold_var::Number = 2,
    kwargs...,
)

    # Run on epochs if available, else run on the raw data
    if haskey(a.processing, "epochs")
        data = reshape(
            a.processing["epochs"],
            size(a.processing["epochs"], 1) * size(a.processing["epochs"], 2),
            size(a.processing["epochs"], 3),
        )
    else
        data = a.data
    end

    valid = vec(channel_rejection(data, threshold_abs, threshold_var))

    @info(
        "Rejected $(sum(.!valid)) channels $(join(channelnames(a)[findall(.!valid)], " "))"
    )

    remove_channel!(a, channelnames(a)[.!valid])

    return a
end


module Tinker

using Gtk.ShortNames, GtkReactive, Graphics, Colors, Images, IntervalSets
using Luxor.isinside, Luxor.Point

img_ctxs = Signal([])

abstract type Shape end

# Rectangle structure
struct Rectangle <: Shape
    x::Number
    y::Number
    w::Number
    h::Number
end

Rectangle() = Rectangle(0,0,-1,-1)
Base.isempty(R::Rectangle) = R.w <= 0 || R.h <= 0

mutable struct ImageContext{T}
    id::Int
    image
    canvas::GtkReactive.Canvas
    zr::Signal{ZoomRegion{T}}
    zl::Int # for tracking zoom level
    mouseactions::Dict{String,Any}
    shape::Signal{<:Shape} # Tracks type of selection in the environment
    points::Signal{<:AbstractArray} # Holds points that define shape outline
    rectview::Signal{<:AbstractArray} # Holds bounding box of selection
end

ImageContext() = ImageContext(-1, nothing, canvas(),
Signal(ZoomRegion((1:10, 1:10))), -1, Dict("dummy"=>Signal(false),"dummy2"=>Signal(4)),
Signal(Rectangle()), Signal([]), Signal([]))

# Returns a view of an image with the given bounding region
function get_view(image,x_min,y_min,x_max,y_max)
    xleft,yleft = Int(floor(Float64(x_min))),Int(ceil(Float64(y_min)))
    xright,yright = Int(floor(Float64(x_max))),Int(ceil(Float64(y_max)))
    min_x, max_x = map(x->x, extrema(indices(image,2)))
    min_y, max_y = map(y->y, extrema(indices(image,1)))
    (xleft < min_x) && (xleft = min_x)
    (yleft < min_y) && (yleft = min_y)
    (xright > max_x) && (xright = max_x)
    (yright > max_y) && (yright = max_y)
    return view(image, yleft:yright, xleft:xright)
end

# Calculates tolerance based on zoom level
function get_tolerance(zr::ZoomRegion)
  tol = 5*(IntervalSets.width(zr.currentview.x)/IntervalSets.width(zr.fullview.x))
  return tol
end

function get_tolerance(ctx::ImageContext)
  zr = value(ctx.zr)
  get_tolerance(zr)
end

# Creates a Rectangle out of any two points
function Rectangle(p1::XY,p2::XY)
    x, w = min(p1.x, p2.x), abs(p2.x - p1.x)
    y, h = min(p1.y, p2.y), abs(p2.y - p1.y)
    return Rectangle(x, y, w, h)
    (p1.x == p2.x) || (p1.y == p2.y) && return Rectangle()
end

# rectangle draw function
function drawrect(ctx, rect, color, width)
    set_source(ctx, color)
    set_line_width(ctx, width)
    rectangle(ctx, rect.x, rect.y, rect.w, rect.h)
    stroke(ctx)
end;

# Polygon shape
struct Polygon <: Shape
    pts::AbstractArray
end

Polygon() = Polygon([])
Base.isempty(P::Polygon) = isempty(P.pts)

# Handles modify rectangles
struct Handle <: Shape
    r::Shape
    pos::String # refers to which side or corner of rectangle handle is on
    x::Float64
    y::Float64
end

Handle() = Handle(Rectangle(),"",0,0)
Base.isempty(H::Handle) = isempty(H.r)

# Creates handle given a Rectangle and a position
function Handle(r::Rectangle, pos::String)
    # Position of handle refers to center coordinate of handle based on rect
    position_coord = Dict("tlc"=>(r.x,r.y),"ts"=>(r.x+(r.w/2),r.y),
                          "trc"=>(r.x+r.w,r.y),"rs"=>(r.x+r.w,r.y+(r.h/2)),
                          "brc"=>(r.x+r.w,r.y+r.h),"bs"=>(r.x+(r.w/2),r.y+r.h),
                          "blc"=>(r.x,r.y+r.h),"ls"=>(r.x,r.y+(r.h/2)))
    xy = get(position_coord, pos, (-Inf,-Inf))
    if xy == (-Inf,-Inf)
        println("Not a valid Handle position.")
        return Handle()
    else
        x = position_coord[pos][1]
        y = position_coord[pos][2]
        return Handle(r,pos,xy[1],xy[2])
    end
end

function Handle(p::Polygon, i::Int)
    pt = p.pts[i]
    return Handle(p, "", pt.x,pt.y)
end

# Draws a handle
function drawhandle(ctx, handle::Handle, d)
  if !isempty(handle)
    rectangle(ctx, handle.x-(d), handle.y-(d),
              d*2, d*2)
    set_source(ctx,colorant"white")
    fill_preserve(ctx)
    set_source(ctx,colorant"black")
    set_line_width(ctx,1.0)
    stroke_preserve(ctx)
  end
end; # like drawrect, but makes x,y refer to center of handle

# A rectangle with handles at all 8 positions
struct RectHandle <: Shape
    r::Rectangle
    h::NTuple{8,Handle}
end

RectHandle() = RectHandle(Rectangle())
Base.isempty(RH::RectHandle) = isempty(RH.r)

# Creates a RectHandle given just a Rectangle
function RectHandle(r::Rectangle)
    # derive all 8 handles from r
    # numbered 1-8, 1=tlc, moving clockwise around rectangle
    h = (Handle(r, "tlc"), Handle(r, "ts"), Handle(r, "trc"), Handle(r, "rs"),
         Handle(r, "brc"), Handle(r, "bs"), Handle(r, "blc"), Handle(r, "ls"))
    return RectHandle(r,h)
end

# A polygon with handles at every vertex
struct PolyHandle <: Shape
    p::Polygon
    h::AbstractArray
end

PolyHandle() = PolyHandle(Polygon([]),[])

function PolyHandle(p::Polygon)
    # makes polyhandle
    h = [Handle(p,1)]
    for i in 2:length(p.pts)
        push!(h,Handle(p,i))
    end
    return PolyHandle(p,h)
end

function PolyHandle(p::AbstractArray)
  return PolyHandle(Polygon(p))
end

# Draws RectHandle
function drawrecthandle(ctx, rh::RectHandle, d, color1, width)
    drawrect(ctx, rh.r, color1, width)
    for n in 1:length(rh.h)
        drawhandle(ctx, rh.h[n], d)
    end
end

# draws PolyHandle
function drawpolyhandle(ctx, ph::PolyHandle, d, color, width)
  drawline(ctx, ph.p.pts, color, width)
  for i in 1:length(ph.h)
    drawhandle(ctx, ph.h[i], d)
  end
end

# Connects an array of points
function drawline(ctx, l, color, width)
    isempty(l) && return
    p = first(l)
    move_to(ctx, p.x, p.y)
    set_source(ctx, color)
    set_line_width(ctx, width)
    for i = 2:length(l)
        p = l[i]
        line_to(ctx, p.x, p.y)
    end
    stroke(ctx)
end

# Set of versatile methods to draw many shapes with the same function
function drawshape(ctx, sh::RectHandle, d, color, width)
  # draw RectHandle
  drawrecthandle(ctx, sh, d, color, width)
end

function drawshape(ctx, sh::Polygon, d, color, width)
  # draw Polygon
  drawline(ctx, sh.pts, color, width)
end

function drawshape(ctx, sh::PolyHandle, d, color, width)
  # draw PolyHandle
  drawpolyhandle(ctx, sh, d, color, width)
end

# Checks if an array of points qualifies as a polygon
function ispolygon(p::AbstractVector)
    length(p) < 4 && return false
    p[1] != p[end] && return false
    for i in 1:(length(p)-1)
        p[i]!=p[i+1] && return true
    end
    return false
end

# Moves polygon to a given location
function move_polygon_to(p::AbstractArray, pt::XY)
    # find diff b/t start & pt; add diff to all in p
    if ispolygon(p)
        diff = XY(pt.x-p[1].x,pt.y-p[1].y)
        map(n -> XY(n.x+diff.x,n.y+diff.y),p)
    else
        #println("Not a polygon")
        return p
    end
end

# Converts an XY to a Point
function Point(p::XY)
    Point(p.x,p.y)
end

include("zoom_interaction.jl")
include("selection_actions.jl")

## Sets up an image in a separate window with the ability to adjust view
function init_image(image::AbstractArray; name="Tinker")
    # set up window
    win = Window(name, size(image,2), size(image,1))
    c = canvas(UserUnit)
    push!(win, c)

    # set up a zoom region
    zr = Signal(ZoomRegion(image))

    # create view
    imagesig = map(zr) do r
        cv = r.currentview
        view(image, UnitRange{Int}(cv.y), UnitRange{Int}(cv.x))
    end;

    # create a view diagram
    viewdim = map(zr) do r
        fvx, fvy = r.fullview.x, r.fullview.y # x, y range of full view
        cvx, cvy = r.currentview.x, r.currentview.y # x, y range of currentview
        xfull, yfull =
            (fvx.right-fvx.left),(fvy.right-fvy.left) # width of full view
        xcurrent, ycurrent =
            (cvx.right-cvx.left),(cvy.right-cvy.left) # width of current view
        # scale
        xsc,ysc = 0.1*(xcurrent/xfull), 0.1*(ycurrent/yfull)
        # offset
        x_off,y_off = cvx.left+(0.01*xcurrent),cvy.left+(0.01*ycurrent)
        # represents full view
        rect1 = Rectangle(x_off, y_off, xsc*xfull, ysc*yfull)
        # represents current view
        rect2 = Rectangle(x_off+(cvx.left*xsc), y_off+(cvy.left*ysc),
                          xsc*xcurrent, ysc*ycurrent)
        return [rect1,rect2]
    end

    # Placeholder dictionary for context
    dummydict = Dict("pandrag"=>Signal(false),"zoomclick"=>Signal(false),
    "select"=>Signal(false),"selmode"=>Signal(rectangle_mode))

    # Context
    imagectx=ImageContext(length(value(img_ctxs))+1,image,c,zr,1,dummydict,
    Signal(Shape, Rectangle()),Signal([]),Signal(view(image,1:size(image,2),
    1:size(image,1))))

    # Mouse actions
    pandrag = init_pan_drag(c, zr)
    zoomclick = init_zoom_click(imagectx)
    select = init_selection_actions(imagectx)

    push!(pandrag["enabled"],false)
    push!(zoomclick["enabled"],false)

    imagectx.mouseactions = Dict("pandrag"=>pandrag["enabled"],"zoomclick"=>
    zoomclick["enabled"],"select"=>select["enabled"],"selmode"=>select["mode"])

    imagectx.rectview = map(imagectx.points) do pts
        if !isempty(pts)
            x_min,x_max = minimum(map(n->n.x,pts)),maximum(map(n->n.x,pts))
            y_min,y_max =minimum(map(n->n.y,pts)),maximum(map(n->n.y,pts))
            get_view(image,x_min,y_min,x_max,y_max)
        else
            view(image,1:size(image,2),1:size(image,1))
        end
    end

    append!(c.preserved, [pandrag,zoomclick,select])

    # draw
    redraw = draw(c, imagesig, zr, viewdim, imagectx.points, imagectx.shape) do cnvs,img,r,vd,pts,sh
        copy!(cnvs, img) # show image on canvas at current zoom level
        set_coordinates(cnvs, r) # set canvas coordinates to zr
        ctx = getgc(cnvs)
        # draw view diagram if zoomed in
        if r.fullview != r.currentview
            drawrect(ctx, vd[1], colorant"blue", 2.0)
            drawrect(ctx, vd[2], colorant"blue", 2.0)
        end
        d = get_tolerance(imagectx)
        if ispolygon(pts)
          # draw shape
          if typeof(sh) == RectHandle
              recthandle = RectHandle(Rectangle(XY(pts[1].x,pts[1].y),XY(pts[3].x,pts[3].y)))
              drawrecthandle(ctx, recthandle, d, colorant"yellow", 1.0)
          elseif typeof(sh) == PolyHandle
              drawpolyhandle(ctx, PolyHandle(pts), d, colorant"yellow", 1.0)
          else
              drawline(ctx, pts, colorant"yellow", 1.0)
          end
        else
          # draw working line
          drawline(ctx, pts, colorant"yellow", 1.0)
          if typeof(sh) == PolyHandle && length(pts) > 0
              drawrect(ctx, Rectangle(pts[1].x-d,pts[1].y-d,2*d,2*d), colorant"yellow", 1.0)
          end
        end
    end

    showall(win);

    # shows measurements
    display_measure(imagectx)

    push!(img_ctxs, push!(value(img_ctxs), imagectx))
    return imagectx
end;

init_image(file::AbstractString) = init_image(load(file); name=file)

@enum Mode zoom_mode rectangle_mode freehand_mode polygon_mode

function set_mode(ctx::ImageContext, mode::Mode)
    push!(ctx.mouseactions["pandrag"], false)
    push!(ctx.mouseactions["zoomclick"], false)
    push!(ctx.mouseactions["select"],false)
    if mode == zoom_mode # turn on zoom controls
        println("Zoom mode")
        push!(ctx.mouseactions["pandrag"], true)
        push!(ctx.mouseactions["zoomclick"], true)
        push!(ctx.mouseactions["select"],false)
    elseif mode == rectangle_mode # turn on rectangular region selection controls
        println("Rectangle mode")
        push!(ctx.mouseactions["select"],true)
        push!(ctx.mouseactions["selmode"],rectangle_mode)
    elseif mode == freehand_mode # freehand select
        println("Freehand mode")
        push!(ctx.mouseactions["select"],true)
        push!(ctx.mouseactions["selmode"],freehand_mode)
    elseif mode == polygon_mode # polygon select
        println("Polygon mode")
        push!(ctx.mouseactions["select"],true)
        push!(ctx.mouseactions["selmode"],polygon_mode)
    end
end

set_mode(sig::Signal, mode) = set_mode(value(sig), mode)

# Sets all contexts
function set_mode_all(mode::Mode)
  for c in value(img_ctxs)
    set_mode(c, mode)
  end
  nothing
end

include("measure.jl")
include("guisetup.jl")

end # module


using MPIParticleInCell
using Test
using MPI

@testset "MPIParticleInCell.jl" begin
    MPI.Init()

    comm = MPI.COMM_WORLD
    rank = MPI.Comm_rank(comm)
    size = MPI.Comm_size(comm)

    num_guard_cells = 1
    g = UniformGrid((0., 0.), (1., 1.), 2 .* (8, 5), num_guard_cells)
    rank == 0 && println(g)

    f = Field(g, "board", Int, true)
    rank == 0 && println(f)

    f.field_values .= rank
    gather_field_on_root!(f)
    rank == 0 && (display(f.field_values); println("\n"))

    f.field_values .= rank
    communicate_halo!(f, comm)
    rank == 0 && (display(f.field_values); println("\n"))

    MPI.Finalize()

    @test true
end


module SMC

using Compat
# using ExpFamily
using Distributions

import Base.length, Base.mean

const Int   = Int64
const Float = Float64

include("hmm.jl")
include("kalman.jl")

include("particles.jl")
include("resample.jl")
include("proposal.jl")
include("particlefilter.jl")
include("particlesmoother.jl")

end # module


function get_ant(;
    timestep=0.05,
    gravity=[0.0; 0.0; -9.81],
    friction_coefficient=0.5,
    spring=0.0,
    damper=1.0,
    contact_feet=true,
    contact_body=true,
    limits::Bool=true,
    joint_limits=[[-30.0, 30.0, -30.0, -70.0, -30.0, -70.0, -30.0, 30.0] * π / 180.0,
                  [ 30.0, 70.0,  30.0, -30.0,  30.0, -30.0,  30.0, 70.0] * π / 180.0],
    T=Float64)

    path = joinpath(@__DIR__, "../deps/ant.urdf")
    mech = Mechanism(path, true, T,
        gravity=gravity,
        timestep=timestep,
        spring=spring,
        damper=damper)

    # joint limits
    joints = deepcopy(mech.joints)

    if limits
        hiparent_vertex = get_joint(mech, :hip_1)
        joints[hiparent_vertex.id] = add_limits(mech, hiparent_vertex,
            rot_limits=[SVector{1}(joint_limits[1][1]), SVector{1}(joint_limits[2][1])])

        ankle1 = get_joint(mech, :ankle_1)
        joints[ankle1.id] = add_limits(mech, ankle1,
            rot_limits=[SVector{1}(joint_limits[1][2]), SVector{1}(joint_limits[2][2])])

        hip2 = get_joint(mech, :hip_2)
        joints[hip2.id] = add_limits(mech, hip2,
            rot_limits=[SVector{1}(joint_limits[1][3]), SVector{1}(joint_limits[2][3])])

        ankle2 = get_joint(mech, :ankle_2)
        joints[ankle2.id] = add_limits(mech, ankle2,
            rot_limits=[SVector{1}(joint_limits[1][4]), SVector{1}(joint_limits[2][4])])

        hip3 = get_joint(mech, :hip_3)
        joints[hip3.id] = add_limits(mech, hip3,
            rot_limits=[SVector{1}(joint_limits[1][5]), SVector{1}(joint_limits[2][5])])

        ankle3 = get_joint(mech, :ankle_3)
        joints[ankle3.id] = add_limits(mech, ankle3,
            rot_limits=[SVector{1}(joint_limits[1][6]), SVector{1}(joint_limits[2][6])])

        hip4 = get_joint(mech, :hip_4)
        joints[hip4.id] = add_limits(mech, hip4,
            rot_limits=[SVector{1}(joint_limits[1][7]), SVector{1}(joint_limits[2][7])])

        ankle4 = get_joint(mech, :ankle_4)
        joints[ankle4.id] = add_limits(mech, ankle4,
            rot_limits=[SVector{1}(joint_limits[1][8]), SVector{1}(joint_limits[2][8])])

        mech = Mechanism(Origin{T}(), mech.bodies, [joints...],
            gravity=gravity,
            timestep=timestep,
            spring=spring,
            damper=damper)
    end

    if contact_feet
        origin = Origin{T}()
        bodies = mech.bodies
        joints = mech.joints

        # foot contact
        normal = [0.0; 0.0; 1.0]
        foot_names = [:front_left_foot, :front_right_foot, :left_back_foot, :right_back_foot]
        foot = [get_body(mech, name) for name in foot_names]
        p = [[0.2; 0.2; 0.0], [-0.2; 0.2; 0.0], [-0.2; -0.2; 0.0], [0.2; -0.2; 0.0]]
        o = [[0.0; 0.0; f.shape.rh[1]] for f in foot]
        contacts = [contact_constraint(foot[i], normal, friction_coefficient=friction_coefficient, contact_point=p[i], offset=o[i]) for i = 1:length(foot_names)]

        if contact_body
            # torso contact
            torso = get_body(mech, :torso)
            p = [0.0; 0.0; 0.0]
            o = [0.0; 0.0; torso.shape.r]
            torso_contacts = contact_constraint(torso, normal, friction_coefficient=friction_coefficient, contact_point=p, offset=o)

            # elbow contact
            elbow_names = [:aux_1, :aux_2, :aux_3, :aux_4]
            elbow = [get_body(mech, e) for e in elbow_names]
            p = [-[0.1; 0.1; 0.0], -[-0.1; 0.1; 0.0], -[-0.1;-0.1; 0.0], -[0.1; -0.1; 0.0]]
            o = [[0.0; 0.0; e.shape.rh[1]] for e in elbow]
            elbow_contacts = [contact_constraint(elbow[i], normal, friction_coefficient=friction_coefficient, contact_point=p[i], offset=o[i]) for i = 1:length(elbow_names)]

            contacts = [contacts..., torso_contacts, elbow_contacts...]
        end
        mech = Mechanism(origin, bodies, joints, contacts,
        gravity=gravity,
        timestep=timestep,
        spring=spring,
        damper=damper)
    end

    return mech
end

function initialize_ant!(mechanism::Mechanism;
        body_position=[0.0; 0.0; 0.63],
        body_orientation=[0.0; 0.0; 0.0 * π],
        ankle_orientation=0.25)

    set_minimal_coordinates!(mechanism, get_joint(mechanism, :floating_base), [body_position; body_orientation])

    for i in [1, 4]
        set_minimal_coordinates!(mechanism, get_joint(mechanism, Symbol("hip_$i")), [0.0 * π])
        set_minimal_coordinates!(mechanism, get_joint(mechanism, Symbol("ankle_$i")), [ankle_orientation * π])
    end

    for i in [2, 3]
        set_minimal_coordinates!(mechanism, get_joint(mechanism, Symbol("hip_$i")), [0.0 * π])
        set_minimal_coordinates!(mechanism, get_joint(mechanism, Symbol("ankle_$i")), [-ankle_orientation * π])
    end

    zero_velocity!(mechanism)
end


using WebAPI
using HTTP
using JSON3
using Sockets
using Test

# for server tests
function bhaskara(a, b, c)
    Δ = b^2 - 4*a*c

    Δ < 0 && (Δ = complex(Δ))

    x₁ = (-b + √Δ) / 2a
    x₂ = (-b - √Δ) / 2a
    return x₁, x₂
end

@testset "WebAPI.jl" begin
    include("params_query.jl")
    include("req_body.jl")
    include("res_constructors.jl")
    include("add_route.jl")
    include("utility_functions.jl")
    include("prints.jl")
    include("server.jl")
end


using PyPlot;
using MAT;
using jInv.Mesh
using jInvVis
file = matopen("hand.mat")
B = read(file, "B");
close(file);
plotModel(convert(Array{Float64,3},B));


using Pkg.TOML, StaticArrays

@testset "ExcitationFields.jl" begin
    params = TOML.parsefile("ExcitationFields.toml")

    Hs = MPISimulations.initialize(MPISimulations.StaticGradientField,"StaticGradientField",params,Float64)
    H1 = MPISimulations.initialize(MPISimulations.SinusoidalField,"SinusoidalField1",params,Float64)
    H2 = MPISimulations.initialize(MPISimulations.SinusoidalField,"SinusoidalField2",params,Float64)
    H3 = MPISimulations.initialize(MPISimulations.SinusoidalField,"SinusoidalField3",params,Float64)
    r = randn(SVector{3,Float64})
    t = randn(Float64)

    excitation = MPISimulations.ExcitationField(Hs,H1)
    @test MPISimulations.magneticFieldStrength(excitation,r,t) ≈ MPISimulations.magneticFieldStrength(Hs,r,t) + MPISimulations.magneticFieldStrength(H1,r,t)

    excitation = MPISimulations.ExcitationField(Hs,H1,H2)
    @test MPISimulations.magneticFieldStrength(excitation,r,t) ≈ MPISimulations.magneticFieldStrength(Hs,r,t) + MPISimulations.magneticFieldStrength(H1,r,t) + MPISimulations.magneticFieldStrength(H2,r,t)

    excitation = MPISimulations.ExcitationField(Hs,H1,H2,H3)
    @test MPISimulations.magneticFieldStrength(excitation,r,t) ≈ MPISimulations.magneticFieldStrength(Hs,r,t) + MPISimulations.magneticFieldStrength(H1,r,t) + MPISimulations.magneticFieldStrength(H2,r,t) + MPISimulations.magneticFieldStrength(H3,r,t)
end


module GeoInterface

using Extents
using Base.Iterators: flatten

export testgeometry, isgeometry, trait, geomtrait, ncoord, getcoord, ngeom, getgeom

# traits
export AbstractGeometryTrait,
    AbstractGeometryCollectionTrait,
    GeometryCollectionTrait,
    AbstractPointTrait,
    PointTrait,
    AbstractCurveTrait,
    AbstractLineStringTrait,
    LineStringTrait,
    LineTrait,
    LinearRingTrait,
    CircularStringTrait,
    CompoundCurveTrait,
    AbstractSurfaceTrait,
    AbstractCurvePolygonTrait,
    CurvePolygonTrait,
    AbstractPolygonTrait,
    PolygonTrait,
    TriangleTrait,
    RectangleTrait,
    QuadTrait,
    PentagonTrait,
    HexagonTrait,
    AbstractPolyhedralSurfaceTrait,
    PolyhedralSurfaceTrait,
    TINTrait,
    AbstractMultiPointTrait,
    MultiPointTrait,
    AbstractMultiCurveTrait,
    MultiCurveTrait,
    AbstractMultiLineStringTrait,
    MultiLineStringTrait,
    AbstractMultiSurfaceTrait,
    MultiSurfaceTrait,
    AbstractMultiPolygonTrait,
    MultiPolygonTrait,
    AbstractFeatureTrait,
    FeatureTrait,
    AbstractCollectionFeatureTrait,
    FeatureCollectionTrait

include("types.jl")
include("interface.jl")
include("fallbacks.jl")
include("utils.jl")
include("base.jl")

end # module


module PkgMPC

# Write your package code here.
using Reexport
@reexport using Plots
@reexport using CSV
@reexport using JuMP
@reexport using Ipopt
@reexport using DataFrames
@reexport using DifferentialEquations
@reexport using PlotlyJS

include("Extra.jl")
include("main.jl")
include("NMPC.jl")
include("Plant.jl")

export F, Model_Par, Simulation_Data, Simulate_Plant
export Simulate_Plant, Plot_Plant, Save_Plant, Export_Plant

end


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# Level-set
# ------------------------------------------------------------
const LVConstraint = CI{AffineDecisionFunction{Float64}, MOI.LessThan{Float64}}
@with_kw mutable struct LVData{T <: AbstractFloat}
    Q̃::T = 1e10
    objective::AffineDecisionFunction{T} = zero(AffineDecisionFunction{T})
    constraint::LVConstraint = LVConstraint(0)
    objective_cached::Bool = false
    incumbent::Int = 1
    major_iterations::Int = 0
    minor_iterations::Int = 0
    levelindex::Int = -1
end

@with_kw mutable struct LVParameters{T <: AbstractFloat}
    λ::T = 0.5
end

"""
    LevelSet

Functor object for using level-set regularization in an L-shaped algorithm. Create by supplying an [`LV`](@ref) object through `regularize` in `LShaped.Optimizer` or by setting the [`Regularizer`](@ref) attribute.

...
# Parameters
- `λ::AbstractFloat = 0.5`: Controls the level position L = (1-λ)*θ + λ*Q̃, a convex combination of the current lower and upper bound.
- `penaltyterm::PenaltyTerm = Quadratic`: Specify penaltyterm variant ([`Quadratic`](@ref), [`InfNorm`](@ref), [`ManhattanNorm`][@ref])
...
"""
struct LevelSet{T <: AbstractFloat, A <: AbstractVector, PT <: AbstractPenaltyTerm} <: AbstractRegularization
    data::LVData{T}
    parameters::LVParameters{T}

    decisions::DecisionMap
    projection_targets::Vector{MOI.VariableIndex}
    ξ::Vector{Decision{T}}

    Q̃_history::A
    levels::A
    incumbents::Vector{Int}

    penaltyterm::PT

    function LevelSet(decisions::DecisionMap, ξ₀::AbstractVector, penaltyterm::AbstractPenaltyTerm; kw...)
        T = promote_type(eltype(ξ₀), Float32)
        A = Vector{T}
        ξ = map(ξ₀) do val
            KnownDecision(val, T)
        end
        PT = typeof(penaltyterm)
        return new{T, A, PT}(LVData{T}(),
                             LVParameters{T}(;kw...),
                             decisions,
                             Vector{MOI.VariableIndex}(undef, length(ξ₀)),
                             ξ,
                             A(),
                             A(),
                             Vector{Int}(),
                             penaltyterm)
    end
end

function initialize_regularization!(lshaped::AbstractLShaped, lv::LevelSet)
    # Add projection targets
    add_projection_targets!(lv, lshaped.master)
    # Sense-correct initial objective
    sense = MOI.get(lshaped.master, MOI.ObjectiveSense())
    coeff = sense == MOI.MIN_SENSE ? 1.0 : -1.0
    lv.data.Q̃ *= coeff
    # Run penalty initialization to check support
    initialize_penaltyterm!(lv.penaltyterm,
                            lshaped.master,
                            1.0,
                            all_decisions(lv.decisions),
                            lv.projection_targets)
    # Delete penalty-term
    remove_penalty!(lv.penaltyterm, lshaped.master)
    return nothing
end

function restore_regularized_master!(lshaped::AbstractLShaped, lv::LevelSet)
    # Delete level-set constraint
    if !iszero(lv.data.constraint.value)
        MOI.delete(lshaped.master, lv.data.constraint)
        lv.data.constraint = LVConstraint(0)
    end
    # Delete penalty-term
    remove_penalty!(lv.penaltyterm, lshaped.master)
    # Delete projection targets
    for var in lv.projection_targets
        MOI.delete(lshaped.master, var)
        StochasticPrograms.remove_decision!(lv.decisions, var)
    end
    empty!(lv.projection_targets)
    return nothing
end

function filter_variables!(lv::LevelSet, list::Vector{MOI.VariableIndex})
    # Filter projection targets
    filter!(vi -> !(vi in lv.projection_targets), list)
    # Filter any auxiliary penaltyterm variables
    remove_penalty_variables!(lv.penaltyterm, list)
    return nothing
end

function filter_constraints!(lv::LevelSet, list::Vector{<:CI})
    # Filter any auxiliary penaltyterm constraints
    remove_penalty_constraints!(lv.penaltyterm, list)
    # Filter level-set constraint
    i = something(findfirst(isequal(lv.data.constraint), list), 0)
    if !iszero(i)
        MOI.deleteat!(list, i)
    end
    return nothing
end

function log_regularization!(lshaped::AbstractLShaped, lv::LevelSet)
    @unpack Q̃,incumbent = lv.data
    push!(lv.Q̃_history, Q̃)
    push!(lv.incumbents, incumbent)
    return nothing
end

function log_regularization!(lshaped::AbstractLShaped, t::Integer, lv::LevelSet)
    @unpack Q̃,incumbent = lv.data
    lv.Q̃_history[t] = Q̃
    lv.incumbents[t] = incumbent
    return nothing
end

function take_step!(lshaped::AbstractLShaped, lv::LevelSet)
    @unpack Q = lshaped.data
    @unpack τ = lshaped.parameters
    @unpack Q̃ = lv.data
    sense = MOI.get(lshaped.master, MOI.ObjectiveSense())
    coeff = sense == MOI.MIN_SENSE ? 1.0 : -1.0
    if coeff*Q <= coeff*Q̃ - τ
        x = current_decision(lshaped)
        for i in eachindex(lv.ξ)
            lv.ξ[i].value = x[i]
        end
        lv.data.Q̃ = Q
        lv.data.incumbent = timestamp(lshaped)
        lv.data.major_iterations += 1
    else
        lv.data.minor_iterations += 1
    end
    if count(active_model_objectives(lshaped)) != num_thetas(lshaped)
        # Only project if all master variables have been added
        return nothing
    end
    if !lv.data.objective_cached
        F = MOI.get(lshaped.master, MOI.ObjectiveFunctionType())
        lv.data.objective = MOI.get(lshaped.master, MOI.ObjectiveFunction{F}())
        lv.data.objective_cached = true
    end
    @unpack Q̃, objective = lv.data
    @unpack λ = lv.parameters
    # Calculate new θ
    # ---------------- #
    # Delete level-set constraint
    if !iszero(lv.data.constraint.value)
        MOI.delete(lshaped.master, lv.data.constraint)
        lv.data.constraint = LVConstraint(0)
    end
    # Delete penalty-term
    remove_penalty!(lv.penaltyterm, lshaped.master)
    # Re-add objective
    F = typeof(objective)
    MOI.set(lshaped.master, MOI.ObjectiveFunction{F}(), objective)
    # Solve unregularized master
    status = solve_master!(lshaped)
    if !(status ∈ AcceptableTermination)
        # Early termination
        return nothing
    end
    # Update model objective
    θ = MOI.get(lshaped.master, MOI.ObjectiveValue())
    lshaped.data.θ = θ
    # Calculate new level
    L = (1-λ)*θ + λ*Q̃
    push!(lv.levels, L)
    # Reformulate projection problem
    F = MOI.ScalarAffineFunction{Float64}
    MOI.set(lshaped.master, MOI.ObjectiveFunction{F}(), zero(F))
    initialize_penaltyterm!(lv.penaltyterm,
                            lshaped.master,
                            1.0,
                            all_decisions(lv.decisions),
                            lv.projection_targets)
    # Add level constraint
    sense = MOI.get(lshaped.master, MOI.ObjectiveSense())
    coeff = sense == MOI.MIN_SENSE ? 1.0 : -1.0
    lv.data.constraint =
        MOI.add_constraint(lshaped.master, coeff*objective,
                           MOI.LessThan(coeff*L))
    return nothing
end

function process_cut!(lshaped::AbstractLShaped, cut::AnyOptimalityCut, lv::LevelSet)
    if lshaped.execution isa AsynchronousExecution
        # Shift level by cut value to prevent infeasibility
        increment = cut(lshaped.x)
        if !iszero(lv.data.constraint.value)
            L = MOI.get(lshaped.master, MOI.ConstraintSet(), lv.data.constraint)
            MOI.set(lshaped.master, MOI.ConstraintSet(), lv.data.constraint, MOIU.shift_constant(L, -increment))
        end
    end
    return nothing
end

function str(::LevelSet)
    return "L-shaped using level sets"
end

# API
# ------------------------------------------------------------
"""
    LV

Factory object for [`LevelSet`](@ref). Pass to `regularize` in `LShaped.Optimizer` or set the [`Regularizer`](@ref) attribute. Equivalent factory calls: `LV`, `WithLV`, `LevelSet`, `WithLevelSets`. See ?LevelSet for parameter descriptions.

"""
mutable struct LV <: AbstractRegularizer
    penaltyterm::AbstractPenaltyTerm
    parameters::LVParameters{Float64}
end
LV(; penaltyterm = Quadratic(), kw...) = LV(penaltyterm, LVParameters(; kw...))
WithLV(; penaltyterm = Quadratic(), kw...) = LV(penaltyterm, LVParameters(; kw...))
LevelSet(; penaltyterm = Quadratic(), kw...) = LV(penaltyterm, LVParameters(; kw...))
WithLevelSets(; penaltyterm = Quadratic(), kw...) = LV(penaltyterm, LVParameters(; kw...))

function MOI.get(lv::LV, ::RegularizationPenaltyTerm)
    return lv.penaltyterm
end

function MOI.set(lv::LV, ::RegularizationPenaltyTerm, penaltyterm::AbstractPenaltyTerm)
    return lv.penaltyterm = penaltyterm
end

function (lv::LV)(decisions::DecisionMap, x::AbstractVector)
    return LevelSet(decisions, x, lv.penaltyterm; type2dict(lv.parameters)...)
end

function str(::LV)
    return "L-shaped using level sets"
end


#############################################################################
#Copyright (c) 2016-2019 Charles Le Losq
#
#The MIT License (MIT)
#
#Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the #Software without restriction, including without limitation the rights to use, copy, #modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, #and to permit persons to whom the Software is furnished to do so, subject to the #following conditions:
#
#The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, #INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
#############################################################################

VERSION >= v"1.0.0" && __precompile__()

module Spectra

using StatsBase
using Statistics
using Random
using Statistics
using SparseArrays
using LinearAlgebra
using PyPlot
using LsqFit
using PyCall
using Dierckx
#using NMF
using Polynomials

# For PyCall modules
const preprocessing = PyNULL()
const model_selection = PyNULL()
const decomposition = PyNULL()
const kernel_ridge = PyNULL()
const svm = PyNULL()
const linear_model = PyNULL()
const gaussian_process = PyNULL()
const pygcvspl = PyNULL()
const signal = PyNULL()
const rampy = PyNULL()

function __init__()

	copy!(preprocessing, pyimport_conda("sklearn.preprocessing", "scikit-learn"))
	copy!(model_selection, pyimport_conda("sklearn.model_selection", "scikit-learn"))
	copy!(decomposition, pyimport_conda("sklearn.decomposition", "scikit-learn"))
	copy!(kernel_ridge, pyimport_conda("sklearn.kernel_ridge", "scikit-learn"))
	copy!(svm, pyimport_conda("sklearn.svm","scikit-learn"))
	copy!(gaussian_process, pyimport_conda("sklearn.gaussian_process", "scikit-learn"))
	copy!(linear_model, pyimport_conda("sklearn.linear_model", "scikit-learn"))
	copy!(pygcvspl, pyimport("gcvspline"))
	copy!(rampy, pyimport("rampy"))
	copy!(signal, pyimport_conda("scipy.signal", "scipy"))

end

include("integrale.jl")
include("functions.jl")
include("baseline.jl")
include("bootstrap.jl")
include("tlcorrection.jl")
include("deprecated.jl")
include("ml_regressor.jl")
include("peakmeasurement.jl")

#From integrale.jl
export trapz, bandarea

#From functions.jl
export poly, polyfit, gaussiennes, lorentziennes, pseudovoigts, pearson7, normal_dist, xshift_inversion, xshift_direct,xshift_correction, smooth, flipsp, resample

#From baseline.jl
export baseline

#From bootstrap
export bootsample, bootperf

#From tlcorrection
export tlcorrection

#From ml_regressor
export mlregressor

# From peakmeasurement
export peakmeas

end # module


using Requests
import Requests: get, post, put, delete, options
using CSV
using DataFrames
using WriteVTK

query = Dict{String, String}(
	"batch" => "l",
	"COMMAND" => "499",
	"CENTER" => "500@10",
	"MAKE_EPHEM" => "YES",
	"TABLE_TYPE" => "VECTORS",
	"START_TIME" => "2017-02-15",
	"STOP_TIME" => "2018-02-15",
	"STEP_SIZE" => "1d",
	"OUT_UNITS" => "KM-S",
	"REF_PLANE" => "ECLIPTIC",
	"REF_SYSTEM" => "J2000",
	"VECT_CORR" => "NONE",
	"VEC_LABELS" => "YES",
	"VEC_DELTA_T" => "NO",
	"CSV_FORMAT" => "YES",
	"OBJ_DATA" => "NO",
	"VEC_TABLE" => "3")

function horizon_data(bodies::Dict{String, String}; url="http://ssd.jpl.nasa.gov/horizons_batch.cgi")
	for (i, k) in enumerate(keys(bodies))
		#download data from horizons
		query["COMMAND"] = k
		response = get(url, query = query)
		println(response)
		str = readstring(response)

		#filter out data
		m = match(r"\$\$SOE\n(?<csv_data>[\s-\S]*)\n\$\$EOE", str)
		data = m[:csv_data]
		data = replace(data, ",\n", "\n")
		df = CSV.read(IOBuffer(data); types=Dict(2=>String), header=["date1", "date2", "X", "Y", "Z", "VX", "VY", "VZ", "LT", "RG", "RR"], nullable=false)
		CSV.write("$(bodies[k]).csv", df)
	end
end

function csv_to_vtu(bodies::Dict{String,String}; file_name="my_vtu_file")
	vectors = Array{DataFrame}(bodies.count)

	for (i, k) in enumerate(keys(bodies))
		vectors[i] = CSV.read("$(bodies[k]).csv", nullable=false)
	end

	cells = WriteVTK.MeshCell[MeshCell(VTKCellTypes.VTK_VERTEX, [i]) for i = 1:length(vectors)]
	for i = 1:size(vectors[1])[1]
		points = Array{Float64}(3, length(vectors))
		for j = 1:length(vectors)
			points[1,j] = vectors[j][:X][i]
			points[2,j] = vectors[j][:Y][i]
			points[3,j] = vectors[j][:Z][i]
		end
		outfile = vtk_grid(string(file_name,i), points, cells) do vtk
			vtk_cell_data(vtk, [1 for i = 1:length(vectors)], "Mass")
		end
	end
end

major_bodies = Dict{String,String}(
	"10" => "Sun",
	"199" => "Mercury",
	"299" => "Venus",
	#"301" => "Moon",
	"399" => "Earth",
	#"401" => "Phobos",
	#"402" => "Deimos",
	"499" => "Mars",
	"599" => "Jupiter",
	"699" => "Saturn",
	"799" => "Uranus",
	"800" => "Neptune")

horizon_data(major_bodies)
csv_to_vtu(major_bodies, file_name="major_bodies")


include("../shared_source/boilerplate.jl")

import ClimateMachine.BalanceLaws:
    vars_state,
    init_state_prognostic!,
    init_state_auxiliary!,
    compute_gradient_argument!,
    compute_gradient_flux!,
    flux_first_order!,
    flux_second_order!,
    source!,
    wavespeed,
    boundary_conditions,
    boundary_state!
import ClimateMachine.DGMethods: DGModel
import ClimateMachine.NumericalFluxes: numerical_flux_first_order!
import ClimateMachine.Orientations: vertical_unit_vector

"""
    ThreeDimensionalCompressibleNavierStokesEquations <: BalanceLaw
A `BalanceLaw` for shallow water modeling.
write out the equations here
# Usage
    ThreeDimensionalCompressibleNavierStokesEquations()
"""
abstract type AbstractFluid3D <: AbstractFluid end
struct Fluid3D <: AbstractFluid3D end

struct ThreeDimensionalCompressibleNavierStokesEquations{
    I,
    D,
    O,
    A,
    T,
    C,
    F,
    BC,
    FT,
} <: AbstractFluid3D
    initial_value_problem::I
    domain::D
    orientation::O
    advection::A
    turbulence::T
    coriolis::C
    forcing::F
    boundary_conditions::BC
    cₛ::FT
    ρₒ::FT
    function ThreeDimensionalCompressibleNavierStokesEquations{FT}(
        initial_value_problem::I,
        domain::D,
        orientation::O,
        advection::A,
        turbulence::T,
        coriolis::C,
        forcing::F,
        boundary_conditions::BC;
        cₛ = FT(sqrt(10)),  # m/s
        ρₒ = FT(1),  #kg/m³
    ) where {FT <: AbstractFloat, I, D, O, A, T, C, F, BC}
        return new{I, D, O, A, T, C, F, BC, FT}(
            initial_value_problem,
            domain,
            orientation,
            advection,
            turbulence,
            coriolis,
            forcing,
            boundary_conditions,
            cₛ,
            ρₒ,
        )
    end
end

CNSE3D = ThreeDimensionalCompressibleNavierStokesEquations

function vars_state(m::CNSE3D, ::Prognostic, T)
    @vars begin
        ρ::T
        ρu::SVector{3, T}
        ρθ::T
    end
end

function init_state_prognostic!(m::CNSE3D, state::Vars, aux::Vars, localgeo, t)
    cnse_init_state!(m, state, aux, localgeo, t)
end

# default initial state if IVP == nothing
function cnse_init_state!(model::CNSE3D, state, aux, localgeo, t)

    x = aux.x
    y = aux.y
    z = aux.z

    ρ = model.ρₒ
    state.ρ = ρ
    state.ρu = ρ * @SVector [-0, -0, -0]
    state.ρθ = ρ

    return nothing
end

# user defined initial state
function cnse_init_state!(
    model::CNSE3D{<:InitialValueProblem},
    state,
    aux,
    localgeo,
    t,
)
    x = aux.x
    y = aux.y
    z = aux.z

    params = model.initial_value_problem.params
    ic = model.initial_value_problem.initial_conditions

    state.ρ = ic.ρ(params, x, y, z)
    state.ρu = ic.ρu(params, x, y, z)
    state.ρθ = ic.ρθ(params, x, y, z)

    return nothing
end

function vars_state(m::CNSE3D, st::Auxiliary, T)
    @vars begin
        x::T
        y::T
        z::T
        orientation::vars_state(m.orientation, st, T)
    end
end

function init_state_auxiliary!(
    m::CNSE3D,
    state_auxiliary::MPIStateArray,
    grid,
    direction,
)
    # update the geopotential Φ in state_auxiliary.orientation.Φ
    init_state_auxiliary!(
        m,
        (m, aux, tmp, geom) ->
            orientation_nodal_init_aux!(m.orientation, m.domain, aux, geom),
        state_auxiliary,
        grid,
        direction,
    )

    # update ∇Φ in state_auxiliary.orientation.∇Φ
    orientation_gradient(m, m.orientation, state_auxiliary, grid, direction)

    # store coordinates and potentially other stuff
    init_state_auxiliary!(
        m,
        (m, aux, tmp, geom) -> cnse_init_aux!(m, aux, geom),
        state_auxiliary,
        grid,
        direction,
    )

    return nothing
end

function orientation_gradient(
    model::CNSE3D,
    ::Orientation,
    state_auxiliary,
    grid,
    direction,
)
    auxiliary_field_gradient!(
        model,
        state_auxiliary,
        ("orientation.∇Φ",),
        state_auxiliary,
        ("orientation.Φ",),
        grid,
        direction,
    )

    return nothing
end

function orientation_gradient(::CNSE3D, ::NoOrientation, _...)
    return nothing
end

function orientation_nodal_init_aux!(
    ::SphericalOrientation,
    domain::Tuple,
    aux::Vars,
    geom::LocalGeometry,
)
    norm_R = norm(geom.coord)
    @inbounds aux.orientation.Φ = norm_R - domain[1]

    return nothing
end

"""
function orientation_nodal_init_aux!(
    ::SuperSphericalOrientation,
    domain::Tuple,
    aux::Vars,
    geom::LocalGeometry,
)
    norm_R = norm(geom.coord)
    @inbounds aux.orientation.Φ = 1 / norm_R^2
end
"""

function orientation_nodal_init_aux!(
    ::FlatOrientation,
    domain::Tuple,
    aux::Vars,
    geom::LocalGeometry,
)
    @inbounds aux.orientation.Φ = geom.coord[3]

    return nothing
end

function orientation_nodal_init_aux!(
    ::NoOrientation,
    domain::Tuple,
    aux::Vars,
    geom::LocalGeometry,
)
    return nothing
end

function cnse_init_aux!(::CNSE3D, aux, geom)
    @inbounds begin
        aux.x = geom.coord[1]
        aux.y = geom.coord[2]
        aux.z = geom.coord[3]
    end

    return nothing
end

function vars_state(m::CNSE3D, ::Gradient, T)
    @vars begin
        ∇ρ::T
        ∇u::SVector{3, T}
        ∇θ::T
        ∇p::T
    end
end

function compute_gradient_argument!(
    model::CNSE3D,
    grad::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ρ = state.ρ
    cₛ = model.cₛ
    ρₒ = model.ρₒ

    grad.∇p = (cₛ * ρ)^2 / (2 * ρₒ)

    compute_gradient_argument!(model.turbulence, grad, state, aux, t)
end

@inline function compute_gradient_argument!(
    ::ConstantViscosity,
    grad::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ρ = state.ρ
    ρu = state.ρu
    ρθ = state.ρθ

    u = ρu
    θ = ρθ

    grad.∇ρ = ρ
    grad.∇u = u
    grad.∇θ = θ

    return nothing
end

function vars_state(m::CNSE3D, ::GradientFlux, T)
    @vars begin
        μ∇ρ::SVector{3, T}
        ν∇u::SMatrix{3, 3, T, 9}
        κ∇θ::SVector{3, T}
        ∇p::SVector{3, T}
    end
end

function compute_gradient_flux!(
    model::CNSE3D,
    gradflux::Vars,
    grad::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)

    gradflux.∇p = grad.∇p

    compute_gradient_flux!(
        model,
        model.turbulence,
        gradflux,
        grad,
        state,
        aux,
        t,
    )
end

@inline function compute_gradient_flux!(
    ::CNSE3D,
    turb::ConstantViscosity,
    gradflux::Vars,
    grad::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    μ = turb.μ * I
    ν = turb.ν * I
    κ = turb.κ * I

    gradflux.μ∇ρ = -μ * grad.∇ρ
    gradflux.ν∇u = -ν * grad.∇u
    gradflux.κ∇θ = -κ * grad.∇θ

    return nothing
end

@inline function flux_first_order!(
    model::CNSE3D,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    ρ = state.ρ
    ρu = state.ρu
    ρθ = state.ρθ

    cₛ = model.cₛ
    ρₒ = model.ρₒ

    flux.ρ += ρu
    # flux.ρu += (cₛ * ρ)^2 / (2 * ρₒ) * I

    advective_flux!(model, model.advection, flux, state, aux, t)

    return nothing
end

advective_flux!(::CNSE3D, ::Nothing, _...) = nothing

@inline function advective_flux!(
    ::CNSE3D,
    ::NonLinearAdvectionTerm,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ρ = state.ρ
    ρu = state.ρu
    ρθ = state.ρθ

    flux.ρu += ρu ⊗ ρu / ρ
    flux.ρθ += ρu * ρθ / ρ

    return nothing
end

function flux_second_order!(
    model::CNSE3D,
    flux::Grad,
    state::Vars,
    gradflux::Vars,
    ::Vars,
    aux::Vars,
    t::Real,
)
    flux_second_order!(model, model.turbulence, flux, state, gradflux, aux, t)
end

@inline function flux_second_order!(
    ::CNSE3D,
    ::ConstantViscosity,
    flux::Grad,
    state::Vars,
    gradflux::Vars,
    aux::Vars,
    t::Real,
)
    flux.ρ += gradflux.μ∇ρ
    flux.ρu += gradflux.ν∇u
    flux.ρθ += gradflux.κ∇θ

    return nothing
end

@inline function source!(
    model::CNSE3D,
    source::Vars,
    state::Vars,
    gradflux::Vars,
    aux::Vars,
    t::Real,
    direction,
)

    source.ρu -= gradflux.∇p

    coriolis_force!(model, model.coriolis, source, state, aux, t)
    forcing_term!(model, model.forcing, source, state, aux, t)

    return nothing
end

coriolis_force!(::CNSE3D, ::Nothing, _...) = nothing

@inline function coriolis_force!(
    model::CNSE3D,
    coriolis::fPlaneCoriolis,
    source,
    state,
    aux,
    t,
)
    # f × u
    f = [-0, -0, coriolis_parameter(model, coriolis, aux.coords)]
    ρu = state.ρu

    source.ρu -= f × ρu

    return nothing
end

@inline function coriolis_force!(
    model::CNSE3D,
    coriolis::SphereCoriolis,
    source,
    state,
    aux,
    t,
)
    # f × u
    Ω = coriolis.Ω
    f = @SVector [-0, -0, 2Ω]
    ρu = state.ρu
    source.ρu -= f × ρu
    return nothing
end

forcing_term!(::CNSE3D, ::Nothing, _...) = nothing

@inline function forcing_term!(
    model::CNSE3D,
    buoy::Buoyancy,
    source,
    state,
    aux,
    t,
)
    α = buoy.α
    g = buoy.g
    ρθ = state.ρθ

    # as temperature increase, density decreases
    B = -α * g * ρθ
    k̂ = vertical_unit_vector(model.orientation, aux)

    # gravity points downward
    source.ρu -= B * k̂
end

@inline vertical_unit_vector(::Orientation, aux) = aux.orientation.∇Φ
@inline vertical_unit_vector(::NoOrientation, aux) = @SVector [0, 0, 1]

@inline wavespeed(m::CNSE3D, _...) = m.cₛ

roe_average(ρ⁻, ρ⁺, var⁻, var⁺) =
    (sqrt(ρ⁻) * var⁻ + sqrt(ρ⁺) * var⁺) / (sqrt(ρ⁻) + sqrt(ρ⁺))

function numerical_flux_first_order!(
    ::RoeNumericalFlux,
    model::CNSE3D,
    fluxᵀn::Vars{S},
    n⁻::SVector,
    state⁻::Vars{S},
    aux⁻::Vars{A},
    state⁺::Vars{S},
    aux⁺::Vars{A},
    t,
    direction,
) where {S, A}
    numerical_flux_first_order!(
        CentralNumericalFluxFirstOrder(),
        model,
        fluxᵀn,
        n⁻,
        state⁻,
        aux⁻,
        state⁺,
        aux⁺,
        t,
        direction,
    )

    FT = eltype(fluxᵀn)

    # constants and normal vectors
    cₛ = model.cₛ
    ρₒ = model.ρₒ

    # - states
    ρ⁻ = state⁻.ρ
    ρu⁻ = state⁻.ρu
    ρθ⁻ = state⁻.ρθ

    # constructed states
    u⁻ = ρu⁻ / ρ⁻
    θ⁻ = ρθ⁻ / ρ⁻
    uₙ⁻ = u⁻' * n⁻

    # in general thermodynamics
    p⁻ = (cₛ * ρ⁻)^2 / (2 * ρₒ)
    c⁻ = cₛ * sqrt(ρ⁻ / ρₒ)

    # + states
    ρ⁺ = state⁺.ρ
    ρu⁺ = state⁺.ρu
    ρθ⁺ = state⁺.ρθ

    # constructed states
    u⁺ = ρu⁺ / ρ⁺
    θ⁺ = ρθ⁺ / ρ⁺
    uₙ⁺ = u⁺' * n⁻

    # in general thermodynamics
    p⁺ = (cₛ * ρ⁺)^2 / (2 * ρₒ)
    c⁺ = cₛ * sqrt(ρ⁺ / ρₒ)

    # construct roe averges
    ρ = sqrt(ρ⁻ * ρ⁺)
    u = roe_average(ρ⁻, ρ⁺, u⁻, u⁺)
    θ = roe_average(ρ⁻, ρ⁺, θ⁻, θ⁺)
    c = roe_average(ρ⁻, ρ⁺, c⁻, c⁺)

    # construct normal velocity
    uₙ = u' * n⁻

    # differences
    Δρ = ρ⁺ - ρ⁻
    Δp = p⁺ - p⁻
    Δu = u⁺ - u⁻
    Δρθ = ρθ⁺ - ρθ⁻
    Δuₙ = Δu' * n⁻

    # constructed values
    c⁻² = 1 / c^2
    w1 = abs(uₙ - c) * (Δp - ρ * c * Δuₙ) * 0.5 * c⁻²
    w2 = abs(uₙ + c) * (Δp + ρ * c * Δuₙ) * 0.5 * c⁻²
    w3 = abs(uₙ) * (Δρ - Δp * c⁻²)
    w4 = abs(uₙ) * ρ
    w5 = abs(uₙ) * (Δρθ - θ * Δp * c⁻²)

    # fluxes!!!
    fluxᵀn.ρ -= (w1 + w2 + w3) * 0.5
    fluxᵀn.ρu -=
        (
            w1 * (u - c * n⁻) +
            w2 * (u + c * n⁻) +
            w3 * u +
            w4 * (Δu - Δuₙ * n⁻)
        ) * 0.5
    fluxᵀn.ρθ -= ((w1 + w2) * θ + w5) * 0.5

    return nothing
end

boundary_conditions(model::CNSE3D) = model.boundary_conditions

"""
    boundary_state!(nf, ::CNSE3D, args...)
applies boundary conditions for the hyperbolic fluxes
dispatches to a function in CNSEBoundaryConditions
"""
@inline function boundary_state!(nf, bc, model::CNSE3D, args...)
    return _cnse_boundary_state!(nf, bc, model, args...)
end

"""
    cnse_boundary_state!(nf, bc::FluidBC, ::CNSE3D)
splits boundary condition application into velocity
"""
@inline function cnse_boundary_state!(nf, bc::FluidBC, m::CNSE3D, args...)
    cnse_boundary_state!(nf, bc.momentum, m, m.turbulence, args...)
    cnse_boundary_state!(nf, bc.temperature, m, args...)

    return nothing
end

include("bc_momentum.jl")
include("bc_temperature.jl")

"""
STUFF FOR ANDRE'S WRAPPERS
"""

function get_boundary_conditions(
    model::SpatialModel{BL},
) where {BL <: AbstractFluid3D}
    bcs = model.boundary_conditions

    west_east = (check_bc(bcs, :west), check_bc(bcs, :east))
    south_north = (check_bc(bcs, :south), check_bc(bcs, :north))
    bottom_top = (check_bc(bcs, :bottom), check_bc(bcs, :top))

    return (west_east..., south_north..., bottom_top...)
end

function DGModel(
    model::SpatialModel{BL};
    initial_conditions = nothing,
) where {BL <: AbstractFluid3D}
    params = model.parameters
    physics = model.physics

    Lˣ, Lʸ, Lᶻ = length(model.grid.domain)
    bcs = get_boundary_conditions(model)
    FT = eltype(model.grid.numerical.vgeo)

    if !isnothing(initial_conditions)
        initial_conditions = InitialValueProblem(params, initial_conditions)
    end

    balance_law = CNSE3D{FT}(
        initial_conditions,
        (Lˣ, Lʸ, Lᶻ),
        physics.orientation,
        physics.advection,
        physics.dissipation,
        physics.coriolis,
        physics.buoyancy,
        bcs,
        ρₒ = params.ρₒ,
        cₛ = params.cₛ,
    )

    numerical_flux_first_order = model.numerics.flux # should be a function

    rhs = DGModel(
        balance_law,
        model.grid.numerical,
        numerical_flux_first_order,
        CentralNumericalFluxSecondOrder(),
        CentralNumericalFluxGradient(),
    )

    return rhs
end


include("../src/benchmark.jl")
include("../src/save.jl")

images, sim_params = main();

beam_energy_map, transmit_time_map = WaveSim.beam_energy_map_and_transmit_time_map(images, sim_params);

saveall(images, beam_energy_map, transmit_time_map, sim_params, ".")


#
# BaseFields are just wrappers around arrays tagged with additional
# information in the form of:
#  * the `B` parameter, for information which is needed at
#    compile-time (right now, just the basis)
#  * the `metadata` field, for information which is only needed at
#    run-time
#

abstract type Proj end
broadcastable(proj::Proj) = Ref(proj)

struct BaseField{B, M<:Proj, T, A<:AbstractArray{T}} <: Field{B, T}
    arr :: A
    metadata :: M
    function (::Type{F})(arr::A, metadata::M) where {B,M<:Proj,T,A<:AbstractArray{T},F<:BaseField{B}}
        check_field_consistency(B(), arr, metadata)
        new{B,M,T,A}(arr, metadata) :: F
    end
end

# can be overrriden for specific types to check things like the
# dimensions of arr being consistent with metadata, etc...
check_field_consistency(::Any, ::Any, ::Any) = ()

typealias_def(::Type{F}) where {B,M,T,A,F<:BaseField{B,M,T,A}} = "BaseField{$(typealias(B)),$(typealias(A)),$(typealias(M))}"

## array interface
# even though the underlying data in BaseFields is multi-dimensional,
# they act as a 1D vector with all their entries splayed out in a
# column
size(f::BaseField) = (length(f.arr),)
lastindex(f::BaseField, i::Int) = lastindex(f.arr, i)
@propagate_inbounds getindex(f::BaseField, I::Union{Int,Colon,AbstractArray}...) = getindex(f.arr, I...)
@propagate_inbounds setindex!(f::BaseField, X, I::Union{Int,Colon,AbstractArray}...) = (setindex!(f.arr, X, I...); f)
similar(f::BaseField{B}, ::Type{T}) where {B,T} = BaseField{B}(similar(f.arr, T), f.metadata)
copy(f::BaseField{B}) where {B} = BaseField{B}(copy(f.arr), f.metadata)

## promotion
function promote(f₁::BaseField{B₁}, f₂::BaseField{B₂}) where {B₁,B₂}
    get_metadata_strict(f₁.metadata, f₂.metadata) # for now just check compatible
    B = typeof(promote_basis_generic(B₁(), B₂()))
    B(f₁), B(f₂)
end

## broadcasting

# any broadcast expression that contains at least one BaseField will
# have a broadcast style of BaseFieldStyle{S,B}. the S is the
# broadcast style for the underlying arrays and B is the B parameter
# of the result BaseField. S and B are promoted recursively across all
# the arguments according to the set of rules below.
struct BaseFieldStyle{S,B} <: AbstractArrayStyle{1} end
BroadcastStyle(::Type{F}) where {B,M,T,A,F<:BaseField{B,M,T,A}} =
    BaseFieldStyle{typeof(BroadcastStyle(A)),B}()
BroadcastStyle(::BaseFieldStyle{S₁,B₁}, ::BaseFieldStyle{S₂,B₂}) where {S₁,B₁,S₂,B₂} =
    BaseFieldStyle{typeof(result_style(S₁(), S₂())), typeof(promote_basis_strict(B₁(),B₂()))}()
BroadcastStyle(S::BaseFieldStyle, ::DefaultArrayStyle{0}) = S

# with the Broadcasted object created, we now compute the answer
function materialize(bc::Broadcasted{BaseFieldStyle{S,B}}) where {S,B}

    # first, recursively go through the broadcast arguments and figure
    # out the metadata of the result, using the
    # promote_metadata_strict rules
    metadata = get_metadata_strict(bc)

    # then "preprocess" all the arguments. this unwraps all of the
    # BaseFields in the expression into just the underlying arrays,
    # and turns things which were ImplicitFields into actual arrays
    # (which are computed by dispatching on the now-known S, B, and
    # metadata of the result)
    bc′ = preprocess((BaseFieldStyle{S,B}(), metadata), bc)

    # the arguments of bc′ are now all normal arrays, so convert it to
    # the broadcast style S that goes along with them
    bc″ = convert(Broadcasted{S}, bc′)

    # run the normal array broadcast, and wrap in the right
    # result type
    BaseField{B}(materialize(bc″), metadata)

end

function materialize!(dst::BaseField{B}, bc::Broadcasted{BaseFieldStyle{S,B′}}) where {B,B′,S}

    (B == B′) || error("Can't broadcast a $(typealias(B′)) into a $(typealias(B))")

    # for inplace broadcasting, we don't need to compute B or the
    # metadata from the broadcasted object, we just take it from the
    # destination BaseField. otherwise its the same as materialize above
    bc′ = preprocess((BaseFieldStyle{S,B}(), dst.metadata), bc)
    bc″ = convert(Broadcasted{S}, bc′)
    materialize!(dst.arr, bc″)
    dst

end

# the default preprocessing, which just unwraps the underlying array.
# this doesn't dispatch on the first argument, but custom BaseFields
# are free to override this and dispatch on it if they need
preprocess(::Any, f::BaseField) = f.arr

# we re-wrap each Broadcasted object as we go through preprocessing
# because some array types do special things here (e.g. CUDA wraps
# bc.f in a CUDA.cufunc)
preprocess(dest::Tuple{BaseFieldStyle{S,B},M}, bc::Broadcasted) where {S,B,M} =
    broadcasted(S(), bc.f, preprocess_args(dest, bc.args)...)

# recursively go through a Broadcasted object's arguments and compute
# the final metadata according to the promote_metadata_strict rules.
# we use the select_known_rule machinery (see util.jl) to make it so
# promote_metadata_strict_rule only needs one argument order defined
# (similar to Base.promote_rule)
get_metadata_strict(x, rest...)      = promote_metadata_strict(get_metadata_strict(x), get_metadata_strict(rest...))
get_metadata_strict(bc::Broadcasted) = get_metadata_strict(bc.args...)
get_metadata_strict(f ::BaseField)   = f.metadata
get_metadata_strict(  ::Any)         = nothing
get_metadata_strict()                = nothing

promote_metadata_strict(x) = x
promote_metadata_strict(x, y) = select_known_rule(promote_metadata_strict_rule, x, y)
promote_metadata_strict_rule(metadata,   ::Nothing) = metadata
promote_metadata_strict_rule(::Nothing,  ::Nothing) = nothing
promote_metadata_strict_rule(::Any,      ::Any) = Unknown()

## properties
getproperty(f::BaseField, s::Symbol)           = getproperty(f,Val(s))
getproperty(f::BaseField,  ::Val{:arr})        = getfield(f,:arr)
getproperty(f::BaseField,  ::Val{:metadata})   = getfield(f,:metadata)
getproperty(f::BaseField,  ::Val{s}) where {s} = getfield(getfield(f,:metadata),s)
propertynames(f::BaseField) = (fieldnames(typeof(f))..., fieldnames(typeof(f.metadata))...)

## CMBLensing-specific stuff
global_rng_for(::Type{BaseField{B,M,T,A}}) where {B,M,T,A} = global_rng_for(A)
fieldinfo(f::BaseField) = f # for backwards compatibility
get_storage(f::BaseField) = typeof(f.arr)
adapt_structure(to, f::BaseField{B}) where {B} = BaseField{B}(adapt(to, f.arr), adapt(to, f.metadata))
hash(f::BaseField, h::UInt64) = foldr(hash, (typeof(f), cpu(f.arr), f.metadata), init=h)

#
default_proj(::Type{F}) where {F<:BaseField{<:Any,<:Proj}} = Base.unwrap_unionall(F).parameters[2].ub
make_field_aliases("Base", Proj)

# useful for enumerating some cases below and in plotting
_sub_components = [
    (:BaseMap,        (:Ix=>:, :I=>:)),
    (:BaseFourier,    (:Il=>:, :I=>:)),
    (:BaseQUMap,      (:Qx=>1, :Ux=>2, :Q =>1, :U=>2, :P=>:)),
    (:BaseQUFourier,  (:Ql=>1, :Ul=>2, :Q =>1, :U=>2, :P=>:)),
    (:BaseEBMap,      (:Ex=>1, :Bx=>2, :E =>1, :B=>2, :P=>:)),
    (:BaseEBFourier,  (:El=>1, :Bl=>2, :E =>1, :B=>2, :P=>:)),
    (:BaseIQUMap,     (:Ix=>1, :Qx=>2, :Ux=>3, :I=>1, :Q=>2, :U=>3, :P=>2:3, :IP=>:)),
    (:BaseIQUFourier, (:Il=>1, :Ql=>2, :Ul=>3, :I=>1, :Q=>2, :U=>3, :P=>2:3, :IP=>:)),
    (:BaseIEBMap,     (:Ix=>1, :Ex=>2, :Bx=>3, :I=>1, :E=>2, :B=>3, :P=>2:3, :IP=>:)),
    (:BaseIEBFourier, (:Il=>1, :El=>2, :Bl=>3, :I=>1, :E=>2, :B=>3, :P=>2:3, :IP=>:)),
]

# sub-components like f.Ix, f.Q, f.P, etc...
for (F, props) in _sub_components
    for (k, I) in props
        if String(k)[end] in "xl"
            if I == (:)
                @eval getproperty(f::$F, ::Val{$(QuoteNode(k))}) = getfield(f,:arr)
            else
                @eval getproperty(f::$F, ::Val{$(QuoteNode(k))}) = view(getfield(f,:arr), pol_slice(f,$I)...)
            end
        else
            if I == (:)
                @eval getproperty(f::$F, ::Val{$(QuoteNode(k))}) = f
            else
                B = (k == :P) ? Basis2Prod{basis(@eval($F)).parameters[end-1:end]...} : basis(@eval($F)).parameters[end]
                @eval getproperty(f::$F, ::Val{$(QuoteNode(k))}) = BaseField{$B}(_reshape_batch(view(getfield(f,:arr), pol_slice(f,$I)...)), f.metadata)
            end
        end
    end
end


abstract type AbstractDimensionalArray{T,N,D<:Tuple} <: AbstractArray{T,N} end

const AbDimArray = AbstractDimensionalArray

dims(a::AbDimArray) = a.dims
label(a::AbDimArray) = ""

# Array interface
Base.size(a::AbDimArray) = size(parent(a))
Base.iterate(a::AbDimArray, args...) = iterate(parent(a), args...)
Base.show(io::IO, a::AbDimArray) = begin
    printstyled(io, "\n", label(a), ": "; color=:red)
    show(io, typeof(a))
    show(io, parent(a))
    printstyled(io, "\n\ndims:\n"; color=:magenta)
    show(io, dims(a))
    show(io, refdims(a))
    printstyled(io, "\n\nmetadata:\n"; color=:cyan)
end

Base.@propagate_inbounds Base.getindex(a::AbDimArray, I::Vararg{<:Integer}) =
    getindex(parent(a), I...)
Base.@propagate_inbounds Base.getindex(a::AbDimArray, I::Vararg{<:Union{AbstractArray,Colon,Integer}}) =
    rebuildsliced(a, getindex(parent(a), I...), I)

Base.@propagate_inbounds Base.view(a::AbDimArray, I::Vararg{<:Union{AbstractArray,Colon,Integer}}) =
    rebuildsliced(a, view(parent(a), I...), I)

Base.convert(::Type{Array{T,N}}, a::AbDimArray{T,N}) where {T,N} =
    convert(Array{T,N}, parent(a))

# Similar. TODO this need a rethink. How do we know what the new dims are?
Base.BroadcastStyle(::Type{<:AbDimArray}) = Broadcast.ArrayStyle{AbDimArray}()
# Need to cover a few type signatures to avoid ambiguity with base
Base.similar(a::AbDimArray, ::Type{T}, I::Vararg{<:Integer}) where T =
    rebuildsliced(a, similar(parent(a), T, I...), I)
Base.similar(a::AbDimArray, ::Type{T}) where T = rebuild(a, similar(parent(a), T))
Base.similar(a::AbDimArray) = rebuild(a, similar(parent(a)))
Base.similar(a::AbDimArray, ::Type{T}, I::Tuple{Union{Integer,OneTo},Vararg{Union{Integer,OneTo},N}}) where {T,N} =
    rebuildsliced(a, similar(parent(a), T, I...), I)
Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{AbDimArray}}, ::Type{ElType}) where ElType =
    rebuildsliced(find_dimensional(bc), similar(Array{ElType}, axes(bc)), axes(bc))

@inline find_dimensional(bc::Base.Broadcast.Broadcasted) = find_dimensional(bc.args)
@inline find_dimensional(ext::Base.Broadcast.Extruded) = find_dimensional(ext.x)
@inline find_dimensional(args::Tuple{}) = error("dimensional array not found")
@inline find_dimensional(args::Tuple) = find_dimensional(find_dimensional(args[1]), tail(args))
@inline find_dimensional(x) = x
@inline find_dimensional(a::AbDimArray, rest) = a
@inline find_dimensional(::Any, rest) = find_dimensional(rest)

"""
A basic DimensionalArray type

Maintains and updates its dimensions through transformations
"""
struct DimensionalArray{T,N,D<:Tuple,R<:Tuple,A<:AbstractArray{T,N}} <: AbstractDimensionalArray{T,N,D}
    data::A
    dims::D
    refdims::R
end
DimensionalArray(a::AbstractArray, dims; refdims=()) =
    DimensionalArray(a, formatdims(a, dims), refdims)

# Getters
refdims(a::DimensionalArray) = a.refdims

# DimensionalArray interface
@inline rebuild(a::DimensionalArray, data, dims, refdims) =
    DimensionalArray(data, dims, refdims)

# Array interface (AbstractDimensionalArray takes care of everything else)
Base.parent(a::DimensionalArray) = a.data

Base.@propagate_inbounds Base.setindex!(a::DimensionalArray, x, I::Vararg{<:Union{AbstractArray,Colon,Real}}) =
    setindex!(parent(a), x, I...)


"""
Wraps an optimizer in a hierarchy which enables one `bridge`.
"""
function lazy_relax(optimizer, bridge::Type{<:MOIB.Constraint.AbstractBridge})
        function construct()
                model = Model(optimizer)
                JuMP.add_bridge(model, bridge)
                backend(model).optimizer
        end
        construct
end


##############################################################################
#
# Basic DocumentTermMatrix type
#
##############################################################################

mutable struct DocumentTermMatrix
    dtm::SparseMatrixCSC{Int, Int}
    terms::Vector{String}
    column_indices::Dict{String, Int}
end

##############################################################################
#
# Construct a DocumentTermMatrix from a Corpus
#
##############################################################################

# create col index lookup dictionary from a (sorted?) vector of terms
function columnindices(terms::Vector{String})
    column_indices = Dict{String, Int}()
    for i in 1:length(terms)
        term = terms[i]
        column_indices[term] = i
    end
    column_indices
end

function DocumentTermMatrix(crps::Corpus, terms::Vector{String})
    column_indices = columnindices(terms)

    m = length(crps)
    n = length(terms)

    rows = Array{Int}(undef, 0)
    columns = Array{Int}(undef, 0)
    values = Array{Int}(undef, 0)
    for i in 1:m
        doc = crps.documents[i]
        ngs = ngrams(doc)
        for ngram in keys(ngs)
            j = get(column_indices, ngram, 0)
            v = ngs[ngram]
            if j != 0
                push!(rows, i)
                push!(columns, j)
                push!(values, v)
            end
        end
    end
    if length(rows) > 0
        dtm = sparse(rows, columns, values, m, n)
    else
        dtm = spzeros(Int, m, n)
    end
    DocumentTermMatrix(dtm, terms, column_indices)
end
DocumentTermMatrix(crps::Corpus) = DocumentTermMatrix(crps, lexicon(crps))

DocumentTermMatrix(crps::Corpus, lex::AbstractDict) = DocumentTermMatrix(crps, sort(collect(keys(lex))))

DocumentTermMatrix(dtm::SparseMatrixCSC{Int, Int},terms::Vector{String}) = DocumentTermMatrix(dtm, terms, columnindices(terms))

##############################################################################
#
# Access the DTM of a DocumentTermMatrix
#
##############################################################################

function dtm(d::DocumentTermMatrix, density::Symbol)
    if density == :sparse
        return d.dtm
    else
        return Matrix(d.dtm)
    end
end

function dtm(d::DocumentTermMatrix)
    return d.dtm
end

function dtm(crps::Corpus)
    dtm(DocumentTermMatrix(crps))
end

tdm(crps::DocumentTermMatrix, density::Symbol) = dtm(crps, density)' #'

tdm(crps::DocumentTermMatrix) = dtm(crps)' #'

tdm(crps::Corpus) = dtm(crps)' #'

##############################################################################
#
# Produce the signature of a DTM entry for a document
#
# TODO: Make this more efficieny by reusing column_indices
#
##############################################################################

function dtm_entries(d::AbstractDocument, lex::Dict{String, Int})
    ngs = ngrams(d)
    indices = Array{Int}(undef, 0)
    values = Array{Int}(undef, 0)
    terms = sort(collect(keys(lex)))
    column_indices = columnindices(terms)

    for ngram in keys(ngs)
        if haskey(column_indices, ngram)
            push!(indices, column_indices[ngram])
            push!(values, ngs[ngram])
        end
    end
    return (indices, values)
end

function dtv(d::AbstractDocument, lex::Dict{String, Int})
    p = length(keys(lex))
    row = zeros(Int, 1, p)
    indices, values = dtm_entries(d, lex)
    for i in 1:length(indices)
        row[1, indices[i]] = values[i]
    end
    return row
end

function dtv(d::AbstractDocument)
    error("Cannot construct a DTV without a pre-existing lexicon")
end

##############################################################################
#
# The hash trick: use a hash function instead of a lexicon to determine the
# columns of a DocumentTermMatrix-like encoding of the data
#
##############################################################################

function hash_dtv(d::AbstractDocument, h::TextHashFunction)
    p = cardinality(h)
    res = zeros(Int, 1, p)
    ngs = ngrams(d)
    for ng in keys(ngs)
        res[1, index_hash(ng, h)] += ngs[ng]
    end
    return res
end

hash_dtv(d::AbstractDocument) = hash_dtv(d, TextHashFunction())

function hash_dtm(crps::Corpus, h::TextHashFunction)
    n, p = length(crps), cardinality(h)
    res = zeros(Int, n, p)
    for i in 1:length(crps)
        doc = crps.documents[i]
        res[i, :] = hash_dtv(doc, h)
    end
    return res
end

hash_dtm(crps::Corpus) = hash_dtm(crps, hash_function(crps))

hash_tdm(crps::Corpus) = hash_dtm(crps)' #'

##############################################################################
#
# Produce entries for on-line analysis when DTM would not fit in memory
#
##############################################################################

mutable struct EachDTV
    crps::Corpus
end

start(edt::EachDTV) = 1

function next(edt::EachDTV, state::Int)
    return (dtv(edt.crps.documents[state], lexicon(edt.crps)), state + 1)
end

done(edt::EachDTV, state::Int) = state > length(edt.crps.documents)

mutable struct EachHashDTV
    crps::Corpus
end

start(edt::EachHashDTV) = 1

function next(edt::EachHashDTV, state::Int)
    (hash_dtv(edt.crps.documents[state]), state + 1)
end

done(edt::EachHashDTV, state::Int) = state > length(edt.crps.documents)

each_dtv(crps::Corpus) = EachDTV(crps)

each_hash_dtv(crps::Corpus) = EachHashDTV(crps)

##
## getindex() methods
##

Base.getindex(dtm::DocumentTermMatrix, k::AbstractString) = dtm.dtm[:, dtm.column_indices[k]]
Base.getindex(dtm::DocumentTermMatrix, i::Any) = dtm.dtm[i]
Base.getindex(dtm::DocumentTermMatrix, i::Any, j::Any) = dtm.dtm[i, j]


# define common PETSc constants
# this excludes configurable constants (e.g. PetscScalar) which are set in lib.jl

const PetscErrorCode      = Cint

struct PetscError <: Exception
    code::PetscErrorCode
end

macro chk(expr)
    :((errcode = $(esc(expr))) == 0 || throw(PetscError(errcode)))
end

const PETSC_DEFAULT = -2
const PETSC_DECIDE = PETSC_DETERMINE = -1

@enum DMBoundaryType begin
    DM_BOUNDARY_NONE
    DM_BOUNDARY_GHOSTED
    DM_BOUNDARY_MIRROR
    DM_BOUNDARY_PERIODIC
    DM_BOUNDARY_TWIST
end

@enum DMStagStencilType begin
    DMSTAG_STENCIL_NONE
    DMSTAG_STENCIL_STAR
    DMSTAG_STENCIL_BOX
end

@enum PetscBool PETSC_FALSE PETSC_TRUE

@enum PetscDataType begin
    PETSC_DATATYPE_UNKNOWN  = 0
    PETSC_DOUBLE            = 1
    PETSC_COMPLEX           = 2
    PETSC_LONG              = 3
    PETSC_SHORT             = 4
    PETSC_FLOAT             = 5
    PETSC_CHAR              = 6
    PETSC_BIT_LOGICAL       = 7
    PETSC_ENUM              = 8
    PETSC_BOOL              = 9
    PETSC_FLOAT128          = 10
    PETSC_OBJECT            = 11
    PETSC_FUNCTION          = 12
    PETSC_STRING            = 13
    PETSC___FP16            = 14
    PETSC_STRUCT            = 15
    PETSC_INT               = 16
    PETSC_INT64             = 17
end

const Petsc64bitInt       = Int64
const PetscLogDouble      = Cdouble

@enum InsertMode NOT_SET_VALUES INSERT_VALUES ADD_VALUES MAX_VALUES MIN_VALUES INSERT_ALL_VALUES ADD_ALL_VALUES INSERT_BC_VALUES ADD_BC_VALUES

@enum NormType begin
    NORM_1 = 0
    NORM_2 = 1
    NORM_FROBENIUS = 2
    NORM_INFINITY = 3
    NORM_1_AND_2 = 4
end

@enum MatAssemblyType MAT_FLUSH_ASSEMBLY=1 MAT_FINAL_ASSEMBLY=0

@enum MatFactorType begin
    MAT_FACTOR_NONE     = 0
    MAT_FACTOR_LU       = 1
    MAT_FACTOR_CHOLESKY = 2
    MAT_FACTOR_ILU      = 3
    MAT_FACTOR_ICC      = 4
    MAT_FACTOR_ILUDT    = 5
end

@enum PetscViewerFormat begin
    PETSC_VIEWER_DEFAULT
    PETSC_VIEWER_ASCII_MATLAB
    PETSC_VIEWER_ASCII_MATHEMATICA
    PETSC_VIEWER_ASCII_IMPL
    PETSC_VIEWER_ASCII_INFO
    PETSC_VIEWER_ASCII_INFO_DETAIL
    PETSC_VIEWER_ASCII_COMMON
    PETSC_VIEWER_ASCII_SYMMODU
    PETSC_VIEWER_ASCII_INDEX
    PETSC_VIEWER_ASCII_DENSE
    PETSC_VIEWER_ASCII_MATRIXMARKET
    PETSC_VIEWER_ASCII_VTK
    PETSC_VIEWER_ASCII_VTK_CELL
    PETSC_VIEWER_ASCII_VTK_COORDS
    PETSC_VIEWER_ASCII_PCICE
    PETSC_VIEWER_ASCII_PYTHON
    PETSC_VIEWER_ASCII_FACTOR_INFO
    PETSC_VIEWER_ASCII_LATEX
    PETSC_VIEWER_ASCII_XML
    PETSC_VIEWER_ASCII_GLVIS
    PETSC_VIEWER_ASCII_CSV
    PETSC_VIEWER_DRAW_BASIC
    PETSC_VIEWER_DRAW_LG
    PETSC_VIEWER_DRAW_LG_XRANGE
    PETSC_VIEWER_DRAW_CONTOUR
    PETSC_VIEWER_DRAW_PORTS
    PETSC_VIEWER_VTK_VTS
    PETSC_VIEWER_VTK_VTR
    PETSC_VIEWER_VTK_VTU
    PETSC_VIEWER_BINARY_MATLAB
    PETSC_VIEWER_NATIVE
    PETSC_VIEWER_HDF5_PETSC
    PETSC_VIEWER_HDF5_VIZ
    PETSC_VIEWER_HDF5_XDMF
    PETSC_VIEWER_HDF5_MAT
    PETSC_VIEWER_NOFORMAT
    PETSC_VIEWER_LOAD_BALANCE
end

@enum MatOperation begin
    MATOP_SET_VALUES=0
    MATOP_GET_ROW=1
    MATOP_RESTORE_ROW=2
    MATOP_MULT=3
    MATOP_MULT_ADD=4
    MATOP_MULT_TRANSPOSE=5
    MATOP_MULT_TRANSPOSE_ADD=6
    MATOP_SOLVE=7
    MATOP_SOLVE_ADD=8
    MATOP_SOLVE_TRANSPOSE=9
    MATOP_SOLVE_TRANSPOSE_ADD=10
end

@enum DMStagStencilLocation begin
    DMSTAG_NULL_LOCATION=0
    DMSTAG_BACK_DOWN_LEFT
    DMSTAG_BACK_DOWN
    DMSTAG_BACK_DOWN_RIGHT
    DMSTAG_BACK_LEFT
    DMSTAG_BACK
    DMSTAG_BACK_RIGHT
    DMSTAG_BACK_UP_LEFT
    DMSTAG_BACK_UP
    DMSTAG_BACK_UP_RIGHT
    DMSTAG_DOWN_LEFT
    DMSTAG_DOWN
    DMSTAG_DOWN_RIGHT
    DMSTAG_LEFT
    DMSTAG_ELEMENT
    DMSTAG_RIGHT
    DMSTAG_UP_LEFT
    DMSTAG_UP
    DMSTAG_UP_RIGHT
    DMSTAG_FRONT_DOWN_LEFT
    DMSTAG_FRONT_DOWN
    DMSTAG_FRONT_DOWN_RIGHT
    DMSTAG_FRONT_LEFT
    DMSTAG_FRONT
    DMSTAG_FRONT_RIGHT
    DMSTAG_FRONT_UP_LEFT
    DMSTAG_FRONT_UP
    DMSTAG_FRONT_UP_RIGHT
end

@enum DMBoundaryType begin
  DM_BOUNDARY_NONE=0
  DM_BOUNDARY_GHOSTED=1
  DM_BOUNDARY_MIRROR=2
  DM_BOUNDARY_PERIODIC=3
  DM_BOUNDARY_TWIST=4
end

@enum DMDAStencilType begin
  DMDA_STENCIL_STAR = 0
  DMDA_STENCIL_BOX = 1
end


push!(LOAD_PATH, "@stdlib")
import Pkg; Pkg.add("Conda");
using Conda
using ADCME

@info "Install Boost"
CONDA = get_conda()
run(`$CONDA install boost`)

@info "Install AMGCL"
UNZIP = joinpath(ADCME.BINDIR, "unzip")
if !isdir("$(@__DIR__)/amgcl")
    download("https://github.com/ddemidov/amgcl/archive/master.zip", "$(@__DIR__)/amgcl.zip")
    run(`$UNZIP -o $(@__DIR__)/amgcl.zip -d $(@__DIR__)`)
    mv("$(@__DIR__)/amgcl-master","$(@__DIR__)/amgcl", force=true)
    rm("$(@__DIR__)/amgcl.zip")
end

@info "Build Custom Operators"
change_directory("CustomOps/build")
require_file("build.ninja") do
    ADCME.cmake()
end
ADCME.make()


################################################################################
function init(::Type{Custom}, config::ConfigCustom, objective::Objective, difftune)
    return Custom(ModelWrappers.ℓDensityResult(objective), config.proposal)
end

function infer(
    _rng::Random.AbstractRNG,
    diagnostics::Type{AbstractDiagnostics},
    kernel::Custom,
    model::ModelWrapper,
    data::D,
) where {D}
    return DiagnosticsCustom
end

################################################################################
# Export
export infer, init


export iauZp
"""
**  Zero a p-vector.
**
**  This function is part of the International Astronomical Union's
**  SOFA (Standards Of Fundamental Astronomy) software collection.
**
**  Status:  vector/matrix support function.
**
**  Returned:
**     p        double[3]      p-vector
**
**  This revision:  2013 June 18
**
**  SOFA release 2018-01-30
**
**  Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# void iauZp(double p[3])

function iauZp()
   p = zeros(Float64, 3)

   ccall((:iauZp, libsofa_c), Cvoid,
           (Ptr{Cdouble},),
           p)

   # Transpose since C call return row-major operation
   return SVector{3}(p)
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    RefinementMethod

A method for refining meshes.
"""
abstract type RefinementMethod end

"""
    refine(mesh, method)

Refine `mesh` with refinement `method`.
"""
function refine end

# ----------------
# IMPLEMENTATIONS
# ----------------

include("refinement/quad.jl")
include("refinement/catmullclark.jl")