using Flux
println("NNPDE_tests")
using DiffEqFlux
println("Starting Soon!")
using ModelingToolkit
using DiffEqBase
using Test, NeuralPDE
println("Starting Soon!")
using GalacticOptim
using Optim
using Quadrature,Cubature, Cuba
using QuasiMonteCarlo
using SciMLBase
using OrdinaryDiffEq
import ModelingToolkit: Interval, infimum, supremum
using DomainSets

using Random
Random.seed!(100)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

## Example 1, 1D ode
function test_ode(strategy_)
    println("Example 1, 1D ode: strategy: $(nameof(typeof(strategy_)))")
    @parameters θ
    @variables u(..)
    Dθ = Differential(θ)

    # 1D ODE
    eq = Dθ(u(θ)) ~ θ^3 + 2*θ + (θ^2)*((1+3*(θ^2))/(1+θ+(θ^3))) - u(θ)*(θ + ((1+3*(θ^2))/(1+θ+θ^3)))

    # Initial and boundary conditions
    bcs = [u(0.) ~ 1.0]

    # Space and time domains
    domains = [θ ∈ Interval(0.0,1.0)]

    # Neural network
    chain = FastChain(FastDense(1,12,Flux.σ),FastDense(12,1))
    initθ = Float64.(DiffEqFlux.initial_params(chain))

    discretization = NeuralPDE.PhysicsInformedNN(chain,
                                                 strategy_;
                                                 init_params = initθ,
                                                 phi = nothing,
                                                 derivative = nothing,
                                                 )

    @named pde_system = PDESystem(eq,bcs,domains,[θ],[u])
    prob = NeuralPDE.discretize(pde_system,discretization)
    sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

    res = GalacticOptim.solve(prob, ADAM(0.1); maxiters=1000)
    prob = remake(prob,u0=res.minimizer)
    res = GalacticOptim.solve(prob, ADAM(0.01); maxiters=500)
    prob = remake(prob,u0=res.minimizer)
    res = GalacticOptim.solve(prob, ADAM(0.001); maxiters=500)
    phi = discretization.phi

    analytic_sol_func(t) = exp(-(t^2)/2)/(1+t+t^3) + t^2
    ts = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1]
    u_real  = [analytic_sol_func(t) for t in ts]
    u_predict  = [first(phi(t,res.minimizer)) for t in ts]

    @test u_predict ≈ u_real atol = 0.1
    # using Plots
    # t_plot = collect(ts)
    # plot(t_plot ,u_real)
    # plot!(t_plot ,u_predict)
end

#TODO There is little meaning in these tests without checking the correctness of the prediction.
#TODO I propose to simply remove them.
function test_heterogeneous_equation(strategy_)
	println("Simple Heterogeneous input PDE, strategy: $(nameof(typeof(strategy_)))")
	@parameters x y
	@variables p(..) q(..) r(..) s(..)
	Dx = Differential(x)
	Dy = Differential(y)

	# 2D PDE
	eq  = p(x) + q(y) + Dx(r(x, y)) + Dy(s(y, x)) ~ 0
	# eq  = Dx(p(x)) + Dy(q(y)) + Dx(r(x, y)) + Dy(s(y, x)) + p(x) + q(y) + r(x, y) + s(y, x) ~ 0

	# Initial and boundary conditions
	bcs = [p(1) ~ 0.f0, q(-1) ~ 0.0f0,
			r(x, -1) ~ 0.f0, r(1, y) ~ 0.0f0,
			s(y, 1) ~ 0.0f0, s(-1, x) ~ 0.0f0]
	# bcs = [s(y, 1) ~ 0.0f0]
	# Space and time domains
	domains = [x ∈ Interval(0.0, 1.0),
				y ∈ Interval(0.0, 1.0)]

	# chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
	numhid = 3
	fastchains = [[FastChain(FastDense(1, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) for i in 1:2];
				[FastChain(FastDense(2, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) for i in 1:2]]
	discretization = NeuralPDE.PhysicsInformedNN(fastchains,
													strategy_)

	@named pde_system = PDESystem(eq, bcs, domains, [x,y], [p(x), q(y), r(x, y), s(y, x)])
	prob = SciMLBase.discretize(pde_system, discretization)
	res = GalacticOptim.solve(prob, BFGS(); maxiters=100)
end

## Heterogeneous system
function test_heterogeneous_system(strategy_)
	println("Heterogeneous input PDE with derivatives, strategy: $(nameof(typeof(strategy_)))")
	@parameters x y
	@variables p(..) q(..)
	Dx = Differential(x)
	Dy = Differential(y)

	# 2D PDE
	#TODO Dx(q(y)) = 0
	#TODO so p(x) = 0, q = const is has only trivial solution
	eq = p(x) + Dx(q(y)) ~ 0

	# Initial and boundary conditions
	bcs = [p(1) ~ 0.f0, q(-1) ~ 0.0f0]

	# Space and time domains
	domains = [x ∈ Interval(0.0, 1.0),
				y ∈ Interval(-1.0, 0.0)]

	# chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
	numhid = 3
	fastchains = [[FastChain(FastDense(1, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) for i in 1:2];
				[FastChain(FastDense(2, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) for i in 1:2]]
	discretization = NeuralPDE.PhysicsInformedNN(fastchains,
													strategy_)

	@named pde_system = PDESystem(eq, bcs, domains, [x,y], [p(x), q(y)])
	prob = SciMLBase.discretize(pde_system, discretization)
	res = GalacticOptim.solve(prob, BFGS(); maxiters=100)
end

grid_strategy = NeuralPDE.GridTraining(0.1)
quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(),
                                                    reltol=1e-3,abstol=1e-3,
                                                    maxiters =50, batch=100)
stochastic_strategy = NeuralPDE.StochasticTraining(100; bcs_points= 50)
quasirandom_strategy = NeuralPDE.QuasiRandomTraining(100;
                                                     sampling_alg = LatinHypercubeSample(),
                                                     resampling =false,
                                                     minibatch = 100
                                                    )
quasirandom_strategy_resampling = NeuralPDE.QuasiRandomTraining(100;
                                                     bcs_points= 50,
                                                     sampling_alg = LatticeRuleSample(),
                                                     resampling = true,
                                                     minibatch = 0)

strategies = [grid_strategy,stochastic_strategy, quadrature_strategy,quasirandom_strategy,quasirandom_strategy_resampling]

map(strategies) do strategy_
    test_ode(strategy_)
end
# map(strategies) do strategy_
#     test_heterogeneous_system(strategy_)
# end
# map(strategies) do strategy_
#     test_heterogeneous_equation(strategy_)
# end

## Heterogeneous system
println("Heterogeneous system")

@parameters x,y,z
@variables u(..), v(..), h(..), p(..)
Dz = Differential(z)
eqs = [
    u(x,y,z) ~ x+y+z,
    v(y,x) ~ x^2 + y^2,
    h(z) ~ cos(z),
    p(x,z) ~ exp(x)*exp(z),
    u(x,y,z) + v(y,x)*Dz(h(z)) - p(x,z) ~ x+y+z - (x^2+y^2)*sin(z) - exp(x)*exp(z)
]

bcs = [u(0,0,0) ~ 0.0]

domains = [x ∈ Interval(0.0, 1.0),
           y ∈ Interval(0.0, 1.0),
           z ∈ Interval(0.0, 1.0)]

chain = [FastChain(FastDense(3,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1)),
         FastChain(FastDense(2,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1)),
         FastChain(FastDense(1,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1)),
         FastChain(FastDense(2,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1))]

initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))

grid_strategy = NeuralPDE.GridTraining(0.1)
quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(),
                                                    reltol=1e-3,abstol=1e-3,
                                                    maxiters =50, batch=100)

discretization = NeuralPDE.PhysicsInformedNN(chain,grid_strategy;init_params = initθ)

@named pde_system = PDESystem(eqs,bcs,domains,[x,y,z],[u(x,y,z),v(y,x),h(z),p(x,z)])

prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

res = GalacticOptim.solve(prob, BFGS();maxiters=1000)
phi = discretization.phi

analytic_sol_func_ =
[
(x,y,z) -> x+y+z ,
(x,y) -> x^2 + y^2,
(z) -> cos(z),
(x,z) -> exp(x)*exp(z)
]

xs,ys,zs = [infimum(d.domain):0.1:supremum(d.domain) for d in domains]

u_real = [analytic_sol_func_[1](x,y,z) for x in xs  for y in ys for z in zs]
v_real = [analytic_sol_func_[2](y,x) for y in ys for x in xs ]
h_real = [analytic_sol_func_[3](z) for z in zs]
p_real = [analytic_sol_func_[4](x,z) for x in xs for z in zs]

real_ = [u_real,v_real,h_real,p_real]

initθ = discretization.init_params
acum =  [0;accumulate(+, length.(initθ))]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
minimizers = [res.minimizer[s] for s in sep]

u_predict = [phi[1]([x,y,z],minimizers[1])[1] for x in xs  for y in ys for z in zs]
v_predict = [phi[2]([y,x],minimizers[2])[1] for y in ys for x in xs ]
h_predict = [phi[3]([z],minimizers[3])[1] for z in zs]
p_predict = [phi[4]([x,z],minimizers[4])[1] for x in xs for z in zs]
predict = [u_predict,v_predict,h_predict,p_predict]

for i in 1:4
    @test predict[i] ≈ real_[i] rtol = 10^-3
end

# x_plot = collect(xs)
# y_plot = collect(ys)
# i=1
# z=0
# u_real = collect(analytic_sol_func_[1](x,y,z) for y in ys, x in xs);
# u_predict = collect(phi[1]([x,y,z],minimizers[1])[1]  for y in ys, x in xs);
# plot(x_plot,y_plot,u_real)
# plot!(x_plot,y_plot,u_predict)

## Example 2, 2D Poisson equation
function test_2d_poisson_equation(chain_, strategy_)
    println("Example 2, 2D Poisson equation, chain: $(nameof(typeof(chain_))), strategy: $(nameof(typeof(strategy_)))")
    @parameters x y
    @variables u(..)
    Dxx = Differential(x)^2
    Dyy = Differential(y)^2

    # 2D PDE
    eq  = Dxx(u(x,y)) + Dyy(u(x,y)) ~ -sin(pi*x)*sin(pi*y)

    # Initial and boundary conditions
    bcs = [u(0,y) ~ 0.0, u(1,y) ~ -sin(pi*1)*sin(pi*y),
           u(x,0) ~ 0.0, u(x,1) ~ -sin(pi*x)*sin(pi*1)]
    # Space and time domains
    domains = [x ∈ Interval(0.0,1.0),
               y ∈ Interval(0.0,1.0)]

    initθ = Float64.(DiffEqFlux.initial_params(chain_))
    discretization = NeuralPDE.PhysicsInformedNN(chain_,
                                                 strategy_;
                                                 init_params = initθ)

    @named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x, y)])
    prob = NeuralPDE.discretize(pde_system,discretization)
    sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)
    res = GalacticOptim.solve(prob, ADAM(0.1); maxiters=500)
    phi = discretization.phi

    xs,ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
    analytic_sol_func(x,y) = (sin(pi*x)*sin(pi*y))/(2pi^2)

    u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
    u_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
    diff_u = abs.(u_predict .- u_real)

    @test u_predict ≈ u_real atol = 2.0

    # p1 = plot(xs, ys, u_real, linetype=:contourf,title = "analytic");
    # p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict");
    # p3 = plot(xs, ys, diff_u,linetype=:contourf,title = "error");
    # plot(p1,p2,p3)
end

fastchain = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
fluxchain = Chain(Dense(2,12,Flux.σ),Dense(12,12,Flux.σ),Dense(12,1))
chains = [fluxchain, fastchain]
for chain in chains
    test_2d_poisson_equation(chain, grid_strategy)
end

for strategy_ in strategies
    chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
    test_2d_poisson_equation(chain_, strategy_)
end

algs = [CubatureJLp()] #CubatureJLh(),
for alg in algs
    chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
    strategy_ =  NeuralPDE.QuadratureTraining(quadrature_alg = alg,reltol=1e-4,abstol=1e-3,maxiters=30, batch=10)
    test_2d_poisson_equation(chain_, strategy_)
end

## Example 3, 3rd-order
println("Example 3, 3rd-order ode")
@parameters x
@variables u(..) ,Dxu(..) ,Dxxu(..) ,O1(..), O2(..)
Dxxx = Differential(x)^3
Dx = Differential(x)

# ODE
eq = Dx(Dxxu(x)) ~ cos(pi*x)

# Initial and boundary conditions
bcs_ = [u(0.) ~ 0.0,
        u(1.) ~ cos(pi),
        Dxu(1.) ~ 1.0]
ep = (cbrt(eps(eltype(Float64))))^2/6

der = [Dxu(x) ~ Dx(u(x)) + ep*O1(x) ,
       Dxxu(x) ~  Dx(Dxu(x))+ ep*O2(x)]

bcs = [bcs_;der]
# Space and time domains
domains = [x ∈ Interval(0.0,1.0)]

# Neural network
chain = [[FastChain(FastDense(1,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1)) for _ in 1:3];
         [FastChain(FastDense(1,4,Flux.tanh),FastDense(4,1)) for _ in 1:2];]
quasirandom_strategy = NeuralPDE.QuasiRandomTraining(100; #points
                                                     sampling_alg = LatinHypercubeSample())
initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))

discretization = NeuralPDE.PhysicsInformedNN(chain,quasirandom_strategy;init_params = initθ)

@named pde_system = PDESystem(eq,bcs,domains,[x],[u(x),Dxu(x),Dxxu(x),O1(x),O2(x)])

prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

pde_inner_loss_functions = prob.f.f.loss_function.pde_loss_function.pde_loss_functions.contents
bcs_inner_loss_functions = prob.f.f.loss_function.bcs_loss_function.bc_loss_functions.contents

cb_ = function (p,l)
    println("loss: ", l )
    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
    return false
end

res = GalacticOptim.solve(prob, BFGS(); maxiters=1000)
phi = discretization.phi[1]

analytic_sol_func(x) = (π*x*(-x+(π^2)*(2*x-3)+1)-sin(π*x))/(π^3)

xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1]
u_real  = [analytic_sol_func(x) for x in xs]
u_predict  = [first(phi(x,res.minimizer)) for x in xs]

@test u_predict ≈ u_real atol = 10^-4

