import Base: *, reduce

import LinearAlgebra
import LinearAlgebra: inv, det, logdet, logabsdet, \, /

using Statistics
using LinearAlgebra: Diagonal, Transpose, Adjoint, diagm, diag

struct TrackedArray{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
  tracker::Tracked{A}
  data::A
  grad::A
  TrackedArray{T,N,A}(t::Tracked{A}, data::A) where {T,N,A} = new(t, data)
  TrackedArray{T,N,A}(t::Tracked{A}, data::A, grad::A) where {T,N,A} = new(t, data, grad)
end

data(x::TrackedArray) = x.data
tracker(x::TrackedArray) = x.tracker

TrackedVector{T,A} = TrackedArray{T,1,A}
TrackedMatrix{T,A} = TrackedArray{T,2,A}
TrackedVecOrMat{T,A} = Union{TrackedVector{T,A},TrackedMatrix{T,A}}

track(c::Call, x::AbstractArray) = TrackedArray(c, x)

TrackedArray(c::Call, x::A) where A <: AbstractArray =
  TrackedArray{eltype(A),ndims(A),A}(Tracked{A}(c), x)

TrackedArray(c::Call, x::A, Δ::A) where A <: AbstractArray =
  TrackedArray{eltype(A),ndims(A),A}(Tracked{A}(c, Δ), x, Δ)

TrackedArray(x::AbstractArray) = TrackedArray(Call(), x, zero(x))

Base.eltype(x::Type{<:TrackedArray{T}}) where T <: Real = TrackedReal{T}

Base.convert(::Type{T}, x::S) where {T<:TrackedArray,S<:T} = x

Base.convert(T::Type{<:TrackedArray}, x::TrackedArray) =
  error("Not implemented: convert $(typeof(x)) to $T")

Base.convert(::Type{<:TrackedArray{T,N,A}}, x::AbstractArray) where {T,N,A} =
  TrackedArray(convert(A, x))

Base.show(io::IO, t::Type{TrackedArray{T,N,A}}) where {T,N,A<:AbstractArray{T,N}} =
  @isdefined(A) ?
    print(io, "TrackedArray{…,$A}") :
    invoke(show, Tuple{IO,DataType}, io, t)

function Base.summary(io::IO, x::TrackedArray)
  print(io, "Tracked ")
  summary(io, data(x))
end

Base.print_array(io::IO, x::TrackedArray) = Base.print_array(io, data(x))

function Base.show(io::IO, x::TrackedArray)
  show(io, data(x))
  print(io, " (tracked)")
end

Base.copy(x::TrackedArray) = x

collect(xs::TrackedArray) = xs

Base.setindex!(xs::TrackedArray, v, i...; kwargs...) =
  error("Can't differentiate `setindex!`")

back!(::TrackedArray) = error("Value is not scalar; use `back!(sum(x))` or `back!(x, Δ)`")

function update!(x::TrackedArray, Δ)
  x.data .+= data(Δ)
  tracker(x).grad .= 0
  return x
end

function update!(x::AbstractArray, Δ)
  x .+= data(Δ)
  return x
end

# Fallthrough methods

for f in :[Base.size, Base.ndims, Base.collect].args
  @eval @inline $f(x::TrackedArray, a...) = $f(data(x), a...)
end

Base.size(x::TrackedArray, i::Integer, j::Integer, is::Integer...) =
  size(data(x), i, j, is...)

Base.similar(x::TrackedArray, T::Type, dims::Dims) = similar(data(x), T, dims)

for op in [:(==), :≈]
    @eval Base.$op(x::TrackedArray, y::AbstractArray) = Base.$op(data(x), y)
    @eval Base.$op(x::AbstractArray, y::TrackedArray) = Base.$op(x, data(y))
    @eval Base.$op(x::TrackedArray, y::TrackedArray) = Base.$op(data(x), data(y))
end

# Array Stdlib

Base.getindex(xs::TrackedArray, i...; kwargs...) = track(getindex, xs, i...; kwargs...)

@grad function getindex(xs::AbstractArray, i...; kwargs...)
  getindex(data(xs), i...; kwargs...), function (Δ)
        Δ′ = zero(xs)
        setindex!(Δ′, data(Δ), i...; kwargs...)
        (nobacksies(:getindex, Δ′), map(_->nothing, i)...)
    end
end

@grad function getindex(xs::AbstractArray, i::Array...)
  data(xs)[i...], function (Δ)
    Δ′ = zero(xs)
    @views Δ′[i...] .+= data(Δ)
    (nobacksies(:getindex, Δ′), map(_->nothing, i)...)
  end
end

Base.view(x::TrackedArray, inds...; kwargs...) = track(Base.view, x, inds...; kwargs...)

@grad function view(x::AbstractArray, inds...; kwargs...)
    view(data(x), inds...; kwargs...), function (Δ)
        grad_output = zero(x)
        subgrad = view(grad_output, inds...; kwargs...)
        subgrad[:] = data(Δ)
        (nobacksies(:view, grad_output), map(_->nothing, inds)...)
    end
end

Base.:-(xs::TrackedArray) = track(-, xs)

@grad -(xs) = -data(xs), Δ -> (-Δ,)

Base.transpose(xs::TrackedArray) = track(transpose, xs)
Base.adjoint(xs::TrackedArray) = track(adjoint, xs)

@grad transpose(xs) = transpose(data(xs)), Δ -> (trim(xs, transpose(Δ)),)
@grad adjoint(xs) = data(xs)', Δ -> (trim(xs, Δ'),)

det(xs::TrackedArray) = track(det, xs)
@grad det(xs) = det(data(xs)), Δ -> (Δ * det(xs) * transpose(inv(xs)),)

logdet(xs::TrackedArray) = track(logdet, xs)
@grad logdet(xs) = logdet(data(xs)), Δ -> (Δ * transpose(inv(xs)),)

logabsdet(xs::TrackedArray) = track(logabsdet, xs)
@grad logabsdet(xs) = logabsdet(data(xs)), Δ -> (Δ[1] * transpose(inv(xs)),)

Base.repeat(xs::TrackedArray; kw...) = track(repeat, xs; kw...)

@grad function repeat(xs; inner=ntuple(x->1, ndims(xs)), outer=ntuple(x->1, ndims(xs)))
  repeat(data(xs), inner = inner, outer = outer), function (Δ)
    Δ′ = zero(xs)
    S = size(xs)

    # Loop through each element of Δ, calculate source dimensions, accumulate into Δ′
    for (dest_idx, val) in pairs(IndexCartesian(), data(Δ))
        # First, round dest_idx[dim] to nearest gridpoint defined by inner[dim], then
        # wrap around based on original size S.
        src_idx = [mod1(div(dest_idx[dim] - 1, inner[dim]) + 1, S[dim]) for dim in 1:length(S)]
        Δ′[src_idx...] += val
    end
    (nobacksies(:repeat, Δ′),)
  end
end

function combinations(xs, n)
  n < 1 && return [[]]
  cs = combinations(xs, n-1)
  [[x, c...] for x in xs, c in cs]
end

for i = 0:2, c = combinations([:AbstractArray, :TrackedArray, :Number], i), f = [:hcat, :vcat]
  cnames = map(_ -> gensym(), c)
  @eval Base.$f($([:($x::$c) for (x, c) in zip(cnames, c)]...), x::Union{TrackedArray,TrackedReal}, xs::Union{AbstractArray,Number}...) =
    track($f, $(cnames...), x, xs...)
end

for i = 0:2, c = combinations([:AbstractVecOrMat, :TrackedVecOrMat], i), f = [:hcat, :vcat]
  cnames = map(_ -> gensym(), c)
  @eval Base.$f($([:($x::$c{T}) for (x, c) in zip(cnames, c)]...), x::TrackedVecOrMat{T}, xs::AbstractVecOrMat{T}...) where T =
    track($f, $(cnames...), x, xs...)
end

for i = 0:2, c = combinations([:AbstractVector, :TrackedVector], i), f = [:hcat, :vcat]
  cnames = map(_ -> gensym(), c)
  @eval Base.$f($([:($x::$c{T}) for (x, c) in zip(cnames, c)]...), x::TrackedVector{T}, xs::AbstractVector{T}...) where T =
    track($f, $(cnames...), x, xs...)
end

@grad function vcat(xs...)
  vcat(data.(xs)...), function (Δ)
    start = 0
    Δs = [begin
      x = map(_ -> :, size(xsi))
      i = isempty(x) ? x : Base.tail(x)
      d = Δ[start+1:start+size(xsi,1), i...]
      start += size(xsi, 1)
      d
    end for xsi in xs]
    return (Δs...,)
  end
end

@grad function hcat(xs...)
  hcat(data.(xs)...), function (Δ)
    start = 0
    Δs = [begin
      d = if ndims(xsi) == 1
        Δ[:, start+1]
      else
        i = map(_ -> :, size(xsi)) |> Base.tail |> Base.tail
        Δ[:, start+1:start+size(xsi,2), i...]
      end
      start += size(xsi, 2)
      d
    end for xsi in xs]
    return (Δs...,)
  end
end

for i = 0:2, c = combinations([:AbstractArray, :TrackedArray], i)
  cnames = map(_ -> gensym(), c)
  @eval Base.cat($([:($x::$c) for (x, c) in zip(cnames, c)]...), x::TrackedArray, xs::AbstractArray...; dims) =
    track(cat, $(cnames...), x, xs..., dims = dims)
end

@grad function cat(Xs...; dims)
  cat(data.(Xs)..., dims = dims), function (Δ)
    start = ntuple(i -> 0, Val(ndims(Δ)))
    Δs = [begin
      dim_xs = 1:ndims(xs)
      till_xs = ntuple((i -> i in dims ? (i in dim_xs ? size(xs,i) : 1) : 0), Val(ndims(Δ)))
      xs_in_Δ = ntuple(i -> i in dims ? (start[i]+1:start[i]+till_xs[i]) : Colon(), Val(ndims(Δ)))
      d = reshape(Δ[xs_in_Δ...],size(xs))
      start = start .+ till_xs
      d
    end for xs in Xs]
    return (Δs...,)
  end
end

function ∇reduce(f::typeof(vcat), Δ, xs)
  start = 0
  Δs = [begin
    x = map(_ -> :, size(xsi))
    i = isempty(x) ? x : Base.tail(x)
    d = Δ[start+1:start+size(xsi,1), i...]
    start += size(xsi, 1)
    d
  end for xsi in xs]
  return Δs
end

function ∇reduce(f::typeof(hcat), Δ, xs)
  start = 0
  Δs = [begin
    d = if ndims(xsi) == 1
      Δ[:, start+1]
    else
      i = map(_ -> :, size(xsi)) |> Base.tail |> Base.tail
      Δ[:, start+1:start+size(xsi,2), i...]
    end
    start += size(xsi, 2)
    d
  end for xsi in xs]
  return Δs
end

for f in [:vcat, :hcat]
  @eval function reduce(::typeof($f), xs::Vector{<:TrackedArray})
    y = reduce($f, data.(xs))
    track(Call(reduce, ($f, xs)), y)
  end
end

function scan(c::Call{typeof(reduce)})
    foreach(scan, tracker.(c.args[2]))
end

function back_(c::Call{typeof(reduce)}, Δ, once)
  f, xs = c.args
  foreach((x, d) -> back(x, d, once), tracker.(xs), data.(∇reduce(f, Δ, xs)))
end

Base.reshape(xs::TrackedArray, dims::Union{Colon,Int}...) = reshape(xs, dims)
Base.reshape(xs::TrackedArray, dims::Tuple{Vararg{Union{Int,Colon}}}) = reshape(xs, Base._reshape_uncolon(xs, dims))
Base.reshape(xs::TrackedArray, dims::Tuple{Vararg{Int}}) = track(reshape, xs, dims)

@grad reshape(xs, dims) = reshape(data(xs), dims), Δ -> (reshape(Δ, size(xs)),nothing)

Base.permutedims(xs::TrackedArray, perm) = track(permutedims, xs, perm)
@grad permutedims(xs, perm) = permutedims(data(xs), perm), Δ -> (permutedims(Δ, invperm(perm)),nothing)

Base.PermutedDimsArray(xs::TrackedArray, perm) = track(PermutedDimsArray, xs, perm)
@grad PermutedDimsArray(xs, perm) = PermutedDimsArray(data(xs), perm), Δ -> (PermutedDimsArray(Δ, invperm(perm)),nothing)

Base.reverse(xs::TrackedArray; dims) = track(reverse, xs, dims = dims)
@grad reverse(xs; dims) = reverse(data(xs), dims = dims), Δ -> (reverse(Δ, dims = dims), nothing)
Base.reverse(xs::TrackedVector) = track(reverse, xs)
@grad reverse(xs::TrackedVector) = reverse(data(xs)), Δ -> (reverse(Δ),)
Base.reverse(xs::TrackedVector, start::Integer, stop::Integer) = track(reverse, xs, start, stop)
@grad reverse(xs, start, stop) = reverse(data(xs), start, stop), Δ -> (reverse(Δ, start, stop), nothing, nothing)

function _kron(mat1::AbstractMatrix,mat2::AbstractMatrix)
    m1, n1 = size(mat1)
    mat1_rsh = reshape(mat1,(1,m1,1,n1))

    m2, n2 = size(mat2)
    mat2_rsh = reshape(mat2,(m2,1,n2,1))

    return reshape(mat1_rsh.*mat2_rsh, (m1*m2,n1*n2))
end

Base.kron(a::TrackedMatrix, b::TrackedMatrix)  = _kron(a, b)
Base.kron(a::TrackedMatrix, b::AbstractMatrix) = _kron(a, b)
Base.kron(a::AbstractMatrix, b::TrackedMatrix) = _kron(a, b)

inv(A::TrackedArray) = Tracker.track(inv, A)
@grad function inv(A)
    return inv(Tracker.data(A)), function (Δ)
        Ainv = inv(A)
        ∇A = - Ainv' * Δ * Ainv'
        return (∇A, )
    end
end

