text:
module GraphPartition

using ..GraphSignal, SparseArrays, LinearAlgebra

include("partition_fiedler.jl")
include("utils.jl")

export GraphPart, partition_tree_fiedler, partition_tree_matrixDhillon

# GraphPart data structure and constructors
"""
    GP = GraphPart(ind::Vector{Int}, rs::Matrix{Int}, tag::Matrix{Int}, compinfo::Matrix{Int}, rsf2c::Matrix{Int}, tagf2c::Matrix{Int}, compinfof2c::Matrix{Int}, inds::Matrix{Int}, method::Vector{Symbol})

is a data structure for a GraphPart object containing the following fields:
* `ind::Vector{Int}`: ordering of the indices on the finest level
* `rs::Matrix{Int}`: regionstarInt (coarse-to-fine) <==> the index in `ind` of the first point in region number `i` is `rs[i]`
* `tag::Matrix{Int}`: tag info for the GHWT coarse-to-fine basis
* `compinfo::Matrix{Int}`: indicates whether the coefficient was formed from 2 coefficenInt (value is nonzero) or from only 1 coefficient (value is zero); when a scaling and Haar-like coefficient are formed, their corresponding values in compinfo indicate the number of nodes in each of the 2 subregions
* `rsf2c::Matrix{Int}`: the fine-to-coarse version of `rs`
* `tagf2c::Matrix{Int}`: the fine-to-coarse version of `tag`
* `compinfof2c::Matrix{Int}`: the fine-to-coarse version of `compinfo`
* `inds::Matrix{Int}`: ordering of the indices on all levels
* `method::Symbol`: how the partition tree was constructed

The unsigned integer depends on the size of the underlying graph.

Copyright 2015 The RegenInt of the University of California

Implemented by Jeff Irion (Adviser: Dr. Naoki Saito) |
Translated and modified by Naoki Saito, Feb. 7, 2017
Revised for two parameters by Naoki Saito, Feb. 24, 2017
"""
mutable struct GraphPart
    ind::Vector{Int}    # ordering of the indices on the finest level
    rs::Matrix{Int}     # `rs[i,j]` = the index in `ind` of the first
                       # point in Region `i` at level j
    tag::Matrix{Int}    # Int <: Int because max(tag) <= log2(max(T))
    compinfo::Matrix{Int}      # coef was from 2 coefs or from 1 coef
    rsf2c::Matrix{Int}         # f2c version of `rs`
    tagf2c::Matrix{Int}        # f2c version of `tag`
    compinfof2c::Matrix{Int}   # f2c version of `compinfo`
    inds::Matrix{Int}   # ordering of the indices on all levels
    method::Symbol           # specification of graph partition method

    # An inner constructor here.
    function GraphPart(ind::Vector{Int}, rs::Matrix{Int};
                       tag::Matrix{Int} = Matrix{Int}(undef, 0, 0),
                       compinfo::Matrix{Int} = Matrix{Int}(undef, 0, 0),
                       rsf2c::Matrix{Int} = Matrix{Int}(undef, 0, 0),
                       tagf2c::Matrix{Int} = Matrix{Int}(undef, 0, 0),
                       compinfof2c::Matrix{Int} = Matrix{Int}(undef, 0, 0),
                       inds::Matrix{Int} = Matrix{Int}(undef, 0, 0),
                       method::Symbol = :unspecified)

        # Sanity checks
        if Base.size(rs, 1) != Base.length(ind) + 1
            @warn("size(rs,1) must be length(ind) + 1")
            @warn("both rs and ind now become null arrays!")
            ind = Vector{Int}(undef, 0)
            rs = Matrix{Int}(undef, 0, 0)
        end
        if Base.length(tag) != 0 && (Base.size(rs, 1) != Base.size(tag, 1) + 1
            || Base.size(rs, 2) != Base.size(tag, 2))
            @warn("tag size is inconsistent with rs size.")
            @warn("tag now becomes a null array!")
            tag = Matrix{Int}(undef, 0, 0)
        end
        if Base.length(compinfo) != 0 &&
            ( Base.size(rs, 1) != Base.size(compinfo, 1) + 1
             || Base.size(rs, 2) != Base.size(compinfo, 2) )
            @warn("compinfo size is inconsistent with rs size.")
            @warn("compinfo now becomes a null array!")
            compinfo = Matrix{Int}(undef, 0, 0)
        end
        if Base.length(rsf2c) != 0 && Base.length(rsf2c) != Base.length(rs)
            @warn("length(rsf2c) must be length(rs)")
            @warn("rsf2c now becomes a null array!")
            rsf2c = Matrix{Int}(undef, 0, 0)
        end
        if Base.length(tagf2c) != 0 && Base.length(tagf2c) != Base.length(tag)
            @warn("length(tagf2c) must be length(tag)")
            @warn("tagf2c now becomes a null array!")
            tagf2c = Matrix{Int}(undef, 0, 0)
        end
        if Base.length(compinfof2c) != 0 &&
            Base.length(compinfof2c) != Base.length(compinfo)
            @warn("length(compinfof2c) must be length(compinfo)")
            @warn("compf2c now becomes a null array!")
            compf2c = Matrix{Int}(undef, 0, 0)
        end
        new(ind, rs, tag, compinfo, rsf2c, tagf2c, compinfof2c, inds, method)
    end # of an inner constructor GraphPart

 end # of type GraphPart

"""
    GP = partition_tree_fiedler(G::GraphSig; method::Symbol = :Lrw, swapRegion::Bool = true)

 Generate a partition tree for a graph using the Fiedler vector of either
 L (the unnormalized Laplacian) or L_rw (the random-walk normalized
 Laplacian).

### Input ArgumenInt
* `G::GraphSig`: an input GraphSig object
* `method::Symbol`: how the partition tree was constructed (default: :Lrw)
* `swapRegion::Bool`: swap the child regions based on the sum of the subgraphs' edge weights (default: true)

### Output Argument
* `GP::GraphPart`: an ouput GraphPart object

Copyright 2015 The RegenInt of the University of California

Implemented by Jeff Irion (Adviser: Dr. Naoki Saito) |
Translated and modified by Naoki Saito, Feb. 7, 2017
"""
function partition_tree_fiedler(G::GraphSignal.GraphSig; method::Symbol = :Lrw, swapRegion::Bool = true)
    #
    # 0. Preliminary stuff
    #
    # constanInt
    N = G.length
    jmax = max(3 * floor(Int, log2(N)), 4) # jmax >= 4 is guaranteed.
    # This jmax is the upper bound of the true jmax; the true jmax
    # is computed as the number of columns of the matrix `rs` in the end.

    # TRACKING VARIABLES

    # `ind` records the way in which the nodes are indexed on each level
    ind = Vector{Int}(1:N)

    # `inds` records the way in which the nodes are indexed on all levels
    inds = zeros(Int, N, jmax)
    inds[:, 1] = ind

    # `rs` stands for regionstarInt, meaning that the index in `ind` of the first
    # point in region number `i` is `rs[i]`
    rs = zeros(Int, N + 1, jmax)
    rs[1, :] .= 1
    rs[2, 1] = N + 1

    #
    # 1. Partition the graph to yield rs, ind, and jmax
    #
    j = 1                       # j=1 in julia <=> j=0 in theory
    regioncount = 0
    rs0 = 0                     # define here for the whole loops,
                                # which differs from MAIntAB.
    while regioncount < N
        #regioncount = countnz(rs[:,j]) - 1
        regioncount = count(!iszero,rs[:, j]) - 1 # the number of regions on level j
        if j == jmax  # add a column to rs and inds for level j+1, if necessary
            rs = hcat(rs, vcat(Int(1), zeros(Int,N)))
            inds = hcat(inds, zeros(Int, N))
            jmax = jmax + 1
        end
        # for tracking the child regions
        rr = 1
        for r = 1:regioncount   # cycle through the parent regions
            rs1 = rs[r, j]      # the start of the parent region
            rs2 = rs[r + 1, j] # 1 node after the end of the parent region
            n = rs2 - rs1   # the number of nodes in the parent region

            if n > 1            # regions with 2 or more nodes
                indrs = ind[rs1:(rs2 - 1)]
                # partition the current region
                (pm, ) = partition_fiedler(G.W[indrs,indrs], method = method)
                # determine the number of poinInt in child region 1
                n1 = sum(pm .> 0)
                # switch regions 1 and 2, if necessary, based on the sum of
                # edge weighInt to the previous child region (i.e., cousin)
                if r > 1 && swapRegion
                    # MAIntAB: sum(sum(...)) instead of sum(...)
                    if sum(G.W[ind[rs0:(rs1 - 1)], indrs[pm .> 0]]) <
                        sum(G.W[ind[rs0:(rs1 - 1)], indrs[pm .< 0]])
                        pm = -pm
                        n1 = n - n1
                    end
                end
                # update the indexing
                ind[rs1:(rs1 + n1 - 1)] = indrs[pm .> 0]
                ind[(rs1 + n1):(rs2 - 1)] = indrs[pm .< 0]
                # update the region tracking
                rs[rr + 1, j + 1] = rs1 + n1
                rs[rr + 2, j + 1] = rs2
                rr = rr + 2
                rs0 = rs1 + n1
            elseif n == 1       # regions with 1 node
                rs[rr + 1, j + 1] = rs2
                rr = rr + 1
                rs0 = rs1
            end # of if n > 1 ... elseif n==1 construct
        end # of for r=1:regioncount
        j = j + 1
        inds[:, j] = ind
    end # of while regioncount < N statement

    #
    # 2. Postprocessing
    #
    # get rid of excess columns in rs
    rs = rs[:, 1:(j - 1)]         # in MAIntAB, it was rs(:,j:end) = [];
    # get rid of excess columns in inds
    inds = inds[:, 1:(j - 1)]
    # create a GraphPart object
    return GraphPart(ind, rs; inds = inds, method = method)
end # of function partition_tree_fiedler

