using Test
using Printf: @sprintf
using StaticArrays: SVector
using LinearAlgebra: dot
using QuantumLattices.Essentials.DegreesOfFreedom
using QuantumLattices.Essentials.QuantumOperators: ID, sequence, id, idtype
using QuantumLattices.Essentials.Spatials: AbstractPID, PID, CPID, Point, Bond, Bonds, Lattice, pidtype, rcoord, icoord
using QuantumLattices.Essentials: kind, update!, reset!
using QuantumLattices.Interfaces: decompose, rank, expand, ⊗
using QuantumLattices.Prerequisites: Float, decimaltostr
using QuantumLattices.Prerequisites.Traits: parameternames, isparameterbound, contentnames, getcontent
using QuantumLattices.Prerequisites.CompositeStructures: NamedContainer

import QuantumLattices.Essentials.DegreesOfFreedom: ishermitian, statistics, latexname, script, couplingcenters, abbr
import QuantumLattices.Prerequisites.VectorSpaces: shape, ndimshape

struct DID{N<:Union{Int, Symbol}} <: SimpleIID
    nambu::N
end
@inline Base.adjoint(sl::DID{Int}) = DID(3-sl.nambu)
@inline statistics(::Type{<:DID}) = :f

struct DFock <: SimpleInternal{DID{Int}}
    nnambu::Int
end
@inline shape(f::DFock) = (1:f.nnambu,)
@inline ndimshape(::Type{DFock}) = 1
@inline DID(i::CartesianIndex, vs::DFock) = DID(i.I...)
@inline CartesianIndex(did::DID{Int}, vs::DFock) = CartesianIndex(did.nambu)
@inline shape(iidspace::IIDSpace{DID{Symbol}, DFock}) = (1:iidspace.internal.nnambu,)
@inline shape(iidspace::IIDSpace{DID{Int}, DFock}) = (iidspace.iid.nambu:iidspace.iid.nambu,)
@inline shape(iidspace::IIDSpace{DID{Symbol}, DFock}) = (1:iidspace.internal.nnambu,)
@inline shape(iidspace::IIDSpace{DID{Int}, DFock}) = (iidspace.iid.nambu:iidspace.iid.nambu,)

function Base.angle(id::OID{<:Index{<:AbstractPID, DID{Int}}}, vectors::AbstractVector{<:AbstractVector{Float}}, values::AbstractVector{Float})
    phase=  (length(vectors) == 1) ? 2pi*dot(decompose(id.icoord, vectors[1]), values) :
            (length(vectors) == 2) ? 2pi*dot(decompose(id.icoord, vectors[1], vectors[2]), values) :
            (length(vectors) == 3) ? 2pi*dot(decompose(id.icoord, vectors[1], vectors[2], vectors[3]), values) :
            error("angle error: not supported number of input basis vectors.")
    (id.index.iid.nambu == 1) ? phase : -phase
end
@inline script(::Val{:site}, index::Index{<:AbstractPID, <:DID}; kwargs...) = index.pid.site
@inline script(::Val{:nambu}, index::Index{<:AbstractPID, <:DID}; kwargs...) = index.iid.nambu==2 ? "\\dagger" : ""

latexname(::Type{<:OID{<:Index{<:AbstractPID, <:DID}}}) = Symbol("OID{Index{AbstractPID, DID}}")
latexformat(OID{<:Index{<:AbstractPID, <:DID}}, LaTeX{(:nambu,), (:site,)}('d'))

const DCoupling{V, I<:ID{DID}, C<:Subscripts, CI<:SubscriptsID} = Coupling{V, I, C, CI}
@inline Base.repr(tc::(Coupling{V, <:ID{DID}} where V)) = @sprintf "%s ph(%s)" decimaltostr(tc.value) join(tc.cid.nambus, ", ")
@inline couplingcenters(::(Coupling{V, <:ID{DID}} where V), ::Bond, ::Val) = (1, 2)
@inline DCoupling(value, nambus::Tuple{Vararg{Int}}) = Coupling(value, ID(DID, nambus), Subscripts((nambu=Subscript(nambus),)))
@inline DCoupling(value, nambus::Subscript) = Coupling(value, ID(DID, convert(Tuple, nambus)), Subscripts((nambu=nambus,)))

@inline abbr(::Type{<:Term{:Mu}}) = :mu
@inline abbr(::Type{<:Term{:Hp}}) = :hp
@inline ishermitian(::Type{<:Term{:Mu}}) = true
@inline ishermitian(::Type{<:Term{:Hp}}) = false

@testset "IID" begin
    did = DID(1)
    @test statistics(did) == statistics(typeof(did)) == :f

    did₁, did₂ = DID(1), DID(2)
    ciid = CompositeIID(did₁, did₂)
    @test length(ciid) == length(typeof(ciid)) == 2
    @test rank(ciid) == rank(typeof(ciid)) == 2
    @test iidtype(ciid, 1) == iidtype(typeof(ciid), 1) == DID{Int}
    @test iidtype(ciid, 2) == iidtype(typeof(ciid), 2) == DID{Int}
    @test ciid[1]==did₁ && ciid[2]==did₂
    @test ciid.nambus == (1, 2)
    @test string(ciid) == "DID(1) ⊗ DID(2)"
    @test did₁⊗did₂ == ciid
    @test did₁⊗ciid == CompositeIID(did₁, did₁, did₂)
    @test ciid⊗did₁ == CompositeIID(did₁, did₂, did₁)
    @test ciid⊗ciid == CompositeIID(did₁, did₂, did₁, did₂)
end

@testset "SimpleInternal" begin
    it = DFock(2)
    @test it|>eltype == DID{Int}
    @test it|>typeof|>eltype == DID{Int}
    @test it == deepcopy(it)
    @test isequal(it, deepcopy(it))
    @test it|>string == "DFock(nnambu=2)"
    @test it|>collect == [DID(1), DID(2)]
    @test statistics(it) == statistics(typeof(it)) == :f
    @test match(DID(1), it) && match(DID, DFock)
    @test filter(DID(1), it) == filter(DID, it) == it
    @test filter(DID(1), DFock) == filter(DID, DFock) == DFock
end

@testset "CompositeInternal" begin
    it₁, it₂ = DFock(2), DFock(3)
    ci = CompositeInternal(it₁, it₂)
    @test eltype(ci) == eltype(typeof(ci)) == CompositeIID{Tuple{DID{Int}, DID{Int}}}
    @test rank(ci) == rank(typeof(ci)) == 2
    @test string(ci) == "DFock(nnambu=2) ⊗ DFock(nnambu=3)"
    @test shape(ci) == (1:2, 1:3)
    @test ndimshape(ci) == 2
    for i = 1:2, j = 1:3
        @test CartesianIndex(CompositeIID(CartesianIndex(i, j), ci), ci) == CartesianIndex(i, j)
    end
    @test it₁⊗it₂ == ci
    @test it₁⊗ci == CompositeInternal(it₁, it₁, it₂)
    @test ci⊗it₁ == CompositeInternal(it₁, it₂, it₁)
    @test ci⊗ci == CompositeInternal(it₁, it₂, it₁, it₂)
    @test filter(DID(1), ci) == filter(DID, ci) == ci
    @test filter(DID(1), typeof(ci)) == filter(DID, typeof(ci)) == typeof(ci)
end