# x_plot = collect(xs)
# plot(x_plot ,u_real)
# plot!(x_plot ,u_predict)

## Example 4, system of pde
println("Example 4, system of pde")
@parameters x, y
@variables u1(..), u2(..)
Dx = Differential(x)
Dy = Differential(y)

# System of pde
eqs = [Dx(u1(x,y)) + 4*Dy(u2(x,y)) ~ 0,
      Dx(u2(x,y)) + 9*Dy(u1(x,y)) ~ 0]
      # 3*u1(x,0) ~ 2*u2(x,0)]

# Initial and boundary conditions
bcs = [u1(x,0) ~ 2*x, u2(x,0) ~ 3*x]

# Space and time domains
domains = [x ∈ Interval(0.0,1.0), y ∈ Interval(0.0,1.0)]

# Neural network
chain1 = FastChain(FastDense(2,15,Flux.tanh),FastDense(15,1))
chain2 = FastChain(FastDense(2,15,Flux.tanh),FastDense(15,1))

quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(),
                                                    reltol=1e-3,abstol=1e-3,
                                                    maxiters =50, batch=100)
chain = [chain1,chain2]
initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))

discretization = NeuralPDE.PhysicsInformedNN(chain,quadrature_strategy; init_params = initθ)

@named pde_system = PDESystem(eqs,bcs,domains,[x,y],[u1(x, y),u2(x, y)])

prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

res = GalacticOptim.solve(prob,BFGS(); maxiters=1000)
phi = discretization.phi

analytic_sol_func(x,y) =[1/3*(6x - y), 1/2*(6x - y)]
xs,ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
u_real  = [[analytic_sol_func(x,y)[i] for x in xs  for y in ys] for i in 1:2]

initθ = discretization.init_params
acum =  [0;accumulate(+, length.(initθ))]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
minimizers = [res.minimizer[s] for s in sep]
u_predict  = [[phi[i]([x,y],minimizers[i])[1] for x in xs  for y in ys] for i in 1:2]

@test u_predict[1] ≈ u_real[1] atol = 0.1
@test u_predict[2] ≈ u_real[2] atol = 0.1

# p1 =plot(xs, ys, u_predict, st=:surface);
# p2 = plot(xs, ys, u_real, st=:surface);
# plot(p1,p2)

## Example 5, 2d wave equation, neumann boundary condition
println("Example 5, 2d wave equation, neumann boundary condition")
#here we use low level api for build solution
@parameters x, t
@variables u(..)
Dxx = Differential(x)^2
Dtt = Differential(t)^2
Dt = Differential(t)

#2D PDE
C=1
eq  = Dtt(u(x,t)) ~ C^2*Dxx(u(x,t))

# Initial and boundary conditions
bcs = [u(0,t) ~ 0.,# for all t > 0
       u(1,t) ~ 0.,# for all t > 0
       u(x,0) ~ x*(1. - x), #for all 0 < x < 1
       Dt(u(x,0)) ~ 0. ] #for all  0 < x < 1]

# Space and time domains
domains = [x ∈ Interval(0.0,1.0),
           t ∈ Interval(0.0,1.0)]

# Neural network
chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))
eltypeθ = eltype(initθ)
parameterless_type_θ = DiffEqBase.parameterless_type(initθ)
phi = NeuralPDE.get_phi(chain,parameterless_type_θ)
derivative = NeuralPDE.get_numeric_derivative()

indvars = [x,t]
depvars = [u(x, t)]
dim = length(domains)

quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(),
                                                    reltol=1e-3,abstol=1e-3,
                                                    maxiters =50, batch=100)

integral = NeuralPDE.get_numeric_integral(quadrature_strategy, indvars, depvars, chain, derivative)
_pde_loss_function = NeuralPDE.build_loss_function(eq,indvars,depvars,phi,derivative,integral,
                                                   chain,initθ,quadrature_strategy)
_pde_loss_function(rand(2,10), initθ)

bc_indvars = NeuralPDE.get_argument(bcs,indvars,depvars)
_bc_loss_functions = [NeuralPDE.build_loss_function(bc,indvars,depvars, phi, derivative,integral,
                                                    chain,initθ,quadrature_strategy,
                                                    bc_indvars = bc_indvar) for (bc,bc_indvar) in zip(bcs,bc_indvars)]
map(loss_f -> loss_f(rand(1,10), initθ),_bc_loss_functions)

dx = 0.1
train_sets = NeuralPDE.generate_training_sets(domains,dx,[eq],bcs,eltypeθ,indvars,depvars)
pde_train_set,bcs_train_set = train_sets
pde_bounds, bcs_bounds = NeuralPDE.get_bounds(domains,[eq],bcs,eltypeθ,indvars,depvars,quadrature_strategy)

lbs,ubs = pde_bounds
pde_loss_functions = [NeuralPDE.get_loss_function(_pde_loss_function,
                                                lbs[1],ubs[1],
                                                eltypeθ, parameterless_type_θ,
                                                quadrature_strategy)]

pde_loss_functions[1](initθ)

lbs,ubs = bcs_bounds
bc_loss_functions = [NeuralPDE.get_loss_function(_loss,lb,ub,
                                                 eltypeθ, parameterless_type_θ,
                                                 quadrature_strategy)
                                                 for (_loss,lb,ub) in zip(_bc_loss_functions, lbs,ubs)]

map(l->l(initθ) ,bc_loss_functions)

loss_functions =  [pde_loss_functions;bc_loss_functions]

function loss_function(θ,p)
    sum(map(l->l(θ) ,loss_functions))
end

f_ = OptimizationFunction(loss_function, GalacticOptim.AutoZygote())
prob = GalacticOptim.OptimizationProblem(f_, initθ)

cb_ = function (p,l)
    println("loss: ", l )
    println("losses: ", map(l -> l(p), loss_functions))
    return false
end

res = GalacticOptim.solve(prob,Optim.BFGS(); maxiters=500,f_abstol=10^-6)

xs,ts = [infimum(d.domain):dx:supremum(d.domain) for d in domains]
analytic_sol_func(x,t) =  sum([(8/(k^3*pi^3)) * sin(k*pi*x)*cos(C*k*pi*t) for k in 1:2:50000])

u_predict = reshape([first(phi([x,t],res.minimizer)) for x in xs for t in ts],(length(xs),length(ts)))
u_real = reshape([analytic_sol_func(x,t) for x in xs for t in ts], (length(xs),length(ts)))

@test u_predict ≈ u_real atol = 0.1

# diff_u = abs.(u_predict .- u_real)
# p1 = plot(xs, ts, u_real, linetype=:contourf,title = "analytic");
# p2 =plot(xs, ts, u_predict, linetype=:contourf,title = "predict");
# p3 = plot(xs, ts, diff_u,linetype=:contourf,title = "error");
# plot(p1,p2,p3)

## Example 6, pde with mixed derivative
println("Example 6, pde with mixed derivative")
@parameters x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
Dx = Differential(x)
Dy = Differential(y)

eq = Dxx(u(x,y)) + Dx(Dy(u(x,y))) - 2*Dyy(u(x,y)) ~  -1.

# Initial and boundary conditions
bcs = [u(x,0) ~ x,
       Dy(u(x,0)) ~ x,
       u(x,0) ~ Dy(u(x,0))]

# Space and time domains
domains = [x ∈ Interval(0.0,1.0), y ∈ Interval(0.0,1.0)]

quadrature_strategy = NeuralPDE.QuadratureTraining()
# Neural network
inner =20
chain = FastChain(FastDense(2,inner,Flux.tanh),FastDense(inner,inner,Flux.tanh),FastDense(inner,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))

discretization = NeuralPDE.PhysicsInformedNN(chain, quadrature_strategy; init_params = initθ)
@named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x, y)])

prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

res = GalacticOptim.solve(prob,BFGS(); maxiters=1500)
@show res.original

phi = discretization.phi

analytic_sol_func(x,y) = x + x*y +y^2/2
xs,ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]

u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
u_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
diff_u = abs.(u_predict .- u_real)

@test u_predict ≈ u_real rtol = 0.1

# p1 = plot(xs, ys, u_real, linetype=:contourf,title = "analytic");
# p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict");
# p3 = plot(xs, ys, diff_u,linetype=:contourf,title = "error");
# plot(p1,p2,p3)

## Example 7, Fokker-Planck equation
println("Example 7, Fokker-Planck equation")
# the example took from this article https://arxiv.org/abs/1910.10503
@parameters x
@variables p(..)
Dx = Differential(x)
Dxx = Differential(x)^2

#2D PDE
α = 0.3
β = 0.5
_σ = 0.5
# Discretization
dx = 0.01
# here we use normalization condition: dx*p(x) ~ 1, in order to get non-zero solution.
#(α - 3*β*x^2)*p(x) + (α*x - β*x^3)*Dx(p(x)) ~ (_σ^2/2)*Dxx(p(x))
eq  = [Dx((α*x - β*x^3)*p(x)) ~ (_σ^2/2)*Dxx(p(x))]
x_0 = -2.2
x_end = 2.2
# Initial and boundary conditions
bcs = [p(x_0) ~ 0. ,p(x_end) ~ 0.]

# Space and time domains
domains = [x ∈ Interval(-2.2,2.2)]

# Neural network
inn = 18
chain = FastChain(FastDense(1,inn,Flux.σ),
                  FastDense(inn,inn,Flux.σ),
                  FastDense(inn,inn,Flux.σ),
                  FastDense(inn,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))

lb = [x_0]
ub = [x_end]
function norm_loss_function(phi,θ,p)
    function inner_f(x,θ)
        dx*phi(x, θ) .- 1
    end
    prob = QuadratureProblem(inner_f, lb, ub, θ)
    norm2 = solve(prob, HCubatureJL(), reltol = 1e-8, abstol = 1e-8, maxiters =10);
    abs(norm2[1])
end
# norm_loss_function(phi,initθ,nothing)

discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             NeuralPDE.GridTraining(dx);
                                             init_params = initθ,
                                             additional_loss=norm_loss_function)

@named pde_system = PDESystem(eq,bcs,domains,[x],[p(x)])
prob = NeuralPDE.discretize(pde_system,discretization)

pde_inner_loss_functions = prob.f.f.loss_function.pde_loss_function.pde_loss_functions.contents
bcs_inner_loss_functions = prob.f.f.loss_function.bcs_loss_function.bc_loss_functions.contents

phi = discretization.phi

cb_ = function (p,l)
    println("loss: ", l )
    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
    println("additional_loss: ", norm_loss_function(phi,p,nothing))
    return false
end

res = GalacticOptim.solve(prob,LBFGS(),maxiters=400)
prob = remake(prob,u0=res.minimizer)
res = GalacticOptim.solve(prob,BFGS(),maxiters=2000)

C = 142.88418699042
analytic_sol_func(x) = C*exp((1/(2*_σ^2))*(2*α*x^2 - β*x^4))
xs = [infimum(d.domain):dx:supremum(d.domain) for d in domains][1]
u_real  = [analytic_sol_func(x) for x in xs]
u_predict  = [first(phi(x,res.u)) for x in xs]

@test u_predict ≈ u_real rtol = 1e-3

# plot(xs ,u_real, label = "analytic")
# plot!(xs ,u_predict, label = "predict")

## Example 8, Lorenz System (Parameter Estimation)
println("Example 8, Lorenz System")

Random.seed!(1234)
@parameters t ,σ_ ,β, ρ
@variables x(..), y(..), z(..)
Dt = Differential(t)
eqs = [Dt(x(t)) ~ σ_*(y(t) - x(t)),
       Dt(y(t)) ~ x(t)*(ρ - z(t)) - y(t),
       Dt(z(t)) ~ x(t)*y(t) - β*z(t)]

bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0]
domains = [t ∈ Interval(0.0,1.0)]
dt = 0.05

input_ = length(domains)
n = 12
chain = [FastChain(FastDense(input_,n,Flux.tanh),FastDense(n,n,Flux.σ),FastDense(n,1)) for _ in 1:3]
#Generate Data
function lorenz!(du,u,p,t)
    du[1] = 10.0*(u[2]-u[1])
    du[2] = u[1]*(28.0-u[3]) - u[2]
    du[3] = u[1]*u[2] - (8/3)*u[3]
end

u0 = [1.0;0.0;0.0]
tspan = (0.0,1.0)
prob = ODEProblem(lorenz!,u0,tspan)
sol = solve(prob, Tsit5(), dt=0.1)
ts = [infimum(d.domain):dt:supremum(d.domain) for d in domains][1]

function getData(sol)
    data = []
    us = hcat(sol(ts).u...)
    ts_ = hcat(sol(ts).t...)
    return [us,ts_]
end

data = getData(sol)

#Additional Loss Function
initθs = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))
acum =  [0;accumulate(+, length.(initθs))]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
(u_ , t_) = data
len = length(data[2])

function additional_loss(phi, θ , p)
    return sum(sum(abs2, phi[i](t_ , θ[sep[i]]) .- u_[[i], :])/len for i in 1:1:3)
end

discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             NeuralPDE.GridTraining(dt);
                                             init_params =initθs,
                                             param_estim=true,
                                             additional_loss=additional_loss)
testθ =reduce(vcat,initθs)
additional_loss(discretization.phi, testθ, nothing)

@named pde_system = PDESystem(eqs,bcs,domains,
                      [t],[x(t), y(t), z(t)],[σ_, ρ, β],
                      defaults=Dict([p => 1.0 for p in [σ_, ρ, β]]))
prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)
prob.f.f.loss_function([testθ;ones(3)])

res = GalacticOptim.solve(prob, Optim.BFGS(); maxiters=6000)
p_ = res.minimizer[end-2:end]
@test sum(abs2, p_[1] - 10.00) < 0.1
@test sum(abs2, p_[2] - 28.00) < 0.1
@test sum(abs2, p_[3] - (8/3)) < 0.1
#Plotting the system
# initθ = discretization.init_params
# acum =  [0;accumulate(+, length.(initθ))]
# sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
# minimizers = [res.minimizer[s] for s in sep]
# ts = [infimum(d.domain):dt/10:supremum(d.domain) for d in domains][1]
# u_predict  = [[discretization.phi[i]([t],minimizers[i])[1] for t in ts] for i in 1:3]
# plot(sol)
# plot!(ts, u_predict, label = ["x(t)" "y(t)" "z(t)"])