#       (/) rdivide
A::TrackedArray     / B::TrackedArray     = Tracker.track(/, A, B)
A::AbstractVecOrMat / B::TrackedArray     = Tracker.track(/, A, B)
A::TrackedArray     / B::AbstractVecOrMat = Tracker.track(/, A, B)
@grad function Base.:/(A, B)
    return Tracker.data(A) / Tracker.data(B), function (Δ)
        ∇A = Δ / B'
        ∇B = - (A / B)' * ∇A
        return (∇A, ∇B)
    end
end

#       (\) ldivide  (left vec divide needs more work to resolve dispatch ambiguity)
A::TrackedArray     \ B::TrackedArray     = Tracker.track(\, A, B)
A::AbstractArray    \ B::TrackedArray     = Tracker.track(\, A, B)
A::TrackedArray     \ B::AbstractVecOrMat = Tracker.track(\, A, B)
A::AbstractMatrix   \ B::TrackedVecOrMat  = Tracker.track(\, A, B)
A::TrackedMatrix    \ B::TrackedVecOrMat  = Tracker.track(\, A, B)
@grad function Base.:\(A, B)
    return Tracker.data(A) \ Tracker.data(B), function (Δ)
        ∇B = A' \ Δ
        ∇A = - ∇B * (A \ B)'
        return (∇A, ∇B)
    end
end

# Reductions

Base.sum(xs::TrackedArray; dims = :) = track(sum, xs, dims = dims)
Base.sum(f::Union{Function,Type},xs::TrackedArray) = sum(f.(xs))

@grad sum(xs; dims = :) = sum(data(xs), dims = dims),
  Δ -> (zero(xs) .+ Δ, )

Base.prod(xs::TrackedArray; dims=:) = track(prod, xs; dims=dims)
Base.prod(f::Union{Function, Type}, xs::TrackedArray) = prod(f.(xs))

@grad function prod(xs; dims=:)
  p = prod(data(xs); dims=dims)
  p, Δ -> (p ./ xs .* Δ,)
end

Base.findfirst(xs::TrackedArray, args...) = findfirst(xs.data, args...)

Statistics.mean(xs::TrackedArray; dims = :) = track(mean, xs, dims = dims)

Base.maximum(xs::TrackedArray; dims = :) = track(maximum, xs, dims = dims)
Base.minimum(xs::TrackedArray; dims = :) = track(minimum, xs, dims = dims)

import LinearAlgebra: dot

dot(xs::TrackedArray, ys::TrackedArray) = track(dot, xs, ys)
dot(xs::AbstractArray, ys::TrackedArray) = track(dot, xs, ys)
dot(xs::TrackedArray, ys::AbstractArray) = track(dot, xs, ys)

@grad dot(xs, ys) = dot(data(xs), data(ys)), Δ -> (Δ .* ys, Δ .* xs)

# Hacks to get std working
Statistics.std(x::TrackedArray; dims = :, mean = Statistics.mean(x, dims = dims), corrected::Bool = true) = _std(x,mean,dims,corrected)
_std(x::TrackedArray, mean, dims, corrected) = sqrt.(sum((x .- mean).^2, dims = dims) ./ (mapreduce(i -> size(x,i),*, dims) - corrected))
_std(x::TrackedArray, mean, ::Colon, corrected) = sqrt.(sum((x .- mean).^2) ./ (length(x) - corrected))

LinearAlgebra.norm(x::TrackedArray{T}, p::Real = 2) where T =
  (sum(abs.(x).^p) + eps(T))^(oneunit(T) / p) # avoid d(sqrt(x))/dx == Inf at 0

@grad mean(xs; dims = :) = mean(data(xs), dims=dims), Δ -> (_backmean(xs,Δ,dims),)
_backmean(xs, Δ, ::Colon) = zero(xs) .+ Δ ./ length(xs)
_backmean(xs, Δ, dims) = zero(xs) .+ Δ ./ mapreduce(i -> size(data(xs),i),*,dims)

@grad function maximum(xs; dims = dims)
  maximum(data(xs), dims = dims), function (Δ)
    Δ′ = zero(xs)
    _, i = findmax(data(xs), dims = dims)
    Δ′[i] = data(Δ)
    return (nobacksies(:maximum, Δ′),)
  end
end

@grad function minimum(xs;  dims = dims)
  minimum(data(xs),  dims = dims), function (Δ)
    Δ′ = zero(xs)
    _, i = findmin(data(xs),  dims = dims)
    Δ′[i] = data(Δ)
    return (nobacksies(:minimum, Δ′),)
  end
end

# BLAS

LinearAlgebra.diagm(x::Pair{<:Integer, <:TrackedVector}) = track(diagm, x...)
@grad diagm(i, x) = diagm(i => data(x)), Δ -> (nothing, diag(Δ, i))

x::TrackedMatrix  * y::AbstractMatrix = track(*, x, y)
x::AbstractMatrix * y::TrackedMatrix  = track(*, x, y)
x::TrackedMatrix  * y::TrackedMatrix  = track(*, x, y)

x::TrackedMatrix  * y::AbstractVector = track(*, x, y)
x::AbstractMatrix * y::TrackedVector  = track(*, x, y)
x::TrackedMatrix  * y::TrackedVector  = track(*, x, y)

x::TrackedVector  * y::AbstractVector = track(*, x, y)
x::AbstractVector * y::TrackedVector  = track(*, x, y)
x::TrackedVector  * y::TrackedVector  = track(*, x, y)

# Ambiguity fixes
Base.:*(x::Transpose{T,<:AbstractVector{T}},y::TrackedMatrix) where {T} = track(*, x, y)
Base.:*(x::TrackedMatrix,y::Transpose{T,<:AbstractVector{T}}) where {T} = track(*, x, y)

Base.:*(x::Transpose{T,<:AbstractMatrix{T}},y::TrackedMatrix) where {T} = track(*, x, y)
Base.:*(x::TrackedMatrix,y::Transpose{T,<:AbstractMatrix{T}}) where {T} = track(*, x, y)

Base.:*(x::Transpose{T,<:AbstractVector{T}},y::TrackedVector) where {T} = track(*, x, y)
Base.:*(x::TrackedVector,y::Transpose{T,<:AbstractVector{T}}) where {T} = track(*, x, y)

Base.:*(x::Transpose{T,<:AbstractMatrix{T}},y::TrackedVector) where {T} = track(*, x, y)
Base.:*(x::TrackedVector,y::Transpose{T,<:AbstractMatrix{T}}) where {T} = track(*, x, y)

Base.:*(x::Adjoint{T,<:AbstractVector{T}},y::TrackedMatrix) where {T} = track(*, x, y)
Base.:*(x::TrackedMatrix,y::Adjoint{T,<:AbstractVector{T}}) where {T} = track(*, x, y)

Base.:*(x::Adjoint{T,<:AbstractMatrix{T}},y::TrackedMatrix) where {T} = track(*, x, y)
Base.:*(x::TrackedMatrix,y::Adjoint{T,<:AbstractMatrix{T}}) where {T} = track(*, x, y)

Base.:*(x::Adjoint{T,<:AbstractVector{T}},y::TrackedVector) where {T} = track(*, x, y)
Base.:*(x::TrackedVector,y::Adjoint{T,<:AbstractVector{T}}) where {T} = track(*, x, y)

Base.:*(x::Adjoint{T,<:AbstractMatrix{T}},y::TrackedVector) where {T} = track(*, x, y)
Base.:*(x::TrackedVector,y::Adjoint{T,<:AbstractMatrix{T}}) where {T} = track(*, x, y)

Base.:*(x::Diagonal, y::TrackedVector) = track(*, x, y)

Base.:*(x::Diagonal, y::TrackedMatrix) = track(*, x, y)
Base.:*(x::TrackedMatrix, y::Diagonal) = track(*, x, y)

@grad a::AbstractVecOrMat * b::AbstractVecOrMat =
  data(a)*data(b), Δ -> (Δ * transpose(b), transpose(a) * Δ)

# NNlib

using NNlib
import NNlib: softmax, ∇softmax, logsoftmax, ∇logsoftmax, conv, ∇conv_data, depthwiseconv, maxpool, meanpool
import NNlib: DenseConvDims, DepthwiseConvDims, PoolDims

softmax(xs::TrackedArray; dims=1) = track(softmax, xs; dims=dims)

@grad softmax(xs; dims=1) = softmax(data(xs); dims=dims), Δ -> (nobacksies(:softmax, ∇softmax(data(Δ), data(xs); dims=dims)),)

logsoftmax(xs::TrackedArray; dims=1) = track(logsoftmax, xs; dims=dims)

@grad logsoftmax(xs; dims=1) = logsoftmax(data(xs); dims=dims), Δ -> (nobacksies(:logsoftmax, ∇logsoftmax(data(Δ), data(xs); dims=dims)),)

depthwiseconv(x::TrackedArray, w::TrackedArray, cdims::DepthwiseConvDims; kw...) = track(depthwiseconv, x, w, cdims; kw...)
depthwiseconv(x::AbstractArray, w::TrackedArray, cdims::DepthwiseConvDims; kw...) = track(depthwiseconv, x, w, cdims; kw...)
depthwiseconv(x::TrackedArray, w::AbstractArray, cdims::DepthwiseConvDims; kw...) = track(depthwiseconv, x, w, cdims; kw...)

@grad depthwiseconv(x, w, cdims::DepthwiseConvDims; kw...) =
  depthwiseconv(data(x), data(w), cdims; kw...),
    Δ -> nobacksies(:depthwiseconv,
      (NNlib.∇depthwiseconv_data(data.((Δ, w))..., cdims; kw...),
       NNlib.∇depthwiseconv_filter(data.((x, Δ))..., cdims; kw...),
       nothing))

conv(x::TrackedArray,  w::TrackedArray, cdims::DenseConvDims;  kw...) = track(conv, x, w, cdims; kw...)
conv(x::AbstractArray, w::TrackedArray, cdims::DenseConvDims;  kw...) = track(conv, x, w, cdims; kw...)
conv(x::TrackedArray,  w::AbstractArray, cdims::DenseConvDims; kw...) = track(conv, x, w, cdims; kw...)

@grad conv(x, w, cdims::DenseConvDims; kw...) =
  conv(data(x), data(w), cdims; kw...),
    Δ -> nobacksies(:conv,
      (NNlib.∇conv_data(data.((Δ, w))..., cdims; kw...),
       NNlib.∇conv_filter(data.((x, Δ))..., cdims; kw...),
       nothing))

∇conv_data(x::TrackedArray,  w::TrackedArray, cdims::DenseConvDims;  kw...) = track(∇conv_data, x, w, cdims; kw...)
∇conv_data(x::AbstractArray, w::TrackedArray, cdims::DenseConvDims;  kw...) = track(∇conv_data, x, w, cdims; kw...)
∇conv_data(x::TrackedArray,  w::AbstractArray, cdims::DenseConvDims; kw...) = track(∇conv_data, x, w, cdims; kw...)

@grad function ∇conv_data(y, w, cdims::DenseConvDims; kw...)
  return (
    ∇conv_data(data(y), data(w), cdims; kw...),
    Δ -> begin
      return nobacksies(:conv,
        (NNlib.conv(data.((Δ, w))..., cdims; kw...),
         NNlib.∇conv_filter(data.((Δ, y))..., cdims; kw...),
         nothing)
      )
    end
  )
end

maxpool(x::TrackedArray, pdims::PoolDims; kw...) = track(maxpool, x, pdims; kw...)

@grad function maxpool(x, pdims::PoolDims; kw...)
  y = maxpool(data(x), pdims; kw...)
  y, Δ -> (nobacksies(:maxpool, NNlib.∇maxpool(data.((Δ, y, x))..., pdims; kw...)), nothing)
end

meanpool(x::TrackedArray, pdims::PoolDims; kw...) = track(meanpool, x, pdims; kw...)

@grad function meanpool(x, pdims::PoolDims; kw...)
  y = meanpool(data(x), pdims; kw...)
  y, Δ -> (nobacksies(:meanpool, NNlib.∇meanpool(data.((Δ, y, x))..., pdims; kw...)), nothing)
end

# Broadcasting

using ForwardDiff: Dual, partials, value

trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val(ndims(x))))

unbroadcast(x::AbstractArray, Δ) =
  size(x) == size(Δ) ? Δ :
  length(x) == length(Δ) ? trim(x, Δ) :
    trim(x, sum(Δ, dims = ntuple(i -> size(x, i) == 1 ? i : ndims(Δ)+1, Val(ndims(Δ)))))

unbroadcast(x::Number, Δ) = sum(Δ)
unbroadcast(x::Base.RefValue, _) = nothing

dual(x, p) = x
dual(x::Real, p) = Dual(x, p)

function partial(f::F, Δ, i, args::Vararg{Any,N}) where {F,N}
  dargs = ntuple(j -> dual(args[j], i==j), Val(N))
  return Δ * f(dargs...).partials[1]
end

@inline function ∇broadcast(f::F, args::Vararg{Any,N}) where {F,N}
  y = broadcast(f, data.(args)...)
  (eltype(y) <: Real && eltype(y) !== Bool) || return y
  function back(Δ)
    Δargs = ntuple(i -> partial.(f, Δ, i, args...), Val(N))
    dxs = map(unbroadcast, args, Δargs)
    return dxs
  end
  # So we can return non-tracked arrays
  track(Call(back, tracker.(args)), y)
end

using Base.Broadcast: BroadcastStyle, ArrayStyle, Broadcasted, broadcasted

struct TrackedStyle <: BroadcastStyle end

Broadcast.BroadcastStyle(::Type{<:Union{TrackedArray,TrackedReal}}) = TrackedStyle()
Broadcast.BroadcastStyle(::TrackedStyle, ::BroadcastStyle) = TrackedStyle()

# We have to re-build the original broadcast struct to get the appropriate array
# style. We need this primarily to support CuArrays' broadcasting fixes.
broadcast_rebuild(xs) = data(xs)

broadcast_rebuild(bc::Broadcasted) =
  Broadcasted(bc.f, map(broadcast_rebuild, bc.args))

preprocess(x) = x

