text:
# mmtd_inference.jl

export ParamsMMTD, PriorMMTD, ModMMTD, λindxMMTD,
  build_λ_indx, ZζtoZandζ, sim_mmtd, symmetricDirPrior_mmtd, transTensor_mmtd,
  counttrans_mmtd,
  mcmc!, timemod!,
  bfact_MC, forecDist_MMTD, llik_MMTD,
  addDecomp_postsims!;

mutable struct ParamsMMTD
  lΛ::Vector{Float64}
  lλ::Vector{Vector{Float64}}
  Zζ::Vector{Int} # will be length TT - L, mapped through ZζtoZandζ(), values start at 0 (intercept)
  lQ0::Union{Vector{Float64}, Nothing}
  lQ::Union{Vector{<:Array{Float64}}, Nothing, Any} # organized so first index is now and lag 1 is the next index; the Any is for compatibility with BSON
  # lQ::Union{Vector{Array{Float64}}, Nothing}
  # lQ::Any # for compatibility with BSON

  ParamsMMTD(lΛ, lλ, Zζ, lQ0, lQ) = new(deepcopy(lΛ), deepcopy(lλ), deepcopy(Zζ),
    deepcopy(lQ0), deepcopy(lQ))
end

mutable struct PriorMMTD
  Λ::Union{Vector{Float64}, SparseDirMix, SBMprior}
  λ::Vector{Union{Vector{Float64}, SparseDirMix, SBMprior}} # may be of mixed type
  Q0::Vector{Float64}
  Q::Union{Vector{<:Array{Float64}}, Vector{<:Array{SparseDirMix}}, Vector{<:Array{SBMprior}}, Any} # MCMC functions depend on this type; the Any is for compatibility with BSON
  # Q::Vector{ Union{ Array{Float64}, Array{SparseDirMix}, Array{SBMprior} } }
  # Q::Any  # for compatibility with BSON

  PriorMMTD(Λ, λ, Q0, Q) = new(deepcopy(Λ), deepcopy(λ), deepcopy(Q0), deepcopy(Q))
end

mutable struct λindxMMTD
    indxs::Vector{Array{Array{Int64,1},1}}
    lens::Vector{Int64}
    Zζindx::Matrix{Int64}
    nZζ::Int64

    λindxMMTD(indxs, lens, Zζindx, nZζ) = new(deepcopy(indxs),
        deepcopy(lens), deepcopy(Zζindx), nZζ)
end

mutable struct ModMMTD
  L::Int # maximal order
  R::Int # largest order considered
  K::Int # number of states
  TT::Int
  S::Vector{Int}
  prior::PriorMMTD
  state::ParamsMMTD
  λ_indx::λindxMMTD
  iter::Int

  ModMMTD(L, R, K, TT, S, prior, state, λ_indx) = new(L, R, K, TT, deepcopy(S),
    deepcopy(prior), deepcopy(state), deepcopy(λ_indx), 0)
end

"""
    build_λ_indx(L, R)

### Example
```julia
λ_indx = build_λ_indx(3, 2)
```
"""
function build_λ_indx(L::Int, R::Int)
  indxs = [ collect(Combinatorics.combinations(1:L, r)) for r in 1:R ]
  lens = [ length(xx) for xx in indxs ]
  ζindx = vcat([collect(1:zz) for zz in lens]...)
  Zindx = vcat([fill(r, lens[r]) for r in 1:R]...)
  nZζ = sum(lens)
  Zζindx = hcat(Zindx, ζindx)
  λindxMMTD(indxs, lens, Zζindx, nZζ)
end

function ZζtoZandζ(Zζ::Vector{Int64}, λ_indx)
    out = Matrix{Int64}(undef, length(Zζ), 2)
    for tt in 1:length(Zζ)
        if Zζ[tt] == 0
            out[tt, :] = [0, 0]
        else
            out[tt, :] = deepcopy(λ_indx.Zζindx[Zζ[tt],:])
        end
    end
    out
end

"""
    sim_mmtd(TT, nburn, L, R, K, λ_indx, Λ, λ, Q0, Q)
"""
function sim_mmtd(TT::Int, nburn::Int, L::Int, R::Int, K::Int, λ_indx::λindxMMTD,
  Λ::Array{Float64}, λ::Array{Vector{Float64}},
  Q0::Vector{Float64}, Q::Vector{<:Array{Float64}}, Xtras::Bool=false)

  Nsim = nburn + TT

  Z = [ StatsBase.sample(Weights(Λ)) - 1 for i in 1:(Nsim-L) ]
  ζ = [ StatsBase.sample(Weights(λ[m])) for i in 1:(Nsim-L), m in 1:R ]

  S = Vector{Int}(undef, Nsim)
  S[1:L] = StatsBase.sample(1:K, L)

  if Xtras
      Pvecs = Matrix{Float64}(undef, Nsim, K)
      Pvecs[1:L, :] .= 0.0
  end

  for tt in (L+1):(Nsim)
    i = tt - L
    if Z[i] == 0
        pvec = deepcopy(Q0)
    else
        Slagrev_now = deepcopy(S[range(tt-1, step=-1, length=L)])
        pvec = deepcopy( Q[Z[i]][:, Slagrev_now[λ_indx.indxs[Z[i]][ζ[i,Z[i]]]]...] )
    end

    if Xtras
        Pvecs[tt,:] = deepcopy(pvec)
    end
    S[tt] = StatsBase.sample(Weights( pvec ))
  end

  if Xtras
      S[(nburn+1):(Nsim)], Z[(nburn+1):(Nsim-L)], ζ[(nburn+1):(Nsim-L),:], Pvecs[(nburn+1):(Nsim),:]
  else
      S[(nburn+1):(Nsim)], Z[(nburn+1):(Nsim-L)], ζ[(nburn+1):(Nsim-L),:]
  end
end

"""
    symmetricPrior_mmtd(size_Λ, size_λ, size_Q0, size_Q, L, R, K, λ_indx)
"""
function symmetricDirPrior_mmtd(size_Λ::Float64, size_λ::Float64,
    size_Q::Float64, size_Q0::Float64,
  L::Int, R::Int, K::Int, λ_indx::λindxMMTD)

  # λ_indx.lens contains a tuple of lengths of λ vectors

  α0_Λ = fill( size_Λ / float(R), R)
  α0_λ = [ fill( size_λ / float(λ_indx.lens[r]), λ_indx.lens[r] ) for r in 1:R ]
  α0_Q0 = fill( size_Q0 / float(K), K )
  a0_Q = size_Q / float(K)
  α0_Q = [ fill( a0_Q, fill(K, (r+1))... ) for r in 1:R ]

  (α0_Λ, α0_λ, α0_Q0, α0_Q)
end

"""
    transTensor_mmtd(L, R, K, λ_indx, Λ, λ, Q0, Q)

Calculate full transition tensor from Λ, λ, Q0 and Q.
"""
# function transTensor_mmtd(L::Int, R::Int, K::Int, λ_indx::λindxMMTD,
#   Λ::Vector{Float64}, λ::Vector{Vector{Float64}},
#   Q0::Vector{Float64}, Q::Vector{<:Array{Float64}}; login::Bool=false)
#
#   froms, nfroms = create_froms(K, L) # in this case, ordered by now, lag1, lag2, etc.
#   Ωmat = zeros(Float64, (K, nfroms))
#
#   if login # all contributing quantities are on log scale
#     for i in 1:nfroms
#       for k in 1:K
#         Ωmat[k,i] += exp(Λ[1] + Q0[k])
#         for r in 1:R
#           for ℓ in 1:λ_indx.lens[r]
#             Ωmat[k,i] += exp( Λ[r+1] + λ[r][ℓ] + Q[r][k,froms[i][[λ_indx.indxs[r][ℓ]]...]...] )
#           end
#         end
#       end
#     end
#   else
#     for i in 1:nfroms
#       for k in 1:K
#         Ωmat[k,i] += Λ[1] * Q0[k]
#         for r in 1:R
#           for ℓ in 1:λ_indx.lens[r]
#             Ωmat[k,i] += Λ[r+1] .* λ[r][ℓ] .* Q[r][k,froms[i][[λ_indx.indxs[r][ℓ]]...]...]
#           end
#         end
#       end
#     end
#   end
#
#   reshape(Ωmat, fill(K,L+1)...)
# end
# below is 5 times faster
function transTensor_mmtd(L::Int, R::Int, K::Int, λ_indx::λindxMMTD,
  Λ::Vector{Float64}, λ::Vector{Vector{Float64}},
  Q0::Vector{Float64}, Q::Vector{<:Array{Float64}}; login::Bool=false)

  out = zeros(fill(K, L+1)...)

  if login
      A0 = exp.(Λ[1] .+ Q0)
  else
      A0 = Λ[1] .* Q0
  end

  out .+= reshape( vec(A0), [K, fill(1, L)...]... )

  for r in 1:R
      for ℓ in 1:λ_indx.lens[r]

          if login
              A = exp.(Λ[r+1] .+ λ[r][ℓ] .+ Q[r])
          else
              A = Λ[r+1] .* λ[r][ℓ] .* Q[r]
          end

          dims = [K, fill(1, L)...]
          for c in λ_indx.indxs[r][ℓ]
              dims[c+1] = K
          end

          Abdcst = reshape(vec(A), dims...)

          out .+= Abdcst

      end
  end

  return out
end

"""
rpost_lΛ_mmtd(α0_Λ, Z, R)
"""
function rpost_lΛ_mmtd(α0_Λ::Vector{Float64}, Z::Vector{Int}, R::Int)
    Nz = StatsBase.counts(Z, 0:R)
    α1_Λ = α0_Λ .+ Nz
    SparseProbVec.rDirichlet(α1_Λ, logout=true)
end
function rpost_lΛ_mmtd(prior::SparseDirMix, Z::Vector{Int}, R::Int)
    Nz = StatsBase.counts(Z, 0:R)
    α1_Λ = prior.α .+ Nz
    d = SparseProbVec.SparseDirMix(α1_Λ, prior.β)
    SparseProbVec.rand(d, logout=true)
end
function rpost_lΛ_mmtd(prior::SBMprior, Z::Vector{Int}, R::Int)
    Nz = StatsBase.counts(Z, 0:R)
    post_lΛ = SparseProbVec.SBM_multinom_post(prior, Nz)
    SparseProbVec.rand(post_lΛ, logout=true)
end

"""
    rpost_lλ_mmtd(prior, Zandζ, λ_lens, R)
"""
function rpost_lλ_mmtd(prior::Vector{Union{Vector{Float64}, SparseDirMix, SBMprior}},
    Zandζ::Matrix{Int},
    λ_lens::Vector{Int}, R::Int)

    lλ_out = [ Vector{Float64}(undef, λ_lens[r]) for r in 1:R ]
    Z = deepcopy(Zandζ[:,1])
    ζ = deepcopy(Zandζ[:,2])

    for r in 1:R
        Zrindx = findall(Z .== r)
        Nζ = StatsBase.counts(ζ[Zrindx], 1:λ_lens[r])

        if typeof(prior[r]) == Vector{Float64}
            α1_λ = prior[r] .+ Nζ
            lλ_out[r] = SparseProbVec.rDirichlet(α1_λ, logout=true)
        elseif typeof(prior[r]) == SparseDirMix
            α1_λ = prior[r].α .+ Nζ
            d = SparseDirMix(α1_λ, prior[r].β)
            lλ_out[r] = SparseProbVec.rand(d, logout=true)
        elseif typeof(prior[r]) == SBMprior
            post_lλ = SparseProbVec.SBM_multinom_post(prior[r], Nζ)
            lλ_out[r] = SparseProbVec.rand(post_lλ, logout=true)
        end
    end

    lλ_out
end

"""
    counttrans_mmtd(S::Vector{Int}, TT::Int, Zandζ::Matrix{Int},
        λ_indx::λindxMMTD, L::Int, R::Int, K::Int)

    Indexing on ouput is [t, t-1, t-2, etc.]

    ### Example
    ```julia
    L = 2
    R = 2
    K = 3
    TT = 12
    S = [1,2,1,3,3,1,2,1,3,2,1,1]
    Zandζ = [1 2; 0 0; 1 2; 1 1; 2 1; 1 2; 2 1; 2 1; 2 1; 1 1]
    λ_indx = build_λ_indx(L, R)
    counttrans_mmtd(S, TT, Zandζ, λ_indx, L, R, K)
    ```
"""
function counttrans_mmtd(S::Vector{Int}, TT::Int, Zandζ::Matrix{Int},
  λ_indx::λindxMMTD, L::Int, R::Int, K::Int)

  Z = deepcopy(Zandζ[:,1])
  ζ = deepcopy(Zandζ[:,2])

  ## initialize
  N0_out = zeros(Int, K)
  N_out = [ zeros(Int, fill(K, r+1)...) for r in 1:R ]

  ## pass through data and add counts
  for tt in (L+1):(TT)
    Z_now = Z[tt-L]
    if Z_now == 0
        N0_out[ S[tt] ] += 1
    else
        Slagrev_now = S[range(tt-1, step=-1, length=L)]
        N_out[Z_now][append!( [deepcopy(S[tt])], deepcopy(Slagrev_now[ λ_indx.indxs[Z_now][ζ[tt-L]] ]) )...] += 1
    end
  end

  N0_out, N_out
end

"""
    rpost_lQ_mmtd(S, TT, prior, Z, ζ, λ_indx, L, R, K)
"""
function rpost_lQ_mmtd(S::Vector{Int}, TT::Int,
  prior_Q0::Vector{Float64}, prior_Q::Vector{<:Array{Float64}},
  Zandζ::Matrix{Int},
  λ_indx::λindxMMTD, L::Int, R::Int, K::Int)

  N0, N = counttrans_mmtd(S, TT, Zandζ, λ_indx, L, R, K) # N is a vector of arrays

  ## initialize
  α0_Q0 = deepcopy(prior_Q0)
  α0_Q = deepcopy(prior_Q) # vector of arrays
  lQ_mats = [ Matrix{Float64}(undef, K, K^r) for r in 1:R ]
  lQ_out = [ reshape(lQ_mats[r], fill(K, r+1)...) for r in 1:R ]

  α1_Q0 = α0_Q0 .+ N0
  lQ0_out = SparseProbVec.rDirichlet(α1_Q0, logout=true)

  for r in 1:R
    ncol = K^r
    α1_Q = α0_Q[r] .+ N[r]
    α1_Q_mat = reshape(α1_Q, (K, ncol))
    for j in 1:ncol
      lQ_mats[r][:,j] = SparseProbVec.rDirichlet(α1_Q_mat[:,j], logout=true)
    end
    lQ_out[r] = reshape(lQ_mats[r], fill(K, r+1)...)
  end

  lQ0_out, lQ_out
