text:
using Test, MPStates, LinearAlgebra, Random

@testset "Operations with Mpo" begin
# Initialize testing Mps.
rtest1 = MPStates.testMps("rtest1")
ctest1 = MPStates.testMps("ctest1")
rtest2 = MPStates.testMps("rtest2")
ctest2 = MPStates.testMps("ctest2")
rtest3 = MPStates.testMps("rtest3")

@testset "initialize empty Mpo" begin
    L = 6
    d = 2
    Op = Mpo(L, d)
    cOp = Mpo(ComplexF64, L, d)
    sOp = Mpo(L, 3)

    @test expected(Op, rtest1) ≈ 0. atol=1e-15
    @test expected(Op, rtest2) ≈ 0. atol=1e-15
    @test expected(sOp, rtest3) ≈ 0. atol=1e-15
    @test expected(cOp, ctest1) ≈ 0. atol=1e-15
    @test expected(cOp, ctest2) ≈ 0. atol=1e-15
end

@testset "Mpo with local terms" begin
    L = 6
    d = 2

    Op = Mpo(L, d)
    weights = cos.(1:L)
    add_ops!(Op, "n", weights)

    res1 = 0.
    res2 = 0.
    for i=1:L
        res1 += expected(rtest1, "n", i)*weights[i]
        res2 += expected(rtest2, "n", i)*weights[i]
    end
    @test expected(Op, rtest1) ≈ res1
    @test expected(Op, rtest2) ≈ res2

    cOp = Mpo(ComplexF64, L, d)
    cweights = complex.(tan.(1:L), sin.(1:L))
    add_ops!(cOp, "n", cweights)

    cres1 = complex(0.)
    cres2 = complex(0.)
    for i=1:L
        cres1 += expected(ctest1, "n", i)*cweights[i]
        cres2 += expected(ctest2, "n", i)*cweights[i]
    end
    @test expected(cOp, ctest1) ≈ cres1
    @test expected(cOp, ctest2) ≈ cres2

    # Test with spin states.
    d = 3

    Op = Mpo(L, d)
    weights1 = cos.(1:L)
    weights2 = tan.(1:L)
    add_ops!(Op, "Sz", weights1)
    add_ops!(Op, "S+", weights2)

    res3 = 0.
    for i=1:L
        res3 += expected(rtest3, "Sz", i)*weights1[i]
        res3 += expected(rtest3, "S+", i)*weights2[i]
    end
    @test expected(Op, rtest3) ≈ res3
end

@testset "Mpo with n_i*n_j terms" begin
    L = 6
    d = 2

    # Operator matrix.
    V = reshape(sin.(1:L^2), L, L)

    # Tests with real Mps.
    Op = Mpo(L, d)
    add_ops!(Op, "n", "n", V)

    res1 = 0.
    res2 = 0.
    for i=1:L, j=1:L
        res1 += expected(rtest1, "n", i, "n", j)*V[i, j]
        res2 += expected(rtest2, "n", i, "n", j)*V[i, j]
    end
    @test expected(Op, rtest1) ≈ res1
    @test expected(Op, rtest2) ≈ res2

    # Tests with complex Mps.
    cOp = Mpo(ComplexF64, L, d)
    cV = complex.(reshape(sin.(1:L^2), L, L), reshape(cos.(1:L^2), L, L))
    add_ops!(cOp, "n", "n", cV)

    cres1 = 0.
    cres2 = 0.
    for i=1:L, j=1:L
        cres1 += expected(ctest1, "n", i, "n", j)*cV[i, j]
        cres2 += expected(ctest2, "n", i, "n", j)*cV[i, j]
    end
    @test expected(cOp, ctest1) ≈ cres1
    @test expected(cOp, ctest2) ≈ cres2
end

@testset "Mpo with n_i*n_j and fermionic c+_i*c_j terms" begin
    L = 6
    d = 2

    # Operator matrices.
    V = reshape(sin.(1:L^2).^2, L, L)
    J = reshape(tan.(1:L^2).^3, L, L)

    # Test with real Mps.
    Op = Mpo(L, d)
    add_ops!(Op, "n", "n", V)
    add_ops!(Op, "c+", "c", J, ferm_op="Z")

    res1 = 0.
    res2 = 0.
    for i=1:L, j=1:L
        res1 += expected(rtest1, "n", i, "n", j)*V[i, j]
        res1 += expected(rtest1, "c+", i, "c", j, ferm_op="Z")*J[i, j]
        res2 += expected(rtest2, "n", i, "n", j)*V[i, j]
        res2 += expected(rtest2, "c+", i, "c", j, ferm_op="Z")*J[i, j]
    end
    @test expected(Op, rtest1) ≈ res1
    @test expected(Op, rtest2) ≈ res2

    # Tests with complex Mps.
    cOp = Mpo(ComplexF64, L, d)
    cV = complex.(reshape(sin.(1:L^2).^2, L, L), reshape(sin.(1:L^2), L, L))
    cJ = complex.(reshape(tan.(1:L^2).^3, L, L), reshape(tan.(1:L^2), L, L))
    add_ops!(cOp, "n", "n", cV)
    add_ops!(cOp, "c+", "c", cJ, ferm_op="Z")

    cres1 = 0.
    cres2 = 0.
    for i=1:L, j=1:L
        cres1 += expected(ctest1, "n", i, "n", j)*cV[i, j]
        cres1 += expected(ctest1, "c+", i, "c", j, ferm_op="Z")*cJ[i, j]
        cres2 += expected(ctest2, "n", i, "n", j)*cV[i, j]
        cres2 += expected(ctest2, "c+", i, "c", j, ferm_op="Z")*cJ[i, j]
    end
    @test expected(cOp, ctest1) ≈ cres1
    @test expected(cOp, ctest2) ≈ cres2
end

@testset "Mpo with non-fermionic b+_i*b_j terms" begin
    L = 6
    d = 2

    # Operator matrix.
    J = reshape(tan.(1:L^2), L, L)

    # Test with real Mps.
    Op = Mpo(L, d)
    add_ops!(Op, "b+", "b", J)

    res1 = 0.
    res2 = 0.
    for i=1:L, j=1:L
        res1 += expected(rtest1, "b+", i, "b", j)*J[i, j]
        res2 += expected(rtest2, "b+", i, "b", j)*J[i, j]
    end
    @test expected(Op, rtest1) ≈ res1
    @test expected(Op, rtest2) ≈ res2

    # Tests with complex Mps.
    cOp = Mpo(ComplexF64, L, d)
    cJ = complex.(reshape(sin.(1:L^2).^2, L, L), reshape(sin.(1:L^2), L, L))
    add_ops!(cOp, "b+", "b", cJ)

    cres1 = 0.
    cres2 = 0.
    for i=1:L, j=1:L
        cres1 += expected(ctest1, "b+", i, "b", j)*cJ[i, j]
        cres2 += expected(ctest2, "b+", i, "b", j)*cJ[i, j]
    end
    @test expected(cOp, ctest1) ≈ cres1
    @test expected(cOp, ctest2) ≈ cres2
end

@testset "Mpo with fermionic c+_i*c_j terms" begin
    L = 6
    d = 2

    # Operator matrices.
    J = reshape(tan.(1:L^2).*cos.(1:L^2), L, L)

    # Test with real Mps.
    Op = Mpo(L, d)
    add_ops!(Op, "c+", "c", J, ferm_op="Z")

    res1 = 0.
    res2 = 0.
    for i=1:L, j=1:L
        res1 += expected(rtest1, "c+", i, "c", j, ferm_op="Z")*J[i, j]
        res2 += expected(rtest2, "c+", i, "c", j, ferm_op="Z")*J[i, j]
    end
    @test expected(Op, rtest1) ≈ res1
    @test expected(Op, rtest2) ≈ res2

    # Tests with complex Mps.
    cOp = Mpo(ComplexF64, L, d)
    cJ = complex.(reshape(tan.(1:L^2).^2, L, L), reshape(sin.(1:L^2), L, L))
    add_ops!(cOp, "c+", "c", cJ, ferm_op="Z")

    cres1 = 0.
    cres2 = 0.
    for i=1:L, j=1:L
        cres1 += expected(ctest1, "c+", i, "c", j, ferm_op="Z")*cJ[i, j]
        cres2 += expected(ctest2, "c+", i, "c", j, ferm_op="Z")*cJ[i, j]
    end
    @test expected(cOp, ctest1) ≈ cres1
    @test expected(cOp, ctest2) ≈ cres2
end

@testset "Mpo with Sz_i*Sz_j and S+_i*S-_j terms" begin
    L = 6
    d = 3

    # Operator matrices.
    Vzz = reshape(sin.(1:L^2).^2, L, L)
    Vpm = reshape(tan.(1:L^2).^3, L, L)

    # Test with real Mps.
    Op = Mpo(L, d)
    add_ops!(Op, "Sz", "Sz", Vzz)
    add_ops!(Op, "S+", "S-", Vpm)

    res3 = 0.
    for i=1:L, j=1:L
        res3 += expected(rtest3, "Sz", i, "Sz", j)*Vzz[i, j]
        res3 += expected(rtest3, "S+", i, "S-", j)*Vpm[i, j]
    end
    @test expected(Op, rtest3) ≈ res3
end

@testset "expected of hermitian Mpo returns real number" begin
    L = 6
    d = 2

    # Let's create some random complex states to make sure that our premade
    # states have nothing in particular that could oversee some bug.
    Random.seed!(0)
    rtest4 = randomMps(L)
    rtest5 = randomMps(L)
    ctest4 = randomMps(ComplexF64, L)
    ctest5 = randomMps(ComplexF64, L)

    # Operator matrices.
    J = Symmetric(reshape(tan.(1:L^2).*cos.(1:L^2), L, L))
    cJ = Hermitian(complex.(reshape(tan.(1:L^2).^2, L, L),
                            reshape(sin.(1:L^2), L, L)))

    # Test with real symmetric Mpo (or complex with zero imaginary part).
    Op = Mpo(L, d)
    add_ops!(Op, "c+", "c", J, ferm_op="Z")
    @test imag(expected(Op, rtest1)) ≈ 0. atol=1e-8
    @test imag(expected(Op, rtest2)) ≈ 0. atol=1e-8
    @test imag(expected(Op, rtest4)) ≈ 0. atol=1e-8
    @test imag(expected(Op, rtest5)) ≈ 0. atol=1e-8

    Op = Mpo(L, d)
    add_ops!(Op, "c+", "c", J.^3)
    @test imag(expected(Op, rtest1)) ≈ 0. atol=1e-8
    @test imag(expected(Op, rtest2)) ≈ 0. atol=1e-8
    @test imag(expected(Op, rtest4)) ≈ 0. atol=1e-8
    @test imag(expected(Op, rtest5)) ≈ 0. atol=1e-8

    Op = Mpo(L, d)
    add_ops!(Op, "n", "n", J)
    @test imag(expected(Op, rtest1)) ≈ 0. atol=1e-8
    @test imag(expected(Op, rtest2)) ≈ 0. atol=1e-8
    @test imag(expected(Op, rtest4)) ≈ 0. atol=1e-8
    @test imag(expected(Op, rtest5)) ≈ 0. atol=1e-8

    cOp = Mpo(ComplexF64, L, d)
    add_ops!(cOp, "c+", "c", convert.(ComplexF64, J), ferm_op="Z")
    @test imag(expected(cOp, ctest1)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest2)) ≈ 0. atol=1e-8
    # Reduce error tolerance here because errors grow a lot with random states.
    @test imag(expected(cOp, ctest4)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest5)) ≈ 0. atol=1e-8

    cOp = Mpo(ComplexF64, L, d)
    add_ops!(cOp, "c+", "c", convert.(ComplexF64, tan.(J)))
    @test imag(expected(cOp, ctest1)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest2)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest4)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest5)) ≈ 0. atol=1e-8

    cOp = Mpo(ComplexF64, L, d)
    add_ops!(cOp, "n", "n", convert.(ComplexF64, tan.(J)))
    @test imag(expected(cOp, ctest1)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest2)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest4)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest5)) ≈ 0. atol=1e-8

    # Test with complex hermitian Mpo.
    cOp = Mpo(ComplexF64, L, d)
    add_ops!(cOp, "c+", "c", cJ, ferm_op="Z")
    @test imag(expected(cOp, ctest1)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest2)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest4)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest5)) ≈ 0. atol=1e-8

    cOp = Mpo(ComplexF64, L, d)
    add_ops!(cOp, "c+", "c", cJ.^2)
    @test imag(expected(cOp, ctest1)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest2)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest4)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest5)) ≈ 0. atol=1e-8

    cOp = Mpo(ComplexF64, L, d)
    add_ops!(cOp, "n", "n", cJ)
    @test imag(expected(cOp, ctest1)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest2)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest4)) ≈ 0. atol=1e-8
    @test imag(expected(cOp, ctest5)) ≈ 0. atol=1e-8

    # Test with spin operators.
    d = 3
    Op = Mpo(L, d)
    add_ops!(Op, "Sz", "Sz", J)
    @test imag(expected(Op, rtest3)) ≈ 0. atol=1e-15
    Op = Mpo(L, d)
    add_ops!(Op, "S+", "S-", J.^2)
    @test imag(expected(Op, rtest3)) ≈ 0. atol=1e-15
    Op = Mpo(L, d)
    add_ops!(Op, "Sz", "Sz", J)
    add_ops!(Op, "S-", "S+", tan.(J))
    @test imag(expected(Op, rtest3)) ≈ 0. atol=1e-15
end

@testset "Contraction of bond dimensions of an Mpo" begin
    L = 10
    d = 2

    # Build an Mpo with first neighbor interactions, bond dimension cannot be
    # greater than 4.
    J = diagm(1 => rand(L-1)) + diagm(-1 => rand(L-1))
    Op = Mpo(L, d)
    add_ops!(Op, "b+", "b", J)
    @test maximum(size.(Op.W, 4)) <= 4

    # Add periodic bond dimensions, this should increase the bond dim to 5.
    Jpbc = zeros(L, L)
    Jpbc[1, L] = 1.
    Jpbc[L, 2] = 3.
    add_ops!(Op, "b+", "b", Jpbc)
    @test maximum(size.(Op.W, 4)) <= 6
