module Rotations

using StaticArrays
using LinearAlgebra: norm, rank, svd, Diagonal
using SHIPs: _mrange, IntS
using SHIPs.SphericalHarmonics: index_y

export ClebschGordan, CoeffArray, single_B

"""
`ClebschGordan: ` storing precomputed Clebsch-Gordan coefficients; see
`?clebschgordan` for the convention that is use.
"""
struct ClebschGordan{T}
	# maxL::Int
	# cg::Array{T, 3}   # rewrite as Dict!
	vals::Dict{Tuple{IntS, IntS, IntS}, T}
end

"""
`CoeffArray: ` storing recursively precomputed coefficients for a
rotation-invariant basis.
"""
struct CoeffArray{T}
   vals::Vector{Dict}
   cg::ClebschGordan{T}
end

# ----------------------------------------------------------------------
#     ClebschGordan code
# ----------------------------------------------------------------------

cg_conditions(j1,m1,j2,m2,j3,m3) =
	cg_l_condition(j1, j2, j3)   &&
	cg_m_condition(m1, m2, m3)   &&
	(abs(m1) <= j1) && (abs(m2) <= j2) && (abs(m3) <= j3)

cg_l_condition(j1, j2, j3) = (abs(j1-j2) <= j3 <= j1 + j2)

cg_m_condition(m1, m2, m3) = (m3 == m1 + m2)

"""
`clebschgordan(j1, m1, j2, m2, j3, m3, T=Float64)` :

A reference implementation of Clebsch-Gordon coefficients based on

https://hal.inria.fr/hal-01851097/document
Equation (4-6)

This heavily uses BigInt and BigFloat and should therefore not be employed
for performance critical tasks.

The ordering of parameters corresponds to the following convention:
```
clebschgordan(j1, m1, j2, m2, j3, m3) = C_{j1m1j2m2}^{j3m3}
```
where
```
   D_{m1k1}^{l1} D_{m2k2}^{l2}}
	=
	∑_j  C_{l1m1l2m2}^{j(m1+m2)} C_{l1k1l2k2}^{j2(k1+k2)} D_{(m1+m2)(k1+k2)}^{j}
```
"""
function clebschgordan(j1, m1, j2, m2, j3, m3, T=Float64)
	if !cg_conditions(j1, m1, j2, m2, j3, m3)
		return zero(T)
	end

   N = (2*j3+1) *
       factorial(big(j1+m1)) * factorial(big(j1-m1)) *
       factorial(big(j2+m2)) * factorial(big(j2-m2)) *
       factorial(big(j3+m3)) * factorial(big(j3-m3)) /
       factorial(big( j1+j2-j3)) /
       factorial(big( j1-j2+j3)) /
       factorial(big(-j1+j2+j3)) /
       factorial(big(j1+j2+j3+1))

   G = big(0)
   # 0 ≦ k ≦ j1+j2-j3
   # 0 ≤ j1-m1-k ≤ j1-j2+j3   <=>   j2-j3-m1 ≤ k ≤ j1-m1
   # 0 ≤ j2+m2-k ≤ -j1+j2+j3  <=>   j1-j3+m2 ≤ k ≤ j2+m2
   lb = (0, j2-j3-m1, j1-j3+m2)
   ub = (j1+j2-j3, j1-m1, j2+m2)
   for k in maximum(lb):minimum(ub)
      bk = big(k)
      G += (-1)^k *
           binomial(big( j1+j2-j3), big(k)) *
           binomial(big( j1-j2+j3), big(j1-m1-k)) *
           binomial(big(-j1+j2+j3), big(j2+m2-k))
   end

   return T(sqrt(N) * G)
end

ClebschGordan(T=Float64) =
	ClebschGordan(Dict{Tuple{IntS,IntS,IntS}, T}())

_cg_key(j1, m1, j2, m2, j3, m3) =
	IntS.((index_y(j1,m1), index_y(j2,m2), index_y(j3,m3)))

function (cg::ClebschGordan{T})(j1, m1, j2, m2, j3, m3) where {T}
	if !cg_conditions(j1,m1,j2,m2,j3,m3)
		return zero(T)
	end
	key = _cg_key(j1, m1, j2, m2, j3, m3)
	if haskey(cg.vals, key)
		return cg.vals[key]
	end
	val = clebschgordan(j1, m1, j2, m2, j3, m3, T)
	cg.vals[key] = val
	return val
end

# ----------------------------------------------------------------------
#     CoeffArray code
# ----------------------------------------------------------------------

dicttype(N::Integer) = dicttype(Val(N))

dicttype(::Val{N}) where {N} =
   Dict{Tuple{SVector{N,Int8}, SVector{N,Int8}, SVector{N,Int8}}, Float64}

CoeffArray(T=Float64) = CoeffArray(Dict[], ClebschGordan(T))

function get_vals(A::CoeffArray, valN::Val{N}) where {N}
	if length(A.vals) < N
		for n = length(A.vals)+1:N
			push!(A.vals, dicttype(n)())
		end
	end
   return A.vals[N]::dicttype(valN)
end

_key(ll::StaticVector{N}, mm::StaticVector{N}, kk::StaticVector{N}) where {N} =
      (SVector{N, Int8}(ll), SVector{N, Int8}(mm), SVector{N, Int8}(kk))

function (A::CoeffArray)(ll::StaticVector{N},
                         mm::StaticVector{N},
                         kk::StaticVector{N}) where {N}
   if       sum(mm) != 0 ||
            sum(kk) != 0 ||
            !all(abs.(mm) .<= ll) ||
            !all(abs.(kk) .<= ll)
      return 0.0
   end
   vals = get_vals(A, Val(N))  # this should infer the type!
   key = _key(ll, mm, kk)
   if haskey(vals, key)
      val  = vals[key]
   else
      val = _compute_val(A, key...)
      vals[key] = val
   end
   return val
end

function (A::CoeffArray)(ll::StaticVector{1},
                         mm::StaticVector{1},
                         kk::StaticVector{1})
   if ll[1] == mm[1] == kk[1] == 0
      return 1.0
   else
      return 0.0
   end
end

function (A::CoeffArray)(ll::StaticVector{2},
                         mm::StaticVector{2},
                         kk::StaticVector{2})
   if ll[1] != ll[2] || sum(mm) != 0 || sum(kk) != 0
      return 0.0
   else
      return 8 * pi^2 / (2*ll[1]+1) * (-1)^(mm[1]-kk[1])
   end
end

function _compute_val(A::CoeffArray, ll::StaticVector{N},
                                     mm::StaticVector{N},
                                     kk::StaticVector{N}) where {N}
	val = 0.0
   llp = ll[1:N-2]
   mmp = mm[1:N-2]
   kkp = kk[1:N-2]
   for j = abs(ll[N-1]-ll[N]):(ll[N-1]+ll[N])
      if abs(kk[N-1]+kk[N]) > j || abs(mm[N-1]+mm[N]) > j
         continue
      end
		cgk = try
			A.cg(ll[N-1], kk[N-1], ll[N], kk[N], j, kk[N-1]+kk[N])
		catch
			@show (ll[N-1], kk[N-1], ll[N], kk[N], j, kk[N-1]+kk[N])
			0.0
		end
		cgm = A.cg(ll[N-1], mm[N-1], ll[N], mm[N], j, mm[N-1]+mm[N])
		if cgk * cgm  != 0
			val += cgk * cgm * A( SVector(llp..., j),
								       SVector(mmp..., mm[N-1]+mm[N]),
								       SVector(kkp..., kk[N-1]+kk[N]) )
		end
   end
   return val
end

_len_mrange(ll) = sum(_ -> 1, _mrange(ll))

function basis(A::CoeffArray{T}, ll; ordered=false) where {T}
	CC = compute_Al(A, ll, Val(ordered))
	svdC = svd(CC)
	rk = rank(Diagonal(svdC.S))
	return svdC.U[:, 1:rk]
end

compute_Al(ll::SVector{N}; ordered = false) where {N} =
		compute_Al(CoeffArray(N, sum(ll)), ll; ordered=ordered)

compute_Al(A::CoeffArray, ll::SVector{N}; ordered = false) where {N} =
		compute_Al(A, ll, Val(ordered))

# unordered
function compute_Al(A::CoeffArray{T}, ll::SVector, ::Val{true}) where {T}
	num_mm_sorted = sum(mm -> issorted(mm), _mrange(ll))
	@show num_mm_sorted
	num_mm = _len_mrange(ll)
   CC = zeros(T, num_mm, num_mm_sorted)
	im = 0
   for mm in _mrange(ll)
		if issorted(mm)
			im += 1
			for (ik, kk) in enumerate(_mrange(ll))
		      CC[ik, im] = A(ll, mm, kk)
			end
		end
	end
   return CC
end

# ordered
function compute_Al(A::CoeffArray{T}, ll::SVector, ::Val{false}) where {T}
	len = _len_mrange(ll)
   CC = zeros(T, len, len)
   for (im, mm) in enumerate(_mrange(ll)), (ik, kk) in enumerate(_mrange(ll))
      CC[ik, im] = A(ll, mm, kk)
   end
   return CC
end

function single_B(A::CoeffArray{T}, ll::SVector) where {T}
	MM = collect(_mrange(ll))
	Is = sortperm([ (sum(mm .!= 0), norm(mm)) for mm in MM ])
	MM = MM[Is]
	CC = zeros(T, length(MM))
	for mm in MM
		for (ik, kk) in enumerate(_mrange(ll))
			CC[ik] = A(ll, mm, kk)
		end
		if norm(CC) > 0
			CC ./= norm(CC)
			break
		end
	end
	return CC
end

end


using Polynomials

"""
    y = evalRPoly(roots, x, k=1)

Compute the value of a polynomial which is given in terms of its `roots`.
"""
function evalRPoly(roots, x, k=1)
    f = poly(roots) * k
    return f(x)
end


function randomSphereFilling(volume,r,seed)

    #Create emptyobjects
    frontSpheresQueue = []
    assemblyList = []
    gridIndex = []
    prevRejectedQueue = []
    newlyRejectedQueue = []
    δbox = 3*r

    #Create seed spheres
    ##First
    push!(frontSpheresQueue,seed)
    push!(assemblyList,seed)
    push!(gridIndex,seed)
    ##Second
    v = rand(Normal(0,1),3)
    v ./= sqrt(sum(v.^2))
    pos2 = seed .+ 2*r .*v
    push!(frontSpheresQueue,pos2)
    push!(assemblyList,pos2)
    push!(gridIndex,pos2)
    ##Third
    mid = seed .+ r .*v
    if v[1] != 0 || v[2] != 0 #Make first orthogonal vector
        orth1 = [v[2],-v[1],0]
    else
        orth1 = [0,v[3],-v[2]]
    end
    orth1 ./= sqrt(sum(orth1.^2))
    orth2 = cross(v,orth1)
    v = rand(Normal(0,1)).*orth1 + rand(Normal(0,1)).*orth2
    v ./= sqrt(sum(v.^2))
    pos2 = mid .+ 2*r*cos(pi/6) .*v
    push!(frontSpheresQueue,pos2)
    push!(assemblyList,pos2)
    push!(gridIndex,pos2)

    #while loop
    while !isempty(frontSpheresQueue)
        #Retrieve one active sphere
        currentSphere = frontSpheresQueue[1]
        #Add neighbour Spheres to currect sphere
        neighboringSpheres = []
        for s in gridIndex
            if sqrt(sum((s-currentSphere).^2)) <= δbox
                push!(neighboringSpheres,s)
            end
        end
        #Candidate point list

    end

    return
end


struct IntegralPotential{T<:AbstractKernel,S}
    kernel::T
    surface::S
end

function IntegralPotential(op::AbstractPDE,surf,ktype::Symbol)
    if ktype == :singlelayer
        k = SingleLayerKernel(op)
    elseif ktype == :doublelayer
        k = DoubleLayerKernel(op)
    else
        error("unknown kernel type: valid options are :singlelayer or :doublelayer")
    end
    return IntegralPotential(k,surf)
end

kernel_type(pot::IntegralPotential) = kernel_type(pot.kernel)

(pot::IntegralPotential)(σ::AbstractVector,x) = pot(kernel_type(pot),σ,x)
function (pot::IntegralPotential)(::SingleLayer,σ::AbstractVector,x)
    f = pot.kernel
    Γ = pot.surface
    iter = zip(qnodes(Γ),qweights(Γ),σ)
    out = mapreduce(+,iter) do (y,w,σ)
        f(x,y)*σ*w
    end
    return out
end
function (pot::IntegralPotential)(::DoubleLayer,σ::AbstractVector,x)
    f = pot.kernel
    Y = pot.surface
    iter = zip(qnodes(Y),qweights(Y),qnormals(Y),σ)
    out  = mapreduce(+,iter) do (y,w,ny,σ)
        f(x,y,ny)*σ*w
    end
    return out
end

Base.getindex(pot::IntegralPotential,σ::AbstractVector) = (x) -> pot(σ,x)

SingleLayerPotential(op::AbstractPDE,surf) = IntegralPotential(SingleLayerKernel(op),surf)
DoubleLayerPotential(op::AbstractPDE,surf) = IntegralPotential(DoubleLayerKernel(op),surf)


using DelimitedFiles

# First Half

exprpt = readdlm("inputs/day01.txt", Int)
[i * j for (i, j) in Iterators.product(exprpt, exprpt) if i + j == 2020][1]

# Second Half

exprpt = readdlm("inputs/day01.txt", Int)
[i * j * k for (i, j, k) in Iterators.product(exprpt, exprpt, exprpt) if i + j + k == 2020][1]


__precompile__()

module Baseline
using Flux

using Common: minibatch, vecvectomat!, vecvectomat, copymodel, copymodel!

#using BSON: @save
using JLD2
using LinearAlgebra
using Distributed
using Random

abstract type AbstractBaseline end

struct Linear{DT<:AbstractFloat} <: AbstractBaseline
    nfeat::Int
    T::Int
    K::Int
    reg_coeff::DT
    coeff::AbstractVector{DT} # states x N

    # workspace
    featnd::Array{DT,3}
    A::AbstractMatrix{DT}
    bline::AbstractVector{DT}
    function Linear{DT}(reg::AbstractFloat, ns::Integer, T::Integer, K::Integer) where DT<:AbstractFloat
        nfeat = 2*ns+3+1
        featnd=zeros(DT, nfeat, T, K)
        A = zeros(DT, nfeat, nfeat)
        bline = zeros(DT, T*K)
        return new(nfeat, T, K, reg, zeros(DT, nfeat),
                   featnd, A, bline)
    end
end

struct Quadratic{DT<:AbstractFloat} <: AbstractBaseline
    nfeat::Int
    T::Int
    K::Int
    reg_coeff::DT
    coeff::AbstractVector{DT} # states x N

    # workspace
    featnd::Array{DT,3}
    A::AbstractMatrix{DT}
    bline::AbstractVector{DT}
    function Quadratic{DT}(reg::AbstractFloat, ns::Integer, T::Integer, K::Integer) where DT<:AbstractFloat
        nfeat = ns+convert(Int,ns*(ns+1)/2)+ns+ns+4+1
        featnd=zeros(DT, nfeat, T, K)
        gettimefeatures!(featnd)

        A = zeros(DT, nfeat, nfeat)
        bline = zeros(DT, T*K)
        return new(nfeat, T, K, reg, zeros(DT, nfeat),
                   featnd, A, bline)
    end
