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

global repeatAlleleSuffix = "__2"
repeatAlleleSuffix_re = Regex("$repeatAlleleSuffix\$")

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

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

Optional arguments:

- file name to write/save resulting CF table. If not specified, then the output
  data frame is not saved to a file.
- column numbers for the taxon names. 1-4 by default.
- any keyword arguments that `CSV.File` would accept.
  For example, delim=',' by default: columns are delimited by commas.
  Unless specified otherwise by the user, `pool`=false
  (to read taxon names as Strings, not levels of a categorical factor,
  for combining the 4 columns with taxon names more easily).
  The same CSV arguments are used to read both input file (mapping file and quartet file)

See also [`mapAllelesCFtable!`](@ref) to input DataFrames instead of file names.

If a `filename` is specified, such as "quartetCF_speciesNames.csv"
in the example below, this file is best read later with the option
`pool=false`. example:

```julia
mapAllelesCFtable("allele-species-map.csv", "allele-quartet-CF.csv";
                  filename = "quartetCF_speciesNames.csv")
df_sp = DataFrame(CSV.File("quartetCF_speciesNames.csv"); copycols=false); # DataFrame object
dataCF_specieslevel = readTableCF!(df_sp); # DataCF object
```
"""
function mapAllelesCFtable(alleleDF::AbstractString, cfDF::AbstractString;
        filename=""::AbstractString, columns=Int[]::Vector{Int}, CSVargs...)
    # force pool=false unless the user wants otherwise
    if :pool ∉ [pair[1] for pair in CSVargs]
        CSVargs = (CSVargs..., :pool=>false)
    end
    d = DataFrame(CSV.File(alleleDF; CSVargs...); copycols=false)
    d2 = DataFrame(CSV.File(cfDF; CSVargs...); copycols=false)
    mapAllelesCFtable!(d2,d, columns, filename != "", filename)
end

"""
    mapAllelesCFtable!(quartet CF DataFrame, mapping DataFrame, columns, write?, filename)

Modify (and return) the quartet concordance factor (CF) DataFrame:
replace each allele name by the species name that the allele maps onto
based on the mapping data frame. This mapping data frame should have columns
named "allele" and "species" (see `rename!` to change column names if need be).

If `write?` is `true`, the modified data frame is written to a file named "filename".

Warning: [`mapAllelesCFtable`](@ref) takes the quartet data file as its second
argument, while `mapAllelesCFtable!` takes the quartet data (which it modifies)
as its first argument.
"""
function mapAllelesCFtable!(cfDF::DataFrame, alleleDF::DataFrame, co::Vector{Int},write::Bool,filename::AbstractString)
    size(cfDF,2) >= 7 || error("CF DataFrame should have 7+ columns: 4taxa, 3CF, and possibly ngenes")
    if length(co)==0 co=[1,2,3,4]; end
    compareTaxaNames(alleleDF,cfDF,co)
    for j in 1:4
        for ia in 1:size(alleleDF,1) # for all alleles
            cfDF[!,co[j]] = map(x->replace(string(x),
                                         Regex("^$(string(alleleDF[ia,:allele]))\$") =>
                                         alleleDF[ia,:species]),
                                cfDF[!,co[j]])
        end
    end
    if write
        filename != "" || error("cannot write quartet CF with mapped alleles: empty filename")
        CSV.write(filename, cfDF)
    end
    return cfDF
end

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

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

# function to merge rows that have repeated taxon names by using the weigthed average of CF
# (if info on number of genes is provided) or simple average
function mergeRows(df::DataFrame, cols::Vector{Int})
    sorttaxa!(df, cols) # sort taxa alphabetically within each row
    colnamtax = DataFrames.propertynames(df)[cols[1:4]]
    colnam = DataFrames.propertynames(df)[cols[5:end]]
    df = combine(groupby(df, colnamtax, sort=false, skipmissing=false),
                 colnam .=> mean .=> colnam)
    # rename!(df, Dict((Symbol(n, "_mean"), n) for n in colnam) )
    n4tax = size(df,1) # total number of 4-taxon sets
    print("$n4tax unique 4-taxon sets were found. CF values of repeated 4-taxon sets will be averaged")
    println((length(cols)>7 ? " (ngenes too)." : "."))
    return df
end

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

# function to compare the taxon names in the allele-species matching table
# and the CF table
function compareTaxaNames(alleleDF::DataFrame, cfDF::DataFrame, co::Vector{Int})
    checkMapDF(alleleDF)
    #println("found $(length(alleleDF[1])) allele-species matches")
    CFtaxa = string.(mapreduce(x -> unique(skipmissing(x)), union, eachcol(cfDF[!,co[1:4]])))
    alleleTaxa = map(string, alleleDF[!,:allele]) # as string, too
    sizeCF = length(CFtaxa)
    sizeAllele = length(alleleTaxa)
    if sizeAllele > sizeCF
        @warn "some alleles in the mapping file do not occur in the quartet CF data. Extra allele names will be ignored"
        alleleTaxa = intersect(alleleTaxa,CFtaxa)
    end
    unchanged = setdiff(CFtaxa,alleleTaxa)
    if length(unchanged) == length(CFtaxa)
        @warn "None of the taxon names in CF table match with allele names in the mapping file"
    end
    if !isempty(unchanged)
        warnmsg = "not all alleles were mapped\n"
        warnmsg *= "alleles not mapped to a species name:"
        for n in unchanged warnmsg *= " $n"; end
        @warn warnmsg
    end
    return nothing
end

# function to check that the allele df has one column labelled alleles and one column labelled species
function checkMapDF(alleleDF::DataFrame)
    size(alleleDF,2) >= 2 || error("Allele-Species matching Dataframe should have at least 2 columns")
    :allele in DataFrames.propertynames(alleleDF) || error("In allele mapping file there is no column named allele")
    :species in DataFrames.propertynames(alleleDF) || error("In allele mapping file there is no column named species")
end

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

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


text:
export record!, new_history, print_history, guess

"""
    record!(h, word, code)
Record the outcome of guessing `word` in the history object `h`. Here `code`
is a five-tuple with the following meanings:
* `0` means this letter is not in the word (gray)
* `1` means this letter is in the word but in the wrong position (yellow)
* `2` means this letter is correct and in the correct position (green)
Colors are as shown in the online Wordle game.

The argument `h` is initialized like this: `h = new_history()`.
"""
function record!(history::Dict{Int,Byte}, word::String, code::NTuple{five,Int})
    w = GUESS_DICT[uppercase(word)]
    c = code_2_byte(code)
    history[w] = c
    guess(history)
end

"""
    new_history()
Start a new solving session by creating a history object. Optional argument is a
scoring function (`min_max_score` is the default).
"""
function new_history(scorer::Function = min_max_score)
    h = Dict{Int,Byte}()
    g = FIRST_GUESS[scorer]
    guess = GUESS_LIST[g]
    println("Suggest you try \"$guess\"")
    return h
end

"""
    print_history(h)
Prints out the history object `h` in color.
"""
function print_history(history::Dict{Int, Byte})
    for w in keys(history)
        wordle_print(w, history[w])
    end
end

"""
    guess(history, scorer = min_max_score)
Provide a guess to try in solving Wordle. Automatically called by
`record!`.
"""
function guess(history::Dict{Int,Byte}, scorer = min_max_score)
    possibles = filter_answers(history)
    if length(possibles) == 0
        println("I give up. No solution possible.")
        return
    end
    guess = GUESS_LIST[best_guess(possibles, scorer)]
    return guess
end

ext: jl


text:
using Whirl
using Plots

gr()

Δt = 1e-4
n_steps = 10000
n_skips = 50

stepper = build(rarefaction_wave_builder(number_of_particles = 1000))

anim = @animate for _ = 1:n_steps
    sphplot(get_positions(stepper), get_velocities(stepper))
    step!(stepper, Δt)
end every n_skips
gif(anim, "rarefaction_wave.gif", fps = 15)

max_stars_repo_name: lars-frogner/Whirl

max_issues_repo_name: lars-frogner/Whirl

max_forks_repo_name: lars-frogner/Whirl


text:
"""
    IntervalConstant(x::SeaPearl.IntVar, lower::Int, upper::Int)

Inequality constraint, `lower <= x <= upper`
"""
struct IntervalConstant <: OnePropagationConstraint
    x::AbstractIntVar
    lower::Int
    upper::Int
    active::StateObject{Bool}
    function IntervalConstant(x, lower, upper, trailer)
        constraint = new(x, lower, upper, StateObject(true, trailer))
        addOnDomainChange!(x, constraint)
        return constraint
    end
end

"""
    propagate!(constraint::IntervalConstant, toPropagate::Set{Constraint}, prunedDomains::CPModification)