end
end # @testset "Operations with Mpo"


text:
using Test

# write your own tests here
include("test_timeframe.jl")

ext: jl


text:
struct StateEstimationPMU
    main::Array{Float64,2}
    flow::Array{Float64,2}
    estimate::Array{Float64,2}
    error::Array{Float64,1}
    baddata::Array{Float64,2}
    observability::Array{Array{Int64,1},1}
end

#### PMU state estimation
function runpmuse(system, measurements, num, numsys, settings, info)
    printstyled("Algorithm: Linear PMU state estimation\n"; bold = true)

    ########## Pre-processing ##########
    busi, type, Gshunt, Bshunt, Tini,
    branchi, resistance, reactance, charging, transTap, transShift, branchOn,
    meanVi, variVi, onVi, meanTi, variTi, onTi,
    meanIij, variIij, onIij, meanDij, variDij, onDij = read_pmusystem(system, measurements)

    main = fill(0.0, numsys.Nbus, 7)
    flow =  fill(0.0, numsys.Nbranch, 14)
    Va, Vm, Pinj, Qinj, Pshunt, Qshunt, Pij, Qij, Pji, Qji, Qcharging, Ploss,
    Qloss, Imij, Iaij, Imji, Iaji, badsave, islands = write_pmusystem(system, numsys, main, flow)

  algtime = @elapsed begin
    ########## Convert in integers ##########
    fromi = convert(Array{Int64,1}, system.branch[:, 2])
    toi = convert(Array{Int64,1}, system.branch[:, 3])
    branchiC = convert(Array{Int64,1}, measurements.pmuCurrent[:, 1])
    fromiC = convert(Array{Int64,1}, measurements.pmuCurrent[:, 2])
    toiC = convert(Array{Int64,1}, measurements.pmuCurrent[:, 3])
    busiV = convert(Array{Int64,1}, measurements.pmuVoltage[:, 1])

    bus = collect(1:numsys.Nbus)
    branch = collect(1:numsys.Nbranch)

    newnumbering = false
    @inbounds for i = 1:numsys.Nbus
        if bus[i] != busi[i]
            newnumbering = true
        end
    end
    if newnumbering
        println("The new bus numbering is running.")
    end
    numlabel = false
    @inbounds for i = 1:numsys.Nbranch
        if branch[i] != branchi[i]
            numlabel = true
            break
        end
    end
    if numlabel
        println("The new branch label numbering is running.")
    end
    from = renumber(fromi, busi, bus, newnumbering)
    to = renumber(toi, busi, bus, newnumbering)
    fromC = renumber(fromiC, busi, bus, newnumbering)
    toC = renumber(toiC, busi, bus, newnumbering)
    busV = renumber(busiV, busi, bus, newnumbering)
    branchC = renumber(branchiC, branchi, branch, newnumbering)

    ########## Jacobian ##########
    Nelement = 0; Nmeasure = 0; Ncur = 0; Nvol = 0;
    @inbounds for i = 1:num.pmuNc
        if onIij[i] == 1 || onDij[i] == 1
            onIij[i] = 1.0; onDij[i] = 1.0
            Nelement += 8; Nmeasure += 2; Ncur += 1
        end
    end
    @inbounds for i = 1:num.pmuNv
        if onVi[i] == 1 || onTi[i] == 1
            onVi[i] = 1.0; onTi[i] = 1.0
            Nelement += 2; Nmeasure += 2; Nvol += 1
        end
    end

    row = fill(0, Nelement); col = similar(row); jac = fill(0.0, Nelement)
    mean = fill(0.0, Nmeasure); variance = similar(mean);
    index = 1; rowindex = 1
    if settings.covariance
        idxco = 1
        covariance = fill(0.0, Ncur + Nvol)
    end

    Ybus, YbusT, Ytt, Yff, Yft, Ytf, admittance, tap, shunt = ybusac(system, numsys, bus, from, to, branchOn,
        Gshunt, Bshunt, resistance, reactance, charging, transTap, transShift)

    @inbounds for i = 1:num.pmuNc
        if onIij[i] == 1 || onDij[i] == 1
            idx = branchC[i]
            if branchOn[idx] == 1
                gij = real(admittance[idx])
                bij = imag(admittance[idx])

                if transTap[idx] == 0
                    tapij = 1
                else
                    tapij = transTap[idx]
                end

                shift = (pi / 180) * transShift[idx]
                bsi = charging[idx] / 2
            else
                gij = 0.0; bij = 0.0; bsi = 0.0; tapij = 1.0; shift = 0.0
            end

            ##### dIreal/dVreal,i and dIimag/dVimag,i  #####
            row[index] = rowindex; col[index] = fromC[i]
            row[index + 1] = rowindex + 1; col[index + 1] = fromC[i] + numsys.Nbus
            if from[idx] == fromC[i]
                jac[index] = gij / (tapij^2)
                jac[index + 1] = jac[index]
            else
                jac[index] = gij
                jac[index + 1] = jac[index]
            end
            index += 2

            ##### dIreal/dVreal,j and dIimag/dVimag,j #####
            row[index] = rowindex; col[index] = toC[i]
            row[index + 1] = rowindex + 1; col[index + 1] = toC[i] + numsys.Nbus
            if from[idx] == fromC[i]
                jac[index] = -(1 / tapij) * (gij * cos(shift) - bij * sin(shift))
                jac[index + 1] = jac[index]
            else
                jac[index] = -(1 / tapij) * (gij * cos(shift) + bij * sin(shift))
                jac[index + 1] = jac[index]
            end
            index += 2

            ##### dIreal/dVimag,i and dIimag/dVreal,i #####
            row[index] = rowindex; col[index] = fromC[i] + numsys.Nbus
            row[index + 1] = rowindex + 1; col[index + 1] = fromC[i]
            if from[idx] == fromC[i]
                jac[index] = -(bij + bsi) / (tapij^2)
                jac[index + 1] = -jac[index]
            else
                jac[index] = -bij - bsi
                jac[index + 1] = -jac[index]
            end
            index += 2

            ##### dIreal/dVimag,j and dIimag/dVreal,j#####
            row[index] = rowindex; col[index] = toC[i] + numsys.Nbus
            row[index + 1] = rowindex + 1; col[index + 1] = toC[i]
            if from[idx] == fromC[i]
                jac[index] = (1 / tapij) * (bij * cos(shift) + gij * sin(shift))
                jac[index + 1] =  -jac[index]
            else
                jac[index] = (1 / tapij) * (bij * cos(shift) - gij * sin(shift))
                jac[index + 1] =  -jac[index]
            end
            index += 2

            ##### Mean Ireal, Iimag #####
            mean[rowindex] = meanIij[i] * cos(meanDij[i])
            mean[rowindex + 1] = meanIij[i] * sin(meanDij[i])

            ##### Variance Ireal, Iimag #####
            variance[rowindex] = variIij[i] * (cos(meanDij[i]))^2 + variDij[i] * (meanIij[i] * sin(meanDij[i]))^2
            variance[rowindex + 1] = variIij[i] * (sin(meanDij[i]))^2 + variDij[i] * (meanIij[i] * cos(meanDij[i]))^2

            ##### Covariance #####
            if settings.covariance
                covariance[idxco] = sin(meanDij[i]) * cos(meanDij[i]) * (variIij[i] - variDij[i] * meanIij[i]^2)
                idxco += 1
            end

            rowindex += 2
        end
    end

    rowindex = 2 * Ncur + 1
    for i = 1:num.pmuNv
        if onVi[i] == 1 || onTi[i] == 1
            ##### dVreal/dVreal,i #####
            row[index] = rowindex
            col[index] = busiV[i]
            jac[index] = 1.0

            ##### Vimag/dVimag,i #####
            row[index + 1] = rowindex + 1
            col[index + 1] = busiV[i] + numsys.Nbus
            jac[index + 1] = 1.0
            index += 2

            ##### Mean Vreal, Vimag #####
            mean[rowindex] = meanVi[i] * cos(meanTi[i])
            mean[rowindex + 1] = meanVi[i] * sin(meanTi[i])

            ##### Variance Vreal, Vimag #####
            variance[rowindex] = variVi[i] * (cos(meanTi[i]))^2 + variTi[i] * (meanVi[i] * sin(meanTi[i]))^2
            variance[rowindex + 1] = variVi[i] * (sin(meanTi[i]))^2 + variTi[i] * (meanVi[i] * cos(meanTi[i]))^2

            ##### Covariance #####
            if settings.covariance
                covariance[idxco] = sin(meanTi[i]) * cos(meanTi[i]) * (variVi[i] - variTi[i] * meanVi[i]^2)
                idxco += 1
            end

            rowindex += 2
        end
    end

    J = sparse(row, col, jac, Nmeasure, 2 * numsys.Nbus)

    ########## Observability analysis ##########
    Npseudo = 0

    ########## LAV ##########
    if settings.lav[:lav] == 1
        VrealVimag = Array{Float64}(undef, 2 * numsys.Nbus)
        Nvar = 2 * numsys.Nbus
        x = lav(J, mean, Nvar, Nmeasure, settings)
        for i = 1:Nvar
            VrealVimag[i] = x[i] - x[Nvar + i]
        end
    end

    ########## WLS and bad data analysis ##########
    if settings.lav[:lav] == 0
        if settings.covariance
            idxVar = collect(1:Nmeasure); Icov = collect(1:2:Nmeasure); Jcov = collect(2:2:Nmeasure)
            R = sparse([idxVar; Icov; Jcov], [idxVar; Jcov; Icov], [variance; covariance; covariance], Nmeasure, Nmeasure)
            W = sparse(inv(Matrix(R)))
            G = J' * W * J
            b = J' * W * mean
        else
            W = spdiagm(0 => 1 ./ sqrt.(variance))
            H = W * J
            G = H' * H
            b = H' * W * mean
        end

        if settings.bad[:bad] == 1
            idx = findall(!iszero, J)
        end

        pass = 1; rbelow = true
        while rbelow
            if settings.bad[:pass] < pass || settings.bad[:bad] == 0
                rbelow = false
            end

            VrealVimag = wls(G, b, settings.solve)

            if rbelow
                J, mean, W, badsave, rbelow = baddata(settings, numsys, VrealVimag, mean, variance, W, G, J, Nmeasure, idx, badsave, rbelow)
                pass += 1
            end
        end
    end

    Vc = zeros(ComplexF64, numsys.Nbus)
    for i = 1:numsys.Nbus
        Vc[i] = complex(VrealVimag[i], VrealVimag[i + numsys.Nbus])
    end

    for i = 1:numsys.Nbus
        Vm[i] = abs(Vc[i])
        Va[i] = (180 / pi) * angle(Vc[i])

        Sshunt = Vc[i] * conj(Vc[i] * shunt[i])
        Pshunt[i] = real(Sshunt) * system.basePower
        Qshunt[i] = imag(Sshunt) * system.basePower

        I = 0.0 + im * 0.0
        for j in Ybus.colptr[i]:(Ybus.colptr[i + 1] - 1)
            row = Ybus.rowval[j]
            I += conj(YbusT[row, i]) * conj(Vc[row])
        end
        Si::ComplexF64 = I * Vc[i]
        Pinj[i] = real(Si) * system.basePower
        Qinj[i] = imag(Si) * system.basePower
    end

    acflow(system.basePower, numsys, from, to, branchOn, Vc, Yff, Yft, Ytt, Ytf,
        admittance, tap, resistance, reactance, charging,
        Pij, Qij, Pji, Qji, Qcharging, Ploss, Qloss, Imij, Iaij, Imji, Iaji)

    estimate_row = 9; error_row = 3; exact = false
    if size(measurements.pmuVoltage, 2) == 9 && size(measurements.pmuCurrent, 2) == 11
        estimate_row = 11; error_row = 6; exact = true
    end
    estimate = zeros(Nmeasure, estimate_row)
    error = zeros(error_row)
    Nm = Nmeasure - size(badsave, 1)
    idx = 1; toRad = pi / 180
    for (i, on) in enumerate(onIij)
        if on != 0
            idy = idx + 1
            estimate[idx, 1:5] = [idx 1.0 2.0 9.0 i]
            estimate[idy, 1:5] = [idx 1.0 2.0 10.0 i]

            if idx in badsave[:, 1]
                estimate[idx, 2] = 2.0; onIij[i] = 0; onDij[i] = 0
            end
            if idy in badsave[:, 1]
                estimate[idy, 2] = 2.0; onIij[i] = 0; onDij[i] = 0
            end

            estimate[idx, 6] = meanIij[i] * cos(meanDij[i])
            estimate[idy, 6] = meanIij[i] * sin(meanDij[i])
            estimate[idx, 7] = variance[idx]
            estimate[idy, 7] = variance[idy]
            k = branchC[i]
            if measurements.pmuCurrent[i, 2] == system.branch[k, 2] && measurements.pmuCurrent[i, 3] == system.branch[k, 3]
                estimate[idx, 8] = Imij[k] * cos(toRad * Iaij[k])
                estimate[idy, 8] = Imij[k] * sin(toRad * Iaij[k])
            else
                estimate[idx, 8] = Imji[k] * cos(toRad * Iaji[k])
                estimate[idy, 8] = Imji[k] * sin(toRad * Iaji[k])
            end
            estimate[idx, 9] = abs(estimate[idx, 6] - estimate[idx, 8])
            estimate[idy, 9] = abs(estimate[idy, 6] - estimate[idy, 8])

            if estimate[idx, 2] != 2.0
                error[1] += estimate[idx, 9] ./ Nm
                error[2] += estimate[idx, 9]^2 ./ Nm
                error[3] += estimate[idx, 9]^2 ./ variance[idx]
            end
            if estimate[idy, 2] != 2.0
                error[1] += estimate[idy, 9] ./ Nm
                error[2] += estimate[idy, 9]^2 ./ Nm
                error[3] += estimate[idy, 9]^2 ./ variance[idy]
            end

            if exact
                estimate[idx, 10] = measurements.pmuCurrent[i, 10] * cos(measurements.pmuCurrent[i, 11])
                estimate[idy, 10] = measurements.pmuCurrent[i, 10] * sin(measurements.pmuCurrent[i, 11])

                estimate[idx, 11] = abs(estimate[idx, 8] - estimate[idx, 10])
                estimate[idy, 11] = abs(estimate[idy, 8] - estimate[idy, 10])

                if estimate[idx, 2] != 2.0
                    error[4] += estimate[idx, 11] / Nm
                    error[5] += estimate[idx, 11]^2 / Nm
                    error[6] += estimate[idx, 11]^2 / variance[idx]
                end
                if estimate[idy, 2] != 2.0
                    error[4] += estimate[idy, 11] / Nm
                    error[5] += estimate[idy, 11]^2 / Nm
                    error[6] += estimate[idy, 11]^2 / variance[idy]
                end
            end
            idx += 2
        end
    end
    for (i, on) in enumerate(onVi)
        if on != 0
            idy = idx + 1
            estimate[idx, 1:5] = [idx 1.0 2.0 11.0 i]
            estimate[idy, 1:5] = [idx 1.0 2.0 12.0 i]

            if idx in badsave[:, 1]
                estimate[idx, 2] = 2.0; onVi[i] = 0; onTi[i] = 0
            end
            if idy in badsave[:, 1]
                estimate[idy, 2] = 2.0; onVi[i] = 0; onTi[i] = 0
            end

            estimate[idx, 6] = meanVi[i] * cos(meanTi[i])
            estimate[idy, 6] = meanVi[i] * sin(meanTi[i])
            estimate[idx, 7] = variance[idx]
            estimate[idy, 7] = variance[idy]
            estimate[idx, 8] = Vm[busV[i]] * cos(toRad * Va[busV[i]])
            estimate[idy, 8] = Vm[busV[i]] * sin(toRad * Va[busV[i]])

            estimate[idx, 9] = abs(estimate[idx, 6] - estimate[idx, 8])
            estimate[idy, 9] = abs(estimate[idy, 6] - estimate[idy, 8])

            if estimate[idx, 2] != 2.0
                error[1] += estimate[idx, 9] ./ Nm
                error[2] += estimate[idx, 9]^2 ./ Nm
                error[3] += estimate[idx, 9]^2 ./ variance[idx]
            end
            if estimate[idy, 2] != 2.0
                error[1] += estimate[idy, 9] ./ Nm
                error[2] += estimate[idy, 9]^2 ./ Nm
                error[3] += estimate[idy, 9]^2 ./ variance[idy]
            end

            if exact
                estimate[idx, 10] = measurements.pmuVoltage[i, 8] * cos(measurements.pmuVoltage[i, 9])
                estimate[idy, 10] = measurements.pmuVoltage[i, 8] * sin(measurements.pmuVoltage[i, 9])

                estimate[idx, 11] = abs(estimate[idx, 8] - estimate[idx, 10])
                estimate[idy, 11] = abs(estimate[idy, 8] - estimate[idy, 10])

                if estimate[idx, 2] != 2.0
                    error[4] += estimate[idx, 11] / Nm
                    error[5] += estimate[idx, 11]^2 / Nm
                    error[6] += estimate[idx, 11]^2 / variance[idx]
                end
                if estimate[idy, 2] != 2.0
                    error[4] += estimate[idy, 11] / Nm
                    error[5] += estimate[idy, 11]^2 / Nm
                    error[6] += estimate[idy, 11]^2 / variance[idy]
                end
            end
            idx += 2
        end
    end
    error[2] = sqrt(error[2])
    if exact
        error[5] = sqrt(error[5])
    end

    pass = collect(1:size(badsave, 1))
    idxbad = trunc.(Int, badsave[:, 1])
    bad = [pass estimate[idxbad, 3:5] badsave[:, 2] estimate[idxbad, 2]]

 end # algtime

    ########## Results ##########
    results = StateEstimationPMU(main, flow, estimate, error, bad, islands)
    header, group = results_estimatepmu(system, numsys, measurements, num, settings, results, idxbad, Npseudo, algtime)
    if !isempty(settings.save)
        savedata(results, measurements, system; info = info, group = group, header = header, path = settings.save)
    end

    return results