"""
    GProws, GPcols = partition_tree_matrixDhillon(matrix::Matrix{Float64})

    Recursively partition the rows and columns of the matrix using the
    bipartite model proposed by Dhillon in "Co-clustering documents and words
    using Bipartite Spectral Graph Partitioning"

### Input Arguments
* `matrix::Matrix{Float64}`: an input matrix

### Output Argument
* `GProws::GraphPart`: partitioning using rows as samples
* `GPcols::GraphPart`: partitioning using cols as samples
"""
function partition_tree_matrixDhillon(matrix::Matrix{Float64})
  ## 0. Preliminary stuff

  # constants
  (rows, cols) = size(matrix)
  N = max(rows,cols)
  jmax = Int(max(3*floor(log2(N)),4))

  ### Tracking variables for the rows

  # order index
  ind_rows = Vector{Int}(1:rows)
  ind_cols = Vector{Int}(1:cols)

  # regionstarts
  rs_rows = zeros(Int, (N+1,jmax))
  rs_rows[1,:] .= 1
  rs_rows[2,1] = rows + 1

  rs_cols = zeros(Int, (N+1,jmax))
  rs_cols[1,:] .= 1
  rs_cols[2,1] = cols + 1

  ## 1. Partition the graph to yield rs, ind, and jmax
  j = 1
  regioncount_rows = 0
  regioncount_cols = 0
  while (regioncount_rows) < rows || (regioncount_cols < cols)
    # the number of true regions (i.e, has 1 or more nodes) on level j
    regioncount_rows = length(unique(vcat(rs_rows[:,j],[0]))) - 2
    regioncount_cols = length(unique(vcat(rs_cols[:,j],[0]))) - 2

    # the number of total regions (including those with no nodes)
    false_regioncount = max(count(!iszero, rs_rows[:,j]),count(!iszero, rs_cols[:,j])) - 1

    # add a column for level j+1, if necessary
    if j == jmax
      rs_rows = hcat(rs_rows,zeros(N+1))
      rs_cols = hcat(rs_cols,zeros(N+1))
      rs_rows[1,j+1] = 1
      rs_cols[1,j+1] = 1
      jmax = jmax + 1
    end

    # for tracking the child regions (for both rows and columns)
    rr  = 1

    # cycle through the parent regions
    for r = 1 : false_regioncount
      # the start of the parent region and 1 node after the end of the parent region
      rs1_rows = rs_rows[r,j]
      rs2_rows = rs_rows[r+1,j]
      rs1_cols = rs_cols[r,j]
      rs2_cols = rs_cols[r+1,j]

      # the number of nodes in the parent region
      n_rows = rs2_rows - rs1_rows
      n_cols = rs2_cols - rs1_cols

      if (n_rows <=1) && (n_cols > 1)
        indrs_rows = ind_rows[rs1_rows]
        indrs_cols = ind_cols[rs1_cols:(rs2_cols-1)]

        # compute the left and right singular vectors

        #(u_rows, v_cols) = second_largest_singular_vectors(reshape(matrix[indrs_rows,indrs_cols],(n_rows,n_cols)))

        # partition the current region
        pm_cols = partition_vector(matrix[indrs_rows,indrs_cols])

        # determine the number of points in child region 1
        n1_cols = sum(pm_cols .>0)

        # update the column indexing
        ind_cols[rs1_cols:(rs1_cols+n1_cols-1)] = indrs_cols[pm_cols.>0]
        ind_cols[(rs1_cols+n1_cols):(rs2_cols-1)] = indrs_cols[pm_cols.<0]

        # update the row region tracking
        rs_rows[rr+1,j+1] = rs1_rows
        rs_rows[rr+2,j+1] = rs2_rows

        # update the column region tracking
        rs_cols[rr+1, j+1] = rs1_cols + n1_cols
        rs_cols[rr+2, j+1] = rs2_cols

        rr = rr + 2

      elseif (n_rows > 1) && (n_cols <= 1)
        indrs_rows = ind_rows[rs1_rows:(rs2_rows-1)]
        indrs_cols = ind_cols[rs1_cols]

        # compute the left and right singular vectors
        #(u_rows, u_cols) = second_largest_singular_vectors(reshape(matrix[indrs_rows, indrs_cols],(n_rows,n_cols)))

        # partition the current region
        pm_rows = partition_vector(matrix[indrs_rows, indrs_cols])

        # determine the number of points in child region 1
        n1_rows = sum(pm_rows.>0)

        # update the row indexing
        ind_rows[rs1_rows:(rs1_rows+n1_rows-1)] = indrs_rows[pm_rows.>0]
        ind_rows[(rs1_rows+n1_rows):(rs2_rows-1)] = indrs_rows[pm_rows.<0]

        # update the row region tracking
        rs_rows[rr+1,j+1] = rs1_rows + n1_rows
        rs_rows[rr+2,j+1] = rs2_rows

        # update the column region tracking
        rs_cols[rr+1,j+1] = rs1_cols
        rs_cols[rr+2,j+1] = rs2_cols

        rr = rr + 2

      elseif (n_rows > 1) && (n_cols > 1)
        indrs_rows = ind_rows[rs1_rows:(rs2_rows-1)]
        indrs_cols = ind_cols[rs1_cols:(rs2_cols-1)]

        # compute the left and right singular vectors
        (u_rows,v_cols) = second_largest_singular_vectors(matrix[indrs_rows,indrs_cols])

        # partition the current region
        pm, = partition_fiedler_pm(vcat(u_rows, v_cols))
        pm_rows = pm[1:n_rows]
        pm_cols = pm[(n_rows+1):end]

        # determine the number of points in child region 1
        n1_rows = sum(pm_rows .> 0)
        n1_cols = sum(pm_cols .> 0)

        # update the row indexing
        ind_rows[rs1_rows:(rs1_rows + n1_rows - 1)] = indrs_rows[pm_rows .> 0]
        ind_rows[(rs1_rows + n1_rows):(rs2_rows - 1)] = indrs_rows[pm_rows .< 0]

        # update the column indexing
        ind_cols[rs1_cols:(rs1_cols + n1_cols - 1)] = indrs_cols[pm_cols .> 0]
        ind_cols[(rs1_cols + n1_cols):(rs2_cols - 1)] = indrs_cols[pm_cols .< 0]

        # update the row region tracking
        rs_rows[rr+1, j+1] = rs1_rows + n1_rows
        rs_rows[rr+2, j+1] = rs2_rows

        # update the column region tracking
        rs_cols[rr+1, j+1] = rs1_cols + n1_cols
        rs_cols[rr+2, j+1] = rs2_cols

        rr = rr + 2

      else
        rs_rows[rr+1, j+1] = rs2_rows
        rs_cols[rr+1, j+1] = rs2_cols
        rr = rr + 1
      end
      if rr > size(rs_rows,1) - 2
        rs_rows = vcat(rs_rows,zeros(Int,(2,size(rs_rows,2))))
      end
      if rr > size(rs_cols,1) - 2
        rs_cols = vcat(rs_cols,zeros(Int,(2,size(rs_cols,2))))
      end
    end

    j = j + 1
  end

  # get rid of excess columns in rs_rows and rs_cols
  rs_rows = rs_rows[:, 1:(j-1)]
  rs_cols = rs_cols[:, 1:(j-1)]
  jmax = size(rs_rows,2)

  #remove duplicates
  M_rows = size(rs_rows,1)
  M_cols = size(rs_cols,1)
  for j = jmax:-1:1
    ### ROWS
    # remove duplicate regionstarts
    unique_rows = unique(rs_rows[:,j])
    if unique_rows[1] == 0
      shift!(unique_rows)
    end
    rs_rows[:,j] = vcat(unique_rows, zeros(M_rows-length(unique_rows)))

    #remove duplicate columns of rs
    if (j < jmax) && (sum(rs_rows[:,j] .!= rs_rows[:,j+1]) == 0)
      rs_rows = rs_rows[:,setdiff(1:end,j+1)]
    end

    ### COLUMNS
    # remove duplicate regionscounts
    unique_cols = unique(rs_cols[:,j])
    if unique_cols[1] == 0
      shift!(unique_cols)
    end
    rs_cols[:,j] = vcat(unique_cols, zeros(M_cols-length(unique_cols)))

    # remove duplicate columns of rs
    if (j < jmax) && (sum(rs_cols[:,j].!=rs_cols[:,j+1]) == 0)
      rs_cols = rs_cols[:,setdiff(1:end,j+1)]
    end
  end

  # remove unnecessary rows and columns in the regionstarts
  if rows < M_rows - 1
    rs_rows = rs_rows[1:(rows+1),:]
  end
  if cols < M_cols - 1
    rs_cols = rs_cols[1:(cols+1),:]
  end

  # create GraphPart Objects
  GProws = GraphPart(ind_rows, rs_rows)
  GPcols = GraphPart(ind_cols, rs_cols)

  return GProws, GPcols

  # make sure the GraphPart objects meet our requirements
  #blablabla

end

function second_largest_singular_vectors(A::Matrix{Float64})
  # Compute the second largest left and right singular vectors of
  # An = D1^(-0.5)*A*D2^(-0.5)

  (rows, cols) = size(A)

  # compute D1 and D2 and make sure there are no zeros
  D1 = sum(A,dims = 2)
  D1[D1.==0] .= max(0.01, minimum(D1[D1.>0]/10))
  D2 = sum(A,dims = 1)
  D2[D2.==0] .= max(0.01, minimum(D2[D2.>0]/10))

  # compute the singular vectors
  (u,D,v) = svd(Diagonal(D1[:].^(-0.5))*A*Diagonal(D2[:].^(-0.5)), full = false)

  if (rows > 1) && (cols > 1)

  # extract the second left singular vector
  u = u[:,2]
  v = v[:,2]

# extract the 2nd singular vectors and multiply by D_i^(-0.5)
  u = D1[:].^(-0.5).*u
  v = D2[:].^(-0.5).*v

  return u, v

  end

end

function partition_vector(v:: Vector{Float64})
  # Partition vector into half, diveded by the mean of the values.
  l = size(v[:],1)
  m = mean(v)
  pm = zeros(l)
  pm[v.>m] .= 1
  pm[v.<=m] .= -1
  if size(unique(pm),1) == 1
    f = UInt(floor(l/2))
    pm[1:f] .= 1
    pm[f+1:end] .= -1
  end
  return pm
end

end # of module GraphPartition


text:
#
println("Ae) Test of the CI part for an internally generated neon multiplet without Breit interaction.")
#
levelSelection = LevelSelection(true, indices = [1,2, 4,5, 7,8])
                                ## symmetries = [ LevelSymmetry(1//2,Basics.plus), LevelSymmetry(1//2,Basics.minus),  LevelSymmetry(5//2,Basics.plus)])

settings = AsfSettings(AsfSettings(), levelSelectionCI = levelSelection, eeInteractionCI=CoulombInteraction(), jjLS = LSjjSettings(false) )

wa = Atomic.Computation(Atomic.Computation(), name="xx", grid=JAC.Radial.Grid(true), nuclearModel=Nuclear.Model(26.),
                        configs=[Configuration("[Ne] 3s^2 3p^5"), Configuration("[Ne] 3s 3p^6"), Configuration("[Ne] 3s^2 3p^4 3d"),
                                 Configuration("[Ne] 3s^2 3p^3 3d^2"), Configuration("[Ne] 3s^2 3p^2 3d^3"), ],  #,
                                 # Configuration("[Ne] 3s^2 3p^2 3d^3"), Configuration("[Ne] 3s 3p^2 3d^4"), Configuration("[Ne] 3p^2 3d^5")],
                        asfSettings=settings  )

@time wb = perform(wa)

ext: jl


text:
using BinaryBuilder

cuda_version = v"10.1.243"  # NOTE: could be less specific

# NOTE: the same build is used for CUDA 10.2, but we can't express that with BinaryBuilder
sources_linux = [
    ArchiveSource("https://developer.download.nvidia.com/assets/gameworks/downloads/secure/cuTensor/libcutensor-linux-x86_64-1.0.1.tar.gz",
                  "ca6122b3f15511cd33a5eb7f911cff1553def1d3ff0b9270e62ef08f1a94f2aa")
]

script = raw"""
cd ${WORKSPACE}/srcdir
if [[ ${target} == x86_64-linux-gnu ]]; then
    cd libcutensor
    find .

    install_license license.pdf

    mv lib/10.1/libcutensor.so* ${libdir}
    mv include/* ${prefix}/include
fi
"""

include("../common.jl")


text:
using NaiveGAflux, Random, Logging, Flux, Test
using NaiveNASlib.Advanced, NaiveNASlib.Extend

@testset "NaiveGAflux.jl" begin

    mutable struct MockRng <:AbstractRNG
        seq::AbstractVector
        ind::Integer
        MockRng(seq) = new(seq, 0)
    end

    function Random.rand(rng::MockRng)
        rng.ind = rng.ind % length(rng.seq) + 1
        return rng.seq[rng.ind]
    end

    mutable struct SeqRng <: AbstractRNG
        ind
        SeqRng(ind) = new(ind)
        SeqRng() = new(0)
    end
    function Random.rand(rng::SeqRng, vec::AbstractArray)
        rng.ind = rng.ind % length(vec) + 1
        return vec[rng.ind]
    end

    @info "Testing util"
    include("util.jl")

    @info "Testing shape"
    include("shape.jl")

    @info "Testing archspace"
    include("archspace.jl")

    @info "Testing mutation"
    include("mutation.jl")

    @info "Testing crossover"
    include("crossover.jl")

    @info "Testing fitness"
    include("fitness.jl")

    @info "Testing candidate"
    include("candidate.jl")

    @info "Testing evolve"
    include("evolve.jl")

    @info "Testing population"
    include("population.jl")

    @info "Testing iterators"
    include("iterators.jl")

    @info "Testing visualization"
    include("visualization/callbacks.jl")

    if VERSION === v"1.6.2"
        @info "Testing README examples"
        include("examples.jl")
    else
        @warn "README examples will only be tested in julia version 1.6.2 due to rng dependency. Skipping..."
    end

    @info "Testing AutoFlux"
    include("app/autoflux.jl")
end


text:
### Abstract parameter IOStreams

abstract type ParameterIOStream{S<:ValueSupport} <: VariableIOStream end


text:
@testset verbose=true "Gradients" begin
    ϵ = 200eps()

    @testset verbose=true "log" begin
        ∂log∂q(q) = [
            ForwardDiff.derivative(ϵ->log(q+ϵ), 0),
            ForwardDiff.derivative(ϵ->log(q+ϵ*imx), 0),
            ForwardDiff.derivative(ϵ->log(q+ϵ*imy), 0),
            ForwardDiff.derivative(ϵ->log(q+ϵ*imz), 0)
        ]
        for q in randn(RotorF64, 10_000)
            ∂1 = ∂log(q)
            ∂2 = ∂log∂q(q)
            for i in 1:4
                @test ∂1[i] ≈ ∂2[i] rtol=ϵ
            end
            l, ∂3 = log∂log(q)
            @test l ≈ log(q) rtol=4eps() atol=10eps()
            for i in 1:4
                @test ∂3[i] ≈ ∂2[i] rtol=ϵ
            end
        end
        ∂1 = ∂log(Rotor(1))
        ∂2 = QuaternionF64[1, 0, 0, imz]
        for i in 1:4
            @test ∂1[i] ≈ ∂2[i] rtol=ϵ
        end
        l, ∂1 = log∂log(Rotor(1))
        ∂2 = QuaternionF64[1, 0, 0, imz]
        @test l ≈ QuaternionF64(0) rtol=4eps() atol=10eps()
        for i in 1:4
            @test ∂1[i] ≈ ∂2[i] rtol=ϵ
        end
    end

    @testset verbose=true "exp" begin
        ∂exp∂q(q) = [
            ForwardDiff.derivative(ϵ->exp(q+(ϵ+0imx)), 0),
            ForwardDiff.derivative(ϵ->exp(q+ϵ*imx), 0),
            ForwardDiff.derivative(ϵ->exp(q+ϵ*imy), 0),
            ForwardDiff.derivative(ϵ->exp(q+ϵ*imz), 0)
        ]
        for q in randn(QuatVecF64, 10_000)
            ∂1 = ∂exp(q)
            ∂2 = ∂exp∂q(q)
            for i in 1:4
                @test ∂1[i] ≈ ∂2[i] rtol=ϵ
            end
            e, ∂3 = exp∂exp(q)
            @test e ≈ exp(q) rtol=4eps()
            for i in 1:4
                @test ∂3[i] ≈ ∂2[i] rtol=ϵ
            end
        end
        ∂1 = ∂exp(QuatVecF64(0))
        ∂2 = QuaternionF64[1, imx, imy, imz]
        for i in 1:4
            @test ∂1[i] ≈ ∂2[i] rtol=ϵ
        end
        e, ∂1 = exp∂exp(QuatVecF64(0))
        @test e ≈ QuaternionF64(1) rtol=4eps() atol=10eps()
        for i in 1:4
            @test ∂1[i] ≈ ∂2[i] rtol=ϵ
        end
    end

    @testset verbose=true "slerp" begin
        ∂slerp(q₁, q₂, τ) = (
            [
                ForwardDiff.derivative(ϵ->slerp(q₁+(ϵ+0imx), q₂, τ), 0),
                ForwardDiff.derivative(ϵ->slerp(q₁+ϵ*imx, q₂, τ), 0),
                ForwardDiff.derivative(ϵ->slerp(q₁+ϵ*imy, q₂, τ), 0),
                ForwardDiff.derivative(ϵ->slerp(q₁+ϵ*imz, q₂, τ), 0)
            ],
            [
                ForwardDiff.derivative(ϵ->slerp(q₁, q₂+(ϵ+0imx), τ), 0),
                ForwardDiff.derivative(ϵ->slerp(q₁, q₂+ϵ*imx, τ), 0),
                ForwardDiff.derivative(ϵ->slerp(q₁, q₂+ϵ*imy, τ), 0),
                ForwardDiff.derivative(ϵ->slerp(q₁, q₂+ϵ*imz, τ), 0)
            ]
        )
        for (q₁, q₂) in zip(randn(RotorF64, 1_000), randn(RotorF64, 1_000))
            for τ ∈ [0.0; 1.0; rand(5)...]
                ∂1, ∂2 = ∂slerp(q₁, q₂, τ)
                sa, ∂a, ∂b, ∂c = slerp∂slerp(q₁, q₂, τ)
                @test sa ≈ slerp(q₁, q₂, τ) rtol=2ϵ
                for i in 1:4
                    @test ∂1[i] ≈ ∂a[i] rtol=5ϵ atol=50eps()
                    @test ∂2[i] ≈ ∂b[i] rtol=5ϵ atol=50eps()
                end
                @test ∂c ≈ log(q₂/q₁) * slerp(q₁, q₂, τ) rtol=2ϵ
                sb, ∂d = slerp∂slerp∂τ(q₁, q₂, τ)
                @test sa ≈ sb rtol=2ϵ
                @test ∂c ≈ ∂d rtol=2ϵ
            end
        end
    end

    @testset verbose=true "squad" begin
        qs = Rotor{Float64}[1, imx, imy, imz, -imy, -imz, -imx, -Rotor(1)]
        ts = Float64.(1:length(qs))
        ∂squad(t) = ForwardDiff.derivative(τ->squad(qs, ts, τ), t)
        for i ∈ 1:length(ts)-1
            qᵢ, qᵢ₊₁ = qs[i], qs[i+1]
            ta, tb = ts[i], ts[i+1]
            for τ ∈ [0.001, 0.1, 0.49, 0.5, 0.51, 0.9, 0.999]
                t = ta + τ*(tb-ta)
                A, B = Quaternionic.squad_control_points(qs, ts, i)
                ∂1 = ∂squad(t)
                s, ∂2 = squad∂squad∂t(qᵢ, A, B, qᵢ₊₁, ta, tb, t)
                @test ∂1 ≈ ∂2 rtol=100ϵ atol=ϵ
                @test s ≈ squad(qs, ts, t) rtol=ϵ atol=ϵ
            end
        end
    end
