module Utils

using Interpolations
using HDF5
using SpecialFunctions
using LinearAlgebra # pinv

export build_interpolation,
	   _readh5_file,
	   multiple_reflections,
	   find_closest,
	   blackbody_radiation,
	   wavelength2rgb,
	   simple_moving_average,
	   exponential_moving_average,
	   savitzky_golay,
	   flatten_arrays,
	   array_arrays,
	   average_polarisation,
	   gaussian,
	   lorentzian,
	   voigtian,
	   _oscillators_input,
	   unfold_bnd,
	   Info

function Info()
    tmp1 = "\n " *
	"\n " *
    "\n Available functions from Utils module:" *
	"\n " *
	"\n     build_interpolation(X)" *
	"\n     multiple_reflections(n)" *
	"\n     find_closest(x, x0)" *
	"\n     blackbody_radiation(λ, T)" *
	"\n     wavelength2rgb(λ)" *
	"\n     simple_moving_average(x)" *
	"\n     exponential_moving_average(x)" *
	"\n     flatten_arrays(x)" *
	"\n     array_arrays(x, y)" *
	"\n     average_polarisation(pol, Xp, Xs)" *
	"\n     gaussian(x, p)" *
	"\n     lorentzian(x, p)" *
	"\n     voigtian(x, p)" *
	"\n "
	"\n     To use any of these functions type: ?Utils.function" *
	"\n "
	return println(tmp1)
end

"""

	Find the index of the closest element in an array of reals x to a given real x0.

		idx = Utils.find_closest(x, x0)

			x: Data array
			x0: Element to search in x

			idx: Index

"""
function find_closest(x::AbstractArray{T1,N1}, x0::T1) where {T1<:Real, N1}
	idmin = findmin(abs.(x .- x0))[2][1]
	return idmin
end

"""Read the HDF5 file with the database."""
function _readh5_file(x::String, y::Symbol)
    datapath = joinpath(@__DIR__, "..", "data/")
	if isequal(y,:RI)
		fid = "RefractiveIndicesDB.h5"
	elseif isequal(y,:SP)
		fid = "ExampleSpectraDB.h5"
	end
    readf = h5open(joinpath(datapath,fid),"r") do file
        read(file,x)
    end
    return readf
end

"""

	Build interpolation objects from Matrix. Uses a Gridded(Linear()) grid.

		itp = Utils.build_interpolation(X)

			X: Matrix array with first column as independent variable and
               second column as dependent one

			itp: interpolation object

"""
function build_interpolation(X::Array{T1,2}) where {T1<:Float64}
	# Sort by ascending wavelength
	X = X[sortperm(X[:,1]),:]
    knots = (sort(vec(X[:,1])),)
	itp = interpolate(knots, vec(X[:,2]), Gridded(Linear()))
    return itp
end

"""

	Computes the reflection of a sequence of indices of refraction without
    interference and absorption effects.

		y = Utils.multiple_reflections(n)

			n: Vector with the indices of refraction of the involved layers

			y: Reflection

	author: Cuchu

"""
function multiple_reflections(n::AbstractArray{T1}) where {T1<:Real}
	r1::Float64 = 0.0
	r::Float64 = 0.0
	for i = 1:length(n)-1
		r2 = ((n[i] - n[i+1])/(n[i] + n[i+1]))^2
		r = (r1 + r2 - 2.0*r1*r2)/(1.0 - r1*r2)
		r1 = r
	end
	return r
end

"""

	Returns the spectral radiance (W·sr^-1·m^-3), given the wavelength (m)
    and the absolute temperature (K).

		B = Utils.blackbody_radiation(λ, T)

			λ: range of wavelengths [m]
			T: range of temperatures [K]

			B: spectral radiance [W·sr^-1·m^-3]

	Source: https://en.wikipedia.org/wiki/Planck%27s_law

"""
function blackbody_radiation(λ::AbstractArray{T1}, T::AbstractArray{T1}) where {T1<:Real}
	h::Float64 = 6.626070040e-34 # Planck's constant, J·s
	c::Float64 =  299792458.0 # speed of light, m·s
	kB::Float64 = 1.38064852e-23 # Boltzmann constant, J·K^-1
	R = Array{Float64,2}(undef,length(λ),length(T))
	num::Array = 2.0*h*c*c./λ.^5
	den::Array = h*c./(kB.*λ)
	for i in eachindex(T)
		R[:,i] .= num./expm1.(den./T[i])
	end
	return R
end

"""

	Returns an RGB matrix-color based on the wavelength input in nm.

		RBG = Utils.wavelength2rgb(λ)

			λ: range of wavelengths [nm]

	Wavelength range: 380nm and 780nm, black otherwise.

	Source: https://stackoverflow.com/questions/3407942/rgb-values-of-visible-spectrum

"""
function wavelength2rgb(λ::AbstractArray{<:Real})
    Γ::Float64 = 0.8
    # warm up
    R = Array{Float64, 1}(undef, length(λ))
    G = similar(R)
    B = similar(R)
    intensityCorrection = similar(R)
    R[380.0 .<= λ .<= 440.0] .= -(λ[380.0 .<= λ .<= 440.0] - 440.)./(440. - 380.)
    G[380.0 .<= λ .<= 440.0] .= 0.0
    B[380.0 .<= λ .<= 440.0] .= 1.0 .+ (0. * λ[380.0 .<= λ .<= 440.0])
    R[440.0 .< λ .<= 490.0] .= 0.0
    G[440.0 .< λ .<= 490.0] .= ((λ[440.0 .< λ .<= 490.0] - 440.)./(490. - 440.))
    B[440.0 .< λ .<= 490.0] .= 1.0
    R[490.0 .< λ .<= 510.0] .= 0.0
    G[490.0 .< λ .<= 510.0] .= 1.0
    B[490.0 .< λ .<= 510.0] .= (-(λ[490.0 .< λ .<= 510.0] - 510.)./(510. - 490.))
    R[510.0 .< λ .<= 580.0] .= ((λ[510.0 .< λ .<= 580.0] - 510.)./(580. - 510.))
    G[510.0 .< λ .<= 580.0] .= 1.0
    B[510.0 .< λ .<= 580.0] .= 0.0
    R[580.0 .< λ .<= 645.0] .= 1.0
    G[580.0 .< λ .<= 645.0] .= (-(λ[580.0 .< λ .<= 645.0] - 645.)./(645. - 580.))
    B[580.0 .< λ .<= 645.0] .= 0.0
    R[645.0 .< λ .<= 780.0] .= 1.0
    G[645.0 .< λ .<= 780.0] .= 0.0
    B[645.0 .< λ .<= 780.0] .= 0.0
    R[λ .> 780] .= 0.0
    R[λ .< 380] .= 0.0
    G[λ .> 780] .= 0.0
    G[λ .< 380] .= 0.0
    B[λ .> 780] .= 0.0
    B[λ .< 380] .= 0.0
    # LET THE INTENSITY SSS FALL OFF NEAR THE VISION LIMITS
    intensityCorrection[λ .> 700.0] .= 0.3 .+ 0.7.*(780.0 .- λ[λ .> 700.0])./(780. - 700.)
    intensityCorrection[λ .< 420.0] .= 0.3 .+ 0.7.*(λ[λ .< 420.0] - 380.0)./(420. - 380.)
    intensityCorrection[420.0 .< λ .< 700.0] .= 1.0.*λ[420.0 .< λ .< 700.0]
    # GAMMA ADJUST
    R .= (intensityCorrection.*R).^Γ
    G .= (intensityCorrection.*G).^Γ
    B .= (intensityCorrection.*B).^Γ
    # # Multiply by 255
    # R = R * 255
    # G = G * 255
    # B = B * 255
    return [R G B]
end

"""

	Smooth vector with simple moving average model.

		y = Utils.simple_moving_average(x; w=1)

			x: Vector to be smoothed out
			w: Integer with the window size

			y: Smoothed output
	Ref: https://en.wikipedia.org/wiki/Moving_average#Simple_moving_average_(boxcar_filter)

"""
function simple_moving_average(y::AbstractVector{T1}; w::T2=1) where {T1<:Real, T2<:Int64}
    _y = [y[1]*ones(w); y; y[end]*ones(w)]
    m = float(2*w + 1)
    w = Int(w + 1)
    y_ = similar(y)
    for i in w:length(_y)-(w-1)
        y_[i-w+1] = sum(_y[i-w+1:i+w-1])
    end
    return y_ ./ m
end

"""

	Smooth vector with exponential moving average model.

		y = Utils.exponential_moving_average(x; α=1.0)

			x: Vector to be smoothed out
			α: Weight

			y: Smoothed output
	Ref: https://en.wikipedia.org/wiki/Moving_average#Exponential_moving_average

"""
function exponential_moving_average(y::AbstractVector{T1}; α::T1=1.0) where {T1<:Real}
    y_ = similar(y)
    y_[1] = y[1]
    for i in 2:length(y)
        y_[i] = α*y[i] + (1.0 - α)*y_[i-1]
    end
    return y_
end

"""

    Convert Array{Array{Real,1},1} to an Array{Float64,1}.

"""
function flatten_arrays(xin::Array{Array{T1,1},1}) where {T1<:Real}
    xout = []
    # Determine the total number of elements of xin
    for i in eachindex(1:length(xin)), j in eachindex(1:length(xin[i]))
        push!(xout, xin[i][j])
    end
    return float.(xout)
end

"""

    Convert Array{Real,1} to Array{Array{Float64,1},1}.

"""
function array_arrays(x::Array{T1,1}, _x::Array{Array{T1,1},1}) where{T1<:Real, T2<:Real}
    xfinal = deepcopy(_x)
    ki::Int64 = 1
    kf::Int64 = 0
    for i in eachindex(collect(1:length(xfinal)))
        kf += length(xfinal[i])
        xfinal[i][1] = x[ki]
        xfinal[i][2:end] = x[ki+1:kf]
        ki += length(xfinal[i])
    end
    return float.(xfinal)
end

"""

    Return the polarisation averaged spectrum.

        X = Utils.average_polarisation(pol, Xp, Xs)

            pol: indicates the polarisation
                 (pol=1.0 => p, pol=0.0 => s, 0.0<=pol<=1.0 => average)
            Xp: p/TM polarisation quantity (e.g. Rp)
            Xs: s/TE polarisation quantity (e.g. Rs)

            Xavg: unpolarized quantity = pol*Xp + (1.0 - pol)*Xs

"""
function average_polarisation(pol::T1, Xp::Array{T1}, Xs::Array{T1}) where {T1<:Float64}
	X = @. pol*Xp + (1.0 - pol)*Xs
    return X
end

"""

	Multipeak Gaussian PDF curve.

		modelg = Utils.gaussian(x, p)

			x = vector with data of the x-axis
			p = [
				  [p0], # offset
				  [A1, μ1, σ1], # first peak
				  [A2, μ2, σ2], # second peak
				  ..., # so on
				]
				p0: offset of the total curve
				Ai: amplitude of peak i
				μi: position of peak i (mean or expectation of the distribution)
				σi: standard deviation of peak i
				i = 1,...,n

	Source: https://en.wikipedia.org/wiki/Normal_distribution

"""
function gaussian(x::AbstractArray{T1}, p::Array{Array{T1,1},1}) where {T1<:Real}
    isequal(length(flatten_arrays(p)[2:end]), 3*(length(p)-1)) || throw("The inputs are not correct. For the Gaussian curve, each peak has three parameters plus the offset.")
    y = ones(length(x)).*p[1][1] # offset
    for i = 2:length(p)
        @. y += p[i][1]*exp(-((x - p[i][2])/p[i][3])^2)
    end
    return y
end

"""

	Multipeak Cauchy-Lorentz PDF curve.

		modell = Utils.lorentzian(x, p)

			x = vector with data of the x-axis
			p = [
				  [p0], # offset
				  [A1, μ1, Γ1], # first peak
				  [A2, μ2, Γ2], # second peak
				  ..., # so on
				]
				p0: offset of the total curve
				Ai: amplitude of peak i
				μi: position of peak i (location parameter)
				Γi: half-width at half-maximum (HWHM) of peak i
				i = 1,...,n

	Source: https://en.wikipedia.org/wiki/Cauchy_distribution

"""
function lorentzian(x::AbstractArray{T1}, p::Array{Array{T1,1},1}) where {T1<:Real}
    isequal(length(flatten_arrays(p)[2:end]), 3*(length(p)-1)) || throw("The inputs are not correct. For the Lorentzian curve, each peak has three parameters plus the offset.")
    y = ones(length(x)).*p[1][1] # offset
    for i = 2:length(p)
        @. y += p[i][1]/(1.0 + ((x - p[i][2])/p[i][3])^2)
    end
    return y
end

"""

	Single peak Voigt PDF curve. The function is computed using the Faddeeva's
    function through the scaled complementary error function of x.

		modelv = Utils.voigtian(x, p)

			x = vector with data of the x-axis
			p = [
				  [p0], # offset
				  [A1, μ1, σ1, Γ1], # first peak
				  [A2, μ2, σ2, Γ2], # second peak
				  ..., # so on
				]
				p0: offset of the total curve
				Ai: amplitude of the curve
				μi: position of peak i (location parameter)
				σi: Gaussian standard deviation of peak i
				Γi: Lorentzian half-width at half-maximum (HWHM) of peak i
				i = 1,...,n

		The μ parameter has been included in a custom fashion to allow the
        distribution to be uncentered.

	Source: https://en.wikipedia.org/wiki/Voigt_profile

"""
function voigtian(x::AbstractArray{T1}, p::Array{Array{T1,1},1}) where {T1<:Real}
    isequal(length(flatten_arrays(p)[2:end]), 4*(length(p)-1)) || throw("The inputs are not correct. For the Voigtian curve, each peak has four parameters plus the offset.")
    y = ones(length(x)).*p[1][1] # offset
    for i = 2:length(p)
        @. y += p[i][1]*real(erfcx(-im*((x - p[i][2] + im*p[i][3])/sqrt(2)/p[i][4])))
    end
    return y
end

