text:
# AffineDecisionFunction #
# ========================== #
MA.mutability(::Type{<:TypedDecisionLike}) = MA.IsMutable()

function MA.mutable_copy(f::AffineDecisionFunction)
    return AffineDecisionFunction(MA.mutable_copy(f.variable_part),
                                  MA.mutable_copy(f.decision_part),
                                  MA.mutable_copy(f.known_part))
end
function MA.mutable_copy(f::LinearPart)
    return LinearPart(
        MA.mutable_copy(f.variable_part),
        MA.mutable_copy(f.decision_part))
end
function MA.mutable_copy(f::QuadraticPart)
    return QuadraticPart(
        MA.mutable_copy(f.variable_part),
        MA.mutable_copy(f.decision_part),
        MA.mutable_copy(f.cross_terms))
end
function MA.mutable_copy(f::QuadraticDecisionFunction)
    return QuadraticDecisionFunction(
        MA.mutable_copy(f.linear_quadratic_terms),
        MA.mutable_copy(f.known_part),
        MA.mutable_copy(f.known_variable_terms),
        MA.mutable_copy(f.known_decision_terms))
end

function MA.isequal_canonical(f::F, g::F) where F<:Union{AffineDecisionFunction, VectorAffineDecisionFunction}
    return MA.isequal_canonical(f.variable_part, g.variable_part) &&
        MA.isequal_canonical(f.decision_part, g.decision_part) &&
        MA.isequal_canonical(f.known_part, g.known_part)
end

function MA.isequal_canonical(f::F, g::F) where F <: LinearPart
    return MA.isequal_canonical(f.variable_part, g.variable_part) &&
        MA.isequal_canonical(f.decision_part, g.decision_part)
end
function MA.isequal_canonical(f::F, g::F) where F <: QuadraticPart
    return MA.isequal_canonical(f.variable_part, g.variable_part) &&
        MA.isequal_canonical(f.decision_part, g.decision_part) &&
        MA.isequal_canonical(f.cross_terms, g.cross_terms)
end
function MA.isequal_canonical(f::F, g::F) where F <: QuadraticDecisionFunction
    return MA.isequal_canonical(f.linear_quadratic_terms, g.linear_quadratic_terms) &&
        MA.isequal_canonical(f.known_part, g.known_part) &&
        MA.isequal_canonical(f.known_variable_terms, g.known_variable_terms) &&
        MA.isequal_canonical(f.known_decision_terms, g.known_decision_terms)
end

function MA.iszero!(f::TypedScalarDecisionLike)
    return iszero!(MOI.constant(f)) && MOIU._is_constant(MOIU.canonicalize!(f))
end

function MA.scaling(f::TypedScalarDecisionLike)
    return MA.scaling(f.variable_part)
end

MA.promote_operation(::Union{typeof(zero), typeof(one)}, F::Type{<:TypedScalarDecisionLike}) = F

const PROMOTE_IMPLEMENTED_OP = Union{typeof(+), typeof(-), typeof(*), typeof(/)}
function MA.promote_operation(op::PROMOTE_IMPLEMENTED_OP,
                              F::Type{<:ScalarDecisionLike{T}},
                              G::Type{<:ScalarDecisionLike{T}}) where T
    MOIU.promote_operation(op, T, F, G)
end
function MA.promote_operation(op::PROMOTE_IMPLEMENTED_OP, F::Type{T},
                              G::Type{<:TypedDecisionLike{T}}) where T
    MOIU.promote_operation(op, T, F, G)
end
function MA.promote_operation(op::PROMOTE_IMPLEMENTED_OP, F::Type{<:TypedDecisionLike{T}},
                              G::Type{T}) where T
    MOIU.promote_operation(op, T, F, G)
end
function MA.promote_operation(op::PROMOTE_IMPLEMENTED_OP, F::Type{<:Number},
                              G::Type{<:Union{SingleDecision, VectorOfDecisions}})
    MOIU.promote_operation(op, F, F, G)
end
function MA.promote_operation(op::PROMOTE_IMPLEMENTED_OP,
                              F::Type{<:Union{SingleDecision, VectorOfDecisions}},
                              G::Type{<:Number})
    MOIU.promote_operation(op, G, F, G)
end

function MA.mutable_operate!(op::Union{typeof(zero), typeof(one)}, f::AffineDecisionFunction)
    MA.mutable_operate!(op, f.variable_part)
    MA.mutable_operate!(zero, f.decision_part)
    MA.mutable_operate!(zero, f.known_part)
    return f
end
function MA.mutable_operate!(op::Union{typeof(zero), typeof(one)}, f::QuadraticDecisionFunction{T,LinearPart{T}}) where T
    lq = f.linear_quadratic_terms
    MA.mutable_operate!(op, lq.variable_part)
    MA.mutable_operate!(zero, lq.decision_part)
    MA.mutable_operate!(zero, f.known_part)
    MA.mutable_operate!(zero, f.known_variable_terms)
    MA.mutable_operate!(zero, f.known_decision_terms)
    return f
end
function MA.mutable_operate!(op::Union{typeof(zero), typeof(one)}, f::QuadraticDecisionFunction{T,QuadraticPart{T}}) where T
    lq = f.linear_quadratic_terms
    MA.mutable_operate!(op, lq.variable_part)
    MA.mutable_operate!(zero, lq.decision_part)
    MA.mutable_operate!(zero, lq.cross_terms)
    MA.mutable_operate!(zero, f.known_part)
    MA.mutable_operate!(zero, f.known_variable_terms)
    MA.mutable_operate!(zero, f.known_decision_terms)
    return f
end

function MA.mutable_operate!(::typeof(-), f::AffineDecisionFunction)
    MA.mutable_operate!(op, f.variable_part)
    MA.mutable_operate!(op, f.decision_part)
    MA.mutable_operate!(op, f.known_part)
    return f
end
function MA.mutable_operate!(::typeof(-), f::QuadraticDecisionFunction{T,LinearPart{T}}) where T
    lq = f.linear_quadratic_terms
    MA.mutable_operate!(op, lq.variable_part)
    MA.mutable_operate!(op, lq.decision_part)
    MA.mutable_operate!(op, f.known_part)
    MA.mutable_operate!(op, f.known_variable_terms)
    MA.mutable_operate!(op, f.known_decision_terms)
    return f
end
function MA.mutable_operate!(::typeof(-), f::QuadraticDecisionFunction{T,QuadraticPart{T}}) where T
    lq = f.linear_quadratic_terms
    MA.mutable_operate!(op, lq.variable_part)
    MA.mutable_operate!(op, lq.decision_part)
    MA.mutable_operate!(op, lq.cross_terms)
    MA.mutable_operate!(op, f.known_part)
    MA.mutable_operate!(op, f.known_variable_terms)
    MA.mutable_operate!(op, f.known_decision_terms)
    return f
end

function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::AffineDecisionFunction{T},
                             g::T) where T
    MA.mutable_operate!(op, f.variable_part, g)
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::AffineDecisionFunction{T},
                             g::MOI.SingleVariable) where T
    _operate_terms!(op, f.variable_part, AffineDecisionFunction{T}(g).variable_part)
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::AffineDecisionFunction{T},
                             g::SingleDecision) where T
    _operate_terms!(op, f.decision_part, AffineDecisionFunction{T}(g).decision_part)
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::AffineDecisionFunction{T},
                             g::SingleKnown) where T
    _operate_terms!(op, f.known_part, AffineDecisionFunction{T}(g).known_part)
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::AffineDecisionFunction{T},
                             g::MOI.ScalarAffineFunction{T}) where T
    _operate_terms!(op, f.variable_part, g)
    f.variable_part.constant = op(f.variable_part.constant, g.constant)
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::AffineDecisionFunction{T},
                             g::AffineDecisionFunction{T}) where T
    # Variable part
    _operate_terms!(op, f.variable_part, g.variable_part)
    f.variable_part.constant = op(f.variable_part.constant, g.variable_part.constant)
    # Decision part
    _operate_terms!(op, f.decision_part, g.decision_part)
    f.decision_part.constant = op(f.decision_part.constant, g.decision_part.constant)
    # Known part
    _operate_terms!(op, f.known_part, g.known_part)
    f.known_part.constant = op(f.known_part.constant, g.known_part.constant)
    return f
end

function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::QuadraticDecisionFunction{T},
                             g::T) where T
    MA.mutable_operate!(op, f.linear_quadratic_terms.variable_part, g)
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::QuadraticDecisionFunction{T},
                             g::MOI.SingleVariable) where T
    _operate_terms!(op, f.linear_quadratic_terms.variable_part, convert(MOI.ScalarAffineFunction{T}, g))
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::QuadraticDecisionFunction{T},
                             g::SingleDecision) where T
    _operate_terms!(op, f.linear_quadratic_terms.decision_part, convert(AffineDecisionFunction{T}, g))
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::QuadraticDecisionFunction{T},
                             g::SingleKnown) where T
    _operate_terms!(op, f.known_part, convert(AffineDecisionFunction{T}, g))
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::QuadraticDecisionFunction{T},
                             g::MOI.ScalarAffineFunction{T}) where T
    lq = f.linear_quadratic_terms
    _operate_terms!(op, lq.variable_part, g)
    lq.variable_part.constant = op(lq.variable_part.constant, g.constant)
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::QuadraticDecisionFunction{T},
                             g::MOI.ScalarQuadraticFunction{T}) where T
    _operate_terms!(op, f.variable_part, g)
    f.variable_part.constant = op(f.variable_part.constant, g.constant)
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::QuadraticDecisionFunction{T},
                             g::AffineDecisionFunction{T}) where T
    lq = f.linear_quadratic_terms
    # Variable_Part
    _operate_terms!(op, lq.variable_part, g.variable_part)
    lq.variable_part.constant = op(lq.variable_part.constant, g.variable_part.constant)
    # Decision part
    _operate_terms!(op, lq.decision_part, g.decision_part)
    lq.decision_part.constant = op(lq.decision_part.constant, lq.decision_part.constant)
    # Known part
    _operate_terms!(op, f.known_part, g.known_part)
    f.known_part.constant = op(f.known_part.constant, g.known_part.constant)
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::QuadraticDecisionFunction{T,LQ},
                             g::QuadraticDecisionFunction{T,LinearPart{T}}) where {T, LQ <: LinearQuadraticPart{T}}
    f_lq = f.linear_quadratic_terms
    g_lq = g.linear_quadratic_terms
    # Variable part
    _operate_terms!(op, f_lq.variable_part, g_lq.variable_part)
    f_lq.variable_part.constant = op(f_lq.variable_part.constant, g_lq.variable_part.constant)
    # Decision part
    _operate_terms!(op, f_lq.decision_part, g_lq.decision_part)
    # Known part
    _operate_terms!(op, f.known_part, g.known_part)
    f.known_part.constant = op(f.known_part.constant, g.known_part.constant)
    _operate_terms!(op, f.known_variable_terms, g.known_variable_terms)
    _operate_terms!(op, f.known_decision_terms, g.known_decision_terms)
    f.known_decision_terms.constant = op(f.known_decision_terms.constant, g.known_decision_terms.constant)
    return f
