using Catalyst, LinearAlgebra, OrdinaryDiffEq, Test, NonlinearSolve

# Repressilator model
@parameters t α₀ α K n δ β μ
@variables m(t) P(t) R(t)
rxs = [
       Reaction(α₀, nothing, [m]),
       Reaction(α / (1 + (R/K)^n), nothing, [m]),
       Reaction(δ, [m], nothing),
       Reaction(β, [m], [m,P]),
       Reaction(μ, [P], nothing)
      ]

specs = [m,P,R]
pars  = [α₀,α,K,n,δ,β,μ]
@named rs = ReactionSystem(rxs, t, specs, pars)

# using ODESystem components
@named sys₁ = convert(ODESystem, rs; include_zero_odes=false)
@named sys₂ = convert(ODESystem, rs; include_zero_odes=false)
@named sys₃ = convert(ODESystem, rs; include_zero_odes=false)
connections = [sys₁.R ~ sys₃.P,
               sys₂.R ~ sys₁.P,
               sys₃.R ~ sys₂.P]
@named connected = ODESystem(connections, t, [], [], systems=[sys₁,sys₂,sys₃])
oderepressilator = structural_simplify(connected)

pvals = [sys₁.α₀ => 5e-4,
         sys₁.α => .5,
         sys₁.K => 40.0,
         sys₁.n => 2,
         sys₁.δ => (log(2)/120),
         sys₁.β => (20*log(2)/120),
         sys₁.μ => (log(2)/600),
         sys₂.α₀ => 5e-4,
         sys₂.α => .5,
         sys₂.K => 40.0,
         sys₂.n => 2,
         sys₂.δ => (log(2)/120),
         sys₂.β => (20*log(2)/120),
         sys₂.μ => (log(2)/600),
         sys₃.α₀ => 5e-4,
         sys₃.α => .5,
         sys₃.K => 40.0,
         sys₃.n => 2,
         sys₃.δ => (log(2)/120),
         sys₃.β => (20*log(2)/120),
         sys₃.μ => (log(2)/600)]
u₀    = [sys₁.m => 0.0, sys₁.P => 20.0, sys₁.R => 0.0, sys₂.m => 0.0, sys₂.P => 0.0, sys₂.R => 0.0, sys₃.m => 0.0, sys₃.P => 0.0, sys₃.R => 0.0]
tspan = (0.0, 100000.0)
oprob = ODEProblem(oderepressilator, u₀, tspan, pvals)
sol = solve(oprob, Tsit5())

# hardcoded network
function repress!(f, y, p, t)
    α = p.α; α₀ = p.α₀; β = p.β; δ = p.δ; μ = p.μ; K = p.K; n = p.n
    f[1] = α / (1 + (y[6] / K)^n) - δ * y[1] + α₀
    f[2] = α / (1 + (y[4] / K)^n) - δ * y[2] + α₀
    f[3] = α / (1 + (y[5] / K)^n) - δ * y[3] + α₀
    f[4] = β * y[1] - μ * y[4]
    f[5] = β * y[2] - μ * y[5]
    f[6] = β * y[3] - μ * y[6]
    nothing
end
ps = (α₀=5e-4, α=.5, K=40.0, n=2, δ=(log(2)/120), β=(20*log(2)/120), μ=(log(2)/600))
u0 = [0.0,0.0,0.0,20.0,0.0,0.0]
oprob2 = ODEProblem(repress!, u0, tspan, ps)
sol2 = solve(oprob2, Tsit5())
tvs = 0:1:tspan[end]
@test all(isapprox.(sol(tvs,idxs=sys₁.P),sol2(tvs, idxs=4),atol=1e-4))

# using ReactionSystem components
@named sys₁ = ReactionSystem(rxs, t, specs, pars)
@named sys₂ = ReactionSystem(rxs, t, specs, pars)
@named sys₃ = ReactionSystem(rxs, t, specs, pars)
connections = [ParentScope(sys₁.R) ~ ParentScope(sys₃.P),
               ParentScope(sys₂.R) ~ ParentScope(sys₁.P),
               ParentScope(sys₃.R) ~ ParentScope(sys₂.P)]
@named csys = ODESystem(connections, t, [], [])
@named repressilator = ReactionSystem(t; systems=[csys,sys₁,sys₂,sys₃])
@named oderepressilator2 = convert(ODESystem, repressilator, include_zero_odes=false)
sys2 = structural_simplify(oderepressilator2)  # FAILS currently
oprob = ODEProblem(sys2, u₀, tspan, pvals)
sol = solve(oprob, Tsit5())
@test all(isapprox.(sol(tvs,idxs=sys₁.P),sol2(tvs, idxs=4),atol=1e-4))

# test conversion to nonlinear system
@named nsys = NonlinearSystem(connections, [], [])
@named ssrepressilator = ReactionSystem(t; systems=[nsys,sys₁,sys₂,sys₃])
@named nlrepressilator = convert(NonlinearSystem, ssrepressilator, include_zero_odes=false)
sys2 = structural_simplify(nlrepressilator)
@test length(equations(sys2)) == 6
nlprob = NonlinearProblem(sys2, u₀, pvals)
sol = solve(nlprob, NewtonRaphson(), tol=1e-9)
@test sol[sys₁.P] ≈ sol[sys₂.P] ≈ sol[sys₃.P]
@test sol[sys₁.m] ≈ sol[sys₂.m] ≈ sol[sys₃.m]
@test sol[sys₁.R] ≈ sol[sys₂.R] ≈ sol[sys₃.R]

# flattening
fsys = Catalyst.flatten(ssrepressilator)
@named nlrepressilator = convert(NonlinearSystem, fsys, include_zero_odes=false)
sys2 = structural_simplify(nlrepressilator)
@test length(equations(sys2)) == 6
nlprob = NonlinearProblem(sys2, u₀, pvals)
sol = solve(nlprob, NewtonRaphson(), tol=1e-9)
@test sol[sys₁.P] ≈ sol[sys₂.P] ≈ sol[sys₃.P]
@test sol[sys₁.m] ≈ sol[sys₂.m] ≈ sol[sys₃.m]
@test sol[sys₁.R] ≈ sol[sys₂.R] ≈ sol[sys₃.R]

# test constraints
connections = [sys₁.R ~ sys₃.P,
               sys₂.R ~ sys₁.P,
               sys₃.R ~ sys₂.P]
@named csys = NonlinearSystem(connections, [sys₁.R,sys₃.P,sys₂.R,sys₁.P,sys₃.R,sys₂.P],[])
@named repressilator2 = ReactionSystem(t; constraints=csys, systems=[sys₁,sys₂,sys₃])
@named nlrepressilator = convert(NonlinearSystem, repressilator2, include_zero_odes=false)
sys2 = structural_simplify(nlrepressilator)
@test length(equations(sys2)) == 6
nlprob = NonlinearProblem(sys2, u₀, pvals)
sol = solve(nlprob, NewtonRaphson(), tol=1e-9)
@test sol[sys₁.P] ≈ sol[sys₂.P] ≈ sol[sys₃.P]
@test sol[sys₁.m] ≈ sol[sys₂.m] ≈ sol[sys₃.m]
@test sol[sys₁.R] ≈ sol[sys₂.R] ≈ sol[sys₃.R]

# test can make ODESystem
@named oderepressilator = convert(ODESystem, repressilator2, include_zero_odes=false)
sys2 = structural_simplify(oderepressilator)  # FAILS currently
oprob = ODEProblem(sys2, u₀, tspan, pvals)
sol = solve(oprob, Tsit5())
@test all(isapprox.(sol(tvs,idxs=sys₁.P),sol2(tvs, idxs=4),atol=1e-4))

# test extending with NonlinearSystem
@named repressilator2 = ReactionSystem(t; systems=[sys₁,sys₂,sys₃])
repressilator2 = Catalyst.flatten(repressilator2)
repressilator2 = extend(csys, repressilator2)
@named nlrepressilator = convert(NonlinearSystem, repressilator2, include_zero_odes=false)
sys2 = structural_simplify(nlrepressilator)
@test length(equations(sys2)) == 6
nlprob = NonlinearProblem(sys2, u₀, pvals)
sol = solve(nlprob, NewtonRaphson(), tol=1e-9)
@test sol[sys₁.P] ≈ sol[sys₂.P] ≈ sol[sys₃.P]
@test sol[sys₁.m] ≈ sol[sys₂.m] ≈ sol[sys₃.m]
@test sol[sys₁.R] ≈ sol[sys₂.R] ≈ sol[sys₃.R]

# TODO add conversion to SDE and JumpSystems once supported

# adding algebraic constraints
@parameters t, r₊, r₋, β
@variables A(t), B(t), C(t), D(t)
rxs1 = [Reaction(r₊, [A,B], [C])]
rxs2 = [Reaction(r₋, [C], [A,B])]
@named rs1 = ReactionSystem(rxs1, t, [A,B,C], [r₊])
@named rs2 = ReactionSystem(rxs2, t, [A,B,C], [r₋])
@named rs  = extend(rs1, rs2)
@test issetequal(states(rs), [A,B,C])
@test issetequal(parameters(rs), [r₊,r₋])
@test issetequal(equations(rs), union(rxs1,rxs2))
A2 = ModelingToolkit.ParentScope(A)
B2 = ModelingToolkit.ParentScope(B)
nseqs = [D ~ 2*A2 + β*B2]
@named ns = ODESystem(nseqs, t, [A2,B2,D], [β])
rs = compose(rs, [ns])
osys = convert(ODESystem, rs; include_zero_odes=false)
p  = [r₊ => 1.0, r₋ => 2.0, ns.β => 3.0]
u₀ = [A => 1.0, B => 2.0, C => 0.0]
oprob = ODEProblem(structural_simplify(osys), u₀, (0.0,10.0), p)
sol = solve(oprob, Tsit5())
@test isapprox(0, norm(sol[ns.D] .- 2*sol[A] - 3*sol[B]), atol=(100*eps()))

# test API functions for composed model
@test issetequal(species(rs), [A,B,C])
@test issetequal(states(rs), [A,B,C,ns.D])
@test issetequal(reactionparams(rs), [r₊,r₋])
@test issetequal(parameters(rs), [r₊,r₋,ns.β])
@test issetequal(reactions(rs), union(rxs1,rxs2))
@test issetequal(filter(eq -> eq isa Reaction, equations(rs)), union(rxs1,rxs2))
@test issetequal(filter(eq -> eq isa Equation, equations(rs)), [ModelingToolkit.namespace_equation(nseqs[1],ns)])

# check several levels of nesting namespace and filter ok for the API functions
@parameters t, p1, p2a, p2b, p3a, p3b
@variables A1(t), A2a(t), A2b(t), A3a(t), A3b(t)
rxs1 = [Reaction(p1, [A1], nothing)]
rxs2 = [Reaction(p2a, [A2a], nothing), Reaction(p2b, [ParentScope(A1)], nothing)]
eqs2 = [ParentScope(A1) ~ ParentScope(p1)*A2b]
rxs3 = [Reaction(p3a, [A3a], nothing), Reaction(ParentScope(p2a), nothing, [ParentScope(A2a)])]
eqs3 = [ParentScope(A2a) ~ p3b*A3b]

@named rs3 = ReactionSystem(rxs3,t,[A3a,ParentScope(A2a)],[p3a,ParentScope(p2a)])
@named ns3 = NonlinearSystem(eqs3,[ParentScope(A2a),A3b],[p3b])
@named rs2 = ReactionSystem(rxs2,t,[A2a,ParentScope(A1)],[p2a,p2b], systems=[rs3,ns3])
@named ns2 = NonlinearSystem(eqs2,[ParentScope(A1),A2b],[ParentScope(p1)])
@named rs1 = ReactionSystem(rxs1,t,[A1],[p1],systems=[rs2,ns2])

# namespaced reactions
nrxs1 = [Reaction(p1, [A1], nothing)]
nrxs2 = [Reaction(rs2.p2a, [rs2.A2a], nothing), Reaction(rs2.p2b, [A1], nothing)]
neqs2 = [0 ~ p1*ns2.A2b - A1]
nrxs3 = [Reaction(rs2.rs3.p3a, [rs2.rs3.A3a], nothing), Reaction(rs2.p2a, nothing, [rs2.A2a])]
neqs3 = [0 ~ rs2.ns3.p3b*rs2.ns3.A3b - rs2.A2a]
rxs = vcat(nrxs1, nrxs2, nrxs3)
eqs = vcat(nrxs1, nrxs2, neqs2, nrxs3, neqs3)

@test issetequal(states(rs1), [A1,rs2.A2a,ns2.A2b,rs2.rs3.A3a,rs2.ns3.A3b])
@test issetequal(species(rs1), [A1,rs2.A2a,rs2.rs3.A3a])
@test issetequal(parameters(rs1), [p1,rs2.p2a,rs2.p2b,rs2.rs3.p3a,rs2.ns3.p3b])
@test issetequal(reactionparams(rs1), [p1, rs2.p2a, rs2.p2b, rs2.rs3.p3a])
@test issetequal(rxs, reactions(rs1))
@test issetequal(eqs, equations(rs1))

# test throw error if there are ODE constraints and convert to NonlinearSystem
rn = @reaction_network rn begin
    (k1,k2), A <--> B
    end k1 k2
@parameters a,b
@variables t,A(t),C(t)
D = Differential(t)
eqs = [D(A) ~ -a*A + C, D(C) ~ -b*C + a*A]
@named osys = ODESystem(eqs,t,[A,C],[a,b])
rn2 = extend(osys,rn)
rnodes = convert(ODESystem,rn2)
@test_throws ErrorException convert(NonlinearSystem,rn2)

# ensure right number of equations are generated
@variables G(t)
eqs = [D(G) ~ -G]
@named osys2 = ODESystem(eqs,t)
rn3 = compose(rn2, osys2)
@test length(equations(rn3)) == 5

# check conversions work with algebraic constraints
eqs = [0 ~ -a*A + C, 0 ~ -b*C + a*A]
@named nlsys = NonlinearSystem(eqs,[A,C],[a,b])
rn2 = extend(nlsys,rn)
rnodes = convert(ODESystem,rn2)
rnnlsys = convert(NonlinearSystem,rn2)
@named nlsys = ODESystem(eqs,t,[A,C],[a,b])
rn2 = extend(nlsys,rn)
rnodes = convert(ODESystem,rn2)
rnnlsys = convert(NonlinearSystem,rn2)

# https://github.com/SciML/ModelingToolkit.jl/issues/1274
@parameters p1 p2
@variables t A(t)
rxs1 = [Reaction(p1, [A], nothing)]
rxs2 = [Reaction(p2, [ParentScope(A)], nothing)]
@named rs1 = ReactionSystem(rxs1, t)
@named rs2 = ReactionSystem(rxs2, t)
rsc = compose(rs1,[rs2])
orsc = convert(ODESystem, rsc)
@test length(equations(orsc)) == 1


# # Compare ADCs
#
# The [apparent diffusion coefficient](@ref adc) (ADC) may be sufficient to describe the
# signal attenuation for small ``b``-values. SpinDoctor comes with multiple approaches for
# computing or estimating the ADC for a `ScalarGradient` ``\vec{g}(t) = f(t) g \vec{d}``:
#
# - Using the free diffusion coefficient ``\vec{d}^\mathsf{T} \mathbf{D} \vec{d}``, which
#   represents unrestricted diffusion in the absence of boundaries;
# - Computing the short diffusion time approximation for the ADC
# - Fitting the log-signal obtained by solving the BTPDE for different ``b``-values
# - Solving a homogenized model (HADC) assuming negligible permeability
# - Using the matrix formalism effective diffusion tensor
#
# In this example we will compare the different approaches.

