using ModelingToolkit, StaticArrays, LinearAlgebra
using OrdinaryDiffEq, Sundials
using DiffEqBase, SparseArrays
using StaticArrays
using Test

using ModelingToolkit: value

# Define some variables
@parameters t σ ρ β
@variables x(t) y(t) z(t)
D = Differential(t)

# Define a differential equation
eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]

ModelingToolkit.toexpr.(eqs)[1]
@named de = ODESystem(eqs; defaults=Dict(x => 1))
@test eval(toexpr(de)) == de

generate_function(de)

function test_diffeq_inference(name, sys, iv, dvs, ps)
    @testset "ODESystem construction: $name" begin
        @test isequal(independent_variables(sys)[1],  value(iv))
        @test length(independent_variables(sys))==1
        @test isempty(setdiff(Set(states(sys)), Set(value.(dvs))))
        @test isempty(setdiff(Set(parameters(sys)), Set(value.(ps))))
    end
end

test_diffeq_inference("standard", de, t, [x, y, z], [ρ, σ, β])
generate_function(de, [x,y,z], [σ,ρ,β])
jac_expr = generate_jacobian(de)
jac = calculate_jacobian(de)
jacfun = eval(jac_expr[2])

for f in [
    ODEFunction(de, [x,y,z], [σ,ρ,β], tgrad = true, jac = true),
    eval(ODEFunctionExpr(de, [x,y,z], [σ,ρ,β], tgrad = true, jac = true)),
]
    # iip
    du = zeros(3)
    u  = collect(1:3)
    p  = collect(4:6)
    f.f(du, u, p, 0.1)
    @test du == [4, 0, -16]

    # oop
    du = @SArray zeros(3)
    u  = SVector(1:3...)
    p  = SVector(4:6...)
    @test f.f(u, p, 0.1) === @SArray [4, 0, -16]

    # iip vs oop
    du = zeros(3)
    g = similar(du)
    J = zeros(3, 3)
    u  = collect(1:3)
    p  = collect(4:6)
    f.f(du, u, p, 0.1)
    @test du == f(u, p, 0.1)
    f.tgrad(g, u, p, t)
    @test g == f.tgrad(u, p, t)
    f.jac(J, u, p, t)
    @test J == f.jac(u, p, t)
end

eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y*t,
       D(z) ~ x*y - β*z]
@named de = ODESystem(eqs)
ModelingToolkit.calculate_tgrad(de)

tgrad_oop, tgrad_iip = eval.(ModelingToolkit.generate_tgrad(de))

u  = SVector(1:3...)
p  = SVector(4:6...)
@test tgrad_oop(u,p,t) == [0.0,-u[2],0.0]
du = zeros(3)
tgrad_iip(du,u,p,t)
@test du == [0.0,-u[2],0.0]

@parameters σ′(t-1)
eqs = [D(x) ~ σ′*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]
@named de = ODESystem(eqs)
test_diffeq_inference("global iv-varying", de, t, (x, y, z), (σ′, ρ, β))

f = eval(generate_function(de, [x,y,z], [σ′,ρ,β])[2])
du = [0.0,0.0,0.0]
f(du, [1.0,2.0,3.0], [x->x+7,2,3], 5.0)
@test du ≈ [11, -3, -7]

@parameters σ(..)
eqs = [D(x) ~ σ(t-1)*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]
@named de = ODESystem(eqs)
test_diffeq_inference("single internal iv-varying", de, t, (x, y, z), (σ(t-1), ρ, β))
f = eval(generate_function(de, [x,y,z], [σ,ρ,β])[2])
du = [0.0,0.0,0.0]
f(du, [1.0,2.0,3.0], [x->x+7,2,3], 5.0)
@test du ≈ [11, -3, -7]

eqs = [D(x) ~ x + 10σ(t-1) + 100σ(t-2) + 1000σ(t^2)]
@named de = ODESystem(eqs)
test_diffeq_inference("many internal iv-varying", de, t, (x,), (σ(t-2),σ(t^2), σ(t-1)))
f = eval(generate_function(de, [x], [σ])[2])
du = [0.0]
f(du, [1.0], [t -> t + 2], 5.0)
@test du ≈ [27561]

# Conversion to first-order ODEs #17
D3 = Differential(t)^3
D2 = Differential(t)^2
@variables u(t) uˍtt(t) uˍt(t) xˍt(t)
eqs = [D3(u) ~ 2(D2(u)) + D(u) + D(x) + 1
       D2(x) ~ D(x) + 2]
@named de = ODESystem(eqs)
de1 = ode_order_lowering(de)
lowered_eqs = [D(uˍtt) ~ 2uˍtt + uˍt + xˍt + 1
               D(xˍt)  ~ xˍt + 2
               D(uˍt)  ~ uˍtt
               D(u)    ~ uˍt
               D(x)    ~ xˍt]

#@test de1 == ODESystem(lowered_eqs)

# issue #219
@test all(isequal.([ModelingToolkit.var_from_nested_derivative(eq.lhs)[1] for eq in equations(de1)], states(@named lowered = ODESystem(lowered_eqs))))

test_diffeq_inference("first-order transform", de1, t, [uˍtt, xˍt, uˍt, u, x], [])
du = zeros(5)
ODEFunction(de1, [uˍtt, xˍt, uˍt, u, x], [])(du, ones(5), nothing, 0.1)
@test du == [5.0, 3.0, 1.0, 1.0, 1.0]

# Internal calculations
@parameters σ
a = y - x
eqs = [D(x) ~ σ*a,
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]
@named de = ODESystem(eqs)
generate_function(de, [x,y,z], [σ,ρ,β])
jac = calculate_jacobian(de)
@test ModelingToolkit.jacobian_sparsity(de).colptr == sparse(jac).colptr
@test ModelingToolkit.jacobian_sparsity(de).rowval == sparse(jac).rowval

f = ODEFunction(de, [x,y,z], [σ,ρ,β])

D = Differential(t)
@parameters A B C
_x = y / C
eqs = [D(x) ~ -A*x,
       D(y) ~ A*x - B*_x]
@named de = ODESystem(eqs)
@test begin
    local f, du
    f = eval(generate_function(de, [x,y], [A,B,C])[2])
    du = [0.0,0.0]
    f(du, [1.0,2.0], [1,2,3], 0.0)
    du ≈ [-1, -1/3]
    f = eval(generate_function(de, [x,y], [A,B,C])[1])
    du ≈ f([1.0,2.0], [1,2,3], 0.0)
end

function lotka(u,p,t)
  x = u[1]
  y = u[2]
  [p[1]*x - p[2]*x*y,
  -p[3]*y + p[4]*x*y]
end

prob = ODEProblem(ODEFunction{false}(lotka),[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0])
de = modelingtoolkitize(prob)
ODEFunction(de)(similar(prob.u0), prob.u0, prob.p, 0.1)

function lotka(du,u,p,t)
  x = u[1]
  y = u[2]
  du[1] = p[1]*x - p[2]*x*y
  du[2] = -p[3]*y + p[4]*x*y
end

prob = ODEProblem(lotka,[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0])

de = modelingtoolkitize(prob)
ODEFunction(de)(similar(prob.u0), prob.u0, prob.p, 0.1)

# automatic state detection for DAEs
@parameters t k₁ k₂ k₃
@variables y₁(t) y₂(t) y₃(t)
D = Differential(t)
# reorder the system just to be a little spicier
eqs = [D(y₁) ~ -k₁*y₁+k₃*y₂*y₃,
       0     ~  y₁ + y₂ + y₃ - 1,
       D(y₂) ~  k₁*y₁-k₂*y₂^2-k₃*y₂*y₃]
@named sys = ODESystem(eqs, defaults=[k₁ => 100, k₂ => 3e7, y₁ => 1.0])
u0 = Pair[]
push!(u0, y₂ => 0.0)
push!(u0, y₃ => 0.0)
p  = [k₁ => 0.04,
      k₃ => 1e4]
p2 = (k₁ => 0.04,
      k₂ => 3e7,
      k₃ => 1e4)
tspan = (0.0,100000.0)
prob1 = ODEProblem(sys,u0,tspan,p)
prob12 = ODEProblem(sys,u0,tspan,[0.04,3e7,1e4])
prob13 = ODEProblem(sys,u0,tspan,(0.04,3e7,1e4))
prob14 = ODEProblem(sys,u0,tspan,p2)
for p in [prob1, prob14]
    @test Set(Num.(parameters(sys)) .=> p.p) == Set([k₁=>0.04, k₂=>3e7, k₃=>1e4])
    @test Set(Num.(states(sys)) .=> p.u0) == Set([y₁=>1, y₂=>0, y₃=>0])
end
prob2 = ODEProblem(sys,u0,tspan,p,jac=true)
prob3 = ODEProblem(sys,u0,tspan,p,jac=true,sparse=true)
@test prob3.f.sparsity isa SparseMatrixCSC
@test_throws ArgumentError ODEProblem(sys,zeros(5),tspan,p)
for (prob, atol) in [(prob1, 1e-12), (prob2, 1e-12), (prob3, 1e-12)]
    local sol
    sol = solve(prob, Rodas5())
    @test all(x->≈(sum(x), 1.0, atol=atol), sol.u)
end

du0 = [
       D(y₁) => -0.04
       D(y₂) => 0.04
       D(y₃) => 0.0
      ]
prob4 = DAEProblem(sys, du0, u0, tspan, p2)
prob5 = eval(DAEProblemExpr(sys, du0, u0, tspan, p2))
for prob in [prob4, prob5]
    local sol
    @test prob.differential_vars == [true, true, false]
    sol = solve(prob, IDA())
    @test all(x->≈(sum(x), 1.0, atol=1e-12), sol.u)
end

@parameters t σ β
@variables x(t) y(t) z(t)
D = Differential(t)
eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x-β*y,
       x + z ~ y]
@named sys = ODESystem(eqs)
@test all(isequal.(states(sys), [x, y, z]))
@test all(isequal.(parameters(sys), [σ, β]))
@test equations(sys) == eqs
@test ModelingToolkit.isautonomous(sys)

# issue 701
using ModelingToolkit
@parameters t a
@variables x(t)
D = Differential(t)
@named sys = ODESystem([D(x) ~ a])
@test equations(sys)[1].rhs isa Sym

# issue 708
@parameters t a
@variables x(t) y(t) z(t)
D = Differential(t)
@named sys = ODESystem([D(x) ~ y, 0 ~ x + z, 0 ~ x - y], t, [z, y, x], [])
sys2 = ode_order_lowering(sys)
M = ModelingToolkit.calculate_massmatrix(sys2)
@test M == Diagonal([1, 0, 0])

# issue #609
@variables t x1(t) x2(t)
D = Differential(t)

eqs = [
   D(x1) ~ -x1,
   0 ~ x1 - x2,
]
@named sys = ODESystem(eqs, t)
@test isequal(ModelingToolkit.get_iv(sys), t)
@test isequal(states(sys), [x1, x2])
@test isempty(parameters(sys))

# one equation ODESystem test
@parameters t r
@variables x(t)
D = Differential(t)
eq = D(x) ~ r*x
@named ode = ODESystem(eq)
@test equations(ode) == [eq]
# issue #808
@testset "Combined system name collisions" begin
    function makesys(name)
        @parameters t a
        @variables x(t) f(t)
        D = Differential(t)

        ODESystem([D(x) ~ -a*x + f], name = name)
    end

    function issue808()
        sys1 = makesys(:sys1)
        sys2 = makesys(:sys1)

        @parameters t
        D = Differential(t)
        @test_throws ArgumentError ODESystem([sys2.f ~ sys1.x, D(sys1.f) ~ 0], t, systems = [sys1, sys2], name=:foo)
    end
    issue808()

end

#Issue 998
@parameters t
pars = []
vars = @variables((u1,))
der = Differential(t)
eqs = [
  der(u1) ~ 1,
]
@test_throws ArgumentError ODESystem(eqs, t, vars, pars, name=:foo)

#Issue 1063/998
pars = [t]
vars = @variables((u1(t),))
@test_throws ArgumentError ODESystem(eqs, t, vars, pars, name=:foo)

@parameters w
der = Differential(w)
eqs = [
  der(u1) ~ t,
]
@test_throws ArgumentError ModelingToolkit.ODESystem(eqs, t, vars, pars, name=:foo)

@variables x(t)
D = Differential(t)
@parameters M b k
eqs = [D(D(x)) ~ -b/M*D(x) - k/M*x]
ps = [M, b, k]
default_u0 = [D(x) => 0.0, x => 10.0]
default_p = [M => 1.0, b => 1.0, k => 1.0]
@named sys = ODESystem(eqs, t, [x], ps, defaults=[default_u0; default_p])
sys = ode_order_lowering(sys)
prob = ODEProblem(sys, [], tspan)
sol = solve(prob, Tsit5())
@test sum(abs, sol[end]) < 1

# check_eqs_u0 kwarg test
@parameters t
@variables x1(t) x2(t)
D = Differential(t)
eqs = [D(x1) ~ -x1]
@named sys = ODESystem(eqs,t,[x1,x2],[])
@test_throws ArgumentError ODEProblem(sys, [1.0,1.0], (0.0,1.0))
prob = ODEProblem(sys, [1.0,1.0], (0.0,1.0), check_length=false)

# check inputs
let
    @parameters t f k d
    @variables x(t) ẋ(t)
    δ = Differential(t)

    eqs = [δ(x) ~ ẋ, δ(ẋ) ~ f - k*x - d*ẋ]
    @named sys = ODESystem(eqs, t, [x, ẋ], [f, d, k]; controls = [f])

    calculate_control_jacobian(sys)

    @test isequal(
        calculate_control_jacobian(sys),
        reshape(Num[0,1], 2, 1)
    )
end

# issue 1109
let
    @variables t x[1:3,1:3](t)
    D = Differential(t)
    @named sys = ODESystem(D.(x) .~ x)
    @test_nowarn structural_simplify(sys)
end

# Array vars
using Symbolics: unwrap, wrap
using LinearAlgebra
@variables t
sts = @variables x[1:3](t) y(t)
ps = @parameters p[1:3] = [1, 2, 3]
D = Differential(t)
eqs = [
       collect(D.(x) ~ x)
       D(y) ~ norm(x)*y
      ]
@named sys = ODESystem(eqs, t, [sts...;], [ps...;])
@test isequal(@nonamespace(sys.x), unwrap(x))
@test isequal(@nonamespace(sys.y), unwrap(y))
@test isequal(@nonamespace(sys.p), unwrap(p))
@test_nowarn sys.x, sys.y, sys.p