function Base.copy(bc::Broadcasted{TrackedStyle})
  bc1 = Broadcast.flatten(bc)
  bc2 = Broadcast.flatten(broadcast_rebuild(bc))
  ∇broadcast(bc2.f, bc1.args...)
end

using Requires

# https://github.com/FluxML/Flux.jl/issues/353
if VERSION < v"1.1.0-DEV.548"
  @init Requires.isprecompiling() || @eval Base.Broadcast begin
    function flatten(bc::Broadcasted{Style}) where {Style}
      isflat(bc) && return bc
      args = cat_nested(bc)
      let makeargs = make_makeargs(bc), f = bc.f
        newf = @inline function(args::Vararg{Any,N}) where N
          f(makeargs(args...)...)
        end
        return Broadcasted{Style}(newf, args, bc.axes)
      end
    end
    @inline function make_makeargs(makeargs, t::Tuple{<:Broadcasted,Vararg{Any}})
      bc = t[1]
      let makeargs = make_makeargs(makeargs, tail(t)), f = bc.f
        let makeargs = make_makeargs(makeargs, bc.args)
          headargs, tailargs = make_headargs(bc.args), make_tailargs(bc.args)
          return @inline function(args::Vararg{Any,N}) where N
            args1 = makeargs(args...)
            a, b = headargs(args1...), tailargs(args1...)
            (f(a...), b...)
          end
        end
      end
    end
  end
end


module AnalysisTools

using LinearAlgebra
using Plots
using Plots.PlotMeasures
using DataFrames
using Glob
using JSON

using ACEtb.FitTools

export plot_bands, plot_bands!, tabulate_results, comp_plot

function plot_bands(fc::FitContext, phase::Symbol, fermi_level::Float64, bands::Matrix; plotobj=nothing, label=nothing, color=:black)
    @assert phase in [:FCC, :BCC]
    phase = String(phase)

    bandpath = JSON.parsefile(joinpath(fc.base_dir, "$(phase)_tb_bandpath.json"))
    kpoints = bandpath["kpoints"]
    special_points = bandpath["special_points"]

    k0 = 0.0
    x_k = []
    tick_p = [k0]
    tick_k = phase == "FCC" ? ["W"] : tick_k = ["G"]
    push!(x_k, k0)
    for xi=2:size(kpoints, 1)
        k0 += norm(kpoints[xi] .- kpoints[xi-1])
        push!(x_k, k0)
        indict = [k for (k, v) in special_points if v==kpoints[xi] ]
        if length(indict) > 0
            push!(tick_p, k0)
            push!(tick_k, indict[1])
        end
    end

    plotobj === nothing && (plotobj = plot())
    for i = 1:size(bands, 2)
        thislabel = i == 1 ? label : false
        plot!(plotobj, x_k, bands[:, i], color=color, lw=1, label=thislabel)
    end
    for p=1:length(tick_p)
        plot!(plotobj, [tick_p[p]], lw=0.5, c=:gray, seriestype=:vline, label=false)
    end
    plot!(plotobj, [fermi_level], lw=0.5, c=:black, ls=:dash, seriestype=:hline, label=false)
    plot!(plotobj; fg_legend=:transparent, bg_legend=:transparent,
          grid=false, framestyle = :box, legend=:bottomleft)
    xticks!(plotobj, ([t for t in tick_p],[k for k in tick_k]), fontsize=14)
    ylabel!(plotobj, "Energy (eV)")

    xlims!(plotobj, 0.0, maximum(x_k))
    ylims!(plotobj, -20, fermi_level + 10)

    plot(plotobj, dpi=120)
    return plotobj
end

function plot_bands(fc::FitContext, phase::Symbol, param_hash::String; label=nothing, kwargs...)
    fermi_level, bands = read_bands(fc, phase, param_hash)
    label === nothing && (label = param_hash)
    plot_bands(fc, phase, fermi_level, bands; kwargs...,
               label=label)
end

plot_bands!(plotobj::Plots.Plot, args...; kwargs...) = plot_bands(args...; plotobj=plotobj, kwargs...)

function tabulate_results(fc::FitContext)
    df = DataFrame(param_hash=String[],
                   order=Int[],
                   degree=Int[],
                   env_deg=Int[],
                   rcut=Float64[],
                   renv=Float64[],
                   pcut=Int[],
                   N_train=Int[],
                   N_bond=Int[],
                   N_basis=Int[],
                   error_train=Float64[],
                   error_test=Union{Missing,Float64}[],
                   band_error_FCC=Union{Missing,Float64}[],
                   band_error_BCC=Union{Missing,Float64}[],
                   pot_file=String[],
                   training_datasets=String[],
                   method=String[],
                   rtol=Union{Missing,Float64}[])
    for pot_file in glob(joinpath(splitpath(fc.output_dir)[end], "pot_*/pot.json"))
        param_hash = match(r"pot_(.*)/pot.json", pot_file).captures[1]
        pot_dict = JSON.parsefile(pot_file)
        band_error = get(pot_dict, "band_error", missing)
        band_error_FCC = band_error_BCC = missing
        if band_error !== missing
            band_error_FCC = band_error["FCC"]
            band_error_BCC = band_error["BCC"]
        end
        fit_params = pot_dict["model"]["fit_params"]
        cutoff_params = pot_dict["model"]["cutoff_params"]
        row = (param_hash,
               fit_params["order"],
               fit_params["degree"],
               fit_params["env_deg"],
               cutoff_params["rcut"],
               cutoff_params["renv"],
               cutoff_params["pcut"],
               length(pot_dict["training_datasets"]),
               pot_dict["size_A"][1],
               pot_dict["size_A"][2],
               pot_dict["error_train"],
               get(pot_dict, "test_error", missing),
               band_error_FCC,
               band_error_BCC,
               pot_file,
               join(pot_dict["training_datasets"], ","),
               get(fit_params, "method", "QR"),
               get(fit_params, "rtol", missing))
        push!(df, row)
    end

    # add a descriptive label for the types of training configs
    df.config = ifelse.(occursin.("005", df.training_datasets),
                        ifelse.(occursin.("010", df.training_datasets), "mixT", "lowT"),
                        "highT")

    return df
end

function ref_plots(fc::FitContext)
    fcc_fermi_level_FHI, fcc_bands_FHI = read_FHI_bands(fc, :FCC)
    bcc_fermi_level_FHI, bcc_bands_FHI = read_FHI_bands(fc, :BCC)
    fcc_plot = plot_bands(fc, :FCC, fcc_fermi_level_FHI, fcc_bands_FHI; label="FHIaims ref", color=:red)
    bcc_plot = plot_bands(fc, :BCC, bcc_fermi_level_FHI, bcc_bands_FHI; label="FHIaims ref", color=:red)
    display(plot(fcc_plot, bcc_plot, title=["FCC" "BCC"]))
    return fcc_plot, bcc_plot
end

function add_plots(fc::FitContext, potential, color, fcc_plot, bcc_plot; label=nothing)
    plot_bands!(fcc_plot, fc, :FCC, potential, color=color, label=label)
    plot_bands!(bcc_plot, fc, :BCC, potential, color=color, label=label)
    p = plot(fcc_plot, bcc_plot, title=["FCC" "BCC"])
    display(p)
    return p
end

function comp_plot(fc::FitContext, potential; color=:blue, label=nothing)
    fcc_plot, bcc_plot = ref_plots(fc)
    add_plots(fc, potential, color, fcc_plot, bcc_plot, label=label)
end

end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 090f6f0f-fd67-419b-a35e-5392a59f39b1
using PlutoTest

# ╔═╡ e89c666e-b0e7-11eb-22cf-7f6fd1e299bb
md"""
# [Day 11: Corporate Policy](https://adventofcode.com/2015/day/11)

Santa's previous password expired, and he needs help choosing a new one.

To help him remember his new password after the old one expires, Santa has devised a method of coming up with a password based on the previous one. Corporate policy dictates that passwords must be exactly eight lowercase letters (for security reasons), so he finds his new password by incrementing his old password string repeatedly until it is valid.

Incrementing is just like counting with numbers: `xx`, `xy`, `xz`, `ya`, `yb`, and so on. Increase the rightmost letter one step; if it was `z`, it wraps around to `a`, and repeat with the next letter to the left until one doesn't wrap around.

Unfortunately for Santa, a new Security-Elf recently started, and he has imposed some additional password requirements:

-   Passwords must include one increasing straight of at least three letters, like `abc`, `bcd`, `cde`, and so on, up to `xyz`. They cannot skip letters; `abd` doesn't count.
-   Passwords may not contain the letters `i`, `o`, or `l`, as these letters can be mistaken for other characters and are therefore confusing.
-   Passwords must contain at least two different, non-overlapping pairs of letters, like `aa`, `bb`, or `zz`.

For example:

-   `hijklmmn` meets the first requirement (because it contains the straight `hij`) but fails the second requirement requirement (because it contains `i` and `l`).
-   `abbceffg` meets the third requirement (because it repeats `bb` and `ff`) but fails the first requirement.
-   `abbcegjk` fails the third requirement, because it only has one double letter (`bb`).
-   The next password after `abcdefgh` is `abcdffaa`.
-   The next password after `ghijklmn` is `ghjaabcc`, because you eventually skip all the passwords that start with `ghi...`, since `i` is not allowed.

**Given Santa's current password (your puzzle input), what should his next password be?**
"""

# ╔═╡ ec540096-d9d1-4b5f-b6a2-16c470322424
puzzle_input = "hxbxwxba"

# ╔═╡ 4e20e380-e928-4764-a982-3c0ca6bb942c
function next(str)
	chars = collect(str)
	vmin = Int('a')
	vmax = Int('z')

	for idx in reverse(1:length(str))
		v = Int(chars[idx]) + 1
		if v > vmax
			v -= vmax-vmin + 1
			chars[idx] = v
		else
			chars[idx] = v
			break
		end
	end

	join(chars)
end

# ╔═╡ 7a68015d-0282-4cb0-bebc-c265d5ca72d3
@assert next("xx") == "xy"

# ╔═╡ c4d6ad5b-f500-4435-a8eb-76c1f33f9775
@assert next("xy") == "xz"

# ╔═╡ 8f2717d2-ed5e-4e4e-abcd-6ff2982d8cd6
@assert next("xz") == "ya"

# ╔═╡ 58f76b47-542f-438b-ac71-694a76d9d77b
@assert next("ya") == "yb"

# ╔═╡ 9a03d3c5-3bc6-48c7-87cc-49ef6cbd31d4
function has_straight(s)
	for idx in 1:length(s)-2
		a = Int(s[idx])
		b = Int(s[idx+1])
		c = Int(s[idx+2])

		if b == a + 1 && c == b + 1
			return true
		end
	end
	false
end

# ╔═╡ b10b13e9-de70-4c09-8b7a-9cd530d1e3ac
@assert has_straight("hijklmmn")

# ╔═╡ 2d9e27ff-5557-4c7e-b6a8-b3accd15f7f0
@assert !has_straight("abbceffg")

# ╔═╡ 26342478-769e-4ef1-a3f3-ddf5aad2e281
has_no_confusing(s) = !occursin("i", s) && !occursin("o", s) && !occursin("l", s)

# ╔═╡ a470ee13-e98c-423b-8398-e6c8774983c5
@assert !has_no_confusing("hijklmmn")

# ╔═╡ dd4c2b0b-d8af-411f-9290-d36398772267
@assert has_no_confusing("abbceffg")

# ╔═╡ a0681275-076f-400f-8862-74c0cba5ffc1
# Passwords must contain at least two different, non-overlapping pairs of letters, like aa, bb, or zz.

# ╔═╡ 2144ecf3-79c2-4de1-8aa4-52b4e6b57208
function has_two_nonoverlapping_pair(s)
	for idx1 in 1:length(s)-1
		if s[idx1] == s[idx1+1]
			for idx2 in idx1+2:length(s)-1
				if s[idx2] == s[idx2+1]
					return true
				end
			end
		end
	end
	false
end

# ╔═╡ d103a1b1-7ade-41da-bc1e-b01337333675
@assert has_two_nonoverlapping_pair("abbceffg")

# ╔═╡ 5a13d4d9-7451-4fb0-8aaf-8d47947c10e4
check1(s) = has_straight(s) && has_no_confusing(s) && has_two_nonoverlapping_pair(s)

# ╔═╡ de4134e2-cc3f-4872-9244-4596d8525292
function next_valid(s)
	while true
		s = next(s)
		if check1(s)
			return s
		end
	end
end

# ╔═╡ 6b671ffc-ccc1-4b84-9ddb-7f7137995348
@assert next_valid("abcdefgh") == "abcdffaa"

# ╔═╡ a33d8ddb-5ef5-4c0c-891c-eef1368fbcc9
@assert next_valid("ghijklmn") == "ghjaabcc"

# ╔═╡ cbcd58ec-5f81-4ca4-8a4f-2a570db79f2a
part1 = next_valid(puzzle_input)

# ╔═╡ 2d3ab22d-11fe-4ee3-9ee1-9410ae03498a
@test part1 == "hxbxxyzz"

# ╔═╡ bd64f7b1-4074-4eeb-925d-eb8cc7a524bf
md"Your puzzle answer was `hxbxxyzz`."

# ╔═╡ aa05c55c-e79d-4846-95e1-823f0a5ea2e3
md"""
# Part Two
Santa's password expired again. **What's the next one?**
"""

# ╔═╡ 1ef77364-44a3-4819-b0c3-11cca842e53c
part2 = next_valid("hxbxxyzz")

# ╔═╡ fb47e75f-a3a0-4ee9-97c8-8a7e3d3f7831
@test part2 == "hxcaabcc"

# ╔═╡ bd67cd4e-af24-403c-a3ba-9753f11dc21c
md"Your puzzle answer was `hxcaabcc`."

