text:
"""
     word2vec(train, output; size=100, window=5, sample=1e-3, hs=0,  negative=5, threads=12, iter=5, min_count=5, alpha=0.025, debug=2, binary=1, cbow=1, save_vocal=Nothing(), read_vocab=Nothing(), verbose=false,)

    Parameters for training:
        train <file>
            Use text data from <file> to train the model
        output <file>
            Use <file> to save the resulting word vectors / word clusters
        size <Int>
            Set size of word vectors; default is 100
        window <Int>
            Set max skip length between words; default is 5
        sample <AbstractFloat>
            Set threshold for occurrence of words. Those that appear with
            higher frequency in the training data will be randomly
            down-sampled; default is 1e-5.
        hs <Int>
            Use Hierarchical Softmax; default is 1 (0 = not used)
        negative <Int>
            Number of negative examples; default is 0, common values are
            5 - 10 (0 = not used)
        threads <Int>
            Use <Int> threads (default 12)
        iter <Int>
            Run more training iterations (default 5)
        min_count <Int>
            This will discard words that appear less than <Int> times; default
            is 5
        alpha <AbstractFloat>
            Set the starting learning rate; default is 0.025
        debug <Int>
            Set the debug mode (default = 2 = more info during training)
        binary <Int>
            Save the resulting vectors in binary moded; default is 0 (off)
        cbow <Int>
            Use the continuous back of words model; default is 1 (skip-gram
            model)
        save_vocab <file>
            The vocabulary will be saved to <file>
        read_vocab <file>
            The vocabulary will be read from <file>, not constructed from the
            training data
        verbose <Bool>
            Print output from training
"""
function word2vec(train::AbstractString, output::AbstractString;
                  size::Int=100, window::Int=5, sample::AbstractFloat=1e-3,
                  hs::Int=0, negative::Int=5, threads::Int=12, iter::Int=5,
                  min_count::Int=5, alpha::AbstractFloat=0.025,
                  debug::Int=2, binary::Int=0, cbow::Int=1,
                  save_vocab=Nothing(), read_vocab=Nothing(),
                  verbose::Bool=false)

    parameters = AbstractString[]
    args = ["-train", "-output", "-size", "-window", "-sample", "-hs",
            "-negative", "-threads", "-iter", "-min-count", "-alpha",
            "-debug", "-binary", "-cbow"]
    values = [train, output, size, window, sample, hs, negative, threads,
              iter, min_count, alpha, debug, binary, cbow]

    for (arg, value) in zip(args, values)
        push!(parameters, arg)
        push!(parameters, string(value))
    end
    if save_vocab != Nothing()
        push!(parameters, "-save-vocab")
        push!(parameters, string(save_vocab))
    end
    if read_vocab != Nothing()
        push!(parameters, "-read-vocab")
        push!(parameters, string(read_vocab))
    end
    Word2Vec_jll.word2vec() do command
        run(`$(command) $(parameters)`)
    end
end

"""
```
     word2clusters(train, output, classes; size=100, window=5, sample=1e-3, hs=0,  negative=5, threads=1, iter=5, min_count=5, alpha=0.025, debug=2, binary=1, cbow=1, save_vocal=Nothing(), read_vocab=Nothing(), verbose=false)
```
    Parameters for training:
        train <file>
            Use text data from <file> to train the model
        output <file>
            Use <file> to save the resulting word vectors / word clusters
        size <Int>
            Set size of word vectors; default is 100
        window <Int>
            Set max skip length between words; default is 5
        sample <AbstractFloat>
            Set threshold for occurrence of words. Those that appear with
            higher frequency in the training data will be randomly
            down-sampled; default is 0 (off), useful value is 1e-5
        hs <Int>
            Use Hierarchical Softmax; default is 1 (0 = not used)
        negative <Int>
            Number of negative examples; default is 0, common values are 5 - 10
            (0 = not used)
        threads <Int>
            Use <Int> threads (default 1)
        iter <Int>
            Run more training iterations (default 5)
        min_count <Int>
            This will discard words that appear less than <Int> times
            (default 5)
        alpha <AbstractFloat>
            Set the starting learning rate; default is 0.025
        classes <Int>
            Number of word classes; if 0, output word classes rather than
            word vectors (default 0)
        debug <Int>
            Set the debug mode (default = 2 = more info during training)
        binary <Int>
            Save the resulting vectors in binary moded; default is 0 (off)
        cbow <Int>
            Use the continuous back of words model; default is 1 (0 for skip-gram
            model)
        save_vocab <file>
            The vocabulary will be saved to <file>
        read_vocab <file>
            The vocabulary will be read from <file>, not constructed from the
            training data
        verbose <Bool>
            Print output from training
"""
function word2clusters(train::AbstractString, output::AbstractString,
                       classes::Int; size::Int=100, window::Int=5,
                       sample::AbstractFloat=0., hs::Int=0,
                       negative::Int=5, threads::Int=1, iter::Int=5,
                       min_count::Int=5, alpha::AbstractFloat=0.025,
                       debug::Int=2, binary::Int=0, cbow::Int=1,
                       save_vocab=Nothing(), read_vocab=Nothing(),
                       verbose::Bool=false)

    parameters = AbstractString[]
    args = ["-train", "-output", "-size", "-window", "-sample", "-hs",
            "-negative", "-threads", "-iter", "-min-count", "-alpha",
            "-debug", "-binary", "-cbow", "-classes"]
    values = [train, output, size, window, sample, hs, negative, threads,
              iter, min_count, alpha, debug, binary, cbow, classes]
    for (arg, value) in zip(args, values)
        push!(parameters, arg)
        push!(parameters, string(value))
    end
    if save_vocab != Nothing()
        push!(parameters, "-save-vocab")
        push!(parameters, string(save_vocab))
    end
    if read_vocab != Nothing()
        push!(parameters, "-read-vocab")
        push!(parameters, string(read_vocab))
    end
    Word2Vec_jll.word2vec() do command
        run(`$(command) $(parameters)`)
    end
end

"""
     word2phrase(train, output; min_count=5, threshold=100, debug=2)

    Parameters for training:
	train <file>
	      Use text data from <file> to train the model
	output <file>
              Use <file> to save the resulting word vectors /
              word clusters / phrases
	min_count <Int>
              This will discard words that appear less than <int> times;
              default is 5
	threshold <AbstractFloat>
      	      The <AbstractFloat> value represents threshold for
              forming the phrases (higher means less phrases); default 100
	debug <Int>
	      Set the debug mode (default = 2 = more info during training)
"""
function word2phrase(train::AbstractString, output::AbstractString;
                     min_count::Int=5, threshold::Int=100, debug::Int=2)
    parameters = AbstractString[]
    args = ["-train", "-output", "-min-count", "-threshold", "-debug"]
    values = [train, output, min_count, threshold, debug]
    for (arg, value) in zip(args, values)
        push!(parameters, arg)
        push!(parameters, string(value))
    end
    Word2Vec_jll.word2phrase() do command
        run(`$(command) $(parameters)`)
    end
end


text:
@testset "889.construct-binary-tree-from-preorder-and-postorder-traversal.jl" begin
    @test construct_from_pre_post([1, 2, 4, 5, 3, 6, 7], [4, 5, 2, 6, 7, 3, 1]) ==
          TreeNode(1, TreeNode(2, TreeNode(4), TreeNode(5)),
                   TreeNode(3, TreeNode(6), TreeNode(7)))

    @test construct_from_pre_post(Int[], Int[]) === nothing
end

ext: jl


text:
using Test
using JSExpr
using JSExpr: @crawl, deparse, JSTerminal, JSString

@testset "array" begin
        @test @js([1, 2, 3]) == JSString("""[1,2,3]""")
        @test @js((1, 2, 3)) == JSString("""[1,2,3]""")

        my_array = ["foo", "bar"]
        @test @js(console.log($my_array)) == JSString(
                """console.log(["foo","bar"])"""
        )

        @test @js([[], []]) == JSString("""[[],[]]""")

        @testset "array indexing" begin
            @test @js([1, 2, 3][0]) == JSString("[1,2,3][0]")
        end

        @testset "tuple unpacking" begin
            @test @js(
                for (key, value) in Object.entries(myObj)
                    console.log(key, value)
                end
            ) == js"""for (let [key,value] of Object.entries(myObj)) { console.log(key, value); }"""
        end
end


text:
@testset "AbstractGrid" begin

    struct SmallGrid{dim,N,C<:Ferrite.AbstractCell} <: Ferrite.AbstractGrid{dim}
        nodes_test::Vector{NTuple{dim,Float64}}
        cells_test::NTuple{N,C}
    end

    Ferrite.getcells(grid::SmallGrid) = grid.cells_test
    Ferrite.getcells(grid::SmallGrid, v::Union{Int, Vector{Int}}) = grid.cells_test[v]
    Ferrite.getncells(grid::SmallGrid{dim,N}) where {dim,N} = N
    Ferrite.getcelltype(grid::SmallGrid) = eltype(grid.cells_test)
    Ferrite.getcelltype(grid::SmallGrid, i::Int) = typeof(grid.cells_test[i])
    Ferrite.getcoordinates(x::NTuple{dim,Float64}) where dim = Vec{dim,Float64}(x)

    Ferrite.getnodes(grid::SmallGrid) = grid.nodes_test
    Ferrite.getnodes(grid::SmallGrid, v::Union{Int, Vector{Int}}) = grid.nodes_test[v]
    Ferrite.getnnodes(grid::SmallGrid) = length(grid.nodes_test)
    Ferrite.nnodes_per_cell(grid::SmallGrid, i::Int=1) = Ferrite.nnodes(grid.cells_test[i])
    Ferrite.n_faces_per_cell(grid::SmallGrid) = nfaces(eltype(grid.cells_test))
    function Ferrite.getcoordinates!(x::Vector{Vec{dim,T}}, grid::SmallGrid, cell::Int) where {dim,T}
        for i in 1:length(x)
            x[i] = Vec{dim,T}(grid.nodes_test[grid.cells_test[cell].nodes[i]])
        end
    end
    function Ferrite.getcoordinates(grid::SmallGrid{dim}, cell::Int) where dim
        nodeidx = grid.cells_test[cell].nodes
        return [Vec{dim,Float64}(grid.nodes_test[i]) for i in nodeidx]::Vector{Vec{dim,Float64}}
    end

    nodes = [(-1.0,-1.0); (0.0,-1.0); (1.0,-1.0); (-1.0,0.0); (0.0,0.0); (1.0,0.0); (-1.0,1.0); (0.0,1.0); (1.0,1.0)]
    cells = (Quadrilateral((1,2,5,4)), Quadrilateral((2,3,6,5)), Quadrilateral((4,5,8,7)), Quadrilateral((5,6,9,8)))
    subtype_grid = SmallGrid(nodes,cells)
    reference_grid = generate_grid(Quadrilateral, (2,2))

    ip = Lagrange{2, RefCube, 1}()
    qr = QuadratureRule{2, RefCube}(2)
    cellvalues = CellScalarValues(qr, ip);

    dhs = [DofHandler(grid) for grid in (subtype_grid, reference_grid)]
    u1 = Vector{Float64}(undef, 9)
    u2 = Vector{Float64}(undef, 9)
    ∂Ω = union(getfaceset.((reference_grid, ), ["left", "right", "top", "bottom"])...)
    dbc = Dirichlet(:u, ∂Ω, (x, t) -> 0)

    function doassemble!(cellvalues::CellScalarValues{dim}, K::SparseMatrixCSC, dh::DofHandler) where {dim}
        n_basefuncs = getnbasefunctions(cellvalues)
        Ke = zeros(n_basefuncs, n_basefuncs)
        fe = zeros(n_basefuncs)
        f = zeros(ndofs(dh))
        assembler = start_assemble(K, f)
        for cellid in 1:getncells(dh.grid)
            fill!(Ke, 0)
            fill!(fe, 0)
            coords = getcoordinates(dh.grid, cellid)
            reinit!(cellvalues, coords)
            for q_point in 1:getnquadpoints(cellvalues)
                dΩ = getdetJdV(cellvalues, q_point)
                for i in 1:n_basefuncs
                    v  = shape_value(cellvalues, q_point, i)
                    ∇v = shape_gradient(cellvalues, q_point, i)
                    fe[i] += v * dΩ
                    for j in 1:n_basefuncs
                        ∇u = shape_gradient(cellvalues, q_point, j)
                        Ke[i, j] += (∇v ⋅ ∇u) * dΩ
                    end
                end
            end
            assemble!(assembler, celldofs(dh,cellid), fe, Ke)
        end
        return K, f
    end

    for (dh,u) in zip(dhs,(u1,u2))
        push!(dh, :u, 1)
        close!(dh)
        ch = ConstraintHandler(dh)
        add!(ch, dbc)
        close!(ch)
        update!(ch, 0.0)
        K = create_sparsity_pattern(dh);
        K, f = doassemble!(cellvalues, K, dh);
        apply!(K, f, ch)
        sol = K \ f
        u .= sol
    end

    @test Ferrite.ndofs_per_cell(dhs[1]) == Ferrite.ndofs_per_cell(dhs[2])
    @test Ferrite.celldofs(dhs[1],3) == Ferrite.celldofs(dhs[2],3)
    @test Ferrite.ndofs(dhs[1]) == Ferrite.ndofs(dhs[2])
    @test isapprox(u1,u2,atol=1e-8)
end


text:
# Author: AJ Tyler
# Date: 24/01/22
module ClusterStructureAnalysis
include("io.jl") # This file contains the IO module.
include("comparison.jl") # This file contains the Comparison module
include("visualisation.jl") # This file contains the Visualisation module.
include("classify.jl") # This file contains the Classification module.
using .Comparison # This module the contains functions which compare configurations.
using .IO # This module contains functions that read in and write out files.
using .Visualisation # This module contains functions which create a .vesta file with helpful colourisations.
using .Classify # This module contains functions which classify configurations based on their symmetries.
using Graphs # This module contains graph theory functions which allow for a graphical representation of configurations.
export classify # This function uses the atomic CNA profile of a configuration to identify symmetries.
export compare # This function computed the CNA profiles of configurations oevr a range of rCut values.