## Approximation of function 1D
println("Approximation of function 1D")

@parameters x
@variables u(..)

func(x) = @. 2 + abs(x - 0.5)

eq = [u(x) ~ func(x)]
bc = [u(0)~u(0)]

x0 = 0
x_end = 2
dx= 0.001
domain = [x ∈ Interval(x0,x_end)]

xs = collect(x0:dx:x_end)
func_s = func(xs)

hidden =10
chain = FastChain(FastDense(1,hidden, Flux.tanh),
                    FastDense(hidden, hidden, Flux.tanh),
                    FastDense(hidden, 1))
initθ = Float64.(DiffEqFlux.initial_params(chain))

strategy = NeuralPDE.GridTraining(0.01)

discretization = NeuralPDE.PhysicsInformedNN(chain,strategy; initial_params=initθ)
@named pdesys = PDESystem(eq,bc,domain,[x],[u(x)])
prob = NeuralPDE.discretize(pdesys,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pdesys,discretization)

res  = GalacticOptim.solve(prob,ADAM(0.05),maxiters=1000)
prob = remake(prob,u0=res.minimizer)
res  = GalacticOptim.solve(prob,BFGS(initial_stepnorm = 0.01),maxiters=500)

@test discretization.phi(xs',res.u) ≈ func(xs') rtol = 0.01

# plot(xs,func(xs))
# plot!(xs, discretization.phi(xs',res.u)')

## Approximation of function 1D 2
println("Approximation of function 1D 2")

@parameters x
@variables u(..)
func(x) =  @. cos(5pi*x)*x
eq = [u(x) ~ func(x)]
bc = [u(0)~u(0)]

x0 = 0
x_end = 4
domain = [x ∈ Interval(x0,x_end)]

hidden =20
chain = FastChain(FastDense(1,hidden, Flux.sin),
                  FastDense(hidden, hidden, Flux.sin),
                  FastDense(hidden, hidden, Flux.sin),
                  FastDense(hidden, 1))
initθ = DiffEqFlux.initial_params(chain)

strategy = NeuralPDE.GridTraining(0.01)

discretization = NeuralPDE.PhysicsInformedNN(chain,strategy; initial_params=initθ)
@named pdesys = PDESystem(eq,bc,domain,[x],[u(x)])
sym_prob = NeuralPDE.symbolic_discretize(pdesys,discretization)
prob = NeuralPDE.discretize(pdesys,discretization)

res  = GalacticOptim.solve(prob,ADAM(0.01),maxiters=500)
prob = remake(prob,u0=res.minimizer)
res  = GalacticOptim.solve(prob,BFGS(),maxiters=1000)

dx= 0.01
xs = collect(x0:dx:x_end)
func_s = func(xs)

@test discretization.phi(xs',res.u) ≈ func(xs') rtol = 0.01

# plot(xs,func(xs))
# plot!(xs, discretization.phi(xs',res.u)')

## Approximation of function 2D
println("Approximation of function 2D")

@parameters x,y
@variables u(..)
func(x,y) =  -cos(x) * cos(y) * exp(-((x - pi)^2 + (y - pi)^2))
eq = [u(x,y) ~ func(x,y)]
bc = [u(0,0) ~ u(0,0)]

x0 = -10
x_end = 10
y0 = -10
y_end = 10
d = 0.4

domain = [x ∈ Interval(x0, x_end), y ∈ Interval(y0, y_end)]

hidden =25
chain = FastChain(FastDense(2,hidden, Flux.tanh),
                  FastDense(hidden, hidden, Flux.tanh),
                  FastDense(hidden, hidden, Flux.tanh),
                  FastDense(hidden, 1))
initθ = Float64.(DiffEqFlux.initial_params(chain))

strategy = NeuralPDE.GridTraining(d)
discretization = NeuralPDE.PhysicsInformedNN(chain,strategy; initial_params=initθ)
@named pdesys = PDESystem(eq,bc,domain,[x,y],[u(x, y)])
prob = NeuralPDE.discretize(pdesys,discretization)
symprob = NeuralPDE.symbolic_discretize(pdesys,discretization)
prob.f.f.loss_function(initθ)

res  = GalacticOptim.solve(prob,ADAM(0.01),maxiters=500)
prob = remake(prob,u0=res.minimizer)
res  = GalacticOptim.solve(prob,BFGS(),maxiters=1000)
prob = remake(prob,u0=res.minimizer)
res  = GalacticOptim.solve(prob,BFGS(),maxiters=500)
phi = discretization.phi

