module VectorSpaces

using Base.Cartesian
using Arbitrary
using Combinatorics: levicivita, permutations
using IterTools: imap
using SimpleTraits

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

# @traitdef VectorSpace{V}
export VectorSpace
abstract type VectorSpace{S} end

# @traitimpl VectorSpace{V} <- isvectorspace(V)
function isvectorspace(::Type{V}) where {V <: VectorSpace}
    try
        # Element type
        S = stype(V)
        S <: Number || return false
        V8 = retype(V, Int8)
        S8 = stype(V8)
        S8 === Int8 || return false
        VS = retype(V8, S)
        VS === V || return false
        # Constructors
        z = zero(V)::V
        e1 = ones(V)::V
        # Collection properties
        @assert eltype(V) === S
        l = length(e1)
        @assert l == length(collect(e1))
        @assert all(==(S(1)), collect(e1))
        @assert all(==(S(1)), e1)
        @assert all(==(S(3)), map(a->2a+1, e1))
        elin = linear(V)
        @assert sum(elin) == l*(l+1)/2
        # Check whether functions exist and return the right types
        (z * S(1))::V
        S <: Union{AbstractFloat, Rational} && (z / S(1))::V
        (S(1) * z)::V
        S <: Union{AbstractFloat, Rational} && (S(1) \ z)::V
        (+e1)::V
        (-e1)::V
        (z + e1)::V
        (z - e1)::V
        (e1 ⋅ e1)::S
        incomplete_norm(e1)::S
        if S <: AbstractFloat
            norm(e1)::S
        end
        # Check vector space laws
        e1 + z == e1 || return false
        z + e1 == e1 || return false
        S(2) * z == z || return false
        e2 = S(2) * e1
        e5 = S(5) * e1
        e1 + (-e1) == z || return false
        e2 + e1 == e1 + e2 || return false
        e1 + (e2 + e5) == (e1 + e2) + e5 || return false
        # Check dot product laws
        z ⋅ z == S(0) || return false
        if e5 == e2
            e5 ⋅ e2 == S(0) || return false
        else
            e5 ⋅ e2 != S(0) || return false
        end
        (S(3) * e5) ⋅ e2 == S(3) * (e5 ⋅ e2) || return false
        # Check norm laws
        incomplete_norm(z) == S(0) || return false
        if e2 == z
            incomplete_norm(e2) == S(0) || return false
        else
            incomplete_norm(e2) > S(0) || return false
        end
        if S <: AbstractFloat
            norm(z) == S(0) || return false
            if e2 == z
                norm(e2) == S(0) || return false
            else
                norm(e2) > S(0) || return false
            end
        end
        return true
    catch
        return false
    end
end

using LinearAlgebra

export stype, retype
export linear
export ⋅, incomplete_norm, partial_mean, norm, mean

function Base.zeros(::Type{V}) where {V <: VectorSpace}
    S = stype(V)
    map(_ -> zero(S), zero(V))::V
end
function Base.ones(::Type{V}) where {V <: VectorSpace}
    S = stype(V)
    map(_ -> one(S), zero(V))::V
end
function linear(::Type{V}) where {V <: VectorSpace}
    S = stype(V)
    idx = 0
    inc(_) = S(idx += 1)
    map(inc, zero(V))::V
end

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

export Vec
struct Vec{D, S} <: VectorSpace{S}
    elts::NTuple{D, S}
    # Vec(elts::NTuple{D, S}) where {D, S<:Number} = new{D, S}(elts)
end

function Base.show(io::IO, x::Vec{D, S}) where {D, S}
    print(io, "[")
    for d in 1:D
        d>1 && print(io, ", ")
        print(io, x.elts[d])
    end
    print(io, "]")
end

# function Vec{D, S}(elts...) where {D, S<:Number}
#     elts :: NTuple{D, S}
#     Vec(elts)
# end
# Vec(elts::NTuple{D, S}) where {D, S} = Vec{D, S}(elts)
Vec(r::AbstractRange) = Vec(tuple(r...))

Base.firstindex(x::Vec) = firstindex(x.elts)
Base.lastindex(x::Vec) = lastindex(x.elts)
Base.getindex(x::Vec, i) = x.elts[i]
# Base.IteratorSize(::Type{Vec{D, S}}) where {D, S} = HasLength()
# Base.IteratorEltype(::Type{Vec{D, S}}) where {D, S} = HasEltype()
Base.eltype(::Type{Vec{D, S}}) where {D, S} = S
# Base.size(x::Vec) = (length(x), )
Base.length(x::Vec) = length(x.elts)
Base.iterate(x::Vec) = iterate(x.elts)
Base.iterate(x::Vec, st) = iterate(x.elts, st)
function Base.map(f, x::Vec{D}) where {D}
    relts = map(f, x.elts)
    R = eltype(relts)
    Vec{D, R}(relts)
end
function Base.map(f, x::Vec{D}, y::Vec{D}) where {D}
    relts = map(f, x.elts, y.elts)
    Vec{D, eltype(relts)}(relts)
end

Base.isequal(x::Vec{D}, y::Vec{D}) where {D} = isequal(x.elts, y.elts)
Base. ==(x::Vec{D}, y::Vec{D}) where {D} = x.elts == y.elts

function Arbitrary.arbitrary(::Type{Vec{D, S}}, ast::ArbState) where {D, S}
    imap(x -> Vec{D, S}(x), arbitrary(NTuple{D, S}, ast))
end

stype(::Type{Vec{D, S}}) where {D, S} = S
retype(::Type{Vec{D, S}}, ::Type{T}) where {D, S, T} = Vec{D, T}

# This needs to accept either a Number or a VectorSpace
Base.zero(::Type{Vec{D, S}}) where {D, S <: Union{Number, VectorSpace}} =
    Vec{D, S}(ntuple(_->zero(S), D))
Base. +(x::Vec{D, S}) where {D, S <: Union{Number, VectorSpace}} =
    Vec{D, S}(map(+, x.elts))
Base. -(x::Vec{D, S}) where {D, S <: Union{Number, VectorSpace}} =
    Vec{D, S}(map(-, x.elts))
Base. *(a::S, x::Vec{D, S}) where {D, S <: Union{Number, VectorSpace}} =
    Vec{D, S}(map(c->a*c, x.elts))
Base. \(a::S, x::Vec{D, S}) where
        {D, S <: Union{AbstractFloat, Rational, VectorSpace}} =
    Vec{D, S}(map(c->a\c, x.elts))
Base. *(x::Vec{D, S}, a::S) where {D, S <: Union{Number, VectorSpace}} =
    Vec{D, S}(map(c->c*a, x.elts))
Base. /(x::Vec{D, S}, a::S) where
        {D, S <: Union{AbstractFloat, Rational, VectorSpace}} =
    Vec{D, S}(map(c->c/a, x.elts))

Base. +(x::Vec{D, S}, y::Vec{D, S}) where {D, S <: Union{Number, VectorSpace}} =
    Vec{D, S}(map(+, x.elts, y.elts))
Base. -(x::Vec{D, S}, y::Vec{D, S}) where {D, S <: Union{Number, VectorSpace}} =
    Vec{D, S}(map(-, x.elts, y.elts))

# @assert isvectorspace(Vec{1, Int})
# @traitimpl VectorSpace{Vec{S}}

LinearAlgebra. ⋅(x::Vec{D, S}, y::Vec{D, S}) where
        {D, S <: Union{Number, VectorSpace}} =
    sum(map(*, x, y))::S

incomplete_norm(x::Vec{D, S}) where {D, S <: Union{Number, VectorSpace}} =
    sum(abs.(x.elts) .^ 2)::S
LinearAlgebra.norm(x::Vec{D, S}) where
        {D, S <: Union{AbstractFloat, VectorSpace}} =
    sqrt(incomplete_norm(x))

function incomplete_norm(x::Vec{D, S}, p::Real) where
        {D, S <: Union{AbstractFloat, VectorSpace}}
    if p == 0
        S(D)
    elseif p == 1
        sum(abs.(x.elts))
    elseif p == 2
        sum(abs.(x.elts) .^ 2)
    elseif p == Inf
        max(abs.(x.elts))
    elseif isinteger(p)
        sum(abs.(x.elts) .^ Int(p))
    else
        sum(abs.(x.elts) .^ p)
    end
end
function LinearAlgebra.norm(x::Vec{D, S}, p::Real) where
        {D, S <: Union{AbstractFloat, VectorSpace}}
    r = incomplete_norm(x, p)
    if p == 0
        r
    elseif p == 1
        r
    elseif p == 2
        sqrt(r)
    else
        r ^ (1/p)
    end
end

function partial_mean(x::Vec{D, S}, p::Real) where
        {D, S <: Union{AbstractFloat, VectorSpace}}
    if p == -Inf
        min(x.elts), S(D)
    elseif p == -1
        sum(inv.(x.elts)), S(D)
    elseif p == 0
        product(x.elts), S(D)
    elseif p == 1
        sum(x.elts), S(D)
    elseif p == 2
        sum(x.elts .^ 2), S(D)
    elseif p == Inf
        max(x.elts), S(D)
    elseif isinteger(p)
        sum(x.elts .^ Int(p)), S(D)
    else
        sum(x.elts .^ p), S(D)
    end
end
function mean(x::Vec{D, S}, p::Real) where
        {D, S <: Union{AbstractFloat, VectorSpace}}
    r, n = partial_mean(x, p)
    if p == -Inf
        r
    elseif p == -1
        inv(r / n)
    elseif p == 0
        r ^ inv(n)
    elseif p == 1
        r / n
    elseif p == 2
        sqrt(r / n)
    elseif p == Inf
        r
    else
        (r / n) ^ (1/p)
    end
end

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

export VProd
struct VProd{V1, V2, S} <: VectorSpace{S}
    v1::V1
    v2::V2
    function VProd{V1, V2, S}(v1::V1, v2::V2) where
            {S, V1 <: VectorSpace{S}, V2 <: VectorSpace{S}}
       new{V1, V2, S}(v1, v2)
   end
end

function VProd(v1::V1, v2::V2) where {V1 <: VectorSpace, V2 <: VectorSpace}
    S1 = stype(V1)
    S2 = stype(V2)
    @assert S1 === S2
    VProd{V1, V2, S1}(v1, v2)
end

Base.show(io::IO, x::VProd) = print(io, "($(x.v1), $(x.v2))")

Base.firstindex(x::VProd) = 1
Base.lastindex(x::VProd) = length(x)
function Base.getindex(x::VProd, i)
    l1 = length(x.v1)
    if i <= l1
        return x.v1[i - 1 + firstindex(x.v1)]
    end
    i = i - l1
    l2 = length(x.v2)
    if i <= l2
        return x.v2[i - 1 + firstindex(x.v2)]
    end
    @assert false
end
Base.eltype(::Type{VProd{V1, V2, S}}) where {V1, V2, S} = S
Base.length(x::VProd) = length(x.v1) + length(x.v2)
function Base.iterate(x::VProd)
    i1 = iterate(x.v1)
    if i1 !== nothing
        el, st1 = i1
        return el, (false, st1)
    end
    i2 = iterate(x.v2)
    if i2 !== nothing
        el, st2 = i2
        return el, (true, st2)
    end
    return nothing
end
function Base.iterate(x::VProd, st)
    b, st12 = st
    if b === false
        i1 = iterate(x.v1, st12)
        if i1 !== nothing
            el, st1 = i1
            return el, (false, st1)
        end
        i2 = iterate(x.v2)
        if i2 !== nothing
            el, st2 = i2
            return el, (true, st2)
        end
        return nothing
    end
    i2 = iterate(x.v2, st12)
    if i2 !== nothing
        el, st2 = i2
        return el, (true, st2)
    end
    return nothing
end
function Base.map(f, x::VProd{V1, V2}) where {V1, V2}
    r1 = map(f, x.v1)
    r2 = map(f, x.v2)
    R = promote_type(stype(typeof(r1)), stype(typeof(r2)))
    VProd{retype(V1, R), retype(V2, R), R}(r1, r2)
end
function Base.map(f, x::VProd{V1, V2}, y::VProd{V1, V2}) where {V1, V2}
    r1 = map(f, x.v1, y.v1)
    r2 = map(f, x.v2, y.v2)
    R = eltype(Tuple{stype(r1), stype(r2)})
    VProd{retype(V1, R), retype(V2, R), R}(r1, r2)
end

ltype(::Type{VProd{V1, V2, S}}) where {V1, V2, S} = V1
rtype(::Type{VProd{V1, V2, S}}) where {V1, V2, S} = V2
stype(::Type{VProd{V1, V2, S}}) where {V1, V2, S} = S
retype(::Type{VProd{V1, V2, S}}, ::Type{T}) where {V1, V2, S, T} =
    VProd{retype(V1, T), retype(V2, T), T}

Base.zero(::Type{VProd{V1, V2, S}}) where
        {V1, V2, S <: Union{Number, VectorSpace}} =
    VProd{V1, V2, S}(zero(V1), zero(V2))

Base. +(x::VProd{V1, V2, S}) where {V1, V2, S} =
    VProd{V1, V2, S}(+x.v1, +x.v2)
Base. -(x::VProd{V1, V2, S}) where {V1, V2, S} =
    VProd{V1, V2, S}(-x.v1, -x.v2)
Base. *(a::S, x::VProd{V1, V2, S}) where {V1, V2, S} =
    VProd{V1, V2, S}(a*x.v1, a*x.v2)
Base. \(a::S, x::VProd{V1, V2, S}) where
        {V1, V2, S <: Union{AbstractFloat, Rational}} =
    VProd{V1, V2, S}(a\x.v1, a\x.v2)
Base. *(x::VProd{V1, V2, S}, a::S) where {V1, V2, S} =
    VProd{V1, V2, S}(x.v1*a, x.v2*a)
Base. /(x::VProd{V1, V2, S}, a::S) where
        {V1, V2, S <: Union{AbstractFloat, Rational}} =
    VProd{V1, V2, S}(x.v1/a, x.v2/a)

Base. +(x::VProd{V1, V2, S}, y::VProd{V1, V2, S}) where {V1, V2, S} =
    VProd{V1, V2, S}(x.v1+y.v1, x.v2+y.v2)
Base. -(x::VProd{V1, V2, S}, y::VProd{V1, V2, S}) where {V1, V2, S} =
    VProd{V1, V2, S}(x.v1-y.v1, x.v2-y.v2)

LinearAlgebra. ⋅(x::VProd{V1, V2, S}, y::VProd{V1, V2, S}) where {V1, V2, S} =
    (x.v1 ⋅ y.v1 + x.v2 ⋅ y.v2)::S

incomplete_norm(x::VProd{V1, V2, S}) where {V1, V2, S <: Number} =
    (incomplete_norm(x.v1) + incomplete_norm(x.v2))::S
LinearAlgebra.norm(x::VProd{V1, V2, S}) where {V1, V2, S <: AbstractFloat} =
    sqrt(incomplete_norm(x))::S

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

export VEmpty
struct VEmpty{S} <: VectorSpace{S} end

Base.show(io::IO, x::VEmpty) = print(io, "()")