# Mixed Difference Differential equations
@parameters t a b c d
@variables x(t) y(t)
δ = Differential(t)
D = Difference(t; dt=0.1)
eqs = [
    δ(x) ~ a*x - b*x*y,
    δ(y) ~ -c*y + d*x*y,
    D(x) ~ y
]
@named de = ODESystem(eqs,t,[x,y],[a,b,c,d])
@test generate_difference_cb(de) isa ModelingToolkit.DiffEqCallbacks.DiscreteCallback

# doesn't work with ODEFunction
# prob = ODEProblem(ODEFunction{false}(de),[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0])

prob = ODEProblem(de,[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0], check_length=false)
@test prob.kwargs[:difference_cb] isa ModelingToolkit.DiffEqCallbacks.DiscreteCallback

sol = solve(prob, Tsit5(); callback=prob.kwargs[:difference_cb], tstops=prob.tspan[1]:0.1:prob.tspan[2])

# Direct implementation
function lotka(du,u,p,t)
    x = u[1]
    y = u[2]
    du[1] = p[1]*x - p[2]*x*y
    du[2] = -p[3]*y + p[4]*x*y
end

prob2 = ODEProblem(lotka,[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0])
function periodic_difference_affect!(int)
    int.u += [int.u[2], 0]
end

difference_cb = ModelingToolkit.PeriodicCallback(periodic_difference_affect!, 0.1)

sol2 = solve(prob2, Tsit5(); callback=difference_cb, tstops=collect(prob.tspan[1]:0.1:prob.tspan[2])[2:end]
)

@test sol(0:0.01:1)[x] ≈ sol2(0:0.01:1)[1,:]
@test sol(0:0.01:1)[y] ≈ sol2(0:0.01:1)[2,:]

using ModelingToolkit

function submodel(;name)
    @variables t y(t)
    @parameters A[1:5]
    A = collect(A)
    D = Differential(t)
    ODESystem(D(y) ~ sum(A) * y; name=name)
end

# Buid system
@named sys1 = submodel()
@named sys2 = submodel()

@variables t
@named sys = ODESystem([0 ~ sys1.y + sys2.y ], t; systems=[sys1, sys2])


@testset "NumField/NfAbs/NfAbs" begin
  cyclo_expl = function(n, m)
    Fn, zn = CyclotomicField(n)
    Fnm, znm = CyclotomicField(n*m)
    x = zn
    x_up = Hecke.force_coerce_cyclo(Fnm, x)
    x_down = Hecke.force_coerce_cyclo(Fn, x_up)
    return (x, x_up, x_down)
  end

  res = cyclo_expl(3, 4)
  @test (res[1]^3, res[2]^3) == (1, 1)

  res = cyclo_expl(3, 2)
  z6 = gen(parent(res[2]))
  # Test that z3 is mapped to z6^2
  @test z6^2 == res[2]
end

@testset "Splitting Field" begin

  Qx, x = PolynomialRing(FlintQQ, "x")
  f = x^3-2
  K = splitting_field(f)
  @test typeof(K) == AnticNumberField
  K1 = number_field([x^3-2, x^2+x+1])[1]
  K1abs = simple_extension(K1)[1]
  @test isisomorphic(K, K1abs)[1]
  K, R = splitting_field(f, do_roots = true)
  for r in R
    @test iszero(f(r))
  end

  K, a = number_field(x^2+1)
  Kt, t = PolynomialRing(K, "t")
  g = t^4-2
  L = splitting_field(g)
  @test typeof(L) == Hecke.NfRel{nf_elem}
  @test absolute_degree(L) == 8
end


# using Revise, Plots
using PseudoArcLengthContinuation, LinearAlgebra, Setfield, SparseArrays
const PALC = PseudoArcLengthContinuation

k = 2
N = 10
F = (x, p) -> p[1] .* x .+ x.^(k+1)/(k+1) .+ 0.01
Jac_m = (x, p) -> diagm(0 => p[1]  .+ x.^k)

normInf = x -> norm(x, Inf)

opts = PALC.ContinuationPar(dsmax = 0.051, dsmin = 1e-3, ds=0.001, maxSteps = 140, pMin = -3., saveSolEveryNsteps = 0, newtonOptions = NewtonPar(tol = 1e-8, verbose = false), computeEigenValues = false, saveEigenvectors = false)
x0 = 0.01 * ones(N)

opts = @set opts.doArcLengthScaling = true
br0, sol, _ = @time PALC.continuation(F,Jac_m,x0, -1.5, (@lens _),opts,verbosity=0) #(15.28 k allocations: 1.055 MiB)

# test with callbacks
br0, sol, _ = @time PALC.continuation(F,Jac_m,x0, -1.5, (@lens _), (@set opts.maxSteps = 3), verbosity=2, callbackN = (x, f, J, res, iteration, itlinear, optionsN; kwargs...)->(@show x;true))

###### Used to check type stability of the methods
# using RecursiveArrayTools
# iter = PALC.PALCIterable(F,Jac_m,x0,-1.5, (@lens _), opts,verbosity=0)
# state = iterate(iter)[1]
# contRes = PALC.initContRes(iter, state)
# @time continuation!(iter, state, contRes)
#
# typeof(contRes)
#
# state = iterate(iter)[1]
#     contRes = PALC.initContRes(iter, state)
#     @code_warntype continuation!(iter, state, contRes)
#####

opts = @set opts.doArcLengthScaling = false
br1, sol, _ = @time PALC.continuation(F,Jac_m,x0,-1.5, (@lens _),opts,verbosity=0) #(14.28 k allocations: 1001.500 KiB)
show(br1)

br2, sol, _ = @time PALC.continuation(F,Jac_m,x0,-1.5, (@lens _),opts,verbosity=0, printSolution = (x,p) -> norm(x,2))

# test for different norms
br3, sol, _ = @time PALC.continuation(F,Jac_m,x0,-1.5, (@lens _),opts,verbosity=0, normC = normInf)

# test for linesearch in Newton method
opts = @set opts.newtonOptions.linesearch = true
br4, sol, _ = @time PALC.continuation(F,Jac_m,x0,-1.5, (@lens _),opts,verbosity=0, normC = normInf) # (15.61 k allocations: 1.020 MiB)

# test for different ways to solve the bordered linear system arising during the continuation step
opts = @set opts.newtonOptions.linesearch = false
br5, sol, _ = @time PALC.continuation(F,Jac_m,x0,-1.5, (@lens _),opts,verbosity=0, normC = normInf, linearAlgo = PALC.BorderingBLS())

br5, sol, _ = @time PALC.continuation(F,Jac_m,x0,-1.5, (@lens _),opts,verbosity=0, normC = normInf, linearAlgo = PALC.MatrixBLS())

# test for stopping continuation based on user defined function
finaliseSolution = (z, tau, step, contResult) -> (step < 20)
br5a, sol, _ = @time PALC.continuation(F,Jac_m,x0,-1.5, (@lens _),opts,verbosity=2, finaliseSolution = finaliseSolution)
@test length(br5a.branch) == 21

# test for different predictors
br6, sol, _ = @time PALC.continuation(F,Jac_m,x0,-1.5, (@lens _),opts,verbosity=0, tangentAlgo = PALC.SecantPred())

optsnat = setproperties(opts; ds = 0.001, dsmax = 0.02, dsmin = 0.0001)
br7, sol, _ = @time PALC.continuation(F,Jac_m,x0,-1.5, (@lens _),optsnat,verbosity=0, tangentAlgo = PALC.NaturalPred())

# tangent prediction with Bordered predictor
br8, sol, _ = @time PALC.continuation(F,Jac_m,x0,-1.5, (@lens _),opts,verbosity=0, tangentAlgo = PALC.BorderedPred())

# further testing with sparse Jacobian operator
Jac_sp_simple = (x, p) -> SparseArrays.spdiagm(0 => p  .+ x.^k)
brsp, sol, _ = @time PALC.continuation(F,Jac_sp_simple,x0,-1.5, (@lens _),opts,verbosity=0)
brsp, sol, _ = @time PALC.continuation(F,Jac_sp_simple,x0,-1.5, (@lens _),opts,verbosity=0, printSolution = (x,p) -> norm(x,2))
brsp, sol, _ = @time PALC.continuation(F,Jac_sp_simple,x0,-1.5, (@lens _),opts,verbosity=0,linearAlgo = PALC.BorderingBLS())
brsp, sol, _ = @time PALC.continuation(F,Jac_sp_simple,x0,-1.5, (@lens _),opts,verbosity=0,linearAlgo = PALC.MatrixBLS())
# plotBranch(br1,marker=:d);title!("")
# plotBranch!(br8,marker=:d);title!("")


# Generic concepts and functions

# a premetric is a function d that satisfies:
#
#   d(x, y) >= 0
#   d(x, x) = 0
#
abstract type PreMetric end

# a semimetric is a function d that satisfies:
#
#   d(x, y) >= 0
#   d(x, x) = 0
#   d(x, y) = d(y, x)
#
abstract type SemiMetric <: PreMetric end

# a metric is a semimetric that satisfies triangle inequality:
#
#   d(x, y) + d(y, z) >= d(x, z)
#
abstract type Metric <: SemiMetric end

# Generic functions

result_type(::PreMetric, ::AbstractArray, ::AbstractArray) = Float64

# Generic column-wise evaluation

function colwise!(r::AbstractArray, metric::PreMetric, a::AbstractVector, b::AbstractMatrix)
    n = size(b, 2)
    length(r) == n || throw(DimensionMismatch("Incorrect size of r."))
    @inbounds for j = 1:n
        r[j] = evaluate(metric, a, view(b, :, j))
    end
    r
end

function colwise!(r::AbstractArray, metric::PreMetric, a::AbstractMatrix, b::AbstractVector)
    n = size(a, 2)
    length(r) == n || throw(DimensionMismatch("Incorrect size of r."))
    @inbounds for j = 1:n
        r[j] = evaluate(metric, view(a, :, j), b)
    end
    r
end

function colwise!(r::AbstractArray, metric::PreMetric, a::AbstractMatrix, b::AbstractMatrix)
    n = get_common_ncols(a, b)
    length(r) == n || throw(DimensionMismatch("Incorrect size of r."))
    @inbounds for j = 1:n
        r[j] = evaluate(metric, view(a, :, j), view(b, :, j))
    end
    r
end

function colwise!(r::AbstractArray, metric::SemiMetric, a::AbstractMatrix, b::AbstractVector)
    colwise!(r, metric, b, a)
end

function colwise(metric::PreMetric, a::AbstractMatrix, b::AbstractMatrix)
    n = get_common_ncols(a, b)
    r = Vector{result_type(metric, a, b)}(undef, n)
    colwise!(r, metric, a, b)
end

function colwise(metric::PreMetric, a::AbstractVector, b::AbstractMatrix)
    n = size(b, 2)
    r = Vector{result_type(metric, a, b)}(undef, n)
    colwise!(r, metric, a, b)
end

function colwise(metric::PreMetric, a::AbstractMatrix, b::AbstractVector)
    n = size(a, 2)
    r = Vector{result_type(metric, a, b)}(undef, n)
    colwise!(r, metric, a, b)
end

# Generic pairwise evaluation

function pairwise!(r::AbstractMatrix, metric::PreMetric, a::AbstractMatrix, b::AbstractMatrix)
    na = size(a, 2)
    nb = size(b, 2)
    size(r) == (na, nb) || throw(DimensionMismatch("Incorrect size of r."))
    @inbounds for j = 1:size(b, 2)
        bj = view(b, :, j)
        for i = 1:size(a, 2)
            r[i, j] = evaluate(metric, view(a, :, i), bj)
        end
    end
    r
end

function pairwise!(r::AbstractMatrix, metric::PreMetric, a::AbstractMatrix)
    pairwise!(r, metric, a, a)
end

function pairwise!(r::AbstractMatrix, metric::SemiMetric, a::AbstractMatrix)
    n = size(a, 2)
    size(r) == (n, n) || throw(DimensionMismatch("Incorrect size of r."))
    @inbounds for j = 1:n
        aj = view(a, :, j)
        for i = (j + 1):n
            r[i, j] = evaluate(metric, view(a, :, i), aj)
        end
        r[j, j] = 0
        for i = 1:(j - 1)
            r[i, j] = r[j, i]   # leveraging the symmetry of SemiMetric
        end
    end
    r
end

function pairwise(metric::PreMetric, a::AbstractMatrix, b::AbstractMatrix)
    m = size(a, 2)
    n = size(b, 2)
    r = Matrix{result_type(metric, a, b)}(undef, m, n)
    pairwise!(r, metric, a, b)
end

function pairwise(metric::PreMetric, a::AbstractMatrix)
    n = size(a, 2)
    r = Matrix{result_type(metric, a, a)}(undef, n, n)
    pairwise!(r, metric, a)
end


using QuDirac
using Base.Test

default_inner(UndefinedInner())

include("constest.jl")
include("abstractinnertests.jl")
include("generaltests.jl")

default_inner(KroneckerDelta())

include("constest.jl")
include("orthotests.jl")
include("generaltests.jl")
include("laddertests.jl")


CONSTRAINT_TYPES = Union{T, Vector{T}} where T <: SamplingConstraint

