text:
""" Algebras of operads of wiring diagrams.
"""
module WiringDiagramAlgebras
export oapply, query

using Requires

using ...Present, ...Theories
using ...Theories: dom_nums, codom_nums, attr, adom, adom_nums
using ...CategoricalAlgebra
import ...CategoricalAlgebra.CSets: homomorphisms, homomorphism, is_homomorphic
using ..UndirectedWiringDiagrams
using ..UndirectedWiringDiagrams: TheoryUWD

""" Compose morphisms according to UWD.

The morphisms corresponding to the boxes, and optionally also the objects
corresponding to the junctions, are given by dictionaries indexed by
box/junction attributes. The default attributes are those compatible with the
`@relation` macro.
"""
function oapply(composite::UndirectedWiringDiagram, hom_map::AbstractDict,
                ob_map::Union{AbstractDict,Nothing}=nothing;
                hom_attr::Symbol=:name, ob_attr::Symbol=:variable)
  # XXX: Julia should be inferring these vector eltypes but isn't on v1.7.
  homs = valtype(hom_map)[ hom_map[name] for name in composite[hom_attr] ]
  obs = isnothing(ob_map) ? nothing :
    valtype(ob_map)[ ob_map[name] for name in composite[ob_attr] ]
  oapply(composite, homs, obs)
end

# UWD algebras of multi(co)spans
################################

function oapply(composite::UndirectedWiringDiagram,
                spans::AbstractVector{<:Multispan}; Ob=nothing, Hom=nothing)
  @assert nboxes(composite) == length(spans)
  # FIXME: This manual type inference is hacky and bad. The right solution is to
  # extend `Multi(co)span` with type parameters that allow abstract types.
  if isnothing(Ob); Ob = typejoin(mapreduce(typeof∘apex, typejoin, spans),
                                  mapreduce(eltype∘feet, typejoin, spans)) end
  if isnothing(Hom); Hom = mapreduce(eltype∘legs, typejoin, spans) end
  junction_feet = Vector{Ob}(undef, njunctions(composite))

  # Create bipartite free diagram whose vertices of types 1 and 2 are the UWD's
  # boxes and junctions, respectively.
  diagram = BipartiteFreeDiagram{Ob,Hom}()
  add_vertices₁!(diagram, nboxes(composite), ob₁=map(apex, spans))
  add_vertices₂!(diagram, njunctions(composite))
  for (b, span) in zip(boxes(composite), spans)
    for (p, leg) in zip(ports(composite, b), legs(span))
      j = junction(composite, p)
      add_edge!(diagram, b, j, hom=leg)
      foot = codom(leg)
      if isassigned(junction_feet, j)
        foot′ = junction_feet[j]
        foot == foot′ || error("Feet of spans are not equal: $foot != $foot′")
      else
        junction_feet[j] = foot
      end
    end
  end
  all(isassigned(junction_feet, j) for j in junctions(composite)) ||
    error("Limits with isolated junctions are not supported")
  diagram[:ob₂] = junction_feet

  # The composite multispan is given by the limit of this diagram.
  lim = limit(diagram)
  outer_legs = map(junction(composite, outer=true)) do j
    e = first(incident(diagram, j, :tgt))
    legs(lim)[src(diagram, e)] ⋅ hom(diagram, e)
  end
  Multispan(ob(lim), outer_legs)
end

function oapply(composite::UndirectedWiringDiagram,
                cospans::AbstractVector{<:StructuredMulticospan{L}},
                junction_feet::Union{AbstractVector,Nothing}=nothing) where L
  @assert nboxes(composite) == length(cospans)
  if isnothing(junction_feet)
    junction_feet = Vector{first(dom(L))}(undef, njunctions(composite))
  else
    @assert njunctions(composite) == length(junction_feet)
  end

  # Create bipartite free diagram whose vertices of types 1 and 2 are the UWD's
  # junctions and boxes, respectively.
  diagram = BipartiteFreeDiagram{codom(L)...}()
  add_vertices₁!(diagram, njunctions(composite))
  add_vertices₂!(diagram, nboxes(composite), ob₂=map(apex, cospans))
  for (b, cospan) in zip(boxes(composite), cospans)
    for (p, leg, foot) in zip(ports(composite, b), legs(cospan), feet(cospan))
      j = junction(composite, p)
      add_edge!(diagram, j, b, hom=leg)
      if isassigned(junction_feet, j)
        foot′ = junction_feet[j]
        foot == foot′ || error("Feet of cospans are not equal: $foot != $foot′")
      else
        junction_feet[j] = foot
      end
    end
  end
  for (j, foot) in enumerate(junction_feet)
    diagram[j, :ob₁] = L(foot)
  end

  # Find, or if necessary create, an outgoing edge for each junction. The
  # existence of such edges is an assumption for colimits of bipartite diagrams.
  # The edges are also needed to construct inclusions for the outer junctions.
  junction_edges = map(junctions(composite)) do j
    out_edges = incident(diagram, j, :src)
    if isempty(out_edges)
      x = ob₁(diagram, j)
      v = add_vertex₂!(diagram, ob₂=x)
      add_edge!(diagram, j, v, hom=id(x))
    else
      first(out_edges)
    end
  end

  # The composite multicospan is given by the colimit of this diagram.
  colim = colimit(diagram)
  outer_js = junction(composite, outer=true)
  outer_legs = map(junction_edges[outer_js]) do e
    hom(diagram, e) ⋅ legs(colim)[tgt(diagram, e)]
  end
  outer_feet = junction_feet[outer_js]
  StructuredMulticospan{L}(Multicospan(ob(colim), outer_legs), outer_feet)
end

# Queries via UWD algebras
##########################

abstract type DataFrameFallback end
abstract type MaybeDataFrame <: DataFrameFallback end

""" Evaluate a conjunctive query on an attributed C-set.

The conjunctive query is represented as an undirected wiring diagram (UWD) whose
boxes and ports are named via the attributes `:name`, `:port_name`, and
`:outer_port_name`. To define such a diagram, use the [`@relation`](@ref) macro
with its named syntax. Parameters to the query may be passed as a collection of
pairs using the optional third argument.

The result is data table whose columns correspond to the outer ports of the UWD.
By default, a data frame is returned if the package
[`DataFrames.jl`](https://github.com/JuliaData/DataFrames.jl) is loaded;
otherwise, a named tuple is returned. To change this behavior, set the keyword
argument `table_type` to a type or function taking two arguments, a vector of
columns and a vector of column names. There is one exceptional case: if the UWD
has no outer ports, the query is a counting query and the result is a vector
whose length is the number of results.

For its implementation, this function wraps the [`oapply`](@ref) method for
multispans, which defines the UWD algebra of multispans.
"""
function query(X::ACSet, diagram::UndirectedWiringDiagram,
               params=(;); table_type=MaybeDataFrame)
  # For each box in the diagram, extract span from ACSet.
  spans = map(boxes(diagram), subpart(diagram, :name)) do b, name
    apex = FinSet(nparts(X, name))
    legs = map(subpart(diagram, ports(diagram, b), :port_name)) do port_name
      FinDomFunction(X, port_name == :_id ? name : port_name)
    end
    Multispan(apex, legs)
  end

  # Add an extra box and corresponding span for each parameter.
  if !isempty(params)
    diagram = copy(diagram)
    spans = vcat(spans, map_pairs(params) do (key, value)
      box = add_part!(diagram, :Box, name=:_const)
      junction = key isa Integer ? key : incident(diagram, key, :variable)
      add_part!(diagram, :Port, port_name=:_value,
                box=box, junction=junction)
      SMultispan{1}(ConstantFunction(value, FinSet(1)))
    end)
  end

  # Call `oapply` and make a table out of the resulting span.
  outer_span = oapply(diagram, spans, Ob=SetOb, Hom=FinDomFunction{Int})
  if nparts(diagram, :OuterPort) == 0
    fill((;), length(apex(outer_span)))
  else
    columns = map(collect, outer_span)
    names = has_subpart(diagram, :outer_port_name) ?
      subpart(diagram, :outer_port_name) : fill(nothing, length(columns))
    make_table(table_type, columns, names)
  end
end

map_pairs(f, collection) = map(f, collect(pairs(collection)))
map_pairs(f, vec::AbstractVector{<:Pair}) = map(f, vec)

""" Create conjunctive query (a UWD) for finding homomorphisms out of a C-set.

Returns the query together with query parameters for the attributes.
Homomorphisms from `X` to `Y` can be computed by:

```julia
query(Y, homomorphism_query(X)...)
```
"""
function homomorphism_query(X::StructACSet{S}; count::Bool=false) where S
  offsets = cumsum([0; [nparts(X,c) for c in ob(S)]])
  nelems = offsets[end]

  diagram = HomomorphismQueryDiagram{Symbol}()
  params = Pair[]
  add_parts!(diagram, :Junction, nelems)
  if !count
    add_parts!(diagram, :OuterPort, nelems, outer_junction=1:nelems)
  end
  for (i, c) in enumerate(ob(S))
    n = nparts(X,c)
    boxes = add_parts!(diagram, :Box, n, name=c)
    add_parts!(diagram, :Port, n, box=boxes, port_name=:_id,
               junction=(1:n) .+ offsets[i])
  end
  for (f, c, i, j) in zip(hom(S), dom(S), dom_nums(S), codom_nums(S))
    n = nparts(X,c)
    add_parts!(diagram, :Port, n, box=(1:n) .+ offsets[i], port_name=f,
               junction=X[:,f] .+ offsets[j])
  end
  for (f, c, i) in zip(attr(S), adom(S), adom_nums(S))
    n = nparts(X,c)
    junctions = add_parts!(diagram, :Junction, n)
    add_parts!(diagram, :Port, n, box=(1:n) .+ offsets[i], port_name=f,
               junction=junctions)
    append!(params, Iterators.map(Pair, junctions, X[:,f]))
  end
  (diagram, params)
end

@present TheoryHomomorphismQueryDiagram <: TheoryUWD begin
  Name::AttrType
  name::Attr(Box, Name)
  port_name::Attr(Port, Name)
end
@acset_type HomomorphismQueryDiagram(TheoryHomomorphismQueryDiagram,
  index=[:box, :junction, :outer_junction]) <: UndirectedWiringDiagram

function homomorphisms(X::ACSet, Y::ACSet, ::HomomorphismQuery)
  columns = query(Y, homomorphism_query(X)...; table_type=AbstractVector)
  map(row -> make_homomorphism(row, X, Y), eachrow(reduce(hcat, columns)))
end
function homomorphism(X::ACSet, Y::ACSet, ::HomomorphismQuery)
  columns = query(Y, homomorphism_query(X)...; table_type=AbstractVector)
  isempty(first(columns)) ? nothing :
    make_homomorphism(map(first, columns), X, Y)
end
function is_homomorphic(X::ACSet, Y::ACSet, ::HomomorphismQuery)
  length(query(Y, homomorphism_query(X, count=true)...)) > 0
end

function make_homomorphism(row, X::StructACSet{S}, Y::StructACSet{S}) where S
  components = let i = 0
    NamedTuple{ob(S)}([row[i+=1] for _ in parts(X,c)] for c in ob(S))
  end
  ACSetTransformation(components, X, Y)
end

make_table(f, columns, names) = f(columns, names)
make_table(::Type{AbstractVector}, columns, names) = columns
make_table(::Type{NamedTuple}, columns, names) =
  NamedTuple{Tuple(names)}(Tuple(columns))
make_table(::Type{<:DataFrameFallback}, columns, names) =
  make_table(NamedTuple, columns, names)

function __init__()
  @require DataFrames="a93c6f00-e57d-5684-b7b6-d8193f3e46c0" begin
    using .DataFrames: DataFrame
    make_table(::Type{MaybeDataFrame}, columns, names) =
      make_table(DataFrame, columns, names)
  end
end

end


text:
using ROCKS
using Documenter

makedocs(;
    modules = [ROCKS],
    authors = "Daymond Ling",
    repo = "https://github.com/DaymondLing/ROCKS.jl/blob/{commit}{path}#L{line}",
    sitename = "ROCKS.jl",
    format = Documenter.HTML(;
        prettyurls = get(ENV, "CI", "false") == "true",
        canonical = "https://DaymondLing.github.io/ROCKS.jl/stable",
        assets = String[],
    ),
    pages = [
        "Home" => "index.md",
        "User's Guide" => [
            "KS Test" => "man/kstest.md",
            "ROC" => "man/roc.md",
            "Plots and Tables" => "man/bcdiag.md",
        ],
        "Function Reference" => "Reference.md",
    ],
)

deploydocs(; repo = "github.com/DaymondLing/ROCKS.jl.git", devbranch = "main")

ext: jl


text:
#tests for SplineModel.jl module
using Test
include("../src/SplineModel.jl")

#####tests for nrbmak()###############
#based on Demonstration 1
pnts = [0.5 1.5 4.5 3.0 7.5 6.0 8.5;
         3.0 5.5 5.5 1.5 1.5 4.0 4.5;
         0.0 0.0 0.0 0.0 0.0 0.0 0.0]
knot_crv = [0, 0, 0, 1/4, 1/2, 3/4, 3/4, 1, 1, 1];
crv = nrbmak(pnts, [knot_crv])
@test crv.number == [7]
@test crv.coefs == [pnts; ones(1, crv.number[1])]
@test crv.knots == [knot_crv]
@test crv.order == [3]

#based on Demonstration 2
pnts = zeros(3,5,5);
pnts[:,:,1] = [ 0.0  3.0  5.0  8.0 10.0;
             0.0  0.0  0.0  0.0  0.0;
             2.0  2.0  7.0  7.0  8.0];
pnts[:,:,2] = [ 0.0  3.0  5.0  8.0 10.0;
             3.0  3.0  3.0  3.0  3.0;
             0.0  0.0  5.0  5.0  7.0];
pnts[:,:,3] = [ 0.0  3.0  5.0  8.0 10.0;
             5.0  5.0  5.0  5.0  5.0;
             0.0  0.0  5.0  5.0  7.0];
pnts[:,:,4] = [ 0.0  3.0  5.0  8.0 10.0;
             8.0  8.0  8.0  8.0  8.0;
             5.0  5.0  8.0  8.0 10.0];
pnts[:,:,5] = [ 0.0  3.0  5.0  8.0 10.0;
            10.0 10.0 10.0 10.0 10.0;
             5.0  5.0  8.0  8.0 10.0];
uknots = [0, 0, 0, 1/3, 2/3, 1, 1, 1]
vknots = [0, 0, 0, 1/3, 2/3, 1, 1, 1]
knots = [uknots, vknots]
srf = nrbmak(pnts,knots);
@test srf.number == [5, 5]
@test srf.coefs == [pnts; ones(1,5,5)]
@test srf.knots == [uknots, vknots]
@test srf.order == [3, 3]