# ╔═╡ Cell order:
# ╠═090f6f0f-fd67-419b-a35e-5392a59f39b1
# ╟─e89c666e-b0e7-11eb-22cf-7f6fd1e299bb
# ╠═ec540096-d9d1-4b5f-b6a2-16c470322424
# ╠═4e20e380-e928-4764-a982-3c0ca6bb942c
# ╠═7a68015d-0282-4cb0-bebc-c265d5ca72d3
# ╠═c4d6ad5b-f500-4435-a8eb-76c1f33f9775
# ╠═8f2717d2-ed5e-4e4e-abcd-6ff2982d8cd6
# ╠═58f76b47-542f-438b-ac71-694a76d9d77b
# ╠═9a03d3c5-3bc6-48c7-87cc-49ef6cbd31d4
# ╠═b10b13e9-de70-4c09-8b7a-9cd530d1e3ac
# ╠═2d9e27ff-5557-4c7e-b6a8-b3accd15f7f0
# ╠═26342478-769e-4ef1-a3f3-ddf5aad2e281
# ╠═a470ee13-e98c-423b-8398-e6c8774983c5
# ╠═dd4c2b0b-d8af-411f-9290-d36398772267
# ╠═a0681275-076f-400f-8862-74c0cba5ffc1
# ╠═2144ecf3-79c2-4de1-8aa4-52b4e6b57208
# ╠═d103a1b1-7ade-41da-bc1e-b01337333675
# ╠═5a13d4d9-7451-4fb0-8aaf-8d47947c10e4
# ╠═de4134e2-cc3f-4872-9244-4596d8525292
# ╠═6b671ffc-ccc1-4b84-9ddb-7f7137995348
# ╠═a33d8ddb-5ef5-4c0c-891c-eef1368fbcc9
# ╠═cbcd58ec-5f81-4ca4-8a4f-2a570db79f2a
# ╠═2d3ab22d-11fe-4ee3-9ee1-9410ae03498a
# ╟─bd64f7b1-4074-4eeb-925d-eb8cc7a524bf
# ╟─aa05c55c-e79d-4846-95e1-823f0a5ea2e3
# ╠═1ef77364-44a3-4819-b0c3-11cca842e53c
# ╠═fb47e75f-a3a0-4ee9-97c8-8a7e3d3f7831
# ╟─bd67cd4e-af24-403c-a3ba-9753f11dc21c


module CloudCovErr

include("cov.jl")

include("perstar.jl")

include("preprocess.jl")

include("decam.jl")
using .decam

end


################################################################################

"""

    extractSignalBin(edfDf::DataFrame, params::Dict)

# Description
Use `extractSignalBin` on EDF file per channel from shell arguments. Returns a dictionary with channel names as keys.

See also: [`extractFFT`](@ref)
"""
function extractSignalBin(edfDf::DataFrame, params::Dict)
  if haskey(params, "window-size") && haskey(params, "bin-overlap")
    return extractSignalBin(edfDf, binSize = params["window-size"], binOverlap = params["bin-overlap"])
  else
    @error "Variables are not defined in dictionary"
  end
end

################################################################################

"""

    extractSignalBin(channel::Array;
    binSize::T, binOverlap::T) where T <: Number

# Description
Bin channel signal.

See also: [`extractFFT`](@ref)
"""
function extractSignalBin(channel::Array; binSize::T, binOverlap::T) where T <: Number
  # define variables
  stepSize = floor(Int64, binSize / binOverlap)
  signalSteps = 1:stepSize:length(channel)
  signalAr = Array{Float64}(
    undef,
    length(signalSteps),
    binSize
  )

  # iterate over signal bins
  for stepIx ∈ 1:length(signalSteps)
    signalBoundry = signalSteps[stepIx] + binSize - 1

    # extract signals
    if signalBoundry <= length(channel)
      signalAr[stepIx, :] = channel[signalSteps[stepIx]:signalBoundry]

    # adjust last bin
    elseif signalBoundry > length(channel)
      signalAr[stepIx, :] = [
        channel[signalSteps[stepIx]:end];
        (signalBoundry - length(channel) |> abs |> zeros)
      ]
    end

  end
  return signalAr
end

################################################################################

"""

    extractSignalBin(edfDf::DataFrames.DataFrame;
    binSize::T, binOverlap::T) where T <: Number

# Description
Use `extractSignalBin` on EDF file per channel. Returns a dictionary with channel names as keys.

See also: [`extractFFT`](@ref)
"""
function extractSignalBin(edfDf::DataFrames.DataFrame; binSize::T, binOverlap::T) where T <: Number
  @info("Binning channels signals...")
  channelDc = Dict()
  signalAr = begin
    stepSize = floor(Int64, binSize / binOverlap)
    signalSteps = 1:stepSize:size(edfDf, 1)
    Array{Float64}(
      undef,
      1,
      binSize,
      length(signalSteps)
    )
  end

  # iterate on dataframe channels
  for (ψ, ε) ∈ enumerate(names(edfDf))
    tmpAr = extractSignalBin(edfDf[:, ψ], binSize = binSize, binOverlap = binOverlap)
    for β ∈ 1:size(tmpAr, 1)
      signalAr[:, :, β] = tmpAr[β, :]
    end
    channelDc[ε] = copy(signalAr)
  end
  return channelDc
end

################################################################################


using SnpParser

using Test

# Single-port Z parameterexample
snp = TouchstoneSnP("example/one_port_z_params")

# Test SnP parameters
@test snp.frequency_unit == "MHz"
@test snp.number_of_ports == 1
@test snp.parameter_type == "Z"
@test snp.number_of_frequencies == 5
@test snp.reference_resistance == 20.0

# Test network data
@test snp.network_data[100.0] == [74.25 -4.0]
@test snp.network_data[200.0] == [60.0 -22.0]
@test snp.network_data[300.0] == [53.025 -45.0]
@test snp.network_data[400.0] == [30.0 -62.0]
@test snp.network_data[500.0] == [0.75 -89.0]


# http://julialang.org/blog/2013/05/graphical-user-interfaces-part2/

using Graphics
using Cairo
using Tk

win = Toplevel("Test", 400, 200)
c = Canvas(win)
pack(c, expand=true, fill="both")
ctx = getgc(c)
# Set coordinates to go from 0 to 10 within a 300x100 centered region
set_coords(ctx, 50, 50, 300, 100, 0, 10, 0, 10)
set_source_rgb(ctx, 0, 0, 1)   # set color to blue
paint(ctx)                     # paint the entire clip region
move_to(ctx, -1, 5)
line_to(ctx, 7, 6)
set_source_rgb(ctx, 1, 0, 0)
set_line_width(ctx, 5)
stroke(ctx)
reveal(c)
Tk.update()
reset_clip(ctx)
rectangle(ctx, 7, 5, 4, 4)
set_source_rgb(ctx, 0, 1, 0)
fill_preserve(ctx)
set_source_rgb(ctx, 1, 0, 1)
stroke(ctx)
reveal(c)
Tk.update()


using DataStructures

"Results structure"
struct UninformedSearchResult{TState}
  status::Symbol
  path::Vector{TState}
  closedsetsize::Int64
  opensetsize::Int64
end

struct UninformedSearchNode{TState}
  data::TState
  depth::Int64
  parent::Union{UninformedSearchNode{TState}, Nothing}
end

mutable struct DepthFirstSearchState{TState, THash}
  openset::Stack{UninformedSearchNode{TState}}
  closedset::Set{THash}
  start_time::Float64
end

function _depthfirst!(
  search_state::DepthFirstSearchState{TState, THash},
  neighbours,
  start,
  goal,
  isgoal,
  hashfn,
  timeout,
  maxdepth,
  enable_closedset,
) where {TState, THash}
  while !isempty(search_state.openset)
    node = pop!(search_state.openset)
    node_data = node.data
    if isgoal(node_data, goal)
      return UninformedSearchResult(
        :success,
        reconstructpath(node),
        length(search_state.closedset),
        length(search_state.openset),
      )
    end

    if timeout < Inf && time() - search_state.start_time >= timeout
      return UninformedSearchResult(
        :timeout,
        reconstructpath(node),
        length(search_state.closedset),
        length(search_state.openset),
      )
    end

    neighbours_data = collect(neighbours(node_data))
    # reverse so that the neighbours are checked in order
    reverse!(neighbours_data)
    for neighbour in neighbours_data
      new_depth = node.depth + 1
      if (enable_closedset && hashfn(neighbour) in search_state.closedset) ||
         new_depth > maxdepth
        continue
      end

      neighbour_node = UninformedSearchNode(neighbour, new_depth, node)
      push!(search_state.openset, neighbour_node)
    end

    if enable_closedset
      push!(search_state.closedset, hashfn(node_data))
    end
  end

  return UninformedSearchResult(
    :nopath,
    [start],
    length(search_state.closedset),
    length(search_state.openset),
  )
end

"""
Execute a depth-first search given the neighbours function, a start and a goal.
It takes the same parameters as astar, but it ignores those about cost and heuristic.
Use the `maxdepth` parameter to limit the search to a certain depth.
"""
function depthfirst(
  neighbours,
  start,
  goal;
  isgoal = defaultisgoal,
  hashfn = hash,
  timeout = Inf,
  maxdepth = typemax(Int64),
  enable_closedset = true,
  kwargs...,
)
  start_time = time()
  start_hash = hashfn(start)

  start_node = UninformedSearchNode(start, 0, nothing)
  stack = Stack{typeof(start_node)}()
  push!(stack, start_node)
  search_state = DepthFirstSearchState(stack, Set{typeof(start_hash)}(), start_time)

  return _depthfirst!(
    search_state,
    neighbours,
    start,
    goal,
    isgoal,
    hashfn,
    timeout,
    maxdepth,
    enable_closedset,
  )
end

"""
Execute iteratively a depth first search starting from depth 0 to maxdepth.
"""
function iterative_deepening(
  neighbours,
  start,
  goal;
  isgoal = defaultisgoal,
  hashfn = hash,
  timeout = Inf,
  maxdepth = typemax(Int64),
  enable_closedset = true,
  kwargs...,
)
  start_time = time()
  end_time = start_time + timeout
  closedsetsize = 0
  opensetsize = 0

  for depth = 0:maxdepth
    depth_first_timeout = max(end_time - time(), 0.0)
    res = depthfirst(
      neighbours,
      start,
      goal;
      isgoal,
      hashfn,
      enable_closedset,
      maxdepth = depth,
      timeout = depth_first_timeout,
      kwargs...,
    )
    if res.status in (:success, :timeout)
      return res
    end
    closedsetsize = res.closedsetsize
    opensetsize = res.opensetsize
  end
  return UninformedSearchResult(:nopath, [start], closedsetsize, opensetsize)
end

mutable struct BreadthFirstSearchState{TState, THash}
  openset::Deque{UninformedSearchNode{TState}}
  closedset::Set{THash}
  start_time::Float64
end

function _breadthfirst!(
  search_state::BreadthFirstSearchState{TState, THash},
  neighbours,
  start,
  goal,
  isgoal,
  hashfn,
  timeout,
  maxdepth,
  enable_closedset,
) where {TState, THash}
  while !isempty(search_state.openset)
    node = popfirst!(search_state.openset)
    node_data = node.data
    if isgoal(node_data, goal)
      return UninformedSearchResult(
        :success,
        reconstructpath(node),
        length(search_state.closedset),
        length(search_state.openset),
      )
    end

    if timeout < Inf && time() - search_state.start_time >= timeout
      return UninformedSearchResult(
        :timeout,
        reconstructpath(node),
        length(search_state.closedset),
        length(search_state.openset),
      )
    end

    for neighbour in neighbours(node_data)
      new_depth = node.depth + 1
      if (enable_closedset && hashfn(neighbour) in search_state.closedset) ||
         new_depth > maxdepth
        continue
      end

      neighbour_node = UninformedSearchNode(neighbour, new_depth, node)
      push!(search_state.openset, neighbour_node)
    end

    if enable_closedset
      push!(search_state.closedset, hashfn(node_data))
    end
  end

  return UninformedSearchResult(
    :nopath,
    [start],
    length(search_state.closedset),
    length(search_state.openset),
  )
end

"""
Executes a breadth first search given the neightbours function, the start and goal state.
"""
function breadthfirst(
  neighbours,
  start,
  goal;
  isgoal = defaultisgoal,
  hashfn = hash,
  timeout = Inf,
  maxdepth = typemax(Int64),
  enable_closedset = true,
  kwargs...,
)
  start_time = time()
  start_hash = hashfn(start)

  start_node = UninformedSearchNode(start, 0, nothing)
  deque = Deque{typeof(start_node)}()
  push!(deque, start_node)
  search_state = BreadthFirstSearchState(deque, Set{typeof(start_hash)}(), start_time)

  return _breadthfirst!(
    search_state,
    neighbours,
    start,
    goal,
    isgoal,
    hashfn,
    timeout,
    maxdepth,
    enable_closedset,
  )
end


"""
    AbstractCache

Abstract supertype used for information storage object the directed acyclic graph.
"""
abstract type AbstractCache end

"""
    AbstractCacheAttribute

Abstract supertype used for attributes stored in a cache.
"""
abstract type AbstractCacheAttribute end

Base.@kwdef mutable struct VariableValues{T<:Real}
    x0::Vector{T}                          = T[]
    x::Vector{T}                           = T[]
    lower_variable_bounds::Vector{T}       = T[]
    upper_variable_bounds::Vector{T}       = T[]
    node_to_variable_map::Vector{Int}      = Int[]
    variable_to_node_map::Vector{Int}      = Int[]
    variable_types::Vector{VariableType}   = VariableType[]
end

@inline val(b::VariableValues{T}, i::Int) where T = @inbounds b.x[i]
@inline lbd(b::VariableValues{T}, i::Int) where T = @inbounds b.lower_variable_bounds[i]
@inline ubd(b::VariableValues{T}, i::Int) where T = @inbounds b.upper_variable_bounds[i]
val(b::VariableValues{T}) where T = b.x
lbd(b::VariableValues{T}) where T = b.lower_variable_bounds
ubd(b::VariableValues{T}) where T = b.upper_variable_bounds

function _get_x!(::Type{BranchVar}, out::Vector{T}, v::VariableValues{T}) where T<:Real
    @inbounds for i = 1:length(v.node_to_variable_map)
        out[i] = v.x[v.node_to_variable_map[i]]
    end
    return nothing
end

