module SpiceUtils

    using Downloads
    using Dates
    using SPICE
    using ModelingToolkit
    using ProgressMeter
    using Pkg.Artifacts
    using ForwardDiff
    using StaticArrays
    using Memoize

    export mass_fraction
    export state_to_synodic, state_from_synodic
    export pos_to_synodic, pos_from_synodic

    const artifacts_toml = find_artifacts_toml(@__DIR__)
    const kernel_bodies = Dict(
        # List of bodies included in each default kernel file (as specified in build.jl)
        "default" => tuple(1:10..., 199, 299, 301, 399),
        "mars"    => tuple(401:402..., 499),
        "jupiter" => tuple(501:505..., 514:516..., 599),
        "saturn"  => tuple(601:609..., 612:614..., 632, 634, 699),
        "uranus"  => tuple(701:705..., 799),
        "neptune" => tuple(801:808..., 814, 899),
        "pluto"   => tuple(901:905..., 999)
    )
    const body_to_kernel = Dict(idx => kernel for (kernel, bodies) in kernel_bodies for idx in bodies)

    function __init__()
        load_kernel_artifact("meta_kernels")
        load_kernel_artifact("default_kernels")
        nothing
    end

    @memoize function load_ephemerides(body_name)
        body_name = lowercase(String(body_name))
        body_ID   = bodn2c(body_name)

        if body_ID ∉ keys(body_to_kernel)
            @warn """No default kernel specified for "$(titlecase(body_name))".
            You may need to manually download relevant kernels and load them with `SPICE.furnsh(path_to_kernel)`."""
            return false
        end

        artifact_name = "$(body_to_kernel[body_ID])_kernels"
        return load_kernel_artifact(artifact_name)
    end

    @memoize function load_kernel_artifact(artifact_name)
        # XXX: This is a bit of a hacky workaround of the existing Pkg.Artifacts system,
        # which supports lazy artifact loading but assumes the files it downloads are
        # gzipped. However, our kernels are just files straight from the NAIF site, so we
        # instead load them lazily here. artifact"(body)_kernels" will only work after this.

        # Get artifact details.
        meta = artifact_meta(artifact_name, artifacts_toml)
        isnothing(meta) && error("Expected to find $(artifact_name) in Artifacts.toml, but could not find it!")

        # Get kernel details from this artifact.
        meta_hash     = Base.SHA1(meta["git-tree-sha1"])
        download_URLs = [d["url"] for d in meta["download"]]
        kernel_dir    = artifact_path(meta_hash)
        kernel_names  = [basename(url) for url in download_URLs]

        # Check if the kernel has already been downloaded before, and if not, download it to a temporary directory.
        temp_dir = artifact_exists(meta_hash) ? nothing : mktempdir(first(Artifacts.artifacts_dirs()))
        if !isnothing(temp_dir)
            max_n = 10000
            for (kernel, url) in zip(kernel_names, download_URLs)
                bar = Progress(max_n; desc="Downloading NAIF kernel: $(kernel)", color=Base.info_color())
                progress = (total, now) -> update!(bar, total > 0 ? round(Int, (now / total) * max_n) : 0)
                try
                    Downloads.download(url, joinpath(temp_dir, kernel); progress)
                finally
                    finish!(bar)
                end
            end

            # All kernels are now downloaded; move them into the final kernel folder.
            mv(temp_dir, kernel_dir)
        end

        # Load each kernel into SPICE.
        for (kernel, url) in zip(kernel_names, download_URLs)
            kernel_path = joinpath(kernel_dir, kernel)
            if !isfile(kernel_path)
                error("Could not find kernel file '$(kernel)'. Delete the $(kernel_dir) directory and try again.")
            end
            furnsh(kernel_path)
        end

        return true
    end

    @doc "Get a target body position from SPICE kernels."
    get_pos(t, target, reference; relative="ECLIPJ2000") = SVector{3,Float64}(spkpos(String(target), t, relative, "none", String(reference))[1])
    get_pos(t::ForwardDiff.Dual, target, reference; kwargs...) = get_pos(ForwardDiff.value(t), String(target), String(reference); kwargs...)
    get_pos(t::Num, tgt, ref) = get_pos(t, Num(tgt), Num(ref))
    ModelingToolkit.@register get_pos(t, target, reference)

    @doc "Get a target body state from SPICE kernels"
    get_state(t, target, reference; relative="ECLIPJ2000") = SVector{6,Float64}(spkezr(String(target), t, relative, "none", String(reference))[1])
    get_state(t::ForwardDiff.Dual, target, reference; kwargs...) = get_state(ForwardDiff.value(t), target, reference; kwargs...)

    @doc "Get body gravitational constant (GM) from SPICE kernels"
    get_GM(body::String) = bodvrd(body, "GM")[1]
    get_GM(body::Symbol) = get_GM(String(body))
    get_GM(body::Number) = get_GM(bodc2n(body))

    @doc "Get the mass fraction of two bodies (smaller mass divided by total mass)"
    mass_fraction(b1::Symbol, b2::Symbol) = mass_fraction(String(b1), String(b2))
    function mass_fraction(body_1::String, body_2::String)
        mass_1 = bodvrd(body_1, "GM")[1]
        mass_2 = bodvrd(body_2, "GM")[1]
        min(mass_1, mass_2) ./ (mass_1 .+ mass_2)
    end

    # Dynamic frame ID should be between 1400000 to 2000000 (added a bit of margin for rand) as per
    # https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/req/frames.html#Selecting%20a%20Frame%20ID
    const INIT_DYNAMIC_FRAME_IDX = rand(1400000:1900000)
    const DYNAMIC_FRAMES_TYPE = Dict{Tuple{Symbol,Symbol},Tuple{Int,String}}
    DYNAMIC_FRAMES = DYNAMIC_FRAMES_TYPE()

    @doc """
    Automatically build and return a dynamic frame for a synodic
    reference (b1 is the observer, b2 is the target).
    """
    function dynamic_synodic_frame(b1::Symbol, b2::Symbol; reference="ECLIPJ2000")
        if (b1, b2) ∉ keys(DYNAMIC_FRAMES::DYNAMIC_FRAMES_TYPE)
            # XXX: Maximum size of the frame name is limited so we use the body numbers.
            bodies = bodn2c.((b1, b2))
            frame_name = "$(bodies[1])_$(bodies[2])_SYNODIC"
            path, io = mktemp()
            ID = INIT_DYNAMIC_FRAME_IDX + length(DYNAMIC_FRAMES::DYNAMIC_FRAMES_TYPE)
            write(io,
            """ \\begindata

                FRAME_$(frame_name)        = $(ID)
                FRAME_$(ID)_NAME           = '$(frame_name)'
                FRAME_$(ID)_CLASS          = 5
                FRAME_$(ID)_CLASS_ID       = $(ID)
                FRAME_$(ID)_CENTER         = $(bodies[1])
                FRAME_$(ID)_RELATIVE       = '$(reference)'
                FRAME_$(ID)_DEF_STYLE      = 'PARAMETERIZED'
                FRAME_$(ID)_ROTATION_STATE = 'ROTATING'
                FRAME_$(ID)_FAMILY         = 'TWO-VECTOR'
                FRAME_$(ID)_PRI_AXIS       = 'X'
                FRAME_$(ID)_PRI_VECTOR_DEF = 'OBSERVER_TARGET_POSITION'
                FRAME_$(ID)_PRI_OBSERVER   = $(bodies[1])
                FRAME_$(ID)_PRI_TARGET     = $(bodies[2])
                FRAME_$(ID)_PRI_ABCORR     = 'NONE'
                FRAME_$(ID)_SEC_AXIS       = 'Y'
                FRAME_$(ID)_SEC_VECTOR_DEF = 'OBSERVER_TARGET_VELOCITY'
                FRAME_$(ID)_SEC_OBSERVER   = $(bodies[1])
                FRAME_$(ID)_SEC_TARGET     = $(bodies[2])
                FRAME_$(ID)_SEC_FRAME      = '$(reference)'
                FRAME_$(ID)_SEC_ABCORR     = 'NONE'
            """)
            close(io)

            (DYNAMIC_FRAMES::DYNAMIC_FRAMES_TYPE)[(b1, b2)] = (ID, frame_name)
            @info "Importing new $frame_name kernel from $(path)"
            furnsh(path)
        else
            ID, frame_name = (DYNAMIC_FRAMES::DYNAMIC_FRAMES_TYPE)[(b1, b2)]
        end
        return frame_name
    end

    # TODO: Fix up the SpiceUtils interface so the arguments are more consistently ordered.
    state_to_synodic(b1, b2, t; reference="ECLIPJ2000") = sxform(reference, dynamic_synodic_frame(b1, b2; reference=reference), t)
    state_to_synodic(b1, b2, t::ForwardDiff.Dual; kwargs...) = state_to_synodic(b1, b2, ForwardDiff.value(t); kwargs...)
    state_from_synodic(args...; kwargs...) = inv(state_to_synodic(args...; kwargs...))

    pos_to_synodic(b1, b2, t; reference="ECLIPJ2000") = pxform(reference, dynamic_synodic_frame(b1, b2; reference=reference), t)
    pos_to_synodic(b1, b2, t::ForwardDiff.Dual; kwargs...) = pos_to_synodic(b1, b2, ForwardDiff.value(t); kwargs...)
    pos_from_synodic(args...; kwargs...) = inv(pos_to_synodic(args...; kwargs...))

end


module Grids

export Center, Face
export AbstractTopology, Periodic, Bounded, Flat, FullyConnected, LeftConnected, RightConnected, topology

export AbstractGrid, AbstractUnderlyingGrid, halo_size, total_size
export AbstractRectilinearGrid, RectilinearGrid
export XRegRectilinearGrid, YRegRectilinearGrid, ZRegRectilinearGrid, HRegRectilinearGrid, RegRectilinearGrid
export AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid
export LatitudeLongitudeGrid, XRegLatLonGrid, YRegLatLonGrid, ZRegLatLonGrid
export ConformalCubedSphereFaceGrid, ConformalCubedSphereGrid
export node, xnode, ynode, znode, xnodes, ynodes, znodes, nodes
export offset_data, new_data
export on_architecture

using CUDA
using Adapt
using OffsetArrays

using Oceananigans
using Oceananigans.Architectures

import Base: size, length, eltype, show
import Oceananigans.Architectures: architecture

#####
##### Abstract types
#####

"""
    Center

A type describing the location at the center of a grid cell.
"""
struct Center end

"""
	Face

A type describing the location at the face of a grid cell.
"""
struct Face end

"""
    AbstractTopology

Abstract supertype for grid topologies.
"""
abstract type AbstractTopology end

"""
    Periodic

Grid topology for periodic dimensions.
"""
struct Periodic <: AbstractTopology end

"""
    Bounded

Grid topology for bounded dimensions, e.g., wall-bounded dimensions.
"""
struct Bounded <: AbstractTopology end

"""
    Flat

Grid topology for flat dimensions, generally with one grid point, along which the solution
is uniform and does not vary.
"""
struct Flat <: AbstractTopology end

"""
    FullyConnected

Grid topology for dimensions that are connected to other models or domains.
"""
struct FullyConnected <: AbstractTopology end

"""
    LeftConnected

Grid topology for dimensions that are connected to other models or domains only on the left (the other direction is bounded)
"""
struct LeftConnected <: AbstractTopology end

"""
    RightConnected

Grid topology for dimensions that are connected to other models or domains only on the right (the other direction is bounded)
"""
struct RightConnected <: AbstractTopology end

"""
    AbstractGrid{FT, TX, TY, TZ}

Abstract supertype for grids with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractGrid{FT, TX, TY, TZ, Arch} end

"""
    AbstractUnderlyingGrid{FT, TX, TY, TZ}

Abstract supertype for "primary" grids (as opposed to grids with immersed boundaries)
with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractUnderlyingGrid{FT, TX, TY, TZ, Arch} <: AbstractGrid{FT, TX, TY, TZ, Arch} end

"""
    AbstractRectilinearGrid{FT, TX, TY, TZ}

Abstract supertype for rectilinear grids with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractRectilinearGrid{FT, TX, TY, TZ, Arch} <: AbstractUnderlyingGrid{FT, TX, TY, TZ, Arch} end

"""
    AbstractCurvilinearGrid{FT, TX, TY, TZ}

Abstract supertype for curvilinear grids with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractCurvilinearGrid{FT, TX, TY, TZ, Arch} <: AbstractUnderlyingGrid{FT, TX, TY, TZ, Arch} end

