# BlockBandedMatrix with block range indexes is also block-banded
const Block1 = Block{1,Int}
const BlockRange1 = BlockRange{1,Tuple{UnitRange{Int}}}
const BlockIndexRange1 = BlockIndexRange{1,Tuple{UnitRange{Int}}}
const SubBlockSkylineMatrix{T,LL,UU,R1,R2} =
    SubArray{T,2,BlockSkylineMatrix{T,LL,UU},Tuple{BlockSlice{R1},BlockSlice{R2}}}

const SubBandedBlockBandedMatrix{T,R1,R2} =
    SubArray{T,2,<:BandedBlockBandedMatrix{T},Tuple{BlockSlice{R1},BlockSlice{R2}}}

getindex(A::BandedBlockBandedMatrix, KR::BlockRange1, JR::BlockRange1) = BandedBlockBandedMatrix(view(A, KR, JR))
getindex(A::BandedBlockBandedMatrix, KR::BlockRange1, J::Block1) = BandedBlockBandedMatrix(view(A, KR, J))
getindex(A::BandedBlockBandedMatrix, K::Block1, JR::BlockRange1) = BandedBlockBandedMatrix(view(A, K, JR))

@lazymul AbstractBlockBandedMatrix

MemoryLayout(A::Type{<:PseudoBlockArray{T,N,R}}) where {T,N,R} = MemoryLayout(R)

#############
# BLAS overrides
#############

function materialize!(M::MatMulVecAdd{<:AbstractBlockBandedLayout,<:AbstractStridedLayout,<:AbstractStridedLayout})
    α, A, x_in, β, y_in = M.α, M.A, M.B, M.β, M.C
    if length(x_in) != size(A,2) || length(y_in) != size(A,1)
        throw(DimensionMismatch())
    end

    # impose block structure
    y = PseudoBlockArray(y_in, BlockSizes((cumulsizes(blocksizes(A),1),)))
    x = PseudoBlockArray(x_in, BlockSizes((cumulsizes(blocksizes(A),2),)))

    _fill_lmul!(β, y)

    for J = Block.(1:nblocks(A,2))
        for K = blockcolrange(A,J)
            muladd!(α, view(A,K,J), view(x,J), one(α), view(y,K))
        end
    end
    y_in
end

function materialize!(M::MatMulMatAdd{<:AbstractBlockBandedLayout,<:AbstractBlockBandedLayout,<:AbstractBlockBandedLayout})
    α, A, X, β, Y = M.α, M.A, M.B, M.β, M.C
    _fill_lmul!(β, Y)
    for J=Block(1):Block(nblocks(X,2)), N=blockcolrange(X,J), K=blockcolrange(A,N)
        muladd!(α, view(A,K,N), view(X,N,J), one(α), view(Y,K,J))
    end
    Y
end

function materialize!(M::MatMulMatAdd{<:AbstractBlockBandedLayout,<:AbstractColumnMajor,<:AbstractColumnMajor})
    α, A, X_in, β, Y_in = M.α, M.A, M.B, M.β, M.C
    _fill_lmul!(β, Y_in)
    X = PseudoBlockArray(X_in, BlockSizes((cumulsizes(blocksizes(A),2),[1,size(X_in,2)+1])))
    Y = PseudoBlockArray(Y_in, BlockSizes((cumulsizes(blocksizes(A),1), [1,size(Y_in,2)+1])))
    for N=Block.(1:nblocks(X,1)), K=blockcolrange(A,N)
        muladd!(α, view(A,K,N), view(X,N,Block(1)), one(α), view(Y,K,Block(1)))
    end
    Y_in
end

function materialize!(M::MatMulMatAdd{<:AbstractColumnMajor,<:AbstractBlockBandedLayout,<:AbstractColumnMajor})
    α, A_in, X, β, Y_in = M.α, M.A, M.B, M.β, M.C
    _fill_lmul!(β, Y_in)
    A = PseudoBlockArray(A_in, BlockSizes(([1,size(A_in,1)+1],cumulsizes(blocksizes(X),1))))
    Y = PseudoBlockArray(Y_in, BlockSizes(([1,size(Y_in,1)+1],cumulsizes(blocksizes(X),2))))
    for J=Block(1):Block(nblocks(X,2)), N=blockcolrange(X,J)
        muladd!(α, view(A,Block(1),N), view(X,N,J), one(α), view(Y,Block(1),J))
    end
    Y_in
end

#############
# * overrides
#############

*(A::BlockBandedMatrix, B::BlockBandedMatrix) = mul(A,B)
*(A::BlockBandedMatrix, B::BandedBlockBandedMatrix) = mul(A,B)
*(A::BandedBlockBandedMatrix, B::BlockBandedMatrix) = mul(A,B)
*(A::BandedBlockBandedMatrix, B::BandedBlockBandedMatrix) = mul(A,B)
*(A::Matrix, B::BlockBandedMatrix) = mul(A,B)
*(A::BlockBandedMatrix, B::Matrix) = mul(A,B)
*(A::BandedBlockBandedMatrix, B::Matrix) = mul(A,B)
*(A::Matrix, B::BandedBlockBandedMatrix) = mul(A,B)

function add_bandwidths(A::AbstractBlockBandedMatrix,B::AbstractBlockBandedMatrix)
    Al,Au = colblockbandwidths(A)
    Bl,Bu = colblockbandwidths(B)

    l = Vector(Bl)
    u = Vector(Bu)

    for (v,Av) in [(l,Al),(u,Au)]
        n = length(v)
        for i = 1:n
            sel = max(i-Bu[i],1):min(i+Bl[i],length(Av))
            isempty(sel) && continue
            v[i] += maximum(Av[sel])
        end
    end

    l,u
end

function add_bandwidths(A::BlockBandedMatrix,B::BlockBandedMatrix)
    l,u = blockbandwidths(A) .+ blockbandwidths(B)
    Fill(l,nblocks(B,2)), Fill(u,nblocks(B,2))
end

function similar(M::MulAdd{<:AbstractBlockBandedLayout,<:AbstractBlockBandedLayout}, ::Type{T}) where T
    A,B = M.A, M.B

    Arows, Acols = A.block_sizes.block_sizes.cumul_sizes
    Brows, Bcols = B.block_sizes.block_sizes.cumul_sizes
    if Acols ≠ Brows
        # diagonal matrices can be converted
        if isdiag(B) && size(A,2) == size(B,1) == size(B,2)
            B = BlockBandedMatrix(B.data, BlockSizes((Acols,Acols)), 0, 0, 0, 0)
        elseif isdiag(A) && size(A,2) == size(B,1) == size(A,1)
            A = BlockBandedMatrix(A.data, BlockSizes((Brows,Brows)), 0, 0, 0, 0)
        else
            throw(DimensionMismatch("*"))
        end
    end
    n,m = size(A,1), size(B,2)

    l,u = add_bandwidths(A,B)
    BlockSkylineMatrix{T}(undef, BlockSkylineSizes(BlockSizes((Arows,Bcols)), l, u))
end

function similar(M::MulAdd{BandedBlockBandedColumnMajor,BandedBlockBandedColumnMajor}, ::Type{T}) where T
    A,B = M.A, M.B

    Arows, Acols = A.block_sizes.block_sizes.cumul_sizes
    Brows, Bcols = B.block_sizes.block_sizes.cumul_sizes
    if Acols ≠ Brows
        # diagonal matrices can be converted
        if isdiag(B) && size(A,2) == size(B,1) == size(B,2)
            # TODO: fix
            B = BandedBlockBandedMatrix(B.data, BlockSizes((Acols,Acols)), 0, 0, 0, 0)
        elseif isdiag(A) && size(A,2) == size(B,1) == size(A,1)
            A = BandedBlockBandedMatrix(A.data, BlockSizes((Brows,Brows)), 0, 0, 0, 0)
        else
            throw(DimensionMismatch("*"))
        end
    end
    n,m = size(A,1), size(B,2)

    bs = BandedBlockBandedSizes(BlockSizes((Arows,Bcols)), A.l+B.l, A.u+B.u, A.λ+B.λ, A.μ+B.μ)

    BandedBlockBandedMatrix{T}(undef, bs)
end

similar(M::MulAdd{<:DiagonalLayout,<:AbstractBlockBandedLayout}, ::Type{T}) where T = similar(M.B,T)
similar(M::MulAdd{<:AbstractBlockBandedLayout,<:DiagonalLayout}, ::Type{T}) where T = similar(M.A,T)

function blocksizes(V::SubBlockSkylineMatrix{<:Any,LL,UU,BlockRange1,BlockRange1}) where {LL,UU}
    A = parent(V)
    Bs = A.block_sizes.block_sizes

    KR = parentindices(V)[1].block.indices[1]
    JR = parentindices(V)[2].block.indices[1]
    shift = Int(KR[1])-Int(JR[1])

    Bs.cumul_sizes[1]
    @assert KR[1] == JR[1] == 1
    BlockSkylineSizes(BlockSizes((Bs.cumul_sizes[1][KR[1]:KR[end]+1] .- Bs.cumul_sizes[1][KR[1]] .+ 1,
                                       Bs.cumul_sizes[2][JR[1]:JR[end]+1] .- Bs.cumul_sizes[2][JR[1]] .+ 1)),
                           colblockbandwidth(A,1)[1:Int(JR[end])] .- shift, colblockbandwidth(A,2)[1:Int(JR[end])] .+ shift)
end

function blockbandwidths(V::SubBlockSkylineMatrix{<:Any,LL,UU,BlockRange1,BlockRange1}) where {LL,UU}
    A = parent(V)
    Bs = A.block_sizes.block_sizes

    KR = parentindices(V)[1].block.indices[1]
    JR = parentindices(V)[2].block.indices[1]

    @assert KR[1] == JR[1] == 1
    blockbandwidths(A)
end

####
# BlockIndexRange subblocks
####

sublayout(::BlockBandedColumnMajor, ::Type{Tuple{BlockSlice{BlockRange1}, BlockSlice{BlockRange1}}}) = BlockBandedColumnMajor()
sublayout(::BlockBandedColumnMajor, ::Type{Tuple{BlockSlice{BlockRange1}, BlockSlice{Block1}}}) = ColumnMajor()
sublayout(::BlockBandedColumnMajor, ::Type{Tuple{BlockSlice{BlockRange1}, BlockSlice{BlockIndexRange1}}}) = ColumnMajor()
sublayout(::BlockBandedColumnMajor, ::Type{Tuple{BlockSlice{BlockIndexRange1}, BlockSlice{BlockIndexRange1}}}) = ColumnMajor()

sublayout(::BandedBlockBandedColumnMajor, ::Type{Tuple{BlockSlice{Block1}, BlockSlice{Block1}}}) = BandedColumnMajor()
sublayout(::BandedBlockBandedColumnMajor, ::Type{Tuple{BlockSlice{BlockRange1}, BlockSlice{BlockRange1}}}) = BandedBlockBandedColumnMajor()
sublayout(::BandedBlockBandedColumnMajor, ::Type{Tuple{BlockSlice{Block1}, BlockSlice{BlockRange1}}}) = BandedBlockBandedColumnMajor()
sublayout(::BandedBlockBandedColumnMajor, ::Type{Tuple{BlockSlice{BlockRange1}, BlockSlice{Block1}}}) = BandedBlockBandedColumnMajor()
sublayout(::BandedBlockBandedColumnMajor, ::Type{Tuple{BlockSlice{BlockRange1}, BlockSlice{BlockIndexRange1}}}) = BandedBlockBandedColumnMajor()
sublayout(::BandedBlockBandedColumnMajor, ::Type{Tuple{BlockSlice{BlockIndexRange1}, BlockSlice{BlockIndexRange1}}}) = BandedBlockBandedColumnMajor()

isbanded(A::SubArray{<:Any,2,<:BandedBlockBandedMatrix}) = MemoryLayout(typeof(A)) == BandedColumnMajor()
isbandedblockbanded(A::SubArray{<:Any,2,<:BandedBlockBandedMatrix}) = MemoryLayout(typeof(A)) == BandedBlockBandedColumnMajor()

subblockbandwidths(V::SubBandedBlockBandedMatrix) = subblockbandwidths(parent(V))

function blockbandwidths(V::SubBandedBlockBandedMatrix{<:Any,BlockRange1,Block1})
    A = parent(V)
    Bs = A.block_sizes.block_sizes

    KR = parentindices(V)[1].block.indices[1]
    J = parentindices(V)[2].block
    shift = Int(KR[1])-Int(J)
    blockbandwidth(A,1) - shift, blockbandwidth(A,2) + shift
end

function blockbandwidths(V::SubBandedBlockBandedMatrix{<:Any,Block1,BlockRange1})
    A = parent(V)
    Bs = A.block_sizes.block_sizes

    K = parentindices(V)[1].block
    JR = parentindices(V)[2].block.indices[1]
    shift = Int(K)-Int(JR[1])

    blockbandwidth(A,1) - shift, blockbandwidth(A,2) + shift
end

function blockbandwidths(V::SubBandedBlockBandedMatrix{<:Any,BlockRange1,BlockRange1})
    A = parent(V)
    Bs = A.block_sizes.block_sizes

    KR = parentindices(V)[1].block.indices[1]
    JR = parentindices(V)[2].block.indices[1]
    shift = Int(KR[1])-Int(JR[1])

    blockbandwidth(A,1) - shift, blockbandwidth(A,2) + shift
end

strides(V::SubBlockSkylineMatrix{<:Any,LL,UU,<:Union{BlockRange1,Block1},Block1}) where {LL,UU} =
    (1,parent(V).block_sizes.block_strides[Int(parentindices(V)[2].block)])

function unsafe_convert(::Type{Ptr{T}}, V::SubBlockSkylineMatrix{T,LL,UU,<:Union{BlockRange1,Block1},Block1}) where {T,LL,UU}
    A = parent(V)
    JR = parentindices(V)[2]
    KR = parentindices(V)[1].block
    J = parentindices(V)[2].block
    p = unsafe_convert(Ptr{T}, view(A, first(KR), J))
end

