export statistics, log_bayes_score_component, log_bayes_score_component_BDeu, net_score, net_score_BDeu

function statistics(
    targetind::Int,
    parents::AbstractVector{Int},
    bincounts::AbstractVector{Int},
    discrete_data::AbstractMatrix{Int}
    )

    q = 1
    if !isempty(parents)
        Np = length(parents)
        q  = prod(bincounts[parents])
        stridevec = fill(1, Np)
        for k = 2:Np
            stridevec[k] = stridevec[k-1] * bincounts[parents[k-1]]
        end
        js = (discrete_data[parents,:] - 1)' * stridevec + 1
    else
        js = fill(1, size(discrete_data,2))
    end
    full(sparse(vec(discrete_data[targetind,:]), vec(js), 1, bincounts[targetind], q))
end
function statistics(
    parents :: Vector{Vector{Int}},
    bincounts :: AbstractVector{Int},
    discrete_data :: AbstractMatrix{Int},
    )

    #=
    Computes sufficient statistics from a discrete dataset
    for a Bayesian Net structure

    INPUT:
        parents:
            list of lists of parent indices
            A variable with index i has bincounts[i]
            and row in discrete_data[i,:]
            No acyclicity checking is done
        bincounts:
            list of variable bin counts, or number of
            discrete values the variable can take on, v ∈ {1 : bincounts[i]}
        discrete_data:
            table of discrete values [n×m]
            where n is the number of nodes
            and m is the number of samples

    OUTPUT:
        N :: Vector{Matrix{Int}}
            a sufficient statistics table for each variable
            Variable with index i has statistics table N[i],
            which is r × q where
            r = bincounts[i] is the number of variable instantiations and
            q is the number of parental instantiations of variable i

            The r-values are ordered from 1 → bincounts[i]
            The q-values are ordered in the same ordering as ind2sub() in Julia Base
                Thus the instantiation of the first parent (by order given in parents[i])
                is varied the fastest.

            ex:
                Variable 1 has parents 2 and 3, with r₁ = 2, r₂ = 2, r₃ = 3
                q for variable 1 is q = r₂×r₃ = 6
                N[1] will be a 6×2 matrix where:
                    N[1][1,1] is the number of time v₁ = 1, v₂ = 1, v₃ = 1
                    N[1][2,1] is the number of time v₁ = 1, v₂ = 2, v₃ = 1
                    N[1][3,1] is the number of time v₁ = 1, v₂ = 1, v₃ = 2
                    N[1][4,1] is the number of time v₁ = 1, v₂ = 2, v₃ = 2
                    N[1][5,1] is the number of time v₁ = 1, v₂ = 1, v₃ = 3
                    N[1][6,1] is the number of time v₁ = 1, v₂ = 2, v₃ = 3
                    N[1][6,2] is the number of time v₁ = 2, v₂ = 1, v₃ = 1
                    ...

    This function uses sparse matrix black magic and was
    mercilessly stolen from Ed Schmerling.
    =#

    n, m = size(discrete_data)
    N = Array(Matrix{Int}, n)
    for i = 1 : n
        # p = parents[i]
        # q = 1
        # if !isempty(p)
        #     Np = length(p)
        #     q  = prod(bincounts[parents[i]])
        #     stridevec = fill(1, Np)
        #     for k = 2:Np
        #         stridevec[k] = stridevec[k-1] * bincounts[p[k-1]]
        #     end
        #     js = (discrete_data[p,:] - 1)' * stridevec + 1
        # else
        #     js = fill(1, m)
        # end
        # N[i] = full(sparse(vec(discrete_data[i,:]), vec(js), 1, bincounts[i], q))
        N[i] = statistics(i, parents[i], bincounts, discrete_data)
    end
    N
end
function statistics(net::Network, bincounts::AbstractVector{Int}, discrete_data::AbstractMatrix{Int})

    n = length(bincounts)
	parentList = [convert(Vector{Int}, collect(get_parents(net, i-1)+1)) for i in 1:n]
	statistics(parentList, bincounts, discrete_data)
end
function statistics(adj::BitMatrix, bincounts::AbstractVector{Int}, discrete_data::AbstractMatrix{Int})

    n = length(bincounts)
    @assert(size(adj,1) == size(adj,2) == n)
    parentList = [find(adj[:,i]) for i in 1:n]
    statistics(parentList, bincounts, discrete_data)
end

function log_bayes_score_component{I<:Integer}(i::Int, parents::AbstractVector{I}, r::AbstractVector{Int}, d::AbstractMatrix{Int})

    #=
    Computes the bayesian score component for the given target variable index
        This assumes a unit dirichlet prior (alpha)

    INPUT:
        i       - index of the target variable
        parents - list of indeces of parent variables (should not contain self)
        r       - list of instantiation counts accessed by variable index
                  r[1] gives number of discrete states variable 1 can take on
        d       - matrix of sufficient statistics / counts
                  d[j,k] gives the number of times the target variable took on its kth instantiation
                  given the jth parental instantiation

    OUTPUT:
        the log bayesian score, Float64
    =#

    alpha = fill(1, r[i], isempty(parents) ? 1 : prod(r[parents]))
    (n, m) = size(d)
    if !isempty(parents)
        Np = length(parents)
        stridevec = fill(1, Np)
        for k = 2:Np
            stridevec[k] = stridevec[k-1] * r[parents[k-1]]
        end
        js = (d[parents,:] - 1)' * stridevec + 1
    else
        js = fill(1, m)
    end
    N = sparse(vec(d[i,:]), vec(js), 1, size(alpha)...) # note: duplicates are added together
    return sum(lgamma(alpha + N)) - sum(lgamma(alpha)) + sum(lgamma(sum(alpha,1))) - sum(lgamma(sum(alpha,1) + sum(N,1)))::Float64
end
function log_bayes_score_component_BDeu{I<:Integer}(i::Int, parents::AbstractVector{I}, r::AbstractVector{Int}, d::AbstractMatrix{Int})

	# obtain BDeu alpha
	if isempty(parents)
		α = 1.0 / r[i]
	else
		α = 1.0 / (r[i] * prod(r[parents]))
	end

    alpha = fill(α, r[i], isempty(parents) ? 1 : prod(r[parents]))
    (n, m) = size(d)
    if !isempty(parents)
        Np = length(parents)
        stridevec = fill(1, Np)
        for k = 2:Np
            stridevec[k] = stridevec[k-1] * r[parents[k-1]]
        end
        js = (d[parents,:] - 1)' * stridevec + 1
    else
        js = fill(1, m)
    end
    N = sparse(vec(d[i,:]), vec(js), 1, size(alpha)...) # note: duplicates are added together
    return sum(lgamma(alpha + N)) - sum(lgamma(alpha)) + sum(lgamma(sum(alpha,1))) - sum(lgamma(sum(alpha,1) + sum(N,1)))
end
function log_bayes_score_component{I<:Integer}(
    i       :: Int,
    parents :: AbstractVector{I},
    r       :: AbstractVector{Int},
    d       :: AbstractMatrix{Int},
    alpha   :: AbstractMatrix{Float64}, # r[i]×prod(r[parents])
    )

    (n, m) = size(d)
    if !isempty(parents)
        Np = length(parents)
        stridevec = fill(1, Np)
        for k = 2:Np
            stridevec[k] = stridevec[k-1] * r[parents[k-1]]
        end
        js = (d[parents,:] - 1)' * stridevec + 1
    else
        js = fill(1, m)
    end
    N = sparse(vec(d[i,:]), vec(js), 1, size(alpha)...) # note: duplicates are added together
    return sum(lgamma(alpha + N)) - sum(lgamma(alpha)) + sum(lgamma(sum(alpha,1))) - sum(lgamma(sum(alpha,1) + sum(N,1)))
end
function log_bayes_score_component{I<:Integer}(i::Int, parents::AbstractVector{I}, r::SharedArray{Int, 1}, d::SharedArray{Int, 2}, cache::Dict{Vector{Int}, Float64})
	if haskey(cache, parents)
        return cache[parents]
    end
    return (cache[parents] = log_bayes_score_component(i, parents, r, d))
end
function log_bayes_score_component{I<:Integer}(i::Int, parents::AbstractVector{I}, r::AbstractVector{Int}, d::Matrix{Int}, cache::Dict{Vector{Int}, Float64})
	if haskey(cache, parents)
        return cache[parents]
    end
    return (cache[parents] = log_bayes_score_component(i, parents, r, d))
end
function log_bayes_score_component{I<:Integer}(
    i       :: Int,
    parents :: AbstractVector{I},
    r       :: AbstractVector{Int},
    d       :: AbstractMatrix{Int},
    alpha   :: AbstractMatrix{Float64}, # r[i]×prod(r[parents])
    cache   :: Dict{Vector{Int}, Float64}
    )

    if haskey(cache, parents)
        return cache[parents]
    end
    return (cache[parents] = log_bayes_score_component(i, parents, r, d, alpha))
end

function net_score( net::Network, r::SharedArray{Int, 1}, d::SharedArray{Int, 2} )
	# compute the log bayes score for the given network on the dataset
	s = @parallel (+) for n in get_all_nodes(net)
		log_bayes_score_component(n + 1, get_parents(net, n)+1, r, d)
	end
	s
end
function net_score( net::Network, r::AbstractVector{Int}, d::Matrix{Int} )
	# compute the log bayes score for the given network on the dataset
	s = @parallel (+) for n in get_all_nodes(net)
		log_bayes_score_component(n + 1, get_parents(net, n)+1, r, d)
	end
	s
end
function net_score_BDeu( net::Network, r::SharedArray{Int, 1}, d::SharedArray{Int, 2} )
	# compute the log bayes score for the given network on the dataset using the BDeu prior
	s = @parallel (+) for n in get_all_nodes(net)
		log_bayes_score_component_BDeu(n + 1, get_parents(net, n)+1, r, d)
	end
	s
end
function net_score_BDeu( net::Network, r::AbstractVector{Int}, d::Matrix{Int} )
	# compute the log bayes score for the given network on the dataset using the BDeu prior
	s = @parallel (+) for n in get_all_nodes(net)
		log_bayes_score_component_BDeu(n + 1, get_parents(net, n)+1, r, d)
	end
	s
end


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ 571c959e-1aee-11eb-0e1e-a1b4d8f7762d
using Distributions, IterTools, LaTeXStrings, PGField, Plots, Plots.Measures, Primes, Random, SparseArrays

# ╔═╡ f7c7716c-3593-11eb-00c3-677a7dbc5fec
md"""
# Plots for the statistical model
"""

# ╔═╡ 1ac80b88-3594-11eb-2a77-7321f527f86a
md"""
## Theoretical analysis
"""

# ╔═╡ 9e393e90-35f0-11eb-25bc-af3fc684f85f
Base.binomial(m::Float64,q::Int) = q == 0 ? 1 : prod(m-i for i in 0:q-1)/factorial(q)

# ╔═╡ fd32d722-3595-11eb-27ea-7d153568cb6a
md"""
Expected number of tests (bound for $q>2$):
"""

# ╔═╡ 1db68090-3594-11eb-3c00-1f948011570b
function effbound(p,n,m,q;α=0,β=1)
	k = n*q/m
	r = 1-p
	u = binomial(m-2,q-2)*ceil(n/binomial(m,q))
	p1 = 1 - β + (β-α)*r^k
	p2 = p1^2 + (β-α)^2*r^(2*k-u)*(1-r^u)
	return (m + n*(1 - q*p1 + binomial(q,2)*p2))/n
end

# ╔═╡ 530c81aa-7bca-11eb-1769-a7fe4d74bb1b
md"""
Sensitivity/specificity:
"""

# ╔═╡ 133cd066-3a9d-11eb-3fa7-19b51f9943ff
sens(p,n,m,q;α=0,β=1) = β^(q+1)

# ╔═╡ 1d15d678-3a9d-11eb-06ef-8fd57eba76be
function spec(p,n,m,q;α=0,β=1)
	k = n*q/m
	r = 1-p
	γ = (β + (α-β)*r^(k-1))^q
	return 1-α*γ
end

# ╔═╡ 215f9960-3595-11eb-313d-ef186d77265b
md"""
Numerical (non-asymptotic) optimization of bound:
"""

# ╔═╡ 2101c74c-3595-11eb-059c-25b8a1a76ff5
mopt(p,n,q;α=0,β=1,maxm=n) = argmin([effbound(p,n,m,q;α=α,β=β) for m in 1:maxm] .|> x->isnan(x) ? Inf : x)

# ╔═╡ 5c9a1260-3595-11eb-0e60-3fe0c174ca5a
effopt(p,n,q;α=0,β=1,maxm=n) = effbound(p,n,mopt(p,n,q;α=α,β=β,maxm=maxm),q;α=α,β=β)

# ╔═╡ 055a9656-35a5-11eb-2094-1d0f43ac2d1d
md"""
## Simulation according to the model
"""

# ╔═╡ fb1e305a-1aef-11eb-0c09-dfdb058d559c
function gen2factors(m)
	(m % 2 == 0) || throw("m must be a multiple of 2, got m=$m.")

	r = m-1
	starter = vcat(
		[[GFElem{r}(0),PGInf{r}()]],
		[[GFElem{r}(i),GFElem{r}(-i)] for i in 1:(r-1)÷2]
	)

	return ([Ai.+GFElem{r}(g) for Ai in starter] for g in 0:r-1)
end

# ╔═╡ fed96cbe-1aef-11eb-18b7-bbf7c302625e
function gen3factors(m)
	(m % 6 == 0) || throw("m must be a multiple of 6, got m=$m.")
	isprime(m-1) || throw("m-1 must be prime, got m=$m.")

	r = m-1
	F = [GFElem{r}(i) for i in 0:r-1]
	PG = [F; PGInf{r}()]

	O = (Set∘map)(PG) do i
		orbit = iterated(x -> -(inv(x)*(one(x)+x)), i)
		Set(collect(Iterators.take(orbit,r+1)))
	end
	ω = F[findfirst(i -> order(i) == r-1, F)]

	return ([(λ.*Ai).+g for Ai in O] for g in F, λ in ω.^(1:(r-1)÷2))
end

# ╔═╡ 0298df42-1af0-11eb-151b-ad0cae6f5363
function graphdesign(n,m,q; reorder=true)
	# Generate factors
	if q == 1
		r = m-1
		factors = vcat([[GFElem{r}(i)] for i in 0:r-1], [[PGInf{r}()]])
	elseif q == 2
		factors = Iterators.flatten(gen2factors(m))
	elseif q == 3
		factors = Iterators.flatten(gen3factors(m))
	else
		throw("q must be 1, 2 or 3, got q=$q.")
	end

	# Form pools and (optionally) reorder to group individuals in identical pools together
	pools = collect(Iterators.take(Iterators.cycle(factors),n))
	if reorder
		nfactors = binomial(m,q)
		pools = collect(Iterators.flatten(pools[i:nfactors:end] for i in 1:nfactors))
	end

	# Convert to pool indices and create sparse array
	poolidx = map(pool -> [num isa GFElem ? num.value+1 : m for num in pool], pools)
	return sparse(collect(Iterators.flatten(poolidx)),repeat(1:n,inner=q),1,m,n)
end

# ╔═╡ c085ecb2-1af4-11eb-1f1b-45f2ced5ef31
graphmatrix(n,m,q; reorder=true) = convert(Array{Bool},graphdesign(n,m,q,reorder=reorder))

# ╔═╡ b7046b98-3aa5-11eb-31c4-95119c9a90fa
struct SimCounts{T<:Integer}
	numtests::T
	falseneg::T
	truepos::T
	trueneg::T
	falsepos::T
end

# ╔═╡ 1038eba8-3aa6-11eb-0bf1-37650421e05e
Base.:+(x::SimCounts,y::SimCounts) = SimCounts(
	x.numtests+y.numtests,
	x.falseneg+y.falseneg,
	x.truepos+y.truepos,
	x.trueneg+y.trueneg,
	x.falsepos+y.falsepos,
)