#based on Demonstration 3
coefs =[ 6.0  0.0  6.0  1;
        -5.5  0.5  5.5  1;
        -5.0  1.0 -5.0  1;
        4.5  1.5 -4.5  1;
        4.0  2.0  4.0  1;
        -3.5  2.5  3.5  1;
        -3.0  3.0 -3.0  1;
        2.5  3.5 -2.5  1;
        2.0  4.0  2.0  1;
        -1.5  4.5  1.5  1;
        -1.0  5.0 -1.0  1;
        0.5  5.5 -0.5  1;
        0.0  6.0  0.0  1]';
knots = [0, 0, 0, 0, .1, .2, .3, .4, .5, .6, .7, .8, .9, 1, 1, 1, 1];

crv3d = nrbmak(coefs,[knots]);
@test crv3d.number == [13]
@test crv3d.coefs == coefs
@test crv3d.knots ==[knots]
@test crv3d.order == [4]

##############tests for findspan()###############

n = 3
U = [0, 0, 0, 1/2, 1, 1, 1]
u = LinRange(0, 1, 10)
s = findspan(n, u, U)
@test s == [2*ones(5); 3*ones(5)]

p = 2
m = 7
n = m - p - 1
U = [zeros(p); LinRange(0,1,m+1-2*p); ones(p)]
u = [0, 0.11880, 0.55118, 0.93141, 0.40068, 0.35492, 0.44392, 0.88360, 0.35414, 0.92186, 0.83085, 1]
s = [2, 2, 3, 4, 3, 3, 3, 4, 3, 4, 4, 4]
@test findspan(n, u, U)==s

#############test for basisfun()##################

n = 3
U = [0, 0, 0, 1/2, 1, 1, 1]
p = 2;
u = LinRange(0, 1, 10);
s = findspan(n, u, U);
Bref = [1.000000000000000                   0                   0
        0.604938271604938   0.370370370370370   0.024691358024691
        0.308641975308642   0.592592592592593   0.098765432098765
        0.111111111111111   0.666666666666667   0.222222222222222
        0.012345679012346   0.592592592592593   0.395061728395062
        0.395061728395062   0.592592592592593   0.012345679012346
        0.222222222222222   0.666666666666667   0.111111111111111
        0.098765432098765   0.592592592592593   0.308641975308642
        0.024691358024691   0.370370370370370   0.604938271604938
        0                   0   1.000000000000000]
B = basisfun(s, u, p, U)
@test B≈Bref
############test for nrbeval()###################
uknots = [0, 0, 0, 1, 1, 1]
vknots = [0, 0, 0, .5, 1, 1, 1]
wknots = [0, 0, 0, 0, 1, 1, 1, 1]
knots = [uknots, vknots, wknots]
cx = [0 0.5 1]
nx = length(cx)
cy = [0 0.25 0.75 1]
ny = length(cy)
cz = [0 1/3 2/3 1]
nz = length(cz)
coefs = zeros(4, nx, ny, nz)
coefs[1,:,:,:] = repeat(reshape(cx,nx,1,1),1, ny, nz)
coefs[2,:,:,:] = repeat(reshape(cy,1,ny,1), nx, 1, nz)
coefs[3,:,:,:] = repeat(reshape(cz,1,1,nz), nx, ny, 1)
coefs[4,:,:,:] .= 1
nurbs = nrbmak(coefs, knots)
x = rand(5,1)
y = rand(5,1)
z = rand(5,1)
tt = [x y z]'
points = nrbeval(nurbs,tt)
@test points ≈ tt

uknots = [0, 0, 0, 1, 1, 1]
vknots = [0, 0, 0, 0, 1, 1, 1, 1]
wknots = [0, 0, 1, 1]
knots = [uknots, vknots, wknots]
cx = [0 0 1]
nx = length(cx)
cy = [0 0 0 1]
ny = length(cy)
cz = [0 1]
nz = length(cz)
coefs = zeros(4, nx, ny, nz)
coefs[1,:,:,:] = repeat(reshape(cx,nx,1,1), 1, ny, nz);
coefs[2,:,:,:] = repeat(reshape(cy,1,ny,1), nx, 1, nz);
coefs[3,:,:,:] = repeat(reshape(cz,1,1,nz), nx, ny, 1);
coefs[4,:,:,:] .= 1
nurbs = nrbmak(coefs, knots)
x = rand(5,1)
y = rand(5,1)
z = rand(5,1)
tt = [x y z]'
points = nrbeval(nurbs,tt)
@test points ≈ [x.^2 y.^3 z]'

uknots = [0, 0, 0, 1, 1, 1]
vknots = [0, 0, 0, 0, 1, 1, 1, 1]
wknots = [0, 0, 1, 1]
knots = [uknots, vknots, wknots]
cx = [0 0 1]
nx = length(cx)
cy = [0 0 0 1]
ny = length(cy)
cz = [0 1]
nz = length(cz)
coefs = zeros(4, nx, ny, nz)
coefs[1,:,:,:] = repeat(reshape(cx,nx,1,1), 1, ny, nz)
coefs[2,:,:,:] = repeat(reshape(cy,1,ny,1), nx, 1, nz)
coefs[3,:,:,:] = repeat(reshape(cz,1,1,nz), nx, ny, 1)
coefs[4,:,:,:] .= 1
coefs = coefs[[2, 1, 3, 4],:,:,:]
nurbs = nrbmak(coefs, knots)
x = rand(5,1)
y = rand(5,1)
z = rand(5,1)
tt = [x y z]'
points = nrbeval(nurbs,tt)
[y.^3 x.^2 z]'
@test points ≈ [y.^3 x.^2 z]'

#test kntrefine()
knots = [[0, 0, 1, 1], [0, 0, 0, 1, 1, 1]]
coefs = zeros(4, 2, 3)
coefs[1,:,:] = [1 sqrt(2)/2 0; 2 sqrt(2) 0]
coefs[2,:,:] = [0 sqrt(2)/2 1; 0 sqrt(2) 2]
coefs[4,:,:] = [1 sqrt(2)/2 1; 1 sqrt(2)/2 1]
nrbs = nrbmak(coefs, knots)
nrbs = nrbkntins(nrbs, [[], [0.5, 0.6, 0.6]])
nrbs = nrbdegelev(nrbs, [0, 1])
nrbs = nrbkntins(nrbs, [[], [0.4]])
rknots, _, _ = kntrefine(nrbs.knots, [1, 1], [1, 1], [0, 0])
@test (rknots[1] == [0, 0, 0.5, 1, 1])
@test (rknots[2] == [0, 0, 0.2, 0.4, 0.45, 0.5, 0.55, 0.6, 0.8, 1, 1])

rknots, _, _ = kntrefine(nrbs.knots, [1, 1], [3, 3], [0, 0])
@test (rknots[1] == [0, 0, 0, 0, 0.5, 0.5, 0.5, 1, 1, 1, 1])
@test (rknots[2] == [0, 0, 0, 0, 0.2, 0.2, 0.2, 0.4, 0.4, 0.4, 0.45, 0.45, 0.45, 0.5,
                0.5, 0.5, 0.55, 0.55, 0.55, 0.6, 0.6, 0.6, 0.8, 0.8, 0.8, 1, 1, 1, 1])

rknots, _, _ = kntrefine(nrbs.knots, [1, 1], [3, 3], [2, 2])
@test (rknots[1] == [0, 0, 0, 0, 0.5, 1, 1, 1, 1])
@test (rknots[2] == [0, 0, 0, 0, 0.2, 0.4, 0.45, 0.5, 0.5, 0.55, 0.6, 0.6, 0.6, 0.8, 1, 1, 1, 1])

#test nrbsquare()
srf = nrbsquare([], 1, 2, 2, 4)
@test (srf.order == [3, 3])
knt = [0, 0, 0, 1/4, 1/2, 3/4, 1, 1, 1]
@test (srf.knots == [knt, knt])
x = LinRange(0, 1, 100)
X = [i for i in x, j in x]
Y = [j for i in x, j in x]
vals = nrbeval(srf, [collect(x), collect(x)])
@test (vals[1,:,:] ≈ X)
@test (vals[2,:,:] ≈ 2*Y)


text:
@testset "Graph size queries" begin

    # Helper methods
    nt(name) = named(name)
    tf(name) = t -> nt(name)(t)
    iv(size, name="in") = inputvertex(name, size)
    tv(in, name) = invariantvertex(identity, in, traitdecoration=tf(name))
    av(in, outsize, name) = absorbvertex(IndMem(MatMul(nout(in), outsize)), in, traitdecoration=tf(name))
    cc(ins...; name) = conc(ins...; dims=2, traitdecoration=tf(name), outwrap=x -> IndMem(x, collect(nout.(ins)), sum(nout, ins)))

    nc(name) = traitconf(nt(name))

    @testset "all_in_Δsize_graph" begin
        using NaiveNASlib: TightΔSizeGraph, LooseΔSizeGraph, all_in_Δsize_graph, Input, Output

        @testset "Concat and elemwise not connected" begin
            v1 = iv(3, "v1")
            v2a = av(v1, 2, "v2a")
            v2b = av(v1, 4, "v2b")
            v3 = cc(v2a, v2b; name="v3")
            v4a = av(v3, 2, "v4a")
            v4b = av(v3, 2, "v4b")
            v5 = "v5" >> v4a + v4b
            v6 = av(v5, 3, "v6")

            @test name.(all_in_Δsize_graph(LooseΔSizeGraph(), v2a, Output())) == ["v2a", "v3", "v2b", "v4a", "v4b"]
            @test name.(all_in_Δsize_graph(TightΔSizeGraph(), v2a, Output())) == ["v2a", "v3", "v4a", "v4b"]

            @test name.(all_in_Δsize_graph(LooseΔSizeGraph(), v2a, Input())) == ["v2a", "v1"]
            @test name.(all_in_Δsize_graph(TightΔSizeGraph(), v2a, Input())) == ["v2a", "v1"]

            @test name.(all_in_Δsize_graph(LooseΔSizeGraph(), v3, Output())) == ["v3", "v2a", "v2b", "v4a", "v4b"]
            @test name.(all_in_Δsize_graph(TightΔSizeGraph(), v3, Output())) == ["v3", "v2a", "v4a", "v4b", "v2b"]

            @test name.(all_in_Δsize_graph(LooseΔSizeGraph(), v3, Input())) == ["v3", "v2a", "v2b", "v4a", "v4b"]
            # When we start from v3, at least one of v2a and v2b must change, and we decide to add both in the set, perhaps we should just pick one?
            @test name.(all_in_Δsize_graph(TightΔSizeGraph(), v3, Input())) == ["v3", "v2a", "v2b", "v4a", "v4b"]

            @test name.(all_in_Δsize_graph(LooseΔSizeGraph(), v4a, Output())) == ["v4a", "v5", "v4b", "v6"]
            @test name.(all_in_Δsize_graph(TightΔSizeGraph(), v4a, Output())) == ["v4a", "v5", "v4b", "v6"]

            @test name.(all_in_Δsize_graph(LooseΔSizeGraph(), v4a, Input())) == ["v4a", "v3", "v2a", "v2b", "v4b"]
            @test name.(all_in_Δsize_graph(TightΔSizeGraph(), v4a, Input())) == ["v4a", "v3", "v2a", "v4b", "v2b"]

            @test name.(all_in_Δsize_graph(LooseΔSizeGraph(), v5, Output())) == ["v5", "v4a", "v4b", "v6"]
            @test name.(all_in_Δsize_graph(TightΔSizeGraph(), v5, Output())) == ["v5", "v4a", "v4b", "v6"]

            @test name.(all_in_Δsize_graph(LooseΔSizeGraph(), v5, Input())) == ["v5", "v4a", "v4b", "v6"]
            @test name.(all_in_Δsize_graph(TightΔSizeGraph(), v5, Input())) == ["v5", "v4a", "v4b", "v6"]
        end

        @testset "Concat and elemwise connected" begin
            v1 = iv(6, "v1")
            v2 = av(v1, 2, "v2")
            v3a = av(v2, 2, "v3a")
            v3b = av(v2, 4, "v3b")
            v4 = cc(v3a, v3b; name="v4")
            v5 = "v5" >> v4 + v1
            v6 = av(v5, 3, "v6")

            @test name.(all_in_Δsize_graph(LooseΔSizeGraph(), v3a, Output())) == ["v3a", "v4", "v3b", "v5", "v1", "v6"]
            @test name.(all_in_Δsize_graph(TightΔSizeGraph(), v3a, Output())) == ["v3a", "v4", "v5", "v1", "v6"]

            @test name.(all_in_Δsize_graph(LooseΔSizeGraph(), v4, Output())) == ["v4", "v3a", "v3b", "v5", "v1", "v6"]
            # When we start from v4, at least one of v3a and v3b must change, and we decide to add both in the set, perhaps we should just pick one?
            @test name.(all_in_Δsize_graph(TightΔSizeGraph(), v4, Output())) == ["v4", "v3a", "v5", "v1", "v6", "v3b"]
        end
    end

    @testset "SizeCycleDetector" begin
        using NaiveNASlib: isinsizecycle

        @testset "Simple cycle$label" for (label, inspre, inspost) in (
            ("", identity, identity),
            (" transparent before", v -> tv(v, "tvpre"), identity),
            (" transparent after", identity, v -> tv(v, "tvpost")),
            (" transparent before and after", v -> tv(v, "tvpre"),v -> tv(v, "tvpost")),
            (" concat before", v -> cc(v; name="ccpre"), identity),
            (" concat after", identity, v -> cc(v;name="ccpost")),
            (" concat before and after", v -> cc(v; name="ccpre"), v -> cc(v;name="ccpost")),
        )
            # Canonical size cycle: if we remove v3b we get the constraint that nout(v2) == nin(v5) == nout(v4) == 2 * nout(v2)
            # Also, if we remove v2 then v5 sees v1 through v4 too and we get the impossible size constraint
            # that nout(v1) == nin(v5) == nout(v4) == nout(v1) + nout(v3b) (where nout(v) > 0 ∀ v).
            v1 = iv(6, "v1")
            v2 = av(v1, 2, "v2")
            v3a = tv(v2, "v3a")
            v3b = av(inspre(v2), 4, "v3b")
            v4 = cc(v3a, inspost(v3b); name="v4")
            v5 = "v5" >> v4 + v1
            v6 = av(v5, 3, "v6")

            @testset "Vertex $(name(v))" for v in all_in_graph(v1)
                if v === v3b || v === v2
                    @test isinsizecycle(v) == true
                else
                    @test isinsizecycle(v) == false
                end
            end
        end

        @testset "Simple cycle with inputvertices" begin
            v1 = iv(6, "v1")
            v2 = iv(2, "v2")
            v3a = tv(v2, "v3a")
            v3b = av(v2, 4, "v3b")
            v4 = cc(v3a, v3b; name="v4")
            v5 = "v5" >> v4 + v1
            v6 = av(v5, 3, "v6")

            @testset "Vertex $(name(v))" for v in all_in_graph(v1)
                if v === v3b # v2 is not in sizecycle as it does not have ancestors
                    @test isinsizecycle(v) == true
                else
                    @test isinsizecycle(v) == false
                end
            end
        end

        @testset "Double absorbing in path -> no cycle" begin
            # Same as above, but now there are two size absorbing vertices in the non-transparent path
            # v3b1 does not see v5 as v3b2 blocks the line of sight and v3b2 does not see v2 as v3b1 block the line of sight
            v1 = iv(6, "v1")
            v2 = av(v1, 2, "v2")
            v3a = tv(v2, "v3a")
            v3b1 = av(v2, 2, "v3b1")
            v3b2 = av(v3b1, 4, "v3b2")
            v4 = cc(v3a, v3b2; name="v4")
            v5 = "v5" >> v4 + v1
            v6 = av(v5, 3, "v6")

            @testset "Vertex $(name(v))" for v in all_in_graph(v1)
                if v === v2
                    @test isinsizecycle(v) == true
                else
                    @test isinsizecycle(v) == false
                end
            end
        end

        @testset "Independent concats" begin
            # Counterexample to Simple cycle. Not a size cycle as the two paths go through different concatenations
            v1 = iv(2, "v1")
            v2 = av(v1, 2, "v2")
            v3a = tv(v2, "v3a")
            v3b = av(v2, 2, "v3b")
            v4a = cc(v3a; name="v4a")
            v4b = cc(v3b; name="v4b")
            v5 = "v5" >> v4a + v4b
            v6 = av(v5, 3, "v6")

            @testset "Vertex $(name(v))" for v in all_in_graph(v1)
                @test isinsizecycle(v) == false
            end
        end

        @testset "Concat and elemwise connected no cycle" begin
            v1 = iv(6, "v1")
            v2 = av(v1, 2, "v2")
            v3a = av(v2, 2, "v3a")
            v3b = av(v2, 4, "v3b")
            v4 = cc(v3a, v3b; name="v4")
            v5 = "v5" >> v4 + v1
            v6 = av(v5, 3, "v6")

            @testset "Vertex $(name(v))" for v in all_in_graph(v1)
                @test isinsizecycle(v) == false
            end

            @test remove!(v3a) == true
            @test isinsizecycle(v3b) == true
        end

        @testset "Concat and elemwise connected in cycle" begin
            v1 = iv(2, "v1")
            v2 = av(v1, 6, "v2")
            v3a = av(v2, 2, "v3a")
            v3b1 = av(v2, 3, "v3b1")
            v3b2 = av(v3b1, 4, "v3b2")
            v4 = cc(v3a, v3b2; name="v4")
            v5 = "v5" >> v4 + v2
            v6 = av(v5, 3, "v6")

            @testset "Vertex $(name(v))" for v in all_in_graph(v1)
                if v === v3a
                    @test isinsizecycle(v) == true
                else
                    @test isinsizecycle(v) == false
                end
            end
        end

        @testset "Size cycle when adding edge" begin
            v1 = av(iv(3, "in"), 3, "v1")
            v2 = av(v1, 5, "v2")
            v3 = tv(v2, "v3")
            ve = tv(v1, "ve")
            v4 = cc(v3; name="v4")
            v5 = "v5" >> v4 + v2

            @testset "Vertex $(name(v))" for v in all_in_graph(v1)
                @test isinsizecycle(v) == false
            end

            # Note: this is basically the default when adding edges to SizeStack
            # We first add it, and then we try to align the sizes
            # In this case alignment becomes impossible due to the created size cycle
            create_edge!(ve, v4; strategy=NoSizeChange())

            @testset "Vertex $(name(v)) after edge" for v in all_in_graph(v1)
                if v === v2 || v === v3
                    @test isinsizecycle(v) == true
                else
                    @test isinsizecycle(v) == false
                end
            end
        end
    end