`IntervalConstant` propagation function. Basically remove the values above `upper` and below `lower`.
"""
function propagate!(constraint::IntervalConstant, toPropagate::Set{Constraint}, prunedDomains::CPModification)
    if !constraint.active.value
        return true
    end
    setValue!(constraint.active, false)

    addToPrunedDomains!(prunedDomains, constraint.x, vcat(removeBelow!(constraint.x.domain, constraint.lower), removeAbove!(constraint.x.domain, constraint.upper)))
    triggerDomainChange!(toPropagate, constraint.x)
    if constraint in toPropagate
        pop!(toPropagate, constraint)
    end
    return !isempty(constraint.x.domain)
end

function Base.show(io::IO, ::MIME"text/plain", con::IntervalConstant)
    println(io, typeof(con), ": ", con.x.id, " ∈ ⟦", con.lower, " ", con.upper, "⟧, active = ", con.active)
    print(io, "   ", con.x)
end

function Base.show(io::IO, con::IntervalConstant)
    print(io, typeof(con), ": ",  con.x.id, " ∈ ⟦", con.lower, " ", con.upper, "⟧")
end


text:
function parse_input(lines)
    function parse_bag_content(content)
        if content == "no other bags."
            []
        elseif !occursin(",", content)
            words = split(content)
            child_bag = join(words[2:3], " ")
            amount = parse(Int64, words[1])
            [(child_bag, amount)]
        else
            map(x -> begin
                    words = split(x)
                    amount = parse(Int64, words[1])
                    node = join(words[2:3], " ")
                    (node, amount)
                    end, split(content, ", "))
        end
    end

    mapping = Dict{AbstractString, Set{Tuple{AbstractString, Int64}}}()
    for line in lines
        first_part, second_part = split(line, " contain ")
        parent_bag = join(split(first_part)[1:2], " ")
        children = parse_bag_content(second_part)
        for child in children
            vals = get(mapping, parent_bag, Set{Tuple{AbstractString, Int64}}())
            push!(vals, child)
            mapping[parent_bag] = vals
        end
    end
    mapping
end

function nr_inner_bags(mapping, start_bag)
    children = get(mapping, start_bag, Set())
    foldl((acc, (child, amount)) -> acc + amount + amount * nr_inner_bags(mapping, child), children; init=0)
end

function main()
    lines = readlines("../resources/input_07.txt")
    children = parse_input(lines)
    result = nr_inner_bags(children, "shiny gold")
    println(result)
end

main()


text:
type BayesTreeNodeData
  frontalIDs::Vector{Int}
  conditIDs::Vector{Int}
  inmsgIDs::Vector{Int}
  potIDs::Vector{Int} # this is likely redundant TODO -- remove
  potentials::Vector{Int}
  cliqAssocMat::Array{Bool,2}
  cliqMsgMat::Array{Bool,2}
  directvarIDs::Vector{Int}
  directFrtlMsgIDs::Vector{Int}
  msgskipIDs::Vector{Int}
  itervarIDs::Vector{Int}
  directPriorMsgIDs::Vector{Int}
  debug
  debugDwn
  BayesTreeNodeData() = new()
  BayesTreeNodeData(x...) = new(x[1],x[2],x[3],x[4],x[5],x[6],x[7],x[8],x[9],x[10],x[11],x[12],x[13],x[14])
end

# TODO -- this should be a constructor
function emptyBTNodeData()
  BayesTreeNodeData(Int[],Int[],Int[],
                    Int[],Int[],Array{Bool}(0,0),
                    Array{Bool}(0,0),Int[],Int[],
                    Int[],Int[],Int[],
                    nothing, nothing)
end

# BayesTree declarations
type BayesTree
  bt
  btid::Int
  cliques::Dict{Int,Graphs.ExVertex}
  frontals::Dict{String,Int}
end

function emptyBayesTree()
    bt =   BayesTree(Graphs.inclist(Graphs.ExVertex,is_directed=true),
                     0,
                     Dict{Int,Graphs.ExVertex}(),
                     #[],
                     Dict{AbstractString, Int}())
    return bt
end

# create a new clique
function addClique!(bt::BayesTree, fg::FactorGraph, varID::Int, condIDs::Array{Int}=Int[])
  bt.btid += 1
  clq = Graphs.add_vertex!(bt.bt, ExVertex(bt.btid,string("Clique",bt.btid)))
  bt.cliques[bt.btid] = clq

  clq.attributes["label"] = ""
  # Specific data container
  clq.attributes["data"] = emptyBTNodeData()

  appendClique!(bt, bt.btid, fg, varID, condIDs)
  return clq
end

# generate the label for particular clique -- graphviz drawing
function makeCliqueLabel(fgl::FactorGraph, bt::BayesTree, clqID::Int)
  clq = bt.cliques[clqID]
  flbl = ""
  clbl = ""
  for fr in clq.attributes["data"].frontalIDs
    flbl = string(flbl,localapi.getvertex(fgl,fr).attributes["label"], ",") #fgl.v[fr].
  end
  for cond in clq.attributes["data"].conditIDs
    clbl = string(clbl, localapi.getvertex(fgl,cond).attributes["label"], ",") # fgl.v[cond].
  end
  clq.attributes["label"] = string(flbl, ": ", clbl)
end

# add a conditional ID to clique
function appendConditional(bt::BayesTree, clqID::Int, condIDs::Array{Int,1})
  clq = bt.cliques[clqID]
  clq.attributes["data"].conditIDs = union(clq.attributes["data"].conditIDs, condIDs)
end

# Add a new frontal variable to clique
function appendClique!(bt::BayesTree, clqID::Int, fg::FactorGraph, varID::Int, condIDs::Array{Int,1}=Int[])
  clq = bt.cliques[clqID]
  var = localapi.getvertex(fg, varID) # fg.v[varID]
  # add frontal variable
  push!(clq.attributes["data"].frontalIDs,varID)
  # total dictionary of frontals for easy access
  bt.frontals[var.attributes["label"]] = clqID#bt.btid

  appendConditional(bt, clqID, condIDs)
  makeCliqueLabel(fg, bt, clqID)
  nothing
end

# instantiate a new child clique in the tree
function newChildClique!(bt::BayesTree, fg::FactorGraph, CpID::Int, varID::Int, Sepj::Array{Int,1})
  chclq = addClique!(bt, fg, varID, Sepj)
  parent = bt.cliques[CpID]
  # Staying with Graphs.jl for tree in first stage
  edge = Graphs.make_edge(bt.bt, parent, chclq)
  Graphs.add_edge!(bt.bt, edge)

  return chclq
end

# post order tree traversal and build potential functions
function findCliqueFromFrontal(bt::BayesTree, frtlID::Int)
  for cliqPair in bt.cliques
    id = cliqPair[1]
    cliq = cliqPair[2]
    for frtl in cliq.attributes["data"].frontalIDs
      if frtl == frtlID
        return cliq
      end
    end
  end
  error("Clique with desired frontal ID not found")
end

# eliminate a variable for new
function newPotential(tree::BayesTree, fg::FactorGraph, var::Int, prevVar::Int, p::Array{Int,1})
    firvert = localapi.getvertex(fg,var)
    if (length(getData(firvert).separator) == 0)
      # warn("newPotential -- sep length is 0")
      if (length(tree.cliques) == 0)
        addClique!(tree, fg, var)
      else
        appendClique!(tree, 1, fg, var) # add to root
      end
    else
      Sj = getData(firvert).separator
      # find parent clique Cp that containts the first eliminated variable of Sj as frontal
      firstelim = 99999999999
      for s in Sj
        temp = findfirst(p, s)
        if (temp < firstelim)
          firstelim = temp
        end
      end
      felbl = localapi.getvertex(fg, p[firstelim]).attributes["label"]
      CpID = tree.frontals[felbl]
      # look to add this conditional to the tree
      unFC = union(tree.cliques[CpID].attributes["data"].frontalIDs, tree.cliques[CpID].attributes["data"].conditIDs)
      if (sort(unFC) == sort(Sj))
        appendClique!(tree, CpID, fg, var)
      else
        newChildClique!(tree, fg, CpID, var, Sj)
      end
    end
end

# build the whole tree in batch format
function buildTree!(tree::BayesTree, fg::FactorGraph, p::Array{Int,1})
  rp = flipdim(p,1)
  prevVar = 0
  for var in rp
    newPotential(tree, fg, var, prevVar, p)
    prevVar = var
  end
end

## Find batch belief propagation solution
function prepBatchTree!(fg::FactorGraph; ordering::Symbol=:qr,drawpdf::Bool=false)
  p = IncrementalInference.getEliminationOrder(fg, ordering=ordering)
  println()
  fge = deepcopy(fg)
  println("Building Bayes net...")
  buildBayesNet!(fge, p)

  tree = emptyBayesTree()
  buildTree!(tree, fge, p)

  println("Bayes Net")
  # sleep(0.1)
  #fid = open("bn.dot","w+")
  #write(fid,to_dot(fge.bn))
  #close(fid)

  # Michael reference -- x2->x1, x2->x3, x2->x4, x2->l1, x4->x3, l1->x3, l1->x4
  println("Bayes Tree")
  if drawpdf
    fid = open("bt.dot","w+")
    write(fid,to_dot(tree.bt))
    close(fid)
    run(`dot bt.dot -Tpdf -o bt.pdf`)
  end

  # GraphViz.Graph(to_dot(tree.bt))
  #Michael reference 3sig -- x2l1x4x3    x1|x2

  println("Find potential functions for each clique")
  cliq = tree.cliques[1] # start at the root
  buildCliquePotentials(fg, tree, cliq); # fg does not have the marginals as fge does

  # now update all factor graph vertices used for this tree
  for v in vertices(fg.g)
    dlapi.updatevertex!(fg, v)
  end

  return tree
end

function resetData!(vdata::VariableNodeData)
  vdata.eliminated = false
  vdata.BayesNetOutVertIDs = Int[]
  vdata.BayesNetVertID = 0
  vdata.separator = Int[]
  nothing
end

function resetData!(vdata::FunctionNodeData)
  vdata.eliminated = false
  vdata.potentialused = false
  nothing
end

function resetFactorGraphNewTree!(fg::FactorGraph)
  for v in vertices(fg.g)
    resetData!(getData(v))
    localapi.updatevertex!(fg, v)
  end

  nothing
end

function wipeBuildNewTree!(fg::FactorGraph; ordering=:qr,drawpdf=false)
  resetFactorGraphNewTree!(fg);
  return prepBatchTree!(fg, ordering=ordering, drawpdf=drawpdf);
end

function whichCliq{T <: AbstractString}(bt::BayesTree, frt::T)
    bt.cliques[bt.frontals[frt]]
end
whichCliq(bt::BayesTree, frt::Symbol) = whichCliq(bt, string(frt))

function appendUseFcts!(usefcts, lblid::Int, fct::Graphs.ExVertex, fid::Int)
  for tp in usefcts
    if tp == fct.index
      return nothing
    end
  end
  tpl = fct.index
  push!(usefcts, tpl )
  nothing
end

function getCliquePotentials!(fg::FactorGraph, bt::BayesTree, cliq::Graphs.ExVertex)
    frtl = cliq.attributes["data"].frontalIDs
    cond = cliq.attributes["data"].conditIDs
    allids = [frtl;cond]
    alldimIDs = Int[]
    for fid in frtl
      alldimIDs = [alldimIDs; getData(localapi.getvertex(fg,fid)).dimIDs]
    end
    for cid in cond
      alldimIDs = [alldimIDs; getData(localapi.getvertex(fg,cid)).dimIDs]
    end

    for fid in frtl
        usefcts = []
        for fct in localapi.outneighbors(fg, localapi.getvertex(fg,fid))
            if getData(fct).potentialused!=true
                loutn = localapi.outneighbors(fg, fct)
                if length(loutn)==1
                    appendUseFcts!(usefcts, fg.IDs[Symbol(loutn[1].label)], fct, fid)
                    # TODO -- make update vertex call
                    fct.attributes["data"].potentialused = true
                    localapi.updatevertex!(fg, fct)
                end
                for sepSearch in loutn
                    sslbl = Symbol(sepSearch.label)
                    if (fg.IDs[sslbl] == fid)
                        continue # skip the fid itself
                    end
                    sea = findmin(abs.(allids-fg.IDs[sslbl]))
                    if sea[1]==0.0
                        appendUseFcts!(usefcts, fg.IDs[sslbl], fct, fid)
                        # usefcts = [usefcts;(fg.IDs[sslbl], fct, fid)]
                        fct.attributes["data"].potentialused = true #fct.attributes["potentialused"] = true
                        localapi.updatevertex!(fg, fct)
                    end
                end
            end
        end
        cliq.attributes["data"].potentials=union(cliq.attributes["data"].potentials,usefcts)
    end
    return nothing
end

function getCliquePotentials!(fg::FactorGraph, bt::BayesTree, chkcliq::Int)
    getCliquePotentials!(fg, bt.cliques[chkcliq])
end

function cliqPotentialIDs(cliq::Graphs.ExVertex)
  potIDs = Int[]
  for idfct in cliq.attributes["data"].potentials
    push!(potIDs,idfct)
  end
  return potIDs
end

function collectSeparators(bt::BayesTree, cliq::Graphs.ExVertex)
  allseps = Int[]
  for child in out_neighbors(cliq, bt.bt)#tree
      allseps = [allseps; child.attributes["data"].conditIDs]
  end
  return allseps
end

function compCliqAssocMatrices!(fgl::FactorGraph, bt::BayesTree, cliq::Graphs.ExVertex)
  frtl = cliq.attributes["data"].frontalIDs
  cond = cliq.attributes["data"].conditIDs
  inmsgIDs = collectSeparators(bt, cliq)
  potIDs = cliqPotentialIDs(cliq)
  # Construct associations matrix here
  # matrix has variables are columns, and messages/constraints as rows
  cols = [frtl;cond]
  cliq.attributes["data"].inmsgIDs = inmsgIDs
  cliq.attributes["data"].potIDs = potIDs
  cliqAssocMat = Array{Bool,2}(length(potIDs), length(cols))
  cliqMsgMat = Array{Bool,2}(length(inmsgIDs), length(cols))
  fill!(cliqAssocMat, false)
  fill!(cliqMsgMat, false)
  for j in 1:length(cols)
    for i in 1:length(inmsgIDs)
      if cols[j] == inmsgIDs[i]
        cliqMsgMat[i,j] = true
      end
    end
    for i in 1:length(potIDs)
      idfct = cliq.attributes["data"].potentials[i]
      if idfct == potIDs[i] # sanity check on clique potentials ordering
        for vertidx in getData(getVert(fgl, idfct, api=localapi)).fncargvID
        # for vertidx in getData(getVertNode(fgl, idfct)).fncargvID
          if vertidx == cols[j]
            cliqAssocMat[i,j] = true
          end
        end
      else
        prtslperr("compCliqAssocMatrices! -- potential ID ordering was lost")
      end
    end
  end
  cliq.attributes["data"].cliqAssocMat = cliqAssocMat
  cliq.attributes["data"].cliqMsgMat = cliqMsgMat
  nothing
end

function getCliqAssocMat(cliq::Graphs.ExVertex)
  cliq.attributes["data"].cliqAssocMat
end
function getCliqMsgMat(cliq::Graphs.ExVertex)
  cliq.attributes["data"].cliqMsgMat
end
function getCliqMat(cliq::Graphs.ExVertex; showmsg=true)
  assocMat = getCliqAssocMat(cliq)
  msgMat = getCliqMsgMat(cliq)
  mat = showmsg ? [assocMat;msgMat] : assocMat
  return mat
end

# function spyCliqMat(cliq::Graphs.ExVertex; showmsg=true)
#   mat = deepcopy(getCliqMat(cliq, showmsg=showmsg))
#   # TODO -- add improved visualization here, iter vs skip
#   mat = map(Float64, mat)*2.0-1.0
#   numlcl = size(getCliqAssocMat(cliq),1)
#   mat[(numlcl+1):end,:] *= 0.9
#   mat[(numlcl+1):end,:] -= 0.1
#   numfrtl1 = floor(Int,length(cliq.attributes["data"].frontalIDs)+1)
#   mat[:,numfrtl1:end] *= 0.9
#   mat[:,numfrtl1:end] -= 0.1
#   @show cliq.attributes["data"].itervarIDs
#   @show cliq.attributes["data"].directvarIDs
#   @show cliq.attributes["data"].msgskipIDs
#   @show cliq.attributes["data"].directFrtlMsgIDs
#   @show cliq.attributes["data"].directPriorMsgIDs
#   sp = Gadfly.spy(mat)
#   push!(sp.guides, Gadfly.Guide.title("$(cliq.attributes["label"]) || $(cliq.attributes["data"].frontalIDs) :$(cliq.attributes["data"].conditIDs)"))
#   push!(sp.guides, Gadfly.Guide.xlabel("fmcmcs $(cliq.attributes["data"].itervarIDs)"))
#   push!(sp.guides, Gadfly.Guide.ylabel("lcl=$(numlcl) || msg=$(size(getCliqMsgMat(cliq),1))" ))
#   return sp
# end
# function spyCliqMat(bt::BayesTree, lbl::Symbol; showmsg=true)
#   spyCliqMat(whichCliq(bt,lbl), showmsg=showmsg)
# end

function countSkips(bt::BayesTree)
  skps = 0
  for cliq in bt.cliques
    m = getCliqMat(cliq[2])
    mi = map(Int,m)
    skps += sum(map(Int,sum(mi,1) .== 1))
  end
  return skps
end

function skipThroughMsgsIDs(cliq::Graphs.ExVertex)
  numfrtl1 = floor(Int,length(cliq.attributes["data"].frontalIDs)+1)
  condAssocMat = cliq.attributes["data"].cliqAssocMat[:,numfrtl1:end]
  condMsgMat = cliq.attributes["data"].cliqMsgMat[:,numfrtl1:end]
  mat = [condAssocMat;condMsgMat];
  mab = sum(map(Int,mat),1) .== 1
  mabM = sum(map(Int,condMsgMat),1) .== 1
  mab = mab .& mabM
  # rang = 1:size(condMsgMat,2)
  msgidx = cliq.attributes["data"].conditIDs[vec(collect(mab))]
  return msgidx
end

function directPriorMsgIDs(cliq::Graphs.ExVertex)
  frtl = cliq.attributes["data"].frontalIDs
  cond = cliq.attributes["data"].conditIDs
  cols = [frtl;cond]
  mat = getCliqMat(cliq, showmsg=true)
  singr = sum(map(Int,mat),2) .== 1
  rerows = collect(1:length(singr))
  b = vec(collect(singr))
  rerows2 = rerows[b]
  sumsrAc = sum(map(Int,mat[rerows2,:]),1)
  sumc = sum(map(Int,mat),1)
  pmSkipCols = (sumsrAc - sumc) .== 0
  return cols[vec(collect(pmSkipCols))]
end

function directFrtlMsgIDs(cliq::Graphs.ExVertex)
  numfrtl = length(cliq.attributes["data"].frontalIDs)
  frntAssocMat = cliq.attributes["data"].cliqAssocMat[:,1:numfrtl]
  frtlMsgMat = cliq.attributes["data"].cliqMsgMat[:,1:numfrtl]
  mat = [frntAssocMat; frtlMsgMat];
  mab = sum(map(Int,mat),1) .== 1
  mabM = sum(map(Int,frtlMsgMat),1) .== 1
  mab = mab .& mabM
  return cliq.attributes["data"].frontalIDs[vec(collect(mab))]
end

function directAssignmentIDs(cliq::Graphs.ExVertex)
  # NOTE -- old version been included in iterated variable stack
  assocMat = cliq.attributes["data"].cliqAssocMat
  msgMat = cliq.attributes["data"].cliqMsgMat
  mat = [assocMat;msgMat];
  mab = sum(map(Int,mat),1) .== 1
  mabA = sum(map(Int,assocMat),1) .== 1
  mab = mab .& mabA
  frtl = cliq.attributes["data"].frontalIDs
  cond = cliq.attributes["data"].conditIDs
  cols = [frtl;cond]
  return cols[vec(collect(mab))]
  # also calculate how which are conditionals
end

function mcmcIterationIDs(cliq::Graphs.ExVertex)
  assocMat = cliq.attributes["data"].cliqAssocMat
  msgMat = cliq.attributes["data"].cliqMsgMat
  mat = [assocMat;msgMat];
  sum(sum(map(Int,mat),1)) == 0 ? error("mcmcIterationIDs -- unaccounted variables") : nothing
  mab = 1 .< sum(map(Int,mat),1)
  frtl = cliq.attributes["data"].frontalIDs
  cond = cliq.attributes["data"].conditIDs
  cols = [frtl;cond]
  # dev code
  # prioritize prior'ed variables, then high degree next
  # get all prior type rows, should be prioritized
  # TODO -- maybe sort according to decreasing degree
  # singleperrows = 1 .== sum(map(Int,mat),2)
  # singleperrows = collect(singleperrows)
  # prioredvarsum = sum(map(Int,mat[singleperrows]),1)
  # prioredperm = sortperm(prioredvarsum, rev=true) # decreasing count

  #end dev code
  # final output
  return cols[vec(collect(mab))]
end

function setCliqMCIDs!(cliq::Graphs.ExVertex)
  cliq.attributes["data"].directPriorMsgIDs = directPriorMsgIDs(cliq)

  # NOTE -- combined larger iter group
  cliq.attributes["data"].directvarIDs = directAssignmentIDs(cliq)
  # cliq.attributes["data"].itervarIDs = mcmcIterationIDs(cliq)
  # NOTE -- fix direct vs itervar issue, DirectVarIDs against Iters should also Iter
  usset = union(mcmcIterationIDs(cliq), directAssignmentIDs(cliq))
  cliq.attributes["data"].itervarIDs = setdiff(usset, cliq.attributes["data"].directPriorMsgIDs)
  cliq.attributes["data"].msgskipIDs = skipThroughMsgsIDs(cliq)
  cliq.attributes["data"].directFrtlMsgIDs = directFrtlMsgIDs(cliq)

  nothing
end

# post order tree traversal and build potential functions
function buildCliquePotentials(fg::FactorGraph, bt::BayesTree, cliq::Graphs.ExVertex)
    for child in out_neighbors(cliq, bt.bt)#tree
        buildCliquePotentials(fg, bt, child)
    end
    println("Get potentials $(cliq.attributes["label"])");
    getCliquePotentials!(fg, bt, cliq);

    compCliqAssocMatrices!(fg, bt, cliq);
    setCliqMCIDs!(cliq);

    nothing
end


text:
with_repo_access() do test_repo, path
    # blob
    @test_throws LibGitError{:Invalid,:NotFound} lookup_tag(test_repo, Oid("fa49b077972391ad58037050f2a75f74e3671e92"))
    # commit
    @test_throws LibGitError{:Invalid,:NotFound} lookup_tag(test_repo, Oid("8496071c1b46c854b31185ea97743be6a8774479"))
    # tree
    @test_throws LibGitError{:Invalid,:NotFound} lookup_tag(test_repo, Oid("c4dc1555e4d4fa0e0c9c3fc46734c7c35b3ce90b"))

    context("test reading a tag") do
        id = Oid("0c37a5391bbff43c37f0d0371823a5509eed5b1d")
        obj = lookup(test_repo, id)
        @test id == Oid(obj)
        @test isa(obj, GitTag)
        @test message(obj) == "test tag message\n"
        @test name(obj) == "v1.0"
        @test target_id(obj) == Oid("5b5b025afb0b4c913b4c338a42934a3863bf3644")
        @test target_id(obj) == Oid(target(obj))
        @test isa(target(obj), GitCommit)
        c = tagger(obj)
        @test name(c) == "Scott Chacon"
        @test time(c) == 1288114383
        @test email(c) == "schacon@gmail.com"
    end

    context("test reading the oid of a tag") do
        id = Oid("0c37a5391bbff43c37f0d0371823a5509eed5b1d")
        obj = lookup(test_repo, id)
        @test Oid(target(obj)) == Oid("5b5b025afb0b4c913b4c338a42934a3863bf3644")
    end
end

with_tmp_repo_access("test writing a tag") do test_repo, path
   sig = Signature("Julia", "julia@julia.com")
   tid = tag!(test_repo,
              name="tag",
              message="test tag message\n",
              target=Oid("5b5b025afb0b4c913b4c338a42934a3863bf3644"),
              tagger=sig)
   t = lookup(test_repo, tid)
   @test isa(t, GitTag)
   @test target_id(t) == Oid("5b5b025afb0b4c913b4c338a42934a3863bf3644")
   @test target_id(t) == Oid(target(t))
   @test message(t) == "test tag message\n"
   @test name(tagger(t)) == "Julia"
   @test email(tagger(t)) == "julia@julia.com"
end

with_tmp_repo_access("test writing a tag without a signature") do test_repo, path
    testname = "Julia"
    testemail = "julia@julia.com"
    GitConfig(test_repo)["user.name"] = testname
    GitConfig(test_repo)["user.email"] = testemail
    id = tag!(test_repo,
              name="tag",
              message="test tag message\n",
              target=Oid("5b5b025afb0b4c913b4c338a42934a3863bf3644"))
    @test isa(id, Oid)
    t = lookup(test_repo, id)
    @test name(tagger(t)) == testname
    @test email(tagger(t)) == testemail
end

with_tmp_repo_access("test invalid message type") do test_repo, path
    sig = Signature("Julia", "julia@julia.com")
    @test_throws TypeError tag!(
                              test_repo,
                              name="tag",
                              message=:error,
                              target=Oid("5b5b025afb0b4c913b4c338a42934a3863bf3644"),
                              tagger=sig)
end

with_tmp_repo_access("test writing light tags") do test_repo, path
    tag!(test_repo, name="tag", target=Oid("5b5b025afb0b4c913b4c338a42934a3863bf3644"))
    t = lookup_ref(test_repo, "refs/tags/tag")
    @test target(t) == Oid("5b5b025afb0b4c913b4c338a42934a3863bf3644")
end


text:
mutable struct FilterDesc
    ptr::Cptr

    function FilterDesc(::Type{T}, dims::Dims{N}) where {T,N}
        ref = Ref{Cptr}()
        @cudnn :cudnnCreateFilterDescriptor (Ptr{Cptr},) ref
        desc = new(ref[])
        push!(ALLOCATED, desc)
        finalizer(desc) do x
            @cudnn :cudnnDestroyFilterDescriptor (Cptr,) x.ptr
        end

        csize = Cint[reverse(dims)...]
        @cudnn(:cudnnSetFilterNdDescriptor,
            (Cptr,Cint,Cint,Cint,Ptr{Cint}),
            desc, datatype(T), CUDNN_TENSOR_NCHW, N, csize)
        desc
    end
end
FilterDesc(::Type{T}, dims::Int...) where T = FilterDesc(T, dims)
FilterDesc(x::CuArray) = FilterDesc(eltype(x), size(x))

Base.cconvert(::Type{Cptr}, desc::FilterDesc) = desc.ptr


text:
# This file produces a figure that shows how particles perform uncertainty propagation, and compares the result to uncertainty propagation through linearization like Measurements.jl does.

using MonteCarloMeasurements, StatsPlots, NNlib, Measurements, KernelDensity
using Measurements: value, uncertainty
default(lab="")
N    = 20 # Number of particles
f    = x -> σ(12x-6) # Nonlinear function
l    = 0 # Left boundary
r    = 1 # Right boundary
d    = Normal(0.5, 0.15) # The probability density of the input
m    = Measurements.:(±)(d.μ, d.σ) # For comparison to Measurements.jl using linear uncertainty propagation
my   = f(m) # output measurement
dm   = Normal(value(my), uncertainty(my)) # Output density according to Measurements
x    = Particles(N,d, permute=false) # Create particles distributed according to d, sort for visualization
y    = f(x).particles # corresponding output particles
x    = x.particles # extract vector to plot manually
xr   = LinRange(l,r,100) # x values for plotting
noll = zeros(N)

plot(f, l, r, legend=:right, xlims=(l,r), ylims=(l,r), axis=false, grid=false, lab="f(x)", xlabel="Input space", ylabel="Output space")
plot!(x->0.2pdf(d,x),l,r, lab="Input dens.")
# Estimate the true output density using a large sample
kdt = kde(f.(rand(d,100000)), npoints=200, bandwidth=0.08)
plot!(l .+ 0.2kdt.density, kdt.x, lab="True output dens.")

# This is the output density as approximated by linear uncertainty propagation
plot!(l .+ 0.2pdf.(Ref(dm),xr), xr, lab="Linear Gaussian propagation")

# Estimate the output density corresponding to the particles
kd = kde(y, npoints=200, bandwidth=0.08)
plot!(l .+ 0.2kd.density, kd.x, lab="Particle kernel dens. est.", l=:dash)

# Draw helper lines that show how particles are transformed from input space to output space
plot!([x x][1:2:end,:]', [noll y][1:2:end,:]', l=(:black, :arrow, :dash, 0.1))
plot!([x fill(l,N).+0.02][1:2:end,:]', [y y][1:2:end,:]', l=(:black, :arrow, :dash, 0.1))

# Plot the particles
scatter!(x, 0y, lab="Input particles")
scatter!(fill(l,N) .+ 0.02, y, lab="Output particles")

# Draw mean lines, these show hoe the mean is transformed using linear uncertainty propagation
plot!([d.μ,d.μ], [0,f(d.μ)], l=(:red, :dash, 0.2))
plot!([l,d.μ], [f(d.μ),f(d.μ)], l=(:red, :dash, 0.2))


text:
module test
using ForwardDiff
include("extra_file.jl")

export my_f,my_derivative

end


text:
println("Testing method dispatch...")

# LinearGaussianSSM
#	- LinearKalmanFilter
@assert method_exists(predict, (LinearGaussianSSM, AbstractMvNormal))
@assert method_exists(update, (LinearGaussianSSM, AbstractMvNormal, Vector))
@assert method_exists(update,
	(LinearGaussianSSM, AbstractMvNormal, Vector, Int))
@assert method_exists(update,
	(LinearGaussianSSM, AbstractMvNormal, Vector, KalmanFilter, Int))

# NonlinearGaussianSSM
#	- NonlinearKalmanFilter
@assert method_exists(update,
	(NonlinearGaussianSSM, AbstractMvNormal, Vector))

for T in [NonlinearKalmanFilter, ExtendedKalmanFilter, UnscentedKalmanFilter]
	println(T)
	@assert method_exists(update,
		(NonlinearGaussianSSM, AbstractMvNormal, Vector, T))
	@assert method_exists(update,
		(NonlinearGaussianSSM, AbstractMvNormal, Vector, T, Int))
end

# Special predict methods that don't return just a Distribution
@assert method_exists(predict, (NonlinearGaussianSSM, AbstractMvNormal, UKF))
@assert method_exists(predict, (NonlinearGaussianSSM, Matrix, EnKF))

@assert method_exists(update,
	(NonlinearGaussianSSM, Matrix, Vector, EnKF))

# NonlinearSSM 			NonlinearFilter

println("Dispatch passed.\n")


text:
using Printf
using SpecialFunctions

include("constants.jl")
include("alias.jl")
include("Atoms.jl")
include("utils.jl")
include("PGBF.jl")
include("CGBF.jl")
include("overlap.jl")
include("kinetic.jl")
include("nuclear.jl")
include("sto3g.jl")
include("BasisSet.jl")

function test_a_terms()
    @assert Aterm(0,0,0,0,0,0.,0.,0.,0.) == 1.0
    @assert Aarray(0,0,0.,0.,0.,1.) == [1.0]
    @assert Aarray(0,1,1.,1.,1.,1.) == [1.0, -1.0]
    @assert Aarray(1,1,1.,1.,1.,1.) == [1.5, -2.5, 1.0]
    @assert Aterm(0,0,0,0,0,0.,0.,0.,1.) == 1.0
    @assert Aterm(0,0,0,0,1,1.,1.,1.,1.) == 1.0
    @assert Aterm(1,0,0,0,1,1.,1.,1.,1.) == -1.0
    @assert Aterm(0,0,0,1,1,1.,1.,1.,1.) == 1.0
    @assert Aterm(1,0,0,1,1,1.,1.,1.,1.) == -2.0
    @assert Aterm(2,0,0,1,1,1.,1.,1.,1.) == 1.0
    @assert Aterm(2,0,1,1,1,1.,1.,1.,1.) == -0.5
    @assert Aterm(2,1,0,1,1,1.,1.,1.,1.) == 0.5

    @printf("test_a_terms passed\n")
end

test_a_terms()

function test_gamma()
    # gammainc test functions. Test values taken from Mathematica
    # println("a=0.5 test")
    @assert maximum([gammainc(0.5,float(x)) for x in 0:10]
        -Any[0, 1.49365, 1.69181, 1.7471, 1.76416, 1.76968,
          1.77151, 1.77213, 1.77234, 1.77241, 1.77244]) < 1e-5

    # println("a=1.5 test")
    @assert maximum([gammainc(1.5,float(x)) for x in 0:10]
        -Any[0, 1.49365, 1.69181, 1.7471, 1.76416, 1.76968,
          1.77151, 1.77213, 1.77234, 1.77241, 1.77244]) < 1e-5
    # println("a=2.5 test")
    @assert maximum([gammainc(2.5,float(x)) for x in 0:10]
        -Any[0, 0.200538, 0.59898, 0.922271, 1.12165, 1.22933,
          1.2831, 1.30859, 1.32024, 1.32542, 1.32768]) < 1e-5
    @printf("test_gamma passed\n")
end

test_gamma()

function test_na()
    s = PGBF(1.0)

    c = init_CGBF(0.0,0.0,0.0)
    push!(c,1.0,1.0)

    println(s)
    println(c)

    Vab = nuclear_attraction( s,s, (0.0,0.0,0.0) )
    @printf("Vab = %18.10f\n", Vab)

    Vab = nuclear_attraction( c,c, (0.0,0.0,0.0) )
    @printf("Vab = %18.10f\n", Vab)
end

#test_na()

function test_fgamma()
    for (x,res) in Any[(0.,1),
            (30.,0.161802),
            (60.,0.1144114),
            (90.,0.0934165),
            (120.,0.08090108),
            (300.,0.051166336)]
        #@assert isapprox(res,Fgamma(0,x))
        @printf("res, Fgamma(0,x) %18.10f %18.10f\n", x, Fgamma(0,x))
    end
    @printf "test_fgamma is passed\n"
end

#test_fgamma()

#todo make into a test
function test_one()
    s1 = PGBF(1.0)
    s2 = PGBF(1.0, (0.0,1.0,0.0) )
    x = y = 0.0
    S = overlap(s1,s2)
    T = kinetic(s1,s2)
    @printf("S = %18.10f\n",S)
    @printf("T = %18.10f\n",T)
    for z in linspace(0,1,5)
        Vab = nuclear_attraction(s1,s2,(x,y,z))
        @printf("z = %18.10f, Vab = %18.10f\n", z, Vab)
    end
end

#test_one()

function test_LiH()
    LiH = init_atoms_xyz("LiH.xyz")
    bfs = build_basis(LiH)

    println(LiH)
    Nbasis = length(bfs)
    for i = 1:Nbasis
        println("--------")
        println("Basis #", i)
        println("--------")
        println(bfs[i])
    end

    Lis1 = bfs[1]
    Lis2 = bfs[2]
    Lipx = bfs[3]
    Lipy = bfs[4]
    Lipz = bfs[5]
    H1s  = bfs[6]
    Vab = nuclear_attraction(Lis1,Lis1,LiH)
    println("Vab = ", Vab)
    #@assert isapprox(nuclear_attraction(Lis1,Lis1,LiH),-8.307532656)
end

test_LiH()

"""
function test_geo_basis()

  @printf("\nCalling test_geo_basis\n")

  E1     = nuclear_repulsion(h2)
  E1_ref = 0.7223600367
  @printf("H2: nuclear_repulsion %18.10f, err = %18.10f\n", E1, abs(E1_ref-E1))

  @assert nel(h2) == 2
  @assert nel(h2o) == 10
  @assert length(sto3g)==10

  bfs = build_basis(h2)
  @assert length(bfs.bfs)==2

  l,r = bfs.bfs
  @assert isapprox(overlap(l,l),1)
  @assert isapprox(overlap(r,r),1)

  #@assert isapprox(overlap(l,r),0.66473625)
  ovl1 = overlap(l,r)
  ovl1_ref = abs(ovl1 - 0.66473625)
  @printf("overlap(l,r), err %18.10f, %18.10f\n", ovl1, ovl1_ref)

  @assert isapprox(kinetic(l,l),0.76003188)
  @assert isapprox(kinetic(r,r),0.76003188)
  @assert isapprox(kinetic(l,r),0.24141861181119084)
  @assert isapprox(coulomb(l,l,l,l), 0.7746059439196398)
  @assert isapprox(coulomb(r,r,r,r), 0.7746059439196398)
  @assert isapprox(coulomb(l,l,r,r), 0.5727937653511646)
  @assert isapprox(coulomb(l,l,l,r), 0.4488373301593464)
  @assert isapprox(coulomb(l,r,l,r), 0.3025451156654606)
  bfs = build_basis(h2o)

  s1,s2,px,py,pz,hl,hr = bfs.bfs
  @assert isapprox(coulomb(s1,s2,hl,hr),0.03855344493645537)
  @assert isapprox(coulomb(s1,pz,hl,hr),-0.0027720110485359053)
  @assert isapprox(coulomb(s1,hl,pz,hr),-0.010049491284827426)
  @assert coulomb(s1,py,hl,hr)==0
  @assert coulomb(s1,hl,py,hr)==0