end
function MA.mutable_operate!(op::Union{typeof(+), typeof(-)},
                             f::QuadraticDecisionFunction{T,QuadraticPart{T}},
                             g::QuadraticDecisionFunction{T,QuadraticPart{T}}) where T
    f_lq = f.linear_quadratic_terms
    g_lq = g.linear_quadratic_terms
    # Variable part
    _operate_terms!(op, f_lq.variable_part, g_lq.variable_part)
    f_lq.variable_part.constant = op(f_lq.variable_part.constant, g_lq.variable_part.constant)
    # Decision part
    _operate_terms!(op, f_lq.decision_part, g_lq.decision_part)
    f_lq.decision_part.constant = op(f_lq.decision_part.constant, g_lq.decision_part.constant)
    # Cross terms
    _operate_terms!(op, f_lq.cross_terms, g_lq.cross_terms)
    # Known part
    _operate_terms!(op, f.known_part, g.known_part)
    f.known_part.constant = op(f.known_part.constant, g.known_part.constant)
    _operate_terms!(op, f.known_variable_terms, g.known_variable_terms)
    _operate_terms!(op, f.known_decision_terms, g.known_decision_terms)
    f.known_decision_terms.constant = op(f.known_decision_terms.constant, g.known_decision_terms.constant)
    return f
end

function _operate_terms!(::typeof(+),
                         f::MOI.ScalarAffineFunction,
                         g::MOI.ScalarAffineFunction)
    for term in g.terms
        add_term!(f.terms, term)
    end
    return nothing
end

function _operate_terms!(::typeof(-),
                         f::MOI.ScalarAffineFunction,
                         g::MOI.ScalarAffineFunction)
    for term in g.terms
        remove_term!(f.terms, term)
    end
    return nothing
end

function _operate_terms!(::typeof(+),
                         f::MOI.ScalarQuadraticFunction,
                         g::MOI.ScalarAffineFunction)
    for term in g.terms
        add_term!(f.affine_terms, term)
    end
    return nothing
end

function _operate_terms!(::typeof(-),
                         f::MOI.ScalarQuadraticFunction,
                         g::MOI.ScalarAffineFunction)
    for term in g.terms
        remove_term!(f.affine_terms, term)
    end
    return nothing
end

function _operate_terms!(::typeof(+),
                         f::MOI.ScalarQuadraticFunction,
                         g::MOI.ScalarQuadraticFunction)
    for term in g.affine_terms
        add_term!(f.affine_terms, term)
    end
    for term in g.quadratic_terms
        add_term!(f.quadratic_terms, term)
    end
    return nothing
end

function _operate_terms!(::typeof(-),
                         f::MOI.ScalarQuadraticFunction,
                         g::MOI.ScalarQuadraticFunction)
    for term in g.affine_terms
        remove_term!(f.affine_terms, term)
    end
    for term in g.quadratic_terms
        remove_term!(f.quadratic_terms, term)
    end
    return nothing
end


text:
using FixedSizeArrays
using GeometryTypes
using GLVisualize

ext: jl


text:
const RealOrMissing = Union{Real,Union{Missing,Real}}

"""
    GeoArray{T::RealOrMissing,A<:AbstractArray{T,3}} <: AbstractArray{T,3}

A GeoArray is an AbstractArray, an AffineMap for calculating coordinates based on the axes and a CRS definition to interpret these coordinates into in the real world.
It's three dimensional and can be seen as a stack (3D) of 2D geospatial rasters (bands), the dimensions are :x, :y, and :bands.
The AffineMap and CRS (coordinates) only operate on the :x and :y dimensions.
"""
mutable struct GeoArray{T<:RealOrMissing,A<:AbstractArray{T,3}} <: AbstractArray{T,3}
    A::A
    f::CoordinateTransformations.AffineMap{StaticArrays.SMatrix{2,2,Float64,4},StaticArrays.SVector{2,Float64}}
    crs::GeoFormatTypes.WellKnownText{GeoFormatTypes.CRS,String}
end

"""
    GeoArray(A::AbstractArray{T,3} where T <: RealOrMissing)

Construct a GeoArray from any Array. A default `AffineMap` and `CRS` will be generated.

# Examples
```julia-repl
julia> GeoArray(rand(10,10,1))
10x10x1 Array{Float64, 3} with AffineMap([1.0 0.0; 0.0 1.0], [0.0, 0.0]) and undefined CRS
```
"""
GeoArray(A::AbstractArray{T,3} where {T<:RealOrMissing}) = GeoArray(A, geotransform_to_affine(SVector(0.0, 1.0, 0.0, 0.0, 0.0, 1.0)), "")
"""
    GeoArray(A::AbstractArray{T,3} where T <: RealOrMissing, f::AffineMap)

Construct a GeoArray from any Array and an `AffineMap` that specifies the coordinates. A default `CRS` will be generated.
"""
GeoArray(A::AbstractArray{T,3} where {T<:RealOrMissing}, f::AffineMap) = GeoArray(A, f, GFT.WellKnownText(GFT.CRS(), ""))

"""
    GeoArray(A::AbstractArray{T,3} where T <: RealOrMissing, f::AffineMap, crs::String)

Construct a GeoArray from any Array and an `AffineMap` that specifies the coordinates and `crs` string in WKT format.
"""
GeoArray(A::AbstractArray{T,3} where {T<:RealOrMissing}, f::AffineMap, crs::String) = GeoArray(A, f, GFT.WellKnownText(GFT.CRS(), crs))

GeoArray(A::AbstractArray{T,3} where {T<:RealOrMissing}, f::AffineMap{Matrix{Float64},Vector{Float64}}, crs::GeoFormatTypes.WellKnownText{GeoFormatTypes.CRS,String}) = GeoArray(A, AffineMap(SMatrix{2,2}(f.linear), SVector{2}(f.translation)), crs)

"""
    GeoArray(A::AbstractArray{T,2} where T <: RealOrMissing)

Construct a GeoArray from any Matrix, a third singleton dimension will be added automatically.

# Examples
```julia-repl
julia> GeoArray(rand(10,10))
10x10x1 Array{Float64, 3} with AffineMap([1.0 0.0; 0.0 1.0], [0.0, 0.0]) and undefined CRS
```
"""
GeoArray(A::AbstractArray{T,2} where {T<:RealOrMissing}, args...) = GeoArray(reshape(A, size(A)..., 1), args...)

"""
    GeoArray(A::AbstractArray{T,3} where T <: RealOrMissing, x::AbstractRange, y::AbstractRange, args...)

Construct a GeoArray any Array and it's coordinates from `AbstractRange`s for each dimension.
"""
function GeoArray(A::AbstractArray{T,3} where {T<:RealOrMissing}, x::AbstractRange, y::AbstractRange, args...)
    size(A)[1:2] != (length(x), length(y)) && error("Size of `GeoArray` $(size(A)) does not match size of (x,y): $((length(x), length(y))). Note that this function takes *center coordinates*.")
    f = unitrange_to_affine(x, y)
    GeoArray(A, f, args...)
end

# Behave like an Array
Base.size(ga::GeoArray) = size(ga.A)
Base.IndexStyle(::Type{<:GeoArray}) = IndexCartesian()
Base.similar(ga::GeoArray, t::Type) = GeoArray(similar(ga.A, t), ga.f, ga.crs)
Base.iterate(ga::GeoArray) = iterate(ga.A)
Base.length(ga::GeoArray) = length(ga.A)
Base.parent(ga::GeoArray) = ga.A
Base.eltype(::Type{GeoArray{T}}) where {T} = T
Base.show(io::IO, ::MIME"text/plain", ga::GeoArray) = show(io, ga)

function Base.show(io::IO, ga::GeoArray)
    crs = GeoFormatTypes.val(ga.crs)
    wkt = length(crs) == 0 ? "undefined CRS" : "CRS $crs"
    print(io, "$(join(size(ga), "x")) $(typeof(ga.A)) with $(ga.f) and $(wkt)")
end

# Getindex
"""
    getindex(ga::GeoArray, i::AbstractRange, j::AbstractRange, k::Union{Colon,AbstractRange,Integer})

Index a GeoArray with `AbstractRange`s to get a cropped GeoArray with the correct `AffineMap` set.

# Examples
```julia-repl
julia> ga[2:3,2:3,1]
2x2x1 Array{Float64, 3} with AffineMap([1.0 0.0; 0.0 1.0], [1.0, 1.0]) and undefined CRS
```
"""
function Base.getindex(ga::GeoArray, i::AbstractRange, j::AbstractRange, k::Union{Colon,AbstractRange,Integer})
    A = getindex(ga.A, i, j, k)
    x, y = first(i) - 1, first(j) - 1
    t = ga.f(SVector(x, y))
    l = ga.f.linear * SMatrix{2,2}([step(i) 0; 0 step(j)])
    GeoArray(A, AffineMap(l, t), ga.crs)
end
Base.getindex(ga::GeoArray, i::AbstractRange, j::AbstractRange) = Base.getindex(ga, i, j, :)

Base.getindex(ga::GeoArray, I::Vararg{<:Integer,2}) = getindex(ga.A, I..., :)
Base.getindex(ga::GeoArray, I::Vararg{<:Integer,3}) = getindex(ga.A, I...)