Base.firstindex(x::VEmpty) = 1
Base.lastindex(x::VEmpty) = 0
Base.eltype(::Type{VEmpty{S}}) where {S} = S
Base.length(x::VEmpty) = 0
Base.iterate(x::VEmpty) = nothing
Base.iterate(x::VEmpty, st) = @assert false
function Base.map(f, x::VEmpty{S}) where {S}
    r = f(S(1))
    R = typeof(r)
    VEmpty{R}()
end
function Base.map(f, x::VEmpty{S}, y::VEmpty{S}) where {S}
    r = f(S(1), S(1))
    R = typeof(r)
    VEmpty{R}()
end

stype(::Type{VEmpty{S}}) where {S} = S
retype(::Type{VEmpty{S}}, ::Type{T}) where {S, T} = VEmpty{T}

Base.zero(::Type{VEmpty{S}}) where {S <: Union{Number, VectorSpace}} =
    VEmpty{S}()

Base. +(x::VEmpty{S}) where {S} = VEmpty{S}()
Base. -(x::VEmpty{S}) where {S} = VEmpty{S}()
Base. *(a::S, x::VEmpty{S}) where {S} = VEmpty{S}()
Base. \(a::S, x::VEmpty{S}) where {S} = VEmpty{S}()
Base. *(x::VEmpty{S}, a::S) where {S} = VEmpty{S}()
Base. /(x::VEmpty{S}, a::S) where {S} = VEmpty{S}()

Base. +(x::VEmpty{S}, y::VEmpty{S}) where {S} = VEmpty{S}()
Base. -(x::VEmpty{S}, y::VEmpty{S}) where {S} = VEmpty{S}()

LinearAlgebra. ⋅(x::VEmpty{S}, y::VEmpty{S}) where {S} = S(0)

incomplete_norm(x::VEmpty{S}) where {S} = S(0)
LinearAlgebra.norm(x::VEmpty{S}) where {S} = S(0)

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

export VComp
struct VComp{V, S} <: VectorSpace{S}
    v::V
    function VComp{V, S}(v::V) where {V, S}
        @assert S === stype(stype(V))
        new{V, S}(v)
    end
end

function VComp(v::V) where {V <: VectorSpace}
    S = stype(stype(V))
    VComp{V, S}(v)
end

Base.show(io::IO, x::VComp) = print(io, "($(x.v))")

function Base.firstindex(x::VComp)
    i0 = firstindex(x.v)
    li = length(x.v)
    li == 0 && return (i0, 1)
    j0 = firstindex(x.v[i0])
    (i0, j0)
end
function Base.lastindex(x::VComp)
    i1 = lastindex(x.v)
    li = length(x.v)
    li == 0 && return (i1, 0)
    j1 = lastindex(x.v[i1])
    (i1, j1)
end
function Base.getindex(x::VComp, ij)
    i = ij[1]
    j = ij[2]
    x.v[i][j]
end
Base.eltype(::Type{VComp{V, S}}) where {V, S} = S
function Base.length(x::VComp)
    li = length(x.v)
    li == 0 && return 0
    lj = length(x.v[end])
    li * lj
end
function Base.iterate(x::VComp)
    iti = iterate(x.v)
    iti === nothing && return nothing
    eli, sti = iti
    itj = iterate(eli)
    while itj === nothing
        iti = iterate(x.v, sti)
        iti === nothing && return nothing
        eli, sti = iti
        itj = iterate(eli)
    end
    elj, stj = itj
    elj, (iti, stj)
end
function Base.iterate(x::VComp, st)
    iti, stj = st
    eli, sti = iti
    itj = iterate(eli, stj)
    while itj === nothing
        iti = iterate(x.v, sti)
        iti === nothing && return nothing
        eli, sti = iti
        itj = iterate(eli)
    end
    elj, stj = itj
    elj, (iti, stj)
end
mapmap(f, x) = map(a -> map(f, a), x)
mapmap(f, x, y) = map((a, b) -> map(f, a, b), x, y)
Base.map(f, x::VComp) = VComp(mapmap(f, x.v))
Base.map(f, x::VComp{V}, y::VComp{V}) where {V} = VComp(mapmap(f, x.v, y.v))

otype(::Type{VComp{V, S}}) where {V, S} = retype(V, S)
itype(::Type{VComp{V, S}}) where {V, S} = stype(V)
stype(::Type{VComp{V, S}}) where {V, S} = S
function retype(::Type{VComp{V, S}}, ::Type{T}) where {V, S, T}
    VT = retype(V, retype(stype(V), T))
    VComp{VT, T}
end

Base.zero(::Type{VComp{V, S}}) where {V, S <: Union{Number, VectorSpace}} =
    VComp{V, S}(zero(V))

Base. +(x::VComp{V, S}) where {V, S} = VComp{V, S}(mapmap(+, x.v))
Base. -(x::VComp{V, S}) where {V, S} = VComp{V, S}(mapmap(-, x.v))
Base. *(a::S, x::VComp{V, S}) where {V, S} = VComp{V, S}(mapmap(c->a*c, x.v))
Base. \(a::S, x::VComp{V, S}) where {V, S <: Union{AbstractFloat, Rational}} =
    VComp{V, S}(mapmap(c->a\c, x.v))
Base. *(x::VComp{V, S}, a::S) where {V, S} = VComp{V, S}(mapmap(c->c*a, x.v))
Base. /(x::VComp{V, S}, a::S) where {V, S <: Union{AbstractFloat, Rational}} =
    VComp{V, S}(mapmap(c->c/a, x.v))

Base. +(x::VComp{V, S}, y::VComp{V, S}) where {V, S} =
    VComp{V, S}(mapmap(+, x.v, y.v))
Base. -(x::VComp{V, S}, y::VComp{V, S}) where {V, S} =
    VComp{V, S}(mapmap(-, x.v, y.v))

LinearAlgebra. ⋅(x::VComp{V, S}, y::VComp{V, S}) where {V, S} =
    sum(map(⋅, x.v, y.v))::S

incomplete_norm(x::VComp{V, S}) where {V, S <: Number} =
    sum(map(incomplete_norm, x.v))::S
LinearAlgebra.norm(x::VComp{V, S}) where {V, S <: AbstractFloat} =
    sqrt(incomplete_norm(x))::S

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

export VUnit
struct VUnit{S} <: VectorSpace{S}
    elt::S
end

# VUnit(elt::S) where {S} = VUnit{S}(elt)

Base.show(io::IO, x::VUnit) = print(io, x.elt)

Base.firstindex(x::VUnit) = 1
Base.lastindex(x::VUnit) = 1
Base.getindex(x::VUnit, i) = x.elt
# Base.IteratorSize(::Type{VUnit{S}}) where {S} = HasLength()
# Base.IteratorEltype(::Type{VUnit{S}}) where {S} = HasEltype()
Base.eltype(::Type{VUnit{S}}) where {S} = S
# Base.size(x::VUnit) = (length(x), )
Base.length(x::VUnit) = 1
Base.iterate(x::VUnit) = (x.elt, ())
Base.iterate(x::VUnit, st) = nothing
function Base.map(f, x::VUnit{S}) where {S}
    r = f(x.elt)
    R = typeof(r)
    VUnit{R}(r)
end
function Base.map(f, x::VUnit{S}, y::VUnit{S}) where {S}
    r = f(x.elt, y.elt)
    R = typeof(r)
    VUnit{R}(r)
end

stype(::Type{VUnit{S}}) where {S} = S
retype(::Type{VUnit{S}}, ::Type{T}) where {S, T} = VUnit{T}

Base.zero(::Type{VUnit{S}}) where {S} = VUnit{S}(zero(S))

Base. +(x::VUnit{S}) where {S<:Number} = VUnit{S}(+x.elt)
Base. -(x::VUnit{S}) where {S<:Number} = VUnit{S}(-x.elt)
Base. *(a::S, x::VUnit{S}) where {S<:Number} = VUnit{S}(a * x.elt)
Base. \(a::S, x::VUnit{S}) where
        {S <: Union{AbstractFloat, Rational, VectorSpace}} =
    VUnit{S}(a \ x.elt)
Base. *(x::VUnit{S}, a::S) where {S<:Number} = VUnit{S}(x.elt * a)
Base. /(x::VUnit{S}, a::S) where
        {S <: Union{AbstractFloat, Rational, VectorSpace}} =
    VUnit{S}(x.elt / a)

Base. +(x::VUnit{S}, y::VUnit{S}) where {S<:Number} =
    VUnit{S}(x.elt + y.elt)
Base. -(x::VUnit{S}, y::VUnit{S}) where {S<:Number} =
    VUnit{S}(x.elt - y.elt)

LinearAlgebra. ⋅(x::VUnit{S}, y::VUnit{S}) where {S<:Number} =
    (x.elt * y.elt)::S

incomplete_norm(x::VUnit{S}) where {S<:Number} = (abs(x.elt) ^ 2)::S
LinearAlgebra.norm(x::VUnit{S}) where {S<:AbstractFloat} =
    sqrt(incomplete_norm(x))::S

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

# See the Cayley–Dickson construction
# <https://en.wikipedia.org/wiki/Cayley–Dickson_construction>
export VComplex
struct VComplex{S} <: VectorSpace{S}
    re::S
    im::S
end

function Base.show(io::IO, x::VComplex{S}) where {S}
    print(io, "($(x.re),$(x.im))")
end

Base.firstindex(x::VComplex) = firstindex((x.re, x.im))
Base.lastindex(x::VComplex) = lastindex((x.re, x.im))
Base.getindex(x::VComplex, i) = getindex((x.re, x.im), i)
Base.eltype(::Type{VComplex{S}}) where {S} = S
Base.length(x::VComplex) = length((x.re, x.im))
Base.iterate(x::VComplex) = iterate((x.re, x.im))
Base.iterate(x::VComplex, st) = iterate((x.re, x.im), st)
function Base.map(f, x::VComplex)
    res = map(f, (x.re, x.im))
    R = eltype(res)
    VComplex{R}(res[1], res[2])
end
function Base.map(f, x::VComplex, y::VComplex)
    res = map(f, (x.re, x.im), (y.re, y.im))
    R = eltype(res)
    VComplex{R}(res[1], res[2])
end

function Base.isequal(x::VComplex{S}, y::VComplex{S}) where {S}
    isequal(x.re, y.re) && isequal(x.im, y.im)
end
function Base. ==(x::VComplex{S}, y::VComplex{S}) where {S}
    x.re == y.re && x.im == y.im
end

function mkcomplex(::Type{T}, arb::Iterators.Stateful) where {T}
    re = popfirst!(arb)::T
    im = popfirst!(arb)::T
    VComplex{T}(re, im)
end
function Arbitrary.arbitrary(::Type{VComplex{T}}, ast::ArbState) where {T}
    iter = Iterators.Stateful(arbitrary(T, ast))
    Generate{VComplex{T}}(() -> mkcomplex(T, iter))
end

stype(::Type{VComplex{S}}) where {S} = S
retype(::Type{VComplex{S}}, ::Type{T}) where {S, T} = VComplex{T}

function Base.zero(::Type{VComplex{S}}) where {S}
    VComplex{S}(zero(S), zero(S))
end
function Base.one(::Type{VComplex{S}}) where {S}
    VComplex{S}(one(S), zero(S))
end

function Base. +(x::VComplex{S}) where {S}
    VComplex{S}(+x.re, +x.im)
end
function Base. -(x::VComplex{S}) where {S}
    VComplex{S}(-x.re, -x.im)
end
function Base.conj(x::VComplex{S}) where {S}
    VComplex{S}(conj(x.re), -x.im)
end
function Base.abs2(x::VComplex{S}) where {S}
    (conj(x) * x).re::S
end
function Base.abs(x::VComplex{S}) where {S}
    sqrt(abs2(x))::S
end
function Base.inv(x::VComplex{S}) where {S}
    conj(x) / abs2(x)
end

function Base. +(x::VComplex{S}, y::VComplex{S}) where {S}
    VComplex{S}(x.re + y.re, x.im + y.im)
end
function Base. -(x::VComplex{S}, y::VComplex{S}) where {S}
    VComplex{S}(x.re - y.re, x.im - y.im)
end

function Base. *(a::S, y::VComplex{S}) where {S}
    VComplex{S}(a * y.re, a * y.im)
end
function Base. *(x::VComplex{S}, b::S) where {S}
    VComplex{S}(x.re * b, x.im * b)
end
function Base. \(a::S, y::VComplex{S}) where {S}
    VComplex{S}(a \ y.re, a \ y.im)
end
function Base. /(x::VComplex{S}, b::S) where {S}
    VComplex{S}(x.re / b, x.im / b)
end

# Note: Do not change the order of arguments, as multiplication might
# not be commutative
function Base. *(x::VComplex{S}, y::VComplex{S}) where {S}
    VComplex{S}(x.re * y.re - conj(y.im) * x.im,
                y.im * x.re + x.im * conj(y.re))
end
function Base. /(x::VComplex{S}, y::VComplex{S}) where {S}
    x * inv(y)
end

LinearAlgebra. ⋅(x::VComplex{S}, y::VComplex{S}) where
        {D, S <: Union{Number, VectorSpace}} =
    (conj(x.re) * y.re + conj(x.im) * y.im)::S

incomplete_norm(x::VComplex{S}) where {S <: Union{Number, VectorSpace}} =
    (abs2(x.re) + abs2(x.im))::S
LinearAlgebra.norm(x::VComplex{S}) where
        {S <: Union{AbstractFloat, VectorSpace}} =
    sqrt(incomplete_norm(x))::S

export VQuaternion
struct VQuaternion{S} <: VectorSpace{S}
    val::VComp{VComplex{VComplex{S}}}
end

function Base.show(io::IO, x::VQuaternion{S}) where {S}
    print(io, x.val)
end

Base.firstindex(x::VQuaternion) = firstindex(x.val)
Base.lastindex(x::VQuaternion) = lastindex(x.val)
Base.getindex(x::VQuaternion, i) = getindex(x.val, i)
Base.eltype(::Type{VQuaternion{S}}) where {S} = S
Base.length(x::VQuaternion) = length(x.val)
Base.iterate(x::VQuaternion) = iterate(x.val)
Base.iterate(x::VQuaternion, st) = iterate(x.val, st)
Base.map(f, x::VQuaternion) = VQuaternion(map(f, x.val))
Base.map(f, x::VQuaternion, y::VQuaternion) = VQuaternion(map(f, x.val, y.val))

Base.isequal(x::VQuaternion, y::VQuaternion) = isequal(x.val.v, y.val.v)
Base. ==(x::VQuaternion, y::VQuaternion) = x.val.v == y.val.v

function Arbitrary.arbitrary(::Type{VQuaternion{T}}, ast::ArbState) where {T}
    imap(x -> VQuaternion(VComp(x)), arbitrary(VComplex{VComplex{T}}, ast))
end

stype(::Type{VQuaternion{S}}) where {S} = S
retype(::Type{VQuaternion{S}}, ::Type{T}) where {S, T} = VQuaternion{T}