end
"""


text:
using MNIST

include("utils/optimizers.jl")
include("models/gaussMMVAE.jl")

function elbo(params, x, prior)
  # prop through autoencoder
  x_recon, posterior, z = fprop(params, x)

  # likelihood term
  expected_nll = sum(-x .* log(x_recon) - (1-x) .* log(1-x_recon), 2)

  # kl divergence term
  kld = computeKLD(prior, posterior, z)

  return mean(expected_nll + kld)
end

function trainVAE(data, params, hyperParams)
  N,d = size(data)
  nBatches = div(N,hyperParams["batchSize"])

  for epoch_idx in 1:hyperParams["nEpochs"]
    elbo_tracker = 0.
    for batch_idx in 1:nBatches

      # get minibatch
      x = data[(batch_idx-1)*hyperParams["batchSize"]+1:batch_idx*hyperParams["batchSize"],:]

      # compute elbo
      elbo_tracker += elbo(params, x, hyperParams["prior"])

      # get elbo gradients
      #grads = elbo_grad(params, x, hyperParams["prior"])

      # perform AdaM update
      #params = adamUpdate(params, grads, hyperParams["adamParams"])

    end
    @printf "Epoch %d. Neg. ELBO: %.3f \n" epoch_idx elbo_tracker/nBatches
  end

  return params
end

function main()
  # load MNIST (values in [0,255])
  data = transpose(traindata()[1])

  # shuffle, normalized, and reduce dataset
  shuffle(vec(data))
  data /= 255
  data = data[1:50000,:]

  # set architecture parameters
  hidden_size = 500
  latent_size = 25
  n_components = 2
  vae_params = init_params(size(data,2), hidden_size, latent_size, n_components)

  # set hyperparams
  adamParams = Dict("lr"=>0.0001, "m"=>init_params(size(data,2), hidden_size, latent_size, n_components, 0.),
                    "v"=>init_params(size(data,2), hidden_size, latent_size, n_components, 0.), "t"=>0)
  hyperParams = Dict("adamParams"=>adamParams, "prior"=>Dict("weights"=>transpose([.5, .5]), "mu"=>[-2., 2.], "sigma"=>[1., 1.]), "nEpochs"=>50, "batchSize"=>100)

  final_vae_params = trainVAE(data, vae_params, hyperParams)
end

main()


text:
@testset "test StopAfterNoImprovement" begin
    env = StateTransformedEnv(CartPoleEnv{Float32}();state_mapping=deepcopy)
    policy = RandomPolicy(action_space(env))

    total_reward_per_episode = TotalRewardPerEpisode()
    patience = 30
    stop_condition =
        StopAfterNoImprovement(() -> total_reward_per_episode.reward, patience, 0.0f0)

    # stop_condition is called between POST_ACT_STAGE & POST_EPISODE_STAGE.
    # total_reward_per_episode.rewards is updated at POST_EPISODE_STAGE.
    # total_reward_per_episode.reward is updated at POST_ACT_STAGE.
    # so the latter one should be used. or the value is from the previous episode.
    hook = ComposedHook(total_reward_per_episode)
    run(policy, env, stop_condition, hook)

    @test argmax(total_reward_per_episode.rewards) != patience
end

@testset "StopAfterNSeconds" begin
    s = StopAfterNSeconds(0.01)
    @test !s()
    sleep(0.02)
    @test s()
end


text:
"""
    get_group_taxa_at_node(nodes::Array{Dict{String,Any}}, node_num::Int64)