end

### View system
function read_pmusystem(system, measurements)
    busi = @view(system.bus[:, 1])
    type = @view(system.bus[:, 2])
    Gshunt = @view(system.bus[:, 5])
    Bshunt =  @view(system.bus[:, 6])
    Tini = @view(system.bus[:, 9])

    branchi = @view(system.branch[:, 1])
    resistance = @view(system.branch[:, 4])
    reactance = @view(system.branch[:, 5])
    charging = @view(system.branch[:, 6])
    transTap = @view(system.branch[:, 10])
    transShift = @view(system.branch[:, 11])
    branchOn = @view(system.branch[:, 12])

    meanVi = @view(measurements.pmuVoltage[:, 2])
    variVi = @view(measurements.pmuVoltage[:, 3])
    onVi = @view(measurements.pmuVoltage[:, 4])
    meanTi = @view(measurements.pmuVoltage[:, 5])
    variTi = @view(measurements.pmuVoltage[:, 6])
    onTi = @view(measurements.pmuVoltage[:, 7])

    meanIij = @view(measurements.pmuCurrent[:, 4])
    variIij = @view(measurements.pmuCurrent[:, 5])
    onIij = @view(measurements.pmuCurrent[:, 6])
    meanDij = @view(measurements.pmuCurrent[:, 7])
    variDij = @view(measurements.pmuCurrent[:, 8])
    onDij = @view(measurements.pmuCurrent[:, 9])

    return busi, type, Gshunt, Bshunt, Tini,
           branchi, resistance, reactance, charging, transTap, transShift, branchOn,
           meanVi, variVi, onVi, meanTi, variTi, onTi,
           meanIij, variIij, onIij, meanDij, variDij, onDij
end

### Write data
function write_pmusystem(system, numsys, main, flow)
    for i = 1:numsys.Nbus
        main[i, 1] = system.bus[i, 1]
    end
    Vm = @view(main[:, 2])
    Va = @view(main[:, 3])
    Pinj = @view(main[:, 4])
    Qinj = @view(main[:, 5])
    Pshunt = @view(main[:, 6])
    Qshunt = @view(main[:, 7])

    for i = 1:numsys.Nbranch
        flow[i, 1] = system.branch[i, 1]
        flow[i, 2] = system.branch[i, 2]
        flow[i, 3] = system.branch[i, 3]
    end
    Pij = @view(flow[:, 4])
    Qij = @view(flow[:, 5])
    Pji = @view(flow[:, 6])
    Qji = @view(flow[:, 7])
    Qcharging = @view(flow[:, 8])
    Ploss = @view(flow[:, 9])
    Qloss = @view(flow[:, 10])
    Imij = @view(flow[:, 11])
    Iaij = @view(flow[:, 12])
    Imji = @view(flow[:, 13])
    Iaji = @view(flow[:, 14])

    badsave = reshape(Float64[], 0, 2)
    islands = [Int64[]]

    return Va, Vm, Pinj, Qinj, Pshunt, Qshunt,
        Pij, Qij, Pji, Qji, Qcharging, Ploss, Qloss, Imij, Iaij, Imji, Iaji,
        badsave, islands
end

max_stars_repo_name: mcosovic/JuliaGrid

max_issues_repo_name: mcosovic/JuliaGrid

max_forks_repo_name: mcosovic/JuliaGrid


text:
using AppliGeneralLedger
using DataFrames
using Test

# TEST MODEL
@testset "Journal" begin
    entry = create_journal_entry("20200700", "APPLI", "A1001", 1300, 8000, 1000.0, 0.0, 210.0, "LS")
    @test entry.id == "20200700"
    @test entry.customer_id == "APPLI"
    @test entry.invoice_nbr == "A1001"
    @test entry.from == 1300
    @test entry.to == 8000
    @test entry.debit == 1000.0
    @test entry.credit == 0
    @test entry.vat == 210.0
    @test entry.descr == "LS"
end

#=
@testset "From" begin
    entry = create_journal_entry("20200700", "APPLI", "A1001", 1300, 8000, 1000.0, 0.0, 210.0, "LS")
    r1 = create_record_from(entry)
    r2 = create_record_to(entry)
    r3 = create_record_vat(entry, 4000)
    @test r1.debit == 1210
    @test r2.credit == 1000
    @test r3.credit == 210
    @test r1.debit == r2.credit + r3.credit
end

@testset "To" begin
    entry = create_journal_entry("20200700", "APPLI", "A1001", 1150, 1300, 1210.0, 0.0, 0.0, "LS")
    r1 = create_record_from(entry)
    r2 = create_record_to(entry)
    r3 = create_record_vat(entry, 4000)
    @test r1.debit == 1210
    @test r2.credit == 1210
    @test r3.credit == 0.0
    @test r1.debit == r2.credit + r3.credit
end
=#

@testset "Database read" begin
    entry1 = create_journal_entry("20200700", "APPLI", "A1001", 1300, 8000, 1000.0, 0.0, 210.0, "LS")
    entry2 = create_journal_entry("20200700", "PRG", "A1002", 1300, 8000, 1000.0, 0.0, 210.0, "LS")
    process([entry1, entry2])

    items =  read_from_file("./test_ledger.txt")
    df = DataFrame(items)
    result = sum(df.debit) - sum(df.credit)
    @test sum(result) == 0

    cmd = `rm test_journal.txt test_ledger.txt`
    run(cmd)
end


text:
using Documenter, MatrixEquations
DocMeta.setdocmeta!(MatrixEquations, :DocTestSetup, :(using MatrixEquations); recursive=true)

makedocs(
  modules  = [MatrixEquations],
  sitename = "MatrixEquations.jl",
  authors  = "Andreas Varga",
  format   = Documenter.HTML(prettyurls = false),
  pages    = [
     "Home"   => "index.md",
     "Library" => [
        "lyapunov.md",
        "plyapunov.md",
        "riccati.md",
        "sylvester.md",
        "sylvkr.md",
        "condest.md",
        "meoperators.md"
     ],
     "Internal" => [
        "lapackutil.md",
        "meutil.md"
     ],
     "Index" => "makeindex.md"
  ]
)

deploydocs(deps = nothing, make = nothing,
  repo = "github.com/andreasvarga/MatrixEquations.jl.git",
  target = "build",
  devbranch = "master"
)


text:
#!/usr/bin/env julia

module PE18

const str=
"75
95 64
17 47 82
18 35 87 10
20 04 82 47 65
19 01 23 75 03 34
88 02 77 73 07 63 67
99 65 04 28 06 16 70 92
41 41 26 56 83 40 80 70 33
41 48 72 33 47 32 37 16 94 29
53 71 44 65 25 43 91 52 97 51 14
70 11 33 28 77 73 17 78 39 68 17 57
91 71 52 38 17 14 91 43 58 50 27 29 48
63 66 04 68 89 53 67 30 73 16 69 87 40 31
04 62 98 27 23 09 70 98 73 93 38 53 60 04 23"

const data = map(split(str,"\n")) do line
    parse.(Int,split(line," "))
end

end

export pe18

function pe18(data=PE18.data)
    foldr(data) do next, prev
        map(max,prev[1:end-1].+next,prev[2:end].+next)
    end |> first
end

if !haskey(ENV,"PROJECT_EULER_WITHOUT_RESULT")
    pe18()|>print
end

max_stars_repo_path: 18. Maximum path sum I/18.jl

max_stars_repo_name: miRoox/ProjectEuler

max_issues_repo_path: 18. Maximum path sum I/18.jl

max_issues_repo_name: miRoox/ProjectEuler

max_forks_repo_path: 18. Maximum path sum I/18.jl

max_forks_repo_name: miRoox/ProjectEuler


text:
# Parameter of neural network
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

export Parameter

mutable struct Parameter
    data
    grad
end

"""
    Class for trainable network parameters.

 *Fields:*

 - `Parameter.data`: weights

 - `Parameter.grad`: gradient

"""
Parameter(x) = Parameter(x, nothing)

size(x::Parameter) = size(x.data)

@Flux.functor Parameter

"""
    clear_grad!(NL::NeuralNetLayer)

or

    clear_grad!(P::AbstractArray{Parameter, 1})

 Set gradients of each `Parameter` in the network layer to `nothing`.
"""
function clear_grad!(P::AbstractArray{Parameter, 1})
    for j=1:length(P)
        P[j].grad = nothing
    end
end


text:
### wiener-time-series.jl
#
# Copyright (C) 2016 Mosè Giordano.
#
# Maintainer: Mosè Giordano <mose AT gnu DOT org>
# Keywords: deconvolution, signal processing
#
# This file is a part of Deconvolution.jl.
#
# License is MIT "Expat".
#
### Commentary:
#
# This file produces an example image used in Deconvolution.jl documentation.
#
### Code:

using Random

Random.seed!(7)

using LombScargle, Deconvolution, Plots, Statistics, FFTW
t = range(0, stop=10, length=1000) # observation times
x = sinpi.(t) .* cos.(5t) - 1.5cospi.(t) .* sin.(2t) # the original signal

# Gaussian blurring kernel
kernel = exp.( - 10 .* (t .- 5).^2)
kernel ./= sum(kernel) # normalize kernel to sum of 1
kernel = ifftshift(kernel) # move center to index pos 1

n = rand(length(x)) # noise to be added
noise = 3 .* (n .- mean(n))
y = x + noise # observed noisy signal
# blurred and noise signal
y_blurred = real(ifft(fft(kernel) .* fft(x))) + noise