"""
This function computed the CNA profiles of configurations over a range of rCut values. These configurations are then compared between
eachother using a similarity score. If these similarity score is high enough over enough rCut values, then the configurations are
considered to be the same. For each input .xyz file, a .txt file is generated inside a rCut_XXX directory for each rCut value (XXX)
containing the total CNA profiles of each configuration. A comparison_NAME.txt file is also generated inside the comparison directory
for each input file, which lists which configurations were most similar to each configuration and for which rCuts. Finally
a unique_NAME.txt file is also created inside the comparison directory which lists each set of distinguishable configurations and their
associated energy statistics.

Inputs:
configDir: (String) The path to (or name of, if in the local directory) the directory containing the .xyz files.
					The .xyz files must have the following format NAME_N_L.xyz, where NAME is unique for all the files in the
					directory, N is the number of atoms in each configuration and L is the number of configurations in the file.
Keyword Arguments:
rCutRange: (Vector{Float64}) A range of cut-off radii to test, in units of EBL. Default is 10 values between 4/3 and 2.
similarityMeasure: (String) The method of comparing configurations, either "atomic" or "total" CNA. Default is total.
similarityThreshold: (Float64) Similarity threshold above which two configurations are considered identical.
EBL: (Float64) Equilibrium Bond Length. Default is 3.1227 Angstroms (from MP2 data for Neon2).
"""
function compare(configDir::String; rCutRange::Vector{Float64} = round.(LinRange(4/3,2,10);digits = 3), similarityMeasure = "total", similarityThreshold = 0.95,EBL=3.1227)
	rm("$configDir/comparison";force = true,recursive = true) # Delete comparision directory if exists from previous run
	mkdir("$configDir/comparison") # Make comparision directory
	M = length(rCutRange) # Number of rCut values to test
	files = readdir(configDir) # Create list of all files/directories in the configuration directory
	for file in files # For each file
		isXYZ,fileName,fp,L,N,B = processFileName(configDir,file) # Extract information from filename
		if isXYZ # If file is a .xyz file
			# Compute CNA profiles of all configurations in file
			energies,totalProfiles, atomicProfiles,_,_,_ = processFile(fp,L,N,B,true,EBL,rCutRange;M)
			# Find most similar configurations for each configuration based on CNA profiles over all rCut values
			similarConfigs, maxSims = compareConfigs(rCutRange,energies,totalProfiles, atomicProfiles, configDir,fileName, similarityMeasure,M,L,N)
			if L == 1 # If only have one configuration in the file
				continue # Don't need to do any comparisons
			end
			sortingArray = sortperm(maxSims[1,:],rev=true) # Sort configurations by maximum similarity values in decreasing order
			# Sort arrays
			maxSims, sortedEnergies, similarConfigs = maxSims[sortingArray], energies[sortingArray], similarConfigs[sortingArray]
			# Write out configuration comparison results
			writeComparison(configDir,fileName, sortingArray, sortedEnergies, maxSims, similarConfigs,L)
			# Group configurations into distinguishable sets
			uniqueConfigs,uniqueEnergies,uniqueSimilarities = groupConfigs(energies,sortingArray,maxSims, similarConfigs, similarityThreshold,M)
			numUnique = length(uniqueSimilarities) # Compute how many distinguishable sets are in the file
			uniqueStats = configStats(uniqueEnergies,numUnique) # Compute statistics on energies of each distinguishable set
			writeUnique(configDir,fileName,numUnique,uniqueConfigs,uniqueSimilarities, uniqueStats) # Write out distinguishable sets file
		end
	end
end

"""
This function uses the atomic CNA profile of a configuration to identify symmetries. These symmetries include icosahedral (ICO),
body centred cubic (BCC), face centered cubic (FCC), hexagonal close packed (HCP) and OTHER. Futhermore, if a core symmetry can be
identified, then cap and capping atoms are also identified. For each input .xyz file, a NAME_classification.txt file is generated inside
the classification directory, with which atoms have which symmertry and if they are part of the cluster's core, a cap or a capping atom.
A NAME_i.vesta file is also generated in the visulisation directory, where i is the configuration number in the .xyz file, which uses the
classification information to which configurations were most similar to each configuration and for which rCuts. Finally a unique_NAME.txt
file is created which colour the atoms so aid visualisation of the cluster's structure.

Inputs:
configDir: (String) The path to (or name of, if in the local directory) the directory containing the .xyz files.
					The .xyz files must have the following format NAME_N_L.xyz, where NAME is unique for all the files in the
					directory, N is the number of atoms in each configuration and L is the number of configurations in the file.
Keyword Arguments:
rCut: (Float64) The cut-off radii used in the CNA analysis in units of EBL. Default is 4/3.
EBL: (Float64) Equilibrium Bond Length. Default is 3.1227 Angstroms (from MP2 data for Neon2).
"""
function classify(configDir::String; rCut = 4/3,EBL=3.1227)
	rm("$configDir/classification";force = true,recursive = true) # Delete classification directory if exists from previous run
	mkdir("$configDir/classification") # Create the classification directory
	rm("$configDir/visualisation";force = true,recursive = true) # Delete visualisation directory if exists from previous run
	mkdir("$configDir/visualisation") # Create the visualisation directory
	files = readdir(configDir) # Create list of all files/directories in the configuration directory
	for file in files # For each file
		isXYZ,fileName,fp,L,N,B = processFileName(configDir,file) # Extract information from filename
		if isXYZ # If file is a .xyz file
			# Compute CNA profiles and atom shell numbers of all configurations in file
			_,_, atomicProfiles,configurations, shells, bondGraphs = processFile(fp,L,N,B,false,EBL,[rCut])
			# Identify symmetries of the clusters using CNA profiles
			classifications, capSymmetries = findSymmetries(shells,atomicProfiles,bondGraphs,L)
			writeClassification(configDir,fileName,classifications,capSymmetries,L) # Output found symmetries
			vestaFile(configDir,fileName,configurations,classifications,capSymmetries,rCut,EBL,L,N) # Create colourised .vesta file
		end
	end
end

end


text:
function probstates(ind1::Tuple, ind2::Tuple)
    ListOfTuples = [ind1, ind2]
    Edge = Dict{Tuple{Int64,Int64},Int64}()
    Ct = 1
    for i in 1:4
       for j in (i+1):4
          Edge[(i,j)] = Ct
          Ct += 1
       end
    end
    return Edge
    A = zeros(6)
    for (i, j) in ListOfTuples
        if i > j
            i, j = j, i
        end
        A[Edge[(i,j)]] = 1
    end
    A == [1,1,1,1,1,1] && return 1
    A == [1,0,0,0,0,1] && return 2
    A == [1,1,0,1,0,0] && return 3
    A == [1,0,0,0,0,0] && return 4
    A == [0,1,1,0,0,1] && return 5
    A == [0,0,0,0,0,1] && return 6
    A == [0,1,0,0,1,0] && return 7
    A == [0,1,0,0,0,0] && return 8
    A == [0,0,0,0,0,0] && return 9
end

if A[1] == 1
    if A[2] == 1
        A[5] == 1 ? return 1 : return 3
    else
        A[6] == 1 ? return 2 : return 4
    end
else
    if A[2] == 1
        A[3] == 1 && return 5
        A[5] == 1 ? return 7 : return 8
    else
        A[6] == 1 ? return 6 : return 9
    end
end

####################
####################
#=
Calculate Pr(Li | Sj)
If the allele identity falls into this class (L1-L9), generate the
probabilities of it belonging to each of the different classes and
return that array of 9 distinct probabilities
=#
function probability_state_table(x::Tuple, y::Tuple, alleles::Dict)
    #TODO Improve how groups are decided based on how similar things are done with moments estimators
    x1, x2 = x
    y1, y2 = y
    ## class L1 -  AᵢAᵢ AᵢAᵢ ##
    if x1 == x2 == y1 == y2
        p = alleles[x1]
        (p, p^2, p^2, p^3, p^2, p^3, p^2, p^3, p^4)

    ## class L2 - AᵢAᵢ AⱼAⱼ ##
    elseif x1 == x2 != y1 == y2
        # (x1 == x2) & (y1 == y2) & (x1 != y1)
        p1,p2 = alleles[x1], alleles[y1]
        prod_p = prod((p1,p2))
        [0, prod_p, 0, prod_p * p2, 0, prod_p * p1, 0, 0, prod_p^2]

    ## class L3a - AᵢAᵢ AᵢAⱼ ## - has issues because of allele order
    elseif x1 == x2 == y1 != y2
        #((x1 == x2 == y1) & (x1 != y2))
        p1,p2 = alleles[x1], alleles[y2]
        prod_p = prod(p)
        [0, 0, prod_p, 2 * prod_p * p1, 0, 0, 0, prod_p * p1, 2 * prod_p * p1^2]

    ## class L3b - AᵢAᵢ AⱼAᵢ ## - has issues because of allele order
    elseif x1 == x2 == y2 != y1
        p1,p2 = alleles[x1], alleles[y1]
        prod_p = prod(p)
        [0, 0, prod_p, 2 * prod_p * p1, 0, 0, 0, prod_p * p1, 2 * prod_p * p1^2]

    ## class L4 - AᵢAᵢ AⱼAₖ ##
    elseif x1 == x2 != y1 != y2 != x1
        #(x1 == x2) & (y1 != y2) & (x1 != y1) & (x1 != y2)
        p = (alleles[x1], alleles[y1], alleles[y2])
        prod_p = prod(p)
        [0, 0, 0, 2 * prod_p, 0, 0, 0, 0, 2 * prod_p * p[1]]

    ## L5a - AiAj AiAi ## - has issues because of allele order
    elseif x1 == y1 == y2 != x2
        p1,p2 = alleles[x1], alleles[x2]
        prod_p = prod(p)
        [0, 0, 0, 0, prod_p, 2 * prod_p * p1, 0, prod_p * p1, 2 * prod_p * p1^2]

    ## L5b - AjAi AiAi ## - has issues because of allele order
    elseif x2 == y1 == y2 != x2
        p1,p2 = alleles[x2], alleles[x1]
        prod_p = prod(p)
        [0, 0, 0, 0, prod_p, 2 * prod_p * p1, 0, prod_p *p1, 2 * prod_p * p1^2]

    ## L6 - AjAk AiAi ##
    elseif y1 != x2 != x1 != y1 == y2
        #(x1 != x2) & (y1 == y2) & (x1 != y1) & (x2 != y1)
        p = (alleles[y1], alleles[x1], alleles[x2])
        prod_p = prod(p)
        [0, 0, 0, 0, 0, 2 * prod_p, 0, 0, 2 * prod_p * p[1]]

    ## L7 - AiAj AiAj ##
    elseif y1 == x1 != x2 == y2
        #(x1 == y1) & (x2 == y2) & (x1 != x2)
        p = (alleles[x1], alleles[x2])
        prod_p = prod(p)
        [0, 0, 0, 0, 0, 0, 2 * prod_p, prod_p * sum(p), 4 * prod_p * prod_p]

    ## L8a - AiAj AiAk ##  - has issues because of allele order
    elseif x2 != x1 == y1 != y2 != x2
        #(x1 == y1) & (x1 != x2) & (y1 != y2) & (x2 != y2)
        p = (alleles[x1], alleles[x2], alleles[y2])
        prod_p = prod(p)
        [0, 0, 0, 0, 0, 0, 0, prod_p, 4 * prod_p * p[1]]

    ## L8b - AjAi AkAi ##  - has issues because of allele order
    elseif x1 != x2 == y2 != y1 != x1
        #(x2 == y2) & (x1 != x2) & (y1 != y2) & (x1 != y1)
        p = (alleles[x2], alleles[x1], alleles[y1])
        prod_p = prod(p)
        [0, 0, 0, 0, 0, 0, 0, prod_p, 4 * prod_p * p[1]]

    ## L8c - AjAi AiAk ##  - has issues because of allele order
    elseif x1 != x2 == y1 != y2 != x1
        #(x2 == y1) & (x1 != x2) & (y1 != y2) & (x1 != y2)
        p = (alleles[x2], alleles[x1], alleles[y2])
        prod_p = prod(p)
        [0, 0, 0, 0, 0, 0, 0, prod_p, 4 * prod_p * p[1]]

    ## L8d - AiAj AkAi ##  - has issues because of allele order
    elseif x2 != x1 == y2 != y2 & (x1 != y1)
        #(x1 == y2) & (x1 != x2) & (y1 != y2) & (x1 != y1)
        p = (alleles[x1], alleles[x2], alleles[y1])
        prod_p = prod(p)
        [0, 0, 0, 0, 0, 0, 0, prod_p, 4 * prod_p * p[1]]

    ## L9 - AiAj AkAl ##
    elseif x1 != x2 != y1 != y2 != x1 != y1 & (y2 != x2)
        #(x1 != x2) & (x1 != y1) & (x1 != y2) & (x2 != y1) & (x2 != y2) & (y1 != x2)
        p = (alleles[x1], alleles[x2], alleles[y1], alleles[y2])
        [0, 0, 0, 0, 0, 0, 0, 0, 4 * prod(p)]
    else
        [-9, -9, -9, -9, -9, -9, -9, -9, -9]
    end
end


text:
input = readlines("day13/input.txt")
# input = readlines("day13/small.txt")
# input = readlines("day13/tiny.txt")
depart_time = parse(Int, input[1])
buses = convert.(String, split(input[2], ","))
numeric_buses = filter(n -> tryparse(Int, n[2]) !== nothing, collect(enumerate(buses)))
numeric_buses = map(n -> (n[1], parse(Int, n[2])), numeric_buses)

println(numeric_buses)
function special_mod(a, b)
    return mod(a, b) == 0 ? b : mod(a,b)
end

prev_buses = [1]
on = 0
for (idx, bus) in numeric_buses
    global prev_buses, on
    search_by = prod(prev_buses)
    condition = (num) -> (bus - special_mod(num, bus) == mod(idx - 1, bus))
    while !condition(on)
        on += search_by
    end
    push!(prev_buses, bus)
end

println(on)

# You gotta love modular arithmetic! Today's was fun!


text:
# This file contains a module
# containing the types of numbers
# used in the package, and
# also their according formats

"""
Module containing the types of numbers used in
the package, and also their according formats.
"""
module Prec

# Export
export IT, RT, RF

"""
Type of integer values. Default value: `Int`.
"""
IT = Int

"""
Type of real values. Default value: `typeof(1.0)`.
"""
RT = typeof(1.0)

"""
Format of real values. Default values:
- for [`RT`](@ref) = `Float64`: `"% .15E"`;
- for [`RT`](@ref) = `Float32`: `"% .6E"`;
- for [`RT`](@ref) = `Float16`: `"% .3E"`.
"""
RF = "% .6E"

# Change the format according to the type of a real number
if (RT == Float64)
    RF = "% .15E"
elseif (RT == Float16)
    RF = "% .3E"
end

end


text:
# Time to wait for the AWS Batch cluster to scale up, spot requests to be fufilled,
# instances to finish initializing, and have the worker instances connect to the manager.
# Note that when using 3 workers, it is not uncommon for the third worker to take 15 - 21
# minutes longer than the other 2 workers to start up
const BATCH_TIMEOUT = Minute(25)

# Note: Communication directly between AWS Batch jobs works since the underlying ECS task
# implicitly uses networkMode: host. If this changes to another networking mode AWS Batch
# jobs may no longer be able to listen to incoming connections.