Gets the sets of taxa for each group at a node.

# Arguments
- `nodes::Array{Dict{String,Any}}`: list of nodes.
- `node_num::Int64`: current node index.
"""
function get_group_taxa_at_node(nodes::Array{Dict{String,Any}}, node_num::Int64)

    #Initialize varibales
    last_group = Array{String}(undef, 0)
    groups = Array{Array{String}}(undef, 0)
    num_groups = nodes[node_num]["Groups"]

    # Iterate through each descendent
    for group in nodes[node_num]["Descendents"]

        # Add taxa taxa to group
        group_taxa = nodes[group]["Taxa"]
        push!(groups, group_taxa)
    end
    if length(groups) < num_groups
        for group in groups
            last_group = setdiff(nodes[node_num]["Taxa"], group)
        end
        if length(groups) == 0
            last_group = nodes[node_num]["Taxa"]
        end
        for idx in 1:num_groups-length(groups)
            if length(last_group) >= idx
                push!(groups, [last_group[idx]])
            end
        end
    end
    return groups
end

"""
    get_group_combos(group_taxa::Array{Array{String}})

Gets all the combinations of group vs non groups.

# Arguments
- `group_taxa::Array{Array{String}}`: list of taxa within a group.
"""
function get_group_combos(group_taxa::Array{Array{String}})
    combos = Array{Dict{String,Array{String}}}(undef, 0)
    for (idx, group) in enumerate(group_taxa)
        non_group = Array{String}(undef, 0)
        group = group
        non_group_groups = group_taxa[1:end .!= idx]
        for non_group_group in non_group_groups
            append!(non_group, non_group_group)
        end
        push!(combos, Dict("group" => group, "Non_group" => non_group))
    end
    return combos
end

"""
    get_sPu_and_sPr(nodes::Array{Dict{String,Any}}, node_num::Int64,
    taxa_labels::Dict{String,String}, character_labels::Dict{String,String})