# Getindex and setindex! with floats
"""
    getindex(ga::GeoArray, I::SVector{2,<:AbstractFloat})

Index a GeoArray with `AbstractFloat`s to automatically get the value at that coordinate, using the function `indices`.
A `BoundsError` is raised if the coordinate falls outside the bounds of the raster.

# Examples
```julia-repl
julia> ga[3.0,3.0]
1-element Vector{Float64}:
 0.5630767850028582
```
"""
function Base.getindex(ga::GeoArray, I::SVector{2,<:AbstractFloat})
    (i, j) = indices(ga, I, Center())
    ga[i, j, :]
end
Base.getindex(ga::GeoArray, I::Vararg{<:AbstractFloat,2}) = getindex(ga, SVector{2}(I))

function Base.setindex!(ga::GeoArray, v, I::SVector{2,AbstractFloat})
    i, j = indices(ga, I, Center())
    ga.A[i, j, :] .= v
end
Base.setindex!(ga::GeoArray, v, I::Vararg{<:AbstractFloat,2}) = setindex!(ga, v, SVector{2}(I))
Base.setindex!(ga::GeoArray, v, I::Vararg{Union{<:Integer,<:AbstractRange{<:Integer}},2}) = setindex!(ga.A, v, I..., :)
Base.setindex!(ga::GeoArray, v, I::Vararg{Union{<:Integer,<:AbstractRange{<:Integer}},3}) = setindex!(ga.A, v, I...)

# Coordinates and indices
abstract type AbstractStrategy end
"""
    Center()

Strategy to use in functions like `indices` and `coords`, in which
it will use the center of the raster cells to do coordinate conversion.
"""
struct Center <: AbstractStrategy
    offset::Float64
    Center() = new(0.5)
end
"""
    Vertex()

Strategy to use in functions like `indices` and `coords`, in which
it will use the top left vertex of the raster cells to do coordinate conversion.
"""
struct Vertex <: AbstractStrategy
    offset::Float64
    Vertex() = new(1.0)
end

"""
    coords(ga::GeoArray, p::SVector{2,<:Integer}, strategy::AbstractStrategy)

Retrieve coordinates of the cell index by `p`.
See `indices` for the inverse function.
"""
function coords(ga::GeoArray, p::SVector{2,<:Integer}, strategy::AbstractStrategy)
    SVector{2}(ga.f(p .- strategy.offset))
end
coords(ga::GeoArray, p::Vector{<:Integer}, strategy::AbstractStrategy=Center()) = coords(ga, SVector{2}(p), strategy)
coords(ga::GeoArray, p::Tuple{<:Integer,<:Integer}, strategy::AbstractStrategy=Center()) = coords(ga, SVector{2}(p), strategy)

"""
    indices(ga::GeoArray, p::SVector{2,<:AbstractFloat}, strategy::AbstractStrategy)

Retrieve logical indices of the cell represented by coordinates `p`.
`strategy` can be used to define whether the coordinates represent the center (`Center`) or the top left corner (`Vertex`) of the cell.
See `coords` for the inverse function.
"""
function indices(ga::GeoArray, p::SVector{2,<:AbstractFloat}, strategy::AbstractStrategy)
    round.(Int, inv(ga.f)(p) .+ strategy.offset)::SVector{2,Int}
end
indices(ga::GeoArray, p::Vector{<:AbstractFloat}, strategy::AbstractStrategy=Center()) = indices(ga, SVector{2}(p), strategy)
indices(ga::GeoArray, p::Tuple{<:AbstractFloat,<:AbstractFloat}, strategy::AbstractStrategy=Center()) = indices(ga, SVector{2}(p), strategy)

# Generate coordinates for complete GeoArray
function coords(ga::GeoArray, strategy::AbstractStrategy=Center())
    (ui, uj) = size(ga)[1:2]
    extra = typeof(strategy) == Center ? 0 : 1
    [coords(ga, SVector{2}(i, j), strategy) for i in 1:ui+extra, j in 1:uj+extra]::Matrix{SVector{2,Float64}}
end

# Generate coordinates for one dimension of a GeoArray
function coords(ga::GeoArray, dim::Symbol, strategy::AbstractStrategy=Center())
    if is_rotated(ga)
        error("This method cannot be used for a rotated GeoArray")
    end
    extra = typeof(strategy) == Center ? 0 : 1
    if dim == :x
        ui = size(ga, 1)
        ci = [coords(ga, SVector{2}(i, 1), strategy)[1] for i in 1:ui+extra]
    elseif dim == :y
        uj = size(ga, 2)
        ci = [coords(ga, SVector{2}(1, j), strategy)[2] for j in 1:uj+extra]
    else
        error("Use :x or :y as second argument")
    end
    return ci
end

"""
    coords!(ga, x::AbstractUnitRange, y::AbstractUnitRange)

Set AffineMap of `GeoArray` by specifying the *center coordinates* for each `x`, `y` dimension by a `UnitRange`.
"""
function coords!(ga, x::AbstractUnitRange, y::AbstractUnitRange)
    size(ga)[1:2] != (length(x), length(y)) && error("Size of `GeoArray` $(size(ga)) does not match size of (x,y): $((length(x), length(y))). Note that this function takes *center coordinates*.")
    ga.f = unitrange_to_affine(x, y)
    ga
end


text:
# Refine elements in the DG solver based on a list of cell_ids that should be refined
function refine!(u_ode::AbstractVector, adaptor, mesh::TreeMesh{1},
                 equations, dg::DGSEM, cache, cells_to_refine)
  # Return early if there is nothing to do
  if isempty(cells_to_refine)
    return
  end

  # Determine for each existing element whether it needs to be refined
  needs_refinement = falses(nelements(dg, cache))

  # The "Ref(...)" is such that we can vectorize the search but not the array that is searched
  elements_to_refine = searchsortedfirst.(Ref(cache.elements.cell_ids[1:nelements(dg, cache)]),
                                          cells_to_refine)
  needs_refinement[elements_to_refine] .= true

  # Retain current solution data
  old_n_elements = nelements(dg, cache)
  old_u_ode = copy(u_ode)
  GC.@preserve old_u_ode begin # OBS! If we don't GC.@preserve old_u_ode, it might be GC'ed
    old_u = wrap_array(old_u_ode, mesh, equations, dg, cache)

    # Get new list of leaf cells
    leaf_cell_ids = local_leaf_cells(mesh.tree)

    # re-initialize elements container
    @unpack elements = cache
    resize!(elements, length(leaf_cell_ids))
    init_elements!(elements, leaf_cell_ids, mesh, dg.basis.nodes)
    @assert nelements(dg, cache) > old_n_elements

    resize!(u_ode, nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache))
    u = wrap_array(u_ode, mesh, equations, dg, cache)

    # Loop over all elements in old container and either copy them or refine them
    element_id = 1
    for old_element_id in 1:old_n_elements
      if needs_refinement[old_element_id]
        # Refine element and store solution directly in new data structure
        refine_element!(u, element_id, old_u, old_element_id,
                        adaptor, equations, dg)
        element_id += 2^ndims(mesh)
      else
        # Copy old element data to new element container
        @views u[:, .., element_id] .= old_u[:, .., old_element_id]
        element_id += 1
      end
    end
    # If everything is correct, we should have processed all elements.
    # Depending on whether the last element processed above had to be refined or not,
    # the counter `element_id` can have two different values at the end.
    @assert element_id == nelements(dg, cache) + 1 || element_id == nelements(dg, cache) + 2^ndims(mesh) "element_id = $element_id, nelements(dg, cache) = $(nelements(dg, cache))"
  end # GC.@preserve old_u_ode

  # re-initialize interfaces container
  @unpack interfaces = cache
  resize!(interfaces, count_required_interfaces(mesh, leaf_cell_ids))
  init_interfaces!(interfaces, elements, mesh)

  # re-initialize boundaries container
  @unpack boundaries = cache
  resize!(boundaries, count_required_boundaries(mesh, leaf_cell_ids))
  init_boundaries!(boundaries, elements, mesh)

  # Sanity check
  if isperiodic(mesh.tree)
    @assert ninterfaces(interfaces) == 1 * nelements(dg, cache) ("For 1D and periodic domains, the number of interfaces must be the same as the number of elements")
  end

  return nothing
end

# TODO: Taal compare performance of different implementations
# Refine solution data u for an element, using L2 projection (interpolation)
function refine_element!(u::AbstractArray{<:Any,3}, element_id,
                         old_u, old_element_id,
                         adaptor::LobattoLegendreAdaptorL2, equations, dg)
  @unpack forward_upper, forward_lower = adaptor

  # Store new element ids
  left_id  = element_id
  right_id = element_id + 1

  @boundscheck begin
    @assert old_element_id >= 1
    @assert size(old_u, 1) == nvariables(equations)
    @assert size(old_u, 2) == nnodes(dg)
    @assert size(old_u, 3) >= old_element_id
    @assert     element_id >= 1
    @assert size(    u, 1) == nvariables(equations)
    @assert size(    u, 2) == nnodes(dg)
    @assert size(    u, 3) >= element_id + 1
  end

  # Interpolate to left element
  for i in eachnode(dg)
    acc = zero(get_node_vars(u, equations, dg, i, element_id))
    for k in eachnode(dg)
      acc += get_node_vars(old_u, equations, dg, k, old_element_id) * forward_lower[i, k]
    end
    set_node_vars!(u, acc, equations, dg, i, left_id)
  end

  # Interpolate to right element
  for i in eachnode(dg)
    acc = zero(get_node_vars(u, equations, dg, i, element_id))
    for k in eachnode(dg)
      acc += get_node_vars(old_u, equations, dg, k, old_element_id) * forward_upper[i, k]
    end
    set_node_vars!(u, acc, equations, dg, i, right_id)
  end

  return nothing
end

