# Compute windowed versions of the Single Layer potential applied to the known density function of the solution of an incident plane wave on a Dirichlet sphere.
using Cubature;
using GSL;

function field(k,the,r)
	res = exp(1im*k*cos(the)*r); # Incident field
	prevTerm = 20*k;
	term = 0;
	for n =0:round(1.4*k+40)
		n = Int64(n);
		term = (2*n+1)*1im^n*sf_bessel_jl(n,k)*(sf_bessel_jl(n,k*r)+1im*sf_bessel_yl(n,k*r))/(sf_bessel_jl(n,k)+1im*sf_bessel_yl(n,k))*sf_legendre_Plm(n,0,cos(the))
		if((abs(term) > abs(prevTerm) ) & (n > 10))
			return res;
		end
		res = res - term;
		prevTerm = term;
	end
	return res
end

function qDens(k,the)
	res = 1im*k*cos(the)*exp(1im*k*cos(the)); # Normal derivative of the incident field
	prevTerm = 20*k;
	term = 0;
	for n =0:round(1.4*k+40)
		n = Int64(n);
		if(n==0)
			derivHank = exp(1im*k)/k*(1+1im/k)
		else
			derivHank = (sf_bessel_jl(n-1,k)+1im*sf_bessel_yl(n-1,k) -(n+1)/k*(sf_bessel_jl(n,k)+1im*sf_bessel_yl(n,k)))
		end
		term = (2*n+1)*1im^n*sf_bessel_jl(n,k)*k*derivHank/(sf_bessel_jl(n,k)+1im*sf_bessel_yl(n,k))*sf_legendre_Plm(n,0,cos(the))
		if((abs(term) > abs(prevTerm) ) & (n > 10))
			return res;
		end
		res = res - term;
		prevTerm = term;
	end
	return res
end

function dis(x,y)
	phis = x[1];
	thetas=x[2];
	rstar = x[3];

	phi = y[1];
	theta=y[2];
	sqrt((cos(theta)-rstar*cos(thetas))^2 + (sin(theta)*sin(phi)-rstar*sin(thetas)*sin(phis))^2 + (sin(theta)*cos(phi)-rstar*sin(thetas)*cos(phis))^2);
end

function inte(x,y,k)
	theta=y[2]; # Jacobian is sin(theta)
	exp(k*1im*dis(x,y) )/4/pi/dis(x,y)*qDens(k,theta)*sin(theta)
end

function windInte(x,y,k,a,b)
	phis = x[1];
	thetas=x[2];
	phi = y[1];
	theta=y[2];
	r = sqrt( (sin(theta)*sin(phi)-sin(thetas)*sin(phis))^2 + (sin(theta)*cos(phi)-sin(thetas)*cos(phis))^2)
	if ((r <= a) || (abs(r-1) < 1e-8))
		return	exp(k*1im*dis(x,y) )/4/pi/dis(x,y)*qDens(k,theta)*sin(theta);
	elseif r < b
		return 	exp(k*1im*dis(x,y) )/4/pi/dis(x,y)*qDens(k,theta)*sin(theta)*exp(2*exp((a-b)/(r-a))/( (r-a)/(b-a)-1));
	else
		return 0;
	end
end

function pCmplx(rea,ima)
	if ima > 0
		return string(rea, " + ", ima, "im ");
	else
		return string(rea, " ",   ima, "im ");
	end
end

function intePol(x,y,k)
	rPol = y[1]
	thetaPol = y[2];
	if((rPol > 2) || (rPol < 0) || (x[3]^2*sin(thetaPol)^2/x[1]^2 + cos(thetaPol)^2 + (sin(thetaPol)-x[2]*cos(thetaPol)/x[3])^2 < 0) )
		print("Error: r=", rPol, ", theta=", thetaPol, ", under sqrt= " , x[3]^2*sin(thetaPol)^2/x[1]^2 + cos(thetaPol)^2 + (sin(thetaPol)-x[2]*cos(thetaPol)/x[3])^2, "\n")
	end
	xCo = -(1/2)*rPol*sqrt(-rPol^2+4)*x[3]*sin(thetaPol)/(sqrt(x[3]^2*sin(thetaPol)^2/x[1]^2 + cos(thetaPol)^2 + (sin(thetaPol)-x[2]*cos(thetaPol)/x[3])^2)*x[1])+(1-(1/2)*rPol^2)*x[1];
	if(abs(xCo) > 1)
		print("Error: r=", rPol, ", theta=", thetaPol, ", xCo= " , xCo, ", x = ", x, "\n")
	end
	exp(k*1im*rPol)/4/pi*qDens(k,acos(xCo))
end

xt = [0.8, pi/4, 1.0];
xtCar = xt[3]*[cos(xt[2]), sin(xt[2])*sin(xt[1]), sin(xt[2])*cos(xt[1])];
xo = [0.8, pi/4, 1.4];
xi = [2*pi/3, -0.1, 0.5];
percDecay = 0.5;

ks = 2.^(3:5);
Ts = linspace(0.9,1.6,2);
tol = 1e-8;

print(field(ks[1],xo[2],xo[3] ), "= total field ext, BC outside = ", exp(1im*ks[1]*cos(xo[2])*xo[3]), "\n")

# Test to check whether the integral without and with window satisfies the boundary condition:
(valr,err) = hcubature(y -> begin real(inte(xt,y,kt)); end, [0,0],[2*pi,pi])
(vali,eri) = hcubature(y -> begin imag(inte(xt,y,kt)); end, [0,0],[2*pi,pi])
print(kt , " =k, integral at x=", xt, " is ", pCmplx(valr,vali), ", b(x) = ", exp(1im*kt*cos(xt[2])), "\n")

bt = 1.9;
at = 1.2;
(valr,err) = hcubature(y -> begin real(windInte(xt,y,kt,at,bt)); end, [0,0],[2*pi,pi])
(vali,eri) = hcubature(y -> begin imag(windInte(xt,y,kt,at,bt)); end, [0,0],[2*pi,pi])
print(kt , " =k, wind integral at x=", xt, " is ", pCmplx(valr,vali), ", b(x) = ", exp(1im*kt*cos(xt[2])), ", a=", at, ", b=", bt, "\n")
# Alternatives to hcubature: hcubature_v, hquadrature, hquadrature_v: the _v should execute in parallel

# Test the integral using polar coordinates
if false
	# This integral should be small as there is no critical point in the range
	(valr,err) = hcubature(y -> begin real(intePol(xtCar,y,ks[ki])); end, [0.67,1.77],[0.7,1.8])
	(valr,err) = hcubature(y -> begin imag(intePol(xtCar,y,ks[ki])); end, [0.67,1.77],[0.7,1.8])
	print(valr, " + (", vali, ") polar coords \n");
elseif false
	(valr,err) = hcubature(y -> begin real(intePol(xtCar,y,ks[ki])); end, [0,0],[2,2*pi])
	(vali,eri) = hcubature(y -> begin imag(intePol(xtCar,y,ks[ki])); end, [0,0],[2,2*pi])
	print(valr, " + (", vali, ") polar coords \n");
end

if false
	# Integrate along some path of steepest descent
	function integr(p) exp(50*im*(-0.72764585274627917e-1*p^2+(.59216178402453730*1im)*p+(.40824829046386303*(-0.72764585274627917e-1*p^2+(.59216178402453730*1im)*p))*sqrt(-(-0.72764585274627917e-1*p^2+(.59216178402453730*1im)*p)^2+4)+(1/3*(1-(1/2)*(-0.72764585274627917e-1*p^2+(.59216178402453730*1im)*p)^2))*sqrt(3)))*(0.45572809000084120e-1-0.72764585274627917e-1*p^2+(.59216178402453730*1im)*p+(-0.72764585274627917e-1*p^2+(.59216178402453730*1im)*p)^2)*(-.14552917054925583*p+.59216178402453730*1im) end # thi=1,sdi=1
	maxP = (17/50)*log(10);
	(valr,err) = hcubature(x -> begin real(integr(x[1])); end, 0, maxP)
	(vali,eri) = hcubature(x -> begin imag(integr(x[1])); end, 0, maxP)
	print(valr+vali*1im, " =val, err = ", err+eri*1im, " , maxP =",maxP ," , integr at maxP = ", integr(maxP), '\n')
end

intVals = (0+0im)*zeros(length(ks),2)
intWindVals = (0+0im)*zeros(length(ks),length(Ts));
extVals = (0+0im)*zeros(length(ks),2)
extWindVals = (0+0im)*zeros(length(ks),length(Ts));
realVals = (0+0im)*zeros(length(ks),2)
windVals = (0+0im)*zeros(length(ks),length(Ts));
realErrs = (0+0im)*zeros(length(ks),1)
windErrs = (0+0im)*zeros(length(ks),length(Ts));

for ki = 1:length(ks)
	print("\n k = ", ks[ki], " ", now());
	valr = 0.0; err = 0.0; vali = 0.0; eri = 0.0;
	(valr,err) = hcubature(y -> begin real(inte(xt,y,ks[ki])); end, [0,0],[2*pi,pi], reltol = tol)
	(vali,eri) = hcubature(y -> begin imag(inte(xt,y,ks[ki])); end, [0,0],[2*pi,pi], reltol = tol)
	realVals[ki,2] = exp(1im*ks[ki]*cos(xt[2]));
	realVals[ki,1] = valr+1im*vali - realVals[ki,2];
	realErrs[ki,1] = err+1im*eri;

	(valr,err) = hcubature(y -> begin real(inte(xo,y,ks[ki])); end, [0,0],[2*pi,pi], reltol = tol)
	(vali,eri) = hcubature(y -> begin imag(inte(xo,y,ks[ki])); end, [0,0],[2*pi,pi], reltol = tol)
	extVals[ki,2] = field(ks[ki],xo[2], xo[3]);
	extVals[ki,1] = valr+1im*vali + extVals[ki,2]; # Maybe also try valr+1im*vali + exp(1im*ks[ki]*cos(xo[2])*xo[3]);
	print("Did BC and exterior, starting interior at ", now());

	(valr,err) = hcubature(y -> begin real(inte(xi,y,ks[ki])); end, [0,0],[2*pi,pi], reltol = tol)
	(vali,eri) = hcubature(y -> begin imag(inte(xi,y,ks[ki])); end, [0,0],[2*pi,pi], reltol = tol)
	intVals[ki,2] = exp(1im*ks[ki]*cos(xi[2])*xi[3]);
	intVals[ki,1] = valr+1im*vali - intVals[ki,2];

	for ti = 1:length(Ts)
		print("Starting windowed version with T = ", Ts[ti], " at ", now());
		valr = 0.0; err = 0.0; vali = 0.0; eri = 0.0;
		(valr,err) = hcubature(y -> begin real(windInte(xt,y,ks[ki],(1-percDecay)*Ts[ti],Ts[ti])); end, [0,0],[2*pi,pi], reltol = tol)
		(vali,eri) = hcubature(y -> begin imag(windInte(xt,y,ks[ki],(1-percDecay)*Ts[ti],Ts[ti])); end, [0,0],[2*pi,pi], reltol = tol)
		windVals[ki,ti] = valr+1im*vali - realVals[ki,2];
		windErrs[ki,ti] = err+1im*eri;

		(valr,err) = hcubature(y -> begin real(windInte(xo,y,ks[ki],(1-percDecay)*Ts[ti],Ts[ti])); end, [0,0],[2*pi,pi], reltol = tol)
		(vali,eri) = hcubature(y -> begin imag(windInte(xo,y,ks[ki],(1-percDecay)*Ts[ti],Ts[ti])); end, [0,0],[2*pi,pi], reltol = tol)
		extWindVals[ki,ti] = valr+1im*vali + extVals[ki,2];
		print(", T ext/int, ", now());

		(valr,err) = hcubature(y -> begin real(windInte(xi,y,ks[ki],(1-percDecay)*Ts[ti],Ts[ti])); end, [0,0],[2*pi,pi], reltol = tol)
		(vali,eri) = hcubature(y -> begin imag(windInte(xi,y,ks[ki],(1-percDecay)*Ts[ti],Ts[ti])); end, [0,0],[2*pi,pi], reltol = tol)
		intWindVals[ki,ti] = valr+1im*vali - intVals[ki,2];
	end
	f = open("simRes.txt","w");
	write(f,"xboundary=", string(xt), "xouter=", string(xo), "xinterior=", string(xi), "\n");
	write(f,"percDecay=", string(percDecay), ", tol=", tol, "\n")
	write(f,"ks=", string(ks), "\n Ts=", string(Ts), "\n")
	write(f,"intVals [scat-inc, inc] ", string(intVals), "\n")
	write(f,"intWindVals [scat-inc] ", string(intWindVals), "\n")
	write(f,"extVals [scat+exact field, exact field] ", string(extVals), "\n")
	write(f,"extWindVals [scat+exact field] ", string(extWindVals), "\n")
	write(f,"realVals [scat-inc, inc] ", string(realVals), "\n")
	write(f,"windVals (only [scattered - incident]) ", string(windVals), "\n")
	write(f,"realErrs", string(realErrs), "\n")
	write(f,"windErrs", string(windErrs), "\n")
	close(f);

	writedlm("intVals", intVals)
	writedlm("intWindVals", intWindVals)
	writedlm("extVals", extVals)
	writedlm("extWindVals", extWindVals)
	writedlm("realVals", realVals)
	writedlm("realErrs", realErrs)
	writedlm("windVals", windVals)
	writedlm("windErrs", windErrs)
end
print("\n\n")
run(`cat simRes.txt`)
print("\n\n")


# Learning Rules.
# We are creating new data types here to hold each learning rule as a concept.
struct hebb end
struct hebbdecay end
struct oja end
struct willshaw end
struct antihebbdecay end

# Implementing different learning methods for different learning rule
# Here we are making use of the multiple dispatch feature in Julia
# Multiple Dispatch provides an unified framework for implementing different methods
learning(X,W,γ,r, rule::hebb;δ) = W + (γ .* X .* r)
learning(X,W,γ,r,rule::hebbdecay;δ) = W + (γ .* X .* r) - (γ .* δ .* W)
learning(X,W,γ,r,rule::oja;δ) = W + (γ .* X .* r) - (γ .* r^2 .* W)
learning(X,W,γ,r,rule::willshaw;δ) = W + (γ .* X .* r) .* (X .- W)
learning(X,W,γ,r,rule::antihebbdecay;δ) = W - (γ .* X .* r) - (γ * δ .* (W .- 1))