Gets all the sPu and sPr for the entire character sequence at a specific node.

# Arguments
- `nodes::Array{Dict{String,Any}}`: list of nodes.
- `node_num::Int64`: current node index.
- `taxa_labels::Dict{String,String}`: a mapping of the taxa labels to the character labels.
- `character_labels::Dict{String,String}`: a mapping of the character labels to the corresponding sequences.
"""
function get_sPu_and_sPr(nodes::Array{Dict{String,Any}}, node_num::Int64, taxa_labels::Dict{String,String}, character_labels::Dict{String,String}; protein::Bool=false)

    # Initialize variables
    group_taxa = get_group_taxa_at_node(nodes, node_num)
    group_combos = get_group_combos(group_taxa)
    sPu = Array{Dict{String,Any}}(undef, 0)
    sPr = Array{Dict{String,Any}}(undef, 0)
    letter_transformations = Dict{Char,Vector{Char}}('A'=>['A'], 'T'=>['T'], 'C'=>['C'], 'G'=>['G'], 'U'=>['U'], 'R'=>['A','G'], 'Y'=>['C','T'], 'S'=>['G','C'], 'W'=>['A','T'], 'K'=>['G','T'], 'M'=>['A','C'], 'B'=>['C','G','T'], 'D'=>['A','G','T'], 'H'=>['A','C','T'], 'V'=>['A','C','G'], 'N'=>['A','T','C','G'], '-'=>['-'])

    # Iterate through all possible group combinations
    for (group_idx, group_combo) in enumerate(group_combos)
        group_taxa = group_combo["group"]
        non_group_taxa = group_combo["Non_group"]
        push!(sPu, Dict{String,Any}("Group_Taxa" => group_taxa, "sPu" => Dict{Int,Tuple{Char,Int}}(), "Num_Non_Group" => length(non_group_taxa)))
        push!(sPr, Dict{String,Any}("Group_Taxa" => group_taxa, "sPr" => Dict{Int,Array{Tuple{Char,Int}}}(), "Num_Non_Group" => length(non_group_taxa)))

        if length(group_taxa) > 0
            # Iterate through each letter in the character sequences
            for idx in 1:length(character_labels[taxa_labels[group_taxa[1]]])

                # Only continue if the non group does not contain any N's at the current index
                non_group_letters = unique([character_labels[taxa_labels[new_character]][idx] for new_character in non_group_taxa])
                if !(('N' in non_group_letters) | ('A' in non_group_letters && 'T' in non_group_letters && 'C' in non_group_letters && 'G' in non_group_letters))
                    # Iterate through each taxon in group
                    for (taxa_idx, taxa) in enumerate(group_taxa)

                        # Get the letter at the index
                        curr_letter = character_labels[taxa_labels[taxa]][idx]

                        # Get the transformation of the current letter
                        if protein
                            curr_letters = curr_letter
                        else
                            curr_letters = letter_transformations[curr_letter]
                        end

                        # Iterate through each letter
                        for letter in curr_letters
                            # If letter is not a gap and does not exist in sPu
                            valid = letter != '-' && !(haskey(sPu[group_idx]["sPu"], idx))

                            # If exists in sPr, check if new letter
                            if haskey(sPr[group_idx]["sPr"], idx)
                                used_sPr = [tup[1] for tup in sPr[group_idx]["sPr"][idx]]
                                valid = !(letter in used_sPr)
                            end
                            if valid
                                pure = true
                                match = false

                                # Iterate through the character of each member of the other groups, to check if candidate CA
                                for new_letter in non_group_letters

                                    if protein
                                        new_letters = new_letter
                                    else
                                        new_letters = letter_transformations[new_letter]
                                    end
                                    match = letter in new_letters
                                    if match
                                        break
                                    end
                                end

                                # If candidate CA, iterate through the character of each member of the same group to determine if pure or private
                                if !match
                                    occurances = 1
                                    for new_character in group_taxa[1:end .!=taxa_idx]
                                        new_letter = character_labels[taxa_labels[new_character]][idx]
                                        if protein
                                            new_letters = new_letter
                                        else
                                            new_letters = letter_transformations[new_letter]
                                        end
                                        match = letter in new_letters
                                        if match
                                            occurances += 1
                                        else
                                            pure = false
                                        end
                                    end

                                    # Add the single pure or private characteristic if found
                                    if pure
                                        sPu[group_idx]["sPu"][idx] = (letter, occurances)
                                    else
                                        if !haskey(sPr[group_idx]["sPr"], idx)
                                            sPr[group_idx]["sPr"][idx] = Array{Tuple{Char,Int}}(undef, 0)
                                        end
                                        push!(sPr[group_idx]["sPr"][idx], (letter, occurances))
                                    end
                                end
                            end
                        end

                    end
                end
            end
        end
    end
    return sPu, sPr
end

# Function to get all the cPu and cPr for the entire character sequence at a specific node ----- function currently not supported for nucleotide options
"""
    get_cPu_and_cPr(nodes::Array{Dict{String,Any}}, node_num::Int64, taxa_labels::Dict{String,String},
    character_labels::Dict{String,String}, sPu::Array{Dict{String,Any}}, sPr::Array{Dict{String,Any}})

Gets all the cPu and cPr for the entire character sequence at a specific node (does not support nucleotide options).

# Arguments
- `nodes::Array{Dict{String,Any}}`: list of nodes.
- `node_num::Int64`: current node index.
- `taxa_labels::Dict{String,String}`: a mapping of the taxa labels to the character labels.
- `character_labels::Dict{String,String}`: a mapping of the character labels to the corresponding sequences.
- `sPu::Array{Dict{String,Any}}`: list of simple pure rules.
- `sPr::Array{Dict{String,Any}}`: list of simple private rules.
"""
function get_cPu_and_cPr(nodes::Array{Dict{String,Any}}, node_num::Int64, taxa_labels::Dict{String,String}, character_labels::Dict{String,String}, sPu::Array{Dict{String,Any}}, sPr::Array{Dict{String,Any}})

    # Figure out which indices have already been matched to an sPu or sPr
    exclude = Vector{Int64}()
    for dict in sPu
        append!(exclude, [key for key in keys(dict["sPu"])])
    end
    for dict in sPr
        append!(exclude, [key for key in keys(dict["sPr"])])
    end
    all_idx = range(1, length=length(character_labels[taxa_labels["1"]]))
    curr_idx = setdiff(all_idx, exclude)

    # Initialize variables
    group_taxa = get_group_taxa_at_node(nodes, node_num)
    group_combos = get_group_combos(group_taxa)
    cPu = Array{Dict{String,Any}}(undef, 0)
    cPr = Array{Dict{String,Any}}(undef, 0)
    num_char = length(character_labels[taxa_labels["1"]])
    letter_transformations = Dict{Char,Vector{Char}}('A'=>['A'], 'T'=>['T'], 'C'=>['C'], 'G'=>['G'], 'U'=>['U'], 'R'=>['A','G'], 'Y'=>['C','T'], 'S'=>['G','C'], 'W'=>['A','T'], 'K'=>['G','T'], 'M'=>['A','C'], 'B'=>['C','G','T'], 'D'=>['A','G','T'], 'H'=>['A','C','T'], 'V'=>['A','C','G'], 'N'=>['A','T','C','G'], '-'=>['-'])

    # Iterate through all possible group combinations
    for (group_idx, group_combo) in enumerate(group_combos)
        group_taxa = group_combo["group"]
        non_group_taxa = group_combo["Non_group"]
        push!(cPu, Dict{String,Any}("Group_Taxa" => group_taxa, "cPu" => Dict{Tuple{Int,Int},Tuple{Tuple{Char,Char},Int}}(), "Num_Non_Group" => length(non_group_taxa)))
        push!(cPr, Dict{String,Any}("Group_Taxa" => group_taxa, "cPr" => Dict{Tuple{Int,Int},Array{Tuple{Tuple{Char,Char},Int}}}(), "Num_Non_Group" => length(non_group_taxa)))

        # Iterate through each letter in the character sequences
        for (iter_idx,idx1) in enumerate(curr_idx)

            # Iterate through each letter that occurs after the first letter to get all combinations
            for idx2 in curr_idx[iter_idx:end]

                # Iterate through each taxon in the group
                for (taxa_idx, taxa) in enumerate(group_taxa)

                    # Extract the two letters of interest
                    letter1 = character_labels[taxa_labels[taxa]][idx1]
                    letter2 = character_labels[taxa_labels[taxa]][idx2]

                    # If letter is not a gap and does not exist in sPu
                    valid = letter1 != '-' && letter2 != '-' && !(haskey(cPu[group_idx]["cPu"], (idx1,idx2)))

                    # If exists in sPr, check if new letter
                    if haskey(cPr[group_idx]["cPr"], (idx1,idx2))
                        used_cPr = [tup[1][1] for tup in cPr[group_idx]["cPr"][(idx1,idx2)]]
                        valid = !((letter1, letter2) in used_cPr)
                    end

                    # If the letter combination is valid, continue check for valid CA
                    if valid
                        pure = true
                        match = false

                        # Iterate through the character of each member of the other groups, to check if candidate CA
                        for new_character in non_group_taxa
                            new_letter1 = character_labels[taxa_labels[new_character]][idx1]
                            new_letter2 = character_labels[taxa_labels[new_character]][idx2]
                            match = (new_letter1 == letter1 && new_letter2 == letter2)
                            if match
                                break
                            end
                        end

                        # If candidate CA, iterate through the character of each member of the same group to determine if pure or private
                        if !match
                            occurances = 1
                            for new_character in group_taxa[1:end .!=taxa_idx]
                                new_letter1 = character_labels[taxa_labels[new_character]][idx1]
                                new_letter2 = character_labels[taxa_labels[new_character]][idx2]
                                match = (new_letter1 == letter1 && new_letter2 == letter2)
                                if match
                                    occurances += 1
                                else
                                    pure = false
                                end
                            end

                            # Add the complex pure or private characteristic if found
                            if pure
                                cPu[group_idx]["cPu"][(idx1,idx2)] = ((letter1,letter2), occurances)
                            else
                                if !haskey(cPr[group_idx]["cPr"], (idx1,idx2))
                                    cPr[group_idx]["cPr"][(idx1,idx2)] = Array{Tuple{Tuple{Char,Char},Int}}(undef, 0)
                                end
                                push!(cPr[group_idx]["cPr"][(idx1,idx2)], ((letter1,letter2), occurances))
                            end
                        end
                    end
                end
            end
        end
    end
    return cPu,cPr
end


text:
#######################################################################
# randomkets
#######################################################################
"""
    randomkets(n, a = r -> cis(2pi*rand()); kw...)