struct ShiftedLayout{T,ML<:MemoryLayout} <: MemoryLayout
    shift::Tuple{Int,Int}  # gives the shift to the start of the memory.
                           # So shift == (0,0) is equivalent to layout
                           # shift == (2,1) has the first two rows and first column all zero
    layout::ML
end

strides(V::SubBlockSkylineMatrix{<:Any,LL,UU,BlockRange1,BlockIndexRange1}) where {LL,UU} =
    (1,parent(V).block_sizes.block_strides[Int(Block(parentindices(V)[2]))])

function unsafe_convert(::Type{Ptr{T}}, V::SubBlockSkylineMatrix{T,LL,UU,BlockRange1,BlockIndexRange1}) where {T,LL,UU}
    A = parent(V)
    JR = parentindices(V)[2]
    K = first(parentindices(V)[1].block)
    J = Block(JR)
    K ∈ blockcolrange(A, J) || throw(ArgumentError("Pointer is only defined when inside colrange"))
    p = unsafe_convert(Ptr{T}, view(A, K, J))
    p + sizeof(T)*(JR.block.indices[1][1]-1)*stride(V,2)
end

function unsafe_convert(::Type{Ptr{T}}, V::SubBlockSkylineMatrix{T,LL,UU,BlockIndexRange1,BlockIndexRange1}) where {T,LL,UU}
    A = parent(V)
    JR = parentindices(V)[2]
    K = parentindices(V)[1].block.block
    kr = parentindices(V)[1].block.indices[1]
    J = parentindices(V)[2].block.block
    jr = parentindices(V)[2].block.indices[1]
    p = unsafe_convert(Ptr{T}, view(A, K, J))
    p + sizeof(T)*(kr[1]-1 + (jr[1]-1)*stride(V,2))
end

strides(V::SubBlockSkylineMatrix{T,LL,UU,BlockIndexRange1,BlockIndexRange1}) where {T,LL,UU} =
    (1,parent(V).block_sizes.block_strides[Int(parentindices(V)[2].block.block)])

MemoryLayout(V::SubBlockSkylineMatrix{T,LL,UU,BlockIndexRange1,BlockIndexRange1}) where {T,LL,UU} = ColumnMajor()


module Optimisers

using Functors: functor, fmap, isleaf
using LinearAlgebra

include("interface.jl")

include("destructure.jl")
export destructure, total, total2

include("rules.jl")
export Descent, ADAM, Momentum, Nesterov, RMSProp,
       ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, ADAMW, RADAM, OADAM, AdaBelief,
       WeightDecay, ClipGrad, ClipNorm, OptimiserChain

"""
    Optimisers.apply!(rule::RuleType, state, parameters, gradient) -> (state, gradient)

This defines the action of any optimisation rule. It should return the modified gradient
which will be subtracted from the parameters, and the updated state (if any) for use at
the next iteration, as a tuple `(state, gradient)`.

For efficiency it is free to mutate the old state, but only what is returned will be used.
Ideally this should check `iswriteable(x)`, which the built-in rules do via [`@..`](@ref).

The initial state is `init(rule::RuleType, parameters)`.

# Example
```jldoctest
julia> Optimisers.init(Descent(0.1), Float32[1,2,3]) === nothing
true

julia> Optimisers.apply!(Descent(0.1), nothing, Float32[1,2,3], [4,5,6])
(nothing, Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}}(*, ([4, 5, 6], 0.1f0)))
```
"""
apply!

"""
    Optimisers.init(rule::RuleType, parameters) -> state

Sets up the initial state for a given optimisation rule, and an array of parameters.
This and [`apply!`](@ref) are the two functions which any new optimisation rule must define.

# Examples
```jldoctest
julia> Optimisers.init(Descent(), Float32[1,2,3])  # is `nothing`

julia> Optimisers.init(Momentum(), [1.0, 2.0])
2-element Vector{Float64}:
 0.0
 0.0
```
"""
init

"""
    Optimisers.setup(rule, model) -> tree

Initialises the given optimiser for every trainable parameter within the model.
Returns a tree of the relevant states, which must be passed to [`update`](@ref)
or [`update!`](@ref).

# Example
```jldoctest
julia> Optimisers.setup(Descent(0.1f0), (x = rand(3), y = (true, false), z = tanh))
(x = Leaf(Descent{Float32}(0.1), nothing), y = (nothing, nothing), z = nothing)
```
"""
setup

"""
    Optimisers.update(tree, model, gradient) -> (tree, model)

Uses the optimiser and the gradient to change the trainable parameters in the model.
Returns the improved model, and the optimiser states needed for the next update.
The initial tree of states comes from [`setup`](@ref).

See also [`update!`](@ref), which will be faster for models of ordinary `Array`s or `CuArray`s.

# Example
```jldoctest
julia> m = (x = Float32[1,2,3], y = tanh);

julia> t = Optimisers.setup(Descent(0.1f0), m)
(x = Leaf(Descent{Float32}(0.1), nothing), y = nothing)

julia> g = (x = [1,1,1], y = nothing);  # fake gradient

julia> Optimisers.update(t, m, g)
((x = Leaf(Descent{Float32}(0.1), nothing), y = nothing), (x = Float32[0.9, 1.9, 2.9], y = tanh))
```
"""
update

"""
    Optimisers.update!(tree, model, gradient) -> (tree, model)

Uses the optimiser and the gradient to change the trainable parameters in the model.
Returns the improved model, and the optimiser states needed for the next update.
The initial tree of states comes from [`setup`](@ref).

This is used in exactly the same manner as [`update`](@ref), but because it may mutate
arrays within the old model (and the old state), it will be faster for models of ordinary
`Array`s or `CuArray`s. However, you should not rely on the old model being fully updated.
"""
update!

end # module


abstract type AbstractNode{T, N, p} <: AbstractArray{T, N} end

@pure Base.size(::Type{Tnode}) where {Tnode <: AbstractNode} = convert.(Int, TreeSize(Tnode)[1])
@pure totalsize(::Type{Tnode}) where {Tnode <: AbstractNode} = totalsize(TreeSize(Tnode))

# null
@pure null(::Type{Tnode}) where {Tnode <: AbstractNode} = Tnode(undef)
null(x::AbstractNode) = null(typeof(x))
isnull(x::AbstractNode) = x === null(x)

# childtype
@pure childtype(::Type{<: AbstractNode}) = nothing
@pure childtype(::Type{<: AbstractNode{T}}) where {T <: AbstractNode} = T
childtype(x::AbstractNode) = childtype(typeof(x))

# leaftype
@pure leaftype(T::Type{<: AbstractNode}) = T
@pure leaftype(::Type{<: AbstractNode{T}}) where {T <: AbstractNode} = leaftype(T)
leaftype(x::AbstractNode) = leaftype(typeof(x))

# leafeltype
@pure leafeltype(T::Type{<: AbstractNode}) = eltype(leaftype(T))
leafeltype(x::AbstractNode) = leafeltype(typeof(x))

# leafblock
nleafblocks(node::AbstractNode) = Int.(totalsize(node) .÷ size(leaftype(node)))
leafblockunit(node::AbstractNode) = TreeSize(node)[end]

getmask(x::AbstractNode) = getmask(x.data)
isactive(x::AbstractNode, i...) = (@_propagate_inbounds_meta; isnull(x) ? false : getmask(x)[i...])

@inline function unsafe_getindex(x::AbstractNode, i::Int)
    @boundscheck checkbounds(x, i)
    @inbounds unsafe_getindex(x.data, i)
end

@inline function unsafe_setindex!(x::AbstractNode, v, i::Int)
    @boundscheck checkbounds(x, i)
    @inbounds unsafe_setindex!(x.data, v, i)
end

@inline function unsafe_getindex(x::AbstractNode, i::Int...)
    @boundscheck checkbounds(x, i...)
    @inbounds unsafe_getindex(x, Base._sub2ind(size(x), i...))
end

@inline function unsafe_setindex!(x::AbstractNode, v, i::Int...)
    @boundscheck checkbounds(x, i...)
    @inbounds unsafe_setindex!(x, v, Base._sub2ind(size(x), i...))
end

function allocate!(node::AbstractNode, mask::AbstractArray{Bool})
    checkbounds(CartesianIndices(totalsize(node)), CartesianIndices(mask))
    dims = totalsize(Base.tail(TreeSize(node)))
    for I in CartesianIndices(node)
        start = CartesianIndex(@. dims * ($Tuple(I) - 1) + 1)
        stop = min(CartesianIndex(@. $Tuple(start) + dims - 1), last(CartesianIndices(mask)))
        childmask = @view mask[start:stop]
        if any(childmask)
            child = allocate!(node, I)
            allocate!(child, childmask)
        end
    end
end

function cleanup!(x::AbstractNode)
    @inbounds for i in eachindex(x)
        if isactive(x, i)
            # if `i` is active, then child node is always not null
            childnode = unsafe_getindex(x, i)
            cleanup!(childnode)
            !any(getmask(childnode)) && delete!(x, i)
        else
            delete!(x, i)
        end
    end
    x
end

function nleaves(x::AbstractNode)
    isnull(x) && return 0
    count = 0
    @inbounds for i in eachindex(x)
        if isactive(x, i)
            child = unsafe_getindex(x, i)
            count += nleaves(child)
        end
    end
    count
end

struct Dynamic end
struct TreeSize{S}
    function TreeSize{S}() where {S}
        new{S::Tuple{Vararg{Union{Tuple{Vararg{Integer}}, Dynamic}}}}()
    end
end
@pure TreeSize(S::Tuple) = TreeSize{S}()
@pure Base.Tuple(::TreeSize{S}) where {S} = S

@pure nodesize(::Type{<: AbstractNode{<: Any, <: Any, Dynamic()}}) = Dynamic()
@pure nodesize(::Type{<: AbstractNode{<: Any, N, p}}) where {N, p} = nfill(Power2(p), Val(N))
@pure TreeSize(::Nothing) = TreeSize{()}()
@pure TreeSize(::Type{Tnode}) where {Tnode <: AbstractNode} = TreeSize{(nodesize(Tnode), Tuple(TreeSize(childtype(Tnode)))...)}()
@pure TreeSize(x::AbstractNode) = TreeSize(typeof(x))

@pure Base.length(::TreeSize{S}) where {S} = length(S)
@pure Base.getindex(::TreeSize{S}, i::Int) where {S} = S[i]
@pure Base.firstindex(::TreeSize) = 1
@pure Base.lastindex(x::TreeSize) = length(x)
@pure Base.tail(::TreeSize{S}) where {S} = length(S) == 1 ? TreeSize{(map(zero, S[1]),)}() : TreeSize{Base.tail(S)}()
@pure totalsize(::TreeSize{S}) where {S} = broadcast(*, S...)

Base.show(io::IO, ::TreeSize{S}) where {S} = print(io, "TreeSize", S)

struct Activated{P, I}
    parent::P
    index::I
end
@inline set!(x::Activated, v) = @inbounds unsafe_setindex!(x.parent, v, x.index)
@inline set!(x::Activated, name::Symbol, v) = @inbounds unsafe_setindex!(getproperty(x.parent, name), v, x.index)
@inline function allocate!(x::Activated, i...)
    @_propagate_inbounds_meta
    child = @inbounds unsafe_getindex(x.parent, x.index)
    allocate!(child, i...)
end
@inline function Base.setindex!(x::Activated, v, i...)
    @_propagate_inbounds_meta
    child = @inbounds unsafe_getindex(x.parent, x.index)
    setindex!(child, v, i...)
end


# Test code on a relatively large molecule; my old friend C60

# Geom: D. Dragoe
# http://www.ccl.net/chemistry/resources/messages/2000/12/07.005-dir/
C60 = Quante.Molecule([
Atom( 6 , 0.726656 , -1.000157 , 3.300459),
Atom( 6 , 1.175755 , 0.382026 , 3.300459),
Atom( 6 , 1.410183 , -1.940951 , 2.581756),
Atom( 6 , 2.281725 , 0.741377 , 2.581756),
Atom( 6 , 2.281725 , 1.977639 , 1.817704),
Atom( 6 , 0.000000 , 1.236262 , 3.300459),
Atom( 6 , 0.000000 , 2.399147 , 2.581756),
Atom( 6 , 1.175755 , 2.781173 , 1.817704),
Atom( 6 , -0.683527 , -2.941108 , 1.817704),
Atom( 6 , -1.410183 , -1.940951 , 2.581756),
Atom( 6 , 0.683527 , -2.941108 , 1.817704),
Atom( 6 , -0.726656 , -1.000157 , 3.300459),
Atom( 6 , -1.175755 , 0.382026 , 3.300459),
Atom( 6 , -2.585938 , -1.558925 , 1.817704),
Atom( 6 , -3.008381 , -0.258779 , 1.817704),
Atom( 6 , -2.281725 , 0.741377 , 2.581756),
Atom( 6 , 0.726656 , -3.399304 , -0.581443),
Atom( 6 , -0.726656 , -3.399304 , -0.581443),
Atom( 6 , 1.410183 , -3.177213 , 0.581443),
Atom( 6 , -1.410183 , -3.177213 , 0.581443),
Atom( 6 , -2.585938 , -2.322977 , 0.581443),
Atom( 6 , -1.175755 , -2.781173 , -1.817704),
Atom( 6 , -2.281725 , -1.977639 , -1.817704),
Atom( 6 , -3.008381 , -1.741534 , -0.581443),
Atom( 6 , 3.008381 , -1.741534 , -0.581443),
Atom( 6 , 2.281725 , -1.977639 , -1.817704),
Atom( 6 , 2.585938 , -2.322977 , 0.581443),
Atom( 6 , 1.175755 , -2.781173 , -1.817704),
Atom( 6 , 0.000000 , -2.399147 , -2.581756),
Atom( 6 , 2.281725 , -0.741377 , -2.581756),
Atom( 6 , 1.175755 , -0.382026 , -3.300459),
Atom( 6 , 0.000000 , -1.236262 , -3.300459),
Atom( 6 , 3.008381 , -0.258779 , 1.817704),
Atom( 6 , 3.457479 , 0.359351 , 0.581443),
Atom( 6 , 2.585938 , -1.558925 , 1.817704),
Atom( 6 , 3.457479 , -0.359351 , -0.581443),
Atom( 6 , 3.008381 , 0.258779 , -1.817704),
Atom( 6 , 3.008381 , 1.741534 , 0.581443),
Atom( 6 , 2.585938 , 2.322977 , -0.581443),
Atom( 6 , 2.585938 , 1.558925 , -1.817704),
Atom( 6 , -0.726656 , 1.000157 , -3.300459),
Atom( 6 , 0.726656 , 1.000157 , -3.300459),
Atom( 6 , 1.410183 , 1.940951 , -2.581756),
Atom( 6 , -1.410183 , 1.940951 , -2.581756),
Atom( 6 , -3.008381 , 0.258779 , -1.817704),
Atom( 6 , -2.281725 , -0.741377 , -2.581756),
Atom( 6 , -1.175755 , -0.382026 , -3.300459),
Atom( 6 , -2.585938 , 1.558925 , -1.817704),
Atom( 6 , -3.008381 , 1.741534 , 0.581443),
Atom( 6 , -3.457479 , 0.359351 , 0.581443),
Atom( 6 , -3.457479 , -0.359351 , -0.581443),
Atom( 6 , -2.585938 , 2.322977 , -0.581443),
Atom( 6 , -0.726656 , 3.399304 , 0.581443),
Atom( 6 , -1.175755 , 2.781173 , 1.817704),
Atom( 6 , -2.281725 , 1.977639 , 1.817704),
Atom( 6 , -1.410183 , 3.177213 , -0.581443),
Atom( 6 , 0.683527 , 2.941108 , -1.817704),
Atom( 6 , 1.410183 , 3.177213 , -0.581443),
Atom( 6 , 0.726656 , 3.399304 , 0.581443),
Atom( 6 , -0.683527 , 2.941108 , -1.817704)
])