function _initialize_or_copy!(y::VariableValues{T}, x::VariableValues{T}, s::Symbol) where {T<:Real}
    isempty(getfield(y, s)) ? setfield!(y, s, copy(getfield(x, s))) : copy!(getfield(y, s), getfield(x, s))
end
function update_box_and_pnt!(y::VariableValues{T}, x::VariableValues{T}, update_box::Bool) where {T<:Real}
    if update_box
        _initialize_or_copy!(y, x, :x0)
        _initialize_or_copy!(y, x, :lower_variable_bounds)
        _initialize_or_copy!(y, x, :upper_variable_bounds)
    end
    _initialize_or_copy!(y, x, :x)
end

"""
    DirectedTree

A tree graph with a single sink node.
"""
Base.@kwdef mutable struct DirectedTree <: AbstractDirectedAcyclicGraph
    "List of nodes"
    nodes::Vector{Node}                       = Node[]
    "List of index of variables in this tree"
    variables::Dict{Int,Int}                  = Dict{Int,Int}()
    "Information on all variables..."
    v::VariableValues{Float64}                = VariableValues{Float64}()
    "List of constant values"
    constant_values::Vector{Float64}          = Float64[]
    "List of constant values"
    parameter_values::Vector{Float64}         = Float64[]
    "Number of nodes"
    node_count::Int                           = 0
    "Number of variables"
    variable_count::Int                       = 0
    "Number of constants"
    constant_count::Int                       = 0
    sink_bnd::Interval{Float64}               = Interval{Float64}(-Inf,Inf)
    ""
    sparsity::Vector{Int}                     = Int[]
    ""
    rev_sparsity::Dict{Int,Int}               = Dict{Int,Int}()
    dependent_variable_count::Int             = 0
    dep_subexpr_count::Int                    = 0
    dependent_subexpressions::Vector{Int}     = Int[]
    linearity::Linearity                      = LIN_CONSTANT
    user_operators::OperatorRegistry          = OperatorRegistry()
    #children::SpraseMatrixCSC{Bool,Int}      = spzeros(Bool,1,1)
    #parents::SparseMatrixCSC{Bool,Int}       = spzeros(Bool,1,1)
end
#DirectedTree(n::Int) = DirectedTree(children = spzeros(Bool,n,n), parents=spzeros(Bool,n,n))
const DAT = DirectedTree

# node property access functions that can be defined at abstract type
node(g::DAT, i)            = g.nodes[i]
nodes(g::DAT)              = g.nodes

variable(g::DAT, i)        = g.variables[i]
variables(g::DAT)          = g.variables

constant_value(g::DAT, i)  = g.constant_values[i]
constant_values(g::DAT)    = g.constant_values
parameter_value(g::DAT, i) = g.parameter_values[i]
parameter_values(g::DAT)   = g.parameter_values

node_class(g::DAT, i)      = node_class(node(g, i))
ex_type(g::DAT, i)         = ex_type(node(g, i))
first_index(g::DAT, i)     = first_index(node(g, i))
secondary_index(g::DAT, i) = secondary_index(node(g, i))
arity(g::DAT, i)           = arity(node(g, i))
children(g::DAT, i)        = children(node(g, i))
child(g::DAT, i, j)        = child(node(g, j), i)

is_binary(g::DAT, i) = arity(g, i) == 2

node_count(g::DAT)     = g.node_count
variable_count(g::DAT) = g.variable_count
constant_count(g::DAT) = g.constant_count

dep_subexpr_count(g::DAT)            = length(g.dependent_subexpressions)
sparsity(g::DAT, i)                  = g.sparsity
rev_sparsity(g::DAT, i::Int, k::Int) = g.rev_sparsity[i]

user_univariate_operator(g::DAT, i) = g.user_operators.univariate_operator_f[i]
user_multivariate_operator(g::DAT, i) = g.user_operators.multivariate_operator_evaluator[i]

function DirectedTree(aux_info, d, op::OperatorRegistry, sub_sparsity::Dict{Int,Vector{Int}}, subexpr_linearity, parameter_values, is_sub, subexpr_indx)

    nd = copy(d.nd)
    adj = copy(d.adj)
    const_values = copy(d.const_values)

    sparsity, dependent_subexpressions = _compute_sparsity(d, sub_sparsity, is_sub, subexpr_indx)
    rev_sparsity = Dict{Int,Int}()
    for (i,s) in enumerate(sparsity)
        rev_sparsity[s] = i
    end

    nodes = _convert_node_list(aux_info, d.nd, op)
    lin = linearity(nd, adj, subexpr_linearity)
    DirectedTree(nodes = nodes,
                    variables = rev_sparsity,
                    constant_values = const_values,
                    parameter_values = parameter_values,
                    node_count = length(nodes),
                    variable_count = length(sparsity),
                    constant_count = length(const_values),
                    sparsity = sparsity,
                    rev_sparsity = rev_sparsity,
                    dependent_variable_count = length(sparsity),
                    dep_subexpr_count = length(dependent_subexpressions),
                    dependent_subexpressions = copy(dependent_subexpressions),
                    linearity = lin[1],
                    user_operators = op
                    )
end

forward_uni = [i for i in instances(AtomType)]
setdiff!(forward_uni, [VAR_ATOM; PARAM_ATOM; CONST_ATOM; SELECT_ATOM; SUBEXPR])
f_switch = binary_switch(forward_uni, is_forward = true)
@eval function fprop!(t::T, ex::Expression, g::DAT, c::AbstractCache , k::Int) where T<:AbstractCacheAttribute
    id = ex_type(g, k)
    $f_switch
    error("fprop! for ex_type = $id not defined.")
    return
end

reverse_uni = [i for i in instances(AtomType)]
setdiff!(reverse_uni, [VAR_ATOM; PARAM_ATOM; CONST_ATOM; SELECT_ATOM; SUBEXPR])
r_switch = binary_switch(reverse_uni, is_forward = false)
@eval function rprop!(t::T, ex::Expression, g::DAT, c::AbstractCache, k::Int) where T<:AbstractCacheAttribute
    id = ex_type(g, k)
    $r_switch
    error("rprop! for ex_type = $id not defined.")
    return
end


"""
Table 1 from Efron-Tibshirani 1996
"""
const pain_scores = [
    3
    7
    6
    1
    2
    3
    3
    1
    7
    5
    4
    4
    1
    3
    3
    5
    0
    1
    0
    0
    2
    2
    0
    0
    0
    1
    0
    0
    1
    1
    1
    0
    0
    0
    0
    0
    0
    0
    0
    0
]


using Partial
include("test-partial.jl")


@testset "Controller: Objective" begin
    s = get_simulation("quadruped", "flat_2D_lc", "flat")
    model = s.model
    env = s.env

    H = 10
    obj = ContactImplicitMPC.TrackingObjective(model, env, H)
    @test length(obj.q) == H
    @test length(obj.u) == H
    @test length(obj.γ) == H
    @test length(obj.b) == H
end


"""
    RemainderModel(main::BalanceLaw, subcomponents::Tuple)

Compute the "remainder" contribution of the `main` model, after subtracting `subcomponents`.
"""
struct RemainderModel{M,S} <: BalanceLaw
  main::M
  subs::S
end

vars_state(rem::RemainderModel, FT) = vars_state(rem.main,FT)
vars_gradient(rem::RemainderModel, FT) = vars_gradient(rem.main,FT)
vars_diffusive(rem::RemainderModel, FT) = vars_diffusive(rem.main,FT)
vars_aux(rem::RemainderModel, FT) = vars_aux(rem.main,FT)
vars_integrals(rem::RemainderModel,FT) = vars_integrals(rem.main,FT)

update_aux!(dg::DGModel, rem::RemainderModel, Q::MPIStateArray, t::Real) =
  update_aux!(dg, rem.main, Q, t)

integrate_aux!(rem::RemainderModel, integ::Vars, state::Vars, aux::Vars) =
  integrate_aux!(rem.main, integ, state, aux)

flux_diffusive!(rem::RemainderModel, flux::Grad, state::Vars, diffusive::Vars, aux::Vars, t::Real) =
  flux_diffusive!(rem.main, flux, state, diffusive, aux, t)

gradvariables!(rem::RemainderModel, transform::Vars, state::Vars, aux::Vars, t::Real) =
  gradvariables!(rem.main, transform, state, aux, t)

diffusive!(rem::RemainderModel, diffusive::Vars, ∇transform::Grad, state::Vars, aux::Vars, t::Real) =
  diffusive!(rem.main, diffusive, ∇transform, state, aux, t)

function wavespeed(rem::RemainderModel, nM, state::Vars, aux::Vars, t::Real)
  ref = aux.ref_state
  return wavespeed(rem.main, nM, state, aux, t) - sum(sub -> wavespeed(sub, nM, state, aux, t), rem.subs)
end

boundary_state!(nf, rem::RemainderModel, x...) = boundary_state!(nf, rem.main, x...)

init_aux!(rem::RemainderModel, aux::Vars, geom::LocalGeometry) = nothing
init_state!(rem::RemainderModel, state::Vars, aux::Vars, coords, t) = nothing

function flux_nondiffusive!(rem::RemainderModel, flux::Grad, state::Vars, aux::Vars, t::Real)
  m = getfield(flux, :array)
  flux_nondiffusive!(rem.main, flux, state, aux, t)

  flux_s = similar(flux)
  m_s = getfield(flux_s, :array)

  for sub in rem.subs
    fill!(m_s, 0)
    flux_nondiffusive!(sub, flux_s, state, aux, t)
    m .-= m_s
  end
  nothing
end

function source!(rem::RemainderModel, source::Vars, state::Vars, aux::Vars, t::Real)
  m = getfield(source, :array)
  source!(rem.main, source, state, aux, t)

  source_s = similar(source)
  m_s = getfield(source_s, :array)

  for sub in rem.subs
    fill!(m_s, 0)
    source!(sub, source_s, state, aux, t)
    m .-= m_s
  end
  nothing
end


#=
Copyright 2017 INSIGNEO Institute for in silico Medicine

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
=#

# Convergence is checked starting from the simulated third cardiac cycle.
# Waveforms are first loaded and subdivided depending on the number of cycles
# already passed. This task is accomplished by function
# [`extractWaveform`](check_convergence.html#extractWaveform). Waveforms are
# then resampled in order to contain a smaller number of points. This would
# speed up the following error calculation. These two steps are taken by
# [`resampleAndCalculateError`](check_convergence.html#resampleAndCalculateError)
# function. Loading, re-sampling, and error calculation operations are done
# for all the quantities calculated by `openBF`. All these functions are
# wrapped into
# [`checkAllQuantities`](check_convergence.html#checkAllQuantities) which is
# called inside [main.jl](main.html#check_convergence) loop.

# *function* __`checkAllQuantities`__ $\rightarrow$ `max_error::Float`
#
# ----------------------------------------------------------------------------
# Parameters:
# --------------- ------------------------------------------------------------
# `v`             `::Vessel` vessel of which the error is being checked.
#
# `passed_cycles` `::Int` number of cardiac cycles already passed.
#
# `n_pts`         `::Int` number of points in the resampled waveform.
# ----------------------------------------------------------------------------
#
function checkAllQuantities(v             :: Vessel,
                            passed_cycles :: Int64,
                            n_pts         :: Int64)
  # --------------------------------------------------------------------------
  # Functioning
  # --------------------------------------------------------------------------
  # `er` is an array which will contain the percentage error for each quantity
  # in the current vessel. `qs` collection contains all the suffixes needed to
  # build `.out` filenames for each quantity.
  # --------------------------------------------------------------------------
  er = zeros(Float64, 5)

  qs = ["_P.out", "_Q.out", "_A.out", "_c.out", "_u.out"]
  # For each quantity `q` listed in `qs`, the `.out` filename is created by
  # joining the `label` inside `v` and the suffix `q`. Then the desired
  # waveforms are loaded by
  # [`extractWaveform`](check_convergence.html#extractWaveform), re-sampled,
  # and the error is calculated by
  # [`resampleAndCalculateError`](check_convergence.html#resampleAndCalculateError).
  # Eventually, the maximum error scored among all the quantities is returned.
  #
  # --------------------------------------------------------------------------
  # Returns:
  # ------------- ------------------------------------------------------------
  # `max_error`   `::Float` maximum error within all the quantities in the
  #               current vessel.
  # --------------------------------------------------------------------------
  idx = 1
  for q in qs
    filename = join([v.label, q])

    W = extractWaveform(filename, passed_cycles)

    er[idx] = resampleAndCalculateError(v, W, passed_cycles, n_pts)

    idx += 1
  end

  return maximum(er)
end

# *function* __`checkAllQuantities`__ $\rightarrow$ `max_error::Float`
#
# ----------------------------------------------------------------------------
# Parameters:
# --------------- ------------------------------------------------------------
# `vs`             `::Array{Vessel, 1}` collection of vessels.
#
# `passed_cycles` `::Int` number of cardiac cycles already passed.
#
# `n_pts`         `::Int` number of points in the resampled waveform.
# ----------------------------------------------------------------------------
#
# ----------------------------------------------------------------------------
# Functioning
# ----------------------------------------------------------------------------
# To compute the error for all the vessels, `checkAllQuantities` is
# re-defined to handle also a vector of `Vessel`s. `checkAllQuantities` is
# called recursively.
# ----------------------------------------------------------------------------
#
# ----------------------------------------------------------------------------
# Returns:
# ------------- --------------------------------------------------------------
# `max_error`   `::Float` maximum error among all the vessels.
# ----------------------------------------------------------------------------
#
function checkAllQuantities(vs            :: Array{Vessel, 1},
                            passed_cycles :: Int64,
                            n_pts         :: Int64)

  ers = zeros(Float64, length(vs))

  idx = 1
  for v in vs

    ers[idx] = checkAllQuantities(v, passed_cycles, n_pts)
    idx += 1

  end

  return maximum(ers)
end

