text:
# PREAMBLE

# PKG_SETUP

# ## Setup

using ContactImplicitMPC# set environment variables to location of Gurobi software and license
ENV["GUROBI_HOME"] = "/gurobi912/linux64/"
ENV["GRB_LICENSE_FILE"] = "/gurobi912/gurobi.lic"

using Gurobi
using JuMP
using LinearAlgebra
using MeshCat
using Plots

include("/methods/structures.jl")
include("/methods/visuals.jl")

# ## Probem data
T = 40 # horizon (T+1 states, T controls)
H = 380 # simulation horizon
dt = 0.04 # step size
x0 = [0.00, 0.0] # initial state
Q = 1.0 # cost on state
Qf = 50.0 # cost on state
R = 1.0 # cost on control
β = 1e3 # extreme values in MIQP

# ## Dynamics model
model = WallPendulum(2, 1, 1.0, 1.0, 10.0, 1e4, 0.1) # m l g k d

A1, B1, c1 = dynamics_model(model, dt, mode = :none)
A2, B2, c2 = dynamics_model(model, dt, mode = :left)
A3, B3, c3 = dynamics_model(model, dt, mode = :right)

A = [A1, A2, A3]
B = [B1, B2, B3]
c = [c1, c2, c3]

# ## Domains
C1 = domain(model, mode = :none)
C2 = domain(model, mode = :left)
C3 = domain(model, mode = :right)
C = [C1, C2, C3]

# ## Problem
prob = WallProblem(model, T, x0, Q, Qf, R, β, A, B, c, C)

# ## Disturbances
i1 = 20
i2 = i1 + 100
i3 = i2 + 40
i4 = i3 + 100
i5 = i4 + 60
ind = [i1, i2, i3, i4, i5]
w = [[-15.5], [+15.5], [+15.5], [-11.5], [-10.5]]

# ## Simulation
x, u, s = simulate!(prob, x0, H, w = w, iw = ind)
sum(s) / (H * dt) # 4.52
maximum(s) / dt #17.6
mean(s) # 0.18
sqrt(mean((s .- mean(s)).^2)) #