# Coarsen elements in the DG solver based on a list of cell_ids that should be removed
function coarsen!(u_ode::AbstractVector, adaptor, mesh::TreeMesh{1},
                  equations, dg::DGSEM, cache, child_cells_to_coarsen)
  # Return early if there is nothing to do
  if isempty(child_cells_to_coarsen)
    return
  end

  # Determine for each old element whether it needs to be removed
  to_be_removed = falses(nelements(dg, cache))
  # The "Ref(...)" is such that we can vectorize the search but not the array that is searched
  elements_to_remove = searchsortedfirst.(Ref(cache.elements.cell_ids[1:nelements(dg, cache)]),
                                          child_cells_to_coarsen)
  to_be_removed[elements_to_remove] .= true

  # Retain current solution data
  old_n_elements = nelements(dg, cache)
  old_u_ode = copy(u_ode)
  GC.@preserve old_u_ode begin # OBS! If we don't GC.@preserve old_u_ode, it might be GC'ed
    old_u = wrap_array(old_u_ode, mesh, equations, dg, cache)

    # Get new list of leaf cells
    leaf_cell_ids = local_leaf_cells(mesh.tree)

    # re-initialize elements container
    @unpack elements = cache
    resize!(elements, length(leaf_cell_ids))
    init_elements!(elements, leaf_cell_ids, mesh, dg.basis.nodes)
    @assert nelements(dg, cache) < old_n_elements

    resize!(u_ode, nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache))
    u = wrap_array(u_ode, mesh, equations, dg, cache)

    # Loop over all elements in old container and either copy them or coarsen them
    skip = 0
    element_id = 1
    for old_element_id in 1:old_n_elements
      # If skip is non-zero, we just coarsened 2^ndims elements and need to omit the following elements
      if skip > 0
        skip -= 1
        continue
      end

      if to_be_removed[old_element_id]
        # If an element is to be removed, sanity check if the following elements
        # are also marked - otherwise there would be an error in the way the
        # cells/elements are sorted
        @assert all(to_be_removed[old_element_id:(old_element_id+2^ndims(mesh)-1)]) "bad cell/element order"

        # Coarsen elements and store solution directly in new data structure
        coarsen_elements!(u, element_id, old_u, old_element_id,
                          adaptor, equations, dg)
        element_id += 1
        skip = 2^ndims(mesh) - 1
      else
        # Copy old element data to new element container
        @views u[:, .., element_id] .= old_u[:, .., old_element_id]
        element_id += 1
      end
    end
    # If everything is correct, we should have processed all elements.
    @assert element_id == nelements(dg, cache) + 1 "element_id = $element_id, nelements(dg, cache) = $(nelements(dg, cache))"
  end # GC.@preserve old_u_ode

  # re-initialize interfaces container
  @unpack interfaces = cache
  resize!(interfaces, count_required_interfaces(mesh, leaf_cell_ids))
  init_interfaces!(interfaces, elements, mesh)

  # re-initialize boundaries container
  @unpack boundaries = cache
  resize!(boundaries, count_required_boundaries(mesh, leaf_cell_ids))
  init_boundaries!(boundaries, elements, mesh)

  # Sanity check
  if isperiodic(mesh.tree)
    @assert ninterfaces(interfaces) == 1 * nelements(dg, cache) ("For 1D and periodic domains, the number of interfaces must be the same as the number of elements")
  end

  return nothing
end

# TODO: Taal compare performance of different implementations
# Coarsen solution data u for two elements, using L2 projection
function coarsen_elements!(u::AbstractArray{<:Any,3}, element_id,
                           old_u, old_element_id,
                           adaptor::LobattoLegendreAdaptorL2, equations, dg)
  @unpack reverse_upper, reverse_lower = adaptor

  # Store old element ids
  left_id  = old_element_id
  right_id = old_element_id + 1

  @boundscheck begin
    @assert old_element_id >= 1
    @assert size(old_u, 1) == nvariables(equations)
    @assert size(old_u, 2) == nnodes(dg)
    @assert size(old_u, 3) >= old_element_id + 1
    @assert     element_id >= 1
    @assert size(    u, 1) == nvariables(equations)
    @assert size(    u, 2) == nnodes(dg)
    @assert size(    u, 3) >= element_id
  end

  for i in eachnode(dg)
    acc = zero(get_node_vars(u, equations, dg, i, element_id))

    # Project from lower left element
    for k in eachnode(dg)
      acc += get_node_vars(old_u, equations, dg, k, left_id) * reverse_lower[i, k]
    end

    # Project from lower right element
    for k in eachnode(dg)
      acc += get_node_vars(old_u, equations, dg, k, right_id) * reverse_upper[i, k]
    end

    # Update value
    set_node_vars!(u, acc, equations, dg, i, element_id)
  end
end

# this method is called when an `ControllerThreeLevel` is constructed
function create_cache(::Type{ControllerThreeLevel}, mesh::TreeMesh{1}, equations, dg::DG, cache)

  controller_value = Vector{Int}(undef, nelements(dg, cache))
  return (; controller_value)
end

function create_cache(::Type{ControllerThreeLevelCombined}, mesh::TreeMesh{1}, equations, dg::DG, cache)

  controller_value = Vector{Int}(undef, nelements(dg, cache))
  return (; controller_value)
end


text:
if( Pkg.installed("PyPlot") == nothing)
    Pkg.add("PyPlot")
end

using PyPlot

function plotfunction(PDATA::PAINTER_Data,OIDATA::PAINTER_Input)
    nx = OIDATA.nx
    nw = OIDATA.nw
    wvl = OIDATA.wvl
    FOV = OIDATA.FOV
    x = PDATA.x
    w = PDATA.w .> 0.

    indpix = linspace(-(FOV / 2), (FOV / 2), nx)
    pos = round(Int, [1, round(Int,nx / 4), round(Int,nx / 2), round(Int,nx * 3 / 4), nx])

    count_y = 0
    count_x = 0
    SubRow  = 6
    SubColumn = 5

    for n in 1:nw
        subplot(SubColumn, SubRow, n)
        imshow(x[:, :, n] .* max(0, w[:, :, n]), origin = "lower")
        titlestring = @sprintf("%2.4f µm", wvl[n] * 1e6)
        title(titlestring)
        xticks([])
        yticks([])
        if( n == (nw + 1 - SubRow + count_x) )
            xticks(collect(pos - 1), round(Int,indpix[pos] * 100000) / 100)
            xlabel("FOV (mas)")
            count_x += 1
        end
        if(n == (1 + count_y * SubRow))
            yticks(collect(pos - 1), round(Int,indpix[pos] * 100000) / 100)
            ylabel("FOV (mas)")
            count_y += 1
        end
    end
end

    PlotFct = plotfunction
    FOV = 0.01
    indwvl = 1:30
    nx = 64
    eps1 = 1e-4
    eps2 = 1e-4
    rho_y = 10.
    alpha = 1e4
    beta = 1e5
    rho_spat = 4.
    rho_ps = rho_spat
    lambda_spat = 1e-5
    rho_spec = 0.5
    lambda_spec = 1e-5
    aff = true     # plot is enabled
    nbitermax = 1
    paral = true     # parallel computing is enabled
    savepath = "data.jld"
    dptype = "sliding"
    dpprm = 5
    Folder = ""
# initialize algorithm and run admm
    OIDATA, PDATA = painter(nbitermax = nbitermax, nx = nx, lambda_spat = lambda_spat,
                            lambda_spec = lambda_spec, rho_y = rho_y, rho_spat = rho_spat,
                            rho_spec = rho_spec, rho_ps = rho_ps, alpha = alpha, beta = beta,
                            eps1 = eps1, eps2 = eps2, FOV = FOV, indwvl = indwvl, admm=true,
                            PlotFct = PlotFct, aff = aff, paral = paral, dptype = dptype,
                            dpprm = dpprm, Folder = Folder)

# save data struture in .jld files
    paintersave(savepath,PDATA,OIDATA)

# load data struture in .jld files
    PDATA, OIDATA = painterload(savepath)

# Warm start of the algorithm
    OIDATA, PDATA = painter(OIDATA,PDATA,100,true, PlotFct = PlotFct)


text:
using Finder
using Geodesy
using Person
using MyGraph

user = Professional("usuario", LLA(-15.836073, -47.912019), "Programador", true)

function listworkers(file::String = "dados.txt", works::Array = Array{String, 1}(); dist::Float64 = 100.0)
    open("test\\" * file) do f
        for (index, line) in enumerate(eachline(f))
            p = Professional(line)
            p.available || continue
            (isempty(works) && distance(user, p)/1000 <= dist) && println("[",index,"] ", p);
            (p.work in works && distance(user, p)/1000 <= dist) && println("[",index,"] ", p);
        end
    end
end

function selectworkers(sw = Int[]) :: FinderGraph
    P = Dict{Int, Professional}(0 => user)
    open("test\\dados2.txt") do f
        for (index, line) in enumerate(eachline(f))
            index in sw ? push!(P, index => Professional(line)) : continue
        end
    end
    FinderGraph(P)
end

function traceroute(F::FinderGraph)
    return makecirclepath(F.graph)
end


text:
using StochasticIntegrals
using Test

# Run tests

println("Test Stochastic Integral Generation")
@time @test include("new_tests.jl")
@time @test include("conditioning.jl")
@time @test include("test_ito_processes.jl")


text:
using DataFrames
using CSV
using Plots
#%%
data_s0 = open("poisson_s0.dat")
df_s0=CSV.read(data_s0, header=false, delim=" ")
r_s0=df_s0[1]
phi_s0=df_s0[2]

data_s02 = open("poisson_s02.dat")
df_s02=CSV.read(data_s02, header=false, delim=" ")
r_s02=df_s02[1]
phi_s02=df_s02[2]

data_s04 = open("poisson_s04.dat")
df_s04=CSV.read(data_s04, header=false, delim=" ")
r_s04=df_s04[1]
phi_s04=df_s04[2]

data_s06 = open("poisson_s06.dat")
df_s06=CSV.read(data_s06, header=false, delim=" ")
r_s06=df_s06[1]
phi_s06=df_s06[2]

data_s08 = open("poisson_s08.dat")
df_s08=CSV.read(data_s08, header=false, delim=" ")
r_s08=df_s08[1]
phi_s08=df_s08[2]

data_s1 = open("poisson_s1.dat")
df_s1=CSV.read(data_s1, header=false, delim=" ")
r_s1=df_s1[1]
phi_s1=df_s1[2]