"""
USAGE:

`hebbianlearn(X,W,b, γ, r, hebb())` for hebbian learning \n
`hebbianlearn(X,W,b, γ, r, hebbdecay(), δ = delata_val)` for hebbian decay learning rule \n
`hebbianlearn(X,W,b, γ, r, oja(), δ = delata_val)` for oja learning rule \n
`hebbianlearn(X,W,b, γ, r, willshaw(), δ = delata_val)` for willshaw learning rule \n
`hebbianlearn(X,W,b, γ, r, antihebbdecay(), δ = delata_val)` for antihebbdecay learning rule

"""
function hebbianlearn(X, W, b, γ, r, rule;δ=0)
    Wb = [W b]
    Xnew = [X ones(size(X,1))]
    Wb = learning(Xnew,Wb, γ,r,rule, δ=δ)
    return Wb[:,1:size(W,2)], Wb[:,size(W,2)+1] # Unpacking W and b from Wb
end

hebbiandemo


# generate a spectra with multiple peaks for calibration testing
dist_energies = [4000,4700,5500,6000,6800,7500,8200,9600,10000]
energy_value_distribution = MixtureModel([Normal(mu,15) for mu in dist_energies])

be = 0:1:20000

events = rand(energy_value_distribution, 10000)
be, c = hist(events, be)
k=UnivariateKDE(be, float(c))

findpeaks(h::Histogram, fwhm) = findpeaks(bin_edges(h), counts(h), fwhm, typemax(Int))
findpeaks(h::Histogram, fwhm,n) = findpeaks(bin_edges(h), counts(h), fhwm,n)

findpeak(bin_edges::Range, x::Vector, fwhm) = findpeaks(bin_edges, x, fwhm, typemax(Int))
function findpeaks(bin_edges::Range, x::Vector, fwhm,n)
	k = UnivariateKDE(bin_edges, convert(Array{Float64,1},x))
	findpeaks(k,fwhm,n)
end

function findpeaks(x::UnivariateKDE, fwhm,n)
	σ = fwhm/sqrt(8*log(2))
	d = Normal(0,σ)
	smoothed = conv(x,d)
	peakinds = findpeaks_(smoothed.density)
	if length(peakinds)>n
		peakinds = peakinds[end-n+1:end]
	end
	peakinds
end

function findpeaks_{T}(x::Vector{T})
	# find all local maxima, return a vector of indicies sorted by value of x at that index
	peakinds = Array(Int,0)
	for i = 2:length(x)-1
		if x[i-1]<x[i]>x[i+1]
			push!(peakinds, i)
		end
	end
	p = sortperm(x[peakinds])
	peakinds[p]
end

function plotpeaks(bin_edges::Range, x::Vector, fwhm, n)
	k = UnivariateKDE(bin_edges, convert(Array{Float64,1},x))
	peakinds = findpeaks(bin_edges, x, fwhm, n)
	peakx = midpoints(bin_edges)[peakinds]
	figure()
	plot(midpoints(bin_edges), x, label="x")
	plot(midpoints(k.x), conv(k,Normal(0,fwhm/2.355)).density, label="smoothed")
	plot(peakx, x[peakinds],"o",label="raw peaks")
	plot(peakx, conv(k,Normal(0,fwhm/2.355)).density[peakinds],"o",label="smoothed peaks")
	legend()
end

peakinds = findpeaks(k,50,length(dist_energies))
peakx = sort(midpoints(be)[peakinds])

plotpeaks(be, c, 10,10)

figure()
plot(midpoints(be),c)
plot(midpoints(k.x), conv(k,Normal(0,10/2.355)).density)
plot(midpoints(be)[peakinds], c[peakinds],"o")
plot(midpoints(k.x)[peakinds], conv(k,Normal(0,10/2.355)).density[peakinds],"o")

@test all(abs(peakx-dist_energies).<5);


function RK4(f, u0, h, t0)
    y0 = u0
    # k0 = f(y0,t0)
    k1 = f(y0,t0)
    k2 = f(y0+h*k1/2,t0+h/2)
    k3 = f(y0+h*k2/2,t0+h/2)
    k4 = f(y0+h*k3,t0+h)
    y1 = y0+(k1+2k2+2k3+k4)*h/6
    return y1
end

function RK2(f, u0, h, t0)
    y0=u0
    k1=f(y0,t0)
    k2=f(y0+k1*h*0.5,t0+h/2)
    y1=y0+k2*h
    return y1
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM, Test

# Local stiffness matrix of plane stress element

element = Element(Quad4, (1, 2, 3, 4))
X = Dict(1 => [0.0, 0.0], 2 => [1.0, 0.0], 3 => [1.0, 1.0], 4 => [0.0, 1.0])
u = Dict(1 => [0.0, 0.0], 2 => [0.0, 0.0], 3 => [0.0, 0.0], 4 => [0.0, 0.0])
update!(element, "geometry", X)
update!(element, "displacement", u)
update!(element, "youngs modulus", 288.0)
update!(element, "poissons ratio", 1/3)
update!(element, "displacement load", [4.0, 8.0])

problem = Problem(Elasticity, "[0x1] x [0x1] block", 2)
problem.properties.formulation  = :plane_stress
add_element!(problem, element)
assemble!(problem, 0.0)

K = Matrix(problem.assembly.K)
f = Vector(problem.assembly.f)

K_expected = [
    144  54 -90   0 -72 -54  18   0
     54 144   0  18 -54 -72   0 -90
    -90   0 144 -54  18   0 -72  54
      0  18 -54 144   0 -90  54 -72
    -72 -54  18   0 144  54 -90   0
    -54 -72   0 -90  54 144   0  18
     18   0 -72  54 -90   0 144 -54
      0 -90  54 -72   0  18 -54 144]

f_expected = [1, 2, 1, 2, 1, 2, 1, 2]

@test isapprox(K, K_expected)
@test isapprox(f, f_expected)


"""
    finite-difference velocity objective
"""
struct VelocityObjective <: Objective
    Q
    n
    h
end

velocity_objective(Q, n; h = 0.0) = VelocityObjective(Q, n, h)

function objective(Z, obj::VelocityObjective, idx, T)
    n = obj.n

    J = 0.0

    for t = 1:T-1
        q⁻ = view(Z, idx.x[t][n .+ (1:n)])
        q⁺ = view(Z, idx.x[t + 1][n .+ (1:n)])
        h = obj.h == 0.0 ? view(Z, idx.u[t][end]) : obj.h
        v = (q⁺ - q⁻) ./ h

        J += v' * obj.Q[t] * v
    end

    return J
end

function objective_gradient!(∇J, Z, obj::VelocityObjective, idx, T)

    n = obj.n

    for t = 1:T-1
        q⁻ = view(Z, idx.x[t][n .+ (1:n)])
        q⁺ = view(Z, idx.x[t + 1][n .+ (1:n)])
        h = obj.h == 0.0 ? view(Z, idx.u[t][end]) : obj.h
        v = (q⁺ - q⁻) ./ h

        dJdv = 2.0 * obj.Q[t] * v
        ∇J[idx.x[t][n .+ (1:n)]] += -1.0 ./ h * dJdv
        ∇J[idx.x[t + 1][n .+ (1:n)]] += 1.0 ./ h * dJdv
        if obj.h == 0.0
            ∇J[idx.u[t][end]] += -1.0 * dJdv' * v ./ h
        end
    end

    return nothing
end


# by Lazaro Alonso
#using GLMakie, Random # HIDE
using GLMakie, Proj4, NPZ, FileIO
using Colors, ColorSchemes
GLMakie.activate!()
let
    # download the data from here:
    # https://github.com/lazarusA/BeautifulMakie/tree/main/_assets/data

    x = npzread("./data/aggregationOSMx.npy")
    y = npzread("./data/aggregationOSMy.npy")
    aggOSM = npzread("./data/aggregationOSM.npy")'
    img = load("./data/osm_kbc.png") # osm_bmy.png

    # getting back to lon and lat coordinates
    source = Projection("+proj=webmerc +datum=WGS84")
    dest = Projection("+proj=longlat +datum=WGS84")
    lonlat = transform(source, dest, [x y])

    function toCartesian(lon, lat; r = 1, cxyz = (0,0,0) )
        lat, lon = lat*π/180, lon*π/180
        cxyz[1] + r * cos(lat) * cos(lon), cxyz[2] + r * cos(lat) * sin(lon), cxyz[3] + r *sin(lat)
    end
    function lonlat3D(lon, lat, data; cxyz = (0,0,0))
        xyzw = zeros(size(data)..., 3)
        for (i,lon) in enumerate(lon), (j,lat) in enumerate(lat)
            x, y, z = toCartesian(lon, lat; cxyz = cxyz)
            xyzw[i,j,1] = x
            xyzw[i,j,2] = y
            xyzw[i,j,3] = z
        end
        xyzw[:,:,1], xyzw[:,:,2], xyzw[:,:,3]
    end

    x, y, z = lonlat3D(lonlat[:,1], lonlat[:,2], aggOSM)

    fig = Figure(figure_padding = -650, resolution = (1400,1400),
        backgroundcolor = "#141414")
    ax = Axis3(fig, aspect = :data, viewmode = :fit, perspectiveness = 0.5,
        azimuth = 6.49150301535897, elevation = 0.8516713267948967)
    mesh!(Sphere(Point3(0),1), color = :black, shading = false)
    surface!(ax, x, y, z, color = rotr90(img), shading = false,
        transparency = false)
    hidedecorations!(ax)
    hidespines!(ax)
    fig[1,1] = ax
    fig
    save(joinpath(@__DIR__, "output", "proj4_to_Sphere.png"), fig) # HIDE
end
using Pkg # HIDE
Pkg.status(["GLMakie","FileIO", "ImageMagick", "QuartzImageIO", "Proj4", "Colors", "ColorSchemes"]) # HIDE


# ==============================================================================
# Motor
#
# Part of the SLICOT benchmarks
#
# system type: continuous LTI system
# state dimension: 8
# input dimension: 2
# ==============================================================================
using SparseArrays

function state_matrix_motor()
    # system matrix
    I = [1, 2, 2, 3, 3, 3, 3, 4, 5, 6, 6, 7, 7, 7, 7, 8]
    J = [2, 3, 2, 1, 2, 3, 4, 1, 6, 7, 6, 5, 6, 7, 8, 5]
    V = [1, 8487.2, -1.0865, -2592.1, -21.119, -698.91, -141399.0, 1.0, 1.0,
         8487.2, -1.0865, -2592.1, -21.119, -698.91, -141399.0, 1.0]
    A = sparse(I, J, V)
end

function initial_state_motor()
    # initial set:
    # xᵢ ∈ [0.002, 0.0025] for i = 1
    # xᵢ ∈ [0.001, 0.0015] for i = 5
    # xᵢ = 0 otherwise
    X0 = Hyperrectangle(low=[0.002, 0.0, 0.0, 0.0, 0.001, 0.0, 0.0, 0.0],
                        high=[0.0025, 0.0, 0.0, 0.0, 0.0015, 0.0, 0.0, 0.0])
end

function motor_homog()
    A = state_matrix_motor()
    X0 = initial_state_motor()

    prob = @ivp(x' = Ax, x(0) ∈ X0) # continuous LTI system
    tspan = (0.0, 20.0)
    return prob, tspan
end

function motor()
    A = state_matrix_motor()
    B = sparse([4, 8], [1, 2], [-1.0, -1.0]) # input matrix
    X = Universe(8) # state domain: unconstrained
    U = Hyperrectangle([0.23, 0.3], [0.07, 0.1]) # input domain
    X0 = initial_state_motor()

    prob = @ivp(x' = Ax + Bu, x ∈ X, u ∈ U, x(0) ∈ X0) # continuous LTI system
    tspan = (0.0, 20.0)
    return prob, tspan
end


export NwtdirectionLDLt, NewtonLDLtAbsLS

include("ldlt_symm.jl")

function NwtdirectionLDLt(H,g; scale_abs :: Bool = true)
    L = Matrix{Float64}
    D = Matrix{Float64}
    pp = Array{Int}
    ρ = Float64
    ncomp = Int64

    try
        (L, D, pp, rho, ncomp) = ldlt_symm(H,'r')
    catch
 	println("*******   Problem in LDLt")
        #res = PDataLDLt()
        #res.OK = false
        #return res
        return (NaN, NaN, NaN, Inf, false, true, :fail)
    end

    # A[pp,pp] = P*A*P' =  L*D*L'

    if true in isnan.(D)
 	println("*******   Problem in D from LDLt: NaN")
        println(" cond (H) = $(cond(H))")
        #res = PDataLDLt()
        #res.OK = false
        #return res
        return (NaN, NaN, NaN, Inf, false, true, :fail)
    end

    Δ, Q = eig(D)

    ϵ2 =  1.0e-8
    if scale_abs
        Γ = max.(abs.(Δ),ϵ2)
    else
        Γ = max.((Δ),ϵ2)
    end
    # Ad = P'*L*Q*Δ*Q'*L'*Pd =    -g
    # replace Δ by Γ to ensure positive definiteness
    d̃ = L\g[pp]
    d̂ = L'\ (Q*(Q'*d̃ ./ Γ))
    d = - d̂[invperm(pp)]

    return d
end

function NewtonLDLtAbsLS(nlp :: AbstractNLPModel;
                         x :: Vector{T}=copy(nlp.meta.x0),
                         stp :: NLPStopping = NLPStopping(nlp,
                                                          NLPAtX(nlp.meta.x0)),
                         LS_algo   :: Function = bracket,
                         LS_logger :: AbstractLogger = Logging.NullLogger(),
                         kwargs...) where T
    return  Newton_G_StopLS(nlp;
                            x = x,
                            stp = stp,
                            LS_algo = LS_algo,
                            LS_logger = LS_logger,
                            NwtDirection = NwtdirectionSpectral,
                            hessian_rep = hessian_dense,
                            kwargs...)
end


## TODO: simplify forward function through forward! (etc...)

export PlanarFlow

import Base: eltype, size, length

@doc raw"""
The PlanarFlow function is defined as

```math
f({\bf{x}}) = {\bf{x}} + {\bf{u}} \tanh({\bf{w}}^\top {\bf{x}} + b)
```

with input and output dimension ``D``. Here ``{\bf{x}}\in \mathbb{R}^D`` represents the input of the function. Furthermore ``{\bf{u}}\in \mathbb{R}^D``, ``{\bf{w}}\in \mathbb{R}^D`` and ``b\in\mathbb{R}`` represent the parameters of the function. The function contracts and expands the input space.

This function has been introduced in:

Rezende, Danilo, and Shakir Mohamed. "Variational inference with normalizing flows." _International conference on machine learning._ PMLR, 2015.
"""
mutable struct PlanarFlow{T1, T2 <: Real} <: AbstractCouplingFlow
    u       :: T1
    w       :: T1
    b       :: T2
    function PlanarFlow(u::T1, w::T1, b::T2) where { T1, T2 <: Real}
        @assert length(u) == length(w) "The parameters u and w in the PlanarFlow structure should have equal length."
        return new{T1,T2}(u, w, float(b))
    end