"""
    ConstrainedValueResampling{N_DATASETS}

Indicates that resampling should be done with constraints on the furnishing distributions/populations.

## Fields

- `constraints`. The constraints for the datasets. The constraints are represented as a tuple of length
    `N_DATASETS`, where the `i`-th tuple element contains the constraints for that dataset.
    Constraints for each dataset must be supplied as either a single sampling constraint,
    or as a vector of sampling constraints with length matching the length of the dataset
    (`Union{SamplingConstraint, Vector{<:SamplingConstraint}}}`). For example, if the `i`-th dataset
    contains 352 observations, then `constraints[i]` must be either a single
    sampling constraint (e.g. `TruncateStd(1.1)`) or a vector of 352 different sampling constraints
    (e.g. `[TruncateStd(1.0 + rand()) for i = 1:352]`).
- `n::Int`. The number of draws.

## Example

Assume we have three collections of uncertain values of, each of length `L = 50`. These should be
resampled `250` times. Before resampling, however, the distributions/populations furnishing
the uncertain values should be truncated:

- For the first collection, truncate each value at `1.5` times its standard deviation
    around its mean. This could simulate measurement errors from an instrument
    that yields stable measurements whose errors are normally distributed, but for
    which we are not interested in outliers or values beyond `1.5` standard devations
    for our analyses.
- For the second collection, truncate each value at the `80`th percentile range.
    This could simulate measurement errors from an instrument that yields stable measurements,
    whose errors are not normally distributed, so that confidence intervals are better
    to use than standard deviations. In this case, we're not interested in outliers,
    and therefore exclude values smaller than the `10`th percentile and larger than the
    `90`th percentile of the data.
- For the third collection, truncate the `i`-th value at an fraction of its standard
    deviation around the mean slightly larger than at the `i-1`-th value, so that
    the standard deviation ranges from `0.5` to `0.5 + L/100`. This could simulate,
    for example, an instrument whose measurement error increases over time.

```julia
L = 50
constraints_d1 = TruncateStd(1.5)
constraints_d2 = TruncateQuantiles(0.1, 0.9)
constraints_d3 = [TruncateStd(0.5 + i/100) for i = 1:L]
```
"""
struct ConstrainedValueResampling{N} <: AbstractUncertainDataResampling
    constraints::Tuple{Vararg{CONSTRAINT_TYPES, N}}
    n::Int
end

# It's tedious for the user to always provide a tuple, so unpack a variable
# number of constraints into a tuple, then call the original constructor.
function ConstrainedValueResampling(constraints::CONSTRAINT_TYPES...) where N
    ConstrainedValueResampling((constraints...,), 1)
end

function ConstrainedValueResampling(n::Int, constraints::CONSTRAINT_TYPES...) where N
    ConstrainedValueResampling((constraints...,), n)
end

Broadcast.broadcastable(c::ConstrainedValueResampling) = Ref(c)

Base.length(c::ConstrainedValueResampling) = length(c.constraints)
Base.firstindex(c::ConstrainedValueResampling) = 1
Base.lastindex(c::ConstrainedValueResampling) = length(c)
Base.getindex(c::ConstrainedValueResampling, i) = c.constraints[i]
Base.iterate(c::ConstrainedValueResampling, state = 1) = iterate(c.constraints, state)
Base.eachindex(c::ConstrainedValueResampling) = Base.OneTo(length(c.constraints))

function Base.show(io::IO, constraints::ConstrainedValueResampling{N_DATASETS}) where {N_DATASETS}
    s = "$(typeof(constraints)) for $N_DATASETS set(s) of value constraints, where n=$(constraints.n)"
    println(io, s)
end

export ConstrainedValueResampling


# tests for sum, mean, and var, which implemented by calling other atoms

x = create_var(5, 3)
s = sum(x)
m = mean(x)
test_expr(s, (x,))
test_expr(m, (x,))

t = sum(x, 1)
n = mean(x, 1)
test_expr(t, (x,))
test_expr(n, (x,))

u = sum(x, 2)
p = mean(x, 2)
test_expr(u, (x,))
test_expr(p, (x,))

v = var(x)
test_expr(v)

info("All sum/mean/var tests passed")


include("euler/euler.jl")
using .Primes: prime_numbers
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

compute(n::Int)::Int = sum(prime_numbers(n))

compute(10)

compute(2_000_000)

@benchmark compute(2_000_000)


# Multinomial Stein Distribution
#
# Represents a Multinomial distribution with parameters ntrials, probs

using Distributions: Multinomial
import Distributions

type SteinMultinomial <: SteinDistribution
    # Number of trials
    ntrials::Int64
    # Probability vector, sums to 1
    probs::Array{Float64}
    # Stein factors, length is 1 less than that of probs
    c1::Float64
    c2::Float64
    c3::Float64
end

# Constructor with (c1,c2,c3) = (1,1,1)
SteinMultinomial(ntrials::Int64, probs::Array{Float64}) = SteinMultinomial(
    ntrials,
    probs,
    1.0,
    1.0,
    1.0
)

# Constructor with (c1,c2,c3) = (1,1,1), all categories equiprobable
SteinMultinomial(ntrials::Int64, ncategories::Int64) = SteinMultinomial(
    ntrials,
    [1.0/ncategories for i=1:ncategories],
    1.0,
    1.0,
    1.0
)

# Draw n independent samples from distribution; last column is dropped
function rand(d::SteinMultinomial, n::Int64)
    Distributions.rand(Multinomial(d.ntrials,d.probs),n)'[:,1:end-1]
end

# Joint distribution evaluated at a point y
function jointdistribution(d::SteinMultinomial, y::Array{Float64})
    Distributions.pdf(Multinomial(d.ntrials,d.probs),[y,d.ntrials-sum(y)])
end

# Conditional distribution of jth coordinate of y, given all other coordinates,
# represented as a probability vector
function conddistribution(d::SteinMultinomial, y::Array{Float64,1}, j::Int64)
    conddist = [jointdistribution(d, [y[1:(j-1)], yj, y[(j+1):end]])
                for yj in 0:d.ntrials]
    conddist / sum(conddist)
end

function supportlowerbound(d::SteinMultinomial, j::Int64)
    0
end

function supportupperbound(d::SteinMultinomial, j::Int64)
    d.ntrials
end

function numdimensions(d::SteinMultinomial)
    length(d.probs) - 1
end

# Cumulative distribution function (only valid when X is univariate)


include("../CovarianceFunctions/CovarianceFunctions.jl")

# For linear algebra
using LinearAlgebra
# logging
using Logging
using PyPlot
## PDEs type
abstract type AbstractPDEs end
struct NonlinElliptic2d{Tα,Tm,TΩ} <: AbstractPDEs
    # eqn: -Δu + α*u^m = f in [Ω[1,1],Ω[2,1]]*[Ω[1,2],Ω[2,2]]
    α::Tα
    m::Tm
    Ω::TΩ
    bdy::Function
    rhs::Function
end

## sample points
function sample_points_rdm(eqn::NonlinElliptic2d, N_domain, N_boundary)
    Ω = eqn.Ω
    x1l = Ω[1,1]
    x1r = Ω[2,1]
    x2l = Ω[1,2]
    x2r = Ω[2,2]

    X_domain = hcat(rand(Float64,(N_domain, 1))*(x1r-x1l).+x1l,rand(Float64,(N_domain, 1))*(x2r-x2l).+x2l)

    N_bd_each=convert(Int64, N_boundary/4)
    if N_boundary != 4 * N_bd_each
        println("[sample points] N_boundary not divided by 4, replaced by ", 4 * N_bd_each)
        N_boundary = 4 * N_bd_each
    end

    X_boundary = zeros((N_boundary, 2))
    # bottom face
    X_boundary[1:N_bd_each, :] = hcat((x1r-x1l)*rand(Float64,(N_bd_each,1)).+x1l, x2l*ones(N_bd_each))
    # right face
    X_boundary[N_bd_each+1:2*N_bd_each, :] = hcat(x1r*ones(N_bd_each),(x2r-x2l)*rand(Float64,(N_bd_each,1)).+x2l)
    # top face
    X_boundary[2*N_bd_each+1:3*N_bd_each, :] = hcat((x1r-x1l)*rand(Float64,(N_bd_each,1)).+x1l, x2r*ones(N_bd_each))

    # left face
    X_boundary[3*N_bd_each+1:N_boundary, :] = hcat(x1l*ones(N_bd_each), (x2r-x2l)*rand(Float64,(N_bd_each,1)).+x2l)
    return X_domain', X_boundary'
end

function sample_points_grid(eqn::NonlinElliptic2d, h_in, h_bd)
    Ω = eqn.Ω
    x1l = Ω[1,1]
    x1r = Ω[2,1]
    x2l = Ω[1,2]
    x2r = Ω[2,2]
    x = x1l + h_in:h_in:x1r-h_in
    y = x2l + h_in:h_in:x2r-h_in
    X_domain = reduce(hcat,[[x[i], y[j]] for i in 1:length(x) for j in 1:length(x)])

    l = length(x1l:h_bd:x1r-h_bd)
    X_boundary = vcat([x1l:h_bd:x1r-h_bd x2l*ones(l)], [x1r*ones(l) x2l:h_bd:x2r-h_bd], [x1r:-h_bd:x1l+h_bd x2r*ones(l)], [x1l*ones(l) x2r:-h_bd:x1l+h_bd])
    return X_domain, X_boundary'
end

## exact algorithms
# assemby Gram matrices
function get_Gram_matrices(eqn::NonlinElliptic2d, cov::AbstractCovarianceFunction, X_domain, X_boundary, sol_now)
    d = 2
    N_domain = size(X_domain,2)
    N_boundary = size(X_boundary,2)
    Δδ_coefs = -1.0
    δ_coefs_int = eqn.α*eqn.m*(sol_now.^(eqn.m-1))

    # get linearized PDEs correponding measurements
    meas_δ = [PointMeasurement{d}(SVector{d,Float64}(X_boundary[:,i])) for i = 1:N_boundary]
    meas_Δδ = [ΔδPointMeasurement{Float64,d}(SVector{d,Float64}(X_domain[:,i]), Δδ_coefs, δ_coefs_int[i]) for i = 1:N_domain]
    meas_test_int = [PointMeasurement{d}(SVector{d,Float64}(X_domain[:,i])) for i = 1:N_domain]

    Theta_train = zeros(N_domain+N_boundary,N_domain+N_boundary)

    measurements = Vector{Vector{<:AbstractPointMeasurement}}(undef,2)
    measurements[1] = meas_δ; measurements[2] = meas_Δδ
    cov(Theta_train, reduce(vcat,measurements))

    Theta_test = zeros(N_domain+N_boundary,N_domain+N_boundary)
    measurements_δ = Vector{Vector{<:AbstractPointMeasurement}}(undef,2)
    measurements_δ[1] = meas_δ; measurements_δ[2] = meas_test_int

    cov(Theta_test, reduce(vcat,measurements_δ), reduce(vcat,measurements))
    # cov(view(Theta_test,1:N_domain,1:N_boundary), meas_test_int, meas_δ)
    # cov(view(Theta_test,1:N_domain,N_boundary+1:N_domain+N_boundary), meas_test_int, meas_Δδ)
    return Theta_train, Theta_test
end

function  get_initial_covariance(cov::AbstractCovarianceFunction, X_domain, X_boundary)
    d = 2
    N_domain = size(X_domain,2)
    N_boundary = size(X_boundary,2)
    meas_δ_bd = [PointMeasurement{d}(SVector{d,Float64}(X_boundary[:,i])) for i = 1:N_boundary]
    meas_bd_int = [PointMeasurement{d}(SVector{d,Float64}(X_domain[:,i])) for i = 1:N_domain]
    Theta_initial = zeros(N_domain+N_boundary,N_domain+N_boundary)
    measurements = Vector{Vector{<:AbstractPointMeasurement}}(undef,2)
    measurements[1] = meas_δ_bd; measurements[2] = meas_bd_int
    cov(Theta_initial, reduce(vcat,measurements))
    return Theta_initial
end