# Lomb-Scargle periodogram
p = lombscargle(t, y, maximum_frequency=2, samples_per_peak=10)
plot(freqpower(p)...)

m1 = LombScargle.model(t, y, findmaxfreq(p, [0, 0.5])[1]) # first model
m2 = LombScargle.model(t, y, findmaxfreq(p, [0.5, 1])[1]) # second model
m3 = LombScargle.model(t, y, findmaxfreq(p, [1, 1.5])[1]) # third model

signal = m1 + m2 + m3
polished = wiener(y, signal, noise)
deblurred = wiener(y_blurred, signal, noise, kernel)

# Plots
plot(t, x, size=(900, 600), label="Original signal", linewidth=2)
plot!(t, y, label="Observed signal")
savefig("time-series-observed.png")

plot(t, x, size=(900, 600), label="Original signal", linewidth=2)
plot!(t, y_blurred, label="Blurred signal")
savefig("time-series-observed-blurred.png")

plot(t, x, size=(900, 600), label="Original signal", linewidth=2)
plot!(t, polished, label="Recovered with Wiener")
plot!(t, signal, label="Lomb–Scargle model")
savefig("time-series-recovered.png")

plot(t, x, size=(900, 600), label="Original signal", linewidth=2)
plot!(t, deblurred, label="Deblurred with Wiener")
plot!(t, signal, label="Lomb–Scargle model")
savefig("time-series-deblurred.png")


text:
# Copyright © nhdalyMadeThis, LLC
# Released under MIT License

module HueClone

using Colors
using Rematch # For matching text-based user input

export play_blink, play_juno

const DEFAULT_ROWS, DEFAULT_COLS = 6,4

function create_colors_grid(rows, cols)
    corners = rand(RGB, 4)
    topleft,topright,bottomleft,bottomright = corners

    leftedge = range(topleft, stop=bottomleft, length=rows)
    rightedge = range(topright, stop=bottomright, length=rows)

    tiles = Array{RGB,2}(undef, (rows,cols));
    for r in 1:rows
        tiles[r,:] = range(leftedge[r], stop=rightedge[r], length=cols)
    end

    tiles
end

macro swap!(a,b)
    quote
        tmp = $(esc(a))
        $(esc(a)) = $(esc(b))
        $(esc(b)) = tmp
    end
end

function user_attempt_swap!(a, b)
    a,b = Tuple(a), Tuple(b)
    (r,c) = size(tiles)
    fixed_tile_positions = [(1,1), (r,1), (1,c), (r,c)]
    if a == b || a in fixed_tile_positions || b in fixed_tile_positions
        return
    end
    @swap!(tiles[a...], tiles[b...])
end

function new_game(rows, cols)
    global goal = create_colors_grid(rows, cols)
    global tiles = deepcopy(goal)
    for _ in 1:100
        user_attempt_swap!(rand(CartesianIndices(tiles)),
                           rand(CartesianIndices(tiles)))
    end
end

function play_game(io_device; size=(DEFAULT_ROWS, DEFAULT_COLS))
    new_game(size...)
    while tiles != goal
        print_board(io_device, tiles)
        a,b = collect_userinput(io_device)
        user_attempt_swap!(a, b)
    end
    print_board(io_device, tiles)
    win_game(io_device)
end

# ------------------------------------
# ------ Juno-REPL based -------------
# ------------------------------------

struct JunoDisplay end

"""
    play_juno()

Play the game in Atom, with the Juno REPL and Juno Plots display
"""
play_juno(;size=(DEFAULT_ROWS, DEFAULT_COLS)) = play_game(JunoDisplay(), size=size)

print_board(::JunoDisplay, tiles) = display(tiles)

function collect_userinput(::JunoDisplay)
    while true
        try
            println("Enter tile coordinates to swap (col,row):")
            # Read a and b from user input
            print("a: ")
            @match Expr(Tuple, a) = Meta.parse(readline())
            print("b: ")
            @match Expr(Tuple, b) = Meta.parse(readline())
            return reverse(a), reverse(b)
        catch e
            e isa InterruptException && rethrow()
            @warn("Must enter valid coordinate tuple (x,y)")
            @error(e)
        end
    end
end
function win_game(::JunoDisplay)
    println("Hooray you win! :)")
end

# ------------------------------------
# ------ Blink based -----------------
# ------------------------------------

using Blink

"""
    play_blink()

Play the game using a Blink html window.
"""
function play_blink(;size=(DEFAULT_ROWS, DEFAULT_COLS))
    w = create_board(size...)
    handle(w, "click") do pos
        handle_click(pos)
    end

    play_game(w; size=size)
end

function create_board(rows, cols)
    sq_size = 50
    w = Blink.Window(Dict(:width => sq_size*cols + 4,
                          :height => sq_size*rows + 40),
                     async=false)
    body!(w, """

    """, async=false)
    return w
end
function print_board(w::Blink.Window, tiles)
    content!(w, ".grid",
            join([
                """

                    $(join([
                        """

"""
                        for r in 1:size(tiles)[1]
                        ], "")
                    )

                """
                for c in 1:size(tiles)[2]
                ], "")
        ; fade=false, async=false)
end

clicksChannel = Channel(2)
function handle_click(pos)
    global clicksChannel
    put!(clicksChannel, Tuple(pos))
end
function collect_userinput(w::Blink.Window)
    return (take!(clicksChannel), take!(clicksChannel))
end

function win_game(w::Blink.Window)
    content!(w, ".messages", "YOU WIN! 🤣")
end

end


text:
"""
    MPI.Comm

An MPI Communicator object.
"""
mutable struct Comm
    val::MPI_Comm
end
Base.:(==)(a::Comm, b::Comm) = a.val == b.val
Base.cconvert(::Type{MPI_Comm}, comm::Comm) = comm
Base.unsafe_convert(::Type{MPI_Comm}, comm::Comm) = comm.val
Base.unsafe_convert(::Type{Ptr{MPI_Comm}}, comm::Comm) = convert(Ptr{MPI_Comm}, pointer_from_objref(comm))

const COMM_NULL = Comm(Consts.MPI_COMM_NULL[])
add_load_time_hook!(() -> COMM_NULL.val = Consts.MPI_COMM_NULL[])

"""
    MPI.COMM_WORLD

A communicator containing all processes with which the local rank can communicate at
initialization. In a typical "static-process" model, this will be all processes.
"""
const COMM_WORLD = Comm(Consts.MPI_COMM_WORLD[])
add_load_time_hook!(() -> COMM_WORLD.val = Consts.MPI_COMM_WORLD[])

"""
    MPI.COMM_SELF

A communicator containing only the local process.
"""
const COMM_SELF = Comm(Consts.MPI_COMM_SELF[])
add_load_time_hook!(() -> COMM_SELF.val = Consts.MPI_COMM_SELF[])

Comm() = Comm(COMM_NULL.val)

function free(comm::Comm)
    if comm != COMM_NULL && !Finalized()
        # int MPI_Comm_free(MPI_Comm *comm)
        @mpichk ccall((:MPI_Comm_free, libmpi), Cint, (Ptr{MPI_Comm},), comm)
    end
    return nothing
end

"""
    Comm_rank(comm::Comm)

The rank of the process in the particular communicator's group.

Returns an integer in the range `0:MPI.Comm_size()-1`.

# See also
- [`MPI.Comm_size`](@ref).

# External links
$(_doc_external("MPI_Comm_rank"))
"""
function Comm_rank(comm::Comm)
    rank = Ref{Cint}()
    @mpichk ccall((:MPI_Comm_rank, libmpi), Cint, (MPI_Comm, Ptr{Cint}), comm, rank)
    Int(rank[])
end

"""
    Comm_size(comm::Comm)

The number of processes involved in communicator.

# See also
- [`MPI.Comm_rank`](@ref).

# External links
$(_doc_external("MPI_Comm_size"))
"""
function Comm_size(comm::Comm)
    size = Ref{Cint}()
    @mpichk ccall((:MPI_Comm_size, libmpi), Cint, (MPI_Comm, Ptr{Cint}), comm, size)
    Int(size[])
end

"""
    Comm_group(comm::Comm)

Accesses the group associated with given communicator.

# External links
$(_doc_external("MPI_Comm_group"))
"""
function Comm_group(comm::Comm)
    newgroup = Group()
    @mpichk ccall((:MPI_Comm_group, libmpi), Cint,
        (MPI_Comm, Ptr{MPI_Group}), comm, newgroup)
    finalizer(free, newgroup)
    newgroup
end

"""
Comm_remote_group(comm::Comm)

Accesses the remote group associated with the given inter-communicator.

# External links
$(_doc_external("MPI_Comm_remote_group"))
"""
function Comm_remote_group(comm::Comm)
    newgroup = Group()
    @mpichk ccall((:MPI_Comm_remote_group, libmpi), Cint,
        (MPI_Comm, Ptr{MPI_Group}), comm, newgroup)
    finalizer(free, newgroup)
    newgroup
end

"""
    Comm_create(comm::Comm, group::Group)

Collectively creates a new communicator.

# See also
- [`MPI.Comm_create_group`](@ref) for the noncollective operation

# External links
$(_doc_external("MPI_Comm_create"))
"""
function Comm_create(comm::Comm, group::Group)
    newcomm = Comm()
    @mpichk ccall((:MPI_Comm_create, libmpi), Cint,
          (MPI_Comm, MPI_Group, Ptr{MPI_Comm}),
          comm, group, newcomm)
    finalizer(free, newcomm)
    newcomm
end

"""
    Comm_create_group(comm::Comm, group::Group, tag::Integer)

Noncollectively creates a new communicator.

# See also
- [`MPI.Comm_create`](@ref) for the noncollective operation

# External links
$(_doc_external("MPI_Comm_create_group"))
"""
function Comm_create_group(comm::Comm, group::Group, tag::Integer)
    newcomm = Comm()
    @mpichk ccall((:MPI_Comm_create_group, libmpi), Cint,
        (MPI_Comm, MPI_Group, Cint, Ptr{MPI_Comm}), comm, group, tag, newcomm)
    finalizer(free, newcomm)
    newcomm
end

"""
    Comm_dup(comm::Comm)

# External links
$(_doc_external("MPI_Comm_dup"))
"""
function Comm_dup(comm::Comm)
    newcomm = Comm()
    @mpichk ccall((:MPI_Comm_dup, libmpi), Cint, (MPI_Comm, Ptr{MPI_Comm}), comm, newcomm)
    finalizer(free, newcomm)
    newcomm
end

"""
    Comm_split(comm::Comm, color::Union{Integer,Nothing}, key::Integer)

Partition the communicator `comm`, one for each value of `color`, returning a
new communicator. Within each group, the processes are ranked in the order of
`key`, with ties broken by the order of `comm`.

`color` should be a non-negative integer, or `nothing`, in which case a null
communicator is returned for that rank.

# External links
$(_doc_external("MPI_Comm_split"))
"""
function Comm_split(comm::Comm, color::Union{Integer, Nothing}, key::Integer)
    if isnothing(color)
        color = Consts.MPI_UNDEFINED[]
    end
    newcomm = Comm()
    @mpichk ccall((:MPI_Comm_split, libmpi), Cint,
        (MPI_Comm, Cint, Cint, Ptr{MPI_Comm}), comm, color, key, newcomm)
    finalizer(free, newcomm)
    newcomm
end

mutable struct SplitType
    val::Cint
end
const COMM_TYPE_SHARED = SplitType(-1)
add_load_time_hook!(() -> COMM_TYPE_SHARED.val     = Consts.MPI_COMM_TYPE_SHARED[])

"""
    Comm_split_type(comm::Comm, split_type, key::Integer; kwargs...)

Partitions the communicator `comm` based on `split_type`, returning a new
communicator. Within each group, the processes are ranked in the order of
`key`, with ties broken by the order of `comm`.

Currently only one `split_type` is provided:

- `MPI.COMM_TYPE_SHARED`: splits the communicator into subcommunicators, each of
  which can create a shared memory region.

# External links
$(_doc_external("MPI_Comm_split_type"))
"""
function Comm_split_type(comm::Comm, split_type, key::Integer; kwargs...)
    if isnothing(split_type)
        split_type = Consts.MPI_UNDEFINED[]
    elseif split_type isa SplitType
        split_type = split_type.val
    end
    newcomm = Comm()
    @mpichk ccall((:MPI_Comm_split_type, libmpi), Cint,
          (MPI_Comm, Cint, Cint, MPI_Info, Ptr{MPI_Comm}),
          comm, split_type, key, Info(kwargs...), newcomm)
    finalizer(free, newcomm)
    newcomm
end

"""
    Comm_get_parent()

# External links
$(_doc_external("MPI_Comm_get_parent"))
"""
function Comm_get_parent()
    comm = Comm()
    @mpichk ccall((:MPI_Comm_get_parent, libmpi), Cint, (Ptr{MPI_Comm},), comm)
    comm
end

"""
    Comm_spawn(command, argv::Vector{String}, nprocs::Integer, comm::Comm[, errors::Vector{Cint}]; kwargs...)

# External links
$(_doc_external("MPI_Comm_spawn"))
"""
function Comm_spawn(command::String, argv::Vector{String}, nprocs::Integer,
                    comm::Comm, errors = Vector{Cint}(undef, nprocs); kwargs...)
    intercomm = Comm()
    # int MPI_Comm_spawn(const char *command, char *argv[], int maxprocs,
    #                    MPI_Info info, int root, MPI_Comm comm, MPI_Comm *intercomm,
    #                    int array_of_errcodes[])
    @mpichk ccall((:MPI_Comm_spawn, libmpi), Cint,
         (Cstring, Ptr{Ptr{Cchar}}, Cint, MPI_Info, Cint, MPI_Comm, Ptr{MPI_Comm}, Ptr{Cint}),
          command, argv, nprocs, Info(kwargs...), 0, comm, intercomm, errors)
    finalizer(free, intercomm)
    return intercomm