#%%
plot(df_s0[1], df_s0[2], label="s 0.0", dpi=600)
plot!(df_s02[1], df_s02[2], label="s 0.2", dpi=600)
plot!(df_s04[1], df_s04[2], label="s 0.4", dpi=600)
plot!(df_s06[1], df_s06[2], label="s 0.6", dpi=600)
plot!(df_s08[1], df_s08[2], label="s 0.8", dpi=600)
plot!(df_s1[1], df_s1[2], label="s 1.0", dpi=600)

plot!(grid=true, legend=:topleft)
xlabel!("steps")
ylabel!("phi")
savefig("poisson_vari_s.png")
# xaxis!(:log10)
# yaxis!(:log10)

max_stars_repo_name: aurelio-amerio/corso-cpp


text:
# -------------------------- Util functions for preprocessing -------------------------- */

# Default header for Tabular datasets is empty String.
"""
    headers(::Tabular)

Returns column names of `Tabular` dataset. Default is empty `String`.
"""
headers(dataset::Tabular) = ""

"""
    target(::Tabular)

Returns order of target column, else throws error.
"""
target(dataset::Tabular) = error("Target value not specified for $dataset dataset.")

"""
    categorical(::Tabular)

Returns indeces of columns containing categorical data. Default is empty range.
"""

categorical(dataset::Tabular) = 1:0
transposed(dataset::Tabular) = false

# ---------------------------- Util functions for splitting ---------------------------- */
"""
    getdata(dataset::Tabular, type::Type{<:Split})

Returns data from dataset. `type` is either `Train`, `Test` or `Valid`.
"""
function getdata(dataset::Tabular, type::Type{Train})
    return CSV.File(joinpath(getpath(dataset) , "data-train.csv"), header=true) |> DataFrame
end

function getdata(dataset::Tabular, type::Type{Valid})
    return CSV.File(joinpath(getpath(dataset) , "data-valid.csv"), header=true) |> DataFrame
end

function getdata(dataset::Tabular, type::Type{Test})
    return CSV.File(joinpath(getpath(dataset) , "data-test.csv"), header=true) |> DataFrame
end
# getpath(dataset) - path to a directory of given datadep

"""
    load(dataset::Tabular, type::Type{<:Split}; kwargs...)

Loads dataset of given type. Type based on filenames in datadeps folder (types available
only if the dataset was splitted before downloading, else the whole dataset is of type `Train`)

- `type::Type{<:Split}=Train`: other possible types are `Test` or `Valid`.

# Keywords
- `toarray::Bool=false`: if false returns `DataFrame`, else returns `Tuple` of arrays
- `header::Bool=false`: if true returnes `DataFrame` has column names belonging to the
dataset (it they are found), else default column naming is returned.
"""
function load(
    dataset::Tabular,
    type::Type{<:Split}=Train;
    toarray::Bool=false,
    header::Bool=false,
    )
    return postprocess(dataset, getdata(dataset, type), toarray, header)
end

"""
    splits(dataset::Tabular, data, indeces1, indeces2)

Devides `data` from `dataset` to two parts.
"""
function splits(dataset::Tabular, data, indeces1, indeces2)
    return data[indeces1,:], data[indeces2,:]
end

"""
    postprocess(dataset::Tabular, df::DataFrame,toarray::Bool, header::Bool)

Returns data of `Tabular` dataset.

#Arguments
toarray::Bool=false: Returns data either as a tuple of two tuples (each has a data array
and labels vector) or as `DataFrame`.
header::Bool=false: If true, searches for header file for the dataset and renames columns
according to the file, else nothing happens.

If both `toarray==true` and `header==true` nothing happens.
"""
function postprocess(
    dataset::Tabular,
    df::DataFrame,
    toarray::Bool=false,
    header::Bool=false,
    )
    header && !toarray && changeheader(dataset, df)

    return df_or_array(df, toarray)
end

function postprocess(
    d::Tabular,
    data1::DataFrame,
    data2::DataFrame;
    toarray::Bool=false,
    header::Bool=false,
    )
    return postprocess(d,data1,toarray,header), postprocess(d,data2,toarray,header)
end

function postprocess(
    d::Tabular,
    data1::DataFrame,
    data2::DataFrame,
    data3::DataFrame;
    toarray::Bool=false,
    header::Bool=false,
    )
    return postprocess(d,data1,toarray,header), postprocess(d,data2,toarray,header), postprocess(d,data3,toarray,header)
end

# ----------------------- Util functions for manipulating header ----------------------- */

"""
    loadheader(path::String, dataset::Tabular)

Loads header for given dataset from a file.

Changes position of target column if target column was not the last column.
"""
function loadheader(path::String, dataset::Tabular)
    if isfile(path)
        df = CSV.File(path, header = false) |> DataFrame
        targ = target(dataset)
        header = Array(df[:,1])
        if targ isa Int
            push!(header, header[targ])
            header = header[Not(targ)]
        end
        return header
    else
        @info "No file with header (column names) found."
        return ""
    end
end

"""
    getheader(dataset::Tabular)

Returns header found in header.csv in folder .../datadeps/dataset. See [`load`](@ref).
"""
function getheader(dataset::Tabular)
    path = joinpath(getpath(dataset), "header.csv")
    return loadheader(path, dataset)
end

"""
    changeheader(dataset::Tabular, df::DataFrame)

Changes header of `DataFrame` of given dataset, id there is a header.
"""
function changeheader(dataset::Tabular, df::DataFrame)
    hds = getheader(dataset)
    if !isempty(hds)
        new_header(hds, df)
    end
end

"""
    new_header(header::Vector{String}, df::DataFrame...)

Change column names of `DataFrames` to column names from header.

Length of `header` must be the same as number of columns of `df`.
"""
function new_header(header::Vector{String}, df::DataFrame...)
    if isempty(header)
        return nothing
    end
    for d in df
        rename!(d, header)
    end
end

# ---------------------- Other functions for manipulating header ---------------------- */

"""
    labels(dataset::Tabular)

Returns train label array for given dataset.
"""
function labels(dataset::Tabular)
    df = load(dataset)
    return df[:,end]
end

max_stars_repo_name: GruAna/PreprocessData

max_issues_repo_name: GruAna/PreprocessData

max_forks_repo_name: GruAna/PreprocessData


text:
@tags button div h1 h2 header img input label li span ul

function findimages(dir=plotsdir(); recursive=false, exts=[".png", ".jpeg", ".jpg", ".svg"])
    dir = abspath(dir)
    paths = if recursive
        paths = String[]
        for (root, dirs, files) in walkdir(dir)
            foreach(file -> push!(paths, joinpath(root, file)), files)
        end
        paths
    else
        readdir(dir; join=true)
    end
    filter(file -> any(endswith.(file, exts)), paths)
end

function dataset(paths; force=false)
    props = parse_savename.(paths)
    prefixes = map(x -> basename(x[1]), props)
    params = map(x -> x[2], props)

    images = vcat(DataFrame.(params)..., cols=:union)

    if hasproperty(images, :path) && !force
        error("image files have a parameter called path; consider passing force=true")
    else
        images.path = "file://" .* paths;
    end

    if hasproperty(images, :dataset) && !force
        error("image files have a parameter called dataset; consider passing force=true")
    else
        images.dataset = prefixes
    end

    if hasproperty(images, :id) && !force
        error("image files have a parameter called id; consider passing force=true")
    else
        images.id = "image" .* string.(eachindex(images.path))
    end

    images
end

function sidebar(values)
    sidebar = []
    for (header, vals) in values
        list = []
        for (i, val) in enumerate(vals)
            push!(list,
                  li(input(type="checkbox", name="$(header)$i", value=string(val),
                           onclick="""Blink.msg("press", ["$header", "$val"]);"""),
                     label(string(val); :for => "$(header)$i")))
        end
        push!(sidebar, h2(header), ul(list...))
    end
    div(id="sidebar",
        header(h1("Filters")),
        sidebar...,
        button(id="reset", onclick="""Blink.msg("reset", []);""", "Reset"))
end

function plots(df, visible)
    combine(groupby(df, [:id, :path])) do group
        (; visible = all(string(group[1,key]) in values for (key, values) in visible))
    end
end

function images(df)
    map(eachrow(df)) do row
        img(src=row.path, id=row.id, title=row.path, loading="lazy", onload="setupImage('#$(row.id)');")
    end
end

function toggle(win, df, visible)
    df = plots(df, visible)
    @js_ win toggle($(df.id), $(df.visible))
end

function explore(dir=plotsdir(); title="ImageExplorer", force=false, kwargs...)
    data = dataset(findimages(dir; kwargs...); force)

    cols = filter(x -> !(x in ["id", "path"]), names(data))
    values = filter(Dict(map(col -> col => sort(unique(data[:, col])), cols)...)) do x
        length(last(x)) > 1
    end

    visible = Dict()

    win = Window(Dict("title" => title, "webPreferences" => Dict("webSecurity" => false)))

    handle(win, "press") do (header, value)
        if haskey(visible, header)
            if value in visible[header]
                delete!(visible[header], value)
                if isempty(visible[header])
                    delete!(visible, header)
                end
            else
                push!(visible[header], value)
            end
        else
            visible[header] = Set([value])
        end

        toggle(win, data, visible)
    end

    handle(win, "reset") do args...
        visible = Dict()
        @js_ win d3.selectAll("input[type='checkbox']").property("checked", false)
        toggle(win, data, visible)
    end

    let loaded = []
        handle(win, "loaded") do id
            push!(loaded, id)
            if isempty(setdiff(data.id, loaded))
                @js_ win d3.selectAll("#overlay").classed("hidden", true)
            end
        end
    end

    loadcss!(win, "file://$(@__DIR__)/../assets/style.css")
    loadjs!(win, "https://d3js.org/d3.v7.min.js")
    loadjs!(win, "file://$(@__DIR__)/../assets/explore.js")

    body!(win, div(div(id="overlay", class="flex",
                       img(src="file://$(@__DIR__)/../assets/spinner.gif"),
                       "Loading images..."),
                   sidebar(values),
                   div(id="content",
                       header(h1("$(nrow(data)) Images"),
                              div(span("Image Width:"),
                                  input(type="range", value=500, min=100, step=10, max=1000, name="width", onchange="resize();"),
                                  label("500px"; :for => "width"))),
                       div(id="plots", images(data)))))

    nothing
end