end
function rpost_lQ_mmtd(S::Vector{Int}, TT::Int,
  prior_Q0::Vector{Float64},
  prior_Q::Vector{<:Array{SparseDirMix}},
  Zandζ::Matrix{Int},
  λ_indx::λindxMMTD, L::Int, R::Int, K::Int)

  N0, N = counttrans_mmtd(S, TT, Zandζ, λ_indx, L, R, K)

  prior_Q_vec = [ reshape(prior_Q[r], K^r ) for r in 1:R ]

  ## initialize
  α0_Q0 = deepcopy(prior_Q0)
  lQ_mats = [ Matrix{Float64}(undef, K, K^r) for r in 1:R ]
  lQ_out = [ reshape(lQ_mats[r], fill(K, r+1)...) for r in 1:R ]

  α1_Q0 = α0_Q0 .+ N0
  lQ0_out = SparseProbVec.rDirichlet(α1_Q0, logout=true)

  for r in 1:R
    ncol = K^r
    Nmat = reshape(N[r], (K, ncol))
    for j in 1:ncol
      α1 = prior_Q_vec[r][j].α .+ Nmat[:,j]
      d = SparseProbVec.SparseDirMix(α1, prior_Q_vec[r][j].β)
      lQ_mats[r][:,j] = SparseProbVec.rand(d, logout=true)
    end
    lQ_out[r] = reshape(lQ_mats[r], fill(K, r+1)...)
  end

  lQ0_out, lQ_out
end
function rpost_lQ_mmtd(S::Vector{Int}, TT::Int,
    prior_Q0::Vector{Float64},
    prior_Q::Vector{<:Array{SBMprior}},
    Zandζ::Matrix{Int}, λ_indx::λindxMMTD, L::Int, R::Int, K::Int)

  N0, N = counttrans_mmtd(S, TT, Zandζ, λ_indx, L, R, K)

  prior_Q_vec = [ reshape(prior[r], K^r ) for r in 1:R ]

  ## initialize
  α0_Q0 = deepcopy(prior_Q0)
  lQ_mats = [ Matrix{Float64}(undef, K, K^r) for r in 1:R ]
  lQ_out = [ reshape(lQ_mats[r], fill(K, r+1)...) for r in 1:R ]

  α1_Q0 = α0_Q0 .+ N0
  lQ0_out = SparseProbVec.rDirichlet(α1_Q0, logout=true)

  for r in 1:R
    ncol = K^r
    Nmat = reshape(N[r], (K, ncol))
    for j in 1:ncol
        d = SparseProbVec.SBM_multinom_post(prior_Q_vec[r][j], Nmat[:,j])
        lQ_mats[r][:,j] = SparseProbVec.rand(d, logout=true)
    end
    lQ_out[r] = reshape(lQ_mats[r], fill(K, r+1)...)
  end

  lQ0_out, lQ_out
end

"""
    rpost_Zζ_fc(S, Zζ_old, λ_indx, lΛ, lλ, lQ0, lQ, TT, L, K)

    Full conditinal updates for Zζ
"""
function rpost_Zζ_fc(S::Vector{Int}, Zζ_old::Vector{Int}, λ_indx::λindxMMTD,
    lΛ::Vector{Float64},
    lλ::Vector{Vector{Float64}},
    lQ0::Vector{Float64},
    lQ::Vector{<:Array{Float64}},
    TT::Int, L::Int, R::Int, K::Int)

  Zζ_out = deepcopy(Zζ_old)
  Zj, ζj = deepcopy(λ_indx.Zζindx[:,1]), deepcopy(λ_indx.Zζindx[:,2])

  for i in 1:(TT-L)  # i indexes Zζ, tt indexes S
    tt = i + L
    Slagrev_now = S[range(tt-1, step=-1, length=L)]

    lw = zeros(Float64, λ_indx.nZζ+1)

    lw[1] += lΛ[1] + lQ0[S[tt]]

    for j in 1:λ_indx.nZζ

      lw[j+1] += lΛ[ Zj[j] ] + lλ[ Zj[j] ][ ζj[j] ]
      lw[j+1] += lQ[ Zj[j] ][ vcat(S[tt], Slagrev_now[ λ_indx.indxs[ Zj[j] ][ ζj[j] ] ]  )... ]

    end

    w = exp.( lw .- maximum(lw) )

    ## Metropolized
    w_cand = deepcopy(w)
    w_cand[Zζ_old[i]+1] = 0.0

    Zζ_cand = sample(StatsBase.Weights(w_cand)) - 1
    lar = logsumexp( lw[1:end .!= (Zζ_old[i]+1)] ) - logsumexp( lw[1:end .!= (Zζ_cand+1)] )

    if log(rand()) < lar
        Zζ_out[i] = deepcopy(Zζ_cand) # already 0-based
    end # otherwise it stays the same

  end

  Zζ_out
end

"""
    rpost_Zζ_marg(S, Zζ_old, λ_indx, prior_Q, lΛ, lλ, TT, L, K)

    Full conditinal updates for Zζ marginalizing over Q
"""
function rpost_Zζ_marg(S::Vector{Int}, Zζ_old::Vector{Int}, λ_indx::λindxMMTD,
    prior_Q0::Vector{Float64}, # always Dirichlet on intercept
    prior_Q::Vector{<:Array{Float64}},
    lΛ::Vector{Float64},
    lλ::Vector{Vector{Float64}},
    TT::Int, L::Int, R::Int, K::Int)

  α0_Q0 = deepcopy(prior_Q0)
  α0_Q = deepcopy(prior_Q) # vector of arrays

  Zζ_out = deepcopy(Zζ_old)
  Zandζ_now = ZζtoZandζ(Zζ_out, λ_indx) # won't get updated

  ## calculate current log marginal likelihood
  N0_now, N_now = counttrans_mmtd(S, TT, Zandζ_now, λ_indx, L, R, K) # vector of arrays, indices in reverse time

  α1_Q0 = α0_Q0 .+ N0_now # gets re-used and updated
  α1_Q = α0_Q .+ N_now # gets re-used and updated
  α1_Q_mats = [ reshape(α1_Q[r], K, K^r) for r in 1:R ] # used only for initial calculation

  llikmarg = SparseProbVec.lmvbeta( α1_Q0 )
  for r in 1:R
      for kk in 1:(K^r)
          llikmarg += SparseProbVec.lmvbeta( α1_Q_mats[r][:,kk] )
      end
  end

  ## for each time point, modify llikmarg accordingly
  for i in 1:(TT-L)  # i indexes Zζ, tt indexes S
    tt = i + L
    Slagrev_now = S[range(tt-1, step=-1, length=L)]
    Zζ_now = deepcopy(Zζ_out[i])

    if Zζ_now == 0
        Zold, ζold = 0, 0
    else
        Zold, ζold = deepcopy(λ_indx.Zζindx[Zζ_now,1]), deepcopy(λ_indx.Zζindx[Zζ_now,2])
    end

    eSt = [ 1.0 * (kk==S[tt]) for kk in 1:K ]

    # remove S_t from llikmarg and α1
    if Zold == 0
        llikmarg -= SparseProbVec.lmvbeta( α1_Q0 )
        α1_Q0 -= eSt
        llikmarg += SparseProbVec.lmvbeta( α1_Q0 )
    else
        SlagrevZζnow = deepcopy(Slagrev_now[ λ_indx.indxs[Zold][ζold] ])
        llikmarg -= SparseProbVec.lmvbeta( α1_Q[Zold][:,SlagrevZζnow...] )
        α1_Q[Zold][:,SlagrevZζnow...] -= eSt
        llikmarg += SparseProbVec.lmvbeta( α1_Q[Zold][:,SlagrevZζnow...] )
    end

    # calculate llikmarg and update probability under each possible Zζ
    llikmarg_cand = fill(llikmarg, λ_indx.nZζ+1)
    lw = zeros(Float64, λ_indx.nZζ+1)

    ## intercept
    llikmarg_cand[1] -= SparseProbVec.lmvbeta( α1_Q0 )
    llikmarg_cand[1] += SparseProbVec.lmvbeta( α1_Q0 .+ eSt )
    lw[1] = llikmarg_cand[1] + lΛ[1]

    ## all others
    for ℓ in 1:λ_indx.nZζ
        Zcand, ζcand = deepcopy(λ_indx.Zζindx[ℓ,1]), deepcopy(λ_indx.Zζindx[ℓ,2])
        SlagrevZζnow = deepcopy(Slagrev_now[ λ_indx.indxs[Zcand][ζcand] ])

        llikmarg_cand[ℓ+1] -= SparseProbVec.lmvbeta( α1_Q[Zcand][:,SlagrevZζnow...] )
        llikmarg_cand[ℓ+1] += SparseProbVec.lmvbeta( α1_Q[Zcand][:,SlagrevZζnow...] .+ eSt )

        lw[ℓ+1] = llikmarg_cand[ℓ+1] + lΛ[Zcand+1] + lλ[Zcand][ζcand]
    end

    w = exp.( lw .- maximum(lw) )

    # newindx = StatsBase.sample(Weights( w )) - 1

    ## Metropolized
    w_cand = deepcopy(w)
    w_cand[Zζ_now+1] = 0.0

    Zζ_cand = sample(StatsBase.Weights(w_cand)) - 1
    lar = logsumexp( lw[1:end .!= (Zζ_now+1)] ) - logsumexp( lw[1:end .!= (Zζ_cand+1)] )

    if log(rand()) < lar
        newindx = deepcopy(Zζ_cand) # already 0-based
      else
        newindx = deepcopy(Zζ_now) # already 0-based
    end

    Zζ_out[i] = deepcopy(newindx)

    ## update running quantities
    if newindx == 0
        Znew, ζnew = 0, 0
        α1_Q0 += eSt
    else
        Znew, ζnew = deepcopy(λ_indx.Zζindx[newindx,1]), deepcopy(λ_indx.Zζindx[newindx,2])
        α1_Q[Znew][:,Slagrev_now[ λ_indx.indxs[Znew][ζnew] ]...] += eSt
    end
    llikmarg = deepcopy(llikmarg_cand[newindx+1])

  end

  Zζ_out
end
function rpost_Zζ_marg(S::Vector{Int}, Zζ_old::Vector{Int},
    λ_indx::λindxMMTD,
    prior_Q0::Vector{Float64}, # always Dirichlet on intercept
    prior_Q::Union{Vector{<:Array{SparseDirMix}}, Vector{<:Array{SBMprior}}},
    lΛ::Vector{Float64},
    lλ::Vector{Vector{Float64}},
    TT::Int, L::Int, R::Int, K::Int)

  α0_Q0 = deepcopy(prior_Q0)

  Zζ_out = deepcopy(Zζ_old)
  Zandζ_now = ZζtoZandζ(Zζ_out, λ_indx) # won't get updated
  prior_Q_vec = [ reshape(prior_Q[r], (K^r)) for r in 1:R ]

  ## calculate current log marginal likelihood
  N0_now, N_now = counttrans_mmtd(S, TT, Zandζ_now, λ_indx, L, R, K) # vector of arrays, indices in reverse time
  N_now_mats = [ reshape(N_now[r], K, K^r) for r in 1:R ] # Nmats won't be updated, only llikmarg, α1_Q0, and N_now

  α1_Q0 = α0_Q0 .+ N0_now # gets re-used and updated

  llikmarg = SparseProbVec.lmvbeta( α1_Q0 )
  for r in 1:R
      for ℓ in 1:(K^r)
          llikmarg += SparseProbVec.logMarginal( prior_Q_vec[r][ℓ], N_now_mats[r][:,ℓ] )
      end
  end

  for i in 1:(TT-L)  # i indexes ζ, tt indexes S
      tt = i + L
      Slagrev_now = S[range(tt-1, step=-1, length=L)]
      Zζ_now = deepcopy(Zζ_out[i])

      if Zζ_now == 0
          Zold, ζold = 0, 0
      else
          Zold, ζold = deepcopy(λ_indx.Zζindx[Zζ_now,1]), deepcopy(λ_indx.Zζindx[Zζ_now,2])
      end

      eSt = [ 1*(kk==S[tt]) for kk in 1:K ]

      # remove S_t from llikmarg
      if Zold == 0
          llikmarg -= SparseProbVec.lmvbeta( α1_Q0 )
          α1_Q0 -= eSt
          llikmarg += SparseProbVec.lmvbeta( α1_Q0 )
      else
          SlagrevZζnow = deepcopy(Slagrev_now[ λ_indx.indxs[Zold][ζold] ])

          llikmarg -= SparseProbVec.logMarginal( prior_Q[Zold][SlagrevZζnow...],
                                      N_now[Zold][:,SlagrevZζnow...] )
          # prior for Q indexed with Slagrev_now[ λ_indx.indxs[Zold][ζold] ]...
          N_now[Zold][:,SlagrevZζnow...] -= eSt
          llikmarg += SparseProbVec.logMarginal( prior_Q[Zold][SlagrevZζnow...],
                                      N_now[Zold][:,SlagrevZζnow...] )
      end

      # calculate llikmarg and update probability under each possible Zζ
      llikmarg_cand = fill(llikmarg, λ_indx.nZζ+1)
      lw = zeros(Float64, λ_indx.nZζ+1)

      ## intercept
      llikmarg_cand[1] -= SparseProbVec.lmvbeta( α1_Q0 )
      llikmarg_cand[1] += SparseProbVec.lmvbeta( α1_Q0 .+ eSt )
      lw[1] = llikmarg_cand[1] + lΛ[1]

      for ℓ in 1:λ_indx.nZζ
          Zcand, ζcand = deepcopy(λ_indx.Zζindx[ℓ,1]), deepcopy(λ_indx.Zζindx[ℓ,2])
          SlagrevZζnow = deepcopy(Slagrev_now[ λ_indx.indxs[Zcand][ζcand] ])

          llikmarg_cand[ℓ+1] -= SparseProbVec.logMarginal( prior_Q[Zcand][SlagrevZζnow...],
                                    N_now[Zcand][:,SlagrevZζnow...] )
          llikmarg_cand[ℓ+1] += SparseProbVec.logMarginal( prior_Q[Zcand][SlagrevZζnow...],
                                    N_now[Zcand][:,SlagrevZζnow...] .+ eSt )

          lw[ℓ+1] = llikmarg_cand[ℓ+1] + lΛ[Zcand+1] + lλ[Zcand][ζcand]
      end

      w = exp.( lw .- maximum(lw) )

      # newindx = StatsBase.sample(Weights( w )) - 1

      ## Metropolized
      w_cand = deepcopy(w)
      w_cand[Zζ_now+1] = 0.0

      Zζ_cand = sample(StatsBase.Weights(w_cand)) - 1
      lar = logsumexp( lw[1:end .!= (Zζ_now+1)] ) - logsumexp( lw[1:end .!= (Zζ_cand+1)] )

      if log(rand()) < lar
          newindx = deepcopy(Zζ_cand) # already 0-based
        else
          newindx = deepcopy(Zζ_now) # already 0-based
      end

      Zζ_out[i] = deepcopy(newindx)

      ## update running quantities
      if newindx == 0
          Znew, ζnew = 0, 0
          α1_Q0 += eSt
      else
          Znew, ζnew = deepcopy(λ_indx.Zζindx[newindx,1]), deepcopy(λ_indx.Zζindx[newindx,2])
          N_now[Znew][:,Slagrev_now[ λ_indx.indxs[Znew][ζnew] ]...] += eSt
      end
      llikmarg = deepcopy(llikmarg_cand[newindx+1])
  end

  Zζ_out
end