# *function* __`extractWaveform`__ $\rightarrow$ `Fx::Array{Array{Float, 2}}`
#
# ----------------------------------------------------------------------------
# Parameters:
# ------------- --------------------------------------------------------------
# `filename`    `::String` current time in seconds.
#
# `cycles`      `::Int` $\Delta t$, current time step.
# ----------------------------------------------------------------------------
#
function extractWaveform(filename :: String, cycles :: Int64)
  # --------------------------------------------------------------------------
  # Functioning
  # --------------------------------------------------------------------------
  # `filename` is a string referring to the `.out` file containing the
  # waveform to be imported. This is read by `readdlm` function and stored
  # into
  # `F` array. `F` is a timesteps $\times$ 6 array; the first column contains
  # the time variable, the remaining columns are one for each monitored node
  # along the current vessel.
  # --------------------------------------------------------------------------
  F = readdlm(filename)
  # These three lines find the cardiac period `t` by knowing how many
  # cycles have already passed.
  sf1 = int(length(F[:,1])*(cycles-1)/cycles)
  T1 = F[sf1:end,1] .- F[sf1,1]
  t = T1[end]
  # `Fx` is an empty array which will contain one array for each cardiac
  # cycle.
  Fx = Array{Float64, 2}[]
  # This loop skims through `F`, separates waveforms, and add them to `Fx`.
  # `si` is the waveform source index, the index from which the waveform
  # starts. `ti` is the waveform terminal index, the index at which the
  # waveforms ends because a new cardiac cycle is due to start.
  si = 1
  ti = 1
  idx = 1
  for i in 1:length(F[:,1])-1
    if F[i,1] < t*idx && F[i+1, 1] >= t*idx
      ti = i
      push!(Fx, F[si:ti,:])
      si = ti
      idx += 1
    end
  end
  # --------------------------------------------------------------------------
  # Returns:
  # ------------- ------------------------------------------------------------
  # `Fx`           `::Array{Array{Float, 2}}`
  # --------------------------------------------------------------------------
  return Fx
end

# *function* __`resampleAndCalculateError`__ $\rightarrow$
# `max_error::Float`
#
# ----------------------------------------------------------------------------
# Parameters:
# ----------------- ----------------------------------------------------------
# `W`               `::Array{Array{Float, 2}}` collection of all waveforms
#                   during the passed cardiac cycles for the current
#                   quantity.
#
# `passed_cycles`   `::Int` number of cardiac cycles already simulated.
#
# `n_pts`           `::Int` number of points to be contained in the
#                   re-sampled waveform.
# ----------------------------------------------------------------------------
#
function resampleAndCalculateError(v, W :: Array{Array{Float64, 2}, 1},
                                   passed_cycles :: Int64,
                                   n_pts         :: Int64)
  # --------------------------------------------------------------------------
  # Functioning
  # --------------------------------------------------------------------------
  # `F` will contain the re-sampled waveform, `er` is an empty variable
  # for the maximum error.
  # --------------------------------------------------------------------------
  F = zeros(Float64, n_pts, 2)
  er = 0
  # The error is computed for all the monitor nodes. Their values are stored
  # in `W`'s columns from 2 to 6. Each waveform is re-sampled by using
  # [`resample`](http://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.resample.html#scipy.signal.resample)
  # function from `scipy.signal` library. In order to compute the error, only
  # the last and the previous waveforms are analysed.
  for j in 2:6

    F[:,1], t = signal.resample(W[end-1][:,j], n_pts, W[end-1][:,1])
    F[:,2], t = signal.resample(W[end  ][:,j], n_pts, W[end  ][:,1])

    er = abs(F[:,1] .- F[:,2])./F[:,1]*100
  end
  # --------------------------------------------------------------------------
  # Returns:
  # ----------- --------------------------------------------------------------
  # `max_error` `::Float` maximum error on the current quantity along the
  #             the current vessel.
  # --------------------------------------------------------------------------
  return maximum(er)
end

function checkConvergence(vessels :: Array{Vessel, 1},
  convergence_tollerance :: Float64)

  for v in vessels
    err = checkConvergence(v, convergence_tollerance)
    if err > convergence_tollerance
      return err
    end
  end
  return err
end

function checkConvergence(v :: Vessel,
  convergence_tollerance :: Float64)

  qs = ["_P", "_Q"]

  for q in qs
    filename_last = join([v.label, q, ".last"])
    filename_temp = join([v.label, q, ".temp"])

    err = checkQuantity(filename_last, filename_temp, convergence_tollerance)
    if err > convergence_tollerance
      return err
    end
  end
  return err
end

function checkQuantity(last, temp, convergence_tollerance)

  s1 = readdlm(last)
  s2 = readdlm(temp)
  n_pts = 1000

  r1 = n_pts/length(s1[:,2])
  r2 = n_pts/length(s2[:,2])
  for j in 2:5
    sr1 = resample(s1[:,j], r1)
    sr2 = resample(s2[:,j], r2)
    for i in 20:length(sr1)-20
      err = abs(100*(sr2[i]-sr1[i])/sr1[i])
      if err > convergence_tollerance
        # println("\n",err)
        return err
      end
    end
  end
  return 9.99
end

function checkConvergence(edge_list, vessels :: Array{Vessel, 1}, passed_cycles :: Int64)
    qs = "_A"
    err = zeros(size(edge_list)[1]).+100
    if passed_cycles >= 2

        for i in 1:size(edge_list)[1]
            v = vessels[i]
            lbl = v.label

            filename_last = join([lbl, qs, ".last"])
            filename_temp = join([lbl, qs, ".temp"])
            w_last = readdlm(filename_last)
            w_temp = readdlm(filename_temp)

            if length(w_last[:,1]) == length(w_temp[:,1])
                err[i] = maximum(abs.((w_last[2:end,:].-w_temp[2:end,:])./w_last[2:end,:])*100)
            end

        end
    end
    return maximum(err)
end


"""
Appends a tracking process to determine the time of the callback to be used in
the reverse pass. The rationale is explain in:

https://github.com/SciML/DiffEqSensitivity.jl/issues/4
"""
track_callbacks(cb,t,u) = track_callbacks(CallbackSet(cb),t,u,p)
track_callbacks(cb::CallbackSet,t,u,p) = CallbackSet(
   map(cb->_track_callback(cb,t,u,p), cb.continuous_callbacks),
   map(cb->_track_callback(cb,t,u,p), cb.discrete_callbacks))

struct TrackedAffect{T,T2,T3,T4}
    event_times::Vector{T}
    uleft::Vector{T2}
    pleft::Vector{T3}
    affect!::T4
end

TrackedAffect(t::Number,u,p,affect!::Nothing) = nothing
TrackedAffect(t::Number,u,p,affect!) = TrackedAffect(Vector{typeof(t)}(undef,0),Vector{typeof(p)}(undef,0),Vector{typeof(u)}(undef,0),affect!)

function (f::TrackedAffect)(integrator)
    uleft = deepcopy(integrator.u)
    pleft = deepcopy(integrator.p)
    f.affect!(integrator)
    if integrator.u_modified
        push!(f.event_times,integrator.t)
        push!(f.uleft,uleft)
        push!(f.pleft,pleft)
    end
end

function _track_callback(cb::DiscreteCallback,t,u,p)
    DiscreteCallback(cb.condition,
                     TrackedAffect(t,u,p,cb.affect!),
                     cb.initialize,
                     cb.finalize,
                     cb.save_positions)
end

function _track_callback(cb::ContinuousCallback,t,u,p)
    ContinuousCallback(
        cb.condition,
        TrackedAffect(t,u,p,cb.affect!),
        TrackedAffect(t,u,p,cb.affect_neg!),
        cb.initialize,
        cb.finalize,
        cb.idxs,
        cb.rootfind,cb.interp_points,
        cb.save_positions,
        cb.dtrelax,cb.abstol,cb.reltol,cb.repeat_nudge)
end

function _track_callback(cb::VectorContinuousCallback,t,u,p)
    VectorContinuousCallback(
               cb.condition,
               TrackedAffect(t,u,p,cb.affect!),
               TrackedAffect(t,u,p,cb.affect_neg!),
               cb.len,cb.initialize,cb.finalize,cb.idxs,
               cb.rootfind,cb.interp_points,
               collect(cb.save_positions),
               cb.dtrelax,cb.abstol,cb.reltol,cb.repeat_nudge)
end

struct FakeIntegrator{uType,P,tType}
    u::uType
    p::P
    t::tType
end

struct CallbackSensitivityFunction{fType,Alg<:DiffEqBase.AbstractSensitivityAlgorithm,C<:AdjointDiffCache,pType} <: SensitivityFunction
    f::fType
    sensealg::Alg
    diffcache::C
    prob::pType
end

getprob(S::CallbackSensitivityFunction) = S.prob
inplace_sensitivity(S::CallbackSensitivityFunction) = true

"""
Sets up callbacks for the adjoint pass. This is a version that has an effect
at each event of the forward pass and defines the reverse pass values via the
vjps as described in https://arxiv.org/pdf/1905.10403.pdf Equation 13.

For more information, see https://github.com/SciML/DiffEqSensitivity.jl/issues/4
"""
setup_reverse_callbacks(cb,sensealg) = setup_reverse_callbacks(CallbackSet(cb),sensealg)
function setup_reverse_callbacks(cb::CallbackSet,sensealg)
    cb = CallbackSet(_setup_reverse_callbacks.(cb.continuous_callbacks,(sensealg,))...,
                     reverse(_setup_reverse_callbacks.(cb.discrete_callbacks,(sensealg,)))...)
    cb
end

function _setup_reverse_callbacks(cb::Union{ContinuousCallback,DiscreteCallback,VectorContinuousCallback},sensealg)

    function affect!(integrator)

        function w(du,u,p,t)
          fakeinteg = FakeIntegrator([x for x in u],[x for x in p],t)
          cb.affect!.affect!(fakeinteg)
          du .= fakeinteg.u
        end

        S = integrator.f.f # get the sensitivity function

        # Create a fake sensitivity function to do the vjps
        fakeS = CallbackSensitivityFunction(w,sensealg,S.diffcache,integrator.sol.prob)

        du = first(get_tmp_cache(integrator))
        λ,grad,y,dλ,dgrad,dy = split_states(du,integrator.u,integrator.t,S)

        update_p = copy_to_integrator!(cb,y,integrator.p,integrator.t)

        if update_p
          # changes in parameters
          if !(sensealg isa QuadratureAdjoint)
            function wp(dp,p,u,t)
              fakeinteg = FakeIntegrator([x for x in u],[x for x in p],t)
              cb.affect!.affect!(fakeinteg)
              dp .= fakeinteg.p
            end

            fakeSp = CallbackSensitivityFunction(wp,sensealg,S.diffcache,integrator.sol.prob)
            #vjp with Jacobin given by dw/dp before event and vector given by grad
            vecjacobian!(dgrad, integrator.p, grad, y, integrator.t, fakeSp;
                                  dgrad=nothing, dy=nothing)
            grad .= dgrad
          end
        end

        vecjacobian!(dλ, y, λ, integrator.p, integrator.t, fakeS;
                              dgrad=dgrad, dy=dy)

        λ .= dλ
        if !(sensealg isa QuadratureAdjoint)
          grad .-= dgrad
        end
    end

    times = if typeof(cb) <: DiscreteCallback
        cb.affect!.event_times
    else
        [cb.affect!.event_times;cb.affect_neg!.event_times]
    end

    PresetTimeCallback(times,
                       affect!,
                       save_positions = (false,false))
end

function copy_to_integrator!(cb::DiscreteCallback, y, p, t)
  indx = searchsortedfirst(cb.affect!.event_times,t)
  copyto!(y, cb.affect!.uleft[indx])
  update_p = (p != cb.affect!.pleft[indx])
  update_p && copyto!(p, cb.affect!.pleft[indx])
  update_p
end

function copy_to_integrator!(cb::Union{ContinuousCallback,VectorContinuousCallback}, y, p, t)
  if !isempty(cb.affect!.event_times)
    indx = searchsortedfirst(cb.affect!.event_times,t)
    if cb.affect!.event_times[indx]!=t
      indx = searchsortedfirst(cb.affect_neg!.event_times,t)
      copyto!(y, cb.affect_neg!.uleft[indx])
      update_p = (p != cb.affect!.pleft[indx])
      update_p && copyto!(p, cb.affect_neg!.pleft[indx])
    else
      copyto!(y, cb.affect!.uleft[indx])
      update_p = (p != cb.affect!.pleft[indx])
      update_p && copyto!(p, cb.affect!.pleft[indx])
    end
  else
    indx = searchsortedfirst(cb.affect_neg!.event_times,t)
    copyto!(y, cb.affect_neg!.uleft[indx])
    update_p = (p != cb.affect_neg!.pleft[indx])
    update_p && copyto!(p, cb.affect_neg!.pleft[indx])
  end
  update_p
end


const MetaUndirectedGraph = MetaGraph{<:Any,<:Any,<:SimpleGraph}

SimpleGraph(g::MetaUndirectedGraph) = g.graph

is_directed(::Type{<:MetaUndirectedGraph}) = false

function arrange(g::MetaUndirectedGraph, label_1, label_2)
    vprops = g.vprops
    arrange(g::MetaUndirectedGraph, label_1, label_2, vprops[label_1], vprops[label_2])
end

function arrange(g::MetaUndirectedGraph, label_1, label_2, u, v)
    if u > v
        (label_2, label_1)
    else
        (label_1, label_2)
    end
end