# ╔═╡ 99ccff42-35a3-11eb-1ef1-bb3d0a6f89c1
function sim(p,n,m,q;α=0,β=1)
	A = graphmatrix(n,m,q)
	x = rand(Bernoulli(p),n)
	y = [any(x[sidx] for sidx in findall(pool)) ? rand(Bernoulli(β)) : rand(Bernoulli(α)) for pool in eachrow(A)]
	xput = [all(y[pidx] for pidx in findall(sample)) for sample in eachcol(A)]
	xhat = fill(false,n)
	for i in findall(xput)
		xhat[i] = x[i] ? rand(Bernoulli(β)) : rand(Bernoulli(α))
	end
	return SimCounts(
		m + count(xput),
		count(  x .& .!xhat),
		count(  x .&   xhat),
		count(.!x .& .!xhat),
		count(.!x .&   xhat),
	)
end

# ╔═╡ c5735898-7bcc-11eb-0e6a-ed0c661835ad
md"""
### Overall simulation efficiency / sensitivity / specificity
"""

# ╔═╡ 2eb94ae2-3a73-11eb-0a46-5585fa02f634
eff(cts::SimCounts) = cts.numtests/(cts.falseneg+cts.truepos+cts.trueneg+cts.falsepos)

# ╔═╡ 62d13952-3a73-11eb-2e22-4db2910fdb48
sens(cts::SimCounts) = cts.truepos/(cts.truepos+cts.falseneg)

# ╔═╡ ddb734be-3a73-11eb-05f6-b30cd2f9e57b
spec(cts::SimCounts) = cts.trueneg/(cts.trueneg+cts.falsepos)

# ╔═╡ 0f4e5394-7bcc-11eb-3521-77679d12146e
md"""
## Figs a-c / d-f

Figs d-f are generated by same code as Figs a-c,
with only the setup changed
(un/comment the relevant line to choose which set).
"""

# ╔═╡ e2a7aad6-7bcc-11eb-3383-65908f153f39
md"""
### Setup
"""

# ╔═╡ 75a6fdb4-3aa0-11eb-38d5-f3c1c4ade957
opchar, figlabels = (α=0.05,β=0.90), (eff="a",sens="b",spec="c")
# opchar, figlabels = (α=0.05,β=0.80), (eff="d",sens="e",spec="f")

# ╔═╡ f576fd34-3a76-11eb-164f-e545ad1f1ba4
prange = 10 .^ range(log10(0.1/100),log10(1/100),length=100)

# ╔═╡ 633999fe-3abc-11eb-113a-ab1c5b66d0fd
designs = [
	(n=3072,m=96, q=2),
	(n=3072,m=192,q=2),
	(n=3072,m=384,q=2),
]

# ╔═╡ 0ce1d1dc-3aaf-11eb-2512-355867435800
Random.seed!(0); results = [[sum(sim(p,des.n,des.m,des.q;α=opchar.α,β=opchar.β) for _ in 1:25) for p in prange] for des in designs]

# ╔═╡ 1a8fd33e-7bcc-11eb-2e76-eb8cb430ec07
md"""
### Fig a / d
"""

# ╔═╡ bcf081c2-506e-11eb-31f2-2743dbbf1bbf
let pticks=[0.1,0.2,0.3,0.4,0.5,1]./100
	effticks = 5:5:25
	plot(size=(300,300),dpi=200,framestyle=:box,
		xscale=:log10,yscale=:log10,legend=:topright,
		xlabel=L"\textrm{Prevalence } (\%)",ylabel=L"\textrm{Efficiency gain}: n/\mathrm{E}\;T",
		xticks=(pticks,string.(pticks*100)),xlim=extrema(pticks),
		yticks=(effticks,string.(effticks)),ylim=(3.8,27))

	for (des,res,lcolor) in zip(designs,results,palette(:default))
		scatter!(prange,inv.(eff.(res)),color=lcolor,markerstrokecolor=:auto,markersize=2,opacity=0.6,label="")
		plot!(prange,p->1/effbound(p,des.n,des.m,des.q;α=opchar.α,β=opchar.β),linewidth=2,color=lcolor,
			label=latexstring("m=$(des.m)"))
	end
	annotate!(0.07/100,32,text(figlabels.eff),top_margin=1mm)
	title!(latexstring("\\beta = $(opchar.β), 1-\\alpha = $(1-opchar.α)"),titlefontsize=9)
	savefig("fig-$(figlabels.eff).png"); plot!()
end

# ╔═╡ 2c59b38c-7bcc-11eb-2a83-458f3d104481
md"""
### Fig b / e
"""

# ╔═╡ d50cdc90-3b1f-11eb-26d5-b9e3c3869715
let pticks=[0.1,0.2,0.3,0.4,0.5,1]./100
	sensticks = 0.5:0.1:0.9
	plot(size=(300,300),dpi=200,framestyle=:box,
		xscale=:log10,legend=:bottomright,
		xlabel=L"\textrm{Prevalence } (\%)",ylabel=L"\textrm{Sensitivity}",
		xticks=(pticks,string.(pticks*100)),xlim=extrema(pticks),
		yticks=(sensticks,string.(sensticks)),ylim=(0.43,0.9))

	for (des,res,lcolor) in zip(designs,results,palette(:default))
		scatter!(prange,sens.(res),color=lcolor,markerstrokecolor=:auto,markersize=2,opacity=0.6,label="")
		plot!(prange,p->sens(p,des.n,des.m,des.q;α=opchar.α,β=opchar.β),linewidth=2,color=lcolor,
			label=latexstring("m=$(des.m)"))
	end
	annotate!(0.065/100,0.94,text(figlabels.sens),top_margin=1mm)
	title!(latexstring("\\beta = $(opchar.β), 1-\\alpha = $(1-opchar.α)"),titlefontsize=9)
	savefig("fig-$(figlabels.sens).png"); plot!()
end

# ╔═╡ 332fda94-7bcc-11eb-1694-7baf75b3a709
md"""
### Fig c / f
"""

# ╔═╡ c588b5de-3b20-11eb-012a-6bfd63359b4d
let pticks=[0.1,0.2,0.3,0.4,0.5,1]./100
	specticks = 0.99:0.002:1
	plot(size=(300,300),dpi=200,framestyle=:box,
		xscale=:log10,legend=:bottomleft,
		xlabel=L"\textrm{Prevalence } (\%)",ylabel=L"\textrm{Specificity}",
		xticks=(pticks,string.(pticks*100)),xlim=extrema(pticks),
		yticks=(specticks,string.(specticks)),ylim=(0.989,1),
	)

	for (des,res,lcolor) in zip(designs,results,palette(:default))
		scatter!(prange,spec.(res),color=lcolor,markerstrokecolor=:auto,markersize=2,opacity=0.6,label="")
		plot!(prange,p->spec(p,des.n,des.m,des.q;α=opchar.α,β=opchar.β),linewidth=2,color=lcolor,
			label=latexstring("m=$(des.m)"))
	end
	annotate!(0.055/100,1.001,text(figlabels.spec),top_margin=1mm)
	title!(latexstring("\\beta = $(opchar.β), 1-\\alpha = $(1-opchar.α)"),titlefontsize=9)
	savefig("fig-$(figlabels.spec).png"); plot!()
end

# ╔═╡ 2383c37a-35a8-11eb-3f37-ed1672ed16dc
md"""
## Figures g-h
"""

# ╔═╡ 9a21144c-7bcb-11eb-3757-a5e3b58e0fde
md"""
### Fig g
"""

# ╔═╡ 249a55a6-3677-11eb-0b55-091be1cbf0a2
let pticks=[0.1,0.2,0.3,0.4,0.5,1]./100
	mticks = 0.02:0.01:0.09
	plot(size=(400,300),dpi=200,framestyle=:box,
		xscale=:log10,legend=:bottomright,
		xlabel=L"\textrm{Prevalence } (\%)",ylabel=L"m/n",
		xticks=(pticks,string.(pticks*100)),xlim=extrema(pticks),
		yticks=(mticks,string.(mticks)),ylim=(0.019,0.095))
	prange = 10 .^ range(log10.(extrema(pticks))...,length=1000)

	for n in [96,384,1536,6144]
		plot!(p->mopt(p,n,2)/n,prange,label=latexstring("n = $n"),linewidth=3)
	end
	plot!(p->2*p^(2/3)-p,prange,label=L"2p^{2/3}-p",color=:black,linewidth=3,linestyle=:dash)
	annotate!(0.075/100,0.095,text("g"),top_margin=1mm)
	savefig("fig-g.png"); plot!()
end

# ╔═╡ a5f9a004-7bcb-11eb-2d97-b3b7a54a54e5
md"""
### Fig h
"""

# ╔═╡ adca3628-506e-11eb-3c52-5d8690c0a184
let pticks=[0.1,0.2,0.3,0.4,0.5,1]./100
	effticks = round.(inv.(3*pticks.^(2/3)),digits=1)
	plot(size=(400,300),dpi=200,framestyle=:box,
		xscale=:log10,yscale=:log10,legend=:topright,
		xlabel=L"\textrm{Prevalence } (\%)",ylabel=L"\textrm{Efficiency gain}: n/\mathrm{E}\;T",
		xticks=(pticks,string.(pticks*100)),xlim=extrema(pticks),
		yticks=(effticks,string.(effticks)))
	prange = 10 .^ range(log10.(extrema(pticks))...,length=1000)

	for n in [96,384,1536,6144]
		plot!(p->1/effopt(p,n,2),prange,label=latexstring("n = $n"),linewidth=3)
	end
	plot!(p->1/(3*p^(2/3)),prange,label=L"p^{-2/3}/3",color=:black,linewidth=3,linestyle=:dash)
	annotate!(0.075/100,36,text("h"))
	savefig("fig-h.png"); plot!()
end

# ╔═╡ Cell order:
# ╟─f7c7716c-3593-11eb-00c3-677a7dbc5fec
# ╠═571c959e-1aee-11eb-0e1e-a1b4d8f7762d
# ╟─1ac80b88-3594-11eb-2a77-7321f527f86a
# ╠═9e393e90-35f0-11eb-25bc-af3fc684f85f
# ╟─fd32d722-3595-11eb-27ea-7d153568cb6a
# ╠═1db68090-3594-11eb-3c00-1f948011570b
# ╟─530c81aa-7bca-11eb-1769-a7fe4d74bb1b
# ╠═133cd066-3a9d-11eb-3fa7-19b51f9943ff
# ╠═1d15d678-3a9d-11eb-06ef-8fd57eba76be
# ╟─215f9960-3595-11eb-313d-ef186d77265b
# ╠═2101c74c-3595-11eb-059c-25b8a1a76ff5
# ╠═5c9a1260-3595-11eb-0e60-3fe0c174ca5a
# ╟─055a9656-35a5-11eb-2094-1d0f43ac2d1d
# ╟─fb1e305a-1aef-11eb-0c09-dfdb058d559c
# ╟─fed96cbe-1aef-11eb-18b7-bbf7c302625e
# ╟─0298df42-1af0-11eb-151b-ad0cae6f5363
# ╟─c085ecb2-1af4-11eb-1f1b-45f2ced5ef31
# ╠═b7046b98-3aa5-11eb-31c4-95119c9a90fa
# ╠═1038eba8-3aa6-11eb-0bf1-37650421e05e
# ╠═99ccff42-35a3-11eb-1ef1-bb3d0a6f89c1
# ╟─c5735898-7bcc-11eb-0e6a-ed0c661835ad
# ╠═2eb94ae2-3a73-11eb-0a46-5585fa02f634
# ╠═62d13952-3a73-11eb-2e22-4db2910fdb48
# ╠═ddb734be-3a73-11eb-05f6-b30cd2f9e57b
# ╟─0f4e5394-7bcc-11eb-3521-77679d12146e
# ╟─e2a7aad6-7bcc-11eb-3383-65908f153f39
# ╠═75a6fdb4-3aa0-11eb-38d5-f3c1c4ade957
# ╟─f576fd34-3a76-11eb-164f-e545ad1f1ba4
# ╟─633999fe-3abc-11eb-113a-ab1c5b66d0fd
# ╠═0ce1d1dc-3aaf-11eb-2512-355867435800
# ╟─1a8fd33e-7bcc-11eb-2e76-eb8cb430ec07
# ╟─bcf081c2-506e-11eb-31f2-2743dbbf1bbf
# ╟─2c59b38c-7bcc-11eb-2a83-458f3d104481
# ╟─d50cdc90-3b1f-11eb-26d5-b9e3c3869715
# ╟─332fda94-7bcc-11eb-1694-7baf75b3a709
# ╟─c588b5de-3b20-11eb-012a-6bfd63359b4d
# ╟─2383c37a-35a8-11eb-3f37-ed1672ed16dc
# ╟─9a21144c-7bcb-11eb-3757-a5e3b58e0fde
# ╟─249a55a6-3677-11eb-0b55-091be1cbf0a2
# ╟─a5f9a004-7bcb-11eb-2d97-b3b7a54a54e5
# ╟─adca3628-506e-11eb-3c52-5d8690c0a184


function test_mf()
    println("-- Testing MF-based recommender")

    m = [NaN 3 NaN 1 2 1 NaN 4
         1 2 NaN NaN 3 2 NaN 3
         NaN 2 3 3 NaN 5 NaN 1]
    data = DataAccessor(m)

    recommender = MF(data, 2)
    build!(recommender, learning_rate=15e-4, max_iter=100)

    # top-4 recommantion list should be same as CF/SVD-based recommender
    rec = recommend(recommender, 1, 4, [i for i in 1:8])
    @test Set([first(r) for r in rec]) == Set([2, 5, 6, 8])
end

test_mf()


# sectors.jl
# Explicit realization of sectors to be used for gradation of vector spaces.

# TODO: Add prettier printing for charges.

# Abelian Sectors:
#-----------------
immutable Parity <: Abelian
    charge::Bool
end
Parity(d::Integer)=Parity(iseven(d))

*(c1::Parity,c2::Parity)=Parity(c1.charge != c2.charge)
Base.conj(c::Parity)=Parity(c.charge)
Base.one(c::Parity)=Parity(false)
Base.one(::Type{Parity})=Parity(false)

immutable ZNCharge{N} <: Abelian
    charge::Int
    ZNCharge(charge::Int)=new(mod(charge,N))
end

*{N}(c1::ZNCharge{N},c2::ZNCharge{N})=ZNCharge{N}(c1.charge+c2.charge)
Base.conj{N}(c::ZNCharge{N})=ZNCharge{N}(-c.charge)
Base.one{N}(c::ZNCharge{N})=ZNCharge{N}(0)
Base.one{N}(::Type{ZNCharge{N}})=ZNCharge{N}(0)

immutable U1Charge <: Abelian
    charge::Int
end

*(c1::U1Charge,c2::U1Charge)=U1Charge(c1.charge+c2.charge)
Base.conj(c::U1Charge)=U1Charge(-c.charge)
Base.one(c::U1Charge)=U1Charge(0)
Base.one(::Type{U1Charge})=U1Charge(0)

# Nonabelian Sectors:
#---------------------
# to be done


using LyapunovExponents

num_rec = 20
sampling_interval = 5
t_renorm = 1
t_clv = t_renorm * num_rec * sampling_interval
prob = CLVProblem(LyapunovExponents.lorenz_63().prob,
                  t_renorm = t_renorm,
                  t_forward_tran = 4000,
                  t_backward_tran = 4000,
                  t_clv = t_clv)
solver = init(prob)

x_history = [Vector{Float64}(3) for _ in 1:num_rec]
G_history = [Matrix{Float64}(3, 3) for _ in 1:num_rec]
C_history = [Matrix{Float64}(3, 3) for _ in 1:num_rec]

forward = @time forward_dynamics!(solver; progress=1)
@time for (i, G) in indexed_forward_dynamics!(forward)
    k, r = divrem(i, sampling_interval)
    j = k + 1
    if r == 1 && j <= num_rec
        x_history[j] .= phase_state(forward)
        G_history[j] .= G
    end
end

@time for (i, C) in @time indexed_backward_dynamics!(solver; progress=1)
    k, r = divrem(i, sampling_interval)
    j = k + 1
    if r == 1
        C_history[j] .= C
    end
end

CLV_history = [G * C for (G, C) in zip(G_history, C_history)]

using DifferentialEquations
sol = solve(ODEProblem(
    prob.phase_prob.f,
    x_history[1],
    (0.0, t_clv * 3),
    prob.phase_prob.p
))