"""
    MetropIndep_ΛλZζ(S, lΛ_old, lλ_old, Zζ_old,
        prior_Λ, prior_λ, prior_Q,
        λ_indx,
        TT, L, R, K)

    Independence Metropolis step for Λ, λ, Z, and ζ.
"""
function MetropIndep_ΛλZζ(S::Vector{Int}, lΛ_old::Vector{Float64},
    lλ_old::Vector{Vector{Float64}}, Zζ_old::Vector{Int},
    prior_Λ::Union{Vector{Float64}, SparseDirMix, SBMprior},
    prior_λ::Vector{Union{Vector{Float64}, SparseDirMix, SBMprior}},
    prior_Q0::Vector{Float64}, # always Dirichlet on intercept
    prior_Q::Union{Vector{<:Array{Float64}}, Vector{<:Array{SparseDirMix}}, Vector{<:Array{SBMprior}}},
    λ_indx::λindxMMTD,
    TT::Int, L::Int, R::Int, K::Int)

  if typeof(prior_Λ) == Vector{Float64}
      lΛ_cand = SparseProbVec.rDirichlet(prior_Λ, logout=true)
    elseif typeof(prior_Λ) <: Union{SparseDirMix, SBMprior}
      lΛ_cand = SparseProbVec.rand(prior_Λ, logout=true)
  end

  lλ_cand = [ Vector{Float64}(undef, λ_indx.lens[r]) for r in 1:R ]

  for r in 1:R
      if typeof(prior_λ[r]) == Vector{Float64}
          lλ_cand[r] = SparseProbVec.rDirichlet(prior_λ[r], logout=true)
      elseif typeof(prior_λ[r]) <: Union{SparseDirMix, SBMprior}
          lλ_cand[r] = SparseProbVec.rand(prior_λ[r], logout=true)
      end
  end

  lΛλ_cand = [ lΛ_cand[λ_indx.Zζindx[j,1]+1] + lλ_cand[λ_indx.Zζindx[j,1]][λ_indx.Zζindx[j,2]] for j in 1:λ_indx.nZζ ]
  pushfirst!(lΛλ_cand, lΛ_cand[1])
  Λλ_cand = exp.(lΛλ_cand)

  Zζ_cand = [ StatsBase.sample( Weights( Λλ_cand ) ) - 1 for i in 1:(TT-L) ]
  Zandζ_cand = ZζtoZandζ(Zζ_cand, λ_indx)
  Zandζ_old = ZζtoZandζ(Zζ_old, λ_indx)

  N0_cand, N_cand = counttrans_mmtd(S, TT, Zandζ_cand, λ_indx, L, R, K)
  N0_old, N_old = counttrans_mmtd(S, TT, Zandζ_old, λ_indx, L, R, K)

  N_cand_mats = [ reshape(N_cand[r], K, K^r) for r in 1:R ]
  N_old_mats = [ reshape(N_old[r], K, K^r) for r in 1:R ]

  llikmarg_cand = SparseProbVec.lmvbeta(prior_Q0 .+ N0_cand)
  llikmarg_old = SparseProbVec.lmvbeta(prior_Q0 .+ N0_old)

  if typeof(prior_Q) <: Vector{<:Array{Float64}}
  # if all([ typeof(prior_Q[r]) <: Array{Float64} for r in 1:R ])

      prior_Q_mats = [ reshape(prior_Q[r], K, K^r) for r in 1:R ]

      for r in 1:R
          for ℓ in 1:(K^r)
            α_now = deepcopy(prior_Q_mats[r][:,ℓ])
            llikmarg_cand += SparseProbVec.lmvbeta( N_cand_mats[r][:,ℓ] + α_now )
            llikmarg_old += SparseProbVec.lmvbeta( N_old_mats[r][:,ℓ] + α_now )
          end
      end

  elseif typeof(prior_Q) <: Union{Vector{<:Array{SparseDirMix}}, Vector{<:Array{SBMprior}}}
  # elseif all([ typeof(prior_Q[r]) <: Union{ <:Array{SparseDirMix}, <:Array{SBMprior} } ])

      prior_Q_vec = [ reshape(prior_Q[r], (K^r)) for r in 1:R ]

      for r in 1:R
          for ℓ in 1:(K^r)
              llikmarg_cand += SparseProbVec.logMarginal( prior_Q_vec[r][ℓ],
                                                N_cand_mats[r][:,ℓ] )
              llikmarg_old += SparseProbVec.logMarginal( prior_Q_vec[r][ℓ],
                                                N_old_mats[r][:,ℓ] )
          end
      end

  end

  lu = log(rand())
  if lu < (llikmarg_cand - llikmarg_old)
      lΛ_out = lΛ_cand
      lλ_out = lλ_cand
      Zζ_out = Zζ_cand
  else
      lΛ_out = lΛ_old
      lλ_out = lλ_old
      Zζ_out = Zζ_old
  end

  (lΛ_out, lλ_out, Zζ_out)
end

"""
mcmc!(model, n_keep[, save=true, report_filename="out_progress.txt",
thin=1, jmpstart_iter=25, report_freq=500, monitor::Vector{Symbol}=[:lΛ, :lλ, :lQ0, :lQ], margQ::true ])
"""
function mcmc!(model::ModMMTD, n_keep::Int;
    save::Bool=true,
    report_filename::String="out_progress.txt",
    thin::Int=1, jmpstart_iter::Int=25,
    report_freq::Int=10000,
    monitor::Union{Vector{Symbol}, Nothing}=[:lΛ, :lλ, :lQ0, :lQ, :Zζ],
    margQ::Bool=true)

    ## fix bson issue of recreating model prior structure
    if typeof(model.prior.Q) <: Array{Any}
      model.prior.Q = [ deepcopy(model.prior.Q[r]) for r in 1:model.R ]
    end

    ## output files
    report_file = open(report_filename, "a+")
    write(report_file, "Commencing MCMC at $(Dates.now()) at iteration $(model.iter) for $(n_keep * thin) iterations.\n")
    close(report_file)

    if save
        sims = postSimsInit(monitor, n_keep, model.state)
    end

    ## sampling
    for i in 1:n_keep
        for j in 1:thin

            jmpstart = (model.iter % jmpstart_iter == 0)

            if jmpstart

                model.state.lΛ, model.state.lλ, model.state.Zζ = MetropIndep_ΛλZζ(model.S,
                    model.state.lΛ, model.state.lλ, model.state.Zζ,
                    model.prior.Λ, model.prior.λ, model.prior.Q0, model.prior.Q,
                    model.λ_indx,
                    model.TT, model.L, model.R, model.K)

                Zandζnow = ZζtoZandζ(model.state.Zζ, model.λ_indx)

                if margQ
                  model.state.lQ0 = nothing # don't update Q unless we need it
                  model.state.lQ = nothing
                else
                  model.state.lQ0, model.state.lQ = rpost_lQ_mmtd(model.S, model.TT,
                      model.prior.Q0, model.prior.Q, Zandζnow,
                      model.λ_indx, model.L, model.R, model.K)
                end

            else

              if margQ

                model.state.Zζ = rpost_Zζ_marg(model.S, model.state.Zζ,
                    model.λ_indx,
                    model.prior.Q0, model.prior.Q,
                    model.state.lΛ, model.state.lλ,
                    model.TT, model.L, model.R, model.K)

                model.state.lQ0 = nothing # don't update Q unless we need it
                model.state.lQ = nothing

              else

                Zandζnow = ZζtoZandζ(model.state.Zζ, model.λ_indx)

                model.state.lQ0, model.state.lQ = rpost_lQ_mmtd(model.S, model.TT,
                    model.prior.Q0, model.prior.Q, Zandζnow,
                    model.λ_indx, model.L, model.R, model.K)

                model.state.Zζ = rpost_Zζ_fc(model.S, model.state.Zζ,
                    model.λ_indx,
                    model.state.lΛ, model.state.lλ,
                    model.state.lQ0, model.state.lQ,
                    model.TT, model.L, model.R, model.K)

              end

              Zandζnow = ZζtoZandζ(model.state.Zζ, model.λ_indx)

              model.state.lΛ = rpost_lΛ_mmtd(model.prior.Λ, Zandζnow[:,1], model.R)

              model.state.lλ = rpost_lλ_mmtd(model.prior.λ, Zandζnow,
                    model.λ_indx.lens, model.R)

            end

            model.iter += 1
            if model.iter % report_freq == 0
                report_file = open(report_filename, "a+")
                write(report_file, "Iter $(model.iter) at $(Dates.now())\n")

                if margQ
                  model.state.lQ0, model.state.lQ = rpost_lQ_mmtd(model.S, model.TT,
                      model.prior.Q0, model.prior.Q, Zandζnow,
                      model.λ_indx, model.L, model.R, model.K) # for likelihood calculation
                end

                llik_now = llik_MMTD(model.S, model.state.lΛ, model.state.lλ,
                    model.state.lQ0, model.state.lQ, model.λ_indx)
                write(report_file, "Log-likelihood $(llik_now)\n")
                write(report_file, "Current order weights: $( round.(exp.(model.state.lΛ), digits=3) )\n\n")
                close(report_file)
            end
        end

        if save

          if margQ && model.iter % report_freq != 0
              model.state.lQ0, model.state.lQ = rpost_lQ_mmtd(model.S, model.TT,
              model.prior.Q0, model.prior.Q, ZζtoZandζ(model.state.Zζ, model.λ_indx),
              model.λ_indx, model.L, model.R, model.K)
          end

          if typeof(monitor) == Vector{Symbol}
            for field in monitor
                sims[i][field] = deepcopy(getfield(model.state, field))
            end
          end

          sims[i][:llik] = llik_MMTD(model.S, model.state.lΛ, model.state.lλ,
                model.state.lQ0, model.state.lQ, model.λ_indx)
        end

    end

    if save
        return sims
    else
        return model.iter
    end

end

"""
Bayes factors for high order Markov chains

    The report options are:
        "lag1" : all factors reported with respect to the first order lag one model.
        "max" : all factors reported with respect to the maximum.
        "recipmax" : (default) all factors reported with respect to maximum (with maximum as numerator).

### Example
```julia
L = 4
R = 2
K = 2
λ_indx = build_λ_indx(L, R)
TT = 200
## simulate a chain with order 2, lags 1 and 3 important
S = sim_mmtd(TT, 100, L, R, K, λ_indx, # non-sparse transitions
    [ 0.0, 1.0 ],
    [ [1.0, 0.0, 0.0, 0.0 ],
      [0.0, 1.0, 0.0, 0.0, 0.0, 0.0] ],
    [ reshape( [1.0, 0.0, 1.0, 0.0], K, K),
      reshape( [0.8, 0.2, 0.3, 0.7, 0.5, 0.5, 0.1, 0.9], K, K, K ) ] )[1]
S = sim_mmtd(TT, 100, L, R, K, λ_indx, # sparse transitions
    [ 0.0, 1.0 ],
    [ [1.0, 0.0, 0.0, 0.0 ],
      [0.0, 1.0, 0.0, 0.0, 0.0, 0.0] ],
    [ reshape( [1.0, 0.0, 1.0, 0.0], K, K),
      reshape( [0.9, 0.1, 0.2, 0.8, 0.15, 0.85, 0.9, 0.1], K, K, K ) ] )[1]
priQ_type = "Dir"
priQ_type = "SBM"
if priQ_type == "Dir"
    prior_Q = symmetricDirPrior_mmtd(1.0, 1.0, 1.0, L, R, K, λ_indx)[3]
elseif priQ_type == "SBM"
end

bf = bfact_MC(S, L, R, K, prior_Q)
```
"""
function bfact_MC(S::Vector{Int}, L::Int, R::Int, K::Int,
    prior_Q0::Vector{Float64},
    prior_Q::Vector{<:Array{Float64}}, report="recipmax")

      TT = length(S)
      n = TT - L

      λ_indx = build_λ_indx(L, R)
      llik = zeros(Float64, λ_indx.nZζ+1)

      Zandζ_now = ZζtoZandζ( fill(0, n) , λ_indx)
      N0_now, N_now = counttrans_mmtd(S, TT, Zandζ_now, λ_indx, L, R, K) # vector of arrays, indices in reverse time
      llik[1] = SparseProbVec.lmvbeta( prior_Q0 .+ N0_now ) - SparseProbVec.lmvbeta( prior_Q0 )

      for i in 1:λ_indx.nZζ

          Znow = deepcopy(λ_indx.Zζindx[i,1])
          Zandζ_now = ZζtoZandζ( fill(i, n) , λ_indx)

          ## calculate current log marginal likelihood
          N0_now, N_now = counttrans_mmtd(S, TT, Zandζ_now, λ_indx, L, R, K) # vector of arrays, indices in reverse time

          α0_Q = deepcopy(prior_Q[Znow])
          α0_Qmat = reshape(α0_Q, K, K^Znow)
          α1_Q = α0_Q .+ N_now[Znow]
          α1_Qmat = reshape(α1_Q, K, K^Znow)

          llikmarg = 0.0
          for kk in 1:(K^Znow)
            llikmarg += SparseProbVec.lmvbeta( α1_Qmat[:,kk] ) - SparseProbVec.lmvbeta( α0_Qmat[:,kk] )
          end

          llik[i+1] = deepcopy(llikmarg)
      end

      if report == "recipmax"
          lbfact = maximum(llik) .- llik
      elseif report == "max"
          lbfact = llik .- maximum(llik)
      elseif report == "lag0"
          lbfact = llik .- llik[1]
      end

      bfact = exp.(lbfact)

      (λ_indx, hcat(vcat([0 0], λ_indx.Zζindx), bfact, llik))
end
function bfact_MC(S::Vector{Int}, L::Int, R::Int, K::Int,
    prior_Q0::Vector{Float64},
    prior_Q::Union{Vector{<:Array{SparseDirMix}}, Vector{<:Array{SBMprior}}},
    report="recipmax")

      prior_Q_vec = [ reshape(prior_Q[r], (K^r)) for r in 1:R ]

      TT = length(S)
      n = TT - L

      λ_indx = build_λ_indx(L, R)
      llik = zeros(Float64, λ_indx.nZζ+1)

      Zandζ_now = ZζtoZandζ( fill(0, n) , λ_indx)
      N0_now, N_now = counttrans_mmtd(S, TT, Zandζ_now, λ_indx, L, R, K) # vector of arrays, indices in reverse time
      llik[1] = SparseProbVec.lmvbeta( prior_Q0 .+ N0_now ) - SparseProbVec.lmvbeta( prior_Q0 )

      for i in 1:λ_indx.nZζ

          Znow, ζnow = deepcopy(λ_indx.Zζindx[i,1:2])
          Zandζ_now = ZζtoZandζ( fill(i, n) , λ_indx)

          ## calculate current log marginal likelihood
          N0_now, N_now = counttrans_mmtd(S, TT, Zandζ_now, λ_indx, L, R, K) # vector of arrays, indices in reverse time
          N_now_mat = reshape(N_now[Znow], K, K^Znow)

          llikmarg = 0.0
          for kk in 1:(K^Znow)
            llikmarg += SparseProbVec.logMarginal( prior_Q_vec[Znow][kk],
                            N_now_mat[:,kk] )
          end

          llik[i+1] = deepcopy(llikmarg)
      end

      if report == "recipmax"
          lbfact = maximum(llik) .- llik
      elseif report == "max"
          lbfact = llik .- maximum(llik)
      elseif report == "lag0"
          lbfact = llik .- llik[1]
      end

      bfact = exp.(lbfact)

      (λ_indx, hcat(vcat([0 0], λ_indx.Zζindx), bfact, llik))
end