end
struct PlanarFlowEmpty{N} <: AbstractCouplingFlowEmpty end
PlanarFlowEmpty(dim::Int) = PlanarFlowEmpty{dim}()
struct PlanarFlowPlaceholder <: AbstractCouplingFlowPlaceholder end
@doc raw"""
The `PlanarFlow()` function creates a planar flow object. Its dimensionality is automatically set when wrapping this object inside a model. Its parameters are initialized during compilation.
"""
PlanarFlow() = PlanarFlowPlaceholder()
function prepare(dim::Int, flow::PlanarFlowPlaceholder)
    return PlanarFlowEmpty(dim)
end

# compile placeholder
compile(f::PlanarFlowEmpty{1})          = PlanarFlow(randn(), randn(), randn())
compile(f::PlanarFlowEmpty)             = PlanarFlow(randn(getdim(f)), randn(getdim(f)), randn())
compile(f::PlanarFlowEmpty{1}, params)  = PlanarFlow(params[1], params[2], params[3])
compile(f::PlanarFlowEmpty, params)     = PlanarFlow(params[1:getdim(f)], params[1+getdim(f):2*getdim(f)], params[2*getdim(f)+1])

@doc raw"""
The `PlanarFlow(dim::Int64)` function creates a mutable `PlanarFlow` structure with parameters corresponding to input of dimensions `dim`. The parameters are each random sampled from a standard (multivariate) normal distribution.
"""
function PlanarFlow(dim::Int64)
    if dim == 1
        return PlanarFlow(randn(), randn(), randn())
    else
        return PlanarFlow(randn(dim), randn(dim), randn())
    end
end

# number of parameters
nr_params(flow::PlanarFlow)      = length(flow.u) + length(flow.w) + length(flow.b)
nr_params(flow::PlanarFlowEmpty{N}) where {N} = 2*N + 1

# get-functions for the PlanarFlow structure.
getu(f::PlanarFlow)                         = return f.u
getw(f::PlanarFlow)                         = return f.w
getb(f::PlanarFlow)                         = return f.b
getall(f::PlanarFlow)                       = return f.u, f.w, f.b
getdim(f::PlanarFlowEmpty{N}) where { N }   = return N

# set-functions for the PlanarFlow structure
function setu!(f::PlanarFlow{T1,T2}, u::T1) where { T1, T2 <: Real}
    @assert length(f.u) == length(u) "The dimensionality of the current value of u and its new value do not match."
    f.u = u
end

function setw!(f::PlanarFlow{T1,T2}, w::T1) where { T1, T2 <: Real}
    @assert length(f.w) == length(w) "The dimensionality of the current value of w and its new value do not match."
    f.w = w
end

function setb!(f::PlanarFlow{T1,T2}, b::T2) where { T1, T2 <: Real }
    f.b = b
end

# custom Base function for the PlanarFlow structure
eltype(f::PlanarFlow{T1,T2}) where { T1 <: Real, T2 <: Real}                 = promote_type(T1, T2)
eltype(f::PlanarFlow{T1,T2}) where { T1 <: AbstractVector, T2 <: Real}       = promote_type(eltype(T1), T2)
eltype(::Type{PlanarFlow{T1,T2}}) where { T1 <: Real, T2 <: Real}            = promote_type(T1, T2)
eltype(::Type{PlanarFlow{T1,T2}}) where { T1 <: AbstractVector, T2 <: Real}  = promote_type(eltype(T1), T2)

size(f::PlanarFlow{T1,T2}) where { T1 <: Real, T2 <: Real}                   = 1
size(f::PlanarFlow{T1,T2}) where { T1 <: AbstractVector, T2 <: Real}         = length(f.u)
size(f::PlanarFlowEmpty{N}) where { N }                                      = return N

length(f::PlanarFlow{T1,T2}) where { T1 <: Real, T2 <: Real}                 = 1
length(f::PlanarFlow{T1,T2}) where { T1 <: AbstractVector, T2 <: Real}       = length(f.u)
length(f::PlanarFlowEmpty{N}) where { N }                                    = return N

# forward pass through the PlanarFlow function (multivariate input)
function _forward(f::PlanarFlow{T1,T2}, input::T1) where { T1, T2 <: Real }

    # fetch values
    u, w, b = getall(f)

    # calculate result
    result = u.*ones(size(u))
    result .*= tanh(dot(w, input) + b)
    result .+= input

    # return result
    return result

end
forward(f::PlanarFlow{T1,T2}, input::T1) where { T1, T2 <: Real } = _forward(f, input)
Broadcast.broadcasted(::typeof(forward), f::PlanarFlow{T1,T2}, input::AbstractVector{T1}) where { T1, T2 <: Real } = broadcast(_forward, Ref(f), input)

# forward pass through the PlanarFlow function (univariate input)
function _forward(f::PlanarFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 <: Real }

    # fetch values
    u, w, b = getall(f)

    # calculate result
    result = copy(u)
    result *= tanh(dot(w, input) + b)
    result += input

    # return result
    return result

end
forward(f::PlanarFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 <: Real } = _forward(f, input)
Broadcast.broadcasted(::typeof(forward), f::PlanarFlow{T1,T2}, input::AbstractVector{ <: Real }) where { T1 <: Real, T2 <: Real } = broadcast(_forward, Ref(f), input)

function _forward(f::PlanarFlow{T1,T2}, input) where { T1 <: Real, T2 <: Real }
    # function when the input is an array with 1 element
    @assert length(input) == 1 "Something is wrong with the dimensionality of the input to the PlanarFlow flow."
    return forward(f, input[1])
end
forward(f::PlanarFlow{T1,T2}, input) where { T1 <: Real, T2 <: Real } = _forward(f, input)
Broadcast.broadcasted(::typeof(forward), f::PlanarFlow{T1,T2}, input::AbstractVector) where { T1 <: Real, T2 <: Real} = broadcast(_forward, Ref(f), input)

# inplace forward pass through the PlanarFlow function (multivariate input)
function forward!(output::T1, f::PlanarFlow{T1,T2}, input::T1) where { T1, T2 <: Real }

    # check dimensionality
    @assert length(output) == length(input) "The length of the preallocated vector does not seem to match the length of the input vector."

    # fetch values
    u, w, b = getall(f)

    # calculate result
    output .= u
    output .*= tanh(dot(w, input) + b)
    output .+= input

end

# jacobian of the PlanarFlow function (multivariate input)
function _jacobian(f::PlanarFlow{T1,T2}, input::T1) where { T1, T2 <: Real}

    # fetch values
    u, w, b = getall(f)

    # calculate result
    result = u*w'
    result .*= dtanh(dot(w, input) + b)
    @inbounds for k = 1:length(input)
        result[k,k] += 1.0
    end

    # return result
    return result

end
jacobian(f::PlanarFlow{T1,T2}, input::T1) where { T1, T2 <: Real } = _jacobian(f, input)
Broadcast.broadcasted(::typeof(jacobian), f::PlanarFlow{T1,T2}, input::AbstractVector{T1}) where { T1, T2 <: Real } = broadcast(_jacobian, Ref(f), input)

# jacobian of the PlanarFlow function (univariate input)
function _jacobian(f::PlanarFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 <: Real }

    # fetch values
    u, w, b = getall(f)

    # calculate result (optimized)
    result = u * w * dtanh(w * input + b) + 1

    # return result
    return result

end
jacobian(f::PlanarFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 <: Real } = _jacobian(f, input)
Broadcast.broadcasted(::typeof(jacobian), f::PlanarFlow{T1,T2}, input::AbstractVector{ <: Real}) where { T1 <: Real, T2 <: Real } = broadcast(_jacobian, Ref(f), input)

function _jacobian(f::PlanarFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 }
    # function when the input is an array with 1 element
    @assert length(input) == 1 "Something is wrong with the dimensionality of the input to the PlanarFlow flow."
    return jacobian(f, input[1])
end
jacobian(f::PlanarFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 } = _jacobian(f, input)
Broadcast.broadcasted(::typeof(jacobian), f::PlanarFlow{T1,T2}, input::AbstractVector) where { T1 <: Real, T2 <: Real } = broadcast(_jacobian, Ref(f), input)

# inplace jacobian of the PlanarFlow function (multivariate input)
function jacobian!(output::AbstractMatrix{T2}, f::PlanarFlow{T1,T2}, input::T1) where { T1, T2 <: Real}

    # check whether the dimensionality is correct
    @assert size(output) == (length(input), length(f.u)) "The dimensionality of the preallocated jacobian matrix seems incorrect."

    # fetch values
    u, w, b = getall(f)

    # calculate result
    for ku = 1:length(u)
        for kw = 1:length(w)
            output[ku,kw] = u[ku]*w[kw]
        end
    end
    output .*= dtanh(dot(w, input) + b)
    @inbounds for k = 1:length(input)
        output[k,k] += 1.0
    end

end

# determinant of the jacobian of the PlanarFlow function (multivariate input)
det_jacobian(f::PlanarFlow{T1,T2}, input::T1) where { T1, T2 <: Real} = det(jacobian(f, input))

# determinant of the jacobian of the PlanarFlow function (univariate input)
function det_jacobian(f::PlanarFlow{T, T}, input::T) where { T <: Real }

    # fetch values
    u, w, b = getall(f)

    # calculate result
    result = 1 + dot(u, w)*dtanh(dot(w, input) + b)

    # return result
    return result

end

# extra utility function (multivariate)
inv_jacobian(f::PlanarFlow{T1,T2}, input::T1) where { T1, T2 <: Real }       = inv(jacobian(f, input))
absdet_jacobian(f::PlanarFlow{T1,T2}, input::T1) where { T1, T2 <: Real }    = abs(det_jacobian(f, input))
logabsdet_jacobian(f::PlanarFlow{T1,T2}, input::T1) where { T1, T2 <: Real}  = logabsdet(jacobian(f, input))

# extra utility functions (univariate)
inv_jacobian(f::PlanarFlow{T,T}, input::T) where { T <: Real }               = 1.0 / jacobian(f, input)
absdet_jacobian(f::PlanarFlow{T,T}, input::T) where { T <: Real }            = abs(det_jacobian(f, input))
logabsdet_jacobian(f::PlanarFlow{T,T}, input::T) where { T <: Real}          = log(absdet_jacobian(f, input))

function inv_jacobian(f::PlanarFlow{T1,T2}, input) where { T1 <: Real, T2 <: Real }
    # function when the input is an array with 1 element
    @assert length(input) == 1 "Something is wrong with the dimensionality of the input to the PlanarFlow flow."
    return inv_jacobian(f, input[1])
end


#=
A diffusion equation with different species coupled through boundary conditions only.
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("coupledbc1d");

useLog("coupledbc1dindlog")

# Configuration setup
domain(1, grid=UNSTRUCTURED)
solverType(CG)
functionSpace(order=2)
timeStepper(EULER_IMPLICIT)

# Mesh
n = 40 # number of elements
mesh(LINEMESH, elsperdim=n, bids=2)

# Variables and BCs
flavor = index("flavor", range = [1,3])
c = variable("c", type=VAR_ARRAY, index = flavor)

boundary(c, 1, DIRICHLET,  ["0.5*cos(pi*t)^2",
                            "0.5 - (c[1]-c[3])/(2+(c[1]-c[3]))",
                            "(c[1] + c[2])/2"])
boundary(c, 2, NEUMANN,  [0,0,0])

testSymbol("v")

# Time interval and initial condition
T = 10.25;
timeInterval(T)
# initial(c, ["sin(pi*x)^2",
#             "sin(pi*2*x)^2",
#             "sin(pi*3*x)^2"])
initial(c, [0.5,0.5,0.5])
# initial(c, ["sin(pi*x)^2",
#             "0"])
# initial(c, "sin(pi*x)^2")

weakForm(c, "Dt(c[flavor]*v) + 0.1 * dot(grad(c[flavor]),grad(v))")

assemblyLoops(c, ["elements", flavor]);

# exportCode("coupledbcind1dcode") # uncomment to export generated code to a file
# importCode("coupledbc1dcode") # uncomment to import code from a file

solve(c)

finalize_finch()

##### Uncomment below to plot
# x = Finch.grid_data.allnodes[1,:];
# using Plots
# pyplot();
# display(plot([x x x], [c.values[1,:] c.values[2,:] c.values[3,:]], marker=:circle))

# display(plot([x x], [c.values[1,:] c.values[2,:]], marker=:circle))
# display(plot([x], [c.values[1,:]], marker=:circle))


#nearest neighbors tour implementation

function nearest_neighbor{T<:Real}(distmat::Matrix{T}, firstcity::Int = 0)
	numCities = size(distmat, 1)
	# put first city on path
	path = Vector{Int}()
	if firstcity == 0
		firstcity = rand(1:numCities)
	end
	push!(path, firstcity)
	# cities to visit
	citiesToVisit = setdiff(1:numCities,path)
	# nearest neighbor loop
	while !isempty(citiesToVisit)
		curCity = path[end]
    dists = distmat[citiesToVisit,curCity]
		_, nextInd = findmin(dists)
		nextCity = citiesToVisit[nextInd]
		push!(path, nextCity)
		deleteat!(citiesToVisit, nextInd)
	end

	return path
end


module TensorFactorizations
using TensorOperations
using LinearAlgebra
export tensorsvd, tensoreig, tensorsplit

"""
    tensorsplit(A, a, b; kwargs...)

Calls tensorsvd with the arguments given to it to decompose the given tensor
A with indices a on one side and indices b on the other.  It then splits
the diagonal matrix of singular values into two with a square root and
multiplies these weights into the isometric tensors.  Thus tensorsplit ends
up splitting A into two parts, which are then returned, possibly together
with auxiliary data such as a truncation error. If the keyword argument
hermitian=true, an eigenvalue decomposition is used in stead of an SVD. All
the keyword arguments are passed to either tensorsvd or tensoreig.

See tensorsvd and tensoreig for further documentation.
"""
function tensorsplit(args...; kwargs...)
    # Find the keyword argument hermitian.
    # TODO This is awful, why do I have to do this?
    hermitian = false
    for (key, value) in kwargs
        key == :hermitian && (hermitian = value)
    end

    if hermitian
        res = tensoreig(args...; kwargs...)
        S, U = res[1:2]
        Vt_perm = [ndims(U), (1:ndims(U)-1)...]
        Vt = conj!(tensorcopy(U, collect(1:ndims(U)), Vt_perm))
        S = Diagonal(S)
        if !isposdef(S)
            S = complex.(S)
        end
        auxdata = res[3:end]
    else
        res = tensorsvd(args...; kwargs...)
        U, S, Vt = res[1:3]
        S = Diagonal(S)
        auxdata = res[4:end]
    end
    S_sqrt = sqrt.(S)
    A1 = tensorcontract(U, (1:ndims(U)-1..., :a), S_sqrt, (:a, :b))
    A2 = tensorcontract(S_sqrt, (:b, :a), Vt, (:a, 1:ndims(Vt)-1...))
    return A1, A2, auxdata...