using Plots
plt = plot(sol, vars=(2, 3), linewidth=0.5, linealpha=0.5, label="")
vec_scale = 3
for (n, (x, V)) in enumerate(zip(x_history, CLV_history))
    for i in 1:3
        plot!(plt,
              [x[2], x[2] + vec_scale * V[2, i]],
              [x[3], x[3] + vec_scale * V[3, i]],
              color = i + 1,
              arrow = 0.4,
              label = (n == 1 ? "CLV$i" : ""))
    end
end
plt


"""
StrateGrid(x::NamedTuple)

Let x = (k1 = v1, k2 = v2, ...., kN = vN) be a NamedTuple of AbstractVectors for state space,
StateGrid(x) returns an AbstractArray M such that
M[i1, i2..., iN] = (k1 = v1[i1], k2 = v2[i2], ...., kN = vN[iN]
"""
struct StateGrid{T, N, C <: NamedTuple} <: AbstractArray{T, N}
    x::C
end
function StateGrid(x::NamedTuple{Names, <: NTuple{N, <: AbstractVector{T}}}) where {Names, N, T}
    StateGrid{T, N, typeof(x)}(x)
end
function Base.eltype(stategrid::StateGrid{T, N, <: NamedTuple{Names, V}}) where {T, N, Names, V}
    NamedTuple{Names, NTuple{N, T}}
end
Base.ndims(stategrid::StateGrid{T, N}) where {T, N} = N
Base.size(stategrid::StateGrid{T, N}) where {T, N} = ntuple(i -> length(stategrid.x[i]), N)
Base.eachindex(stategrid::StateGrid) = CartesianIndices(size(stategrid))
function Base.getindex(stategrid::StateGrid, args::CartesianIndex)
    eltype(stategrid)(ntuple(i -> stategrid.x[i][args[i]], ndims(stategrid)))
end
Base.getindex(stategrid::StateGrid, x::Symbol) = stategrid.x[x]

#========================================================================================

Derive

========================================================================================#
# 1 state variable
@generated function differentiate(::Type{Tsolution}, grid::StateGrid{T1, 1, <: NamedTuple{N}}, y::AbstractArray{T}, icar, bc) where {Tsolution, T1, N, T}
    statename = N[1]
    expr = Expr[]
    for k in 1:length(Tsolution.parameters[1])
        solname = Tsolution.parameters[1][k]
        push!(expr, Expr(:(=), solname, :(y[i, $k])))
        push!(expr, Expr(:(=), Symbol(solname, statename, :_up), :((i < size(y, 1)) ? (y[i+1, $k] - y[i, $k]) / Δxp : convert($T, bc[end, $k]))))
        push!(expr, Expr(:(=), Symbol(solname, statename, :_down), :((i > 1) ? (y[i, $k] - y[i-1, $k]) / Δxm : convert($T, bc[1, $k]))))
        push!(expr, Expr(:(=), Symbol(solname, statename, statename), :((1 < i < size(y, 1)) ? (y[i + 1, $k] / (Δxp * Δx) + y[i - 1, $k] / (Δxm * Δx) - 2 * y[i, $k] / (Δxp * Δxm)) : ((i == 1) ? (y[2, $k] / (Δxp * Δx) + (y[1, $k] - bc[1, $k] * Δxm) / (Δxm * Δx) - 2 * y[1, $k] / (Δxp * Δxm)) : ((y[end, $k] + bc[end, $k] * Δxp) / (Δxp * Δx) + y[end - 1, $k] / (Δxm * Δx) - 2 * y[end, $k] / (Δxp * Δxm))))))
    end
    quote
        $(Expr(:meta, :inline))
        i = icar[1]
        grida = grid.x[1]
        @inbounds Δxm = grida[max(i, 2)] - grida[max(i-1, 1)]
        @inbounds Δxp = grida[min(i+1, size(y, 1))] - grida[min(i, size(y, 1) - 1)]
        Δx = (Δxm + Δxp) / 2
        @inbounds $(Expr(:tuple, expr...))
    end
end

# 2 state variables
@generated function differentiate(::Type{Tsolution}, grid::StateGrid{T1, 2, <: NamedTuple{N}}, y::AbstractArray{T}, icar, bc) where {Tsolution, T1, N, T}
    statename1 = N[1]
    statename2 = N[2]
    expr = Expr[]
    for k in 1:length(Tsolution.parameters[1])
        solname = Tsolution.parameters[1][k]
        push!(expr, Expr(:(=), solname, :(y[i1, i2, $k])))
        push!(expr, Expr(:(=), Symbol(solname, statename1, :_up), :((i1 < size(y, 1)) ? (y[i1+1, i2, $k] - y[i1, i2, $k]) / Δx1p : convert($T, bc[end, i2, $k]))))
        push!(expr, Expr(:(=), Symbol(solname, statename1, :_down), :((i1 > 1) ? (y[i1, i2, $k] - y[i1-1, i2, $k]) / Δx1m : convert($T, bc[1, i2, $k]))))
        push!(expr, Expr(:(=), Symbol(solname, statename2, :_up), :((i2 < size(y, 2)) ? (y[i1, i2+1, $k] - y[i1, i2, $k]) / Δx2p : convert($T, bc[i1, end, $k]))))
        push!(expr, Expr(:(=), Symbol(solname, statename2, :_down),  :((i2 > 1) ? (y[i1, i2, $k] - y[i1, i2-1, $k]) / Δx2m : convert($T, bc[i1, 1, $k]))))
        push!(expr, Expr(:(=), Symbol(solname, statename1, statename1), :((1 < i1 < size(y, 1)) ? (y[i1 + 1, i2, $k] / (Δx1p * Δx1) + y[i1 - 1, i2, $k] / (Δx1m * Δx1) - 2 * y[i1, i2, $k] / (Δx1p * Δx1m)) : ((i1 == 1) ? (y[2, i2, $k] / (Δx1p * Δx1) + (y[1, i2, $k] - bc[1, i2, $k] * Δx1m) / (Δx1m * Δx1) - 2 * y[1, i2, $k] / (Δx1p * Δx1m)) : ((y[end, i2, $k] + bc[end, i2, $k] * Δx1p) / (Δx1p * Δx1) + y[end - 1, i2, $k] / (Δx1m * Δx1) - 2 * y[end, i2, $k] / (Δx1p * Δx1m))))))
        push!(expr, Expr(:(=), Symbol(solname, statename2, statename2), :((1 < i2 < size(y, 2)) ? (y[i1, i2 + 1, $k] / (Δx2p * Δx2) + y[i1, i2 - 1, $k] / (Δx2m * Δx2) - 2 * y[i1, i2, $k] / (Δx2p * Δx2m)) : ((i2 == 1) ? (y[i1, 2, $k] / (Δx2p * Δx2) + (y[i1, 1, $k] - bc[i1, 1, $k] * Δx2m) / (Δx2m * Δx2) - 2 * y[i1, 1, $k] / (Δx2p * Δx2m)) : ((y[i1, end, $k] + bc[i1, end, $k] * Δx2p) / (Δx2p * Δx2) + y[i1, end - 1, $k] / (Δx2m * Δx2) - 2 * y[i1, end, $k] / (Δx2p * Δx2m))))))
        push!(expr, Expr(:(=), Symbol(solname, statename1, statename2), :((y[min(i1 + 1, size(y, 1)), min(i2 + 1, size(y, 2)), $k] - y[min(i1 + 1, size(y, 1)), max(i2 - 1, 1), $k] - y[max(i1 - 1, 1), min(i2 + 1, size(y, 2)), $k] + y[max(i1 - 1, 1), max(i2 - 1, 1), $k]) / (4 * Δx1 * Δx2))))
    end
    quote
        $(Expr(:meta, :inline))
        i1, i2 = icar[1], icar[2]
        grid1, grid2 = grid.x[1], grid.x[2]
        @inbounds Δx1m = grid1[max(i1, 2)] - grid1[max(i1-1, 1)]
        @inbounds Δx1p = grid1[min(i1+1, size(y, 1))] - grid1[min(i1, size(y, 1) - 1)]
        Δx1 = (Δx1m + Δx1p) / 2
        @inbounds Δx2m = grid2[max(i2, 2)] - grid2[max(i2-1, 1)]
        @inbounds Δx2p = grid2[min(i2+1, size(y, 2))] - grid2[min(i2, size(y, 2) - 1)]
        Δx2 = (Δx2m + Δx2p) / 2
        @inbounds $(Expr(:tuple, expr...))
    end
end


## Turing Complete Interpreter
### A Very Tiny Turing Complete Programming Language defined with denotational semantics

# semantica dalle dispense degano

Mem = Dict{Symbol, Union{Bool, Int}}

read_mem = @theory begin
	(v::Symbol, mem) |> mem[v]
end

@testset "Reading Memory" begin
	@test 2 == rewrite(:((x), $(Mem(:x => 2))), read_mem; order=:inner, m=@__MODULE__)
end

arithm_rules = @theory begin
	(a + b, σ) 		  => (a, σ) + (b, σ)
	(a * b, σ) 		  => (a, σ) * (b, σ)
	(a - b, σ) 		  => (a, σ) - (b, σ)
	(a::Int + b::Int) |> a + b
	(a::Int * b::Int) |> a * b
	(a::Int - b::Int) |> a - b
	(n::Int, σ) |> n
end

eval_arithm(ex, mem) = (@rewriter(read_mem ∪ arithm_rules, :inner))(:($ex, $mem))

@testset "Arithmetic" begin
	@test 5 == eval_arithm(:(2 + 3), Mem())
	@test 4 == eval_arithm(:(2 + x), Mem(:x => 2))
end

# don't need to access memory
bool_rules = @theory begin
	(a::Bool ∨ b::Bool) |> (a || b)
	(a::Bool ∧ b::Bool) |> (a && b)
	(a::Int < b::Int) 	|> (a < b)
	¬a::Bool 			|> !a
	(bv::Bool, σ) 		|> bv
	(a < b, mem1) |> (eval_arithm(a, mem1) < eval_arithm(b, mem1))
	(¬b, σ) |> !eval_bool(b, σ)
	(a ∨ b, σ) => (a, σ) ∨ (b, σ)
	(a ∧ b, σ) => (a, σ) ∧ (b, σ)
end

eval_bool(ex, mem) = (@rewriter(bool_rules, :inner))(:($ex, $mem))

@testset "Booleans" begin
	@test false == eval_bool(:(false ∨ false), Mem())
	@test true == eval_bool(:((false ∨ false) ∨ ¬(false ∨ false)), Mem(:x => 2))
	@test true == eval_bool(:((2 < 3) ∧ (3 < 4)), Mem(:x => 2))
	@test false == eval_bool(:((2 < x) ∨ ¬(3 < 4)), Mem(:x => 2))
	@test true == eval_bool(:((2 < x) ∨ ¬(3 < 4)), Mem(:x => 4))
end

if_rules = @theory begin
	(if guard; t end, σ) => (if guard; t else skip end, σ)
	(if guard; t else f end, σ) |>
		(eval_bool(guard, σ) ? :($t, $σ) : :($f, $σ))
end

eval_if(ex::Expr, mem::Mem) = (@rewriter(read_mem ∪ arithm_rules ∪ if_rules, :inner))(:($ex, $mem))

@testset "If Semantics" begin
	@test 2 == eval_if(:(if true x else 0 end), Mem(:x => 2))
	@test 0 == eval_if(:(if false x else 0 end), Mem(:x => 2))
	@test 2 == eval_if(:(if ¬(false) x else 0 end), Mem(:x => 2))
	@test 0 == eval_if(:(if ¬(2 < x) x else 0 end), Mem(:x => 3))
end

while_rules = @theory begin
	(:skip, σ) => σ
	((:skip; c2), σ) => (c2, σ)
	((c1; c2), σ) |> begin
		r = eval_while(c1, σ);
		(r isa Mem) ? :($c2, $r) : :($c2, $σ)
	end
	(while guard body end, σ) =>
		(if guard; (body; while guard body end) else skip end, σ)
end

write_mem = @theory begin
	(sym::Symbol = val, memory) |>
		(memory[sym] = eval_arithm(val, memory); memory)
end

while_language = @compile_theory write_mem ∪ read_mem ∪ arithm_rules ∪ if_rules ∪ while_rules;

eval_while(ex, mem) = (@rewriter(while_language, :inner))(:($ex, $mem))

@testset "While Semantics" begin
	@test Mem(:x => 3) == eval_while(:((x = 3)), Mem(:x => 2))
	@test Mem(:x => 5) == eval_while( :(x = 4; x = x + 1) , Mem(:x => 3))
	@test Mem(:x => 4) == eval_while( :( if x < 10; x = x + 1 end  ) , Mem(:x => 3))
	@test 10 == eval_while( :( while x < 10; x = x + 1 end ; x ) , Mem(:x => 3))
	@test 50 == eval_while( :( while x < y; (x = x + 1; y = y - 1) end ; x ) , Mem(:x => 0, :y => 100))
end


using DelimitedFiles

function calc_window_sum(inputvals, index)
    inputvals[index] + inputvals[index+1] + inputvals[index+2]
end

inputs = readdlm("input.txt", '\n', Int)
bigger_count = 0
for i in 2:size(inputs)[1]-2
    if calc_window_sum(inputs,i) > calc_window_sum(inputs,i-1)
        global bigger_count += 1
    end
end

print("Number of increases ",bigger_count)


using PyPlot, Revise, Random, ProgressMeter
#includet("2NNlogistic-param-predict.jl")
using3D()

##

"""
Gradient ascent to train a 2-layers ReLU neural net for the soft-min loss
INPUT: X (training input), Y (training output), m (nb neurons), both: training both layers or just the output
OUTPUT: Ws (training trajectory)
"""
function logistic_reg(X, Y, stepsize, niter)
    (n,d) = size(X)
    a     = zeros(d+1)
    as    = zeros(d+1, niter) # store optimization path
    betas = zeros(niter)
    margins = zeros(niter)
    Z = Y .* cat(X, 1*ones(n), dims=2)/10
    for iter = 1:niter
        as[:,iter] = a
        perf = Z*a
        margin = minimum(perf)
        p = exp.( margin .- Z*a)
        p = p./ sum(p)
        grad = Z'*p
        betas[iter] = maximum([1,sqrt(sum(a.^2))])
        margins[iter] = margin/sqrt(sum(a.^2))
        a = a + betas[iter] * stepsize * grad /(sqrt(iter+1))
    end
    as, betas, margins
end

##

"""
Gradient ascent to train a 2-layers ReLU neural net for the soft-min loss
INPUT: X (training input), Y (training output), m (nb neurons), both: training both layers or just the output
OUTPUT: Ws (training trajectory)
"""
function logistic_reg_lin(X, Y, stepsize, niter)
    (n,d) = size(X)
    a     = zeros(d)
    as    = zeros(d, niter) # store optimization path
    betas = zeros(niter)
    margins = zeros(niter)
    Z = Y .* X

    for iter = 1:niter
        as[:,iter] = a
        perf = Z*a
        margin = minimum(perf)
        p = exp.( margin .- Z*a)
        p = p./ sum(p)
        grad = Z'*p
        betas[iter] = maximum([1,sqrt(sum(a.^2))])
        margins[iter] = margin/sqrt(sum(a.^2))
        a = a + betas[iter] * stepsize * grad /(sqrt(iter+1))
    end
    as, betas, margins
end

# Random.seed!(4)
# nx = 3
# nxx = 50
# n = 2*(nx + nxx)
# XX = rand(nx,2)
# A1 = (XX[:,1] .+ max.(XX[:,2] .- XX[:,1], 0) .+0.0)/2.5
# A2 = XX[:,2] .- max.(XX[:,2] .- XX[:,1], 0) .- 0.5
# AA = [0; -0.5]' .+ 0.03*randn(nxx,2)
# XX = rand(nx,2)
# B1 = (XX[:,1] .+ min.(XX[:,2] .- XX[:,1], 0) .-1.0)/2.5
# B2 = XX[:,2] .- min.(XX[:,2] .- XX[:,1], 0) .-0.5
# BB = [-0.5; -0.5]' .+ 0.03*randn(nxx,2)
# X = [A1 A2;AA; B1 B2;BB]
# Y = cat(ones(nx+nxx), -ones(nx+nxx), dims=1)
# plot(X[Y.==1,1], X[Y.==1,2],"oC0",label=L"Y=+1")
# plot(X[Y.==-1,1], X[Y.==-1,2],"oC1",label=L"Y=-1")
# legend()