Create a lazy collection of `n` `KetModel`s of amplitude `a`, and `normalization = 1/√n`.
Other keyword arguments are forwarded to `ketmodel`.

This type of ket model is useful e.g. in stochastic trace evaluation of KPM methods, where
the amplitude is chosen as a random function with `⟨a⟩ = 0`, `⟨aa⟩ = 0` and `⟨a'a⟩ = 1`. The
default `a` is a uniform random phase on each site, which satisties these conditions. Then,
the normalized trace of an operator `A` can be estimated with `Tr[A]/N₀ ≈ ∑⟨ket|A|ket⟩`,
where the sum is taken over the `n` random kets `|ket⟩` of norm `1/√n` produced by
`randomkets`, and `N₀` is the total number of orbitals in the full unit cell.

To apply it to a multiorbital system with a maximum of `N` orbitals per site, `a` must in
general be adapted to produce the desired random `SVector{N}` (unless `maporbitals = true`),
with the above statistical properties for each orbital. Example: to have independent,
complex, normally-distributed random components of two orbitals use `randomkets(n, r ->
randn(SVector{2,ComplexF64}))`, or alternatively `randomkets(n, r -> randn(ComplexF64),
maporbitals = true)`.

# See also
    `ket`
"""
randomkets(n::Int, a = r -> cis(2pi*rand()); kw...) =
    Iterators.repeated(ketmodel(a; normalization = 1/√n, kw...), n)

#######################################################################
# basiskets
#######################################################################
"""
    basiskets(a = I; kw...)

Create a multicolumn ket model that represents a basis for sites selected by
`siteselctor(kw...)`, with amplitude `a` on each site. For hamiltonians with `N` orbitals
per site, `a` will need to be either `I` or a matrix with `N` columns.

`basiskets(a; kw...)` is equivalent to `ketmodel(a; singlesitekets = true, kw...)`.

# See also
    `ket`
"""
basiskets(a = I; kw...) = ketmodel(a; singlesitekets = true, kw...)

#######################################################################
# Kernel Polynomial Method : momenta
#######################################################################
using Base.Threads

struct MomentaKPM{T,B<:Tuple}
    mulist::Vector{T}
    bandbracket::B
end

struct KPMBuilder{AM,HM<:AbstractMatrix,T,H<:Hamiltonian,K<:Ket,B}
    h::H
    hmat::HM
    Amat::AM
    ket0::K
    ket1::K
    ket2::K
    bandbracket::Tuple{B,B}
    order::Int
    mulist::Vector{T}
end

function KPMBuilder(h::Hamiltonian{<:Lattice,L}, A, ket_or_ketmodels, order, bandrange, flat) where {L}
    iszero(L) ||
        throw(ArgumentError("Hamiltonian is defined on an infinite lattice. Reduce it to zero-dimensions with `wrap` or `unitcell`."))
    eh = eltype(eltype(h))
    eA = eltype(eltype(A))
    mulist = zeros(promote_type(eh, eA), order + 1)
    bandbracket = bandbracketKPM(h, bandrange)
    hmat = matrix_KPM(flat, h)
    Amat = matrix_KPM(flat, A, orbitalstructure(h))
    zeroket = maybe_flatten_KPM(flat, firstket_KPM(ket_or_ketmodels, h))
    builder = KPMBuilder(h, hmat, Amat, zeroket, similar(zeroket), similar(zeroket), bandbracket, order, mulist)
    return builder
end

firstket_KPM(k::Ket, h) =  k
firstket_KPM(k::KetModel, h) = ket(k, h)
firstket_KPM(kets, h) = firstket_KPM(first(kets), h)

maybe_flatten_KPM(::Val{true}, k) = flatten(k)
maybe_flatten_KPM(::Val{false}, k) = k
maybe_flatten_KPM(flat, k) = flat ? flatten(k) : k

matrix_KPM(::Val{true}, h::Hamiltonian, x...) = bloch!(similarmatrix(h, flatten), h)
matrix_KPM(::Val{false}, h::Hamiltonian, x...) = bloch!(similarmatrix(h), h)
matrix_KPM(flat, h::Hamiltonian, x...) = bloch!(flat ? similarmatrix(h, flatten) : similarmatrix(h), h)
matrix_KPM(::Val{true}, A::AbstractArray, o::OrbitalStructure) = flatten(A, o)
matrix_KPM(::Val{false}, A::AbstractArray, ::OrbitalStructure) = A
matrix_KPM(flat, A::AbstractArray, ::OrbitalStructure) = flat ? flatten(A, o) : A
matrix_KPM(flat, A, x...) = A

"""
    momentaKPM(h::Hamiltonian, A = I; ket = randomkets(1), order = 10, bandrange = missing, flat = Val(true))

Compute the Kernel Polynomial Method (KPM) momenta `μₙ = ⟨k|Tₙ(h) A|k⟩`, where `|k⟩ =
ket(ket::KetModel, h)` or `|k⟩ = ket::Ket` (depending on the type of `ket`), `A` is an
observable (`Hamiltonian` or `AbstractMatrix`) and `Tₙ(h)` is the order-`n` Chebyshev
polynomial of the Hamiltonian `h`.

`ket` can be a single `KetModel` or `Ket`, as above, or a collection of them, as in the
default `ket = randomkets(n)`. In the latter case, `μₙ` is summed over all models/kets, `μₙ
= ∑_k ⟨k|Tₙ(h) A|k⟩`. If `ket::Ket` is a multi-column ket, a sum over columns will similarly
be performed. A `ket = randomkets(n)` produces a lazy collection of `n` random `KetModel`s
that is useful to estimate momenta of normalized traces using the stochastic trace approach,
whereby `μ_n = Tr[A T_n(h)]/N₀ ≈ ∑ₖ⟨k|A T_n(h)|k⟩`. Here the `|k⟩`s are `n` random kets of
norm `1/√n` and `N₀` is the total number of orbitals per unit cell of `h` (see
`randomkets`).

The order of the Chebyshev expansion is `order`. The `bandbrange = (ϵmin, ϵmax)` should
completely encompass the full bandwidth of `hamiltonian`. If `missing` it is computed
automatically using `ArnoldiMethod` (must be loaded `using ArnoldiMethod`). `flat` indicates
whether, in the case of multiorbital systems, the internal computations are to be performed
using flattened arrays, typically increasing performace by making use of external
linear algebra libraries (e.g. MKL or OpenBLAS).

# Examples

```
julia> h = LatticePresets.cubic() |> hamiltonian(hopping(1)) |> unitcell(region = RegionPresets.sphere(10));

julia> momentaKPM(h, bandrange = (-6,6)).mulist |> length
11
```
"""
function momentaKPM(h::Hamiltonian, A = I; ket = randomkets(1), order = 10, bandrange = missing, flat = Val(false))
    builder = KPMBuilder(h, A, ket, order, bandrange, flat)
    momentaKPM!(builder, ket)
    jackson!(builder.mulist)
    return MomentaKPM(builder.mulist, builder.bandbracket)
end

function momentaKPM!(b::KPMBuilder, ket::Ket)
    pmeter = Progress(b.order, "Computing moments: ")
    addmomentaKPM!(b, pmeter)
    return nothing
end

momentaKPM!(b::KPMBuilder, model::KetModel) = momentaKPM!(b, ket!(b.ket0, model, b.h))

momentaKPM!(b::KPMBuilder, models) = foreach(model -> momentaKPM!(b, model), models)

# This iterates bras <psi_n| = <psi_0|AT_n(h) instead of kets (faster CSC multiplication)
# In practice we iterate their conjugate |psi_n> = T_n(h') A'|psi_0>, and do the projection
# onto the start ket, |psi_0>
function addmomentaKPM!(b::KPMBuilder{<:AbstractMatrix,<:AbstractSparseMatrix}, pmeter)
    mulist = b.mulist
    kmat, kmat0, kmat1 = parent.((b.ket0, b.ket1, b.ket2))
    h, A, bandbracket = b.hmat, b.Amat, b.bandbracket
    order = length(mulist) - 1
    mul!(kmat0, A', kmat)
    mulscaled!(kmat1, h', kmat0, bandbracket)
    mulist[1] += proj(kmat0, kmat)
    mulist[2] += proj(kmat1, kmat)
    for n in 3:(order+1)
        ProgressMeter.next!(pmeter; showvalues = ())
        iterateKPM!(kmat0, h', kmat1, bandbracket)
        mulist[n] += proj(kmat0, kmat)
        kmat0, kmat1 = kmat1, kmat0
    end
    return mulist
end

function addmomentaKPM!(b::KPMBuilder{<:UniformScaling, <:AbstractSparseMatrix,T}, pmeter) where {T}
    mulist = b.mulist
    kmat, kmat0, kmat1 = parent.((b.ket0, b.ket1, b.ket2))
    h, A, bandbracket = b.hmat, b.Amat, b.bandbracket
    order = length(mulist) - 1
    kmat0 .= kmat
    mulscaled!(kmat1, h', kmat0, bandbracket)
    mulist[1] += μ0 = proj(kmat0, kmat0)
    mulist[2] += μ1 = proj(kmat1, kmat0)
    # This is not used in the currently activated codepath (BLAS mul!), but is needed in the
    # commented out @threads codepath
    thread_buffers = (zeros(T, Threads.nthreads()), zeros(T, Threads.nthreads()))
    for n in 3:2:(order+1)
        ProgressMeter.next!(pmeter; showvalues = ())
        ProgressMeter.next!(pmeter; showvalues = ()) # twice because of 2-step
        proj11, proj10 = iterateKPM!(kmat0, h', kmat1, bandbracket, thread_buffers)
        mulist[n] += 2 * proj11 - μ0
        n + 1 > order + 1 && break
        mulist[n + 1] += 2 * proj10 - μ1
        kmat0, kmat1 = kmat1, kmat0
    end
    A.λ ≈ 1 || (mulist .*= A.λ)
    return mulist
end

function mulscaled!(y, h´, x, (center, halfwidth))
    mul!(y, h´, x)
    invhalfwidth = 1/halfwidth
    @. y = (y - center * x) * invhalfwidth
    return y
end

function iterateKPM!(kmat0, h´, kmat1, (center, halfwidth), buff = ())
    α = 2/halfwidth
    β = 2center/halfwidth
    mul!(kmat0, h´, kmat1, α, -1)
    @. kmat0 = kmat0 - β * kmat1
    return proj_or_nothing(buff, kmat0, kmat1)
end

proj_or_nothing(::Tuple{}, kmat0, kmat1) = nothing
proj_or_nothing(buff, kmat0, kmat1) = (proj(kmat1, kmat1), proj(kmat1, kmat0))

# This is equivalent to tr(kmat1'*kmat2) for matrices
proj(kmat1, kmat2) = dot(kmat1, kmat2)

# function iterateKPM!(ket0, h´, ket1, (center, halfwidth), thread_buffers = ())
#     h = parent(h´)
#     nz = nonzeros(h)
#     rv = rowvals(h)
#     α = -2 * center / halfwidth
#     β = 2 / halfwidth
#     reset_buffers!(thread_buffers)
#     @threads for row in 1:size(ket0, 1)
#         ptrs = nzrange(h, row)
#         @inbounds for col in 1:size(ket0, 2)
#             k1 = ket1[row, col]
#             tmp = α * k1 - ket0[row, col]
#             for ptr in ptrs
#                 tmp += β * adjoint(nz[ptr]) * ket1[rv[ptr], col]
#             end
#             # |k0⟩ → (⟨k1|2h - ⟨k0|)' = 2h'|k1⟩ - |k0⟩
#             ket0[row, col] = tmp
#             update_buffers!(thread_buffers, k1, tmp)
#         end
#     end
#     return sum_buffers(thread_buffers)
# end

# reset_buffers!(::Tuple{}) = nothing
# function reset_buffers!((q, q´))
#     fill!(q, zero(eltype(q)))
#     fill!(q´, zero(eltype(q´)))
#     return nothing
# end

# @inline update_buffers!(::Tuple{}, k1, tmp) = nothing
# @inline function update_buffers!((q, q´), k1, tmp)
#     q[threadid()]  += dot(k1, k1)
#     q´[threadid()] += dot(tmp, k1)
#     return nothing
# end

# @inline sum_buffers(::Tuple{}) = nothing
# @inline sum_buffers((q, q´)) = (sum(q), sum(q´))

function jackson!(μ::AbstractVector)
    order = length(μ) - 1
    @inbounds for n in eachindex(μ)
        μ[n] *= ((order - n + 1) * cos(π * n / (order + 1)) +
                sin(π * n / (order + 1)) * cot(π / (order + 1))) / (order + 1)
    end
    return μ
end

function bandbracketKPM(h, ::Missing)
    bandbracketKPM(h, bandrangeKPM(h))
end

bandbracketKPM(h, (ϵmin, ϵmax)::Tuple{T,T}, pad = float(T)(0.01)) where {T} = ((ϵmax + ϵmin) / 2, (ϵmax - ϵmin) / (2 - pad))

bandrangeKPM(h::Hamiltonian; kw...) = bandrangeKPM(matrix_KPM(Val(true), h); kw...)

function bandrangeKPM(h::AbstractMatrix{T}; quiet = false) where {T}
    if quiet == true
        nothing
    else @warn "Computing spectrum bounds... Consider using the `bandrange` option for faster performance."
    end
    checkloaded(:ArnoldiMethod)
    R = real(T)
    decompl, _ = Main.ArnoldiMethod.partialschur(h, nev=1, tol=1e-4, which = Main.ArnoldiMethod.LR());
    decomps, _ = Main.ArnoldiMethod.partialschur(h, nev=1, tol=1e-4, which = Main.ArnoldiMethod.SR());
    ϵmax = R(real(decompl.eigenvalues[1]))
    ϵmin = R(real(decomps.eigenvalues[1]))
    if quiet == true nothing
    else @warn("Computed bandrange = ($ϵmin, $ϵmax)") end
    return (ϵmin, ϵmax)
end

#######################################################################
# Kernel Polynomial Method : observables
#######################################################################
"""
    dosKPM(h::Hamiltonian; resolution = 2, ket = randomkets(1), kw...)