function Base.zero(::Type{VQuaternion{S}}) where {S}
    VQuaternion(VComp(zero(VComplex{VComplex{S}})))
end
function Base.one(::Type{VQuaternion{S}}) where {S}
    VQuaternion(VComp(one(VComplex{VComplex{S}})))
end

function Base. +(x::VQuaternion{S}) where {S}
    VQuaternion(+x.val)
end
function Base. -(x::VQuaternion{S}) where {S}
    VQuaternion(-x.val)
end
function Base.conj(x::VQuaternion{S}) where {S}
    VQuaternion(VComp(conj(x.val.v)))
end
function Base.abs2(x::VQuaternion{S}) where {S}
    abs2(x.val.v)::S
end
function Base.abs(x::VQuaternion{S}) where {S}
    abs(x.val.v)::S
end
function Base.inv(x::VQuaternion{S}) where {S}
    VQuaternion(VComp(inv(x.val.v)))
end

function Base. +(x::VQuaternion{S}, y::VQuaternion{S}) where {S}
    VQuaternion(x.val + y.val)
end
function Base. -(x::VQuaternion{S}, y::VQuaternion{S}) where {S}
    VQuaternion(x.val - y.val)
end

function Base. *(a::S, y::VQuaternion{S}) where {S}
    VQuaternion(a * y.val)
end
function Base. *(x::VQuaternion{S}, b::S) where {S}
    VQuaternion(x.val * b)
end
function Base. \(a::S, y::VQuaternion{S}) where {S}
    VQuaternion(a \ y.val)
end
function Base. /(x::VQuaternion{S}, b::S) where {S}
    VQuaternion(x.val / b)
end

function Base. *(x::VQuaternion{S}, y::VQuaternion{S}) where {S}
    VQuaternion(VComp(x.val.v * y.val.v))
end
function Base. /(x::VQuaternion{S}, y::VQuaternion{S}) where {S}
    VQuaternion(VComp(x.val.v / y.val.v))
end

LinearAlgebra. ⋅(x::VQuaternion{S}, y::VQuaternion{S}) where
        {D, S <: Union{Number, VectorSpace}} =
    (x.val ⋅ y.val)::S

incomplete_norm(x::VQuaternion{S}) where {S <: Union{Number, VectorSpace}} =
    incomplete_norm(x.val)::S
LinearAlgebra.norm(x::VQuaternion{S}) where
        {S <: Union{AbstractFloat, VectorSpace}} =
    norm(x.val)::S

export VOctonion
struct VOctonion{S} <: VectorSpace{S}
    val::VComp{VComplex{VQuaternion{S}}}
end

function Base.show(io::IO, x::VOctonion{S}) where {S}
    print(io, x.val)
end

Base.firstindex(x::VOctonion) = firstindex(x.val)
Base.lastindex(x::VOctonion) = lastindex(x.val)
Base.getindex(x::VOctonion, i) = getindex(x.val, i)
Base.eltype(::Type{VOctonion{S}}) where {S} = S
Base.length(x::VOctonion) = length(x.val)
Base.iterate(x::VOctonion) = iterate(x.val)
Base.iterate(x::VOctonion, st) = iterate(x.val, st)
Base.map(f, x::VOctonion) = VOctonion(map(f, x.val))
Base.map(f, x::VOctonion, y::VOctonion) = VOctonion(map(f, x.val, y.val))

Base.isequal(x::VOctonion, y::VOctonion) = isequal(x.val.v, y.val.v)
Base. ==(x::VOctonion, y::VOctonion) = x.val.v == y.val.v

function Arbitrary.arbitrary(::Type{VOctonion{T}}, ast::ArbState) where {T}
    imap(x -> VOctonion(VComp(x)), arbitrary(VComplex{VQuaternion{T}}, ast))
end

stype(::Type{VOctonion{S}}) where {S} = S
retype(::Type{VOctonion{S}}, ::Type{T}) where {S, T} = VOctonion{T}

function Base.zero(::Type{VOctonion{S}}) where {S}
    VOctonion(VComp(zero(VComplex{VQuaternion{S}})))
end
function Base.one(::Type{VOctonion{S}}) where {S}
    VOctonion(VComp(one(VComplex{VQuaternion{S}})))
end

function Base. +(x::VOctonion{S}) where {S}
    VOctonion(+x.val)
end
function Base. -(x::VOctonion{S}) where {S}
    VOctonion(-x.val)
end
function Base.conj(x::VOctonion{S}) where {S}
    VOctonion(VComp(conj(x.val.v)))
end
function Base.abs2(x::VOctonion{S}) where {S}
    abs2(x.val.v)::S
end
function Base.abs(x::VOctonion{S}) where {S}
    abs(x.val.v)::S
end
function Base.inv(x::VOctonion{S}) where {S}
    VOctonion(VComp(inv(x.val.v)))
end

function Base. +(x::VOctonion{S}, y::VOctonion{S}) where {S}
    VOctonion(x.val + y.val)
end
function Base. -(x::VOctonion{S}, y::VOctonion{S}) where {S}
    VOctonion(x.val - y.val)
end

function Base. *(a::S, y::VOctonion{S}) where {S}
    VOctonion(a * y.val)
end
function Base. *(x::VOctonion{S}, b::S) where {S}
    VOctonion(x.val * b)
end
function Base. \(a::S, y::VOctonion{S}) where {S}
    VOctonion(a \ y.val)
end
function Base. /(x::VOctonion{S}, b::S) where {S}
    VOctonion(x.val / b)
end

function Base. *(x::VOctonion{S}, y::VOctonion{S}) where {S}
    VOctonion(VComp(x.val.v * y.val.v))
end
function Base. /(x::VOctonion{S}, y::VOctonion{S}) where {S}
    VOctonion(VComp(x.val.v / y.val.v))
end

LinearAlgebra. ⋅(x::VOctonion{S}, y::VOctonion{S}) where
        {D, S <: Union{Number, VectorSpace}} =
    (x.val ⋅ y.val)::S

incomplete_norm(x::VOctonion{S}) where {S <: Union{Number, VectorSpace}} =
    incomplete_norm(x.val)::S
LinearAlgebra.norm(x::VOctonion{S}) where
        {S <: Union{AbstractFloat, VectorSpace}} =
    norm(x.val)::S

export VSedenion
struct VSedenion{S} <: VectorSpace{S}
    val::VComp{VComplex{VOctonion{S}}}
end

function Base.show(io::IO, x::VSedenion{S}) where {S}
    print(io, x.val)
end

Base.firstindex(x::VSedenion) = firstindex(x.val)
Base.lastindex(x::VSedenion) = lastindex(x.val)
Base.getindex(x::VSedenion, i) = getindex(x.val, i)
Base.eltype(::Type{VSedenion{S}}) where {S} = S
Base.length(x::VSedenion) = length(x.val)
Base.iterate(x::VSedenion) = iterate(x.val)
Base.iterate(x::VSedenion, st) = iterate(x.val, st)
Base.map(f, x::VSedenion) = VSedenion(map(f, x.val))
Base.map(f, x::VSedenion, y::VSedenion) = VSedenion(map(f, x.val, y.val))

Base.isequal(x::VSedenion, y::VSedenion) = isequal(x.val.v, y.val.v)
Base. ==(x::VSedenion, y::VSedenion) = x.val.v == y.val.v

function Arbitrary.arbitrary(::Type{VSedenion{T}}, ast::ArbState) where {T}
    imap(x -> VSedenion(VComp(x)), arbitrary(VComplex{VOctonion{T}}, ast))
end

stype(::Type{VSedenion{S}}) where {S} = S
retype(::Type{VSedenion{S}}, ::Type{T}) where {S, T} = VSedenion{T}

function Base.zero(::Type{VSedenion{S}}) where {S}
    VSedenion(VComp(zero(VComplex{VOctonion{S}})))
end
function Base.one(::Type{VSedenion{S}}) where {S}
    VSedenion(VComp(one(VComplex{VOctonion{S}})))
end

function Base. +(x::VSedenion{S}) where {S}
    VSedenion(+x.val)
end
function Base. -(x::VSedenion{S}) where {S}
    VSedenion(-x.val)
end
function Base.conj(x::VSedenion{S}) where {S}
    VSedenion(VComp(conj(x.val.v)))
end
function Base.abs2(x::VSedenion{S}) where {S}
    abs2(x.val.v)::S
end
function Base.abs(x::VSedenion{S}) where {S}
    abs(x.val.v)::S
end
function Base.inv(x::VSedenion{S}) where {S}
    VSedenion(VComp(inv(x.val.v)))
end

function Base. +(x::VSedenion{S}, y::VSedenion{S}) where {S}
    VSedenion(x.val + y.val)
end
function Base. -(x::VSedenion{S}, y::VSedenion{S}) where {S}
    VSedenion(x.val - y.val)
end

function Base. *(a::S, y::VSedenion{S}) where {S}
    VSedenion(a * y.val)
end
function Base. *(x::VSedenion{S}, b::S) where {S}
    VSedenion(x.val * b)
end
function Base. \(a::S, y::VSedenion{S}) where {S}
    VSedenion(a \ y.val)
end
function Base. /(x::VSedenion{S}, b::S) where {S}
    VSedenion(x.val / b)
end

function Base. *(x::VSedenion{S}, y::VSedenion{S}) where {S}
    VSedenion(VComp(x.val.v * y.val.v))
end
function Base. /(x::VSedenion{S}, y::VSedenion{S}) where {S}
    VSedenion(VComp(x.val.v / y.val.v))
end

LinearAlgebra. ⋅(x::VSedenion{S}, y::VSedenion{S}) where
        {D, S <: Union{Number, VectorSpace}} =
    (x.val ⋅ y.val)::S

incomplete_norm(x::VSedenion{S}) where {S <: Union{Number, VectorSpace}} =
    incomplete_norm(x.val)::S
LinearAlgebra.norm(x::VSedenion{S}) where
        {S <: Union{AbstractFloat, VectorSpace}} =
    norm(x.val)::S

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

# dim   rank   elems

#   0      0       1

#   1      0       1
#   1      1       1

#   2      0       1
#   2      1       2
#   2      2       1

#   3      0       1
#   3      1       3
#   3      2       3
#   3      3       1

#   4      0       1
#   4      1       4
#   4      2       6
#   4      3       4
#   4      4       1

#   5      0       1
#   5      1       5
#   5      2      10
#   5      3      10
#   5      4       5
#   5      5       1

#   6      0       1
#   6      1       6
#   6      2      15
#   6      3      20
#   6      4      15
#   6      5       6
#   6      6       1

nelems(d::Int, n::Int) = binomial(d, n)

export Form
struct Form{D, N, S, L} <: VectorSpace{S}
    elts::Vec{L, S}

    function Form{D, N, S, L}() where {D, N, S, L}
        @assert L == nelems(D, N)
        new{D, N, S, L}()
    end
    function Form{D, N, S}() where {D, N, S}
        Form{D, N, S, nelems(D, N)}()
    end
    function Form{D, N}() where {D, N}
        Form{D, N, Missing}()
    end

    function Form{D, N, S, L}(elts::Vec{L, S}) where {D, N, S, L}
        @assert L == nelems(D, N)
        new{D, N, S, L}(elts)
    end
    function Form{D, N, S}(elts::Vec{L, S}) where {D, N, S, L}
        Form{D, N, S, L}(elts)
    end
    function Form{D, N}(elts::Vec{L, S}) where {D, N, S, L}
        Form{D, N, S, L}(elts)
    end
end

Base.firstindex(x::Form) = firstindex(x.elts)
Base.lastindex(x::Form) = lastindex(x.elts)
Base.getindex(x::Form, i) = x.elts[i]
# Base.IteratorSize(::Type{Form{D, S}}) where {D, S} = HasLength()
# Base.IteratorEltype(::Type{Form{D, S}}) where {D, S} = HasEltype()
Base.eltype(::Type{Form{D, N, S}}) where {D, N, S} = S
# Base.size(x::Form) = (length(x), )
Base.length(x::Form) = length(x.elts)
Base.iterate(x::Form) = iterate(x.elts)
Base.iterate(x::Form, st) = iterate(x.elts, st)
function Base.map(f, x::Form{D, N}) where {D, N}
    relts = map(f, x.elts)
    R = eltype(relts)
    Form{D, N, R}(relts)
end
function Base.map(f, x::Form{D, N}, y::Form{D, N}) where {D, N}
    relts = map(f, x.elts, y.elts)
    Form{D, N, eltype(relts)}(relts)
end

Base.isequal(x::Form{D, N}, y::Form{D, N}) where {D, N} =
    isequal(x.elts, y.elts)
Base. ==(x::Form{D, N}, y::Form{D, N}) where {D, N} = x.elts == y.elts

function Arbitrary.arbitrary(::Type{Form{D, N, S, L}}, ast::ArbState) where {
        D, N, S, L}
    ET = Vec{L, S}
    VT = Form{D, N, S, L}
    imap(x -> VT(x), arbitrary(ET, ast))
end
function Arbitrary.arbitrary(::Type{<: Form{D, N, S}}, ast::ArbState) where {
        D, N, S}
    arbitrary(Form{D, N, S, nelems(D, N)})
end

stype(::Type{<: Form{D, N, S}}) where {D, N, S} = S
retype(::Type{Form{D, N, S, L}}, ::Type{T}) where {D, N, S, L, T} =
    Form{D, N, T, L}

# This needs to accept either a Number or a VectorSpace
Base.zero(::Type{<: Form{D, N, S}}) where {
        D, N, S <: Union{Number, VectorSpace}} =
    Form{D, N, S}(zero(Vec{nelems(D, N), S}))
Base. +(x::Form{D, N, S}) where {D, N, S <: Union{Number, VectorSpace}} =
    Form{D, N, S}(+ x.elts)
Base. -(x::Form{D, N, S}) where {D, N, S <: Union{Number, VectorSpace}} =
    Form{D, N, S}(- x.elts)
Base. *(a::S, x::Form{D, N, S}) where {D, N, S <: Union{Number, VectorSpace}} =
    Form{D, N, S}(a * x.elts)
Base. \(a::S, x::Form{D, N, S}) where
        {D, N, S <: Union{AbstractFloat, Rational, VectorSpace}} =
    Form{D, N, S}(a \ x.elts)
Base. *(x::Form{D, N, S}, a::S) where {D, N, S <: Union{Number, VectorSpace}} =
    Form{D, N, S}(x.elts * a)
Base. /(x::Form{D, N, S}, a::S) where
        {D, N, S <: Union{AbstractFloat, Rational, VectorSpace}} =
    Form{D, N, S}(x.elts / a)

Base. +(x::Form{D, N, S}, y::Form{D, N, S}) where {
        D, N, S <: Union{Number, VectorSpace}} =
    Form{D, N, S}(x.elts + y.elts)
Base. -(x::Form{D, N, S}, y::Form{D, N, S}) where {
        D, N, S <: Union{Number, VectorSpace}} =
    Form{D, N, S}(x.elts - y.elts)