text:
mutable struct ElementaryPathAlgosSolver <: ElementaryPathSolver
  graph::SimpleDiGraph
  source::Int
  destination::Int

  function ElementaryPathAlgosSolver()
    return new(DiGraph(0), -1, -1)
  end
end

has_lp_formulation(::ElementaryPathAlgosSolver) = false
approximation_ratio(::ElementaryPathAlgosSolver) = 1.0
approximation_term(::ElementaryPathAlgosSolver) = 0.0
approximation_ratio_budgeted(::ElementaryPathAlgosSolver) = 1.0
approximation_term_budgeted(::ElementaryPathAlgosSolver) = 0.0
supports_solve_budgeted_linear(::ElementaryPathAlgosSolver) = true
supports_solve_all_budgeted_linear(::ElementaryPathAlgosSolver) = true

_tuple_dict_to_edge_dict(d::Dict{Tuple{Int, Int}, T}) where T =
  Dict(Edge(k...) => v for (k, v) in d)

function build!(solver::ElementaryPathAlgosSolver, graph::SimpleDiGraph, source::Int, destination::Int)
  solver.graph = graph
  solver.source = source
  solver.destination = destination
end

function solve_linear(solver::ElementaryPathAlgosSolver, reward::Dict{Tuple{Int, Int}, Float64})
  i = ElementaryPathInstance(solver.graph, _tuple_dict_to_edge_dict(reward),
    solver.source, solver.destination)
  s = lp_dp(i).path
  return [(src(e), dst(e)) for e in s]
end

function solve_budgeted_linear(solver::ElementaryPathAlgosSolver,
                               reward::Dict{Tuple{Int, Int}, Float64},
                               weights::Dict{Tuple{Int, Int}, Int},
                               budget::Int)
  i = BudgetedElementaryPathInstance(solver.graph,
    _tuple_dict_to_edge_dict(reward), _tuple_dict_to_edge_dict(weights),
    solver.source, solver.destination, budget=budget)
  s = budgeted_lp_dp(i).path
  return [(src(e), dst(e)) for e in s]
end

function solve_all_budgeted_linear(solver::ElementaryPathAlgosSolver,
                                   reward::Dict{Tuple{Int, Int}, Float64},
                                   weights::Dict{Tuple{Int, Int}, Int},
                                   max_budget::Int)
  i = BudgetedElementaryPathInstance(solver.graph,
    _tuple_dict_to_edge_dict(reward), _tuple_dict_to_edge_dict(weights),
    solver.source, solver.destination, budget=max_budget)
  return paths_all_budgets_as_tuples(budgeted_lp_dp(i), max_budget)
end


text:
using DelimitedFiles
using FFTW

function prepareFWIDataFiles(m,Minv::RegularMesh,mref,boundsHigh,boundsLow,
							filenamePrefix::String,omega::Array{Float64,1},waveCoef::Array{ComplexF64,1},
							pad::Int64,ABLpad::Int64,jumpSrc::Int64,jumpRcv::Int64,offset::Int64=prod(Minv.n+1),workerList = workers(),
							maxBatchSize::Int64=48, Ainv::AbstractSolver = getMUMPSsolver([],0,0,2),useFilesForFields::Bool = false,calcTravelTime::Bool=false)

########################## m is in Velocity here. ###################################
RCVfile = string(filenamePrefix,"_rcvMap.dat");
SRCfile = string(filenamePrefix,"_srcMap.dat");
writeSrcRcvLocFile(SRCfile,Minv,ABLpad,jumpSrc);
writeSrcRcvLocFile(RCVfile,Minv,ABLpad,jumpRcv);

dataFullFilenamePrefix = string(filenamePrefix,"_freq");
gamma = prepareFWIDataFiles2(m, Minv, filenamePrefix,dataFullFilenamePrefix,omega,waveCoef,pad,ABLpad,offset,workerList,maxBatchSize,Ainv,useFilesForFields);

HO = false;
if calcTravelTime
	dataFullFilenamePrefix = string(filenamePrefix,"_travelTime");
	prepareTravelTimeDataFiles(m, Minv, filenamePrefix,dataFullFilenamePrefix,offset,HO,useFilesForFields);
end

file = matopen(string(filenamePrefix,"_PARAM.mat"), "w");
write(file,"boundsLow",boundsLow);
write(file,"boundsHigh",boundsHigh);
write(file,"mref",mref);
write(file,"domain",Minv.domain);
write(file,"n",Minv.n);
write(file,"gamma",gamma);
write(file,"pad",pad);
write(file,"omega",omega);
write(file,"waveCoef",waveCoef);
if calcTravelTime
	write(file,"HO",HO);
end
close(file);
end

function prepareFWIDataFiles2(m, Minv::RegularMesh, filenamePrefix::String,dataFullFilenamePrefix::String, omega::Array{Float64,1},
								waveCoef::Array{ComplexF64,1}, pad::Int64,ABLpad::Int64,offset::Int64,workerList::Array{Int64,1},maxBatchSize::Int64,
								Ainv::AbstractSolver,useFilesForFields::Bool = false)
########################## m is in Velocity here. ###################################
println("maxOmega*maxKappaSq*h: should be below 0.6");
println(omega[end]*maximum(Minv.h)*(1.0./minimum(m)));

RCVfile = string(filenamePrefix,"_rcvMap.dat");
SRCfile = string(filenamePrefix,"_srcMap.dat");
srcNodeMap = readSrcRcvLocationFile(SRCfile,Minv);
rcvNodeMap = readSrcRcvLocationFile(RCVfile,Minv);

Q = generateSrcRcvProjOperators(Minv.n.+1,srcNodeMap);
P = generateSrcRcvProjOperators(Minv.n.+1,rcvNodeMap);
Q = Q.*1.0./(norm(Minv.h)^2);
println("We have ",size(Q,2)," sources");
# compute observed data

ABLamp = getMaximalFrequency(1.0./(minimum(m).^2),Minv);
gamma = getABL(Minv.n,true,ones(Int64,Minv.dim)*ABLpad,ABLamp);
attenuation = 0.01*4*pi;
gamma .+= attenuation; # adding Attenuation.

println("~~~~~~~ Getting data FWI: ~~~~~~~");

# Solve forward problem (should be relaced in reading data from file)

batch = min(size(Q,2),maxBatchSize);
(pFor,contDiv,SourcesSubInd) = getFWIparam(omega,waveCoef,vec(gamma),Q,P,Minv,Ainv,workerList,batch,useFilesForFields);

(D,pFor) = getData(velocityToSlowSquared(m[:])[1],pFor,ones(length(pFor)),true);

nsrc = size(Q,2);
nrcv = size(P,2);

for k = 1:length(omega)
	I = contDiv[k] : contDiv[k+1] - 1
	Dobsk = Array{Array{ComplexF64,2}}(undef,length(I));
	for i = 1:length(I)
		Dobsk[i] = fetch(D[I[i]]);
	end
	Dobsk = arrangeRemoteCallDataIntoLocalData(Dobsk);
	Dobsk += 0.01*mean(abs.(Dobsk))*(randn(size(Dobsk,1),size(Dobsk,2)) + 1im*randn(size(Dobsk,1),size(Dobsk,2)));
	omRound = string(round((omega[k]/(2*pi))*100.0)/100.0);
	# Wd_k = (1.0./(abs.(real.(Dobsk)).+0.1*mean(abs.(Dobsk)))) + 1im*(1.0./(abs.(imag.(Dobsk)).+0.1*mean(abs.(Dobsk))));
	# Wd_k = (1.0./(0.0*abs(real(Dobsk)).+1.0*mean(abs(Dobsk)))) .+ 1im*(1.0./(0.0*abs(imag(Dobsk)).+1.0*mean(abs(Dobsk))));
	Wd_k = (ones(size(Dobsk))./mean(abs.(Dobsk))).*(1+1im);
	# Wd_k = limitDataToOffset(Wd_k,srcNodeMap,rcvNodeMap,offset);
	filename = string(dataFullFilenamePrefix,omRound,".dat");
	writeDataFile(filename,Dobsk,Wd_k,srcNodeMap,rcvNodeMap);
end
return gamma;
end

# function prepareFWIDataFilesFromTime(m,Minv::RegularMesh,mref,boundsHigh,boundsLow,
							# filenamePrefix::String,omega::Array{Float64,1},
							# pad::Int64,ABLpad::Int64,jump::Int64,offset::Int64,workerList::Array{Int64},
							# T::Float64,dt::Float64,fm::Float64,sigmaNoise::Float64 = 0.0,pickFirstArival::Bool = true)

# ########################## m is in Velocity here. ###################################
# RCVfile = string(filenamePrefix,"_rcvMap.dat");
# SRCfile = string(filenamePrefix,"_srcMap.dat");
# writeSrcRcvLocFile(SRCfile,Minv,ABLpad,jump);
# writeSrcRcvLocFile(RCVfile,Minv,ABLpad,1);

# dataFullFilenamePrefix = string(filenamePrefix,"_freq");
# DobsTimeNoisy,gamma,rickt,waveCoef,Mask,zeroTimeOffset,rickerWidth = prepareFWIDataFilesFromTime(m,Minv,filenamePrefix,dataFullFilenamePrefix,omega,pad,ABLpad,offset,workerList,T,dt,fm,sigmaNoise);

# if pickFirstArival
	# recInfo = readdlm(RCVfile); x_rcv_loc = vec(recInfo[:,2])
	# srcInfo = readdlm(SRCfile); x_src_loc = vec(srcInfo[:,2])
	# # getFirstArival(dataTraces::Array{Array{Float64}},mask::Array,x_rcv_loc::Array{Float64},x_src_loc::Float64,vmax::Float64,vmin::Float64,window::Int64,dt::Float64,offset)
	# DobsPicked = getFirstArival(DobsTimeNoisy,Mask,x_rcv_loc,x_src_loc,maximum(boundsHigh),minimum(boundsLow),rickerWidth,dt,zeroTimeOffset);
	# Wd = (1.0./(abs.(DobsPicked) .+ 0.1*mean(abs.(DobsPicked))));
	# Wd = Wd.*Mask;
	# srcNodeMap = readSrcRcvLocationFile(SRCfile,Minv);
	# rcvNodeMap = readSrcRcvLocationFile(RCVfile,Minv);
	# writeDataFile(string(filenamePrefix,"_travelTime.dat"),DobsPicked,Wd,srcNodeMap,rcvNodeMap);