"""
    AWSBatchManager(max_workers; kwargs...)
    AWSBatchManager(min_workers:max_workers; kwargs...)
    AWSBatchManager(min_workers, max_workers; kwargs...)

A cluster manager which spawns workers via [Amazon Web Services Batch](https://aws.amazon.com/batch/)
service. Typically used within an AWS Batch job to add additional resources. The number of
workers spawned may be potentially be lower than the requested `max_workers` due to resource
contention. Specifying `min_workers` can allow the launch to succeed with less than the
requested `max_workers`.

## Arguments
- `min_workers::Int`: The minimum number of workers to spawn or an exception is thrown
- `max_workers::Int`: The number of requested workers to spawn

## Keywords
- `definition::AbstractString`: Name of the AWS Batch job definition which dictates
  properties of the job including the Docker image, IAM role, and command to run
- `name::AbstractString`: Name of the job inside of AWS Batch
- `queue::AbstractString`: The job queue in which workers are submitted. Can be either the
  queue name or the Amazon Resource Name (ARN) of the queue. If not set will default to
  the environmental variable "WORKER_JOB_QUEUE".
- `memory::Integer`: Memory limit (in MiB) for the job container. The container will be killed
  if it exceeds this value.
- `region::AbstractString`: The region in which the API requests are sent and in which new
  worker are spawned. Defaults to "us-east-1". [Available regions for AWS batch](http://docs.aws.amazon.com/general/latest/gr/rande.html#batch_region)
  can be found in the AWS documentation.
- `timeout::Period`: The maximum number of seconds to wait for workers to become available
  before attempting to proceed without the missing workers.

## Examples
```julia
julia> addprocs(AWSBatchManager(3))  # Needs to be run from within a running AWS batch job
```
"""
AWSBatchManager

struct AWSBatchManager <: ContainerManager
    min_workers::Int
    max_workers::Int
    job_definition::AbstractString
    job_name::AbstractString
    job_queue::AbstractString
    job_memory::Integer
    region::AbstractString
    timeout::Second
    min_ip::IPv4
    max_ip::IPv4

    function AWSBatchManager(
        min_workers::Integer,
        max_workers::Integer,
        definition::AbstractString,
        name::AbstractString,
        queue::AbstractString,
        memory::Integer,
        region::AbstractString,
        timeout::Period=BATCH_TIMEOUT,
        min_ip::IPv4=ip"0.0.0.0",
        max_ip::IPv4=ip"255.255.255.255",
    )
        min_workers >= 0 || throw(ArgumentError("min workers must be non-negative"))
        min_workers <= max_workers ||
            throw(ArgumentError("min workers exceeds max workers"))

        # Default the queue to using the WORKER_JOB_QUEUE environmental variable.
        if isempty(queue)
            queue = get(ENV, "WORKER_JOB_QUEUE", "")
        end

        region = isempty(region) ? "us-east-1" : region

        return new(
            min_workers,
            max_workers,
            definition,
            name,
            queue,
            memory,
            region,
            Second(timeout),
            min_ip,
            max_ip,
        )
    end
end

function AWSBatchManager(
    min_workers::Integer,
    max_workers::Integer;
    definition::AbstractString="",
    name::AbstractString="",
    queue::AbstractString="",
    memory::Integer=-1,
    region::AbstractString="",
    timeout::Period=BATCH_TIMEOUT,
    min_ip::IPv4=ip"0.0.0.0",
    max_ip::IPv4=ip"255.255.255.255",
)
    return AWSBatchManager(
        min_workers,
        max_workers,
        definition,
        name,
        queue,
        memory,
        region,
        timeout,
        min_ip,
        max_ip,
    )
end

function AWSBatchManager(workers::UnitRange{<:Integer}; kwargs...)
    return AWSBatchManager(first(workers), last(workers); kwargs...)
end

function AWSBatchManager(workers::Integer; kwargs...)
    return AWSBatchManager(workers, workers; kwargs...)
end

launch_timeout(mgr::AWSBatchManager) = mgr.timeout
desired_workers(mgr::AWSBatchManager) = mgr.min_workers, mgr.max_workers

function Base.:(==)(a::AWSBatchManager, b::AWSBatchManager)
    return (
        a.min_workers == b.min_workers &&
        a.max_workers == b.max_workers &&
        a.job_definition == b.job_definition &&
        a.job_name == b.job_name &&
        a.job_queue == b.job_queue &&
        a.job_memory == b.job_memory &&
        a.region == b.region &&
        a.timeout == b.timeout &&
        a.min_ip == b.min_ip &&
        a.max_ip == b.max_ip
    )
end

function spawn_containers(mgr::AWSBatchManager, override_cmd::Cmd)
    min_workers, max_workers = desired_workers(mgr)
    max_workers < 1 && return nothing

    queue = @mock JobQueue(mgr.job_queue)
    max_compute = @mock max_vcpus(queue)

    if min_workers > max_compute
        error(
            string(
                "Unable to launch the minimum number of workers ($min_workers) as the ",
                "minimum exceeds the max VCPUs available ($max_compute).",
            ),
        )
    elseif max_workers > max_compute
        # Note: In addition to warning the user about the VCPU cap we could also also reduce
        # the number of worker we request. Unfortunately since we don't know how many jobs
        # are currently running or how long they will take we'll leave `max_workers`
        # untouched.
        warn(
            LOGGER,
            string(
                "Due to the max VCPU limit ($max_compute) most likely only a partial amount ",
                "of the requested workers ($max_workers) will be spawned.",
            ),
        )
    end

    # Since each batch worker can only use one cpu we override the vcpus to one.
    job = @mock run_batch(;
        name=mgr.job_name,
        definition=mgr.job_definition,
        queue=mgr.job_queue,
        region=mgr.region,
        vcpus=1,
        memory=mgr.job_memory,
        cmd=override_cmd,
        num_jobs=max_workers,
        allow_job_registration=false,
    )

    if max_workers > 1
        notice(LOGGER, "Spawning array job: $(job.id) (n=$(mgr.max_workers))")
    else
        notice(LOGGER, "Spawning job: $(job.id)")
    end

    return nothing
end


text:
doc"""
    SymTriangularDist(μ,σ)

The *Symmetric triangular distribution* with location `μ` and scale `σ` has probability density function

$f(x; \mu, \sigma) = \frac{1}{\sigma} \left( 1 - \left| \frac{x - \mu}{\sigma} \right| \right), \quad \mu - \sigma \le x \le \mu + \sigma$

```julia
SymTriangularDist()         # Symmetric triangular distribution with zero location and unit scale
SymTriangularDist(u)        # Symmetric triangular distribution with location u and unit scale
SymTriangularDist(u, s)     # Symmetric triangular distribution with location u and scale s

params(d)       # Get the parameters, i.e. (u, s)
location(d)     # Get the location parameter, i.e. u
scale(d)        # Get the scale parameter, i.e. s
```
"""
immutable SymTriangularDist <: ContinuousUnivariateDistribution
    μ::Float64
    σ::Float64

    function SymTriangularDist(μ::Real, σ::Real)
        @check_args(SymTriangularDist, σ > zero(σ))
        new(μ, σ)
    end
    SymTriangularDist(μ::Real) = new(μ, 1.0)
    SymTriangularDist() = new(0.0, 1.0)
end

@distr_support SymTriangularDist d.μ - d.σ d.μ + d.σ

#### Parameters

location(d::SymTriangularDist) = d.μ
scale(d::SymTriangularDist) = d.σ

params(d::SymTriangularDist) = (d.μ, d.σ)

#### Statistics

mean(d::SymTriangularDist) = d.μ
median(d::SymTriangularDist) = d.μ
mode(d::SymTriangularDist) = d.μ

var(d::SymTriangularDist) = d.σ^2 / 6.0
skewness(d::SymTriangularDist) = 0.0
kurtosis(d::SymTriangularDist) = -0.6

entropy(d::SymTriangularDist) = 0.5 + log(d.σ)

#### Evaluation

zval(d::SymTriangularDist, x::Float64) = (x - d.μ) / d.σ
xval(d::SymTriangularDist, z::Float64) = d.μ + z * d.σ

pdf(d::SymTriangularDist, x::Float64) = insupport(d, x) ? (1.0 - abs(zval(d, x))) / scale(d) : 0.0

logpdf(d::SymTriangularDist, x::Float64) = insupport(d, x) ? log((1.0 - abs(zval(d, x))) / scale(d)) : -Inf

function cdf(d::SymTriangularDist, x::Float64)
    (μ, σ) = params(d)
    x <= μ - σ ? 0.0 :
    x <= μ ? 0.5 * (1.0 + zval(d, x))^2 :
    x < μ + σ ? 1.0 - 0.5 * (1.0 - zval(d, x))^2 : 1.0
end

function ccdf(d::SymTriangularDist, x::Float64)
    (μ, σ) = params(d)
    x <= μ - σ ? 1.0 :
    x <= μ ? 1.0 - 0.5 * (1.0 + zval(d, x))^2 :
    x < μ + σ ? 0.5 * (1.0 - zval(d, x))^2 : 0.0
end

function logcdf(d::SymTriangularDist, x::Float64)
    (μ, σ) = params(d)
    x <= μ - σ ? -Inf :
    x <= μ ? loghalf + 2.0 * log1p(zval(d, x)) :
    x < μ + σ ? log1p(-0.5 * (1.0 - zval(d, x))^2) : 0.0
end

function logccdf(d::SymTriangularDist, x::Float64)
    (μ, σ) = params(d)
    x <= μ - σ ? 0.0 :
    x <= μ ? log1p(-0.5 * (1.0 + zval(d, x))^2) :
    x < μ + σ ? loghalf + 2.0 * log1p(-zval(d, x)) : -Inf
end

quantile(d::SymTriangularDist, p::Float64) = p < 0.5 ? xval(d, sqrt(2.0 * p) - 1.0) :
                                                       xval(d, 1.0 - sqrt(2.0 * (1.0 - p)))

cquantile(d::SymTriangularDist, p::Float64) = p > 0.5 ? xval(d, sqrt(2.0 * (1.0-p)) - 1.0) :
                                                        xval(d, 1.0 - sqrt(2.0 * p))

invlogcdf(d::SymTriangularDist, lp::Float64) = lp < loghalf ? xval(d, expm1(0.5*(lp - loghalf))) :
                                                              xval(d, 1.0 - sqrt(-2.0 * expm1(lp)))

invlogccdf(d::SymTriangularDist, lp::Float64) = lp > loghalf ? xval(d, sqrt(-2.0 * expm1(lp)) - 1.0) :
                                                               xval(d, -(expm1(0.5 * (lp - loghalf))))

function mgf(d::SymTriangularDist, t::Float64)
    (μ, σ) = params(d)
    a = σ * t
    a == zero(a) && return one(a)
    4.0 * exp(μ * t) * (sinh(0.5 * a) / a)^2
end

function cf(d::SymTriangularDist, t::Float64)
    (μ, σ) = params(d)
    a = σ * t
    a == zero(a) && return complex(one(a))
    4.0 * cis(μ * t) * (sin(0.5 * a) / a)^2
end

#### Sampling

rand(d::SymTriangularDist) = xval(d, rand() - rand())


text:
function update_fgen_conf!(fgen_conf, fgen::GFG3015, refresh_cnt, base)
    (refresh_cnt==base*1) && (fgen_conf.crt_func = get_wfm(fgen, fgen_conf.crt_chan))

	# this is the most frequent source of timeouts
	(refresh_cnt==base*3) && (fgen_conf.amplit_unit = get_amplit_unit(fgen, fgen_conf.crt_chan))

	(refresh_cnt==base*6) && (fgen_conf.freq = string(get_freq(fgen, fgen_conf.crt_chan)))
	(refresh_cnt==base*9) && (fgen_conf.amplit = string(get_amplit(fgen, fgen_conf.crt_chan)))
	(refresh_cnt==base*12) && (fgen_conf.offs = string(get_offs(fgen, fgen_conf.crt_chan)))

    (refresh_cnt==base*15) && (fgen_conf.duty = string(get_duty(fgen, fgen_conf.crt_chan)))
    return nothing
end

function toggle_func!(fgen::GFG3015, fgen_conf)
    idx_sel = parse(Int64, fgen_conf.idx_sel)
    idx_sel += 1
    idx_sel > 3 ? idx_sel=1 : idx_sel
    crt_fct = fgen.rev_wfm_dict["$idx_sel"]
    set_wfm(fgen, fgen_conf.crt_chan, crt_fct)
    fgen_conf.idx_sel = string(idx_sel)
    return nothing
end

max_stars_repo_name: iuliancioarca/TIVM

max_issues_repo_name: iuliancioarca/TIVM

max_forks_repo_name: iuliancioarca/TIVM


text:
# Test basic visualization of a single pose
model = RobotZoo.Quadrotor()
vis = Visualizer()
if !haskey(ENV, "CI")
    open(vis)
    wait(vis)
end
TrajOptPlots.set_mesh!(vis, model)
x = RBState([1,1,1], expm(SA[1,0,0]*deg2rad(30)), zeros(3), zeros(3))
x0 = zero(RBState)
visualize!(vis, x)
visualize!(vis, x0)
visualize!(vis, model, SVector(x))

# Visualize a trajectory
delete!(vis)
prob = DubinsCar(:escape)
model = prob.model
TrajOptPlots.set_mesh!(vis, model)
N = prob.N
times = TrajectoryOptimization.get_times(prob)
X = [@SVector [sin(t), cos(t), -t] for t in times]
U = [@SVector zeros(2) for k = 1:N-1]
Z = Traj(X, U, push!(diff(times),0))

visualize!(vis, model, Z)
visualize!(vis, model, times[end]/4, states(Z))

initial_states!(prob, X)
initial_controls!(prob, U)
visualize!(vis, prob)

traj3!(vis, prob)
traj3!(vis, prob, inds=1:3)
traj3!(vis, prob, inds=SA[1,2,3])
traj3!(vis, prob, inds=[1,2,3])
traj3!(vis, prob, linewidth=4)
traj3!(vis, prob, linewidth=4, color=colorant"red")

# Vis 2 at the same time
X2 = [@SVector [2sin(t), 2cos(t), -t] for t in times]
Z2 = Traj(X2, U, push!(diff(times),0))
visualize!(vis, model, times[end], X2)

visualize!(vis, model, Z, Z2)
visualize!(vis, model, times[end], X, X2)
prob2 = Problem(prob.model, prob.obj, prob.xf, prob.tf, X0=X2, U0=U)
visualize!(vis, prob, prob2)
delete!(vis["robot_copies"])

# Waypoints
@test_throws ArgumentError TrajOptPlots.waypoints!(vis, model, Z)
TrajOptPlots.waypoints!(vis, model, Z, length=5)
TrajOptPlots.waypoints!(vis, model, Z, length=10)
TrajOptPlots.waypoints!(vis, model, Z, inds=[1,51,75,101])
TrajOptPlots.waypoints!(vis, model, Z, length=11, color=colorant"red", color_end=colorant"green")
delete!(vis["waypoints"])
TrajOptPlots.clear_waypoints!(vis)