Base.one(::Type{<: Form{D, N, S}}) where {
        D, N, S <: Union{Number, VectorSpace}} =
    Form{D, N, S}(zero(Vec{nelems(D, N), S}))

# multi-indices range 0:D, linear indices range 1:L
const linear2multi = Dict{Tuple{Int, Int}, Any}()
const multi2linear = Dict{Tuple{Int, Int}, Any}()

function l2m(ms::Vector{NTuple{N, Int}}, l::Int)::NTuple{N, Int} where {N}
    ms[l]
end

function m2l(ls::Array{Int, N}, m::NTuple{N, Int})::Tuple{Int, Int} where {N}
    sl = ls[CartesianIndex(map(i->i+1, m))]
    sl == 0 && return 0, 0
    return sign(sl), abs(sl)
end

function makel2m(::Val{D}, ::Val{N}) where {D, N}
    multi = Vector{NTuple{N, Int}}()
    for ind in CartesianIndices(ntuple(_->D, N))
        m = map(i->i-1, ind.I)::NTuple{N, Int}
        ismono = true
        for n in 2:N
            ismono &= m[n] > m[n-1]
        end
        if ismono
            push!(multi, m)
        end
    end
    @assert length(multi) == nelems(D, N)
    multi
end

function makem2l(::Val{D}, ::Val{N}) where {D, N}
    ms = linear2multi[(D, N)]::Vector{NTuple{N, Int}}
    linear = zeros(Int, ntuple(_->D, N))
    for (l, m) in enumerate(ms)
        for p in permutations(collect(1:N))
            s = levicivita(p)
            m′ = m[p]::NTuple{N, Int}
            @assert abs(s) == 1
            linear[CartesianIndex(map(i->i+1, m′))] = s * l
        end
    end
    linear
end

function checklm(::Val{D}, ::Val{N}) where {D, N}
    ms = linear2multi[(D, N)]::Vector{NTuple{N, Int}}
    ls = multi2linear[(D, N)]::Array{Int, N}
    L = length(ms)
    for l in 1:L
        m = l2m(ms, l)
        @assert all(i >= 0 && i <= D for i in m)
        s′, l′ = m2l(ls, m)
        @assert s′ == 1
        @assert l′ == l
    end
    for ind in CartesianIndices(ls)
        m = map(i->i-1, ind.I)::NTuple{N, Int}
        s, l = m2l(ls, m)
        if s == 0
            @assert l == 0
        else
            @assert l >= 1 && l <= L
            m′ = l2m(ms, l)
            @assert [m′...] == sort([m...])
        end
    end
end

for D in 0:6, N in 0:D
    linear2multi[(D, N)] = makel2m(Val(D), Val(N))
    multi2linear[(D, N)] = makem2l(Val(D), Val(N))
    checklm(Val(D), Val(N))
end

export formindex
@generated function formindex(x::Form{D, N, S}, is::NTuple{N, Int})::S where {
        D, N, S}
    ls = multi2linear[(D, N)]::Array{Int, N}
    quote
        $(Expr(:meta, :inline))
        s, l = m2l($ls, is)
        s == 0 && return zero(S)
        y = @inbounds x.elts[l]
        s >= 0 ? y : -y
    end
end

# Hodge ∗ ast
# inner product

export ∧
# TODO: Number is not enough -- need Union{AbstractFloat, Rational}.
# What about Complex? Matrix?
@generated function ∧(x::Form{D, N1, S}, y::Form{D, N2, S}) where {
        D, N1, N2, S <: Union{Number, Form}}
    N = N1 + N2
    @assert N <= D
    L = nelems(D, N)
    ms = linear2multi[(D, N)]::Vector{NTuple{N, Int}}
    ls1 = multi2linear[(D, N1)]::Array{Int, N1}
    ls2 = multi2linear[(D, N2)]::Array{Int, N2}
    elts = Expr[]
    for l in 1:L
        m = l2m(ms, l)
        expr = :0
        ps = permutations(collect(1:N))
        np = length(ps)
        for p in ps
            s = levicivita(p)
            m′ = m[p]::NTuple{N, Int}
            m1′ = m′[1:N1]::NTuple{N1, Int}
            m2′ = m′[N1+1:N]::NTuple{N2, Int}
            s1′, l1′ = m2l(ls1, m1′)
            s2′, l2′ = m2l(ls2, m2′)
            @assert abs(s1′) == abs(s2′) == 1
            op = s * s1′ * s2′ > 0 ? :+ : :-
            expr = :($op($expr, x.elts[$l1′] * y.elts[$l2′]))
        end
        push!(elts, :($expr / S($np)))
    end
    quote
        Form{D, $N, S}(Vec(tuple($(elts...))))
    end
end

end


export fit_statespace_dp

# Piecewise constant segmentation ==============================================
@inline function argmin_const(y,w,t1,t2)
    # y'w/sum(w) weighted mean of y
    s = 0.
    dp = 0.
    for t = t1:t2
        dp += y[t]*w[t]
        s += w[t]
    end
    dp/s
end
@inline function cost_const(y,w,t1,t2,argminfun)
    # w'(y-dpamin)^2 weigted variance of y
    dp = 0.
    dpm = argminfun(y,w,t1,t2)
    for t = t1:t2
        dp += w[t]*(y[t]-dpm)^2
    end
    dp
end

# Piecewise linear segmentation ============================================
@inline function argmin_lin(input,w,t1,t2)
    y   = input[1]
    A   = input[2]
    yminus = A[t1:t2,1] .- y[t1]
    yreg = y[t1:t2] .- yminus
    ones(t2-t1+1,1)\yreg # TODO: does not take care of w yet
end

# inuti costfun måste hela y användas, dvs simulation i stället för prediction. Annars hamnar alla knutpunkter på kurvan.

@inline function cost_lin(input,w,t1,t2,argminfun)
    if t2 == t1
        return Inf
    end
    y   = input[1]
    A   = input[2]
    n   = size(y,1) ÷ length(w)
    k   = argminfun(input,w,t1,t2)
    yminus = A[t1:t2,1] .- y[t1]

    e   = y[t1:t2] .- yminus .- k
    return e⋅e
end

# Piecewise statespace segmentation ============================================
@inline function argmin_ss(input,w,t1,t2)
    y   = input[1]
    A   = input[2]
    n   = size(y,1) ÷ length(w)
    ii  = (t1-1)*n+1
    ii2 = t2*n-1
    A[ii:ii2,:]\y[ii:ii2] # TODO: does not take care of w yet
end
@inline function cost_ss(input,w,t1,t2,argminfun)
    y   = input[1]
    A   = input[2]
    n   = size(y,1) ÷ length(w)
    ii  = (t1-1)*n+1
    ii2 = t2*n-1
    k   = argminfun(input,w,t1,t2)
    e   = y[ii:ii2]-A[ii:ii2,:]*k
    return e⋅e
end

using Base.Threads
function seg_bellman(y,M,w, costfun=cost_const, argminfun=argmin_const; doplot=false)
    doplot = @static isinteractive() && doplot
    n = length(w)
    @assert (M < n) "M must be smaller than the length of the data sequence"
    B = zeros(Int,M-1,n) # back-pointer matrix

    # initialize Bellman iteration
    fi = zeros(n)
    for jj = M:n
        fi[jj] = costfun(y,w,jj,n,argminfun)
    end
    doplot && plot(fi, lab="fi init")
    memorymat = fill(Inf,n,n)

    # iterate Bellman iteration
    fnext = Vector{Float64}(undef,n)
    for jj = M-1:-1:1
        for kk = jj:n-(M-jj)
            opt   = Inf
            optll = 0
            for ll = kk+1:n-(M-jj-1)
                if memorymat[kk,ll-1] == Inf
                    memorymat[kk,ll-1] = costfun(y,w,kk,ll-1,argminfun)
                end
                cost = memorymat[kk,ll-1] + fi[ll]
                if cost < opt
                    opt   = cost
                    optll = ll
                end
            end
            fnext[kk] = opt
            B[jj,kk]  = optll-1
        end
        fi .= fnext
        fnext = Vector{Float64}(undef,n)
        doplot && plot!(fi, lab="fi at $jj")
    end

    # last Bellman iterate
    V = [(memorymat[1,jj] == Inf ? costfun(y,w,1,jj,argminfun) : memorymat[1,jj])+fi[jj+1] for jj = 1:n-M]
    if isempty(V) || length(V) <= M
        error("Failed to perform segmentation, try lowering M")
    end
    doplot && plot!(V, lab="V")

    # backward pass
    t = Vector{Int}(undef,M);
    a = Vector{typeof(argminfun(y,w,1,2))}(undef,M+1);
    _,t[1] = findmin(V[1:end-M])
    a[1] = argminfun(y,w,1,t[1])

    for jj = 2:M
        t[jj] = B[jj-1,t[jj-1]]
        a[jj] = argminfun(y,w,t[jj-1]+1,t[jj])
    end
    a[M+1] = argminfun(y,w,t[M]+1,n)

    V,t,a
end

"""
model = fit_statespace_dp(x,u, M; extend=true)

Fit model using dynamic programming. `M` is the number of changepoints.
´x` and `u` should have time in the second dimension.

See Bagge Carlson et al. 2018 section IID
Usage example in function `benchmark_ss`
"""
function fit_statespace_dp(x,u, M; extend=true, kwargs...)
    n,T = size(x)
    m = size(u,1)
    d = iddata(x[:,2:end],u, x[:,1:end-1])
    input = matrices(SimpleLTVModel(d), d)
    @show size.(input)
    # input = matrices(x,u)
    V,bps,a = seg_bellman(input,M, ones(T-1), cost_ss, argmin_ss, doplot=false)
    At,Bt = segments2full(a,bps,n,m,T)

    SimpleLTVModel(At, Bt,extend)
end

# Tests
# Dynamic programming ==========================================================
function benchmark_const(N, M=1, doplot=false)
    doplot = @static isinteractive() && doplot
    n = 3N
    y = [0.1randn(N); 10 .+ 0.1randn(N); 20 .+ 0.1randn(N).+range(1, stop=10, length=N)]
    V,t,a = @time seg_bellman(y,M, ones(length(y)))
    if doplot
        tplot = [1;t;n];
        aplot = [a;a[end]];
        plot(y)
        plot!(tplot, aplot, l=:step);gui()
    end
    # yhat, x, a, b
end

function benchmark_lin(T_, M, doplot=false)
    doplot = @static isinteractive() && doplot

    # M        = 1
    # T_       = 400
    x = sin.(range(0, stop=2π, length=T_))
    u = ones(1,T_)
    d = iddata(x',u, x')
    input = matrices(SimpleLTVModel(d), d)
    @time V,t,a = seg_bellman(input,M, ones(T_-1), cost_lin, argmin_lin, doplot=false)
    # if doplot
    #     k = reduce(hcat, a)'
    #     @show size(a), size(k)
    #     At = reshape(k[:,1:n^2]',n,n,M+1)
    #     At = permutedims(At, [2,1,3])
    #     tplot = [1;t;T_]
    #     plot()
    #     for i = 1:M+1
    #         plot!(tplot[i:i+1],flatten(At)[i,:]'.*ones(2), c=[:blue :green :red :magenta], xlabel="Time index", ylabel="Model coefficients")
    #     end
    #     plot!([1,T_÷2-1], [0.95 0.1; 0 0.95][:]'.*ones(2), ylims=(-0.1,1), l=(:dash,:black, 1))
    #     plot!([T_÷2,T_], [0.5 0.05; 0 0.5][:]'.*ones(2), l=(:dash,:black, 1), grid=false)
    #     gui()
    # end
    V,t,a
end