# ## Building a biological model
#
# We start by loading SpinDoctor and a Makie plotting backend.

# LSP indexing solution                                                          #src
# https://github.com/julia-vscode/julia-vscode/issues/800#issuecomment-650085983 #src
if isdefined(@__MODULE__, :LanguageServer)                                       #src
    include("../src/SpinDoctor.jl")                                              #src
    using .SpinDoctor                                                            #src
end                                                                              #src

using SpinDoctor
using LinearAlgebra

if haskey(ENV, "GITHUB_ACTIONS")
    using CairoMakie
else
    using GLMakie
end

# Here we create a recipe for five stacked plates with isotropic diffusion tensors. They
# should allow for free diffusion in the horizontal direction, but a rather restricted
# vertical diffusion with the permeable membranes.

ncell = 5
setup = PlateSetup(;
    name = "Plates",
    width = 50.0,
    depth = 50.0,
    heights = fill(5.0, ncell),
    bend = 0.0,
    twist = 0.0,
    refinement = 10.0,
)
coeffs = coefficients(
    setup;
    D = [0.002 * I(3) for _ = 1:ncell],
    T₂ = fill(Inf, ncell),
    ρ = fill(1.0, ncell),
    κ = (; interfaces = fill(1e-4, ncell - 1), boundaries = fill(0.0, ncell)),
    γ = 2.67513e-4,
)

# We then proceed to build the geometry and finite element mesh.

mesh, = create_geometry(setup; recreate = true)
plot_mesh(mesh)

# The mesh looks good, so we may then proceed to assemble the biological model and the
# associated finite element matrices.

model = Model(; mesh, coeffs...)
matrices = assemble_matrices(model);

# We may also compute some useful quantities, including a scalar diffusion coefficient from
# the diffusion tensors.

volumes = get_cmpt_volumes(model.mesh)
D_avg = 1 / 3 * tr.(model.D)' * volumes / sum(volumes)
ncompartment = length(model.mesh.points)

# The gradient pulse sequence will be a PGSE with both vertical and horizontal components.
# This allows for both restricted vertical diffusion and almost unrestricted horizontal
# diffusion. The different approaches should hopefully confirm this behaviour.

dir = [1.0, 0.0, 1.0]
profile = PGSE(2500.0, 4000.0)
b = 1000
g = √(b / int_F²(profile)) / model.γ
gradient = ScalarGradient(dir, profile, g)

# ## Computing the ADC using different methods

# ### Short term approximation

# A simple model for the ADC is given by the [short term approximation](@ref STA).

adc_sta_cmpts = compute_adc_sta(model, gradient)
adc_sta = volumes'adc_sta_cmpts / sum(volumes)

# ### Fitting the ADC

# A more robust approach is to directly fit the BTPDE signal to a series of b-values. This
# is however more computationally expensive. We start by building a set of gradients.

bvalues = 0:400:4000
gvalues = map(b -> √(b / int_F²(profile)) / coeffs.γ, bvalues)
gradients = [ScalarGradient(gradient.dir, gradient.profile, g) for g ∈ gvalues]

# The [`solve_multigrad`](@ref) function computes the magnetization for each gradient. Since
# the gradients are interval-wise constant, we can use a specialized solver.

btpde = BTPDE(; model, matrices)
solver = IntervalConstantSolver(; θ = 0.5, timestep = 5.0)
ξ, = solve_multigrad(btpde, gradients, solver)

# The signals can be computed from the magnetization fields.

signals = [compute_signal(matrices.M, ξ) for ξ ∈ ξ]
signals_cmpts = [compute_signal.(matrices.M_cmpts, split_field(model.mesh, ξ)) for ξ ∈ ξ]

# Fitting the ADC is straightforward.

adc_fit = fit_adc(bvalues, signals)
adc_fit_cmpts =
    [fit_adc(bvalues, [s[icmpt] for s ∈ signals_cmpts]) for icmpt = 1:ncompartment]

# ### Homogenized ADC model

# The [HADC](@ref)-model uses homogenization and assumes negligible permeability between the
# compartments. This does require solving an ODE involving all the degrees of freedom.

hadc = HADC(; model, matrices)
adc_homogenized_cmpts = solve(hadc, gradient)
adc_homogenized = volumes'adc_homogenized_cmpts / sum(volumes)

# ### Matrix Formalism

# If a Laplace eigendecomposition is available, the HADC can be approximated with little
# additional computational expense.

laplace = Laplace(; model, matrices, neig_max = 400)
lap_eig = solve(laplace)

# An effective diffusion tensor can be computed.

D_mf = compute_mf_diffusion_tensor(model.mesh, matrices.M, lap_eig, gradient)

# We observe that ``D_{x x}`` and ``D_{y y}`` are very close to the intrinsic diffusion
# coefficient ``D = 0.002``, which confirms that diffusion in the horizontal direction is
# almost unrestricted. ``D_{z z}`` is significantly smaller, confirming the presence of
# membranes along the vertical direction.

# In particular, we may deduce the MF-ADC in our direction.

adc_mf = dir'D_mf * dir / dir'dir

# ## Comparing results

# ### Compartment ADCs

# Here we make a comparison between the compartment ADCs. For MF, only a global ADC was computed.

n = ncompartment
fig = Figure()
ax = Axis(fig[1, 1];
    xticks = (1:4, ["STA", "Fit BTPDE", "HADC", "MF"]),
    ylabel = "ADC / D",
    title = "Compartment ADCs",
)
barplot!(ax, fill(1, n), adc_sta_cmpts ./ D_avg; dodge = 1:n)
barplot!(ax, fill(2, n), adc_fit_cmpts ./ D_avg; dodge = 1:n)
barplot!(ax, fill(3, n), adc_homogenized_cmpts ./ D_avg; dodge = 1:n)
barplot!(ax, [4], [adc_mf / D_avg])
fig

# The STA and HADC are the same for all compartments, as they consider them separately and
# all the compartments have the same size. The fitted ADC is larger for the three inner
# compartments, as they all have permeable membranes both below and above, in contrast to
# the top and bottom compartments that have hard walls.

# ### Signal attenuation

# We may also inspect the resulting signal attenuations.

fig = Figure()
ax = Axis(fig[1,1]; xlabel = "b", yscale = log10, title = "Signal attenuation")
lines!(ax, [0, bvalues[end]], [1, exp(-adc_sta * bvalues[end])]; linestyle = :dash, label = "ADC STA")
lines!(ax, [0, bvalues[end]], [1, exp(-adc_fit * bvalues[end])]; linestyle = :dash, label = "ADC Fit")
lines!(ax, [0, bvalues[end]], [1, exp(-adc_homogenized * bvalues[end])]; linestyle = :dash, label = "HADC")
lines!(ax, [0, bvalues[end]], [1, exp(-adc_mf * bvalues[end])]; linestyle = :dash, label = "ADC MF")
lines!(ax, [0, bvalues[end]], [1, exp(-D_avg * bvalues[end])]; linestyle = :dash, label = "Free diffusion")
scatterlines!(ax, bvalues, abs.(signals) ./ abs(signals[1]); label = "BTPDE Signal")
axislegend(ax)
fig

# We observe the following:
#
# - The exact signal starts to visually deviate from the log-linear regime after ``b =
#   2000``. For higher ``b``-values, the ADC is no longer sufficient to describe the
#   attenuation, as higher order terms can no longer be neglected.
# - The fit-ADC signal coincides with the exact signal for the lowest ``b``-values. This
#   makes sense since that is how this ADC was obtained to begin with. This is considered to
#   be the "reference" ADC.
# - The free diffusion signal attenuates more thant the exact signal by many orders of
#   magnitude, which confirms the presence of restrictive membranes and boundaries in the
#   gradient direction.
# - The HADC signal attenuates less than the exact signal, as it assumes a more severe
#   restriction with impermeable membranes.
# - The MF-ADC signal coincides with the HADC signal, as the former is simply an MF
#   approximation of the latter.


module Utils

using StaticArrays
using ProgressMeter

function toMatrix(mat::Matrix{SVector{3,Float64}})
    n,m = size(mat)
    smat = fill(0.0, 3*n, m )
    for i=1:n
        for j =1:m
            smat[(3*(i-1)+1):(3*i), j] = mat[i,j]
        end
    end
    return smat
end

function toMatrix(mat::Matrix{SMatrix{3, 3, Float64, 9}})
    n,m = size(mat)
    smat = fill(0.0, 3*n, 3*m )
    for i=1:n
        for j =1:m
            smat[(3*(i-1)+1):(3*i), (3*(j-1)+1):(3*j)] = mat[i,j]
        end
    end
    return smat
end

function dot2(A::Matrix{Float64}, B::Matrix{SMatrix{3, 3, Float64, 9}})
    return A
end

function stack_Γ(Γ_list::Vector{Matrix{Float64}})
    Γ_tensor = zeros(size(Γ_list[1])...,length(Γ_list))
    for (i,Γ) in enumerate(Γ_list)
        Γ_tensor[:,:,i] = Γ
    end
    return Γ_tensor
end

function stack_B(B_list::Vector{Vector{Matrix{Float64}}} )

    B_tensor = zeros(size(B_list[1][1])...,length(B_list[1]),length(B_list))
    for (i,B) in enumerate(B_list)
        B_tensor[:,:,:,i] = stack_Γ(B)
    end
    return B_tensor
end

using LinearAlgebra
#n = 4
#a = [i+j for i in 1:n, j =1:n]
#B = [@SVector [i+j,2*i+j,3*i+j] for i in 1:n, j =1:n]
#b = [2*i+j for i in 1:n, j =1:n]
#a
#B
#C = a*B
#sum(a.*b)
#dot(a,b)
end


function matrix_mul(y, x, iflag::Ti; A::Matrix)

    # forward
    if iflag == 1
       mul!(y, A, x)

    # adjoint
    elseif iflag == 2
       mul!(y, transpose(A), x)

    else
       error("undefined operation")
    end

    return y
end

# set up the problem
for T in (Float32, Float64, ComplexF32, ComplexF64)
end

m = 100; n = 100;
A  = rand(T, m, n)
x0 = rand(T, n)
d  = A * x0

(x, his) = cgls(matrix_mul!, y; op_params=params, maxIter=m)

"""
   modified for memory saving
"""
function cgls(A::Tf, b::Vector{Tv}; op_params::NamedTuple=NamedTuple(),
              x0=[], shift=0.0, tol=1e-6, maxIter=50, print_flag=true) where {Tf<:Union{Function, Matrix}, Tv<:Number}

    # determine the type of A(matrix or operator)
    if isa(A, Matrix)
       explicitA = true

    elseif isa(A, Function)
       explicitA = false

    else
       error("A can be either a matrix or a function")
    end

    # initial guess of the unknowns
    if explicitA

       (m, n) = size(A)

       # provided initial guess
       if x0 != 0.0
          x = copy(x0)

          r = b - A * x
          s = A' * r - shift * x

       # zero vector as initial guess
       else
          x = zeros(eltype(b), n)
          r = copy(b)
          s = A' * r
       end

    else

       m = length(b)

       # provided initial guess
       if x0 != 0.0
          x = copy(x0)
          r = b - A(x, 1; op_params...)
          s = A(r, 2; op_params...) - shift * x
          n = length(s)
       # zero vector as initial guess
       else
          r = copy(b)
          s = A(b, 2; params...)
          n = length(s)
          x = zeros(eltype(b), n)
       end

    end

    data_fitting = dot(r,r)
    constraint   = 0.0
    cost0 = data_fitting + constraint

    # initialize some intermidiate vectors
    p = copy(s)
    norms0= vecnorm(s)
    gamma = norms0^2
    normx = vecnorm(x)
    xmax  = normx     # keep the initial one
    resNE = 1.0

    gamma0= copy(gamma)
    delta = 0.0

    # iteration counter and stop condition
    k = 0
    run_flag = true

    if print_flag
       header = "  k         data_fitting           constraint             normx                resNE"
       println(""); println(header);
       @printf("%3.0f %20.10e %20.10e %20.10e %20.10e\n", k, data_fitting, constraint, normx, resNE);
    end

    while (k < maxIter) && run_flag

          k = k + 1

          if explicitA
             q = A * p
          else
             q = A(p, 1; params...)
          end

          delta = (vecnorm(q))^2 + shift * (vecnorm(p))^2
          indefinite = delta <= 0 ? true : false
          delta      = delta == 0.? eps(): delta

          alpha = gamma / delta

          x = x + alpha * p
          r = r - alpha * q

          data_fitting = dot(r,r)
          constraint   = shift * dot(x, x)
          cost = data_fitting + constraint

          if explicitA
             s = A' * r - shift * x
          else
             s = A(r, 2; params...) - shift * x
          end

          norms  = norm(s)
          gamma0 = copy(gamma)
          gamma  = norms^2
          beta   = gamma / gamma0

          p = s + beta * p

          # check the stopping crietia
          normx = norm(x)
          xmax  = normx > xmax ? normx : xmax
          if norms <= norms0 * tol || normx * tol >= 1.0
             run_flag = false
          end

          # print information
          # resNE = norms / norms0
          resNE = cost  / cost0
          if print_flag
             @printf("%3.0f %20.10e %20.10e %20.10e %20.10e\n", k, data_fitting, constraint, normx, resNE);
          end

    end

    return x, resNE

end

"""
   CGLS algorithm for linear operators whose input and output can be save in memory
"""
function cgls(A::Tf, b::Vector{Tv}; op_params::NamedTuple=NamedTuple(),
              x0=[], shift=0.0, tol=1e-6, maxIter=50, print_flag=true) where {Tf<:Union{Function, Matrix}, Tv<:Number}

    # determine the type of A(matrix or operator)
    if isa(A, Matrix)
       explicitA = true

    elseif isa(A, Function)
       explicitA = false

    else
       error("A can be either a matrix or a function")
    end

    # initial guess of the unknowns
    if explicitA

       (m, n) = size(A)

       # provided initial guess
       if x0 != 0.0
          x = copy(x0)
          r = b - A * x
          s = A' * r - shift * x
       # zero vector as initial guess
       else
          x = zeros(eltype(b), n)
          r = copy(b)
          s = A' * r
       end

    else

       m = length(b)

       # provided initial guess
       if x0 != 0.0
          x = copy(x0)
          r = b - A(x, 1; op_params...)
          s = A(r, 2; op_params...) - shift * x
          n = length(s)
       # zero vector as initial guess
       else
          r = copy(b)
          s = A(b, 2; params...)
          n = length(s)
          x = zeros(eltype(b), n)
       end

    end

    data_fitting = dot(r,r)
    constraint   = 0.0
    cost0 = data_fitting + constraint

    # initialize some intermidiate vectors
    p = copy(s)
    norms0= vecnorm(s)
    gamma = norms0^2
    normx = vecnorm(x)
    xmax  = normx     # keep the initial one
    resNE = 1.0

    gamma0= copy(gamma)
    delta = 0.0

    # iteration counter and stop condition
    k = 0
    run_flag = true

    if print_flag
       header = "  k         data_fitting           constraint             normx                resNE"
       println(""); println(header);
       @printf("%3.0f %20.10e %20.10e %20.10e %20.10e\n", k, data_fitting, constraint, normx, resNE);
    end

    while (k < maxIter) && run_flag

          k = k + 1

          if explicitA
             q = A * p
          else
             q = A(p, 1; params...)
          end

          delta = (vecnorm(q))^2 + shift * (vecnorm(p))^2
          indefinite = delta <= 0 ? true : false
          delta      = delta == 0.? eps(): delta

          alpha = gamma / delta

          x = x + alpha * p
          r = r - alpha * q

          data_fitting = dot(r,r)
          constraint   = shift * dot(x, x)
          cost = data_fitting + constraint

          if explicitA
             s = A' * r - shift * x
          else
             s = A(r, 2; params...) - shift * x
          end

          norms  = norm(s)
          gamma0 = copy(gamma)
          gamma  = norms^2
          beta   = gamma / gamma0

          p = s + beta * p

          # check the stopping crietia
          normx = norm(x)
          xmax  = normx > xmax ? normx : xmax
          if norms <= norms0 * tol || normx * tol >= 1.0
             run_flag = false
          end

          # print information
          # resNE = norms / norms0
          resNE = cost  / cost0
          if print_flag
             @printf("%3.0f %20.10e %20.10e %20.10e %20.10e\n", k, data_fitting, constraint, normx, resNE);
          end

    end

    return x, resNE