"""

	Build the input for oscillators models. Each index of refraction associated with
	Lorentz can be modelled with arbitrary numbers of oscillators. This function provides a
	way to construct a valid input for them.

	The input for the following models can be constructed are:
		RI.drudelorentz, RI.tauclorentz, RI.forouhibloomer,
		RI.lorentzplasmon, RI.codylorentz

		x = _oscillators_input(argin, n)

			argin: Array with two elements.
				   The first element is an array with the parameters not associated to
				   oscillators, while the second element contains an array with all the
				   parameters of the oscillators.
			n: number of parameters in each oscillator
			x: Array to input in RI modelled with Lorentz.

	For instance, the Tauc-Lorentz model accepts inputs with two parameters
	(length(argin[1]) == 2) not associated with the oscillator and three (n=3) forming
	each oscillator. Say you want to pass a three oscillator model input to the
	Tauc-Lorentz:
		argin = [
			[4.0, 3.5],  # ϵinf, Eg
			[15.0, 6.0, 1.05, # osc 1, A1, E01, Γ1
			 8.41, 4.5, 0.2, # osc 2, A2, E02, Γ2
			 1.5, 4.6, 0.1, # osc 3, A3, E03, Γ3
			],
		]
		x = _oscillators_input(argin, 3)
		x = [
			  [4.0, 3.5],  # ϵinf, Eg
			  [15.0, 6.0, 1.05], # osc 1, A1, E01, Γ1
			  [8.41, 4.5, 0.2], # osc 2, A2, E02, Γ2
			  [1.5, 4.6, 0.1], # osc 3, A3, E03, Γ3
		]
		RI.tauclorentz(x, λ)

	Notice: Careful with the number of parameteres in each model.

"""
## Build the input for oscillators models.
# Each element of x has the parameters for the corresponding oscillator.
# function _oscillators_input(argin, n::Int64)
#     x = argin[1]
#     numel::Int64 = Int(length(argin[2:end])/n + 1)
#     k::Int64 = 0
#     for i = 2:numel
#         push!(x, argin[i+k:n+i-1+k])
#         k += n-1
#     end
#     return x
# end
function _oscillators_input(argin, n::Int64)
    x = [argin[1]]
    numel::Int64 = Int(length(argin[2])/n)
    k::Int64 = 0
    for i = 1:numel
        push!(x, argin[2][i+k:n+i-1+k])
        k += n-1
    end
    return x
end

"""

	unfold_bnd(lb, ub)

Returns two arrays with the lower and upper bounds in accepted format for BBO package.

"""
function unfold_bnd(lb, ub)
    return (collect(zip(flatten_arrays(lb), flatten_arrays(ub))))
end

end # Utils


module ExtremalOptimization
using Random, StatsBase, LinearAlgebra

struct EOState{vP,vC,vO,tW,tN}
    P::Vector{vP}
    C::Vector{vC}
    order::Vector{vO}
    W::tW
    N::tN
end

struct EOProblem{tF,tS}
    f::tF
    s::tS
end

function approxequality(a, b, atol, rtol)
    Δ = norm(a - b)
    M = hypot(norm(a), norm(b))
    Δ <= max(atol, M * rtol)
end

function hasconverged(state::EOState, atol, rtol, f_atol, f_rtol)
    f_a = state.C[state.order[1]]
    f_b = state.C[state.order[2]]
    sol_a = state.P[state.order[1]]
    sol_b = state.P[state.order[2]]
    approxequality(f_a, f_b, f_atol, f_rtol) || approxequality(sol_a, sol_b, atol, rtol)
end

function eostate(problem::EOProblem, N; τ, kw...)
    P = [problem.s(i) for i = 1:N]
    C = [problem.f(P[i]) for i = 1:N]
    for i in eachindex(P)
        isfinite(C[i]) && continue
        error("The cost function has produced non-finite values for the initial point: $(P[i])")
    end
    order = sortperm(C)
    wS = [i^(-τ) for i = N:-1:1]
    wS[end] += wS[1]
    wS[1] = 0
    W = Weights(wS)
    best = order[1]
    P[best], C[best], EOState(P, C, order, W, N)
end

function sample_distinct(rng, N, set, notset)
    N <= 0 && return ()
    @label resample
    s = rand(rng, set)
    (s ∈ notset) && @goto resample
    return (s, sample_distinct(rng, N - 1, set, (notset..., s))...)
end

function eostate(problem::EOProblem, S::EOState; β, rng, verbose, kw...)
    i = S.order[sample(rng, S.W)]
    @label regen
    j, k, q = sample_distinct(rng, 3, 1:S.N, (i,))
    nP = S.P[j] + (β/sqrt(2)) * randn(rng) * (S.P[k] - S.P[q])
    nC = problem.f(nP)
    isfinite(nC) || begin
        verbose && println("Failed function evaluation at $nP")
        @goto regen
    end
    S.P[i] = nP
    S.C[i] = nC
    sortperm!(S.order, S.C)
    best = S.order[1]
    S.P[best], S.C[best], EOState(S.P, S.C, S.order, S.W, S.N)
end

"""
```julia
function optimize(
    f,
    s,
    N;
    reps_per_particle = 100,
    β = 1.5,
    A = 1.0,
    atol = 0.0,
    rtol = sqrt(eps(1.0)),
    f_atol = 0.0,
    f_rtol = sqrt(eps(1.0)),
    verbose = false,
    rng = Random.GLOBAL_RNG,
    callback = state -> nothing,
)
```

- `f` : cost function to minimize, whose argument is either a scalar or a vector, must returns a scalar value.
- `s` : function whose input is the particle number and output is a random initial point to be ranked by `f`.
- `N` : number of particles to use, choose a number greater than `d+4` where `d` is the number of dimensions.
- `reps_per_particle` : maximum number of iterations per particle.

## Usage example:

```julia
using ExtremalOptimization
rosenbrock2d(x) = (x[1] - 1)^2 + 100*(x[2] - x[1]^2)^2
initpoint(i) = randn(2)
optimize(rosenbrock2d, initpoint, 20)
```
output
```
(x = [1.0000008, 1.0000016], fx = 6.50e-13, f_nevals = 1109)
```
as expected the algorithm has found the optimum at `(1, 1)`, up to the specified tolerance.
"""
function optimize(
    f,
    s,
    N;
    reps_per_particle = 100,
    β = 1.5,
    A = 1.0,
    atol = 0.0,
    rtol = sqrt(eps(1.0)),
    f_atol = 0.0,
    f_rtol = sqrt(eps(1.0)),
    verbose = false,
    rng = Random.GLOBAL_RNG,
    callback = state -> nothing,
)
    problem = EOProblem(f, s)
    τ = 1 + A / log(N)
    best, cost, state = eostate(problem, N; τ)
    callback(state)
    fnevals = N
    for i = 1:(N*reps_per_particle)
        best, cost, state = eostate(problem, state; β, rng, verbose)
        fnevals += 1
        callback(state)
        verbose && println(i, " ", cost)
        hasconverged(state, atol, rtol, f_atol, f_rtol) && break
    end
    (x = best, fx = cost, f_nevals = fnevals)
end

export optimize
end


syms lon lat r
syms x1 y1 z1
syms xx1 yy1 zz1
syms lam1 theta1 r1

% xx1=x1/sqrt(x1^2+y1^2+z1^2);
% yy1=y1/sqrt(x1^2+y1^2+z1^2);
% zz1=z1/sqrt(x1^2+y1^2+z1^2);

% xx1=r*cos(lam)*cos(theta)
% yy1=r*sin(lam)*cos(theta)
% zz1=r*sin(theta);
% yy1/xx1=sin(lam)/cos(lam)
% lam=atan(yy1/xx1)
% theta=asin(zz1/sqrt(xx1^2+yy1^2+zz1^2))
% r=sqrt(xx1^2+yy1^2+zz1^2);
%
% d atan(s) / ds = 1/(1+s^2)
% d asin(s) / ds = 1/sqrt(1-s^2)

lam=atan(yy1/xx1);
theta=asin(zz1);

D1=[diff(lam,xx1) diff(lam,yy1) diff(lam,zz1)
    diff(theta,xx1) diff(theta,yy1) diff(theta,zz1)];
D11=subs(D1,{xx1,yy1,zz1},{cos(lam1)*cos(theta1),sin(lam1)*cos(theta1),sin(theta1)});
D111=simplify(D11);

C=[cos(theta1) 0
   0           1 ];
D1111=C*D111;

lam=atan(yy1/xx1);
theta=asin(zz1/sqrt(xx1^2+yy1^2+zz1^2));
r=sqrt(xx1^2+yy1^2+zz1^2);
E1=[diff(lam,xx1) diff(lam,yy1) diff(lam,zz1)
    diff(theta,xx1) diff(theta,yy1) diff(theta,zz1)
    diff(r,xx1) diff(r,yy1) diff(r,zz1)];
E11=subs(E1,{xx1,yy1,zz1},{r1*cos(lam1)*cos(theta1),r1*sin(lam1)*cos(theta1),r1*sin(theta1)});
E111=simplify(E11);

C=[r1*cos(theta1) 0  0
   0              r1 0
   0              0  1];

E1111=C*E111;

% [-sin(lam1)                cos(lam1)                     0
%  -cos(lam1)*sin(theta1)   -sin(lam1)*sin(theta1)    cos(theta1)
%   cos(lam1)*cos(theta1)    sin(lam1)*cos(theta1)    sin(theta1)]

A(1,1)=diff(xx1,x1);
A(1,2)=diff(xx1,y1);
A(1,3)=diff(xx1,z1);
A(2,1)=diff(yy1,x1);
A(2,2)=diff(yy1,y1);
A(2,3)=diff(yy1,z1);
A(3,1)=diff(zz1,x1);
A(3,2)=diff(zz1,y1);
A(3,3)=diff(zz1,z1);

AA=simplify(subs(A,{x1,y1,z1},{r*cos(lon)*cos(lat),r*sin(lon)*cos(lat),r*sin(lat)}));
B=[r*cos(lat) 0 0
  0         r 0
  0         0 r];
C=simplify(B*AA);
bb=77;


Base.:+(x::ExactReal, y::ExactReal) = ExactReal(n -> (x(n+2) + y(n+2) + 2) >> 2)

Base.:-(x::ExactReal) = ExactReal(n->-x(n))
Base.:-(x::ExactReal, y::ExactReal) = x + (-y)
Base.:-(x::Union{Rational, Integer}, y::ExactReal) = ExactReal(x) - y

"Most significant digit"
function msd(x::ExactReal)
    x.msd > -1 && return x.msd

    n = 0

    while x(n) == 0
        n += 1
    end

    x.msd = n

    return n
end

function Base.:*(x::ExactReal, y::ExactReal)
    return ExactReal(n -> begin
        p = max(n - msd(y) + 4, n + 3 >> 1)
        q = max(n - msd(x) + 4, n + 3 >> 1)

        (1 + x(p) * y(q)) >> (p + q - n)
    end
    )
end

Base.:*(x::Rational, y::ExactReal) = ExactReal(x) * y

Base.sqrt(x::ExactReal) = ExactReal(n -> isqrt(x(2n)))

"Returns -1 if x < y and +1 if x > y"
function compare(x::ExactReal, y::ExactReal)
    n = 0

    while true
        xn, yn = x(n), y(n)
        xn < yn - 1 && return -1  # x < y
        xn > yn + 1 && return 1   # x > y

        n += 1
    end
end

Base.:<(x::ExactReal, y::ExactReal) = compare(x, y) == -1


export M4, M5, M6

const σ4 = (2/3, 10/7/π, 1/π)
const σ5 = (1/24, 96/1199/π, 1/20/π)
const σ6 = (1/120, 7/478/π, 1/120/π)

const M4 = σ4 .* @piecewise_polynomial begin
    0.0 => p"0.0"
    1.0 => p"0.25*(2-q)^3 - (1 - q)^3"
    2.0 => p"0.25*(2-q)^3"
    _   => p"0.0"
end

const M5 = σ5 .* @piecewise_polynomial begin
    0.0 => p"0.0"
    0.5 => p"(2.5 - q)^4 - 5*(1.5 - q)^4 + 10*(0.5 - q)^4"
    1.5 => p"(2.5 - q)^4 - 5*(1.5 - q)^4"
    2.5 => p"(2.5 - q)^4"
    _   => p"0.0"
end

const M6 = σ6 .* @piecewise_polynomial begin
    0.0 => p"0.0"
    1.0 => p"(3 - q)^5 - 6*(2 - q)^5 + 15*(1 - q)^5"
    2.0 => p"(3 - q)^5 - 6*(2 - q)^5"
    3.0 => p"(3 - q)^5"
    _   => p"0.0"
end


# Fractional Brownian motion
immutable FBM <: ContinuousUnivariateStochasticProcess
  t::Vector{Float64}
  n::Int64
  h::Float64 # Hurst index

  function FBM(t::Vector{Float64}, n::Int64, h::Float64)
    t[1] == 0. || error("Starting time point must be equal to 0.0.")
    issorted(t, lt=<=) || error("The time points must be strictly sorted.")
    int64(length(t)) == n || error("Number of time points must be equal to the vector holding the time points.")
    0. < h < 1. || error("Hurst index must be between 0 and 1.")
    new(t, n, h)
  end
end

FBM(t::Vector{Float64}, h::Float64) = FBM(t, int64(length(t)), h)
FBM(t::Ranges, h::Float64) = FBM(collect(t), int64(length(t)), h)
FBM(t::Float64, n::Int64, h::Float64) = FBM(t/n:t/n:t, h)
FBM(t::Float64, h::Float64) = FBM([t], 1, h)

FBM(t::Matrix{Float64}, h::Float64) = FBM[FBM(t[:, i], h) for i = 1:size(t, 2)]
FBM(t::Ranges, np::Int, h::Float64) = FBM[FBM(t, h) for i = 1:np]
FBM(t::Float64, n::Int64, np::Int, h::Float64) = FBM[FBM(t, n, h) for i = 1:np]

# Fractional Gaussian noise
immutable FGN <: ContinuousUnivariateStochasticProcess
  σ::Float64
  h::Float64 # Hurst index

  function FGN(σ::Float64, h::Float64)
    σ > 0. || error("Standard deviation must be positive.")
    0. < h < 1. || error("Hurst index must be between 0 and 1.")
    new(σ, h)
  end
end

FGN(h::Float64) = FGN(1., h)

convert(::Type{FBM}, p::BrownianMotion) = FBM(p.t, p.n, 0.5)
convert(::Type{FGN}, p::BrownianMotion) = FGN((p.t[end]/(p.n-1))^0.5, 0.5)
convert(::Type{FGN}, p::FBM) = FGN((p.t[end]/(p.n-1))^p.h, p.h)

function autocov(p::FGN, i::Int64, j::Int64)
  twoh::Float64 = 2*p.h
  0.5*abs2(p.σ)*(abs(j-i+1)^twoh+abs(j-i-1)^twoh-2*abs(j-i)^twoh)
end

function autocov!(c::Matrix{Float64}, p::FGN)
  n::Int64 = size(c, 1)

  for i = 1:n
    for j = 1:i
      c[i, j] = autocov(p, i, j)
    end
  end

  for i = 1:n
    for j = (i+1):n
      c[i, j] = c[j, i]
    end
  end

  c
end

autocov(p::FGN, maxlag::Int64) = autocov!(Array(Float64, maxlag, maxlag), p)

function autocov!(y::Vector{Float64}, p::FGN, lags::IntegerVector)
  nlags = length(lags)
  twoh::Float64 = 2*p.h

  for i = 1:nlags
    y[i] = 0.5*abs2(p.σ)*(abs(lags[i]+1)^twoh+abs(lags[i]-1)^twoh-2*abs(lags[i])^twoh)
  end

  y
end

autocov(p::FGN, lags::IntegerVector) = autocov!(Array(Float64, length(lags)), p, lags)

function autocov(p::FBM, i::Int64, j::Int64)
  twoh::Float64 = 2*p.h
  0.5*((p.t[i])^twoh+(p.t[j])^twoh-abs(p.t[i]-p.t[j])^twoh)
end

function autocov!(c::Matrix{Float64}, p::FBM)
  n::Int64 = p.n-1

  for i = 1:n
    for j = 1:i
      c[i, j] = autocov(p, i+1, j+1)
    end
  end

  for i = 1:n
    for j = (i+1):n
      c[i, j] = c[j, i]
    end
  end

  c