end


text:
using YaoEGraph
using YaoHIR, YaoLocations
using YaoHIR: X, Y, Z, H
circ = Chain(Gate(X, Locations(1)), Gate(X, Locations(1)), Gate(Y, Locations(2)))
H2 = Chain(Gate(H, Locations(1)), Gate(H, Locations(1)))
simplify!(circ)
simplify!(H2)


text:
# The tests are based on the test suite of the `particle` Python package from
# the SciKit-HEP group: https://github.com/scikit-hep/particle with the Git
# commit 6fe8211181d6a0b21933b2d1042eeb4e87d77b3c from 2020-11-27
using Test
using InteractiveUtils
using Corpuscles
using Unitful

const DATA_DIR = joinpath(@__DIR__, "../data")

include("particles.jl")

@testset "inits" begin
    for T in vcat(map(subtypes, [Signed, Unsigned])...)
        p = Particle(T(11))
        @test PDGID(11) == p.pdgid
    end
end

@testset "pdg digits" begin
    p = Particle(9020113)
    @test p.pdgid.Nj == 3
    @test p.pdgid.Nq3 == 1
    @test p.pdgid.Nq2 == 1
    @test p.pdgid.Nq1 == 0
    @test p.pdgid.Nl == 2
    @test p.pdgid.Nr == 0
    @test p.pdgid.N == 9
    @test p.pdgid.N8 == 0
    @test p.pdgid.N9 == 0
    @test p.pdgid.N10 == 0

    # digits should be positive
    p = Particle(-321)
    @test p.pdgid.Nj == 1
    @test p.pdgid.Nq3 == 2
    @test p.pdgid.Nq2 == 3
end

@testset "conversions" begin

    # Particle Identites
    # PDGID <-> Geant3ID
    geant_id = Geant3ID(1)
    pdg_id = convert(PDGID, geant_id)
    @test isequal(22, pdg_id.value)
    geant_id = Geant3ID(3)
    pdg_id = convert(PDGID, geant_id)
    @test isequal(11, pdg_id.value)
    geant_id = Geant3ID(3)
    pdg_id = convert(PDGID, geant_id)
    @test !isequal(13, pdg_id.value)
    pdg_id = PDGID(22)
    geant_id = convert(Geant3ID, pdg_id)
    @test isequal(1, geant_id.value)
    pdg_id = PDGID(11)
    geant_id = convert(Geant3ID, pdg_id)
    @test isequal(3, geant_id.value)
    pdg_id = PDGID(13)
    geant_id = convert(Geant3ID, pdg_id)
    @test !isequal(1, geant_id.value)

    @test PDGID(13) === convert(PDGID, PDGID(13))

    try
        convert(PDGID, Geant3ID(5000))
        @test false
    catch e
        @test e isa Corpuscles.IDException
        msg = sprint(showerror, e)
        @test msg == "ParticleID Error: No corresponding PDGID for Geant3ID(5000) found!"
    end

    try
        convert(Geant3ID, PDGID(2222212))
        @test false
    catch e
        @test e isa Corpuscles.IDException
        msg = sprint(showerror, e)
        @test occursin("2222212", msg)
    end

    for particle in particles()
        if particle.pdgid in [PDGID(-9000321), PDGID(9000321), PDGID(9020113),
                              PDGID(-9000311), PDGID(9020213), PDGID(9000311),
                              PDGID(-9020213)]
            continue
        end
        @test convert(Geant3ID, particle.pdgid) !== nothing
    end

end

@testset "catalogs" begin
    catalog_files = map(basename, Corpuscles.available_catalog_files())
    @test "particle2008.csv" in catalog_files
    @test "particle2018.csv" in catalog_files
    @test "particle2019.csv" in catalog_files

    Corpuscles.use_catalog_file(joinpath(DATA_DIR, "particle2008.csv"))
    @test 0.054u"MeV/c^2" == Particle(553).width.value
    Corpuscles.use_catalog_file(joinpath(DATA_DIR, "particle2019.csv"))
    @test 0.05402u"MeV/c^2" == Particle(553).width.value

end

@testset "comparison" begin
    a = Corpuscles.MeasuredValue(1u"m", 10u"cm", 10u"cm")
    b = Corpuscles.MeasuredValue(1u"m", 10u"cm", 10u"cm")
    c = Corpuscles.MeasuredValue(1u"m", 5u"cm", 5u"cm")
    d = Corpuscles.MeasuredValue(0.8u"m", 5u"cm", 5u"cm")
    e = Corpuscles.MeasuredValue(0.9u"m", 5u"cm", 5u"cm")
    f = Corpuscles.MeasuredValue(90u"s", 5u"s", 5u"s")
    g = 91u"cm"
    h = 91u"s"

    @test a == b
    @test !(a == c)
    @test d < a
    @test !(e < a)
    @test !(e < c)
    @test g < c
    @test !(f < h)
    @test isapprox(g, e)
    @test isapprox(g, a)
    @test isapprox(a, b)
    @test isapprox(a, c)
    @test isapprox(a, e)
    @test !isapprox(a, d)

    try
        e < f
        @test false
    catch e
        @test e isa Unitful.DimensionError
    end
end

@testset "show and print" begin
    io = IOBuffer(write=true)
    show(io, Particle(1))
    seekstart(io)
    @test "Particle(1) d" == String(read(io))

    io = IOBuffer(write=true)
    print(io, Particle(1))
    seekstart(io)
    output = String(read(io))
    @test occursin(r"PDG ID:\s*1", output)
    @test occursin(r"Name:\s*d\n", output)
end

@testset "Particle(String) helper and -(Particle)" begin
    @test Particle("photon") == Particle(22)
    @test Particle("higgs") == Particle(25)
    @test Particle(-13) == -Particle(13)
    @test Particle("mu") == Particle(13)
    for p in ("mu", "tau")
        @test Particle("$p~") == -Particle(p)
        @test Particle("anti-$p") == -Particle(p)
        # alias
        @test Particle("$(p)-") == Particle(p)
        @test Particle("$(p)+") == -Particle(p)
        @test Particle("$(p)on") == Particle(p)
        @test Particle("anti-$(p)on") == -Particle(p)
    end

    @test Particle("D_sst_plus") == Particle(433)
    @test Particle("anti-D_sst_plus") == -Particle(433)
    @test_throws ErrorException Particle("SUSY is fake")

    # self anti-particle
    @test Particle("H") == -Particle("H")
end