# Plot the cylinders
TrajOptPlots.add_cylinders!(vis, prob, height=1)
TrajOptPlots.add_cylinders!(vis, prob, height=0.25, color=colorant"forestgreen")

# Add a point
TrajOptPlots.add_point!(vis, [1,1,1], radius=1)
TrajOptPlots.add_point!(vis, [1,1,1], radius=1, color=colorant"purple")
TrajOptPlots.add_point!(vis, [3,1,1], radius=0.5, color=colorant"cyan", name="point2")


text:
#=
exp_xform.jl

see also MIRT/time/exp_xform.jl

2019-12-08 Connor Martin
Based on exp_xform_mex.c
Copyright 2004-9-23, Jeff Fessler, University of Michigan
=#

export exp_xform

"""
    exp_xform(x, u, v ; mode::Symbol = :matrix)

in:
* `x [N L]` possibly complex vector(s)
* `u [D N]` possibly complex vectors
* `v [D M]` possibly complex vectors
* `mode::Symbol` `:matrix` (default) | `:element` | `:row` | `:column`

out:
* `y [M L]` typically complex vector
 `y[m,l] = sum_n x[n,l] exp(-sum(u[:,n] .* v[:,m]))`

Iterates through subsets of the ML matrix designated by `:mode`
(i.e. row, column, element, or just computing the entire matrix)
This is the 'slow' 'exact' transform model for MRI.

Output type will depend on input types.
"""
function exp_xform(x::AbstractMatrix{<:Number},
        u::AbstractMatrix{<:Number},
        v::AbstractMatrix{<:Number}
        ; mode::Symbol = :matrix)

    T = promote_type(eltype(u), eltype(v), eltype(x), ComplexF32)

    N = size(u,2)
    M = size(v,2)
    out = zeros(T, M, size(x,2))

    if (mode === :matrix)
        return exp.(-(transpose(v) * u)) * x # [M D] * [D N] * [N L]

    elseif (mode === :element)
        # avoid generating intermediate steps by iterating thru N and M
        # dot the columns
        for n = 1:N
            for m = 1:M
                t = transpose(u[:,n]) * v[:,m] #calculate one spot in the array
                out[m,:] .+= (exp.(-t) .* x[n,:])
            end
        end
        return out

    elseif (mode === :row) # iterate through rows of the large matrix
        for m = 1:M
            t = transpose(v[:,m]) * u # [1 D] * [D N] -> [1 N]
            # m == 1 && @show size(t), size(exp.(-t) * x)
            # product is a 2d matrix (row), but should be a column vector:
            out[m,:] .+= (exp.(-t) * x)[1,:]
        end
        return out

    elseif (mode === :column)
        for n = 1:N
            t = transpose(v) * u[:,n] # [M D] * [D .] -> [M]
            # n == 1 && @show size(t), size(x[n,:])
            out .+= (exp.(-t) * transpose(x[n,:])) # outer product [M 1] * [1 L]
        end
        return out

    else
        throw("Invalid mode parameter $mode")
    end
end

# handle "vector" case where L=1
exp_xform(x::AbstractVector{<:Number},
        u::AbstractMatrix{<:Number},
        v::AbstractMatrix{<:Number}
        ; kwargs...) = exp_xform(reshape(x, :, 1), u, v ; kwargs...)[:,1]


text:
using TimeData
module WorldBankDataTd

using HTTPClient.HTTPC
using JSON
using DataArrays
using DataFrames
using Dates
using TimeData

export #
loadWBMeta,
getWBMeta,
search_wdi,
wdi

include("download.jl")
include("ctrs_indicators.jl")
include("search.jl")
include("utils.jl")

end


text:
################################################################################
#
#	CONCRETE TYPE
#
#   Unitcell{D,N,L,S,B} <: AbstractUnitcell{S,B}
#   --> S is the site type (<: AbstractSite{LS,D})
#       --> D is the dimension of embedding space
#       --> LS is the label type of sites
#   --> B is the bond type (<: AbstractBond{LB,N})
#       --> N is the dimension of the Bravais lattice that the bond is located in
#       --> LB is the label type of bonds
#
#   FILE CONTAINS
#       - concrete struct definition
#       - interface implementation
#
################################################################################

################################################################################
#
#   CONCRETE STRUCT DEFINITION
#
################################################################################
mutable struct Unitcell{S,B} <: AbstractUnitcell{S,B}

    # basis vectors of the Bravais lattice
    lattice_vectors	:: Vector{Vector{Float64}}

    # basis sites within the unitcell
    sites			:: Vector{S}

    # list of bonds
    bonds			:: Vector{B}

end

# export the concrete type
export Unitcell

################################################################################
#
#	INTERFACING / ACCESSING UNITCELLS
#	(functions have to be overwritten by concrete types)
#
################################################################################

