module KnetLayers

using Knet
using UnicodePlots
using Random
using Distributed

import Base.+, Base.*, Base.==

export NeuralNet, NullLayer, LinearLayer, NonlinearLayer, PoolLayer, ConvLayer
export DropoutLayer, BatchnormLayer, ActivationLayer, BiasLayer, DeconvLayer
export Layer, Dense, Pool, Conv1D, Conv2D, Conv3D, Bias, Conv, Deconv
export ReLU, Sigm, Tanh, Dropout, Batchnorm
export operator, weights, repflow, randlayer, depth, train
export isvalid, outdims, flops, nparams, dimflow

#
# LAYER DEFINITIONS AND CONVENIENCES
#

abstract type Layer end
abstract type NullLayer <: Layer end
abstract type ReductionLayer <: Layer end

# TODO: Just const NeuralNet = Array{Layer}  ?
# no, in future want other topologies than just linear, so will need to
# implement trees or graph to hold Layer objects at vertices
mutable struct NeuralNet  # TODO: make some of these types immutable?
    #indims::NTuple{4,Int}
    layers::Array{Layer}
end
NeuralNet() = NeuralNet(Layer[LinearLayer(1)])

depth(N::NeuralNet) = length(N.layers)

#
# Fully Connected (Linear) Layers
#
mutable struct LinearLayer <: Layer
    size::Int
end
# mat(x) = reshape(x, (X1*X2*...*X[D-1],XD))
operator(L::LinearLayer) = (x, w, b) -> w*mat(x) .+ b   # Fully connected layer
nparams(L::LinearLayer) = 2
nparams(L::LinearLayer, indims) =  L.size * prod(indims[1:end-1])
string(L::LinearLayer) = "Dense($(L.size))"
outdims(L::LinearLayer, indims) = (L.size, indims[end])
weights(L::LinearLayer, indims) = (xavier(L.size, prod(indims[1:end-1])),zeros(L.size))
flops(L::LinearLayer, indims) = prod(indims[1:end-1]).^2*L.size + L.size
==(a::LinearLayer, b::LinearLayer) = a.size == b.size
const Linear = LinearLayer
const Dense = LinearLayer

#
#  Nonlinear (Activation) Layers
#
mutable struct NonlinearLayer <: Layer
    op::Function
end
NonlinearOpList = Any[relu, sigm, tanh]
NonlinearOpNames = Dict{Function,String}(relu=>"ReLU",sigm=>"Sigm",tanh=>"Tanh")
ReLU() = NonlinearLayer(relu)
Sigm() = NonlinearLayer(sigm)
Tanh() = NonlinearLayer(tanh)
function leaky_relu(x, alpha=0.2)
    pos = max(0,x)
    neg = min(0,x) * alpha
    return pos + neg
end
LeakyReLU() = NonlinearLayer(leaky_relu)  # TODO
operator(L::NonlinearLayer) = (x) -> (L.op).(x)
nparams(L::NonlinearLayer) = 0
nparams(L::NonlinearLayer, indims) = 0
string(L::NonlinearLayer) = NonlinearOpNames[L.op]*"()"
outdims(L::NonlinearLayer, indims) = indims
weights(L::NonlinearLayer, indims) = ()
flops(L::NonlinearLayer, indims)  = prod(indims)
==(a::NonlinearLayer, b::NonlinearLayer) = a.op == b.op
const ActivationLayer = NonlinearLayer

#
# Convolution Layers
#
mutable struct ConvLayer <: ReductionLayer
    window::Tuple{Int,Int}
    outchan::Int
    stride::Int
    padding::Int
end
ConvLayer(window::Tuple{Int,Int}, outchan) =
    ConvLayer(window::Tuple{Int,Int}, outchan, 1, 0)  # default: stride=1, padding=0
ConvLayer(w::Int, outchan) = ConvLayer((w,w), outchan)
ConvLayer(w::Int, outchan, stride, padding) = ConvLayer((w,w), outchan, stride, padding)
operator(L::ConvLayer) = (x, w, b)-> conv4(w, x, padding=L.padding, stride=L.stride) .+ b
nparams(L::ConvLayer) = 2   # TODO: rename this to nweights?
nparams(L::ConvLayer, indims) = prod(L.window)* indims[3] * L.outchan
string(L::ConvLayer) = "Conv$(length(L.window))D($(L.window), $(L.outchan), $(L.stride), $(L.padding))"
function outdims(L::ConvLayer, X)
    Y1 = 1 + floor(Int, (X[1] + 2*L.padding - L.window[1]) / L.stride)
    Y2 = 1 + floor(Int, (X[2] + 2*L.padding - L.window[2]) / L.stride)
    O = L.outchan
    N = X[3]
    (Y1, Y2, O, N)
end
weights(L::ConvLayer, indims) = (xavier(L.window[1], L.window[2], indims[3], L.outchan),zeros(1,1,L.outchan,1)) # TODO: remove trailing singleton in bias weights?
function flops(L::ConvLayer, indims)
    mw = prod(L.window)
    mw = mw + mw - 1
    return prod(indims[1:2])*mw*L.outchan
end
==(a::ConvLayer, b::ConvLayer) = (a.window==b.window) && (a.outchan==b.outchan) && (a.stride==b.stride) && (a.padding==b.padding)
const Conv = ConvLayer
const Conv1D = ConvLayer
const Conv2D = ConvLayer
const Conv3D = ConvLayer

#
#  Pooling Layers
#
mutable struct PoolLayer <: ReductionLayer
    window::Int
    padding::Int
    stride::Int
    mode::Int
end
PoolLayer() = PoolLayer(2, 0, 2, 0)
operator(L::PoolLayer) = (x) -> pool(x, window=L.window, padding=L.padding, stride=L.stride, mode=L.mode)
nparams(L::PoolLayer) = 0
nparams(L::PoolLayer, indims) = 0
string(L::PoolLayer) = "Pool($(L.window), $(L.padding), $(L.stride), $(L.mode))"
function outdims(L::PoolLayer, X)
    Y1 = 1 + floor(Int, (X[1] + 2*L.padding - L.window) / L.stride)
    Y2 = 1 + floor(Int, (X[2] + 2*L.padding - L.window) / L.stride)
    (Y1, Y2, X[3], X[4])
end
weights(L::PoolLayer, indims) = ()
flops(L::PoolLayer, indims) = prod(indims) # TODO
==(a::PoolLayer, b::PoolLayer) = (a.window==b.window) && (a.mode==b.mode) && (a.stride==b.stride) && (a.padding==b.padding)
const Pool = PoolLayer

#
# Dropout Layers
#
mutable struct DropoutLayer <: Layer
    prob::Float64
end
DropoutLayer() = DropoutLayer(0.1)
operator(L::DropoutLayer) = (x) -> dropout(x, L.prob)
nparams(L::DropoutLayer) = 0
nparams(L::DropoutLayer, indims) = 0
string(L::DropoutLayer) = "Dropout($(L.prob))"
outdims(L::DropoutLayer, indims) = indims
weights(L::DropoutLayer, indims) = ()
flops(l::DropoutLayer, indims) = prod(indims) # TODO
==(a::DropoutLayer, b::DropoutLayer) = a.prob == b.prob
const Dropout = DropoutLayer

#
# Batchnorm Layers
#
mutable struct BatchnormLayer <: Layer
    moments::Knet.BNMoments
end
BatchnormLayer() = BatchnormLayer(bnmoments())
operator(L::BatchnormLayer) = (x, w) -> batchnorm(x, L.moments, w)
nparams(L::BatchnormLayer) = 1
nparams(L::BatchnormLayer, indims) = indims[3]
string(L::BatchnormLayer) = "Batchnorm()"
outdims(L::BatchnormLayer, indims) = indims
weights(L::BatchnormLayer, indims) = (bnparams(indims[3]),)
flops(L::BatchnormLayer, indims) =  prod(indims)  # TODO
==(a::BatchnormLayer, b::BatchnormLayer) = true
const Batchnorm = BatchnormLayer

#
# Bias layer (add vector)
#
struct  BiasLayer <: Layer end
operator(L::BiasLayer) = (x, w) -> x .+ w
nparams(L::BiasLayer) = 1
nparams(L::BiasLayer, indims) =  prod(indims[1:end-1])
string(L::BiasLayer) = "Bias()"
outdims(L::BiasLayer, indims) = indims
weights(L::BiasLayer, indims) = (zeros(indims[1:end-1]...),)
flops(L::BiasLayer, indims) = prod(indims)
==(a::BiasLayer, b::BiasLayer) = true
const Bias = BiasLayer

#
# Deconvolution Layers  (TODO)
#
mutable struct DeconvLayer <: Layer
    window::Tuple{Int,Int}
    outchan::Int
    stride::Int
    padding::Int
end
DeconvLayer(window::Tuple{Int,Int}, outchan) =
    DeconvLayer(window::Tuple{Int,Int}, outchan, 1, 0)  # default: stride=1, padding=0
DeconvLayer(w::Int, outchan) = DeconvLayer((w,w), outchan)
DeconvLayer(w::Int, outchan, stride, padding) = DeconvLayer((w,w), outchan, stride, padding)
operator(L::DeconvLayer) = (x, w, b)-> deconv4(w, x, padding=L.padding, stride=L.stride) .+ b
nparams(L::DeconvLayer) = 2   # TODO: rename this to nweights?
nparams(L::DeconvLayer, indims) = prod(L.window)* indims[3] * L.outchan
string(L::DeconvLayer) = "Deconv$(length(L.window))D($(L.window), $(L.outchan), $(L.stride), $(L.padding))"
function outdims(L::DeconvLayer, X)
# If w has dimensions (W1,W2,...,O,I) and x has
  # dimensions (X1,X2,...,I,N), the result y will have dimensions (Y1,Y2,...,O,N) where
    Y1 = L.window[1] + L.stride*(X[1]-1)-2L.padding
    Y2 = L.window[2] + L.stride*(X[2]-1)-2L.padding
    O = L.outchan
    N = X[3]
    (Y1, Y2, O, N)
end
weights(L::DeconvLayer, indims) = (xavier(L.window[1], L.window[2], indims[3], L.outchan),zeros(1,1,L.outchan,1))
function flops(L::DeconvLayer, indims)
    mw = prod(L.window)
    mw = mw + mw - 1
    return prod(indims[1:2])*mw*L.outchan
end
==(a::DeconvLayer, b::DeconvLayer) = (a.window==b.window) && (a.outchan==b.outchan) && (a.stride==b.stride) && (a.padding==b.padding)
const Deconv = DeconvLayer
const Deconv1D = DeconvLayer
const Deconv2D = DeconvLayer
const Deconv3D = DeconvLayer

#
# NeuralNet and Layer methods
#
operator(A::NeuralNet) = function (w, x)
    i = 1
    for L in A.layers
        n = nparams(L)
        x = operator(L)(x, w[i:i+n-1]...)
        i += n
    end
    return x
end
function weights(net::NeuralNet, indims; atype=Array{Float32})
    #kaiming(h, w, i, o) = sqrt(2/(w*h*o)) .* randn(h, w, i, o)
    W = Any[]
    D = indims  # running value of data size as it passes through net
    for L in net.layers
        push!(W, weights(L, D)...)
        #info(L, " -> ", size(W[end]))
        D = outdims(L, D)
    end
    np = sum([length(w) for w in W])
    return (map(a -> convert(atype, a), W), np)
end

function +(A::Array{Layer,1}, B::Array{Layer,1})
    m = length(A)
    n = length(B)
    C = Array{Layer}(undef, m+n)
    C[1:m] = copy(A)
    C[m+1:m+n] = copy(B)
    return C
end
+(A::NeuralNet, B::NeuralNet) = NeuralNet(A.layers + B.layers)
function *(n::Int, L::Layer)
    A = Array{Layer}(undef, n)
    for i in 1:n
        A[i] = L
    end
    return A
end
function *(n::Int, A::Array{Layer,1})
    m = n*length(A)
    B = Array{Layer}(undef, m)
    for i in 1:m
        @info "$i <- $((i-1)%length(A)+1)"
        B[i] = A[(i-1) % length(A) + 1]
    end
    return B
end
*(n::Int, A::NeuralNet) = NeuralNet(n*A.layers)
function ==(A::NeuralNet, B::NeuralNet)
    if length(A.layers) != length(B.layers)
        return false
    else
        return all(A.layers .== B.layers)
    end
end

string(A::NeuralNet) = "NeuralNet(Layer[" * join([string(x) for x in A.layers], ",") * "])"
Base.show(io::IO, A::NeuralNet) = print(io, string(A))

function Base.show(io::IO, ::MIME"text/plain", A::NeuralNet)
    println(io, "$(depth(A))-layer NeuralNet:")
    for (i, L) in enumerate(A.layers)
        i < depth(A) ? println(io, "  ├ $(string(L)),") : print(io, "  └ $(string(L))")
    end
end

function repflow(A::NeuralNet, indims)
    n = zeros(depth(A)+1)
    Din = indims
    n[1] = prod(Din[1:end-1])
    for (i, L) in enumerate(A.layers)
        Dout = outdims(L, Din)
        n[i+1] = prod(Dout[1:end-1])
        Din = Dout
    end
    return n ./ n[1]
end
function dimflow(A::NeuralNet, indims)
    n = zeros(depth(A)+1)
    D = indims
    @info D
    for L in A.layers
        D = outdims(L, D)
        @info "$L -> $D"
    end
end
function flops(A::NeuralNet, indims::NTuple{4,Int})
    # approximate FLOPS of network
    D = indims
    n = 0
    for (k, L) in enumerate(A.layers)
        n += flops(L, D)
        D = outdims(L, D)
    end
    return n
end

Base.isvalid(dims::NTuple{N,Int}) where N = all(dims .> 0)
Base.isvalid(L::Layer, indims) = isvalid(outdims(L, indims))
function Base.isvalid(A::Array{Layer}, indims)
    dims = indims
    gonelinear = false
    for L in A
        if gonelinear && (typeof(L) == ConvLayer || typeof(L) == PoolLayer ||
            typeof(L) == BatchnormLayer)
            #warn("Conv, Pool, or Batchnorm after Linear is INVALID.")
            return false
        elseif !gonelinear && typeof(L) == LinearLayer
            gonelinear = true
        end
        dims = outdims(L, dims)
        isvalid(dims) || return false
    end
    return true