end


#**************************************************************************************
# GeoEfficiency.jl package
# ========================
#
# represent a fast and flexible tool to calculate in batch or individually the geometrical efficiency
# for a set of common radiation detectors shapes (cylindrical, Bore-hole, Well-type) as seen form a sources.
# the source can be a point, a disc or even a cylinder.
#
#**************************************************************************************

module GeoEfficiency
print_with_color(:black,"""\n
    \t *************************************************
    \t **            -=) GeoEfficiency (=-            **
    \t **  Accurate Geometrical Efficiency Calculator **
    \t **                                             **
    \t *************************************************

  @author: Mohamed Krar
  @Profile: https://www.researchgate.net/profile/Mohamed_Krar3
  @repositry: https://github.com/DrKrar/GeoEfficiency.jl/
  @documentation: http://geoefficiencyjl.readthedocs.org
  @version: v"0.8.3"
  Created on Fri Aug 14 20:12:01 2015

  loading the Package:-
  ====================
  Pkg.add("GeoEfficiency")
  using GeoEfficiency

  Quick usage:-
  ============
  geoEff()	: calculate once and return the value.
  calc()	: calculate once and display full information on the console.
  calcN()	: calculate untill you quit.
  batch()	: calculate using data in the "GeoEfficiency" folder.
""")

export 	input,
		getfloat,
		getDetectors,

		Point,
		source,
		RadiationDetector,
		CylDetector,
		BoreDetector,
		WellDetector,
		detectorFactory,

		geoEff,

		calc,
		calcN,
		batch

include("Physical_model.jl")
include("Input_Interface.jl")
include("calculations.jl")
include("Output_Interface.jl")

end #module


"""
```
function solve_slm_system(Mdes::AbstractMatrix{S}, rhs::AbstractVector{S},
                          Mreg::AbstractMatrix{S}, rhsreg::AbstractVector{S},
                          λ::S, Meq::AbstractMatrix{S}, rhseq::AbstractVector{S},
                          Mineq::AbstractMatrix{S}, rhsineq::AbstractVector{S};
                          method::Symbol = :ipopt, atol::S = NaN, rtol::S = NaN,
                          max_iter::Int = 0, max_time::S = 30.,
                          init::AbstractVector{S} = Vector{S}(undef, 0),
                          use_lls::Bool = true, use_sparse::Bool = true,
                          verbose::Symbol = :low) where {S <: Real}
```

estimates the coefficient of a least squares spline by solving the
constrained linear least squares problem

```
min ½ || Mfit * x - rhsfit ||²
```

subject to the constraints

```
Meq * x = rhseq,    Mineq * x ≤ rhsineq,
```

where `Mfit = vcat(Mdes, λ .* Mreg)` and `rhsfit = vcat(rhs, λ .* rhsreg)`.

### Inputs
- `Mdes`: design matrix in system of equations defining the spline
- `rhs`: RHS of the system of equations defining the spline
- `Mreg`: matrix in system of equations regularizing the spline's smoothness
- `rhsreg`: RHS of the system of equations regularizing the spline smoothnes
- `λ`: degree of smoothness (smaller ⇒ smoother)
- `Meq`: matrix for equality constraints
- `rhseq`: the RHS of equality constraints
- `Mineq`: matrix for inequality constraints
- `rhsineq`: upper bound for inequality constraints

### Keywords
- `method`: minimization algorithm, currently only the interior point optimizer is used (Ipopt).
    New methods are welcome! One algorithm that may be added is adapting the matrices to the trust-region method
    via the calculation of pseudo-inverses
- `atol`: absolute tolerance for minimization algorithm
- `rtol`: relative tolerance for minimization algorithm
- `max_iter`: maximum number of evaluations
- `max_time`: maximum time (in minutes) for optimization
- `init`: initial guess of coefficients for the spline
- `use_lls`: use an `ADNLSModel` to define the problem. Currently, using an `LLSModel` (or equivalent) does not work.
    See the `NLPModels.jl` package.
- `use_sparse`: use sparse matrices (assuming the inputs are not already sparse matrices)
- `verbose`: how much information will be printed. Can be `:low` or `:high`
"""
function solve_slm_system(Mdes::AbstractMatrix{S}, rhs::AbstractVector{S},
                          Mreg::AbstractMatrix{S}, rhsreg::AbstractVector{S},
                          λ::S, Meq::AbstractMatrix{S}, rhseq::AbstractVector{S},
                          Mineq::AbstractMatrix{S}, rhsineq::AbstractVector{S};
                          method::Symbol = :ipopt, atol::S = 1e-8, rtol::S = 1e-6,
                          max_iter::Int = 3000, max_time::S = 3600.,
                          init::AbstractVector{S} = Vector{S}(undef, 0),
                          use_lls::Bool = true, use_sparse::Bool = true,
                          verbose::Symbol = :low, kwargs...) where {S <: Real}

    if use_sparse && all([issparse(x) for x in [Mdes, Mreg, Meq, Mineq]])
        # Don't spend time converting already sparse matrices
        use_sparse = false
    end

    Mfit = use_sparse ? vcat(sparse(Mdes), sparse(Mreg) .* λ) : vcat(Mdes, Mreg .* λ)
    rhsfit = vcat(rhs, rhsreg .* λ)

    if verbose == :high
        println("Condition number of the regression: $(cond(Mdes))")
    end

    if isempty(Mineq) && isempty(Meq)
        # Backslash suffices in this case
        coef = Mfit \ rhsfit
    else
        # Set up NLPModel and constraints
        if isempty(Mineq)
            lcon = rhseq
            ucon = rhseq
            C = use_sparse ? sparse(Meq) : Meq
        else
            lcon = vcat(fill(-Inf, length(rhsineq)), rhseq) # -Inf for the inequalities since that's how lsqlin is written in MATLAB
            ucon = vcat(rhsineq, rhseq)
            C = use_sparse ? vcat(sparse(Mineq), sparse(Meq)) : vcat(Mineq, Meq)
        end

        if use_lls
            nls = LLSModel(Mfit, rhsfit; C = C, lcon = lcon, ucon = ucon)
        else
            if isempty(init)
                init = zeros(S, size(Mfit, 2))
            end

            nls = ADNLSModel(x -> Mfit * x - rhsfit, init, size(Mfit, 1), c = x -> C * x, lcon = lcon, ucon = ucon)
        end

        # Decide solver for minimizing the constrained linear least-squares problem
        if method == :ipopt
            if use_lls
                # Feasibility Form adds additional coefficient which we don't want, hence
                # the [1:size(Mdes, 2)] slicing.
                coef = ipopt(FeasibilityFormNLS(nls); print_level = verbose == :high ? 3 : 0,
                             tol = atol, acceptable_tol = rtol, max_iter = max_iter,
                             max_cpu_time = max_time, kwargs...).solution[1:size(Mdes, 2)]
            else
                #
                coef = ipopt(nls; print_level = verbose == :high ? 3 : 0, tol = atol, acceptable_tol = rtol,
                             max_iter = max_iter, max_cpu_time = max_time, kwargs...).solution
            end
        else
            error("Method $method not is not implemented.")
        end
    end

    if verbose == :high
        println("Solver employed: $(out_info[:solver])")
        if !isempty(Mineq)
            println(nls)
        end
    end

    return coef
end


struct adVAE
	encoder
	generator
	transformer
end

function adVAE(
		;idim::Int=1,
		zdim::Int=1,
		hdim::Int=64,
		activation = "relu",
		nlayers::Int=3,
		init_seed=nothing,
		kwargs...
		)
	(init_seed !== nothing) ? Random.seed!(init_seed) : nothing

	encoder = Flux.Chain(
		build_mlp(idim, hdim, hdim, nlayers-1, activation=activation),
		ConditionalDists.SplitLayer(hdim, [zdim,zdim], [identity,softplus])
	)
	transformer = Flux.Chain(
		build_mlp(zdim*2, hdim, hdim, 2, activation=activation),
		ConditionalDists.SplitLayer(hdim, [zdim,zdim], [identity,softplus])
	)
	generator = build_mlp(zdim, hdim, idim, nlayers, activation=activation, lastlayer="linear")

	(init_seed !== nothing) ? Random.seed!() : nothing
	return adVAE(encoder, generator, transformer)
end

function Conv_adVAE(
		;isize::Int=1,
		in_ch::Int=1,
		zdim::Int=64,
		nf::Int=16,
		extra_layers::Int=0,
		depth::Int=3,
		activation="relu",
		init_seed=nothing,
		kwargs...
		)

	(init_seed !== nothing) ? Random.seed!(init_seed) : nothing
	encoder = Flux.Chain(
		ConvEncoder(isize, in_ch, zdim, nf, extra_layers),
		x->reshape(x, (zdim,size(x)[end])), # equivalent to flatten
		ConditionalDists.SplitLayer(zdim, [zdim, zdim], [identity, softplus])
	)
	transformer = Flux.Chain(
		build_mlp(zdim*2, zdim, zdim, depth, activation=activation),
		ConditionalDists.SplitLayer(zdim, [zdim,zdim], [identity,softplus])
	)
	generator = Flux.Chain(
		Flux.Dense(zdim, zdim, eval(:($(Symbol(activation)))) ),
		x->reshape(x, (1,1,zdim,size(x)[end])), # inverse to flatten
		ConvDecoder(isize, zdim, in_ch, nf, extra_layers)
	)
	(init_seed !== nothing) ? Random.seed!() : nothing
	return adVAE(encoder, generator, transformer)
end

Flux.@functor adVAE

function (advae::adVAE)(x)
	μ, Σ = advae.encoder(x)
	z = μ + Σ * randn(Float32)
	xᵣ = advae.generator(z)
end

kl_divergence(μ, Σ) = - Flux.mean(0.5f0 * sum(1f0 .+ log.(Σ.^2) - μ.^2  - Σ.^2, dims=1))
kl_divergence(μ₁, Σ₁, μ₂, Σ₂) = Flux.mean(sum(log.(Σ₂) - log.(Σ₁) + (Σ₁.^2 + (μ₁ - μ₂).^2) ./ (2*Σ₂.^2) .- 0.5f0, dims=1))

function loss(advae::adVAE, x; γ=1e-3, λ=1e-2, mx=1, mz=1)
	μ, Σ = advae.encoder(x)
	μₜ, Σₜ = advae.transformer(cat(μ, Σ, dims=1))

	z = μ + Σ * randn(Float32)
	zₜ = μₜ + Σₜ * randn(Float32)

	xᵣ = advae.generator(z)
	xₜᵣ = advae.generator(zₜ)

	μᵣ, Σᵣ = advae.encoder(xᵣ)
	μₜᵣ, Σₜᵣ = advae.encoder(xₜᵣ)

	# 𝓛 for generator => 𝓛 = 𝓛_z + 𝓛_zₜ
	𝓛_z = Flux.Losses.mse(x, xᵣ) .+ γ * kl_divergence(μᵣ, Σᵣ)
	𝓛_zₜ = max(0f0, (mx - Flux.Losses.mse(xᵣ, xₜᵣ))) + γ * max(0f0, (mz - kl_divergence(μₜᵣ, Σₜᵣ)))
	𝓛 = 𝓛_z + 𝓛_zₜ
	# 𝓛ₜ for transformer
	𝓛ₜ = kl_divergence(μ, Σ, μₜ, Σₜ)
	# 𝓛ₑ for encoder
	𝓛ₑ = Flux.Losses.mse(x, xᵣ) .+ λ * kl_divergence(μ, Σ)
		+ γ * max(0f0, (mz - kl_divergence(μᵣ, Σᵣ)))
		+ γ * max(0f0, (mz - kl_divergence(μₜᵣ, Σₜᵣ)))

	return 𝓛 + λ*𝓛ₜ, 𝓛ₑ, 𝓛, 𝓛ₜ
end

function anomaly_score(advae::adVAE, real_input; L=100, dims=3, batch_size=64, to_testmode::Bool=true)
	real_input = Flux.Data.DataLoader(real_input, batchsize=batch_size)
	(to_testmode == true) ? Flux.testmode!(advae) : nothing
	advae = advae |> gpu
	output = Array{Float32}([])
	for x in real_input
        x = x |> gpu
		X = Array{Float32}(undef, size(x)[end], 1)
		μ, Σ = advae.encoder(x)
		for l=1:L
			z = μ + Σ * randn(Float32)
			xᵣ = advae.generator(z)
			rec_loss = vec(Flux.mse(x, xᵣ, agg=x->Flux.mean(x, dims=dims))) |> cpu # loss per batch
			X = cat(X, rec_loss, dims=2)
		end
		output = cat(output, Flux.mean(X, dims=2), dims=1)
	end
	(to_testmode == true) ? Flux.testmode!(advae, false) : nothing
	return output
end