end

struct NN{DT<:AbstractFloat} <: AbstractBaseline
    nhidden::Int
    nfeat::Int
    T::Int
    K::Int

    features::Function
    loss::Function
    opt::Function
    epochs::Int

    #workspace
    featnd::Array{DT,2}
    m::Flux.Chain
    bline::AbstractVector{DT}
    function NN{DT}(ns::Integer, T::Integer, K::Integer;
                    nhidden=64,
                    epochs=3,
                    step=0.001,
                    bias=0.0, # l2 regularization is broken when bias vec is 0...
                    priorreg=false,
                    featurefunc=getNNfeaturesNoTime
                   ) where DT<:AbstractFloat

        nfeat = size(featurefunc(rand(ns,2)), 1)
        featnd = zeros(DT, nfeat, T*K)
        m = Chain(Dense(nfeat, nhidden, tanh),
                  #Dense(nhidden, nhidden, tanh),
        Dense(nhidden, nhidden, tanh),
        Dense(nhidden, 1, initb=(x)->fill(bias,x)) )
        #Dense(nhidden, 1, (x)->bias+x))

        #m[end].b.data .= bias
        m0 = deepcopy(params(m))

        bline = zeros(DT, T*K)
        #opt = ADAMW(params(m), step)

        if priorreg
            println("doing Prior function")
            my_vecnorm(m) = vecnorm(m.+0.00001)
            opt = ADAM(params(m), step)
            return new(nhidden, nfeat, T, K,
                       featurefunc,
                       (a,b)->Flux.mse(m(a), b)+sum(vecnorm, params(m)-m0.+0.00001)/length(b),
                       #(a,b)->sum(abs.(m(a)-b))/length(b)+sum(my_vecnorm, params(m)-m0)/length(b),
            opt,
            epochs,
            featnd, m, bline)
        else
            println("doing MSE Loss function")
            opt = ADAM(params(m), step)
            return new(nhidden, nfeat, T, K,
                       featurefunc,
                       (a,b)->Flux.mse(m(a), b), #+0.1*(sum(vecnorm, params(m).+0.00001)/length(b)),
                       #(a,b)->sum(abs.(m(a)-b))/length(b),
            opt,
            epochs,
            featnd, m, bline)
        end
    end
end

######################################################################## common

