text:
using LinearAlgebra

"""
    overlaps(a::HilbertSpace, b::HilbertSpace)

Calculate the overlaps between the basis vectors of each Hilbert space.
"""
function overlaps(a::HilbertSpace, b::HilbertSpace)
    @assert a.weighting == b.weighting
    LHS = a.data
    RHS = b.weighting * b.data
    return [dot(LHS[:,i],RHS[:,j]) for i in 1:size(a.data,2), j in 1:size(b.data,2)]
end

"""
    unitary_approx(M::AbstractArray)

Calculate the best unitary approxmation of `M` using singular value
decomposition.
"""
function unitary_approx(M::AbstractArray)
    F = svd(M)
    return F.U * F.Vt
end

"""
    unitary_overlaps(a::HilbertSpace, b::HilbertSpace)

Shortcut for unitary_approx(overlaps(a, b))
"""
function unitary_overlaps(a::HilbertSpace, b::HilbertSpace)
    unitary_overlaps = overlaps(a, b)
    return unitary_approx(unitary_overlaps)
end

# for FDFD solver
function overlaps(a::HilbertSpace_FDFD, b::HilbertSpace_FDFD)
    @assert a.weighting == b.weighting
    LHS = a.data
    RHS = b.weighting .* b.data
    return [dot(LHS[:,i],RHS[:,j]) for i in 1:size(a.data,2), j in 1:size(b.data,2)]
end

function unitary_overlaps(a::HilbertSpace_FDFD, b::HilbertSpace_FDFD)
    unitary_overlaps = overlaps(a, b)
    return Peacock.unitary_approx(unitary_overlaps)
end

"""
    wilson_matrix(spaces::Array{HilbertSpace,1}; closed::Bool=true)

Calculate the Wilson loop matrix through the Hilbert spaces.

The `closed` keyword will assert that `spaces[1] == spaces[end]`.
Otherwise, it will be assumed that the Wilson loop begins and finishes
at the same `k0`, but in different Brillouin zones.
"""
function wilson_matrix(spaces::Array{HilbertSpace,1}; closed::Bool=true)
    @assert length(spaces) > 1
    # Assume all spaces have the same plane wave expansion
    basis = spaces[1].basis
    for space in spaces
        @assert space.basis == basis
    end
    # However, k0 is still changing...
    if closed
        # In a closed Wilson loop...
        @assert spaces[1] == spaces[end]
    else
        # If not closed, then the loop begins and finishes at
        # the same k0 but in different Brillouin zones
        delta_k0 = spaces[end].k0 - spaces[1].k0
        b1 = spaces[1].basis.b1
        b2 = spaces[1].basis.b2
        B = [b1 b2]
        dp, dq = B \ delta_k0
        # Round to integers if approximately correct
        if isapprox(round(dp), dp, atol=1e-6) && isapprox(round(dq), dq, atol=1e-6)
            dp = Int(round(dp))
            dq = Int(round(dq))
        else
            @show dp, dq
            throw("First and last spaces do not... (TODO)")
        end
        # Duplicate spaces[1] and shift its k0 so that it matches spaces[end]
        space_1_at_end = shift_k0(spaces[1], dp, dq)
        spaces = [spaces; space_1_at_end]
    end
    W = I
    for (a,b) in zip(spaces, spaces[2:end])
        W = W * unitary_overlaps(a, b)
    end
    @assert W * W' ≈ I
    return W
end

"""
    wilson_eigvals(spaces::AbstractArray{HilbertSpace,1}; closed=true)

Return the eigenvalues of [`wilson_matrix`](@ref), sorted by phase angle.
"""
function wilson_eigvals(spaces::AbstractArray{HilbertSpace,1}; closed=true)
    W = wilson_matrix(spaces, closed=closed)
    vals = eigvals(W)
    return sort(vals, by=angle)
end

"""
    wilson_eigvals(spaces::AbstractArray{HilbertSpace,1}; closed=true)

Return the eigenvalues and eigenvectors of [`wilson_matrix`](@ref),
sorted by the phase angle of the eigenvalues.
"""
function wilson_eigen(spaces::AbstractArray{HilbertSpace,1}; closed=true)
    W = wilson_matrix(spaces, closed=closed)
    vals, vecs = eigen(W)
    idx = sortperm(vals, by=angle)
    vals = vals[idx]
    vecs = vecs[:,idx]
    vecs = unitary_approx(vecs) # or other orthonormalisation?
    return vals, vecs
end

"""
    wilson_gauge(spaces::AbstractArray{HilbertSpace,1}; closed=true)

Return the eigenvalues, eigenvectors, and gauge of the Wilson loop through
the Hilbert space, sorted by the phase angle of the eigenvalues.
"""
function wilson_gauge(spaces::AbstractArray{HilbertSpace,1}; closed=true)
    vals, vecs = wilson_eigen(spaces, closed=closed)
    gauge = copy(spaces)
    gauge[1] = HilbertSpace(gauge[1].k0, gauge[1].data*vecs, gauge[1].weighting, gauge[1].basis)
    for i in 1:length(gauge)-1
        a = gauge[i]
        b = gauge[i+1]
        mixing = unitary_overlaps(b, a)
        gauge[i+1] = HilbertSpace(b.k0, b.data*mixing, b.weighting, b.basis)
    end
    return vals, vecs, gauge
end

"""
    plot_wilson_loop_winding(solver::Solver, ks, polarisation, bands::AbstractVector{<:Int}, <keyword arguments>)

Perform a series of Wilson loops along `ks`, and plot the spectra on a band diagram.

Keyword arguments
- `dk_outer=nothing`: maximum distance between each loop (resolution of the scan)
- `dk_inner=nothing`: maximum distance between points along a loop (resolution of the loop)
- `delta_brillouin_zone=BrillouinZoneCoordinate(0,1)`: each Wilson loop starts at and finishes in at the same `k` in different Brillouin zones
- `labels=[]`: overwrite the labels for each `k` in `ks`
- `markersize=nothing`: overwrite the size of each point
"""
function plot_wilson_loop_winding(solver::Solver, ks, polarisation, bands::AbstractVector{<:Int};
                dk_outer=nothing, dk_inner=nothing, delta_brillouin_zone=BrillouinZoneCoordinate(0,1),
                labels=[], markersize=nothing)
    # Convert BrillouinZoneCoordinate to labelled positions in k space
    if labels == []
        labels = [hasproperty(x,:label) ? x.label : "" for x in ks]
    end
    ks = [typeof(x)==BrillouinZoneCoordinate ? get_k(x,solver.basis) : x for x in ks]
    # Function to return the Wilson inner eigenvalues
    function my_solve(k)
        spaces = HilbertSpace[]
        delta_k = get_k(delta_brillouin_zone, solver.basis)
        ks_inner = [k, k+delta_k]
        ks_inner, _ = sample_path(ks_inner, dk=dk_inner)
        for k_inner in ks_inner
            modes = solve(solver, k_inner, polarisation)
            space = HilbertSpace(modes[bands])
            push!(spaces, space)
        end
        vals = wilson_eigvals(spaces, closed=false)
        angles = angle.(vals)
        angles = [angles.-2pi angles angles.+2pi]
        return angles
    end
    plot_band_diagram(my_solve, ks, dk=dk_outer, labels=labels, show_vlines=false, markersize=markersize)
    ylim(-pi,pi)
    yticks((-1:1)*pi, labels=["-π","0","+π"])
    ylabel("Wilson spectrum")
end

function plot_wilson_loop_winding(solver::Peacock.FDFD.Solver, ks, polarisation, bands::AbstractVector{<:Int};
    dk_outer=nothing, dk_inner=nothing, delta_brillouin_zone=BrillouinZoneCoordinate(0,1),
    labels=[], markersize=nothing)
    if labels == []
        labels = [hasproperty(x,:label) ? x.label : "" for x in ks]
    end
    ks = [typeof(x)==BrillouinZoneCoordinate ? get_k(x,solver.basis) : x for x in ks]
    function my_solve(k)
        spaces = HilbertSpace_FDFD[]
        delta_k = get_k(delta_brillouin_zone, solver.basis)
        ks_inner = [k, k+delta_k]
        ks_inner, _ = sample_path(ks_inner, dk=dk_inner)
        for k in ks_inner
            modes = Peacock.FDFD.solve(solver, k, polarisation)
            space = HilbertSpace_FDFD(modes[bands])
            push!(spaces, space)
        end
        spaces = [spaces; spaces[1]]
        W = I
        for (a, b) in zip(spaces, spaces[2:end])
            W = W * unitary_overlaps(a, b)
        end
        @assert W * W' ≈ I
        vals = eigvals(W)
        vals = sort(vals, by=angle)
        return angle.(vals)
    end
    plot_band_diagram(my_solve, ks, dk=dk_outer, labels=labels, show_vlines=false, markersize=markersize)
    ylim(-pi,pi)
    yticks((-1:1)*pi, labels=["-π","0","+π"])
    ylabel("Wilson spectrum")
end


text:
using DataStructures

#####################
# Nodes and Graphs
#####################

abstract type Node end
abstract type DagNode <: Node end
abstract type TreeNode <: DagNode end

const DiGraph = AbstractVector{<:Node}
const Dag = AbstractVector{<:DagNode}
const Tree = AbstractVector{<:TreeNode}

Base.eltype(::Type{DiGraph}) = Node
Base.eltype(::Type{Dag}) = DagNode
Base.eltype(::Type{Tree}) = TreeNode

#####################
# traits
#####################

"""
A trait hierarchy denoting types of nodes
`NodeType` defines an orthogonal type hierarchy of node types, so we can dispatch on node type regardless of the graph type.
See @ref{https://docs.julialang.org/en/v1/manual/methods/#Trait-based-dispatch-1}
"""

abstract type NodeType end

struct Leaf <: NodeType end
struct Inner <: NodeType end

@inline NodeType(instance::Node) = NodeType(typeof(instance))

#####################
# required methods
#####################

"Get the children of a given inner node"
@inline children(n::Node)::Vector{<:Node} = children(NodeType(n), n)
@inline children(::Inner, n::Node)::Vector{<:Node} = error("Each inner node should implement a `children` method; one is missing for $(typeof(n))")

"Does the node have children?"
@inline has_children(n::Node)::Bool = has_children(NodeType(n), n)
@inline has_children(::Inner, n::Node)::Bool = !isempty(children(n))
@inline has_children(::Leaf, n::Node)::Bool = false

"Get the number of children of a given inner node"
@inline num_children(n::Node)::Int = num_children(NodeType(n), n)
@inline num_children(::Inner, n::Node)::Int = length(children(n))
@inline num_children(::Leaf, n::Node)::Int = 0

#####################
# traversal
#####################

"Flip the bit field throughout this circuit (or ensure it is set to given value)"
function flip_bit(node::DagNode, ::Val{Bit} = Val(!node.bit)) where Bit
    if node.bit != Bit
        node.bit = Bit
        if isinner(node)
            for c in children(node)
                flip_bit(c, Val(Bit))
            end
        end
    end
    nothing # returning nothing helps save some allocations and time
end

import Base.foreach #extend

"Apply a function to each node in a circuit, bottom up"
function foreach(f::Function, node::DagNode)
    foreach_rec(f, node)
    flip_bit(node)
    nothing # returning nothing helps save some allocations and time
end

function foreach(node::Union{DagNode,Dag}, f_leaf::Function, f_inner::Function)
    foreach(node) do n
        isinner(n) ? f_inner(n) : f_leaf(n)
    end
    nothing # returning nothing helps save some allocations and time
end

function foreach_rec(f::Function, node::DagNode, ::Val{Bit} = Val(!node.bit)) where Bit
    if node.bit != Bit
        node.bit = Bit
        if isinner(node)
            for c in children(node)
                foreach_rec(f, c, Val(Bit))
            end
        end
        f(node)
    end
    nothing # returning nothing helps save some allocations and time
end

function foreach(f::Function, node::TreeNode)
    if isinner(node)
        for c in children(node)
            foreach(f, c)
        end
    end
    f(node)
    nothing
end

import Base.filter #extend

"""Retrieve list of nodes in circuit matching predicate `p`"""
function filter(p::Function, root::DagNode, ::Type{T} = Union{})::Vector where T
    results = Vector{T}()
    foreach(root) do n
        if p(n)
            if !(n isa eltype(results))
                results = copy_with_eltype(results, typejoin(eltype(results),typeof(n)))
            end
            push!(results,n)
        end
    end
    results
end

"""
Compute a function bottom-up on the circuit.
`f_leaf` is called on leaf nodes, and `f_inner` is called on inner nodes.
Values of type `T` are passed up the circuit and given to `f_inner` as a function on the children.
"""
function foldup(node::Dag, f_leaf::Function, f_inner::Function, ::Type{T})::T where {T}
    foldup(node[end], f_leaf, f_inner, T)
end

function foldup(node::DagNode, f_leaf::Function, f_inner::Function, ::Type{T})::T where {T}
    @assert node.bit == false
    v = foldup_rec(node, f_leaf, f_inner, T)
    flip_bit(node)
    v
end

function foldup_rec(node::DagNode, f_leaf::Function, f_inner::Function,
                    ::Type{T}, ::Val{Bit} = Val(!node.bit))::T where {T,Bit}
    if node.bit == Bit
        return node.data::T
    else
        node.bit = Bit
        v = if isinner(node)
            callback(c) = (foldup_rec(c, f_leaf, f_inner, T, Val(Bit))::T)
            f_inner(node, callback)::T
        else
            f_leaf(node)::T
        end
        node.data = v
        return v
    end
end

