#========================================================================================#
"""
	QuBits

Module QuBits: My Very Own Quantum Computer

Author: Niall Palfreyman, 03/05/22
"""
module QuBits

# Externally available names:
export Tensor, ⊗, kron, ⊚, kpow, isclose, ishermitian, isunitary
export State, qubit, ampl, phase, prob, maxprob, nbits

# Imports:
import Base:kron

using LinearAlgebra

#========================================================================================#
# Tensor definitions:
#-----------------------------------------------------------------------------------------
"""
	Tensor

Tensor is a synonym for AbstractArray.
"""
Tensor = AbstractArray

#-----------------------------------------------------------------------------------------
"""
	Kronecker product ⊗

Define ⊗ as an infix operator for the Kronecker product of two Tensors
"""
⊗ = kron

#-----------------------------------------------------------------------------------------
# Tensor methods:

#-----------------------------------------------------------------------------------------
"""
	kpow( t::Tensor, n::Int)

Kronecker the Tensor t n times with itself.
"""
function kpow( t::Tensor, n::Int)
	if n==0
		# Base case:
		return 1.0
	end

	# Build Kronecker product:
	tn = t
	for _ in 1:(n-1)
		tn = tn ⊗ t
	end
	tn
end
# Use circled ring as operator for kpow:
⊚ = kpow

#-----------------------------------------------------------------------------------------
"""
	isclose( t::Tensor, u::Tensor, atol::Float64=1e-6)

Check whether all elements of the Tensor t are close to those of the Tensor u.
"""
function isclose( t, u, atol::Float64=1e-6)
	all(abs.(t.-u) .< atol)
end

