function ct_kalman_simple(T, Z, Q, E, mean_0, var_0, data_y, dt)
    # Compute Kalman filter of a continuous time process sampled at discrete time
    # ported from SeHyoun Ahn
    #
    #   & dx_t = Tx_t dt + R dW_t \\
    #   & y_{t_k} = Zx_{t_k} + v_{t_k} \\
    #   & x_0 \sim N(mean_0, var_0) \\
    #   & k\in\{1, 2, \ldots, n\} \\
    #   & W_t = \text{Brownian Motion} \\
    #   & v_k \sim N(0, E) \\
    #   & dt_i = t_{i} - t_{i-1}
    #
    # Args:
    #   T (matrix): see equation above
    #   Z (matrix): see equation above
    #   Q (matrix): :math:`Q = RR'` i.e., variance
    #   E (matrix): see equation above
    #   mean_0 (vector): see equation above
    #   var_0 (matrix): see equation above
    #   data_y (vector): :math:`y_k` above
    #   dt (vector): time period between data measurement (see above)
    #
    # Note:
    #   - The data does not have to come with uniform step size.
    #   - There is a `kalman <https://www.mathworks.com/help/control/ref/kalman.html>`_ function in MATLAB as well, but it is under control system toolbox.
    #
    # Returns:
    #   : [mean_out, var_out, prob_out]
    #
    #   - **mean_out** (vector): :math:`E(x_{t_n}| y_{t_1},\ldots, y_{t_n})`
    #   - **var_out** (matrix): :math:`var(x_{t_n}| y_{t_1},\ldots, y_{t_n})`
    #   - **prob_out** (double): log-likelihood of the data
    #
    # Reference:
    #
    #   - `Wikipedia article on Kalman Filtering <https://en.wikipedia.org/wiki/Kalman_filter#Hybrid_Kalman_filter>`_
    #   - :cite:`grewal2011kalman`
    #
    # .. Syntax:
    #
    # ..   [mean_out, var_out, prob_out] = kalman_ct(T, Z, Q, E, mean_0, var_0, data_y, dt)
    #

    # Get data size
    n_data = size(data_y, 1)
    if (length(dt) < n_data)
        dt = repeat(dt, n_data)
    end
    n_dim = length(mean_0)
    prob_out = 0
    mean_out = mean_0
    var_out = var_0

    # Step through data points

    for iter_data = 1:n_data

        #       tspan = (0.0,dt[iter_data][1])
        # s0 = mean_0 #initial values
        # P0 = var_out
        #        f(u,p,t) = T*u
        #     g(u,p,t) = T*u + u*T' + Q
        #           s_prob = ODEProblem(f, mean_out, tspan)
        #     P_prob = ODEProblem(g, P0, tspan)
        # sol = solve(s_prob,Tsit5(),reltol=1e-4,abstol=1e-4)

        ## forecasting state (this is different for ct)
        mean_out = exp(dt[iter_data][1]*T)*mean_out
        AB = exp(dt[iter_data][1]*[hcat(T, Q); hcat(zeros(n_dim,n_dim), -T')])*[var_out; Matrix{Float64}(I,n_dim,n_dim)]
        var_inter = AB[1:n_dim, :]/AB[n_dim+1:end, :]
        ##########

        ## the rest is the same
        var_obs = Z*var_inter*Z' .+ E

        v_t = data_y[iter_data,:] .- Z*mean_out
        prob_out = prob_out - .5*n_dim*log(2*pi) -.5*log(abs(det(var_obs))) - 0.5* v_t'*(var_obs\v_t)
        mean_out = mean_out .+ var_inter*Z'*(var_obs\v_t)
        var_out = var_inter .- var_inter*Z'*(var_obs\Z)*var_inter

    end

    return prob_out
end

#=

function [mean_out, var_out, prob_out] = kalman_ct(F, H, Q, R, mean_0, var_0, data_y, dt)
  % Compute Kalman filter of a continuous time process sampled at discrete time
  %
  % .. math::
  %
  %   & dx_t = Fx_t dt + J dW_t \\
  %   & y_{t_k} = Hx_{t_k} + v_{t_k} \\
  %   & x_0 \sim N(mean_0, var_0) \\
  %   & k\in\{1, 2, \ldots, n\} \\
  %   & W_t = \text{Brownian Motion} \\
  %   & v_k \sim N(0, R) \\
  %   & dt_i = t_{i} - t_{i-1}
  %
  % Args:
  %   F (matrix): see equation above
  %   H (matrix): see equation above
  %   Q (matrix): :math:`Q = JJ'` i.e., variance
  %   R (matrix): see equation above
  %   mean_0 (vector): see equation above
  %   var_0 (matrix): see equation above
  %   data_y (vector): :math:`y_k` above
  %   dt (vector): time period between data measurement (see above)
  %
  % Note:
  %   - The data does not have to come with uniform step size.
  %   - There is a `kalman <https://www.mathworks.com/help/control/ref/kalman.html>`_ function in MATLAB as well, but it is under control system toolbox.
  %
  % Returns:
  %   : [mean_out, var_out, prob_out]
  %
  %   - **mean_out** (vector): :math:`E(x_{t_n}| y_{t_1},\ldots, y_{t_n})`
  %   - **var_out** (matrix): :math:`var(x_{t_n}| y_{t_1},\ldots, y_{t_n})`
  %   - **prob_out** (double): log-likelihood of the data
  %
  % Reference:
  %
  %   - `Wikipedia article on Kalman Filtering <https://en.wikipedia.org/wiki/Kalman_filter#Hybrid_Kalman_filter>`_
  %   - :cite:`grewal2011kalman`
  %
  % .. Syntax:
  %
  % ..   [mean_out, var_out, prob_out] = kalman_ct(F, H, Q, R, mean_0, var_0, data_y, dt)
  %

  % Get data size
  n_data = size(data_y, 2);
  if (length(dt) < n_data) & (length(n_data) == 1)
    dt = repmat(dt, n_data, 1);
  end
  n_dim = length(mean_0);

  prob_out = 0;
  mean_out = mean_0;
  var_out = var_0;

  % Step through data points
  for iter_data = 1:n_data
    mean_out = expm(dt(iter_data)*F)*mean_out;
    AB = expm(dt(iter_data)*[F, Q; H'*(R\H), -F])*[var_out; speye(n_dim)];
    var_inter = AB(1:n_dim, :)*inv(AB((n_dim+1):end, :));

    var_obs = H*var_inter*H' + R;
    v_t = data_y(:, iter_data) - H*mean_out;

    prob_out = prob_out - n_dim*log(2*pi)/2 ...
               -(log(abs(det(var_obs))) + v_t'*(var_obs\v_t))/2;

    mean_out = mean_out + var_inter*H'*(var_obs\v_t);
    var_out = var_inter - (var_inter*H')*(var_obs\H)*var_inter;
  end
end

=#


function evolve()
sinfo = Params()

#== start template ==#
#fundamental constants/units
  ħ = 1.0545718e-34
  kB = 1.38064852e-23
  amu = 1.660339040e-27
  Bohr = 5.29e-11
#Rb87 mass and scattering length
  m = 86.909180527*amu
  as = 100*Bohr
#trap frequencies
  ωx = 2π
  ωy = 4ωx
  ωz = 0.
#choice of time, length, energy units
  t0 = 1.0/ωy
  x0 = sqrt(ħ*t0/m)
  E0 = ħ/t0
#interactions
  #g  = (4*pi*ħ^2*as/m)*x0^3/E0 #dimensionless 3D
  g = 0.1 #test 2D
#damping parameters (dimensionless)
  γ  = 0.1
  ℳ  = 0.0
#chemical potential (dimensionless)
  μ  = 12.0
#time evolution parameters
  ti = 0.0
  tf = 1.0/γ  #evolve for 2 damping times
  Nt = 50
  t = collect(linspace(ti,tf,Nt))
  dt = 0.01π/μ #integrate step size [ - should have dt ≪ 2π/μ]
#== end template ==#

  @pack sinfo = ωx,ωy,ωz,γ,ℳ,g,t0,E0,x0,μ,ti,tf,Nt,t,dt

  #Initialize CField
  basis = "Hermite"
  ecut = 30 #units of ħ*ωy
  Ω = [ωx; ωy]*t0
  cinfo = makecinfo(ecut,Ω,basis)
  tinfo = maketinfo(cinfo,4)
  @unpack en, P, M = cinfo; Mx, My = M;
  #x,wx,Tx,y,wy,Ty = makealltrans(M,Ω,n=4,basis=basis)
  #W = wx.*wy'
#test transform
  c0   = randn(Mx,My)+im*randn(Mx,My);
  c0   = P.*c0;   #Project
  ψ0   = c2x(c0,tinfo)
  ψ    = simlar(ψ0) #a field to write to in place if needed

c0    = randn(Mx,My) + im*randn(Mx,My); #create new random initial state
tspan = (t[1],t[end])
prob = ODEProblem(Lgp!,c0,tspan)
alg = DP5()
#abstol = 1e-3!
#reltol = 1e-3
println("Started evolution ...")
@time sol = solve(prob,alg,dt=dt,saveat=t);
println("... Finished.")
return sinfo,cinfo,sol
end


abstract Spline{T} <: ScalarFunc{T}

# cubic 1D spline
type Spline1D{T,V} <: Spline{T}
    x::V
    y::V
    y2::V
    Spline1D(x::AbstractVector{T}, y::AbstractVector{T}, y2::AbstractVector{T}) = new(x,y,y2)
end

Spline1D{T}(x::AbstractVector{T}, y::AbstractVector{T}, y2::AbstractVector{T}) =
    Spline1D{eltype(x),typeof(x)}(x,y,y2)

typealias Spline1Df Spline1D{Float64,Vector{Float64}}

@fastmath function Spline1D(x::AbstractVector{Float64}, y::AbstractVector{Float64})
    @assert length(x) == length(y)
    local sig,p
    y2 = similar(y)
    n = length(x)
    u = Array{Float64}(n-1)
    y2[1] = u[1] = 0.0
    @inbounds for i in 2:n-1
        sig = (x[i]-x[i-1])/(x[i+1]-x[i-1])
        p = sig*y2[i-1]+2.0
        y2[i]=(sig-1.0)/p
        u[i]=(y[i+1]-y[i])/(x[i+1]-x[i]) - (y[i]-y[i-1])/(x[i]-x[i-1])
        u[i]=(6.0*u[i]/(x[i+1]-x[i-1])-sig*u[i-1])/p
    end
    qn=un=0.0
    y2[n]=(un-qn*u[n-1])/(qn*y2[n-1]+1.0)
    @inbounds for k in n-1:-1:1  # This is the backsubstitution loop of the tridiagonal
        y2[k]=y2[k]*y2[k+1]+u[k];
    end
    Spline1D(x,y,y2)
end

function Spline1D(x::AbstractVector{Float64}, y::AbstractVector{Float64}, num::Int)
    n = length(x)
    ind = num < n ? round(Int,linspace(1,n,num)) : Colon()
    Spline1D(x[ind],y[ind])
end

Base.show(io::IO, f::Spline1D) = write(io, "Spline1D $(length(f.x))-points on interval [$(f.x[1]),$(f.x[end])]")

@fastmath function locate(xx::AbstractVector{Float64}, x::Float64)
    local k
    klo,khi = 1,length(xx)
    while khi-klo > 1
        k = (khi+klo) >> 1  # mean of khi and klo
        if xx[k] > x
            khi = k
        else
            klo = k
        end
    end
    klo,khi
end

@fastmath function Base.call(f::Spline1D, x::Float64)
    klo,khi = locate(f.x,x)
    h = f.x[khi] - f.x[klo]  # @assert h != 0.0
    a = (f.x[khi]-x)/h
    b = (x-f.x[klo])/h
    a*f.y[klo] + b*f.y[khi] + ((a*a*a-a)*f.y2[klo]+(b*b*b-b)*f.y2[khi])*(h*h)/6.
end

@vectorize2 Base.call Spline1D Float64 Float64

# function Base.call(f::Spline1D, dst::AbstractVector{Float64}, x::AbstractVector{Float64})
#     @assert length(dst) == length(x)
#     for i in 1:length(x)
#         dst[i] = f(x[i])
#     end
# end

# function Base.call(f::Spline1D, x::AbstractVector{Float64})
#     dst = similar(x)
#     f(dst,x)
#     dst
# end

@fastmath function deriv1(f::Spline1D, x::Float64)
    klo,khi = locate(f.x,x)
    h = f.x[khi] - f.x[klo]  # @assert h != 0.0
    a = (f.x[khi]-x)/h
    b = (x-f.x[klo])/h
    (f.y[khi] - f.y[klo])/h + ((1.-3.*a*a)*f.y2[klo] + (3*b*b-1.)*f.y2[khi])*h/6.
end

@vectorize2 deriv1 Spline1D Float64 Float64

# function deriv1(f::Spline1D, dst::AbstractVector{Float64}, x::AbstractVector{Float64})
#     @assert length(dst) == length(x)
#     @inbounds for i in 1:length(x)
#         dst[i] = deriv1(f,x[i])
#     end
# end

# function deriv1(f::Spline1D, x::AbstractVector{Float64})
#     dst = similar(x)
#     deriv1(f,dst,x)
#     dst
# end

@fastmath function deriv2(f::Spline1D, x::Float64)
    klo,khi = locate(f.x,x)
    h = f.x[khi] - f.x[klo]  # @assert h != 0.0
    a = (f.x[khi]-x)/h
    b = (x-f.x[klo])/h
    a*f.y2[klo] + b*f.y2[khi]
end

@vectorize2 deriv2 Spline1D Float64 Float64

# function deriv2(f::Spline1D, dst::AbstractVector{Float64}, x::AbstractVector{Float64})
#     @assert length(dst) == length(x)
#     @inbounds for i in 1:length(x)
#         dst[i] = deriv2(f,x[i])
#     end
# end

# function deriv2(f::Spline1D, x::AbstractVector{Float64})
#     dst = similar(x)
#     deriv2(f,dst,x)
#     dst
# end

@fastmath function eval_deriv1(f::Spline1D, x::Float64)
    klo,khi = locate(f.x,x)
    h = f.x[khi] - f.x[klo]  # @assert h != 0.0
    a = (f.x[khi]-x)/h
    b = (x-f.x[klo])/h
    a*f.y[klo] + b*f.y[khi] + ((a*a*a-a)*f.y2[klo]+(b*b*b-b)*f.y2[khi])*(h*h)/6.,
        (f.y[khi] - f.y[klo])/h + ((1.-3.*a*a)*f.y2[klo] + (3*b*b-1.)*f.y2[khi])*h/6.
end

@fastmath function eval_deriv2(f::Spline1D, x::Float64)
    klo,khi = locate(f.x,x)
    h = f.x[khi] - f.x[klo]  # @assert h != 0.0
    a = (f.x[khi]-x)/h
    b = (x-f.x[klo])/h
    a*f.y[klo] + b*f.y[khi] + ((a*a*a-a)*f.y2[klo]+(b*b*b-b)*f.y2[khi])*(h*h)/6.,
        a*f.y2[klo] + b*f.y2[khi]
end

@fastmath function eval_deriv12(f::Spline1D, x::Float64)
    klo,khi = locate(f.x,x)
    h = f.x[khi] - f.x[klo]  # @assert h != 0.0
    a = (f.x[khi]-x)/h
    b = (x-f.x[klo])/h
    a*f.y[klo] + b*f.y[khi] + ((a*a*a-a)*f.y2[klo]+(b*b*b-b)*f.y2[khi])*(h*h)/6.,
        (f.y[khi] - f.y[klo])/h + ((1.-3.*a*a)*f.y2[klo] + (3*b*b-1.)*f.y2[khi])*h/6.,
            a*f.y2[klo] + b*f.y2[khi]
end

precompile(Spline1D, (Vector{Float64}, Vector{Float64}))
precompile(locate, (Vector{Float64}, Float64))
precompile(Base.call, (Spline1D, Float64))
precompile(Base.call, (Spline1D, Vector{Float64}, Float64))
precompile(Base.call, (Spline1D, Vector{Float64}))
precompile(deriv1, (Spline1D, Float64))
precompile(deriv1, (Spline1D, Vector{Float64}, Float64))
precompile(deriv1, (Spline1D, Vector{Float64}))
precompile(deriv2, (Spline1D, Float64))
precompile(deriv2, (Spline1D, Vector{Float64}, Float64))
precompile(deriv2, (Spline1D, Vector{Float64}))
precompile(eval_deriv1, (Spline1D, Float64))
precompile(eval_deriv2, (Spline1D, Float64))
precompile(eval_deriv12, (Spline1D, Float64))


using Polynomials
using ARMA

"""`noisePSD(tes, freq, [SI_amp])`

Returns (noise,A,B,C,D) where `noise` is a noise power spectral
density in A^2 per Hz, at each input frequency (Hz) given by `freq`.
A,B,C,D are the components of the noise, such that A+B+C+D=noise.
In order, they are TES Johnson noise; load resistor Johnson noise;
thermal fluctuation noise; and amplifier noise.

You can input the amplifier current noise `SI_amp`, in A^2/Hz, or you
can take the default value."""
function noisePSD(tes::IrwinHiltonTES, freq::AbstractVector, SI_amp=5e-22)
    # This term F goes from 0 to one and depends on whether the thermal conductivity is ballistic
    # or diffusive, hardcoded as 1 for now
    F  = 1
    SP_TFN = 4kb*tes.T0^2*tes.G0*F
    SV_TES = 4kb*tes.T0*tes.R0*(1+2*tes.beta) # TES voltage noise
    SV_L   = 4kb*tes.T0*tes.Rl  # Load voltage noise

    omega = 2pi*freq  # Radians / sec
    sIomeg = (1-tes.tauplus/tes.taucc)*(1-tes.tauminus/tes.taucc)./((1 .+im*omega*tes.tauplus).*(1 .+im*omega*tes.tauminus)) /(tes.I0*tes.R0*(2+tes.beta))
    sIomeg2 = abs2.(sIomeg)

    Inoise_TES = SV_TES*tes.I0^2/tes.loopgain^2 * (1 .+(tes.tauthermal*omega).^2) .* sIomeg2
    Inoise_load = SV_L*tes.I0^2*(tes.loopgain-1)^2/tes.loopgain^2 * (1 .+(tes.taucc*omega).^2) .* sIomeg2
    Inoise_TFN = SP_TFN*sIomeg2
    Inoise_amp = SI_amp

    Inoise = Inoise_TFN .+ Inoise_amp .+ Inoise_TES .+ Inoise_load
    Inoise, Inoise_TES, Inoise_load, Inoise_TFN, Inoise_amp .+ zeros(Float64, length(Inoise))
end

noisePSD(tes::BiasedTES, freq::AbstractVector, SI_amp::Float64=5e-22) =
    noisePSD(IrwinHiltonTES(tes), freq, SI_amp)

"""`NoiseModel(tes, sampleTime, [SI_amp])`

Returns noise of an IrwinHiltonTES modeled as an ARMA(2,2) process, specifically as an
ARMAModel object from package ARMA.jl"""
function NoiseModel(tes::IrwinHiltonTES, sampleTime::Float64, SI_amp=5e-22)
    # This term F goes from 0 to one and depends on whether the thermal conductivity is ballistic
    # or diffusive, hardcoded as 1 for now
    F  = 1
    SP_TFN = 4kb*tes.T0^2*tes.G0*F
    SV_TES = 4kb*tes.T0*tes.R0*(1+2*tes.beta) # TES voltage noise
    SV_L   = 4kb*tes.T0*tes.Rl  # Load voltage noise

    A = SI_amp
    B = abs2((1-tes.tauplus/tes.taucc)*(1-tes.tauminus/tes.taucc) /(tes.I0*tes.R0*(2+tes.beta)))
    C = SP_TFN
    D = SV_TES*tes.I0^2/tes.loopgain^2
    E = SV_L*tes.I0^2*(tes.loopgain-1)^2/tes.loopgain^2

    # (u,v) are the solutions to the quadratic equation ax^2+bx+c=0
    a = A * (tes.tauplus*tes.tauminus)^2
    b = -(A*(tes.tauplus^2+tes.tauminus^2)+B*(D*tes.tauthermal^2+E*tes.taucc^2))
    c = A+B*(C+D+E)
    discr = sqrt(b^2-4a*c)
    u2 = (-b-discr)/(2a)
    v2 = (-b+discr)/(2a)

    u,v = sqrt(u2), sqrt(v2)
    K = sqrt(a)
    theta = [(sampleTime*u*.5+1)*(sampleTime*v*0.5+1),
            sampleTime^2*u*v*.5-2,
            (sampleTime*u*.5-1)*(sampleTime*v*.5-1)]
    phi = [(2tes.tauplus/sampleTime+1)*(2tes.tauminus/sampleTime+1),
            2-8tes.tauplus*tes.tauminus/sampleTime^2,
            (1-2tes.tauplus/sampleTime)*(1-2tes.tauminus/sampleTime)]

    # Rescale so that phi has unit coefficients for the z^0 term; theta[1] is sigma.
    sigma = theta[1]/phi[1]*4K/(sampleTime^2)
    sigma /= sqrt(2sampleTime)
    theta .*= sigma/theta[1]
    phi ./= phi[1]

    ARMAModel(real(theta), real(phi))
end

NoiseModel(tes::BiasedTES, sampleTime::Float64, SI_amp::Float64=5e-22) =
    NoiseModel(IrwinHiltonTES(tes), sampleTime, SI_amp)


module TrussComplianceDemo3D1

using TopOpt, LinearAlgebra, StatsFuns
# using Makie
# using TopOpt.TrussTopOptProblems.TrussVisualization: visualize

# 3D
ndim = 3
node_points, elements, mats, crosssecs, fixities, load_cases = load_truss_json(
    joinpath(@__DIR__, "tim_$(ndim)d.json")
)
ndim, nnodes, ncells = length(node_points[1]), length(node_points), length(elements)
loads = load_cases["0"]
problem = TrussProblem(
    Val{:Linear}, node_points, elements, loads, fixities, mats, crosssecs
)

xmin = 0.0001 # minimum density
x0 = fill(1.0, ncells) # initial design
p = 4.0 # penalty
V = 0.5 # maximum volume fraction

solver = FEASolver(Direct, problem; xmin=xmin)
comp = TopOpt.Compliance(problem, solver)

function obj(x)
    # minimize compliance
    return comp(x)
end
function constr(x)
    # volume fraction constraint
    return sum(x) / length(x) - V
end

m = Model(obj)
addvar!(m, zeros(length(x0)), ones(length(x0)))
Nonconvex.add_ineq_constraint!(m, constr)

options = MMAOptions(; maxiter=1000, tol=Tolerance(; kkt=1e-4, f=1e-4))
TopOpt.setpenalty!(solver, p)
@time r = Nonconvex.optimize(
    m, MMA87(; dualoptimizer=ConjugateGradient()), x0; options=options
)

@show r.minimizer
@show obj(r.minimizer)
@show constr(r.minimizer)
# fig = visualize(
#     problem, solver.u; topology = r.minimizer,
#     default_exagg_scale=0.0, exagg_range=100.0,
# )
# Makie.display(fig)

end


function petk05()
#
# M. Petkovic testing polynomials, p142
#
    p = [1.0;-7;20;-28;-18;110;-92;-44;345;225];
    z = [3.00000000000000 + 0.00000000000000im
         1.00000000000000 + 2.00000000000000im
         1.00000000000000 - 2.00000000000000im
        -1.00000000000000 - 0.00000000000000im];
    z = [z [2;2;2;3]];

    p, PolyZeros(z)
end


using LazIO

const c_air, c_water = 2.997*10^8, 2.25*10^8
const n_air_water = c_water/c_air

function correct!(ws::LazIO.LazPoint, uw::LazIO.LazPoint)
    x, y, z = correct((ws.X, ws.Y, ws.Z), (uw.X, uw.Y, uw.Z))
    uw.X = round(Int32, x)
    uw.Y = round(Int32, y)
    uw.Z = round(Int32, z)
    uw
end

"Correct an underwater point according to a watersurface point from the refraction factor by using Snell's Law."
function correct(watersurface::Tuple{Real,Real,Real}, underwater::Tuple{Real,Real,Real})
    x0, y0, z0 = watersurface
    x1, y1, z1 = underwater
    Δx, Δy, Δz = (x1-x0), (y1-y0), (z1-z0)

    if Δz >= 0.
        throw(DomainError(Δz, "The underwater point can not be equal or greater in the z dimension to the watersurface point."))
    end

    xcor, ycor, zcor = correct(Δx, Δy, Δz)

    # Define the new corrected coordinates of the points
    xnew = x0 + xcor
    ynew = y0 + ycor
    znew = z0 - zcor

    return (xnew, ynew, znew)
end

function correct(Δx::Real, Δy::Real, Δz::Real)

    # Calculate the distance of the two points (ρ), the vertical-incidence angle (ϕ) and horizontal angle (θ) between the two points
    ρ, θ, ϕ = cart2pol(Δx, Δy, -Δz)

    # Calculate the underwater angle due to the refraction effect.
    # Using Snell's Law: n(air)*sin(ϕ_air) = n(water)*sin(ϕ_water),
    # where ϕ_air: the angle of incidence (I know that the angle of the scan is 20)
    #      ϕ_water: the angle of refraction
    #      n(air): the refraction idex of medium containing the incident ray, value=1.000293
    #      n(water): the refraction idex of medium containing the transmitted ray, value=1.333
    ϕ_water = asind(n_air_water * sind(ϕ))

    # Distance under the water has been influenced by the different speed of light in the water.
    # In a specific moment, ρ = c_air * t and ρ_water = c_water * t.
    # So, by diving these two equations, the ratio (ρ1/ρ2) = n_air_water
    ρ_water = (ρ * n_air_water)

    # Calculate the correction in the 3d space using the pol2cart function
    pol2cart(ρ_water, θ, ϕ_water)
end


using Base.Test
using SymEngine
CDenseMatrix = SymEngine.CDenseMatrix

@vars x

# constructors
A = [x 1 2; 3 x 4; 5 6 x]
M = convert(CDenseMatrix, A)
M = CDenseMatrix(A)

# abstract interface
@test M[1,1] == A[1,1]
for i in eachindex(M)
    @test M[i] == A[i]
end
M[1,1] = x^2
@test M[1,1] == x^2
M[1,1] = x

# inherits from DenseArray
@test all(M + 3 .== A + 3)
@test all(M * 3 .== A * 3)
@test all(M + M .== A + A)
@test all(M * M .== A * A)
@test all(M' .== A')

# generic det
@test prod([subs(det(M) - det(A), x, i) == 0 for i in 2:10]) == true
@test inv(M) - inv(A) == zeros(Basic, 3,3)

# factorizations
@test lu(M) == lu(A)

A = [x 1 2; 0 x 4; 0 0 x]
b = [1, 2, 3]
M = convert(CDenseMatrix, A)
out = M \ b
@test M * out - b == zeros(Basic, 3, 1)


"""
Created on Thu Apr 23 2020

@author: Yoann Pradat

Cost function for nmf_α
"""

"""
    _α_cost(V, W, H, params)

Compute the cost function that is minimized by the NMF alogrithm running with parameters params. It adds up
the α_divergence and regularization terms.

See also: [`α_divergence`](@ref), [`CParams`](@ref)
"""
function _α_cost(V::Matrix{T}, W::Matrix{T}, H::Matrix{T}, params::CParams) where T <: Real
    regu_H = params.α_H * (params.l₁ratio_H * norm(H, 1) + (1. - params.l₁ratio_H) * norm(H, 2)) / length(H)
    regu_W = params.α_W * (params.l₁ratio_W * norm(W, 1) + (1. - params.l₁ratio_W) * norm(W, 2)) / length(W)
    α_div  = α_divergence(V, W*H, params.α)

    α_div + regu_H + regu_W
end


using StatsBase
function randeig(t, n; xlims=(-2.1*√n, 2.1*√n), bins=100)
    const xgrid = linspace(xlims..., bins)
    counts = @parallel (+) for _ = 1:t
        fit(Histogram, eigvals(Symmetric(randn(n, n))), xgrid).weights
    end
    Histogram(xgrid, counts)
end


# Julia version of Boyd et al's Matlab script available on:
#   https://web.stanford.edu/~boyd/papers/admm/nonconvex/regressor_sel.html

# regressor_sel  Solve lasso problem via ADMM
#
#  x, history = regressor_sel(A, b, K, rho, alpha)
#
# Attempts to solve the following problem via ADMM:
#
#   minimize || Ax - b ||_2^2
#   subject to card(x) <= K
#
# where card() is the number of nonzero entries.
#
# The solution is returned in the vector x.
#
# history is a structure that contains the objective value, the primal and
# dual residual norms, and the tolerances for the primal and dual residual
# norms at each iteration.
#
# rho is the augmented Lagrangian parameter.
#
# alpha is the over-relaxation parameter (typical values for alpha are
# between 1.0 and 1.8).
#
# More information can be found in the paper linked at:
# http://www.stanford.edu/~boyd/papers/distr_opt_stat_learning_admm.html
#
function regressor_sel(A, b, K, rho;
    MaxIterations = 1000, Quiet = true, ABSTOL = 1e-4, RELTOL = 1e-2)

    # Data preprocessing and setup
    K = iceil(K)
    m, n = size(A)
    Atb = A' * b            # save a matrix-vector multiply

    local r_norm::Float64, s_norm::Float64, eps_pri::Float64, eps_dual::Float64
    x = zeros(n)
    z = zeros(n)
    u = zeros(n)
    L, U = factor(A, rho)   # cache the factorization

    for k in 1:MaxIterations
        # x-update
        q = Atb + rho*(z - u)    # temporary value
        if ( m >= n )    # if skinny
            x = U \ (L \ q)
        else            # if fat
            x = q/rho - (A'*(U \ ( L \ (A*q) )))/rho^2
        end
        x = vec(x)

        # z-update with relaxation
        zold = z
        z = keep_largest!(x + u, K)

        # u-update
        u = u + (x - z)

        # Check convergence
        r_norm  = norm(x - z)
        s_norm  = norm(-rho*(z - zold))
        eps_pri = sqrt(n) * ABSTOL + RELTOL * max(norm(x), norm(-z))
        eps_dual= sqrt(n) * ABSTOL + RELTOL * norm(rho*u)

        if !Quiet
            println(@sprintf("%3d\t%10.4f\t%10.4f\t%10.4f\t%10.4f\t%10.2f\n", k,
                 r_norm, eps_pri, s_norm, eps_dual, objective(A,b,x)))
        end

        if (r_norm < eps_pri) && (s_norm < eps_dual)
            break
        end
    end

    return z
end

function objective(A, b, x)
    sum( (A*x - b).^2 )
end

function keep_largest!(z, K)
    pos = sortperm(abs(z), rev=true)
    z[pos[K+1:end]] = 0
    z
end

import Base.factor
function factor(A, rho)
    m, n = size(A)
    if ( m >= n )  # if skinny
        L = chol( A'*A + rho*speye(n), :L )
    else            # if fat
        L = chol( speye(m) + 1/rho*(A*A'), :L )
    end

    return L, L'
end


# this file contains some information theoretic functionals of probability distributions
# should ultimately spin this off into separate package
# functional benchmarked just as fast as direct methods!!!

# TODO make this work for functions with ℵ₀ support
function functional(d::Distribution{T, Discrete}, f::Function, st::Function=(x -> true)) where T
    o = 0.0
    for k ∈ support(d)
        if st(k)
            o += f(k)
        end
    end
    o
end

function functional(d::SampleDistribution, f::Function, st::Function=(x -> true))
    o = 0.0
    for i ∈ 1:ncategories(d)
        if st(getsample(d, i))
            o += f(pdfbyindex(d, i))
        end
    end
    o
end
export functional

Hfunc(p::Real) = -p*log(p)
H(d::SampleDistribution, st::Function=(x -> true)) = functional(d, Hfunc, st)
export H


using ReservoirComputing

#model parameters
const res_size = 30
const radius = 1.2
const activation = tanh
const degree = 6
const sparsity = 0.5
const sigma = 0.1
const beta = 0.0
const alpha = 1.0
const nla_type = NLADefault()
const in_size = 3
const extended_states = false
const h_steps = 2
const max_value = 0.8

const weight = 0.8
const fb_weight = 0.05
const jumps = 2

W_in = init_input_layer(res_size, in_size, sigma)

train_len = 50
predict_len = 12
data = ones(Float64, in_size, 100)
train = data[:, 1:1+train_len-1]
test = data[:, train_len:train_len+predict_len-1]

#test input layers functions
input_layer = [init_reservoir_givendeg, init_reservoir_givensp, DLR, DLRB, SCR, CRJ]

for t in input_layer

    if t == init_reservoir_givendeg
        W = t(res_size, radius, degree)
    elseif t == init_reservoir_givensp
        W = t(res_size, radius, sparsity)
    elseif t == DLRB
        W = t(res_size, weight, fb_weight)
    elseif t == CRJ
        W = t(res_size, weight, fb_weight, jumps)
    else
        W = t(res_size, weight)
    end
    esn = ESN(W,
        train,
        W_in,
        activation = activation,
        alpha = alpha,
        nla_type = nla_type,
        extended_states = extended_states)

    @test size(W, 1) == res_size
end

W = pseudoSVD(res_size, max_value, sparsity)
@test size(W, 1) == res_size


export *,+, basis_matrix, ambient_space, base_ring, base_field, root_lattice,
kernel_lattice, invariant_lattice
# scope & verbose scope: :Lattice

basis_matrix(L::ZLat) = L.basis_matrix

ambient_space(L::ZLat) = L.space

base_ring(L::ZLat) = FlintZZ

base_field(L::ZLat) = base_ring(gram_matrix(ambient_space(L)))

################################################################################
#
#  Creation
#
################################################################################

@doc Markdown.doc"""
    Zlattice([B::MatElem]; gram) -> ZLat

Return the Z-lattice with basis matrix $B$ inside the quadratic space with
Gram matrix `gram`.

If the keyword `gram` is not specified, the Gram matrix is the identity matrix.
If $B$ is not specified, the basis matrix is the identity matrix.

# Examples

```jldoctest
julia> L = Zlattice(matrix(QQ, 2, 2, [1//2, 0, 0, 2]));

julia> gram_matrix(L) == matrix(QQ, 2, 2, [1//4, 0, 0, 4])
true

julia> L = Zlattice(gram = matrix(ZZ, [2 -1; -1 2]));

julia> gram_matrix(L) == matrix(ZZ, [2 -1; -1 2])
true
```
"""
function Zlattice(B::fmpq_mat; gram = identity_matrix(FlintQQ, ncols(B)))
  @req issymmetric(gram) "Gram matrix must be symmetric"
  V = quadratic_space(FlintQQ, gram)
  return lattice(V, B)
end

function Zlattice(B::fmpz_mat; gram = identity_matrix(FlintQQ, ncols(B)))
  @req issymmetric(gram) "Gram matrix must be symmetric"
  V = quadratic_space(FlintQQ, gram)
  return lattice(V, B)
end

function Zlattice(;gram)
  @req issymmetric(gram) "Gram matrix must be symmetric"
  n = nrows(gram)
  return lattice(quadratic_space(FlintQQ, gram), identity_matrix(FlintQQ, n))
end

@doc Markdown.doc"""
    lattice(V::QuadSpace, B::MatElem) -> ZLat

Return the Z-lattice with basis matrix $B$ inside the quadratic space $V$.
"""
function lattice(V::QuadSpace{FlintRationalField, fmpq_mat}, B::MatElem; isbasis::Bool = true, check::Bool = true)
  @req dim(V) == ncols(B) "Ambient space and the matrix B have incompatible dimension"
  local Gc
  try
    Gc = change_base_ring(FlintQQ, B)
  catch e
    throw(error("Cannot convert entries of the matrix to the rationals"))
  end
  if typeof(Gc) !== fmpq_mat
    throw(error("Cannot convert entries of the matrix to the rationals"))
  end

  # We need to produce a basis matrix

  if !isbasis
    BB = fmpq_mat(hnf(FakeFmpqMat(B), :upper_right))
    i = nrows(BB)
    while i > 0 && iszero_row(BB, i)
      i = i - 1
    end
    return ZLat(V, BB[1:i, :])
  else
    return ZLat(V, Gc)
  end
end

function lattice_in_same_ambient_space(L::ZLat, B::MatElem)
  V = L.space
  return lattice(V,B)
end

function rescale(G::ZLat, r)
  B = basis_matrix(G)
  gram_space = gram_matrix(G.space)
  Vr = quadratic_space(QQ, r*gram_space)
  return lattice(Vr, B)
end

################################################################################
#
#  Gram matrix
#
################################################################################

@doc Markdown.doc"""
    gram_matrix(L::ZLat) -> fmpq_mat

Return the gram matrix of $L$.

# Examples
```jldoctest
julia> L = Zlattice(matrix(ZZ, [2 0; -1 2]));

julia> gram_matrix(L)
[ 4   -2]
[-2    5]
```
"""
function gram_matrix(L::ZLat)
  if isdefined(L, :gram_matrix)
    return L.gram_matrix
  end
  b = basis_matrix(L)
  G = b * gram_matrix(ambient_space(L)) * transpose(b)
  L.gram_matrix = G
  return G
end

gram_matrix_of_rational_span(L::ZLat) = gram_matrix(L)

################################################################################
#
#  Rational span
#
################################################################################

@doc Markdown.doc"""
    rational_span(L::ZLat) -> QuadSpace

Return the rational span of $L$, which is the quadratic space with Gram matrix
equal to `gram_matrix(L)`.

# Examples
```jldoctest
julia> L = Zlattice(matrix(ZZ, [2 0; -1 2]));

julia> rational_span(L)
Quadratic space over
Rational Field
with Gram matrix
[4 -2; -2 5]
```
"""
function rational_span(L::ZLat)
  if isdefined(L, :rational_span)
    return L.rational_span
  else
    G = gram_matrix(L)
    V = quadratic_space(FlintQQ, G)
    L.rational_span = V
    return V
  end
end

################################################################################
#
#  Orthogonal sum
#
################################################################################

@doc Markdown.doc"""
    orthogonal_sum(L1::ZLat, L2::ZLat)

Return the orthogonal direct sum of the lattices `L1` and `L2`.
"""
function orthogonal_sum(L1::ZLat, L2::ZLat)
  V1 = ambient_space(L1)
  V2 = ambient_space(L2)
  B1 = basis_matrix(L1)
  B2 = basis_matrix(L2)
  B = diagonal_matrix(B1, B2)
  V, f1, f2 = orthogonal_sum(V1, V2)
  return lattice(V, B), f1, f2
end

@doc Markdown.doc"""
    orthogonal_submodule(L::ZLat, S::ZLat) -> ZLat

Return the largest submodule of `L` orthogonal to `S`.
"""
function orthogonal_submodule(L::ZLat, S::ZLat)
  @assert ambient_space(L)==ambient_space(S) "L and S must have the same ambient space"
  B = basis_matrix(L)
  C = basis_matrix(S)
  V = ambient_space(L)
  G = gram_matrix(V)
  M = B * G * transpose(C)
  _, K = left_kernel(M)
  K = change_base_ring(ZZ, K*denominator(K))
  Ks = saturate(K)
  return lattice(V, Ks*B)
end
################################################################################
#
#  String I/O
#
################################################################################

function Base.show(io::IO, L::ZLat)
  print(io, "Quadratic lattice of rank ", rank(L),
            " and degree ", degree(L), " over the rationals")
end

################################################################################
#
#  Automorphism groups
#
################################################################################

# This is an internal function, which sets
# L.automorphism_group_generators
# L.automorphism_group_order
function assert_has_automorphisms(L::ZLat; redo::Bool = false,
                                           try_small::Bool = true)

  if !redo && isdefined(L, :automorphism_group_generators)
    return nothing
  end

  if rank(L) == 0
    L.automorphism_group_generators = fmpz_mat[identity_matrix(ZZ, 0)]
    L.automorphism_group_order = one(fmpz)
    return nothing
  end

  V = ambient_space(L)
  GL = gram_matrix(L)
  d = denominator(GL)
  res = fmpz_mat[change_base_ring(FlintZZ, d * GL)]
  # So the first one is either positive definite or negative definite
  # Make it positive definite. This does not change the automorphisms.
  if res[1][1, 1] < 0
    res[1] = -res[1]
  end
  Glll, T = lll_gram_with_transform(res[1])
  Ttr = transpose(T)
  res_orig = copy(res)
  res[1] = Glll

  bm = basis_matrix(L)

  # Make the Gram matrix small

  C = ZLatAutoCtx(res)
  if try_small
    fl, Csmall = try_init_small(C)
    if fl
      auto(Csmall)
      _gens, order = _get_generators(Csmall)
      gens = fmpz_mat[matrix(FlintZZ, g) for g in _gens]
    else
      init(C)
      auto(C)
      gens, order = _get_generators(C)
    end
  else
    init(C)
    auto(C)
    gens, order = _get_generators(C)
  end

  # Now translate back
  Tinv = inv(T)
  for i in 1:length(gens)
    gens[i] = Tinv * gens[i] * T
  end

  # Now gens are with respect to the basis of L
  @hassert :Lattice 1 all(change_base_ring(FlintQQ, gens[i]) * GL *
                          transpose(change_base_ring(FlintQQ, gens[i])) == GL
                                                        for i in 1:length(gens))

  L.automorphism_group_generators = gens
  L.automorphism_group_order = order

  return nothing
end

# documented in ../Lattices.jl

function automorphism_group_generators(L::ZLat; ambient_representation::Bool = true)

  @req rank(L) == 0 || isdefinite(L) "The lattice must be definite"
  assert_has_automorphisms(L)

  gens = L.automorphism_group_generators
  if !ambient_representation
    return fmpq_mat[ change_base_ring(FlintQQ, g) for g in gens]
  else
    # Now translate to get the automorphisms with respect to basis_matrix(L)
    bm = basis_matrix(L)
    gens = L.automorphism_group_generators
    V = ambient_space(L)
    if rank(L) == rank(V)
      bminv = inv(bm)
      res = fmpq_mat[bminv * change_base_ring(FlintQQ, g) * bm for g in gens]
    else
      # Extend trivially to the orthogonal complement of the rational span
      !isregular(V) &&
        throw(error(
          """Can compute ambient representation only if ambient space is
             regular"""))
      C = orthogonal_complement(V, basis_matrix(L))
      C = vcat(basis_matrix(L), C)
      Cinv = inv(C)
      D = identity_matrix(FlintQQ, rank(V) - rank(L))
      res = fmpq_mat[Cinv * diagonal_matrix(change_base_ring(FlintQQ, g), D) * C for g in gens]
    end
    @hassert :Lattice 1 all(g * gram_matrix(V) * transpose(g) == gram_matrix(V)
                            for g in res)
    return res
  end
end

# documented in ../Lattices.jl

function automorphism_group_order(L::ZLat)
  assert_has_automorphisms(L)
  return L.automorphism_group_order
end

################################################################################
#
#  Isometry
#
################################################################################

# documented in ../Lattices.jl

function isisometric(L::ZLat, M::ZLat; ambient_representation::Bool = true)
  @req isdefinite(L) && isdefinite(M) "The lattices must be definite"

  if rank(L) != rank(M)
    return false, zero_matrix(FlintQQ, 0, 0)
  end

  if rank(L) == 0
    return true, identity_matrix(FlintQQ, 0, 0)
  end

  i = sign(gram_matrix(L)[1,1])
  j = sign(gram_matrix(M)[1,1])
  @req i==j "The lattices must have the same signatures"

  if i < 0
    L = rescale(L,-1)
    M = rescale(M,-1)
  end

  GL = gram_matrix(L)
  dL = denominator(GL)
  GLint = change_base_ring(FlintZZ, dL * GL)
  cL = content(GLint)
  GLint = divexact(GLint, cL)

  GM = gram_matrix(M)
  dM = denominator(GM)
  GMint = change_base_ring(FlintZZ, dM * GM)
  cM = content(GMint)
  GMint = divexact(GMint, cM)

  # GLint, GMint are integral, primitive scalings of GL and GM
  # If they are isometric, then the scalars must be identitcal.
  if dL//cL != dM//cM
    return false, zero_matrix(FlintQQ, 0, 0)
  end

  # Now compute LLL reduces gram matrices

  GLlll, TL = lll_gram_with_transform(GLint)
  @hassert :Lattice 1 TL * change_base_ring(FlintZZ, GL) * transpose(TL) * dL == GLlll *cL
  GMlll, TM = lll_gram_with_transform(GMint)
  @hassert :Lattice 1 TM * change_base_ring(FlintZZ, GM) * transpose(TM) * dM == GMlll *cM

  # Setup for Plesken--Souvignier

  G1 = fmpz_mat[GLlll]
  G2 = fmpz_mat[GMlll]

  fl, CLsmall, CMsmall = _try_iso_setup_small(G1, G2)
  if fl
    b, _T = isometry(CLsmall, CMsmall)
    T = matrix(FlintZZ, _T)
  else
    CL, CM = _iso_setup(fmpz_mat[GLlll], fmpz_mat[GMlll])
    b, T = isometry(CL, CM)
  end

  if b
    T = change_base_ring(FlintQQ, inv(TL)*T*TM)
    if !ambient_representation
      @hassert :Lattice 1 T * gram_matrix(M) * transpose(T) == gram_matrix(L)
      return true, T
    else
      V = ambient_space(L)
      W = ambient_space(M)
      if rank(L) == rank(V)
        T = inv(basis_matrix(L)) * T * basis_matrix(M)
      else
        (!isregular(V) || !isregular(W)) &&
          throw(error(
            """Can compute ambient representation only if ambient space is
               regular"""))
          (rank(V) != rank(W)) &&
          throw(error(
            """Can compute ambient representation only if ambient spaces
            have the same dimension."""))

        CV = orthogonal_complement(V, basis_matrix(L))
        CV = vcat(basis_matrix(L), CV)
        CW = orthogonal_complement(W, basis_matrix(M))
        CW = vcat(basis_matrix(M), CW)
        D = identity_matrix(FlintQQ, rank(V) - rank(L))
        T = inv(CV) * diagonal_matrix(T, D) * CW
      end
      @hassert :Lattice 1 T * gram_matrix(ambient_space(M))  * transpose(T) ==
                  gram_matrix(ambient_space(L))
      return true, T
    end
  else
    return false, zero_matrix(FlintQQ, 0, 0)
  end
end

################################################################################
#
#  Is sublattice?
#
################################################################################

function issublattice(M::ZLat, N::ZLat)
  if ambient_space(M) != ambient_space(N)
    return false
  end

  hassol, _rels = can_solve_with_solution(basis_matrix(M), basis_matrix(N), side=:left)

  if !hassol || !isone(denominator(_rels))
    return false
  end

  return true
end

@doc Markdown.doc"""
    issublattice_with_relations(M::ZLat, N::ZLat) -> Bool, fmpq_mat

Returns whether $N$ is a sublattice of $M$. In this case, the second return
value is a matrix $B$ such that $B B_M = B_N$, where $B_M$ and $B_N$ are the
basis matrices of $M$ and $N$ respectively.
"""
function issublattice_with_relations(M::ZLat, N::ZLat)
   if ambient_space(M) != ambient_space(N)
     return false, basis_matrix(M)
   end

   hassol, _rels = can_solve_with_solution(basis_matrix(M), basis_matrix(N), side=:left)

   if !hassol || !isone(denominator(_rels))
     return false, basis_matrix(M)
   end

   return true, _rels
 end

################################################################################
#
#  Root lattice
#
################################################################################

@doc Markdown.doc"""
    root_lattice(R::Symbol, n::Int)

Return the root lattice of type `R` given by ``:A``, ``:D`` or ``:E`` with parameter `n`.
"""
function root_lattice(R::Symbol, n::Int)
  if R === :A
    return Zlattice(gram = _root_lattice_A(n))
  elseif R === :E
    return Zlattice(gram = _root_lattice_E(n))
  elseif R === :D
    return Zlattice(gram = _root_lattice_D(n))
  else
    error("Type (:$R) must be :A, :D or :E")
  end
end

function _root_lattice_A(n::Int)
  n < 0 && error("Parameter ($n) for root lattice of type :A must be positive")
  z = zero_matrix(FlintQQ, n, n)
  for i in 1:n
    z[i, i] = 2
    if i > 1
      z[i, i - 1] = -1
    end
    if i < n
      z[i, i + 1] = -1
    end
  end
  return z
end

function _root_lattice_D(n::Int)
  n < 2 && error("Parameter ($n) for root lattices of type :D must be greater or equal to 2")
  if n == 2
    G = matrix(ZZ, [2 0 ;0 2])
  elseif n == 3
    return _root_lattice_A(n)
  else
    G = zero_matrix(ZZ, n, n)
    G[1,3] = G[3,1] = -1
    for i in 1:n
      G[i,i] = 2
      if 2 <= i <= n-1
        G[i,i+1] = G[i+1,i] = -1
      end
    end
  end
  return G
end

function _root_lattice_E(n::Int)
  n in [6,7,8] || error("Parameter ($n) for lattice of type :E must be 6, 7 or 8")
  if n == 6
    G = [2 -1 0 0 0 0;
        -1 2 -1 0 0 0;
        0 -1 2 -1 0 -1;
        0 0 -1 2 -1 0;
        0 0 0 -1 2 0;
        0 0 -1 0 0 2]
  elseif n == 7
    G = [2 -1 0 0 0 0 0;
        -1 2 -1 0 0 0 0;
        0 -1 2 -1 0 0 -1;
        0 0 -1 2 -1 0 0;
        0 0 0 -1 2 -1 0;
        0 0 0 0 -1 2 0;
        0 0 -1 0 0 0 2]
  else
    G = [2 -1 0 0 0 0 0 0;
        -1 2 -1 0 0 0 0 0;
        0 -1 2 -1 0 0 0 -1;
        0 0 -1 2 -1 0 0 0;
        0 0 0 -1 2 -1 0 0;
        0 0 0 0 -1 2 -1 0;
        0 0 0 0 0 -1 2 0;
        0 0 -1 0 0 0 0 2]
  end
  return matrix(ZZ, G)
end

################################################################################
#
#  Dual lattice
#
################################################################################

# documented in ../Lattices.jl

function dual(L::ZLat)
  G = gram_matrix(L)
  new_bmat = inv(G)*basis_matrix(L)
  return lattice(ambient_space(L), new_bmat)
end

################################################################################
#
#  Scale
#
################################################################################

function scale(L::ZLat)
  if isdefined(L, :scale)
    return L.scale
  end
  G = gram_matrix(L)
  s = zero(fmpq)
  for i in 1:nrows(G)
    for j in 1:i
      s = gcd(s, G[i, j])
    end
  end
  L.scale = s
  return s
end

################################################################################
#
#  Norm
#
################################################################################

function norm(L::ZLat)
  if isdefined(L, :norm)
    return L.norm
  end
  n = 2 * scale(L)
  G = gram_matrix(L)
  for i in 1:nrows(G)
    n = gcd(n, G[i, i])
  end
  L.norm = n
  return n
end

################################################################################
#
#  Eveness
#
################################################################################

iseven(L::ZLat) = isintegral(L) && iseven(numerator(norm(L)))

################################################################################
#
#  Discriminant
#
################################################################################

function discriminant(L::ZLat)
  d = det(gram_matrix(L))
  return d
end

################################################################################
#
#  Determinant
#
################################################################################

function det(L::ZLat)
  return det(gram_matrix(L))
end

################################################################################
#
#  Rank
#
################################################################################

function rank(L::ZLat)
  return nrows(basis_matrix(L))
end

################################################################################
#
#  Local basis matrix
#
################################################################################

# so that abstract lattice functions also work with Z-lattices

local_basis_matrix(L::ZLat, p) = basis_matrix(L)

################################################################################
#
#  Intersection
#
################################################################################

function intersect(M::ZLat, N::ZLat)
  @req ambient_space(M) === ambient_space(N) "Lattices must have same ambient space"
  BM = basis_matrix(M)
  BN = basis_matrix(N)
  dM = denominator(BM)
  dN = denominator(BN)
  d = lcm(dM, dN)
  BMint = change_base_ring(FlintZZ, d * BM)
  BNint = change_base_ring(FlintZZ, d * BN)
  H = vcat(BMint, BNint)
  k, K = left_kernel(H)
  BI = divexact(change_base_ring(FlintQQ, hnf(view(K, 1:k, 1:nrows(BM)) * BMint)), d)
  return lattice(ambient_space(M), BI)
end

################################################################################
#
#  Sum
#
################################################################################

function +(M::ZLat, N::ZLat)
  @req ambient_space(M) === ambient_space(N) "Lattices must have same ambient space"
  BM = basis_matrix(M)
  BN = basis_matrix(N)
  B = fmpq_mat(hnf(FakeFmpqMat(vcat(BM, BN))))
  i = 1
  while iszero_row(B, i)
    i += 1
  end
  return lattice(ambient_space(M), B[i:end, 1:ncols(B)])
end

################################################################################
#
#  Local isometry
#
################################################################################

# TODO: This is slow. We need a dedicated implementation

function islocally_isometric(L::ZLat, M::ZLat, p::Int)
  return islocally_isometric(L, M, fmpz(p))
end

function islocally_isometric(L::ZLat, M::ZLat, p::fmpz)
  LL = _to_number_field_lattice(L)
  K = base_field(LL)
  MM = _to_number_field_lattice(L, K = K)
  OK = maximal_order(K)
  P = prime_decomposition(OK, p)[1][1]
  return islocally_isometric(LL, MM, P)
end

################################################################################
#
#  Conversion between ZLat and QuadLat
#
################################################################################

function _to_number_field_lattice(L::ZLat, K, V)
  LL = lattice(V, change_base_ring(K, basis_matrix(L)))
  return LL
end

function _to_number_field_lattice(L::ZLat;
                                  K::AnticNumberField = rationals_as_number_field()[1],
                                  V::QuadSpace = quadratic_space(K, gram_matrix(ambient_space(L))))
  return _to_number_field_lattice(L, K, V)
end

function _to_ZLat(L::QuadLat, K, V)
  pm = pseudo_matrix(L)
  cm = coefficient_ideals(pm)
  pmm = matrix(pm)
  bm = zero_matrix(FlintQQ, rank(L), dim(V))
  for i in 1:nrows(pm)
    a = norm(cm[i])
    for j in 1:ncols(pm)
      bm[i, j] = a * FlintQQ(pmm[i, j])
    end
  end
  return lattice(V, bm)
end

function _to_ZLat(L::QuadLat;
                  K::FlintRationalField = FlintQQ,
                  V::QuadSpace = quadratic_space(K, map_entries(FlintQQ, gram_matrix(ambient_space(L)))))
  return _to_ZLat(L, K, V)
end

################################################################################
#
#  Mass
#
################################################################################

mass(L::ZLat) = mass(_to_number_field_lattice(L))

################################################################################
#
#  Genus representatives
#
################################################################################

function genus_representatives(L::ZLat)
  LL = _to_number_field_lattice(L)
  K = base_field(L)
  G = genus_representatives(LL)
  res = ZLat[]
  for N in G
    push!(res, _to_ZLat(N, K = K))
  end
  return res
end

################################################################################
#
#  Maximal integral lattice
#
################################################################################

function maximal_integral_lattice(L::ZLat)
  LL = _to_number_field_lattice(L)
  M = maximal_integral_lattice(LL)
  return _to_ZLat(M, V = ambient_space(L))
end

################################################################################
#
#  Scalar multiplication
#
################################################################################

@doc Markdown.doc"""
    *(a, L::ZLat) -> ZLat

Returns the lattice $aM$ inside the ambient space of $M$.
"""
function Base.:(*)(a, L::ZLat)
  @assert has_ambient_space(L)
  B = a*L.basis_matrix
  return lattice_in_same_ambient_space(L, B)
end

function Base.:(*)(L::ZLat, a)
  return a * L
end

################################################################################
#
#  Equality
#
################################################################################

@doc Markdown.doc"""
Return ``true`` if both lattices have the same ambient quadratic space
and the same underlying module.
"""
function Base.:(==)(L1::ZLat, L2::ZLat)
  V1 = ambient_space(L1)
  V2 = ambient_space(L2)
  if V1 != V2
    return false
  end
  B1 = basis_matrix(L1)
  B2 = basis_matrix(L2)
  return hnf(FakeFmpqMat(B1)) == hnf(FakeFmpqMat(B2))
end

@doc Markdown.doc"""
    local_modification(M::ZLat, L::ZLat, p)

Return a local modification of `M` that matches `L` at `p`.

INPUT:

- ``M`` -- a `\ZZ_p`-maximal lattice
- ``L`` -- the a lattice
            isomorphic to `M` over `\QQ_p`
- ``p`` -- a prime number

OUTPUT:

an integral lattice `M'` in the ambient space of `M` such that `M` and `M'` are locally equal at all
completions except at `p` where `M'` is locally isometric to the lattice `L`.
"""
function local_modification(M::ZLat, L::ZLat, p)
  # notation
  d = denominator(inv(gram_matrix(L)))
  level = valuation(d,p)
  d = p^(level+1) # +1 since scale(M) <= 1/2 ZZ

  @req isisometric(L.space, M.space, p) "quadratic spaces must be locally isometric at m"
  L_max = maximal_integral_lattice(L)

  # invert the gerstein operations
  GLm, U = padic_normal_form(gram_matrix(L_max), p, prec=level+3)
  B1 = inv(U*basis_matrix(L_max))

  GM, UM = padic_normal_form(gram_matrix(M), p, prec=level+3)
  # assert GLm == GM at least modulo p^prec
  B2 = B1 * UM * basis_matrix(M)
  Lp = lattice(M.space, B2)

  # the local modification
  S = intersect(Lp, M) + d * M
  # confirm result
  @hassert :Lattice 2 genus(S, p)==genus(L, p)
  return S
end

################################################################################
#
#  Kernel lattice
#
################################################################################

@doc Markdown.doc"""
    kernel_lattice(L::ZLat, f::MatElem;
                   ambient_representation::Bool = true) -> ZLat

Given a $\mathbf{Z}$-lattice $L$ and a matrix $f$ inducing an endomorphism of
$L$, return $\ker(f)$ is a sublattice of $L$.

If `ambient_representation` is `true` (the default), the endomorphism is
represented with respect to the ambient space of $L$. Otherwise, the
endomorphism is represented with respect to the basis of $L$.
"""
function kernel_lattice(L::ZLat, f::MatElem; ambient_representation::Bool = true)
  bL = basis_matrix(L)
  if ambient_representation
    if !issquare(bL)
      fl, finL = can_solve_with_solution(bL, bL * f, side = :left)
      @req fl "f must preserve the lattice L"
    else
      finL = bL * f * inv(bL)
    end
  else
    finL = f
  end
  k, K = left_kernel(change_base_ring(ZZ, finL))
  return lattice(ambient_space(L), K * basis_matrix(L))
end

################################################################################
#
#  Invariant lattice
#
################################################################################

@doc Markdown.doc"""
    invariant_lattice(L::ZLat, G::Vector{MatElem};
                      ambient_representation::Bool = true) -> ZLat

Given a $\mathbf{Z}$-lattice $L$ and a list of matrices $G$ inducing
endomorphisms of $L$, return the lattice $L^G$, consisting of elements fixed by
$G$.

If `ambient_representation` is `true` (the default), the endomorphism is
represented with respect to the ambient space of $L$. Otherwise, the
endomorphism is represented with respect to the basis of $L$.
"""
function invariant_lattice(L::ZLat, G::Vector{<:MatElem};
                           ambient_representation::Bool = true)
  if length(G) == 0
    return L
  end

  M = kernel_lattice(L, G[1] - 1,
                     ambient_representation = ambient_representation)
  for i in 2:length(G)
    N = kernel_lattice(L, G[i] - 1,
                       ambient_representation = ambient_representation)
    M = intersect(M, N)
  end
  return M
end

function invariant_lattice(L::ZLat, G::MatElem;
                           ambient_representation::Bool = true)
  return kernel_lattice(L, G - 1, ambient_representation = ambient_representation)
end

################################################################################
#
#  Membership check
#
################################################################################

@doc Markdown.doc"""
    Base.in(v::Vector, L::ZLat) -> Bool

  Check if the vector `v` lies in the lattice `L` or not.
"""
function Base.in(v::Vector, L::ZLat)
  @assert length(v) == degree(L) "The vector should have the same length as the degree of the lattice."
  V = matrix(QQ, 1, length(v), v)
  return V in L
end

@doc Markdown.doc"""
    Base.in(v::fmpq_mat, L::ZLat) -> Bool

  Check if the row span of `v` lies in the lattice `L` or not.
"""
function Base.in(v::fmpq_mat, L::ZLat)
  @assert ncols(v)==degree(L) "The vector should have the same length as the degree of the lattice."
  B = basis_matrix(L)
  fl, w = can_solve_with_solution(B, v, side=:left)
  return fl && isone(denominator(w))
end


function huber_loss(ŷ, y)
    δ = 0.5
    error = abs.(ŷ - y)
    cond = error .< δ
    squared_loss = 0.5 * error.^2
    linear_loss = δ .* (error .- δ^2 .* 0.5)

    is_squared = Int.(cond)
    is_linear = Int.(.!cond)

    # why this line doesn't work?
    # Int.(cond) .* squared_loss + Int.(.!cond) .* (10 .* linear_loss)
    #=
    ERROR: LoadError: MethodError: no method matching !(::ForwardDiff.Dual{Nothing,Bool,1})
    Closest candidates are:
        !(!Matched::Missing) at missing.jl:83
        !(!Matched::Bool) at bool.jl:35
        !(!Matched::Function) at operators.jl:853
    =#
    is_squared .* squared_loss + is_linear .* (10 .* linear_loss)
end

function huber_loss_mean(ŷ, y)
    mean(huber_loss(ŷ, y))
end

function mse_rnn(x, y)
    l = Flux.mse(m(x), y)
    Flux.reset!(m)
    return l
end


"""
Imbibition cell, IMPES
Buckley Leverett equation
dependent variables: pressure and water saturation
Prepared for educational purposes by ** AAE **
The only input variable is 0<=SF0<=1,
is the degree of wettability
"""
function forced_imb_impes(SF0, SW0)
Nx = 100 # number of cells in x direction
Ny = 40 # number of cells in y direction
W = 0.02 # [m] length of the domain in x direction
H = 0.07 # [m] length of the domain in y direction
A_core=π*W^2/4
V_core=H*A_core
m = createMesh1D(Nx, H)
# m = createMeshCylindrical2D(Nx, Ny, W, H) # creates a 2D mesh
## define the physical parametrs
# all the Corey-type relperm parametrs are defined for the oil-wet and water-wet
# cases
k0 = 0.001e-12 # [m^2] average reservoir permeability
phi0 = 0.35 # average porosity
mu_oil = 2e-3 # [Pa.s] oil viscosity
mu_water = 1e-3 # [Pa.s] water viscosity
krw0_ww = 0.3
krw0_ow = 1.0
kro0_ww = 0.6
kro0_ow = 0.76
nw_ww = 2.4
nw_ow= 2.4
no_ww = 2.0
no_ow= 2.0
sor_ww=0.15
sor_ow=0.2
sor_mw=0.08
swc_ww=0.09
swc_ow=0.09
teta_ow=deg2rad(180)
teta_ww=deg2rad(0)
teta_mw=deg2rad(90)
teta_prime=teta_mw-π/2 # see my report
a=(sor_mw-0.5*(sor_ww-sor_ow)*cos(teta_mw)-0.5*(sor_ww+sor_ow))/(cos(teta_mw)^2-1)
b=0.5*(sor_ww-sor_ow)
c=0.5*(sor_ww+sor_ow)-a

gama_ow=0.03 # N/m
labda=2.4 # for Ekofisk chalk
r_ave=sqrt(k0/phi0) #  meter average pore diameter
pce0= 5e2 #2*gama_ow*cos(teta_ww)/r_ave/10 # Pa capillary entry pressure
# SF0=1.0
SF=createFaceVariable(m, SF0) # 1 is water wet, 0 is oil wet
krw0=krw0_ww*SF+krw0_ow*(1-SF)
kro0=kro0_ww*SF+kro0_ow*(1-SF)
# sor=sor_ww*SF+sor_ow*(1-SF)
swc=swc_ww*SF+swc_ow*(1-SF)
no= no_ww*SF+no_ow*(1-SF)
nw= nw_ww*SF+nw_ow*(1-SF)
teta0=teta_ww*SF0+teta_ow*(1-SF0)
# sor0=(sor_ow-sor_mw)*cos(teta0-teta_prime)^2+sor_mw
sor0=a*cos(teta0)^2+b*cos(teta0)+c
teta=createFaceVariable(m, teta0)
sor=createFaceVariable(m, sor0)
pce=createFaceVariable(m, pce0)

p0 = 100e5 # [bar] pressure
pin = 150e5 # [bar] injection pressure at the left boundary
u_in= 1.0/(24*3600) # [m/s] equal to 1 m/day
sw0 = swc_ww+0.01
# sw0(10:end-10, 10:end-10)=swc+0.2
# sw0 = swc+0.1 # initial water saturation
sw_in = 1

eps1=1e-6
clx=0.05
cly=0.05
V_dp=0.01 # Dykstra-Parsons coef.
perm_val= k0 #field2d(Nx,Ny,k0,V_dp,clx,cly)

k=createCellVariable(m, perm_val)
phi=createCellVariable(m, phi0)

lw = geometricMean(k)/mu_water
lo = geometricMean(k)/mu_oil

## Define the boundaries: all fixed Sw=1, fixed pressure everywhere(?)
BCp = createBC(m) # Neumann BC for pressure
BCs = createBC(m) # Neumann BC for saturation
# left boundary pressure gradient
# BCp.left.a[:]=(krw(sw_in)*lw.xvalue(1,:)+kro(sw_in)*lo.xvalue(1,:)) BCp.left.b[:]=0 BCp.left.c[:]=-u_in
# change the right boandary to constant pressure (Dirichlet)
BCp.left.a[:]=-k0/mu_water
BCp.left.b[:]=0
BCp.left.c[:]=u_in
BCp.right.a[:]=0.0
BCp.right.b[:]=1.0
BCp.right.c[:]=p0
# BCp.top.a[:]=0.0
# BCp.top.b[:]=1.0
# BCp.top.c[:]=p0
# BCp.bottom.a[:]=0.0
# BCp.bottom.b[:]=1.0
# BCp.bottom.c[:]=p0
# change the left boundary to constant saturation (Dirichlet)
# BCs.left.a[:]=0 BCs.left.b[:]=1 BCs.left.c[:]=1.0-sor
# calculate the water saturation at the boundaries
f1(sw)=pc_imb(sw, pce0, swc_ww*SF0+swc_ow*(1-SF0), sor_ww*SF0+sor_ow*(1-SF0),
  teta=teta_ww*SF0+teta_ow*(1-SF0))
sw_b=fzero(f1, 0.5)
println(sw_b)
BCs.right.a[:]=0.0
BCs.right.b[:]=1.0
BCs.right.c[:]=sw_b
BCs.left.a[:]=0.0
BCs.left.b[:]=1.0
BCs.left.c[:]=sw_in

# BCs.top.a[:]=0.0
# BCs.top.b[:]=1.0
# BCs.top.c[:]=sw_b
# BCs.bottom.a[:]=0.0
# BCs.bottom.b[:]=1.0
# BCs.bottom.c[:]=sw_b
## define the time step and solver properties
# dt = 1000 # [s] time step
# dt=(W/Nx)/u_in/20 # [s]
pv_inj=5
t_end = pv_inj*V_core/(u_in*A_core) # [s] final time
dt=H/Nx/u_in
eps_p = 1e-7 # pressure accuracy
eps_sw = 1e-7 # saturation accuracy
## define the variables
sw_old = createCellVariable(m, SW0, BCs)
p_old = createCellVariable(m, p0, BCp)
sw = copyCell(sw_old)
oil_init=domainInt(1-sw_old)
p = copyCell(p_old)
uw = -gradientTerm(p_old) # an estimation of the water velocity
## start the main loop
# generate intial pressure profile (necessary to initialize the fully
# implicit solver)
rec_fact=zeros(1)
t_day=zeros(1)
t = 0.0
dt0=dt
dsw_alwd= 0.01
dp_alwd= 100.0 # Pa
labdaw=createFaceVariable(m,1.0) # initialize
# prog_1=ProgressThresh(t_end)
while (t<t_end)
  print("progress is $((t/t_end*100)) %", "\u1b[1G")
  # for i=1:5
    error_p = 1e5
    error_sw = 1e5
    # Implicit loop
#     while ((error_p>eps_p) || (error_sw>eps_sw))
    while(true)
        # calculate parameters
        pgrad = gradientTerm(p)
#         pcgrad=gradientTerm(pc(sw))
        sw_face = upwindMean(sw, -labdaw.*pgrad) # average value of water saturation
        sw_grad=gradientTerm(sw)
        sw_ave=arithmeticMean(sw)
        pcgrad=faceEval(dpc_imb, sw_ave, pce, swc, sor, teta).*sw_grad
        # solve for pressure at known Sw
        labdao = lo.*faceEval(kro, sw_face, kro0, sor, swc, no)
        labdaw = lw.*faceEval(krw, sw_face, krw0, sor, swc, nw)
        labda = labdao+labdaw
        # compute [Jacobian] matrices
        Mdiffp1 = diffusionTerm(-labda)
        RHSpc1=divergenceTerm(labdao.*pcgrad)
        Mbcp, RHSbcp = boundaryConditionTerm(BCp)
        RHS1 = RHSpc1+RHSbcp # with capillary
        p_new=solvePDE(m, Mdiffp1+Mbcp, RHS1)

        # solve for Sw
        pgrad = gradientTerm(p_new)
        uw=-labdaw.*pgrad
        Mbcsw, RHSbcsw = boundaryConditionTerm(BCs)
        RHS_sw=-divergenceTerm(uw)
        sw_new=solveExplicitPDE(sw_old, dt, RHS_sw, BCs, phi)

        error_p = maximum(abs((internalCells(p_new)-internalCells(p))./internalCells(p_new)))
        error_sw = maximum(abs(internalCells(sw_new)-internalCells(sw)))
        dt_new=dt*min(dp_alwd/error_p, dsw_alwd/error_sw)
        # print("sw_error = $error_sw \n")
        # print("new time step = $dt \n")
        # print(internalCells(sw_new))
        # print(internalCells(sw_old))
        # sleep(1.0)
        # assign new values of p and sw
        if error_sw>dsw_alwd
            dt=dt*(dsw_alwd/error_sw)
            # print("new time step = $dt \n")
        else
            t=t+dt
            p = copyCell(p_new)
            sw = copyCell(sw_new)
            p_old = copyCell(p)
            sw_old = copyCell(sw)
            dt=min(dt*(dsw_alwd/error_sw), 15*dt, t_end-t-eps())
            break
        end
    end

    rec_fact=push!(rec_fact, (oil_init-domainInt(1-sw))/oil_init)
    t_day=push!(t_day, t)
    # print(t)
    # GR.plot(sw.value[2:end-1])
    # GR.imshow(sw.value[2:end-1,2:end-1])
    # visualizeCells(1-sw)
    # GR.plot(t_day/3600/24, rec_fact)
    # xlabel('time [day]')
    # ylabel('recovery factor')
    # title([num2str(t/3600/24) ' day']) drawnow
end
t_day, rec_fact, sw
end # imb_impes


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Functions related to the conversion from quaternion to DCM.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export quat_to_dcm

"""
    quat_to_dcm(q::Quaternion)

Convert the quaternion `q` to a Direction Cosine Matrix (DCM).

# Examples

```jldoctest
julia> q = Quaternion(cosd(45/2), sind(45/2), 0, 0);

julia> quat_to_dcm(q)
3×3 StaticArrays.SMatrix{3, 3, Float64, 9} with indices SOneTo(3)×SOneTo(3):
 1.0   0.0       0.0
 0.0   0.707107  0.707107
 0.0  -0.707107  0.707107
```
"""
function quat_to_dcm(q::Quaternion)
    # Auxiliary variables.
    q0 = q.q0
    q1 = q.q1
    q2 = q.q2
    q3 = q.q3

    return DCM(
        q0^2 + q1^2 - q2^2 - q3^2,   2(q1 * q2 + q0 * q3)   ,   2(q1 * q3 - q0 * q2),
          2(q1 * q2 - q0 * q3)   , q0^2 - q1^2 + q2^2 - q3^2,   2(q2 * q3 + q0 * q1),
          2(q1 * q3 + q0 * q2)   ,   2(q2 * q3 - q0 * q1)   , q0^2 - q1^2 - q2^2 + q3^2
    )'
end


using Convex, SCS
set_default_solver(SCSSolver(verbose = 0, max_iters = 1000000))

println(
    "small test to compile and see if we throw warnings when warmstarting what can't be warmstarted.",
)
n = 10
y = rand(n)
x = Variable(n)

lambda = [100]
problem = minimize(sumsquares(y - x) + lambda[1] * sumsquares(x - 10))
@time solve!(problem, warmstart = true)

println("\nnow try first solving and then warmstarting")
n = 1000
y = rand(n)
x = Variable(n)

lambda = [100]
problem = minimize(sumsquares(y - x) + lambda[1] * sumsquares(x - 10))
@time solve!(problem)
@show problem.optval
lambda[1] = 105
println(
    "this run should be faster than the last and have slightly higher optimal value",
)
@time solve!(problem, warmstart = true)
@show problem.optval

println("\nnow try with scalar variables")
lambda = 100
problem = minimize(sumsquares(y - x) + lambda * sumsquares(x - 10))
@time solve!(problem)
@show problem.optval
lambda = 105
println(
    "this run should be faster than the last and have slightly higher optimal value",
)
@time solve!(problem, warmstart = true)
@show problem.optval


"""
**Paired Differences Index**

    pdi{T<:Number}(x::Array{T, 1})

This function will range the values of each row, so that the strongest link has
a value of one. This works for deterministic and quantitative networks.

> Poisot, T., Canard, E., Mouquet, N., Hochberg, M.E., 2012. A comparative study
> of ecological specialization estimators. Methods in Ecology and Evolution 3,
> 537–544. doi:10.1111/j.2041-210X.2011.00174.x

"""
function pdi(x::Vector{T}) where {T<:Number}
    s = reverse(sort(x))
    s = s ./ maximum(s)
    p = s[1] .- s
    return sum(p[2:end])/(length(s)-1)
end

"""
**Resource range**

    specificity(N::DeterministicNetwork)

Measure of specificity in a deterministic network. This returns a value between
0 and 1, where 1 indicates maximal specificity.

> Poisot, T., Canard, E., Mouquet, N., Hochberg, M.E., 2012. A comparative study
> of ecological specialization estimators. Methods in Ecology and Evolution 3,
> 537–544. doi:10.1111/j.2041-210X.2011.00174.x

"""
function specificity(N::BinaryNetwork)
    A = map(Int64, N.A)
    return Dict(zip(species(N; dims=1), vec(mapslices(pdi, A; dims=2))))
end

"""
**Paired Differences Index**

    specificity(N::QuantitativeNetwork)

Measure of specificity in a quantitative network. This returns a value between 0
and 1, where 1 indicates maximal specificity. Note that the PDI is measured
species-wise, and the maximal interaction strength of every species is set to 1.

> Poisot, T., Canard, E., Mouquet, N., Hochberg, M.E., 2012. A comparative study
> of ecological specialization estimators. Methods in Ecology and Evolution 3,
> 537–544. doi:10.1111/j.2041-210X.2011.00174.x

"""
function specificity(N::QuantitativeNetwork)
    return Dict(zip(species(N; dims=1), vec(mapslices(pdi, N.A; dims=2))))
end


using LinearAlgebra

"""
    Plane(c::Vector{Float64}, a::Vector{Float64})

Construct a `Plane` from a point `c` which lies in the plane and a unit normal `a`.
"""
struct Plane
    c::Vector{Float64}
    a::Vector{Float64}
    Plane(c,a) = begin
        if length(c) != 3 || length(a) != 3
            error("3-vectors please")
        else
            new(c,a/norm(a))
        end
    end
end

"""
    Cylinder(R::Float64, c::Vector{Float64}, a::Vector{Float64})

Construct a `Cylinder` with radius `R` and unit axis `a` which passes through point `c`.
"""
struct Cylinder
    R::Real          # radius
    c::Vector{Float64}  # center point
    a::Vector{Float64}  # axis (unit) (NOTE: c + h*a yields a point on the opposite cap => this is an inward normal axis)
    Cylinder(R,c,a) = begin
        if length(c) != 3 || length(a) != 3
            error("3-vectors please")
        else
            new(R,c,a/norm(a))
        end
    end
end

"""
    Paraboloid(b::Float64, c::Vector{Float64}, a::Vector{Float64})

Construct a `Paraboloid` at vertex `c` with growth rate `b` and unit axis `a`.
"""
struct Paraboloid
    b::Real
    c::Vector{Float64}
    a::Vector{Float64}
    Paraboloid(b,c,a) = begin
        if length(c) != 3 || length(a) != 3
            error("3-vectors please")
        else
            new(b,c,a/norm(a))
        end
    end
end

"""
    Hyperboloid(R::Float64, b::Float64, c::Vector{Float64}, a::Vector{Float64})

Construct a `Hyperboloid` of one sheet centered at `c` with minimum radius `R`, radial growth rate `b`, and unit axis `a`.
"""
struct Hyperboloid
    R::Real
    b::Real
    c::Vector{Float64}
    a::Vector{Float64}
    Hyperboloid(R,b,c,a) = begin
        if length(c) != 3 || length(a) != 3
            error("3-vectors please")
        else
            new(R,b,c,a/norm(a))
        end
    end
end

"""
    Quadric(Q::Matrix{Float64})

Construct a `Quadric` surface from a 4x4 symmetric quadric matrix.
"""
struct Quadric
    Q::Matrix{Float64}
    Quadric(Q) = begin
        if issymmetric(Q)
            new(Q)
        else
            error("quadric matrix must be symmetric")
        end
    end
end

"""
    TruncatedQuadric(q::Quadric, p::Vector{Plane}, caps)

Construct a `TruncatedQuadric` surface by slicing a non-reducible `Quadric` q with two planes `p`. Planes may not contain quadric axis. Argument `caps=true` closes the volume between the planes, `caps=false` defines an open-ended tube bounded by the `Quadric` and the planes.
"""
struct TruncatedQuadric
    q::Quadric
    p::Vector{Plane}
    caps::Bool
    TruncatedQuadric(q,p,caps) = begin
        Qr = q.Q[1:3,1:3]
        if all(eigvals(Qr) .!= 0)
            for plane in p
                if all(q.Q[1:3,1:3]*plane.a .== 0)
                    error("planes must be oblique to quadric axis")
                end
            end
        end
        new(q,p,caps)
    end
end

@doc raw"""
    Quadric(s::Plane)

Construct a `Quadric` from a `Plane` (with attributes axis = `a` = ``\boldsymbol{a}``; and point in plane = `c` = ``\boldsymbol{c}``) using the definition:

``\boldsymbol{Q}_s = \begin{bmatrix} \boldsymbol{0} & \frac{1}{2}\boldsymbol{a} \\ \frac{1}{2}\boldsymbol{a}^\mathsf{T} & -\boldsymbol{a}^\mathsf{T}\boldsymbol{c} \end{bmatrix}``
"""
function Quadric(s::Plane)
    Q = [zeros(3,3)      0.5*s.a;
         0.5*s.a'        -s.a'*s.c]
    return Quadric(Q)
end

@doc raw"""
    Quadric(s::Cylinder)

Construct a `Quadric` from a `Cylinder` (which has attributes axis = `a` = ``\boldsymbol{a}``; point on axis = `c` = ``\boldsymbol{c}``; and radius = `R` = ``R``) using the definitions:

``\boldsymbol{Q}_r = \boldsymbol{a}\boldsymbol{a}^\mathsf{T} - \boldsymbol{1}``

``\boldsymbol{Q}_c = \begin{bmatrix} \boldsymbol{Q}_r & -\boldsymbol{Q}_r\boldsymbol{c} \\ -\boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r & R^2 + \boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r\boldsymbol{c} \end{bmatrix}``
"""
function Quadric(s::Cylinder)
    Q = [s.a*s.a' - I            (I - s.a*s.a')*s.c;
         ((I - s.a*s.a')*s.c)'   s.R^2 + s.c'*(s.a*s.a' - I)*s.c]
    return Quadric(Q)
end

@doc raw"""
    Quadric(s::Paraboloid)

Construct a `Quadric` from a `Paraboloid` (which has attributes axis = `a` = ``\boldsymbol{a}``; point on axis = `c` = ``\boldsymbol{c}``; and quadratic parameter = `b` = ``b``) using the definitions:

``\boldsymbol{Q}_r = \boldsymbol{a}\boldsymbol{a}^\mathsf{T} - \boldsymbol{1}``

``\boldsymbol{Q}_p = \begin{bmatrix} \boldsymbol{Q}_r & \frac{b^2}{2}\boldsymbol{a} - \boldsymbol{Q}_r\boldsymbol{c} \\ \frac{b^2}{2}\boldsymbol{a}^\mathsf{T} - \boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r & -b^2\boldsymbol{a}^\mathsf{T}\boldsymbol{c} + \boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r\boldsymbol{c} \end{bmatrix}``
"""
function Quadric(s::Paraboloid)
    Q = [s.a*s.a' - I   s.b^2/2*s.a + (I - s.a*s.a')*s.c;
         (s.b^2/2*s.a + (I - s.a*s.a')*s.c)'    -s.b^2*s.a'*s.c + s.c'*(s.a*s.a' - I)*s.c]
    return Quadric(Q)
end

@doc raw"""
    Quadric(s::Hyperboloid)

Construct a `Quadric` from a `Hyperboloid` (which has attributes axis = `a` = ``\boldsymbol{a}``; point on axis = `c` = ``\boldsymbol{c}``; neck radius = `R` = ``R``; and slope parameter = `b` = ``b``) using the definitions:

``\gamma = 1 + \frac{R^2}{b^2}``

``\boldsymbol{Q}_r = \gamma\boldsymbol{a}\boldsymbol{a}^\mathsf{T} - \boldsymbol{1}``

``\boldsymbol{Q}_h = \begin{bmatrix} \boldsymbol{Q}_r & -\boldsymbol{Q}_r\boldsymbol{c} \\ -\boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r & R^2 + \boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r\boldsymbol{c} \end{bmatrix}``
"""
function Quadric(s::Hyperboloid)
    γ = (1 + s.R^2/s.b^2)
    Q = [γ*s.a*s.a' - I     (I - γ*s.a*s.a')*s.c;
         ((I - γ*s.a*s.a')*s.c)'  s.R^2 + s.c'*(γ*s.a*s.a' - I)*s.c]
    return Quadric(Q)
end

"""
    TruncatedQuadric(s::Cylinder, p::Vector{Plane}, caps=true)

Construct a `TruncatedQuadric` surface from a `Cylinder` and set of `Plane`s by converting `Cylinder` directly to a `Quadric`.
"""
TruncatedQuadric(s::Cylinder, p::Vector{Plane}, caps=true) = TruncatedQuadric(Quadric(s), p, caps)

"""
    TruncatedQuadric(s::Paraboloid, p::Vector{Plane}, caps=true)

Construct a `TruncatedQuadric` surface from a `Paraboloid` and set of `Plane`s by converting `Paraboloid` directly to a `Quadric`.
"""
TruncatedQuadric(s::Paraboloid, p::Vector{Plane}, caps=true) = TruncatedQuadric(Quadric(s), p, caps)

"""
    TruncatedQuadric(s::Hyperboloid, p::Vector{Plane}, caps=true)

Construct a `TruncatedQuadric` surface from a `Hyperboloid` and set of `Plane`s by converting `Hyperboloid` directly to a `Quadric`.
"""
TruncatedQuadric(s::Hyperboloid, p::Vector{Plane}, caps=true) = TruncatedQuadric(Quadric(s), p, caps)

"""
    TruncatedQuadric(q::Quadric, h1::Vector{Float64}, h2::Vector{Float64}, caps=true)

A convenience constructor for a `TruncatedQuadric`. Specify `h1` and `h2` as points along the axis of `q`, and the constructor will return `q` truncated by planes passing through `h1` and `h2` sharing a normal vector equal to the axis of `q`.
"""
function TruncatedQuadric(q::Quadric, h1::Vector{Float64}, h2::Vector{Float64}, caps=true)
    s = changerepresentation(q)
    if typeof(s) == Plane
        display("TruncatedQuadric just a bunch of planes")
    end

    if abs(((h1 - h2)/norm(h1 - h2))'*s.a/norm(s.a)) == 1
        p1 = Plane(h1, s.a)
        p2 = Plane(h2, s.a)
    else
        error("h1, h2 must lie on Quadric axis")
    end

    return TruncatedQuadric(q, [p1, p2], caps)
end

@doc raw"""
    changerepresentation(q::Quadric)

Converts a `Quadric` to a `Plane`, `Cylinder`, `Paraboloid`, or `Hyperboloid` based on its matrix structure.

!!! note
    This function is not fast. Minimize conversion between explicit surface and quadric matrix in heavy-lifting code.
"""
function changerepresentation(q::Quadric)
    ε = 1e-15

    Qr = q.Q[1:3, 1:3]
    qd = q.Q[1:3, end]
    q0 = q.Q[end, end]
    E = eigvals(Qr)
    if all(E .== 0)
        # plane
        a = 2*qd                    # Qr is zero, qd is parallel to plane normal
        c = -q0*a                   # q0 = -a'*c: choose c to lie on a (free choice)
        return Plane(c, a)

    elseif all(abs.(E) .>= ε)
        # hyperboloid
        v = eigvecs(Qr)
        a = v[:,end]                # axis is eigenvector for largest eigenvalue
        c = (-Qr)\qd                # center of hyperboloid
        γ = tr(Qr + I)              # Qr = γ*a*a' - I, and trace(a*a') == 1
        R = sqrt(q0 - c'*Qr*c)      # q0 = sum(c*c' .* Qr) + R^2
        b = 1/sqrt(γ - 1)*R         # γ = 1 + (R/b)^2
        return Hyperboloid(R, b, c, a)

    else
        v = eigvecs(Qr)
        a = v[:,end]        # axis is eigenvector for largest eigenvalue

        if rank(q.Q) == 4
            # paraboloid
            a = sign(qd'*a)*a       # correct for antiparallel axis
            b = sqrt(2*qd'*a)

            cb1 = v[:,1]'*qd
            cb2 = v[:,2]'*qd
            cb3 = -(q0 + cb1^2 + cb2^2)/b^2

            c = v*[cb1; cb2; cb3]

            return Paraboloid(b, c, a)

        elseif rank(q.Q) == 3
            # cylinder
            c = (-Qr)\qd
            badI = findfirst(isnan.(c) .| isinf.(c))
            if !isnothing(badI)
                c[badI] = 1.0
            end
            R = sqrt(q0 - c'*Qr*c)
            return Cylinder(R, c, a)

        else
            error("unclassified shape")

        end
    end
end

"""
    changerepresentation(s::Plane)

Convert a `Plane` to `Quadric` representation via the `Quadric(s::Plane)` constructor.
"""
changerepresentation(s::Plane) = Quadric(s)

"""
    changerepresentation(s::Cylinder)

Convert a `Cylinder` to `Quadric` representation via the `Quadric(s::Cylinder)` constructor.
"""
changerepresentation(s::Cylinder) = Quadric(s)

"""
    changerepresentation(s::Paraboloid)

Convert a `Paraboloid` to `Quadric` representation via the `Quadric(s::Paraboloid)` constructor.
"""
changerepresentation(s::Paraboloid) = Quadric(s)

"""
    changerepresentation(s::Hyperboloid)

Convert a `Hyperboloid` to `Quadric` representation via the `Quadric(s::Hyperboloid)` constructor.
"""
changerepresentation(s::Hyperboloid) = Quadric(s)

"""
    normal(q::Quadric, r::Vector{Float64})

Returns unit normal vector to a `Quadric` at position `r` on the surface.

!!! warning
    This will return a normal vector even if `r` does not fall on the surface. It is up to *you* to input an `r` you know to lie on the surface of `q`.
"""
function normal(q::Quadric, r::Vector{Float64})
    n = 2*q.Q*[r; 1]
    n = n[1:3]
    n = n/norm(n)
    return n
end

"""
    normal(s::Plane, r::Vector{Float64})

Returns unit normal vector to a `Plane`. This will return the `Plane`'s normal vector regardless of the `r` supplied; it takes the `r` argument nonetheless to maintain the same pattern as `normal()` functions for nonreducible `Quadric`s.
"""
function normal(s::Plane, r::Vector{Float64})
    return s.a
end

"""
    normal(s::Cylinder, r::Vector{Float64})

Returns unit normal vector to a `Cylinder` at position `r` on the surface.

!!! warning
    This will return a normal vector even if `r` does not fall on the surface. It is up to *you* to input an `r` you know to lie on the surface of `s`.
"""
function normal(s::Cylinder, r::Vector{Float64})
    n = 2*(s.a - (r - s.c)./norm(r - s.c))
    return n/norm(n)
end

"""
    normal(s::Paraboloid, r::Vector{Float64})

Returns unit normal vector to a `Paraboloid` at position `r` on the surface.

!!! warning
    This will return a normal vector even if `r` does not fall on the surface. It is up to *you* to input an `r` you know to lie on the surface of `s`.
"""
function normal(s::Paraboloid, r::Vector{Float64})
    n = s.b^2*s.a + 2*(s.a - (r - s.c)./norm(r - s.c))
    return n/norm(n)
end

"""
    normal(s::Hyperboloid, r::Vector{Float64})

Returns normal vector to a `Hyperboloid` at position `r` on the surface.

!!! warning
    This will return a normal vector even if `r` does not fall on the surface. It is up to *you* to input an `r` you know to lie on the surface of `q`.
"""
function normal(s::Hyperboloid, r::Vector{Float64})
    n = 2*((1 + s.R^2/s.b^2)*s.a - (r - s.c)./norm(r - s.c))
    return n/norm(n)
end

function classify(q::Quadric)

    Qr = q.Q[1:3,1:3]

    E = eigvals(Qr)

    p = sum(E .> 0)
    n = sum(E .< 0)
    z = sum(E .== 0)

    # options:
    #   [p, n, z]
    #   [0, 0, 3]:  plane
    #   [0, 1, 2]:  parabolic cylinder, parallel planes
    #   [0, 2, 1]:  imaginary cylinder
    #   [0, 3, 0]:  imag. ellipsoid
    #   [1, 0, 2]:  parabolic cylinder, parallel planes
    #   [1, 1, 1]:  hyperbolic paraboloid, hyperbolic cylinder
    #   [1, 2, 0]:  hyperboloid 2 sheet
    #   [2, 0, 1]:  elliptic paraboloid, elliptic cylinder
    #   [2, 1, 0]:  hyperboloid 1 sheet, cone @ origin
    #   [3, 0, 0]:  ellipsoid, imag. context
    shapes0 = [
        "none"  "none"  "none"  "ellipsoid";
        "none"  "none"  "hyperboloid"   "none";
        "none"  "hyperboloid"  "none"   "none";
        "ellipsoid" "none"  "none"  "none"
    ]
    shapes1 = [
        "none"  "none"  "elliptic cylinder/paraboloid"    "none";
        "none"  "hyperbolic paraboloid/cylinder"  "none"  "none";
        "elliptic cylinder/paraboloid" "none"   "none"  "none";
        "none"  "none"  "none"  "none"
    ]
    shapes2 = [
        "none"  "parabolic cylinder/planes" "none"  "none";
        "parabolic cylinder/planes" "none"  "none"  "none";
        "none"  "none"  "none"  "none";
        "none"  "none"  "none"  "none"
    ]
    shapes3 = [
        "plane"  "none"  "none"  "none";
        "none"   "none"  "none"  "none";
        "none"   "none"  "none"  "none";
        "none"   "none"  "none"  "none";
    ]

    shapes = cat(shapes0, shapes1, shapes2, shapes3, dims=3)

    shapename = shapes[p+1, n+1, z+1]
    return shapename
end