# default constructor interface
# used for creation of new unitcells
function newUnitcell(
            :: Type{Unitcell{S,B}},
            lattice_vectors :: Vector{<:Vector{<:Real}},
            sites           :: Vector{S},
            bonds           :: Vector{B}
        ) :: Unitcell{S,B} where {D,N,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # return a newly created object
    Unitcell{S,B}(lattice_vectors, sites, bonds)
end

# accessing a list of lattice vectors
function latticeVectors(
            unitcell :: Unitcell{S,B}
        ) :: Vector{Vector{Float64}} where {D,N,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # return the list of lattice vectors
    return unitcell.lattice_vectors
end
# setting a list of lattice vectors
function latticeVectors!(
            unitcell        :: Unitcell{S,B},
            lattice_vectors :: Vector{<:Vector{<:Real}}
        ) where {D,N,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # set the list of lattice vectors
    unitcell.lattice_vectors = lattice_vectors
end

# accessing a list of sites
function sites(
            unitcell :: Unitcell{S,B}
        ) :: Vector{S} where {D,N,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # return the list of sites
    return unitcell.sites
end
# setting a list of sites
function sites!(
            unitcell :: Unitcell{S,B},
            sites    :: Vector{S}
        ) where {D,N,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # set the list of sites
    unitcell.sites = sites
end

# accessing a list of bonds
function bonds(
            unitcell :: Unitcell{S,B}
        ) :: Vector{B} where {D,N,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # return the list of bonds
    return unitcell.bonds
end
# setting a list of bonds
function bonds!(
            unitcell :: Unitcell{S,B},
            bonds    :: Vector{B}
        ) where {D,N,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N}}

    # set the list of bonds
    unitcell.bonds = bonds
end

function saveUnitcell(
        uc :: U,
        fn :: AbstractString,
        group :: AbstractString = "unitcell"
        ;
        append :: Bool = false
    ) where {LS,LB,D,N,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N},U<:Unitcell{S,B}}

    # determine the mode based on if one wants to append stuff
    if append
        mode = "r+"
    else
        mode = "w"
    end

    # determine the site and bond group names
    group_sites = group*"/sites"
    group_bonds = group*"/bonds"

    # open the file in mode
    h5open(fn, mode) do ucfile
        # create the top level group
        group_uc = g_create(ucfile, group)
        # save into the attributes in which groups the sites and bonds are stored
        attrs(group_uc)["sites"] = group_sites
        attrs(group_uc)["bonds"] = group_bonds
        # save types into the attributes
        attrs(group_uc)["site_type"] = string(S)
        attrs(group_uc)["bond_type"] = string(B)
        # number of bravais lattice dimensions
        attrs(group_uc)["N"] = Int64(N)
        # save further attributes

        # write bravais lattice vectors
        for i in 1:N
            group_uc["a"*string(i)] = latticeVectors(uc)[i]
        end
    end

    # save sites and bonds into the file
    saveSites(sites(uc), fn, group_sites, append=true)
    saveBonds(bonds(uc), fn, group_bonds, append=true)

    # return nothing
    return nothing
end

function loadUnitcell(
        ::Type{UC},
        fn :: AbstractString,
        group :: AbstractString = "unitcell"
    ) where {SUC<:AbstractSite,BUC<:AbstractBond,UC<:Union{Unitcell{SUC,BUC}, Unitcell}}

    # read the attribute data
    attr_data = h5readattr(fn, group)

    # determine the bravais lattice dimension
    N = attr_data["N"]

    # determine the site and bond group names
    group_sites = attr_data["sites"]
    group_bonds = attr_data["bonds"]

    # load bravais lattice vectors
    lattice_vectors = Vector{Float64}[]
    for i in 1:N
        push!(lattice_vectors, h5read(fn, group*"/a"*string(i)))
    end

    # load the sites
    site_list = loadSites(fn, group_sites)
    # load the bonds
    bond_list = loadBonds(fn, group_bonds)

    # determine site and bond type
    S = typeof(site_list[1])
    B = typeof(bond_list[1])

    # return new unitcell
    return newUnitcell(
        Unitcell{S,B},
        lattice_vectors,
        site_list,
        bond_list
    )
end

# convinience function for standard type
function loadUnitcell(
        fn :: AbstractString,
        group :: AbstractString = "unitcell"
    )

    return loadUnitcell(Unitcell, fn, group)
end


text:
include("common.jl")

using CSV
using StatsBase

df = CSV.read(result_path)

for (algo_name, algo_fn) in ALGORITHMS
    slice = df[df.algorithm .== algo_name, :]
    println("$algo_name")
    println("-" ^ length(algo_name))
    println("  * solve time:    ", mean(slice.solve_time), " ± ", 2*StatsBase.std(slice.solve_time)/sqrt(num_instances), " sec")
    println("  * is_reward:     ", mean(slice.is_reward), " ± ", 2*StatsBase.std(slice.is_reward)/sqrt(num_instances))
    println("  * is_prop_sold:  ", mean(slice.is_proportion_sold), " ± ", 2*StatsBase.std(slice.is_proportion_sold)/sqrt(num_instances))
    println("  * os_reward:     ", mean(slice.os_reward), " ± ", StatsBase.std(slice.os_reward))
    println("  * os_prop_sold:  ", mean(slice.os_proportion_sold), " ± ", 2*StatsBase.std(slice.os_proportion_sold)/sqrt(num_instances))
end

STATS = ["solve_time", "is_reward", "os_reward", "is_proportion_sold", "os_proportion_sold"]
max_algo_name_length = maximum(length(algo_name) for (algo_name,algo_fn) in ALGORITHMS)

for stat in STATS
    println("$stat")
    println("-" ^ length(stat))
    for (algo_name, algo_fn) in ALGORITHMS
        slice = df[df.algorithm .== algo_name, :]
        println(rpad(string(algo_name, ": "), max_algo_name_length+2), mean(slice[!, Symbol(stat)]), " ± ", 2*StatsBase.std(slice[!, Symbol(stat)])/sqrt(num_instances))
    end
    slice = df[df.algorithm .== "clairvoyant", :]
    println(rpad("clairv.:", max_algo_name_length+2), mean(slice[!, Symbol(stat)]), " ± ", 2*StatsBase.std(slice[!, Symbol(stat)])/sqrt(num_instances))
    println()
end

winners = Dict(algo_name => 0 for (algo_name,algo_fn) in ALGORITHMS)
winners["tie"] = 0
for instance in 1:num_instances
    slice = df[df.instance .== instance, :]
    best_os_reward = maximum(slice.os_reward)
    # NOTE: This excludes ties from the tally!
    if size(slice[slice.os_reward .== best_os_reward, :], 1) != 1
        winners["tie"] += 1
        continue
    end
    for (algo_name, algo_fn) in ALGORITHMS
        if maximum(slice[slice.algorithm .== algo_name, :os_reward]) == best_os_reward
            winners[algo_name] += 1
            # break
        end
    end
end
@show winners

gap(lb, ub) = (ub - lb) / abs(ub)

function gap_closed(cl, dc, mip)
    mip_gap = gap(mip, cl)
    dc_gap = gap(dc, cl)
    return 1 - mip_gap / dc_gap
end

is_gaps_closed = Float64[]
for instance in 1:num_instances
    slice = df[df.instance .== instance, :]
    mip_is = maximum(slice[slice.algorithm .== "mip", :].is_reward)
    dc_is = maximum(slice[slice.algorithm .== "dc", :].is_reward)
    cl_is = maximum(slice[slice.algorithm .== "clairvoyant", :].is_reward)
    push!(is_gaps_closed, gap_closed(cl_is, dc_is, mip_is))
end
@show is_gaps_closed .* 100
@show mean(is_gaps_closed) * 100

os_gaps_closed = Float64[]
for instance in 1:num_instances
    slice = df[df.instance .== instance, :]
    mip_os = maximum(slice[slice.algorithm .== "mip", :].os_reward)
    dc_os = maximum(slice[slice.algorithm .== "dc", :].os_reward)
    cl_os = maximum(slice[slice.algorithm .== "clairvoyant", :].os_reward)
    push!(os_gaps_closed, gap_closed(cl_os, dc_os, mip_os))
end
@show os_gaps_closed .* 100
@show mean(os_gaps_closed) * 100

max_stars_repo_name: joehuchette/reserve-price-optimization

max_issues_repo_name: joehuchette/reserve-price-optimization

max_forks_repo_name: joehuchette/reserve-price-optimization


text:
module Moments

export  eᵢ,
        var_kron,
        var_kron_C,
        make_mon_expo,
        make_xxᵀ⨂yyᵀ,
        get_ℝ_block_diag,
        make_xx̄ᵀ⨂yȳᵀ,
        get_γδ_dict,
        get_ℝℂ_coef_expo,
        get_coef_expoᴿ,
        get_ℂ_block_diag

## Utils
"""The standard basis vector eᵢ in dimension n"""
eᵢ(n::Int,i::Int) = [Int(j==i) for j in 1:n]

""" input: A,B (arrays of integer tupples) output: A ⊗ B """
function var_kron(A,B)
   C_temp = [ a + b for  b in B, a ∈ A]
   rcs =  [(i,j) for i in 1:size(A)[1] for j in 1:size(B)[1]]
   return [C_temp[ij[2],kl[2],ij[1],kl[1]] for ij in rcs , kl in rcs ]
end

function var_kron_C(A,B,fun) # still a shitty fix
    C_temp = [isempty(𝐴) || isempty(𝐵) ? [] : fun(𝐴,𝐵)  for 𝐵 in B, 𝐴 in A]
    rcs =  [(i,j) for i in 1:size(A)[1] for j in 1:size(B)[1]]
    return  [C_temp[ij[2],kl[2],ij[1],kl[1]] for ij in rcs , kl in rcs ] # is the form correct ?
end

## Moments
"""[x]≦ₜ or [x]₌ₜ"""
function make_mon_expo(n::Int,t::Int; isle::Bool = true)
    @assert typeof(t) == Int64
    t == 0 ? (return [eᵢ(n,0)]) : 0
    tmp = make_mon_expo(n,t-1;isle=isle)
    M_vec = reshape([m + eᵢ(n,i) for i ∈ 1:n, m ∈ tmp],:,1)
    return unique(isle ? vcat(tmp,M_vec) : M_vec)
end

"""[x]≦ₜ[x]ᵀ≦ₜ or [x]₌ₜ[x]ᵀ₌ₜ"""
function make_mon_expo(n::Int,t::Tuple{Int,Int}; isle::Bool = true)
    M_vec1      = make_mon_expo(n,t[1]; isle=isle)
    M_vec2      = make_mon_expo(n,t[2]; isle=isle)
    M_mat_flat  = [mi+mj for mi in M_vec1, mj in M_vec2]
end

## Complex moments
"""[x,̄x]≦ₜ or [x,̄x]ᵀ₌ₜ"""
make_mon_expo(d::Tuple{Int,Int},t::Int; isle::Bool = true) = make_mon_expo(sum(2 .* d),t; isle=isle)

"""[x,̄x]≦ₜ[x,̄x]*≦ₜ or [x,̄x]₌ₜ[x,̄x]*₌ₜ but is actually [xᵣₑ,xᵢₘ]≦ₜ[xᵣₑ,xᵢₘ]ᵀ≦ₜ or [xᵣₑ,xᵢₘ]₌ₜ[xᵣₑ,xᵢₘ]ᵀ₌ₜ """
make_mon_expo(d::Tuple{Int,Int},t::Tuple{Int,Int}; isle::Bool = true) = make_mon_expo(sum(2 .* d),t; isle=isle)

## Real
"""xxᵀ⊗yyᵀ"""
function make_xxᵀ⨂yyᵀ(d)
    n = sum(d); p_ex = make_mon_expo(n,1;isle = false)
    x, y = p_ex[1:d[1]], p_ex[d[1]+1:n]
    return var_kron(x .+ reshape(x,1,:), y .+ reshape(y,1,:))
end

"""Returns diagonal block of the momoment matrix with partition:
    xᵃ⁺ᶜyᵇ⁺ᵈ where |a+c|,|b+d| ∈ 2N """
function get_ℝ_block_diag(d::Tuple{Int,Int},t::Tuple{Int,Int})
    mom_mat = Moments.make_mon_expo(sum(d),t)
    mom_vec = mom_mat[1,:]
    halfsum(arr) = [sum(arr[1:d[1]]),sum(arr[(d[1]+1):end])]
    tf(arr)      = map(x -> (isodd(x[1]),isodd(x[2])), halfsum.(arr))
    tf2(arr)     = map(p -> p[1], findall(arr))

    isevev(pair) = pair == (false,false)
    isodod(pair) = pair == (true,true)
    isevod(pair) = pair == (false,true)
    isodev(pair) = pair == (true,false)

    ee = tf2(isevev.(tf(mom_vec)))
    oo = tf2(isodod.(tf(mom_vec)))
    eo = tf2(isevod.(tf(mom_vec)))
    oe = tf2(isodev.(tf(mom_vec)))

    ℝ_block_diag = Dict("ee" => mom_mat[ee,ee],
                        "oo" => mom_mat[oo,oo],
                        "eo" => mom_mat[eo,eo],
                        "oe" => mom_mat[oe,oe])

    for key in ["ee","oo","eo","oe"]
         isempty(ℝ_block_diag[key]) ? delete!(ℝ_block_diag,key) : continue
    end
    return ℝ_block_diag
end

## Complex
""" ℜe(xx*⊗yy*), ℑm(xx*⊗yy*)
xx*⊗yy* = (xᵣₑ + i xᵢₘ)(xᵣₑ - i xᵢₘ)ᵀ ⊗ (yᵣₑ + i yᵢₘ)(yᵣₑ - i yᵢₘ)ᵀ
= (xᵣₑxᵣₑᵀ + xᵢₘxᵢₘᵀ + i (xᵢₘxᵣₑᵀ - xᵣₑxᵢₘᵀ)) ⊗ (yᵣₑyᵣₑᵀ + yᵢₘyᵢₘᵀ + i (yᵢₘyᵣₑᵀ - yᵣₑyᵢₘᵀ))
⟹ cℝxx̄ᵀ⨂yȳᵀ ∘ eℝxx̄ᵀ⨂yȳᵀ := ℜe(xx*⊗yy*) = (xᵣₑxᵣₑᵀ + xᵢₘxᵢₘᵀ ) ⊗ (yᵣₑyᵣₑᵀ + yᵢₘyᵢₘᵀ ) - (xᵢₘxᵣₑᵀ - xᵣₑxᵢₘᵀ) ⊗ (yᵢₘyᵣₑᵀ - yᵣₑyᵢₘᵀ)
⟹ c𝕀xx̄ᵀ⨂yȳᵀ ∘ e𝕀xx̄ᵀ⨂yȳᵀ := ℑm(xx*⊗yy*) = (xᵣₑxᵣₑᵀ + xᵢₘxᵢₘᵀ) ⊗ (yᵢₘyᵣₑᵀ - yᵣₑyᵢₘᵀ) + (xᵢₘxᵣₑᵀ - xᵣₑxᵢₘᵀ)⊗(yᵣₑyᵣₑᵀ + yᵢₘyᵢₘᵀ) """

function make_xx̄ᵀ⨂yȳᵀ(d,γδ_dict)
    p_ex = make_mon_expo(sum(2 .*d),1;isle = false)
    x, x̄, y, ȳ = split_expo(p_ex,d...)
    xx̄ᵀ⨂yȳᵀ = var_kron(x .+ reshape(x̄,1,:), y .+ reshape(ȳ,1,:))
end

## Block diagonalization
# The complex blocks
"""Splits a vector of exponents into the register components ααᶥββᶥ →  (α,αᶥ,β,βᶥ)"""
split_expo(ααᶥββᶥ,d₁,d₂) = (ααᶥββᶥ[1:d₁],ααᶥββᶥ[d₁+1:2*d₁],
                            ααᶥββᶥ[1+2*d₁:2*d₁+d₂],ααᶥββᶥ[1+2*d₁+d₂:end])

"""xᵅ̄xᵃyᵝ̄yᵇ -> xᵃx̄ᵅyᵇ̄yᵝ̄"""
function conj_expo(ααᶥββᶥ,d₁,d₂)
    α,αᶥ,β,βᶥ = split_expo(ααᶥββᶥ,d₁,d₂) ; return vcat(αᶥ,α,βᶥ,β)
end

"""Gets the principal block of the moment matrix after block diagonalization"""
function get_xx̄yȳMM_blocks(d,t)
    Iᵗ = make_mon_expo(d,t[1])
    Iᵗ_temp = map(x ->  sum.(split_expo(x,d...)), Iᵗ)
    r_list  = map(v -> v[1]-v[2],Iᵗ_temp) ; s_list  = map(v -> v[3]-v[4],Iᵗ_temp)
    P = Dict()
    for r ∈ -t[1]:t[1], s ∈ -t[1]:t[1]
        T = Iᵗ[(r_list .== r) .& (s_list .== s)]
        isempty(T) ? nothing : P[r,s] = T
    end
    build_block(k) = P[k] .+ reshape(map(x->conj_expo(x,d...),P[k]),1,:)
    Dict(zip(keys(P), [build_block(k) for k in keys(P)]))
end

## The real analogue
"""Returns all  (γ,γᶥ,ζ,ζᶥ),(δ,δᶥ,η,ηᶥ) ∈ (ℕᵈ)⁴ s.t.
(γ,γᶥ,ζ,ζᶥ)+(δ,δᶥ,η,ηᶥ)=(α,αᶥ,β,βᶥ)  ∀  ααᶥββᶥ ∈ Moment matrixₜ
"""
function get_γδ(m)
    # nz_coord = findall(.!(m .== 0)) ; n = length(m)
    # γ = [m] ; δ = [eᵢ(n,0)]
    # for k in nz_coord
    #     n_count = m[k]
    #     while n_count > 0
    #         mtemp = eᵢ(n,k)
    #         push!(γ, γ[end] - mtemp)
    #         push!(δ, δ[end] + mtemp)
    #         n_count -= 1
    #     end
    # end
    # return unique(vcat(γ,δ))
    n = length(m)
    nz_coord = findall(.!(m .== 0))
    isempty(nz_coord) ? (return [m]) : nothing
    c = pop!(nz_coord)
    Nar_old = [ j*eᵢ(n,c) for j in 0:m[c]]
    global Nar_old
    while  !isempty(nz_coord)
        c = pop!(nz_coord)
        Nar_new = [ j*eᵢ(n,c) for j in 0:m[c]]
        Nar_old = [ n + m for n in Nar_old for m in Nar_new ]
    end
    return Nar_old
end

function get_γδ_dict(d,t)
    MM_vec = make_mon_expo(d,t[1]*2)
    return Dict(zip(MM_vec,map(m ->([m] .- get_γδ(m),get_γδ(m)) ,MM_vec)))
end

# function get_γδ_dict(d,t)
#     M = Moments.make_mon_expo(d,2 .* t)
#     M_vec = M[:,1]
#     γδ_dict = Dict()
#     for k in Moments.make_mon_expo(d,2*t[1])
#         fd = hcat(map(x -> [x[1],x[2]],findall(M .== [k]))...)
#         γδ_dict[k] = [M_vec[fd[1,:]], M_vec[fd[2,:]]]
#     end
#     return γδ_dict
# end

"""The coefficients function
-1^|δᶥ+ηᶥ|
i^|δ+δᶥ+η+ηᶥ|
α!αᶥ!β!βᶥ! / γ!γᶥ!ζ!ζᶥ!δ!δᶥ!η!ηᶥ!
"""
f_temp(X) = prod(map(x-> prod(factorial.(x)),X))
function get_coef(p1,p2,d)
    γ,γᶥ,ζ,ζᶥ = Moments.split_expo(p1,d...) ; δ,δᶥ,η,ηᶥ = Moments.split_expo(p2,d...)
    a = ((-1.0)^sum(sum.([δᶥ,ηᶥ])))*((1.0im)^sum(sum.([δ,δᶥ,η,ηᶥ])))*f_temp([[γ,γᶥ,ζ,ζᶥ]+[δ,δᶥ,η,ηᶥ]...])
    return a/f_temp([γ,γᶥ,ζ,ζᶥ,δ,δᶥ,η,ηᶥ])
end
get_coef(p1p2_arr,d) = map((x,y)->get_coef(x,y,d),p1p2_arr...)

""" get_expo(γγᶥδδᶥ,ζζᶥηηᶥ,d) = [γ+γᶥ,δ+δᶥ,η+ηᶥ,ζ+ζᶥ]"""
function get_expo(p1,p2,d)
    γ,γᶥ,ζ,ζᶥ = split_expo(p1,d...) ; δ,δᶥ,η,ηᶥ = split_expo(p2,d...)
    return vcat(γ+γᶥ,δ+δᶥ,ζ+ζᶥ,η+ηᶥ)
end
get_expo(p1p2_arr,d) = map((x,y)->get_expo(x,y,d),p1p2_arr...)

"""f"""
function get_coef_expo(d,ααᶥββᶥ::Array{Int64,1},γδ_dict)
    γδ = γδ_dict[ααᶥββᶥ]
    return get_coef(γδ,d),get_expo(γδ,d)
end

function get_ℝℂ_coef_expo(coef,expo)
    ℝmask = isreal.(coef)
    ℝcoef = real(coef[ℝmask])   ; ℝexpo = expo[ℝmask]
    ℂcoef = imag(coef[.!ℝmask]) ; ℂexpo = expo[.!ℝmask]
    return (ℝcoef,ℝexpo),(ℂcoef,ℂexpo)
end

function get_ℝℂ_coef_expo(d,B,γδ_dict)
    s₁,s₂ = size(B)
    MM = map(x-> get_coef_expo(d,x,γδ_dict),B)
    MMℝℂ = map(x-> get_ℝℂ_coef_expo(x...),MM)

    ℝcoef = [MMℝℂ[i,j][1][1] for i in 1:s₁, j in 1:s₂]
    ℝexpo = [MMℝℂ[i,j][1][2] for i in 1:s₁, j in 1:s₂]
    ℂcoef = [MMℝℂ[i,j][2][1] for i in 1:s₁, j in 1:s₂]
    ℂexpo = [MMℝℂ[i,j][2][2] for i in 1:s₁, j in 1:s₂]
    return (ℝcoef,ℝexpo),(ℂcoef,ℂexpo)
end

"""(ℝcoef,ℝexpo),(ℂcoef,ℂexpo) ->
MMCoefᴿ = (ℝcoef -ℂcoef)        MMexᴿ :=     (ℝexpo ℂexpo)
          (ℂcoef  ℝcoef)    and              (ℂexpo ℝexpo)
 """
function get_coef_expoᴿ(d,B,γδ_dict)
    (ℝcoef,ℝexpo),(ℂcoef,ℂexpo) = get_ℝℂ_coef_expo(d,B,γδ_dict)
    MMCoefᴿ = hcat(vcat(ℝcoef,ℂcoef),vcat((-1.0)*ℂcoef,ℝcoef))
    MMexᴿ   = hcat(vcat(ℝexpo,ℂexpo),vcat(ℂexpo,ℝexpo))
    return MMCoefᴿ,MMexᴿ
end

""""""
function get_ℂ_block_diag(d,t;noBlock = false)
    if noBlock
        MMexᴿ   = Dict("Default" => map(x->[x],make_mon_expo(d,t) .+ reshape(make_mon_expo(d,t),1,:) ))
        MMCoefᴿ = Dict("Default" => fill([1.0],size(MMexᴿ["Default"])...))
        return MMCoefᴿ,MMexᴿ
    end
    xx̄yȳMM_blocks = get_xx̄yȳMM_blocks(d,t)
    γδ_dict = get_γδ_dict(d,t) ## This should be moved outside so that it can be precomputed.

    MMexᴿ= Dict() ; MMCoefᴿ = Dict()
    for B in keys(xx̄yȳMM_blocks)
        MMCoefᴿ[B],MMexᴿ[B] = get_coef_expoᴿ(d,xx̄yȳMM_blocks[B],γδ_dict)
    end
    return MMCoefᴿ,MMexᴿ
end

end

# U(vec) = vec .+ reshape(vec,1,:)
# W(vec1,vec2) = vec1 .+ reshape(vec2,1,:)
# function make_xx̄ᵀ⨂yȳᵀ(d)
#     d₁,d₂ = d
#     n = sum(2 .* d)
#     pre_expo = make_mon_expo(n,1;isle =false)
#     xᵣₑ, xᵢₘ = pre_expo[1:d₁], pre_expo[d₁+1:2*d₁]
#     yᵣₑ, yᵢₘ = pre_expo[2*d₁+1:2*d₁+d₂], pre_expo[2*d₁+d₂+1:2*d₁+2*d₂]
#
#     xᵣₑxᵣₑᵀ, xᵢₘxᵢₘᵀ, yᵣₑyᵣₑᵀ, yᵢₘyᵢₘᵀ = U(xᵣₑ), U(xᵢₘ), U(yᵣₑ), U(yᵢₘ)
#     xᵢₘxᵣₑᵀ, xᵣₑxᵢₘᵀ, yᵢₘyᵣₑᵀ, yᵣₑyᵢₘᵀ = W(xᵢₘ,xᵣₑ), W(xᵣₑ,xᵢₘ), W(yᵢₘ,yᵣₑ), W(yᵣₑ,yᵢₘ)
#
#     Real_dict =  Dict(("+ 1")  => var_kron(xᵣₑxᵣₑᵀ,yᵣₑyᵣₑᵀ),
#                       ("+ 2")  => var_kron(xᵣₑxᵣₑᵀ,yᵢₘyᵢₘᵀ),
#                       ("+ 3")  => var_kron(xᵢₘxᵢₘᵀ,yᵣₑyᵣₑᵀ),
#                       ("+ 4")  => var_kron(xᵢₘxᵢₘᵀ,yᵢₘyᵢₘᵀ),
#                       ("- 5")  => var_kron(xᵢₘxᵣₑᵀ,yᵢₘyᵣₑᵀ),
#                       ("+ 6")  => var_kron(xᵢₘxᵣₑᵀ,yᵣₑyᵢₘᵀ),
#                       ("+ 7")  => var_kron(xᵣₑxᵢₘᵀ,yᵢₘyᵣₑᵀ),
#                       ("- 8")  => var_kron(xᵣₑxᵢₘᵀ,yᵣₑyᵢₘᵀ))
#
#     Imag_dict =  Dict(("+ 1")  => var_kron(xᵣₑxᵣₑᵀ,yᵢₘyᵣₑᵀ),
#                       ("- 2") => var_kron(xᵣₑxᵣₑᵀ,yᵣₑyᵢₘᵀ),
#                       ("+ 3")  => var_kron(xᵢₘxᵢₘᵀ,yᵢₘyᵣₑᵀ),
#                       ("- 4") => var_kron(xᵢₘxᵢₘᵀ,yᵣₑyᵢₘᵀ),
#                       ("+ 5")  => var_kron(xᵢₘxᵣₑᵀ,yᵣₑyᵣₑᵀ),
#                       ("+ 6")  => var_kron(xᵢₘxᵣₑᵀ,yᵢₘyᵢₘᵀ),
#                       ("- 7") => var_kron(xᵣₑxᵢₘᵀ,yᵣₑyᵣₑᵀ),
#                       ("- 8") => var_kron(xᵣₑxᵢₘᵀ,yᵢₘyᵢₘᵀ))
#
#     return Dict("real" => Real_dict,
#                 "imag" => Imag_dict)
#
#
# end

max_stars_repo_name: JAndriesJ/sep-rank

max_issues_repo_name: JAndriesJ/sep-rank

max_forks_repo_name: JAndriesJ/sep-rank


text:
using Dates

const SHOW_DETAILED_ERROR = false

"""
`DataConsumer` is the abstract type for all the concrete structs that have an implementation
for digesting the data parsed by `LazyCSV`
"""
abstract type DataConsumer end

"""
    consume(consumer::DataConsumer, f::File)

This function is only called once per `LazyCSV.File` for a given `DataConsumer` instance.

The purpose of this functions is to iterate over the whole file.

First, the header is consumed (via `consume_header` based on consumer's knowledge of whether
       it exists of not) and the the corresponding header is produced inside the consumer if
       necessary (via `produce_header`).
Then, each line of the body is traversed and consumed separately (via `consume_rec`).
      And then inside each record, each field is consumed (via `consume_field`).
      If any error is found in processing the records or field, the corresponding error
      handling function (either `consume_rec_error` or `consume_field_error`
      correspondingly) is called.
Finally, after all records are processed, the consumption is finalized (via a call to
      `consume_finalize`)

If this sequence of operations is adequate for a consumer, it's suggested to overload the
    inner functions (i.e., consume_header, produce_header, consume_rec, consume_field,
                           consume_rec_error, consume_field_error, consume_finalize).
    Otherwise, a consumer can overload the `consume` method directly and change its behavior.
"""
function consume(consumer::DataConsumer, f::File)
    iter_res = iterate(f)
    (iter_res, new_consumer) = consume_header(consumer, f, iter_res)
    while iter_res !== nothing
        ((pos, line), state) = iter_res
        consume_rec(new_consumer, f, pos, line, first_rec_fields(f))
        iter_res = iterate(f, state)
    end
    consume_finalize(new_consumer)
end
function consume_header end
function produce_header end
function consume_rec end
function consume_field end
function consume_rec_error end
function consume_field_error end
function consume_finalize end

function consume_header(fn::Function, consumer::DataConsumer, file::File, iter_res)
    new_iter_res = if iter_res !== nothing && file.header_exists
        ((pos, line), state) = iter_res
        i = 1
        for field_header in first_rec_fields(file)
            fn(field_header, i)
            i += 1
        end
        iterate(file, state)
    else
        iter_res
    end
    updated_consumer = produce_header(consumer, file)
    (new_iter_res, updated_consumer)
end

"""
Creates a key object for each row for a given `FilePos` (which is the byte sequence number
of the first character in the line of that record in the CSV source and it's a unique ID for
each row)
"""
create_key(pc::DataConsumer, pos::FilePos) = pos

function consume_rec(pc::DataConsumer, f::File, pos::FilePos, line, fields)
    key = create_key(pc, pos)
    if !validate_record(pc.rec_type, length(fields))
        consume_rec_error(pc, f, key, line)
    else
        for (i, field) in enumerate(fields)
            if !consume_field(pc, f, key, line, field, i)
                consume_field_error(pc, f, key, line, field, i)
            end
        end
    end
end

"""
This function is applied on the consumer, when processing a CSV file is done
"""
function consume_finalize(consumer::DataConsumer) end

"""
An abstract type for specifying the record length.

Its concrete extensions are `UnknownRecordLength` and `FixedRecordLength`.

The only method that has to be overloaded for it is `Base.length`.
"""
abstract type RecordLength end
struct UnknownRecordLength <: RecordLength end
struct FixedRecordLength <: RecordLength
    len::Int
end
Base.length(::UnknownRecordLength) = -1
Base.length(rec_len::FixedRecordLength) = rec_len.len

# Currently, only the first byte in LSB is used for storing the type information
# The second byte is currently reserved (to be used by user libraries for adding
# a few more user defined types).
# The third and fourth bytes can be used for storing additional type information
# if required.
const FieldType = UInt64

# The list of all supported types in CSV data loading
const NO_TYPE       = 0b000000000 % UInt64
const MISSING_TYPE  = 0b100000000 % UInt64
const INT_TYPE      = 0b000000001 % UInt64
const FLOAT_TYPE    = 0b000000010 % UInt64
const DATE_TYPE     = 0b000000100 % UInt64
const DATETIME_TYPE = 0b000001000 % UInt64
const BOOL_TYPE     = 0b000010000 % UInt64
const CHAR_TYPE     = 0b000100000 % UInt64
const STRING_TYPE   = 0b001000000 % UInt64

typecode(::Type{Int64}) = INT_TYPE
typecode(::Type{Float64}) = FLOAT_TYPE
typecode(::Type{Date}) = DATE_TYPE
typecode(::Type{DateTime}) = DATETIME_TYPE
typecode(::Type{Bool}) = BOOL_TYPE
typecode(::Type{String}) = STRING_TYPE
typecode(::Type{Char}) = CHAR_TYPE
typecode(::Type{Tuple{Ptr{UInt8}, Int}}) = STRING_TYPE
typecode(::Type{Missing}) = MISSING_TYPE
typecode(::Type) = NO_TYPE

"""
An abstract type for specifying the type of a record in a CSV file.

Two methods have to be defined for each subtype of it:
    - record_length: length of the record, which should be of type `RecordLength`
    - record_fields: a vector of `FieldType`, one for each column
"""
abstract type RecordType end
function record_length end
function record_fields end

function validate_record(rec_type::RecordType, len::UInt)
    rec_len = record_length(rec_type)
    rec_len == FixedRecordLength(len) || rec_len == UnknownRecordLength()
end

struct UntypedRecord <: RecordType
    num_fields::Int
end
function record_length(rec_type::UntypedRecord)
    rec_type.num_fields <= 0 ? UnknownRecordLength() : FixedRecordLength(rec_type.num_fields)
end
function record_fields(rec_type::UntypedRecord)
    field_types = Vector{FieldType}()
    for i in 1:rec_type.num_fields
        push!(field_types, NO_TYPE)
    end
    field_types
end
function typecodes(tp::Type{Union{A,B}}) where {A,B}
    if tp isa DataType
        typecode(tp)
    elseif tp.b isa DataType
        typecode(tp.b) | typecode(tp.a)
    else
        typecodes(tp.b) | typecode(tp.a)
    end
end

function get_field_union_types(::Type{T}) where {T <: Tuple}
    res = Vector{FieldType}()
    for field_tp in T.parameters
        field_types = typecodes(field_tp)
        push!(res, field_types)
    end
    res
end

struct FieldHandler{HANDLER <: Function}
    fn::HANDLER
end

struct TypedRecord{FIELD_HANDLERS <: Tuple} <: RecordType
    field_types::Vector{FieldType}
    fallback_type::FieldType
    other_type_handlers::FIELD_HANDLERS
    function TypedRecord(field_types::Vector{FieldType}=Vector{FieldType}(),
                         fallback_type::FieldType=NO_TYPE,
                         other_type_handlers::FIELD_HANDLERS=Tuple{}()) where {FIELD_HANDLERS}
        new{FIELD_HANDLERS}(field_types, fallback_type, other_type_handlers)
    end
end

function TypedRecord(field_types::Type{T},
                     fallback_type::FieldType=NO_TYPE,
                     other_type_handlers::FIELD_HANDLERS=Tuple{}()) where {FIELD_HANDLERS, T <: Tuple}
    TypedRecord(get_field_union_types(field_types), fallback_type, other_type_handlers)
end

function record_length(rec_type::TypedRecord)
    length(rec_type.field_types) == 0 ? UnknownRecordLength() : FixedRecordLength(length(rec_type.field_types))
end
function record_fields(rec_type::TypedRecord)
    rec_type.field_types
end

function field_type(rec_type::TypedRecord, i::Int)
    if i <= length(rec_type.field_types)
        rec_type.field_types[i]
    elseif i <= length(rec_type.other_type_handlers)
        NO_TYPE
    else
        rec_type.fallback_type
    end
end

ismissing_value(field_value::AbstractString) = isempty(field_value)

function has_valid_type(tp::FieldType, field_value::AbstractString)
    tp == NO_TYPE && return (false, NO_TYPE)
    tp & MISSING_TYPE  != 0 && ismissing_value(field_value) && return (true, MISSING_TYPE)
        if tp & INT_TYPE      != 0 && tryparse(Int, field_value) !== nothing; return (true, INT_TYPE)
    elseif tp & FLOAT_TYPE    != 0 && tryparse(Float64, field_value) !== nothing; return (true, FLOAT_TYPE)
    elseif tp & DATE_TYPE     != 0 && tryparse(Date, field_value) !== nothing; return (true, DATE_TYPE)
    elseif tp & DATETIME_TYPE != 0 && tryparse(DateTime, field_value) !== nothing; return (true, DATETIME_TYPE)
    elseif tp & BOOL_TYPE     != 0 && tryparse(Bool, field_value) !== nothing; return (true, BOOL_TYPE)
    elseif tp & CHAR_TYPE     != 0 && length(field_value) == 1; return (true, CHAR_TYPE)
    elseif tp & STRING_TYPE   != 0; return (true, STRING_TYPE)
    end
    return (false, NO_TYPE)
end

include("print_consumer.jl")

export DataConsumer, consume, consume_header, consume_rec, consume_field, consume_rec_error,
       consume_field_error
export RecordLength, UnknownRecordLength, FixedRecordLength
export FieldType, RecordType, UntypedRecord, record_length, csv_string, record_fields
export MISSING_TYPE, INT_TYPE, FLOAT_TYPE, DATE_TYPE, DATETIME_TYPE
export BOOL_TYPE, CHAR_TYPE, STRING_TYPE


text:
println("\n\n\nPerformance of FinanceLib.FixedIncomes functions")

using Dates: Date
using FinanceLib: yearFrac

import FinanceLib as Fl

import FinanceLib.FixedIncomes as FI
import FinanceLib.FixedIncomes.MoneyMarkets as MM

println("\n\n\nPerformance of FinanceLib.FixedIncomes functions")

# @code_warntype MM.tBillR(150,98000,100000)
# @code_warntype MM.tBillD(0.048, 150, 100_000)

# @code_warntype MM.holdingPerYield(98, 95, 5)
# @code_warntype MM.effAnnYield(150, 98, 95, 5)
# @code_warntype MM.moneyMktYield(150, 98, 95, 5)

# @code_warntype MM.twrr([4,6,5.775,6.72,5.508],[1,-0.5,0.225,-0.6])
# @code_warntype MM.twrr(1.0, [100, 112, 142.64], [0, 20.])

# println("Benchmarks\n")

include("Bonds/perf.jl")


text:
"""
    checkdata!(args...)

Check `data` is `Tables.AbstractColumns`-compatible
and find valid rows for options `subset` and `weightname`.
See also [`CheckData`](@ref).
"""
function checkdata!(data, subset::Union{BitVector, Nothing}, weightname::Union{Symbol, Nothing})
    checktable(data)
    nrow = Tables.rowcount(data)
    if subset !== nothing
        length(subset) == nrow || throw(DimensionMismatch(
            "data contain $(nrow) rows while subset has $(length(subset)) elements"))
        # Do not modify subset in-place
        esample = copy(subset)
    else
        esample = trues(nrow)
    end

    # A cache that makes updating BitVector (esample or tr_rows) faster
    # See https://github.com/JuliaData/DataFrames.jl/pull/2726
    aux = BitVector(undef, nrow)

    if weightname !== nothing
        colweights = getcolumn(data, weightname)
        if Missing <: eltype(colweights)
            aux .= .!ismissing.(colweights)
            esample .&= aux
        end
        aux[esample] .= view(colweights, esample) .> 0
        esample[esample] .&= view(aux, esample)
    end
    sum(esample) == 0 && error("no nonmissing data")
    return (esample=esample, aux=aux)
end

"""
    CheckData <: StatsStep

Call [`DiffinDiffsBase.checkdata!`](@ref)
for some preliminary checks of the input data.
"""
const CheckData = StatsStep{:CheckData, typeof(checkdata!), true}

required(::CheckData) = (:data,)
default(::CheckData) = (subset=nothing, weightname=nothing)

"""
    grouptreatintterms(treatintterms)

Return the argument without change for allowing later comparisons based on object-id.
See also [`GroupTreatintterms`](@ref).
"""
grouptreatintterms(treatintterms::TermSet) = (treatintterms=treatintterms,)

"""
    GroupTreatintterms <: StatsStep

Call [`DiffinDiffsBase.grouptreatintterms`](@ref)
to obtain one of the instances of `treatintterms`
that have been grouped by equality (`hash`)
for allowing later comparisons based on object-id.

This step is only useful when working with [`@specset`](@ref) and [`proceed`](@ref).
"""
const GroupTreatintterms = StatsStep{:GroupTreatintterms, typeof(grouptreatintterms), false}

default(::GroupTreatintterms) = (treatintterms=TermSet(),)

"""
    groupxterms(xterms)

Return the argument without change for allowing later comparisons based on object-id.
See also [`GroupXterms`](@ref).
"""
groupxterms(xterms::TermSet) = (xterms=xterms,)

"""
    GroupXterms <: StatsStep

Call [`DiffinDiffsBase.groupxterms`](@ref)
to obtain one of the instances of `xterms`
that have been grouped by equality (`hash`)
for allowing later comparisons based on object-id.

This step is only useful when working with [`@specset`](@ref) and [`proceed`](@ref).
"""
const GroupXterms = StatsStep{:GroupXterms, typeof(groupxterms), false}

default(::GroupXterms) = (xterms=TermSet(),)

"""
    groupcontrasts(contrasts)

Return the argument without change for allowing later comparisons based on object-id.
See also [`GroupContrasts`](@ref).
"""
groupcontrasts(contrasts::Union{Dict{Symbol,Any},Nothing}) = (contrasts=contrasts,)

"""
    GroupContrasts <: StatsStep

Call [`DiffinDiffsBase.groupcontrasts`](@ref)
to obtain one of the instances of `contrasts`
that have been grouped by equality (`hash`)
for allowing later comparisons based on object-id.

This step is only useful when working with [`@specset`](@ref) and [`proceed`](@ref).
"""
const GroupContrasts = StatsStep{:GroupContrasts, typeof(groupcontrasts), false}

default(::GroupContrasts) = (contrasts=nothing,)

function _checkscales(col1::AbstractArray, col2::AbstractArray, treatvars::Vector{Symbol})
    if col1 isa ScaledArrOrSub || col2 isa ScaledArrOrSub
        col1 isa ScaledArrOrSub && col2 isa ScaledArrOrSub ||
            throw(ArgumentError("time fields in both columns $(treatvars[1]) and $(treatvars[2]) must be ScaledArrays; see settime and aligntime"))
        first(DataAPI.refpool(col1)) == first(DataAPI.refpool(col2)) &&
            scale(col1) == scale(col2) || throw(ArgumentError(
            "time fields in columns $(treatvars[1]) and $(treatvars[2]) are not aligned; see aligntime"))
    else
        eltype(col1) <: Integer || throw(ArgumentError(
            "columns $(treatvars[1]) and $(treatvars[2]) must be stored in ScaledArrays; see settime and aligntime"))
    end
end

function checktreatvars(::DynamicTreatment{SharpDesign},
        pr::TrendOrUnspecifiedPR{Unconditional}, treatvars::Vector{Symbol}, data)
    # treatvars should be cohort and time variables
    col1 = getcolumn(data, treatvars[1])
    col2 = getcolumn(data, treatvars[2])
    T1 = nonmissingtype(eltype(col1))
    T2 = nonmissingtype(eltype(col2))
    T1 == T2 || throw(ArgumentError(
        "nonmissing elements from columns $(treatvars[1]) and $(treatvars[2]) have different types $T1 and $T2"))
    T1 <: ValidTimeType ||
        throw(ArgumentError("column $(treatvars[1]) has unaccepted element type $(T1)"))
    if pr isa TrendParallel
        eltype(pr.e) == T1 || throw(ArgumentError("element type $(eltype(pr.e)) of control cohorts from $pr does not match element type $T1 from data; expect $T1"))
    end
    if T1 <: RotatingTimeValue
        col1 isa RotatingTimeArray && col2 isa RotatingTimeArray ||
            throw(ArgumentError("columns $(treatvars[1]) and $(treatvars[2]) must be RotatingTimeArrays; see settime"))
        _checkscales(col1.time, col2.time, treatvars)
    else
        _checkscales(col1, col2, treatvars)
    end
end

function _overlaptime(tr::DynamicTreatment, esample::BitVector, tr_rows::BitVector, data)
    timeref = refarray(getcolumn(data, tr.time))
    control_time = Set(view(timeref, .!tr_rows.&esample))
    treated_time = Set(view(timeref, tr_rows.&esample))
    return intersect(control_time, treated_time), control_time, treated_time
end

function overlap!(esample::BitVector, tr_rows::BitVector, aux::BitVector, tr::DynamicTreatment,
        ::NeverTreatedParallel{Unconditional}, treatname::Symbol, data)
    overlap_time, control_time, treated_time = _overlaptime(tr, esample, tr_rows, data)
    if !(length(control_time)==length(treated_time)==length(overlap_time))
        aux[esample] .= view(refarray(getcolumn(data, tr.time)), esample) .∈ (overlap_time,)
        esample[esample] .&= view(aux, esample)
    end
    tr_rows .&= esample
end

function overlap!(esample::BitVector, tr_rows::BitVector, aux::BitVector, tr::DynamicTreatment,
        pr::NotYetTreatedParallel{Unconditional}, treatname::Symbol, data)
    timecol = getcolumn(data, tr.time)
    # First exclude cohorts not suitable for comparisons
    if !(eltype(timecol) <: RotatingTimeValue)
        invpool = invrefpool(timecol)
        e = invpool === nothing ? Set(pr.e) : Set(invpool[c] for c in pr.e)
        ecut = invpool === nothing ? pr.ecut[1] : invpool[pr.ecut[1]]
        isvalidcohort = x -> x < ecut || x in e
    else
        invpool = invrefpool(timecol.time)
        if invpool === nothing
            e = Set(pr.e)
            ecut = IdDict(e.rotation=>e.time for e in pr.ecut)
        else
            e = Set(RotatingTimeValue(c.rotation, invpool[c.time]) for c in pr.e)
            ecut = IdDict(e.rotation=>invpool[e.time] for e in pr.ecut)
        end
        isvalidcohort = x -> x.time < ecut[x.rotation] || x in e
    end
    aux[esample] .= isvalidcohort.(view(refarray(getcolumn(data, treatname)), esample))
    esample[esample] .&= view(aux, esample)
    # Check overlaps among the remaining cohorts only
    overlap_time, _c, _t = _overlaptime(tr, esample, tr_rows, data)
    if !(eltype(timecol) <: RotatingTimeValue)
        filter!(x -> x < ecut, overlap_time)
    else
        filter!(x -> x.time < ecut[x.rotation], overlap_time)
    end
    aux[esample] .= view(refarray(timecol), esample) .∈ (overlap_time,)
    esample[esample] .&= view(aux, esample)
    tr_rows .&= esample
end

overlap!(esample::BitVector, tr_rows::BitVector, aux::BitVector, tr::DynamicTreatment,
    pr::UnspecifiedParallel{Unconditional}, treatname::Symbol, data) = nothing

"""
    checkvars!(args...)

Exclude rows with missing data or violate the overlap condition
and find rows with data from treated units.
See also [`CheckVars`](@ref).
"""
function checkvars!(data, pr::AbstractParallel,
        yterm::AbstractTerm, treatname::Symbol, esample::BitVector, aux::BitVector,
        treatintterms::TermSet, xterms::TermSet, ::Type, @nospecialize(trvars::Tuple),
        tr::AbstractTreatment)
    # Do not check eltype of treatintterms
    treatvars = union([treatname], trvars, termvars(pr))
    checktreatvars(tr, pr, treatvars, data)

    allvars = union(treatvars, termvars(yterm), termvars(xterms))
    for v in allvars
        col = getcolumn(data, v)
        if Missing <: eltype(col)
            aux .= .!ismissing.(col)
            esample .&= aux
        end
    end
    # Values of treatintterms from untreated units are ignored
    tr_rows = copy(esample)
    if !(pr isa UnspecifiedParallel)
        istreated!(view(aux, esample), pr, view(getcolumn(data, treatname), esample))
        tr_rows[esample] .&= view(aux, esample)
    end
    treatintvars = termvars(treatintterms)
    for v in treatintvars
        col = getcolumn(data, v)
        if Missing <: eltype(col)
            aux[tr_rows] .= .!ismissing.(view(col, tr_rows))
            esample[tr_rows] .&= view(aux, tr_rows)
        end
    end
    isempty(treatintvars) || (tr_rows[tr_rows] .&= view(esample, tr_rows))

    overlap!(esample, tr_rows, aux, tr, pr, treatname, data)
    sum(esample) == 0 && error("no nonmissing data")
    return (esample=esample, tr_rows=tr_rows::BitVector)
end

"""
    CheckVars <: StatsStep

Call [`DiffinDiffsBase.checkvars!`](@ref) to exclude invalid rows for relevant variables.
"""
const CheckVars = StatsStep{:CheckVars, typeof(checkvars!), true}

required(::CheckVars) = (:data, :pr, :yterm, :treatname, :esample, :aux,
    :treatintterms, :xterms)
transformed(::CheckVars, @nospecialize(nt::NamedTuple)) =
    (typeof(nt.tr), (termvars(nt.tr)...,))

combinedargs(step::CheckVars, allntargs) =
    combinedargs(step, allntargs, typeof(allntargs[1].tr))

combinedargs(::CheckVars, allntargs, ::Type{DynamicTreatment{SharpDesign}}) =
    (allntargs[1].tr,)

copyargs(::CheckVars) = (5,)

"""
    groupsample(esample)

Return the argument without change for allowing later comparisons based on object-id.
See also [`GroupSample`](@ref).
"""
groupsample(esample::BitVector) = (esample=esample,)

"""
    GroupSample <: StatsStep

Call [`DiffinDiffsBase.groupsample`](@ref)
to obtain one of the instances of `esample`
that have been grouped by equality (`hash`)
for allowing later comparisons based on object-id.

This step is only useful when working with [`@specset`](@ref) and [`proceed`](@ref).
"""
const GroupSample = StatsStep{:GroupSample, typeof(groupsample), false}

required(::GroupSample) = (:esample,)

"""
    makeweights(args...)

Construct a generic `Weights` vector.
See also [`MakeWeights`](@ref).
"""
function makeweights(data, esample::BitVector, weightname::Symbol)
    weights = Weights(convert(Vector{Float64}, view(getcolumn(data, weightname), esample)))
    all(isfinite, weights) || error("data column $weightname contain not-a-number values")
    return (weights=weights,)
end

function makeweights(data, esample::BitVector, weightname::Nothing)
    weights = uweights(sum(esample))
    return (weights=weights,)
end

"""
    MakeWeights <: StatsStep

Call [`DiffinDiffsBase.makeweights`](@ref) to create a generic `Weights` vector.
"""
const MakeWeights = StatsStep{:MakeWeights, typeof(makeweights), true}

required(::MakeWeights) = (:data, :esample)
default(::MakeWeights) = (weightname=nothing,)


text:
using Qaintessent

"""
Identity map as sparse matrix with real-valued entries.
"""
sparse_identity(n) = sparse(1.0*I, n, n)

"""
Construct a parametrized quantum layer consisting of single qubit rotation gates.
"""
function single_qubit_layer(L::Integer)
    # local gates
    # [circuit_gate(i, RotationGate(0.05*π*randn(3))) for i in 1:L]
    cgc = CircuitGate[]
    append!(cgc, [circuit_gate(i, RxGate(0.05*π*randn())) for i in 1:L])
    append!(cgc, [circuit_gate(i, RyGate(0.05*π*randn())) for i in 1:L])
    append!(cgc, [circuit_gate(i, RzGate(0.05*π*randn())) for i in 1:L])
    cgc
end

"""
Construct a parametrized quantum layer consisting of two qubit entanglement gates for 1D chain.
"""
function trotterized_layer_1d(L::Integer, istart::Integer; boundary::Bool=false)
    # interaction gates
    vcat(
        [circuit_gate(i, i+1, EntanglementXXGate(0.05π*randn())) for i in istart:2:L-1],
        ifelse(boundary, [circuit_gate(L, 1, EntanglementXXGate(0.05π*randn()))], []),
        [circuit_gate(i, i+1, EntanglementYYGate(0.05π*randn())) for i in istart:2:L-1],
        ifelse(boundary, [circuit_gate(L, 1, EntanglementYYGate(0.05π*randn()))], []),
        [circuit_gate(i, i+1, EntanglementZZGate(0.05π*randn())) for i in istart:2:L-1],
        ifelse(boundary, [circuit_gate(L, 1, EntanglementZZGate(0.05π*randn()))], [])
    )
end

function create_cgc(L::Int; isboundary::Bool=false)
    cgc = CircuitGate[]
    layers = 3
    for _ in 1:layers
        append!(cgc, single_qubit_layer(L::Integer))
        append!(cgc, trotterized_layer_1d(L, 1; boundary=isboundary))
        append!(cgc, trotterized_layer_1d(L, 2; boundary=isboundary))

    end

    cgc
end

max_stars_repo_name: cmendl/density-matrix-bloch-qcircuit

max_issues_repo_name: cmendl/density-matrix-bloch-qcircuit

max_forks_repo_name: cmendl/density-matrix-bloch-qcircuit


text:
using BenchmarkBot

runbenchmarks(
    package="YaoArrayRegister",
    target=BenchmarkConfig(
        id="origin/multithreading",
        env = Dict("JULIA_NUM_THREADS"=>4),
        juliacmd=`julia -O3`
    ),
    baseline=BenchmarkConfig(
        id="master",
        env = Dict("JULIA_NUM_THREADS"=>1),
        juliacmd=`julia -O3`
    )
)

runbenchmarks(
    package="YaoArrayRegister",
    target=BenchmarkConfig(
        id="origin/multithreading",
        env = Dict("JULIA_NUM_THREADS"=>1),
        juliacmd=`julia -O3`
    ),
    baseline=BenchmarkConfig(
        id="master",
        env = Dict("JULIA_NUM_THREADS"=>1),
        juliacmd=`julia -O3`
    )
)


text:
using FFTransforms
import FFTransforms as FF
using FFTW
using Test

@testset "𝕎  constructors" begin

	# 1-d constructors
	# ------------------
	ns  = [7, 12]

	Ia  = 𝕀(ns[1])
	Ib  = 𝕀(ns[2])
	Ic  = 𝕀(ns[1], 2π)
	Id  = 𝕀(ns[2], 2π)

	Wa  = 𝕎(ns[1])
	Wb  = 𝕎(ns[2])     * 2.0
	Wc  = 𝕎(ns[1], 2π) * 2.0f0
	Wd  = 𝕎(ns[2], 2π)

	W32a  = 𝕎(Float32, ns[1])
	W32b  = 𝕎(Float32, ns[2]) * 2.0
	W32c  = 𝕎(Float32, ns[1], 2π) * 2.0f0
	W32d  = 𝕎(Float32, ns[2], 2π)

	rWa  = 𝕎(Float64, ns[1])
	rWb  = 𝕎(Float64, ns[2]) * 2.0
	rWc  = 𝕎(Float64, ns[1], 2π) * 2.0f0
	rWd  = 𝕎(Float64, ns[2], 2π)

	rW32a  = 𝕎(Float32, ns[1])
	rW32b  = 𝕎(Float32, ns[2]) * 2.0
	rW32c  = 𝕎(Float32, ns[1], 2π) * 2.0f0
	rW32d  = 𝕎(Float32, ns[2], 2π)

	# kron of 1-d
	# ------------------

	W = Ia ⊗ Wa ⊗ rW32d
	W = Ib ⊗ rW32a ⊗ Wb
	W = Ic ⊗ rW32a ⊗ Wb
	W = rW32a ⊗ Wb ⊗ Ic

	szi  = @inferred size_in(W)
	szo  = @inferred size_out(W)
	WTf  = @inferred eltype_in(W)
	WTi  = @inferred eltype_out(W)
	Δx   = @inferred FF.Δpix(W)
	Δk   = @inferred FF.Δfreq(W)
	nq   = @inferred FF.nyq(W)
	ωx   = @inferred FF.Ωpix(W)
	ωk   = @inferred FF.Ωfreq(W)
	invs = @inferred FF.inv_scale(W)
	us   = @inferred FF.unitary_scale(W)
	os   = @inferred FF.ordinary_scale(W)
	xvecs   = @inferred FF.pix(W)
	kvecs   = @inferred FF.freq(W)
	xarrays = @inferred FF.fullpix(W)
	karrays = @inferred FF.fullfreq(W)
	λmat    = @inferred FF.wavenum(W)

	sW = @inferred W * (1/√(2π))
	uW = @inferred FF.unitary_scale(W) * W
	oW = @inferred FF.ordinary_scale(W) * W
	W′ = @inferred FF.unscale(sW)
	rW = @inferred FF.real(W)
	cW = @inferred FF.complex(W)

	# TODO get these inferred correctly
	P  = plan(W)
	sP = plan(sW)
	uP = plan(uW)
	oP = plan(oW)
	P′ = plan(W′)
	rP = plan(rW)
	cP = plan(cW)

	X  = rand(eltype_in(W), size_in(W))
	sX = rand(eltype_in(sW), size_in(sW))
	uX = rand(eltype_in(uW), size_in(uW))
	oX = rand(eltype_in(oW), size_in(oW))
	X′ = rand(eltype_in(W′), size_in(W′))
	rX = rand(eltype_in(rW), size_in(rW))
	cX = rand(eltype_in(cW), size_in(cW))

	Pout  = @inferred P  * X
	sPout = @inferred sP * sX
	uPout = @inferred uP * uX
	oPout = @inferred oP * oX
	P′out = @inferred P′ * X′
	rPout = @inferred rP * rX
	cPout = @inferred cP * cX

	@test sum(abs2, X  - P  \ Pout ) ./ length(X ) ≈ 0 atol = 1e-7
	@test sum(abs2, sX - sP \ sPout) ./ length(sX) ≈ 0 atol = 1e-7
	@test sum(abs2, uX - uP \ uPout) ./ length(uX) ≈ 0 atol = 1e-7
	@test sum(abs2, oX - oP \ oPout) ./ length(oX) ≈ 0 atol = 1e-7
	@test sum(abs2, X′ - P′ \ P′out) ./ length(X′) ≈ 0 atol = 1e-7
	@test sum(abs2, rX - rP \ rPout) ./ length(rX) ≈ 0 atol = 1e-7
	@test sum(abs2, cX - cP \ cPout) ./ length(cX) ≈ 0 atol = 1e-7

end

@testset "𝕌, 𝔽" begin

	Ia  = 𝕀(100)
	for (n,T,p) in zip(8:11, (Float32, Float64, ComplexF32, ComplexF64), (1, π, 0.1, 2π))
		@test 𝕌(n)         == 𝕎(n) * FF.unitary_scale(𝕎(n))
		@test 𝕌(n)⊗Ia      == (𝕎(n)⊗Ia) * FF.unitary_scale(𝕎(n)⊗Ia)
		@test Ia⊗𝕌(n)      == (Ia⊗𝕎(n)) * FF.unitary_scale(Ia⊗𝕎(n))
		@test 𝕌(T,n)       == 𝕎(T,n) * FF.unitary_scale(𝕎(T,n))
		@test 𝕌(T,n)⊗Ia    == (𝕎(T,n)⊗Ia) * FF.unitary_scale(𝕎(T,n)⊗Ia)
		@test Ia⊗𝕌(T,n)    == (Ia⊗𝕎(T,n)) * FF.unitary_scale(Ia⊗𝕎(T,n))
		@test 𝕌(T,n,p)     == 𝕎(T,n,p) * FF.unitary_scale(𝕎(T,n,p))
		@test 𝕌(T,n,p)⊗Ia  == (𝕎(T,n,p)⊗Ia) * FF.unitary_scale(𝕎(T,n,p)⊗Ia)
		@test Ia⊗𝕌(T,n,p)  == (Ia⊗𝕎(T,n,p)) * FF.unitary_scale(Ia⊗𝕎(T,n,p))
		@test real(𝕌(n))   == 𝕌(Float64,n)

		@test 𝔽(n)         == 𝕎(n) * FF.ordinary_scale(𝕎(n))
		@test 𝔽(n)⊗Ia      == (𝕎(n)⊗Ia) * FF.ordinary_scale(𝕎(n)⊗Ia)
		@test Ia⊗𝔽(n)      == (Ia⊗𝕎(n)) * FF.ordinary_scale(Ia⊗𝕎(n))
		@test 𝔽(T,n)       == 𝕎(T,n) * FF.ordinary_scale(𝕎(T,n))
		@test 𝔽(T,n)⊗Ia    == (𝕎(T,n)⊗Ia) * FF.ordinary_scale(𝕎(T,n)⊗Ia)
		@test Ia⊗𝔽(T,n)    == (Ia⊗𝕎(T,n)) * FF.ordinary_scale(Ia⊗𝕎(T,n))
		@test 𝔽(T,n,p)     == 𝕎(T,n,p) * FF.ordinary_scale(𝕎(T,n,p))
		@test 𝔽(T,n,p)⊗Ia  == (𝕎(T,n,p)⊗Ia) * FF.ordinary_scale(𝕎(T,n,p)⊗Ia)
		@test Ia⊗𝔽(T,n,p)  == (Ia⊗𝕎(T,n,p)) * FF.ordinary_scale(Ia⊗𝕎(T,n,p))
		@test real(𝔽(n))   == 𝔽(Float64,n) == 𝕎(Float64, n) * FF.ordinary_scale(𝕎(Float64, n))
		@test real(𝔽(n))   == 𝔽(Float64,n)
	end

end

@testset "plan and adjoint (low level)" begin

	T_forward_arg  = (Float32, Float64, Complex{Float32}, Complex{Float64})
	sz_forward_arg = ((1024,1024), (5,), (7,8,16), (14,4))
	region         = ((false,true), (true,), (false,true,false), (true,false))
	scale_forward =  (1, true, 1.0, 0.1f0)

	for indx = 1:length(T_forward_arg)
		#@show indx
		Tf  = T_forward_arg[indx]
		szf = sz_forward_arg[indx]
		reg = region[indx]
		scf = scale_forward[indx]
		FT = FF.plan(
			Tf,
			FF.SizeInt{szf},
			FF.RegionBool{reg},
			scf
		)

		X = rand(Tf, szf)
		Y = FT.scale_forward .* (FT.unscaled_forward_transform * X)
		X′ = FT.scale_inverse .* (FT.unscaled_inverse_transform * Y)
		@test sum(abs2, X .- X′) / sum(szf) ≈ 0 atol=1e-7

		reg_tp = tuple(findall(reg)...)
		if Tf <: Real
			@test sum(abs2, Y .- scf .* rfft(X,reg_tp)) / sum(szf) ≈ 0 atol=1e-7
		else
			@test sum(abs2, Y .- scf .* fft(X,reg_tp)) / sum(szf) ≈ 0 atol=1e-7
		end
	end

end


text:
#####
##### Transition rate (Q) matrices.
#####

####
#### utilities
####

"""
$(SIGNATURES)

Assert that both axes are the same, and return a single one.

Internal, not part of the API. Works with generalized indexing.
"""
function checked_square_axis(A::AbstractMatrix)
    a1, a2 = axes(A)
    @argcheck a1 == a2 "Mismatching axes for matrix."
    a1
end

"""
$(SIGNATURES)

Check that off-diagonal elements are nonnegative, and return their sum by row.

Internal, not part of the API. Works with generalized indexing.
"""
function off_diagonal_checked_row_sum(A::AbstractMatrix{T}) where {T}
    _off_diagonal_checked_row_sum!(zeros(result_ring(T), checked_square_axis(A)), A)
end

"""
$(SIGNATURES)

Accumulate the sum of off-diagonal elements into `d` by row. Return `d`.

Internal, not part of the API. Works with generalized indexing.
"""
function _off_diagonal_checked_row_sum!(d, A::AbstractMatrix)
    @inbounds for ι in CartesianIndices(A)
        if ι[1] ≠ ι[2]
            a = A[ι]
            @argcheck a ≥ 0 DomainError(a, "Element at $(ι) is negative.")
            d[ι[1]] += a
        end
    end
    d
end

function _off_diagonal_checked_row_sum!(d, A::SparseMatrixCSC)
    rows = rowvals(A)
    vals = nonzeros(A)
    m, n = size(A)
    for j in 1:n
        for k in nzrange(A, j)
            i = rows[k]
            if i ≠ j
                a = vals[k]
                @argcheck a ≥ 0 DomainError(a, "Element at $(CartesianIndex(i,j)) is negative.")
                d[i] += a
            end
        end
    end
    d
end

"""
$(SIGNATURES)

Set the diagonal of the argument (which is modified) so that rows sum to `0`.

Internal, not part of the API. Works with generalized indexing.
"""
function normalize_diagonal!(A::AbstractMatrix)
    d = off_diagonal_checked_row_sum(A)
    @inbounds for i in checked_square_axis(A)
        A[i, i] = -d[i]
    end
    A
end

####
#### Transition rate matrices
####

struct TransitionRateMatrix{T, S <: AbstractMatrix} <: AbstractMatrix{T}
    matrix::S
    # NOTE: this is the unchecked constructor. Don't use unless you are sure that
    # 1. `matrix` is square, with the same `axis` 1 and 2,
    # 2. its rows sum to `0`,
    # 3. it will not be modified later (ie does not share structure).
    function TransitionRateMatrix(::Val{:trust}, matrix::S) where {T, S <: AbstractMatrix{T}}
        checked_square_axis(matrix)
        new{T, S}(matrix)
    end
end

Base.size(m::TransitionRateMatrix) = size(m.matrix)

Base.getindex(m::TransitionRateMatrix, I...) = Base.getindex(m.matrix, I...)

Base.IndexStyle(::Type{TransitionRateMatrix{T,S}}) where {T,S} = IndexStyle(S)

"""
$(SIGNATURES)

Create a transition rate matrix from the argument. This makes a copy, checks that
off-diagonal elements are nonnegative, and sets the diagonal so that rows sum to `0`,
striving to preserve sparsity and structure.
"""
function TransitionRateMatrix(Q::AbstractMatrix)
    TransitionRateMatrix(Val{:trust}(), normalize_diagonal!(copy(Q)))
end

"""
$(SIGNATURES)

Create a `TransitionRateMatrix` from a matrix `Q`. Modifies the argument (to normalize the
diagonal).
"""
function TransitionRateMatrix!(Q::AbstractMatrix)
    TransitionRateMatrix(Val{:trust}(), normalize_diagonal!(Q))
end

# NOTE: remove when we don't support ≤ v"1.7-"
import LinearAlgebra
const _LU_NoPivot = VERSION >= v"1.7.0-DEV.1188" ? LinearAlgebra.NoPivot : Val{false}

"""
$(SIGNATURES)

Return the stationary distribution of a transition rate matrix as a vector.
"""
function stationary_distribution(m::TransitionRateMatrix)
    @unpack matrix = m
    @assert !Base.has_offset_axes(matrix) "Generalized indexing version not implemented."
    LU = lu(matrix, _LU_NoPivot(); check = false) # not pivoted, allow singular
    L = LU.L
    x = vcat(L[1:(end-1), 1:(end-1)]' \ -L[end, 1:(end-1)], 1)
    normalize!(x, 1)
    x
end


text:
include("intcode.jl")

function execute_and_show(prog, script)
    result = run_intcode!(copy(prog), Int8.(collect(replace(script, "\r" => ""))))
    if result[end] <= 127
        println(string(Char.(result)...))
    else
        println("Damage: $(result[end])")
    end
end

# Part 1: write the springscript program iteratively
prog = open(f -> parse.(Int64, split(readline(f), ",")), "day21_input.txt")
# simple description: (!a || (d && (!b || (b && !c)))
script = """NOT C J
AND B J
NOT B T
OR T J
AND D J
NOT A T
OR T J
WALK
"""
# Patterns we need to cross:
# .### -> jump on .### -> !a
# ..#. -> jump on #..# -> !b && d
# .#.. -> jump on ?#.# -> b && d && !c
#execute_and_show(prog, script)

# Part 2: much bigger input sensor range
# Again iteratively learn about which patterns we need to cross, writing scripts as we go
#
# ABCDEFGHI
# 011111111 -> jump on 011111111 -> !a
# 001011111 -> jump on 100101111 -> !b && d
# 010011111 -> jump on 110100111 -> b && !c && d && h
# 010100010 -> jump on 010100010 -> !a
#              DO NOT jump on 1101010000 -> modified rule for third option above
# 010001001 -> jump on 110100010 -> modified rule for third option above again
#
# simple description: (!a || (d && (!b || (b && !c && g)))
script = """NOT C J
AND B J
AND H J
NOT B T
OR T J
AND D J
NOT A T
OR T J
RUN
"""
execute_and_show(prog, script)

max_stars_repo_name: ericgreveson/adventofcode

max_issues_repo_name: ericgreveson/adventofcode

max_forks_repo_name: ericgreveson/adventofcode