#-----------------------------------------------------------------------------------------
"""
	ishermitian( t::Tensor)

Check whether the Tensor t is hermitian.
"""
function ishermitian( t::Tensor)
	sz = size(t)
	if length(sz) != 2 || sz[1] != sz[2]
		return false
	end

	isclose(t,t')
end

#-----------------------------------------------------------------------------------------
"""
	isunitary( t::Tensor)

Check whether the Tensor t is unitary.
"""
function isunitary( t::Tensor)
	sz = size(t)
	if length(sz) != 2 || sz[1] != sz[2]
		return false
	end

	isclose(t*t',Matrix(I,size(t)))
end

#========================================================================================#
# State definitions:
#-----------------------------------------------------------------------------------------
"""
	State

A State is a collection of one or more qubits: a wrapper for Vector{Complex}.
"""
struct State <: AbstractVector{Complex}
	amp::Vector{Complex}								# The State amplitudes

	function State(amp::Vector)
		if isclose(amp,0.0)
			error("State vectors must be non-zero")
		end
		new(amp/norm(amp))
	end
end

#-----------------------------------------------------------------------------------------
# Delegated methods:
Base.length(s::State) = length(s.amp)
Base.size(s::State) = size(s.amp)
Base.setindex(s::State,i) = setindex(s.amp,i)
Base.getindex(s::State,i) = getindex(s.amp,i)

#-----------------------------------------------------------------------------------------
"""
	kron( s1::State, s2::State)

Return Kronecker product of two States.
"""
function kron( s1::State, s2::State)
	State(kron(s1.amp,s2.amp))
end

#-----------------------------------------------------------------------------------------
"""
	ampl( s::State, bits::BitVector)

Return amplitude for qubit indexed by bits.
"""
function ampl( s::State, bits::BitVector)
	s.amp[1+bits2dec(bits)]					# Add 1 to index because bits counts from 0
end

ampl( s::State, bits...) = ampl(s,BitVector(bits))

#-----------------------------------------------------------------------------------------
"""
	phase( s::State, bits::BitVector)

Return phase of the amplitude of the qubit indexed by bits.
"""
function phase( s::State, bits::BitVector)
	angle(ampl(s,bits))
end

phase( s::State, bits...) = phase(s,BitVector(bits))

#-----------------------------------------------------------------------------------------
"""
	prob( s::State, bits::BitVector)

Return probability for qubit indexed by bits.
"""
function prob( s::State, bits::BitVector)
	amp = ampl(s,bits)						# Retrieve own amplitude
	real(amp'*amp)							# Calculate its squared norm
end

prob( s::State, bits...) = prob(s,BitVector(bits))

#-----------------------------------------------------------------------------------------
"""
	normalise( s::State)

Normalise the state s (but throw an error if it's a zero state).
"""
function normalise(s::State)
	nrm = norm(s)
	if isclose(nrm,0.0)
		error("Attempting to normaise zero-probability state.")
	end
	s /= nrm
end

#-----------------------------------------------------------------------------------------
"""
	qubit( alpha::Complex, beta::Complex)

Create a single qubit State from the given amplitudes.
"""
function qubit( alpha::Number, beta=nothing)
	if beta===nothing
		beta = sqrt(1.0 - alpha'*alpha)
	end

#=	if !isclose( alpha'*alpha+beta'*beta, 1.0)
		error("Qubit probabilities do not add up to 1.")
	end=#

	State([alpha,beta])
end

function qubit(; alpha=nothing, beta=nothing)
	if alpha === beta === nothing
		# Neither amplitude is given:
		error("alpha, beta or both are required.")
	end

	if alpha === nothing
		# alpha not given:
		alpha = sqrt(1.0 - beta'*beta)
	end

	qubit(alpha,beta)
end

#-----------------------------------------------------------------------------------------
"""
	nbits( stt::State)

Return number of qubits in this state.
"""
function nbits( s::State)
	Int(log2(length(s)))
end

#-----------------------------------------------------------------------------------------
"""
	maxprob( stt::State)

Return tuple (bitindex,prob) of highest probability qubit in this State.
"""
function maxprob( s::State)
	mxbindex, mxprob = BitVector([]), 0.0
	for bindex in bitprod(nbits(s))
		thisprob = prob(s,bindex)
		if thisprob > mxprob
			mxbindex, mxprob = bindex, thisprob
		end
	end

	(mxbindex, mxprob)
end

#-----------------------------------------------------------------------------------------
"""
	Base.show( s::State)

Display the given State.
"""
function Base.show( io::IO, s::State)
	print( io, "State{", nbits(s), "}[")
	for i in 1:length(s)-1
		print( round(s[i],digits=4), ", ")
	end
	print( io, round(s[end],digits=4), "]")
end

#========================================================================================#
# Helper methods:
#-----------------------------------------------------------------------------------------
"""
	bits2dec( bits::BitVector)

Compute decimal representation of a BitVector.
"""
function bits2dec( bits::BitVector)
    s = 0; v = 1
   for i in view(bits,length(bits):-1:1)
        s += v*i
        v <<= 1
    end
    s
end

#-----------------------------------------------------------------------------------------
"""
	dec2bits( dec::Int, nbits::Int)

Compute a hi2lo nbit binary representation of a decimal value.
"""
function dec2bits( dec::Int, nbits::Int)
    BitVector(reverse( digits(dec, base=2, pad=nbits)))
end

#-----------------------------------------------------------------------------------------
"""
	bitprod( nbits::Int)

Construct in numerical order a list of all binary numbers containing nbits.
"""
function bitprod( nbits::Int)
    [dec2bits(i,nbits) for i in 0:2^nbits-1]
end

#-----------------------------------------------------------------------------------------
function demo()
	p1 = qubit(alpha=rand())
	x1 = qubit(beta=rand())
	psi = p1 ⊗ x1

	println( isclose( psi'*psi, 1.0))
	println( isclose( p1'*p1*x1'*x1, 1.0))
	println( ampl(psi,1,0))
	println( prob(psi,1,0))
end

end		# ... of module QuBits


include("../include.jl")

using JuMP
#using Ipopt

function formulation_3D(lambda,w,mu,mu1, A, B, C, alpha)
    T = [(i,j,k) for i in A for j in B for k in C];

    m = Model()
    @variable(m, c[i in A,j in B,k in C] >= alpha[k])# for (i,j) in T)
    @variable(m, y[i in A,j in B,k in C] >= alpha[k])# for (i,j) in T )
    @variable(m, obj_val)

    @NLobjective(m, Min, -sum( lambda[i,j,k] * (log(c[i,j,k]) - (y[i,j,k]/w[i])^mu1[j] / mu1[j]) for (i,j,k) = T))
    #@NLobjective(m, Min, obj_val)
    #@NLconstraint(m, -sum( lambda[i,j,k] * (log(c[i,j,k]) - (y[i,j,k]/w[i])^mu1[j] / mu1[j]) for (i,j,k) = T) <= obj_val)

    for (i,j,k) in T
      for (p,q,r) in T
        if i != p || j != q || k != r
          @NLconstraint(m,(log(c[i,j,k]) - (y[i,j,k]/w[i])^mu1[j] / mu1[j]) - (log(c[p,q,r]) - (y[p,q,r]/w[i])^mu1[j] / mu1[j]) >= 0.0)
        end
      end
    end

    @NLconstraint(m, sum(lambda[i,j,k] * (y[i,j,k] - c[i,j,k]) for (i,j,k) in T) >= 0.0);
    return m
end

na = 21
nb = 3
nc = 3

A = 1:na;
B = 1:nb;
C = 1:nc;

lambda = ones(na,nb,nc);
wmin = 2;
#wmax = wmin + na - 1;
wmax = 4
w = [wmin + ((wmax-wmin)/(na-1))*(i-1) for i = 1:na]
#mu = [1.0, 2.0, 3.0, 5.0, 8.0]
#mu = [1.0, 2.0, 4.0] #, 5.0, 8.0]

mu = [0.5, 1.0, 2.0]

gamma = 0.0

alpha = [0.0, 1.0, 1.5]

mu1 = mu + 1.0

srand(1)
m = formulation_3D(lambda,w,mu,mu1,A,B,C, alpha)
nlp_raw = MathProgNLPModel(m);
my_pars = Class_parameters();
my_pars.init.mu_scale = 1.0
iter = one_phase_solve(nlp_raw, my_pars);

#=
ORG_STDOUT = STDOUT
file = open("test.txt", "w")
redirect_stdout(file)

srand(1)
m = create_tax_problem(10)
one_phase_solve(m)

STDOUT = ORG_STDOUT

close(file)=#


# FunctionalForm.jl - explicit form of pair-potentials

# Goodwin--Skinner--Pettifor
macro GSP(A,n, nc,r0,rc)
    return :(r -> $A*($r0/r)^$n * exp( $n*(-(r/$rc)^$nc + ($r0/$rc)^$nc)) )
end

# Exponential x power law
#   Written as a crazy summation expansion for the joy of it.
macro EPL(r0,An,mn,pn)
    sum=:(0.0)
    for (A,m,p) in zip(An.args,mn.args,pn.args)
        sum=:($sum + $A * ($r0/r)^$m * exp(-$p *(r-$r0)) )
    end
    println("sum: ",sum)
    return :(r -> $sum)
end

# Chadi's quadratic
macro QUAD(r0,A1,A2)
    return :(r -> $A1*(r-$r0) + $A2*(r-$r0)*(r-$r0) )
end


"""
    chebyexp(A)

Compute the exponential of `A`. The exponential is calculated with a Chebyshev series expansion. For \$|A| ≤ 1\$, we have that

```math
\\exp(A) = J₀(i) + 2 ∑_{k=1}^∞ iᵏ Jₖ(-i) Tₖ(A)
```

where \$Jₖ\$ is the Bessel function of the first kind, and \$Tₖ\$ is a Chebyshev polynomial. If \$|A| > 1\$, scaling and squaring is used.
"""
chebyexp(A::AbstractMatrix{T}) where T<:Number = chebyexp!(Matrix{float(T)}(size(A)),A)

function chebyexp!(R::AbstractMatrix{<:Number},A::AbstractMatrix{<:Number})
    # Run alg
    normA = norm(A,1)
    if normA <= 1
        _chebyexp!(R,A)
    else
        p = ceil(Int,log2(normA))
        _chebyexp!(R,A/2^p)
        Q = similar(R)
        for i = 1:p
            copy!(Q,R)
            A_mul_B!(R,Q,Q)
        end
    end
    return R
end

function _chebyexp!(R::AbstractMatrix{T},A::AbstractMatrix{<:Number}) where T
    # only valid for |A|<1
    n = checksquare(A)
    if T <: BigFloat
        order, coeff0, coeff = 32, cheby_coeff0_big, cheby_coeff_big
    else
        order, coeff0, coeff = 16, cheby_coeff0, cheby_coeff
    end
    # Initialize required matrices
    fill!(R,zero(T))
    Tₖ₋₂ = similar(R); copy!(Tₖ₋₂,A)
    Tₖ₋₁ = eye(T,n,n); A_mul_B!(T(2),A,Tₖ₋₂,T(-1),Tₖ₋₁)
    # Compute expansion
    for i=1:n; R[i,i] = coeff0; end # R₀ = J₀(i)*I, the 0th order
    R .+= coeff[1].*Tₖ₋₂ .+ coeff[2].*Tₖ₋₁ # 1st and 2nd order
    for k in 3:order # higher orders
        A_mul_B!(T(2),A,Tₖ₋₁,T(-1),Tₖ₋₂) # result Tₖ is stored in Tₖ₋₂
        R .+= coeff[k].*Tₖ₋₂ # coeff[k].*Tₖ
        Tₖ₋₂, Tₖ₋₁ = Tₖ₋₁, Tₖ₋₂ # Tₖ₋₁->Tₖ₋₂ and Tₖ->Tₖ₋₁ for next order
    end
    return R
end

# Bessel function values precomputed in Mathematica

# J₀(i)
const cheby_coeff0_big = big"1.2660658777520083355982446252147175376076703113550"
const cheby_coeff0 = Float64(cheby_coeff0_big)

# 2 * iᵏ * Jₖ(-i)
const cheby_coeff_big = BigFloat[
    big"1.1303182079849700544153920552197266146577992432422"
    big"0.27149533953407656236570513998998184589974213622556"
    big"0.044336849848663804952571495259799231058830698339959"
    big"0.0054742404420937326502761684311864595467579461858053"
    big"0.00054292631191394375036214781030755468476712885351629"
    big"0.000044977322954295146654690328110912699086657650642406"
    big"3.1984364624019905058638729766022957272370458074163E-6"
    big"1.9921248066727957259610643848055890533900933857774E-7"
    big"1.1036771725517344326169960913353241812896390172496E-8"
    big"5.5058960796737472504714204020055270687431547280674E-10"
    big"2.4979566169849825227120109342187675410080716361415E-11"
    big"1.0391522306785700504996346724238478525397128556091E-12"
    big"3.9912633564144015128877204015326949127607826796533E-14"
    big"1.4237580108256571488273680253471752219093588992517E-15"
    big"4.7409261025614961710899305606042914145777617483995E-17"
    big"1.4801800572082975003888571658877975360303747318760E-18"
    big"4.3499194949441698455876297633392992805626063964777E-20"
    big"1.2074289272797528890630463524357806390885570735305E-21"
    big"3.1753567370594449606628945704889798438009317678357E-23"
    big"7.9336719716380401114641564996829844420300175294498E-25"
    big"1.8879484042289160772319706157860669889247560558152E-26"
    big"4.2886738765925870898799133815030885460420950260008E-28"
    big"9.3189852817775768480872792470802866623424437483364E-30"
    big"1.9406469749017397597649278461566813645709017660684E-31"
    big"3.8798022492260012156255855282161124021152712081106E-33"
    big"7.4585028873915195213508204862516351326616201311347E-35"
    big"1.3807477824110645231588753652621331312287399205215E-36"
    big"2.4648623717710962928935138361162240264245603184147E-38"
    big"4.2485421925059913850761703704767643098614220928398E-40"
    big"7.0790011762128954933502123970072670493550456762707E-42"
    big"1.1414867782540890660429322724040802483946870774259E-43"
    big"1.7831510375432838840323081019695093079857962301911E-45"
]
const cheby_coeff = Vector{Float64}(cheby_coeff_big)


using RNLA, Test, LinearAlgebra

@testset "rpca_test1" begin
    m = 10
    k = 10
    X = rand(m, k)
    U, S, W1 = svd(X)
    Z1 = U * Diagonal(S)
    Λ1 = (S .^ 2) ./ (m - 1)

    W2, Λ2, Z2 = rpca(X, k)

    @test norm(abs.(W1) - abs.(W2)) < sqrt(eps())
    @test norm(abs.(Λ1) - abs.(Λ2)) < sqrt(eps())
    @test norm(abs.(Z1) - abs.(Z2)) < sqrt(eps())
end

@testset "rpca_test2" begin
    m = 100
    n = 1000
    k = 50
    X = rand(m, n)
    U, S, W1 = svd(X)
    W1 = W1[:, 1:k]
    Z1 = U[:, 1:k] * Diagonal(S[1:k])
    Λ1 = (S[1:k] .^ 2) ./ (m - 1)

    W2, Λ2, Z2 = rpca(X, k)

    @test norm(abs.(W1) - abs.(W2)) < sqrt(eps())
    @test norm(abs.(Λ1) - abs.(Λ2)) < sqrt(eps())
    @test norm(abs.(Z1) - abs.(Z2)) < sqrt(eps())
end

@testset "rpca_test3" begin
    m = 1000
    n = 100
    k = 50
    X = rand(m, n)
    U, S, W1 = svd(X)
    W1 = W1[:, 1:k]
    Z1 = U[:, 1:k] * Diagonal(S[1:k])
    Λ1 = (S[1:k] .^ 2) ./ (m - 1)

    W2, Λ2, Z2 = rpca(X, k)

    @test norm(abs.(W1) - abs.(W2)) < sqrt(eps())
    @test norm(abs.(Λ1) - abs.(Λ2)) < sqrt(eps())
    @test norm(abs.(Z1) - abs.(Z2)) < sqrt(eps())
end


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 5163d91e-7e01-11eb-3d85-1bde1b52bcc2
using Distributions

# ╔═╡ 4cac0536-7e01-11eb-0f19-29257e61fef3
using StatsPlots

# ╔═╡ 9bcb7778-7e01-11eb-38f5-0d934dfa6201
x = [0:0.01:1;]

# ╔═╡ 5605125a-7e01-11eb-13a6-596920930e6a
begin
	α = @bind a html"<input type='range' min='0.1' max='10' step='0.1' value='5'>"
	β = @bind b html"<input type='range' min='0.1' max='10' step='0.1' value='5'>"
	md"""**distribution parameters:**
	$$\alpha$$: $α
	$$\beta$$: $β
	"""
end

# ╔═╡ 0132d818-7e02-11eb-1a70-bdb0bbcefa7a
md"""
 $$\alpha$$ = $a

 $$\beta$$ = $b
"""

# ╔═╡ 8cb866a0-7e02-11eb-1fdb-eff565b95d5f
begin
	apb = a + b
	m = a / (apb)
	v = a * b / apb * apb * (apb + 1)
	md"""
	Mean: $$\frac{\alpha}{\alpha + \beta}$$ = $m

	Variance: $$\frac{\alpha\beta}{(\alpha + \beta)^2 * (\alpha + \beta +1)}$$ = $v
	"""
end

# ╔═╡ 8fa79d8c-7e01-11eb-19a7-d33fcde6ff4c
B = Beta(a, b)

# ╔═╡ 995a9474-7e01-11eb-348b-cb8a8233b489
plot(x, pdf.(B, x),
	label = "Beta($a, $b)", linewidth = 2)

# ╔═╡ Cell order:
# ╟─9bcb7778-7e01-11eb-38f5-0d934dfa6201
# ╟─5605125a-7e01-11eb-13a6-596920930e6a
# ╟─0132d818-7e02-11eb-1a70-bdb0bbcefa7a
# ╟─8cb866a0-7e02-11eb-1fdb-eff565b95d5f
# ╠═8fa79d8c-7e01-11eb-19a7-d33fcde6ff4c
# ╟─995a9474-7e01-11eb-348b-cb8a8233b489
# ╟─5163d91e-7e01-11eb-3d85-1bde1b52bcc2
# ╟─4cac0536-7e01-11eb-0f19-29257e61fef3


function gauss(x, w)
    1/sqrt(2.0*pi)./w.*exp(-x.^2/2.0./w.^2)
end


"""
$(TYPEDEF)

One sample superiority test for proportion

Constructors
------------
* `OneSamplePropSuperior(p1::Real, p0::Real, delta::Real)`

Arguments
---------
* `p1`: True proportion

* `p0`: Proportion under null hypothesis

* `delta`: Superiority margin

Fields
------
$(FIELDS)
"""
type OneSamplePropSuperior <: TrialTest
    p1::Real
    p0::Real
    delta::Real

    # Validator
    function OneSamplePropSuperior(p1, p0, delta)

        if !((0 < p1 < 1) & (0 < p0 < 1))

            error("Proportion value must be in (0, 1)")

        end # end if

        if !(0 <= delta < 1)

            error("The superiority margin δ must be in [0, 1)")

        end # end if
        new(p1, p0, delta)

    end # end function

end # type

# One sample superiority test for proportion
function hypotheses{T <: OneSamplePropSuperior}(test::T, n::Real, std::Void, alpha::Real, side::String)

    diff = test.p1 - test.p0 - test.delta
    se = sqrt(1 / n * test.p1 * (1 - test.p1))
    z = diff / se
    p = cdf(ZDIST, z - quantile(ZDIST, 1 - alpha)) + cdf(ZDIST, -z - quantile(ZDIST, 1 - alpha))
    return p

end # end function


function grabzonalstates(file; ghost = 3, startind = 1)
    # ghost should be 3 WENO
    zonalstatistics = jldopen(file)
    tkeys = keys(zonalstatistics["timeseries"]["t"])[1:1]
    t = [zonalstatistics["timeseries"]["t"][tkey] for tkey in tkeys]
    yC = zonalstatistics["grid"]["yC"][ghost+1:end-ghost]
    zC = zonalstatistics["grid"]["zC"][ghost+1:end-ghost]
    yF = zonalstatistics["grid"]["yF"][ghost+1:end-ghost]
    zF = zonalstatistics["grid"]["zF"][ghost+1:end-ghost]
    fields = [:u, :v, :w, :b, :vb, :vv, :wb, :ww]
    for field in fields
        label = string(field)
        @eval $field = mean([zonalstatistics["timeseries"][$label][tkeys[i]][1,:,:] for i in $startind:length(tkeys)])
    end
    close(zonalstatistics)

    vpbp =  average1(vb) - average1(v) .* b
    uz = Δ2(u) ./ reshape((zC[2:end] - zC[1:end-1]), (1, length(zC)-1))
    bz = Δ2(b) ./ reshape((zC[2:end] - zC[1:end-1]), (1, length(zC)-1))
    by = Δ1(b) ./ reshape((yC[2:end] - yC[1:end-1]), (length(yC)-1, 1))
    coriolis = -1e-4 .+ 1e-11 * reshape(yC, (length(yC), 1))
    κ =  -average1(vpbp) ./ by
    ν = coriolis .* average2(vpbp) ./ bz ./ uz

    #bz = sum(bz, dims = 1) ./ size(bz)[1]
    zCA = (zC[2:end] + zC[1:end-1])/2

    states = [eval(field) for field in fields]
    statenames = [string(field) for field in fields]
    units = ["[m/s]", "[m/s]", "[m/s]", "[m/s²]", "[m²/s³]", "[m²/s²]", "[m²/s³]", "[m²/s²]"]
    # bz
    push!(states, bz)
    push!(statenames, "∂ᶻb")
    push!(units, "[1/s²]")

    # by
    push!(states, by)
    push!(statenames, "∂ʸb")
    push!(units, "[1/s²]")

    # v'b'
    push!(states, vpbp)
    push!(statenames, "v'b'")
    push!(units, "[m²/s³]")
    # κ
    push!(states, κ)
    push!(statenames, "κ ≈  v'b' / ∂ʸb")
    push!(units, "[m²/s]")
    # ν
    push!(states, ν)
    push!(statenames, "ν ≈  f v'b' / ∂ᶻb / ∂ᶻu")
    push!(units, "[m²/s]")
    domain = [yF, zF]
    return states, statenames, units, domain
end


## Plot recipe
## brute force this to avoid vertical asymptotes

@recipe function f(pq::AbstractRationalFunction{T}, a=nothing, b=nothing) where {T <: Real}

    p,q = lowest_terms(convert(RationalFunction,pq), method=:numerical)

    if a==nothing && b==nothing
        u=Polynomials.poly_interval(p)
        v=Polynomials.poly_interval(q)
        a,b = min(first(u), first(v)), max(last(u), last(v))
    end

    n = 251
    Δ = (b-a)/n

    # identify vertical asumptotes
    zs = filter(z -> a+Δ < z < b-Δ, sort(real.(filter(isreal, roots(q)))))
    xs = Any[]
    ys = Any[]
    if length(zs) > 0
        z₁ = popfirst!(zs)
        xs′ = range(a, stop=z₁ - Δ, length=50)
        ys′ = pq.(xs′)
        push!(xs, xs′)
        push!(ys, ys′)
        push!(xs, z₁)
        push!(ys, NaN)
        for z₂ ∈ zs
            xs′ = range(z₁ + Δ, stop=z₂ - Δ, length=50)
            ys′ = pq.(xs′)
            push!(xs, xs′)
            push!(ys, ys′)
            push!(xs, z₂)
            push!(ys, NaN)
            z₁ = z₂
        end
        xs′ = range(z₁+Δ, stop=b, length=50)
        ys′ = pq.(xs′)
        push!(xs, xs′)
        push!(ys, ys′)

        xs = vcat(xs...)
        ys = vcat(ys...)
    else
        xs = range(a, stop=b, length=251)
        ys = pq.(xs)
    end

    xs, ys
end


using CompScienceMeshes
using Test

function isoriented(m::Mesh)

    @assert dimension(m) >= 0
    edges = skeleton(m, dimension(m)-1)

    D = connectivity(edges, m)
    S = (sum(D,1) .== 0)
    return all(S)
end

fn = joinpath(dirname(@__FILE__),"assets","cube2.in")
m = readmesh(fn)
@test isoriented(m)

fn = joinpath(dirname(@__FILE__),"assets","sphere2.in")
m = readmesh(fn)
@test isoriented(m)


using Distributed
@everywhere using Mamba, LinearAlgebra

## Data
birats = Dict{Symbol, Any}(
  :N => 30, :T => 5,
  :x => [8.0, 15.0, 22.0, 29.0, 36.0],
  :Y => [151 199 246 283 320
         145 199 249 293 354
         147 214 263 312 328
         155 200 237 272 297
         135 188 230 280 323
         159 210 252 298 331
         141 189 231 275 305
         159 201 248 297 338
         177 236 285 350 376
         134 182 220 260 296
         160 208 261 313 352
         143 188 220 273 314
         154 200 244 289 325
         171 221 270 326 358
         163 216 242 281 312
         160 207 248 288 324
         142 187 234 280 316
         156 203 243 283 317
         157 212 259 307 336
         152 203 246 286 321
         154 205 253 298 334
         139 190 225 267 302
         146 191 229 272 302
         157 211 250 285 323
         132 185 237 286 331
         160 207 257 303 345
         169 216 261 295 333
         157 205 248 289 316
         137 180 219 258 291
         153 200 244 286 324],
  :mean => [0.0, 0.0],
  :var => [1.0e6 0.0
           0.0   1.0e6],
  :Omega => [200.0 0.0
               0.0 0.2]
)

## Model Specification
model = Model(

  Y = Stochastic(2,
    (beta, x, sigmaC, N, T) ->
      UnivariateDistribution[
        Normal(beta[i, 1] + beta[i, 2] * x[j], sigmaC)
        for i in 1:N, j in 1:T
      ],
    false
  ),

  beta = Stochastic(2,
    (mu_beta, Sigma, N) ->
      MultivariateDistribution[
        MvNormal(mu_beta, Sigma)
        for i in 1:N
      ],
    false
  ),

  mu_beta = Stochastic(1,
    (mean, var) -> MvNormal(mean, var)
  ),

  Sigma = Stochastic(2,
    Omega -> InverseWishart(2, Omega),
    false
  ),

  sigma2C = Stochastic(
    () -> InverseGamma(0.001, 0.001),
    false
  ),

  sigmaC = Logical(
    sigma2C -> sqrt(sigma2C)
  )

)

## Initial Values
inits = [
  Dict(:Y => birats[:Y], :beta => repeat([100 6], birats[:N], 1),
       :mu_beta => [0, 0], :Sigma => Matrix{Float64}(I, 2, 2), :sigma2C => 1.0),
  Dict(:Y => birats[:Y], :beta => repeat([50 3], birats[:N], 1),
       :mu_beta => [10, 10], :Sigma => 0.3 * Matrix{Float64}(I, 2, 2), :sigma2C => 10.0)
]

## Sampling Scheme
scheme = [AMWG([:beta, :mu_beta], repeat([10.0, 1.0], birats[:N] + 1)),
          AMWG(:Sigma, 1.0),
          Slice(:sigma2C, 10.0)]
setsamplers!(model, scheme)

## MCMC Simulations
sim = mcmc(model, birats, inits, 10000, burnin=2500, thin=2, chains=2)
describe(sim)


@testset "Magnetostatics Functional Tests" begin
    @testset "Repelling magnetic dipoles" begin
        d1 = 0.01 # m
        m1 = 5e-6 # kg
        m2 = 5e-6 # kg
        ρ = 7800 # kg/m^3
        M = 1.2e6 # A/m
        mm1 = SVector(0.0, 0.0, M * m1 / ρ)
        mm2 = SVector(0.0, 0.0, M * m2 / ρ)
        μ_4π = 1e-7

        t1 = 0.0  # s
        t2 = 1.0 # s
        τ = (t2 - t1) / 100

        p1 = MagneticParticle(SVector(-d1 / 2, 0.0, 0.0), SVector(0.0, 0.0, 0.0), m1, mm1)
        p2 = MagneticParticle(SVector(d1 / 2, 0.0, 0.0), SVector(0.0, 0.0, 0.0), m2, mm2)
        parameters = MagnetostaticParameters(μ_4π)
        system = PotentialNBodySystem([p1, p2], Dict(:magnetic => parameters))
        simulation = NBodySimulation(system, (t1, t2))
        sim_result = run_simulation(simulation, VelocityVerlet(), dt=τ)

        d2 = norm(get_position(sim_result, t2, 2) - get_position(sim_result, t2, 1))
        v_expected = sqrt(μ_4π * ( dot(mm1, mm2) * (1 / d1^3 - 1 / d2^3) ) / m1)
        v_actual = norm(get_velocity(sim_result, t2, 2))

        ε = 0.001
        @test v_expected ≈ v_actual atol = ε
    end

    @testset "Constructing magnetostatic potential parameters entity" begin
        default_potential = MagnetostaticParameters()
        @test 1e-7 == default_potential.μ_4π

        io = IOBuffer()
        potential1 = MagnetostaticParameters()
        potential2 = MagnetostaticParameters(4π)
        @test sprint(io -> show(io, potential1)) == "Magnetostatic:\n\tμ/4π:1.0e-7\n"
        @test sprint(io -> show(io, potential2)) == "Magnetostatic:\n\tμ/4π:12.566370614359172\n"
    end
end


# Generate 3 Gaussian peak at position : 5, 10, 20 with difference
# shape factors σ 1, 1.5, 2
#
using DelimitedFiles,Random
using NLS_Fit

Random.seed!(1234)

model = Gaussian_Peak() + Gaussian_Peak() + Gaussian_Peak()
θ1 = Float64[1,5,1]
θ2 = Float64[1,10,1.5]
θ3 = Float64[1,20,2]
θ = vcat(θ1,θ2,θ3)

X=Float64[1:0.25:30;]
n=length(X)
Y=eval_y(model,X,θ) + 0.1*(rand(n) .- 0.5)

rootDir = joinpath(dirname(pathof(NLS_Fit)), "..")
dataDir = joinpath(rootDir,"data")
dataFile = joinpath(dataDir,"mapped_parameters.txt")
writedlm(dataFile,hcat(X,Y))


using AccuracyAtTopKernels

using Test, Distributions, LinearAlgebra, Random
import AccuracyAtTopKernels.MLKernels

include("./tests_surrogates.jl")
include("./tests_datasets.jl")
include("./tests_projections.jl")
include("./tests_primal_problems.jl")
include("./tests_dual_problems.jl")
include("./tests_kernels.jl")
include("./tests_scores.jl")

@time @testset "all tests" begin
    @testset "tests surrogates" begin
        test_hinge()
        test_quadratic()
    end

    @testset "tests datasets" begin
        test_primal()
        test_dual()
    end

    @testset "tests projections" begin
        test_projections()
    end

    @testset "tests primal problems" begin
        test_primal_problems()
    end

    @testset "tests dual problems" begin
        test_dual_problems()
    end

    @testset "tests kernels" begin
        test_kernels()
    end

    @testset "tests scores" begin
        test_scores()
    end
end;


# Developed by Jason Ma
# This is a Julia implementation of the strangle model given to us
# This includes the updated version of the finutil_stu file and the
# strangle_stu.py file. My Julia environment had issues importing the
# functions from finutil_stu to strangle_stu, so I merged the files here
# The split between finutil and strangle_stu occurs around line 265.
# Note, I have also included what the files should look like when
# importing works correctly. The only change from this merged file
# to what it should look like is by adding a "fu." prefix to any functions
# from finutil in the strangle_stu.jl file

# This is finutil_stu.jl
# This is a frozen and reduced version of finutil.py. This file is not likely
# to be current with funutil and should not be used in place of finutil (see log).
# This should only be used with models with the _stu appendage.
# These are the utilities that are used throughout the various financial models
# for such things as calculating options prices and volatilities. These are
# designed to be used with Version 3.6 of Python.
# This version 1.2stu, April 11, 2018  NOTE: This version has not been completely
# tested.
# Maintained by Professor Evans
#
#
# daysto calculates the number of days between now and some event, such as days2expiry
# as an integer float.
#
using SpecialFunctions

function daysto(eyear,emonth,eday)
    tnow = Dates.today()
    expiry = Dates.Date(eyear, emonth, eday)
    days2expiry = abs(expiry - tnow)
    return float(days2expiry.days)
end
#
function monthname(mo)
    monthlabel = ["JAN", "FEB", "MAR", "APR", "MAY", "JUN", "JUL", "AUG",
                 "SEP", "OCT", "NOV", "DEC" ]
    return monthlabel[mo]
end
#
# csnd integrates a standard normal distribution up to some sigma.
#
function csnd(point)
    return (1.0 + erf(point/sqrt(2.0)))/2.0
end
#
# cnd integrates a Gaussian distribution up to some value.
#
function cnd(center,point,stdev)
    return (1.0 + erf((point - center)/(stdev*sqrt(2.0))))/2.0
end
#
# Calculating the BSM call option price, traditional model. This requires the
# user to provide stock price, strike price, daily volatility, risk-free interest
# rate and days to expiry. (To calculate days use method daysto above).
# This returns the call price, the delta, and duration volatility as a tuple
# array. See popo below for puts.
#
function copo(stock,strike,dayvol,days,rfir)
    d1 = log(stock/strike)+((rfir/365)+(dayvol^2)/2)*days
    durvol = dayvol*sqrt(days)
    delta = csnd(d1/durvol)
    cumd2 = csnd((d1/durvol) - durvol)
    discount = exp(-rfir*days/365)
    callpr = (stock*delta)-(strike*discount*cumd2)
    return [callpr,delta,durvol]
end
#
# An elementary function for calculating the stock price adjusted for drift.
#
function drift(alpha,time)
    return 1.0*exp(alpha*time)
end
#
# An elementary multiplier function for converting daily volatility to duration
# volatility.
#
function durvol(time)
    return 1.0*sqrt(time)
end
#
# An elementary price expected-mean-value adjustment multiplier for log
# distributed prices. The mean of a log-distributed pdf is adjusted by minus
# one-half variance.
#
function lnmeanshift(sigma)
    return 1.0*exp(-1.0*(sigma*sigma/2))
end
#
# A time-discount function to discount the value of a future payment (like an
# option) discounted at the risk-free interest rate. The variable riskfreerate
# is annual and time is in days.
#
function dcount(riskfreerate,time)
    return 1.0*exp(-1.0*(riskfreerate/365)*time)
end
#
# This is an option tranche value calculator function that assumes you have a
# stock price and strike price, adjusted externally (for example, drift is
# adjusted with drift above). This will calculate the strike-price adjusted
# tranche from either -5 sigma to the strike or from the strike to +5 sigma.
# Sigma used here is duration sigma, adjusted outside using durvol above.
# Main program must set call to true if a call, false if a put.
#
function otranche(stock,strike,dursigma,call)
    sspread = (log(strike/stock))/dursigma
    if call
        binborder = linspace(sspread, 5.00, 24)
    else
        binborder = linspace(-5.0, sspread, 24)
    end
    num_bins = 24
    binedgeprob = zeros(num_bins)
#
    for i = 1:num_bins
        binedgeprob[i] = csnd(binborder[i])
    end
    num_bins = num_bins - 1
    binprob = zeros(num_bins)
    binmidprice = zeros(num_bins)
    binvalue = zeros(num_bins)
#
    for i = 1:num_bins
        binprob[i] = binedgeprob[i+1] - binedgeprob[i]
        binmidprice[i] = ((stock*exp(((binborder[i+1]+binborder[i])/2.0)*dursigma))*lnmeanshift(dursigma)) - strike
        binvalue[i] = binmidprice[i]*binprob[i]
    end
    if call
        optionprice = sum(binvalue)
    else
        optionprice = (sum(binvalue))*-1.0
    end
#
    return optionprice
end
#
# This is full tranche value calculator function that assumes you have a stock
# price and reference price, usually a strike price (and still called that
# here). This will calculate the tranche value from either -5 sigma to the
# reference price (left is true) or from the reference to +5 sigma (left
# is false). This is similar to otranche except that it does not subtract the
# strike price and therefore gives the full value of the tranche. It is
# designed to be used primarily by the Aruba Model to calculate the value of
# the remaining stock tranche if the covered call you wrote has been exercised.
#
function ftranche(stock,strike,sigma,left)
    sspread = (log(strike/stock))/sigma
    if left
        binborder = linspace(-5.0, sspread, n=24)
    else
        binborder = linspace(sspread, 5.0, n=24)
    end
    num_bins = 24
    binedgeprob = zeros(24)
#
    for i = 1:num_bins
        binedgeprob[i] = csnd(binborder[i])
    end
    num_bins = num_bins - 1
    binprob = zeros(num_bins)
    binmidprice = np.zeros(num_bins)
    binvalue = np.zeros(num_bins)
#
    for i = 1:num_bins
        binprob[i] = binedgeprob[i+1] - binedgeprob[i]
        binmidprice[i] = ((stock*exp(((binborder[i+1]+binborder[i])/2.0)
        *sigma))*lnmeanshift(sigma))
        binvalue[i] = binmidprice[i]*binprob[i]
    end
#
    trancheprice = sum(binvalue)
    return trancheprice
end
#
# oidv calculates implied daily and duration volatility for a call or a put
# using divide and conquer (the default for most models). Also see oidvnm.
# This uses an iterative process that uses otranche (above) to calculate the
# sigma, here an implied sigma, from the existing option value (ovalue).
# The call variable is True for a call, False for a put. The convergence is
# within the while loop. This function returns a tuple of two values, daily
# IDV and duration IDV.
#
function oidv(stock,strike,ovalue,days,call)
    epsilon = 1e-4
    low = 0.0
    high = 1.0
    daysigma = (high+low)/2
    dursigma = daysigma*durvol(days)
    tempop = otranche(stock,strike,dursigma,call)
    while tempop<=(ovalue-epsilon) || tempop>=(ovalue+epsilon)
        if tempop >= (ovalue+epsilon)
            high = daysigma
        else
            low = daysigma
        end
        daysigma = (high+low)/2
        dursigma = daysigma*durvol(days)
        tempop = otranche(stock,strike,dursigma,call)
    end
    # End of Loop!
    return [daysigma,dursigma]
end
#
# oidvnm calculates implied daily and duration volatility for a call or a put
# using Newton's Method for convergence (default is divide and conquer).
# This uses an iterative process that uses otranche (above) to calculate the
# sigma, here an implied sigma, from the existing option value (ovalue).
# The call variable is True for a call, False for a put. The convergence is
# within the while loop. This function returns a tuple of two values, daily
# IDV and duration IDV.
#
function oidvnm(stock,strike,ovalue,days,call)
    seedsigma = 1e-6
    durseed = seedsigma*durvol(days)
    # NOTE: daysigma is supposed to be set at a reasonable estimate of the
    # actual idv. The Newton method can explode if it is not (especially if you
    # enter option price, strike, and value data that are very unrealistic).
    # Consider setting daysigma at sqrt*time*ln(strike/stock). It was
    # originally set at 0.05
    daysigma = (log(strike/stock))*sqrt(days)
    dursigma = daysigma*durvol(days)
    cutoff = 1e-4
    tempop = 0.00
    #
    # The loop starts here. You start with a test sigma and converge to the
    # actual sigma. The convergence shown here (Newton's method) was designed
    # by Alec Griffith '17
    #
    while abs(tempop - ovalue) > cutoff
        tempop = otranche(stock,strike,dursigma,call)
        price2 = otranche(stock,strike,dursigma+durseed,call)
        deriv = (price2-tempop)/seedsigma
        daysigma -= (tempop-ovalue)/deriv
        dursigma = daysigma*durvol(days)
    end
    # End of Loop!
    return [daysigma,dursigma]
end
#
# Calculating the BSM put option price, traditional model. This requires the
# user to provide stock price, strike price, daily volatility, risk-free interest
# rate and days to expiry. (To calculate days use method daysto above).
# This returns the put price, the delta, and duration volatility as a tuple
# array. See copo below for calls.
#
function popo(stock,strike,dayvol,days,rfir)
    d1 = log(stock/strike)+((rfir/365)+(dayvol^2)/2)*days
    durvol = dayvol*sqrt(days)
    delta = csnd(-d1/durvol)
    cumd2 = csnd(-(d1/durvol - durvol))
    discount = exp(-rfir*days/365)
    putpr = -(stock*delta)+(strike*discount*cumd2)
    return [putpr,delta,durvol]
end
#
# tdfecay adds an extension to the otranche calculator (Taboga model)
# to calculate one-day time decay.
#

function tdecay(stock, strike, daysigma, oprice, days, call)
    days = days - 1.0
    dursigma = daysigma*durvol(days)
    oprice1d = otranche(stock,strike,dursigma,call)
    timedecay = oprice - oprice1d
    return timedecay
end

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

# strangle_stu.py
#
# This is a MOS model: Mudd Open Source only.
# This is a deep strangle model based upon the Taboga concept (student version).
# Much of the documentation can be found in the core Taboga model or finutil.py
# Designed in Python 3.6 on for Economics 136
# Original version developed on May 30, 2017 by Prof Evans
# This is version 2.1stu, dated April 11, 2017
#

# This model draws upon utility model finutil_stu. That model must either be in
# the same directory as this model or the path command below must be used to
# point to its location.
# import sys
# sys.path.append('c:/Users/Prof Gary Evans/Dropbox/PyGo/PyFi')

# This version calculates IDV for the put and call, but also requires that you
# supply historical volatility (252 and 30 day). You can include an earnings
# date if date is relevant.
#
# Load data
aname = "Prof Evans"
stosym = "MSFT"
tnow = Dates.today()
# Set earnings to True if this is an earnings-based strangle.
earnings = true
# If you don't need to calculate the days to expiry and earnings (optional)
# then calcdays is False (default).
calcdays = false
# BEGIN of optional time data. Load below ONLY if calcdays = True
# Expiration date:
exyear = 2018
exmonth = 2
exday = 16
# Earnings date (load only if relevant):
eayear = 2018
eamonth = 1
eaday = 31
# END of optional time data load
#
# This model does not use the alpha.
# stock = 15.15
# callstrike = 16
# callbid = 0.90
# callask = 0.93
# putstrike = 15
# putbid = 1.150
# putask = 1.18
# sigma252 = 0.043534
# sigma30 = 0.025
###########################
stock = 166.86
callstrike = 170
callbid = 3.20
callask = 3.25
putstrike = 165
putbid = 4.50
putask = 4.60
sigma252 = 0.011358
sigma30 = 0.016387
if calcdays
    days = daysto(exyear,exmonth,exday)
    if earnings
        daystoearn = daysto(eayear,eamonth,eaday)
    end
else
    # Here you have to supply the days if you don't calculate them.
    days = 14.0
    if earnings
        daystoearn = 11.0
    end
end
#
moname = monthname(exmonth)
# Set effective options prices (PEG)
peg = 0.6
cspread = callask - callbid
callpr = callbid + (peg*cspread)
pspread = putask - putbid
putpr = putbid + (peg*pspread)
cost = callpr + putpr
# Note: callbe and putbe should be calced based upon the prices they would
# have to be the day after earnings, especially for more distant options.
callbe = stock + cost
putbe = stock - cost
callbegr = log(callbe/stock)*100
putbegr = log(putbe/stock)*100
# Consider adding log percent breakeven to current value
#
# Calculate call and put ideal prices based upon historical volatilities.
hv252 = sigma252*durvol(days)
cop252 = otranche(stock,callstrike,hv252,true)
pop252 = otranche(stock,putstrike,hv252,false)
hv30 = sigma30*durvol(days)
cop30 = otranche(stock,callstrike,hv30,true)
pop30 = otranche(stock,putstrike,hv30,false)
#
# Calculate call and put IDV based upon actual prices
# See the documentation for finutil and tabogacidv to see what these are doing.
# Start with the call:
#
seedsigma = 1e-6
cutoff = 1e-4
tempop = 0.00
price2 = 0.00
derive = 0.00
cdurseed = seedsigma*durvol(days)
cdaysigma = 0.06
cdursigma = cdaysigma*durvol(days)
#
# The convergence shown here (Newton's method) was designed by Alec Griffith '17
#
while abs(tempop - callpr) > cutoff
    tempop = otranche(stock,callstrike,cdursigma,true)
    price2 = otranche(stock,callstrike,cdursigma+cdurseed,true)
    deriv = (price2-tempop)/seedsigma
    cdaysigma -= (tempop-callpr)/deriv
    cdursigma = cdaysigma*durvol(days)
end
# End of Loop!
#
# Repeat for the put:
#
seedsigma = 1e-6
cutoff = 1e-4
tempop = 0.00
price2 = 0.00
derive = 0.00
pdurseed = seedsigma*durvol(days)
pdaysigma = 0.01
pdursigma = pdaysigma*durvol(days)
call = false
#
#
while abs(tempop - putpr) > cutoff
    tempop = otranche(stock,putstrike,pdursigma,false)
    price2 = otranche(stock,putstrike,pdursigma+pdurseed,false)
    deriv = (price2-tempop)/seedsigma
    pdaysigma -= (tempop-putpr)/deriv
    pdursigma = pdaysigma*durvol(days)
end
# End of Loop!
callsigmaratio = cdaysigma/sigma252
putsigmaratio = pdaysigma/sigma252
#
dateString = Dates.format(now(), "e u d y")
println()
println("Strangle Model")
println("Analyst's name: $(aname)")
println( "Date: $(dateString)")
println( "Stock symbol: $(stosym)")
println("Option chain: $(moname), $(exday)")
println("Days to expiry: $(days)")
println( "Stock price: $(stock)")
println( "Call strike price: $(callstrike)")
println( "Call Bid: $(callbid)")
println( "Call Ask: $(callask)")
println( "Call price: $(callpr)")
println( "Put strike price: $(putstrike)")
println( "Put Bid: $(putbid)")
println( "Put Ask: $(putask)")
println( "Put price: $(putpr)")
println( "Position cost: $(cost)")
if earnings
    println("Days to earnings: $(daystoearn)")
println( "Call breakeven price: $(callbe)")
println( "Call breakeven ln percent: $(callbegr)")
println( "Put breakeven price: $(putbe)")
println( "Put breakeven ln percent: $(putbegr)")
println( "Call value based upon 252-day HV ($(sigma252)) : $(cop252)")
println( "Call value based upon 30-day HV ($(sigma30)) : $(cop30)")
println( "Put value based upon 252-day HV ($(sigma252)) : $(pop252)")
println( "Put value based upon 30-day HV ($(sigma30)) : $(pop30)")
println( "Call implied duration volatility: $(cdursigma)")
println( "Put implied duration volatility: $(pdursigma)")
println("CRITICAL")
println( "Call implied daily volatility: $(cdaysigma)")
println( "Put implied daily volatility: $(pdaysigma)")
println( "Call sigma ratio: $(callsigmaratio)")
println( "Put sigma ratio: $(putsigmaratio)")
#
# Calculate one day time decay using our calculated IDV.
#
if days >= 2.0
    days = days - 1.0
    cdursigma = cdaysigma*durvol(days)
    cprice1d = otranche(stock,callstrike,cdursigma,true)
    pdursigma = pdaysigma*durvol(days)
    pprice1d = otranche(stock,putstrike,pdursigma,false)
    timedecay = (callpr - cprice1d) + (putpr - pprice1d)
    println( "One day time decay: $(timedecay)")
else
    println( "No time decay with one day remaining.")
end
end
#


using GeometryRandomFields
using Base.Test

gf = GaussianCovarianceFunction(10,1)
@test gf(0) == 1
@test_approx_eq gf(10) exp(-1/2)

mf = MatérnCovarianceFunction(1/2,10,1)
@test mf(0) == 1
@test_approx_eq mf(10) exp(-1)

m,n = 1024,512
M = grf(m,n,gf)
@test size(M) == (m,n)

GeometryRandomFields.acf(real(M))

Z1,Z2 = reim(M)

@test Γ(Z1,10) == 0


"""
    HarmonicAngle(i, j, k, th0, cth)

A bond angle between three atoms.
"""
struct HarmonicAngle{T} <: SpecificInteraction
    i::Int
    j::Int
    k::Int
    th0::T
    cth::T
end

# Sometimes domain error occurs for acos if the value is > 1.0 or < -1.0
acosbound(x::Real) = acos(clamp(x, -1, 1))

@inline @inbounds function force(a::HarmonicAngle,
                                  coords,
                                  s::Simulation)
    ba = vector(coords[a.j], coords[a.i], s.box_size)
    bc = vector(coords[a.j], coords[a.k], s.box_size)
    pa = normalize(ba × (ba × bc))
    pc = normalize(-bc × (ba × bc))
    angle_term = -a.cth * (acosbound(dot(ba, bc) / (norm(ba) * norm(bc))) - a.th0)
    fa = (angle_term / norm(ba)) * pa
    fc = (angle_term / norm(bc)) * pc
    fb = -fa - fc
    return [a.i, a.j, a.k], [fa, fb, fc]
end

@inline @inbounds function potential_energy(a::HarmonicAngle,
                                            s::Simulation)
    ba = vector(s.coords[a.j], s.coords[a.i], s.box_size)
    bc = vector(s.coords[a.j], s.coords[a.k], s.box_size)
    pa = normalize(ba × (ba × bc))
    pc = normalize(-bc × (ba × bc))

    a.cth / 2 * (acosbound(dot(ba, bc) / (norm(ba) * norm(bc))) - a.th0)^2
end


##################################################################
# READ TROUBLESOME FILES WITH TOO MUCH WHITESPACE IN PLACES
##################################################################
read_and_unspace() = strip(replace(readline(), r"\s+", " "))
read_ints() = [parse(Int, x) for x in split(read_and_unspace(), " ")]

##################################################################
# PRIMES AND DIVISORS
##################################################################

# prime sieve with prime factors dictionary
function primesAndFactors(n::Int)
    sieve = [true for i in 1:n]
    sieve[1] = false
    prime_factors = [Dict{Int,Int}() for i in 1:n]
    for i in 2:n
        if sieve[i]
            prime_factors[i][i] = 1
            for j in (2i):i:n
                sieve[j] = false
                mul = 1
                k = j/i
                while k%i == 0
                    mul += 1
                    k/=i
                end
                prime_factors[j][i] = mul
            end
        end
    end
    primes = [i for i in 1:n if sieve[i]]
    return (sieve, primes, prime_factors)
end

# get the divisors of a number given its prime factorisation
function divisors_from_factors(factors::Dict{Int,Int})
    if length(factors)==0
        return Array{Int,1}()
    end
    ndivisors = prod(1+collect(values(factors)))

    bases = collect(keys(factors))
    exponents = [factors[x] for x in bases]
    divisors = zeros(Int, ndivisors)
    offset = 1
    N = length(bases)

    function enumerate_divs(cur_idx, cur_val)
        if cur_idx > N
            divisors[offset] = cur_val
            offset+=1
        else
            mul = cur_val
            for i in 0:exponents[cur_idx]
                enumerate_divs(cur_idx+1, mul)
                mul *= bases[cur_idx]
            end
        end
    end
    enumerate_divs(1, 1)
    return divisors
end

##################################################################
# Least unused natural number structure
##################################################################
type least_unused
    start_pos::Int
    cur_pos::Int
    unused::Base.Collections.PriorityQueue{Int, Int, Base.Order.ForwardOrdering}
    least_unused(sp::Int=0) = new(sp, sp, Base.Collections.PriorityQueue{Int, Int, Base.Order.ForwardOrdering}())
end

function unuse!(l::least_unused, val::Int)
    l.unused[val] = val
end

function use!(l::least_unused)
    if(isempty(l.unused))
        l.cur_pos += 1
        return l.cur_pos
    else
        return Base.Collections.dequeue!(l.unused)
    end
end

function use!(l::least_unused, n::Int)
    if n > l.cur_pos
        for k in (l.cur_pos+1):(n-1)
            unuse!(l, k)
        end
        l.cur_pos = n
    else
        l.unused[n] = -1
        Base.Collections.dequeue!(l.unused)
    end
end

function top(l::least_unused)
    if(isempty(l.unused))
        return l.cur_pos + 1
    else
        return Base.Collections.peek(l.unused)[1]
    end
end

##################################################################
# Interval tree to aggregate any function over an interval
##################################################################

type itreeNode{T<:Number}
    value::T
    pos::Int # position at which the value is obtained
    left::Int
    right::Int
    children::Vector{itreeNode{T}}
    itreeNode() = new(0, 0, 0, 0, Vector{itreeNode{T}}())
    itreeNode(v::T, p::Int, l::Int, r::Int) = new(v,p,l,r,Vector{itreeNode{T}}())
    itreeNode(v::T, p::Int, l::Int, r::Int, lc::itreeNode{T}, rc::itreeNode{T}) = new(v,p, l, r, [lc,rc])
    itreeNode(v::T, p::Int, l::Int, r::Int, ch::Vector{itreeNode{T}}) = new(v,p,l,r,ch)
end
itreeNode() = itreeNode{Int}()

function create_itree_closure{T}(v::Vector{T}, agg::Function, isleft::Function)
    function create_interval_tree(l::Int, r::Int)
        if l == r
            return itreeNode{T}(v[l], l, l, r)
        else
            mid_l = (l+r)>>1
            mid_r = mid_l + 1
            lhs = create_interval_tree(l, mid_l)
            rhs = create_interval_tree(mid_r, r)
            val = agg(lhs.value, rhs.value)
            left_side = isleft(lhs.value, rhs.value)
            pos = 0
            if left_side
                pos = lhs.pos
            else
                pos = rhs.pos
            end
            return itreeNode{T}(val, pos, l, r, lhs, rhs)
        end
    end
    return create_interval_tree
end

type itree{T}
    fn::Function
    isleft::Function
    x::Vector{T}
    root::itreeNode{T}
    default::T
    itree(v::Vector{T}, fn::Function, def::T) = new(fn, (x,y)->true, v, create_itree_closure(v,fn, (x,y)->true)(1, length(v)), def)
    itree(v::Vector{T}, fn::Function, leftfn::Function, def::T) = new(fn, leftfn, v, create_itree_closure(v,fn, leftfn)(1, length(v)), def)
end
itree(v::Vector{Int}, fn::Function, leftfn::Function, def::Int) = itree{Int}(v,fn,leftfn,def)
mintree(v::Vector{Int}) = itree{Int}(v, min, (x,y)->(x<=y), maximum(v)+1)
maxtree(v::Vector{Int}) = itree{Int}(v, max, (x,y)->(x>=y), minimum(v)-1)

# Query an itree on an interval
function query{T}(it::itree{T}, left::Int, right::Int)
    function qry_inner(cur::itreeNode{T}, l::Int, r::Int)
        if r < cur.left || l > cur.right
            return it.default
        elseif l <= cur.left && cur.right <= r
            return cur.value
        else
            mid_l = cur.children[1].right
            mid_r = cur.children[2].left
            if r <= mid_l
                return qry_inner(cur.children[1], l, r)
            elseif l >= mid_r
                return qry_inner(cur.children[2], l, r)
            else
                return it.fn(qry_inner(cur.children[1], l, min(r,mid_l)), qry_inner(cur.children[2], max(mid_r, l), r))
            end
        end
    end
    return qry_inner(it.root, left, right)
end

# Query an itree on an interval
function query_pos{T}(it::itree{T}, left::Int, right::Int)
    function qry_inner(cur::itreeNode{T}, l::Int, r::Int)
        if r < cur.left || l > cur.right
            return it.default
        elseif l <= cur.left && cur.right <= r
            return (cur.pos, cur.value)
        else
            mid_l = cur.children[1].right
            mid_r = cur.children[2].left
            if r <= mid_l
                return qry_inner(cur.children[1], l, r)
            elseif l >= mid_r
                return qry_inner(cur.children[2], l, r)
            else
                x = qry_inner(cur.children[1], l, min(r,mid_l))
                y = qry_inner(cur.children[2], max(mid_r, l), r)
                if it.isleft(x[2],y[2])
                    return (x[1], it.fn(x[2],y[2]))
                else
                    return (y[1], it.fn(x[2], y[2]))
                end
            end
        end
    end
    return qry_inner(it.root, left, right)[1]
end

##################################################################
# next smaller index algorithm in linear time
##################################################################

# Find the index of the next number in an array to be smaller than
# each element and to the right
function get_next_smaller(A::Vector{Int}, default_elt::Int=0)
    n = length(A)
    unmatched = Base.Collections.PriorityQueue{Int, Int, Base.Order.ForwardOrdering}()
    next_smaller = [default_elt for i in 1:n]

    for i in 1:n
        while !isempty(unmatched)
            (idx, val) = Base.Collections.peek(unmatched)
            val = -val
            if val > A[i]
                next_smaller[idx] = i
                Base.Collections.dequeue!(unmatched)
            else
                break
            end
        end
        unmatched[i] = -A[i]
    end
    return next_smaller
end

##################################################################
# Strongly connected components
##################################################################

# 1-indexed directed adjacency set graph
type directed_adjset_graph
    n::Int
    adj::Array{Set{Int},1}
    directed_adjset_graph(n_vert::Int) = new(n_vert, [Set{Int}() for i in 1:n_vert])
end

function add_edge!(G::directed_adjset_graph, u::Int, v::Int)
    if u == G.n+1 || v == G.n + 1
        push!(G.adj, Set{Int}())
        G.n += 1
    end

    push!(G.adj[u], v)
end

function is_edge(G::directed_adjset_graph, u::Int, v::Int)
    return (v in G.adj[u])
end

# Tarjan's algorithm to compute strongly connected components & condensation
function tarjan(G::directed_adjset_graph)
    index = 0
    dfs_stack   = Vector{Tuple}()
    ddone       = falses(G.n)
    ddescending = trues(G.n)
    S        = Vector{Int}()
    vindex   = zeros(Int, G.n)
    vlowlink = zeros(Int, G.n)
    visited  = falses(G.n)
    onStack  = falses(G.n)
    SCC_map  = zeros(Int, G.n)
    cur_SCC  = 0

    for vertex in 1:G.n
        if !visited[vertex]
            push!(dfs_stack, (vertex, vertex))
            while(!isempty(dfs_stack))
                v_prev,v = dfs_stack[end]
                if ddone[v]
                    pop!(dfs_stack)
                    continue
                end

                visited[v] = true

                # Set the depth index for v to the smallest unused index
                if ddescending[v]
                    ddescending[v] = false

                    vindex[v] = index
                    vlowlink[v] = index
                    index += 1
                    push!(S, v)
                    onStack[v] = true

                    for w in G.adj[v]
                        if !visited[w]
                            push!(dfs_stack, (v,w))
                        elseif onStack[w]
                            vlowlink[v] = min(vlowlink[v], vindex[w])
                        end
                    end
                else
                    v_prev, v = pop!(dfs_stack)
                    ddone[v] = true
                    vlowlink[v_prev] = min(vlowlink[v_prev], vlowlink[v])

                    # If v is a root node, pop the stack and generate an SCC
                    if vlowlink[v] == vindex[v]
                        cur_SCC += 1
                        while !isempty(S)
                            w = pop!(S)
                            onStack[w] = false
                            SCC_map[w] = cur_SCC
                            if w == v
                                break
                            end
                        end
                    end
                end
            end
        end
    end

    condensation = directed_adjset_graph(cur_SCC)
    for i in 1:G.n
        for j in G.adj[i]
            if SCC_map[i] != SCC_map[j]
               add_edge!(condensation, SCC_map[i], SCC_map[j])
           end
        end
    end

    return (SCC_map, condensation)
end

#################################################################
#          Connected components                                 #
#################################################################

# 1-indexed undirected adjacency set graph
type adjset_graph
    n::Int
    adj::Array{Set{Int},1}
    adjset_graph(n_vert::Int) = new(n_vert, [Set{Int}() for i in 1:n_vert])
end

function add_edge!(G::adjset_graph, u::Int, v::Int)
    if u == G.n+1 || v == G.n + 1
        push!(G.adj, Set{Int}())
        G.n += 1
    end

    push!(G.adj[u], v)
    push!(G.adj[v], u)
end

function is_edge(G::adjset_graph, u::Int, v::Int)
    return (v in G.adj[u])
end

immutable Edge
    u::Integer
    v::Integer
end

# Simple DFS (can be done much more efficiently with union-find)
function get_components(g::adjset_graph)
    components = zeros(Int, g.n)
    v0 = 1
    cpt = 1
    while v0 <= g.n
        if components[v0] == 0
            S = [v0]
            while !isempty(S)
                v = pop!(S)
                components[v] = cpt
                for w in g.adj[v]
                    if components[w] == 0
                        push!(S, w)
                    end
                end
            end
            cpt += 1
        end
        v0 += 1
    end
    return components
end

#################################################################
#          Biconnected components (adapted from LightGraphs.jl) #
#################################################################

type Biconnections
    low::Vector{Int}
    depth::Vector{Int}
    stack::Vector{Edge}
    biconnected_comps::Vector{Vector{Edge}}
    is_articulation::Vector{Bool}
    id::Int
end

function Biconnections(g::adjset_graph)
    n = g.n
    return Biconnections(zeros(Int, n), zeros(Int, n), Vector{Edge}(), Vector{Vector{Edge}}(), falses(n),0)
end

# Biconnected component dfs
function bc_dfs!(g::adjset_graph, state::Biconnections, u::Integer, v::Integer)
    children = 0
    state.id += 1
    state.depth[v] = state.id
    state.low[v] = state.depth[v]

    for w in g.adj[v]
        if state.depth[w] == 0
            children += 1
            push!(state.stack, Edge(min(v, w), max(v, w)))
            bc_dfs!(g, state, v, w)
            state.low[v] = min(state.low[v], state.low[w])

            if (u == v && children > 1) || (u != v && state.low[w] >= state.depth[v])
                state.is_articulation[v] = true
                e = Edge(0, 0)
                st = Vector{Edge}()
                while e != Edge(min(v, w),max(v, w))
                    e = pop!(state.stack)
                    push!(st, e)
                end
                push!(state.biconnected_comps, st)
            end

        elseif w != u && state.low[v] > state.depth[w]
            push!(state.stack, Edge(min(v, w), max(v, w)))
            state.low[v] = state.depth[w]
        end
    end
end

function do_bc_dfs(g::adjset_graph)
    state = Biconnections(g)
    for u in 1:g.n
        if state.depth[u] == 0
            bc_dfs!(g, state, u, u)
        end

        if !isempty(state.stack)
            push!(state.biconnected_comps, reverse(state.stack))
            empty!(state.stack)
        end
    end
    return state
end

function biconnected_components(g::adjset_graph)
    state = do_bc_dfs(g)
    return state.biconnected_comps
end

# Root a forest at the best possible roots, creating a parent array
# Root a forest at the best possible roots, creating a parent array
function root_forest(T::adjset_graph)
    vis_1 = falses(T.n)
    vis_2 = falses(T.n)
    vis_3 = falses(T.n)
    par = zeros(Int, T.n)
    dep = zeros(Int, T.n)
    v0 = 1
    while v0 <= T.n
        if par[v0] == 0
            v = v0

            # BFS 1: random node (v) to furthest point (w)
            vis_1[v] = true
            S = [v]
            w = v
            while !isempty(S)
                w = S[1]
                deleteat!(S, 1)
                for nbr in T.adj[w]
                    if !vis_1[nbr]
                        vis_1[nbr] = true
                        push!(S, nbr)
                    end
                end
            end

            # BFS 2: previous furthest point (v <- w_old) to furthest point from it (w)
            v = w
            S = [(v, 0)]
            vis_2[v] = true
            d = 0

            while !isempty(S)
                w, d = S[1]
                deleteat!(S,1)
                for nbr in T.adj[w]
                    if !vis_2[nbr]
                        vis_2[nbr] = true
                        push!(S, (nbr, d+1))
                    end
                end
            end

            # DFS: go from v to w & find half-way mark - the root of the current tree
            S = [v]
            npops = div(d+(d%2),2)

            root = v
            hst = []
            while !isempty(S)
                wcur = pop!(S)
                push!(hst, wcur)
                vis_3[wcur] = true

                if wcur == w
                    root = pop!(hst)
                    for idx in 1:npops
                        root = pop!(hst)
                    end
                    break
                end
                for nbr in T.adj[wcur]
                    if !vis_3[nbr]
                        vis_3[nbr] = true
                        push!(S, nbr)
                    end
                end
            end

            # Root the parent array at root with... yet another DFS!
            S = [root]
            par[root] = root # parent of root is itself by definition
            while !isempty(S)
                w = pop!(S)
                for nbr in T.adj[w]
                    if par[nbr] == 0
                        par[nbr] = w
                        dep[nbr] = dep[w] + 1
                        push!(S, nbr)
                    end
                end
            end
        end

        v0 += 1
    end
    return (par, dep)
end

# Get the block cut tree of an undirected graph
function block_cut_tree(g::adjset_graph)
    state = do_bc_dfs(g)
    a_pts = find(state.is_articulation)
    N0 = length(state.biconnected_comps)
    N1 = length(a_pts)
    N  = N0 + N1
    a_lkup = Dict{Int,Int}(a_pts[i] => N0+i for i in 1:N1)
    bctree = adjset_graph(N)
    v_lkup = zeros(Int,g.n)

    for cpt in 1:N0
        for e in state.biconnected_comps[cpt]
            if state.is_articulation[e.u]
                add_edge!(bctree, a_lkup[e.u], cpt)
                v_lkup[e.u] = a_lkup[e.u]
            else
                v_lkup[e.u] = cpt
            end
            if state.is_articulation[e.v]
                add_edge!(bctree, a_lkup[e.v], cpt)
                v_lkup[e.v] = a_lkup[e.v]
            else
                v_lkup[e.v] = cpt
            end
        end
    end

    return (v_lkup, bctree)
end

function scratch_bc()
    # sample code
    g = adjset_graph(9)
    add_edge!(g, 1,2)
    add_edge!(g, 1,3)
    add_edge!(g, 2,4)
    add_edge!(g, 3,4)
    add_edge!(g, 4,5)
    add_edge!(g, 4,6)
    add_edge!(g, 5,7)
    add_edge!(g, 6,7)
    add_edge!(g, 5,9)
    add_edge!(g, 7,8)
    add_edge!(g, 8,9)
    add_edge!(g, 7,10)
    add_edge!(g, 10,11)
    add_edge!(g, 10,12)

    biconnected_components(g)
end

#################################################################
#          Fast intersection of two sorted lists                #
#################################################################

function intersect_sorted(x::Vector{Int}, y::Vector{Int})
    if isempty(x) || isempty(y)
        return Vector{Int}
    end
    px = 1
    py = 1
    sol = Vector{Int}()
    while px <= length(x) && py <= length(y)
        while px <= length(x) && x[px] < y[py]
            px += 1
        end
        if px > length(x)
            break
        end

        while py <= length(y) && y[py] < x[px]
            py += 1
        end
        if py > length(y)
            break
        end

        while px <= length(x) && py <= length(y) && x[px] == y[py]
            push!(sol, x[px])
            px += 1
            py += 1
        end
    end
    return sol
end

##################################################################
# TRIE
##################################################################
type MinTrie
    value::Int
    children::Dict{Char,MinTrie}
end


"""
    bivariate_gauss(rho, num_elements)

returns a bivariate gauss series of lengths `num_elements` with correlation
coefficient `rho`

# Parameters
    * len : seriese length
    * rho : correlation coefficient, 0 <= rho < 1, and rho = exp(-1/tau_exp)
"""
function bivariate_gauss(rng, rho, num_elements::Int)
    series = Vector{Float64}(undef, num_elements)

    series[1] = bivariate_gauss(rng)
    for i in 2:num_elements
        series[i] = bivariate_gauss(rng, rho, series[i-1])
    end

    return series
end

function bivariate_gauss(rng, rho, x_::Float64)
    return rho*x_ + sqrt(1-rho^2)*randn(rng)
end
function bivariate_gauss(rng)
    return randn(rng)
end


using DDD
using Test
using DDD: makeInstanceDict, inclusiveComparison

cd(@__DIR__)

@testset "Geometry" begin
    arr = Int[3; 4; 6]
    @test isapprox(internalAngle(arr[1]), π / 3)
    @test isapprox(internalAngle(arr[2]), π / 2)
    @test isapprox(internalAngle(arr[3]), 2π / 3)
    @test isapprox(externalAngle(arr[1]), 2π / 3)
    @test isapprox(externalAngle(arr[2]), π / 2)
    @test isapprox(externalAngle(arr[3]), π / 3)
    xyz = [1.0, 0.0, 0.0]
    θ = pi / 2
    uvw = [0.0, 5.0, 0.0]
    abc = [0.0, 0.0, 0.0]
    p = rot3D(xyz, uvw, abc, θ)
    @test isapprox(p, [0.0, 0.0, -1.0])
    uvw = [0.0, 0.0, 20.0]
    abc = [0.0, 0.0, 0.0]
    p = rot3D(xyz, uvw, abc, θ)
    @test isapprox(p, [0.0, 1.0, 0.0])
    uvw = [1.0, 0.0, 0.0]
    abc = [0.0, 0.0, 0.0]
    p = rot3D(xyz, uvw, abc, θ)
    @test isapprox(p, xyz)
    xyz = [-23.0, 29.0, -31.0]
    uvw = [11.0, -13.0, 17.0]
    abc = [-2.0, 5.0, 7.0]
    θ = 37 / 180 * pi
    p = rot3D(xyz, uvw, abc, θ)
    @test isapprox(p, [-21.1690, 31.0685, -30.6029]; atol = 1e-4)
    @test compStruct(1, 1.2) == false

    planenorm = Float64[0, 0, 1]
    planepnt = Float64[0, 0, 5]
    raydir = Float64[0, -1, -2]
    raypnt = Float64[0, 0, 10]

    ψ = linePlaneIntersect(planenorm, planepnt, raydir, raypnt)
    @test isapprox(ψ, [0, -2.5, 5.0])

    planenorm = Float64[0, 2, 1]
    planepnt = Float64[0, 0, 5]
    raydir = Float64[0, -1, -2]
    raypnt = Float64[0, 0, 10]
    ψ = linePlaneIntersect(planenorm, planepnt, raydir, raypnt)
    @test isapprox(ψ, [0.0, -1.25, 7.5])

    planenorm = Float64[0, 0, 1]
    planepnt = Float64[0, 0, 5]
    raydir = Float64[0, 1, 2]
    raypnt = Float64[0, 0, 10]

    ψ = linePlaneIntersect(planenorm, planepnt, raydir, raypnt)
    @test isapprox(ψ, [0, -2.5, 5.0])

    planenorm = Float64[0, 0, 1]
    planepnt = Float64[0, 0, 5]
    raydir = Float64[0, 1, -2]
    raypnt = Float64[0, 0, 10]

    ψ = linePlaneIntersect(planenorm, planepnt, raydir, raypnt)
    @test isapprox(ψ, [0, 2.5, 5.0])

    planenorm = Float64[0, 0, 1]
    planepnt = Float64[0, 0, 5]
    raydir = Float64[0, 1, 0]
    raypnt = Float64[0, 0, 5]
    ψ = linePlaneIntersect(planenorm, planepnt, raydir, raypnt)
    @test isinf(ψ)

    planenorm = Float64[0, 0, 1]
    planepnt = Float64[0, 0, 5]
    raydir = Float64[0, 1, 0]
    raypnt = Float64[0, 0, 6]
    ψ = linePlaneIntersect(planenorm, planepnt, raydir, raypnt)
    @test isnothing(ψ)

    x0, x1 = zeros(3), ones(3)
    y0, y1 = zeros(3), zeros(3)
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (0, 0, 0, 0)

    x0, x1 = zeros(3), ones(3)
    y0, y1 = ones(3), ones(3)
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (0, 0, 1, 0)

    y0, y1 = zeros(3), ones(3)
    x0, x1 = zeros(3), zeros(3)
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (0, 0, 0, 0)

    y0, y1 = zeros(3), ones(3)
    x0, x1 = ones(3), ones(3)
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (0, 0, 0, 1)

    x0, x1 = zeros(3), ones(3)
    y0, y1 = 0.5 * ones(3), 0.5 * ones(3)
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (0, 0, 0.5, 0)

    y0, y1 = zeros(3), ones(3)
    x0, x1 = 0.5 * ones(3), 0.5 * ones(3)
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (0, 0, 0, 0.5)

    x0, x1 = zeros(3), ones(3)
    y0, y1 = zeros(3), ones(3) .+ eps(Float64)^2
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (0, 0, 0, 0)

    x0, x1 = zeros(3), ones(3)
    y0, y1 = [1, 0, 0], [2, 1, 1]
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (1, 0, 0, 0)

    x1, x0 = zeros(3), ones(3)
    y0, y1 = [1, 0, 0], [2, 1, 1]
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (6, 0, 1, 1)

    x0, x1 = zeros(3), ones(3)
    y1, y0 = [1, 0, 0], [2, 1, 1]
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (2, 0, 1, 1)

    x1, x0 = zeros(3), ones(3)
    y0, y1 = [1, 1, 0], [1.5, 0.5, 0.5]
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (0.5, 0, 0.25, 0.5)

    x0, x1 = zeros(3), zeros(3)
    y0, y1 = zeros(3), zeros(3)
    vx0, vx1 = zeros(3), zeros(3)
    vy0, vy1 = zeros(3), zeros(3)
    distSq, dDistSqDt, L1, L2 = minimumDistance(x0, x1, y0, y1, vx0, vx1, vy0, vy1)
    @test (distSq, dDistSqDt, L1, L2) == (0, 0, 0, 0)
end

@testset "Auxiliary" begin
    dict = Dict(
        "intFixDln" => nodeTypeDln(2),
        "noneDln" => nodeTypeDln(0),
        "intMobDln" => nodeTypeDln(1),
        "srfFixDln" => nodeTypeDln(4),
        "extDln" => nodeTypeDln(5),
        "srfMobDln" => nodeTypeDln(3),
        "tmpDln" => nodeTypeDln(6),
    )

    @test makeInstanceDict(nodeTypeDln) == dict
    data = rand(5)
    @test inclusiveComparison(data[rand(1:5)], data...)
    @test !inclusiveComparison(data, data[rand(1:5)] * 6)
end

@testset "Quadrature" begin
    n = 17
    a = -13
    b = 23
    x, w = gausslegendre(n, a, b)

    bma = (b - a) * 0.5
    bpa = (b + a) * 0.5
    x1, w1 = gausslegendre(n)
    x1 = bma * x1 .+ bpa
    w1 = bma * w1

    @test x ≈ x1
    @test w ≈ w1
end


type GDA
    n_class::Integer
    class_mean::Matrix
    class_cov::Matrix
    class_priors::Vector
    class_::Vector
end

function GDA(;
            class_mean=zeros(2,2),
            class_cov=zeros(2,2),
            class_priors=zeros(2),
            n_class = 2,
            class_ = zeros(2))
    return GDA(n_class, class_mean, class_cov, class_priors,class_)
end

function train!(model::GDA, X::Matrix, y::Vector)
    n_feature = size(X,2)
    n_sample = size(X,1)
    model.class_ = sort(unique(y))
    model.n_class = length(unique(y))
    model.class_mean = zeros(n_feature, model.n_class)
    model.class_cov = cov(X)
    model.class_priors = zeros(model.n_class)
    for (i,class) in enumerate(model.class_)
        X_temp = X[vec(y.==class),:]
        model.class_mean[:,i] = mean(X_temp,1)
        model.class_priors[i] = size(X_temp,1)/n_sample
    end
end

function predict(model::GDA,
                 X::Matrix)
    n_sample = size(X,1)
    res = zeros(n_sample)
    for i = 1:n_sample
        res[i] = predict(model, vec(X[i,:]))
    end
    return res
end

function predict(model::GDA,
                 X::Vector)
    temp = zeros(model.n_class)
    for (i, class) in enumerate(model.n_class)
        p = log_pdf(model, X, i)
        temp[i] = p[1]
    end
    res = indmax(exp(temp - log(sum(exp(temp)))))
    @show res
    class = model.class_[res]
    return class
end

function log_pdf(model::GDA,
                 X::Vector, class::Integer)
    t = size(X,1)
    c = -0.5 * (X - model.class_mean[:, class])' * pinv(model.class_cov) * (X- model.class_mean[:, class])
    d = -0.5 * log(det(model.class_cov))
    p = -t/2*log(2*pi)
    pp = log(model.class_priors[class])
    return c + d + p + pp
end

# therer is some problem to be fixed
function test_GDA()
    X_train, X_test, y_train, y_test = make_cla()
    model = GDA()
    train!(model,X_train, y_train)
    predictions = predict(model,X_test)
    print("classification accuracy", accuracy(y_test, predictions))
end


#=

This file contains helper functions that are used throughout the module

=#

""" Convenience function to perform (-1)^x using x's parity """
exp_m1(x) = iseven(x) ? 1 : -1

"""
    $(FUNCTIONNAME)(size_parameter)
Computes the number of required Legendre functions  for a given size parameter.
See eq 6 in Sanghavi 2014
- `size_parameter` size parameter of the aerosol (2πr/λ)
The function returns a rounded integer, following conventions by BH, Rooj/Stap, Siewert
"""
get_n_max(size_parameter) = round(Int, size_parameter + 4.05 * size_parameter^(1/3) + 10)

"""
    $(FUNCTIONNAME)(size_parameter,refractive_idx::Number,an,bn,Dn)
Computes Mie coefficients `an` and `bn` as a function of size parameter and complex
refractive index. See eq 4.88 in Bohren and Huffman
- `size_parameter` size parameter of the aerosol (2πr/λ)
- `refractive_idx` refractive index of the aerosol (complex number)
- `an` and `bn` pre-allocated arrays, need to match at least n_max for the given size parameter
- `Dn` pre-allocated array for the logarithmic derivative (see BH, eq 4.88)
(need to check whether Dn can be created internally without causing too many allocations)

The function returns a rounded integer, following conventions by BH, Rooj/Stap, Siewert
"""
function compute_mie_ab!(size_param, refractive_idx::Number, an, bn, Dn)
    FT = eltype(refractive_idx)
    FT2 = eltype(an)

    # Compute y
    y = size_param * refractive_idx

    # Maximum expansion (see eq. A17 from de Rooij and Stap, 1984)
    n_max = get_n_max(size_param)

    # Make sure downward recurrence starts higher up
    # (at least 15, check eq. A9 in de Rooij and Stap, 1984, may need to check what is needed)
    nmx = length(Dn)
    @assert size(an)[1] >= n_max
    @assert size(an) == size(bn)
    fill!(Dn, 0);

    # Dn as in eq 4.88, Bohren and Huffman, to calculate an and bn
    # Downward Recursion, eq. 4.89, Bohren and Huffman
    [Dn[n] = ((n+1) / y) - (1 / (Dn[n+1] + (n+1) / y)) for n = (nmx - 1):-1:1]

    # Get recursion for bessel functions ψ and ξ
    ψ₀, ψ₁, χ₀, χ₁ =  (cos(size_param), sin(size_param), -sin(size_param), cos(size_param))
    ξ₁ = FT(ψ₁, -χ₁)

    # This solves Bohren and Huffman eq. 4.88 for an and bn, computing updated ψ and ξ on the fly
    for n = 1:n_max
        fn = (2n + 1) / (n * (n + 1))
        ψ  = (2n - 1) * ψ₁ / size_param - ψ₀
        χ  = (2n - 1) * χ₁ / size_param - χ₀
        ξ   = FT(ψ, -χ)
        t_a = Dn[n] / refractive_idx + n / size_param
        t_b = Dn[n] * refractive_idx + n / size_param

        an[n] = (t_a * ψ - ψ₁) / (t_a * ξ - ξ₁)
        bn[n] = (t_b * ψ - ψ₁) / (t_b * ξ - ξ₁)

        ψ₀ = ψ₁
        ψ₁ = ψ
        χ₀ = χ₁
        χ₁ = χ
        ξ₁ = FT(ψ₁, -χ₁)
    end
end

"""
    $(FUNCTIONNAME)(model::MieModel, λ, radius)
Compute all an, bn using compute_mie_ab!
Input: MieModel, wavelength (λ), radius
Output: an, bn. Both of shape (aerosol.nquad_radius, N_max) (N_max from aerosol.r_max)
"""
function compute_anbn(model::MieModel, λ, radius)

    @unpack computation_type, aerosol, r_max, nquad_radius, λ, polarization_type, truncation_type, wigner_A, wigner_B = model
    @unpack size_distribution, nᵣ, nᵢ = aerosol

    FT = eltype(λ)
    FT2 = eltype(nᵣ)

    # Find overall N_max from the maximum radius
    N_max = Scattering.get_n_max(2 * π * r_max / λ)

    # Where to store an, bn, computed over size distribution
    an = zeros(Complex{FT2}, nquad_radius, N_max)
    bn = zeros(Complex{FT2}, nquad_radius, N_max)

    # Loop over the size distribution, and compute an, bn, for each size
    for i in 1:nquad_radius

        # Get current radius and size parameter
        r = radius[i]
        size_param = 2 * π * r / λ

        # Pre-allocate Dn:
        y = size_param * (nᵣ - nᵢ);
        nmx = round(Int, max(N_max, abs(y)) + 51)
        Dn = zeros(Complex{FT2}, nmx)

        # Compute an, bn
        Scattering.compute_mie_ab!(size_param, nᵣ + nᵢ * im,
                                      view(an, i, :),
                                      view(bn, i, :), Dn)
    end

    return an, bn;
end

"""
    $(FUNCTIONNAME)(an, bn, ab_pairs, w, Nmax, N_max_)
From the an, bn matrices, precompute all (an✶)am, (an✶)bm, (bn✶)am, (bn✶)bm
This allows quick computation of (an✶ + bn✶) × (am + bm)
"""
function compute_avg_anbns!(an, bn, ab_pairs, w, Nmax, N_max_)
    FT2 = eltype(an)

    # Unpack ab_pairs
    mat_anam, mat_anbm, mat_bnam, mat_bnbm = ab_pairs

    # Fill all matrices with 0
    [fill!(mat, 0) for mat in [mat_anam, mat_bnbm, mat_anbm, mat_bnam]]

    @inbounds for n in 1:Nmax, m in n:Nmax
        anam = bnbm = anbm = bnam = FT2(0);
        @inbounds for i = 1:size(an, 1)
            if m < N_max_[i] && n < N_max_[i]
                anam += w[i] * an[i,n]' * an[i,m]
                bnbm += w[i] * bn[i,n]' * bn[i,m]
                anbm += w[i] * an[i,n]' * bn[i,m]
                bnam += w[i] * bn[i,n]' * an[i,m]
            end
        end
        @inbounds mat_anam[m,n] = anam;
        @inbounds mat_bnbm[m,n] = bnbm;
        @inbounds mat_anbm[m,n] = anbm;
        @inbounds mat_bnam[m,n] = bnam;
    end
end

"""
    $(FUNCTIONNAME)(an, bn, π_, τ_, S₁, S₂)
Determines the amplitude functions `S₁`,`S₂` in Mie theory
- `an` and `bn` pre-calculated Mie coefficients `an` and `bn`, see [`compute_mie_ab!`](@ref) function
- `π` and `τ` pre-calculated associated Legendre functions `π` and `τ`, see [`compute_mie_π_τ!`](@ref) function
The function returns `S₁`,`S₂` as a function of the cosine of the scattering angle `ξ`.
Users need to make sure `an` and `bn`, `π` and `τ` are pre-computed.
"""
function compute_mie_S₁S₂!(an, bn, π_, τ_, S₁, S₂)
    FT = eltype(an)
    nmax = size(an, 1);
    nμ   = size(π_, 1);

    # Verify sizes
    @assert size(S₁) == size(S₂)
    @assert length(S₁) == nμ

    for l in 1:nmax, iμ in 1:nμ
            S₁[iμ] += (2l + 1) / (l * (l + 1)) * (an[l] * τ_[iμ,l] + bn[l] * π_[iμ,l])
            S₂[iμ] += (2l + 1) / (l * (l + 1)) * (an[l] * π_[iμ,l] + bn[l] * τ_[iμ,l])
    end
end

"""
    $(FUNCTIONNAME)(n,xmin,xmax; norm=false)
Returns the `n` Gauss-Legendre quadrature points and weights with a change of interval between xmin and xmax
- `n` number of quadrature points
- `xmin`,`xmax` lower and upper bound of integral
- `norm`: if `true`, normalizes the weights so that a mean can be computed instead of full integration
The function returns `n` quadrature points ξ within [xmin,xmax] with associated weightes `w`
"""
function gauleg(n, xmin, xmax; norm=false)
    ξ, w = gausslegendre(n)
    ξ = (xmax - xmin) / 2 * ξ .+ (xmin + xmax) / 2
    norm ? w /= sum(w) : w *= (xmax - xmin) / 2
    return ξ, w
end

@doc raw"""
    $(FUNCTIONNAME)(greek_coefs, μ; returnLeg = false)
Returns the reconstructed elements of the 4x4 scattering matrix at positions
f₁₁, f₁₂, f₂₂, f₃₃, f₃₄, f₄₄ from the greek coefficients

f₁₁ represents the phase function p for the Intensity (first Stokes Vector element) and is normalized as follows:
```math
\frac{1}{4\pi}\int_0^{2\pi}d\phi \int_{-1}^1 p(\mu) d\mu  = 1
```

- `greek_coefs` greek coefficients (Domke Type)
- `returnLeg` if `false` (default), just return `f₁₁, f₁₂, f₂₂, f₃₃, f₃₄, f₄₄`, if `true`,
- return `f₁₁, f₁₂, f₂₂, f₃₃, f₃₄, f₄₄, P, P²` (i.e. also the two legendre polynomials as matrices)
"""
function reconstruct_phase(greek_coefs, μ; returnLeg=false)

    FT = eltype(greek_coefs.α)
    l_max = length(greek_coefs.α);
    nμ = length(μ)

    # Compute legendre polynomials
    P, P², R², T² = compute_legendre_poly(μ, l_max)

    # To stay general, we also don't assume f₂₂=f₁₁ or f₄₄=f₃₃
    # which only holds for spherical
    f₁₁, f₃₃, f₁₂, f₃₄, f₂₂, f₄₄ = (zeros(FT, nμ), zeros(FT, nμ), zeros(FT, nμ),
                                    zeros(FT, nμ), zeros(FT, nμ), zeros(FT, nμ))

    # Compute prefactor
    fac = zeros(l_max);
    [fac[l + 1] = sqrt(1 / ((l-1) * l * (l+1) * (l+2))) for l = 2:(l_max-1)]

    # In matrix form:
    f₁₁[:] = P * greek_coefs.β                                           # a₁ in Rooij notation
    f₄₄[:] = P * greek_coefs.δ                                           # a₄ in Rooij notation
    f₁₂[:] = P² * (fac .* greek_coefs.γ)                                 # b₁ in Rooij notation
    f₃₄[:] = P² * (fac .* greek_coefs.ϵ)                                 # b₂ in Rooij notation
    f₂₂[:] = R² * (fac .* greek_coefs.α) .+ T² * (fac .* greek_coefs.ζ)  # a₂ in Rooij notation
    f₃₃[:] = R² * (fac .* greek_coefs.ζ) .+ T² * (fac .* greek_coefs.α)  # a₃ in Rooij notation

    # Put elements into a struct
    scattering_matrix = ScatteringMatrix(f₁₁, f₁₂, f₂₂, f₃₃, f₃₄, f₄₄)

    # For truncation in δ-BGE, we need P and P² as well, convenient to return here:
    return returnLeg ? (scattering_matrix, P, P²) : scattering_matrix
end

"""
    $(FUNCTIONNAME)(depol)
Returns the greek coefficients (as [`GreekCoefs`](@ref)) of the Rayleigh phase function given
depolarization value.
- `depol` Depolarization (best use 0 as default )
"""
function get_greek_rayleigh(depol::Number)
    FT = eltype(depol)
    # Rayleigh Greek Parameters
    dpl_p = (1 - depol)  / (1 + depol / 2)
    dpl_q = (1 + depol)  / (1 - depol)
    dpl_r = (1 - 2depol) / (1 - depol)

    α  =  FT[0.0, 0.0,             3dpl_p]
    β  =  FT[1.0, 0.0,             0.5 * dpl_p]
    γ  =  FT[0.0, 0.0,             dpl_p * sqrt(1.5)]
    δ  =  FT[0.0, dpl_p * dpl_r * 1.5, 0.0]
    ϵ  =  FT[0.0, 0.0,             0.0]
    ζ  =  FT[0.0, 0.0,             0.0]
    return GreekCoefs(α, β, γ, δ, ϵ, ζ)
end

"""
    $(FUNCTIONNAME)(k, an, bn, w)
Calculate the average Scattering and Extinction Cross Section
Eqn. 1, averaged over size distribution
"""
function compute_avg_C_scatt_ext(k, an, bn, w)
    n_ = collect(1:size(an)[2]);
    n_ = 2n_ .+ 1
    coef = 2π / k^2 * n_'
    return (coef * (w' * (abs2.(an') + abs2.(bn'))')', coef * (w' * real(an + bn))')
end

""" Compute probability weights of radii """
function compute_wₓ(size_distribution, wᵣ, r, r_max)

    wₓ = pdf.(size_distribution,r)      # Weights from distribution
    wₓ .*= wᵣ                           # pre multiply with wᵣ to get proper means eventually:

    # normalize (could apply a check whether cdf.(size_distribution,r_max) is larger than 0.99:
    #println("Test")
    @info "Fraction of size distribution cut by max radius: $((1-cdf.(size_distribution,r_max))*100) %"
    wₓ /= sum(wₓ)
    return wₓ
end

#####
##### Π-matrix construction methods (Sanghavi 2014, eq. 15)
#####

"""
    $(FUNCTIONNAME)(mo::Stokes_IQUV, P, R, T, l::Int, m::Int; sign_change=false)
Compute Π matrix for all stokes vector elements used in computations of the phase matrix
See Sanghavi 2014, eq. 15
"""
function construct_Π_matrix(mo::Stokes_IQUV, P, R, T, l::Int, m::Int; sign_change=false)
    if sign_change # (basically gets it for -μ due to symmetries on P,R,T)
        if isodd(l-m)
            Π = [SMatrix{4,4}([-P[i,l,m] 0 0 0 ; 0 -R[i,l,m] -T[i,l,m] 0; 0 -T[i,l,m] -R[i,l,m] 0; 0 0 0 -P[i,l,m]]) for i in 1:size(P, 1)]
        else
            Π = [SMatrix{4,4}([P[i,l,m] 0 0 0 ; 0 R[i,l,m] T[i,l,m] 0; 0 T[i,l,m] R[i,l,m] 0; 0 0 0 P[i,l,m]]) for i in 1:size(P, 1)]
        end
    else
        Π = [SMatrix{4,4}([P[i,l,m] 0 0 0 ; 0 R[i,l,m] -T[i,l,m] 0; 0 -T[i,l,m] R[i,l,m] 0; 0 0 0 P[i,l,m]]) for i in 1:size(P, 1)]
    end
    return Π
end

"""
    $(FUNCTIONNAME)(mo::Stokes_IQU, P, R, T, l::Int, m::Int; sign_change=false)
Compute Π matrix for  stokes vector elements I,Q,U used in computations of the phase matrix
See Sanghavi 2014, eq. 15
"""
function construct_Π_matrix(mo::Stokes_IQU, P, R, T, l::Int, m::Int; sign_change=false)
    if sign_change # (basically gets it for -μ due to symmetries on P,R,T)
        if isodd(l-m)
            Π = [SMatrix{3,3}([-P[i,l,m] 0 0  ; 0 -R[i,l,m] -T[i,l,m] ; 0 -T[i,l,m] -R[i,l,m] ]) for i in 1:size(P, 1)]
        else
            Π = [SMatrix{3,3}([P[i,l,m] 0 0  ; 0 R[i,l,m] T[i,l,m] ; 0 T[i,l,m] R[i,l,m] ]) for i in 1:size(P, 1)]
        end
    else
        Π = [SMatrix{3,3}([P[i,l,m] 0 0  ; 0 R[i,l,m] -T[i,l,m] ; 0 -T[i,l,m] R[i,l,m] ]) for i in 1:size(P, 1)]
    end
    return Π
end

"""
    $(FUNCTIONNAME)(mo::Stokes_I, P, R, T, l::Int, m::Int; sign_change=false)
Compute Π matrix for  stokes vector elements I used in computations of the phase matrix

"""
construct_Π_matrix(mod::Stokes_I, P, R, T, l::Int, m::Int; sign_change=false) = sign_change ? -P[:,l,m] : P[:,l,m]

#####
##### B-matrix construction methods (Sanghavi 2014, eq. 16)
#####

"""
    $(FUNCTIONNAME)(mod::Stokes_IQUV, α, β, γ, δ, ϵ, ζ, l::Int)
Compute B matrix for all stokes vector elements used in computations of the phase matrix
See Sanghavi 2014, eq. 16
"""
construct_B_matrix(mod::Stokes_IQUV, α, β, γ, δ, ϵ, ζ, l::Int) = SMatrix{4,4}([β[l] γ[l] 0 0 ; γ[l] α[l] 0 0; 0 0 ζ[l] ϵ[l]; 0 0 -ϵ[l] δ[l]])

"""
    $(FUNCTIONNAME)(mod::Stokes_IQU, α, β, γ, δ, ϵ, ζ, l::Int)
Compute B matrix for stokes vector elements I,Q,U used in computations of the phase matrix
    See Sanghavi 2014, eq. 16
"""
construct_B_matrix(mod::Stokes_IQU, α, β, γ, δ, ϵ, ζ, l::Int) = SMatrix{3,3}([β[l] γ[l] 0 ; γ[l] α[l] 0 ; 0 0 ζ[l]])

"""
$(FUNCTIONNAME)(mod::Stokes_I, α, β, γ, δ, ϵ, ζ, l::Int)
Compute Π matrix for stokes vector elements I used in computations of the phase matrix
See Sanghavi 2014, eq. 16
"""
construct_B_matrix(mod::Stokes_I, α, β, γ, δ, ϵ, ζ, l::Int) = β[l]

"""
    $(FUNCTIONNAME)(mod::AbstractPolarizationType, μ, α, β, γ, δ, ϵ, ζ, m::Int)
Compute moments of the phase matrix
"""
function compute_Z_moments(mod::AbstractPolarizationType, μ, greek_coefs::GreekCoefs, m::Int ; arr_type = Array)
    @unpack α, β, γ, δ, ϵ, ζ = greek_coefs
    FT = eltype(β)
    n = length(μ)

    # Change from 0-index to 1-index (i.e. the lowest m is 0 ),
    # make more logical later to avoid confusion later (m=0 has a meaning!!)
    m = m+1

    # Set prefactor for moments (note 1-notation for `m` here):
    fact = (m == 1) ? 0.5 : 1.0

    # get l_max just from length of array:
    l_max = length(β)

    # Check that all μ are positive here ([0,1])
    # @show μ
    @assert all(0 .< μ .≤ 1) "all μ's within compute_Z_moments have to be ∈ ]0,1]"

    # Compute legendre Polynomials at μ and up to lmax
    P, R, T    = Scattering.compute_associated_legendre_PRT(μ, l_max)
    P⁻, R⁻, T⁻ = Scattering.compute_associated_legendre_PRT(-μ, l_max)

    # Pre-compute all required B matrices
    𝐁_all = [construct_B_matrix(mod, α, β, γ, δ, ϵ, ζ, i) for i in 1:l_max]

    # Get dimension of square matrix (easier for Scalar/Stokes dimensions)
    B_dim = Int(sqrt(length(𝐁_all[1])))

    # Create matrices:
    nb = B_dim * n
    𝐙⁺⁺, 𝐙⁻⁺ = (zeros(FT, nb, nb), zeros(FT, nb, nb))
    A⁺⁺, A⁻⁺ = (zeros(FT, B_dim, B_dim, n, n), zeros(FT, B_dim, B_dim, n, n))

    # Iterate over l
    for l = m:l_max

        # B matrix for l
        𝐁 = 𝐁_all[l];

        # Construct Π matrix for l,m pair (change to in place later!)
        # See eq. 15 in Sanghavi 2014, note that P,R,T are already normalized
        Π  = construct_Π_matrix(mod, P, R, T, l, m)
        Π⁻ = construct_Π_matrix(mod, P⁻, R⁻, T⁻, l, m)

        # Iterate over angles
        for i in eachindex(μ), j in eachindex(μ)
            if B_dim == 1
                A⁺⁺[B_dim,B_dim,i,j] += Π[i] * 𝐁 * Π[j]
                A⁻⁺[B_dim,B_dim,i,j] += Π[i] * 𝐁 * Π⁻[j]
            else
                A⁺⁺[:,:,i,j] += Π[i] * 𝐁 * Π[j]
                A⁻⁺[:,:,i,j] += Π[i] * 𝐁 * Π⁻[j]
            end
        end
    end

    # Now get to the Z part:
    for imu in eachindex(μ), jmu in eachindex(μ)

        # Indices adjusted for size of A
        ii, jj = ((imu - 1) * B_dim, (jmu - 1) * B_dim)

        # This is equivalent to Z̄ = 1/(1+δ) * C̄m+S̄m = 1/(1+δ) * (A+DAD+AD-DA)
        # (see eq 11 in Sanghavi et al, 2013)
        for i in 1:B_dim, j in 1:B_dim
            𝐙⁺⁺[ii + i,jj + j] = 2fact * A⁺⁺[i,j,imu,jmu]
            if i <= 2 && j >= 3
                𝐙⁻⁺[ii + i,jj + j] = -2fact * A⁻⁺[i,j,imu,jmu]
            elseif i >= 3 && j <= 2
                𝐙⁻⁺[ii + i,jj + j] = -2fact * A⁻⁺[i,j,imu,jmu]
            else
                𝐙⁻⁺[ii + i,jj + j] = 2fact * A⁻⁺[i,j,imu,jmu]
            end
        end
    end

    # Return Z-moments
    return arr_type(𝐙⁺⁺), arr_type(𝐙⁻⁺)
end


struct BeamSection
    id::String
    hid::Int
    A::Float64
    I₂::Float64
    I₃::Float64
    J::Float64
    As₂::Float64
    As₃::Float64
    W₂::Float64
    W₃::Float64
    sec_type::SectionType
    sizes::Vector
    BeamSection(id,hid,A,I₂,I₃,J,As₂,As₃,W₂,W₃,sec_type=SectionType(0),sizes=[])=new(string(id),hid,A,I₂,I₃,J,As₂,As₃,W₂,W₃,sec_type,sizes)
end

# function BeamSection(id,hid,A,I₂,I₃,J,As₂,As₃,W₂,W₃)
#     id=string(id)
# end

function RectangleSection(id,hid,h,b)::BeamSection
    id=string(id)

    A=h*b
    I₃=b*h^3/12
    I₂=h*b^3/12

    bb,aa=sort([h,b])
    J=aa*bb^3*(1/3-0.21*bb/aa*(1-bb^4/12/aa^4))

    As₂=A#wrong
    As₃=A#wrong
    W₃=I₃/h*2
    W₂=I₂/b*2

    BeamSection(id,hid,A,I₂,I₃,J,As₂,As₃,W₂,W₃,Enums.RECTANGLE,[h,b])
end

function HSection(id,hid,h,b,tw,tf)::BeamSection
    id=string(id)

    A=b*tf*2+tw*(h-2*tf)
    J=(b*tf^3*2+(h-tf)*tw^3)/3 #should be confirm!!!!!!!!!!
    I₃=b*h^3/12-(b-tw)*(h-2*tf)^3/12
    I₂=2*tf*b^3/12+(h-2*tf)*tw^3/12
    As₂=tw*(h-2tf)#wrong
    As₃=2*b*tf#wrong
    W₃=I₃/h*2
    W₂=I₂/b*2
    BeamSection(id,hid,A,I₂,I₃,J,As₂,As₃,W₂,W₃,Enums.HSECTION,[h,b,tw,tf])
end

function ISection(id,hid,h,b1,b2,tw,tf1,tf2)::BeamSection
    id=string(id)

    hw=h-tf1-tf2
    A=b1*tf1+b2*tf2+tw*hw
    y0=(b1*tf1*(h-tf1/2)+b2*tf2*tf2/2+hw*tw*(hw/2+tf2))/A

    J=(b1*tf1^3+b2*tf2^3+(h-tf1/2-tf2/2)*tw^3)/3 #should be confirm!!!!!!!!!!

    I₃=tw*hw^3/12
    I₃+=b1*tf1^3/12+b1*tf1*(hw/2+tf1/2)^2
    I₃+=b2*tf2^3/12+b2*tf2*(hw/2+tf2/2)^2
    I₃-=A*(y0-h/2)^2

    I₂=b1^3*tf1/12+b2^3*tf2/12+tw^3*hw/12

    As₂=tw*hw#wrong
    As₃=b1*tf1+b2*tf2#wrong
    W₃=I₃/max(y0,h-y0)
    W₂=I₂/max(b1/2,b2/2)

    BeamSection(id,hid,A,I₂,I₃,J,As₂,As₃,W₂,W₃,Enums.I_SECTION,[h,b1,b2,tw,tf1,tf2])
end

function BoxSection(id,hid,h,b,tw,tf)::BeamSection
    id=string(id)

    A=h*b-(h-2*tf)*(b-2*tw)
    J=(2*tw*(h-tf)/tw+2*tf*(b-tw)/tf)*2*A
    I₃=b*h^3/12-(b-2*tw)*(h-2*tf)^3/12
    I₂=h*b^3/12-(h-2*tf)*(b-2*tw)^3/12

    As₂=2*h*tw#wrong
    As₃=2*b*tf#wrong
    W₃=I₃/h*2
    W₂=I₂/b*2
    BeamSection(id,hid,A,I₂,I₃,J,As₂,As₃,W₂,W₃,Enums.BOX,[h,b,tw,tf])
end

function PipeSection(id,hid,d,t)::BeamSection
    id=string(id)

    A=π*d^2/4-π*(d-2*t)^2/4
    J=π*(d-t)/t*2*A
    I₃=π*d^4/64*(1-((d-2*t)/d)^4)
    I₂=I₃

    As₂=A#wrong
    As₃=A#wrong
    W₃=I₃/d*2
    W₂=W₃
    r=d/2
    J=2/3*π*r*t^3
    BeamSection(id,hid,A,I₂,I₃,J,As₂,As₃,W₂,W₃,Enums.PIPE,[d,t])
end

function CircleSection(id,hid,d)::BeamSection
    id=string(id)

    A=π*d^2/4
    J=π*d^4/32
    I₃=π*d^4/64
    I₂=I₃

    As₂=A#wrong
    As₃=A#wrong
    W₃=I₃/d*2
    W₂=W₃
    J=π*d^4/32
    BeamSection(id,hid,A,I₂,I₃,J,As₂,As₃,W₂,W₃,Enums.CIRCLE,[d])
end

# end


using QuadGK

"""
    diagnose_buoyancy_flux(b, z, time_period, depth)

Returns the buoyancy flux (m²/s³) required to extract all the buoyancy from a 1D buoyancy profile `b(z)`
down to some z-coordinate `depth` (< 0) over a `time_period` τ (s) by integrating the buoyancy profile
and converting the integrated buoyancy content to an equivalent buoyancy flux:

          1  ⌠ z_surface
    Qb = --- |           [b(z) - b(depth)] dz
          τ  ⌡ depth
"""
function diagnose_buoyancy_flux(b, z, time_period, depth)

    Δz = z[2] - z[1]
    z_surface = last(z)
    @assert first(z) < depth < z_surface

    ℑb = InterpolatedProfile(b, z, Δz, z_surface)
    b₀ = ℑb(depth)

    integrand(z) = ℑb(z) - b₀
    B, err = quadgk(integrand, depth, z_surface)

    return B / time_period
end


using NLsolve
using ForwardDiff
using LinearAlgebra
using Printf
using Plots
include("../src/algorithms/linesearches.jl")

# ELEMENTAL FUNCTIONS

function gfun!(F, x, u, p)

  # retrieve variables
  VM3 = x[1]
  VA3 = x[2]
  VA2 = x[3]
  P1  = x[4]

  VM1 = u[1]
  P2 = u[2]
  VM2 = u[3]

  VA1 = p[1]
  P3 = p[2]
  Q3 = p[3]

  # intermediate quantities
  VA23 = VA2 - VA3
  VA31 = VA3 - VA1
  VA32 = VA3 - VA2
  VA13 = VA1 - VA3

  F[1] = 4.0*VM2*VM2 + VM2*VM3*(-4*cos(VA23) + 10*sin(VA23)) - P2
  F[2] = (8.0*VM3*VM3 + VM3*VM1*(-4*cos(VA31) + 5*sin(VA31))
          + VM3*VM2*(-4*cos(VA32) + 10*sin(VA32)) + P3)
  F[3] = (15.0*VM3*VM3 + VM3*VM1*(-4*sin(VA31) - 5*cos(VA31))
          + VM3*VM2*(-4*sin(VA32) - 10*cos(VA32)) + Q3)
  F[4] = 4.0*VM1*VM1 + VM1*VM3*(-4*cos(VA13) + 5*sin(VA13)) - P1
end

function gfun(x, u, p, T=typeof(x))
  # Get Float64 of Vectors{Float64}, that is the first parameter
  F = zeros(T.parameters[1], 4)

  # retrieve variables
  VM3 = x[1]
  VA3 = x[2]
  VA2 = x[3]
  P1  = x[4]

  VM1 = u[1]
  P2 = u[2]
  VM2 = u[3]

  VA1 = p[1]
  P3 = p[2]
  Q3 = p[3]

  # intermediate quantities
  VA23 = VA2 - VA3
  VA31 = VA3 - VA1
  VA32 = VA3 - VA2
  VA13 = VA1 - VA3

  F[1] = 4.0*VM2*VM2 + VM2*VM3*(-4*cos(VA23) + 10*sin(VA23)) - P2
  F[2] = (8.0*VM3*VM3 + VM3*VM1*(-4*cos(VA31) + 5*sin(VA31))
          + VM3*VM2*(-4*cos(VA32) + 10*sin(VA32)) + P3)
  F[3] = (15.0*VM3*VM3 + VM3*VM1*(-4*sin(VA31) - 5*cos(VA31))
          + VM3*VM2*(-4*sin(VA32) - 10*cos(VA32)) + Q3)
  F[4] = 4.0*VM1*VM1 + VM1*VM3*(-4*cos(VA13) + 5*sin(VA13)) - P1
  return F
end

function cfun(x, u, p)

  VM3 = x[1]
  VA3 = x[2]
  P1  = x[4]

  VM1 = u[1]
  P2 = u[2]

  VA1 = p[1]

  # we fix generation weights inside the
  # function to simplify the script and
  # follow the paper closely.
  w1 = 1.0
  w2 = 1.0

  cost = P1 + P2

  return cost
end

# OPF COMPUTATION

function solve_pf(x, u, p, verbose=true)

  fun_pf!(F, x) = gfun!(F, x, u, p)
  x0 = copy(x)
  res = nlsolve(fun_pf!, x0)
  if verbose
    show(res)
    println("")
  end

  return res.zero

end

# initial parameters
x = zeros(4)
u = zeros(3)
p = zeros(3)

# this is an initial guess
x[1] = 1.0 #VM3
x[2] = 0.0 #VA3
x[3] = 0.0 #VA2
x[4] = 0.5 #P1

# this is given by the problem data, but might be "controlled" via OPF
u[1] = 1.0 #VM1
u[2] = 1.7 #P2
u[3] = 1.0 #VM2

# these parameters are fixed through the computation
p[1] = 0.0 #VA1, slack angle
p[2] = 2.0 #P3
p[3] = 1.0 #Q3

# print initial guesses
println(x)
println(u)

# copy to iteration vectors
xk = copy(x)
uk = copy(u)

iterations = 0

for i = 1:1000
  global xk
  global uk
  println("Iteration ", i)

  # solve power flow
  println("Solving power flow")
  xk = solve_pf(xk, uk, p, false)

  # jacobian
  gx_x(x) = gfun(x, uk, p, typeof(x))
  gx = x -> ForwardDiff.jacobian(gx_x, x)

  # gradient
  cfun_x(x) = cfun(x, uk, p)
  fx = x ->ForwardDiff.gradient(cfun_x, x)

  # lamba calculation
  println("Computing Lagrange multipliers")
  lambda = -inv(gx(xk)')*fx(xk)

  # compute g_u, g_u
  cfun_u(u) = cfun(xk, u, p)
  fu = u ->ForwardDiff.gradient(cfun_u, u)
  gx_u(u) = gfun(xk, u, p, typeof(u))
  gu = u -> ForwardDiff.jacobian(gx_u, u)

  # compute gradient of cost function
  Lu = u -> cfun_u(u) + gx_u(u)'*lambda
  grad_Lu = u -> (fu(u) + gu(u)'*lambda)
  grad_L = grad_Lu(uk)
  println("Norm of gradient ", norm(grad_L))

  # step
  println("Computing new control vector")
  c_par = 1.0
  # Optional linesearch
  c_par = ls(uk, grad_L, Lu, grad_Lu)
  @show c_par
  uk = uk - c_par*grad_L

  #@printf("VM3 %3.2f. VA3 %2.2f. VA2 %2.2f.\n", xk[1], xk[2], xk[3])
  #@printf("VM1 %3.2f. P2 %2.2f. VM2 %2.2f.\n", uk[1], uk[2], uk[3])
end