end


text:
function print_alr_path{T<:Number}(gamma_list::Vector,Z_list::Array{T,3})
  @assert length(gamma_list) == size(Z_list,1)
  for i in 1:length(gamma_list)
    Z = Z_list[i,:,:]
    num_zeros = length(Z)-length(find(Z))
    L0_Z = num_nonzeros = length(find(Z))
    L1_Z = sumabs(Z)
    L2_Z = sum(abs2,Z)
    #rank_Z = rank(Z)
    println(i, "  ", gamma_list[i], ": ", L0_Z, "  ", L1_Z, " ", L2_Z)# " rank= ", rank_Z, " Z=",Z)
  end
end

function print_alr_path{T<:Number}(gamma_list::Vector,Z_list::Array{T,3},X::Array{T,2},Y::Array{T,2})
  @assert length(gamma_list) == size(Z_list,1)
  for i in 1:length(gamma_list)
    Z = reshape(Z_list[i,:,:],(size(Z_list,2),size(Z_list,3)))
    num_zeros = length(Z)-length(find(Z))
    L0_Z = num_nonzeros = length(find(Z))
    L1_Z = sum(abs,Z)
    L2_Z = sum(abs2,Z)
    #rank_Z = rank(Z)
    chisq = sum(abs2,Y-X*Z)
    println(i, "  ", gamma_list[i], ": ", L0_Z, "  ", L1_Z, " ", L2_Z, " chi^2= ",chisq)# " rank= ", rank_Z, " Z=",Z)
  end
end

function compute_alr_path_stats{T<:Number}(Z_list::Array{T,3}; target=nothing)
  num_it = size(Z_list,1)
  L0_Z_list = zeros(T,num_it)
  L1_Z_list = zeros(T,num_it)
  L2_Z_list = zeros(T,num_it)
  rmsdelta_list = zeros(T,num_it)
  for i in 1:num_it
    Z =  reshape(Z_list[i,:,:],(size(Z_list,2),size(Z_list,3)))
    num_zeros = length(Z)-length(find(Z))
    L0_Z_list[i] = num_nonzeros = length(find(Z))
    L1_Z_list[i] = sum(abs,Z)
    L2_Z_list[i] = sum(abs2,Z)
    if target != nothing
      rmsdelta_list[i] = sqrt(sum(abs2,Z.-target)/length(target))
    end
  end
  return(L0_Z_list,L1_Z_list,L2_Z_list,rmsdelta_list)
end

function compute_alr_path_chisq!{T<:Number}(Z_list::Array{T,3},X::Array{T,2},Y::Array{T,2})
  num_it = size(Z_list,1)
  chisq_list = zeros(T,num_it)
  chisq_refit_list = zeros(T,num_it)
  for i in 1:num_it
    Z =  reshape(Z_list[i,:,:],(size(Z_list,2),size(Z_list,3)))
    chisq_list[i] = sum(abs2,Y-X*Z)
    (Zp, chisq_refit) = refit_arp_slr_2d(X,Y,Z)
    Z_list[i,:,:] = Zp
    chisq_refit_list[i] = chisq_refit
  end
  return (chisq_list, chisq_refit_list)
end

function compute_alr_path_chisq!{T<:Number}(Z_list::Array{T,3}, X::Array{T,2}, Y::Array{T,2}, Xcv::Array{T,2}, Ycv::Array{T,2} )
  num_it = size(Z_list,1)
  chisq_list = zeros(T,num_it)
  chisq_refit_list = zeros(T,num_it)
  chisq_cv_list = zeros(T,num_it)
  for i in 1:num_it
    Z =  reshape(Z_list[i,:,:],(size(Z_list,2),size(Z_list,3)))
    (Zp, chisq_refit, chisq_cv) = refit_arp_slr_2d(X,Y,Z,Xcv,Ycv)
    Z_list[i,:,:] = Zp
    chisq_refit_list[i] = chisq_refit
    chisq_cv_list[i] = chisq_cv
  end
  return (chisq_list, chisq_refit_list, chisq_cv_list)
end

max_stars_repo_name: eford/RVSpectraKitLearn

max_issues_repo_name: eford/RVSpectraKitLearn

max_forks_repo_name: eford/RVSpectraKitLearn


text:
# Example Julia script with typical workflow to populate a MESH2UMLS database
# table relating all concets associated with all MESH terms in the input database
# Date: September 6, 2016
# Authors: Isabel Restrepo
# BCBI - Brown University
# Version: Julia 0.4.5

using SQLite
using MySQL
using BioMedQuery.Processes
using BioMedQuery.UMLS

user = ENV["UMLS_USER"]
psswd = ENV["UMLS_PSSWD"]
credentials = Credentials(user, psswd)
append = false
db = nothing

#Database backend
using_sqlite=false
using_mysql=true

if using_sqlite
    #************************ LOCALS TO CONFIGURE!!!! **************************
    db_path="./pubmed_obesity_2010_2012.db"
    #***************************************************************************
    db = SQLite.DB(db_path)
elseif using_mysql
    #************************ LOCALS TO CONFIGURE!!!! **************************
    host="localhost" #If want to hide - use enviroment variables instead
    mysql_usr="root"
    mysql_pswd=""
    dbname="pubmed_obesity_2010_2012"
    #***************************************************************************
    db = mysql_connect(host, mysql_usr, mysql_pswd, dbname)
else
    error("Unsupported database backend")
end

@time begin
    map_mesh_to_umls_async!(db, credentials; append_results=append)
end

println("-------------------------------------------------------------")
println("Done Mapping Mesh to UMLS")
println("-------------------------------------------------------------")


text:
using OrdinaryDiffEq, Flux, Optim, Random, Plots
using Zygote
using ForwardDiff
using LinearAlgebra
using Statistics
using ProgressBars, Printf
using Flux.Optimise: update!
using Flux.Losses: mae, mse
using BSON: @save, @load
using DiffEqSensitivity

Random.seed!(1234);

# Argments
is_restart = false;
n_epoch = 1000000;
n_plot = 1;

n_exp_train = 20;
n_exp_val = 10;
n_exp = n_exp_train + n_exp_val;

noise = 1.f-3;

ntotal = 300;
batch_min = 32;
tstep = 5 / ntotal;

ns = 7;
ns_ = 12;
nr = 12;

alg = AutoTsit5(TRBDF2(autodiff=false));
atol = 1e-5;
rtol = 1e-2;
lb = atol;
ub = 1.f2;

opt = Flux.Optimiser(ExpDecay(5e-3, 0.5, 100 * n_exp_train, 1e-5),
                     ADAMW(0.005, (0.9, 0.999), 1.f-6));

q = 4
K1 = 0.52
A = 4
N = 1
J0 = 2.5
phi = 0.1

function trueODEfunc(dsdt, s, k, t)
    r1 = k[1] * s[1] * s[6] / (1 + (s[6] / K1)^q)
    r2 = k[2] * s[2] * (N - s[5])
    r3 = k[3] * s[3] * (A - s[6])
    r4 = k[4] * s[4] * s[5]
    r5 = k[5] * s[6]
    r6 = k[6] * s[2] * s[5]
    r7 = 13 * s[7]
    r8 = 13 * (s[4] - s[7])

    dsdt[1] = J0 - r1
    dsdt[2] = 2 * r1 - r2 - r6
    dsdt[3] = r2 - r3
    dsdt[4] = r3 - r4 - r8
    dsdt[5] = r2 - r4 - r6
    dsdt[6] = -2 * r1 + 2 * r3 - r5
    dsdt[7] = phi * r8 - r7
end

# Generate data sets
u0_list = rand(Float32, (n_exp, ns));
ic_lb = [0.15, 1.19, 0.04, 0.10, 0.08, 0.14, 0.05];
ic_ub = [1.60, 2.16, 0.20, 0.35, 0.30, 2.67, 0.10];
for i in 1:ns
    u0_list[:, i] .= rand(n_exp) .* (ic_ub[i] - ic_lb[i]) .+ ic_lb[i]
end

tspan = Float32[0.0, ntotal * tstep];
tsteps = range(tspan[1], tspan[2], length=ntotal);
k = [100, 6, 16, 100, 1.28, 12];
ode_data_list = zeros(Float32, (n_exp, ns, ntotal));
yscale_list = [];

u0 = u0_list[1, :];
prob_trueode = ODEProblem(trueODEfunc, u0, tspan, k);
ode_data = Array(solve(prob_trueode, alg, saveat=tsteps));

function max_min(ode_data)
    return maximum(ode_data, dims=2) .- minimum(ode_data, dims=2) .+ lb
end

for i = 1:n_exp
    u0 = u0_list[i, 1:ns]
    prob_trueode = ODEProblem(trueODEfunc, u0, tspan, k)
    ode_data = Array(solve(prob_trueode, alg, saveat=tsteps))

    ode_data += randn(size(ode_data)) .* ode_data .* noise
    ode_data_list[i, :, :] = ode_data

    push!(yscale_list, std(ode_data, dims=2))
end