@testset "helpers" begin
    Corpuscles.use_catalog_file(joinpath(DATA_DIR, "particle2020.csv"))

    """
    Runs the given is/has-function on a set of candidates and also checks every
    other noncandidate in PDGIDs for the opposite outcome.
    """
    function check_candidates(f, candidates)
        noncandidates = setdiff(Set(instances(PDGIDS)), Set(candidates))
        for candidate ∈ candidates
            @test f(candidate)
        end
        for noncandidate ∈ noncandidates
            @test !f(noncandidate)
        end
    end

    for _particles ∈ [particles(), [p.pdgid for p in particles()], [p.pdgid.value for p in particles()]]
        @test sum(map(isstandard, _particles)) > 0
        @test sum(map(isfundamental, _particles)) > 0
        @test sum(map(Corpuscles.isstandard, _particles)) == 610
        @test sum(map(Corpuscles.fundamentalid, _particles)) == 413

        @test sum(map(isquark, _particles)) == 12
        @test sum(map(islepton, _particles)) == 16
        @test sum(map(ismeson, _particles)) == 234
        @test sum(map(isbaryon, _particles)) == 292
        @test sum(map(ishadron, _particles)) == 526
        @test sum(map(isRhadron, _particles)) == 1
        @test sum(map(isSUSY, _particles)) == 0
        @test sum(map(ispentaquark, _particles)) == 0
        @test sum(map(isgaugebosonorhiggs, _particles)) == 6
        @test sum(map(issmgaugebosonorhiggs, _particles)) == 6
        @test sum(map(isdyon, _particles)) == 0
        @test sum(map(isnucleus, _particles)) == 4
        @test sum(map(isdiquark, _particles)) == 50
        @test sum(map(istechnicolor, _particles)) == 0
        @test sum(map(iscompositequarkorlepton, _particles)) == 0
        @test sum(map(isgeneratorspecific, _particles)) == 0
        @test sum(map(isspecial, _particles)) == 0
        @test sum(map(isQball, _particles)) == 0
        @test sum(map(hasfundamentalanti, _particles)) == 30

        @test sum(map(hasdown, _particles)) == 328
        @test sum(map(hasup, _particles)) == 346
        @test sum(map(hascharm, _particles)) == 107
        @test sum(map(hasstrange, _particles)) == 257
        @test sum(map(hasbottom, _particles)) == 68
        @test sum(map(hastop, _particles)) == 0
    end

    @testset "isquark" begin
        candidates = [DQuark, UQuark, SQuark, CQuark, BQuark, TQuark, BPrimeQuark, TPrimeQuark]
        check_candidates(isquark, candidates)
    end

    @testset "islepton" begin
        candidates = [Electron, Positron, Muon, AntiMuon, Tau, TauPrime, Nu_e, NuBar_tau, AntiElectronStar]
        check_candidates(islepton, candidates)
    end

    @testset "ismeson" begin
        candidates = [jpsi, psi_2S, Upsilon_1S, Upsilon_4S, Pi0, PiPlus, eta,
                  eta_prime, a_0_1450_plus, KL, KS, KMinus, phi, omega,
                  rho_770_minus, K1_1270_0, K1_1400_0, rho_1700_0, a2_1320_minus,
                  omega_3_1670, f_4_2300, D0, DPlus, DsPlus, B0, BPlus, Bs,
                  BcPlus, Pi0TC, PiMinusTC, T0, Reggeon, Pomeron, Odderon,
                  RPlus_TTildeDbar, R0_GTildeG]
        check_candidates(ismeson, candidates)
    end

    @testset "isbaryon" begin
        candidates = [Proton, AntiNeutron, HydrogenNucleus, Lambda, Sigma0,
                      SigmaPlus, SigmaMinus, Xi0, AntiXiMinus, OmegaMinus,
                      LcPlus, Lb, LtPlus, RPlusPlus_GTildeUUU,
                      UCbarCUDPentaquark, AntiUCbarCUDPentaquark]
        check_candidates(isbaryon, candidates)
    end

    @testset "ishadron" begin
        @test all(ishadron(p) == (ismeson(p) || isbaryon(p)) for p in instances(PDGIDS))
    end

    @testset "ispentaquark" begin
        candidates = [UCbarCUDPentaquark, AntiUCbarCUDPentaquark]
        check_candidates(ispentaquark, candidates)
    end

    @testset "isgaugebosonorhiggs" begin
        candidates = [Gluon, Photon, Z0, WMinus, HiggsBoson, ZPrime, Graviton]
        check_candidates(isgaugebosonorhiggs, candidates)
    end

    @testset "issmgaugebosonorhiggs" begin
        candidates = [Gluon, Photon, Z0, WMinus, HiggsBoson]
        check_candidates(issmgaugebosonorhiggs, candidates)
    end

    @testset "isgeneratorspecific" begin
        candidates = [AntiCHadron]
        check_candidates(isgeneratorspecific, candidates)
    end

    @testset "isspecial" begin
        candidates = [Graviton, Reggeon, Pomeron, Odderon, AntiCHadron]
        check_candidates(isspecial, candidates)
    end

    @testset "isnucleus" begin
        candidates = [Proton, AntiNeutron, HydrogenNucleus, Carbon12]
        check_candidates(isnucleus, candidates)
    end

    @testset "isdiquark" begin
        candidates = [DD1, SD0]
        check_candidates(isdiquark, candidates)
    end

    @testset "isRhadron" begin
        candidates = [RPlus_TTildeDbar, R0_GTildeG, RPlusPlus_GTildeUUU]
        check_candidates(isRhadron, candidates)
    end

    @testset "isQball" begin
        candidates = [QBall1, QBall2]
        check_candidates(isQball, candidates)
    end

    @testset "isdyon" begin
        candidates = [DyonSameMagElecChargeSign, DyonOppositeMagElecChargeSign]
        check_candidates(isdyon, candidates)
    end

    @testset "isSUSY" begin
        candidates = [Gluino, Gravitino, STildeL, CTildeR]
        check_candidates(isSUSY, candidates)
    end

    @testset "istechnicolor" begin
        candidates = [Pi0TC, PiMinusTC]
        check_candidates(istechnicolor, candidates)
    end

    @testset "iscompositequarkorlepton" begin
        candidates = [UQuarkStar, AntiElectronStar]
        check_candidates(iscompositequarkorlepton, candidates)
    end

    @testset "hasdown" begin
        f = hasdown
        @test f(Photon) == false
        @test f(Gluon) == false
        @test f(Electron) == false
        @test f(AntiMuon) == false
        @test f(jpsi) == false
        @test f(Upsilon_1S) == false
        @test f(PiPlus) == true
        @test f(KMinus) == false
        @test f(D0) == false
        @test f(DPlus) == true
        @test f(DsPlus) == false
        @test f(B0) == true
        @test f(Bs) == false
        @test f(BcPlus) == false
        @test f(Proton) == true
        @test f(LcPlus) == true
        @test f(Lb) == true
        @test f(DD1) == true
        @test f(SD0) == true
        @test f(Invalid1) == false
        @test f(Invalid2) == false
    end

    @testset "hasup" begin
        f = hasup
        @test f(Photon) == false
        @test f(Gluon) == false
        @test f(Electron) == false
        @test f(AntiMuon) == false
        @test f(jpsi) == false
        @test f(Upsilon_1S) == false
        @test f(PiPlus) == true
        @test f(KMinus) == true
        @test f(D0) == true
        @test f(DPlus) == false
        @test f(DsPlus) == false
        @test f(B0) == false
        @test f(Bs) == false
        @test f(BcPlus) == false
        @test f(Proton) == true
        @test f(LcPlus) == true
        @test f(Lb) == true
        @test f(DD1) == false
        @test f(SD0) == false
        @test f(Invalid1) == false
        @test f(Invalid2) == false
    end

    @testset "hasstrange" begin
        f = hasstrange
        @test f(Photon) == false
        @test f(Gluon) == false
        @test f(Electron) == false
        @test f(AntiMuon) == false
        @test f(jpsi) == false
        @test f(Upsilon_1S) == false
        @test f(PiPlus) == false
        @test f(KMinus) == true
        @test f(D0) == false
        @test f(DPlus) == false
        @test f(DsPlus) == true
        @test f(B0) == false
        @test f(Bs) == true
        @test f(BcPlus) == false
        @test f(Proton) == false
        @test f(LcPlus) == false
        @test f(Lb) == false
        @test f(DD1) == false
        @test f(SD0) == true
        @test f(Invalid1) == false
        @test f(Invalid2) == false
    end

    @testset "hascharm" begin
        candidates = [jpsi, psi_2S, D0, DPlus, DsPlus, BcPlus, LcPlus, UCbarCUDPentaquark, AntiUCbarCUDPentaquark]
        check_candidates(hascharm, candidates)
    end

    @testset "hasbottom" begin
        candidates = [Upsilon_1S, Upsilon_4S, B0, BPlus, Bs, BcPlus, Lb]
        check_candidates(hasbottom, candidates)
    end

    @testset "hastop" begin
        candidates = [T0, LtPlus]
        check_candidates(hastop, candidates)
    end

    @testset "hasfundamentalanti" begin
        candidates = [WMinus, Electron, Positron, Muon, AntiMuon, Tau, TauPrime,
                      Nu_e, NuBar_tau, DQuark, UQuark, SQuark, CQuark, BQuark,
                      TQuark, BPrimeQuark, TPrimeQuark, UQuarkStar,
                      AntiElectronStar, STildeL, CTildeR, AntiCHadron]
        check_candidates(hasfundamentalanti, candidates)
    end

    @testset "A" begin
        candidates = Dict(Proton => 1, AntiNeutron => 1, HydrogenNucleus => 1, Carbon12 => 12)
        for (candidate, value) ∈ candidates
            @test A(candidate) == value
        end
        noncandidates = setdiff(Set(keys(candidates)), Set(instances(PDGIDS)))
        for noncandidate in noncandidates
            @test isnothiung(A(candidate))
        end
    end

    @testset "Z" begin
        candidates = Dict(Proton => 1, AntiNeutron => 0, HydrogenNucleus => 1, Carbon12 => 6)
        for (candidate, value) ∈ candidates
            @test Z(candidate) == value
        end
        noncandidates = setdiff(Set(keys(candidates)), Set(instances(PDGIDS)))
        for noncandidate in noncandidates
            @test isnothiung(Z(candidate))
        end
    end

    @testset "isvalid" begin
        f = isvalid
        candidates = [Photon, Gluon, Electron, AntiMuon, jpsi, Upsilon_1S,
                      PiPlus, KMinus, D0, DPlus, DsPlus, B0, Bs, BcPlus, Proton,
                      LcPlus, Lb, DD1, SD0]
        noncandidates = [Invalid1, Invalid2]
        for candidate ∈ candidates
            @test f(Corpuscles.pdgid(candidate))
        end
        for noncandidate in noncandidates
            @test !f(Corpuscles.pdgid(noncandidate))
        end
    end

    @testset "charge" begin
        @test charge(Photon) == 0
        @test charge(Gluon) == 0
        @test charge(Electron) == -1
        @test charge(AntiMuon) == +1
        @test charge(jpsi) == 0
        @test charge(Upsilon_1S) == 0
        @test charge(PiPlus) == +1
        @test charge(KMinus) == -1
        @test charge(D0) == 0
        @test charge(DPlus) == +1
        @test charge(DsPlus) == +1
        @test charge(B0) == 0
        @test charge(Bs) == 0
        @test charge(BcPlus) == +1
        @test charge(Proton) == +1
        @test charge(LcPlus) == +1
        @test charge(Lb) == 0
        @test charge(DD1) == -2 / 3
        @test charge(SD0) == -2 / 3
        @test Corpuscles.isnothing(charge(Invalid1))
        @test Corpuscles.isnothing(charge(Invalid2))
    end

    @testset "threecharge" begin
        @test threecharge(Photon) == 0
        @test threecharge(Electron) == -3
        @test threecharge(jpsi) == 0
        @test threecharge(Upsilon_1S) == 0
        @test threecharge(KMinus) == -3
        @test threecharge(D0) == 0
        @test threecharge(Proton) == +3
        @test threecharge(LcPlus) == +3
        @test threecharge(Lb) == 0
        @test threecharge(DD1) == -2
        @test threecharge(SD0) == -2
        @test Corpuscles.isnothing(threecharge(Invalid1))
        @test Corpuscles.isnothing(threecharge(Invalid2))
    end

    JSLstatelist = Dict(
        "000" => [Pi0, PiPlus, eta, eta_prime, KL, KS, KMinus, D0, DPlus,
                  DsPlus, B0, BPlus, Bs, BcPlus, T0],
        "011" => [a_0_1450_plus],
        "101" => [K1_1270_0],
        "110" => [rho_770_minus],
        "111" => [K1_1400_0],
        "112" => [rho_1700_0],
        "211" => [a2_1320_minus],
        "312" => [omega_3_1670],
    )

    @testset "JSL" begin
        @testset "badly known meons" begin
            @test jspin(f_4_2300) == 9
            @test sspin(f_4_2300) === nothing
            @test sspin(f_4_2300) === nothing
        end

        for (states, particles) ∈ JSLstatelist
            for p in particles
                J_, S_, L_ = [parse(Int, s) for s ∈ states]
                @test J(p) == J_
                @test S(p) == S_
                @test L(p) == L_
            end
        end
    end

    @testset "J non-mesons" begin
        Jlist = Dict(
            1 => [Gluon, Photon, Z0, jpsi, psi_2S, Upsilon_1S, Upsilon_4S, K1_1270_0,],
            1//2 => [Electron, Positron, Muon, AntiMuon, Tau, Nu_e, NuBar_tau,
                     DQuark, UQuark, SQuark, CQuark, BQuark, TQuark, Proton,
                     AntiNeutron, Lambda, Sigma0, SigmaPlus, SigmaMinus, Xi0,
                     AntiXiMinus, LcPlus, Lb, LtPlus, STildeL, CTildeR],
            3//2 => [OmegaMinus],
            nothing => [Invalid1, Invalid2],
            nothing => [TauPrime, BPrimeQuark, TPrimeQuark]
        )
        for (J_, particles) ∈ Jlist
            for p in particles
                @test J(p) == J_
            end
        end
    end

    @testset "S non-mesons" begin
        Slist = Dict(nothing => [Gluon, Photon, Z0])
        for (S_, particles) ∈ Slist
            for p in particles
                @test S(p) == S_
            end
        end
    end

    @testset "L non-mesons" begin
        Llist = Dict(nothing => [Gluon, Photon, Z0])
        for (L_, particles) ∈ Llist
            for p in particles
                @test L(p) == L_
            end
        end
    end
end


text:
function setup_folder(p::Param)

    println("Checking folder structure...")
    setup = false

    fdata   = "./data"
    folder  = fdata * p.folder
    fname   = fdata * p.folder *  "$(p.fname).h5"

     # if no name specified error out
     if folder == ""
        println("ERROR: Output folder cannot be empty string.")
        return setup
    end

    #Create data folder
   if !isdir(fdata)
        println(" './data' folder did not exist, creating it...")
        mkdir(fdata)
    end

    #Create specified folder inside data folder if it does not exist
    if !isdir(   folder )
        println("Output folder did not exist, creating it...")
        mkdir(folder)
    end

    if isfile( fname )

        println("File " * fname * " already exists" )
        return setup

    end

    setup = true
    println("OK")

    return setup
end


text:
import ..UncertainValues: minimum, maximum

"""
    AbstractUncertainValueDataset

A dataset of uncertain values, with fixed indices. Concrete implementations
must as a minimum implement the following fields:

- **`values::UncertainDataset`**: The (uncertain) values of the dataset.
"""
abstract type AbstractUncertainValueDataset <: AbstractUncertainDataset end

##########################
# Indexing and iteration
#########################
Base.getindex(uvd::AbstractUncertainValueDataset, i) = uvd.values[i]
Base.length(uvd::AbstractUncertainValueDataset) = length(uvd.values)
Base.size(uvd::AbstractUncertainValueDataset) = length(uvd)
Base.firstindex(uvd::AbstractUncertainValueDataset) = 1
Base.lastindex(uvd::AbstractUncertainValueDataset) = length(uvd.values)

Base.eachindex(ud::AbstractUncertainValueDataset) = Base.OneTo(length(ud.values))
Base.iterate(ud::AbstractUncertainValueDataset, state = 1) = iterate(ud.values, state)

Base.minimum(udata::AbstractUncertainValueDataset) = minimum([minimum(uval) for uval in udata])
Base.maximum(udata::AbstractUncertainValueDataset) = maximum([maximum(uval) for uval in udata])

###################
# Pretty printing
###################
function summarise(ud::AbstractUncertainValueDataset)
    _type = typeof(ud)
    n_values = length(ud.values)
    summary = "$_type with $n_values values"
    return summary
end

Base.show(io::IO, ud::AbstractUncertainValueDataset) = print(io, summarise(ud))

###################
# Various useful functions
###################
"""
    distributions(ud::UncertainDataset)

Returns the distributions for all the uncertain values of the dataset.
"""
distributions(ud::AbstractUncertainValueDataset) = [ud[i].distribution for i = 1:length(ud)]

export
AbstractUncertainValueDataset,
distributions


text:
@testset "cat" begin
    Ws = [randn(3,rand(1:10)) for _ in 1:4]
    ψs = [randn(size(Ws[i], 2)) for i in 1:4]
    As = PreImputingMatrix.(Ws, ψs)
    for is in powerset(1:4), is in permutations(is)
        length(is) > 0 || continue
        @test hcat(As[is]...) == PreImputingMatrix(hcat(Ws[is]...), vcat(ψs[is]...))
        @inferred hcat(As[is]...)
        @test_throws ArgumentError vcat(As[is]...)
    end

    Ws = [randn(rand(1:10), 3) for _ in 1:4]
    ψs = [randn(size(Ws[i], 1)) for i in 1:4]
    As = PostImputingMatrix.(Ws, ψs)
    for is in powerset(1:4), is in permutations(is)
        length(is) > 0 || continue
        @test vcat(As[is]...) == PostImputingMatrix(vcat(Ws[is]...), vcat(ψs[is]...))
        @inferred vcat(As[is]...)
        @test_throws ArgumentError hcat(As[is]...)
    end
end

@testset "correct pre imputing behavior for standard vector (maybe missing)" begin
    function _test_imput(W, ob::Vector, b::Vector)
        A = PreImputingMatrix(W, ob)
        @test A * b == W * ob
        @test eltype(A * b) === promote_type(eltype(W), eltype(ob))
        @inferred A * b
    end

    function _test_imput(W, ob::Vector, t=10, p=0.2)
        for _ in 1:t
            b = [rand() < 0.2 ? missing : x for x in ob]
            idcs = rand([true, false], length(ob))
            _test_imput(W, ob, b)
        end
    end
    _test_imput(randn(3,3), randn(3))
    _test_imput(reshape(1:9 |> collect, (3,3)), rand(1:3, 3))
    _test_imput(randn(3,3), randn(3), fill(missing, 3))
    _test_imput(reshape(1:9 |> collect, (3,3)), rand(1:3, 3), fill(missing, 3))