# end

# file = matopen(string(filenamePrefix,"_PARAM.mat"), "w");
# write(file,"boundsLow",boundsLow);
# write(file,"boundsHigh",boundsHigh);
# write(file,"mref",mref);
# write(file,"domain",Minv.domain);
# write(file,"n",Minv.n);
# write(file,"gamma",gamma);
# write(file,"pad",pad);
# write(file,"omega",omega);
# write(file,"T",T);
# write(file,"dt",dt);
# write(file,"fm",fm);
# write(file,"sigmaNoise",sigmaNoise);
# write(file,"waveCoef",waveCoef);
# write(file,"dataMask",Mask);
# write(file,"HO",false);
# write(file,"zeroTimeOffset",zeroTimeOffset);
# write(file,"rickerWidth",rickerWidth)
# close(file);
# return DobsTimeNoisy
# end

# function prepareFWIDataFilesFromTime(m, Minv::RegularMesh, filenamePrefix::String,dataFullFilenamePrefix::String, omega::Array{Float64,1},
									# pad::Int64,ABLpad::Int64,offset::Int64,workerList::Array{Int64,1},T::Float64,dt::Float64,fm::Float64,sigmaNoise)
# ########################## m is in Velocity here. ###################################
# println("maxOmega*maxKappaSq*h: should be below 0.6");
# println(omega[end]*maximum(Minv.h)*(1.0./minimum(m)));

# RCVfile = string(filenamePrefix,"_rcvMap.dat");
# SRCfile = string(filenamePrefix,"_srcMap.dat");
# srcNodeMap = readSrcRcvLocationFile(SRCfile,Minv);
# rcvNodeMap = readSrcRcvLocationFile(RCVfile,Minv);

# Q = generateSrcRcvProjOperators(Minv.n.+1,srcNodeMap);
# P = generateSrcRcvProjOperators(Minv.n.+1,rcvNodeMap);
# Q = Q.*(1.0./(norm(Minv.h)^2));
# println("We have ",size(Q,2)," sources");
# # compute observed data

# nsrc = size(Q,2);
# nrcv = size(P,2);

# ABLamp = getMaximalFrequency(1.0./(minimum(m).^2),Minv);
# gamma = getABL(Minv.n,true,ones(Int64,Minv.dim)*ABLpad,ABLamp);
# attenuation = 0.01*4*pi;;
# gamma .+= attenuation; # adding Attenuation.

# println("~~~~~~~ Getting time domain FWI data: ~~~~~~~");

# Mask = ones(nrcv,nsrc);
# Mask = limitDataToOffset(Mask,srcNodeMap,rcvNodeMap,offset);
# Mask = convert(Array{Int8},Mask);

# rickt,zeroOffset,rickWidth = getRickerFunction(T,dt,fm);
# computeTimeDomain = true;

# if computeTimeDomain
	# pFor,SourcesSubInd = getTimeDomainFWIParam(gamma[:],Q,P,Mask,Minv,rickt,dt,T,workerList);
	# (DRF,pFor) = getData(velocityToSlowSquared(m[:])[1],pFor,ones(length(pFor)),true);
	# DobsTime = Array{Array{Float32,2}}(undef,size(Q,2));
	# for k=1:length(pFor)
		# Dt = fetch(DRF[k]);
		# for i = 1:length(SourcesSubInd[k])
			# DobsTime[SourcesSubInd[k][i]] = Dt[i];
		# end
		# Dt = 0;
	# end
	# file = matopen(string(filenamePrefix,"_timeDomain.mat"),"w");
	# write(file,"DobsTime",convert(Array{Array{Float32,2},1},DobsTime));
	# close(file);
# else
	# file = matopen(string(filenamePrefix,"_timeDomain.mat"));
	# DobsTime = read(file,"DobsTime");
	# DobsTime = convert(Array{Array{Float32,2},1},DobsTime);
	# close(file);
# end

# DobsOmHat = Array{Array{ComplexF64,2}}(undef,length(omega));
# for k=1:length(omega)
	# DobsOmHat[k] = zeros(ComplexF64,nrcv,nsrc);
# end

# nw = length(rickt);
# rickh = fft(rickt);
# rickhsub = zeros(ComplexF64,length(omega));

# for k = 1:length(DobsTime)
	# Dk = DobsTime[k];
	# Dk = Dk .+ sigmaNoise*mean(abs.(Dk))*randn(size(Dk));
	# DobsTime[k] = Dk;
	# Dhat = fft(Dk,2);
	# for jj = 1:length(omega)
		# iw = (nw*dt*omega[jj])/(2*pi) + 1
		# iw = convert(Int64,round(iw));
		# DobsOmHat[jj][:,k] = Dhat[:,iw]./rickh[iw];
		# rickhsub[jj] = rickh[iw];
	# end
# end
# for k = 1:length(omega)
	# omRound = string(round((omega[k]/(2*pi))*100.0)/100.0);
	# Dobsk = DobsOmHat[k];
	# Wd_k = (1.0./(abs.(real(Dobsk)).+0.1*mean(abs.(Dobsk)))) .+ 1im*(1.0./(abs.(imag(Dobsk)).+0.1*mean(abs.(Dobsk))));
	# # Wd_k = (1.0./(0.0*abs(real(Dobsk)).+1.0*mean(abs(Dobsk)))) + 1im*(1.0./(0.0*abs(imag(Dobsk))+1.0*mean(abs(Dobsk))));
	# Wd_k .*= Mask;
	# filename = string(dataFullFilenamePrefix,omRound,".dat");
	# writeDataFile(filename,Dobsk,Wd_k,srcNodeMap,rcvNodeMap);
# end
# return DobsTime,gamma,rickt,rickhsub,Mask,zeroOffset,rickWidth;
# end

# function prepareTravelTimeDataFiles(m, Minv::RegularMesh, filenamePrefix::String,dataFullFilename::String, offset::Int64,HO::Bool,useFilesForFields::Bool = false)

# ########################## m is in Velocity here. ###################################

# RCVfile = string(filenamePrefix,"_rcvMap.dat");
# SRCfile = string(filenamePrefix,"_srcMap.dat");
# srcNodeMap = readSrcRcvLocationFile(SRCfile,Minv);
# rcvNodeMap = readSrcRcvLocationFile(RCVfile,Minv);

# Q = generateSrcRcvProjOperators(Minv.n.+1,srcNodeMap);
# Q = Q.*(1.0./(norm(Minv.h)^2));
# P = generateSrcRcvProjOperators(Minv.n.+1,rcvNodeMap);

# # compute observed data

# println("~~~~~~~ Getting data Eikonal: ~~~~~~~");
# (pForEIK,contDivEIK,SourcesSubIndEIK) = getEikonalInvParam(Minv,Q,P,HO,nworkers(),useFilesForFields);

# (D,pForEIK) = getData(velocityToSlowSquared(m[:])[1],pForEIK,ones(length(pForEIK)),true);

# Dobs = Array{Array{Float64,2}}(undef,length(pForEIK))
# for k = 1:length(pForEIK)
	# Dobs[k] = fetch(D[k]);
# end
# Dobs = arrangeRemoteCallDataIntoLocalData(Dobs);

# # D should be of length 1 becasue constMUSTBeOne = 1;
# # Dobs += 0.01*mean(abs(Dobs))*randn(size(Dobs,1),size(Dobs,2));
# Wd = (1.0./(abs.(Dobs).+ 0.1*mean(abs.(Dobs))));
# Wd = limitDataToOffset(Wd,srcNodeMap,rcvNodeMap,offset);
# writeDataFile(string(dataFullFilename,".dat"),Dobs,Wd,srcNodeMap,rcvNodeMap);

# return (Dobs,Wd)
# end


text:
"""
    circle(; center::Union{Point, Position}, radius::Real=5., steps::Integer=64, units::String="kilometers")

Take a Point or a Position and calculate the circle polygon given a radius in degrees, radians, miles, or kilometers; and steps for precision.

# Examples
```jldoctest
julia> point = Point([35, 45])
Point([35.0, 45.0])

julia> circle(center=point, steps=5)
Polygon(Array{Array{Float64,1},1}[[[34.9395, 45.0139], [34.9626, 44.9636], [35.0374, 44.9636], [35.0605, 45.0139], [35.0, 45.045], [34.9395, 45.0139]]])

julia> circle(center=point, radius=2.5, steps=5, units="degrees")
Polygon(Array{Array{Float64,1},1}[[[31.5893, 45.7231], [32.99, 42.9571], [37.01, 42.9571], [38.4107, 45.7231], [35.0, 47.5029], [31.5893, 45.7231]]])
```
"""
function circle(; center::Union{Point, Position}, radius::Real=5., steps::Integer=64, units::String="kilometers")
    coords = []
    position::Position = []

    if geotype(center) === :Point
        position = Position(center.coordinates)
    else
        position = center
    end

    for i in 1:steps
        push!(coords, destination(position, Float64(radius), i * (-360 / steps), units).coordinates)
    end
    push!(coords, coords[1])

    return Polygon([coords])
end

"""
    sector(center::Point, radius::Real, bearing1::Real, bearing2::Real, steps::Integer=64, units::String="kilometers")

Creates a circular sector of a circle of given radius and center Point,
between (clockwise) bearing1 and bearing2; 0 bearing is North of center point, positive clockwise.

# Examples
```jldoctest
julia> point = Point([35, 45])
Point([35.0, 45.0])

julia> sector(point, 5, 0, 0.5, 5)
Polygon(Array{Array{Float64,1},1}[[[35.0, 45.0], [35.0, 45.045], [35.0006, 45.045], [35.0, 45.0]]])

julia> sector(point, 5, 0, 270, 5)
Polygon(Array{Array{Float64,1},1}[[[35.0, 45.0], [35.0, 45.045], [35.0605, 45.0139], [35.0374, 44.9636], [34.9626, 44.9636], [34.9364, 45.0], [35.0, 45.0]]])
```
"""
function sector(center::Point, radius::Real, bearing1::Real, bearing2::Real, steps::Integer=64, units::String="kilometers")
    complem(bearing1) === complem(bearing2) && return circle(center=center, radius=radius, steps=steps, units=units)

    coords = center.coordinates
    arc = linearc(center, radius, bearing1, bearing2, steps, units)
    sliceCoords = [[coords]]

    for i in eachindex(arc.coordinates)
        push!(sliceCoords[1], arc.coordinates[i])
    end
    push!(sliceCoords[1], coords)

    return Polygon(sliceCoords)