# ##
DTu = range(0.0, step = dt, length = size(u)[1])
plot(DTu, hcat(u...)', label = "u")

# ##
DTx = range(0.0, step = dt, length = size(x)[1])
plot(DTx[1:end], hcat(x...)'[1:end,:], label = ["θ" "θd"])

# ## Visualizer
vis = Visualizer()
open(vis)

# ## Visualise
build_robot!(vis, model)
set_robot!(vis, model, [0.10, 0.1])
anim = visualize_robot!(vis, model, x, h = dt)
pθ_right = generate_pusher_traj(ind, w, x, side=:right)
pθ_left  = generate_pusher_traj(ind, w, x, side=:left)
visualize_disturbance!(vis, model, pθ_right, anim=anim, sample=1, offset=0.02, name=:PusherRight)
visualize_disturbance!(vis, model, pθ_left,  anim=anim, sample=1, offset=0.02, name=:PusherLeft)

ext: jl


text:
#Image
export loadImage, saveImage, rotate, scale, transform, skew, translate,
        regions, rgb2gray, gray2rgb, rgb2hsv, hsv2rgb, rgb2ycbcr, ycbcr2rgb

#Image Utils
"Load image as an AFArray"
function loadImage(path::AbstractString; color = false)
	return AFArray{Float32}(af_loadImage(path, color))
end

"Save AFArray as image"
function saveImage(path::AbstractString, a::AFAbstractArray)
    af_saveImage(path, a)
end

#Image Transformation
"Rotate image by theta radians"
function rotate{T}(a::AFAbstractArray{T}, theta)
	AFArray{T}(af_rotate(a,theta))
end

"Scale image in either dimension"
function scale{T}(a::AFAbstractArray{T}, scale1::Real, scale2::Real)
	AFArray{T}(af_scale(a,scale1,scale2))
end

"Skew an image in either dimension by angles"
function skew{T}(a::AFAbstractArray{T}, skew1::Real, skew2::Real)
	AFArray{T}(af_skew(a,b,c))
end

"Translate a given image in either dimension"
function translate{T}(a::AFAbstractArray{T}, trans1::Real, trans2::Real)
	AFArray{T}(af_translate(a,trans1, trans2))
end

function transform{T}(a::AFAbstractArray{T}, b::AFAbstractArray)
	AFArray{T}(af_transform(a,b))
end

#Image labelling
regions{T}(a::AFAbstractArray{T}) = AFArray{T}(af_regions(a))

#Create Gaussian Kernel
gaussiankernel(rows::Integer, cols::Integer) = AFArray{Float32}(af_gaussianKernel(rows,cols))

#Colorspace conversions
gray2rgb{T}(a::AFAbstractArray{T}) = AFArray{T}(af_gray2rgb(a))
hsv2rgb{T}(a::AFAbstractArray{T}) = AFArray{T}(af_hsv2rgb(a))
rgb2gray{T}(a::AFAbstractArray{T}) = AFArray{T}(af_rgb2gray(a))
rgb2hsv{T}(a::AFAbstractArray{T}) = AFArray{T}(af_rgb2hsv(a))
rgb2ycbcr{T}(a::AFAbstractArray{T}) = AFArray{T}(af_rgb2ycbcr(a))
ycbcr2rgb{T}(a::AFAbstractArray{T}) = AFArray{T}(af_ycbcr2rgb(a))
colorspace{T}(a::AFAbstractArray{T}, from::Cxx.CppEnum{:af_cspace_t}, to::Cxx.CppEnum{:af_cspace_t}) = af_colorspace(a, to, from)

#Filters
export SAT, bilateral, maxfilt, medfilt, minfilt, sobel, meanShift

"Summed Area Tables"
SAT{T}(a::AFAbstractArray{T}) = AFArray{T}(af_SAT(a))

"Apply bilateral filter on image"
function bilateral{T}(a::AFAbstractArray{T}, spatial_sigma::Real, chromatic_sigma::Real)
    AFArray{T}(af_bilateral(a, spatial_sigma, chromatic_sigma))
end

"For every pixel, find max value from a window"
maxfilt{T}(a::AFAbstractArray{T}) = AFArray{T}(af_maxfilt(a))

"For every pixel, find median value from a window"
medfilt{T}(a::AFAbstractArray{T}) = AFArray{T}(af_medfilt(a))

"For every pixel, find minimum value from a window"
minfilt{T}(a::AFAbstractArray{T}) = AFArray{T}(af_minfilt(a))

"Sobel operator on image"
sobel{T}(a::AFAbstractArray{T}) = AFArray{T}(af_sobel(a))

"Meanshift filter"
function meanShift{T}(a::AFAbstractArray{T}, spatial_sigma::Real, chromatic_sigma::Real, iter::Integer = 100)
    AFArray{T}(af_meanShift(a, spatial_sigma, chromatic_sigma, iter))
end

#Histogram
import Base: hist
export histEqual
hist{T}(a::AFAbstractArray{T}, nbins::Integer) = AFArray{T}(af_histogram(a, nbins))
function hist{T}(a::AFAbstractArray{T}, nbins::Integer, minval::Real, maxval::Real)
    AFArray{T}(af_histogram(a, nbins, minval, maxval))
end

"Histogram equalization of input image"
histEqual{T}(a::AFAbstractArray{T}, hist::AFAbstractArray{T}) = AFArray{T}(af_histEqual(a, hist))

#Morphological ops
export dilate, dilate3, erode, erode3

dilate{T}(a::AFAbstractArray{T}, mask::AFAbstractArray{T}) = AFArray{T}(icxx"af::dilate($a, $mask);")
dilate3{T}(a::AFAbstractArray{T}, mask::AFAbstractArray{T}) = AFArray{T}(icxx"af::dilate3($a, $mask);")
erode{T}(a::AFAbstractArray{T}, mask::AFAbstractArray) = AFArray{T}(icxx"af::erode($a, $mask);")
erode3{T}(a::AFAbstractArray{T}, mask::AFAbstractArray{T}) = AFArray{T}(icxx"af::erode3($a, $mask);")

#Computer Vision

#Constants
export matchTemplate, DiffOfGaussians

function matchTemplate{T}(searchImg::AFAbstractArray{T}, template::AFAbstractArray{T}, matchType = AF_SAD)
    AFArray{T}(af_matchTemplate(searchImg, template, matchType))
end

DiffOfGaussians{T}(img::AFAbstractArray{T}, radius1::Integer, radius2::Integer) = AFArray{T}(dog(img, radius1, radius2))

import Base:show
show(io::IO, f::AFFeatures) = show(io, f.feat)
show(io::IO, f::vcpp"af::features") = print(io, "ArrayFire Feature")

function getX(f::AFFeatures)
    out = af_getX(f)
    AFArray{backend_eltype(out)}(out)
end
function getY(f::AFFeatures)
    out = af_getY(f)
    AFArray{backend_eltype(out)}(out)
end
function getScore(f::AFFeatures)
    out = af_getScore(f)
    AFArray{backend_eltype(out)}(out)
end
function getSize(f::AFFeatures)
    out = af_getSize(f)
    AFArray{backend_eltype(out)}(out)
end
function getOrientation(f::AFFeatures)
    out = af_getOrientation(f)
    AFArray{backend_eltype(out)}(out)
end

function ORB(img::AFAbstractArray; fast_thr = 20.0, max_feat = 400, scl_fctr = 1.5, levels = 4, blur_img = false)
    out_features = AFFeatures()
    out_desc = AFArray()
    af_orb(out_features, out_desc, img, fast_thr, max_feat, scl_fctr, levels, blur_img)
    AFFeatures(out_features), AFArray{backend_eltype(out_desc)}(out_desc)
end

function sift(img::AFAbstractArray; n_layers = 3, constant_thr = 0.04,
                edge_thr = 10.0, init_sigma = 1.6, double_input = true,
                intensity_scale = 0.00390625, feature_ratio = 0.05)
    out_features = AFFeatures()
    out_desc = AFArray()
    af_sift(out_features, out_desc, img, n_layers, constant_thr,
            edge_thr, init_sigma, double_input, intensity_scale, feature_ratio)
    AFFeatures(out_features), AFArray{backend_eltype(out_desc)}(out_desc)
end

function fast(img::AFAbstractArray; thr = 20.0, arc_length = 9, non_max = true, feature_ratio = 0.05, edge = 3)
    AFFeatures(af_fast(img, thr, arc_length, non_max, feature_ratio, edge))
end

function harris(img::AFAbstractArray; max_corners = 500, min_response = 1e5, sigma = 1.0, block_size = 0, k_thr = 0.04)
    AFFeatures(af_harris(img, max_corners, min_response, sigma, block_size, k_thr))
end

function susan(img::AFAbstractArray; radius = 3, diff_thr = 32.0, geom_thr = 10.0, feature_ratio = 0.05, edge = 3)
    AFFeatures(af_susan(img, radius, diff_thr, geom_thr, feature_ratio, edge))
end

function hammingMatcher(query::AFAbstractArray, train::AFAbstractArray; dim = 1, n_dist = 1)
    out_idx = AFArray()
    out_dist = AFArray()
    af_hammingMatcher(out_idx, out_dist, query, train, dim-1, dist)
    AFArray{backend_eltype(out_idx)}(out_idx), AFArray{backend_eltype(out_dist)}(out_dist)
end

function nearnestNeighbour(query::AFAbstractArray, train::AFAbstractArray; dim = 1, n_dist = 1, dist_type = AF_SSD)
    out_idx = AFArray()
    out_dist = AFArray()
    af_nearestNeighbour(out_idx, out_dist, query, train, dim, n_dist, dist_type)
    AFArray{backend_eltype(out_idx)}(out_idx), AFArray{backend_eltype(out_dist)}(out_dist)
end


text:
export prime, sub, sdd_size, sdd_num_nodes, sdd_num_nodes_leafs, mgr,
       compress, unique⋁, canonicalize, negate,
       entails, equivalent, (≡)

#####################
# SDDmgr
#####################

@inline children(n::SddMgrInnerNode) = [n.left, n.right]

@inline variable(n::SddMgrLeafNode)::Var = n.var
@inline variables(n::SddMgrLeafNode)::BitSet = BitSet(n.var)
@inline variables(n::SddMgrInnerNode)::BitSet = n.variables

#####################
# Sdd
#####################

"Get the manager of a `Sdd` node, which is its `SddMgr` vtree"
mgr(s::Sdd) = s.vtree

@inline children(n::Sdd⋀Node) = [n.prime,n.sub]
@inline children(n::Sdd⋁Node) = n.children

@inline varsubset(n::Sdd, m::Sdd) = varsubset(mgr(n), mgr(m))
@inline varsubset_left(n::Sdd, m::Sdd)::Bool = varsubset_left(mgr(n), mgr(m))
@inline varsubset_right(n::Sdd, m::Sdd)::Bool = varsubset_right(mgr(n), mgr(m))

"Count the number of elements in the decision nodes of the SDD"
sdd_size(sdd) = mapreduce(n -> num_children(n), +, ⋁_nodes(sdd); init=0) # defined as the number of `elements`; length(⋀_nodes(sdd)) also works but undercounts in case the compiler decides to cache elements

"Count the number of decision nodes in the SDD"
sdd_num_nodes(sdd) = length(⋁_nodes(sdd)) # defined as the number of `decisions`

"Count the number of decision and leaf nodes in the SDD"
sdd_num_nodes_leafs(sdd) = sdd_num_nodes(sdd) + num_leafnodes(sdd) # defined in sdd file format

Base.show(io::IO, n::SddConstantNode) =
    print(io, (isfalse(n) ? "⊥" : "⊤"))
Base.show(io::IO, c::SddLiteralNode) = print(io, literal(c))
Base.show(io::IO, c::Sdd⋀Node) = begin
    recshow(c::Union{SddConstantNode,SddLiteralNode}) = "$c"
    recshow(c::Sdd⋁Node) = "D$(hash(c))"
    print(io, "($(recshow(prime(c))),$(recshow(sub(c))))")
end
Base.show(io::IO, c::Sdd⋁Node) = begin
    elems = ["$e" for e in children(c)]
    print(io, "[$(join(elems,','))]")
end

#############
# Compilation
#############

compile(::Type{<:Sdd}, mgr::SddMgr, arg::Bool) = compile(mgr, arg)
compile(::Type{<:Sdd}, mgr::SddMgr, arg::Lit) = compile(mgr, arg)
compile(::Type{<:Sdd}, mgr::SddMgr, arg::LogicCircuit) = compile(mgr, arg)

"Compile a circuit (e.g., CNF or DNF) into an SDD, bottom up by distributing circuit nodes over vtree nodes"

compile(mgr::SddMgr, c::LogicCircuit) = begin
    scopes = Dict{LogicCircuit, BitSet}()
    variables(c, scopes) # populate scopes dict
    compile(mgr, c, GateType(c), scopes)
end
compile(mgr::SddMgr, c::LogicCircuit, scopes) =
    compile(mgr, c, GateType(c), scopes)
compile(mgr::SddMgr, c::LogicCircuit, ::ConstantGate, _) =
    compile(mgr, constant(c))
compile(mgr::SddMgr, c::LogicCircuit, ::LiteralGate, _) =
    compile(mgr, literal(c))
compile(mgr::SddMgr, c::LogicCircuit, gt::InnerGate, scopes) =
    compile(mgr, NodeType(mgr), children(c), gt, scopes)
compile(mgr::SddMgr, children::Vector{<:LogicCircuit}, gt::InnerGate, scopes) =
    compile(mgr, NodeType(mgr), children, gt, scopes)

function compile(mgr::SddMgr, ::Leaf, children::Vector{<:LogicCircuit}, gt::InnerGate, scopes)
    isempty(children) && return compile(mgr, neutral(gt))
    mapreduce(x -> compile(mgr,x,scopes), op(gt), children)
end

function compile(mgr::SddMgr, ::Inner, children::Vector{<:LogicCircuit}, gt::InnerGate, scopes)
    isempty(children) && return compile(mgr, neutral(gt))

    # partition children according to vtree
    left_children = filter(x -> subseteq_fast(scopes[x], variables(mgr.left)), children)
    right_children = filter(x -> subseteq_fast(scopes[x], variables(mgr.right)), children)
    middle_children = setdiff(children, left_children, right_children)

    # separately compile left and right vtree children
    left = compile(mgr.left, left_children, gt, scopes)
    right = compile(mgr.right, right_children, gt, scopes)

    mapreduce(x -> compile(mgr,x,scopes), op(gt), middle_children; init=op(gt)(left, right))
end

"""
Compile a given variable, literal, or constant
"""
function compile(n::SddMgrLeafNode, l::Lit)::SddLiteralNode
    @assert n.var == lit2var(l) "Cannot compile literal $l respecting vtree leaf for variable $(n.var)"
    l>0 ? n.positive_literal::SddLiteralNode : n.negative_literal::SddLiteralNode
end

function compile(n::SddMgrInnerNode, l::Lit)
    if lit2var(l) ∈ variables(n.left)
        compile(n.left, l)::SddLiteralNode
    else
        @assert lit2var(l) ∈ variables(n.right) "$l is not contained in this vtree $n with scope $(variables(n))"
        compile(n.right, l)::SddLiteralNode
    end
end

compile(::SddMgr, constant::Bool) =
    constant ? true_sdd : false_sdd

"""
Negate an SDD
"""
@inline negate(c::SddConstantNode) =
    (c === false_sdd) ? true_sdd : false_sdd

function negate(s::SddLiteralNode)
    if ispositive(s)
        mgr(s).negative_literal::SddLiteralNode
    else
        mgr(s).positive_literal::SddLiteralNode
    end
end

negate(node::Sdd⋁Node) = node.negation::Sdd⋁Node

@inline Base.:!(c::LogicCircuit) = negate(c)

"""
Get the canonical compilation of the given XY Partition
"""
function canonicalize(xy::XYPartition, mgr::SddMgrInnerNode)
    # @assert !isempty(xy)
    return canonicalize_compressed(compress(xy), mgr)::Sdd
end

"""
Compress a given XY Partition (merge elements with identical subs)
"""
function compress(xy::XYPartition)
    compressed = true
    for i in eachindex(xy), j in i+1:length(xy)
        if (sub(xy[i]) === sub(xy[j]))
            compressed = false
            break
        end
    end
    compressed && return xy
    # make it compressed
    out = XYPartition()
    sizehint!(out, length(xy))
    mask = falses(length(xy))
    for i in eachindex(xy)
        if !mask[i]
            prime_all = prime(xy[i])
            sub_i = sub(xy[i])
            for j in i+1:length(xy)
                sub_j = sub(xy[j])
                if !mask[j] && (sub_i === sub_j)
                    prime_all = prime_all | prime(xy[j])
                    mask[j] = true
                end
            end
            push!(out,Element(prime_all,sub_i))
        end
    end
    return out::XYPartition
end

"""
Get the canonical compilation of the given compressed XY Partition
"""
function canonicalize_compressed(xy::XYPartition, mgr::SddMgrInnerNode)
    # @assert !isempty(xy)
    # trim
    if length(xy) == 1 && (prime(first(xy)) === true_sdd)
        return sub(first(xy))
    elseif length(xy) == 2
        if (sub(xy[1]) === true_sdd) && (sub(xy[2]) === false_sdd)
            return prime(xy[1])
        elseif (sub(xy[2]) === true_sdd) && (sub(xy[1]) === false_sdd)
            return prime(xy[2])
        end
    end
    # get unique node representation
    return unique⋁(xy, mgr)::Sdd⋁Node
end

"""
Construct a unique decision gate for the given vtree
"""
function unique⋁(xy::XYPartition, mgr::SddMgrInnerNode)
    #TODO add finalization trigger to remove from the cache when the node is gc'ed + weak value reference
    get!(mgr.unique⋁cache, xy) do
        xynodes = [Sdd⋀Node(prime(e), sub(e), mgr) for e in xy]
        node = Sdd⋁Node(xynodes, mgr)
        # some memory allocations can be saved here, by not allocating the intermediate vector of elements
        # however, that does not appear to speed things up...
        not_xy = [Element(prime(e), !sub(e)) for e in xy]
        not_xynodes = [Sdd⋀Node(prime(e), sub(e), mgr) for e in not_xy]
        not_node = Sdd⋁Node(not_xynodes, mgr, node)
        node.negation = not_node
        mgr.unique⋁cache[not_xy] = not_node
        node
    end::Sdd⋁Node
end

# specialize mode count to extract number of variables from vtree global scope
model_count(root::StructLogicCircuit)::BigInt =
    model_count(root, length(global_scope(vtree(root))))

"Decide whether one sentence logically entails another"
entails(x::Sdd, y::Sdd) = ((x & !y) === false_sdd)

"Decide whether two sentences are logically equivalent"
equivalent(x::Sdd, y::Sdd) = begin
    # note that ≡ has the wrong operator prescedence for use in logic
    @assert root(vtree(x)) === root(vtree(y))
    return (x === y)
end

import Base: xor

"Exclusive logical disjunction (XOR)"
xor(x::Sdd,y::Sdd) = (x ∧ ¬y) ∨ (y ∧ ¬x)

"Smooth an sdd to a StructLogicCircuit"
function smooth(root::Sdd)::StructLogicCircuit
    lc = LogicCircuit(root)
    plc = propagate_constants(lc, remove_unary=true)
    slc = StructLogicCircuit(root.vtree, plc)
    smooth(slc)
end

max_issues_repo_name: Juice-jl/LogicCircuits

max_forks_repo_name: Juice-jl/LogicCircuits


text:
module TMGlib
# wrap some routines from the LAPACK test matrix generator library
export latme!, latmr!

using LinearAlgebra
import LinearAlgebra.BlasInt
import LinearAlgebra.BLAS.@blasfunc

# CHECKME: routines linked here are included in many BLAS/LAPACK distributions
# (OpenBLAS, MKL), but do we need to allow for others?

const liblapack = Base.liblapack_name

const latm_modes = Dict(:specified => 0, :clustered_small => 1,
                    :clustered_large => 2, :exponential => 3,
                    :arithmetic => 4, :random_log => 5, :random => 6)

function _check_iseed(iseed::Vector{T}) where {T<:Integer}
    if length(iseed) < 4
        throw(ArgumentError("iseed must have (at least) 4 elements"))
    end
    if any(iseed .< 0) || any(iseed .> 4095)
        throw(ArgumentError("iseed entries must be in 0:4095"))
    end
    isodd(iseed[4]) || throw(ArgumentError("iseed[4] must be odd"))
    nothing
end

"""
construct a random non-symmetric square matrix with specified eigenvalues

`latme!(A::Matrix{<:BlasFloat}, anorm::Real, imode::Int, rcond::Real,
simrcond::Real = 1.0;
iseed::Vector{BlasInt} = [2518, 3899, 995, 397],
kl::Int = size(A,2), ku::Int = size(A,2),
dist::AbstractChar = 'D', rsign::Bool = true,
upper::Bool = true, simtrans::Bool = true,
simmode::Int = 4, d::Vector, ds::Vector)`

- `imode`: 1 - clustered small, 2 - clustered large, 3 - exponential,
     4 - arithmetic, 5 - random log, 6 - random (<0 for reversed order),
     0 - eigenvalues explicitly specified in `d`.
- `rcond`: ratio of extreme eigenvalue magnitudes.
- `anorm`: if positive, finally scale the matrix so its maximum-element-norm is `anorm`.
- `kl,ku`: lower and upper bandwidths
- If `upper`, fill the upper triangle (make result non-normal).
- If `rsign`, eigenvalues have random signs (real) or angles (complex).
- If `simtrans`, apply a similarity transform with singular values
      specified by `simmode, simrcond, ds`, similarly to `imode, rcond, d`.

- `dist`: character key for distribution to be sampled: `'U'` - uniform(0,1),
   `'S'` - symmetric uniform(-1,1), `'N'` - normal, `'D'` - uniform on complex unit disk.
"""
function latme! end

for (fname, elty, relty) in ((:zlatme_, :ComplexF64, :Float64),
                             (:clatme_, :ComplexF32, :Float32))
    @eval begin
        function latme!(A::Matrix{$elty}, anorm::Real, imode::Int, rcond::Real,
                simrcond::Real = 1.0;
                iseed::Vector{BlasInt} = BlasInt[2518, 3899, 995, 397],
                kl::Int = size(A,2), ku::Int = size(A,2),
                dist::AbstractChar = 'D', rsign::Bool = true,
                upper::Bool = true, simtrans::Bool = true,
                simmode::Int = 4,
                dmax::$elty = one($elty),
                d::Vector{$elty} = Vector{$elty}(undef,size(A,2)),
                ds::Vector{$relty} = Vector{$relty}(undef,size(A,2))
                )
            n = size(A,2)
            lda = size(A,1)

            if imode == 0
                (length(d) == n) || throw(ArgumentError("imode = 0 needs explicit d of length n"))
                if simtrans
                    (length(ds) == n) || throw(ArgumentError("imode = 0 needs explicit ds of length n"))
                end
            end
            (dist ∈ ['D','U','S','N']) || throw(ArgumentError("allowed values of dist are 'D','U','S', and 'N'"))

            _check_iseed(iseed)
            work = Vector{$elty}(undef, 3*n)
            info = Ref{BlasInt}(0)

            ccall((@blasfunc($fname), liblapack), Cvoid,
                  (Ref{BlasInt}, # n
                   Ref{UInt8}, # dist
                   Ptr{BlasInt}, # iseed
                   Ptr{$elty}, # d(w)
                   Ref{BlasInt}, # mode
                   Ref{$relty}, # cond
                   Ref{$elty}, # dmax
                   Ref{UInt8}, # rsign
                   Ref{UInt8}, # upper
                   Ref{UInt8}, # sim
                   Ptr{$relty}, # ds(w)
                   Ref{BlasInt}, # modes
                   Ref{$relty}, # conds
                   Ref{BlasInt}, # kl
                   Ref{BlasInt}, # ku
                   Ref{$relty}, # anorm
                   Ptr{$elty}, # A
                   Ref{BlasInt}, # lda
                   Ptr{$elty}, # work(w)
                   Ptr{BlasInt}, # info(w)
                   Clong, Clong, Clong, Clong # char lengths
                   ),
                  n, dist, iseed, d, imode, rcond, dmax,
                  rsign ? 'T' : 'F', upper ? 'T' : 'F', simtrans ? 'T' : 'F',
                  ds, simmode, simrcond, kl, ku, anorm, A, lda, work, info, 1,1,1,1)
            (info[] == 0) || error("info = $(info[]) from latme!")
            A, d, ds, iseed
        end
    end
end

# Note LAPACK routine has different arg list for the real versions:
for (fname, celty, elty) in ((:dlatme_, :ComplexF64, :Float64),
                             (:slatme_, :ComplexF32, :Float32))
    @eval begin
        function latme!(A::Matrix{$elty}, anorm::Real, imode::Int, rcond::Real,
                simrcond::Real = 1.0;
                iseed::Vector{BlasInt} = BlasInt[2518, 3899, 995, 397],
                kl::Int = size(A,2), ku::Int = size(A,2),
                dist::AbstractChar = 'S', rsign::Bool = true,
                upper::Bool = true, simtrans::Bool = true,
                simmode::Int = 4,
                dmax::$elty = one($elty),
                d = nothing,
                ds::Vector{$elty} = Vector{$elty}(undef,size(A,2))
                )
            n = size(A,2)
            lda = size(A,1)

            d1 = Vector{$elty}(undef,size(A,2))

            if imode == 0
                (d isa Vector) || throw(ArgumentError("imode = 0 needs vector d"))
                (length(d) == n) || throw(ArgumentError("imode = 0 needs explicit d of length n"))
                if simtrans
                    (length(ds) == n) || throw(ArgumentError("imode = 0 needs explicit ds of length n"))
                end
            end
            (dist ∈ ['U','S','N']) || throw(ArgumentError("allowed values of dist are 'U','S', and 'N'"))

            _check_iseed(iseed)
            work = Vector{$elty}(undef, 3*n)
            info = Ref{BlasInt}(0)

            # if specifying eigenvalues
#           EI is CHARACTER*1 array, dimension ( N )
#            If MODE is 0, and EI(1) is not ' ' (space character),
#            this array specifies which elements of D (on input) are
#            real eigenvalues and which are the real and imaginary parts
#            of a complex conjugate pair of eigenvalues.  The elements
#            of EI may then only have the values 'R' and 'I'.  If
#            EI(j)='R' and EI(j+1)='I', then the j-th eigenvalue is
#            CMPLX( D(j) , D(j+1) ), and the (j+1)-th is the complex
#            conjugate thereof.  If EI(j)=EI(j+1)='R', then the j-th
#            eigenvalue is D(j) (i.e., real).  EI(1) may not be 'I',
#            nor may two adjacent elements of EI both have the value 'I'.
#            If MODE is not 0, then EI is ignored.  If MODE is 0 and
#            EI(1)=' ', then the eigenvalues will all be real.

            if imode == 0
                if eltype(d) <: Real
                    ei = UInt8[' ']
                    d1 .= d[1:n]
                else
                    ei = Vector{UInt8}(undef, n)
                    ie = 1
                    id = 1
                    while id <= n
                        ei[ie] = 'R'
                        d1[ie] = real(d[id])
                        ie += 1
                        if imag(d[id]) != 0
                            if d[id+1] != conj(d[id])
                                throw(ArgumentError("conjugate pairs in d must be adjacent"))
                            end
                            ei[ie] = 'I'
                            d1[ie] = imag(d[id])
                            ie += 1
                            id += 1
                        end
                        id += 1
                    end
                end
            else
                ei = UInt8[' ']
            end

            ccall((@blasfunc($fname), liblapack), Cvoid,
                  (Ref{BlasInt}, # n
                   Ref{UInt8}, # dist
                   Ptr{BlasInt}, # iseed
                   Ptr{$elty}, # d(w)
                   Ref{BlasInt}, # mode
                   Ref{$elty}, # cond
                   Ref{$elty}, # dmax
                   Ref{UInt8}, # ei
                   Ref{UInt8}, # rsign
                   Ref{UInt8}, # upper
                   Ref{UInt8}, # sim
                   Ptr{$elty}, # ds(w)
                   Ref{BlasInt}, # modes
                   Ref{$elty}, # conds
                   Ref{BlasInt}, # kl
                   Ref{BlasInt}, # ku
                   Ref{$elty}, # anorm
                   Ptr{$elty}, # A
                   Ref{BlasInt}, # lda
                   Ptr{$elty}, # work(w)
                   Ptr{BlasInt}, # info(w)
                   Clong, Clong, Clong, Clong, Clong # char lengths
                   ),
                  n, dist, iseed, d1, imode, rcond, dmax, ei,
                  rsign ? 'T' : 'F', upper ? 'T' : 'F', simtrans ? 'T' : 'F',
                  ds, simmode, simrcond, kl, ku, anorm, A, lda, work, info,
                  1,1,1,1,1)
            (info[] == 0) || error("info = $(info[]) from latme!")
            A, d1, ds, iseed
        end
    end
end

"""
construct a random matrix of various types

latmr!(A::Matrix{T<:BLAStype}, anorm::Real, imode::Int, rcond::Real,
                rcondr::Real = 1.0, rcondl::Real = 1.0;
                sym::AbstractChar = 'N',
                iseed::Vector{BlasInt} = [2518, 3899, 995, 397],
                kl::Int = size(A,2), ku::Int = size(A,2),
                dist::AbstractChar = 'D', rsign::Bool = true,
                grade::AbstractChar = 'N', model::Int = 4, moder::Int = 4,
                pivtng::AbstractChar = 'N', ipivot::Vector{BlasInt} = BlasInt[],
                sparse::Real = 0.0,
                pack::AbstractChar = 'N'
                )

- `imode`: 1 - clustered small, 2 - clustered large, 3 - exponential,
       4 - arithmetic, 5 - random log, 6 - random (<0 for reversed order)
- `imode = 0`, explicitly specified diagonal, is not yet implemented.
- `sym`: `'S'` - symmetric, `'H'` - Hermitian, `'N'` - nonsymmetric
- `rcond`: ratio of extreme diagonal elements
- `kl,ku`: lower and upper bandwidths
- `sparse`: fraction to be nulled out
- `dist`: character key for distribution to be sampled: `'U'` - uniform(0,1),
   `'S'` - symmetric uniform(-1,1), `'N'` - normal, `'D'` - uniform on complex unit disk.
"""
function latmr! end

for (fname, elty, relty) in ((:zlatmr_, :ComplexF64, :Float64),
                             (:clatmr_, :ComplexF32, :Float32),
                             (:dlatmr_, :Float64, :Float64),
                             (:slatmr_, :Float32, :Float32))
    @eval begin
        function latmr!(A::Matrix{$elty}, anorm::Real, imode::Int, rcond::Real,
                        rcondr::Real = 1.0, rcondl::Real = 1.0;
                        sym::AbstractChar = 'N',
                        iseed::Vector{BlasInt} = BlasInt[2518, 3899, 995, 397],
                        kl::Int = size(A,2), ku::Int = size(A,2),
                        dist::AbstractChar = $elty <: Complex ? 'D' : 'S',
                        rsign::Bool = true,
                        grade::AbstractChar = 'N', model::Int = 1, moder::Int = 1,
                        pivtng::AbstractChar = 'N',
                        ipivot::Vector{BlasInt} = Vector{BlasInt}(undef,size(A,2)),
                        sparse::Real = 0.0,
                        pack::AbstractChar = 'N',
                        dmax = one($elty),
                        d = Vector{$elty}(undef, size(A,2)),
                        dl = Vector{$relty}(undef, size(A,2)),
                        dr = Vector{$relty}(undef, size(A,2))
                        )
            n = size(A,2)
            lda = size(A,1)
            m = lda # CHECKME: is it worth generalizing?

            (imode == 0) && throw(ArgumentError("imode = 0 needs explicit d,ds; not yet implemented"))
            _check_iseed(iseed)

            iwork = Vector{BlasInt}(undef, max(m,n))
            info = Ref{BlasInt}(0)

#=
            println("zlatmr arg list:")
            println(
                  " m $m, n $n, dist $dist, iseed $iseed, sym $sym, d $d, imode $imode, rcond $rcond, dmax $dmax,
                  rsign $(rsign ? 'T' : 'F'), grade $grade, dl $dl, model $model, rcondl $rcondl,
                  dr $dr, moder $moder, rcondr $rcondr, pivtng $pivtng, ipivot $ipivot,
                  kl $kl, ku $ku, sparse $sparse, anorm $anorm, pack $pack, A, lda $lda")
=#
            ccall((@blasfunc($fname), liblapack), Cvoid,
                  (Ref{BlasInt}, # m
                   Ref{BlasInt}, # n
                   Ref{UInt8}, # dist
                   Ptr{BlasInt}, # iseed
                   Ref{UInt8}, # sym
                   Ptr{$elty}, # d(w)
                   Ref{BlasInt}, # mode
                   Ref{$relty}, # cond
                   Ref{$elty}, # dmax
                   Ref{UInt8}, # rsign
                   Ref{UInt8}, # grade
                   Ptr{$relty}, # dl(w)
                   Ref{BlasInt}, # model
                   Ref{$relty}, # condl
                   Ptr{$relty}, # dr(w)
                   Ref{BlasInt}, # moder
                   Ref{$relty}, # condr
                   Ref{UInt8}, # pivtng
                   Ptr{BlasInt}, # ipivot
                   Ref{BlasInt}, # kl
                   Ref{BlasInt}, # ku
                   Ref{$relty}, # sparse
                   Ref{$relty}, # anorm
                   Ref{UInt8}, # pack
                   Ptr{$elty}, # A
                   Ref{BlasInt}, # lda
                   Ptr{BlasInt}, # iwork(w)
                   Ptr{BlasInt}, # info(w)
                   Clong, Clong, Clong, Clong, Clong, Clong
                   ),
                  m, n, dist, iseed, sym, d, imode, rcond, dmax,
                  rsign ? 'T' : 'F', grade, dl, model, rcondl,
                  dr, moder, rcondr, pivtng, ipivot,
                  kl, ku, sparse, anorm, pack, A, lda, iwork, info,
                  1,1,1,1,1,1)
            (info[] == 0) || error("info = $(info[]) from latmr")
            A, d, dl, dr, iseed
        end
    end
end

end # module


text:
#!/usr/bin/env julia
include("se_dmon.jl")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\usepackage{fontspec}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\setmainfont{Asana-Math}")
flumonths = ["oktober", "november", "december", "januari",
   "februari", "mars", "april"]
pgfsave("../../images/se_dmon_20_19_kv.svg",
   plotmonths("kvinnor", flumonths, 1920, 2019))
pgfsave("../../images/se_dmon_20_19_m.svg",
   plotmonths("män", flumonths, 1920, 2019))

max_stars_repo_name: klpn/static-dust

max_issues_repo_name: klpn/static-dust

max_forks_repo_name: klpn/static-dust


text:
export Register, report
import Base: length

"""
`Register(n::Int=1)` creates a new quantum register holding `n`
qubits, all set to `Q0`.

`Register(q1,q2,...)` or `Register([q1,q2,...])` creates a new quantum
register holding the qubits `q1`, `q2`, etc.

`Register(b1,b2,...)` where `b1,b2,...` are 0s and 1s creates a new register
holding qubits `Qubit(b1)`, `Qubit(b2)`, etc.
"""
mutable struct Register <: QuantumState
    vec::Vector{CC}
    function Register(nbits::Integer = 1)
        if nbits < 0
            throw(DomainError(nbits, "Number of bits must be positive."))
        end
        if nbits == 0
            nbits = 1
        end
        v = zeros(CC, 1 << nbits)
        v[1] = one(CC)
        new(v)
    end

    # This constructor generally should not be used the user.
    function Register(v::Vector{T}) where {T<:Number}
        new(normalize(v))
    end
end

function Register(bits::Integer...)
    qlist = Qubit.(bits)
    return Register(qlist...)
end

function Register(Qlist::Qubit...)
    nbits = length(Qlist)
    v = Qlist[1].vec
    for j = 2:nbits
        v = v ⊗ Qlist[j].vec
    end
    return Register(v)
end

function Register(QQ::Vector{Qubit})
    Qtuple = Tuple(q for q in QQ)
    return Register(Qtuple...)
end

"""
`length(R::Register)` is the number of qubits in `R`.
(The number of complex numbers in its vector is `2^n` where `n`
is the number of qubits.)
"""
length(R::Register)::Int = lg(length(R.vec))

"""
`measure!(R::Register)` returns a vector of 0s and 1s from measuring
the qubits in `R`. Then `R` is overwritten with `Q0`s and `Q1`s
respectively.
"""
function measure!(R::Register)::Vector{Int}
    pv = prob_vector(R)
    k = random_choice(pv)
    bits = reverse(digits(k-1,base=2, pad=length(R)))
    R.vec = Register(Tuple(bits)...).vec
    return bits
end

function binary_string(n::Int, bits::Int)
    dlist = reverse(digits(n,base=2,pad=bits))
    dd = Tuple(string.(dlist))
    return *(dd...)
end

"""
`report(R::Register)` prints a table of the possible register states
and their probabilities.
"""
function report(R::QuantumState)
    pvec = prob_vector(R)
    digs = length(R)
    for k=1:1<<digs
        println(binary_string(k-1,digs),"\t", pvec[k])
    end
    nothing
end


text:
using Base.Meta

# http://docs.julialang.org/en/release-0.5/manual/metaprogramming/

# expression structure
ex = :(1+1)
dump(ex)
show_sexpr(ex)

# rewriting
macro noplus(ex)
    @assert isexpr(ex, :call)
    if ex.args[1] == :+
        ex.args[1] = :-
    end
    return ex
end

@noplus 1+1

# recursively
function plustominus(ex::Expr)
    ex = copy(ex) # Don't modify the input
    if ex.head == :call && ex.args[1] == :+
        ex.args[1] = :-
    end
    for i in 1:length(ex.args)
        ex.args[i] = plustominus(ex.args[i])
    end
    return ex
end
plustominus(ex) = ex # If not expression, do nothing

plustominus(:(1+1))
plustominus(:(1+(1+1)))

macro noplus_recursive(ex)
    return plustominus(ex)
end

@noplus_recursive 2+(1+2)
macroexpand(:(@noplus_recursive 2+(1+2)))

function foo()
    x = 10
    @noplus_recursive x+5
end

@code_lowered foo()


text:
"""

    cm97(setting; maxiter = 1000)

Perform the Chatterjee and Mächler (1997) algorithm for the given regression setting.

# Arguments
- `setting::RegressionSetting`: RegressionSetting object with a formula and dataset.

# Examples
```julia-repl
julia> myreg = createRegressionSetting(@formula(stackloss ~ airflow + watertemp + acidcond), stackloss)
julia> result = cm97(myreg)
Dict{String,Any} with 3 entries:
  "betas"      => [-37.0007, 0.839285, 0.632333, -0.113208]
  "iterations" => 22
  "converged"  => true
```

# References
Chatterjee, Samprit, and Martin Mächler. "Robust regression: A weighted least squares approach."
Communications in Statistics-Theory and Methods 26.6 (1997): 1381-1394.
"""
function cm97(setting::RegressionSetting; maxiter::Int=1000)

    X, y = @extractRegressionSetting setting
    return cm97(X, y)

end

function cm97(X::Array{Float64,2}, y::Array{Float64,1}; maxiter::Int=1000)

    n, p = size(X)
    pbar::Float64 = p / n
    hat = hatmatrix(X)

    w_is  = Array{Float64}(undef, n)
    betas = Array{Float64}(undef, p)

    converged::Bool = false
    iter::Int = 0

    # initial weights
    @inbounds for i in 1:n
        w_is[i] = 1.0 / max(hat[i, i], pbar)
    end

    while iter <= maxiter
        oldbetas = betas

        wols = wls(X, y, w_is)
        betas = coef(wols)
        r = y - X * betas
        medi = median(abs.(r))

        @inbounds for i in 1:n
            w_is[i] = (1.0 - hat[i, i])^2.0 / max(abs(r[i]), medi)
        end

        if isapprox(oldbetas, betas)
            converged = true
            break
        end

        iter += 1

    end

    result = Dict{String,Any}(
        "converged" => converged,
        "betas" => betas,
        "iterations" => iter
    )

    return result
end

max_stars_repo_name: angeris/LinRegOutliers

max_issues_repo_name: angeris/LinRegOutliers

max_forks_repo_name: angeris/LinRegOutliers


text:
"""
    tfisinghamiltonian(h)
return the DTensor-representation of the two-body part of the
transverse-field ising Hamiltonian in two dimensions, i.e.

    H = -σx⊗σx - h(𝟙⊗σz + σz⊗𝟙)

Note that depending on the lattice geometry,
`h` needs to be scaled appropriately (e.g. on the square, `h → h/4`).

"""
function tfisinghamiltonian(h, T = ComplexF64)
    σx = T[0 1; 1  0]
    σz = T[1 0; 0 -1]
    id = T[1 0; 0  1]
    @tensor H[1,2,3,4] := -σx[1,3] * σx[2,4] -
                           h * id[1,3] * σz[2,4] -
                           h * id[2,4] * σz[1,3]
    return DTensor(H)
end

"""
    tfisingpropagator(h)
return the rank-4 propagator for a two-site ising-hamlitonian with
transverse field `h`, see `tfisinghamiltonian`.
"""
function tfisingpropagator(β,h)
    H = tfisinghamiltonian(h)
    Hmat, rs = fuselegs(H, ((1,2),(3,4)))
    apply!(Hmat, x -> x .= exp(-β .* x))
    U = splitlegs(Hmat, ((1,1,1),(1,1,2),(2,2,1),(2,2,2)), rs...)
    return U
end


text:
using TableTraits
using IteratorInterfaceExtensions
using Test

@testset "TableTraits" begin

    table_array = [(a = 1,), (a = 2,)]
    any_table_array = Any[(a = 1,), (a = 2,)]
    other_array = [1,2,3]
    without_eltype = (i for i in table_array)

    @test isiterabletable(table_array)
    @test !isiterabletable(other_array)
    @test isiterabletable(without_eltype) === missing
    @test isiterabletable(any_table_array) === missing
    @test !supports_get_columns_copy(table_array)
    @test !supports_get_columns_view(table_array)
    @test !supports_get_columns_copy_using_missing(table_array)

end


text:
# UnitCommitment.jl: Optimization Package for Security-Constrained Unit Commitment
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

"""
Formulation described in:

    Knueven, B., Ostrowski, J., & Watson, J. P. (2018). Exploiting identical
    generators in unit commitment. IEEE Transactions on Power Systems, 33(4),
    4496-4507. DOI: https://doi.org/10.1109/TPWRS.2017.2783850
"""
module KnuOstWat2018

import ..PiecewiseLinearCostsFormulation

struct PwlCosts <: PiecewiseLinearCostsFormulation end

end


text:
push!(LOAD_PATH, string(pwd(),"/src"))

using Compat.LinearAlgebra
using Compat.Test
using Compat.Random

using ModifiedStokesSolver

@testset "Kernels" begin
    srand(1)
    r = rand(2)
    f = rand(2)
    n = rand(2)
    alpha = rand()

    # Stokeslet functions are same
    u1 = stokeslet(r, f, alpha)
    K11, K12, K21, K22 = ModifiedStokesSolver.stokeslet_kernel_noarrays(r[1], r[2], alpha)
    K = [K11 K12; K21 K22]
    u2 = K*f
    @test norm(u1-u2, Inf) / norm(u1, Inf) < 1e-15

    # Stresslet functions are same
    u1 = stresslet(r, f, n, alpha)
    u2 = dblkernel(r, n, alpha) * f
    @test norm(u1-u2, Inf) / norm(u1, Inf) < 1e-15

    # Test log split
    (uS, uL, uC, uQ) = stresslet_split(r, f, n, alpha)
    KL = dblkernel_log(r, n, alpha)
    uL2 = KL * f
    @test norm(uL-uL2, Inf) / norm(uL, Inf) < 1e-15

    @testset "Ti derivatives" begin
        # Finite difference test
        h = 1e-8
        z = 2.0
        f = ModifiedStokesSolver.Ti_direct
        dFD = (f(z+h).-f(z-h))./(2*h)
        d = ModifiedStokesSolver.Ti_direct_prime(z)
        @test maximum(abs.((dFD .- d)./d)) < 1e-7
    end

end


text:
# # Ordinary random walk

# Here is the prototypical application. We repeatedly perform an ordinary random walk with a fixed
# number of steps on $\mathbb{Z}$, collecting statistics on the final position.
# We call each walk, or trajectory, a *sample*. An ensemble of samples is a *trial*.

# We begin by loading packages that we will use.

using RandomWalks
using EmpiricalCDFs: EmpiricalCDF, get_data
using Statistics

# First, we show an uncommented version of the code that specifies and runs the experiment.
# Note that it is very compact!

trial_loop = SampleLoopActor(10^6, ECDFActor(get_x, EmpiricalCDF()));
trial!(WalkPlan(WalkF(), StepLimitActor(10^3)), trial_loop);

# In the following we go through the same code step-by-step.

# ### Generating the samples and the trial

# A sample is a single trajectory, or walk. We will take `nsamples` samples.
nsamples = 10^6;

# Choose the number of steps per walk.
nsteps = 10^3;

# Create an emtpy empirical cumulative distribution function for floating point data.
ecdf = EmpiricalCDF()

# A trial is a collection of samples (walks) from which we collect statistics.
# The following instance of a `SampleLoopActor` specifies that we collect `nsamples` samples and that we
# build an empirical CDF of the final position (on the x-axis) of the walker after each walk. Below, we will
# use a one-dimensional walk, so `get_x` returns the position.
trial_loop = SampleLoopActor(nsamples, ECDFActor(get_x, ecdf));

# We used two `AbstractActors` above. `Actors` are treated (for the most part)
# in a uniform way within a sample, and in (a different)
# uniform way within a trial. This gives flexibility in designing experiments.

# We create a `WalkPlan` and run a trial of `nsamples` samples.
# The `WalkPlan` consists of  two parts;
# The first, a walk specification `WalkF()`, which is the default one-dimension walk on the integers.
# The second an action `StepLimitActor(nsteps)` to take after each step.
# `StepLimitActor(nsteps)` stops the walk if the number of steps exceeds `nsteps`.
# The `WalkPlan` specifies the sample and how it is collected.

# The `Actor` `trial_loop` was described above. It describes how to take several samples and actions to take
# afer each sample.
# In this case, which statistics to collect.
# So, the following line builds an ECDF of the final position of a one-dimensional walk after `nsteps` steps.
@time trial!(WalkPlan(WalkF(), StepLimitActor(nsteps)), trial_loop);

# We rescale the data so that the expected standard deviation is equal to `1`.
x = get_data(ecdf);
x ./= sqrt(nsteps);

# Test that the standard deviation and mean are as expected

## The mean should be near 0 and the standard deviation near 1
@show mean(ecdf), std(ecdf);

# We test the above results with a generous tolerance.
tolerance = 1e-3;

isapprox(std(ecdf), 1; atol = tolerance)
#-
isapprox(mean(ecdf), 0; atol = tolerance)

# We approximate the probability density function with a normalized `Histogram`
# and see that it looks reasonable.
using StatsBase: Histogram, fit, normalize
h = normalize(fit(Histogram, get_data(ecdf), -5.5:1.0:5.5));

println(h.weights);


text:
function interface_path(library::AbstractString)
    return abspath(INTERFACE_PATH, interface_name(library))
end

function interface_name(library::AbstractString)
    if Sys.iswindows()
        return "$library.dll"
    elseif Sys.isapple()
        return "lib$library.dylib"
    else
        return "lib$library.so"
    end
end

# We can later replace this with:
# using Pkg.Artifacts
# const INTERFACE_PATH = artifact"brainflow"
# Right now using a hardcoded path during refactoring, with manual /lib location
const INTERFACE_PATH = abspath(joinpath(@__DIR__, "../lib"))

# we should instead use dlopen(interface_path())?
const DATA_HANDLER_INTERFACE = interface_path("DataHandler")
const BOARD_CONTROLLER_INTERFACE = interface_path("BoardController")
const ML_MODULE_INTERFACE = interface_path("MLModule")

max_stars_repo_name: matthijscox/brainflow

max_issues_repo_name: matthijscox/brainflow

max_forks_repo_name: matthijscox/brainflow


text:
type Legend
	label::String
end

type DataGroup
	data::DataFrame
	markerScale::Vector{Float64}
	markerLineWidth::Float64
	errorLineWidth::Float64
	lineColor::Union{String, RGBA{Float64}}
	lineStyle::String
	markerColor::Union{String, RGBA{Float64}}
	markerType::String
	legend::Legend
	plotPoints::Bool
end

DataGroup(xvals, yvals) = DataGroup(DataFrame(x = vec(xvals), y = vec(yvals)), [0.7, 0.7], 1.2, 1.2, "blank", "solid", "blank", "blank", Legend("Data Group "), true)
DataGroup(xvals, yvals, yerr) = DataGroup(DataFrame(x = vec(xvals), y = vec(yvals), yerr = vec(yerr)), [0.7, 0.7], 1.2, 1.2, "blank", "solid", "blank", "blank", Legend("Data Group "), true)
DataGroup(df::DataFrame) = DataGroup(df, [0.7, 0.7], 1.2, 1.2, "blank", "solid", "blank", "blank", Legend("Data Group "), true)


text:
@testsuite "construct/direct" AT->begin
    for T in supported_eltypes()
        B = AT{T}(undef, 10)
        @test B isa AT{T,1}
        @test size(B) == (10,)
        @test eltype(B) == T

        B = AT{T}(undef, 10, 10)
        @test B isa AT{T,2}
        @test size(B) == (10, 10)
        @test eltype(B) == T

        B = AT{T}(undef, (10, 10))
        @test B isa AT{T,2}
        @test size(B) == (10, 10)
        @test eltype(B) == T
    end

    # compare against Array
    for typs in [(), (Int,), (Int,1), (Int,2), (Float32,), (Float32,1), (Float32,2)],
        args in [(), (1,), (1,2), ((1,),), ((1,2),),
                 (undef,), (undef, 1,), (undef, 1,2), (undef, (1,),), (undef, (1,2),),
                 (Int,), (Int, 1,), (Int, 1,2), (Int, (1,),), (Int, (1,2),),
                 ([1,2],), ([1 2],)]
        cpu = try
            Array{typs...}(args...)
        catch ex
            isa(ex, MethodError) || rethrow()
            nothing
        end

        gpu = try
            AT{typs...}(args...)
        catch ex
            isa(ex, MethodError) || rethrow()
            cpu == nothing || rethrow()
            nothing
        end

        if cpu == nothing
            @test gpu == nothing
        else
            @test typeof(cpu) == typeof(convert(Array, gpu))
        end
    end
end

@testsuite "construct/similar" AT->begin
    for T in supported_eltypes()
        B = AT{T}(undef, 10)

        B = similar(B, Int32, 11, 15)
        @test B isa AT{Int32,2}
        @test size(B) == (11, 15)
        @test eltype(B) == Int32

        B = similar(B, T)
        @test B isa AT{T,2}
        @test size(B) == (11, 15)
        @test eltype(B) == T

        B = similar(B, (5,))
        @test B isa AT{T,1}
        @test size(B) == (5,)
        @test eltype(B) == T

        B = similar(B, 7)
        @test B isa AT{T,1}
        @test size(B) == (7,)
        @test eltype(B) == T

        B = similar(AT{Int32}, (11, 15))
        @test B isa AT{Int32,2}
        @test size(B) == (11, 15)
        @test eltype(B) == Int32

        B = similar(AT{T}, (5,))
        @test B isa AT{T,1}
        @test size(B) == (5,)
        @test eltype(B) == T

        B = similar(AT{T}, 7)
        @test B isa AT{T,1}
        @test size(B) == (7,)
        @test eltype(B) == T

        B = similar(Broadcast.Broadcasted(*, (B, B)), T)
        @test B isa AT{T,1}
        @test size(B) == (7,)
        @test eltype(B) == T

        if VERSION >= v"1.5"
            B = similar(Broadcast.Broadcasted(*, (B, B)), Int32, (11, 15))
            @test B isa AT{Int32,2}
            @test size(B) == (11, 15)
            @test eltype(B) == Int32
        end
    end
end

@testsuite "construct/convenience" AT->begin
    for T in supported_eltypes()
        A = AT(rand(T, 3))
        b = rand(T)
        fill!(A, b)
        @test A isa AT{T,1}
        @test Array(A) == fill(b, 3)

        A = zero(AT(rand(T, 2)))
        @test A isa AT{T,1}
        @test Array(A) == zero(rand(T, 2))

        A = zero(AT(rand(T, 2, 2)))
        @test A isa AT{T,2}
        @test Array(A) == zero(rand(T, 2, 2))

        A = zero(AT(rand(T, 2, 2, 2)))
        @test A isa AT{T,3}
        @test Array(A) == zero(rand(T, 2, 2, 2))

        A = one(AT(rand(T, 2, 2)))
        @test A isa AT{T,2}
        @test Array(A) == one(rand(T, 2, 2))

        A = oneunit(AT(rand(T, 2, 2)))
        @test A isa AT{T,2}
        @test Array(A) == oneunit(rand(T, 2, 2))
    end
end

@testsuite "construct/conversions" AT->begin
    for T in supported_eltypes()
        Bc = round.(rand(10, 10) .* 10.0)
        B = AT{T}(Bc)
        @test size(B) == (10, 10)
        @test eltype(B) == T
        @test Array(B) ≈ Bc

        Bc = rand(T, 10)
        B = AT(Bc)
        @test size(B) == (10,)
        @test eltype(B) == T
        @test Array(B) ≈ Bc

        Bc = rand(T, 10, 10)
        B = AT{T, 2}(Bc)
        @test size(B) == (10, 10)
        @test eltype(B) == T
        @test Array(B) ≈ Bc

        Bc = rand(Int32, 3, 3, 3)
        B = convert(AT{T, 3}, Bc)
        @test size(B) == (3, 3, 3)
        @test eltype(B) == T
        @test Array(B) ≈ Bc
    end
end

@testsuite "construct/uniformscaling" AT->begin
    for T in supported_eltypes()
        x = Matrix{T}(I, 4, 2)

        x1 = AT{T, 2}(I, 4, 2)
        x2 = AT{T}(I, (4, 2))
        x3 = AT{T, 2}(I, (4, 2))

        @test Array(x1) ≈ x
        @test Array(x2) ≈ x
        @test Array(x3) ≈ x

        x = Matrix(T(3) * I, 2, 4)
        x1 = AT(T(3) * I, 2, 4)
        @test eltype(x1) == T
        @test Array(x1) ≈ x
    end
end


text: include("./src/SparseRF.jl")


text:
@registerN  Meixner AbstractCDOP2 γ μ
export  Meixner

"""
    Meixner{γ,μ}

References: [Koekoek and Swarttouw §1.9](https://arxiv.org/pdf/math/9602214.pdf)
"""
Meixner

basis_symbol(::Type{<:Meixner{γ,μ}}) where {γ,μ} = "m⁽ᵞᵝ⁾"
Polynomials.domain(::Type{<:Meixner{γ, μ}}) where {γ, μ} = Polynomials.Interval(0, Inf)
abcde(::Type{<:Meixner{γ,μ}})  where {γ,μ} = NamedTuple{(:a,:b,:c,:d,:e)}((0,1,0,μ-1,γ*μ))

function kn(P::Type{<:Meixner{γ,μ}}, n::Int) where {γ,μ}
    (1 - 1/μ)^n
end

function k1k0(P::Type{<:Meixner{γ,μ}}, n::Int) where {γ,μ}
    (1 - 1/μ)
end
function k1k_1(P::Type{<:Meixner{γ,μ}}, n::Int) where {γ,μ}
    (1 - 1/μ)^(n>0 ? 2 : 1)
end

function  classical_hypergeometric(P::Type{<:Meixner{γ,μ}}, n::Int,   x) where {γ,μ}
    Pochhammer(γ,n) * pFq((-n, -x), γ, 1 - 1/μ)
end

# Overrides


text:
using Latexify
import MacroTools
using LaTeXStrings

const transition_idx_script = Ref(:^)

"""
    transition_superscript(::Bool)

Specify whether the indices in a [`Transition`](@ref) operator should be
printed as superscript. Default is `true`. If set to `false`, the indices
corresponding to the levels are printed as subscript.
"""
function transition_superscript(x::Bool)
    if x
        transition_idx_script[] = :^
    else
        transition_idx_script[] = :_
    end
    return x
end

function _postwalk_func(x)
    if x==:𝟙
        return "\\mathbb{1}"
    elseif x==:im
        return :i
    elseif MacroTools.@capture(x, dagger(arg_))
        s = "$(arg)^\\dagger"
        return s
    elseif MacroTools.@capture(x, Transition(arg_,i_,j_))
        s = "{$(arg)}$(transition_idx_script[]){{$(i)$(j)}}"
        return s
    else
        return x
    end
end

function _postwalk_average(x)
    if MacroTools.@capture(x, AVG(arg_))
        arg = MacroTools.postwalk(_postwalk_func, arg)
        # TODO: clean up; tricky because of nested string conversion of eg average(dagger(a))
        s = string(arg)
        s = replace(s, "\"" => "")
        s = replace(s, "\\\\" => "\\")
        s = replace(s, "*" => "")
        s = "\\langle " * s * "\\rangle "
        return s
    end
    return x
end

@latexrecipe function f(de::AbstractMeanfieldEquations)
    # Options
    env --> :align
    cdot --> false

    lhs = getfield.(de.equations, :lhs)
    rhs = getfield.(de.equations, :rhs)
    lhs, rhs = _latexify(lhs, rhs, de.iv)
    return lhs, rhs
end

function _latexify(lhs_, rhs_, t)

    # Convert eqs to Exprs
    rhs = _to_expression.(rhs_)
    rhs = [MacroTools.postwalk(_postwalk_func, eq) for eq=rhs]
    rhs = [MacroTools.postwalk(_postwalk_average, eq) for eq=rhs]
    lhs = _to_expression.(lhs_)
    lhs = [MacroTools.postwalk(_postwalk_func, eq) for eq=lhs]
    lhs = [MacroTools.postwalk(_postwalk_average, eq) for eq=lhs]

    tsym = Symbol(t)
    dt = Symbol(:d, tsym)
    ddt = :(d/$(dt))
    lhs = [:($ddt * ($l)) for l=lhs]

    return lhs, rhs
end
_latexify(lhs,rhs) = _latexify([lhs],[rhs])

@latexrecipe function f(op::QNumber)
    # Options
    cdot --> false

    ex = _to_expression(op)
    ex = MacroTools.postwalk(_postwalk_func, ex)
    ex = MacroTools.postwalk(_postwalk_average, ex)
    ex = isa(ex,String) ? latexstring(ex) : ex
    return ex
end

@latexrecipe function f(s::SymbolicUtils.Symbolic{<:CNumber})
    # Options
    cdot --> false

    ex = _to_expression(s)
    ex = MacroTools.postwalk(_postwalk_func, ex)
    ex = MacroTools.postwalk(_postwalk_average, ex)
    ex = isa(ex,String) ? latexstring(ex) : ex
    return ex
end

@latexrecipe function f(c::CorrelationFunction)
    avg = average(c.op1*c.op2)
    return latexify(avg)
end

@latexrecipe function f(S::Spectrum)
    l = latexstring(L"\mathcal{F}(", latexify(S.corr), L")(\omega)")
    return l
end

_to_expression(x::Number) = x
function _to_expression(x::Complex) # For brackets when using latexify
    iszero(x) && return x
    if iszero(real(x))
        return :( $(imag(x))*im )
    elseif iszero(imag(x))
        return real(x)
    else
        return :( $(real(x)) + $(imag(x))*im )
    end
end
_to_expression(op::QSym) = op.name
_to_expression(op::Create) = :(dagger($(op.name)))
_to_expression(op::Transition) = :(Transition($(op.name),$(op.i),$(op.j)) )
function _to_expression(t::QMul)
    args = if SymbolicUtils._isone(t.arg_c)
        t.args_nc
    else
        SymbolicUtils.arguments(t)
    end
    return :( *($(_to_expression.(args)...)) )
end
_to_expression(t::QAdd) = :( +($(_to_expression.(t.arguments)...)) )

_to_expression(p::Parameter) = p.name
function _to_expression(s::SymbolicUtils.Symbolic)
    if SymbolicUtils.istree(s)
        f = SymbolicUtils.operation(s)
        fsym = if f === sym_average
            :AVG
        else
            Symbol(f)
        end
        args = map(_to_expression, SymbolicUtils.arguments(s))
        return :( $(fsym)($(args...)) )
    else
        return nameof(s)
    end
end


text:
# This file is a part of RadiationDetectorDSP.jl, licensed under the MIT License (MIT).

import Test
import RadiationDetectorDSP
import Documenter

Test.@testset "Package RadiationDetectorDSP" begin
    include("test_filters.jl")

    # doctests
    Documenter.DocMeta.setdocmeta!(
        RadiationDetectorDSP,
        :DocTestSetup,
        :(using RadiationDetectorDSP);
        recursive=true,
    )
    Documenter.doctest(RadiationDetectorDSP)
end # testset


text:
export sliceColorDim, sliceTimeDim, sliceSpatialDim, threeSlices,
       getColoredSlices, getColoredSlicesMovie

### sliceColorDim ###

sliceColorDim(A::AbstractArray{T,3}, proj::Integer) where T = A[proj,:,:]
sliceColorDim(A::AbstractArray{T,4}, proj::Integer) where T = A[proj,:,:,:]
sliceColorDim(A::AbstractArray{T,5}, proj::Integer) where T = A[proj,:,:,:,:]

### sliceTimeDim ###

sliceTimeDim(A::ImageMeta{T,D}, proj::String) where {T,D} =
  copyproperties(A, sliceTimeDim(A.data, proj))

function sliceTimeDim(A::AbstractArray{T,D}, proj::String) where {T,D}
  if  proj == "MIP"
    return dropdims(maximum(A, dims=D), dims=D)
  elseif proj == "TTP"
    error("not yet implemented!")
    #data_ = [maximum(d, 4) for d in m.data]
    #data_ = [  mip(d,4) for d in m.data]
  #elseif params[:frameProj] == 2 && ndims(m.data[1]) == 4
  #data_ = [timetopeak(d, alpha=params[:TTPThresh], alpha2=params[:TTPThresh]) for d in m.data]
  end
end

function sliceTimeDim(A::AbstractArray{T,5}, proj::Integer) where T
  return getindex(A,:,:,:,:,proj)
end

function sliceTimeDim(A::AbstractArray{T,4}, proj::Integer) where T
  return getindex(A,:,:,:,proj)
end

### sliceSpatialDim ###

function sliceSpatialDim(A::AbstractArray{T,3}, proj::Union{Tuple,Vector}, dim::Integer) where T
  if dim == 1
    return A[proj[1],:,:]
  elseif dim == 2
    return A[:,proj[2],:]
  else
    return A[:,:,proj[3]]
  end
end

function sliceSpatialDim(A::AbstractArray{T,4}, proj::Union{Tuple,Vector}, dim::Integer) where T
  if dim == 1
    return A[:,proj[1],:,:]
  elseif dim == 2
    return A[:,:,proj[2],:]
  else
    return A[:,:,:,proj[3]]
  end
end

function sliceSpatialDim(A::AbstractArray{T,D}, proj::String, dim::Integer) where {T,D}
  d = (D == 4) ? dim + 1 : dim
  if proj == "MIP"
    return AbstractArray{T,D-1}(dropdims(maximum(A, dims=d),dims=d))
  elseif proj == "SUM"
    return AbstractArray{T,D-1}(dropdims(sum(A, d),dims=d))
  else
    error("proj=$proj not available!")
  end
end

### threeSlices ###

function threeSlices(data::AbstractArray{T,3}, proj) where T
  zx = reverse(permutedims(sliceSpatialDim(data, proj, 2),(2,1)),dims=2)
  zy = permutedims(sliceSpatialDim(data, proj, 1), (2,1))
  xy = reverse(sliceSpatialDim(data, proj, 3), dims=1)
  return zx, zy, xy
end

function threeSlices(data::AbstractArray{T,4}, proj) where T
  zx = sliceSpatialDim(data, proj, 2)
  zy = sliceSpatialDim(data, proj, 1)
  xy = sliceSpatialDim(data, proj, 3)
  return _changeDims(zx, zy, xy)
end

function _changeDims(zx, zy, xy)
  zx_ = reverse(permutedims(zx,(1,3,2)),dims=3)
  zy_ = permutedims(zy, (1,3,2))
  xy_ = reverse(xy, dims=2)
  return zx_, zy_, xy_
end

### getColoredSlices ###

function getColoredSlices(data::AbstractArray{T,4}, dataBG, coloring,
                          minval, maxval, params) where {T}
  slices = (min(params[:sliceX], size(data,2)),
	    min(params[:sliceY], size(data,3)),
	    min(params[:sliceZ], size(data,4)))
  proj = params[:spatialMIP] ? "MIP" : slices

  return getColoredSlices(data, dataBG, coloring, minval, maxval, proj, params[:blendChannels],
                          params[:complexBlending], params[:activeChannel],
                          get(params,:coloringBG,ColoringParams(0,1,"gray")),
                          get(params,:hideFG, false), get(params,:hideBG, false),
                          get(params,:translucentBlending, false),
                          get(params,:spatialMIPBG, false) ? "MIP" : slices)
end

function getColoredSlices(data::AbstractArray{T,4}, dataBG, coloring,
                          minval, maxval, proj, blendChannels, complexBlending,
                          activeChannels, coloringBG, hideFG, hideBG,
                          translucentBlending, projBG) where {T}

  zx, zy, xy = threeSlices(data, proj)

  cdata_zx = colorize(zx,coloring,minval,maxval,blendChannels,activeChannels,complexBlending)
  cdata_zy = colorize(zy,coloring,minval,maxval,blendChannels,activeChannels,complexBlending)
  cdata_xy = colorize(xy,coloring,minval,maxval,blendChannels,activeChannels,complexBlending)

  if dataBG != nothing #!isempty(dataBG)
    zxBG, zyBG, xyBG = threeSlices(dataBG, projBG)  #no MIP for BG data!!!

    minval,maxval = extrema(dataBG)
    cdataBG_zx = colorize(zxBG,coloringBG,minval,maxval)
    cdataBG_zy = colorize(zyBG,coloringBG,minval,maxval)
    cdataBG_xy = colorize(xyBG,coloringBG,minval,maxval)

    if !hideFG && !hideBG
      if translucentBlending
        cdata_zx = blend(cdataBG_zx, cdata_zx)
        cdata_zy = blend(cdataBG_zy, cdata_zy)
        cdata_xy = blend(cdataBG_xy, cdata_xy)
      else
        #all_maps = existing_cmaps()
        colormap = cmap( coloring[1].cmap)

        cdata_zx = overlay(cdataBG_zx, cdata_zx, first(colormap))
        cdata_zy = overlay(cdataBG_zy, cdata_zy, first(colormap))
        cdata_xy = overlay(cdataBG_xy, cdata_xy, first(colormap))
      end
    elseif hideFG
      cdata_zx = cdataBG_zx
      cdata_zy = cdataBG_zy
      cdata_xy = cdataBG_xy
    end
  end

  return cdata_zx, cdata_zy, cdata_xy
end

export getInterpAndColoredSlices

function getInterpAndColoredSlices(data::AbstractArray{T,4}, dataBG,
                                   coloring, minval, maxval, params) where {T}
  if dataBG != nothing
    dataFG = interpolateToRefImage(dataBG, data, params)
    dataBG = interpolateToRefImage(dataBG, params)
  else
    dataFG = data
    dataBG = nothing
  end

  cdata_zx, cdata_zy, cdata_xy = getColoredSlices(dataFG, dataBG, coloring, minval, maxval, params)
  return cdata_zx, cdata_zy, cdata_xy
end

function getColoredSlicesMovie(data::ImageMeta{T,5}, dataBG,
                               coloring, params) where {T}

    xx, yy, zz = _getColoredSlicesMovie(data, dataBG, coloring, params)

    return [copyproperties(data, xx), copyproperties(data, yy), copyproperties(data, zz)]

end

function getColoredSlicesMovie(data::AbstractArray{T,5}, dataBG,
                               coloring, params) where {T}
   return _getColoredSlicesMovie(data, dataBG, coloring, params)
end

function _getColoredSlicesMovie(data::AbstractArray{T,5}, dataBG,
                               coloring, params) where {T}
    L = size(data, 5)

    maxval = [maximum(sliceColorDim(data,d)) for d=1:size(data,Axis{:color})]
    minval = [minimum(sliceColorDim(data,d)) for d=1:size(data,Axis{:color})]

    cdata_zx, cdata_zy, cdata_xy = getInterpAndColoredSlices(sliceTimeDim(data,1), dataBG, coloring, minval, maxval, params)

    cdata_xy_all = AxisArray(zeros(eltype(cdata_xy),size(cdata_xy,1),size(cdata_xy,2),L),
                             AxisArrays.axes(cdata_xy)..., AxisArrays.axes(data)[5])

    cdata_zx_all = AxisArray(zeros(eltype(cdata_zx),size(cdata_zx,1),size(cdata_zx,2),L),
                             AxisArrays.axes(cdata_zx)..., AxisArrays.axes(data)[5])

    cdata_zy_all = AxisArray(zeros(eltype(cdata_zy),size(cdata_zy,1),size(cdata_zy,2),L),
                             AxisArrays.axes(cdata_zy)..., AxisArrays.axes(data)[5])

    cdata_xy_all[:,:,1] .= cdata_xy
    cdata_zy_all[:,:,1] .= cdata_zy
    cdata_zx_all[:,:,1] .= cdata_zx

    for l=2:L
      cdata_zx, cdata_zy, cdata_xy = getInterpAndColoredSlices(sliceTimeDim(data,l), dataBG, coloring, minval, maxval, params)

      cdata_xy_all[:,:,l] .= cdata_xy
      cdata_zy_all[:,:,l] .= cdata_zy
      cdata_zx_all[:,:,l] .= cdata_zx
    end
    return [cdata_xy_all, cdata_zx_all, cdata_zy_all]
end


text:
println("Compile libraries on main processor..."); flush(stdout)
include("imports.jl")
println("Finished loading libraries."); flush(stdout)

using Distributed
rmprocs(workers())
addprocs(32)
println("Compile libraries on workers ..."); flush(stdout)
@everywhere include("imports.jl")
println("Finished loading libraries."); flush(stdout)

markers = [:CD3z, :EOMES, :Perforin, :Granzyme_A, :Siglec7, :LAG3, :CD56, :CD57]
println("markers: ", markers); flush(stdout)

Ks = collect(2:9)
println("K: ", Ks)

skewtmix = [true, false]
println("skewtmix: ", skewtmix); flush(stdout)

sims = dict_list(Dict(:K => Ks, :marker => markers, :skewtmix => skewtmix))

# Run simulations in parallel.
println("Running analyses ..."); flush(stdout)
@time res = pmap(run, sims, on_error=identity)

# Print results status.
foreach(z -> println("$(z[1]) => $(z[2])"), zip(sims, res))
flush(stdout)

# Post process
println("Post process ..."); flush(stdout)
@time pp_res = pmap(postprocess, sims)

# Print number of small clusters
@time nc_res = pmap(sim -> print_number_of_small_clusters(sim, p=0.02),
                    cdd.MCMC.ProgressBar(sims), on_error=identity)

# Combine results, plot DIC.
println("Compute DIC ..."); flush(stdout)
foreach(marker -> plot_dic(marker, Ks, skewtmix), markers)
foreach(marker -> plot_dic(marker, Ks, skewtmix, calibrate=true), markers)

# Send results to S3
println("Send results to s3 ..."); flush(stdout)
Util.s3sync(from="$(Info.resultsdir_datastudy)/$(simname)",
            to="$(Info.awsbucket_datastudy)/$(simname)",
            tags=`--exclude '*.nfs'`)


text:
using COVID19SchoolReopening
using Test

@testset "COVID19SchoolReopening.jl" begin
    # Write your tests here.
end


text:
## v updates

"""
    updatev!(M::HePPCATModel,Y,method)

Update the noise variances `M.v` with respect to data `Y` using `method`.
Internally calls [`updatevl(_)`](@ref) to update each entry.
"""
function updatev!(M::HePPCATModel,Y,method)
    for (l,Yl) in enumerate(Y)
        M.v[l] = updatevl(M.v[l],M.U,M.λ,Yl,method)
    end
    return M
end

"""
    updatevl(vl,U,λ,Yl,method)

Compute update for `l`th noise variance `vl` with respect to data `Yl` using `method`.
"""
function updatevl end

# Update method: Global maximization via root-finding
"""
    updatevl(vl,U,λ,Yl,::RootFinding)

Root-finding (global maximization) update of `vl`.
"""
function updatevl(vl,U,λ,Yl,::RootFinding)
    d, k = size(U)
    nl = size(Yl,2)

    # Compute coefficients and check edge case
    UYl = [sum(abs2,Yl'U[:,j]) for j in 1:k]
    UYl0 = nonnegative(sum(abs2,Yl)-sum(UYl))
    α = [j == 0 ? d-k : 1 for j in IdentityRange(0:k)]
    β = [j == 0 ? UYl0/nl : UYl[j]/nl for j in IdentityRange(0:k)]
    γ = [j == 0 ? zero(eltype(λ)) : λ[j] for j in IdentityRange(0:k)]
    all(iszero(β[j]) for j in 0:k if iszero(γ[j])) && return zero(vl)

    # Find nonnegative critical points
    tol = 1e-8  # todo: choose tolerance adaptively
    vmin, vmax = extrema(β[j]/α[j]-γ[j] for j in 0:k if !(iszero(α[j]) && iszero(β[j])))
    vcritical = roots(interval(vmin,vmax) ∩ interval(zero(vl),Inf),Newton,tol) do v
        - sum(α[j]/(γ[j]+v) for j in 0:k if !iszero(α[j])) + sum(β[j]/(γ[j]+v)^2 for j in 0:k if !iszero(β[j]))
    end

    # Return maximizer
    return _argmax(mid.(interval.(vcritical))) do v
        -(sum(α[j]*log(γ[j]+v) for j in 0:k if !iszero(α[j])) + sum(β[j]/(γ[j]+v) for j in 0:k if !iszero(β[j])))
    end
end

# Update method: Expectation Maximization
"""
    updatevl(vl,U,λ,Yl,::ExpectationMaximization)

Expectation maximization update of `vl`.
"""
function updatevl(vl,U,λ,Yl,::ExpectationMaximization)
    d, k = size(U)
    nl = size(Yl,2)

    UYl = [sum(abs2,Yl'U[:,j]) for j in 1:k]
    UYl0 = nonnegative(sum(abs2,Yl)-sum(UYl))
    β = [j == 0 ? UYl0/nl : UYl[j]/nl for j in IdentityRange(0:k)]
    γ = [j == 0 ? zero(eltype(λ)) : λ[j] for j in IdentityRange(0:k)]

    ρ = sum((oneunit(eltype(λ)).-γ./(γ.+vl)).^2 .* β) + vl*sum(λ./(λ.+vl))
    return ρ/d
end

# Update method: Difference of concave approach
"""
    updatevl(vl,U,λ,Yl,::DifferenceOfConcave)

Difference of concave update of `vl`.
"""
function updatevl(vl,U,λ,Yl,::DifferenceOfConcave)
    d, k = size(U)
    nl = size(Yl,2)

    # Compute coefficients and check edge case
    UYl = [sum(abs2,Yl'U[:,j]) for j in 1:k]
    UYl0 = nonnegative(sum(abs2,Yl)-sum(UYl))
    α = [j == 0 ? d-k : 1 for j in IdentityRange(0:k)]
    β = [j == 0 ? UYl0/nl : UYl[j]/nl for j in IdentityRange(0:k)]
    γ = [j == 0 ? zero(eltype(λ)) : λ[j] for j in IdentityRange(0:k)]
    affslope = -sum(α[j]/(γ[j]+vl) for j in 0:k if !iszero(α[j]))
    Ltp = v -> affslope + sum(β[j]/(γ[j]+v)^2 for j in 0:k if !iszero(β[j]))
    if affslope == -Inf || Ltp(zero(vl)) <= zero(vl)
        return zero(vl)
    end

    # Return nonnegative critical point
    tol = 1e-8  # todo: choose tolerance adaptively
    vmax = maximum(sqrt(β[j]/α[j]*(γ[j]+vl)) - γ[j] for j in 0:k if !(iszero(α[j]) && iszero(β[j])))
    return find_zero(Ltp, (zero(vmax),vmax))
end

# Update method: Quadratic solvable minorizer
"""
    updatevl(vl,U,λ,Yl,::QuadraticSolvableMinorizer)

Minorize maximize update of `vl` with quadratic solvable minorizer.
"""
function updatevl(vl,U,λ,Yl,::QuadraticSolvableMinorizer)
    d, k = size(U)
    nl = size(Yl,2)

    # Compute coefficients
    UYl = [sum(abs2,Yl'U[:,j]) for j in 1:k]
    UYl0 = nonnegative(sum(abs2,Yl)-sum(UYl))
    α = [j == 0 ? d-k : 1 for j in IdentityRange(0:k)]
    β = [j == 0 ? UYl0/nl : UYl[j]/nl for j in IdentityRange(0:k)]
    γ = [j == 0 ? zero(eltype(λ)) : λ[j] for j in IdentityRange(0:k)]
    J0 = findall(iszero,γ)
    αtl = sum(α[j] for j in J0)
    βtl = sum(β[j] for j in J0)
    ζtl = sum(α[j]/(γ[j]+vl) for j in 0:k if j ∉ J0)
    B = βtl + sum(β[j]*vl^2/(γ[j]+vl)^2 for j in 0:k if j ∉ J0)

    return (-αtl + sqrt(αtl^2 + 4*ζtl*B))/(2*ζtl)
end

# Update method: Cubic solvable minorizer
"""
    updatevl(vl,U,λ,Yl,::CubicSolvableMinorizer)

Minorize maximize update of `vl` with cubic solvable minorizer.
"""
function updatevl(vl,U,λ,Yl,::CubicSolvableMinorizer)
    d, k = size(U)
    nl = size(Yl,2)

    # Compute coefficients
    UYl = [sum(abs2,Yl'U[:,j]) for j in 1:k]
    UYl0 = nonnegative(sum(abs2,Yl)-sum(UYl))
    α = [j == 0 ? d-k : 1 for j in IdentityRange(0:k)]
    β = [j == 0 ? UYl0/nl : UYl[j]/nl for j in IdentityRange(0:k)]
    γ = [j == 0 ? zero(eltype(λ)) : λ[j] for j in IdentityRange(0:k)]
    c = [-2*β[j]/γ[j]^3 for j in IdentityRange(0:k)]
    J0 = findall(iszero,γ)
    αtl = sum(α[j] for j in J0)
    βtl = sum(β[j] for j in J0)
    ζtl = sum(α[j]/(γ[j]+vl) for j in 0:k if j ∉ J0)
    γtl = -ζtl + sum(β[j]/(γ[j]+vl)^2 for j in 0:k if j ∉ J0)
    ctl = sum(c[j] for j in 0:k if j ∉ J0)

    complexroots = PolynomialRoots.solve_cubic_eq(complex.([βtl,-αtl,γtl-ctl*vl,ctl]))
    vcritical = [real(v) for v in complexroots if real(v) ≈ v && real(v) >= zero(vl)]

    isempty(vcritical) && return zero(vl)
    return _argmax(vcritical) do v
        -αtl*log(v) - βtl/v - ζtl*v + sum(β[j]/(γ[j]+vl)^2*v + (1/2)*c[j]*(v-vl)^2 for j in 0:k if j ∉ J0)
    end
end

# Utilities
nonnegative(x) = max(zero(x),x)


text:
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Defines the data structures used to record solution quality and what happened
# on each iteration.

"""
RestartChoice specifies whether a restart was performed on a given iteration.

# Values

- `RESTART_CHOICE_UNSPECIFIED`: Default value.
- `RESTART_CHOICE_NO_RESTART`: No restart on this iteration.
- `RESTART_CHOICE_WEIGHTED_AVERAGE_RESET`: The weighted average of iterates is
  cleared and reset to the current point. Note that from a mathematical
  perspective this can be equivalently viewed as restarting the algorithm but
  picking the restart point to be the current iterate.
- `RESTART_CHOICE_RESTART_TO_AVERAGE`: The algorithm is restarted at the average
  of iterates since the last restart.
"""
@enum RestartChoice begin
  RESTART_CHOICE_UNSPECIFIED
  RESTART_CHOICE_NO_RESTART
  RESTART_CHOICE_WEIGHTED_AVERAGE_RESET
  RESTART_CHOICE_RESTART_TO_AVERAGE
end

"""
PointType identifies the type of point used to compute the fields in a given
struct; see ConvergenceInformation and InfeasibilityInformation.

# Values

- `POINT_TYPE_UNSPECIFIED`: Default value.
- `POINT_TYPE_CURRENT_ITERATE`: Current iterate (x_k, y_k).
- `POINT_TYPE_ITERATE_DIFFERENCE`: Difference of iterates
  (x_{k+1} - x_k, y_{k+1} - y_k).
- `POINT_TYPE_AVERAGE_ITERATE`: Average of iterates since the last restart.
- `POINT_TYPE_NONE`: There is no corresponding point.
"""
@enum PointType begin
  POINT_TYPE_UNSPECIFIED
  POINT_TYPE_CURRENT_ITERATE
  POINT_TYPE_ITERATE_DIFFERENCE
  POINT_TYPE_AVERAGE_ITERATE
  POINT_TYPE_NONE
end

"""
Information measuring how close a candidate is to establishing feasibility and
optimality; see also TerminationCriteria.
"""
mutable struct ConvergenceInformation
  """
  Type of the candidate point described by this ConvergenceInformation.
  """
  candidate_type::PointType

  """
  The primal objective. The primal need not be feasible.
  """
  primal_objective::Float64

  """
  The dual objective. The dual need not be feasible. The dual objective should
  include the contributions from the reduced costs.
  """
  dual_objective::Float64

  """
  If possible (e.g., when all primal variables have lower and upper bounds), a
  correct dual bound. The field should be set to negative infinity if no
  corrected dual bound is available.
  """
  corrected_dual_objective::Float64

  """
  The maximum violation of any primal constraint, i.e., the l_∞ norm of the
  violations.
  """
  l_inf_primal_residual::Float64

  """
  The l_2 norm of the violations of primal constraints.
  """
  l2_primal_residual::Float64

  """
  The maximum violation of any dual constraint, i.e., the l_∞ norm of the
  violations.
  """
  l_inf_dual_residual::Float64

  """
  The l_2 norm of the violations of dual constraints.
  """
  l2_dual_residual::Float64

  # Relative versions of the residuals, defined as
  #   relative_residual = residual / (eps_ratio + norm),
  # where
  #   eps_ratio = eps_optimal_absolute / eps_optimal_relative
  #   residual = one of the residuals (l{2,_inf}_{primal,dual}_residual)
  #   norm = the relative norm (l{2,_inf} norm of
  #          {constraint_bounds,primal_linear_objective} respectively).
  # If eps_optimal_relative = 0.0, these will all be 0.0.
  #
  # If eps_optimal_relative > 0.0, the absolute and relative termination
  # criteria translate to relative_residual <= eps_optimal_relative.
  # NOTE: The usefulness of these relative residuals is based on their
  # relationship to TerminationCriteria. If the TerminationCriteria change
  # consider adding additional iteration measures here.

  relative_l_inf_primal_residual::Float64

  relative_l2_primal_residual::Float64

  relative_l_inf_dual_residual::Float64

  relative_l2_dual_residual::Float64

  """
  Relative optimality gap:
    |primal_objective - dual_objective| /
    (eps_ratio + |primal_objective| + |dual_objective|)
  """
  relative_optimality_gap::Float64

  """
  The maximum absolute value of the primal variables, i.e., the l_∞ norm. This
  is useful to detect when the primal iterates are diverging. Divergence of the
  primal variables could be an algorithmic issue, or indicate that the dual is
  infeasible.
  """
  l_inf_primal_variable::Float64

  """
  The l_2 norm of the primal variables.
  """
  l2_primal_variable::Float64

  """
  The maximum absolute value of the dual variables, i.e., the l_∞ norm. This is
  useful to detect when the dual iterates are diverging. Divergence of the dual
  variables could be an algorithmic issue, or indicate the primal is infeasible.
  """
  l_inf_dual_variable::Float64

  """
  The l_2 norm of the dual variables.
  """
  l2_dual_variable::Float64
end

function ConvergenceInformation()
  return ConvergenceInformation(POINT_TYPE_UNSPECIFIED, zeros(16)...)
end

"""
Information measuring how close a point is to establishing primal or dual
infeasibility (i.e. has no solution); see also TerminationCriteria.
"""
mutable struct InfeasibilityInformation
  """
  Type of the point used to compute the InfeasibilityInformation.
  """
  candidate_type::PointType

  """
  Let x_ray be the algorithm's estimate of the primal extreme ray where x_ray is
  a vector scaled such that its infinity norm is one. A simple and typical
  choice of x_ray is x_ray = x / | x |_∞ where x is the current primal iterate.
  For this value compute the maximum absolute error in the primal linear program
  with the right hand side set to zero.
  """
  max_primal_ray_infeasibility::Float64

  """
  The value of the linear part of the primal objective (ignoring additive
  constants) evaluated at x_ray, i.e., c' * x_ray where c is the objective
  coefficient vector.
  """
  primal_ray_linear_objective::Float64

  """
  The l_∞ norm of the vector resulting from taking the quadratic matrix from
  primal objective and multiplying it by the primal variables. For linear
  programming problems this is zero.
  """
  primal_ray_quadratic_norm::Float64

  """
  Let (y_ray, r_ray) be the algorithm's estimate of the dual and reduced cost
  extreme ray where (y_ray, r_ray) is a vector scaled such that its infinity
  norm is one. A simple and typical choice of y_ray is
  (y_ray, r_ray) = (y, r) / max(| y |_∞, | r |_∞) where y is the current dual
  iterate and r is the current dual reduced costs. Consider the quadratic
  program we are solving but with the objective (both quadratic and linear
  terms) set to zero. This forms a linear program (label this linear program
  (1)) with no objective. Take the dual of (1) and compute the maximum absolute
  value of the constraint error for (y_ray, r_ray) to obtain the value of
  max_dual_ray_infeasibility.
  """
  max_dual_ray_infeasibility::Float64

  """
  The objective of the linear program labeled (1) in the previous paragraph.
  """
  dual_ray_objective::Float64
end

function InfeasibilityInformation()
  return InfeasibilityInformation(POINT_TYPE_UNSPECIFIED, zeros(5)...)
end

# All values in IterationStats assume that the primal quadratic program is a
# minimization problem and the dual is a maximization problem. Problems should
# be transformed to this form if they are not already in this form. The dual
# vector is defined to be the vector of multipliers on the linear constraints,
# that is, excluding dual multipliers on variable bounds (reduced costs).
mutable struct IterationStats
  """
  The iteration number at which these stats were recorded. By convention,
  iteration counts start at 1, and the stats correspond to the solution *after*
  the iteration. Therefore stats from iteration 0 are the stats at the starting
  point.
  """
  iteration_number::Int32

  """
  A set of statistics measuring how close a point is to establishing primal and
  dual feasibility and optimality. This field is repeated since there might be
  several different points that are considered.
  """
  convergence_information::Vector{ConvergenceInformation}

  """
  A set of statistics measuring how close a point is to establishing primal or
  dual infeasibility (i.e., has no solution). This field is repeated since there
  might be several different points that could establish infeasibility.
  """
  infeasibility_information::Vector{InfeasibilityInformation}

  """
  The cumulative number of passes through the KKT matrix since the start of the
  solve. One pass is a multply by the constraint matrix, its transpose and the
  matrix that defines the quadratic part of the objective.

  For example, each iteration of mirror saddle prox contributes 2.0 to this sum.
  This is a float because it can include fractional passes through the data.
  For example, in an active set method we may only use a submatrix with 20% of
  the nonzeros of the KKT matrix at each iteration in which case 0.2 would be
  added to the total.
  """
  cumulative_kkt_matrix_passes::Float64

  """
  The total number of rejected steps (e.g., within a line search procedure)
  since the start of the solve.
  """
  cumulative_rejected_steps::Int32

  """
  The amount of time passed since we started solving the problem (see solver log
  solve_time_sec) which records total time.
  """
  cumulative_time_sec::Float64

  """
  The kind of restart that occurred at this iteration, or
  RESTART_CHOICE_NO_RESTART if a restart did not occur.
  """
  restart_used::RestartChoice

  """
  Step size used at this iteration. Note that the step size used for the primal
  update is step_size / primal_weight, while the one used for the dual update
  is step_size * primal_weight.
  """
  step_size::Float64

  """
  Primal weight controlling the relation between primal and dual step sizes. See
  field 'step_size' for a detailed description.
  """
  primal_weight::Float64

  method_specific_stats::Dict{String,Float64}
end

function IterationStats()
  return IterationStats(
    0,
    ConvergenceInformation[],
    InfeasibilityInformation[],
    0.0,
    Int32(0),
    0.0,
    RESTART_CHOICE_UNSPECIFIED,
    0.0,
    0.0,
    Dict{String,Float64}(),
  )
end

"""
TerminationReason explains why the solver stopped. See termination.jl for the
precise criteria used to check termination.

# Values

- `TERMINATION_REASON_UNSPECIFIED`: Default value.
- `TERMINATION_REASON_OPTIMAL`
- `TERMINATION_REASON_PRIMAL_INFEASIBLE`: Note in this situation the dual could
  be either unbounded or infeasible.
- `TERMINATION_REASON_DUAL_INFEASIBLE`: Note in this situation the primal could
  either unbounded or infeasible.
- `TERMINATION_REASON_TIME_LIMIT`
- `TERMINATION_REASON_KKT_MATRIX_PASS_LIMIT`
- `TERMINATION_REASON_NUMERICAL_ERROR`
- `TERMINATION_REASON_INVALID_PROBLEM`: Indicates that the solver detected
  invalid problem data, e.g., inconsistent bounds.
- `TERMINATION_REASON_OTHER`
"""
@enum TerminationReason begin
  TERMINATION_REASON_UNSPECIFIED
  TERMINATION_REASON_OPTIMAL
  TERMINATION_REASON_PRIMAL_INFEASIBLE
  TERMINATION_REASON_DUAL_INFEASIBLE
  TERMINATION_REASON_TIME_LIMIT
  TERMINATION_REASON_ITERATION_LIMIT
  TERMINATION_REASON_KKT_MATRIX_PASS_LIMIT
  TERMINATION_REASON_NUMERICAL_ERROR
  TERMINATION_REASON_INVALID_PROBLEM
  TERMINATION_REASON_OTHER
end

mutable struct SolveLog
  """
  The name of the optimization problem.
  """
  instance_name::String

  """
  The command line used to invoke the solve. This is used to record the solve
  parameters. Does not apply if the solve is invoked through an API.
  """
  command_line_invocation::String

  """
  The reason that the solve terminated.
  """
  termination_reason::TerminationReason

  """
  Optional extra information about the termination reason.
  """
  termination_string::String

  """
  The total number of iterations during the solve.
  """
  iteration_count::Int32

  """
  The runtime of the solve. Note: This should not be used for comparing methods
  unless care is taken to control for noise in runtime measurement.
  """
  solve_time_sec::Float64

  """
  The IterationStats corresponding to the solution returned by the solver.
  """
  solution_stats::IterationStats

  """
  The type of the output point. This type specifies the information entry that
  prompted termination. If TerminationReason is TERMINATION_REASON_OPTIMAL, this
  type matches the PointType of the ConvergenceInformation entry (in
  solution_stats) that caused termination. Similarly, if TerminationReason is
  either TERMINATION_REASON_PRIMAL_INFEASIBLE or
  TERMINATION_REASON_DUAL_INFEASIBLE this type identifies the
  InfeasibilityInformation entry that caused termination.
  """
  solution_type::PointType

  """
  A history of iteration stats for the solve. The iteration_number fields should
  be in increasing order. The frequency at which these stats should be recorded
  is not specified. This field is "more" optional than the others because it
  often significantly increases the size of the message, and because the
  information may not be available for third-party solvers.
  """
  iteration_stats::Vector{IterationStats}
end

function SolveLog()
  return SolveLog(
    "",
    "",
    TERMINATION_REASON_UNSPECIFIED,
    "",
    Int32(0),
    0.0,
    IterationStats(),
    POINT_TYPE_UNSPECIFIED,
    IterationStats[],
  )
end

# Needed for JSON serialization.
StructTypes.StructType(::Type{ConvergenceInformation}) = StructTypes.Mutable()
StructTypes.StructType(::Type{InfeasibilityInformation}) = StructTypes.Mutable()
StructTypes.StructType(::Type{IterationStats}) = StructTypes.Mutable()
StructTypes.StructType(::Type{SolveLog}) = StructTypes.Mutable()


text:
"Calculates the angular momentum between two wavefunctions and around the center."
function calc_angmom(wfc1::WannierFunction{N, T}, wfc2::WannierFunction{N, T}, center::Point3{T}, cutoff=Inf) where {N, T <: AbstractFloat}
	points = wfc1.points
    origin = points[1, 1, 1]
    da     = points[2, 1, 1] - origin
    db     = points[1, 2, 1] - origin
    dc     = points[1, 1, 2] - origin
    V      = SMatrix{3,3}(inv([convert(Array, da) convert(Array, db) convert(Array, dc)])')
    L      = zero(Point3{Complex{T}})
    c2 = cutoff^2
    @inbounds for i2 in 2:size(wfc1, 3)
        for i1 in 2:size(wfc1, 2)
            for i in 2:size(wfc1, 1)
                r       = points[i, i1, i2] - center
                if dot(r, r) < c2
                    dw_cryst = Point3(wfc2.values[i, i1, i2] - wfc2.values[i-1, i1,   i2],
                                      wfc2.values[i, i1, i2] - wfc2.values[i,   i1-1, i2],
      	  	  	  		              wfc2.values[i, i1, i2] - wfc2.values[i,   i1,   i2-1])

                    dw_cart = V * dw_cryst
                    L      += (wfc1.values[i, i1, i2]',) .* cross(r, dw_cart)
                end
	        end
	    end
	end
	return -1im * L
end

#this doesn't work I think
# function calc_angmom_squared(wfc1::WannierFunction{N, T}, wfc2::WannierFunction{N, T}, center::Point3{T}) where {N, T <: AbstractFloat}
# 	points = wfc1.points
#     origin = points[1, 1, 1]
#     da     = points[2, 1, 1] - origin
#     db     = points[1, 2, 1] - origin
#     dc     = points[1, 1, 2] - origin
#     V      = SMatrix{3,3}(inv([convert(Array, da) convert(Array, db) convert(Array, dc)])')
#     Lsq      = zero(Complex{T})

#     @inbounds for i2 = 2:size(wfc1)[3]
#         for i1 = 2:size(wfc1)[2]
#             for i = 2:size(wfc1)[1]
#                 dw_cryst = Point3(wfc2.values[i, i1, i2] - wfc2.values[i-1, i1,   i2],
#                                   wfc2.values[i, i1, i2] - wfc2.values[i,   i1-1, i2],
#   	  	  	  		              wfc2.values[i, i1, i2] - wfc2.values[i,   i1,   i2-1])
# 	            dw_cryst_sq = map(x->x .^2,dw_cryst)

#                 r       = points[i, i1, i2] - center
#                 dw_cart = V * dw_cryst
#                 Lsq    += wfc1.values[i, i1, i2] ⋅ (r[2]^2 * (dw_cryst_sq[1] + dw_cryst_sq[3]) +
#                                                     r[1]^2 * (dw_cryst_sq[2] + dw_cryst_sq[3]) +
#                                                     r[3]^2 * (dw_cryst_sq[1] + dw_cryst_sq[2]) -
#                                                     2 * (r ⋅ dw_cryst +
#                                                          r[2] * r[3] * dw_cryst[2] .* dw_cryst[3] +
#                                                          r[1] * r[3] * dw_cryst[1] .* dw_cryst[3] +
#                                                          r[1] * r[2] * dw_cryst[1] .* dw_cryst[2]))

# 	        end
# 	    end
# 	end
# 	return Lsq
# end

function calc_spin(wfc1::WannierFunction{2, T}, wfc2::WannierFunction{2, T}) where T <: AbstractFloat
	S = Point3(SMatrix{2, 2}(0, 1, 1, 0)/2,
	           SMatrix{2, 2}(0, -1im, 1im, 0)/2,
	           SMatrix{2, 2}(1, 0, 0, -1)/2)

	outS = zero(Point3{Complex{T}})
	for (w1, w2) in zip(wfc1.values, wfc2.values)
		outS += (w1',) .* S .* (w2,)
	end
	return outS
end

"Calculates the dipole term between two wavefunctions. Make sure the wavefunctions are normalized!"
function calc_dip(wfc1::WannierFunction{N, T}, wfc2::WannierFunction{N, T}) where {N, T <: AbstractFloat}
	out = zero(Point3{Complex{T}})
	for (w1, w2, p) in zip(wfc1, wfc2, wfc1.points)
		out += w1' * w2 * p
	end
	return real(out)
end

# Is this code actually correct?
# "Calculates the dipoles from the supplied wannier dipole output."
# function calc_k_dips(dip_raw::Array{Tuple{Int, Int, Int, Int, Int, Point3{T}}}, kpoints::AbstractArray) where T<:AbstractFloat
# 	dim = 0
# 	for i=1:length(dip_raw)
# 		d = dip_raw[i][4]
# 		if d > dim
# 			dim = d
# 		else
# 			break
# 		end
# 	end
# 	out = zeros(Point3{T}, dim, dim)
# 	tmp = [[zeros(Complex{T}, 3)] for i=1:dim, i1=1:dim]
# 	for i=1:size(dip_raw)[1]
# 		d = dip_raw[i]
# 		complex_part = 2π*(kpoints[1]*d[1]+kpoints[2]*d[2]+kpoints[3]*d[3])
# 		factor = exp(-1im * complex_part)
# 		tmp[d[4],d[5]][1] += d[6][1] * factor
# 		tmp[d[4],d[5]][2] += d[6][2] * factor
# 		tmp[d[4],d[5]][3] += d[6][3] * factor
# 	end
# 	for i in eachindex(out)
# 		out[i] = Point3(real(tmp[i][1]),real(tmp[i][2]),real(tmp[i][3]))
# 	end
# 	return Mat{2*dim, 2*dim, Point3{T}}([out zeros(out);zeros(out) out])
# end


text:
struct SMultHandler end

struct SMatMultHandler end

gethandler(::Bool, ::Type{typeof(*)}, ::Type{<:SBitstream}, ::Type{<:SBitstream}) = SMultHandler()
init_state(::SMultHandler) = (id = 0,)

gethandler(broadcasted::Bool,
           ::Type{typeof(*)},
           ::Type{<:AbstractArray{<:SBitstream}},
           ::Type{<:AbstractArray{<:SBitstream}}) =
    broadcasted ? SMultHandler() : SMatMultHandler()
init_state(::SMatMultHandler) = (id = 0,)

function (handler::SMultHandler)(buffer, netlist, state, inputs, outputs)
    # update netlist with inputs
    setsigned!(netlist, inputs[1], true)
    setsigned!(netlist, inputs[2], true)

    # compute output size
    lname, rname = handle_broadcast_name(name(inputs[1]), name(inputs[2]),
                                         netsize(inputs[1]), netsize(inputs[2]))
    outsize = netsize(outputs[1])

    # add internal nets to netlist
    push!(netlist, Net(name = "mult$(state.id)_out_pp", size = outsize))
    push!(netlist, Net(name = "mult$(state.id)_out_pm", size = outsize))
    push!(netlist, Net(name = "mult$(state.id)_out_mp", size = outsize))
    push!(netlist, Net(name = "mult$(state.id)_out_mm", size = outsize))
    push!(netlist, Net(name = "mult$(state.id)_out_11", size = outsize))
    push!(netlist, Net(name = "mult$(state.id)_out_12", size = outsize))
    push!(netlist, Net(name = "mult$(state.id)_out_13", size = outsize))
    push!(netlist, Net(name = "mult$(state.id)_out_14", size = outsize))

    # add output net to netlist
    setsigned!(netlist, outputs[1], true)

    write(buffer, """
        $stdcomment
        // BEGIN mult$(state.id)
        assign mult$(state.id)_out_pp = $(lname("_p")) & $(rname("_p"))
        assign mult$(state.id)_out_pm = $(lname("_p")) & $(rname("_m"))
        assign mult$(state.id)_out_mp = $(lname("_m")) & $(rname("_p"))
        assign mult$(state.id)_out_mm = $(lname("_m")) & $(rname("_m"))
        stoch_sat_sub_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mult$(state.id)_11 (
                .CLK(CLK),
                .nRST(nRST),
                .A(mult$(state.id)_out_pp),
                .B(mult$(state.id)_out_pm),
                .Y(mult$(state.id)_out_11)
            );
        stoch_sat_sub_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mult$(state.id)_12 (
                .CLK(CLK),
                .nRST(nRST),
                .A(mult$(state.id)_out_pm),
                .B(mult$(state.id)_out_pp),
                .Y(mult$(state.id)_out_12)
            );
        stoch_sat_sub_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mult$(state.id)_13 (
                .CLK(CLK),
                .nRST(nRST),
                .A(mult$(state.id)_out_mp),
                .B(mult$(state.id)_out_mm),
                .Y(mult$(state.id)_out_13)
            );
        stoch_sat_sub_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mult$(state.id)_14 (
                .CLK(CLK),
                .nRST(nRST),
                .A(mult$(state.id)_out_mm),
                .B(mult$(state.id)_out_mp),
                .Y(mult$(state.id)_out_14)
            );
        stoch_add_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mult$(state.id)_p (
                .CLK(CLK),
                .nRST(nRST),
                .A(mult$(state.id)_out_11),
                .B(mult$(state.id)_out_14),
                .Y($(name(outputs[1]))_p)
            );
        stoch_add_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mult$(state.id)_m (
                .CLK(CLK),
                .nRST(nRST),
                .A(mult$(state.id)_out_12),
                .B(mult$(state.id)_out_13),
                .Y($(name(outputs[1]))_m)
            );
        // END mult$(state.id)
        \n""")

    return buffer, (id = state.id + 1,)
end

function (handler::SMatMultHandler)(buffer, netlist, state, inputs, outputs)
    # update netlist with inputs
    setsigned!(netlist, inputs[1], true)
    setsigned!(netlist, inputs[2], true)

    # compute output size
    m, n = netsize(inputs[1])
    _, p = netsize(inputs[2])
    outsize = netsize(outputs[1])

    # add output net to netlist
    setsigned!(netlist, outputs[1], true)

    # add internal nets to netlist
    push!(netlist, Net(name = "mmult$(state.id)_out_pp", size = outsize))
    push!(netlist, Net(name = "mmult$(state.id)_out_pm", size = outsize))
    push!(netlist, Net(name = "mmult$(state.id)_out_mp", size = outsize))
    push!(netlist, Net(name = "mmult$(state.id)_out_mm", size = outsize))
    push!(netlist, Net(name = "mmult$(state.id)_out_11", size = outsize))
    push!(netlist, Net(name = "mmult$(state.id)_out_12", size = outsize))
    push!(netlist, Net(name = "mmult$(state.id)_out_13", size = outsize))
    push!(netlist, Net(name = "mmult$(state.id)_out_14", size = outsize))

    write(buffer, """
        $stdcomment
        // BEGIN mmult$(state.id)
        stoch_matrix_mult #(
                .NUM_ROWS($m),
                .NUM_MID($n),
                .NUM_COLS($p)
            ) mmult$(state.id)_pp (
                .CLK(CLK),
                .nRST(nRST),
                .A($(name(inputs[1]))_p),
                .B($(name(inputs[2]))_p),
                .Y(mmult$(state.id)_out_pp)
            );
        stoch_matrix_mult #(
                .NUM_ROWS($m),
                .NUM_MID($n),
                .NUM_COLS($p)
            ) mmult$(state.id)_pm (
                .CLK(CLK),
                .nRST(nRST),
                .A($(name(inputs[1]))_p),
                .B($(name(inputs[2]))_m),
                .Y(mmult$(state.id)_out_pm)
            );
        stoch_matrix_mult #(
                .NUM_ROWS($m),
                .NUM_MID($n),
                .NUM_COLS($p)
            ) mmult$(state.id)_mp (
                .CLK(CLK),
                .nRST(nRST),
                .A($(name(inputs[1]))_m),
                .B($(name(inputs[2]))_p),
                .Y(mmult$(state.id)_out_mp)
            );
        stoch_matrix_mult #(
                .NUM_ROWS($m),
                .NUM_MID($n),
                .NUM_COLS($p)
            ) mmult$(state.id)_mm (
                .CLK(CLK),
                .nRST(nRST),
                .A($(name(inputs[1]))_m),
                .B($(name(inputs[2]))_m),
                .Y(mmult$(state.id)_out_mm)
            );
        stoch_sat_sub_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mmult$(state.id)_11 (
                .CLK(CLK),
                .nRST(nRST),
                .A(mmult$(state.id)_out_pp),
                .B(mmult$(state.id)_out_pm),
                .Y(mmult$(state.id)_out_11)
            );
        stoch_sat_sub_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mmult$(state.id)_12 (
                .CLK(CLK),
                .nRST(nRST),
                .A(mmult$(state.id)_out_pm),
                .B(mmult$(state.id)_out_pp),
                .Y(mmult$(state.id)_out_12)
            );
        stoch_sat_sub_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mmult$(state.id)_13 (
                .CLK(CLK),
                .nRST(nRST),
                .A(mmult$(state.id)_out_mp),
                .B(mmult$(state.id)_out_mm),
                .Y(mmult$(state.id)_out_13)
            );
        stoch_sat_sub_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mmult$(state.id)_14 (
                .CLK(CLK),
                .nRST(nRST),
                .A(mmult$(state.id)_out_mm),
                .B(mmult$(state.id)_out_mp),
                .Y(mmult$(state.id)_out_14)
            );
        stoch_add_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mmult$(state.id)_p (
                .CLK(CLK),
                .nRST(nRST),
                .A(mmult$(state.id)_out_11),
                .B(mmult$(state.id)_out_14),
                .Y($(name(outputs[1]))_p)
            );
        stoch_add_mat #(
                .NUM_ROWS($(outsize[1])),
                .NUM_COLS($(outsize[2]))
            ) mmult$(state.id)_m (
                .CLK(CLK),
                .nRST(nRST),
                .A(mmult$(state.id)_out_12),
                .B(mmult$(state.id)_out_13),
                .Y($(name(outputs[1]))_m)
            );
        // END mmult$(state.id)
        \n""")

    return buffer, (id = state.id + 1,)
end