Random.seed!(4)
nx = 3
nxx = 50
n = 2*(nx + nxx)
XX = rand(nx,2)
A1 = (XX[:,1] .+ max.(XX[:,2] .- XX[:,1], 0) .+0.0)/2.5
A2 = XX[:,2] .- max.(XX[:,2] .- XX[:,1], 0) .- 0.5
AA = [0; -0.5]' .+ 0.03*randn(nxx,2)
XX = rand(nx,2)
B1 = (XX[:,1] .+ min.(XX[:,2] .- XX[:,1], 0) .-1.0)/2.5
B2 = XX[:,2] .- min.(XX[:,2] .- XX[:,1], 0) .-0.5
BB = [-0.5; -0.5]' .+ 0.03*randn(nxx,2)
X = [A1 A2;AA; B1 B2;BB]
Y = cat(ones(nx+nxx), -ones(nx+nxx), dims=1)
plot(X[Y.==1,1], X[Y.==1,2],"oC0",label=L"Y=+1")
plot(X[Y.==-1,1], X[Y.==-1,2],"oC1",label=L"Y=-1")
legend()

stepsize = 0.05
niter = 100000
as, betas, margins = logistic_reg(X, Y, stepsize, niter)
abars = as ./ betas'
plot(margins[end] .- margins)

# Create a sphere
fig = figure(figsize=[10,5])
ax1 = subplot(121, projection="3d")

res = 50
phi = range(0.0, 2*π, length=res)#.* ones(res)'

#theta =  ones(res) .* range(0.0, 2π, length=res)'
#x = sin.(phi) .* cos.(theta)
#y = sin.(phi) .* sin.(theta)
#z = cos.(phi)

plot3D(sin.(phi), cos.(phi),zeros(res),".k",ms=1)
plot3D(sin.(phi), zeros(res),cos.(phi),".k",ms=1)
plot3D(zeros(res),sin.(phi), cos.(phi),".k",ms=1)
#plot_surface(x, y, z,  rstride=1, cstride=1, color="C0", alpha=0.1, linewidth=0,shade=true)
ax1.set_xlim3d(-1.2, 1.2)
ax1.set_ylim3d(-1.2, 1.2)
ax1.set_zlim3d(-1.2, 1.2)
ax1.view_init(15,-80)
plot3D([0],[0],[0],"ok")
#plot3D(abars[1,:], abars[2,:],abars[3,:],"C0",lw=3)
N = max.(atan.(betas)./atan(1),1)
plot3D(abars[1,:] .* N, abars[2,:].* N,abars[3,:],"C1",lw=3)

ax2 = subplot(122)
plot(X[Y.==1,1], X[Y.==1,2],"oC0",label=L"Y=+1")
plot(X[Y.==-1,1], X[Y.==-1,2],"oC1",label=L"Y=-1")
ts = -2:0.1:2
iter=1000
    plot(-(abars[2,iter].*ts .+ abars[3,iter])/abars[1,iter], ts,"k")
iter=niter
    plot(-(abars[2,iter].*ts .+ abars[3,iter])/abars[1,iter], ts,"r")

legend()

stepsize = 0.5
niter = 500000
as, betas, margins = logistic_reg(X, Y, stepsize, niter)

nframes=200
resolution = 0.005
# define the sequence of time steps ts to be plotted
a = (niter-1)/(nframes-1)^4
ts = setdiff(Int.(floor.(a*(0:nframes-1).^4)) .+ 1,1)
as = as[:,ts]
N =tanh.(0.5*sqrt.(sum(as.^2, dims=1)))
alogs = as .* N./ sqrt.(sum(as.^2, dims=1))

@showprogress 1 "Plotting images..." for k = 1:length(ts)
       ioff() # turns off interactive plotting
    fig = figure(figsize=[7,4])
    ax1 = subplot(121, projection="3d")
    ax1.set_position([0,0.1,0.5,0.8])

    if k<21
        indt = 1:k
    else
        indt = 1:k #(k-20):k
    end

    plot3D(alogs[1,indt], alogs[2,indt],alogs[3,indt],"k",lw=1) # tail
    plot3D([alogs[1,k]], [alogs[2,k]],[alogs[3,k]],"o",color="C3", markersize=3)

    res = 50
    phi = range(0.0, 2π, length=res)#.* ones(res)'
    r = 0.5
    plot3D(r*sin.(phi), r*cos.(phi),zeros(res),".k",ms=0.3)
    plot3D(r*sin.(phi), zeros(res),r*cos.(phi),".k",ms=0.3)
    plot3D(zeros(res),r*sin.(phi), r*cos.(phi),".k",ms=0.3)

    ax1.set_xlim3d(-1, 1)
    ax1.set_ylim3d(-1, 1)
    ax1.set_zlim3d(-1, 1)
    ax1.set_xticks([-1/2, 0, 1/2])
    ax1.set_yticks([-1/2, 0, 1/2])
    ax1.set_zticks([-1/2, 0, 1/2])
    ax1.view_init(25-20*sinpi(k/length(ts)),-45*cospi(k/length(ts)))
   # ax1.view_init(25-20*sinpi(0.2),45-45*cospi(0.2))

    ax2 = subplot(122)
    ax2.set_position([0.45,0.25,0.5,0.5])

    f(x1,x2,k) = sum( as[:,k] .* [x1;x2;1.0]) # prediction function

    xs = -0.8:resolution:0.8
    tab = [f(xs[i],xs[j],k) for i=1:length(xs), j=1:length(xs)]
    pcolormesh(xs', xs, tanh.(tab'), cmap="coolwarm", shading="gouraud", vmin=-1.0, vmax=1.0, edgecolor="face")

    xs = -0.8:resolution:0.8
    tab = [f(xs[i],xs[j],k) for i=1:length(xs), j=1:length(xs)]
    contour(xs', xs, tanh.(tab'), levels =0, colors="k", antialiased = true, linewidths=2)

    # plot training set
    X1 = X[(Y .== 1),:]
    X2 = X[(Y .== -1),:]
    plot(X1[:,1],X1[:,2],"+k")
    plot(X2[:,1],X2[:,2],"_k")
    axis("equal");axis("off");
    ax2.set_xticks([-1/2, 0, 1/2])
    ax2.set_yticks([-1/2, 0, 1/2])

    savefig("dynamics_$(k).png",bbox_inches="tight", dpi=300)
    close(fig)
end

ion()

Random.seed!(7)
X = rand(4,2)
X = [-3 -1; -1 1; 1 -2; 2 1]
Y = cat(-ones(2), ones(2), dims=1)
as, betas, margins = logistic_reg_lin(X, Y, 0.1, 1000)
abars = as ./ betas'
a = abars[:,end]

fig = figure(figsize=[4,4*6/7])
plot(X[Y.==1,1], X[Y.==1,2],"oC3",label=L"y=+1")
plot(X[Y.==-1,1], X[Y.==-1,2],"oC0",label=L"y=-1")
ts = -5:0.01:5
fill_betweenx(ts, -a[2].*ts /a[1],4,alpha=0.1,color="C3")
fill_betweenx(ts, -4,-a[2].*ts /a[1],alpha=0.1,color="C0")
arrow(0,0,0.9*a[1],0.9*a[2],lw=2,head_width=0.15,color="k",overhang=0)
text(0.6,-0.2,L"a")
plot([0],[0],"ok")
#text(-0.95,0.1,L"(0,0)")
ts=0.0:0.01:2.0
plot(1.2*cospi.(ts),1.2*sinpi.(ts),":k",label="unit circle")
f = 2.2
plot([X[1,1],X[1,1]+f*a[1]],[X[1,2], X[1,2] + f*a[2]],"--C0")
f = 1.3
plot([X[2,1],X[2,1]+f*a[1]],[X[2,2], X[2,2] + f*a[2]],"--C0")
f = -1.8
plot([X[3,1],X[3,1]+f*a[1]],[X[3,2], X[3,2] + f*a[2]],"--C3")
f = -1.4
plot([X[4,1],X[4,1]+f*a[1]],[X[4,2], X[4,2] + f*a[2]],"--C3")
text(X[1,1]+0.2,X[1,2],L"x_1")
text(X[2,1]-0.5,X[2,2]-0.1,L"x_2")
text(X[3,1]+0.2,X[3,2],L"x_3")
text(X[4,1]+0.2,X[4,2],L"x_4")
axis([-3.5,3.5,-2.5,3.5])
arrow(1.2,1.7,0.8*a[1],0.8*a[2],lw=1,head_width=0.15,color="k",overhang=0)
f =0.2
arrow(1.2+f*a[1],1.7+f*a[2],-f*a[1],-f*a[2],lw=1,head_width=0.15,color="k",overhang=0)
text(1.4,1.8,L"margin of $a$")
ts = -5:0.01:5
plot(-a[2].*ts /a[1], ts,"k",lw=3)
axis("off")
legend(loc=2)
savefig("max_margin.png",dpi=200, bbox_inches="tight")

?arrow

SS(u,v) = - log(exp(log(1+exp(-u))/2 + log(1+exp(-v))/2) - 1)
S(u,v) = - log(exp(-u)/2 +exp(-v)/2)
u = -20:0.05:20
v = -20:0.05:20

plot(u,SS.(u,10))

plot(u[1:end-1],SS.(u,10)[2:end].- SS.(u,10)[1:end-1])

surf(u,v,SS.(u,v'))

show()


@doc raw"""
    function aic(t :: Trace)

Computes the Akaike Information Criterion for a *single trace* (thus replacing the definition) with
"maximum likelihood" by one with "likelihood". The formula is

```math
\text{AIC}(t)/2 = |\text{params}(t)| - \ell(t),
```

where ``\text{params}(t)|`` is the sum of the dimensionalities of non-observed
and non-deterministic sample nodes.
"""
function aic(t :: Trace)
    ll = 0.0
    k = 0
    for v in values(t)
        if !(v.interpretation == DETERMINISTIC || v.observed)
            length(size(v.value)) == 0 ? k += 1 : k += length(v.value)
        end
        if v.observed
            ll += v.logprob_sum
        end
    end
    2.0 * (k - ll)
end

function numparams(t::Trace)
    k = 0
    for v in values(t)
        if !(v.interpretation == DETERMINISTIC || v.observed)
            length(size(v.value)) == 0 ? k += 1 : k += length(v.value)
        end
    end
    k
end

@doc raw"""
    function aic(r :: SamplingResults{I}) where I <: InferenceType

Computes an empirical estimate of the Akaike Information Criterion from a `SamplingResults`.
The formula is

```math
\text{AIC}(r)/2 = \min_{t \in \text{traces}(r)}|\text{params}(t)| - \hat\ell(t),
```

where ``\text{params}(t)|`` is the sum of the dimensionalities of non-observed
and non-deterministic sample nodes and
``\hat\ell(t)`` is the empirical maximum likelihood.
"""
function aic(r :: SamplingResults{I}) where I <: InferenceType
    min_aic = Inf
    for t in r.traces
        a = aic(t)
        if a < min_aic
            min_aic = a
        end
    end
    min_aic
end

@doc raw"""
    function hpds(r::SamplingResults{I}, pct::Float64) where I <: InferenceType

Computes the highest posterior density set (HPDS) of the `SamplingResults` object.
Let ``\mathcal T`` be the set of traces. The ``100\times Q \%``-percentile HPDS is defined as the
set that satisfies ``\sum_{t \in \mathrm{HPDS}} p(t) = Q`` and, for all ``t \in \mathrm{HPDS}``,
``p(t) > p(s)`` for every ``s \in \mathcal T  - \mathrm{HPDS}``.
It is possible to compute the HPDS using the full joint density ``p(t) \equiv p(x, z)``, where ``x`` is the
set of observed rvs and ``z`` is the set of latent rvs, since ``p(z|x) \propto p(x, z)``.

`pct` should be a float in (0.0, 1.0). E.g., `pct = 0.95` returns the 95% HPDS.
"""
function hpds(r::SamplingResults{I}, pct::Float64) where I <: InferenceType
    (pct > 0.0 && pct < 1.0) || error("pct must be in (0, 1)")
    ix = Int(floor(length(r.traces) * pct))
    for t in r.traces
        if t.logprob_sum == 0.0
            logprob!(t)
        end
    end
    # p(z|x) \propto p(x, z), so we can sort by the joint density, which is easy to
    # calculate, instead of needing to approximate the posterior density
    tr_ix = partialsortperm(r.traces, 1:ix, by=(t,) -> t.logprob_sum, rev=true)
    NonparametricSamplingResults{typeof(r.interpretation)}(
        r.interpretation,
        r.log_weights[tr_ix],
        r.return_values[tr_ix],
        r.traces[tr_ix]
    )
end

@doc raw"""
    function hpdi(r::SamplingResults{I}, pct::Float64, addresses::AbstractArray{T}) where {I <: InferenceType, T}

Computes the highest posterior density interval(s) for a univariate variable. Does *not* check that the
data corresponding to each address in `addresses` is actually univariate; if in doubt, use `hpds` instead.
"""
function hpdi(r::SamplingResults{I}, pct::Float64, addresses::AbstractArray{T}) where {I <: InferenceType, T}
    r = hpds(r, pct)
    Dict(a => (minimum(r[a]), maximum(r[a])) for a in addresses)
end

export aic, hpds, hpdi, numparams


export draw_curve

"""
`draw_curve(plist)` draws a closed curve through points in the
`plist` which is a 1-dimensional list of complex numbers.
Use `draw_curve(plist,false)` for an open curve.
"""
function draw_curve(pts::Array{Complex{T},1}, closed::Bool = true; opts...) where {T}
    x = real.(pts)
    y = imag.(pts)
    np = length(pts)

    kind = closed ? :closed : :open
    maxt = closed ? np + 1 : np
    Sx = Spline(x, kind)
    Sy = Spline(y, kind)

    fx = funk(Sx)
    fy = funk(Sy)

    plot!(fx, fy, 1, maxt; opts...)
end


function build_simple_lp!(pre_opt::MOIU.CachingOptimizer)
    optim = MOIB.full_bridge_optimizer(pre_opt, Float64)
    MOI.empty!(optim)
    @test MOI.is_empty(optim)

    # add 10 variables - only diagonal is relevant
    X = MOI.add_variables(optim, 2)

    # add sdp constraints - only ensuring positivenesse of the diagonal
    vov = MOI.VectorOfVariables(X)

    c1 = MOI.add_constraint(optim,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(2.0, X[1]),
            MOI.ScalarAffineTerm(1.0, X[2])
        ], 0.0), MOI.EqualTo(4.0))

    c2 = MOI.add_constraint(optim,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[1]),
            MOI.ScalarAffineTerm(2.0, X[2])
        ], 0.0), MOI.EqualTo(4.0))

    b1 = MOI.add_constraint(optim,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[1])
        ], 0.0), MOI.GreaterThan(0.0))

    b2 = MOI.add_constraint(optim,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[2])
        ], 0.0), MOI.GreaterThan(0.0))

    MOI.set(optim,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0, -3.0], [X[1], X[2]]), 0.0)
        )
    MOI.set(optim, MOI.ObjectiveSense(), MOI.MIN_SENSE)
end

const optimizer_maxiter = MOIU.CachingOptimizer(cache, ProxSDP.Optimizer(max_iter = 1, log_verbose = false))
const optimizer_timelimit = MOIU.CachingOptimizer(cache, ProxSDP.Optimizer(time_limit = 0.0001, log_verbose = false))

@testset "MOI status" begin
    @testset "MOI.OPTIMAL" begin
        build_simple_lp!(optimizer)
        MOI.optimize!(optimizer)
        @test MOI.get(optimizer, MOI.TerminationStatus()) == MOI.OPTIMAL
        MOI.empty!(optimizer)
    end

    @testset "MOI.ITERATION_LIMIT" begin
        build_simple_lp!(optimizer_maxiter)
        MOI.optimize!(optimizer_maxiter)
        @test MOI.get(optimizer_maxiter, MOI.TerminationStatus()) == MOI.ITERATION_LIMIT
        MOI.empty!(optimizer_maxiter)
    end

    # @testset "MOI.TIME_LIMIT" begin
    #     build_simple_lp!(optimizer_timelimit)
    #     MOI.optimize!(optimizer_timelimit)
    #     @test MOI.get(optimizer_timelimit, MOI.TerminationStatus()) == MOI.TIME_LIMIT
    #     MOI.empty!(optimizer_timelimit)
    # end