end
Base.isvalid(A::NeuralNet) = Base.isvalid(A.layers)

function nparams(net::NeuralNet, indims)
    n = 0
    D = indims  # running value of data size as it passes through net
    for L in net.layers
        n += nparams(L, D)
        D = outdims(L, D)
    end
    return n
end

#
# TRAINING AND DATA LOADING HELPERS
#

include("mnist.jl")

xtrn, ytrn, xtst, ytst = mnist()

function train(net::NeuralNet; epochs=10, fast=true, lr=0.001, seed=0xc0ffee,
    optfunc=Adam, ftrain=1.0, atype=KnetArray{Float32})
    Random.seed!(seed)
    gpuid = 0
    if gpu() >= 0 && atype == KnetArray{Float32}
        gpuid = myid() % Knet.gpuCount()
        gpu(gpuid) # activate appropriate GPU
        device="gpu$gpuid"
        atype = KnetArray{Float32}
    else
        device="cpu"
        atype = Array{Float32}
    end
    #infoprefix="$(gethostname())$(myid())$device: "
    @info "received $net"
    batchsize = 100
    Ntrain = round(Int, size(xtrn,4)*ftrain)
    dtrn = minibatch(xtrn[:,:,1,1:Ntrain], ytrn[1:Ntrain], batchsize; xtype=atype)  # keep this on CPU for Augmentor
    dtst = minibatch(xtst, ytst, batchsize; xtype=atype)
    nx, ny, nz = size(xtrn)[1:3]
    indims = (nx, ny, nz, 1)
    (w, np) = weights(net, indims, atype=atype)

    predict = operator(net)
    loss(w, x, ygold) = nll(predict(w, x), ygold)
    lossgradient = grad(loss)
    params = Array{Any}(undef, length(w))
    for k in 1:length(w)
        params[k] = optfunc(lr=lr)
    end

    if !fast
        @info "epoch 0: $(accuracy(w,dtst,predict))"
    end
    xaug = zeros(Float32, nx, ny, 1, batchsize)
    accvtime = zeros(epochs)
    traintime = time()
    for epoch in 1:epochs
        for (x, y) in dtrn
            #x = xtmp
            x = reshape(x, (nx, ny, nz, batchsize))
            #info("type(x): $(typeof(x)), type(y): $(typeof(y)), type(w): $(typeof(w))")
            #info(size(x))
            g = lossgradient(w, x, y)
            update!(w, g, params)
        end
        if !fast
            accvtime[epoch] = accuracy(w,dtst,predict)
            @info "epoch $epoch: $(accvtime[epoch])"
        end
    end
    if gpu() >= 0
        Knet.cudaDeviceSynchronize()
    end
    traintime = time() - traintime
    testacc = accuracy(w, dtst, predict)
    @info "Test accuracy = $testacc, training time = $traintime"
    if !fast
        @info "Accuracy vs Epoch:"
        println(lineplot(accvtime, color=:green))
    end
    return (testacc, traintime, np)
end
train(A::Array{Layer}; kwargs...) = train(NeuralNet(A); kwargs...)

# Example NeuralNets

mnist_mlp = NeuralNet(Layer[Dense(128), Bias(), ReLU(), Dense(64), Bias(), ReLU(), Dense(10)])

mnist_lenet = NeuralNet(Layer[
        Conv2D(5, 20), Tanh(), Pool(),
        Conv2D(5, 50), Tanh(), Pool(),
        Dense(500), Tanh(),
        Dense(10)
    ])

function test(net::NeuralNet = mnist_lenet; kwargs...)
    #show(STDOUT, "text/plain", net)
    acc, t = train(net; kwargs...)
end

end

if PROGRAM_FILE == "KnetLayers.jl"
    #KnetLayers.dimflow(KnetLayers.mnist_lenet,(28,28,1,1))
    KnetLayers.test(KnetLayers.mnist_mlp, fast=false)
    KnetLayers.test(KnetLayers.mnist_lenet, fast=false)
end


using ConstraintSolver, JuMP, MathOptInterface, JSON

if !@isdefined CS
    const CS = ConstraintSolver
end
const MOI = MathOptInterface
const MOIU = MOI.Utilities
include("../../test/sudoku_fcts.jl")
# include("../../visualizations/plot_search_space.jl")

function parseJSON(json_sums)
    sums = []
    for s in json_sums
        indices = Tuple[]
        for ind in s["indices"]
            push!(indices, tuple(ind...))
        end

        push!(sums, (result = s["result"], indices = indices, color = s["color"]))
    end
    return sums
end

function solve_all(filenames; benchmark = false, single_times = true)
    ct = time()
    for (i, filename) in enumerate(filenames)
        sums = parseJSON(JSON.parsefile("data/$(filename)"))

        # plot_killer(zeros(Int, (9,9)), sums, filename; fill=false)
        # continue

        m = CS.Optimizer(logging = [], time_limit = 20)

        x = [[MOI.add_constrained_variable(m, MOI.Integer()) for i in 1:9] for j in 1:9]
        for r in 1:9, c in 1:9
            MOI.add_constraint(m, x[r][c][1], MOI.GreaterThan(1.0))
            MOI.add_constraint(m, x[r][c][1], MOI.LessThan(9.0))
        end

        for s in sums
            saf = MOI.ScalarAffineFunction{Float64}(
                [MOI.ScalarAffineTerm(1.0, x[ind[1]][ind[2]][1]) for ind in s.indices],
                0.0,
            )
            MOI.add_constraint(m, saf, MOI.EqualTo(convert(Float64, s.result)))
            # MOI.add_constraint(m, [x[ind[1]][ind[2]][1] for ind in s.indices], CS.CPE.AllDifferent(length(s.indices)))
        end

        # sudoku constraints
        moi_add_sudoku_constr!(m, x)

        if single_times
            GC.enable(false)
            MOI.optimize!(m)
            status = MOI.get(m, MOI.TerminationStatus())
            GC.enable(true)
            println(i - 1, ", ", MOI.get(m, MOI.SolveTime()))
        else
            GC.enable(false)
            MOI.optimize!(m)
            status = MOI.get(m, MOI.TerminationStatus())
            GC.enable(true)
        end
        if !benchmark
            println("Status: ", status)
            var_x = fill(MOI.VariableIndex(0), (9, 9))
            for r in 1:9
                var_x[r, :] = [x[r][c][1] for c in 1:9]
            end
            if status == MOI.OPTIMAL
                solution = zeros(Int, 9, 9)
                for r in 1:9
                    solution[r, :] =
                        [MOI.get(m, MOI.VariablePrimal(), x[r][c][1]) for c in 1:9]
                end
                @assert jump_fulfills_sudoku_constr(solution)
            else
                println("NOT SOLVED TO OPTIMALITY")
            end
        end
        com = nothing
        GC.gc()
    end
    println("")
    tt = time() - ct
    println("total time: ", tt)
    println("avg: ", tt / length(filenames))
end

function main(; benchmark = false, single_times = true)
    solve_all(
        [
            "niallsudoku_5500",
            # "niallsudoku_5501",
            # "niallsudoku_5502",
            # "niallsudoku_5503",
            # "niallsudoku_6417",
            # "niallsudoku_6249",
        ];
        benchmark = benchmark,
        single_times = single_times,
    )
    # solve_all(from_file("hardest.txt"), "hardest")
end


using CellwiseEditingDifferentiation
using Distributions
using Lora

import Base: rand

immutable MvBeta <: ContinuousMultivariateDistribution
  betas::Vector{Beta}
end

Base.rand(d::MvBeta) = [rand(d.betas[i]) for i in 1:length(d.betas)]

ncells = 20

λ = Hyperparameter(:λ)

μ = Data(:μ)
coverage = Data(:coverage)
edited = Data(:edited)
rate = Data(:rate)

function wlogtarget(w::Float64, s::Vector)
  v = s[2]*(1-s[2])*inv_logit(w)
  a, b = beta_pars_from_mv(s[2], v)
  mapreduce(x -> logpdf(Beta(a, b), x), +, s[7])+vpcprior(s[2], v, a, b, s[1])+logdvdw(w, s[2])
end

w = BasicContUnvParameter(:w, logtarget=wlogtarget, nkeys=7)

function psetpdf(p::Vector{Float64}, s::Vector)
  a, b = beta_pars_from_mv(s[2], s[2]*(1-s[2])*inv_logit(s[6]))
  MvBeta([Beta(a+s[4][i], b+s[3][i]-s[4][i]) for i in 1:ncells])
end

p = BasicContMuvParameter(:p, setpdf=psetpdf, nkeys=7)

model = GenericModel(
  [λ, μ, coverage, edited, rate, w, p],
  [λ w; μ w; μ p; coverage p; edited p; w p],
  isindexed=false
)

init = Dict(
  :λ => 10.,
  :μ => 0.5,
  :coverage => fill(10.0, 20),
  :edited => rand([1., 2., 3., 8., 9.], 20),
  :w => 0.
)
init[:rate] = float(init[:edited]./init[:coverage])
init[:p] = rand(psetpdf(fill(NaN, ncells), [init[k] for k in [:λ, :μ, :coverage, :edited, :rate, :w]]))

wjob = BasicMCJob(
  model,
  RAM(),
  BasicMCRange(nsteps=100, burnin=10),
  init;
  pindex=6,
  resetpstate=false,
  verbose=false
)

mcrange = BasicMCRange(nsteps=1100, burnin=100)

outopts = [Dict(:monitor => [:value]), Dict(:monitor => [:value])]

gibbsjob = GibbsJob(model, Dict(:w => wjob), mcrange, init)

@time run(gibbsjob)

output(gibbsjob)

chains = Dict(gibbsjob)

v = map(x -> init[:μ]*(1-init[:μ])*inv_logit(x), chains[:w].value)

writecsv("pchains.csv", chains[:p].value)
writecsv("vchain.csv", v)