function gettimefeatures!(feat::Array{DT,3}) where DT<:AbstractFloat
    T    = size(feat, 2)
    K = size(feat, 3)
    al = convert(Array{DT}, linspace(1, T, T)') ./ DT(T)
    for t=1:K
        feat[end-4,:,t] = al
        feat[end-3,:,t] = al.^2
        feat[end-2,:,t] = al.^3
        feat[end-1,:,t] = al.^4
        feat[end,:,t]   .= DT(1.0)
    end
end

function getlinearfeatures(obs::AbstractMatrix{Float64})
    T = size(obs,2)
    # TODO add clipping
    o = max.(obs, -10.0)
    o = min.(o, 10.0)
    al = linspace(1, T, T)' / T
    feat = vcat(o, o.^2, al, al.^2, al.^3, ones(1, T))
end

function getquadfeatures!(feat::Array{DT,3},
                          obs::AbstractMatrix{DT},
                          t::Integer) where DT<:AbstractFloat
    n = size(obs, 1)
    T = size(feat, 2)

    feat[1:n,:,t] = obs[:,(t-1)*T+1:t*T] #implicite conversion between 32, 64
    feat[1:n,:,t] = max.(feat[1:n,:,t], DT(-10.0))
    feat[1:n,:,t] = min.(feat[1:n,:,t], DT(10.0))
    feat[1:n,:,t] /= DT(10.0)

    qfsize = convert(Int, n*(n+1)/2)
    qf = view(feat, (n+1):(n+qfsize) , :, t) #Matrix{Float64}(qfsize, T)
    k = 1
    @inbounds for i=1:n
        for j=i:n
            for x=1:T
                qf[k,x] = feat[i,x,t].*feat[j,x,t]
            end
            k += 1
        end
    end

    i0 = n+qfsize
    i1 = i0+n
    @inbounds for j = 1:n
        for i=1:T
            feat[i0+j,i,t] = sin(pi*feat[j,i,t])
            feat[i1+j,i,t] = cos(pi*feat[j,i,t])
        end
    end
end

function getNNfeatures(obs::AbstractMatrix{Float64})
    T = size(obs,2)
    o = max.(obs, -10.0)
    o = min.(o, 10.0)
    al = linspace(1, T, T)' / 1000.0
    feat = vcat(o, o.^2, al, al.^2, al.^3, ones(1, T))
    #feat = vcat(o, al, al.^2, al.^3, ones(1, T))
end

function getNNfeaturesNoTime(obs::AbstractMatrix{Float64})
    T = size(obs,2)
    o = max.(obs, -10.0)
    o = min.(o, 10.0)
    feat = vcat(o, o.^2, sin.(o), cos.(o), ones(1, T))
end

function prefit!(b::AbstractBaseline, returns::AbstractVector{DT}) where DT<:AbstractFloat
    featmat = reshape(b.featnd, b.nfeat, b.K*b.T)
    #b.A[:] = (featmat*featmat')
    mul!(b.A, featmat, transpose(featmat))

    return featmat*returns
end

# returns: T*K
function fit!(b::AbstractBaseline, returns::AbstractVector{DT}) where DT<:AbstractFloat

    target = prefit!(b, returns) # sets A matrix, returns target

    for i=1:b.nfeat # set diag
        b.A[i,i] += b.reg_coeff
    end
    #b.coeff[:] = b.A\(target)
    if isposdef(b.A)
        A_ldiv_B!(b.coeff, cholfact!(b.A), target) # needs A to be factorized
    else
        A_ldiv_B!(b.coeff, lufact!(convert(Array{Float64}, b.A)), target) # needs A to be factorized
    end
end

function fit!(b::NN, returns::AbstractVector{DT}) where DT<:AbstractFloat
    nsamples = length(returns)
    nbatch = 50

    X = b.featnd |> gpu
    Y = returns |> gpu

    rndidx = Flux.chunk(shuffle(1:nsamples), nbatch)

    datas = [ (X[:,i], Y[i]') for i in rndidx ]

    #opt = ADAM(params(b.m))
    #evalcb() = @show( norm(b.loss(b.featnd, returns')) / norm(returns) )
    evalcb() = @show( norm(b.loss(datas[1]...)) / norm(datas[1][2]) )
    evalcb(x, y) = norm(b.loss(x, y')) / norm(y)

    for i=1:b.epochs
        Flux.train!(b.loss, datas, b.opt)#, cb = Flux.throttle(evalcb, 1))
        #if evalcb(b.featnd, returns) < 0.5
        #if evalcb() < 0.5
        #   break
        #end
    end
end

function bellmanfit!(b::NN,
                     features::AbstractMatrix{DT},
                     nstep_features::AbstractMatrix{DT},
                     returns::AbstractVector{DT},
                     discount::Float64;
                     batch=16, epochs=20,
                     nbatch=ceil(Int, N/batch),
                     rndidx=minibatch(length(returns), batch)) where DT<:AbstractFloat
    N = length(returns)

    X = b.features(features) #|> gpu
    Y = returns #|> gpu
    Z = nstep_features #|> gpu

    #evalcb() = @show( norm(b.loss(datas[1]...)) / norm(datas[1][2]) )
    #evalcb(x, y) = norm(b.loss(x, y')) / norm(y)

    idx = 1:batch
    values = zeros(DT, batch)
    #for i=1:epochs
    for n=1:nbatch
        #idx = rndidx[n]
        values[:] = discount*Baseline.predict(b, Z[:,idx])
        data = [( @view( X[:,idx] ), (Y[idx]+values)' )]

        Flux.train!(b.loss, data, b.opt)

        idx += batch
    end
    #   rndidx = minibatch(N, batch)
    #end
end

######################################################################## linear

function predict!(b::Linear, obs::AbstractMatrix{T}) where T<:AbstractFloat
    nthread = min.(b.K, Threads.nthreads()) # cant have more threads than data
    Threads.@threads for tid=1:nthread
        #for tid=1:nthread
        thread_range = Distributed.splitrange(b.K, nthread)[tid]

        for t=thread_range
            b.featnd[:,:,t] = getlinearfeatures(obs[:,(t-1)*b.T+1:t*b.T])
        end
    end
    b.bline[:] = reshape(b.featnd, b.nfeat, b.K*b.T)'*b.coeff
end

##################################################################### quadratic
# TODO this is hacky and this whole module is a disgrace
function predict(b::Quadratic, obs::AbstractMatrix{DT}) where DT<:AbstractFloat
    nsamples = size(obs, 2)
    featnd=zeros(DT, b.nfeat, 1, nsamples)
    gettimefeatures!(featnd)
    for t=1:nsamples
        getquadfeatures!(featnd, obs, t)
    end
    values = reshape(featnd, b.nfeat, nsamples)'*b.coeff
    return values
end

function predict!(b::Quadratic, obs::AbstractMatrix{DT}) where DT<:AbstractFloat
    nthread = min.(b.K, Threads.nthreads()) # cant have more threads than data
    Threads.@threads for tid=1:nthread
        #for tid=1:nthread
        thread_range = Distributed.splitrange(b.K, nthread)[tid]

        for t=thread_range
            getquadfeatures!(b.featnd, obs, t)
        end
    end
    b.bline[:] = reshape(b.featnd, b.nfeat, b.K*b.T)'*b.coeff
end

##################################################################### Neural Network

function predict(b::NN, obs::AbstractMatrix{DT}) where DT<:AbstractFloat
    return b.m( b.features(obs) ).data[1,:]
end

function predict!(b::NN, obs::AbstractMatrix{DT}) where DT<:AbstractFloat
    b.featnd[:,:] = b.features(obs)
    b.bline .= b.m(b.featnd).data[1,:]
end

end # module


module Testing
    export run_test

    include("Utils.jl")
    include("Spa.jl")

    struct SynthesizedSignal
        f::Real
        Fs::Int
        amplitudes::Array{Real}
    end

    struct Test
        f::Real
        signal::Array{Real}
    end

    function training_signal(f, Fs; amplitudes = .01 .* rand(50:100, num_harmonics))
        return SynthesizedSignal(f, Fs, amplitudes)
    end

    function get_signal_step(signal::SynthesizedSignal, n::Int)
        return sum([sin(2*π*n*i * signal.f / signal.Fs) * signal.amplitudes[i] / i for i in 1:length(signal.amplitudes)])
    end

    function generate_test(f, length::Int, Fs; num_harmonics = 20, amplitudes = 0.01 .* rand(50:100, num_harmonics))
        test = zeros(length)
        signal = training_signal(f, Fs, amplitudes = amplitudes)
        for n = 0:length-1
            test[n+1] = get_signal_step(signal, n)
        end
        return Test(f, test)
    end

    function get_error(guess, answer)
        return abs(answer - guess) / answer
    end

    function run_test(detect_pitch::Function; N = 10000, Fs = 48000, Wl = 2048, low_note = 80, high_note = Wl / 2)
        note_range = Utils.getRange(low_note, high_note)

        err = zeros(N)
        Threads.@threads for n in 1:N
            test = generate_test(note_range[rand(1:length(note_range))], Wl, Fs)
            guess, _ = detect_pitch(test.signal, Fs)

            err[n] = get_error(guess, test.f)
        end

        return sum(err) / N
    end
end


using Jecco.KG_3_1

grid = SpecCartGrid3D(
    x_min            = -5.0,
    x_max            =  5.0,
    x_nodes          =  64,
    y_min            = -5.0,
    y_max            =  5.0,
    y_nodes          =  64,
    u_max            =  1.001,
    u_domains        =  1,
    u_nodes          =  48,
)

potential = SquarePotential()

evoleq = AffineNull(
    potential      = potential,
)

id = Sine2D(
    kx   = 1,
    ky   = 2,
    Lx   = grid.x_max - grid.x_min,
    Ly   = grid.y_max - grid.y_min,
)

io = InOut(
    out_bulk_every_t            = 0.04,
    checkpoint_every_walltime_hours = 1,
    remove_existing   = true,
)

integration = Integration(
    dt              = 0.002,
    ODE_method      = KG_3_1.RK4(),
    adaptive        = false,
    tmax            = 4.0,
)

run_model(grid, id, evoleq, integration, io)


# István Maros. Computational Techniques of the Simplex Method. Kluwer Academic Publishers, Boston, 2003. [CTSM]
# Robert J. Vanderbei. Linear Programming: Foundations and Extensions. Springer, second edition, 2001. [LPFE]
# Suhl, Suhl. A fast LU update for linear programming. [LULP]

#=
	List of algorithms

	Linear Programming
		I. Solve, Revised, Canonical Simplex.
			1. revised simplex, basis inversion, no reuse
			2. revised simplex, basis factorization, no reuse
			3. revised simplex, basis factorization, reuse, no re-factorization
			4. revised simplex, basis factorization, reuse, re-factorization
			5. revised simplex, basis factorization, reuse, re-factorization, degen (CTSM p230), sing
			6. 5. with 'textbook' Phase-I
			7. 5. with Maros' Phase-I (p203)

		II. Solve, Revised, Sparse, Canonical Simplex.
			1. revised sparse simplex

		III. Generalized Simplex variants of I, then II (Presolve, Postsolve)
			1. generalized version of I.1, using Maros' CF-2 (p13-17), but without translation, calling it CF-2b. (Handwritten p29).
			2. 1. with variable transformation for less lohi calculations, the proper CF-2 formulation
			3. 2. with infeasibility tolerance, type-0 handling and Harris ratio. (CSTM p178)

	LU Factorization
		I. Suhl-suhl (CTSM p136), and original [LULP]
=#

module Jad

module lp_I_1
	using Dcd
	using Lp

	type Dat{T}
		prob::Lp.Cf0_problem
		sol::Lp.Solution
		maxit::Int
		n::Int
		m::Int
		status::Symbol

		# Using CTSM notation.
		# LPFE notation translation.
		# αq - ΔxB, β - xB*, q - j, i - p, t - θ

		iB::Array{Int, 1}
		iR::Array{Int, 1}
		B::Array{T, 2}
		Binv::Array{T, 2}
		β::Array{T, 1}
		π::Array{T, 1}
		cBT::Array{T, 2}
		dJ::Array{T, 1}
		αq::Array{T, 1}
		z::T
		zero::T

		dcd::Dcd.Session
		it::Int
		q::Int
		p::Int
		θ::T

		Dat() = new()
	end

	function construct_dat(T::DataType) return Dat{T}() end

	function fill_dat{T}(prob::Lp.Cf0_problem, dat::Dat{T})
		n = prob.n; m = prob.m;
		dat.maxit = get(prob.params, "maxit", 0)
		dat.prob = prob
		dat.n = n
		dat.m = m
		dat.iB = Array(Int, m)
		dat.iR = Array(Int, n)
		dat.B = Array(T, (m,m))
		#dat.Binv = Array(T, (m,m))
		dat.β = Array(T, m)
		#dat.π = Array(T, m+n)
		#dat.cBT = Array(T, m)
		dat.dJ = Array(T, n)
		#dat.αq = Array(T, n)
		dat.zero = zero(T)
	end

	function warn{T}(dat::Dat{T}, warning::Symbol)
		println("Warning:", warning)
	end

	function succeed{T}(dat::Dat{T})
		dat.status = :Ended
		sol = dat.sol
		sol.iter = dat.it
		sol.solved = true
		sol.status = :Optimal
		sol.z = dat.z
		sol.x = zeros(dat.prob.conv.t, dat.n)
		for i = 1:length(dat.iB)
			xi = dat.iB[i]
			if (xi <= dat.n) sol.x[xi] = dat.β[i]; end
		end
	end

	function fail{T}(dat::Dat{T}, status::Symbol)
		dat.status = :Ended
		sol = dat.sol
		sol.iter = dat.it
		sol.solved = false
		sol.status = status
	end

	function sel_cols{T}(src::Matrix{T}, dst::Matrix{T}, cols::Vector{Int})
		for i = 1:length(cols)
			c = cols[i]
			for r = 1:size(dst)[1]
				dst[r,i] = src[r,c]
			end
		end
	end

	function set_basis_logical{T}(dat::Dat{T})
		dat.iB = [i for i in dat.n+1:dat.n+dat.m]
		dat.iR = [i for i in 1:dat.n]
	end

	function comp_B_R{T}(dat::Dat{T}) sel_cols(dat.prob.A, dat.B, dat.iB) end
	function comp_Binv{T}(dat::Dat{T}) dat.Binv = inv(dat.B) end
	function comp_cBT{T}(dat::Dat{T}) dat.cBT = transpose(dat.prob.c[dat.iB]) end
	function init_β{T}(dat::Dat{T}) dat.β = dat.Binv * (dat.prob.b) end
	function update_β{T}(dat::Dat{T}, p::Int, θ::T) dat.β = dat.β - (θ*dat.αq); dat.β[p] = θ; end
	function init_z{T}(dat::Dat{T}) dat.z = dot(reshape(dat.cBT, length(dat.cBT)), dat.β); end
	function update_z{T}(dat::Dat{T}, q::Int, θ::T) dat.z = dat.z + θ * dat.dJ[q]; end
	function comp_π{T}(dat::Dat{T}) dat.π = reshape(dat.cBT * dat.Binv, length(dat.cBT)); end
	function calc_dj{T}(dat::Dat{T}, j::Int) i = dat.iR[j]; dj = dat.prob.c[i] - dot(dat.π, dat.prob.A[:,i]); return dj; end
	function comp_dJ{T}(dat::Dat{T}) for j = 1:dat.n dat.dJ[j] = calc_dj(dat, j) end; end
	function comp_αq{T}(dat::Dat{T}, q::Int) aq = dat.iR[q]; dat.αq = dat.Binv * (dat.prob.A[:,aq]); end
	function check_optimal_dJ{T}(dat::Dat{T}) return all( dj->(dj >= dat.zero), dat.dJ ); end
	function check_feasible_β{T}(dat::Dat{T}) return all( β->(β >= dat.zero), dat.β ); end

	function pivot_iB_iR{T}(dat::Dat{T}, q::Int, p::Int) dat.iR[q], dat.iB[p] = dat.iB[p], dat.iR[q] end

	function price_full_dantzig{T}(dat::Dat{T})
		# [CTSM].p187
		min_i = 0; min_dj = dat.zero;
		for i = 1:length(dat.dJ)
			dj = dat.dJ[i]
			if (dj < dat.zero && dj <= min_dj)
				if (dj == min_dj)
					warn(dat, :Degen)
				end
				min_i = i; min_dj = dj;
			end
		end
		return min_i
	end

	function chuzro{T}(dat::Dat{T})
		min_i = 0; min_ratio = dat.zero;
		for i = 1:length(dat.αq)
			if (dat.αq[i] > dat.zero)
				ratio = dat.β[i] / dat.αq[i]
				if (min_i == 0 || ratio <= min_ratio)
					if (ratio == min_ratio)
						warn(dat, :Degen)
					end
					min_i = i; min_ratio = ratio;
				end
			end
		end
		return (min_i, min_ratio)
	end

	function solve_init{T}(dat::Dat{T}, sol::Lp.Solution{T})
		dat.sol = sol
		dat.dcd = sol.dcd
		Dcd.@it(dcd)
		dat.it = 0
		dat.status = :Iterating
	end

	function solve_step0{T}(dat::Dat{T})
		set_basis_logical(dat)
		Dcd.@set(dcd, "iB", dat.iB)
		#Initializations
		comp_cBT(dat); comp_B_R(dat); comp_Binv(dat);
		init_β(dat); init_z(dat);
		Dcd.@set(dcd, "β0", dat.β)
	end

	function solve_step_phaseI{T}(dat::Dat{T})
		if (check_feasible_β(dat) == false)
			warn(degen, :PhaseI)
			fail(dat, sol, :Infeasible)
		end
	end

	function solve_iter{T}(dat::Dat{T})
		return (dat.status == :Iterating && (dat.maxit == 0 || dat.it < dat.maxit))
	end

	function solve_next{T}(dat::Dat{T})
		dat.it = dat.it + 1
	end

	function solve_step1{T}(dat::Dat{T})
		Dcd.@it(dcd)
		#Step 1
		Dcd.@set(dcd, "B", dat.B); Dcd.@set(dcd, "Binv", dat.Binv);
		Dcd.@set(dcd, "cBT", dat.cBT);
		comp_π(dat); Dcd.@set(dcd, "π", dat.π);
	end

	function solve_step_price{T}(dat::Dat{T})
		comp_dJ(dat); Dcd.@set(dcd, "dJ", dat.dJ);
		if check_optimal_dJ(dat) succeed(dat); return; end
		dat.q = price_full_dantzig(dat); Dcd.@set(dcd, "q", dat.q);
	end

	function solve_step3{T}(dat::Dat{T})
		comp_αq(dat, dat.q); Dcd.@set(dcd, "αq", dat.αq);
	end

	function solve_step_chuzro{T}(dat::Dat{T})
		dat.p, dat.θ = chuzro(dat); Dcd.@set(dcd, "p", dat.p); Dcd.@set(dcd, "θ", dat.θ);
		if (dat.p == 0) fail(dat, :Unbounded); return; end
		Dcd.@set(dcd, "pivot", (dat.q, dat.p))
	end

	function solve_step5{T}(dat::Dat{T})
		pivot_iB_iR(dat, dat.q, dat.p)
	end

	function solve_update{T}(dat::Dat{T})
		update_β(dat, dat.p, dat.θ); Dcd.@set(dcd, "β", dat.β);
		update_z(dat, dat.q, dat.θ); Dcd.@set(dcd, "z", dat.z);
		comp_cBT(dat); comp_B_R(dat); comp_Binv(dat);
	end

	function solve_dat{T}(dat::Dat{T}, sol::Lp.Solution{T})
		# [CTSM].p33,p30, but with naive basis reinversion instead of update.
		solve_init(dat, sol)
		solve_step0(dat)
		solve_step_phaseI(dat)

		while solve_iter(dat)
			solve_step1(dat)
			solve_step_price(dat);
			if (dat.status != :Iterating) break; end
			solve_step3(dat)
			solve_step_chuzro(dat); if (dat.status != :Iterating) break; end
			solve_step5(dat)
			solve_update(dat)
			solve_next(dat)
		end
		if (dat.status == :Iterating) fail(dat, :Maxit); end
	end

end

module lp_III_1
	using Dcd
	using Lp
	using Jad.lp_I_1

	type Dat{T}
		Base::Module
		prob::Lp.Cf2b_problem{T}
		bdat::lp_I_1.Dat{T}
		is_hi::Vector{Bool}

		p_hi::Bool

		Dat() = new()
	end

	function construct_dat(T::DataType) return Dat{T}() end

	function fill_dat{T}(prob::Lp.Cf2b_problem{T}, dat::Dat{T})
		dat.Base = lp_I_1
		dat.prob = prob
		dat.bdat = dat.Base.construct_dat(T)
		Base.fill_dat(prob.base, dat.bdat)
		dat.is_hi = zeros(Bool, dat.prob.n+dat.prob.m) # all variables at low.
	end

	function init_β{T}(dat::Dat{T}, xlo::Vector{T})
		R = dat.base.prob.A[:,dat.iR]
		dat.β = dat.Binv * (dat.prob.b - R*xlo)
	end

	function init_z{T}(dat::Dat{T}, xlo::Vector{T})
		cR = dat.prob.c[dat.iR]
		zR = dot(cR, xlo)
		zB = dot(reshape(dat.cBT, length(dat.cBT)), dat.β)
		dat.z = zB + zR
	end

	function solve_step0{T}(sdat::Dat{T})
		Base = dat.Base; dat = sdat.bdat;
		Base.set_basis_logical(dat)
		Dcd.@set(dcd, "iB", dat.iB)
		#Initializations
		Base.comp_cBT(dat); Base.comp_B_R(dat); Base.comp_Binv(dat);
		xlo = dat.prob.l[dat.iR]
		init_β(sdat, xlo); init_z(sdat, xlo);
		Dcd.@set(dcd, "β0", dat.β)
	end

	function price_full_dantzig{T}(sdat::Dat{T})
		# [CTSM].p187
		Base = sdat.Base; dat = sdat.bdat;
		min_i = 0; min_dj = dat.zero;
		for i = 1:length(dat.dJ)
			dj = dat.dJ[i]
			if (dj < dat.zero && dj <= min_dj && (sdat.is_hi[dat.iB[dj]] == false) )
				if (dj == min_dj)
					Base.warn(dat, :Degen)
				end
				min_i = i; min_dj = dj;
			end
		end
		return min_i
	end

	function solve_step_price{T}(sdat::Dat{T})
		Base = sdat.Base; dat = sdat.bdat;
		Base.comp_dJ(dat); Dcd.@set(dcd, "dJ", dat.dJ);
		if Base.check_optimal_dJ(dat) Base.succeed(dat); return; end
		dat.q = price_full_dantzig(sdat); Dcd.@set(dcd, "q", dat.q);
	end

	function chuzro{T}(dat::Dat{T})
		# Handwritten notes (p29)
		glob_i = 0; glob_max_θ = dat.zero; glob_hi = false;
		for i = 1:length(dat.αq)
			set_hi = false
			if (dat.αq[i] > dat.zero)
				lo = dat.prob.l[dat.iB[i]]
				max_θ = (dat.β[i] - lo) / dat.αq[i]
				set_hi = false
			elseif (dat.αq[i] < dat.zero)
				hi = dat.prob.h[dat.iB[i]]
				max_θ = (dat.β[i] - hi) / dat.αq[i]
				set_hi = true
			else
				continue
			end
			if (glob_i == 0 || max_θ >= glob_max_θ)
				if (max_θ == globa_max_θ) warn(dat, :Degen); end
				glob_i = i; glob_max_θ = max_θ; glob_hi = set_hi;
			end
		end

		return (glob_i, glob_max_θ, glob_hi)
	end

	function solve_step_chuzro{T}(dat::Dat{T})
		dat.p, dat.θ, dat.p_hi = chuzro(dat); Dcd.@set(dcd, "p", dat.p); Dcd.@set(dcd, "θ", dat.θ);
		if (dat.p == 0) fail(dat, :Unbounded); return; end
		Dcd.@set(dcd, "pivot", (dat.q, dat.p))
	end

	function solve_update{T}(dat::Dat{T})
		dat.Base.solve_update(dat.bdat)
		dat.is_hi[dat.p] = dat.p_hi
	end

	function solve_dat{T}(sdat::Dat{T}, sol::Lp.Solution{T})
		# lp_I_1, with generalized forumulation (w/o translation)
		Base = sdat.Base;	bdat = sdat.bdat;
		Base.solve_init(bdat, sol)
		solve_step0(dat)
		Base.solve_step_phaseI(bdat)

		while solve_iter(bdat)
			Base.solve_step1(bdat)
			solve_step_price(dat); if (bdat.status != :Iterating) break; end
			Base.solve_step3(bdat)
			solve_step_chuzro(dat); if (bdat.status != :Iterating) break; end
			Base.solve_step5(bdat)
			solve_update(dat)
			Base.solve_next(bdat)
		end
		if (dat.status == Iterating) Base.fail(bdat, :Maxit); end
	end

end

module lu_I
	#function factorize{}(B::Array{T, 2})
	#end
end

#=
module lp_I_2
	using Dcd
	using Lp

	type Dat{T}
		prob::Lp.Cf0_problem
		maxit::Int
		n::Int
		m::Int
		# Using CTSM notation.
		iB::Array{Int, 1}
		iR::Array{Int, 1}
		B::Array{T, 2}
		BL::Array{T, 2}
		BU::Array{T, 2}
		BP::Array{T, 1}
		β::Array{T, 1}
		π::Array{T, 1}
		cBT::Array{T, 2}
		dJ::Array{T, 1}
		αq::Array{T, 1}
		z::T
		zero::T

		# LPFE notation translation.
		# αq - ΔxB, β - xB*, q - j, i - p, t - θ

		Dat() = new()
	end

	function fill_dat{T}(prob::Lp.Cf0_problem, dat::Dat{T})
		n = prob.n; m = prob.m;
		dat.maxit = get(prob.params, "maxit", 0)
		dat.prob = prob
		dat.n = n
		dat.m = m
		dat.iB = Array(Int, m)
		dat.iR = Array(Int, n)
		dat.B = Array(T, (m,m))
		#dat.Binv = Array(T, (m,m))
		dat.β = Array(T, m)
		#dat.π = Array(T, m+n)
		#dat.cBT = Array(T, m)
		dat.dJ = Array(T, n)
		#dat.αq = Array(T, n)
		dat.zero = zero(T)
	end

	function sel_cols{T}(src::Matrix{T}, dst::Matrix{T}, cols::Vector{Int})
		for i = 1:length(cols)
			c = cols[i]
			for r = 1:size(dst)[1]
				dst[r,i] = src[r,c]
			end
		end
	end

	function set_basis_logical{T}(dat::Dat{T})
		dat.iB = [i for i in dat.n+1:dat.n+dat.m]
		dat.iR = [i for i in 1:dat.n]
	end

	function comp_B_R{T}(dat::Dat{T}) sel_cols(dat.prob.A, dat.B, dat.iB) end
	function comp_BLU{T}(dat::Dat{T}) dat.BL, dat.BU, dat.BP = lu(dat.B) end
	function comp_cBT{T}(dat::Dat{T}) dat.cBT = transpose(dat.prob.c[dat.iB]) end
	function init_β{T}(dat::Dat{T}) dat.β = dat.Binv * (dat.prob.b) end
	function update_β{T}(dat::Dat{T}, p::Int, θ::T) dat.β = dat.β - (θ*dat.αq); dat.β[p] = θ; end
	function init_z{T}(dat::Dat{T}) dat.z = dot(reshape(dat.cBT, length(dat.cBT)), dat.β); end
	function update_z{T}(dat::Dat{T}, q::Int, θ::T) dat.z = dat.z + θ * dat.dJ[q]; end
	function comp_π{T}(dat::Dat{T}) dat.π = reshape(dat.cBT * dat.Binv, length(dat.cBT)); end
	function calc_dj{T}(dat::Dat{T}, j::Int) i = dat.iR[j]; dj = dat.prob.c[i] - dot(dat.π, dat.prob.A[:,i]); return dj; end
	function comp_dJ{T}(dat::Dat{T}) for j = 1:dat.n dat.dJ[j] = calc_dj(dat, j) end; end
	function comp_αq{T}(dat::Dat{T}, q::Int) aq = dat.iR[q]; dat.αq = dat.Binv * (dat.prob.A[:,aq]); end
	function check_optimal_dJ{T}(dat::Dat{T}) return all( dj->(dj >= dat.zero), dat.dJ ); end
	function check_feasible_β{T}(dat::Dat{T}) return all( β->(β >= dat.zero), dat.β ); end

	function pivot_iB_iR{T}(dat::Dat{T}, q::Int, p::Int) dat.iR[q], dat.iB[p] = dat.iB[p], dat.iR[q] end

	function price_full_dantzig{T}(dat::Dat{T})
		# [CTSM].p187
		min_i = 0; min_dj = dat.zero;
		for i = 1:length(dat.dJ)
			dj = dat.dJ[i]
			if (dj < dat.zero && dj <= min_dj)
				if (dj == min_dj)
					println("Warning: degeneracy.")
				end
				min_i = i; min_dj = dj;
			end
		end
		return min_i
	end

	function chuzro{T}(dat::Dat{T})
		min_i = 0; min_ratio = dat.zero;
		for i = 1:length(dat.αq)
			if (dat.αq[i] > dat.zero)
				ratio = dat.β[i] / dat.αq[i]
				if (min_i == 0 || ratio <= min_ratio)
					if (ratio == min_ratio)
						println("Warning: degeneracy.")
					end
					min_i = i; min_ratio = ratio;
				end
			end
		end
		return (min_i, min_ratio)
	end

	function succeed{T}(dat::Dat{T})
		sol.solved = true
		sol.status = :Optimal
		sol.z = dat.z
		sol.x = eval(parse( "zeros($(dat.prob.type_s), $(dat.n))" ))
		for i = 1:length(dat.iB)
			xi = dat.iB[i]
			if (xi <= dat.n) sol.x[xi] = dat.β[i]; end
		end
	end

	function fail(dat, sol::Lp.Solution, status::Symbol) sol.solved = false; sol.status = status; end

	function solve_dat{T}(dat::Dat{T}, sol::Lp.Solution{T})
		# [CTSM].p33,p30, but with naive basis reinversion instead of update.
		dcd = sol.dcd
		Dcd.@it(dcd)
		it = 0
		#Step 0
		set_basis_logical(dat)
		Dcd.@set(dcd, "iB", dat.iB)
		#Initializations
		comp_cBT(dat); comp_B_R(dat); comp_Binv(dat);
		init_β(dat); init_z(dat);
		Dcd.@set(dcd, "β0", dat.β)
		#todo phaseI
		if (check_feasible_β(dat) == false) println("Warning: phaseI."); fail(dat, sol, :Infeasible); return; end

		while(dat.maxit == 0 || it < dat.maxit)
			Dcd.@it(dcd)
			#Step 1
			Dcd.@set(dcd, "B", dat.B); Dcd.@set(dcd, "Binv", dat.Binv);
			Dcd.@set(dcd, "cBT", dat.cBT);
			comp_π(dat); Dcd.@set(dcd, "π", dat.π);
			#Step 2
			comp_dJ(dat); Dcd.@set(dcd, "dJ", dat.dJ);
			if check_optimal_dJ(dat) succeed(dat, sol); return; end
			q = price_full_dantzig(dat); Dcd.@set(dcd, "q", q);
			#Step 3
			comp_αq(dat, q); Dcd.@set(dcd, "αq", dat.αq);
			#Step 4
			p,θ = chuzro(dat); Dcd.@set(dcd, "p", p); Dcd.@set(dcd, "θ", θ);
			if (p == 0) fail(dat, sol, :Unbounded); return; end
			Dcd.@set(dcd, "pivot", (q, p))
			#Step 5
			pivot_iB_iR(dat, q, p)
			#Updates
			update_β(dat, p, θ); Dcd.@set(dcd, "β", dat.β);
			update_z(dat, q, θ); Dcd.@set(dcd, "z", dat.z);
			comp_cBT(dat); comp_B_R(dat); comp_Binv(dat);
			it = it + 1
		end
		fail(dat, sol, :Maxit)
	end

end
=#

end #module Jad


@testset "Confidence intervals ................. " begin

  ss = SampleSpace(5:20, 50)

  params = MESS(
    ss, .2, .4, alpha = .05, beta = .2, pess = .4
  )

  design = optimaldesign(10, params, solver)
  supp   = support(design)

  cpe    = OCEstimator(design, qsolver)

  ci     = CPInterval(design, confidence = .9)
  @test minimum(coverage.(ci, linspace(0, 1))) > .85

  lim  = hcat(limits.(ci, supp[:, 1], supp[:, 2])...)'
  @test all((0 .<= lim[:, 1]) .& (1 .>= lim[:, 2]))

  ci     = ECPInterval(cpe, confidence = .9)
  @test minimum(coverage.(ci, linspace(0, 1))) >= .9
  lim  = hcat(limits.(ci, supp[:, 1], supp[:, 2])...)'
  @test all((0 .<= lim[:, 1]) .& (1 .>= lim[:, 2]))

  # do not test mmwinterval, not recommended for use

end


# represent the   basis x^(mbar) = x⋅(x-1)⋯(x-m+1) ∈ Πm
"""
    FallingFactorial{T}

Construct  a  polynomial with   respect to the basis `x⁰̲,  x¹̲, x²̲, …` where
`xⁱ̲ = x  ⋅  (x-1) ⋅  (x-2)  ⋯ (x-i+1)` is the falling Pochhammer  symbol.  See
[Falling factorial](https://en.wikipedia.org/wiki/Falling_and_rising_factorials)  for several  facts
about this  polynomial basis.

In [Koepf and Schmersau](https://arxiv.org/pdf/math/9703217.pdf)
connection coefficients between the falling factorial polynomial
system and classical discrete orthogonal polynomials are given.

## Examples

```jldoctest
julia> using Polynomials, SpecialPolynomials

julia> p = basis(FallingFactorial, 3)
FallingFactorial(1.0⋅x³̲)

julia> x = variable(Polynomial)
Polynomials.Polynomial(1.0*x)

julia> p(x) ≈ x*(x-1)*(x-2)
true
```

"""
struct FallingFactorial{T<:Number,X} <: AbstractSpecialPolynomial{T,X}
    coeffs::Vector{T}
    function FallingFactorial{T,X}(coeffs::Vector{T}) where {T<:Number,X}
        length(coeffs) == 0 && return new{T,X}(zeros(T, 1))
        last_nz = findlast(!iszero, coeffs)
        last = max(1, last_nz === nothing ? 0 : last_nz)
        return new{T,X}(coeffs[1:last])
    end
    function FallingFactorial{T}(
        coeffs::Vector{T},
        var::Polynomials.SymbolLike=:x,
    ) where {T<:Number}
        FallingFactorial{T,Symbol(var)}(coeffs)
    end
end

Polynomials.@register FallingFactorial
export FallingFactorial

# modified from  Polynomials.unicode_exponent
function unicode_exponent(io, var, j)
    print(io, var)
    a = ("⁻", "", "", "⁰", "¹", "²", "³", "⁴", "⁵", "⁶", "⁷", "⁸", "⁹")
    for i in string(j)
        print(io, a[Int(i) - 44])
    end
end

# show as 1⋅(x)₀ + 2⋅(x)₁ + 3⋅(x)₂
function Polynomials.showterm(
    io::IO,
    ::Type{P},
    pj::T,
    var,
    j,
    first::Bool,
    mimetype,
) where {N,T,P<:FallingFactorial}
    iszero(pj) && return false
    !first && print(io, " ")
    print(io, Polynomials.hasneg(T) && Polynomials.isneg(pj) ? "- " : (!first ? "+ " : ""))
    print(io, "$(abs(pj))⋅")
    #    print(io,"$(var)")
    unicode_exponent(io, var, j)
    print(io, "̲")
    return true
end

function Polynomials.evalpoly(x, p::FallingFactorial)
    d = degree(p)
    d <= 0 && return p[0] * one(x)

    xⁱ̲ = one(x)
    tot = p[0] * xⁱ̲
    for i in 1:d
        xⁱ̲ *= x - (i - 1)
        tot = muladd(xⁱ̲, p[i], tot)
    end

    tot
end

function Base.one(::Type{P}) where {P<:FallingFactorial}
    T, X = eltype(P), Polynomials.indeterminate(P)
    ⟒(P){T,X}(ones(T, 1))
end

function Polynomials.variable(::Type{P}) where {P<:FallingFactorial}
    T, X = eltype(P), Polynomials.indeterminate(P)
    ⟒(P){T,X}([zero(T), one(T)])
end

k0(P::Type{<:FallingFactorial}) = one(eltype(P))

Base.:*(p::FallingFactorial, q::FallingFactorial) =
    convert(FallingFactorial, convert(Polynomial, p) * convert(Polynomial, q))

## Connect FallingFactorial with AbstractCDOP
function Base.convert(::Type{P}, q::Q) where {P<:FallingFactorial,Q<:AbstractCDOP}
    _convert_cop(P, q)
end
function Base.convert(::Type{P}, q::Q) where {P<:AbstractCDOP,Q<:FallingFactorial}
    _convert_cop(P, q)
end

##  Koepf and Schmersa Thm 6 connection for P, Q=x^n̲
function connection_m(
    ::Type{P},
    ::Type{Q},
    m,
    n,
) where {P<:AbstractCDOP,Q<:FallingFactorial}
    a, b, c, d, e = abcde(P)

    c₀ = (a * n + a * m - a + d) * (n - m)

    c₁ = (m + 1) * (a * n^2 - 2a * m^2 - a * n - a * m + n * d - 2d * m - b * m - d - e)

    c₂ = -(m + 1) * (m + 2) * (a * m^2 + 2a * m + d * m + b * m + a + d + b + c + e)

    (c₀, c₁, c₂)
end

##  Koepf and Schmersa Thm 7, p25 connection for P=x^n̲, Q
function connection_m(
    ::Type{P},
    ::Type{Q},
    m,
    n,
) where {P<:FallingFactorial,Q<:AbstractCDOP}
    ā, b̄, c̄, d̄, ē = abcde(Q)

    c₀ =
        (2m * ā + ā + d̄) *
        (2m * ā + 3ā + d̄) *
        (2m * ā + 2ā + d̄)^2 *
        (2m * ā + d̄) *
        (n - m)

    c₁ = (2m * ā + ā + d̄) * (2m * ā + 3ā + d̄) * (2m * ā + 2ā + d̄) * (m + 1)
    c₁a = 2m^2 * n * ā^2 - 2m^2 * ā^2
    c₁a +=
        m^2 * ā * d̄ + 2m^2 * ā * b̄ + 2m * n * ā^2 + 2m * n * ā * d̄ - 2m * ā^2 -
        m * ā * d̄ +
        2m * ā * b̄ +
        m * d̄^2
    c₁a +=
        2m * d̄ * b̄ + n * ā * d̄ + 2n * ā * ē - n * d̄ * b̄ - ā * d̄ +
        d̄ * b̄ +
        d̄ * ē
    c₁ *= c₁a

    c₂ = (m + 1) * (2m * ā + d̄)
    c₂a = m^4 * ā^3 + 4m^3 * ā^3 + 2m^3 * ā^2 * d̄ + 6m^2 * ā^3 + 6m * ā^2 * d̄
    c₂a +=
        4m^2 * ā^2 * c̄ + 2m^2 * ā^2 * ē + m^2 * ā * d̄^2 - m^2 * ā * d̄ * b̄ -
        m^2 * ā * b̄^2 +
        4m * ā^3 +
        6m^2 * ā^2 * d̄
    c₂a +=
        8m * ā^2 * c̄ + 4m * ā^2 * ē + 2m * ā * d̄^2 - 2m * ā * d̄ * b̄ +
        4m * ā * d̄ * c̄ +
        2m * ā * d̄ * ē
    c₂a +=
        -2m * ā * b̄^2 - m * d̄^2 * b̄ - m * d̄ * b̄^2 +
        ā^3 +
        2ā^2 * d̄ +
        4ā^2 * c̄ +
        2ā^2 * ē +
        ā * d̄^2
    c₂a +=
        -ā * d̄ * b̄ + 4ā * d̄ * c̄ + 2 * ā * d̄ * ē - ā * b̄^2 + ā * ē^2 -
        d̄^2 * b̄ + d̄^2 * c̄
    c₂a += -d̄ * b̄^2 - d̄ * b̄ * ē
    c₂ *= c₂a * (m + 2) * (m * ā + n * ā + ā + d̄)

    (c₀, c₁, c₂)
end

function Base.convert(
    ::Type{P},
    q::Q,
) where {P<:Polynomials.StandardBasisPolynomial,Q<:FallingFactorial}
    q(variable(P))
end
function Base.convert(
    ::Type{P},
    q::Q,
) where {P<:FallingFactorial,Q<:Polynomials.StandardBasisPolynomial}
    connection(P, q)
end

# stirling of 2nd kind
@memoize function sterling²(n::Int, k::Int)
    (iszero(n) && iszero(k)) && return 1
    (iszero(n) || iszero(k)) && return 0

    k * sterling²(n - 1, k) + sterling²(n - 1, k - 1)
end

# xʲ = ∑ᵢ sterling²(j,i) (x)ᵢ [wikipedia](https://en.wikipedia.org/wiki/Falling_and_rising_factorials)
function Base.iterate(
    o::Connection{P,Q},
    state=nothing,
) where {P<:FallingFactorial,Q<:Polynomials.StandardBasisPolynomial}
    n, k = o.n, o.k
    if state == nothing
        j = k
    else
        j = state
        j += 1
    end
    j > n && return nothing
    val = sterling²(j, k)
    (j, val), j
end


using
    BenchmarkTools,
    LinearAlgebra,
    KalmanFilter,
    Plots,
    LazyArrays,
    Colors

function init_mu(num_states, num_measures)
    x = randn(num_states)
    PL = randn(num_states, num_states)
    P = PL'PL
    RL = randn(num_measures, num_measures)
    R = RL'RL
    y = randn(num_measures)
    H = randn(num_measures, num_states)
    P_chol = cholesky(P)
    R_chol = cholesky(R)
    h(x) = H * x
    h!(y, x) = mul!(y, H, x)
    h(x, noise) = H * x .+ noise
    h!(y, x, noise) = y .= @~ H * x .+ noise
    return x, y, P, H, R, P_chol, R_chol, h, h!
end

function init_tu(num_states)
    x = randn(num_states)
    A = randn(num_states, num_states)
    P = A'A
    B = randn(num_states, num_states)
    Q = B'B
    F = randn(num_states, num_states)
    P_chol = cholesky(P)
    Q_chol = cholesky(Q)
    f(x) = F * x
    f!(y, x) = mul!(y, F, x)
    f(x, noise) = F * x .+ noise
    f!(y, x, noise) = y.= @~ F * x .+ noise
    return x, P, Q, F, P_chol, Q_chol, f, f!
end

function run_measurement_update_benchmarks(num_state_tests, num_measurement_tests; allocation = false)

    num_measurements = length(num_measurement_tests)
    kf_types = (:kf, :srkf, :ukf, :srukf, :aukf, :sraukf)
    buffers = [(inplace = zeros(length(num_state_tests), num_measurements), allocating = zeros(length(num_state_tests), num_measurements)) for _ = 1:length(kf_types)]
    results = NamedTuple{kf_types}(buffers)

    for (i, num_states) in enumerate(num_state_tests)
        for (j, num_measures) in enumerate(num_measurement_tests)
            x, y, P, H, R, P_chol, R_chol, h, h! = init_mu(num_states, num_measures)

            results.kf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P, y, H, R)
            else
                @belapsed measurement_update($x, $P, $y, $H, $R)
            end
            kf_inter = KFMUIntermediate(num_states, num_measures)
            results.kf.inplace[i,j] = if allocation
                @allocated measurement_update!(kf_inter, x, P, y, H, R)
            else
                @belapsed measurement_update!($kf_inter, $x, $P, $y, $H, $R)
            end

            results.srkf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P_chol, y, H, R_chol)
            else
                @belapsed measurement_update($x, $P_chol, $y, $H, $R_chol)
            end
            srkf_inter = SRKFMUIntermediate(num_states, num_measures)
            results.srkf.inplace[i,j] = if allocation
                @allocated measurement_update!(srkf_inter, x, P_chol, y, H, R_chol)
            else
                @belapsed measurement_update!($srkf_inter, $x, $P_chol, $y, $H, $R_chol)
            end

            results.ukf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P, y, h, R)
            else
                @belapsed measurement_update($x, $P, $y, $h, $R)
            end
            ukf_inter = UKFMUIntermediate(num_states, num_measures)
            results.ukf.inplace[i,j] = if allocation
                @allocated measurement_update!(ukf_inter, x, P, y, h!, R)
            else
                @belapsed measurement_update!($ukf_inter, $x, $P, $y, $h!, $R)
            end

            results.srukf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P_chol, y, h, R_chol)
            else
                @belapsed measurement_update($x, $P_chol, $y, $h, $R_chol)
            end
            srukf_inter = SRUKFMUIntermediate(num_states, num_measures)
            results.srukf.inplace[i,j] = if allocation
                @allocated measurement_update!(srukf_inter, x, P_chol, y, h!, R_chol)
            else
                @belapsed measurement_update!($srukf_inter, $x, $P_chol, $y, $h!, $R_chol)
            end

            results.aukf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P, y, h, Augment(R))
            else
                @belapsed measurement_update($x, $P, $y, $h, $(Augment(R)))
            end
            aukf_inter = AUKFMUIntermediate(num_states, num_measures)
            results.aukf.inplace[i,j] = if allocation
                @allocated measurement_update!(aukf_inter, x, P, y, h!, Augment(R))
            else
                @belapsed measurement_update!($aukf_inter, $x, $P, $y, $h!, $(Augment(R)))
            end

            results.sraukf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P_chol, y, h, Augment(R_chol))
            else
                @belapsed measurement_update($x, $P_chol, $y, $h, $(Augment(R_chol)))
            end
            sraukf_inter = SRAUKFMUIntermediate(num_states, num_measures)
            results.sraukf.inplace[i,j] = if allocation
                @allocated measurement_update!(sraukf_inter, x, P_chol, y, h!, Augment(R_chol))
            else
                @belapsed measurement_update!($sraukf_inter, $x, $P_chol, $y, $h!, $(Augment(R_chol)))
            end
        end
    end
    results
end

function run_time_update_benchmarks(num_state_tests; allocation = false)

    kf_types = (:kf, :srkf, :ukf, :srukf, :aukf, :sraukf)
    buffers = [(inplace = zeros(length(num_state_tests)), allocating = zeros(length(num_state_tests))) for _ = 1:length(kf_types)]
    results = NamedTuple{kf_types}(buffers)

    for (i, num_states) in enumerate(num_state_tests)
        x, P, Q, F, P_chol, Q_chol, f, f! = init_tu(num_states)

        results.kf.allocating[i] = if allocation
            @allocated time_update(x, P, F, Q)
        else
            @belapsed time_update($x, $P, $F, $Q)
        end
        kf_inter = KFTUIntermediate(num_states)
        results.kf.inplace[i] = if allocation
            @allocated time_update!(kf_inter, x, P, F, Q)
        else
            @belapsed time_update!($kf_inter, $x, $P, $F, $Q)
        end

        results.srkf.allocating[i] = if allocation
            @allocated time_update(x, P_chol, F, Q_chol)
        else
            @belapsed time_update($x, $P_chol, $F, $Q_chol)
        end
        srkf_inter = SRKFTUIntermediate(num_states)
        results.srkf.inplace[i] = if allocation
            @allocated time_update!(srkf_inter, x, P_chol, F, Q_chol)
        else
            @belapsed time_update!($srkf_inter, $x, $P_chol, $F, $Q_chol)
        end

        results.ukf.allocating[i] = if allocation
            @allocated time_update(x, P, f, Q)
        else
            @belapsed time_update($x, $P, $f, $Q)
        end
        ukf_inter = UKFTUIntermediate(num_states)
        results.ukf.inplace[i] = if allocation
            @allocated time_update!(ukf_inter, x, P, f!, Q)
        else
            @belapsed time_update!($ukf_inter, $x, $P, $f!, $Q)
        end

        results.srukf.allocating[i] = if allocation
            @allocated time_update(x, P_chol, f, Q_chol)
        else
            @belapsed time_update($x, $P_chol, $f, $Q_chol)
        end
        srukf_inter = SRUKFTUIntermediate(num_states)
        results.srukf.inplace[i] = if allocation
            @allocated time_update!(srukf_inter, x, P_chol, f!, Q_chol)
        else
            @belapsed time_update!($srukf_inter, $x, $P_chol, $f!, $Q_chol)
        end

        results.aukf.allocating[i] = if allocation
            @allocated time_update(x, P, f, (Augment(Q)))
        else
            @belapsed time_update($x, $P, $f, $(Augment(Q)))
        end
        aukf_inter = AUKFTUIntermediate(num_states)
        results.aukf.inplace[i] = if allocation
            @allocated time_update!(aukf_inter, x, P, f!, Augment(Q))
        else
            @belapsed time_update!($aukf_inter, $x, $P, $f!, $(Augment(Q)))
        end

        results.sraukf.allocating[i] = if allocation
            @allocated time_update(x, P_chol, f, Augment(Q_chol))
        else
            @belapsed time_update($x, $P_chol, $f, $(Augment(Q_chol)))
        end
        sraukf_inter = SRAUKFTUIntermediate(num_states)
        results.sraukf.inplace[i] = if allocation
            @allocated time_update!(sraukf_inter, x, P_chol, f!, Augment(Q_chol))
        else
            @belapsed time_update!($sraukf_inter, $x, $P_chol, $f!, $(Augment(Q_chol)))
        end
    end
    results
end

function plot_benchmarks(results, num_state_tests; num_measurement_tests = 1:1, ylabel = "Time (μs)", scale = 10^-6, logscale = false)
    colors = distinguishable_colors(length(results), [RGB(1,1,1), RGB(0,0,0)], dropseed = true)
    num_rows = ceil(Int, length(num_measurement_tests) / 2)
    p = plot(
        num_state_tests,
        results.kf.allocating / scale,
        layout = length(num_measurement_tests) > 1 ? (num_rows, 2) : 1,
        seriescolor = colors[1],
        legend = false,
        title = length(num_measurement_tests) > 1 ? permutedims(map(x -> "$x measurements", num_measurement_tests)) : "Time update",
        xlabel = "# States",
        ylabel = ylabel,
        yscale = logscale ? :log10 : :identity
    )
    plot!(num_state_tests, results.kf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[1], linestyle = :dash)

    plot!(num_state_tests, results.srkf.allocating / scale, layout = length(num_measurement_tests), seriescolor = colors[2])
    plot!(num_state_tests, results.srkf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[2], linestyle = :dash)

    plot!(num_state_tests, results.ukf.allocating / scale, layout = length(num_measurement_tests), seriescolor = colors[3])
    plot!(num_state_tests, results.ukf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[3], linestyle = :dash)

    plot!(num_state_tests, results.srukf.allocating / scale, layout = length(num_measurement_tests), seriescolor = colors[4])
    plot!(num_state_tests, results.srukf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[4], linestyle = :dash)

    plot!(num_state_tests, results.aukf.allocating / scale, layout = length(num_measurement_tests), seriescolor = colors[5])
    plot!(num_state_tests, results.aukf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[5], linestyle = :dash)

    plot!(num_state_tests, results.sraukf.allocating / scale, layout = length(num_measurement_tests), seriescolor = colors[6], label = "SRAUKF")
    plot!(num_state_tests, results.sraukf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[6], linestyle = :dash, label = "SRAUKF inplace")

    # Hack for single legend
    legend_plot = plot([1], framestyle = :none, seriescolor = colors[1], label = "KF")
    plot!([1], seriescolor = colors[1], linestyle = :dash, label = "KF inplace")

    plot!([1], seriescolor = colors[2], label = "SRKF")
    plot!([1], seriescolor = colors[2], linestyle = :dash, label = "SRKF inplace")

    plot!([1], seriescolor = colors[3], label = "UKF")
    plot!([1], seriescolor = colors[3], linestyle = :dash, label = "UKF inplace")

    plot!([1], seriescolor = colors[4], label = "SRUKF")
    plot!([1], seriescolor = colors[4], linestyle = :dash, label = "SRUKF inplace")

    plot!([1], seriescolor = colors[5], label = "AUKF")
    plot!([1], seriescolor = colors[5], linestyle = :dash, label = "AUKF inplace")

    plot!([1], seriescolor = colors[6], label = "SRAUKF")
    plot!([1], seriescolor = colors[6], linestyle = :dash, label = "SRAUKF inplace")
    plot(p, legend_plot, layout = grid(2, 1, heights = [num_rows, 1.2] / (num_rows + 1.2)), size = (800, num_rows * 440))
end

num_state_tests = [1, 5, 10, 20, 30, 40, 50, 60]
num_measurement_tests = [2, 4, 8, 16, 32, 64]
tu_time = run_time_update_benchmarks(num_state_tests)
tu_allocations = run_time_update_benchmarks(num_state_tests, allocation = true)

tu_time_plot = plot_benchmarks(tu_time, num_state_tests)
tu_alloc_plot = plot_benchmarks(tu_allocations, num_state_tests, ylabel = "Allocations (kB)", scale = 10^3)

png(tu_time_plot, "tu_time")
png(tu_alloc_plot, "tu_alloc")

mu_time = run_measurement_update_benchmarks(num_state_tests, num_measurement_tests)
mu_allocations = run_measurement_update_benchmarks(num_state_tests, num_measurement_tests, allocation = true)

mu_time_plot = plot_benchmarks(mu_time, num_state_tests, num_measurement_tests = num_measurement_tests)
mu_alloc_plot = plot_benchmarks(mu_allocations, num_state_tests, num_measurement_tests = num_measurement_tests, ylabel = "Allocations (kB)", scale = 10^3)

png(mu_time_plot, "mu_time")
png(mu_alloc_plot, "mu_alloc")


using KitBase, LinearAlgebra
using KitBase.FastGaussQuadrature, KitBase.JLD2, KitBase.Plots
using ProgressMeter: @showprogress

function flux_wall!(
    ff::T1,
    f::T2,
    u::T3,
    dt,
    rot = 1,
) where {
    T1<:AbstractVector{<:AbstractFloat},
    T2<:AbstractVector{<:AbstractFloat},
    T3<:AbstractVector{<:AbstractFloat},
}
    δ = heaviside.(u .* rot)
    fWall = 0.5 .* δ .+ f .* (1.0 .- δ)
    @. ff = u * fWall * dt

    return nothing
end

begin
    # setup
    set = Setup("radiation", "linesource", "1d1f1v", "kfvs", "bgk", 1, 2, "vanleer", "extra", 0.5, 0.3)

    # physical space
    x0 = 0
    x1 = 1
    nx = 100
    nxg = 0
    ps = PSpace1D(x0, x1, nx, nxg)

    # velocity space
    nu = 100
    points, weights = gausslegendre(nu)
    vs = VSpace1D(points[1], points[end], nu, points, ones(nu) .* (points[end] - points[1]) / (nu - 1), weights)

    # material
    σs = ones(Float32, nx)
    σa = zeros(Float32, nx)
    σt = σs + σa
    σq = zeros(Float32, nx)

    # time
    dt = set.cfl * ps.dx[1]
    nt = set.maxTime / dt |> floor |> Int
    global t = 0.0

    # solution
    f0 = 0.0001 * ones(nu)
    phi = zeros(nu, nx)
    for i = 1:nx
        phi[:, i] .= f0
    end

    flux = zeros(nu, nx+1)

    cd(@__DIR__)
end

for iter = 1:nt
    println("iteration $iter of $nt")

    ff = @view flux[:, 1]
    flux_wall!(ff, phi[:, 1], points, dt, 1.0)

    @inbounds for i = 2:nx
        ff = @view flux[:, i]
        KitBase.flux_kfvs!(ff, phi[:, i-1], phi[:, i], points, dt)
    end

    @inbounds for i = 1:nx-1
        for q = 1:nu
            phi[q, i] =
                phi[q, i] +
                (flux[q, i] - flux[q, i+1]) / ps.dx[i] +
                (σs[i] * phi[q, i] - σt[i] * phi[q, i]) * dt +
                σq[i] * dt
        end
    end
    phi[:, nx] .=  phi[:, nx-1]

    global t += dt
end

ρ = [sum(phi[:, i] .* weights) for i = 1:nx]
plot(ps.x[1:nx], ρ)


using SynchrotronKernel
#using QuadGK
global const X_MAX = 70.0  # -> F(x) < 2e-6
global const X_MIN = 1e-20 # -> F(x) < 5e-6

"""
    min_synch_loud_momentum(ν0::Real, B::Real)

Computes the minimum synchrotron loud momentum as described in Donnert+16, MNRAS 462, 2014–2032 (2016), Eq. 22.
"""
min_synch_loud_momentum(ν0::Real, B::Real) = √( ν0 / (C_crit * B)  )

"""
    ν_over_ν_crit(p::Real, B::Real, ν::Real)

Computes the fraction ``x = \frac{ν}{ν_c}`` needed for the synchrotron Kernel.
See Donnert+16, MNRAS 462, 2014–2032 (2016), Eq. 19.
"""
ν_over_ν_crit(p::Real, B::Real, ν::Real, sinθ::Real=1.0) = ν / ( C_crit * B * sinθ * p^2 )

"""
    dsa_spectral_index_momentum(M::Real)

Compute the spectral index from DSA for a distribution in momentum space.
"""
function dsa_spectral_index_momentum(M::Real)
    xs = ( γ + 1 ) * M^2 / ( ( γ - 1 ) * M^2 + 2 )
    return 3.0 * xs / (xs - 1.0 )
end

"""
    spectrum(p::Real, p_min::Real, n0::Real, s::Real)

Number density `n` for a given momentum `p` interpolated from `n0` at minimum momentum `p_min`.
"""
spectrum(p::Real, p_min::Real, n0::Real, s::Real) = p < p_min ? 0.0 : n0 * (p/p_min)^(-s)

"""
    integrate_θ(x_in::Real, θ_steps::Integer=100)

Pitch angle integration in Donnert+16, Eq. 17.
"""
function integrate_θ(x_in::Real, θ_steps::Integer=64)

    dθ = 0.5π / θ_steps
    K = 0.0
    @inbounds for θ ∈ LinRange(0.0, 0.5π, θ_steps)
        sinθ = sin(θ)
        x  = x_in / sinθ
        K += sinθ^2 * synchrotron_kernel(x) * dθ
    end
    return K
end

"""
    integrate_θ(x_in::Real, θ_steps::Integer=100)

Pitch angle integration in Donnert+16, Eq. 17 using Simpson's rule.
"""
function integrate_θ_simpson(x_in::Real, θ_steps::Integer=50)

    dθ = 0.5π / θ_steps
    #θ = LinRange(0.0, 0.5π, θ_steps)

    F     = Vector{Float64}(undef, θ_steps)
    F_mid = Vector{Float64}(undef, θ_steps)

    # sin(0.0) = 0.0
    F[1]     = 0.0
    F_mid[1] = sin(0.5dθ)^2 * synchrotron_kernel(x_in/sin(0.5dθ))
    K = 0.0
    @inbounds for i = 2:θ_steps
        # boundary
        sinθ = sin((i-1)*dθ)
        x  = x_in / sinθ
        F[i] = sinθ^2 * synchrotron_kernel(x)

        # mid point
        sinθ = sin((i-0.5)*dθ)
        x    = x_in / sinθ
        F_mid[i] = sinθ^2 * synchrotron_kernel(x)

        # Simpson rule: https://en.wikipedia.org/wiki/Simpson%27s_rule
        K += dθ / 6.0 * ( F[i] + F[i-1] + 4F_mid[i] )
    end
    return K
end

"""
    get_F( p::Real, p_inj::Real, ϵ_cr0::Real, s::Real, B_cgs::Real, ν0::Real)

Helper function to get synchrotron emission at momentum `p`.
"""
function get_F( p::Real, p_inj::Real, ϵ_cr0::Real, s::Real, B_cgs::Real, ν0::Real, integrate_pitch_angle::Bool)

    # energy density at momentum p
    N_E = spectrum(p, p_inj, ϵ_cr0, s)
    # x from Eq. 19 (without sinθ -> for integration)
    x   = ν_over_ν_crit(p, B_cgs, ν0)

    if integrate_pitch_angle
        K = integrate_θ_simpson(x)
    else
        K = synchrotron_kernel(x)
    end

    # integral Donnert+16, Eq. 17 evaluated at momenum p
    return N_E * K

end

"""
    get_log_mid(p_low::Real, p_up::Real)

Helper function to compute the momentum in the middle of the bin.
"""
get_log_mid(q_low::Real, q_up::Real) = 10.0^( 0.5 * ( log10(q_low) + log10(q_up) ) )

"""
    spectral_synchrotron_emission(rho_cgs::Real, B_cgs::Real,
                                  T_K::Real, Mach::Real;
                                  xH::Real=0.76, dsa_model::Integer=1,
                                  ν0::Real=1.44e9,
                                  K_ep::Real=0.01,
                                  Emin::Real=5.0e4,
                                  Emax::Real=1.e10,
                                  p_inj::Real=0.1, # in [me*c]
                                  integrate_pitch_angle::Bool=true,
                                  convert_to_mJy::Bool=false,
                                  N_sample_bins::Integer=128)

Computes the synchrotron emission for a powerlaw spectrum as described in Donnert+16, MNRAS 462, 2014–2032 (2016).

# Arguments
- `rho_cgs::Array{<:Real}`: Density in ``g/cm^3``.
- `B_cgs::Array{<:Real}`:   Magnetic field in Gauss.
- `T_K::Array{<:Real}`:     Temperature in Kelvin.
- `Mach::Array{<:Real}`:    Mach number.

# Keyword Arguments
- `xH::Float64 = 0.76`:               Hydrogen fraction of the simulation, if run without chemical model.
- `dsa_model::Integer=1`:             Diffuse-Shock-Acceleration model. Takes values `0...4`, see next section.
- `ν0::Real=1.44e9`:                  Observation frequency in ``Hz``.
- `K_ep::Real=0.01`:                  Ratio of CR proton to electron energy density.
- `integrate_pitch_angle::Bool=true`: Explicitly integrates over the pitch angle. IF `false` assumes ``sin(θ) = 1``.
- `convert_to_mJy::Bool=false`:       Convert the result from ``[erg/cm^3/Hz/s]`` to ``mJy/cm``.

# DSA models
- `0`: [`KR07_acc`](@ref). Efficiency model from Kang, Ryu, Cen, Ostriker 2007, http://arxiv.org/abs/0704.1521v1.
- `1`: [`KR13_acc`](@ref). Efficiency model from Kang&Ryu 2013, doi:10.1088/0004-637X/764/1/95 .
- `2`: [`Ryu19_acc`](@ref). Efficiency model from Ryu et al. 2019, https://arxiv.org/abs/1905.04476 .
- `3`: [`CS14_acc`](@ref). Efficiency model from Caprioli&Spitkovsky 2015, doi: 10.1088/0004-637x/783/2/91 .
- `4`: [`P16_acc`](@ref). Constant efficiency as in Pfrommer+ 2016, doi: 10.1093/mnras/stw2941 .

"""
function spectral_synchrotron_emission(rho_cgs::Real, B_cgs::Real,
                                       T_K::Real, Mach::Real;
                                       xH::Real=0.76, dsa_model::Integer=1,
                                       ν0::Real=1.44e9,
                                       K_ep::Real=0.01,
                                       Emin::Real=5.0e4,
                                       Emax::Real=1.e10,
                                       p_inj::Real=0.1, # in [me*c]
                                       integrate_pitch_angle::Bool=true,
                                       convert_to_mJy::Bool=false,
                                       N_sample_bins::Integer=128)

    if ( B_cgs == 0 || Mach < 2 )
        return 0
    end

    if dsa_model == 0
        acc_function = KR07_acc
    elseif dsa_model == 1
        acc_function = KR13_acc
    elseif dsa_model == 2
        acc_function = Ryu19_acc
    elseif dsa_model == 3
        acc_function = CS14_acc
    elseif dsa_model == 4
        acc_function = P16_acc
    else
        error("Invalid DSA model selection!")
    end

    # in dimensionless momentum!
    p_inj = 0.1
    p_min = Emin * eV2cgs / ( m_e * c_light^2 )
    p_max = Emax * eV2cgs / ( m_e * c_light^2 )

    # prefactor to Eq. 17
    j_ν_prefac = q_e * √(3) / (m_e * c_light^2)

    # include magnetic field into this
    j_ν_prefac *= B_cgs

    # if run without pitch angle integration
    # sinθ = 1.0 -> integral factor π/2
    if !integrate_pitch_angle
        j_ν_prefac *= 0.5π
    end

    # get spectral index s ∈ ]-∞, 4.0[
    s  = dsa_spectral_index_momentum(Mach)

    # energy density of thermal gas [erg/cm^3]
    ϵ_th = EpsNtherm(rho_cgs, T_K, xH=xH)

    # # CR Energy density of whole powerlaw
    # ϵ_cr0 = K_ep * get_rel_energy_density(Mach, acc_function) * ϵ_th
    # # compensate for injection momentum
    # ϵ_cr0 *= (s - 2) / (p_inj^( 2 - s ) )
    # ! This gives the same result as in Smac2

    # this gives the same result as the analytic solution
    ϵ_cr0 = cre_spec_norm_particle(Mach, acc_function) * ϵ_th

    # width of momentum bins
    di    = log10(p_max/p_min) / (N_sample_bins-1)

    # momentum bins
    p     = @. p_min * 10.0^( di * 0:N_sample_bins-1)

    # width of momentum bins
    dp = zeros(N_sample_bins)
    dp[1] = p[1] - p_min * 10.0^(-1di)
    @inbounds for i = 2:N_sample_bins
        dp[i] = p[i] - p[i-1]
    end

    # storage array for integrand of Eq. 17
    F        = Vector{Float64}(undef, N_sample_bins)
    F_mid    = Vector{Float64}(undef, N_sample_bins)

    F[1]     = get_F(p[1], p_inj, ϵ_cr0, s, B_cgs, ν0, integrate_pitch_angle)
    # not needed
    F_mid[1] = 0.0

    # store total synchrotron emissivity
    jν = 0.0

    # Simpson rule
    @inbounds for i = 2:N_sample_bins

        # end of bin
        F[i]     = get_F(p[i], p_inj, ϵ_cr0, s, B_cgs, ν0, integrate_pitch_angle)

        # middle of bin
        p_mid    = get_log_mid(p[i-1], p[i])
        F_mid[i] = get_F(p_mid, p_inj, ϵ_cr0, s, B_cgs, ν0, integrate_pitch_angle)

        # Simpson rule: https://en.wikipedia.org/wiki/Simpson%27s_rule
        jν += dp[i] / 6.0 * ( F[i] + F[i-1] + 4F_mid[i] )

    end

    if convert_to_mJy
        j_ν_prefac *= mJy_factor
    end

    return j_ν_prefac * jν

end

"""
    spectral_synchrotron_emission( n_p::Vector{<:Real},
                                   p::Vector{<:Real},
                                   B_cgs::Real;
                                   ν0::Real=1.44e9
                                   integrate_pitch_angle::Bool=false,
                                   convert_to_mJy::Bool=false)

Computes the synchrotron emission (in ``[erg/cm^3/Hz/s]``) for a CR spectrum as described in Donnert+16, MNRAS 462, 2014–2032 (2016), Eq. 17.

# Arguments
- `n_p::Vector{<:Real}`: Number density in ``1/cm^3`` for momenta `p`.
- `p::Vector{<:Real}`:   Momenta `p` for number densities.
- `B_cgs::Real`:         Magnetic field strength (absolute value) in Gauss.

# Keyword Arguments
- `ν0::Real=1.44e9`:                  Observation frequency in ``Hz``.
- `integrate_pitch_angle::Bool=false`: Explicitly integrates over the pitch angle. If `false` assumes ``sin(θ) = 1``.
- `convert_to_mJy::Bool=false`:       Convert the result from ``[erg/cm^3/Hz/s]`` to ``mJy/cm``.

"""
function spectral_synchrotron_emission(n_p::Vector{<:Real},
                                       p::Vector{<:Real},
                                       B_cgs::Real;
                                       ν0::Real=1.44e9,
                                       integrate_pitch_angle::Bool=false,
                                       convert_to_mJy::Bool=false)

    if ( B_cgs == 0 ) || sum(n_p) == 0
        return 0
    end

    # avoid error for n_p = -0.0
    n_p[ n_p .== -0.0 ] .= 0.0

    # if there are Infs or NaNs present we get a wrong result, so skip this particle
    # ToDo: Check if there is a better way to fix this!
    if (length(findall( isnan.(n_p) )) > 0) ||
       (length(findall( isinf.(n_p) )) > 0) ||
       (length(findall(   n_p .< 0  )) > 0)
        return 0
    end

    # get the number of momentums for which the energy density is defined
    Nbins = length(n_p)

    # prefactor to Eq. 17
    j_ν_prefac = q_e * √(3) / (m_e * c_light^2)

    # include magnetic field into this
    j_ν_prefac *= B_cgs

    # if run without pitch angle integration
    # sinθ = 1.0 -> integral factor π/2
    if !integrate_pitch_angle
        j_ν_prefac *= 0.5π
    end

    # width of momentum bins
    dp = Vector{Float64}(undef, Nbins)
    @inbounds for i = 1:Nbins
        dp[i] = p[i+1] - p[i]
    end

    # storage array for integrand of Eq. 17
    F        = Vector{Float64}(undef, Nbins)
    F_mid    = Vector{Float64}(undef, Nbins-1)

    @inbounds for i = 1:Nbins-1

        # beginning of bin

        # x from Eq. 19
        x   = ν_over_ν_crit(p[i], B_cgs, ν0)
        # pitch-angle integral
        if integrate_pitch_angle
            K = integrate_θ_simpson(x)
        else
            K = synchrotron_kernel(x)
        end

        # energy density at momentum p * integrated synchrotron kernel
        F[i] = n_p[i] * K

        # middle of bin
        N_E_mid  = get_log_mid( n_p[i],  n_p[i+1] )

        p_mid = get_log_mid( p[i], p[i+1] )
        x     = ν_over_ν_crit(p_mid, B_cgs, ν0)

        if integrate_pitch_angle
            K_mid = integrate_θ_simpson(x)
        else
            K_mid = synchrotron_kernel(x)
        end

        F_mid[i] = N_E_mid * K_mid

    end

    # last bin seperate
    x   = ν_over_ν_crit(p[end], B_cgs, ν0)
    # pitch-angle integral
    if integrate_pitch_angle
        K = integrate_θ_simpson(x)
    else
        K = synchrotron_kernel(x)
    end

    # energy density at momentum p * integrated synchrotron kernel
    F[end] = n_p[end] * K

    # store total synchrotron emissivity
    jν = 0.0
    @inbounds for i = 1:Nbins-1
        # Simpson rule: https://en.wikipedia.org/wiki/Simpson%27s_rule
        jν += dp[i] / 6.0 * ( F[i] + F[i+1] + 4F_mid[i] )
    end

    if convert_to_mJy
        j_ν_prefac *= mJy_factor
    end

    return j_ν_prefac * jν

end


"""
    ReluVal(max_iter::Int64, tree_search::Symbol)

ReluVal combines symbolic reachability analysis with iterative interval refinement to minimize over-approximation of the reachable set.

# Problem requirement
1. Network: any depth, ReLU activation
2. Input: hyperrectangle
3. Output: hpolytope

# Return
`CounterExampleResult` or `ReachabilityResult`

# Method
Symbolic reachability analysis and iterative interval refinement (search).
- `max_iter` default `10`.
- `tree_search` default `:DFS` - depth first search.

# Property
Sound but not complete.

# Reference
[S. Wang, K. Pei, J. Whitehouse, J. Yang, and S. Jana,
"Formal Security Analysis of Neural Networks Using Symbolic Intervals,"
*CoRR*, vol. abs/1804.10829, 2018. arXiv: 1804.10829.](https://arxiv.org/abs/1804.10829)

[https://github.com/tcwangshiqi-columbia/ReluVal](https://github.com/tcwangshiqi-columbia/ReluVal)
"""
@with_kw struct ReluVal
    max_iter::Int64     = 10
    tree_search::Symbol = :DFS # only :DFS/:BFS allowed? If so, we should assert this.
end

struct SymbolicInterval
    Low::Matrix{Float64}
    Up::Matrix{Float64}
    interval::Hyperrectangle
end

SymbolicInterval(x::Matrix{Float64}, y::Matrix{Float64}) = SymbolicInterval(x, y, Hyperrectangle([0],[0]))

# Data to be passed during forward_layer
struct SymbolicIntervalMask
    sym::SymbolicInterval
    LΛ::Vector{Vector{Int64}}
    UΛ::Vector{Vector{Int64}}
end

function solve(solver::ReluVal, problem::Problem)
    reach = forward_network(solver, problem.network, problem.input)
    result = check_inclusion(reach.sym, problem.output, problem.network)
    result.status == :Unknown || return result
    reach_list = SymbolicIntervalMask[reach]
    for i in 2:solver.max_iter
        length(reach_list) > 0 || return BasicResult(:SAT)
        reach = pick_out!(reach_list, solver.tree_search)
        LG, UG = get_gradient(problem.network, reach.LΛ, reach.UΛ)
        feature = get_smear_index(problem.network, reach.sym.interval, LG, UG)
        intervals = split_interval(reach.sym.interval, feature)
        for interval in intervals
            reach = forward_network(solver, problem.network, interval)
            result = check_inclusion(reach.sym, problem.output, problem.network)
            result.status == :UNSAT && return result
            result.status == :SAT || (push!(reach_list, reach))
        end
    end
    return BasicResult(:Unknown)
end

function pick_out!(reach_list, tree_search)
    if tree_search == :BFS
        reach = reach_list[1]
        deleteat!(reach_list, 1)
    else
        n = length(reach_list)
        reach = reach_list[n]
        deleteat!(reach_list, n)
    end
    return reach
end

function symbol_to_concrete(reach::SymbolicInterval)
    n_output = size(reach.Low, 1)
    upper = zeros(n_output)
    lower = zeros(n_output)
    for i in 1:n_output
        lower[i] = lower_bound(reach.Low[i, :], reach.interval)
        upper[i] = upper_bound(reach.Low[i, :], reach.interval)
    end
    return Hyperrectangle(low = lower, high = upper)
end

function check_inclusion(reach::SymbolicInterval, output::AbstractPolytope, nnet::Network)
    reachable = symbol_to_concrete(reach)
    issubset(reachable, output) && return BasicResult(:SAT)
    is_intersection_empty(reachable, output) && return BasicResult(:UNSAT)

    # Sample the middle point
    middle_point = (high(reach.interval) + low(reach.interval))./2
    y = compute_output(nnet, middle_point)
    ∈(y, output) || return CounterExampleResult(:UNSAT, middle_point)

    return BasicResult(:Unknown)
end

function forward_layer(solver::ReluVal, layer::Layer, input::Union{SymbolicIntervalMask, Hyperrectangle})
    return forward_act(forward_linear(input, layer))
end

# Symbolic forward_linear for the first layer
function forward_linear(input::Hyperrectangle, layer::Layer)
    (W, b) = (layer.weights, layer.bias)
    sym = SymbolicInterval(hcat(W, b), hcat(W, b), input)
    LΛ = Vector{Vector{Int64}}(undef, 0)
    UΛ = Vector{Vector{Int64}}(undef, 0)
    return SymbolicIntervalMask(sym, LΛ, UΛ)
end

# Symbolic forward_linear
function forward_linear(input::SymbolicIntervalMask, layer::Layer)
    (W, b) = (layer.weights, layer.bias)
    output_Up = max.(W, 0) * input.sym.Up + min.(W, 0) * input.sym.Low
    output_Low = max.(W, 0) * input.sym.Low + min.(W, 0) * input.sym.Up
    output_Up[:, end] += b
    output_Low[:, end] += b
    sym = SymbolicInterval(output_Low, output_Up, input.sym.interval)
    return SymbolicIntervalMask(sym, input.LΛ, input.UΛ)
end

# Symbolic forward_act
function forward_act(input::SymbolicIntervalMask)
    n_node, n_input = size(input.sym.Up)
    output_Low, output_Up = input.sym.Low[:, :], input.sym.Up[:, :]
    mask_lower, mask_upper = zeros(Int, n_node), ones(Int, n_node)
    for i in 1:n_node
        if upper_bound(input.sym.Up[i, :], input.sym.interval) <= 0.0
            # Update to zero
            mask_lower[i], mask_upper[i] = 0, 0
            output_Up[i, :] = zeros(n_input)
            output_Low[i, :] = zeros(n_input)
        elseif lower_bound(input.sym.Low[i, :], input.sym.interval) >= 0
            # Keep dependency
            mask_lower[i], mask_upper[i] = 1, 1
        else
            # Concretization
            mask_lower[i], mask_upper[i] = 0, 1
            output_Low[i, :] = zeros(n_input)
            if lower_bound(input.sym.Up[i, :], input.sym.interval) < 0
                output_Up[i, :] = zeros(n_input)
                output_Up[i, end] = upper_bound(input.sym.Up[i, :], input.sym.interval)
            end
        end
    end
    sym = SymbolicInterval(output_Low, output_Up, input.sym.interval)
    LΛ = push!(input.LΛ, mask_lower)
    UΛ = push!(input.UΛ, mask_upper)
    return SymbolicIntervalMask(sym, LΛ, UΛ)
end

# Return the splited intervals
function get_smear_index(nnet::Network, input::Hyperrectangle, LG::Matrix, UG::Matrix)
    largest_smear = - Inf
    feature = 0
    r = input.radius
    for i in 1:dim(input)
        smear = sum(max.(abs.(UG[:, i]), abs.(LG[:, i]))) * r[i]
        if smear > largest_smear
            largest_smear = smear
            feature = i
        end
    end
    return feature
end

# Get upper bound in concretization
function upper_bound(map::Vector{Float64}, input::Hyperrectangle)
    bound = map[dim(input)+1]
    input_upper = high(input)
    input_lower = low(input)
    for i in 1:dim(input)
        bound += ifelse( map[i]>0, map[i]*input_upper[i], map[i]*input_lower[i])
    end
    return bound
end

# Get lower bound in concretization
function lower_bound(map::Vector{Float64}, input::Hyperrectangle)
    bound = map[dim(input)+1]
    input_upper = high(input)
    input_lower = low(input)
    for i in 1:dim(input)
        bound += ifelse(map[i]>0, map[i]*input_lower[i], map[i]*input_upper[i])
    end
    return bound
end

# Concrete forward_linear
# function forward_linear_concrete(input::Hyperrectangle, W::Matrix{Float64}, b::Vector{Float64})
#     n_output = size(W, 1)
#     output_upper = zeros(n_output)
#     output_lower = zeros(n_output)
#     for j in 1:n_output
#         output_upper[j] = upper_bound(W[j, :], input) + b[j]
#         output_lower[j] = lower_bound(W[j, :], input) + b[j]
#     end
#     output = Hyperrectangle(low = output_lower, high = output_upper)
#     return output
# end

# Concrete forward_act
# function forward_act(input::Hyperrectangle)
#     input_upper = high(input)
#     input_lower = low(input)
#     output_upper = zeros(dim(input))
#     output_lower = zeros(dim(input))
#     mask_upper = ones(Int, dim(input))
#     mask_lower = zeros(Int, dim(input))
#     for i in 1:dim(input)
#         if input_upper[i] <= 0
#             mask_upper[i] = mask_lower[i] = 0
#             output_upper[i] = output_lower[i] = 0.0
#         elseif input_lower[i] >= 0
#             mask_upper[i] = mask_lower[i] = 1
#         else
#             output_lower[i] = 0.0
#         end
#     end
#     return (output_lower, output_upper, mask_lower, mask_upper)
# end

# This function is similar to forward_linear
# function backward_linear(Low::Matrix{Float64}, Up::Matrix{Float64}, W::Matrix{Float64})
#     n_output, n_input = size(W)
#     n_symbol = size(Low, 2) - 1

#     output_Low = zeros(n_output, n_symbol + 1)
#     output_Up = zeros(n_output, n_symbol + 1)
#     for k in 1:n_symbol + 1
#         for j in 1:n_output
#             for i in 1:n_input
#                 output_Up[j, k] += ifelse(W[j, i]>0, W[j, i] * Up[i, k], W[j, i] * Low[i, k])
#                 output_Low[j, k] += ifelse(W[j, i]>0, W[j, i] * Low[i, k], W[j, i] * Up[i, k])
#             end
#         end
#     end
#     return (output_Low, output_Up)
# end


#
# Returns the pearson coefficient of the fit
#
function pearson(X,Y,model,fit)
  ypred = similar(Y)
  ypred = model(X,fit.param)
  R = Statistics.cor(Y,ypred)
  return R
end


#Program using Agatz et al. method to make a TSP-D using a TSP solution
#See "Optimization Approaches for the Traveling SalesmanProblem with Drone" - Niels AgatzPaul BoumanMarie Schmidt
#######################
#Constantes

#######################
include("tsp.jl")

#=
structure of a TSP-D solution :
tsp : TSP solution used to create the TSP-D
drone : nodes where the drones are assigned
z : value of the objective function for this solution
=#
mutable struct Solution
	tsp::Array{Integer}
	drone::Array{Integer}
	combined::Array{Integer}
	z::Integer
	Solution(tsp::Array{Integer},drone::Array{Integer},combined::Array{Integer},z::Integer) = new(sol,drone,combined,z)
	Solution(tsp::Array{Integer},z::Integer) = new(sol,Array{Integer}(),Array{Integer}(),z)
end

struct Constants
	drone_speed::Integer
	truck_speed::Integer
	drone_range::Integer
end

function evaluate!(sol::Solution)

end
#=
function that assign the drone to the node
=#
function greedy_partitionning!(sol::Solution,con::Constants)
	simple = sol.tsp
	util = zeros(Integer, simple.length)
	for node in util
		node = cout(node-1 % length(sol.tsp),node,con.truck_speed) + cout(node,node+1 % length(sol.tsp),con.truck_speed) - max(cout(node-1% length(sol.tsp),node,con.drone_speed)+cout(node,node+1% length(sol.tsp),con.drone_speed),cout(node-1% length(sol.tsp),node+1% length(sol.tsp),con.truck_speed))
	end

	while !isempty(simple)
		select = whereis(max(util))
		makefly!(sol select)
		if cout(select-1% length(simple),select,1) <= con.drone_range && cout(select,select+1,% length(simple),1) <= con.drone_range
				push!(sol.drone,select)
				push!(sol.combine,select+1%length(sol.tsp))
				push!(sol.combine,select-1%length(sol.tsp))
		else

		pop!(simple,select)

	end

end

#=
function that assign a drone to a node (cf the article p.18)
PRE: node != 1 && node != length(sol.tsp)
=#
function make_fly!(sol::Solution,node::Integer)

	push!(sol.drone,node)
	push!(sol.combined,node+1)
	push!(sol.combined,node-1)

	#we save the performing of the operation
	sol.z -= cout(node-1% length(sol.tsp),node,con.truck_speed) + cout(node,node+1,% length(sol.tsp)con.truck_speed) - max(cout(node-1% length(sol.tsp),node,con.drone_speed)+cout(node,node+1% length(sol.tsp),con.drone_speed),cout(node-1% length(sol.tsp),node+1% length(sol.tsp),con.truck_speed))

end

#=

=#
function push_left!()

end

#=

=#
function push_right!()

end


# speed of light in a vacuum in cm s^-1
const c0 = 2.99792458e10

# Gravity constant in g^-1 cm^3 s^-2
const G = 6.6738399e-8

# atomic mass unit in g
const m_H = 1.660539040e-24

# Boltzmann constant in erg K^-1
const k_B = 1.38064852e-16

# Stefan-Boltzmann constant in erg cm^-2 s^-1 K^-4
const sigma_sb = 5.6704e-5

# radiation density constant in erg cm^-3 K^-4
const a_rad = 4 * sigma_sb / c0

# solar mass in g
const Msun = 1.98e33

# solar radius in cm
const Rsun = 6.95508e10

# solar luminosity in erg s^-1
const Lsun = 3.845999e33


module ModelTES
export
    transitionwidth,
    BiasedTES,
    TESParams,
    getlinearparams,
    NoiseModel,
    TESRecord,
    times,
    rk8,
    IrwinHiltonTES,
    stochastic,
    pulse

using Roots, ForwardDiff, DifferentialEquations
include("rk8.jl")
include("tes_models.jl")

const J_per_eV = 1.602177e-19 #unitless
const kb = 1.38064852e-23 #k boltzmann (J/K)

abstract type AbstractRIT end

# following Irwin-Hilton figure 3
mutable struct TESParams{RITType<:AbstractRIT}
    n       ::Float64   # thermal conductance exponent (unitless)
    Tc      ::Float64  # TES critical temperature (K)
    Tbath   ::Float64   # bath temperature (K)

    k       ::Float64   # Pbath pre-factor (W/K^n)
    C       ::Float64   # heat capacity of TES (J/K)

    L       ::Float64  # inductance of SQUID (H)
    Rl      ::Float64   # Thevenin-equivalent resistance Rload = (Rshunt+Rparasitic)(ohms)
                        # note Ibias = V/Rshunt
    Rp      ::Float64   # Rparastic; Rshunt = Rl-Rparasitic (ohms)
    Rn      ::Float64   # normal resistance of TES (ohms)

    RIT     ::RITType   # RIT surface
end

mutable struct ShankRIT <: AbstractRIT
    Tw::Float64 # transition width (K)
    A::Float64  # current dependence for R(T,I) (A/K^(3/2))
end
transitionwidth(RIT::ShankRIT)=RIT.Tw
transitionwidth(p::TESParams)=transitionwidth(p.RIT)

"Constructor that fixes `Tw` and `A` for a ShankRIT to have the given `alpha` and `beta`
parameters when biased at resistance `R0`."
function ShankRIT(alpha, beta, n, Tc, Tbath, k, R0, Rn)
    T0 = Tc /(1 + 3*beta/(2*alpha) - 2*(Rn-R0)/(Rn*alpha)*atanh(2*R0/Rn-1))
    I0 = sqrt(k*(T0^n-Tbath^n)/R0)
    Tw = T0*(Rn-R0)/(Rn*log(2)*alpha)
    A = I0*(2*alpha/(3*T0*beta))^(3/2)
    ShankRIT(Tw,A)
end

mutable struct BiasedTES{T}
    p::TESParams{T}
    I0::Float64 # intial current for diff eq, aka current through TES (A)
    T0::Float64 # initial temperature for diff equations, aka temperature of TES (K)
    V ::Float64 # Thévinen equivalent voltage V = I0*(p.Rl+p.R0), R0=quiescent resistance
                # also equal to Ibias*Rshunt. Careful! This V is a constant and is NOT
                # the voltage drop across the TES.
end

"For a given T0, the difference R-targetR. Use to solve numerically for T0."
function getR0error(T0, targetR, p::TESParams)
   I0 = sqrt(p.k*(T0^p.n-p.Tbath^p.n)/targetR)
   R(I0,T0,p)-targetR
end

"Find the initial conditions (I0, T0, V) that cause `p` to have resistance `targetR`."
function initialconditions(p::TESParams, targetR)
   T00 = fzero((t)->getR0error(t, targetR,p),p.Tc-10*transitionwidth(p), p.Tc+10*transitionwidth(p))
   I00 = sqrt(p.k*(T00^p.n-p.Tbath^p.n)/targetR)
   R00 = R(I00,T00,p)
   V00 = I00*(p.Rl+R00)
   # now evolve these conditions through integration to really lock them in.
   # shouldn't hard code step size here
   out = pulse(10,1e-1, BiasedTES(p, I00, T00, V00), 0)
   T0 = out.T[end]
   I0 = out.I[end]
   R0 = R(I0,T0,p)
   V = I0*(p.Rl+R0)
   I0,T0,V
end

"Created a biased tes with quiescent state resistance R0"
function BiasedTES(p::TESParams{T}, R0::Float64) where T
   @assert 0 < R0 < p.Rn
   I0, T0, V = initialconditions(p,R0)
   BiasedTES(p,I0,T0,V)
end

"iv_point(p::TESParams, V, I0, T0)
takes thevinin voltage `V`, and initial current `I0`, and intial temperature `T0`
evolves a pulse for 1 second, and takes the final values
returns I,T,V,R"
function iv_point(p::TESParams, V, I0, T0)
    # solve with an adapative algorithm that is fast for large time steps,
    # ask for very long time steps
    # we probably shouldn't hardcode the time, but 1 second is long for all TESs I know of
    out = pulse(2,1.0, BiasedTES(p, I0, T0, V), 0, method=DifferentialEquations.Rosenbrock23())
    T = out.T[end]
    I = out.I[end]
    R = out.R[end]
    V = I*(p.Rl+R)
    I,T,R,V
end

"iv_curve(p::TESParams, Vs)
takes a sorted array of V thevinin values `Vs`, calculates ivs points by
evolving a pulse for 1 second, and taking the last value
returns Is, Ts, Vs_out, Rs"
function iv_curve(p::TESParams, Vs)
    Is=Vector{Float64}(undef, length(Vs))
    Ts=Vector{Float64}(undef, length(Vs))
    Rs=Vector{Float64}(undef, length(Vs))
    Vs_out=Vector{Float64}(undef, length(Vs))
    @assert issorted(Vs)
    for i in length(Vs):-1:1
            if Vs[i]==0
                I,T,R,V=0.0, p.Tbath, 0.0, 0.0
            elseif i==length(Vs)
                # provide guesses that guaranteed to be resistive
                I,T,R,V = iv_point(p, Vs[i], Vs[i]/p.Rn, p.Tc)
            elseif i<length(Vs)
                #provide last solution as guesses
                # I got a speedup by providing nearby starting points, it was less than a factor of 2
                I,T,R,V = iv_point(p, Vs[i], Is[i+1], Ts[i+1])
            end
            Is[i]=I
            Ts[i]=T
            Rs[i]=R
            Vs_out[i]=V
    end
    Is,Ts,Rs,Vs_out
end

function dT_and_dI_iv_point(p,I,T,R,V)
    bt = BiasedTES(p, I, T, V)
    du = zeros(2)
    u = [T,I]
    bt(0.0, u, du)
    du
end

"Calculate `R0` the quiescent resistance of `tes`."
getR0(tes::BiasedTES) = R(tes.I0, tes.T0, tes.p)

"Calculate thermal conductivity from TES heat capacity to Tbath in small signal limit."
function getG0(tes::BiasedTES)
   p=tes.p
   p.k*p.n*tes.T0^(p.n-1)
end

"Calculate paramaters in Irwin-Hilton table 1."
function getlinearparams(bt::BiasedTES)
   f(x) = R(x[1], x[2], bt.p)
   p = bt.p
   R0 = getR0(bt)
   G0 = getG0(bt)
   drdi,drdt = ForwardDiff.gradient(f, [bt.I0, bt.T0])
   alpha = drdt*bt.T0/R0
   beta = drdi*bt.I0/R0
   PJ = bt.I0^2*R0
   loopgain = PJ*alpha/G0/bt.T0
   tauthermal = p.C/G0
   taucc = tauthermal/(1-loopgain) # constant current time constant
   r = p.Rl/R0
   taueff = (1+beta+r)/(1+beta+r+(1-r)*loopgain) # zero inductance effective thermal time constant
   tauelectrical = p.L/(p.Rl+R0*(1+beta))
   invtau = 1/(2*tauelectrical)+1/(2*taucc)
   a = (1/tauelectrical-1/taucc)^2
   b = -4*(R0/p.L)*loopgain*(2+beta)/tauthermal
   lcritical = -p.L*b/a
   invtaupm = 0.5*sqrt(a+b+0*im) # make it complex, so I can get a complex answer
   tauplus = 1/(invtau+invtaupm)
   tauminus = 1/(invtau-invtaupm)
   c = loopgain*(3+beta-r)+(1+beta+r)
   d = 2*sqrt(loopgain*(2+beta)*(loopgain*(1-r)+(1+beta+r)))
   f = R0*tauthermal/(loopgain-1)^2
   lcritplus = (c+d)*f
   lcritminus = (c-d)*f
   bt.I0, bt.T0, bt.V, p.Rl, p.Tbath, p.Tbath, p.L, R0, G0, p.C, alpha, beta, loopgain,
        tauthermal, taucc, taueff, tauelectrical, tauplus, tauminus, lcritplus, lcritminus, lcritical
end

"Paramters from Irwin-Hilton table one for modeling a linear TES. Defined in Table 1 of Irwin-Hilton chapter."
mutable struct IrwinHiltonTES
   I0::Float64
   T0::Float64
   V::Float64
   Rl::Float64
   Tl::Float64 # temperature of the load resistor, usually modeled as =Tbath, but really should come from modeling ep coupling in the load resistor
   Tbath::Float64
   L::Float64
   R0::Float64
   G0::Float64
   C0::Float64
   alpha::Float64
   beta::Float64 #βI
   loopgain::Float64 #ℒI
   tauthermal::Float64
   taucc::Float64 #τI
   taueff::Float64
   tauelectrical::Float64 #τel
   tauplus::Complex{Float64}
   tauminus::Complex{Float64}
   lcritplus::Float64
   lcritminus::Float64
   Lcritical::Float64
end

"Constructor for IrwinHiltonTES model"
function IrwinHiltonTES(tes::BiasedTES)
    linearparams = getlinearparams(tes)
    IrwinHiltonTES(linearparams...)
end

"`Z(tes::IrwinHiltonTES, f)`

Returns the impedance of the `tes` at frequency `f`.
Implements equation 42 of Irwin-Hilton chapter."
function Z(tes::IrwinHiltonTES, f)
  ω=2π*f
  tes.R0*(1+tes.beta) .+ tes.R0*tes.loopgain*(2+tes.beta)./((1-tes.loopgain)*(1 .+im*ω*tes.taucc))
end

"`Zcircuit(tes::IrwinHiltonTES, f)`

Returns impedance of complete circuit of `tes` at frequency `f`."
function Zcircuit(tes::IrwinHiltonTES, f)
  ω=2π*f
  tes.R0 .+ im*ω*tes.L .+ Z(tes,ω)
end

include("tes_noise.jl")
include("stochastic_integration.jl")

isoverdamped(tes::IrwinHiltonTES) = (isreal(tes.tauplus) && isreal(tes.tauminus) &&
                                    tes.tauplus<tes.tauminus)
isunderdamped(tes::IrwinHiltonTES) = (!isoverdamped(tes) && tes.tauplus!=tes.tauminus)

mutable struct TESRecord
    T::Vector{Float64}  # temperature (K)
    I::Vector{Float64}  # current (A)
    R::Vector{Float64}  # TES resistance (Ohm)
    dt::Float64   # seconds between samples (seconds)
end
times(r::TESRecord) = range(0, step=r.dt, length=length(r.I))
Base.length(r::TESRecord) = length(r.I)

# I want to re-write this code so that
# A and Tw are the TES params, alpha and beta are derived
# and R0 and be changed to change the bias conditions
# then I0 and T0 and derived from those 3 parametrs

"find bias parameters T0, I0, V for to make the TES have resistane R0 in steady state"
function findbiasparams(tes::TESParams{ShankRIT},R0)

    T0 = Tc /(1 + 3*beta/(2*alpha) - 2*(Rn-R0)/(Rn*alpha)*atanh(2*R0/Rn-1))

    # V = I0*(p.Rl+p.R0)
end

"TES resistance model (Shank et al. 2014)"
R(I, T, RIT::ShankRIT, Tc, Rn) = Rn/2*(1+tanh.((T-Tc+(max.(I,0.0)/RIT.A).^(2/3))/(2*log(2)*RIT.Tw)))
R(I,T, p::TESParams) = R(I,T, p.RIT, p.Tc, p.Rn)

"thermal TES equation
C*dT/dt = I^2*R - k*(T^n-Tbath^n)"
function dT(I, T, k, n, Tbath, C, R)
    # if you run into domain errors, try uncommenting the following line
    # but it probably represents a mistake in your timesteps that should be fixed.
    #T=max(T,0.0) # avoid domain errors when raising T to a power
    Q=-k*(T^n-Tbath^n)+I^2*R
    Q/C

end

"electrical TES equation
L*di/dt = (IBias-I)*Rs+I*Rs-I*Rtes"
function dI(I, T, V, Rl, L, R)
    (V-I*(Rl+R))/L
end

"Calling a BiasedTES gives the dI and dT terms for integration in an in place manner."
function (bt::BiasedTES)(du,u,p_,t) # use DifferentialEquation 4.0+ API
    T,I = u[1],u[2]
    p = bt.p
    r = R(I,T,p)
    # dT(I, T, p.k, p.n, p.Tbath, p.C, r)
    # dI(I,T, bt.V, p.Rl, p.L, r)
    du[1] = dT(I, T, p.k, p.n, p.Tbath, p.C, r)
    du[2] = dI(I,T, bt.V, p.Rl, p.L, r)
end
function (bt::BiasedTES)(t, u, du) # legacy API for rk8
    bt(du,u,nothing,t)
end

function rk8(nsample::Int, dt::Float64, bt::BiasedTES, E::Vector, npresamples::Int=0)
    out = Vector{TESRecord}(length(E))
    for i in 1:length(E)
        out[i] = rk8(nsample, dt, bt, E[i], npresamples)
    end
    out
end

function rk8(nsample::Int, dt::Float64, bt::BiasedTES, E::Number, npresamples::Int=0)
    # Pair of differential equations y' = f(t,y), where y=[T,I]
    p = bt.p
    # Integrate pair of ODEs for all energies EE
    T = Array{Float64}(undef, nsample)
    I = Array{Float64}(undef, nsample)
    T[1:npresamples] .= bt.T0
    I[1:npresamples] .= bt.I0 # set T0, I0 for presamples
    y = [bt.T0+E*J_per_eV/p.C, bt.I0]; ys = similar(y); work = Array{Float64}(undef, 14)
    T[npresamples+1]=y[1]
    I[npresamples+1]=y[2]
    # npresamples+1 is the point at which initial conditions hold (T differs from T0)
    # npresamples+2 is the first point at which I differs from I0
    for i = npresamples+2:nsample
        rk8!(bt, 0.0, dt, y, ys, work)
        y[:] = ys
        T[i] = y[1]
        I[i] = y[2]
    end
    Rout = [R(I[i],T[i],bt.p) for i=1:length(T)]

    TESRecord(T, I, Rout, dt)
end

# example of using the DifferentialEquations API to solve the relevant equations
function adaptive_solve(bt::BiasedTES, dt::Float64, tspan::Tuple{Float64,Float64}, E::Number, method, abstol, reltol, saveat)
    u0 = [bt.T0+E*ModelTES.J_per_eV/bt.p.C, bt.I0]
    prob = ODEProblem(bt, u0, tspan)
    sol = solve(prob,method,dt=dt,abstol=abstol,reltol=reltol, saveat=saveat, save_everystep=false, dense=false)
end
"pulse(nsample::Int, dt::Float64, bt::BiasedTES, E::Number, npresamples::Int=0; dtsolver=1e-9, method=DifferentialEquations.Tsit5(), abstol=1e-9, reltol=1e-9)"
function pulse(nsample::Int, dt::Float64, bt::BiasedTES, E::Number, npresamples::Int=0; dtsolver=1e-9, method=DifferentialEquations.Tsit5(), abstol=1e-9, reltol=1e-9)
    u0 = [bt.T0+E*ModelTES.J_per_eV/bt.p.C, bt.I0]
    saveat = range(0, step=dt, length=nsample-npresamples)
    prob = ODEProblem(bt, u0, (0.0, last(saveat)))
    sol = solve(prob,method,dt=dtsolver,abstol=abstol,reltol=reltol, saveat=saveat, save_everystep=false, dense=false)
    # npresamples+1 is the point at which initial conditions hold (T differs from T0) (sol[1])
    # npresamples+2 is the first point at which I differs from I0
    T = Vector{Float64}(undef, nsample)
    I = Vector{Float64}(undef, nsample)
    T[npresamples+1:end] = sol[1,:]
    I[npresamples+1:end] = sol[2,:]
    T[1:npresamples] .= bt.T0
    I[1:npresamples] .= bt.I0
    Rout = [R(I[i],T[i],bt.p) for i=1:length(T)]
    TESRecord(T,I, Rout,dt)
end

function pulses(nsample::Int, dt::Float64, bt::BiasedTES, Es::Vector, arrivaltimes::Vector; dtsolver=1e-9, method=DifferentialEquations.Tsit5(), abstol=1e-9, reltol=1e-9)
  u0 = [bt.T0, bt.I0]
  saveat = range(0, step=dt, length=nsample)
  prob = ODEProblem(bt, u0, (0.0, last(saveat)))
  Esdict = Dict([(at,E) for (at,E) in zip(arrivaltimes,Es)])
  # this defines a callback that is evaluated when t equals a value in arrival times
  # when evaluated it discontinuously changes the temperature (u[1])
  # the last (true,true) argument has to do with which points are saved
  # it doesn't appear to add extra points beyond saveat, so maybe save_everystep overrides it?
  function cbfun(integrator)
    integrator.u[1]+=Esdict[integrator.t]*ModelTES.J_per_eV/bt.p.C
    # modify the integrator timestep back to dtsolver, to take small steps on the rising edge of the pulse
    integrator.dtpropose=dtsolver # in future use modify_proposed_dt!, see http://docs.juliadiffeq.org/latest/basics/integrator.html#Stepping-Controls-1
  end
  cb = DiscreteCallback((u,t,integrator)->(t in arrivaltimes), cbfun; save_positions=(false,false))
  # tstops is used to make sure the integrator checks each time in arrivaltimes
  sol = solve(prob,method,dt=dtsolver,abstol=abstol,reltol=reltol, saveat=saveat, save_everystep=false, dense=false,callback=cb, tstops=arrivaltimes)

  T = sol[1,:]
  I = sol[2,:]
  Rout = [R(I[i],T[i],bt.p) for i=1:length(T)]
  TESRecord(T,I, Rout,dt)
end

end # module