end

"""
    Intercomm_merge(intercomm::Comm, flag::Bool)

# External links
$(_doc_external("MPI_Intercomm_merge"))
"""
function Intercomm_merge(intercomm::Comm, flag::Bool)
    newcomm = Comm()
    @mpichk ccall((:MPI_Intercomm_merge, libmpi), Cint,
        (MPI_Comm, Cint, Ptr{MPI_Comm}), intercomm, Cint(flag), newcomm)
    finalizer(free, newcomm)
    return newcomm
end

function unsafe_get_attr(comm::Comm, keyval::Integer)
    flag = Ref{Cint}()
    result = Ref(C_NULL)
    # int MPI_Comm_get_attr(MPI_Comm comm, int comm_keyval, void *attribute_val, int *flag)
    @mpichk ccall((:MPI_Comm_get_attr, libmpi), Cint,
        (MPI_Comm, Cint, Ptr{Cvoid}, Ptr{Cint}), comm, keyval, result, flag)
    if flag[] == 0
        return nothing
    end
    return result[]
end

"""
    universe_size()

The total number of available slots, or `nothing` if it is not defined. This is determined by the `MPI_UNIVERSE_SIZE` attribute of `COMM_WORLD`.

This is typically dependent on the MPI implementation: for MPICH-based implementations, this is specified by the `-usize` argument. OpenMPI defines a default value based on the number of processes available.
"""
function universe_size()
    ptr = unsafe_get_attr(COMM_WORLD, Consts.MPI_UNIVERSE_SIZE[])
    isnothing(ptr) && return nothing
    return Int(unsafe_load(Ptr{Cint}(ptr)))
end

"""
    tag_ub()

The maximum value tag value for point-to-point operations.
"""
function tag_ub()
    ptr = something(unsafe_get_attr(COMM_WORLD, Consts.MPI_TAG_UB[]))
    return Int(unsafe_load(Ptr{Cint}(ptr)))
end

"""
    Comm_compare(comm1::Comm, comm2::Comm)::MPI.Comparison

Compare two communicators and their underlying groups, returning an element of the [`Comparison`](@ref) enum.

# External links
$(_doc_external("MPI_Comm_compare"))
"""
function Comm_compare(comm1::Comm, comm2::Comm)
    result = Ref{Cint}()
    @mpichk ccall((:MPI_Comm_compare, libmpi), Cint,
                  (MPI_Comm, MPI_Comm, Ptr{Cint}), comm1, comm2, result)
    return Comparison(result[])
end


text:
# # Heat Equation + Slab Tutorial

# In this tutorial, we demonstrate simple sequential coupling
# of two PDE models using the `ClimaCore.jl` backends.

# # Model 1
# Model 1 represents a simplified atmosphere (atm)
# and solves the [heat
# equation](https://en.wikipedia.org/wiki/Heat_equation)
# in a one-column domain:

# ``
# \frac{∂ T}{∂ t} + ∇ ⋅ (-μ ∇T) = 0
# ``

#    with top and bottom boundary conditions set to fixed-temperature (non-zero Dirichlet) and fixed-flux (non-zero Neumann) conditions, respectively:

# ``
#    T_{top}  = 280 K, \,\,\,\,\,\,\, \frac{∂ T_{bottom}}{∂ t} = - ∇ F_{sfc}
# ``

# where
#  - `t` is time
#  - `μ` is the thermal diffusivity
#  - `T` is the temperature
#  - `F_sfc` is the thermal boundary flux (see below for calculation)

# # Model 2
# Model 2 assumes a simplified soil (lnd) domain to be a slab, represented by an ODE for its surface temperature:

# ``
#    \frac{dT_{sfc}}{dt} = - (F_{accumulated} + G ) / h_{lnd}
# ``

#    where

# ``
#    F_{accumulated} = {F_{integrated}} / Δt_{coupler}
# ``

# where
#  - `Δt_coupler` is the duration of the coupling cycle
#  - `T_sfc` is the temperature
#  - `h_lnd` is the slab thickness
#  - `F_integrated` and `F_accumulated` thermal boundary fluxes, respectively (see below for calculation)
#  - `G` represents soil physics, which we assume to be 0 for now

# # Coupling and Flux Calculation

# We use this Model 1 (usually this is done by the model with the shortest timestep)
# to calculate and accumulate the downward surface fluxes, `F_sfc`:

# ``
#    F_{sfc} = - λ (T_{sfc} - T1)
# ``

# ``
#    d(F_{integrated})/dt  = F_{sfc}
# ``

# where
#  - `T1` is the atm temperature near the surface (here assumed equal to the first model level)
#  - `λ` a constant relaxation timescale

# Note that in a more realistic setup the above equations would be weighted by their domains' densities and thermal heat capacities, so that the thermal flux would have the units of W m$^{-2}$.
# Here we assume these are unity for both domains.

# Sequential coupling has the following steps:
# 1) pre-Model 1: supply Model 1 with `T_sfc` for the `F_sfc` calculation; reset `F_integrated` to zero
# 2) run Model 1: step forward for all Model 1 timesteps within one coupling cycle using `F_sfc` as the bottom boundary condition; accumulate `F_integrated` at each (sub-)step
# 3) post-Model 1: pass `F_integrated` into coupler and convert to `F_accumulated` for the correct units.
# 4) pre-Model 2: supply Model 2 with `F_accumulated`
# 5) run Model 2: step forward for all Model 2 timesteps within one coupling cycle;
# 6) post-Model 2: state variable, `T_sfc` of Model 2 into coupler.
# 7) repeat steps 1-6 for all coupling timesteps.

# # Implementation

# ## [Loading Packages](@id Loading-code-coupler)

# First, we'll load our pre-requisites:
#  - load CliMA packages under development - you may need to add unregistered packages in Pkg, e.g.:
## import Pkg; Pkg.add(url="https://github.com/CliMA/ClimaCore.jl",rev="main")

#  - load external packages:
import LinearAlgebra, UnPack
import ClimaCore: Fields, Domains, Topologies, Meshes, DataLayouts, Operators, Geometry, Spaces

using Base: show_supertypes
using OrdinaryDiffEq: ODEProblem, solve, SSPRK33

using Logging: global_logger
using TerminalLoggers: TerminalLogger

using RecursiveArrayTools

using OrdinaryDiffEq

using Statistics

#src ## Setup Logging Information
global_logger(TerminalLogger()) #src
const CI = !isnothing(get(ENV, "CI", nothing)) #src

# ## Define Parameters
#  - Global Constants
const FT = Float64;

#  - Experiment-specific Parameters
parameters = (
    ## atmos parameters
    zmin_atm = FT(0.0), # height of atm stack bottom [m]
    zmax_atm = FT(1.0), # height of atm stack top [m]
    n = 15,  # number of elements in atm stack
    μ = FT(0.0001), # diffusion coefficient [m^2 / s]
    T_top = FT(280.0), # fixed temperature at the top of the domain_atm [K]
    T_atm_ini = FT(280.0), # initial condition of at temperature (isothermal) [K]
    ## slab parameters
    h_lnd = FT(0.5), # depth of slab layer [m]
    T_lnd_ini = FT(260.0), # initial condition of at temperature (isothermal) [K]
    ## coupling parameters
    λ = FT(1e-5), # transfer coefficient
)

# ## Define Model Functions

# - Model 1 (atm) Equations
"""
    ∑tendencies_atm!(du, u, (parameters, T_sfc), t)

Heat diffusion equation
    dT/dt =  ∇ μ ∇ T
    where
        T  = 280 K              at z = zmax_atm
        dT/dt = - ∇ F_sfc       at z = zmin_atm

We also use this model to calculate and accumulate the downward surface fluxes, F_sfc:
    F_sfc = - λ * (T_sfc - T1)
    d(F_integrated)/dt  = F_sfc
    where
        F_integrated is reset to 0 at the beginning of each coupling cycle
        T1 = atm temperature near the surface (here assumed equal to the first model level)
"""
function ∑tendencies_atm!(du, u, (parameters, T_sfc), t)
    T = u.x[1] # u.x = vector of prognostic variables from DifferentialEquations
    F_sfc = calculate_flux(T_sfc[1], parent(T)[1], parameters)
    ## set BCs
    bcs_bottom = Operators.SetValue(Geometry.Cartesian3Vector(F_sfc)) # F_sfc is converted to a Cartesian vector in direction 3 (vertical)
    bcs_top = Operators.SetValue(FT(parameters.T_top))

    gradc2f = Operators.GradientC2F(top = bcs_top) # Dirichlet BC (center-to-face)
    gradf2c = Operators.DivergenceF2C(bottom = bcs_bottom) # Neumann BC (face-to-center)

    ## tendency calculations
    @. du.x[1] = gradf2c(parameters.μ * gradc2f(T)) # dT/dt
    du.x[2] .= -F_sfc[1] # d(F_integrated)/dt
end

# - Model 2 (lnd) Equations
"""
    ∑tendencies_lnd!(dT_sfc, T_sfc, (parameters, F_accumulated), t)

Slab layer equation
    lnd d(T_sfc)/dt = - (F_accumulated + G) / h_lnd
    where
        F_accumulated = F_integrated / Δt_coupler
"""
function ∑tendencies_lnd!(dT_sfc, T_sfc, (parameters, F_accumulated), t)
    G = 0.0 # place holder for soil dynamics
    @. dT_sfc = (-F_accumulated + G) / parameters.h_lnd
end

# - Surface Flux Calculation (coarse bulk formula)
calculate_flux(T_sfc, T1, parameters) = -parameters.λ * (T_sfc - T1);

# - Coupler Communication Functions
# These functions export / import / transform variables
# These functions are now just place holders for coupler transformations (e.g. regridding, masking, etc)
coupler_get_(x) = x;
coupler_put_(x) = x;

# ## Model Initialization
# - initialize atm model domain and grid
domain_atm = Domains.IntervalDomain(
    Geometry.ZPoint{FT}(parameters.zmin_atm),
    Geometry.ZPoint{FT}(parameters.zmax_atm);
    boundary_tags = (:bottom, :top),
);
mesh_atm = Meshes.IntervalMesh(domain_atm, nelems = parameters.n); # struct, allocates face boundaries to 5,6: atmos
center_space_atm = Spaces.CenterFiniteDifferenceSpace(mesh_atm); # collection of the above, discretises space into FD and provides coords

# - initialize prognostic variables, either as ClimaCore's Field objects or as Arrays
T_atm_0 = Fields.ones(FT, center_space_atm) .* parameters.T_atm_ini; # initiates a spatially uniform atm progostic var
T_lnd_0 = [parameters.T_lnd_ini]; # initiates lnd progostic var
ics = (; atm = T_atm_0, lnd = T_lnd_0)

# - specify timestepping information
stepping = (;
    Δt_min = 0.02,
    timerange = (0.0, 6.0),
    Δt_coupler = 1.0,
    odesolver = SSPRK33(),
    nsteps_atm = 8, # number of timesteps of atm per coupling cycle
    nsteps_lnd = 1, # number of timesteps of lnd per coupling cycle
);

# ## Define the sequential coupling loop
function coupler_solve!(stepping, ics, parameters)
    t = 0.0
    Δt_min = stepping.Δt_min
    Δt_coupler = stepping.Δt_coupler
    t_start = stepping.timerange[1]
    t_end = stepping.timerange[2]

    ## init coupler fields
    coupler_F_sfc = [0.0]
    coupler_T_lnd = copy(ics.lnd)

    ## atmos copies of coupler variables
    atm_T_lnd = copy(coupler_T_lnd)
    atm_F_sfc = copy(coupler_F_sfc)

    ## SETUP ATMOS
    ## put all prognostic variable arrays into a vector and ensure that solve can partition them
    T_atm = ics.atm
    Y_atm = ArrayPartition((T_atm, atm_F_sfc))
    prob_atm = ODEProblem(∑tendencies_atm!, Y_atm, (t_start, t_end), (parameters, atm_T_lnd))
    integ_atm = init(prob_atm, stepping.odesolver, dt = Δt_min, saveat = 10 * Δt_min)

    ## land copies of coupler variables
    T_lnd = ics.lnd
    lnd_F_sfc = copy(coupler_F_sfc)

    ## SETUP LAND
    prob_lnd = ODEProblem(∑tendencies_lnd!, T_lnd, (t_start, t_end), (parameters, lnd_F_sfc))
    integ_lnd = init(prob_lnd, stepping.odesolver, dt = Δt_min, saveat = 10 * Δt_min)

    ## coupler stepping
    for t in (t_start:Δt_coupler:t_end)

        ## STEP ATMOS
        ## pre_atmos
        integ_atm.p[2] .= coupler_get_(coupler_T_lnd) # integ_atm.p is the parameter vector of an ODEProblem from DifferentialEquations
        integ_atm.u.x[2] .= [0.0] # surface flux to be accumulated

        ## run atmos
        ## NOTE: use (t - integ_atm.t) here instead of Δt_coupler to avoid accumulating roundoff error in our timestepping.
        OrdinaryDiffEq.step!(integ_atm, t - integ_atm.t, true)

        ## post_atmos
        coupler_F_sfc .= coupler_put_(integ_atm.u.x[2]) / Δt_coupler

        ## STEP LAND
        ## pre_land
        lnd_F_sfc .= coupler_get_(coupler_F_sfc)

        ## run land
        OrdinaryDiffEq.step!(integ_lnd, t - integ_lnd.t, true)

        ## post land
        coupler_T_lnd .= coupler_put_(integ_lnd.u) # update T_sfc
    end

    return integ_atm, integ_lnd
end;

# ## Run the Coupler Model Simulation
integ_atm, integ_lnd = coupler_solve!(stepping, ics, parameters);
sol_atm, sol_lnd = integ_atm.sol, integ_lnd.sol;

# ## Postprocessing and Visualization