# iterative GPR
function iterGPR_exact(eqn, cov, X_domain, X_boundary, sol_init, noise, GNsteps)
    N_domain = size(X_domain,2)
    N_boundary = size(X_boundary,2)
    # get the rhs and bdy data
    rhs = [eqn.rhs(X_domain[:,i]) for i in 1:N_domain] .+ sqrt(noise)*randn(N_domain)
    bdy = [eqn.bdy(X_boundary[:,i]) for i in 1:N_boundary]
    v = sol_init

    Theta_train = zeros(N_domain+N_boundary, N_domain+N_boundary)
    Theta_test = zeros(N_domain,N_domain+N_boundary)
    for _ in 1:GNsteps
        Theta_train, Theta_test = get_Gram_matrices(eqn, cov, X_domain, X_boundary, v)
        rhs_now = vcat(bdy, rhs.+eqn.α*(eqn.m-1)*v[N_boundary+1:end].^eqn.m)

        v = Theta_test*(((Theta_train+noise*LinearAlgebra.I))\rhs_now)
    end

    Cov = get_initial_covariance(cov, X_domain, X_boundary)
    Cov_posterior = Cov .- Theta_test*(((Theta_train+noise*LinearAlgebra.I))\Theta_test')
    return v[N_boundary+1:end], Cov_posterior
end

α = 10.0
m = 3
Ω = [[0,1] [0,1]]
h_in = 0.02
h_bd = 0.02
lengthscale = 0.3
kernel = "Matern5half"
cov = MaternCovariance5_2(lengthscale)
noise = 0.0001
GNsteps = 3

# ground truth solution
freq = 1000
s = 4
function fun_u(x)
    ans = 0
    @inbounds for k = 1:freq
        ans += sin(pi*k*x[1])*sin(pi*k*x[2])/k^s
        # H^t norm squared is sum 1/k^{2s-2t}, so in H^{s-1/2}
    end
    return ans
end

# right hand side
function fun_rhs(x)
    ans = 0
    @inbounds for k = 1:freq
        ans += (2*k^2*pi^2)*sin(pi*k*x[1])*sin(pi*k*x[2])/k^s
    end
    return ans + α*fun_u(x)^m
end

# boundary value
function fun_bdy(x)
    return fun_u(x)
end

@info "[solver started] NonlinElliptic2d"
@info "[equation] -Δu + $α u^$m = f"
eqn = NonlinElliptic2d(α,m,Ω,fun_bdy,fun_rhs)

# N_domain = 2000
# N_boundary = 400
# X_domain, X_boundary = sample_points_rdm(eqn,N_domain, N_boundary)

X_domain, X_boundary = sample_points_grid(eqn, h_in, h_bd)
N_domain = size(X_domain,2)
N_boundary = size(X_boundary,2)
@info "[sample points] grid size $h_in"
@info "[sample points] N_domain is $N_domain, N_boundary is $N_boundary"
@info "[kernel] choose $kernel, lengthscale $lengthscale\n"
@info "[noise] $noise"
@info "[GNsteps] $GNsteps"

sol_init = zeros(N_domain+N_boundary) # initial solution
truth = [fun_u(X_domain[:,i]) for i in 1:N_domain]
@time sol_exact, sol_postvar = iterGPR_exact(eqn, cov, X_domain, X_boundary, sol_init, noise, GNsteps)
pts_accuracy_exact = sqrt(sum((truth-sol_exact).^2)/sum(truth.^2))
@info "[L2 accuracy: exact method] $pts_accuracy_exact"
pts_max_accuracy_exact = maximum(abs.(truth-sol_exact))/maximum(abs.(truth))
@info "[Linf accuracy: exact method] $pts_max_accuracy_exact"

sol_std = [sqrt(sol_postvar[i,i]) for i in N_boundary+1:N_boundary+N_domain]

Nh = convert(Int,sqrt(N_domain))
figure()
plot_surface(reshape(X_domain[1,:],Nh,Nh), reshape(X_domain[2,:],Nh,Nh), reshape(truth,Nh,Nh), label = "Reference")
plot_surface(reshape(X_domain[1,:],Nh,Nh), reshape(X_domain[2,:],Nh,Nh), reshape(sol_exact - sol_std,Nh,Nh), color="C1", label = "Lower")
plot_surface(reshape(X_domain[1,:],Nh,Nh), reshape(X_domain[2,:],Nh,Nh), reshape(sol_exact + sol_std,Nh,Nh), color="C1", label = "Upper")
display(gcf())

figure()
contourf(reshape(X_domain[1,:],Nh,Nh), reshape(X_domain[2,:],Nh,Nh), reshape(abs.(truth-sol_exact),Nh,Nh))
colorbar()
display(gcf())

figure()
contourf(reshape(X_domain[1,:],Nh,Nh), reshape(X_domain[2,:],Nh,Nh), reshape(sol_std,Nh,Nh))
colorbar()
display(gcf())


module day20

include("inputs.jl")

struct Tile1
    id::Int
    tile::Array{Int,2}
    edges::Dict{String,Int}
end

ex = inputs.ex_d20();
input = inputs.input_d20();

tiles_raw = split(ex, "\n\n")
tile_raw = tiles_raw[1]

function parse_tile(tile_raw)
    lines = split(tile_raw, "\n")
    id = parse(Int, lines[1][6:9])

    tile_s = replace.(lines[2:end], "."=>"0") .|>x->replace(x, "#"=>"1")
    tile = permutedims(vcat(hcat([[parse(Int, c) for c in line] for line in tile_s]...)))

    edges = Dict{}()
    edges["t"] = parse(Int, tile_s[1,:][1], base=2)
    edges["b"] = parse(Int, tile_s[end,:][1], base=2)
    edges["l"] = parse(Int, String([c[1] for c in tile_s]), base=2)
    edges["r"] = parse(Int, String([c[end] for c in tile_s]), base=2)

    return Tile1(id, tile, edges)
end

function rotate(tile)
    edges = Dict{}()
    edges["t"] = tile.edges["l"]
    edges["b"] = tile.edges["r"]
    edges["l"] = tile.edges["b"]
    edges["r"] = tile.edges["t"]
    return Tile1(tile.id, rotr90(tile.tile), edges)
end

function flip(tile)
    edges = Dict{}()
    edges["t"] = parse(Int, reverse(bitstring(tile.edges["t"])[end-9:end]), base=2)
    edges["b"] = parse(Int, reverse(bitstring(tile.edges["b"])[end-9:end]), base=2)
    edges["l"] = tile.edges["r"]
    edges["r"] = tile.edges["l"]
    return Tile1(tile.id, reverse(tile.tile, dims=2), edges)
end

apply(f, tile) = f(tile)

# parse(Int, reverse(bitstring(6)[end-11:end]), base=2)
end


# This is essentially the definition of sampling from a dirac delta.
Base.rand(RNG::AbstractRNG, P::Gaussian{U, Zero}) where U = P.μ

function kalman_sample(m::Model, inputs, rng::AbstractRNG, start_state,
                       N::Integer=length(inputs))
    # This code was optimized in 0.6, but sampling doesn't usually have to be
    # hyper-efficient, and it would look cleaner with a loop.
    result = accumulate(1:N; init=(start_state, nothing)) do v, t
        state, _ = v
        next_state = transition_mat(m, inputs, t) * state +
            rand(rng, transition_noise(m, inputs, t))
        return (next_state,
                observation_mat(m, inputs, t) * next_state +
                rand(rng, observation_noise(m, inputs, t)))
    end
    return (map(first, result), map(last, result))
end

################################################################################
# sample_and_recover

""" A data structure containing the results of `sample_and_recover`. """
struct RecoveryResults
    true_model       # the true model, passed by the user
    estimated_model  # the model estimated from maximum likelihood
    true_samples     # the samples generated from the true model
    estimated_state  # the P(state|observation, estimated_model) distributions
    obs              # the sampled observations
    optim            # the result from calling `optimize` (contains convergence details)
end
parameter_accuracy_ratios(rr::RecoveryResults) =
    [f=>getfield(rr.estimated_model, f) ./ getfield(rr.true_model, f)
     for f in parameters(rr.estimated_model)]

function Base.show(io::IO, rr::RecoveryResults)
    println(io, "Ratio of estimated/true parameters (1.0 is best): ")
    for (f, ratio) in parameter_accuracy_ratios(rr)
        println(io, "  ",
                f, " => ", round.(ratio*1000)/1000)
    end
    # This functionality was neat, but requiring Plots is not nice.
    # show(io, MIME"text/html"(),
    #      plot_hidden_state(1:length(rr.obs), rr.estimated_state;
    #                        true_samples=rr.true_samples))
end

""" See if we can recover the model parameters _and_ the true parameters using
data generated from the model.

Concretely, we sample observations and hidden state from `true_model` for the
given `inputs`, then call `optimize` on `true_model * fuzz_factor`.

If `start_model` isn't specified, we start from a model in the neighborhood of
`true_model` (with `fuzz_factor ~= 1.0` controlling how far we start).

We return a `RecoveryResults` object. See its definition for details. """
function sample_and_recover(true_model::Model, inputs, N=nothing; rng=GLOBAL_RNG,
                            parameters_to_optimize=parameters(true_model),
                            fuzz_factor=exp.(randn(rng, length(get_params(true_model, parameters_to_optimize)))),
                            initial_state::Gaussian=initial_state(true_model),
                            start_model=nothing)
    rng = rng isa AbstractRNG ? rng : MersenneTwister(rng::Integer)
    inputs2 = get_inputs(true_model, inputs)
    N = N === nothing ? length(inputs2) : N
    true_samples, obs =
        kalman_sample(true_model, inputs2, rng, rand(rng, initial_state), N)
    log_likelihood(true_model, inputs2, obs)
    if start_model === nothing
        true_params = get_params(true_model, parameters_to_optimize)
        start_model = set_params(true_model, (true_params .* fuzz_factor),
                                 parameters_to_optimize)
    end
    (best_model, o) = optimize(start_model, inputs, obs;
                               initial_state=initial_state,
                               parameters_to_optimize=parameters_to_optimize)
    estimated_state = kalman_smoother(best_model, get_inputs(best_model, inputs), obs,
                                      initial_state=initial_state)
    return RecoveryResults(true_model, best_model, true_samples, estimated_state, obs, o)
end


function toh02()
#
#  generalization of K.C. Toh and L. N. Trefethen
#
    k = [(20:-1:0)...]
    p1 = exp10.(k)
    p2 = ones(21)
    for j = 20:-1:0, p2[21-j] = factorial(j); end
    p = p1 ./ p2;
    z = 2 .* (k .+0.5) ./ 19 + im .* sin.((2*pi) .* (k .+ 0.5) ./ 19)
    z = [1.08045842459084 + 0.92291979046775im;
        1.08045842459084 - 0.92291979046775im;
        0.57624370699834 + 0.97554880155012im;
        0.57624370699834 - 0.97554880155012im;
        0.23672740822527 + 0.94133582606706im;
        0.23672740822527 - 0.94133582606706im;
        -0.01684064074376 + 0.86388496107331im;
        -0.01684064074376 - 0.86388496107331im;
        -0.21255469312515 + 0.76040546396802im;
        -0.21255469312515 - 0.76040546396802im;
        -0.36406222857634 + 0.63986739119691im;
        -0.36406222857634 - 0.63986739119691im;
        -0.47903273599273 + 0.50773158889098im;
        -0.47903273599273 - 0.50773158889098im;
        -0.56209888019876 + 0.36770724207462im;
        -0.56209888019876 - 0.36770724207462im;
        -0.64273026408505 + 0.07449713941656im;
        -0.64273026408505 - 0.07449713941656im;
        -0.61611009709266 + 0.22255230874526im;
        -0.61611009709266 - 0.22255230874526im;
        ];
   z = [z ones(20)];
   p, PolyZeros(z)
end


"""
LatentND <: AbstractLatent

# Description
A mutable describing a multivariate latent variable.

# Fields
- **`names::Vector{String}`**
- **`val::Vector{Float64}`**
- **`bounds::Vector{Vector{Float64}}`**
- **`prior::Distributions.MultivariateDistribution`**
- **`posterior::Distributions.MultivariateDistribution`**
- **`chain::Vector{Vector{Float64}}`**
- **`expected_information::Matrix{Float64}`**

# Factories
LatentND(names, val, bounds, prior, posterior, chain, expected_information) = new(names, val, bounds, prior, posterior, chain, expected_information)
Creates a new N-dimensional latent variable with custom fields.

# Random Initializers
LatentND()

Randomly generates a value for the N-dimensional latent variable and assigns a default standardized Gaussian prior and posterior
LatentND(dist, bounds)

Randomly generates a value for the N-dimensional latent variable and assigns a custom multivariate distribution to prior and posterior with specific bounds.
"""
mutable struct LatentND <: AbstractLatent
    names::Vector{String}
    val::Vector{Float64}
    bounds::Vector{Vector{Float64}}
    prior::Distributions.MultivariateDistribution
    posterior::Distributions.MultivariateDistribution
    chain::Vector{Vector{Float64}}
    expected_information::Matrix{Float64}

    function LatentND(names, val, bounds, prior, posterior, chain, expected_information)
        new(names, val, bounds, prior, posterior, chain, expected_information)
    end

    # Random Initializers

    function LatentND(
        dist::Distributions.MultivariateDistribution,
        bounds::Vector{Vector{Float64}},
    )
        N = size(bounds, 1)
        val = truncate_rand(dist, bounds)
        new(string.("L_", collect(1:N)), val, bounds, dist, dist, [val], LinearAlgebra.I(N))
    end
end


@testset "Parallel.Johnson" begin
    g3 = PathGraph(5)
    d = Symmetric([0 1 2 3 4; 1 0 6 7 8; 2 6 0 11 12; 3 7 11 0 16; 4 8 12 16 0])
    for g in testgraphs(g3)
        z = @inferred(Parallel.johnson_shortest_paths(g, d))
        @test z.dists[3, :][:] == [7, 6, 0, 11, 27]
        @test z.parents[3, :][:] == [2, 3, 0, 3, 4]
        @test @inferred(enumerate_paths(z))[2][2] == []
        @test @inferred(enumerate_paths(z))[2][4] == enumerate_paths(z, 2)[4] == enumerate_paths(z, 2, 4) == [2, 3, 4]
    end

    g4 = PathDiGraph(4)
    for g in testdigraphs(g4)
        z = @inferred(Parallel.johnson_shortest_paths(g))
        @test length(enumerate_paths(z, 4, 3)) == 0
        @test length(enumerate_paths(z, 4, 1)) == 0
        @test length(enumerate_paths(z, 2, 3)) == 2
    end

    g5 = DiGraph([1 1 1 0 1; 0 1 0 1 1; 0 1 1 0 0; 1 0 1 1 0; 0 0 0 1 1])
    d = [0 3 8 0 -4; 0 0 0 1 7; 0 4 0 0 0; 2 0 -5 0 0; 0 0 0 6 0]
    for g in testdigraphs(g5)
        z = @inferred(Parallel.johnson_shortest_paths(g, d))
        @test z.dists == [0 1 -3 2 -4; 3 0 -4 1 -1; 7 4 0 5 3; 2 -1 -5 0 -2; 8 5 1 6 0]
    end
end


using Test, LinearAlgebra
using DiffEqSensitivity, StochasticDiffEq
using ForwardDiff, Zygote
using Random

@info "SDE Non-Diagonal Noise Adjoints"

seed = 100
Random.seed!(seed)

tstart = 0.0
tend = 0.1
dt = 0.005
trange = (tstart, tend)
t = tstart:dt:tend
tarray = collect(t)

function g(u,p,t)
  sum(u.^2.0/2.0)
end

function dg!(out,u,p,t,i)
  (out.=-u)
end

# non-diagonal noise
@testset "Non-diagonal noise tests" begin
    Random.seed!(seed)

    u₀ = [0.75,0.5]
    p = [-1.5,0.05,0.2, 0.01]

    dtnd = tend/1e3

    # Example from Roessler, SIAM J. NUMER. ANAL, 48, 922–952 with d = 2; m = 2
    function f_nondiag!(du,u,p,t)
      du[1] = p[1]*u[1] + p[2]*u[2]
      du[2] = p[2]*u[1] + p[1]*u[2]
      nothing
    end

    function g_nondiag!(du,u,p,t)
      du[1,1] = p[3]*u[1] + p[4]*u[2]
      du[1,2] = p[3]*u[1] + p[4]*u[2]
      du[2,1] = p[4]*u[1] + p[3]*u[2]
      du[2,2] = p[4]*u[1] + p[3]*u[2]
      nothing
    end

    function f_nondiag(u,p,t)
      dx = p[1]*u[1] + p[2]*u[2]
      dy = p[2]*u[1] + p[1]*u[2]
      [dx,dy]
    end

    function g_nondiag(u,p,t)
      du11 = p[3]*u[1] + p[4]*u[2]
      du12 = p[3]*u[1] + p[4]*u[2]
      du21 = p[4]*u[1] + p[3]*u[2]
      du22 = p[4]*u[1] + p[3]*u[2]

      [du11  du12
        du21  du22]
    end

    function f_nondiag_analytic(u0,p,t,W)
      A = [[p[1], p[2]] [p[2], p[1]]]
      B = [[p[3], p[4]] [p[4], p[3]]]
      tmp = A*t + B*W[1] + B*W[2]
      exp(tmp)*u0
    end

    noise_matrix = similar(p,2,2)
    noise_matrix .= false

    Random.seed!(seed)
    prob = SDEProblem(f_nondiag!,g_nondiag!,u₀,trange,p,noise_rate_prototype=noise_matrix)
    sol = solve(prob, EulerHeun(), dt=dtnd, save_noise=true)

    noise_matrix = similar(p,2,2)
    noise_matrix .= false
    Random.seed!(seed)
    proboop = SDEProblem(f_nondiag,g_nondiag,u₀,trange,p,noise_rate_prototype=noise_matrix)
    soloop = solve(proboop,EulerHeun(), dt=dtnd, save_noise=true)

    res_sde_u0, res_sde_p = adjoint_sensitivities(soloop,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=BacksolveAdjoint())

    @info res_sde_p

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ReverseDiffNoise()))

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

    @info res_sde_pa

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=BacksolveAdjoint(noise=false))

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

    @info res_sde_pa

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=InterpolatingAdjoint(noise=DiffEqSensitivity.ZygoteNoise()))

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-4)

    @info res_sde_pa

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=InterpolatingAdjoint(noise=DiffEqSensitivity.ReverseDiffNoise()))

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-4)

    @info res_sde_pa

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=InterpolatingAdjoint(noise=false))

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-4)

    @info res_sde_pa

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=BacksolveAdjoint())

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

    @info res_sde_pa

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ZygoteNoise()))

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

    @info res_sde_pa

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=BacksolveAdjoint(noise=false))

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

    @info res_sde_pa

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=InterpolatingAdjoint())

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-4)

    @info res_sde_pa

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=InterpolatingAdjoint(noise=DiffEqSensitivity.ZygoteNoise()))

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-4)

    @info res_sde_pa

    res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,Array(t)
        ,dt=dtnd,adaptive=false,sensealg=InterpolatingAdjoint(noise=false))

    @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
    @test isapprox(res_sde_pa, res_sde_p, rtol=1e-4)

    @info res_sde_pa

    function compute_grads_nd(sol)
      xdis = sol(tarray)

      mat1 = Matrix{Int}(I, 2, 2)
      mat2 = ones(2,2)-mat1

      tmp1 = similar(p)
      tmp1 *= false

      tmp2 = similar(xdis.u[1])
      tmp2 *= false

      for (i, u) in enumerate(xdis)
        tmp1[1]+=xdis.t[i]*u'*mat1*u
        tmp1[2]+=xdis.t[i]*u'*mat2*u
        tmp1[3]+=sum(sol.W(xdis.t[i])[1])*u'*mat1*u
        tmp1[4]+=sum(sol.W(xdis.t[i])[1])*u'*mat2*u

        tmp2 += u.^2
      end

      return tmp2 ./ xdis.u[1], tmp1
    end

    res1, res2 = compute_grads_nd(soloop)

    @test isapprox(res1, res_sde_u0, rtol=1e-4)
    @test isapprox(res2, res_sde_p', rtol=1e-4)

end

@testset "diagonal but mixing noise tests" begin

  Random.seed!(seed)
  u₀ = [0.75,0.5]
  p = [-1.5,0.05,0.2, 0.01]
  dtmix = tend/1e3

  # Example from Roessler, SIAM J. NUMER. ANAL, 48, 922–952 with d = 2; m = 2
  function f_mixing!(du,u,p,t)
    du[1] = p[1]*u[1] + p[2]*u[2]
    du[2] = p[2]*u[1] + p[1]*u[2]
    nothing
  end

  function g_mixing!(du,u,p,t)
    du[1] = p[3]*u[1] + p[4]*u[2]
    du[2] = p[3]*u[1] + p[4]*u[2]
    nothing
  end

  function f_mixing(u,p,t)
    dx = p[1]*u[1] + p[2]*u[2]
    dy = p[2]*u[1] + p[1]*u[2]
    [dx,dy]
  end

  function g_mixing(u,p,t)
    dx = p[3]*u[1] + p[4]*u[2]
    dy = p[3]*u[1] + p[4]*u[2]
    [dx,dy]
  end

  Random.seed!(seed)
  prob = SDEProblem(f_mixing!,g_mixing!,u₀,trange,p)

  soltsave = collect(trange[1]:dtmix:trange[2])
  sol = solve(prob, EulerHeun(), dt=dtmix, save_noise=true, saveat=soltsave)

  Random.seed!(seed)
  proboop = SDEProblem(f_mixing,g_mixing,u₀,trange,p)
  soloop = solve(proboop,EulerHeun(), dt=dtmix, save_noise=true, saveat=soltsave)

  #oop

  res_sde_u0, res_sde_p = adjoint_sensitivities(soloop,EulerHeun(),dg!,tarray
    ,dt=dtmix,adaptive=false,sensealg=BacksolveAdjoint(noisemixing=true))

  @info res_sde_p

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,Array(t)
    ,dt=dtmix,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ZygoteNoise(), noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,Array(t)
    ,dt=dtmix,adaptive=false,sensealg=BacksolveAdjoint(noise=false,noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,Array(t)
      ,dt=dtmix,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ReverseDiffNoise(), noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,tarray
    ,dt=dtmix,adaptive=false,sensealg=InterpolatingAdjoint(noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-5)

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,tarray
    ,dt=dtmix,adaptive=false,sensealg=InterpolatingAdjoint(noisemixing=true, noise=DiffEqSensitivity.ZygoteNoise()))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-5)

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,tarray
    ,dt=dtmix,adaptive=false,sensealg=InterpolatingAdjoint(noise=false, noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-5)

  @info res_sde_pa

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(soloop,EulerHeun(),dg!,tarray
    ,dt=dtmix,adaptive=false,sensealg=InterpolatingAdjoint(noise=DiffEqSensitivity.ReverseDiffNoise(),noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-5)

  @info res_sde_pa

  res_sde_u0, res_sde_p = adjoint_sensitivities(sol,EulerHeun(),dg!,tarray
      ,dt=dtmix,adaptive=false,sensealg=BacksolveAdjoint(noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-5)

  @info res_sde_p

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,tarray
      ,dt=dtmix,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ZygoteNoise(), noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,tarray
      ,dt=dtmix,adaptive=false,sensealg=BacksolveAdjoint(noise=false,noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,tarray
      ,dt=dtmix,adaptive=false,sensealg=BacksolveAdjoint(noise=DiffEqSensitivity.ReverseDiffNoise(), noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-6)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-6)

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,tarray
      ,dt=dtmix,adaptive=false,sensealg=InterpolatingAdjoint(noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-5) # would pass with 1e-4 but last noise value is off

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,tarray
      ,dt=dtmix,adaptive=false,sensealg=InterpolatingAdjoint(noisemixing=true, noise=DiffEqSensitivity.ZygoteNoise()))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-5)

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,tarray
    ,dt=dtmix,adaptive=false,sensealg=InterpolatingAdjoint(noise=false, noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-5)

  @info res_sde_pa

  res_sde_u0a, res_sde_pa = adjoint_sensitivities(sol,EulerHeun(),dg!,tarray
      ,dt=dtmix,adaptive=false,sensealg=InterpolatingAdjoint(noise=DiffEqSensitivity.ReverseDiffNoise(),noisemixing=true))

  @test isapprox(res_sde_u0a, res_sde_u0, rtol=1e-5)
  @test isapprox(res_sde_pa, res_sde_p, rtol=1e-5)

  @info res_sde_pa

  function GSDE(p)
    Random.seed!(seed)
    tmp_prob = remake(prob,u0=eltype(p).(prob.u0),p=p,
                  tspan=eltype(p).(prob.tspan))
    _sol = solve(tmp_prob,EulerHeun(),dt=dtmix,adaptive=false,saveat=Array(t))
    A = convert(Array,_sol)
    res = g(A,p,nothing)
  end

  res_sde_forward = ForwardDiff.gradient(GSDE,p)

  @test isapprox(res_sde_p', res_sde_forward, rtol=1e-5)

  function GSDE2(u0)
    Random.seed!(seed)
    tmp_prob = remake(prob,u0=u0,p=eltype(p).(prob.p),
                      tspan=eltype(p).(prob.tspan))
    _sol = solve(tmp_prob,EulerHeun(),dt=dtmix,adaptive=false,saveat=Array(t))
    A = convert(Array,_sol)
    res = g(A,p,nothing)
  end

  res_sde_forward = ForwardDiff.gradient(GSDE2,u₀)

  @test isapprox(res_sde_forward, res_sde_u0, rtol=1e-5)

end

@testset "mixing noise inplace/oop tests" begin
  Random.seed!(seed)
  u₀ = [0.75,0.5]
  p = [-1.5,0.05,0.2, 0.01]
  dtmix = tend/1e3

  # Example from Roessler, SIAM J. NUMER. ANAL, 48, 922–952 with d = 2; m = 2
  function f_mixing!(du,u,p,t)
    du[1] = p[1]*u[1] + p[2]*u[2]
    du[2] = p[2]*u[1] + p[1]*u[2]
    nothing
  end

  function g_mixing!(du,u,p,t)
    du[1] = p[3]*u[1] + p[4]*u[2]
    du[2] = p[3]*u[1] + p[4]*u[2]
    nothing
  end

  function f_mixing(u,p,t)
    dx = p[1]*u[1] + p[2]*u[2]
    dy = p[2]*u[1] + p[1]*u[2]
    [dx,dy]
  end

  function g_mixing(u,p,t)
    dx = p[3]*u[1] + p[4]*u[2]
    dy = p[3]*u[1] + p[4]*u[2]
    [dx,dy]
  end

  Random.seed!(seed)
  prob = SDEProblem(f_mixing!,g_mixing!,u₀,trange,p)
  soltsave = collect(trange[1]:dtmix:trange[2])
  sol = solve(prob, EulerHeun(), dt=dtmix, save_noise=true, saveat=soltsave)

  Random.seed!(seed)
  proboop = SDEProblem(f_mixing,g_mixing,u₀,trange,p)
  soloop = solve(proboop,EulerHeun(), dt=dtmix, save_noise=true, saveat=soltsave)

  @test sol.u ≈ soloop.u atol = 1e-14

  # BacksolveAdjoint

  res_sde_u0, res_sde_p = adjoint_sensitivities(soloop,EulerHeun(),dg!,tarray
    ,dt=dtmix,adaptive=false,sensealg=BacksolveAdjoint(noisemixing=true))

  res_sde_u02, res_sde_p2 = adjoint_sensitivities(sol,EulerHeun(),dg!,tarray
    ,dt=dtmix,adaptive=false,sensealg=BacksolveAdjoint(noisemixing=true))

  @test res_sde_u0 ≈ res_sde_u02 atol = 1e-14
  @test res_sde_p ≈ res_sde_p2 atol = 1e-14

  @show res_sde_u0

  adjproboop = SDEAdjointProblem(soloop,BacksolveAdjoint(noisemixing=true),dg!,tarray, nothing)
  adj_soloop = solve(adjproboop,EulerHeun(); dt=dtmix, tstops=soloop.t, adaptive=false)

  @test adj_soloop[end][length(p)+length(u₀)+1:end] == soloop.u[1]
  @test - adj_soloop[end][1:length(u₀)] == res_sde_u0
  @test adj_soloop[end][length(u₀)+1:end-length(u₀)] == res_sde_p'

  adjprob = SDEAdjointProblem(sol,BacksolveAdjoint(noisemixing=true, checkpointing=true),dg!,tarray, nothing)
  adj_sol = solve(adjprob,EulerHeun(); dt=dtmix, adaptive=false,tstops=soloop.t)

  @test adj_soloop[end] ≈  adj_sol[end]  rtol=1e-15

  adjprob = SDEAdjointProblem(sol,BacksolveAdjoint(noisemixing=true, checkpointing=false),dg!,tarray, nothing)
  adj_sol = solve(adjprob,EulerHeun(); dt=dtmix, adaptive=false,tstops=soloop.t)

  @test adj_soloop[end] ≈  adj_sol[end]  rtol=1e-8

  # InterpolatingAdjoint

  res_sde_u0, res_sde_p = adjoint_sensitivities(soloop,EulerHeun(),dg!,tarray
    ,dt=dtmix,adaptive=false,sensealg=InterpolatingAdjoint(noisemixing=true))

  res_sde_u02, res_sde_p2 = adjoint_sensitivities(sol,EulerHeun(),dg!,tarray
    ,dt=dtmix,adaptive=false,sensealg=InterpolatingAdjoint(noisemixing=true))

  @test res_sde_u0 ≈ res_sde_u02 atol = 1e-8
  @test res_sde_p ≈ res_sde_p2 atol = 5e-8

  @show res_sde_u0

  adjproboop = SDEAdjointProblem(soloop,InterpolatingAdjoint(noisemixing=true),dg!,tarray, nothing)
  adj_soloop = solve(adjproboop,EulerHeun(); dt=dtmix, tstops=soloop.t, adaptive=false)

  @test - adj_soloop[end][1:length(u₀)] ≈ res_sde_u0  atol = 1e-14
  @test adj_soloop[end][length(u₀)+1:end] ≈ res_sde_p' atol = 1e-14

  adjprob = SDEAdjointProblem(sol,InterpolatingAdjoint(noisemixing=true, checkpointing=true),dg!,tarray, nothing)
  adj_sol = solve(adjprob,EulerHeun(); dt=dtmix, adaptive=false,tstops=soloop.t)

  @test adj_soloop[end] ≈ adj_sol[end]  rtol=1e-8

  adjprob = SDEAdjointProblem(sol,InterpolatingAdjoint(noisemixing=true, checkpointing=false),dg!,tarray, nothing)
  adj_sol = solve(adjprob,EulerHeun(); dt=dtmix, adaptive=false,tstops=soloop.t)

  @test adj_soloop[end] ≈  adj_sol[end]  rtol=1e-8
end

@testset "mutating non-diagonal noise" begin
  a!(du,u,_p,t) = (du .= -u)
  a(u,_p,t) = -u

  function b!(du,u,_p,t)
    KR, KI = _p[1:2]

    du[1,1] = KR
    du[2,1] = KI
  end

  function b(u,_p,t)
    KR, KI = _p[1:2]

    [ KR zero(KR)
      KI zero(KR) ]
  end

  p = [1.,0.]

  prob! = SDEProblem{true}(a!,b!,[0.,0.],(0.0,0.1),p,noise_rate_prototype=eltype(p).(zeros(2,2)))
  prob = SDEProblem{false}(a,b,[0.,0.],(0.0,0.1),p,noise_rate_prototype=eltype(p).(zeros(2,2)))

  function loss(p;SDEprob=prob,sensealg=BacksolveAdjoint())
    _prob = remake(SDEprob,p=p)
    sol = solve(_prob, EulerHeun(), dt=1e-5, sensealg=sensealg)
    return sum(Array(sol))
  end

  function compute_dp(p, SDEprob, sensealg)
    Random.seed!(seed)
    Zygote.gradient(p->loss(p,SDEprob=SDEprob,sensealg=sensealg), p)[1]
  end

  # test mutating against non-mutating

  # non-mutating

  dp1 = compute_dp(p, prob, ForwardDiffSensitivity())
  dp2 = compute_dp(p, prob, BacksolveAdjoint())
  dp3 = compute_dp(p, prob, InterpolatingAdjoint())

  @show dp1 dp2 dp3

  # different vjp choice
  _dp2 = compute_dp(p, prob, BacksolveAdjoint(autojacvec=ReverseDiffVJP()))
  @test dp2 ≈ _dp2 rtol=1e-8
  _dp3 = compute_dp(p, prob, InterpolatingAdjoint(autojacvec=ReverseDiffVJP()))
  @test dp3 ≈ _dp3 rtol=1e-8

  # mutating
  _dp1 = compute_dp(p, prob!, ForwardDiffSensitivity())
  _dp2 = compute_dp(p, prob!, BacksolveAdjoint())
  _dp3 = compute_dp(p, prob!, InterpolatingAdjoint())

  @test dp1 ≈ _dp1 rtol=1e-8
  @test dp2 ≈ _dp2 rtol=1e-8
  @test dp3 ≈ _dp3 rtol=1e-8
end


# Test code for identifiying inside/outside points

using Base.Test

push!(LOAD_PATH, string(pwd(),"/src"))
using PyPlot
using ModifiedStokesSolver
import AnalyticDomains

@testset "InsideOutside" begin
    # Discretize
    numpanels = 10
    panelorder = 16
    curve = AnalyticDomains.starfish(n_arms = 7, amplitude=0.3)
    dcurve = discretize(curve, numpanels, panelorder)

    # Easy test
    rgrid = linspace(0.01, 0.7-1e-8, 50)
    tgrid = linspace(0, 2*pi, 500)
    R, T = ndgrid(rgrid, tgrid)
    X = R.*cos.(T)
    Y = R.*sin.(T)
    zt = [X[:] Y[:]]'
    interior = interior_points(dcurve, zt)
    exterior = .!interior
    @test all(interior)

    # clf()
    # plot(zt[1,interior], zt[2,interior],".b")
    # plot(zt[1,exterior], zt[2,exterior],".r")
    # plot(dcurve.points[1,:], dcurve.points[2,:], "-k")
    # axis("image")

    # Hard test
    rgrid = linspace(0, 2*pi, 1000)
    igrid = linspace(-5e-2, 5e-2, 100)
    R, I = ndgrid(rgrid, igrid)
    Z = @. curve.tau(R + 1im*I)
    X = real(Z)
    Y = imag(Z)

    zt = [X[:] Y[:]]'
    interior = interior_points(dcurve, zt)
    exterior = .!interior

    ref_interior = I[:] .> 0
    ref_exterior = I[:] .< 0

    # clf()
    # plot(zt[1,interior], zt[2,interior],".b")
    # plot(zt[1,exterior], zt[2,exterior],".r")
    # plot(dcurve.points[1,:], dcurve.points[2,:], "-k")
    # axis("image")

    @test !any(ref_interior .& exterior)
    @test !any(ref_exterior .& interior)
end


include("../examples/lorenz.jl")
using Test
#using PyPlot
function test_attractor()
	s = [10., 28., 8.0/3]
	u = rand(3)
	n_runup = 20000
	for i = 1:n_runup
		u = next(u,s)
	end
	@show u
	n = 500000
	u_trj = step(u, s, n)
	u_trj = u_trj[:,1:10:end]
	x, y, z = view(u_trj,1,:), view(u_trj,2,:), view(u_trj,3,:)
	fig, ax = subplots(1,1)
	ax.plot(x,y,".",ms=2)
	ax.set_xlabel("x",fontsize=36)
	ax.set_ylabel("y",fontsize=36)
	ax.tick_params(axis="both",labelsize=36)
	fig, ax = subplots(1,1)
	ax.plot(x,z,".",ms=2)
	ax.set_xlabel("x",fontsize=36)
	ax.set_ylabel("z",fontsize=36)
	ax.tick_params(axis="both",labelsize=36)
end
function test_dstep()
	eps = 1.e-8
	s = [10., 28., 8.0/3]
	u = rand(3)
	u1px = u .+ eps*[1,0,0]
	u1mx = u .- eps*[1,0,0]
	u1py = u .+ eps*[0,1,0]
	u1my = u .- eps*[0,1,0]
	u1pz = u .+ eps*[0,0,1]
	u1mz = u .- eps*[0,0,1]

	du_dx_fd = (next(u1px,s) - next(u1mx,s))/(2*eps)
	du_dy_fd = (next(u1py,s) - next(u1my,s))/(2*eps)
	du_dz_fd = (next(u1pz,s) - next(u1mz,s))/(2*eps)

	du = dstep(u, s)
	du_dx = du[:,1]
	du_dy = du[:,2]
	du_dz = du[:,3]

	@show abs.(du_dx_fd .- du_dx)
	@show abs.(du_dy_fd .- du_dy)
	@show abs.(du_dz_fd .- du_dz)
	@test du_dx_fd ≈ du_dx atol=1.e-8
	@test du_dy_fd ≈ du_dy atol=1.e-8
	@test du_dz_fd ≈ du_dz atol=1.e-8
end
function test_dstep_arr()
	eps = 1.e-8
	s = [10., 28., 8.0/3]
	u = rand(3,1)
	u1px = u[:,1] .+ eps*[1,0,0]
	u1mx = u[:,1] .- eps*[1,0,0]
	u1py = u[:,1] .+ eps*[0,1,0]
	u1my = u[:,1] .- eps*[0,1,0]
	u1pz = u[:,1] .+ eps*[0,0,1]
	u1mz = u[:,1] .- eps*[0,0,1]

	du_dx_fd = (next(u1px,s) - next(u1mx,s))/(2*eps)
	du_dy_fd = (next(u1py,s) - next(u1my,s))/(2*eps)
	du_dz_fd = (next(u1pz,s) - next(u1mz,s))/(2*eps)

	du = dstep(u, s)
	@show du
	du_dx = du[:,1,1]
	du_dy = du[:,2,1]
	du_dz = du[:,3,1]

	@test du_dx_fd ≈ du_dx atol=1.e-8
	@test du_dy_fd ≈ du_dy atol=1.e-8
	@test du_dz_fd ≈ du_dz atol=1.e-8
end
function test_d2step()
	eps = 1.e-6
	s = [10., 28., 8.0/3]
	u = rand(3)
	u1px = u .+ eps*[1,0,0]
	u1mx = u .- eps*[1,0,0]
	u1py = u .+ eps*[0,1,0]
	u1my = u .- eps*[0,1,0]
	u1pz = u .+ eps*[0,0,1]
	u1mz = u .- eps*[0,0,1]

	ddu_dx_fd = (dstep(u1px,s) - dstep(u1mx,s))/(2*eps)
	ddu_dy_fd = (dstep(u1py,s) - dstep(u1my,s))/(2*eps)
	ddu_dz_fd = (dstep(u1pz,s) - dstep(u1mz,s))/(2*eps)

	ddu = d2step(u, s)
	@show ddu[:,:,1] .- ddu_dx_fd
	@show ddu[:,:,2] .- ddu_dy_fd
	@show ddu[:,:,3] .- ddu_dz_fd

	@test ddu[:,:,1] ≈ ddu_dx_fd atol=1.e-8
	@test ddu[:,:,2] ≈ ddu_dy_fd atol=1.e-8
	@test ddu[:,:,3] ≈ ddu_dz_fd atol=1.e-8

end
function test_d2step_arr()
	eps = 1.e-6
	s = [10., 28., 8.0/3]
	u = rand(3,1)
	u1px = u[:,1] .+ eps*[1,0,0]
	u1mx = u[:,1] .- eps*[1,0,0]
	u1py = u[:,1] .+ eps*[0,1,0]
	u1my = u[:,1] .- eps*[0,1,0]
	u1pz = u[:,1] .+ eps*[0,0,1]
	u1mz = u[:,1] .- eps*[0,0,1]

	ddu_dx_fd = (dstep(u1px,s) - dstep(u1mx,s))/(2*eps)
	ddu_dy_fd = (dstep(u1py,s) - dstep(u1my,s))/(2*eps)
	ddu_dz_fd = (dstep(u1pz,s) - dstep(u1mz,s))/(2*eps)

	ddu = d2step(u, s)
	@show ddu[:,:,1,1] .- ddu_dx_fd
	@show ddu[:,:,2,1] .- ddu_dy_fd
	@show ddu[:,:,3,1] .- ddu_dz_fd

	@test ddu[:,:,1,1] ≈ ddu_dx_fd atol=1.e-8
	@test ddu[:,:,2,1] ≈ ddu_dy_fd atol=1.e-8
	@test ddu[:,:,3,1] ≈ ddu_dz_fd atol=1.e-8

end
function test_pert()
	eps = 1.e-6
	s = [10., 28., 8.0/3]
	u = rand(3)
	sp = s .+ [0., eps, 0.]
	sm = s .- [0., eps, 0.]

	pp_fd = (next(u, sp) - next(u, sm))/(2*eps)
	pp = pert(u, s)

	@show pp_fd, pp
	@test pp_fd ≈ pp atol=1.e-8
end
function test_pert_arr()
	eps = 1.e-6
	s = [10., 28., 8.0/3]
	u = rand(3)
	sp = s .+ [0., eps, 0.]
	sm = s .- [0., eps, 0.]

	pp_fd = (next(u,sp) - next(u, sm))/(2*eps)
	u = reshape(u, 3, 1)
	pp = pert(u, s)[:,1]

	@show pp_fd, pp
	@test pp_fd ≈ pp atol=1.e-8
end
function test_dpert_next()
	eps = 1.e-6
	s = [10., 28., 8.0/3]
	u = rand(3)
	u1px = u .+ eps*[1,0,0]
	u1mx = u .- eps*[1,0,0]
	u1py = u .+ eps*[0,1,0]
	u1my = u .- eps*[0,1,0]
	u1pz = u .+ eps*[0,0,1]
	u1mz = u .- eps*[0,0,1]

	dpert_dx_fd = (pert(u1px,s) - pert(u1mx,s))/(2*eps)
	dpert_dy_fd = (pert(u1py,s) - pert(u1my,s))/(2*eps)
	dpert_dz_fd = (pert(u1pz,s) - pert(u1mz,s))/(2*eps)
	dpp = dpert_next(u, s)

	@show dpp[:,1] .- dpert_dx_fd
	@show dpp[:,2] .- dpert_dy_fd
	@show dpp[:,3] .- dpert_dz_fd

	@test dpp[:,1] ≈ dpert_dx_fd atol=1.e-8
	@test dpp[:,2] ≈ dpert_dy_fd atol=1.e-8
	@test dpp[:,3] ≈ dpert_dz_fd atol=1.e-8
end
function test_dpert_next_arr()
	eps = 1.e-6
	s = [10., 28., 8.0/3]
	u = rand(3,1)
	u1px = u[:,1] .+ eps*[1,0,0]
	u1mx = u[:,1] .- eps*[1,0,0]
	u1py = u[:,1] .+ eps*[0,1,0]
	u1my = u[:,1] .- eps*[0,1,0]
	u1pz = u[:,1] .+ eps*[0,0,1]
	u1mz = u[:,1] .- eps*[0,0,1]

	dpert_dx_fd = (pert(u1px,s) - pert(u1mx,s))/(2*eps)
	dpert_dy_fd = (pert(u1py,s) - pert(u1my,s))/(2*eps)
	dpert_dz_fd = (pert(u1pz,s) - pert(u1mz,s))/(2*eps)
	dpp = dpert_next(u, s)

	@show dpp[:,1,1] .- dpert_dx_fd
	@show dpp[:,2,1] .- dpert_dy_fd
	@show dpp[:,3,1] .- dpert_dz_fd

	@test dpp[:,1,1] ≈ dpert_dx_fd atol=1.e-8
	@test dpp[:,2,1] ≈ dpert_dy_fd atol=1.e-8
	@test dpp[:,3,1] ≈ dpert_dz_fd atol=1.e-8
end
function test_dmag_flow()
	eps = 1.e-6
	s = [10., 28., 8.0/3]
	u = rand(3)
	dmag_dx = (mag_flow(u + eps*[1.0,0,0],s) -
			   mag_flow(u - eps*[1.0,0,0],s))/(2*eps)
	dmag_dy = (mag_flow(u + eps*[0,1.0,0],s) -
			   mag_flow(u - eps*[0,1.0,0],s))/(2*eps)
	dmag_dz = (mag_flow(u + eps*[0,0,1.0],s) -
			   mag_flow(u - eps*[0,0,1.0],s))/(2*eps)

	dmag = dmag_flow(u, s)
	@test dmag[1] ≈ dmag_dx atol=1.e-8
	@test dmag[2] ≈ dmag_dy atol=1.e-8
	@test dmag[3] ≈ dmag_dz atol=1.e-8
end
function test_dmag_flow_arr()
	eps = 1.e-6
	s = [10., 28., 8.0/3]
	u = rand(3,1)
	dmag_dx = (mag_flow(u .+ eps*reshape([1.0,0,0],3,1),s) -
			   mag_flow(u .- eps*reshape([1.0,0,0],3,1),s))/(2*eps)
	dmag_dy = (mag_flow(u .+ eps*reshape([0,1.0,0],3,1),s) -
			   mag_flow(u .- eps*reshape([0,1.0,0],3,1),s))/(2*eps)
	dmag_dz = (mag_flow(u .+ eps*reshape([0,0,1.0],3,1),s) -
			   mag_flow(u .- eps*reshape([0,0,1.0],3,1),s))/(2*eps)

	dmag = dmag_flow(u, s)
	@test dmag[1,:] ≈ dmag_dx atol=1.e-8
	@test dmag[2,:] ≈ dmag_dy atol=1.e-8
	@test dmag[3,:] ≈ dmag_dz atol=1.e-8
end
function test_dunit_flow()
	eps = 1.e-6
	s = [10., 28., 8/3]
	u = rand(3)
	upx = u .+ eps*[1.0,0,0]
	umx = u .- eps*[1.0,0,0]
	upy = u .+ eps*[0,1.0,0]
	umy = u .- eps*[0,1.0,0]
	upz = u .- eps*[0,0,1.0]
	umz = u .- eps*[0,0,1.0]
	dfdx = (flow(upx,s)/mag_flow(upx,s) .-
			flow(umx,s)/mag_flow(umx,s))/(2*eps)
	dfdy = (flow(upy,s)/mag_flow(upy,s) .-
			flow(umy,s)/mag_flow(umy,s))/(2*eps)
	dfdz = (flow(upz,s)/mag_flow(upz,s) .-
			flow(umz,s)/mag_flow(umz,s))/(2*eps)
	df = dunit_flow(u, s)
	@show dfdx, df[:,1]

	@test df[:,1] ≈ dfdx atol=1.e-8
	@test df[:,2] ≈ dfdy atol=1.e-8
	@test df[:,3] ≈ dfdz atol=1.e-8

end


function _likelihood!(dR, D, L, dt, er)
    d = Diffusion.(D, dt, er)
    @inbounds for i = 1:track_num
        @inbounds for t = 1:track_length[i]+1
            @inbounds for s = 1:K
                L[t, s, i] = pdf.(d[s], ifelse(dR[t, s, i] < 1e-4, 1e-4, dR[t, s, i]))
            end
        end
    end
end

# function likelihood!(
#     observations::AbstractArray,
#     L::AbstractArray,
#     D,
#     dt,
#     er,
#     track_length,
# )
#     @argcheck size(observations) == size(L)
#     K, N = size(L, 2), size(L, 3)
#     d = Diffusion.(D, dt, er)
#     @inbounds for i = 1:N
#         @inbounds for t = 1:track_length[i]+1
#             @inbounds for j = 1:K
#                 L[t, j, i] =
#                     pdf.(
#                         d[j],
#                         ifelse(observations[t, j, i] < 1e-4, 1e-4, observations[t, j, i]),
#                     )
#             end
#         end
#     end
# end

function loglikelihood!(observations, L, d)
    @argcheck size(observations, 1) == size(L, 1)
    @argcheck size(observations, 2) == size(L, 3)
    K, N = size(L, 2), size(L, 3)
    fill!(L, nothing)
    @inbounds for n = 1:N
        T = length(filter(!isnothing, observations[:, n]))
        for t = 1:T
            for k = 1:K
                L[t, k, n] = pdf(d[k], observations[t, n])
            end
        end
    end
end


module LE11NAFEMS_examples
using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsDeforLinear.AlgoDeforLinearModule

function LE11NAFEMS_Q8_algo()
    # NAFEMS LE11 benchmark with Q8 elements.
    # # This is a test recommended by the National Agency for Finite Element
    # # Methods and Standards (U.K.): Test LE11 from NAFEMS Publication TNSB,
    # # Rev. 3, “The Standard NAFEMS Benchmarks,” October 1990.
    # #
    # # Target solution: Direct stress,   =  –105 MPa at point A.
    #function  LE11NAFEMS()
    # Parameters:
    Ea =  210000*phun("MEGA*Pa")
    nua =  0.3;
    alphaa = 2.3e-4;              # thermal expansion coefficient
    sigmaA = -105*phun("MEGA*Pa")
    nref =  1;                        # how many times should we refine the mesh?
    X = [1.     0.;#A
    1.4    0.;#B
    0.995184726672197   0.098017140329561;
    1.393258617341076 0.137223996461385;
    0.980785  0.195090;#
    1.37309939 0.27312645;
    0.956940335732209   0.290284677254462
    1.339716470025092 0.406398548156247
    0.9238795  0.38268;#C
    1.2124  0.7;#D
    0.7071  0.7071;#E
    1.1062  1.045;#F
    0.7071  (0.7071+1.79)/2;#(E+H)/2
    1.      1.39;#G
    0.7071  1.79;#H
    1.      1.79;#I
    ]*phun("M")
    tolerance  = 1.e-6*phun("M")
    ##
    # Note that the material object needs to be created with the proper
    # model-dimension reduction in mind.  In this case that is the axial symmetry
    # assumption.
    MR  =  DeforModelRed2DAxisymm

    fens = FENodeSet(X);
    fes = FESetQ4([1 2 4 3; 3 4 6 5; 5 6 8 7; 7 8 10 9; 9 10 12 11; 11 12 14 13; 13 14 16 15]);
    for ref = 1:nref
        fens,fes = Q4refine(fens,fes);
        list = selectnode(fens,distance = 1.0+0.1/2^nref, from = [0. 0.], inflate = tolerance);
        fens.xyz[list,:] =  FinEtools.MeshUtilModule.ontosphere(fens.xyz[list,:],1.0);
    end
    fens,fes = Q4toQ8(fens,fes)
    list = selectnode(fens,distance = 1.0+0.1/2^nref, from = [0. 0.], inflate = tolerance);
    fens.xyz[list,:] =  FinEtools.MeshUtilModule.ontosphere(fens.xyz[list,:],1.0);

    # EBC's
    l1 = selectnode(fens,box = [-Inf Inf 0 0],inflate = tolerance)
    e1 = FDataDict("node_list"=>l1, "component"=>2, "displacement"=>x -> 0.0)
    l1 = selectnode(fens,box = [-Inf Inf 1.79  1.79],inflate = tolerance)
    e2 = FDataDict("node_list"=>l1, "component"=>2, "displacement"=>x -> 0.0)

    # Temperature field
    dtemp = FDataDict("temperature"=>x -> x[1] + x[2])

    # Property and material
    material = MatDeforElastIso(MR, 0.0, Ea, nua, alphaa)

    femm  =  FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(2, 3), true), material)

    # Make region 1
    region = FDataDict("femm"=>femm)
    # Make model data
    modeldata =  FDataDict(
    "fens"=> fens, "regions"=>  [region],
    "essential_bcs"=>[e1, e2], "temperature_change"=>dtemp)

    # Call the solver
    modeldata = AlgoDeforLinearModule.linearstatics(modeldata)
    geom = modeldata["geom"]
    u = modeldata["u"]
    dT = modeldata["temp"]

    nA  = selectnode(fens,box = FFlt[1.0  1.0 0.0 0.0], inflate = tolerance);

    fld =  fieldfromintegpoints(femm, geom, u, dT, :Cauchy, 2)

    File  =   "LE11NAFEMS_Q8_sigmay.vtk"
    vtkexportmesh(File, fens, fes; scalars = [("sigmay", fld.values)],
    vectors = [("u", u.values)])
    println("range of  sigmay = $((minimum(fld.values), maximum(fld.values)))")
    @async run(`"paraview.exe" $File`)

    sA  =  fld.values[nA]/phun("MEGA*Pa")
    sAn  =  fld.values[nA]/sigmaA
    println("Stress at point A: $(sA) i. e.  $( sAn*100  )% of reference value")

    fen2fe  = FENodeToFEMap(fes.conn, nnodes(geom))
    function inspector(idat, elnum, conn, xe,  out,  xq)
        println("loc = $(  xq  ) : $(  transpose(out)/phun("MEGA*Pa")  )")
        return idat
    end

    inspectintegpoints(femm, geom, u, dT,  fen2fe.map[nA[1]],
    inspector, []; quantity = :Cauchy)