function benchmark_ss(T_, M, doplot=false)
    doplot = @static isinteractive() && doplot

    # M        = 1
    # T_       = 400
    x,xm,u,n,m = testdata(T_)
    dn = iddata(xm,u,xm)
    input = matrices(SimpleLTVModel(dn), dn)

    @time V,t,a = seg_bellman(input,M, ones(T_-1), cost_ss, argmin_ss, doplot=false)
    if doplot
        k = reduce(hcat, a)'
        At = reshape(k[:,1:n^2]',n,n,M+1)
        Bt = reshape(k[:,n^2+1:end]',m,n,M+1)
        At = permutedims(At, [2,1,3])
        Bt = permutedims(Bt, [2,1,3])
        tplot = [1;t;T_]
        plot()
        for i = 1:M+1
            plot!(tplot[i:i+1],flatten(At)[i,:]'.*ones(2), c=[:blue :green :red :magenta], xlabel="Time index", ylabel="Model coefficients")
        end
        plot!([1,T_÷2-1], [0.95 0.1; 0 0.95][:]'.*ones(2), ylims=(-0.1,1), l=(:dash,:black, 1))
        plot!([T_÷2,T_], [0.5 0.05; 0 0.5][:]'.*ones(2), l=(:dash,:black, 1), grid=false)
        gui()
    end
    V,t,a
end
# T_vec = round(Int,logspace(2, 4,10))
# times = map(T_vec) do T
#     println(T)
#     benchmark_ss(T, 1)
# end
# plot(T_vec, times)
#
# M_vec = [2,4,8,16,32]
# times = map(M_vec) do M
#     println(M)
#     @elapsed benchmark_ss(110, M)
# end
# plot(M_vec, times)


@testset "Comparison of monomials" begin
    Mod.@ncpolyvar x y
    @test x*y != y*x
    @test x^2 == x*x
end


"""

ДАНО: Робот - в произвольной клетке ограниченного прямоугольной рамкой поля без внутренних перегородок и маркеров.

РЕЗУЛЬТАТ: Робот - в исходном положении в центре косого креста (в форме X) из маркеров.

(Подсказка: решение будет подобно решению задачи 1, если направление премещения Робота задавать кортежами пары значений типа HorizonSide)

"""

using HorizonSideRobots
originRobot = Robot(animate = true, "13example.sit")
#show!(originRobot)
include("../AbstractType.jl")
include("../functions.jl")

function get_line_of_cross!(robot::Robot, sides)
    path = []
    while ! ( isborder(robot,sides[1]) || isborder(robot, sides[2]) )
        for side in sides
            if (!isborder(robot,side))
                move!(robot, side)
                push!(path,Int(side))
            else
                continue
            end
        end
        putmarker!(robot)
    end
    get_back!(robot, path)
end

function master13!(robot::Robot)::Nothing

    for sides in [ [West,Sud], [West,Nord], [Ost,Sud], [Ost,Nord] ]
        get_line_of_cross!(robot, sides)
    end

    putmarker!(robot)
end

master13!(originRobot)


using Makeitso

@target A ()->1:10
@target B ()->[-4,-3,-2,-1,0,1,2,3,4,5]
@target C (A,B)->A.+B
@target D (A,B,C)->A.+B.+C

x = (@make D)[end]
@assert x == 30

@target B ()->pi
println("--- Recipe for B modified! ---")

x = (@make D)[end]
@assert x ≈ (20+2pi)


# Solving ESPPRC
# Modified from the example given by Google OR-Tools https://developers.google.com/optimization/routing/vrp

using Routing
using Random

# This example is modifed from the VRPTW example from OR-Tools
# But ESPPRC can be used outside the VRP context
# This example describes non-VRP context

# Number of nodes
n_nodes = 17

# cost matrix:
# The objective of ESPPRC is to minimize the total cost
# Randomly generated in this example
# cost allows negative values
cost = - rand(n_nodes, n_nodes)

# travel_time matrix:
# Some elements are Inf, meaning disconnected
# travel_time is used to check the time windows constraints
travel_time = Float64[
    0 8 3 2 6 8 4 8 8 13 7 5 8 12 10 Inf 6;
    8 0 11 10 6 3 9 5 8 4 15 14 13 9 18 9 9;
    3 11 0 1 7 10 6 10 10 14 6 7 9 14 6 16 8;
    2 10 1 0 6 9 4 8 9 13 4 6 8 12 8 14 7;
    6 6 7 6 0 2 3 2 2 7 9 7 7 6 12 8 3;
    8 3 10 9 2 0 6 2 5 4 Inf 10 10 6 Inf 5 6;
    4 9 6 4 3 6 0 4 4 8 5 4 3 7 8 10 2;
    8 5 10 8 2 2 4 0 3 4 9 8 7 3 13 6 3;
    8 8 10 9 2 5 4 3 0 4 6 5 4 3 9 5 2;
    13 4 14 13 7 4 8 4 4 0 10 9 8 4 13 4 6;
    7 Inf 6 4 9 Inf 5 9 6 10 0 1 3 7 3 10 6;
    5 14 7 6 7 10 4 8 5 9 1 0 2 6 4 8 4;
    8 13 9 8 7 10 3 7 4 8 3 2 0 4 5 6 4;
    12 9 14 12 6 6 7 3 3 4 7 6 4 0 9 2 5;
    10 18 6 8 12 15 8 13 9 13 3 4 5 9 0 9 9;
    Inf 9 16 14 8 5 10 6 5 4 10 8 6 2 9 0 7;
    6 9 8 7 3 6 2 3 2 6 6 4 4 5 9 7 0
]

# load matrix
# The load information is for each arc, hence given as a matrix
# If you have load on each node, then it should be converted to a matrix
# For example, load[i, j] = load_on_node[j]
load_on_node = [1, 1, 2, 4, 2, 4, 8, 8, 1, 2, 1, 2, 4, 4, 8, 8, 0]
load = zeros(n_nodes, n_nodes)
for i in 1:n_nodes, j in 1:n_nodes
    load[i, j] = load_on_node[j]
end

# Time windows, for all nodes, including origin and destination
time_windows = [
    (0, 12),    #  1
    (0, 15),    #  2
    (16, 28),   #  3
    (10, 13),   #  4
    (0, 5),     #  5
    (5, 10),    #  6
    (0, 4),     #  7
    (5, 10),    #  8
    (0, 3),     #  9
    (10, 16),   # 10
    (10, 15),   # 11
    (0, 5),     # 12
    (5, 10),    # 13
    (7, 8),     # 14
    (10, 15),   # 15
    (11, 15),   # 16
    (0, 35),    # 17
]
# If early_time is set for the origin, it means the earliest possible departure time.
early_time = [time_windows[i][1] for i in 1:n_nodes]
late_time = [time_windows[i][2] for i in 1:n_nodes]

# vehicle capacity
capacity = 15

# service_time is just set to all zeros in this example
service_time = zeros(n_nodes)

# origin and destination
origin, destination = 6, 3

pg = ESPPRC_Instance(
    origin,
    destination,
    capacity,
    cost,
    travel_time,
    load,
    early_time,
    late_time,
    service_time
)

# Solve the ESPPRC
solution_label = solveESPPRC(pg)

@show solution_label.path
@show solution_label.cost
@show solution_label.time
@show solution_label.load


function FastConv(x, y)
    Lx = length(x); Ly = size(y, 2); problem_size = size(y, 1)

    r = Lx
    z = zeros(Number, problem_size, r)
    X = ourfft(x, r)
    for i = 1:problem_size
        Y = ourfft(y[i, :]', r)
        Z = X.*Y
        z[i, :] = ourifft(Z, r)
    end
    return z
end

function ourfft(x, n)
    s=length(x)
    x=x[:]
    if s > n
        return fft(x[1:n])
    elseif s < n
        return fft([x; zeros(n-s)])
    else
        return fft(x)
    end
end

function ourifft(x, n)
    s=length(x)
    x=x[:]
    if s > n
        return ifft(x[1:n])
    elseif s < n
        return ifft([x; zeros(n-s)])
    else
        return ifft(x)
    end
end

function rowfft(x::AbstractMatrix, n)
    result = zeros(Complex, size(x)[1], n)
    for i=1:size(x)[1]
        result[i, :] = ourfft(x[i, :], n)
    end
    return result
end


include("../src/PBCM.jl")

cell_general = cell()
cell_design_op = cylindrical_cell_designer(cell_general)
cost = cost_default()
cell_general, cost = convert_all(cell_general, cost, mult)

using DiffEqSensitivity
using PyPlot

per = 30

no_cells_yr = 10
no_cells_yr_intrvl = [no_cells_yr*(1 - (per/100)) , no_cells_yr*(1 + (per/100))]

cost_pos_AM = 18.0
cost_pos_AM_intrvl = [cost_pos_AM*(1 - (per/100)) , cost_pos_AM*(1 + (per/100))]

cost_neg_AM = 12.5
cost_neg_AM_intrvl = [cost_neg_AM*(1 - (per/100)) , cost_neg_AM*(1 + (per/100))]

cost_sep = 1.1
cost_sep_intrvl    = [cost_sep  * (1 - (per/100)) , cost_sep * (1  + (per/100))]

function gsa_cost(arr)

    cost.general_costs.no_units_mfg   = converter(arr[1], mult.units_mfg)
    cost.cell_costs.cathode.AM[1]     = arr[2]
    cost.cell_costs.anode.AM[1]       = arr[3]
    cost.cell_costs.seperator_cost[1] = arr[4]

    dollars_kWh = cost_calc(cell_general, cost, system="Cell", cost_verbosity=0)[1]

    return dollars_kWh
end

# gsa_cost([0.25, 150.0, 130.0, 25.0, 19.0])

gsa_result = gsa(gsa_cost, Sobol(order = [2]), [no_cells_yr_intrvl, cost_pos_AM_intrvl, cost_neg_AM_intrvl, cost_sep_intrvl], N = 10000)

x = [0, 1, 2, 3, 1, 2, 3, 2, 3, 3]
y = [0, 0, 0, 0, 1, 1, 1, 2, 2, 3]

print("\n")
S1 = gsa_result.S1
S2 = gsa_result.S2
print(S1)
print(S2)

z = [S1[1],S2[1],S2[2],S2[3],S1[2],S2[4],S2[5],S1[3],S2[6],S1[4]]
clf()
r = bar3D(x, y, z*0, 0.4, 0.4, z)
xticks([0,1,2,3], ["No Cells","Cost Pos AM", "Cost Neg AM", "Cost Sep"])
yticks([1,2,3,4], ["No Cells","Cost Pos AM", "Cost Neg AM", "Cost Sep"])

title(string("GSA: 2nd Order: ", per, "%"))
zlabel("Sobol Indices")
figure(2)

# savefig("GSA.png", dpi=500)


# This file is part of the TaylorIntegration.jl package; MIT licensed

using TaylorIntegration
using Test
using LinearAlgebra: Transpose, norm
using Logging
import Logging: Warn

@testset "Testing `many_ode.jl`" begin

    local _order = 28
    local _abstol = 1.0E-20
    local tT = Taylor1(_order)

    max_iters_reached() = "Maximum number of integration steps reached; exiting.\n"

    @testset "Tests: dot{x}=x.^2, x(0) = [3, 1]" begin
        function eqs_mov!(Dx, x, p, t)
            for i in eachindex(x)
                Dx[i] = x[i]^2
            end
            nothing
        end
        exactsol(t, x0) = x0/(1.0-x0*t)
        t0 = 0.0
        q0 = [3, 1]
        q0T = Taylor1.(Float64.(q0), _order)
        xdotT = Array{Taylor1{Float64}}(undef, length(q0))
        xaux = Array{Taylor1{Float64}}(undef, length(q0))
        tT[1] = t0
        TaylorIntegration.jetcoeffs!(eqs_mov!, tT, q0T, xdotT, xaux, nothing)
        @test q0T[1].coeffs[end] == 3.0^(_order+1)
        @test q0T[2].coeffs[end] == 1.0
        δt = (_abstol/q0T[1].coeffs[end-1])^inv(_order-1)
        @test TaylorIntegration.stepsize(q0T, _abstol) == δt

        tv, xv = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            eqs_mov!, q0, 0.0, 0.5, _order, _abstol, nothing))
        @test length(tv) == 501
        @test isa(xv, SubArray)
        @test xv[1,:] == q0
        @test tv[end] < 1/3

        trange = 0.0:1/8:1.0
        xv = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            eqs_mov!, q0, trange, _order, _abstol))
        @test size(xv) == (9,2)
        @test q0 == [3.0, 1.0]
        @test typeof(xv) == Transpose{Float64, Array{Float64,2}}
        @test xv[1,1:end] == q0
        @test (isnan(xv[4,1]) && isnan(xv[4,2]))
        @test (isnan(xv[end,1]) && isnan(xv[end,2]))
        @test abs(xv[3,2] - 4/3) ≤ eps(4/3)
        @test abs(xv[2,1] - 4.8) ≤ eps(4.8)

        tarray = vec(trange)
        xv2 = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            eqs_mov!, q0, tarray, _order, _abstol, nothing))
        @test xv[1:3,:] == xv2[1:3,:]
        @test xv2[1:3,:] ≈ xv[1:3,:] atol=eps() rtol=0.0
        @test size(xv2) == (9,2)
        @test q0 == [3.0, 1.0]
        @test typeof(xv2) == Transpose{Float64, Array{Float64,2}}
        @test xv2[1,1:end] == q0
        @test (isnan(xv2[4,1]) && isnan(xv2[4,2]))
        @test (isnan(xv2[end,1]) && isnan(xv2[end,2]))
        @test abs(xv2[3,2] - 4/3) ≤ eps(4/3)
        @test abs(xv2[2,1] - 4.8) ≤ eps(4.8)

        # Output includes Taylor polynomial solution
        tv, xv, polynV = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            eqs_mov!, q0, 0, 0.5, _order, _abstol, Val(true), nothing, maxsteps=2))
        @test size(polynV) == (3, 2)
        @test xv[1,:] == q0
        @test polynV[1,:] == Taylor1.(q0, _order)
        @test xv[2,:] == evaluate.(polynV[2, :], tv[2]-tv[1])
        @test polynV[2,2] == Taylor1(ones(_order+1))
    end

    @testset "Tests: dot{x}=x.^2, x(0) = [3, 3]" begin
        function eqs_mov!(Dx, x, p, t)
            for i in eachindex(x)
                Dx[i] = x[i]^2
            end
            nothing
        end
        exactsol(t, x0) = x0/(1.0-x0*t)

        q0 = [3.0, 3.0]
        tmax = 0.3
        tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
            eqs_mov!, q0, 0, tmax, _order, _abstol, nothing))
        @test length(tv) < 501
        @test length(xv[:,1]) < 501
        @test length(xv[:,2]) < 501
        @test length(tv) == 14
        @test length(xv[:,1]) == 14
        @test length(xv[:,2]) == 14
        @test xv[1,1:end] == q0
        @test tv[end] < 1/3
        @test tv[end] == tmax
        @test xv[end,1] == xv[end,2]
        @test abs(xv[end,1]-exactsol(tv[end], xv[1,1])) < 5e-14
        @test abs(xv[end,2]-exactsol(tv[end], xv[1,2])) < 5e-14

        tmax = 0.33
        tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
            eqs_mov!, [3, 3], 0.0, tmax, _order, _abstol))
        @test length(tv) < 501
        @test length(xv[:,1]) < 501
        @test length(xv[:,2]) < 501
        @test length(tv) == 28
        @test length(xv[:,1]) == 28
        @test length(xv[:,2]) == 28
        @test xv[1,1:end] == q0
        @test tv[end] < 1/3
        @test tv[end] == tmax
        @test xv[end,1] == xv[end,2]
        @test abs(xv[end,1]-exactsol(tv[end], xv[1,1])) < 1.0e-11
        @test abs(xv[end,2]-exactsol(tv[end], xv[1,2])) < 1.0e-11

        # Output includes Taylor polynomial solution
        tv, xv, polynV = taylorinteg(eqs_mov!, q0, 0, tmax, _order, _abstol, Val(true), nothing)
        @test size(polynV) == size(xv)
        @test polynV[1,:] == Taylor1.(q0, _order)
        @test xv[end,1] == evaluate.(polynV[end, 1], tv[end]-tv[end-1])
        @test polynV[:,1] == polynV[:,2]
    end

    @testset "Test non-autonomous ODE (2): dot{x}=cos(t)" begin
        function f!(Dx, x, p, t)
            Dx[1] = one(t)
            Dx[2] = cos(t)
            nothing
        end
        t0 = 0//1
        tmax = 10.25*(2pi)
        abstol = 1e-20
        order = 25
        x0 = [t0, 0.0] #initial conditions such that x(t)=sin(t)

        tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
            f!, x0, t0, tmax, order, abstol))
        @test length(tv) < 501
        @test length(xv[:,1]) < 501
        @test length(xv[:,2]) < 501
        @test xv[1,1:end] == x0
        @test tv[1] < tv[end]
        @test tv[end] == xv[end,1]
        @test abs(sin(tmax)-xv[end,2]) < 1e-14

        # Backward integration
        tb, xb = (@test_logs min_level=Logging.Warn taylorinteg(
            f!, [tmax, sin(tmax)], tmax, t0, order, abstol))
        @test length(tb) < 501
        @test length(xb[:,1]) < 501
        @test length(xb[:,2]) < 501
        @test tb[1] > tb[end]
        @test xb[1,1:end] == [tmax, sin(tmax)]
        @test tb[end] == xb[end,1]
        @test abs(sin(t0)-xb[end,2]) < 5e-14

        # Tests with a range, for comparison with backward integration
        tmax = 15*(2pi)
        Δt = (tmax-t0)/1024
        tspan = t0:Δt:tmax
        xv = (@test_logs min_level=Logging.Warn taylorinteg(
            f!, x0, tspan, order, abstol, nothing))
        @test xv[1,1:end] == x0
        @test tmax == xv[end,1]
        @test abs(sin(tmax)-xv[end,2]) < 1e-14

        # Backward integration
        xback = (@test_logs min_level=Logging.Warn taylorinteg(
            f!, xv[end, :], reverse(tspan), order, abstol, nothing))
        @test xback[1,:] == xv[end, :]
        @test abs(xback[end,1]-x0[1]) < 5.0e-14
        @test abs(xback[end,2]-x0[2]) < 5.0e-14
        @test abs(sin(t0)-xback[end,2]) < 5.0e-14
        @test norm((xv[:,:]-xback[end:-1:1,:]), Inf) < 5.0e-14

        # Tests if trange is properly sorted
        @test_throws AssertionError taylorinteg(f!, x0, rand(t0:Δt:tmax, 100), order, abstol)
    end

    @testset "Falling ball (stepsize)" begin
        function fallball!(dx, x, p, t)
            dx[1] = x[2]
            dx[2] = -one(x[1])
            nothing
        end
        exactsol(t, x0, v0) = x0 + v0*t - 0.5*t^2
        t0 = 0.0
        tmax = 5.0
        abstol = 1e-20
        order = 10
        x0 = [10.0, 0.0] #initial conditions such that x(t)=sin(t)
        tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
            fallball!, x0, t0, tmax, order, abstol))
        @test length(tv) < 501
        @test length(xv[:,1]) < 501
        @test exactsol.(tv, x0[1], x0[2]) ≈ xv[:,1]
    end