end

@testset "Wrong dimensions" begin
    A = PreImputingMatrix(rand(2,3), rand(3))
    @test_throws DimensionMismatch A * []
    @test_throws DimensionMismatch A * [[]]
    @test_throws DimensionMismatch A * [1, 2]
    @test_throws DimensionMismatch A * [missing, missing]
    @test_throws DimensionMismatch A * [1, 2, 3, 4]
    @test_throws DimensionMismatch A * [1, 2, 3, missing]

    A = PostImputingMatrix(rand(2,3), rand(2))
    @test_throws DimensionMismatch A * MaybeHotVector(2, 2)
    @test_throws DimensionMismatch A * MaybeHotVector(missing, 4)
    @test_throws DimensionMismatch A * MaybeHotMatrix(Int[], 1)
    @test_throws DimensionMismatch A * MaybeHotMatrix([1, 2], 5)
    @test_throws DimensionMismatch A * MaybeHotMatrix([missing, 5, 8], 10)
    @test_throws DimensionMismatch A * MaybeHotMatrix([missing, missing], 2)
end

@testset "correct pre imputing behavior for standard matrix (maybe missing)" begin
    W = [1 2; 3 4]
    ψ = [2, 1]
    A = PreImputingMatrix(W, ψ)
    B1 = [4 3; 2 1]
    B2 = [missing missing; missing missing]
    B3 = [missing 3; 1 missing]
    B4 = [3 missing; missing 1]

    @test A * B1 == W * B1
    @test A * B2 == [W * ψ W * ψ]
    @test A * B3 == W * [2 3; 1 1]
    @test A * B4 == W * [3 2; 1 1]

    @inferred A * B1
    @inferred A * B2
    @inferred A * B3
    @inferred A * B4

    @test eltype(A * B1) === eltype(A * B2) === eltype(A * B3) === eltype(A * B4) === Int64
end

@testset "correct post imputing behavior for full standard arrays" begin
    W = [1 2; 3 4]
    ψ = [2, 1]
    A = PostImputingMatrix(W, ψ)
    B = [4 3; 2 1]
    b = [3, 1]
    @test A * B == W * B
    @test A * b == W * b
    @inferred A * B
    @inferred A * b
    @test eltype(A * B) === Int64
    @test eltype(A * b) === Int64

    B = [missing 3; missing 1]
    b = [3, missing]
    @test isequal(A * B, W * B)
    @test isequal(A * b, W * b)
    @inferred A * B
    @inferred A * b
    @test eltype(A * B) === Union{Int64, Missing}
    @test eltype(A * b) === Union{Int64, Missing}

    W = randn(2,3)
    ψ = randn(2)
    A = PostImputingMatrix(W, ψ)

    B = randn(3, 5)
    b = randn(3)
    @test A * B ≈ W * B
    @test A * b ≈ W * b
    @inferred A * B
    @inferred A * b
    @test eltype(A * B) === Float64
    @test eltype(A * b) === Float64

    B = [missing 3.0; missing 1.0; 2.0 5.0]
    b = [3.0, missing, 1.0]
    @test isequal(A * B, W * B)
    @test isequal(A * b, W * b)
    @inferred A * B
    @inferred A * b
    @test eltype(A * B) === Union{Float64, Missing}
    @test eltype(A * b) === Union{Float64, Missing}

    B = fill(missing, 3, 5)
    b = fill(missing, 3)
    @test isequal(A * B, W * B)
    @test isequal(A * b, W * b)
    @inferred A * B
    @inferred A * b
    @test eltype(A * B) === Missing
    @test eltype(A * b) === Missing
end

@testset "correct post imputing behavior for maybe hot vector" begin
    for (m, n) in product(fill((2, 5, 10), 2)...)
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)
        i = rand(1:n)
        b = MaybeHotVector(i, n)
        @test A * b == W * onehot(b)
        @inferred A * b
        @test eltype(A * b) === eltype(W)
        b = MaybeHotVector(missing, n)
        @test A * b == ψ
        @inferred A * b
        @test eltype(A * b) === eltype(ψ)
    end
end

@testset "correct post imputing behavior for maybe hot matrix" begin
    for (m, n, k) in product(fill((2, 5, 10), 3)...)
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)
        i1 = rand(1:n, k)
        i2 = fill(missing, k)
        i3 = [isodd(i) ? missing : rand(1:n) for i in 1:k]
        B1 = MaybeHotMatrix(i1, n)
        B2 = MaybeHotMatrix(i2, n)
        B3 = MaybeHotMatrix(i3, n)
        @test A * B1 == W * onehotbatch(B1)
        @inferred A * B1
        @test eltype(A * B1) === eltype(W)
        @test all(isequal(ψ), eachcol(A * B2))
        @inferred A * B2
        @test eltype(A * B2) === eltype(ψ)
        C = A * B3
        @inferred A * B3
        @test eltype(A * B3) === promote_type(eltype(W), eltype(ψ))
        @test all(isequal(ψ), eachcol(C[:, ismissing.(i3)]))
        @test C[:, .!ismissing.(i3)] == W * onehotbatch(skipmissing(i3), 1:n)
    end
end

@testset "correct post imputing behavior for ngram matrix" begin
    for (m, n, k) in product(fill((2, 5, 10), 3)...)
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)

        S = [randstring(rand(1:100)) for _ in 1:k] |> PooledArray
        B = NGramMatrix(S, 3, 256, n)
        @test A * B ≈ W * Matrix(SparseMatrixCSC(B))
        @inferred A * B
        @test eltype(A * B) === eltype(W)

        S = fill(missing, k) |> PooledArray
        B = NGramMatrix(S, 3, 256, n)
        @test all(isequal(ψ), eachcol(A * B))
        @inferred A * B
        @test eltype(A * B) === eltype(ψ)

        if k > 1
            S = [isodd(i) ? missing : randstring(rand(1:10)) for i in 1:k]
            B = NGramMatrix(S, 3, 256, n)
            C = A * B
            @test all(isequal(ψ), eachcol(C[:, ismissing.(S)]))
            @test C[:, .!ismissing.(S)] ≈ W * NGramMatrix(skipmissing(S) |> collect, 3, 256, n)
            @inferred A * B
            @test eltype(A * B) === promote_type(eltype(W), eltype(ψ))
        end
    end
end

@testset "imputing matrix * full vector gradient testing" begin
    for (m, n) in product(fill((2, 5, 10), 2)...)
        b = randn(n)
        W = randn(m, n)

        ψ = randn(n)
        A = PreImputingMatrix(W, ψ)

        @test gradtest((W, ψ, b) -> PreImputingMatrix(W, ψ) * b, W, ψ, b)

        (dW, dψ), db = gradient(sum ∘ *, A, b)
        @test dW ≈ gradient(W -> sum(W * b), W) |> only
        @test isnothing(dψ)
        @test db ≈ gradient(b -> sum(W * b), b) |> only

        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)

        @test gradtest((W, ψ, b) -> PostImputingMatrix(W, ψ) * b, W, ψ, b)

        (dW, dψ), db = gradient(sum ∘ *, A, b)
        @test dW ≈ gradient(W -> sum(W * b), W) |> only
        @test isnothing(dψ)
        @test db ≈ gradient(b -> sum(W * b), b) |> only
    end
end

@testset "imputing matrix * full matrix gradient testing" begin
    for (m, n, k) in product(fill((2, 5, 10), 3)...)
        B = randn(n, k)
        W = randn(m, n)

        ψ = randn(n)
        A = PreImputingMatrix(W, ψ)

        @test gradtest((W, ψ, B) -> PreImputingMatrix(W, ψ) * B, W, ψ, B)

        (dW, dψ), dB = gradient(sum ∘ *, A, B)
        @test dW ≈ gradient(W -> sum(W * B), W) |> only
        @test isnothing(dψ)
        @test dB ≈ gradient(B -> sum(W * B), B) |> only

        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)

        @test gradtest((W, ψ, B) -> PostImputingMatrix(W, ψ) * B, W, ψ, B)

        (dW, dψ), dB = gradient(sum ∘ *, A, B)
        @test dW ≈ gradient(W -> sum(W * B), W) |> only
        @test isnothing(dψ)
        @test dB ≈ gradient(B -> sum(W * B), B) |> only
    end
end

@testset "post imputing matrix * full maybe hot vector gradient testing" begin
    for (m, n) in product(fill((2, 5, 10), 2)...), i in [rand(1:n) for _ in 1:3]
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)
        b = MaybeHotVector(i, n)

        @test gradtest((W, ψ) -> PostImputingMatrix(W, ψ) * b, W, ψ)

        (dW, dψ), db = gradient(sum ∘ *, A, b)
        @test dW ≈ gradient(W -> sum(W * b), W) |> only
        @test isnothing(dψ)
        @test isnothing(db)
    end
end

@testset "post imputing matrix * full maybe hot matrix gradient testing" begin
    for (m, n, k) in product(fill((2, 5, 10), 3)...), I in [rand(1:n, k) for _ in 1:3]
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)
        B = MaybeHotMatrix(I, n)

        @test gradtest((W, ψ) -> PostImputingMatrix(W, ψ) * B, W, ψ)

        (dW, dψ), dB = gradient(sum ∘ *, A, B)
        @test dW ≈ gradient(W -> sum(W * B), W) |> only
        @test isnothing(dψ)
        @test isnothing(dB)
    end
end

@testset "post imputing matrix * full ngram matrix gradient testing" begin
    for (m, n, k) in product(fill((2, 5, 10), 3)...)
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)
        S = [randstring(rand(1:100)) for _ in 1:k] |> PooledArray
        B = NGramMatrix(S, 3, 256, n)

        @test gradtest((W, ψ) -> PostImputingMatrix(W, ψ) * B, W, ψ)

        (dW, dψ), dB = gradient(sum ∘ *, A, B)
        @test dW ≈ gradient(W -> sum(W * B), W) |> only
        @test isnothing(dψ)
        @test isnothing(dB)
    end
end

@testset "post imputing matrix * empty maybe hot vector gradient testing" begin
    for (m, n) in product(fill((2, 5, 10), 2)...)
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)
        b = MaybeHotVector(missing, n)

        @test gradtest((W, ψ) -> PostImputingMatrix(W, ψ) * b, W, ψ)

        (dW, dψ), db = gradient(sum ∘ *, A, b)
        @test isnothing(dW)
        @test all(dψ .== 1)
        @test isnothing(db)
    end
end

@testset "post imputing matrix * empty maybe hot matrix gradient testing" begin
    for (m, n, k) in product(fill((2, 5, 10), 3)...)
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)
        B = MaybeHotMatrix(fill(missing, k), n)

        @test gradtest((W, ψ) -> PostImputingMatrix(W, ψ) * B, W, ψ)

        (dW, dψ), dB = gradient(sum ∘ *, A, B)
        @test isnothing(dW)
        @test all(dψ .== k)
        @test isnothing(dB)
    end
end

@testset "post imputing matrix * empty ngram matrix gradient testing" begin
    for (m, n, k) in product(fill((2, 5, 10), 3)...)
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)
        S = fill(missing, k)
        B = NGramMatrix(S, 3, 256, n)

        @test gradtest((W, ψ) -> PostImputingMatrix(W, ψ) * B, W, ψ)

        (dW, dψ), dB = gradient(sum ∘ *, A, B)
        @test isnothing(dW)
        @test all(dψ .== k)
        @test isnothing(dB)
    end
end

@testset "post imputing matrix * mixed maybe hot matrix gradient testing" begin
    for (m, n, k) in product(fill((2, 5, 10), 3)...)
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)
        I = [isodd(i) ? missing : rand(1:n) for i in 1:k]
        B = MaybeHotMatrix(I, n)

        @test gradtest((W, ψ) -> PostImputingMatrix(W, ψ) * B, W, ψ)

        (dW, dψ), dB = gradient(sum ∘ *, A, B)
        Bskip = MaybeHotMatrix(skipmissing(I) |> collect, n)
        @test dW ≈ gradient(W -> sum(PostImputingMatrix(W, ψ) * Bskip), W) |> only
        Bmiss = MaybeHotMatrix(filter(ismissing, I) |> collect, n)
        @test dψ ≈ gradient(ψ -> sum(PostImputingMatrix(W, ψ) * Bmiss), ψ) |> only
        @test all(dψ .== count(ismissing, I))
        @test isnothing(dB)
    end
end

@testset "post imputing matrix * mixed ngram matrix gradient testing" begin
    for (m, n, k) in product(fill((2, 5, 10), 3)...)
        W = randn(m, n)
        ψ = randn(m)
        A = PostImputingMatrix(W, ψ)
        S = [isodd(i) ? missing : randstring(rand(1:100)) for i in 1:k] |> PooledArray
        B = NGramMatrix(S, 3, 256, n)

        @test gradtest((W, ψ) -> PostImputingMatrix(W, ψ) * B, W, ψ)

        (dW, dψ), dB = gradient(sum ∘ *, A, B)
        Bskip = NGramMatrix(skipmissing(S) |> collect, 3, 256, n)
        @test dW ≈ gradient(W -> sum(PostImputingMatrix(W, ψ) * Bskip), W) |> only
        Bmiss = NGramMatrix(filter(ismissing, S) |> collect, 3, 256, n)
        @test dψ ≈ gradient(ψ -> sum(PostImputingMatrix(W, ψ) * Bmiss), ψ) |> only
        @test all(dψ .== count(ismissing, S))
        @test isnothing(dB)
    end
end