end

function autocov(p::FBM)
  n::Int64 = p.n-1
  autocov!(Array(Float64, n, n), p)
end

### rand_chol generates FBM using the method based on Cholesky decomposition.
### T. Dieker, Simulation of Fractional Brownian Motion, master thesis, 2004.
### The complexity of the algorithm is O(n^3), where n is the number of FBM samples.
function rand_chol(p::FBM; fbm::Bool=true)
  w = chol(autocov(p), :L)*randn(p.n-1)

  # If fbm is true return FBM, otherwise return FGN
  w = [0., w]
  if !fbm
    w = diff(w)
  end

  w
end

function rand_chol(p::Vector{FBM}; fbm::Bool=true)
  n::Int64 = fbm ? p[1].n : p[1].n-1
  np = length(p)
  x = Array(Float64, n, np)

  for i = 1:np
    x[:, i] = rand_chol(p[i], fbm=fbm)
  end

  x
end

### rand_fft generates FBM using fast Fourier transform (FFT).
### The time interval of FBM is [0, 1] with a stepsize of 2^p, where p is a natural number.
### The algorithm is known as the Davies-Harte method or the method of circular embedding.
### R.B. Davies and D.S. Harte, Tests for Hurst Effect, Biometrika, 74 (1), 1987, pp. 95-102.
### For a more recent publication, see P.F. Craigmile, Simulating a Class of Stationary Gaussian Processes Using the
### Davies–Harte Algorithm, With Application to Long Memory Processes, Journal of Time Series Analysis, 24 (5), 2003,
### pp. 505-511.
### The complexity of the algorithm is O(n*log(n)), where n=2^p is the number of FBM samples.
function rand_fft(p::FBM; fbm::Bool=true)
  # Determine number of points of simulated FBM
  pnmone::Int64 = p.n-1
  n::Int64 = 2^ceil(log2(pnmone))

  # Compute autocovariance sequence of underlying FGN
  c = Array(Float64, n+1)
  autocov!(c, FGN((1/n)^p.h, p.h), 0:n)

  # Compute square root of eigenvalues of circular autocovariance sequence
  l = real(fft([c, c[end-1:-1:2]]))
  all(i->(i>0), l) || error("Non-positive eigenvalues encountered.")
  lsqrt = sqrt(l)

  # Simulate standard random normal variables
  twon::Int64 = 2*n
  z = randn(twon)

  # Generate fractional Gaussian noise (retain only the first p.n-1 values)
  x = sqrt(0.5)*lsqrt[2:n].*complex(z[2*(2:n)-2], z[2*(2:n)-1])
  w = real(bfft([lsqrt[1]*z[1], x, lsqrt[n+1]*z[twon], conj(reverse(x))]))[1:pnmone]/sqrt(twon)

  # If fbm is true return FBM, otherwise return FGN
  if fbm
    w = [0., cumsum(w)]
  end

  w
end

function rand_fft(p::Vector{FBM}; fbm::Bool=true)
  n::Int64 = fbm ? p[1].n : p[1].n-1
  np = length(p)
  x = Array(Float64, n, np)

  for i = 1:np
    x[:, i] = rand_fft(p[i], fbm=fbm)
  end

  x
end

function rand(p::FBM; fbm::Bool=true, rtype::Symbol=:fft)
  if rtype == :chol
    rand_chol(p; fbm=fbm)
  elseif rtype == :fft
    rand_fft(p; fbm=fbm)
  else
  end
end

function rand(p::Vector{FBM}; fbm::Bool=true, rtype::Symbol=:fft)
  if rtype == :chol
    rand_chol(p; fbm=fbm)
  elseif rtype == :fft
    rand_fft(p; fbm=fbm)
  else
  end
end


@test isregular(rand(4))
@test isregular(rand(4,4,4,4))
@test !isregular(rand(4,5))

@test mirror([1]) == [1]
@test mirror([2,3]) == [2,-3]
@test mirror([2,3,4]) == [2,-3,4]
@test mirror(@SVector([2,3,4])) == @SVector [2,-3,4]
@test mirror([4.9,5,6,7]) == [4.9,-5,6,-7]

@test upsample([1]) == [1,0]
@test upsample([1]; skipfirst=true) == [0,1]
@test upsample([1,2]) == [1,0,2,0]
@test upsample([1,2]; skipfirst=true) == [0,1,0,2]
@test downsample([1,2]) == [1]
@test downsample([1,2]; skipfirst=true) == [2]
@test downsample([1,2,3,4]) == [1,3]
@test downsample([1,2,3,4]; skipfirst=true) == [2,4]

signalsize = 64

@test_throws ArgumentError coefficientcount(ones(signalsize); level=0)
@test coefficientcount(ones(signalsize)) == signalsize
@test coefficientcount(rand(signalsize); threshold=1.01) == 0
@test coefficientcount(rand(signalsize).-2.5; threshold=1.01) == signalsize
@test coefficientcount([0,2,6,7,8]; threshold=6.6) == 2
@test coefficientcount([10,-11,6,7,8.0,-5,-8,0]; threshold=3, level=4) == 3
@test coefficientcount(ones(signalsize,signalsize)) == signalsize * signalsize
@test coefficientcount(rand(signalsize,signalsize); threshold=1.01) == 0
@test coefficientcount([-1 2;3 4]; threshold=2.5) == 2

signal = [1,2]
output = [1,2]
upto = 2
@test partition(signal) == output
@test partition(signal; upto) == output
@test combine(signal) == output
@test combine(signal; upto) == output

signal = [-1,2,3.3,4]
output = [-1,3.3,2,4]
upto = 4
@test partition(signal) == output
@test partition(signal; upto) == output

signal = [-1,3.3,2,4]
output = [-1,2,3.3,4]
@test combine(signal) == output
@test combine(signal; upto) == output

signal = randn(128)
@test signal == combine(partition(signal))

# with strides
signal = [1,2,3,4,5,6,7,8,9,10,11,12,13,14]

stride(s, i, l) = s:i:(s + (l - 1) * i)

@test partition(signal, stride(3, 2, 2))[1:2] == [3,5]
@test partition(signal, stride(3, 2, 4))[1:4] == [3,7,5,9]
@test partition(signal, stride(1, 3, 4))[1:4] == [1,7,4,10]
@test partition(signal, stride(1, 1, 8)) ==  partition(signal; upto=8)

range = stride(3, 2, 2)
@test combine(signal, range)[range] == [1,2]

range = stride(3, 2, 4)
@test combine(signal, range)[range] == [1,3,2,4]

range = stride(1, 3, 4)
@test combine(signal, range)[range] == [1,3,2,4]
@test combine(signal, stride(1, 1, 8)) ==  combine(signal; upto=8)


@show filter(iseven, 1:10)


export quick_sort!

function quick_sort!(arr::Vector, h::Int64, t::Int64)
    len = t-h+1
    if len < 2
        return arr
    end

    n, i = h, h
    pivot = arr[t]
    for n = h:t
        if arr[n] <= pivot
            arr[i], arr[n] = arr[n], arr[i]
            i += 1
        end
    end

    if i > t
        i = t
    end

    quick_sort!(arr, h, i-1)
    quick_sort!(arr, i, t)

    return arr
end


@testset "$TEST $G" begin

r1 = G(10,20)
r2 = DG(5,10)
@test typeof(r1) == G
@test typeof(r2) == DG
@test nv(r1) == 10
@test ne(r1) == 20
@test nv(r2) == 5
@test ne(r2) == 10

@test G(10,20,seed=3) == G(10, 20, seed=3)
@test DG(10,20,seed=3) == DG(10, 20, seed=3)
@test G(10,20,seed=3) == erdos_renyi(10, 20, G, seed=3)
@test ne(G(10,40,seed=3)) == 40
@test ne(DG(10,80,seed=3)) == 80

er = erdos_renyi(10, 0.5, G)
@test nv(er) == 10
@test ne(er) < 40
@test is_directed(er) == false

er = erdos_renyi(10, 0.5, DG)
@test nv(er) == 10
@test is_directed(er) == true

er = erdos_renyi(10, 0.5, G, seed=17)
@test nv(er) == 10
@test is_directed(er) == false

ws = watts_strogatz(10, 4, 0.2,  G)
@test nv(ws) == 10
@test ne(ws) == 20
@test is_directed(ws) == false

ws = watts_strogatz(10, 4, 0.2, DG)
@test nv(ws) == 10
@test ne(ws) == 20
@test is_directed(ws) == true

ba = barabasi_albert(10, 2,  G)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == false

ba = barabasi_albert(10, 2, 2,  G)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == false

ba = barabasi_albert(10, 4, 2, G)
@test nv(ba) == 10
@test ne(ba) == 12
@test is_directed(ba) == false