end


@testset "utils tests" begin
    @testset "combine_fuzzy_sets tests" begin
        A = Float32[1.0 0.1; 0.4 1.0]
        union_res = Float32[1.0 0.46; 0.46 1.0]
        res = combine_fuzzy_sets(A, 1f0)
        @test isapprox(res, union_res)
        inter_res = [1.0 0.04; 0.04 1.0]
        res = combine_fuzzy_sets(A, 0f0)
        @test isapprox(res, inter_res)

    end

    @testset "fit_ab tests" begin
        @test all((1, 2) .== fit_ab(-1, 0, 1, 2))
    end
end


using Stheno, Zygote, BenchmarkTools
using Stheno: @model, BlockData, cross

const x = collect(range(-5.0, 5.0; length=50));
const x1 = randn(25)
const x2 = randn(25)

θ = (m=5.0, l=1.3);

@model foo(θ) = GP(ConstMean(θ.m), kernel(SEKernel(); l=θ.l));

#
# Test single-process versions of operations on foo.
#

y = rand(foo(θ)(x, 0.1));

@code_warntype mean(foo(θ)(x, 0.1))
@benchmark mean(foo($θ)($x, 0.1))

@code_warntype cov(foo(θ)(x, 0.1))
@benchmark cov(foo($θ)($x, 0.1))

@code_warntype rand(foo(θ)(x, 0.1))
@benchmark rand(foo($θ)($x, 0.1))

@code_warntype logpdf(foo(θ)(x, 0.1), y)
@benchmark logpdf(foo($θ)($x, 0.1), $y)

#
# Test multi-process versions of operations on foo.
#

function test_cross(θ)
    f = foo(θ)
    return cross([f, f])
end

@code_warntype test_cross(θ)

test_kernels(f) = kernel.([f, f], permutedims([f, f]))
@code_warntype test_kernels(f)

test_kernel(f) = Stheno.kernel(f, f)
@code_warntype test_kernel(f)

function test_multi_mean(θ, x1, x2)
    f = foo(θ)
    return mean(cross([f, f])(BlockData([x1, x2]), 0.1))
end

@code_warntype test_multi_mean(θ, x1, x2)

function test_multi_rand(θ, x1, x2)
    f = foo(θ)
    fx1 = f(x1, 0.1)
    fx2 = f(x2, 0.1)
    return rand([fx1, fx2])
end

y1, y2 = test_multi_rand(θ, x1, x2);

@code_warntype test_multi_rand(θ, x, x)
@benchmark test_multi_rand($θ, $x1, $x2)

function test_multi_logpdf(θ, x1, x2, y1, y2)
    f = foo(θ)
    fx1 = f(x1, 0.1)
    fx2 = f(x2, 0.1)
    return logpdf([fx1 ← y1, fx2 ← y2])
end

@code_warntype test_multi_logpdf(θ, x1, x2, y1, y2)
@benchmark test_multi_logpdf($θ, $x1, $x2, $y1, $y2)


using DataFrames
using CSV
using Calculus
using Plots
using Statistics

plotly()
@doc "
    traffic_assignment(;networkName, tol=0.001, maxIters=20)
"
function traffic_assignment(;networkName, tol=0.001, maxIters=20)
    networkFile = "Network\\$networkName\\network.csv"
    demandFile = "Network\\$networkName\\demand.csv"

    𝒩 = Int64[]                            # Nodes
    𝒜 = Array{Int64,1}[]                   # Arcs as adjacency list
    v₀ = Array{Float64,1}[]                 # Link volume capcity
    t₀ = Array{Float64,1}[]                 # Link free flow travel time
    α₀ = Array{Float64,1}[]                 # BPR parameters
    β₀ = Array{Float64,1}[]                 # BPR parameters

    W = Array{Int64,1}[]                    # OD pairs
    q = Float64[]                           # Demand between OD pairs
    P = Array{Array{Int64,1},1}[]           # Paths between OD pairs

    cⁱ = Float64[]                  # Stores total cost from each iteration
    xⁱₐ = Array{Array{Float64,1}}[] # Stores link flows from each iteration
    eⁱ = Float64[]                  # Stores convergent from each iteration

    # Enumerates all paths from pivot node i using nodes from a vector of open nodes 𝕏
    function enumerate(i, 𝕏)
        paths = Array{Int64,1}[]          # Vector of paths
        𝕏 = filter(x -> (x ≠ i), 𝕏)

        # Dead ends: if there is no outgoing arc from node i or if all outgoing nodes from i are closed
        if isempty(𝒜[i]) push!(paths, [i]) end
        if sum([if 𝒜[i][j] in 𝕏 1 else 0 end for j in 1:length(𝒜[i])]) == 0 push!(paths, [i]) end

        # Dveveloping paths recursively with subpaths
        for j in 1:length(𝒜[i])
            if 𝒜[i][j] in 𝕏
                for subpath in enumerate(𝒜[i][j], 𝕏)
                    push!(paths, [i])
                    append!(paths[end], subpath)
                end
            end
        end
        return paths
    end

    # Network build
    function build()
        println()
        # File import
        network = CSV.File(networkFile, types=[Int64, Int64, Float64, Float64, Float64, Float64, Float64])
        df = DataFrame(network)
        head = df[1]::Array{Int64,1}
        tail = df[2]::Array{Int64,1}
        linkcapacity = df[3]::Array{Float64,1}
        linklength = df[4]::Array{Float64,1}
        linkfft = df[5]::Array{Float64,1}
        alpha = df[6]::Array{Float64,1}
        beta = df[7]::Array{Float64,1}

        n = max(maximum(head), maximum(tail))
        for i in 1:n
            append!(𝒩, i)
            push!(𝒜, [])
            push!(v₀, [])
            push!(t₀, [])
            push!(α₀, [])
            push!(β₀, [])
        end

        for i in 1:length(head)
            append!(𝒜[head[i]], tail[i])
            append!(v₀[head[i]], linkcapacity[i])
            append!(t₀[head[i]], linkfft[i])
            append!(α₀[head[i]], alpha[i])
            append!(β₀[head[i]], beta[i])
        end

        demand = CSV.File(demandFile, types=[Int64, Int64, Float64])
        df′ = DataFrame(demand)
        origin = df′[1]::Array{Int64,1}
        destination = df′[2]::Array{Int64,1}
        demand = df′[3]::Array{Float64,1}

        for i in 1:nrow(df′)
            push!(W, [origin[i], destination[i]])
            append!(q, demand[i])
            push!(P, [])
        end

        for r in unique(origin)
            S = [W[k][2] for k in findall(x -> (x == r), origin)]
            for s in S
                index = findfirst(x -> (x == [r,s]), W)
                paths = Array{Int64,1}[]
                for path in enumerate(r, copy(𝒩))
                    if s in path
                        i = findfirst(x -> (x == s), path)
                        push!(paths, path[1:i])
                    end
                end
                append!(P[index], paths)
            end
        end
    end

    # Arc cost
    function 𝐶ₐ(i, j, x)
        α = α₀[i][j]
        β = β₀[i][j]
        t = t₀[i][j]
        v = v₀[i][j]
        return t * (1 + α * (x/v) ^ β)
    end

    # Path cost for path p given arc flows x and arc cost function cₐ
    function 𝐶ₚ(p, x, cₐ)
        cost = 0.0
        for k in 2:length(p)
            i = p[k-1]
            j = findfirst(x -> (x == p[k]), 𝒜[i])::Int64
            cost += cₐ(i, j, x[i][j])
        end
        return cost
    end

    # Djikstra's shortest path between r and s
    function sp(𝒞, r, s)
        p = Int64[]
        ℒ = [[if i==r r else -1 end, if i==r 0.0 else Inf end] for i in 𝒩]    # predecessor and cost
        𝕏 = copy(𝒩)   # Set of open nodes
        i = r
        deleteat!(𝕏, i)
        while !isempty(𝕏)
            for j in 1:length(𝒜[i])
                k = 𝒜[i][j]
                C = ℒ[i][2] + 𝒞[i][j]
                if C < ℒ[k][2] && k in 𝕏
                    ℒ[k][1] = i
                    ℒ[k][2] = C
                end
            end
            index = argmin([ℒ[i][2] for i in 𝕏])
            i = 𝕏[index]
            deleteat!(𝕏, index)
        end

        i = s
        append!(p, i)
        while i ≠ r
            i = Int(ℒ[i][1])
            append!(p, i)
        end
        reverse!(p)
        return p
    end

    # General Equilibration Algorithm
    function gea(;cₐ, ε, kₘₐₓ)
        cₚ = 𝐶ₚ
        ε′ = 0.000001
        k′ₘₐₓ = 20

        xᵏₐ = Array{Float64,1}[[0.0 for j in 𝒜[i]] for i in 𝒩]
        fᵏₚ = Array{Float64,1}[[0.0 for _ in P[n]] for n in 1:length(W)]

        # Intialization
        cᵒₐ = [[cₐ(i, j, 0.0) for j in 1:length(𝒜[i])] for i in 𝒩]
        for n in 1:length(W)
            r, s = W[n]
            qᵣₛ = q[n]
            pᵣₛ = sp(cᵒₐ, r, s)
            m = findfirst(x -> (x == pᵣₛ), P[n])
            fᵏₚ[n][m] = qᵣₛ
            for l in 2:length(pᵣₛ)
                i = pᵣₛ[l-1]
                j = findfirst(x -> (x == pᵣₛ[l]), 𝒜[i])
                xᵏₐ[i][j] += fᵏₚ[n][m]
            end
        end

        # Iteration
        allConverged = false
        k = 1
        while !allConverged
            xᵏ⁺¹ₐ = deepcopy(xᵏₐ)
            fᵏ⁺¹ₚ = deepcopy(fᵏₚ)

            for n in 1:length(W)
                odConverged=false
                k′ = 1
                while !odConverged
                    c = [cₚ(P[n][m], xᵏ⁺¹ₐ, cₐ) for m in 1:length(P[n])]
                    indices = findall(x -> (x > 0), fᵏ⁺¹ₚ[n])

                    # Active path with highest costs
                    m₁ = indices[argmax([c[m] for m in indices])]
                    # Path with least cost
                    m₂ = argmin(c)

                    X = [[P[n][m₁][l-1], P[n][m₁][l]] for l in 2:length(P[n][m₁])]
                    Y = [[P[n][m₂][l-1], P[n][m₂][l]] for l in 2:length(P[n][m₂])]
                    Z = intersect(X, Y)
                    num = c[m₁] - c[m₂]
                    den = Float64(length(X) + length(Y) - 2*length(Z))
                    if num == den == 0 δ = 0.0
                    else δ = min(num/den, fᵏ⁺¹ₚ[n][m₁])
                    end

                    # OD convergence test
                    if δ < ε′ || k′ == k′ₘₐₓ odConverged = true end

                    # Updation
                    fᵏ⁺¹ₚ[n][m₁] -= δ
                    fᵏ⁺¹ₚ[n][m₂] += δ
                    for l in 2:length(P[n][m₁])
                        i = P[n][m₁][l-1]
                        j = findfirst(x -> (x == P[n][m₁][l]), 𝒜[i])
                        xᵏ⁺¹ₐ[i][j] -= δ
                    end
                    for l in 2:length(P[n][m₂])
                        i = P[n][m₂][l-1]
                        j = findfirst(x -> (x == P[n][m₂][l]), 𝒜[i])
                        xᵏ⁺¹ₐ[i][j] += δ
                    end

                    k′ += 1
                end
            end

            # Overall convergence test
            num = 0.0
            den = 0.0
            for i in 𝒩
                for j in 1:length(𝒜[i])
                    num += (xᵏ⁺¹ₐ[i][j] - xᵏₐ[i][j])^2
                    den += xᵏₐ[i][j]
                end
            end
            Δ = sqrt(num)/den
            if Δ < ε || k == kₘₐₓ allConverged = true end

            # Updation
            xᵏₐ = deepcopy(xᵏ⁺¹ₐ)
            fᵏₚ = deepcopy(fᵏ⁺¹ₚ)

            k += 1
        end
        return xᵏₐ
    end

    # Variational Inequality: Projection method
    function vi(;cₐ=𝐶ₐ, ε=tol, kₘₐₓ=maxIters)
        xᵏₐ = Array{Float64,1}[[0.0 for j in 𝒜[i]] for i in 𝒩]

        # Intialization
        cₒ = [[cₐ(i, j, 0.0) for j in 1:length(𝒜[i])] for i in 𝒩]
        for n in 1:length(W)
            r, s = W[n]
            qᵣₛ = q[n]
            pᵣₛ = sp(cₒ, r, s)
            for l in 2:length(pᵣₛ)
                i = pᵣₛ[l-1]
                j = findfirst(x -> (x == pᵣₛ[l]), 𝒜[i])
                xᵏₐ[i][j] += qᵣₛ
            end
        end

        f′(i,j,x) = derivative(x -> cₐ(i,j,x))(x)
        L = sum([(f′(i, j, xᵏₐ[i][j]))^2 for i in 𝒩 for j in 1:length(𝒜[i])])^0.5
        γ = 1/L
        println("gamma: $γ")

        println("\niter: 0")
        println("total flow: ", sum(sum.(xᵏₐ)))
        println("total cost: ", sum([xᵏₐ[i][j] * cₐ(i, j, xᵏₐ[i][j]) for i in 𝒩 for j in 1:length(𝒜[i])]))
        println("convergent: NA")

        append!(cⁱ, sum([xᵏₐ[i][j] * cₐ(i, j, xᵏₐ[i][j]) for i in 𝒩 for j in 1:length(𝒜[i])]))
        push!(xⁱₐ, xᵏₐ)

        # Iteration
        converged = false
        k = 1
        while !converged
            xᵏ⁺¹ₐ = Array{Float64,1}[[0.0 for j in 𝒜[i]] for i in 𝒩]

            # Auxilary problem
            c(i,j,x) = x - (xᵏₐ[i][j] - γ*cₐ(i,j,xᵏₐ[i][j]))
            xᵏ⁺¹ₐ = gea(cₐ=c, ε=0.000001, kₘₐₓ=20)

            # Convergence test
            Δ = sqrt(sum(vcat((xᵏ⁺¹ₐ - xᵏₐ)...) .^2))/sum(vcat(xᵏₐ...))
            if Δ < ε || k == kₘₐₓ converged = true end

            # Updating link flows and cost
            xᵏₐ = deepcopy(xᵏ⁺¹ₐ)

            println("\niter: $k")
            println("total flow: ", sum(sum.(xᵏₐ)))
            println("total cost: ", sum([xᵏₐ[i][j] * cₐ(i, j, xᵏₐ[i][j]) for i in 𝒩 for j in 1:length(𝒜[i])]))
            println("convergent: $Δ")

            append!(cⁱ, sum([xᵏₐ[i][j] * cₐ(i, j, xᵏₐ[i][j]) for i in 𝒩 for j in 1:length(𝒜[i])]))
            push!(xⁱₐ, xᵏₐ)
            append!(eⁱ, Δ)

            k += 1
        end
        return xᵏₐ
    end

    build()
    println(vi())
    return cⁱ, xⁱₐ, eⁱ