"""
    AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ}

Abstract supertype for horizontally-curvilinear grids with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch} <: AbstractCurvilinearGrid{FT, TX, TY, TZ, Arch} end

isrectilinear(grid) = false

include("grid_utils.jl")
include("zeros.jl")
include("new_data.jl")
include("inactive_node.jl")
include("automatic_halo_sizing.jl")
include("input_validation.jl")
include("grid_generation.jl")
include("rectilinear_grid.jl")
include("conformal_cubed_sphere_face_grid.jl")
include("latitude_longitude_grid.jl")

end # module


using KCenters, KNearestCenters, TextSearch, TextClassification, SearchModels, UnicodePlots, StatsBase, Statistics
using LIBLINEAR, JLD2

function score_by_name(score, y, ypred)
    if score === :macrorecall
        recall_score(y, ypred, weight=:macro)
    elseif score === :macrof1
        f1_score(y, ypred, weight=:macro)
    elseif score === :accuracy
        accuracy_score(y, ypred)
    elseif score isa NamedTuple
        if score.name == :classf1
            classification_scores(y, ypred).classf1[score.label]
        elseif score.name == :classrecall
            classification_scores(y, ypred).classrecall[score.label]
        else
            error("Unknown score $score")
        end
    else
        error("Unknown score $score")
    end
end

function search_models_for_task(D, paramsfile, space, search_kwargs; verbose=true)
    cv = pop!(search_kwargs, :cv)

    if cv.kind == :folds
        partitions = stratified_folds(D.labels; folds=cv.folds, shuffle=true)
    elseif cv.kind == :montecarlo
        partitions = [stratified_split(D.labels; p=cv.ratio, shuffle=true) for i in 1:cv.repeat]
    elseif cv.kind == :holdout
        partitions = [stratified_split(D.labels; p=cv.ratio, shuffle=true)]
    else
        error("Unknown $cv")
    end

    score = pop!(search_kwargs, :score, :recall)
    initialpopulation = pop!(search_kwargs, :initialpopulation, 32)
    function error_function(config::MicroTC_Config)
        S = Float64[]
        for (itrain, itest) in partitions
            tc = MicroTC(config, D.corpus[itrain], D.labels[itrain]; verbose=true)
            ypred = predict_corpus(tc, D.corpus[itest])
            push!(S, score_by_name(score, D.labels[itest], ypred))
        end
        s = mean(S)
        verbose && println(stderr, "score: $s, $(typeof(config)), config: $(JSON3.write(config))")
        (1.0 - s, S)
    end

    iter = 0

    function inspect_population(space, params, population)
        iter += 1
        println(stderr, "===== inspecting population iter: $iter, popsize: $(length(population))")
        filename = paramsfile * ".iter=$iter"
        println(stderr, "saving iteration to $filename")
        jldsave(filename; population)
        println(lineplot([first(last(p)) for p in population], title="Error", xlabel="configurations", ylabel="error"))
    end

    search_kwargs[:inspect_population] = inspect_population
    best_list = search_models(space, error_function, initialpopulation; by=first, search_kwargs...)

    for (i, b) in enumerate(best_list)
        @info "-- perf best_lists[$i]:", b[1] => b[2]
    end

    best_list
end

function run_params(D, paramsfile, space, search_kwargs)
    if !isfile(paramsfile)
        population = search_models_for_task(D, paramsfile, space, search_kwargs)
        jldsave(paramsfile; population)
        population
    else
        load(paramsfile, "population")
    end
end

function run_train(D, config, modelfile; tok=Tokenizer(config.textconfig, invmap=nothing))
    if !isfile(modelfile)
        model = MicroTC(config, D.corpus, D.labels; tok)
        jldsave(modelfile; model)
        model
    else
        load(modelfile, "model")
    end
end

function run_predict(T, model, predictedfile; textkey="text", labelkey="klass", levels_=levels(T.labels))
    ypred = predict_corpus(model, T.corpus)
    open(predictedfile, "w") do f
        for (message, ilabel) in zip(T.corpus, ypred)
            println(f, JSON3.write(Dict(textkey => message, labelkey => levels_[ilabel])))
        end
    end

    ypred
end


"""
    function check_par(par)
---
## Description
Check if enough and right.
"""
function check_par(par)
    @info "Check if the simulation parameters are valid"
    must_keys = Dict(
        :nSire => "Number of sires",
        :nDam  => "Number of dams",
        :nSib  => "Number of sibs",
        :nC7e  => "Number of sibs for challenge(c7e)",
        :nG8n  => "Number of generation(g8n) to breed",
        :h²    => "Heritability for traits",
        :w4t   => "Weight(w4t) for the binary trait",
        :cv    => "Covariance of pleiotropic QTL effects for 2 traits",
        :p17h  => "Percentage-of-death(p17h) after challenge",
        :e19e  => "edit-successful-rate(e19e), [0, 1]",
        :nK3n  => "Number of known(k3n) QTL of the binary trait",
        :dd    => "An ϵ to be added to diag(LHS matrix) to avoid singular",
        :edit  => "Whether to edit known QTL",
        :fix   => "Whether to emphasize known QTL",
        :b4y   => "Whether the binary(b4y) trait to be recorded as binary",
        :u38y  => "use-genetypes-of-current-generation-only(u38y)"
    )
    missingkey = []
    for key in keys(must_keys)
        haskey(par, key) || push!(missingkey, key)
    end
    isempty(missingkey) || begin
        for key in missingkey
            println("Key $key \"$(must_keys[key])\" is missing")
        end
        error("Some must-have key(s) is/are missing")
    end
    haskey(par, :b4y) && !par.b4y &&
        @warn "Continuous disease phenotype was for debugging" par.b4y
    # More to be added
end


################################################################################
### Mandatory Interface
################################################################################

### Model
########################################

"""
    hopping_matrix(mc::DQMC, m::Model)