AtoBohr!(C60) # convert from Angstrom to Bohr

@time Energy,E,U = rhf(C60,verbose=true)
println("C60 Energy: $Energy")


import Base.copy

Base.@kwdef mutable struct NestedInterval{T₁ <: AbstractRange} <: AbstractModel
    start::Float64
    stop::Float64
    interval::T₁
end

Base.@kwdef struct NestedIntervalDisplayProperties{T₁ <: PlotContext, T₂ <: NTuple} <: AbstractDisplayProperties
    id::String
    caption::String
    col::ImVec4 = ImVec4(0, 0, 0, 1)
    createwindow::Bool = true
    layout::RectangularLayout
    plotcontext::T₁
    padding::T₂ = (0, 0, 0 ,0)
end

function copy(ni::NestedInterval)
    NestedInterval(get_start(ni), get_stop(ni), get_interval(ni))
end

function get_id(p::NestedIntervalDisplayProperties)
    p.id
end

function get_caption(p::NestedIntervalDisplayProperties)
    p.caption
end

function get_col(p::NestedIntervalDisplayProperties)
    p.col
end

function is_new_window(p::NestedIntervalDisplayProperties)
    p.createwindow
end

function get_layout(p::NestedIntervalDisplayProperties)
    p.layout
end

function get_padding(p::NestedIntervalDisplayProperties)
    p.padding
end

function get_plotcontext(p::NestedIntervalDisplayProperties)
   p.plotcontext
end

function get_start(ni::NestedInterval)
    ni.start
end

function set_start!(ni::NestedInterval, val::Number)
    ni.start = val > first(ni.interval) ? val : first(ni.interval)
end

function get_stop(ni::NestedInterval)
    ni.stop
end

function set_stop!(ni::NestedInterval, val::Number)
    ni.stop = val < last(ni.interval) ? val : last(ni.interval)
end

function get_interval(ni::NestedInterval)
    ni.interval
end

mutable struct NestedIntervalControl <: AbstractControl
    isenabled::Bool
end

struct NestedIntervalContext{T₁ <: AbstractControl,   T₂ <: AbstractModel,  T₃ <: AbstractDisplayProperties} <: AbstractPlotContext
    control::T₁
    model::T₂
    display_properties::T₃
end

function (context::NestedIntervalContext{<: NestedIntervalControl,   <: NestedInterval,  <: NestedIntervalDisplayProperties})()
        control = context.control
        model = context.model
        display_properties = context.display_properties
        id = get_id(display_properties)
        caption = get_caption(display_properties)
        is_new_window(display_properties) ? CImGui.Begin(caption*id,C_NULL, CImGui.ImGuiWindowFlags_NoBringToFrontOnFocus) : nothing
        isenabled(control) ? control(model, display_properties) : nothing
        is_new_window(display_properties) ? CImGui.End() : nothing
end

function (control::NestedIntervalControl)(model::NestedInterval, properties::NestedIntervalDisplayProperties)
    id = get_id(properties)
    caption = get_caption(properties)
    col = get_col(properties)
    rectangle = get_layout(properties)
    # pos = get_pos(rectangle) % TODO Remove this
    totalwidth = get_width(rectangle)
    totalheight = get_height(rectangle)
    padding = get_padding(properties)
    interval = get_interval(model)
    start = get_start(model)
    stop = get_stop(model)
    CImGui.Dummy(ImVec2(0, padding[1]))
    CImGui.Indent(padding[2])
    width = totalwidth - padding[2]
    height = totalheight
    draw_list = CImGui.GetWindowDrawList()
    pos = CImGui.GetCursorScreenPos()
    CImGui.InvisibleButton("###Invisible Nested Interval Button", ImVec2(width, height))
    x₀ = stretch_linearly(start, first(interval), last(interval), pos.x, pos.x + width)
    x₁ = stretch_linearly(stop, first(interval), last(interval), pos.x, pos.x + width)
    io = CImGui.GetIO()
    if CImGui.IsItemHovered()
       mousepos = io.MousePos
       # Relax threshold when dragging to make things smoother
       if (abs(x₀ - mousepos.x) <=  25)
           CImGui.SetMouseCursor(CImGui.ImGuiMouseCursor_ResizeEW)
           if CImGui.IsItemActive()
               i₀ = stretch_linearly(io.MousePos.x, pos.x, pos.x + width, first(interval), last(interval))
               # Round to the nearest multiple of the interval stepsize
               s = step(interval)
               i₀′ = div(i₀ + (s / 2), s) * s
               set_start!(model, i₀′)
           end
       elseif (abs(x₁ - mousepos.x) <= 25)
           CImGui.SetMouseCursor(CImGui.ImGuiMouseCursor_ResizeEW)
           if CImGui.IsItemActive()
               i₁ = stretch_linearly(io.MousePos.x, pos.x, pos.x + width, first(interval), last(interval))
               # Round to the nearest multiple of the interval stepsize
               s = step(interval)
               i₁′ = div(i₁ + (s / 2), s) * s
               set_stop!(model,  i₁′)
           end
       end
       # Make sure that start is always less than or equal to stop
       if get_start(model) > get_stop(model)
           x₀′ = get_start(model)
           set_start!(model, get_stop(model))
           set_stop!(model, x₀′)
       end
      #CImGui.Text(@sprintf("Application average %.3f ms/frame (%.1f FPS)", 1000 / CImGui.GetIO().Framerate, CImGui.GetIO().Framerate))
      #CImGui.Text(@sprintf("Start: %.3f  / Stop: %.3f", get_start(model), get_stop(model)))
    end
    CImGui.Unindent(padding[2])
    CImGui.SetCursorScreenPos(pos.x, pos.y)
    #CImGui.AddRectFilled(draw_list, ImVec2(pos.x, pos.y), ImVec2(pos.x + width,  pos.y + height), Base.convert(ImU32, col));
    # The plotcontext allows us to plot data that appears underneath the nested interval control.
    plotcontext = get_plotcontext(properties)
    plotcontext()
    CImGui.AddRectFilled(draw_list, ImVec2(x₀, pos.y + 5), ImVec2(x₁,  pos.y + 5 + height), Base.convert(ImU32, ImVec4(0.0, 0, 0.99, 0.1)));
end


mutable struct HybridTauLeapSimulator{M1,M2} <: AbstractSimulator
  exact::M1
  tauleap::M2
  is_critical::Bool
  t::Float64
end

##### constructors #####
function HybridTauLeapSimulator(exact::M1, tauleap::M2) where {M1,M2}
  return HybridTauLeapSimulator(exact, tauleap, false, 0.0)
end

##### main simulation routines #####

# TODO: this is an ugly hack
@inline cumulative_intensity(simulator::HybridTauLeapSimulator) = first(simulator.exact.algorithm.total_rate)

@inline function initialize!(simulator::HybridTauLeapSimulator, state, model, tfinal)
  # reset the current system time
  simulator.t = 0.0

  # initialize each simulator
  initialize!(simulator.exact, state, model, tfinal)
  initialize!(simulator.tauleap, state, model, tfinal)

  # check for critical reaction channels
  simulator.is_critical = check_critical(state, model)

  # seed the simulator with events
  # TODO: right now, this seeds both a leap update and jump update
  generate_next_step!(simulator)

  return nothing
end

# carry out the next leap
@inline function step!(simulator::HybridTauLeapSimulator, state, model)
  # unpack information
  exact = simulator.exact
  tauleap = simulator.tauleap

  # TODO: right now, each simulator will update by default
  # we should find a way to trigger the correct update
  # based on whether any reaction channel is critical
  if simulator.is_critical
    # perform an update using an exact method
    no_update_step!(exact, state, model)

    simulator.t = exact.t
    tauleap.t   = exact.t
  else
    # otherwise it is safe to leap
    no_update_step!(tauleap, state, model)

    simulator.t = tauleap.t
    exact.t     = tauleap.t
  end

  # check if simulation is critical
  simulator.is_critical = check_critical(state, model)

  # update the correct simulator
  update!(simulator, state, model)

  # figure out what type of step should be taken next
  generate_next_step!(simulator)

  return nothing
end

@inline function update!(simulator::HybridTauLeapSimulator, state, model)
  exact = simulator.exact
  tauleap = simulator.tauleap

  if simulator.is_critical
    update!(exact.algorithm, state, model, 0)
    tauleap.algorithm.total_rate = cumulative_intensity(exact)
  else
    update!(tauleap.algorithm, state, model)
    exact.algorithm.total_rate = cumulative_intensity(tauleap)
  end

  return nothing
end

@inline function generate_next_step!(simulator::HybridTauLeapSimulator)
  if simulator.is_critical
    generate_next_jump!(simulator.exact)
  else
    generate_next_leap!(simulator.tauleap)
  end

  return nothing
end

@inline function get_new_time(simulator::HybridTauLeapSimulator)
  if simulator.is_critical
    t = get_new_time(simulator.exact)
  else
    t = get_new_time(simulator.tauleap)
  end
  return t
end

# TODO
@inline function check_critical(state, model)
  return any(x -> x < 10, state)
end


"""
    KLocalOperatorSum

A KLocalOperator representing the sum of several KLocalOperator-s. Internally,
the sum is stored as a vector of local operators acting on some sites.
"""
struct KLocalOperatorSum{H<:AbstractHilbert, T, N, VS<:AbstractVector, Op<:AbsLinearOperator{T,N}} <: AbsLinearOperator{T,N}
    hilb::H

    # list of sites in this sum
    sites::VS

    # List of operators
    operators::Vector{Op}
end

KLocalOperatorSum(op::KLocalOperator) = KLocalOperatorSum(basis(op), [sites(op)], [op])
KLocalOperatorSum(op::AbsLinearOperator) = KLocalOperatorSum(basis(op), [sites(op)], [op])
KLocalOperatorSum(op::KLocalOperatorSum) = duplicate(op)

## Accessors
QuantumOpticsBase.basis(op::KLocalOperatorSum) = op.hilb
operators(op::KLocalOperatorSum) = op.operators
conn_type(op::KLocalOperatorSum) = conn_type(eltype(operators(op)))
conn_type(op::Type{KLocalOperatorSum{H,T,N,Vs,Op}}) where {H,T,N,Vs,Op} = conn_type(Op)
sites(op::KLocalOperatorSum) = op.sites

# Copy
function duplicate(op::KLocalOperatorSum)
    new_sites = deepcopy(op.sites)
    ops = eltype(op.operators)[];
    for _op=operators(op)
        push!(ops, duplicate(_op))
    end

    return KLocalOperatorSum(basis(op), new_sites, ops)
end

#
function _row_valdiff!(conn::AbsOpConnection, op::KLocalOperatorSum, v)
    #@boundscheck checkbounds(v, nsites(basis(op)))

    for _op=operators(op)
        @inbounds _row_valdiff!(conn, _op, v)
    end
    return conn
end

function row_valdiff_index!(conn::OpConnectionIndex, op::KLocalOperatorSum, v)
    for _op=operators(op)
        row_valdiff_index!(conn, _op, v)
    end
    return conn
end

function map_connections(fun::Function, ∑Ô::KLocalOperatorSum, v)
    #@boundscheck checkbounds(v, nsites(basis(∑Ô)))

    for Ô=operators(∑Ô)
        @inbounds map_connections(fun, Ô, v)
    end
    return nothing
end

function accumulate_connections!(acc::AbstractAccumulator, ∑Ô::KLocalOperatorSum, v)
    #@boundscheck checkbounds(v, nsites(basis(∑Ô)))

    for Ô=operators(∑Ô)
        @inbounds accumulate_connections!(acc, Ô, v)
    end
    return acc
end

function _add!(op_sum::KLocalOperatorSum, op::AbsLinearOperator)
    id = findfirst(isequal(sites(op)), op_sum.sites)

    if isnothing(id)
        push!(op_sum.sites, sites(op))
        push!(op_sum.operators, op)
    else
        _add_samesite!(op_sum.operators[id], op)
    end

    return op_sum
end

function _add!(op_l::KLocalOperatorSum, op_r::KLocalOperatorSum)
    for op=operators(op_r)
        _add!(op_l, op)
    end
    op_l
end

Base.:+(op::AbsLinearOperator, ops::KLocalOperatorSum) = ops + op
Base.:-(op::AbsLinearOperator, ops::KLocalOperatorSum) = ops - op

Base.:+(op_l::KLocalOperatorSum, op::AbsLinearOperator) = _add!(duplicate(op_l), op)
Base.:-(op_l::KLocalOperatorSum, op::AbsLinearOperator) = _add!(duplicate(op_l), -op)