Compute, using the Kernel Polynomial Method (KPM), the local density of states `ρₖ(ϵ) =
⟨k|δ(ϵ-h)|k⟩` for a ket `|k⟩ = ket(ket::KetModel, h)` or `|k⟩ = ket::Ket` (depending on the
type of `ket`). The result is a tuple of energy points `ϵᵢ::Vector` spanning the band range,
and real `ρₖ(ϵᵢ)::Vector` values (any residual imaginary part in `ρₖ` is dropped). The
number of energy points `ϵᵢ` is `order * resolution`, rounded to the closest integer.

If `ket` is a collection of `KetModel`s, the sum `∑ₖρₖ(ε)` over all models will be computed.
In the case of the default `ket = randomkets(n)`, this results in an estimate of the total
density of states per orbital, computed through an stochastic trace, `ρ(ϵ) =
∑ₖ⟨k|δ(ϵ-h)|k⟩/n ≈ Tr[δ(ϵ-h)]/N₀`, where `N₀` is the total number of orbitals in the unit
cell.

`dosKPM` is a particular case of `densityKPM` for an operator `A = I` and with any residual
imaginary parts dropped

    dosKPM(μ::MomentaKPM; resolution = 2)

Same as above with KPM momenta `μ` as input.

# See also
    `momentaKPM`, `densityKPM`, `averageKPM`
"""
dosKPM(h; resolution = 2, ket = randomkets(1), kw...) =
    dosKPM(momentaKPM(h, I; ket = ket, kw...), resolution = resolution)

dosKPM(μ::MomentaKPM; resolution = 2) = real.(densityKPM(μ; resolution = resolution))

"""
    densityKPM(h::Hamiltonian, A; resolution = 2, ket = randomkets(1), kw...)

Compute, using the Kernel Polynomial Method (KPM), the spectral density of `A`, `ρᴬₖ(ϵ) =
⟨k|A δ(ϵ-h)|k⟩` for a ket `|k⟩ = ket(ket::KetModel, h)` or `|k⟩ = ket::Ket` (depending on
the type of `ket`). The result is a tuple of energy points `ϵᵢ::Vector` spanning the band
range, and real `ρᴬₖ(ϵᵢ)::Vector` values. The number of energy points `ϵᵢ` is `order *
resolution`, rounded to the closest integer.

If `ket` is a collection of `KetModel`s, the sum `∑ₖρᴬₖ(ε)` over all models will be computed. In
the case of the default `ket = randomkets(n)`, this results in an estimate of the average
spectral density per orbital, computed through an stochastic trace, `ρᴬ(ϵ) =
∑ₖ⟨k|δ(ϵ-h)A|k⟩/n ≈ Tr[δ(ϵ-h)A]/N₀`, where `N₀` is the total number of orbitals in the unit
cell.

    densityKPM(μ::MomentaKPM; resolution = 2)

Same as above with KPM momenta `μ` as input.

# See also
    `dosKPM`, `momentaKPM`, `averageKPM`
"""
densityKPM(h, A; resolution = 2, kw...) =
    densityKPM(momentaKPM(h, A; kw...); resolution = resolution)

function densityKPM(momenta::MomentaKPM{T}; resolution = 2) where {T}
    checkloaded(:FFTW)
    (center, halfwidth) = momenta.bandbracket
    numpoints = round(Int, length(momenta.mulist) * resolution)
    ρlist = zeros(T, numpoints)
    copyto!(ρlist, momenta.mulist)
    Main.FFTW.r2r!(ρlist, Main.FFTW.REDFT01, 1)  # DCT-III in FFTW
    xk = [cos(π * (k + 0.5) / numpoints) for k in 0:numpoints - 1]
    @. ρlist = center + halfwidth * ρlist / (π * sqrt(1.0 - xk^2))
    @. xk = center + halfwidth * xk
    return xk, ρlist
end

"""
    averageKPM(h::Hamiltonian, A; kBT = 0, Ef = 0, kw...)

Compute, using the Kernel Polynomial Method (KPM), the thermal expectation value `⟨A⟩ = Σ_k
f(E_k) ⟨k|A|k⟩ = ∫dE f(E) Tr [A δ(E-H)]/N₀ = Tr [A f(H)]/N₀` for a given hermitian operator
`A` and a zero-dimensional hamiltonian `h` with a total of `N₀` orbitals (see `momentaKPM`
and its options `kw` for further details). `f(E)` is the Fermi-Dirac distribution function,
`|k⟩` are `h` eigenstates with energy `E_k`, kBT` is the temperature in energy units and
`Ef` the Fermi energy.

    averageKPM(μ::MomentaKPM, A; kBT = 0, Ef = 0)

Same as above with the KPM momenta as input (see `momentaKPM`).

# See also
    `dosKPM`, `momentaKPM`, `averageKPM`
"""
averageKPM(h, A; kBT = 0, Ef = 0, kw...) = averageKPM(momentaKPM(h, A; kw...); kBT = kBT, Ef = Ef)

function averageKPM(momenta::MomentaKPM; kBT = 0.0, Ef = 0.0)
    (center, halfwidth) = momenta.bandbracket
    order = length(momenta.mulist) - 1
    # if !iszero(kBT)
    #     @warn "Finite temperature requires numerical evaluation of the integrals"
    #     checkloaded(:QuadGK)
    # end
    average = sum(n -> momenta.mulist[n + 1] * fermicheby(n, Ef, kBT, center, halfwidth), 0:order)
    return average
end

# Pending issue: Unexpected behaviour with center != 0.
function fermicheby(n, Ef, kBT, center, halfwidth)
    kBT´ = kBT / halfwidth;
    Ef´ = (Ef - center) / halfwidth;
    T = typeof(Ef´)
    if kBT´ == 0
        int = n == 0 ? 0.5+asin(Ef´)/π : -2.0*sin(n*acos(Ef´))/(n*π)
    else
        throw(error("Finite temperature not yet implemented"))
        # η = 1e-10
        # int = Main.QuadGK.quadgk(E´ -> _intfermi(n, E´, Ef´, kBT´), -1.0+η, 1.0-η, atol= 1e-10, rtol=1e-10)[1]
    end
    return T(int)
end

_intfermi(n, E´, Ef´, kBT´) = fermifun(E´, Ef´, kBT´) * 2/(π*(1-E´^2)^(1/2)) * chebypol(n, E´) / (1+(n == 0 ? 1 : 0))

fermifun(E´, Ef´, kBT´) = kBT´ == 0 ? (E´<Ef´ ? 1.0 : 0.0) : (1/(1+exp((E´-Ef´)/(kBT´))))

function chebypol(m::Int, x::T) where {T<:Number}
    cheby0 = one(T)
    cheby1 = x
    if m == 0
        chebym = cheby0
    elseif m == 1
        chebym = cheby1
    else
        for i in 2:m
            chebym = 2x * cheby1 - cheby0
            cheby0, cheby1 = cheby1, chebym
        end
    end
    return chebym
end

max_stars_repo_name: pablosanjose/Quantica

max_issues_repo_name: pablosanjose/Quantica

max_forks_repo_name: pablosanjose/Quantica


text:
module NonlinearBandits

using Flux:
    relu, Chain, Dense, gpu, @epochs, train!, ADAM, DataLoader, throttle, params, cpu
using Distributions: Uniform, Normal, loggamma, InverseGamma, MvNormal
using GLMNet: glmnet
using LinearAlgebra: Hermitian, diagm, logdet
using Random: randperm
using Suppressor: @suppress

include("utils.jl")
include("BayesLM.jl")
include("BayesPM.jl")
include("PartitionedBayesPM.jl")
include("bandits.jl")
include("contexts.jl")
include("rewards.jl")
include("metrics.jl")
include("drivers.jl")
include("RandomPolicy.jl")
include("PolynomialThompsonSampling.jl")
include("NeuralLinear.jl")
include("NearestNeighboursUCB.jl")

"""
    fit!(model, X::AbstractMatrix, y::AbstractMatrix)

Update the parameters of `model`.

# Arguments

- `X::AbstractMatrix`: A matrix with observations stored as columns.
- `y::AbstractMatrix`: A matrix with 1 row of response variables.
"""
function fit!(model, X::AbstractMatrix, y::AbstractMatrix)
    throw(ErrorException("must implement fit! for $(typeof(model))"))
end

"""
    shape_scale(model::AbstractBayesianLM)

Return the shape/scale of model.
"""
function shape_scale(model::AbstractBayesianLM)
    throw(ErrorException("must implement shape_scale for $(typeof(model))"))
end

"""
    update!(pol::AbstractPolicy, X::AbstractMatrix, a::AbstractVector{<:Int},
            r::AbstractMatrix)

Update `pol` with a batch of data.
"""
function update!(
    pol::AbstractPolicy, X::AbstractMatrix, a::AbstractVector{<:Int}, r::AbstractMatrix
)
    throw(ErrorException("must implement update! for $(typeof(pol))"))
end

export AbstractBayesianLM,
    AbstractContextSampler,
    AbstractDriver,
    AbstractMetric,
    AbstractPolicy,
    AbstractRewardSampler,
    append_data!,
    arm_data,
    BanditDataset,
    BayesLM,
    BayesPM,
    expand,
    fit!,
    FunctionalRegret,
    GaussianRewards,
    Index,
    lasso_selection,
    LatentDriver,
    locate,
    nearest_neighbours,
    NearestNeighboursUCB,
    NeuralEncoder,
    NeuralLinear,
    Partition,
    PartitionedBayesPM,
    PolynomialThompsonSampling,
    RandomPolicy,
    run!,
    shape_scale,
    split!,
    StandardDriver,
    std,
    tpbasis,
    UniformContexts,
    update!

end


text:
##
## The non-flaming sheep.
##
"""
    function sheep()

The non-flaming sheep.

# Example
```jldoctest
julia> cowsay("Baa", cow=Cowsay.sheep)
 _____
