using DiffEqSensitivity, OrdinaryDiffEq, ForwardDiff, Calculus
using Test
function fb(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -t*p[3]*u[2] + t*u[1]*u[2]
end
function jac(J,u,p,t)
  (x, y, a, b, c) = (u[1], u[2], p[1], p[2], p[3])
  J[1,1] = a + y * b * -1
  J[2,1] = t * y
  J[1,2] = b * x * -1
  J[2,2] = t * c * -1 + t * x
end
function paramjac(pJ,u,p,t)
  (x, y, a, b, c) = (u[1], u[2], p[1], p[2], p[3])
  pJ[1,1] = x
  pJ[2,1] = 0.0
  pJ[1,2] = - x * y
  pJ[2,2] = 0.0
  pJ[1,3] = 0.0
  pJ[2,3] = - t * y
end

f = ODEFunction(fb,jac=jac,paramjac=paramjac)
p = [1.5,1.0,3.0]
prob = ODEForwardSensitivityProblem(f,[1.0;1.0],(0.0,10.0),p)
probInpl = ODEForwardSensitivityProblem(fb,[1.0;1.0],(0.0,10.0),p)
probnoad = ODEForwardSensitivityProblem(fb,[1.0;1.0],(0.0,10.0),p,
                                        ForwardSensitivity(autodiff=false))
probnoad2 = ODEForwardSensitivityProblem(f,[1.0;1.0],(0.0,10.0),p,
                                        ForwardSensitivity(autodiff=false))
probvecmat = ODEForwardSensitivityProblem(fb,[1.0;1.0],(0.0,10.0),p,
                                          ForwardSensitivity(autojacvec=false,autojacmat=true))
sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14)
@test_broken solve(probInpl,KenCarp4(),abstol=1e-14,reltol=1e-14).retcode == :Success
solInpl = solve(probInpl,KenCarp4(autodiff=false),abstol=1e-14,reltol=1e-14)
solInpl2 = solve(probInpl,Rodas4(autodiff=false),abstol=1e-10,reltol=1e-10)
solnoad = solve(probnoad,KenCarp4(autodiff=false),abstol=1e-14,reltol=1e-14)
solnoad2 = solve(probnoad,KenCarp4(autodiff=false),abstol=1e-14,reltol=1e-14)
solvecmat = solve(probvecmat,Tsit5(),abstol=1e-14,reltol=1e-14)

x = sol[1:sol.prob.f.numindvar,:]

@test sol(5.0) ≈ solnoad(5.0)
@test sol(5.0) ≈ solnoad2(5.0)
@test sol(5.0) ≈ solInpl(5.0)
@test isapprox(solInpl(5.0), solInpl2(5.0),rtol=1e-5)
@test sol(5.0) ≈ solvecmat(5.0)

# Get the sensitivities

da = sol[sol.prob.f.numindvar+1:sol.prob.f.numindvar*2,:]
db = sol[sol.prob.f.numindvar*2+1:sol.prob.f.numindvar*3,:]
dc = sol[sol.prob.f.numindvar*3+1:sol.prob.f.numindvar*4,:]

sense_res1 = [da[:,end] db[:,end] dc[:,end]]

prob = ODEForwardSensitivityProblem(f.f,[1.0;1.0],(0.0,10.0),p,
                                    ForwardSensitivity(autojacvec=true))
sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14,saveat=0.01)
x = sol[1:sol.prob.f.numindvar,:]

# Get the sensitivities

res = sol[1:sol.prob.f.numindvar,:]
da = sol[sol.prob.f.numindvar+1:sol.prob.f.numindvar*2,:]
db = sol[sol.prob.f.numindvar*2+1:sol.prob.f.numindvar*3,:]
dc = sol[sol.prob.f.numindvar*3+1:sol.prob.f.numindvar*4,:]

sense_res2 = [da[:,end] db[:,end] dc[:,end]]

function test_f(p)
  prob = ODEProblem(f,eltype(p).([1.0,1.0]),(0.0,10.0),p)
  solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14,save_everystep=false)[end]
end

p = [1.5,1.0,3.0]
fd_res = ForwardDiff.jacobian(test_f,p)
calc_res = Calculus.finite_difference_jacobian(test_f,p)

@test sense_res1 ≈ sense_res2 ≈ fd_res
@test sense_res1 ≈ sense_res2 ≈ calc_res

## Check extraction

xall, dpall = extract_local_sensitivities(sol)
@test xall == res
@test dpall[1] == da

_,dpall_matrix = extract_local_sensitivities(sol,Val(true))
@test mapreduce(x->x[:, 2], hcat, dpall) == dpall_matrix[2]

x, dp = extract_local_sensitivities(sol,length(sol.t))
sense_res2 = reduce(hcat,dp)
@test sense_res1 == sense_res2

@test extract_local_sensitivities(sol,sol.t[3]) == extract_local_sensitivities(sol,3)

tmp = similar(sol[1])
@test extract_local_sensitivities(tmp,sol,sol.t[3]) == extract_local_sensitivities(sol,3)

# asmatrix=true
@test extract_local_sensitivities(sol, length(sol), true) == (x, sense_res2)
@test extract_local_sensitivities(sol, sol.t[end], true) == (x, sense_res2)
@test extract_local_sensitivities(tmp, sol, sol.t[end], true) == (x, sense_res2)

# Return type inferred
@inferred extract_local_sensitivities(sol, 1)
@inferred extract_local_sensitivities(sol, 1, Val(true))
@inferred extract_local_sensitivities(sol, sol.t[3])
@inferred extract_local_sensitivities(sol, sol.t[3], Val(true))
@inferred extract_local_sensitivities(tmp, sol, sol.t[3])
@inferred extract_local_sensitivities(tmp, sol, sol.t[3], Val(true))

### ForwardDiff version

prob = ODEForwardSensitivityProblem(f.f,[1.0;1.0],(0.0,10.0),p,
                                    ForwardDiffSensitivity())
sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14,saveat=0.01)

xall, dpall = extract_local_sensitivities(sol)
@test xall ≈ res
@test dpall[1] ≈ da atol=1e-9

_,dpall_matrix = extract_local_sensitivities(sol,Val(true))
@test mapreduce(x->x[:, 2], hcat, dpall) == dpall_matrix[2]

x, dp = extract_local_sensitivities(sol,length(sol.t))
sense_res2 = reduce(hcat,dp)
@test fd_res == sense_res2

@test extract_local_sensitivities(sol,sol.t[3]) == extract_local_sensitivities(sol,3)

tmp = similar(sol[1])
@test extract_local_sensitivities(tmp,sol,sol.t[3]) == extract_local_sensitivities(sol,3)

# asmatrix=true
@test extract_local_sensitivities(sol, length(sol), true) == (x, sense_res2)
@test extract_local_sensitivities(sol, sol.t[end], true) == (x, sense_res2)
@test extract_local_sensitivities(tmp, sol, sol.t[end], true) == (x, sense_res2)

# Return type inferred
@inferred extract_local_sensitivities(sol, 1)
@inferred extract_local_sensitivities(sol, 1, Val(true))
@inferred extract_local_sensitivities(sol, sol.t[3])
@inferred extract_local_sensitivities(sol, sol.t[3], Val(true))
@inferred extract_local_sensitivities(tmp, sol, sol.t[3])
@inferred extract_local_sensitivities(tmp, sol, sol.t[3], Val(true))

# Test Float32