Base.:+(op_l::KLocalOperatorSum, op::KLocalOperatorSum) = _add!(duplicate(op_l), op)
Base.:-(op_l::KLocalOperatorSum, op::KLocalOperatorSum) = _add!(duplicate(op_l), -op)

Base.:+(op_l::KLocalOperator, op_r::KLocalOperator) = begin
    sites(op_l) == sites(op_r) && return _add_samesite(op_l, op_r)
    return KLocalOperatorSum(op_l) + op_r
end

Base.:-(op_l::KLocalOperator, op_r::KLocalOperator) = begin
    sites(op_l) == sites(op_r) && return _add_samesite(op_l, -op_r)
    return KLocalOperatorSum(op_l) - op_r
end

function Base.transpose(ops::KLocalOperatorSum)
    new_sites  = copy(ops.sites)

    new_ops    = eltype(operators(ops))[]
    for op = operators(ops)
        push!(new_ops, transpose(op))
    end
    return KLocalOperatorSum(basis(ops), new_sites, new_ops)
end

function Base.conj(ops::KLocalOperatorSum)
    new_sites  = similar(ops.sites)
    new_sites .= ops.sites

    new_ops    = eltype(operators(ops))[]
    for op = operators(ops)
        push!(new_ops, conj(op))
    end
    return KLocalOperatorSum(basis(ops), new_sites, new_ops)
end

function Base.conj!(ops::KLocalOperatorSum)
    for op = operators(ops)
        conj!(op)
    end
    return ops
end

Base.adjoint(ops::KLocalOperatorSum) = conj!(transpose(ops))

function Base.:*(opl::KLocalOperatorSum, opr::KLocalOperator)
    ∑op =  duplicate(opl)
    for (i,op)=enumerate(operators(∑op))
        op_new = op*opr
        ∑op.operators[i] = op_new
        ∑op.sites[i] = sites(op_new)
    end
    return ∑op
end

function Base.:*(opl::KLocalOperator, opr::KLocalOperatorSum)
    ∑op =  duplicate(opr)
    for (i,op)=enumerate(operators(∑op))
        op_new = opl*op
        ∑op.operators[i] = op_new
        ∑op.sites[i] = sites(op_new)
    end
    return ∑op
end

Base.show(io::IO, ::MIME"text/plain", op::KLocalOperatorSum) = print(io,
    "KLocalOperatorSum: \n\t -sites: $(op.sites)")

Base.show(io::IO, op::KLocalOperatorSum) = print(io,
    "KLocalOperatorSum: on sites: $(op.sites)")

Base.eltype(::T) where {T<:KLocalOperatorSum} = eltype(T)
Base.eltype(::Type{KLocalOperatorSum{H,T,N,Vec,Op}}) where {H,T,N,Vec,Op} =
    eltype(Op)

Base.:*(a::Number, b::KLocalOperatorSum) =
    _op_alpha_prod(b,a)

function _op_alpha_prod(ops::KLocalOperatorSum, a::Number)
    op_all = [a*op for op=operators(ops)]
    return sum(op_all)
end

function Base.:(==)(l::KLocalOperatorSum, r::KLocalOperatorSum)
    p_l = sortperm(sites(l))
    p_r = sortperm(sites(r))

    sites(l)[p_l] == sites(r)[p_r] || return false

    for (op_l,op_r)=zip(operators(l)[p_l], operators(r)[p_r])
        op_l == op_r || return false
    end
    return true
end

function Base.isapprox(l::KLocalOperatorSum, r::KLocalOperatorSum; kwargs...)
    p_l = sortperm(sites(l))
    p_r = sortperm(sites(r))

    sites(l)[p_l] == sites(r)[p_r] || return false

    for (op_l,op_r)=zip(operators(l)[p_l], operators(r)[p_r])
        isapprox(op_l, op_r; kwargs...) || return false
    end
    return true
end


#!/usr/bin/env julia

using RobotOS
@rosimport geometry_msgs.msg: Point, Pose, Pose2D, PoseStamped, Vector3
@rosimport std_srvs.srv: Empty, SetBool
@rosimport nav_msgs.srv.GetPlan
@rosimport mavs_control.msg: Control

@rosimport gazebo_msgs.msg: ModelState
@rosimport gazebo_msgs.srv: SetModelState, GetModelState, GetWorldProperties

rostypegen()
using geometry_msgs.msg
using std_srvs.srv
using nav_msgs.srv.GetPlan
using gazebo_msgs.msg
using gazebo_msgs.srv
using mavs_control.msg
using PyCall
@pyimport tf.transformations as tf
using VehicleModels # for Linear_Spline()

# TODO
# 1) check to see if model is paused
# 2) make a while loop that publishes the current message until there is a new one ready or the current one ran out -> in this case stop the vehicle
# 3) interpolate message in time, is there a while! new message

function callback(ctr_msg::Control, get_state::ServiceProxy{GetModelState}, set_state::ServiceProxy{SetModelState})
    # interpolate optimal controls
    sp_X=Linear_Spline(ctr_msg.t,ctr_msg.x);
    sp_Y=Linear_Spline(ctr_msg.t,ctr_msg.y);
    sp_PSI=Linear_Spline(ctr_msg.t,ctr_msg.psi);

    modelName = RobotOS.get_param("robotName")
    println(" service...")

    t=to_sec(get_rostime())
    t0=to_sec(get_rostime())
    while((t-t0) < 0.5)
        println(t-t0)
        # Get the current position of the Gazebo model
        gs = GetModelStateRequest()
        gs.model_name = modelName
        gs_r = get_state(gs)

        if !gs_r.success
            error(string(" calling /gazebo/get_model_state service: ", gs_r.status_message))
        end

        # Define the current time
        t = to_sec(get_rostime())

        # Define position to move robot
        vehPose = gs_r.pose  # use the current position
        vehPose.position.x = sp_X[t]
        vehPose.position.y = sp_Y[t]
        roll = 0; pitch = 0; yaw = sp_PSI[t];
        Q = tf.quaternion_from_euler(roll, pitch, yaw)
        vehPose.orientation.x = Q[1]
        vehPose.orientation.y = Q[2]
        vehPose.orientation.z = Q[3]
        vehPose.orientation.w = Q[4]

        # Define the robot state
        ms = ModelState()
        ms.model_name = modelName
        ms.pose = vehPose

        # Set the state of the Gazebo model
        ss = SetModelStateRequest()
        ss.model_state = ms
    #    println("Calling 'gazebo/set_model_state' service...")
        ss_r = set_state(ss)

        if !ss_r.success
            error(string(" calling /gazebo/set_model_state service: ", ss_r.status_message))
        end
    end
end

function main()
    init_node("rosjl_move_hmmwv")

    # Set up service to get Gazebo model state
    const get_state = ServiceProxy("/gazebo/get_model_state", GetModelState)
    println("Waiting for 'gazebo/get_model_state' service...")
    wait_for_service("gazebo/get_model_state")

    # Set up service to set Gazebo model state
    const set_state = ServiceProxy("/gazebo/set_model_state", SetModelState)
    println("Waiting for 'gazebo/set_model_state' service...")
    wait_for_service("gazebo/set_model_state")

    sub_control = Subscriber{Control}("/mavs/optimal_control", callback, (get_state, set_state), queue_size = 2)

    spin()
end

if ! isinteractive()
    main()
end


import Base: ==, convert, vec

import JuLIP
using JuLIP: Atoms, JVec, JMat, AbstractCalculator, mat, vecs,
             numbers, positions, cell, pbc

using JuLIP.MLIPs: IPBasis

import JuLIP: read_dict, write_dict

export LsqDB, Dat, basis, configs, config

"""
`Dat`: store one atomistic configuration and the associated
observations (e.g. E, F, V, ...)
"""
mutable struct Dat
   at::Atoms                         # configuration
   configtype::String                # group identifier
   D::Dict{String, Vector{Float64}}  # list of observations
   rows::Dict{String, Vector{Int}}   # row indices for LSQ system
   info::Dict{String, Any}           # anything else...
end

# should D.info also be compared?
==(d1::Dat, d2::Dat) = (
      (d1.configtype == d2.configtype) && (d1.D == d2.D) &&
      (d1.info == d2.info) &&
      all( f(d1.at) == f(d2.at) for f in (positions, numbers, cell, pbc) )
   )

function Dat(at::Atoms, config_type::AbstractString; kwargs...)
   dat = Dat(at, config_type, Dict{String, Vector{Float64}}(),
             Dict{String, Vector{Int}}(), Dict{String, Any}())
   for (key, val) in kwargs
      str_key = string(key)
      if !(ismissing(val) || (val == nothing))
         # convert the observation to a basic vector and store it
         dat.D[str_key] = vec_obs(str_key, val)
      end
   end
   return dat
end

write_dict(d::Dat) =
   Dict("__id__" => "IPFitting.Dat",
         "at" => write_dict(d.at),
         "configtype" => d.configtype,
         "D" => d.D,
         "rows" => d.rows,
         "info" => d.info)

function Dat(D::Dict)
   at = Atoms(D["at"])
   return Dat(at,
              D["configtype"],
              Dict{String, Any}(D["D"]),
              Dict{String, Any}(D["rows"]),
              Dict{String, Any}(D["info"]))
end

read_dict(::Val{Symbol("IPFitting.Dat")}, D::Dict) = Dat(D)

observation(d::Dat, key::String) = d.D[key]
hasobservation(d::Dat, key::String) = haskey(d.D, key)
observation(key::String, d::Dat) = d.D[key]

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

"""
`mutable struct LsqDB{TD, TB}`

A representation of a least-squares system stored on disk, which can
be extended by adding data, or basis functions
"""

mutable struct LsqDB
   basis::IPBasis
   configs::Vector{Dat}
   Ψ::Matrix{Float64}
   dbpath::String
end

basis(db::LsqDB) = db.basis
basis(db::LsqDB, i::Integer) = db.basis[i]
configs(db::LsqDB) = db.configs
config(db::LsqDB, i::Integer) = db.configs[i]

# prototypes

"""
convert some real data, in some generic format, into a vector to be stored
in a `Dat` or Lsq system. E.g.,
```
F = forces(...)::Vector{JVecF}
vec_obs(::Val{:F}, F) = mat(F)[:]
```
or equivalently, `vec_obs("F", F)`
"""
function vec_obs end

"""
convert a Vector{T} to some real (atomistic) data, e.g.,
```
x::Vector{Float64}
devec_obs(::Val{:F}, x) = vecs( resize(x, 3, length(x) ÷ 3) )
```
"""
function devec_obs end

"""
evaluate a specific observation type
"""
function eval_obs end

"""
create special weights for different observations
"""
weighthook(::Val, d::Dat) = 1.0
weighthook(s::String, d::Dat) = weighthook(Val(Symbol(s)), d)

"""
create special weights for different observations, similar to
`weighthooks` but this is used for computing errors rather than
for fitting
"""
err_weighthook(::Val, d::Dat) = 1.0
err_weighthook(s::String, d::Dat) = err_weighthook(Val(Symbol(s)), d)


module TopOpt

const PENALTY_BEFORE_INTERPOLATION = true
using Requires, Reexport

macro cuda_only(mod, code)
    return esc(
        quote
            @init @require CUDASupport = "97986420-7ec3-11e9-24cd-4f0e301eb539" @eval $mod begin
                $code
            end
        end,
    )
end

abstract type AbstractDevice end
struct CPU <: AbstractDevice end
struct GPU <: AbstractDevice end
whichdevice(::Any) = CPU()

# GPU utilities
module GPUUtils end

@reexport using Nonconvex, NonconvexMMA, NonconvexSemidefinite, NonconvexPercival

# Utilities
include(joinpath("Utilities", "Utilities.jl"))
using .Utilities

# Topopology optimization problem definitions
include(joinpath("TopOptProblems", "TopOptProblems.jl"))

using LinearAlgebra, Statistics
using Reexport, Parameters, Setfield
@reexport using .TopOptProblems

# Truss Topopology optimization problem definitions
include(joinpath("TrussTopOptProblems", "TrussTopOptProblems.jl"))
@reexport using .TrussTopOptProblems

using Ferrite, StaticArrays

using ForwardDiff, IterativeSolvers#, Preconditioners
@reexport using VTKDataTypes

const DEBUG = Base.RefValue(false)

# FEA solvers
include(joinpath("FEA", "FEA.jl"))
using .FEA

# Chequeurboard filter
include(joinpath("CheqFilters", "CheqFilters.jl"))
using .CheqFilters

# Objective and constraint functions
@reexport using Flux
include(joinpath("Functions", "Functions.jl"))
@reexport using .Functions

# Various topology optimization algorithms
include(joinpath("Algorithms", "Algorithms.jl"))
using .Algorithms

macro init_cuda()
    return esc(
        quote
            const CuArrays = CUDASupport.CuArrays
            const CUDAdrv = CUDASupport.CUDAdrv
            const CUDAnative = CUDASupport.CUDAnative
            const GPUArrays = CUDASupport.GPUArrays
            CuArrays.allowscalar(false)
            const dev = CUDAdrv.device()
            const ctx = CUDAdrv.CuContext(dev)
            using .CuArrays, .CUDAnative
            using .GPUArrays: GPUVector, GPUArray
        end,
    )
end

@cuda_only GPUUtils include("GPUUtils/GPUUtils.jl")
@cuda_only Utilities include("Utilities/gpu_utilities.jl")
@cuda_only TopOptProblems include("TopOptProblems/gpu_support.jl")
@cuda_only FEA include("FEA/gpu_solvers.jl")
@cuda_only CheqFilters include("CheqFilters/gpu_cheqfilter.jl")
@cuda_only Functions include("Functions/gpu_support.jl")
@cuda_only Algorithms include("Algorithms/SIMP/gpu_simp.jl")