Calculates the hopping matrix \$T_{i\\sigma, j\\sigma '}\$ where \$i, j\$ are
site indices and \$\\sigma , \\sigma '\$ are flavor indices (e.g. spin indices).
The hopping matrix should also contain potential chemical potential terms on the
diagonal.

A matrix element is the hopping amplitude for a hopping process: \$j,\\sigma '
\\rightarrow i,\\sigma\$.

Regarding the order of indices, if `T[i, σ, j, σ']` is your desired 4D hopping
array, then `reshape(T, (n_sites * n_flavors, :))` is the hopping matrix.
"""
hopping_matrix(mc::DQMC, m::Model) = throw(MethodError(hopping_matrix, (mc, m)))

nflavors(m::Model) = throw(MethodError(nflavors, (m,)))

### Field
########################################

"""
    interaction_matrix_exp!(
        mc::DQMC, m::Model, field::AbstractField,
        result::AbstractArray, slice::Int, power::Float64 = 1.0
    )

Calculate the, exponentiated interaction matrix
`exp(- power * delta_tau * V(slice))` and store it in `result::AbstractArray`.

This only includes terms with 4 operators, i.e. not the chemical potential or
any hoppings. By default the calculation will be performed by the appropriate
field type (i.e. by `interaction_matrix_exp!(field, result, slice, power)`)

By default this function will call
`interaction_matrix_exp!(field, result, slice, power)`.

This is a performance critical method and one might consider efficient in-place
(in `result`) construction.
"""
@inline function interaction_matrix_exp!(
        mc, model, field, result, slice, power = +1.0
    )
    interaction_matrix_exp!(field, result, slice, power)
end
function interaction_matrix_exp!(f, A, s, p)
    throw(MethodError(interaction_matrix_exp!, (f, A, s, p)))
end

"""
    propose_local(mc::DQMC, m::Model, field::AbstractField, i::Int, slice::Int)

Propose a local move for lattice site `i` at time slice `slice` for a `field`
holding the current configuration. Returns the Green's function determinant
ratio, the boson energy difference `ΔE_boson = E_boson_new - E_boson`,
and any extra information `passthrough` that might be useful in `accept_local`.

By default this function will call `propose_local(mc, field, i, slice)`.

See also [`accept_local!`](@ref).
"""
@inline propose_local(mc, m, field, i, slice) = propose_local(mc, field, i, slice)
propose_local(mc, f, i, s) = throw(MethodError(propose_local, (mc, f, i, s)))

"""
    accept_local!(
        mc::DQMC, m::Model, field::AbstractField, i::Int, slice::Int,
        detratio, ΔE_boson, passthrough
    )

Accept a local move for site `i` at imaginary time slice `slice` of current
configuration in `field`. Arguments `detratio`, `ΔE_boson` and `passthrough`
correspond to output of `propose_local` for that local move.

By default this function will call
`accept_local!(mc, field, i, slice, detration, ΔE_boson, passthrough)`

See also [`propose_local`](@ref).
"""
@inline function accept_local!(
        mc, m, field, i, slice, detratio, ΔE_boson, passthrough
    )
    accept_local!(mc, field, i, slice, detratio, ΔE_boson, passthrough)
end
function accept_local!(mc, f, i, s, d, ΔE, pt)
    throw(MethodError(accept_local!, (mc, f, i, s, d, ΔE, pt)))
end

Base.rand(f::AbstractField) = throw(MethodError(rand, (f,)))
Random.rand!(f::AbstractField) = throw(MethodError(rand!, (f,)))

nflavors(f::AbstractField) = throw(MethodError(nflavors, (f,)))

# For ConfigRecorder
compress(f::AbstractField) = throw(MethodError(compress, (f, )))
compressed_conf_type(f::AbstractField) = throw(MethodError(compressed_conf_type, (f, )))
decompress(f::AbstractField, c) = throw(MethodError(decompress, (f, c)))
decompress!(f::AbstractField, c) = throw(MethodError(decompress!, (f, c)))

################################################################################
### Optional Interface
################################################################################

### Mixed (Field + Model)
########################################

"""
    greens_eltype(field, model)

Returns the type of the elements of the Green's function matrix. Defaults to
Float64 if both the hopping and interaction matrix contain floats and ComplexF64
otherwise.
"""
function greens_eltype(field::AbstractField, model::Model)
    generalized_eltype(interaction_eltype(field), hopping_eltype(model))
end

"""
    greens_matrix_type(field, model)

Returns the (matrix) type of the greens and most work matrices. Defaults to
`Matrix{greens_eltype(T, m)}`.
"""
greens_matrix_type(f::AbstractField, m::Model) = Matrix{greens_eltype(f, m)}

### Model
########################################

"""
    hopping_eltype(model)

Returns the type of the elements of the hopping matrix. Defaults to `Float64`.
"""
hopping_eltype(::Model) = Float64

"""
    hopping_matrix_type(field, model)

Returns the (matrix) type of the hopping matrix. Defaults to
`Matrix{hopping_eltype(model)}`.
"""
hopping_matrix_type(::AbstractField, m::Model) = Matrix{hopping_eltype(m)}

### Field
########################################

"""
    interaction_eltype(model)

Returns the type of the elements of the interaction matrix. Defaults to `Float64`.
"""
interaction_eltype(::AbstractField) = Float64

"""
    interaction_matrix_type(field, model)

Returns the (matrix) type of the interaction matrix. Defaults to
`Matrix{interaction_eltype(model)}`.
"""
interaction_matrix_type(f::AbstractField, m::Model) = Matrix{interaction_eltype(f)}

"""
    init_interaction_matrix(field::AbstractField, model::Model)

Returns an initial interaction matrix. This only used to allocate a correctly
sized matrix.

By default this uses the matrix type from `interaction_matrix_type` and uses
`max(nflavors(field), nflavors(model)) * length(lattice(model))` as the size.
"""
function init_interaction_matrix(f::AbstractField, m::Model)
    flv = max(nflavors(f), nflavors(m))
    N = length(lattice(m))
    FullT = interaction_matrix_type(f, m)

    if FullT <: BlockDiagonal
        MT = matrix_type(interaction_eltype(f))
        return BlockDiagonal(MT(undef, N, N), MT(undef, N, N))
    elseif FullT <: Diagonal
        VT = vector_type(interaction_eltype(f))
        Diagonal(VT(undef, flv * N))
    else
        FullT(undef, flv * N, flv * N)
    end
end

"""
    energy_boson(mc::DQMC, model::Model, conf)

Calculate bosonic part (non-Green's function determinant part) of energy for
configuration `conf` for Model `m`.

This is required for global and parallel updates as well as boson energy
measurements, but not for local updates. By default calls
`energy_boson(field(mc), conf)`
"""
energy_boson(mc::DQMC, m::Model, c = nothing) = energy_boson(field(mc), c)
energy_boson(f::AbstractField, c = nothing) = throw(MethodError(energy_boson, (f, c)))

conf(f::AbstractField) = f.conf
conf!(f::AbstractField, c) = conf(f) .= c
temp_conf(f::AbstractField) = f.temp_conf


using SALTBase
using AbbreviatedTypes
using Test

# Define a concrete subtype of AbstractMaxwellSolver for the test purspose.
struct DefaultMXS <: AbstractMaxwellSolver end
Base.similar(::DefaultMXS) = DefaultMXS()
Base.size(::DefaultMXS) = (0,0)

function SALTBase.init_mxsolver!(mxsolver::DefaultMXS, ω::Number, ε::AbsVecNumber) end
function SALTBase.mxsolve!(x::AbsVecComplexF, mxsolver::DefaultMXS, b::AbsVecNumber) end
function SALTBase.mxsolve_transpose!(x::AbsVecComplexF, mxsolver::DefaultMXS, b::AbsVecNumber) end
function SALTBase.mxapply!(b::AbsVecComplexF, mxsolver::DefaultMXS, x::AbsVecNumber) end
function SALTBase.abs2_interp!(abs2ψ::AbsVecReal, mxsolver::DefaultMXS, ψ::AbsVecNumber) end

include("util.jl")
include("gain.jl")
include("lasing.jl")


module TestExamples2DEuler

using Test
using Trixi

include("test_trixi.jl")

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "tree_2d_dgsem")

@testset "Compressible Euler" begin
  @trixi_testset "elixir_euler_source_terms.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_source_terms.jl"),
      l2   = [8.517808508019351e-7, 1.2350203856098537e-6, 1.2350203856728076e-6, 4.277886946638239e-6],
      linf = [8.357848139128876e-6, 1.0326302096741458e-5, 1.0326302101404394e-5, 4.496194024383726e-5])
  end

  @trixi_testset "elixir_euler_convergence_pure_fv.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_convergence_pure_fv.jl"),
      l2   = [0.026440292358506527, 0.013245905852168414, 0.013245905852168479, 0.03912520302609374],
      linf = [0.042130817806361964, 0.022685499230187034, 0.022685499230187922, 0.06999771202145322])
  end

  @trixi_testset "elixir_euler_density_wave.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_density_wave.jl"),
      l2   = [0.0010600778457964775, 0.00010600778457634275, 0.00021201556915872665, 2.650194614399671e-5],
      linf = [0.006614198043413566, 0.0006614198043973507, 0.001322839608837334, 0.000165354951256802],
      tspan = (0.0, 0.5))
  end

  @trixi_testset "elixir_euler_source_terms_nonperiodic.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_source_terms_nonperiodic.jl"),
      l2   = [2.3653424742011065e-6, 2.1388875092652046e-6, 2.1388875093792834e-6, 6.010896863407163e-6],
      linf = [1.4080465934984687e-5, 1.7579850582816192e-5, 1.757985059525069e-5, 5.95689353266593e-5])
  end

  @trixi_testset "elixir_euler_ec.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_ec.jl"),
      l2   = [0.061751715597716854, 0.05018223615408711, 0.05018989446443463, 0.225871559730513],
      linf = [0.29347582879608825, 0.31081249232844693, 0.3107380389947736, 1.0540358049885143])
  end

  @trixi_testset "elixir_euler_ec.jl with flux_kennedy_gruber" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_ec.jl"),
      l2   = [0.03481471610306124, 0.027694280613944234, 0.027697905866996532, 0.12932052501462554],
      linf = [0.31052098400669004, 0.3481295959664616, 0.34807152194137336, 1.1044947556170719],
      maxiters = 10,
      surface_flux = flux_kennedy_gruber,
      volume_flux = flux_kennedy_gruber)
  end

  @trixi_testset "elixir_euler_ec.jl with flux_chandrashekar" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_ec.jl"),
      l2   = [0.03481122603050542, 0.027662840593087695, 0.027665658732350273, 0.12927455860656786],
      linf = [0.3110089578739834, 0.34888111987218107, 0.3488278669826813, 1.1056349046774305],
      maxiters = 10,
      surface_flux = flux_chandrashekar,
      volume_flux = flux_chandrashekar)
  end

  @trixi_testset "elixir_euler_shockcapturing.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_shockcapturing.jl"),
      l2   = [0.05380378209319121, 0.04696482040740498, 0.046967440141927635, 0.19686385916646665],
      linf = [0.1852693811520424, 0.24028641582658528, 0.2326441454434102, 0.6874069967414047])
  end

  @trixi_testset "elixir_euler_blast_wave.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_blast_wave.jl"),
      l2   = [0.1399704878290245, 0.11580885255188444, 0.115808885313645, 0.3391969791404581],
      linf = [1.4473441276982646, 1.3355435317505784, 1.335543531751034, 1.8203831482774346],
      maxiters = 30)
  end

  @trixi_testset "elixir_euler_blast_wave_neuralnetwork_perssonperaire.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_blast_wave_neuralnetwork_perssonperaire.jl"),
      l2   = [0.470793846459279, 0.20772743006225688, 0.20772649286847356, 0.7011946441199758],
      linf = [1.459980543578883, 0.976674580177999, 0.9759214032734075, 2.960406757135902],
      initial_refinement_level = 4,
      maxiters = 50)
  end

  @trixi_testset "elixir_euler_blast_wave_neuralnetwork_rayhesthaven.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_blast_wave_neuralnetwork_rayhesthaven.jl"),
      l2   = [0.4667685563568254, 0.20611874363505156, 0.20588110990842964, 0.6993986008683889],
      linf = [1.4753582207406488, 0.9904209156187197, 0.9841399353882521, 2.9346283218997806],
      initial_refinement_level = 4,
      maxiters = 50)
  end

  @trixi_testset "elixir_euler_blast_wave_neuralnetwork_rayhesthaven.jl with mortars" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_blast_wave_neuralnetwork_rayhesthaven.jl"),
      l2   = [0.016398689755743465, 0.03074685501652015, 0.03078660152836722, 0.15176702545248866],
      linf = [0.27612812614937343, 0.5598351620373528, 0.5601415330040513, 2.48372589562021],
      refinement_patches=(
          (type="box", coordinates_min=(-0.25, -0.25), coordinates_max=(0.25, 0.25)),
          (type="box", coordinates_min=(-0.125, -0.125), coordinates_max=(0.125, 0.125)),),
      initial_refinement_level = 4,
      maxiters = 5)
  end

  @trixi_testset "elixir_euler_blast_wave_neuralnetwork_cnn.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_blast_wave_neuralnetwork_cnn.jl"),
      l2   = [0.47425390570612547, 0.20990014031106777, 0.2103622732820483, 0.7022475294762982],
      linf = [1.6366874508175182, 0.9734578279971777, 0.9412921347930632, 2.949655726363399],
      initial_refinement_level = 4,
      maxiters = 50)
  end

  @trixi_testset "elixir_euler_blast_wave_pure_fv.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_blast_wave_pure_fv.jl"),
      l2   = [0.39957047631960346, 0.21006912294983154, 0.21006903549932, 0.6280328163981136],
      linf = [2.20417889887697, 1.5487238480003327, 1.5486788679247812, 2.4656795949035857],
      tspan = (0.0, 0.5))
  end

  @trixi_testset "elixir_euler_blast_wave_amr.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_blast_wave_amr.jl"),
      l2   = [0.6777351803934385, 0.28143110338521693, 0.28143035507881886, 0.7211949078634373],
      linf = [2.881942488461436, 1.8037294505391677, 1.803677581960154, 3.037915614619938],
      tspan = (0.0, 1.0))
  end

  @trixi_testset "elixir_euler_sedov_blast_wave.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_sedov_blast_wave.jl"),
      l2   = [0.4819231474652106, 0.1655453843075888, 0.16554538430758944, 0.6182451347074112],
      linf = [2.4849541392059202, 1.2816207838889486, 1.2816207838888944, 6.4743208586529875],
      tspan = (0.0, 1.0))
  end

  @trixi_testset "elixir_euler_sedov_blast_wave_neuralnetwork_perssonperaire.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_sedov_blast_wave_neuralnetwork_perssonperaire.jl"),
      l2   = [8.29516192e-02, 9.18870398e-02, 9.18870398e-02, 4.37449170e-01],
      linf = [1.23402442e+00, 1.63254552e+00, 1.63254552e+00, 4.21977821e+00],
      maxiters = 30)
  end

  @trixi_testset "elixir_euler_positivity.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_positivity.jl"),
      l2   = [0.48401054859424775, 0.16613328471131006, 0.16613328471131006, 0.6181793382171535],
      linf = [2.4903327485761544, 1.2898210694161085, 1.2898210694161072, 6.4723873993158385],
      tspan = (0.0, 1.0))
  end

  @trixi_testset "elixir_euler_blob_mortar.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_blob_mortar.jl"),
      l2   = [0.22206706482833133, 0.6296849068200803, 0.24351238064525166, 2.9118777552679345],
      linf = [9.584083314003589, 27.72635365994101, 9.343044146443756, 101.53747277352629],
      tspan = (0.0, 0.5))
  end

  @trixi_testset "elixir_euler_blob_amr.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_blob_amr.jl"),
      l2   = [0.20150296718390992, 1.1816675092370863, 0.10129652123574924, 5.230357225596295],
      linf = [14.174958287930844, 71.16876942972455, 7.2799739019809415, 291.67608863630244],
      tspan = (0.0, 0.12))
  end

  @trixi_testset "elixir_euler_kelvin_helmholtz_instability.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_kelvin_helmholtz_instability.jl"),
      l2   = [0.055689867638859906, 0.032985369421826734, 0.052243738479375024, 0.08009259073134713],
      linf = [0.24051443859674326, 0.16611572070980668, 0.123559477474734, 0.2694484247345663],
      tspan = (0.0, 0.2))
  end

  @trixi_testset "elixir_euler_kelvin_helmholtz_instability_amr.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_kelvin_helmholtz_instability_amr.jl"),
      l2   = [0.05569140599245201, 0.033108757924402134, 0.05223526059785845, 0.08007239985228143],
      linf = [0.2540392117514094, 0.17432468822204936, 0.12323401271048683, 0.26897998287356195],
      tspan = (0.0, 0.2))
  end

  @trixi_testset "elixir_euler_kelvin_helmholtz_instability_amr_neuralnetwork_perssonperaire.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_kelvin_helmholtz_instability_amr_neuralnetwork_perssonperaire.jl"),
      l2   = [0.0009932162527962933, 0.004946122604416242, 0.004863174253792848, 0.005019630958200973],
      linf = [0.008603757244273202, 0.02090493316272357, 0.01712529102069524, 0.028966987799588395],
      maxiters = 30)
  end

  @trixi_testset "elixir_euler_colliding_flow.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_colliding_flow.jl"),
      l2   = [7.23705116e-03,   4.48861301e-02,   8.47273157e-07,   6.62667920e-01],
      linf = [1.93754293e-01,   5.52253475e-01,   4.98055004e-05,   1.50879674e+01],
      tspan = (0.0, 0.1))
  end

  @trixi_testset "elixir_euler_colliding_flow_amr.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_colliding_flow_amr.jl"),
      l2   = [6.76960244e-03,   3.21735698e-02,   2.63316434e-07,   6.78416591e-01],
      linf = [2.50119529e-01,   4.07493507e-01,   9.96993324e-05,   2.23204482e+01],
      tspan = (0.0, 0.1))
  end

  @trixi_testset "elixir_euler_astro_jet_amr.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_astro_jet_amr.jl"),
      l2   = [0.011338493575859507, 10.097552800226566, 0.0038670992882799088, 4031.828021246965],
      linf = [3.316670548834158, 2992.6286682433883, 7.919515633718346, 1.1914747265545786e6],
      tspan = (0.0, 1.0e-7))
  end

  @trixi_testset "elixir_euler_vortex.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_vortex.jl"),
      l2   = [3.6343141788303172e-6, 0.0032111379945554378, 0.0032111482803927763, 0.004545715899533244],
      linf = [7.901851921288117e-5, 0.030561510550305093, 0.03050266851613559, 0.042876690674344076])
  end

  @trixi_testset "elixir_euler_vortex_mortar.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_vortex_mortar.jl"),
      l2   = [2.1203067597905396e-6, 2.792922844344066e-5, 3.7593422075434476e-5, 8.813644239472934e-5],
      linf = [5.932046017209647e-5, 0.0007491265252463908, 0.0008165690091972433, 0.0022122638482677814])
  end

  @trixi_testset "elixir_euler_vortex_mortar_split.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_vortex_mortar_split.jl"),
      l2   = [2.1201025902895562e-6, 2.8056524989293713e-5, 3.759500393896144e-5, 8.84137216637438e-5],
      linf = [5.934103282212444e-5, 0.0007552316670187409, 0.0008152449093751235, 0.0022069874183294758])
  end

  @trixi_testset "elixir_euler_vortex_shockcapturing.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_vortex_shockcapturing.jl"),
      l2   = [3.7998676433961623e-6, 5.563687203373343e-5, 5.561628017766063e-5, 0.00015707075214260475],
      linf = [8.522682141076654e-5, 0.0009610342345369727, 0.0009656468947373265, 0.0030861481615822584])
  end

  @trixi_testset "elixir_euler_vortex_mortar_shockcapturing.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_vortex_mortar_shockcapturing.jl"),
      l2   = [2.1203212494737245e-6, 2.8056721454951553e-5, 3.759510818832351e-5, 8.842116438932004e-5],
      linf = [5.9340649220973596e-5, 0.0007552376287922602, 0.0008152432005811283, 0.0022069702233196153])
  end

  @trixi_testset "elixir_euler_vortex_amr.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_vortex_amr.jl"),
      l2   = [2.1325938318497577e-6, 0.003281545066125044, 0.0032806321457482615, 0.00464605975938406],
      linf = [4.667437348149228e-5, 0.03175420871507906, 0.0318039789241531, 0.04561735256198318])
  end

  @trixi_testset "elixir_euler_ec.jl with boundary_condition_slip_wall" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_ec.jl"),
      l2   = [0.06176845600430613, 0.05020624043492084, 0.05021389111189423, 0.22592682624517807],
      linf = [0.29347582879609024, 0.3108124923286465, 0.3107380389949771, 1.054035804988522],
      periodicity = false, boundary_conditions = boundary_condition_slip_wall)
  end
end

# Coverage test for all initial conditions
@testset "Compressible Euler: Tests for initial conditions" begin
  @trixi_testset "elixir_euler_vortex.jl one step with initial_condition_density_pulse" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_vortex.jl"),
      l2   = [0.003489659044164644, 0.0034896590441646494, 0.0034896590441646502, 0.003489659044164646],
      linf = [0.04761180654650543, 0.04761180654650565, 0.047611806546505875, 0.04761180654650454],
      maxiters = 1,
      initial_condition = initial_condition_density_pulse)
  end

  @trixi_testset "elixir_euler_vortex.jl one step with initial_condition_pressure_pulse" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_vortex.jl"),
      l2   = [0.00021747693728234874, 0.0022010142997830533, 0.0022010142997830485, 0.010855273768135729],
      linf = [0.005451116856088789, 0.03126448432601536, 0.03126448432601536, 0.14844305553724624],
      maxiters = 1,
      initial_condition = Trixi.initial_condition_pressure_pulse)
  end

  @trixi_testset "elixir_euler_vortex.jl one step with initial_condition_density_pressure_pulse" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_vortex.jl"),
      l2   = [0.003473649182284682, 0.005490887132955628, 0.005490887132955635, 0.015625074774949926],
      linf = [0.046582178207169145, 0.07332265196082899, 0.07332265196082921, 0.2107979471941368],
      maxiters = 1,
      initial_condition = Trixi.initial_condition_density_pressure_pulse)
  end

  @trixi_testset "elixir_euler_vortex.jl one step with initial_condition_constant" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_vortex.jl"),
      l2   = [7.89034964747135e-17, 8.095575651413758e-17, 1.0847287658433571e-16, 1.2897732640029767e-15],
      linf = [2.220446049250313e-16, 3.191891195797325e-16, 4.163336342344337e-16, 3.552713678800501e-15],
      maxiters = 1,
      initial_condition = initial_condition_constant)
  end

  @trixi_testset "elixir_euler_sedov_blast_wave.jl one step" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_sedov_blast_wave.jl"),
      l2   = [0.0021037031798961914, 0.010667428589443025, 0.01066742858944302, 0.10738893384136498],
      linf = [0.11854059147646778, 0.7407961272348982, 0.7407961272348981, 3.92623931433345],
      maxiters=1)
  end

  @trixi_testset "elixir_euler_sedov_blast_wave.jl one step with initial_condition_medium_sedov_blast_wave" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_sedov_blast_wave.jl"),
      l2   = [0.0021025532272874827, 0.010661548568022292, 0.010661548568022284, 0.10734939168392313],
      linf = [0.11848345578926645, 0.7404217490990809, 0.7404217490990809, 3.9247328712525973],
      maxiters=1, initial_condition=initial_condition_medium_sedov_blast_wave)
  end
end

end # module


############################################################################################
# Basic Functionality
_modelProb = ModelWrapper(ProbModel(), val_dist)
@testset "Models - basic functionality" begin
    ## Type Check 1 - Constrain/Unconstrain
    theta_unconstrained_vec = randn(length(_modelProb))
    theta_unconstrained = unflatten(_modelProb, theta_unconstrained_vec)
    @test typeof(theta_unconstrained) == typeof(_modelProb.val)
    theta_constrained = constrain(_modelProb.info.transform, theta_unconstrained)
    theta_constrained2 = unflatten_constrain(_modelProb, theta_unconstrained_vec)
    @test typeof(theta_constrained) == typeof(_modelProb.val)
    @test typeof(theta_constrained2) == typeof(_modelProb.val)
    ## Type Check 2 - Flatten/Unflatten
    _θ1, _ = flatten(_modelProb.info.reconstruct, theta_constrained)
    _θ2, _ = flatten(_modelProb.info.reconstruct, theta_constrained2)
    @test sum(abs.(_θ1 - _θ2)) ≈ 0 atol = _TOL
    ## Check if densities match
    @test log_prior(_modelProb) + log_abs_det_jac(_modelProb) ≈
          log_prior_with_transform(_modelProb)
    ## Utility functions
    unconstrain(_modelProb)
    flatten(_modelProb)
    unconstrain_flatten(_modelProb)
    simulate(_modelProb)
    fill(_modelProb, _modelProb.val)
    fill!(_modelProb, _modelProb.val)
    subset(_modelProb, keys(_modelProb.val))
    unflatten!(_modelProb, flatten(_modelProb))
    unflatten_constrain!(_modelProb, unconstrain_flatten(_modelProb))
    sample(_RNG, _modelProb)
    sample(_modelProb)
    sample!(_RNG, _modelProb)
    sample!(_modelProb)
end

############################################################################################
## Model with transforms in lower dimensions

_modelExample = ModelWrapper(ExampleModel(), _val_examplemodel)
_tagged = Tagged(_modelExample)
@testset "Models - Model with transforms in lower dimensions" begin
    ## Model Length accounting discrete parameter
    unconstrain(_modelExample)
    flatten(_modelExample)
    unconstrain_flatten(_modelExample)
    ## Type Check 1 - Constrain/Unconstrain
    theta_unconstrained_vec = randn(length(_modelExample))
    theta_unconstrained = unflatten(_modelExample, theta_unconstrained_vec)
    @test typeof(theta_unconstrained) == typeof(_modelExample.val)
    theta_constrained = constrain(_modelExample.info.transform, theta_unconstrained)
    theta_constrained2 = unflatten_constrain(_modelExample, theta_unconstrained_vec)
    @test typeof(theta_constrained) == typeof(_modelExample.val)
    @test typeof(theta_constrained2) == typeof(_modelExample.val)
    ## Type Check 2 - Flatten/Unflatten
    _θ1, _ = flatten(_modelExample.info.reconstruct, theta_constrained)
    _θ2, _ = flatten(_modelExample.info.reconstruct, theta_constrained2)
    @test sum(abs.(_θ1 - _θ2)) ≈ 0 atol = _TOL
    ## Check if densities match
    @test log_prior(_modelExample) + log_abs_det_jac(_modelExample) ≈
          log_prior_with_transform(_modelExample)
    ## Check utility functions
    @test length(_modelExample) == 23
    @test ModelWrappers.paramnames(_modelExample) == keys(_val_examplemodel)
    fill(_modelExample, _tagged, _modelExample.val)
    fill!(_modelExample, _tagged, _modelExample.val)
end

############################################################################################


struct InitialInvestments{R,G}
    options::Matrix{Int} # Existing (buildable) options (r x g)
    builds::Matrix{Int}  # Existing (dispatchable) units

    function InitialInvestments{}(options::Matrix, builds::Matrix)
        R, G = size(options)
        @assert (R, G) == size(builds)
        new{R,G}(options, builds)
    end

end

struct InitialConditions{R,G1,G2,G3}
    thermal::InitialInvestments{R,G1}
    variable::InitialInvestments{R,G2}
    storage::InitialInvestments{R,G3}
end

mutable struct ResourceInvestments{R,G}

    # Unit state machine:
    # --> vesting --> building --> retired
    #            \      ^     \       ^
    #             \     |      \      |
    #              > holding    > dispatching

    # Parameters

    # Capital costs are considered sunk and so modeled as one-time expenses,
    # although in terms of cash flow they may actually be amortized over time
    optioncost::Matrix{Float64}     # ($/unit, r x g)
    buildcost::Matrix{Float64}
    retirementcost::Matrix{Float64}

    optionleadtime::Matrix{Int}     # investment periods, r x g
    buildleadtime::Matrix{Int}

    newoptionslimit::Matrix{Int} # investment periods, r x g
    newbuildslimit::Matrix{Int}
    newforcedretirements::Matrix{Int}

    # Variables

    newoptions::Matrix{VariableRef} # New options purchased (r x g)
    newholding::Matrix{VariableRef} # Hold a newly-vested option
    newbuilds_vesting::Matrix{VariableRef}  # Excercise a newly-vested option
    newbuilds_holding::Matrix{VariableRef} # Start construction on held project
    newdispatching::Matrix{VariableRef} # Start dispatching newly-constructed units
    newretirements_dispatching::Matrix{VariableRef}  # Retire a dispatchable unit
    newretirements_building::Matrix{VariableRef}  # Retire a newly-constructed unit

    # Expressions

    optionsvested::Matrix{ExpressionRef} # Newly-vested options (r x g)
    buildsstarted::Matrix{ExpressionRef}  # Newly-started construction
    buildsfinished::Matrix{ExpressionRef}  # Newly-completed construction
    newretirements::Matrix{ExpressionRef}  # New retirements

    vesting::Matrix{ExpressionRef} # Options vesting (r x g)
    holding::Matrix{ExpressionRef} # Units buildable (option vested but not used)
    building::Matrix{ExpressionRef} # Units under construction (option used)
    dispatching::Matrix{ExpressionRef} # Units dispatchable (finished construction)
    retired::Matrix{ExpressionRef} # Units retired

    totalforcedretirements::Matrix{ExpressionRef} # Minimum cumulative units retired

    ucap::Matrix{ExpressionRef} # Unforced dispatchable capacity (MW, r x g)

    investmentcosts::Matrix{ExpressionRef} # $, r x g

    # Constraints

    minnewoptions::Matrix{GreaterThanConstraintRef} # r x g
    maxnewoptions::Matrix{LessThanConstraintRef}

    minnewholding::Matrix{GreaterThanConstraintRef}

    minnewbuilds_vesting::Matrix{GreaterThanConstraintRef}
    minnewbuilds_holding::Matrix{GreaterThanConstraintRef}
    maxnewbuilds::Matrix{LessThanConstraintRef}

    minnewdispatching::Matrix{GreaterThanConstraintRef}

    minnewretirements_dispatching::Matrix{GreaterThanConstraintRef}
    minnewretirements_building::Matrix{GreaterThanConstraintRef}

    optionsvested_balance::Matrix{EqualToConstraintRef}
    buildsfinished_balance::Matrix{EqualToConstraintRef}

    minvesting::Matrix{GreaterThanConstraintRef}
    minholding::Matrix{GreaterThanConstraintRef}
    minbuilding::Matrix{GreaterThanConstraintRef}
    mindispatching::Matrix{GreaterThanConstraintRef}
    minretired::Matrix{GreaterThanConstraintRef}

    mintotalretirements::Matrix{GreaterThanConstraintRef}

    function ResourceInvestments{}(
        optioncost::Matrix{Float64}, buildcost::Matrix{Float64},
        retirementcost::Matrix{Float64},
        optionleadtime::Matrix{Int}, buildleadtime::Matrix{Int},
        maxnewoptions::Matrix{Int}, maxnewbuilds::Matrix{Int},
        forcedretirements::Matrix{Int})

        R, G = size(maxnewoptions)
        @assert size(maxnewbuilds) == (R,G)

        @assert size(optioncost) == (R,G)
        @assert size(buildcost) == (R,G)
        @assert size(retirementcost) == (R,G)

        @assert size(optionleadtime) == (R,G)
        @assert size(buildleadtime) == (R,G)

        @assert size(maxnewoptions) == (R,G)

        new{R,G}(optioncost, buildcost, retirementcost,
                 optionleadtime, buildleadtime,
                 maxnewoptions, maxnewbuilds, forcedretirements)

    end

end

function ResourceInvestments(
    tech::AbstractTechnology{G},
    regionlookup::Dict{String,Int},
    resourcetypefolder::String
) where G

    R = length(regionlookup)
    techdata = DataFrame!(CSV.File(joinpath(resourcetypefolder, "parameters.csv"),
                                   types=scenarios_resource_param_types))
    genlookup = Dict(zip(tech.name, 1:G))

    optioncost = zeros(Float64, R, G)
    buildcost = zeros(Float64, R, G)
    retirementcost = zeros(Float64, R, G)

    optionleadtime = zeros(Int, R, G)
    buildleadtime = zeros(Int, R, G)

    newoptionslimit = zeros(Int, R, G)
    newbuildslimit = zeros(Int, R, G)
    newforcedretirements = zeros(Int, R, G)

    for row in eachrow(techdata)

        region_idx = regionlookup[row.region]
        gen_idx = genlookup[row.class]

        optioncost[region_idx, gen_idx] = row.optioncost
        buildcost[region_idx, gen_idx] = row.buildcost
        retirementcost[region_idx, gen_idx] = row.retirementcost

        optionleadtime[region_idx, gen_idx] = row.optionleadtime
        buildleadtime[region_idx, gen_idx] = row.buildleadtime

        newoptionslimit[region_idx, gen_idx] = row.newoptionslimit
        newbuildslimit[region_idx, gen_idx] = row.newbuildslimit
        newforcedretirements[region_idx, gen_idx] = row.newforcedretirements

    end

    return ResourceInvestments(
        optioncost, buildcost, retirementcost,
        optionleadtime, buildleadtime,
        newoptionslimit, newbuildslimit, newforcedretirements)

end

function setup!(
    s::AbstractScenario,
    invtype::Symbol,
    m::Model,
    initconds::InitialInvestments{R,G}
) where {R, G}

    invprob = s.investmentproblem
    Rs = invprob.regionnames
    Gs = getfield(invprob.technologies, invtype).name

    invs = getfield(s.investments, invtype)

    # Variables

    invs.newoptions = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newoptions, "newoptions_$(invtype)_$(s.name)", Rs, Gs)

    invs.newholding = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newholding, "newholding_$(invtype)_$(s.name)", Rs, Gs)

    invs.newbuilds_vesting = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newbuilds_vesting, "newbuilds_vesting_$(invtype)_$(s.name)", Rs, Gs)

    invs.newbuilds_holding = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newbuilds_holding, "newbuilds_holding_$(invtype)_$(s.name)", Rs, Gs)

    invs.newdispatching = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newdispatching, "newdispatching_$(invtype)_$(s.name)", Rs, Gs)

    invs.newretirements_dispatching = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newretirements_dispatching, "newretirements_dispatching_$(invtype)_$(s.name)", Rs, Gs)

    invs.newretirements_building = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newretirements_building, "newretirements_building_$(invtype)_$(s.name)", Rs, Gs)

    # Expressions

    invs.optionsvested = @expression(m, [r in 1:R, g in 1:G],
        invs.newholding[r,g] + invs.newbuilds_vesting[r,g])

    invs.buildsstarted = @expression(m, [r in 1:R, g in 1:G],
        invs.newbuilds_holding[r,g] + invs.newbuilds_vesting[r,g])

    invs.buildsfinished = @expression(m, [r in 1:R, g in 1:G],
        invs.newretirements_building[r,g] + invs.newdispatching[r,g])

    invs.newretirements = @expression(m, [r in 1:R, g in 1:G],
        invs.newretirements_dispatching[r,g] + invs.newretirements_building[r,g])

    invs.vesting = @expression(m, [r in 1:R, g in 1:G],
        invs.newoptions[r,g] - invs.optionsvested[r,g])

    invs.holding = @expression(m, [r in 1:R, g in 1:G],
        invs.newholding[r,g] - invs.newbuilds_holding[r,g])

    invs.building = @expression(m, [r in 1:R, g in 1:G],
        invs.buildsstarted[r,g] - invs.buildsfinished[r,g])

    invs.dispatching = @expression(m, [r in 1:R, g in 1:G],
        invs.newdispatching[r,g] - invs.newretirements_dispatching[r,g])

    invs.retired = @expression(m, [r in 1:R, g in 1:G],
        invs.newretirements[r,g])

    invs.totalforcedretirements = @expression(m, [r in 1:R, g in 1:G],
        invs.newforcedretirements[r,g])

    if isnothing(s.parent)
        invs.holding .+= initconds.options
        invs.dispatching .+= initconds.builds
    else
        parentinvs = getfield(s.parent.investments, invtype)
        invs.vesting .+= parentinvs.vesting
        invs.holding .+= parentinvs.holding
        invs.building .+= parentinvs.building
        invs.dispatching .+= parentinvs.dispatching
        invs.retired .+= parentinvs.retired
        invs.totalforcedretirements .+= parentinvs.totalforcedretirements
    end

    invs.investmentcosts =
        @expression(m, [r in 1:R, g in 1:G],
                    invs.optioncost[r,g] * invs.newoptions[r,g] +
                    invs.buildcost[r,g] * invs.buildsstarted[r,g] +
                    invs.retirementcost[r,g] * invs.newretirements[r,g])

    invs.minnewoptions = @constraint(m, [r in 1:R, g in 1:G],
        invs.newoptions[r,g] >= 0)

    invs.minnewholding = @constraint(m, [r in 1:R, g in 1:G],
        invs.newholding[r,g] >= 0)

    invs.minnewbuilds_vesting = @constraint(m, [r in 1:R, g in 1:G],
        invs.newbuilds_vesting[r,g] >= 0)

    invs.minnewbuilds_holding = @constraint(m, [r in 1:R, g in 1:G],
        invs.newbuilds_holding[r,g] >= 0)

    invs.minnewdispatching = @constraint(m, [r in 1:R, g in 1:G],
        invs.newdispatching[r,g] >= 0)

    invs.minnewretirements_dispatching = @constraint(m, [r in 1:R, g in 1:G],
        invs.newretirements_dispatching[r,g] >= 0)

    invs.minnewretirements_building = @constraint(m, [r in 1:R, g in 1:G],
        invs.newretirements_building[r,g] >= 0)

    invs.minvesting = @constraint(m, [r in 1:R, g in 1:G],
        invs.vesting[r,g] >= 0)

    invs.minholding = @constraint(m, [r in 1:R, g in 1:G],
        invs.holding[r,g] >= 0)

    invs.minbuilding = @constraint(m, [r in 1:R, g in 1:G],
        invs.building[r,g] >= 0)

    invs.mindispatching = @constraint(m, [r in 1:R, g in 1:G],
        invs.dispatching[r,g] >= 0)

    invs.minretired = @constraint(m, [r in 1:R, g in 1:G],
        invs.retired[r,g] >= 0)

    invs.maxnewoptions =
        @constraint(m, [r in 1:R, g in 1:G],
                    invs.newoptions[r,g] <= invs.newoptionslimit[r,g])

    invs.maxnewbuilds =
        @constraint(m, [r in 1:R, g in 1:G],
                    invs.buildsstarted[r,g] <= invs.newbuildslimit[r,g])

    invs.optionsvested_balance = @constraint(m, [r in 1:R, g in 1:G],
        invs.optionsvested[r,g] ==
        maturing(s, r, g, invtype, :optionleadtime, :newoptions))

    invs.buildsfinished_balance = @constraint(m, [r in 1:R, g in 1:G],
        invs.buildsfinished[r,g] ==
        maturing(s, r, g, invtype, :buildleadtime, :buildsstarted))

    invs.mintotalretirements =
        @constraint(m, [r in 1:R, g in 1:G],
                    invs.retired[r,g] >= invs.totalforcedretirements[r,g])

    return

end

welfare(x::ResourceInvestments) = -sum(x.investmentcosts)

struct Investments{R,G1,G2,G3}
    thermal::ResourceInvestments{R,G1}
    variable::ResourceInvestments{R,G2}
    storage::ResourceInvestments{R,G3}
end

function loadinvestments(
    techs::Technologies, regions::Dict{String,Int}, resourcepath::String)

    thermal =
        ResourceInvestments(techs.thermal, regions,
                            joinpath(resourcepath, "thermal"))
    variable =
        ResourceInvestments(techs.variable, regions,
                            joinpath(resourcepath, "variable"))
    storage =
        ResourceInvestments(techs.storage, regions,
                            joinpath(resourcepath, "storage"))

    return Investments(thermal, variable, storage)

end

welfare(inv::Investments) =
    welfare(inv.thermal) + welfare(inv.variable) + welfare(inv.storage)


module Shapefile
    import Base: read, show, +, /, ./

    export Handle, RGBGradient, LogLinearRGBGradient, LinearRGBGradient

    type Rect{T}
        top::T
        left::T
        bottom::T
        right::T
    end

    abstract ESRIShape

    type NullShape <: ESRIShape
    end

    type Interval{T}
        left::T
        right::T
    end

    type Point{T} <: ESRIShape
        x::T
        y::T
    end

    +{T}(a::Point{T},b::Point{T}) = Point{T}(a.x+b.x,a.y+b.y)
    /{T}(a::Point{T},val::Real) = Point{T}(a.x/val,a.y/val)
    ./{T}(a::Point{T},val::Real) = Point{T}(a.x./val,a.y./val)

    type PointM{T,M} <: ESRIShape
        x::T
        y::T
        m::M # measure
    end

    type PointZ{T,M} <: ESRIShape
        x::T
        y::T
        z::T
        m::M # measure
    end

    type Polyline{T} <: ESRIShape
        MBR::Rect{T}
        parts::Vector{Int32}
        points::Vector{Point{T}}
    end

    type PolylineM{T,M} <: ESRIShape
        MBR::Rect{T}
        parts::Vector{Int32}
        points::Vector{Point{T}}
        measures::Vector{M}
    end

    type PolylineZ{T,M} <: ESRIShape
        MBR::Rect{T}
        parts::Vector{Int32}
        points::Vector{Point{T}}
        zvalues::Vector{T}
        measures::Vector{M}
    end

    type Polygon{T} <: ESRIShape
        MBR::Rect{T}
        parts::Vector{Int32}
        points::Vector{Point{T}}
    end

    show{T}(io::IO,p::Polygon{T}) = print(io,"Polygon(",length(p.points)," ",T," Points)")

    type PolygonM{T,M} <: ESRIShape
        MBR::Rect{T}
        parts::Vector{Int32}
        points::Vector{Point{T}}
        measures::Vector{M}
    end

    type PolygonZ{T,M} <: ESRIShape
        MBR::Rect{T}
        parts::Vector{Int32}
        points::Vector{Point{T}}
        zvalues::Vector{T}
        measures::Vector{M}
    end

    type MultiPoint{T} <: ESRIShape
        MBR::Rect{T}
        points::Vector{Point{T}}
    end

    type MultiPointM{T,M} <: ESRIShape
        MBR::Rect{T}
        points::Vector{Point{T}}
        measures::Vector{M}
    end

    type MultiPointZ{T,M} <: ESRIShape
        MBR::Rect{T}
        points::Vector{Point{T}}
        zvalues::Vector{T}
        measures::Vector{M}
    end

    type MultiPatch{T,M} <: ESRIShape
        MBR::Rect{T}
        parts::Vector{Int32}
        parttypes::Vector{Int32}
        points::Vector{Point{T}}
        zvalues::Vector{T}
        # measures::Vector{M} # (optional)
    end

    type Handle
        code::Int32
        length::Int32
        version::Int32
        shapeType::Int32
        MBR::Rect{Float64}
        zrange::Interval{Float64}
        mrange::Interval{Float64}
        shapes::Vector{ESRIShape}
    end

    function read{T}(io::IO,::Type{Rect{T}})
        minx = read(io,T)
        miny = read(io,T)
        maxx = read(io,T)
        maxy = read(io,T)
        Rect{T}(miny,minx,maxy,maxx)
    end

    read(io::IO,::Type{NullShape}) = NullShape()

    function read{T}(io::IO,::Type{Point{T}})
        x = read(io,T)
        y = read(io,T)
        Point{T}(x,y)
    end

    function read{T,M}(io::IO,::Type{PointM{T,M}})
        x = read(io,T)
        y = read(io,T)
        m = read(io,M)
        PointM{T,M}(x,y,m)
    end

    function read{T,M}(io::IO,::Type{PointZ{T,M}})
        x = read(io,T)
        y = read(io,T)
        z = read(io,T)
        m = read(io,M)
        PointZ{T,M}(x,y,z,m)
    end

    function read{T}(io::IO,::Type{Polyline{T}})
        box = read(io,Rect{T})
        numparts = read(io,Int32)
        numpoints = read(io,Int32)
        parts = Array(Int32,numparts)
        read!(io, parts)
        points = Array(Point{T},numpoints)
        read!(io, points)
        Polyline{T}(box,parts,points)
    end

    function read{T,M}(io::IO,::Type{PolylineM{T,M}})
        box = read(io,Rect{T})
        numparts = read(io,Int32)
        numpoints = read(io,Int32)
        parts = Array(Int32,numparts)
        read!(io, parts)
        points = Array(Point{T},numpoints)
        read!(io, points)
        mrange = Array(M,2)
        read!(io, mrange)
        measures = Array(M,numpoints)
        read!(io, measures)
        PolylineM{T,M}(box,parts,points,measures)
    end

    function read{T,M}(io::IO,::Type{PolylineZ{T,M}})
        box = read(io,Rect{T})
        numparts = read(io,Int32)
        numpoints = read(io,Int32)
        parts = Array(Int32,numparts)
        read!(io, parts)
        points = Array(Point{T},numpoints)
        read!(io, points)
        zrange = Array(T,2)
        read!(io, zrange)
        zvalues = Array(T,numpoints)
        read!(io, zvalues)
        mrange = Array(M,2)
        read!(io, mrange)
        measures = Array(M,numpoints)
        read!(io, measures)
        PolylineZ{T,M}(box,parts,points,zvalues,measures)
    end

    function read{T}(io::IO,::Type{Polygon{T}})
        box = read(io,Rect{Float64})
        numparts = read(io,Int32)
        numpoints = read(io,Int32)
        parts = Array(Int32,numparts)
        read!(io, parts)
        points = Array(Point{T},numpoints)
        read!(io, points)
        Polygon{T}(box,parts,points)
    end

    function read{T,M}(io::IO,::Type{PolygonM{T,M}})
        box = read(io,Rect{Float64})
        numparts = read(io,Int32)
        numpoints = read(io,Int32)
        parts = Array(Int32,numparts)
        read!(io, parts)
        points = Array(Point{T},numpoints)
        read!(io, points)
        mrange = Array(M,2)
        read!(io, mrange)
        measures = Array(M,numpoints)
        read!(io, measures)
        PolygonM{T,M}(box,parts,points,measures)
    end

    function read{T,M}(io::IO,::Type{PolygonZ{T,M}})
        box = read(io,Rect{Float64})
        numparts = read(io,Int32)
        numpoints = read(io,Int32)
        parts = Array(Int32,numparts)
        read!(io, parts)
        points = Array(Point{T},numpoints)
        read!(io, points)
        zrange = Array(T,2)
        read!(io, zrange)
        zvalues = Array(T,numpoints)
        read!(io, zvalues)
        mrange = Array(M,2)
        read!(io, mrange)
        measures = Array(M,numpoints)
        read!(io, measures)
        PolygonZ{T,M}(box,parts,points,zvalues,measures)
    end

    function read{T}(io::IO,::Type{MultiPoint{T}})
        box = read(io,Rect{Float64})
        numpoints = read(io,Int32)
        points = Array(Point{T},numpoints)
        read!(io, points)
        MultiPoint{T}(box,points)
    end

    function read{T,M}(io::IO,::Type{MultiPointM{T,M}})
        box = read(io,Rect{Float64})
        numpoints = read(io,Int32)
        points = Array(Point{T},numpoints)
        read!(io, points)
        mrange = Array(M,2)
        read!(io, mrange)
        measures = Array(M,numpoints)
        read!(io, measures)
        MultiPointM{T,M}(box,points,measures)
    end

    function read{T,M}(io::IO,::Type{MultiPointZ{T,M}})
        box = read(io,Rect{Float64})
        numpoints = read(io,Int32)
        points = Array(Point{T},numpoints)
        read!(io, points)
        zrange = Array(T,2)
        read!(io, zrange)
        zvalues = Array(T,numpoints)
        read!(io, zvalues)
        mrange = Array(M,2)
        read!(io, mrange)
        measures = Array(M,numpoints)
        read!(io, measures)
        MultiPointZ{T,M}(box,points,zvalues,measures)
    end

    function read{T,M}(io::IO,::Type{MultiPatch{T,M}})
        box = read(io,Rect{Float64})
        numparts = read(io,Int32)
        numpoints = read(io,Int32)
        parts = Array(Int32,numparts)
        read!(io, parts)
        parttypes = Array(Int32,numparts)
        read!(io, parttypes)
        points = Array(Point{T},numpoints)
        read!(io, points)
        zrange = Array(T,2)
        read!(io, zrange)
        zvalues = Array(T,numpoints)
        read!(io, zvalues)
        # mrange = Array(M,2)
        # read!(io, mrange)
        # measures = Array(M,numpoints)
        # read!(io, measures)
        MultiPatch{T,M}(box,parts,parttypes,points,zvalues) #,measures)
    end

    function read(io::IO,::Type{ESRIShape})
        num = bswap(read(io,Int32))
        rlength = bswap(read(io,Int32))
        shapeType = read(io,Int32)
        if(shapeType == 0)
            return read(io,NullShape)
        elseif(shapeType == 1)
            return read(io,Point{Float64})
        elseif(shapeType == 3)
            return read(io,Polyline{Float64})
        elseif(shapeType == 5)
            return read(io,Polygon{Float64})
        elseif(shapeType == 8)
            return read(io,MultiPoint{Float64})
        elseif(shapeType == 11)
            return read(io,PointZ{Float64,Float64})
        elseif(shapeType == 13)
            return read(io,PolylineZ{Float64,Float64})
        elseif(shapeType == 15)
            return read(io,PolygonZ{Float64,Float64})
        elseif(shapeType == 18)
            return read(io,MultiPointZ{Float64,Float64})
        elseif(shapeType == 21)
            return read(io,PointM{Float64,Float64})
        elseif(shapeType == 23)
            return read(io,PolylineM{Float64,Float64})
        elseif(shapeType == 25)
            return read(io,PolygonM{Float64,Float64})
        elseif(shapeType == 28)
            return read(io,MultiPointM{Float64,Float64})
        elseif(shapeType == 31)
            return read(io,MultiPatch{Float64,Float64})
        else
            error("Unknown shape type $shapeType")
        end
    end

    function read(io::IO,::Type{Handle})
        code = bswap(read(io,Int32))
        read(io,Int32,5)
        fileSize = bswap(read(io,Int32))
        version = read(io,Int32)
        shapeType = read(io,Int32)
        MBR = read(io,Rect{Float64})
        zmin = read(io,Float64)
        zmax = read(io,Float64)
        mmin = read(io,Float64)
        mmax = read(io,Float64)
        shapes = Array(ESRIShape,0)
        file = Handle(code,fileSize,version,shapeType,MBR,Interval(zmin,zmax),Interval(mmin,mmax),shapes)
        while(!eof(io))
            push!(shapes,read(io,ESRIShape))
        end
        file
    end

    #If Compose.jl is present, define useful interconversion functions
    isdefined(:Compose) && isa(Compose, Module) && include("compose.jl")
end # module


push!(LOAD_PATH, pwd())
const assets = "assets"
using Images, ImageMagick, Colors, Photopolarimetry
import Tk.ChooseDirectory
#include("Photopolar.jl")
include("PhotopolarGUIfunctions.jl")

path = Input(ChooseDirectory())
flip = Input(false)
flop = Input(false)
torun = Input{Any}(leftbutton)
rotate = Input(0)
cropright = Input(100)
cropbottom = Input(100)
cropleft = Input(0)
croptop = Input(0)
scale = Input(20)
topI = Input(100)
topdolp = Input(100)
topdolcp = Input(100)
topdorcp = Input(100)
botprofile = Input(0)
topprofile = Input(100)
Sp = Input(5)
ellipse_window = Input(0.5)
minpol = Input(0.2)
maxpol = Input(1.0)
eccentricity = Input(0.01)

torun2 = Input{Any}(leftbutton)

fnames = lift(getfnames,path)
SZ = lift(x -> getsz(x[1]),fnames)
lift(x -> write_originalRGB(x[1]),fnames)

row1 = ["Flip", checkbox(false) >>> flip, "Flop", checkbox(false) >>> flop, button("Done!") >>> torun]
row2 = ["Rotate", slider(-90:90,value=0) >>> rotate, "Scale", slider(1:100,value=20) >>> scale]
row3 = ["Crop top", slider(0:99,value=0) >>> croptop, "Crop bottom", slider(1:100,value=100) >>> cropbottom]
row4 = ["Crop left", slider(0:99,value=0) >>> cropleft, "Crop right", slider(1:100,value=100) >>> cropright]

mymin(u,l) = u > l ? l : u - 1
botprofile2 = lift(mymin, topprofile, botprofile)

s = lift(_ -> path2stokes(value(path),value(flip),value(flop),value(rotate),value(cropright),value(cropbottom),value(cropleft),value(croptop),value(scale), angleoffset = deg2rad(60)), torun, typ=Any, init=empty)
polarimgs = lift(calculatepolar, s, typ=Any, init=empty)

I = lift((x,t) -> recolor_I(x[1],t),polarimgs,topI,typ=Any, init=empty)
dolp = lift((x,t) -> recolor_dolp(x[2],t),polarimgs,topdolp,typ=Any, init=empty)
docp = lift((x,b,t) -> recolor_docp(x[4],b,t),polarimgs,topdolcp,topdorcp,typ=Any, init=empty)

profileimg = lift(stokes2profile, s, Sp, typ=Any, init=empty)
ellipse_img_name = lift(latexellipse, s, ellipse_window, maxpol, minpol, eccentricity, typ=Any, init=empty)

orientationtab = vbox(map(pad(0.5em),[
	hbox(map(pad(0.5em), row1)),
	hbox(map(pad(0.5em), row2)),
	hbox(map(pad(0.5em), row3)),
	hbox(map(pad(0.5em), row4)),
	consume(flip,flop,rotate,cropright,cropbottom,cropleft,croptop,scale) do flp, flP, rtt, crprght, crpbttm, crplft, crptp, scl
		@async cleanassets()
		name = updateRGB(flp,flP,rtt,crprght,crpbttm,crplft,crptp,scl)
		image("$assets/$name.jpg")
	end,
	]))

intensitytab = vbox(map(pad(0.5em),[
	hbox("Intensity", slider(0:100, value = 100) >>> topI),
	hbox("DoLP", slider(0:100, value = 100) >>> topdolp),
	hbox("DoLCP", slider(0:100, value = 100) >>> topdolcp),
	hbox("DoRCP", slider(0:100, value = 100) >>> topdorcp),
	consume(I, dolp, docp, polarimgs, typ=Any, init=empty) do i, dlp, dcp, p
		buff = Image(ones(RGB{FixedPointNumbers.UFixed{UInt8,8}},10,size(i,2)))
		cat(1,i, buff, dlp, buff, p[3], buff, dcp)
	end
	]))

profiletab = vbox(map(pad(0.5em),[
	hbox(map(pad(0.5em),["upper", slider(2:100, value = 100) >>> topprofile])),
	consume(topprofile) do x
		hbox(map(pad(0.5em),["Lower", slider(0:(x - 1), value = 0) >>> botprofile]))
	end,
	hbox(map(pad(0.5em),[tex("S_p"), slider(2:10, value = 5) >>> Sp])),
iiimg = lift(recolor_profile, profileimg, botprofile2, topprofile, typ=Any, init=empty)
]))

lift(_ -> println(join([splitdir(value(path))[2],value(flip),value(flop),value(rotate),value(cropright),value(cropbottom),value(cropleft),value(croptop),value(scale),value(ellipse_window), value(minpol), value(maxpol), value(eccentricity)],",")), torun2, typ=Any, init=empty)

ellipsetab = vbox(map(pad(0.5em),[
	hbox(map(pad(0.5em),["Cell size (cm)", slider(0.1:0.05:2, value = 0.5) >>> ellipse_window])),
	hbox(map(pad(0.5em),["Min pol.", slider(0:.01:1, value = 0.2) >>> minpol])),
	hbox(map(pad(0.5em),["Max pol.", slider(0:.01:1, value = 1.0) >>> maxpol])),
	hbox(map(pad(0.5em),["Eccentricity", slider(0:.01:1, value = 0.01) >>> eccentricity])),
	button("Done!") >>> torun2,
	lift(image,ellipse_img_name, typ=Any, init=empty)
]))

#=ok = Input{Any}(leftbutton)
donetab = vbox(map(pad(0.5em),[
	button("Print", raised = true) >>> ok,
	consume(ok, typ=Any, init=empty) do o
		colorimg = load("$assets/RGB.jpg")
		aopimg = convert(Image, value(X)[3])
		Iimg = convert(Image, value(I))
		dolpimg = convert(Image, value(dolp))
		docpimg = convert(Image, value(docp))
		buff = round(Int,0.05*mean(widthheight(colorimg)))
		buffimg = Iimg[1:buff,:]
		buffimg[:] = RGB{U8}(1,1,1)
		img = cat(1,colorimg, buffimg, Iimg, buffimg, dolpimg, buffimg, aopimg, buffimg, docpimg)
		name = string(abs(rand(Int)))
		save( "$assets/$name.jpg",img)
		image("assets/$name.jpg")
	end
	]))=#

function main(window)
	push!(window.assets, "layout2")
	push!(window.assets, "widgets")
	push!(window.assets, "tex")

	tabbar = tabs([ hbox("Orientation"), hbox("Photopolarimetry"), hbox("Interneuron"), hbox("Ellipse")])
	tabcontent = pages([ orientationtab, intensitytab, profiletab, ellipsetab])

	t, p = wire(tabbar, tabcontent, :tab_channel, :selected)
	vbox(t, p)

end


function extract_paths(tour, graph, graph_type)

    X = tour["X"]
    Y = tour["Y"]

    path = DataFrame(
            start_node = Int64[],
            end_node = Int64[],
            nb_passage = Float64[],
            serv = Bool[],
            )

    if nrow(graph["req_arcs"])>0
      for arc in 1:nrow(graph["req_arcs"])
        start_node = graph["req_arcs"][arc, :start_node]
        end_node = graph["req_arcs"][arc, :end_node]
        nb_passage = X["req_arcs"][arc] + Y["req_arcs"][arc]
        serv = X["req_arcs"][arc] >= 1
        if nb_passage > 10^-5
          push!(path, (start_node, end_node, nb_passage, serv))
        end
      end
    end

    if nrow(graph["noreq_arcs"])>0
      for arc in 1:nrow(graph["noreq_arcs"])
        start_node = graph["noreq_arcs"][arc, :start_node]
        end_node = graph["noreq_arcs"][arc, :end_node]
        nb_passage = Y["noreq_arcs"][arc]
        if nb_passage > 10^-5
          push!(path, (start_node, end_node, nb_passage, false))
        end
      end
    end

    if nrow(graph["req_edges"])>0
      for arc in 1:nrow(graph["req_edges"])
        start_node = graph["req_edges"][arc, :start_node]
        end_node = graph["req_edges"][arc, :end_node]
        nb_passage = X["req_edges"][arc] + Y["req_edges"][arc]
        serv = X["req_edges"][arc] >= 1
        if nb_passage > 10^-5
          push!(path, (start_node, end_node, nb_passage, serv))
        end
      end
    end

    if (nrow(graph["req_edges"])>0)&(graph_type=="directed")
      for arc in 1:nrow(graph["req_edges"])
        start_node = graph["req_edges"][arc, :end_node]
        end_node = graph["req_edges"][arc, :start_node]
        nb_passage = X["req_edges_swap"][arc] + Y["req_edges_swap"][arc]
        serv = X["req_edges_swap"][arc] >= 1
        if nb_passage > 10^-5
          push!(path, (start_node, end_node, nb_passage, serv))
        end
      end
    end

    if nrow(graph["noreq_edges"])>0
      for arc in 1:nrow(graph["noreq_edges"])
        start_node = graph["noreq_edges"][arc, :start_node]
        end_node = graph["noreq_edges"][arc, :end_node]
        nb_passage = Y["noreq_edges"][arc]
        if nb_passage > 10^-5
          push!(path, (start_node, end_node, nb_passage, false))
        end
      end
    end

    if (nrow(graph["req_edges"])>0)&(graph_type=="directed")
      for arc in 1:nrow(graph["noreq_edges"])
        start_node = graph["noreq_edges"][arc, :end_node]
        end_node = graph["noreq_edges"][arc, :start_node]
        nb_passage = Y["noreq_edges_swap"][arc]
        if nb_passage > 10^-5
          push!(path, (start_node, end_node, nb_passage, false))
        end
      end
    end

    return(path)
end

function display_results(tours, graph, data, solving_time)

  println("Solving time : ", solving_time)
  println("Total cost : ", sum([tour["cost"] for tour in tours]))
  println("Number of vehicule(s) used : ", length([tour for tour in tours if tour["cost"]!=0]))
  println("--------------")
  for (i,tour) in enumerate(tours)
    if tour["cost"] != 0
      println("Vehicle ", i)
      println("\tCost : ", tour["cost"])
      println("\tCapacity used : ", tour["capacity"])
      if (nrow(graph["req_arcs"])+nrow(graph["noreq_arcs"])==0) & (length(data[1, :depots])==1) & (data[1, :nb_vehicles]==1)
        graph_type = "undirected"
      else
        graph_type = "directed"
      end
      path = extract_paths(tour, graph, graph_type)
      println(path)
    end
  end

end


"""
    UncertainIndexDataset

Generic dataset containing uncertain indices.

## Fields
- **`indices::AbstractVector{AbstractUncertainValue}`**: The uncertain values.
"""
struct UncertainIndexDataset <: AbstractUncertainIndexDataset
    indices::AbstractVector{<:AbstractUncertainValue}
end

"""
    ConstrainedUncertainIndexDataset

Constrained dataset containing uncertain indices.

## Fields
- **`indices::AbstractVector{AbstractUncertainValue}`**: The uncertain indices.
"""
struct ConstrainedUncertainIndexDataset <: AbstractUncertainIndexDataset
    indices::AbstractVector{<:AbstractUncertainValue}
end

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

export
UncertainIndexDataset,
ConstrainedUncertainIndexDataset,
distributions


using TableTransforms
using Distributions
using Tables
using TypedTables
using CategoricalArrays
using ScientificTypes: Count, Multiclass
using LinearAlgebra
using Statistics
using Test, Random, Plots
using ReferenceTests, ImageIO

# set default configurations for plots
gr(ms=2, mc=:black, aspectratio=:equal,
   label=false, size=(600,400))

# workaround GR warnings
ENV["GKSwstype"] = "100"

# environment settings
isCI = "CI" ∈ keys(ENV)
islinux = Sys.islinux()
visualtests = !isCI || (isCI && islinux)
datadir = joinpath(@__DIR__,"data")

# for functor tests in Functional testset
struct Polynomial{T<:Real}
  coeffs::Vector{T}
end
Polynomial(args::T...) where {T<:Real} = Polynomial(collect(args))
(p::Polynomial)(x) = sum(a * x^(i-1) for (i, a) in enumerate(p.coeffs))

function isequalmissing(a, b)
  length(a) == length(b) || return false
  for (x, y) in zip(a, b)
    x === y || return false
  end
  return true
end

# list of tests
testfiles = [
  "distributions.jl",
  "colspec.jl",
  "transforms.jl"
]

@testset "TableTransforms.jl" begin
  for testfile in testfiles
    include(testfile)
  end
end


export softmax_regression

import ..SoftmaxRegression

using mlpack._Internal.io

import mlpack_jll
const softmax_regressionLibrary = mlpack_jll.libmlpack_julia_softmax_regression

# Call the C binding of the mlpack softmax_regression binding.
function softmax_regression_mlpackMain()
  success = ccall((:softmax_regression, softmax_regressionLibrary), Bool, ())
  if !success
    # Throw an exception---false means there was a C++ exception.
    throw(ErrorException("mlpack binding error; see output"))
  end
end

" Internal module to hold utility functions. "
module softmax_regression_internal
  import ..softmax_regressionLibrary

import ...SoftmaxRegression

# Get the value of a model pointer parameter of type SoftmaxRegression.
function IOGetParamSoftmaxRegression(paramName::String)::SoftmaxRegression
  SoftmaxRegression(ccall((:IO_GetParamSoftmaxRegressionPtr, softmax_regressionLibrary), Ptr{Nothing}, (Cstring,), paramName))
end

# Set the value of a model pointer parameter of type SoftmaxRegression.
function IOSetParamSoftmaxRegression(paramName::String, model::SoftmaxRegression)
  ccall((:IO_SetParamSoftmaxRegressionPtr, softmax_regressionLibrary), Nothing, (Cstring, Ptr{Nothing}), paramName, model.ptr)
end

# Serialize a model to the given stream.
function serializeSoftmaxRegression(stream::IO, model::SoftmaxRegression)
  buf_len = UInt[0]
  buf_ptr = ccall((:SerializeSoftmaxRegressionPtr, softmax_regressionLibrary), Ptr{UInt8}, (Ptr{Nothing}, Ptr{UInt}), model.ptr, Base.pointer(buf_len))
  buf = Base.unsafe_wrap(Vector{UInt8}, buf_ptr, buf_len[1]; own=true)
  write(stream, buf)
end
# Deserialize a model from the given stream.
function deserializeSoftmaxRegression(stream::IO)::SoftmaxRegression
  buffer = read(stream)
  SoftmaxRegression(ccall((:DeserializeSoftmaxRegressionPtr, softmax_regressionLibrary), Ptr{Nothing}, (Ptr{UInt8}, UInt), Base.pointer(buffer), length(buffer)))
end
end # module

"""
    softmax_regression(; [input_model, labels, lambda, max_iterations, no_intercept, number_of_classes, test, test_labels, training, verbose])

This program performs softmax regression, a generalization of logistic
regression to the multiclass case, and has support for L2 regularization.  The
program is able to train a model, load  an existing model, and give predictions
(and optionally their accuracy) for test data.

Training a softmax regression model is done by giving a file of training points
with the `training` parameter and their corresponding labels with the `labels`
parameter. The number of classes can be manually specified with the
`number_of_classes` parameter, and the maximum number of iterations of the
L-BFGS optimizer can be specified with the `max_iterations` parameter.  The L2
regularization constant can be specified with the `lambda` parameter and if an
intercept term is not desired in the model, the `no_intercept` parameter can be
specified.

The trained model can be saved with the `output_model` output parameter. If
training is not desired, but only testing is, a model can be loaded with the
`input_model` parameter.  At the current time, a loaded model cannot be trained
further, so specifying both `input_model` and `training` is not allowed.

The program is also able to evaluate a model on test data.  A test dataset can
be specified with the `test` parameter. Class predictions can be saved with the
`predictions` output parameter.  If labels are specified for the test data with
the `test_labels` parameter, then the program will print the accuracy of the
predictions on the given test set and its corresponding labels.

For example, to train a softmax regression model on the data `dataset` with
labels `labels` with a maximum of 1000 iterations for training, saving the
trained model to `sr_model`, the following command can be used:

```julia
julia> using CSV
julia> dataset = CSV.read("dataset.csv")
julia> labels = CSV.read("labels.csv"; type=Int)
julia> sr_model, _ = softmax_regression(labels=labels,
            training=dataset)
```

Then, to use `sr_model` to classify the test points in `test_points`, saving the
output predictions to `predictions`, the following command can be used:

```julia
julia> using CSV
julia> test_points = CSV.read("test_points.csv")
julia> _, predictions = softmax_regression(input_model=sr_model,
            test=test_points)
```

# Arguments

 - `input_model::SoftmaxRegression`: File containing existing model
      (parameters).
 - `labels::Array{Int, 1}`: A matrix containing labels (0 or 1) for the
      points in the training set (y). The labels must order as a row.
 - `lambda::Float64`: L2-regularization constant  Default value `0.0001`.

 - `max_iterations::Int`: Maximum number of iterations before termination.
       Default value `400`.

 - `no_intercept::Bool`: Do not add the intercept term to the model.
      Default value `false`.

 - `number_of_classes::Int`: Number of classes for classification; if
      unspecified (or 0), the number of classes found in the labels will be
      used.  Default value `0`.

 - `test::Array{Float64, 2}`: Matrix containing test dataset.
 - `test_labels::Array{Int, 1}`: Matrix containing test labels.
 - `training::Array{Float64, 2}`: A matrix containing the training set
      (the matrix of predictors, X).
 - `verbose::Bool`: Display informational messages and the full list of
      parameters and timers at the end of execution.  Default value `false`.

# Return values

 - `output_model::SoftmaxRegression`: File to save trained softmax
      regression model to.
 - `predictions::Array{Int, 1}`: Matrix to save predictions for test
      dataset into.

"""
function softmax_regression(;
                            input_model::Union{SoftmaxRegression, Missing} = missing,
                            labels = missing,
                            lambda::Union{Float64, Missing} = missing,
                            max_iterations::Union{Int, Missing} = missing,
                            no_intercept::Union{Bool, Missing} = missing,
                            number_of_classes::Union{Int, Missing} = missing,
                            test = missing,
                            test_labels = missing,
                            training = missing,
                            verbose::Union{Bool, Missing} = missing,
                            points_are_rows::Bool = true)
  # Force the symbols to load.
  ccall((:loadSymbols, softmax_regressionLibrary), Nothing, ());

  IORestoreSettings("Softmax Regression")

  # Process each input argument before calling mlpackMain().
  if !ismissing(input_model)
    softmax_regression_internal.IOSetParamSoftmaxRegression("input_model", convert(SoftmaxRegression, input_model))
  end
  if !ismissing(labels)
    IOSetParamURow("labels", labels)
  end
  if !ismissing(lambda)
    IOSetParam("lambda", convert(Float64, lambda))
  end
  if !ismissing(max_iterations)
    IOSetParam("max_iterations", convert(Int, max_iterations))
  end
  if !ismissing(no_intercept)
    IOSetParam("no_intercept", convert(Bool, no_intercept))
  end
  if !ismissing(number_of_classes)
    IOSetParam("number_of_classes", convert(Int, number_of_classes))
  end
  if !ismissing(test)
    IOSetParamMat("test", test, points_are_rows)
  end
  if !ismissing(test_labels)
    IOSetParamURow("test_labels", test_labels)
  end
  if !ismissing(training)
    IOSetParamMat("training", training, points_are_rows)
  end
  if verbose !== nothing && verbose === true
    IOEnableVerbose()
  else
    IODisableVerbose()
  end

  IOSetPassed("output_model")
  IOSetPassed("predictions")
  # Call the program.
  softmax_regression_mlpackMain()

  return softmax_regression_internal.IOGetParamSoftmaxRegression("output_model"),
         IOGetParamURow("predictions")
end


using .Zygote

Zygote.@adjoint function (ev::Eval)(args...)
    Z = ev.fwd(args...)
    Z, Δ -> begin
        isnothing(ev.rev) && error("no gradient definition here!")
        tuple(nothing, ev.rev(Δ, Z, args...)...)
    end
end

Tullio.promote_storage(T::Type, ::Type{<:Zygote.Fill}) = T
Tullio.promote_storage(::Type{<:Zygote.Fill}, T::Type) = T


export SAGA, SVRG, LSVRG, SVAG, SAG, QSAGA, ILSVRG, ASVAG, SG

export primiterates, dualiterates

################################################################################
# Index Samplings
################################################################################
abstract type IndexSampling end

struct UniformSampling{RNG} <: IndexSampling
	rng::RNG
	n::Int
end
(s::UniformSampling)() = (rand(s.rng, 1:s.n), 1/s.n)

struct WeightedSampling{RNG, CP} <: IndexSampling
	rng::RNG
	p::CP
	cp::CP
	function WeightedSampling(rng, w)
		p = w./sum(w)
		cp = cumsum(p)
		new{typeof(rng), typeof(cp)}(rng, p, cp)
	end
end
function (s::WeightedSampling)()
	u = rand(s.rng)
	i = searchsortedfirst(s.cp, u)
	return i, s.p[i]
end

################################################################################
# Variance Reduced Algorithms
################################################################################
abstract type VRAlgorithm end

update!(alg::VRAlgorithm) = nothing
update!(alg::VRAlgorithm, iter) = update!(alg::VRAlgorithm)
update!(alg::VRAlgorithm, iter, stage) = update!(alg::VRAlgorithm, iter)

stageupdate!(alg::VRAlgorithm) =
	error("stageupdate! need to return the stage length")
stageupdate!(alg::VRAlgorithm, stage) = stageupdate!(alg::VRAlgorithm)
stageupdate!(alg::VRAlgorithm, iter, stage) =
	stageupdate!(alg::VRAlgorithm, stage)

initialize!(alg::VRAlgorithm) = nothing

primiterates(alg::VRAlgorithm) = nothing
dualiterates(alg::VRAlgorithm) = nothing

expectedstagelen(alg::VRAlgorithm) = nothing

##############################
# SAGA
function SAGA(vrg, x0, stepsize, rng; weights=nothing, prox_f=IndFree())
	innoweight = length(vrg)
	SVAG(
		vrg, x0, stepsize, innoweight, rng, weights=weights, prox_f=prox_f)
end

##############################
# SAG
function SAG(vrg, x0, stepsize, rng; weights=nothing, prox_f=IndFree())
	innoweight = 1
	SVAG(
		vrg, x0, stepsize, innoweight, rng, weights=weights, prox_f=prox_f)
end

##############################
# SVAG
struct SVAG{X,VRG,PF,T,S} <: VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	innoweight::T
	sampling::S
end
# Constructor
function SVAG(
		vrg, x0, stepsize, innoweight, rng; weights=nothing, prox_f=IndFree())

	T = eltype(x0)
	stepsize = T(stepsize)
	innoweight = T(innoweight)

	x = similar(x0)
	x_tmp = similar(x0)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)

	SVAG(x,x_tmp,x0,vrg,prox_f,stepsize,innoweight,sampling)
end
# Interface
function initialize!(alg::SVAG)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::SVAG, iter, stage)
	(i,p) = alg.sampling()

	n = length(alg.vrg)
	innoweight_i = alg.innoweight/(n*n*p)
	vrgrad_store!(alg.x_tmp, alg.vrg, alg.x, i, innoweight_i)
	alg.x_tmp .= alg.x .- alg.stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, alg.stepsize)
end
stageupdate!(alg::SVAG) = 1
primiterates(alg::SVAG) = alg.x
dualiterates(alg::SVAG) = alg.vrg
expectedstagelen(alg::SVAG) = 1

##############################
# Adaptive SVAG
struct ASVAG{X,VRG,PF,T,S,F} <: VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	innoweight::Base.RefValue{T}
	innov::X
	expinnov::X
	decay::T
	biascorr::Base.RefValue{T}
	sampling::S
	stepcorrection::F
end
# Constructor
function ASVAG(
		vrg, x0, stepsize, rng;
		weights=nothing, prox_f=IndFree(), enforce_step=false)

	T = eltype(x0)
	stepsize = T(stepsize)
	innoweight = T(length(vrg))

	x = similar(x0)
	x_tmp = similar(x0)

	innov = similar(x0)
	expinnov = zero(x0)
	biascorr = zero(T)
	decay = T(1/10)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)
	if enforce_step
		stepcorrection = (n,theta) -> max(
			2/(2 + abs(n-theta)),
			2/(2 + (n-theta)*(theta-1)/n*((theta-1)/n-1+sign(theta-1)*sqrt(2)) )
			)
	else
		stepcorrection = (n,theta) -> 1
	end
	ASVAG(
		x, x_tmp, x0, vrg, prox_f, stepsize, Ref(innoweight),
		innov, expinnov, decay, Ref(biascorr),
		sampling, stepcorrection
		)
end
# Interface
function initialize!(alg::ASVAG)
	reset!(alg.vrg)
	alg.x .= alg.x0
	alg.innoweight[] = length(alg.vrg)
	alg.expinnov .= zero(eltype(alg.expinnov))
	alg.biascorr[] = zero(alg.biascorr[])
end
function update!(alg::ASVAG, iter, stage)
	(i,p) = alg.sampling()
	n = length(alg.vrg)

	# Update expected innovation
	state = innov!(alg.innov, alg.vrg, alg.x, i)
	alg.expinnov .= alg.innov .+ (1-alg.decay).*alg.expinnov
	alg.biascorr[] = 1 + (1-alg.decay)*alg.biascorr[]

	# Calc innovation weight and make stepsize correction
	alg.innoweight[] =
		n*dot(alg.expinnov, alg.innov) /
		(alg.biascorr[]*norm(alg.innov)^2 + 1e-8)
	lim = n
	alg.innoweight[] = min(lim, alg.innoweight[])
	alg.innoweight[] = max(-lim, alg.innoweight[])

	innoweight_i = alg.innoweight[] / (n*n*p)
	stepsize = alg.stepsize * alg.stepcorrection(n,alg.innoweight[])

	# Take VRSG step
	grad_approx!(alg.x_tmp, alg.vrg)
	alg.x_tmp .= alg.x .- stepsize.*(innoweight_i.*alg.innov .+ alg.x_tmp)
	prox!(alg.x, alg.prox_f, alg.x_tmp, stepsize)
	store_from_innov!(alg.vrg, alg.innov, state)
end
stageupdate!(alg::ASVAG) = 1
primiterates(alg::ASVAG) = alg.x
dualiterates(alg::ASVAG) = alg.vrg
expectedstagelen(alg::ASVAG) = 1

##############################
# SVRG
struct SVRG{X,VRG,PF,T,I,S} <: VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	stagelen::I
	sampling::S
end
# Constructor
function SVRG(vrg, x0, stepsize, stagelen, rng;weights=nothing,prox_f=IndFree())
	T = eltype(x0)
	stepsize = T(stepsize)

	x = similar(x0)
	x_tmp = similar(x0)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)
	SVRG(x,x_tmp,x0,vrg,prox_f,stepsize,stagelen,sampling)
end
# Interface
function initialize!(alg::SVRG)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::SVRG, iter, stage)
	(i,p) = alg.sampling()
	innoweight = 1/(length(alg.vrg)*p)
	vrgrad!(alg.x_tmp, alg.vrg, alg.x, i, innoweight)
	alg.x_tmp .= alg.x .- alg.stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, alg.stepsize)
end
function stageupdate!(alg::SVRG)
	store_from_point!(alg.vrg, alg.x)
	return alg.stagelen
end
primiterates(alg::SVRG) = alg.x
dualiterates(alg::SVRG) = alg.vrg
expectedstagelen(alg::SVRG) = alg.stagelen

##############################
# Loopless SVRG
struct LSVRG{X,VRG,PF,T,Q,RNG,S} <:VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	q::Q
	rng::RNG
	sampling::S
end
# Constructor
function LSVRG(vrg, x0, stepsize, q, rng; weights=nothing,prox_f=IndFree())
	T = eltype(x0)
	stepsize = T(stepsize)

	x = similar(x0)
	x_tmp = similar(x0)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)
	LSVRG(x,x_tmp,x0,vrg,prox_f,stepsize,q,rng,sampling)
end
# Interface
function initialize!(alg::LSVRG)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::LSVRG, iter, stage)
	(i,p) = alg.sampling()

	innoweight = 1/(length(alg.vrg)*p)
	vrgrad!(alg.x_tmp, alg.vrg, alg.x, i, innoweight)

	rand(alg.rng) < alg.q && store_from_point!(alg.vrg, alg.x)

	alg.x_tmp .= alg.x .- alg.stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, alg.stepsize)
end
stageupdate!(alg::LSVRG) = 1
primiterates(alg::LSVRG) = alg.x
dualiterates(alg::LSVRG) = alg.vrg
expectedstagelen(alg::LSVRG) = 1

##############################
# Incoherent Loopless SVRG
struct ILSVRG{X,VRG,PF,T,Q,RNG,S} <:VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	q::Q
	rng::RNG
	sampling::S
end
# Constructor
function ILSVRG(vrg, x0, stepsize, q, rng; weights=nothing,prox_f=IndFree())
	T = eltype(x0)
	stepsize = T(stepsize)

	x = similar(x0)
	x_tmp = similar(x0)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)
	ILSVRG(x,x_tmp,x0,vrg,prox_f,stepsize,q,rng,sampling)
end
# Interface
function initialize!(alg::ILSVRG)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::ILSVRG, iter, stage)
	(i,p) = alg.sampling()
	n = length(alg.vrg)

	innoweight = 1/(n*p)
	vrgrad!(alg.x_tmp, alg.vrg, alg.x, i, innoweight)

	for j = 1:n
		rand(alg.rng) < alg.q && store_from_point!(alg.vrg, alg.x, j)
	end

	alg.x_tmp .= alg.x .- alg.stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, alg.stepsize)
end
stageupdate!(alg::ILSVRG) = 1
primiterates(alg::ILSVRG) = alg.x
dualiterates(alg::ILSVRG) = alg.vrg
expectedstagelen(alg::ILSVRG) = 1

##############################
# qSAGA with or without replacement in the dual update
struct QSAGA{X,VRG,PF,T,RNG,S} <:VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	dualindex::Vector{Int}
	replace::Bool
	rng::RNG
	sampling::S
end
# Constructor
function QSAGA(
		vrg, x0, stepsize, q, rng;
		weights=nothing, prox_f=IndFree(), replace=true)
	T = eltype(x0)
	stepsize = T(stepsize)

	x = similar(x0)
	x_tmp = similar(x0)
	dualindex = Vector{Int}(undef,q)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)
	QSAGA(x,x_tmp,x0,vrg,prox_f,stepsize,dualindex,replace,rng,sampling)
end
# Interface
function initialize!(alg::QSAGA)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::QSAGA, iter, stage)
	n = length(alg.vrg)
	(i,p) = alg.sampling()

	innoweight = 1/(n*p)
	vrgrad!(alg.x_tmp, alg.vrg, alg.x, i, innoweight)

	sample!(alg.rng, 1:n, alg.dualindex, replace=alg.replace)
	store_from_point!(alg.vrg, alg.x, alg.dualindex)

	alg.x_tmp .= alg.x .- alg.stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, alg.stepsize)
end
stageupdate!(alg::QSAGA) = 1
primiterates(alg::QSAGA) = alg.x
dualiterates(alg::QSAGA) = alg.vrg
expectedstagelen(alg::QSAGA) = 1

##############################
# SG, ordinary stoachastic gradient, no variance reduction
struct SG{X,T,VRG,PF,SF,RNG,S} <: VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize_f::SF
	iter::Base.RefValue{T}
	rng::RNG
	sampling::S
end
# Constructor
function SG(vrg, x0, stepsize_f, rng; weights=nothing, prox_f=IndFree())
	iter = zero(eltype(x0))

	x = similar(x0)
	x_tmp = similar(x0)

	sampling = (weights == nothing) ?
		UniformSampling(rng, length(vrg)) :
		WeightedSampling(rng, weights)

	SG(x,x_tmp,x0,vrg,prox_f,stepsize_f,Ref(iter),rng,sampling)
end
# Interface
function initialize!(alg::SG)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::SG, iter, stage)
	(i,p) = alg.sampling()

	stepsize = alg.stepsize_f(alg.iter[])
	alg.iter[] += 1

	grad!(alg.x_tmp, alg.vrg, alg.x, i)
	alg.x_tmp .= alg.x .- stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, stepsize)
end
stageupdate!(alg::SG) = 1
primiterates(alg::SG) = alg.x
dualiterates(alg::SG) = alg.vrg
expectedstagelen(alg::SG) = 1


"""
    AlleeExtinction(minfounders)
    AlleeExtinction(; minfounders=5.0)
    AlleeExtinction{R,W}(minfounders)

Causes extinction in a cell when a population is below the
minimum number of individuals required to maintain it.

Pass grid name `Symbol`s to `R` and `W` type parameters to use specific grids.

$(FIELDDOCTABLE)
"""
@columns struct AlleeExtinction{R,W,MF} <: CellRule{R,W}
    # Field         | Default | Flatten | Limits       | Description
    minfounders::MF | 5.0     | true    | (1.0, 200.0) | "Minimum founding individuals required to to start an ongoing population"
end

@inline applyrule(data, model::AlleeExtinction, state, args...) =
    (state >= model.minfounders ? state : zero(state))