@testset "pre imputing matrix * empty vector gradient testing" begin
    for (m, n) in product(fill((1, 5, 10), 2)...)
        W = randn(m, n)
        ψ = randn(n)
        A = PreImputingMatrix(W, ψ)
        b = fill(missing, n)

        @test gradtest((W, ψ) -> PreImputingMatrix(W, ψ) * b, W, ψ)

        (dW, dψ), db = gradient(sum ∘ *, A, b)
        @test dW ≈ gradient(W -> sum(W * ψ), W) |> only
        @test dψ ≈ gradient(ψ -> sum(W * ψ), ψ) |> only
        @test isnothing(db)
    end
end

@testset "pre imputing matrix * empty missing matrix gradient testing" begin
    for (m, n, k) in product(fill((1, 5, 10), 3)...)
        W = randn(m, n)
        ψ = randn(n)
        A = PreImputingMatrix(W, ψ)
        B = fill(missing, n, k)

        @test gradtest((W, ψ) -> PreImputingMatrix(W, ψ) * B, W, ψ)

        (dW, dψ), dB = gradient(sum ∘ *, A, B)
        @test dW ≈ gradient(W -> sum(W * repeat(ψ, 1, k)), W) |> only
        @test dψ ≈ gradient(ψ -> sum(W * repeat(ψ, 1, k)), ψ) |> only
        @test isnothing(dB)
    end
end

@testset "pre imputing matrix * mixed vector gradient testing" begin
    for (m, n) in product(fill((1, 5, 10), 2)...)
        W = randn(m, n)
        ψ = randn(n)
        b = Vector{Union{Float64, Missing}}(randn(n))
        b[rand(eachindex(b), rand(1:n))] .= missing

        @test gradtest((W, ψ) -> PreImputingMatrix(W, ψ) * b, W, ψ)
    end
end

@testset "pre imputing matrix * mixed matrix gradient testing" begin
    for (m, n, k) in product(fill((1, 5, 10), 3)...)
        W = randn(m, n)
        ψ = randn(n)
        B = Matrix{Union{Float64, Missing}}(randn(n, k))
        B[rand(eachindex(B), rand(1:n*k))] .= missing

        @test gradtest((W, ψ) -> PreImputingMatrix(W, ψ) * B, W, ψ)
    end
end

@testset "Parameters get indeed updated with Flux/Zygote" begin
    function check(d, X, pW, pψ)
        for opt in [Flux.Descent(), Flux.ADAM()]
            for it in 1:3
                dc = deepcopy(d)
                Flux.train!(X -> sum(dc(X)), Flux.params(dc), [(X,)], opt)
                @test any(dc.weight.W .!= d.weight.W) == pW
                @test any(dc.weight.ψ .!= d.weight.ψ) == pψ
            end
        end
    end

    for (m, n, k) in product(fill((1, 5, 10), 3)...)
        # https://github.com/FluxML/Flux.jl/issues/1479
        if m == 1 || n == 1 || k == 1
            continue
        end

        d = preimputing_dense(n, m)

        # https://github.com/FluxML/Flux.jl/issues/1479
        # check(d, randn(n), true, false)
        # x = Vector{Union{Float64, Missing}}(randn(n))
        # x[rand(eachindex(x), rand(1:n))] .= missing
        # check(d, x, true, true)
        # check(d, fill(missing, n), false, true)

        check(d, randn(n, k), true, false)
        X = Matrix{Union{Float64, Missing}}(randn(n, k))
        X[rand(eachindex(X), rand(1:n*k))] .= missing
        check(d, X, true, true)
        check(d, fill(missing, n, k), false, true)

        d = postimputing_dense(n, m)

        S = [randstring(rand(1:100)) for _ in 1:k]
        X = NGramMatrix(S, 3, 256, n)
        check(d, X, true, false)
        if k > 1
            S = [isodd(i) ? missing : randstring(rand(1:10)) for i in 1:k] |> PooledArray
            X = NGramMatrix(S, 3, 256, n)
            check(d, X, true, true)
        end
        S = fill(missing, k)
        X = NGramMatrix(S, 3, 256, n)
        check(d, X, false, true)

        # https://github.com/FluxML/Flux.jl/issues/1479
        # x = maybehot(rand(1:n), 1:n)
        # check(d, x, true, false)
        # x = maybehot(missing, 1:n)
        # check(d, x, false, true)

        S = [rand(1:n) for _ in 1:k]
        X = maybehotbatch(S, 1:n)
        check(d, X, true, false)
        if k > 1
            S = [isodd(i) ? missing : rand(1:n) for i in 1:k]
            X = maybehotbatch(S, 1:n)
            check(d, X, true, true)
        end
        S = fill(missing, k)
        X = maybehotbatch(S, 1:n)
        check(d, X, false, true)
    end
end

@testset "broadcasting" begin
    function check(A::T, W, ψ) where T
        @test A .+ A == T(W .+ W, ψ .+ ψ)
        @test A .- A == T(W .- W, ψ .- ψ)
        @test A .- 1 == T(W .- 1, ψ .- 1) == -1 .+ A
        @test 2 .* A == T(2 .* W, 2 .* ψ) == A .* 2
        @test A .+ A .+ A == 3 .* A == A .* 3
        A .+= A
        @test A == T(W .+ W, ψ .+ ψ)
        A .-= 2 .* A
        @test A == T(-2 .* W, -2 .* ψ)
        A .*= -1
        @test A == T(2 .* W, 2 .* ψ)
    end

    W = randn(3, 3)
    ψ = randn(3)
    A1 = PreImputingMatrix(W, ψ)
    check(A1, copy(W), copy(ψ))
    A2 = PostImputingMatrix(W, ψ)
    check(A2, copy(W), copy(ψ))

    @test_throws MethodError A1 .+ W
    @test_throws MethodError W .+ A1
    @test_throws MethodError A2 .+ W
    @test_throws MethodError W .+ A2
    @test A1 .+ A2 isa Matrix
    @test A2 .* A1 isa Matrix
end

@testset "imputing Dense construction" begin
    A = preimputing_dense(2, 3)
    @test size(A.weight) == (3, 2)
    @test A.weight isa PreImputingMatrix

    A = postimputing_dense(2, 3)
    @test size(A.weight) == (3, 2)
    @test A.weight isa PostImputingMatrix

    d = Dense(4, 5)
    @test preimputing_dense(d).weight |> size == d.weight |> size
    @test postimputing_dense(d).weight |> size == d.weight |> size
end


text:
######################
### Parser Methods ###
######################

###################################################
## good but expensive test for create_or_update  ##
## (insert after get_edge and get_cons)          ##
## @assert id(edge) in keys(agenda) string(edge) ##
## @assert id(cons) in keys(agenda) string(cons) ##
###################################################

function create_or_update!(key::EdgeKey, trav, agenda, logbook, grammar)
    if isdiscovered(logbook, key)
        edge = get_edge(logbook, key)
        add!(edge, trav)
    else
        edge = Edge(key, trav, 0, grammar) # 0 is just a placeholder for the real id
        discover!(logbook, edge)           # that is corrected by `discover!`
    end
    enqueue!(agenda, edge, false) #::Void
    nothing
end

function create_or_update!(key::ConsKey, comp, agenda, logbook, grammar)
    if isdiscovered(logbook, key)
        cons = get_cons(logbook, key)
        add!(cons, comp)
    else
        cons = Constituent(key, comp, 0, grammar) # 0 is just a placeholder for the real id
        discover!(logbook, cons)                  # that is corrected by `discover!`
    end
    enqueue!(agenda, cons, false)
    nothing
end

function initialize(input, grammar, parsing_method, epsilon)
    n = length(input)

    C = category_type(grammar)
    T = terminal_type(grammar)
    CR = category_rule_type(grammar)
    TR = terminal_rule_type(grammar)
    St = state_type(grammar)
    S = score_type(grammar)

    chart = Chart([ChartCell(C, St) for i in 1:n+1])
    agenda = Agenda(parsing_method)
    logbook = ParserLogbook(C, T, CR, TR, St, S)
    for i in 1:n
        for (cat, rule, sco) in completions(grammar, input[i])
            create_or_update!(ConsKey(i,i+1,cat), TerminalCompletion{T,TR,S}(input[i], rule, sco), agenda, logbook, grammar)
        end
        if !ismissing(epsilon)
            for (cat, rule, sco) in completions(grammar, epsilon)
                create_or_update!(ConsKey(i,i,cat), TerminalCompletion(epsilon, rule, sco), agenda, logbook, grammar)
            end
        end
    end
    if !ismissing(epsilon)
        for (cat, rule, sco) in completions(grammar, epsilon)
            create_or_update!(ConsKey(n+1,n+1,cat), TerminalCompletion(epsilon, rule, sco), agenda, logbook, grammar)
        end
    end
    chart, agenda, logbook
end

function do_fundamental_rule!(edge::Edge, chart, agenda, logbook, grammar)
    for cat in keys(consids(chart, edge))
        if is_possible_transition(grammar, state(edge), cat)
            for consid in consids(chart, edge)[cat]
                cons = get_cons(logbook, consid)
                new_state = transition(grammar, state(edge), cat)
                new_key = EdgeKey(start(edge), tail(cons), new_state)
                traversal = Traversal(id(edge), consid, score(edge) * score(cons))
                create_or_update!(new_key, traversal, agenda, logbook, grammar)
            end
        end
    end
end

function do_fundamental_rule!(cons::Constituent, chart, agenda, logbook, grammar)
    for state in keys(edgeids(chart, cons))
        if is_possible_transition(grammar, state, cat(cons))
            for edgeid in edgeids(chart, cons)[state]
                edge = get_edge(logbook, edgeid)
                new_state = transition(grammar, state, cat(cons))
                new_key = EdgeKey(start(edge), tail(cons), new_state)
                traversal = Traversal(edgeid, id(cons), score(edge) * score(cons))
                create_or_update!(new_key, traversal, agenda, logbook, grammar)
            end
        end
    end
end

function introduce_edge!(cons::Constituent, agenda, logbook, grammar)
    if is_possible_transition(grammar, startstate(grammar), cat(cons))
        state = transition(grammar, startstate(grammar), cat(cons))
        key = EdgeKey(start(cons), tail(cons), state)
        traversal = Traversal(0, id(cons), score(cons))
        create_or_update!(key, traversal, agenda, logbook, grammar)
    end
    nothing
end

function complete_edge!(edge::Edge, agenda, logbook, grammar)
    for (cat, rule, sco) in completions(grammar, state(edge))
        key = ConsKey(start(edge), tail(edge), cat)
        CR , S = category_rule_type(grammar), score_type(grammar)
        comp = EdgeCompletion{CR, S}(id(edge), rule, score(edge) * sco)
        create_or_update!(key, comp, agenda, logbook, grammar)
    end
    nothing
end

Base.isapprox(x::LogProb, y::LogProb) = abs(x.log-y.log) < LOGSPACE_THRESHOLD

function finish!(edge::Edge, chart, agenda, logbook, grammar)
    edge.insidepopnumber += 1
    s = score(edge)
    if s ≈ lastpopscore(edge) || insidepopnumber(edge) == MAX_POP_NUMBER
        if !haskey(chart.cells[tail(edge)].edgeids, state(edge)) || !(id(edge) in chart.cells[tail(edge)].edgeids[state(edge)])
            if !isfinal(state(edge))
                insert!(chart, edge)
            end
            edge.isfinished = true
            edge.score = s
            do_fundamental_rule!(edge, chart, agenda, logbook, grammar)
        end
    else
        complete_edge!(edge, agenda, logbook, grammar)
        edge.lastpopscore = score(edge)
        enqueue!(agenda, edge, true)
    end
    nothing
end

function finish!(cons::Constituent, chart, agenda, logbook, grammar)
    cons.insidepopnumber += 1
    s = score(cons)
    if s ≈ lastpopscore(cons) || insidepopnumber(cons) == MAX_POP_NUMBER
        if !haskey(chart.cells[start(cons)].consids, cat(cons)) || !(id(cons) in chart.cells[start(cons)].consids[cat(cons)]) string(start(cons), tail(cons), cat(cons))
            insert!(chart, cons)
            cons.isfinished = true
            cons.score = s
            do_fundamental_rule!(cons, chart, agenda, logbook, grammar)
        end
    else
        introduce_edge!(cons, agenda, logbook, grammar)
        cons.lastpopscore = score(cons)
        enqueue!(agenda, cons, true)
    end
    nothing
end

function run_chartparser(treestring::AbstractString, grammar, parsing_method=:full, epsilon=missing)
    run_chartparser(tree(treestring), grammar, parsing_method, epsilon=epsilon)
end

run_chartparser(tree::Tree, grammar, parsing_method=:full; epsilon=missing) =
    run_chartparser(leaf_data(tree), grammar, boolean_dependency_matrix(tree), parsing_method, epsilon=epsilon)
    # run_chartparser(leaf_data(tree), grammar, trues(length(leaf_data(tree))+1, length(leaf_data(tree))+1))

run_chartparser(input::Vector, grammar, parsing_method=:full; epsilon=missing) =
    run_chartparser(input, grammar, trues(length(input)+1, length(input)+1), parsing_method, epsilon=epsilon)

function run_chartparser(input::Vector, grammar, dependency_matrix::AbstractMatrix{Bool}, parsing_method=:full; epsilon=missing)
    # input = map(terminal_type(grammar), input) :: Vector{terminal_type(grammar)}
    chart, agenda, logbook = initialize(input, grammar, parsing_method, epsilon)
    while !isempty(agenda)
        # finish = do inference and insert or accumulate
        id = dequeue!(agenda)
        if id > 0 # then id is an edge id
            finish!(get_edge(logbook, id), chart, agenda, logbook, grammar)#::Void
        else
            cons = get_cons(logbook, id)
            # println(dependency_matrix[start(cons), tail(cons)])
            if dependency_matrix[start(cons), tail(cons)]
                # println(cons)
                finish!(cons, chart, agenda, logbook, grammar)#::Void

                if parsing_method == :viterbi && length(cons) == length(input) && cat(cons) in startsymbols(grammar)
                    break
                end
            end
        end
    end
    ParseForest(chart, logbook, input, grammar)