export TopOpt,
    simulate,
    TopOptTrace,
    DirectDisplacementSolver,
    PCGDisplacementSolver,
    StaticMatrixFreeDisplacementSolver,
    SensFilter,
    DensityFilter,
    Displacement,
    Compliance,
    CG,
    Direct,
    Assembly,
    MatrixFree,
    FEASolver,
    Optimizer,
    SIMP,
    ContinuationSIMP,
    BESO,
    GESO,
    PowerContinuation,
    ExponentialContinuation,
    LogarithmicContinuation,
    CubicSplineContinuation,
    SigmoidContinuation,
    Continuation,
    save_mesh,
    CPU,
    GPU,
    DefaultCriteria,
    EnergyCriteria,
    PowerPenalty,
    RationalPenalty,
    SinhPenalty,
    MMA87,
    MMA02,
    HeavisideProjection,
    ProjectedPenalty,
    PowerPenalty
end


export NeuralNetworkQ, update!

using Flux
using CUDAapi

import Flux:params
import Zygote:gradient
import .Utils:to_device, to_host

# TODO: Add static size check for legal input data type.

"""
    NeuralNetworkQ(;kwargs...) -> NeuralNetworkQ{D, M, O, P}

Use neural networks to generate estimations of state-action values.

# Keywords

- `model::M`: describes the network structure.
- `optimizer::O`: defines how to update parameters given grads.
- `parameters::P=params(model)`: the parameters of `model`.
- `device::Symbol=:cpu`: the param `D` of `NeuralNetworkQ`, specify where to run the model. Supported keywords are:
    - `:cpu`
    - `:gpu`, if the specified device is `:gpu`, then we will automatically change the device to one of the following device types according to the `backend(model)`:
        - `:Knet_gpu`, means the model is `Knet` that runs on gpu.
        - `:Zygote_gpu`, means the model is `Flux` with `Zygote` that runs on gpu.

# Fields

- `model::M`
- `optimizer::O`
- `params::P`

"""
struct NeuralNetworkQ{D, M, O, P} <: AbstractQApproximator
    model::M
    optimizer::O
    params::P
end

function NeuralNetworkQ(;model::M, optimizer::O, parameters::P=params(model), device::Symbol=:cpu) where {M, O, P}
    if device == :gpu
        has_cuda() || throw(ArgumentError("the specified device is gpu but can not find cuda!"))
        device = Symbol(backend(model), "_gpu")
    end
    m = to_device(Val(device), model)
    ps = params(m)
    NeuralNetworkQ{device, typeof(m), O, typeof(ps)}(m, optimizer, ps)
end

params(Q::NeuralNetworkQ) = Q.params

to_device(Q::NeuralNetworkQ{D}, x) where D = to_device(Val(D), x)

"get Q value of the batch"
batch_estimate(Q::NeuralNetworkQ, states) = Q.model(states)

"get Q value of all actions"
(Q::NeuralNetworkQ)(s) = Q.model(to_device(Q, reshape(s, size(s)..., 1))) |> to_host |> drop_last_dim

"get Q value of some specific action"
(Q::NeuralNetworkQ)(s, a) = Q(s)[a]

function update!(Q::NeuralNetworkQ, gs)
    Flux.Optimise.update!(Q.optimizer, Q.params, gs)
end

Base.copyto!(dest::NeuralNetworkQ, src::NeuralNetworkQ) = Flux.loadparams!(dest.model, src.params)

gradient(f, Q::NeuralNetworkQ) = gradient(f, Val(backend(Q.model)), params(Q))


include("Technology.jl")

mutable struct SystemState <: ScenarioNode
    year::Int64
    capacityMix::Array{Technology,1}
    installedCapacity::Float64
    installedLimitByTechnologyType::Dict{Symbol,Real} #MW
    installedByTechnologyType::Dict{Symbol,Real} #MW
    emissionPrice::Real

    CalculateInstalledCapacity::Function
    TimeSeries::Function

    function SystemState(year,capacityMix)
        self = new()
        self.year = year
        self.capacityMix = capacityMix
        self.installedCapacity = 0

        self.TimeSeries = function(day::Int64, hour::Int64, systemLoad::Load, yearlyTechnologyMix::Array{Technology,1})
            systemLoad.TimeSeries(day,hour)
            for technology in yearlyTechnologyMix
                if typeof(technology) == RenewablePlants
                    technology.TimeSeries(day,hour)
                end
            end
        end

        self.CalculateInstalledCapacity = function()
            for technology in self.capacityMix
                self.installedCapacity += technology.designProperties.installedCapacity
            end
        end

        return self
    end
end

mutable struct SystemHistory <: Buffer
    systemCapacityMix::Array{Array{Technology,1},1}
    systemElectricityPrice::Array{Float64,1}
    systemDecisionSequence::Array{Array{Array{Float64,1},1},1}
    capacityMixProjection::Array{Array,1}
    InvestmentHistory::Array{Symbol,1}

    function SystemHistory()
        self = new()
        self.systemCapacityMix = Array{Array{Technology,1},1}(undef,0)
        self.systemElectricityPrice = Array{Float64,1}(undef,0)
        self.systemDecisionSequence = Array{Array{Array{Float64,1},1},1}(undef,0)
        self.capacityMixProjection = Array{Array,1}(undef,0)
        self.InvestmentHistory = Array{Symbol,1}(undef,0)

        return self
    end

end


struct HydroDispatchRunOfRiverOutage <: PSI.AbstractHydroDispatchFormulation end
struct HydroDispatchReservoirStorageOutage <: PSI.AbstractHydroReservoirFormulation end
struct HydroDispatchPumpedStoragewReservationOutage <: PSI.AbstractHydroReservoirFormulation end

function outage_constraints!(
    optimization_container::PSI.OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    model::PSI.DeviceModel{T, D},
    ::Type{S},
    feedforward::Union{Nothing, PSI.AbstractAffectFeedForward},
) where {T <: PSY.HydroGen, S <: PM.AbstractPowerModel, D <: PSI.AbstractHydroFormulation}
    parameters = PSI.model_has_parameters(optimization_container)
    resolution = PSI.model_resolution(optimization_container)

    forecast_label = "outage"
    constraint_infos = Vector{DeviceOutageConstraintInfo}()
    for d in devices
        name = PSY.get_name(d)
        info = DeviceOutageConstraintInfo(
            name,
            nothing,
            1.0,
            PSI.get_time_series(optimization_container, d, forecast_label),
        )
        push!(constraint_infos, info)
    end

    if !(isempty(devices))
        if parameters
            device_outage_ub_parameter!(
                optimization_container,
                constraint_infos,
                PSI.make_constraint_name(OUTAGE, T),
                PSI.make_variable_name(PSI.ActivePowerVariable, T),
                PSI.UpdateRef{T}(OUTAGE, forecast_label),
            )
        else
            device_outage_ub!(
                optimization_container,
                constraint_infos,
                PSI.make_constraint_name(OUTAGE, T),
                PSI.make_variable_name(PSI.ActivePowerVariable, T),
            )
        end
    else
        @warn "Data doesn't contain generators of type $T, consider adjusting your formulation"
    end

    return
end

function outage_constraints!(
    optimization_container::PSI.OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    model::PSI.DeviceModel{T, D},
    ::Type{S},
    feedforward::Union{Nothing, PSI.AbstractAffectFeedForward},
) where {
    T <: PSY.HydroPumpedStorage,
    S <: PM.AbstractPowerModel,
    D <: PSI.AbstractHydroFormulation,
}
    parameters = PSI.model_has_parameters(optimization_container)
    resolution = PSI.model_resolution(optimization_container)

    forecast_label = "outage"
    constraint_infos = Vector{DeviceOutageConstraintInfo}()
    for d in devices
        name = PSY.get_name(d)
        info = DeviceOutageConstraintInfo(
            name,
            nothing,
            1.0,
            PSI.get_time_series(optimization_container, d, forecast_label),
        )
        push!(constraint_infos, info)
    end

    if !(isempty(devices))
        if parameters
            device_outage_ub_parameter!(
                optimization_container,
                constraint_infos,
                PSI.make_constraint_name(OUTAGE, T),
                PSI.make_variable_name(PSI.ActivePowerInVariable, T),
                PSI.UpdateRef{T}(OUTAGE, forecast_label),
            )
            device_outage_ub_parameter!(
                optimization_container,
                constraint_infos,
                PSI.make_constraint_name(OUTAGE, T),
                PSI.make_variable_name(PSI.ActivePowerOutVariable, T),
                PSI.UpdateRef{T}(OUTAGE, forecast_label),
            )
        else
            device_outage_ub!(
                optimization_container,
                constraint_infos,
                PSI.make_constraint_name(OUTAGE, T),
                PSI.make_variable_name(PSI.ActivePowerInVariable, T),
            )
            device_outage_ub!(
                optimization_container,
                constraint_infos,
                PSI.make_constraint_name(OUTAGE, T),
                PSI.make_variable_name(PSI.ActivePowerOutVariable, T),
            )
        end
    else
        @warn "Data doesn't contain generators of type $T, consider adjusting your formulation"
    end

    return
end


using MPIReco

@testset "system matrix center estimation" begin
  #dirs = ["8.mdf", "9.mdf", "10.mdf", "11.mdf"]   8.mdf is wrong. BG is subtracted twice
  dirs = ["9.mdf", "10.mdf", "11.mdf"]
  SFs = joinpath.(datadir, "calibrations", dirs)
  #center = [[9.0,23.0,1.0],[9.0,10.0,1.0],[22.0,23.0,1.0],[22.0,10.0,1.0]]
  center = [[9.0,10.0,1.0],[22.0,23.0,1.0],[22.0,10.0,1.0]]

  for (l,SF) in enumerate(SFs)
    bSF = MPIFile(SF)
    a = findCenterOfDfFov(bSF)
    @test a == center[l]
    S = getSF(bSF,2,0,0,2)
    exportImage(joinpath(imgdir, "Center1.png"), abs.(S)[:,:,1])
    #p = imagesc(abs.(S)[:,:,1])
    #add(p, Points([a[2]-1],[a[1]-1],color="yellow",lw=5))
    #savefig(p, "$imgdir/Center1.png")
  end

  for (l,SF) in enumerate(SFs)
    bSF = MPIFile(SF)
    a = findCenterOfDfFov(bSF)
    @test a == center[l]
    S = getSF(bSF,5,6,0,1)
    exportImage(joinpath(imgdir, "Center2.png"), abs.(S)[:,:,1])
    #p = imagesc(abs.(S)[:,:,1])
    #add(p, Points([a[2]-1],[a[1]-1],color="yellow",lw=5))
    #savefig(p, "$imgdir/Center2.png")
  end
end


struct Wind
    id::Int
    pwm::UInt8
end
Wind(id::Int, ::Missing) = Wind(id, 0x00)
Wind(id::Int, pwm::String) = Wind(id, parse(Int, pwm))

tosecond(t::T) where {T <: TimePeriod}= t/convert(T, Second(1))
sincestart(t) = tosecond(t - t₀[])

function toint(msg)
    y = zero(UInt32)
    for c in msg
        y <<= 8
        y += c
    end
    return y
end

function getrpm(t)
    t4 = 15000000
    shortest_t = t4/1.1top_rpm[]
    t < shortest_t ?  missing : t4/t
end

function set_pwm!(sp, c, pwm)
    lock(c) do
        sp_flush(sp, SP_BUF_INPUT)
        encode(sp, pwm)
    end
end

function update_rpm!(sp, c, pwm, msg, rpm)
    if pwm[] > 19
        lock(c) do
            sp_flush(sp, SP_BUF_OUTPUT)
            decode!(msg, sp, 12)
        end
        for (i, x) in enumerate(Iterators.partition(msg, 4))
            rpm[i] = getrpm(toint(x))
        end
    else
        fill!(rpm, 0.0)
    end
end

function isconnected(port)
    try
        sp = LibSerialPort.open(port, baudrate[])
        close(sp)
        true
    catch ex
        false
    end
end

mutable struct FanArduino <: AbstractArduino
    id::Int
    c::ReentrantLock
    port::String
    sp::SerialPort
    msg::Vector{UInt8}
    rpm::Vector{Union{Missing, Float64}}
    pwm::Observable{UInt8}
    function FanArduino(id::Int, port::String)
        c = ReentrantLock()
        sp = LibSerialPort.open(port, baudrate[])
        pwm = Observable(0x00)
        on(pwm) do x
            set_pwm!(sp, c, x)
        end
        msg = Vector{UInt8}(undef, 12)
        rpm = Vector{Float64}(undef, 3)
        new(id, c, port, sp, msg, rpm, pwm)
    end
end

update_rpm!(a::FanArduino) = update_rpm!(a.sp, a.c, a.pwm, a.msg, a.rpm)

get_rpm(a::FanArduino) = a.rpm

mutable struct AllWind
    arduinos::Vector{FanArduino}
    io::IOStream
    framerate::Int
    function AllWind(arduinos::Vector{FanArduino}, framerate::Int)
        io = open(tempname(), "w")
        close(io)
        new(arduinos, io, framerate)
    end
end

function get_rpms(allwind::AllWind)
    @sync for a in allwind.arduinos
        @async update_rpm!(a)
    end
    now() => get_rpm.(allwind.arduinos)
end

function record(allwind::AllWind, folder)
    !isdir(folder) && mkpath(folder)
    isopen(allwind.io) && close(allwind.io)
    allwind.io = open(folder / "fans.csv", "w")
    println(allwind.io, "time,", join([join(["fan$(a.id)_speed$j" for j in 1:3], ",") for a in allwind.arduinos], ","))
    @async while isopen(allwind.io)
        println(allwind.io, now(), ",",join(Iterators.flatten(get_rpm.(allwind.arduinos)), ","))
        sleep(1/allwind.framerate)
    end
end

Base.isopen(allwind::AllWind) = all(isopen, allwind.arduinos)
Base.close(allwind::AllWind) = close.(allwind.arduinos)