end # LE11NAFEMS_Q8_algo

function LE11NAFEMS_Q8_algo2()
    # NAFEMS LE11 benchmark with Q8 elements.
    # # This is a test recommended by the National Agency for Finite Element
    # # Methods and Standards (U.K.): Test LE11 from NAFEMS Publication TNSB,
    # # Rev. 3, “The Standard NAFEMS Benchmarks,” October 1990.
    # #
    # # Target solution: Direct stress,   =  –105 MPa at point A.
    #function  LE11NAFEMS()
    # Parameters:
    Ea =  210000*phun("MEGA*Pa")
    nua =  0.3;
    alphaa = 2.3e-4;              # thermal expansion coefficient
    sigmaA = -105*phun("MEGA*Pa")
    nref =  1;                        # how many times should we refine the mesh?
    X = [1.     0.;#A
    1.4    0.;#B
    0.995184726672197   0.098017140329561;
    1.393258617341076 0.137223996461385;
    0.980785  0.195090;#
    1.37309939 0.27312645;
    0.956940335732209   0.290284677254462
    1.339716470025092 0.406398548156247
    0.9238795  0.38268;#C
    1.2124  0.7;#D
    0.7071  0.7071;#E
    1.1062  1.045;#F
    0.7071  (0.7071+1.79)/2;#(E+H)/2
    1.      1.39;#G
    0.7071  1.79;#H
    1.      1.79;#I
    ]*phun("M")
    tolerance  = 1.e-6*phun("M")
    ##
    # Note that the material object needs to be created with the proper
    # model-dimension reduction in mind.  In this case that is the axial symmetry
    # assumption.
    MR  =  DeforModelRed2DAxisymm

    fens = FENodeSet(X);
    fes = FESetQ4([1 2 4 3; 3 4 6 5; 5 6 8 7; 7 8 10 9; 9 10 12 11; 11 12 14 13; 13 14 16 15]);
    for ref = 1:nref
        fens,fes = Q4refine(fens,fes);
        list = selectnode(fens,distance = 1.0+0.1/2^nref, from = [0. 0.], inflate = tolerance);
        fens.xyz[list,:] =  FinEtools.MeshUtilModule.ontosphere(fens.xyz[list,:],1.0);
    end
    fens,fes = Q4toQ8(fens,fes)
    list = selectnode(fens,distance = 1.0+0.1/2^nref, from = [0. 0.], inflate = tolerance);
    fens.xyz[list,:] =  FinEtools.MeshUtilModule.ontosphere(fens.xyz[list,:],1.0);

    # EBC's
    l1 = selectnode(fens,box = [-Inf Inf 0 0],inflate = tolerance)
    e1 = FDataDict("node_list"=>l1, "component"=>2, "displacement"=>x -> 0.0)
    l1 = selectnode(fens,box = [-Inf Inf 1.79  1.79],inflate = tolerance)
    e2 = FDataDict("node_list"=>l1, "component"=>2, "displacement"=>x -> 0.0)

    # Temperature field
    dtemp = FDataDict("temperature"=>x -> x[1] + x[2])

    # Property and material
    material = MatDeforElastIso(MR, 0.0, Ea, nua, alphaa)

    femm  =  FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(2, 3), true), material)

    # Make region 1
    region = FDataDict("femm"=>femm)
    # Make model data
    modeldata =  FDataDict(
    "fens"=> fens, "regions"=>  [region],
    "essential_bcs"=>[e1, e2], "temperature_change"=>dtemp)

    # Call the solver
    modeldata = AlgoDeforLinearModule.linearstatics(modeldata)
    geom = modeldata["geom"]
    u = modeldata["u"]
    dT = modeldata["temp"]

    modeldata["postprocessing"] = FDataDict("boundary_only"=> true,
    "file"=>"LE11NAFEMS_Q8_deformation.vtk")
    modeldata = AlgoDeforLinearModule.exportdeformation(modeldata)
    @async run(`"paraview.exe" $(modeldata["postprocessing"]["exported"][1]["file"])`)

    nA  = selectnode(fens,box = FFlt[1.0  1.0 0.0 0.0], inflate = tolerance);

    modeldata["postprocessing"] = FDataDict("boundary_only"=> true,
    "file"=>"LE11NAFEMS_Q8_sigmay.vtk", "quantity"=>:Cauchy,
    "component"=>2)
    modeldata = AlgoDeforLinearModule.exportstress(modeldata)
    modeldata["postprocessing"] = FDataDict("boundary_only"=> false,
    "file"=>"LE11NAFEMS_Q8_sigmay.vtk", "quantity"=>:Cauchy,
    "component"=>2)
    modeldata = AlgoDeforLinearModule.exportstress(modeldata)
    @async run(`"paraview.exe" $(modeldata["postprocessing"]["exported"][1]["file"])`)
    fld =  modeldata["postprocessing"]["exported"][1]["field"]

    sA  =  fld.values[nA]/phun("MEGA*Pa")
    sAn  =  fld.values[nA]/sigmaA
    println("Stress at point A: $(sA) i. e.  $( sAn*100  )% of reference value")

    # Loop over only those elements that share the node nA
    fen2fe  = FENodeToFEMap(fes.conn, nnodes(geom))
    function inspector(idat, elnum, conn, xe,  out,  xq)
        println("loc = $(  xq  ) : $(  transpose(out)/phun("MEGA*Pa")  )")
        return idat
    end

    inspectintegpoints(femm, geom, u, dT,  fen2fe.map[nA[1]],
    inspector, []; quantity = :Cauchy)

    modeldata["postprocessing"] = FDataDict("boundary_only"=> false,
    "file"=>"LE11NAFEMS_Q8_sigmay_ew.vtk", "quantity"=>:Cauchy,
    "component"=>2)
    modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)
    @async run(`"paraview.exe" $(modeldata["postprocessing"]["exported"][1]["file"])`)