end

function run_chartparser(input::Vector, grammar, dependency_dict::AbstractDict{Tuple{C, Int, Int},Bool}, parsing_method=:full; epsilon=missing) where C
    # input = map(terminal_type(grammar), input) :: Vector{terminal_type(grammar)}
    chart, agenda, logbook = initialize(input, grammar, parsing_method, epsilon)
    while !isempty(agenda)
        # finish = do inference and insert or accumulate
        id = dequeue!(agenda)
        if id > 0 # then id is an edge id
            finish!(get_edge(logbook, id), chart, agenda, logbook, grammar)#::Void
        else
            cons = get_cons(logbook, id)
            # println(dependency_matrix[start(cons), tail(cons)])
            if get(dependency_dict, (cat(cons), start(cons), tail(cons)), false)
                # println(cons)
                finish!(cons, chart, agenda, logbook, grammar)#::Void

                if parsing_method == :viterbi && length(cons) == length(input) && cat(cons) in startsymbols(grammar)
                    break
                end
            end
        end
    end
    ParseForest(chart, logbook, input, grammar)
end


text:
typedict(x::T) where {T} = Dict(fn=>getfield(x, fn) for fn ∈ fieldnames(T))

h5read(filename, keys...) = read(h5open(filename, "r"), keys...)

function qr_data(d::Array{Float32,2}, ps::Integer; seed=nothing, ts=0)
    !isnothing(seed) && Random.seed!(seed)
    S = rand([-1f0, 1f0], size(d, 1), ps)
    ts > 0 && (d = circshift(d, (-ts, 0)))
    AS = d * (d' * S)
    Q, _ = qr(AS)
    return Matrix(Q)
end

simil(x, y) = dot(x, y)/(norm(x)*norm(y))

_wsave(s, fig::Figure; dpi::Int=150) = fig.savefig(s, bbox_inches="tight", dpi=dpi)

function normalize!(d::judiVector)
    for i=1:d.nsrc
        d.data[i] ./= mapslices(norm, d.data[i], dims=1)
    end
end

normalize(d::Vector) = d / norm(d)

function loss(d0::judiVector, d1::judiVector)
    loss = 0
    for i=1:d0.nsrc
        for r=1:size(d0.data[i], 2)
            loss += dot(d0.data[i][:, r], normalize(d0.data[i][:, r]) - normalize(d1.data[i][:, r]))
        end
    end
    loss
end

max_stars_repo_name: slimgroup/TimeProbeSeismic

max_issues_repo_name: slimgroup/TimeProbeSeismic

max_forks_repo_name: slimgroup/TimeProbeSeismic


text:
## Generic matrix operation
abstract AbstractMatOp{T}

## Dense matrix multiplication
immutable DenseMatProd{T} <: AbstractMatOp{T}
    mat::Matrix{T}
end
## The operation
performop{T}(A::DenseMatProd{T}, x::Vector{T}) = A.mat * x
## Size of the matrix
nrow{T}(A::DenseMatProd{T}) = size(A.mat, 1)

## Arnoldi factorization starting from step-k
function arnoldifac!{T}(k::Int, m::Int, V::Matrix{T}, H::Matrix{T}, f::Vector{T}, A::AbstractMatOp{T}, prec::T)
    if m <= k
        error("m must be greater than k")
    end

    ## Keep the upperleft k x k submatrix of H and
    ## set other elements to 0
    H[:, (k + 1):end] = zero(T)
    H[(k + 1):end, 1:k] = zero(T)

    beta::T = norm(f)

    for i = k:(m - 1)
        ## V{i+1} <- f / ||f||
        V[:, i + 1] = f / beta
        H[i + 1, i] = beta

        ## w <- A * V{i+1}
        w::Vector{T} = performop(A, V[:, i + 1])

        H[i, i + 1] = beta
        H[i + 1, i + 1] = dot(V[:, i + 1], w)

        ## f <- w - V * V' * w
        f[:] = w - beta * V[:, i] - H[i + 1, i + 1] * V[:, i + 1]
        beta = norm(f)

        ## f/||f|| is going to be the next column of V, so we need to test
        ## whether V' * (f/||f||) ~= 0
        Vf::Vector{T} = V[:, 1:(i + 1)]' * f
        count = 0
        while count < 5 && maximum(abs(Vf)) > prec * beta
            ## f <- f - V * Vf
            f[:] -= V[:, 1:(i + 1)] * Vf
            ## h <- h + Vf
            H[i, i + 1] += Vf[i]
            H[i + 1, i] = H[i, i + 1]
            H[i + 1, i + 1] += Vf[i + 1]
            ## beta <- ||f||
            beta = norm(f)

            Vf[:] = V[:, 1:(i + 1)]' * f
        end
    end
end

## Apply shifts on V, H and f
function applyshifts!{T}(k::Int, V::Matrix{T}, H::Matrix{T}, f::Vector{T}, shifts::Vector{T})
    n = size(V, 1)
    ncv = size(V, 2)
    Q::Matrix{T} = eye(T, ncv)

    for i = (k + 1):ncv
        ## QR decomposition of H-mu*I, mu is the shift
        for j = 1:ncv
            H[j, j] -= shifts[i]
        end
        qr = tridiagqr!(H)  ## H -> RQ
        applyright!(qr, Q)  ## Q -> Q * Qi
        for j = 1:ncv
            H[j, j] += shifts[i]
        end
    end

    ## V -> VQ, only need to update the first k+1 columns
    ## Q has some elements being zero
    ## The first (ncv - k + i) elements of the i-th column of Q are non-zero
    # Vs = zeros(T, n, k + 1)
    # for i = 1:k
    #     nnz = ncv - k + i
    #     Vs[:, i] = V[:, 1:nnz] * Q[1:nnz, i]
    # end
    # Vs[:, k + 1] = V * Q[:, k + 1]
    # V[:, 1:(k + 1)] = Vs
    ## However this seems to be slower than a direct multiplication

    ## V -> VQ
    V[:, :] = V * Q
    f[:] = f * Q[ncv, k] + V[:, k + 1] * H[k + 1, k]
end

## Retrieve Ritz values and Ritz vectors
function ritzpairs!{T}(which::Symbol, H::Matrix{T}, ritzval::Vector{T}, ritzvec::Matrix{T}, ritzest::Vector{T})
    ncv = size(ritzvec, 1)
    nev = size(ritzvec, 2)
    ## Eigen decomposition on H, which is symmetric and tridiagonal
    decomp = eigfact(SymTridiagonal(diag(H), diag(H, -1)))

    ## Sort Ritz values according to "which"
    if which == :LM
        trans = abs
        rev = true
    elseif which == :SM
        trans = abs
        rev = false
    elseif which == :LA || which == :BE
        trans = (x -> x)
        rev = true
    else which == :SA
        trans = (x -> x)
        rev = false
    end

    ix = sortperm(decomp.values, by = trans, rev = rev)

    if which == :BE
        ixcp = copy(ix)
        for i = 1:ncv
            if i % 2 == 1
                ix[i] = ixcp[(i + 1) / 2]
            else
                ix[i] = ixcp[ncv - i / 2 + 1]
            end
        end
    end

    ritzval[:] = decomp.values[ix]
    ritzest[:] = decomp.vectors[ncv, ix]
    ritzvec[:, :] = decomp.vectors[:, ix[1:nev]]
end

## Adjusted nev
function nevadjusted{T}(nev::Int, ncv::Int, nconv::Int, ritzest::Vector{T}, prec::T)
    nevnew = nev + sum(abs(ritzest[(nev + 1):ncv]) .< prec)
    nevnew += min(nconv, div(ncv - nevnew, 2))
    if nevnew == 1 && ncv >= 6
        nevnew = div(ncv, 2)
    elseif nevnew == 1 && ncv > 2
        nevnew = 2
    end

    return nevnew
end

function symeigs{T}(A::AbstractMatOp{T};
                    nev::Int = 6, ncv::Int = min(nrow(A), max(2 * nev + 1, 20)),
                    which::Symbol = :LM,
                    tol::T = 1e-8, maxiter::Int = 300,
                    returnvec::Bool = true,
                    v0::Vector{T} = rand(T, nrow(A)) - convert(T, 0.5))
    ## Size of matrix
    n = nrow(A)
    ## Check arguments
    if nev < 1 || nev > n - 1
        error("nev must satisfy 1 <= nev <= n - 1, n is the size of matrix")
    end
    if ncv <= nev || ncv > n
        error("ncv must satisfy nev < ncv <= n, n is the size of matrix")
    end

    ## Number of matrix operations called
    nmatop::Int = 0
    ## Number of restarting iteration
    niter::Int = maxiter
    ## Number of converged eigenvalues
    nconv::Int = 0

    ## Matrices and vectors in the Arnoldi factorization
    V::Matrix{T} = zeros(T, n, ncv)
    H::Matrix{T} = zeros(T, ncv, ncv)
    f::Vector{T} = zeros(T, n)
    ritzval::Vector{T} = zeros(T, ncv)
    ritzvec::Matrix{T} = zeros(T, ncv, nev)
    ritzest::Vector{T} = zeros(T, ncv)
    ritzconv::Vector{Bool} = zeros(Bool, nev)

    ## Precision parameter used to test convergence
    prec::T = eps(T)^(convert(T, 2.0) / 3)

    ## Initialize vectors
    v0norm = norm(v0)
    if v0norm < prec
        error("initial residual vector cannot be zero")
    end
    v0[:] /= v0norm
    w::Vector{T} = performop(A, v0)
    nmatop += 1
    H[1, 1] = dot(v0, w)
    f[:] = w - v0 * H[1, 1]
    V[:, 1] = v0

    ## First Arnoldi factorization
    arnoldifac!(1, ncv, V, H, f, A, prec)
    nmatop += (ncv - 1)
    ritzpairs!(which, H, ritzval, ritzvec, ritzest)
    ## Restarting
    thresh::Vector{T} = zeros(T, nev)
    resid::Vector{T} = zeros(T, nev)
    for i = 1:maxiter
        ## Convergence test
        thresh[:] = tol * max(abs(ritzval[1:nev]), prec)
        resid[:] = abs(ritzest[1:nev]) * norm(f)
        ritzconv[:] = (resid .< thresh)
        nconv = sum(ritzconv)

        if nconv >= nev
            niter = i
            break
        end

        nevadj::Int = nevadjusted(nev, ncv, nconv, ritzest, prec)

        applyshifts!(nevadj, V, H, f, ritzval)
        arnoldifac!(nevadj, ncv, V, H, f, A, prec)
        nmatop += (ncv - nevadj)
        ritzpairs!(which, H, ritzval, ritzvec, ritzest)
    end

    ## Final sorting of Ritz values
    ix = sortperm(ritzval[1:nev], by = abs, rev = true)
    converged = ix[ritzconv[ix]]
    eigval = ritzval[converged]
    if returnvec
        eigvec = V * ritzvec[:, converged]
    else
        eigvec = nothing
    end

    return eigval, eigvec, nconv, niter, nmatop
end


text:
@fastmath function outerloop!(
    pcs::PotentialCalculationSetup{T},
    use_nthreads::Int,
    update_even_points::Val{even_points},
    depletion_handling::Val{depletion_handling_enabled},
    is_weighting_potential::Val{_is_weighting_potential},
    only2d::Val{only_2d}
)::Nothing where {T, even_points, depletion_handling_enabled, _is_weighting_potential, only_2d}
    @inbounds begin
        rb_tar_idx::Int, rb_src_idx::Int = even_points ? (rb_even::Int, rb_odd::Int) : (rb_odd::Int, rb_even::Int)

        @onthreads 1:use_nthreads for i3 in workpart(2:1:(size(pcs.potential, 3) - 1), 1:use_nthreads, Base.Threads.threadid())
            middleloop!( i3, rb_tar_idx, rb_src_idx, pcs,
                         update_even_points, depletion_handling,
                         is_weighting_potential, only2d,
                         Val(iseven(i3))::Union{Val{true}, Val{false}},
                         Val(i3 == 2)::Union{Val{true}, Val{false}}
            )
        end
    end
    nothing
end


text:
""" Data structure for connecting one layer to another by wires.

This module defines a generic data structure to represent a wiring between one
layer of input ports and another layer of output ports. A wiring layer forms a
bipartite graph with independent edge sets the input ports and the output ports.

Wiring layers are an auxillary data structure. They are not very interesting in
their own right, but they can be a useful intermediate representation. For
example, a morphism expression comprising generators, compositions, products,
and wiring layers is intermediate between a pure GAT expression (which has no
wiring layers, but may have identities, braidings, copies, etc.) and a wiring
diagram, which is purely graphical.
"""
module WiringLayers
export WiringLayer, NLayer, nwires, wires, has_wire, add_wire!, add_wires!,
  rem_wire!, rem_wires!, in_wires, out_wires,
  complete_layer, to_wiring_diagram, wiring_layer_between,
  dom, codom, id, compose, otimes, munit, braid, mcopy, delete, mmerge, create

using AutoHashEquals
using DataStructures: OrderedDict

using ...GAT, ...Syntax
import ...Doctrines: ObExpr, HomExpr, MonoidalCategoryWithBidiagonals,
  dom, codom, id, compose, otimes, munit, braid, mcopy, delete, mmerge, create
using ..WiringDiagramCore
import ..WiringDiagramCore: nwires, wires, has_wire, add_wire!, add_wires!,
  rem_wire!, rem_wires!, in_wires, out_wires

# Data types
############

""" Connection by wires of one layer to another.

Morphism in the category of wiring layers.
"""
@auto_hash_equals struct WiringLayer
  wires::OrderedDict{Int}{OrderedDict{Int}{Int}} # src -> tgt -> multiplicity
  ninputs::Int
  noutputs::Int
  function WiringLayer(ninputs::Int, noutputs::Int)
    new(OrderedDict{Int}{OrderedDict{Int}{Int}}(), ninputs, noutputs)
  end
end

""" Number of input or output ports in a layer.

Object in the category of wiring layers.
"""
@auto_hash_equals struct NLayer
  n::Int
end

# Imperative interface
######################

# Constructors.

WiringLayer(inputs::NLayer, outputs::NLayer) = WiringLayer(inputs.n, outputs.n)

function WiringLayer(wires, ninputs, noutputs)
  f = WiringLayer(ninputs, noutputs)
  add_wires!(f, wires)
  f
end

""" Completely connected wiring layer.

The layer's underlying graph is the complete bipartite graph.
"""
function complete_layer(ninputs::Int, noutputs::Int)
  f = WiringLayer(ninputs, noutputs)
  add_wires!(f, i=>j for j in 1:noutputs for i in 1:ninputs)
  f
end

# Basic accessors.

has_wire(f::WiringLayer, wire) = nwires(f, wire) > 0

function nwires(f::WiringLayer, wire)
  src, tgt = wire
  tgts = get(f.wires, src) do; Dict() end
  get(tgts, tgt, 0)
end

nwires(f::WiringLayer) = sum(Int[ sum(values(d)) for d in values(f.wires) ])
wires(f::WiringLayer) = vcat((out_wires(f, src) for src in keys(f.wires))...)

function out_wires(f::WiringLayer, src::Int)
  tgts = get(f.wires, src) do; Dict() end
  vcat((repeat([src => tgt], n) for (tgt, n) in tgts)...)
end

# Layer mutation.

function add_wire!(f::WiringLayer, wire)
  check_wire_bounds(f, wire)
  src, tgt = wire
  tgts = get!(f.wires, src) do; OrderedDict{Int}{Int}() end
  if haskey(tgts, tgt)
    tgts[tgt] += 1
  else
    tgts[tgt] = 1
  end
end

function add_wires!(f::WiringLayer, wires)
  for wire in wires
    add_wire!(f, wire)
  end
end

function rem_wire!(f::WiringLayer, wire)
  src, tgt = wire
  tgts = f.wires[src]
  if tgts[tgt] > 1
    tgts[tgt] -= 1
  else
    delete!(tgts, tgt)
  end
  if isempty(tgts)
    delete!(f.wires, src)
  end
end

function rem_wires!(f::WiringLayer, wire)
  src, tgt = wire
  tgts = get(f.wires, src) do; Dict() end
  pop!(tgts, tgt, nothing)
  if isempty(tgts)
    pop!(f.wires, src, nothing)
  end
end

function check_wire_bounds(f::WiringLayer, wire)
  src, tgt = wire
  if !(1 <= src <= f.ninputs)
    throw(BoundsError("layer inputs 1:$(f.ninputs)", src))
  end
  if !(1 <= tgt <= f.noutputs)
    throw(BoundsError("layer outputs 1:$(f.noutputs)", tgt))
  end
end

# Categorical interface
#######################

NLayer(ob::ObExpr) = NLayer(ndims(ob))
WiringLayer(inputs::ObExpr, outputs::ObExpr) =
  WiringLayer(ndims(inputs), ndims(outputs))

""" Wiring layers as *monoidal category with diagonals and codiagonals*.
"""
@instance MonoidalCategoryWithBidiagonals(NLayer, WiringLayer) begin
  dom(f::WiringLayer) = NLayer(f.ninputs)
  codom(f::WiringLayer) = NLayer(f.noutputs)

  function id(A::NLayer)
    f = WiringLayer(A, A)
    add_wires!(f, i=>i for i in 1:A.n)
    f
  end

  function compose(f::WiringLayer, g::WiringLayer)
    if codom(f) != dom(g)
      error("Incompatible domains $(codom(f)) and $(dom(g))")
    end
    h = WiringLayer(dom(f), codom(g))
    for (src, mid) in wires(f)
      for (_, tgt) in out_wires(g, mid)
        add_wire!(h, src => tgt)
      end
    end
    h
  end

  otimes(A::NLayer, B::NLayer) = NLayer(A.n + B.n)
  munit(::Type{NLayer}) = NLayer(0)

  function otimes(f::WiringLayer, g::WiringLayer)
    h = WiringLayer(otimes(dom(f),dom(g)), otimes(codom(f),codom(g)))
    m, n = dom(f).n, codom(f).n
    add_wires!(h, wires(f))
    add_wires!(h, src+m => tgt+n for (src, tgt) in wires(g))
    h
  end

  function braid(A::NLayer, B::NLayer)
    f = WiringLayer(otimes(A,B), otimes(B,A))
    add_wires!(f, i => i+B.n for i in 1:A.n)
    add_wires!(f, i+A.n => i for i in 1:B.n)
    f
  end

  mcopy(A::NLayer) = mcopy(A, 2)
  mmerge(A::NLayer) = mmerge(A, 2)
  delete(A::NLayer) = WiringLayer(A, munit(NLayer))
  create(A::NLayer) = WiringLayer(munit(NLayer), A)
end

function mcopy(A::NLayer, n::Int)
  f = WiringLayer(A, otimes([A for j in 1:n]))
  add_wires!(f, i => i+A.n*(j-1) for i in 1:A.n, j in 1:n)
  f
end

function mmerge(A::NLayer, n::Int)
  f = WiringLayer(otimes([A for j in 1:n]), A)
  add_wires!(f, i+A.n*(j-1) => i for i in 1:A.n, j in 1:n)
  f
end

# Wiring diagrams
#################

""" Convert a wiring layer into a wiring diagram.
"""
function to_wiring_diagram(layer::WiringLayer, inputs::Vector, outputs::Vector)
  @assert length(inputs) == layer.ninputs && length(outputs) == layer.noutputs
  diagram = WiringDiagram(inputs, outputs)
  add_wires!(diagram, ((input_id(diagram), src) => (output_id(diagram), tgt)
                       for (src, tgt) in sort!(wires(layer))))
  diagram
end

""" Wiring layer representing the wires between two boxes in a wiring diagram.
"""
function wiring_layer_between(diagram::WiringDiagram, v1::Int, v2::Int)::WiringLayer
  inputs, outputs = output_ports(diagram, v1), input_ports(diagram, v2)
  layer = WiringLayer(length(inputs), length(outputs))
  add_wires!(layer, (wire.source.port => wire.target.port
                     for wire in wires(diagram, v1, v2)))
  layer
end

end


text:
"""
    full_matrix(system)

    returns matrix from a simulation step's linear system (i.e., A from "Ax = b")

    system: System
"""
full_matrix(system::System) = full_matrix(system.matrix_entries, system.dimrow, system.dimcol)

function full_matrix(matrix_entries::SparseMatrixCSC, dimrow, dimcol)
    range_row = [1:dimrow[1]]
    for (i,dim) in enumerate(collect(Iterators.rest(dimrow, 2)))
        push!(range_row, sum(dimrow[1:i])+1:sum(dimrow[1:i])+dim)
    end
    range_col = [1:dimcol[1]]
    for (i,dim) in enumerate(collect(Iterators.rest(dimcol, 2)))
        push!(range_col, sum(dimcol[1:i])+1:sum(dimcol[1:i])+dim)
    end

    A = zeros(sum(dimrow),sum(dimcol))

    for (i,row) in enumerate(matrix_entries.rowval)
        col = findfirst(x->i<x,matrix_entries.colptr)-1
            A[range_row[row],range_col[col]] = matrix_entries[row,col].value
    end
    return A
end

# assemble gradient (i.e., b from "Ax = b")
"""
    full_vector(system)

    returns residual vector from a simulation step's linear system (i.e., b from "Ax = b")

    system: System
"""
full_vector(system) = vcat(getfield.(system.vector_entries,:value)...)


text:
using StatsBase: sample

export pair_ring, pair_square, cnot_entangler
export rotor, merged_rotor, rotorset
export random_diff_circuit
export rand_single_gate, rand_gate, rand_circuit

################## Entangler ###################
"""
    pair_ring(n::Int) -> Vector

Pair ring.
"""
pair_ring(n::Int) = [i=>mod(i, n)+1 for i=1:n]

"""
    pair_square(m::Int, n::Int) -> Vector

Pair square.
"""
function pair_square(m::Int, n::Int)
    nsite = m*n
    res = Vector{Pair{Int, Int}}(undef, 2*nsite)
    li = LinearIndices((m, n))
    k = 1
    for i = 1:2:m, j=1:n
        res[k] = li[i, j] => li[i%m+1, j]
        k+=1
    end
    for i = 2:2:m, j=1:n
        res[k] = li[i, j] => li[i%m+1, j]
        k+=1
    end
    for i = 1:m, j=1:2:n
        res[k] = li[i, j] => li[i, j%n+1]
        k+=1
    end
    for i = 1:m, j=2:2:n
        res[k] = li[i, j] => li[i, j%n+1]
        k+=1
    end
    res
end

"""
    cnot_entangler([n::Int, ] pairs::Vector{Pair}) = ChainBlock

Arbitrary rotation unit, support lazy construction.
"""
cnot_entangler(n::Int, pairs) = chain(n, control(n, [ctrl], target=>X) for (ctrl, target) in pairs)
cnot_entangler(pairs) = n->cnot_entangler(n, pairs)

###################### rotor and rotorset #####################
"""
    merged_rotor(noleading::Bool=false, notrailing::Bool=false) -> ChainBlock{1, ComplexF64}

Single qubit arbitrary rotation unit, set parameters notrailing, noleading true to remove trailing and leading Z gates.

!!! note

    Here, `merged` means `Rz(η)⋅Rx(θ)⋅Rz(ξ)` are multiplied first, this kind of operation if now allowed in differentiable
    circuit with back-propagation (`:BP`) mode (just because we are lazy to implement it!).
    But is a welcoming component in quantum differentiation.
"""
merged_rotor(noleading::Bool=false, notrailing::Bool=false) = noleading ? (notrailing ? Rx(0) : chain(Rx(0), Rz(0))) : (notrailing ? chain(Rz(0), Rx(0)) : chain(Rz(0), Rx(0), Rz(0)))

"""
    rotor(nbit::Int, ibit::Int, noleading::Bool=false, notrailing::Bool=false) -> ChainBlock{nbit, ComplexF64}

Arbitrary rotation unit (put in `nbit` space), set parameters notrailing, noleading true to remove trailing and leading Z gates.
"""
function rotor(nbit::Int, ibit::Int, noleading::Bool=false, notrailing::Bool=false)
    rt = chain(nbit, [put(nbit, ibit=>Rz(0.0)), put(nbit, ibit=>Rx(0.0)), put(nbit, ibit=>Rz(0.0))])
    noleading && popfirst!(rt)
    notrailing && pop!(rt)
    rt
end

rotorset(::Val{:Merged}, nbit::Int, noleading::Bool=false, notrailing::Bool=false) = chain(nbit, [put(nbit, j=>merged_rotor(noleading, notrailing)) for j=1:nbit])
rotorset(::Val{:Split}, nbit::Int, noleading::Bool=false, notrailing::Bool=false) = chain(nbit, [rotor(nbit, j, noleading, notrailing) for j=1:nbit])
rotorset(mode::Symbol, nbit::Int, noleading::Bool=false, notrailing::Bool=false) = rotorset(Val(mode), nbit, noleading, notrailing)

"""
    random_diff_circuit(nbit, nlayer, pairs; mode=:Split, do_cache=false)

A kind of widely used differentiable quantum circuit, angles in the circuit is randomely initialized.

ref:
    1. Kandala, A., Mezzacapo, A., Temme, K., Takita, M., Chow, J. M., & Gambetta, J. M. (2017).
       Hardware-efficient Quantum Optimizer for Small Molecules and Quantum Magnets. Nature Publishing Group, 549(7671), 242–246.
       https://doi.org/10.1038/nature23879.
"""
function random_diff_circuit(nbit, nlayer, pairs; mode=:Split, do_cache=false)
    circuit = chain(nbit)

    ent = cnot_entangler(pairs)
    if do_cache
        ent = ent |> cache
    end
    for i = 1:(nlayer + 1)
        i!=1 && push!(circuit, ent)
        push!(circuit, rotorset(mode, nbit, i==1, i==nlayer+1))
    end
    circuit
end

############### Completely random circuits (for testing and demo) ################
randlocs(nbit::Int, mbit::Int) = sample(1:nbit, mbit, replace=false)
const SINGLE_GATES = [X, Y, Z, H, Rx, Ry, Rz, shift, phase]

rand_single_gate(ngate::Int) = [rand_single_gates() for i=1:ngate]
function rand_single_gate()
    gate = rand(SINGLE_GATES)
    gate isa AbstractBlock ? gate : gate(rand()*2π)
end

"""
    rand_gate(nbit::Int, mbit::Int, [ngate::Int]) -> MatrixBlock

random nbit gate.
"""
rand_gate(nbit::Int, mbit::Int) = rand_gate(nbit, Val(mbit))
rand_gate(nbit::Int, mbit::Int, ngate::Int) = [rand_gate(nbit, mbit) for i=1:ngate]
rand_gate(nbit::Int, ::Val{1}) = put(nbit, rand(1:nbit)=>rand_single_gate())
function rand_gate(nbit::Int, ::Val{M}) where M
    locs = randlocs(nbit, M)
    control(nbit, locs[1:M-1], last(locs)=>rand_single_gate())
end

function rand_circuit(nbit::Int; p1 = (nbit==1 ? 1.0 : 0.66), ngate=5*nbit)
    c = chain(nbit)
    for i=1:ngate
        if rand() < p1
            push!(c, rand_gate(nbit, 1))
        else
            push!(c, rand_gate(nbit, 2))
        end
    end
    c
end


text:
# Run the tests
testdir = joinpath(dirname(dirname(@__DIR__)), "test")
include("$testdir/runtests.jl")