"""
    minimalsetup1d(N)

Creates a minimal setup of all important variables.

Such it generates the moments `height`, `velocity` as well as a dummy distribution with its speeds `f0`, `f1` and `f2`.
Further a set of forcing is supplied, they can be used with `foces.slip`, `.h∇p`, `.bathymetry` and `.thermal`.

# Example
```jldoctest
julia> using JuSwalbe

julia> constants, mom, f, distribution = minimalsetup1d(10)
(JuSwalbe.Inputconstants
  lx: Int64 10
  ly: Int64 512
  maxruntime: Int64 100000
  dumping: Int64 1000
  τ: Float64 1.0
  gravity: Float64 0.0
  γ: Float64 0.01
  δ: Float64 1.0
  μ: Float64 0.16666666666666666
  kbt: Float64 0.0
, JuSwalbe.Macroquant{Array{Float64,1},Array{Float64,1}}
  height: Array{Float64}((10,)) [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
  velocity: Array{Float64}((10,)) [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]
  pressure: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  energy: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
, JuSwalbe.Forces{Array{Float64,1}}
  slip: Array{Float64}((10,)) [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]
  h∇p: Array{Float64}((10,)) [-0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1]
  bathymetry: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  thermal: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
, JuSwalbe.DistributionD1Q3{Array{Float64,1}}
  f0: Array{Float64}((10,)) [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
  f1: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  f2: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
)

julia> mom.height
10-element Array{Float64,1}:
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
```
See also: [`minimalsetup2d`](@ref)
"""
function minimalsetup1d(N::Int)
    input = Inputconstants(lx=N)
    mom = Macroquant(height = ones(N), velocity=fill(0.1, N), pressure=zeros(N), energy=zeros(N))
    forces = Forces(slip=fill(0.1, N), thermal=zeros(N), h∇p=fill(-0.1, N), bathymetry=zeros(N))
    dist = DistributionD1Q3(f0 = ones(N), f1 = zeros(N), f2 = zeros(N))
    return input, mom, forces, dist
end

"""
    minimalsetup2d(N,M)

Creates a minimal working example set of macroscopic moments, forces and a distribution.

# Example
```jldoctest
julia> using JuSwalbe

julia> input, mom, f, dist = minimalsetup2d(5,5)
(JuSwalbe.Inputconstants
  lx: Int64 5
  ly: Int64 5
  maxruntime: Int64 100000
  dumping: Int64 1000
  τ: Float64 1.0
  gravity: Float64 0.0
  γ: Float64 0.01
  δ: Float64 1.0
  μ: Float64 0.16666666666666666
  kbt: Float64 0.0
, JuSwalbe.Macroquant{Array{Float64,2},JuSwalbe.Twovector{Array{Float64,2}}}
  height: Array{Float64}((5, 5)) [1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0; … ; 1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0]
  velocity: JuSwalbe.Twovector{Array{Float64,2}}
  pressure: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  energy: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
, JuSwalbe.Forces{JuSwalbe.Twovector{Array{Float64,2}}}
  slip: JuSwalbe.Twovector{Array{Float64,2}}
  h∇p: JuSwalbe.Twovector{Array{Float64,2}}
  bathymetry: JuSwalbe.Twovector{Array{Float64,2}}
  thermal: JuSwalbe.Twovector{Array{Float64,2}}
, JuSwalbe.DistributionD2Q9{Array{Float64,2}}
  f0: Array{Float64}((5, 5)) [1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0; … ; 1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0]
  f1: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f2: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f3: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f4: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f5: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f6: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f7: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f8: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
)

```

See also: [`minimalsetup1d`](@ref)
"""
function minimalsetup2d(N::Int, M::Int)
    input = Inputconstants(lx=N, ly=M)
    vel = Twovector(x=fill(0.02, (N,M)), y=fill(-0.02, (N,M)))
    zerovec = Twovector(x=zeros(N,M), y=zeros(N,M))
    mom = Macroquant(height = ones(N,M), velocity=vel, pressure=zeros(N,M), energy=zeros(N,M))
    forces = Forces(slip=vel, thermal=zerovec, h∇p=vel, bathymetry=zerovec)
    dist = DistributionD2Q9(f0 = ones(N,M), f1 = zeros(N,M), f2 = zeros(N,M),
                            f3 = zeros(N,M), f4 = zeros(N,M), f5 = zeros(N,M),
                            f6 = zeros(N,M), f7 = zeros(N,M), f8 = zeros(N,M))
    return input, mom, forces, dist
end

"""
  simplemoment2d(n,m,type)

Creates moments of type JuSwalbe.Macroquant with dimensions (n,m).

# Example
```jldoctest
julia> using JuSwalbe

julia> mom = simplemoment2d(5,5)
JuSwalbe.Macroquant{Array{Float64,2},JuSwalbe.Twovector{Array{Float64,2}}}
  height: Array{Float64}((5, 5)) [1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0; … ; 1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0]
  velocity: JuSwalbe.Twovector{Array{Float64,2}}
  pressure: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  energy: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]

julia> mom.height
5×5 Array{Float64,2}:
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0

```
"""
function simplemoment2d(n::Int, m::Int; T=Float64)
  mom = JuSwalbe.Macroquant(height=ones(T, (n,m)),
                            velocity=JuSwalbe.Twovector(x=fill(T(0.1),(n,m)), y=fill(T(-0.1),(n,m))),
                            pressure=zeros(T, (n,m)),
                            energy=zeros(T,(n,m)))
  return mom
end

"""
  simplemoment1d(n,type)

Creates moments of type JuSwalbe.Macroquant with length n.

# Example
```jldoctest
julia> using JuSwalbe

julia> mom = simplemoment1d(5)
JuSwalbe.Macroquant{Array{Float64,1},Array{Float64,1}}
  height: Array{Float64}((5,)) [1.0, 1.0, 1.0, 1.0, 1.0]
  velocity: Array{Float64}((5,)) [0.1, 0.1, 0.1, 0.1, 0.1]
  pressure: Array{Float64}((5,)) [0.0, 0.0, 0.0, 0.0, 0.0]
  energy: Array{Float64}((5,)) [0.0, 0.0, 0.0, 0.0, 0.0]

julia> mom.height
5-element Array{Float64,1}:
 1.0
 1.0
 1.0
 1.0
 1.0

```
"""
function simplemoment1d(n::Int; T=Float64)
  mom = JuSwalbe.Macroquant(height=ones(T, n), velocity=fill(T(0.1),n), pressure=zeros(T, n), energy=zeros(T, n))
  return mom
end

"""
  simpleTwovector(n,m,T)

Generates a JuSwalbe.Twovector of type T and dimension (n,m).

# Example
```jldoctest
julia> using JuSwalbe

julia> twovec = simpleTwovector(5,5)
JuSwalbe.Twovector{Array{Float64,2}}
  x: Array{Float64}((5, 5)) [0.1 0.1 … 0.1 0.1; 0.1 0.1 … 0.1 0.1; … ; 0.1 0.1 … 0.1 0.1; 0.1 0.1 … 0.1 0.1]
  y: Array{Float64}((5, 5)) [-0.1 -0.1 … -0.1 -0.1; -0.1 -0.1 … -0.1 -0.1; … ; -0.1 -0.1 … -0.1 -0.1; -0.1 -0.1 … -0.1 -0.1]

julia> twovec.x
5×5 Array{Float64,2}:
 0.1  0.1  0.1  0.1  0.1
 0.1  0.1  0.1  0.1  0.1
 0.1  0.1  0.1  0.1  0.1
 0.1  0.1  0.1  0.1  0.1
 0.1  0.1  0.1  0.1  0.1
```
"""
function simpleTwovector(n::Int,m::Int; T=Float64)
  xy = JuSwalbe.Twovector(x=fill(T(0.1),(n,m)), y=fill(T(-0.1),(n,m)))
  return xy
end

function zeroTwovector(n::Int,m::Int; T=Float64)
  xy = JuSwalbe.Twovector(x=zeros(T ,(n,m)), y=zeros(T, (n,m)))
  return xy
end

function zeroTwovectorCU(n::Int,m::Int; T=Float32)
  jaja = zeros(n,m)
  xy = JuSwalbe.Twovector(x=CuArray(jaja), y=CuArray(jaja))
  return xy
end

"""
  simpledistD2Q9(n, m, T)

Generates a `D2Q9` distribution function of dimension `n`,`m` with defined fill statements.

# Example
```jldoctest
julia> using JuSwalbe

julia> dist = simpledistD2Q9(5,5)
JuSwalbe.DistributionD2Q9{Array{Float64,2}}
  f0: Array{Float64}((5, 5)) [1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0; … ; 1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0]
  f1: Array{Float64}((5, 5)) [0.1 0.1 … 0.1 0.1; 0.1 0.1 … 0.1 0.1; … ; 0.1 0.1 … 0.1 0.1; 0.1 0.1 … 0.1 0.1]
  f2: Array{Float64}((5, 5)) [-0.1 -0.1 … -0.1 -0.1; -0.1 -0.1 … -0.1 -0.1; … ; -0.1 -0.1 … -0.1 -0.1; -0.1 -0.1 … -0.1 -0.1]
  f3: Array{Float64}((5, 5)) [0.01 0.01 … 0.01 0.01; 0.01 0.01 … 0.01 0.01; … ; 0.01 0.01 … 0.01 0.01; 0.01 0.01 … 0.01 0.01]
  f4: Array{Float64}((5, 5)) [-0.01 -0.01 … -0.01 -0.01; -0.01 -0.01 … -0.01 -0.01; … ; -0.01 -0.01 … -0.01 -0.01; -0.01 -0.01 … -0.01 -0.01]
  f5: Array{Float64}((5, 5)) [0.2 0.2 … 0.2 0.2; 0.2 0.2 … 0.2 0.2; … ; 0.2 0.2 … 0.2 0.2; 0.2 0.2 … 0.2 0.2]
  f6: Array{Float64}((5, 5)) [0.02 0.02 … 0.02 0.02; 0.02 0.02 … 0.02 0.02; … ; 0.02 0.02 … 0.02 0.02; 0.02 0.02 … 0.02 0.02]
  f7: Array{Float64}((5, 5)) [0.02 0.02 … 0.02 0.02; 0.02 0.02 … 0.02 0.02; … ; 0.02 0.02 … 0.02 0.02; 0.02 0.02 … 0.02 0.02]
  f8: Array{Float64}((5, 5)) [0.5 0.5 … 0.5 0.5; 0.5 0.5 … 0.5 0.5; … ; 0.5 0.5 … 0.5 0.5; 0.5 0.5 … 0.5 0.5]

julia> dist.f0
5×5 Array{Float64,2}:
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0

julia> typeof(dist)
JuSwalbe.DistributionD2Q9{Array{Float64,2}}
```

See also: [`simpleTwovector`](@ref), [`simplemoment2d`](@ref)

"""
function simpledistD2Q9(n::Int, m::Int; T=Float64)
  dist = JuSwalbe.DistributionD2Q9(f0 = ones(T, (n,m)), f1 = fill(T(0.1), (n,m)), f2 = fill(T(-0.1), (n,m)),
                                   f3 = fill(T(0.01), (n,m)), f4 = fill(T(-0.01), (n,m)), f5 = fill(T(0.2), (n,m)),
                                   f6 = fill(T(0.02), (n,m)), f7 = fill(T(0.02), (n,m)), f8 = fill(T(0.5), (n,m)))

  return dist
end


#src execute = true