"""
    forecDist_MMTD(Slagrev::Vector{Int}, Λ::Vector{Float64},
      λ::Vector{Vector{Float64}}, Q0::Vector{Float64}, Q::Vector{<:Array{Float64}},
      λ_indx::λindxMMTD)

Computes the forecast distribution given lagged values and parameters.

### Example
```julia
    Λ = [0.0, 0.0, 1.0]
    λ = [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0]]
    Q0 = [0.2, 0.8]
    Q = [ reshape([0.7, 0.3, 0.25, 0.75], (2,2)),
          reshape([0.1, 0.9, 0.4, 0.6, 0.5, 0.5, 0.2, 0.8], (2,2,2)) ]
    λ_indx = build_λ_indx(3,2)
    Slagrev = [1,1,2]
    forecDist_MMTD(Slagrev, Λ, λ, Q0, Q, λ_indx)
```
"""
function forecDist_MMTD(Slagrev::Vector{Int}, Λ::Vector{Float64},
  λ::Vector{Vector{Float64}}, Q0::Vector{Float64}, Q::Vector{<:Array{Float64}},
  λ_indx::λindxMMTD)
    K = length(Q0)
    L = length(λ[1])
    R = length(Λ) - 1
    @assert length(Slagrev) == L
    w = zeros(Float64, K)
    for k in 1:K
        w[k] += Λ[1] * Q0[k]
        for ℓ in 1:λ_indx.nZζ
            Znow, ζnow = λ_indx.Zζindx[ℓ, 1:2]
            w[k] += Λ[Znow+1] * λ[Znow][ζnow] * Q[Znow][k, Slagrev[λ_indx.indxs[Znow][ζnow]]...]
        end
    end
    w
end

"""
    llik_MMTD(S::Vector{Int}, lΛ::Vector{Float64}, lλ::Vector{Vector{Float64}},
        lQ0::Vector{Float64}, lQ::Vector{<:Array{Float64}}, λ_indx::λindxMMTD)

Computes the (conditional on intial values) log-likelihood for a time series.

### Example
```julia
    L = 3
    R = 2
    K = 2
    lΛ = log.([0.0, 0.0, 1.0])
    lλ = [log.([0.0, 0.0, 1.0]), log.([0.0, 1.0, 0.0])]
    lQ0 = log.([0.2, 0.8])
    lQ = [ log.(reshape([0.7, 0.3, 0.25, 0.75], (2,2))),
          log.(reshape([0.1, 0.9, 0.4, 0.6, 0.5, 0.5, 0.2, 0.8], (2,2,2))) ]
    λ_indx = build_λ_indx(L,R)
    S = [1,1,2,1,1,2,2,1,2,2,1,2,1,2,1,1,1,1,2,1,2]
    llik_MMTD(S, lΛ, lλ, lQ0, lQ, λ_indx)
```
"""
function llik_MMTD(S::Vector{Int}, lΛ::Vector{Float64}, lλ::Vector{Vector{Float64}},
    lQ0::Vector{Float64}, lQ::Vector{<:Array{Float64}}, λ_indx::λindxMMTD)

    TT = length(S)
    L = length(lλ[1])

    n = TT - L

    lpt = Matrix{Float64}(undef, n, λ_indx.nZζ+1)

    for i in 1:n
        tt = i + L
        Slagrev_now = S[range(tt-1, step=-1, length=L)]

        lpt[i,1] = lΛ[1] + lQ0[S[tt]]
        for ℓ in 1:λ_indx.nZζ
            Znow, ζnow = λ_indx.Zζindx[ℓ, 1:2]
            lpt[i,ℓ+1] = lΛ[Znow+1] + lλ[Znow][ζnow] + lQ[Znow][S[tt], Slagrev_now[λ_indx.indxs[Znow][ζnow]]...]
        end
    end

    lptvec = [ SparseProbVec.logsumexp(lpt[i,:]) for i in 1:n ]
    llik = sum(lptvec)

    return llik
end

"""
    addDecomp_postsims!(sims::Vector{Any}, model::ModMMTD; totalweight_thresh=0.001, maxiter=100, leftoversum_tol=1.0e-9, Qfc=false)

Reconstruct and decompose transition probability tensor from MMTD posterior simulations. Adds aggregated level weights and
lag involvment to each sample.

If Qfc=true, the decomposition is based on the full conditional mean of Q0 and Q, which requires membership indicators.
"""
function addDecomp_postsims!(sims, model::ModMMTD; totalweight_thresh=0.001, maxiter=100, leftoversum_tol=1.0e-9, Qfc=false)

  if Qfc # calculate using full conditions of Q

    addDecomp_postsims_Qfc!(sims, model; totalweight_thresh=totalweight_thresh, maxiter=maxiter, leftoversum_tol=leftoversum_tol)

  else

    for ii in 1:length(sims)

      if typeof(sims[ii][:lλ]) <: Array{Any}
        sims[ii][:lλ] = [ deepcopy(sims[ii][:lλ][j]) for j = 1:length(sims[ii][:lλ]) ] # correctly convert from BSON
      end
      if typeof(sims[ii][:lQ]) <: Array{Any}
        sims[ii][:lQ] = [ deepcopy(sims[ii][:lQ][j]) for j = 1:length(sims[ii][:lQ]) ] # correctly convert from BSON
      end

      Omeg = transTensor_mmtd(model.L, model.R, model.K, model.λ_indx,
      sims[ii][:lΛ], sims[ii][:lλ], sims[ii][:lQ0], sims[ii][:lQ], login=true)
      dcp = reduce_transTens(TransTens(Omeg), totalweight_thresh=totalweight_thresh, maxiter=maxiter, leftoversum_tol=leftoversum_tol)
      sims[ii][:Dlaginvolv] = deepcopy(dcp.involv)
      sims[ii][:Dlevelweight] = deepcopy(dcp.level_lam_total)
    end

  end

  return nothing
end

function fc_mean_lQ_mmtd(S::Vector{Int}, TT::Int,
  prior_Q0::Vector{Float64}, prior_Q::Vector{<:Array{Float64}},
  Zandζ::Matrix{Int},
  λ_indx::λindxMMTD, L::Int, R::Int, K::Int)

  N0, N = counttrans_mmtd(S, TT, Zandζ, λ_indx, L, R, K) # N is a vector of arrays

  ## initialize
  α0_Q0 = deepcopy(prior_Q0)
  α0_Q = deepcopy(prior_Q) # vector of arrays
  lmQ_mats = [ Matrix{Float64}(undef, K, K^r) for r in 1:R ]
  lmQ_out = [ reshape(lmQ_mats[r], fill(K, r+1)...) for r in 1:R ]

  α1_Q0 = α0_Q0 .+ N0
  lmQ0_out = log.( α1_Q0 ./ sum(α1_Q0) )

  for r in 1:R
    ncol = K^r
    α1_Q = α0_Q[r] .+ N[r]
    α1_Q_mat = reshape(α1_Q, (K, ncol))
    for j in 1:ncol
      lmQ_mats[r][:,j] = log.(α1_Q_mat[:,j] ./ sum(α1_Q_mat[:,j]))
    end
    lmQ_out[r] = reshape(lmQ_mats[r], fill(K, r+1)...)
  end

  lmQ0_out, lmQ_out
end

function addDecomp_postsims_Qfc!(sims, model::ModMMTD; totalweight_thresh=0.001, maxiter=100, leftoversum_tol=1.0e-9)

  for ii in 1:length(sims)

    if typeof(sims[ii][:lλ]) <: Array{Any}
      sims[ii][:lλ] = [ deepcopy(sims[ii][:lλ][j]) for j = 1:length(sims[ii][:lλ]) ] # correctly convert from BSON
    end
    if typeof(sims[ii][:lQ]) <: Array{Any}
      sims[ii][:lQ] = [ deepcopy(sims[ii][:lQ][j]) for j = 1:length(sims[ii][:lQ]) ] # correctly convert from BSON
    end
    if typeof(model.prior.Q) <: Array{Any}
      model.prior.Q = [ deepcopy(model.prior.Q[j]) for j = 1:length(model.prior.Q) ] # correctly convert from BSON
    end

    # calculate full conditional mean for Q0, Q

    Zandζnow = ZζtoZandζ(sims[ii][:Zζ], model.λ_indx)

    lmQ0, lmQ = fc_mean_lQ_mmtd(model.S, model.TT,
      model.prior.Q0, model.prior.Q, Zandζnow,
      model.λ_indx, model.L, model.R, model.K)

    Omeg = transTensor_mmtd(model.L, model.R, model.K, model.λ_indx,
      sims[ii][:lΛ], sims[ii][:lλ], lmQ0, lmQ, login=true)
    dcp = reduce_transTens(TransTens(Omeg), totalweight_thresh=totalweight_thresh, maxiter=maxiter, leftoversum_tol=leftoversum_tol)
    sims[ii][:Dlaginvolv] = deepcopy(dcp.involv)
    sims[ii][:Dlevelweight] = deepcopy(dcp.level_lam_total)
  end

  return nothing

end

max_stars_repo_name: mheiner/MTD

max_issues_repo_name: mheiner/MTD

max_forks_repo_name: mheiner/MTD


text:
@testset "Test utility functions" begin
    concrete_types = IS.get_all_concrete_subtypes(IS.InfrastructureSystemsType)
    @test length([x for x in concrete_types if isconcretetype(x)]) == length(concrete_types)
end

@testset "Test strip_module_name" begin
    @test IS.strip_module_name("PowerSystems.HydroDispatch") == "HydroDispatch"

    @test IS.strip_module_name(
        "InfrastructureSystems.Deterministic{PowerSystems.HydroDispatch}") ==
        "Deterministic{PowerSystems.HydroDispatch}"
end

@testset "Test exported names" begin
    @test IS.validate_exported_names(IS)
end

ext: jl


text:
import MPI
import MLStyle: @match

function hello(config::Config)
    worldrank = MPI.Comm_rank(MPI.COMM_WORLD)
    worldsize = MPI.Comm_size(MPI.COMM_WORLD)
    # TODO: Should this be a command-line argument for the hello sub-command?
    backend = config[:datasets][1][:io_backend]
    basename = "hello"
    data = "Hello world, I am rank $worldrank of $worldsize\n"

    @match backend begin
        "HDF5" => hello_hdf5(data)
        "ADIOS2" => hello_adios2(data)
        "Julia" => write(basename * "_$worldrank.dat", data)
        _ => @error "Unsupported I/O strategy"
    end

    MPI.Barrier(MPI.COMM_WORLD)
end

hello(filename::AbstractString) = hello(read_config(filename))

hello(::Nothing) = hello(default_config())


text:
module DifferenceEquations

using ChainRulesCore
using Distributions
using LinearAlgebra
using CommonSolve
using UnPack
using PDMats
using RecursiveArrayTools
using SciMLBase: @add_kwonly, NullParameters, promote_tspan, AbstractRODESolution
using DiffEqBase
using DiffEqBase: __solve

include("matrix_vector_of_vectors.jl")
include("utilities.jl")
include("problems/state_space_problems.jl")
include("solutions/state_space_solutions.jl")
include("solve.jl")
include("algorithms/linear.jl")
include("algorithms/quadratic.jl")

# Exports
export MatrixVectorOfArray

export AbstractStateSpaceProblem, LinearStateSpaceProblem, QuadraticStateSpaceProblem
export StateSpaceSolution, DirectIteration, KalmanFilter

export solve

end # module


text:
#-------------------------------------------------------------------- # DNSS.jl
# Soham 03-2022
# Dump the data using NPZ files to plot using Python
#--------------------------------------------------------------------

using NPZ

function save(fname::String, w::Field{S}) where {S}
    u   = Field(w.space, u->u)
    wlm = basistransform(w)
    npzwrite("output/data/$fname", Dict("u" => u.value, "w" => w.value, "wlm" => wlm.value))
end

function save(fname::String, w::Field{ProductSpace{S1, S2}}) where {S1, S2}
    u   = Field(w.space.S1, u->u)
    v   = Field(w.space.S2, v->v)
    wlm = basistransform(w)
    npzwrite("$fname", Dict("u" => u.value, "v" => v.value, "w" => w.value, "wlm" => wlm.value))
end

function save(fname::String, AoF::T) where {T<:Array{Field, 2}}
    for index in CartesianIndices(AoF)
        if eltype(AoF[index].value) <: Number
            (a, b) = index.I
            save(join(["$fname", "$(10 + a)", "$(10 + b)"], "_"), AoF[index])
        else
            @show index
            "We found a element that does not contain a number?"
        end
    end
end

function save(fname::String, n::Vector, e::Vector)
    npzwrite("$fname", Dict("n" => n, "e" => e))
end


text:
"""
    compare(s1, s2, dist)

return a similarity score between 0 and 1 for the strings `s1` and
`s2` based on the distance `dist`.

### Examples
```julia-repl
julia> compare("martha", "marhta", Levenshtein())
0.6666666666666667
```
"""
compare(s1, s2, dist::StringDistance; min_score = 0.0) = 1 - normalize(dist)(s1, s2, 1 - min_score)

"""
    findnearest(s, itr, dist::StringDistance; min_score = 0.0) -> (x, index)

`findnearest` returns the value and index of the element of `itr` that has the
highest similarity score with `s` according to the distance `dist`.
It returns `(nothing, nothing)` if none of the elements has a similarity score
higher or equal to `min_score` (default to 0.0).

It is particularly optimized for [`Levenshtein`](@ref) and [`DamerauLevenshtein`](@ref) distances
(as well as their modifications via [`Partial`](@ref), [`TokenSort`](@ref), [`TokenSet`](@ref), or [`TokenMax`](@ref)).

### Examples
```julia-repl
julia> using StringDistances
julia> s = "Newark"
julia> iter = ["New York", "Princeton", "San Francisco"]
julia> findnearest(s, iter, Levenshtein())
("NewYork", 1)
julia> findnearest(s, iter, Levenshtein(); min_score = 0.9)
(nothing, nothing)
```
"""
function findnearest(s, itr, dist::StringDistance; min_score = 0.0)
    min_score_atomic = Threads.Atomic{Float64}(min_score)
    scores = [0.0 for _ in 1:Threads.nthreads()]
    is = [0 for _ in 1:Threads.nthreads()]
    s = _helper(s, dist)
    # need collect since @threads requires a length method
    Threads.@threads for i in collect(eachindex(itr))
        score = compare(s, _helper(itr[i], dist), dist; min_score = min_score_atomic[])
        score_old = Threads.atomic_max!(min_score_atomic, score)
        if score >= score_old
            scores[Threads.threadid()] = score
            is[Threads.threadid()] = i
        end
    end
    imax = is[argmax(scores)]
    imax == 0 ? (nothing, nothing) : (itr[imax], imax)
end

function _helper(s, dist::QGramDistance)
    s !== missing ? QGramSortedVector(s, dist.q) : s
end
_helper(s, dist::StringDistance) = s

function Base.findmax(s, itr, dist::StringDistance; min_score = 0.0)
    @warn "findmax(s, itr, dist; min_score) is deprecated. Use findnearest(s, itr, dist; min_score)"
    findnearest(s, itr, dist; min_score = min_score)