function StatsBase.fit!(advae::adVAE, data, params)
	# prepare batches & loaders
	train_loader, val_loader = prepare_dataloaders(data, params)
	# training info logger
	history = MVHistory()
	# prepare for early stopping
	best_adVAE = deepcopy(advae)
	patience = params.patience
	best_val_loss = 1e10
	val_batches = length(val_loader)

	# ADAMW(η = 0.001, β = (0.9, 0.999), decay = 0) = Optimiser(ADAM(η, β), WeightDecay(decay))
	opt_step1 = haskey(params, :decay) ? ADAMW(params.lr, (0.9, 0.999), params.decay) : ADAM(params.lr)
	opt_step2 = haskey(params, :decay) ? ADAMW(params.lr, (0.9, 0.999), params.decay) : ADAM(params.lr)

	ps_step1 = Flux.params(advae.generator, advae.transformer)
	ps_step2 = Flux.params(advae.encoder)

	progress = Progress(length(train_loader))
	for (iter, X) in enumerate(train_loader)
		# training step 1
		loss_1, back = Flux.pullback(ps_step1) do
			loss(advae|>gpu , getobs(X)|>gpu, γ=params.gamma, λ=params.lambda, mx=params.mx, mz=params.mz)
		end
		grad_step1 = back((1f0, 0f0, 0f0, 0f0))
		Flux.Optimise.update!(opt_step1, ps_step1, grad_step1)
		# training step 2
		loss_2, back = Flux.pullback(ps_step2) do
			loss(advae|>gpu , getobs(X)|>gpu, γ=params.gamma, λ=params.lambda, mx=params.mx, mz=params.mz)
		end
		grad_step2 = back((0f0, 1f0, 0f0, 0f0))
		Flux.Optimise.update!(opt_step1, ps_step2, grad_step2)

		push!(history, :loss_step_1, iter, loss_1[1])
		push!(history, :encoder_loss, iter, loss_2[2])
		push!(history, :generator_loss, iter, loss_1[3])
		push!(history, :transformer_loss, iter, loss_1[4])

		next!(progress; showvalues=[
			(:iters, "$(iter)/$(params.iters)"),
			(:loss_step_1, loss_1[1]),
			(:loss_step_2, loss_2[2])
			])
		# TODO: check and prepare early stopping
		if mod(iter, params.check_every) == 0
			total_val_loss = 0
			Flux.testmode!(advae)
			for X_val in val_loader
				X_val = X_val |> gpu
				(𝓛, 𝓛ₑ, _, _) = loss(advae|>gpu , X_val, γ=params.gamma, λ=params.lambda, mx=params.mx, mz=params.mz)
				total_val_loss += 𝓛 + 𝓛ₑ
			end
			Flux.testmode!(advae, false)
			push!(history, :validation_loss, iter, total_val_loss/val_batches)

			if total_val_loss < best_val_loss
				best_val_loss = total_val_loss
				patience = params.patience
				best_adVAE = deepcopy(advae)
			else
				patience -= 1
				if patience == 0
					@info "Stopped training after $(iter) iterations"
					global iters = iter - params.check_every*params.patience
					break
				end
			end
		end
		if iter == params.iters
			global iters = params.iters
		end
	end
	return history, best_adVAE, sum(map(p->length(p), Flux.params(advae))), iters
end


export contour

"""
    contour(f; options...)

Compute the contour 𝐶(𝑓) of the amoeba ``\\mathcal{A}(f)``.

## Example
```julia
@polyvar x y

contour(x^2+y^2+1)

# custom domain
contour(x^2+y^2+1, domain=(-5, 5, -5, 5))
```

## Options
* `domain`: A tuple in the form `(xmin, xmax, ymin, ymax)` which defines a section Ω
for which the contour 𝐶(𝑓) is computed.
* `membership_options`: The options for the membership test
* `res=(600,600)`: The resolution with which starting points are sampled
* `samples_off_axis=2*MP.maxdegree(p)^2`: The number of sample points per off-axis.
"""
function contour(f::MP.AbstractPolynomial;
    domain=amoeba_carcase_domain_heuristic(f, factor=2.0),
    res::NTuple{2, Int}=(600, 600),
    samples_off_axis::Int=2*MP.maxdegree(f),
    membership_options=MembershipTestOptions(ntries=10, maxiters=30, tol=1e-8))

    grid = Grid2D(domain..., res...)

    _contour2D(ContourFiber2D(f), grid, samples_off_axis, membership_options)
end

# only internally used
"""
    grid_contour(F::ContourFiber2D, grid::Grid2D; kwargs...)

Instead of returning a list of contour points it fills a bitmap with the nearest contour
points.
"""
function grid_contour(p::MP.AbstractPolynomial; res=600, kwargs...)
    grid_contour(p, Grid2D(p, res); kwargs...)
end

function grid_contour(p::MP.AbstractPolynomial, grid::Grid2D;
    every_main_axis::Int=10, samples_off_axis::Int=8, options=MembershipTestOptions(ntries=3, maxiters=40, tol=1e-8))

    fiber = ContourFiber2D(p)

    _contour2D_grid(fiber, grid, every_main_axis, samples_off_axis, options)
end


push!(LOAD_PATH, pwd())
using Day3_module

wire_paths = readlines("input.txt")
#wire_paths = ['R75,D30,R83,U83,L12,D49,R71,U7,L72',
#              'U62,R66,U55,R34,D71,R55,D58,R83']
#wire_paths = ['R98,U47,R26,D63,R33,U87,L62,D20,R33,U53,R51',
#              'U98,R91,D20,R16,D67,R40,U7,R15,U6,R7']

function get_steps(wire_path)
    return split(wire_path, ',')
end

t0 = time()
pos_a, pos_b = [get_positions(get_steps(path)) for path in wire_paths]
crossings = intersect(Set(pos_a), Set(pos_b))

function manhattan_distance(pos)
    sum(map(abs, pos))
end

manhattan_distances = [manhattan_distance(pos) for pos in crossings]
println(minimum(manhattan_distances))

t1 = time()
println(t1 - t0)

println(minimum(indexin(crossings, pos_b) + indexin(crossings, pos_a)))

t2 = time()
println(t2 - t1)


# transportMatrixTests.jl

# Age test, should be ≈ 1000 years
function computeIdealMeanAge(wet3d,T)
    l3d = copy(wet3d)
    l3d[:,:,2:end] .= 0
    iwet = find(wet3d)
  L = spdiagm(l3d[iwet])
    nwet = length(iwet)
  return (T + L) \ ones(nwet)
end
meanΓyrs = mean(computeIdealMeanAge(wet3d,T)) / (365 * spd)
@test 100 < meanΓyrs < 1000

# T 1 ≈ 0 test
Tcolmax = vec(maximum(abs.(T),1))
Trowmax = vec(maximum(abs.(T),2))
T1 = T * ones(nwet)
maximum(abs.(T1) ./ Tcolmax)
maximum(abs.(T1) ./ Trowmax)

# vᵀ T ≈ 0 test
vᵀT = T.' * v
maximum(abs.(vᵀT) ./ (v .* Tcolmax))
maximum(abs.(vᵀT) ./ (v .* Trowmax))

# S 1 ≈ 0 test
Scolmax = vec(maximum(abs.(S),1))
Srowmax = vec(maximum(abs.(S),2))
S1 = S * ones(nwet)
maximum(abs.(S1) ./ Scolmax)
maximum(abs.(S1) ./ Srowmax)
S13d = NaN * wet3d
S13d[iwet] .= S1
using Plots
gr()
plot(S13d[50,54,:],line = :scatter)

# vᵀ S ≈ 0 test
vᵀS = S.' * v
maximum(abs.(vᵀS) ./ (v .* Scolmax))
maximum(abs.(vᵀS) ./ (v .* Srowmax))

function isMassConserving(x,v)
  vᵀx = v.' * x
  return vᵀx


@testset "broadcast" begin
    i = Indices([1,2,3,4,5])

    @test isequal((i .+ 1)::Dictionary, dictionary([1=>2, 2=>3, 3=>4, 4=>5, 5=>6]))
    @test isequal(Base.Broadcast.broadcasted(+, i, 1)::BroadcastedDictionary, dictionary([1=>2, 2=>3, 3=>4, 4=>5, 5=>6]))

    @test isequal((i .+ i)::Dictionary, dictionary([1=>2, 2=>4, 3=>6, 4=>8, 5=>10]))
    @test isequal(Base.Broadcast.broadcasted(+, i, i)::BroadcastedDictionary, dictionary([1=>2, 2=>4, 3=>6, 4=>8, 5=>10]))

    @test isequal((i .+ copy(i))::Dictionary, dictionary([1=>2, 2=>4, 3=>6, 4=>8, 5=>10]))
    @test isequal(Base.Broadcast.broadcasted(+, i, copy(i))::BroadcastedDictionary, dictionary([1=>2, 2=>4, 3=>6, 4=>8, 5=>10]))

    @test_throws IndexError Indices([1,2]) .+ Indices([2,3])

    d = i .+ 1

    @test isequal((d .+ 1)::Dictionary, dictionary([1=>3, 2=>4, 3=>5, 4=>6, 5=>7]))
    @test isequal(Base.Broadcast.broadcasted(+, d, 1)::BroadcastedDictionary, dictionary([1=>3, 2=>4, 3=>5, 4=>6, 5=>7]))

    @test isequal((d .+ d)::Dictionary, dictionary([1=>4, 2=>6, 3=>8, 4=>10, 5=>12]))
    @test isequal(Base.Broadcast.broadcasted(+, d, d)::BroadcastedDictionary, dictionary([1=>4, 2=>6, 3=>8, 4=>10, 5=>12]))

    @test isequal((d .+ copy(d))::Dictionary, dictionary([1=>4, 2=>6, 3=>8, 4=>10, 5=>12]))
    @test isequal(Base.Broadcast.broadcasted(+, d, copy(d))::BroadcastedDictionary, dictionary([1=>4, 2=>6, 3=>8, 4=>10, 5=>12]))

    d2 = similar(d)
    d2 .= d .+ d
    @test isequal(d2, dictionary([1=>4, 2=>6, 3=>8, 4=>10, 5=>12]))

    @test_throws IndexError Dictionary([1,2],[1,2]) .+ Dictionary([2,3],[2,3])
end


function rankfractional{T<:Real}(a::Array{T,1})
    r = float64(collect(1:length(a)))
    1.0 < r[end] || return r
    for i in ties(a)
        r[a.==i] = mean(r[a.==i])
    end
    return r
end


using Distributions
using Test, Random

@test typeof(rand(Bernoulli())) == Int
@test typeof(rand(Bernoulli(), 10)) == Vector{Int}


#
#--------------------------------------#
export diver
#--------------------------------------#
"""
 for q ∈ H¹₀(Ω)  (test  pressure)
 for u⃗ ∈ H¹₀(Ω)² (trial velocity)

 (q,∇⋅u⃗) = (J*q)' * Bᵥ * [Dx Dy][ux]
                                [uy]

 Bᵥ is mass matrix on velocity grid
 J  is the interpolation matrix
 from pressure grid to velocity grid.

"""
function diver(ux  ::Array
              ,uy  ::Array
              ,mshV::Mesh
              ,Jr,Js)

    uxdx,_    = grad(ux,msh)
    _   ,uydy = grad(uy,msh)

    div = uxdx + uydy

    Bdiv  = mass(div,msh)
    JBdiv = ABu(Js',Jr',Bdiv)

    return JBdiv # pressure grid
end
#--------------------------------------#
export diverᵀ
#--------------------------------------#
"""
 for p ∈ H¹₀(Ω)  (trial pressure)
 for v⃗ ∈ H¹₀(Ω)² (test  velocity)

 -(v⃗,∇p) = -([vx] .* [px])
            ([vy]    [py])

 integrating by parts,

 -(∇v⃗,p) = ([Dx Dy] * [vx])' Bᵥ * J * p
           (          [vy])

         + surface terms (vanish)

         = [vx vy] *  [Dx'] Bᵥ * J * p
                      [Dy']

"""
function diverᵀ(pr  ::Array
               ,mshV::Mesh
               ,Jr,Js)

    Jp  = ABu(Js,Jr,pr) # pr -> vel gird
    BJp = mass(Jp,mshV)

    qx,qy = gradᵀ(BJp,mshV)

    return qx,qy
end
#--------------------------------------#
export stokesOp
#--------------------------------------#
"""
 Schur complement + operator splitting
 operation for stokes solve

 EE = -DD AA^-1 * DD'
"""
function stokesOp(q   ::Array
                 ,mshV::Mesh
                 ,Mvx,Mvy,Jr,Js)

    # DD'
    qx,qy = diverᵀ(q,mshV,Jr,Js)

    # HHinv
    qx = approxHlmzInv(qx,mshV)
    qy = approxHlmzInv(qy,mshV)

    # DD
    Eq = diver(qx,qy,mshV,Jr,Js)

  return -Eq
end
#--------------------------------------#
export approxHlmzInv
#--------------------------------------#
function approxHlmzInv(u::Array,b0::Number
                      ,mshV::Mesh)

    v = gatherScatter(u,mshV)
    v = mask(v,Mvx)

    v = v .* mshV.Bi ./ b0

    v = gatherScatter(v,mshV)
    v = mask(v,Mvx)

    return v
end
#--------------------------------------#


using ApproxOperator, LinearAlgebra

aps = import_msh("./msh/cube_pfm.msh")
nₚ = length(aps["surface"][1].nodes)
set_integration_rule!(aps["surface"],Val(:PhaseFieldFriction))

E = 2.0e11
ν = 0.3
c = 40
𝜙 = 15/180*π
𝜙ᵣ = 15/180*π
θ = 0.
u₀ = 5.0e-3
nls = 50
k̄ = 2700
l = 0.01
tol = 1e-10
maxiter = 50

b = (x,y,z)->(0.0,0.0)
n = (x,y,z)->(1.0,0.0,1.0)
ū = (x,y,z)->(0.0,0.0,0.0)

ops = [
    NonlinearPlaneStress_C_Ω(b),
    NonlinearPlaneStress_Γᵍ_penalty(ū,n,E*1e7),
    SecondOrderPhaseField(k̄,l),
    Update_Friction_PhaseField_PlaneStress(E,ν,c,𝜙,𝜙ᵣ,1e-6),
    VTKExport(
                filename="figure/movie00.vtk",
                topic="phase field modeling fracture",
                parameters=Dict{String,Float64}("E"=>E,"ν"=>ν)
            )
]

d = zeros(2*nₚ)
k = zeros(2*nₚ,2*nₚ)
f = zeros(2*nₚ)
kᵥ = zeros(nₚ,nₚ)
fᵥ = zeros(nₚ)
dᵥ = zeros(nₚ)
dᵗ = zeros(nₚ)

ops[5](aps["surface"],d,dᵥ,Val(:PFM_PlaneStress))

for t in 1:nls
    # define essential BC
    u = (x,y,z)->(cos(θ)*u₀/nls*t,sin(θ)*u₀/nls*t)
    op_EBC = NonlinearPlaneStress_Γᵍ_penalty(u,n,E*1e7)
    iter = 0

    # Alternate minimization algorithm
    while iter ≤ maxiter
        dᵗ .= dᵥ
        iter += 1

        # update history field
        ops[4](aps["surface"],d,dᵥ)

        # loop of elasticity
        iteru = 0
        while iteru ≤ maxiter
            iteru += 1
            fill!(k,0.0)
            fill!(f,0.0)
            ops[1](aps["surface"],k,f)
            ops[2](aps["disp_fix"],k,f,d)
            op_EBC(aps["disp"],k,f,d)
            erru = norm(f)/E/1e7
            print("u: time step = $t, alternate iter = $iter, disp iter = $iteru, erru = $erru\n")
            if erru < tol;break;end
            d .+= k\f
        end

        # loop of phase field
        iterv = 0
        while iterv ≤ maxiter
            iterv += 1
            fill!(kᵥ,0.0)
            fill!(fᵥ,0.0)
            ops[3](aps["surface"],kᵥ,fᵥ,dᵥ)
            errv = norm(fᵥ)/k̄
            print("v: time step = $t, alternate iter = $iter, damage iter = $iterv, errv = $errv\n")
            if errv < tol;break;end
            dᵥ .+= kᵥ\fᵥ
        end

        err = norm(dᵗ .- dᵥ)/norm(dᵥ)
        print("err = $err\n")
        if err < tol;break;end
    end
    for ap in aps["surface"]
        for qw in ap.qw
            qw.ℋₙ = qw.ℋ
            qw.ℋₘ = max(qw.ℋₘ,qw.ℋₙ)
        end
    end
    # export to VTK file
    ops[5].filename = "figure/movie"*string(t,pad=2)*".vtk"
    ops[5](aps["surface"],d,dᵥ,Val(:PFM_PlaneStress))
end