xs = collect(x0:0.1:x_end)
ys = collect(y0:0.1:y_end)
u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
u_real = reshape([func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
diff_u = abs.(u_predict .- u_real)

@test u_predict ≈ u_real rtol = 0.05

# p1 = plot(xs, ys, u_real, st=:surface,title = "analytic");
# p2 = plot(xs, ys, u_predict, st=:surface,title = "predict");
# p3 = plot(xs, ys, diff_u,st=:surface,title = "error");
# plot(p1,p2,p3)

## approximation from data
println("Approximation of function from data and additional_loss")

@parameters x
@variables u(..)
eq = [u(0) ~ u(0)]
bc = [u(0) ~ u(0)]
x0 = 0
x_end = pi
dx =pi/10
domain = [x ∈ Interval(x0,x_end)]

hidden =10
chain = FastChain(FastDense(1,hidden, Flux.tanh),
                  FastDense(hidden, hidden, Flux.sin),
                  FastDense(hidden, hidden, Flux.tanh),
                  FastDense(hidden, 1))

initθ = Float64.(DiffEqFlux.initial_params(chain))

strategy = NeuralPDE.GridTraining(dx)
xs = collect(x0:dx:x_end)'
aproxf_(x) = @. cos(pi*x)
data =aproxf_(xs)

function additional_loss_(phi, θ , p)
    sum(abs2,phi(xs,θ) .- data)
end

discretization = NeuralPDE.PhysicsInformedNN(chain,strategy;
                                             initial_params=initθ,
                                             additional_loss=additional_loss_)

phi = discretization.phi
phi(xs, initθ)
additional_loss_(phi, initθ , nothing)

@named pdesys = PDESystem(eq,bc,domain,[x],[u(x)])
prob = NeuralPDE.discretize(pdesys,discretization)

res  = GalacticOptim.solve(prob,ADAM(0.01),maxiters=500)
prob = remake(prob,u0=res.minimizer)
res  = GalacticOptim.solve(prob,BFGS(),maxiters=500)

@test phi(xs,res.u) ≈ aproxf_(xs) rtol = 0.01

# xs_ = xs'
# plot(xs_,data')
# plot!(xs_, phi(xs,res.u)')

# func(x,y) = -20.0 * exp(-0.2 * sqrt(0.5 * (x^2 + y^2))) - exp(0.5 * (cos(2 * pi * x) + cos(2 * pi * y))) + e + 20
# func(x,y) = -abs(sin(x) * cos(y) * exp(abs(1 - (sqrt(x^2 + y^2)/pi))))


@testset "maximumlikelihood_gp.jl" begin
    n = 100

    x₁ = Variable("x₁", randn(n) / 3)
    x₂ = Variable("x₂", randn(n) / 10)

    ϕ = -.05 .+ x₁.value
    ξ = x₂.value

    σ = exp.(ϕ)
    θ = [-0.05; 1.0; 0.0; 1.0]

    pd = GeneralizedPareto.(σ, ξ)
    y = rand.(pd)

    @testset "gpfit(y; logscalecov, shapecov)" begin
        # model building with non-stationary logscale and shape
        fm = Extremes.gpfit(y,
            logscalecov = [x₁],
            shapecov = [x₂])

        # data is y
        @test fm.model.data.value ≈ y

        # logscale is x₁
        @test length(fm.model.logscale.covariate) == 1
        @test fm.model.logscale.covariate[1].value ≈ x₁.value

        # shape is x₂
        @test length(fm.model.shape.covariate) == 1
        @test fm.model.shape.covariate[1].value ≈ x₂.value

    end

    @testset "gpfit(y, initialvalues; logscalecov, shapecov)" begin
        # model building with non-stationary logscale and shape
        @test_logs Extremes.gpfit(y, zeros(Float64, 4),
            logscalecov = [x₁],
            shapecov = [x₂])

    end

    @testset "gpfit(df, datacol; logscalecovid, shapecovid)" begin
        # model building with non-stationary location, logscale and shape
        df = DataFrame(y = y, x1 = x₁.value, x2 = x₂.value)

        fm = Extremes.gpfit(df, :y, logscalecovid = [:x1], shapecovid = [:x2])

        # data is y
        @test fm.model.data.value ≈ y

        # logscale is x₁
        @test length(fm.model.logscale.covariate) == 1
        @test fm.model.logscale.covariate[1].value ≈ x₁.value

        # shape is x₂
        @test length(fm.model.shape.covariate) == 1
        @test fm.model.shape.covariate[1].value ≈ x₂.value

    end

    @testset "gpfit(df, datacol, initialvalues; logscalecovid, shapecovid)" begin
        # model building with non-stationary location, logscale and shape
        df = DataFrame(y = y, x1 = x₁.value, x2 = x₂.value)

        @test_logs Extremes.gpfit(df, :y, zeros(Float64, 4), logscalecovid = [:x1], shapecovid = [:x2])

    end

    @testset "gpfit(model, initialvalues)" begin
        # non-stationary location, logscale and shape
        model = Extremes.ThresholdExceedance(Variable("y", y),
            logscalecov = [x₁],
            shapecov = [x₂])

        fm = Extremes.gpfit(model, zeros(Float64, 4))

        # data is y
        @test fm.model.data.value ≈ y

        # logscale is x₁
        @test length(fm.model.logscale.covariate) == 1
        @test fm.model.logscale.covariate[1].value ≈ x₁.value

        # shape is x₂
        @test length(fm.model.shape.covariate) == 1
        @test fm.model.shape.covariate[1].value ≈ x₂.value

    end

end


"""
    integerfeasibility(A,d,nullVecs=false)

Given a linear system `Ax=d`, return an integer vector `x` which satisfies the
system.  This is the *integer programming feasibility problem*.

If `nullVecs==true`, then as well as returning a solution `x`, also return a
matrix `xNull` of vectors in the null space of `A` which could be added to the
`x` vector to find a solution which satisfies a constraint such as `0 .≤ x
.≤ u`; see the paper below.

This is not a robust tool, just a demo.

"Solving A System Of Diophantine Equations With Bounds On The Variables" by
Karen Aardal, Cor Hurkens, and Arjen Lenstra in Integer Programming and
Combinatorial Optimization, 6th International IPCO Conference, vol 1412, pp
229-242, 1998. See http://softlib.rice.edu/pub/CRPC-TRs/reports/CRPC-TR98782.pdf

# Examples
```jldoctest
julia> A=[10 1 -9; 1 8 8]; xtrue=[0; 2; 9]; d=A*xtrue;

julia> integerfeasibility(A,d)
3-element Array{Int64,1}:
 0
 2
 9

julia> n=20;m=30; A = rand(-10:10,n,m); xtrue = rand(0:10,m); d=A*xtrue;

julia> sum(abs.(xtrue - integerfeasibility(A,d) ))
0

```
"""
function integerfeasibility(A::AbstractArray{Td,2},
                            d::AbstractArray{Td,1}, flag=false) where {Td<:Number}
    m,n=size(A)
    N1=10^3
    N2=10^4
    B = [[I           zeros(Td,n,1)];
         zeros(Td,1,n) N1;
         N2*A          -N2*d]

    Bhat,_ = lll(B)
    Bprime = Bhat[1:n+1,1:n-m+1]
    x_d = Bprime[1:n,end]

    if flag
        xNull=Bprime[1:n,1:n-m]
        return x_d,xNull
    else
        return x_d
    end
end

"""
    x = subsetsum(a,s)

For a vector of integers `a`, and an integer `s`, try to find a binary
vector `x` such that `x'*a=s`. We use the LLL algorithm to find the
solution. This is not a robust tool, just a demo.

This follows the technique described by Lagarias and Odlyzko  in
"Solving Low-Density Subset Sum Problems"  in Journal of ACM, Jan 1985.
Code based on http://web.eecs.umich.edu/~cpeikert/lic15/lec05.pdf
We can likely get better results using techniques described and referenced in
https://www-almasty.lip6.fr/~joux/pages/papers/ToolBox.pdf

It's odd that permuting the `a` vector in the second example given below
causes the alg to often not find a solution. The example doesn't have the
required density, so maybe we should be asking why it finds an answer at
all.

# Examples
```jldoctest
julia> a=[1,2,4,9,20,38]; s=30; x=subsetsum(a,s); s-x'*a
0.0

julia> a=[32771,65543,131101,262187,524387,1048759, # from Bremner p 117
          2097523,4195057,8390143,16780259,33560539,
          67121039,134242091,268484171,536968403];

julia> s=891221976; x=subsetsum(a,s); s-x'*a
0.0

julia> N=40;a=rand(1:2^BigInt(256),N);xtrue=rand(Bool,N); s=a'*xtrue;

julia> setprecision(BigFloat,300); x=subsetsum(a,s); s-x'*a
0.0

```
"""
function subsetsum(a::AbstractArray{Ti,1},ss::Ti,returnBinary=false) where {Ti<:Integer}
    # page numbers below refer to lecture note above

    n = length(a)

    if ss<sum(a)/2
        flag=true
        s = sum(a)-ss
    else
        flag = false
        s = ss
    end
    # b is the "B" parameter defined at bottom of page 2
    nt = Ti(n)
    b = ceil(sqrt(nt*2^nt))
    # BB is the 'B' matrix defined also at bottom of page 2
    BB = [Matrix{Ti}(I,n,n+1); -b*a' b*s]
    B,T = lll(BB)

    # Pick off the interesting parts of the B matrix
    Bp = B[1:n,1:n+1]
    # Check if any of the columns satisfy the requirement
    sums = abs.(a'*Bp)
    ixMatch = findall(sums.==s)
    # d = diag(Bp'*Bp)  # solution generally matches a small element of d

    binarySolution = false
    global xb
    for ix=1:length(ixMatch)
        # Pick off a column of B
        x = Bp[:,ixMatch[ix][2]]
        if !(any(x.<0)|any(x.>1))
            binarySolution= true
            xb = flag ? mod.(x .+ 1,2) : x
        end
        x = -x
        if !(any(x.<0)|any(x.>1))
            binarySolution= true
            xb = flag ? mod.(x .+ 1,2) : x
        end
    end

    if binarySolution
        return xb
    else
        x = mdsubsetsum(a,ss,.5,1)
        if ~ismissing(x)
            print("A solution was found via mdsubsetsum\n")
            return x
        end

        if maximum(a)<2^(n^2/2)
            density = n/maximum(log2.(a))
            @printf("The density (%4.2f) of the 'a' vector is not as low as required\n",
                    density)
            @printf("(%4.2f) for Lagarias-Odlyzko to work. \n",2/n)
        end

        if returnBinary && length(ixMatch)>=1
            print("A non-binary solution was found; check that it's correct\n")
            return Bp[:,ixMatch[1][2]]
        else
            print("A solution was not found\n")
        end
    end
    return missing.*Bp[:,1]
end

"""
    x = mdsubsetsum(a,sM,ratio=.5,Kpm=3)

For a vector of integers `a`, and an integer `sM`, try to find a binary
vector `x` such that `x'*a=s` using the technique from "Multidimensional
subset sum problem" [1][2]. A major goal of the technique is to solve
problems in there are about 50% ones in `x`; other ratios of ones to zeros
can be specified in `ratio`.  The thesis also suggests searching `Kpm=3`
values around the nominal k. This technique is related to that in
[`subsetsum`](@ref) in that both use the LLL algorithm.  This is not a
robust tool, just a demo.

[1] https://scholarworks.rit.edu/theses/64/
[2] https://pdfs.semanticscholar.org/21a7/c2f9ff29507f1153aefcca04d1cd308e45c0.pdf

# Examples
```jldoctest
julia> a=[1,2,4,9,20,38]; s=30; x=mdsubsetsum(a,s); s-x'*a
0.0

julia> a=[32771,65543,131101,262187,524387,1048759, # from Bremner p 117
          2097523,4195057,8390143,16780259,33560539,
          67121039,134242091,268484171,536968403];

julia> sM=891221976; x=mdsubsetsum(a,sM); sM-x'*a

julia> N=40;a=rand(1:2^BigInt(256),N);xtrue=rand(Bool,N); s=a'*xtrue;

julia> setprecision(BigFloat,300); x=mdsubsetsum(a,s); s-x'*a
```
"""
function mdsubsetsum(a::AbstractArray{Ti,1},sM::Ti,ratio=.5,Kpm=3) where {Ti<:Integer}

    n = Ti(length(a))

    # r is heuristic value for r from end of Section 11 of thesis above
    r = Ti(2^round(.65*log2(maximum(a))))
    c = 2^10 # a heuristic; look right below eq (10)

    p = a .÷ r # See start of Section 10
    s = a-p.*r
    k0 = Ti(round(sum(s)*ratio/r)) # Equation (9) with ratio=.5 and j=1.
    p0 = sM ÷ r
    m = sM - p0*r

    for k = k0-Kpm:k0+Kpm
        # Equation (14)
        Bt = [Matrix{Ti}(I,n,n)*2 c*s       c*p      zeros(Ti,n);
              ones(Ti,1,n)        c*(k*r+m) c*(p0-k) 1]
        BB = Bt'  # Notation in paper is row-based; my brain is column-based
        B,T = lll(BB)

        for nx = 1:n
            if abs(B[n+3,nx])==1 && B[n+2,nx]==0 && B[n+1,nx]==0 &&
                      all((B[1:n,nx] .==1) .| (B[1:n,nx] .==-1))
                xhat = abs.(B[1:n,nx] .- B[n+3,nx])/2
                return Ti.(xhat)
            end
        end
    end
    @warn "Solution not found"
    return missing
end

"""
    rationalapprox(x::AbstractArray{<:Real,1},M,Ti=BigInt,verbose=false)

For a vector of Reals `x`, and an integer `M`, find an integer q such that
`maximum(abs.(x*q-round.(x*q)))` is small; the vector `x` is approximated by
`round.(x*q)//q`.  The integer `q` is less than or equal to `M` and the
approximation satisfies `max(abs.(x*q-round.(x*q)))≤sqrt(5)*2^(n/4 -
5)*M^(-1/n)`; this equation comes from the paper below.  The LLL algorithm
reduction is used to find the solution. The approximation vector is returned.
This is also known as "simultaneous diophantine approximation"; see for
example the title of the Hanrot paper below.

This is not a robust tool, just a demo.

"LLL: A Tool for Effective Diophantine Approximation" by Guillaume Hanrot in
the book "The LLL Algorithm: Survey and Applications" edited by Phong
Q. Nguyen and Brigitte Vallée, Springer, Heidelberg, 2010.

See also Chapter 9 of M. R. Bremner, "Lattice Basis Reduction: An
Introduction to the LLL Algorithm and Its Applications" CRC Press, 2012.

# Examples
```jldoctest
julia> x = [0.3912641745333527; 0.5455179974014548; 0.1908698210882469];
julia> rationalapprox(x,1e4,Int64)
3-element Array{Rational{Int64},1}:
 43//110
  6//11
 21//110
```
"""
function rationalapprox(x::AbstractArray{<:Real,1},M,Ti=BigInt,verbose=false)
    # Follows the multivariate technique in Theorem 8 of the
    # Nguyen and Vallee paper above

    n = length(x)
    Q = (M/2^(1/4))^(1/n)
    C = Ti(ceil(Q^(n+1)))
    Cx = Ti.(round.(C*x))

    X = [one(Ti) zeros(Ti,1,n);
         [Cx     C*I] ]
    B,_ = lll(X)
    q = abs(B[1,1])
    maxErr = maximum(abs.(x-round.(x*q)/q))
    bound = sqrt(5)*2^(n/4 - 1)/Q /q
    if maxErr>bound
        warning("The max error in the approximation $(maxErr) is unexpectedly "*
                "greater than the bound $(bound).")
    end
    p = Ti.(round.(x*q));
    if verbose
        println("q        = $(q)\n"*
                "maxErr/q = $(maxErr)\n"*
                "bound/q  = $(bound)")
        display([x round.(x*q)/q abs.(x-round.(x*q)/q)])
    end
    return p//q
end


 addprocs(3)
@everywhere using Base.Test, Compat, IntervalArithmetic, ValidatedTransfer

# module Gauss
#   using IntervalArithmetic, ValidatedTransfer, Compat
#   # println(names(IntervalArithmetic))
#   # println(names(ValidatedTransfer))
#
#   λ = @biginterval 2
#   λ̌ = sqrt(λ)
#   C1 = @biginterval log(2)
#   δ = @biginterval acosh(16/5)
#   P = @biginterval acosh(2log2(1+2^1.1)-1)
#   H = @biginterval 2
#
#   sp = ChebyshevZeroOneBV(BigInterval)
#   entrybounds = AnalyticEntryBounds(sp,δ,P,H)
#   lybounds = FirstOrderLYBounds(sp,λ,C1)
#
#
#   # const biglog2 = (@biginterval(log(2)))::BigInterval
#   Compat.@compat immutable g end; Compat.@compat (::g){T}(x::T) = (1/(exp2(x)-1))::T
#   Compat.@compat immutable dg end; Compat.@compat (::dg){T}(x::T) = begin e2x = exp2(x); e2x1 = e2x-1; (-log(T(2))*e2x/(e2x1*e2x1))::T end
#   Compat.@compat immutable γa end; Compat.@compat (::γa){T}(x::T) = (log2(1+1/x))::T
#   Compat.@compat immutable h end; Compat.@compat (::h){T}(x::T) = (exp2(x)-1)::T
#   Compat.@compat immutable dh end; Compat.@compat (::dh){T}(x::T) = (log(T(2))*exp2(x))::T
#   α = IntervalArithmetic.@biginterval 1;
#   α1 = IntervalArithmetic.@biginterval -1;
#   G0 = IntervalArithmetic.@biginterval 1/log(2);
#   G1 = IntervalArithmetic.@biginterval 0.1
#   saferad = IntervalArithmetic.@biginterval 0.4;

# LANFORD TEST

@everywhere @eval begin
  N = 100
  N_interp = nextpow2(N)

  prec = 8round(Int,(N+512)*0.2/8)
  @compat setprecision(BigFloat,prec)

  import Base.sinpi
  Base.sinpi(x::IntervalArithmetic.Interval) = sin(pi*x)

  ## MAP CONSTANTS
  @compat lan_v(x,b) = (5 - sqrt.(25-8(x+b)))/2
  @compat lan_dv(x,b) = 2./sqrt.(25-8(x+b))
  λ = @biginterval 3/2
  λ̌ = sqrt(λ)
  C1 = @biginterval 4/9
  δ = @biginterval acosh( 7/4)
  P = @biginterval acosh(4 - sqrt(6))
  H = @biginterval sqrt((7+sqrt(33))/2) #5/2sqrt(39-24sqrt(2)) + 3/2sqrt(40sqrt(2)-55)

  # Pointwise aliasing error
  t(j) = (1+(j!=0))
  aliaserror(j,k,N) = (H * t(j-1) * exp(P*(k-1)) * 2cosh(δ*(j-1)) /(exp(2δ*N)-1)).hi
  entrybound(j,k) = (H*t(j)*exp((k-1)*P-(j-1)*δ)).hi

  @compat @assert precision(BigFloat) == prec

  function acos_lan_v_norm{T}(A::Vector{T},b)
    B = similar(A)
    for i in eachindex(B)
      B[i] = acos(2lan_v((A[i]+1)/2,b)-1)
    end
    B
  end

  function lan_dv_norm{T}(A::Vector{T},b)
    B = similar(A)
    for i in eachindex(B)
      B[i] = lan_dv((A[i]+1)/2,b)
    end
    B
  end

  cheb_pts = ValidatedTransfer.chebyshevpoints(N_interp,BigInterval)
  acos_v0_cheb_pts = acos_lan_v_norm(cheb_pts,0); acos_v1_cheb_pts = acos_lan_v_norm(cheb_pts,1);
  dv0_cheb_pts = lan_dv_norm(cheb_pts,0); dv1_cheb_pts = lan_dv_norm(cheb_pts,1);
  helperdata = [acos_v0_cheb_pts,dv0_cheb_pts,acos_v1_cheb_pts,dv1_cheb_pts]

  println(size(acos_v0_cheb_pts))
  lanford_transferfn(i,k,helperdata) = cos((k-1)*helperdata[1][i])*helperdata[2][i] + cos((k-1)*helperdata[3][i])*helperdata[4][i]
end

@time M = maketransfer(N,lanford_transferfn,helperdata,aliaserror,entrybound,workers(),N_interp)
@time S = getsolution(M,workers())
@time savezip(tempname(),S)


using BandedMatrices, LinearAlgebra, FillArrays, Test
import BandedMatrices: _BandedMatrix, DefaultBandedMatrix

@testset "misc tests" begin
    @testset "Diagonal of banded" begin
        D = Diagonal([BandedMatrix(Eye(5),(2,3)), BandedMatrix(Eye(6),(1,1))])
        @test @inferred(D[1,1]) == I
        @test @inferred(D[1,2]) == zeros(5,6)
        @test bandwidths(D[1,2]) == (2,1)
    end

    @time @testset "BandedMatrix of BandedMatrix" begin
        A = BandedMatrix{DefaultBandedMatrix{Float64}}(undef, (1, 2), (0, 1))
        A[1,1] = BandedMatrix(Eye(1),(0,1))
        A[1,2] = BandedMatrix(Zeros(1,2),(0,1))
        A[1,2][1,1] = -1/3
        A[1,2][1,2] = 1/3
        B = BandedMatrix{DefaultBandedMatrix{Float64}}(undef, (2, 1), (1, 1))
        B[1,1] = 0.2BandedMatrix(Eye(1),(0,1))
        B[2,1] = BandedMatrix(Zeros(2,1), (1,0))
        B[2,1][1,1] = -2/30
        B[2,1][2,1] = 1/3

        @test (A*B)[1,1][1,1] ≈ 1/3

        A = BandedMatrix(Diagonal([randn(2,3), randn(3,4)]), (0,0))
        @test @inferred(A[1,2]) == zeros(2,4)

        A = BandedMatrix(Diagonal([brand(5,6,2,1), brand(7,8,1,1)]), (0,0))
        @test @inferred(A[1,2]) == zeros(5,8)
    end

    @time @testset "dense overrides" begin
        A = rand(10,11)
        @test bandwidths(A) == (9,10)
        A = rand(10)
        @test bandwidths(A) == (9,0)
        @test bandwidths(A') == (0,9)
    end

    @time @testset "trivial convert routines" begin
        A = brand(3,4,1,2)
        @test isa(BandedMatrix{Float64}(A), BandedMatrix{Float64})
        @test isa(AbstractMatrix{Float64}(A), BandedMatrix{Float64})
        @test isa(AbstractArray{Float64}(A), BandedMatrix{Float64})
        @test isa(BandedMatrix(A), BandedMatrix{Float64})
        @test isa(AbstractMatrix(A), BandedMatrix{Float64})
        @test isa(AbstractArray(A), BandedMatrix{Float64})
        @test isa(BandedMatrix{ComplexF16}(A), BandedMatrix{ComplexF16})
        @test isa(AbstractMatrix{ComplexF16}(A), BandedMatrix{ComplexF16})
        @test isa(AbstractArray{ComplexF16}(A), BandedMatrix{ComplexF16})
    end
    @time @testset "show" begin
      @test occursin("10×10 BandedMatrix{Float64,Array{Float64,2},"*string(Base.OneTo{Int})*"}",
         sprint() do io
            show(io, MIME"text/plain"(), brand(10, 10, 3, 3))
         end)
        needle = "1.0  0.0   ⋅ \n 0.0  1.0  0.0\n  ⋅   0.0  1.0"
        @test occursin(needle, sprint() do io
             show(io, MIME"text/plain"(), BandedMatrix(Eye(3),(1,1)))
          end)
    end
    @time @testset "Issue #27" begin
        A=brand(1,10,0,9)
        B=brand(10,10,255,255)
        @test Matrix(A*B)  ≈ Matrix(A)*Matrix(B)
    end

    @testset "defaultdot" begin
        A = randn(5)
        @test BandedMatrices.dot(A,A) ≡ LinearAlgebra.dot(A,A)
    end

    @testset "Offset axes" begin
        A = _BandedMatrix(Ones((Base.OneTo(3),Base.Slice(-5:5),)), Base.Slice(-4:2), 1,1)
        @test size(A) == (7,11)
        @test A[-4,-5] == 1
        @test A[-4,-5+2] == 1
        @test A[-4,-5+3] == 0
        @test_throws BoundsError A[-5,-5+3]
    end
end


using DoubleFloats, AccurateArithmetic
using Test

using Random
phi = Base.MathConstants.golden

include("basic.jl")
include("concrete_accuracy.jl")
include("function_accuracy.jl")
include("arithmetic.jl")


using FLoops, Distances

export threaded_pairwise, non_threaded_pairwise, threaded_inbounds_pairwise

function non_threaded_pairwise(d::Metric,x,y)
    C = zeros(length(x),length(y))
    for i in 1:size(C,1)
        for j in 1:size(C,2)
            C[i,j] = d(x[i],y[j])
        end
    end
    return C
end

function threaded_pairwise(d::Metric, x, y)
    C = zeros(length(x),length(y))
    @floop for (i,j) in Iterators.product(eachindex(x), eachindex(y))
            C[i,j] = d(x[i],y[j])
    end

    return C
end

function threaded_inbounds_pairwise(d::Metric, x, y)
    C = zeros(length(x),length(y))
    @inbounds begin
        @floop for (i,j) in Iterators.product(eachindex(x), eachindex(y))
            C[i,j] = d(x[i],y[j])
        end
    end
    return C
end


# This code presents an example to illustrate how NLboxsolve.jl can be used

using NLboxsolve

function test_nlboxsolve_outplace()

    function gershwin(x)

        f = similar(x)

        f[1] = 2*x[1] - x[5]
        f[2] = 2*x[2] - x[5]
        f[3] = 2*x[3] - x[5]
        f[4] = 2*x[4] - x[5] + x[6]
        f[5] = x[1] + x[2] + x[3] + x[4] - 1.0
        f[6] = x[6]*(x[4]-0.2)

        return f

    end

    # Solve Gershwin system

    x0 = [0.25,0.25,0.45,0.05,0.1,0.1]
    l = [-Inf,-Inf,-Inf,-Inf,0.0,0.0]
    u = [Inf,Inf,Inf,0.2,Inf,Inf]

    soln_newton     = constrained_newton(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm         = constrained_levenberg_marquardt(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_kyf     = constrained_levenberg_marquardt_kyf(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_ar      = constrained_levenberg_marquardt_ar(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg     = constrained_dogleg_solver(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg_bmp = constrained_dogleg_bmp_solver(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_nk         = constrained_newton_krylov(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_jfnk       = constrained_jacobian_free_newton_krylov(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)

    soln_newton_sparse     = constrained_newton_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_sparse         = constrained_levenberg_marquardt_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_kyf_sparse     = constrained_levenberg_marquardt_kyf_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_ar_sparse      = constrained_levenberg_marquardt_ar_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_sparse     = constrained_dogleg_solver_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_bmp_sparse = constrained_dogleg_solver_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_nk_sparse         = constrained_newton_krylov_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)

    soln1 = nlboxsolve(gershwin,x0,l,u,method = :dogleg,xtol=1e-15,ftol=1e-15)
    soln2 = nlboxsolve(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)

    test_one   = maximum(abs,soln_newton.zero       - soln2.zero) < 1e-14
    test_two   = maximum(abs,soln_lm.zero           - soln2.zero) < 1e-14
    test_three = maximum(abs,soln_lm_kyf.zero       - soln2.zero) < 1e-14
    test_four  = maximum(abs,soln_lm_ar.zero        - soln2.zero) < 1e-14
    test_five  = maximum(abs,soln_dogleg.zero       - soln2.zero) < 1e-14
    test_six   = maximum(abs,soln_dogleg_bmp.zero   - soln2.zero) < 1e-14
    test_seven = maximum(abs,soln_nk.zero           - soln2.zero) < 1e-14
    test_eight = maximum(abs,soln_jfnk.zero         - soln2.zero) < 1e-14

    test_nine     = maximum(abs,soln_newton_sparse.zero       - soln2.zero) < 1e-14
    test_ten      = maximum(abs,soln_lm_sparse.zero           - soln2.zero) < 1e-14
    test_eleven   = maximum(abs,soln_lm_kyf_sparse.zero       - soln2.zero) < 1e-14
    test_twelve   = maximum(abs,soln_lm_ar_sparse.zero        - soln2.zero) < 1e-14
    test_thirteen = maximum(abs,soln_dogleg_sparse.zero       - soln2.zero) < 1e-14
    test_fourteen = maximum(abs,soln_dogleg_bmp_sparse.zero   - soln2.zero) < 1e-14
    test_fifteen  = maximum(abs,soln_nk_sparse.zero           - soln2.zero) < 1e-14

    return test_one, test_two, test_three, test_four, test_five, test_six, test_seven, test_eight, test_nine, test_ten, test_eleven,
           test_twelve, test_thirteen, test_fourteen, test_fifteen

end

function test_nlboxsolve_inplace()

    function gershwin(f,x)

        f[1] = 2*x[1] - x[5]
        f[2] = 2*x[2] - x[5]
        f[3] = 2*x[3] - x[5]
        f[4] = 2*x[4] - x[5] + x[6]
        f[5] = x[1] + x[2] + x[3] + x[4] - 1.0
        f[6] = x[6]*(x[4]-0.2)

    end

    # Solve Gershwin system

    x0 = [0.25,0.25,0.45,0.05,0.1,0.1]
    l = [-Inf,-Inf,-Inf,-Inf,0.0,0.0]
    u = [Inf,Inf,Inf,0.2,Inf,Inf]

    soln_newton     = constrained_newton(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm         = constrained_levenberg_marquardt(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_kyf     = constrained_levenberg_marquardt_kyf(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_ar      = constrained_levenberg_marquardt_ar(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg     = constrained_dogleg_solver(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg_bmp = constrained_dogleg_bmp_solver(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_nk         = constrained_newton_krylov(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_jfnk       = constrained_jacobian_free_newton_krylov(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)

    soln_newton_sparse     = constrained_newton_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_sparse         = constrained_levenberg_marquardt_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_kyf_sparse     = constrained_levenberg_marquardt_kyf_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_ar_sparse      = constrained_levenberg_marquardt_ar_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_sparse     = constrained_dogleg_solver_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_bmp_sparse = constrained_dogleg_solver_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_nk_sparse         = constrained_newton_krylov_sparse(gershwin,x0,l,u,ftol=1e-15,xtol=1e-15)

    soln1 = nlboxsolve(gershwin,x0,l,u,method = :dogleg,xtol=1e-15,ftol=1e-15)
    soln2 = nlboxsolve(gershwin,x0,l,u,xtol=1e-15,ftol=1e-15)

    test_one   = maximum(abs,soln_newton.zero       - soln2.zero) < 1e-14
    test_two   = maximum(abs,soln_lm.zero           - soln2.zero) < 1e-14
    test_three = maximum(abs,soln_lm_kyf.zero       - soln2.zero) < 1e-14
    test_four  = maximum(abs,soln_lm_ar.zero        - soln2.zero) < 1e-14
    test_five  = maximum(abs,soln_dogleg.zero       - soln2.zero) < 1e-14
    test_six   = maximum(abs,soln_dogleg_bmp.zero   - soln2.zero) < 1e-14
    test_seven = maximum(abs,soln_nk.zero           - soln2.zero) < 1e-14
    test_eight = maximum(abs,soln_jfnk.zero         - soln2.zero) < 1e-14

    test_nine     = maximum(abs,soln_newton_sparse.zero       - soln2.zero) < 1e-14
    test_ten      = maximum(abs,soln_lm_sparse.zero           - soln2.zero) < 1e-14
    test_eleven   = maximum(abs,soln_lm_kyf_sparse.zero       - soln2.zero) < 1e-14
    test_twelve   = maximum(abs,soln_lm_ar_sparse.zero        - soln2.zero) < 1e-14
    test_thirteen = maximum(abs,soln_dogleg_sparse.zero       - soln2.zero) < 1e-14
    test_fourteen = maximum(abs,soln_dogleg_bmp_sparse.zero   - soln2.zero) < 1e-14
    test_fifteen  = maximum(abs,soln_nk_sparse.zero           - soln2.zero) < 1e-14

    return test_one, test_two, test_three, test_four, test_five, test_six, test_seven, test_eight, test_nine, test_ten, test_eleven,
           test_twelve, test_thirteen, test_fourteen, test_fifteen

end

test_nlboxsolve_outplace()
test_nlboxsolve_inplace()


using SparseArrays
using ArgCheck

"""
    Permute{A} <: Bijector{1}

A bijector implementation of a permutation. The permutation is performed
using a matrix of type `A`. There are a couple of different ways to construct `Permute`:

    Permute([0 1; 1 0])          # will map [1, 2] => [2, 1]
    Permute([2, 1])              # will map [1, 2] => [2, 1]
    Permute(2, 2 => 1, 1 => 2)   # will map [1, 2] => [2, 1]
    Permute(2, [1, 2] => [2, 1]) # will map [1, 2] => [2, 1]

If this is not clear, the examples might be of help.

# Examples
A simple example is permuting a vector of size 3.

```julia-repl
julia> b1 = Permute([
           0 1 0;
           1 0 0;
           0 0 1
       ])
Permute{Array{Int64,2}}([0 1 0; 1 0 0; 0 0 1])

julia> b2 = Permute([2, 1, 3])           # specify all elements at once
Permute{SparseArrays.SparseMatrixCSC{Float64,Int64}}(

  [2, 1]  =  1.0
  [1, 2]  =  1.0
  [3, 3]  =  1.0)

julia> b3 = Permute(3, 2 => 1, 1 => 2)    # element-wise
Permute{SparseArrays.SparseMatrixCSC{Float64,Int64}}(
  [2, 1]  =  1.0
  [1, 2]  =  1.0
  [3, 3]  =  1.0)

julia> b4 = Permute(3, [1, 2] => [2, 1])  # block-wise
Permute{SparseArrays.SparseMatrixCSC{Float64,Int64}}(
  [2, 1]  =  1.0
  [1, 2]  =  1.0
  [3, 3]  =  1.0)

julia> b1.A == b2.A == b3.A == b4.A
true

julia> b1([1., 2., 3.])
3-element Array{Float64,1}:
 2.0
 1.0
 3.0

julia> b2([1., 2., 3.])
3-element Array{Float64,1}:
 2.0
 1.0
 3.0

julia> b3([1., 2., 3.])
3-element Array{Float64,1}:
 2.0
 1.0
 3.0

julia> b4([1., 2., 3.])
3-element Array{Float64,1}:
 2.0
 1.0
 3.0

julia> inv(b1)
Permute{LinearAlgebra.Transpose{Int64,Array{Int64,2}}}([0 1 0; 1 0 0; 0 0 1])

julia> inv(b1)(b1([1., 2., 3.]))
3-element Array{Float64,1}:
 1.0
 2.0
 3.0
```
"""
struct Permute{A} <: Bijector{1}
    A::A
end

Base.:(==)(b1::Permute, b2::Permute) = b1.A == b2.A

function Permute(indices::AbstractVector{Int})
    # construct a sparse-matrix for use in the multiplication
    n = length(indices)
    A = spzeros(n, n)

    for (i, idx) in enumerate(indices)
        A[idx, i] = 1.0
    end

    return Permute(A)
end

function Permute(n::Int, indices::Pair{Int, Int}...)
    A = spdiagm(0 => ones(n))

    dests = Set{Int}()
    sources = Set{Int}()

    for (src, dst) in indices
        @argcheck dst ∉ dests
        @argcheck src ∉ sources

        push!(dests, dst)
        push!(sources, src)

        A[dst, src] = 1.0
        A[src, src] = 0.0  # <= remove `src => src`
    end

    @argcheck (sources ∩ dests) == (sources ∪ dests) "$sources ∩ $dests ≠ $sources ∪ $dests"

    dropzeros!(A)
    return Permute(A)
end

function Permute(n::Int, indices::Pair{Vector{Int}, Vector{Int}}...)
    A = spdiagm(0 => ones(n))

    dests = Set{Int}()
    sources = Set{Int}()

    for (srcs, dsts) in indices
        @argcheck length(srcs) == length(dsts) "$srcs => $dsts is not bijective"

        for (src, dst) in zip(srcs, dsts)
            @argcheck dst ∉ dests "$dst used more than once"
            @argcheck src ∉ sources "$src used more than once"

            push!(dests, dst)
            push!(sources, src)

            A[dst, src] = 1.0
            A[src, src] = 0.0  # <= remove `src => src`
        end
    end

    @argcheck (sources ∩ dests) == (sources ∪ dests) "$sources ∩ $dests ≠ $sources ∪ $dests"

    dropzeros!(A)
    return Permute(A)
end

@inline (b::Permute)(x::AbstractVecOrMat) = b.A * x
@inline inv(b::Permute) = Permute(transpose(b.A))

logabsdetjac(b::Permute, x::AbstractVector) = zero(eltype(x))
logabsdetjac(b::Permute, x::AbstractMatrix) = zero(eltype(x), size(x, 2))


# Find the largest possible axes-aligned rectangle of true values within an array.
function largest_rect(X::AbstractArray{Bool})
	m,n = size(X)
	H = count.(X[1,:])
	col,height,areamax = largest_rect(H)
	rowrange = 1:1
	colrange = col .+ (0:height-1)
	for i in 2:m
		# in each column, heights of consecutive ones from row 1 to i
		H = [ X[i,j]*(1+H[j]) for j in 1:n ]
		newcol,newheight,newarea = largest_rect(H)
		if newarea > areamax
			areamax = newarea
			rowrange = i-newheight+1:i
			colrange = newcol .+ (0:newarea÷newheight-1)
		end
	end
	return rowrange,colrange,areamax

end

# Find the largest rectangle underneath a histogram-count vector.
function largest_rect(x::AbstractVector{T} where T <: Number)
    maxarea = 0
    x = [-1;x;-1]
    stack = [1]  # index into x

	idx = 0
	height = 0
    for i in 1:length(x)
        while x[i] < x[stack[end]]
			h = x[pop!(stack)]
			area = h*(i-stack[end]-1)
			if area > maxarea
				idx,height = stack[end],h
				maxarea = area
			end
		end
		push!(stack,i)
	end
	return idx,height,maxarea
end

# Grow a rectangular region to include all the adjacent pixels with similar value.
function grow_rectangle(X::AbstractMatrix,irange,jrange)
	m,n = size(X)
	rect = Array(CartesianIndices((irange,jrange)))
	isempty(rect) && (return rect)

	# Go far outside to get an idea of the intensity of non-belonging pixels in the vicinity.
	h,w = length(irange),length(jrange)
	mid = round.(Int,(median(irange),median(jrange)))
	ii = clamp.(mid[1]-h:mid[1]+h,1,m)
	jj = clamp.(mid[2]-w:mid[2]+w,1,n)
	outerval = median( [ vec(X[ii[[1,end]],jj]); vec(X[ii,jj[[1,end]]]) ] )

	# Define region and average value in it.
	region = falses(size(X))
	region[irange,jrange] .= true
	checked = copy(region)

	# Include those more like the average of the interior than the sampled exterior.
	avg = mean(X[region])
	thresh = max(0.2,0.5*(avg-outerval))
	nbrs = CartesianIndex.([(-1,0),(1,0),(0,-1),(0,1)])

	for kk = 1:1000
		added = 0
		for nbr in nbrs, r in findall(region)
			newpt = r + nbr
			if checkbounds(Bool,X,newpt) && !checked[newpt]
				if abs(X[newpt] - avg) < thresh
					region[newpt] = true
					added += 1
				end
				checked[newpt] = true
			end
		end
		#@show added
		(added == 0) && break
	end
	return findall(region)
end

"""
	findpurkinje(img,thresh=0.5)
Return indices of pixels believed to be in the purkinje of the image `img`. Decrease `thresh` toward zero to include more pixels, or increase to one to include fewer. Returns a vector of `CartesianIndex`.
"""
function findpurkinje(img::AbstractMatrix{T} where T<:Colorant,thresh=0.5;options=get_defaults(),rectangle=false)
	X = options.purkinje_channel.(img)
	idx = findpurkinje(X,thresh;options,rectangle=rectangle)
	return idx
end

function findpurkinje(X::AbstractMatrix{T} where T<:Number,thresh=0.5;options=get_defaults(),rectangle=false)
	m,n = size(X)
	θ,area = 0.95,0
	iran = jran = cen = NaN
	B = X.>θ
	ctr = 0
	removed = CartesianIndex{2}[]
	while θ >= thresh
		iran,jran,area = largest_rect(B)
		@debug "iran = $iran, jran = $jran"
		# If area is too small, there aren't enough pixels at this brightness.
		if (ctr > 4) || (area ≤ m*n*options.purkinje_minarea)
			θ -= 0.05
			ctr = 0
			B = X.>θ
			#B[removed] .= false
			@debug "θ = $θ"
		else
			h,w = (iran[end]-iran[1]+1,jran[end]-jran[1]+1)
			# Must have appropriate aspect ratio. Bright lines at eyelids can give short, skinny rectangles.
			if h > 0.8w
				# Must have contrast with neighboring pixels.
				cen = round.(Int,(median(iran),median(jran)))
				# Expand to larger box
				rows = clamp.(cen[1]-h:cen[1]+h,1,m)
				cols = clamp.(cen[2]-w:cen[2]+w,1,n)
				inner = sum(X[iran,jran])
				outer = sum(X[rows,cols]) - inner
				inneravg = inner/area
				outeravg = outer/((2h+1)*(2w+1)-area)
				@debug "inner = $inneravg, outer = $outeravg"
				if inneravg - outeravg > 0.25
					# we've got it
					break
				end
			end
			# Remove this region from consideration.
			bad = CartesianIndices((iran,jran))
			#bad = grow_rectangle(B,iran,jran)
			B[bad] .= false
			append!(removed,bad)
			@debug length(removed)
			ctr += 1
		end
	end
	return rectangle ? CartesianIndices((iran,jran)) : grow_rectangle(X,iran,jran)
end

"""
	findpeaks(index,sz,dim)
Look for peaks in the distribution of a vector `index` of indexes. Tuple `sz` is the size of the image that the indexes are taken from, and `dim` is the dimension (1=rows,2=columns) to operate in. Returns a pair of real values at the peaks of the distributions in the two halves of the range of the index.
"""
function findpeaks(idx,sz,dim)
	# look for the peaks in a distribution of indices
	vals = getindex.(idx,dim)

	edges = LinRange(1,sz[dim]+1,101)
	hist = fit(Histogram,vals,edges)
	w = hist.weights

	## search left/right sides separately
	k = argmax(w[1:50])
	jleft = mean(edges[k:k+1])
	wleft = w[k]
	k = argmax(w[51:end])
	jright = mean(edges[50+k:50+k+1])
	wright = w[50+k]

	return (jleft,wleft),(jright,wright)
end


# # Clocked thermostat

# ## Model

module clocked_thermostat #jl

using ReachabilityAnalysis, ModelingToolkit

const A1 = 40.0
const A2 = 30.0
const B = 0.5
const t1 = 75
const t2 = 65

const var′ = @variables t t′

function thermostat_on′()
    invariant = HalfSpace(t <= t1, var′)
    @system(x' = [-B 0; 0 0]*x + [A1, 1], x ∈ invariant)
end

function thermostat_off′()
    invariant = HalfSpace(t >= t2, var′)
    @system(x' = [-B 0; 0 0]*x + [A2, 1], x ∈ invariant)
end

function thermostat_hybrid′()
    automaton = LightAutomaton(2)
    add_transition!(automaton, 1, 2, 1)
    add_transition!(automaton, 2, 1, 2)

    mode1 = thermostat_on′()
    mode2 = thermostat_off′()
    modes = [mode1, mode2]

    ## transition on -> off
    guard = HalfSpace(t >= t1, var′)
    trans1 = ConstrainedIdentityMap(1, guard)
    ## transition off -> on
    guard = HalfSpace(t <= t2, var′)
    trans2 = ConstrainedIdentityMap(1, guard)
    resetmaps = [trans1, trans2]

    return HybridSystem(automaton, modes, resetmaps, [AutonomousSwitching()])
end

function model(X0)
    H = thermostat_hybrid′()
return IVP(H, X0)
end

end  #jl

# ## References


"""
    randstate([T=ComplexF64,] basis)

Calculate a random normalized ket state.
"""
function randstate(::Type{T}, b::Basis) where T
    psi = Ket(b, rand(T, length(b)))
    normalize!(psi)
    psi
end
randstate(b) = randstate(ComplexF64, b)

"""
    randoperator([T=ComplexF64,] b1[, b2])

Calculate a random unnormalized dense operator.
"""
randoperator(::Type{T}, b1::Basis, b2::Basis) where T = DenseOperator(b1, b2, rand(T, length(b1), length(b2)))
randoperator(b1::Basis, b2::Basis) = randoperator(ComplexF64, b1, b2)
randoperator(::Type{T}, b::Basis) where T = randoperator(T, b, b)
randoperator(b) = randoperator(ComplexF64, b)

"""
    thermalstate(H,T)

Thermal state ``exp(-H/T)/Tr[exp(-H/T)]``.
"""
function thermalstate(H,T)
    return normalize(exp(-dense(H)/T))
end

"""
    coherentthermalstate([C=ComplexF64,] basis::FockBasis,H,T,alpha)

Coherent thermal state ``D(α)exp(-H/T)/Tr[exp(-H/T)]D^†(α)``.
"""
function coherentthermalstate(::Type{C},basis::B,H::AbstractOperator{B,B},T,alpha) where {C,B<:FockBasis}
    D = displace(C,basis,alpha)
    return D*thermalstate(H,T)*dagger(D)
end
coherentthermalstate(basis::B,H::AbstractOperator{B,B},T,alpha) where B<:FockBasis = coherentthermalstate(ComplexF64,basis,H,T,alpha)

"""
    phase_average(rho)

Returns the phase-average of ``ρ`` containing only the diagonal elements.
"""
function phase_average(rho)
    return Operator(basis(rho),diagm(0 => diag(rho.data)))
end

"""
    passive_state(rho,IncreasingEigenenergies=true)

Passive state ``π`` of ``ρ``. IncreasingEigenenergies=true means that higher indices correspond to higher energies.
"""
function passive_state(rho,IncreasingEigenenergies=true)
    return DenseOperator(basis(rho),diagm(0 => sort!(abs.(eigvals(rho.data)),rev=IncreasingEigenenergies)))
end


export LinearMassSpringChain
export LinearMassSpringDamperChain

mutable struct LinearMassSpringDamperChain <: ControlBenchmarkProblem
    n::Integer          # The number of masses in the chain
    k::Vector{Float64}  # The spring constants in the chain
    β::Vector{Float64}  # The damping constants in the chain
    m::Vector{Float64}  # The masses in the chain
    inputs::Vector{Int} # The masses which have an acceleration input

    function LinearMassSpringDamperChain( n::Integer, k, β, m; inputs = [1; n] )
        if( n < 1 )
            throw( DomainError( n, "Must specify more than one mass" ) )
        end

        if isa( k, Number )
            k = fill( k, n+1 )
        end

        if isa( β, Number )
            β = fill( β, n+1 )
        end

        if isa( m, Number )
            m = fill( m, n )
        end

        # Create vectors
        k = reshape( k, : )
        β = reshape( β, : )
        m = reshape( m, : )

        if size( k, 1 ) != n+1
            error( "k must have n+1 values" )
        end

        if size( β, 1 ) != n+1
            error( "β must have n+1 values" )
        end

        if size( m, 1 ) != n
            error( "m must have n values" )
        end

        # Expand ranges and ensure no inputs are duplicated
        inArray = collect( inputs )
        inArray = reshape( inArray, : )
        unique!( inArray )

        if size( inArray, 1 ) > n
            error( inArray, "More inputs requested than masses" )
        end

        if !all( inArray .<= n )
            throw( DomainError( inputs, "Inputs must be on masses" ) )
        end

        return new(n, k, β, m, inArray)
    end
end

LinearMassSpringChain( n::Integer, k, m; inputs = [1; n] ) = LinearMassSpringDamperChain( n, k, 0, m; inputs )

function controlbenchmark( chain::LinearMassSpringDamperChain )
    # Extract into temp variables for more easy reference
    n = chain.n
    k = chain.k
    β = chain.β
    m = chain.m
    inputs = chain.inputs

    # Create the diagonal elements for the matrix part for the springs
    springUpperDiag = k[2:end-1] ./ m[1:end-1]
    springMainDiag  = -( k[1:end-1] + k[2:end] ) ./ m[1:end]
    springLowerDiag = k[2:end-1] ./ m[2:end]

    # Create the state matrix
    S = spdiagm(  0 => springMainDiag,    # The main diagonal matrix elements
                 -1 => springLowerDiag,   # The elements below the diagonal
                  1 => springUpperDiag )  # The elements above the diagonal

    # Create the diagonal elements for the matrix part for the damper
    damperUpperDiag = β[2:end-1] ./ m[1:end-1]
    damperMainDiag  = -( β[1:end-1] + β[2:end] ) ./ m[1:end]
    damperLowerDiag = β[2:end-1] ./ m[2:end]

    Beta = spdiagm(  0 => damperMainDiag,    # The main diagonal matrix elements
                    -1 => damperLowerDiag,   # The elements below the diagonal
                     1 => damperUpperDiag )  # The elements above the diagonal

    Z   = spzeros( n, n )
    Eye = sparse( 1.0*I, n, n )

    A = [ Z Eye
          S Beta ]

    # Create the input matrix
    dim = size( inputs, 1 )
    ind = CartesianIndex.( inputs .+ n, 1:dim )
    B   = spzeros( 2*n, dim )
    B[ind] .= - 1 ./ m[inputs]

    C = sparse( 1.0*I, 2*n, 2*n )
    D = spzeros( 2*n, dim )

    return (; :sys => HeteroStateSpace( A, B, C, D, Continuous() ) )
end


using Oceananigans
using Printf
using Statistics
using GLMakie

using Oceananigans.ImmersedBoundaries: mask_immersed_field!
using Oceananigans.Utils: prettysummary

# Monin-Obukhov drag coefficient
z₀ = 1e-2 # Charnock roughness
κ = 0.4 # Von Karman constant
Cᴰ(Δz) = (κ / log(Δz / 2z₀))^2

@inline bottom_drag_u(x, y, t, u, w, Cᴰ) = - Cᴰ * u * sqrt(u^2 + w^2)
@inline bottom_drag_w(x, y, t, u, w, Cᴰ) = - Cᴰ * w * sqrt(u^2 + w^2)
@inline bottom_drag_u(x, y, z, t, u, w, Cᴰ) = - Cᴰ * u * sqrt(u^2 + w^2)
@inline bottom_drag_w(x, y, z, t, u, w, Cᴰ) = - Cᴰ * w * sqrt(u^2 + w^2)

function hilly_simulation(; Nx = 64,
                            Nz = Nx,
                            h = 0.1,
                            Re = 1e3,
                            N² = 1e-2,
                            boundary_condition = :no_slip,
                            stop_time = 1,
                            save_interval = 0.1,
                            architecture = CPU(),
                            filename = "flow_over_hills")

    underlying_grid = RectilinearGrid(architecture, size = (Nx, Nz), halo = (3, 3),
                                      x = (0, 2π), z = (0, 1),
                                      topology = (Periodic, Flat, Bounded))

    min_Δz = 1 / 32
    if h > 0
        hills(x, y) = h * (1 + sin(x)) / 2 + 2min_Δz # ensure that the bottom boundary is immersed
        grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(hills))
    else # no hills
        grid = underlying_grid
    end

    closure = isfinite(Re) ? ScalarDiffusivity(ν=1/Re, κ=1/Re) : nothing

    if boundary_condition == :no_slip
        no_slip = ValueBoundaryCondition(0)
        u_bcs = FieldBoundaryConditions(bottom=no_slip, immersed=no_slip)
        w_bcs = FieldBoundaryConditions(immersed=no_slip)
        boundary_conditions = (; u = u_bcs, w = w_bcs)
    elseif boundary_condition == :bottom_drag
        Δz = 1 / Nz
        Δx = 2π / Nz
        u_drag_bc = FluxBoundaryCondition(bottom_drag_u, field_dependencies=(:u, :w), parameters=Cᴰ(Δz))
        w_drag_bc = FluxBoundaryCondition(bottom_drag_w, field_dependencies=(:u, :w), parameters=Cᴰ(Δx))
        u_bcs = FieldBoundaryConditions(bottom=u_drag_bc, immersed=u_drag_bc)
        w_bcs = FieldBoundaryConditions(immersed=w_drag_bc)
        boundary_conditions = (; u = u_bcs, w = w_bcs)
        @info string("Using a bottom drag with coefficient ", Cᴰ(Δz))
    else
        boundary_conditions = NamedTuple()
    end

    model = NonhydrostaticModel(; grid, closure, boundary_conditions,
                                advection = WENO5(),
                                timestepper = :RungeKutta3,
                                tracers = :b,
                                buoyancy = BuoyancyTracer())

    # Steady flow + perturbations
    δh = 0.1
    ∂z_ψᵋ(x, z) = 4π * sin(4x) * cos(4π * z) * exp(-(z - h)^2 / 2δh^2)
    ∂x_ψᵋ(x, z) = 4  * cos(4x) * sin(4π * z) * exp(-(z - h)^2 / 2δh^2)
    bᵢ(x, y, z) = N² * z + 1e-9 * rand()
    uᵢ(x, y, z) = 1.0 + ∂z_ψᵋ(x, z)
    wᵢ(x, y, z) = - ∂x_ψᵋ(x, z)
    set!(model, b=bᵢ, u=uᵢ, w=wᵢ)

    Δx = 2π / Nx
    Δt = 0.1 * Δx
    simulation = Simulation(model; Δt, stop_time)

    u, v, w = model.velocities
    Uᵢ = mean(u)

    wall_clock = Ref(time_ns())

    function progress(sim)
        δU = mean(u) / Uᵢ
        elapsed = 1e-9 * (time_ns() - wall_clock[])
        @info @sprintf("Iter: %d, time: %.2e, δU: %.2e, max|w|: %.2e, wall time: %s",
                       iteration(sim), time(sim), δU, maximum(abs, w), prettytime(elapsed))
        wall_clock[] = time_ns()
        return nothing
    end

    simulation.callbacks[:progress] = Callback(progress, IterationInterval(100))

    wizard = TimeStepWizard(cfl=0.5)
    simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

    U = Average(u, dims=(1, 2, 3))
    ξ = ∂z(u) - ∂x(w)

    ke = @at (Center, Center, Center) (u^2 + v^2 + w^2) / 2
    KE = Average(ke, dims=(1, 2, 3))

    simulation.output_writers[:fields] =
        JLD2OutputWriter(model, merge(model.velocities, model.tracers, (; ξ, U, KE));
                         schedule = TimeInterval(save_interval),
                         with_halos = true,
                         filename,
                         overwrite_existing = true)

    @info "Made a simulation of"
    @show model

    @info "The grid is"
    @show model.grid

    @info "The x-velocity immersed and bottom boundary conditions are"
    @show boundary_condition
    @show model.velocities.u.boundary_conditions.immersed
    @show model.velocities.u.boundary_conditions.bottom

    return simulation
end

#####
##### Run them!
#####

Nx = 128
stop_time = 20.0

experiments = ["reference", "no_slip", "free_slip", "bottom_drag"]
Nexp = length(experiments)

function filename(exp, Nx, h)
    exp == "reference" && (h = 0.0)
    return @sprintf("hills_%s_%d_h%d", exp, Nx, 10h)
end

for exp in experiments
    h = exp == "reference" ? 0.0 : 0.3
    boundary_condition = exp == "reference" ? :no_slip : Symbol(exp)
    reference_sim = hilly_simulation(; stop_time, Nx, filename=filename(exp, Nx, h), h, boundary_condition)
    run!(reference_sim)
end

#####
##### Plot results
#####

h = 0.3
ξ  = Dict(exp => FieldTimeSeries(filename(exp, Nx, h) * ".jld2", "ξ")  for exp in experiments)
U  = Dict(exp => FieldTimeSeries(filename(exp, Nx, h) * ".jld2", "U")  for exp in experiments)
KE = Dict(exp => FieldTimeSeries(filename(exp, Nx, h) * ".jld2", "KE") for exp in experiments)

Nt = minimum(length(ξ[exp].times) for exp in experiments)
t = ξ["reference"].times[1:Nt]
δU_series(U) = [(U[1, 1, 1, n] - U[1, 1, 1, 1]) / U[1, 1, 1, 1] for n = 1:Nt]
δK_series(K) = [(K[1, 1, 1, n] - K[1, 1, 1, 1]) / K[1, 1, 1, 1] for n = 1:Nt]
δU = Dict(exp => δU_series(u) for (exp, u) in U)
δKE = Dict(exp => δK_series(k) for (exp, k) in KE)

fig = Figure(resolution=(1800, 1200))
ax = Dict(exp => Axis(fig[i+1, 2:4], aspect=2π, xlabel="x", ylabel="z", title=exp)
          for (i, exp) in enumerate(experiments))

axu = Axis(fig[2:3, 5], xlabel="t", ylabel="Total momentum")
axe = Axis(fig[4:5, 5], xlabel="t", ylabel="Total kinetic energy")

slider = Slider(fig[Nexp+2, 2:4], range=1:Nt, startvalue=1)
n = slider.value

# Title
title = @lift string("Flow over hills at t = ", prettysummary(t[$n]))
Label(fig[1, 1:5], title)

# Vorticity heatmaps
ξi(ξ) = @lift begin
    ξn = ξ[$n]
    mask_immersed_field!(ξn, NaN)
    interior(ξn, :, 1, :)
end

ξⁱ = Dict(exp => ξi(ξ[exp]) for exp in experiments)
x, y, z = nodes(ξ["reference"])
ξmax = maximum(abs, ξ["reference"])
ξlim = ξmax / 50

hm = Dict(exp => heatmap!(ax[exp], x, z, ξⁱ[exp], colorrange=(-ξlim, ξlim), colormap=:redblue)
          for exp in experiments)

cb = Colorbar(fig[2:5, 1], hm["reference"], vertical=true, flipaxis=true, label="Vorticity, ∂z(u) - ∂x(w)")

# Momentum and energy plots
for exp in experiments
    lines!(axu, t, δU[exp], label=exp)
    lines!(axe, t, δKE[exp], label=exp)
end

axislegend(axu, position=:lb)
axislegend(axe, position=:rt)

tn = @lift t[$n]
min_δU = minimum(minimum(δ) for δ in values(δU))
min_δK = minimum(minimum(δ) for δ in values(δKE))
vlines!(axu, tn, ymin=min_δU, ymax=1.0)
vlines!(axe, tn, ymin=min_δK, ymax=1.0)

display(fig)

moviename = @sprintf("flow_over_hills_%dd_h%d.mp4", Nx, 10h)
record(fig, moviename, 1:Nt, framerate=24) do nn
    n[] = nn
end


export quadratic_estimate

"""

    quadratic_estimate(ds::DataSet, which; wiener_filtered=true)
    quadratic_estimate((ds1::DataSet,ds2::DataSet), which; wiener_filtered=true)

Compute quadratic estimate of ϕ given data.

The `ds` or `(ds1,ds2)` tuple contain the DataSet object(s) which houses the
data and covariances used in the estimate. Note that only the Fourier-diagonal
approximations for the beam, mask, and noise,, i.e. `ds.B̂`, `ds.M̂`, and
`ds.Cn̂`, are accounted for. To account full operators (if they are not actually
Fourier-diagonal), you should compute the impact using Monte Carlo.

If a tuple is passed in, the result will come from correlating the data from
`ds1` with that from `ds2`, which can be useful for debugging / isolating
various noise terms.

An optional keyword argument `AL` can be passed in in case the QE normalization
was already computed, in which case it won't be recomputed during the
calculation.

Returns a NamedTuple `(ϕqe, AL, Nϕ)` where `ϕqe` is the (possibly Wiener
filtered, depending on `wiener_filtered` option) quadratic estimate, `AL` is the
normalization (which is already applied to ϕqe, it does not need to be applied
again), and `Nϕ` is the analytic N0 noise bias (Nϕ==AL if using unlensed
weights, currently only Nϕ==AL is always returned, no matter the weights)
"""
function quadratic_estimate((ds1,ds2)::Tuple{DataSet{F1},DataSet{F2}}, which; wiener_filtered=true, AL=nothing, weights=:unlensed) where {F1,F2}
    @assert weights in (:lensed, :unlensed) "weights should be :lensed or :unlensed"
    @assert (which in [:TT, :EE, :EB]) "which='$which' not implemented"
    @assert (ds1.Cf===ds2.Cf && ds1.Cf̃===ds2.Cf̃ && ds1.Cn̂===ds2.Cn̂ && ds1.Cϕ===ds2.Cϕ && ds1.B̂===ds2.B̂) "operators in `ds1` and `ds2` should be the same"
    @assert spin(F1)==spin(F2)
    check_hat_operators(ds1)
    @unpack Cf, Cf̃, Cn̂, Cϕ, B̂, M̂ = ds1()
    (quadratic_estimate_func, pol) = @match which begin
        :TT => (quadratic_estimate_TT, :I)
        :EE => (quadratic_estimate_EE, :P)
        :EB => (quadratic_estimate_EB, :P)
        _   => error("`which` argument to `quadratic_estimate` should be one of (:TT, :EE, :EB)")
    end
    quadratic_estimate_func((ds1.d[pol], ds2.d[pol]), Cf[pol], Cf̃[pol], Cn̂[pol], Cϕ, (M̂*B̂)[pol], wiener_filtered, weights, AL)
end
quadratic_estimate(ds::DataSet, which; kwargs...) = quadratic_estimate((ds,ds), which; kwargs...)
quadratic_estimate(ds::DataSet{<:FlatS0}; kwargs...) = quadratic_estimate(ds, :TT; kwargs...)
quadratic_estimate(ds::DataSet{<:FlatS2}; kwargs...) = quadratic_estimate(ds, :EB; kwargs...) # somewhat arbitraritly make default P estimate be EB

quadratic_estimate_TT(d::Field, args...) = quadratic_estimate_TT((d,d), args...)
quadratic_estimate_EB(d::Field, args...) = quadratic_estimate_EB((d,d), args...)
quadratic_estimate_EE(d::Field, args...) = quadratic_estimate_EE((d,d), args...)

@doc doc"""

All of the terms in the quadratic estimate and normalization expressions look like

    C * l[i] * l̂[j] * l̂[k] * ...

where C is some field or diagonal covariance. For example, there's a term in the EB
estimator that looks like:

    (CE * (CẼ+Cn) \ d[:E])) * l[i] * l̂[j] * l̂[k]

(where note that `l̂[j]` and `l̂[k]` are unit vectors, but `l[i]` is not).  The
function `get_term_memoizer` returns a function `term` which could be called in
the following way to compute this term:

    term((CE * (CẼ+Cn) \ d[:E])), [i], j, k)

(note that the fact that `l[i]` is not a unit vector is specified by putting the
`[i]` index in brackets).

Additionally, all of these terms are symmetric in their indices, i.e. in
`(i,j,k)` in this case. The `term` function is smart about this, and is memoized
so that each unique set of indices is only computed once. This leads to a pretty
drastic speedup for terms with many indices like those that arize in the EE and
EB normalizations, and lets us write code which is both clear and fast without
having to think too hard about these symmetries.

"""
function get_term_memoizer(f)
    Ðf = Ð(f)
    Lfactors() = 1
    Lfactors(inds...) = broadcast!(*, similar(Ðf), getproperty.(getindex.(Ref(∇),inds),:diag)...)
    term(C::Diagonal, inds...) = term(C.diag, inds...)
    term(C::Field, inds...) = term(count((x->x isa Int),inds)/2f0, C, first.(inds)...)
    @sym_memo term(n, C::Field, @sym(inds...)) = Map(nan2zero.(C .* Lfactors(inds...) ./ ∇².diag.^n))
    term
end
ϵ(x...) = levicivita([x...])
inds(n) = collect(product(repeated(1:2,n)...))[:]

function quadratic_estimate_TT((d1,d2)::NTuple{2,FlatS0}, Cf, Cf̃, Cn, Cϕ, TF, wiener_filtered, weights, AL=nothing)

    term = get_term_memoizer(d1)
    ΣTtot = TF^2 * Cf̃ + Cn
    CT = (weights==:unlensed) ? Cf : Cf̃
    @unpack Ωpix = fieldinfo(d1)

    # unnormalized estimate
    ϕqe_unnormalized = @subst -sum(∇[i] * Fourier(term($(ΣTtot\(TF*d1))) * term($(CT*(ΣTtot\(TF*d2))), [i])) for i=1:2)

    # normalization
    if AL == nothing
        AL = @subst begin
            A(i,j) = (
                term($(@. TF^2 * CT^2 / ΣTtot), [i], [j]) * term($(@. TF^2      / ΣTtot)     )
              + term($(@. TF^2 * CT   / ΣTtot), [i]     ) * term($(@. TF^2 * CT / ΣTtot), [j])
            )
            Ωpix * pinv(Diagonal(sum(∇[i].diag .* ∇[j].diag .* Fourier(A(i,j)) for (i,j) in inds(2))))
        end
    end
    Nϕ = AL # true only for unlensed weights

    ϕqe = (wiener_filtered ? (Cϕ*pinv(Cϕ+Nϕ)) : 1) * (AL*ϕqe_unnormalized)
    @namedtuple ϕqe AL Nϕ

end

function quadratic_estimate_EE((d1,d2)::NTuple{2,FlatS2}, Cf, Cf̃, Cn, Cϕ, TF, wiener_filtered, weights, AL=nothing)

    term = get_term_memoizer(d1[:E])
    TF² = TF[:E]^2
    ΣEtot = TF² * Cf̃[:E] + Cn[:E]
    CE = ((weights==:unlensed) ? Cf : Cf̃)[:E]
    @unpack Ωpix = fieldinfo(d1)

    # unnormalized estimate
    ϕqe_unnormalized = @subst begin
        I(i) = -(
            2sum(term($(CE * (ΣEtot \ (TF*d1)[:E])), [i], j, k) * term($((ΣEtot \ (TF*d2)[:E])), j, k) for (j,k) in inds(2))
               - term($(CE * (ΣEtot \ (TF*d1)[:E])), [i]      ) * term($((ΣEtot \ (TF*d2)[:E]))      )
        )
        sum(∇[i] * Fourier(I(i)) for i=1:2)
    end

    # normalization
    if AL == nothing
        AL = @subst begin
            A1(i,j) = -4 * sum( ϵ(m,p,3) * ϵ(n,q,3) * (
                  term($(@. TF² * CE^2 / ΣEtot), [i], [j], k, l, m, n) * term($(@. TF²      / ΣEtot),      k, l, p, q)
                + term($(@. TF² * CE   / ΣEtot), [i],      k, l, m, n) * term($(@. TF² * CE / ΣEtot), [j], k, l, p, q))
                for (k,l,m,n,p,q) in inds(6)
            )
            A2(i,j) = (
                  term($(@. TF² * CE^2 / ΣEtot), [i], [j]) * term($(@. TF²      / ΣEtot)     )
                + term($(@. TF² * CE   / ΣEtot), [i]     ) * term($(@. TF² * CE / ΣEtot), [j])
            )
            Ωpix * pinv(Diagonal(sum(∇[i].diag .* ∇[j].diag .* Fourier(A1(i,j) + A2(i,j)) for i=1:2,j=1:2)))
        end
    end
    Nϕ = AL # true only for unlensed weights

    ϕqe = (wiener_filtered ? (Cϕ*pinv(Cϕ+Nϕ)) : 1) * (AL*ϕqe_unnormalized)
    @namedtuple ϕqe AL Nϕ

end

function quadratic_estimate_EB((d1,d2)::NTuple{2,FlatS2}, Cf, Cf̃, Cn, Cϕ, TF, wiener_filtered, weights, AL=nothing; zeroB=false)

    term = get_term_memoizer(d1[:E])
    CE, CB = getindex.(Ref((weights==:unlensed) ? Cf : Cf̃),(:E,:B))
    TF²E, TF²B = TF[:E]^2, TF[:B]^2
    ΣEtot = TF²E * Cf̃[:E] + Cn[:E]
    ΣBtot = TF²B * Cf̃[:B] + Cn[:B]
    @unpack Ωpix = fieldinfo(d1)

    # unnormalized estimate
    ϕqe_unnormalized = @subst begin
        I(i) = 2 * sum(  ϵ(k,l,3) * (
                           term($(CE * (ΣEtot \ (TF*d1)[:E])), [i], j, k) * term($(     (ΣBtot \ (TF*d2)[:B])),      j, l)
            - (zeroB ? 0 : term($(      ΣEtot \ (TF*d1)[:E]),       j, k) * term($(CB * (ΣBtot \ (TF*d2)[:B])), [i], j, l)))
            for (j,k,l) in inds(3)
        )
        sum(∇[i] * Fourier(I(i)) for i=1:2)
    end

    # normalization
    if AL == nothing
        AL = @subst begin
            @sym_memo A(@sym(i,j)) = 4 * sum( ϵ(m,p,3) * ϵ(n,q,3) * (
                                 term($(@. TF²E * CE^2 / ΣEtot), [i], [j], k, l, m, n) * term($(@. TF²B        / ΣBtot),           k, l, p, q)
                + (zeroB ? 0 : -2term($(@. TF²E * CE   / ΣEtot), [i],      k, l, m, n) * term($(@. TF²B * CB   / ΣBtot),      [j], k, l, p, q))
                + (zeroB ? 0 :   term($(@. TF²E        / ΣEtot),           k, l, m, n) * term($(@. TF²B * CB^2 / ΣBtot), [i], [j], k, l, p, q)))
                for (k,l,m,n,p,q) in inds(6)
            )
            Ωpix * pinv(Diagonal(sum(∇[i].diag .* ∇[j].diag .* Fourier(A(i,j)) for i=1:2,j=1:2)))
        end
    end
    Nϕ = AL # true only for unlensed weights

    ϕqe = (wiener_filtered ? (Cϕ*pinv(Cϕ+Nϕ)) : 1) * (AL * ϕqe_unnormalized)
    @namedtuple ϕqe AL Nϕ

end


include("../examples/smiley_examples.jl")
using .SmileyExample22, .SmileyExample52, .SmileyExample54, .SmileyExample55

using BenchmarkTools
using ForwardDiff
using IntervalArithmetic
using IntervalRootFinding

import Random

const SUITE = BenchmarkGroup()

Random.seed!(0)  # Seed the RNG to get consistent results
tol = 1e-15

include("dietmar_ratz_functions.jl")

S = SUITE["Smiley"] = BenchmarkGroup()
for example in (SmileyExample22, SmileyExample52, SmileyExample54) #, SmileyExample55)
    s = S[example.title] = BenchmarkGroup()
    for method in (Newton, Krawczyk)
        s[string(method)] = @benchmarkable roots($(example.f), $(example.region), $method, $tol)
    end
end

S = SUITE["Rastigrin stationary points"] = BenchmarkGroup()

# Rastrigin function:
const A = 10

f(x, y) = 2A + x^2 - A*cos(2π*x) + y^2 - A*cos(2π*y)
f(X) = f(X...)

ForwardDiff.gradient(f, X::IntervalBox) = ForwardDiff.gradient(f, X.v)
∇f = X -> ForwardDiff.gradient(f, X)

L = 5.0
X = IntervalBox(-L..(L+1), 2)

for method in (Newton, Krawczyk)
    S[string(method)] = @benchmarkable roots($(∇(f)), $X, $method, 1e-5)
end

S = SUITE["Linear equations"] = BenchmarkGroup()

sizes = (2, 5, 10)

for n in sizes
    s = S["n = $n"] = BenchmarkGroup()
    A = Interval.(randn(n, n))
    b = Interval.(randn(n))

    s["Gauss seidel"] = @benchmarkable gauss_seidel_interval($A, $b)
    s["Gauss seidel contractor"] = @benchmarkable gauss_seidel_contractor($A, $b)
    s["Gauss elimination"] = @benchmarkable gauss_elimination_interval($A, $b)
end

S = SUITE["Dietmar-Ratz"] = BenchmarkGroup()
X = Interval(0.75, 1.75)

for (k, dr) in enumerate(dr_functions)
    s = S["Dietmar-Ratz $k"] = BenchmarkGroup()

    if k != 8  # dr8 is excluded as it has too many roots
        for method in (Newton, Krawczyk)
            s[string(method)] = @benchmarkable roots($dr, $X, $method, $tol)
        end
    end

    s["Automatic differentiation"] = @benchmarkable ForwardDiff.derivative($dr, $X)
    s["Slope expansion"] = @benchmarkable slope($dr, $X, $(mid(X)))
end


"""
    College(j, f, t, ft, omf)

Contains a college's index `j`, admissions probability `f`, utility value `t`,
their product `ft`, and `1 - f = omf`. Used only by `applicationorder()`.
"""
struct College
    j::Int
    f::Float64
    t::Float64
    ft::Float64
    omf::Float64
end

# Overload isless() so that the heap is ordered by expected utility.
isless(c1::College, c2::College) = isless(c1.ft, c2.ft)

"""
    applicationorder_list(mkt::SameCostsMarket) -> X, V

Produce the optimal application order `X` and associated valuations `V`
for the [`SameCostsMarket`](@ref) defined by `mkt`. Uses a list data structure;
typically faster than the equivalent [`applicationorder_heap`](@ref).

All `SameCostsMarket`s satisfy a nestedness property, meaning that the optimal
portfolio when `mkt.h = h` is given by the first `h` entries of the optimal portfolio
when `mkt.h = mkt.m`. For example:

```julia-repl
julia> mkt = SameCostsMarket([0.2, 0.5, 0.1, 0.6, 0.1], [1, 4, 9, 1, 8], 5);

julia> x, v = applicationorder_list(mkt)
([2, 3, 5, 4, 1], [2.0, 2.7, 3.24, 3.483, 3.5154])

julia> x[1:4], v[4] # optimal portfolio and valuation for h = 4
([2, 3, 5, 4], 3.483)
```
"""
function applicationorder_list(
    mkt::SameCostsMarket;
    verbose::Bool=false
)::Tuple{Vector{Int},Vector{Float64}}
    apporder = zeros(Int, mkt.h)
    v = zeros(mkt.h)
    mkt_list = College[College(j, mkt.f[j], mkt.t[j], mkt.ft[j], mkt.omf[j]) for j in 1:mkt.m]

    c_best::College, idx_best::Int = findmax(mkt_list)
    @inbounds for j in 1:mkt.h
        if verbose
            println("Iteration $j")
            println("  ft: ", [mkt_list[i].ft for i in 1:length(mkt_list)])
            println("  Add school $(c_best.j)")
        end

        v[j] = get(v, j - 1, 0) + c_best.ft
        apporder[j] = c_best.j

        next_c_best::College = College(0, 1.0, -1.0, -1.0, 0.0)
        next_idx_best::Int = 0

        for i in 1:idx_best-1
            mkt_list[i] =
                College(
                    mkt_list[i].j,
                    mkt_list[i].f,
                    mkt_list[i].t * c_best.omf,
                    mkt_list[i].ft * c_best.omf,
                    mkt_list[i].omf
                )

            if isless(next_c_best, mkt_list[i])
                next_c_best = mkt_list[i]
                next_idx_best = i
            end
        end

        for i in idx_best+1:length(mkt_list)
            mkt_list[i-1] =
                College(
                    mkt_list[i].j,
                    mkt_list[i].f,
                    mkt_list[i].t - c_best.ft,
                    mkt_list[i].ft - mkt_list[i].f * c_best.ft,
                    mkt_list[i].omf
                )

            if isless(next_c_best, mkt_list[i-1])
                next_c_best = mkt_list[i-1]
                next_idx_best = i - 1
            end
        end

        c_best, idx_best = next_c_best, next_idx_best

        pop!(mkt_list)
    end

    return mkt.perm[apporder], v
end

"""
    applicationorder_list(mkt::SameCostsMarket) -> X, V

Produce the optimal application order `X` and associated valuations `V`
for the [`SameCostsMarket`](@ref) defined by `mkt`. Uses a heap data structure;
typically the equivalent [`applicationorder_list`](@ref) is faster.

All `SameCostsMarket`s satisfy a nestedness property, meaning that the optimal
portfolio when `mkt.h = h` is given by the first `h` entries of the optimal portfolio
when `mkt.h = mkt.m`. For example:

```julia-repl
julia> mkt = SameCostsMarket([0.2, 0.5, 0.1, 0.6, 0.1], [1, 4, 9, 1, 8], 5);

julia> x, v = applicationorder_heap(mkt)
([2, 3, 5, 4, 1], [2.0, 2.7, 3.24, 3.483, 3.5154])

julia> x[1:4], v[4] # optimal portfolio and valuation for h = 4
([2, 3, 5, 4], 3.483)
```
"""
function applicationorder_heap(mkt::SameCostsMarket)::Tuple{Vector{Int},Vector{Float64}}
    apporder = zeros(Int, mkt.h)
    v = zeros(mkt.h)

    mkt_heap = BinaryMaxHeap{College}(collect(College(j, mkt.f[j], mkt.t[j], mkt.ft[j], mkt.omf[j]) for j in 1:mkt.m))

    @inbounds for j in 1:mkt.h
        c_k = first(mkt_heap)
        v[j] = get(v, j - 1, 0) + c_k.ft
        apporder[j] = c_k.j

        # These two implementations perform almost identically
        mkt_heap = BinaryMaxHeap{College}(
            # College{T}[
            #     College(
            #         c.j,
            #         c.f,
            #         c.j < c_k.j ? c.t * c_k.omf : c.t - c_k.ft,
            #         c.j < c_k.j ? c.ft * c_k.omf : c.ft - c.f * c_k.ft,
            #         c.omf
            #     )
            #     for c in mkt_heap.valtree if c.j != c_k.j
            # ]
            map(filter(c -> c.j != c_k.j, mkt_heap.valtree)) do c
                College(
                    c.j,
                    c.f,
                    c.j < c_k.j ? c.t * c_k.omf : c.t - c_k.ft,
                    c.j < c_k.j ? c.ft * c_k.omf : c.ft - c.f * c_k.ft,
                    c.omf
                )
            end
        )
    end

    return mkt.perm[apporder], v
end

# Used by dynamic program below
@inbounds function V_recursor!(
    V_dict::Dict{Tuple{Int,Int},Float64},
    j::Int,
    h::Int,
    mkt::VariedCostsMarket
)::Float64
    return get(V_dict, (j, h)) do
        if j == 0 || h == 0
            return 0.0
        elseif h < mkt.g[j]
            push!(V_dict, (j, h) => V_recursor!(V_dict, j - 1, h, mkt))
            return V_dict[(j, h)]
        else
            jmo = j - 1
            push!(V_dict, (j, h) => max(
                V_recursor!(V_dict, jmo, h, mkt),
                mkt.omf[j] * V_recursor!(V_dict, jmo, h - mkt.g[j], mkt) + mkt.ft[j]
            ))
            return V_dict[(j, h)]
        end
    end
end

"""
    optimalportfolio_dynamicprogram(mkt::VariedCostsMarket) -> X, v

Use the dynamic program on application costs to produce the optimal portfolio `X` and associated
value `v` for the [`VariedCostsMarket`](@ref) defined by `mkt`.

```julia-repl
julia> mkt = VariedCostsMarket([0.2, 0.5, 0.1, 0.6, 0.1], [1, 4, 9, 1, 8], [2, 4, 2, 5, 1], 8);

julia> optimalportfolio_dynamicprogram(mkt)
([3, 5, 2], 3.24)
```
"""
function optimalportfolio_dynamicprogram(
    mkt::VariedCostsMarket;
    verbose::Bool=false
)::Tuple{Vector{Int},Float64}
    V_dict = Dict{Tuple{Int,Int},Float64}()
    sizehint!(V_dict, mkt.m * mkt.m ÷ 2)

    v = V_recursor!(V_dict, mkt.m, mkt.H, mkt)

    h = mkt.H
    X = Int[]

    for j in reverse(1:mkt.m)
        if V_recursor!(V_dict, j - 1, h, mkt) < V_recursor!(V_dict, j, h, mkt)
            push!(X, j)
            h -= mkt.g[j]
        end
    end

    if verbose
        V_table = Array{Union{Missing,Float64}}(missing, mkt.m, mkt.H)
        for (j, h) in keys(V_dict)
            if 0 < j ≤ mkt.m && 0 < h ≤ mkt.H
                V_table[j, h] = V_dict[(j, h)]
            end
        end

        display(V_table)
    end

    return mkt.perm[X], v
end