end
traffic_assignment(networkName="Braess-Example", tol=0.001, maxIters=20)


function dis2spt(w::Array{Float64,1}, dis::Array{Float64,1}, nz::Int64, nx::Int64, ext::Int64, iflag::Int64, dt::Float64, it::Int64)
    if length(w) < it
       error("out of source time range")
    end
    if length(dis) != nz*nx
       error("length of dis does not much model size")
    end
    if iflag == 1
       zupper = ext
       Nz = nz + 2*ext
    elseif iflag == 2
       zupper = 0
       Nz = nz +   ext
    end
    Nx = nx + 2*ext
    spt = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, it)
    tmp = zeros(Nz, Nx); dis_tmp = reshape(dis, nz, nx)
    tmp[zupper+1: zupper+nz, ext+1:ext+nx] = 1/2 * w[it] * dis_tmp
    tmp = vec(tmp)
    spt.pz[:] = tmp[:]
    spt.px[:] = tmp[:]
    return spt
end

function spt2dis(w::Array{Float64,1}, spt::SnapShot)
    nz = spt.nz ; nx = spt.nx;
    ext= spt.ext; iflag = spt.iflag;
    it = spt.it
    if length(w) < spt.it
       error("out of source time range")
    end
    if iflag == 1
       zupper = ext
       Nz = nz + 2*ext
    elseif iflag == 2
       zupper = 0
       Nz = nz +   ext
    end
    Nx = nx + 2*ext
    pz = reshape(spt.pz, Nz, Nx)[zupper+1:zupper+nz, ext+1:ext+nx]
    px = reshape(spt.px, Nz, Nx)[zupper+1:zupper+nz, ext+1:ext+nx]
    dis = 1/2 * w[it] * pz + 1/2 * w[it] * px
    return vec(dis)
end

function spt2wlet(w::Array{Float64,1}, dis::Array{Float64,1}, spt::SnapShot)
    nz = spt.nz ; nx = spt.nx;
    ext= spt.ext; iflag = spt.iflag;
    it = spt.it
    if length(w) < it
       error("out of source time range")
    end
    if iflag == 1
       zupper = ext
       Nz = nz + 2*ext
    elseif iflag == 2
       zupper = 0
       Nz = nz +   ext
    end
    Nx = nx + 2*ext
    p = reshape(spt.pz, Nz, Nx)[zupper+1:zupper+nz, ext+1:ext+nx] + reshape(spt.px, Nz, Nx)[zupper+1:zupper+nz, ext+1:ext+nx]
    p = vec(p)
    w[it] = 1/2 * dot(dis, p)
    return nothing
end


function get_dudxe!(solver, u, Ke, xe, penalty, E0, xmin, global_dofs)
    E, dE = get_E_dE(xe, penalty, E0, xmin)
    solver.rhs .= 0
    for i in 1:size(Ke, 2)
        solver.rhs[global_dofs] .-= dE .* Ke[:, i] .* u[global_dofs[i]]
    end
    solver(assemble_f = false)
    return solver.lhs
end

get_sigma_vm(E_e, utMu_e) = E_e * sqrt(utMu_e)

function get_E(x_e::T, penalty, E0, xmin) where T
    return E0 * density(penalty(x_e), xmin)
end

function get_E_dE(x_e::T, penalty, E0, xmin) where T
    d = ForwardDiff.Dual{T}(x_e, one(T))
    p = density(penalty(d), xmin)
    g = p.partials[1] * E0
    return p.value * E0, g
end

struct CorrectedStress <: Function end
function (s::CorrectedStress)(x_e, sigma_vm_e, sigma_bar)
    return x_e * (sigma_vm_e - sigma_bar)
end

struct OffsetCorrectedStress <: Function end
function (s::OffsetCorrectedStress)(x_e, sigma_vm_e, sigma_bar)
    return x_e * sigma_vm_e + (1 - x_e) * sigma_bar
end

function get_s(x_e, sigma_vm_e, sigma_bar)
    return x_e * sigma_vm_e + (1 - x_e) * sigma_bar
end

@inline function get_ϵ(u, ∇ϕ, i, j)
	return 1/2*(u[i]*∇ϕ[j] + u[j]*∇ϕ[i])
end
@inline function apply_T!(Tu, u, dh, cellidx, global_dofs, cellvalues, ν, ::Val{2})
    # assumes cellvalues is initialized before passing
    Tu[:] .= 0
    # Generalize to higher order field basis functions
    q_point = 1
    n_basefuncs = getnbasefunctions(cellvalues)
    dim = 2
    for a in 1:n_basefuncs
        ∇ϕ = shape_gradient(cellvalues, q_point, a)
        _u = @view u[(@view global_dofs[dim*(a-1) + 1 : a*dim])]
        ϵ_11 = get_ϵ(_u, ∇ϕ, 1, 1)
        ϵ_22 = get_ϵ(_u, ∇ϕ, 2, 2)
        ϵ_12 = get_ϵ(_u, ∇ϕ, 1, 2)

        ϵ_sum = ϵ_11 + ϵ_22

        temp1 = ν/(1-ν^2)
        temp2 = ν*(1+ν)

        Tu[1] += temp1*ϵ_sum + temp2*ϵ_11 # σ[1,1] / E
        Tu[2] += temp1*ϵ_sum + temp2*ϵ_22 # σ[2,2] / E
        Tu[3] += temp2*ϵ_12 # σ[1,2] / E
    end
    return Tu
end
@inline function apply_T!(Tu, u, dh, global_dofs, cellvalues, ν, ::Val{3})
    # assumes cellvalues is initialized before passing
    Tu[:] .= 0
    q_point = 1
    n_basefuncs = getnbasefunctions(cellvalues)
    dim = 3
    for a in 1:n_basefuncs
        ∇ϕ = shape_gradient(cellvalues, q_point, a)
        _u = @view u[(@view global_dofs[dim*(a-1) + 1 : a*dim])]
        ϵ_11 = get_ϵ(_u, ∇ϕ, 1, 1)
        ϵ_22 = get_ϵ(_u, ∇ϕ, 2, 2)
        ϵ_33 = get_ϵ(_u, ∇ϕ, 3, 3)
        ϵ_12 = get_ϵ(_u, ∇ϕ, 1, 2)
        ϵ_23 = get_ϵ(_u, ∇ϕ, 2, 3)
        ϵ_31 = get_ϵ(_u, ∇ϕ, 3, 1)

        ϵ_sum = ϵ_11 + ϵ_22 + ϵ_33

        temp1 = ν/(1-ν^2)
        temp2 = ν*(1+ν)

        Tu[1] += temp1*ϵ_sum + temp2*ϵ_11 # σ[1,1] / E
        Tu[2] += temp1*ϵ_sum + temp2*ϵ_22 # σ[2,2] / E
        Tu[3] += temp1*ϵ_sum + temp2*ϵ_33 # σ[3,3] / E
        Tu[4] += temp2*ϵ_12 # σ[1,2] / E
        Tu[5] += temp2*ϵ_23 # σ[2,3] / E
        Tu[6] += temp2*ϵ_31 # σ[3,1] / E
    end
    return Tu
end

function fill_T!(T, ::Val{3}, cellvalues, ν)
    # assumes cellvalues is initialized before passing
    dim = 3
    temp1 = ν/(1-ν^2)
    temp2 = ν*(1+ν)
    q_point = 1
    n_basefuncs = getnbasefunctions(cellvalues)
    @assert size(T) == (6, dim*n_basefuncs)
    for a in 1:n_basefuncs
        ∇ϕ = shape_gradient(cellvalues, q_point, a)
        cols = dim * (a - 1) + 1 : dim * a
        T[1, cols[1]] = (temp1 + temp2) * ∇ϕ[1]
        T[2, cols[1]] = temp1 * ∇ϕ[1]
        T[3, cols[1]] = temp1 * ∇ϕ[1]
        T[4, cols[1]] = temp2 * ∇ϕ[2] / 2
        T[5, cols[1]] = 0
        T[6, cols[1]] = temp2 * ∇ϕ[3] / 2

        T[1, cols[2]] = temp1 * ∇ϕ[2]
        T[2, cols[2]] = (temp1 + temp2) * ∇ϕ[2]
        T[3, cols[2]] = temp1 * ∇ϕ[2]
        T[4, cols[2]] = temp2 * ∇ϕ[1] / 2
        T[5, cols[2]] = temp2 * ∇ϕ[3] / 2
        T[6, cols[2]] = 0

        T[1, cols[3]] = temp1 * ∇ϕ[3]
        T[2, cols[3]] = temp1 * ∇ϕ[3]
        T[3, cols[3]] = (temp1 + temp2) * ∇ϕ[3]
        T[4, cols[3]] = 0
        T[5, cols[3]] = temp2 * ∇ϕ[2] / 2
        T[6, cols[3]] = temp2 * ∇ϕ[1] / 2
    end
    return T
end

function fill_T!(T, ::Val{2}, cellvalues, ν)
    # assumes cellvalues is initialized before passing
    dim = 2
    temp1 = ν/(1-ν^2)
    temp2 = ν*(1+ν)
    q_point = 1
    n_basefuncs = getnbasefunctions(cellvalues)
    @assert size(T) == (3, dim*n_basefuncs)
    for a in 1:n_basefuncs
        ∇ϕ = shape_gradient(cellvalues, q_point, a)
        cols = dim * (a - 1) + 1 : dim * a
        T[1, cols[1]] = (temp1 + temp2) * ∇ϕ[1]
        T[2, cols[1]] = temp1 * ∇ϕ[1]
        T[3, cols[1]] = temp2 * ∇ϕ[2] / 2

        T[1, cols[2]] = temp1 * ∇ϕ[2]
        T[2, cols[2]] = (temp1 + temp2) * ∇ϕ[2]
        T[3, cols[2]] = temp2 * ∇ϕ[1] / 2
    end
    return T
end

@params struct StressTemp{T}
    VTu::AbstractVector{T}
    Tu::AbstractVector{T}
    Te::AbstractMatrix{T}
    global_dofs::AbstractVector{Int}
end
function StressTemp(solver)
    @unpack u, problem = solver
    @unpack dh = problem.ch
    T = eltype(u)
    dim = TopOptProblems.getdim(problem)
    k = dim == 2 ? 3 : 6
    VTu = zeros(T, k)
    Tu = similar(VTu)
    Te = zeros(T, k, dim * getnbasefunctions(solver.elementinfo.cellvalues))
    global_dofs = zeros(Int, ndofs_per_cell(dh))

    return StressTemp(VTu, Tu, Te, global_dofs)
end

function fill_Mu_utMu!(Mu, utMu, solver, stress_temp)
    @unpack problem, elementinfo, u = solver
    @unpack ch, ν = problem
    @unpack dh = ch
    @unpack VTu, Tu, Te, global_dofs = stress_temp

    _fill_Mu_utMu!(Mu, utMu, dh, elementinfo, u, ν, global_dofs, Tu, VTu, Te)
    return Mu, utMu
end