end


Distributions.rand(rng::AbstractRNG, d::Dirac, n::Int) = hcat(fill(d.value, n)...)

"""
    BudgetSet{T<:Real, D<:ContinuousMultivariateSampleable} <: CenteredAmbiguitySet{T,D}

Bertsimas's uncertainty set:

```math
\\left\\{ \\mu \\; \\middle| \\begin{array}{ll}
s.t.  \\quad \\mu_i \\leq \\hat{r}_i + z_i \\Delta_i \\quad \\forall i = 1:\\mathcal{N} \\\\
\\quad \\quad \\mu_i \\geq \\hat{r}_i - z_i \\Delta_i  \\quad \\forall i = 1:\\mathcal{N} \\\\
\\quad \\quad z_i \\geq 0 \\quad \\forall i = 1:\\mathcal{N} \\\\
\\quad \\quad z_i \\leq 1 \\quad \\forall i = 1:\\mathcal{N} \\\\
\\quad \\quad \\sum_{i}^{\\mathcal{N}} z_i \\leq \\Gamma \\\\
\\end{array}
\\right\\} \\\\
```

Attributes:
- `d::Sampleable{Multivariate, Continous}`: The parent distribution with an uncertain mean
- `Δ::Array{Float64,1}`: Uncertainty around mean. (default: std(d) / 5)
- `Γ::Float64`: Number of assets in worst case. (default: 0.1 * length(d))

References:
- Bertsimas, D. e Sim, M. (2004). The price of robustness. Operations research, 52(1):35–53.

"""
struct BudgetSet{T<:Real, D<:Sampleable} <: CenteredAmbiguitySet{T,D}
    d::D
    Δ::Vector{T}
    Γ::T
    # Inner constructor for validating arguments
    function BudgetSet{T, D}(d::D, Δ::Vector{T}, Γ::T) where {T<:Real, D<:Sampleable}
        length(d) == length(Δ) || throw(ArgumentError(
            "Distribution ($(length(d))) and Δ ($(length(Δ))) are not the same length"
        ))
        all(>=(0), Δ) || throw(ArgumentError("All uncertainty deltas must be >= 0"))
        Γ >= 0 || throw(ArgumentError("Budget must be >= 0"))
        Γ <= length(d) || @warn "Budget should not exceed the distribution length"
        return new{T, D}(d, Δ, Γ)
    end
end

distribution(s::BudgetSet) = s.d

# Default outer constructor
BudgetSet(d::D, Δ::Vector{T}, Γ::T) where {T<:Real, D<:Sampleable} = BudgetSet{T, D}(d, Δ, Γ)

"""
    calculate_measure!(measure::ExpectedReturn{BudgetSet,WorstCase}, w)

Returns worst case return (WCR) in BudgetSet's uncertainty set ([`BudgetSet`](@ref)).

For further reading:
- BudgetSet' paper: BudgetSet, D. e Sim, M. (2004). The price of robustness. Operations research, 52(1):35–53.
- Original implementation from: https://github.com/andrewrosemberg/PortfolioOpt.jl.

WCR is defined by the following primal problem:

```math
\\begin{aligned}
\\min_{\\mu, z} \\quad & \\mu ' w \\\\
s.t. \\quad & \\mu_i \\leq \\hat{r}_i + z_i \\Delta_i \\quad \\forall i = 1:\\mathcal{N} \\quad &: \\pi^-_i \\\\
& \\mu_i \\geq \\hat{r}_i - z_i \\Delta_i  \\quad \\forall i = 1:\\mathcal{N} \\quad &: \\pi^+_i \\\\
& z_i \\geq 0 \\quad \\forall i = 1:\\mathcal{N} \\\\
& z_i \\leq 1 \\quad \\forall i = 1:\\mathcal{N} \\quad &: \\theta_i \\\\
& \\sum_{i}^{\\mathcal{N}} z_i \\leq \\Gamma \\quad : \\lambda \\\\
\\end{aligned}
```

Which is equivalent to the following dual problem:

```math
\\begin{aligned}
\\max_{\\lambda, \\pi^-, \\pi^+, \\theta} \\quad  \\sum_{i}^{\\mathcal{N}} (\\hat{r}_i (\\pi^+_i - \\pi^-_i) - \\theta_i ) - \\Gamma \\lambda\\\\
s.t. \\quad &  w_i = \\pi^+_i - \\pi^-_i  \\quad \\forall i = 1:\\mathcal{N} \\\\
&  \\Delta_i (\\pi^+_i + \\pi^-_i) - \\theta_i \\leq \\lambda \\quad \\forall i = 1:\\mathcal{N} \\\\
& \\lambda \\geq 0 , \\; \\pi^- \\geq 0 , \\; \\pi^+ \\geq 0 , \\; \\theta \\geq 0 \\\\
\\end{aligned}
```

To avoid solving an optimization problem we enforece the dual constraints in
the upper level problem and return the objective expression (a lower bound of the optimum).

Arguments:
 - `w`: portfolio optimization investment variable ("weights").
 - `s::BudgetSet`: Struct containing atributes of BudgetSet's uncertainty set.
"""
function calculate_measure!(measure::ExpectedReturn{S,WorstCase}, w)  where {S<:BudgetSet}
    model = owner_model(w)
    s = ambiguityset(measure)

    # parameters
    means = Vector(Distributions.mean(s.d))
    n = length(s)
    Δ = s.Δ
    Γ = s.Γ
    # dual variables
    @variable(model, λ >= 0)
    @variable(model, π_neg[i=1:n] >= 0)
    @variable(model, π_pos[i=1:n] >= 0)
    @variable(model, θ[i=1:n] >= 0)
    # constraints: from duality
    @constraints(
        model,
        begin
            constrain_dual1[i=1:n], w[i] == π_pos[i] - π_neg[i]
        end
    )
    @constraints(
        model,
        begin
            constrain_dual2[i=1:n], Δ[i] * (π_pos[i] + π_neg[i]) - θ[i] <= λ
        end
    )

    return return dot(means, w) - sum(θ) - λ * Γ
end

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

"""
    EllipticalSet <: CenteredAmbiguitySet{T,D}

```math
\\left\\{ \\mu \\; \\middle| \\begin{array}{ll}
s.t.  \\quad \\sqrt{(\\hat{r} - \\mu) ' \\Sigma^{-1} (\\hat{r} - \\mu)} \\leq \\delta \\\\
\\end{array}
\\right\\} \\\\
```

Atributes:
- `d::Sampleable{Multivariate, Continous}`: The parent distribution with an uncertain mean
- `Δ::Array{Float64,1}`: Uniform uncertainty around mean. (default: 0.025)

References:
- Ben-Tal, A. e Nemirovski, A. (2000). Robust solutions of linear programming problems contaminated with uncertain data. Mathematical programming, 88(3):411–424.

For more information on how BenTal uncertainty sets are used for RO, please review
the PortfolioOptimization.jl [docs](https://invenia.pages.invenia.ca/PortfolioOptimization.jl/).
"""
struct EllipticalSet{T<:Real, D<:ContinuousMultivariateSampleable} <: CenteredAmbiguitySet{T,D}
    d::D
    Δ::T

    # Inner constructor for validating arguments
    function EllipticalSet{T, D}(d::D, Δ::T) where {T<:Real, D<:ContinuousMultivariateSampleable}
        Δ >= 0 || throw(ArgumentError("Uncertainty delta must be >= 0"))
        return new{T, D}(d, Δ)
    end
end

# Default outer constructor
EllipticalSet(d::D, Δ::T) where {T<:Real, D<:ContinuousMultivariateSampleable} = EllipticalSet{T, D}(d, Δ)

distribution(s::EllipticalSet) = s.d

"""
    calculate_measure!(measure::ExpectedReturn{EllipticalSet,WorstCase}, w)

Returns worst case return (WCR) in EllipticalSet's uncertainty set ([`EllipticalSet`](@ref)).

WCR is defined by the following primal problem:

```math
\\begin{aligned}
\\min_{\\mu} \\quad & \\mu ' w \\\\
s.t. \\quad & ||Σ^{-\\frac{1}{2}} (\\mu - \\hat{r}) || \\leq \\delta \\quad &: \\theta \\\\
\\end{aligned}
```

Which is equivalent to the following dual problem:

```math
\\begin{aligned}
\\max_{\\theta} \\quad &  w ' \\hat{r} - \\theta \\delta \\\\
s.t. \\quad & ||Σ^{\\frac{1}{2}}  w || \\leq \\theta \\\\
\\end{aligned}
```

To avoid solving an optimization problem we enforece the dual constraints in
the upper level problem and return the objective expression (a lower bound of the optimum).

Arguments:
 - `w`: portfolio optimization investment variable ("weights").
 - `s:: EllipticalSet `: Struct containing atributes of EllipticalSet's uncertainty set.
"""
function calculate_measure!(measure::ExpectedReturn{S,WorstCase}, w)   where {S<:EllipticalSet}
    model = owner_model(w)
    s = ambiguityset(measure)

    # parameters
    means = Vector(Distributions.mean(s.d))
    Σ = PDMat(Symmetric(cov(s.d)))
    sqrt_Σ = collect(Σ.chol.U)

    Δ = s.Δ
    # dual variables
    @variable(model, θ)
    # constraints: from duality
    @constraint(model, [θ; sqrt_Σ * w] in JuMP.SecondOrderCone())

    return dot(w, means) - θ * Δ
end


function completes_successfully(swap_line)
    instructions = split.(readlines("day8/input.txt"), " ")

    if instructions[swap_line][1] == "acc"
        return false
    end
    instructions[swap_line][1] = instructions[swap_line][1] == "jmp" ? "nop" : "jmp"
    visited_lines = Set{Int}()

    function execute_line(line_num, instruction, num, acc)
        if instruction == "nop"
            return line_num+1, acc
        elseif instruction == "jmp"
            return line_num+num, acc
        elseif instruction == "acc"
            return line_num+1, acc+num
        end
    end

    num, acc = execute_line(1, instructions[1][1], parse(Int, instructions[1][2]), 0)
    push!(visited_lines, num)
    while true
        acc
        num, acc = execute_line(num, instructions[num][1], parse(Int, instructions[num][2]), acc)
        if num in visited_lines || !(1 <= num <= 652)
            return false
        end
        if num == 652
            return acc
        end
        push!(visited_lines, num)
    end
    return false
end

println([completes_successfully(n) for n in 1:651])

# see if you can spot the answer in the output!
# that was a bit messy :P


using IntervalWavelets
using Test

include("DyadicRationals.jl")
include("Filters.jl")
include("InteriorScalingFunction.jl")
include("BoundaryScalingFunction.jl")
include("Basis.jl")


using LinearAlgebra, StaticArrays, Attitude, SatelliteDynamics
using Infiltrator, ProgressMeter

function initial_conditions()

    # epc0 = Epoch(2019, 1, 1, 12, 0, 0, 0.0)

    # Declare initial state in terms of osculating orbital elements
    rp = R_EARTH + 200e3;ra = R_EARTH + 35000e3
    a = (ra+rp)/2
    e = (ra-rp)/(ra+rp)
    i = deg2rad(27)
    # i = 0
    ω = deg2rad(34)
    Ω = deg2rad(133)
    M = deg2rad(0)
    oe0  = [a,e,i,Ω,ω,M]

    # Convert osculating elements to Cartesean state
    eci0 = sOSCtoCART(oe0, use_degrees=false)

    # modified equinoctial elements
    p = a*(1-e^2)
    f = e*cos(ω + Ω)
    g = e*sin(ω + Ω)
    h = tan(i/2)*cos(Ω)
    k = tan(i/2)*sin(Ω)
    L = wrap_to_2pi(ω + Ω + M)

    eoe0 = [p;f;g;h;k;L]

    r_eci = eci0[1:3]
    v_eci = eci0[4:6]

    p_0 = u_from_x(r_eci)                          # u
    p_prime_0 = uprime_from_xdot(v_eci,p_0)     # du/ds
    h_0 = GM_EARTH/norm(r_eci) - .5*norm(v_eci)^2         # m^2/s^2
    t_0 = 0

    ks0 = [p_0;p_prime_0;h_0;t_0]

    return eci0, eoe0, ks0
end
function rk4(f::Function, x_n, dt)
    k1 = dt*f(x_n)
    k2 = dt*f(x_n+k1/2)
    k3 = dt*f(x_n+k2/2)
    k4 = dt*f(x_n+k3)
    return (x_n + (1/6)*(k1+ 2*k2 + 2*k3 + k4))
end
function fode_j2(r_eci)
    r = norm(r_eci)
    x,y,z = r_eci
    J2 = J2_EARTH
    μ = GM_EARTH

    # precompute repeated stuff
    Re_r_sqr = 1.5*J2*(R_EARTH/r)^2
    five_z_sqr = 5*z^2/r^2

    return  SA[(μ/r^3)*x*(Re_r_sqr*(five_z_sqr - 1)),
               (μ/r^3)*y*(Re_r_sqr*(five_z_sqr - 1)),
               (μ/r^3)*z*(Re_r_sqr*(five_z_sqr - 3))]
end
function cartesian_dynamics(x)
    r_eci = SA[x[1],x[2],x[3]]
    v_eci = SA[x[4],x[5],x[6]]

    # relavent parameters
    J2 = J2_EARTH
    μ = GM_EARTH

    # eci position stuff
    r = norm(r_eci)
    x,y,z = r_eci

    # precompute repeated stuff
    Re_r_sqr = 1.5*J2*(R_EARTH/r)^2
    five_z_sqr = 5*z^2/r^2

    return  SA[v_eci[1],v_eci[2],v_eci[3],
               (-μ/r^3)*x*(1 - Re_r_sqr*(five_z_sqr - 1)),
               (-μ/r^3)*y*(1 - Re_r_sqr*(five_z_sqr - 1)),
               (-μ/r^3)*z*(1 - Re_r_sqr*(five_z_sqr - 3))]
end
function J2_from_EOE(x)
    """J2 acceleration in RTN from equinoctial orbital elements.
    Args:
        p: a*(1-e^2) semi-latus rectum       (m)
        f: e*cos(ω + Ω)                      (non-dimensional)
        g: e*sin(ω + Ω)                      (non-dimensional)
        h: tan(i/2)*cos(Ω)                   (non-dimensional)
        k: tan(i/2)*sin(Ω)                   (non-dimensional)
        L: ω + Ω + θ argument of latitude    (rad)
    Returns:
        a_rtn: acceleration in the RTN frame  (m/s)
    """

    # unpack state
    p,f,g,h,k,L = x
    J2 =  J2_EARTH

    # precompute trig functions
    sinL = sin(L)
    cosL = cos(L)

    # constants for EOE
    mu = GM_EARTH
    w = 1+f*cosL + g*sinL
    r = p/w
    Re = R_EARTH

    # accelerations
    a_r = -((3*mu*J2*Re^2)/(2*r^4))*(1 -   (  (12*(h*sinL-k*cosL)^2 )
                                                    /   ((1+h^2 + k^2)^2))  )
    a_t = -((12*mu*J2*Re^2)/(r^4))*((  ((h*sinL - k*cosL)*(h*cosL+k*sinL) )
                                                    /  ((1+h^2 + k^2)^2))  )
    a_n = -((6*mu*J2*Re^2)/(r^4))*((  ((h*sinL - k*cosL)*(1-h^2-k^2) )
                                                    /  ((1+h^2 + k^2)^2))  )
    return a_r, a_t, a_n