@testset "Index" begin
    index = Index(CPID('S', 4), DID(1))
    @test index|>pidtype == CPID{Char}
    @test index|>typeof|>pidtype == CPID{Char}
    @test index|>iidtype == DID{Int}
    @test index|>typeof|>iidtype == DID{Int}
    @test index|>adjoint == Index(CPID('S', 4), DID(2))
    @test statistics(index) == statistics(typeof(index)) == :f
    @test ishermitian(ID(index', index)) == true
    @test ishermitian(ID(index, index)) == false
end

@testset "OID" begin
    @test contentnames(CompositeOID) == (:index,)
    @test contentnames(OID) == (:index, :rcoord, :icoord)

    oid = OID(Index(PID(1), DID(1)), [0.0, -0.0], [0.0, 0.0])
    @test indextype(oid) == indextype(typeof(oid)) == Index{PID, DID{Int}}
    @test statistics(oid) == statistics(typeof(oid)) == :f
    @test oid' == OID(Index(PID(1), DID(2)), rcoord=SVector(0.0, 0.0), icoord=SVector(0.0, 0.0))
    @test hash(oid, UInt(1)) == hash(OID(Index(PID(1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 1.0)), UInt(1))
    @test propertynames(ID(oid)) == (:indexes, :rcoords, :icoords)
    @test fieldnames(OID) == (:index, :rcoord, :icoord)
    @test string(oid) == "OID(Index(PID(1), DID(1)), [0.0, 0.0], [0.0, 0.0])"
    @test ID(oid', oid)' == ID(oid', oid)
    @test ishermitian(ID(oid', oid)) == true
    @test ishermitian(ID(oid, oid)) == false
    @test oidtype(DFock, Point{2, PID, Float}, Val(:default)) == OID{Index{PID, DID{Int}}, SVector{2, Float}}
end

@testset "Operator" begin
    opt = Operator(1.0im, ID(
        OID(Index(PID(2), DID(2)), SVector(1.0, 0.0), SVector(2.0, 0.0)),
        OID(Index(PID(1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 0.0))
        ))
    @test opt' == Operator(-1.0im, ID(
        OID(Index(PID(1), DID(2)), SVector(0.0, 0.0), SVector(0.0, 0.0)),
        OID(Index(PID(2), DID(1)), SVector(1.0, 0.0), SVector(2.0, 0.0))
        ))
    @test ishermitian(opt) == false
    @test string(opt) == "Operator(1.0im, ID(OID(Index(PID(2), DID(2)), [1.0, 0.0], [2.0, 0.0]), OID(Index(PID(1), DID(1)), [0.0, 0.0], [0.0, 0.0])))"

    opt = Operator(1.0, ID(
        OID(Index(PID(1), DID(2)), SVector(0.5, 0.5), SVector(1.0, 1.0)),
        OID(Index(PID(1), DID(1)), SVector(0.5, 0.5), SVector(1.0, 1.0))
        ))
    @test opt' == opt
    @test ishermitian(opt) == true

    opt = Operator(1.0, ID(
        OID(Index(PID(1), DID(2)), SVector(0.5, 0.5), SVector(1.0, 1.0)),
        OID(Index(PID(1), DID(1)), SVector(0.0, 0.5), SVector(0.0, 1.0))
        ))
    @test rcoord(opt) == SVector(0.5, 0.0)
    @test icoord(opt) == SVector(1.0, 0.0)

    opt = Operator(1.0, ID(OID(Index(PID(1), DID(2)), SVector(0.5, 0.0), SVector(1.0, 0.0))))
    @test rcoord(opt) == SVector(0.5, 0.0)
    @test icoord(opt) == SVector(1.0, 0.0)
end

@testset "Operators" begin
    opt1 = Operator(1.0im, ID(
        OID(Index(PID(2), DID(2)), SVector(1.0, 0.0), SVector(0.0, 0.0)),
        OID(Index(PID(1), DID(1)), SVector(0.0, 0.0), SVector(2.0, 0.0))
        ))
    opt2 = Operator(1.0, ID(
        OID(Index(PID(1), DID(2)), SVector(0.0, 0.0), SVector(0.0, 0.0)),
        OID(Index(PID(1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 0.0))
        ))
    opts = Operators(opt1, opt2)
    @test summary(opts) == "Operators{$(opts|>valtype)}"
    @test opts' == Operators(opt1', opt2')
    @test opts'+opts == Operators(opt1, opt1', opt2*2)
    @test ishermitian(opts) == false
    @test ishermitian(opts'+opts) == true
end

@testset "Hilbert" begin
    map = pid->DFock((pid.site-1)%2+1)
    hilbert = Hilbert(map, [CPID(1, 1), CPID(1, 2)])
    @test hilbert == Hilbert{DFock}(map, [CPID(1, 1), CPID(1, 2)])
    @test convert(Dict, hilbert) == Dict(CPID(1, 1)=>DFock(1), CPID(1, 2)=>DFock(2))
    reset!(hilbert, (CPID(2, 1), CPID(2, 2)))
    @test convert(Dict, hilbert) == Dict(CPID(2, 1)=>DFock(1), CPID(2, 2)=>DFock(2))

    hilbert = Hilbert(pid=>DFock(2) for pid in [PID(1), PID(2)])
    @test hilbert[PID(1)] == hilbert[PID(2)] == DFock(2)

    hilbert = Hilbert(PID(1)=>DFock(2), PID(2)=>DFock(3))
    @test hilbert[PID(1)]==DFock(2) && hilbert[PID(2)]==DFock(3)
end

@testset "IIDSpace" begin
    DID₁, DID₂, it = DID(2), DID(:σ), DFock(2)
    iidspace = IIDSpace(DID₁⊗DID₂, it⊗it)
    @test eltype(iidspace) == eltype(typeof(iidspace)) == CompositeIID{Tuple{DID{Int}, DID{Int}}}
    @test kind(iidspace) == kind(typeof(iidspace)) == :info
    @test shape(iidspace) == (2:2, 1:2)
    @test ndimshape(iidspace) == ndimshape(typeof(iidspace)) == 2
    for i = 1:length(iidspace)
        iid = iidspace[i]
        @test iidspace[CartesianIndex(iid, iidspace)] == iid
    end
    @test expand((DID₁, DID₂), (it, it)) == iidspace
    @test collect(iidspace) == [DID(2)⊗DID(1), DID(2)⊗DID(2)]
end

@testset "Metric" begin
    valtype(Metric, AbstractOID) = AbstractOID

    index = Index(CPID('S', 4), DID(1))
    oid = OID(index, SVector(0.5, 0.0), SVector(1.0, 0.0))

    m = OIDToTuple((:scope, :site, :nambu))
    @test m == OIDToTuple(:scope, :site, :nambu)
    @test isequal(m, OIDToTuple(:scope, :site, :nambu))
    @test keys(m) == keys(typeof(m)) == (:scope, :site, :nambu)
    @test filter(≠(:nambu), m) == OIDToTuple(:scope, :site)
    @test OIDToTuple(Index{PID, DID{Int}}) == OIDToTuple(:site, :nambu)
    @test OIDToTuple(OID{Index{CPID{Int}, DID{Int}}}) == OIDToTuple(:scope, :site, :nambu)
    @test OIDToTuple(Hilbert{DFock, CPID{Int}}) == OIDToTuple(:scope, :site, :nambu)
    @test m(index) == ('S', 4, 1) == m(oid)
    @test filter(≠(:scope), m)(index) == (4, 1)
    @test filter(≠(:nambu), m)(index) == ('S', 4)
    @test filter(∉((:site, :nambu)), m)(index) == ('S',)
end

@testset "Table" begin
    @test contentnames(Table) == (:by, :contents)

    by = filter(≠(:nambu), OIDToTuple(Index{PID, DID{Int}}))

    table = Table([Index(PID(1), DID(1)), Index(PID(1), DID(2))], by)
    @test empty(table) == Table{Index{PID, DID{Int}}}(by)
    @test table[Index(PID(1), DID(1))]==1 && table[Index(PID(1), DID(2))]==1

    hilbert = Hilbert(pid=>DFock(2) for pid in [PID(1), PID(2)])
    inds1 = (Index(PID(1), iid) for iid in DFock(2))|>collect
    inds2 = (Index(PID(2), iid) for iid in DFock(2))|>collect
    @test Table(hilbert, by) == Table([inds1; inds2], by)
    @test Table(hilbert, by) == union(Table(inds1, by), Table(inds2, by))

    opt = Operator(1.0im, ID(
        OID(Index(PID(1), DID(2)), SVector(0.0, 0.0), SVector(1.0, 0.0)),
        OID(Index(PID(1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 0.0))
        ))
    @test sequence(opt, table) == (1, 1)
    @test haskey(table, opt.id) == (true, true)

    table = Table(hilbert)
    @test table == Table([inds1; inds2])
    @test reset!(empty(table), [inds1; inds2]) == table
    @test reset!(empty(table), hilbert) == table
end

@testset "Subscript" begin
    sub = Subscript(4)
    @test contentnames(typeof(sub)) == (:contents, :rep, :constraint)
    @test getcontent(sub, :contents) == sub.pattern
    @test sub==deepcopy(sub) && isequal(sub, deepcopy(sub))
    @test string(sub) == "[* * * *]"
    @test rank(sub) == rank(typeof(sub)) == 4
    @test match(sub, (2, 2, 2, 2))

    sub = Subscript((1, 2, 3, 4))
    @test string(sub) == "[1 2 3 4]"
    @test sub == subscript"[1 2 3 4]" == subscript"[1 2 3 4]; false"
    @test rank(sub) == 4
    @test match(sub, (1, 3, 3, 4))

    sub = Subscript((1, 2, 2, 1), true)
    @test string(sub) == "[1 2 2 1]"
    @test sub == subscript"[1 2 2 1]; true"
    @test rank(sub) == 4
    @test match(sub, (1, 2, 2, 1)) && !match(sub, (1, 1, 2, 1))

    sub = subscript"[x₁ x₂ x₁ x₂]"
    @test string(sub) == "[x₁ x₂ x₁ x₂]"
    @test rank(sub) == 4
    @test match(sub, (1, 2, 1, 2)) && !match(sub, (1, 2, 2, 1))

    sub = subscript"[x₁ 4 4 x₂](x₁ < x₂)"
    @test string(sub) == "[x₁ 4 4 x₂](x₁ < x₂)"
    @test rank(sub) == rank(typeof(sub)) == 4
    @test match(sub, (1, 4, 4, 2)) && !match(sub, (2, 4, 4, 1))
end

@testset "Subscripts" begin
    subscripts = Subscripts((nambu=subscript"[x y](x > y)",), (nambu=subscript"[x x y y](x < y)",))
    @test string(subscripts) == "nambu[x y](x > y) × nambu[x x y y](x < y)"
    @test repr(subscripts, 1:2, :nambu) == "[x y](x > y)×[x x y y](x < y)"
    @test repr(subscripts, 1, :nambu) == "[x y](x > y)"
    @test repr(subscripts, 2, :nambu) == "[x x y y](x < y)"

    @test rank(subscripts) == rank(typeof(subscripts)) == 6
    @test rank(subscripts, 1) == rank(typeof(subscripts), 1) == 2
    @test rank(subscripts, 2) == rank(typeof(subscripts), 2) == 4
    @test match(subscripts, (DID(2), DID(1), DID(2), DID(2), DID(3), DID(3)))
    @test match(subscripts, DID(2)⊗DID(1)⊗DID(2)⊗DID(2)⊗DID(3)⊗DID(3))
    @test !match(subscripts, DID(1)⊗DID(2)⊗DID(2)⊗DID(2)⊗DID(3)⊗DID(3))
    @test !match(subscripts, DID(2)⊗DID(1)⊗DID(3)⊗DID(3)⊗DID(2)⊗DID(2))
    @test subscripts == Subscripts((nambu=subscript"[x y](x > y)",))*Subscripts((nambu=subscript"[x x y y](x < y)",))

    subscriptsid = SubscriptsID(subscripts)
    @test subscriptsid == SubscriptsID((1:2=>("[x y](x > y)",), 3:6=>("[x x y y](x < y)",)))
    @test idtype(subscripts) == idtype(typeof(subscripts)) == SubscriptsID{NTuple{2, Pair{UnitRange{Int}, Tuple{String}}}}
end

@testset "couplings" begin
    @test parameternames(Coupling) == (:value, :cid, :subscripts, :subscriptsid)
    @test contentnames(Coupling) == (:value, :id, :subscripts)
    @test isparameterbound(Coupling, :cid, Tuple{DID{Int}}) == false
    @test isparameterbound(Coupling, :cid, ID{DID{Int}}) == true
    @test isparameterbound(Coupling, :subscripts, Subscripts{Tuple{NamedTuple{(:nambu,), Tuple{Subscript{Tuple{Int}, typeof(noconstrain)}}}}}) == false
    @test isparameterbound(Coupling, :subscripts, Subscripts) == true
    @test isparameterbound(Coupling, :subscriptsid, SubscriptsID{Tuple{}}) == false
    @test isparameterbound(Coupling, :subscriptsid, SubscriptsID) == true

    tc = DCoupling(2.0, (2,))
    @test id(tc) == ID(CompositeIID(tc.cid), SubscriptsID(tc.subscripts))
    @test rank(tc) == rank(typeof(tc)) == 1
    @test tc == Coupling(2.0, id(tc), tc.subscripts)
    @test ID{SimpleIID}(tc) == tc.cid
    @test Subscripts(tc) == tc.subscripts

    point = Point(CPID(1, 1), (0.0, 0.0), (0.0, 0.0))
    hilbert = Hilbert(point.pid=>DFock(2))
    @test couplingcenters(tc, point, Val(:Mu)) == (1,)
    @test couplingpoints(tc, point, Val(:Mu)) == (point,)
    @test couplinginternals(tc, point, hilbert, Val(:Mu)) == (DFock(2),)

    tc₁ = DCoupling(1.5, (1, 2))
    tc₂ = DCoupling(2.0, subscript"[a b](a < b)")
    @test tc₁*tc₂ == Coupling(3.0, ID(DID(1), DID(2), DID(:a), DID(:b)), Subscripts((nambu=Subscript((1, 2)),), (nambu=subscript"[a b](a < b)",)))

    ex = expand(tc₁, point, hilbert, Val(:info))
    @test eltype(ex) == eltype(typeof(ex)) == Tuple{Float64, NTuple{2, OID{Index{CPID{Int}, DID{Int}}, SVector{2, Float64}}}}
    @test collect(ex) == [
        (1.5, ID(OID(Index(CPID(1, 1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 0.0)),
                 OID(Index(CPID(1, 1), DID(2)), SVector(0.0, 0.0), SVector(0.0, 0.0))
                 ))
        ]

    ex = expand(tc₂, point, hilbert, Val(:info))
    @test eltype(ex) == eltype(typeof(ex)) == Tuple{Float64, NTuple{2, OID{Index{CPID{Int}, DID{Int}}, SVector{2, Float64}}}}
    @test collect(ex) == [
        (2.0, ID(OID(Index(CPID(1, 1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 0.0)),
                    OID(Index(CPID(1, 1), DID(2)), SVector(0.0, 0.0), SVector(0.0, 0.0))
                    ))
        ]

    @test @couplings(DCoupling(1.0, (1, 1))) == Couplings(DCoupling(1.0, (1, 1)))
    @test @couplings(DCoupling(1.0, (1,))+DCoupling(1.0, (2,))) == DCoupling(1.0, (1,))+DCoupling(1.0, (2,))
end

@testset "TermFunction" begin
    ta = TermAmplitude()
    @test ta() == 1

    ta = TermAmplitude(x->x+3.0)
    @test ta(1.0) == 4.0

    tcs = DCoupling(1.0, (1, 1)) + DCoupling(2.0, (2, 2))
    termcouplings = TermCouplings(tcs)
    @test termcouplings == deepcopy(TermCouplings(tcs))
    @test isequal(termcouplings, deepcopy(TermCouplings(tcs)))
    @test valtype(termcouplings) == valtype(typeof(termcouplings)) == typeof(tcs)
    @test termcouplings() == tcs

    fx = i -> i%2==1 ? (DCoupling(1.0, (1, 1))+DCoupling(1.0, (2, 2))) : (DCoupling(1.0, (2, 1))+DCoupling(1.0, (1, 2)))
    termcouplings = TermCouplings(fx)
    @test termcouplings == TermCouplings{typejoin(typeof(fx(1)), typeof(fx(2)))}(fx)
    @test valtype(termcouplings) == valtype(typeof(termcouplings)) == typejoin(typeof(fx(1)), typeof(fx(2)))
    @test termcouplings(1) == fx(1)
    @test termcouplings(2) == fx(2)

    @test ismodulatable(TermModulate{Val{true}}) == ismodulatable(TermModulate{<:Function}) == true
    @test ismodulatable(TermModulate{Val{false}}) == false
    termmodulate = TermModulate(:t)
    @test termmodulate(t=1) == 1
    @test isnothing(termmodulate(mu=1))
    @test isnothing(termmodulate())
    @test ismodulatable(termmodulate) == true

    termmodulate = TermModulate(:t, t->t*2.0)
    @test ismodulatable(termmodulate) == true
    @test termmodulate(2) == 4
end

@testset "Term" begin
    point = Point(CPID(1, 1), (0.0, 0.0), (0.0, 0.0))
    hilbert = Hilbert(point.pid=>DFock(2))
    term = Term{:Mu}(:mu, 1.5, 0, couplings=@couplings(DCoupling(1.0, (2, 1))), amplitude=(bond->3), modulate=false)
    @test term|>kind == term|>typeof|>kind == :Mu
    @test term|>id == term|>typeof|>id == :mu
    @test term|>valtype == term|>typeof|>valtype == Float
    @test term|>rank == term|>typeof|>rank == 2
    @test term|>abbr == term|>typeof|>abbr == :mu
    @test term|>ismodulatable == term|>typeof|>ismodulatable == false
    @test term|>ishermitian == term|>typeof|>ishermitian == true
    @test term == deepcopy(term)
    @test isequal(term, deepcopy(term))
    @test string(term) == "Mu{2}(id=mu, value=1.5, bondkind=0, factor=1.0)"
    @test repr(term, point, hilbert) == "mu: 4.5 ph(2, 1)"
    @test +term == term
    @test -term == term*(-1) == replace(term, factor=-term.factor)
    @test 2*term == term*2 == replace(term, factor=term.factor*2)
    @test term/2 == term*(1/2) == replace(term, factor=term.factor/2)

    p1 = Point(PID(1), (0.0, 0.0), (0.0, 0.0))
    p2 = Point(PID(2), (0.0, 0.0), (0.0, 0.0))
    hilbert = Hilbert(pid=>DFock(2) for pid in [PID(1), PID(2)])
    term = Term{:Mu}(:mu, 1.5, 0,
        couplings=bond->bond.pid.site%2==0 ? Couplings(DCoupling(1.0, (2, 2))) : Couplings(DCoupling(1.0, (1, 1))),
        amplitude=bond->3,
        modulate=true
    )
    @test term|>ismodulatable == term|>typeof|>ismodulatable == true
    @test repr(term, p1, hilbert) == "mu: 4.5 ph(1, 1)"
    @test repr(term, p2, hilbert) == "mu: 4.5 ph(2, 2)"
    @test one(term) == replace(term, value=1.0)
    @test zero(term) == replace(term, value=0.0)
    @test term.modulate(mu=4.0) == 4.0
    @test isnothing(term.modulate(t=1.0))
    @test update!(term, mu=4.25) == replace(term, value=4.25)
    @test term.value == 4.25
end

@testset "expand" begin
    point = Point(PID(1), (0.0, 0.0), (0.0, 0.0))
    hilbert = Hilbert(point.pid=>DFock(2))
    term = Term{:Mu}(:mu, 1.5, 0, couplings=@couplings(DCoupling(1.0, (2, 1))), amplitude=bond->3.0, modulate=true)
    operators = Operators(
        Operator(+2.25, ID(
            OID(Index(PID(1), DID(2)), SVector(0.0, 0.0), SVector(0.0, 0.0)),
            OID(Index(PID(1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 0.0))
            ))
        )
    @test expand(term, point, hilbert, half=true) == operators
    @test expand(term, point, hilbert, half=false) == operators*2

    bond = Bond(1, Point(PID(2), (1.5, 1.5), (1.0, 1.0)), Point(PID(1), (0.5, 0.5), (0.0, 0.0)))
    hilbert = Hilbert(pid=>DFock(2) for pid in [PID(1), PID(2)])
    term = Term{:Hp}(:t, 1.5, 1, couplings=@couplings(DCoupling(1.0, (2, 1))), amplitude=bond->3.0, modulate=true)
    operators = Operators(
        Operator(4.5, ID(
            OID(Index(PID(1), DID(2)), SVector(0.5, 0.5), SVector(0.0, 0.0)),
            OID(Index(PID(2), DID(1)), SVector(1.5, 1.5), SVector(1.0, 1.0))
            ))
        )
    @test expand(term, bond, hilbert, half=true) == operators
    @test expand(term, bond, hilbert, half=false) == operators+operators'

    lattice = Lattice("Tuanzi", [Point(CPID(1, 1), (0.0, 0.0), (0.0, 0.0))], vectors=[[1.0, 0.0]], neighbors=1)
    bonds = Bonds(lattice)
    hilbert = Hilbert(pid=>DFock(2) for pid in lattice.pids)
    term = Term{:Mu}(:mu, 1.5, 0, couplings=@couplings(DCoupling(1.0, (2, 1))), amplitude=bond->3.0, modulate=true)
    operators = Operators(
        Operator(+2.25, ID(
            OID(Index(CPID(1, 1), DID(2)), SVector(0.0, 0.0), SVector(0.0, 0.0)),
            OID(Index(CPID(1, 1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 0.0))
            ))
        )
    @test expand(term, bonds, hilbert, half=true) == operators
    @test expand(term, bonds, hilbert, half=false) == operators*2
end

@testset "LaTeX" begin
    @test latexname(OID) == :OID
    @test latexformat(OID{<:Index{<:AbstractPID, <:DID}}) == LaTeX{(:nambu,), (:site,)}('d')

    latex = LaTeX{(:nambu,), (:site,)}('c', vectors=(SVector(1.0, 0.0), SVector(0.0, 1.0)))
    @test superscript(latex|>typeof) == (:nambu,)
    @test subscript(latex|>typeof) == (:site,)

    oid = OID(Index(CPID('d', 1), DID(2)), SVector(0.0, 0.0), SVector(1.0, 0.0))
    @test script(Val(:rcoord), oid) == "[0.0, 0.0]"
    @test script(Val(:icoord), oid) == "[1.0, 0.0]"
    @test script(Val(:integralicoord), oid; vectors=get(latex.options, :vectors, nothing)) == "[1, 0]"
    @test script(Val(:BD), oid, latex) == 'c'
    @test script(Val(:SP), oid, latex) == ("\\dagger",)
    @test script(Val(:SB), oid, latex) == (1,)
    @test repr(oid, latex) == "c^{\\dagger}_{1}"

    opt = Operator(1.0, ID(
        OID(Index(CPID('d', 2), DID(2)), SVector(1.0, 0.0), SVector(2.0, 0.0)),
        OID(Index(CPID('d', 2), DID(2)), SVector(1.0, 0.0), SVector(2.0, 0.0))
        ))
    @test repr(opt) == "(d^{\\dagger}_{2})^2"

    opt = Operator(1.0, ID(
        OID(Index(CPID('d', 2), DID(2)), SVector(1.0, 0.0), SVector(2.0, 0.0)),
        OID(Index(CPID('d', 1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 0.0))
        ))
    @test repr(opt) == "d^{\\dagger}_{2}d^{}_{1}"

    latexformat(OID{<:Index{<:AbstractPID, <:DID}}, LaTeX{(:nambu,), (:site,)}('c'))
    @test repr(opt) == "c^{\\dagger}_{2}c^{}_{1}"

    latexformat(OID{<:Index{<:AbstractPID, <:DID}},  LaTeX{(:nambu,), (:site,)}('d'))
    @test repr(opt) == "d^{\\dagger}_{2}d^{}_{1}"

    latexformat(OID{<:Index{<:AbstractPID, <:DID}}, LaTeX{(:nambu,), (:rcoord,)}('d'))
    @test repr(opt) == "d^{\\dagger}_{[1.0, 0.0]}d^{}_{[0.0, 0.0]}"

    latexformat(OID{<:Index{<:AbstractPID, <:DID}}, LaTeX{(:nambu,), (:integralicoord,)}('d', vectors=(SVector(1.0, 0.0), SVector(0.0, 1.0))))
    @test repr(opt) == "d^{\\dagger}_{[2, 0]}d^{}_{[0, 0]}"

    opts = Operators(
            Operator(1.0-1.0im, ID(
                OID(Index(PID(2), DID(2)), SVector(0.0, 0.0), SVector(0.0, 0.0)),
                OID(Index(PID(1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 0.0))
                )),
            Operator(-1.0, ID(
                OID(Index(PID(1), DID(2)), SVector(0.0, 0.0), SVector(0.0, 0.0)),
                OID(Index(PID(1), DID(1)), SVector(0.0, 0.0), SVector(0.0, 0.0))
                ))
            )
    latexformat(OID{<:Index{<:AbstractPID, <:DID}}, LaTeX{(:nambu,), (:site,)}('c'))
    @test repr(opts) == "-c^{\\dagger}_{1}c^{}_{1}+(1.0-1.0im)c^{\\dagger}_{2}c^{}_{1}"

    latexformat(OID{<:Index{<:AbstractPID, <:DID}}, LaTeX{(:nambu,), (:site,)}('d'))
    io = IOBuffer()
    show(io, MIME"text/latex"(), opt)
    @test String(take!(io)) == "\$d^{\\dagger}_{2}d^{}_{1}\$"
    show(io, MIME"text/latex"(), opts)
    @test String(take!(io)) == "\$-d^{\\dagger}_{1}d^{}_{1}+(1.0-1.0im)d^{\\dagger}_{2}d^{}_{1}\$"
end

@testset "Boundary" begin
    opt = Operator(4.5, ID(
        OID(Index(PID(1), DID(2)), SVector(0.5, 0.5), SVector(0.0, 0.0)),
        OID(Index(PID(2), DID(1)), SVector(1.5, 1.5), SVector(1.0, 1.0))
        ))
    bound = Boundary{(:θ₁, :θ₂)}([0.1, 0.2], [[1.0, 0.0], [0.0, 1.0]])
    @test keys(bound) == keys(typeof(bound)) == (:θ₁, :θ₂)
    @test bound==deepcopy(bound) && isequal(bound, deepcopy(bound))
    @test angle(bound, opt) ≈ angle(opt.id, bound.vectors, bound.values) ≈ 0.6pi
    @test bound(opt) ≈ twist(opt, bound.vectors, bound.values) ≈ replace(opt, 4.5*exp(2im*pi*0.3))
    update!(bound, θ₁=0.3)
    @test bound(opt) ≈ replace(opt, 4.5*exp(2im*pi*0.5))

    @test angle(plain, opt) == 0
    @test plain(opt) == opt
    @test update!(plain) == plain
end


abstract type AbstractRaggedMatrix{T,I} <: AbstractMatrix{T} end

# struct RaggedMatrix{T,I} <: AbstractRaggedMatrix{T,I}
#     data::Vector{T} # length(data) == sum(column_lengths) == column_offsets[end]
#     column_offsets::Vector{I} # length(column_offsets) == size(A,2)
#     column_lengths::Vector{I} # length(column_lengths) == size(A,2)
#     nrow::Int
# end

# Column offsets go from 0 (col 1 offset), col 2 offset,...,2nd last col offset
struct RaggedMatrix{T,I,VI<:AbstractVector{I},VT<:AbstractVector{T}} <: AbstractRaggedMatrix{T,I}
    data::VT # length(data) == sum(column_lengths) == column_offsets[end]
    column_offsets::VI # length(column_offsets) == size(A,2)
    column_lengths::VI # length(column_lengths) == size(A,2)
    nrow::Int
end
struct FixedSizeRaggedMatrix{M,N,P,T,I} <: AbstractRaggedMatrix{T,I}
    data::FixedSizeVector{P,T,P}
    column_offsets::FixedSizeVector{N,I,N}
    column_lengths::FixedSizeVector{N,I,N}
end
function FixedSizeRaggedMatrix{M,N,P}(
    data::AbstractVector{T},
    column_offsets::AbstractVector{I},
    column_lengths::AbstractVector{I}
) where {M,N,P,T,I}
    FixedSizeRaggedMatrix{M,N,P,T,I}(
        FixedSizeVector{P,T,P}(data),
        FixedSizeVector{N,I,N}(column_offsets),
        FixedSizeVector{N,I,N}(column_lengths)
    )
end
# const FixedSizeRaggedMatrix{T,I,NC,ND} = RaggedMatrix{T,I,FixedSizeVector{NC,I,NC},FixedSizeVector{ND,T,ND}}
nrow(A::AbstractRaggedMatrix) = A.nrow
ncol(A::AbstractRaggedMatrix) = length(A.column_lengths)
nrow(::FixedSizeRaggedMatrix{M}) where {M} = M
ncol(::FixedSizeRaggedMatrix{M,N}) where {M,N} = N

Base.size(A::AbstractRaggedMatrix) = (nrow(A),ncol(A))
number_not_structurally_zero(A::AbstractRaggedMatrix) = length(A.data)
number_not_structurally_zero(::FixedSizeRaggedMatrix{M,N,P}) where {M,N,P} = P

@Base.propagate_inbounds Base.getindex(A::AbstractRaggedMatrix, i::Integer) = A.data[i]
@inline function Base.getindex(A::AbstractRaggedMatrix{T,I}, i::Integer, j::Integer) where {T,I}
    @boundscheck begin
        ( (i > nrow(A)) || ( (j > ncol(A)) || (min(i,j) < 1) ) ) && PaddedMatrices.ThrowBoundsError("Tried to index array of size $(size(A)) with index ($i,$j).")
    end
    @inbounds begin
        col_j_offset = A.column_offsets[j] #j == one(j) ? zero(I) : A.column_offsets[j-1]#; col_j_length = A.column_descriptions[j,2]
        col_j_nextoffset = j == length(A.column_offsets) ? Base.unsafe_trunc(I, length(A.data)) : A.column_offsets[j+1]
        ind = col_j_offset + Base.unsafe_trunc(I, i)
        col_j_nextoffset >= ind ? A.data[ind] : zero(T)
    end
end

### Iteration returns (row, column, value)
@inline function Base.iterate(A::AbstractRaggedMatrix)
    (1, 1, 1), (2, 2, 1)
end
@inline function Base.iterate(A::AbstractRaggedMatrix, (k,i,j)::NTuple{3,Int})
    if i > A.column_lengths[j]
        j == ncol(A) && return nothing
        i, j = 1, j+1
    end
    # k = 1,...,P; i = 1,...,M; j = 1,...,N
    (k, i, j), (k+1, i+1, j)
end

function expand_by_x_quote!(q, N, ncol, nrl, T, sp::Bool, def::Bool=true)
    push!(q.args, ind = 1)
    loop_body = quote end
    a_getindex = quote end
    return_expression = Expr(:tuple)
    for n in 1:N
        if N > 1
            asym = Symbol(:a_,n)
            push!(q.args, :($asym = @inbounds a[$n]))
        else
            asym = :a
        end
        asymi = Symbol(:a_i_,n)
        bsym = Symbol(:b_,n)
        def && push!(q.args, pointer_vector_expr( bsym, nrl, T, sp, :sptr) )
        push!(a_getindex.args, :($asymi = $asym[i]))
        push!(loop_body.args, :($bsym[ind] = $asymi))
        push!(return_expression.args, bsym)
    end
    return_expression = sp ? :((sptr,$return_expression)) : return_expression
    loop_quote = quote
        @inbounds for i in 1:$ncol
            $a_getindex
            nr = column_lengths[i]
            for r in one(nr):nr
                $loop_body
                ind += 1
            end
        end
        $return_expression
    end
    push!(q.args, loop_body)
    q
end
function contract_by_x_quote!(q, N, ncol, nrl, T, sp::Bool)
    push!(q.args, ind = 1)
    loop_body = quote end
    bset_quote = quote end
    loop_suffix = quote end
    return_expression = Expr(:tuple)
    for n in 1:N
        asym = Symbol(:a_,n)
        # asymi = Symbol(:a_i_,n)
        bsym = Symbol(:b_,n)
        bsymi = Symbol(:b_i_,n)
        push!(q.args, pointer_vector_expr( bsym, ncol, T, sp, :sptr) )
        push!(bset_quote.args, :($bsymi = zero(T)))
        push!(loop_body.args, :($bsymi += $asym[i]))
        push!(loop_suffix.args, :($bsym[ind] = $bymi))
        push!(return_expression.args, bsym)
    end
    return_expression = sp ? :((sptr,$return_expression)) : return_expression
    loop_quote = quote
        @inbounds for i in 1:$ncol
            $bset_quote
            nr = column_lengths[i]
            for r in one(nr):nr
                $loop_body
                ind += 1
            end
            $loop_suffix
        end
        $return_expression
    end
    push!(q.args, loop_body)
    q
end
function expand_contract_prequote(VI, VT)
    q = quote column_lengths = x.column_lengths end
    if VI <: AbstractFixedSizeVector
        ncol = full_length(VI)
    else
        ncol = :ncol
        push!(q.args, :($ncol = :(ncol(x))))
    end
    if VT <: AbstractFixedSizeVector
        nrl = full_length(VI)
    else
        nrl = :nrl
        push!(q.args, :($nrl = :(number_not_structurally_zero(x))))
    end
    ncol, nrl, q
end

@generated function expand_by_x(x::RaggedMatrix{T,I,VI,VT}, a::Vararg{<:AbstractVector,N}) where {N,T,I,VI,VT}
    ncol, nrl, q = expand_contract_prequote(VI, VT)
    expand_by_x_quote!(q, N, ncol, nrl, T, false)
end
@generated function contract_by_x(x::RaggedMatrix{T,I,VI,VT}, a::Vararg{<:AbstractVector,N}) where {N,T,I,VI,VT}
    ncol, nrl, q = expand_contract_prequote(VI, VT)
    contract_by_x_quote!(q, N, ncol, nrl, T, false)
end
function expand_by_x!(b::AbstractVector, x::RaggedMatrix{T,I,VI,VT}, a::AbstractVector) where {T,I,VI,VT}
    ncol, nrl, q = expand_contract_prequote(VI, VT)
    contract_by_x_quote!(q, 1, ncol, nrl, T, false, false)
end
@generated function expand_by_x(sptr::StackPointer, x::RaggedMatrix{T,I,VI,VT}, a::Vararg{<:AbstractVector,N}) where {N,I,T,VI,VT}
    ncol, nrl, q = expand_contract_prequote(VI, VT)
    expand_by_x_quote!(q, N, ncol, nrl, T, true)
end
@generated function contract_by_x(sptr::StackPointer, x::RaggedMatrix{T,I,VI,VT}, a::Vararg{<:AbstractVector,N}) where {N,I,T,VI,VT}
    ncol, nrl, q = expand_contract_prequote(VI, VT)
    contract_by_x_quote!(q, N, ncol, nrl, T, true)
end


function readlabelset(io::IO, bytecomp::Bool, bias::Integer)
    n = read(io, Int32)
    dictf = sizehint!(Dict{Float64, String}(), n)
    dictui = sizehint!(Dict{UInt8, String}(), n)
    for i = 1:n
        k = read(io, Float64)
        nchar = read(io, Int8)
        label = String(read(io, nchar))
                     # advance io to next multiple of 8 bytes
        extras = rem(nchar + 1, 8)   # + 1 b/c of reading nchar
        if extras ≠ 0
            skip(io, 8 - extras)
        end
        dictf[k] = label
        if bytecomp && isinteger(k)
            kb = Int(k) + bias
            if 0 < kb < 252
                dictui[kb] = label
            end
        end
    end
    read(io, Int32) == 4 || throw(ArgumentError("Malformed stream"))
    inds = read!(io, Vector{Int32}(read(io, Int32)))
    dictf, dictui, inds
end

function readvaluelabels!(io::IO, vdict::VariableDictionary, bytecomp::Bool, bias::Integer)
    rectyp = Int32(3)
    nms = vdict.nms
    typs = vdict.typs
    vlabs = vdict.valuelabels
    while rectyp == 3
        dictf, dictui, inds = readlabelset(io, bytecomp, bias)
        if bytecomp && length(dictui) == length(dictf)  ## switch to compressed bytes
            for j in inds
                typs[j] = UInt8
                vlabs[nms[j]] = dictui
            end
#=        else
            for j in inds
                vlabs[j] = dictf
            end
=#
        end
        rectyp = read(io, Int32)
    end
    return rectyp
end


@testset "Comparison" begin
    for mod in (StructProg, SDDP)
        @test mod._isapproxzero(0, 1e-10)
        @test !mod._isapproxzero(1, 1e-10)
        @test mod._isapproxzero(1e-5, 1e-4)
        @test !mod._isapproxzero(1e-3, 1e-4)
        @test mod._isapproxzero([1e-5, 1e-5], 1e-4)
        @test !mod._isapproxzero([1e-5, 1e-3], 1e-4)
        @test mod._isapprox(0, 0, 1e-3)
        @test mod._isapprox(1, 1, 1e-3)
        @test !mod._isapprox(0, 1, 1e-3)
        @test mod._lt(0, 1, 1e-3)
        @test !mod._lt(1, 1, 1e-3)
        @test !mod._lt(1, 0, 1e-3)
        @test !mod._lt(1, BigInt(1), 1e-3)
    end
end


# Evaluates the ITMIX runs by creating a `data` DataFrame containing summary info and, optionally, plots.
#
# NOTE: the run-plots show the mean thickness field not the one produced by using mtheta in the forward model.

error("remove Rhat")

# run this first to make sure all precompilation this through
include("itmix-setup.jl")
using FileIO, DataFrames, StatsBase
import Plots. PyPlot, Interpolations
global pl

const make_plots = false
const make_plots_thetas = make_plots
parallel = false

n_1d_btilde = 3 # number of spatial parameters
n_1d_fsl = 3 # number of spatial parameters
n_1d_temp = 1 # number of spatial parameters
subdir = ["",
          "res-7oct-2",
          "res-23-oct",
          "res-28-oct",
          "res-30-oct",
          "res-31-oct-rmiv",
          "res-03-nov-rmiv",
          "esa-14nov"
          ][end]

if !isdefined(:runit)
    runit = true
end

"""Construct an empty DataFrame"""
function make_data()
    error("remove Rhat")
    DataFrame(person = Symbol[],
              gl = Symbol[],
              prio = Int[],
              expnr = Int[],
              rm_iv = Bool[],
              nl = Int[],
              mean_h = Float64[],
              volume_area = Float64[],
              rmse = Float64[],
              rmse_fitting = Float64[],
              max = Float64[],
              max_fitting = Float64[],
              Rhat = Float64[],
              Rhat_max = Float64[],
              mtheta = Vector{Float64}[],
              err_h = VAWTools.Traj{Float64}[]
              )
end

function load_one(gid, expnr, rm_iv;
                  fit_target=BM.FitTarget.h_iv)
    gl,gb,pp,pm,pn,pl = gls[gid]
    priority = find([gid in pr for pr in I2.priorities])[1]
    fl, fls, fliv, flivs, flj = I2.filepaths(gl, pl, expnr, rm_iv, "results_Werder/"*subdir)
    dir_,fl_ = splitdir(fl)
    fl_ = splitext(fl_)[1]
    if !isfile(fl)
        println("File $fl not found")
        return nothing
    end

    jldloaded, jldd = try
        true, FileIO.load(flj)
    catch e
        println("File $flj not loading with $e")
        false, nothing
    end

    if jldloaded
        thetas, Rhat, sample_size = map(k->jldd[k], ("thetas","Rhat","sample_size"))
        mtheta = squeeze(mean(thetas,2),2)

        gln = BM.ITMIX2.make_ITMIX2_glacier(gl, expnr)
        gbn = BM.Bands(gb, gl=gln)
        th0d = BM.theta0_dict_defaults(pm, n_1d_btilde, n_1d_fsl, n_1d_temp)
        if length(mtheta)==8
            # make it compatible with older runs
            pop!(th0d, :iv_dist_exp2)
        end
        (theta0, logposterior, logposterior1d, logprior, logprior_debug,
         loglikelihood, loglikelihood1d, fwdm_fn, fwdm1d_fn,
         pmcmc_defaults, fit_target) =
             BM.init_inverse(gbn, pp, pm, pn,
                             theta0_dict=th0d,
                             fit_target=fit_target)
        fwdsol = fwdm_fn(theta0.th0)
    end

    r_lines_fitting = I2.get_rlines(gid, expnr)
    hs2d = VAWTools.read_agr(fl)
    hs2ds = VAWTools.read_agr(fls)
    iv2d = try VAWTools.read_agr(fliv) end
    iv2ds = try VAWTools.read_agr(flivs) end
    if gid==BITEModel.ITMIXGlacier(:Columbia, 2) || gid==BITEModel.ITMIXGlacier(:Elbrus, 2)
        # downsample agr
        hs2d = downsample(hs2d, 2)
        hs2ds = downsample(hs2ds, 2)
    end

    # tuples (control, fitting, all)
    RMSE, max_err, quantiles, smaller_20, smaller_std, smaller_2std, err_h = BM.compare_to_radar(gl, hs2d, hs2ds, r_lines_fitting)

    mean_ice_thickness, volume_area = VAWTools.meannan(hs2d.v), BM.volume_area_mean_h(gl)

    data_out = if jldloaded
        [:Werder, gid.id, priority, expnr, rm_iv, length(r_lines_fitting),
         map(x->signif(x,3), (mean_ice_thickness, volume_area, RMSE[1], RMSE[2], (max_err[1]), (max_err[2]), (mean(Rhat)), (maximum(Rhat))))...,
         mtheta, err_h]
    else
        [:Werder, gid.id, priority, expnr, rm_iv, length(r_lines_fitting),
         map(x->signif(x,3), (mean_ice_thickness, volume_area, RMSE[1], RMSE[2], (max_err[1]), (max_err[2]), 0.0, 0.0))...,
         Float64[], err_h]
    end

    fwdsol = fwdm_fn(mtheta);

    return data_out, hs2d, hs2ds, iv2d, iv2ds, r_lines_fitting, dir_, fl_, fwdsol, jldd, mtheta, thetas, gbn
end

# """
# Evaluate one ITMIX2 experiment and store/save relevant input.

# Returns the data to be appended to data-DataFrame
# """
@everywhere function evaluate_one(expnr, rm_iv, gb, pp, pm, pn, pl;
                                  n_1d_btilde=3, n_1d_fsl=3, n_1d_temp=1,
                                  make_plots = false,
                                  make_plots_thetas = make_plots,
                                  subdir="",
                                  fit_target=BM.FitTarget.h_iv )
    gl = gb.gl
    gid = gl.gid
    (data_out, hs2d, hs2ds, iv2d, iv2ds, r_lines_fitting, dir_,
     fl_, fwdsol, jldd, mtheta, thetas, gbn) = load_one(gid, expnr, rm_iv, fit_target=fit_target)

    # plot it
    if make_plots
        BM.plot2d_h(gl, hs2d.v; clims=:radar, plotline=I2.get_rlines(gid, expnr), reuse=true,
                    outlines=true)
        Plots.savefig(joinpath(dir_, "plots", "$(fl_)-hs2d.png"))
        BM.plot2d_h(gl, hs2ds.v; clims=:model, plotline=I2.get_rlines(gid, expnr), plottraj=false,
                    colorbar_title="h std (m)", reuse=true, outlines=true)
        Plots.savefig(joinpath(dir_, "plots", "$(fl_)-hs2d-std.png"))
        if iv2d!=nothing
            BM.plot2d_h(gl, iv2d.v; plotiv=true, reuse=true, clims=gl.iv.vals isa VAWTools.Traj ? :radar : :model,
                        outlines=true)
            Plots.savefig(joinpath(dir_, "plots", "$(fl_)-iv2d.png"))
            if !(gl.iv isa BM.SomeData{Void})
                if gl.iv.vals isa VAWTools.Gridded
                    err = iv2d.v-gl.iv.vals.v
                    err[.!gl.ivmask] = NaN
                    mi, ma = VAWTools.extremanan(err)
                    clim = max(abs(mi), abs(ma)) # set clims symmetric because of colorscale
                    BM.plot2d_h(gl, err; plotiv=true, colorbar_title="iv model-data (m)",
                                seriescolor=:pu_or, clims=(-clim,clim), reuse=true, outlines=true)
                    Plots.savefig(joinpath(dir_, "plots", "$(fl_)-iv2d-err.png"))
                elseif gl.iv.vals isa VAWTools.Traj
                    #BM.plot2d_h(gl, iv2d.v; clims=:radar, plotiv=true, colorbar_title="iv model-data (m)")
                else
                    error()
                end
            end
            BM.plot2d_h(gl, iv2ds.v; plotiv=true, colorbar_title="iv std (m)", reuse=true,
                        plottraj=false, clims=:model, outlines=true)
            Plots.savefig(joinpath(dir_, "plots", "$(fl_)-iv2d-std.png"))
        end
        # this isn't quite correct as, e.g., hs1d of mtheta is not equal to expected hs1d
        #
        # NOTE this plots the mean hs2d, but the others are from running the model with mean paras
        BM.plot_run(fwdsol,
                    reuse=true, plot2d=[:h,:iv,:ivlog][1], clims=[:model,:radar][2],
                    plotline=I2.get_rlines(gid, expnr),
                    radarline_gl=gl)
        Plots.savefig(joinpath(dir_, "plots", "$(fl_)-run.png"))
    end
    if make_plots_thetas
        names = jldd["theta0_names"]
        theta0 = BM.Theta0{Float64}(mtheta, names,
                                    OrderedDict{Symbol,LinSpace{Float64}}(), gbn, pp, pm, pn)
        BM.plottheta(thetas, theta0, fsize=(1000,1000), reuse=true)
        Plots.savefig(joinpath(dir_, "plots", "$(fl_)-thetas.png"))
    end
    # PyPlot.close("all") does not help
    return data_out
end

# Good queries
# data[data[:gl].==:Unteraar,:]
# by(data, :gl, d->DataFrame( mean_rmse = mean(d[:rmse]), max_err = maximum(d[:max]), runs=length(d[:expnr])))

function summarize_data(data; expnrs=0:16, person=false)
    d = data[ in.(data[:expnr], (expnrs,)), :]
    bycols = person ? [:person, :gl] : [:gl]
    by(d, bycols, d->DataFrame(prio = round(Int, mean(d[:prio])),
                               median_mean_h =  median(d[:mean_h]),
                               quantiles_mean_h =  (summarystats(d[:mean_h]).q25, summarystats(d[:mean_h]).q75),
                               extrema_mean_h =  extrema(d[:mean_h]),
                               mean_rmse = mean(d[:rmse]),
                               rel_mean_rmse = mean(d[:rmse])/median(d[:mean_h]),
                               mean_max_err = mean(d[:max]),
                               maximum_max_err = maximum(d[:max]),
                               runs=length(d[:expnr]),
                               Rhat_median=median(d[:Rhat_max]),
                               Rhat_max=maximum(d[:Rhat_max])
                               ),
       sort = true)

    # Easier:
    # flds = [:gl, :rmse, :rmse_fitting, :max, :max_fitting, :Rhat, :Rhat_max]
    # sort(aggregate(d[flds], :gl, [mean, median, maximum]))
end

@everywhere begin
    include("itmix-setup.jl")
    using FileIO, DataFrames
    import Plots. PyPlot, Interpolations
    pl_kws = Dict{Symbol,Any}()
    PyPlot.matplotlib[:interactive](false) # memory leak https://github.com/JuliaPy/PyPlot.jl/issues/111

# global priority,gid,gl,gb,pp,pm,pn,pl,fl,fls,flj,hs2d,hs2ds,RMSE,max_err,smaller_20,thetas,theta0
# global fwdm_fn, jldd, gbn, mtheta, names

    gids = [BM.ITMIXv2_glaciers, BM.ITMIXGlacier.(BM.itmix_glaciers_iv,2)][2]
    if !isdefined(:gls) || length(gls)<length(gids)
        gls = try
            open("data/ITMIX2/results_Werder/$subdir/gls.jls", "r") do io
                deserialize(io)
            end
        catch
            gls = Dict()
            for gid in gids
                gls[gid] = BM.init_forward(gid; pl_kws...);
            end
            gls
        end
    end
end
# # load and save gls with
# open("data/ITMIX2/results_Werder/$subdir/gls.jls", "w") do io
#     serialize(io, gls)
# end
# gls = open("data/ITMIX2/results_Werder/$subdir/gls.jls", "r") do io
#     deserialize(io)
# end

if runit
    println("""
    ==================
    Running over directory $(joinpath(first(gls)[2][end].data_root,"results_Werder", subdir))
    ==================
    """)

    data = make_data()

    #gid = BITEModel.ITMIXGlacier(:Austfonna, 2)
    for priority=1:4, rm_iv in [false,true]
        println("\n\n++++++++++++++++ PRIORITY $priority ; rm_iv = $rm_iv ++++++++++++++++++++++++++\n\n")
        for gid in I2.priorities[priority] #[BM.ITMIXGlacier(:SouthGlacier,2)] #
            if !haskey(gls, gid)
                continue
            end
            gl,gb,pp,pm,pn,pl = gls[gid]

            dir_ = splitdir(I2.filepaths(gl, pl, 0, rm_iv, "results_Werder/"*subdir)[1])[1]
            !isdir(joinpath(dir_, "plots")) && mkdir(joinpath(dir_, "plots"))

            println("\n\n============================  Glacier $(BM.getname(gid))")
            #println("expnr, nl, rmse, rmse_fitting, max, Rhat, Rhat_max")
            if parallel
                out = pmap(expnr -> evaluate_one(expnr, rm_iv, gb, pp, pm, pn, pl;
                                                 make_plots=make_plots,
                                                 make_plots_thetas=make_plots,
                                                 subdir=subdir),
                           0:I2.nexps)
                [push!(data, o) for o in out if o!=nothing]
            else
                for expnr=0:I2.nexps

                println("\n\n                              Exp $expnr")
                    #try
                    data_out = evaluate_one(expnr, rm_iv, gb, pp, pm, pn, pl;
                                            make_plots=make_plots,
                                            make_plots_thetas=make_plots,
                                            subdir=subdir)
                    data_out!=nothing && push!(data, data_out)
                    # catch e
                    #     e isa InterruptException && rethrow(e)
                    #     print(e)
                # end
                end
            end
        end
    end
    # save data:
    FileIO.save(joinpath(pl.data_root,"results_Werder", subdir, "data.jld"), "data", data,
                "data_summarized", summarize_data(data)) #, "gls", gls) saving gls takes 3.5GiB!!!

    ## Load with
    # using DataFrames, FileIO, BITEModel
    # data, data_summarized = FileIO.load("data/ITMIX2/results_Werder/res-30-oct/data.jld", "data", "data_summarized")

end

# read others results

function read_others_results(gid, expnr, dir, person; prefix="thickness", thin=1, start=1)
    gl,gb,pp,pm,pn,pl = gls[gid]
    r_lines_fitting = I2.get_rlines(gid, expnr)

    expnr_ = expnr>9 ? Symbol("exp$expnr") : Symbol("exp0$expnr")
    fl = joinpath(dir, "$(prefix)_$(BM.getname(gid))_$(person)_$(expnr_).asc")
    hs2d = VAWTools.read_agr(fl)
    hs2d.v[hs2d.v.==0] = NaN
    if thin>1
        hs2d = VAWTools.downsample(hs2d, thin, start)
    end

    RMSE, max_err, quantiles, smaller_20, smaller_std, smaller_2std, err_h = BM.compare_to_radar(gl, hs2d, nothing, r_lines_fitting)
    mean_ice_thickness, volume_area = VAWTools.meannan(hs2d.v), BM.volume_area_mean_h(gl)

    data_out = [person, gid.id, -1, expnr, true, length(r_lines_fitting),
                map(x->signif(x,3), (mean_ice_thickness, volume_area, RMSE[1],
                                     RMSE[2], (max_err[1]), (max_err[2]), NaN, NaN))..., Float64[], err_h]
    return data_out
end


using Test
using LogicCircuits

@testset "Unbatch test" begin # wrapping all the statements into a function also trigers signal (4) due to `XDataset`
    for name in ["msnbc", "jester", "plants", "ad"]
        one_batch_set = dataset(twenty_datasets(name); do_shuffle=false, batch_size=-1)
        batches_set = dataset(twenty_datasets(name); do_shuffle=false, batch_size=100)
        unbatched_set = XDataset(unbatch(batches_set)...)
        for f in [train, valid, test]
            # unweighted data
            @test feature_matrix(f(unbatched_set)) == feature_matrix(f(one_batch_set))

            # weighted data
            w = rand(num_examples(f(one_batch_set)))
            w_one_batch = WXData(f(one_batch_set), w)

            w_batches = batch(w_one_batch, 100)
            w_unbatched = unbatch(w_batches)

            @test w_one_batch.xd.x == w_unbatched.xd.x
            @test w_one_batch.w == w_unbatched.w
        end
    end
end


"""
GenX: An Configurable Capacity Expansion Model
Copyright (C) 2021,  Massachusetts Institute of Technology
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
A complete copy of the GNU General Public License v2 (GPLv2) is available
in LICENSE.txt.  Users uncompressing this from an archive may not have
received this license file.  If not, see <http://www.gnu.org/licenses/>.
"""

@doc raw"""
	write_charge(path::AbstractString, inputs::Dict, setup::Dict, EP::Model)

Function for writing the charging energy values of the different storage technologies.
"""
function write_charge(path::AbstractString, inputs::Dict, setup::Dict, EP::Model)
	dfGen = inputs["dfGen"]
	G = inputs["G"]     # Number of resources (generators, storage, DR, and DERs)
	T = inputs["T"]     # Number of time steps (hours)
	# Power withdrawn to charge each resource in each time step
	dfCharge = DataFrame(Resource = inputs["RESOURCES"], Zone = dfGen[!,:Zone], AnnualSum = Array{Union{Missing,Float32}}(undef, G))
	charge = zeros(G,T)
	for i in 1:G
		if setup["ParameterScale"] ==1
			if i in inputs["STOR_ALL"]
				charge[i,:] = value.(EP[:vCHARGE])[i,:] * ModelScalingFactor
			elseif i in inputs["FLEX"]
				charge[i,:] = value.(EP[:vCHARGE_FLEX])[i,:] * ModelScalingFactor
			end
		else
			if i in inputs["STOR_ALL"]
				charge[i,:] = value.(EP[:vCHARGE])[i,:]
			elseif i in inputs["FLEX"]
				charge[i,:] = value.(EP[:vCHARGE_FLEX])[i,:]
			end
		end
		dfCharge[!,:AnnualSum][i] = sum(inputs["omega"].* charge[i,:])
	end
	dfCharge = hcat(dfCharge, DataFrame(charge, :auto))
	auxNew_Names=[Symbol("Resource");Symbol("Zone");Symbol("AnnualSum");[Symbol("t$t") for t in 1:T]]
	rename!(dfCharge,auxNew_Names)
	total = DataFrame(["Total" 0 sum(dfCharge[!,:AnnualSum]) fill(0.0, (1,T))], :auto)
	for t in 1:T
		total[:,t+3] .= sum(dfCharge[:,Symbol("t$t")][union(inputs["STOR_ALL"],inputs["FLEX"])])
	end
	rename!(total,auxNew_Names)
	dfCharge = vcat(dfCharge, total)
	CSV.write(joinpath(path, "charge.csv"), dftranspose(dfCharge, false), writeheader=false)
	return dfCharge
end


using Test, InteractiveUtils

using JuMP
using ECOS, SCS
using LightGraphs, MetaGraphs

using Seleroute

struct TestEdgeWise <: EdgeWiseObjectiveFunction end

@testset "Seleroute.jl" begin
    @testset "Data" begin
        include("data_modeltype.jl")
        include("data_parameters.jl")
        include("data_solution.jl")
        include("data_topology.jl")
    end
    @testset "Base models" begin
        include("basemodels_knowncapacities.jl")
        include("basemodels_objectives.jl")
    end
    @testset "Compute" begin
        include("compute_certain.jl")
        include("compute_oblivious_demand.jl")
    end
end


"""
    fire!(actr)

Selects a production rule and registers conflict resolution and new events for selected production rule

# Arguments

- `actr`: an ACT-R model object
"""
function fire!(actr::AbstractACTR)
    actr.procedural.state.busy ? (return nothing) : nothing
    rules = select_rule(actr)
    isempty(rules) ? (return nothing) : nothing
    rule = rules[1]
    description = "Selected "*rule.name
    type = "model"
    id = get_name(actr)
    tΔ = rnd_time(.05)
    resolving(actr, true)
    f(r, a, v) = (resolving(a, v), r.action())
    register!(actr, f, after, tΔ, rule, actr, false; description, type, id)
    return nothing
end

resolving(actr, v) = actr.procedural.state.busy = v

function select_rule(actr)
    rules = get_rule_set(actr)
    isempty(rules) ? (return rules) : nothing
    compute_utilities!(actr, rules)
    _,idx = findmax(x -> x.utility, rules)
    return rules[idx:idx]
end

function get_rule_set(actr)
    @unpack mmp_utility = actr.parms
    rules = get_rules(actr)
    if mmp_utility
        # use @code_string to filter out conditions that do not contain
        # slots. Consider caching a boolean that indicates whether a production
        # rule is eligible for partial matching
        return rules
    end
    return filter(r->match(actr, r), rules)
end

get_rules(actr) = actr.procedural.rules

function match(actr, rule)
    return all_match(actr, rule.conditions)
end

"""
    all_match(actr, conditions)

Checks whether all conditions of a production rule are satisfied.

# Arguments

- `actr`: an ACT-R model object
- `conditions`: a tuple of functions representing production rule conditions.
"""
function all_match(actr, conditions)
    for c in conditions
        !c(actr) ? (return false) : nothing
    end
    return true
end

function compute_utilities!(actr, rules)
    for rule in rules
        compute_utility!(actr, rule)
    end
    return nothing
end

function compute_utility!(actr, rule)
    @unpack utility_noise, mmp_utility = actr.parms
    mmp_utility ? compute_penalties!(actr, rule) : nothing
    utility_noise ? add_noise!(actr, rule) : nothing
    total_utility!(rule)
end

function compute_penalties!(actr, rule)
    @unpack model_trace = actr.scheduler
    @unpack δu,util_mmp_fun = actr.parms
    model_trace ? println("rule: ", rule.name) : nothing
    penalty = 0.0
    for c in rule.conditions
        penalty += util_mmp_fun(actr, c)
    end
    model_trace ? println("") : nothing
    rule.utility_penalty = penalty * δu
    return nothing
end

function utility_match(actr::ACTR, condition)
    @unpack model_trace = actr.scheduler
    penalty = condition(actr) ? 0.0 : 1.0
    model_trace ? condition_trace(actr, condition, penalty) : nothing
    return penalty
end

function condition_trace(actr, condition, penalty)
    str = @code_string condition(actr)
    println(str)
    println("penalty: $penalty")
    return nothing
end

function add_noise!(actr, rule)
    @unpack σu = actr.parms
    rule.utility_noise = rand(Normal(0, σu))
end

function total_utility!(r::Rule)
    r.utility_mean = r.initial_utility + r.utility_penalty
    r.utility = r.utility_mean + r.utility_noise
    return nothing
end

function Rule(;
    utility = 0.0,
    initial_utility = 0.0,
    utility_mean = 0.0,
    utility_penalty = 0.0,
    utility_noise = 0.0,
    conditions,
    name = "",
    actr,
    task,
    action,
    args = (),
    kwargs...
    )
    funs = conditions(actr, args...; kwargs...)
    can_pm = can_partial_match(funs, actr)
    Rule(
        utility,
        initial_utility,
        utility_mean,
        utility_penalty,
        utility_noise,
        conditions(actr, args...; kwargs...),
        ()->action(actr, task; kwargs...),
        can_pm,
        name
     )
end

function can_partial_match(conditions, actr)
    for c in conditions
        str = @code_string c(actr)
        !occursin("slots", str) ? (return false) : nothing
    end
    return true
end


"""
    SimpleAxis(v)

Povides an `AbstractAxis` interface for any `AbstractUnitRange`, `v `. `v` will
be considered both the `values` and `keys` of the return instance.

## Examples

A `SimpleAxis` is useful for giving a standard set of indices the ability to use
the filtering syntax for indexing.
```jldoctest
julia> using AxisIndices, StaticRanges

julia> x = SimpleAxis(2:10)
SimpleAxis(2:10)

julia> x[2]
2

julia> x[==(2)]
2

julia> x[2] == x[==(2)]  # keys and values are same
true

julia> x[>(2)]
SimpleAxis(3:10)

julia> x[>(2)]
SimpleAxis(3:10)

julia> x[1]
ERROR: BoundsError: attempt to access SimpleAxis(2:10) at index [1]
[...]
```
"""
struct SimpleAxis{I,Inds<:AbstractUnitRange{I}} <: AbstractAxis{I,Inds}
    parent::Inds

    function SimpleAxis{I,Inds}(inds::AbstractUnitRange) where {I,Inds}
        if inds isa Inds
            return new{I,Inds}(inds)
        else
            return SimpleAxis{I,Inds}(Inds(inds))
        end
    end

    SimpleAxis{I}(inds::AbstractUnitRange{I}) where {I} = SimpleAxis{I,typeof(inds)}(inds)
    SimpleAxis{I}(inds::AbstractUnitRange) where {I} = SimpleAxis{I}(AbstractUnitRange{I}(inds))
    function SimpleAxis{I}(inds::SimpleAxis) where {I}
        if eltype(inds) <: I
            return inds
        else
            return SimpleAxis{I}(parent(inds))
        end
    end

    SimpleAxis(inds::SimpleAxis) = inds
    SimpleAxis(inds::AbstractUnitRange{I}) where {I} = SimpleAxis{I}(inds)
    SimpleAxis(inds::IdentityUnitRange) = SimpleAxis(inds.indices)
    SimpleAxis() = new{Int,OneToMRange{Int}}(OneToMRange(0))

    """
        SimpleAxis(start::Integer, stop::Integer) -> SimpleAxis{UnitRange{Integer}}

    Passes `start` and `stop` arguments to `UnitRange` to construct the values of `SimpleAxis`.

    ## Examples
    ```jldoctest
    julia> using AxisIndices

    julia> SimpleAxis(1, 10)
    SimpleAxis(1:10)
    ```
    """
    SimpleAxis(start::Integer, stop::Integer) = SimpleAxis(start:stop)

    """
        SimpleAxis(stop::Integer) -> SimpleAxis{Base.OneTo{Integer}}

    Passes `stop` `Base.OneTo` to construct the values of `SimpleAxis`.

    ## Examples
    ```jldoctest
    julia> using AxisIndices

    julia> SimpleAxis(10)
    SimpleAxis(Base.OneTo(10))
    ```
    """
    SimpleAxis(stop::Integer) = SimpleAxis(StaticInt(1):stop)
end

@inline Base.getproperty(axis::SimpleAxis, k::Symbol) = getproperty(parent(axis), k)

ArrayInterface.unsafe_reconstruct(axis::SimpleAxis, inds; kwargs...) = SimpleAxis(inds)

# FIXME this should be deleted once https://github.com/SciML/ArrayInterface.jl/issues/79 is resolved
@propagate_inbounds function Base.getindex(axis::SimpleAxis, arg::StepRange{I}) where {I<:Integer}
    @boundscheck checkbounds(axis, arg)
    return maybe_unsafe_reconstruct(axis, arg)
end

function print_axis(io::IO, axis::SimpleAxis)
    print(io, "SimpleAxis($(Int(first(axis)):Int(last(axis))))")
end


# Digital signal plots
function PlotDigital(x::AbstractArray{T}, y::AbstractArray{T}; label_id::String="",
                     count::Integer=min(length(x), length(y)), offset::Integer=0,
                     stride::Integer=1) where {T<:ImPlotData}
    return PlotDigital(label_id, x, y, count, offset, stride * sizeof(T))
end

function PlotDigital(x::AbstractArray{T1}, y::AbstractArray{T2};
                     kwargs...) where {T1<:Real,T2<:Real}
    return PlotDigital(promote(x, y)...; kwargs...)
end

# xfield, yfield should be propertynames of eltype(structvec)
function PlotDigital(structvec::Vector{T}, xfield::Symbol, yfield::Symbol;
                     count::Integer=length(structvec), offset::Integer=0, stride::Integer=1,
                     label_id::String="") where {T}
    Tx = fieldtype(T, xfield)
    Ty = fieldtype(T, yfield)
    x_offset = fieldoffset(T, Base.fieldindex(T, xfield))
    y_offset = fieldoffset(T, Base.fieldindex(T, yfield))
    x_ptr = Ptr{Tx}((pointer(structvec, 1) + x_offset))
    y_ptr = Ptr{Ty}((pointer(structvec, 1) + y_offset))

    if !T.mutable
        # this is somewhat illegal and is used only to pass a pointer through AbstractArray argument into ccall
        x = unsafe_wrap(Vector{Tx}, x_ptr, size(structvec); own=false)
        y = unsafe_wrap(Vector{Ty}, y_ptr, size(structvec); own=false)
        stride = stride * sizeof(T)
    else # two new vectors every 1/60 second...
        x = Vector{Tx}(undef, length(structvec))
        y = Vector{Ty}(undef, length(structvec))
        for (i, val) in enumerate(structvec)
            x[i] = getproperty(val, xfield)
            y[i] = getproperty(val, yfield)
        end
        if Tx !== Ty
            x, y = promote(x, y)
        end
        stride = stride * sizeof(eltype(x))
    end

    return PlotDigital(label_id, x, y, count, offset, stride)
end


module Fst

include("semiring.jl")
include("wfst.jl")
include("io.jl")
include("compose.jl")
include("topological_sort.jl")
include("trim.jl")

end


using ImageCore, Colors, FixedPointNumbers, ColorVectorSpace, MappedArrays, OffsetArrays
using Test
using ImageCore: Pixel, NumberLike, GenericImage, GenericGrayImage, default_names

@testset "Image traits" begin
    for (B, swap) in ((rand(UInt16(1):UInt16(20), 3, 5), false),
                      (rand(Gray{Float32}, 3, 5), false),
                      (rand(RGB{Float16}, 3, 5), false),
                      (bitrand(3, 5), false),
                      (rand(UInt32, 3, 5), false),
                      (view(rand(3, 2, 5), :, 1, :), false),
                      (OffsetArray(rand(3, 5), -1:1, -2:2), false),
                      (PermutedDimsArray(rand(5, 3), (2, 1)), true),
                      (mappedarray(identity, PermutedDimsArray(rand(5, 3), (2, 1))), true),
                      (colorview(RGB, zeros(3, 5), zeroarray, zeros(3, 5)), false))
        @test pixelspacing(B) == (1,1)
        if !isa(B, SubArray)
            @test spacedirections(B) == (swap ? ((0,1),(1,0)) : ((1,0),(0,1)))
        else
            @test spacedirections(B) == ((1,0,0), (0,0,1))
        end
        @test sdims(B) == 2
        @test coords_spatial(B) == (swap ? (2,1) : (1,2))
        @test nimages(B) == 1
        @test size_spatial(B) == (3,5)
        if isa(B, OffsetArray)
            @test indices_spatial(B) == (-1:1, -2:2)
        else
            @test indices_spatial(B) == (Base.OneTo(3), Base.OneTo(5))
        end
        assert_timedim_last(B)
        @test width(B) == 5
        @test height(B) == 3
    end
end

# delibrately written in a redundant way
@testset "*Like traits" begin
    @testset "Pixel" begin
        @test NumberLike <: Pixel
        @test Number <: Pixel
        @test Gray <: Pixel
        @test RGB <: Pixel

        @test isa(oneunit(Gray), Pixel)
        @test isa(RGB(1.0, 0.0, 0.0), Pixel)
    end

    @testset "NumberLike" begin
        @test Number <: NumberLike
        @test Real <: NumberLike
        @test AbstractFloat <: NumberLike
        @test FixedPoint <: NumberLike
        @test Integer <: NumberLike
        @test Bool <: NumberLike

        @test Gray <: NumberLike
        @test Gray{<:AbstractFloat} <: NumberLike
        @test Gray{<:Bool} <: NumberLike

        @test isa(oneunit(Gray), NumberLike)
    end

    @testset "GenericImage" begin
        @test GenericGrayImage <: GenericImage
        for sz in Any[(3, 3), (3, 3, 3)]
            @test isa(rand(Bool, sz), GenericImage)
            @test isa(rand(N0f8, sz), GenericImage)
            @test isa(rand(Float32, sz), GenericImage)

            @test isa(rand(Gray, sz), GenericImage)
            @test isa(rand(Gray{Bool}, sz), GenericImage)
            @test isa(rand(Gray{N0f8}, sz), GenericImage)
            @test isa(rand(Gray{Float32}, sz), GenericImage)

            @test isa(rand(GrayA, sz), GenericImage)
            @test isa(rand(GrayA{N0f8}, sz), GenericImage)
            @test isa(rand(GrayA{Float32}, sz), GenericImage)

            @test isa(rand(RGB, sz), GenericImage)
            @test isa(rand(RGB{N0f8}, sz), GenericImage)
            @test isa(rand(RGB{Float32}, sz), GenericImage)

            @test isa(rand(RGBA, sz), GenericImage)
            @test isa(rand(RGBA{N0f8}, sz), GenericImage)
            @test isa(rand(RGBA{Float32}, sz), GenericImage)

            @test isa(rand(Lab, sz), GenericImage)
            @test isa(rand(Lab{Float32}, sz), GenericImage)
        end
    end

    @testset "GrayImage" begin
        for sz in [(3, 3), (3, 3, 3)]
            @test isa(rand(Bool, sz), GenericGrayImage)
            @test isa(rand(N0f8, sz), GenericGrayImage)
            @test isa(rand(Float32, sz), GenericGrayImage)

            @test isa(rand(Gray, sz), GenericGrayImage)
            @test isa(rand(Gray{Bool}, sz), GenericGrayImage)
            @test isa(rand(Gray{N0f8}, sz), GenericGrayImage)
            @test isa(rand(Gray{Float32}, sz), GenericGrayImage)
        end
    end

    @testset "dispatch" begin
        begin
            whatis(::GenericImage) = "GenericImage"
            whatis(::GenericGrayImage) = "GenericGrayImage"
            whatis(::GenericImage{<:AbstractRGB}) = "GenericRGBImage"

            whatis(::GenericImage{<:Pixel, 2}) = "Generic2dImage"
            whatis(::GenericGrayImage{<:NumberLike, 2}) = "Gray2dImage"
            whatis(::GenericImage{<:AbstractRGB, 2}) = "RGB2dImage"

            @test whatis(rand(Lab, 2, 2, 2)) == "GenericImage"

            @test whatis(rand(Lab, 2, 2)) == "Generic2dImage"

            @test whatis(rand(2, 2, 2)) == "GenericGrayImage"
            @test whatis(rand(N0f8, 2, 2, 2)) == "GenericGrayImage"
            @test whatis(rand(Bool, 2, 2, 2)) == "GenericGrayImage"
            @test whatis(rand(Float32, 2, 2, 2)) == "GenericGrayImage"
            @test whatis(rand(Int64, 2, 2, 2)) == "GenericGrayImage"

            @test whatis(rand(Gray, 2, 2, 2)) == "GenericGrayImage"
            @test whatis(rand(Gray{N0f8}, 2, 2, 2)) == "GenericGrayImage"
            @test whatis(rand(Gray{Bool}, 2, 2, 2)) == "GenericGrayImage"
            @test whatis(rand(Gray{Float32}, 2, 2, 2)) == "GenericGrayImage"

            @test whatis(rand(2, 2)) == "Gray2dImage"
            @test whatis(rand(N0f8, 2, 2)) == "Gray2dImage"
            @test whatis(rand(Bool, 2, 2)) == "Gray2dImage"
            @test whatis(rand(Float32, 2, 2)) == "Gray2dImage"
            @test whatis(rand(Int64, 2, 2)) == "Gray2dImage"

            @test whatis(rand(Gray, 2, 2)) == "Gray2dImage"
            @test whatis(rand(Gray{N0f8}, 2, 2)) == "Gray2dImage"
            @test whatis(rand(Gray{Bool}, 2, 2)) == "Gray2dImage"
            @test whatis(rand(Gray{Float32}, 2, 2)) == "Gray2dImage"

            @test whatis(rand(RGB, 2, 2, 2)) == "GenericRGBImage"
            @test whatis(rand(RGB{N0f8}, 2, 2, 2)) == "GenericRGBImage"
            @test whatis(rand(RGB{Float32}, 2, 2, 2)) == "GenericRGBImage"

            @test whatis(rand(BGR, 2, 2, 2)) == "GenericRGBImage"
            @test whatis(rand(BGR{N0f8}, 2, 2, 2)) == "GenericRGBImage"
            @test whatis(rand(BGR{Float32}, 2, 2, 2)) == "GenericRGBImage"

            @test whatis(rand(RGB, 2, 2)) == "RGB2dImage"
            @test whatis(rand(RGB{N0f8}, 2, 2)) == "RGB2dImage"
            @test whatis(rand(RGB{Float32}, 2, 2)) == "RGB2dImage"

            @test whatis(rand(BGR, 2, 2)) == "RGB2dImage"
            @test whatis(rand(BGR{N0f8}, 2, 2)) == "RGB2dImage"
            @test whatis(rand(BGR{Float32}, 2, 2)) == "RGB2dImage"
        end
    end
end

struct RowVector{T,P} <: AbstractVector{T}
    v::Vector{T}
    p::P
end

ImageCore.HasDimNames(::Type{<:RowVector}) = HasDimNames{true}()

ImageCore.HasProperties(::Type{<:RowVector}) = HasProperties{true}()

Base.names(::RowVector) = (:row,)
Base.axes(rv::RowVector) = axes(rv.v)

@testset "Trait Interface" begin
    img = reshape(1:24, 2,3,4)
    @test @inferred(namedaxes(img)) == NamedTuple{(:dim_1, :dim_2, :dim_3)}(axes(img))
    @test @inferred(HasDimNames(img)) == HasDimNames{false}()
    @test @inferred(HasProperties(img)) == HasProperties{false}()

    rv = RowVector([1:10...], Dict{String,Any}())
    @test @inferred(HasDimNames(rv)) == HasDimNames{true}()
    @test @inferred(HasProperties(rv)) == HasProperties{true}()
    @test @inferred(namedaxes(rv)) == NamedTuple{(:row,)}((Base.OneTo(10),))

    # default names
    @test @inferred(default_names(Val(3))) == (:dim_1, :dim_2, :dim_3)
end

nothing


# __precompile__()

module StabTrussTopOpt

# dependencies
using JSON
using Makie
using LinearAlgebra
using SparseArrays

# export TOProblemPar

include("common_types.jl")
include("draw_utils.jl")
include("parse_json.jl")
include("tto_define.jl")
include("solve.jl")

end # module


"""
    FullSpaceModel{T, S} <: AbstractNLPModel{T, S}

A data type to store the full space `BlockNLPModel` as an `AbstractNLPModel`.
"""
mutable struct FullSpaceModel{T,S} <: AbstractNLPModel{T,S}
    meta::NLPModelMeta{T,S}
    counters::Counters
    blocknlp::AbstractBlockNLPModel
end

"""
    FullSpaceModel(
        m::AbstractBlockNLPModel
        x0::Union{Vector{Float64}, Nothing} = nothing,
        y0::Union{Vector{Float64}, Nothing} = nothing
    ) where {T}

Converts a `BlockNLPModel` to a `AbstractNLPModel` that can be solved with any standard NLP solver.

# Arguments

- `m::AbstractBlockNLPModel`: name of the BlockNLPModel
- `x0::Union{Vector{Float64}, Nothing}`: warm start primal solution
- `y0::Union{Vector{Float64}, Nothing}`: warm start dual solution
"""
function FullSpaceModel(
    m::AbstractBlockNLPModel;
    x0::Union{Vector{Float64},Nothing} = nothing,
    y0::Union{Vector{Float64},Nothing} = nothing,
)
    nb = m.problem_size.block_counter
    n_var = sum(m.blocks[i].meta.nvar for i = 1:nb)

    l_con::Vector{Float64} = []
    u_con::Vector{Float64} = []
    for i = 1:nb
        if m.blocks[i].meta.ncon > 0
            l_con = vcat(l_con, m.blocks[i].meta.lcon)
            u_con = vcat(u_con, m.blocks[i].meta.ucon)
        end
    end
    l_con = vcat(l_con, get_rhs_vector(m))
    u_con = vcat(u_con, get_rhs_vector(m))

    l_var::Vector{Float64} = []
    u_var::Vector{Float64} = []
    for i = 1:nb
        l_var = vcat(l_var, m.blocks[i].meta.lvar)
        u_var = vcat(u_var, m.blocks[i].meta.uvar)
    end

    isnothing(x0) && (x0 = NLPModels.get_x0(m))
    isnothing(y0) && (y0 = NLPModels.get_y0(m))

    meta = NLPModelMeta(
        n_var,
        ncon = n_constraints(m),
        nnzh = sum(m.blocks[i].meta.nnzh for i = 1:nb),
        nnzj = sum(m.blocks[i].meta.nnzj for i = 1:nb) + nnz(get_linking_matrix(m)),
        x0 = x0,
        y0 = y0,
        lvar = l_var,
        uvar = u_var,
        lcon = l_con,
        ucon = u_con,
        minimize = true,
        name = "full_space",
    )
    return FullSpaceModel(meta, Counters(), m)
end

function NLPModels.get_x0(m::AbstractBlockNLPModel)
    nb = m.problem_size.block_counter
    x0 = [NLPModels.get_x0(m.blocks[i]) for i = 1:nb]
    return x0 = reduce(vcat, x0)
end

function NLPModels.get_y0(m::AbstractBlockNLPModel)
    nb = m.problem_size.block_counter
    y0 = [NLPModels.get_y0(m.blocks[i]) for i = 1:nb]
    return y0 = reduce(vcat, push!(y0, zeros(Float64, m.problem_size.link_counter)))
end

function NLPModels.obj(nlp::FullSpaceModel, x::AbstractVector)
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nvar x
    increment!(nlp, :neval_obj)
    return sum(
        obj(nlp.blocknlp.blocks[i].problem_block, x[nlp.blocknlp.blocks[i].var_idx]) for
        i = 1:nb
    )
end

function NLPModels.grad!(nlp::FullSpaceModel, x::AbstractVector, gx::AbstractVector)
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nvar x gx
    increment!(nlp, :neval_grad)
    for i = 1:nb
        v = @view gx[nlp.blocknlp.blocks[i].var_idx]
        grad!(nlp.blocknlp.blocks[i].problem_block, x[nlp.blocknlp.blocks[i].var_idx], v)
    end
    return gx
end

function NLPModels.hess_structure!(
    nlp::FullSpaceModel,
    rows::AbstractVector{T},
    cols::AbstractVector{T},
) where {T}
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nnzh rows cols
    idx = 0
    for i = 1:nb
        block = nlp.blocknlp.blocks[i]
        temp_idx = idx + 1
        idx += block.meta.nnzh
        v1 = @view rows[temp_idx:idx]
        v2 = @view cols[temp_idx:idx]
        hess_structure!(block.problem_block, v1, v2)
        v1 .+= (block.var_idx[1] - 1)
        v2 .+= (block.var_idx[1] - 1)
    end
    return rows, cols
end

function NLPModels.hess_coord!(
    nlp::FullSpaceModel,
    x::AbstractVector{T},
    vals::AbstractVector{T};
    obj_weight = one(T),
) where {T}
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nvar x
    @lencheck nlp.meta.nnzh vals
    increment!(nlp, :neval_hess)
    idx = 0
    for i = 1:nb
        block = nlp.blocknlp.blocks[i]
        temp_idx = idx + 1
        idx += block.meta.nnzh
        v = @view vals[temp_idx:idx]
        hess_coord!(block.problem_block, x[block.var_idx], v, obj_weight = obj_weight)
    end
    return vals
end

function NLPModels.hess_coord!(
    nlp::FullSpaceModel,
    x::AbstractVector{T},
    y::AbstractVector{T},
    vals::AbstractVector{T};
    obj_weight = one(T),
) where {T}
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nvar x
    @lencheck nlp.meta.nnzh vals
    @lencheck nlp.meta.ncon y
    increment!(nlp, :neval_hess)
    idx = 0
    for i = 1:nb
        block = nlp.blocknlp.blocks[i]
        if block.meta.ncon == 0
            temp_idx = idx + 1
            idx += block.meta.nnzh
            v = @view vals[temp_idx:idx]
            hess_coord!(block.problem_block, x[block.var_idx], v, obj_weight = obj_weight)
        else
            temp_idx = idx + 1
            idx += block.meta.nnzh
            v = @view vals[temp_idx:idx]
            hess_coord!(
                block.problem_block,
                x[block.var_idx],
                y[block.con_idx],
                v,
                obj_weight = obj_weight,
            )
        end
    end
    return vals
end

function NLPModels.cons!(nlp::FullSpaceModel, x::AbstractVector, cx::AbstractVector)
    nb = nlp.blocknlp.problem_size.block_counter
    @lencheck nlp.meta.nvar x
    @lencheck nlp.meta.ncon cx
    increment!(nlp, :neval_cons)
    idx = 0
    for i = 1:nb
        if nlp.blocknlp.blocks[i].meta.ncon > 0
            temp_idx = idx + 1
            idx += nlp.blocknlp.blocks[i].meta.ncon
            v = @view cx[temp_idx:idx]
            cons!(
                nlp.blocknlp.blocks[i].problem_block,
                x[nlp.blocknlp.blocks[i].var_idx],
                v,
            )
        end
    end
    idx += 1
    A = get_linking_matrix_blocks(nlp.blocknlp)
    cx[idx:end] .= sum(A[j] * x[nlp.blocknlp.blocks[j].var_idx] for j = 1:nb)
    return cx
end

function NLPModels.jac_structure!(
    nlp::FullSpaceModel,
    rows::AbstractVector{T},
    cols::AbstractVector{T},
) where {T}
    @lencheck nlp.meta.nnzj rows cols
    nb = nlp.blocknlp.problem_size.block_counter
    idx = 0
    for i = 1:nb
        block = nlp.blocknlp.blocks[i]
        if block.meta.ncon > 0
            temp_idx = idx + 1
            idx += block.meta.nnzj
            v1 = @view rows[temp_idx:idx]
            v2 = @view cols[temp_idx:idx]
            jac_structure!(block.problem_block, v1, v2)
            v1 .+= (block.con_idx[1] - 1)
            v2 .+= (block.var_idx[1] - 1)
        end
    end
    r, c, v = findnz(get_linking_matrix(nlp.blocknlp))
    rows[idx+1:end] .= r .+ nlp.blocknlp.problem_size.con_counter
    cols[idx+1:end] .= c
    return rows, cols
end

function NLPModels.jac_coord!(nlp::FullSpaceModel, x::AbstractVector, vals::AbstractVector)
    @lencheck nlp.meta.nvar x
    # @lencheck nlp.meta.nnzj vals # for madnlp
    increment!(nlp, :neval_jac)
    nb = nlp.blocknlp.problem_size.block_counter
    idx = 0
    for i = 1:nb
        block = nlp.blocknlp.blocks[i]
        if block.meta.ncon > 0
            temp_idx = idx + 1
            idx += block.meta.nnzj
            v1 = @view vals[temp_idx:idx]
            jac_coord!(block.problem_block, x[block.var_idx], v1)
        end
    end
    r, c, v = findnz(get_linking_matrix(nlp.blocknlp))
    # vals[idx+1:end] .= v
    vals[idx+1:nlp.meta.nnzj] .= v # for madnlp
    return vals
end


import .Flux

"""
    transform(model)

Convert a Flux Model to Lux Model.

# Examples
```julia
using Lux, Metalhead, Random

m = ResNet(18)
m2 = Lux.transform(m.layers)

x = randn(Float32, 224, 224, 3, 1);

ps, st = Lux.setup(Random.default_rng(), m2);

m2(x, ps, st)
```
"""
transform(::T) where {T} = error("Transformation for type $T not implemented")

transform(model::Flux.Chain) = Chain(transform.(model.layers)...)

function transform(model::Flux.BatchNorm)
    return BatchNorm(model.chs, model.λ; affine=model.affine,
                     track_stats=model.track_stats, epsilon=model.ϵ,
                     momentum=model.momentum)
end

function transform(model::Flux.Conv)
    return Conv(size(model.weight)[1:(end - 2)],
                size(model.weight, ndims(model.weight) - 1) * model.groups => size(model.weight,
                                                                                ndims(model.weight)),
                model.σ;
                stride=model.stride,
                pad=model.pad,
                bias=model.bias isa Bool ? model.bias : !(model.bias isa Flux.Zeros),
                dilation=model.dilation,
                groups=model.groups)
end

function transform(model::Flux.SkipConnection)
    return SkipConnection(transform(model.layers), model.connection)
end

function transform(model::Flux.Dense)
    return Dense(size(model.weight, 2), size(model.weight, 1), model.σ)
end

function transform(model::Flux.MaxPool)
    return MaxPool(model.k, model.pad, model.stride)
end

function transform(model::Flux.MeanPool)
    return MeanPool(model.k, model.pad, model.stride)
end

function transform(::Flux.GlobalMaxPool)
    return GlobalMaxPool()
end

function transform(::Flux.GlobalMeanPool)
    return GlobalMeanPool()
end

function transform(p::Flux.AdaptiveMaxPool)
    return AdaptiveMaxPool(p.out)
end

function transform(p::Flux.AdaptiveMeanPool)
    return AdaptiveMeanPool(p.out)
end

function transform(model::Flux.Parallel)
    return Parallel(model.connection, transform.(model.layers)...)
end

function transform(d::Flux.Dropout)
    return Dropout(Float32(d.p); dims=d.dims)
end

transform(::typeof(identity)) = NoOpLayer()

transform(::typeof(Flux.flatten)) = FlattenLayer()

transform(f::Function) = WrappedFunction(f)