# # [Parsing Data](@id parsing)
#
# `PowerSystems.jl` supports the creation of a `System` from a variety of common data formats:
#
# - [MATPOWER](@ref pm_data) (code copied with permission from [`PowerModels.jl`](https://github.com/lanl-ansi/PowerModels.jl))
# - [PSS/e RAW Files](@ref pm_data) (code copied with permission from [`PowerModels.jl`](https://github.com/lanl-ansi/PowerModels.jl))
# - [PSS/e DYR Files](@ref dyr_data)
# - [PowerSystems table data (CSV Files)](@ref table_data)
#
# ## [MATPOWER / PSS/e](@id pm_data)
#
# The following code will create a System from a MATPOWER or PSS/e file:
#
# ```julia
# sys = System(joinpath(data_dir, "case5.m"))
# ```
#
# ## [PSS/e dynamic data parsing](@id dyr_data)
#
# PSS/e's dynamic model library is extensive, we currently support parsing a limited amount
# of models out of the box.
#
# | Machine models | AVR Models | Prime Movers | PSS models |
# |----------------|------------|--------------|------------|
# | GENSAE         | IEEET1     | HYGOV        | IEEEST     |
# | GENSAL         | ESDC1A     | IEEEG1       |            |
# | GENROE         | ESAC1A     | GGOV1        |            |
# | GENCLS         | ESST4B     |              |            |
# | GENROU         | EXAC2      |              |            |
# |                | EXPIC1     |              |            |
# |                | ESAC6A     |              |            |
# |                | EXAC1      |              |            |
# |                | SCRX       |              |            |
# |                | ESDC2A     |              |            |
#
# ### Creating a Dynamic System using .RAW and .DYR data
#
# A `PowerSystems.jl` system can be created using a .RAW and a .DYR file. In this example we will create the following three bus system using the following RAW file:
#
# ```raw
# 0, 100, 33, 0, 0, 60  / 24-Apr-2020 19:28:39 - MATPOWER 7.0.1-dev
#
#
#      101, 'BUS 1       ',       138, 3,    1,    1, 1,           1.02,        0,  1.1,  0.9,  1.1,  0.9
#      102, 'BUS 2       ',       138, 2,    1,    1, 1,           1.0142,           0,  1.1,  0.9,  1.1,  0.9
#      103, 'BUS 3       ',       138, 2,    1,    1, 1,           1.0059,           0,  1.1,  0.9,  1.1,  0.9
# 0 / END OF BUS DATA, BEGIN LOAD DATA
#      101,  1, 1,    1,    1,       100,       20, 0, 0, 0, 0, 1, 1, 0
#      102,  1, 1,    1,    1,       70,       10, 0, 0, 0, 0, 1, 1, 0
#      103,  1, 1,    1,    1,       50,       10, 0, 0, 0, 0, 1, 1, 0
# 0 / END OF LOAD DATA, BEGIN FIXED SHUNT DATA
# 0 / END OF FIXED SHUNT DATA, BEGIN GENERATOR DATA
#      101,  1,       20,         0,       100,      -100,    1.02, 0,     100, 0, 0, 0, 0, 1, 1, 100,       318,         0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1
#      102,  1,       100,         0,       100,      -100,   1.0142, 0,     100, 0, 0.7, 0, 0, 1, 1, 100,       318,         0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1
#      103,  1,       100,         0,       100,      -100,   1.0059, 0,     100, 0, 0.2, 0, 0, 1, 1, 100,       318,         0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1
# 0 / END OF GENERATOR DATA, BEGIN BRANCH DATA
#      101,      103, 1,  0.01000,     0.12,      0.0,     250,     250,     250, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1
#      101,      102, 1,  0.01000,     0.12,      0.0,     250,     250,     250, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1
#      102,      103, 1,  0.01000,     0.12,      0.0,     250,     250,     250, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1
# 0 / END OF BRANCH DATA, BEGIN TRANSFORMER DATA
# 0 / END OF TRANSFORMER DATA, BEGIN AREA DATA
# 0 / END OF AREA DATA, BEGIN TWO-TERMINAL DC DATA
# 0 / END OF TWO-TERMINAL DC DATA, BEGIN VOLTAGE SOURCE CONVERTER DATA
# 0 / END OF VOLTAGE SOURCE CONVERTER DATA, BEGIN IMPEDANCE CORRECTION DATA
# 0 / END OF IMPEDANCE CORRECTION DATA, BEGIN MULTI-TERMINAL DC DATA
# 0 / END OF MULTI-TERMINAL DC DATA, BEGIN MULTI-SECTION LINE DATA
# 0 / END OF MULTI-SECTION LINE DATA, BEGIN ZONE DATA
# 0 / END OF ZONE DATA, BEGIN INTER-AREA TRANSFER DATA
# 0 / END OF INTER-AREA TRANSFER DATA, BEGIN OWNER DATA
# 0 / END OF OWNER DATA, BEGIN FACTS CONTROL DEVICE DATA
# 0 / END OF FACTS CONTROL DEVICE DATA, BEGIN SWITCHED SHUNT DATA
# 0 / END OF SWITCHED SHUNT DATA, BEGIN GNE DEVICE DATA
# 0 / END OF GNE DEVICE DATA, BEGIN INDUCTION MACHINE DATA
# 0 / END OF INDUCTION MACHINE DATA
# Q
# ```
#
# This system is a three bus system with three generators, three loads and three branches. The dynamic data for the generators is provided in the DYR file:
#
# ```raw
#   101 'GENROE' 1   8.000000  0.030000  0.400000  0.050000  6.500000  0.000000  1.800000
#   1.700000  0.300000  0.550000  0.250000  0.200000  0.039200  0.267200  /
#   101 'ESST1A' 1   1  1  0.01  99  -99  1  10  1  1  200  0  4  -4  4  -4  0  0  1  0  3  /
#   102 'GENCLS' 1   0.0   0.0 /
#   103 'GENCLS' 1   3.1   2.0 /
# ```
#
# That assigns a GENROU generator and a ESST1A voltage regulator at the generator located at bus 101, while classic machine models for the generators located at bus 102 and 103.
#
# To create the system we can do it passing both files directories:
#
# ```@example raw_dyr_system
# using PowerSystems
# data_dir = "../../../data" # hide
# RAW_dir = joinpath(data_dir, "PSSE_test/ThreeBusNetwork.raw")
# DYR_dir = joinpath(data_dir, "PSSE_test/TestGENCLS.dyr")
# dyn_system = System(RAW_dir, DYR_dir, runchecks = false)
# ```
#
# ## [PowerSystems Table Data](@id table_data)
#
# This is a custom format that allows users to define power system component data
# by category and column with custom names, types, and units.
#
# ### Categories
#
# Components for each category must be defined in their own CSV file. The
# following categories are currently supported:
#
# - branch.csv
# - bus.csv (required)
#   - columns specifying `area` and `zone` will create a corresponding set of `Area` and `LoadZone` objects.
#   - columns specifying `max_active_power` or `max_reactive_power` will create `PowerLoad` objects when nonzero values are encountered and will contribute to the `peak_active_power` and `peak_reactive_power` values for the
#   corresponding `LoadZone` object.
# - dc_branch.csv
# - gen.csv
# - load.csv
# - reserves.csv
# - storage.csv
#
# These must reside in the directory passed when constructing PowerSystemTableData.
#
# ### [Adding Time Series Data](@id parsing_time_series)
#
# PowerSystems requires a metadata file that maps components to their time series
# data in order to be able to automatically construct time_series from raw data
# files. The following fields are required for each time array:
#
# - `simulation`:  User description of simulation
# - `resolution`:  Resolution of time series in seconds
# - `module`:  Module that defines the abstract type of the component
# - `category`:  Type of component. Must map to abstract types defined by the "module"
#   entry (Bus, ElectricLoad, Generator, LoadZone, Reserve)
# - `component_name`:  Name of component
# - `name`:  User-defined name for the time series data.
# - `normalization_factor`:  Controls normalization of the data. Use 1.0 for
#   pre-normalized data. Use 'Max' to divide the time series by the max value in the
#   column. Use any float for a custom scaling factor.
# - `scaling_factor_multiplier_module`:  Module that defines the accessor function for the
# scaling factor
# - `scaling_factor_multiplier`:  Accessor function of the scaling factor
# - `data_file`:  Path to the time series data file
#
# Notes:
#
# - The "module", "category", and "component_name" entries must be valid arguments to retrieve
# a component using `get_component(${module}.${category}, sys, $name)`.
# - The "scaling_factor_multiplier_module" and the "scaling_factor_multiplier" entries must
# be sufficient to return the scaling factor data using
# `${scaling_factor_multiplier_module}.${scaling_factor_multiplier}(component)`.
#
# PowerSystems supports this metadata in either CSV or JSON formats. Refer to
# [RTS_GMLC](https://github.com/GridMod/RTS-GMLC/blob/master/RTS_Data/FormattedData/SIIP/timeseries_pointers.json)
# for an example.
#
# #### Performance considerations
#
# By default PowerSystems stores time series data in HDF5 files. It does not keep
# all of the data in memory. This means that every time you access a time_series
# PowerSystems will have to read the data from storage, which will add latency. If
# you know ahead of time that all of your data will fit in memory then you can
# change this behavior by passing `time_series_in_memory = true` when you create
# the System.
#
# If the time series data is stored in HDF5 then PowerSystems will use the tmp filesystem by
# default. You can change this by passing `time_series_directory = X` when you create the
# System. This is required if the time series data is larger than the amount of tmp space
# available. You can also override the location by setting the environment variable
# SIIP_TIME_SERIES_DIRECTORY to another directory.
#
# ### Customization
#
# The tabular data parser in `PowerSystems.jl` can be customized to read a variety of
# datasets by configuring:
#
#  - [which type of generator (`<:Generator`) to create based on the fuel and prime mover specifications](@ref csv_genmap)
#  - [property names](@ref csv_columns), [units](@ref csv_units), and per units conversions](@ref csv_per_unit) in *.csv files
#
# Here is an example of how to construct a System with all customizations listed in this section:
#
# ```julia
# data_dir = "/data/my-data-dir"
# base_power = 100.0
# descriptors = "./user_descriptors.yaml"
# timeseries_metadata_file = "./timeseries_pointers.json"
# generator_mapping_file = "./generator_mapping.yaml"
# data = PowerSystemTableData(
#     data_dir,
#     base_power,
#     descriptors;
#     timeseries_metadata_file = timeseries_metadata_file,
#     generator_mapping_file = generator_mapping_file,
# )
# sys = System(data, time_series_in_memory = true)
# ```
#
# Examples configuration files can be found in the [RTS-GMLC](https://github.com/GridMod/RTS-GMLC/) repo:
#
#  - [user_descriptors.yaml](https://github.com/GridMod/RTS-GMLC/blob/master/RTS_Data/FormattedData/SIIP/user_descriptors.yaml)
#  - [generator_mapping.yaml](https://github.com/GridMod/RTS-GMLC/blob/master/RTS_Data/FormattedData/SIIP/generator_mapping.yaml)
#
# #### [CSV Data Configurations](@id csv_data)
#
# ##### [Custom construction of generators](@id csv_genmap)
#
# PowerSystems supports custom construction of subtypes of the abstract type Generator based
# on `fuel` and `type`. The parsing code detects these fields in the raw data and then
# constructs the concrete type listed in the passed generator mapping file. The default file
# is `src/parsers/generator_mapping.yaml`. You can override this behavior by specifying your
# own file when constructing PowerSystemTableData.
#
# ##### [Column names](@id csv_columns)
#
# PowerSystems provides am input mapping capability that allows you to keep your own
# column names.
#
# For example, when parsing raw data for a generator the code expects a column
# called `name`. If the raw data instead defines that column as `GEN UID` then
# you can change the `custom_name` field under the `generator` category to
# `GEN UID` in your YAML file.
#
# To enable the parsing of a custom set of csv files, you can generate a configuration
# file (such as `user_descriptors.yaml`) from the defaults, which are stored
# in `src/descriptors/power_system_inputs.json`.
#
# ```python
# python ./bin/generate_config_file.py ./user_descriptors.yaml
# ```
#
# Next, edit this file with your customizations.
#
# Note that the user-specific customizations are stored in YAML rather than JSON
# to allow for easier editing. The next few sections describe changes you can
# make to this YAML file.  Do not edit the default JSON file.
#
# A list of `PowerSystemTableData` accepted columns for each CSV is available
# [here](@ref tabledata_input_config).
#
# ##### [Per-unit conversion](@id csv_per_unit)
#
# For more info on the per-unit conventions in `PowerSystems.jl`, refer to the [per-unit
# section of the system documentation](@ref per_unit).
#
# PowerSystems defines whether it expects a column value to be per-unit system base,
# per-unit device base, or natural units in `power_system_inputs.json`
# ([see default values](@ref tabledata_input_config)). If it expects a
# per-unit convention that differs from your values then you can set the `unit_system` in
# `user_descriptors.yaml` and PowerSystems will automatically convert the values. For
# example, if you have a `max_active_power` value stored in natural units (MW), but
# `power_system_inputs.json` specifies `unit_system: device_base`, you can enter
# `unit_system: natural_units` in `user_descriptors.yaml` and PowerSystems will divide
# the value by the value of the corresponding entry in the column identified by the
# `base_reference` field in `power_system_inputs.json`. You can also override the
# `base_reference` setting by adding `base_reference: My Column` to make device base
# per-unit conversion by dividing the value by the entry in `My Column`. System base
# per-unit conversions always divide the value by the system `base_power` value
# instantiated when constructing a `System`.
#
# ##### [Unit conversion](@id csv_units)
#
# PowerSystems provides a limited set of unit conversions. For example, if
# `power_system_inputs.json` indicates that a value's unit is degrees but
# your values are in radians then you can set `unit: radian` in
# your YAML file. Other valid `unit` entries include `GW`, `GWh`, `MW`, `MWh`, `kW`,
# and `kWh`.
#

# ## [`PowerSystemTableData` Accepted CSV Columns](@id tabledata_input_config)
#
# The following tables describe default CSV column definitions accepted by the
# `PowerSystemeTableData` parser defined by `src/descriptors/power_system_inputs.json`:
#

using PowerSystems # hide
function create_md() # hide
    descriptor = PowerSystems._read_config_file(joinpath( # hide
        dirname(pathof(PowerSystems)), # hide
        "descriptors", # hide
        "power_system_inputs.json", # hide
    )) # hide
    columns = [ # hide
        "name", # hide
        "description", # hide
        "unit", # hide
        "unit_system", # hide
        "base_reference", # hide
        "default_value", # hide
        "value_options", # hide
        "value_range", # hide
    ] # hide
    header = "| " * join(columns, " | ") * " |\n" * repeat("|----", length(columns)) * "|\n" # hide
    s = "" # hide
    for (cat, items) in descriptor # hide
        csv = "" # hide
        for name in PowerSystems.INPUT_CATEGORY_NAMES # hide
            if name[2] == cat # hide
                csv = name[1] # hide
                break # hide
            end # hide
        end # hide
        csv == "" && continue # hide
        s = string(s, "### $csv.csv:\n\n") # hide
        s = string(s, header) # hide
        for item in items # hide
            extra_cols = setdiff(keys(item), columns) # hide
            if !isempty(extra_cols) # hide
                throw(@error "config file fields not included in header" extra_cols) # hide
            end # hide
            row = [] # hide
            for col in columns # hide
                val = string(get(item, col, " ")) # hide
                if col == "default_value" && val == " " # hide
                    val = "*REQUIRED*" # hide
                end # hide
                push!(row, val) # hide
            end # hide
            s = string(s, "|" * join(row, "|") * "|\n") # hide
        end # hide
        s = string(s, "\n") # hide
    end # hide
    s = replace(s, r"[_$]" => s"\\\g<0>"); # hide
    return s # hide
end # hide
txt = create_md(); # hide
fname = joinpath(dirname(dirname(pathof(PowerSystems))), "docs", "src", "modeler_guide", "generated_inputs_tables.md") # hide
open(fname, "w") do io # hide
      write(io, txt) # hide
end # hide
println(" "); # hide

# APPEND_MARKDOWN("docs/src/modeler_guide/generated_inputs_tables.md")


################################################################################
#
#	ABSTRACT TYPE
#
#   BrillouinZone{R}
#   --> R is the reciprocal unitcell type (<: AbstractReciprocalUnitcell{P,B})
#       --> P is the reciprocal site type (<: AbstractReciprocalPoint{D})
#           --> D is the dimension of embedding space
#       --> B is the bond type (<: AbstractBond{L,N})
#           --> N is the dimension of the Bravais lattice that the bond is located in
#           --> L is the label type of bonds
#
#   FILE CONTAINS
#       - abstract type definition
#       - interface definition
#       - TODO interface testing
#
################################################################################

################################################################################
#
#   ABSTRACT TYPE DEFINITION
#
################################################################################
abstract type AbstractBrillouinZone{
        R <: AbstractReciprocalUnitcell{P,B} where {
            P <: AbstractReciprocalPoint{D} where {D},
            B <: AbstractBond{L,N} where {L,N},
        }
    } end

# export the type
export AbstractBrillouinZone

################################################################################
#
#	INTERFACING / ACCESSING BRILLOUIN ZONES
#	(functions that have to be overwritten by concrete types)
#
################################################################################