end
"""
    findall(s, itr , dist::StringDistance; min_score = 0.8)

`findall` returns the vector of indices for elements of `itr` that have a
similarity score higher or equal than `min_score` according to the distance `dist`.
If there are no such elements, return an empty array.

It is particularly optimized for [`Levenshtein`](@ref) and [`DamerauLevenshtein`](@ref) distances
(as well as their modifications via `Partial`, `TokenSort`, `TokenSet`, or `TokenMax`).

### Examples
```julia-repl
julia> using StringDistances
julia> s = "Newark"
julia> iter = ["Newwark", "Princeton", "San Francisco"]
julia> findall(s, iter, Levenshtein())
1-element Array{Int64,1}:
 1
julia> findall(s, iter, Levenshtein(); min_score = 0.9)
0-element Array{Int64,1}
```
"""
function Base.findall(s, itr, dist::StringDistance; min_score = 0.8)
    out = [Int[] for _ in 1:Threads.nthreads()]
    s = _helper(s, dist)
    # need collect since @threads requires a length method
    Threads.@threads for i in collect(eachindex(itr))
        score = compare(s, _helper(itr[i], dist), dist; min_score = min_score)
        if score >= min_score
            push!(out[Threads.threadid()], i)
        end
    end
    vcat(out...)
end


text:
# See https://github.com/tlienart/Franklin.jl/issues/326

@testset "Issue 326" begin
    h1 = "
Blah
"
    h1e = Markdown.htmlesc(h1)
    @test F.is_html_escaped(h1e)
    @test F.html_unescape(h1e) == h1
end


text:
using Test
using MultivariateBases
using DynamicPolynomials
@polyvar x y

@testset "Polynomials" begin
    basis = FixedPolynomialBasis([1 - x^2, x^2 + 2])
    @test polynomialtype(basis, Int) == polynomialtype(x, Int)
    @test polynomial(one, basis) == 3
end
@testset "Monomial" begin
    basis = FixedPolynomialBasis([x, x^2])
    @test polynomialtype(basis, Int) == polynomialtype(x, Int)
    @test polynomial(i -> i^2, basis) == 4x^2 + x
end
@testset "Terms" begin
    basis = FixedPolynomialBasis([1, x, x^2])
    @test polynomialtype(basis, Int) == polynomialtype(x, Int)
    @test polynomial(i -> (-1)^i, basis) == -x^2 + x - 1
end
@testset "Linear" begin
    basis = FixedPolynomialBasis([x, y])
    @test polynomialtype(basis, Int) == polynomialtype(x, Int)
    @test polynomial(identity, basis) == x + 2y
end
@testset "One variable" begin
    basis = FixedPolynomialBasis([x])
    @test polynomialtype(basis, Int) == polynomialtype(x, Int)
    @test polynomial(i -> 5, basis) == 5x
    @test typeof(polynomial(i -> 5, basis)) == polynomialtype(basis, Int)
    @test typeof(polynomial(ones(Int, 1, 1), basis, Int)) <: AbstractPolynomial{Int}
    @test typeof(polynomial(ones(Int, 1, 1), basis, Float64)) <: AbstractPolynomial{Float64}
end
@testset "Complex" begin
    basis = FixedPolynomialBasis([(1 + 2im) * x])
    @test 5x^2 == polynomial(ones(Int, 1, 1), basis, Complex{Int})
    @test 5x^2 == polynomial(ones(Int, 1, 1), basis, Int)
    # TODO not inferred on Julia v1.0
    #@test 5x^2 == @inferred polynomial(ones(Int, 1, 1), basis, Complex{Int})
    #@test 5x^2 == @inferred polynomial(ones(Int, 1, 1), basis, Int)
end

max_stars_repo_name: JuliaAlgebra/PolynomialBases

max_issues_repo_name: JuliaAlgebra/PolynomialBases

max_forks_repo_name: JuliaAlgebra/PolynomialBases


text:
using PyPlot: ColorMap, register_cmap, imshow, imsave

"""Call this function before loading first call to imshow"""
function register_ims_cmap()
    r = [(0.0,  1.0, 1.0),
         (0.05, 0.0, 0.0),
         (0.1,  1.0, 1.0),
         (0.2,  1.0, 1.0),
         (0.6,  1.0, 1.0),
         (1.0,  1.0, 1.0)]

    g = [(0.0,  1.0, 1.0),
         (0.05, 0.0, 0.0),
         (0.1,  0.0, 0.0),
         (0.2,  0.0, 0.0),
         (0.6,  1.0, 1.0),
         (1.0,  1.0, 1.0)]

    b = [(0.0,  1.0, 1.0),
         (0.05, 1.0, 1.0),
         (0.1,  1.0, 1.0),
         (0.2,  0.0, 0.0),
         (0.6,  0.0, 0.0),
         (1.0,  0.0, 0.0)]

    cm = ColorMap("ims", r, g, b)
    register_cmap("ims", cm)
end

"""Show a measurement matrix using an appropriate colormap"""
imsshow{T <: Number}(s::Matrix{T}; vmin=0, vmax=100) =
    imshow(s; cmap="ims", vmin=vmin, vmax=vmax, interpolation="none", origin="lower")
export imsshow

"""Like `imshow` but save an image using the `imsave` method of `PyPlot`"""

lsimssave{T <: Number}(s::Matrix{T}, f::AbstractString; vmin=0, vmax=100) =
    imsave(f, s; cmap="ims", vmin=vmin, vmax=vmax, origin="lower")


text:
import ViscousFlow: Systems

using LinearAlgebra

@testset "Navier-Stokes" begin

  @testset "Lamb-Oseen vortex" begin

    woseen(x::Tuple{Real,Real},t;Re=1.0,x0::Tuple{Real,Real}=(0,0),t0=1) =
                          exp(-((x[1]-x0[1])^2+(x[2]-x0[2])^2)/(4(t+t0)/Re))/(1+t/t0)

    #wexact(x, y, t, Re, σ, U, Δx) = Δx*exp(-((x-U*t)^2+y^2)/(σ^2+4t/Re))/(1+4t/(Re*σ^2))

    #nx = 260; ny = 130
    Re = 200 + 50rand()
    U = 1.0 + 0.2randn()
    U∞ = (U,0.0)
    #Lx = 2.0
    Δx = 0.015  #Lx/(ny-1)
    Δt = min(0.5*Δx,0.5*Δx^2*Re)

    #sys = Systems.NavierStokes((nx,ny),Re,Δx,Δt,U∞ = U∞)
    sys = Systems.NavierStokes(Re,Δx,(0.0,3.0),(0.0,2.0),Δt,U∞ = U∞)

    w₀ = Nodes(Dual,size(sys))
    xg,yg = coordinates(w₀,dx=Δx)
    x0 = (1,1)
    t0 = 1
    wexact(t) = [Δx*woseen((x,y),t;Re=Re,x0=x0.+U∞.*t,t0=t0) for x in xg, y in yg]

    ifrk = IFRK(w₀,sys.Δt,
                (t,w) -> plan_intfact(t,w,sys),
                (w,t) -> r₁(w,t,sys) ,rk=TimeMarching.RK31)

    t = 0.0
    w₀ .= wexact(t)

    w = deepcopy(w₀)
    tf = 1.0
    T = 0:Δt:tf

    for ti in T
      t, w = ifrk(t,w)
    end

    @test norm(w-wexact(t),Inf) < 1e-1
    @test abs(sum(w)-sum(wexact(t))) < 1e-3

  end

  @testset "Flow past cylinder" begin

    Re = 200
    U = 1.0
    U∞ = (U,0.0)

    n = 100
    body = Bodies.Circle(0.5,n)

    xlim = (-1.0,3.0)
    ylim = (-1.0,1.0)

    X = VectorData(body.x,body.y)

    Δx = 0.02
    Δt = min(0.5*Δx,0.5*Δx^2*Re)

    sys = Systems.NavierStokes(Re,Δx,xlim,ylim,Δt,U∞ = U∞, X̃ = X, isstore = true)

    @test size(sys,1) == 208
    @test size(sys,2) == 104
    @test size(sys) == (208,104)

    @test Systems.origin(sys) == (54,52)

    wf = Systems.PointForce(Nodes(Dual,size(sys)),(1.5,0.0),10.0,1.5,1.0,sys)
    @test sum(wf(1.5)) ≈ 10

    qf = Systems.PointForce(Edges(Primal,size(sys)),(1.5,0.0),(10.0,-10.0),1.5,1.0,sys)
    @test sum(qf(1.5).u) ≈ 10
    @test sum(qf(1.5).v) ≈ -10

  end

  @testset "Histories" begin

    a = 1.0
    b = 0.0
    c = 5.0

    h = History(Nodes(Dual,(5,5)))
    d = Nodes(Dual,(5,5))
    fill!(d,a)
    push!(h,deepcopy(d))
    fill!(d,b)
    push!(h,deepcopy(d))
    fill!(d,c)
    push!(h,deepcopy(d))

    @test length(h.r) == length(h)

    dh = diff(h)

    @test typeof(dh) <: History

    @test dh[2][1,1] == c-b

    # test periodic history
    hp = History(h.vec,htype=PeriodicHistory)
    @test hp[4] == hp[1]

    # test the differencing of periodic history
    dhp = diff(hp)
    @test typeof(dhp) <: History
    @test dhp[1][1,1] == b-a
    @test dhp[5] == dhp[2]

    # test staggered indexing
    hp2 = History(hp[1:2:5],htype=PeriodicHistory)
    @test hp2[2][1,1] == c

    # test circular shifting
    hpshift = circshift(hp,1)
    @test hpshift[1][1,1] == c

    # testing arithmetic
    h3 = 3*h
    @test h3[3][1,1] == 3*c

    h4 = h+h3
    @test h4[1][1,1] == 4*a

    h5 = -h
    @test h5[3][1,1] == -c

    h6 = h/2
    @test h6[1][1,1] == a/2

    # test setting ghosts on regular history
    h_pre = deepcopy(h)
    h_post = deepcopy(h)
    set_first_ghost!(h,h_pre)
    @test h[1][1,1] == c

    set_last_ghost!(h,h_post)
    @test h[end][1,1] == a

    @test length(h) == length(h.r)+2
    @test h.r.start == 2
    @test h.r.stop == 4

  end

end


text:
autocor(dv::DataVector{T}, lag::Int) where {T} = cor(dv[1:(end - lag)], dv[(1 + lag):end])
autocor(dv::DataVector{T}) where {T} = autocor(dv, 1)

"""
    gl(n::Integer, k::Integer, l::Integer = n*k) -> PooledDataArray

Generate a [`PooledDataArray`](@ref) with `n` levels and `k` replications, optionally
specifying an output length `l`. If specified, `l` must be a multiple of `n*k`.

# Examples

```jldoctest
julia> gl(2, 1)
2-element DataArrays.PooledDataArray{Int64,UInt8,1}:
 1
 2

julia> gl(2, 1, 4)
4-element DataArrays.PooledDataArray{Int64,UInt8,1}:
 1
 2
 1
 2
```
"""
function gl(n::Integer, k::Integer, l::Integer)
    nk = n * k
    d, r = divrem(l, nk)
    r == 0 || throw(ArgumentError("length out must be a multiple of n * k"))
    aa = Array(Int, l)
    for j = 0:(d - 1), i = 1:n
        aa[j * nk + (i - 1) * k + (1:k)] = i
    end
    compact(PooledDataArray(aa))
end

gl(n::Integer, k::Integer) = gl(n, k, n*k)

StatsBase.describe(X::DataVector) = StatsBase.describe(stdout, X)

function StatsBase.describe(io::IO, X::AbstractDataVector{T}) where T<:Real
    nacount = sum(ismissing, X)
    pna = 100nacount/length(X)
    if pna != 100 # describe will fail if skipmissing returns an empty vector
        describe(io, collect(skipmissing(X)))
    else
        println(io, "Summary Stats:")
        println(io, "Type:           $(T)")
    end
    println(io, "Number Missing: $(nacount)")
    @printf(io, "%% Missing:      %.6f\n", pna)
    return
end

function StatsBase.describe(io::IO, X::AbstractDataVector)
    nacount = sum(ismissing, X)
    pna = 100nacount/length(X)
    println(io, "Summary Stats:")
    println(io, "Length:         $(length(X))")
    println(io, "Type:           $(Missings.T(eltype(X)))")
    println(io, "Number Unique:  $(length(unique(X)))")
    println(io, "Number Missing: $(nacount)")
    @printf(io, "%% Missing:      %.6f\n", pna)
    return
end


text:
using Distributed
# */ --------------------------------------------------------------------------------------------- /* #

function get_case()::Int64

    println("")
    println(" case 0 -> 0 uncertain parameters")
    println(" case 1 -> 1 uncertain parameters")
    println(" case 2 -> 2 uncertain parameters")
    println(" case 3 -> 3 uncertain parameters")
    println(" select case study:")
    cs = ""
    cs = readline()
    while (cs ∉ ["0","1","2","3"])
        println(" error: select 0, 1, 2, or 3")
        println(" select case:")
        cs = readline()
    end

    return parse(Int64,cs)
end

function get_algm()::Int64

    println("")
    println(" algorithm 1 -> Stand_Bend (Standard Benders)")
    println(" algorithm 2 -> Adapt_Bend (Benders with Adaptive Oracles)")
    println(" select Benders-type algorithm:")
    al = ""
    al = readline()
    while (al ∉ ["1","2"])
        println(" error: select 1 or 2")
        println(" select algorithm:")
        al = readline()
    end

    return parse(Int64,al)
end

function get_params()::Tuple{Int64,Int64}

    c =  get_case()
    a =  get_algm()

    return c,a
end

# */ --------------------------------------------------------------------------------------------- /* #

function get_wrkr(n::String,d::Dict,case::Int64)::Int64

    d["hostip"] == "" ? hip = "local_machine" : hip = d["hostip"]
    case == 0 ? (workers()[1] == 1 ? nw = 2   - 1 : nw = 2   - 1 - nworkers()) : nothing
    case == 1 ? (workers()[1] == 1 ? nw = 12  - 1 : nw = 12  - 1 - nworkers()) : nothing
    case == 2 ? (workers()[1] == 1 ? nw = 90  - 1 : nw = 90  - 1 - nworkers()) : nothing
    case == 3 ? (workers()[1] == 1 ? nw = 756 - 1 : nw = 756 - 1 - nworkers()) : nothing
    println("")
    println(" select number of workers on $(n) ($(hip))")
    println(" select 0 <= w <= $(nw) :")
    w = ""
    w = readline()
    while (parse(Int64,w) < 0 || parse(Int64,w) > nw )
        println(" error ")
        println(" select 0 <= w <= $(nw) :")
        w = readline()
    end

    return parse(Int64,w)
end

# */ --------------------------------------------------------------------------------------------- /* #

# */ --------------------------------------------------------------------------------------------- /* #


text:
module GraphSignals

using LinearAlgebra: issymmetric, diag, diagm, Transpose

using CUDA: CuArray, cu
using FillArrays
using GraphLaplacians
using LightGraphs
using LightGraphs: AbstractSimpleGraph, outneighbors
using SimpleWeightedGraphs: AbstractSimpleWeightedGraph, outneighbors
using Zygote

import GraphLaplacians: degrees, degree_matrix, inv_sqrt_degree_matrix, laplacian_matrix,
    normalized_laplacian, scaled_laplacian
import LightGraphs: nv, ne, adjacency_matrix, is_directed