function f32(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + u[1]*u[2]
end
p = [1.5f0,1.0f0,3.0f0]
prob = ODEForwardSensitivityProblem(f32,[1.0f0;1.0f0],(0.0f0,10.0f0),p)
sol = solve(prob,Tsit5())


using Estimation
using Base.Test

function noiseless_bias()
	n = 65536
	N = 100

	result = []
	for trial=1:N
		freq = 0.1 + 0.3*rand()
		sig = [cos(2.0*pi*freq*i) for i=1:n]

		est_freq = Estimation.quinn_fernandes(sig, 2.0*pi*0.1)/(2.0*pi)
		result = [(freq-est_freq),result]
	end

	result
end

function noisy_bias()
	n = 65536
	N = 100

	result = []
	for trial=1:N
		freq = 0.1 + 0.3*rand()
		sig = [cos(2.0*pi*freq*i) for i=1:n] + 0.1*randn(n)

		est_freq = Estimation.quinn_fernandes(sig, 2.0*pi*0.1)/(2.0*pi)
		result = [(freq-est_freq),result]
	end

	result
end


##
## L1 regularization
##
@with_kw mutable struct regularizer_l1 <: Regularizer
    λ::Float64
end

function wholespace_manifold(reg::regularizer_l1, x)
    return l1Manifold(ones(size(x)))
end

## 0th order
g(reg::regularizer_l1, x) = reg.λ * norm(x, 1)

## 1st order
function prox_αg!(reg::regularizer_l1, res, x, α)
    res .= softthresh.(x, reg.λ * α)
    return l1Manifold(map(x -> abs(x) > 0, res))
end
function prox_αg(reg::regularizer_l1, x, α)
    res = softthresh.(x, reg.λ * α)
    return res, l1Manifold(map(x -> abs(x) > 0, res))
end

∇M_g!(reg::regularizer_l1, ::l1Manifold, ∇M_g, x) = (@. ∇M_g = reg.λ * sign.(x))
function ∇M_g(reg::regularizer_l1, M::l1Manifold, x)
    ∇M_g = zeros(representation_size(M))
    return ∇M_g!(reg, M, ∇M_g, x)
end

## 2nd order
∇²M_g_ξ(::regularizer_l1, M::l1Manifold, x, ξ) = zeros(size(ξ))
∇²M_g_ξ!(::regularizer_l1, M::l1Manifold, res, x, ξ) = (res .= 0)

#
### Optimality status
#
function model_g_subgradient!(model, regularizer::regularizer_l1, M, x)
    n = size(x, 1)
    λ = regularizer.λ

    d = manifold_dimension(M)

    @variable(model, ḡ_normal[1:n-d])
    @constraint(model, -λ .<= ḡ_normal .<= λ)

    return ḡ_normal
end

function build_subgradient_from_normalcomp(regularizer::regularizer_l1, M, x, ḡ_normal)
    n = size(x, 1)
    ḡ = Vector(undef, n)
    λ = regularizer.λ
    ind_norm = 1
    for i in 1:n
        if M.nnz_coords[i]
            ḡ[i] = sign(x[i]) * λ
        else
            ḡ[i] = ḡ_normal[ind_norm]
            ind_norm += 1
        end
    end
    return ḡ
end

function build_normalcomp_from_subgradient(::regularizer_l1, M, x, ḡ)
    return ḡ[.!M.nnz_coords]
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.3'
#       jupytext_version: 0.8.6
#   kernelspec:
#     display_name: Julia 1.0.3
#     language: julia
#     name: julia-1.0
# ---

module ExactRecursiveExpressions
using Plots

#Get Final Size for SIR Model
@fastmath function getSIR_final_size(N,β,γ)
    final_size = zeros(N+1,1);
    final_size[2] = 1;
    for Z2 in 0:N
        @inbounds for Z1 in Z2+1:N-1
            p1 = 1 / ( 1 + γ/(β*(N-Z1)));
            final_size[Z1+2] = final_size[Z1+2] + final_size[Z1+1]*p1;
            final_size[Z1+1] = final_size[Z1+1]*(1-p1);
        end
    end
    return final_size;
end

#Sample Size
N = 20;
#Infected per-Capita Rate
β = 2/(N-1);
#Recover per-capita rate
γ = 1.0;

#Get Final Size for SIR Model
final_size = getSIR_final_size(N,β,γ);

#Get Time Final Size Function
@time final_size = getSIR_final_size(N,β,γ);

bar(0:N,final_size)

@fastmath function getSI4R_final_size(N::Int64,β::Float64,γ::Float64)
    Psi = (N+1)*(N+2)*(N+3)*(N+4)/24;
    p_vec = zeros(Float64,Int(Psi),1);
    p_vec[2]=1;
    @inbounds for Z5 in 0:N
        w::Int64 = Psi - (N-Z5+1)*(N-Z5+2)*(N-Z5+3)*(N-Z5+4)/24 + 1;
        @inbounds for Z4 in Z5:N
            a5 = 4*γ*(Z4-Z5);
            @inbounds for Z3 in Z4:N
                a4 = 4*γ*(Z3-Z4);
                @inbounds for Z2 in Z3:N;
                    a3 = 4*γ*(Z2-Z3);
                    @inbounds for Z1 in Z2:N
                    a1 = β*(N-Z1)*(Z1-Z5);
                    a2 = 4*γ*(Z1-Z2);
                    tot = a1+a2+a3+a4+a5;
                    if Z1-Z5 == 0
                        final_size[Z5+1] = p_vec[w];
                    end
                    if a1 > 0
                       p_vec[w+1] = p_vec[w+1]+ p_vec[w]*a1/tot;
                    end
                    if a2 > 0
                        p_vec[w+N-Z2] = p_vec[w+N-Z2]+ p_vec[w]*a2/tot;
                    end
                    if a3 > 0
                        place3::Int64 = (N-Z3)*(N-Z3+1)/2;
                        p_vec[w+place3] = p_vec[w+place3]+ p_vec[w]*a3/tot;
                    end
                    if a4 > 0
                        place4::Int64 = (N-Z4)*(N-Z4+1)*(N-Z4+2)/6;
                        p_vec[w+place4] = p_vec[w+place4] + p_vec[w]*a4/tot;
                    end
                    if a5 > 0
                        p_vec[w] = p_vec[w]*a5/tot;
                    end
                    w = w + 1;
                end
            end
        end
    end
end
    return final_size
end

#Get Final Size for SI(4)R model
final_size = getSI4R_final_size(N,β,γ);

#Get Run time of SI(4)R Final Size Model
@time final_size = getSI4R_final_size(N,β,γ);

bar(0:N,final_size)

end


using Distributions
using LinearAlgebra
using Random
using StatsFuns: softplus
using Roots # for inverse

################################################################################
#                            Planar and Radial Flows                           #
#             Ref: Variational Inference with Normalizing Flows,               #
#               D. Rezende, S. Mohamed(2015) arXiv:1505.05770                  #
################################################################################

###############
# PlanarLayer #
###############

mutable struct PlanarLayer{T1,T2} <: Bijector{1}
    w::T1
    u::T1
    b::T2
end

function get_u_hat(u, w)
    # To preserve invertibility
    return (
        u + (planar_flow_m(w' * u) - w' * u)[1]
        * w / (norm(w[:,1],2) ^ 2)
    )   # from A.1
end

function PlanarLayer(dims::Int, container=Array)
    w = container(randn(dims, 1))
    u = container(randn(dims, 1))
    b = container(randn(1))
    return PlanarLayer(w, u, b)
end

planar_flow_m(x) = -1 .+ softplus.(x)   # for planar flow from A.1
dtanh(x) = 1 .- (tanh.(x)) .^ 2         # for planar flow
ψ(z, w, b) = dtanh(w' * z .+ b) .* w    # for planar flow from eq(11)

# An internal version of transform that returns intermediate variables
function _transform(flow::PlanarLayer, z)
    u_hat = get_u_hat(flow.u, flow.w)
    transformed = z + u_hat * tanh.(flow.w' * z .+ flow.b) # from eq(10)
    return (transformed=transformed, u_hat=u_hat)
end

(b::PlanarLayer)(z) = _transform(b, z).transformed

function _forward(flow::PlanarLayer, z)
    transformed, u_hat = _transform(flow, z)
    # Compute log_det_jacobian
    psi = ψ(z, flow.w, flow.b)
    log_det_jacobian = log.(abs.(1.0 .+ psi' * u_hat))          # from eq(12)
    return (rv=transformed, logabsdetjac=vec(log_det_jacobian)) # from eq(10)
end

forward(flow::PlanarLayer, z::AbstractMatrix) = _forward(flow, z)

function forward(flow::PlanarLayer, z::AbstractVector{<: Real})
    res = _forward(flow, z)
    return (rv=vec(res.rv), logabsdetjac=res.logabsdetjac[1])
end

function (ib::Inversed{<: PlanarLayer})(y::AbstractMatrix{<: Real})
    flow = ib.orig
    u_hat = get_u_hat(flow.u, flow.w)
    # Define the objective functional; implemented with reference from A.1
    f(y) = alpha -> (flow.w' * y)[1] - alpha - (flow.w' * u_hat)[1] * tanh(alpha+flow.b[1])
    # Run solver
    alphas_ = [find_zero(f(y[:,i:i]), 0.0, Order16()) for i in 1:size(y, 2)]
    alphas = alphas_'
    z_para = (flow.w ./ norm(flow.w,2)) * alphas
    z_per = y - z_para - u_hat * tanh.(flow.w' * z_para .+ flow.b)

    return z_para + z_per
end

function (ib::Inversed{<: PlanarLayer})(y::AbstractVector{<: Real})
    return vec(ib(reshape(y, (length(y), 1))))
end

logabsdetjac(flow::PlanarLayer, x) = forward(flow, x).logabsdetjac
isclosedform(b::Inversed{<:PlanarLayer}) = false

###############
# RadialLayer #
###############

# FIXME: using `TrackedArray` for the parameters, we end up with
# nested tracked structures; don't want this.
mutable struct RadialLayer{T1,T2} <: Bijector{1}
    α_::T1
    β::T1
    z_0::T2
end

function RadialLayer(dims::Int, container=Array)
    α_ = container(randn(1))
    β = container(randn(1))
    z_0 = container(randn(dims, 1))
    return RadialLayer(α_, β, z_0)
end

h(α, r) = 1 ./ (α .+ r)     # for radial flow from eq(14)
dh(α, r) = - h(α, r) .^ 2   # for radial flow; derivative of h()

# An internal version of transform that returns intermediate variables
function _transform(flow::RadialLayer, z)
    α = softplus(flow.α_[1])            # from A.2
    β_hat = -α + softplus(flow.β[1])    # from A.2
    r = sqrt.(sum((z .- flow.z_0).^2; dims = 1))
    transformed = z + β_hat .* h(α, r) .* (z .- flow.z_0)   # from eq(14)
    return (transformed=transformed, α=α, β_hat=β_hat, r=r)
end

(b::RadialLayer)(z::AbstractMatrix{<:Real}) = _transform(b, z).transformed
(b::RadialLayer)(z::AbstractVector{<:Real}) = vec(_transform(b, z).transformed)

function _forward(flow::RadialLayer, z)
    transformed, α, β_hat, r = _transform(flow, z)
    # Compute log_det_jacobian
    d = size(flow.z_0, 1)
    h_ = h(α, r)
    log_det_jacobian = @. (
        (d - 1) * log(1.0 + β_hat * h_)
        + log(1.0 +  β_hat * h_ + β_hat * (- h_ ^ 2) * r)
    )   # from eq(14)
    return (rv=transformed, logabsdetjac=vec(log_det_jacobian))
end

forward(flow::RadialLayer, z) = _forward(flow, z)

function forward(flow::RadialLayer, z::AbstractVector{<:Real})
    res = _forward(flow, z)
    return (rv=vec(res.rv), logabsdetjac=res.logabsdetjac[1])
end

function (ib::Inversed{<:RadialLayer})(y)
    flow = ib.orig
    α = softplus(flow.α_[1])            # from A.2
    β_hat = - α + softplus(flow.β[1])   # from A.2
    # Define the objective functional
    f(y) = r -> norm(y - flow.z_0, 2) - r * (1 + β_hat / (α + r))   # from eq(26)
    # Run solver
    rs = [find_zero(f(y[:,i:i]), 0.0, Order16()) for i in 1:size(y, 2)]'    # from A.2
    z_hat = (y .- flow.z_0) ./ (rs .* (1 .+ β_hat ./ (α .+ rs)))            # from eq(25)
    z = flow.z_0 .+ rs .* z_hat # from A.2
    return z
end

function (ib::Inversed{<: RadialLayer})(y::AbstractVector{<:Real})
    return vec(ib(reshape(y, (length(y), 1))))
end

logabsdetjac(flow::RadialLayer, x) = forward(flow, x).logabsdetjac
isclosedform(b::Inversed{<:RadialLayer}) = false


__precompile__(true)

module MinMaxFilter

# This is a port of the Lemire min max filter as implemented by Bruno Luong
# http://arxiv.org/abs/cs.DS/0610046
# http://lemire.me/
# http://www.mathworks.com/matlabcentral/fileexchange/24705-min-max-filter

export minmax_filter

using Base.Cartesian

type Wedge{A <: AbstractArray}
    buffer::A
    size::Int
    n::Int
    first::Int
    last::Int
    mxn::Int
end

for N = 2:4
    @eval begin
    function minmax_filter{T <: Number}(A::Array{T, $N}, window::Array{Int, 1})

        maxval_temp = copy(A); minval_temp = copy(A)

        for dim = 1:$N

            # For all but the last dimension
            @nloops $(N-1) i maxval_temp begin

                # Create index for full array (fa) length
                @nexprs $(N)   j->(fa_{j} = 1:size(maxval_temp)[j])
                @nexprs $(N-1) j->(fa_{j} = i_j)

                # Create index for short array (sa) length
                @nexprs $(N)   j->(sa_{j} = 1:size(maxval_temp)[j] - window[dim] + 1)
                @nexprs $(N-1) j->(sa_{j} = i_j)

                # Filter the last dimension
                (@nref $N minval_temp sa) = min_filter(vec( @nref $N minval_temp fa), window[dim])
                (@nref $N maxval_temp sa) = max_filter(vec( @nref $N maxval_temp fa), window[dim])

            end

            # Circular shift the dimensions
            maxval_temp = permutedims(maxval_temp, mod(collect(1:$N), $N)+1)
            minval_temp = permutedims(minval_temp, mod(collect(1:$N), $N)+1)

        end

        # The dimensions to extract
        @nexprs $N j->(a_{j} = 1:size(A, j)-window[j]+1)

        # Extract set dimensions
        maxval_out = @nref $N maxval_temp a
        minval_out = @nref $N minval_temp a

        return minval_out, maxval_out
    end
    end
end

function minmax_filter{T <: Number}(a::AbstractArray{T}, window::Int)

    n = length(a)

    # Initialise the output variables
    # This is the running minimum and maximum over the specified window length
    minval = zeros(T, 1, n-window+1)
    maxval = zeros(T, 1, n-window+1)

    # Initialise the internal wedges
    # U[1], L[1] are the location of the global maximum and minimum
    # U[2], L[2] are the maximum and minimum over (U1, inf)
    L = Wedge(zeros(Int,1,window+1), window+1, 0, 1, 0, 0)          # Min
    U = Wedge(zeros(Int,1,window+1), window+1, 0, 1, 0, 0)

    for i = 2:n
        if i > window
            if !wedgeisempty(U)
                maxval[i-window] = a[getfirst(U)]
            else
                maxval[i-window] = a[i-1]
            end
            if !wedgeisempty(L)
                minval[i-window] = a[getfirst(L)]
            else
                minval[i-window] = a[i-1]
            end
        end # window

        if a[i] > a[i-1]
            pushback!(L, i-1)
            if i==window+getfirst(L); L=popfront(L); end
            while !wedgeisempty(U)
                if a[i] <= a[getlast(U)]
                    if i == window+getfirst(U); U = popfront(U); end
                    break
                end
                U = popback(U)
            end

        else

            pushback!(U, i-1)
            if i==window+getfirst(U); U=popfront(U); end

            while !wedgeisempty(L)
                if a[i] >= a[getlast(L)]
                    if i == window+getfirst(L); L = popfront(L); end
                    break
                end
                L = popback(L)
            end

        end  # a>a-1

    end # for i

    i = n+1
    if !wedgeisempty(U)
        maxval[i-window] = a[getfirst(U)]
    else
        maxval[i-window] = a[i-1]
    end

    if !wedgeisempty(L)
        minval[i-window] = a[getfirst(L)]
    else
        minval[i-window] = a[i-1]
    end

    return minval, maxval
end

function min_filter(a::AbstractArray, window::Int)

    minval, maxval = minmax_filter(a, window)

    return minval
end

function max_filter(a::AbstractArray, window::Int)

    minval, maxval = minmax_filter(a, window)

    return maxval
end

function wedgeisempty(X::Wedge)
    X.n <= 0
end

function pushback!(X::Wedge, v)
    X.last += 1
    if X.last > X.size
        X.last = 1
    end
    X.buffer[X.last] = v
    X.n = X.n+1
    X.mxn = max(X.mxn, X.n)
end

function getfirst(X::Wedge)
    X.buffer[X.first]
end

function getlast(X::Wedge)
    X.buffer[X.last]
end

function popfront(X::Wedge)
    X.n = X.n-1
    X.first = mod(X.first, X.size) + 1
    return X
end

function popback(X::Wedge)
    X.n = X.n-1
    X.last = mod(X.last-2, X.size) + 1
    return X
end

end # module


module SymbolicML

using CSV
using DataFrames
using Plots
using StatsPlots

include("dataset.jl")
include("functions.jl")

export discrete_derivative, discrete_derivatives
export siemensDataInfo, siemensData

export select_attributes
export abs_energy
export absolute_maximum
export absolute_sum_of_changes
export agg_autocorrelation
export autocorrelation
export benford_correlation
export binned_entropy
export c3
export change_quantiles
export cid_ce
export count_above
export count_above_mean
export count_below
export count_below_mean
export first_location_of_maximum
export first_location_of_minimum
export has_duplicate
export has_duplicate_max
export has_duplicate_min
export large_standard_deviation
export last_location_of_maximum
export last_location_of_minimum
export mean_abs_change
export mean_change
export mean_n_absolute_max
export longest_strike_above_mean
export longest_strike_below_mean
export mean_second_derivative_central
export number_crossing_m
export percentage_of_reoccurring_datapoints_to_all_datapoints
export percentage_of_reoccurring_values_to_all_values
export range_count
export ratio_beyond_r_sigma
export ratio_value_number_to_time_series_length
export root_mean_square
export sample_entropy
export sum_of_reoccurring_data_points

end


@testset "utilities" begin
    @testset "affine/projective" begin
        @test projective([1, 2, 3]) == [1, 1, 2, 3]
        @test affine(projective([1, 2, 3])) == [1, 2, 3]
    end

    @testset "totaldegree" begin
        @PolyImpl.polyvar x y

        F = PolySystem([
            3.0x^2 + 2.0x * y - 1.0 * y^2,
            (-2.5 + 2im)x^2 + -3.0x * y + 4.0 * y^3])

        G, solutions = totaldegree(F)

        @test length(solutions) == 6
        @test degrees(G) == [2, 3]

        for sol in solutions
            @test norm(evaluate(G, sol)) ≈ 0.0 atol = 1e-12
        end

        @test any(x -> abs(1 - x) > 1e-8, norm.(solutions)) == true

        G, solutions = totaldegree(F, unit_roots=true)
        @test all(x -> norm(1 - abs2.(x)) < 1e-8, solutions)

        @test_throws ErrorException totaldegree(PolySystem([x-1, x-2]))
    end

    @testset "randomsystem" begin
        F = randomsystem(Complex64, 2, 3, mindegree=2, maxdegree=5)

        @test nvariables(F) == 3
        @test length(F) == 2
        @test F isa FixedPolySystem.PolySystem{Complex64}

        @test all(d -> 2 ≤ d ≤ 5, degrees(F))

        @test randomsystem(2, 3) isa FixedPolySystem.PolySystem{Complex128}

        G = randomsystem([2, 3, 4], [:x, :y])
        @test nvariables(G) == 2
        @test length(G) == 3
        @test degrees(G) == [2, 3, 4]
    end
end


using ADCME
using LinearAlgebra
using MAT
using AdFem
using PyPlot
using SparseArrays

function Q_exact(x, y)
    # (2*x^2*y*(x - 1)*(2*y - 1) + 1.0*(y^2 + 1)^2*(x*y*(x - 1) + x*y*(y - 1) + x*(x - 1)*(y - 1) + y*(x - 1)*(y - 1)) - (y^2 + 1)*(2*x*(x - 1)*(x + (x^2 + 1)*(y^2 + 1)) + 2*y*(x + (x^2 + 1)*(y^2 + 1))*(y - 1) + y*(2*x - 1)*(y - 1)*(2*x*(y^2 + 1) + 1)))/(y^2 + 1)^2
    # (2*x^2*y*(x - 1)*(2*y - 1) + x*(y^2 + 1)^2*(-y^2*(2*x - 1)*(y - 1)^2 + 1.0*y*(x - 1) + 1.0*(x - 1)*(y - 1)) - (y^2 + 1)*(2*x*(x - 1)*(x + (x^2 + 1)*(y^2 + 1)) + 2*y*(x + (x^2 + 1)*(y^2 + 1))*(y - 1) + y*(2*x - 1)*(y - 1)*(2*x*(y^2 + 1) + 1)))/(y^2 + 1)^2
    (2*x^2*y*(x - 1)*(2*y - 1)*(x + y + 1) - x*y^2*(2*x - 1)*(y - 1)^2*(y^2 + 1)^2*(x + y + 1) + x*(x - 1)*(2*y - 1)*(y^2 + 1)^2 - (y^2 + 1)*(x + y + 1)*(2*x*(x - 1)*(x + (x^2 + 1)*(y^2 + 1)) + 2*y*(x + (x^2 + 1)*(y^2 + 1))*(y - 1) + y*(2*x - 1)*(y - 1)*(2*x*(y^2 + 1) + 1)))/((y^2 + 1)^2*(x + y + 1))
end

function T_exact(x, y)
    x * (1-x) * y * (1-y)
end

function u_exact(x, y)
    x * y * (1-y)
end

function v_exact(x, y)
    1 / (1 + x + y)
end

function k_exact(x, y)
    1 + x^2 + x / (1 + y^2)
end

function k_nn(x, y)
    out = fc([x y], [20,20,20,1])^2 + 0.1 # N x 1
    squeeze(out)
end

#---------------------------------------------
# grid setup
m = 20
n = 20
h = 1/n

#---------------------------------------------
xy = fem_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
T0 = @. T_exact(x, y)
u = @. u_exact(x,y)
v = @. v_exact(x,y)
Q = @. Q_exact(x,y)
k = @. k_nn(x, y); k=stack(k)
# k0 = @. k_nn(xy)

# ---------------------------------------------

bd = bcnode("all", m, n, h)

ugauss = fem_to_gauss_points(u, m, n, h)
vgauss = fem_to_gauss_points(v, m, n, h)
kgauss = fem_to_gauss_points(k, m, n, h)
Qgauss = fem_to_gauss_points(Q, m, n, h)

Advection = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h))
Laplace = compute_fem_laplace_matrix1(kgauss, m, n, h)
A = Advection + Laplace
A, _ = fem_impose_Dirichlet_boundary_condition1(A, bd, m, n, h)
b = constant(compute_fem_source_term1(constant(Qgauss), m, n, h))
b = scatter_update(b, bd, zeros(length(bd)))
T_computed = A\b

T_data = matread("chtenergy_data.mat")["T"]
# loss = mean((T_computed .- T_data)^2)

sample_size = 40
idx = rand(1:(m+1)*(n+1), sample_size)
observed_data = T_data[idx]
loss = mean((T_computed[idx] .- observed_data)^2)

noise = false
noise_level = 0.01
if noise
    noise_ratio = (1 - noise_level) .+ 2 * noise_level * rand(Float64, size(observed_data)) # uniform on (1-noise_level, 1+noise_level)
    observed_data = observed_data .* noise_ratio
end

loss = loss * 1e10
# ---------------------------------------------------
# create a session and run
max_iter = 100
sess = Session(); init(sess)
loss_ = BFGS!(sess, loss, max_iter)
figure(); semilogy(loss_); savefig("chtenergy_loss.png")

figure(figsize=(14,4));
subplot(131)
visualize_scalar_on_fem_points(k_exact.(x,y), m, n, h); title("conductivity exact")
subplot(132)
visualize_scalar_on_fem_points(run(sess, k), m, n, h); title("conductivity prediction")
subplot(133)
visualize_scalar_on_fem_points(k_exact.(x,y).-run(sess, k), m, n, h); title("conductivity difference")
savefig("chtenergy_k.png")


function relative_attractor(F::BoxMap, B::BoxSet{<:AbstractBoxPartition{Box{N,T}}}; steps=12) where {N,T}
    for k = 1:steps
        B = subdivide(B, (k%N)+1)
        B = B ∩ F(B)
    end
    return B
end

function unstable_set!(F::BoxMap, B::BoxSet)
    B_new = B
    while !isempty(B_new)
        B_new = F(B_new)
        setdiff!(B_new, B)
        union!(B, B_new)
    end
    return B
end

function chain_recurrent_set(F::BoxMap, B::BoxSet{<:AbstractBoxPartition{Box{N,T}}}; steps=12) where {N,T}
    for k in 1:steps
        B = subdivide(B, (k%N)+1)
        P = TransferOperator(F, B)
        B = strongly_connected_components(P)
    end
    return B
end

@muladd function adaptive_newton_step(g, g_jacobian, x, k)
    function armijo_rule(g, x, α, σ, ρ)
        Dg = g_jacobian(x)
        d = Dg\g(x)
        while any(g(x + α * d) .> g(x) + σ * α * Dg' * d) && α > 0.1
            α = ρ * α
        end
        return α
    end
    h = armijo_rule(g, x, 1.0, 1e-4, 0.8)

    expon(ϵ, σ, h, δ) = Int(ceil(log(ϵ * (1/2)^σ)/log(maximum((1 - h, δ)))))
    n = expon(0.2, k, h, 0.1)

    for _ in 1:n
        Dg = g_jacobian(x)
        x = x - h * (Dg \ g(x))
    end

    return x
end

function cover_roots(g, Dg, B::BoxSet{<:AbstractBoxPartition{Box{N,T}}}; steps::Int=12) where {N,T}
    domain = B.partition.domain
    for k in 1:steps
        B = subdivide(B, (k%N)+1)
        f = x -> adaptive_newton_step(g, Dg, x, k)
        F_k = BoxMap(f, domain, no_of_points = 40)
        B = F_k(B)
    end
    return B
end

"""
    relative_attractor(boxset::BoxSet, g; T, num_points=20, ϵ=1e-6) -> BoxFun

Compute the Finite Time Lyapunov Exponent for every box in boxset.

Arguments:
* `boxset` is the starting BoxSet.
* `g` is a point map (not a BoxMap) describing the dynamics.
* `T` is the length of the time interval under consideration.
* `num_points` is the number of points used in each box
* `ϵ` is the maximum size of a perturbation in each coordinate
"""
function finite_time_lyapunov_exponents(boxset::BoxSet, g; T, num_points=20, ϵ=1e-6)
    # TODO with some refactoring, this could maybe also use (some form of)
    # ParallelBoxIterator to implement parallelisation
    d = size(boxset.partition.domain.center)[1]
    dict = Dict{keytype(typeof(boxset.partition)),Float64}()
    sizehint!(dict,length(boxset))

    for box in boxset
        g_center = g(box.center)

        perturbations = ϵ * (2 * rand(Float64,(d,num_points)) .- 1)
        sample_points = box.center .+ perturbations

        distances = [norm(perturbations[:,i]) for i in 1:num_points]
        g_distances = [norm(g_center .- g(sample_points[:,i])) for i in 1:num_points]
        maximal_ftle = log(maximum(g_distances ./ distances)) / abs(T)
        # calculating the minimal ftle would also be straightforward (TODO?)
        dict[point_to_key(boxset.partition,box.center)] = maximal_ftle  # pretty sure this could be more efficient/idiomatic
    end

    return BoxFun(boxset.partition,dict)
end

# Runge-Kutta scheme of 4th order
const half, sixth, third = 1/2, 1/6, 1/3
@muladd @propagate_inbounds function rk4(f, x, τ)
    τ½ = τ * half

    k = f(x)
    dx = @. k * sixth

    k = f(@. x + τ½ * k)
    dx = @. dx + k * third

    k = f(@. x + τ½ * k)
    dx = @. dx + k * third

    k = f(@. x + τ * k)
    dx = @. dx + k * sixth

    return @. x + τ * dx
end

@propagate_inbounds function rk4_flow_map(v, x; step_size=0.01, steps=20)
    for _ in 1:steps
        x = rk4(v, x, step_size)
    end
    return x
end


# surrogate_ground_truth_nodal_mala-logistic
#
# This script runs a long chain based on the
# Metropolis-adjusted Langevin algorithm (MALA). This is used as a gold
# standard as this procedure yields an asymptotically unbiased chain
# and thus is used to compare the efficiency of the biased chains produced
# in the experiment bias_variance-arwmh-logistic.
#
# The target distribution is a Bayesian logistic regression posterior
# under independent Gaussian priors based on a prostate cancer dataset with
# binary outcomes indicating whether cancer has spread to surrounding lymph
# nodes.

using Distributions: Normal
using RDatasets
using StatsFuns: logistic
using SteinDistributions: SteinLogisticRegressionGaussianPrior
using SteinSamplers: runmala

include("experiment_utils.jl")

# Set random number generator seed
@parseintcli seed "s" "seed" 7
# setup the data
data_set = "nodal"
X = y = nothing
nodal = dataset("boot", "nodal")
y = 2 * convert(Array{Float64},nodal[:R]) - 1
# X includes an intercept term
X = convert(Array{Float64, 2}, nodal[[:M, :Aged, :Stage, :Grade, :Xray, :Acid]])

# N is number of samples, p is dimension of parameter
N, d = size(X)
# the max number of samples to run in the solver
n = 10000000
# setup distribution
target = SteinLogisticRegressionGaussianPrior(X, y)
# name of prior + posterior
distname = "logisticgaussianprior"
# intiial guess for beta
beta0 = zeros(d)
# the epsilons to run the experiemnt
epsilons = [1e-2]

for epsilon = epsilons
    # Draw MALA samples
    (betas_mala, numgrad_mala) =
        @setseed runmala(target, beta0;
                         epsilonfunc=t -> epsilon,
                         numiter=n)
    # estimate the final beta moments
    betahat = vec(mean(betas_mala, 1))
    betavcov = cov(betas_mala)
    betasd = sqrt(diag(betavcov))
    # compute the running estimated mu (i.e. P(Y_i = 1 | X_i))
    # when using the first k values from the chain
    ks = round(Int, linspace(0, n, 1001))[2:end]
    meanpredprobs = Array(Float64, N, length(ks))
    for (i,k) in enumerate(ks)
        meanbeta = mean(betas_mala[1:k,:], 1)
        predprobs = logistic(X * meanbeta')
        meanpredprobs[:,i] = vec(predprobs)
    end
    # scale the mu's by the l_infty norm of each X_i
    scaledpredprobs = meanpredprobs ./ vec(maximum(X, 2))
    # fetch constants for computing confidence intervals of beta
    phi = Normal()
    zscore = quantile(phi, 1 - 0.05/2)
    # Package and save results
    summary_data = Dict{Any, Any}(
        "n" => n,
        "seed" => seed,
        "distname" => distname,
        "epsilon" => epsilon,
        "beta.hat" => betahat,
        "beta.vcov" => betavcov,
        "beta.sd" => betasd,
        "beta.hat.low" => betahat - zscore .* betasd,
        "beta.hat.high" => betahat + zscore .* betasd,
        "scaled.pred.probs" => scaledpredprobs,
        "scaled.pred.probs.i" => ks,
        "numgrad" => numgrad_mala,
        "sampler" => "MALA-long-nodal"
    )
    save_json(
        summary_data;
        dir="bias_variance-arwmh-logistic-MALA",
        distname=distname,
        dataset=data_set,
        n=n,
        epsilon=epsilon,
        d=d,
        seed=seed
    )
end


##############################################################################
#
# Wishart Distribution
#
#  Parameters nu and S such that E(X) = nu * S
#  See the rwish and dwish implementation in R's MCMCPack
#  This parametrization differs from Bernardo & Smith p 435
#  in this way: (nu, S) = (2.0 * alpha, 0.5 * beta^-1)
#
##############################################################################

immutable Wishart <: ContinuousMatrixDistribution
    nu::Float64
    Schol::Cholesky{Float64}
    function Wishart(n::Real, Sc::Cholesky{Float64})
        if n > size(Sc, 1) - 1
            new(float64(n), Sc)
        else
            error("Wishart parameters must be df > p - 1")
        end
    end
end

Wishart(nu::Real, S::Matrix{Float64}) = Wishart(nu, cholfact(S))

function insupport(W::Wishart, X::Matrix{Float64})
    return size(X, 1) == size(X, 2) && isApproxSymmmetric(X) &&
           size(X, 1) == size(W.Schol, 1) && hasCholesky(X)
end
# This just checks if X could come from any Wishart
function insupport(::Type{Wishart}, X::Matrix{Float64})
    return size(X, 1) == size(X, 2) && isApproxSymmmetric(X) && hasCholesky(X)
end

mean(w::Wishart) = w.nu * w.Schol[:U]' * w.Schol[:U]

pdf(W::Wishart, X::Matrix{Float64}) = exp(logpdf(W, X))

size(W::Wishart) = size(W.Schol)
dim(W::Wishart) = size(W.Schol, 1)

function expected_logdet(W::Wishart)
    logd = 0.
    d = dim(W)

    for i=1:d
        logd += digamma(0.5 * (W.nu + 1 - i))
    end

    logd += d * log(2)
    logd += logdet(W.Schol)

    return logd
end

function lognorm(W::Wishart)
    d = dim(W)
    return (W.nu / 2) * logdet(W.Schol) + (d * W.nu / 2) * log(2) + lpgamma(d, W.nu / 2)
end

function logpdf(W::Wishart, X::Matrix{Float64})
    if !insupport(W, X)
        return -Inf
    else
        d = dim(W)
        logd = -lognorm(W)
        logd += 0.5 * (W.nu - d - 1.0) * logdet(X)
        logd -= 0.5 * trace(W.Schol \ X)
        return logd
    end
end

function rand(w::Wishart)
    p = size(w.Schol, 1)
    X = zeros(p, p)
    for ii in 1:p
        X[ii, ii] = sqrt(rand(Chisq(w.nu - ii + 1)))
    end
    if p > 1
        for col in 2:p
            for row in 1:(col - 1)
                X[row, col] = randn()
            end
        end
    end
    Z = X * w.Schol[:U]
    return Z' * Z
end

function entropy(W::Wishart)
    d = dim(W)
    return lognorm(W) - (W.nu - d - 1) / 2 * expected_logdet(W) + W.nu * d / 2
end

var(w::Wishart) = error("Not yet implemented")


module Blackscholes

using GPUBenchmarks, BenchmarkTools
import CUDAdrv

using CUDAnative
const cu = CUDAnative

const description = """
Blackschole is a nice benchmark for broadcasting performance.
It's a medium heavy calculation per array element, where the calculation is completely
independant from each other.
The CuArray package is a bit slower here compared to GPUArrays, which should be straightforward to fix.
I suspect that it's due to more promotions between integer types in the indexing code.
"""

function blackscholes(
        sptprice,
        strike,
        rate,
        volatility,
        time
    )
    logterm = log( sptprice / strike)
    powterm = .5f0 * volatility * volatility
    den = volatility * sqrt(time)
    d1 = (((rate + powterm) * time) + logterm) / den
    d2 = d1 - den
    NofXd1 = cndf2(d1)
    NofXd2 = cndf2(d2)
    futureValue = strike * exp(-rate * time)
    c1 = futureValue * NofXd2
    call_ = sptprice * NofXd1 - c1
    put  = call_ - futureValue + sptprice
    return put
end
function blackscholes_af(
        sptprice,
        strike,
        rate,
        volatility,
        time
    )
    logterm = log( sptprice / strike)
    powterm = .5f0 * volatility * volatility
    den = volatility * sqrt(time)
    d1 = (((rate + powterm) * time) + logterm) / den
    d2 = d1 - den
    NofXd1 = cndf2(d1)
    NofXd2 = cndf2(d2)
    futureValue = strike * exp(-rate * time)
    c1 = futureValue * NofXd2
    call_ = sptprice * NofXd1 - c1
    put  = call_ - futureValue + sptprice
    return put
end

@inline function cndf2(x)
    0.5f0 + 0.5f0 * erfc(0.707106781f0 * x)
end

function cu_blackscholes(sptprice, strike, rate, volatility, time)
    logterm = cu.log( sptprice / strike)
    powterm = .5f0 * volatility * volatility
    den = volatility * cu.sqrt(time)
    d1 = (((rate + powterm) * time) + logterm) / den
    d2 = d1 - den
    NofXd1 = cu_cndf2(d1)
    NofXd2 = cu_cndf2(d2)
    futureValue = strike * cu.exp(- rate * time)
    c1 = futureValue * NofXd2
    call_ = sptprice * NofXd1 - c1
    put  = call_ - futureValue + sptprice
    return put
end

function cu_cndf2(x)
    0.5f0 + 0.5f0 * cu.erfc(0.707106781f0 * x)
end

function execute_broadcast(f, res, a, b, c, d, e)
    res .= f.(a, b, c, d, e)
    synchronize(res)
end

function execute_broadcast_af(f, res, a, b, c, d, e)
    res .= f.(a, b, c, d, e)
    synchronize(res)
end

function makeresult(bench, N, device, hardware, mdiff)
    BenchResult(
        "Blackscholes",
        bench,
        N,
        Complex64,
        string(device),
        hardware,
        @__FILE__,
        mdiff
    )
end

function execute(device)
    hardware, array_type = init(device)
    T = Float32
    results = BenchResult[]
    func = is_cudanative(device) ? cu_blackscholes : blackscholes
    func = is_arrayfire(device) ? blackscholes_af : func
    for x in 1:7
        N = 10^x
        jl_arrays = (
            zeros(T, N),
            T[42.0 for i = 1:N],
            T[40.0 + (i / N) for i = 1:N],
            T[0.5 for i = 1:N],
            T[0.2 for i = 1:N],
            T[0.5 for i = 1:N]
        )
        gpu_arrays = array_type.(jl_arrays)
        res, a, b, c, d, e = gpu_arrays
        br_func = is_arrayfire(device) ? execute_broadcast_af : execute_broadcast
        bench = @benchmark $(br_func)($func, $res, $a, $b, $c, $d, $e)
        jl_arrays[1] .= blackscholes.(jl_arrays[2:end]...)
        mdiff = meandifference(jl_arrays[1], res)
        push!(results, makeresult(bench, N, device, hardware, mdiff))
        for elem in gpu_arrays
            free(elem)
            println("freeing element: ", CUDAdrv.Mem.used() / 10^7)
        end
        gc()
    end
    return results
end

end


function collatzLength(n)
  i = 1
  while n != 1
    if mod(n,2) == 0
      n = n / 2
    else
      n = 3n + 1
    end
    i = i + 1
  end
  i
end

max = 1
maxAt = 1
for i in 1:1000000
  if collatzLength(i) > max
    max = collatzLength(i)
    maxAt = i
  end
end

println(maxAt)


"""
RationalQuadraticKernel([ρ=1.0[,α=2.0]])
The rational-quadratic kernel is an isotropic Mercer kernel given by the formula:
```
    κ(x,y)=(1+ρ²||x−y||²/α)^(-α)
```
where `α` is a shape parameter of the Euclidean distance. Check [`GammaRationalQuadraticKernel`](@ref) for a generalization.
"""
struct RationalQuadraticKernel{T,Tr,Tα<:Real} <: Kernel{T,Tr}
    transform::Tr
    metric::SqEuclidean
    α::Tα
    function RationalQuadraticKernel{T,Tr,Tα}(t::Tr,α::Tα) where {T,Tr,Tα<:Real}
        new{T,Tr,Tα}(t,SqEuclidean(),α)
    end
end

function RationalQuadraticKernel(ρ::T₁=1.0,α::T₂=2.0) where {T₁<:Real,T₂<:Real}
    @check_args(RationalQuadraticKernel, α, α > zero(T₂), "α > 1")
    RationalQuadraticKernel{T₁,ScaleTransform{Base.RefValue{T₁}},T₂}(ScaleTransform(ρ),α)
end

function RationalQuadraticKernel(ρ::A,α::T=2.0) where {A<:AbstractVector{<:Real},T<:Real}
    @check_args(RationalQuadraticKernel, α, α > zero(T), "α > 1")
    RationalQuadraticKernel{eltype(A),ScaleTransform{A},T}(ScaleTransform(ρ),α)
end

function RationalQuadraticKernel(t::Tr,α::T=2.0) where {Tr<:Transform,T<:Real}
    @check_args(RationalQuadraticKernel, α, α > zero(T), "α > 1")
    RationalQuadraticKernel{eltype(t),Tr,T}(t,α)
end

params(k::RationalQuadraticKernel) = (params(transform(k)),k.α)

@inline kappa(κ::RationalQuadraticKernel, d²::T) where {T<:Real} = (one(T)+d²/κ.α)^(-κ.α)

"""
`GammaRationalQuadraticKernel([ρ=1.0[,α=2.0[,γ=2.0]]])`
The Gamma-rational-quadratic kernel is an isotropic Mercer kernel given by the formula:
```
    κ(x,y)=(1+ρ^(2γ)||x−y||^(2γ)/α)^(-α)
```
where `α` is a shape parameter of the Euclidean distance and `γ` is another shape parameter.
"""
struct GammaRationalQuadraticKernel{T,Tr,Tα<:Real,Tγ<:Real} <: Kernel{T,Tr}
    transform::Tr
    metric::SqEuclidean
    α::Tα
    γ::Tγ
    function GammaRationalQuadraticKernel{T,Tr,Tα,Tγ}(t::Tr,α::Tα,γ::Tγ) where {T,Tr,Tα<:Real,Tγ<:Real}
        new{T,Tr,Tα,Tγ}(t,SqEuclidean(),α,γ)
    end
end

function GammaRationalQuadraticKernel(ρ::T₁=1.0,α::T₂=2.0,γ::T₃=2.0) where {T₁<:Real,T₂<:Real,T₃<:Real}
    @check_args(GammaRationalQuadraticKernel, α, α > one(T₂), "α > 1")
    @check_args(GammaRationalQuadraticKernel, γ, γ >= one(T₂), "γ >= 1")
    GammaRationalQuadraticKernel{T₁,ScaleTransform{Base.RefValue{T₁}},T₂,T₃}(ScaleTransform(ρ),α,γ)
end

function GammaRationalQuadraticKernel(ρ::A,α::T₁=2.0,γ::T₂=2.0) where {A<:AbstractVector{<:Real},T₁<:Real,T₂<:Real}
    @check_args(GammaRationalQuadraticKernel, α, α > one(T₁), "α > 1")
    @check_args(GammaRationalQuadraticKernel, γ, γ >= one(T₂), "γ >= 1")
    GammaRationalQuadraticKernel{eltype(A),ScaleTransform{A},T₁,T₂}(ScaleTransform(ρ),α,γ)
end

function GammaRationalQuadraticKernel(t::Tr,α::T₁=2.0,γ::T₂=2.0) where {Tr<:Transform,T₁<:Real,T₂<:Real}
    @check_args(GammaRationalQuadraticKernel, α, α > one(T₁), "α > 1")
    @check_args(GammaRationalQuadraticKernel, γ, γ >= one(T₂), "γ >= 1")
    GammaRationalQuadraticKernel{eltype(t),Tr,T₁,T₂}(t,α,γ)
end

params(k::GammaRationalQuadraticKernel) = (params(k.transform),k.α,k.γ)

@inline kappa(κ::GammaRationalQuadraticKernel, d²::T) where {T<:Real} = (one(T)+d²^κ.γ/κ.α)^(-κ.α)


@testset "Selections" begin

    @testset "Rank" begin
        Random.seed!(2);
        fitness = [5, 1, 3, 4, 2.0]
        s = ranklinear(1.0)
        @test s(fitness, 2) == [2,3]

        s = ranklinear(2.0)
        @test s([1,2], 2) == [2,2]
    end

    @testset "Uniform" begin
        s = uniformranking(2)
        @test sort(unique(s([1.0,2.0,3.0], 10))) == [2,3]
        @test sort(unique(s([5,2,3], 5))) == [1,3]
        @test_throws AssertionError s([1.,2.], 2)
    end

    @testset "Roulette" begin
        @test roulette([0.0,0.0,3.0], 2) == [3, 3]
        @test roulette([0,2,0], 2) == [2,2]
    end

    @testset "Truncation" begin
        @test truncation([0.0,2.0,3.0], 2) == [3,2]
        @test truncation([0,2,1], 1) == [2]
        @test_throws AssertionError truncation([1.,2.], 5)
    end

    @testset "Tournament" begin
        Random.seed!(2);
        @test_throws AssertionError tournament(0)
        t = tournament(2)
        @test t([0,2,0],2) == [2,2]
        @test t([0.0,0.0,1.0],2) == [3,3]
    end

    @testset "SUS" begin
        @test sort(unique(sus([1.0,0.0,2.0], 5))) == [1,3]
        @test sort(unique(sus([0,1,2], 5))) == [2,3]
    end

end


# d x 2 matrix
type LineSegment
    ab::Array{Float64,2}
end

function create_line( ls::Array{Float64,2})
    li = Vector{LineSegment}(size(ls,2)-1)
    for zi in 1:size(ls,2)-1; li[zi] = LineSegment(ls[:,[zi,zi+1]]); end
    return li
end

function dist2line(line::Vector{LineSegment},x::Array{Float64,2})
    dd = zeros(size(line,1),size(x,2))
    for i in 1:size(line,1)
        dd[i,:] = dist2LS(line[i],x)
    end
    #print(dd)
    d  = minimum(dd,1)
    return d[:]
end

function dist2LS(ls::LineSegment,x::Array{Float64,2})
    u = ls.ab[:,2] - ls.ab[:,1]
    v = broadcast( - , x , ls.ab[:,1] )

    uv = RowVector(u) * v
    dec = uv / (norm(u)^2)

    p = broadcast( + , ls.ab[:,1] , broadcast( * , dec , u  ) )
    d = zeros(1,size(x,2))

    #println(u)
    #println(v)
    #println(uv)
    #println(dec)
    #println(p)

    for i=1:size(x,2)
        if (dec[i]<0)
            d[i] = norm(ls.ab[:,1] - x[:,i])
        elseif (dec[i]>1)
            d[i] = norm(ls.ab[:,2] - x[:,i])
        else
                d[i] = norm( p[:,i] - x[:,i] )
        end
    end
    return d[:]
end


using Base.Test
using BlackScholes

tol = 1.0e-6
@test_approx_eq_eps(BSCall(10.0,10.0,0.15,0.0323,7.0,0.02),1.699938,tol)
@test_approx_eq_eps(BSPut(10.0,10.0,0.15,0.0323,7.0,0.02),0.982739,tol)

@test BSImpliedVol(10.0,10.0,0.0323,7.0,0.02,BSCall(10.0,10.0,0.55,0.0323,7.0,0.02),true) == 0.55
@test BSImpliedVol(10.0,10.0,0.0323,7.0,0.02,BSPut(10.0,10.0,0.55,0.0323,7.0,0.02),false) == 0.55

@test BSImpliedVol(10.0,10.0,0.0323,7.0,0.02,BSCall(10.0,10.0,0.55,0.0323,7.0,0.02),true) == 0.55


using Poisson2D
using LinearAlgebra
using BenchmarkTools

@show n_threads = Threads.nthreads();
# without the next line, it will use the max number of threads available
BLAS.set_num_threads(n_threads);
# set problem parameters
a = 0;
b = 1;
c = 0;
d = 1;
n = 999;

kx = 5;
ky = 4;
f = (x,y)->((kx*π/(b-a))^2 + (ky * π/(d-c))^2) * sin(kx*π*(x-a)/(b-a)) * sin(ky*π*(y-c)/(d-c));

# set up and solve problem
problem = FDPoisson2DProblem(a,b,c,d,n,f);
assemble_system!(problem);
@btime U = direct_solve_poisson2d(problem);


# If you want a scatter plot of the spectra, define `plotresults = true`
# prior to running this script.

using CompScienceMeshes, BEAST
using LinearAlgebra

# Γ = readmesh(joinpath(@__DIR__,"sphere2.in"))
Γ = meshrectangle(1.0, 1.0, 0.2, 3)
println("Mesh with $(numvertices(Γ)) vertices and $(numcells(Γ)) cells.")
X = raviartthomas(Γ)
Y = BEAST.buffachristiansen3(Γ)

κ = ω = 1.0; γ = κ*im
T = MWSingleLayer3D(γ)
N = NCross()

Txx = assemble(T,X,X); println("primal discretisation assembled.")
Tyy = assemble(T,Y,Y); println("dual discretisation assembled.")
Nxy = assemble(N,X,Y); println("duality form assembled.")

iNxy = inv(Nxy); println("duality form inverted.")
A = iNxy' * Tyy * iNxy * Txx

@show cond(Txx)
@show cond(Tyy)
@show cond(A)

wx, wy, wp = eigvals(Txx), eigvals(Tyy), eigvals(A); println("eigvals found.")

(@isdefined plotresults) || (plotresults = false)
if plotresults
    @eval begin
        using Plots
        plot()
        scatter!(wx,c=:blue,label="primal")
        scatter!(wy,c=:red,label="dual")
        scatter!(wp,c=:green,label="CMP")
        plot!(xlim=(-0.5,0.5),ylim=(-10.0,10.0))
    end
end


using ITensors

function ising(n; h)
  H = OpSum()
  for j in 1:(n - 1)
    H -= "Z", j, "Z", j + 1
  end
  for j in 1:n
    H += h, "X", j
  end
  return H
end


# transpose
function ad_transpose(x::AD)
  x.value = (x.value)'
  x.grad = (x.grad)'
  return x
end
transpose(x::AD) = ad_transpose(x)

# addition
function ad_add(x::AD, y::AD)
  gradSize = (size(x.value) > size(y.value))? size(x.value):size(y.value)
  result = AD("+", x.value + y.value, true, zeros(gradSize), ad_addD, (x.level > y.level ? x.level:y.level)+1)
  push!(result.parents, x)
  push!(result.parents, y)
  return result
end
function ad_addD(this::AD)
  prevGradSize = size(this.grad)
  adNode1Size = size(this.parents[1].value)
  adNode2Size = size(this.parents[2].value)

  if adNode1Size == prevGradSize
    this.parents[1].grad += 1 * this.grad
  elseif adNode1Size[1] == prevGradSize[1]
    this.parents[1].grad += 1 * sum(this.grad,2)
  else
    this.parents[1].grad += 1 * sum(this.grad,1)
  end

  if adNode2Size == prevGradSize
    this.parents[2].grad += 1 * this.grad
  elseif adNode2Size[1] == prevGradSize[1]
    this.parents[2].grad += 1 * sum(this.grad,2)
  else
    this.parents[2].grad += 1 * sum(this.grad,1)
  end
  return
end
+(x::AD, y::AD) = ad_add(x, y)

function ad_eladd(x::AD, y::AD)
  gradSize = (size(x.value) > size(y.value))? size(x.value):size(y.value)
  result = AD(".+", x.value .+ y.value, true, zeros(gradSize), ad_eladdD, (x.level > y.level ? x.level:y.level)+1)
  push!(result.parents, x)
  push!(result.parents, y)
  return result
end
function ad_eladdD(this::AD)
  prevGradSize = size(this.grad)
  adNode1Size = size(this.parents[1].value)
  adNode2Size = size(this.parents[2].value)

  if adNode1Size == prevGradSize
    this.parents[1].grad += 1 * this.grad
  elseif adNode1Size[1] == prevGradSize[1]
    this.parents[1].grad += 1 * sum(this.grad,2)
  else
    this.parents[1].grad += 1 * sum(this.grad,1)
  end

  if adNode2Size == prevGradSize
    this.parents[2].grad += 1 * this.grad
  elseif adNode2Size[1] == prevGradSize[1]
    this.parents[2].grad += 1 * sum(this.grad,2)
  else
    this.parents[2].grad += 1 * sum(this.grad,1)
  end

  return
end
.+(x::AD, y::AD) = ad_eladd(x, y)

# multiplication
function ad_mul(x::AD, y::AD)
  if size(x.value) == (1,1) || size(y.value) == (1,1)
    return ad_elmul(x, y)
  else
    resultValue = x.value * y.value
    result = AD("*", resultValue, true, zeros(size(resultValue)), ad_mulD, (x.level > y.level ? x.level:y.level)+1)
    push!(result.parents, x)
    push!(result.parents, y)
    return result
  end
end
function ad_mulD(this::AD)
  rowNode1, colNode1 = size(this.parents[1].value)
  rowNode2, colNode2 = size(this.parents[2].value)

  for i = 1:rowNode1, j = 1:colNode2
    pd = this.grad[i,j]
    for k = 1:colNode1
      this.parents[1].grad[i,k] += this.parents[2].value[k,j] * pd
      this.parents[2].grad[k,j] += this.parents[1].value[i,k] * pd
    end
  end
end
*(x::AD, y::AD) = ad_mul(x, y)

# element-wise multiplication
function ad_elmul(x::AD, y::AD)
  gradSize = (size(x.value) > size(y.value))? size(x.value):size(y.value)
  result = AD(".*", x.value .* y.value, true, zeros(gradSize), ad_elmulD, (x.level > y.level ? x.level:y.level)+1)
  push!(result.parents, x)
  push!(result.parents, y)
  return result
end
function ad_elmulD(this::AD)
  this.parents[1].grad += this.parents[2].value .* this.grad
  this.parents[2].grad += this.parents[1].value .* this.grad
  return
end
.*(x::AD, y::AD) = ad_elmul(x, y)

# relu
function reluMat(z::Matrix{Float64})
    return max(0.0, z);
end
function relu(x::AD)
  result = AD("relu", reluMat(x.value), true, zeros(size(x.value)), ad_reluD, x.level + 1)
  push!(result.parents, x)
  return result
end
function reluGradientMat(z::Matrix{Float64})
    grad = ones(z);
    grad[z.<=0] = 0;
    return grad;
end
function ad_reluD(this::AD)
  this.parents[1].grad += reluGradientMat(this.parents[1].value) .* this.grad
  return
end

# sigmoid
function sigmoidMat(z::Matrix{Float64})
    g = 1.0 ./ (1.0 + exp(-z));
    return g;
end
function sigmoid(x::AD)
  result = AD("sigmoid", sigmoidMat(x.value), true, zeros(size(x.value)), ad_sigmoidD, x.level + 1)
  push!(result.parents, x)
  return result
end
function sigmoidGradientMat(z::Matrix{Float64})
  return z.*(1.0 - z);
end
function ad_sigmoidD(this::AD)
  this.parents[1].grad += sigmoidGradientMat(this.parents[1].value) .* this.grad
  return
end

# tanh
function tanh(x::AD)
  result = AD("tanh", tanh(x.value), true, zeros(size(x.value)), ad_tanhD, x.level + 1)
  push!(result.parents, x)
  return result
end
function tanhGradientMat(z::Matrix{Float64})
  return (1.0 - z^2)
end
function ad_tanhD(this::AD)
  this.parents[1].grad += tanhGradientMat(this.parents[1].value) .* this.grad
  return
end

# softmax
function softmaxMat(z::Matrix{Float64}, dimOrder::Int64)
    g = exp(z)
    g ./=sum(g, dimOrder)
    return g;
end

# Using cross-entropy loss function.
# evaluating d(loss)/d(loss)*d(loss)/d(softmax)*d(softmax)/d(z)
# because calculating d(softmax)/d(z) individually requires too much memory.
function softmaxLoss(y::Matrix{Float64}, x::AD, dimOrder::Int64, computeLoss::Bool, reg::Float64)
   # dimOrder: along which the dimentions/classes are present
   probs = softmaxMat(x.value, dimOrder)
   # compute the loss: average cross-entropy loss and regularization
   loss = 0.0
   if computeLoss
     dataInstanceOrder = (dimOrder == 1? 2 : 1)
     numInstance = size(y, dataInstanceOrder)
     correctLogProbs = zeros(size(y))
     for j in 1:numInstance
       correctLogProbs[j, :] = -log(probs[j,Int64(dataInstanceOrder==1?y[j, 1]:y[1,j])]);
     end
     loss = sum(correctLogProbs)/numInstance + reg
   end

   # return node
   result = AD("softmaxLoss", probs, true, zeros(size(x.value)), ad_softmaxLossD, x.level + 1, y, dataInstanceOrder, loss)
   push!(result.parents, x)
   return result
end

softmaxLoss(y::Matrix{Float64}, x::AD, dimOrder::Int64, computeLoss::Bool) = softmaxLoss(y, x, dimOrder, computeLoss, 0.0)
softmaxLoss(y::Matrix{Float64}, x::AD, dimOrder::Int64) = softmaxLoss(y, x, dimOrder, false, 0.0)

function ad_softmaxLossD(this::AD)
  this.parents[1].grad += this.value .* this.grad
  if this.outputOrder == 1
    for j in 1:size(this.parents[1].grad, this.outputOrder)
      this.parents[1].grad[j,Int64(this.output[j,1])] -= 1;
    end
  elseif this.outputOrder == 2
    for j in 1:size(this.parents[1].grad, this.outputOrder)
      this.parents[1].grad[Int64(this.output[1,j]), j] -= 1;
    end
  end
  this.parents[1].grad /= size(this.output, this.outputOrder);
  return
end


using CSV, GLM, Plots; pyplot()

data = CSV.read("../data/polynomialData.csv")
data.X2 = abs2.(data.X)
model1 = lm(@formula(Y ~ X + X2), data)
model2 = lm(@formula(Y ~ X*X), data)
println(model1)
println(model2)

yHat(x) = coef(model1)[1] + coef(model1)[2]*x + coef(model1)[3]*x^2

xGrid = -5:0.1:5
scatter(data.X, data.Y, c=:blue, msw=0)
plot!(xGrid, yHat.(xGrid),
    xlabel="X", ylabel="Y", c=:red, legend=:none)


function generatetransformation(g::Geometry{3},t::TetrahedronTopology,k::Int64)

	@inbounds begin
		dx1=x(g,t,k,2)-x(g,t,k,1)
		dy1=y(g,t,k,2)-y(g,t,k,1)
		dz1=z(g,t,k,2)-z(g,t,k,1)

		dx2=x(g,t,k,3)-x(g,t,k,1)
		dy2=y(g,t,k,3)-y(g,t,k,1)
		dz2=z(g,t,k,3)-z(g,t,k,1)

		dx3=x(g,t,k,4)-x(g,t,k,1)
		dy3=y(g,t,k,4)-y(g,t,k,1)
		dz3=z(g,t,k,4)-z(g,t,k,1)

		# determinant on each tetrahedron
		edet = dx1*dy2*dz3+dx2*dy3*dz1+dx3*dy1*dz2-dx3*dy2*dz1-dx1*dy3*dz2-dx2*dy1*dz3

		# transformation on each tetrahedron
		dFinv = zeros(Float64,3,3)

		dFinv[1,1] = (dy2*dz3-dy3*dz2)/edet
		dFinv[1,2] = (dx3*dz2-dx2*dz3)/edet
		dFinv[1,3] = (dx2*dy3-dx3*dy2)/edet

		dFinv[2,1] = (dy3*dz1-dy1*dz3)/edet
		dFinv[2,2] = (dx1*dz3-dx3*dz1)/edet
		dFinv[2,3] = (dx3*dy1-dx1*dy3)/edet

		dFinv[3,1] = (dy1*dz2-dy2*dz1)/edet
		dFinv[3,2] = (dx2*dz1-dx1*dz2)/edet
		dFinv[3,3] = (dx1*dy2-dx2*dy1)/edet
	end

	return edet,dFinv

end

function generatetransformation(g::Geometry{2},t::TriangleTopology,k::Int64)

	@inbounds begin
		dx1 = x(g,t,k,2)-x(g,t,k,1)
		dy1 = y(g,t,k,2)-y(g,t,k,1)

		dx2 = x(g,t,k,3)-x(g,t,k,1)
		dy2 = y(g,t,k,3)-y(g,t,k,1)
	end

	edet  = dx1*dy2 - dx2*dy1
	dFinv = [dy2 -dx2;-dy1 dx1] / edet

	return edet,dFinv

end

function generatetransformation(g::Geometry{3},t::TriangleTopology,k::Int64)

	@inbounds begin
		dx1 = x(g,t,k,2)-x(g,t,k,1)
		dy1 = y(g,t,k,2)-y(g,t,k,1)
		dz1 = z(g,t,k,2)-z(g,t,k,1)

		dx2 = x(g,t,k,3)-x(g,t,k,1)
		dy2 = y(g,t,k,3)-y(g,t,k,1)
		dz2 = z(g,t,k,3)-z(g,t,k,1)
	end

	n1   = sqrt(dx1^2+dy1^2+dz1^2)       # norm v1
	n22  = dx2^2+dy2^2+dz2^2             # squared norm v2
	s12  = (dx1*dx2+dy1*dy2+dz1*dz2)/n1  # scalar product v1*v2

	# project vp = v2 - (v1*v2)*v1/norm(v1)
	dx2p = dx2 - s12*dx1
	dy2p = dy2 - s12*dy1
	dz2p = dz2 - s12*dz1

	# set 2D coordinates based on e1 = v1/|v1| and e2=vp/|vp|
	dxp1 = n1
	dyp1 = 0.0
	dxp2 = s12
	dyp2 = sqrt(n22-dxp2^2)

	# compute determinant and transformation
	edet  = dxp1*dyp2 - dxp2*dyp1
	dFinv = [dyp2 -dxp2;-dyp1 dxp1] / edet

	return edet,dFinv

end

function generatetransformation(g::Geometry{1},t::IntervalTopology,k::Int64)

	@inbounds begin
		dx = x(g,t,k,2)-x(g,t,k,1)
	end

	edet  = dx
	dFinv = [1 / edet]'
	return edet,dFinv
end

function generatetransformation(g::Geometry{2},t::IntervalTopology,k::Int64)

	@inbounds begin
		dx = x(g,t,k,2)-x(g,t,k,1)
		dy = y(g,t,k,2)-y(g,t,k,1)
	end

	edet  = sqrt(dx^2+dy^2)
	dFinv = reshape([1 / edet],(1,1))  # should be this if we assume to only return the tangential
	# derivative and t. transformation. question is if we also need a tangential
	# projector in this case. for this, however, all "generatetransformation" should
	# have this syntax

	return edet,dFinv

end

function generatetransformation(g::Geometry{3},t::IntervalTopology,k::Int64)

	@inbounds begin
		dx = x(g,t,k,2)-x(g,t,k,1)
		dy = y(g,t,k,2)-y(g,t,k,1)
		dz = z(g,t,k,2)-z(g,t,k,1)
	end

	edet  = sqrt(dx^2+dy^2+dz^2)
	dFinv = reshape([1 / edet],(1,1))  # see comment Geometry{2} IntervalTopology

	return edet,dFinv

end


# unit vectors and identity matrix
const e1 = SVector(1, 0, 0)
const e2 = SVector(0, 1, 0)
const e3 = SVector(0, 0, 1)
const I3 = @SMatrix [1 0 0; 0 1 0; 0 0 1]
const Z3 = @SMatrix zeros(3,3)

"""
    tilde(x)

Construct the cross product operator matrix
"""
@inline tilde(x) = @SMatrix [0 -x[3] x[2]; x[3] 0 -x[1]; -x[2] x[1] 0]

"""
    wiener_milenkovic(c)

Construct a Wiener-Milenkovic transformation matrix, given the three Wiener-Milenkovic
parameters in `c`.

Note that the corresponding rotation matrix is the transpose of this
transformation matrix.
"""
@inline function wiener_milenkovic(c)

    c0 = 2 - c'*c/8

    C = 1/(4-c0)^2*(@SMatrix [
        c0^2 + c[1]^2 - c[2]^2 - c[3]^2  2*(c[1]*c[2] + c0*c[3])          2*(c[1]*c[3]-c0*c[2]);
        2*(c[1]*c[2] - c0*c[3])          c0^2 - c[1]^2 + c[2]^2 - c[3]^2  2*(c[2]*c[3] + c0*c[1]);
        2*(c[1]*c[3] + c0*c[2])          2*(c[2]*c[3] - c0*c[1])          c0^2 - c[1]^2 - c[2]^2 + c[3]^2
        ]
    )

    return C
end

"""
    transform_properties(K, T)

Applies the transformation `T` to the stiffness or mass matrix `K`
"""
transform_properties(X, T) = [T' Z3; Z3 T'] * X * [T Z3; Z3 T]

"""
    rotation_parameter_scaling(θ)

Extract a scaling parameter which may be multiplied by the angular parameters
to yield the Wiener-Milenkovic rotation parameters.  Use of this scaling
parameter allows deflections greater than 360 degrees.
"""
@inline function rotation_parameter_scaling(θ)

    θ = real(θ)

    θnorm = sqrt(θ'*θ)

    scaling = 1

    m = div(div(θnorm, 4) + 1, 2) # number of times to add/subtract

    if θnorm > 4
        scaling -= 8*m/θnorm
    end

    return scaling
end

@inline function rotation_parameter_scaling_θ(θ)

    return rotation_parameter_scaling_θ(rotation_parameter_scaling(θ), θ)
end

@inline function rotation_parameter_scaling_θ(scaling, θ)

    return ifelse(isone(scaling), zero(θ), (1 - scaling)*θ/(θ'*θ))
end

@inline function rotation_parameter_scaling_θ_θ(θ)

    return rotation_parameter_scaling_θ_θ(rotation_parameter_scaling(θ), θ)
end

@inline function rotation_parameter_scaling_θ_θ(scaling, θ)

    return ifelse(isone(scaling), (@SMatrix zeros(eltype(θ),3,3)), (1 - scaling)/(θ'*θ)*(I - 3*θ*θ'/(θ'*θ)))
end

"""
    get_C(θ)

Returns the transformation matrix `C` given the three angular parameters in `θ`.
"""
@inline function get_C(θ)

    scaling = rotation_parameter_scaling(θ)

    c = scaling*θ

    return wiener_milenkovic(c)
end

"""
    get_C_θ([C, ] θ)

Calculate the derivative of the Wiener-Milenkovic transformation matrix `C` with
respect to each of the rotation parameters in `θ`.
"""
@inline get_C_θ(c) = get_C_θ(get_C(c), c)

@inline function get_C_θ(C, θ)

    scaling = rotation_parameter_scaling(θ)
    scaling_θ = rotation_parameter_scaling_θ(scaling, θ)

    c = scaling*θ
    c_θ = scaling_θ*θ' + scaling*I3

    c0 = 2 - c'*c/8
    c0_c = -c/4

    tmp = 1/(4-c0)^2
    tmp_c = 2/(4-c0)^3*c0_c

    C_c1 = tmp_c[1]*C/tmp + tmp*(@SMatrix [
        2*c0*c0_c[1] + 2*c[1]    2*(c[2] + c0_c[1]*c[3]) 2*(c[3] - c0_c[1]*c[2]);
         2*(c[2] - c0_c[1]*c[3])  2*c0*c0_c[1] - 2*c[1]   2*(c0_c[1]*c[1] + c0);
         2*(c[3] + c0_c[1]*c[2]) -2*(c0_c[1]*c[1] + c0)   2*c0*c0_c[1] - 2*c[1]
        ]
    )

    C_c2 = tmp_c[2]*C/tmp + tmp*(@SMatrix [
        2*c0*c0_c[2] - 2*c[2]   2*(c[1] + c0_c[2]*c[3]) -2*(c0_c[2]*c[2] + c0);
        2*(c[1] - c0_c[2]*c[3]) 2*c0*c0_c[2] + 2*c[2]    2*(c[3] + c0_c[2]*c[1]);
         2*(c0_c[2]*c[2] + c0)   2*(c[3] - c0_c[2]*c[1])  2*c0*c0_c[2] - 2*c[2]
        ]
    )

    C_c3 = tmp_c[3]*C/tmp + tmp*(@SMatrix [
         2*c0*c0_c[3] - 2*c[3]   2*(c0_c[3]*c[3] + c0)    2*(c[1] - c0_c[3]*c[2]);
        -2*(c0_c[3]*c[3] + c0)   2*c0*c0_c[3] - 2*c[3]    2*(c[2] + c0_c[3]*c[1]);
         2*(c[1] + c0_c[3]*c[2]) 2*(c[2] - c0_c[3]*c[1])  2*c0*c0_c[3] + 2*c[3]
        ]
    )

    C_θ1 = C_c1*c_θ[1,1] + C_c2*c_θ[2,1] + C_c3*c_θ[3,1]
    C_θ2 = C_c1*c_θ[1,2] + C_c2*c_θ[2,2] + C_c3*c_θ[3,2]
    C_θ3 = C_c1*c_θ[1,3] + C_c2*c_θ[2,3] + C_c3*c_θ[3,3]

    return C_θ1, C_θ2, C_θ3
end

"""
    get_C_t([C, ] θ, θ_t)

Calculate the derivative of the Wiener-Milenkovic transformation matrix `C` with
respect to the scalar parameter `t`. `θ_t` is the derivative of the angular parameter
`θ` with respect to `t`.
"""
@inline get_C_t(θ, θ_t) = get_C_t(get_C(θ), θ, θ_t)

@inline function get_C_t(C, θ, θ_t)

    scaling = rotation_parameter_scaling(θ)
    scaling_θ = rotation_parameter_scaling_θ(scaling, θ)

    c = scaling*θ
    c_t = scaling_θ'*θ_t*θ + scaling*θ_t

    c0 = 2 - c'*c/8
    c0dot = -c'*c_t/4
    tmp = 1/(4-c0)^2
    tmpdot = 2/(4-c0)^3*c0dot

    Cdot = tmpdot*C/tmp + tmp*(@SMatrix [
        2*c0*c0dot + 2*c[1]*c_t[1] - 2*c[2]*c_t[2] - 2*c[3]*c_t[3] (
        2*(c_t[1]*c[2] + c[1]*c_t[2] + c0dot*c[3] + c0*c_t[3])) (
        2*(c_t[1]*c[3] + c[1]*c_t[3] - c0dot*c[2] - c0*c_t[2]));
        2*(c_t[1]*c[2] + c[1]*c_t[2] - c0dot*c[3] - c0*c_t[3]) (
        2*c0*c0dot - 2*c[1]*c_t[1] + 2*c[2]*c_t[2] - 2*c[3]*c_t[3]) (
        2*(c_t[2]*c[3] + c[2]*c_t[3] + c0dot*c[1] + c0*c_t[1]));
        2*(c_t[1]*c[3] + c[1]*c_t[3] + c0dot*c[2] + c0*c_t[2]) (
        2*(c_t[2]*c[3] + c[2]*c_t[3] - c0dot*c[1] - c0*c_t[1])) (
        2*c0*c0dot - 2*c[1]*c_t[1] - 2*c[2]*c_t[2] + 2*c[3]*c_t[3])]
    )

    return Cdot
end

"""
    get_C_t_θ(θ, θ_t)

Calculate the derivative of the time derivative of the Wiener-Milenkovic
transformation matrix `C` with respect to `θ`.
"""
@inline function get_C_t_θ(θ, θ_t)

    θ1, θ2, θ3 = θ
    C_t_θ1 = ForwardDiff.derivative(θ1 -> get_C_t(SVector(θ1, θ2, θ3), θ_t), θ1)
    C_t_θ2 = ForwardDiff.derivative(θ2 -> get_C_t(SVector(θ1, θ2, θ3), θ_t), θ2)
    C_t_θ3 = ForwardDiff.derivative(θ3 -> get_C_t(SVector(θ1, θ2, θ3), θ_t), θ3)

    return C_t_θ1, C_t_θ2, C_t_θ3
end

"""
    get_C_t_θdot(θ)

Calculate the derivative of the time derivative of the Wiener-Milenkovic
transformation matrix `C` with respect to each of the time derivatives of `θ`.
"""
@inline function get_C_t_θdot(θ)

    Cdot_θdot1 = ForwardDiff.derivative(cdot1 -> GXBeam.get_C_t(θ, SVector(cdot1, 0, 0)), 0)
    Cdot_θdot2 = ForwardDiff.derivative(cdot2 -> GXBeam.get_C_t(θ, SVector(0, cdot2, 0)), 0)
    Cdot_θdot3 = ForwardDiff.derivative(cdot3 -> GXBeam.get_C_t(θ, SVector(0, 0, cdot3)), 0)

    return Cdot_θdot1, Cdot_θdot2, Cdot_θdot3
end

"""
    get_Q(θ)

Calculate the matrix Q as defined in the paper "Geometrically nonlinear analysis
of composite beams using Wiener-Milenković parameters" by Qi Wang and Wenbin Yu
given the rotational parameters in `θ`.
"""
@inline function get_Q(θ)

    scaling = rotation_parameter_scaling(θ)

    c = scaling*θ

    c0 = 2 - c'*c/8

    return 1/(4-c0)^2*((4 - 1/4*c'*c)*I - 2*tilde(c) + 1/2*c*c')
end

"""
    get_Q_θ(θ)
    get_Q_θ(Q, θ)

Calculate the derivative of the matrix `Q` with respect to each of the rotation
parameters in `θ`.
"""
@inline get_Q_θ(θ) = get_Q_θ(get_Q(θ), θ)

@inline function get_Q_θ(Q, θ)

    scaling = rotation_parameter_scaling(θ)
    scaling_θ = rotation_parameter_scaling_θ(scaling, θ)

    c = scaling*θ
    c_θ = scaling_θ*θ' + scaling*I3

    c0 = 2 - c'*c/8
    c0_c = -c/4

    tmp = 1/(4-c0)^2
    tmp_c = 2/(4-c0)^3*c0_c

    Q_c1 = tmp_c[1]*Q/tmp + tmp*(-c[1]/2*I - 2*tilde(e1) + (e1*c' + c*e1')/2)

    Q_c2 = tmp_c[2]*Q/tmp + tmp*(-c[2]/2*I - 2*tilde(e2) + (e2*c' + c*e2')/2)

    Q_c3 = tmp_c[3]*Q/tmp + tmp*(-c[3]/2*I - 2*tilde(e3) + (e3*c' + c*e3')/2)

    Q_θ1 = Q_c1*c_θ[1,1] + Q_c2*c_θ[2,1] + Q_c3*c_θ[3,1]
    Q_θ2 = Q_c1*c_θ[1,2] + Q_c2*c_θ[2,2] + Q_c3*c_θ[3,2]
    Q_θ3 = Q_c1*c_θ[1,3] + Q_c2*c_θ[2,3] + Q_c3*c_θ[3,3]

    return Q_θ1, Q_θ2, Q_θ3
end

"""
    get_Qinv(θ)

Calculate the matrix inverse `Qinv` as defined in the paper "Geometrically
nonlinear analysis of composite beams using Wiener-Milenković parameters" by
Qi Wang and Wenbin Yu given the rotational parameters in `θ`.
"""
@inline function get_Qinv(θ)

    scaling = rotation_parameter_scaling(θ)

    c = scaling*θ

    Qinv = (1 - 1/16*c'*c)*I3 + 1/2*tilde(c) + 1/8*c*c'

    return Qinv
end

"""
    get_Qinv_θ(θ)

Calculate the derivative of the matrix inverse `Qinv` with respect to each of the
rotation parameters in `θ`.
"""
@inline function get_Qinv_θ(θ)

    scaling = rotation_parameter_scaling(θ)
    scaling_θ = rotation_parameter_scaling_θ(scaling, θ)

    c = scaling*θ
    c_θ = scaling_θ*θ' + scaling*I3

    Qinv_c1 = -c[1]/8*I3 + 1/2*tilde(e1) + 1/8*(e1*c' + c*e1')
    Qinv_c2 = -c[2]/8*I3 + 1/2*tilde(e2) + 1/8*(e2*c' + c*e2')
    Qinv_c3 = -c[3]/8*I3 + 1/2*tilde(e3) + 1/8*(e3*c' + c*e3')

    Qinv_θ1 = Qinv_c1*c_θ[1,1] + Qinv_c2*c_θ[2,1] + Qinv_c3*c_θ[3,1]
    Qinv_θ2 = Qinv_c1*c_θ[1,2] + Qinv_c2*c_θ[2,2] + Qinv_c3*c_θ[3,2]
    Qinv_θ3 = Qinv_c1*c_θ[1,3] + Qinv_c2*c_θ[2,3] + Qinv_c3*c_θ[3,3]

    return Qinv_θ1, Qinv_θ2, Qinv_θ3
end

"""
    mul3(A_1, A_2, A_3, b)

Return the product of a 3x3x3 tensor represented by `A_1`, `A_2`, and `A_3` with
the vector `b`.
"""
@inline mul3(A_1, A_2, A_3, b) = hcat(A_1*b, A_2*b, A_3*b)

"""
    gauss_quadrature(f, a, b)

Default gauss-quadrature function used for integrating distributed loads.
"""
@inline function gauss_quadrature(f, a, b)
    h = b - a
    c = (a + b)/2
    x = h/2*GAUSS_NODES .+ c
    return h/2*GAUSS_WEIGHTS'*f.(x)
end

# this function is not used
# @inline function linf_norm_scaling(A)
#
#     droptol!(A, eps(eltype(A)))
#
#     m, n = size(A)
#     rows = rowvals(A)
#     vals = nonzeros(A)
#
#     # initialize outputs
#     rowmax = fill(typemin(eltype(A)), n)
#     rowmin = fill(typemax(eltype(A)), n)
#     @inbounds for j = 1:n
#         @inbounds for i in nzrange(A, j)
#            row = rows[i]
#            val = abs(vals[i])
#            # perform sparse wizardry...
#            if val > rowmax[row]
#                rowmax[row] = val
#            end
#            if val < rowmin[row]
#                rowmin[row] = val
#            end
#         end
#     end
#
#     # use storage already allocated for rowmax for r
#     r = rowmax
#     @inbounds for i = 1:m
#         if i in rows
#             r[i] = nextpow(2, 1/sqrt(rowmax[i]*rowmin[i]))
#         else
#             r[i] = 1
#         end
#     end
#
#     # use storage already allocated for rowmin for s
#     s = rowmin
#     @inbounds for j = 1:n
#         colrange = nzrange(A, j)
#         if isempty(colrange)
#             s[j] = 1
#         else
#             colmax = maximum(i -> r[rows[i]]*abs(vals[i]), colrange)
#             colmin = minimum(i -> r[rows[i]]*abs(vals[i]), colrange)
#             s[j] = nextpow(2, 1/sqrt(colmax*colmin))
#         end
#     end
#
#     return r, s
# end


using Clines, SimpleDrawing, Plots

function originals()
    X = Line(-1, 7)

    C1 = Circle(1, 1, 1)
    C2 = Circle(2.5, 1, 0.5)
    C3 = Circle(4, 1, 1)

    pts = [1, 2 + im, 3 + im]
    SS = Cline(pts...)

    return [X, C1, C2, C3, SS]
end

function fourway(F::LFT)
    x = originals()
    y = F.(x)
    newdraw()
    draw.(y[1:4])
    try
        draw(y[5], linestyle = :dash)
    catch
        draw(y[5])
    end
    finish()
end

function fourway(z::Number = 2 + 4im)
    F = LFT(0, 1, 1, -z)
    fourway(F)
end


# -----------------------------------------------------------
#
# Evaluating the 1D Galerkin Basis functions
#
# -----------------------------------------------------------

const K_max = 10;

# Efficiency criticality: HIGH
# Central to the efficiency of the package

# Performing fastmath on an array of coefficients
# of x^k and |x|^k, giving the value at a specified point
# supported on [-1, 1]
@fastmath function array2poly{T<:Real}(v::Array{T,1}, x::Real)
	if abs(x)>1
		return zero(T)
	end
	n = length(v)
	k = div(n,2)
	s = zero(T)
	xi = one(T)
	sgn = sign(x)
	@inbounds for i in k:-1:1
		# Using Horner's method
		s *= x
		s += v[i] + flipsign(v[i+k], x)#v[i+k]*sign(x)#
	end
	return s
end

function array2poly{T<:Real}(v::Array{T,1})
	return (x-> array2poly(v,x))
end

#precompute Legendre
leg_coeffs = legendre(K_max)

# Legendre polynomial supported on [-1, 1]
function LegendreP(k, x)
	k<=K_max || throw(DomainError())
	return array2poly(leg_coeffs[k+1], x)
end

function LegendreP(k)
	k<=K_max || throw(DomainError())
	return array2poly(leg_coeffs[k+1])
end

# precomputing the DG functions
# TODO Make this a 2D array:
dg_coeffs = Array{Array{Array{Float64,1},1}}(K_max)

for i in 1:K_max
	dg_coeffs[i] = dg_basis(i)
end

# This is the dg basis function corresponding a given mode
# supported on [-1, 1]
function h(k, mode, x)
	mode<=k || throw(DomainError())
	return array2poly((dg_coeffs[k])[mode],x)
end

function h(k, mode)
	mode<=k || throw(DomainError())
	return array2poly((dg_coeffs[k])[mode])
end

# ----------------------------------------------
#
# Methods relating to the 1D Position Basis:
#
# ----------------------------------------------

# Legendre polynomial on [0, 1]
function leg{T<:Real}(mode::Int, x::T)
	return sqrt(2.0)*LegendreP(mode-1, 2*x-1)
end

# level >= 1
# cell in 1:1<<level
# mode in 1:k
# Legendre polynomial on [(cell-1)*h, cell*h] with h = 1/(1<<level)
function basis{T<:Real}(level::Int, cell::Int, mode::Int, x::T)
	return leg(mode, (1<<level)*x - (cell-1)) * (2.0)^(level/2)
end

function basis(level::Int, cell::Int, mode::Int)
	return x->basis(level, cell, mode, x)
end

# Perform numerical integration to get position basis coeffs
#
# Never used except if user demands to evolve wave equation
# using only the position basis, never going through heir
function pos_vcoeffs_DG(k::Int, level::Int, f::Function;
						rel_tol = REL_TOL, abs_tol = ABS_TOL, max_evals = MAX_EVALS)
	vcoeffs = Array{Float64}((1<<level)*(k))
	i = 1
	for cell in 1:(1<<level)
		for mode in 1:k
			fcn			= x->(basis(level, cell, mode,x)*f(x))
			left_bound  = (cell-1)/(1<<level)
			right_bound = (cell)/(1<<level)
			vcoeffs[i]  = hquadrature(fcn, left_bound, right_bound; abstol=abs_tol)[1]
			i += 1
		end
	end
	return vcoeffs
end

# ----------------------------------------------
# Building the hier2pos matrix:
# ----------------------------------------------

# Given an array of the type above, of coefficients
# for both both x^k and |x|^k, then given a choice of
# side (right vs. left) this becomes just a pure
# polynomial array of half the length
function convert_polyarray{T<:Real}(v::Array{T,1}, side="left")
    n = Int(length(v)//2)
	if side == "left"
		return [v[i]-v[n+i] for i in 1:n]
	elseif side == "right"
		[v[i]+v[n+i] for i in 1:n]
	else throw(ArgumentError())
	end
end

# Gives the polynomial array corresponding to p(x/c)
function scale_polyarray{T<:Real}(v::Array{T,1}, c::Real)
	n = length(v)
	return [v[i] * (1//c^(i-1)) for i in 1:n]
end

# Gives the polynomial array corresponding to p(x-a)
function shift_polyarray{T<:Real}(v::Array{T,1}, a::Real)
	n = length(v)
	v_new = zeros(v)
	for i::Int in 1:n
		for j::Int in 1:i
			v_new[j] += v[i] * (-a)^(i-j) * binomial(i-1, j-1)
		end
	end
	return v_new
end

# Performs symbolic integration of the polynomials
# represented by v and w on the interval [a, b]
function integrate_polyarray{T<:Real}(v::Array{T,1}, w::Array{T,1}; a::Real=0, b::Real=1)
	ans = zero(T)
	for i::Int in 1:length(v)
		for j::Int in 1:length(w)
			ans += v[i] * w[j] * ((b^(i+j-1)-a^(i+j-1))//(i+j-1))
		end
	end
	return ans
end

# With a bit of case work, we can do an explicit
# symbolic integration of the elements in the
# hierarchical basis against the legendre polynomials
# defining the position basis:
function hier2pos(k::Int, max_level::Int, level::Int, cell::Int, mode::Int)

	ans = Real[]

	# Compute the support and discontinuity
	# of the hierarchical basis element
	left_point  = (cell-1)//(1<<max(0, level-1))
	mid_point   = (cell- 1//2)//(1<<max(0, level-1))
	right_point = cell//(1<<max(0, level-1))

	# Level 0 of the hierarchical basis is a continous
	# Legendre polynomial
	if level == 0
		vl = vr = sqrt(2)*shift_polyarray(
								scale_polyarray(leg_coeffs[mode][1:k],
												1//2),
											1//2)

	# The higher DG levels are given by two polynomials
	# defined on neighboring (right & left) intervals of size h/2,
	# which would need to be integrated against individually
	else
		h = 1//(1<<max(0, level-1))
		vl = shift_polyarray(
				scale_polyarray(
					convert_polyarray(dg_coeffs[k][mode], "left"),
								1//2),
							1//2)
		vl = shift_polyarray(scale_polyarray(vl, h), (cell-1)*h)*sqrt(1<<(level))

		vr = shift_polyarray(
				scale_polyarray(
					convert_polyarray(dg_coeffs[k][mode], "right"),
					 			1//2),
							1//2)
		vr = shift_polyarray(scale_polyarray(vr, h), (cell-1)*h)*sqrt(1<<(level))
	end

	h = 1//(1<<max_level)
	for i in 1:1<<max_level
		for p in 1:k
			# Check if our hier function is supported on this interval
			if (i-1)*h < left_point || i*h > right_point
				push!(ans, 0)
				continue
			end

			# Obtain the polynomial array for the position basis function
			unit_leg = shift_polyarray(
							scale_polyarray(leg_coeffs[p][1:k], 1//2), 1//2)
			pos_element = shift_polyarray(
								scale_polyarray(unit_leg, h), (i-1)*h)
			pos_element *= sqrt(1<<(max_level+1)) # (2.0)^((max_level+1)/2)

			# Perform symbolic integration depending on which
			# interval the position basis function overlaps with the
			# DG basis element
			if (i-1)*h < mid_point
				val = integrate_polyarray(vl, pos_element; a=(i-1)*h, b=i*h)
			else
				val = integrate_polyarray(vr, pos_element; a=(i-1)*h, b=i*h)
			end
			push!(ans, val)
		end
	end
	return ans
end

# Construct the change of basis matrix from hierarchical
# to position basis
function hier2pos(k::Int, max_level::Int; abs_tol=ABS_TOL)
	j = 1
	I = Int[]
	J = Int[]
	V = Float64[]
	for level in 0:max_level
		for cell in 1:(1<<max(0, level-1))
			for mode in 1:k
				ans = pos_vcoeffs_DG(k, max_level, v(k, level, cell, mode))
				for i in 1:length(ans)
					if abs(ans[i]) > abs_tol
						push!(I, i)
						push!(J, j)
						push!(V, ans[i])
					end
				end
				j += 1
			end
		end
	end
	return sparse(I, J, V, k * (1<<max_level), k * (1<<max_level), +)
end

function pos2hier(k::Int, max_level::Int; abs_tol=ABS_TOL)
	return hier2pos(k, max_level; abs_tol=abs_tol)'
end


function max_digit_sum(limit)
    max_sum = 0

    for a in 1:limit - 1
        for b in 1:limit - 1
            digit_sum = sum(digits(BigInt(a) ^ b))
            if digit_sum > max_sum
                max_sum = digit_sum
            end
        end
    end

    return max_sum
end

println(max_digit_sum(100))


######### StanVariational Bernoulli example  ###########

using StanVariational, Test

bernoulli_model = "
data {
  int<lower=1> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  y ~ bernoulli(theta);
}
";

bernoulli_data = Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1])

# Use tmpdir across multiple runs to prevent re-compilation
#tmpdir = joinpath(@__DIR__, "tmp")

sm = VariationalModel("variational", bernoulli_model)

rc = stan_variational(sm; data=bernoulli_data)

if success(rc)

  (samples, cnames) = read_variational(sm)

end


using LuxurySparse

using Yao
using Yao.Intrinsics
using Random
using BenchmarkTools
using Statistics
using LinearAlgebra
using SparseArrays

vreg = randn(ComplexF64, 1<<16)
mreg = randn(ComplexF64, 1<<16, 2)

randmat(::Type{Val{:Dense}}, n) = randn(ComplexF64,n,n)
randmat(::Type{Val{:Diag}}, n) = Diagonal(randn(ComplexF64,n))
randmat(::Type{Val{:Perm}}, n) = pmrand(ComplexF64, n)
randmat(::Type{Val{:CSC}}, n) = SparseMatrixCSC(pmrand(ComplexF64, n))

vs = [vreg, mreg]

for m in [:Diag, :Dense, :Perm, :Diag]
    for v in [1,2]
        token = "$m * $(v ==1 ? "Vec" : "Mat")"
        for ng in [1,2]
            n=1<<ng
            un = randmat(Val{m}, n)
            sun = un |> staticize
            inds =  randperm(n) .+18
            sinds = SVector{n}(inds)

            if ng==1 && m == :Dense
                println("$token, u1rows($ng) => ", median(@benchmark u1rows!($(vs[v]), $(inds[1]), $(inds[2]), $(un[1]), $(un[3]), $(un[2]), $(un[4]))))
            end
            #println("$token, dynamic($ng) => ", median(@benchmark unrows!($(vs[v]), $inds, $un) seconds=1))
            res = @benchmark unrows!($(vs[v]), $sinds, $sun)
            println("$token, static($ng) => ", median(res))
            display(res)
        end
    end
end

using StaticArrays: SVector, SMatrix, SDiagonal

@inline function unrows!(state::Vector, inds::AbstractVector, A::Union{SSparseMatrixCSC, SparseMatrixCSC}, work::Vector)
    work .= 0
    @inbounds for col = 1:length(inds)
        xj = state[inds[col]]
        @inbounds @simd for j = A.colptr[col]:(A.colptr[col + 1] - 1)
            work[A.rowval[j]] += A.nzval[j]*xj
        end
    end
    state[inds] = work
    state
end

n=2
inds =  randperm(n) .+18
sinds = SVector{n}(inds)
un = Diagonal(randn(ComplexF64, n))
#un = randmat(Val{:CSC}, n)
sun = un |> staticize;

work = Vector{ComplexF64}(undef, n)
#@benchmark unrows!($vreg, $sinds, $sun, $work)
@benchmark unrows!($mreg, $sinds, $sun)


using StatsBase

function random_subspaces(n_subspaces, min_dim = 1, max_dim = 2; min_ss_size = 2, max_ss_size=2)
    @assert min_dim < max_dim && min_dim >= 1
    @assert min_ss_size <= max_ss_size && min_ss_size >=2
    @assert max_ss_size <= max_dim - min_dim + 1

    max_possible_subspaces = sum([binomial(max_dim - min_dim + 1, k) for k in min_ss_size:max_ss_size])
    n_subspaces > max_possible_subspaces && error("Maximum number subspaces within the specification (min/max dim and ss_size) is $max_possible_subspaces,
        but method is called with n_subspaces = $n_subspaces.")
    subspaces = Set{Vector{Int}}()
    while length(subspaces) < n_subspaces
        subspace_size = rand(min_ss_size:max_ss_size)
        push!(subspaces, sort(sample(min_dim:max_dim, subspace_size; replace=false), rev=false))
    end
    return [s for s in subspaces]
end

function generate_randsub_from_data(n::Int; kwargs...)
    (data::Array{Float64, 2}) -> random_subspaces(n, 1, size(data, 1); kwargs...)
end


using PolyChaos, Test

deg = 5
op = genLaguerreOrthoPoly(deg,1.23)

n = 10
X = [ exp(3 + rand()), exp.(3 .+ rand(n)) ]

@testset "evaluate univariate basis at point(s)" begin
    for x in X
        @test evaluate(collect(0:deg),x,op.α,op.β) == evaluate(collect(0:deg),x,op)
        for d in 0:deg
            if typeof(x) <: Real
                @test evaluate(d,x,op.α,op.β) == evaluate(d,x,op) == evaluate(d,[x],op.α,op.β)
            else
                @test evaluate(d,x,op.α,op.β) == evaluate(d,x,op)
            end
        end

        res = hcat(map(i->evaluate(i,x,op.α,op.β),collect(0:op.deg))...)
        if typeof(x) <: Real
            @test evaluate(x,op) == evaluate([x], op) == evaluate(collect(0:op.deg),[x],op) == evaluate(collect(0:op.deg),[x],op.α,op.β) == res
        else
            @test evaluate(x, op) == evaluate(collect(0:op.deg),x,op) == evaluate(collect(0:op.deg),x,op.α,op.β) == res
        end
    end
end

ops = [  genLaguerreOrthoPoly(deg,1.23),
        genHermiteOrthoPoly(deg,3.),
        LegendreOrthoPoly(deg),
        JacobiOrthoPoly(deg,4.3,10.)
     ]

mop = MultiOrthoPoly(ops,deg)

nop = length(ops)
X = [ exp.(3 .+ rand(nop)) , exp.(3 .+ rand(n,nop)) ]

@testset "evaluate multivariate basis at point(s)" begin
    for x in X
        @test evaluate(mop.ind,x,coeffs(mop)...) == evaluate(mop.ind,x,mop) == evaluate(x,mop)
        if typeof(x) <: Vector{<:Real}
            @test evaluate(mop.ind,x,coeffs(mop)...) == evaluate(mop.ind,x,mop) == evaluate(x,mop) == evaluate(mop.ind,reshape(x,1,length(x)),coeffs(mop)...)
            @test evaluate(x,mop) == evaluate(mop.ind,reshape(x,1,length(x)),mop) == evaluate(mop.ind,reshape(x,1,length(x)),coeffs(mop)...)
        else
            @test evaluate(x,mop) == evaluate(mop.ind,x,mop) == evaluate(mop.ind,x,coeffs(mop)...)
        end
        for d in eachrow(mop.ind[1:3,:])
            @test evaluate(collect(d),x,coeffs(mop)...) == evaluate(collect(d),x,mop)
        end
    end
end


function rand_poisson(lambda)
	x::Int64 = 0
	t::Float64 = 0.0

	while t < lambda
		x = x + 1
		t = t + randexp()
	end

	return x - 1
end


using Equate
using Statistics, StatsBase

# methods = [
#     "Tucker"
#     "LevineCongeneric"
#     "ChainedLinear"
#     "FrequencyEstimation"
#     "BraunHolland"
#     "ChainedEquipercentile"
# ]

@testset "Tucker" begin
    res = Tucker(neattabX, neattabY)
    @test round(coef(res)[1]; digits = 4) == 1.0292
    @test round(coef(res)[2]; digits = 4) == .5378
end

@testset "LevineObservedScore (external, not show in K&B) " begin
    res = LevineObservedScore(neattabX, neattabY; common = :external)
    @test round(coef(res)[1]; digits = 4) == 1.0168
    @test round(coef(res)[2]; digits = 4) == .3257
end

@testset "LevineObservedScore (internal) " begin
    res = LevineObservedScore(neattabX, neattabY; common = :internal)
    @test round(coef(res)[1]; digits = 4) == 1.0110
    @test round(coef(res)[2]; digits = 4) == .2514
end

@testset "LevineTrueScore (external, not show in K&B) " begin
    res = LevineTrueScore(neattabX, neattabY; common = :external)
    @test round(coef(res)[1]; digits = 4) == 1.0165
    @test round(coef(res)[2]; digits = 4) == .331
end

@testset "LevineTrueScore (internal) " begin
    res = LevineTrueScore(neattabX, neattabY; common = :internal)
    @test round(coef(res)[1]; digits = 4) == 1.0086
    @test round(coef(res)[2]; digits = 4) == .2912
end

@testset "ChainedLinear" begin
    res = ChainedLinear(neattabX, neattabY)
    @test round(coef(res)[1]; digits = 4) == 1.0213
    @test round(coef(res)[2]; digits = 4) == .3937
end

# Invoke the percentile rank methods
# The results in Kolen & Brennan seems to be wrong.
@testset "BraunHolland (R-equate)" begin
    res = BraunHolland(neattabX, neattabY)
    @test round(coef(res)[1]; digits = 4) == 1.0067
    @test round(coef(res)[2]; digits = 4) == 0.8976
    # popXonY = ExpandTable(res.table.lYx, neattabX.tabX.freq)
    # @test round(mean(popXonY); digits = 4) == 16.8329
    # @test round(std(popXonY; corrected = false); digits = 4) == 6.6017
    # @test round(skewness(popXonY); digits = 4) == .5799
end

# @testset "BraunHolland (Kolen & Brennan)" begin
#     res = BraunHolland(dummytabX, dummytabY)
#     @test round(coef(res)[1]; digits = 4) == 1.0331
#     @test round(coef(res)[2]; digits = 4) == -.1927
#     popXonY = ExpandTable(res.table.lYx, Int64.(dummytabX.tabX.freq))
#     @test round(mean(popXonY); digits = 4) == 16.8329
#     @test round(std(popXonY; corrected = false); digits = 4) == 6.6017
#     @test round(skewness(popXonY); digits = 4) == .5799
# end

@testset "FrequencyEstimation" begin
    res = FrequencyEstimation(dummytabX, dummytabY; case = :upper, w₁ = 1.0)
    @test round.(res.table.eYx; digits = 2) == [-0.02, 0.83, 1.76, 2.92, 3.98, 5.00]
end

@testset "ChainedEquipercentile" begin
    res = ChainedEquipercentile(neattabX, neattabY; case = :upper)
    # Largest 3 values in the concordance table are matched.
    @test round.(res.table.eYx[end-2:end]; digits = 4) == [34.3125, 35.4125, 36.0938]
end


abstract type TullyModel <: DiabaticModel end

NQCModels.ndofs(::TullyModel) = 1
NQCModels.nstates(::TullyModel) = 2

"""
    TullyModelOne(a=0.01, b=1.6, c=0.005, d=1.0)

Tully's simple avoided crossing model from [J. Chem. Phys. 93, 1061 (1990)](https://doi.org/10.1063/1.459170).
"""
Parameters.@with_kw struct TullyModelOne{A,B,C,D} <: TullyModel
    a::A = 0.01
    b::B = 1.6
    c::C = 0.005
    d::D = 1.0
end

function NQCModels.potential(model::TullyModelOne, q::Real)
    Parameters.@unpack a, b, c, d = model
    if q > 0
        V11 = a * (1 - exp(-b*q))
    else
        V11 = -a * (1 - exp(b*q))
    end
    V22 = -V11
    V12 = c * exp(-d*q^2) # sets both off-diagonals as V is Hermitian
    return Hermitian(SMatrix{2,2}(V11, V12, V12, V22))
end

function NQCModels.derivative(model::TullyModelOne, q::Real)
    Parameters.@unpack a, b, c, d = model
    D11 = a * b * exp(-b * abs(q))
    D22 = -D11
    D12 = -2 * c * d * q * exp(-d*q^2)
    return Hermitian(SMatrix{2,2}(D11, D12, D12, D22))
end

"""
    TullyModelTwo(a=0.1, b=0.28, c=0.015, d=0.06, e=0.05)

Tully's dual avoided crossing model from [J. Chem. Phys. 93, 1061 (1990)](https://doi.org/10.1063/1.459170).
"""
Parameters.@with_kw struct TullyModelTwo{A,B,C,D,E} <: TullyModel
    a::A = 0.1
    b::B = 0.28
    c::C = 0.015
    d::D = 0.06
    e::E = 0.05
end

function NQCModels.potential(model::TullyModelTwo, q::Real)
    Parameters.@unpack a, b, c, d, e = model
    V11 = 0
    V22 = -a*exp(-b*q^2) + e
    V12 = c * exp(-d*q^2)
    return Hermitian(SMatrix{2,2}(V11, V12, V12, V22))
end

function NQCModels.derivative(model::TullyModelTwo, q::Real)
    Parameters.@unpack a, b, c, d  = model
    D11 = 0
    D22 = 2*a*b*q*exp(-b*q^2)
    D12 = -2*c*d*q*exp(-d*q^2)

    Hermitian(SMatrix{2,2}(D11, D12, D12, D22))
end

"""
    TullyModelThree(a=0.0006, b=0.1, c=0.9)

Tully's extended coupling with reflection model from [J. Chem. Phys. 93, 1061 (1990)](https://doi.org/10.1063/1.459170).
"""
Parameters.@with_kw struct TullyModelThree{A,B,C} <: TullyModel
    a::A = 0.0006
    b::B = 0.1
    c::C = 0.9
end

function NQCModels.potential(model::TullyModelThree, q::Real)
    Parameters.@unpack a, b, c = model
    V11 = a
    V22 = -a
    if q > 0
        V12 = b * (2 - exp(-c*q))
    else
        V12 = b * exp(c*q)
    end
    return Hermitian(SMatrix{2,2}(V11, V12, V12, V22))
end

function NQCModels.derivative(model::TullyModelThree, R::Real)
    Parameters.@unpack a, b, c = model
    q = R[1]
    D11 = 0
    D22 = 0
    if q > 0
        D12 = b * c * exp(-c*q)
    else
        D12 = b * c * exp(c*q)
    end
    return Hermitian(SMatrix{2,2}(D11, D12, D12, D22))
end


# # One-dimensional cellular automata
# Agents.jl provides a module (CA1D) to create and plot 1D cellular automata.

using Agents
using Agents.CA1D
using AgentsPlots

# ## 1. Define the rule
# Here is Wolfram's rule 22

rules = Dict("111"=>"0", "110"=>"0", "101"=>"0", "100"=>"1", "011"=>"0",
            "010"=>"1", "001"=>"1", "000"=>"0")

# ## 2. Build the model
# All the cells are by default initially "off"
model = CA1D.build_model(rules=rules, ncols=21)

# This turns on the middle cell
model.agents[11].status="1"

# ## 3. Run the model and collect data.
runs = 20
data = CA1D.ca_run(model, runs);

# ## 4. Visualize the data
p = plot_CA1D(data, nodesize=2)


using JuMP
using CutPruners
using StructJuMP
using StructDualDynProg
using Compat, Compat.LinearAlgebra, Test
using StochOptInterface
const SOI = StochOptInterface

# solver independent tests
include("comp.jl")
include("stopcrit.jl")
include("paths.jl")

# load a solver
include("solvers.jl")

# solver dependent tests
include("optimize_stock.jl")
include("hydro_thermal_scheduling.jl")
include("prob5.2.jl")


"""
SamplerState for stochastic gradient Hamiltonian Monte Carlo following

Chen, T., Fox, E., & Guestrin, C. (2014). Stochastic Gradient Hamiltonian Monte Carlo. In Proceedings of The 31st International Conference on Machine Learning (pp. 1683–1691).

Fields:

    - x: the current state
    - p: current momentum
    - t: iterations counter
    - stepsize: s.stepsize(s.t) gives the stepsize to be used in iteration s.t
    - mass: mass matrix
    - D: friction parameters
"""

type SGHMCState <: SamplerState
    x::Array{Float64}
    p::Array{Float64}

    t::Int64
    stepsize::Function
    mass::Array{Float64,1}
    D::Array{Float64,1}
    """
    Constructs an SGHMCState given an initial state, a momentum and a fixed stepsize
    """
    function SGHMCState(x::Array{Float64,1}, p::Array{Float64,1}, stepsize::Float64;  mass=ones(length(x)), D=[1.])
    # momentum, mass and state need to have the same size.
    @assert size(x) == size(p)
    @assert size(mass) == size(x)
    # D matrix should be scalar or have the same size as x
    @assert size(x) == size(D) || length(D) == 1

    # stepsize is used as a function
    s(niters) = stepsize
    new(x,p,0,s,mass,D)
    end
    """
    Constructs an SGHMCState given an initial state and a fixed stepsize
    """
    function SGHMCState(x::Array{Float64,1}, stepsize::Float64;  mass=ones(length(x)), D=[1.])
    # mass and
    @assert size(mass) == size(x)
    # D matrix should be scalar or have the same size as x
    @assert size(x) == size(D) || length(D) == 1
    #stepsize is used as a function
    s(niters) = stepsize
    new(x,sqrt(mass).*randn(length(x)),0,s,mass,D)
    end
    """
    Constructs an SGHMCState given an initial state, a momentum and a stepsize function
    """
    function SGHMCState(x::Array{Float64,1}, p::Array{Float64,1}, stepsize::Function;  mass=ones(length(x)), D=[1.])
    # momentum, mass and state need to have the same size.
    @assert size(x) == size(p)
    @assert size(mass) == size(x)
    # D matrix should be scalar or have the same size as x
    @assert size(x) == size(D) || length(D) == 1

    new(x,p,0,stepsize,mass,D)
    end
    """
    Constructs an SGHMCState given an initial state and a stepsize function
    """
    function SGHMCState(x::Array{Float64,1}, stepsize::Function;  mass=ones(length(x)), D=[1.])
    # mass and
    @assert size(mass) == size(x)
    # D matrix should be scalar or have the same size as x
    @assert size(x) == size(D) || length(D) == 1
    #stepsize is used as a function
    new(x,sqrt(mass).*randn(length(x)),0,stepsize,mass,D)
    end
end

"""
sample! performs one SGHMC update and returns the updated SGHMCState.
"""
function sample!(s::SGHMCState, grad)
    #increase iteration counter
    s.t += 1

    # improve readability
    x = s.x
    p = s.p
    D = s.D
    t = s.t
    stepsize = s.stepsize(t)
    m = s.mass

    # inplace updates
    add_inplace!(stepsize,-D.*p./m,p)
    add_inplace!(stepsize,grad(x),p)
    add_inplace!(sqrt(2*stepsize),sqrt(D).*randn(length(x)),p)
    add_inplace!(stepsize,p ./ m,x)
    s
end


using BasisFunctions, BasisFunctions.Test, DomainSets, DoubleFloats

using Test

types = (Float64, LargeFloat)

#####
# Orthogonal polynomials
#####

function test_chebyshevT(T)
    println("- ChebyshevT polynomials ($T)")
    bc = ChebyshevT{T}(20)
    test_ops_generic(bc)
    x1 = T(4//10)
    @test bc[4](x1) ≈ cos(3*acos(x1))

    @test support(bc) == ChebyshevInterval{T}()

    n1 = 160
    b1 = ChebyshevT{T}(n1)

    n1 = 160
    b1 = ChebyshevT{T}(n1)
    f = exp
    e = approximate(b1, exp)
    x0 = T(1//2)
    @test abs(e(T(x0))-f(x0)) < sqrt(eps(T))

    @test hastransform(b1)
    @test hastransform(b1, interpolation_grid(b1))
    @test hastransform(b1, ChebyshevExtremae(n1))
    @test hastransform(b1, ChebyshevNodes(n1))
    @test ~hastransform(b1, PeriodicEquispacedGrid(n1, -1, 1))
    @test ~hastransform(b1, ChebyshevExtremae(n1+5))
    @test ~hastransform(b1, ChebyshevNodes(n1+5))
    @test BasisFunctions.jacobi_α(b1) ≈ -.5 ≈ BasisFunctions.jacobi_β(b1)
end

function test_chebyshevU(T)
    println("- ChebyshevU polynomials ($T)")
    n1 = 20
    b1 = ChebyshevU{T}(n1)
    test_ops_generic(b1)
    x1 = T(4//10)
    @test b1[4](x1) ≈ sin(4*acos(x1))/sin(acos(x1))

    @test support(b1) == ChebyshevInterval{T}()
    @test BasisFunctions.jacobi_α(b1) ≈ .5 ≈ BasisFunctions.jacobi_β(b1)

    test_orthogonality_orthonormality(b1, true, false, BasisFunctions.ChebyshevUWeight{T}())
    test_orthogonality_orthonormality(b1, true, false, gauss_rule(b1))
    test_orthogonality_orthonormality(b1, true, false, gauss_rule(resize(b1,length(b1)+1)))
end

function test_legendre(T)
    println("- Legendre polynomials ($T)")
    bl = Legendre{T}(10)
    test_ops_generic(bl)
    x1 = T(4//10)
    @test abs(bl[6](x1) - 0.27064) < 1e-5

    test_orthogonality_orthonormality(bl, true, false, BasisFunctions.LegendreWeight{T}())
    test_orthogonality_orthonormality(bl, true, false, gauss_rule(bl))
    test_orthogonality_orthonormality(bl, true, false, gauss_rule(resize(bl,length(bl)+1)))
end

function test_laguerre(T)
    println("- Laguerre polynomials ($T)")
    bl = Laguerre(10, T(1//3))
    test_ops_generic(bl)
    x1 = T(4//10)
    @test abs(bl[6](x1) + 0.08912346) < 1e-5

    test_orthogonality_orthonormality(bl, true, false, BasisFunctions.LaguerreWeight(T(1//3)))
    test_orthogonality_orthonormality(Laguerre(10,T(0)), true, true, BasisFunctions.LaguerreWeight(T(0)))
    test_orthogonality_orthonormality(bl, true, false, gauss_rule(bl))
    test_orthogonality_orthonormality(bl, true, false, gauss_rule(resize(bl,length(bl)+1)))
end

function test_hermite(T)
    println("- Hermite polynomials ($T)")
    bh = Hermite{T}(10)
    test_ops_generic(bh)
    x1 = T(4//10)
    @test abs(bh[6](x1) - 38.08768) < 1e-5

    test_orthogonality_orthonormality(bh, true, false, BasisFunctions.HermiteWeight{T}())
    test_orthogonality_orthonormality(bh, true, false, gauss_rule(bh))
    test_orthogonality_orthonormality(bh, true, false, gauss_rule(resize(bh,length(bh)+1)))
end

function test_jacobi(T)
    println("- Jacobi polynomials ($T)")
    bj = Jacobi(10, T(2//3), T(3//4))
    test_ops_generic(bj)
    x1 = T(4//10)
    @test abs(bj[6](x1) - 0.335157) < 1e-5
    @test BasisFunctions.jacobi_α(bj) ≈ T(2//3)
    @test BasisFunctions.jacobi_β(bj) ≈ T(3//4)

    test_orthogonality_orthonormality(bj, true, false, BasisFunctions.JacobiWeight(T(2//3), T(3//4)))
    test_orthogonality_orthonormality(bj, true, false, gauss_rule(bj))
    test_orthogonality_orthonormality(bj, true, false, gauss_rule(resize(bj,length(bj)+1)))
end

function test_ops_generic(ops)
    T = codomaintype(ops)
    tol = test_tolerance(T)

    x = fixed_point_in_domain(ops)
    z1 = BasisFunctions.unsafe_eval_element(ops, length(ops), x)
    z2 = recurrence_eval(ops, length(ops), x)
    @test abs(z1-z2) < tol
    a,b = monic_recurrence_coefficients(ops)
    z3 = monic_recurrence_eval(a, b, length(ops), x)
    γ = leading_order_coefficient(ops, length(ops))
    @test abs(z1 - γ*z3) < tol

    d1 = BasisFunctions.unsafe_eval_element_derivative(ops, length(ops), x, 1)
    d2 = recurrence_eval_derivative(ops, length(ops), x)
    @test abs(d1-d2) < tol

    # if prectype(ops) == Float64
        # We only do these tests for Float64 because eig currently does not support BigFloat
        r = roots(ops)
        @test maximum(abs.(BasisFunctions.unsafe_eval_element.(Ref(ops), length(ops)+1, r))) < 100tol

        m = gauss_rule(ops)
        x = points(m)
        w = BasisFunctions.weights(m)
        @test abs(sum(w) - first_moment(ops)) < tol
    # end
end

# for T in types
for T in (Float64, LargeFloat)
    @testset "Orthogonal polynomials ($T)" begin
        println("Classical orthogonal polynomials ($T):")
        test_chebyshevT(T)
        test_chebyshevU(T)
        test_legendre(T)
        test_laguerre(T)
        test_hermite(T)
        test_jacobi(T)
    end
    println()
end

@testset "Orthogonality of orthogonal polynomials" begin
    OPSs = [ChebyshevT, ChebyshevU, Legendre, Hermite, Jacobi, Laguerre]
    for ops in OPSs, n in (5,6), T in (Float64,LargeFloat)
        B = ops{T}(n)
        test_orthogonality_orthonormality(B, gauss_rule(B))
        test_orthogonality_orthonormality(B, gauss_rule(resize(B,2n)))
        test_orthogonality_orthonormality(B, gauss_rule(resize(B,n-1)))
    end
end


using CollectiveSpins

using PyCall
@pyimport matplotlib.pyplot as plt

function plot_geometry_xy(positions)
    for (i,p)=enumerate(positions)
        plt.text(p[1], p[2], "($i)", color="black", va="center", ha="center")
    end
    xmin = minimum([p[1] for p=positions])
    xmax = maximum([p[1] for p=positions])
    ymin = minimum([p[2] for p=positions])
    ymax = maximum([p[2] for p=positions])
    dx = xmax - xmin
    dy = ymax - ymin
    marginx = max(dx/5, 0.1)
    marginy = max(dy/5, 0.1)
    plt.xlim([xmin-marginx, xmax+marginx])
    plt.ylim([ymin-marginy, ymax+marginy])
end

d = 0.1
geomN = 5
positions = geometry.hexagonal(d; Nr=geomN)
positions = geometry.chain(d, 2*geomN+1)
plot_geometry_xy(positions)

plt.show()