function foldup(node::TreeNode, f_leaf::Function, f_inner::Function, ::Type{T})::T where T
    v = if isinner(node)
        callback(c) = (foldup(c, f_leaf, f_inner, T)::T)
        f_inner(node, callback)::T
    else
        f_leaf(node)::T
    end
    return v
end

"""
Compute a function bottom-up on the circuit.
`f_leaf` is called on leaf nodes, and `f_inner` is called on inner nodes.
Values of type `T` are passed up the circuit and given to `f_inner` in aggregate as a vector from the children.
"""
function foldup_aggregate(node::Dag, f_leaf::Function, f_inner::Function, ::Type{T})::T where {T}
    foldup_aggregate(node[end], f_leaf, f_inner, T)
end

function foldup_aggregate(node::DagNode, f_leaf::Function, f_inner::Function, ::Type{T})::T where {T}
    @assert node.bit == false
    v = foldup_aggregate_rec(node, f_leaf, f_inner, T)
    flip_bit(node)
    return v
end

function foldup_aggregate_rec(node::DagNode, f_leaf::Function, f_inner::Function,
                    ::Type{T}, ::Val{Bit} = Val(!node.bit))::T where {T,Bit}
    if node.bit == Bit
        return node.data::T
    else
        node.bit = Bit
        v = if isinner(node)
            child_values = Vector{T}(undef, num_children(node))
            map!(c -> foldup_aggregate_rec(c, f_leaf, f_inner, T, Val(Bit))::T, child_values, children(node))
            f_inner(node, child_values)::T
        else
            f_leaf(node)::T
        end
        node.data = v
        return v
    end
end

function foldup_aggregate(node::TreeNode, f_leaf::Function, f_inner::Function, ::Type{T})::T where T
    v = if isinner(node)
        child_values = Vector{T}(undef, num_children(node))
        map!(c -> foldup_aggregate(c, f_leaf, f_inner, T)::T, child_values, children(node))
        f_inner(node, child_values)::T
    else
        f_leaf(node)::T
    end
    return v
end

"""
Compute a function top down on the circuit.
`f_root` is called on the root node, `f_leaf` is called on leaf nodes, and `f_inner` is called on inner nodes.
Values of type `T` are passed down the circuit and given to `f_inner` as a value from the parent.
"""
mutable struct Downpass
    passup
    passdown
    Downpass(T, passup) = new(passup, Vector{T}())
end
function folddown_aggregate(node::Dag, f_root::Function, f_leaf::Function, f_inner::Function, ::Type{T})::Nothing where {T}
    # folddown_aggregate(node[end], f_root, f_leaf, f_inner, T)
    @inline isroot(n) = n === node[end]
    @inline function folddown_init(n::DagNode)
        n.data = Downpass(T, n.data)
        nothing
    end
    @inline function data_to_childern(c::DagNode, data)
        push!(c.data.passdown, data)
    end

    foreach(folddown_init, node)
    n = node[end]
    n.data = x = f_root(n, n.data.passup)
    map(c -> data_to_childern(c, x), children(n))

    inners = filter(isinner, @view node[1:end-1])
    map(Iterators.reverse(inners)) do n
        n.data = @inbounds f_inner(n, n.data.passup, n.data.passdown)::T
        map(c -> data_to_childern(c, n.data), children(n))
        nothing
    end

    leafs = filter(isleaf, node)
    map(leafs) do n
        n.data = f_leaf(n, n.data.passup, n.data.passdown)::T
        nothing
    end
    flip_bit(node[end], Val(false))
end

#####################
# other methods
#####################

# When you suspect there is a bug but execution halts, it may be because of
# pretty printing a huge recursive graph structure.
# To safeguard against that case, we set a default show:
Base.show(io::IO, c::Node) = print(io, "$(typeof(c))($(hash(c))))")

"Number of nodes in the graph"
@inline num_nodes(c::DiGraph) = length(c)
function num_nodes(node::DagNode)
    count::Int = 0
    foreach(node) do n
        count += 1
    end
    count
end

"Number of edges in the graph"
@inline num_edges(c::DiGraph) = sum(n -> num_children(n), c)
function num_edges(node::DagNode)
    count::Int = 0
    foreach(node) do n
        count += num_children(n)
    end
    count
end

"Is the node a leaf node?"
@inline isleaf(n::Node) = NodeType(n) isa Leaf

"Is the node an Inner node?"
@inline isinner(n::Node) = NodeType(n) isa Inner

"Get the list of inner nodes in a given graph"
inodes(c::Union{DagNode,DiGraph}) = filter(isinner, c)

"Get the list of inner nodes in a given graph"
innernodes(c::Union{DagNode,DiGraph}) = inodes(c)

"Get the list of leaf nodes in a given graph"
leafnodes(c::DiGraph) = filter(isleaf, c)

"""
Compute the number of nodes in of a tree-unfolding of the DAG.
"""
function tree_num_nodes(dag::Dag)::BigInt
    size = Dict{DagNode,BigInt}()
    for node in dag
        if has_children(node)
            size[node] = one(BigInt) + sum(c -> size[c], children(node))
        else
            size[node] = one(BigInt)
        end
    end
    size[dag[end]]
end

function tree_num_nodes(node::DagNode)::BigInt
    @inline f_leaf(n) = zero(BigInt)
    @inline f_inner(n, call) = (num_children(n) + mapreduce(c -> call(c), +, children(n)))
    foldup(node, f_leaf, f_inner, BigInt)
end

"Rebuild a DAG's linear bottom-up order from a new root node"
@inline node2dag(r::DagNode, ::Type{T} = Union{}) where T = filter(x -> true, r, typejoin(T,typeof(r)))

@inline dag2node(dag::Dag)::DagNode = dag[end]

"Get the type of node contained in this graph"
grapheltype(circuit::DiGraph)::Type{<:Node} = eltype(circuit)
grapheltype(::Type{T}) where {T<:DiGraph} = eltype(T)

"Is one node equal to another locally, ignoring children?"
function isequal_local end

import Base.isequal

"Is one ordered tree equal to another?"
isequal(t1::Tree, t2::Tree)::Bool =
    isequal(t1[end], t2[end])
isequal(n1::TreeNode, n2::TreeNode)::Bool =
    isequal_local(n1,n2) && isequal_rec(NodeType(n1), NodeType(n2), n1, n2)

isequal_rec(::Leaf, ::Leaf, ::TreeNode, ::TreeNode)::Bool = true
function isequal_rec(::Inner, ::Inner, n1::TreeNode, n2::TreeNode)::Bool
    foreach(children(n1), children(n2)) do c1, c2 # we need all to support varagrs!
        if !isequal(c1, c2)
            return false
        end
    end
    return true
end

"Is one unordered tree equal to another?"
isequal_unordered(t1::Tree, t2::Tree)::Bool =
    isequal_unordered(t1[end], t2[end])
isequal_unordered(n1::TreeNode, n2::TreeNode)::Bool =
    isequal_local(n1,n2) && isequal_unordered_rec(NodeType(n1), NodeType(n2), n1, n2)

isequal_unordered_rec(::Leaf, ::Leaf, ::TreeNode, ::TreeNode)::Bool = true
function isequal_unordered_rec(::Inner, ::Inner, n1::TreeNode, n2::TreeNode)::Bool
    @assert num_children(n1) == 2 && num_children(n2) == 2 "`isequal_unordered` is only implemented for binary trees"
    c1 = children(n1)
    c2 = children(n2)
    return ((isequal_unordered(c1[1],c2[1]) &&  isequal_unordered(c1[2],c2[2]))
            || (isequal_unordered(c1[1],c2[2]) && isequal_unordered(c1[2],c2[1])))
end

"""
Return the leftmost child.
"""
function left_most_child(root::DagNode)::DagNode
    while isinner(root)
        root = children(root)[1]
    end
    root
end

"""
Return the rightmost child.
"""
function right_most_child(root::DagNode)::DagNode
    while isinner(root)
        root = children(root)[end]
    end
    root
end

"""
Find the least common ancestor (assumes the graph has a parent pointer and a list of descendents)
"""
lca(v::DagNode)::DagNode = v
lca(v::DagNode, w::DagNode)::DagNode = begin
    if v == w
        return v
    end
    if descends_from(w,v)
        return v
    end
    candidate::Union{DagNode,Nothing} = w
    while issomething(candidate)
        if descends_from(v,candidate)
            return candidate
        end
        candidate = parent(candidate)
    end
    error("First argument is not contained in the root of second argument. There is no LCA.")
end
lca(v::DagNode, w::DagNode, u::DagNode, r::DagNode...)::DagNode = lca(lca(v,w), u, r...)

function descends_from end
function parent end

#####################
# debugging methods (not performance critical)
#####################

"Give count of types and fan-ins of all nodes in the graph"
node_stats(c::Union{DiGraph,DagNode}) = merge(leaf_stats(c), inode_stats(c))

"Give count of types and fan-ins of inner nodes in the graph"
function inode_stats(c::Union{DiGraph,DagNode})
    groups = groupby(e -> (typeof(e),num_children(e)), inodes(c))
    map_values(v -> length(v), groups, Int)
end

"Give count of types of leaf nodes in the graph"
function leaf_stats(c::Union{DiGraph,DagNode})
    groups = groupby(e -> typeof(e), leafnodes(c))
    map_values(v -> length(v), groups, Int)
end

ext: jl


text:
####################################################

#  testRRT.jl

#  This file is used for development and testing of RRT.jl
#  nouyang 2017

####################################################
# MAIN
####################################################
include("RRT.jl")
#plotly()

using GeometryTypes
using Distributions

using plotfxn
using algT
using algfxn

function main()
#  runs the RRT once and plots the results.
    startstate = Point(1.,1)
    goalstate = Point(20.,20)

    numSamples = 10
    connectRadius = 5
	flagOptimal = false
    param = algT.AlgParameters(numSamples, connectRadius)

    print("\n ---- Running RRT ------\n")
    ## Define obstacles
    obs1 = HyperRectangle(Vec(8,3.), Vec(2,2.)) #Todo
    obs2 = HyperRectangle(Vec(4,4.), Vec(2,10.)) #Todo
    obs3 = HyperRectangle(Vec(14,14.), Vec(3,4.)) #Todo

    obstacles = Vector{HyperRectangle}()
    push!(obstacles, obs1, obs3, obs2)

    # Define walls
    walls = Vector{LineSegment}()
    w,h  = 20,20
    perimeter = HyperRectangle(Vec(0.,0), Vec(w,h))
    roomPerimeter = algfxn.decompRect(perimeter)
    #internalwalls =  (linesegs
    for l in roomPerimeter
        push!(walls, l)
    end

    r = algT.Room(w,h,walls,obstacles)
    roomPlot = plotfxn.plotRoom(r)

    #!
    nodeslist, edgeslist, isPathFound, pathNodes, pathcost = rrtPlan(r, param, startstate, goalstate, obstacles)

    ## Plot path found
    #title = "PRM with # samples =$numSamples, \nPathfound = $isPathFound, \npathcost = $pathcost"

    timestamp = Base.Dates.now()
    timestamp = Base.Dates.format(timestamp, "dd u yyyy HH:MM:SS")
    title = " Week 11 RRT [ $timestamp ] \n\n$numSamples samples, connectRadius = $connectRadius, pathcost = $(ceil(pathcost,2))"

    graph = algT.Graph(startstate, goalstate, nodeslist, edgeslist)
    # this is the gold
    #!
    rrtPlot = plotfxn.plotRRT(roomPlot, graph, pathNodes, title::String);

####
#startGoal = algT.GraphNode(0, Point(0,0))
#endNode = algT.GraphNode(0, Point(0,0))

end


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder

name = "tree_sitter_highlight"
version = v"0.20.6"