y_scale = maximum(hcat(yscale_list...), dims=2) .+ lb;
show(stdout, "text/plain", round.(y_scale', digits=8))

np = nr * (ns_ + 1) + ns + 1;
pcrnn = (rand(Float32, np) .- 0.5) * 2 * sqrt(6 / (ns_ + nr));
pcrnn[end] = 0.1

function p2vec(p)
    slope = p[end] * 100
    w_b = p[1:nr] .* slope
    w_out = reshape(@view(p[nr + 1:nr * (ns_ + 1)]), ns_, nr)
    w_in = clamp.(-w_out, 0.f0, 4.f0)
    w_J = p[nr * (ns_ + 1) + 1:np - 1]
    return w_in, w_b, w_out, w_J
end

function display_p(p)
    w_in, w_b, w_out, w_J = p2vec(p)
    println("\nspecies (column) reaction (row) w_out | w_b")
    w_in_ = vcat(w_out, exp.(w_b'), )
    show(stdout, "text/plain", round.(w_in_', digits=2))

    println("\nw_J")
    show(stdout, "text/plain", round.(w_J', digits=2))
end
display_p(pcrnn)

node = ns_ - ns;
dudt2 = Chain(x -> x,
              Dense(ns, node, gelu),
              Dense(node, node, gelu),
              Dense(node, node, gelu),
              Dense(node, ns_ - ns, softplus));
pnn, re = Flux.destructure(dudt2);
rep = re(pnn);
p = vcat(pcrnn, pnn);

function crnn(du, u, pcrnn, t)
    u_ = vcat(u, rep(u))
    w_in_x = w_in' * @. log(clamp(u_, lb, ub))
    du .= (w_out * (@. exp(w_in_x + w_b)))[1:ns] .+ w_J
end

i_exp = 1;
u0 = u0_list[i_exp, :];
prob = ODEProblem(crnn, u0, tspan, saveat=tsteps, atol=atol, rtol=rtol)

sense = BacksolveAdjoint(checkpointing=true; autojacvec=ZygoteVJP());
function predict_neuralode(u0, p, batch=ntotal)
    pcrnn = p[1:np]
    pnn = p[np + 1:end]
    global w_in, w_b, w_out, w_J = p2vec(pcrnn)
    global rep = re(pnn)
    prob_ = remake(prob, u0=u0, p=pcrnn)
    pred = clamp.(Array(solve(prob_, alg, saveat=tsteps[1:batch], sensalg=sense)), lb, ub)
    return pred
end
predict_neuralode(u0, p)

function loss_neuralode(p, i_exp, batch=ntotal)
    pred = predict_neuralode(u0_list[i_exp, :], p, batch)
    ode_data = clamp.(ode_data_list[i_exp, 1:ns, 1:size(pred)[2]], lb, ub)
    loss = mae(ode_data ./ y_scale, pred[1:ns, :] ./ y_scale)
    return loss
end
loss_neuralode(p, 1)

# Callback function to observe training
cbi = function (p, i_exp)
    ode_data = ode_data_list[i_exp, :, :]
    pred = predict_neuralode(u0_list[i_exp, :], p)
    rep = re(pnn)
    pred_ = rep(pred)
    l_plt = []
    for i = 1:ns
        plt = scatter(
            tsteps,
            ode_data[i, :],
            markercolor=:transparent
        )
        plot!(plt, tsteps, pred[i, :], lw=3)
        xlabel!(plt, "Time (minutes)")
        ylabel!(plt, "S$i (mM)")
        push!(l_plt, plt)
    end
    for i = 1:ns_ - ns
        plt = plot(tsteps, pred_[i, :], lw=3)
        xlabel!(plt, "Time (minutes)")
        ylabel!(plt, "vS_$i (mM)")
        push!(l_plt, plt)
    end
    plt_all = plot(l_plt..., legend=false, size=(800, 800))
    png(plt_all, string("figs/i_exp_", i_exp))

    return false
end
cbi(p, 1)

l_loss_train = []
l_loss_val = []
l_grad = []
iter = 1
cb = function (p, loss_train, loss_val, g_norm)
    global l_loss_train, l_loss_val, l_grad, iter
    push!(l_loss_train, loss_train)
    push!(l_loss_val, loss_val)
    push!(l_grad, g_norm)

    if iter % n_plot == 0
        display_p(p)
        l_exp = [1, n_exp]
        println("\n plot ", l_exp)
        for i_exp in l_exp
            cbi(p, i_exp)
        end

        plt_loss = plot(l_loss_train, xscale=:identity, yscale=:log10, label="Training")
        plot!(plt_loss, l_loss_val, xscale=:identity, yscale=:log10, label="Validation")
        plt_grad = plot(l_grad, xscale=:identity, yscale=:log10, label="grad_norm")
        xlabel!(plt_loss, "Epoch")
        xlabel!(plt_grad, "Epoch")
        ylabel!(plt_loss, "Loss")
        ylabel!(plt_grad, "Gradient Norm")
        # ylims!(plt_loss, (-Inf, 1e0))
        plt_all = plot([plt_loss, plt_grad]..., legend=:top)
        png(plt_all, "figs/loss_grad")

        @save "./checkpoint/mymodel.bson" p opt l_loss_train l_loss_val l_grad iter
    end
    iter += 1
end

if is_restart
    @load "./checkpoint/mymodel.bson" p opt l_loss_train l_loss_val l_grad iter
    iter += 1
    # opt = ADAMW(0.0005, (0.9, 0.999), 1.f-6);
end

epochs = ProgressBar(iter:n_epoch);
loss_epoch = zeros(Float32, n_exp);
grad_norm = zeros(Float32, n_exp_train);
for epoch in epochs
    global p
    for i_exp in randperm(n_exp_train)
        batch = rand(batch_min:ntotal)
        grad = ForwardDiff.gradient(x -> loss_neuralode(x, i_exp, batch), p);
        # loss, back = Zygote.pullback(x -> loss_neuralode(x, i_exp, batch), p)
        # grad = back(one(loss))[1]
        grad_norm[i_exp] = norm(grad, 2)
        update!(opt, p, grad)
    end
    for i_exp in 1:n_exp
        loss_epoch[i_exp] = loss_neuralode(p, i_exp)
    end
    loss_train = mean(loss_epoch[1:n_exp_train]);
    loss_val = mean(loss_epoch[n_exp_train + 1:end]);
    set_description(epochs,
                    string(@sprintf("Loss train %.4e val %.4e lr %.1e", loss_train, loss_val, opt[1].eta)))
    cb(p, loss_train, loss_val, mean(grad_norm));
end

@printf("min loss train %.4e val %.4e\n", minimum(l_loss_train), minimum(l_loss_val))


text:
using Knet
using AutoGrad
import AutoGrad: gcheck
using Taarruz
using Test

@testset "mnist" begin
    include(Knet.dir("data", "mnist.jl"))
    global dtrn, dtst = mnistdata()
    global atype = dtrn.xtype
    @test dtrn.length == 60000
    @test dtst.length == 10000
    @test dtrn.xsize[1:end-1] == dtst.xsize[1:end-1] == (28, 28, 1)
    @test length(dtrn) == div(dtrn.length, dtrn.batchsize)
    @test length(dtst) == div(dtst.length, dtst.batchsize)
end

@testset "linear" begin
    global linear = Taarruz.Chain(
        (Taarruz.Dense(784,10,identity; atype=atype), ))
    x, y = first(dtst); x = param(x)
    @test gcheck(linear, x, y; atol=0.05)
    @test typeof(linear) == Taarruz.Chain
    @test size(linear(x)) == (10, dtst.batchsize)
    @test accuracy(linear, dtst) < 0.50
    progress!(adam(linear, dtrn))
    global tstacc = accuracy(linear, dtst)
    @test tstacc > 0.75
end

@testset "lenet" begin
    lenet = Lenet(; atype=atype)
    x, y = first(dtst); x = param(x)
    @test gcheck(lenet, x, y; atol=0.05)
    @test typeof(lenet) == Taarruz.Chain
    @test size(lenet(x)) == (10, dtst.batchsize)
end

@testset "fgsm" begin
    x, y = first(dtst)

    ϵ = 0.2
    x̂s = FGSM(linear, ϵ, x, y)
    x̂ = x̂s[1]
    @test length(x̂s) == 1
    @test size(x̂) == size(x)
    @test typeof(x̂) == typeof(x)
    @test maximum(x̂) == 1
    @test minimum(x̂) == 0

    example(x,y,ϵ=0.2,f=linear) = FGSM(f,ϵ,x,y)[1]
    abuse(x,y,ϵ=0.2,f=linear; o...) = accuracy(f(example(x,y)), y; o...)
    abuse(d::Knet.Data, ϵ=0.2, f=linear) = sum(abuse(x,y,ϵ,f; average=false) for (x,y) in d) / d.length

    fgsmacc = abuse(dtst)
    @test tstacc - fgsmacc > 0.2

    fgsmacc = map(ϵi->abuse(x,y,ϵi), 0.1:0.05:0.5)
    @test issorted(fgsmacc, rev=true)

    ϵ = 0.2
    xp = param(x)
    J = @diff linear(xp, y)
    ∇x = grad(J, xp)
    x̂ = min.(1, max.(0, x + ϵ * sign.(∇x)))
    @test length(Taarruz.∇xJ(linear, x, y)) == 1
    @test ∇x ≈ Taarruz.∇xJ(linear, x, y)[1]
    @test x̂ ≈ FGSM(x, ∇x, ϵ)
    @test x̂ ≈ FGSM(linear, ϵ, x, y)[1]
end


text:
@var container;
@var camera, scene, renderer;
@var geometry, material, mesh;
function init()
    renderer = new THREE.WebGLRenderer();
    container = document.createElement("div");
    document.body.appendChild(container);
    camera = new THREE.PerspectiveCamera( 50, window.innerWidth / window.innerHeight, 1, 5000 );
    camera.position.z = 1400;
    scene = new THREE.Scene();
    @var circleGeometry = new THREE.CircleBufferGeometry( 1, 6 );
    geometry = new THREE.InstancedBufferGeometry();
    geometry.index = circleGeometry.index;
    geometry.attributes = circleGeometry.attributes;
    @var particleCount = 75000;
    @var translateArray = new Float32Array( particleCount * 3 );
    for ( @var i = 0, i3 = 0, l = particleCount; i < l; i ++, i3 += 3 ) {
        translateArray[ i3 + 0 ] = Math.random() * 2 - 1;
        translateArray[ i3 + 1 ] = Math.random() * 2 - 1;
        translateArray[ i3 + 2 ] = Math.random() * 2 - 1;
    }
    geometry.addAttribute( 'translate', new THREE.InstancedBufferAttribute( translateArray, 3, 1 ) );
    material = new THREE.RawShaderMaterial( {
        uniforms: {
            map: { value: new THREE.TextureLoader().load( 'file:///home/s/Desktop/circle.png' ) },
            time: { value: 0.0 }
        },
        vertexShader: document.getElementById( 'vshader' ).textContent,
        fragmentShader: document.getElementById( 'fshader' ).textContent,
        depthTest: true,
        depthWrite: true
    } );
    mesh = new THREE.Mesh( geometry, material );
    mesh.scale.set( 500, 500, 500 );
    scene.add( mesh );
    renderer.setPixelRatio( window.devicePixelRatio );
    renderer.setSize( window.innerWidth, window.innerHeight );
    container.appendChild( renderer.domElement );
    window.addEventListener( 'resize', onWindowResize, false );
    return true;
}
function onWindowResize( event ) {
    camera.aspect = window.innerWidth / window.innerHeight;
    camera.updateProjectionMatrix();
    renderer.setSize( window.innerWidth, window.innerHeight );
}
function animate() {
    requestAnimationFrame( animate );
    render();
}
function render() {
    @var time = performance.now() * 0.0005;
    material.uniforms.time.value = time;
    mesh.rotation.x = time * 0.2;
    mesh.rotation.y = time * 0.4;
    renderer.render( scene, camera );
}
if ( init() ) {
    animate();
}


text:
using Base.Test
using ToeplitzMatrices

ns = 101
nl = 2000

xs = randn(ns, 5)
xl = randn(nl, 5)

println("Toeplitz")
for (As, Al, st) in ((Toeplitz(0.9.^(0:ns-1), 0.4.^(0:ns-1)),
                        Toeplitz(0.9.^(0:nl-1), 0.4.^(0:nl-1)),
                            "Real general square"),
                     (Toeplitz(complex(0.9.^(0:ns-1)), complex(0.4.^(0:ns-1))),
                        Toeplitz(complex(0.9.^(0:nl-1)), complex(0.4.^(0:nl-1))),
                            "Complex general square"),
                     (Circulant(0.9.^(0:ns - 1)),
                        Circulant(0.9.^(0:nl - 1)),
                            "Real circulant"),
                     (Circulant(complex(0.9.^(0:ns - 1))),
                        Circulant(complex(0.9.^(0:nl - 1))),
                            "Complex circulant"),
                     (TriangularToeplitz(0.9.^(0:ns - 1), :U),
                        TriangularToeplitz(0.9.^(0:nl - 1), :U),
                            "Real upper triangular"),
                     (TriangularToeplitz(complex(0.9.^(0:ns - 1)), :U),
                        TriangularToeplitz(complex(0.9.^(0:nl - 1)), :U),
                            "Complex upper triangular"),
                     (TriangularToeplitz(0.9.^(0:ns - 1), :L),
                        TriangularToeplitz(0.9.^(0:nl - 1), :L),
                            "Real lower triangular"),
                     (TriangularToeplitz(complex(0.9.^(0:ns - 1)), :L),
                        TriangularToeplitz(complex(0.9.^(0:nl - 1)), :L),
                            "Complex lower triangular"))
    print("$st: ")
    @test_approx_eq As * xs full(As) * xs
    @test_approx_eq Al * xl full(Al) * xl
    @test_approx_eq A_ldiv_B!(As, LinAlg.copy_oftype(xs, eltype(As))) full(As) \ xs
    @test_approx_eq A_ldiv_B!(Al, LinAlg.copy_oftype(xl, eltype(Al))) full(Al) \ xl
    println("OK!")
end

print("Real general rectangular: ")
Ar1 = Toeplitz(0.9.^(0:nl-1), 0.4.^(0:ns-1))
Ar2 = Toeplitz(0.9.^(0:ns-1), 0.4.^(0:nl-1))
@test_approx_eq Ar1 * xs full(Ar1) * xs
@test_approx_eq Ar2 * xl full(Ar2) * xl
println("OK!")

print("Complex general rectangular: ")
Ar1 = Toeplitz(complex(0.9.^(0:nl-1)), complex(0.4.^(0:ns-1)))
Ar2 = Toeplitz(complex(0.9.^(0:ns-1)), complex(0.4.^(0:nl-1)))
@test_approx_eq Ar1 * xs full(Ar1) * xs
@test_approx_eq Ar2 * xl full(Ar2) * xl
println("OK!")

print("Symmetric Toeplitz: ")
As = SymmetricToeplitz(0.9.^(0:ns-1))
Ab = SymmetricToeplitz(abs(randn(ns)))
Al = SymmetricToeplitz(0.9.^(0:nl-1))
@test_approx_eq As * xs full(As) * xs
@test_approx_eq Ab * xs full(Ab) * xs
@test_approx_eq Al * xl full(Al) * xl
@test_approx_eq A_ldiv_B!(As, copy(xs)) full(As) \ xs
@test_approx_eq A_ldiv_B!(Ab, copy(xs)) full(Ab) \ xs
@test_approx_eq A_ldiv_B!(Al, copy(xl)) full(Al) \ xl
@test_approx_eq StatsBase.levinson(As, xs) full(As) \ xs
@test_approx_eq StatsBase.levinson(Ab, xs) full(Ab) \ xs
@test_approx_eq StatsBase.levinson(Al, xl) full(Al) \ xl
println("OK!")

println("\nHankel")
print("Real square: ")
H = Hankel([1.0,2,3,4,5],[5.0,6,7,8,0])
x = ones(5)
@test_approx_eq full(H)*x H*x

Hs = Hankel(0.9.^(ns-1:-1:0), 0.4.^(0:ns-1))
Hl = Hankel(0.9.^(nl-1:-1:0), 0.4.^(0:nl-1))
@test_approx_eq Hs * xs[:,1] full(Hs) * xs[:,1]
@test_approx_eq Hs * xs full(Hs) * xs
@test_approx_eq Hl * xl full(Hl) * xl
println("OK!")

print("Complex square: ")
H = Hankel(complex([1.0,2,3,4,5]), complex([5.0,6,7,8,0]))
x = ones(5)
@test_approx_eq full(H)*x H*x

Hs = Hankel(complex(0.9.^(ns-1:-1:0)), complex(0.4.^(0:ns-1)))
Hl = Hankel(complex(0.9.^(nl-1:-1:0)), complex(0.4.^(0:nl-1)))
@test_approx_eq Hs * xs[:,1] full(Hs) * xs[:,1]
@test_approx_eq Hs * xs full(Hs) * xs
@test_approx_eq Hl * xl full(Hl) * xl
println("OK!")

print("Real rectangular: ")
Hs = Hankel(0.9.^(ns-1:-1:0), 0.4.^(0:nl-1))
Hl = Hankel(0.9.^(nl-1:-1:0), 0.4.^(0:ns-1))
@test_approx_eq Hs * xl[:,1] full(Hs) * xl[:,1]
@test_approx_eq Hs * xl full(Hs) * xl
@test_approx_eq Hl * xs full(Hl) * xs
println("OK!")

print("Complex rectangular: ")
Hs = Hankel(complex(0.9.^(ns-1:-1:0)), complex(0.4.^(0:nl-1)))
Hl = Hankel(complex(0.9.^(nl-1:-1:0)), complex(0.4.^(0:ns-1)))
@test_approx_eq Hs * xl[:,1] full(Hs) * xl[:,1]
@test_approx_eq Hs * xl full(Hs) * xl
@test_approx_eq Hl * xs full(Hl) * xs
println("OK!")

if isdir(Pkg.dir("FastTransforms"))
    print("\nBigFloat")
    using FastTransforms
    T = Toeplitz(BigFloat[1,2,3,4,5], BigFloat[1,6,7,8,0])
    @test_approx_eq T*ones(BigFloat,5) [22,24,19,16,15]

    n = 512
    r = map(BigFloat,rand(n))
    T = Toeplitz(r,[r[1];map(BigFloat,rand(n-1))])
    @test_approx_eq T*ones(BigFloat,n) full(T)*ones(BigFloat,n)

    T = TriangularToeplitz(BigFloat[1,2,3,4,5],:L)
    @test_approx_eq T*ones(BigFloat,5) full(T)*ones(BigFloat,5)

    n = 512
    r = map(BigFloat,rand(n))
    T = TriangularToeplitz(r,:L)
    @test_approx_eq T*ones(BigFloat,n) full(T)*ones(BigFloat,n)

    T = TriangularToeplitz(BigFloat[1,2,3,4,5],:U)
    @test_approx_eq T*ones(BigFloat,5) full(T)*ones(BigFloat,5)

    n = 512
    r = map(BigFloat,rand(n))
    T = TriangularToeplitz(r,:U)
    @test_approx_eq T*ones(BigFloat,n) full(T)*ones(BigFloat,n)
    println("OK!")
end

println("Convert")

T = Toeplitz(ones(2),ones(2))

@test isa(convert(Matrix{Complex128},T),Matrix{Complex128})
@test isa(convert(AbstractMatrix{Complex128},T),Toeplitz{Complex128})
@test isa(convert(AbstractArray{Complex128},T),Toeplitz{Complex128})
@test isa(convert(ToeplitzMatrices.AbstractToeplitz{Complex128},T),Toeplitz{Complex128})
@test isa(convert(ToeplitzMatrices.Toeplitz{Complex128},T),Toeplitz{Complex128})

T = SymmetricToeplitz(ones(2))

@test isa(convert(Matrix{Float32},T),Matrix{Float32})
@test isa(convert(AbstractMatrix{Float32},T),SymmetricToeplitz{Float32})
@test isa(convert(AbstractArray{Float32},T),SymmetricToeplitz{Float32})
@test isa(convert(ToeplitzMatrices.AbstractToeplitz{Float32},T),SymmetricToeplitz{Float32})
@test isa(convert(ToeplitzMatrices.SymmetricToeplitz{Float32},T),SymmetricToeplitz{Float32})

T = Circulant(ones(2))

@test isa(convert(Matrix{Complex128},T),Matrix{Complex128})
@test isa(convert(AbstractMatrix{Complex128},T),Circulant{Complex128})
@test isa(convert(AbstractArray{Complex128},T),Circulant{Complex128})
@test isa(convert(ToeplitzMatrices.AbstractToeplitz{Complex128},T),Circulant{Complex128})
@test isa(convert(ToeplitzMatrices.Circulant{Complex128},T),Circulant{Complex128})

T = TriangularToeplitz(ones(2),:U)

@test isa(convert(Matrix{Complex128},T),Matrix{Complex128})
@test isa(convert(AbstractMatrix{Complex128},T),TriangularToeplitz{Complex128})
@test isa(convert(AbstractArray{Complex128},T),TriangularToeplitz{Complex128})
@test isa(convert(ToeplitzMatrices.AbstractToeplitz{Complex128},T),TriangularToeplitz{Complex128})
@test isa(convert(ToeplitzMatrices.TriangularToeplitz{Complex128},T),TriangularToeplitz{Complex128})

T = Hankel(ones(2),ones(2))

@test isa(convert(Matrix{Complex128},T),Matrix{Complex128})
@test isa(convert(AbstractMatrix{Complex128},T),Hankel{Complex128})
@test isa(convert(AbstractArray{Complex128},T),Hankel{Complex128})
@test isa(convert(ToeplitzMatrices.Hankel{Complex128},T),Hankel{Complex128})


text:
myunpack(a) = a
myunpack(a::ArrayPartition) = (a.x[1], a.x[2], a.x[3][])
mypack(a,b,c) = ArrayPartition(a,b,[c])
mypack(a::Number...) = [a...]
mypack(a::SArray,b::SArray,c::Number) = ArrayPartition(a,b,@SVector[c])
mypack(a::SArray,b::SArray,c::SArray) = ArrayPartition(a,b,c)

function backwardfilter(k::SDEKernel, p::WGaussian{(:μ, :Σ, :c)};
    filter=CovarianceFilter(), alg=Euler(),
    inplace=false, apply_timechange=false, abstol=1e-6, reltol=1e-3)

  message, solend = _backwardfilter(filter, k::SDEKernel, (p.c, p.μ, p.Σ);
    alg=alg, inplace=inplace, apply_timechange=apply_timechange, abstol=abstol, reltol=reltol)

  return message, WGaussian{(:μ, :Σ, :c)}(myunpack(solend)...)
end

function backwardfilter(k::SDEKernel, (c, μ, Σ)::NamedTuple{(:logscale, :μ, :Σ)};
    filter=CovarianceFilter(), alg=Euler(), inplace=false, apply_timechange=false, abstol=1e-6,reltol=1e-3)

  return _backwardfilter(filter, k::SDEKernel, (c, μ, Σ);
    alg=alg, inplace=inplace, apply_timechange=apply_timechange, abstol=abstol,reltol=reltol)
end

# covariance filter ODE Eqs.
compute_dP(B,P,σtil) = convert(typeof(P), B*P + P*B' - outer_(σtil))
compute_dP(B,P::SArray,σtil::Number) = B*P + P*B' - σtil*σtil'*similar_type(P, Size(size(P,1),size(P,1)))(I)

compute_dν(B,ν,β::Number) = B*ν .+ β
compute_dν(B,ν,β) = convert(typeof(ν),B*ν + β)

function compute_dP!(dP,B,P,σtil)
  dP .= B*P + P*B' - outer_(σtil)
  return nothing
end

function compute_dν!(dν,B,ν,β::Number)
  dν .= B*ν .+ β
  return nothing
end

function compute_dν!(dν,B,ν,β)
  dν .= B*ν + β
  return nothing
end

function CovarianceFilterODE(u, p, t)
  B, β, σtil = p

  # take care for multivariate case here if P isa Matrix, ν  isa Vector, c isa Scalar
  ν, P, c = myunpack(u)

  # H = inv(P)
  # F = H*ν

  dν = compute_dν(B,ν,β)
  dP = compute_dP(B,P,σtil)
  dc = tr(B)

  return mypack(dν, dP, dc)
end

function CovarianceFilterODE(du, u, p, t)
  B, β, σtil = p

  # take care for multivariate case here if P isa Matrix, ν  isa Vector, c isa Scalar
  ν, P, c = myunpack(u)

  #  H = inv(P)
  #  F = H*ν

  compute_dν!(du.x[1],B,ν,β)
  compute_dP!(du.x[2],B,P,σtil)
  du.x[3] .= tr(B)

  return nothing
end

function _backwardfilter(filter::CovarianceFilter,k::SDEKernel, (c, ν, P);
    alg=Euler(), inplace=false, apply_timechange=false, abstol=1e-6,reltol=1e-3)
  @unpack trange, p = k

  # Initialize OD
  u0 = mypack(ν, P, c)
  prob = ODEProblem{inplace}(CovarianceFilterODE, u0, reverse(get_tspan(trange)), p)

  if !apply_timechange
    _ts = trange # use collect() here?
  else
    _ts = timechange(trange)
  end

  if !OrdinaryDiffEq.isadaptive(alg)
    sol = solve(prob, alg, tstops=_ts)
  else
    sol = solve(prob, alg, dt=get_dt(k.trange), tstops=_ts, abstol=abstol, reltol=reltol)
  end
  message = Message(sol, k, filter, apply_timechange)

  return message, sol[end]
end

# information filter ODE Eqs.
compute_dH(H,B,σtil) = -B'*H - H*B + H*outer_(σtil)*H
compute_dF(F,H::AbstractArray,B,σtil,β::Number) = -B'*F + H*outer_(σtil)*F + H*fill(β,size(H,1))
compute_dF(F,H,B,σtil,β) = -B'*F + H*outer_(σtil)*F + H*β

function compute_dH!(dH,H,B,σtil)
  dH .= -B'*H - H*B + H*outer_(σtil)*H
  return nothing
end

function compute_dF!(dF,F,H,B,σtil,β)
  dF .= -B'*F + H*outer_(σtil)*F + H*β
  return nothing
end

function compute_dF!(dF,F,H,B,σtil,β::Number)
  dF .= -B'*F + H*outer_(σtil)*F + H*fill(β,size(H,1))
  return nothing
end

function InformationFilterODE(u, p, t)
  B, β, σtil = p

  # take care for multivariate case here if H isa Matrix, F  isa Vector, c isa Scalar
  F, H, c = myunpack(u)

  dF = compute_dF(F,H,B,σtil,β)
  dH = compute_dH(H,B,σtil)
  dc = tr(B)

  return mypack(dF, dH, dc)
end

function InformationFilterODE(du, u, p, t)
  B, β, σtil = p

  # take care for multivariate case here if H isa Matrix, F  isa Vector, c isa Scalar
  F, H, c = myunpack(u)

  #  H = inv(P)
  #  F = H*ν

  compute_dF!(du.x[1],F,H,B,σtil,β)
  compute_dH!(du.x[2],H,B,σtil)
  du.x[3] .= tr(B)

  return nothing
end

function _backwardfilter(filter::InformationFilter,k::SDEKernel, (c, F, H);
    alg=Euler(), inplace=false, apply_timechange=false, abstol=1e-6,reltol=1e-3)
  @unpack trange, p = k

  # Initialize ODE
  u0 = mypack(F, H, c)
  prob = ODEProblem{inplace}(InformationFilterODE, u0, reverse(get_tspan(trange)), p)

  if !apply_timechange
    _ts = trange # use collect() here?
  else
    _ts = timechange(trange)
  end

  if !OrdinaryDiffEq.isadaptive(alg)
    sol = solve(prob, alg, tstops=_ts, abstol=abstol, reltol=reltol)
  else
    sol = solve(prob, alg, dt=get_dt(k.trange), tstops=_ts, abstol=abstol, reltol=reltol)
  end
  message = Message(sol, k, filter, apply_timechange)

  return message, sol[end]
end

Φfunc(t,T,B) = exp(-(T-t)*B)
function solP(t,T,B,PT,Σ)
  Φ = Φfunc(t,T,B)
  Φ*(Σ+PT)*Φ'-Σ
end

function solν(t,T,B,β,νT)
  #TODO
  # β time-dependent
  Φfunc(t,T,B)*(νT + inv(B)*β) - inv(B)*β
end

function solc(t,T,B,cT)
  cT-tr(B)*(T-t)
end

function (G::solLyapunov)(t)
  @unpack ktilde, Σ, νT, PT, cT = G
  @unpack trange, p = ktilde
  B, β, σtil = p
  T = last(trange)

  ν = solν(t,T,B,β,νT)
  P = solP(t,T,B,PT,Σ)
  c = solc(t,T,B,cT)
  return mypack(ν, P, c)
end

function _backwardfilter(filter::LyapunovFilter,k::SDEKernel, (c, ν, P); apply_timechange=false, kwargs...)
  @unpack trange, p = k

  # solve Lyapunov equation
  B, β, σtil = p
  atil = construct_a(σtil,P)
  (B isa Symmetric) && (B = Array(B))
  Σ = lyap(B,atil)

  if !apply_timechange
    _ts = trange
  else
    _ts = timechange(trange)
  end

  sol = solLyapunov(k, Σ, ν, P, c)

  u = sol.(_ts)
  T = eltype(eltype(u))
  N = length((size(u[1])..., length(u)))
  SciMLsol = SciMLBase.ODESolution{T,N,typeof(u),Nothing,Nothing,typeof(_ts),Nothing,
                        Nothing,Nothing,Nothing,Nothing}(
                        u,nothing,nothing,
                        _ts,nothing,nothing,nothing,nothing,true,0,nothing,:Success)

  soldis = construct_discrete_sol(u)

  message = Message(k, sol, soldis, _ts, filter)

  return message, SciMLsol[1]
end


text:
using Juriba
using Test

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


text:
module Shearheating

# This implements different methods to specify Shearheating of rocks
#
# If you want to add a new method here, feel free to do so.
# Remember to also export the function name in GeoParams.jl (in addition to here)

using Parameters, LaTeXStrings, Unitful
using ..Units
using GeoParams: AbstractMaterialParam
import Base.show, GeoParams.param_info
using ..MaterialParameters: MaterialParamsInfo

abstract type AbstractShearheating{T} <: AbstractMaterialParam end

export  ConstantShearheating,   # constant
        compute_shearheating,   # calculation routines
        compute_shearheating!,
        param_info

# Constant Shearheating -------------------------------------------------------
"""
    ConstantShearheating(Χ=0.0NoUnits)

Set the shear heating efficiency [0-1] parameter
```math
Χ  = cst
```
where ``\\Chi`` is the shear heating efficiency [NoUnits]

Shear heating is computed as
```math
H_s = \\Chi \\cdot \\tau_{ij}(\\dot{\\varepsilon}_{ij} - \\dot{\\varepsilon}^{el}_{ij})
```

"""
@with_kw_noshow struct ConstantShearheating{T,U} <: AbstractShearheating{T}
    Χ::GeoUnit{T,U}         =   0.0*NoUnits
end
ConstantShearheating(args...) = ConstantShearheating(convert.(GeoUnit,args)...)

function param_info(s::ConstantShearheating) # info about the struct
    return MaterialParamsInfo(Equation =  L"\H_s = \Chi \tau_{ij}(\dot{\varepsilon}_{ij} - \dot{\varepsilon}^{el}_{ij})")
end

# In-place routine (H_s can't take new value)
function compute_shearheating!(H_s, s::ConstantShearheating{_T}, τ, ε, ε_el) where _T
    @unpack_val Χ   = s

    if isnothing(ε_el)
        H_s = Χ*sum( τ .* ε  )
    else
        H_s = Χ*sum( τ .* (ε .- ε_el) )
    end

end

# Calculation routine
function compute_shearheating(s::ConstantShearheating{_T}, τ, ε, ε_el) where _T
    @unpack_val Χ   = s

    if isnothing(ε_el)
        H_s = Χ*sum( τ .* ε  )
    else
        H_s = Χ*sum( τ .* (ε .- ε_el) )
    end

    return H_s
end

# Print info
function show(io::IO, g::ConstantShearheating)
    print(io, "Shear heating: H_s = $(UnitValue(g.Χ)) τ_ij*(ε_ij - ε^el_ij)")
end
#-------------------------------------------------------------------------

# Help info for the calculation routines
"""
    H_s = compute_shearheating(s:<AbstractShearheating, τ, ε, ε_el)

Computes the shear heating source term

```math
H_s = \\Chi \\cdot \\tau_{ij} ( \\dot{\\varepsilon}_{ij} - \\dot{\\varepsilon}^{el}_{ij})
```

# Parameters
- ``\\Chi`` : The efficiency of shear heating (between 0-1)
- ``\\tau_{ij}`` : The full deviatoric stress tensor [4 components in 2D; 9 in 3D]
- ``\\dot{\\varepsilon}_{ij}`` : The full deviatoric strainrate tensor
- ``\\dot{\\varepsilon}^{el}_{ij}`` : The full elastic deviatoric strainrate tensor

"""
compute_shearheating(s::AbstractShearheating, τ, ε, ε_el)

"""
    H_s = ComputeShearheating(s:<AbstractShearheating, τ, ε)

Computes the shear heating source term when there is no elasticity

```math
H_s = \\Chi \\cdot \\tau_{ij}  \\dot{\\varepsilon}_{ij}
```

# Parameters
- ``\\Chi`` : The efficiency of shear heating (between 0-1)
- ``\\tau_{ij}`` : The full deviatoric stress tensor [4 components in 2D; 9 in 3D]
- ``\\dot{\\varepsilon}_{ij}`` : The full deviatoric strainrate tensor
"""
compute_shearheating(s::AbstractShearheating{_T}, τ::Any, ε::Any) where _T = compute_shearheating(s, τ, ε, nothing)

"""
    compute_shearheating!(H_s, s:<AbstractShearheating,  τ, ε, ε_el)

Computes the shear heating source term in-place

```math
H_s = \\Chi \\cdot \\tau_{ij} ( \\dot{\\varepsilon}_{ij} - \\dot{\\varepsilon}^{el}_{ij})
```

# Parameters
- ``\\Chi`` : The efficiency of shear heating (between 0-1)
- ``\\tau_{ij}`` : The full deviatoric stress tensor [4 components in 2D; 9 in 3D]
- ``\\dot{\\varepsilon}_{ij}`` : The full deviatoric strainrate tensor
- ``\\dot{\\varepsilon}^{el}_{ij}`` : The full elastic deviatoric strainrate tensor

*NOTE:*
The shear heating terms require the full deviatoric stress & strain rate tensors, i.e.:

```math
2D: \\tau_{ij} = \\left(
                \\begin{matrix}
                    \\tau_{xx} & \\tau_{xz} \\\\
                    \\tau_{zx} & \\tau_{zz}
                \\end{matrix}
            \\right)
```
Since ``\\tau_{zx}=\\tau_{xz}``, most geodynamic codes only take one of the terms into account; shear heating requires all components to be used!
"""
compute_shearheating!(H_s, s::AbstractShearheating{_T}, τ, ε, ε_el) where _T

"""
    compute_shearheating!(H_s, s:<AbstractShearheating, τ, ε)

Computes the shear heating source term `H_s` in-place when there is no elasticity

```math
H_s = \\Chi \\cdot \\tau_{ij}  \\dot{\\varepsilon}_{ij}
```

# Parameters
- ``\\Chi`` : The efficiency of shear heating (between 0-1)
- ``\\tau_{ij}`` : The full deviatoric stress tensor [4 components in 2D; 9 in 3D]
- ``\\dot{\\varepsilon}_{ij}`` : The full deviatoric strainrate tensor
"""
compute_shearheating!(H_s::Any, s::AbstractShearheating{_T}, τ::Any, ε::Any) where _T = compute_shearheating!(H_s, s, τ, ε, nothing)

end


text:
# Tests for Entry type and associated methods

@testset "Entry" begin
    @testset "Constructors" begin
        @test_throws ArgumentError Entry(-5, 10)
        @test_throws ArgumentError Entry(10, -5)
        @test_throws ArgumentError Entry(15, 20)

        @test_throws InexactError Entry(2.5, 1.0, 4)
        @test_throws InexactError Entry(2.0, 1.5, 4)

        @test Entry(4,2,6) == Entry(4,2)

        a = Entry(6,5)
        b = Entry(a)
        c = Entry(6 => 5)
        d = Entry{Int}(6, 5)
        e = Entry(Int32(6), Int64(5))
        f = Entry(UInt(6), Int(5))
        @test a == b == c == d == e == f

        g = Entry(6,5,6.0)
        h = Entry{Float64}(6, 5)
        @test g == h
    end

    @testset "Accessors" begin
        a = Entry(15,6)
        @test rownumber(a) == 15
        @test rowposition(a) == 6
        @test value(a) == 5005
    end

    @testset "Comparators" begin
        a = Entry(4,1)
        b = Entry{Float64}(4, 1)
        c = Entry{Float64}(10, 5)
        d = Entry(12, 6)
        e = Entry(14, 7)
        @test a ≤ b < c < d < e

        @test a ≈ b
    end

    @testset "Operators" begin
        a = Entry(7,3,35)
        b = Entry(7,4,35.0)
        c = Entry(8,4,70)
        @test a + b == c
        @test_throws NonAdjacentError a + c
        @test c - a == b
        @test_throws NonAdjacentError a - b
    end

    @testset "Checks" begin
        a = Entry(0,0)
        b = Entry(9,0)
        c = Entry(9,9)
        d = Entry(15,3)
        e = Entry(15,4)
        f = Entry(16,4)
        @test isfirst(a)
        @test isatleft(a)
        @test isatleft(b)
        @test !isatleft(c)
        @test isatright(c)
        @test !isatright(b)
        @test isadjacent(d, e)
        @test !isadjacent(a, b)
        @test isinterior(d)
        @test !isinterior(c)
        @test issubtractable(f, d)
        @test issubtractable(f, e)
        @test !issubtractable(d, e)
        @test PT.isvalid(d)
        @test !PT.isvalid(Entry(15,3,0))
    end

    @testset "Movement" begin
        a = Entry(0,0)
        b = Entry(20,20)
        c = Entry(21,0)
        d = Entry(30,10)
        e = Entry(30,11)
        f = Entry(30,12)
        g = Entry(29,11)
        h = Entry(31,11)

        @test_throws OutOfBoundsError prev(a)
        @test_throws OutOfBoundsError up(a)
        @test_throws OutOfBoundsError left(a)
        @test_throws OutOfBoundsError right(a)
        @test down(a) == Entry(1,0)

        @test next(b) == c
        @test b == prev(c)

        @test left(e) == prev(e) == d
        @test right(e) == next(e) == f
        @test up(e) == g
        @test down(e) == h
    end
end


text:
module InvertedIndices

export InvertedIndex, Not

using Base: tail

struct InvertedIndex{S}
    skip::S
end
const Not = InvertedIndex
# Support easily inverting multiple indices without a temporary array in Not([...])
InvertedIndex(i₁::Integer, i₂::Integer, iₓ::Integer...) = InvertedIndex(TupleVector((i₁, i₂, iₓ...)))

"""
    InvertedIndex(idx)
    Not(idx)

Construct an inverted index, selecting all indices not in the passed `idx`.

Upon indexing into an array, the `InvertedIndex` behaves like a 1-dimensional
collection of the indices of the array that are not in `idx`. Bounds are
checked to ensure that all indices in `idx` are within the bounds of the array
— even though they are skipped. The `InvertedIndex` behaves like a
1-dimensional collection of its inverted indices. If `idx` spans multiple
dimensions (like a multidimensional logical mask or `CartesianIndex`), then the
inverted index will similarly span multiple dimensions.
"""
InvertedIndex, Not

"""
    BroadcastedInvertedIndex(x::InvertedIndex)

A wrapper for `InvertedIndex` if it is used in broadcasting context.
Since `InvertedIndex` does not have a reference to the collection it applies to
it is impossible to define its `axes` eagerly.
Therefore it is the responsiblity of the caller to resolve the handling
of the broadcast result.

# Examples

julia> Not(1) .=> sin
InvertedIndices.BroadcastedInvertedIndex(InvertedIndex{Int64}(1)) => sin

julia> Not(:col) .=> [minimum, maximum]
2-element Vector{Pair{InvertedIndices.BroadcastedInvertedIndex, _A} where _A}:
 InvertedIndices.BroadcastedInvertedIndex(InvertedIndex{Symbol}(:col)) => minimum
 InvertedIndices.BroadcastedInvertedIndex(InvertedIndex{Symbol}(:col)) => maximum
"""
struct BroadcastedInvertedIndex
    sel::InvertedIndex
end

Base.Broadcast.broadcastable(x::InvertedIndex) = Ref(BroadcastedInvertedIndex(x))

# A very simple and primitive static array to avoid allocations for Not(1,2,3) while fulfilling the indexing API
struct TupleVector{T<:Tuple} <: AbstractVector{Int}
    data::T
end
Base.size(::TupleVector{<:NTuple{N}}) where {N} = (N,)
@inline function Base.getindex(t::TupleVector, i::Int)
    @boundscheck checkbounds(t, i)
    return @inbounds t.data[i]
end

# Like Base.LogicalIndex, the InvertedIndexIterator is a pseudo-vector that is
# just used as an iterator and does not support getindex.
struct InvertedIndexIterator{T,S,P} <: AbstractVector{T}
    skips::S
    picks::P
end
InvertedIndexIterator(skips, picks) = InvertedIndexIterator{eltype(picks), typeof(skips), typeof(picks)}(skips, picks)
Base.size(III::InvertedIndexIterator) = (length(III.picks) - length(III.skips),)

@inline Base.iterate(I::InvertedIndexIterator) = iterate(I, (iterate(I.skips), iterate(I.picks)))
Base.iterate(I::InvertedIndexIterator, ::Tuple{Any, Nothing}) = nothing
@inline function Base.iterate(I::InvertedIndexIterator, (skipitr, pickitr))
    while should_skip(skipitr, pickitr)
        skipitr = iterate(I.skips, tail(skipitr)...)
        pickitr = iterate(I.picks, tail(pickitr)...)
        pickitr === nothing && return nothing
    end
    return (pickitr[1], (skipitr, iterate(I.picks, tail(pickitr)...)))
end
Base.collect(III::InvertedIndexIterator) = [i for i in III]

should_skip(::Nothing, ::Any) = false
should_skip(s::Tuple, p::Tuple) = _should_skip(s[1], p[1])
_should_skip(s, p) = s == p
_should_skip(s::Integer, p::CartesianIndex{1}) = s == p.I[1]
_should_skip(s::CartesianIndex{1}, p::Integer) = s.I[1] == p

@inline Base.checkbounds(::Type{Bool}, A::AbstractArray, I::InvertedIndexIterator{<:Integer}) =
    checkbounds(Bool, A, I.skips) && eachindex(IndexLinear(), A) == eachindex(IndexLinear(), I.picks)
@inline Base.checkbounds(::Type{Bool}, A::AbstractArray, I::InvertedIndexIterator) = checkbounds(Bool, A, I.skips) && axes(A) == axes(I.picks)
@inline Base.checkindex(::Type{Bool}, indx::AbstractUnitRange, I::InvertedIndexIterator) = checkindex(Bool, indx, I.skips) && (indx,) == axes(I.picks)
@inline Base.checkindex(::Type{Bool}, inds::Tuple, I::InvertedIndexIterator) = checkindex(Bool, indx, I.skips) && inds == axes(I.picks)

@inline Base.ensure_indexable(I::Tuple{InvertedIndexIterator, Vararg{Any}}) = (collect(I[1]), Base.ensure_indexable(tail(I))...)

# This is a little hacky, but we display InvertedIndexIterators like the `Not`s they come from
function Base.show(io::IO, I::InvertedIndexIterator)
    print(io, "Not(")
    show(io, I.skips)
    print(io, ")")
end

# Inverted indices must be sorted and unique to ensure that iterating over
# them and the axes simultaneously will work appropriately. Doing this fully
# generically is a challenge. It's a little annoying to need to take a pass
# through the inverted index before actually doing the indexing.
uniquesort(A::AbstractArray) = uniquesort(vec(A))
uniquesort(A::DenseVector) = issorted(A, lt=(<=)) ? A : (unique! ∘ sort)(A)
uniquesort(A::AbstractVector) = issorted(A, lt=(<=)) ? A : (unique ∘ sort)(A)
uniquesort(r::AbstractRange) = step(r) > 0 ? r : step(r) == 0 ? r[end:end] : reverse(r)
uniquesort(A::Base.LogicalIndex) = A
uniquesort(x) = x

@inline function Base.to_indices(A, inds, I::Tuple{InvertedIndex, Vararg{Any}})
    new_indices = to_indices(A, inds, (I[1].skip, tail(I)...))
    skips = uniquesort(new_indices[1])
    picks = spanned_indices(inds, skips)[1]
    return (InvertedIndexIterator(skips, picks), tail(new_indices)...)
end

struct ZeroDArray{T} <: AbstractArray{T,0}
    x::T
end
Base.size(::ZeroDArray) = ()
Base.getindex(Z::ZeroDArray) = Z.x

# Be careful with CartesianIndex as they splat out to a variable number of new indices and do not iterate
function Base.to_indices(A, inds, I::Tuple{InvertedIndex{<:CartesianIndex}, Vararg{Any}})
    skips = ZeroDArray(I[1].skip)
    picks, tails = spanned_indices(inds, skips)
    return (InvertedIndexIterator(skips, picks), to_indices(A, tails, tail(I))...)
end

# Either return a CartesianRange or an axis vector
@inline spanned_indices(inds, ::Any) = inds[1], tail(inds)
const NIdx{N} = Union{CartesianIndex{N}, AbstractArray{CartesianIndex{N}}, AbstractArray{Bool,N}}
@inline spanned_indices(inds, ::NIdx{0}) = CartesianIndices(()), inds
@inline spanned_indices(inds, ::NIdx{1}) = inds[1], tail(inds)
@inline function spanned_indices(inds, ::NIdx{N}) where N
    heads, tails = Base.IteratorsMD.split(inds, Val(N))
    return CartesianIndices(heads), tails
end

# It's possible more indices were specified than there were axes
@inline spanned_indices(inds::Tuple{}, ::Any) = Base.OneTo(1), ()
@inline spanned_indices(inds::Tuple{}, ::NIdx{0}) = CartesianIndices(()), ()
@inline spanned_indices(inds::Tuple{}, ::NIdx{1}) = Base.OneTo(1), ()
@inline spanned_indices(inds::Tuple{}, ::NIdx{N}) where N = CartesianIndices(ntuple(i->Base.OneTo(1), N)), ()

# This is an interesting need — we need this because otherwise indexing with a
# single multidimensional boolean array ends up comparing a multidimensional cartesian
# index to a linear index. Does this need addressing in Base, too?
@inline Base.to_indices(A, I::Tuple{Not{<:NIdx{1}}}) = to_indices(A, (eachindex(IndexLinear(), A),), I)
@inline Base.to_indices(A, I::Tuple{Not{<:NIdx}}) = to_indices(A, axes(A), I)
# Arrays of Bool are even more confusing as they're sometimes linear and sometimes not
@inline Base.to_indices(A, I::Tuple{Not{<:AbstractArray{Bool, 1}}}) = to_indices(A, (eachindex(IndexLinear(), A),), I)
@inline Base.to_indices(A, I::Tuple{Not{<:Union{Array{Bool}, BitArray}}}) = to_indices(A, (eachindex(A),), I)

end # module


text:
module Polyester

using ThreadingUtilities
using ArrayInterface: static_length, static_step, static_first, size
using StrideArraysCore: object_and_preserve
using ManualMemory: Reference
using Static
using Requires
using PolyesterWeave: request_threads, free_threads!, mask, UnsignedIteratorEarlyStop, assume
using CPUSummary: num_threads, num_cores

export batch, @batch, num_threads

include("batch.jl")
include("closure.jl")

# y = rand(1)
# x = rand(1)
# @batch for i ∈ eachindex(y,x)
#   y[i] = sin(x[i])
# end
end


text:
import NSBEntropy
import Spiketrains
import Information
import StatsBase

function example1()
    nreps = 1000
    #20 ms window
    t = linspace(0,.02,1000)
    dt = mean(diff(t))
    #constant firing rate of 260 Hz
    fr = 260.*ones(length(t))
    T,repidx = Spiketrains.generateSpikes(fr,dt,nreps)
    #create an aligned spiketrain
    at1 = Information.AlignedSpiketrain(T,repidx,nreps,length(T))
    #create a histogram using 0.5 ms bins
    countsT = Information.getTrialSpikeCount(at1,[0:0.0005:0.015])
    #ensure that we have a maxiumum of 1 spike per bin
    countsT.counts[countsT.counts.>1] = 1
    #hask the words
    Q = countsT.counts'*(2.^[0:29])
    #get the counts of each unique value
    q = StatsBase.countmap(Q)
    #compute entropy, making use of the fact that the refractory period ensures L <= 10^16
    Entropy.NSBEntropy.saddlepoint(collect(values(q)),10^16)
end


text:
# BASIC FUNCTIONS

"""
    num_qubits(code::QuantumCode) -> Int

Return the number of physical qubits of the code.
"""
function num_qubits(code::QuantumCode)
    return length(code.stabilizers) == 0 ? 0 : length(code.stabilizers[1])
end

"""
    verify(code::QuantumCode; log_warn=true) -> Bool

Return true if the code satisfied the properties of a valid code, or false otherwise.  If the
code is not valid and `log_warn` is true then a warning is logged with the specific reason.

The following checks are performed:
* Number of stabilizers, pure errors and logicals are consistent.
* Stabilizers are independent and mutually commute.
* Pure errors anticommute with corresponding stabilizers and commute with other stabilizers.
* Logicals commute with stabilizers.

The following checks are not yet performed:
* Logical commutation relations.
"""
function verify(code::QuantumCode; log_warn=true)
    n = num_qubits(code)
    r = length(code.stabilizers)
    p = length(code.pure_errors)
    l = length(code.logicals)

    # Check we have the right number of operators.
    if p != r
        log_warn && @warn "number of stabilizers and pure errors don't match!"
        return false
    end
    if n != r + l / 2
        log_warn && @warn "number of stabilizers and logicals don't add up!"
        return false
    end
    # Check stabilizers are independent.
    if !pauli_are_independent(code.stabilizers)
        log_warn && @warn "stabilizers aren't independent!"
        return false
    end
    # Check stabilizers mutually commute.
    if !pauli_are_commuting(code.stabilizers)
        log_warn && @warn "stabilizers don't commute!"
        return false
    end
    # Check pure errors anticommute with corresponding stabilizers.
    if !all(==(1), pauli_commutation.(code.stabilizers, code.pure_errors))
        log_warn && @warn "pure errors don't anticommute with corresponding stabilizers!"
        return false
    end
    # Check pure error do not anticommute with other stabilizers.
    for α in 1:r, β in 1:r
        if α != β && pauli_commutation(code.stabilizers[α], code.pure_errors[β]) == 1
            log_warn && @warn "pure errors anticommute with the wrong stabilizers!"
            return false
        end
    end
    # Check logicals commute with stabilizers.
    for logical in code.logicals
        if !all(==(0), pauli_commutation.(code.stabilizers, Ref(logical)))
            log_warn && @warn "logicals don't commute with stabilizers!"
            return false
        end
    end
    # TODO: check logical commutation relations.

    return true
end

# EVALUATION FUNCTIONS

"""
    find_distance_logicals(code::Quantum_code; max_distance=5) -> Int, Vector{Vector{Int}}

Return the distance of the code and all minimum-weight logical operators.

This method works by brute force. It searches for operators of increasing weight so it works
well for low-distance codes but will be slow for high-distance codes. If during the search
`max_distance` is exceeded then an `ErrorException` is thrown.

# Examples
```jldoctest
julia> d, ls = find_distance_logicals(five_qubit_code());

julia> d, length(ls), ls[1]  # distance, number and example of minimum-weight logicals
(3, 30, [1, 2, 1, 0, 0])
```
"""
function find_distance_logicals(code::QuantumCode; max_distance=5)
    n = num_qubits(code)
    r = length(code.stabilizers)

    lowest_weight_logicals = Vector{Int}[]

    if n == 0  # no qubits so distance is 0
        return 0, lowest_weight_logicals
    end
    if n == r  # no logicals, so distance is n (or ∞?)
        return n, lowest_weight_logicals
    end

    # find distance and lowest_weight_logicals
    distance = n - 1
    locations_iterator = combinations(1:n)
    for locations in locations_iterator

        L = length(locations)

        if L > distance && length(lowest_weight_logicals) > 0  # we have finished
            return distance, lowest_weight_logicals
        end

        if L > max_distance  # we have exceeded max_distance, so throw exception
            error("distance > max_distance = $max_distance; stopping search")
        end

        for l in 1:4^L - 1
            paulis = digits!(zeros(Int, L), l, base=4) # a nice iterator would be better

            operator = zeros(Int, n)
            for index in 1:L
                operator[locations[index]] = paulis[index]
            end

            if all(==(0), pauli_commutation.(Ref(operator), code.stabilizers))
                if (operator in code.logicals
                        || any(==(1), pauli_commutation.(Ref(operator), code.logicals)))
                    distance = L  # = pauli_weight(operator)
                    push!(lowest_weight_logicals, operator)
                end
            end
        end
    end
end

"""
    _find_product_indices(
        operators::AbstractVector{<:AbstractVector{Int}},
        target_operator::AbstractVector{Int}
    ) -> Vector{Int}

Return the indices of the smallest number of operators that as a product give the target
operator. If no such product exists an `ErrorException` is thrown.
"""
function _find_product_indices(operators, target_operator)
    identity = zeros(Int, length(target_operator))
    if target_operator == identity
        return Int[]
    end
    for i in combinations(collect(enumerate(operators)))
        ids, ops = collect(zip(i...))
        op = reduce((a, b)->pauli_product.(a, b), ops; init=identity)
        if op == target_operator
            return collect(ids)
        end
    end
    error("no product of operators yields operator")
end

"""
    find_pure_error(code::QuantumCode, syndrome::AbstractVector{Int}) -> AbstractVector{Int}

Return a pure error which yields the given syndrome with the given code.

The pure error is formed from a product of `code.pure_errors` and is not unique nor
necessarily the lowest-weight error corresponding to the syndrome.

See also [`find_syndrome`](@ref).

# Examples
```jldoctest
julia> code = five_qubit_code();

julia> syndrome = find_syndrome(code, [0, 1, 3, 0, 1])  # error = IXZIX
4-element Vector{Int64}:
 1
 0
 0
 1

julia> pure_error = find_pure_error(code, syndrome)
5-element Vector{Int64}:
 1
 1
 0
 0
 0

julia> find_syndrome(code, pure_error) == syndrome
true
```
"""
function find_pure_error(code::QuantumCode, syndrome::AbstractVector{Int})
    (length(code.pure_errors) == length(syndrome)) || error("invalid syndrome length")
    return pauli_product_pow(code.pure_errors, syndrome)
end

"""
    _find_pure_errors_disordered(stabilizers::AbstractVector{<:AbstractVector{Int}})
        -> Vector{Vector{Int}}

Return pure errors corresponding to a list of stabilizers. The returned list of pure errors
can be used (combined) to find a list of pure errors such that each pure error anticommutes
with precisely one stabilizer; this can be achieved by passing the output of this method to
[`_fix_pure_errors!`](@ref).

See also [`_fix_pure_errors!`](@ref), [`find_pure_errors`](@ref).
"""
function _find_pure_errors_disordered(stabilizers::AbstractVector{<:AbstractVector{Int}})
    stabilizers = copy(stabilizers)  # avoid modifying incoming stabilizers
    pure_errors = Vector{Vector{Int}}()
    r = length(stabilizers)
    if r == 0  # empty stabilizers, so return empty pure_errors
        return pure_errors
    end
    n = length(stabilizers[1])
    remaining = collect(1:r)

    for qubit in 1:r  # TODO: should be 1:n ?
        paulis = [1,2,3]
        indices = Int[]
        for α in remaining
            if stabilizers[α][qubit] in paulis
                setdiff!(paulis, stabilizers[α][qubit])
                push!(indices, α)
            end
            if length(paulis) == 1
                new_pure_error1 = zeros(Int64, n)
                new_pure_error2 = zeros(Int64, n)
                new_pure_error1[qubit] = stabilizers[indices[end]][qubit]
                new_pure_error2[qubit] = stabilizers[indices[end - 1]][qubit]
                push!(pure_errors, new_pure_error1)
                push!(pure_errors, new_pure_error2)
                break
            end
            if α == remaining[end] && length(paulis) == 2
                new_pure_error1 = zeros(Int64, n)
                new_pure_error1[qubit] = paulis[1]
                push!(pure_errors, new_pure_error1)
            end
        end

        remaining = setdiff(remaining, indices)
        for α in remaining
            if stabilizers[α][qubit] == 0
                continue
            end
            new_operator = pauli_product.(stabilizers[α], stabilizers[indices[1]])
            if new_operator[qubit] == 0
                stabilizers[α] = deepcopy(new_operator)
                continue
            end
            if length(paulis) == 2
                continue
            end
            new_operator = pauli_product.(stabilizers[α], stabilizers[indices[2]])
            if new_operator[qubit] == 0
                stabilizers[α] = deepcopy(new_operator)
                continue
            end
            new_operator = pauli_product.(stabilizers[α], stabilizers[indices[1]])
            new_operator = pauli_product.(new_operator, stabilizers[indices[2]])
            if new_operator[qubit] == 0
                stabilizers[α] = deepcopy(new_operator)
                continue
            end
        end
    end

    return pure_errors
end

"""
    _fix_pure_errors!(
        pure_errors::AbstractVector{<:AbstractVector{Int}},
        stabilizers::AbstractVector{<:AbstractVector{Int}}
    )

Reorder and take products of the pure errors so that each pure error anticommutes with
precisely one stabilizer and the order of pure errors respects that of the stabilizers. An
`ErrorException` is thrown if the function cannot succeed; for example if the stabilizers
are not linearly independent.

See also [`_find_pure_errors_disordered`](@ref), [`find_pure_errors`](@ref).
"""
function _fix_pure_errors!(
    pure_errors::AbstractVector{<:AbstractVector{Int}},
    stabilizers::AbstractVector{<:AbstractVector{Int}}
)
    r = length(pure_errors)

    success = 0
    for α in 1:r
        for β in α:r
            if pauli_commutation(pure_errors[β], stabilizers[α]) != 0
                pure_errors[α], pure_errors[β] = pure_errors[β], pure_errors[α]
                success += 1
                break
            end
        end
        for β in 1:r
            if α != β && pauli_commutation(pure_errors[β], stabilizers[α]) != 0
                pure_errors[β] = pauli_product.(pure_errors[α], pure_errors[β])
            end
        end
    end

    # Check if enough pure_errors were found
    if success != r
        error("pure errors cannot be fixed to correspond to stabilizers")
    end
end

"""
    find_pure_errors(stabilizers::AbstractVector{<:AbstractVector{Int}})
        -> Vector{Vector{Int}}

Return pure errors corresponding to a list of stabilizers, such that each pure error
anticommutes with precisely one stabilizer and the order of pure errors respects that of the
stabilizers.

This function is efficient but does not give lowest weight pure errors (you cannot have both
of these properties). An `ErrorException` is thrown if the function cannot succeed; for
example if the stabilizers are not linearly independent.

# Examples
```jldoctest
julia> stabilizers = [[1, 3, 3, 1, 0], [0, 1, 3, 3, 1], [1, 0, 1, 3, 3], [3, 1, 0, 1, 3]];

julia> pure_errors = find_pure_errors(stabilizers)
4-element Vector{Vector{Int64}}:
 [0, 1, 0, 0, 0]
 [1, 3, 0, 0, 0]
 [3, 1, 0, 0, 0]
 [1, 0, 0, 0, 0]

julia> [pauli_commutation(s, p) for s in stabilizers, p in pure_errors]  # commutations
4×4 Matrix{Int64}:
 1  0  0  0
 0  1  0  0
 0  0  1  0
 0  0  0  1
```
"""
function find_pure_errors(stabilizers::Array{Array{Int64,1},1})
    pure_errors = _find_pure_errors_disordered(stabilizers)
    _fix_pure_errors!(pure_errors, stabilizers)
    return pure_errors
end

"""
    find_syndrome(code::QuantumCode, error::AbstractVector{Int}) -> AbstractVector{Int}

Return the syndrome yielded by the given error with the given code.

The syndrome is a list of 1 and 0 of the same length as `code.stabilizers`, where 1
indicates the error anticommutes with the corresponding stabilizer.

See also [`find_pure_error`](@ref).

# Examples
```jldoctest
julia> syndrome = find_syndrome(five_qubit_code(), [0, 1, 3, 0, 1])  # error = IXZIX
4-element Vector{Int64}:
 1
 0
 0
 1
```
"""
function find_syndrome(code::QuantumCode, error_operator::AbstractVector{Int})
    num_qubits(code) == length(error_operator) || error("invalid error operator length")
    return pauli_commutation.(code.stabilizers, Ref(error_operator))
end

# TRANSFORMATION FUNCTIONS

"""
    gauge(code::SimpleCode, logical_qubit::Int, logical_pauli::Int) -> SimpleCode

    gauge(code::TensorNetworkCode, logical_qubit::Int, logical_pauli::Int)
        -> TensorNetworkCode

Given a code with ``k`` logicals on ``n`` physical qubits, return a new code with ``k - 1``
logicals on ``n`` physical qubits by adding a logical operator as a stabilizer, where
`logical_qubit` indexes which logical qubit is gauged and `logical_pauli` indicates which
logical Pauli is added to the stabilizers.

A `ErrorException` is thrown if `logical_qubit` indexes a non-existant logical qubit, or if
`logical_pauli` is not in `1:3` (logical identity does not fix a gauge).

# Examples
```jldoctest
julia> code = five_qubit_code();

julia> code.stabilizers
4-element Vector{Vector{Int64}}:
 [1, 3, 3, 1, 0]
 [0, 1, 3, 3, 1]
 [1, 0, 1, 3, 3]
 [3, 1, 0, 1, 3]

julia> code.logicals
2-element Vector{Vector{Int64}}:
 [1, 1, 1, 1, 1]
 [3, 3, 3, 3, 3]

julia> new_code = gauge(code, 1, 3);  # gauge logical qubit 1 using logical Z

julia> new_code.stabilizers
5-element Vector{Vector{Int64}}:
 [1, 3, 3, 1, 0]
 [0, 1, 3, 3, 1]
 [1, 0, 1, 3, 3]
 [3, 1, 0, 1, 3]
 [3, 3, 3, 3, 3]

julia> new_code.logicals
Vector{Int64}[]
```
"""
function gauge(code::SimpleCode, logical_qubit::Int, logical_pauli::Int)
    k = Int(length(code.logicals) / 2)
    # preconditions
    (logical_qubit in 1:k) || error("logical qubit index out of bounds!")
    (logical_pauli in 0:3) || error("unknown logical pauli operator!")
    (logical_pauli != 0) || error("logical identity doesn't fix a gauge!")
    # new code fields
    name = "$logical_qubit/$(pauli_rep_change(logical_pauli)) gauged $(code.name)"
    output_stabilizers = deepcopy(code.stabilizers)
    output_logicals = deepcopy(code.logicals)
    output_pure_errors = deepcopy(code.pure_errors)
    # remove gauged logicals
    gauged_logicals = [
        popat!(output_logicals, 2*logical_qubit-1),
        popat!(output_logicals, 2*logical_qubit-1),
    ]
    # new stabilizer from gauged logicals
    logical_powers = [[0, 0], [1, 0], [1, 1], [0, 1]][logical_pauli + 1]
    new_stabilizer = pauli_product_pow(gauged_logicals, logical_powers)
    push!(output_stabilizers, new_stabilizer)
    # new pure error from gauged logicals
    logical_powers = [[0, 0], [1, 1], [0, 1], [1, 1]][logical_pauli + 1]
    new_pure_error = pauli_product_pow(gauged_logicals, logical_powers)
    push!(output_pure_errors, new_pure_error)
    _fix_pure_errors!(output_pure_errors, output_stabilizers)

    return SimpleCode(name, output_stabilizers, output_logicals, output_pure_errors)
end
function gauge(code::TensorNetworkCode, logical_qubit::Int, logical_pauli::Int)
    new_code = SimpleCode(code)
    new_code = gauge(new_code, logical_qubit, logical_pauli)
    return TensorNetworkCode(new_code, code.code_graph, code.seed_codes)
end

"""
    permute(code::SimpleCode, permutation) -> SimpleCode

Return a new simple code with the physical qubits permuted relative to the given code,
according to the permutation.

The `permutation` is expected in the format used for `Base.permute!` and it is applied to
each stabilizer, logical and pure error of the code. No checking is done to verify that
`permuation` is a valid.

# Examples
```jldoctest
julia> code = five_qubit_code();

julia> code.name
"Five qubit code"

julia> code.stabilizers
4-element Vector{Vector{Int64}}:
 [1, 3, 3, 1, 0]
 [0, 1, 3, 3, 1]
 [1, 0, 1, 3, 3]
 [3, 1, 0, 1, 3]

julia> new_code = permute(code, [2, 1, 3, 4, 5]);

julia> new_code.name
"Five qubit code [2, 1, 3, 4, 5]"

julia> new_code.stabilizers
4-element Vector{Vector{Int64}}:
 [3, 1, 3, 1, 0]
 [1, 0, 3, 3, 1]
 [0, 1, 1, 3, 3]
 [1, 3, 0, 1, 3]
```
"""
function permute(code::SimpleCode, permutation)
    output_code = SimpleCode("$(code.name) $(permutation)",
        deepcopy(code.stabilizers),
        deepcopy(code.logicals),
        deepcopy(code.pure_errors)
    )
    permute!.(output_code.stabilizers, Ref(permutation))
    permute!.(output_code.logicals, Ref(permutation))
    permute!.(output_code.pure_errors, Ref(permutation))
    return output_code
end

"""
    purify(code::SimpleCode) -> SimpleCode

Given a simple code with ``k`` logicals on ``n`` physical qubits, return a new simple code
with ``0`` logicals on ``n + k`` physical qubits.

# Examples
```jldoctest
julia> code = purify(five_qubit_code());

julia> num_qubits(code), length(code.logicals)
(6, 0)
```
"""
function purify(code::SimpleCode)
    g = code.stabilizers
    l = code.logicals
    K = length(l)
    k = Int(K / 2)
    n = num_qubits(code)

    output_stabilizers = vcat.(Ref(zeros(Int, k)), g)
    new_stabilizers = vcat.(Ref(zeros(Int, k)), l)
    output_pure_errors = vcat.(Ref(zeros(Int, k)), code.pure_errors)
    new_pure_errors = [zeros(Int, n + k) for _ in 1:K]

    for α in 1:length(new_stabilizers)
        if isodd(α)
            β = Int((α + 1) ÷ 2)
            new_stabilizers[α][β] = 1
            new_pure_errors[α][β] = 3
        elseif iseven(α)
            β = Int(α ÷ 2)
            new_stabilizers[α][β] = 3
            new_pure_errors[α][β] = 1
        end
    end

    name = "Purified $(code.name)"
    output_stabilizers = vcat(output_stabilizers, new_stabilizers)
    output_logicals = Vector{Int}[]
    output_pure_errors = vcat(output_pure_errors, new_pure_errors)
    _fix_pure_errors!(output_pure_errors, output_stabilizers)

    return SimpleCode(name, output_stabilizers, output_logicals, output_pure_errors)
end


text:
include("../nanomsg.jl")

function node0(url)
    sock = nn_socket(NN_PULL)
    check_err("nn_socket", sock)
    endpoint = nn_bind(s, url)
    check_err("nn_bind", endpoint)
    while true
        msg = nn_recv(s)
        check_err("nn_recv", msg)
        println("NODE0: RECEIVED ", msg)
    end
end

function node1(url, msg)
    sock = nn_socket(NN_PUSH)
    check_err("nn_socket", sock)
    rv = nn_connect(sock, url)
    check_err("nn_connect", rv)
    bytes = nn_send(sock, msg)
    check_err("nn_send", bytes)
    sleep(1) #  wait for messages to flush before shutting down
    rv = nn_shutdown(s, 0)
    check_err("nn_shutdown", rv)
end

node0("ipc:///tmp/pipeline.ipc")
node1("ipc:///tmp/pipeline.ipc", "hello, world its me")


text:
mutable struct Estep{Tv,MatT}
    etotal::Tv
    eb::Vector{Tv}
    ey::Vector{Tv}
    ez::Vector{Tv}
    en::MatT
end

function Estep(ph::GPH{Tv,MatT}) where {Tv,MatT}
    Estep(Tv(0),
        Vector{Tv}(undef, ph.dim),
        Vector{Tv}(undef, ph.dim),
        Vector{Tv}(undef, ph.dim),
        similar(ph.T))
end

function estep!(ph::CF1{Tv}, data::WeightedSample{Tv}, eres::Estep{Tv,MatT};
    eps::Tv = Tv(1.0e-8), ufact::Tv = Tv(1.01)) where {Tv,MatT}
    estep!(GPH(ph, MatT), data, eres, eps=eps, ufact=ufact)
end

@origin (vf => 0, vb => 0, vc => 0) function estep!(ph::GPH{Tv,MatT},
    data::WeightedSample{Tv}, eres::Estep{Tv,MatT}; eps::Tv = Tv(1.0e-8), ufact::Tv = Tv(1.01)) where {Tv,MatT}

    dim, alpha, P, tau, qv = _phunif(ph, ufact)

    llf = Tv(0)
    tllf = Tv(0)
    # eres.etotal = Tv(0)
    eres.eb .= Tv(0)
    eres.ey .= Tv(0)
    # eres.ez .= Tv(0)
    for i = 1:length(eres.en)
        eres.en[i] = Tv(0)
    end
    right = rightbound(qv * data.maxtime, eps) + 1
    poi = Vector{Tv}(undef, right + 1)

    m = data.length
    blf = Vector{Tv}(undef, m)
    vf = Vector{Vector{Tv}}(undef, m+1)
    vb = Vector{Vector{Tv}}(undef, m+1)
    vc = Vector{Vector{Tv}}(undef, m+1)
    vf[0] = alpha
    vb[0] = tau
    for k = 1:m
        right = rightbound(qv * data.tdat[k], eps) + 1
        weight = poipmf!(qv * data.tdat[k], poi, left=0, right=right)
        vf[k] = zero(alpha)
        unifstep!(:T, P, poi, (0, right), weight, copy(vf[k-1]), vf[k])
        scale = @dot(vf[k], tau)
        @scal(1.0/scale, vf[k])
        @axpy(data.wdat[k], vf[k], eres.ey)

        blf[k] = scale
        vb[k] = zero(alpha)
        unifstep!(:N, P, poi, (0, right), weight, copy(vb[k-1]), vb[k])
        scale = @dot(alpha, vb[k])
        @scal(1.0/scale, vb[k])
        @axpy(data.wdat[k], vb[k], eres.eb)

        tllf += log(blf[k])
        llf += data.wdat[k] * tllf
    end

    vc[m] = zeros(Tv, dim)
    @axpy(data.wdat[m]/blf[m], alpha, vc[m])
    for k = m-1:-1:1
        right = rightbound(qv * data.tdat[k+1], eps) + 1
        weight = poipmf!(qv * data.tdat[k+1], poi, left=0, right=right)
        vc[k] = zeros(Tv, dim)
        unifstep!(:T, P, poi, (0, right), weight, copy(vc[k+1]), vc[k])
        @scal(1.0/blf[k], vc[k])
        @axpy(data.wdat[k]/blf[k], alpha, vc[k])
    end

    tmpb = similar(alpha)
    tmpn = zero(P)
    for k = 1:m
        right = rightbound(qv * data.tdat[k], eps) + 1
        weight = poipmf!(qv * data.tdat[k], poi, left=0, right=right)
        tmpb .= Tv(0)
        convunifstep!(:T, :N, P, poi, (0, right), weight, qv*weight,
            vc[k], vb[k-1], tmpb, tmpn)
        for i = 1:length(tmpn)
            eres.en[i] += tmpn[i]
            tmpn[i] = Tv(0)
        end
    end

    eres.eb .*= alpha
    eres.ey .*= tau
    eres.ez .= spdiag(eres.en)
    for i = 1:length(eres.en)
        eres.en[i] *= ph.T[i]
    end
    eres.etotal = sum(eres.eb)

    return llf
end

function mstep!(ph::GPH{Tv,Matrix{Tv}}, eres::Estep{Tv,Matrix{Tv}}) where Tv
    dim = ph.dim
    tmp = zeros(Tv, dim)
    for j = 1:dim
        for i = 1:dim
            if i != j
                ph.T[i,j] = eres.en[i,j] / eres.ez[i]
                tmp[i] += ph.T[i,j]
            end
        end
    end
    for i = 1:dim
        ph.alpha[i] = eres.eb[i] / eres.etotal
        ph.tau[i] = eres.ey[i] / eres.ez[i]
        tmp[i] += ph.tau[i]
        ph.T[i,i] = -tmp[i]
    end
    nothing
end

function mstep!(ph::GPH{Tv,SparseCSR{Tv,Ti}}, eres::Estep{Tv,SparseCSR{Tv,Ti}}) where {Tv,Ti}
    dim = ph.dim
    tmp = zeros(Tv, dim)
    d = zeros(Ti, dim)
    for i = 1:dim
        for z = eres.en.rowptr[i]:eres.en.rowptr[i+1]-1
            j = eres.en.colind[z]
            if i != j
                ph.T[z] = eres.en[z] / eres.ez[i]
                tmp[i] += ph.T[z]
            else
                d[i] = z
            end
        end
    end
    for i = 1:dim
        ph.alpha[i] = eres.eb[i] / eres.etotal
        ph.tau[i] = eres.ey[i] / eres.ez[i]
        tmp[i] += ph.tau[i]
        ph.T.val[d[i]] = -tmp[i]
    end
    nothing
end

function mstep!(ph::GPH{Tv,SparseCSC{Tv,Ti}}, eres::Estep{Tv,SparseCSC{Tv,Ti}}) where {Tv,Ti}
    dim = ph.dim
    tmp = zeros(Tv, dim)
    d = zeros(Ti, dim)
    for j = 1:dim
        for z = eres.en.colptr[j]:eres.en.colptr[j+1]-1
            i = eres.en.rowind[z]
            if i != j
                ph.T[z] = eres.en[z] / eres.ez[i]
                tmp[i] += ph.T[z]
            else
                d[i] = z
            end
        end
    end
    for i = 1:dim
        ph.alpha[i] = eres.eb[i] / eres.etotal
        ph.tau[i] = eres.ey[i] / eres.ez[i]
        tmp[i] += ph.tau[i]
        ph.T.val[d[i]] = -tmp[i]
    end
    nothing
end

function mstep!(ph::GPH{Tv,SparseMatrixCSC{Tv,Ti}}, eres::Estep{Tv,SparseMatrixCSC{Tv,Ti}}) where {Tv,Ti}
    dim = ph.dim
    tmp = zeros(Tv, dim)
    d = zeros(Ti, dim)
    for j = 1:dim
        for z = eres.en.colptr[j]:eres.en.colptr[j+1]-1
            i = eres.en.rowval[z]
            if i != j
                ph.T[z] = eres.en[z] / eres.ez[i]
                tmp[i] += ph.T[z]
            else
                d[i] = z
            end
        end
    end
    for i = 1:dim
        ph.alpha[i] = eres.eb[i] / eres.etotal
        ph.tau[i] = eres.ey[i] / eres.ez[i]
        tmp[i] += ph.tau[i]
        ph.T.nzval[d[i]] = -tmp[i]
    end
    nothing
end

function mstep!(ph::GPH{Tv,SparseCOO{Tv,Ti}}, eres::Estep{Tv,SparseCOO{Tv,Ti}}) where {Tv,Ti}
    dim = ph.dim
    tmp = zeros(Tv, dim)
    d = zeros(Ti, dim)
    for z = 1:nnz(eres.en)
        i = eres.en.rowind[z]
        j = eres.en.colind[z]
        if i != j
            ph.T[z] = eres.en[z] / eres.ez[i]
            tmp[i] += ph.T[z]
        else
            d[i] = z
        end
    end
    for i = 1:dim
        ph.alpha[i] = eres.eb[i] / eres.etotal
        ph.tau[i] = eres.ey[i] / eres.ez[i]
        tmp[i] += ph.tau[i]
        ph.T.val[d[i]] = -tmp[i]
    end
    nothing
end

function mstep!(cf1::CF1{Tv}, eres::Estep{Tv,MatT}; cf1sort = true) where {Tv,MatT}
    dim = cf1.dim
    total = Tv(0)
    for i = 1:dim
        total += eres.eb[i]
        cf1.alpha[i] = eres.eb[i] / eres.etotal
        cf1.rate[i] = total / eres.ez[i]
    end
    cf1sort && _cf1_sort!(cf1.alpha, cf1.rate)
    nothing
end