end
function eoe_dynamics(x)
    """Equations of motion for equinoctial elements in 2-body problem.
    Args:
        epc: time(::epoch)
        x: state
            - p: a*(1-e^2) semi-latus rectum       (m)
            - f: e*cos(ω + Ω)                      (non-dimensional)
            - g: e*sin(ω + Ω)                      (non-dimensional)
            - h: tan(i/2)*cos(Ω)                   (non-dimensional)
            - k: tan(i/2)*sin(Ω)                   (non-dimensional)
            - L: ω + Ω + θ argument of latitude    (rad)
            - mass: spacecraft mass                (kg)
        u: thruster acceleration                   (m/s^2)
        thruster_on: bool for thruster             ()
    Returns:
        ̇x: state derivative
    """

    # unpack state
    p,f,g,h,k,L = x

    # equinoctial parameters for GVE's
    sinL = sin(L)
    cosL = cos(L)
    alpha_2 = h^2 - k^2
    s2 = 1+h^2 + k^2
    w = 1+f*cosL + g*sinL
    r = p/w

    # precompute to save calculations
    sqrtpu = sqrt(p/GM_EARTH)

    # J2
    u_r, u_t, u_n = J2_from_EOE(x)

    # GVE's for equinoctial orbital elements
    p_dot = (2*p/w)*sqrtpu*u_t
    f_dot = sqrtpu*( u_r*sinL    +    ((w+1)*cosL + f)*u_t/w    -
                                            (h*sinL - k*cosL)*g*u_n/w  )
    g_dot = sqrtpu*(-u_r*cosL   +    ((w+1)*sinL + g)*u_t/w    +
                                            (h*sinL - k*cosL)*f*u_n/w  )
    h_dot = sqrtpu*(s2*u_n/(2*w))*cosL
    k_dot = sqrtpu*(s2*u_n/(2*w))*sinL
    L_dot = sqrt(GM_EARTH*p)*(w/p)^2 + (1/w)*sqrtpu*(h*sinL - k*cosL)*u_n

    # return state derivative
    return SA[p_dot,f_dot,g_dot,h_dot,k_dot,L_dot]
end
function rv_from_eoe(eoe_vec)
    """Position and Velocity resolved in ECI
    Args:
        p: a*(1-e^2) semi-latus rectum     (m)
        f: e*cos(ω + Ω)                    (non-dimensional)
        g: e*sin(ω + Ω)                    (non-dimensional)
        h: tan(i/2)*cos(Ω)                 (non-dimensional)
        k: tan(i/2)*sin(Ω)                 (non-dimensional)
        L: ω + Ω + θ argument of latitude  (rad)
    Returns:
        r: position vector from eci origin to spacecraft, resolved in eci. (m)
        v: velocity vector of spacecraft in eci, resolved in eci.          (m/s)
    """

    # unpack eoe
    p,f,g,h,k,L = eoe_vec

    # precompute sin and cosine
    sinL = sin(L)
    cosL = cos(L)

    # various eoe specific parameters
    alf2 = h^2 - k^2
    s2 = 1 + h^2 + k^2
    w = 1 + f*cosL + g*sinL
    r = p/w

    # eci position vector
    r = SA[(r/s2) * (  cosL + alf2*cosL + 2*h*k*sinL),
         (r/s2) * (  sinL - alf2*sinL + 2*h*k*cosL),
         (2*r/s2)*(h*sinL -   k *cosL)]

    # ECI velocity vector
    v = sqrt(GM_EARTH/p)*SA[-(1/(s2)) * (sinL + alf2*sinL - 2*h*k*cosL + g - 2*f*h*k + alf2*g),
         -(1/(s2))* (-cosL + alf2*cosL + 2*h*k*sinL - f + 2*g*h*k + alf2*f),
          (2/s2) * (h*cosL +    k*sinL + f*h + g*k)]

    # @infiltrate
    return SA[r[1],r[2],r[3],v[1],v[2],v[3]]
end