end # LE11NAFEMS_Q8_algo2

function LE11NAFEMS_Q8_export_stress()
    # NAFEMS LE11 benchmark with Q8 elements.
    # # This is a test recommended by the National Agency for Finite Element
    # # Methods and Standards (U.K.): Test LE11 from NAFEMS Publication TNSB,
    # # Rev. 3, “The Standard NAFEMS Benchmarks,” October 1990.
    # #
    # # Target solution: Direct stress,   =  –105 MPa at point A.
    #function  LE11NAFEMS()
    # Parameters:
    Ea =  210000*phun("MEGA*Pa")
    nua =  0.3;
    alphaa = 2.3e-4;              # thermal expansion coefficient
    sigmaA = -105*phun("MEGA*Pa")
    nref = 2;                        # how many times should we refine the mesh?
    X = [1.     0.;#A
    1.4    0.;#B
    0.995184726672197   0.098017140329561;
    1.393258617341076 0.137223996461385;
    0.980785  0.195090;#
    1.37309939 0.27312645;
    0.956940335732209   0.290284677254462
    1.339716470025092 0.406398548156247
    0.9238795  0.38268;#C
    1.2124  0.7;#D
    0.7071  0.7071;#E
    1.1062  1.045;#F
    0.7071  (0.7071+1.79)/2;#(E+H)/2
    1.      1.39;#G
    0.7071  1.79;#H
    1.      1.79;#I
    ]*phun("M")
    tolerance  = 1.e-6*phun("M")
    ##
    # Note that the material object needs to be created with the proper
    # model-dimension reduction in mind.  In this case that is the axial symmetry
    # assumption.
    MR  =  DeforModelRed2DAxisymm

    fens = FENodeSet(X);
    fes = FESetQ4([1 2 4 3; 3 4 6 5; 5 6 8 7; 7 8 10 9; 9 10 12 11; 11 12 14 13; 13 14 16 15]);
    for ref = 1:nref
        fens,fes = Q4refine(fens,fes);
        list = selectnode(fens,distance = 1.0+0.1/2^nref, from = [0. 0.], inflate = tolerance);
        fens.xyz[list,:] =  FinEtools.MeshUtilModule.ontosphere(fens.xyz[list,:],1.0);
    end
    fens,fes = Q4toQ8(fens,fes)
    list = selectnode(fens,distance = 1.0+0.1/2^nref, from = [0. 0.], inflate = tolerance);
    fens.xyz[list,:] =  FinEtools.MeshUtilModule.ontosphere(fens.xyz[list,:],1.0);

    #     File  =   "mesh.vtk"
    # vtkexportmesh(File, fens, fes)

    # now we create the geometry and displacement fields
    geom  =  NodalField(fens.xyz)
    u  =  NodalField(zeros(size(fens.xyz,1),2)) # displacement field

    # Apply EBC's
    lbottom = selectnode(fens,box = [-Inf Inf 0 0],inflate = tolerance)
    setebc!(u, lbottom, true, 2, 00.0)
    ltop = selectnode(fens,box = [-Inf Inf 1.79  1.79],inflate = tolerance)
    setebc!(u, ltop, true, 2, 00.0)
    applyebc!(u)
    numberdofs!(u)

    # Temperature field
    dT  = NodalField(reshape(fens.xyz[:,1]+fens.xyz[:,2],size(fens.xyz,1),1));

    # Property and material
    material = MatDeforElastIso(MR, 0.0, Ea, nua, alphaa)

    femm  =  FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(2, 3), true), material)

    K  = stiffness(femm, geom, u)
    F  =  thermalstrainloads(femm, geom, u, dT)
    #K = cholesky(K)
    U =   K\F
    scattersysvec!(u, U[:])

    nA  = selectnode(fens,box = FFlt[1.0  1.0 0.0 0.0], inflate = tolerance);

    fld =  fieldfromintegpoints(femm, geom, u, dT, :Cauchy, 2)

    File  =   "LE11NAFEMS_Q8_sigmay.vtk"
    vtkexportmesh(File, fens, fes; scalars = [("sigmay", fld.values)],
    vectors = [("u", u.values)])
    println("range of  sigmay = $((minimum(fld.values), maximum(fld.values)))")
    @async run(`"paraview.exe" $File`)

    sA  =  fld.values[nA]/phun("MEGA*Pa")
    sAn  =  fld.values[nA]/sigmaA
    println("Stress at point A: $(sA) i. e.  $( sAn*100  )% of reference value")

    fen2fe  = FENodeToFEMap(fes.conn, nnodes(geom))
    function inspector(idat, elnum, conn, xe,  out,  xq)
        # println("loc = $(  xq  ) : $(  transpose(out)/phun("MEGA*Pa")  )")
        return idat
    end

    inspectintegpoints(femm, geom, u, dT,  fen2fe.map[nA[1]],
    inspector, []; quantity = :Cauchy)

    fld =  fieldfromintegpoints(femm, geom, u, dT, :Pressure, 1)
    File  =   "LE11NAFEMS_Q8_pressure.vtk"
    vtkexportmesh(File, fens, fes; scalars = [("pressure", fld.values)],
    vectors = [("u", u.values)])
    println("range of  pressure = $((minimum(fld.values), maximum(fld.values)))")
    @async run(`"paraview.exe" $File`)

    fld =  fieldfromintegpoints(femm, geom, u, dT, :vm, 1)
    File  =   "LE11NAFEMS_Q8_vm.vtk"
    vtkexportmesh(File, fens, fes; scalars = [("pressure", fld.values)],
    vectors = [("u", u.values)])
    println("range of von Mises = $((minimum(fld.values), maximum(fld.values)))")
    @async run(`"paraview.exe" $File`)

    AE = AbaqusExporter("LE11NAFEMS_Q8_export_stress");
    HEADING(AE, "NAFEMS LE11 benchmark with Q8 elements.");
    COMMENT(AE, "sigmaA = -105 MPa ");
    PART(AE, "part1");
    END_PART(AE);
    ASSEMBLY(AE, "ASSEM1");
    INSTANCE(AE, "INSTNC1", "PART1");
    NODE(AE, fens.xyz);
    ELEMENT(AE, "cax8", "AllElements", 1, connasarray(fes))
    NSET_NSET(AE, "ltop", ltop)
    NSET_NSET(AE, "lbottom", lbottom)
    ORIENTATION(AE, "GlobalOrientation", vec([1. 0 0]), vec([0 1. 0]));
    SOLID_SECTION(AE, "elasticity", "GlobalOrientation", "AllElements");
    SURFACE_SECTION(AE, "TractionElements")
    END_INSTANCE(AE);
    END_ASSEMBLY(AE);
    MATERIAL(AE, "elasticity")
    ELASTIC(AE, Ea, nua)
    EXPANSION(AE, alphaa)
    STEP_PERTURBATION_STATIC(AE)
    BOUNDARY(AE, "ASSEM1.INSTNC1.ltop", 2)
    BOUNDARY(AE, "ASSEM1.INSTNC1.lbottom", 2)
    TEMPERATURE(AE, "ASSEM1.INSTNC1.", collect(1:count(fens)), vec(dT.values))
    END_STEP(AE)
    close(AE)