# Collection of sources required to complete build
sources = [
    GitSource("https://github.com/tree-sitter/tree-sitter/",
              "ccd6bf554d922596ce905730d98a77af368bba5c"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/

# Use cargo-rustc to build tree-sitter-highlight as a dynamic library here
cargo rustc --release --manifest-path highlight/Cargo.toml -- --crate-type=cdylib

install -Dvm 755 target/${rust_target}/release/deps/*tree_sitter_highlight-*.${dlext} "${libdir}/libtree_sitter_highlight.${dlext}"
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()
# Rust toolchain is unusable on i686-w64-mingw32
filter!(p -> !(Sys.iswindows(p) && arch(p) == "i686"), platforms)
# Also, can't build cdylib for Musl systems
filter!(p -> libc(p) != "musl", platforms)

# The products that we will ensure are always built
products = [
    LibraryProduct("libtree_sitter_highlight", :libtree_sitter_highlight),
]

# Dependencies that must be installed before this package can be built
dependencies = Dependency[
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies;
               compilers=[:c, :rust], julia_compat="1.6", lock_microarchitecture=false)

max_stars_repo_name: MichelJuillard/Yggdrasil

max_issues_repo_name: MichelJuillard/Yggdrasil

max_forks_repo_name: MichelJuillard/Yggdrasil


text:
__precompile__()

@doc read(joinpath(dirname(@__DIR__), "README.md"), String)
module Atom

using Base.StackTraces, InteractiveUtils, Logging
using Juno, Lazy, JSON, MacroTools, Media, Requires
import Media: @dynamic

function __init__()
  Juno.activate()

  atreplinit() do repl
    fixdisplayorder()
    instantiate_repl_keybindings(repl)
  end

  # HACK: overwrite Pkg mode `on_done` function after Pkg mode has been Initialized
  f = (repl) -> begin
    i = findfirst(Base.active_repl.interface.modes) do mode
      mode isa LineEdit.Prompt &&
      isdefined(mode, :on_done) &&
      let ms = collect(methods(mode.on_done))
        length(ms) > 0 && parentmodule(ms[1].module) == Pkg
      end
    end
    i === nothing && return # if no Pkg mode somehow
    pkg_mode = Base.active_repl.interface.modes[i]
    pkg_mode.on_done = let original = pkg_mode.on_done
      (s, buf, ok) -> begin
        original(s, buf, ok)
        update_project()
      end
    end
  end
  push!(Base.repl_hooks, f)

  start_eval_backend()

  Atom.handle("connected") do
    if isREPL(before_run_repl = true)
      reset_repl_history()
      fixdisplayorder()

      # HACK: overloading this allows us to open remote files
      InteractiveUtils.eval(quote
        function InteractiveUtils.edit(path::AbstractString, line::Integer = 0)
          if endswith(path, ".jl")
            f = Base.find_source_file(path)
            f !== nothing && (path = f)
          end
          $(msg)("openFile", Base.abspath(path), line - 1)
        end
      end)
    end

    update_project()

    nothing
  end

  @require WebIO="0f1e0344-ec1d-5b48-a673-e5cf874b6c29" include("display/webio.jl")
  @require Traceur="37b6cedf-1f77-55f8-9503-c64b63398394" include("profiler/traceur.jl")
end

# basics
include("comm.jl")
include("utils.jl")
include("misc.jl")
include("environments.jl")
include("display/display.jl")
include("progress.jl")
include("static/static.jl")
include("modules.jl")

include("eval.jl")
include("repl.jl")
include("workspace.jl")
include("outline.jl")
include("docs.jl")
include("completions.jl")
include("goto.jl")
include("datatip.jl")
include("formatter.jl")
include("frontend.jl")
include("debugger/debugger.jl")
include("profiler/profiler.jl")

include("precompile.jl")
_precompile_()

end # module


text:
using Agents
using AgentsX

include("LayeredAgentActions.jl")

"""
    Economic <:ParamLayer
The Economic layer for LayeredAgent2
"""
mutable struct Economic <: ParamLayer
    rich::Bool
end

"""
    Cognitive <:ParamLayer
The Cognitive Layer for LayeredAgent2
"""
mutable struct Cognitive <: ParamLayer
    tired::Bool
end

"""
    Physiological <:ParamLayer
The Physiological Layer for LayeredAgent2
"""
mutable struct Physiological <: ParamLayer
    hungry::Bool
end

"""
    LayeredAgent2 <: AbstractAgent
A LayeredAgent in a GridSpace
This agent has Social, Economic, and Cognitive Layers
"""
mutable struct LayeredAgent2 <: AbstractAgent
    id::Int
    pos::NTuple{2, Int}
    Economic::Economic
    Cognitive::Cognitive
    Physiological::Physiological
end


text:
function HarmonicParametricAcceleration(arg0::Vector3D, arg1::AttitudeProvider, arg2::JString, arg3::AbsoluteDate, arg4::jdouble, arg5::jint)
    return HarmonicParametricAcceleration((Vector3D, AttitudeProvider, JString, AbsoluteDate, jdouble, jint), arg0, arg1, arg2, arg3, arg4, arg5)
end

function HarmonicParametricAcceleration(arg0::Vector3D, arg1::jboolean, arg2::JString, arg3::AbsoluteDate, arg4::jdouble, arg5::jint)
    return HarmonicParametricAcceleration((Vector3D, jboolean, JString, AbsoluteDate, jdouble, jint), arg0, arg1, arg2, arg3, arg4, arg5)
end

function depends_on_position_only(obj::HarmonicParametricAcceleration)
    return jcall(obj, "dependsOnPositionOnly", jboolean, ())
end

function get_parameters_drivers(obj::HarmonicParametricAcceleration)
    return jcall(obj, "getParametersDrivers", Vector{ParameterDriver}, ())
end

function init(obj::HarmonicParametricAcceleration, arg0::SpacecraftState, arg1::AbsoluteDate)
    return jcall(obj, "init", void, (SpacecraftState, AbsoluteDate), arg0, arg1)
end


text:
using Revise
using MendelIHT
using SnpArrays
using Random
using GLM
using DelimitedFiles
using Distributions
using LinearAlgebra
using CSV
using DataFrames
using StatsBase
using TraitSimulation
BLAS.set_num_threads(1)

"""
    βi ~ N(0, 0.1) chosen randomly across genome

k = Number of causal SNPs
p = Total number of SNPs
traits = Number of traits (phenotypes)
overlap = number of causal SNPs shared in each trait
"""
function simulate_random_beta(k::Int, p::Int, traits::Int; overlap::Int=0, βσ=1.0)
    d = Normal(0, βσ)
    true_b = zeros(p, traits)
    if overlap == 0
        causal_snps = sample(1:(traits * p), k, replace=false)
        true_b[causal_snps] = rand(d, k)
    else
        shared_snps = sample(1:p, overlap, replace=false)
        weight_vector = aweights(1 / (traits * (p - overlap)) * ones(traits * p))
        for i in 1:traits
            weight_vector[i*shared_snps] .= 0.0 # avoid sampling from shared snps
        end
        @assert sum(weight_vector) ≈ 1.0 "sum(weight_vector) = $(sum(weight_vector)) != 1.0"
        # simulate β for shared predictors
        for i in 1:traits
            true_b[shared_snps, i] = rand(d, overlap)
        end
        # simulate β for none shared predictors
        nonshared_snps = sample(1:(traits * p), weight_vector, k - (traits * overlap), replace=false)
        true_b[nonshared_snps] = rand(d, k - (traits * overlap))
    end

    return true_b
end

"""
    βi ~ Uniform([-0.5, -0.45, ..., -0.05, 0.05, ..., 0.5]) chosen uniformly across genome

k = Number of causal SNPs
p = Total number of SNPs
traits = Number of traits (phenotypes)
overlap = number of pleiotropic SNPs (affects each trait with effect size βi / r)
"""
function simulate_fixed_beta(k::Int, p::Int, traits::Int; overlap::Int=0)
    true_b = zeros(p, traits)
    effect_sizes = collect(0.05:0.05:0.5)
    # num_causal_snps = (k - (traits - 1) * overlap) :: Int
    # @assert num_causal_snps > 0 "number of causal SNPs should be positive but was $num_causal_snps"
    # idx_causal_snps = sample(1:p, num_causal_snps, replace=false)
    # @assert length(idx_causal_snps) == num_causal_snps "length(idx_causal_snps) = $(length(idx_causal_snps)) != num_causal_snps"
    # shuffle!(idx_causal_snps)
    k_indep = k - 2overlap # pleiotropic SNPs affect 2 phenotypes
    num_causal_snps = k_indep + overlap
    @assert num_causal_snps > 0 "number of causal SNPs should be positive but was $num_causal_snps"
    idx_causal_snps = sample(1:p, num_causal_snps, replace=false)
    @assert length(idx_causal_snps) == num_causal_snps "length(idx_causal_snps) = $(length(idx_causal_snps)) != num_causal_snps"
    shuffle!(idx_causal_snps)

    # pleiotropic SNPs affect 2 phenotypes
    for i in 1:overlap
        j = idx_causal_snps[i]
        rs = sample(1:traits, 2, replace=false)
        for r in rs
            true_b[j, r] = rand(-1:2:1) * effect_sizes[rand(1:10)]
        end
    end
    # non pleiotropic SNPs affect only 1 phenotype
    for i in (overlap+1):length(idx_causal_snps)
        idx = idx_causal_snps[i]
        true_b[idx, rand(1:traits)] = rand(-1:2:1) * effect_sizes[rand(1:10)]
    end

    @assert count(!iszero, true_b) == k "count(!iszero, true_b) = $(count(!iszero, true_b)) != k = $k"

    return true_b
end

"""
    every causal SNP affects every phenotype, chosen fixed distance with each other across genome
    Their total effect size sums up to 0.5

k = Number of causal SNPs (each SNP affects every trait)
p = Total number of SNPs
traits = Number of traits (phenotypes)
"""
function simulate_pleiotropic_beta(k::Int, p::Int, traits::Int)
    true_b = zeros(p, traits)
    effect_size = 0.5 / traits
    idx_causal_snps = collect(div(p, k):div(p, k):p)
    shuffle!(idx_causal_snps)

    # pleiotropic SNPs affect every phenotype
    for i in idx_causal_snps
        for r in 1:traits
            true_b[i, r] = effect_size
        end
    end

    @assert count(!iszero, true_b) == k*traits "count(!iszero, true_b) should be equal to $(k*traits)"

    return true_b
end

"""
# Arguments
xla = simulated genotype matrix (converted to a SnpLinAlg)
k = number of causal SNPs
r = number of traits
Φ = estimated GRM (using GEMMA)

# Optional arguments
seed = random seed for reproducibility
σ2 = contribution of GRM
σe = random environmental effect
βoverlap = number of causal SNPs shared in all traits
"""
function simulate_multivariate_polygenic(
    plinkname::String, n::Int, p::Int, k::Int, r::Int;
    seed::Int=2021, σg=0.6, σe=0.4, βoverlap=2
    )
    # set seed
    Random.seed!(seed)

    # simulate `.bed` file with no missing data
    x = simulate_random_snparray("sim$seed/" * plinkname * ".bed", n, p)
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, impute=true, center=true, scale=true)
    make_bim_fam(xla, "sim$seed/" * plinkname)

    # intercept is the only nongenetic covariate
    Z = ones(n, 1)
    intercepts = zeros(r)' # each trait have 0 intercept

    # simulate β
    # B = simulate_pleiotropic_beta(k, p, r)
    B = simulate_fixed_beta(k, p, r, overlap=βoverlap)
    writedlm("sim$(seed)/trueb.txt", B)

    # between trait covariance matrix
    Σ = random_covariance_matrix(r)

    # between sample covariance is identity + GRM (2x because OpenMendel always uses half the GRM)
    run(`./gemma -bfile sim$(seed)/$plinkname -gk 1 -o $plinkname`)
    run(`mv ./output/$(plinkname).cXX.txt sim$(seed)`)
    Φ = readdlm("sim$(seed)/" * plinkname * ".cXX.txt")
    V = σg * Φ + σe * I

    # simulate y using TraitSimulations.jl (https://github.com/OpenMendel/TraitSimulation.jl/blob/master/src/modelframework.jl#L137)
    vc = @vc Σ ⊗ V
    μ = zeros(n, r)
    μ_null = zeros(n, r)
    LinearAlgebra.mul!(μ_null, Z, intercepts)
    mul!(μ, xla, B)
    BLAS.axpby!(1.0, μ_null, 1.0, μ)
    VCM_model = VCMTrait(Z, intercepts, xla, B, vc, μ)
    Y = Matrix(Transpose(simulate(VCM_model)))

    # simulate using Distributions.jl
    # μ = z * intercepts + xla * B
    # Y = rand(MatrixNormal(μ', Σ, V))

    return xla, Matrix(Z'), B, Σ, Y
end

function simulate_multivariate_sparse(
    plinkname::String, n::Int, p::Int, k::Int, r::Int;
    seed::Int=2021, σg=0.6, σe=0.4, βoverlap=2
    )
    # set seed
    Random.seed!(seed)

    # simulate `.bed` file with no missing data
    x = simulate_random_snparray("sim$seed/" * plinkname * ".bed", n, p)
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, impute=true, center=true, scale=true)
    make_bim_fam(xla, "sim$seed/" * plinkname)

    # intercept is the only nongenetic covariate
    Z = ones(n, 1)
    intercepts = zeros(r)' # each trait have 0 intercept

    # simulate β
    # B = simulate_pleiotropic_beta(k, p, r)
    B = simulate_fixed_beta(k, p, r, overlap=βoverlap)
    writedlm("sim$(seed)/trueb.txt", B)

    # between trait covariance matrix (but GEMMA GRM still needed)
    run(`./gemma -bfile sim$(seed)/$plinkname -gk 1 -o $plinkname`)
    run(`mv ./output/$(plinkname).cXX.txt sim$(seed)`)
    Σ = random_covariance_matrix(r)
    writedlm("sim$(seed)/trueCovariance.txt", Σ)

    # simulate multivariate normal phenotype for each sample
    μ = xla * B + Z*intercepts

    # simulate response
    Y = zeros(n, r)
    for i in 1:n
        μi = @view(μ[i, :])
        Y[i, :] = rand(MvNormal(μi, Σ))
    end

    return xla, Matrix(Z'), B, Σ, Matrix(Y')
end

"""
simulate under IHT's model
"""
function simulate_NFBC1966_sparse(
    plinkname::String, k::Int, r::Int;
    seed::Int=2021, σg=0.6, σe=0.4, βoverlap=2
    )
    # set seed
    Random.seed!(seed)

    # simulate `.bed` file with no missing data
    x = SnpArray(plinkname * ".bed")
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, impute=true, center=true, scale=true)
    n, p = size(xla)

    # intercept is the only nongenetic covariate
    Z = ones(n, 1)
    intercepts = zeros(r)' # each trait have 0 intercept

    # simulate β
    # B = simulate_pleiotropic_beta(k, p, r)
    B = simulate_fixed_beta(k, p, r, overlap=βoverlap)
    writedlm("sim$(seed)/trueb.txt", B)

    # between trait covariance matrix
    Σ = random_covariance_matrix(r)
    writedlm("sim$(seed)/trueCovariance.txt", Σ)

    # simulate multivariate normal phenotype for each sample
    μ = xla * B + Z*intercepts

    # simulate response
    Y = zeros(n, r)
    for i in 1:n
        μi = @view(μ[i, :])
        Y[i, :] = rand(MvNormal(μi, Σ))
    end

    return xla, Matrix(Z'), B, Σ, Matrix(Y')
end

"""
Trait covariance matrix is σg * Φ + σe * I where Φ is the GRM.
"""
function simulate_NFBC1966_polygenic(
    plinkname::String, k::Int, r::Int;
    seed::Int=2021, σg=0.1, σe=0.9, βoverlap=2,
    )
    # set seed
    Random.seed!(seed)

    # simulate `.bed` file with no missing data
    x = SnpArray(plinkname * ".bed")
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, impute=true, center=true, scale=true)
    n, p = size(x)

    # intercept is the only nongenetic covariate
    Z = ones(n, 1)
    intercepts = zeros(r)' # each trait have 0 intercept

    # simulate β
    # B = simulate_pleiotropic_beta(k, p, r)
    B = simulate_fixed_beta(k, p, r, overlap=βoverlap)
    writedlm("sim$(seed)/trueb.txt", B)

    # between trait covariance matrix
    Σ = random_covariance_matrix(r)
    writedlm("sim$(seed)/true_cov.txt", Σ)

    # between sample covariance is identity + GRM
    Φ = readdlm(plinkname * ".cXX.txt")
    V = σg * Φ + σe * I

    # simulate y using TraitSimulations.jl (https://github.com/OpenMendel/TraitSimulation.jl/blob/master/src/modelframework.jl#L137)
    vc = @vc Σ ⊗ V
    μ = zeros(n, r)
    μ_null = zeros(n, r)
    LinearAlgebra.mul!(μ_null, Z, intercepts)
    mul!(μ, xla, B)
    BLAS.axpby!(1.0, μ_null, 1.0, μ)
    VCM_model = VCMTrait(Z, intercepts, xla, B, vc, μ)
    Y = Matrix(Transpose(simulate(VCM_model)))

    # simulate using Distributions.jl
    # μ = z * intercepts + xla * B
    # Y = rand(MatrixNormal(μ', Σ, V))

    return xla, Matrix(Z'), B, Σ, Y
end

function make_bim_fam(x::AbstractMatrix, name::String)
    n, p = size(x)

    #create .bim file structure: https://www.cog-genomics.org/plink2/formats#bim
    open(name * ".bim", "w") do f
        for i in 1:p
            write(f, "1\tsnp$i\t0\t1\t1\t2\n")
        end
    end

    open(name * ".fam", "w") do f
        for i in 1:n
            write(f, "$i\t1\t0\t0\t1\t1.0\n")
        end
    end
end

function make_GEMMA_fam_file(x::AbstractMatrix, y::AbstractVecOrMat, name::String)
    ly = size(y, 1)
    n, p = size(x)

    # put 1st phenotypes in 6th column, 2nd phenotype in 7th column ... etc
    traits = size(y, 1)
    open(name * ".fam", "w") do f
        for i in 1:n
            write(f, "$i\t1\t0\t0\t1")
            for j in 1:traits
                write(f, "\t$(y[j, i])")
            end
            write(f, "\n")
        end
    end
end

function make_MVPLINK_fam_and_phen_file(x::AbstractMatrix, y::AbstractVecOrMat, name::String)
    ly = size(y, 1)
    n, p = size(x)

    # put a random phenotype in fam file
    traits = size(y, 1)
    open(name * ".fam", "w") do f
        for i in 1:n
            println(f, "$i\t1\t0\t0\t1\t1")
        end
    end

    # save phenotypes in separate `.phen` file
    open(name * ".phen", "w") do io
        print(io, "FID\tIID")
        for j in 1:traits
            print(io, "\tT$j")
        end
        print(io, "\n")
        for i in 1:n
            print(io, "$i\t1")
            for j in 1:traits
                write(io, "\t$(y[j, i])")
            end
            print(io, "\n")
        end
    end
end

"""
Computes power and false positive rates
- p: total number of SNPs
- pleiotropic_snps: Indices (or ID) of the true causal SNPs that affect >1 phenotype
- independent_snps: Indices (or ID) of the true causal SNPs that affect exactly 1 phenotype
- signif_snps: Indices (or ID) of SNPs that are significant after testing

returns: pleiotropic SNP's power, independent SNP's power, number of false positives, and false positive rate
"""
function power_and_fpr(p::Int, pleiotropic_snps, independent_snps, signif_snps)
    pleiotropic_power = length(signif_snps ∩ pleiotropic_snps) / length(pleiotropic_snps)
    independent_power = length(signif_snps ∩ independent_snps) / length(independent_snps)
    correct_snps = pleiotropic_snps ∪ independent_snps
    FP = length(signif_snps) - length(signif_snps ∩ correct_snps) # number of false positives
    TN = p - length(signif_snps) # number of true negatives
    FPR = FP / (FP + TN)
    return pleiotropic_power, independent_power, FP, FPR
end

# https://github.com/OpenMendel/MendelPlots.jl/blob/master/src/gwasplots.jl#L108
function genomic_inflation(pvalues::AbstractVecOrMat)
    λ = median(quantile.(Chisq(1), 1 .- pvalues) ./ quantile(Chisq(1), 0.5))
end

"""
- filename: gemma's output file name
- pleiotropic_snpid: ID of causal SNPs that affect > 1 phenotypes
- independent_snpid: ID of causal SNPs that affect 1 phenotypes

returns: power, number of false positives, and false positive rate
"""
function process_gemma_result(filename, pleiotropic_snpid, independent_snpid)
    # read GEMMA result
    gemma_df = CSV.read(filename, DataFrame)
    snps = size(gemma_df, 1)

    # pvalues
    pval_wald = gemma_df[!, :p_wald]
#    pval_lrt = gemma_df[!, :p_lrt]
#     pval_score = gemma_df[!, :p_score]

    # estimated beta
#     estim_β1 = gemma_df[!, :beta_1]
#     estim_β2 = gemma_df[!, :beta_2]

#     # estimated covariance matrix
#     estim_σ11 = gemma_df[!, :Vbeta_1_1]
#     estim_σ12 = gemma_df[!, :Vbeta_1_2]
#     estim_σ22 = gemma_df[!, :Vbeta_2_2];

    # check how many real SNPs were recovered
    signif_snps = findall(x -> x ≤ 0.05 / snps, pval_wald) # gemma's selected snps
    signif_snpid = gemma_df[signif_snps, :rs]

    # return power, false positives, false positive rate, and genomic inflation
    pleiotropic_power, independent_power, FP, FPR = power_and_fpr(snps,
        pleiotropic_snpid, independent_snpid, signif_snpid)
    λ = genomic_inflation(pval_wald)
    return pleiotropic_power, independent_power, FP, FPR, λ
end

"""
- filename: mvPLINK's output file name
- pleiotropic_snps: indices for causal SNPs that affect > 1 phenotypes
- independent_snps: indices for causal SNPs that affect 1 phenotypes

returns: power, number of false positives, and false positive rate
"""
function process_mvPLINK(filename, pleiotropic_snps, independent_snps)
    # read mvPLINK result
    mvplink_df = CSV.read(filename, DataFrame, delim=' ', ignorerepeated=true)
    snps = size(mvplink_df, 1)

    # get pvalues, possibly accounting for "NA"s
    if eltype(mvplink_df[!, :P]) == Float64
        pval = mvplink_df[!, :P]
    else
        mvplink_df[findall(x -> x == "NA", mvplink_df[!, :P]), :P] .= "1.0"
        pval = parse.(Float64, mvplink_df[!, :P])
    end

    # SNPs passing threshold
    signif_snps = findall(x -> x ≤ 0.05 / snps, pval)

    # compute power, false positives, and false positive rate
    pleiotropic_power, independent_power, FP, FPR = power_and_fpr(snps,
        pleiotropic_snps, independent_snps, signif_snps)
    λ = genomic_inflation(pval)
    return pleiotropic_power, independent_power, FP, FPR, λ
end

function one_NFBC_simulation(
    k::Int, r::Int;
    seed::Int=2021, σg=0.1, σe=0.9, βoverlap=2,
    path=5:5:50, init_beta=false, model=:polygenic, debias=100
    )
    isdir("sim$seed") ? (return nothing) : mkdir("sim$seed")
    plinkname = "NFBC.qc.imputeBy0.chr.1"

    # simulate data
    Random.seed!(seed)
    if model == :polygenic
        xla, Z, B, Σ, Y = simulate_NFBC1966_polygenic("../" * plinkname, k, r,
            seed=seed, σg=σg, σe=σe, βoverlap=βoverlap)
    elseif model == :sparse
        xla, Z, B, Σ, Y = simulate_NFBC1966_sparse("../" * plinkname, k, r,
            seed=seed, σg=σg, σe=σe, βoverlap=βoverlap)
    else
        error("model misspecified!")
    end

    correct_snps = unique([x[1] for x in findall(!iszero, B)])
    pleiotropic_snps, independent_snps = Int[], Int[]
    for snp in correct_snps
        count(x -> abs(x) > 0, @view(B[snp, :])) > 1 ?
            push!(pleiotropic_snps, snp) : push!(independent_snps, snp)
    end
    snpdata = SnpData("../" * plinkname)
    pleiotropic_snp_rsid = snpdata.snp_info[pleiotropic_snps, :snpid]
    independent_snp_rsid = snpdata.snp_info[independent_snps, :snpid]

    # run GEMMA (GRM is precomputed already)
    cd("sim$seed")
    run(`cp ../../$(plinkname).bed .`)
    run(`cp ../../$(plinkname).bim .`)
    run(`cp ../../$(plinkname).cXX.txt .`)
    make_GEMMA_fam_file(xla, Y, plinkname)
    pheno_columns = [string(ri) for ri in 1:r]
    gemma_time = @elapsed begin
        run(`../../gemma -bfile $plinkname -k $(plinkname).cXX.txt -notsnp -lmm 1 -n $pheno_columns -o gemma.sim$seed`)
    end
    gemma_pleiotropic_power, gemma_independent_power, gemma_FP, gemma_FPR, gemma_λ =
        process_gemma_result("output/gemma.sim$seed.assoc.txt", pleiotropic_snp_rsid, independent_snp_rsid)
    println("GEMMA time = $gemma_time, pleiotropic power = $gemma_pleiotropic_power, independent power = $gemma_independent_power, FP = $gemma_FP, FPR = $gemma_FPR, gemma_λ=$gemma_λ")
    mv("output/gemma.sim$seed.assoc.txt", "gemma.sim$seed.assoc.txt")
    mv("output/gemma.sim$seed.log.txt", "gemma.sim$seed.log.txt")

    # run multivariate IHT
    mIHT_time = @elapsed begin
        mses = cross_validate(plinkname, MvNormal, path=path, phenotypes=collect(1:r).+5;
            init_beta=init_beta, debias=debias)
        k_rough_guess = path[argmin(mses)]
        dense_path = (k_rough_guess - 4):(k_rough_guess + 4)
        mses_new = cross_validate(plinkname, MvNormal, path=dense_path, phenotypes=collect(1:r).+5;
            init_beta=init_beta, debias=debias, cv_summaryfile="miht.cviht.summary.txt")
        iht_result = iht(plinkname, dense_path[argmin(mses_new)], MvNormal, phenotypes=collect(1:r).+5;
            init_beta=init_beta, debias=debias, summaryfile="miht.summary.txt")
    end
    detected_snps = Int[]
    for i in 1:r # save each beta separately
        β = iht_result.beta[i, :]
        detected_snps = detected_snps ∪ findall(!iszero, β)
        writedlm("multivariate_iht_beta$i.txt", β)
    end
    mIHT_pleiotropic_power, mIHT_independent_power, mIHT_FP, mIHT_FPR = power_and_fpr(size(B, 1), pleiotropic_snps, independent_snps, detected_snps)
    println("multivariate IHT time = $mIHT_time, pleiotropic power = $mIHT_pleiotropic_power, independent power = $mIHT_independent_power, FP = $mIHT_FP, FPR = $mIHT_FPR")

    # run multiple univariate IHT
    detected_snps = Int[]
    uIHT_time = @elapsed begin
        for trait in 1:r
            mses = cross_validate(plinkname, Normal, path=path, phenotypes=trait+5;
                init_beta=init_beta, debias=debias)
            k_rough_guess = path[argmin(mses)]
            dense_path = (k_rough_guess == 5) ? (0:5) : ((k_rough_guess - 4):(k_rough_guess + 4))
            mses_new = cross_validate(plinkname, Normal, path=dense_path, phenotypes=trait+5;
                init_beta=init_beta, debias=debias, cv_summaryfile="uiht.cviht.summary$trait.txt")
            best_k = dense_path[argmin(mses_new)]
            if best_k > 0
                iht_result = iht(plinkname, best_k, Normal, phenotypes=trait+5;
                    init_beta=init_beta, debias=debias, summaryfile="uiht.summary$trait.txt")
                β = iht_result.beta
            else
                β = zeros(size(B, 2))
            end

            # save results
            detected_snps = detected_snps ∪ findall(!iszero, β)
            writedlm("univariate_iht_beta$trait.txt", β)
        end
    end
    uIHT_pleiotropic_power, uIHT_independent_power, uIHT_FP, uIHT_FPR = power_and_fpr(size(B, 1), pleiotropic_snps, independent_snps, detected_snps)
    println("univariate IHT time = $uIHT_time, pleiotropic power = $uIHT_pleiotropic_power, independent power = $uIHT_independent_power, FP = $uIHT_FP, FPR = $uIHT_FPR")

    # run MVPLINK
    phenofile = plinkname * ".phen"
    make_MVPLINK_fam_and_phen_file(xla, Y, plinkname)
    mvplink_time = @elapsed run(`../../plink.multivariate --bfile $plinkname --noweb --mult-pheno $phenofile --mqfam`)
    mvPLINK_pleitropic_power, mvPLINK_independent_power, mvPLINK_FP, mvPLINK_FPR, mvPLINK_λ =
        process_mvPLINK("plink.mqfam.total", pleiotropic_snps, independent_snps)
    println("mvPLINK time = $mvplink_time, pleiotropic power = $mvPLINK_pleitropic_power, independent power = $mvPLINK_independent_power, FP = $mvPLINK_FP, FPR = $mvPLINK_FPR, mvPLINK_λ=$mvPLINK_λ")

    # clean up
    # rm("plink.hh", force=true)
    rm("$(plinkname).fam", force=true)
    rm("$(plinkname).bed", force=true)
    rm("$(plinkname).bim", force=true)
    rm("$(plinkname).cXX.txt", force=true)

    # save summary stats
    n, p = size(xla)
    open("summary.txt", "w") do io
        println(io, "Simulation $seed summary")
        println(io, "n = $n, p = $p, k = $k, r = $r, βoverlap=$βoverlap")
        println(io, "debias=$debias, init_beta=$init_beta")
        model == :polygenic ? println(io, "model = $model, σg=$σg, σe=$σe") : println(io, "model = $model")
        println(io, "")
        println(io, "mIHT time = $mIHT_time seconds, pleiotropic power = $mIHT_pleiotropic_power, independent power = $mIHT_independent_power, FP = $mIHT_FP, FPR = $mIHT_FPR, λ = NaN")
        println(io, "uIHT time = $uIHT_time seconds, pleiotropic power = $uIHT_pleiotropic_power, independent power = $uIHT_independent_power, FP = $uIHT_FP, FPR = $uIHT_FPR, λ = NaN")
        println(io, "mvPLINK time = $mvplink_time seconds, pleiotropic power = $mvPLINK_pleitropic_power, independent power = $mvPLINK_independent_power, FP = $mvPLINK_FP, FPR = $mvPLINK_FPR, λ = $mvPLINK_λ")
        println(io, "GEMMA time = $gemma_time seconds, pleiotropic power = $gemma_pleiotropic_power, independent power = $gemma_independent_power, FP = $gemma_FP, FPR = $gemma_FPR, λ = $gemma_λ")
    end
    cd("../")

    return nothing
end

#
# polygenic/independent beta = sign * Uniform{0.05, …, 0.25}
# max condition number = 10
# q = 5, iterates ≥5 times, init_beta=true, debias=false)
# sim set 1 are for k = 10, r = 2, βoverlap = 3 each affects 2 traits, path = 5:5:50 (then search around best k)
# sim set 2 are for k = 20, r = 3, βoverlap = 5 each affects 2 traits, path = 5:5:50 (then search around best k)
# sim set 3 are for k = 30, r = 5, βoverlap = 7 each affects 2 traits, path = 5:5:50 (then search around best k)
# sim set 4 are for k = 10, r = 10, βoverlap = 3 each affects 2 traits, path = 5:5:50 (then search around best k)
# sim set 5 are for k = 20, r = 50, βoverlap = 5 each affects 2 traits, path = 5:5:50 (then search around best k)
# sim set 6 are for k = 30, r = 100, βoverlap = 7 each affects 2 traits, path = 5:5:50 (then search around best k)
#

function run_simulation(set::Int, model::Symbol)
    σg = 0.1
    σe = 0.9
    init_beta = true
    debias = 100
    path = 5:5:50
    βoverlap = [3, 5, 7, 3, 5, 7]
    k = [10, 20, 30, 10, 20, 30]
    r = [2, 3, 5, 10, 50, 100]

    println("Simulation model = $model, set $set has k = $(k[set]), r = $(r[set]), βoverlap = $(βoverlap[set])")

    cur_dir = pwd() * "/set$set"
    isdir(cur_dir) || mkdir(cur_dir)
    k_cur = k[set]
    r_cur = r[set]
    βoverlap_cur = βoverlap[set]

    for seed in 1:100
        try
            cd(cur_dir)
            one_NFBC_simulation(k_cur, r_cur, seed = seed, path = path, βoverlap=βoverlap_cur,
                σg=σg, σe=σe, init_beta=init_beta, model=model, debias=debias)
        catch e
            bt = catch_backtrace()
            msg = sprint(showerror, e, bt)
            println("set $set sim $seed threw an error!")
            println(msg)
            continue
        end
    end
end

set = parse(Int, ARGS[1])
model = :polygenic
run_simulation(set, model)


text:
module test_julia_array

using Test

A = [6, 5, 4, 3, 2, 1]
@test 2 == splice!(A, 5)
@test [6, 5, 4, 3, 1] == A

A = [6, 5, 4, 3, 2, 1]
@test 2 == splice!(A, 5, [])
@test [6, 5, 4, 3, 1] == A

A = [6, 5, 4, 3, 2, 1]
@test 2 == splice!(A, 5, 10)
@test [6, 5, 4, 3, 10, 1] == A

A = [6, 5, 4, 3, 2, 1]
@test 2 == splice!(A, 5, [10, 11])
@test [6, 5, 4, 3, 10, 11, 1] == A

@test circshift(1:3, 1) == [3, 1, 2]

end # module test_julia_array

module test_julia_array_hcat

using Test

@test hcat([1,2,3]...) == reduce(hcat, [1,2,3]) == [1 2 3]

end # module test_julia_array_hcat

module test_julia_array_broadcast

using Test

@test (["a","b","c"] .=> [1,2,3]) == ["a"=>1, "b"=>2, "c"=>3]

end # module test_julia_array_broadcast

module test_julia_array_bool

using Test

@test !([] isa Array{Bool})
@test !([1] isa Array{Bool})
@test Bool[] isa Array{Bool}
@test Bool[1] isa Array{Bool}
@test [true] isa Array{Bool}
@test [true false] isa Array{Bool}
@test_throws InexactError Bool[2] isa Array{Bool}

end # module test_julia_array_bool

module test_julia_array_indices

using Test

A = Matrix(reshape(1:16, (4,4)))

@test LinearIndices(A)[2,2] == LinearIndices(A)[6] == 6
@test CartesianIndices(A)[2,2] == CartesianIndices(A)[CartesianIndex(2,2)] == CartesianIndex(2,2)

end #  module test_julia_array_indices

module test_julia_array_math

using Test

A = [1 2;
     3 4]

@test reverse(A, dims=1) == [3 4; 1 2]
@test reverse(A, dims=2) == [2 1; 4 3]

end # module test_julia_array_math

module test_julia_array_only_first

using Test

@test_throws ArgumentError only([])
@test 1 == only([1])
@test_throws ArgumentError only([1, 2])

@test_throws BoundsError   first([])

end # module test_julia_array_only_first


text:
@testset "utils.jl" begin

    # softplus
    s = (3,4,2)
    x = GenerativeModels.softplus_safe.(randn(Float32, s))
    @test size(x) == s
    @test all(x .> 0)
    @test eltype(x) == Float32
    @test eltype(softplus_safe.(x, Float64)) == Float64

    # stack layers
    s = (8, 6)
    x = randn(Float32, s)
    zdim = 2
    lsize = (s[1], 4, zdim)
    m = GenerativeModels.stack_layers(lsize, relu)
    y = m(x)

    @test length(m.layers) == 2
    @test size(y) == (zdim, s[2])
    @test eltype(y) == Float32
    @test size(m.layers[1].W) == (lsize[2], lsize[1])

    m = GenerativeModels.stack_layers(lsize, relu, last = sigmoid)
    y = m(x)
    @test size(y) == (2, 6)

    m = GenerativeModels.layer_builder([3,3,2], "tanh", "linear")
    @test length(m.layers) == 2
    @test m.layers[1] isa Dense
    @test m.layers[1].σ == tanh
    @test m.layers[2] isa Dense
    @test m.layers[2].σ == identity

end


text:
module ranking

export initSunStrategy, SunStrategy, Bucket, evaluate

using util
using Dates
using DataStructures
using preprocessing
using similarity_scorer
using Metrics
using Method

mutable struct Bucket
    master::Report
    reports::MutableLinkedList{Report}
    newestReport::Report
end

struct SunStrategy
    buckets::Dict{Int64, Bucket}
    window::Int32
    docFreq::DocFreq
    funcRemoval::DocFreqRemoval
    updDocFreq::Bool
end

function addReport!(strategy::SunStrategy, report; updateDocFreq=true)
    #update DocFreq
    if updateDocFreq
        updateDocFreq!(strategy.docFreq, [report])
    end

    buckets = strategy.buckets

    if !haskey(buckets, report.dup_id)
        # New bucket
        buckets[report.dup_id] =  Bucket(report, MutableLinkedList{Report}(report), report)
    else
        bucket = buckets[report.dup_id]
        push!(bucket.reports, report)

        if submittedEarlier(bucket.newestReport, report)
            bucket.newestReport = report
        end

        if submittedEarlier(report, bucket.master)
            bucket.master = report
        end
    end
end

function initSunStrategy(queries, reportid2report::Dict{Int64, Report}, scorer, funcRemoval::DocFreqRemoval, vocab, window, freq_by_stacks, static_df_ukn, upd_doc_freq=false, smooth=false)
    oldest = reportid2report[queries[1]]

    for (idx, report_id) in enumerate(queries)
        idx == 1 && continue
        cur =  reportid2report[report_id]

        if submittedEarlier(cur, oldest)
            oldest = cur
        end
    end
    @debug "Oldest query $(oldest)"

    @info "Update document frequency: $(upd_doc_freq)"

    if smooth
        @info "Smoothing Document frequency "
        docfreqObj = DocFreq(ones(length(vocab.vocab)), 1.0, freq_by_stacks, static_df_ukn, vocab)
    else
        docfreqObj = DocFreq(zeros(length(vocab.vocab)), 0.0, freq_by_stacks, static_df_ukn, vocab)
    end

    strategy = SunStrategy(Dict{Int64, Bucket}(), window, docfreqObj, funcRemoval, upd_doc_freq)
    initialDbReports = MutableLinkedList{Report}()
    bucketId2Size = Dict{Int64, Int64}()

    for report in values(reportid2report)
        !submittedEarlier(report, oldest) && continue

        get!(bucketId2Size, report.dup_id, 0)
        bucketId2Size[report.dup_id]+=1

        addReport!(strategy, report, updateDocFreq=false)

        push!(initialDbReports, report)
    end

    # Check buckets
    for (bucket_id, bucket) in strategy.buckets
        if bucket_id != bucket.master.id
            @error "Bucket id and master id is not the same: $(bucket_id) $(bucket.master.id)"
        end

        length(bucket.reports) != bucketId2Size[bucket_id] && @error "Bucket size is wrong: $(bucket_id) $(length(bucket.reports)) $(bucketId2Size[bucket_id])"
    end

    # Update function document frequency
    updateDocFreq!(docfreqObj, initialDbReports)

    for report in initialDbReports
        updateIndex!(scorer, report, docfreqObj)
    end

    # Remove common function in the stacktrace
    for report in initialDbReports
        removeCommonFunction!(report, strategy.funcRemoval, docfreqObj)
    end

    return strategy
end

function getCandidates(strategy::SunStrategy, query::Report)
    candidates = MutableLinkedList{Report}()

    for bucket in values(strategy.buckets)
        if strategy.window > 0 &&  strategy.window < (query.creation_day_ts - bucket.newestReport.creation_day_ts )
            continue
        end

        !submittedEarlier(bucket.newestReport, query) && continue

        for report in  bucket.reports
            submittedEarlier(query, report) && error("$(report.id) is in the bucket but it was submitted after query $(query.id) ")
            query.id == report.id && error("query $(query.id) is in the bucket")
            push!(candidates, report)
        end
    end

    return collect(candidates)
end

function evaluate(strategy, queries, reportid2report::Dict{Int64, Report},  scorer, metrics)
    total = 0.0
    total_time_sim = 0.0
    total_time_update = 0.0

    for (idx, query_id) in enumerate(queries)
        timeScore = @elapsed begin
            query = reportid2report[query_id]
            candidates = getCandidates(strategy, query)

            removeCommonFunction!(query, strategy.funcRemoval, strategy.docFreq)

            for candidate in candidates
                candidate.id == query_id && error("Query $(query_id) is within candidate list")
            end

            if length(candidates) == 0
                @warn "Query $(query_id) has 0 candidates!"
            end

            timesim = @elapsed scores = computeSimilarity(scorer, query, candidates, strategy.docFreq)
            recommendation_list = map((i,j)->(i,j), candidates, scores)

            # Sort  in descending order the bugs by probability of being duplicate
            sort!(recommendation_list, by=x->(x[2],x[1].id), rev=true)

            # println("$(query_id)\t$(length(recommendation_list))\t$([(recommendation_list[i][1].id, recommendation_list[i][2]) for i in 1:10])")

            for metric in metrics
                update_metric!(metric, query, recommendation_list)
            end

            # Add report to buckets and add to index
            addReport!(strategy, query; updateDocFreq=strategy.updDocFreq)
            timeupdate = @elapsed updateIndex!(scorer, query, strategy.docFreq)

            total_time_sim += timesim
            total_time_update += timeupdate
        end

        if idx > 0 && idx % 500 == 0
            @info "It generated recommendation of $(idx) reports"
        end

        total += timeScore
    end

    @info "Time to generate recommendation of $(total) reports: $(length(queries))"
    @info "Time to related to the method: similarity=$(total_time_sim) update=$(total_time_update)   reports: $(length(queries))"

    return [compute(m) for m in metrics]
end

end

max_stars_repo_name: irving-muller/FaST

max_issues_repo_name: irving-muller/FaST

max_forks_repo_name: irving-muller/FaST


text:
"""
    permute_loci!(data::PopData)
Edits `PopData` in place with loci permuted across populations within
the `.loci` dataframe.
"""
@inline function permute_loci!(data::PopData)
    @inbounds for locus in groupby(data.loci, :locus)
        shuffle!(locus.population)
    end
    data
end

"""
    permute_samples!(data::PopData; meta::Bool = false)
Edits `PopData` in place with samples permuted across populations within
the `.loci` dataframe. Since performance is important for many permutations,
the default is to only edit the `.loci` table in place; use `meta = true`
if you also require the `.meta` dataframe edited in place.
"""
@inline function permute_samples!(data::PopData; meta::Bool = false)
    pops = shuffle(data.meta.population)

    if meta == true
        meta_pops = deepcopy(pops)
        @inbounds for name in groupby(data.meta, :name)
            @inbounds name.population .= pop!(meta_pops)
        end
    end

    @inbounds for name in groupby(data.loci, :name)
        @inbounds name.population .= pop!(pops)
    end
    data
end

@inline function permute_samples!(data::AbstractDataFrame, popnames::Vector{String})
    pops = shuffle(popnames)

    @inbounds for name in groupby(data, :name)
        @inbounds name.population .= pop!(pops)
    end
    data
end

"""
    permute_genotypes!(data::PopData; by::String = "locus")
Edits `PopData` in place with genotypes permuted across individuals within
the `.loci` dataframe. Use `by = "population"` (or `"pop"`) to permute genotypes
within populations.
"""
@inline function permute_genotypes!(data::PopData; by::String = "locus")
    #establish mode of operation
    if by in ["locus", "loci"]
        groupings = :locus
    elseif by in ["population", "pop"]
        groupings = [:locus, :population]
    else
        error("Please choose from either \"locus\" or \"population\" run methods.")
    end
    @inbounds for grp in groupby(data.loci, groupings)
        grp.genotype .= strict_shuffle!(grp.genotype)
    end
    data
end

"""
    permute_alleles!(data::PopData; ploidy::Union{Nothing, Int} = nothing, by::String = "locus")
Edits `PopData` in place with alleles permuted and reconstructed into genotypes
for each locus within the `.loci` dataframe. Use `by = "population"` (or `"pop"`)
to permute alleles within populations. If `ploidy` is not provided (default `ploidy = nothing`),
then ploidy will be identified from the PopData. If performance is important,
it would be best to identify ploidy in advance and set it to a specific integer.
"""
@inline function permute_alleles!(data::PopData; ploidy::Union{Nothing, Int} = nothing, by::String = "locus")
    #establish mode of operation
    if by in ["locus", "loci"]
        groupings = :locus
    elseif by in ["population", "pop"]
        groupings = [:locus, :population]
    else
        error("Please choose from either \"locus\" or \"population\" run methods.")
    end

    if ploidy == nothing
        tmp = unique(data.meta.ploidy)
        length(tmp) > 1 && error("This permutation method is not appropriate for mixed-ploidy data")
        ploidy = tmp[1]
    end

    @inbounds for grp in groupby(data.loci, groupings)
        alle = shuffle(alleles(grp.genotype))
        new_genos = Tuple.(Base.Iterators.partition(alle, ploidy))
        (@view grp.genotype[.!ismissing.(grp.genotype)]) .= new_genos
    end
    data
end


text:
import StatsBase: sample
using Random, KernelDensity

function sample(rng::AbstractRNG, chn::MCMCChains.AbstractChains, n;
    replace=true, ordered=false)
  indxs = sample(rng,
    range(chn),
    n,
    replace=replace, ordered=ordered)
  return Chains(chn.value[indxs, :, :],
    names(chn),
    chn.name_map,
    info = chn.info)
end

function sample(chn::MCMCChains.AbstractChains, n;
    replace=true, ordered=false)
  indxs = sample(Random.GLOBAL_RNG,
    range(chn),
    n,
    replace=replace, ordered=ordered)
  return Chains(chn.value[indxs, :, :],
    names(chn),
    chn.name_map,
    info = chn.info)
end

function sample(rng::AbstractRNG, chn::MCMCChains.AbstractChains,
    wv::AbstractWeights, n)
  indxs = sample(rng,
    range(chn),
    wv,
    n)
  return Chains(chn.value[indxs, :, :],
    names(chn),
    chn.name_map,
    info = chn.info)
end

function sample(chn::MCMCChains.AbstractChains,
    wv::AbstractWeights, n)
  indxs = sample(Random.GLOBAL_RNG,
    range(chn),
    wv,
    n)
    return Chains(chn.value[indxs, :, :],
      names(chn),
      chn.name_map,
      info = chn.info)
end


text:
using Sockets
using Base.Threads: @spawn
using Dates
using MAT
using JLD2

const R = u"Ω"
const V = u"V"
const A = u"A"
const Hz = u"Hz"

using Unitful: Current, Voltage, Frequency, Time

raw(a::Current)   = Float64(ustrip(uconvert(A, a)))
raw(a::Voltage)   = Float64(ustrip(uconvert(V, a)))
raw(a::Frequency) = Float64(ustrip(uconvert(Hz, a)))
raw(a::Time)      = Float64(ustrip(uconvert(u"s", a)))

function elapsed_time(start_time)
    seconds = floor(time() - start_time)
    return Time(0) + Second(seconds)
end

function elapsed_time(func, start_time)
    seconds = floor(time() - start_time)
    return Time(0) + Second(func(seconds))
end

"""
    save(data)
    save(data; format=:matlab)
    save(data; filename="custom_file_name.ext")

Save data to a file

By default saves to julia format (.jld2) but can also export
data to matlab by using the format=:matlab keyword argument
"""
function save(data; filename = "", format = :julia)
    if isempty(filename)
        t = Dates.format(Dates.now(), "yy-mm-dd_HH:MM:SS")
        filename = "scan_" * t
    end
    if format == :julia
        @save (filename * ".jld2") data
    elseif format == :matlab
        if isa(data, ScopeData)
            data = ScopeData(data.info, ustrip(data.volt), data.time)
        end
        file = matopen(filename * ".mat", "w")
        write(file, "data", data)
        close(file)
    end
end

"""
    data = load("file.jld2")

Loads saved data from a file
"""
function load(filename)
    ext = split(filename, '.')[end]
    if ext == "jld2"
        jldopen(filename)["data"]
    else
        error("unsupported file type: $ext")
    end
end

"""
    scan_prologix(ip::AbstractString)

Given an IP Address to a prologix adapter return a
Dict of GPIB bus numbers that map to the connected device names
"""
function scan_prologix(ip::AbstractString)
    instr = initialize(Instrument, ip)
    return scan_prologix(instr)
end

"""
    scan_prologix(instr)

Given any instrument with a connection to a prologix adapter that
supports query and write. Return a Dict of GPIB bus numbers
that map to the connected device names
"""
function scan_prologix(obj)
    devices = Dict()
    for i in 0:15
        write(obj, "++addr $i")
        try
            devices[i] = query(obj, "*IDN?"; timeout = 0.5)
        catch

        end
    end
    return devices
end

"""
    scan_network(; network_id="10.1.30.0", host_range=1:255, v=false)
Will scan your network and report all found devices.

By searches for devices connected on port:
    - 5025 (scpi)
    - 1234 (prologix)
"""
function scan_network(; network="10.1.30.", host_range=1:255)
    network = ensure_ending_dot(network)
    @info "Scanning $network$(host_range[1])-$(host_range[end])"

    # Scan for SCPI devices
    ips_scpi = asyncmap(x-> _get_info_from_ip(x), [network*"$ip" for ip in host_range])
    # Scan for Prologix device
    ips_prlx = asyncmap(x-> _get_info_from_ip(x; port=1234), [network*"$ip" for ip in host_range])

    ips_all = vcat(ips_scpi, ips_prlx)
    print("\n")
    return [ip for ip in ips_all if !isempty(ip)]
end
ensure_ending_dot(network) = network[end] != '.' ? network*'.' : network

function _get_info_from_ip(ip_str; port = 5025)
    temp_ip = ip_str * ":$port"
    print(".")
    proc = @spawn temp_ip => _get_instr_info_and_close(temp_ip)
    sleep(2)
    if proc.state == :runnable
        print("t") # for time out
        kill_task(proc)
        return ""
    elseif proc.state == :done
        printstyled("s"; color = :green) # for success
        return fetch(proc)
    elseif proc.state == :failed
        print("f") # for failed
        return ""
    else
        error("Uncaught state: $(proc.state)")
    end
end

kill_task(proc) = schedule(proc, ErrorException("Timed out"), error=true)

function _get_instr_info_and_close(ip)

    obj = initialize(Instrument, ip)
    info_str = info(obj)
    close(obj)
    return info_str
end

"""
    scan_prologix(ip::AbstractString)
    Scans all GPIB addresses on a prologix device having the ip-address `ip`.
"""
function scan_prologix(ip::AbstractString)
    devices = Dict()
    prologix_port = ":1234"
    full_ip = ip * prologix_port
    obj = initialize(Instrument, ip)

    for i in 0:15
        write(obj, "++addr $i")
        try
            devices[i] = query(obj, "*IDN?"; timeout=0.5)
        catch

        end
    end
    return devices
end

udef(func) =  error("$(func) not implemented")

macro codeLocation()
    return quote
        st = stacktrace(backtrace())
        myf = ""
        for frm in st
            funcname = frm.func
            if frm.func != :backtrace && frm.func != Symbol("macro expansion")
                myf = frm.func
                break
            end
        end
        println(
            "Running function ",
            $("$(__module__)"),
            ".$(myf) at ",
            $("$(__source__.file)"),
            ":",
            $("$(__source__.line)"),
        )

        myf
    end
end

function alias_print(msg)
    printstyled("[ Aliasing: ", color = :blue, bold = true)
    println(msg)
end

"""
	split_str_into_host_and_port(str)
Splits a string like "192.168.1.1:5056" into ("192.168.1.1", 5056)
"""
function split_str_into_host_and_port(str::AbstractString)
	spl_str = split(str, ":")
    isempty(spl_str) && error("IP address string is empty!")
    host = spl_str[1]
    if length(spl_str) == 1
        port = 0
    else
        port = parse(Int, spl_str[2])
    end
    return (host, port)
end


text:
"""
The first part of make transform matrix, usually used by the `learn_paths` function to
save time and computing resources.
"""
function make_transform_fac end

"""
Using Cholesky decomposition to calculate the transformation matrix from S to C
or from C to S.
"""
function make_transform_matrix end

"""
    make_transform_fac(X::SparseMatrixCSC)

Calculate the first step of Cholesky decomposition for sparse matrices.
"""
function make_transform_fac(
    X::SparseMatrixCSC;
    method = :additive,
    shift = 0.02,
    multiplier = 1.01,
)

    XtX = X'X

    if method == :additive
        fac = cholesky(XtX, shift = shift)
    else
        # convert value to Float64
        # otherwise multiplier would raise error
        XtX = convert(SparseMatrixCSC{Float64,Int64}, XtX)
        for i = 1:size(XtX, 2)
            XtX[i, i] *= multiplier
        end
        fac = cholesky(XtX)
    end

    fac
end

"""
    make_transform_fac(X::Matrix)

Calculate the first step of Cholesky decomposition for dense matrices.
"""
function make_transform_fac(
    X::Matrix;
    method = :additive,
    shift = 0.02,
    multiplier = 1.01,
)

    XtX = X'X

    if method == :additive
        fac = cholesky(XtX + shift * I)
    else
        for i = 1:size(XtX, 2)
            XtX[i, i] *= multiplier
        end
        fac = cholesky(XtX)
    end

    fac
end

"""
    make_transform_matrix(fac::Union{LinearAlgebra.Cholesky, SuiteSparse.CHOLMOD.Factor}, X::Union{SparseMatrixCSC, Matrix}, Y::Union{SparseMatrixCSC, Matrix})

Second step in calculating the Cholesky decomposition for the transformation matrix.
"""
function make_transform_matrix(
    fac::Union{LinearAlgebra.Cholesky,SuiteSparse.CHOLMOD.Factor},
    X::Union{SparseMatrixCSC,Matrix},
    Y::Union{SparseMatrixCSC,Matrix};
    output_format = :auto,
    sparse_ratio = 0.05,
    verbose = false,
)

    M = fac \ (X'Y)
    format_matrix(
        M,
        output_format,
        sparse_ratio = sparse_ratio,
        verbose = verbose,
    )
end

"""
    make_transform_matrix(X::SparseMatrixCSC, Y::Matrix)

Use Cholesky decomposition to calculate the transformation matrix from X to Y,
where X is a sparse matrix and Y is a dense matrix.

# Obligatory Arguments
- `X::SparseMatrixCSC`: the X matrix, where X is a sparse matrix
- `Y::Matrix`: the Y matrix, where Y is a dense matrix

# Optional Arguments
- `method::Symbol = :additive`: whether :additive or :multiplicative decomposition is required
- `shift::Float64 = 0.02`: shift value for :additive decomposition
- `multiplier::Float64 = 1.01`: multiplier value for :multiplicative decomposition
- `output_format::Symbol = :auto`: to force output format to dense(:dense) or sparse(:sparse), make it auto(:auto) to determined by the program
- `sparse_ratio::Float64 = 0.05`: the ratio to decide whether a matrix is sparse
- `verbose::Bool = false`: if true, more information will be printed out

# Examples
```julia
# additive mode
JudiLing.make_transform_matrix(
    C,
    S,
    method = :additive,
    shift = 0.02,
    verbose = false)

# multiplicative mode
JudiLing.make_transform_matrix(
    C,
    S,
    method = :multiplicative,
    multiplier = 1.01,
    verbose = false)

# further control of sparsity ratio
JudiLing.make_transform_matrix(
  ...
    output_format = :auto,
    sparse_ratio = 0.05,
  ...)
```
"""
function make_transform_matrix(
    X::SparseMatrixCSC,
    Y::Matrix;
    method = :additive,
    shift = 0.02,
    multiplier = 1.01,
    output_format = :auto,
    sparse_ratio = 0.05,
    verbose = false,
)

    XtX = X'X

    if method == :additive
        fac = cholesky(XtX, shift = shift)
    else
        XtX = convert(SparseMatrixCSC{Float64,Int64}, XtX)
        for i = 1:size(XtX, 2)
            XtX[i, i] *= multiplier
        end
        fac = cholesky(XtX)
    end

    # M is in sparse format
    # but sometimes it is actually a dense matrix
    M = fac \ (X'Y)
    format_matrix(
        M,
        output_format,
        sparse_ratio = sparse_ratio,
        verbose = verbose,
    )
end

"""
    make_transform_matrix(X::Matrix, Y::Union{SparseMatrixCSC, Matrix})

Use the Cholesky decomposition to calculate the transformation matrix from X to Y,
where X is a dense matrix and Y is either a dense matrix or a sparse matrix.

# Obligatory Arguments
- `X::Matrix`: the X matrix, where X is a dense matrix
- `Y::Union{SparseMatrixCSC, Matrix}`: the Y matrix, where Y is either a sparse or a dense matrix

# Optional Arguments
- `method::Symbol = :additive`: whether :additive or :multiplicative decomposition is required
- `shift::Float64 = 0.02`: shift value for :additive decomposition
- `multiplier::Float64 = 1.01`: multiplier value for :multiplicative decomposition
- `output_format::Symbol = :auto`: to force output format to dense(:dense) or sparse(:sparse), make it auto(:auto) to determined by the program
- `sparse_ratio::Float64 = 0.05`: the ratio to decide whether a matrix is sparse
- `verbose::Bool = false`: if true, more information will be printed out

# Examples
```julia
# additive mode
JudiLing.make_transform_matrix(
    C,
    S,
    method = :additive,
    shift = 0.02,
    verbose = false)

# multiplicative mode
JudiLing.make_transform_matrix(
    C,
    S,
    method=:multiplicative,
    multiplier = 1.01,
    verbose = false)

# further control of sparsity ratio
JudiLing.make_transform_matrix(
    ...
    output_format = :auto,
    sparse_ratio = 0.05,
    ...)
```
"""
function make_transform_matrix(
    X::Matrix,
    Y::Union{SparseMatrixCSC,Matrix};
    method = :additive,
    shift = 0.02,
    multiplier = 1.01,
    output_format = :auto,
    sparse_ratio = 0.05,
    verbose = false,
)

    XtX = X'X

    if method == :additive
        fac = cholesky(XtX + shift * I)
    else
        XtX = convert(SparseMatrixCSC{Float64,Int64}, XtX)
        for i = 1:size(XtX, 2)
            XtX[i, i] *= multiplier
        end
        fac = cholesky(XtX)
    end

    # M is in sparse format
    # but sometimes it is actually a dense matrix
    M = fac \ (X'Y)
    format_matrix(
        M,
        output_format,
        sparse_ratio = sparse_ratio,
        verbose = verbose,
    )
end

"""
    make_transform_matrix(X::SparseMatrixCSC, Y::SparseMatrixCSC)

Use the Cholesky decomposition to calculate the transformation matrix from X to Y,
where X is a sparse matrix and Y is a sparse matrix.

# Obligatory Arguments
- `X::SparseMatrixCSC`: the X matrix, where X is a sparse matrix
- `Y::SparseMatrixCSC`: the Y matrix, where Y is a sparse matrix

# Optional Arguments
- `method::Symbol = :additive`: whether :additive or :multiplicative decomposition is required
- `shift::Float64 = 0.02`: shift value for :additive decomposition
- `multiplier::Float64 = 1.01`: multiplier value for :multiplicative decomposition
- `output_format::Symbol = :auto`: to force output format to dense(:dense) or sparse(:sparse), make it auto(:auto) to determined by the program
- `sparse_ratio::Float64 = 0.05`: the ratio to decide whether a matrix is sparse
- `verbose::Bool = false`: if true, more information will be printed out

# Examples
```julia
# additive mode
JudiLing.make_transform_matrix(
    C,
    S,
    method = :additive,
    shift = 0.02,
    verbose = false)

# multiplicative mode
JudiLing.make_transform_matrix(
    C,
    S,
    method = :multiplicative,
    multiplier = 1.01,
    verbose = false)

# further control of sparsity ratio
JudiLing.make_transform_matrix(
    ...
    output_format = :auto,
    sparse_ratio = 0.05,
    ...)
```
"""
function make_transform_matrix(
    X::SparseMatrixCSC,
    Y::SparseMatrixCSC;
    method = :additive,
    shift = 0.02,
    multiplier = 1.01,
    output_format = :auto,
    sparse_ratio = 0.05,
    verbose = false,
)

    XtX = X'X

    if method == :additive
        fac = cholesky(XtX, shift = shift)
    else
        # convert value to Float64
        # otherwise multiplier would raise error
        XtX = convert(SparseMatrixCSC{Float64,Int64}, XtX)
        for i = 1:size(XtX, 2)
            XtX[i, i] *= multiplier
        end
        fac = cholesky(XtX)
    end

    # M is in sparse format
    # but sometimes it is actually a dense matrix
    M = fac \ (X'Y)
    format_matrix(
        M,
        output_format,
        sparse_ratio = sparse_ratio,
        verbose = verbose,
    )
end

"""
    format_matrix(M::Union{SparseMatrixCSC, Matrix}, output_format=:auto)

Convert output matrix format to either a dense matrix or a sparse matrix.
"""
function format_matrix(
    M::Union{SparseMatrixCSC,Matrix},
    output_format = :auto;
    sparse_ratio = 0.05,
    verbose = false,
)

    if output_format == :dense
        verbose && println("Returning a dense matrix format")
        Array(M)
    elseif output_format == :sparse
        verbose && println("Returning a sparse matrix format")
        sparse(M)
    else
        verbose && print("Auto mode: ")
        if is_truly_sparse(M, threshold = sparse_ratio, verbose = verbose)
            verbose && println("Returning a sparse matrix format")
            return M
        else
            verbose && println("Returning a dense matrix format")
            return Array(M)
        end
    end
end


text:
##########################################################
# Rafał Nowak
#
# Instytut Informatyki
# Uniwersytet Wrocławski
#
#
# Pomocniczne materiały do wykładu z analizy numerycznej
#
##########################################################

f(x) = ( 1.0 - cos(x) ) ./ x.^2

# Zadanie: obliczyć f(1.2e-8)

# Rozwiązanie
x = 1.2e-8
y = f(x)    # 7.709882115452e-01

@printf("f(%.1e)=%.5e\n", x,y);

# Trochę dziwny ten wynik, bo łatwo udowodnić, że f(x) <= 0.5

max_stars_repo_name: Magikis/Uniwersity

max_issues_repo_name: Magikis/Uniwersity

max_forks_repo_name: Magikis/Uniwersity


text:
#=
We can determine how "out of order" an array A is by counting the number of inversions it has. Two elements A[i] and A[j] form an inversion if A[i] > A[j] but i < j. That is, a smaller element appears after a larger element.

Given an array, count the number of inversions it has. Do this faster than O(N^2) time.

You may assume each element in the array is distinct.

For example, a sorted list has zero inversions. The array [2, 4, 1, 3, 5] has three inversions: (2, 1), (4, 1), and (4, 3). The array [5, 4, 3, 2, 1] has ten inversions: every distinct pair forms an inversion.
=#

using Test
include("Solutions/problem44_count_inversions.jl")

@test inversions([2, 4, 1, 3, 5]) == 3
@test inversions([5, 4, 3, 2, 1]) == 10

max_stars_repo_name: DominiqueCaron/daily-coding-problem

max_issues_repo_name: DominiqueCaron/daily-coding-problem

max_forks_repo_name: DominiqueCaron/daily-coding-problem


text:
module OscarDevTools

export oscar_develop, oscar_update, oscar_branch, oscar_add_remotes

include("Helpers.jl")
include("OscarCI.jl")
include("DevUtils.jl")

using .DevUtils

# for old CI yml files
import .OscarCI: parse_meta, ci_matrix, github_json, parse_job

end # module


text:
struct VectorField{T,N,N2,L} <: AbstractVecArray{Complex{T},N}
    rr::VecArray{T, N, Array{T,N}, Array{T,N}, Array{T,N}}
    c::VecArray{Complex{T}, N, ScalarField{T,N,N2,L}, ScalarField{T,N,N2,L}, ScalarField{T,N,N2,L}}
    r::SubArray{Vec{T},N,VecArray{T,N,Array{T,N},Array{T,N},Array{T,N}},Tuple{Base.OneTo{Int},Vararg{Base.Slice{Base.OneTo{Int}},N2}},L}
    function VectorField{T,N,N2,L}(x::ScalarField{T,N,N2,L},y::ScalarField{T,N,N2,L},z::ScalarField{T,N,N2,L}) where {T,N,N2,L}
        c = VecArray(x,y,z)
        rr = VecArray(InplaceRealFFT.data(x),InplaceRealFFT.data(y),InplaceRealFFT.data(z))
        r = view(rr, Base.OneTo(size(x.field.r, 1)), ntuple(i->Colon(), Val(N2))...)
        return new{T,N,N2,L}(rr,c,r)
    end
end

@inline function Base.getproperty(a::S,s::Symbol) where {S<:VectorField}
    if (s === :kx || s === :ky || s === :kz || s === :k)
        return getproperty(getfield(getfield(a,:c),:x),s)
    else
        return getfield(a,s)
    end
end

@inline FluidTensors.xvec(v::VectorField) =
    FluidTensors.xvec(v.c)
@inline FluidTensors.yvec(v::VectorField) =
    FluidTensors.yvec(v.c)
@inline FluidTensors.zvec(v::VectorField) =
    FluidTensors.zvec(v.c)

VectorField(x::ScalarField{T,N,N2,L},y::ScalarField{T,N,N2,L},z::ScalarField{T,N,N2,L}) where {T,N,N2,L} = VectorField{T,N,N2,L}(x,y,z)

VectorField{T}(dims::NTuple{3,Int},l::NTuple{3,Real}) where T = VectorField(ScalarField{T}(dims,l),ScalarField{T}(dims,l),ScalarField{T}(dims,l))
VectorField(dims::NTuple{3,Int},l::NTuple{3,Real}) = VectorField(ScalarField(dims,l),ScalarField(dims,l),ScalarField(dims,l))

Base.similar(a::VectorField) = VectorField(similar(a.c.x),similar(a.c.y),similar(a.c.z))

function InplaceRealFFT.rfft!(v::VectorField)
    rfft!(v.c.x)
    rfft!(v.c.y)
    rfft!(v.c.z)
    return v
end

function InplaceRealFFT.brfft!(v::VectorField)
    brfft!(v.c.x)
    brfft!(v.c.y)
    brfft!(v.c.z)
    return v.r
end

function InplaceRealFFT.irfft!(v::VectorField)
    irfft!(v.c.x)
    irfft!(v.c.y)
    irfft!(v.c.z)
    return v.r
end

isrealspace(a::VectorField) =
    isrealspace(a.c.x) && isrealspace(a.c.y) && isrealspace(a.c.z)

Base.show(io::IO,m::MIME"text/plain",a::VectorField) =
    isrealspace(a) ? show(io,m,a.r) : show(io,m,a.c)

setreal!(a::VectorField) = (a.c.x.realspace[] = true; a.c.y.realspace[] = true; a.c.z.realspace[] = true)
setfourier!(a::VectorField) = (a.c.x.realspace[] = false; a.c.y.realspace[] = false; a.c.z.realspace[] = false)

fourier!(a::VectorField) = (fourier!(a.c.x); fourier!(a.c.y); fourier!(a.c.z); a)

real!(a::VectorField) = (real!(a.c.x); real!(a.c.y); real!(a.c.z); a.r)


text:
"""
PriorityQueue datastructure (based on a implicit binary heap).
Code copied from DataStructures.jl (long time ago)
"""
module Heap
import Base.Order: Ordering, Forward, lt

struct PriorityQueue{K,V,O <: Ordering}
    xs::Vector{Pair{K,V}}
    index::Vector{Int}    # Map elements to their index in xs
    o::O

    function PriorityQueue{K,V}(o::O=Forward) where {K,V,O <: Ordering}
        new{K,V,O}(Vector{Pair{K,V}}(), Vector{Int}(), o)
    end

    function PriorityQueue{K,V}(n::Integer, o::O=Forward) where {K,V,O <: Ordering}
        q = PriorityQueue{K,V}(o)
        sizehint!(q, n)
        q
    end
end

@inline Base.length(pq::PriorityQueue) = length(pq.xs)

@inline Base.isempty(pq::PriorityQueue) = isempty(pq.xs)

@inline Base.haskey(pq::PriorityQueue, key) = pq.index[key] > 0

function Base.sizehint!(pq::PriorityQueue, n::Integer)
    sizehint!(pq.xs, n)
    resize!(pq.index, n)
    pq.index .= 0
end

@inline function percolate_down!(pq::PriorityQueue{K,V}, i::Integer) where {K,V}
    heapleft(i::Integer) = 2i
    heapright(i::Integer) = 2i + 1

    x = pq.xs[i]
    @inbounds while (l = heapleft(i)) <= length(pq)
        r = heapright(i)
        j = r > length(pq) || lt(pq.o, pq.xs[l].second, pq.xs[r].second) ? l : r
        if lt(pq.o, pq.xs[j].second, x.second)
            pq.index[pq.xs[j].first] = i
            pq.xs[i] = pq.xs[j]
            i = j
        else
            break
        end
    end
    pq.index[x.first] = i
    pq.xs[i] = x
end

@inline function percolate_up!(pq::PriorityQueue{K,V}, i::Integer) where {K,V}
    heapparent(i::Integer) = div(i, 2)

    x = pq.xs[i]
    @inbounds while i > 1
        j = heapparent(i)
        if lt(pq.o, x.second, pq.xs[j].second)
            pq.index[pq.xs[j].first] = i
            pq.xs[i] = pq.xs[j]
            i = j
        else
            break
        end
    end
    pq.index[x.first] = i
    pq.xs[i] = x
end

"""
Change the priority of an existing element, or enqueue it if it isn't present.
"""
function Base.setindex!(pq::PriorityQueue{K,V}, value::V, key::K) where {K,V}
    if haskey(pq, key)
        i = pq.index[key]
        oldvalue = pq.xs[i].second
        pq.xs[i] = Pair{K,V}(key, value)
        if lt(pq.o, oldvalue, value)
            percolate_down!(pq, i)
        else
            percolate_up!(pq, i)
        end
    else
        enqueue!(pq, key, value)
    end
    value
end

enqueue!(pq::PriorityQueue, key::K, value::V) where {K,V} =
    enqueue!(pq, key => value)

function enqueue!(pq::PriorityQueue{K,V}, pair::Pair{K,V}) where {K,V}
    key = pair.first
    if haskey(pq, key)
        throw(ArgumentError("PriorityQueue keys must be unique"))
    end
    push!(pq.xs, pair)
    pq.index[key] = length(pq)
    percolate_up!(pq, length(pq))

    return pq
end

function dequeue!(pq::PriorityQueue{K,V}) where {K,V}
    x = pq.xs[1]
    y = pop!(pq.xs)
    if !isempty(pq)
        pq.xs[1] = y
        pq.index[y.first] = 1
        percolate_down!(pq, 1)
    end
    # delete!(pq.index, x.first)
    pq.index[x.first] = 0
    x.first
end

end

max_stars_repo_name: EQt/graphidx

max_issues_repo_name: EQt/graphidx

max_forks_repo_name: EQt/graphidx


text:
"""
"""
function writevtk(args...;kwargs...)
  map(visualization_data(args...;kwargs...)) do visdata
    write_vtk_file(
    visdata.grid,visdata.filebase,celldata=visdata.celldata,nodaldata=visdata.nodaldata)
  end
end

"""
"""
function createvtk(args...;kwargs...)
  v = visualization_data(args...;kwargs...)
  @notimplementedif length(v) != 1
  visdata = first(v)
  create_vtk_file(
    visdata.grid,visdata.filebase,celldata=visdata.celldata,nodaldata=visdata.nodaldata)
end

"""
"""
function createpvd(args...;kwargs...)
  paraview_collection(args...;kwargs...)
end

function createpvd(parts::Nothing,args...;kwargs...)
  createpvd(args...;kwargs...)
end

function createpvd(f,parts::Nothing,args...;kwargs...)
  createpvd(f,args...;kwargs...)
end

"""
"""
function savepvd(pvd::WriteVTK.CollectionFile)
  vtk_save(pvd)
end

"""

    write_vtk_file(
      trian::Grid,
      filebase;
      celldata=Dict(),
      nodaldata=Dict())

Low level entry point to vtk. Other vtk-related routines in Gridap eventually call this one.

"""
function write_vtk_file(
  trian::Grid, filebase; celldata=Dict(), nodaldata=Dict())
  vtkfile = create_vtk_file(trian, filebase, celldata=celldata, nodaldata=nodaldata)
  outfiles = vtk_save(vtkfile)
end

"""

    create_vtk_file(
      trian::Grid,
      filebase;
      celldata=Dict(),
      nodaldata=Dict())

Low level entry point to vtk. Other vtk-related routines in Gridap eventually call this one.
This function only creates the vtkFile, without writing to disk.

"""
function create_vtk_file(
  trian::Grid, filebase; celldata=Dict(), nodaldata=Dict())

  points = _vtkpoints(trian)
  cells = _vtkcells(trian)
  vtkfile = vtk_grid(filebase, points, cells, compress=false)

  for (k,v) in celldata
    vtk_cell_data(vtkfile, _prepare_data(v), k)
  end
  for (k,v) in nodaldata
    vtk_point_data(vtkfile, _prepare_data(v), k)
  end

  return vtkfile
end

function create_pvtk_file(
  trian::Grid, filebase;
  part, nparts, ismain=(part==1), celldata=Dict(), nodaldata=Dict())

  points = _vtkpoints(trian)
  cells = _vtkcells(trian)
  vtkfile = pvtk_grid(filebase, points, cells, compress=false;
                      part=part, nparts=nparts, ismain=ismain)

  for (k,v) in celldata
    vtkfile[k,VTKCellData()] = _prepare_data(v)
  end
  for (k,v) in nodaldata
    vtkfile[k,VTKPointData()] = _prepare_data(v)
  end

  return vtkfile
end

function _vtkpoints(trian)
  D = num_point_dims(trian)
  x = get_node_coordinates(trian)
  xflat = collect(x)
  reshape(reinterpret(Float64,xflat),(D,length(x)))
end

function _vtkcells(trian)

  type_to_reffe = get_reffes(trian)
  cell_to_type = get_cell_type(trian)
  type_to_vtkid = map(get_vtkid, type_to_reffe)
  type_to_vtknodes = map(get_vtknodes, type_to_reffe)
  cell_to_nodes = get_cell_node_ids(trian)
  cache = array_cache(cell_to_nodes)

  _generate_vtk_cells(
    cache,
    cell_to_nodes,
    cell_to_type,
    type_to_vtkid,
    type_to_vtknodes)
end

function _generate_vtk_cells(
  cache,
  cell_to_nodes,
  cell_to_type,
  type_to_vtkid,
  type_to_vtknodes)

  V = eltype(cell_to_nodes)
  meshcells = MeshCell{WriteVTK.VTKCellTypes.VTKCellType,V}[]

  d = _vtkcelltypedict()

  cells = 1:length(cell_to_type)
  for cell in cells

    t = cell_to_type[cell]
    vtkid = type_to_vtkid[t]
    vtknodes = type_to_vtknodes[t]

    nodes = getindex!(cache,cell_to_nodes,cell)
    meshcell = MeshCell(d[vtkid], nodes[vtknodes])

    push!(meshcells,meshcell)

  end

  meshcells

end

_prepare_data(v) = v

function _prepare_data(v::AbstractArray{<:MultiValue})
  a = collect(v)
  reinterpret(a)
end

function _prepare_data(v::AbstractArray{<:VectorValue{2,T}}) where T
  a = collect(v)
  b = reshape(reinterpret(T,a),(2,length(a)))
  z = zeros((1,size(b,2)))
  vcat(b,z)
end

"""
"""
function get_vtkid(reffe::LagrangianRefFE)
  basis = get_prebasis(reffe)
  p = get_polytope(reffe)
  get_vtkid(p,basis)
end

"""
"""
function get_vtknodes(reffe::LagrangianRefFE)
  basis = get_prebasis(reffe)
  p = get_polytope(reffe)
  get_vtknodes(p,basis)
end

function get_vtkid(p::ExtrusionPolytope, basis::MonomialBasis)
  exponents = get_exponents(basis)
  vtkid, _ = _vtkinfo_extrusion_polytope(p,exponents)
  vtkid
end

function get_vtknodes(p::ExtrusionPolytope, basis::MonomialBasis)
  exponents = get_exponents(basis)
  _, vtknodes = _vtkinfo_extrusion_polytope(p,exponents)
  vtknodes
end

function get_vtkid(p::SerendipityPolytope,basis::MonomialBasis)
  get_vtkid(p.hex,basis)
end

function get_vtknodes(p::SerendipityPolytope,basis::MonomialBasis)
  get_vtknodes(p.hex,basis)
end

function _vtkinfo_extrusion_polytope(p,exponents)

  # Taken from the vtk specification
  # https://vtk.org/wp-content/uploads/2015/04/file-formats.pdf

  n_nodes = length(exponents)

  if p == VERTEX
    if n_nodes == 1
      vtkid = 1
      vtknodes = [1,]
    else
      @notimplemented
    end

  elseif p == SEGMENT
    if n_nodes == 2
      vtkid = 3
      vtknodes = [1,2]
    elseif n_nodes == 3
      vtkid = 21
      vtknodes = [1,2,3]
    else
      @notimplemented
    end

  elseif p == TRI
    if n_nodes == 3
      vtkid = 5
      vtknodes = [1,2,3]
    elseif n_nodes == 6
      vtkid = 22
      vtknodes = [1,2,3,4,6,5]
    elseif n_nodes == 10
      vtkid = 69
      vtknodes = [1,2,3,4,5,8,9,7,6,10]
    elseif n_nodes == 15
      vtkid = 69
      vtknodes = [1,2,3,4,5,6,10,11,12,9,8,7,13,14,15]
    elseif n_nodes == 21
      vtkid = 69
      vtknodes = [1,2,3,4,5,6,7,12,13,14,15,11,10,9,8,16,18,21,17,20,19]
    else
      @notimplemented
    end

  elseif p == QUAD
    if n_nodes == 4
      vtkid = 9
      vtknodes = [1,2,4,3]
    elseif n_nodes == 8
      vtkid = 23
      vtknodes = [1,2,4,3,5,8,6,7]
    elseif n_nodes == 9
      vtkid = 28
      vtknodes = [1,2,4,3,5,8,6,7,9]
    elseif n_nodes == 16
      vtkid = 70
      vtknodes = [1,2,4,3,5,6,11,12,7,8,9,10,13,14,15,16]
    elseif n_nodes == 25
      vtkid = 70
      vtknodes = [1,2,4,3,5,6,7,14,15,16,8,9,10,11,12,13,17,18,19,20,21,22,23,24,25]
    else
      @notimplemented
    end

  elseif p == TET
    if n_nodes == 4
      vtkid = 10
      vtknodes = [1,2,3,4]
    elseif n_nodes == 10
      vtkid = 24
                 #0,1,2,3,4,5,6,7,8,9
      vtknodes = [1,2,3,4,5,7,6,8,9,10]
    elseif n_nodes == 20
      vtkid = 71
      vtknodes = [1,2,3,4,5,6,9,10,7,8,11,12,13,14,15,16,18,20,19,17]
    else
      @notimplemented
    end

  elseif p == HEX
    if n_nodes == 8
      vtkid = 12
      vtknodes = [1,2,4,3,5,6,8,7]
    elseif n_nodes == 27
      vtkid = 72
      vtknodes = [1,2,4,3,5,6,8,7,9,14,10,13,11,16,12,15,17,18,19,20,25,26,23,24,21,22,27]
    elseif n_nodes == 64
      vtkid = 72
      vtknodes = [1,2,4,3,5,6,8,7,9,10,19,20,11,12,17,18,13,14,23,24,15,16,21,22,25,26,27,
                  28,29,30,31,32,49,50,51,52,53,54,55,56,41,42,43,44,45,46,47,48,33,34,35,
                  36,37,38,39,40,57,58,59,60,61,62,63,64]
    else
      @notimplemented
    end

  elseif p == WEDGE
    if n_nodes == 6
      vtkid = 13
      vtknodes = [1,3,2,4,6,5]
    else
      @notimplemented
    end

  elseif p == PYRAMID
    if n_nodes == 5
      vtkid = 14
      vtknodes = [1,2,4,3,5]
    else
      @notimplemented
    end

  else
    @notimplemented "vtkid not implemented for given ExtrusionPolytope"
  end

  (vtkid, vtknodes)
end

function _vtkcelltypedict()
  d = Dict{Int,WriteVTK.VTKCellTypes.VTKCellType}()
  d[VTK_VERTEX.vtk_id] = VTK_VERTEX
  d[VTK_LINE.vtk_id] = VTK_LINE
  d[VTK_LINE.vtk_id] = VTK_LINE
  d[VTK_TRIANGLE.vtk_id] = VTK_TRIANGLE
  d[VTK_QUAD.vtk_id] = VTK_QUAD
  d[VTK_TETRA.vtk_id] = VTK_TETRA
  d[VTK_HEXAHEDRON.vtk_id] = VTK_HEXAHEDRON
  d[VTK_WEDGE.vtk_id] = VTK_WEDGE
  d[VTK_QUADRATIC_QUAD.vtk_id] = VTK_QUADRATIC_QUAD
  d[VTK_BIQUADRATIC_QUAD.vtk_id] = VTK_BIQUADRATIC_QUAD
  d[VTK_QUADRATIC_TRIANGLE.vtk_id] = VTK_QUADRATIC_TRIANGLE
  d[VTK_QUADRATIC_TETRA.vtk_id] = VTK_QUADRATIC_TETRA
  d[VTK_QUADRATIC_EDGE.vtk_id] = VTK_QUADRATIC_EDGE
  d[VTK_QUADRATIC_HEXAHEDRON.vtk_id] = VTK_QUADRATIC_HEXAHEDRON
  d[VTK_PYRAMID.vtk_id] = VTK_PYRAMID
  d[VTK_LAGRANGE_TRIANGLE.vtk_id] = VTK_LAGRANGE_TRIANGLE
  d[VTK_LAGRANGE_QUADRILATERAL.vtk_id] = VTK_LAGRANGE_QUADRILATERAL
  d[VTK_LAGRANGE_TETRAHEDRON.vtk_id] = VTK_LAGRANGE_TETRAHEDRON
  d[VTK_LAGRANGE_HEXAHEDRON.vtk_id] = VTK_LAGRANGE_HEXAHEDRON
  #d[VTK_BIQUADRATIC_HEXAHEDRON.vtk_id] = VTK_BIQUADRATIC_HEXAHEDRON
  d
end