module CbfComparison

    using LinearAlgebra
    using Plots
    using Logging

    struct TestCondition
        x_0::Array{Float64}
        x_goal::Array{Float64}
        x_o1::Array{Float64}
        x_o2::Array{Float64}
        d_obs::Float64
        sampling_time::Float64
    end

    mutable struct RobotState
        x::Array{Float64}
        dx::Array{Float64}
        x_history::Matrix{Float64}
        dx_history::Matrix{Float64}
        time_history::Array{Float64}
    end

    function min_dist(param::TestCondition, state::RobotState)

        dist_o1 = norm(state.x - param.x_o1)
        dist_o2 = norm(state.x - param.x_o2)

        if dist_o1 < dist_o2
            dist = dist_o1 - param.d_obs
            return dist, state.x-param.x_o1
        else
            dist = dist_o2 - param.d_obs
            return dist, state.x-param.x_o2
        end

    end

    function delta_h(param::TestCondition, state::RobotState)

        dist_o1 = norm(state.x - param.x_o1)
        dist_o2 = norm(state.x - param.x_o2)

        if dist_o1 < dist_o2
            dist = dist_o1
            return (state.x-param.x_o1)/dist
        else
            dist = dist_o2
            return (state.x-param.x_o2)/dist
        end

    end

    function check_goal(param::TestCondition, state::RobotState)

        dist_goal = norm(state.x - param.x_goal)

        if dist_goal < 0.01
            return true
        else
            return false
        end
    end

    function update_state(param::TestCondition, state::RobotState, u, time)

        state.dx = u
        state.x = state.x + state.dx * param.sampling_time
        state.x_history = vcat(state.x_history, reshape(state.x, (1, 2)))
        state.dx_history = vcat(state.dx_history, reshape(state.dx, (1, 2)))
        push!(state.time_history, time)

    end

    function controller(param::TestCondition, state::RobotState, k_att, k_rep, rho0)

            # apf

            # attractive potential
            delta_u_att = k_att * (state.x - param.x_goal)

            # repulsive potential
            rho, xor = min_dist(param, state)

            if rho > rho0
                delta_u_rep = zeros(2)
            else
                delta_u_rep = - k_rep/rho^2 * (1/rho - 1/rho0) * xor
            end

            # artificail potential
            f_apf = - delta_u_att - delta_u_rep

            return f_apf
    end

    function controller(param::TestCondition, state::RobotState, k, alpha)

            # cbf

            # nominal controller
            v_des = k * (param.x_goal - state.x)

            # h(x)
            h, xor = min_dist(param, state)

            # Psi
            lgh_t = delta_h(param, state)
            psi = lgh_t' * v_des + alpha * h

            # input for satisfying barrier function
            if psi < 0
                v_safe = -lgh_t / (lgh_t' * lgh_t) * psi
            else
                v_safe = [0; 0]
            end

            return v_des + v_safe
    end

    function apf_cbf()
        # place holder
    end

    function simulate(param::TestCondition, gains, name)

        # initialize
        x = param.x_0
        dx = zeros(2)
        x_history = reshape(x, (1, 2))
        dx_history = reshape(dx, (1, 2))
        time_history = [0.0]
        state = RobotState(x, dx, x_history, dx_history, time_history)

        # start
        count = 0
        while true

            f = controller(param, state, gains...)

            # update robot state
            time = param.sampling_time * count
            update_state(param, state, f, time)

            # check goal state
            if check_goal(param, state) == true #|| count > 1000
                break
            end

            count = count + 1
        end

        # plot results
        plot_pic(param, state, name)
        plot_gif(param, state, name)

    end

    function circle_shape(x, y, r)
        theta = LinRange(0, 2pi, 500)
        x .+ r * sin.(theta), y .+ r * cos.(theta)
    end

    function plot_pic(param::TestCondition, state::RobotState, name)
        plot(state.x_history[:, 1], state.x_history[:, 2], label="trajectory", xlabel="x [m]", ylabel="y [m]")
        plot!(circle_shape(param.x_o1[1], param.x_o1[2], param.d_obs), seriestype=[:shape,], c=:blue, linecolor=:black, fillalpha=0.2, aspectratio=1, label="Obstacle 1")
        plot!(circle_shape(param.x_o2[1], param.x_o2[2], param.d_obs), seriestype=[:shape,], c=:purple, linecolor=:black, fillalpha=0.2, aspectratio=1, label="Obstacle 2")
        png(name)

        p1 = plot(state.time_history[:, 1], state.x_history[:, 1], ylabel="x [m]")
        p2 = plot(state.time_history[:, 1], state.x_history[:, 2], xlabel="time [s]", ylabel="y [m]")
        plot(p1, p2, layout = (2, 1), legend = false)
        png(name*"_pos_time")

        p3 = plot(state.time_history[:, 1], state.dx_history[:, 1], ylabel="dx [m/s]")
        p4 = plot(state.time_history[:, 1], state.dx_history[:, 2], xlabel="time [s]", ylabel="dy [m/s]")
        plot(p3, p4, layout = (2, 1), legend = false)
        png(name*"_vel_time")
    end

    function plot_gif(param::TestCondition, state::RobotState, name)
        # initialize a plot
        plt = plot(
            1,
            label="trajectory",
            xlabel="x [m]",
            ylabel="y [m]"
        )

        # add obstacles
        plot!(circle_shape(param.x_o1[1], param.x_o1[2], param.d_obs), seriestype=[:shape,], c=:blue, linecolor=:black, fillalpha=0.2, aspectratio=1, label="Obstacle 1")
        plot!(circle_shape(param.x_o2[1], param.x_o2[2], param.d_obs), seriestype=[:shape,], c=:purple, linecolor=:black, fillalpha=0.2, aspectratio=1, label="Obstacle 2")

        # build an animated gif by pushing new points to the plot
        anim = @animate for i=1:length(state.time_history)
            push!(plt, state.x_history[i, 1], state.x_history[i, 2])
        end every 10
        gif(anim, name*"_xy_2d.gif", fps = 30)
    end

end # module

using .CbfComparison

# test condition
x_0 = [0.0, 0.0]
x_goal = [3.0, 5.0]
x_o1 = [1.0, 2.0]
x_o2 = [2.5, 3.0]
d_obs = 0.5
sampling_time = 0.01

param = CbfComparison.TestCondition(x_0, x_goal, x_o1, x_o2, d_obs, sampling_time)

# simulation for APF
k_att = 1.0
k_rep = 1.0
rho0 = 0.5
gains = [k_att, k_rep, rho0]
CbfComparison.simulate(param, gains, "apf")

# simulation for CBF
# parameters
k = 1.0
alpha = 1.0
gains = [k, alpha]
CbfComparison.simulate(param, gains, "cbf")


using DOIHS
using ValidatedNumerics
using Distributions
using Base.Test

@testset "CRRA utility" begin

    @test_throws DomainError crra_utility_function(0.1)
    @test_throws DomainError crra_utility(0.1, 0.1)
    @test_throws DomainError crra_utility(-1, 1)
    @test crra_utility(0, 1) ≡ -∞

    u = crra_utility_function(1)
    @test_throws DomainError u(-10)
    @test u(0) ≡ -∞
    @test u(1) == 0.0
    @test isa(u(1), Float64)

    u = crra_utility_function(2)
    @test_throws DomainError u(-10)
    @test u(0.0) ≡ -∞
    @test u(1) == 0.0
    @test isa(u(1), Float64)

    @test_throws DomainError crra_u′(-0.1, 0.1)
    @test crra_u′(2, 1) == 1/2

    @test crra_inv_u′(crra_u′(9.0, 3.0), 3.0) ≈ 9.0
end

@testset "function approximation" begin

    function approxerror(basis, f, n=100)
        α = basis_matrix(basis) \ f(collocation_points(basis))
        x = linspace(basis, n)
        y = basis_matrix(basis, x) * α
        maximum(abs(y-f(x)) for (x,y) in zip(x,y))
    end

    @test approxerror(ChebyshevBasis(10), exp) ≤ 1e-9

    @test approxerror(IntervalAB(0..pi, ChebyshevBasis(10)), sin) ≤ 1e-7

end

@testset "Chebyshev extrapolation" begin
    basis = IntervalAB(0..10, ChebyshevBasis(10; extrapolate = true))
    f(x) = x^2 - 5*x + 6
    g = basis \ f
    @test g(11) ≈ f(11)
    @test g(-1) ≈ f(-1)
end

@testset "linear interpolation" begin
    @test_throws Exception LinearInterpolation([2.0,1.0])
    li = LinearInterpolation(Float64[1,2,4])
    @test collocation_points(li) == [1.0,2.0,4.0]
    @test domain(li) == 1.0..4.0
    @test degf(li) == 3
    y = [5.0, 3.0, 7.0]
    lif = li \ y
    @test collocation_values(lif) == y
    @test lif(1.0) == 5.0
    @test lif(1.5) == 4.0
    @test lif(1.75) == 3.5
    @test lif(3.0) == 5.0
    @test lif(3.5) == 6.0
    @test lif(3.75) == 6.5
    @test lif(4.0) == 7.0
    @test_throws Exception lif(-1.0)
    @test_throws Exception lif(5.0)
end

@testset "extrapolate level" begin
    inner_basis = IntervalAB(1..10, ChebyshevBasis(3))
    basis = ExtrapolateLevel(inner_basis)
    @test collocation_points(inner_basis) == collocation_points(basis)
    @test degf(inner_basis) == degf(basis)
    y = [2.0, 3.0, 7.0]
    ipf = basis \ y
    @test collocation_values(ipf) ≈ y
    a = ipf(1.0)
    b = ipf(10.0)
    @test ipf(-1.0) ≈ a
    @test ipf(11.0) ≈ b
end

@testset "quadrature" begin

    d = Truncated(Normal(0,1), -1, 2)
    q = quadrature(10, d)
    m = mean(d)
    @test isapprox(q(identity), m; rtol=1e-6)
    @test isapprox(q(x->(x-m)^2), var(d); rtol=1e-6)

end

@testset "quadrature normal" begin

    μ = 1.0
    σ = 2.0
    q = quadrature_normal(10, μ, σ)

    n = Normal(μ,σ)

    @test mean(n) ≈ q(identity)
    @test var(n) ≈ q(x->(x-μ)^2)

end


#####################################################
### Project of IN480 - From Python 2.7 to Julia 0.6
### --- SIMPLEXN - serial version (everywhere) ---
### Authors: Fabio Fatelli, Emanuele Loprevite
#####################################################

# Generation of the output model vertices in a multiple extrusion of a LAR model
@everywhere function larExtrude1{T<:Real}(model::Tuple{Array{Array{T,1},1},Array{Array{Int64,1},1}},pattern::Array{Int64,1})
	V, FV = model
	d, m = length(FV[1]), length(pattern)
	coords = cumsum(append!([0],abs.(pattern))) # built-in function cumsum
	offset, outcells, rangelimit = length(V), Array{Int64}(m,0), d*m
	for cell in FV
		tube = [v+k*offset for k in 0:m for v in cell]
		celltube = Int64[]
		for k in 1:rangelimit
			append!(celltube,tube[k:k+d])
		end
		outcells = hcat(outcells,reshape(celltube,d*(d+1),m)')
	end
	cellGroups = Int64[]
	for k in 1:m
		if pattern[k]>0
			cellGroups = vcat(cellGroups,outcells[k,:])
		end
	end
	outVertices = [vcat(v,z) for z in coords for v in V]
	outCellGroups = Array{Int64,1}[]
	for k in 1:d+1:length(cellGroups)
		append!(outCellGroups,[cellGroups[k:k+d]])
	end
	return outVertices, outCellGroups
end

# Generation of simplicial grids of any dimension and shape
@everywhere function larSimplexGrid1(shape::Array{Int64,1})
	model = [Int64[]],[[0]] # the empty simplicial model
	for item in shape
		model = larExtrude1(model,repmat([1],item))
	end
	return model
end

# Extraction of non-oriented (d-1)-facets of d-dimensional simplices
@everywhere using Combinatorics # for the function combinations()

@everywhere function larSimplexFacets(simplices::Array{Array{Int64,1},1})
	out = Array{Int64,1}[]
	d = length(simplices[1])
	for simplex in simplices
		append!(out,collect(combinations(simplex,d-1)))
	end
	return sort!(unique(out),lt=lexless)
end

# Transformation to triangles by sorting circularly the vertices of faces
@everywhere function quads2tria{T<:Real}(model::Tuple{Array{Array{T,1},1},Array{Array{Int64,1},1}})
	V, FV = model
	if typeof(V) != Array{Array{Float64,1},1}
		V = convert(Array{Array{Float64,1},1},V)
	end
	out = Array{Int64,1}[]
	nverts = length(V)-1
	for face in FV
		arr = [V[v+1] for v in face]
		centroid = sum(arr)/length(arr)
		append!(V,[centroid])
		nverts += 1
		v1, v2 = V[face[1]+1]-centroid, V[face[2]+1]-centroid
		v3 = cross(v1,v2)
		if norm(v3) < 1/(10^3)
			v1, v2 = V[face[1]+1]-centroid, V[face[3]+1]-centroid
			v3 = cross(v1,v2)
		end
		transf = inv(hcat(v1,v2,v3)')
		verts = [(V[v+1]'*transf)'[1:end-1] for v in face]
		tcentroid = sum(verts)/length(verts)
		tverts = [v-tcentroid for v in verts]
		iterator = collect(zip(tverts,face))
		rverts = [[atan2(reverse(iterator[i][1])...),iterator[i][2]] for i in 1:length(iterator)]
		rvertsS = sort(rverts,lt=(x,y)->isless(x[1],y[1]))
		ord = [pair[2] for pair in rvertsS]
		append!(ord,ord[1])
		edges = [[i[2],ord[i[1]+1]] for i in enumerate(ord[1:end-1])]
		triangles = [prepend!(edge,nverts) for edge in edges]
		append!(out,triangles)
	end
	return V, out
end


using PlotlyJS

function contour1()
    x = y = [-2*pi + 4*pi*i/100 for i in 1:100]
    z = [sin(x[i]) * cos(y[j]) * sin(x[i]*x[i]+y[j]*y[j])/log(x[i]*x[i]+y[j]*y[j]+1)
         for i in 1:100 for j in 1:100]
    z_ = [z[i:i+99] for i in 1:100:10000]

    data = contour(;z=z_, x=x, y=y)

    plot(data)
end

function contour2()
    z =  [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    data = contour(;z=z)

    layout = Layout(;title="Basic Contour Plot")
    plot(data, layout)
end

function contour3()
    x = [-9, -6, -5 , -3, -1]
    y = [0, 1, 4, 5, 7]
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    trace = contour(x=x, y=y, z=z)

    layout = Layout(title="Setting the X and Y Coordinates in a Contour Plot")
    plot(trace, layout)
end

function contour4()
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    data = contour(;z=z, colorscale="Jet")

    layout = Layout(;title="Colorscale for Contour Plot")
    plot(data, layout)
end

function contour5()
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    data = contour(;z=z,
                   colorscale="Jet",
                   autocontour=false,
                   contours=Dict(:start=>0, :end=>8, :size=>2))

    layout = Layout(;title="Customizing Size and Range of Contours")
    plot(data, layout)
end

function contour6()
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    data = contour(;z=z, colorscale="Jet", dx=10, x0=5, dy=10, y0=10)

    layout = Layout(;title="Customizing Spacing Between X and Y Axis Ticks")

    plot(data, layout)
end

function contour7()
    z = [NaN NaN NaN 12  13  14  15  16
         NaN 1   NaN 11  NaN NaN NaN 17
         NaN 2   6   7   NaN NaN NaN 18
         NaN 3   NaN 8   NaN NaN NaN 19
         5   4   10  9   NaN NaN NaN 20
         NaN NaN NaN 27  NaN NaN NaN 21
         NaN NaN NaN 26  25  24  23  22]'

    p1 = plot(contour(;z=z, showscale=false))
    p2 = plot(contour(;z=z, connectgaps=true, showscale=false))
    p3 = plot(heatmap(;z=z, zsmooth="best",showscale=false))
    p4 = plot(heatmap(;z=z, zsmooth="best", connectgaps=true, showscale=false))

    p = [p1 p2; p3 p4]

    relayout!(p, title="Connect the Gaps Between Null Values in the Z Matrix")

    p
end

function contour8()
    z = [2  4   7   12  13  14  15  16
         3  1   6   11  12  13  16  17
         4  2   7   7   11  14  17  18
         5  3   8   8   13  15  18  19
         7  4   10  9   16  18  20  19
         9  10  5   27  23  21  21  21
         11 14  17  26  25  24  23  22]

    p1 = plot(contour(;z=z, line_smoothing=0))
    p2 = plot(contour(;z=z, line_smoothing=0.85))

    p = [p1 p2]

    relayout!(p, title="Smoothing Contour Lines")

    p
end

function contour9()
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    data = contour(;z=z, contours_coloring="heatmap")

    layout = Layout(;title="Smooth Contour Coloring")
    plot(data, layout)
end

function contour10()
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    data = contour(;z=z, colorscale="Jet", contours_coloring="lines")

    layout = Layout(;title="Contour Lines")
    plot(data, layout)
end

function contour11()
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
   data = contour(;z=z,
                  colorscale=[[0, "rgb(166,206,227)"],
                              [0.25, "rgb(31,120,180)"],
                              [0.45, "rgb(178,223,138)"],
                              [0.64, "rgb(51,160,44)"],
                              [0.85, "rgb(251,154,153)"],
                              [1, "rgb(227,26,28)"]])

   layout = Layout(;title="Custom Contour Plot Colorscale")

   plot(data, layout)
end

function contour12()
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    data = contour(;z=z,
                   colorbar=attr(;title="Color Bar Title",titleside="right",
                                 titlefont=attr(;size=14,
                                                family="Arial, sans-serif")))

    layout = Layout(;title="Colorbar with Title")
    plot(data,layout)
end

function contour13()
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    data = contour(;z=z,
                   colorbar=attr(;thickness=75, thicknessmode="pixels",
                                 len=0.9, lenmode="fraction",
                                 outlinewidth=0))

    layout = Layout(;title="Colorbar Size for Contour Plots")
    plot(data,layout)
end

function contour14()
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    data = contour(;z=z,
                   colorbar=attr(;ticks="outside", dtick=1,
                                 tickwidth=2, ticklen=10,
                                 tickcolor="grey", showticklabels=true,
                                 tickfont_size=15, xpad=50))

    layout = Layout(;title="Styling Color Bar Ticks for Contour Plots")
    plot(data,layout)
end


export initGMM, fitGMM

using CUDA
using ParallelKMeans

# NOTE:
#  - Julia is column major by default => X ∈ R^{d×n}
#  - ω are the weights of the data points
function initGMM(X::AbstractArray{T}, ω::AbstractArray{T}, K::Int) where {T<:Real}
    dims = size(X)
    @assert length(dims) == 2

    d = dims[1]
    n = dims[2]

    @assert n > 0
    @assert d > 0

    @assert K > 0

    gmm = GMM{T}(d, K)

    # start by a round of k-means to quickly generate good guesses for the cluster centroids
    # (and potentially the number of components)
    kmeans_results = kmeans(Yinyang(), X, K; weights=ω)

    # seed the centroids
    copyto!(gmm.μ, kmeans_results.centers)

    gmm
end

function fitGMM(in_X::AbstractArray{T}, in_ω::AbstractArray{T}, K; max_iterations=100, tolerance=1e-8) where {T<:Real}
    gmm = initGMM(in_X, in_ω, K)

    X = CuArray(in_X)
    ω = CuArray(in_ω)

    n = size(X)[2]

    iteration = 0

    # todo: make this not necessarily dependent on CUDA
    # NOTE: this matrix can easily get HUGE
    log_γ = CUDA.zeros(n, K)

    last_log_likelihood = typemin(T)

    while iteration < max_iterations
        # 1.) E-Step

		# we use logarithms in the following
		log_π = transpose(log.(gmm.π))

        # NOTE: what we would like to do is compute all of those in parallel, but this could have a huge memory footprint
        # TODO: spawn in parallel! wait for completion (see https://cuda.juliagpu.org/stable/usage/multitasking/), pin memory
        for k in 1:K
            # compute log(p(X|μ, Σ))
            #
            # since p(X|μ,Σ) = 1/sqrt((2π)^d|Σ|) * exp(-1/2 (x-μ)^t Σ^{-1} (x-μ))
            #   => log(p(X|μ,Σ)) = -0.5 * (d * log(2π) + log(|Σ|)) - 0.5 log((x-μ)^t Σ^{-1} (x-μ))
            # and since right now, we only allow diagonal covariances, Σ^{-1} is the diagonal matrix with
            # 1/σ_i on the diagonal
            #
            # since each point occurs ω_i times, we can multiply the terms in the sum by ω_i
            #
            # TODO: once Distributions.jl is GPU friendly, we should replace this entirely
            log_γ[:, k] = get_normalizer(gmm.σ[:, k], gmm.d) .- 0.5 * sum(1 ./ gmm.σ[:, k] .* transpose(ω) .* (X .- gmm.μ[:, k]).^2; dims=1)
        end

        # for each point, find the maximum log probability from log(π_k) + log(p(X_i|μ_k, Σ_k)) - log(p(X_i)) = log(π_k) + log_γ[:, k]
        max_log_p_k_per_point = maximum(log_π .+ log_γ, dims=2)

        # this seems necessary for numerical stability
        log_p_x = log.(sum(exp.(log_γ .+ log_π .- max_log_p_k_per_point), dims=2)) .+ max_log_p_k_per_point

        log_γ = log_γ .- log_p_x

        # at this point, log_γ[i, k] = p(X_i | μ_k, Σ_k) / p(X_i)
        current_log_likelihood = sum(log_p_x)

        if abs(current_log_likelihood - last_log_likelihood) < tolerance || current_log_likelihood < last_log_likelihood
            break
        end

        last_log_likelihood = current_log_likelihood

        # 2.) M-Step
        max_log_γ_per_x = maximum(log_γ, dims=1)
        log_Γ = max_log_γ_per_x .+ log.(sum(exp.(log_γ .- max_log_γ_per_x), dims=1))

		# TODO: check whether we should a loop to distribute that over K loops
		gmm.μ = ((transpose(ω) .* X) * exp.(log_γ)) ./ exp.(log_Γ)

		# TODO: use streams and distribute?
		# NOTE: this here runs into allowscalar - problems
		gmm.σ = reduce(hcat, [sum(transpose(exp.(log_γ[:, k]) .* ω) .* (X .- gmm.μ[:, k]).^2; dims=2) / exp(log_Γ[k]) for k in 1:K])

		log_π_log_Γ = log_π .+ log_Γ
		max_log_π_log_Γ	= maximum(log_π_log_Γ)

		gmm.π = transpose(exp.(log_π .+ log_Γ .- max_log_π_log_Γ .- log(sum(exp.(log_π_log_Γ .- max_log_π_log_Γ)))))

		# TODO: handle this correctly!
		gmm.π = clamp.(gmm.π, 1e-8, 1.0f0-1e-8)

		iteration += 1
	end

	gmm
end


using Random
using Distributions

include("utils.jl")

abstract type Signal <: AbstractArray{Float64,1} end # maybe extend to N-d SPRIGHT?
Base.IndexStyle(::Signal) = IndexLinear()

struct TestSignal <: Signal
    n::Int64
    k::Int64
    locs::Array{Int64,1}
    strengths::Array{Float64,1}
    noise_sd::Float64
    signal_t::Array{Float64,1}

    function TestSignal(n::Int64, locs::Array{Int64,1}, strengths::Array{Float64,1} = ones(length(locs)), noise_sd::Float64 = 0.0)
        N = 2^n
        k = locs |> length |> log2 |> ceil |> Int64
        noise = rand(Normal(0.0, noise_sd), N)
        wht = zeros(N)
        for (l, s) in zip(locs, strengths)
            wht[l+1] = s
        end
        signal_t = fwht(wht) + noise
        new(n, k, locs, strengths, noise_sd, signal_t)
    end
end

Base.size(s::TestSignal) = size(s.signal_t)
Base.getindex(s::TestSignal, i::Int64) = s.signal_t[i]
Base.setindex!(s::TestSignal, v, i::Int64) = (s.signal_t[i] = v)

struct LazySignal <: Signal
    n::Int64
    k::Int64
    locs::Array{Int64,1}
    strengths::Array{Float64,1}
    noise_sd::Float64
    signal_t::Dict

    function LazySignal(n::Int64, locs::Array{Int64,1}, strengths::Array{Float64,1} = ones(length(locs)), noise_sd::Float64 = 0.0)
        N = 2^n
        k = locs |> length |> log2 |> ceil |> Int64
        return new(n, k, locs, strengths, noise_sd, Dict{Int64,Float64}())
    end
end

Base.size(s::LazySignal) = length(s.signal_t)

function Base.getindex(s::LazySignal, i::Int)
    if haskey(s.signal_t, i)
        return s.signal_t[i]
    else
        m = dec_to_bin(i-1, s.n)
        sgns = (-1) .^ [m ⋅ dec_to_bin(l, s.n) for l in s.locs]
        v = s.strengths ⋅ sgns + rand(Normal(0.0, s.noise_sd))
        s.signal_t[i] = v
        return v
    end
end

function Base.getindex(s::Signal, r::AbstractUnitRange)
    return map(i -> s[i], r)
end

function Base.getindex(s::Signal, a::Array{<: Int})
    return map(i -> s[i], a)
end

function Base.getindex(s::Signal, a::Array{Array{<: Int}})
    return map(i -> s[i], a)
end

function Base.setindex!(s::LazySignal, v, i::Int)
    @warn("LazySignal is set based on transformed modes: explicitly setting a value is not supported.")
end

function Base.show(io::IO, s::LazySignal)
    println("Lazy signal with locations at ", s.locs)
end

struct InputSignal <: Signal
    n::Int64
    signal_t::Array{Float64,1}
    noise_sd::Float64

    function InputSignal(signal_t::Array{Float64,1}, noise_sd::Float64)
        n = length(signal_t) |> log2 |> ceil |> Int64
        new(n, signal_t, noise_sd)
    end
end

Base.size(s::InputSignal) = size(s.signal_t)
Base.getindex(s::InputSignal, i::Int) = s.signal_t[i]
Base.setindex!(s::InputSignal, v, i::Int) = (s.signal_t[i] = v)

function get_random_sparse_signal(n::Int, K::Int, σ::Float64, minpower::Float64, maxpower::Float64; lazy::Bool=true)
    locs = sample(0:(2^n-1), K, replace=false)
    strengths = Float64.(rand(Uniform(minpower, maxpower), K)) .* (-1) .^ rand(Bool, K)
    if lazy
        return LazySignal(n, locs, strengths, σ)
    else
        return TestSignal(n, locs, strengths, σ)
    end
end

function get_random_delta_sparse_signal(n::Int, σ::Float64, delta::Float64, c::Int, minpower::Float64, maxpower::Float64)
    k = c * n ^ delta |> ceil |> Int64
    return get_random_sparse_signal(n, k, σ, minpower, maxpower)
end


# functions relating to more general input to the generalized bunching conjecture
# instead of pure state inputs

"""

    schur_matrix(H)

computes the Schur matrix as defined in Eq. 1 of Linear Algebra and its Applications 490 (2016) 196–201
"""
function schur_matrix(H)

    @test is_positive_semidefinite(H)

    n = size(H,1)
    schur_mat = Matrix{eltype(H)}(undef,(factorial(n), factorial(n)))

    @showprogress for (i,sigma) in enumerate(permutations(collect(1:n)))
        for (j,rho) in enumerate(permutations(collect(1:n)))
            schur_mat[i,j] = prod([H[sigma[k], rho[k]] for k in 1:n])
        end
    end

    schur_mat

end

function find_permutation_index(this_perm, permutation_array)
    findfirst(x->x == this_perm, permutation_array)
end

function multiply_permutations(a,b)
    a[b[:]]
end

"""

    J_array(theta, n)

returns the `J` as defined in in Eq.10 of [Universality of Generalized Bunching and Efficient Assessment of Boson Sampling](https://arxiv.org/abs/1509.01561), with the permutations coming in the order given by `permutations(collect(1:n))`
"""
function J_array(theta, n)

    J = zeros(eltype(theta), factorial(n))

    permutation_array = collect(permutations(collect(1:n)))

    @showprogress for (i,sigma) in enumerate(permutations(collect(1:n)))
        for (j,tau) in enumerate(permutations(collect(1:n)))
            J[i] += conj(theta[j]) * theta[find_permutation_index(multiply_permutations(tau, sigma), permutation_array)]
        end
    end

    J

end

"""
    density_matrix_from_J(J,n)

density matrix associated to a J function as defined in [Universality of Generalized Bunching and Efficient Assessment of Boson Sampling](https://arxiv.org/abs/1509.01561), computed through Eq. 46.
"""
function density_matrix_from_J(J,n)

    density_matrix = zeros(eltype(J), factorial(n), factorial(n))
    permutation_array = collect(permutations(collect(1:n)))

    @showprogress for (i,sigma) in enumerate(permutations(collect(1:n)))
        for (j,tau) in enumerate(permutations(collect(1:n)))
            density_matrix[j,find_permutation_index(multiply_permutations(sigma, tau), permutation_array)] += J[i]
        end
    end

    density_matrix/factorial(n)

end


using ForwardDiff
using Rotations
using StaticArrays
using LinearAlgebra

using ConstrainedDynamics
using ConstrainedDynamics: vrotate, Lmat, Vᵀmat, Lᵀmat, VLᵀmat, ∂g∂pos, ∂g∂vel, skew

function transfunc3pos(vars)
    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    pa = vars[15:17]
    pb = vars[18:20]

    (xb + vrotate(SVector{3}(pb...), qb)) - (xa + vrotate(SVector{3}(pa), qa))
end

function transfunc3vel(vars)
    Δt = 0.01

    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    va = vars[15:17]
    wa = vars[18:20]
    vb = vars[21:23]
    wb = vars[24:26]

    pa = vars[27:29]
    pb = vars[30:32]

    xa = xa + Δt * va
    xb = xb + Δt * vb
    qa = Quaternion(Δt / 2 * (qa * Quaternion(SVector(sqrt((2 / Δt)^2 - wa' * wa), wa...))))
    qb = Quaternion(Δt / 2 * (qb * Quaternion(SVector(sqrt((2 / Δt)^2 - wb' * wb), wb...))))

    (xb + vrotate(SVector{3}(pb...), qb)) - (xa + vrotate(SVector{3}(pa...), qa))
end

function transfunc2pos(vars)
    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    pa = vars[15:17]
    pb = vars[18:20]

    V1 = vars[21:23]
    V2 = vars[24:26]
    V12 = [V1';V2']

    V12 * vrotate(SVector{3}(((xb + vrotate(SVector{3}(pb...), qb)) - (xa + vrotate(SVector{3}(pa...), qa)))...), inv(qa))
end

function transfunc2vel(vars)
    Δt = 0.01

    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    va = vars[15:17]
    wa = vars[18:20]
    vb = vars[21:23]
    wb = vars[24:26]

    pa = vars[27:29]
    pb = vars[30:32]

    V1 = vars[33:35]
    V2 = vars[36:38]
    V12 = [V1';V2']

    xa = xa + Δt * va
    xb = xb + Δt * vb
    qa = Quaternion(Δt / 2 * (qa * Quaternion(SVector(sqrt((2 / Δt)^2 - wa' * wa), wa...))))
    qb = Quaternion(Δt / 2 * (qb * Quaternion(SVector(sqrt((2 / Δt)^2 - wb' * wb), wb...))))

    V12 * vrotate(SVector{3}(((xb + vrotate(SVector{3}(pb...), qb)) - (xa + vrotate(SVector{3}(pa...), qa)))...), inv(qa))
end

function transfunc1pos(vars)
    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    pa = vars[15:17]
    pb = vars[18:20]

    v = vars[21:23]

    v' * vrotate(SVector{3}(((xb + vrotate(SVector{3}(pb...), qb)) - (xa + vrotate(SVector{3}(pa...), qa)))...), inv(qa))
end

function transfunc1vel(vars)
    Δt = 0.01

    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    va = vars[15:17]
    wa = vars[18:20]
    vb = vars[21:23]
    wb = vars[24:26]

    pa = vars[27:29]
    pb = vars[30:32]

    v = vars[33:35]

    xa = xa + Δt * va
    xb = xb + Δt * vb
    qa = Quaternion(Δt / 2 * (qa * Quaternion(SVector(sqrt((2 / Δt)^2 - wa' * wa), wa...))))
    qb = Quaternion(Δt / 2 * (qb * Quaternion(SVector(sqrt((2 / Δt)^2 - wb' * wb), wb...))))

    v' * vrotate(SVector{3}(((xb + vrotate(SVector{3}(pb...), qb)) - (xa + vrotate(SVector{3}(pa...), qa)))...), inv(qa))
end

function rotfunc3pos(vars)
    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    offset = Quaternion(SVector(vars[15:18]...))

    VLᵀmat(offset) * Lᵀmat(qa) * qb
end

function rotfunc3vel(vars)
    Δt = 0.01

    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    va = vars[15:17]
    wa = vars[18:20]
    vb = vars[21:23]
    wb = vars[24:26]

    offset = Quaternion(SVector(vars[27:30]...))

    xa = xa + Δt * va
    xb = xb + Δt * vb
    qa = Quaternion(Δt / 2 * (qa * Quaternion(SVector(sqrt((2 / Δt)^2 - wa' * wa), wa...))))
    qb = Quaternion(Δt / 2 * (qb * Quaternion(SVector(sqrt((2 / Δt)^2 - wb' * wb), wb...))))

    VLᵀmat(offset) * Lᵀmat(qa) * qb
end

function rotfunc2pos(vars)
    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    offset = Quaternion(SVector(vars[15:18]...))

    V1 = vars[19:21]
    V2 = vars[22:24]
    V12 = [V1';V2']

    V12 * (VLᵀmat(offset) * Lᵀmat(qa) * qb)
end

function rotfunc2vel(vars)
    Δt = 0.01

    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    va = vars[15:17]
    wa = vars[18:20]
    vb = vars[21:23]
    wb = vars[24:26]

    offset = Quaternion(SVector(vars[27:30]...))

    V1 = vars[31:33]
    V2 = vars[34:36]
    V12 = [V1';V2']

    xa = xa + Δt * va
    xb = xb + Δt * vb
    qa = Quaternion(Δt / 2 * (qa * Quaternion(SVector(sqrt((2 / Δt)^2 - wa' * wa), wa...))))
    qb = Quaternion(Δt / 2 * (qb * Quaternion(SVector(sqrt((2 / Δt)^2 - wb' * wb), wb...))))

    V12 * (VLᵀmat(offset) * Lᵀmat(qa) * qb)
end

function rotfunc1pos(vars)
    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    offset = Quaternion(SVector(vars[15:18]...))

    V3 = vars[19:21]

    V3' * (VLᵀmat(offset) * Lᵀmat(qa) * qb)
end

function rotfunc1vel(vars)
    Δt = 0.01

    xa = vars[1:3]
    qa = Quaternion(SVector(vars[4:7]...))
    xb = vars[8:10]
    qb = Quaternion(SVector(vars[11:14]...))

    va = vars[15:17]
    wa = vars[18:20]
    vb = vars[21:23]
    wb = vars[24:26]

    offset = Quaternion(SVector(vars[27:30]...))

    V3 = vars[31:33]

    xa = xa + Δt * va
    xb = xb + Δt * vb
    qa = Quaternion(Δt / 2 * (qa * Quaternion(SVector(sqrt((2 / Δt)^2 - wa' * wa), wa...))))
    qb = Quaternion(Δt / 2 * (qb * Quaternion(SVector(sqrt((2 / Δt)^2 - wb' * wb), wb...))))

    V3' * (VLᵀmat(offset) * Lᵀmat(qa) * qb)
end

function transtest3()
    xa = rand(3)
    qa = Quaternion(rand(RotMatrix{3}))
    xb = rand(3)
    qb = Quaternion(rand(RotMatrix{3}))

    va = rand(3)
    wa = rand(3)
    vb = rand(3)
    wb = rand(3)

    pa = rand(3)
    pb = rand(3)

    origin = Origin{Float64}()
    link1 = Body(Box(1., 1., 1., 1.))
    link2 = Body(Box(1., 1., 1., 1.))

    oc1 = EqualityConstraint(OriginConnection(origin, link1))
    oc2 = EqualityConstraint(OriginConnection(origin, link2))
    joint1 = EqualityConstraint(ConstrainedDynamics.Translational3(link1, link2, pa, pb))

    bot = Mechanism(origin, [link1;link2], [oc1;oc2;joint1])

    setPosition!(bot, link1, x = xa, q = qa)
    setPosition!(bot, link2, x = xb, q = qb)
    link1.s1 = SVector([va;wa]...)
    link2.s1 = SVector([vb;wb]...)

    res = ForwardDiff.jacobian(transfunc3pos, [xa;qa;xb;qb;pa;pb])
    X1 = res[1:3,1:3]
    Q1 = res[1:3,4:7] * Lmat(Quaternion(qa)) * Vᵀmat()
    X2 = res[1:3,8:10]
    Q2 = res[1:3,11:14] * Lmat(Quaternion(qb)) * Vᵀmat()

    n11 = norm(X1 - ∂g∂pos(joint1, 1, bot)[1:3,1:3])
    n12 = norm(Q1 - ∂g∂pos(joint1, 1, bot)[1:3,4:6])
    n21 = norm(X2 - ∂g∂pos(joint1, 2, bot)[1:3,1:3])
    n22 = norm(Q2 - ∂g∂pos(joint1, 2, bot)[1:3,4:6])

    res = ForwardDiff.jacobian(transfunc3vel, [xa;qa;xb;qb;va;wa;vb;wb;pa;pb])
    V1 = res[1:3,15:17]
    W1 = res[1:3,18:20]
    V2 = res[1:3,21:23]
    W2 = res[1:3,24:26]

    n31 = norm(V1 - ∂g∂vel(joint1, 1, bot)[1:3,1:3])
    n32 = norm(W1 - ∂g∂vel(joint1, 1, bot)[1:3,4:6])
    n41 = norm(V2 - ∂g∂vel(joint1, 2, bot)[1:3,1:3])
    n42 = norm(W2 - ∂g∂vel(joint1, 2, bot)[1:3,4:6])

    display((n11, n12, n21, n22))
    display((n31, n32, n41, n42))
    return
end

function transtest2()
    xa = rand(3)
    qa = Quaternion(rand(RotMatrix{3}))
    xb = rand(3)
    qb = Quaternion(rand(RotMatrix{3}))

    va = rand(3)
    wa = rand(3)
    vb = rand(3)
    wb = rand(3)

    pa = rand(3)
    pb = rand(3)

    v = rand(3)

    origin = Origin{Float64}()
    link1 = Body(Box(1., 1., 1., 1.))
    link2 = Body(Box(1., 1., 1., 1.))

    oc1 = EqualityConstraint(OriginConnection(origin, link1))
    oc2 = EqualityConstraint(OriginConnection(origin, link2))
    joint1 = EqualityConstraint(ConstrainedDynamics.Translational2(link1, link2, pa, pb, v))
    V12 = joint1.constraints[1].V12

    bot = Mechanism(origin, [link1;link2], [oc1;oc2;joint1])

    setPosition!(bot, link1, x = xa, q = qa)
    setPosition!(bot, link2, x = xb, q = qb)
    link1.s1 = SVector([va;wa]...)
    link2.s1 = SVector([vb;wb]...)

    res = ForwardDiff.jacobian(transfunc2pos, [xa;qa;xb;qb;pa;pb;V12[1,:];V12[2,:]])
    X1 = res[1:2,1:3]
    Q1 = res[1:2,4:7] * Lmat(Quaternion(qa)) * Vᵀmat()
    X2 = res[1:2,8:10]
    Q2 = res[1:2,11:14] * Lmat(Quaternion(qb)) * Vᵀmat()

    n11 = norm(X1 - ∂g∂pos(joint1, 1, bot)[1:2,1:3])
    n12 = norm(Q1 - ∂g∂pos(joint1, 1, bot)[1:2,4:6])
    n21 = norm(X2 - ∂g∂pos(joint1, 2, bot)[1:2,1:3])
    n22 = norm(Q2 - ∂g∂pos(joint1, 2, bot)[1:2,4:6])

    res = ForwardDiff.jacobian(transfunc2vel, [xa;qa;xb;qb;va;wa;vb;wb;pa;pb;V12[1,:];V12[2,:]])
    V1 = res[1:2,15:17]
    W1 = res[1:2,18:20]
    V2 = res[1:2,21:23]
    W2 = res[1:2,24:26]

    n31 = norm(V1 - ∂g∂vel(joint1, 1, bot)[1:2,1:3])
    n32 = norm(W1 - ∂g∂vel(joint1, 1, bot)[1:2,4:6])
    n41 = norm(V2 - ∂g∂vel(joint1, 2, bot)[1:2,1:3])
    n42 = norm(W2 - ∂g∂vel(joint1, 2, bot)[1:2,4:6])

    display((n11, n12, n21, n22))
    display((n31, n32, n41, n42))
    return
end

function transtest1()
    xa = rand(3)
    qa = Quaternion(rand(RotMatrix{3}))
    xb = rand(3)
    qb = Quaternion(rand(RotMatrix{3}))

    va = rand(3)
    wa = rand(3)
    vb = rand(3)
    wb = rand(3)

    pa = rand(3)
    pb = rand(3)

    v = rand(3)

    origin = Origin{Float64}()
    link1 = Body(Box(1., 1., 1., 1.))
    link2 = Body(Box(1., 1., 1., 1.))

    oc1 = EqualityConstraint(OriginConnection(origin, link1))
    oc2 = EqualityConstraint(OriginConnection(origin, link2))
    joint1 = EqualityConstraint(ConstrainedDynamics.Translational1(link1, link2, pa, pb, v))
    v = joint1.constraints[1].V3'

    bot = Mechanism(origin, [link1;link2], [oc1;oc2;joint1])

    setPosition!(bot, link1, x = xa, q = qa)
    setPosition!(bot, link2, x = xb, q = qb)
    link1.s1 = SVector([va;wa]...)
    link2.s1 = SVector([vb;wb]...)

    res = ForwardDiff.gradient(transfunc1pos, [xa;qa;xb;qb;pa;pb;v])
    X1 = res[1:3]'
    Q1 = res[4:7]' * Lmat(Quaternion(qa)) * Vᵀmat()
    X2 = res[8:10]'
    Q2 = res[11:14]' * Lmat(Quaternion(qb)) * Vᵀmat()

    n11 = norm(X1 - ∂g∂pos(joint1, 1, bot)[1:3]')
    n12 = norm(Q1 - ∂g∂pos(joint1, 1, bot)[4:6]')
    n21 = norm(X2 - ∂g∂pos(joint1, 2, bot)[1:3]')
    n22 = norm(Q2 - ∂g∂pos(joint1, 2, bot)[4:6]')

    res = ForwardDiff.gradient(transfunc1vel, [xa;qa;xb;qb;va;wa;vb;wb;pa;pb;v])
    V1 = res[15:17]'
    W1 = res[18:20]'
    V2 = res[21:23]'
    W2 = res[24:26]'

    n31 = norm(V1 - ∂g∂vel(joint1, 1, bot)[1:3]')
    n32 = norm(W1 - ∂g∂vel(joint1, 1, bot)[4:6]')
    n41 = norm(V2 - ∂g∂vel(joint1, 2, bot)[1:3]')
    n42 = norm(W2 - ∂g∂vel(joint1, 2, bot)[4:6]')

    display((n11, n12, n21, n22))
    display((n31, n32, n41, n42))
    return
end

function rottest3()
    xa = rand(3)
    qa = Quaternion(rand(RotMatrix{3}))
    xb = rand(3)
    qb = Quaternion(rand(RotMatrix{3}))

    va = rand(3)
    wa = rand(3)
    vb = rand(3)
    wb = rand(3)

    offset = Quaternion(rand(RotMatrix{3}))

    origin = Origin{Float64}()
    link1 = Body(Box(1., 1., 1., 1.))
    link2 = Body(Box(1., 1., 1., 1.))

    oc1 = EqualityConstraint(OriginConnection(origin, link1))
    oc2 = EqualityConstraint(OriginConnection(origin, link2))
    joint1 = EqualityConstraint(ConstrainedDynamics.Rotational3(link1, link2, offset = offset))

    bot = Mechanism(origin, [link1;link2], [oc1;oc2;joint1])

    setPosition!(bot, link1, x = xa, q = qa)
    setPosition!(bot, link2, x = xb, q = qb)
    link1.s1 = SVector([va;wa]...)
    link2.s1 = SVector([vb;wb]...)

    res = ForwardDiff.jacobian(rotfunc3pos, [xa;qa;xb;qb;offset])
    X1 = res[1:3,1:3]
    Q1 = res[1:3,4:7] * Lmat(Quaternion(qa)) * Vᵀmat()
    X2 = res[1:3,8:10]
    Q2 = res[1:3,11:14] * Lmat(Quaternion(qb)) * Vᵀmat()

    n11 = norm(X1 - ∂g∂pos(joint1, 1, bot)[1:3,1:3])
    n12 = norm(Q1 - ∂g∂pos(joint1, 1, bot)[1:3,4:6])
    n21 = norm(X2 - ∂g∂pos(joint1, 2, bot)[1:3,1:3])
    n22 = norm(Q2 - ∂g∂pos(joint1, 2, bot)[1:3,4:6])

    res = ForwardDiff.jacobian(rotfunc3vel, [xa;qa;xb;qb;va;wa;vb;wb;offset])
    V1 = res[1:3,15:17]
    W1 = res[1:3,18:20]
    V2 = res[1:3,21:23]
    W2 = res[1:3,24:26]

    n31 = norm(V1 - ∂g∂vel(joint1, 1, bot)[1:3,1:3])
    n32 = norm(W1 - ∂g∂vel(joint1, 1, bot)[1:3,4:6])
    n41 = norm(V2 - ∂g∂vel(joint1, 2, bot)[1:3,1:3])
    n42 = norm(W2 - ∂g∂vel(joint1, 2, bot)[1:3,4:6])

    display((n11, n12, n21, n22))
    display((n31, n32, n41, n42))
    return
end

function rottest2()
    xa = rand(3)
    qa = Quaternion(rand(RotMatrix{3}))
    xb = rand(3)
    qb = Quaternion(rand(RotMatrix{3}))

    va = rand(3)
    wa = rand(3)
    vb = rand(3)
    wb = rand(3)

    offset = Quaternion(rand(RotMatrix{3}))

    v = rand(3)

    origin = Origin{Float64}()
    link1 = Body(Box(1., 1., 1., 1.))
    link2 = Body(Box(1., 1., 1., 1.))

    oc1 = EqualityConstraint(OriginConnection(origin, link1))
    oc2 = EqualityConstraint(OriginConnection(origin, link2))
    joint1 = EqualityConstraint(ConstrainedDynamics.Rotational2(link1, link2, v, offset = offset))
    V12 = joint1.constraints[1].V12

    bot = Mechanism(origin, [link1;link2], [oc1;oc2;joint1])

    setPosition!(bot, link1, x = xa, q = qa)
    setPosition!(bot, link2, x = xb, q = qb)
    link1.s1 = SVector([va;wa]...)
    link2.s1 = SVector([vb;wb]...)

    res = ForwardDiff.jacobian(rotfunc2pos, [xa;qa;xb;qb;offset;V12[1,:];V12[2,:]])
    X1 = res[1:2,1:3]
    Q1 = res[1:2,4:7] * Lmat(Quaternion(qa)) * Vᵀmat()
    X2 = res[1:2,8:10]
    Q2 = res[1:2,11:14] * Lmat(Quaternion(qb)) * Vᵀmat()

    n11 = norm(X1 - ∂g∂pos(joint1, 1, bot)[1:2,1:3])
    n12 = norm(Q1 - ∂g∂pos(joint1, 1, bot)[1:2,4:6])
    n21 = norm(X2 - ∂g∂pos(joint1, 2, bot)[1:2,1:3])
    n22 = norm(Q2 - ∂g∂pos(joint1, 2, bot)[1:2,4:6])

    res = ForwardDiff.jacobian(rotfunc2vel, [xa;qa;xb;qb;va;wa;vb;wb;offset;V12[1,:];V12[2,:]])
    V1 = res[1:2,15:17]
    W1 = res[1:2,18:20]
    V2 = res[1:2,21:23]
    W2 = res[1:2,24:26]

    n31 = norm(V1 - ∂g∂vel(joint1, 1, bot)[1:2,1:3])
    n32 = norm(W1 - ∂g∂vel(joint1, 1, bot)[1:2,4:6])
    n41 = norm(V2 - ∂g∂vel(joint1, 2, bot)[1:2,1:3])
    n42 = norm(W2 - ∂g∂vel(joint1, 2, bot)[1:2,4:6])

    display((n11, n12, n21, n22))
    display((n31, n32, n41, n42))
    return
end

function rottest1()
    xa = rand(3)
    qa = Quaternion(rand(RotMatrix{3}))
    xb = rand(3)
    qb = Quaternion(rand(RotMatrix{3}))

    va = rand(3)
    wa = rand(3)
    vb = rand(3)
    wb = rand(3)

    offset = Quaternion(rand(RotMatrix{3}))

    v = rand(3)

    origin = Origin{Float64}()
    link1 = Body(Box(1., 1., 1., 1.))
    link2 = Body(Box(1., 1., 1., 1.))

    oc1 = EqualityConstraint(OriginConnection(origin, link1))
    oc2 = EqualityConstraint(OriginConnection(origin, link2))
    joint1 = EqualityConstraint(ConstrainedDynamics.Rotational1(link1, link2, v, offset = offset))
    V3 = joint1.constraints[1].V3'

    bot = Mechanism(origin, [link1;link2], [oc1;oc2;joint1])

    setPosition!(bot, link1, x = xa, q = qa)
    setPosition!(bot, link2, x = xb, q = qb)
    link1.s1 = SVector([va;wa]...)
    link2.s1 = SVector([vb;wb]...)

    res = ForwardDiff.gradient(rotfunc1pos, [xa;qa;xb;qb;offset;V3])
    X1 = res[1:3]'
    Q1 = res[4:7]' * Lmat(Quaternion(qa)) * Vᵀmat()
    X2 = res[8:10]'
    Q2 = res[11:14]' * Lmat(Quaternion(qb)) * Vᵀmat()

    n11 = norm(X1 - ∂g∂pos(joint1, 1, bot)[1:3]')
    n12 = norm(Q1 - ∂g∂pos(joint1, 1, bot)[4:6]')
    n21 = norm(X2 - ∂g∂pos(joint1, 2, bot)[1:3]')
    n22 = norm(Q2 - ∂g∂pos(joint1, 2, bot)[4:6]')

    res = ForwardDiff.gradient(rotfunc1vel, [xa;qa;xb;qb;va;wa;vb;wb;offset;V3])
    V1 = res[15:17]'
    W1 = res[18:20]'
    V2 = res[21:23]'
    W2 = res[24:26]'

    n31 = norm(V1 - ∂g∂vel(joint1, 1, bot)[1:3]')
    n32 = norm(W1 - ∂g∂vel(joint1, 1, bot)[4:6]')
    n41 = norm(V2 - ∂g∂vel(joint1, 2, bot)[1:3]')
    n42 = norm(W2 - ∂g∂vel(joint1, 2, bot)[4:6]')

    display((n11, n12, n21, n22))
    display((n31, n32, n41, n42))
    return
end


# ExpSmoothing.jl - Tipo para computar suavizamiento exponencial de medidas de inflación

"""
    InflationExpSmoothing{F <: InflationFunction} <: InflationFunction
Función de inflación que computa la inflación con decaimiento exponencial con el
parámetro ```\\alpha`` (el parámetro del decaimiento ) de la trayectoria
interanual dada por la medida de inflación `inflfn` almacenada.

## Ejemplo

Para computar la inflación con decaimiento exponencial de la variación
interanual del IPC:
```julia-repl
julia> inflfn = InflationExpSmoothing(InflationTotalCPI(), 0.8)
(::InflationExpSmoothing{InflationTotalCPI}) (generic function with 5 methods)
```
"""
struct InflationExpSmoothing{F <: InflationFunction} <: InflationFunction
    inflfn::F
    alpha::Float64
end

# 2. Extender el método de nombre
measure_name(esfn::InflationExpSmoothing) = "Suavizamiento exponencial con parámetro $(round(esfn.alpha, digits=4)) parámetros de " * measure_name(esfn.inflfn)

# Función de parámetros
CPIDataBase.params(esfn::InflationExpSmoothing) = (esfn.alpha, )

# Método de conveniencia para definir sobre parámetro de suavizamiento Enteros
InflationExpSmoothing(inflfn::InflationFunction,alpha::Int) =
    InflationExpSmoothing(inflfn::InflationFunction, convert(Float64,alpha))

# Método que opera sobre CountryStructure: computa la trayectoria de inflación
# con la función inflfn y luego computa el promedio con suavizamiento exponencial con el parámetro
# alpha (el parámetro del decaimiento )
function (esfn::InflationExpSmoothing)(cs::CountryStructure)

    # Cómputo usual de inflación
    tray_infl = esfn.inflfn(cs)

    # Algoritmo de promedio con suavizamiento exponencial
    k = esfn.alpha
    k == 1 && return tray_infl
    k == 0 && return ones(length(tray_infl))*tray_infl[1]

    # Computar el promedio con suavizamiento exponencial (in-place)
    es_tray_infl = smoothing_exponential(tray_infl, k)

    # Devolver la trayectoria de inflación
    es_tray_infl
end

# Actualmente, se define el resumen intermensual por el obtenido con la función
# de inflación interna
# function (mafn::InflationExpSmoothing)(base::VarCPIBase)
#     # Obtener resumen intermensual y aplicarle el suavizamiento exponencial
#     vm = mafn.inflfn(base)
#     vm
# end

# Función para computar el promedio con suavizamiento exponencial de la serie `v` con `k`
# como valor de suavizamiento que pondera la historia 0<=k<=1.
# Esta función modifica los elementos de `v` con el promedio de suavizmiento exponencial
# con alpha dada.
# El tipo de suavizamiento que se utiliza es SES (Simple Exponential Smoothing, o
# Suavizamiento Exponencial Simple en español)

function smoothing_exponential(v, k)
    es = similar(v)
    n = length(es)
    es[1] = v[1]
    for j = 2:(n)
        es[j] = es[j-1] + (v[j] - es[(j-1)])*k
    end
    es
end


"Abstract type that deals with global boundary-boundary matrices"
abstract BBSolver

type BBMatrix{T<:Number} <: BBSolver
    nb::Int
    nbslv::Int
    A::Array{T,2}
    fact::Base.LinAlg.LU{T,Array{T,2}}
    function BBMatrix(dof::DofMap)
        nb = nbmodes(dof)
        nbslv = nbslvmodes(dof)
        new(nb, nbslv, zeros(T, nbslv, nbslv))
    end
end

function assemble!{T<:Number}(Ag::BBMatrix{T}, Ae, m)

    nm = length(m)
    n1 = Ag.nbslv+1
    for i = 1:nm
        ig = m[i]
        if ig < n1
            for k = 1:nm
                kg = m[k]
                if kg < n1
                    Ag.A[kg,ig] += Ae[k,i]
                end
            end
        end
    end
end

function trf!(Ag::BBMatrix)
    Ag.fact = lufact!(Ag.A)
end
trs!(Ag::BBMatrix, x) = A_ldiv_B!(Ag.fact, x)

"Global boundary-boundary Symmetric Positive-Definite matrix"
abstract BBSym <: BBSolver

type BBSymMatrix{T<:Number} <: BBSym
    nb::Int
    nbslv::Int
    A::Array{T,2}
    fact::Base.LinAlg.Cholesky{T,Array{T,2}}
    function BBSymMatrix(dof::DofMap)
        nb = nbmodes(dof)
        nbslv = nbslvmodes(dof)
        new(nb, nbslv, zeros(T, nbslv, nbslv))
    end

end

function assemble!{T<:Number}(Ag::BBSymMatrix{T}, Ae, m)

    nm = length(m)
    n1 = Ag.nbslv+1
    for i = 1:nm
        ig = m[i]
        if ig < n1
            for k = 1:nm
                kg = m[k]
                if kg < n1
                    Ag.A[kg,ig] += Ae[k,i]
                end
            end
        end
    end
end

function trf!(Ag::BBSymMatrix)
    Ag.fact = cholfact!(Ag.A)
end
trs!(Ag::BBSymMatrix, x) = A_ldiv_B!(Ag.fact, x)

"Global boundary-boundary symmetric tridiagonal matrices"
type BBSymTri{T<:Number} <: BBSym
    "Number of boundary modes"
    nb::Int
    "Number of boundary modes that should be solved"
    nbslv::Int
    "Diagonal elements of the matrix"
    D::Array{T,1}
    "Sub-diagonal elements"
    E::Array{T,1}

    "Tridiagonal Matrix"
    tri::SymTridiagonal{T}
    "Factorization of symmetric tridiagonal matrix"
    fact:: LDLt{T,SymTridiagonal{T}}

    function BBSymTri(dof::DofMap)
        nb = nbmodes(dof)
        nbslv = nbslvmodes(dof)
        D = zeros(T,nbslv)
        E = zeros(T, max(nbslv-1,0))
        tri = SymTridiagonal(D, E)
        new(nb, nbslv, D, E, tri)
    end
end

"LU (or Choleksy) decomposition of boundary-boundary system"
function trf!(Ag::BBSymTri)
    Ag.fact = ldltfact!(Ag.tri)
end

"Solving linear system using LU (or Cholesky) decomposition"
trs!(Ag::BBSymTri, x) = A_ldiv_B!(Ag.fact, x)

"""
Assembles the global boundary-boundary matrix.
"""
function assemble!{T}(Ag::BBSymTri{T}, Ae, m)
    np1 = Ag.nbslv + 1
    D = Ag.D
    E = Ag.E

    ig = m[1]
    kg = m[2]

    if ig < np1
        D[ig] += Ae[1,1]
        if kg < np1
            D[kg] += Ae[2,2]
            E[min(ig,kg)] += Ae[1,2]
        end
    elseif kg < np1
        D[kg] += Ae[2,2]
    end

end

"Global boundary-boundary tridiagonal matrices"
type BBTri{T<:Number} <: BBSolver
    "Number of boundary modes"
    nb::Int
    "Number of boundary modes that should be solved"
    nbslv::Int
    "Diagonal elements of the matrix"
    D::Array{T,1}
    "Lower sub-diagonal elements"
    Dl::Array{T,1}
    "Upper sub-diagonal elements"
    Du::Array{T,1}
    "Tridiagonal object"
    tri::Tridiagonal{T}
    fact::Base.LinAlg.LU{T,Tridiagonal{T}}
    """
    Inner constructor. First the matrix must be assembled.
    """
    function BBTri(dof::DofMap)
        nb = nbmodes(dof)
        nbslv = nbslvmodes(dof)
        D = zeros(T,nbslv)
        Dl = zeros(T, max(nbslv-1, 0))
        Du = zeros(T, max(nbslv-1, 0))
        tri = Tridiagonal(Dl, D, Du)
        new(nb, nbslv, D, Dl, Du, tri)
    end
end

function assemble!{T}(Ag::BBTri{T}, Ae, m)
    np1 = Ag.nbslv + 1
    D = Ag.D
    Dl = Ag.Dl
    Du = Ag.Du

    ig = m[1]
    kg = m[2]

    if ig < np1
        D[ig] += Ae[1,1]
        if kg < np1
            D[kg] += Ae[2,2]
            if ig < kg
                Du[ig] += Ae[1,2]
                Dl[ig] += Ae[2,1]
            elseif kg < ig
                Du[kg] = Ae[2,1]
                Dl[kg] = Ae[1,2]
            end
        end
    elseif kg < np1
        D[kg] += Ae[2,2]
    end

end

function trf!(Ag::BBTri)
    Ag.fact = lufact!(Ag.tri)
    return
end
trs!(Ag::BBTri, x) = A_ldiv_B!(Ag.fact, x)

type BBTriP{T<:Number} <: BBSolver
    "Number of independent boundary modes"
    nb::Int
    "Number of boundary modes that should be solved"
    nbslv::Int
    "Lower diagonal"
    Dl::Vector{T}
    "Main Diagonal"
    D::Vector{T}
    "Upper diagonal"
    Du::Vector{T}
    "Auxiliary memory"
    x2::Vector{T}
    an1::T
    b1::T
    bn1::T
    cn::T
    cn1::T
    tri::Tridiagonal{T}
    fact::LinAlg.LU{T,Tridiagonal{T}}

    function BBTriP(dof::DofMap)
        nb = nbmodes(dof)
        nbslv = nbslvmodes(dof)

        nbslv = nb-1
        D = zeros(T,nbslv)
        Dl = zeros(T,nbslv-1)
        Du = zeros(T,nbslv-1)
        x2 = zeros(T,nbslv)
        tri = Tridiagonal(Dl, D, Du)
        new(nb, nbslv, Dl, D, Du, x2, zero(T), zero(T), zero(T), zero(T), zero(T), tri)

    end

end

function assemble!{T}(Ag::BBTriP{T}, Ae, m)

    m1 = m[1]
    m2 = m[2]

    nb = Ag.nb
    nbslv = Ag.nb-1

    ig = m1
    kg = m2

    D = Ag.D
    Dl = Ag.Dl
    Du = Ag.Du

    if m2 < nb && m2 != 1 #Not the last element
        np1 = Ag.nbslv + 1

        if ig < np1
            D[ig] += Ae[1,1]
            if kg < np1
                D[kg] += Ae[2,2]
                if ig < kg
                    Du[ig] += Ae[1,2]
                    Dl[ig] += Ae[2,1]
                elseif kg < ig
                    Du[kg] = Ae[2,1]
                    Dl[kg] = Ae[1,2]
                end
            end
        elseif kg < np1
            D[kg] += Ae[2,2]
        end
    elseif m2 == nb # It is the next to last
        Ag.an1   += Ae[2,2]
        Ag.cn    += Ae[1,2]
        Ag.bn1   += Ae[2,1]
        D[nbslv] += Ae[1,1]
    elseif m2 == 1
        Ag.b1    += Ae[2,1]
        Ag.cn1   += Ae[1,2]
        D[1]     += Ae[2,2]
        Ag.an1   += Ae[1,1]
    end

end

function trf!(Ag::BBTriP)
    Ag.fact = lufact!(Ag.tri)
    return
end
function trs!(Ag::BBTriP, x)
    n = Ag.nb
    n1 = n-1
    qn1 = view(x, 1:n1, :)
    nrhs = size(x,2)
    A_ldiv_B!(Ag.fact, qn1)
    x2 = Ag.x2
    for i = 1:nrhs
        fill!(x2, 0)
        x2[1] = -Ag.b1
        x2[n1] = -Ag.cn
        A_ldiv_B!(Ag.fact, x2)
        x[n,i] = (x[n,i] - Ag.cn1*qn1[1,i] - Ag.bn1*qn1[n1,i]) /
            (Ag.an1 + Ag.cn1*x2[1] + Ag.bn1*x2[n1])
        for k = 1:n1
            x[k,i] += x[n,i]*x2[k]
        end
    end
    x

end

type BBSymTriP{T<:Number} <: BBSym
    "Number of independent boundary modes"
    nb::Int
    "Number of boundary modes that should be solved"
    nbslv::Int
    "Main Diagonal"
    D::Vector{T}
    "Upper diagonal"
    Du::Vector{T}
    "Auxiliary memory"
    x2::Vector{T}
    an1::T
    cn::T
    cn1::T
    "Tridiagonal Matrix"
    tri::SymTridiagonal{T}
    "Factorization of symmetric tridiagonal matrix"
    fact:: LDLt{T,SymTridiagonal{T}}

    function BBSymTriP(dof::DofMap)
        nb = nbmodes(dof)
        nbslv = nbslvmodes(dof)

        nbslv = nb-1
        D = zeros(T,nb-1)
        Du = zeros(T,nb-2)
        x2 = zeros(T,nb-1)
        tri = SymTridiagonal(D, Du)
        new(nb, nbslv, D, Du, x2, zero(T), zero(T), zero(T), tri)
    end
end

function assemble!{T<:Number}(Ag::BBSymTriP{T}, Ae, m)

    m1 = m[1]
    m2 = m[2]

    nb = Ag.nb
    nbslv = Ag.nb-1

    ig = m1
    kg = m2

    D = Ag.D
    Du = Ag.Du

    if m2 < nb && m2 != 1 # Not the last element (if periodic)

        np1 = Ag.nbslv + 1

        if ig < np1
            D[ig] += Ae[1,1]
            if kg < np1
                D[kg] += Ae[2,2]
                if ig < kg
                    Du[ig] += Ae[1,2]
                elseif kg < ig
                    Du[kg] = Ae[2,1]
                end
            end
        elseif kg < np1
            D[kg] += Ae[2,2]
        end
    elseif m2 == nb # It is  the next to last
        Ag.an1   += Ae[2,2]
        Ag.cn    += Ae[1,2]
        D[nbslv] += Ae[1,1]
    elseif m2 == 1
        Ag.cn1   += Ae[1,2]
        D[1]     += Ae[2,2]
        Ag.an1   += Ae[1,1]
    end

end

function trf!(Ag::BBSymTriP)
    Ag.fact = ldltfact!(Ag.tri)
    return
end
function trs!(Ag::BBSymTriP, x)
    n = Ag.nb
    n1 = n-1
    qn1 = view(x, 1:n1, :)
    nrhs = size(x,2)
    A_ldiv_B!(Ag.fact, qn1)
    x2 = Ag.x2
    for i = 1:nrhs
        fill!(x2, 0)
        x2[1] = -Ag.cn1
        x2[n1] = -Ag.cn
        A_ldiv_B!(Ag.fact, x2)
        x[n,i] = (x[n,i] - Ag.cn1*qn1[1,i] - Ag.cn*qn1[n1,i]) /
            (Ag.an1 + Ag.cn1*x2[1] + Ag.cn*x2[n1])
        for k = 1:n1
            x[k,i] += x[n,i]*x2[k]
        end
    end
    x
end

@enum MatType FULLMAT TRIMAT TRIPMAT

"""
Generic matrix interface for 1d problems.

If the problem is very small, a full matrix will be used.
If it is larger and periodic, a BBTriP will be used.
Else a BBTri is used.
"""
type BBMatrix1d{T<:Number} <: BBSolver
    nb::Int
    nbslv::Int
    mtype::MatType
    periodic::Bool
    mat::BBMatrix{T}
    tri::BBTri{T}
    trip::BBTriP{T}
    #BBMatrix1d(nb, nbslv, mtype, periodic) = new(nb, nbslv, mtype, periodic)
    function BBMatrix1d(dof::DofMap)
        nb = nbmodes(dof)
        nbslv = nbslvmodes(dof)
        periodic = false
        if dof.bmap[1,1] == dof.bmap[2,end]
            periodic = true
        end

        if nbslv <= 3
            mtype = FULLMAT
        elseif periodic == false
            mtype = TRIMAT
        else
            mtype = TRIPMAT
        end

        bb = new(nb, nbslv, mtype, periodic)

        if mtype == FULLMAT
            bb.mat = BBMatrix{T}(dof)
        elseif mtype == TRIMAT
            bb.tri = BBTri{T}(dof)
        else
            bb.trip = BBTriP{T}(dof)
        end

        return bb
    end
end

"""
Generic symmetric matrix interface for 1d problems.

If the problem is very small, a full matrix will be used.
If it is larger and periodic, a BBSymTriP will be used.
Else a BBSymTri is used.
"""
type BBSymMatrix1d{T<:Number} <: BBSolver
    nb::Int
    nbslv::Int
    mtype::MatType
    periodic::Bool
    mat::BBSymMatrix{T}
    tri::BBSymTri{T}
    trip::BBSymTriP{T}
    function BBSymMatrix1d(dof::DofMap)
        nb = nbmodes(dof)
        nbslv = nbslvmodes(dof)
        periodic = false
        if dof.bmap[1,1] == dof.bmap[2,end]
            periodic = true
        end

        if nbslv <= 3
            mtype = FULLMAT
        elseif periodic == false
            mtype = TRIMAT
        else
            mtype = TRIPMAT
        end

        bb = new(nb, nbslv, mtype, periodic)

        if mtype == FULLMAT
            bb.mat = BBSymMatrix{T}(dof)
        elseif mtype == TRIMAT
            bb.tri = BBSymTri{T}(dof)
        else
            bb.trip = BBSymTriP{T}(dof)
        end

        return bb
    end
end

function assemble!{T<:Number}(Ag::Union{BBMatrix1d{T},BBSymMatrix1d{T}}, Ae, m)
    mt = Ag.mtype
    if mt == FULLMAT
        assemble!(Ag.mat, Ae, m)
    elseif mt == TRIMAT
        assemble!(Ag.tri, Ae, m)
    else
        assemble!(Ag.trip, Ae, m)
    end

end

function trf!(Ag::Union{BBMatrix1d,BBSymMatrix1d})
    mt = Ag.mtype
    if mt == FULLMAT
        trf!(Ag.mat)
    elseif mt == TRIMAT
        trf!(Ag.tri)
    else
        trf!(Ag.trip)
    end
end

function trs!(Ag::Union{BBMatrix1d,BBSymMatrix1d}, x)
    mt = Ag.mtype
    if mt == FULLMAT
        trs!(Ag.mat,x)
    elseif mt == TRIMAT
        trs!(Ag.tri,x)
    else
        trs!(Ag.trip,x)
    end

end


module splines

    export filter_coeffs, interpolant_cspline
    export eval_UC_spline, eval_UC_spline_G, eval_UC_multi_spline, eval_UC_multi_spline!
    include("csplines.jl")
    include("splines_filter.jl")

function interpolant_cspline(a, b, orders, V)

    coefs = filter_coeffs(a, b, orders, V)

    function fun(s::Array{Float64,2})
        return eval_UC_spline(a, b, orders, coefs, s)
    end

    function fun(p::Float64...)
        return fun([p...]')
    end

    return fun

end

end


# HellwarthTheory.jl

#####
"""
	HellwarthBScheme(LO)

    Multiple phonon mode reduction to a single effective frequency.
	Hellwarth et al. 1999 PRB, 'B scheme'; the athermal method.
    Averaging procedure is constructed by considering the average effect of the action of multiple branches.

	Follows Eqn (58) in this paper, assuming typo on LHS, should actually be W_e.
"""
function HellwarthBScheme(LO)
    println("Hellwarth B Scheme... (athermal)")
    H58 = sum( (LO[:,2].^2)./ LO[:,1].^2 )
    println("Hellwarth (58) summation: ",H58)

    H59 = sum( LO[:,2].^2 ) # sum of total ir activity squarred
    println("Hellwarth (59) summation (total ir activity ^2): ", H59)
    println("Hellwarth (59) W_e (total ir activity ): ", sqrt(H59))

    omega = sqrt(H59 / H58)
    println("Hellwarth (61) Omega (freq): ",omega)

	return(omega)
end

# More complex scheme, involving thermodynamic Beta
# Hellwarth(50), RHS
"""
	HellwarthAScheme(LO,T=295)

    Multiple phonon mode reduction to a single effective frequency.
	Temperature dependent, defaults to T=295 K.

    Follows Hellwarth et al. 1999 PRB 'A' scheme, Eqn 50 RHS.

	UNTESTED AND UNCERTAIN CODE.
"""
function HellwarthAScheme(LO; T=295)
    println("Hellwarth A scheme...T=$T K")
    omega=LO[:,1] # No unit conversion
    #omega=LO[:,1].*2*pi*1E12 # THz --> Hz
    #omega=LO[:,1].*2*pi*0.02998*1E12 # cm^-1 --> Hz
    β=(omega.*ħ)/(kB*T) #assuming units SI
    H50 = sum( ((LO[:,2].^2).*coth.(β))./omega )
    println("Hellwarth (50) summation: ",H50)

    H51= sum( LO[:,2].^2 ) # sum of total ir activity squarred
    println("Hellwarth (51) summation (total ir activity ^2): ", H51)
    println("Hellwarth (51) W_e (total ir activity ): ", sqrt(H51))

    # OK; so this is deriving Omega / coth(Beta/2)
    omegacoth=H51/H50
    println("omegacoth: ",omegacoth)

    solnfreq=0.0 #required for Julia 0.5 so it realises variable still required for return.
    # Very primitive manner to decouple Omega from both sides of the eqn.
	# Should really rewrite as a bisection (at least!)
    maxfreq=maximum(omega)
    for freq in 0.1:maxfreq/2500:maxfreq
        pseudo_omega=omegacoth/coth(freq * ħ/(2*kB*T))
        if freq>pseudo_omega
            println("freq: $freq pseudo-omega: $pseudo_omega")
            println("freq: ",freq/(2*pi*0.02998*1E12), " pseudo-omega: ",pseudo_omega/(2*pi*0.02998*1E12))
            solnfreq=freq
            break
        end
    end
	return solnfreq
end


@register Base.getindex(x,i::Integer)
@register Base.getindex(x,i)
@register Base.binomial(n,k)

@register Base.signbit(x)
ModelingToolkit.derivative(::typeof(signbit), args::NTuple{1,Any}, ::Val{1}) = 0
ModelingToolkit.derivative(::typeof(abs), args::NTuple{1,Any}, ::Val{1}) = IfElse.ifelse(signbit(args[1]),-one(args[1]),one(args[1]))
function ModelingToolkit.derivative(::typeof(min), args::NTuple{2,Any}, ::Val{1})
    x, y = args
    IfElse.ifelse(x < y, one(x), zero(x))
end
function ModelingToolkit.derivative(::typeof(min), args::NTuple{2,Any}, ::Val{2})
    x, y = args
    IfElse.ifelse(x < y, zero(y), one(y))
end
function ModelingToolkit.derivative(::typeof(max), args::NTuple{2,Any}, ::Val{1})
    x, y = args
    IfElse.ifelse(x > y, one(x), zero(x))
end
function ModelingToolkit.derivative(::typeof(max), args::NTuple{2,Any}, ::Val{2})
    x, y = args
    IfElse.ifelse(x > y, zero(y), one(y))
end

@register IfElse.ifelse(x,y,z::Any)
ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{1}) = 0
ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{2}) = IfElse.ifelse(args[1],1,0)
ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{3}) = IfElse.ifelse(args[1],0,1)

ModelingToolkit.@register Base.rand(x)
ModelingToolkit.@register Base.randn(x)

ModelingToolkit.@register Distributions.pdf(dist,x)
ModelingToolkit.@register Distributions.logpdf(dist,x)
ModelingToolkit.@register Distributions.cdf(dist,x)
ModelingToolkit.@register Distributions.logcdf(dist,x)
ModelingToolkit.@register Distributions.quantile(dist,x)

ModelingToolkit.@register Distributions.Uniform(mu,sigma)
ModelingToolkit.@register Distributions.Normal(mu,sigma)

@register ∈(x::Num, y::AbstractArray)
@register ∪(x, y)
@register ∩(x, y)
@register ∨(x, y)
@register ∧(x, y)
@register ⊆(x, y)

function LinearAlgebra.det(A::AbstractMatrix{<:Num}; laplace=true)
    if laplace
        n = LinearAlgebra.checksquare(A)
        if n == 1
            return A[1, 1]
        elseif n == 2
            return A[1, 1] * A[2, 2] - A[1, 2] * A[2, 1]
        else
            temp = 0
            # Laplace expansion along the first column
            M′ = A[:, 2:end]
            for i in axes(A, 1)
                M = M′[(1:n) .!= i, :]
                d′ = A[i, 1] * det(M)
                if iseven(i)
                    temp = iszero(temp) ? d′ : temp - d′
                else
                    temp = iszero(temp) ? d′ : temp + d′
                end
            end
        end
        return temp
    else
        if istriu(A) || istril(A)
            return det(UpperTriangular(A))
        end
        return det(lu(A; check = false))
    end
end


"""
Coordinate transform function `ctparent` convert the coordinates of `point = (x1, y1)` on a child frame `c1` to its corresponding coordinates `(x0, y0)` on a parent frame `c0`.

For example,
saying you have a circle centered at `(x1, y1)` of the coordinate system `c1`, where `c1` is located on `contextdim` of `c0`:
```julia
c0 = context();
c1 = context(contextdim...;units=UnitBox(0,0,2,2));
circle1 = circle(x1,y1,0.1);
compose(c0, (c1, circle1));
```
Using the transformed coordinates `(x0,y0)=ctparent((x1,y1),contextdim;UnitBox=(0,0,2,2))` we can draw the same circle as `circle1` by `compose(c0, circle(x0,y0,0.1))`.

Broadcasting example:
```julia
ctparent.([(x1,y1), (x2,y2)], [contextdim]; UnitBox=(0,0,2,2))
```
"""
function ctparent(point::NTuple{2,<:Real},contextdim::NTuple{4,<:Real};UnitBox::NTuple{4,<:Real}=(0,0,1,1))
    x1, y1 = point;
    xstart1a, ystart1a, width1a, height1a = contextdim;
    xstart1b, ystart1b, width1b, height1b = UnitBox;
    Δx1 = x1-xstart1b;
    Δy1 = y1-ystart1b;
    cfactorx = width1a/1; # default context()'s width/height is 1
    cfactory = height1a/1;
    ubfactorx = 1/width1b; # default UnitBox()'s width/height is 1
    ubfactory = 1/height1b;
    x0 = xstart1a + Δx1*cfactorx*ubfactorx;
    y0 = ystart1a + Δy1*cfactory*ubfactory;
    return (x0, y0)
end

"""
`ctchildwidth(width0, contextdim::NTuple{4,<:Real};UnitBox::NTuple{4,<:Real}=(0,0,1,1))` convert the width in `c0` to the length it should be in `c1` that the line:
```
c0 = context();
c1 = context(contextdim...;units=UnitBox(0,0,2,2));
width1=ctchildwidth(width0, contextdim;UnitBox=(0,0,2,2));
compose(c0, (c1, line([(0,0),(width1,0)])) )
```
has the length looks visually the same as:
```
compose(c0, line([(0,0),(width0,0)]) )
```
"""
function ctchildwidth(width0, contextdim::NTuple{4,<:Real};UnitBox::NTuple{4,<:Real}=(0,0,1,1))
	# width0/contextdim[3] = width1/UnitBox[3]
	return width0*UnitBox[3]/contextdim[3] # width1
end

"""
`height1 = ctchildheight(height0, contextdim::NTuple{4,<:Real};UnitBox::NTuple{4,<:Real}=(0,0,1,1))` works the same as `ctchildwidth()` but rescale in y direction. See `ctchildwidth`'s documentation.
"""
function ctchildheight(height0, contextdim::NTuple{4,<:Real};UnitBox::NTuple{4,<:Real}=(0,0,1,1))
	# height0/contextdim[4] = height1/UnitBox[4]
	return height0*UnitBox[4]/contextdim[4] # height1
end

dfthw = 0.015; # default head width
dfthh = 0.025; # default head height
"""
`myarrow(startpoint, arrangle, arrlength; headwidth=$dfthw, headheight=$dfthh) |> polygon`
"""
function myarrow(startpoint, arrangle, arrlength; headwidth=dfthw, headheight=dfthh)
	rot(θ,point) = (point[1]*cos(θ) - point[2]*sin(θ), point[1]*sin(θ) + point[2]*cos(θ));
	xo,yo = (0.0,0.0);
	head_x =xo + arrlength;
	headpoint = (head_x, yo);
	sidepoint_a = headpoint.+(-headheight,yo-headwidth/2);
	sidepoint_b = headpoint.+(-headheight,yo+headwidth/2);
	function lim_x(sidepoint_a)
		sidepoint_a = maximum.(([sidepoint_a[1], xo], [sidepoint_a[2]]))
	end
	sidepoint_a = lim_x(sidepoint_a);
	sidepoint_b = lim_x(sidepoint_b);

	eq_triangle = [sidepoint_a,
				   sidepoint_b,
				   headpoint]

	tailwidth = headwidth/5;
	tail_a = (xo, yo+tailwidth);
	tail_b = (xo, yo-tailwidth);
	tail_c = (sidepoint_b[1], yo-tailwidth); # sidepoint_a[1] should equal sidepoint_b[1]
	tail_d = (sidepoint_a[1], yo+tailwidth);

	pts = rot.([arrangle], [tail_a, tail_b, tail_c, sidepoint_a, headpoint, sidepoint_b, tail_d]);
	xs = [pt[1] for pt in pts] .+ startpoint[1];
	ys = [pt[2] for pt in pts] .+ startpoint[2];

	return tuple.(xs, ys)
end

"""
`myarrow(startpoint, arrangle, arrlength, contextdim::NTuple{4,<:Real}; UnitBox::NTuple{4,<:Real}=(0,0,1,1))` is the same as `myarrow(startpoint, arrangle, arrlength; headwidth=$dfthw, headheight=$dfthh)` but the arrows `headwidth` and `headheight` will be converted to make sure their size is consistent.
"""
function myarrow(startpoint, arrangle, arrlength, contextdim::NTuple{4,<:Real}; UnitBox::NTuple{4,<:Real}=(0,0,1,1), headwidth=dfthw, headheight=dfthh)
	hh = ctchildwidth(headheight, contextdim; UnitBox=UnitBox); # headheight is "along x axis" at zero degree, so we have to use `ctchildwidth` to covert.
	hw = ctchildheight(headwidth, contextdim; UnitBox=UnitBox);
	return myarrow(startpoint, arrangle, arrlength; headwidth=hw, headheight=hh)
end

"""
`composelabels(xstr, ystr, contextdim)` add x and y labels on the parent frame of `contextdim`.

```julia
composelabels(labels::Vector{<:AbstractString}, contextdim::NTuple{4,<:Real}; labelloc::Vector{Symbol}=[:bottom,:left], labeldir::Vector{Symbol}=[:h,:h], labelcolor=("black","black"))
```
- supported `labelloc`: `:bottom`,`:top`,`:left`,`:right`.
- `labeldir=:h` means the label is horizontally placed; `labeldir=:h`, vertically placed.

# Example:
```julia
contextdim = (0.1,0.4,0.8,0.55); # the position of `cplot1`
cplot1 = (context(contextdim..., units=UnitBox(UB1...)),
blablabla...); # a plot perhaps

cplot1withlabels = (context(), composelabels(["time","velocity"],contextdim;labelloc=[:left,:top], labeldir=[:v,:h])..., cplot1);
compose(context(), cplot1withlabels)
```
"""
function composelabels(labels::Vector{<:AbstractString}, contextdim::NTuple{4,<:Real}; labelloc::Vector{Symbol}=[:bottom,:left], labeldir::Vector{Symbol}=[:h,:h], labelcolor=("black","black"))
    labh = 0.03; # labelheight
    textpos = (0.5,0.5); # of created label context
    (pad1x, pad1y, width1, height1) = contextdim;

    lenl = length(labels);

    if length(labelloc) != lenl || length(labeldir) != lenl
        error("`labelloc`/`labeldir` has to be the same length as `labels` for specifying labels' location/orientation.");
    end

    contexts = [];
    for i = 1:lenl
        textrot = ();
        textalign = ();
        # txtpt = (1.0,0.5);
        # (pad1x, pad1y, width1, height1) = contextdim;
        # labh = 0.03;
        #
        # leftlabel = (a1,pad1y,pad1x-a1,height1);
        # cplot1a = (context(), (context(leftlabel...),Compose.text(txtpt...,"hello world",hright,vcenter,Rotation(-1/2*pi,0.5,0.5)) ,(context(), rectangle(), stroke("pink"),fillopacity(0.3)))
        # ,cplot1);

        if labelloc[i] == :left
            a1 = maximum([0.0, pad1x-labh]);
            contextdim_txt = (a1,pad1y,pad1x-a1,height1);
            textpos = (0.9,0.5);
            textalign= (hright, vcenter);
        end

        if labelloc[i] == :bottom
            a1 = minimum([1.0-pad1y-height1, labh]);
            contextdim_txt = (pad1x, pad1y+height1, width1, a1);
            textpos = (0.5, 0.1);
            textalign= (hcenter,vtop);
        end

        if labelloc[i] == :top
            a1 = maximum([0.0, pad1y-labh]);
            contextdim_txt = (pad1x, a1, width1, pad1y-a1);
            textpos = (0.5,0.9);
            textalign= (hcenter,vbottom);
        end

        if labelloc[i] == :right
            a1 = minimum([1.0-pad1x-width, labh]);
            contextdim_txt = (pad1x+width1, pad1y, a1, height1);
            textpos = (0.1, 0.5);
            textalign= (hleft,vcenter);
        end

        if labeldir[i] == :v
            textrot = (Rotation(-pi/2,0.5,0.5),);
            textpos = (0.5, 0.5);
            textalign= (hcenter,vcenter);
        end
        # convert dictionary to named tuple
        # TODO: write this technique to Notes
        # opts = NamedTuple{Tuple(keys(contextopt_d))}(values(contextopt_d));

        push!(contexts,
              (context(contextdim_txt...),
               Compose.text(textpos..., labels[i], textalign..., textrot...),fill(labelcolor[i])),
        );

    end
    return contexts
end

"""
`pos2points(xstart, ystart, width, height)` convert the position to a vector of points `[(xstart, ystart), (xstart, ystart+height), (xstart+width, ystart+height), (xstart+width,ystart)]`.
"""
function pos2points(xstart, ystart, width, height)
    return [(xstart, ystart), (xstart, ystart+height), (xstart+width, ystart+height), (xstart+width,ystart)]
end

function pos2points(pos::NTuple{4,<:Real})
    pos2points(pos...);
end

"""
`pointshift(point::NTuple{2,<:Real}, shift2::NTuple{2,<:Real})` returns `(point[1]+shift2[1], point[2]+shift2[2])`.
"""
function pointshift(point::NTuple{2,<:Real}, shift2::NTuple{2,<:Real})
    return (point[1]+shift2[1], point[2]+shift2[2])
end


using ArgParse

s=ArgParseSettings()
@add_arg_table s begin
    "--part2"
        help="Do part 2"
        action=:store_true
    "--test"
        help="Run test data"
        action=:store_true
end
parsed_args = parse_args(ARGS, s)

function parseLine(line)
    parse(Int32, line)
end

function growcounts!(counts)
    n = counts[1]
    for i in 1:8
        counts[i] = counts[i+1]
    end
    counts[7] += n
    counts[9] = n
end

test_suffix = parsed_args["test"] ? "_test" : ""
lines = readlines("input/day_6$test_suffix.txt")
fish = [parse(Int8, f) for f in split(lines[1], ",")]
census = zeros(UInt64, 9)
function countfish(age)
    ageindex = age + 1
    census[ageindex] += 1
end
foreach(countfish, fish)
for _ in 1:(parsed_args["part2"] ? 256 : 80)
    growcounts!(census)
end
answer = sum(census)
println("The answer is: $answer")


"""
```jldoctest
julia> parse(Rational{Int}, "1/4")
1//4
julia> parse(Rational, "1//4")
1//4
```
"""
function Base.parse(::Type{Rational{T}}, x::AbstractString) where {T <: Integer}
    ms, ns = split(x, '/', keepempty = false)
    m = parse(T, ms)
    n = parse(T, ns)
    return m // n
end

Base.parse(::Type{Rational}, x::AbstractString) = parse(Rational{Int32}, x)

"""
    get_ratio(x::AbstractString)

Parse x as either a Float64 or a Rational and return a Float64 ratio.

```jldoctest
julia> get_ratio("1/4")
0.25
julia> get_ratio("0.5")
0.5
```
"""
function get_ratio(x::AbstractString)
    r = tryparse(Float64, x)
    return r === nothing ? convert(Float64, parse(Rational, x)) : r
end

"""
    get_ratio(x::Float64)

For Float64, simply returns x.

```jldoctest
julia> get_ratio(0.8)
0.8
```
"""
get_ratio(x::Float64) = x