# Each integrator output (`sol_atm`, `sol_lnd`), contains the DifferentialEquations variable `.u` (the name is hard coded).
# If `ArrayPartition` was used for combining multiple prognostic variables, `u` will include an additional variable `x` (also hard coded)
# `parent()` accesses the `Field` values.
# So, for example, the structure of `u` from Model 1 is:

# `parent(sol_atm.u[<time-index>].x[<ArrayPartition-index>])[<z-index>,<variable-index>]`

ENV["GKSwstype"] = "nul"
import Plots
Plots.GRBackend()

show_plots = isdefined(Main, :SHOWPLOTS) ? SHOWPLOTS : true

path = joinpath(dirname(@__FILE__), "images/")
mkpath(path);

# - Vertical profile at start and end
t0_ = parent(sol_atm.u[1].x[1])[:, 1];
tend_ = parent(sol_atm.u[end].x[1])[:, 1];
z_centers = parent(Fields.coordinate_field(center_space_atm))[:, 1];
show_plots ?
Plots.png(
    Plots.plot(
        [t0_ tend_],
        z_centers,
        title = "model 1: atm",
        labels = ["t=0" "t=end"],
        xlabel = "T (K)",
        ylabel = "z (m)",
    ),
    joinpath(path, "tc1_f1.png"),
) : nothing
# ![](images/tc1_f1.png)

# - Conservation: absolute "energy" of both models with time
# convert to the same units (analogous to energy conservation, assuming that is both domains density=1 and thermal capacity=1)
lnd_sfc_u_t = [u[1] for u in sol_lnd.u] .* parameters.h_lnd;
atm_sum_u_t = [sum(parent(u.x[1])[:]) for u in sol_atm.u] .* (parameters.zmax_atm - parameters.zmin_atm) ./ parameters.n;
v1 = lnd_sfc_u_t .- lnd_sfc_u_t[1];
v2 = atm_sum_u_t .- atm_sum_u_t[1];
show_plots ?
Plots.png(
    Plots.plot(
        sol_lnd.t,
        [v1 v2 v1 + v2],
        labels = ["lnd" "atm" "tot"],
        xlabel = "time (s)",
        ylabel = "pseudo-energy (J / m2)",
    ),
    joinpath(path, "tc1_f2.png"),
) : nothing
# ![](images/tc1_f2.png)

# - Conservation: relative error with time
total = atm_sum_u_t + lnd_sfc_u_t;
rel_error = (total .- total[1]) / mean(total);
show_plots ?
Plots.png(
    Plots.plot(sol_lnd.t, rel_error, labels = ["tot"], xlabel = "time (s)", ylabel = "relative error"),
    joinpath(path, "tc1_f3.png"),
) : nothing
# ![](images/tc1_f3.png)

#src # - Animation
#src anim = Plots.@animate for u in sol_atm.u
#src     Plots.plot(u.x[1], xlim=(220,280))
#src end
#src Plots.mp4(anim, joinpath(path, "heat.mp4"), fps = 10)
#src
#src function linkfig(figpath, alt = "")
#src     # buildkite-agent upload figpath
#src     # link figure in logs if we are running on CI
#src     if get(ENV, "BUILDKITE", "") == "true"
#src         artifact_url = "artifact://$figpath"
#src         print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
#src     end
#src end
#src
#src dirname = "heat"
#src linkfig("output/$(dirname)/heat_end.png", "Heat End Simulation")

max_stars_repo_name: CliMA/ClimaCoupler


text:
# Simple script that automates setting up the documentation project environment
# before building the docs.
import Pkg
Pkg.activate(@__DIR__)
Pkg.develop([
    Pkg.PackageSpec(path = normpath(joinpath(@__DIR__, ".."))),
    Pkg.PackageSpec(
        path = normpath(joinpath(@__DIR__, "..", "lib", "ClimaCoreVTK")),
    ),
    Pkg.PackageSpec(
        path = normpath(joinpath(@__DIR__, "..", "lib", "ClimaCorePlots")),
    ),
    Pkg.PackageSpec(
        path = normpath(joinpath(@__DIR__, "..", "lib", "ClimaCoreMakie")),
    ),
])
Pkg.instantiate()
include(joinpath(@__DIR__, "make.jl"))


text:
using Test

using FileIO: Stream, File, @format_str, load, save
using OpticalFlowUtils

@testset "io" begin
    data = UInt32[0x48454950, 0x00000002, 0x00000001,
        0x00000000, 0x3f800000, 0x7f7fffff, 0x00000000]

    file, io = mktemp()
    write(io, data)
    close(io)

    x = load(file)
    @test size(x) == (2, 1, 2)
    @test all(x[:, 1, 1] .== (1., 0))
    @test all(ismissing, x[:, 1, 2])

    save(File{format"FLO"}(file), x)
    saveddata = reinterpret(eltype(data), read(file))
    # last entry will be written as being missing
    @test saveddata[1:end-1] == data[1:end-1]
    @test saveddata[end] == 0x7f7fffff
end


text:
#    Copyright (c) 2021 J.A. Duffek
#    Copyright (c) 2000 D.M. Spink
#
#    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.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, see <http://www.gnu.org/licenses/>.

"""
    vectrans(tvec::Vector{F})::Matrix{F} where {F<:AbstractFloat}

Transformation matrix for a translation.

Returns a (4x4) Transformation matrix for translation.

The matrix is:

    [ 1   0   0  tx ]
    [ 0   1   0  ty ]
    [ 0   0   1  tz ]
    [ 0   0   0   1 ]

# Arguments:
- `tvec`: A vectors defining the translation along the x,y and z axes.
  i.e. `[tx; ty; ty]`

# Output:
- `st`: Translation Transformation Matrix

# Examples:
```julia
julia> st = vectrans(tvec)
```

Translate the NURBS line `[0.0;0.0;0.0]` - `[1.0;1.0;1.0]` by 3 along the
x-axis, 2 along the y-axis and 4 along the z-axis.

```julia
julia> line = nrbline(vec([0.0 0.0 0.0]),vec([1.0 1.0 1.0]));
julia> trans = vectrans(vec([3.0 2.0 4.0]));
julia> tline = nrbtform(line, trans);
```

# Reference
- see also: [`nrbtform`](@ref)
"""
function vectrans(tvec::Vector{F})::Matrix{F} where {F<:AbstractFloat}
if length(tvec) < 1 || length(tvec) > 3
  throw(ArgumentError("Translation vector required"));
end
v = [tvec;0;0];
st = [1.0   0   0 v[1];
        0 1.0   0 v[2];
        0   0 1.0 v[3];
        0   0   0 1.0];
return st;
end # vectrans


text:
"""
    ParamGroups(grouper, m)

A logical grouping of parameters in `m` created by [`ParamGrouper`](#)
`grouper`. Parameters in `m` are assigned a group that can be queried
using `getgroup(paramgroups, param)`. If a parameter is not assigned a
group, `getgroup` returns `nothing`.

## Examples

```julia
using Flux: Chain, Dense, params
using FastAI: ParamGroups, IndexGrouper, getgroup

model = Chain(Dense(3, 5), Dense(5, 3))
paramgroups = ParamGroups(IndexGrouper([1, 2]), model)

getgroup(paramgroups, model[1].weight) == 1
getgroup(paramgroups, model[2].weight) == 2
getgroup(paramgroups, rand(10)) === nothing
```
"""
struct ParamGroups
    map::IdDict
end

ParamGroups() = ParamGroups(IdDict())

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

getgroup(pg::ParamGroups, x::AbstractArray) = get(pg.map, x, nothing)

function assigngroups!(pg::ParamGroups, grouper, m)
    for (group, m_) in group(grouper, m)
        for p in params(m_)
            pg.map[p] = group
        end
    end
end

abstract type ParamGrouper end

struct IndexGrouper <: ParamGrouper
    idxs
end

group(grouper::IndexGrouper, m) = Dict(i => m[is] for (i, is) in enumerate(grouper.idxs))

function ParamGroups(grouper::ParamGrouper, m)
    pg = ParamGroups()
    assigngroups!(pg, grouper, m)
    return pg
end

# ## Tests

@testset "ParamGroups" begin
    model = Chain(Dense(3, 5), Dense(5, 3))
    paramgroups = ParamGroups(IndexGrouper([1, 2]), model)

    @test getgroup(paramgroups, model[1].weight) == 1
    @test getgroup(paramgroups, model[2].weight) == 2
    @test getgroup(paramgroups, rand(10)) === nothing
end


text:
header("DOCUMENT")

@testset "Show" begin
str =
    """
<?xml version="1.0"?><!-- https://github.com/daemontus/pnml-parser -->
<pnml xmlns="http://www.pnml.org/version-2009/grammar/pnml">
  <net id="small-net" type="http://www.pnml.org/version-2009/grammar/ptnet">
    <name> <text>P/T Net with one place</text> </name>
    <page id="page0">
      <place id="place1">
        <name> <text>Some place</text> </name>
        <initialMarking> <text>100</text> </initialMarking>
      </place>
      <transition id="transition1">
        <name> <text>Some transition </text> </name>
      </transition>
      <arc source="transition1" target="place1" id="arc1">
        <inscription> <text>12 </text> </inscription>
      </arc>
      <arc source="place1" target="transition1" id="arc2">
        <inscription> <text> 13 </text> </inscription>
      </arc>
    </page>
  </net>
</pnml>
    """
    pnml_ir = parse_pnml(root(parsexml(str)); reg=PNML.IDRegistry())
    @test typeof(pnml_ir) <: PNML.PnmlModel
    @show pnml_ir
end

header("### Registry")
@testset "Document & IDRegistry" begin
    str = """
    <?xml version="1.0"?>
    <pnml xmlns="http://www.pnml.org/version-2009/grammar/pnml">
      <net id="net" type="pnmlcore"> <page id="page"/> </net>
    </pnml>
    """
    reg = PNML.IDRegistry()
    @test !PNML.isregistered(reg, :net)
    @test :net ∉ reg.ids

    parse_pnml(root(parsexml(str)); reg)
    @test_call parse_pnml(root(parsexml(str)); reg)
    #@show reg

    @test PNML.isregistered(reg, :net)
    @test :net ∈ reg.ids
end

@testset "multiple net type" begin
    str = """
    <?xml version="1.0"?>
    <pnml xmlns="http://www.pnml.org/version-2009/grammar/pnml">
      <net id="net1" type="http://www.pnml.org/version-2009/grammar/ptnet"> <page id="page1"/> </net>
      <net id="net2" type="pnmlcore"> <page id="page2"/> </net>
      <net id="net3" type="ptnet"> <page id="page3"/> </net>
      <net id="net4" type="hlcore"> <page id="page4"/> </net>
      <net id="net5" type="pt_hlpng"> <page id="page5"/> </net>
    </pnml>
    """

    @show model = parse_str(str)

    v1 = PNML.find_nets(model, :ptnet)
    printnode(v1, label="v1")
    @test_call PNML.PnmlTypes.pnmltype(:ptnet)
    foreach(v1) do net
        @test net.type === PNML.PnmlTypes.pnmltype(:ptnet)
    end
    v2 = PNML.find_nets(model, "ptnet")
    printnode(v2, label="v2")
    foreach(v2) do net
        @test net.type === PNML.PnmlTypes.pnmltype(:ptnet)
    end

    @test v1 == v2
    @test length(v1) == 2

    v3 = PNML.find_nets(model, :pnmlcore)
    printnode(v3, label="v3")
    foreach(v3) do net
        @test net.type === PNML.PnmlTypes.pnmltype(:pnmlcore)
    end

    @test !isempty(v3)
    @test v3 != v1

    @testset for t in [:ptnet, :pnmlcore, :hlcore, :pt_hlpng, :hlnet, :symmetric, :stochastic, :timednet]
        foreach(PNML.find_nets(model, t)) do net
            @test net.type === PNML.PnmlTypes.pnmltype(t)
        end
    end
end

@testset "Empty" begin
    str = """
    <?xml version="1.0"?>
    <pnml xmlns="http://www.pnml.org/version-2009/grammar/pnml">
      <net id="net" type="pnmlcore">
        <page id="page">
        </page>
      </net>
    </pnml>
    """

    model = parse_str(str)
    @test model isa PNML.PnmlModel
end


text:
"""
    function is_constraint_violated(
        com::CoM,
        constraint::BoolConstraint,
        fct,
        set::OrSet,
    )

Check if both of the inner constraints are violated
"""
function is_constraint_violated(
    com::CoM,
    constraint::BoolConstraint,
    fct,
    set::OrSet,
)
    return is_lhs_constraint_violated(com, constraint) && is_rhs_constraint_violated(com, constraint)
end

"""
    still_feasible(com::CoM, constraint::OrConstraint, fct, set::OrSet, vidx::Int, value::Int)

Return whether the constraint can be still fulfilled when setting a variable with index `vidx` to `value`.
**Attention:** This assumes that it isn't violated before.
"""
function still_feasible(
    com::CoM,
    constraint::OrConstraint,
    fct,
    set::OrSet,
    vidx::Int,
    value::Int,
)
    lhs_feasible = !is_constraint_violated(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set)
    if lhs_feasible
        lhs_indices = constraint.lhs.indices
        for i in 1:length(lhs_indices)
            if lhs_indices[i] == vidx
                lhs_feasible = still_feasible(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set, vidx, value)
                lhs_feasible && return true
                break
            end
        end
    end
    rhs_feasible = !is_constraint_violated(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set)
    if rhs_feasible
        rhs_indices = constraint.rhs.indices
        for i in 1:length(rhs_indices)
            if rhs_indices[i] == vidx
                rhs_feasible = still_feasible(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set, vidx, value)
                rhs_feasible && return true
                break
            end
        end
    end
    return rhs_feasible || lhs_feasible
end