module ParserCombinator
  export Success,
         Failure,
         Parser,
         @rule,
         @define_rule,
         or,
         cat,
         repeat,
         chainl,
         map,
         p_string,
         regex

  type Success
    value
    next
    success
    Success(value, next) = new(value, next, true)
  end

  type Failure
    next
    success
    Failure(next) = new(next, false)
  end

  type Parser
    parse
    Parser(parse) = new(parse)
  end

  macro rule(body, name)
    esc(quote
      function $(name)()
        Parser((input) -> ($body)().parse(input))
      end
    end)
  end

  function or(lhs :: Parser, rhs :: Parser)
    Parser() do input
      r = lhs.parse(input)
      if !r.success
        rhs.parse(input)
      else
        r
      end
    end
  end

  function cat(lhs :: Parser, rhs :: Parser)
    Parser() do input
      r1 = lhs.parse(input)
      if r1.success
        r2 = rhs.parse(r1.next)
        if r2.success
          Success((r1.value, r2.value), r2.next)
        else
          Failure(r2.next)
        end
      else
        Failure(r1.next)
      end
    end
  end

  function map(parser :: Parser, f :: Function)
    Parser() do input
      r = parser.parse(input)
      if r.success
        Success(f(r.value), r.next)
      else
        r
      end
    end
  end

  function repeat(parser :: Parser)
    Parser() do input
      rest = input
      values = []
      while true
        r = parser.parse(rest)
        if !r.success
          return Success(values, rest)
        end
        values = vcat(values, r.value)
        rest = r.next
      end
    end
  end

  function chainl(p :: Parser, q :: Parser)
    value = map(cat(p, repeat(cat(q, p))),
                  function(value)
                    x = value[1]
                    xs = value[2]
                    a = x
                    while length(xs) > 0
                      f = xs[1][1]
                      b = xs[1][2]
                      a = f(a, b)
                      xs = xs[2:end]
                    end
                    return a
                  end
                 )
    return value
  end

  function p_string(literal)
    Parser() do input
      range = search(input, literal)
      if range.start != 1
        Failure(input)
      else
        Success(literal, input[length(literal)+1:end])
      end
    end
  end

  function regex(regex :: Regex)
    Parser() do input
      r = match(regex, input)
      if r != nothing
        Success(r.match, input[length(r.match)+1:end])
      else
        Failure(input)
      end
    end
  end

  Base.:+(p1::Parser, p2::Parser) = cat(p1, p2)
  Base.:/(p1::Parser, p2::Parser) = or(p1, p2)
  Base.:>(p::Parser, f::Function) = map(p, f)
end


# helper function for testing to load the example results
function load_example(path, ncol, delim = ',', skiplines = 1) #make sure to use nscenario+1 cols

    nlines = countlines(path) - skiplines
    output = Array{Float64, 2}(undef, nlines, ncol)
    filestream = open(path, "r")

    try
        for skip in 1:skiplines
            readline(filestream)
        end

        for (index, line) in enumerate(eachline(filestream))
            parsedline = parse.(Float64, split(line, delim, keepempty = false))
            output[index, :] = parsedline
        end
    finally
        close(filestream)
    end

    return output
end

# helper function for testing to interpolate the results to match wellbore segmentation
function match_example(well::Wellbore, example::Array{Float64,2})
    nrow = length(well.md)
    ncol = size(example)[2]
    output = Array{Float64, 2}(undef, nrow, ncol)
    output[:, 1] = well.md

    output[1, 2:end] = example[1, 2:end]
    for depth_index in 2:length(well.md)
        depth = well.md[depth_index]
        row_above = example[example[:,1] .<= depth, :][end,:]

        if row_above[1] == depth
            output[depth_index,2:end] = row_above[2:end]
        else
            row_below = example[example[:,1] .> depth, :][1,:]
            #interpolated_row = y1 + (y2 - y1)/(x2 - x1) * (depth - x1) :
            interpolated_values = [row_above[i] + (row_below[i] - row_above[i])/(row_below[1] - row_above[1]) * (depth - row_above[1]) for i in 2:ncol]
            output[depth_index,2:end] = interpolated_values
        end
    end

    return output
end

#%% general parameters
dp_est = 10 #psi
error_tolerance = 0.1 #psi
outlet_pressure = 220 - 14.65 #WHP in psig
oil_API = 35
sg_gas = 0.65
CO2 = 0.005
sg_water = 1.07
roughness = 0.0006500
BHT = 165
id = 2.441

#%% scenario parameters
scenarios = (:A, :B, :C, :D, :E)

parameters =    (rate = (A = 500, B = 250, C = 1000, D = 3000, E = 50),
                 WC = (A = 0.5, B = 0.25, C = 0.75, D = 0.85, E = 0.25),
                 GLR = (A = 4500, B = 6000, C = 3000, D = 1200, E = 10000),
                 WHT = (A = 100, B = 90, C = 105, D = 115, E = 80))

#%% load test Wellbore
testpath = joinpath(dirname(dirname(pathof(PressureDrop))), "test/testdata/Sawgrass_9_32")
surveypath = joinpath(testpath, "Test_survey_Sawgrass_9.csv")
testwell = read_survey(path = surveypath, id = id)

#%% generate linear temperature profiles
function create_temps(scenario)
        WHT = parameters[:WHT][scenario]

        return linear_wellboretemp(WHT = WHT, BHT = BHT, wellbore = testwell)
end

temps = [create_temps(s) for s in scenarios]
temp_profiles = NamedTuple{scenarios}(temps) #temp profiles labelled by scenario

@testset "Beggs and Brill with Palmer correction - scenarios" begin

# load test results to compare & generate interpolations from expected results at same depths as test well segmentation
examplepath = joinpath(testpath,"Perform_BandB_with_Palmer_correction.csv")
test_example = load_example(examplepath, length(scenarios)+1)
matched_example = match_example(testwell, test_example)
matched_example[:,2:end] = matched_example[:,2:end] .- 14.65 #convert to psig

# generate test data -- map across all examples
corr = BeggsAndBrill
test_results = Array{Float64, 2}(undef, length(testwell.md), 1 + length(scenarios))
test_results[:,1] = testwell.md
for (index, scenario) in enumerate(scenarios)
    WC = parameters[:WC][scenario]
    rate = parameters[:rate][scenario]
    q_o = rate * (1 - WC)
    q_w = rate * WC
    GLR = parameters[:GLR][scenario]

    test_results[:,index+1] = traverse_topdown(wellbore = testwell, roughness = roughness, temperatureprofile = temp_profiles[scenario],
                             pressurecorrelation = corr, dp_est = dp_est, error_tolerance = error_tolerance,
                             q_o = q_o, q_w = q_w, GLR = GLR, APIoil = oil_API, sg_water = sg_water, sg_gas = sg_gas,
                             WHP = outlet_pressure, molFracCO2 = CO2)
end

# compare test data
# NOTE: points after the first survey point at > 90° inclination are discarded;
# the reference data appears to force positive frictional effects.
hz_index = findnext(x -> x >= 90, testwell.inc, 1)
compare_tolerance = 40 #every scenario except the high-rate scenario can take a 15-psi tolerance
for index in 2:length(scenarios)+1
    println("Comparing scenario ", index-1, " on index ", index)
    @test all( abs.(test_results[1:hz_index,index] .- matched_example[1:hz_index,index])  .<= compare_tolerance )
    println("Max difference : ", maximum(abs.(test_results[1:hz_index,index] .- matched_example[1:hz_index,index])))
end

end #testset for B&B scenarios

@testset "Hagedorn & Brown with G&W correction - scenarios" begin

# load test results to compare & generate interpolations from expected results at same depths as test well segmentation
examplepath = joinpath(testpath,"Perform_HandB_with_GriffithWallis.csv")
test_example = load_example(examplepath, length(scenarios)+1)
matched_example = match_example(testwell, test_example)
matched_example[:,2:end] = matched_example[:,2:end] .- 14.65 #convert to psig

# generate test data -- map across all examples
corr = HagedornAndBrown
test_results = Array{Float64, 2}(undef, length(testwell.md), 1 + length(scenarios))
test_results[:,1] = testwell.md
for (index, scenario) in enumerate(scenarios)
    WC = parameters[:WC][scenario]
    rate = parameters[:rate][scenario]
    q_o = rate * (1 - WC)
    q_w = rate * WC
    GLR = parameters[:GLR][scenario]

    test_results[:,index+1] = traverse_topdown(wellbore = testwell, roughness = roughness, temperatureprofile = temp_profiles[scenario],
                             pressurecorrelation = corr, dp_est = dp_est, error_tolerance = error_tolerance,
                             q_o = q_o, q_w = q_w, GLR = GLR, APIoil = oil_API, sg_water = sg_water, sg_gas = sg_gas,
                             WHP = outlet_pressure, molFracCO2 = CO2)
end

# compare test data
# NOTE: points after the first survey point at > 90° inclination are discarded
hz_index = findnext(x -> x >= 90, testwell.inc, 1)
compare_tolerance = 85 #lposer tolerance to H&B due to wide range of methods to calculate correlating groups & reynolds numbers
for index in 2:length(scenarios)+1
    println("Comparing scenario ", index-1, " on index ", index)
    @test all( abs.(test_results[1:hz_index,index] .- matched_example[1:hz_index,index])  .<= compare_tolerance )
    println("Max difference : ", maximum(abs.(test_results[1:hz_index,index] .- matched_example[1:hz_index,index])))
end

end #H&B testset


module ITensorParallel

using Distributed
using MPI
using ITensors
using ITensors.NDTensors

import ITensors: product, position!, noiseterm, lproj, rproj

include("partition.jl")
include("threaded_projmposum.jl")
include("distributed_projmposum.jl")
include("mpi_projmposum.jl")

export ThreadedProjMPOSum, partition, DistributedSum, MPISum

end


#!/usr/bin/env julia

pnglist = filter!(r"\.png", readdir())

function expand(locci)
    list = []
    pos = split(locci, ':')
        chr = pos[1]
        region = split(pos[2], '-')
        n = parse(Int64, region[1])
        while n <= parse(Int64, region[end])
             out = "$chr" * ":" * "$n"
            push!(list, out)
            n +=  1
        end
    return list
end

function getSNPs(pnglist=pnglist)
	snps = []
	for img in pnglist
		img = replace(img, ".png", "")
		parts = split(img, '\_')
		chr = replace(parts[1], "chr", "")
		start = parse(Int, replace(parts[2], ",", ""))
		endp  = parse(Int, replace(parts[3], ",", ""))
		pos = "$chr" * ":" * "$start" * "-" * "$endp"
		vars = expand(pos)
		append!(snps, vars)
	end
	snps = unique(snps)
	for snp in snps
		println(snp)
	end
end

getSNPs()


using Test, SafeTestsets

@safetestset "Political Democracy" begin include("political_democracy/political_democracy.jl") end
@safetestset "Recover Parameters" begin include("recover_parameters/recover_parameters_twofact.jl") end
@safetestset "Multigroup" begin include("multigroup/multigroup.jl") end


BlackBoxOptim.MethodNames
BlackBoxNames=Array{String}(undef,0)
push!(BlackBoxNames,"Adapt. Diff. Evol. (rand 1 bin)")
push!(BlackBoxNames,"Adapt. Diff. Evol. (rand 1 bin, radious limited)")
push!(BlackBoxNames,"Borg Multiobj. Evol. Alg.")
push!(BlackBoxNames,"Diff. Evol. (rand 1 bin)")
push!(BlackBoxNames,"Diff. Evol. (rand 1 bin, radious limited)")
push!(BlackBoxNames,"Diff. Evol. (rand 2 bin)")
push!(BlackBoxNames,"Diff. Evol. (rand 2 bin, radious limited)")
push!(BlackBoxNames,"Distance-weighted Exponential Nat. Evol. Strat.")
push!(BlackBoxNames,"Compass Coordinate Direct Search")
push!(BlackBoxNames,"Probabilistic Descent Direct search")
push!(BlackBoxNames,"Random Search")
push!(BlackBoxNames,"Resampling Inheritance Memetic Search")
push!(BlackBoxNames,"Resampling Memetic Search")
push!(BlackBoxNames,"Separable Natural Evolution Strategies")
push!(BlackBoxNames,"Simultaneous Perturbation Stochastic Approximation")
push!(BlackBoxNames,"Exponential Natural Evolution Strategies")

ibb2=[ibb 17 18 19 20]
history_plot=history_all[1:iter,ibb2[1,:],:] # func_evals,n_optimizers,length(funcs)
n_optimizers=length(ibb2)
history_plotINIT=copy(history_plot)
for i=1:n_optimizers
    for j=1:size(history_plot,3)
        history_plotINIT[1:3,i,j].=maximum(history_plotINIT[1:3,:,j])
        mi=minimum(history_plotINIT[:,:,j])
        ma=maximum(history_plotINIT[:,:,j])
        # println(mi," ",ma)
        history_plot[:,i,j]=(history_plotINIT[:,i,j].-mi)./(ma.-mi)
    end
end
history_plot=mean(history_plot, dims=3)[:,:,1]
mi=minimum(history_plot[end,:])
ma=maximum(history_plot[end,:])
sR=history_plot[end,length(ibb)+1]
ssR=@sprintf("mi= %.2E ma= %.2E ITS= %.2E",mi,ma,sR)
println(ssR)
# display(plot(history_plot,width=1,label=[string.(ibb) "oht"],legend=:topleft,title=ssR))
display(plot(history_plot[:,1:length(ibb)],width=1,label=BlackBoxNames[ibb],legend=:topright,title=ssR,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+2],width=1,label="Nelder-Mead",legend=:topright,title=ssR,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+3],width=1,label="Particle Swarm",legend=:topright,title=ssR,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+4],width=1,label="Simulated Annealing",legend=:topright,title=ssR,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+1],width=2,label="Proposed-BNO",legendfontsize=8,color=:black,size = (sizex, sizey)))
display(xlims!(0,sf1*evals1))
display(xticks!(0:1000:evals1))
xaxis!("Function Evaluations")
yaxis!("Normalized Responces")
for i=[1:length(ibb);length(ibb)+2;length(ibb)+3;length(ibb)+4]
    println(minimum(history_plot[:,i]))
end
println(minimum(history_plot[:,length(ibb)+1]))

history_plot=history_all[1:iter,ibb2[1,:],:] # func_evals,n_optimizers,length(funcs)
n_optimizers=length(ibb2)
history_plotINIT=copy(history_plot)
for i=1:n_optimizers
    for j=1:size(history_plot,3)
        history_plotINIT[1:3,i,j].=maximum(history_plotINIT[1:3,:,j])
        mi=minimum(history_plotINIT[:,:,j])
        ma=maximum(history_plotINIT[:,:,j])
        history_plot[:,i,j]=(history_plotINIT[:,i,j].-mi)./(ma.-mi)
    end