end

"""
    tensoreig(A, a, b; chis=nothing, eps=0,
              return_error=false, print_error=false,
              break_degenerate=false, degeneracy_eps=1e-6,
              norm_type=:frobenius, hermitian=false)

Finds the "right" eigenvectors and eigenvalues of A. The indices of A are
permuted so that the indices listed in the Array/Tuple a are on the "left"
side and indices listed in b are on the "right".  The resulting tensor is
then reshaped to a matrix, and eig is called on this matrix to get the vector
of eigenvalues E and matrix of eigenvectors U. Finally, U is reshaped to
a tensor that has as its last index the one that enumerates the eigenvectors
and the indices in a as its first indices.

Truncation and error printing work as with tensorsvd.

Note that no iterative techniques are used, which means that choosing to
truncate provides no performance benefits: All the eigenvalues are computed
in any case.

The keyword argument hermitian (false by default) tells the algorithm
whether the reshaped matrix is Hermitian or not. If hermitian=true, then A =
U*diagm(E)*U' up to the truncation error.

Output is E, U, and possibly error, if return_error=true. Here E is a
vector of eigenvalues values and U[:,...,:,k] is the kth eigenvector.
"""
function tensoreig(A, a, b; chis=nothing, eps=0, return_error=false,
                   print_error=false, break_degenerate=false,
                   degeneracy_eps=1e-6, norm_type=:frobenius,
                   hermitian=false)
    # Create the matrix and decompose it.
    A, shp_a, shp_b = to_matrix(A, a, b; return_tensor_shape=true)
    if hermitian
        A = Hermitian((A+A')/2)
    end
    fact = eigen(A)
    E, U = fact.values, fact.vectors
    # Sort the by largest magnitude eigenvalue first.
    perm = sortperm(abs.(E), rev=true)
    if perm != collect(1:length(E))
        E = E[perm]
        U = U[:,perm]
    end

    # Find the dimensions to truncate to and the error caused in doing so.
    chi, error = find_trunc_dim(E, chis, eps, break_degenerate, degeneracy_eps,
                                norm_type)

    # Truncate
    E = E[1:chi]
    U = U[:,1:chi]

    if print_error
        println("Relative truncation error ($norm_type norm) in eig: $error")
    end

    # Reshape U to a tensor with a shape matching the shape of A and
    # return.
    dim = size(E)[1]
    U_tens = reshape(U, shp_a..., dim)
    retval = (E, U_tens)
    if return_error
        retval = (retval..., error)
    end
    return retval
end

"""
    tensorsvd(A, a, b;
              chis=nothing, eps=0,
              return_error=false, print_error=false,
              break_degenerate=false, degeneracy_eps=1e-6,
              norm_type=:frobenius)

Singular valued decomposes a tensor A. The indices of A are
permuted so that the indices listed in the Array/Tuple a are on the "left"
side and indices listed in b are on the "right".  The resulting tensor is
then reshaped to a matrix, and this matrix is SVDed into U*diagm(S)*Vt.
Finally, the unitary matrices U and Vt are reshaped to tensors so that
they have a new index coming from the SVD, for U as the last index and for
Vt as the first, and U has indices a as its first indices and V has
indices b as its last indices.

If eps>0 then the SVD may be truncated if the relative error can be kept
below eps. For this purpose different dimensions to truncate to can be tried,
and these dimensions should be listed in chis. If chis is nothing (the
default) then the full range of possible dimensions is tried. If
break_degenerate=false (the default) then the truncation never cuts between
degenerate singular values. degeneracy_eps controls how close the values need
to be to be considered degenerate.

norm_type specifies the norm used to measure the error. This defaults to
:frobenius, which means that the error measured is the Frobenius norm of the
difference between A and the decomposition, divided by the Frobenius norm of
A.  This is the same thing as the 2-norm of the singular values that are
truncated out, divided by the 2-norm of all the singular values. The other
option is :trace, in which case a 1-norm is used instead.

If print_error=true the truncation error is printed. The default is false.

If return_error=true then the truncation error is also returned.
The default is false.

Note that no iterative techniques are used, which means choosing to truncate
provides no performance benefits: The full SVD is computed in any case.

Output is U, S, Vt, and possibly error. Here S is a vector of
singular values and U and Vt are isometric tensors (unitary if the matrix
that is SVDed is square and there is no truncation) such that  U*diag(S)*Vt =
A, up to truncation errors.
"""
function tensorsvd(A, a, b; chis=nothing, eps=0, return_error=false,
                   print_error=false, break_degenerate=false,
                   degeneracy_eps=1e-6, norm_type=:frobenius)
    # Create the matrix and SVD it.
    A, shp_a, shp_b = to_matrix(A, a, b; return_tensor_shape=true)
    fact = svd(A)
    U, S, Vt = fact.U, fact.S, fact.Vt

    # Find the dimensions to truncate to and the error caused in doing so.
    chi, error = find_trunc_dim(S, chis, eps, break_degenerate, degeneracy_eps,
                                norm_type)

    # Truncate
    S = S[1:chi]
    U = U[:,1:chi]
    Vt = Vt[1:chi,:]

    if print_error
        println("Relative truncation error ($norm_type norm) in SVD: $error")
    end

    # Reshape U and V to tensors with shapes matching the shape of A and
    # return.
    dim = size(S)[1]
    U_tens = reshape(U, shp_a..., dim)
    Vt_tens = reshape(Vt, dim, shp_b...)
    retval = (U_tens, S, Vt_tens)
    if return_error
        retval = (retval..., error)
    end
    return retval
end

"""
Format the bond dimensions listed in chis to a standard format.
"""
function format_trunc_chis(v, chis, eps)
    max_dim = length(v)
    if chis == nothing
        if eps > 0
            # Try all possible chis.
            chis = collect(1:max_dim)
        else
            # No truncation.
            chis = [max_dim]
        end
    else
        if isa(chis, Number)
            # Wrap an individual number in an Array.
            chis = [chis]
        else
            # Make sure chis is an Array, and not, say, a tuple.
            chis = collect(chis)
        end
        if eps == 0
            chis = [maximum(chis)]
        else
            sort!(chis)
        end
    end
    # If some of the chis are larger than max_dim, get rid of them.
    for (i, chi) in enumerate(chis)
        if chi >= max_dim
            chis[i] = max_dim
            chis = chis[1:i]
            break
        end
    end
    return chis
end

"""
Transpose A so that the indices listed in a are on the left and the indices
listed in b on the right, and reshape A into a matrix.

a and b should be arrays of Integers that together include the numbers from 1
to ndims(A). Alternatively they can be just individual Integers, if only one
index is left on one side.

If return_tensor_shape is true (by default it's not) return, in addition to the
matrix, the shape of the tensor after the transpose but before the reshape.
"""
function to_matrix(A, a, b; return_tensor_shape=false)
    # Make sure a and b are Arrays.
    if isa(a, Number)
        a = [a]
    elseif !isa(a, Array)
        a = collect(a)
    end
    if isa(b, Number)
        b = [b]
    elseif !isa(b, Array)
        b = collect(b)
    end

    # Permute the indices of A to the right order
    perm = vcat(a, b)
    A = tensorcopy(A, collect(1:ndims(A)), perm)
    # The lists shp_a and shp_b list the dimensions of the bonds in a and b
    shp = size(A)
    shp_a = shp[1:length(a)]
    shp_b = shp[end+1-length(b):end]

    # Compute the dimensions of the the matrix that will be formed when
    # indices of a and b are joined together.
    dim_a = prod(shp_a)
    dim_b = prod(shp_b)

    A = reshape(A, dim_a, dim_b)

    if return_tensor_shape
        return A, shp_a, shp_b
    else
        return A
    end
end

"""
Finds the dimension to which v should be truncated, and the error caused in
this truncation. See documentation for tensorsvd for the meaning of the
different arguments.
"""
function find_trunc_dim(v, chis, eps,
                        break_degenerate=false, degeneracy_eps=1e-6,
                        norm_type=:frobenius)
    # Put chis in a standard format.
    chis = format_trunc_chis(v, chis, eps)
    v = abs.(v)
    if !issorted(v, rev=true)
        sort!(v, rev=true)
    end

    if norm_type==:frobenius
        v = v.^2
        eps = eps^2
    elseif norm_type == :trace
        # Nothing to be done
    else
        throw(ArgumentError("Unknown norm_type $norm_type."))
    end
    sum_all = sum(v)
    # Find the smallest chi for which the error is small enough.
    # If none is found, use the largest chi.
    if sum_all != 0
        for i in 1:length(chis)
            chi = chis[i]
            if !break_degenerate
                # Make sure that we don't break degenerate singular values by
                # including one but not the other by decreasing chi if
                # necessary.
                while 0 < chi < length(v)
                    last_in = v[chi]
                    last_out = v[chi+1]
                    rel_diff = abs(last_in - last_out)/last_in
                    if rel_diff < degeneracy_eps
                        chi -= 1
                    else
                        break
                    end
                end
            end
            sum_disc = sum(v[chi+1:end])
            error = sum_disc/sum_all
            if error <= eps
                break
            end
        end
        if norm_type==:frobenius
            error = sqrt(error)
        end
    else
        error = 0
        chi = minimum(chis)
    end
    return chi, error
end

end


# This code is part of QuantumCircuits.
#
# (C) Copyright Rafał Pracht 2022.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

module Math

using LinearAlgebra

export unitary_error, matrix_norm, eye, observe_unitary_error,
       min_observe_unitary_error

"Calculate the unitary matrix error"
unitary_error(exp_unit, mat_unit) = _unitary_error(exp_unit, mat_unit)
function unitary_error(exp_unit, mat_unit, exp_scale, mat_scale)
    return _unitary_error(exp_unit, mat_unit, x -> x * exp_scale, x -> x * mat_scale)
end
function _unitary_error(exp_unit, mat_unit, fa=identity, fb=identity)
    return sum(((a, b),) -> abs2(fa(a) - fb(b)), zip(exp_unit, mat_unit))
end

"Calculate the observe unitary matrix error, the error which we can obserwe
from quantum state"
function observe_unitary_error(exp_unit, mat_unit, index=1)
    # Remove phase
    exp_scale = cis(-angle(exp_unit[index]))
    mat_scale = cis(-angle(mat_unit[index]))

    return unitary_error(exp_unit, mat_unit, exp_scale, mat_scale)
end

"Calculate the minimum observe unitary matrix error."
function min_observe_unitary_error(exp_unit, mat_unit)
    l = length(exp_unit)
    return minimum(i -> observe_unitary_error(exp_unit, mat_unit, i), 1:l)
end

"Calculate the matrix norm"
function matrix_norm(mat)
    return abs(tr(adjoint(mat) * mat))
end

"Return the identity matrix."
eye(n) = Matrix(I, n, n)

end  # module Math


export constructtriad
export constructtorus
export constructsphere
export constructcylinder
export constructbox
export constructwhirl
export constructframe
export getplane
export constructfiber
export constructtwospinor
export transformg

"""
    constructtriad(q [; length])

Construct a triad with the given configuration `q` and `length`.
"""
function constructtriad(q::Biquaternion; length::Float64 = 1.0)
    map(x -> x + gettranslation(q), rotate(ℝ³.([[0; 0; 0], [length; 0; 0],
                                                [0; 0; 0], [0; length; 0],
                                                [0; 0; 0], [0; 0; length]]),
                                           getrotation(q)))
end

"""
    constructtorus(q, r, R)

Construct a torus of revolution with the given configuration `q`, the smaller radius `r`
and the bigger radius `R`.
"""
function constructtorus(q::Biquaternion,
                        r::Real,
                        R::Real;
                        segments::Int = 36)
    array = Array{ℝ³,2}(undef, segments, segments)
    for i in 1:segments
        for j in 1:segments
            ϕ = i * 2pi / (segments - 1)
            θ = j * 2pi / (segments - 1)
            x₁ = (R + r * cos(ϕ)) * cos(θ)
            x₂ = (R + r * cos(ϕ)) * sin(θ)
            x₃ = r * sin(ϕ)
            array[i, j] = ℝ³(x₁, x₂, x₃)
        end
    end
    map(x -> x + gettranslation(q), rotate(array, getrotation(q)))
end

"""
    constructsphere(q, radius)

Construct a sphere with the given configuration `q` and `radius`.
"""
function constructsphere(q::Biquaternion,
                         radius::Real;
                         segments::Int = 36)
    array = Array{ℝ³,2}(undef, segments, segments)
    lspace = collect(range(float(pi), stop = float(-pi), length = segments))
    for i in 1:segments
        for j in 1:segments
            ϕ = lspace[i]
            θ = lspace[j] / 2
            array[j, i] = ℝ³(Cartesian(Geographic(radius, ϕ, θ)))
        end
    end
    map(x -> x + gettranslation(q), rotate(array, getrotation(q)))
end

"""
    constructcylinder(q, height, radius)

Construct a cylinder with the given configuration `q`, `height` and `radius`.
"""
function constructcylinder(q::Biquaternion,
                           height::Float64,
                           radius::Float64;
                           segments::Int = 36)
    array = Array{ℝ³,2}(undef, segments, segments)
    lspace = collect(range(float(-pi), stop = float(pi), length = segments))
    for i in 1:segments
        for j in 1:segments
            α = lspace[j]
            x₁ = radius * cos(α)
            x₂ = radius * sin(α)
            x₃ = (i - 1) * height / (segments - 1)
            array[i, j] = ℝ³(x₁, x₂, x₃)
        end
    end
    map(x -> x + gettranslation(q), rotate(array, getrotation(q)))
end

"""
    constructbox(p, q, s)

Construct a box with the given translation `p`, rotation `q` and scale `s`.
"""
function constructbox(p::ℝ³, q::Quaternion, s::ℝ³)
    x, y, z = vec(s) ./ 2
    points = [ℝ³(-x, -y, -z) ℝ³(x, -y, -z) ℝ³(-x, -y, -z) ℝ³(x, -y, -z);
              ℝ³(-x, -y, +z) ℝ³(x, -y, +z) ℝ³(-x, -y, +z) ℝ³(x, -y, +z);
              ℝ³(-x, +y, +z) ℝ³(x, +y, +z) ℝ³(-x, +y, +z) ℝ³(x, +y, +z);
              ℝ³(-x, +y, -z) ℝ³(x, +y, -z) ℝ³(-x, +y, -z) ℝ³(x, +y, -z)]
    map(x -> x + p, rotate(points, q))
end

"""
    getactions(point,
               s2tos3map,
               s2tos2map;
               [top, [bottom, [s3rotation, [segments [, scale]]]]])

Take the integral along the Hopf fiber for getting a linear action space
under stereographic projection, with the given 'point' in the base space,
a map from the base space to the total space `s2tos3map`, a conformal map that
takes the base space to itself `s2tos2map`, `top` and `bottom` U(1) actions,
S³ rotation `s3rotation`, the number of `segments` and `scale`.
"""
function getactions(point::S²,
                    s2tos3map,
                    s2tos2map;
                    top::S¹ = U1(-pi),
                    bottom::S¹ = U1(pi),
                    s3rotation::S³ = Quaternion(1, 0, 0, 0),
                    segments::Int = 30,
                    scale::Real = 1.0)
    segments2 = 1000
    array = Array{Float64,2}(undef, segments2, 2)
    minaction = -π
    maxaction = π
    lspace = collect(range(minaction, stop = maxaction, length = segments2))
    array[:, begin] = lspace
    accumulator = 0
    getp(point, α, rotation, scale) = compressedλmap(rotate(S¹action(s2tos3map(s2tos2map(point)), U1(α)), rotation)) * scale
    α₁ = lspace[begin]
    p₁ = getp(point, α₁, s3rotation, scale)
    for i in 1:segments2
        α₂ = lspace[i]
        p₂ = getp(point, α₂, s3rotation, scale)
        accumulator += norm(p₂ - p₁)
        array[i, end] = accumulator
        α₁ = α₂
        p₁ = p₂
    end

    # decide based on cumulative length and then pick the corresponding actions from the array
    totalspan = maxaction - minaction
    normalizedtop = angle(top) / totalspan
    normalizedbottom = angle(bottom) / totalspan
    normalizedsegmentlength = (normalizedbottom - normalizedtop) / (segments - 1)
    actions = Array{Float64,1}(undef, segments)
    for i in 1:segments
        target = (normalizedtop + normalizedsegmentlength * (i - 1)) * accumulator
        for j in 1:segments2
            if array[j, end] ≥ target
                actions[i] = array[j, begin]
                break
            end
        end
    end
    actions
end

"""
    transformg(p, g₁, g₂, segments)

Transform G-acions by finding and applying affine parameters with the given point `p`
in the principal bundle, beginning gauge potential `g₁`, ending gauge potential `g₂`
and the number of segments.
"""
function transformg(p::ComplexPlane, g₁::U1, g₂::U1, segments::Int)
    intervals = 1000
    array = Array{Float64,2}(undef, intervals, 2)
    minaction, maxaction = 0, 2π
    lspace = collect(range(minaction, stop = maxaction, length = intervals))
    array[:, begin] = lspace
    accumulator = 0
    α₁ = lspace[begin]
    z = λ⁻¹map(πmap(Quaternion(p)))
    p₁ = compressedλmap(S¹action(z, U1(α₁)))
    array[begin, end] = 0
    for i in 2:intervals
        α₂ = lspace[i]
        p₂ = compressedλmap(S¹action(p, U1(α₂)))
        accumulator += norm(p₂ - p₁)
        array[i, end] = accumulator
        α₁ = α₂
        p₁ = p₂
    end

    # decide based on cumulative length and then pick the corresponding actions from the array
    totalspan = maxaction - minaction
    n₁ = angle(g₁) / totalspan
    n₂ = angle(g₂) / totalspan
    normalizedsegmentlength = (n₂ - n₁) / (segments - 1)
    actions = Array{U1,1}(undef, segments)
    for i in 1:segments
        target = (n₁ + normalizedsegmentlength * (i - 1)) * accumulator
        for j in 1:intervals
            if array[j, end] ≥ target
                actions[i] = U1(array[j, begin])
                break
            end
        end
    end
    actions
end

"""
    constructwhirl(points, gauge1, gauge2, configuration, segments)

Construct a whirl with the given `points`, `gauge1`, `gauge2`, `configuration` and `segments`.
"""
function constructwhirl(points::Array{ComplexPlane,1},
                        gauge1::Array{U1,1},
                        gauge2::Array{U1,1},
                        configuration::Biquaternion,
                        segments::Int)
    array = Array{ℝ³,2}(undef, segments, length(points))
    for (j, p) in enumerate(points)
        g₁ = gauge1[j]
        g₂ = gauge2[j]
        g = transformg(p, g₁, g₂, segments)
        for (i, α) in enumerate(g)
            array[i, j] = compressedλmap(S¹action(p, α))
        end
    end
    applyconfig(array, configuration)
end

"""
    constructfiber(point,
                   s2tos3map,
                   s2tos2map;
                   [radius, [top, [bottom, [s3rotation, [config, [segments [, scale]]]]]]])

Construct a Hopf fiber with the given `point` in the base space, , a map from the base space to
the total space `s2tos3map`, a conformal map that takes the base space to itself `s2tos2map`,
fiber sectional `radius`, `top` and `bottom` U(1) actions, S³ rotation `s3rotation`,
configuration `config`, the number of `segments` and `scale`.
"""
function constructfiber(point::S²,
                        s2tos3map,
                        s2tos2map;
                        radius::Float64 = 0.02,
                        top::S¹ = U1(-pi),
                        bottom::S¹ = U1(pi),
                        s3rotation::S³ = Quaternion(1, 0, 0, 0),
                        config::Biquaternion = Biquaternion(ℝ³(0, 0, 0)),
                        segments::Int = 30,
                        scale::Real = 1.0)
    segments2 = 10
    α₁ = U1(min(angle(top), angle(bottom)))
    α₂ = U1(max(angle(top), angle(bottom)))
    circle = [ℝ³(cos(α), sin(α), 0) for α in range(0, stop = 2π, length = segments2)] .* radius
    actions = getactions(point, s2tos3map, s2tos2map, top = α₁, bottom = α₂, s3rotation = s3rotation, segments = segments, scale = scale)
    getp(point, α, rotation, scale) = compressedλmap(rotate(S¹action(s2tos3map(s2tos2map(point)), U1(α)), rotation)) * scale
    array = Array{ℝ³,2}(undef, segments, segments2)
    n = ℝ³(0, 0, 1)
    for (i, α) in enumerate(actions)
        p₁ = getp(point, α, s3rotation, scale)
        p₂ = getp(point, α + α * 1e-2, s3rotation, scale)
        w = normalize(p₂ - p₁)
        q = Biquaternion(getrotation(n, w), p₁)
        array[i, :] = applyconfig(circle, q)
    end
    applyconfig(array, config)
end

"""
    constructtwospinor(point, gauge1, gauge2, configuration, radius, segments1, segments2)

Construct a 2-spinor with the given `point` in the principal bundle, `gauge1`, `gauge2`,
`configuration`, `radius`, `segments1` and `segments2`.
"""
function constructtwospinor(point::ComplexPlane, gauge1::U1, gauge2::U1, configuration::Biquaternion,
                            radius::Float64, segments1::Int, segments2::Int)
    angles = angle.([gauge1; gauge2])
    g₁, g₂ = U1(min(angles...)), U1(max(angles...))
    circle = [ℝ³(cos(α), sin(α), 0) for α in range(0, stop = 2π, length = segments2)] .* radius
    actions = transformg(point, g₁, g₂, segments1)
    getp(p::ComplexPlane, g::U1) = compressedλmap(S¹action(p, g))
    array = Array{ℝ³,2}(undef, segments1, segments2)
    n = ℝ³(0, 0, 1)
    ϵ = 1e-2
    for (i, g) in enumerate(actions)
        p₁ = getp(point, g)
        p₂ = getp(point, U1(angle(g) + ϵ))
        w = normalize(p₂ - p₁)
        q = Biquaternion(getrotation(n, w), p₁)
        array[i, :] = applyconfig(circle, q)
    end
    applyconfig(array, configuration)
end

"""
    constructframe(section, configuration, segments)

Construct a frame that is like an S³ cross-section under stereographic projection with the
given `section`, `configuration` and `segments`.
"""
function constructframe(section::Any,
                        configuration::Biquaternion = Biquaternion(ℝ³(0, 0, 0)),
                        segments::Int = 36)
    array = Array{ℝ³,2}(undef, segments, segments)
    lspaceϕ = collect(range(float(-pi), stop = float(pi), length = segments))
    lspaceθ = collect(range(float(pi / 2), stop = float(-pi / 2), length = segments))
    for (i, θ) in enumerate(lspaceθ)
        for (j, ϕ) in enumerate(lspaceϕ)
            p = Geographic(1, ϕ, θ)
            array[i, j] = compressedλmap(section(p))
        end
    end
    applyconfig(array, configuration)
end

getplane(p::Array{Float64}, q::Quaternion, s::Array{Float64}) = begin
    plane = Array{Float64,3}(undef, 2, 2, 3)
    plane[1, 1, :] = vec(rotate(ℝ³([0.0; -1.0; 1.0]), q)) .* s + p
    plane[1, 2, :] = vec(rotate(ℝ³([0.0; 1.0; 1.0]), q)) .* s + p
    plane[2, 1, :] = vec(rotate(ℝ³([0.0; -1.0; -1.0]), q)) .* s + p
    plane[2, 2, :] = vec(rotate(ℝ³([0.0; 1.0; -1.0]), q)) .* s + p
    plane
end


function hess{T <: Number}(
    f::Function,
    x::T;
    method::Symbol = :finite,
    direction::Symbol = :central,
) # -> T
    if method == :finite
        return finitehess(f, x, direction)
    else
        throw(ArgumentError("method must be :finite"))
    end
end

function hess!{T <: Number}(
    H::Matrix{T},
    f::Function,
    x::Vector{T};
    method::Symbol = :finite,
    direction::Symbol = :central,
) # -> Matrix{T}
    if method == :finite
        finitehess!(H, f, x, direction)
    else
        throw(ArgumentError("method must be :finite"))
    end
    return H
end

function hess{T <: Number}(
    f::Function,
    x::Vector{T};
    method::Symbol = :finite,
    direction::Symbol = :central,
) # -> Matrix{T}
    n = length(x)
    H = Array(T, n, n)
    hess!(H, f, x, method = method, direction = direction)
    return H
end


"""
    getARXregressor(y::AbstractVector,u::AbstractVecOrMat, na, nb)
Returns a shortened output signal `y` and a regressor matrix `A` such that the least-squares ARX model estimate of order `na,nb` is `y\\A`
Return a regressor matrix used to fit an ARX model on, e.g., the form
`A(z)y = B(z)f(u)`
with output `y` and input `u` where the order of autoregression is `na` and
the order of input moving average is `nb`
# Example
Here we test the model with the Function `f(u) = √(|u|)`
```julia
A     = [1,2*0.7*1,1] # A(z) coeffs
B     = [10,5] # B(z) coeffs
u     = randn(100) # Simulate 100 time steps with Gaussian input
y     = filt(B,A,u)
yr,A  = getARXregressor(y,u,3,2) # We assume that we know the system order 3,2
x     = A\\yr # Estimate model polynomials
plot([yr A*x], lab=["Signal" "Prediction"])
```
For nonlinear ARX-models, see [BasisFunctionExpansions.jl](https://github.com/baggepinnen/BasisFunctionExpansions.jl/). See also `arx`
"""
function getARXregressor(y::AbstractVector,u::AbstractVecOrMat, na, nb)
    length(nb) == size(u,2) || throw(ArgumentError("Length of nb must equal number of input signals"))
    m    = max(na,maximum(nb))+1 # Start of yr
    @assert m >= 1
    n    = length(y) - m + 1 # Final length of yr
    @assert n <= length(y)
    A    = toeplitz(y[m:m+n-1],y[m:-1:m-na])
    @assert size(A,2) == na+1
    y    = A[:,1] # extract yr
    A    = A[:,2:end]
    for i = 1:length(nb)
        nb[i] <= 0 && continue
        s = m-1
        A = [A toeplitz(u[s:s+n-1,i],u[s:-1:s-nb[i]+1,i])]
    end
    return y,A
end

function getARregressor(y::AbstractVector, na)
    m    = na+1 # Start of yr
    n    = length(y) - m + 1 # Final length of yr
    A    = toeplitz(y[m:m+n-1],y[m:-1:m-na])
    @assert size(A,2) == na+1
    y    = A[:,1] # extract yr
    A    = A[:,2:end]
    return y,A
end

@userplot Find_na
"""
    find_na(y::AbstractVector,n::Int)
Plots the RMSE and AIC For model orders up to `n`. Useful for model selection
"""
find_na
@recipe function find_na(p::Find_na)
    y,n = p.args[1:2]
    error = zeros(n,2)
    for i = 1:n
        yt,A = getARXregressor(y,0y,i,0)
        e = yt-A*(A\yt)
        error[i,1] = rms(e)
        error[i,2] = aic(e,i)
    end
    layout --> 2
    title --> ["RMS error" "AIC"]
    seriestype --> :scatter
    @series begin
        error
    end
end

@userplot Find_nanb
"""
    find_nanb(y::AbstractVector,u,na,nb)
Plots the RMSE and AIC For model orders up to `n`. Useful for model selection
"""
find_nanb
@recipe function find_nanb(p::Find_nanb; logrms=false)
    y,u,na,nb = p.args[1:4]
    error = zeros(na,nb, 2)
    for i = 1:na, j = 1:nb
        yt,A = getARXregressor(y,u,i,j)
        e = yt-A*(A\yt)
        error[i,j,1] = logrms ? log10.(rms(e)) : rms(e)
        error[i,j,2] = aic(e,i+j)
    end
    layout --> 2
    seriestype --> :heatmap
    xticks := (1:nb, 1:nb)
    yticks := (1:na, 1:na)
    ylabel := "na"
    xlabel := "nb"
    @series begin
        title := "RMS error"
        subplot := 1
        error[:,:,1]
    end
    @series begin
        title := "AIC"
        subplot := 2
        error[:,:,2]
    end
end

"""
    Gtf, Σ = arx(h,y, u, na, nb; λ = 0, estimator=\\)

Fit a transfer Function to data using an ARX model and equation error minimization.
`nb` and `na` are the length of the numerator and denominator polynomials. `h` is the sample time of the data. `λ > 0` can be provided for L₂ regularization. `estimator` defaults to \\ (least squares), alternatives are `estimator = tls` for total least-squares estimation. `arx(Δt,yn,u,na,nb, estimator=wtls_estimator(y,na,nb)` is potentially more robust in the presence of heavy measurement noise.
The number of free parameters is `na-1+nb`
`Σ` is the covariance matrix of the parameter estimate. See `bodeconfidence` for visualiztion of uncertainty.

Supports MISO estimation by supplying a matrix `u` where times is first dim, with nb = [nb₁, nb₂...]
"""
function arx(h,y::AbstractVector, u::AbstractVecOrMat, na, nb; λ = 0, estimator=\)
    all(nb .<= na) || throw(DomainError(nb,"nb must be <= na"))
    na >= 1 || throw(ArgumentError("na must be positive"))
    y_train, A = getARXregressor(y,u, na, nb)
    w = ls(A,y_train,λ,estimator)
    a,b = params2poly(w,na,nb)
    model = tf(b,a,h)
    Σ = parameter_covariance(y_train, A, w, λ)
    return model, Σ
end

function ar(h,y::AbstractVector, na; λ = 0, estimator=\)
    na >= 1 || throw(ArgumentError("na must be positive"))
    y_train, A = getARregressor(y, na)
    w = ls(A,y_train,λ,estimator)
    a,b = params2poly(w,na)
    model = tf(b,a,h)
    Σ = parameter_covariance(y_train, A, w, λ)
    return model, Σ
end

function ls(A,y,λ=0,estimator=\)
    if λ == 0
        w = estimator(A,y)
    else
        w = estimator([A; λ*I], [y;zeros(size(A,2))])
    end
    w
end

"""
G, Gn = plr(h,y,u,na,nb,nc; initial_order = 20)

Perform pseudo-linear regression to estimate a model on the form
`Ay = Bu + Cw`
The residual sequence is estimated by first estimating a high-order arx model, whereafter the estimated residual sequence is included in a second estimation problem. The return values are the estimated system model, and the estimated noise model. `G` and `Gn` will always have the same denominator polynomial.
"""
function plr(h,y,u,na,nb,nc; initial_order = 20)
    all(nb .<= na) || throw(DomainError(nb,"nb must be <= na"))
    na >= 1 || throw(ArgumentError("na must be positive"))
    na -= 1
    y_train, A = getARXregressor(y,u,initial_order,initial_order)
    w1 = A\y_train
    yhat = A*w1
    ehat = yhat - y_train
    ΔN = length(y)-length(ehat)
    y_train, A = getARXregressor(y[ΔN+1:end-1],[u[ΔN+1:end-1,:] ehat[1:end-1]],na,[nb;nc])
    w = A\y_train
    a,b = params2poly(w,na,nb)
    model = tf(b,a,h)
    c = w[na+sum(nb)+1:end]
    noise_model = tf(c,a,h)
    model, noise_model
end
# Helper constructor to make a MISO system after MISO arx estimation
function ControlSystems.tf(b::AbstractVector{<:AbstractVector{<:Number}}, a::AbstractVector{<:Number}, h)
    tfs = map(b) do b
        tf(b,a,h)
    end
    hcat(tfs...)
end

"""
wtls_estimator(y,na,nb)
Create an estimator function for estimation of arx models in the presence of measurement noise.
"""
function wtls_estimator(y,na,nb)
    rowQ = [diagm(0=>[ones(na);zeros(nb);1]) for i = 1:length(y)-(na)]
    Qaa,Qay,Qyy = rowcovariance(rowQ)
    (A,y)->wtls(A,y,Qaa,Qay,Qyy)
end

"""
    a,b = params2poly(params,na,nb)
Used to get numerator and denominator polynomials after arx fitting
"""
function params2poly(w,na,nb)
    a = [1; -w[1:na]]
    w = w[na+1:end]
    b = map(nb) do nb
        b = w[1:nb]
        w = w[nb+1:end]
        b
    end
    a,b
end
function params2poly(w,na)
    a = [1; -w]
    b = 1
    a,b
end
"""
w, a, b = params(G::TransferFunction)
w = [a;b]
"""
function params(G::TransferFunction)
    am, bm = -denpoly(G)[1].a[1:end-1], G.matrix[1].num.a
    wm     = [am; bm]
    wm, am, bm
end

"""
    Σ = parameter_covariance(y_train, A, w, λ=0)
"""
function parameter_covariance(y_train, A, w, λ=0)
    σ² = var(y_train .- A*w)
    iATA = if λ == 0
        inv(A'A)
    else
        ATA = A'A
        ATAλ = ATA + λ*I
        ATAλ\ATA/ATAλ
    end
    iATA = (iATA+iATA')/2
    Σ = σ²*iATA + sqrt(eps())*Matrix(LinearAlgebra.I,size(iATA))
end

"""
TransferFunction(T::Type{<:AbstractParticles}, G::TransferFunction, Σ, N=500)

Create a `TransferFunction` where the coefficients are `Particles` from [`MonteCarloMeasurements.jl`](https://github.com/baggepinnen/MonteCarloMeasurements.jl) that can represent uncertainty.
# Example
```julia
using MonteCarloMeasurements
Gls,Σls = arx(Δt,y,u,na,nb)
Glsp    = TransferFunction(Particles, Gls, Σls)
w       = exp10.(LinRange(-3,log10(π/Δt),100))
mag     = bode(Glsp,w)[1][:]
errorbarplot(w,mag,0.01; yscale=:log10, xscale=:log10, layout=3, subplot=1, lab="ls")

See full example [here](https://github.com/baggepinnen/MonteCarloMeasurements.jl/blob/master/examples/controlsystems.jl)
```
"""
function ControlSystems.TransferFunction(T::Type{<:MonteCarloMeasurements.AbstractParticles}, G::TransferFunction, Σ, N=500)
      wm, am, bm = ControlSystemIdentification.params(G)
      na,nb      = length(am), length(bm)
      p          = T(N, MvNormal(wm, Σ))
      a,b        = ControlSystemIdentification.params2poly(p,na,nb)
      arxtf      = tf(b,a,G.Ts)
end

function ControlSystems.TransferFunction(T::Type{<:MonteCarloMeasurements.AbstractParticles}, G::TransferFunction, p::AbstractMatrix)
      wm, am, bm = ControlSystemIdentification.params(G)
      na,nb      = length(am), length(bm)
      p          = T(p .+ wm')
      a,b        = ControlSystemIdentification.params2poly(p,na,nb)
      arxtf      = tf(b,a,G.Ts)
end

function DSP.filt(tf::ControlSystems.TransferFunction, y)
    b,a = numvec(tf)[], denvec(tf)[]
    filt(b,a,y)
end

"""
    bodeconfidence(arxtf::TransferFunction, Σ::Matrix, ω = logspace(0,3,200))
Plot a bode diagram of a transfer function estimated with [`arx`](@ref) with confidence bounds on magnitude and phase.
"""
bodeconfidence

@userplot BodeConfidence

@recipe function BodeConfidence(p::BodeConfidence)
    arxtfm = p.args[1]
    Σ      = p.args[2]
    ω      = length(p.args) >= 3 ? p.args[3] : exp10.(LinRange(-2,3,200))
    L      = cholesky(Hermitian(Σ)).L
    wm, am, bm = params(arxtfm)
    na,nb  = length(am), length(bm)
    mc     = 100
    res = map(1:mc) do _
        w             = L*randn(size(L,1)) .+ wm
        a,b           = params2poly(w,na,nb)
        arxtf         = tf(b,a,arxtfm.Ts)
        mag, phase, _ = bode(arxtf, ω)
        mag[:], phase[:]
    end
    magmc      = reduce(hcat, getindex.(res,1))
    phasemc    = reduce(hcat, getindex.(res,2))
    mag        = mean(magmc,dims=2)[:]
    phase      = mean(phasemc,dims=2)[:]
    # mag,phase,_ = bode(arxtfm, ω) .|> x->x[:]
    uppermag   = getpercentile(magmc,0.95)[:]
    lowermag   = getpercentile(magmc,0.05)[:]
    upperphase = getpercentile(phasemc,0.95)[:]
    lowerphase = getpercentile(phasemc,0.05)[:]
    layout := (2,1)

    @series begin
        subplot := 1
        title --> "ARX estimate"
        ylabel --> "Magnitude"
        fillrange := [lowermag, uppermag]
        yscale --> :log10
        xscale --> :log10
        alpha --> 0.3
        ω, mag
    end
    @series begin
        subplot := 2
        fillrange := [lowerphase, upperphase]
        ylabel --> "Phase [deg]"
        xlabel --> "Frequency [rad/s]"
        xscale --> :log10
        alpha --> 0.3
        ω, phase
    end
    nothing

end

"""
    getpercentile(x,p)

calculates the `p`th percentile along dim 2
"""
function getpercentile(mag,p)
    uppermag = mapslices(mag, dims=2) do magω
        sort(magω)[round(Int,length(magω)*p)]
    end
end


# auxiliary function to compute necessary contributions for refined Simpson rule
function residual!(
    f!   :: Function,
    buff :: Matrix{Float64},
    a    :: Float64,
    b    :: Float64,
    n    :: Int64
    )    :: Nothing

    # compute refined step size
    h = 0.5 * (b - a) / n

    # compute integral contributions for new nodes
    for i in 1 : n
        f!(buff, a + (2.0 * i - 1.0) * h, +4.0 * h / 3.0)
    end

    # compute integral contributions for reweighted nodes
    for i in 1 : n ÷ 2
        f!(buff, a + (4.0 * i - 2.0) * h, -2.0 * h / 3.0)
    end

    return nothing
end

# integrate inplace function over an interval [a, b], with b >= a, using an adaptive Simpson rule with Richardson extrapolation on the converged result
function simps!(
    f!    :: Function,
    buff1 :: Matrix{Float64},
    buff2 :: Matrix{Float64},
    a     :: Float64,
    b     :: Float64,
    atol  :: Float64,
    rtol  :: Float64
    )     :: Nothing

    # reset result buffer
    @turbo buff1 .= 0.0

    # compute initial approximation
    m = 0.5 * (a + b)
    f!(buff1, a, 1.0 * (b - a) / 6.0)
    f!(buff1, m, 2.0 * (b - a) / 3.0)
    f!(buff1, b, 1.0 * (b - a) / 6.0)

    # compute improved approximation
    @turbo buff2 .= 0.5 .* buff1
    residual!((b, x, dx) -> f!(b, x, dx), buff2, a, b, 2)
    n = 4

    # continue improving the integral approximation until convergence is reached
    while true
        # perform Richardson extrapolation
        @turbo buff1 .*= -1.0 / 15.0
        @turbo buff1 .+= 16.0 / 15.0 .* buff2

        # compute errors
        norm1 = norm(buff1)
        norm2 = norm(buff2)
        @turbo buff1 .-= buff2
        adiff = norm(buff1)
        rdiff = adiff / max(norm1, norm2)

        # if result has converged or maximum depth (= 20) is reached, terminate
        if adiff < atol || rdiff < rtol || n > 2^20
            @turbo buff1 .+= buff2
            break
        end

        # update current solution
        @turbo buff1 .= buff2

        # compute improved approximation
        @turbo buff2 .= 0.5 .* buff1
        residual!((b, x, dx) -> f!(b, x, dx), buff2, a, b, n)
        n *= 2
    end

    return nothing
end

# integrate inplace function over an interval [a, b], with b >= a, by pre-discretizing the integration domain linearly and applying an adaptive Simpson rule to each subdomain
# note: tbuff[1] is not reset for convenience in flow integration
function integrate_lin!(
    f!    :: Function,
    tbuff :: NTuple{3, Matrix{Float64}},
    a     :: Float64,
    b     :: Float64,
    eval  :: Int64,
    atol  :: Float64,
    rtol  :: Float64
    )     :: Nothing

    @assert b >= a "Upper integration bound must be larger than or equal to lower bound."

    if b > a
        # split integration domain in subdomains of equal length
        h = (b - a) / eval

        # iterate over subdomains and apply adaptive Simpson rule
        for i in 1 : eval
            simps!((b, x, dx) -> f!(b, x, dx), tbuff[2], tbuff[3], a + (i - 1) * h, a + i * h, atol, rtol)
            tbuff[1] .+= tbuff[2]
        end
    end

    return nothing
end

# integrate inplace function over an interval [a, b], with b >= a and a > 0, by pre-discretizing the integration domain logarithmically and applying an adaptive Simpson rule to each subdomain
# note: tbuff[1] is not reset for convenience in flow integration
# note: the sgn keyword can be usend to map [a, b] -> [-b, -a]
function integrate_log!(
    f!    :: Function,
    tbuff :: NTuple{3, Matrix{Float64}},
    a     :: Float64,
    b     :: Float64,
    eval  :: Int64,
    atol  :: Float64,
    rtol  :: Float64
    ;
    sgn   :: Float64 = 1.0
    )     :: Nothing

    @assert b >= a  "Upper integration bound must be larger than or equal to lower bound."
    @assert a > 0.0 "Lower bound must be larger zero."

    if b > a
        # determine logarithmic factor
        ξ = (b / a)^(1.0 / eval)

        # iterate over subdomains and apply adaptive Simpson rule
        for i in 1 : eval
            simps!((b, x, dx) -> f!(b, sgn * x, dx), tbuff[2], tbuff[3], ξ^(i - 1) * a, ξ^i * a, atol, rtol)
            tbuff[1] .+= tbuff[2]
        end
    end

    return nothing
end


using DataStructures
using Plots; pyplot(fmt="png")
PyPlot.matplotlib.rc("mathtext", fontset="cm");
PyPlot.matplotlib.rc("text", usetex=true);
PyPlot.matplotlib.rc("font", family="serif", size=12);
using LaTeXStrings
using LsqFit
using Statistics
using JLD

data_dir = "/home/user/thesis/data/scaling";
savepath = "/home/user/thesis/latex/images";
save_bool = true
dpi = 144
if save_bool
	dpi = 300
end
save_bool

# %%

Ns = 2 .^ collect(15:24);
save_file = joinpath(data_dir, "cluster_size_distribution_avgs.jld")
distribution_avgs = load(save_file)["distribution_avgs"]

# %% Plot t_0 data

distribution_avg = distribution_avgs["t_0"]
# distribution_avg = Dict(2^k => distributions["t_0"][2^k][1] for k in 15:24)
plot_0 = plot(dpi=dpi);
Ns = sort(collect(keys(distribution_avg)));
colors = RGB[cgrad(:plasma)[z] for z=range(0, stop=1, length=length(Ns))];
reverse!(colors);
markers = filter((m->begin m in Plots.supported_markers() end), Plots._shape_keys);

for i in 1:length(Ns)
	N = Ns[i]
	distribution = distribution_avg[N]
	s = collect(keys(distribution))
	n_s = collect(values(distribution))
	x = log10.(s)
	y = log10.(n_s)
	scatter!(x, y,
		legend=:topright,
		xaxis=(L"\log_{10} (s)", (-0.2, 5)),
		yaxis=(L"\log_{10} (n_s(t_0))", (-12, 0)),
		marker=(2, colors[i], markers[i], 0.8, stroke(colors[i])),
		label=latexstring("N = 2^{$(Int(log2(N)))}"),
	)
end

@. function f_linear(x, α)
	return α[1] * x + α[2]
end

αs    = [];
σs    = [];
ratio = 0.01;
for N in Ns[end:end]
	d = distribution_avg[N];
	x = log10.(collect(keys(d)));
	y = log10.(collect(values(d)));
	order = sortperm(x);
	x = x[order][1:Int(floor(ratio * length(x)))];
	y = y[order][1:Int(floor(ratio * length(y)))];
	fit = curve_fit(f_linear, x, y, [-0.1, 0])
	push!(αs, fit.param)
	push!(σs, stderror(fit))
end

α = sum(αs) / length(αs)
σ = sum(σs) / length(σs)
τ₀ = -α[1]

d = distribution_avg[2^24];
x = log10.(collect(keys(d)));
order = sortperm(x);
x = x[order][1:Int(floor(ratio * length(x)))];
y = f_linear(x, α);
plot!(x, y,
	label=latexstring("\\mathrm{Fit} (N = 2^{24}) \\sim $(round(α[1], digits=3)) ($(Int(round(σ[1] * 1e3, digits=0)))) \\cdot \\log_{10}(s)"),
	line=(1, :black),
)

if save_bool
	savefig(plot_0, joinpath(savepath, "n_s_t_0.png"));
end

plot_0

# %% Plot t_1 data

distribution_avg = distribution_avgs["t_1"];
plot_1 = plot(dpi=dpi);
Ns = sort(collect(keys(distribution_avg)));
colors = RGB[cgrad(:plasma)[z] for z=range(0, stop=1, length=length(Ns))];
reverse!(colors);
markers = filter((m->begin m in Plots.supported_markers() end), Plots._shape_keys);

for i in 1:length(Ns)
	N = Ns[i]
	distribution = distribution_avg[N]
	s = collect(keys(distribution))
	n_s = collect(values(distribution))
	x = log10.(s ./ N)
	y = log10.(n_s)
	scatter!(x, y,
		legend=:topright,
		xaxis=(L"\log_{10} (s / N)", (-8, 0)),
		yaxis=(L"\log_{10} (n_s(t_1))", (-12, 0)),
		marker=(2, colors[i], markers[i], 0.8, stroke(colors[i])),
		label=latexstring("N = 2^{$(Int(log2(N)))}"),
	)
end

@. function f_linear(x, α)
	return α[1] * x + α[2]
end

αs    = [];
σs    = [];
ratio = 0.10;
for N in Ns[end:end]
	d = distribution_avg[N];
	x = log10.(collect(keys(d)) ./ N);
	y = log10.(collect(values(d)));
	order = sortperm(x);
	x = x[order][1:Int(floor(ratio * length(x)))];
	y = y[order][1:Int(floor(ratio * length(y)))];
	fit = curve_fit(f_linear, x, y, [-0.1, 0])
	push!(αs, fit.param)
	push!(σs, stderror(fit))
end

α = sum(αs) / length(αs)
σ = sum(σs) / length(σs)
τ₁ = -α[1]

d = distribution_avg[2^24];
x = log10.(collect(keys(d)) ./ 2^24);
order = sortperm(x);
x = x[order][1:Int(floor(ratio * length(x)))];
y = f_linear(x, α)
plot!(x, y,
	label=latexstring("\\mathrm{Fit} (N = 2^{24}) \\sim $(round(α[1], digits=3)) ($(Int(round(σ[1] * 1e3, digits=0)))) \\cdot \\log_{10}(s/N)"),
	line=(1, :black),
)

if save_bool
	savefig(plot_1, joinpath(savepath, "n_s_t_1.png"));
end

plot_1

# %% Plot FSS collapse for τ₀

distribution_avg = distribution_avgs["t_0"]
plot_a = plot(dpi=dpi);
Ns = sort(collect(keys(distribution_avg)));
colors = RGB[cgrad(:plasma)[z] for z=range(0, stop=1, length=length(Ns))]
reverse!(colors)
markers = filter((m->begin m in Plots.supported_markers() end), Plots._shape_keys)

for i in 1:length(Ns)
	N = Ns[i]
	distribution = distribution_avg[N]
	s = collect(keys(distribution))
	n_s = collect(values(distribution))
	x = log10.(s ./ N^(1/τ₁))
	y = log10.(s.^τ₁ .* n_s)
	scatter!(x, y,
		legend=:topleft,
		xaxis=(L"\log_{10} (s / N^{1/\tau})", (-4, 2), -4:1:2),
		yaxis=(L"\log_{10} (s^\tau n_s(t_0))", (-1.5, 0.5), -1.5:0.5:0.5),
		marker=(2, colors[i], markers[i], 0.8, stroke(colors[i])),
		label=latexstring("N = 2^{$(Int(log2(N)))}"),
	)
	annotate!(-2, 0, text(latexstring("\\tau_1 \\approx $(round(τ₁, digits=3))")))
end

if save_bool
	savefig(plot_a, joinpath(savepath, "fss_collapse.png"));
end

plot_a

# %% Plot FSS collapse for τ₀

distribution_avg = distribution_avgs["t_0"]
plot_b = plot(dpi=dpi);
Ns = sort(collect(keys(distribution_avg)));
colors = RGB[cgrad(:plasma)[z] for z=range(0, stop=1, length=length(Ns))]
reverse!(colors)
markers = filter((m->begin m in Plots.supported_markers() end), Plots._shape_keys)

for i in 1:length(Ns)
	N = Ns[i]
	distribution = distribution_avg[N]
	s = collect(keys(distribution))
	n_s = collect(values(distribution))
	x = log10.(s ./ N^(1/τ₀))
	y = log10.(s.^τ₀ .* n_s)
	scatter!(x, y,
		legend=false,
		xaxis=(L"\log_{10} (s / N^{1/\tau})", (-4, 2), -4:1:2),
		yaxis=(L"\log_{10} (s^\tau n_s(t_0))", (-1.5, 0.5), -1.5:0.5:0.5),
		marker=(2, colors[i], markers[i], 0.8, stroke(colors[i])),
		label=latexstring("N = 2^{$(Int(log2(N)))}"),
	)
	annotate!(-2, 0, text(latexstring("\\tau_0 \\approx $(round(τ₀, digits=3))")))
end

plot_b

# %% Plot FSS collapse for τ₀

τ_ref = τ₁ + 0.05
distribution_avg = distribution_avgs["t_0"]
plot_c = plot(dpi=dpi);
Ns = sort(collect(keys(distribution_avg)));
colors = RGB[cgrad(:plasma)[z] for z=range(0, stop=1, length=length(Ns))]
reverse!(colors)
markers = filter((m->begin m in Plots.supported_markers() end), Plots._shape_keys)

for i in 1:length(Ns)
	N = Ns[i]
	distribution = distribution_avg[N]
	s = collect(keys(distribution))
	n_s = collect(values(distribution))
	x = log10.(s ./ N^(1/τ_ref))
	y = log10.(s.^τ_ref .* n_s)
	scatter!(x, y,
		legend=false,
		xaxis=(L"\log_{10} (s / N^{1/\tau})", (-4, 2), -4:1:2),
		yaxis=(L"\log_{10} (s^\tau n_s(t_0))", (-1.5, 0.5), -1.5:0.5:0.5),
		marker=(2, colors[i], markers[i], 0.8, stroke(colors[i])),
		label=latexstring("N = 2^{$(Int(log2(N)))}"),
	)
	annotate!(-2, 0, text(latexstring("\\tau_1 + 0.05")))
end

plot_c

# %% Plot FSS comparison

l = @layout [a{0.6h}; [b{0.5w} c]]
plot_triple = plot(plot_a, plot_b, plot_c, layout=l, size=(800, 600))

if save_bool
	savefig(plot_triple, joinpath(savepath, "fss_collapse_triple.png"));
end

plot_triple

# %% Maximual decay exponent calculation

δ = 0.302;
exponent_of_decay = -δ * (τ₁ - 2) / (τ₁ - 1)


using AugmentedGaussianProcesses
using Distributions, LinearAlgebra
using Random: seed!;
seed!(42);
using Plots;
pyplot();

N = 200;#Number of training points
N_test = 1000 # Number of points for predictions
X = sort(rand(N)) # Training points
shift = 0.1 * (maximum(X) - minimum(X))
X_test = collect(range(minimum(X) - shift; length=N_test, stop=maximum(X) + shift))
## Noise parameters
σ = 5e-1 #Gaussian noise
ν = 4.0; #Degrees of freedom for Student T
β = 5e-1; # Shape for Laplace
mu0 = 0.3; # Constant for heteroscedastic gp
## Kernel
k = SqExponentialKernel() ∘ ScaleTransform(10.0) # Squared exponential kernel
K = kernelmatrix(k, X) + 1e-9 * I #Kernel matrix
f = rand(MvNormal(K)) #Random GP sampled from the prior
g = rand(MvNormal(mu0 * ones(N), K)) #Random GP sampled from the prior for Heteroscedasticity
y_gaussian = f .+ rand(Normal(0, σ), N)
y_studentt = f .+ rand(TDist(ν), N)
y_laplace = f .+ rand(Laplace(0, β), N)
y_hetero = f .+ randn(N) .* sqrt.(inv.(AGP.logistic.(g)))

gpmodel = GP(X, y_gaussian, k; noise=σ^2);
train!(gpmodel, 50)
gppred, gppred_cov = proba_y(gpmodel, X_test)

stumodel = VGP(X, y_studentt, k, StudentTLikelihood(ν), AnalyticVI());
train!(stumodel, 50)#,callback=intplot)
stupred, stupred_cov = proba_y(stumodel, X_test)

lapmodel = VGP(X, y_laplace, k, LaplaceLikelihood(β), AnalyticVI());
train!(lapmodel, 50)#,callback=intplot)
lappred, lappred_cov = proba_y(lapmodel, X_test)

hetmodel = VGP(X, y_hetero, k, HeteroscedasticLikelihood(), AnalyticVI())
train!(hetmodel, 100)
hetpred, hetpred_cov = proba_y(hetmodel, X_test)

## Plotting
lw = 4.0 #Width of the lines
lwt = 2.0
nsig = 2 #Number of sigmas for the uncertainty
falpha = 0.3
pstudent = scatter(X, y_studentt; t=:scatter, lab="")
plot!(
    X_test,
    stupred;
    ribbon=nsig .* sqrt.(stupred_cov),
    title="StudentT Regression",
    lw=lw,
    color=1,
    lab="",
    fillalpha=falpha,
    linecolor=2,
)
plot!(X, f; color=:black, lw=lwt, lab="")
plot!(X, f .+ nsig * sqrt(ν / (ν - 2)); lab="", color=:black, lw=lwt, linestyle=:dash)
plot!(X, f .- nsig * sqrt(ν / (ν - 2)); lab="", color=:black, lw=lwt, linestyle=:dash)
# ylims!(maxlims)

pgauss = scatter(X, y_gaussian; lab="")
plot!(
    X_test,
    gppred;
    ribbon=nsig * gppred_cov,
    title="Gaussian Regression",
    lw=lw,
    color=1,
    lab="",
    fillalpha=falpha,
    linecolor=2,
)
plot!(X, f; color=:black, lw=lwt, lab="")
plot!(X, f .+ nsig * σ; lab="", color=:black, lw=lwt, linestyle=:dash)
plot!(X, f .- nsig * σ; lab="", color=:black, lw=lwt, linestyle=:dash)

plaplace = scatter(X, y_laplace; lab="")
plot!(
    X_test,
    lappred;
    ribbon=nsig * sqrt.(lappred_cov),
    title="Laplace Regression",
    lw=lw,
    color=1,
    lab="",
    fillalpha=falpha,
    linecolor=2,
)
plot!(X, f; color=:black, lw=lwt, lab="")
plot!(X, f .+ nsig * β * sqrt(2); lab="", color=:black, lw=lwt, linestyle=:dash)
plot!(X, f .- nsig * β * sqrt(2); lab="", color=:black, lw=lwt, linestyle=:dash)

phetero = plot(X, y_hetero; t=:scatter, lab="")
plot!(
    X_test,
    lappred;
    ribbon=nsig .* sqrt.(hetpred_cov),
    title="Heteroscedastic Regression",
    lw=lw,
    color=1,
    lab="",
    fillalpha=falpha,
    linecolor=2,
)
plot!(X, f; color=:black, lw=lwt, lab="")
plot!(
    X,
    f .+ nsig * sqrt.(inv.(AGP.logistic.(g)));
    lab="",
    color=:black,
    lw=lwt,
    linestyle=:dash,
)
plot!(
    X,
    f .- nsig * sqrt.(inv.(AGP.logistic.(g)));
    lab="",
    color=:black,
    lw=lwt,
    linestyle=:dash,
)

# default(legendfontsize=10.0,xtickfontsize=10.0,ytickfontsize=10.0)
p = plot(pgauss, pstudent, plaplace, phetero; ylims=(-5.0, 5.0))
display(p)
savefig(p, joinpath(@__DIR__, "Regression.png"))


using ReachabilityAnalysis, SparseArrays, JLD2

LazySets.set_ztol(Float64, 1e-14)

const x25 = [zeros(24); 1.0; zeros(23)]
const x25e = vcat(x25, 0.0)

examples_dir = normpath(@__DIR__, "..", "..", "..", "examples")
building_path = joinpath(examples_dir, "Building", "building.jld2")

function building_BLDF01()
    @load building_path A B
    n = size(A, 1)
    U = Interval(0.8, 1.0)
    S = @system(x' = Ax + Bu, u ∈ U, x ∈ Universe(n))

    # initial states
    center_X0 = [fill(0.000225, 10); fill(0.0, 38)]
    radius_X0 = [fill(0.000025, 10); fill(0.0, 14); 0.0001; fill(0.0, 23)]
    X0 = Hyperrectangle(center_X0, radius_X0)

    prob_BLDF01 = InitialValueProblem(S, X0)
end

using ReachabilityAnalysis: add_dimension

function building_BLDC01()
    @load building_path A B
    n = size(A, 1)
    U = Interval(0.8, 1.0)

    # initial states
    center_X0 = [fill(0.000225, 10); fill(0.0, 38)]
    radius_X0 = [fill(0.000025, 10); fill(0.0, 14); 0.0001; fill(0.0, 23)]
    X0 = Hyperrectangle(center_X0, radius_X0)

    Ae = add_dimension(A)
    Ae[1:n, end] = B
    prob_BLDC01 = @ivp(x' = Ae * x, x(0) ∈ X0 × U)
end

using Plots

prob_BLDF01 = building_BLDF01()

sol_BLDF01_dense = solve(prob_BLDF01, T=20.0, alg=LGG09(δ=0.004, vars=(25), n=48));

plot(sol_BLDF01_dense, vars=(0, 25), linecolor=:blue, color=:blue, alpha=0.8, lw=1.0, xlab="t", ylab="x25")

ρ(x25, sol_BLDF01_dense)

ρ(x25, sol_BLDF01_dense) <= 5.1e-3 # BLDF01 - BDS01

ρ(x25, sol_BLDF01_dense) <= 4e-3 # BLDF01 - BDU01

ρ(x25, sol_BLDF01_dense(20.0))

ρ(x25, sol_BLDF01_dense(20.0)) <= -0.78e-3 # BLDF01 - BDU02

sol_BLDF01_discrete = solve(prob_BLDF01, T=20.0, alg=LGG09(δ=0.01, vars=(25), n=48, approx_model=NoBloating()));

plot(sol_BLDF01_discrete, vars=(0, 25), linecolor=:blue, color=:blue, alpha=0.8, lw=1.0, xlab="t", ylab="x25")

ρ(x25, sol_BLDF01_discrete)

ρ(x25, sol_BLDF01_discrete) <= 5.1e-3 # BLDF01 - BDS01

ρ(x25, sol_BLDF01_discrete)

ρ(x25, sol_BLDF01_discrete) <= 4e-3 # BLDF01 - BDU01

ρ(x25, sol_BLDF01_discrete(20.0))

ρ(x25, sol_BLDF01_discrete(20.0)) <= -0.78e-3 # BLDF01 - BDU02

prob_BLDC01 = building_BLDC01()

sol_BLDC01_dense = solve(prob_BLDC01, T=20.0, alg=LGG09(δ=0.005, vars=(25), n=49))

plot(sol_BLDC01_dense, vars=(0, 25), linecolor=:blue, color=:blue, alpha=0.8, lw=1.0, xlab="t", ylab="x25")

ρ(x25e, sol_BLDC01_dense)

ρ(x25e, sol_BLDC01_dense) <= 5.1e-3 # BLDC01 - BDS01

ρ(x25e, sol_BLDC01_dense) <= 4e-3 # BLDC01 - BDU01

ρ(x25, sol_BLDF01_discrete(20.0))

ρ(x25e, sol_BLDC01_dense(20.0)) <= -0.78e-3 # BLDC01 - BDU02

sol_BLDC01_discrete = solve(prob_BLDC01, T=20.0, alg=LGG09(δ=0.01, vars=(25), n=49, approx_model=NoBloating()))

plot(sol_BLDC01_discrete, vars=(0, 25), linecolor=:blue, color=:blue, alpha=0.8, lw=1.0, xlab="t", ylab="x25")

ρ(x25e, sol_BLDC01_discrete)

ρ(x25e, sol_BLDC01_discrete) <= 5.1e-3 # BLDC01 - BDS01

ρ(x25e, sol_BLDC01_discrete)

ρ(x25e, sol_BLDC01_discrete) <= 4e-3 # BLDC01 - BDU01

ρ(x25e, sol_BLDC01_discrete(20.0))

ρ(x25e, sol_BLDC01_discrete(20.0)) <= -0.78e-3 # BLDC01 - BDU02


module AdventOfCode

using Combinatorics
mutable struct Moon pos; vel end

gravity!(m1::Moon, m2::Moon) = (dv = (m1.pos .< m2.pos) - (m1.pos .> m2.pos); m1.vel += dv; m2.vel -= dv)
gravity!(moons) = map(ms -> gravity!(ms...), combinations(moons, 2))
velocity!(m::Moon) = m.pos += m.vel
energy(m::Moon) = sum(abs.(m.pos)) * sum(abs.(m.vel))

# Part 1: compute total energy after 1000 simulation steps
initial_moons() = [
    Moon([-7 -8 9], [0 0 0]),
    Moon([-12 -3 -4], [0 0 0]),
    Moon([6 -17 -9], [0 0 0]),
    Moon([4 -10 -6], [0 0 0])
]
jupiter_moons = initial_moons()
for step in 1:1000
    gravity!(jupiter_moons)
    velocity!.(jupiter_moons)
end
println("Total energy: $(sum(energy.(jupiter_moons)))")

# Part 2: find how long it takes before first repeated step
# Approach: x, y and z dimension equations are independent, maybe each one has a shorter period?
dim_history = [Dict() Dict() Dict()]
cycle_ends = Array{Any}([nothing nothing nothing])
jupiter_moons = initial_moons()
for step in 0:10000000
    gravity!(jupiter_moons)
    velocity!.(jupiter_moons)
    for dim in 1:3
        dim_state = map(m -> (m.pos[dim], m.vel[dim]), jupiter_moons)
        if isnothing(cycle_ends[dim])
            if dim_state ∈ keys(dim_history[dim])
                cycle_ends[dim] = (dim_history[dim][dim_state], step)
            else
                dim_history[dim][dim_state] = step
            end
        end
    end
    if !any(isnothing.(cycle_ends)) break end
end

if all(0 .== first.(cycle_ends))
    period = lcm(last.(cycle_ends)...)
    println("Loop start/end: $cycle_ends, repeat period: $period")
else
    println("Ugh, the maths is harder, I have to think more to implement this")
end

end


"Container type for elements of product (semi-)rings."
struct Product{S,T}
    first  :: S
    second :: T
end
make_accesses(Product)

Product(first, second, more...) = Product(first, Product(second, more...))
components(p::Product{S,T}) where {S, T} = (p.first, p.second)
components(p::Product{S,Product{T1,T2}}) where {S,T1,T2} = (p.first, components(p.second)...)

show(io::IO, p::Product) = print(io, "Product", components(p))

for op in (:+,:-,:*,:/)
    eval(quote
        ($op)(p1::Product{S,T}, p2::Product{S,T}) where {S,T} =
            Product(($op)(p1.first, p2.first), ($op)(p1.second, p2.second))
    end)
end

zero(::Product{S,T}) where {S,T} = Product(zero(S), zero(T))
one(::Product{S,T}) where {S,T} = Product(one(S), one(T))

<(p1::Product, p2::Product) = p1.first < p2.first && p1.second < p2.second


# By starting at the top of the triangle below and moving to adjacent numbers
# on the row below, the maximum total from top to bottom is 23.
#
# 3
# 7 4
# 2 4 6
# 8 5 9 3
#
# That is, 3 + 7 + 4 + 9 = 23.
#
# Find the maximum total from top to bottom of the triangle below:
#
# 75
# 95 64
# 17 47 82
# 18 35 87 10
# 20 04 82 47 65
# 19 01 23 75 03 34
# 88 02 77 73 07 63 67
# 99 65 04 28 06 16 70 92
# 41 41 26 56 83 40 80 70 33
# 41 48 72 33 47 32 37 16 94 29
# 53 71 44 65 25 43 91 52 97 51 14
# 70 11 33 28 77 73 17 78 39 68 17 57
# 91 71 52 38 17 14 91 43 58 50 27 29 48
# 63 66 04 68 89 53 67 30 73 16 69 87 40 31
# 04 62 98 27 23 09 70 98 73 93 38 53 60 04 23
#
# NOTE: As there are only 16384 routes, it is possible to solve this problem by
# trying every route. However, Problem 67, is the same challenge with a triangle
# containing one-hundred rows; it cannot be solved by brute force, and requires
# a clever method! ;o)

using ProjectEulerSolutions

# Convert a string representation of a pyramid of numbers into an integer
# representation.
function pyramid_from_string(input::String="32")::Vector{Vector{Integer}}
    pyramid = Vector{Vector{Integer}}()
    for row in split(input, "\n")
        if length(row) > 0
            cols = map(x->convert(Integer, parse(Int, x)), split(row, " "))
            push!(pyramid, cols)
        end
    end
    return pyramid
end

p018 = Problems.Problem(max_path_through_pyramid)

input =  """
         75
         95 64
         17 47 82
         18 35 87 10
         20 04 82 47 65
         19 01 23 75 03 34
         88 02 77 73 07 63 67
         99 65 04 28 06 16 70 92
         41 41 26 56 83 40 80 70 33
         41 48 72 33 47 32 37 16 94 29
         53 71 44 65 25 43 91 52 97 51 14
         70 11 33 28 77 73 17 78 39 68 17 57
         91 71 52 38 17 14 91 43 58 50 27 29 48
         63 66 04 68 89 53 67 30 73 16 69 87 40 31
         04 62 98 27 23 09 70 98 73 93 38 53 60 04 23
         """

Problems.benchmark(p018, pyramid_from_string(input))


# This file is part of the ValidatedNumerics.jl package; MIT licensed

## Powers

# CRlibm does not contain a correctly-rounded ^ function for Float64
# Use the BigFloat version from MPFR instead, which is correctly-rounded:

# Write explicitly like this to avoid ambiguity warnings:
for T in (:Integer, :Rational, :Float64, :BigFloat, :Interval)

    @eval ^(a::Interval{Float64}, x::$T) = convert(Interval{Float64}, big53(a)^x)
end

# Integer power:

function ^(a::Interval{BigFloat}, n::Integer)
    isempty(a) && return a
    n == 0 && return one(a)
    n == 1 && return a
    n == 2 && return sqr(a)
    n < 0 && a == zero(a) && return emptyinterval(a)

    T = BigFloat

    if isodd(n) # odd power
        isentire(a) && return a
        if n > 0
            a.lo == zero(T) && return @round(T, zero(T), a.hi^n)
            a.hi == zero(T) && return @round(T, a.lo^n, zero(T))
            return @round(T, a.lo^n, a.hi^n)
        else
            if a.lo ≥ zero(T)
                a.lo == zero(T) && return @round(T, a.hi^n, convert(T, Inf))
                return @round(T, a.hi^n, a.lo^n)

            elseif a.hi ≤ zero(T)
                a.hi == zero(T) && return @round(T, convert(T, -Inf), a.lo^n)
                return @round(T, a.hi^n, a.lo^n)
            else
                return entireinterval(a)
            end
        end

    else # even power
        if n > 0
            if a.lo ≥ zero(T)
                return @round(T, a.lo^n, a.hi^n)
            elseif a.hi ≤ zero(T)
                return @round(T, a.hi^n, a.lo^n)
            else
                return @round(T, mig(a)^n, mag(a)^n)
            end

        else
            if a.lo ≥ zero(T)
                return @round(T, a.hi^n, a.lo^n)
            elseif a.hi ≤ zero(T)
                return @round(T, a.lo^n, a.hi^n)
            else
                return @round(T, mag(a)^n, mig(a)^n)
            end
        end
    end
end

function sqr{T<:Real}(a::Interval{T})
    isempty(a) && return a
    if a.lo ≥ zero(T)
        return @round(T, a.lo^2, a.hi^2)

    elseif a.hi ≤ zero(T)
        return @round(T, a.hi^2, a.lo^2)
    end

    return @round(T, mig(a)^2, mag(a)^2)
end

# Floating-point power of a BigFloat interval:
function ^(a::Interval{BigFloat}, x::AbstractFloat)
    T = BigFloat
    domain = Interval(zero(T), convert(T, Inf))

    if a == zero(a)
        a = a ∩ domain
        x > zero(x) && return zero(a)
        return emptyinterval(a)
    end

    isinteger(x) && return a^(round(Int, x))
    x == one(T)/2 && return sqrt(a)

    a = a ∩ domain
    (isempty(x) || isempty(a)) && return emptyinterval(a)

    xx = convert(Interval{BigFloat}, x)

    lo = @round(T, a.lo^xx.lo, a.lo^xx.lo)
    lo1 = @round(T, a.lo^xx.hi, a.lo^xx.hi)
    hi = @round(T, a.hi^xx.lo, a.hi^xx.lo)
    hi1 = @round(T, a.hi^xx.hi, a.hi^xx.hi)

    lo = hull(lo, lo1)
    hi = hull(hi, hi1)

    return hull(lo, hi)
end

function ^{T<:Integer,}(a::Interval{Rational{T}}, x::AbstractFloat)
    a = Interval(a.lo.num/a.lo.den, a.hi.num/a.hi.den)
    a = a^x
    convert(Interval{Rational{T}}, a)
end

# Rational power
function ^{S<:Integer}(a::Interval{BigFloat}, r::Rational{S})
    T = BigFloat
    domain = Interval(zero(a.lo), convert(T, Inf))

    if a == zero(a)
        a = a ∩ domain
        r > zero(r) && return zero(a)
        return emptyinterval(a)
    end

    isinteger(r) && return convert(Interval{T}, a^round(S,r))
    r == one(S)//2 && return sqrt(a)

    a = a ∩ domain
    (isempty(r) || isempty(a)) && return emptyinterval(a)

    y = convert(Interval{BigFloat}, r)

    a^y
end

# Interval power of an interval:
function ^(a::Interval{BigFloat}, x::Interval)
    T = BigFloat
    domain = Interval(zero(T), convert(T, Inf))

    a = a ∩ domain

    (isempty(x) || isempty(a)) && return emptyinterval(a)

    lo1 = a^x.lo
    lo2 = a^x.hi
    lo1 = hull(lo1, lo2)

    hi1 = a^x.lo
    hi2 = a^x.hi
    hi1 = hull(hi1, hi2)

    hull(lo1, hi1)
end

function sqrt{T}(a::Interval{T})
    domain = Interval(zero(T), convert(T, Inf))
    a = a ∩ domain

    isempty(a) && return a

    @round(T, sqrt(a.lo), sqrt(a.hi))  # `sqrt` is correctly-rounded
end

for f in (:exp, :expm1)
    @eval begin
        function ($f){T}(a::Interval{T})
            isempty(a) && return a
            Interval( ($f)(a.lo, RoundDown), ($f)(a.hi, RoundUp) )
        end
    end
end

for f in (:exp2, :exp10)

    @eval function ($f)(x::BigFloat, r::RoundingMode)  # add BigFloat functions with rounding:
            setrounding(BigFloat, r) do
                ($f)(x)
            end
        end

    @eval ($f)(a::Interval{Float64}) = float($f(big53(a)))  # no CRlibm version

    @eval function ($f)(a::Interval{BigFloat})
            isempty(a) && return a
            Interval( ($f)(a.lo, RoundDown), ($f)(a.hi, RoundUp) )
        end
end

for f in (:log, :log2, :log10, :log1p)

    @eval function ($f){T}(a::Interval{T})
            domain = Interval(zero(T), convert(T, Inf))
            a = a ∩ domain

            (isempty(a) || a.hi ≤ zero(T)) && return emptyinterval(a)

            Interval( ($f)(a.lo, RoundDown), ($f)(a.hi, RoundUp) )
        end
end