export
    # featuredgraph
    AbstractFeaturedGraph,
    NullGraph,
    FeaturedGraph,
    graph,
    node_feature,
    edge_feature,
    global_feature,
    mask,
    has_graph,
    has_node_feature,
    has_edge_feature,
    has_global_feature,

    # graph
    nv,
    ne,
    adjacency_list,
    is_directed,
    fetch_graph,

    # linalg
    adjacency_matrix,
    degrees,
    degree_matrix,
    inv_sqrt_degree_matrix,
    laplacian_matrix,
    normalized_laplacian,
    scaled_laplacian,
    laplacian_matrix!,
    normalized_laplacian!,
    scaled_laplacian!

include("cuda.jl")
include("featuredgraph.jl")
include("graph.jl")
include("linalg.jl")
include("simplegraph.jl")
include("weightedgraph.jl")
include("utils.jl")

end


text:
"""
    KernelFormulation

Abstract supertype for functors that calculate the probability density based on distance.

Concrete implementations must define functor methods with the form:

```julia
(k::SomeKernel)(distance) = # do something with `distance` and `k`"
```

Using an anonymous function would not allow rebuildable model parameters.
"""
abstract type KernelFormulation end

"""
    ExponentialKernel <: KernelFormulation

    ExponentialKernel(λ)

Probability density function of distance ``d``.

```math
y = e^{-d/λ}
```

where λ is a shape parameter.
"""
Base.@kwdef struct ExponentialKernel{P} <: KernelFormulation
    λ::P = Param(1.0, bounds=(0.0, 2.0))
end
(f::ExponentialKernel)(d) = exp(-d / f.λ)

"""
    GeometricKernel <: KernelFormulation

    GeometricKernel(α)

Probability density function of distance ``d``.

The Geometric kernel has a power-law decrease.

```math
y = (1+d)^α (α+1)(α+2) / (2 π)
```

where α is a shape parameter.
"""
Base.@kwdef struct GeometricKernel{P} <: KernelFormulation
    α::P = Param(1.0, bounds=(-1000.0, 1000.0))
end
(f::GeometricKernel)(d) = (1 + d)^f.α * ((f.α + 1)*(f.α + 2)) / (2 * π)

"""
    GaussianKernel <: KernelFormulation

    GaussianKernel(α)

Probability density function of distance ``d``.

```math
y = 1/ (π α^2) e^{-d^2/α^2}
```

where α is a positive parameter.
"""
Base.@kwdef struct GaussianKernel{P} <: KernelFormulation
    α::P = Param(1.0, bounds=(0.0, 1000.0))
end
(f::GaussianKernel)(d) = 1 / (π * f.α^2) * exp(- d^2 / f.α^2)

"""
    WeibullKernel <: KernelFormulation

    WeibullKernel(α,β)

Probability density function of distance ``d``.

```math
y =β /(2 π α^2) d^{β-2} e^{ -d^β/α^β}
```

where α and β are positive parameters.
"""
Base.@kwdef struct WeibullKernel{A,B} <: KernelFormulation
    α::A = Param(1.0, bounds=(0.0, 1000.0))
    β::B = Param(2.0, bounds=(0.0, 1000.0))
end
(f::WeibullKernel)(d) = f.β / (2 * π * f.α^2) * d^(f.β - 2) * exp(- d^f.β / f.α^f.β)


text:
struct Figure9
    error::Matrix{Float64}
    Ω::Vector{Float64}
    dates::Vector{Float64}
    function Figure9()
        @load "dat/anonomous_data.jld"
        wavetime = wavetime ./24 ./3600
        errorgram, Ω, spectime = PaperAnalysis.compute_errorgram(waveseries,wavetime)
        new(errorgram, Ω, spectime)
    end
end

@recipe function f(f::Figure9)
    seriestype := :heatmap
    seriescolor := :balance
    clims := (-12,12)
    xticks := (1:5,map(i->"day $i",1:5))
    f.dates, f.Ω, f.error
end


text:
using NLsolve

"""
    saha_ion_weights(T, nₑ, atom, ionization_energies, partition_functions)

Returns `(wII, wIII)`, where `wII` is the ratio of singly ionized to neutral atoms of a given
element, and `wIII` is the ration of doubly ionized to neutral atoms.

arguments:
- temperature `T` [K]
- electron number density `nₑ` [cm^-3]
- atom, the atomic number of the element
- `ionization_energies` is a collection indexed by integers (e.g. a `Vector`) mappping elements'
   atomic numbers to their first three ionization energies
- `partition_funcs` is a `Dict` mapping species to their partition functions
"""
function saha_ion_weights(T, nₑ, atom, ionization_energies, partition_funcs::Dict)
    χI, χII, χIII = ionization_energies[atom]
    atom = Formula(atom)
    UI = partition_funcs[Species(atom, 0)](T)
    UII = partition_funcs[Species(atom, 1)](T)

    k = kboltz_eV
    transU = translational_U(electron_mass_cgs, T)

    wII =  2.0/nₑ * (UII/UI) * transU * exp(-χI/(k*T))
    wIII = if atom == Formula(1) # hydrogen
        0.0
    else
        UIII = partition_funcs[Species(atom, 2)](T)
        (wII*2.0/nₑ * (UIII/UII) * transU * exp(-χII/(k*T)))
    end
    wII, wIII
end

"""
    translational_U(m, T)

The (inverse) contribution to the partition function from the free movement of a particle.
Used in the Saha equation.

arguments
- `m` is the particle mass
- `T` is the temperature in K
"""
function translational_U(m, T)
     k = kboltz_cgs
     h = hplanck_cgs
     (2π*m*k*T/h^2)^1.5
end

"""
    molecular_equilibrium_equations(absolute_abundances, ironization_energies, partiation_fns, equilibrium_constants)

Returns a NamedTuple representing the system of equations specifying molecular equilibrium.

arguments:
- A Dict of `absolute_abundances`, N_X/N_total
- a Dict of ionization energies, `ionization_energies`.  The keys of act as a list of all atoms.
- a Dict of partition functions, `partition_fns`
- a Dict of log molecular equilibrium constants, `equilibrium_constants`, in partial pressure form.
  The keys of `equilibrium_constants` act as a list of all molecules.

The system of equations is specified with the number densities of the neutral atoms as free
parameters.  Each equation specifies the conservation of a particular species, e.g. (simplified)

    n(O) = n(CO) + n(OH) + n(O I) + n(O II) + n(O III).

In this equation:
- `n(O)`, the number density of oxygen atoms in any form comes `absolute_abundances` and the total
number density (supplied later)
- `n(O I)` is a free parameter.  The numerical solver is varying this to satisfy the system of
  equations.
- `n(O II)`, and `n(O III)` come from the Saha (ionization) equation given `n(O I)`
- `n(CO)` and `n(OH)` come from the molecular equilibrium constants K, which are precomputed
  over a range of temperatures.

Equilibrium constants are defined in terms of partial pressures, so e.g.

    K(OH)  ==  (p(O) p(H)) / p(OH)  ==  (n(O) n(H)) / n(OH)) kT

This could also be computed via a Saha equation involving the molecular partition function and
dissolution energy.
"""
function molecular_equilibrium_equations(absolute_abundances, ionization_energies, partition_fns,
                                         equilibrium_constants)
    atoms = 0x01:Natoms
    molecules = keys(equilibrium_constants)

    #the residuals of the molecular equilibrium equations parametrized by T, electron number density
    #and number densities of each element [cm^-3]
    function system(T, nₜ, nₑ)
        atom_number_densities = nₜ .* absolute_abundances
        ion_factors = map(atoms) do elem
            wII, wIII = saha_ion_weights(T, nₑ, elem, ionization_energies, partition_fns)
            (1 + wII + wIII)
        end
        #`residuals!` puts the residuals the system of molecular equilibrium equations in `F`
        #`x` is a vector containing the number density of the neutral species of each element
        function residuals!(F, x)
            #LHS: total number of atoms, RHS: first through third ionization states
            F .= atom_number_densities .- ion_factors .* x
            for m in molecules
                el1, el2 = get_atoms(m.formula)
                nₘ = x[el1] * x[el2] * kboltz_cgs * T / 10^equilibrium_constants[m](T)
                #RHS: subtract atoms which are part of mollecules
                F[el1] -= nₘ
                F[el2] -= nₘ
            end
        end
    end

    #passing atoms and molecules might seem a little weird architecturally, but it's partly in
    #anticipation of potentially culling the list of species considered
    (atoms=atoms, molecules=molecules, equations=system, absolute_abundances=absolute_abundances,
     ionization_energies=ionization_energies, partition_fns=partition_fns,
     equilibrium_constants=equilibrium_constants)
end

"""
    molecular_equilibrium(MEQS, T, nₜ, nₑ; x0)

Iteratively solve for the number density of each species. Returns a `Dict` mapping species to number
densities.

arguments:
- the system of molecular equilibrium equations `MEQs` (the thing returned by
`molecular_equilibrium_equations`)
- the temperature `T`,
- the number density of non-electron particles `nₜ`
- the electron number density `nₑ`
- optionally, `x0`, a starting point for the solver
"""
function molecular_equilibrium(MEQs, T, nₜ, nₑ; x0=nothing)
    if x0 == nothing
        #compute good first guess by neglecting molecules
        x0 = map(MEQs.atoms) do atom
            wII, wIII =  saha_ion_weights(T, nₑ, atom, MEQs.ionization_energies,
                                                  MEQs.partition_fns)
            nₜ*MEQs.absolute_abundances[atom] / (1 + wII + wIII)
        end
    end

    #numerically solve for equlibrium.
    sol = nlsolve(MEQs.equations(T, nₜ, nₑ), x0; iterations=200, store_trace=true, ftol=nₜ*1e-12,
                  autodiff=:forward)
    if !sol.f_converged
        error("Molecular equlibrium unconverged", sol, "\n", sol.trace)
    end

    #start with the neutral atomic species
    number_densities = Dict(Species.(Formula.(MEQs.atoms), 0) .=> sol.zero)
    #now the ionized atomic species
    for a in MEQs.atoms
        wII, wIII = saha_ion_weights(T, nₑ, a, MEQs.ionization_energies, MEQs.partition_fns)
        number_densities[Species(Formula(a), 1)] = wII  * number_densities[Species(Formula(a), 0)]
        number_densities[Species(Formula(a), 2)] = wIII * number_densities[Species(Formula(a), 0)]
    end
    #now the molecules
    for m in MEQs.molecules
        el1, el2 = get_atoms(m.formula)
        n₁ = number_densities[Species(Formula(el1), 0)]
        n₂ = number_densities[Species(Formula(el2), 0)]
        logK = MEQs.equilibrium_constants[m]
        #add 1 to logK to conver to to cgs from mks
        number_densities[m] = n₁ * n₂ * kboltz_cgs * T / 10^logK(T)
    end

    number_densities
end


text:
@testset "IERS" begin
    @test obliquity_of_ecliptic_06(2.4578265e6, 0.30434616919175345) ≈ ERFA.obl06(2.4578265e6, 0.30434616919175345)
    @test mean_obliquity_of_ecliptic(2.4578265e6, 0.30434616919175345) ≈ ERFA.obl80(2.4578265e6, 0.30434616919175345)
    @test earth_rotation_angle(2.4578265e6, 0.30434616919175345) ≈ ERFA.era00(2.4578265e6, 0.30434616919175345)

    let (w1, x1, y1, z1) = precession_fukushima_williams06(2.4578265e6, 0.30434616919175345)
        (w2, x2, y2, z2) = ERFA.pfw06(2.4578265e6, 0.30434616919175345)
        @test w1 ≈ w2
        @test x1 ≈ x2
        @test y1 ≈ y2
        @test z1 ≈ z2
    end

    # Celestial to intermediate frame of date
    @test celestial_to_intermediate(0.2, 0.2, 0.2) ≈ ERFA.c2ixys(0.2, 0.2, 0.2)

    # Polar motion
    @test polar_motion(30, 30, 30) ≈ ERFA.pom00(30, 30, 30)
    @test polar_motion(20, 30, 50) ≈ ERFA.pom00(20, 30, 50)

    @test tio_locator(2.4578265e6, 0.30434616919175345) ≈ ERFA.sp00(2.4578265e6, 0.30434616919175345)

    @test AstroBase.EarthAttitude.mean_anomaly(luna, 1)  ≈ ERFA.fal03(1)
    @test AstroBase.EarthAttitude.mean_anomaly(sun, 1) ≈ ERFA.falp03(1)
    @test mean_longitude_minus_lan(luna, 1) ≈ ERFA.faf03(1)
    @test mean_elongation(luna, 1) ≈ ERFA.fad03(1)
    @test mean_longitude_ascending_node(luna, 1) ≈ ERFA.faom03(1)
    @test mean_longitude(mercury, 1) ≈ ERFA.fame03(1)
    @test mean_longitude(venus, 1) ≈ ERFA.fave03(1)
    @test mean_longitude(earth, 1) ≈ ERFA.fae03(1)
    @test mean_longitude(mars, 1) ≈ ERFA.fama03(1)
    @test mean_longitude(jupiter, 1) ≈ ERFA.faju03(1)
    @test mean_longitude(saturn, 1) ≈ ERFA.fasa03(1)
    @test mean_longitude(uranus, 1) ≈ ERFA.faur03(1)
    @test mean_longitude(neptune, 1) ≈ ERFA.fane03(1)
    @test general_precession_in_longitude(1) ≈ ERFA.fapa03(1)
    let (x1, y1) = xy06(2.4578265e6, 0.30440190993249416)
        (x2, y2) = ERFA.xy06(2.4578265e6, 0.30440190993249416)
        @test x1 ≈ x2
        @test y1 ≈ y2
    end

    let (x1, y1) = nutation_00b(2.4578265e6, 0.30440190993249416)
        (x2, y2) = ERFA.nut00b(2.4578265e6, 0.30440190993249416)
        @test x1 ≈ x2
        @test y1 ≈ y2
    end

    @test fukushima_williams_matrix(0.2,0.3,0.5,0.6) ≈ ERFA.fw2m(0.2,0.3,0.5,0.6)
    @test numat(0.7, 1.4, 1.3) ≈ ERFA.numat(0.7, 1.4, 1.3)

    let (ap, ao) = precession_rate_part_of_nutation(2.4578265e6, 0.30434616919175345)
        (ep, eo) = ERFA.pr00(2.4578265e6, 0.30434616919175345)
        @test ap ≈ ep
        @test ao ≈ eo
    end

    @test greenwich_mean_sidereal_time00(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851) ≈  ERFA.gmst00(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851)
    @test greenwich_mean_sidereal_time06(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851) ≈  ERFA.gmst06(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851)
    @test greenwich_mean_sidereal_time82(2.4578265e6, 0.30434616919175345) ≈ ERFA.gmst82(2.4578265e6, 0.30434616919175345)
    @test greenwich_mean_sidereal_time82(0.30434616919175345, 2.4578265e6) ≈ ERFA.gmst82(0.30434616919175345, 2.4578265e6)
    let (b1, p1, bp1) = bias_precession_matrix_00(2.4578265e6, 0.30434616919175345)
        (b2, p2, bp2) = ERFA.bp00(2.4578265e6, 0.30434616919175345)
        @test b1 ≈ b2
        @test p1 ≈ p2
        @test bp1 ≈ bp2
    end

    let a = [0.9999989440476103608 -0.1332881761240011518e-2 -0.5790767434730085097e-3;
         0.1332858254308954453e-2 0.9999991109044505944 -0.4097782710401555759e-4;
         0.5791308472168153320e-3 0.4020595661593994396e-4 0.9999998314954572365]
        @test equation_of_origins(a, 0.3) ≈ ERFA.eors(a, 0.3)
    end

    let (x1, y1) = nutation(2.4578265e6, 0.30434616919175345),
        (x2, y2) = ERFA.nut80(2.4578265e6, 0.30434616919175345)
        @test x1 ≈ x2
        @test y1 ≈ y2
    end

    @test equation_of_equinoxes_complementary_terms(2.4578265e6, 0.30434616919175345) ≈ ERFA.eect00(2.4578265e6, 0.30434616919175345)
    @test equation_of_equinoxes_00(2.4578265e6, 0.30440190993249416, 1.5, 1.7) ≈ ERFA.ee00(2.4578265e6, 0.30440190993249416, 1.5, 1.7)

    @test s00(2.4578265e6, 0.30434616919175345, 0, 0) ≈ ERFA.s00(2.4578265e6, 0.30434616919175345, 0, 0)
    @test s06(2.4578265e6, 0.30434616919175345, 0, 0) ≈ ERFA.s06(2.4578265e6, 0.30434616919175345, 0, 0)

    let (x1, y1) = nutation_00a(2.4578265e6, 0.30440190993249416)
        (x2, y2) = ERFA.nut00a(2.4578265e6, 0.30440190993249416)
        @test x1 ≈ x2
        @test y1 ≈ y2
    end

    let (x1, y1) = nutation_a06(2.4578265e6, 0.30434616919175345)
        (x2, y2) = ERFA.nut06a(2.4578265e6, 0.30434616919175345)
        @test x1 ≈ x2
        @test y1 ≈ y2
    end

    let t = [0.9999989440476103608 -0.1332881761240011518e-2 -0.5790767434730085097e-3;
            0.1332858254308954453e-2 0.9999991109044505944 -0.4097782710401555759e-4;
            0.5791308472168153320e-3 0.4020595661593994396e-4 0.9999998314954572365]
        @test greenwich_apparent_sidereal_time06(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851, t) ≈ ERFA.gst06(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851, t)
    end
    @test equation_of_equinoxes_94(2.4578265e6, 0.30434616919175345) ≈ equation_of_equinoxes_94(2.4578265e6, 0.30434616919175345)