# constructor
function newBrillouinZone(
            ::Type{BZ},
            reciprocal_unitcell :: RU,
            corners             :: Vector{<:Vector{<:Real}},
            faces               :: Vector{<:Vector{<:Integer}}
        ) :: BZ where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    # throw an error because implementation for concrete type is missing
    error(  "not implemented function 'newBrillouinZone' for concrete Brillouin zone type " *
            string(BZ) * " with reciprocal unitcell type " * string(RU) )
end

# export constructor
export newBrillouinZone

# accessing the corner points of the BZ
function corners(
            bz :: BZ
        ) :: Vector{Vector{Float64}} where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    # throw an error because implementation for concrete type is missing
    error(  "not implemented function 'corners' for concrete Brillouin zone type " *
            string(BZ) * " with reciprocal unitcell type " * string(RU) )
end

# setting the corner points of the BZ
function corners!(
            bz      :: BZ,
            corners :: Vector{<:Vector{<:Real}}
        ) where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    # throw an error because implementation for concrete type is missing
    error(  "not implemented function 'corners!' for concrete Brillouin zone type " *
            string(BZ) * " with reciprocal unitcell type " * string(RU) )
end

# export corner interface
export corners, corners!

# accessing the faces of the BZ
function faces(
            bz :: BZ
        ) :: Vector{Vector{Int64}} where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    # throw an error because implementation for concrete type is missing
    error(  "not implemented function 'faces' for concrete Brillouin zone type " *
            string(BZ) * " with reciprocal unitcell type " * string(RU) )
end

# setting the faces of the BZ
function faces!(
            bz    :: BZ,
            faces :: Vector{<:Vector{<:Integer}}
        ) where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    # throw an error because implementation for concrete type is missing
    error(  "not implemented function 'faces!' for concrete Brillouin zone type " *
            string(BZ) * " with reciprocal unitcell type " * string(RU) )
end

# export faces interface
export faces, faces!

# accessing the reciprocal unitcell of the BZ
function reciprocalUnitcell(
            bz :: BZ
        ) :: RU where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    # throw an error because implementation for concrete type is missing
    error(  "not implemented function 'reciprocalUnitcell' for concrete Brillouin zone type " *
            string(BZ) * " with reciprocal unitcell type " * string(RU) )
end

# setting the reciprocal unitcell of the BZ
function reciprocalUnitcell!(
            bz  :: BZ,
            ruc :: RU
        ) where {D,N,L,P<:AbstractReciprocalPoint{D},B<:AbstractBond{L,N},RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}

    # throw an error because implementation for concrete type is missing
    error(  "not implemented function 'reciprocalUnitcell!' for concrete Brillouin zone type " *
            string(BZ) * " with reciprocal unitcell type " * string(RU) )
end

# export reciprocalUnitcell interface
export reciprocalUnitcell, reciprocalUnitcell!

# CUSTOM SHOW FUNCTION

# single Brillouin Zone
function show(io::IO, bz::BZ) where {P,B,RU<:AbstractReciprocalUnitcell{P,B},BZ<:AbstractBrillouinZone{RU}}
    print(io, "Brillouin Zone object\n--> reciprocal unitcell type ", RU, "\n--> type ", BZ, "\n--> ", length(corners(bz)), " BZ corners\n--> ", length(faces(bz)), " BZ faces")
end


# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

@doc raw"""Spec to control the desired behavior of daemon set rolling update.

    IoK8sApiAppsV1RollingUpdateDaemonSet(;
        maxUnavailable=nothing,
    )

    - maxUnavailable::IoK8sApimachineryPkgUtilIntstrIntOrString : The maximum number of DaemonSet pods that can be unavailable during the update. Value can be an absolute number (ex: 5) or a percentage of total number of DaemonSet pods at the start of the update (ex: 10%). Absolute number is calculated from percentage by rounding up. This cannot be 0. Default value is 1. Example: when this is set to 30%, at most 30% of the total number of nodes that should be running the daemon pod (i.e. status.desiredNumberScheduled) can have their pods stopped for an update at any given time. The update starts by stopping at most 30% of those DaemonSet pods and then brings up new DaemonSet pods in their place. Once the new pods are available, it then proceeds onto other DaemonSet pods, thus ensuring that at least 70% of original number of DaemonSet pods are available at all times during the update.
"""
mutable struct IoK8sApiAppsV1RollingUpdateDaemonSet <: SwaggerModel
    maxUnavailable::Any # spec type: Union{ Nothing, IoK8sApimachineryPkgUtilIntstrIntOrString } # spec name: maxUnavailable

    function IoK8sApiAppsV1RollingUpdateDaemonSet(;maxUnavailable=nothing)
        o = new()
        validate_property(IoK8sApiAppsV1RollingUpdateDaemonSet, Symbol("maxUnavailable"), maxUnavailable)
        setfield!(o, Symbol("maxUnavailable"), maxUnavailable)
        o
    end
end # type IoK8sApiAppsV1RollingUpdateDaemonSet

const _property_map_IoK8sApiAppsV1RollingUpdateDaemonSet = Dict{Symbol,Symbol}(Symbol("maxUnavailable")=>Symbol("maxUnavailable"))
const _property_types_IoK8sApiAppsV1RollingUpdateDaemonSet = Dict{Symbol,String}(Symbol("maxUnavailable")=>"IoK8sApimachineryPkgUtilIntstrIntOrString")
Base.propertynames(::Type{ IoK8sApiAppsV1RollingUpdateDaemonSet }) = collect(keys(_property_map_IoK8sApiAppsV1RollingUpdateDaemonSet))
Swagger.property_type(::Type{ IoK8sApiAppsV1RollingUpdateDaemonSet }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiAppsV1RollingUpdateDaemonSet[name]))}
Swagger.field_name(::Type{ IoK8sApiAppsV1RollingUpdateDaemonSet }, property_name::Symbol) =  _property_map_IoK8sApiAppsV1RollingUpdateDaemonSet[property_name]

function check_required(o::IoK8sApiAppsV1RollingUpdateDaemonSet)
    true
end

function validate_property(::Type{ IoK8sApiAppsV1RollingUpdateDaemonSet }, name::Symbol, val)
end


"""
Accumulator-like data structure for storing frequencies of CJK words (although other tokens
can be stored as well).

You generally don't need to explicitly call this struct's constructor yourself; rather, use
the `charfreq` function.
"""
struct CJKFrequency{S <: AbstractString, C <: Number}
    freq::Accumulator{S, C}
    size::Ref{C}
    function CJKFrequency(frequencies=Dict{String, Int}())
        key_type = eltype(keys(frequencies))
        val_type = eltype(values(frequencies))
        new{key_type, val_type}(counter(frequencies), Ref(sum(values(frequencies))))
    end
end

# mutation
function DataStructures.inc!(cf::CJKFrequency, key, count=1)
    inc!(cf.freq, key, count)
    cf.size[] += count
end
function DataStructures.dec!(cf::CJKFrequency, key, count=1)
    if key in keys(cf.freq)
        cf.size[] -= cf.freq[key] - max(0, cf.freq[key] - count)
        dec!(cf.freq, key, min(count, cf.freq[key]))
    end
end
function DataStructures.reset!(cf::CJKFrequency, key)
    key in keys(cf.freq) && (size[] -= cf.freq[key])
    reset!(cf.freq, key)
end

# collections operations
Base.isempty(cf::CJKFrequency) = isempty(cf.freq)
function Base.empty!(cf::CJKFrequency)
     empty!(cf.freq)
     size = 0
     return cf
end
Base.in(item, cf::CJKFrequency) = in(item, cf.freq)

# dictionaries
Base.haskey(cf::CJKFrequency, key) = haskey(cf.freq, key)
Base.keys(cf::CJKFrequency) = keys(cf.freq)
Base.values(cf::CJKFrequency) = values(cf.freq)

# indexing
Base.getindex(cf::CJKFrequency, i) = getindex(cf.freq, i)
Base.firstindex(cf::CJKFrequency) = firstindex(cf.freq)
Base.lastindex(cf::CJKFrequency) = lastindex(cf.freq)

# iteration
Base.iterate(cf::CJKFrequency) = iterate(cf.freq)
Base.iterate(cf::CJKFrequency, state) = iterate(cf.freq, state)
Base.length(cf::CJKFrequency) = length(cf.freq)
Base.size(cf::CJKFrequency) = cf.size[]

"""
    charfreq(text)
    charfreq(charfreq_type)

Create a character frequency mapping from either text or load it from a default location for
pre-specified character frequency datasets (e.g. `SimplifiedLCMC`, `SimplifiedJunDa`, etc.).

## Examples
When creating a character frequency from text, this method behaves almost exactly like
`DataStructures.counter` except that the return value always has type `CharacterFrequency`
(`Accumulator{String, Int}`).
```julia-repl
julia> text = split("王老师性格内向，沉默寡言，我除在课外活动小组“文学研究会”听过他一次报告，并听-邓知识渊博，是“老师的老师”外，对他一无所知。所以，研读他的作", "");

julia> charfreq(text)
CJKFrequency{SubString{String}, Int64}(Accumulator(除 => 1, 报 => 1, 是 => 1, 知 => 2, 并 => 1, 性 => 1, ， => 6, 言 => 1, 邓 => 1, 外 => 2, 所 => 2, 对 => 1, 动 => 1, 寡 => 1, 。 => 1, 渊 => 1, 学 => 1, - => 1, 听 => 2, 我 => 1, 次 => 1, 一 => 2, 读 => 1, 作 => 1, 格 => 1, “ => 2, 博 => 1, 课 => 1, 老 => 3, 会 => 1, 告 => 1, 无 => 1, 活 => 1, 组 => 1, 内 => 1, 师 => 3, 的 => 2, 小 => 1, 文 => 1, 默 => 1, 究 => 1, 过 => 1, 在 => 1, 以 => 1, ” => 2, 研 => 2, 他 => 3, 向 => 1, 沉 => 1, 王 => 1), Base.RefValue{Int64}(71))
```

See the documentation for individual character frequency dataset structs for examples of the
second case.
"""
function charfreq end
charfreq(tokens) = CJKFrequency(counter(tokens))
charfreq(text::AbstractString) = CJKFrequency(counter(split(text, "")))


export Episode, get_inputs, step!, run!

struct Episode
    robot::Robot
    grid::Grid
    reward::Function
    meta::Dict
end

no_reward(e::Episode) = 0

function Episode(g::Grid; reward::Function=no_reward, meta::Dict=Dict())
    robot = Robot(g)
    meta["step"] = get(meta, "step", 0)
    meta["max_step"] = get(meta, "max_step", 100)
    meta["time"] = get(meta, "time", 0.0)
    meta["max_time"] = get(meta, "max_time", 10.0)
    meta["memory"] = get(meta, "memory", 0)
    meta["max_memory"] = get(meta, "max_memory", 1073741824)
    meta["reward"] = get(meta, "reward", 0.0)
    Episode(robot, g, reward, meta)
end

function get_inputs(e::Episode)
    cids = Array{Int64}(undef, (3, 3))
    linds = LinearIndices(e.grid.cells)
    i = 1
    for dy in -1:1
        for dx in -1:1
            cids[i] = get(linds, (e.robot.y + dy, e.robot.x + dx), 0)
            i += 1
        end
    end
    if e.robot.d > 0
        cids = [rotr90, rot180, rotl90][e.robot.d](cids)
    end
    inputs = Array{Float64}(undef, 10)
    inputs[1] = e.meta["reward"]
    i = 2
    for cid in cids
        c = get(e.grid.cells, cid, WallCell)
        inputs[i] = c.obj
        # inputs[i+1] = c.color
        i += 1
    end
    inputs
end

function step!(e::Episode, action::Int64)
    if action == 1
        move!(e.robot, e.grid)
    elseif action == 2
        move!(e.robot, e.grid; forward=false)
    elseif action == 3
        turn_right!(e.robot)
    elseif action == 4
        turn_left!(e.robot)
    elseif action == 5
        e.grid.cells[e.robot.y, e.robot.x].color = color_to_float("red")
    elseif action == 6
        e.grid.cells[e.robot.y, e.robot.x].color = color_to_float("green")
    end
    e.meta["reward"] = e.reward(e)
    e.meta["reward"]
end

function step!(e::Episode, actions::AbstractArray)
    step!(e, argmax(actions))
end

function step!(e::Episode, controller::Function)
    inputs = get_inputs(e)
    output = controller(inputs)
    step!(e, output)
end

function terminate(e::Episode)
    (e.meta["step"] >= e.meta["max_step"] ||
     e.meta["time"] >= e.meta["max_time"] ||
     e.meta["memory"] >= e.meta["max_memory"])
end

function run!(e::Episode, controller::Function; terminate::Function=terminate)
    linds = LinearIndices(e.grid.cells)
    total_reward = 0.0
    while true
        rew, t, bytes, gctime, memallocs = @timed step!(e, controller)
        total_reward += rew
        e.meta["step"] += 1
        e.meta["time"] += t
        e.meta["memory"] += bytes
        if (terminate(e) || (linds[e.robot.y, e.robot.x] in e.grid.exits))
            break
        end
    end
    total_reward
end


@everywhere push!(LOAD_PATH,"solver")
@everywhere push!(LOAD_PATH,"mdp")
@everywhere using ParallelValueIteration
@everywhere using VerticalCAS
using HDF5

### OPTIONS ###
saveFile = "./VertCAS_qvals_parallel_v7_tauMax40.h5"
# Run in parallel on all available processors
nProcs = nprocs()-1
mdp = VerticalCAS_MDP()
approx_solver = ParallelValueIterationSolver(interp, verbose=true, max_iterations=1,n_procs=nProcs)

# Solve
println("Solving...")
qvals = solve(approx_solver, mdp)

# Write Qvalues to a file
println("Writing Qvalues")
h5open(saveFile, "w") do file
    write(file, "q", qvals)
end
end


# MIRT.jl

"""
`MIRT` is the Michigan Image Reconstruction Toolbox
"""
module MIRT
	using Reexport

	@reexport using MIRTio # make I/O routines available

    const RealU = Number # Union{Real, Unitful.Length}
	include("z-all.jl")

end # module