end # LE11NAFEMS_Q8_export_stress

function LE11NAFEMS_Q8_stress()
    # NAFEMS LE11 benchmark with Q8 elements.
    # # This is a test recommended by the National Agency for Finite Element
    # # Methods and Standards (U.K.): Test LE11 from NAFEMS Publication TNSB,
    # # Rev. 3, “The Standard NAFEMS Benchmarks,” October 1990.
    # #
    # # Target solution: Direct stress,   =  –105 MPa at point A.
    #function  LE11NAFEMS()
    # Parameters:
    Ea =  210000*phun("MEGA*Pa")
    nua =  0.3;
    alphaa = 2.3e-4;              # thermal expansion coefficient
    sigmaA = -105*phun("MEGA*Pa")
    nref =  1;                        # how many times should we refine the mesh?
    X = [1.     0.;#A
    1.4    0.;#B
    0.995184726672197   0.098017140329561;
    1.393258617341076 0.137223996461385;
    0.980785  0.195090;#
    1.37309939 0.27312645;
    0.956940335732209   0.290284677254462
    1.339716470025092 0.406398548156247
    0.9238795  0.38268;#C
    1.2124  0.7;#D
    0.7071  0.7071;#E
    1.1062  1.045;#F
    0.7071  (0.7071+1.79)/2;#(E+H)/2
    1.      1.39;#G
    0.7071  1.79;#H
    1.      1.79;#I
    ]*phun("M")
    tolerance  = 1.e-6*phun("M")
    ##
    # Note that the material object needs to be created with the proper
    # model-dimension reduction in mind.  In this case that is the axial symmetry
    # assumption.
    MR  =  DeforModelRed2DAxisymm

    fens = FENodeSet(X);
    fes = FESetQ4([1 2 4 3; 3 4 6 5; 5 6 8 7; 7 8 10 9; 9 10 12 11; 11 12 14 13; 13 14 16 15]);
    for ref = 1:nref
        fens,fes = Q4refine(fens,fes);
        list = selectnode(fens,distance = 1.0+0.1/2^nref, from = [0. 0.], inflate = tolerance);
        fens.xyz[list,:] =  FinEtools.MeshUtilModule.ontosphere(fens.xyz[list,:],1.0);
    end
    fens,fes = Q4toQ8(fens,fes)
    list = selectnode(fens,distance = 1.0+0.1/2^nref, from = [0. 0.], inflate = tolerance);
    fens.xyz[list,:] =  FinEtools.MeshUtilModule.ontosphere(fens.xyz[list,:],1.0);

    #     File  =   "mesh.vtk"
    # vtkexportmesh(File, fens, fes)

    # now we create the geometry and displacement fields
    geom  =  NodalField(fens.xyz)
    u  =  NodalField(zeros(size(fens.xyz,1),2)) # displacement field

    # Apply EBC's
    l1 = selectnode(fens,box = [-Inf Inf 0 0],inflate = tolerance)
    setebc!(u, l1, true, 2, 00.0)
    l1 = selectnode(fens,box = [-Inf Inf 1.79  1.79],inflate = tolerance)
    setebc!(u, l1, true, 2, 00.0)
    applyebc!(u)
    numberdofs!(u)

    # Temperature field
    dT  = NodalField(reshape(fens.xyz[:,1]+fens.xyz[:,2],size(fens.xyz,1),1));

    # Property and material
    material = MatDeforElastIso(MR, 0.0, Ea, nua, alphaa)

    femm  =  FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(2, 3), true), material)

    K  = stiffness(femm, geom, u)
    F  =  thermalstrainloads(femm, geom, u, dT)
    #K = cholesky(K)
    U =   K\F
    scattersysvec!(u, U[:])

    nA  = selectnode(fens,box = FFlt[1.0  1.0 0.0 0.0], inflate = tolerance);

    fld =  fieldfromintegpoints(femm, geom, u, dT, :Cauchy, 2)

    File  =   "LE11NAFEMS_Q8_sigmay.vtk"
    vtkexportmesh(File, fens, fes; scalars = [("sigmay", fld.values)],
    vectors = [("u", u.values)])
    println("range of  sigmay = $((minimum(fld.values), maximum(fld.values)))")
    @async run(`"paraview.exe" $File`)

    sA  =  fld.values[nA]/phun("MEGA*Pa")
    sAn  =  fld.values[nA]/sigmaA
    println("Stress at point A: $(sA) i. e.  $( sAn*100  )% of reference value")

    fen2fe  = FENodeToFEMap(fes.conn, nnodes(geom))
    function inspector(idat, elnum, conn, xe,  out,  xq)
        println("loc = $(  xq  ) : $(  transpose(out)/phun("MEGA*Pa")  )")
        return idat
    end

    inspectintegpoints(femm, geom, u, dT,  fen2fe.map[nA[1]],
    inspector, []; quantity = :Cauchy)