< Baa >
 -----
  \\
   \\
       __
      UooU\\.'@@@@@@`.
      \\__/(@@@@@@@@@@)
           (@@@@@@@@)
           `YY~~~~YY'
            ||    ||
```
"""
function sheep(;eyes="oo", tongue="  ", thoughts="\\")
the_cow =
"""
  $thoughts
   $thoughts
       __
      U$(eyes)U\\.'@@@@@@`.
      \\__/(@@@@@@@@@@)
           (@@@@@@@@)
           `YY~~~~YY'
            ||    ||
"""
return the_cow
end


text:
function digit_conversion(digit::Integer, one::Char, five::Char, ten::Char)
    if digit == 9
        return one*ten
    end
    if 9 > digit >= 5
        digit = digit % 5
        return five*one^digit
    end
    if digit == 4
        return one*five
    end
    if digit < 4
        return one^digit
    end
end

function digit_conversion(digit::Integer, one::Char)
    if digit >= 4
        error("Input number exceeds maximum.")
    end

    return one^digit
end

function to_roman(number)
    if number <= 0
        error("Input number must be larger than 0.")
    end

    output = ""

    thousand = div(number, 1000)
    output *= digit_conversion(thousand, 'M')

    hundred = div(number % 1000, 100)
    output *= digit_conversion(hundred, 'C', 'D', 'M')

    ten = div(number % 100, 10)
    output *= digit_conversion(ten, 'X', 'L', 'C')

    one = number % 10
    output *= digit_conversion(one, 'I', 'V', 'X')

end

to_roman(1024)

max_stars_repo_name: jfcliang/jl-exercism-solutions

max_issues_repo_name: jfcliang/jl-exercism-solutions

max_forks_repo_name: jfcliang/jl-exercism-solutions


text:
using Test
using CartesianGeneticProgramming
import YAML
using Statistics

@testset "Processing" begin
    cfg = get_config("test.yaml"; functions=["f_avg", "f_max", "f_min", "f_med"])
    ind = CGPInd(cfg)

    # test that f(0, 0, 0, 0) = 0
    inputs = zeros(cfg.n_in)
    set_inputs(ind, inputs)
    for i in 1:cfg.n_in
        @test ind.buffer[i] == 0.0
    end
    output = process(ind)
    @test output[1] == 0.0
    for i in eachindex(ind.nodes)
        if ind.nodes[i].active
            @test ind.buffer[i] == 0.0
        end
    end

    inputs = rand(cfg.n_in)

    functions = [
        (CGPFunctions.f_avg, mean),
        (CGPFunctions.f_max, maximum),
        (CGPFunctions.f_min, minimum),
        (CGPFunctions.f_med, median)
    ]

    for f in functions
        ind.nodes[cfg.n_in + 1] = CartesianGeneticProgramming.Node(1, cfg.n_in, f[1], Float64[], true)
        ind.outputs[1] = cfg.n_in + 1
        output = process(ind, inputs)
        @test output[1] == f[2](inputs[1:cfg.n_in])
    end
end


text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

function OpticSim.Vis.draw!(scene::Makie.LScene, h::ParaxialAnalysis.HeadEye.Head; draw_head=true, kwargs...) where {T<:Real}
    if (draw_head)
        scale_factor = 60.0;
        face = ParaxialAnalysis.HeadEye.transform_mesh(ParaxialAnalysis.HeadEye.get_face_model(), ParaxialAnalysis.HeadEye.tr(h)*Geometry.scale(scale_factor, scale_factor, scale_factor))
        Makie.mesh!(scene, face, color=RGBA(1.0,0.2,0.2, 0.1), transparency=true)
    end

    ParaxialAnalysis.HeadEye.draw_local_frame(scene, ParaxialAnalysis.HeadEye.tr(h), "Head Transform")

    for e in ParaxialAnalysis.HeadEye.eyes(h)
        Vis.draw!(scene, e, parent_transform = ParaxialAnalysis.HeadEye.tr(h))
    end
end

function OpticSim.Vis.draw!(scene::Makie.LScene, e::ParaxialAnalysis.HeadEye.Eye; parent_transform::Geometry.Transform = Transform(), kwargs...) where {T<:Real}

    ParaxialAnalysis.HeadEye.draw_local_frame(scene, parent_transform*ParaxialAnalysis.HeadEye.tr(e), ParaxialAnalysis.HeadEye.text(e), scale=2)
    Vis.draw!(scene, ParaxialAnalysis.HeadEye.pupil(e), parent_transform = parent_transform*ParaxialAnalysis.HeadEye.tr(e))
end

function OpticSim.Vis.draw!(scene::Makie.LScene, p::ParaxialAnalysis.HeadEye.Pupil; parent_transform::Geometry.Transform = Transform(), kwargs...) where {T<:Real}

    ParaxialAnalysis.HeadEye.draw_local_frame(scene, parent_transform*ParaxialAnalysis.HeadEye.tr(p), "pupil", scale=1.0)

    # half_size = size(p) / 2.0
    # eli = Ellipse(half_size, half_size, rotation(parent_transform)*unitZ3(), parent_transform*origin(tr(p)))
    # Vis.draw!(scene, eli, transparency=true, color=RGBA(0.2, 0.2, 0.2, 0.4))
end


text:
using Distributions
using Lora
using PGUManifoldMC

DATADIR = "../../data"
SUBDATADIR = "mala"

nchains = 10
nmcmc = 110000
nburnin = 10000

function C(n::Int, c::Float64)
  X = eye(n)
  [(j <= n-i) ? X[i+j, i] = X[i, i+j] = c^j : nothing for i = 1:(n-1), j = 1:(n-1)]
  X
end

n = 20
μ = zeros(n)
Σ = C(n, 0.9)
ν = 30.

Σt = (ν-2)*Σ/ν
Σtinv = inv(Σt)

function plogtarget(p::Vector{Float64}, v::Vector)
  hdf = 0.5*ν
  hdim = 0.5*n
  shdfhdim = hdf+hdim
  v = lgamma(shdfhdim)-lgamma(hdf)-hdim*log(ν)-hdim*log(pi)-0.5*logdet(Σt)
  z = p-μ
  v-shdfhdim*log1p(dot(z, Σtinv*z)/ν)
end

sampler = MALA(0.05)

mcrange = BasicMCRange(nsteps=nmcmc, burnin=nburnin)

outopts = Dict{Symbol, Any}(:monitor=>[:value], :diagnostics=>[:accept])

times = Array(Float64, nchains)
stepsizes = Array(Float64, nchains)
i = 1

while i <= nchains
  v0 = Dict(:p=>rand(Normal(0, 2), n))

  p = BasicContMuvParameter(
    :p,
    logtarget=plogtarget,
    nkeys=1,
    autodiff=:reverse,
    init=Any[(:p, v0[:p]), (:v, Any[v0[:p]])],
    order=1
  )

  model = likelihood_model([p], isindexed=false)

  job = BasicMCJob(
    model,
    sampler,
    mcrange,
    v0,
    tuner=AcceptanceRateMCTuner(0.574, score=x -> logistic_rate_score(x, 3.), verbose=false),
    outopts=outopts
  )

  tic()
  run(job)
  runtime = toc()

  chain = output(job)
  ratio = acceptance(chain)

  if 0.5 < ratio < 0.65
    writedlm(joinpath(DATADIR, SUBDATADIR, "chain"*lpad(string(i), 2, 0)*".csv"), chain.value, ',')
    writedlm(joinpath(DATADIR, SUBDATADIR, "diagnostics"*lpad(string(i), 2, 0)*".csv"), vec(chain.diagnosticvalues), ',')

    times[i] = runtime
    stepsizes[i] = job.sstate.tune.step

    println("Iteration ", i, " of ", nchains, " completed with acceptance ratio ", ratio)
    i += 1
  end
end

writedlm(joinpath(DATADIR, SUBDATADIR, "times.csv"), times, ',')
writedlm(joinpath(DATADIR, SUBDATADIR, "stepsizes.csv"), stepsizes, ',')


text:
function saveplot(pl::py.SyncPlot,format::Symbol,outPath::AbstractString; js=:remote)
    filename = string(outPath,'.',format)
    if format==:html
        py.savehtml(pl, filename, js)
    elseif format in [:png, :pdf, :svg]
        py.savefig(pl, filename)
    else
        error("Unknown image format $format")
    end
    filename
end

saveplot(pl::py.SyncPlot,formats::AbstractArray{Symbol},outPath::AbstractString;kwargs...) = map(fmt->saveplot(pl,fmt,outPath;kwargs...), formats)


text:
struct SteadyStateAdjointSensitivityFunction{C<:AdjointDiffCache,Alg<:SteadyStateAdjoint,uType,SType,fType<:ODEFunction,CV,λType,VJPType} <: SensitivityFunction
  diffcache::C
  sensealg::Alg
  discrete::Bool
  y::uType
  sol::SType
  f::fType
  colorvec::CV
  λ::λType
  vjp::VJPType
#  linsolve::LS
end

function SteadyStateAdjointSensitivityFunction(g,sensealg,discrete,sol,dg,colorvec)
  @unpack f, p, u0 = sol.prob

  diffcache, y = adjointdiffcache(g,sensealg,discrete,sol,dg,f;quad=false)

  λ = zero(y) # solution of f_x^T λ = g_x^T, Eq. (2) in  https://math.mit.edu/~stevenj/18.336/adjoint.pdf
  #linsolve = sensealg.linsolve(Val{:init},diffcache.uf,y)
  vjp = similar(λ, length(p))
  SteadyStateAdjointSensitivityFunction(diffcache,sensealg,discrete,y,sol,f,colorvec,λ,vjp)
end

@noinline function SteadyStateAdjointProblem(sol,sensealg::SteadyStateAdjoint,g,dg;save_idxs=nothing)
  @unpack f, p = sol.prob

  discrete = false

  p === DiffEqBase.NullParameters() && error("Your model does not have parameters, and thus it is impossible to calculate the derivative of the solution with respect to the parameters. Your model must have parameters to use parameter sensitivity calculations!")

  sense = SteadyStateAdjointSensitivityFunction(g,sensealg,discrete,sol,dg,f.colorvec)
  @unpack diffcache, y, sol, λ, vjp = sense
  if DiffEqBase.has_jac(f)
     f.jac(diffcache.J,y,p,nothing)
  else
     if DiffEqBase.isinplace(sol.prob)
       jacobian!(diffcache.J, diffcache.uf, y, diffcache.f_cache, sensealg, diffcache.jac_config)
     else
        temp = jacobian(diffcache.uf, y, sensealg)
        @. diffcache.J = temp
     end
  end

  _save_idxs = save_idxs === nothing ? Colon() : save_idxs
  if dg != nothing
    if g!= nothing
      dg(vec(diffcache.dg_val),y,p,nothing,nothing)
    else
      if typeof(_save_idxs) <: Number
        diffcache.dg_val[_save_idxs] = dg[_save_idxs]
      elseif typeof(dg) <: Number
        @. diffcache.dg_val[_save_idxs] = dg
      else
        @. diffcache.dg_val[_save_idxs] = dg[_save_idxs]
      end
    end
  else
    if g != nothing
      gradient!(vec(diffcache.dg_val),diffcache.g,y,sensealg,diffcache.g_grad_config)
    end
  end

  λ .= diffcache.J'\vec(diffcache.dg_val') # use linsolve here
  vecjacobian!(vec(diffcache.dg_val), y, λ, p, nothing, sense, dgrad=vjp, dy=nothing)

  if g != nothing
    # compute del g/del p
    dg_dp_val = zero(p)
    dg_dp = ParamGradientWrapper(g,nothing,y)
    dg_dp_config = build_grad_config(sensealg,dg_dp,p,y)
    gradient!(dg_dp_val,dg_dp,p,sensealg,dg_dp_config)

    @. dg_dp_val = dg_dp_val - vjp
    return dg_dp_val

  else
    return -vjp
  end
end


text:
include("Density.jl")

function calc_design_matrix(umbrella_centers, sim_all_datas, sigma_matrix)
    design_matrix = zeros(Float64, size(umbrella_centers, 1) * size(sim_all_datas[1], 1), size(umbrella_centers, 1))

    # i, jのループで全シミュレーションデータ点にアクセス
    for i = 1:size(umbrella_centers, 1)
        for j = 1:size(sim_all_datas[i], 1)
            for k = 1:size(umbrella_centers, 1)
                diff = calc_gaussian(sim_all_datas[i][j, :], umbrella_centers[k, :], sigma_matrix)
                     - calc_gaussian(umbrella_centers[i, :], umbrella_centers[k, :], sigma_matrix)

                design_matrix[j + size(sim_all_datas[i], 1) * (i - 1), k] = diff
            end
        end
    end

    return design_matrix
end