end

function complem(a::Real)
    b = a % 360
    b < 0 && (b += 360)

    return b
end


text:
include("capacity.jl")
include("rec.jl")
include("energy.jl")
include("raisereserve.jl")
include("lowerreserve.jl")

struct Markets{R,T,P}
   capacity::CapacityMarket
   rec::RECMarket
   energy::EnergyMarket{R,T,P}
   raisereserve::RaiseReserveMarket{R,T,P}
   lowerreserve::LowerReserveMarket{R,T,P}
end

function loadmarkets(
    regions::Dict{String,Int}, periods::Dict{String,Int}, n_timesteps::Int,
    resourcepath::String
)
    capacity = CapacityMarket(joinpath(resourcepath, "capacity"))

    rec = RECMarket(joinpath(resourcepath, "rec"))

    energy = loadmarket(
        EnergyMarket, regions, periods, n_timesteps,
        joinpath(resourcepath, "energy"))

    raisereserve = loadmarket(
        RaiseReserveMarket, regions, periods, n_timesteps,
        joinpath(resourcepath, "raisereserve"))

    lowerreserve = loadmarket(
        LowerReserveMarket, regions, periods, n_timesteps,
        joinpath(resourcepath, "lowerreserve"))

    return Markets(capacity, rec, energy, raisereserve, lowerreserve)

end

function loadmarket(
    market::Type, regions::Dict{String,Int}, periods::Dict{String,Int}, T::Int,
    marketpath::String
)

    R = length(regions)
    P = length(periods)

    rulesdata = DataFrame!(CSV.File(joinpath(marketpath, "rules.csv"),
                                    types=scenarios_market_param_types))
    demanddata = DataFrame!(CSV.File(joinpath(marketpath, "demand.csv")))
    demanddata = stack(demanddata, Not([:period, :timestep]),
                       variable_name=:region, value_name=:demand)

    pricecap = zeros(Float64, R)
    demand = zeros(Float64, R, T, P)

    for row in eachrow(rulesdata)
        r_idx = regions[row.region]
        pricecap[r_idx] = row.pricecap
    end

    for row in eachrow(demanddata)
        r_idx = regions[row.region]
        t = row.timestep
        p_idx = periods[row.period]
        demand[r_idx,t,p_idx] = row.demand
    end

    return market(demand, pricecap)

end

welfare(x::Markets) =
    welfare(x.capacity) + welfare(x.rec) + welfare(x.energy) +
    welfare(x.raisereserve) + welfare(x.lowerreserve)


text:
using MyPkg2
using Test

@testset "MyPkg2.jl" begin
    # Write your own tests here.

    @test π == f()
    @test 100 == g()
    @testset "func" begin
        @test 3 == func(1, 2)
        @test 4 == func(1, 3)
    end

end


text:
# Value Function Iteration with IID Income
# Greg Kaplan 2017
# Translated by Tom Sweeney Dec 2020

using Random, NLsolve, Plots
include("discrete_normal.jl")

# PARAMETERS

## preferences
risk_aver = 2
beta = 0.95

## returns
r = 0.03
R = 1+r

## income risk: discretized N(mu,sigma^2)
mu_y = 1
sd_y = 0.2
ny = 5

## asset grids
na = 500
amax = 20
borrow_lim = 0
agrid_par = 1 # 1 for linear, 0 for L-shaped

## computation
max_iter = 1000
tol_iter = 1.0e-6
Nsim = 50000
Tsim = 500

# OPTIONS
Display = 1
DoSimulate = 1
MakePlots = 1

# DRAW RANDOM NUMBERS

Random.seed!(2020)
yrand = rand(Nsim,Tsim)

# SET UP GRIDS

## assets
agrid = range(0,1,length=na)
agrid = agrid.^(1 ./ agrid_par)
agrid = borrow_lim .+ (amax.-borrow_lim).*agrid

## income: disretize normal distribution
width = nlsolve(x -> discrete_normal(ny,mu_y,sd_y,x...)[1],[2.0]).zero
temp, ygrid, ydist = discrete_normal(ny,mu_y,sd_y,width...)
ycumdist = cumsum(ydist)

# UTILITY FUNCTION

if risk_aver==1
    u(c) = log.(c)
else
    u(c) = (c.^(1-risk_aver).-1)./(1-risk_aver)
end

# INITIALIZE VALUE FUNCTION

Vguess = zeros(na,ny)
for iy = 1:ny
    Vguess[:,iy] .= u(r.*agrid[1]+ygrid[iy])./(1-beta)
end

### Vguess = ones(na,ny)

# ITERATE ON VALUE FUNCTION

V = copy(Vguess)

Vdiff = 1
iter = 0

while iter <= max_iter && Vdiff > tol_iter
    iter = iter + 1
    Vlast = copy(V)
    V = zeros(na,ny)
    global sav = zeros(na,ny)
    global savind = zeros(Int,na,ny)
    global con = zeros(na,ny)

    ## loop over assets
    for ia = 1:na

        ## loop over income
        for iy = 1:ny
            cash = R.*agrid[ia] + ygrid[iy]
            Vchoice = u(max.(cash.-agrid,1.0e-10)) .+ beta.*(Vlast*ydist)
            V[ia,iy] = maximum(Vchoice)
            savind[ia,iy] = argmax(Vchoice)[1]
            sav[ia,iy] = agrid[savind[ia,iy]]
            con[ia,iy] = cash .- sav[ia,iy]
        end
    end

    Vdiff = maximum(abs.(V-Vlast))
    if Display >= 1
        println("Iteration no. " * string(iter), " max val fn diff is " * string(Vdiff))
    end
end

# SIMULATE
if DoSimulate == 1
    yindsim = zeros(Int,Nsim,Tsim)
    aindsim = zeros(Int,Nsim,Tsim)

    ## initial assets
    aindsim[:,1] .= 1

    ## loop over time periods
    for it = 1:Tsim
        if Display >= 1 && mod(it,100)==0
            println(" Simulating, time period " * string(it))
        end

        ### income realization: note we vectorize simulations at once because
        ### of matlab, in other languages we would loop over individuals
        yindsim[yrand[:,it].<=ycumdist[1],it] .= 1
        for iy = 2:ny
            yindsim[(yrand[:,it].>ycumdist[iy-1]) .& (yrand[:,it].<=ycumdist[iy]),it] .= iy
        end

        ## asset choice
        if it < Tsim
            for iy = 1:ny
                aindsim[yindsim[:,it].==iy,it+1] = savind[aindsim[yindsim[:,it].==iy,it],iy]
            end
        end
    end

    ## assign actual asset and income values
    asim = agrid[aindsim]
    ysim = ygrid[yindsim]
end

# MAKE PLOTS
if MakePlots==1

    ## consumption policy function
    p1 = plot(agrid, [con[:,1] con[:,ny]], xlims=(0,amax), title="Consumption", color=[:blue :red], label=["Lowest income state" "Highest income state"])
    display(p1)

    ## savings policy function
    p2 = plot(agrid, [sav[:,1].-agrid[:,1] sav[:,ny].-agrid[:,1]], xlims=(0,amax), title="Savings", color=[:blue :red], legend=false)
    plot!(agrid, zeros(na,1), color=:black, lw=0.5)
    display(p2)

    ## nice zoom
    xlimits = (0,1)
    xlimind = trues(na)
    if minimum(agrid) < xlimits[1]
        xlimind = xlimind .& (agrid.>=maximum(agrid[agrid<xlimits[1]]))
    elseif minimum(agrid) > xlimits[2]
        xlimind .= 0
    end
    if maximum(agrid) > xlimits[2]
        xlimind = xlimind .& (agrid.<=minimum(agrid[agrid.>xlimits[2]]))
    elseif maximum(agrid) < xlimits[1]
        xlimind .= 0
    end

    ## consumption policy function: zoomed in
    p3 = plot(agrid[xlimind], [con[xlimind,1] con[xlimind,ny]], xlims=xlimits, title="Consumption: Zoomed", marker=:circle, color=[:blue :red], linewidth=2, legend=false)
    plot!(show=true)
    display(p3)

    ## savings policy function: zoomed in
    p4 = plot(agrid[xlimind], [sav[xlimind,1].-agrid[xlimind] sav[xlimind,ny].-agrid[xlimind]], xlims=xlimits, title="Savings: Zoomed (a'-a)", marker=:circle, color=[:blue :red], linewidth=2, legend=false)
    plot!(agrid, zeros(na,1), color=:black, lw=0.5)
    display(p4)

    ## income distribution
    p5 = histogram(ysim[:,Tsim], bins=[2*ygrid[1]-ygrid[2];ygrid].+(ygrid[2]-ygrid[1])/2, title="Income distribution", color=RGB(0,0.5,0.5), linecolor=:blue, legend=false)
    display(p5)

    ## asset distribution
    p6 = histogram(asim[:,Tsim], bins=-0.025:0.05:1.975, title="Asset distribution", color=RGB(.7,.7,.7), linecolor=:black, legend=false)
    display(p6)

    ## convergence check
    p7 = plot(1:Tsim, mean(asim,dims=1)', title="Mean Asset Convergence", xlabel="Time Period", color=:black, lw=1.5, legend=false)
    display(p7)

    ## asset distribution statistics
    aysim = asim[:,Tsim]./mean(ysim[:,Tsim])
    println("Mean assets: " * string(mean(aysim)))
    println("Fraction borrowing constrained: " * string(sum(aysim.==borrow_lim)/Nsim * 100) * '%')
    println("10th Percentile: " * string(quantile(aysim,.1)))
    println("50th Percentile: " * string(quantile(aysim,.5)))
    println("90th Percentile: " * string(quantile(aysim,.9)))
    println("99th Percentile: " * string(quantile(aysim,.99)))
end