ba = barabasi_albert!(CompleteGraph(2, G), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 17
@test is_directed(ba) == false

ba = barabasi_albert!(CompleteGraph(2, G), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 17
@test is_directed(ba) == false

ba = barabasi_albert!(CompleteGraph(4, G), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 18
@test is_directed(ba) == false

ba = barabasi_albert(10, 2, DG)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == true

ba = barabasi_albert(10, 2, 2, DG)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == true

ba = barabasi_albert(10, 4, 2, DG)
@test nv(ba) == 10
@test ne(ba) == 12
@test is_directed(ba) == true

ba = barabasi_albert!(CompleteDiGraph(2, DG), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 18
@test is_directed(ba) == true

ba = barabasi_albert!(CompleteDiGraph(2, DG), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 18
@test is_directed(ba) == true

ba = barabasi_albert!(CompleteDiGraph(4, DG), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 24
@test is_directed(ba) == true

fm = static_fitness_model(20, rand(10), G)
@test nv(fm) == 10
@test ne(fm) == 20
@test is_directed(fm) == false

fm = static_fitness_model(20, rand(10), rand(10), DG)
@test nv(fm) == 10
@test ne(fm) == 20
@test is_directed(fm) == true

sf = static_scale_free(10, 20, 2.0, G)
@test nv(sf) == 10
@test ne(sf) == 20
@test is_directed(sf) == false

sf = static_scale_free(10, 20, 2.0, 2.0, DG)
@test nv(sf) == 10
@test ne(sf) == 20
@test is_directed(sf) == true

rr = random_regular_graph(5, 0, G)
@test nv(rr) == 5
@test ne(rr) == 0
@test is_directed(rr) == false

rd = random_regular_digraph(10,0)
@test nv(rd) == 10
@test ne(rd) == 0
@test is_directed(rd)

rr = random_regular_graph(6, 3, G, seed=1)
@test nv(rr) == 6
@test ne(rr) == 9
@test is_directed(rr) == false

rr = random_regular_graph(100, 5, G)
@test nv(rr) == 100
@test ne(rr) == 250
@test is_directed(rr) == false
for v in vertices(rr)
    @test degree(rr, v) == 5
end

rr = random_configuration_model(10, repeat([2,4] ,5), G, seed=3)
@test nv(rr) == 10
@test ne(rr) == 15
@test is_directed(rr) == false
num2 = 0; num4 = 0
for v in vertices(rr)
    d = degree(rr, v)
    @test  d == 2 || d == 4
    d == 2 ? num2 += 1 : num4 += 1
end
@test num4 == 5
@test num2 == 5

rr = random_configuration_model(100, zeros(Int,100), G)
@test nv(rr) == 100
@test ne(rr) == 0
@test is_directed(rr) == false

rr = random_configuration_model(3, [2,2,2], G, check_graphical=true)
@test nv(rr) == 3
@test ne(rr) == 3
@test is_directed(rr) == false

rd = random_regular_digraph(100, 4, DG)
@test nv(rd) == 100
@test ne(rd) == 400
@test is_directed(rd)
@test std(out_degree(rd)) == 0

rd = random_regular_digraph(100, 4, DG, dir=:in)
@test nv(rd) == 100
@test ne(rd) == 400
@test is_directed(rd)
@test std(in_degree(rd)) == 0

rr = random_regular_graph(10, 8, G, seed=4)
@test nv(rr) == 10
@test ne(rr) == 40
@test is_directed(rr) == false
for v in vertices(rr)
    @test degree(rr, v) == 8
end

rd = random_regular_digraph(10, 8, DG, dir=:out, seed=4)
@test nv(rd) == 10
@test ne(rd) == 80
@test is_directed(rd)

g = stochastic_block_model(2., 3., [100,100], G)
@test  3.5 < mean(degree(g)) < 6.5
g = stochastic_block_model(3., 4., [100,100,100])
@test  9.5 < mean(degree(g)) < 12.5

end # testset


# author: Tommaso Sartor

include("./riccati.jl")
using .riccati

# main script
length(ARGS) != 2 && error("Usage: julia test_riccati.jl <nmass> <nrep>")

nmass = parse(Int, ARGS[1])
nrep = parse(Int, ARGS[2])

nx = 2*nmass
nu = nmass
N = 5

# data

# mass spring system
Ts = 0.5

Ac = zeros(nx, nx)
for ii in 1:nmass
	Ac[ii, nmass+ii] = 1.0
end
for ii in 1:nmass
	Ac[nmass+ii, ii] = -2.0
end
for ii in 1:nmass-1
	Ac[nmass+ii+1, ii] = 1.0
end
for ii in 1:nmass-1
	Ac[nmass+ii, ii+1] = 1.0
end

Bc = zeros(nx, nu)
for ii in 1:nu
	Bc[nmass+ii, ii] = 1.0
end

#display(Ac)
#println()
#display(Bc)
#println()

MM = [ Ts*Ac Ts*Bc; zeros(nu, nx+nu) ]

#display(MM)
#println()

#MM = exp( MM )
MM = randn(nu+nx, nu+nx)

#display(MM)
#println()

A = MM[1:nx,1:nx]
B = MM[1:nx, nx+1:end]

#display(A)
#println()
#display(B)
#println()

Q = zeros(nx, nx)
for ii in 1:nx
	Q[ii, ii] = 1.0
end

R = zeros(nu, nu)
for ii in 1:nu
	R[ii, ii] = 2.0
end

#display(Q)
#println()
#display(R)
#println()

x0 = zeros(nx, 1)
x0[1] = 3.5
x0[2] = 3.5

#display(x0)
#println()

# work matrices

BAt = [transpose(B); transpose(A)]
#display(BAt)
#println()
RSQ = [R zeros(nu,nx); zeros(nx,nu) Q]

BAtL = zeros(nu+nx, nx)
L = zeros(nu+nx, nu+nx, N)
LN = zeros(nx, nx)
M = zeros(nu+nx, nu+nx)

# riccati recursion, square root algorithm

for rep in 1:nrep
	riccati_trf(N, nx, nu, BAt, RSQ, L, LN, BAtL, M)
end

time_start = time_ns()

for rep in 1:nrep
	riccati_trf(N, nx, nu, BAt, RSQ, L, LN, BAtL, M)
end

time_end = time_ns()

println((time_end-time_start)/nrep*1e-9)

#display(L)
#println()
#display(LN)
#println()


module CartPole

using MDPs
using DelimitedFiles
using PGFPlots

export state_space, action_space, transition, reward, simulate,
       viz_policy, viz_trajectory
export THETAMIN, THETAMAX, THETADOTMIN, THETADOTMAX, XMIN, XMAX, XDOTMIN, XDOTMAX, T

# mdp
const MTHETA = 10
const MTHETADOT = 10
const MX = 10
const MXDOT = 10
const N = 1000

# state
const THETAMIN = float(-pi / 2)
const THETAMAX = float(pi / 2)
const THETADOTMIN = -3.0
const THETADOTMAX = 3.0
const XMIN = -2.4
const XMAX = 2.4
const XDOTMIN = -3.5
const XDOTMAX = 3.5

# action
const AMIN = -10
const AMAX = 10

# transition
const GRAVITY = 9.8
const MASSCART = 1.0
const MASSPOLE = 0.1
const LENGTH = 0.5
# const FORCE_MAG = 10
const DT = 0.1  # time intervals

# reward
const K = 15

# test
const T = 200

# visualization
const EPS = 1e-4

function state_space(mtheta::Int64=MTHETA, mthetadot::Int64=MTHETADOT, mx::Int64=MX, mxdot::Int64=MXDOT)
    theta = range(THETAMIN, stop=THETAMAX, length=mtheta)
    thetadot = range(THETADOTMIN, stop=THETADOTMAX, length=mthetadot)
    x = range(XMIN, stop=XMAX, length=mx)
    xdot = range(XDOTMIN, stop=XDOTMAX, length=mxdot)
    return RectangleGrid(theta, thetadot, x, xdot)
end

function action_space(n::Int64=N)
    return range(AMIN, stop=AMAX, length=n)
end

function wrap_around(angle::Float64)
    while angle < -pi / 2
        angle = -pi / 2
    end
    while angle > pi / 2
        angle = pi / 2
    end
    return angle
end

function clip(val::Float64, minval::Float64, maxval::Float64)
    return min(maxval, max(minval, val))
end

function transition(state::Vector{Float64}, action::Float64)
    temp = (action + MASSPOLE * LENGTH * state[2] * state[2] * sin(state[1])) / (MASSPOLE + MASSCART)
    theta_acc = (GRAVITY * sin(state[1]) - cos(state[1]) * temp) / (LENGTH * (4/3 - MASSPOLE * cos(state[1]) * cos(state[1]) / (MASSPOLE + MASSCART)))
    x_acc = temp - MASSPOLE * LENGTH * theta_acc * cos(state[1]) / (MASSPOLE + MASSCART)

    snext = zeros(length(state))
    snext[1] = wrap_around(state[1] + DT * state[2])
    snext[2] = clip(state[2] + DT * theta_acc, THETADOTMIN, THETADOTMAX)
    snext[3] = clip(state[3] + DT * state[4], XMIN, XMAX)
    snext[4] = clip(state[4] + DT * x_acc, XDOTMIN, XDOTMAX)
    return snext
end

function reward(state::Vector{Float64}, action::Float64)
    return (cos(K * state[1]))^4
end

function simulate(mdp::MDP, policy::Policy, state::Vector{Float64})
    trajectory = zeros(T, dimensions(mdp.S))
    actions = zeros(T)
    for t = 1:T
        trajectory[t, :] = state
        if state[1] == 0.0
            trajectory = trajectory[1:t, :]
            actions = actions[1:t - 1]
            break
        end
        action = policy!(policy, get_belief(mdp, state))
        state = mdp.transition(state, action)
        actions[t] = action
    end
    return trajectory, actions
end

function viz_policy(mdp::MDP, policy::Policy, title::String="Policy heatmap", saveplot::Bool=false, name::String="")
    function getmap(p::Float64, v::Float64)
        return policy!(policy, get_belief(mdp, [p, v]))
    end

    p = Axis([Plots.Image(getmap, (THETAMIN + EPS, THETAMAX - EPS),
                      (THETADOTMIN + EPS, THETADOTMAX - EPS),
                      xbins = 250, ybins = 250,
                      colormap = ColorMaps.Named("Blues"))],
          width="14cm", height="14cm",
          xlabel="angle", ylabel="angular speed",
          title=title)
    p
    if saveplot
      PGFPlots.save(name, p)
    end
end

function rad2deg_vec(trajectory::Vector{Float64})
    l = length(trajectory)
    out = zeros(l)
    for i = 1:l
        out[i] = rad2deg(trajectory[i])
    end
    return out
end

function viz_trajectory(trajectory::Matrix{Float64}, actions::Vector{Float64},
                        title1::String="angular position over time", title2::String="control input over time",
                        saveplot::Bool=false, name::String="")
    g = GroupPlot(2, 1, groupStyle="horizontal sep=1.5cm")
    push!(g, Axis([Plots.Linear(rad2deg_vec(trajectory[:, 1]))],
                  width="12cm", height="12cm",
                  xlabel="time", ylabel="angle (deg)",
                  title=title1))
    push!(g, Axis([Plots.Linear(actions)],
                  width="12cm", height="12cm",
                  xlabel="time", ylabel="input",
                  title=title2))
    g
    if saveplot
      PGFPlots.save(name, g)
    end
end

end


## Exercise 8-5
## Encapsulate this code in a function named count, and generalize it so that it accepts the string and the letter as arguments.
println("Ans 1: ")

function count(word::AbstractString, to_count::Char)
    counter = 0
    for letter in word
        if letter == to_count
            counter = counter + 1
        end
    end
    counter
end

println(count("banana", 'a'))

## Then rewrite the function so that instead of traversing the string, it uses the three-parameter version of find from the previous section.
println("Ans 2: ")

function find(word, letter, start::Int = 1)
    word = word[start : end]
    index = firstindex(word)
    while index <= sizeof(word)
        if word[index] == letter
            return index + start - 1
        end
        index = nextind(word, index)
    end
    -1
end

function count_find(word::AbstractString, letter::Char)
    counter = 0
    last_ind = 1
    while true
        ind = find(word, letter, last_ind)
        if ind > 0
            counter += 1
            last_ind = ind + 1
        end

        if last_ind > length(word) || ind == -1
            break;
        end
    end
    counter
end

println(count_find("bananana", 'a'))

println("End.")


struct BudgetedBipartiteMatchingInstance{T, U} <: CombinatorialInstance
  matching::BipartiteMatchingInstance{T}
  weight::Dict{Edge{T}, U}
  budget::U

  function BudgetedBipartiteMatchingInstance(graph::AbstractGraph{T}, reward::Dict{Edge{T}, Float64}, weight::Dict{Edge{T}, U}, budget::U) where {T, U}
    matching = BipartiteMatchingInstance(graph, reward)
    return new{T, U}(matching, weight, budget)
  end
end

function _budgeted_bipartite_matching_compute_value(i::BudgetedBipartiteMatchingInstance{T, U}, solution::Vector{Edge{T}}) where {T, U}
  return _budgeted_bipartite_matching_compute_value(i.matching, solution)
end

function _budgeted_bipartite_matching_compute_value(i::BipartiteMatchingInstance{T}, solution::Vector{Edge{T}}) where T
  return sum(i.reward[(e in keys(i.reward)) ? e : reverse(e)] for e in solution)
end

function _budgeted_bipartite_matching_compute_weight(i::BudgetedBipartiteMatchingInstance{T, U}, solution::Vector{Edge{T}}) where {T, U}
  return sum(i.weight[(e in keys(i.weight)) ? e : reverse(e)] for e in solution)
end

abstract type BudgetedBipartiteMatchingSolution{T, U} <: CombinatorialSolution
  # instance::BudgetedBipartiteMatchingInstance{T, U}
  # solution::Vector{Edge{T}}
  # value::Float64 # TODO: remove me, only useful for Lagrangian.
end


#!/usr/bin/env julia

module chunk_u

#=
   Chunking Utilities - chunk_u.jl
=#

export fetch_chunk
export bound_shape, bounds, zip_bounds, intersect_bounds
export chunk_bounds, chunk_bounds2D

export vol_shape

"""

    vol_shape( bounds )

  Returns the overall shape of a subvolume which ranges
  between the two coordinates contained within the Pair bounds.
"""
function bound_shape( bounds::Pair )
  bounds.second - bounds.first + 1
end

@inline function zip_bounds( bounds::Pair )

  #need Any for possible Colon()s
  bz = Vector{Any}(length(bounds.first))

  for i in eachindex(bounds.first)
    bz[i] = bounds.first[i]:bounds.second[i]
  end

  bz
end

"""

    fetch_chunk( d, bounds::Pair, offset )

  Generalized function to fetch a 3d chunk
  from within an Nd Array, H5Dataset, etc.
"""
function fetch_chunk( d, bounds::Pair, offset=[0,0,0] )

  shifted = (bounds.first  + offset) => (bounds.second + offset);
  zipped = zip_bounds(shifted)

  while length(zipped) < length(size(d)) push!(zipped,Colon()) end

  d[zipped...]
end

"""
    chunk_bounds( vol_size, chunk_size )

  Computes the index bounds for splitting a volume
  into chunks of a maximum size. If the dimensions
  don't match to allow for perfect splitting, smaller
  are included at the end of each dimension. The offset coordinate
  is added to each resulting bound.
"""
function chunk_bounds( vol_size, chunk_size, offset=[0,0,0] )

  x_bounds = bounds1D( vol_size[1], chunk_size[1] )
  y_bounds = bounds1D( vol_size[2], chunk_size[2] )
  z_bounds = bounds1D( vol_size[3], chunk_size[3] )

  num_bounds = prod(
    (length(x_bounds), length(y_bounds), length(z_bounds))
    )
  bounds = Vector{Pair{Vector{Int},Vector{Int}}
                 }(num_bounds)
  i=0;

  for z in z_bounds
    for y in y_bounds
      for x in x_bounds
        i+=1;
        bounds[i] = [x.first,  y.first,  z.first] + offset =>
                    [x.second, y.second, z.second] + offset
      end
    end
  end

  bounds
end

"""

    chunk_bounds2D( vol_size, chunk_size, offset=[0,0] )

  Computes the index bounds for splitting a 2D image
  into chunks of a maximum size. If the dimensions
  don't match to allow for perfect splitting, smaller chunks
  are included at the end of each dimension. The offset coordinate
  is added to each resulting bound.
"""
function chunk_bounds2D( vol_size, chunk_size, offset=[0,0] )

  x_bounds = bounds1D( vol_size[1], chunk_size[1] )
  y_bounds = bounds1D( vol_size[2], chunk_size[2] )

  num_bounds = prod(
    (length(x_bounds), length(y_bounds))
  )

  bounds = Vector{Pair{Vector{Int},Vector{Int}}
                }(num_bounds)
  i=0

  for y in y_bounds
    for x in x_bounds
      i+=1;
      bounds[i] = [x.first,  y.first] + offset =>
                  [x.second, y.second] + offset
    end
  end

  bounds
end

"""

    bounds1D( full_width, step_size )

  Returns the 1D index bounds for splitting an interval
  of a given width into increments of a given step size.
  Includes smaller increments at the end if the interval
  isn't evenly divided
"""
function bounds1D( full_width, step_size )

  @assert step_size > 0
  @assert full_width > 0

  start = 1
  ending = step_size

  num_bounds = convert(Int,ceil(full_width / step_size));
  bounds = Vector{Pair{Int,Int}}( num_bounds )

  i=1;
  while ending < full_width
    bounds[i] = start => ending

    i += 1;
    start  += step_size
    ending += step_size
  end

  #last window for remainder
  bounds[end] = start => full_width

  bounds
end

"""
Tests whether a voxel coordinate is within 3d bounds
"""
function in_bounds( voxel, bounds )

  #BROADCASTING IS THE DEVIL
  # ( all(voxel .>= bounds.first) &&
  #   all(voxel .<= bounds.second)  )

  b_beg, b_end = bounds

  ( voxel[1] >= b_beg[1] &&
    voxel[2] >= b_beg[2] &&
    voxel[3] >= b_beg[3] &&

    voxel[1] <= b_end[1] &&
    voxel[2] <= b_end[2] &&
    voxel[3] <= b_end[3] )
end

function in_bounds2D( voxel, bounds )

  b_beg, b_end = bounds

  ( voxel[1] >= b_beg[1] &&
    voxel[2] >= b_beg[2] &&

    voxel[1] <= b_end[1] &&
    voxel[2] <= b_end[2] )
end

"""

    bounds( d, offset=[0,0,0] )

  Extracts the index bounds for general purpose objects
  which have size and a given offset
"""
function bounds( d, offset=[0,0,0] )
  offset + 1 => collect(size(d)[1:3]) + offset
end

"""

    intersect_bounds( bounds1, bounds2, bounds2_offset )

  Takes two index bounds defined by Pairs and returns their
  intersection. Performs no checking for validity of the result
  (i.e. the result can specify bounds with 0 or negative volume).
"""
function intersect_bounds( bounds1, bounds2, bounds2_offset )
  b1_beg, b1_end = bounds1
  b2_beg = bounds2.first  + bounds2_offset
  b2_end = bounds2.second + bounds2_offset

  max( b1_beg, b2_beg ) => min( b1_end, b2_end )
end

end#module


# initalize an array
function Init!(X,L)
    size=length(X)
    h=L/size
    for i=0:size-1
        if i>floor(size//8) && i<floor(size//2)+floor(size//8)
            X[i+1]=1.-2*(i-floor(size/8))*h/L;
        else
            X[i+1]=0.0
        end
    end
end

# proc[1-3]! functions make the same computation, with different programing
# styles.

function proc1!(In::Array{Float64,1},Out::Array{Float64,1},niter::Int64)
    const size=length(In)
    h2::Float64= (1./size)^2
    for it=1:niter
        Out[2:size-1]= h2*(In[1:size-2]- 2.0*In[2:size-1]+ In[3:size])
        In,Out=Out,In
    end
end
function proc2!(In::Array{Float64,1},Out::Array{Float64,1},niter::Int64)
    const size=length(In)
    h2::Float64= (1./size)^2
    for it=1:niter
        @.  Out[2:size-1]= h2*(In[1:size-2]- 2.0*In[2:size-1]+ In[3:size])
        In,Out=Out,In
    end
end
function proc3!(In::Array{Float64,1},Out::Array{Float64,1},niter::Int64)
    const size=length(In)
    h2::Float64= (1./size)^2
    for it=1:niter
        @simd for i=2:size-1
            Out[i]= h2*(In[i-1]- 2.0*In[i]+ In[i+1])
        end
        In,Out=Out,In
    end
end
# measure as precisely as possible the execution time for a given
# p= proc[1-5]!
function test(p,In::Array{Float64,1},Out::Array{Float64,1},nit::Int)

    niter=nit
    # be sure to run once before actually running the benchmark!
    Init!(In,1.)
    Init!(Out,1.)
    p(In,Out,niter)
    #
    T=0.
    while true
        Init!(In,1.)
        Init!(Out,1.)
        t1 = time_ns()
        p(In,Out,niter)
        t = (time_ns() -t1)/niter
        if abs(t-T)/t<0.025
            break
    else
            T=t
            if niter> 150000
                error("too many iterations",niter)
            else
                niter*=2
            end
    end

end
    T,niter
end
DD=Dict("proc1!"=>"Vectorisation with braces         ",
        "proc2!"=>"Vectorisation with @. and braces  ",
        "proc3!"=>"Unrolled loop                     "
        )

fw=open("RunningOn"*gethostname()*"_lapl_1","w")

# computation starts here:
size=32
sizemax=10^5
const niter=2
while size<sizemax
    println("size: ",size)
    In=Array{Float64}(size)
    Out=Array{Float64}(size)
    tbest=10.^29
    best=0
    t=0.0
    for p= [proc1!,proc2!,proc3!]
        t,it=test(p,In,Out,niter)
        if t<tbest
            tbest=t
            best=p
        end
        t*=10.0^(-9)
        println(DD[string(p)]," : t= ",t," seconds ")
    end

    nflops=size*4
    flops=nflops/tbest
    println("\nbest: ",DD[string(best)])

    write(fw,string(size)," ",string(tbest*10.0^(-9)),"\n")

    println("nb. flops: ",nflops, ", Gflops/s (best): ",flops)
    println("-------")

    size*=2
    println()
end
close(fw)


struct ConstMean{T}
    c::T
end

(m::ConstMean)(x) = m.c.*ones(size(x, 2))

struct LinearMean{T}
    a::T
    b::T
end

function (m::LinearMean)(x)
    F, N = size(x, 1), size(x, 2)
    y = Array{eltype(m.a)}(undef, N)
    for i in 1:N
        tmp = zero(eltype(m.a))
        for j in 1:F
            tmp += m.a[j]*x[j,i]
        end
        y[i] = tmp
    end
    return y
end


function get_standardized_matrix(mat::AbstractArray{<:Real,2})
    m = mat
    Threads.@threads for col in 1:size(m)[2]
        summ = 0.0f0
        rows = size(m)[1]
        for row in 1:rows
            summ += m[row, col]
        end
        mean = summ/rows
        sums = 0.0f0
        for row in 1:rows
            sums += (m[row,col] - mean)^2
        end
        std = sqrt(sums/rows)
        for row in 1:rows
            m[row,col] = (m[row,col]-mean)/std
        end
    end
    return m
end

function get_standardized_matrix_gpu(m::AbstractArray{<:Real,2})
    return (m .- mean(m, dims=1)) ./ std(m, corrected=false, dims=1)
end

using Statistics
using Test

mat = rand(1000, 1000)

correct_std_mat = get_standardized_matrix(mat)

new_std_mat = get_standardized_matrix_gpu(mat)

compare_res = isapprox.(correct_std_mat, new_std_mat, atol=1e-5)

@test sum(compare_res) == prod(size(mat))


# coding: utf-8

# In[1]:

#okay. the epic one!
#quick sort, mmm, not so quick
#the idea of quick sort is be the quickest
#first thing first, we gotta pick a pivot number for the list
#normally, people use the first element as pivot number
#however, we may encounter a case that the first element is the largest or the smallest
#that would make our sorting a total failure
#here i use the median of 3 approach
#take the first, the last and the one in the middle
#and get the median of the three
#we use the median as a pivot number
#after that, we do the same trick as merge sort
#we have a left index and a right index
#we do traversal on both indices
#say the left index comes from the left part
#the right index comes from the right part
#we compare two elements with pivot number at the same time
#there are four cases assuming that we dont have any duplicate values in the list
#first case, left is larger than pivot, right is smaller than pivot
#so we swap both
#second case, left is larger than pivot, right is larger than pivot as well
#we freeze the left, and move right side index one step backwards
#right=right-1
#then if left is larger than pivot, right is smaller than pivot
#we back to the first case
#if it is still the second case, we repeat this procedure until we move to first case
#or left and right indices cross, we stop the sorting
#third case, left is smaller than pivot, right is smaller than pivot
#it is the opposite case of the second case
#fourth case, left is smaller than pivot, right is larger than pivot
#great, we do nothing but moving both indices closer to the centre
#these are four basic scenarios when we do both traversals
#when left and right indices cross, we stop the sorting
#and we insert the pivot number in where indices cross
#the next step is just like merge sort
#we divide the list in two halves (excluding the pivot number)
#we perform the same trick on both halves recursively
#until we reach the base case when there are only two elements in the list
#we sort those two elements with simple comparison

# In[2]:

function quick_sort(arr)

    #the first step is to get a pivot number
    #it only works on list with more than two elements
    #otherwise there is no point of getting a pivot number
    if length(arr)>2

        #we take three elements, first, last and middle
        #we create a new list
        three_element=[arr[1],arr[end],arr[Int32(floor((1+length(arr))/2))]]

        #get the median
        #remove the max and the min
        #what is left is the pivot
        maxvalue=max(three_element[1],three_element[2],three_element[3])
        minvalue=min(three_element[1],three_element[2],three_element[3])
        setdiff!(three_element,maxvalue)
        setdiff!(three_element,minvalue)
        pivot=three_element[1]

        #move the pivot to the beginning of the list
        setdiff!(arr,pivot)
        insert!(arr,1,pivot)

        #initialize
        #first index is at 2, cuz we wanna exclude pivot number
        left=2
        right=length(arr)

        #here comes the real deal
        #when left and right dont cross
        #excluding a case when two indices equal to each other
        while left-1<right

            #case 4
            if (arr[left]<pivot) && (arr[right]>pivot)
                left+=1
                right-=1

            #case 3
            elseif (arr[left]<pivot) && (arr[right]<pivot)
                left+=1

            #case 2
            elseif (arr[left]>pivot) && (arr[right]>pivot)
                right-=1

            #case 1
            else
                arr[left],arr[right]=arr[right],arr[left]
            end

        end

        #when left and right indices cross
        #when indices cross, we are one step after the input list is sorted
        #therefore, we insert pivot at left-1 instead of left
        if left>=right
            setdiff!(arr,pivot)
            insert!(arr,left-1,pivot)
        end

        #the recursive part
        #we do the same trick on two halves
        arr[left:end]=quick_sort(arr[left:end])
        arr[1:left-2]=quick_sort(arr[1:left-2])

    end

    #the base case
    #when we are left with two elements in a sublist
    #we just compare and return in reverse order
    #u might ask, what about one element
    #well, we dont have to do anything so no codes needed
    if length(arr)==2
        if arr[1]>arr[2]
            arr[1],arr[2]=arr[2],arr[1]
        end
    end

    return arr

end

# In[3]:

#there is a constraint for quick sort
#duplicate values would jeopardize everything we have built
#to solve that issue, we must amend the criteria for selecting pivot number
#for simplicity, i remove duplicates from our test list
#to handle duplicate values which might affect pivot
#my idea is to insert an extra if function
#if elements not at position 0 equal to pivot
#we create a temporary list to collect them
#when we plan to insert pivot
#we insert the temporary list to the full list

#since we have recursive functions
#this approach to get a quick sort
#is kinda different from my textbook
#to see the alternative version
#plz click the link below
# http://interactivepython.org/runestone/static/pythonds/SortSearch/TheQuickSort.html

# In[4]:

for _ in 1:100

    test_arr=rand(1000)

    if !(merge_sort(test_arr)==sort(test_arr))
        printstyled("Erreur",color=:red)
    end

end


"""
    $(FUNCTIONNAME)(model::HitranModel, grid::AbstractRange{<:Real}, pressure::Real, temperature::Real; wavelength_flag::Bool=false)

Calculate absorption cross-section at the given pressure, temperature, and grid of wavelengths
(or wavenumbers), and have the option to perform auto-differentiation

"""
function absorption_cross_section(model::AbstractCrossSectionModel,          # Model to use
                                  grid::Union{AbstractRange{<:Real}, AbstractArray}, # Wavelength [nm] or wavenumber [cm-1] grid
                                  pressure::Real,              # actual pressure [hPa]
                                  temperature::Real;
                                  autodiff::Bool=false,           # actual temperature [K]
                                  wavelength_flag::Bool=false)

    # This function takes in the "x-vector" along with the input model so that ForwardDiff will work
    function absorption_cross_section_autodiff(x ; model::AbstractCrossSectionModel = model, grid=grid, pressure=pressure,
                                               temperature=temperature, wavelength_flag=wavelength_flag)

        if length(x) !== 2
            @error "Must receive two cross-section parameters for auto-differentiation (p, T)" x
        end

        # Make sure that 𝐱 and parameter match
        @assert (pressure== x[1])
        @assert (temperature == x[2])

        cross_section = compute_absorption_cross_section(model, grid, x[1], x[2], wavelength_flag=wavelength_flag);

        return cross_section
    end

    if (autodiff)

        x = [pressure, temperature]

        result = DiffResults.JacobianResult(zeros(length(collect(grid))), x);
        ForwardDiff.jacobian!(result, absorption_cross_section_autodiff, x);

        return (result.value, result.derivs[1]);

    else
        return compute_absorption_cross_section(model, grid, pressure, temperature, wavelength_flag=wavelength_flag)
    end

end


@testset "Non-commutative Monomial" begin
    Mod.@ncpolyvar x
    X = constantmonomial(typeof(x))
    @test nvariables(X) == 0
    @test isempty(X.vars)
    @test isempty(X.z)
    X = constantmonomial(x)
    @test nvariables(X) == 1
    @test variables(X)[1] == x
    @test collect(exponents(X)) == [0]
    X = monomial(x)
    @test nvariables(X) == 1
    @test variables(X)[1] == x
    @test collect(exponents(X)) == [1]
    X = x^0
    Y = X * x^2
    @test nvariables(Y) == 1
    @test variables(Y)[1] == x
    @test collect(exponents(Y)) == [2]
    Y = x^2 * X
    @test nvariables(Y) == 1
    @test variables(Y)[1] == x
    @test collect(exponents(Y)) == [2]
end
@testset "Non-commutative MonomialVector" begin
    Mod.@ncpolyvar x y
    X = emptymonovec(typeof(x))
    @test iszero(nvariables(X))
    @test isempty(variables(X))
    X = emptymonovec(x)
    @test nvariables(X) == 1
    @test variables(X)[1] == x
    X = monovec([y, x])
    @test X[1] == x
    @test X[2] == y
    X = monomials([x, y], 2)
    @test variables(X)[1] == x
    @test variables(X)[2] == y
    @test variables(X)[3] == x
    @test collect(exponents(X[1])) == [2, 0, 0]
    @test collect(exponents(X[2])) == [1, 1, 0]
    @test collect(exponents(X[3])) == [0, 2, 0]
    @test collect(exponents(X[4])) == [0, 1, 1]
    X0 = [x^3, x^2*y, x*y^2, x*y*x, y^3, y^2*x, y*x^2, y*x*y]
    X1 = monomials([x, y], 3)
    X2 = monovec(X0)
    Z = [[3,0,0,0],[2,1,0,0],[1,2,0,0],[1,1,1,0],[0,3,0,0],[0,2,1,0],[0,1,2,0],[0,1,1,1]]
    @test X1 isa AbstractVector{<:AbstractMonomial}
    @test X2 isa AbstractVector{<:AbstractMonomial}
    @test length(X1) == length(X2) == length(Z)
    @test issorted(X1, rev=true)
    @test issorted(X2, rev=true)
    @test X1 == X0
    for i in 1:length(Z)
        @test variables(X1[i])[1] == variables(X2[i])[1] == x
        @test variables(X1[i])[2] == variables(X2[i])[2] == y
        @test variables(X1[i])[3] == variables(X2[i])[3] == x
        @test variables(X1[i])[4] == variables(X2[i])[4] == y
        @test collect(exponents(X1[i])) == collect(exponents(X2[i])) == Z[i]
    end
end


using Pipe: @pipe

mutable struct Coordinate
    x::Int
    y::Int
end

struct Vent
    from::Coordinate
    to::Coordinate
end

function input()
    s = read(stdin, String)
    return @pipe (split(s, "\n", keepempty = false)
        |> map( l -> match(r"(\d+),(\d+) -> (\d+),(\d+)", l), _ )
        |> map( m -> Vent( Coordinate(Base.parse(Int, m.captures[1]), Base.parse(Int, m.captures[2])), Coordinate(Base.parse(Int, m.captures[3]), Base.parse(Int, m.captures[4])) ), _)
    )
end
function expand_vent(v::Vent)
    result = Coordinate[]
    inc = Coordinate( v.from.x == v.to.x ? 0 : (v.to.x - v.from.x) / abs(v.to.x - v.from.x), v.from.y == v.to.y ? 0 : (v.to.y - v.from.y) / abs(v.to.y - v.from.y) )
    cursor = v.from
    while cursor.x != v.to.x || cursor.y != v.to.y
        push!(result, Coordinate(cursor.x, cursor.y))
        cursor.x += inc.x
        cursor.y += inc.y
    end
    push!(result, v.to)
    return result
end
@pipe (input()
    |> map( expand_vent, _ )
    |> Iterators.flatten( _ )
    |> collect( _ )
    |> sort( _, lt = (a,b) -> a.x < b.x || (a.x == b.x && a.y < b.y) )
    |> map( c -> (1, c), _)
    |> reduce( function(acc, c)
        if length(acc) > 0 && acc[1][2].x == c[2].x && acc[1][2].y == c[2].y
            acc[1] = (acc[1][1] + c[1], acc[1][2])
        else
            pushfirst!(acc, c)
        end
        return acc
    end, _, init=[])
    |> filter( t -> t[1] > 1, _)
    |> length(_)
    |> println(_) )


using NDTools # this is needed for the select_region! function below.

@testset "chirp z-transformation" begin
    @testset "czt" begin
        x = randn(ComplexF32, (5,6,7))
        y = randn(ComplexF32, (5,6))
        zoom = (1.0,1.0,1.0)
        @test ≈(czt(x, zoom), ft(x),rtol=1e-4)
        @test ≈(czt(y, zoom), ft(y),rtol=1e-5)
        @test ≈(iczt(czt(y,zoom),zoom),  y, rtol=1e-5)
        zoom = (2.0,2.0)
        @test ≈(czt(y,zoom),  select_region(upsample2(ft(y), fix_center=true),new_size=size(y)), rtol=1e-5)
        # @vt czt(y,zoom)  select_region(upsample2(ft(y), fix_center=true), new_size=size(y))
    end
end


import gmsh

gmsh.initialize()

gmsh.model.add("t4")

cm = 1e-02
e1 = 4.5 * cm; e2 = 6 * cm / 2; e3 =  5 * cm / 2
h1 = 5 * cm; h2 = 10 * cm; h3 = 5 * cm; h4 = 2 * cm; h5 = 4.5 * cm
R1 = 1 * cm; R2 = 1.5 * cm; r = 1 * cm
Lc1 = 0.01
Lc2 = 0.003

function hypot(a, b)
    return sqrt(a * a + b * b)
end

ccos = (-h5*R1 + e2 * hypot(h5, hypot(e2, R1))) / (h5*h5 + e2*e2)
ssin = sqrt(1 - ccos*ccos)

factory = gmsh.model.geo
factory.addPoint(-e1-e2, 0    , 0, Lc1, 1)
factory.addPoint(-e1-e2, h1   , 0, Lc1, 2)
factory.addPoint(-e3-r , h1   , 0, Lc2, 3)
factory.addPoint(-e3-r , h1+r , 0, Lc2, 4)
factory.addPoint(-e3   , h1+r , 0, Lc2, 5)
factory.addPoint(-e3   , h1+h2, 0, Lc1, 6)
factory.addPoint( e3   , h1+h2, 0, Lc1, 7)
factory.addPoint( e3   , h1+r , 0, Lc2, 8)
factory.addPoint( e3+r , h1+r , 0, Lc2, 9)
factory.addPoint( e3+r , h1   , 0, Lc2, 10)
factory.addPoint( e1+e2, h1   , 0, Lc1, 11)
factory.addPoint( e1+e2, 0    , 0, Lc1, 12)
factory.addPoint( e2   , 0    , 0, Lc1, 13)

factory.addPoint( R1 / ssin, h5+R1*ccos, 0, Lc2, 14)
factory.addPoint( 0        , h5        , 0, Lc2, 15)
factory.addPoint(-R1 / ssin, h5+R1*ccos, 0, Lc2, 16)
factory.addPoint(-e2       , 0.0       , 0, Lc1, 17)

factory.addPoint(-R2 , h1+h3   , 0, Lc2, 18)
factory.addPoint(-R2 , h1+h3+h4, 0, Lc2, 19)
factory.addPoint( 0  , h1+h3+h4, 0, Lc2, 20)
factory.addPoint( R2 , h1+h3+h4, 0, Lc2, 21)
factory.addPoint( R2 , h1+h3   , 0, Lc2, 22)
factory.addPoint( 0  , h1+h3   , 0, Lc2, 23)

factory.addPoint( 0, h1+h3+h4+R2, 0, Lc2, 24)
factory.addPoint( 0, h1+h3-R2,    0, Lc2, 25)

factory.addLine(1 , 17, 1)
factory.addLine(17, 16, 2)

factory.addCircleArc(14,15,16, 3)

factory.addLine(14,13, 4)
factory.addLine(13,12, 5)
factory.addLine(12,11, 6)
factory.addLine(11,10, 7)
factory.addCircleArc(8,9,10, 8)
factory.addLine(8,7, 9)
factory.addLine(7,6, 10)
factory.addLine(6,5, 11)
factory.addCircleArc(3,4,5, 12)
factory.addLine(3,2, 13)
factory.addLine(2,1, 14)
factory.addLine(18,19, 15)
factory.addCircleArc(21,20,24, 16)
factory.addCircleArc(24,20,19, 17)
factory.addCircleArc(18,23,25, 18)
factory.addCircleArc(25,23,22, 19)
factory.addLine(21,22, 20)

factory.addCurveLoop([17,-15,18,19,-20,16], 21)
factory.addPlaneSurface([21], 22)
factory.addCurveLoop([11,-12,13,14,1,2,-3,4,5,6,7,-8,9,10], 23)

factory.addPlaneSurface([23,21], 24)

factory.synchronize()

v = gmsh.view.add("comments")

gmsh.view.addListDataString(v, [10, -10], ["Created with Gmsh"])

gmsh.view.addListDataString(v, [0, 0.11, 0], ["Hole"],
                            ["Align", "Center", "Font", "Helvetica"])

gmsh.option.setString("View[0].DoubleClickedCommand",
                      "Printf('View[0] has been double-clicked!');")
gmsh.option.setString(
    "Geometry.DoubleClickedLineCommand",
    "Printf('Curve %g has been double-clicked!', Geometry.DoubleClickedEntityTag);")

gmsh.model.setColor([(2, 22)], 127, 127, 127)
gmsh.model.setColor([(2, 24)], 160, 32, 240)
gmsh.model.setColor([(1, i) for i in 1:14], 255, 0, 0)
gmsh.model.setColor([(1, i) for i in 15:20], 255, 255, 0)

gmsh.model.mesh.generate(2)

gmsh.write("t4.msh")

gmsh.finalize()

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


#!/usr/bin/env julia

using SphericalFourierBesselDecompositions
SFB = SphericalFourierBesselDecompositions

using Test
using LinearAlgebra

@testset "Mixing matrices" begin
    @testset "No window" begin
        # Note: inaccuracies in this test are dominated by inaccuracies in healpix.
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(3, 5, rmin, rmax)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :separable)

        # wmix
        wmix = SFB.calc_wmix(win, wmodes, amodes)
        diff = wmix - I
        for i′=1:size(diff,2), i=1:size(diff,1)
            @test isapprox(diff[i,i′], 0, atol=1e-2)
        end
        @test isapprox(wmix, I, atol=1e-2)

        # wmix_negm
        wmix_negm = SFB.calc_wmix(win, wmodes, amodes, neg_m=true)
        for i′=1:size(diff,2), i=1:size(diff,1)
            n, l, m = SFB.getnlm(amodes, i)
            n′, l′, m′ = SFB.getnlm(amodes, i′)
            correct = (n==n′ && l==l′ && m==-m′) ? 1 : 0
            @test isapprox(wmix_negm[i,i′], correct, atol=1e-2)
        end

        # M
        cmodes = SFB.ClnnModes(amodes, Δnmax=1)
        mmix1 = SFB.power_win_mix(wmix, wmix_negm, cmodes)
        @test isapprox(mmix1, I, atol=1e-2)
        M = SFB.power_win_mix(win, wmodes, cmodes)
        @test M ≈ mmix1
        @test isapprox(M, I, atol=1e-2)

        # CWlnn
        pk(k) = 1e4 * (k/1e-2)^(-3.1)
        Clnn = SFB.gen_Clnn_theory(pk, cmodes)
        CnlmNLM = SFB.Clnn2CnlmNLM(Clnn, cmodes)
        CWlnn1 = SFB.sum_m_lmeqLM(wmix * CnlmNLM * wmix', cmodes)
        CWlnn2 = M * Clnn
        @test CWlnn1 ≈ CWlnn2

        # Nshot
        nbar = 3e-4
        Nshotobs = SFB.win_lnn(win, wmodes, cmodes) ./ nbar
        Nshotobs2 = SFB.sum_m_lmeqLM(wmix, cmodes) ./ nbar
        @test isapprox(Nshotobs, Nshotobs2, rtol=1e-2)
    end

    # Inhomogeneous masks
    @testset "Inhomogeneous Window sep & insep" begin
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(2, 5, rmin, rmax)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)

        win1 = SFB.make_window(wmodes, :separable)
        @. win1.mask = rand()
        sidx = (1:length(win1.mask))[win1.mask .> 0]
        sidx = 1:(length(sidx) ÷ 2)
        @. win1.mask[sidx] = 0.5 * win1.mask[sidx]

        win2 = [win1[i,j] for i=1:size(win1,1), j=1:size(win1,2)]
        @show sum(win1.mask) ./ length(win1.mask)
        @show extrema(win1.mask)
        @show typeof(win1) typeof(win2)
        @show size(win1) size(win2)

        # wmix
        wmix1 = SFB.calc_wmix(win1, wmodes, amodes)
        wmix1_negm = SFB.calc_wmix(win1, wmodes, amodes, neg_m=true)
        wmix2 = SFB.calc_wmix(win2, wmodes, amodes)
        wmix2_negm = SFB.calc_wmix(win2, wmodes, amodes, neg_m=true)
        @show typeof(wmix1) typeof(wmix2)
        @test wmix1 == wmix2

        # M
        cmodes = SFB.ClnnModes(amodes, Δnmax=1)
        M = SFB.power_win_mix(wmix1, wmix1_negm, cmodes)
        M1 = SFB.power_win_mix(win1, wmodes, cmodes)
        M2 = SFB.power_win_mix(win2, wmodes, cmodes)
        @show extrema(M .- M1)
        @show extrema(M .- M2)
        @test M ≈ M1
        @test M ≈ M2

        # CWlnn
        pk(k) = 1e4 * (k/1e-2)^(-3.1)
        Clnn = SFB.gen_Clnn_theory(pk, cmodes)
        CnlmNLM = SFB.Clnn2CnlmNLM(Clnn, cmodes)
        CWlnn1 = SFB.sum_m_lmeqLM(wmix1 * CnlmNLM * wmix1', cmodes)
        CWlnn2 = M * Clnn
        @show extrema((CWlnn1 .- CWlnn2) ./ CWlnn1)
        @test CWlnn1 ≈ CWlnn2 rtol=1e-2

        # Nshot
        nbar = 3e-4
        Nshotobs1 = SFB.win_lnn(win1, wmodes, cmodes) ./ nbar
        Nshotobs2 = SFB.win_lnn(win2, wmodes, cmodes) ./ nbar
        @test Nshotobs1 ≈ Nshotobs2

        # Bandpower binning
        fsky = sum(win1[1,:]) / size(win1,2)
        Δℓ = round(Int, 1 / fsky)
        w̃, v = SFB.bandpower_binning_weights(cmodes; Δℓ=Δℓ, Δn=1)
        bcmodes = SFB.ClnnBinnedModes(w̃, v, cmodes)
        Nmix1 = SFB.power_win_mix(win1, w̃, v, wmodes, bcmodes)
        Nmix2 = SFB.power_win_mix(win2, w̃, v, wmodes, bcmodes)
        @show extrema(Nmix1 .- Nmix2)
        @test Nmix1 ≈ Nmix2
        w̃M1 = SFB.power_win_mix(win1, w̃, I, wmodes, bcmodes)
        Mv1 = SFB.power_win_mix(win1, I, v, wmodes, bcmodes)
        w̃M2 = SFB.power_win_mix(win2, w̃, I, wmodes, bcmodes)
        Mv2 = SFB.power_win_mix(win2, I, v, wmodes, bcmodes)
        @show extrema(w̃M1 .- w̃M2)
        @test w̃M1 ≈ w̃M2
        @test inv(Nmix1) * w̃M1 ≈ inv(Nmix2) * w̃M2
        @test Mv1 * inv(Nmix1) ≈ Mv2 * inv(Nmix2)
    end

    # complex windows
    win_descriptions = [
                        (:ang_75,),
                        (:ang_75, :radial),
                        (:ang_half, :radial),
                        (:ang_quarter, :radial),
                        (:separable, :ang_75,),
                        (:separable, :ang_75, :radial),
                        (:separable, :ang_half, :radial),
                        (:separable, :ang_quarter, :radial),
                      ]
    @testset "Window $(win_features...)" for win_features in win_descriptions
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(2, 5, rmin, rmax)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, win_features...)

        # wmix
        wmix = SFB.calc_wmix(win, wmodes, amodes)
        wmix_negm = SFB.calc_wmix(win, wmodes, amodes, neg_m=true)

        # M
        cmodes = SFB.ClnnModes(amodes, Δnmax=1)
        mmix1 = SFB.power_win_mix(wmix, wmix_negm, cmodes)
        M = SFB.power_win_mix(win, wmodes, cmodes)
        @test M ≈ mmix1

        # CWlnn
        pk(k) = 1e4 * (k/1e-2)^(-3.1)
        Clnn = SFB.gen_Clnn_theory(pk, cmodes)
        CnlmNLM = SFB.Clnn2CnlmNLM(Clnn, cmodes)
        CWlnn1 = SFB.sum_m_lmeqLM(wmix * CnlmNLM * wmix', cmodes)
        CWlnn2 = M * Clnn
        @test CWlnn1 ≈ CWlnn2

        # Nshot
        nbar = 3e-4
        Nshotobs = SFB.win_lnn(win, wmodes, cmodes) ./ nbar
        Nshotobs2 = SFB.sum_m_lmeqLM(wmix, cmodes) ./ nbar
        @test isapprox(Nshotobs, Nshotobs2, rtol=1e-2)

        # Bandpower binning basics
        fsky = sum(win[1,:]) / size(win,2)
        Δℓ = round(Int, 1 / fsky)
        w̃, v = SFB.bandpower_binning_weights(cmodes; Δℓ=Δℓ, Δn=1)
        N = w̃ * M * v
        w = inv(N) * w̃ * M
        ṽ = M * v * inv(N)
        @test w * v ≈ I  # these really just test linear algebra
        @test w̃ * ṽ ≈ I  # these really just test linear algebra

        # Bandpower binning optimizations
        bcmodes = SFB.ClnnBinnedModes(w̃, v, cmodes)
        Nmix = SFB.power_win_mix(win, w̃, v, wmodes, bcmodes)
        @test Nmix ≈ N
        w̃M = SFB.power_win_mix(win, w̃, I, wmodes, bcmodes)
        Mv = SFB.power_win_mix(win, I, v, wmodes, bcmodes)
        @test w̃M ≈ w̃ * M
        @test inv(Nmix) * w̃M ≈ w
        @test Mv * inv(Nmix) ≈ ṽ
    end
end

# vim: set sw=4 et sts=4 :


@testset "Arithmetic" begin
    x = rand(Complex{Double64})
    y = rand(Double64)

    @test x + y isa Complex{Double64}
    @test y + x isa Complex{Double64}
    @test x - y isa Complex{Double64}
    @test y - x isa Complex{Double64}
    @test x * y isa Complex{Double64}
    @test y * x isa Complex{Double64}
    @test x / y isa Complex{Double64}
    @test y / x isa Complex{Double64}
    @test x ^ y isa Complex{Double64}
    @test y ^ x isa Complex{Double64}
    @test x ^ x isa Complex{Double64}
    @test y ^ y isa Double64

    x, y = rand(2)
    @test add2(x, y) == (Double64(x) + Double64(y))
    @test sub2(x, y) == (Double64(x) - Double64(y))
    @test mul2(x, y) == (Double64(x) * Double64(y))
    @test div2(x, y) == (Double64(x) / Double64(y))

    @test x ⊕ y == (Double64(x) + Double64(y))
    @test x ⊖ y == (Double64(x) - Double64(y))
    @test x ⊗ y == (Double64(x) * Double64(y))
    @test x ⊘ y == (Double64(x) / Double64(y))
end

@testset "Trig functions" begin

    x = rand(Double64)
    sinx, cosx = sincos(x)
    @test sinx ≈ sin(x)
    @test cosx ≈ cos(x)
end


#= using UnicodePlots =#

# Set up model
function build_model()
    F = diagm([1.0])
    V = diagm([2.0])
    G = reshape([1, 2, -0.5], 3, 1)
    W = diagm([8.0, 2.5, 4.0])
    x0 = randn(1)
    P0 = diagm([1e7])
    StateSpaceModel(F, V, G, W, x0, P0)
end

# Time varying model
function sinusoid_model(omega::Real; fs::Int=256, x0=[0.5, -0.5], W::AbstractFloat=0.1)
    F(n)  = [1.0 0; 0 1.0]
    V(n)  = diagm([1e-10, 1e-10])
    G(n)  = [cos(2*pi*omega*(1/fs)*n) -sin(2*pi*omega*(1/fs)*n)]
    w(n)  = diagm([W])
    P0    = diagm([1e-1, 1e-1])
    StateSpaceModel(F, V, G, w, x0, P0)
end

facts("Kalman Filter") do

    srand(1)

    context("Time invariant models") do

        context("Building model") do
            build_model()
        end

        context("Simulations") do
            mod1 = build_model()
            x, y = TimeModels.simulate(mod1, 100)
        end

        context("Filtering") do
            mod1 = build_model()
            x, y = TimeModels.simulate(mod1, 100)
            filt = kalman_filter(y, mod1)
        end

        context("Smoothing") do
            mod1 = build_model()
            x, y = TimeModels.simulate(mod1, 100)

            smooth = kalman_filter(y, mod1) |> kalman_smooth
            smooth2 = kalman_smooth(y, mod1)

            @fact smooth.smoothed --> roughly(smooth2.smoothed, atol=1e-2)
        end

        context("Missing data") do
            mod1 = build_model()
            x, y = simulate(mod1, 100)
            y[1:9:end] = NaN
            y[100] = NaN
            filt = kalman_filter(y, mod1)
            smooth = kalman_smooth(filt)
            smooth2 = kalman_smooth(y, mod1)
            @fact any(isnan(filt.filtered)) --> false
            @fact any(isnan(smooth.filtered)) --> false
            @fact any(isnan(smooth2.smoothed)) --> false
        end

        context("Return sizes") do
            mod1 = build_model()
            x, y = TimeModels.simulate(mod1, 100)
            filt = kalman_filter(y, mod1)
            smooth = kalman_smooth(filt)
            smooth2 = kalman_smooth(y, mod1)

            @fact size(filt.filtered) --> size(filt.predicted)
            @fact size(filt.filtered) --> size(smooth.filtered)
            @fact size(filt.filtered) --> size(smooth.smoothed)
            @fact size(filt.error_cov) --> size(smooth.error_cov)
            @fact size(filt.filtered) --> size(smooth2.smoothed)
            @fact size(filt.error_cov) --> size(smooth2.error_cov)
        end

        context("Linear regression test") do
            m, b, s, dt = 5, 2, 2, .1
            t = 0:dt:10
            y_true = m*t + b
            input = 100*[sin(t/2) sin(t/4) cos(t/2) cos(t/4)] + 10
            y_noisy = [y_true zeros(length(t)) -y_true] +
                        100*[sin(t/2)+sin(t/4) sin(t/2)+cos(t/2) cos(t/2)+cos(t/4)] + 10 + randn(length(t), 3)
            lm = StateSpaceModel([1 dt; 0 1], zeros(2,2),
                                  [1. 0; 0 0; -1 0], s*eye(3),
                                  zeros(2), 100*eye(2),
                                  B=zeros(2, 4), D=[1. 1 0 0; 1 0 1 0; 0 0 1 1])
            lm_filt = kalman_filter(y_noisy, lm, u=input)
            @fact lm_filt.filtered[end,1] --> roughly(y_true[end, 1], atol=3*sqrt(lm_filt.error_cov[1,1,end]))

            lm_smooth = kalman_smooth(lm_filt)
            stderr = sqrt(lm_smooth.error_cov[1,1,:][:])
            @fact lm_filt.filtered[end,:] --> lm_smooth.smoothed[end,:]
            @fact all(abs(y_true - lm_smooth.smoothed[:,1]) .< 3*stderr) --> true
            @fact ones(t) * lm_smooth.smoothed[1,2] --> roughly(lm_smooth.smoothed[:, 2], atol=1e-12)

            # Repeat with DK smoother
            lm_smooth = kalman_smooth(y_noisy, lm, u=input)
            stderr = sqrt(lm_smooth.error_cov[1,1,:][:])
            @fact all(abs(y_true - lm_smooth.smoothed[:,1]) .< 3*stderr) --> true
            @fact ones(t) * lm_smooth.smoothed[1,2] --> roughly(lm_smooth.smoothed[:, 2], atol=1e-12)

        end

    end

    context("Time varying models") do

        context("Building Model") do
            sinusoid_model(40)
        end

        context("Simulations") do
            srand(1)
            fs = 256
            mod2 = sinusoid_model(4, fs = 256)
            x, y = TimeModels.simulate(mod2, fs*2)

            #= display(lineplot(collect(1:size(x, 1)) / fs, vec(y), width = 120, title = "Original Data")) =#
        end

        context("Filtering") do
            srand(1)
            fs = 256
            mod2 = sinusoid_model(4, fs = 256)
            x, y = TimeModels.simulate(mod2, fs*2)

            context("Correct initial guess") do
                filt = kalman_filter(y, mod2)
                @fact filt.predicted[end, :] --> roughly([0.5 -0.5]; atol= 0.3)

                #= x_est = round(filt.predicted[end, :], 3) =#
                #= display(lineplot(collect(1:size(x, 1)) / fs, vec(filt.predicted[:, 1]), width = 120, title="Filtered State 1: $(x_est[1])")) =#
                #= display(lineplot(collect(1:size(x, 1)) / fs, vec(filt.predicted[:, 2]), width = 120, title="Filtered State 2: $(x_est[2])")) =#
            end

            context("Incorrect initial guess") do
                mod3 = sinusoid_model(4, fs = 256, x0=[1.7, -0.2])
                filt = kalman_filter(y, mod3)
                @fact filt.predicted[end, :] --> roughly([0.5 -0.5]; atol= 0.3)

            end

            context("Model error") do
                mod4 = sinusoid_model(4, fs = 256, x0=[1.7, -0.2], W=3.0)
                filt = kalman_filter(y, mod4)
                @fact filt.predicted[end, :] --> roughly([0.5 -0.5]; atol= 0.3)
            end

        end

        context("Smoothing") do
            srand(1)
            fs = 256
            mod2 = sinusoid_model(4, fs = 8192)
            x, y = TimeModels.simulate(mod2, fs*10)
            smooth = kalman_smooth(y, sinusoid_model(4, fs = 8192, x0=[1.7, -0.2]) )
            @fact mean(smooth.smoothed, 1) --> roughly([0.5 -0.5]; atol= 0.1)

            #= x_est = round(smooth.smoothed[end, :], 3) =#
            #= display(lineplot(collect(1:size(x, 1)) / fs, vec(smooth.smoothed[1:end, 1]), width = 120, title="Smoothed State 1: $(x_est[1])")) =#
            #= display(lineplot(collect(1:size(x, 1)) / fs, vec(smooth.smoothed[1:end, 2]), width = 120, title="Smoothed State 2: $(x_est[2])")) =#
        end

    end

    context("Correct failure") do
        @fact_throws StateSpaceModel([1 0.1; 0 1], zeros(2,3), zeros(2,2),
                [1. 0; 0 0; -1 0], [1. 1 0 0; 1 0 1 0; 0 0 1 1], s*eye(3), zeros(2), 100*eye(2))
    end

end


@testset "linear objective function tests" begin
    PR.silence!()
    f, partition = a -> a^3, collect(-1.0:0.25:1.0)
    ufd = UnivariateFunctionData(
        f,
        a -> 3 * a^2,
        partition,
        NaN64,
        1e-3,
        1e-3,
        0,
        length(partition),
    )

    # create LP relaxation of the MILP relaxation of the univariate function.
    milp = Model(cbc_optimizer)
    @variable(milp, -1.0 <= x <= 1.0)
    @variable(milp, y)
    construct_univariate_relaxation!(milp, f, x, y, partition, true)

    # create LP relaxation of the univariate function using the convex hull formulation.
    lp = Model(cbc_optimizer)
    @variable(lp, -1.0 <= x_lp <= 1.0)
    @variable(lp, y_lp)
    construct_univariate_relaxation!(lp, f, x_lp, y_lp, partition, false)

    # univariate function data

    # Collect possible solutions of relaxations.
    sec_verts, tan_verts = PR._collect_vertices(ufd)
    sln_verts = PR.Vertex2d[]
    push!(sln_verts, sec_verts[1])
    append!(sln_verts, tan_verts)
    push!(sln_verts, sec_verts[end])

    tol = 1e-5

    for α in collect(π/20:π/20:(2*π/5))
        set_objective_function(milp, y - x * tan(α))
        set_objective_function(lp, y_lp - x_lp * tan(α))

        for s in [MOI.MIN_SENSE, MOI.MAX_SENSE]
            # Solve MILP relaxation.
            set_objective_sense(milp, s)
            optimize!(milp)
            @test termination_status(milp) == MOI.OPTIMAL

            # Verify that the solution is one of the candidate solutions in `sln_verts`.
            milp_obj = objective_value(milp)
            sln_found = false
            x_sln = value.(x)
            y_sln = value.(y)
            for v in sln_verts
                if isapprox(v[1], x_sln, atol = tol) &&
                   isapprox(v[2], y_sln, atol = tol)
                    sln_found = true
                    break
                end
            end
            @test sln_found == true

            # Verify that the LP relaxation obtained by dropping binary restrictions gives the
            # same objective value as the MILP relaxation.
            JuMP.relax_integrality(milp)
            optimize!(milp)
            @test termination_status(milp) == MOI.OPTIMAL
            lp1_obj = objective_value(milp)
            @test milp_obj ≈ lp1_obj atol = tol

            # Solve LP relaxation formulated in the convex hull form.
            # Verify that the LP relaxation formulated in the convex hull form has the same
            # objective as the above 2 relaxations.
            set_objective_sense(lp, s)
            optimize!(lp)
            @test lp1_obj ≈ objective_value(lp) atol = tol
        end
    end
end


# Plots of quadrature geometry

push!(LOAD_PATH, string(pwd(),"/src"))

using Revise
using PyPlot
import IterativeSolvers
import LinearMaps

import ModifiedStokesSolver
import AnalyticDomains

mss = ModifiedStokesSolver

# Params
alpha = 200

# Discretize
numpanels = 20
panelorder = 16
curve = AnalyticDomains.starfish(n_arms = 3, amplitude=0.1)
dcurve = mss.discretize(curve, numpanels, panelorder)

N = dcurve.numpoints
panel_length = sum(dcurve.dS[1:panelorder])

n_sub = 32
dt_max = mss.LARGEALPHA_LIMIT_16*2/(alpha*panel_length)
glpoints1, glweights1 = mss.gausslegendre(panelorder)
glpoints2, glweights2 = mss.gausslegendre(n_sub)

######################## Near boundary point
panel_idx = 1
point_idx = 6

t1 = dcurve.t_edges[:,panel_idx]
idx = panelorder*(panel_idx-1) + (1:panelorder)
xpanel = dcurve.points[1,idx]
ypanel = dcurve.points[2,idx]
zj = xpanel + 1im*ypanel
nj = dcurve.normals[1,idx] + 1im*dcurve.normals[2,idx]

h1 = t1[2]-t1[1]
@show t = glpoints1[point_idx] + 0.08im;
ztmp = curve.tau( t1[1] + (t+1)*h1/2 )
zt = [real(ztmp), imag(ztmp)]

coeffs = CurveDiscretization.map_panels(dcurve)
zcpx = zt[1] + 1im*zt[2]
troot, _, converged =
    CurveDiscretization.invert_map(dcurve, coeffs, panel_idx, zcpx,
                                   fall_back_to_initial=false)
@show troot

paperplots = true
if paperplots
    close("all")
    figure(1, figsize=(3,3))
else
    figure(1)
    clf()
end

plot(dcurve.points[1,:], dcurve.points[2,:], "-k", label=L"\Gamma")

cpx2vec(z) = [real(z),imag(z)]
function plot_edge(x, n, h, style="b")
    plot(x[1]+n[1]*[-1,1]*h, x[2]+n[2]*[-1,1]*h, style)
end

h1 = 0.06*panel_length
h2 = 0.03*panel_length

for i=1:numpanels
    x = dcurve.edges[1:2,i]
    n = dcurve.n_edges[1:2,i]
    plot_edge(x, n, h1)
end

plot(zt[1], zt[2], ".", label=L"x \in \Omega")

# Plot subintervals
bcweights = mss.bclag_interp_weights(glpoints1)
@show intervals = mss.subdivide_interval_with_bisection(dt_max, troot, n_sub)
#intervals = mss.subdivide_interval(dt_max, troot)
for idx = 1:length(intervals)-1
    ta = intervals[idx]
    tb = intervals[idx+1]
    troot_sec = (troot-ta)*2/(tb-ta) - 1.0
    # New source points, in original frame
    tsrc = ta + (1+glpoints2)*(tb-ta)/2
    # Interpolation to new source points
    Psec = mss.bclag_interp_matrix(glpoints1, tsrc, bcweights)
    zj_sec = Psec*zj
    za_sec = mss.bclag_interp_eval(glpoints1, zj, ta, bcweights)
    zb_sec = mss.bclag_interp_eval(glpoints1, zj, tb, bcweights)
    nj_sec = Psec*nj
    na_sec = mss.bclag_interp_eval(glpoints1, nj, ta, bcweights)
    nb_sec = mss.bclag_interp_eval(glpoints1, nj, tb, bcweights)

    if idx>1
        plot_edge(cpx2vec(za_sec), cpx2vec(na_sec), h2, "r")
    end
    if idx<length(intervals)-1
        plot_edge(cpx2vec(zb_sec), cpx2vec(nb_sec), h2, "r")
    end
end

xticks([])
yticks([])
minorticks_off()
axis("equal")
axis((0.9153349805587189, 1.2063410174030949, -0.02194527272134522, 0.3516612032549363))
legend()

function write_fig(num, name)
    figure(num)
    savefig(name)
    println("Saved $name")
end

filename = joinpath(dirname(@__DIR__),"figs", "quad_geo_near.pdf")
write_fig(1, filename)
run(`cp $filename ../docs/paper/figs`)

######################## On boundary point
zt = dcurve.points[:,point_idx]

coeffs = CurveDiscretization.map_panels(dcurve)
zcpx = zt[1] + 1im*zt[2]
troot = glpoints1[point_idx]

paperplots = true
if paperplots
    figure(2, figsize=(3,3))
else
    figure(2)
    clf()
end

plot(dcurve.points[1,:], dcurve.points[2,:], "-k", label=L"\Gamma")

for i=1:numpanels
    x = dcurve.edges[1:2,i]
    n = dcurve.n_edges[1:2,i]
    plot_edge(x, n, h1)
end

plot(zt[1], zt[2], ".", label=L"x \in \Gamma")

# Plot subintervals
intervals = mss.subdivide_interval_with_bisection(dt_max, troot, 32)
#intervals = mss.subdivide_interval(dt_max, troot)
for idx = 2:length(intervals)-2
    ta = intervals[idx]
    tb = intervals[idx+1]
    troot_sec = (troot-ta)*2/(tb-ta) - 1.0
    # New source points, in original frame
    tsrc = ta + (1+glpoints2)*(tb-ta)/2
    # Interpolation to new source points
    Psec = mss.bclag_interp_matrix(glpoints1, tsrc, bcweights)
    zj_sec = Psec*zj
    za_sec = mss.bclag_interp_eval(glpoints1, zj, ta, bcweights)
    zb_sec = mss.bclag_interp_eval(glpoints1, zj, tb, bcweights)
    nj_sec = Psec*nj
    na_sec = mss.bclag_interp_eval(glpoints1, nj, ta, bcweights)
    nb_sec = mss.bclag_interp_eval(glpoints1, nj, tb, bcweights)

    plot_edge(cpx2vec(za_sec), cpx2vec(na_sec), h2, "r")
    plot_edge(cpx2vec(zb_sec), cpx2vec(nb_sec), h2, "r")
end

xticks([])
yticks([])
minorticks_off()
axis("equal")
axis((0.9153349805587189, 1.2063410174030949, -0.02194527272134522, 0.3516612032549363))
legend()

filename = joinpath(dirname(@__DIR__),"figs", "quad_geo_on.pdf")
write_fig(2, filename)
run(`cp $filename ../docs/paper/figs`)


"""
```julia
vvar(A; dims=:, mean=nothing, corrected=true)
```
Compute the variance of all elements in `A`, optionally over dimensions specified by `dims`.
As `Statistics.var`, but vectorized and (optionally) multithreaded.

A precomputed `mean` may optionally be provided, which results in a somewhat faster
calculation. If `corrected` is `true`, then _Bessel's correction_ is applied, such
that the sum is divided by `n-1` rather than `n`.

## Examples
```julia
julia> using VectorizedStatistics

julia> A = [1 2; 3 4]
2×2 Matrix{Int64}:
 1  2
 3  4

julia> vvar(A, dims=1)
1×2 Matrix{Float64}:
 2.0  2.0

julia> vvar(A, dims=2)
2×1 Matrix{Float64}:
 0.5
 0.5
```
"""
function vvar(A; dims=:, mean=nothing, corrected=true, multithreaded=:auto)
    if (multithreaded===:auto && length(A) > 4095) || multithreaded===true
        _vtvar(mean, corrected, A, dims)
    else
        _vvar(mean, corrected, A, dims)
    end
end
export vvar

## Singlethreaded implementation
# If dims is an integer, wrap it in a tuple
_vvar(μ, corrected::Bool, A, dims::Int) = _vvar(μ, corrected, A, (dims,))

# If the mean isn't known, compute it
_vvar(::Nothing, corrected::Bool, A, dims::Tuple) = _vvar!(_vmean(A, dims), corrected, A, dims)
function _vvar(::Nothing, corrected::Bool, A, ::Colon)
    # Reduce all the dims!
    n = length(A)
    Tₒ = Base.promote_op(/, eltype(A), Int)
    Σ = zero(Tₒ)
    @turbo for i ∈ eachindex(A)
            Σ += A[i]
    end
    μ = Σ / n
    σ² = zero(typeof(μ))
    @turbo for i ∈ eachindex(A)
            δ = A[i] - μ
            σ² += δ * δ
    end
    return σ² / (n-corrected)
end

# If the mean is known, pass it on in the appropriate form
_vvar(μ, corrected::Bool, A, dims::Tuple) = _vvar!(collect(μ), corrected, A, dims)
_vvar(μ::Array, corrected::Bool, A, dims::Tuple) = _vvar!(copy(μ), corrected, A, dims)
_vvar(μ::Number, corrected::Bool, A, dims::Tuple) = _vvar!([μ], corrected, A, dims)
function _vvar(μ::Number, corrected::Bool, A, ::Colon)
    # Reduce all the dims!
    n = length(A)
    σ² = zero(typeof(μ))
    @turbo for i ∈ eachindex(A)
        δ = A[i] - μ
        σ² += δ * δ
    end
    return σ² / (n-corrected)
end

# Chris Elrod metaprogramming magic:
# Generate customized set of loops for a given ndims and a vector
# `static_dims` of dimensions to reduce over
function staticdim_var_quote(static_dims::Vector{Int}, N::Int)
  M = length(static_dims)
  # `static_dims` now contains every dim we're taking the var over.
  Bᵥ = Expr(:call, :view, :B)
  reduct_inds = Int[]
  nonreduct_inds = Int[]
  # Firstly, build our expressions for indexing each array
  Aind = :(A[])
  Bind = :(Bᵥ[])
  inds = Vector{Symbol}(undef, N)
  len = Expr(:call, :*)
  for n ∈ 1:N
    ind = Symbol(:i_,n)
    inds[n] = ind
    push!(Aind.args, ind)
    if n ∈ static_dims
      push!(reduct_inds, n)
      push!(Bᵥ.args, :(firstindex(B,$n)))
      push!(len.args, :(size(A, $n)))
    else
      push!(nonreduct_inds, n)
      push!(Bᵥ.args, :)
      push!(Bind.args, ind)
    end
  end
  # Secondly, build up our set of loops
  if !isempty(nonreduct_inds)
    firstn = first(nonreduct_inds)
    block = Expr(:block)
    loops = Expr(:for, :($(inds[firstn]) = indices((A,B),$firstn)), block)
    if length(nonreduct_inds) > 1
      for n ∈ @view(nonreduct_inds[2:end])
        newblock = Expr(:block)
        push!(block.args, Expr(:for, :($(inds[n]) = indices((A,B),$n)), newblock))
        block = newblock
      end
    end
    rblock = block
    # Push more things here if you want them at the beginning of the reduction loop
    push!(rblock.args, :(μ = $Bind))
    push!(rblock.args, :(σ² = zero(eltype(Bᵥ))))
    # Build the reduction loop
    for n ∈ reduct_inds
      newblock = Expr(:block)
      push!(block.args, Expr(:for, :($(inds[n]) = axes(A,$n)), newblock))
      block = newblock
    end
    # Push more things here if you want them in the innermost loop
    push!(block.args, :(δ = $Aind - μ))
    push!(block.args, :(σ² += δ * δ))
    # Push more things here if you want them at the end of the reduction loop
    push!(rblock.args, :($Bind = σ² * invdenom))
    # Put it all together
    return quote
      invdenom = inv(($len) - corrected)
      Bᵥ = $Bᵥ
      @turbo $loops
      return B
    end
  else
    firstn = first(reduct_inds)
    block = Expr(:block)
    loops = Expr(:for, :($(inds[firstn]) = axes(A,$firstn)), block)
    if length(reduct_inds) > 1
      for n ∈ @view(reduct_inds[2:end])
        newblock = Expr(:block)
        push!(block.args, Expr(:for, :($(inds[n]) = axes(A,$n)), newblock))
        block = newblock
      end
    end
    # Push more things here if you want them in the innermost loop
    push!(block.args, :(δ = $Aind - μ))
    push!(block.args, :(σ² += δ * δ))
    # Put it all together
    return quote
      invdenom = inv(($len) - corrected)
      Bᵥ = $Bᵥ
      μ = Bᵥ[]
      σ² = zero(eltype(Bᵥ))
      @turbo $loops
      Bᵥ[] = σ² * invdenom
      return B
    end
  end
end

# Chris Elrod metaprogramming magic:
# Turn non-static integers in `dims` tuple into `StaticInt`s
# so we can construct `static_dims` vector within @generated code
function branches_var_quote(N::Int, M::Int, D)
  static_dims = Int[]
  for m ∈ 1:M
    param = D.parameters[m]
    if param <: StaticInt
      new_dim = _dim(param)::Int
      @assert new_dim ∉ static_dims
      push!(static_dims, new_dim)
    else
      t = Expr(:tuple)
      for n ∈ static_dims
        push!(t.args, :(StaticInt{$n}()))
      end
      q = Expr(:block, :(dimm = dims[$m]))
      qold = q
      ifsym = :if
      for n ∈ 1:N
        n ∈ static_dims && continue
        tc = copy(t)
        push!(tc.args, :(StaticInt{$n}()))
        qnew = Expr(ifsym, :(dimm == $n), :(return _vvar!(B, corrected, A, $tc)))
        for r ∈ m+1:M
          push!(tc.args, :(dims[$r]))
        end
        push!(qold.args, qnew)
        qold = qnew
        ifsym = :elseif
      end
      # Else, if dimm ∉ 1:N, drop it from list and continue
      tc = copy(t)
      for r ∈ m+1:M
        push!(tc.args, :(dims[$r]))
      end
      push!(qold.args, Expr(:block, :(return _vvar!(B, corrected, A, $tc))))
      return q
    end
  end
  staticdim_var_quote(static_dims, N)
end

# Efficient @generated in-place var
@generated function _vvar!(B::AbstractArray{Tₒ,N}, corrected::Bool, A::AbstractArray{T,N}, dims::D) where {Tₒ,T,N,M,D<:Tuple{Vararg{Integer,M}}}
  branches_var_quote(N, M, D)
end
@generated function _vvar!(B::AbstractArray{Tₒ,N}, corrected::Bool, A::AbstractArray{T,N}, dims::Tuple{}) where {Tₒ,T,N}
  :(fill!(B, Tₒ(NaN)); return B)
end

## Multithreaded implementation
# If dims is an integer, wrap it in a tuple
_vtvar(μ, corrected::Bool, A, dims::Int) = _vtvar(μ, corrected, A, (dims,))

# If the mean isn't known, compute it
_vtvar(::Nothing, corrected::Bool, A, dims::Tuple) = _vtvar!(_vtmean(A, dims), corrected, A, dims)
function _vtvar(::Nothing, corrected::Bool, A, ::Colon)
    # Reduce all the dims!
    n = length(A)
    Tₒ = Base.promote_op(/, eltype(A), Int)
    Σ = zero(Tₒ)
    @tturbo for i ∈ eachindex(A)
            Σ += A[i]
    end
    μ = Σ / n
    σ² = zero(typeof(μ))
    @tturbo for i ∈ eachindex(A)
            δ = A[i] - μ
            σ² += δ * δ
    end
    return σ² / (n-corrected)
end

# If the mean is known, pass it on in the appropriate form
_vtvar(μ, corrected::Bool, A, dims::Tuple) = _vtvar!(collect(μ), corrected, A, dims)
_vtvar(μ::Array, corrected::Bool, A, dims::Tuple) = _vtvar!(copy(μ), corrected, A, dims)
_vtvar(μ::Number, corrected::Bool, A, dims::Tuple) = _vtvar!([μ], corrected, A, dims)
function _vtvar(μ::Number, corrected::Bool, A, ::Colon)
    # Reduce all the dims!
    n = length(A)
    σ² = zero(typeof(μ))
    @tturbo for i ∈ eachindex(A)
        δ = A[i] - μ
        σ² += δ * δ
    end
    return σ² / (n-corrected)
end

# Chris Elrod metaprogramming magic:
# Generate customized set of loops for a given ndims and a vector
# `static_dims` of dimensions to reduce over
function staticdim_tvar_quote(static_dims::Vector{Int}, N::Int)
  M = length(static_dims)
  # `static_dims` now contains every dim we're taking the var over.
  Bᵥ = Expr(:call, :view, :B)
  reduct_inds = Int[]
  nonreduct_inds = Int[]
  # Firstly, build our expressions for indexing each array
  Aind = :(A[])
  Bind = :(Bᵥ[])
  inds = Vector{Symbol}(undef, N)
  len = Expr(:call, :*)
  for n ∈ 1:N
    ind = Symbol(:i_,n)
    inds[n] = ind
    push!(Aind.args, ind)
    if n ∈ static_dims
      push!(reduct_inds, n)
      push!(Bᵥ.args, :(firstindex(B,$n)))
      push!(len.args, :(size(A, $n)))
    else
      push!(nonreduct_inds, n)
      push!(Bᵥ.args, :)
      push!(Bind.args, ind)
    end
  end
  # Secondly, build up our set of loops
  if !isempty(nonreduct_inds)
    firstn = first(nonreduct_inds)
    block = Expr(:block)
    loops = Expr(:for, :($(inds[firstn]) = indices((A,B),$firstn)), block)
    if length(nonreduct_inds) > 1
      for n ∈ @view(nonreduct_inds[2:end])
        newblock = Expr(:block)
        push!(block.args, Expr(:for, :($(inds[n]) = indices((A,B),$n)), newblock))
        block = newblock
      end
    end
    rblock = block
    # Push more things here if you want them at the beginning of the reduction loop
    push!(rblock.args, :(μ = $Bind))
    push!(rblock.args, :(σ² = zero(eltype(Bᵥ))))
    # Build the reduction loop
    for n ∈ reduct_inds
      newblock = Expr(:block)
      push!(block.args, Expr(:for, :($(inds[n]) = axes(A,$n)), newblock))
      block = newblock
    end
    # Push more things here if you want them in the innermost loop
    push!(block.args, :(δ = $Aind - μ))
    push!(block.args, :(σ² += δ * δ))
    # Push more things here if you want them at the end of the reduction loop
    push!(rblock.args, :($Bind = σ² * invdenom))
    # Put it all together
    return quote
      invdenom = inv(($len) - corrected)
      Bᵥ = $Bᵥ
      @tturbo $loops
      return B
    end
  else
    firstn = first(reduct_inds)
    block = Expr(:block)
    loops = Expr(:for, :($(inds[firstn]) = axes(A,$firstn)), block)
    if length(reduct_inds) > 1
      for n ∈ @view(reduct_inds[2:end])
        newblock = Expr(:block)
        push!(block.args, Expr(:for, :($(inds[n]) = axes(A,$n)), newblock))
        block = newblock
      end
    end
    # Push more things here if you want them in the innermost loop
    push!(block.args, :(δ = $Aind - μ))
    push!(block.args, :(σ² += δ * δ))
    # Put it all together
    return quote
      invdenom = inv(($len) - corrected)
      Bᵥ = $Bᵥ
      μ = Bᵥ[]
      σ² = zero(eltype(Bᵥ))
      @tturbo $loops
      Bᵥ[] = σ² * invdenom
      return B
    end
  end
end

# Chris Elrod metaprogramming magic:
# Turn non-static integers in `dims` tuple into `StaticInt`s
# so we can construct `static_dims` vector within @generated code
function branches_tvar_quote(N::Int, M::Int, D)
  static_dims = Int[]
  for m ∈ 1:M
    param = D.parameters[m]
    if param <: StaticInt
      new_dim = _dim(param)::Int
      @assert new_dim ∉ static_dims
      push!(static_dims, new_dim)
    else
      t = Expr(:tuple)
      for n ∈ static_dims
        push!(t.args, :(StaticInt{$n}()))
      end
      q = Expr(:block, :(dimm = dims[$m]))
      qold = q
      ifsym = :if
      for n ∈ 1:N
        n ∈ static_dims && continue
        tc = copy(t)
        push!(tc.args, :(StaticInt{$n}()))
        qnew = Expr(ifsym, :(dimm == $n), :(return _vtvar!(B, corrected, A, $tc)))
        for r ∈ m+1:M
          push!(tc.args, :(dims[$r]))
        end
        push!(qold.args, qnew)
        qold = qnew
        ifsym = :elseif
      end
      # Else, if dimm ∉ 1:N, drop it from list and continue
      tc = copy(t)
      for r ∈ m+1:M
        push!(tc.args, :(dims[$r]))
      end
      push!(qold.args, Expr(:block, :(return _vtvar!(B, corrected, A, $tc))))
      return q
    end
  end
  staticdim_tvar_quote(static_dims, N)
end

# Efficient @generated in-place var
@generated function _vtvar!(B::AbstractArray{Tₒ,N}, corrected::Bool, A::AbstractArray{T,N}, dims::D) where {Tₒ,T,N,M,D<:Tuple{Vararg{Integer,M}}}
  branches_tvar_quote(N, M, D)
end
@generated function _vtvar!(B::AbstractArray{Tₒ,N}, corrected::Bool, A::AbstractArray{T,N}, dims::Tuple{}) where {Tₒ,T,N}
  :(fill!(B, Tₒ(NaN)); return B)
end