import Base: vec

export
ruleSPNonlinearEOutNG,
ruleSPNonlinearEOutNGX,
ruleSPNonlinearEIn1GG,
ruleSPNonlinearEInGX,
ruleMNonlinearEInGX

"""
Concatenate a vector (of vectors and floats) and return with original dimensions (for splitting)
"""
function concatenate(xs::Vector)
    ds = size.(xs) # Extract dimensions
    x = vcat(vec.(xs)...)

    return (x, ds)
end

ForneyLab.vec(d::Float64) = [d] # Extend vectorization to Float

"""
Return local linearization of g around expansion point x_hat
for Nonlinear node with single input interface
"""
function localLinearizationSingleIn(g::Function, x_hat::Float64)
    a = ForwardDiff.derivative(g, x_hat)
    b = g(x_hat) - a*x_hat

    return (a, b)
end

function localLinearizationSingleIn(g::Function, x_hat::Vector{Float64})
    A = ForwardDiff.jacobian(g, x_hat)
    b = g(x_hat) - A*x_hat

    return (A, b)
end

"""
Return local linearization of g around expansion point x_hat
for Nonlinear node with multiple input interfaces
"""
function localLinearizationMultiIn(g::Function, x_hat::Vector{Float64})
    g_unpacked(x::Vector) = g(x...)
    A = ForwardDiff.gradient(g_unpacked, x_hat)'
    b = g(x_hat...) - A*x_hat

    return (A, b)
end

function localLinearizationMultiIn(g::Function, x_hat::Vector{Vector{Float64}})
    (x_cat, ds) = concatenate(x_hat)
    g_unpacked(x::Vector) = g(split(x, ds)...)
    A = ForwardDiff.jacobian(g_unpacked, x_cat)
    b = g(x_hat...) - A*x_cat

    return (A, b)
end

#-----------------------
# Extended Update Rules
#-----------------------

# Forward rule
function ruleSPNonlinearEOutNG(g::Function,
                               msg_out::Nothing,
                               msg_in1::Message{<:Gaussian})

    (m_in1, V_in1) = unsafeMeanCov(msg_in1.dist)
    (A, b) = localLinearizationSingleIn(g, m_in1)
    m = A*m_in1 + b
    V = A*V_in1*A'

    return Message(variateType(m), GaussianMeanVariance, m=m, v=V)
end

# Multi-argument forward rule
function ruleSPNonlinearEOutNGX(g::Function, # Needs to be in front of Vararg
                                msg_out::Nothing,
                                msgs_in::Vararg{Message{<:Gaussian}})

    (ms_fw_in, Vs_fw_in) = collectStatistics(msgs_in...) # Returns arrays with individual means and covariances
    (A, b) = localLinearizationMultiIn(g, ms_fw_in)
    (m_fw_in, V_fw_in, _) = concatenateGaussianMV(ms_fw_in, Vs_fw_in)
    m = A*m_fw_in + b
    V = A*V_fw_in*A'

    return Message(variateType(m), GaussianMeanVariance, m=m, v=V)
end

# Backward rule with given inverse
function ruleSPNonlinearEIn1GG(g::Function,
                               g_inv::Function,
                               msg_out::Message{<:Gaussian},
                               msg_in1::Nothing)

    (m_out, V_out) = unsafeMeanCov(msg_out.dist)
    (A, b) = localLinearizationSingleIn(g_inv, m_out)
    m = A*m_out + b
    V = A*V_out*A'

    return Message(variateType(m), GaussianMeanVariance, m=m, v=V)
end

# Multi-argument backward rule with given inverse
function ruleSPNonlinearEInGX(g::Function, # Needs to be in front of Vararg
                              g_inv::Function,
                              msg_out::Message{<:Gaussian},
                              msgs_in::Vararg{Union{Message{<:Gaussian}, Nothing}})

    (ms, Vs) = collectStatistics(msg_out, msgs_in...) # Returns arrays with individual means and covariances
    (A, b) = localLinearizationMultiIn(g_inv, ms)
    (mc, Vc) = concatenateGaussianMV(ms, Vs)
    m = A*mc + b
    V = A*Vc*A'

    return Message(variateType(m), GaussianMeanVariance, m=m, v=V)
end

# Backward rule with unknown inverse
function ruleSPNonlinearEIn1GG(g::Function,
                               msg_out::Message{<:Gaussian},
                               msg_in1::Message{<:Gaussian})

    m_in1 = unsafeMean(msg_in1.dist)
    (m_out, W_out) = unsafeMeanPrecision(msg_out.dist)
    (A, b) = localLinearizationSingleIn(g, m_in1)
    xi = A'*W_out*(m_out - b)
    W = A'*W_out*A

    return Message(variateType(xi), GaussianWeightedMeanPrecision, xi=xi, w=W)
end

# Multi-argument backward rule with unknown inverse
function ruleSPNonlinearEInGX(g::Function,
                              inx::Int64, # Index of inbound interface inx
                              msg_out::Message{<:Gaussian},
                              msgs_in::Vararg{Message{<:Gaussian}})

    # Approximate joint inbounds
    (ms_fw_in, Vs_fw_in) = collectStatistics(msgs_in...) # Returns arrays with individual means and covariances
    (A, b) = localLinearizationMultiIn(g, ms_fw_in)

    (m_fw_in, V_fw_in, ds) = concatenateGaussianMV(ms_fw_in, Vs_fw_in)
    m_fw_out = A*m_fw_in + b
    V_fw_out = A*V_fw_in*A'
    C_fw = V_fw_in*A'

    # RTS Smoother
    (m_bw_out, V_bw_out) = unsafeMeanCov(msg_out.dist)
    (m_in, V_in) = smoothRTS(m_fw_out, V_fw_out, C_fw, m_fw_in, V_fw_in, m_bw_out, V_bw_out)

    # Marginalize joint belief on in's
    (m_inx, V_inx) = marginalizeGaussianMV(m_in, V_in, ds, inx)
    W_inx = cholinv(V_inx) # Convert to canonical statistics
    xi_inx = W_inx*m_inx

    # Divide marginal on inx by forward message
    (xi_fw_inx, W_fw_inx) = unsafeWeightedMeanPrecision(msgs_in[inx].dist)
    xi_bw_inx = xi_inx - xi_fw_inx
    W_bw_inx = W_inx - W_fw_inx # Note: subtraction might lead to posdef violations

    return Message(variateType(xi_bw_inx), GaussianWeightedMeanPrecision, xi=xi_bw_inx, w=W_bw_inx)
end

function ruleMNonlinearEInGX(g::Function,
                             msg_out::Message{<:Gaussian},
                             msgs_in::Vararg{Message{<:Gaussian}})

    # Approximate joint inbounds
    (ms_fw_in, Vs_fw_in) = collectStatistics(msgs_in...) # Returns arrays with individual means and covariances
    (A, b) = localLinearizationMultiIn(g, ms_fw_in)

    (m_fw_in, V_fw_in, _) = concatenateGaussianMV(ms_fw_in, Vs_fw_in)
    m_fw_out = A*m_fw_in + b
    V_fw_out = A*V_fw_in*A'
    C_fw = V_fw_in*A'

    # RTS Smoother
    (m_bw_out, V_bw_out) = unsafeMeanCov(msg_out.dist)
    (m_in, V_in) = smoothRTS(m_fw_out, V_fw_out, C_fw, m_fw_in, V_fw_in, m_bw_out, V_bw_out)

    return ProbabilityDistribution(Multivariate, GaussianMeanVariance, m=m_in, v=V_in)
end


module CoupledFields

using LinearAlgebra
using Statistics
using StatsBase: sample, zscore
export InputSpace, ModelObj
export KernelParameters, GaussianKP, PolynomialKP
export gradvecfield
export bf, cca, gKCCA
export CVfn, Rsq_adj
export whiten
include("MLKernels.jl")

"""
    KernelParameters: An abstract type.
All KernelParameters types contain certain parameters which are later passed to internal functions `Kf` and `∇Kf`. \n
A KernelParameters type is set using e.g. `PolynomialKP(X::Matrix{Float64})` or `GaussianKP(X::Matrix{Float64})`.
"""
abstract type KernelParameters end

"""
    GaussianKP: For the gaussian kernel
"""
struct GaussianKP <: KernelParameters
    xx::Matrix{Float64}
    varx::Float64
end

function GaussianKP(X::Matrix{Float64})
    xx = kernelmatrix(SquaredDistanceKernel(1.0), X, X)
    varx = Statistics.median(xx[xx.>10.0^-9])
    return GaussianKP(xx, varx)
end

    function Kf(par::Array{T}, X::Matrix{T}, kpars::GaussianKP) where T<:Float64
        sx2 = 2*par[1]*par[1]*kpars.varx
        return exp.(-kpars.xx/sx2)
    end

    function ∇Kf(par::Array{T}, X::Matrix{T}, kpars::GaussianKP) where T<:Float64
        n,p = size(X)
        sx2 = 2*par[1]*par[1]*kpars.varx
        Gx = Kf(par, X, kpars)
        ∇K = Float64[2.0*Gx[i,k]*(X[i,j]-X[k,j])/sx2  for i in 1:n, k in 1:n, j in 1:p ]
        return ∇K
    end

"""
    PolynomialKP: For the polynomial kernel
"""
struct PolynomialKP <: KernelParameters
    xx::Matrix{Float64}
    function PolynomialKP(X::Matrix{Float64})
        xx = kernelmatrix(LinearKernel(1.0, 1.0), X, X)
        return new(xx)
    end
end

    function Kf(par::Array{T}, X::Matrix{T}, kpars::PolynomialKP) where T<:Float64
        return (kpars.xx).^par[1]
    end

    function ∇Kf(par::Array{T}, X::Matrix{T}, kpars::PolynomialKP) where T<:Float64
        n,p = size(X)
        m = par[1]
        ∇K = Float64[m*X[k,j]*kpars.xx[i,k]^(m-1.0)  for i in 1:n, k in 1:n, j in 1:p ]
        return ∇K
    end

"""
    ModelObj: A type to hold statistical model results
Such as the matrices `W, R, A, T`, where `R=XW` and `T=YA`.
"""
struct ModelObj
    W::Matrix{Float64}
    R::Matrix{Float64}
    A::Matrix{Float64}
    T::Matrix{Float64}
    evals::Vector{Float64}
    pars::Array{Float64}
    method::String
end

"""
    InputSpace: A type to hold the `X` and `Y` fields of the Input space
InputSpace(X, Y, d, lat): The fields are whitened if `d=[d1, d2]` is supplied.
Area weighting is applied if `lat` is supplied.
"""
struct InputSpace
    X::Matrix{Float64}
    Y::Matrix{Float64}
    function InputSpace(a::T, b::T) where T<:Matrix{Float64}
        new(zscore(a,1), zscore(b,1))
    end
    function InputSpace(a::T, b::T, d::Vector{Float64}) where T<:Matrix{Float64}
        new(whiten(a, d[1]), whiten(b, d[2]))
    end
    function InputSpace(a::T, b::T, d::V, lat::V) where {V<:Matrix{Float64}, T<:Vector{Float64}}
        new(whiten(a, d[1], lat=lat), whiten(b, d[2], lat=lat))
    end
end