end

@testset "Dependencies" begin
    @test nutation_matrix80(2.4578265e6, 0.30434616919175345) ≈ ERFA.nutm80(2.4578265e6, 0.30434616919175345)

    let (a1, b1, c1, d1, e1, f1) = precession_nutation00(2.4578265e6, 0.30434616919175345, 0.2, 0.2)
        (a2, b2, c2, d2, e2, f2) = ERFA.pn00(2.4578265e6, 0.30434616919175345, 0.2, 0.2)
        @test a1 ≈ a2
        @test b1 ≈ b2
        @test c1 ≈ c2
        @test d1 ≈ d2
        @test e1 ≈ e2
        @test f1 ≈ f2
    end

    let (a1, b1, c1) = precession_nutation_a00(2.4578265e6, 0.30434616919175345)
        (a2, b2, c2, d2, e2, f2, g2, h2) = ERFA.pn00a(2.4578265e6, 0.30434616919175345)
        @test a1 ≈ a2
        @test b1 ≈ b2
        @test c1[1] ≈ c2
        @test c1[2] ≈ d2
        @test c1[3] ≈ e2
        @test c1[4] ≈ f2
        @test c1[5] ≈ g2
        @test c1[6] ≈ h2
    end

    let (a1, b1, c1) = precession_nutation_b00(2.4578265e6, 0.30434616919175345)
        (a2, b2, c2, d2, e2, f2, g2, h2) = ERFA.pn00b(2.4578265e6, 0.30434616919175345)
        @test a1 ≈ a2
        @test b1 ≈ b2
        @test c1[1] ≈ c2
        @test c1[2] ≈ d2
        @test c1[3] ≈ e2
        @test c1[4] ≈ f2
        @test c1[5] ≈ g2
        @test c1[6] ≈ h2
    end

    @test precession_nutation_matrix_a00(2.4578265e6, 0.30434616919175345) ≈ ERFA.pnm00a(2.4578265e6, 0.30434616919175345)
    @test precession_nutation_matrix_b00(2.4578265e6, 0.30434616919175345) ≈ ERFA.pnm00b(2.4578265e6, 0.30434616919175345)
    @test precession_nutation_matrix_a06(2.4578265e6, 0.30434616919175345) ≈ ERFA.pnm06a(2.4578265e6, 0.30434616919175345)

    @test nutation_matrix_day_a00(2.4578265e6, 0.30434616919175345) ≈ ERFA.num00a(2.4578265e6, 0.30434616919175345)
    @test nutation_matrix_day_b00(2.4578265e6, 0.30434616919175345) ≈ ERFA.num00b(2.4578265e6, 0.30434616919175345)

    @test nutation_matrix_day(2.4578265e6, 0.30434616919175345) ≈ ERFA.num06a(2.4578265e6, 0.30434616919175345)
    @test s06a(2.4578265e6, 0.30434616919175345) ≈ ERFA.s06a(2.4578265e6, 0.30434616919175345)

    @test equation_of_equinoxes_a00(2.4578265e6, 0.30434616919175345) ≈ ERFA.ee00a(2.4578265e6, 0.30434616919175345)
    @test equation_of_equinoxes_b00(2.4578265e6, 0.30434616919175345) ≈ ERFA.ee00b(2.4578265e6, 0.30434616919175345)

    @test greenwich_mean_sidereal_time_a06(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851) ≈ ERFA.gst06a(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851)
    @test equation_of_equinoxes_a06(2.4578265e6, 0.30434616919175345) ≈ ERFA.ee06a(2.4578265e6, 0.30434616919175345)
    @test greenwich_apparent_sidereal_time94(2.4578265e6, 0.30434616919175345) ≈ ERFA.gst94(2.4578265e6, 0.30434616919175345)
    @test greenwich_apparent_sidereal_time_a00(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851) ≈ ERFA.gst00a(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851)
    @test greenwich_apparent_sidereal_time_b00(2.4579405e6, 0.0) ≈ ERFA.gst00b(2.4579405e6, 0.0)
    @test greenwich_apparent_sidereal_time_a06(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851) ≈ ERFA.gst06a(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851)
    @test celestial_to_intermediate_frame_of_date(2.4578265e6, 0.30434616919175345, 0.2, 0.3) ≈ ERFA.c2ixy(2.4578265e6, 0.30434616919175345, 0.2, 0.3)
    let t = [9.999962358680738e-1 -2.516417057665452e-3 -1.093569785342370e-3;
            2.516462370370876e-3 9.999968329010883e-1 4.006159587358310e-5;
            1.093465510215479e-3 -4.281337229063151e-5 9.999994012499173e-1]
        @test celestial_to_intermediate_matrix(2.4578265e6, 0.30434616919175345, t) ≈ ERFA.c2ibpn(2.4578265e6, 0.30434616919175345, t)
    end
    @test celestial_to_intermediate_matrix_a00(2.4578265e6, 0.30434616919175345) ≈ ERFA.c2i00a(2.4578265e6, 0.30434616919175345)
    @test celestial_to_intermediate_matrix_b00(2.4578265e6, 0.30434616919175345) ≈ ERFA.c2i00b(2.4578265e6, 0.30434616919175345)
    @test celestial_to_terrestrial_matrix(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851, 0.1, 0.3, 0.5, 0.7) ≈ ERFA.c2txy(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851, 0.1, 0.3, 0.5, 0.7)
    @test celestial_to_terrestrial_a00(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851, 0.1, 0.3) ≈ ERFA.c2t00a(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851, 0.1, 0.3)
    @test celestial_to_terrestrial_b00(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851, 0.1, 0.3) ≈ ERFA.c2t00b(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851, 0.1, 0.3)
    @test c2tpe(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851, 0.1, 0.3, 0.5, 0.7) ≈ ERFA.c2tpe(2.4579405e6, 0.0, 2.4579405e6, -0.0007966009351851851, 0.1, 0.3, 0.5, 0.7)

    @test s00a(2.4578265e6, 0.30434616919175345) ≈ ERFA.s00a(2.4578265e6, 0.30434616919175345)
    @test s00b(2.4578265e6, 0.30434616919175345) ≈ ERFA.s00b(2.4578265e6, 0.30434616919175345)

    let (x1, y1, z1) = xys00a(2.4578265e6, 0.30434616919175345)
        (x2, y2, z2) = ERFA.xys00a(2.4578265e6, 0.30434616919175345)
        @test x1 ≈ x2
        @test y1 ≈ y2
        @test z1 ≈ z2
    end
    let (x1, y1, z1) = xys00b(2.4578265e6, 0.30434616919175345)
        (x2, y2, z2) = ERFA.xys00b(2.4578265e6, 0.30434616919175345)
        @test x1 ≈ x2
        @test y1 ≈ y2
        @test z1 ≈ z2
    end
end


text:
struct QNVars{T,Ty}
    d::T # search direction
    s::T # change in x
    y::Ty # change in successive gradients
end
function QNVars(x, g)
    QNVars(copy(g), copy(x), copy(x))
end
struct QNObjVars{T,Tf,Tn,Tb,Tpg}
    x::T
    F::Tf
    ∇f::Tn
    B::Tb
    Pg::Tpg
end
function preallocate_qn_caches(mstyle, x0)
    if mstyle === InPlace()
        # Maintain gradient and state pairs in QNVars
        cache = QNVars(x0, x0)
        return cache
    else
        return nothing
    end
end

function solve(
    problem::OptimizationProblem,
    x0,
    approach::LineSearch,
    options::OptimizationOptions,
    cache = preallocate_qn_caches(mstyle(problem), x0),
)
    _solve(mstyle(problem), problem, (x0, nothing), approach, options, cache)
end
function solve(
    problem::OptimizationProblem,
    s0::Tuple,
    approach::LineSearch,
    options::OptimizationOptions,
    cache = preallocate_qn_caches(mstyle(problem), first(s0)),
)
    _solve(mstyle(problem), problem, s0, approach, options, cache)
end

function _solve(
    mstyle,
    problem::OptimizationProblem,
    s0::Tuple,
    approach::LineSearch,
    options::OptimizationOptions,
    cache,
)
    t0 = time()

    #==============
         Setup
    ==============#
    x0, B0 = s0
    T = eltype(x0)

    objvars = prepare_variables(problem, approach, x0, copy(x0), B0)
    P = initial_preconditioner(approach, x0)

    f0, ∇f0 = objvars.fz, norm(objvars.∇fz, Inf) # use user norm

    if any(initial_converged(approach, objvars, ∇f0, options))
        x, fx, ∇fx, z, fz, ∇fz, B, Pg = objvars
        return ConvergenceInfo(
            approach,
            (
                P = P,
                B = B,
                ρs = norm(x .- z),
                ρx = norm(x),
                solution = z,
                fx = fx,
                minimum = fz,
                ∇fz = ∇fz,
                f0 = f0,
                ∇f0 = ∇f0,
                iter = 0,
                time = time() - t0,
            ),
            options,
        )
    end
    qnvars = QNVars(copy(objvars.∇fz), copy(objvars.∇fz), copy(objvars.∇fz))

    #==============================
             First iteration
    ==============================#
    objvars, P, qnvars = iterate(mstyle, qnvars, objvars, P, approach, problem, options)
    iter = 1
    # Check for gradient convergence
    is_converged = converged(approach, objvars, ∇f0, options)
    print_trace(approach, options, iter, t0, objvars)
    while iter < options.maxiter && !any(is_converged)
        iter += 1
        #==============================
                     iterate
        ==============================#
        objvars, P, qnvars =
            iterate(mstyle, qnvars, objvars, P, approach, problem, options, false)

        #==============================
                check convergence
        ==============================#
        is_converged = converged(approach, objvars, ∇f0, options)
        print_trace(approach, options, iter, t0, objvars)
    end
    x, fx, ∇fx, z, fz, ∇fz, B, Pg = objvars
    return ConvergenceInfo(
        approach,
        (
            P = P,
            B = B,
            ρs = norm(x .- z),
            ρx = norm(x),
            solution = z,
            fx = fx,
            minimum = fz,
            ∇fz = ∇fz,
            f0 = f0,
            ∇f0 = ∇f0,
            iter = iter,
            time = time() - t0,
        ),
        options,
    )
end
function print_trace(approach::LineSearch, options, iter, t0, objvars)
    if false
        println(
            @sprintf(
                "iter: %d   time: %.4f   obj: %.4e   ||∇f||: %.4e    α: %.4e",
                iter,
                time() - t0,
                objvars.fz,
                norm(objvars.∇fz, Inf),
                objvars.α
            )
        )
    end
end
function iterate(
    mstyle::InPlace,
    cache,
    objvars,
    P,
    approach::LineSearch,
    problem::OptimizationProblem,
    options::OptimizationOptions,
    is_first = nothing,
)
    # split up the approach into the hessian approximation scheme and line search
    x, fx, ∇fx, z, fz, ∇fz, B, Pg = objvars

    Tf = typeof(fx)
    scheme, linesearch = modelscheme(approach), algorithm(approach)
    y, d, s = cache.y, cache.d, cache.s

    # Move nexts into currs
    fx = fz
    copyto!(x, z)
    copyto!(∇fx, ∇fz)

    # Update preconditioner
    P = update_preconditioner(scheme, x, P)
    # Update current gradient and calculate the search direction
    d = find_direction!(d, B, P, ∇fx, scheme) # solve Bd = -∇fx
    # real is needed to convert complex dots to actually being real
    φ = _lineobjective(mstyle, problem, ∇fz, z, x, d, fx, real(dot(∇fx, d)))

    # Perform line search along d
    # Also returns final step vector and update the state
    α, f_α, ls_success = find_steplength(mstyle, linesearch, φ, Tf(1))

    @. s = α * d
    z = retract(problem, z, x, s)

    # Update approximation
    fz, ∇fz, B, s, y = update_obj!(problem, s, y, ∇fx, z, ∇fz, B, scheme, is_first)
    return (x = x, fx = fx, ∇fx = ∇fx, z = z, fz = fz, ∇fz = ∇fz, B = B, Pg = Pg, α = α),
    P,
    QNVars(d, s, y)
end
function print_trace(approach::LineSearch, options, iter, t0, objvars, Δ)
    if false
        println(
            @sprintf(
                "iter: %d   time: %f   f: %.4e   ||∇f||: %.4e    Δ: %.4e",
                iter,
                time() - t0,
                objvars.fz,
                norm(objvars.∇fz, Inf),
                Δ
            )
        )
    end
end

