text:
"""
Ensemble Kalman Filter (EnKF) using (sparse) tensor graphical models for state
    covariance / inverse covariance estimation

Author: Wayne Wang
Last modified: 10/06/2021
"""

function enkf(Y::AbstractMatrix{<:Real},
    inv_cov_method::Union{SG_PALM, TERALASSO, KGLASSO, GLASSO},
    dynamic_type::AbstractString,
    H::AbstractArray,
    px::Tuple,
    py::Tuple,
    N::Int,
    obs_noise::Real,
    process_noise::Real,
    add_process_noise::Bool)
    # initialization
    T = size(Y, 2)
    X = zeros((T + 1, prod(px), N)) # T × p × N
    Xbar = zeros((T + 1, prod(px))) # T × p: mean latest process each time
    X0 = randn((prod(px), N))
    X[1, :, :] .= X0
    Omega = spzeros(prod(px), prod(px))
    obs_noise_vec = zeros(size(H, 1))
    Ht_Rinv_H = (obs_noise^-1) * copy(H') * H #assuming R is diagonal w/ obs_noise*I
    X_shift = similar(X[1, :, 1])

    # evolution
    for t = 1:T
        @printf("###### Time step: %d ######\n", t)
        # estiamte state precision matrix
        state_inv_cov_est!(Omega, inv_cov_method, X, Xbar, px, t, N)

        for i = 1:N
            # forecast step
            X[t + 1, :, i] .=  kalmanfilter_dynamic_update(dynamic_type,
                                    X[t, :, i], px, add_process_noise, process_noise)
            # update step
            rand!(MvNormal(size(obs_noise_vec, 1), obs_noise), obs_noise_vec)
            X_shift .= (obs_noise^-1) * copy(H') * (view(Y, :, t) .+ obs_noise_vec .- H * view(X, t + 1, :, i))
            X[t + 1, :, i] .+= (Omega .+ Ht_Rinv_H) \ X_shift
        end
    end

    return X, Xbar, Omega
end

function enkf(Y::AbstractMatrix{<:Real},
    cov_method::KPCA,
    dynamic_type::AbstractString,
    H::AbstractArray,
    px::Tuple,
    py::Tuple,
    N::Int,
    obs_noise::Real,
    process_noise::Real,
    add_process_noise::Bool)
    # initial ensemble
    T = size(Y, 2)
    X = zeros((T + 1, prod(px), N)) # T × p × N
    Xbar = zeros((T + 1, prod(px))) # T × p: mean latext process each time
    X0 = randn((prod(px), N))
    X[1, :, :] .= X0
    obs_noise_vec = zeros(size(H, 1))
    K = zeros((size(H, 1), size(H, 1)))
    R = ScalMat(size(H, 1), obs_noise)
    Sigma = zeros((prod(px), prod(px)))
    X_shift = similar(Y[:, 1])

    # evolution
    for t = 1:T
        @printf("###### Time step: %d ######\n", t)
        # estiamte state precision matrix
        state_cov_est!(Sigma, cov_method, X, Xbar, px, t, N)

        ## Kalman gain matrix: Sigma H^T (R + H Sigma H^T)^-1
        ## form the invertible part of the Kalman gain matrix
        K .= H * Sigma * copy(H') .+ R
        # droptol!(K, sp_thres) # threshold K since it has many small values

        for i = 1:N
            # forecast step
            X[t + 1, :, i] .=  kalmanfilter_dynamic_update(dynamic_type,
                                    X[t, :, i], px, add_process_noise, process_noise)
            # update step
            rand!(MvNormal(size(obs_noise_vec, 1), obs_noise), obs_noise_vec)
            X_shift .= K \ (view(Y,:,t) .+ obs_noise_vec .- H * view(X, t + 1, :, i))
            X[t + 1, :, i] .+= Sigma * copy(H') * X_shift
        end
    end

    return X, Xbar, Sigma
end

function state_inv_cov_est!(Omega::AbstractArray, method::SG_PALM, X::AbstractArray, Xbar::AbstractArray,
    px::Tuple, t::Int, N::Int;
    K::Int = 2,
    λ::AbstractVector{<:Real} = Array{Float64, 1}(),
    regtype::String = "L1",
    a::Real = 3,
    niter::Int = 100,
    ninner::Int = 10,
    η0::Real = 0.01,
    c::Real = 0.01,
    lsrule::String = "constant",
    ϵ::Real = 1e-5)
    ## compute mode-k Gram matrices
    Xbar[t, :] .= mean(view(X, t, :, :), dims = 2)[:]
    X_kGram = [zeros(px[k], px[k]) for k = 1:K]
    Xk = [zeros(px[k], Int(prod(px) / px[k])) for k = 1:K]
    for k = 1:K
        for i = 1:N
            copy!(Xk[k], tenmat(reshape(view(X, t, :, i) - Xbar[t, :], Tuple(px)), k))
            mul!(X_kGram[k], Xk[k], copy(transpose(Xk[k])), 1.0 / N, 1.0)
        end
    end

    ## run PALM for precision matrix estimation
    Psi0 = [sparse(eye(px[k])) for k = 1:K]
    fun = (iter, Psi) -> [1, time()] # NULL func
    lambda = (length(λ) == 0) ? [sqrt(px[k] * log(prod(px)) / N) for k = 1:K] : λ
    PsiH, _ = syglasso_palm(X[t, :, :] .- Xbar[t, :], X_kGram, lambda, Psi0,
                    regtype = regtype, a = a, niter = niter, ninner = ninner,
                    η0 = η0, c = c, lsrule = lsrule, ϵ = ϵ, fun = fun)

    ## form Omega
    Omega .= kron(I(px[2]), PsiH[1]^2) .+ kron(PsiH[2]^2, I(px[1])) .+ 2 * kron(PsiH[2], PsiH[1])
    return nothing
end

function state_inv_cov_est!(Omega::AbstractArray, method::TERALASSO, X::AbstractArray, Xbar::AbstractArray,
    px::Tuple, t::Int, N::Int;
    K::Int = 2,
    niter::Int = 50,
    ninner::Int = 11,
    λ::AbstractVector{<:Real} = Array{Float64, 1}())
    ## compute mode-k Gram matrices
    Xbar[t, :] .= mean(view(X, t, :, :), dims = 2)[:]
    X_kGram = [zeros(px[k], px[k]) for k = 1:K]
    Xk = [zeros(px[k], Int(prod(px) / px[k])) for k = 1:K]
    for k = 1:K
        for i = 1:N
            copy!(Xk[k], tenmat(reshape(view(X, t, :, i) - Xbar[t, :], Tuple(px)), k))
            mul!(X_kGram[k], Xk[k], copy(transpose(Xk[k])), 1.0 / N, 1.0)
        end
    end

    ## run teralasso
    # lambda = (length(λ) == 0) ? [10 * sqrt(log(prod(px)) / (N * px[k])) for k = 1:K] : λ
    PsiH, _ = teralasso(X_kGram)

    ## form Omega
    Omega .= kroneckersum(PsiH[1], PsiH[2])
    return nothing
end

function state_inv_cov_est!(Omega::AbstractArray, method::KGLASSO, X::AbstractArray, Xbar::AbstractArray,
    px::Tuple, t::Int, N::Int;
    K::Int = 2,
    λ::AbstractVector{<:Real} = Array{Float64, 1}())
    ## convert data matrix to multi-dimensional array
    # str_to_eval = string("(", join([string("px[", i, "], ") for i = 1:K]), "N)")
    # Xt = zeros(eval(Meta.parse(str_to_eval))) #d_1 × … × d_K × N tensor
    ## K=2 only! TODO: modify to work with various K ##
    Xt = zeros(px[1], px[2], N)
    for i = 1:N
        ## K=2 only! TODO: modify to work with various K ##
        Xt[:, :, i] .= reshape(view(X, t, :, i), px)
    end
    # run kglasso algorithm
    lambda = (λ == 0) ? [5 * sqrt(px[k] * log(prod(px)) / N) for k = 1:K] : λ
    Ψ_hat_list = kglasso(Xt, λ_list = lambda)
    Omega .= kron(Ψ_hat_list...)
    return nothing
end

function state_inv_cov_est!(Omega::AbstractArray, method::GLASSO, X::AbstractArray, Xbar::AbstractArray,
    px::Tuple, t::Int, N::Int;
    λ::Real = 0,
    penalize_diag::Bool = false,
    tol::Real = 1e-5)
    ## compute the sample covariance
    Xbar[t,:] .= mean(view(X, t, :, :), dims = 2)[:]
    S = zeros(prod(px), prod(px))
    for i = 1:N
        mul!(S,
            view(X, t, :, i) - view(Xbar, t, :),
            copy(transpose(view(X, t, :, i) - view(Xbar, t, :))), 1.0 / N, 1.0)
    end

    ## run Glasso for precision matrix estimation
    lambda = (λ == 0) ? 100 * sqrt(log(prod(px)) / N) : λ
    Omega .= glasso(S, lambda; tol = tol, penalize_diag = penalize_diag)
    return nothing
end

function state_cov_est!(Sigma::AbstractArray, method::KPCA, X::AbstractArray, Xbar::AbstractArray,
    px::Tuple, t::Int, N::Int;
    r::Int = 5,
    λ::Real = 0,
    lambdaL::Real = 0,
    lambdaS::Real = 0,
    tau::Real = 0.5,
    robust_pca_method::String = "SVT",
    iter::Int = 5)
    ## compute the sample covariance
    Xbar[t, :] .= mean(view(X, t, :, :), dims = 2)[:]
    S = zeros(prod(px), prod(px))
    for i = 1:N
        mul!(S,
            view(X, t, :, i) - view(Xbar, t, :),
            copy(transpose(view(X, t, :, i) - view(Xbar, t, :))), 1.0 / N, 1.0)
    end

    lambdaL = (lambdaL == 0) ? (px[1]^2+px[2]^2+log(max(px[1],px[2],N)))/N : lambdaL
    lambdaS = (lambdaS == 0) ? 5*sqrt(log(px[1]*px[2])/N) : lambdaS

    ## run Glasso for precision matrix estimation
    Sigma .= robust_kron_pca(S, px[1], px[2], lambdaL, lambdaS, robust_pca_method;
                             tau = tau, iter = iter, r = r)
    return nothing
end


text:
export MonoidalCategory, otimes, munit, ⊗, collect, ndims,
  SymmetricMonoidalCategory, FreeSymmetricMonoidalCategory, braid,
  MonoidalCategoryWithDiagonals, CartesianCategory, FreeCartesianCategory,
  mcopy, delete, pair, proj1, proj2, Δ, ◇,
  MonoidalCategoryWithCodiagonals, CocartesianCategory, FreeCocartesianCategory,
  mmerge, create, copair, incl1, incl2, ∇, □,
  MonoidalCategoryWithBidiagonals, BiproductCategory, FreeBiproductCategory,
  CartesianClosedCategory, FreeCartesianClosedCategory, hom, ev, curry,
  CompactClosedCategory, FreeCompactClosedCategory, dual, dunit, dcounit,
  DaggerCategory, FreeDaggerCategory, dagger,
  DaggerCompactCategory, FreeDaggerCompactCategory

import Base: collect, ndims

# Monoidal category
###################

""" Doctrine of *monoidal category*

To avoid associators and unitors, we assume the monoidal category is *strict*.
By the coherence theorem there is no loss of generality, but we may add a
signature for weak monoidal categories later.
"""
@signature Category(Ob,Hom) => MonoidalCategory(Ob,Hom) begin
  otimes(A::Ob, B::Ob)::Ob
  otimes(f::Hom(A,B), g::Hom(C,D))::Hom(otimes(A,C),otimes(B,D)) <=
    (A::Ob, B::Ob, C::Ob, D::Ob)
  munit()::Ob

  # Unicode syntax
  ⊗(A::Ob, B::Ob) = otimes(A, B)
  ⊗(f::Hom, g::Hom) = otimes(f, g)
end

# Convenience constructors
otimes(xs::Vector{T}) where T = isempty(xs) ? munit(T) : foldl(otimes, xs)
otimes(x, y, z, xs...) = otimes([x, y, z, xs...])

""" Collect generators of object in monoidal category as a vector.
"""
collect(expr::ObExpr) = [ expr ]
collect(expr::ObExpr{:otimes}) = vcat(map(collect, args(expr))...)
collect(expr::ObExpr{:munit}) = roottypeof(expr)[]

# XXX: We shouldn't have to do this.
roottype(T) = T isa UnionAll ? T : T.name.wrapper
roottypeof(x) = roottype(typeof(x))

""" Number of "dimensions" of object in monoidal category.
"""
ndims(expr::ObExpr) = 1
ndims(expr::ObExpr{:otimes}) = sum(map(ndims, args(expr)))
ndims(expr::ObExpr{:munit}) = 0

function show_unicode(io::IO, expr::ObExpr{:otimes}; kw...)
  Syntax.show_unicode_infix(io, expr, "⊗"; kw...)
end
function show_unicode(io::IO, expr::HomExpr{:otimes}; kw...)
  Syntax.show_unicode_infix(io, expr, "⊗"; kw...)
end
show_unicode(io::IO, expr::ObExpr{:munit}; kw...) = print(io, "I")

function show_latex(io::IO, expr::ObExpr{:otimes}; kw...)
  Syntax.show_latex_infix(io, expr, "\\otimes"; kw...)
end
function show_latex(io::IO, expr::HomExpr{:otimes}; kw...)
  Syntax.show_latex_infix(io, expr, "\\otimes"; kw...)
end
show_latex(io::IO, expr::ObExpr{:munit}; kw...) = print(io, "I")

# Symmetric monoidal category
#############################

""" Doctrine of *symmetric monoidal category*

The signature (but not the axioms) is the same as a braided monoidal category.
"""
@signature MonoidalCategory(Ob,Hom) => SymmetricMonoidalCategory(Ob,Hom) begin
  braid(A::Ob, B::Ob)::Hom(otimes(A,B),otimes(B,A))
end

@syntax FreeSymmetricMonoidalCategory(ObExpr,HomExpr) SymmetricMonoidalCategory begin
  otimes(A::Ob, B::Ob) = associate_unit(Super.otimes(A,B), munit)
  otimes(f::Hom, g::Hom) = associate(Super.otimes(f,g))
  compose(f::Hom, g::Hom) = associate(Super.compose(f,g; strict=true))
end

function show_latex(io::IO, expr::HomExpr{:braid}; kw...)
  Syntax.show_latex_script(io, expr, "\\sigma")
end

# Cartesian category
####################

""" Doctrine of *monoidal category with diagonals*

A monoidal category with diagonals is a symmetric monoidal category equipped
with coherent collections of copying and deleting morphisms (comonoids).
Unlike in a cartesian category, the naturality axioms need not be satisfied.

References:

- Selinger, 2010, "A survey of graphical languages for monoidal categories",
  Section 6.6: "Cartesian center"
- Selinger, 1999, "Categorical structure of asynchrony"
"""
@signature SymmetricMonoidalCategory(Ob,Hom) => MonoidalCategoryWithDiagonals(Ob,Hom) begin
  mcopy(A::Ob)::Hom(A,otimes(A,A))
  delete(A::Ob)::Hom(A,munit())

  # Unicode syntax
  Δ(A::Ob) = mcopy(A)
  ◇(A::Ob) = delete(A)
end

""" Doctrine of *cartesian category*

Actually, this is a cartesian *symmetric monoidal* category but we omit these
qualifiers for brevity.
"""
@signature MonoidalCategoryWithDiagonals(Ob,Hom) => CartesianCategory(Ob,Hom) begin
  pair(f::Hom(A,B), g::Hom(A,C))::Hom(A,otimes(B,C)) <= (A::Ob, B::Ob, C::Ob)
  proj1(A::Ob, B::Ob)::Hom(otimes(A,B),A)
  proj2(A::Ob, B::Ob)::Hom(otimes(A,B),B)
end

""" Syntax for a free cartesian category.

In this syntax, the pairing and projection operations are defined using
duplication and deletion, and do not have their own syntactic elements.
Of course, this convention could be reversed.
"""
@syntax FreeCartesianCategory(ObExpr,HomExpr) CartesianCategory begin
  otimes(A::Ob, B::Ob) = associate_unit(Super.otimes(A,B), munit)
  otimes(f::Hom, g::Hom) = associate(Super.otimes(f,g))
  compose(f::Hom, g::Hom) = associate(Super.compose(f,g; strict=true))

  pair(f::Hom, g::Hom) = compose(mcopy(dom(f)), otimes(f,g))
  proj1(A::Ob, B::Ob) = otimes(id(A), delete(B))
  proj2(A::Ob, B::Ob) = otimes(delete(A), id(B))
end

function show_latex(io::IO, expr::HomExpr{:mcopy}; kw...)
  Syntax.show_latex_script(io, expr, "\\Delta")
end
function show_latex(io::IO, expr::HomExpr{:delete}; kw...)
  Syntax.show_latex_script(io, expr, "\\lozenge")
end

# Cocartesian category
######################

""" Doctrine of *monoidal category with codiagonals*

A monoidal category with codiagonals is a symmetric monoidal category equipped
with coherent collections of merging and creating morphisms (monoids).
Unlike in a cocartesian category, the naturality axioms need not be satisfied.

For references, see `MonoidalCategoryWithDiagonals`.
"""
@signature SymmetricMonoidalCategory(Ob,Hom) => MonoidalCategoryWithCodiagonals(Ob,Hom) begin
  mmerge(A::Ob)::Hom(otimes(A,A),A)
  create(A::Ob)::Hom(munit(),A)

  # Unicode syntax
  ∇(A::Ob) = mmerge(A)
  □(A::Ob) = create(A)
end

""" Doctrine of *cocartesian category*

Actually, this is a cocartesian *symmetric monoidal* category but we omit these
qualifiers for brevity.
"""
@signature MonoidalCategoryWithCodiagonals(Ob,Hom) => CocartesianCategory(Ob,Hom) begin
  copair(f::Hom(A,C), g::Hom(B,C))::Hom(otimes(A,B),C) <= (A::Ob, B::Ob, C::Ob)
  incl1(A::Ob, B::Ob)::Hom(A,otimes(A,B))
  incl2(A::Ob, B::Ob)::Hom(B,otimes(A,B))
end

""" Syntax for a free cocartesian category.

In this syntax, the copairing and inclusion operations are defined using
merging and creation, and do not have their own syntactic elements.
Of course, this convention could be reversed.
"""
@syntax FreeCocartesianCategory(ObExpr,HomExpr) CocartesianCategory begin
  otimes(A::Ob, B::Ob) = associate_unit(Super.otimes(A,B), munit)
  otimes(f::Hom, g::Hom) = associate(Super.otimes(f,g))
  compose(f::Hom, g::Hom) = associate(Super.compose(f,g; strict=true))

  copair(f::Hom, g::Hom) = compose(otimes(f,g), mmerge(codom(f)))
  incl1(A::Ob, B::Ob) = otimes(id(A), create(B))
  incl2(A::Ob, B::Ob) = otimes(create(A), id(B))
end

function show_latex(io::IO, expr::HomExpr{:mmerge}; kw...)
  Syntax.show_latex_script(io, expr, "\\nabla")
end
function show_latex(io::IO, expr::HomExpr{:create}; kw...)
  Syntax.show_latex_script(io, expr, "\\square")
end

# Biproduct category
####################

""" Doctrine of *monoidal category with bidiagonals*

The terminology is nonstandard (is there any standard terminology?) but is
intended to mean a monoidal category with coherent diagonals and codiagonals.
Unlike in a biproduct category, the naturality axioms need not be satisfied.

FIXME: This signature should extend both `MonoidalCategoryWithDiagonals` and
`MonoidalCategoryWithCodiagonals`, but we don't support multiple inheritance.
"""
@signature SymmetricMonoidalCategory(Ob,Hom) => MonoidalCategoryWithBidiagonals(Ob,Hom) begin
  mcopy(A::Ob)::Hom(A,otimes(A,A))
  mmerge(A::Ob)::Hom(otimes(A,A),A)
  delete(A::Ob)::Hom(A,munit())
  create(A::Ob)::Hom(munit(),A)

  # Unicode syntax
  ∇(A::Ob) = mmerge(A)
  Δ(A::Ob) = mcopy(A)
  ◇(A::Ob) = delete(A)
  □(A::Ob) = create(A)
end

""" Doctrine of *bicategory category*

Also known as a *semiadditive category*.

FIXME: This signature should extend `MonoidalCategoryWithBidiagonals`,
`CartesianCategory`, and `CocartesianCategory`, but we don't support multiple
inheritance.
"""
@signature MonoidalCategoryWithBidiagonals(Ob,Hom) => BiproductCategory(Ob,Hom) begin
  pair(f::Hom(A,B), g::Hom(A,C))::Hom(A,otimes(B,C)) <= (A::Ob, B::Ob, C::Ob)
  copair(f::Hom(A,C), g::Hom(B,C))::Hom(otimes(A,B),C) <= (A::Ob, B::Ob, C::Ob)
  proj1(A::Ob, B::Ob)::Hom(otimes(A,B),A)
  proj2(A::Ob, B::Ob)::Hom(otimes(A,B),B)
  incl1(A::Ob, B::Ob)::Hom(A,otimes(A,B))
  incl2(A::Ob, B::Ob)::Hom(B,otimes(A,B))
end

@syntax FreeBiproductCategory(ObExpr,HomExpr) BiproductCategory begin
  otimes(A::Ob, B::Ob) = associate_unit(Super.otimes(A,B), munit)
  otimes(f::Hom, g::Hom) = associate(Super.otimes(f,g))
  compose(f::Hom, g::Hom) = associate(Super.compose(f,g; strict=true))

  pair(f::Hom, g::Hom) = compose(mcopy(dom(f)), otimes(f,g))
  copair(f::Hom, g::Hom) = compose(otimes(f,g), mmerge(codom(f)))
  proj1(A::Ob, B::Ob) = otimes(id(A), delete(B))
  proj2(A::Ob, B::Ob) = otimes(delete(A), id(B))
  incl1(A::Ob, B::Ob) = otimes(id(A), create(B))
  incl2(A::Ob, B::Ob) = otimes(create(A), id(B))
end

# Cartesian closed category
###########################

""" Doctrine of *cartesian closed category* (aka, CCC)

A CCC is a cartesian category with internal homs (aka, exponential objects).
"""
@signature CartesianCategory(Ob,Hom) => CartesianClosedCategory(Ob,Hom) begin
  # Internal hom of A and B, an object representing Hom(A,B)
  hom(A::Ob, B::Ob)::Ob

  # Evaluation map
  ev(A::Ob, B::Ob)::Hom(otimes(hom(A,B),A),B)

  # Currying (aka, lambda abstraction)
  curry(A::Ob, B::Ob, f::Hom(otimes(A,B),C))::Hom(A,hom(B,C)) <= (C::Ob)
end

""" Syntax for a free cartesian closed category.

See also `FreeCartesianCategory`.
"""
@syntax FreeCartesianClosedCategory(ObExpr,HomExpr) CartesianClosedCategory begin
  otimes(A::Ob, B::Ob) = associate_unit(Super.otimes(A,B), munit)
  otimes(f::Hom, g::Hom) = associate(Super.otimes(f,g))
  compose(f::Hom, g::Hom) = associate(Super.compose(f,g; strict=true))

  pair(f::Hom, g::Hom) = compose(mcopy(dom(f)), otimes(f,g))
  proj1(A::Ob, B::Ob) = otimes(id(A), delete(B))
  proj2(A::Ob, B::Ob) = otimes(delete(A), id(B))
end

function show_latex(io::IO, expr::ObExpr{:hom}; kw...)
  show_latex(io, last(expr))
  print(io, "^{")
  show_latex(io, first(expr))
  print(io, "}")
end
function show_latex(io::IO, expr::HomExpr{:ev}; kw...)
  Syntax.show_latex_script(io, expr, "\\mathrm{eval}")
end
function show_latex(io::IO, expr::HomExpr{:curry}; kw...)
  print(io, "\\lambda ")
  show_latex(io, last(expr))
end

# Compact closed category
#########################

""" Doctrine of *compact closed category*
"""
@signature SymmetricMonoidalCategory(Ob,Hom) => CompactClosedCategory(Ob,Hom) begin
  # Dual A^* of object A
  dual(A::Ob)::Ob

  # Unit of duality, aka the coevaluation map
  dunit(A::Ob)::Hom(munit(), otimes(dual(A),A))

  # Counit of duality, aka the evaluation map
  dcounit(A::Ob)::Hom(otimes(A,dual(A)), munit())

  # Closed monoidal category
  hom(A::Ob, B::Ob) = otimes(B, dual(A))
  ev(A::Ob, B::Ob) = otimes(id(B), compose(braid(dual(A),A), dcounit(A)))
  curry(A::Ob, B::Ob, f::Hom) = compose(
    otimes(id(A), compose(dunit(B), braid(dual(B),B))),
    otimes(f, id(dual(B)))
  )
end

@syntax FreeCompactClosedCategory(ObExpr,HomExpr) CompactClosedCategory begin
  dual(A::Ob) = anti_involute(Super.dual(A), dual, otimes, munit)
  otimes(A::Ob, B::Ob) = associate_unit(Super.otimes(A,B), munit)
  otimes(f::Hom, g::Hom) = associate(Super.otimes(f,g))
  compose(f::Hom, g::Hom) = associate(Super.compose(f,g; strict=true))
end

function show_latex(io::IO, expr::ObExpr{:dual}; kw...)
  show_latex(io, first(expr))
  print(io, "^*")
end
function show_latex(io::IO, expr::HomExpr{:dunit}; kw...)
  Syntax.show_latex_script(io, expr, "\\eta")
end
function show_latex(io::IO, expr::HomExpr{:dcounit}; kw...)
  Syntax.show_latex_script(io, expr, "\\varepsilon")
end

# Dagger category
#################

""" Doctrine of *dagger category*
"""
@signature Category(Ob,Hom) => DaggerCategory(Ob,Hom) begin
  dagger(f::Hom(A,B))::Hom(B,A) <= (A::Ob,B::Ob)
end

@syntax FreeDaggerCategory(ObExpr,HomExpr) DaggerCategory begin
  compose(f::Hom, g::Hom) = associate(Super.compose(f,g; strict=true))
  dagger(f::Hom) = anti_involute(Super.dagger(f), dagger, compose, id)
end

""" Doctrine of *dagger compact category*

FIXME: This signature should extend both `DaggerCategory` and
`CompactClosedCategory`, but we don't support multiple inheritance yet.
"""
@signature CompactClosedCategory(Ob,Hom) => DaggerCompactCategory(Ob,Hom) begin
  dagger(f::Hom(A,B))::Hom(B,A) <= (A::Ob,B::Ob)
end

@syntax FreeDaggerCompactCategory(ObExpr,HomExpr) DaggerCompactCategory begin
  dual(A::Ob) = anti_involute(Super.dual(A), dual, otimes, munit)
  otimes(A::Ob, B::Ob) = associate_unit(Super.otimes(A,B), munit)
  otimes(f::Hom, g::Hom) = associate(Super.otimes(f,g))
  compose(f::Hom, g::Hom) = associate(Super.compose(f,g; strict=true))
  function dagger(f::Hom)
    f = anti_involute(Super.dagger(f), dagger, compose, id)
    distribute_unary(f, dagger, otimes)
  end
end

function show_latex(io::IO, expr::HomExpr{:dagger}; kw...)
  f = first(expr)
  if (head(f) != :generator) print(io, "\\left(") end
  show_latex(io, f)
  if (head(f) != :generator) print(io, "\\right)") end
  print(io, "^\\dagger")
end

ext: jl


text:
module TrueSkill
    import Base: show, isless, *, /
    export TrueSkillEnv, Rating, expose, rate, mu, sigma
    include("gaussian.jl")
    include("factorgraph.jl")
    include("distribution.jl")
    include("mainbody.jl")
end


text:
module GPUMarker

using ..LIKWID: LibLikwid, NvMon, gpusupport, capture_stderr!

"""
Initialize the NvMon Marker API of the LIKWID library. Must be called previous to all other functions.
"""
function init()
    gpusupport() || error(
        "liblikwid hasn't been compiled with GPU support (i.e. `NVIDIA_INTERFACE=true`).",
    )
    return LibLikwid.likwid_gpuMarkerInit()
end

"""
Register a region with name `regiontag` to the GPU Marker API. On success, `true` is returned.

This is an optional function to reduce the overhead of region registration at `Marker.startregion`.
If you don't call `registerregion`, the registration is done at `startregion`.
"""
function registerregion(regiontag::AbstractString)
    ret = LibLikwid.likwid_gpuMarkerRegisterRegion(regiontag)
    return ret == 0
end

"""
Start measurements under the name `regiontag`. On success, `true` is returned.
"""
function startregion(regiontag::AbstractString)
    ret = LibLikwid.likwid_gpuMarkerStartRegion(regiontag)
    return ret == 0
end

"""
Stop measurements under the name `regiontag`. On success, `true` is returned.
"""
function stopregion(regiontag::AbstractString)
    ret = LibLikwid.likwid_gpuMarkerStopRegion(regiontag)
    return ret == 0
end

"""
    getregion(regiontag::AbstractString) -> nevents, events, time, count
Get the intermediate results of the region identified by `regiontag`. On success, it returns
    * `nevents`: the number of events in the current group,
    * `events`: a list with all the aggregated event results,
    * `time`: the measurement time for the region and
    * `count`: the number of calls.
"""
function getregion(regiontag::AbstractString)
    current_group = NvMon.get_id_of_active_group()
    nevents = Ref(NvMon.get_number_of_events(current_group))
    events_ref = Ref{Ptr{Float64}}()
    time = Ref(0.0)
    count = Ref(Int32(0))
    ngpus = Ref(Int32(0))
    LibLikwid.likwid_gpuMarkerGetRegion(regiontag, ngpus, nevents, events_ref, time, count)
    events = unsafe_wrap(Array, events_ref[], nevents[])
    return ngpus[], nevents[], events, time[], count[]
end

"""
Switch to the next event set in a round-robin fashion.
If you have set only one event set on the command line, this function performs no operation.
"""
nextgroup() = LibLikwid.likwid_gpuMarkerNextGroup()

"""
Reset the values stored using the region name `regiontag`.
On success, `true` is returned.
"""
function resetregion(regiontag::AbstractString)
    ret = LibLikwid.likwid_gpuMarkerResetRegion(regiontag)
    return ret == 0
end

"""
Close the connection to the LIKWID GPU Marker API and write out measurement data to file.
This file will be evaluated by `likwid-perfctr`.
"""
close() = LibLikwid.likwid_gpuMarkerClose()

"""
Checks whether the NVIDIA GPU Marker API is active, i.e. julia has been started under `likwid-perfctr -G ... -W ... -m`.
"""
function isactive()
    buf = IOBuffer()
    capture_stderr!(GPUMarker.init, buf)
    s = String(take!(buf))
    return !startswith(s, "Running without GPU Marker API")
end

"""
    gpuregion(f, regiontag::AbstractString)
Adds a LIKWID GPU marker region around the execution of the given function `f` using [`GPUMarker.startregion`](@ref),
[`GPUMarker.stopregion`](@ref) under the hood.
Note that `LIKWID.GPUMarker.init()` and `LIKWID.GPUMarker.close()` must be called before and after, respectively.

# Examples
```julia
julia> using LIKWID, CUDA

julia> GPUMarker.init()

julia> gpuregion("sleeping...") do
           sleep(1)
       end
true

julia> gpuregion(()->CUDA.rand(100), "create rand vec")
true

julia> GPUMarker.close()

```
"""
function gpuregion(f, regiontag::AbstractString)
    GPUMarker.startregion(regiontag)
    f()
    return GPUMarker.stopregion(regiontag)
end

"""
Convenience macro for flanking code with [`GPUMarker.startregion`](@ref) and [`GPUMarker.stopregion`](@ref).

# Examples
```julia
julia> using LIKWID, CUDA

julia> GPUMarker.init()

julia> @gpuregion "sleeping..." sleep(1)
true

julia> @gpuregion "create rand vec" CUDA.rand(100)
true

julia> GPUMarker.close()

```
"""
macro gpuregion(regiontag, expr)
    q = quote
        LIKWID.GPUMarker.startregion($regiontag)
        $(expr)
        LIKWID.GPUMarker.stopregion($regiontag)
    end
    return esc(q)
end

end # module


text:
module AugmentorBase

using ColorTypes
using MappedArrays
using Interpolations
using StaticArrays
using OffsetArrays
using IdentityRanges
using MLDataPattern
using ComputationalResources
using FileIO
using LinearAlgebra

export
    testpattern,
    CPU1,
    CPUThreads,

    Reshape,

    ConvertEltype,
    MapFun,
    AggregateThenMapFun,

    Crop,
    CropNative,
    CropSize,
    CropRatio,
    RCropRatio,

    CacheIntermediate,

    Resize,

    NoOp,
    Either,

    augment,
    augment!,
    augmentbatch!

include("utils.jl")
include("types.jl")
include("operation.jl")

include("operations/convert.jl")
include("operations/mapfun.jl")

include("operations/noop.jl")
include("operations/cache.jl")

include("operations/crop.jl")
include("operations/either.jl")

include("pipeline.jl")
include("codegen.jl")
include("augment.jl")
include("augmentbatch.jl")

function __init__()
    rand_mutex[] = ReentrantLock()
end

end # module


text:
include("scalar-creation.jl")
include("tensor-creation.jl")
include("tensor-arrayif.jl")
include("tensor-indexing.jl")

include("grad.jl")

include("simple-script.jl")

include("issues.jl")


text:
# functions to extend snaq to multiple alleles case
# Claudia November 2015

global repeatAlleleSuffix = "__2"

"""
`mapAllelesCFtable(mapping file, CF file; filename, columns, delim)`

Create a new DataFrame containing the same concordance factors as in the input CF file,
but with modified taxon names. Each allele name in the input CF table is replaced by the
species name that the allele maps onto, based on the mapping file. The mapping file should have column names: allele and species.

Optional arguments:

- file name to write/save resulting CF table. If not specified, then the output
  data frame is not saved to a file.
- column numbers for the taxon names. 1-4 by default.
- delim=',': how columns are delimited. comma by default: for csv files.
"""
function mapAllelesCFtable(alleleDF::AbstractString, cfDF::AbstractString;
        filename=""::AbstractString, columns=Int[]::Vector{Int}, delim=','::Char)
    d = CSV.read(alleleDF, delim=delim)
    d2 = CSV.read(cfDF, delim=delim)
    if(filename=="")
        mapAllelesCFtable!(d,d2,columns,false,filename)
    else
        mapAllelesCFtable!(d,d2,columns,true,filename)
    end
end

# function to read a table of allele-species matchs (dataframe with 2 columns)
# and a table of CF in the allele names, and replace all the allele names
# to the species names
# this will create a new CF table, will not rewrite on the original one
# filename is the name to give to the new table, if write=true
function mapAllelesCFtable!(alleleDF::DataFrame, cfDF::DataFrame, co::Vector{Int},write::Bool,filename::AbstractString)
    size(cfDF,2) >= 7 || error("CF DataFrame should have 7+ columns: 4taxa, 3CF, and possibly ngenes")
    if length(co)==0 co=[1,2,3,4]; end
    compareTaxaNames(alleleDF,cfDF,co)
    for j in 1:4
        for ia in 1:size(alleleDF,1) # for all alleles
            cfDF[co[j]] = map(x->replace(string(x),
                                         Regex("^$(string(alleleDF[ia,:allele]))\$"),
                                         alleleDF[ia,:species]),
                              cfDF[co[j]])
        end
    end
    if(write)
        filename != "" || error("want to write table of CF with alleles mapped but filename is empty")
        CSV.write(filename, cfDF)
    end
    return cfDF
end

# function to clean a df after changing allele names to species names
# inside mapAllelesCFtable
# by deleting rows that are not informative like sp1 sp1 sp1 sp2
# keepOne=true: we only keep one allele per species
function cleanAlleleDF!(newdf::DataFrame, cols::Vector{Int};keepOne=false::Bool)
    global DEBUG
    withngenes = (length(cols)==8)
    delrows = Int[] # indices of rows to delete
    repSpecies = String[]
    if(isa(newdf[1,cols[1]],Integer)) #taxon names as integers: we need this to be able to add __2
        newdf[cols[1]] = map(x->string(x),newdf[cols[1]])
        newdf[cols[2]] = map(x->string(x),newdf[cols[2]])
        newdf[cols[3]] = map(x->string(x),newdf[cols[3]])
        newdf[cols[4]] = map(x->string(x),newdf[cols[4]])
    end
    row = Vector{String}(4)
    for i in 1:size(newdf,1) #check all rows
        DEBUG && println("row number: $i")
        # fixit: check for no missing value, or error below
        map!(j -> newdf[i,cols[j]], row, 1:4)
        DEBUG && println("row $(row)")
        uniq = unique(row)
        DEBUG && println("unique $(uniq)")

        keep = false # default: used if 1 unique name, or 2 in some cases
        if(length(uniq) == 4)
            keep = true
        else
            if(!keepOne)
                if(length(uniq) == 3) #sp1 sp1 sp2 sp3
                    keep = true
                    for u in uniq
                        DEBUG && println("u $(u), typeof $(typeof(u))")
                        ind = row .== u #taxon names matching u
                        DEBUG && println("taxon names matching u $(ind)")
                        if(sum(ind) == 2)
                            push!(repSpecies,string(u))
                            found = false
                            for k in 1:4
                                if(ind[k])
                                    if(found)
                                        DEBUG && println("found the second one in k $(k), will change newdf[i,cols[k]] $(newdf[i,cols[k]]), typeof $(typeof(newdf[i,cols[k]]))")
                                        newdf[i,cols[k]] = string(u, repeatAlleleSuffix)
                                        break
                                    else
                                        found = true
                                    end
                                end
                            end
                            break
                        end
                    end
                elseif(length(uniq) == 2)
                    # keep was initialized to false
                    for u in uniq
                        DEBUG && println("length uniq is 2, u $(u)")
                        ind = row .== u
                        if(sum(ind) == 1 || sum(ind) == 3)
                            DEBUG && println("ind $(ind) is 1 or 3, should not keep")
                            break
                        elseif(sum(ind) == 2)
                            DEBUG && println("ind $(ind) is 2, should keep")
                            keep = true
                            found = false
                            push!(repSpecies,string(u))
                            for k in 1:4
                                if(ind[k])
                                    if(found)
                                        newdf[i,cols[k]] = string(u, repeatAlleleSuffix)
                                        break
                                    else
                                        found = true
                                    end
                                end
                            end
                        end
                    end
                end
                DEBUG && println("after if, keep is $(keep)")
            end
        end
        keep || push!(delrows, i)
        DEBUG && (@show keep)
    end
    DEBUG && (@show delrows)
    DEBUG && (@show repSpecies)
    nrows = size(newdf,1)
    nkeep = nrows - length(delrows)
    if(nkeep < nrows)
        print("""found $(length(delrows)) 4-taxon sets uninformative about between-species relationships, out of $(nrows).
              These 4-taxon sets will be deleted from the data frame. $nkeep informative 4-taxon sets will be used.
              """)
        nkeep > 0 || warn("All 4-taxon subsets are uninformative, so the dataframe will be left empty")
        deleterows!(newdf, delrows)
    end
    # @show size(newdf)
    return unique(repSpecies)
end

# function to merge rows that have repeated taxon names by using the weigthed average of CF
# (if info on number of genes is provided) or simple average
function mergeRows!(df::DataFrame, cols::Vector{Int})
    sorttaxa!(df, cols)
    n4tax  = size(df,1) # total number of 4-taxon sets
    delrows = Int[] # indices of rows to delete
    nrows  =  ones(Int ,n4tax) # total # of rows that row i combines. 0 if row i is to be deleted.
    for i in 1:n4tax # for each row / 4-taxon set
        if nrows[i]>0
            for j in (i+1):n4tax # rows with larger index
                nrows[j]>0 || continue # skip j if it matched a 4-taxon set earlier
                rowmatch = true
                for k in 1:4
                    rowmatch *= (df[j,cols[k]] == df[i,cols[k]])
                end
                rowmatch || continue   # skip j if doesn't match i^th taxon set
                nrows[j]=0
                push!(delrows, j)
                for k in 5:length(cols)
                    df[i,cols[k]] += df[j,cols[k]]
                end
                nrows[i] += 1
            end
        end
    end
    # @show head(df); @show delrows[1:10]; @show length(delrows); @show sum(nrows); println("number with nrows>0: $(sum(map(x -> x>0, nrows)))")
    for ir in delrows
        if nrows[ir]>0
            println("problem: ir=$ir, nrows=$(nrows[ir])")
            @show df[ir,:]
        end
    end
    length(delrows)>0 || return df
    sort!(delrows)
    for i in 1:size(df,2)
        deleteat!(df[i], delrows) # more efficient than deleterows!(df, delrows)
    end
    deleteat!(nrows, delrows)
    n4tax = size(df,1) # re-defined
    print("$n4tax unique 4-taxon sets were found. CF values of repeated 4-taxon sets will be averaged")
    println((length(cols)>7 ? " (ngenes too)." : "."))
    if length(cols)>7 && eltype(df[cols[8]])<: Integer # ngenes is present: integer. Need to convert to float
        df[cols[8]] = convert(Array{Float64},df[cols[8]])
    end
    for i in 1:n4tax
        nrows[i]>0 || error("Original $(delrows[i])) was retained (now row $i) but has nrows=0")
        for k in 5:length(cols)
            df[i,cols[k]] /= nrows[i]
        end
    end
    return df
end

# function to expand leaves in tree to two individuals
# based on cf table with alleles mapped to species names
function expandLeaves!(repSpecies::Union{Vector{String},Vector{Int}},tree::HybridNetwork)
    for sp in repSpecies
        for n in tree.node
            if(n.name == sp) #found leaf with sp name
                n.leaf || error("name $(sp) should correspond to a leaf, but it corresponds to an internal node")
                length(n.edge) == 1 || error("leaf $(sp) should have only one edge attached and it has $(length(n.edge))")
                if(n.edge[1].length == -1.0)
                    setLength!(n.edge[1],1.0)
                end
                removeLeaf!(tree,n)
                n.leaf = false
                n.edge[1].istIdentifiable = true
                n.name = ""
                max_node = maximum([e.number for e in tree.node]);
                max_edge = maximum([e.number for e in tree.edge]);
                e1 = Edge(max_edge+1,0.0)
                e2 = Edge(max_edge+2,0.0)
                n1 = Node(max_node+1,true,false,[e1])
                n2 = Node(max_node+2,true,false,[e2])
                setNode!(e1,n1)
                setNode!(e1,n)
                setNode!(e2,n2)
                setNode!(e2,n)
                setEdge!(n,e1)
                setEdge!(n,e2)
                pushNode!(tree,n1)
                pushNode!(tree,n2)
                pushEdge!(tree,e1)
                pushEdge!(tree,e2)
                n1.name = string(sp)
                n2.name = string(sp,repeatAlleleSuffix)
                push!(tree.names,n2.name)
                break
            end
        end
    end
end

# function to compare the taxon names in the allele-species matching table
# and the CF table
function compareTaxaNames(alleleDF::DataFrame, cfDF::DataFrame, co::Vector{Int})
    checkMapDF(alleleDF)
    println("Allele map: found $(length(alleleDF[1])) allele-species matches")
    CFtaxa = convert(Array, unique(stack(cfDF[co[1:4]], 1:4)[:value]))
    CFtaxa = map(x->string(x),CFtaxa) #treat as string
    alleleTaxa = map(x->string(x),alleleDF[:allele]) #treat as string
    sizeCF = length(CFtaxa)
    sizeAllele = length(alleleTaxa)
    if(sizeAllele > sizeCF)
        println("Allele map: more alleles in the mapping file: $(sizeAllele) than in the CF table: $(sizeCF). Extra allele names will be ignored")
        alleleTaxa = intersect(alleleTaxa,CFtaxa)
    elseif(sizeAllele < sizeCF)
        println("Allele map: fewer alleles in the mapping file: $(sizeAllele) than in the CF table: $(sizeCF). Some names in the CF table will remain unchanged")
    end
    unchanged = setdiff(CFtaxa,alleleTaxa)
    if(length(unchanged) == length(CFtaxa))
        warn("None of the taxon names in CF table match with allele names in the mapping file")
    end
    if(isempty(unchanged))
        println("All allele names in the CF table were found in the allele-species mapping file.")
    else
        warn("not all alleles were mapped")
        println("The following taxa in the CF table were not found in the allele-species mapping file:\n$(unchanged)")
    end
    return nothing
end

# function to check that the allele df has one column labelled alleles and one column labelled species
function checkMapDF(alleleDF::DataFrame)
    size(alleleDF,2) <= 2 || error("Allele-Species matching Dataframe should have at least 2 columns")
    size(alleleDF,2) >= 2 || warn("allele mapping file contains more than two columns: will ignore all columns not labelled allele or species")
    try
        alleleDF[:allele]
    catch
        error("In allele mapping file there is no column named allele")
    end
    try
        alleleDF[:species]
    catch
        error("In allele mapping file there is no column named species")
    end
end

## function to check if a new proposed topology satisfies the condition for
## multiple alleles: no gene flow to either allele, and both alleles as sister
## returns false if the network is not ok
function checkTop4multAllele(net::HybridNetwork)
    for n in net.leaf
        if(endswith(n.name, repeatAlleleSuffix))
            n.leaf || error("weird node $(n.number) not leaf in net.leaf list")
            length(n.edge) == 1 || error("weird leaf with $(length(n.edge)) edges")
            par = getOtherNode(n.edge[1],n)
            if(par.hybrid) ## there is gene flow into n
                return false
            end
            nameOther = replace(n.name,repeatAlleleSuffix,"")
            foundOther = false
            for i in 1:3
                other = getOtherNode(par.edge[i],par)
                if(other.leaf && other.name == nameOther)
                    foundOther = true
                end
            end
            foundOther || return false
        end
    end
    return true
end

## function to merge the two alleles into one
function mergeLeaves!(net::HybridNetwork)
    leaves = copy(net.leaf) # bc we change this list
    for n in leaves
        if(endswith(n.name, repeatAlleleSuffix))
            n.leaf || error("weird node $(n.number) not leaf in net.leaf list")
            length(n.edge) == 1 || error("weird leaf with $(length(n.edge)) edges")
            par = getOtherNode(n.edge[1],n)
            foundOther = false
            other = Node()
            nameOther = replace(n.name,repeatAlleleSuffix,"")
            for i in 1:3
                other = getOtherNode(par.edge[i],par)
                if(other.leaf && other.name == nameOther)
                    foundOther = true
                    break
                end
            end
            if(!foundOther)
                checkTop4multAllele(net) || error("current network does not comply with multiple allele condition")
                error("strange network that passes checkTop4multAllele, but cannot find the other allele for $(n.name)")
            end
            removeEdge!(par,n.edge[1])
            removeEdge!(par,other.edge[1])
            deleteNode!(net,n)
            deleteNode!(net,other)
            deleteEdge!(net,n.edge[1])
            deleteEdge!(net,other.edge[1])
            par.name = other.name
            par.leaf = true
            push!(net.leaf,par)
            net.numTaxa += 1
        end
    end
end


text:
using RunTests
using Base.Test

@testmodule TestModuleWithOnePassingTest begin

  function test_hello()
    @test true
  end

end


text:
"""
    second-order cone
"""
function second_order_cone_projection(z)
    n = length(z)

    z0 = z[1]
    z1 = view(z, 2:n)

    if norm(z1) <= z0
        return z, true
    elseif norm(z1) <= -z0
        return zero(z), false
    else
        a = 0.5 * (1.0 + z0 / norm(z1))
        z_proj = zero(z)
        z_proj[1] = a * norm(z1)
        z_proj[2:end] = a * z1
        return z_proj, false
    end
end

function second_order_cone_product(z, s)
    n = length(z)
    SVector{n}([z' * s; z[1] * view(s, 2:n) + s[1] * view(z, 2:n)])
end

# check that second-order cone constraints are satisfied
function second_order_cone_check(x, idx_soc)
    for idx in idx_soc
        !second_order_cone_projection(view(x, idx))[2] && return false
    end
    return true
end

# check that inequality (non-negative orthant) constraints are satisfied
function inequality_check(x, idx_ineq)
    for i in idx_ineq
        if x[i] <= 0.0
            return false
        end
    end
    return true
end

function cone_check(x, idx_ineq, idx_soc)
    !inequality_check(x, idx_ineq) && return true
    !second_order_cone_check(x, idx_soc) && return true
end

max_issues_repo_name: thowell/DirectMotionPlanning


text:
using Documenter, JlrsReflect
makedocs(sitename="JlrsReflect")


text:
"""
    simulate(
        scene::Scene{E}, roadway::R, models::Dict{I,M}, nticks::Int64, timestep::Float64;
        rng::AbstractRNG = Random.GLOBAL_RNG, callbacks = nothing
    ) where {E<:Entity,A,R,I,M<:DriverModel}

Simulate a `scene`. For detailed information, consult the documentation of `simulate!`.
By default, returns a vector containing one scene per time step.
"""
function simulate(
    scene::Scene{E},
    roadway::R,
    models::Dict{I,M},
    nticks::Int64,
    timestep::Float64;
    rng::AbstractRNG = Random.GLOBAL_RNG,
    callbacks = nothing,
) where {E<:Entity,A,R,I,M<:DriverModel}
    scenes = [Scene(E, length(scene)) for i=1:nticks+1]
    n = simulate!(
        scene, roadway, models, nticks, timestep, scenes, nothing,
        rng=rng, callbacks=callbacks
    )
    return scenes[1:(n+1)]
end

"""

    simulate!(
        scene::Scene{E}, roadway::R, models::Dict{I,M},
        nticks::Int64, timestep::Float64,
        scenes::Vector{Scene{E}}, actions::Union{Nothing, Vector{Scene{A}}} = nothing;
        rng::AbstractRNG = Random.GLOBAL_RNG, callbacks = nothing
    ) where {E<:Entity,A<:EntityAction,R,I,M<:DriverModel}

Simulate the entities in `scene` along a `roadway` for a maximum of
`nticks` time steps of size `timestep`.
Returns the number of successfully performed timesteps.

At each time step, `models` is used to determine the action for each agent.
`scenes` and `actions` are pre-allocated vectors of `Scene`s containing either
`Entity`s (for scenes) or `EntityAction`s (for actions).
If `actions` is equal to `nothing` (default), the action history is not tracked.
`scenes` must always be provided.

`callbacks` is an array of callback functions which are invoked before
the simulation starts and after every simulation step.
Any callback function can cause an early termination by returning `true`
(the default return value for callback functions should be `false`).
The random number generator for the simulation can be provided using the `rng`
keyword argument, it defaults to `Random.GLOBAL_RNG`.
"""
function simulate!(
    scene::Scene{E},
    roadway::R,
    models::Dict{I,M},
    nticks::Int64,
    timestep::Float64,
    scenes::Vector{Scene{E}},
    actions::Union{Nothing, Vector{Scene{A}}} = nothing;
    rng::AbstractRNG = Random.GLOBAL_RNG,
    callbacks = nothing
    ) where {E<:Entity,A<:EntityAction,R,I,M<:DriverModel}

    copyto!(scenes[1], scene)

    # potential early out right off the bat
    if (callbacks !== nothing) && _run_callbacks(callbacks, scenes, actions, roadway, models, 1)
        return 0
    end

    for tick in 1:nticks

        empty!(scenes[tick + 1])
        if (actions !== nothing) empty!(actions[tick]) end

        for (i, veh) in enumerate(scenes[tick])

            observe!(models[veh.id], scenes[tick], roadway, veh.id)
            a = rand(rng, models[veh.id])

            veh_state_p  = propagate(veh, a, roadway, timestep)

            push!(scenes[tick + 1], Entity(veh_state_p, veh.def, veh.id))
            if (actions !== nothing) push!(actions[tick], EntityAction(a, veh.id)) end

        end

        if !(callbacks === nothing) && _run_callbacks(callbacks, scenes, actions, roadway, models, tick+1)
            return tick
        end
    end
    return nticks
end


text:
using QuantumTomography, Cliffords, LinearAlgebra, StatsBase

"""
    squeeze(A::AbstractArray)

Drops singleton dimentions from array structures.
Helper function to be used with caution!

https://stackoverflow.com/questions/52505760/dropping-singleton-dimensions-in-julia
"""
function squeeze(A::AbstractArray)
    singleton_dims = tuple((d for d in 1:ndims(A) if size(A, d) == 1)...)
    return dropdims(A, dims=singleton_dims)
end

"""
    tomo_gate_set(nbrQubits, nbrAxes; pulse_type::String="Clifford", prep_meas::Integer=1)

Return a set of state preparation or readout unitary gates for a give
`nbrQubits` along a given `nbrAxes`.
See http://arxiv.org/abs/quant-ph/0308098v1 for more information

# Arguments
- `nbrAxes::Integer`: number of single-qubit axes ∈ [4,6,12]
- `pulse_type::String="Clifford"`: prepared states/meas. axes
- `prep_meas::Integer=1`: 1 for prep gates, 2 for meas. gates

# Examples
```julia-repl
julia> tomo_gate_set(2, 4)
16-element Array{Array{Complex{Float64},2},1}:
[1.0 - 0.0im -0.0 + 0.0im -0.0 + 0.0im 0.0 + 0.0im; -0.0 + 0.0im 1.0 - 0.0im
 0.0 + 0.0im -0.0 + 0.0im; -0.0 + 0.0im 0.0 + 0.0im 1.0 - 0.0im -0.0 + 0.0im;
 0.0 + 0.0im -0.0 + 0.0im -0.0 + 0.0im 1.0 - 0.0im]
[-0.7071067811865477 + 0.0im 0.0 + 0.7071067811865475im ...
```
"""
function tomo_gate_set(nbrQubits, nbrAxes; pulse_type="Clifford", prep_meas = 1)
    if nbrAxes==4
        # Four pulse set
        if pulse_type == "Clifford"
            Uset1Q = [complex(RI),
                      exp(-im*pi/4*X),
                      exp(-im*pi/4*Y),
                      -im*X]
        elseif pulse_type == "Tetra"
            if prep_meas == 1
                Uset1Q = [complex(RI),
                          exp(-im*acos(-1/3)*X),
                          exp(-im*2pi/3*Z)*exp(-im*acos(-1/3)*X),
                          exp(+im*2pi/3*Z)*exp(-im*acos(-1/3)*X)]
            else
                Uset1Q = [complex(RI),
                          exp(-im*acos(-1/3)*X),
                          exp(+im*acos(-1/3)*X)*exp(+im*2pi/3*Z),
                          exp(+im*acos(-1/3)*X)*exp(-im*2pi/3*Z)]
            end
        else
            error("Invalid prep./meas. pulse pulse_type.
            Must be ∈ {\"Tetra\", \"Clifford\"}")
        end
    elseif nbrAxes==6
        # Six pulse set
        Uset1Q = [complex(RI),
                  exp(-im*pi/4*X),
                  exp(+im*pi/4*X),
                  exp(-im*pi/4*Y),
                  exp(+im*pi/4*Y),
                  -im*X]
    elseif nbrAxes==12
        # 12 pulse set
        Uset1Q = [complex(RI),
                  -im*X,
                  -im*Y,
                  -im*Z,
                  exp(-im*pi/3*(+X+Y-Z)/sqrt(3)),  #X+Y-Z 120
                  exp(-im*pi/3*(+X-Y+Z)/sqrt(3)),  #X-Y+Z 120
                  exp(-im*pi/3*(-X+Y+Z)/sqrt(3)),  #-X+Y+Z 120
                  exp(-im*pi/3*(-X-Y-Z)/sqrt(3)),  #X+Y+Z -120 (equivalent to -X-Y-Z 120)
                  exp(-im*pi/3*(+X+Y+Z)/sqrt(3)),   #X+Y+Z 120
                  exp(-im*pi/3*(-X+Y-Z)/sqrt(3)),  #X-Y+Z -120 (equivalent to -X+Y-Z 120)
                  exp(-im*pi/3*(+X-Y-Z)/sqrt(3)),  #-X+Y+Z -120 (equivalent to X-Y-Z 120)
                  exp(-im*pi/3*(-X-Y+Z)/sqrt(3))]  #X+Y-Z -120 (equivalent to -X-Y+Z 120)
    else
        error("Invalid number of pulses.  Must be ∈ [4,6,12]");
    end

    # Now the gate set is the cartesian product of the 1Q gate set over the
    # number of qubits. Unfornately, Julia's default product is anti-
    # lexicographic (first index is fastest), so we need to reverse the
    # gate order before taking the kronecker product.
    gateSet = Array{ComplexF64,2}[]
    for gates in Base.product([Uset1Q for _ in 1:nbrQubits]...)
        push!(gateSet, kron(1, reverse(gates)...))
    end
    return gateSet
end

"""
    QST_FLSQ(expResults, varMat, measPulseMap, measOpMap, measPulseUs, measOps)

Function to perform unconstrained least-squares inversion of state
tomography data.

# Arguments
- `expResults::Array{Number}`: data array
- `varmat::Array{Number}`: covariance array for the expResults
- `measPulseMap::`: array mapping each experiment outcome to the corresponding
                    measPulseUs
- `measOpMap::`: array mapping each experiment outcome to a measurement operator
- `measPulseUs::`: array of unitaries applied before measurement pulses
- `measOps::`: array of measurement operators for each channel

# Examples
```julia-repl
julia> QST_LSQ(2, 4)

```
"""
function QST_FLSQ(expResults,
                  varMat,
                  measPulseMap,
                  measOpMap,
                  measPulseUs,
                  measOps)
    # construct the vector of observables for each experiment
    obs = Matrix{Complex{Float64}}[]
    for ct in 1:length(expResults)
        U = measPulseUs[measPulseMap[ct]]
        op = measOps[measOpMap[ct]]
        obs_ct = U' * op * U
        obs_ct = (obs_ct + obs_ct')/2 # force to be Hermitian
        push!(obs, obs_ct)
    end

    #return measOps
    # in order to constrain the trace to unity, add an identity observable
    # and a corresponding value to expResults
    push!(obs, Diagonal((1.0)*fill(I.λ, size(measOps[1])))) # this can be replaced with I(size(measOps[1])) for Julia>=1.3
    expResults2 = [expResults; 1]
    # corresponding variance chosen arbitrarily (it should be very small)
    varMat2 = [varMat; minimum(varMat)]
    tomo = FreeLSStateTomo(obs)

    ρest, obj, status = fit(tomo, expResults2, varMat2)
    if status != :Optimal
        println("FreeLSStateTomo fit return status: $status")
    end
    return ρest
end
"""
    QST_LSQ(expResults, varMat, measPulseMap, measOpMap, measPulseUs, measOps)

Function to perform constrained least-squares quantum state tomography.

# Arguments
- `expResults::Array{Number}`: data array.  This is the list of expectation
                               values for each of the measurements obserables.
- `varmat::Array{Number}`: variance matrix for data in expResults
- `measPulseMap::`: array mapping each experiment to a measurement readout
                    pulse
- `measOpMap::`: array mapping each experiment to a measurement channel
- `measPulseUs::`: array of unitaries of measurement pulses.  These are
                   pulses applied before measurement in the experimental data,
                   mapping the measurement to the correct axis
- `measOps::`: array of measurement operators for each channel.  In the case of
               ML tomography, these need to be text book projectors to diagonal states.

# Returns
- `ρest::Array{Complex{Float64},d}`: d dimensional estimate of the density
                                     matrix obtained by constrained
                                     least-squares.

# Examples
```julia-repl
julia> QST_LSQ(expResults, measPulseMap, measOpMap, measPulseUs, measOps)
2×2 Array{Complex{Float64},2}:
 0.295239+2.94858e-14im  -0.21073+0.3957im
 -0.21073-0.3957im       0.704761-3.64935e-15im
```
"""
function QST_LSQ(expResults,
                 varMat,
                 measPulseMap,
                 measOpMap,
                 measPulseUs,
                 measOps)
    # construct the vector of observables for each experiment
    obs = Matrix{Complex{Float64}}[]
    for ct in 1:length(expResults)
        U = measPulseUs[measPulseMap[ct]]
        op = measOps[measOpMap[ct]]
        obs_ct = U' * op * U
        obs_ct = (obs_ct + obs_ct')/2 # force to be Hermitian
        push!(obs, obs_ct)
    end
    tomo = LSStateTomo(obs)

    ρest, obj, status = fit(tomo, expResults, varMat)
    if status != :Optimal
        println("LSStateTomo fit return status: $status")
    end
    return ρest
end

"""
    QST_ML(expResults, measPulseMap, measOpMap, measPulseUs, measOps)

Function to perform maximum-likelihood quantum state tomography.  This function
is usually wrapped by the analyzeStateTomo function.

# Arguments
- `expResults::Array{Number}`: data array.  This is a list of expectation
                               values for each of the measurements obserables.
- `measPulseMap::`: array mapping each experiment to a measurement
                    readout pulse
- `measOpMap::`: array mapping each experiment to a measurement channel
- `measPulseUs::`: array of unitaries of measurement pulses.  These are
                   pulses applied before measurement in the experimental data,
                   mapping the measurement to the correct axis
- `measOps::`: array of measurement operators for each channel.  In the case of
               ML tomography, these need to be text book projectors to
               diagonal states.

# Returns
- `ρest::Array{Complex{Float64},d}`: d dimensional estimate of the density
                                     matrix obtained by maximum likelihood.

# Examples
```julia-repl
julia> QST_ML(expResults, measPulseMap, measOpMap, measPulseUs, measOps)
2×2 Array{Complex{Float64},2}:
 0.295239+2.94858e-14im  -0.21073+0.3957im
 -0.21073-0.3957im       0.704761-3.64935e-15im
```
"""
function QST_ML(expResults,
                measPulseMap,
                measOpMap,
                measPulseUs,
                measOps; n=100_000, β=0.0, maxiter=5000, ϵ=1000)
    # construct the vector of observables for each experiment
    obs = Matrix{Complex{Float64}}[]
    for ct in 1:length(expResults)
        U = measPulseUs[measPulseMap[ct]]
        op = measOps[measOpMap[ct]]
        push!(obs, Hermitian(U' * op * U)) # force to be Hermitian
    end
    # ML obs must be POVMs -> Hermitian, possitive-semidefinite, and trace 1
    if length(obs) < 6
        @warn("State observations do not form a POVM!")
    end

    tomo = MLStateTomo(obs)

    ρest, obj, status = fit(tomo,
                            expResults,
                            maxiter=maxiter,
                            δ=1/ϵ,
                            λ=β)
    if status != :Optimal
        println("MLStateTomo fit return status: $status")
    end
    return ρest
end

"""
    analyzeStateTomo(data::Dict{String,Dict{String,Array{Any,N} where N}},
                          nbrQubits::Int,
                          nbrAxes::Int ∈ [4,6,12];
                          nbrCalRepeats::Int=2)

Function to setup and run quantum state tomography for a given number of qubits
and measurement axes.

# Arguments
- `data`: data array strucured as a set of data with a string name with a 'data'
          key and a 'variance' key.  The variance is required for all tomography
          reconstructions except for the free-LSQ tomo.  Also, for two-qubit
          tomography, the correlation data between the two qubit data sets is
          required for reconstruction.
- `nbrQubits`: number of qubits
- `nbrAxes`: number of measurements.  Either 4 or 6.  12 is possible but left
             for the user to do manually.
- `nbrCalRepeats`: number of repeated calibration points per calibration state

# Returns
- `rhoLSQ` : a 2^nbrQubits x 2^nbrQubits complex density matrix reconstructed
            with least-squares
- `rhoML` : (optional) a 2 x 2 complex density matrix reconstructed
            with least-squares.  Note this is only supported with single qubit
            data and six axes measurements.

# Examples
julia> datapath = "/path/to/data/folder"
julia> data, desc = load_data(datapath,24,"200315",load_var=true);
julia> rhoLSQ  = Qlab.analyzeStateTomo(data[1],2,4);
julia> Qlab.pauli_set_plot(rhoLSQ)
"""
function analyzeStateTomo(data::Dict{String,Dict{String,Array{Any,N} where N}},
                          nbrQubits,
                          nbrAxes;
                          nbrCalRepeats=2)

    measOps = Matrix{Float64}[]
    tomoData = Float64[]
    varData = Float64[]
    numMeas = length(data)

    for data_q in values(data)
        # Average over calibration repeats
        data_ql = data_q["data"]
        if nbrCalRepeats == 0
            # In this case, assume data is already calibrated and insert
            # standard projectors
            append!(measOps, real(Qlab._create_ml_POVM(nbrQubits)))
            append!(tomoData, real(data_ql[1:end]))
            append!(varData, real(data_q["variance"])[1:end])
        else
            calData = real(data_ql[end-nbrCalRepeats*(2^nbrQubits )+1:end])
            avgCalData = mean(reshape(calData, nbrCalRepeats, 2^nbrQubits), dims=1)
            # Pull out the calibrations as diagonal measurement operators
            push!(measOps, diagm(0 => avgCalData[:]))

            #The data to invert
            append!(tomoData, real(data_ql[1:end-nbrCalRepeats*(2^nbrQubits)]))

            #variance
            append!(varData, real(data_q["variance"])[1:end-nbrCalRepeats*(2^nbrQubits)])
        end
    end
    # Map each experiment to the appropriate readout pulse
    # These are integers 1:length(data), each repeated numAxes^nbrQubits times
    measOpMap = repeat(1:numMeas, inner=nbrAxes^nbrQubits)
    # These are integers 1:nbrAxes^nbrQubits, unrolled length(data) times
    measPulseMap = repeat(1:nbrAxes^nbrQubits, outer=numMeas)
    # Use a helper to get the measurement unitaries.
    measPulseUs = tomo_gate_set(nbrQubits, nbrAxes)

    # Now call the inversion routines
    # First least squares
    rhoLSQ = QST_LSQ(tomoData,
                     varData,
                     measPulseMap,
                     measOpMap,
                     measPulseUs,
                     measOps)

    # Constrained maximum-likelihood is currently unsupported in general
    #
    # The reason for this has to do with the delicate nature of constructing
    # the correlations data and make operators like POVMs physical etc... If
    # this becomes a dire need, we can always revisit.  For most applications
    # of interest, LSQ is perfectly good

    # The one exception to this is single qubit data where measurements form a
    # proper POVM.  In experimental language this is the case where we take six
    # data points for the state reconstruction
    if nbrAxes == 6 && nbrQubits == 1
        rhoML = QST_ML(tomoData,
                       measPulseMap,
                       measOpMap,
                       measPulseUs,
                       measOps)

        return rhoLSQ, rhoML
    else
        return rhoLSQ, []
    end
    # plotting to be implemented    pauliSetPlot(rho2pauli(rhoLSQ), newplot)
end

"""
    _create_ml_POVM(numQubits::Int)

Create the textbook POVM for a given qubit system.

# Returns
- `mlPOVM::Array{ComplexF64}` : an array of numQubits^2 x numQubits^2
                                dimentional set of projectors that form a POVM

"""
function _create_ml_POVM(numQubits::Int)
    mlPOVM = Array{ComplexF64}[]
    if numQubits == 1
        push!(mlPOVM, diagm([1.,0.]))
        push!(mlPOVM, diagm([0.,1.]))
    elseif numQubits == 2
        push!(mlPOVM, diagm([1.,0.,0.,0.]))
        push!(mlPOVM, diagm([0.,1.,0.,0.]))
        push!(mlPOVM, diagm([0.,0.,1.,0.]))
        push!(mlPOVM, diagm([0.,0.,0.,1.]))
    end
    return mlPOVM
end

"""
    _parese_exp_num(n::Int, n_qubits::Int)

Determine the number of calibration points, calibration repeats, axes in a
given tomography data set.  Helper function for the StateTomo structure.  This
function makes many assumptions about the structure of the data and is built
using heuristics.  Your milage may vary.

# Arguments
- `n` : total number of experimental data points including
        calibration points
- `n_qubits` : number of qubits represented in the data set.  Only one
               and two qubit tomography is supported.
# Returns
- `numCalRepeats::Int` : number of repeats for each calibration point
- `numCals::Int` : the total number of calibration points in an experiment
- `numAxes::Int` : the number of of axes were observations were made.
                   This must be 4 or 6.
"""
function _parese_exp_num(numDataPoints::Int, numQubits::Int)

    # Assume zero cals as a base case for calibration or manual data scaling
    numCalRepeats = 0
    numCals = 0
    nbr_basis_states = (numQubits == 1) ? 2 : 4

    # Determine the cal repeats number
    # Given the number of qubits and the limited possible number of observables
    # search for a possible number of calibration repeats that matches the data
    for i in [4,6].^numQubits
        numCals_guess = numDataPoints - i
        nbrRepeats_guess = numCals_guess/nbr_basis_states
        if numDataPoints <= 6
            # catch the case where there are no cals
            nbrRepeats_guess = 0
        end
        if nbrRepeats_guess % 1 != 0
            # correct number will be a whole number
            continue
        end
        if nbrRepeats_guess < 0
            # filter negitive guesses
            continue
        end
        if !iseven(Int(nbrRepeats_guess))
            # This is a very safe assumption
            @warn("Assuming numCalRepeats is even!")
            continue
        end
        if !ispow2(Int(nbrRepeats_guess))
            # This is likely the case but warn the user
            @warn("Assuming nbr repeats is a power of 2!")
            continue
        end
        numCalRepeats = nbrRepeats_guess
    end
    if numQubits == 1 && numCalRepeats != 0
        numCals = numCalRepeats * 2
    elseif numQubits == 2 && numCalRepeats != 0
        numCals = numCalRepeats * 4
    end

    #determine the number of axes
    if numQubits == 2
        numAxes = sqrt(numDataPoints-numCals)
    elseif numQubits ==1
        numAxes = numDataPoints-numCals
    end

    # assert numAxes must equal [4,6]
    if !(numAxes in [4,6])
        error("Obervables must be 4 or 6.  Please check your data!")
    end

    return numCals, numCalRepeats, numAxes
end

"""
    _pre_process_data(data::Dict{String,Dict{String,Array{Any,N} where N}},
                           desc::Dict{String,Any})

Preprocess the data

Determine the number of qubits, and return organized datasets based on the
structure of the data.  Single dimensional data is assumed to be averaged
and 2D data is assumed to be (experiment, shot) data.  This is designed to be
a private function of the StateTomo structure.  Note, any shot data found is
left to the user to process manually.
"""
function _pre_process_data(data::Dict{String,Dict{String,Array{Any,N} where N}},
                           desc::Dict{String,Any})
    numDataPoints = 0
    corrData = false
    varData = false

    # load  and parse the data
    println("Preprocessing data")
    qubit_data_keys = filter(x -> occursin(r"([qQ]\d?\d?\d?)[ -_](\w*)", x), keys(data))

    qubits = []
    labels = []
    for i in qubit_data_keys
        ql = match(r"([qQ]\d?\d?\d?)[ -_](\w*)", i)
        push!(qubits, string(ql[1]))
        push!(labels, string(ql[2]))
    end
    # filter out any repeated entries
    unique!(qubits)
    unique!(labels)
    println("Found $(length(qubits)) sets of qubit data: " * string([string(i) * " " for i in qubits]...))
    println("Found $(length(labels)) sets of qubit data labels: " * string([string(i) * " " for i in labels]...))
    numQubits = length(qubits)
    numDatasets = length(labels)

    #pull out any correlation data
    correlation_data_sets = []
    qubit_correlation_keys = filter(x -> occursin(r"([Cc]orrelated?)", x), keys(data))
    if length(qubit_correlation_keys) != 0
        println("Correlation data found... ✓")
        corrData = true
    else
        println("Correlation data found... no")
        corrData = false
        if numQubits > 1
            @error("This appears to be two-qubit data but no correlation data
                    is provided!  Tomography will not work!  If you have
                    correlation data, please add it manually to the data
                    dictionary.")
        end
    end

    # check that at least one variance dataset exists
    # Least-squares reconstruction will not be possible without variance data
    variance_data = []
    for i in keys(data)
        if length(filter(x -> occursin(r"([Vv]ariance)", x), keys(data[i]))) != 0
            println("Variance data found for dataset: $(i)")
            varData = true
        else
            println("Variance data for $(i) found... no")
            varData = false
            @warn("You may not be able to do state tomography")
        end
    end

    # get the data size and classify assuming multi-dimensional data is raw,
    # integrated shot data
    tomo_data_idx = empty([], String)
    shot_data_idx = empty([], String)
    for i in keys(data)
        dims = size(squeeze(data[i]["data"]))
        if length(dims) == 1
            println("Main data set: " * string(i))
            push!(tomo_data_idx, i)
            numDataPoints = dims[1]
        elseif length(dims) > 1
            println("Shots data set: " * string(i))
            push!(shot_data_idx, i)
        end
    end

    tomoDataSets = filter(p -> p.first in tomo_data_idx, data)
    shotDataSets = filter(p -> p.first in shot_data_idx, data)

    return numQubits, numDatasets, corrData, varData, numDataPoints,
                                                      tomoDataSets,
                                                      shotDataSets
end

"""
State tomography object

This holds all the infromation necessary to do state tomography.

The object is constructed by passing it a tomography dataset and its descriptor
loaded from load_data.  Once created, this object can be passed directly to the
any of the tomographyic reconstruction methods.
"""
struct StateTomo
    numQubits::Int
    numDatasets::Int
    corrData::Bool
    varData::Bool

    tomoDataSets::Dict{String,Dict{String,Array{Any,N} where N}}
    shotDataSets::Dict{String,Dict{String,Array{Any,N} where N}}

    numAxes::Int
    numCals::Int
    numCalRepeats::Int
    numDataPoints::Int

    mlPOVM::Array{Matrix{ComplexF64}}

    """
    Basic constructor
    """
    function StateTomo(data::Dict{String,Dict{String,Array{Any,N} where N}},
                        desc::Dict{String,Any})

        numQubits,
        numDatasets,
        corrData,
        varData,
        numDataPoints,
        tomoDataSets,
        shotDataSets = _pre_process_data(data, desc)

        ############################################################
        numCals, numCalRepeats, numAxes = _parese_exp_num(numDataPoints,
                                                          numQubits)
        ################################################################
        mlPOVM = _create_ml_POVM(numQubits)
        ################################################################
        new(numQubits, numDatasets, corrData, varData, tomoDataSets,
                                                       shotDataSets,
                                                       numAxes,
                                                       numCals,
                                                       numCalRepeats,
                                                       numDataPoints,
                                                       mlPOVM)
    end
end

"""
    analyzeStateTomo(tomo::StateTomo)

Function to setup and run quantum state tomography for a given number of qubits
and measurement axes.

# Arguments
- `tomo::Qlab.StateTomo`: State tomography object constructed from the data

# Returns
- `rhoLSQ`: a 2^nbrQubits x 2^nbrQubits complex density matrix reconstructed
            with least-squares
- `rhoML` : (optional) a 2 x 2 complex density matrix reconstructed
            with least-squares.  Note this is only supported with single qubit
            data and six axes measurements.

# Examples
julia> datapath = "/path/to/data/folder"
julia> data, desc = load_data(datapath,24,"200315",load_var=true);
julia> tomo = Qlab.StateTomo(data, desc);
julia> rhoLSQ  = Qlab.analyzeStateTomo(tomo);
julia> Qlab.pauli_set_plot(rhoLSQ)
"""
function analyzeStateTomo(tomoObj::StateTomo)

    data = tomoObj.tomoDataSets
    nbrQubits = tomoObj.numQubits
    nbrAxes = tomoObj.numAxes
    nbrCalRepeats = tomoObj.numCalRepeats
    measOps = Matrix{Float64}[]
    tomoData = Float64[]
    varData = Float64[]
    numMeas = length(data)

    # mlTomoData = Float64[]
    # mlVarData = Float64[]

    # For LSQ reconstruction - make the observables out of un-scaled
    # calibration points.  Note these do not form a POVM
    for data_q in values(data)
        # Average over calibration repeats
        data_ql = data_q["data"]
        if nbrCalRepeats == 0
            # In this case, assume data is already calibrated and insert
            # standard projectors
            append!(measOps, real(Qlab._create_ml_POVM(nbrQubits)))
            append!(tomoData, real(data_ql[1:end]))
            append!(varData, real(data_q["variance"][1:end]))
        else
            calData = real(data_ql[end-nbrCalRepeats*(2^nbrQubits )+1:end])
            avgCalData = mean(reshape(calData, nbrCalRepeats, 2^nbrQubits), dims=1)
            # Pull out the calibrations as diagonal measurement operators
            push!(measOps, diagm(0 => avgCalData[:]))

            #The data to invert
            append!(tomoData, real(data_ql[1:end-nbrCalRepeats*(2^nbrQubits)]))

            #variance
            append!(varData, real(data_q["variance"])[1:end-nbrCalRepeats*(2^nbrQubits)])
        end
    end
    # Map each experiment to the appropriate readout pulse
    # These are integers 1:length(data), each repeated numAxes^nbrQubits times
    measOpMap = repeat(1:numMeas, inner=nbrAxes^nbrQubits)
    # These are integers 1:nbrAxes^nbrQubits, unrolled length(data) times
    measPulseMap = repeat(1:nbrAxes^nbrQubits, outer=numMeas)
    # Use a helper to get the measurement unitaries.
    measPulseUs = Qlab.tomo_gate_set(nbrQubits, nbrAxes)

    # Now call the inversion routines
    # First least squares
    rhoLSQ = QST_LSQ(tomoData,
                     varData,
                     measPulseMap,
                     measOpMap,
                     measPulseUs,
                     measOps)

    # Constrained maximum-likelihood is currently unsupported in general
    #
    # The reason for this has to do with the delicate nature of constructing
    # the correlation data and making operators, like POVMs, physical etc... If
    # this becomes a dire need, we can always revisit.  For most applications
    # of interest, LSQ is perfectly good.
    #
    # The one exception to this is single qubit data where measurements form a
    # proper POVM.  In experimental language this is the case where we take six
    # data points for the state reconstruction
    if nbrAxes == 6 && nbrQubits == 1
        rhoML = QST_ML(tomoData,
                       measPulseMap,
                       measOpMap,
                       measPulseUs,
                       measOps)

        return rhoLSQ, rhoML
    else
        return rhoLSQ, []
    end
end

"""
    rho2pauli(ρ)

Convert a density matrix, ρ, to a Pauli set vector of Pauli expectation values.

# Arguments
- `ρ`: State tomography object constructed from the data

# Returns
- `paulivec:Array{Float64}`: array of Pauli expectation vaules
- `paulis:Array{Pauli}`: array of Pauli operators as defined in Cliffords.jl

# Examples
```julia-repl
julia> using Cliffords,
julia> foo = randn(ComplexF64, 2, 2)
2×2 Array{Complex{Float64},2}:
  0.310878-1.27868im    0.28776+1.87167im
 -0.904544+0.015083im  0.272544+0.620223im
julia> rho2pauli(foo)
([0.5834214596710264, -0.616783378678339, -1.8565830050748344,
 0.03833362063824197], Pauli{1}[+I, +X, +Y, +Z])
```
"""
function rho2pauli(ρ)
    n = round(Int, log2(size(ρ,1)))
    if n == 1 # special case for vectors
        paulis = sort(allpaulis(n), by=weight)
    else
        paulis = sort(allpaulis(n), dims=1, by=weight)
    end
    paulivec = [real(tr(ρ * p)) for p in paulis]
    return paulivec, paulis
end

## Process tomography is still a WIP

"""
Process tomography object

This holds all the infromation necessary to do process tomography.

The object is constructed by passing it a tomography dataset and its descriptor
loaded from load_data.  Once created, this object can be passed directly to the
any of the tomographyic reconstruction methods.
"""
struct ProcessTomo
    numQubits::Int
    numDatasets::Int
    corrData::Bool
    varData::Bool

    tomoDataSets::Dict{String,Dict{String,Array{Any,N} where N}}
    shotDataSets::Dict{String,Dict{String,Array{Any,N} where N}}

    numAxes::Int
    numCals::Int
    numCalRepeats::Int
    numDataPoints::Int

    mlPOVM::Array{Matrix{ComplexF64}}

    """
    Basic constructor
    """
    function ProcessTomo(data::Dict{String,Dict{String,Array{Any,N} where N}},
                        desc::Dict{String,Any})

        numQubits,
        numDatasets,
        corrData,
        varData,
        numDataPoints,
        tomoDataSets,
        shotDataSets = _pre_process_data(data, desc)

        ############################################################
        numCals, numCalRepeats, numAxes = _parese_exp_num(numDataPoints,
                                                          numQubits)
        ################################################################
        mlPOVM = _create_ml_POVM(numQubits)
        ################################################################
        new(numQubits, numDatasets, corrData, varData, tomoDataSets,
                                                       shotDataSets,
                                                       numAxes,
                                                       numCals,
                                                       numCalRepeats,
                                                       numDataPoints,
                                                       mlPOVM)
    end
end

"""
    analyzeProcessTomo(tomo::ProcessTomo)

Function to setup and run quantum process tomography for a given number
of qubits and measurement axes.

# Arguments
- `tomo::Qlab.ProcessTomo`: Process tomography object constructed from the data

# Returns
- `choiLSQ`: a 2^nbrQubits x 2^nbrQubits complex density matrix reconstructed
        with least-squares

# Examples
julia> datapath = "/path/to/data/folder"
julia> data, desc = load_data(datapath,24,"200315",load_var=true);
julia> tomo = Qlab.StateTomo(data, desc);
julia> rhoLSQ  = Qlab.analyzeStateTomo(tomo);
julia> Qlab.pauli_set_plot(rhoLSQ)
"""
function analyzeProcessTomo(data::Dict{String,Dict{String,Array{Any,N} where N}},
                            nbrQubits,
                            nbrPrepPulses,
                            nbrReadoutPulses,
                            nbrCalRepeats=2)

    measOps = Matrix{Float64}[]
    tomoData = Float64[]
    varData = Float64[]
    numMeas = length(data)
    numPreps = nbrPrepPulses^nbrQubits
    numExps = numPreps*numMeas*nbrReadoutPulses^nbrQubits

    for data_q in values(data)
        # Average over calibration repeats
        calData = real(data_q["data"][end-nbrCalRepeats*(2^nbrQubits )+1:end])
        avgCalData = mean(reshape(calData, nbrCalRepeats, 2^nbrQubits), dims=1)
        # Pull out the calibrations as diagonal measurement operators
        push!(measOps, diagm(avgCalData[:]))

        #The data to invert
        append!(tomoData, real(data_q["data"][1:end-nbrCalRepeats*(2^nbrQubits)]) )

        #variance
        append!(varData, real(data_q["variance"])[1:end-nbrCalRepeats*(2^nbrQubits)] )
    end

    #weightMat = 1./sqrt(varData)
    #weightMat/=sum(weightMat)
    # Map each experiment to the appropriate readout pulse
    measOpMap = repeat(1:numMeas, inner=nbrPulses^nbrQubits)
    measPulseMap = repeat(1:nbrPulses^nbrQubits, outer=numMeas)
    # Use a helper to get the measurement unitaries.
    measPulseUs = Qlab.tomo_gate_set(nbrQubits, nbrPulses)
    prepPulseUs = measPulseUs # for now, assume that preps and meas Us are the same

    # Now call the inversion routines
    # First least squares
    choiLSQ = QPT_LSQ(tomoData, varData, measPulseMap, measOpMap, prepPulseUs, measPulseUs, measOps, nbrQubits)

    # calculate the overlap with the ideal process

    return choiLSQ
end

"""
    QPT_LSQ(expResults, varMat, measPulseMap, measOpMap, measPulseUs, measOps, n)

Function to perform least-squares inversion of process tomography data.

   + expResults : data array
   + varmat : covariance matrix for data
   + measPulseMap : array mapping each experiment to a measurement readout
     pulse
   + measOpMap: array mapping each experiment to a measurement channel
   + prepPulseUs : array of unitaries of preparation pulses
   + measPulseUs : array of unitaries of measurement pulses
   + measOps : array of measurement operators for each channel
"""
function QPT_LSQ(expResults, varMat, measPulseMap, measOpMap, prepPulseUs, measPulseUs, measOps, nbrQubits)
    d = 2^nbrQubits
    # construct the vector of observables for each experiment
    obs = Matrix{}[]
    preps = Matrix{}[]
    for ct in 1:length(expResults)
        Uprep = prepPulseUs[mod1(ct, 16)]
        Umeas = measPulseUs[measPulseMap[mod1(ct, 48)]]
        rhoIn = zeros(d,d)
        rhoIn[1,1] = 1
        op = measOps[measOpMap[mod1(ct, 48)]]

        preps_ct = Uprep' * rhoIn * Uprep
        preps_ct = (preps_ct + preps_ct')/2 # force to be Hermitian
        push!(preps, preps_ct)
        #println(LinearAlgebra.tr(preps_ct))

        meas_ct = Umeas' * op * Umeas
        meas_ct = (meas_ct + meas_ct')/2 # force to be Hermitian
        push!(obs, meas_ct)
        # note the trace of these measurement operators will NOT be close to
        # 1 given the way the data is scaled.
    end
    # # in order to constrain the trace to unity, add an identity observable
    # # and a corresponding value to expResults
    # push!(obs, eye(Complex128, size(measOps[1])...))
    # expResults2 = [expResults; 1]
    # # corresponding variance chosen arbitrarily (it should be very small)
    # varMat2 = [varMat; minimum(varMat)]
    tomo = LSProcessTomo(obs, preps)

    choiLSQ, obj, status = fit(tomo, expResults, varMat)
    if status != :Optimal
        println("LSProcessTomo fit return status: $status")
    end
    return choiLSQ
end