function _fill_Mu_utMu!( Mu::AbstractMatrix{T},
                        utMu::AbstractVector{T},
                        dh::DofHandler{3},
                        elementinfo,
                        u,
                        ν,
                        global_dofs = zeros(Int, ndofs_per_cell(dh)),
                        Tu = zeros(T, 6),
                        VTu = zeros(T, 6),
                        Te = zeros(T, 6, ndofs_per_cell(dh))
                      ) where {T}
    dim = 3
    @unpack cellvalues = elementinfo
    for (cellidx, cell) in enumerate(CellIterator(dh))
        reinit!(cellvalues, cell)
        Mu_e = @view Mu[:, cellidx]
        fill_T!(Te, Val(3), cellvalues, ν)
        celldofs!(global_dofs, dh, cellidx)
        @views mul!(Tu, Te, u[global_dofs])

		VTu[1] = Tu[1] - Tu[2]/2 - Tu[3]/2
		VTu[2] = -Tu[1]/2 + Tu[2] - Tu[3]/2
		VTu[3] = -Tu[1]/2 - Tu[2]/2 + Tu[3]
		VTu[4] = 3*Tu[4]
		VTu[5] = 3*Tu[5]
		VTu[6] = 3*Tu[6]

        utMu_e = dot(Tu, VTu)
        @assert utMu_e > 0
        utMu[cellidx] = utMu_e
        mul!(Mu_e, Te', VTu)
	end
	return Mu, utMu
end

function _fill_Mu_utMu!( Mu::AbstractMatrix{T},
                        utMu::AbstractVector{T},
                        dh::DofHandler{2},
                        elementinfo,
                        u,
                        ν,
                        global_dofs = zeros(Int, ndofs_per_cell(dh)),
                        Tu = zeros(T, 3),
                        VTu = zeros(T, 3),
                        Te = zeros(T, 3, ndofs_per_cell(dh))
                      ) where {T}
    dim = 2
    @unpack cellvalues = elementinfo
    for (cellidx, cell) in enumerate(CellIterator(dh))
        reinit!(cellvalues, cell)
        Mu_e = @view Mu[:, cellidx]
        fill_T!(Te, Val(2), cellvalues, ν)
        celldofs!(global_dofs, dh, cellidx)
        @views mul!(Tu, Te, u[global_dofs])

        VTu[1] = Tu[1] - Tu[2]/2
		VTu[2] = -Tu[1]/2 + Tu[2]
		VTu[3] = 3*Tu[3]

        utMu_e = dot(Tu, VTu)
        @assert utMu_e > 0
        utMu[cellidx] = utMu_e
        mul!(Mu_e, Te', VTu)
	end
	return Mu, utMu
end

@params mutable struct GlobalStress{T} <: AbstractFunction{T}
    reducer
    reducer_g::AbstractVector{T}
    corrected_stress
    utMu::AbstractVector{T}
    Mu::AbstractArray{T}
    s::AbstractVector{T}
    sigma_bar::T
    solver
    global_dofs::AbstractVector{<:Integer}
    buffer::AbstractVector{T}
    stress_temp::StressTemp
    fevals::Int
    reuse::Bool
    maxfevals::Int
end
function GlobalStress(solver, sigma_bar, reducer = WeightedKNorm(4, 1/length(solver.vars)), stress = OffsetCorrectedStress(); reuse = false, maxfevals = 10^8)
    T = eltype(solver.u)
    dh = solver.problem.ch.dh
    k = ndofs_per_cell(dh)
    N = getncells(dh.grid)
    global_dofs = zeros(Int, k)
    Mu = zeros(T, k, N)
    utMu = zeros(T, N)
    stress_temp = StressTemp(solver)
    s = similar(utMu)
    buffer = zeros(T, ndofs_per_cell(dh))
    reducer_g = similar(utMu)

    return GlobalStress(reducer, reducer_g, stress, utMu, Mu, s, sigma_bar, solver, global_dofs, buffer, stress_temp, 0, reuse, maxfevals)
end

function (gs::GlobalStress)(x, g)
    @unpack s, sigma_bar, Mu, utMu, buffer, stress_temp, corrected_stress, reuse = gs
    @unpack solver, global_dofs, buffer, reducer, reducer_g = gs
    @unpack elementinfo, u, penalty, problem, xmin = solver
    @unpack Kes = elementinfo
    @unpack dh = problem.ch
    E0 = problem.E
    gs.fevals += 1

    @assert length(global_dofs) == ndofs_per_cell(solver.problem.ch.dh)
    if !reuse
        solver()
        fill_Mu_utMu!(Mu, utMu, solver, stress_temp)
    end
    s .= corrected_stress.(x, get_sigma_vm.(get_E.(x, Ref(penalty), E0, xmin), utMu), sigma_bar)
    reduced = reducer(s, reducer_g)
    g .= 0
    for e1 in 1:length(g)
        if !reuse
            celldofs!(global_dofs, dh, e1)
            dudxe = get_dudxe!(solver, u, Kes[e1], x[e1], penalty, E0, xmin, global_dofs)
            @views buffer .= dudxe[global_dofs]
        end
        for e2 in 1:length(g)
            utMu_e2 = utMu[e2]
            Ee2, dEe2 = get_E_dE(x[e2], penalty, E0, xmin)
            t1 = (e1 == e2) * (dEe2 * sqrt(utMu_e2))
            if reuse
                t2 = zero(T)
            else
                @views t2 = (Ee2 / sqrt(utMu_e2)) * dot(buffer, Mu[:, e2])
            end
            dsigmae2_dxe1 = t1 + t2
            dse2_dxe1 = (e1 == e2) * (Ee2 * sqrt(utMu_e2) - sigma_bar) + x[e2] * dsigmae2_dxe1
            g[e1] += dse2_dxe1 * reducer_g[e2]
        end
    end
    return reduced
end

struct LogSumExp <: Function end
function (lse::LogSumExp)(s, g_s)
    out = logsumexp(s)
    g_s .= exp.(s .- out)
    return out
end

struct OffsetLogSumExp{T} <: Function
    k::T
end
function (olse::OffsetLogSumExp)(s, g_s)
    lse = logsumexp(s)
    g_s .= exp.(s .- lse)
    return out - olse.k * log(length(s))
end

struct KNorm <: Function
    k::Int
end
function (knorm::KNorm)(s, g_s)
    out = norm(s, knorm.k)
    g_s .= (s ./ out).^(k-1)
    return out
end

struct WeightedKNorm{T} <: Function
    k::Int
    w::T
end
function (wknorm::WeightedKNorm{T})(s, g_s) where {T}
    @unpack k, w = wknorm
    if T <: AbstractVector
        mw = MappedArray(w -> w^(1/k), w)
        out = norm(MappedArray(*, s, mw), k)
    else
        mw = w^(1/k)
        out = norm(BroadcastArray(*, s, mw), k)
    end
    g_s .= (s ./ out).^(k-1) .* mw
    return out
end


using BlockSparseMatrices
using Base.Test

include("test_SparseMatrixBCS.jl")


module PSIS

using Turing, DynamicPPL

export psislw
export psisloo

# Copyright (c) 2015 Aki Vehtari, Tuomas Sivula
# Original Matlab version by Aki Vehtari. Translation to Python
# by Tuomas Sivula.

# This software is distributed under the GNU General Public
# License (version 3 or later); please refer to the file
# License.txt, included with the software, for details.

"""
    psisloo(log_lik, [wcpp, wtrunc])
PSIS leave-one-out log predictive densities.
Computes the log predictive densities given posterior samples of the log
likelihood terms p(y_i|\theta^s) in input parameter `log_lik`. Returns a
sum of the leave-one-out log predictive densities `loo`, individual
leave-one-out log predictive density terms `loos` and an estimate of Pareto
tail indeces `ks`. If tail index k>0.5, variance of the raw estimate does
not exist and if tail index k>1 the mean of the raw estimate does not exist
and the PSIS estimate is likely to have large variation and some bias.
# Arguments
* `log_lik::Union{AbstractArray, Mamba.Chains}`: Array of size n x m containing n posterior samples of the log likelihood terms p(y_i|\theta^s).
* `wcpp::Real`: Percentage of samples used for GPD fit estimate (default is 20).
* `wtrunc::Float64`: Positive parameter for truncating very large weights to n^wtrunc. Providing False or 0 disables truncation. Default values is 3/4.
# Returns
* `loo::Real`: sum of the leave-one-out log predictive densities.
* `loos::AbstractArray`: Individual leave-one-out log predictive density terms.* `ks::AbstractArray`: Estimated Pareto tail indeces.
"""
# Compute LOO and standard error
function psisloo(model:: Model, chain:: Chains, wcpp::Int64=20, wtrunc::Float64=3/4)

    # log raw weights from log_lik

    lw = pointwise_loglikelihoods(model, chain)

    # compute Pareto smoothed log weights given raw log weights
    lwp, ks = psislw(-lw, wcpp, wtrunc)

    lwp += lw
    loos = logsumexp(lwp, 1)
    loo = sum(loos)

    return loo, loos, ks
end

"""
    psislw(lw, [wcpp, wtrunc])
Compute the Pareto smoothed importance sampling (PSIS).
# Arguments
* `lw::Union{AbstractArray, Mamba.Chains}`: Array of size n x m containing m sets of n log weights. It is also possible to provide one dimensional array of length n.
* `wcpp::Real`: Percentage of samples used for GPD fit estimate (default is 20).
* `wtrunc::Float64`: Positive parameter for truncating very large weights to n^wtrunc. Providing False or 0 disables truncation. Default values is 3/4.
# Returns
* `lw_out::AbstractArray`: Smoothed log weights
* `kss::AbstractArray`: Pareto tail indices
"""
function psislw(lw::Union{AbstractArray, Mamba.Chains},
                wcpp::Int64=20, wtrunc::Float64=3/4)

    if isa(lw, Mamba.Chains)
        lw_out = Mamba.combine(lw)
    else
        lw_out = copy(lw)
    end

    if ~(1 <= ndims(lw_out) <= 2)
        throw(DimensionMismatch("Argument `lw` must be 1 or 2 dimensional."))
    end
    if size(lw_out,1) <= 1
        error("More than one log-weight needed.")
    end
    return _psislw(lw_out, wcpp, wtrunc)
end

function _psislw(lw_out::Array{Float64}, wcpp::Int64, wtrunc::Float64)

    if ndims(lw_out) == 2
        n, m = size(lw_out)
    elseif ndims(lw_out) == 1
        n = length(lw_out)
        m = 1
    end
    kss = zeros(Float64, m)

    # precalculate constants
    cutoffmin = log(realmin(Float64))
    logn = log(n)

    # loop over sets of log weights
    for i in 1:m
        x = view(lw_out, :, i)
        x .-= maximum(x)
        xcutoff = max(quantile(x, 1 - wcpp/100), cutoffmin)
        expxcutoff = exp(xcutoff)
        x2 = x[x .> xcutoff]
        n2 = length(x2)
        if n2 <= 4
            # not enough tail samples for gpdfitnew
            k = Inf
        else
            # order of tail samples
            x2si = sortperm(x2)
            # fit generalized Pareto distribution to the right tail samples
            x2 .= exp.(x2) .- expxcutoff
            k, sigma = gpdfitnew(x2)
            # compute ordered statistic for the fit
            sti = collect(0.5:n2) / n2
            qq = gpinv(sti, k, sigma)
            qq .= log.(qq .+ expxcutoff)
            # place the smoothed tail into the output array
            x[x .> xcutoff] = qq[x2si]
        end
        if wtrunc > 0
            # truncate too large weights
            lwtrunc = wtrunc * logn - logn + logsumexp(x)
            x[x .> lwtrunc] = lwtrunc
        end
        # renormalize weights
        x .-= logsumexp(x)
        # store tail index k
        kss[i] = k
    end
    # If the provided input array is one dimensional, return kss as scalar.
    if ndims(lw_out) == 1
        kss = kss[1]
    end
    return lw_out, kss
end

"""
    gpdfitnew(x)
Estimate the paramaters for the Generalized Pareto Distribution (GPD). Returns empirical Bayes estimate for the parameters of the two-parameter generalized Parato distribution given the data.
# Arguments
* `x::AbstractArray`: One dimensional data array.
# Returns
* `k::Float64, sigma::Float64`: Estimated parameter values.
# Notes
* This function returns a negative of Zhang and Stephens's k, because it is more common parameterization.
"""
function gpdfitnew{T<:Real}(x::AbstractArray{T})
    if ndims(x) != 1 || length(x) <= 1
        throw(DimensionMismatch())
    end
    sort!(x)

    n = length(x)
    m = 80 + trunc(Int, sqrt(n))
    prior = 3

    bs = collect(1:m) - 0.5
    bs .= (1 .- sqrt.(m ./ bs)) ./ prior ./ x[trunc(Int64, n/4 + 0.5)] .+ 1 / x[end]

    ks = mean(log1p.(-bs .* x'), 2)[:,1]
    L = n * (log.(bs ./ -ks) .- ks .- 1)
    w = similar(L)
    w .= (1 ./ sum(exp.(L .- L'), 1)')[:,1]

    # remove negligible weights
    dii = w .>= (10 * eps(Float64))
    if ~all(dii)
        w = w[dii]
        bs = bs[dii]
    end
    # normalise w
    w ./= sum(w)
    # posterior mean for b
    b = sum(bs .* w)
    # Estimate for k, note that we return a negative of Zhang and
    # Stephens's k, because it is more common parameterisation.
    k = mean(log1p.(-b .* x))
    # estimate for sigma
    sigma = -k / b

    return k, sigma

end

"""
    gpinv(p, k, sigma)
Inverse Generalised Pareto distribution function.
"""
function gpinv(p, k, sigma)
    x = similar(p)
    if sigma <= 0
        return x
    end
    ok = (p .> 0) & (p .< 1)
    pv = view(p, ok)
    xv = view(x,ok)
    if abs(k) < eps(Float64)
        xv .= -log1p.(-pv)
    else
        xv .= expm1.(log1p.(-pv) .* (-k)) ./ k
    end
    x .*= sigma
    if ~all(ok)
        x[p .== 0] = 0
        if k >= 0
            x[p .== 1] = Inf
        else
            x[p .== 1] = -sigma / k
        end
    end
    x
end

"""
    logsumexp(x[, d])
Compute `log(sum(exp(x), d))` of `x`.
"""
function logsumexp{T<:Real}(x::AbstractArray{T}, d::Int64=1)
    u = maximum(x)
    abs(u) == Inf && return any(isnan, x) ? T(NaN) : u
    result = u .+ log.(sum(exp.(x-u), d))
    if length(result) == 1
        result = result[1]
    end
    result
end

end


# This code calculates the total energies for all atoms Z = 1..92
# in the LDA(SVWN) DFT approximation and compares them with
# the NIST data available at https://www.nist.gov/document/dftdata-targz

const ENIST = [-0.445671;     -2.834836;     -7.335195;    -14.447209;    -24.344198;
              -37.425749;    -54.025016;    -74.473077;    -99.099648;   -128.233481;
             -161.440060;   -199.139406;   -241.315573;   -288.198397;   -339.946219;
             -396.716081;   -458.664179;   -525.946195;   -598.200590;   -675.742283;
             -758.679275;   -847.277216;   -941.678904;  -1042.030238;  -1148.449372;
            -1261.093056;  -1380.091264;  -1505.580197;  -1637.785861;  -1776.573850;
            -1921.846456;  -2073.807332;  -2232.534978;  -2398.111440;  -2570.620700;
            -2750.147940;  -2936.337293;  -3129.453161;  -3329.520604;  -3536.737751;
            -3751.196175;  -3973.013235;  -4202.188857;  -4438.981228;  -4683.301031;
            -4935.368406;  -5195.031215;  -5462.390982;  -5737.309064;  -6019.953353;
            -6310.376268;  -6608.631413;  -6914.773092;  -7228.856107;  -7550.557710;
            -7880.111578;  -8217.575230;  -8563.360285;  -8917.664369;  -9280.311037;
            -9651.484134; -10031.259090; -10419.710775; -10816.653877; -11222.941975;
           -11637.869664; -12061.770549; -12494.718304; -12936.786494; -13388.048594;
           -13848.230375; -14317.493720; -14795.884202; -15283.448822; -15780.236024;
           -16286.295408; -16801.677471; -17326.576377; -17860.790943; -18404.274220;
           -18956.957627; -19518.993145; -20090.414449; -20671.256539; -21261.555215;
           -21861.346869; -22470.319655; -23088.688083; -23716.496952; -24353.832231;
           -25001.291382; -25658.417889];

using AtomEnergyLevels
using Test, Printf, Statistics

function test_all(N = length(ENIST))
  E = zeros(N)
  it = zeros(N)
  for at_number=1:N
    element = keys(atom)[at_number]
    @info "Calculating $element:"
    results = lda(at_number, -35:0.1:5, Α = 1e5)
    E[at_number] = results.energy.total
    it[at_number] = results.iterations
    @info @sprintf("Energy levels for the %s atom are:", element)
    for (quantum_numbers, ψ) in sort(collect(results.orbitals), by = x -> last(x).ϵᵢ)
      nᵣ, l = quantum_numbers
      n = nᵣ + l + 1
      @info @sprintf("\t%i%s\t(%4.1f)\t%14.6f", n, shells[l + 1], ψ.nᵢ, ψ.ϵᵢ)
    end
    @info @sprintf("ΔE = %12.6f", ENIST[at_number]-E[at_number])
  end
  return E, N, it
end

function summary(E, N, it)
  # Median absolute deviation
  MAD(X) = median(abs.(X .- median(X)))

  @info "== RESULTS SUMMARY =="
  for Z=1:N
    element = keys(atom)[Z]
    @info @sprintf("%2s: ΔE = %12.6f", element, ENIST[Z]-E[Z])
  end

  err, atn = findmax(abs.(ENIST[1:N] .- E))
  el = keys(atom)[atn]

  @info "== STATISTICS =="
  @info @sprintf("Median absolute deviation:\t%4.1e", MAD(ENIST[1:N] .- E))
  @info @sprintf("Maximum absolute deviation:\t%4.1e is for %2s", err, el)
  @info @sprintf("Average number of iterations:\t%i", mean(it))
end

summary(test_all(18)...)