"""
    prune_constraint!(com::CS.CoM, constraint::OrConstraint, fct, set::OrSet; logs = true)

Reduce the number of possibilities given the `OrConstraint` by pruning both parts
Return whether still feasible
"""
function prune_constraint!(
    com::CS.CoM,
    constraint::OrConstraint,
    fct,
    set::OrSet;
    logs = true,
)
    lhs_violated = is_constraint_violated(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set)
    rhs_violated = is_constraint_violated(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set)
    if lhs_violated && rhs_violated
        return false
    end
    if lhs_violated
        activate_rhs!(com, constraint)
        return prune_constraint!(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set; logs=logs)
    end
    if rhs_violated
        activate_lhs!(com, constraint)
        return prune_constraint!(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set; logs=logs)
    end
    return true
end


text:
@reexport module SquareLattice
export Square

using Lattices: BoundaryCondition, Fixed, Periodic, BoundedLattice
import Lattices

const SiteType = Tuple{Int, Int}
const EdgeType = Tuple{SiteType, SiteType}

"""
    Square{Shape, BC} <: BoundedLattice{2, BC}

Square lattice.
"""
struct Square{Shape, BC <: BoundaryCondition} <: BoundedLattice{2, BC} end

# alias
const FixedSquare{h, w} = Square{Tuple{h, w}, Fixed} where {h, w}
const PeriodicSquare{h, w} = Square{Tuple{h, w}, Periodic} where {h, w}

"""
    Square(height, width; [boundary=Periodic])

Returns a square lattice with given `height` and `width`, `boundary` is set to `Periodic`
by default.
"""
Square(height::Int, width::Int; boundary=Periodic) = Square{Tuple{height, width}, boundary}()
"""
    Square(extent; [boundary=Periodic])

Returns a square lattice with height and width set to `extent`.
"""
Square(extent::Int; boundary=Periodic) = Square{Tuple{extent, extent}, boundary}()

Base.size(::Square{Tuple{height, width}}) where {height, width} = (height, width)
Base.length(::Square{Tuple{height, width}}) where {height, width} = height * width
Base.nameof(::Square) = "Square Lattice"

Base.getindex(lattice::Square, inds::Int...) = getindex(lattice, inds)
Base.getindex(lattice::FixedSquare{h, w}, inds::NTuple{2, Int}) where h where w =
    inds[1] > h ? throw(BoundsError(lattice, inds)) :
    inds[2] > w ? throw(BoundsError(lattice, inds)) :
    inds[1] + (inds[2] - 1) * h
Base.getindex(lattice::PeriodicSquare{h, w}, inds::NTuple{2, Int}) where h where w =
    mod1(inds[1], h) + mod(inds[2] - 1, w) * h

(ltc::FixedSquare{h, w})(s::Int) where {h, w} =
    s > length(ltc) || s < 1 ? throw(BoundsError(ltc, s)) :
    mod1(s, h), div(s - 1, h) + 1

(ltc::PeriodicSquare{h, w})(s::Int) where {h, w} =
    mod1(s, h), div(s - 1, h) + 1

# iterators
include("sites.jl")
include("edges.jl")
# include("surround.jl")

end


text:
T=3
eps_exp=(1.)/(3.)
tsname="joBlockDiag - basic"
@testset "$tsname" begin
for t=1:T # start test loop

    s=rand(3:9,3)
    a=rand(ComplexF64,s[1],s[1])
    A=joMatrix(a;DDT=ComplexF32,RDT=ComplexF64)
    b=rand(ComplexF64,s[2],s[2])
    B=joMatrix(b;DDT=ComplexF32,RDT=ComplexF64)
    c=rand(ComplexF64,s[3],s[3])
    C=joMatrix(c;DDT=ComplexF32,RDT=ComplexF64)
    bd=cat(a,b,c,dims=[1 2])
    BD=joBlockDiag(A,B,C)
    BDm=BD.m
    BDn=BD.n

    verbose && println("$tsname ($BDm,$BDn)")
    @testset "$BDm x $BDn" begin
        vn=rand(ComplexF32,BDn);
        vm=rand(ComplexF64,BDm);
        mvn=rand(ComplexF32,BDn,2);
        mvm=rand(ComplexF64,BDm,2);
        @test norm(bd*vn-BD*vn)<eps(norm(bd*vn))^eps_exp
        @test norm(bd*mvn-BD*mvn)<eps(norm(bd*mvn))^eps_exp
        @test norm(transpose(bd)*vm-transpose(BD)*vm)<eps(norm(transpose(bd)*vm))^eps_exp
        @test norm(transpose(bd)*mvm-transpose(BD)*mvm)<eps(norm(transpose(bd)*mvm))^eps_exp
        @test norm(adjoint(bd)*vm-adjoint(BD)*vm)<eps(norm(adjoint(bd)*vm))^eps_exp
        @test norm(adjoint(bd)*mvm-adjoint(BD)*mvm)<eps(norm(adjoint(bd)*mvm))^eps_exp
        @test norm(conj(bd)*vn-conj(BD)*vn)<eps(norm(conj(bd)*vn))^eps_exp
        @test norm(conj(bd)*mvn-conj(BD)*mvn)<eps(norm(conj(bd)*mvn))^eps_exp
        @test norm((-bd)*vn-(-BD)*vn)<eps(norm((-bd)*vn))^eps_exp
        @test norm((-bd)*mvn-(-BD)*mvn)<eps(norm((-bd)*mvn))^eps_exp
    end

end # end test loop
end
tsname="joBlockDiag - weighted"
@testset "$tsname" begin
for t=1:T # start test loop

    s=rand(3:9,3)
    w=rand(ComplexF64,3)
    a=rand(ComplexF64,s[1],s[1])
    A=joMatrix(a;DDT=ComplexF32,RDT=ComplexF64)
    b=rand(ComplexF64,s[2],s[2])
    B=joMatrix(b;DDT=ComplexF32,RDT=ComplexF64)
    c=rand(ComplexF64,s[3],s[3])
    C=joMatrix(c;DDT=ComplexF32,RDT=ComplexF64)
    bd=cat(w[1]*a,w[2]*b,w[3]*c,dims=[1 2])
    BD=joBlockDiag(A,B,C;weights=w)
    BDm=BD.m
    BDn=BD.n

    verbose && println("$tsname ($BDm,$BDn)")
    @testset "$BDm x $BDn" begin
        vn=rand(ComplexF32,BDn);
        vm=rand(ComplexF64,BDm);
        mvn=rand(ComplexF32,BDn,2);
        mvm=rand(ComplexF64,BDm,2);
        @test norm(bd*vn-BD*vn)<eps(norm(bd*vn))^eps_exp
        @test norm(bd*mvn-BD*mvn)<eps(norm(bd*mvn))^eps_exp
        @test norm(transpose(bd)*vm-transpose(BD)*vm)<eps(norm(transpose(bd)*vm))^eps_exp
        @test norm(transpose(bd)*mvm-transpose(BD)*mvm)<eps(norm(transpose(bd)*mvm))^eps_exp
        @test norm(adjoint(bd)*vm-adjoint(BD)*vm)<eps(norm(adjoint(bd)*vm))^eps_exp
        @test norm(adjoint(bd)*mvm-adjoint(BD)*mvm)<eps(norm(adjoint(bd)*mvm))^eps_exp
        @test norm(conj(bd)*vn-conj(BD)*vn)<eps(norm(conj(bd)*vn))^eps_exp
        @test norm(conj(bd)*mvn-conj(BD)*mvn)<eps(norm(conj(bd)*mvn))^eps_exp
        @test norm((-bd)*vn-(-BD)*vn)<eps(norm((-bd)*vn))^eps_exp
        @test norm((-bd)*mvn-(-BD)*mvn)<eps(norm((-bd)*mvn))^eps_exp
    end

end # end test loop
end
tsname="joBlockDiag - replicated"
@testset "$tsname" begin
for t=1:T # start test loop

    s=rand(3:9)
    l=3
    w=rand(ComplexF64,l)
    a=rand(ComplexF64,s,s)
    A=joMatrix(a;DDT=ComplexF32,RDT=ComplexF64)
    bd=cat(w[1]*a,w[2]*a,w[3]*a,dims=[1 2])
    BD=joBlockDiag(l,A;weights=w)
    BDm=BD.m
    BDn=BD.n

    verbose && println("$tsname ($BDm,$BDn)")
    @testset "$BDm x $BDn" begin
        vn=rand(ComplexF32,BDn);
        vm=rand(ComplexF64,BDm);
        mvn=rand(ComplexF32,BDn,2);
        mvm=rand(ComplexF64,BDm,2);
        @test norm(bd*vn-BD*vn)<eps(norm(bd*vn))^eps_exp
        @test norm(bd*mvn-BD*mvn)<eps(norm(bd*mvn))^eps_exp
        @test norm(transpose(bd)*vm-transpose(BD)*vm)<eps(norm(transpose(bd)*vm))^eps_exp
        @test norm(transpose(bd)*mvm-transpose(BD)*mvm)<eps(norm(transpose(bd)*mvm))^eps_exp
        @test norm(adjoint(bd)*vm-adjoint(BD)*vm)<eps(norm(adjoint(bd)*vm))^eps_exp
        @test norm(adjoint(bd)*mvm-adjoint(BD)*mvm)<eps(norm(adjoint(bd)*mvm))^eps_exp
        @test norm(conj(bd)*vn-conj(BD)*vn)<eps(norm(conj(bd)*vn))^eps_exp
        @test norm(conj(bd)*mvn-conj(BD)*mvn)<eps(norm(conj(bd)*mvn))^eps_exp
        @test norm((-bd)*vn-(-BD)*vn)<eps(norm((-bd)*vn))^eps_exp
        @test norm((-bd)*mvn-(-BD)*mvn)<eps(norm((-bd)*mvn))^eps_exp
    end

end # end test loop
end


text:
export BitString

typealias BitString UInt64


text:
function check(numbers, idx)
    for i in idx - 25:idx - 2
        for j in i:idx - 1
            if numbers[i] + numbers[j] == numbers[idx]
                return true
            end
        end
    end
    return false
end

function run(s)
    numbers = readlines(IOBuffer(s))
    numbers = [parse(Int, n) for n in numbers]
    for idx in 26:length(numbers)
        s = numbers[idx]
        if !check(numbers, idx)
            return numbers[idx]
        end
    end
    return 0
end

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

function main()
    res, time, memory = @timed run(ARGS[1])
    println("_duration:$(time * 1000)")
    println(res)
end

main()


text:
module SingleStackUtils

export get_vars_from_nodal_stack,
    get_vars_from_element_stack,
    get_horizontal_variance,
    get_horizontal_mean,
    reduce_nodal_stack,
    reduce_element_stack,
    horizontally_average!,
    dict_of_nodal_states

using OrderedCollections
using StaticArrays
import KernelAbstractions: CPU

using ..BalanceLaws
using ..DGMethods
using ..DGMethods.Grids
using ..MPIStateArrays
using ..VariableTemplates

"""
    get_vars_from_nodal_stack(
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars;
        vrange::UnitRange = 1:size(Q, 3),
        i::Int = 1,
        j::Int = 1,
        exclude::Vector{String} = String[],
        interp = false,
    ) where {T, dim, N}

Return a dictionary whose keys are the `flattenednames()` of the variables
specified in `vars` (as returned by e.g. `vars_state`), and
whose values are arrays of the values for that variable along the vertical
dimension in `Q`. Only a single element is expected in the horizontal as
this is intended for the single stack configuration and `i` and `j` identify
the horizontal nodal coordinates.

Variables listed in `exclude` are skipped.
"""
function get_vars_from_nodal_stack(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars;
    vrange::UnitRange = 1:size(Q, 3),
    i::Int = 1,
    j::Int = 1,
    exclude::Vector{String} = String[],
    interp = false,
) where {T, dim, N}

    # extract grid information and bring `Q` to the host if needed
    FT = eltype(Q)
    Nq = N .+ 1
    # Code assumes the same polynomial order in all horizontal directions
    @inbounds begin
        Nq1 = Nq[1]
        Nq2 = Nq[2]
        Nqk = dim == 2 ? 1 : Nq[dim]
    end
    Np = dofs_per_element(grid)
    state_data = array_device(Q) isa CPU ? Q.realdata : Array(Q.realdata)

    # set up the dictionary to be returned
    var_names = flattenednames(vars)
    stack_vals = OrderedDict()
    num_vars = varsize(vars)
    vars_wanted = Int[]
    @inbounds for vi in 1:num_vars
        if !(var_names[vi] in exclude)
            stack_vals[var_names[vi]] = FT[]
            push!(vars_wanted, vi)
        end
    end
    elemtobndy = convert(Array, grid.elemtobndy)
    vmap⁻ = convert(Array, grid.vmap⁻)
    vmap⁺ = convert(Array, grid.vmap⁺)
    vgeo = convert(Array, grid.vgeo)
    # extract values from `state_data`
    @inbounds for ev in vrange, k in 1:Nqk, v in vars_wanted
        if interp && k == 1 && elemtobndy[5, ev] == 0
            # Get face degree of freedom number
            n = i + Nq1 * ((j - 1))
            # get the element numbers
            ev⁻ = ev
            # Get neighboring id data
            id⁻, id⁺ = vmap⁻[n, 5, ev⁻], vmap⁺[n, 5, ev⁻]
            ev⁺ = ((id⁺ - 1) ÷ Np) + 1
            # get the volume degree of freedom numbers
            vid⁻, vid⁺ = ((id⁻ - 1) % Np) + 1, ((id⁺ - 1) % Np) + 1

            J⁺, J⁻ = vgeo[vid⁻, Grids._M, ev⁻], vgeo[vid⁺, Grids._M, ev⁺]
            state_local = J⁻ * state_data[vid⁻, v, ev⁻]
            state_local += J⁺ * state_data[vid⁺, v, ev⁺]
            state_local /= (J⁻ + J⁺)
            push!(stack_vals[var_names[v]], state_local)
        elseif interp && k == Nqk && elemtobndy[6, ev] == 0
            # Get face degree of freedom number
            n = i + Nq1 * ((j - 1))
            # get the element numbers
            ev⁻ = ev
            # Get neighboring id data
            id⁻, id⁺ = vmap⁻[n, 6, ev⁻], vmap⁺[n, 6, ev⁻]
            # periodic and need to handle this point (otherwise handled above)
            if id⁺ == id⁻
                vid⁻ = ((id⁻ - 1) % Np) + 1

                state_local = state_data[vid⁻, v, ev⁻]
                push!(stack_vals[var_names[v]], state_local)
            end
        else
            ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))
            state_local = state_data[ijk, v, ev]
            push!(stack_vals[var_names[v]], state_local)
        end
    end

    return stack_vals