function iterate(
    mstyle::OutOfPlace,
    cache,
    objvars,
    P,
    approach::LineSearch,
    problem::OptimizationProblem,
    options::OptimizationOptions,
    is_first = nothing,
)
    # split up the approach into the hessian approximation scheme and line search
    x, fx, ∇fx, z, fz, ∇fz, B, Pg = objvars
    Tf = typeof(fx)
    scheme, linesearch = modelscheme(approach), algorithm(approach)
    # Move nexts into currs
    fx = fz
    x = copy(z)
    ∇fx = copy(∇fz)

    # Update preconditioner
    P = update_preconditioner(scheme, x, P)
    # Update current gradient and calculate the search direction
    d = find_direction(B, P, ∇fx, scheme) # solve Bd = -∇fx
    # real is needed to convert complex dots to actually being real
    φ = _lineobjective(mstyle, problem, ∇fz, z, x, d, fx, real(dot(∇fx, d)))

    # Perform line search along d
    α, f_α, ls_success = find_steplength(mstyle, linesearch, φ, Tf(1))

    # # Calculate final step vector and update the state
    s = @. α * d
    z = retract(problem, z, x, s)

    # Update approximation
    fz, ∇fz, B, s, y = update_obj(problem, s, ∇fx, z, ∇fz, B, scheme, is_first)

    return (x = x, fx = fx, ∇fx = ∇fx, z = z, fz = fz, ∇fz = ∇fz, B = B, Pg = Pg),
    P,
    QNVars(d, s, y)
end

max_stars_repo_name: JuliaNLSolvers/NLSolvers

max_issues_repo_name: JuliaNLSolvers/NLSolvers

max_forks_repo_name: JuliaNLSolvers/NLSolvers


text:
import Distributions
import DifferentialEquations: SOSRI

abstract type AbstractEntDet end
struct Entrainment <: AbstractEntDet end
struct Detrainment <: AbstractEntDet end

stochastic_closure(::APS, ::sde_struct{NoneClosureType}, ::AbstractEntDet) = 1
function stochastic_closure(param_set::APS, ::sde_struct{LogNormalClosureType}, term::AbstractEntDet)
    u = lognormal_closure(param_set, term)
    if u < 0
        @warn("Negative stochastic parameter ⟹ negative $(typeof(term)). Setting to 0")
        u = 0
    end
    return u
end
function stochastic_closure(param_set::APS, sde_model::sde_struct{SDEClosureType}, term::AbstractEntDet)
    u = sde_closure(param_set, sde_model, term)
    if u < 0
        @warn("Negative stochastic parameter ⟹ negative $(typeof(term)). Setting to 0")
        u = 0
    end
    return u
end

#
## Lognormal closure functions
#
stoch_lognormal_var(param_set, ::Entrainment) = ICP.stoch_ε_lognormal_var(param_set)
stoch_lognormal_var(param_set, ::Detrainment) = ICP.stoch_δ_lognormal_var(param_set)

function lognormal_closure(param_set::APS, term::AbstractEntDet)
    lognormal_var = stoch_lognormal_var(param_set, term)
    return lognormal_sampler(1.0, lognormal_var)
end

function lognormal_sampler(m, var)
    μ = log(m^2 / √(m^2 + var))
    σ = √(log(1 + var / m^2))
    return rand(Distributions.LogNormal(μ, σ))
end

# Stochastic differential equation

sde_θ(param_set, ::Entrainment) = ICP.sde_ϵ_θ(param_set)
sde_σ(param_set, ::Entrainment) = ICP.sde_ϵ_σ(param_set)
sde_θ(param_set, ::Detrainment) = ICP.sde_δ_θ(param_set)
sde_σ(param_set, ::Detrainment) = ICP.sde_δ_σ(param_set)

function sde_closure(param_set::APS, sde_model::sde_struct{SDEClosureType}, term::AbstractEntDet)
    θ = sde_θ(param_set, term)
    σ = sde_σ(param_set, term)
    u0 = sde_model.u0
    dt = sde_model.dt
    u = sde(θ, σ, u0, dt)
    sde_model.u0 = u
    return u
end

"""
    Solve the Ornstein-Uhlenbeck process `du = f(u,p,t)⋅dt + g(u,p,t)⋅dW` numerically

This formulation solves the Cox-Ingersoll-Ross model, a variant of the Vasicek model
which ensures that `u` remains positive so long as `θ, μ > 0`,

    `du = θ⋅(μ - u)⋅dt + σ⋅√u⋅dW`.

ref: https://en.wikipedia.org/wiki/Cox–Ingersoll–Ross_model
"""
function sde(θ::FT, σ::FT, u0::FT, dt::FT) where {FT <: Real}
    μ = FT(1)                       # μ :: long-term mean (fixed to 1)
    u_pos(u) = max(u, FT(0))        # ensure u remains non-negative
    f(u, p, t) = θ * (μ - u)        # θ :: speed of reversion
    g(u, p, t) = σ * √(u_pos(u))    # σ :: standard deviation
    tspan = (0.0, dt)
    prob = SDEProblem(f, g, u0, tspan)
    sol = solve(prob, SOSRI())
    return u_pos(sol[end])
end


text:
## MNIST data using Non-negative matrix Factorization
using MLDatasets
train_x, train_y = MNIST.traindata()
test_x,  test_y  = MNIST.testdata()

## Create the input for NMF
m = size(train_x,1)*size(train_x,2)
n = size(train_x,3)
X = Float64.(reshape(train_x, m, n))
A = copy(X') # input

## Show a random set of images
imgs = Gray.(reshape(A[1:25,:]',28,28,25))
p = plot([heatmap(imgs[:,:,i]') for i=1:25]...,framestyle=:none,yflip=true)
##
heatmap(Gray.(reshape(sum(A;dims=1)/size(A,1),28,28)'),framestyle=:none,yflip=true,colorbar=true)
##
#using ImageMagick
using NMF
using Images
## Make a rank 5 factorization
k = 5
W, H = NMF.nndsvd(A, k; variant= :ar)
# alginst = NMF.MultUpdate{Float64}(obj=:mse, maxiter=1000, verbose=false)
alginst = NMF.ALSPGrad{Float64}(maxiter=100)
r = NMF.solve!(alginst, A, W, H)
## show the images
imgs = Gray.(reshape(r.H',28,28,k))
p = plot([heatmap(imgs[:,:,i]) for i=1:k]...,framestyle=:none,yflip=true)
savefig("nmf-mnist-$k.pdf")

## Turn this into a method
function nmf(A,k)
  W, H = NMF.nndsvd(A, k; variant= :ar)
  alginst = NMF.ALSPGrad{Float64}(maxiter=100)
  r = NMF.solve!(alginst, A, W, H)
  imgs = Gray.(reshape(r.H',28,28,k))
  plot([heatmap(imgs[:,:,i]) for i=1:k]...,framestyle=:none,yflip=true), r
  savefig("nmf-mnist-$k.pdf")
end
nmf(A,8)

# alginst = NMF.MultUpdate{Float64}(obj=:mse, maxiter=1000, verbose=false)

##
nmf(A,15)

##
nmf(A,25)

##
nmf(A,50)

##
nmf(A,100)


text:
# Unit tests for functions in data.jl

@testset "Tests: building elementary universal gate" begin
    test_angle = π/3
    pauli_Y = QCO.YGate()
    H = QCO.HGate()
    U3_1 = [(1/sqrt(2))+0.0im     0.0-(1/sqrt(2))im
             0.0+(1/sqrt(2))im   -(1/sqrt(2))+0.0im]

    test_U3_1 = QCO.U3Gate(π/2,π/2,π/2)
    @test isapprox(U3_1, test_U3_1)
    # @test ceil.(real(U3_1), digits=6) == ceil.(real(test_U3_1), digits=6)
    # @test ceil.(imag(U3_1), digits=6) == ceil.(imag(test_U3_1), digits=6)

    test_U3_2 = QCO.U3Gate(π,π/2,π/2)
    @test isapprox(test_U3_2, pauli_Y)

    test_U2_2 = QCO.U2Gate(0,π)
    @test isapprox(test_U2_2, H)

    test_U3_3 = QCO.U3Gate(test_angle, -π/2, π/2)
    @test isapprox(QCO.RXGate(test_angle), test_U3_3)

    test_U3_4 = QCO.U3Gate(test_angle, 0, 0)
    @test isapprox(QCO.RYGate(test_angle), test_U3_4)

    test_U1_1 = exp(-((test_angle)/2)im)*QCO.U1Gate(test_angle)
    @test isapprox(QCO.RZGate(test_angle), test_U1_1)
end

@testset "Tests: get_full_sized_gate" begin

    # 2-qubit gates
    params = Dict{String, Any}(
    "num_qubits" => 2,
    "maximum_depth" => 2,
    "elementary_gates" => ["T_1", "T_2", "Tdagger_1", "Tdagger_2", "S_1", "S_2", "Sdagger_1", "Sdagger_2", "SX_1", "SX_2", "SXdagger_1", "SXdagger_2", "X_1", "X_2", "Y_1", "Y_2", "Z_1", "Z_2", "CZ_1_2", "CH_1_2", "CV_1_2", "Phase_1", "Phase_2", "CSX_1_2", "DCX_1_2", "Sycamore_1_2"],
    "Phase_discretization" => [π],
    "target_gate" => QCO.IGate(2),
    )

    data = QCO.get_data(params, eliminate_identical_gates = false)
    @test length(keys(data["gates_dict"])) == 26

    # kron gates
    params = Dict{String, Any}(
    "num_qubits" => 3,
    "maximum_depth" => 2,
    "elementary_gates" => ["I_1xH_2xT_3", "Tdagger_1xS_2xSdagger_3", "SX_1xSXdagger_2xX_3", "Y_1xCNot_2_3", "CNot_2_1xZ_3", "CV_2_1xI_3", "I_1xCV_2_3", "CVdagger_2_1xI_3", "I_1xCVdagger_2_3", "CX_2_1xI_3", "I_1xCX_2_3", "CY_2_1xI_3", "I_1xCY_2_3", "CZ_2_1xI_3", "I_1xCZ_2_3", "CH_2_1xI_3", "I_1xCH_2_3", "CSX_2_1xI_3", "I_1xCSX_2_3", "Swap_2_1xI_3", "I_1xiSwap_2_3", "DCX_2_1xI_3", "I_1xSycamore_2_3"],
    "target_gate" => QCO.IGate(3),
    )

    data = QCO.get_data(params, eliminate_identical_gates = false)
    @test length(keys(data["gates_dict"])) == 23

    # 3-qubit gates
    params = Dict{String, Any}(
    "num_qubits" => 3,
    "maximum_depth" => 2,
    "elementary_gates" => ["H_3", "T_3", "Tdagger_3", "Sdagger_3", "SX_3", "SXdagger_3", "X_3", "Y_3", "Z_3", "CNot_1_2", "CNot_2_3", "CNot_2_1", "CNot_3_2", "CNot_1_3", "CNot_3_1"],
    "target_gate" => QCO.IGate(3)
    )

    # 4-qubit gates
    params = Dict{String, Any}(
        "num_qubits" => 4,
        "maximum_depth" => 2,
        "elementary_gates" => ["CU3_1_3", "CRX_3_4", "CNot_2_4", "CH_1_4"],
        "CRX_discretization" => [π],
        "CU3_θ_discretization" => [π/2],
        "CU3_ϕ_discretization" => [-π/2],
        "CU3_λ_discretization" => [π/4],
        "target_gate" => QCO.IGate(4)
        )
    data = QCO.get_data(params)
    @test length(keys(data["gates_dict"])) == 4

    #5-qubit gates
    params = Dict{String, Any}(
        "num_qubits" => 5,
        "maximum_depth" => 2,
        "elementary_gates" => ["H_1", "T_2", "Tdagger_3", "Sdagger_4", "SX_5", "CX_1_2", "CX_2_1", "CY_1_2", "CY_2_1", "CRX_2_3", "CNot_3_4", "CH_5_4", "CRY_1_3", "CRZ_2_4", "CV_3_5", "CZ_4_1", "CSX_5_1"],
        "CRX_discretization" => [π],
        "CRY_discretization" => [π],
        "CRZ_discretization" => [π],
        "target_gate" => QCO.IGate(5)
        )

    data = QCO.get_data(params)
    @test length(keys(data["gates_dict"])) == 17
end

@testset "Tests: get_input_circuit_dict" begin

    function input_circuit_1()
        # [(depth, gate)]
        return [(1, "CNot_2_1"),
                (2, "S_1"),
                (3, "H_2"),
                (4, "S_2")
                ]
    end

    params = Dict{String, Any}(
    "num_qubits" => 2,
    "maximum_depth" => 5,
    "elementary_gates" => ["S_1", "S_2", "H_1", "H_2", "CNot_1_2", "CNot_2_1", "Identity"],
    "target_gate" => QCO.MGate(),
    "input_circuit" => input_circuit_1(),
    )

    data = QCO.get_data(params)

    @test "input_circuit" in keys(data)
    @test data["input_circuit"]["1"]["depth"] == 1
    @test data["input_circuit"]["1"]["gate"] == "CNot_2_1"
    @test data["input_circuit"]["2"]["depth"] == 2
    @test data["input_circuit"]["2"]["gate"] == "S_1"
    @test data["input_circuit"]["3"]["depth"] == 3
    @test data["input_circuit"]["3"]["gate"] == "H_2"
    @test data["input_circuit"]["4"]["depth"] == 4
    @test data["input_circuit"]["4"]["gate"] == "S_2"

    function input_circuit_2()
        # [(depth, gate)]
        return [(1, "CNot_2_1"),
                (2, "S_1"),
                (3, "H_2"),
                (4, "T_1")
                ]
    end

    params["input_circuit"] = input_circuit_2()
    data = QCO.get_data(params)
    @test !("input_circuit" in keys(data))

    function input_circuit_3()
        # [(depth, gate)]
        return [(1, "CNot_2_1"),
                (2, "S_1"),
                (2, "H_2"),
                (4, "S_2")
                ]
    end

    params["input_circuit"] = input_circuit_3()
    data = QCO.get_data(params)
    @test !("input_circuit" in keys(data))

    function input_circuit_4()
        # [(depth, gate)]
        return [(1, "CNot_2_1"),
                (2, "S_1"),
                (3, "H_2"),
                (4, "S_2"),
                (5, "Identity"),
                (6, "S_1"),
                ]
    end

    params["input_circuit"] = input_circuit_4()
    data = QCO.get_data(params)
    @test !("input_circuit" in keys(data))

end


text:
global std = stdout

function lg(minlevel::Int, vals... ;
            prefix::String="", suffix::String="", doflush=true)
    global args
    if args["verbosity"] >= minlevel
        if !isempty(vals)
            print(prefix, join(vals, ""), suffix)
            doflush && flush(stdout)
        end
        return true
    else
        return false
    end
end

"""
Disables log messages
"""
function stdoff()

    global args
    global std
    std = stdout
    redirect_stdout()

end

"""
Enables log messages
"""
function stdon()

    global args
    global std
    redirect_stdout(std)

end

# functions depending on the verbosity level
lg0(vals...; kwargs... ) = lg(0, vals...; kwargs... )
lg1(vals...; kwargs... ) = lg(1, vals...; kwargs... )
lg2(vals...; kwargs... ) = lg(2, vals...; kwargs... )

ln0(vals...; kwargs... ) = lg(0, vals...; kwargs..., suffix="\n" )
ln1(vals...; kwargs... ) = lg(1, vals...; kwargs..., suffix="\n" )
ln2(vals...; kwargs... ) = lg(2, vals...; kwargs..., suffix="\n" )

max_stars_repo_name: antcrepin/flight-route-optimizer

max_issues_repo_name: antcrepin/flight-route-optimizer

max_forks_repo_name: antcrepin/flight-route-optimizer