"""
    gradvecfield(par::Array{Float64}, X::T, Y::T, kpars::KernelParameters) where T<:Matrix{Float64}:
Compute the gradient vector or gradient matrix at each instance of the `X` and `Y` fields, by making use of a kernel feature space.
"""
function gradvecfield(par::Array{Float64}, X::T, Y::T, kpars::KernelParameters ) where T<:Matrix{Float64}
    n,p = size(X)
    Gx = Kf(par, X, kpars) + (10.0^par[2])*n*I
    ∇K = ∇Kf(par, X, kpars)
    return [∇K[i,:,:]' * (Gx \ Y) for i in 1:n]
end

"""
    bf(x::Vector{Float64}, df::Int):
Compute a piecewise linear basis matrix for the vector x.
"""
function bf(x::Vector{Float64}, df::Int)
    if df<2 return x[:,1:1] end
    bp = Statistics.quantile(x, range(0, stop=1, length=df+1))
    n = length(x)
    a1 = repeat(x, inner=(1,df))
    a2 = repeat( permutedims(bp[1:df]), inner=(n,1))

    for j in 1:df
        ix1, ix2 = x.<bp[j], x.>bp[j+1]
        a1[ix1,j] .= bp[j]
        a1[ix2,j] .= bp[j+1]
    end

    return a1-a2
end

"""
    cca(v::Array{Float64}, X::T,Y::T) where T<:Matrix{Float64}:
Regularized Canonical Correlation Analysis using SVD.
"""
function cca(v::Array{Float64}, X::T,Y::T) where T<:Matrix{Float64}
#    n,p = size(X)
    q = size(Y, 2)

    Cxx_fac_inv = (Statistics.cov(X)+(10.0^v[1])*I)^-0.5
    Cyy_fac_inv = (q>1) ? (Statistics.cov(Y)+(10.0^v[2])*I)^-0.5 : Statistics.cov(Y)^(-0.5)
    M = Statistics.cov(X*Cxx_fac_inv, Y*Cyy_fac_inv)
    U, D, V = svd(M)

    Wx = Cxx_fac_inv * U
    Wy = Cyy_fac_inv * V
    Tx = X * Wx
    Ty = Y * Wy

    L = D.^2
    return ModelObj(Wx, Tx, Wy, Ty, L, v, "CCA")
end

cca(v::Array{Float64}, X::T,Y::T, kpars::KernelParameters) where {T<:Matrix{Float64}} = cca(v, X, Y)

"""
    gKCCA(par::Array{Float64}, X::Matrix{Float64}, Y::Matrix{Float64}, kpars::KernelParameters):
Compute the projection matrices and components for gKCCA.
"""
function gKCCA(par::Array{Float64}, X::Matrix{Float64}, Y::Matrix{Float64}, kpars::KernelParameters)
    q = size(Y,2)
    n, p = size(X)
    ∇g = gradvecfield(par, X, Y, kpars)
    M = mapreduce(x -> x*x', +, ∇g)/n
    values, vectors = eigen(Symmetric(M))
    Wx = reverse(vectors, dims=2)
    R = X * Wx

    dmin = minimum([p q 3])
    Ay = zeros(q, dmin)
    T = zeros(n, dmin)
    for j in 1:dmin
        ϕRj = bf(R[:,j], Int(par[3]))
        cc = cca([-9.0 -9.0], ϕRj, Y)
        Ay[:,j] = cc.A[:,1]
        T[:,j] = cc.T[:,1]
    end

    return ModelObj(Wx, R, Ay, T, reverse(values), par, "gKCCA")
end

"""
    whiten(x::Matrix{Float64}, d::Float64; lat=nothing): Whiten matrix
`d` (0-1) Percentage variance of components to retain. \n
`lat` Latitudinal area-weighting.
"""
function whiten(X::Matrix{Float64}, d::Float64; lat=nothing)
    m1 = zscore(X, 1)
#    lat != nothing && scale!(m1, 1.0./sqrt(cos(pi*lat/180)) )
    U,S,V = svd(m1)
    l = cumsum(S.^2)/sum(abs2, S)
    U = U[:, l.≤d]
    U ./= Statistics.std(U, dims=1)
    return U
end

"""
    CVfn(parm::T, X::T, Y::T, modelfn::Function, kerneltype::DataType; verbose::Bool=true, dcv::Int64=2) where T<:Matrix{Float64}
Cross-validation function
"""
function CVfn(parm::T, X::T, Y::T, modelfn::Function, kerneltype::DataType; verbose::Bool=true, dcv::Int64=2) where T<:Matrix{Float64}

    # free parameters
    # df = par[3] # Number of segments in PWLM
    # dcv = 2 # Number of components used in CV

    # other params
    nrg = size(parm, 1)
    pargrid = Float64[parm zeros(nrg)]
    n, p = size(X)

    # Setup
    yr = div(0:n-1,12)+1
    yrs = unique(yr)
    ta = findall(in(sample(yrs, div(length(yrs),2)+1, replace=false, ordered=true)), yr)
    tb = setdiff(1:n,ta)

    xm1, xm2 = mean(X[ta,:], dims=1), mean(X[tb,:], dims=1)
    ym1, ym2 = mean(Y[ta,:], dims=1), mean(Y[tb,:], dims=1)

    X1train, Y1train = X[ta,:].-xm1, Y[ta,:].-ym1
    X2train, Y2train = X[tb,:].-xm2, Y[tb,:].-ym2
    X1test, Y1test = X[ta,:].-xm2, Y[ta,:].-ym2
    X2test, Y2test = X[tb,:].-xm1, Y[tb,:].-ym1

    k1pars = kerneltype(X1train)
    k2pars = kerneltype(X2train)

    # Main loop

    for i in 1:nrg
        percent = round(Int, 100*i/nrg)
        if verbose
           print(string(percent,"% "))
        end
        par = parm[i,:]
        model1 = modelfn(par, X1train, Y1train, k1pars)
        model2 = modelfn(par, X2train, Y2train, k2pars)
        W1, W2 = model1.W[:,1:dcv], model2.W[:,1:dcv]
        A1, A2 = model1.A[:,1:dcv], model2.A[:,1:dcv]
        W1 .*= sign.(W1[1:1,:])
        W2 .*= sign.(W2[1:1,:])
        A1 .*= sign.(A1[1:1,:])
        A2 .*= sign.(A2[1:1,:])
        Rcv = [X2test * W1; X1test * W2]
        Tcv = [Y2test * A1; Y1test * A2]
        pargrid[i,end] = Rsq_adj(Rcv, Tcv, Int(par[3]))
    end

    imax = argmax(pargrid[:,end])
    return pargrid[imax,:]
end

"""
    Rsq_adj(Tx::T, Ty::T, df::Int) where T<:Array{Float64}:
Cross-validation metric
"""
function Rsq_adj(Tx::T, Ty::T, df::Int) where {T<:Array{Float64}}
    y = vec(Ty)
    n,p = size(Tx)
    o = ones(n*p)
    mL = [ bf(Tx[:,j], df) for j in 1:p  ]
#    mL = [ [ bf(Tx[:,j], df) o] for j in 1:p  ]
    Xs =  [ cat([1,2], mL...) o]   # Block diagonal
    px = p*df
    nx = p*n
    Σyhat = (y'*Xs / (Xs' * Xs) ) * (Xs' * y)
    Rsq = ((Σyhat) / (y'y))[1]
    return Rsq - (1-Rsq)*px/(nx-px-1)
end

end


function stable_logpdf(approx_d::Distribution; logϵ = log(ϵ_lower_bound))
    _logpdf(x) = map(s -> logsumexp(s, logϵ), logpdf(approx_d, x))
    return _logpdf
end

function stable_logpdf(approx_d::Distribution, x; logϵ = log(ϵ_lower_bound))
    return map(s -> logsumexp(s, logϵ), logpdf(approx_d, x))
end

function stable_ll(ll::Function; logϵ = log(ϵ_lower_bound))
    _ll(x) = map(s -> logsumexp(s, logϵ), ll(x))
    return _ll
end

function stable_ll(ll::Function, x; logϵ = log(ϵ_lower_bound))
    return map(s -> logsumexp(s, logϵ), ll(x))
end

function stable_negative_log_residual(ll::Function, approx_d::Distribution; logϵ = log(ϵ_lower_bound))
    _negative_log_res(x) = -ll(x) + stable_logpdf(approx_d, x; logϵ = logϵ)
    return _negative_log_res
end


module Reversible

using Yao
using NiLang
using TropicalNumbers
using StaticArrays
using LinearAlgebra: Diagonal
using LuxurySparse
using ..TropicalYao: TropicalMatrixBlock, TropicalTypes, tropicalblock, copytensor, resettensor

export Tropical, CountingTropical, TropicalTypes
export LogLikeNumbers

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

include("instructs.jl")
include("spinglass_gates.jl")

end


using Convex, SCS, ToeplitzMatrices, LinearAlgebra

function structHMirls(y,N,k,omega,sigma)
    # Line spectral estimator based on matrix completion via an iteratively
    # reweighted least-squares approach.
    # See: Kümmerle, Verdun '19: "Completion of Structured Low-Rank Matrices via
    # Iteratively Reweighted Least Squares".
    #
    # Inputs:
    # Y: Data matrix
    # N: ambient dimension \geq M
    # k: number of sinusoids
    # omega: indices of the observed entries of the data in the ambient space
    # sigma: noise variance of the additive white Gaussian noise with mean zero
    #
    # M = size(Y,1): number of measurements per snapshot
    # L = size(Y,2): number of snapshots
    #
    # Output:
    # f_hat: estimate of the frequencies of the K sinusoids in [rad]

    maxiters = 100
    #lambda = 1e-3
    lambda = sqrt(N*log(length(omega))*sigma)
    #lambda = 1e-1
    alpha = 0
    d1::Int = ceil(N/2)
    d2 = N + 1 - d1

    phi(y) = y[omega]

    function phiStar(z,N,omega)
        x = zeros(Complex{Float64},N)
        x[omega] = z
    end

    function truncSVD(A,k)
        svdA = svd(A)
        return svdA.U[:,1:k]*Diagonal(svdA.S[1:k])*svdA.Vt[1:k,:]
    end

    function kronSum(A,B)
        d1 = size(A)[1]
        d2 = size(B)[1]
        return kron(Matrix{Float64}(I, d2, d2),A) + kron(B,Matrix{Float64}(I, d1, d1))
    end

# Optimization variables
z = ComplexVariable(N)
H = ComplexVariable(d1,d2)
Hvec = ComplexVariable(d1*d2)

# initialize
z1 = phiStar(y,N,omega)
epsilon = svdvals(Hankel(z1[1:d1],z1[d1:end]))[1]
W = Matrix{Float64}(I, d1*d2, d1*d2)
error = Inf
ii = 0
z_new = Inf*ones(N)

# Define OP
constraints = [H[i,j] == z[i+j-1] for i = 1:d1, j = 1:d2]
constraints += [Hvec == vec(H)]
objective = abs(vec(H)'*W*vec(H)) + 1/(2*lambda)*sumsquares(z[omega]-y)
problem = minimize(objective, constraints)

while error > 1e-3 && ii < maxiters
ii += 1
z_old = z_new
if ii == 1
    problem1 = minimize(epsilon^2*sumsquares(z) + 1/(2*lambda)*sumsquares(z[omega]-y), constraints)
    solve!(problem1, SCS.Optimizer)
else
    objective = abs(2*Hvec'*W*Hvec) + 1/(2*lambda)*sumsquares(z[omega]-y)
    problem = minimize(objective, constraints)
    solve!(problem, SCS.Optimizer)
end
z_new = z.value
epsilon = min(epsilon, norm(z_new - z_old)+alpha^(2*k))
error = norm(z_new-z_old)/norm(z_new)
Hz = Hankel(z_new[1:d1],z_new[d1:end])
#Hz = Array{Complex{Float64},2}(undef,d1,d2)
#for i = 1:d1, j = 1:d2
#Hz[i,j] = z_new[i+j-1]
#end
W = inv(kronSum(truncSVD(Hz*Hz',k),truncSVD(Hz'*Hz,k)) + (epsilon^2)*Matrix{Float64}(I, d1*d2, d1*d2))
#println(error)
end

# Retrieve Frequencies
L = 1
clmnsY0 = setdiff(1:d2*L,(1:L).*d2)
clmnsY1 = clmnsY0.+1

Y0 = H.value[:,clmnsY0]
Y1 = H.value[:,clmnsY1]

F = svd(Y0)
p = sortperm(F.S,rev=true)[1:k]

f = eigvals(Diagonal(ones(K)./F.S[p])*F.U[:,p]'*Y1*F.V[:,p])
f_hat = (angle.(f)/(2*pi) .+ 1).%1
#println(f_hat)
#println(svdvals(H.value))
return f_hat

end


function coords(idx::Int, s::Int)::Tuple
    row = ((idx - 1) ÷ s) + 1
    col = ((idx - 1) % s) + 1
    return (row, col)
end

function coord_index(row::Int, col::Int, s::Int)::Int
    s * (row - 1) + col
end

function rot(idx::Int, s::Int, angle::Int)::Int
    row, col = coords(idx, s)
    if angle == 0
        return idx
    elseif angle == 90
        rot_row = s + 1 - col
        rot_col = row
    elseif angle == 180
        rot_row = s + 1 - row
        rot_col = s + 1 - col
    elseif angle == 270
        rot_row = col
        rot_col = s + 1 - row
    else
        throw(DomainError(angle, "Angle must be 0º, 90º, 180º, or 270º"))
    end
    coord_index(rot_row, rot_col, s)
end

function rot(combo::AbstractArray, s::Int, angle::Int)::AbstractArray
    [rot(v, s, angle) for v in combo]
end

function reflect(idx::Int, s::Int, horizontal::Bool)::Int
    row, col = coords(idx, s)
    if horizontal
        ref_row = row
        ref_col = s + 1 - col
    else
        ref_row = s + 1 - row
        ref_col = col
    end
    coord_index(ref_row, ref_col, s)
end

function reflect(combo::AbstractArray, s::Int, horizontal::Bool)::AbstractArray
    [reflect(v, s, horizontal) for v in combo]
end

function transpose(idx::Int, s::Int)::Int
    row, col = coords(idx, s)
    coord_index(col, row, s)
end

function transpose(combo::AbstractArray, s::Int)::AbstractArray
    [transpose(v, s) for v in combo]
end

"""
Determines if `left` is less than `right`--that is, `left` becomes
`right` when the vote distributions are in sorted order.
"""
function voter_dist_lt(left::AbstractArray, right::AbstractArray)
    left_s = sort(left)
    right_s = sort(left)
    for (l, r) in zip(left_s, right_s)
        if l > r
            return false
        elseif l < r
            return true
        end
    end
    false  # exactly equal
end

function combo_duplicated(combo::AbstractArray, s::Int)
    voter_dist_lt(reflect(combo, s, true), s) ||
    voter_dist_lt(reflect(combo, s, false), s) ||
    voter_dist_lt(rot(combo, s, 90), s) ||
    voter_dist_lt(rot(combo, s, 180), s) ||
    voter_dist_lt(rot(combo, s, 270), s) ||
    voter_dist_lt(transpose(combo, s), s) ||
    voter_dist_lt(rot(transpose(combo, s), s, 180), s)
end

"""
Determines the number of unique combinations resulting from
transforming a combination (reflection, rotation, transpose).
"""
function combo_count(combo::AbstractArray)
    combos = Set([
        sort(combos),
        sort(reflect(combo, s, true)),
        sort(reflect(combo, s, false)),
        sort(rot(combo, s, 90)),
        sort(rot(combo, s, 180)),
        sort(rot(combo, s, 270)),
        sort(transpose(combo, s)),
        sort(rot(transpose(combo, s), s, 180))
    ])
    length(combos)
end


using ApproxManifoldProducts
# using Distributions

##

# these do work
ptsA = [-pi/2; pi/2]

# these dont work with difftheta, but fine with `-`
ptsB = [-pi/2; pi/2+1e-10]

pA = kde!(ptsA, [0.1], (addtheta,), (difftheta,))
pB = kde!(ptsB, [0.1], (addtheta,), (difftheta,))

##

ept = ones(1,1)
ept[1,1] = pi/2

resA = evaluateDualTree(pA, ept, false, 1e-3, (addtheta,), (difftheta,))[1]

println("")
println("")

resB = evaluateDualTree(pB, ept, false, 1e-3, (addtheta,), (difftheta,))[1]

##

pA.bt.centers

pB.bt.centers

#


module ElementarySymmetricFunctions
    import SpecialFunctions.logabsbinomial

    include("util.jl")
    include("esf.jl")
    include("poisbin.jl")

    function esf(x::AbstractVector{T}; scale::Symbol=:identity, method=:recursive, k::D=2) where {T <: Real, D <: Integer}
        if scale == :identity
            if method == :recursive
                return esf_sum(x)
            elseif method == :DCFFT
                return esf_dc_fft(x,k)
            elseif method == :FFT
                return poisbin_fft(x ./ (1 .+ x)) ./ prod(1 ./(1 .+ x))
            else
                error("method $method not implemented with scaling $scale")
            end
        elseif scale == :norm
            p = x ./ (1 .+ x)
            if method == :recursive
                return poisbin_sum_taub(p)
            elseif method == :DCFFT
                return poisbin_dc_fft(p,k)
            elseif method == :FFT
                return poisbin_fft(p)
            else
                error("method $method not implemented with scaling $scale")
            end
        elseif scale == :log
            if method == :recursive
        end
    end

    export esf_sum, esf_sum_reg, esf_dc_fft, esf_dc_fft_reg, esf_sum_log
    export esf_sum_dervs_1, esf_sum_dervs_2, esf_sum_dervs_1_reg, esf_sum_dervs_2_reg
    export esf_diff, esf_diff_reg, esf_diff_updown
    export poisbin_sum_taub, poisbin_fft, poisbin_chen, poisbin_dc_fft, poisbin_fft_cf, poisbin_sum_taub_log
end


using DEBmicroTrait
using CSV, DataFrames, Statistics
using JLD
########################################
# I/O
df_isolates     = CSV.read("/Users/glmarschmann/.julia/dev/DEBmicroTrait/final_manuscript/files/isolates2traits.csv", DataFrame, missingstring="N/A")
########################################

########################################
Genome_size     = convert(Array{Float64,1}, df_isolates.Genome_size)
V_cell          = DEBmicroTrait.genome_size_to_cell_volume(Genome_size)
rrn_copies      = convert(Array{Float64,1}, df_isolates.rRNA_genes)
Min_gen_time    = df_isolates.Min_gen_time
Gram_stain      = convert(Array{String,1}, df_isolates.gram_stain)
########################################

########################################
k_M    = DEBmicroTrait.cell_volume_to_specific_maintenance_rate(V_cell, Min_gen_time, Gram_stain)
k_M_med = median(k_M)
########################################

########################################
# I/O
save("/Users/glmarschmann/.julia/dev/DEBmicroTrait/final_manuscript/files/isolates_maintenance.jld", "kM", k_M)


function Project(Fun,CG,Param)
OrdPoly=CG.OrdPoly;
nz=Param.Grid.nz;
p=zeros(CG.NumG,nz,1);
fLoc=zeros(OrdPoly+1,OrdPoly+1);
X = Param.cache.X
JC = Param.cache.JC
for iz=1:nz
  for iF=1:Param.Grid.NumFaces
    for j=1:OrdPoly+1
      for i=1:OrdPoly+1
        x=0.5*(X[i,j,1,:,iF,iz]+X[i,j,2,:,iF,iz]);
        det=JC[i,j,iF,iz];
        fLoc[i,j]=Fun(x,Param)*det;
      end
    end
    p[CG.Faces[iF].Glob,iz]=p[CG.Faces[iF].Glob,iz]+reshape(fLoc,(OrdPoly+1)*(OrdPoly+1),1);
  end
end
p=p./CG.M;
return p
end


abstract type AbstractMobilityModel

end

abstract type AbstractSaturationVelocity
end

struct ConstantVsat <: AbstractSaturationVelocity
    vsat::Number
end

Vsat(v::ConstantVsat,T) = v.vsat

struct SwatzVsat <: AbstractSaturationVelocity
    α::Number
    θ::Number
    T₀::Number
end

Vsat(v::SwatzVsat,T) = v.α/(1+v.θ*exp(T/v.T₀))

struct ConstantMobility <: AbstractMobilityModel
    μ₀::Number
    sat::AbstractSaturationVelocity
    β₀::Number
    βexp::Number
    T₀::Number
end

β(C::ConstantMobility,T) = C.β₀*(T/C.T₀)^C.βexp
Vsat(C::ConstantMobility,T) = Vsat(C.sat,T)
Gsat(C::ConstantMobility,T,ξ) = (1/(1+(C.μ₀*ξ/Vsat(C,T))^β(C,T)))^(1/β(C,T))
μ(C::ConstantMobility,S::SemiconductorData) = C.μ₀*Gsat(C,Temperature(S),ξ(S,ψ(S)))

struct CaugheyThomasMobility <: AbstractMobilityModel
    μ₁::Number
    μ₂::Number
    α::Number
    β::Number
    γ::Number
    δ::Number
    Ncrit::Number
    sat::AbstractSaturationVelocity
    β₀::Number
    βexp::Number
    T₀::Number
end

μ₀(C::CaugheyThomasMobility,T,N) = C.μ₁*(T/C.T₀)^C.α + (C.μ₂*(T/C.T₀)^C.β - C.μ₁*(T/C.T₀)^C.α)/(1+(T/C.T₀)^C.γ*(N/C.Ncrit)^C.δ)
β(C::CaugheyThomasMobility,T) = C.β₀*(T/300.0)^C.βexp
Vsat(C::CaugheyThomasMobility,T) = Vsat(C.sat,T)
Gsat(C::CaugheyThomasMobility,T,N,ξ) = (1/(1+(μ₀(C,T,N)*ξ/Vsat(C,T))^β(C,T)))^(1/β(C,T))
μ(muM::CaugheyThomasMobility,S::SemiconductorData) = μ₀(muM,Temperature(S),abs(C(S)))*Gsat(muM,Temperature(S),abs(C(S)),abs(ξ(S,ψ(S))))
μ(muM::CaugheyThomasMobility,S::SemiconductorData,ψ) = μ₀(muM,Temperature(S),abs(C(S)))*Gsat(muM,Temperature(S),abs(C(S)),abs(ξ(S,ψ)))
μ(muM::CaugheyThomasMobility,S::SemiconductorData,ψ,ϕ) = μ₀(muM,Temperature(S),abs(C(S)))*Gsat(muM,Temperature(S),abs(C(S)),abs(ξ(S,ψ,ϕ)))

μₙCT() = CaugheyThomasMobility(55.24,1429.23,0.0,-2.3,-3.8,0.73,1.072e17,SwatzVsat(2.4e7,0.8,600),2.0,0.01,300.0)
μₚCT() = CaugheyThomasMobility(49.7,479.37,0,-2.2,-3.7,0.7,1.606e17,SwatzVsat(2.4e7,0.8,600),0.6,0.01,300.0)

struct PSPMobility <: AbstractMobilityModel
    FACTUO::Number
    BETN::Number
    Cox::Number
    T₀::Number
    STBET::Number
    THEMU::Number
    STTHEMU::Number
    MUE::Number
    STMUE::Number
    XCOX::Number
    STXCOX::Number
    CS::Number
    STCS::Number
    THCS::Number
    STTHECS::Number
    EFF0::Number
    FETA::Number
    ETAMUSC::Number
    εsi::Number
end

β(PSPm::PSPMobility,T) = PSPm.FACTUO*PSPm.BETN*PSPm.Cox*(T/PSPm.T₀)^PSPm.STBET
θμ(PSPm::PSPMobility,T) = PSPm.THEMU*(T/PSPm.T₀)^PSPm.STTHEMU
μe(PSPm::PSPMobility,T) = PSPm.MUE*(T/PSPm.T₀)^PSPm.STMUE
xcor(PSPm::PSPMobility,T) = PSPm.XCOX*(T/PSPm.T₀)^PSPm.STXCOX
Cs(PSPm::PSPMobility,T) = PSPm.CS*(T/PSPm.T₀)^PSPm.STCS
θCS(PSPm::PSPMobility,T) = PSPm.THCS*(T/PSPm.T₀)^PSPm.STTHECS
ξff0(PSPm::PSPMobility) = PSPm.EFF0*PSPm.Cox/PSPm.εsi
ημ(PSPm::PSPMobility) = PSPm.ETAMUSC*PSPm.FETA
ξeff(PSPm::PSPMobility,ξ) = ξeff0(PSPm)*ξ

#Gmob(PSPm::PSPMobility,ξ,T) = 1+( μe(PSPm,T)*ξeff(PSPm,ξ) )^θμ(PSPm,T) + Cs(PSPm,T) *

struct CVTMobilityATLAS <: AbstractMobilityModel
    B::Number
    C::Number
    D::Number
    K::Number
    τ::Number
    γ::Number
    μ0::Number
    μ1::Number
    μmax::Number
    CR::Number
    CS::Number
    α::Number
    β::Number
    PC::Number
    δ::Number
    E::Number
    A::Number
    AL::Number
    η::Number
end

μAC(ξ,C,T,m::CVTMobilityATLAS) = m.B/(ξ^m.E)+(m.C*C^m.τ)/(T*ξ^m.D)
μsr(ξ,n,p,C,m::CVTMobilityATLAS) = m.δ/(ξ^m.K) #m.δ/(ξ^γ(m,n,p,C))
γ(m::CVTMobilityATLAS,n,p,C) = m.A+m.AL*(n+p)/abs(C^m.η)
μb(C,T,m::CVTMobilityATLAS) = m.μ0*exp(-m.PC/C)+(m.μmax*(T/300.0)^(-m.γ)-m.μ0)/(1+(C/m.CR)^m.α)-m.μ1/(1+(m.CS/C)^m.β)
μ(m::CVTMobilityATLAS,ξ,n,p,C,T) = 1.0/(1.0/μAC(ξ,C,T,m)+1.0/μsr(ξ,n,p,C,m)+1.0/μb(C,T,m))
μ(m::CVTMobilityATLAS,S::SemiconductorData,ψ,ϕ_n=0.0,ϕ_p=0.0) = μ(m,abs(ξ(S,ψ)),n(S,ψ-ϕ_n),p(S,ψ-ϕ_p),abs(C(S)),Temperature(S))
μcvn() = CVTMobilityATLAS(4.75e7,1.74e5,0.333,2.0,0.125,2.5,52.2,43.4,1417.0,9.68e16,3.43e20,0.68,2.0,0.0,5.82e14,1.0,2.58,5.85e-21,0.0767)

struct CVTMobility <: AbstractMobilityModel
    B
    C
    μ0
    μmax
    μ1
    γ
    CR
    CS
    α
    β
    δ
end

μAC(ξ,C,T,m::CVTMobility) = m.B*T/(ξ)+m.C*(1/ξ^(1/3.0))/(1/T)
μsr(ξ,n,p,C,m::CVTMobility) = m.δ/(ξ^2.0) #m.δ/(ξ^γ(m,n,p,C))
γ(m::CVTMobility,n,p,C) = m.A+m.AL*(n+p)/abs(C^m.η)
μb(C,T,m::CVTMobility) = m.μ0+(m.μmax*(T/300.0)^(-m.γ)-m.μ0)/(1+(C/m.CR)^m.α)-m.μ1/(1+(m.CS/C)^m.β)
μ(m::CVTMobility,ξ,n,p,C,T) = 1.0/(1.0/μAC(ξ,C,T,m)+1.0/μsr(ξ,n,p,C,m)+1.0/μb(C,T,m))
μ(m::CVTMobility,S::SemiconductorData,ψ,ϕ_n=0.0,ϕ_p=0.0) = μ(m,abs(ξ(S,ψ)),n(S,ψ-ϕ_n),p(S,ψ-ϕ_p),abs(C(S)),Temperature(S))
μcvn(C) = CVTMobility(4.75e7,1.75e5*C,52.2,1417,23.4,2.5,96.8e16,3.43e20,0.68,2.0,5.82e14)

struct TaschMobility <: AbstractMobilityModel
    RN#2
    BETAS#2
    MUB#1150
    TMUB#2.5
    D#3.2e-9
    P1#0.09
    B1#1.75
    P2#4.53e-8
    PE#3.17e-7
    B2#-0.25
    Z11#0.0388
    Z22#1.73e-5
    ESR#2.449e7
    BETA#2
    N2#1.1e21
    N1#2e19
    ALPHA#2.0
end
Γ(m::TaschMobility,Ep) = μeff(m)/(1+(μeff(m)*Ep/m.VSAT)^m.BETAS)^(1/m.BETAS)
μph(m::TaschMobility,T,ξ,n) = 1.0/(m.MUBN*(T/300.0)^(-m.TMUBN))+1.0/(Z(m,T,ξ)/m.DN*Y(m,T,n)*(T/300.0)^0.5)
Z(m::TaschMobility,T,ξ) = m.Z11*(T/300.0)/ξ+m.Z22*ξ^(-1/3.0)
Eff(m::TaschMobility,ξ) = ξ+(m.RN-1)*E0(m)/m.RN
Y(m::TaschMobility,T,n)= m.P1*(T/300.0)^B1 + m.p2*n+m.B2*0.0
μsr(m::TaschMobility,ξ) = (m.ESR/ξ)^m.BETA
μc(m::TaschMobility,C) = m.N2N*(T/300)^1.5/(C*log(1+γB(m,T))-γB(m,T)/(1+γB(m,T)))
γB(m::TaschMobility,T) = m.N1*(T/300.0)^m.ALPHA


# LLᵀ (Choleski) factorization for symmetric systems


# type StepReg{W<:Weight} <: OnlineStat
#     β0::Float64
#     β::VecF
#     intercept::Bool
#     c::CovMatrix{W}
#     s::MatF         # placeholder for swept version of c
#     set::IntSet
#     weight::W
#     n::Int
#     nups::Int
# end
# function StepReg(p::Integer, wgt::Weight = EqualWeight(); intercept::Bool = true)
#     StepReg(
#         0.0, zeros(p), intercept, CovMatrix(p), zeros(p, p), IntSet([]),
#         wgt, 0, 0
#     )
# end
# function StepReg(x::AMat, y::AVec, wgt::Weight = EqualWeight())
#     o = StepReg(size(x, 2), wgt)
#     fit!(o, x, y)
#     o
# end
#
# function StatsBase.fit!(o::StepReg, x::AVec, y::Real)
#     fit!(o.c, vcat(x, y))
# end
# function value(o::StepReg)
#     o.intercept ? vcat(o.β0, o.β) : o.β
# end
#
# ## TESTING
# n, p = 1000, 10
# x = randn(n, p)
# β = collect(1.:p)
# y = x * β + randn(n)
# o = StepReg(x, y)
# display(o)
#
# # # Online Stepwise Regression
# # #
# # # With each updated batch, there is the possibility of one variable
# # # entering or leaving the model based on Mallow's Cp
# #
# # #-------------------------------------------------------# Type and Constructors
# # """
# # Online stepwise regression.
# #
# # At each update, there is the possible of one variable entering or leaving the model.
# # Performs best using `fitbatch!()` with large batches.
# # """
# # type StepwiseReg{W <: Weight} <: OnlineStat
# #     C::CovarianceMatrix{W}  # Cov([X y])
# #     s::MatF                 # "Swept" version of [X y]' [X y] / n
# #     set::IntSet             # set of coefficients included in the model
# #     n::Int
# # end
# #
# # function StepwiseReg(x::AMatF, y::AVecF, wgt::Weight = EqualWeight())
# #     n, p = size(x)
# #     o = StepwiseReg(p, wgt)
# #     fitbatch!(o, x, y)
# #     o
# # end
# #
# # function StepwiseReg(p::Integer, wgt::Weighting = default(Weighting))
# #     c = CovarianceMatrix(p + 1, wgt)
# #     StepwiseReg(c, zeros(p + 1, p + 1), IntSet([]), 0)
# # end
# #
# #
# # #-----------------------------------------------------------------------# state
# # statenames(o::StepwiseReg) = [:β, :nobs]
# # state(o::StepwiseReg) = Any[coef(o), nobs(o)]
# #
# # function StatsBase.coef(o::StepwiseReg)
# #     β = vec(o.s[end, 1:end - 1])
# #     for i in setdiff(1:length(β), o.set)
# #         β[i] = 0.
# #     end
# #     β
# # end
# #
# # #----------------------------------------------------------------------# fit!
# # function fit!(o::StepwiseReg, x::AVecF, y::Float64)
# #     fitbatch!(o, x', collect(y))
# # end
# #
# # function fitbatch!(o::StepwiseReg, x::AMatF, y::AVecF)
# #     n, p = size(x)
# #     o.n += n
# #     fitbatch!(o.C, hcat(x, y))
# #     copy!(o.s, o.C.A)
# #
# #     # get average squared error using all predictors
# #     sweep!(o.s, 1:p)
# #     ase = o.s[end, end]
# #
# #     copy!(o.s, o.C.A)
# #     sweep!(o.s, collect(o.set))
# #
# #     # Find best index to add/remove
# #     s = 1
# #     val = o.s[end, 1] ^ 2 / (o.s[1, 1] * ase) * (o.n - in(1, o.set)) + 2.0 - in(1, o.set)
# #     for i in 2:p
# #         newval = o.s[end, i] ^ 2 / (o.s[i, i] * ase) * (o.n - in(i, o.set)) - 2.0 * in(i, o.set)
# #         if newval > val
# #             val = newval
# #             s = i
# #         end
# #     end
# #
# #     if s in o.set
# #         delete!(o.set, s)
# #     else
# #         push!(o.set, s)
# #     end
# #
# #     DEBUG("Active set: ", o.set)
# # end
# #
# #
# # ########## TEST code
# # # log_severity!(DebugSeverity)
# # #
# # # n,p = 10_000, 10
# # # x = randn(n,p)
# # # β = collect(1.:p)
# # # β[5] = 0.
# # # y = x*β + randn(n)
# # #
# # # o = StepwiseReg(p)
# # # @time fit!(o, x, y, 500)
# # # print(coef(o))


using HOODESolver
using DifferentialEquations
using LinearAlgebra
using Plots
using Random

function getmindif(tab::Vector{Vector{BigFloat}})
    nmmin = Inf
    ret = (0, 0)
    for i=1:size(tab,1)
        for j=(i+1):size(tab,1)
            nm = norm(tab[i]-tab[j], Inf)
            if nm < nmmin
                nmmin = nm
                ret = i, j
            end
        end
    end
    return ret, nmmin
end
function fct(u, p, t)
    s1, c1 = sincos(u[1]/2)
    s2, c2 = sincos(u[2])
    s3, c3 = sincos(u[3])
    return [0, 0, u[6], c1*s2*s3/2, s1*c2*s3, s1*s2*c3]
end

function fctmain(n_tau, prec)
    setprecision(512)
    Random.seed!(561909)
    u0 = 2rand(BigFloat,6)-ones(BigFloat,6)
    setprecision(prec)
    u0 = BigFloat.(u0)
    println("u0=$u0")
    t_max = big"1.0"
    epsilon=big"1e-7"
    println("epsilon=$epsilon")
    nbmaxtest=15
    ordmax=9
    debord=3
    pasord=1
    y = ones(Float64, nbmaxtest, div(ordmax-debord,pasord)+1 )
    res_gen = Array{ Array{BigFloat,1}, 2}(undef, nbmaxtest, div(ordmax-debord,pasord)+1 )
    x=zeros(Float64,nbmaxtest)
    A = [0 0 0 1 0 0; 0 0 0 0 1 0;0 0 0 0 0 0; 0 0 0 0 1 0; 0 0 0 -1 0 0; 0 0 0 0 0 0]
    t_0=big"0.0"
    t_max=big"1.0"
    prob = HOODEProblem(fct, u0, (t_0, t_max), missing, A, epsilon)
    nm = NaN
    ordmax += 1
    ordprep=undef
    nb = 10*2^(nbmaxtest)
    solref=undef
    while isnan(nm)
        ordmax -= 1
        @time sol = solve(prob, nb_t=nb, order=ordmax, getprecision=false, nb_tau=n_tau, dense=false)
        solref = sol[end]
	    nm = norm(solref, Inf)
    end

    tabsol = Array{Array{BigFloat,1},1}(undef,1)

    tabsol[1] = solref

    println("ordmax=$ordmax")

    solref_gen = solref
    indref = 1

    ind=1
    for order=debord:pasord:ordmax
        ordprep=order+2
        println("eps=$epsilon solRef=$solref order=$order")
        nb = 10
        indc = 1
        labels=Array{String,2}(undef, 1, order-debord+1)
        resnorm=0
        resnormprec=1
        sol =undef
        par_u0=missing
        println("preparation ordre $order + 2")
        while indc <= nbmaxtest
            @time solall = solve(prob, nb_t=nb, order=order, getprecision=false, nb_tau=n_tau, par_u0=par_u0, dense=false)
            par_u0 = solall.par_u0
            sol = solall[end]
            push!(tabsol, sol)
            res_gen[indc,ind] = sol
            (a, b), nm = getmindif(tabsol)
            if a != indref
                println("New solref !!!! a=$a, b=$b nm=$nm")
                indref = a
                solref = tabsol[a]
                for i=1:ind
                    borne = (i <ind) ? size(res_gen,1) : indc
                    for j = 1:borne
                        nm2 = min( norm(res_gen[j,i] - solref, Inf), 1.1)
                        y[j,i] = nm2 == 0 ? nm : nm2
                   end
                end
            else
                diff=solref-sol
                y[indc,ind] = min(norm(diff,Inf), 1.1)
            end
            x[indc] = 1.0/nb
            println("result=$y")
            println("log2(y)=$(log2.(y))")
            nb *= 2
            indc += 1
        end
        for i=debord:pasord:order
            labels[1,(i-debord)÷pasord+1] = " eps,order=$(convert(Float32,epsilon)),$i "
        end
        p=Plots.plot(
                        x,
                        view(y,:,1:ind),
                        xlabel="delta t",
                        xaxis=:log,
                        ylabel="error",
                        yaxis=:log,
                        legend=:topleft,
                        label=labels,
                        marker=2
                    )
        prec_v = precision(BigFloat)
        eps_v = convert(Float32,epsilon)
        Plots.savefig(p,"out/r4_$(prec_v)_$(eps_v)_$(order)_$(n_tau)_particle.pdf")
        if resnorm > resnormprec
            break
        end
        resnormprec = resnorm
        ind+= 1
    end
end
fctmain(32,512)


using CayleyDickson
using Base.Test: @test, @test_throws

@test begin
    a = CayleyDickson.randomBigFloat()
    isreal(Exo4Real(Exo3Real(a)))
end

@test begin
    a = 1
    !isreal(Exo4Real(Exo3Real(a, a, a, a)))
end

@test begin
    z = random(Exo4Real{Int})
    z == +(z)
end

@test begin
    a = CayleyDickson.randomBigFloat()
    real(Exo4Real(Exo3Real(a))) == a
end

@test begin
    io = IOBuffer()
    show(io, Exo4Real(Exo3Real(1.0, 2, 3, 4, 5, 6, 7, 8)))
    l = String(take!(io))
    r = "[1: 1.0, a: 2.0, b: 3.0, ab: 4.0, c: 5.0, ac: 6.0, bc: 7.0, (ab)c: 8.0, d: 0.0, ad: 0.0, bd: 0.0, (ab)d: 0.0, cd: 0.0, (ac)d: 0.0, (bc)d: 0.0, ((ab)c)d: 0.0]"
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigInt()
    c = CayleyDickson.randomBigInt()
    d = CayleyDickson.randomBigInt()
    l = Exo4Real(Exo3Real(a, b, c, d))
    r = Exo4Real(Exo2Real(a, b, c, d))
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigInt()
    l = Exo4Real(Exo1Real(a, b))
    r = Exo4Real(Exo3Real(a, b))
    l == r
end

@test begin
    x = random(Exo4Real{BigFloat})
    y = random(Exo4Real{BigFloat})
    x != y
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo4Real{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = random(Exo3Real{BigInt})
    z = random(Exo4Real{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo4Real{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    a = random(Exo3Real{BigInt})
    z = random(Exo4Real{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    l = x * y
    r = y * x
    l != r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    z = random(Exo4Real{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    z = random(Exo4Real{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l != r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    l = alternatorL(x, y)
    r = zero(Exo4Real{BigInt})
    l != r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    l = alternatorR(x, y)
    r = zero(Exo4Real{BigInt})
    l != r
end

@test begin
    x = random(Exo4Real{BigInt})
    conj(conj(x)) == x
end

@test begin
    x = random(Exo4Real{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    l = conj(x * y)
    r = conj(y) * conj(x)
    l == r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    l = conj(x + y)
    r = conj(x) + conj(y)
    l == r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    l = conj(x - y)
    r = conj(x) - conj(y)
    l == r
end

# @test begin
#     x = random(Exo4Real{Rational{BigInt}})
#     y = random(Exo4Real{Rational{BigInt}})
#     z = x * y
#     z / y == x
# end

# @test begin
#     x = random(Exo4Real{Rational{BigInt}})
#     y = random(Exo4Real{Rational{BigInt}})
#     z = x * y
#     x \ z == y
# end

# @test begin
#     x = random(Exo4Real{Rational{BigInt}})
#     y = random(Exo4Real{Rational{BigInt}})
#     l = inv(x * y)
#     r = inv(y) * inv(x)
#     l == r
# end

# @test begin
#     x = random(Exo4Real{Rational{BigInt}})
#     y = random(Exo4Real{Rational{BigInt}})
#     l = inv(x / y)
#     r = y / x
#     l == r
# end

# @test begin
#     x = random(Exo4Real{Rational{BigInt}})
#     y = random(Exo4Real{Rational{BigInt}})
#     l = inv(x \ y)
#     r = y \ x
#     l == r
# end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    z = random(Exo4Real{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo4Real{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    a = random(Exo3Real{BigInt})
    x = random(Exo4Real{BigInt})
    l = commutator(a, x)
    r = zero(Exo4Real{BigInt})
    l != r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    z = random(Exo4Real{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    z = random(Exo4Real{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    z = random(Exo4Real{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(Exo4Real{BigInt})
    quadrance(x) > big(0) || quadrance(x) == big(0)
end

@test begin
    x = random(Exo4Real{BigInt})
    y = random(Exo4Real{BigInt})
    l = quadrance(x * y)
    r = quadrance(x) * quadrance(y)
    l == r
end

@test_throws ErrorException begin
    inv(Exo4Real(Exo3Real(0)))
end

@test_throws ErrorException begin
    random(Exo4Real{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(Exo4Real{Int})
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo4Real{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo4Real{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(Exo4Real{BigInt})
    z == z + zero(Exo4Real{BigInt})
end

@test begin
    z = random(Exo4Real{BigInt})
    z == z * one(Exo4Real{BigInt})
end

@test begin
    z = random(Exo4Real{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(Exo4Real{BigInt})
    z == z * one(z)
end

@test begin
    x = random(Exo4Real{BigInt})
    cloak(cloak(x)) == x
end

@test begin
    x = random(Exo4Real{BigInt})
    dagger(dagger(x)) == x
end

@test begin
    x = random(Exo4Real{BigInt})
    star(star(x)) == dagger(x)
end

@test begin
    x = random(Exo4Real{Rational{BigInt}})
    selfstar(selfstar(x)) == selfstar(x)
end

@test begin
    x = random(Exo4Real{Rational{BigInt}})
    anti_selfstar(anti_selfstar(x)) == anti_selfstar(x)
end

@test begin
    l = Exo4Real(1,2,3,4)
    r = Exo4Real(Exo2Real(1,2,3,4))
    l == r
end

@test begin
    l = Exo4Real(1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17)
    r = Exo4Real(1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16)
    l == r
end


using ArrayInterfaceBlockBandedMatrices
using ArrayInterfaceCore
using BlockBandedMatrices
using Test

BB=BlockBandedMatrix(Ones(10,10),[1,2,3,4],[4,3,2,1],(1,0))
BB[Block(1,1)].=[1 2 3 4]
BB[Block(2,1)].=[5 6 7 8;9 10 11 12]
rowind,colind=ArrayInterfaceCore.findstructralnz(BB)
@test [BB[rowind[i],colind[i]] for i in 1:length(rowind)]==
    [1,5,9,2,6,10,3,7,11,4,8,12,
    1,1,1,1,1,1,1,1,1,1,1,1,1,1,
    1,1,1,1,1,1,1,1,1,1,1,1,1,1,
    1,1,1,1,1]
@test ArrayInterfaceCore.isstructured(typeof(BB))
@test ArrayInterfaceCore.has_sparsestruct(typeof(BB))
@test ArrayInterfaceCore.fast_matrix_colors(typeof(BB))

dense=collect(Ones(8,8))
for i in 1:8
    dense[:,i].=[1,2,3,4,5,6,7,8]
end
BBB=BandedBlockBandedMatrix(dense, [4, 4] ,[4, 4], (1, 1), (1, 1))
rowind,colind=ArrayInterfaceCore.findstructralnz(BBB)
@test [BBB[rowind[i],colind[i]] for i in 1:length(rowind)]==
    [1,2,3,1,2,3,4,2,3,4,5,6,7,5,6,7,8,6,7,8,
     1,2,3,1,2,3,4,2,3,4,5,6,7,5,6,7,8,6,7,8]
@test ArrayInterfaceCore.isstructured(typeof(BBB))
@test ArrayInterfaceCore.has_sparsestruct(typeof(BBB))
@test ArrayInterfaceCore.fast_matrix_colors(typeof(BBB))


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    RotationRule(order, motion, radian, main, extrinsic)

Creates a rule for the ellipsoidal rotation

## Parameters

    * order     - sequence of three axes by which the rotations are made; e.g. `:ZXZ`
    * motion    - inform for each of the three rotations if it is clockwise (`:CW`) or
                  counterclockwise (`:CCW`). Right-hand rule: motion defined looking
                  towards the negative direction of the axis
    * radian    - `true` if the input angles are in radians or `false` if in degrees
    * main      - inform if the main semiaxis is `:x` or `:y`
    * extrinsic - `true` if rotation is extrinsic or `false` if it is intrinsic

## Example

Rotation rule to reproduce GSLIB rotation:

```julia
julia> rule = RotationRule(:ZXY,[:CW,:CCW,:CCW],false,:y,false)
```
"""
struct RotationRule
  order::Symbol
  motion::Vector{Symbol}
  radian::Bool
  main::Symbol
  extrinsic::Bool
end

# list of available rotation rules
rules = Dict(
  :TaitBryanExtr => RotationRule(:ZXY,[:CCW,:CCW,:CCW],true,:x,true),
  :TaitBryanIntr => RotationRule(:ZXY,[:CCW,:CCW,:CCW],true,:x,false),
  :EulerExtr     => RotationRule(:ZXZ,[:CCW,:CCW,:CCW],true,:x,true),
  :EulerIntr     => RotationRule(:ZXZ,[:CCW,:CCW,:CCW],true,:x,false),
  :GSLIB         => RotationRule(:ZXY,[:CW,:CCW,:CCW],false,:y,false),
  :Leapfrog      => RotationRule(:ZXZ,[:CW,:CW,:CW],false,:x,false),
  :Datamine      => RotationRule(:ZXZ,[:CW,:CW,:CW],false,:x,false)
)

"""
    aniso2distance(semiaxes, angles; convention=:TaitBryanExtr)

A distance defined by an ellipsoid with given `semiaxes` and rotation `angles`.

- For 2D ellipsoids, there are two semiaxes and one rotation angle.
- For 3D ellipsoids, there are three semiaxes and three rotation angles.

## Conventions

Different rotation conventions can be passed via the `convention` option:

- `:TaitBryanExtr` => Extrinsic right-handed rotation by the ZXY axes
- `:TaitBryanIntr` => Intrinsic right-handed rotation by the ZXY axes
- `:EulerExtr`     => Extrinsic right-handed rotation by the ZXZ axes
- `:EulerIntr`     => Intrinsic right-handed rotation by the ZXZ axes
- `:GSLIB`         => GSLIB software rotation convention
- `:Leapfrog`      => Leapfrog software rotation convention
- `:Datamine`      => Datamine software rotation convention (fixed to ZXZ axes)

Tait-Bryan and Euler conventions expect angles in radians.
The other conventions expect them in degrees.

Default convention is `:TaitBryanExtr`.

## Examples

2D ellipsoid making 45ᵒ with the horizontal axis:

```julia
julia> aniso2distance([1.0,0.5], [π/2])
```

3D ellipsoid rotated by 45ᵒ in the xy plane:

```julia
julia> aniso2distance([1.0,0.5,0.5], [π/2,0.0,0.0])
```
"""
function aniso2distance(semiaxes::AbstractVector, angles::AbstractVector;
  convention::Symbol=:TaitBryanExtr)

  N = length(semiaxes)
  @assert all(semiaxes .> 0) "semiaxes must be positive"
  @assert N ∈ [2,3] "dimension must be either 2 or 3"

  rule = rules[convention]

  # invert x and y if necessary
  if rule.main == :y
     semiaxes[1], semiaxes[2] = semiaxes[2], semiaxes[1]
  end

  # scaling matrix
  Λ = Diagonal(SVector{N}(one(eltype(semiaxes))./semiaxes.^2))

  # convert to radian and invert sign if necessary
  !rule.radian && (angles = deg2rad.(angles))
  _0 = zero(eltype(angles))
  N == 2 && (angles = [angles[1], _0, _0])
  intr = @. (rule.motion == :CW) & !rule.extrinsic
  extr = @. (rule.motion == :CCW) & rule.extrinsic
  angles[intr .| extr] *= -1

  # rotation matrix
  P = angle_to_dcm(angles..., rule.order)[SOneTo(N),SOneTo(N)]

  # ellipsoid matrix
  Q = rule.extrinsic ? P*Λ*P' : P'*Λ*P

  Mahalanobis(Q)
end