end

"""
    get_vars_from_element_stack(
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars;
        vrange::UnitRange = 1:size(Q, 3),
        exclude::Vector{String} = String[],
        interp = false,
    ) where {T, dim, N}

Return an array of [`get_vars_from_nodal_stack()`](@ref)s whose dimensions
are the number of nodal points per element in the horizontal plane.

Variables listed in `exclude` are skipped.
"""
function get_vars_from_element_stack(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars;
    vrange::UnitRange = 1:size(Q, 3),
    exclude::Vector{String} = String[],
    interp = false,
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds Nq1 = Nq[1]
    @inbounds Nq2 = Nq[2]

    return [
        get_vars_from_nodal_stack(
            grid,
            Q,
            vars,
            vrange = vrange,
            i = i,
            j = j,
            exclude = exclude,
            interp = interp,
        ) for i in 1:Nq1, j in 1:Nq2
    ]
end

"""
    get_horizontal_mean(
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars;
        vrange::UnitRange = 1:size(Q, 3),
        exclude::Vector{String} = String[],
        interp = false,
    ) where {T, dim, N}

Return a dictionary whose keys are the `flattenednames()` of the variables
specified in `vars` (as returned by e.g. `vars_state`), and
whose values are arrays of the horizontal averages for that variable along
the vertical dimension in `Q`. Only a single element is expected in the
horizontal as this is intended for the single stack configuration.

Variables listed in `exclude` are skipped.
"""
function get_horizontal_mean(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars;
    vrange::UnitRange = 1:size(Q, 3),
    exclude::Vector{String} = String[],
    interp = false,
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds Nq1 = Nq[1]
    @inbounds Nq2 = Nq[2]

    vars_avg = OrderedDict()
    vars_sq = OrderedDict()
    for i in 1:Nq1
        for j in 1:Nq2
            vars_nodal = get_vars_from_nodal_stack(
                grid,
                Q,
                vars,
                vrange = vrange,
                i = i,
                j = j,
                exclude = exclude,
                interp = interp,
            )
            vars_avg = merge(+, vars_avg, vars_nodal)
        end
    end
    map!(x -> x ./ Nq1 / Nq1, values(vars_avg))
    return vars_avg
end

"""
    get_horizontal_variance(
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars;
        vrange::UnitRange = 1:size(Q, 3),
        exclude::Vector{String} = String[],
        interp = false,
    ) where {T, dim, N}

Return a dictionary whose keys are the `flattenednames()` of the variables
specified in `vars` (as returned by e.g. `vars_state`), and
whose values are arrays of the horizontal variance for that variable along
the vertical dimension in `Q`. Only a single element is expected in the
horizontal as this is intended for the single stack configuration.

Variables listed in `exclude` are skipped.
"""
function get_horizontal_variance(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars;
    vrange::UnitRange = 1:size(Q, 3),
    exclude::Vector{String} = String[],
    interp = false,
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds Nq1 = Nq[1]
    @inbounds Nq2 = Nq[2]

    vars_avg = OrderedDict()
    vars_sq = OrderedDict()
    for i in 1:Nq1
        for j in 1:Nq2
            vars_nodal = get_vars_from_nodal_stack(
                grid,
                Q,
                vars,
                vrange = vrange,
                i = i,
                j = j,
                exclude = exclude,
                interp = interp,
            )
            vars_nodal_sq = OrderedDict(vars_nodal)
            map!(x -> x .^ 2, values(vars_nodal_sq))
            vars_avg = merge(+, vars_avg, vars_nodal)
            vars_sq = merge(+, vars_sq, vars_nodal_sq)
        end
    end
    map!(x -> (x ./ Nq1 / Nq1) .^ 2, values(vars_avg))
    map!(x -> x ./ Nq1 / Nq1, values(vars_sq))
    vars_var = merge(-, vars_sq, vars_avg)
    return vars_var
end

"""
    reduce_nodal_stack(
        op::Function,
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars::NamedTuple,
        var::String;
        vrange::UnitRange = 1:size(Q, 3),
    ) where {T, dim, N}

Reduce `var` from `vars` within `Q` over all nodal points in the specified
`vrange` of elements with `op`. Return a tuple `(result, z)` where `result` is
the final value returned by `op` and `z` is the index within `vrange` where the
`result` was determined.
"""
function reduce_nodal_stack(
    op::Function,
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars::Type,
    var::String;
    vrange::UnitRange = 1:size(Q, 3),
    i::Int = 1,
    j::Int = 1,
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds begin
        Nq1 = Nq[1]
        Nq2 = Nq[2]
        Nqk = dim == 2 ? 1 : Nq[dim]
    end

    var_names = flattenednames(vars)
    var_ind = findfirst(s -> s == var, var_names)
    if var_ind === nothing
        return
    end

    state_data = array_device(Q) isa CPU ? Q.realdata : Array(Q.realdata)
    z = vrange.start
    result = state_data[i + Nq1 * (j - 1), var_ind, z]
    for ev in vrange
        for k in 1:Nqk
            ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))
            new_result = op(result, state_data[ijk, var_ind, ev])
            if !isequal(new_result, result)
                result = new_result
                z = ev
            end
        end
    end

    return (result, z)
end

"""
    reduce_element_stack(
        op::Function,
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars::NamedTuple,
        var::String;
        vrange::UnitRange = 1:size(Q, 3),
    ) where {T, dim, N}

Reduce `var` from `vars` within `Q` over all nodal points in the specified
`vrange` of elements with `op`. Return a tuple `(result, z)` where `result` is
the final value returned by `op` and `z` is the index within `vrange` where the
`result` was determined.
"""
function reduce_element_stack(
    op::Function,
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars::Type,
    var::String;
    vrange::UnitRange = 1:size(Q, 3),
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds Nq1 = Nq[1]
    @inbounds Nq2 = Nq[2]

    return [
        reduce_nodal_stack(
            op,
            grid,
            Q,
            vars,
            var,
            vrange = vrange,
            i = i,
            j = j,
        ) for i in 1:Nq1, j in 1:Nq2
    ]
end

"""
    horizontally_average!(
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        i_vars,
    ) where {T, dim, N}

Horizontally average variables, from variable
indexes `i_vars`, in `MPIStateArray` `Q`.

!!! note
    These are not proper horizontal averages-- the main
    purpose of this method is to ensure that there are
    no horizontal fluxes for a single stack configuration.
"""
function horizontally_average!(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    i_vars,
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds begin
        Nq1 = Nq[1]
        Nq2 = Nq[2]
        Nqk = dim == 2 ? 1 : Nq[dim]
    end

    ArrType = typeof(Q.data)
    state_data = array_device(Q) isa CPU ? Q.realdata : Array(Q.realdata)

    for ev in 1:size(state_data, 3), k in 1:Nqk, i_v in i_vars
        Q_sum = 0
        for i in 1:Nq1, j in 1:Nq2
            Q_sum += state_data[i + Nq1 * ((j - 1) + Nq2 * (k - 1)), i_v, ev]
        end
        Q_ave = Q_sum / (Nq1 * Nq2)
        for i in 1:Nq1, j in 1:Nq2
            ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))
            state_data[ijk, i_v, ev] = Q_ave
        end
    end
    Q.realdata .= ArrType(state_data)
end

get_data(solver_config, ::Prognostic) = solver_config.Q
get_data(solver_config, ::Auxiliary) = solver_config.dg.state_auxiliary
get_data(solver_config, ::GradientFlux) = solver_config.dg.state_gradient_flux

"""
    dict_of_nodal_states(
        solver_config,
        state_types = (Prognostic(), Auxiliary());
        aux_excludes = [],
        interp = false,
        )

A dictionary of single stack prognostic and auxiliary
variables at the `i=1`,`j=1` node given
 - `solver_config` a `SolverConfiguration`
 - `aux_excludes` a vector of strings containing the
    variables to exclude from the auxiliary state.
"""
function dict_of_nodal_states(
    solver_config,
    state_types = (Prognostic(), Auxiliary());
    aux_excludes = String[],
    interp = false,
)
    FT = eltype(solver_config.Q)
    all_state_vars = []
    for st in state_types
        state_vars = get_vars_from_nodal_stack(
            solver_config.dg.grid,
            get_data(solver_config, st),
            vars_state(solver_config.dg.balance_law, st, FT),
            exclude = st isa Auxiliary ? aux_excludes : String[],
            interp = interp,
        )
        push!(all_state_vars, state_vars...)
    end
    return OrderedDict(all_state_vars...)
end

end # module

max_forks_repo_name: jm-c/CLIMA


text:
"""
  msm_slices(sMMProblem::MSMProblem, paramValues::Vector; nbPoints::Int64 = 5, offset::Float64 = 0.001)

Function to "slice" the objective function. That is, to hold the variables constant,
except for one dimension.
"""
function msm_slices(sMMProblem::MSMProblem, paramValues::Vector; nbPoints::Int64 = 5, offset::Float64 = 0.001)

    # Create a grid in the neighborhood of the minimizer
    lb_slice = paramValues .- abs.(paramValues) .* offset;
    ub_slice = paramValues .+ abs.(paramValues) .* offset;

    #Checks
    if nbPoints < 3
      error("nbPoints must be >= 3")
    end
    if mod(nbPoints,2) == 0
      error("nbPoints must be odd")
    end

    nbPointsBelow = round(Int, ceil(nbPoints/2))
    nbPointAbove = nbPoints - nbPointsBelow

    vXGrid = zeros(nbPoints, length(keys(sMMProblem.priors)))
    vYGrid = zeros(nbPoints, length(keys(sMMProblem.priors)))

    # Loop over parameter values
    #---------------------------
    for (keyIndex, keyValue) in enumerate(keys(sMMProblem.priors))

        #Create one grid below the minimizer (which also contains the minimizer)
        grid_below = linspace(lb_slice[keyIndex], paramValues[keyIndex], nbPointsBelow);
        #Create one grid above the minimizer (which does not contain the minimizer)
        grid_above = linspace(paramValues[keyIndex], ub_slice[keyIndex], nbPointsBelow);
        grid_above = grid_above[2:end]; #exclude the minimizer from the grid (already in grid_below)
        # store grid points:
        vXGrid[:,keyIndex] = vcat(grid_below, grid_above);
        vYGrid[:,keyIndex] = zeros(nbPoints)

        info("slicing along $(keyValue)")

        # Move along one dimension, keep other values constant
        localParamValues = transpose(repeat(paramValues,outer=[1,nbPoints]))
        localParamValues[:, keyIndex] = vXGrid[:, keyIndex]

        # Use pmap to use several workers in parallel:
        vYGrid[:, keyIndex] = pmap(sMMProblem.objective_function, eachrow(localParamValues))

    end

    return vXGrid, vYGrid

end

#= OLD VERSION THAT USES SharedArrays
"""
  msm_slices(sMMProblem::MSMProblem, paramValues::Vector; nbPoints::Int64 = 5, offset::Float64 = 0.001)
Function to "slice" the objective function. That is, to hold the variables constant,
except for one dimension.
"""
function msm_slices(sMMProblem::MSMProblem, paramValues::Vector; nbPoints::Int64 = 5, offset::Float64 = 0.001)

    # Create a grid in the neighborhood of the minimizer
    lb_slice = paramValues .- abs.(paramValues) .* offset;
    ub_slice = paramValues .+ abs.(paramValues) .* offset;

    #Checks
    if nbPoints < 5
      error("nbPoints must be >= 5")
    end
    if mod(nbPoints,2) == 0
      error("nbPoints must be odd")
    end

    vXGrid = SharedArray(zeros(round(Int, 2*floor(nbPoints/2) - 1), length(keys(sMMProblem.priors))))
    vYGrid = SharedArray(zeros(round(Int, 2*floor(nbPoints/2) - 1), length(keys(sMMProblem.priors))))

    # Loop over parameter values
    #---------------------------
    for (keyIndex, keyValue) in enumerate(keys(sMMProblem.priors))

        #Create one grid below the minimizer (which also contains the minimizer)
        grid_below = linspace(lb_slice[keyIndex], paramValues[keyIndex], round(Int, floor(nbPoints/2)));
        #Create one grid abive the minimizer (which does not contain the minimizer)
        grid_above = linspace(paramValues[keyIndex], ub_slice[keyIndex], round(Int, floor(nbPoints/2)));
        grid_above = grid_above[2:end]; #exclude the minimizer from the grid (already in grid_below)
        # store grid points:
        vXGrid[:,keyIndex] = vcat(grid_below, grid_above);

        #vXGrid[:,keyIndex] = collect(linspace(sMMProblem.priors[keyValue][2], sMMProblem.priors[keyValue][3], nbPoints))
        vYGrid[:,keyIndex] = zeros(round(Int, 2*floor(nbPoints/2) - 1))

        info("slicing along $(keyValue)")
        # Options A.
        #-----------
        @sync @distributed for xIndex = 1:1:size(vXGrid,1)

          # Move along one dimension, keeping other values constant
          #--------------------------------------------------------
          localParamValues = copy(paramValues)
          localParamValues[keyIndex] = copy(vXGrid[xIndex, keyIndex])

          vYGrid[xIndex, keyIndex] = sMMProblem.objective_function(localParamValues)

        end

    end

    return vXGrid, vYGrid

end
=#