end
# history_plot=(mean(history_plot, dims=3)[:,:,1].^(1/10))
history_plot=mean(history_plot, dims=3)[:,:,1].^(1/10)
# mi=minimum(history_plot[end,:])
# ma=maximum(history_plot[end,:])
# sR=history_plot[end,length(ibb)+1]
# ssR=@sprintf("mi= %.2E ma= %.2E ITS= %.2E",mi,ma,sR)
# println(ssR)
# display(plot(history_plot,width=1,label=[string.(ibb) "oht"],legend=:topleft,title=ssR))
display(plot(history_plot[:,1:length(ibb)],width=1,label=BlackBoxNames[ibb],legend=:false,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+2],width=1,label="Nelder-Mead",legend=:false,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+3],width=1,label="Particle Swarm",legend=:false,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+4],width=1,label="Simulated Annealing",legend=:false,titlefont = font("Times", 9),size = (sizex, sizey)))
display(plot!(history_plot[:,length(ibb)+1],width=2,label="Proposed-BNO",legend=:false,color=:black,size = (sizex, sizey)))
# display(xlims!(0,sf1*evals1))
display(xticks!(0:1000:evals1))
xaxis!("Function Evaluations")
yaxis!("Normalized Responces")
# xlims()


using HTTP, Joseki, JSON, Sockets, Test

@testset "Integration tests" begin

    # Simple server
    function pow(req::HTTP.Request)
        j = HTTP.queryparams(HTTP.URI(req.target))
        if !(haskey(j, "x")&haskey(j, "y"))
            return error_responder(req, "You need to specify values for x and y!")
        end
        x = parse(Float32, j["x"])
        y = parse(Float32, j["y"])
        json_responder(req, x^y)
    end

    function test_error_responder(req::HTTP.Request, e::Exception)
        error_responder(req, "Unhandled error!")
    end

    endpoints = [
        (pow, "GET", "/pow")
    ]
    # Default middleware but a custom error responder
    router = Joseki.router(endpoints; error_fn=test_error_responder)

    @info "Starting test server"
    server = Sockets.listen(Sockets.InetAddr(ip"127.0.0.1", 8000))
    server_task = @async HTTP.serve(router, ip"127.0.0.1", 8000; verbose=false,
                                    server = server)
    sleep(1.0)

    # A valid request
    res = HTTP.get("http://localhost:8000/pow/?x=2&y=3")
    @test res.status == 200
    @test HTTP.hasheader(res, "Access-Control-Allow-Methods", "GET,PUT,POST,DELETE,PATCH,OPTIONS")

    # An invalid request
    res = HTTP.get("http://localhost:8000/pow/?x=2&y=nothing")
    @test res.status == 200 # Our handeled errors return 200 responses
    @test occursin("Unhandled error!", String(copy(res.body)))

    @info "Shutting down test server"

    close(server)
    sleep(2)
    @test istaskdone(server_task)
end


# Base.summary(io::IO, x::CachedNode) = (print(io, "(Tracked) "); summary(io, x.output))
# Base.show(io::IO, m::MIME"text/plain", x::CachedNode{<:Node, <:AbstractArray}) = (print(io, "(tracked) "); show(io, m, x.output))
# Base.show(io::IO, x::CachedNode) = show(io, x.output)

# printing
"""
    PrintTrait(node) -> Trait
"""
function PrintTrait end

struct ScalarVariable{T}
    value::T
end

# a trait type to forward variable to array's type tree
struct ArrayVariable{T, N, AT <: AbstractArray{T, N}} <: AbstractArray{T, N}
    value::AT
end

struct Other{NT}
    node::NT
end

# forward to default
PrintTrait(var::Number) = ScalarVariable(var)
PrintTrait(var::Value{<:Number}) = ScalarVariable(value(var))
PrintTrait(var::Value{AT}) where {T, N, AT <: AbstractArray{T, N}} = ArrayVariable{T, N, AT}(value(var))
PrintTrait(var::Value{T}) where {T} = Other(var)

Base.size(x::ArrayVariable) = size(x.value)
Base.getindex(x::ArrayVariable, i...) = getindex(x.value, i...)

Base.show(io::IO, x::Value) = show(io, "text/plain", PrintTrait(x))
Base.show(io::IO, m::MIME"text/plain", x::Value) = show(io, m, PrintTrait(x))
Base.summary(io::IO, x::Union{Variable, CachedNode, Node}) = summary(io, PrintTrait(x))
Base.show(io::IO, x::ScalarVariable) = print(io, "(Tracked) ",x.value)
Base.summary(io::IO, x::ArrayVariable) = (print(io, "(Tracked) "); summary(io, x.value))

function Base.show(io::IO, ::MIME"text/plain", x::Other{<:CachedNode})
    println(io, "CachedNode:")
    println(io, "       f: ", operator(x.node))
    println(io, "    args: ", args(x.node))
    println(io, "  kwargs: ", kwargs(x.node))
      print(io, "  output: ", x.node.output)
end

function Base.show(io::IO, ::MIME"text/plain", x::Other{<:Variable})
    println(io, "Variable:")
    println(io, "  value: ", x.node.value)

    if isdefined(x.node, :grad)
        println(io, "   grad: ", x.node.grad)
    end
end


################################################# FILE DESCRIPTION #########################################################

# This file contains tests for the Mutation API

############################################################################################################################

@testset "Mutation Addvertex" begin
   g = completegraph(10)

   # Set vertex properties
   setvprop!(g, :, 1:10, :p1)
   setvprop!(g, :, 1:10, :p2)

   # Set edge properties
   seteprop!(g, :, 1:90, :p1)
   seteprop!(g, :, 1:90, :p2)

   g1 = deepcopy(g)
   # Addvertex without labels
   # Check for NA row
   # Check for type of LabelMap
   @test addvertex!(g1) == 11
   @test isequal(getvprop(g1, 11, :p1), NA)
   @test isequal(getvprop(g1, 11, :p2), NA)
   @test isa(g1.lmap, IdentityLM)

   g2 = deepcopy(g)
   # Addvertex with label
   # Check for NA row
   # Check for type of LabelMap
   @test addvertex!(g2, 15) == 11
   @test isequal(getvprop(g2, 11, :p1), NA)
   @test isequal(getvprop(g2, 11, :p2), NA)
   @test isa(g2.lmap, DictLM)

   g3 = deepcopy(g)
   # Addvertex with existing label
   # Check for no new row
   # Check for type of LabelMap
   @test addvertex!(g3, 10) == 10
   @test size(g3.vdata) == (10,2)
   @test isa(g3.lmap, IdentityLM)

   g4 = deepcopy(g)
   g4 + 11
   @test g4 == g1

   g5 = deepcopy(g)
   g5 + 15
   @test g5 == g2
end

@testset "Mutation Addedge" begin
   g = completegraph(10)

   # Set vertex properties
   setvprop!(g, :, 1:10, :p1)
   setvprop!(g, :, 1:10, :p2)

   # Set edge properties
   seteprop!(g, :, 1:90, :p1)
   seteprop!(g, :, 1:90, :p2)

   g1 = deepcopy(g)
   # Addedge 1=>1
   # Check for NA row at the END.
   # Check for mapping
   @test addedge!(g1, 1=>1) == true
   @test g1.indxs[1=>1] == 91

   g2 = deepcopy(g)
   g2[1] = 1
   @test g2 == g1

   g3 = deepcopy(g)
   g3[1] = [1]
   g3[5] = [5]
   g3[10] = [10]
   @test g3.indxs[[1=>1, 5=>5, 10=>10]] == [91, 92, 93]
end

@testset "Mutation Rmvertex" begin
   g = completegraph(10)
   setlabel!(g, collect(1:10))

   # Set vertex properties
   setvprop!(g, :, 1:10, :p1)
   setvprop!(g, :, 1:10, :p2)

   # Set edge properties
   seteprop!(g, :, 1:90, :p1)
   seteprop!(g, :, 1:90, :p2)

   g1 = deepcopy(g)
   # Remove vertex 10
   # Check for size of vdata
   # Check for edata entries
   rmvertex!(g1, 10)
   @test size(g1.vdata) == (9, 2)
   @test size(g1.edata) == (72, 2)

   g2 = deepcopy(g)
   g2 - 10
   @test g2 == g1

   g3 = deepcopy(g)
   # Remove vertices 5, 8
   # Check for size of vdata
   # Check for edata entries
   rmvertex!(g3, [5,8])
   @test nv(g3) == 8
   @test size(vdata(g3)) == (8,2)
   @test size(edata(g3)) == (56,2)

   g4 = deepcopy(g)
   g4 - [5,8]
   @test g4 == g3
end

@testset "Mutation Rmedge" begin
   g = Graph(completeindxs(10))
   # Set edge properties
   seteprop!(g, :, 1:90, :p1)
   seteprop!(g, :, 1:90, :p2)

   g1 = deepcopy(g)
   # Remove edge 1=>2
   # Check for size of edata
   # Check indexing table
   rmedge!(g1, 1=>2)
   @test size(g1.edata) == (89, 2)
   @test g1.indxs.nzval == collect(1:89)
   @test hasedge(g1, 1=>2) == false

   g2 = deepcopy(g)
   # Remove edges 1=>2, 5=>4, 10=>9
   # Check for size of edata
   # Check indexing table
   rmedge!(g2, 1=>2)
   rmedge!(g2, 5=>4)
   rmedge!(g2, 10=>9)
   @test size(g2.edata) == (87,2)
   @test g2.indxs.nzval == collect(1:87)
   @test hasedge(g1, 1=>2) == false
   @test hasedge(g2, 5=>4) == false
   @test hasedge(g2, 10=>9) == false
end


module SeisPlot

using PyCall, Statistics, Requires, FFTW
using PyPlot
import SeisMain

#@require PyPlot
const pl    = PyNULL()
const lines = PyNULL()
const anim  = PyNULL()

export SeisPlotTX,
SeisPlotAmplitude,
SeisPlotCoordinates,
SeisPlotFK,
SeisAnimation

include("SeisAnimation.jl")
include("SeisPlotTX.jl")
include("SeisPlotCoordinates.jl")
include("SeisPlotAmplitude.jl")
include("SeisPlotFK.jl")

export pl, lines

function __init__()
    copy!(pl,    pyimport_conda("matplotlib.pyplot",    "matplotlib"))
    copy!(lines, pyimport_conda("matplotlib.lines",     "matplotlib"))
    copy!(anim,  pyimport_conda("matplotlib.animation", "matplotlib"))
end
end


module Blind

using Random

using ArgParse
using ProgressMeter

using FileIO, CSVFiles

export blind

DEFAULT_NUMERIC = 4
DEFAULT_SETS = 3

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table! s begin
        "--sets", "-s"
            help = "Number of numeric sets."
            arg_type = Int
            default = DEFAULT_SETS
        "--numeric", "-n"
            help = "Number of numeric characters."
            arg_type = Int
            default = DEFAULT_NUMERIC
        "input"
            help = "Input directory."
            required = true
        "output"
            help = "Output directory."
            required = false
        "key"
            help = "Key file (.csv)."
            required = false
    end

    return parse_args(s)
end

function generate_names(n::Int, num_sets::Int, num_numeric::Int)
    strs = Vector{String}()

    @info "Generating $n strings with $num_sets sets of $num_numeric numerics."
    while length(strs) < n

        sets = [String(rand('0':'9', num_numeric)) for i in 1:num_sets]

        str = join(sets, "-")

        if !in(str, strs)
            push!(strs, str)
        end
    end

    return strs
end

function _nt(path_original::AbstractString, path_blind::AbstractString)
    return NamedTuple{(:original, :blind),Tuple{String,String}}((path_original, path_blind))
end

function blind(dir_input::AbstractString, dir_output::AbstractString, path_key::AbstractString; num_sets::Int = DEFAULT_SETS, num_numeric::Int = DEFAULT_NUMERIC)

    isdir(dir_input) || error("Input $dir_input is not a directory.")

    mkpath(dir_output)

    @info "Reading input $dir_input."
    files_original = readdir(dir_input)

    strs = generate_names(length(files_original), num_sets, num_numeric)
    exts = files_original .|> (str) ->(splitext(str)[2])

    files_blind = string.(strs, exts)

    paths_original = joinpath.(dir_input, files_original)
    paths_blind = joinpath.(dir_output, files_blind)

    @info "Copying to output $dir_output"
    @showprogress 1 for (path_original, path_blind) in zip(paths_original, paths_blind)
        cp(path_original, path_blind)
    end

    @info "Saving key to $path_key"
    nts = _nt.(paths_original, paths_blind)

    nts |> save(path_key)
end

function blind(dict::Dict)

    if dict["key"] === nothing
        dict["key"] = joinpath(dict["input"], "key.csv")
    end

    if dict["output"] === nothing
        dict["output"] = string(dict["input"], "-blind")
    end

    return blind(dict["input"], dict["output"], dict["key"], num_sets = dict["sets"], num_numeric = dict["numeric"])
end

end # module


# wrappers

function cudnnSoftmaxForward(x::DenseCuArray{T,4}, y::DenseCuArray{T,4}=x;
                             algo=CUDNN_SOFTMAX_FAST, # or CUDNN_SOFTMAX_ACCURATE
                             mode=CUDNN_SOFTMAX_MODE_INSTANCE, # or CUDNN_SOFTMAX_MODE_CHANNEL
                             alpha=1.0, beta=0.0) where T
    cudnnSoftmaxForward(handle(), algo, mode,
                        scalingParameter(T, alpha), TensorDesc(x), x,
                        scalingParameter(T, beta ), TensorDesc(y), y)
    return y
end

function cudnnSoftmaxBackward(y::DenseCuArray{T,4}, dy::DenseCuArray{T,4}, dx::DenseCuArray{T,4}=dy;
                              algo=CUDNN_SOFTMAX_FAST, # or CUDNN_SOFTMAX_ACCURATE
                              mode=CUDNN_SOFTMAX_MODE_INSTANCE, # or CUDNN_SOFTMAX_MODE_CHANNEL
                              alpha=1.0, beta=0.0) where T
    cudnnSoftmaxBackward(handle(), algo, mode,
                         scalingParameter(T, alpha), TensorDesc(y), y,
                         TensorDesc(dy), dy,
                         scalingParameter(T, beta ), TensorDesc(dx), dx)
    return dx
end