end # LE11NAFEMS_Q8_stress

function allrun()
    println("#####################################################")
    println("# LE11NAFEMS_Q8_algo ")
    LE11NAFEMS_Q8_algo()
    println("#####################################################")
    println("# LE11NAFEMS_Q8_algo2 ")
    LE11NAFEMS_Q8_algo2()
    println("#####################################################")
    println("# LE11NAFEMS_Q8_export_stress ")
    LE11NAFEMS_Q8_export_stress()
    println("#####################################################")
    println("# LE11NAFEMS_Q8_stress ")
    LE11NAFEMS_Q8_stress()
    return true
end # function allrun

end # module LE11NAFEMS_examples


using DiffEqFlux, OrdinaryDiffEq, Flux, Plots
using Zygote
using DifferentialEquations
using ForwardDiff
using LinearAlgebra
using BSON: @save
using BSON: @load
LinearAlgebra.BLAS.set_num_threads(16)

function rober(du,u,p,t)
    y₁,y₂,y₃ = u
    k₁,k₂,k₃ = p
    du[1] = -k₁*y₁+k₃*y₂*y₃
    du[2] =  k₁*y₁-k₂*y₂^2-k₃*y₂*y₃
    du[3] =  k₂*y₂^2
    nothing
end

function generate_data(u0, pz, tspan, solver)
    ode_prob = ODEProblem(rober,u0,tspan,pz)
    ode_sol = solve(ode_prob, solver)
    data = Array(ode_sol)
    return data, ode_sol.t
end

function predict_n_ode(k)
    Array(solve(prob, Rosenbrock23(), u0=u0,p=k, saveat=tsteps))
end

function loss_n_ode(k)
    #loss for stiff-net
      pred = predict_n_ode(k)
      loss = sum(abs2,ode_data .- pred)
      return loss
end

function grab_∇(par)
    #Grabs gradients of NN loss w.r.t params
    ForwardDiff.gradient(loss_n_ode, par)
end

function dθ_dt!(du,u,k,t)
    #ODE used to evolve parmaeters based on the negative gradient definition (evolving in the direction of steepest descent)
    du .= -1 .* grab_∇(u)
end

#################################################################################################################################################################################
#Generate the data
pz = [0.04,3e7,1e4]
u0 = [1.0,0.0,0.0]
tspan = (1.0e-5,1.0e5)
solver = Rosenbrock23(autodiff=false,diff_type=Val{:forward})
ode_data, tsteps = generate_data(u0, pz, tspan, solver)
plot(tsteps, ode_data')

#Create the NN
dudt2 = Chain(Dense(3,3, tanh),
             Dense(3, 3, tanh),
             Dense(3, 3, tanh),
             Dense(3,3))

@load "stiff-net.bson" dudt2

p,re = Flux.destructure(dudt2) # use this p as the initial condition!
dudt(u,p,t) = re(p)(u) # need to restrcture for backprop!
prob = ODEProblem(dudt,u0,tspan)
#Test predicion functions
predict_n_ode(p)
#Test loss functions, check for good initialisation (loss must be <100)
loss_n_ode(p)
#Save well behaved initialisation
#@save "stiff-net.bson" dudt2

#Train with solver
tspan_θ = (0.0,0.125)
prob_θ = ODEProblem(dθ_dt!,p,tspan_θ)
integrator = init(prob_θ,Rosenbrock23(autodiff=false))
loss_arr = []
params_arr = []
for (u,t) in tuples(integrator)
    loss = loss_n_ode(u)
    push!(loss_arr, loss)
    push!(params_arr, u)
    println(loss)
end

plot(loss_arr)

#Plotting first concentration
plotly()
plot(tsteps, predict_n_ode(param_his[:, end])[1, :])
scatter!(tsteps, ode_data[1, :], title = "First concentration", label = ["NN-pred" "Ground truth"])

#Plotting second concentration
plot(tsteps, predict_n_ode(param_his[:, end])[2, :])
scatter!(tsteps, ode_data[2, :], title = "Second concentration", label = ["NN-pred" "Ground truth"])

#Plotting third concentration
plot(tsteps, predict_n_ode(param_his[:, end])[3, :])
scatter!(tsteps, ode_data[3, :], title = "Third concentration", label = ["NN-pred" "Ground truth"])

#Plotting losss
loss = [loss_n_ode(param_his[:,i]) for i in 1:size(param_his)[2]]
plot(loss, title = "Loss")