function ks_dynamics(x)

    # unpack state
    p = SVector(x[1],x[2],x[3],x[4])
    p_prime = SVector(x[5],x[6],x[7],x[8])
    h = x[9]
    t = x[10]

    # position in eci
    r_eci = x_from_u(p)

    # j2 accel
    a_j2 = fode_j2(r_eci)

    L = L_fx(p)

    # append 0 to end of P_j2 acceleration vector
    P_j2 = SA[a_j2[1],a_j2[2],a_j2[3],0]

    p_dp = -(h/2)*p + ((dot(p,p))/(2))*L'*P_j2

    return SVector(p_prime[1],p_prime[2],p_prime[3],p_prime[4],
                   p_dp[1],p_dp[2],p_dp[3],p_dp[4],
                   -2*dot(p_prime,L'*P_j2),
                   dot(p,p))
end
function rv_from_ks(p,p_prime)
    x = x_from_u(p)
    xdot = xdot_from_u(p,p_prime)
    return SA[x[1],x[2],x[3],xdot[1],xdot[2],xdot[3]]
end
function ϵ(rv)
    r = SA[rv[1],rv[2],rv[3]]
    v = SA[rv[4],rv[5],rv[6]]
    return dot(v,v)/2 - GM_EARTH/norm(r)
end
# function driver()
#     eci0, eoe0, ks0 = initial_conditions()
#
#     tf = 1*24*3600
#     N = 100000
#     dt = tf/N
#
#     X_eci = [@SVector zeros(6) for i = 1:N]
#     X_eci[1] = SVector{6}(eci0)
#     X_eoe = [@SVector zeros(6) for i = 1:N]
#     X_eoe[1] = eoe0
#     X_ks = [@SVector zeros(10) for i = 1:N]
#     X_ks[1] = ks0
#
#     # one day in fictitious time (this was guess and check)
#     sf = 0.003709551289
#     ds = sf/N
#
#     @inbounds for i = 1:(N-1)
#         X_eci[i+1] = rk4(cartesian_dynamics,X_eci[i],dt)
#         X_eoe[i+1] = rk4(eoe_dynamics,X_eoe[i],dt)
#         X_ks[i+1] = rk4(ks_dynamics,X_ks[i],ds)
#     end
#
#     ksm = mat_from_vec(X_ks)
#
#     ps = ksm[1:4,:]
#
#     # mat"
#     # figure
#     # hold on
#     # plot($ps')
#     # hold off
#     # "
#
#     # @show tf - X_ks[end][10]
#     X_eoe_eci = [rv_from_eoe(X_eoe[i]) for i = 1:length(X_eoe)]
#     X_ks_eci = [rv_from_ks(X_ks[i][1:4],X_ks[i][5:8]) for i = 1:length(X_ks)]
#
#     eci_m = mat_from_vec(X_eci)
#     ks_m = mat_from_vec(X_ks_eci)
#
#     # energies
#     ks_ϵ = ϵ(X_ks_eci[end])
#     eci_ϵ = ϵ(X_eci[end])
#     eoe_ϵ = ϵ(X_eoe_eci[end])
#
#     @show ks_ϵ
#     @show eci_ϵ
#     @show eoe_ϵ
#
#
#     # mat"
#     # figure
#     # hold on
#     # plot3($eci_m(1,:),$eci_m(2,:),$eci_m(3,:))
#     # plot3($ks_m(1,:),$ks_m(2,:),$ks_m(3,:))
#     # hold off
#     # "
#
#
#     # # @show "yes"
#     # eci_m = mat_from_vec(X_eci)
#     # eoe_m = mat_from_vec(X_eoe_eci)
#     # mat"
#     # figure
#     # hold on
#     # plot3($eci_m(1,:),$eci_m(2,:),$eci_m(3,:))
#     # plot3($eoe_m(1,:),$eoe_m(2,:),$eoe_m(3,:))
#     # hold off
#     # "
#     #
#     #
#     # @show norm(X_eci[end] - X_eoe_eci[end])
#     # @show norm(X_eci[end] - X_ks_eci[end])
# end
function driver(N,eci0, eoe0, ks0 )
    # eci0, eoe0, ks0 = initial_conditions()

    # N = 100

    X_eci = [@SVector zeros(6) for i = 1:N]
    X_eci[1] = SVector{6}(eci0)
    X_eoe = [@SVector zeros(6) for i = 1:N]
    X_eoe[1] = eoe0
    X_ks = [@SVector zeros(10) for i = 1:N]
    X_ks[1] = ks0

    # one day in fictitious time (this was guess and check)
    sf = 10*0.0003709551289
    ds = sf/(N)

    @inbounds for i = 1:(N-1)
        X_ks[i+1] = rk4(ks_dynamics,X_ks[i],ds)
    end
    tf = X_ks[end][10]
    # @show tf
    tf += tf/N
    dt = tf/N
    @inbounds for i = 1:(N-1)
        X_eci[i+1] = rk4(cartesian_dynamics,X_eci[i],dt)
        X_eoe[i+1] = rk4(eoe_dynamics,X_eoe[i],dt)
    end

    X_eoe_eci = [rv_from_eoe(X_eoe[i]) for i = 1:length(X_eoe)]
    X_ks_eci = [rv_from_ks(X_ks[i][1:4],X_ks[i][5:8]) for i = 1:length(X_ks)]

    # eci_m = mat_from_vec(X_eci)
    # ks_m = mat_from_vec(X_ks_eci)

    # true energy
    # true_ϵ = -8.336482081923751e6
    true_ϵ = -8.3364820836748555e6
    # energies
    ks_ϵ = ϵ(X_ks_eci[end])
    eci_ϵ = ϵ(X_eci[end])
    eoe_ϵ = ϵ(X_eoe_eci[end])

    # @show ks_ϵ
    # @show eci_ϵ
    # @show eoe_ϵ
    # error()

    ks_error = -100*abs(ks_ϵ - true_ϵ)/true_ϵ
    eci_error = -100*abs(eci_ϵ- true_ϵ)/true_ϵ
    eoe_error = -100*abs(eoe_ϵ- true_ϵ)/true_ϵ
    # ks_error = abs(ks_ϵ - true_ϵ)
    # eci_error = abs(eci_ϵ- true_ϵ)
    # eoe_error = abs(eoe_ϵ- true_ϵ)
    # @show ks_ϵ
    # @show eci_ϵ
    # @show eoe_ϵ
    #
    # @show ks_error
    # @show eci_error
    # @show eoe_error
    #
    #
    # mat"
    # figure
    # hold on
    # %plot($eci_m(1,:),$eci_m(2,:),'o')
    # plot($eci_m(1,:),$eci_m(2,:),'o')
    # plot($ks_m(1,:),$ks_m(2,:))
    # hold off
    # "
    #
    # @show norm(X_eci[end] - X_ks_eci[end])
    # @show norm(X_eci[end] - X_eoe_eci[end])
    #
    # @show X_eci[end]
    return SA[ks_error, eci_error, eoe_error]
end

function grid_search()
    eci0, eoe0, ks0 = initial_conditions()

    N_vec = 15:1:300
    errors = [@SVector zeros(3) for i = 1:length(N_vec)]

    @showprogress "simulating" for i = 1:length(N_vec)
        errors[i] = driver(N_vec[i],eci0, eoe0, ks0 )
        # errors[i] = driver(10000000,eci0, eoe0, ks0 )
    end

    errors = mat_from_vec(errors)
    mat"
    figure
    hold on
    plot($errors')
    set(gca, 'YScale', 'log')
    set(gca, 'XScale', 'log')
    ylim([0 10000])
    xlim([15 54])
    legend('KS','ECI','EOE')
    hold off
    "
end
grid_search()


## CPU

# manually vectorized version of trilinear kernel
function _vscattering(p::NTuple{3,A}, xi::NTuple{3,B}, F::Array{Float64,3}) where {A,B}
    # unpack tuples
    dx, dy, dz = @. 1 / (x[2] - x[1], y[2] - y[1], z[2] - z[1])
    px, py, pz = p
    x, y, z = xi

    # indices of lowermost-left corner of the cell
    # containing the particle
    idx_x, idx_y, idx_z = parent_cell(p, dxi)

    # distance from particle to lowermost-left corner of the cell
    dx_particle = px - x[idx_x]
    dy_particle = py - y[idx_y]
    dz_particle = pz - z[idx_z]

    a1 = VectorizationBase.Vec(
        (
            F[idx_x + 1, idx_y, idx_z], # 3
            F[idx_x + 1, idx_y + 1, idx_z], # 4
            F[idx_x, idx_y + 1, idx_z], # 2
            F[idx_x, idx_y, idx_z], # 1
        )...,
    )
    a2 = VectorizationBase.Vec(
        (
            F[idx_x + 1, idx_y, idx_z + 1], # 3
            F[idx_x + 1, idx_y + 1, idx_z + 1], # 4
            F[idx_x, idx_y + 1, idx_z + 1], # 2
            F[idx_x, idx_y, idx_z + 1], # 1
        )...,
    )

    b1 = VectorizationBase.Vec(
        (
            (1 - dx_particle * dx) * (dy_particle * dy) * (1 - dz_particle * dz),     # 3
            (dx_particle * dx) * (dy_particle * dy) * (1 - dz_particle * dz),       # 4
            (dx_particle * dx) * (1 - dy_particle * dy) * (1 - dz_particle * dz),     # 2
            (1 - dx_particle * dx) * (1 - dy_particle * dy) * (1 - dz_particle * dz),   # 1
        )...,
    )
    b2 = VectorizationBase.Vec(
        (
            (1 - dx_particle * dx) * (dy_particle * dy) * (dz_particle * dz),   # 3
            (dx_particle * dx) * (dy_particle * dy) * (dz_particle * dz),     # 4
            (dx_particle * dx) * (1 - dy_particle * dy) * (dz_particle * dz),   # 2
            (1 - dx_particle * dx) * (1 - dy_particle * dy) * (dz_particle * dz), # 1
        )...,
    )

    # Interpolate field F onto particle
    Fp =
        VectorizationBase.vsum(VectorizationBase.vmul(a1, b1)) +
        VectorizationBase.vsum(VectorizationBase.vmul(a2, b2))

    return Fp
end

function _vscattering(p::NTuple{3,A}, xi::NTuple{3,B}, F::Array{Float32,3}) where {A,B}
    # unpack tuples
    dx, dy, dz = @. 1 / (x[2] - x[1], y[2] - y[1], z[2] - z[1])
    px, py, pz = p
    x, y, z = xi

    # indices of lowermost-left corner of the cell
    # containing the particle
    idx_x, idx_y, idx_z = parent_cell(p, dxi)

    # distance from particle to lowermost-left corner of the cell
    dx_particle = px - x[idx_x]
    dy_particle = py - y[idx_y]
    dz_particle = pz - z[idx_z]

    a = VectorizationBase.Vec(
        (
            F[idx_x + 1, idx_y, idx_z], # 3
            F[idx_x + 1, idx_y + 1, idx_z], # 4
            F[idx_x, idx_y + 1, idx_z], # 2
            F[idx_x, idx_y, idx_z], # 1
            F[idx_x + 1, idx_y, idx_z + 1], # 3
            F[idx_x + 1, idx_y + 1, idx_z + 1], # 4
            F[idx_x, idx_y + 1, idx_z + 1], # 2
            F[idx_x, idx_y, idx_z + 1], # 1
        )...,
    )

    b = VectorizationBase.Vec(
        (
            (1 - dx_particle * dx) * (dy_particle * dy) * (1 - dz_particle * dz),     # 3
            (dx_particle * dx) * (dy_particle * dy) * (1 - dz_particle * dz),       # 4
            (dx_particle * dx) * (1 - dy_particle * dy) * (1 - dz_particle * dz),     # 2
            (1 - dx_particle * dx) * (1 - dy_particle * dy) * (1 - dz_particle * dz),   # 1
            (1 - dx_particle * dx) * (dy_particle * dy) * (dz_particle * dz),       # 3
            (dx_particle * dx) * (dy_particle * dy) * (dz_particle * dz),         # 4
            (dx_particle * dx) * (1 - dy_particle * dy) * (dz_particle * dz),       # 2
            (1 - dx_particle * dx) * (1 - dy_particle * dy) * (dz_particle * dz),     # 1
        )...,
    )

    # Interpolate field F onto particle
    Fp = VectorizationBase.vsum(VectorizationBase.vmul(a, b))

    return Fp
end

function _scattering(
    p::NTuple{3,A}, dxi::NTuple{3,A}, xi::NTuple{3,B}, F::Array{C,3}
) where {A,B,C}
    # unpack tuples
    dx, dy, dz = dxi
    px, py, pz = p
    x, y, z = xi

    # check that the particle is inside the grid
    isinside(px, py, pz, x, y, z)

    # indices of lowermost-left corner of the cell
    # containing the particle
    idx_x, idx_y, idx_z = parent_cell(p, dxi)

    # distance from particle to lowermost-left corner of the cell
    tx = (px - x[idx_x]) / dx
    ty = (py - y[idx_y]) / dy
    tz = (pz - z[idx_z]) / dz

    # Interpolate field F onto particle
    Fp = trilinear(
        tx,
        ty,
        tz,
        F[idx_x, idx_y, idx_z],   # v000
        F[idx_x + 1, idx_y, idx_z],   # v100
        F[idx_x, idx_y, idx_z + 1], # v001
        F[idx_x + 1, idx_y, idx_z + 1], # v101
        F[idx_x, idx_y + 1, idx_z],   # v010
        F[idx_x + 1, idx_y + 1, idx_z],   # v110
        F[idx_x, idx_y + 1, idx_z + 1], # v011
        F[idx_x + 1, idx_y + 1, idx_z + 1], # v111
    )
    return Fp
end

function scattering(xi, F::Array{T,3}, particle_coords) where {T}
    # unpack tuples
    dxi = grid_size(xi)
    np = length(particle_coords[1])
    Fp = zeros(T, np)

    Threads.@threads for i in 1:np
        @inbounds Fp[i] = _scattering(
            (particle_coords[1][i], particle_coords[2][i], particle_coords[3][i]),
            dxi,
            xi,
            F,
        )
    end

    return Fp
end

function scattering!(Fp, xi, F::Array{T,3}, particle_coords) where {T}
    # unpack tuples
    x, y, z = xi
    dxi = (x[2] - x[1], y[2] - y[1], z[2] - z[1])
    np = length(particle_coords[1])

    Threads.@threads for i in 1:np
        @inbounds Fp[i] = _scattering(
            (particle_coords[1][i], particle_coords[2][i], particle_coords[3][i]),
            dxi,
            xi,
            F,
        )
    end
end

## CUDA

function _scattering!(
    Fp::CuDeviceVector{T,1},
    p::NTuple{3,CuDeviceVector{T,1}},
    dxi::NTuple{3,T},
    xi::NTuple{3,A},
    F::CuDeviceArray{T,3},
) where {A,T}
    ix = (blockIdx().x - 1) * blockDim().x + threadIdx().x

    # unpack tuples
    dx, dy, dz = dxi
    px, py, pz = p
    x, y, z = xi

    @inbounds if ix ≤ length(px)

        # check that the particle is inside the grid
        isinside(px[ix], py[ix], pz[ix], x, y, z)

        # indices of lowermost-left corner of the cell
        # containing the particle
        idx_x, idx_y, idx_z = parent_cell((px[ix], py[ix], pz[ix]), dxi)

        # # distance from particle to lowermost-left corner of the cell
        tx = (px[ix] - x[idx_x]) / dx
        ty = (py[ix] - y[idx_y]) / dy
        tz = (pz[ix] - z[idx_z]) / dz

        # Interpolate field F onto particle
        Fp[ix] = trilinear(
            tx,
            ty,
            tz,
            F[idx_x, idx_y, idx_z],   # v000
            F[idx_x + 1, idx_y, idx_z],   # v100
            F[idx_x, idx_y, idx_z + 1], # v001
            F[idx_x + 1, idx_y, idx_z + 1], # v101
            F[idx_x, idx_y + 1, idx_z],   # v010
            F[idx_x + 1, idx_y + 1, idx_z],   # v110
            F[idx_x, idx_y + 1, idx_z + 1], # v011
            F[idx_x + 1, idx_y + 1, idx_z + 1], # v111
        )
    end

    return nothing
end

function scattering(
    xi, Fd::CuArray{T,3}, particle_coords::NTuple{3,CuArray}; nt=512
) where {T}
    x, y, z = xi
    dxi = (x[2] - x[1], y[2] - y[1], z[2] - z[1])
    N = length(particle_coords[1])
    Fpd = CuArray{T,1}(undef, N)

    numblocks = ceil(Int, N / nt)
    CUDA.@sync begin
        @cuda threads = nt blocks = numblocks _scattering!(
            Fpd, particle_coords, dxi, xi, Fd
        )
    end

    return Fpd
end

function scattering!(
    Fpd, xi, Fd::CuArray{T,3}, particle_coords::NTuple{3,CuArray}; nt=512
) where {T}
    x, y, z = xi
    dxi = (x[2] - x[1], y[2] - y[1], z[2] - z[1])
    N = length(particle_coords[1])
    numblocks = ceil(Int, N / nt)
    CUDA.@sync begin
        @cuda threads = nt blocks = numblocks _scattering!(
            Fpd, particle_coords, dxi, xi, Fd
        )
    end
end


using SparseArrays
using LinearAlgebra
using QuadGK

mutable struct modelParameters
    N
    c1
    c2
    K
    xf
    τ
    A1
    solDir
    T
    tMax
    solAdj
    phaseSensNorm
end

function setParams(N,c1,c2,K,xf,τ,print=false)
	D  = zeros(N,N)
	Om = zeros(N,N)
	Id = zeros(N,N)
    for i in 1:N
        D[i,i]  = 2*damping(c1,c2,i)*i*π
        Om[i,i] = i*i*π*π
        Id[i,i] = 1
    end
    A1 = [zeros(N,N) Id;-Om -D]
    A1 = sparse(A1)
    params = modelParameters(N,c1,c2,K,xf,τ,A1,false,false,false,false,false)
    if(print)
	println("#####################################")
	println("## Rijke tube Galerkin modes model ##")
	println("#####################################")
	println("\n##  Parameters ##")
	println("N  = ",params.N)
	println("c1 = ",params.c1)
	println("c2 = ",params.c2)
	println("K  = ",params.K)
	println("xf = ",params.xf)
	println("τ  = ",params.τ)
	println("#################\n")
    end
    return params
end

function damping(c1,c2,i)
    return (1/(2π))*(c1*i+c2*sqrt(1/i))
end

function flameVel(xd,p)
    uf = 0
    for i in 1:p.N
        uf = uf + xd[i]*cos(i*π*p.xf)
    end
    return uf
end

function RHS(u,h,p,t)

    uf = flameVel(h(p,t-p.τ),p)

    y = p.A1*u

    for i in 1:p.N
        y[N+i] = y[N+i] - i*π*p.K*(sqrt(abs(1/3+uf))-sqrt(1/3))*sin(i*π*p.xf)
    end
    return y
end

function pressureVelocity(modes,x,params)
    # Returns the pressure and velocity from the Galerkin modes at x
    u = zeros(length(x),)
    p = zeros(length(x),)
    for i in 1:params.N
	    u += modes[i]*cos.(i*π*x)
        p -= modes[i+params.N]*sin.(i*π*x)*γ*Ma/(i*π)
    end

    return u,p
end

function energy(modes,params)
    # Calculates the non-dimensional acoustic energy per unit volume
    # from the Galerkin modes
    E = 0
    for i=1:params.N
        E += modes[i]^2+(modes[params.N+i]/(i*π))^2
    end
    E *= 0.5
    return E
end

## Linear and Adjoint routines
function linHeatRelease(t,p)
    # Creates the time-delayed system matrix A2 from the direct solution at t.
    Q = zeros(p.N,p.N)
    z = p.solDir(t-p.τ)
    ufb = flameVel(z,p)

    for i in 1:p.N
        for j in 1:p.N
            Q[i,j]=0.5*i*π*p.K*sin(i*π*p.xf)*sign(1/3+ufb)*cos(j*π*p.xf)/sqrt(abs(1/3+ufb))
        end
    end
    A2 = -[zeros(p.N,p.N) zeros(p.N,p.N);Q zeros(p.N,p.N)]
    return A2
end

function linHeatReleaseMat(base,p)
    # Creates the linearised matrix A2 from the state base
    Q = zeros(p.N,p.N)
    z = base
    ufb = flameVel(z,p)

    for i in 1:p.N
        for j in 1:p.N
            Q[i,j]=0.5*i*π*p.K*sin(i*π*p.xf)*sign(1/3+ufb)*cos(j*π*p.xf)/sqrt(abs(1/3+ufb))
        end
    end
    A2 = -[zeros(p.N,p.N) zeros(p.N,p.N);Q zeros(p.N,p.N)]
    return A2
end

function adjTime(tAdj)
    # Allows us to solve the adjoint forwards in time
    return params.tMax .- tAdj
end

function RHSadj(u,h,p,t)
    A2 = linHeatRelease(adjTime(t-p.τ),p)
    return p.A1'*u + A2'*h(p,t-p.τ)
end

# Misc functions

function bilinearForm(t,p)
    # Inner product part
    innerProduct = dot(p.solAdj(adjTime(t)),p.solDir(t,Val{1}))
    # Integral part
    integrand(ξ) = dot(p.solAdj(adjTime(t+ξ+p.τ)),linHeatRelease(t+ξ+p.τ,p)*p.solDir(t+ξ,Val{1}))
    integral,error = quadgk(integrand, -p.τ, 0., rtol=1e-4)

    norm = innerProduct + integral
    return norm, innerProduct, integral
end

function phaseCouplingFunction(ϕ,ωf,f,p)
    Tf = 2*π/ωf
    integrand(s) = dot(p.solAdj(adjTime(ϕ+params.T*s/Tf)),f(s))
    integral,error = quadgk(integrand, 0, Tf, rtol=1e-4)
    Γ = integral / Tf
    return Γ
end

function phaseCouplingFunctionmn(ϕ,ωf,f,m,n,p)
    Tf = 2*π/ωf
    integrand(s) = dot(p.solAdj(adjTime(ϕ+(n/m)*params.T*s/Tf)),f(s))
    integral,error = quadgk(integrand, 0, m*Tf, rtol=1e-4)
    Γ = integral / (m*Tf)
    return Γ
end


# using FloatFloat

type FittingMethodCommons{T <: Real}
  solution::Vector{T}
  guessSolution::Vector{T}
  numberOfIterations::Int
  minimumCostValue::Float64
  weights::Vector{T}
  costFunction::FittingCost
end

function FittingMethodCommons(size::Int, gsize::Int)
  solution = zeros(size)
  # solution = Vector{DD}(size)
  guessSolution = zeros(gsize)
  # guessSolution = Vector{DD}(gsize)
  numberOfIterations = 0
  minimumCostValue = 0.0
  weights = zeros(size)
  # weights = Vector{DD}(size)
  curve = NullCurve()
  costFunction = FittingCost(size, curve)

  return FittingMethodCommons(solution, guessSolution, numberOfIterations, minimumCostValue, weights, costFunction)
end

type ExponentialSplinesFitting <: FittingMethod
  constrainAtZero::Bool
  size::Int
  commons::FittingMethodCommons
end

function ExponentialSplinesFitting(constrainAtZero::Bool, size::Int)
  if constrainAtZero
    gsize = 9
  else
    gsize = 10
  end

  commons = FittingMethodCommons(size, gsize)

  return ExponentialSplinesFitting(constrainAtZero, gsize, commons)
end

type SimplePolynomialFitting <: FittingMethod
  constrainAtZero::Bool
  degree::Int
  size::Int
  commons::FittingMethodCommons
end

function SimplePolynomialFitting(constrainAtZero::Bool, degree::Int, size::Int)
  if constrainAtZero
    gsize = degree
  else
    gsize = degree + 1
  end

  commons = FittingMethodCommons(size, gsize)

  return SimplePolynomialFitting(constrainAtZero, degree, gsize, commons)
end

type NelsonSiegelFitting <: FittingMethod
  constrainAtZero::Bool
  size::Int
  commons::FittingMethodCommons
end

function NelsonSiegelFitting(size::Int)
  constrainAtZero = true
  gsize = 4

  commons = FittingMethodCommons(size, gsize)

  return NelsonSiegelFitting(constrainAtZero, gsize, commons)
end

type SvenssonFitting <: FittingMethod
  constrainAtZero::Bool
  size::Int
  commons::FittingMethodCommons
end

function SvenssonFitting(size::Int)
  constrainAtZero = true
  gsize = 6

  commons = FittingMethodCommons(size, gsize)

  return SvenssonFitting(constrainAtZero, gsize, commons)
end

type CubicBSplinesFitting <: FittingMethod
  constrainAtZero::Bool
  size::Int
  knots::Vector{Float64}
  splines::BSpline
  N::Int
  commons::FittingMethodCommons

  function CubicBSplinesFitting(constrainAtZero::Bool, knots::Vector{Float64}, size::Int)
    m = length(knots)
    m >= 8 || error("At least 8 knots are required")

    splines = BSpline(3, m - 4, knots)
    basis_functions = m - 4
    if constrainAtZero
      gsize = basis_functions - 1

      N = 2
      abs(spline_oper(splines, N, 0.0) > QuantLib.Math.EPS_VAL) || error("N_th cubic B-spline must be nonzero at t=0")
    else
      gsize = basis_functions
      N = 1
    end

    commons = FittingMethodCommons(size, gsize)

    new(constrainAtZero, gsize, knots, splines, N, commons)
  end
end

# CubicBSplinesFitting(constrainAtZero::Bool, knots::Vector{Float64}, size::Int) = CubicBSplinesFitting(constrainAtZero, knots, size)

# function ExponentialSplinesFitting(constrainAtZero::Bool, size::Integer)
#   solution = zeros(size)
#   if constrainAtZero
#     gsize = 9
#   else
#     gsize = 10
#   end
#   guessSolution = zeros(gsize)
#   numberOfIterations = 0
#   minimumCostValue = 0.0
#   weights = zeros(size)
#   curve = NullCurve()
#   costFunction = FittingCost(size, curve)
#
#   return ExponentialSplinesFitting(constrainAtZero, gsize,
#           FittingMethodCommons(solution, guessSolution, numberOfIterations, minimumCostValue, weights, costFunction))
# end

guess_size(fitting::ExponentialSplinesFitting) = fitting.constrainAtZero ? 9 : 10
guess_size(fitting::SimplePolynomialFitting) = fitting.constrainAtZero ? fitting.degree : fitting.degree + 1

# Discount functions
function discount_function{T <: Real}(method::ExponentialSplinesFitting, x::Vector{T}, t::Float64)
  d = 0.0
  N = guess_size(method)
  kappa = x[N]
  coeff = 0.0
  if !method.constrainAtZero
    for i = 1:N
      @inbounds d += x[i] * exp(-kappa * (i) * t)
    end
  else
    for i = 1:N - 1
      @inbounds d += x[i]  * exp(-kappa * (i + 1) * t)
      @inbounds coeff += x[i]
    end
    coeff = 1.0 - coeff
    d += coeff * exp(-kappa * t)
  end
  return d
end

function discount_function{T <: Real}(method::SimplePolynomialFitting, x::Vector{T}, t::Float64)
  d = 0.0
  N = method.size

  if !method.constrainAtZero
    for i = 1:N
      @inbounds d += x[i] * get_polynomial(BernsteinPolynomial(), i-1, i-1, t)
    end
  else
    d = 1.0
    for i = 1:N
      @inbounds d += x[i] * get_polynomial(BernsteinPolynomial(), i, i, t)
    end
  end

  return d
end

function discount_function{T <: Real}(method::NelsonSiegelFitting, x::Vector{T}, t::Float64)
  kappa = x[method.size]
  @inbounds zero_rate = x[1] + (x[2] + x[3]) * (1.0 - exp(-kappa * t)) / ((kappa + QuantLib.Math.EPS_VAL) * (t + QuantLib.Math.EPS_VAL)) - (x[3]) * exp(-kappa * t)
  d = exp(-zero_rate * t)

  return d
end

function discount_function{T <: Real}(method::SvenssonFitting, x::Vector{T}, t::Float64)
  kappa = x[method.size - 1]
  kappa_1 = x[method.size]
  eps_v = QuantLib.Math.EPS_VAL

  zero_rate = x[1] + (x[2] + x[3]) * (1.0 - exp(-kappa * t)) / ((kappa + eps_v) * (t + eps_v)) - (x[3]) * exp(-kappa * t) + x[4] * (((1.0 - exp(-kappa_1 * t)) / ((kappa_1 + eps_v) * (t + eps_v))) - exp(-kappa_1 * t))

  d = exp(-zero_rate * t)
  return d
end

function discount_function{T <: Real}(method::CubicBSplinesFitting, x::Vector{T}, t::Float64)
  d = 0.0
  if !method.constrainAtZero
    for i = 1:method.size
      @inbounds d += x[i] * spline_oper(method.splines, i, t)
    end
  else
    t_star = 0.0
    sum = 0.0
    for i = 1:method.size
      if i < method.N
        @inbounds d += x[i] * spline_oper(method.splines, i, t)
        @inbounds sum += x[i] * spline_oper(method.splines, i, t_star)
      else
        @inbounds d += x[i] * spline_oper(method.splines, i + 1, t)
        @inbounds sum += x[i] * spline_oper(method.splines, i + 1, t_star)
      end
    end
    coeff = 1.0 - sum
    coeff /= spline_oper(method.splines, method.N, t_star)
    d += coeff * spline_oper(method.splines, method.N, t)
  end

  return d
end


module DualArray

import TensorOperations: IndexTuple
import TensorOperations: contract!
import ForwardDiff: Dual
import ForwardDiff

struct DualTensor{N, T}
    A::Array{T}
    ∂::NTuple{N, Array{T}}

    "Construct an instance of `DualTensor` from `Array` components."
    DualTensor(A::Array{T},
               ∂::NTuple{N, Array{T}}) where {N, T} = begin
        all(map(sz -> sz == size(A), size.(∂))) ||
            throw(DimensionMismatch("Tensor data and its derivatives must have the same size."))

        new{N, T}(A, ∂)
    end
end

"Convert `Array` of `Dual`s into `DualTensor`."
DualTensor(A::Array{Dual{Tag, T, N}}) where {N, T, Tag} = begin
    AA = ForwardDiff.value.(A)
    ∂A = (( (x -> x[i∂]).(ForwardDiff.partials.(A)) for i∂=1:N )...,)

    DualTensor(AA, ∂A)
end

include("tensor_opr.jl")

end


# Script to replicate figure 7

using DataFrames, StatsBase, Statistics, LinearAlgebra, XLSX, StatsPlots

# -------------- Previous scripts -----------------------------------------------------------------------------------------------------------------------

include("src/data_wrangling.jl") # formats the raw input data obtained from WIOD

# -------------- Sectoral aggregation for country level I/O-Table ---------------------------------------------------------------------------------------

# the function "country_level_head_reis_index" uses the output from the function "obtain_matrices"
# it prepares the data to replicate figure 7 by aggregating the I/O table from a country-sector level to a country level I/O-Table
# and calculates the Head-Reis index (trade costs) for intermediate and final goods
# outputs:
#

function country_level_head_reis_index(df::DataFrame, year::Int64)

    # beginning of the function is the same as in "figure_2_3.jl"

    ID, FD, GO, VA, IV, other = obtain_matrices(df, year) # N*S×N*S, N*S×N, N*S×1, N*S×1, N*S×N, N*S×1

    # aggregate to country level
    ID_ctry = [sum(ID[i:i+n_sectors-1, j:j+n_sectors-1]) for i in 1:n_sectors:n_sectors*n_ctrys, j in 1:n_sectors:n_sectors*n_ctrys] # N×N
    FD_ctry = [sum(FD[i:i+n_sectors-1, j]) for i in 1:n_sectors:n_sectors*n_ctrys, j in 1:n_ctrys] # N×N, sum over origin sector
    IV_ctry = [sum(IV[i:i+n_sectors-1, j]) for i in 1:n_sectors:n_sectors*n_ctrys, j in 1:n_ctrys] # N×N, sum over origin sector
    GO_ctry = [sum(GO[i:i+n_sectors-1]) for i in 1:n_sectors:n_sectors*n_ctrys] # N×1, sum over destination countries
    VA_ctry = [sum(VA[i:i+n_sectors-1]) for i in 1:n_sectors:n_sectors*n_ctrys] # N×1

    # net inventory correction
    iv_correction = [GO_ctry[i]/(GO_ctry[i]-IV_ctry[i, j]) for i in 1:n_ctrys, j in 1:n_ctrys] # N×N, to match country wise inventory changes

    ID_ctry = ID_ctry .* iv_correction # N×N, element wise multiplication
    FD_ctry = FD_ctry .* iv_correction # N×N

    # new section of function starts

    # footnote 22 on page 21: substitute 1e-18 for any zero entries otherwise trade costs explode to infinity!
    replace!(x -> x == 0.0 ? 1e-18 : x, ID_ctry)
    replace!(x -> x == 0.0 ? 1e-18 : x, FD_ctry)

    # claculate Head-Reis index
    θ = 5 # assumption on trade elasticity
    τ_ID = [((ID_ctry[i, j] * ID_ctry[j, i]) / (ID_ctry[i, i] * ID_ctry[j, j]))^(-1/(2θ)) for i in 1:n_ctrys, j in 1:n_ctrys] # bilateral trade costs
    τ_FD = [((FD_ctry[i, j] * FD_ctry[j, i]) / (FD_ctry[i, i] * FD_ctry[j, j]))^(-1/(2θ)) for i in 1:n_ctrys, j in 1:n_ctrys]

    # simple average over all country pairs where i<j => upper triangular matrix without within country trade costs (i.e. diagonal = 1.0)
    τ_avg_id = filter(x -> x != 1.0 && x != 0.0, UpperTriangular(τ_ID)) # filter out 1.0 (within country trade costs) & 0.0 (lower triangular matrix)
    τ_avg_fd = filter(x -> x != 1.0 && x != 0.0, UpperTriangular(τ_FD))

    # footnote 23 on page 21: drop largest 1% of values
    sort!(τ_avg_id)
    sort!(τ_avg_fd)
    τ_avg_id = mean(τ_avg_id[1:end-10]) # use 10/820 instead of 1%
    τ_avg_fd = mean(τ_avg_fd[1:end-10])

    return τ_avg_id, τ_avg_fd
end

t_fig_7 = DataFrame(year=Int[], τ_ID=Float64[], τ_FD=Float64[])

for i in years
    τ_avg_id, τ_avg_fd = country_level_head_reis_index(df, i)
    push!(t_fig_7, [i τ_avg_id τ_avg_fd])
end

XLSX.writetable("clean/t_fig_7.xlsx", t_fig_7, overwrite=true) # export table to folder as input for plotting figure 7

# -------------- Figure 7 -------------------------------------------------------------------------------------------------------------------

t_fig_7 = DataFrame(XLSX.readtable("clean/t_fig_7.xlsx", "Sheet1")...) # load directly from file
transform!(t_fig_7, :year .=> ByRow(Int) .=> :year, renamecols=false) # change type to Int64
transform!(t_fig_7, [:τ_ID, :τ_FD] .=> ByRow(Float64) .=> [:τ_ID, :τ_FD], renamecols=false) # change type to Float64

p_fig7 = plot(t_fig_7.year, [t_fig_7.τ_ID t_fig_7.τ_FD], xlabel="Year", ylabel="Iceberg τ", title="Figure 7: Head-Reis τ's (country level) over time",
                label=["Input-use mean τ" "Final-use mean τ"], ylims=(2.45,5.05), leg=:bottomleft, lw=2)

savefig(p_fig7, "images/figure7.png") # export image to folder


abstract type LimiterFunction end
abstract type minmod <:LimiterFunction end
abstract type superbee <:LimiterFunction end
abstract type vanAlbada <:LimiterFunction end
abstract type vanAlbada2 <:LimiterFunction end
abstract type vanLeer <:LimiterFunction end
abstract type ospre <:LimiterFunction end
abstract type hquick <:LimiterFunction end
abstract type koren <:LimiterFunction end

struct FluxLimiter{T<:LimiterFunction} end

struct KnickPointsParameters{T, N}
    Kr::T
    h::T
    kappa_a::T
    m::T
    n::T
    φ::Type{FluxLimiter{N}}

    function KnickPointsParameters{T, N}(;
        Kr = 5e-7,    # whipple and tucker 1999_ dimensional coeff of erosion
        h = 1.92,    # whipple and tucker 1999
        kappa_a = 4.6071,  # whipple and tucker 1999_ area-length coeff
        m = 1.0,
        n = 1.1,
        φ = FluxLimiter{vanAlbada},
        ) where {T,N}

        new{T, N}(
            Kr,
            h,
            kappa_a,
            m,
            n,
            φ,
        )

    end

end

struct KnickPointsArrays{T}
    L::T
    r::T
    rf::T

    function KnickPointsArrays{T}(River, nn, RiverPhysics) where T
        x, z, Δx = River.x, River.x, River.Δx
        L = riverlength(x, z, nn)
        r = @. -RiverPhysics.Kr * (RiverPhysics.kappa_a^RiverPhysics.m) * (L^(RiverPhysics.h * RiverPhysics.m)) / Δx # EQUATION
        rf = zero(z)
        new{T}(L,r,rf)
    end

end

struct TerraceParameters{T}
    βz::T
    h_wb::T
    P0::T
    P_off::T

    function TerraceParameters{T}(;
        βz = 1e-6,
        h_wb = 15.0,
        P0 = 5e-5,
        P_off = 5e-2,
    ) where T
        new{T}(
            βz,
            h_wb,
            P0,
            P_off,
        )
    end
end

struct HillSlopeParameters{T}
    κa::T
    diffusion::T
    alpha_dif::T

    function HillSlopeParameters{T}(;
        κa = 5e-5,
        diffusion = 0.1 ,
        alpha_dif = 1.0 ,
    ) where T
        new{T}(
            κa,
            diffusion,
            alpha_dif,
        )
    end

end

struct SubmarineParameters{T}
    K_s::T
    λ::T

    function SubmarineParameters{T}(;
        K_s = 1e2,
        λ = 5e-4,
    ) where T
        new{T}(
            K_s,
            λ,
        )
    end

end


#include("HPDE_options.jl")
##########################################################
"""
$(TYPEDEF)

Data type for PDE equations (i.e., problem)
"""
abstract type AbstractHPDEProblem end

abstract type Dimension end
abstract type _1D <: Dimension end
abstract type _2D <: Dimension end
abstract type _3D <: Dimension end

function checkDim(D::Type{<:Dimension})
    if D === _1D
        return 1
    elseif D === _2D
        return 2
    elseif D === _3D
        return 3
    end
    return 0
end

mutable struct FVMPDEProblem
    """
    Flux function F in x direction

    """
    F::Function

    """
    Flux function G in y direction
    !!! compat
        Will be added in future versions
    """
    G::Function

    """
    Flux function H in z direction
    !!! compat
        Will be added in future versions
    """
    H::Function

    """
    Source function
    !!! compat
        Will be added in future versions
    """
    S::Function

    nvars::Int8

    param::Any

    U0

    U_min
    U_max

    dimension::Type{<:Dimension}

    function FVMPDEProblem(func_F::Function, U0; kwargs...)

        zeroFunc(U) = zeros(size(U))
        F = func_F
        D = _1D
        (G,D) = haskey(kwargs,:func_G) ? (kwargs[:func_G],_2D) : (zeroFunc,D)
        (H,D) = haskey(kwargs,:func_G) && haskey(kwargs,:func_H) ? (kwargs[:func_H],_3D) : (zeroFunc,D)
        S = haskey(kwargs,:func_S) ? kwargs[:func_S] : zeroFunc

        @assert size(F(U0)) == size(U0)
        @assert size(G(U0)) == size(U0)
        @assert size(H(U0)) == size(U0)
        @assert size(S(U0)) == size(U0)

        θ = haskey(kwargs,:parameters) ? kwargs[:parameters] : nothing

        U_min = haskey(kwargs,:ul) ? kwargs[:ul] : -Inf * ones(size(U0))
        U_max = haskey(kwargs,:uh) ? kwargs[:uh] : Inf * ones(size(U0))

        @assert size(U_min) == size(U0)
        @assert size(U_max) == size(U0)

        nvars = size(U0,1)

        #@assert checkDim(D) == size(U0,)

        return new(F, G, H, S, nvars, θ, U0, U_min, U_max, D )
    end

end

    # for 2D it should be Array{Float64,4}
    # for 3D it should be Array{Float64,5}
    # this will be fixed later
F(u) = zeros(size(u))
A = FVMPDEProblem(F,[1])


using Base.SparseArrays.CHOLMOD

function chm_analyze_ordering{Tv<:CHOLMOD.VTypes}(A::CHOLMOD.Sparse{Tv}, ordering::Cint, Perm::Vector{Int})
    s = unsafe_load(A.p)
    Parent   = Array(Int,s.nrow)
    Post     = Array(Int,s.nrow)
    ColCount = Array(Int,s.nrow)
    First    = Array(Int,s.nrow)
    Level    = Array(Int,s.nrow)

    ccall((@CHOLMOD.cholmod_name("analyze_ordering",Int),:libcholmod), Cint,
          (Ptr{CHOLMOD.C_Sparse{Tv}}, Cint, Ptr{Int}, Ptr{Int}, Csize_t,
           Ptr{Int}, Ptr{Int}, Ptr{Int},
           Ptr{Int}, Ptr{Int}, Ptr{UInt8}),
           A.p, ordering, Perm, C_NULL, zero(Csize_t), Parent, Post, ColCount, First, Level, CHOLMOD.common())
    Parent, Post, ColCount
end

function pothen_sun{T<:Integer}(par::Vector{T}, post::Vector{T}, colcount::Vector{T})
    n = length(par)
    flag = -ones(T,n)
    ns = n
    for j=post
        mdeg = colcount[j] - 1
        if par[j] != 0 && mdeg == colcount[par[j]] && flag[par[j]] == -1
            # par[j] not assigned to supernode
            ns -= 1
            if flag[j] < 0    # j is a repr. vertex
                flag[par[j]] = j
                flag[j] -= 1
            else              # j is not a repr. vertex
                flag[par[j]] = flag[j]
                flag[flag[j]] -= 1
            end
        end
    end
    ns, flag
end

function chordal_embedding{Tv<:Number,Ti<:Int}(A::SparseMatrixCSC{Tv,Ti}, Perm::Vector{Ti})
    m, n = size(A)
    S = CHOLMOD.Sparse(round(Float64,A))

    F = cholfact(round(Float64,A), shift=n, perm=Perm)
    s = unsafe_load(F.p)
    p = round(Int64,[ unsafe_load(s.Perm, i) for i=1:s.n])

    par, post, colcount = chm_analyze_ordering(S, s.ordering, p)
    ns, flag = pothen_sun(par+1, post+1, colcount)

    # extract cliques
    L = sparse(CHOLMOD.Sparse(F))
    Array{Int,1}[ sort(1+p[L.rowval[L.colptr[i]:L.colptr[i+1]-1]]) for i = find(flag .< 0) ]
end

# If no permutation is specified, use the CHOLMOD ordering
chordal_embedding{Tv<:Number,Ti<:Int}(A::SparseMatrixCSC{Tv,Ti}) = chordal_embedding(A, Array(Int,0))

function clique_index(I::Array{Array{Int,1},1}, S::Array{Int,1}, n::Int)

    J = Int[]
    mask = zeros(Int, n)
    s = sparsevec(S, 1, n)
    for k=1:length(I)
        if length(S) <= length(I[k])
            mask[I[k]] = 1
            if dot(mask, s) == length(S)
                push!(J, k)
            end
            mask[I[k]] = 0
        end
    end
    J
end


mutable struct QP{T <: AbstractFloat} <: AbstractLoss
    Q::Matrix{T}
    q::Vector{T}
    temp::Vector{T}

    function (::Type{QP})(d::Integer, t::Real, ::Type{T} = Float64) where T
        Q = randn(T,d,d)
        q = randn(T,d)
        F = qr(Q)
        temp = t .^ rand(d-2)
        D = Diagonal(vcat(1/t,temp,t))
        mul!(Q, D, Matrix(F.Q))
        lmul!(F.Q', Q)

        return new{T}(Q, q, copy(q))
    end
end

nfeatures(loss::QP) = size(loss.Q,2)

function loss!(loss::QP{T}, x::Vector{T}, g::Vector{T}) where T <: AbstractFloat
    mul!(loss.temp,loss.Q,x)
    result = convert(T, 0.5) * x⋅loss.temp + loss.q⋅x
    g .= loss.temp .+ loss.q
    return result
end