function solve_nde(nde, NN, T₀, alg, nde_params)
    nn_weights, _ = Flux.destructure(NN)
    # See: https://github.com/SciML/DiffEqFlux.jl/blob/449efcecfc11f1eab65d0e467cf57db9f5a5dbec/src/neural_de.jl#L68
    return solve(nde, alg, reltol=1e-4, u0=T₀, p=[nn_weights; nde_params],
                 sense=InterpolatingAdjoint(autojacvec=ZygoteVJP(), checkpointing=true))
end

function solve_nde(ds, NN, NDEType, algorithm, T_scaling, wT_scaling; T₀=nothing)

    T = ds["T"]
    wT = ds["wT"]
    Nz = size(T, 3)
    zc = znodes(T)
    zf = znodes(wT)

    H = abs(zf[1]) # Domain depth/height

    Δẑ = diff(zc)[1] / H  # Non-dimensional grid spacing
    Dzᶠ = Dᶠ(Nz, Δẑ) # Differentiation matrix operator

    nde_params = FreeConvectionNDEParameters(ds, T_scaling, wT_scaling)
    nde = NDEType(NN, ds)

    if isnothing(T₀)
        T₀ = T_scaling.(interior(T)[1, 1, :, 1])
    end

    T = solve_nde(nde, NN, T₀, algorithm, nde_params) |> Array

    bottom_flux, top_flux, _ = nde_params

    Nz, Nt = size(T)
    wT = zeros(Nz+1, Nt)

    for n in 1:Nt
        T_n = T[:, n]

        wT_NN_interior = NN(T_n)
        wT_NN_n = cat(bottom_flux, wT_NN_interior, top_flux, dims=1)

        if NDEType == FreeConvectionNDE
            @. wT[:, n] = wT_NN_n
        elseif NDEType == ConvectiveAdjustmentNDE
            ∂T∂z_n = Dzᶠ * T_n
            K∂T∂z_n = min.(0, 10 * ∂T∂z_n)
            @. wT[:, n] = wT_NN_n - K∂T∂z_n
        end
    end

    return (T=inv(T_scaling).(T), wT=inv(wT_scaling).(wT))
end


using Flux, Zygote, ForwardDiff, Random, Optim
using DifferentialEquations, OrdinaryDiffEq
using Plots
using Flux.Optimise: update!
using BSON: @save, @load
using ProgressBars
using Printf
using LinearAlgebra
using Statistics
using DiffEqFlux

n_epoch = 1000
nplot = 100
npoint = 100
ns = 5
opt = ADAMW(0.005, (0.9, 0.999), 1.f-6)

function f(y, k, t)
    dydt = [-2 * k[1] * y[1]^2 - k[2] * y[1],
            k[1] * y[1]^2 - k[4] * y[2] * y[4],
            k[2] * y[1] - k[3] * y[3],
            k[3] * y[3] - k[4] * y[2] * y[4],
            k[4] * y[2] * y[4]]
    return dydt
end

u0 = Float32[1, 1, 0, 0, 0]
k = Float32[0.1, 0.2, 0.13, 0.3]
tspan = Float32[0, 40.0]
dt = tspan[2] / npoint
ts = tspan[1]:dt:tspan[2]

prob = ODEProblem(f, u0, tspan, k)
sol = solve(prob, saveat=ts);

plt = scatter(sol, yscale=:identity);
png(plt, "figs/sol");

hls = 8;
phi_ = Flux.Chain(Dense(1, hls, tanh),
                Dense(hls, hls, tanh),
                Dense(hls, hls, tanh),
                Dense(hls, hls, tanh),
                Dense(hls, ns))

θ, re = Flux.destructure(phi_)

function phi(t, θ)
    return u0 .+ t .* re(θ)([t])
end
phi(1, θ)

dfdx = (t, θ) -> ForwardDiff.derivative(t -> phi(t, θ), t)
dfdx(1, θ)

function inner_loss(t, θ)
    mean(abs2, dfdx(t, θ) - f(phi(t, θ), k, t))
end

loss(θ) = mean(abs2, [inner_loss(t, θ) for t in ts])

l_loss_train = []
l_loss_val = []
l_grad = []
iter = 1
cb = function (p, loss_train, loss_val, g_norm)
    global iter
    global l_loss_train, l_loss_val, l_grad, iter
    push!(l_loss_train, loss_train)
    push!(l_loss_val, loss_val)
    push!(l_grad, g_norm)

    if iter % nplot == 0
        pred = hcat([phi(t, θ) for t in ts]...)
        l_plt = []
        for i in 1:ns
            plt = scatter(ts, sol[i, :], label="label");
            plot!(plt, ts, pred[i, :], label="PINN")
            xlabel!(plt, "Time")
            ylabel!(plt, "x$i")
            push!(l_plt, plt)
        end
        plt_all = plot(l_plt..., legend=:best, size=(800, 800))
        png(plt_all, "figs/pinn")

        plt_loss = plot(l_loss_train, xscale=:identity, yscale=:log10, label="Training")
        plot!(plt_loss, l_loss_val, xscale=:identity, yscale=:log10, label="Validation")
        plt_grad = plot(l_grad, xscale=:identity, yscale=:log10, label="grad_norm")
        xlabel!(plt_loss, "Epoch")
        xlabel!(plt_grad, "Epoch")
        ylabel!(plt_loss, "Loss")
        ylabel!(plt_grad, "Gradient Norm")
        plt_all = plot([plt_loss, plt_grad]..., legend=:top)
        png(plt_all, "figs/loss_grad")
    end
    iter = iter + 1
end

epochs = ProgressBar(1:n_epoch);
for epoch in epochs
    global θ
    l = loss(θ)
    grad = ForwardDiff.gradient(x -> loss(x), θ)
    update!(opt, θ, grad)
    set_description(epochs, string(@sprintf("Loss %.4e", l)))
    cb(θ, l, l, norm(grad, 2))
end


r = @theory begin
    foo(x,y) => 2*x%y
    foo(x,y) => sin(x)
    sin(x) => foo(x,x)
end
@testset "Basic Equalities 1" begin
    @test (@areequal r foo(b,c) foo(d,d)) == false
end

r = @theory begin
    a * 1 => foo
    a * 2 => bar
    1 * a => baz
    2 * a => mag
end
@testset "Matching Literals" begin
    @test (@areequal r a * 1 foo) == true
    @test (@areequal r a * 2 foo) == false
    @test (@areequal r a * 1 bar) == false
    @test (@areequal r a * 2 bar) == true

    @test (@areequal r 1 * a baz) == true
    @test (@areequal r 2 * a baz) == false
    @test (@areequal r 1 * a mag) == false
    @test (@areequal r 2 * a mag) == true
end

comm_monoid = @commutative_monoid (*) 1
@testset "Basic Equalities - Commutative Monoid" begin
    @test true == (@areequal comm_monoid a*(c*(1*d)) c*(1*(d*a)) )
    @test true == (@areequal comm_monoid x*y y*x )
    @test true == (@areequal comm_monoid (x*x)*(x * 1) x*(x*x) )
end

comm_group = @abelian_group (+) 0 inv
t = comm_monoid ∪ comm_group ∪ distrib(:(*), :(+))
@testset "Basic Equalities - Comm. Monoid, Abelian Group, Distributivity" begin
    @test true == (@areequal t (a * b) + (a * c) a*(b+c) )
    @test true == (@areequal t a*(c*(1*d)) c*(1*(d*a)) )
    @test true == (@areequal t a+(b*(c*d)) ((d*c)*b)+a )
    @test true == (@areequal t (x+y)*(a+b) ((a*(x+y)) + b*(x+y)) ((x*(a+b)) + y*(a+b)) )
    @test true == (@areequal t (((x*a + x*b) + y*a) + y*b) (x+y)*(a+b) )
    @test true == (@areequal t a+(b*(c*d)) ((d*c)*b)+a )
    @test true == (@areequal t a+inv(a) 0 (x*y)+inv(x*y) 1*0 )
end

@testset "Basic Equalities - False statements" begin
    @test false == (@areequal t (a * b) + (a * c) a*(b+a))
    @test false == (@areequal t (a * c) + (a * c) a*(b+c) )
    @test false == (@areequal t a*(c*c) c*(1*(d*a)) )
    @test false == (@areequal t c+(b*(c*d)) ((d*c)*b)+a )
    @test false == (@areequal t (x+y)*(a+c) ((a*(x+y)) + b*(x+y)) )
    @test false == (@areequal t ((x*(a+b)) + y*(a+b)) (x+y)*(a+c) )
    @test false == (@areequal t (((x*a + x*b) + y*a) + y*b) (x+y)*(a+x) )
    @test false == (@areequal t a+(b*(c*a)) ((d*c)*b)+a )
    @test false == (@areequal t a+inv(a) a )
    @test false == (@areequal t (x*y)+inv(x*y) 1 )
end

# Issue 21
simp_theory = @theory begin
    munit() => :foo
end
G = EGraph( :(munit()) )
saturate!(G, simp_theory, timeout=1)

module Bar
   foo = 42
   using Metatheory
   @metatheory_init
   export foo
end

module Foo
   foo = 12
   using Metatheory
   @metatheory_init
end

t = @theory begin
   :woo |> foo
end

g = EGraph(:woo);
saturate!(g, t; mod=Bar);
saturate!(g, t; mod=Foo);
foo = 12

@testset "Different modules" begin
    @test @areequalg g t 42 12
end

# expr = cleanast(:(1 * 1 * 1 * 1 * 1 * zoo * 1 * 1 * foo * 1))
#
# G = EGraph(expr)
#
# @time saturate!(G, comm_monoid)
#
# G.H |> display


using Test
using LinearAlgebra

points2D = [SVector(0.1, 0.1), #1
            SVector(1.0, 1.0), #2
            SVector(0.2, 0.2), #3
            SVector(1.1, 1.1)] #4

root = create_tree(points2D, KMeansTreeOptions())

@test length(root.children) == 2
@test root.children[1].data[1] == 1
@test root.children[1].data[2] == 3
@test root.children[2].data[1] == 2
@test root.children[2].data[2] == 4
@test root.children[1].radius ≈ norm([0.05 0.05]) atol=1e-15


# Load modules
using Plots, LaTeXStrings
pyplot()
using ThinFilmsTools
using Optim

# Wavelength range [nm]
λ = 400:1000
# Angle of incidence [degrees]
θ = [5.]
# Polarisation (1.0 = p, 0.0 = s, between 0.0 and 1.0 = average)
pol = 0.5
beam = PlaneWave(λ, θ; p=pol)

# Refractive indices of incident (0) and substrate (2)
incident = RIdb.air(beam.λ)
emergent = RIdb.silicon(beam.λ)

# Define the RI model to use
layers = [
    LayerTMMO(incident),
    ModelFit(:looyenga; N=(ninc=incident, nhost=emergent)),
    ModelFit(:looyenga; N=(ninc=incident, nhost=emergent)),
    LayerTMMO(emergent),
]

# Set the order of the layers (build the system) to fit
order = [
    1 # incident medium
    repeat([2, 3], 9) # top DBR of the MC
    4 # substrate
]

# Absolute reflectance spectrum to fit stored in Utils
Rexp = SpectraDB.bragg_spectrum(beam.λ)
# plot(beam.λ, Rexp)
# gui()

options = Optim.Options(
    g_abstol=1e-8, g_reltol=1e-8, iterations=10^5, show_trace=true, store_trace=true,
);

# Seeds for each ModelFit layer defined above (without alpha)
seed = [
    [69, 0.54], # layers[2]
    [119.0, 0.77], # layers[3]
]

solOptim = fit_tmm_optics(
    Reflectance(Rexp), seed, beam, layers;
    order=order, options=options, alg=LBFGS(),
)

plot(FitSpectrum(), solOptim.beam.λ, solOptim.spectrumExp, solOptim.spectrumFit)
gui()

# Seeds for each ModelFit layer defined above plus the alpha (with alpha)
seed2 = [
    [69, 0.54], # layers[2]
    [119.0, 0.77], # layers[3]
    [0.995], # alpha
]

solOptim2 = fit_tmm_optics(
    Reflectance(Rexp), seed2, beam, layers;
    order=order, options=options, alg=LBFGS(), alpha=true,
)

plot(FitSpectrum(), solOptim2.beam.λ, solOptim2.spectrumExp, solOptim2.spectrumFit)
gui()


module TestBalls

println("\nCollidingBalls: Demonstrating the use of allInstances to set up contact force between any number of balls")

using Modia
using ModiaMath

const k=10000
const d=0
const radius=0.05

springDamper(r, v) = k*r + d*v

function getForce(r, v, positions,
  velocities, contactLaw)
  # Calculate the forces on this body
  force = zeros(2)
  for i in 1:length(positions)
    pos = positions[i]
    vel = velocities[i]
    if r != pos
      delta = r - pos
      deltaV = v - vel
      f = if norm(delta) < 2*radius;
        -contactLaw((norm(delta)-
          2*radius)*delta/norm(delta), deltaV) else
        zeros(2) end
      force += f
    end
  end
  return force
end

@model Ball begin
  r = Var()
  v = Var()
  f = Var()
  m = 1.0
  event = Boolean(size=(), start=false)
@equations begin
  der(r) = v
  m*der(v) = f
  f = getForce(r, v, allInstances(r),
#    allInstances(v), (r,v) -> (k*r + d*v))
    allInstances(v), springDamper)
    event = Modia.Synchronous.positive(f[1])
  end
end

@model Balls1 begin
  b1 = Ball(r = Var(start=[0.0,2]), v = Var(start=[1,0]))
  b2 = Ball(r = Var(start=[0.5,2]), v = Var(start=[-1,0]))
  b3 = Ball(r = Var(start=[1.0,2]), v = Var(start=[0,0]))
end

#=
result = checkSimulation(Balls1, 1.5, "b3.r", 2.0, expandArrayIncidence=true, storeEliminated=false)
ModiaMath.plot(result, ("b1.r", "b2.r", "b3.r"), figure=1)
=#

# -----------------------------

@model Balls2 begin
  b1 = Ball(r = Var(start=[0.0,2]), v = Var(start=[1,0]))
  b2 = Ball(r = Var(start=[0.5,2]), v = Var(start=[0,0]))
  b3 = Ball(r = Var(start=[1.0,2]), v = Var(start=[-1,0]))
end

#=
result = checkSimulation(Balls2, 1.5, "b3.r", 2.0, expandArrayIncidence=true, storeEliminated=false)
ModiaMath.plot(result, ("b1.r", "b2.r", "b3.r"), figure=2)
=#

# -----------------------------

@model Balls3 begin
  b1 = Ball(r = Var(start=[1.0,0]), v = Var(start=[1,0]))
  b2 = Ball(r = Var(start=[2.0,0]), v = Var(start=[0,0]))
  b3 = Ball(r = Var(start=[3.0,0]), v = Var(start=[0,0]))
  b4 = Ball(r = Var(start=[4.0,0]), v = Var(start=[0,0]))
  b5 = Ball(r = Var(start=[5.0,0]), v = Var(start=[0,0]))
  b6 = Ball(r = Var(start=[10.0,0]), v = Var(start=[-1,0]))
end

result = simulate(Balls3, 8, expandArrayIncidence=true, storeEliminated=false)
ModiaMath.plot(result, Tuple(["b$i.r" for i in 1:5]), figure=1, heading="Colliding balls", grid=false, figure=19)

# -----------------------------

#=
nBalls=4
@model Balls begin
#  balls = [Ball(r = Var(start=[i,2]), v = Var(start=if i==1; [1,0] else [0,0] end)) for i in 1:nBalls]
  balls = [Ball() for i in 1:nBalls]
end

result = simulate(Balls, 1.5, expandArrayIncidence=true, storeEliminated=false, logTranslation=true, logOnFile=false)
=#

end


"""
    simulate linear feedback policy
"""
function simulate_linear_feedback(
        model,
        K,
        x̄, ū,
        x_ref, u_ref,
        Q, R,
        T_sim, Δt,
        x1,
        w;
        ul = -Inf * ones(length(ū[1])),
        uu = Inf * ones(length(ū[1])))

    T = length(x̄)
    times = [(t - 1) * Δt for t = 1:T-1]
    tf = Δt * (T-1)
    t_sim = range(0, stop = tf, length = T_sim)
    t_ctrl = range(0, stop = tf, length = T)
    dt_sim = tf / (T_sim - 1)

    x_sim = [x1]
    u_sim = []

    J = 0.0
    Jx = 0.0
    Ju = 0.0

    A_state = Array(hcat(x̄...))
    A_ctrl = Array(hcat(ū...))

    for tt = 1:T_sim-1
        t = t_sim[tt]
        k = searchsortedlast(times, t)
        x = x_sim[end]

        x_cubic = Array(zero(x̄[1]))

        for i = 1:length(x̄[1])
            interp_cubic = CubicSplineInterpolation(t_ctrl, A_state[i,:])
            x_cubic[i] = interp_cubic(t)
        end

        u = ū[k] + K[k] * (x - x_cubic)

        # clip controls
        u = max.(u, ul)
        u = min.(u, uu)

        push!(x_sim, fd(model, x, u, w[k], dt_sim, tt))
        push!(u_sim, u)

        J += (x_sim[end] - x_ref[k])' * Q[k + 1] * (x_sim[end] - x_ref[k])
        J += (u_sim[end] - u_ref[k])' * R[k] * (u_sim[end] - u_ref[k])
        Jx += (x_sim[end] - x_ref[k])' * Q[k + 1] * (x_sim[end] - x_ref[k])
        Ju += (u_sim[end] - u_ref[k])' * R[k] * (u_sim[end] - u_ref[k])
    end

    return x_sim, u_sim, J / (T_sim - 1), Jx / (T_sim - 1), Ju / (T_sim - 1)
end

function simulate_policy(
        model,
        θ,
        x_ref, u_ref,
        Q, R,
        T_sim, Δt,
        x1,
        w;
        ul = -Inf * ones(length(ū[1])),
        uu = Inf * ones(length(ū[1])))

    T = length(x_ref)
    times = [(t - 1) * Δt for t = 1:T-1]
    tf = Δt * (T-1)
    t_sim = range(0, stop = tf, length = T_sim)
    t_ctrl = range(0, stop = tf, length = T)
    dt_sim = tf / (T_sim - 1)

    x_sim = [x1]
    u_sim = []

    J = 0.0
    Jx = 0.0
    Ju = 0.0

    # A_state = Array(hcat(x̄...))
    # A_ctrl = Array(hcat(ū...))

    for tt = 1:T_sim-1
        t = t_sim[tt]
        k = searchsortedlast(times, t + 1.0e-16)
        x = x_sim[end]

        # println("t $t")
        # println("k $k")
        # println("")

        # x_cubic = Array(zero(x̄[1]))
        #
        # for i = 1:length(x̄[1])
        #     interp_cubic = CubicSplineInterpolation(t_ctrl, A_state[i,:])
        #     x_cubic[i] = interp_cubic(t)
        # end

        # # clip controls
        # u = max.(u, ul)
        # u = min.(u, uu)
        push!(x_sim, fd(model, x, [zeros(model.m); θ[k]], w[k], dt_sim, tt))
        push!(u_sim, policy(θ[k], x, nothing, model.n, model.m))

        J += (x_sim[end] - x_ref[k])' * Q[k + 1] * (x_sim[end] - x_ref[k])
        J += (u_sim[end] - u_ref[k])' * R[k] * (u_sim[end] - u_ref[k])
        Jx += (x_sim[end] - x_ref[k])' * Q[k + 1] * (x_sim[end] - x_ref[k])
        Ju += (u_sim[end] - u_ref[k])' * R[k] * (u_sim[end] - u_ref[k])
    end

    return x_sim, u_sim, J / (T_sim - 1), Jx / (T_sim - 1), Ju / (T_sim - 1)
end


@testset "LocalField" begin

  @testset "Creation" begin
    Qx, x = FlintQQ["x"]
    f = x^2-2*x+2
    K, a = local_field(f, 2, 10, "a", Hecke.EisensteinLocalField, check = false)
    @test precision(K) == 20
    @test characteristic(K) == 0
    @test prime(K) == 2

    Kt, t = PolynomialRing(K, "t")
    g = t^2+2
    L, b = local_field(g, "b", Hecke.EisensteinLocalField, check = false)
    @test precision(L) == 2*precision(g)
    @test precision(b) == precision(L)
    @test degree(L) == 2
    @test absolute_degree(L) == 4
    @test prime(L) == 2

    Q2 = PadicField(2, 10)
    Q2s, s = PolynomialRing(Q2, "s")
    f = s^2+s+1
    Ku, c = local_field(f, "s", Hecke.UnramifiedLocalField, check = false)
    @test precision(Ku) == precision(f)
    @test precision(f) == precision(c)
    @test degree(Ku) == 2
    @test absolute_degree(Ku) == 2
    @test inertia_degree(Ku) == 2
    @test ramification_index(Ku) == 1
    @test absolute_ramification_index(Ku) == 1
    @test absolute_inertia_degree(Ku) == 2

    Lu, u = PolynomialRing(L, "u")
    Lu, d = local_field(u^2+u+1, "s", Hecke.UnramifiedLocalField, check = false)
    @test absolute_degree(Lu) == 8
    @test ramification_index(Lu, K) == 2
    @test inertia_degree(Lu, K) == 2

  end

  @testset "Norm" begin
    K = QadicField(3, 4, 10)[1]
    Kx, x = PolynomialRing(K, "x")
    L = eisenstein_extension(x^20+3)[1]
    b = @inferred basis(L)
    for i = 1:10
      r = 1+2*uniformizer(L)^i * sum([rand(1:10)*b[i] for i in 1:5])
      M = @inferred representation_matrix(r)
      n = @inferred norm(r)
      @test n == det(M)
      @test trace(r) == trace(M)
    end
  end

  @testset "Completions" begin
    K, gK = cyclotomic_field(15)
    OK = maximal_order(K)
    lP = prime_decomposition(OK, 3)
    P = lP[1][1]
    Kp, mKp = @inferred Hecke.generic_completion(K, P, 100)
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(elem_in_nf(uniformizer(P)))) == fmpq(1, 2)
    @test valuation(mKp(gen(K)))*2 == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    P = prime_decomposition(OK, 7)[1][1]
    Kp, mKp = @inferred Hecke.unramified_completion(K, P, 100)
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test isone(valuation(mKp(elem_in_nf(uniformizer(P)))))
    @test valuation(mKp(gen(K))) == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    K, gK = cyclotomic_field(9)
    OK = maximal_order(K)
    lP = prime_decomposition(OK, 3)
    P = lP[1][1]
    Kp, mKp = @inferred Hecke.totally_ramified_completion(K, P, 100)
    @test valuation(mKp(elem_in_nf(uniformizer(P)))) == fmpq(1, 6)
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(gen(K)))*2 == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    K, gK = cyclotomic_field(55)
    OK = maximal_order(K)
    lP = prime_decomposition(OK, 11)
    P = lP[1][1]
    Kp, mKp = @inferred Hecke.generic_completion(K, P, 100)
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(elem_in_nf(uniformizer(P)))) == fmpq(1, 10)
    @test valuation(mKp(gen(K)))*10 == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    lP = prime_decomposition(OK, 7)
    P = lP[1][1]
    Kp, mKp = @inferred Hecke.unramified_completion(K, P, 100)
    @test isone(valuation(mKp(elem_in_nf(uniformizer(P)))))
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(gen(K))) == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    Qx, x = PolynomialRing(FlintQQ, "x", cached = false)
    f = x^10 + 5*x^9 - 633*x^8 + 38157*x^7 + 1360601*x^6 - 39808345*x^5 - 464252491*x^4 - 17622187401*x^3 + 2826886632714*x^2 - 96335539805104*x + 1313743135204448
    K, a = number_field(f, check = false, cached = false)
    OK = maximal_order(K)
    P = prime_decomposition(OK, 2)[1][1]
    Kp, mKp = @inferred Hecke.unramified_completion(K, P, 100)
    @test isone(valuation(mKp(elem_in_nf(uniformizer(P)))))
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(gen(K))) == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    lP = prime_decomposition(OK, 11)
    if lP[1][2] == 5
      P = lP[1][1]
    else
      P = lP[2][1]
    end
    Kp, mKp = @inferred Hecke.totally_ramified_completion(K, P, 100)
    @test valuation(mKp(elem_in_nf(uniformizer(P)))) == fmpq(1, 5)
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(gen(K)))*5 == valuation(gen(K), P)
    #@test iszero(defining_polynomial(K)(mKp(gen(K))))

  end

  @testset "Exp and Log" begin
    K = PadicField(2, 100)
    Kx, x = PolynomialRing(K, "x", cached = false)
    L, b = Hecke.eisenstein_extension(x^7+2, "a")
    pi = uniformizer(L)
    @test iszero(log(pi))
    B = basis(L)
    for i = 15:20
      el = sum([rand(FlintZZ, 0:10)*B[j] for j = 1:7])*pi^i
      explog = exp(log(1+el))
      logexp = log(exp(el))
      @test iszero(setprecision(explog, precision(explog)-35) - 1 - el)
      @test iszero(setprecision(logexp, precision(logexp)-35)-el)
    end
  end

end


import Pkg
Pkg.develop(path = joinpath(pwd(), "..", "..", "connectors", "OmegaDistributions"))

using OmegaCore, OmegaMH, OmegaDistributions
using Test
using Distributions
using Random

# function test_bernoulli()
#   rng = MersenneTwister(1)
#   β = 1 ~ Beta(0.5)
#   # bern = 2 ~ Bernoulli(β)
#   bern_(id, ω) = Bernoulli(β(ω))(id, ω)
#   bern = 3 ~ bern_
#   βc = β |ᶜ bern ==ₚ 1.0
#   logdensity(ω) = logpdf(ω)
#   mh(rng, defΩ(), logdensity, 1000)
# end  # stdnormalpdf(x) = logpdf(Normal(0, 1), x)

# function test_beta_bernoulli_constraint()
#   β = 1 ~ Beta(0.5)
#   bern_(id, ω) = Bernoulli(β(ω))(id, ω)
#   bern = 3 ~ bern_
#   βc = β |ᶜ ((bern ==ₚ 1.0) & (β in 0.2..0.4))
#   randsample(βc, 1000; alg = OmegaMH)
# end

# function test_ss_propose(rng = MersenneTwister(0))
#   ω = defω()
#   x = 1 ~ Normal(0, 1)
#   y = 2 ~ Normal(0, 1)
#   xy = @joint x y
#   xy(ω)
#   propose(rng, xy, ω, OmegaMH.SSProposal())
# end

function auto_logenergy(rng = MersenneTwister(0))
  μ = 1 ~ Normal(0, 1)
  x = 2 ~ Normal(μ, 1)
  x_ = 3.1234
  μᵪ = μ |ᶜ x ==ₚ x_

  function prop(rng, ω)
    ω = update(ω, μ, rand(rng, Normal(ω[μ], 1)))
    (ω, 0.0)
  end
  ωinit = defω()
  ωinit[μ] = 0.3
  samples = mh(rng, ω -> logenergy(μᵪ, ω), 1000, ωinit, prop)
end

function test_manual_proposal(rng = MersenneTwister(0))
  μ = 1 ~ Normal(0, 1)
  x = 2 ~ Normal(μ, 1)
  x_ = 3.234
  μᵪ = μ |ᶜ x ==ₚ x_

  function prop(rng, ω)
    ω = update(ω, μ, rand(rng, Normal(ω[μ], 1)))
    (ω, 0.0)
  end
  ωinit = defω()
  ωinit[μ] = 0.3
  logenergy(ω) = logpdf(Normal(0, 1), ω[μ]) + logpdf(Normal(ω[μ], 1), x_)
  samples = mh(rng, logenergy, 1000, ωinit, prop)
end

function test_custom_proposal(rng = MersenneTwister(0))
  x = 1 ~ Normal(0, 1)
  ϵ = 2 ~ Normal(0, 1)
  y(ω) = x(ω) + ϵ(ω)
  function prop(rng, ω)
    q = Normal(0, 1)
    θ = rand(rng, q)
    y_ = ω.y
    xn_, ϵn_ = y_ - θ, θ
    # qlogpdf = logpdf(q, θ)
    ωn = (x = xn_, ϵ = ϵn_, y = y_)
    (ωn, 1)
  end
  ωinit = (x = 2.0, ϵ = 1.0, y = 3.0)
  stdnormalpdf(x) = logpdf(Normal(0, 1), x)
  logenergy(ω) = stdnormalpdf(ω.x) + stdnormalpdf(ω.ϵ)
  samples = mh(rng, logenergy, 10000, ωinit, prop)
end

# TODO
# 1. Convince myself that this is correct
# 2. Reduce / remove dependence of Mh on omega
# 3. Need more flexible interface to create LazyΩ
# 4. Fix allocations
# 5. Automatic proposals

@testset "OmegaMH" begin
  # test_bernoulli()
  # test_beta_bernoulli_constraint()
  # test_ss_propose()
  test_manual_proposal()
  test_custom_proposal()
end


using Roots, Dierckx, NLopt, CoolProp, PyPlot, LsqFit
import JSON
include("rel_perms.jl")
include("frac_flow_funcs.jl")

# read the data files
lowsal=JSON.parsefile("../juliaphreeqc/data/Lowsal_data.json")
paper_title = "fathi2010"
exp_type = "coreFlood"
exp_name = ["SF20-a", "SF20-b"] # formation water and sea water
recovery_data = "coreFloodData" # -> time(day), recovery
core_flood_condition = lowsal["fathi2010"]["coreFlood"]["SF20-a"]
L_core = core_flood_condition["L(mm)"]/1000
D_core = core_flood_condition["D(mm)"]/1000
poros  = core_flood_condition["porosity"]/100
swi     = core_flood_condition["swi"]/100
soi     = core_flood_condition["soi"]/100 # 1-swi
T_c     = core_flood_condition["T(C)"]
oil_type = core_flood_condition["oil"]
mu_oil_25 = lowsal["fathi2010"]["oil"][oil_type]["viscosity(cp)"]/100 # Pa.s
t_day = lowsal["fathi2010"]["coreFloodData"]["SF20-a"]["time(day)"]
R_oil = lowsal["fathi2010"]["coreFloodData"]["SF20-a"]["recovery"]/100
injection_rate_pv_day = core_flood_condition["injectionRate(PV/day)"]
recovery_final = core_flood_condition["recovery(%OOIP)"]/100

so_final = (1-recovery_final)*soi # final oil saturation
pv_inj = t_day[end]*injection_rate_pv_day
t_sec = t_day*24*3600 # s
# conversion of the input data to the right unit
A_core  = π*D_core^2/4
PV_core = L_core*A_core*poros
u_inj   = injection_rate_pv_day/(24*3600)*PV_core/A_core # Darcy velocity
T_K = T_c + 273.15
perm_ave = 0.002e-12 # m^2 permeability
mu_water = PropsSI("viscosity", "T", T_K, "P", 5e5, "H2O")
# My formulation for oil viscosity at higher Temperature
# source: P.Daučík et al., Temperature Dependence of the Viscosity of Hydrocarbon Fractions
# Acta Chimica Slovaca, Vol.1, No. 1, 2008, 43 – 57
# log(mu/mu0) = B(1/T - 1/T0)
# find B for n-Dodecane, and use it for extrapolation
mu_1   = PropsSI("viscosity", "T", 300, "P", 5e5, "n-Dodecane")
mu_2   = PropsSI("viscosity", "T", T_K, "P", 5e5, "n-Dodecane")
B_visc = log(mu_1/mu_2)/(1/300-1/T_K)
mu_oil = exp(B_visc*(1/T_K-1/(25+273.15)))*mu_oil_25 # Pa.s

# Corrections to make the data consistent;
# The recovery line must be a straight line at the begining of the experiment.
# I shift the time axis a bit to make it a straight line:
# plot(t_sec, R_oil, "-o")
(a,b) = linreg(t_sec[2:6], R_oil[2:6])
# 0 = a + b*t_shift
t_shift = -a/b
t_sec_cor = [0; t_sec[2:end]-t_shift]
# plot(t_sec_cor, R_oil, "-o")

# define the objective functions
# (xt_shock_ww, sw_shock_ww, xt_prf_ww, sw_prf_ww, t_ww, p_inj_ww, R_oil_ww)=
# frac_flow_wf(muw=mu_water, muo=mu_oil, ut=inj_vel, phi=poros_res,
#   k=perm_res, swc=swc_ww, sor=sor_ww, kro0=kro0_ww, no=no_ww, krw0=krw0_ww,
#   nw=nw_ww, sw0=swc_ww, sw_inj=1.0, L=L_res, pv_inj=inj_pv)

# define the objective functions
# param = [sor, swc, kro0, krw0, no, nw]
function rec_fact(param, muw, muo, ut, phi, k, sw0, sw_inj, L, pv_inj, t_exp, R_exp)
  (sor, swc, kro0, krw0, no, nw) = param
  (xt_shock, sw_shock, xt_prf, sw_prf, t, p_inj, R)=
  frac_flow_wf(muw=muw, muo=muo, ut=ut, phi=phi,
    k=k, swc=swc, sor=sor, kro0=kro0, no=no, krw0=krw0,
    nw=nw, sw0=sw0, sw_inj=sw_inj, L=L, pv_inj=pv_inj)
  R_int = Spline1D(t, R, k=1)
  return R_int(t_exp)
end

model = (t_exp, param) -> rec_fact(param, mu_water, mu_oil, u_inj, poros, perm_ave, swi, 1.0, L_core, pv_inj, t_sec_cor, R_oil)

function error_calc(param, muw, muo, ut, phi, k, sw0, sw_inj, L, pv_inj, t_exp, R_exp)
  (sor, swc, kro0, krw0, no, nw) = param
  (xt_shock, sw_shock, xt_prf, sw_prf, t, p_inj, R)=
  frac_flow_wf(muw=muw, muo=muo, ut=ut, phi=phi,
    k=k, swc=swc, sor=sor, kro0=kro0, no=no, krw0=krw0,
    nw=nw, sw0=sw0, sw_inj=sw_inj, L=L, pv_inj=pv_inj)
  R_int = Spline1D(t, R, k=1)
  error_vals = R_int(t_exp) - R_exp
  return error_vals
end

function error_calculation(param, param_ind, muw, muo, ut, phi, k, sw0, sw_inj, L, pv_inj,
  t_exp, R_exp; grad_calc=true)
  eps1 = 1e-8
  error_vals = error_calc(param, muw, muo, ut, phi, k, sw0, sw_inj, L, pv_inj, t_exp, R_exp)
  if grad_calc
    jac  = zeros(length(error_vals), length(param_ind))
  else
    jac = zeros(0)
  end
  if grad_calc
    for j in eachindex(param_ind)
      param2 = copy(param)
      param2[param_ind[j]]+=eps1
      error_val2 = error_calc(param2, muw, muo, ut, phi, k, sw0, sw_inj, L, pv_inj, t_exp, R_exp)
      jac[:, j] = (error_val2 - error_vals) / eps1
    end
  end
  return error_vals, jac
end

function objective_function(param::Vector{Float64}, grad::Vector{Float64}, param_all,
  param_ind, muw, muo, ut, phi, k, sw0, sw_inj, L, pv_inj, t_exp, R_exp; w=1.0)
  param2 = copy(param_all)
  param2[param_ind] = param
  grad_calc = false
  if length(grad)>0
    grad_calc = true
  end
  error_val, jac = error_calculation(param2, param_ind, muw, muo, ut, phi, k, sw0,
    sw_inj, L, pv_inj, t_exp, R_exp, grad_calc = grad_calc)
  obj_func_value = sum(abs2, w.*error_val)
  if length(grad)>0
      grad[:] = 2.0*sum(w.*error_val.*jac, 1)
  end
  return obj_func_value
end

function plot_results(param, muw, muo, ut, phi, k, sw0, sw_inj, L, pv_inj, t_exp, R_exp)
  (sor, swc, kro0, krw0, no, nw) = param
  (xt_shock, sw_shock, xt_prf, sw_prf, t, p_inj, R)=
  frac_flow_wf(muw=muw, muo=muo, ut=ut, phi=phi,
    k=k, swc=swc, sor=sor, kro0=kro0, no=no, krw0=krw0,
    nw=nw, sw0=sw0, sw_inj=sw_inj, L=L, pv_inj=pv_inj)
  R_int = Spline1D(t, R, k=1)
  error_vals = R_int(t_exp) - R_exp
  plot(t_exp, R_int(t_exp), t_exp, R_exp, "o")
end

plot_quick = param -> plot_results(param, mu_water, mu_oil, u_inj, poros, perm_ave, swi, 1.0, L_core, pv_inj, t_sec_cor, R_oil)
# Optimization problem:

# parameter values
# [sor, swc, kro0, krw0, no, nw]
param_all = [so_final, 0.25, 0.7, 0.3, 2, 2.4]
param_ind = [1, 2,3,4,5,6]
w = ones(length(R_oil))
w[end]=5
w[end-1]=15
w[7] = 3
obj_fun = (param, grad) -> objective_function(param::Vector{Float64}, grad::Vector{Float64}, param_all,
  param_ind, mu_water, mu_oil, u_inj, poros, perm_ave, swi, 1.0, L_core, pv_inj, t_sec_cor, R_oil, w=w)

# x_init = [0.2, 0.8, 0.4, 2, 2]
# x_lb = [0.1, 0.1, 0.1, 0.5, 0.5]
# x_ub = [0.5, 1.0, 1.0, 5.0, 5.0]

x_init = copy(param_all)
x_lb = [0.1, 0.1, 0.05, 0.05, 1.0, 1.0]
x_ub = [0.5, 0.5, 1.0, 1.0, 4.0, 4.0]

fit = curve_fit(model, t_sec_cor, R_oil, w, x_init, lower = x_lb, upper = x_ub)
#
# # algorithms
# # :LD_MMA
# # :LN_COBYLA
# # :LD_LBFGS
# # :GN_DIRECT
# # :GN_DIRECT_L
# opt_alg=:LN_COBYLA
#
# opt1 = Opt(opt_alg, length(x_init)) # choose the algorithm
# lower_bounds!(opt1, x_lb)
# upper_bounds!(opt1, x_ub)
# ftol_rel!(opt1, 1e-10)
# ftol_abs!(opt1, 1e-10)
#
# min_objective!(opt1, obj_fun)
# (fObjOpt, paramOpt, flag) = optimize(opt1, x_init)
#
# param_plot = copy(param_all)
# param_plot[param_ind] = paramOpt
plot_quick(fit.param)


module SfCoriolis

using Unitful, UnitfulAngles, Documenter
import LinearAlgebra: cross

import ..SfGeometry: radius, unit_x, unit_y, unit_z
import ..SfUnits: AngularVelocity, AngularMomentum, Angle, to_angle

export coriolis_acceleration, coriolis_force, tangential_velocity, centrifugal_acceleration, centrifugal_force,
	angular_momentum, angular_velocity, radius

DocMeta.setdocmeta!(SfCoriolis, :DocTestSetup, :(using Unitful, UnitfulAstro, ..SfCoriolis, ..SfGeometry); recursive=true)

# this can't be a jldoctest example right now, as there's some regex/string comparison error that breaks it.
"""
    coriolis_acceleration(Ω::Vector, v::Vector)

The apparent acceleration due to coriolis effects for an object moving with velolcity
vector `v` in a rotating reference frame with angular velocity vector `Ω`.

# Example

```julia-repl
julia> coriolis_acceleration(0.001u"rad/s" * unit_y, 400u"m/s" * unit_x)
3-element Array{Quantity{Float64,𝐋  𝐓 ^-2,Unitful.FreeUnits{(m, s^-2),𝐋  𝐓 ^-2,nothing}},1}:
  0.0 m s^-2
  0.0 m s^-2
 -0.8 m s^-2
```
"""
coriolis_acceleration(Ω::Vector, v::Vector) = upreferred.(2cross(Ω, v))

"""
    coriolis_acceleration(Ω::Vector, v::Vector)

The strength of the coriolis force apparently acting on an object with mass `m` moving with velolcity
vector `v` in a rotating reference frame with angular velocity vector `Ω`.
"""
coriolis_force(m::Unitful.Mass, Ω::Vector, v::Vector) = m * coriolis_acceleration(Ω, v)

"""
    tangential_velocity(ω::AngularVelocity, r::Unitful.Length)

Calculate the tangential velocity for an object rotating about a point at a distance `r` with angular velocity `ω`.
"""
tangential_velocity(ω::AngularVelocity, r::Unitful.Length) = upreferred(ω * r)

"""
    tangential_velocity(ω::AngularVelocity, r::Unitful.Length, λ::Angle)

Calculate the tangential velocity of an object at latitude `λ` on a rotating sphere of radius `r` with angular velocity `ω`.
"""
tangential_velocity(ω::AngularVelocity, r::Unitful.Length, λ::Angle) = upreferred(ω * r * cos(λ))

"""
    centrifugal_acceleration(ω::AngularVelocity, r::Unitful.Length)

Calculate the centrifugal acceleration for an object rotating about a point at a distance `r` with angular velocity `ω`.
"""
centrifugal_acceleration(ω::AngularVelocity, r::Unitful.Length) = upreferred(ω^2 * r)

"""
    angular_velocity(ca::Unitful.Acceleration, r::Unitful.Length)

Calculate the angular velocity required to develop a centrifugal acceleration of `ca` at a radius of `r`.
"""
angular_velocity(ca::Unitful.Acceleration, r::Unitful.Length) = sqrt(ca / r) |> u"rad/s"

"""
    radius(ω::AngularVelocity, ca::Unitful.Acceleration)

Calculate the radius of rotation of a body with angular velocity `ω` to develop a centrifugal acceleration of `ca`.
"""
radius(ω::AngularVelocity, ca::Unitful.Acceleration) = upreferred(ca / ω^2)

"""
    centrifugal_acceleration(ω::AngularVelocity, r::Unitful.Length, λ::Angle)

Calculate the radial part of the centrifugal acceleration of object at latitude `λ` on a rotating sphere of radius `r` with angular velocity `ω`.
"""
centrifugal_acceleration(ω::AngularVelocity, r::Unitful.Length, λ::Angle) = upreferred(ω^2 * r * cos(λ))

"""
    centrifugal_force(m::Unitful.Mass, ω::AngularVelocity, r::Unitful.Length)

Calculate the centrifugal force acting on an object of mass `m` rotating about a point at a distance `r` with angular velocity `ω`.
"""
centrifugal_force(m::Unitful.Mass, ω::AngularVelocity, r::Unitful.Length) = upreferred(m * centrifugal_acceleration(ω, r))

"""
    angular_velocity(m::Unitful.Mass, cf::Unitful.Force, r::Unitful.Length)

Calculate the angular velocity required to develop a centrifugal force of `cf` on a body of mass `m` at a radius of `r`.
"""
angular_velocity(m::Unitful.Mass, cf::Unitful.Force, r::Unitful.Length) = angular_velocity(cf / m, r)

"""
    centrifugal_force(m::Unitful.Mass, ω::AngularVelocity, r::Unitful.Length, λ::Angle)

Calculate the radial part of the centrifugal force acting of object of mass `m` at latitude `λ` on a rotating sphere of radius `r` with angular velocity `ω`.
"""
centrifugal_force(m::Unitful.Mass, ω::AngularVelocity, r::Unitful.Length, λ::Angle) = upreferred(m * centrifugal_acceleration(ω, r, λ))

"""
    angular_momentum(m::Unitful.Mass, ω::AngularVelocity, r::Unitful.Length)

Calculate the angular momentum of an object of mass `m` rotating about a point at a distance `r` with angular velocity `ω`.
"""
angular_momentum(m::Unitful.Mass, ω::AngularVelocity, r::Unitful.Length) = upreferred(m * r^2 * ω)

"""
	angular_velocity(m::Unitful.Mass, L::AngularMomentum, r::Unitful.Length)

Calculate the angular velocity of an object with mass `m` and angular momentum `L` rotating about a point at a distance of `r`.
"""
angular_velocity(m::Unitful.Mass, L::AngularMomentum, r::Unitful.Length) = L / (m * r^2) |> u"rad/s"

end


using MetaGraphs: AbstractMetaGraph

## Linear algebra API for AbstractMetaGraph

function degrees(mg::AbstractMetaGraph, T::DataType=eltype(mg); dir::Symbol=:out)
    degrees(adjacency_matrix(mg.graph, T; dir=dir), T; dir=dir)
end

function degree_matrix(mg::AbstractMetaGraph, T::DataType=eltype(mg); dir::Symbol=:out)
    degree_matrix(adjacency_matrix(mg.graph, T; dir=dir), T; dir=dir)
end

function inv_sqrt_degree_matrix(mg::AbstractMetaGraph, T::DataType=eltype(mg); dir::Symbol=:out)
    inv_sqrt_degree_matrix(adjacency_matrix(mg.graph, T; dir=dir), T; dir=dir)
end

function laplacian_matrix(mg::AbstractMetaGraph, T::DataType=eltype(mg); dir::Symbol=:out)
    laplacian_matrix(adjacency_matrix(mg.graph, T; dir=dir), T; dir=dir)
end

function normalized_laplacian(mg::AbstractMetaGraph, T::DataType=eltype(mg); selfloop::Bool=false)
    adj = adjacency_matrix(mg.graph, T)
    selfloop && (adj += I)
    normalized_laplacian(adj, T)
end

## Convolution layers accepting AbstractMetaGraph

GCNConv(g::AbstractMetaGraph, ch::Pair{<:Integer,<:Integer}, σ=identity; kwargs...) =
    GCNConv(g.graph, ch, σ; kwargs...)

ChebConv(g::AbstractMetaGraph, ch::Pair{<:Integer,<:Integer}, k::Integer; kwargs...) =
    ChebConv(g.graph, ch, k; kwargs...)

GraphConv(g::AbstractMetaGraph, ch::Pair{<:Integer,<:Integer}, aggr=:add; kwargs...) =
    GraphConv(g.graph, ch, aggr; kwargs...)

GATConv(g::AbstractMetaGraph, ch::Pair{<:Integer,<:Integer}; kwargs...) =
    GATConv(g.graph, ch; kwargs...)

GatedGraphConv(g::AbstractMetaGraph, out_ch::Integer, num_layers::Integer; kwargs...) =
    GatedGraphConv(g.graph, out_ch, num_layers; kwargs...)

EdgeConv(g::AbstractMetaGraph, nn; kwargs...) = EdgeConv(g.graph, nn; kwargs...)


using Test
using MultivariateBases
using DynamicPolynomials
@polyvar x y

@testset "Polynomials" begin
    basis = FixedPolynomialBasis([1 - x^2, x^2 + 2])
    @test polynomialtype(basis, Int) == polynomialtype(x, Int)
    @test polynomial(one, basis) == 3
end
@testset "Monomial" begin
    basis = FixedPolynomialBasis([x, x^2])
    @test polynomialtype(basis, Int) == polynomialtype(x, Int)
    @test polynomial(i -> i^2, basis) == 4x^2 + x
end
@testset "Terms" begin
    basis = FixedPolynomialBasis([1, x, x^2])
    @test polynomialtype(basis, Int) == polynomialtype(x, Int)
    @test polynomial(i -> (-1)^i, basis) == -x^2 + x - 1
end
@testset "Linear" begin
    basis = FixedPolynomialBasis([x, y])
    @test polynomialtype(basis, Int) == polynomialtype(x, Int)
    @test polynomial(identity, basis) == x + 2y
end
@testset "One variable" begin
    basis = FixedPolynomialBasis([x])
    @test polynomialtype(basis, Int) == polynomialtype(x, Int)
    @test polynomial(i -> 5, basis) == 5x
    @test typeof(polynomial(i -> 5, basis)) == polynomialtype(basis, Int)
    @test typeof(polynomial(ones(Int, 1, 1), basis, Int)) <: AbstractPolynomial{Int}
    @test typeof(polynomial(ones(Int, 1, 1), basis, Float64)) <: AbstractPolynomial{Float64}
end
@testset "Complex" begin
    basis = FixedPolynomialBasis([(1 + 2im) * x])
    @test 5x^2 == polynomial(ones(Int, 1, 1), basis, Complex{Int})
    @test 5x^2 == polynomial(ones(Int, 1, 1), basis, Int)
    # TODO not inferred on Julia v1.0
    #@test 5x^2 == @inferred polynomial(ones(Int, 1, 1), basis, Complex{Int})
    #@test 5x^2 == @inferred polynomial(ones(Int, 1, 1), basis, Int)
end


"""
    struct MultiDiscreteSpace{N} <:AbstractDiscreteSpace
        counts::Array{Int, N}
    end

The element in `MultiDiscreteSpace{N}` is a multi-dimension array.
The number of each dimension is specified by `counts`.
"""
struct MultiDiscreteSpace{N} <:AbstractDiscreteSpace
    counts::Array{Int, N}
end

size(s::MultiDiscreteSpace) = Tuple(s.counts)
eltype(s::MultiDiscreteSpace) = typeof(s.counts)
sample(s::MultiDiscreteSpace) = map(x -> rand(1:x), s.counts)
in(xs::Array{Int}, s::MultiDiscreteSpace) = all(map((e, x) -> 1 ≤ x < e + 1 , s.counts, xs))
in(xs::NTuple{N, Int}, s::MultiDiscreteSpace{N}) where N = all(map((e, x) -> 1 ≤ x < e + 1 , s.counts, xs))


using Pkg

Pkg.activate(homedir() * "/Documents/Coding/gpublockbanded")
using BandedMatrices, BlockBandedMatrices, SharedArrays, ArrayLayouts, BlockArrays, FillArrays
import BlockBandedMatrices: _BandedBlockBandedMatrix, BandedBlockBandedSizes, BlockSizes, blockcolsupport
import BandedMatrices: AbstractBandedMatrix, bandwidths, BandedStyle
import Base: getindex, size

struct FiniteDifference{T} <: AbstractBandedMatrix{T}
    n::Int
end

FiniteDifference(n) = FiniteDifference{Float64}(n)

getindex(F::FiniteDifference{T}, k::Int, j::Int) where T =
    if k == j
        -2*one(T)*F.n^2
    elseif abs(k-j) == 1
        one(T)*F.n^2
    else
        zero(T)
    end

bandwidths(F::FiniteDifference) = (1,1)
size(F::FiniteDifference) = (F.n,F.n)

Base.BroadcastStyle(::Type{<:Eye}) = BandedStyle()

function blockcolbuild!(ret, A, JR)
    for J in JR, K in blockcolsupport(A,J)
        view(ret,K,J) .= view(A,K,J)
    end
end


input = read(joinpath(@__DIR__, "input"), String)
coords, instructions = split(input, "\n\n")

paper = Set()
for coord in split(coords, "\n")
    x, y = parse.(Int, match(r"(\d+),(\d+)", coord).captures)
    push!(paper, (x, y))
end

folds = []
for i in split(instructions, "\n", keepempty=false)
    dim, n = match(r".+([xy])=(\d+)$", i).captures
    push!(folds, (dim, parse(Int, n)))
end

function translatex(coord, n)
    x, y = coord
    x > n && return (n - (x - n), y)
    return (x, y)
end

function translatey(coord, n)
    x, y = coord
    y > n && return (x, n - (y - n))
    return (x, y)
end

function fold(paper, folds)
    paper = copy(paper)
    for (dim, n) in folds
        if dim == "x"
            paper = map(idx -> translatex(idx, n), collect(paper))
            paper = filter(idx -> first(idx) <= n, paper)
            paper = unique(paper)
        else
            paper = map(idx -> translatey(idx, n), collect(paper))
            paper = filter(idx -> last(idx) <= n, paper)
            paper = unique(paper)
        end
    end
    return paper
end

firstfold = first(folds)
p1 = length(fold(paper, [firstfold]))

paper = fold(paper, folds)
x = maximum(map(idx -> first(idx), collect(paper)))
y = maximum(map(idx -> last(idx), collect(paper)))

println("-----------------------------------------------------------------------")
println("transparent origami -- part one :: $p1")
println("transparent origami -- part two ::")
println()
for j in 0:y
    print("    ")
    for i in 0:x
        (i, j) in paper ? print("█") : print(" ")
    end
    println()
end
println()
println("-----------------------------------------------------------------------")

@assert(p1 == 759)


export PatchedSystem, PatchedSystemCache

"""
    PatchedSystem(F::AbstractSystem, patch, v::PVector)

Augment the system `F` with the given patch `v`. This results in the system `[F(x,t); v ⋅ x - 1]`
"""
struct PatchedSystem{F<:AbstractSystem, PS<:AbstractAffinePatchState} <: AbstractSystem
    system::F
    patch::PS
end

function PatchedSystem(F::AbstractSystem, p::AbstractAffinePatch, x)
    PatchedSystem(F, state(p, x))
end

function Base.size(F::PatchedSystem)
    m, n = size(F.system)
    (m + nequations(F.patch), n)
end

struct PatchedSystemCache{FC, T} <: AbstractSystemCache
    cache::FC
    A::Matrix{T} # intermediate storage of the jacobian
    b::Vector{T} # intermediate storage for the evaluation
end

function cache(ph::PatchedSystem, x)
    F = ph.system
    c = cache(F, x)
    b = evaluate(F, x, c)
    A = similar(b, size(F))
    PatchedSystemCache(c, A, b)
end

function evaluate!(u, F::PatchedSystem, x, c::PatchedSystemCache)
    M, N = size(F.system)

    # F(x)
    evaluate!(c.b, F.system, x, c.cache)
    @inbounds for i=1:M
        u[i] = c.b[i]
    end
    evaluate!(u, F.patch, x)
    u
end

function jacobian!(U, F::PatchedSystem, x::PVector, c::PatchedSystemCache)
    M, N = size(F.system)
    # J_F(x)
    jacobian!(c.A, F.system, x, c.cache)
    @inbounds for j=1:N, i=1:M
        U[i, j] = c.A[i, j]
    end
    jacobian!(U, F.patch, x)
    U
end

function evaluate_and_jacobian!(u, U, F::PatchedSystem, x::PVector, c::PatchedSystemCache)
    M, N = size(F.system)
    evaluate_and_jacobian!(c.b, c.A, F.system, x, c.cache)

    @inbounds for j=1:N, i=1:M
        U[i, j] = c.A[i, j]
    end
    jacobian!(U, F.patch, x)

    @inbounds for i=1:M
        u[i] = c.b[i]
    end
    evaluate!(u, F.patch, x)

    nothing
end

function evaluate(F::PatchedSystem, x, c::PatchedSystemCache)
    M, N = size(F)
    evaluate!(similar(c.b, M), F, x, c)
end

function jacobian(F::PatchedSystem, x, c::PatchedSystemCache)
    M, N = size(F)
    jacobian!(similar(c.A, M, N), F, x, c)
end


using GraphOrdering, Test

G = Dict(
    1 => (3, 8, 9),
    2 => (3, 8, 7),
    3 => (1, 2),
    4 => (8, 9),
    5 => (7, 8),
    6 => (2, 7),
    7 => (5, 2, 6),
    8 => (1, 2, 4, 5),
    9 => (1, 4))

@test GraphOrdering.bandwidth(G) == 17


using Rotations, StaticArrays, Test
using Unitful

@testset "2d Rotations" begin

    #################################
    # Traits, types, and construction
    #################################

    @testset "Core" begin
        r = one(RotMatrix{2,Float32})
        @test RotMatrix((1,0,0,1)) == RotMatrix(@SMatrix [1 0; 0 1])
        @test_throws DimensionMismatch RotMatrix((1,0,0,1,0))
    end

    @testset "Unitful" begin
        # Make sure rotations created from unitful angles
        # don't extraneously contain those units (see issue #55)
        @test eltype(Angle2d(10u"°")) <: Real
        @test eltype(Angle2d(20u"rad")) <: Real
    end

    ###############################
    # Check fixed relationships
    ###############################

    @testset "Identity rotation checks" begin
        I = one(SMatrix{2,2,Float64})
        I32 = one(SMatrix{2,2,Float32})
        for R in [RotMatrix{2}, Angle2d]
            @test @inferred(size(R)) == (2,2)
            @test @inferred(size(R{Float32})) == (2,2)
            @test one(R)::R == I
            @test one(R{Float32})::R{Float32} == I32
        end
    end

    ###############################
    # Check zero function
    ###############################

    @testset "zero checks" begin
        for R in (RotMatrix{2}, Angle2d)
            # zero
            @test zero(R) == zero(R{Float64}) == zero(one(R))
            @test zero(R) isa SMatrix
            @test zero(R{Float64}) isa SMatrix
            @test zero(one(R)) isa SMatrix
            # zeros
            @test zeros(R)[1] == zeros(R,3)[1] == zeros(R,3,3)[1] == zeros(R,(3,3,3))[1] == zero(R)
            @test zeros(R) isa Array{<:SMatrix,0}
            @test zeros(R,3) isa Array{<:SMatrix,1}
            @test zeros(R,3,3) isa Array{<:SMatrix,2}
            @test zeros(R,(3,3,3)) isa Array{<:SMatrix,3}
        end
    end

    ################################
    # check on the inverse function
    ################################

    @testset "Testing inv()" begin
        repeats = 100
        for R in [RotMatrix{2,Float64}, Angle2d{Float64}]
            I = one(R)
            Random.seed!(0)
            for i = 1:repeats
                r = rand(R)
                @test isrotation(r)
                @test inv(r) == adjoint(r)
                @test inv(r) == transpose(r)
                @test inv(r)*r ≈ I
                @test r*inv(r) ≈ I
            end
        end
    end

    ################################
    # check on the norm functions
    ################################

    @testset "Testing norm() and normalize()" begin
        repeats = 100
        for R in [RotMatrix{2,Float64}, Angle2d{Float64}]
            I = one(R)
            Random.seed!(0)
            for i = 1:repeats
                r = rand(R)
                @test norm(r) ≈ norm(Matrix(r))
                if VERSION ≥ v"1.5"
                    @test normalize(r) ≈ normalize(Matrix(r))
                    @test normalize(r) isa SMatrix
                end
            end
        end
    end

    #########################################################################
    # Rotate some stuff
    #########################################################################

    # a random rotation of a random point
    @testset "Rotate Points" begin
        repeats = 100
        for R in [RotMatrix{2}, Angle2d]
            Random.seed!(0)
            for i = 1:repeats
                r = rand(R)
                m = SMatrix(r)
                v = randn(SVector{2})

                @test r*v ≈ m*v
            end

            # Test Base.Vector also
            r = rand(R)
            m = SMatrix(r)
            v = randn(2)

            @test r*v ≈ m*v
        end
    end

    # a random rotation of a random unitful point
    @testset "Rotate Unitful Points" begin
        repeats = 100
        for R in [RotMatrix{2}, Angle2d]
            Random.seed!(0)
            for i = 1:repeats
                r = rand(R)
                m = SMatrix(r)
                v = randn(SVector{2}) * u"m"

                @test r*v ≈ m*v
                @test eltype(r*v) <: Unitful.Length
                @test eltype(m*v) <: Unitful.Length
            end

            # Test Base.Vector also
            r = rand(R)
            m = SMatrix(r)
            v = randn(2) * u"m"

            @test r*v ≈ m*v
            @test eltype(r*v) <: Unitful.Length
            @test eltype(m*v) <: Unitful.Length
        end
    end

    # compose two random rotations
    @testset "Compose rotations" begin
        repeats = 100
        for R in [RotMatrix{2}, Angle2d]
            Random.seed!(0)
            for i = 1:repeats
                r1 = rand(R)
                m1 = SMatrix(r1)

                r2 = rand(R)
                m2 = SMatrix(r2)

                @test r1*r2 ≈ m1*m2
                θ1, θ2 = atan(r1[2,1],r1[1,1]), atan(r2[2,1],r2[1,1])
                @test r1*r2 ≈ RotMatrix(θ1+θ2)
                @test r1/r2 ≈ RotMatrix(θ1-θ2)
                @test r1\r2 ≈ RotMatrix(θ2-θ1)
            end
        end
    end

    #########################################################################
    # Test conversions between rotation types
    #########################################################################
    @testset "Convert rotations" begin
        repeats = 100
        for R in [RotMatrix{2}, Angle2d]
            Random.seed!(0)
            for _ in 1:repeats
                r1 = rand(R)
                @test R(r1) == r1
            end
        end
    end

    @testset "Types and products" begin
        for (R,T) in ((RotMatrix(pi/4), Float64),
                      (RotMatrix(Float32(pi/4)), Float32),
                      (RotMatrix{2,Float32}(pi/4), Float32))
            @test eltype(R) == T
            @test size(R) == (2,2)
            @test R == T[cos(pi/4) -sin(pi/4); sin(pi/4) cos(pi/4)]
            @test R * R ≈ T[0 -1; 1 0]
        end
    end

    @testset "RotMatrix{2} vs Angle2d" begin
        repeats = 100
        Random.seed!(0)
        for _ in 1:repeats
            theta = randn()
            r1 = RotMatrix{2}(theta)
            r2 = Angle2d(theta)
            v = randn(2)
            @test r1 ≈ r2
            @test r1 * v ≈ r2 * v
            @test RotMatrix{2}(r2) ≈ r1
            @test Angle2d(r1) ≈ r2
        end
    end

    @testset "Angle2d" begin
        repeats = 100
        Random.seed!(0)
        for _ in 1:repeats
            r = rand(Angle2d)
            @test r^(-1) ≈ inv(r)
            @test r^(-1.0) ≈ inv(r)
            @test r^1 ≈ r
            @test r^1.0 ≈ r
            @test r^2 ≈ r*r
            @test r^2.0 ≈ r*r

            x = randn()
            y = randn()
            @test r^(x+y) ≈ r^x * r^y
        end
    end

    @testset "angle" begin
        repeats = 100
        Random.seed!(0)
        for _ in 1:repeats
            theta = 2pi*rand() - pi
            a2d = Angle2d(theta)
            r2 = RotMatrix{2}(theta)
            @test rotation_angle(r2) ≈ theta
            @test rotation_angle(a2d) == theta
            @test rotation_angle(a2d) == Rotations.params(a2d)[1]
        end
    end
end

nothing


# ellipse_im.jl
# Copyright 2019-03-05, Jeff Fessler, University of Michigan

#using MIRT: MIRT_image_geom, jim
using Plots

"""
`phantom = ellipse_im(ig, params;
rot=0, oversample=1, hu_scale=1, replace=false, return_params=false)`

Generate ellipse phantom image from parameters:

`[x_center y_center x_radius y_radius angle_degrees amplitude]`

in
* `ig`		from `image_geom()`
* `params`	[ne 6]	ellipse parameters

option
* `rot`			rotate ellipses by this amount [degrees]
* `oversample`	oversampling factor, for grayscale boundaries
* `hu_scale`	use 1000 to scale shepp-logan to HU
* `replace`		replace ellipse values if true, else add
* `how`			:fast is the only option
* `return_params` if true return (phantom, params)

out
* `phantom`		[nx ny]	image
* `params`		[ne 6] parameters (only if return_params=true)

note: `op ellipse` in aspire with `nsub=3` is `oversample=4 = 2^(3-1)` here

"""
function ellipse_im(ig::MIRT_image_geom,
	params::AbstractArray{<:Real,2};
	rot::Real=0,
	oversample::Integer=1,
	hu_scale::Real=1,
	replace::Bool=false,
	how::Symbol=:fast, # todo
	return_params::Bool=false)

	params[:,6] .*= hu_scale

#if how == :fast && oversample == 1
#	warning("ignoring :fast option for oversample=1")
#	how = :slow
#end

	if oversample > 1
		ig = ig.over(oversample)
	end
	args = (ig.nx, ig.ny, params, ig.dx, ig.dy, ig.offset_x, ig.offset_y,
	rot, oversample, replace)

	if how == :fast
		phantom = ellipse_im_fast(args...)
#	elseif how == :slow
#		phantom = ellipse_im_slow(args...)
	else
		error("bad how")
	end

	if oversample > 1
		phantom = downsample2(phantom, oversample)
	end

	if return_params
		return (phantom, params)
	end
	return phantom
end

#%
#% ellipse_im_slow()
#% brute force fine grid - can use lots of memory
#%
#function phantom = ellipse_im_slow(nx, ny, params, dx, dy, ...
	#offset_x, offset_y, rot, over, replace)
#
#if size(params,2) ~= 6
	#fail('bad ellipse parameter vector size')
#end

#% optional rotation of ellipse parameters
#if rot ~= 0
	#th = deg2rad(rot);
	#cx = params[:,1);
	#cy = params[:,2);
	#params[:,1) = cx * cos(th) + cy * sin(th);
	#params[:,2) = -cx * sin(th) + cy * cos(th);
	#params[:,5) = params[:,5) + rot;
	#clear cx cy th
#end
#
#wx = (nx*over-1)/2 + offset_x * over;
#wy = (ny*over-1)/2 + offset_y * over;
#xx = ((0:nx*over-1) - wx) / over * dx;
#yy = ((0:ny*over-1) - wy) / over * dy;
#[xx yy] = ndgrid(xx, yy); % fine grid, equally spaced
#
#phantom = zeros(nx*over, ny*over, 'single'); % fine array
#
#ticker reset
#ne = nrow(params);
#for ie = 1:ne
	#ticker(mfilename, ie, ne)
#
	#ell = params[ie, :);
	#cx = ell(1);	rx = ell(3);
	#cy = ell(2);	ry = ell(4);
	#theta = deg2rad(ell(5));
	#[xr yr] = rot2(xx-cx, yy-cy, theta);
	#tmp = (xr / rx).^2 + (yr / ry).^2 <= 1;
#
	#if replace
		#phantom(tmp > 0) = ell(6);
	#else
		#phantom = phantom + ell(6) * tmp;
	#end
#end
#
#phantom = downsample2(phantom, over);

"""
`phantom = ellipse_im_fast()`
"""
function ellipse_im_fast(nx, ny, params_in, dx, dy,
	offset_x, offset_y, rot, over, replace)

	params = copy(params_in)
	if size(params,2) != 6
		error("bad ellipse parameter vector size")
	end

	# optional rotation
	if rot != 0
		th = rot * 180/pi
		cx = params[:,1]
		cy = params[:,2]
		params[:,1] = cx * cos(th) + cy * sin(th)
		params[:,2] = -cx * sin(th) + cy * cos(th)
		params[:,5] .+= rot
	end

	phantom = zeros(Float32, nx, ny)

	wx = (nx-1)/2 + offset_x
	wy = (ny-1)/2 + offset_y
	x1 = ((0:nx-1) .- wx) * dx
	y1 = ((0:ny-1) .- wy) * dy
	(xx, yy) = ndgrid(x1, y1)

	hx = abs(dx) / 2
	hy = abs(dy) / 2

	for ie = 1:size(params,1)

		ell = params[ie, :]
		cx = ell[1];	rx = ell[3]
		cy = ell[2];	ry = ell[4]
		theta = ell[5] * pi/180

		xs = xx .- cx # shift per ellipse center
		ys = yy .- cy

		# coordinates of "outer" corner of each pixel, relative to ellipse center
		xo = xs + sign.(xs) * hx
		yo = ys + sign.(ys) * hy

		# voxels that are entirely inside the ellipse:
		(xr, yr) = rot2(xo, yo, theta)
		is_inside = (xr / rx).^2 + (yr / ry).^2 .<= 1

		if replace
			phantom(is_inside) .= ell[6]
		else
			phantom += ell[6] * Float32.(is_inside)
		end

	end # ie loop

	return phantom
end # ellipse_im_fast()

"""
`phantom = ellipse_im(nx, dx, params; kwarg...)`

square image of size `nx` by `nx`,
specifying pixel size `dx` and ellipse `params`
"""
function ellipse_im(nx::Integer, dx::Real, params; kwarg...)
	ig = image_geom(nx=nx, dx=1)
	return ellipse_im(ig, params; kwarg...)
end

"""
`phantom = ellipse_im(nx::Integer, params; kwarg...)`

square image of size `nx` by `nx` with
pixel size `dx=1` and ellipse `params`
"""
function ellipse_im(nx::Integer, params; kwarg...)
	return ellipse_im(nx, 1., params; kwarg...)
end

"""
`phantom = ellipse_im(nx::Integer; ny::Integer=nx, dx::Real=1, kwarg...)`

image of size `nx` by `ny` (default `nx`) with specified `dx` (default 1),
defaults to `:shepplogan_emis`
"""
function ellipse_im(nx::Integer; ny::Integer=nx, dx::Real=1, kwarg...)
	ig = image_geom(nx=nx, ny=ny, dx=dx)
	return ellipse_im(ig, :shepplogan_emis; kwarg...)
end

"""
`phantom = ellipse_im(nx::Integer, ny::Integer; kwarg...)`

`:shepplogan_emis` of size `nx` by `ny`
"""
function ellipse_im(nx::Integer, ny::Integer; kwarg...)
	return ellipse_im(nx, ny=ny, dx=1.; kwarg...)
end

"""
`phantom = ellipse_im(ig, code, kwarg...)`

`code = :shepplogan | :shepplogan_emis | :shepplogan_brainweb`
"""
function ellipse_im(ig::MIRT_image_geom, params::Symbol; kwarg...)
	fov = ig.fov
	if params == :shepplogan
		params = shepp_logan_parameters(fov, fov, case=:kak)
	elseif params == :shepplogan_emis
		params = shepp_logan_parameters(fov, fov, case=:emis)
	elseif params == :shepplogan_brainweb
		params = shepp_logan_parameters(fov, fov, case=:brainweb)
	else
		error("bad phantom symbol")
	end
	return ellipse_im(ig, params; kwarg...)
end

"""
`phantom = ellipse_im(ig; kwarg...)`

`:shepplogan` (default) for given image geometry `ig`
"""
function ellipse_im(ig::MIRT_image_geom; kwarg...)
	return ellipse_im(ig, :shepplogan; kwarg...)
end

"""
`(xx,yy) = ndgrid(x,y)`
"""
function ndgrid(x::AbstractVector{<:Number},
				y::AbstractVector{<:Number})
	return (repeat(x, 1, length(y)), repeat(y', length(x), 1))
end

"""
`(xr,yr) = rot2(x, y, theta)`
2D rotation
"""
function rot2(x, y, theta)
	xr = cos(theta) * x + sin(theta) * y
	yr = -sin(theta) * x + cos(theta) * y
	return (xr, yr)
end

"""
`params = shepp_logan_parameters(xfov, yfov)`

parameters from Kak and Slaney text, p. 255

the first four columns are unitless "fractions of field of view"
"""
function shepp_logan_parameters(xfov::Real, yfov::Real; case::Symbol=:kak)
	params = [
	0		0		0.92	0.69	90	2;
	0		-0.0184	0.874	0.6624	90	-0.98;
	0.22	0		0.31	0.11	72	-0.02;
	-0.22	0		0.41	0.16	108	-0.02;
	0		0.35	0.25	0.21	90	0.01;
	0		0.1		0.046	0.046	0	0.01;
	0		-0.1	0.046	0.046	0	0.01;
	-0.08	-0.605	0.046	0.023	0	0.01;
	0		-0.605	0.023	0.023	0	0.01;
	0.06	-0.605	0.046	0.023	90	0.01];

	params[:,[1,3]] .*= xfov/2
	params[:,[2,4]] .*= yfov/2

	if case == :emis
		params[:,6] = [1, 1, -2, 2, 3, 4, 5, 6, 1, 1]
	elseif case == :brainweb
		params[:,6] = [1, 0, 2, 3, 4, 5, 6, 7, 8, 9] # brainweb uses index 1-10
	elseif case != :kak
		error("bad phantom case")
	end

	return params
end

#%
#% ellipse_im_profile()
#%
#function ellipse_im_profile
#ig = image_geom('nx', 2^9, 'ny', 2^9+2', 'fov', 250);
#profile on
#x0 = ellipse_im(ig, [], 'oversample', 3, 'type', 'fast');
#profile off
#profile report

"""
`ellipse_im()`

show docstring(s)
"""
function ellipse_im()
	@doc ellipse_im
end

"""
`ellipse_im_show()`
"""
function ellipse_im_show()
	ig = image_geom(nx=2^8, ny=2^8+2, fov=250)
#	ig.offset_y = 75.6 / ig.dy

	over = 2^2
	x0 = ellipse_im(ig, oversample=over)
	p1 = jim(ig.x(), ig.y(), x0, title="Shepp Logan", clim=[0.9,1.1])

	x1 = ellipse_im(ig, :shepplogan_emis, oversample=over)
	p4 = jim(x1, title="Shepp Logan Emission")

	x3 = ellipse_im(ig, :shepplogan_brainweb)
	p5 = jim(x3, title="Shepp Logan Brainweb")

	plot(p1,p4,p5)
end

# compare to aspire
function ellipse_im_aspire()
	nx = 2^6;
	ig = image_geom(nx=nx, ny=nx+2, fov=2^7)
	ell = [10, 20, 30, 40, 50, 1]

	ir_test_dir = "/tmp/" # todo
	file = ir_test_dir * "t.fld"
	over = 2^2
	com = sprintf("echo y | op -chat 99 ellipse %s %d %d  %g %g %g %g %g %g %d",
		file, ig.nx, ig.ny, ell ./ [ig.dx, ig.dx, ig.dx, ig.dx, 1, 1],
		log2(over)+1)
	os_run(com) # todo: not done
	asp = fld_read(file)
	jim(asp, title="aspire")

	area_asp = sum(asp[:]) * abs(ig.dx * ig.dy)

	jul = ellipse_im(ig, ell, oversample=over) # 'type', types{ii})
	area_jul = sum(jul[:]) * abs(ig.dx * ig.dy)

	jim(jul, title="jul")
	jim((jul-asp)*over^2, title="difference: (jul-asp)*over^2")

	@show maximum(abs.(jul - mat)) / ell[6] * over^2
	@assert isapprox(jul, mat)
#	max_percent_diff(mat, asp)

	area_real = pi * ell[3] * ell[4] * ell[6]
	@show area_real, area_asp, area_jul

	true
end

# ellipse_im_test()
function ellipse_im_test()
	fov = 100
	shepp_logan_parameters(fov, fov)
	shepp_logan_parameters(fov, fov, case=:kak)
	shepp_logan_parameters(fov, fov, case=:emis)
	shepp_logan_parameters(fov, fov, case=:brainweb)

	# test various ways of calling
	ellipse_im(20)
	ellipse_im(30, :shepplogan_emis, oversample=2)

	ig = image_geom(nx=80, dx=1)
	fov = ig.fov
	params = shepp_logan_parameters(fov, fov, case=:brainweb)
	ellipse_im(ig, params)
	ellipse_im(ig, params, oversample=2)

	if false
		x1 = ellipse_im(100, :shepplogan_emis)
		x2 = ellipse_im(100, :shepplogan_emis, oversample=2)
		plot(jim(x1), jim(x2), jim(x2-x1))
	end

	ellipse_im_show()

#	if has_aspire
#		ellipse_im_aspire() # todo
#	end
	true
end

"""
`ellipse_im(:test)`

`ellipse_im(:show)`

run tests
"""
function ellipse_im(test::Symbol)
	if test == :show
		return ellipse_im_show()
	end
	@assert test == :test
	ellipse_im()
	ellipse_im(:show)
	ellipse_im_test()
end


# Copyright 2019 Tobias Frilling
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

@inline initMethod(settings::Settings = defaultSettings) =
    initMethod(Val(settings.initMethod), settings)

"""
    initMethod(::Val{:uniform_min_max})

Returns an initializer function.

The values are uniformly distributed over ``[\\min(c), \\max(c)]`` of the data
column.
"""
function initMethod(::Val{:uniform_min_max}, settings::Settings = defaultSettings)
    col -> rand(Uniform(minimum(col), maximum(col)), prod(settings.latticeSize))
end

"""
    initMethod(::Val{:uniform_mean_std})

Returns an initializer function.

The values are uniformly distributed between ``[μ-2σ, μ+2σ]`` of the data
column.
"""
function initMethod(::Val{:uniform_mean_std}, settings::Settings = defaultSettings)
    m = mean(col)
    s = std(col)
    col -> rand(Uniform(m - 2s, m + 2s), prod(settings.latticeSize))
end

"""
    initMethod(::Val{:normal_mean_std})

Returns an initializer function.

The values are Normal distributed with the same mean and std as the data column.
"""
function initMethod(::Val{:normal_mean_std}, settings::Settings = defaultSettings)
    col -> rand(Normal(mean(col), std(col)), prod(settings.latticeSize))
end

"""
    initMethod(::Val{:zeros})

Returns an initializer function.

The values are set to 0.
"""
function initMethod(::Val{:zeros}, settings::Settings = defaultSettings)
    col -> zeros(prod(settings.latticeSize))
end


using VortexDistributions

## Basic unit testing for coverage
@test scalar_ansatz(1) == scalar_ansatz(-1)
f = Ansatz()
x = LinRange(0, 10, 100)
y = f.(x)
@test length(y) == length(x)
@test f.(3,4) == f(5)

y, ψ, res = VortexDistributions.gpecore_exact(1,2,100)
@test length(y) == 100


using FMMTrees
using StaticArrays
using AdaptiveCrossApproximation

p = rand(SVector{2,Float64}, 100)
q, tree = AdaptiveCrossApproximation.clustertree(p)

using Plots
gr()

rectangle(ll, ur) = [ll[1],ur[1],ur[1],ll[1],ll[1]], [ll[2],ll[2],ur[2],ur[2],ll[2]]

plot()
level = 1
clrs = [:blue, :red, :green, :black, :cyan, :yellow]
state = [(0,tree[1].num_children)]
for (i,b) in enumerate(tree)
    ll, ur = FMMTrees.boundingbox(q[b.data.begin_idx:b.data.end_idx-1])
    x, y = rectangle(ll, ur)
    level == 5 && plot!(x,y,color=clrs[level])
    level == 4 && println(level, ": ", b.data.begin_idx:b.data.end_idx-1)
    state[end] = (state[end][1] + b.num_children + 1, state[end][2])
    if b.num_children != 0
         level += 1
         push!(state, (0,b.num_children))
     else
         while state[end][1] == state[end][2]
             level -= 1
             pop!(state)
         end
     end
     plot!()
end

x = [r[1] for r in q]
y = [r[2] for r in q]
scatter!(x,y)
plot!(axis=:none,legend=:none)
plot!(title="level 4 ACA clusters")

FMMTrees.print_tree(tree)


@testset "pool" begin
    @testset "GlobalPool" begin
        p = GlobalPool(+)
        n = 10
        chin = 6
        X = rand(Float32, 6, n)
        g = GNNGraph(random_regular_graph(n, 4), ndata=X, graph_type=GRAPH_T)
        u = p(g, X)
        @test u ≈ sum(X, dims=2)

        ng = 3
        g = Flux.batch([GNNGraph(random_regular_graph(n, 4),
                                 ndata=rand(Float32, chin, n),
                                 graph_type=GRAPH_T)
                        for i=1:ng])
        u = p(g, g.ndata.x)
        @test size(u) == (chin, ng)
        @test u[:,[1]] ≈ sum(g.ndata.x[:,1:n], dims=2)
        @test p(g).gdata.u == u

        test_layer(p, g, rtol=1e-5, exclude_grad_fields = [:aggr], outtype=:graph)
    end

    @testset "GlobalAttentionPool" begin
        n = 10
        chin = 6
        chout = 5
        ng = 3

        fgate = Dense(chin, 1)
        ffeat = Dense(chin, chout)
        p = GlobalAttentionPool(fgate, ffeat)
        @test length(Flux.params(p)) == 4

        g = Flux.batch([GNNGraph(random_regular_graph(n, 4),
                                 ndata=rand(Float32, chin, n),
                                 graph_type=GRAPH_T)
                        for i=1:ng])

        test_layer(p, g, rtol=1e-5, outtype=:graph, outsize=(chout, ng))
    end

    @testset "TopKPool" begin
        N = 10
        k, in_channel = 4, 7
        X = rand(in_channel, N)
        for T = [Bool, Float64]
            adj = rand(T, N, N)
            p = TopKPool(adj, k, in_channel)
            @test eltype(p.p) === Float32
            @test size(p.p) == (in_channel,)
            @test eltype(p.Ã) === T
            @test size(p.Ã) == (k, k)
            y = p(X)
            @test size(y) == (in_channel, k)
        end
    end

    @testset "topk_index" begin
        X = [8,7,6,5,4,3,2,1]
        @test topk_index(X, 4) == [1,2,3,4]
        @test topk_index(X', 4) == [1,2,3,4]
    end
end


@testset "Example α_offset" begin

    α0=LinearFilter(Float64[0,1,0],0)
    α1=LinearFilter(Float64[0,1,0],1)
    β=collect(Float64,1:6)
    γ1=directConv(α0,1,β,ZeroPaddingBE,ZeroPaddingBE)
    γ2=directConv(α1,1,β,ZeroPaddingBE,ZeroPaddingBE)

    @test γ1 ≈ [2.0, 3.0, 4.0, 5.0, 6.0, 0.0]
    @test γ2 ≈ [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]

end;

@testset "Adjoint operator" begin

    α=LinearFilter(rand(4),2);
    β=rand(10);

    vβ=rand(length(β))
    d1=dot(directConv(α,-3,vβ,ZeroPaddingBE,ZeroPaddingBE),β)
    d2=dot(directConv(α,+3,β,ZeroPaddingBE,ZeroPaddingBE),vβ)

    @test isapprox(d1,d2)

    d1=dot(directConv(α,-3,vβ,PeriodicBE,PeriodicBE),β)
    d2=dot(directConv(α,+3,β,PeriodicBE,PeriodicBE),vβ)

    @test isapprox(d1,d2)

end;

@testset "Convolution commutativity" begin

    α=rand(4);
    αf=LinearFilter(α,0)
    β=rand(10);
    βf=LinearFilter(β,0)
    v1=zeros(20)
    v2=zeros(20)
    directConv!(αf,-1,
                β,v1,UnitRange(1,20),ZeroPaddingBE,ZeroPaddingBE)
    directConv!(βf,-1,
                α,v2,UnitRange(1,20),ZeroPaddingBE,ZeroPaddingBE)

    @test isapprox(v1,v2)

end;

@testset "Interval split" begin

    α=LinearFilter(rand(4),3)
    β=rand(10);

    γ=directConv(α,2,β,MirrorBE,PeriodicBE) # global computation
    Γ=zeros(length(γ))
    Ω1=UnitRange(1:3)
    Ω2=UnitRange(4:length(γ))
    directConv!(α,2,β,Γ,Ω1,MirrorBE,PeriodicBE) # compute on Ω1
    directConv!(α,2,β,Γ,Ω2,MirrorBE,PeriodicBE) # compute on Ω2

    @test isapprox(γ,Γ)

end;

@testset "2D convolution" begin
    β=rand(5,8)
    β_save=deepcopy(β)

    α_I=LinearFilter(Float64[0,2,0])
    α_J=LinearFilter(Float64[0,0,2,0,0])

    directConv2D!(α_I,1,α_J,-1,β)

    @test isapprox(β,4*β_save)
end;

@testset "2D crossCorrelation" begin
    β=rand(5,8)

    α_I=LinearFilter(Float64[0,2,0])
    α_J=LinearFilter(Float64[0,0,2,0,0])

    γ=directCrossCorrelation2D(α_I,α_J,β)

    @test isapprox(γ,4*β)
end;


struct Akima
    n::Int
    x::Vector{Int}
    y::Vector{Float64}
    b::Vector{Float64}
    c::Vector{Float64}
    d::Vector{Float64}
end

function Akima(x, y)
    n = length(x)
    @assert n == length(y)
    dx = diff(x)
    m = Array{Float64}(undef, n + 3)
    m[3:end-2] = diff(y) ./ dx
    m[2] = 2m[3] - m[4]
    m[1] = 2m[2] - m[3]
    m[end-1] = 2m[end-2] - m[end-3]
    m[end] = 2m[end-1] - m[end-2]

    # if m1 == m2 != m3 == m4, the slope at the breakpoint is not defined.
    # This is the fill value:
    b = 0.5 .* (m[4:end] .+ m[1:end-3])
    # get the denominator of the slope t
    dm = abs.(diff(m))
    f1 = dm[3:n+2]
    f2 = dm[1:n]
    f12 = f1 + f2
    # These are the mask of where the the slope at breakpoint is defined:
    ind = findall(f12 .> 1e-9 * maximum(f12))
    # Set the slope at breakpoint
    b[ind] = (f1[ind] .* m[ind .+ 1] .+
              f2[ind] .* m[ind .+ 2]) ./ f12[ind]
    # calculate the higher order coefficients
    c = (3.0 .* m[3:end-2] .- 2.0 .* b[1:end-1] .- b[2:end]) ./ dx
    d = (b[1:end-1] .+ b[2:end] .- 2.0 .* m[3:end-2]) ./ dx.^2

    Akima(n, x, y, b, c, d)
end

@inline function interpolate(a::Akima, xi)
    idx = searchsortedlast(a.x, trunc(Int, xi))
    idx == 0 && return a.y[1]
    idx == a.n && return a.y[end]

    wj = xi - a.x[idx]
    @evalpoly wj a.y[idx] a.b[idx] a.c[idx] a.d[idx]
end


branches(::Type{AbstractBinomialTree}) = 2

mutable struct BinomialTreeCommon
  x0::Float64
  driftPerStep::Float64
  dt::Float64
  columns::Int
end

get_x0(t::AbstractBinomialTree) = t.common.x0
get_drift_per_step(t::AbstractBinomialTree) = t.common.driftPerStep
get_dt(t::AbstractBinomialTree) = t.common.dt

function BinomialTreeCommon(process::StochasticProcess1D, endTime::Float64, steps::Int)
  x0 = get_x0(process)
  dt = endTime/steps
  driftPerStep = drift(process, 0.0, x0) * dt
  columns = steps + 1
  BinomialTreeCommon(x0, driftPerStep, dt, columns)
end

struct Joshi4 <: BinomialTreeType end
struct LeisenReimer <: BinomialTreeType end
struct Tian <: BinomialTreeType end

mutable struct BinomialTree{T <: BinomialTreeType} <: AbstractBinomialTree
  up::Float64
  down::Float64
  pu::Float64
  pd::Float64
  common::BinomialTreeCommon
  treeType::T
end

const Joshi4BinomialTree = BinomialTree{Joshi4}
const LeisenReimerBinomialTree = BinomialTree{LeisenReimer}
const TianBinomialTree = BinomialTree{Tian}

function Tian(process::StochasticProcess1D, endTime::Float64, steps::Int, ::Float64)
  common = BinomialTreeCommon(process, endTime, steps)
  q = exp(variance(process, 0.0, common.x0, common.dt))
  r = exp(common.driftPerStep) * sqrt(q)

  up = 0.5 * r * q * (q + 1.0 + sqrt(q * q + 2.0 * q - 3.0))
  down = 0.5 * r * q * (q + 1.0 - sqrt(q * q + 2.0 * q - 3.0))

  pu = (r - down) / (up - down)
  pd = 1.0 - pu

  pu <= 1.0 || error("negative probability")
  pu >= 0.0 || error("negative probability")

  return BinomialTree{Tian}(up, down, pu, pd, common, Tian())
end

function LeisenReimer(process::StochasticProcess1D, endTime::Float64, steps::Int, strike::Float64)
  common = BinomialTreeCommon(process, endTime, steps)
  strike > 0.0 || error("strike must be positive")
  oddSteps = isodd(steps) ? steps : steps + 1
  variance_ = variance(process, 0.0, common.x0, endTime)
  ermqdt = exp(common.driftPerStep + 0.5 * variance_ / oddSteps)
  d2 = (log(common.x0 / strike) + common.driftPerStep * oddSteps) / sqrt(variance_)

  pu = peizer_pratt_method_2_inversion(d2, oddSteps)
  pd = 1.0 - pu
  pdash = peizer_pratt_method_2_inversion(d2 + sqrt(variance_), oddSteps)

  up = ermqdt * pdash / pu
  down = (ermqdt - pu * up) / (1.0 - pu)

  return BinomialTree{LeisenReimer}(up, down, pu, pd, common, LeisenReimer())
end

function compute_up_prob(k::Float64, dj::Float64)
  alpha = dj / sqrt(8.0)
  # alpha2 = alpha * alpha
  alpha3 = alpha ^ 3
  alpha5 = alpha ^ 5
  alpha7 = alpha ^ 7
  beta = -0.375 * alpha - alpha3
  gamma = (5.0 / 6.0) * alpha5 + (13.0/12.0) * alpha3 + (25.0 / 128.0) * alpha
  delta = -0.1025 * alpha - 0.9285 * alpha3 - 1.43 * alpha5 - 0.5 * alpha7
  p = 0.5
  rootk = sqrt(k)
  p += alpha / rootk
  p += beta / (k * rootk)
  p += gamma / (k * k * rootk)
  # delete next line to get results for j-three tree
  p += delta / (k * k * k * rootk)
  return p
end

function Joshi4(process::StochasticProcess1D, endTime::Float64, steps::Int, strike::Float64)
  common = BinomialTreeCommon(process, endTime, steps)
  strike > 0.0 || error("strike must be positive")
  oddSteps = isodd(steps) ? steps : steps + 1
  variance_ = variance(process, 0.0, common.x0, endTime)
  ermqdt = exp(common.driftPerStep + 0.5 * variance_ / oddSteps)
  d2 = (log(common.x0 / strike) + common.driftPerStep * oddSteps) / sqrt(variance_)

  pu = compute_up_prob((oddSteps - 1.0) / 2.0, d2)
  pd = 1.0 - pu
  pdash = compute_up_prob((oddSteps - 1.0) / 2.0, d2 + sqrt(variance_))

  up = ermqdt * pdash / pu
  down = (ermqdt - pu * up) / (1.0 - pu)

  return BinomialTree{Joshi4}(up, down, pu, pd, common, Joshi4())
end

struct JarrowRudd <: EqualProbabilitiesBinomialTreeType end
struct AdditiveEQP <: EqualProbabilitiesBinomialTreeType end

mutable struct EqualProbabilitiesBinomialTree{T <: EqualProbabilitiesBinomialTreeType} <: AbstractBinomialTree
  up::Float64
  common::BinomialTreeCommon
  treeType::T
end

const JarrowRuddBinomialTree = EqualProbabilitiesBinomialTree{JarrowRudd}
const AdditiveEQPBinomialTree = EqualProbabilitiesBinomialTree{AdditiveEQP}

# overloaded to build this alias
function JarrowRudd(process::StochasticProcess1D, endTime::Float64, steps::Int, ::Float64)
  common = BinomialTreeCommon(process, endTime, steps)
  up = std_deviation(process, 0.0, common.x0, common.dt)

  return EqualProbabilitiesBinomialTree{JarrowRudd}(up, common, JarrowRudd())
end

function AdditiveEQP(process::StochasticProcess1D, endTime::Float64, steps::Int, ::Float64)
  common = BinomialTreeCommon(process, endTime, steps)
  up = -0.5 * common.driftPerStep + 0.5 * sqrt(4.0 * variance(process, 0.0, common.x0, common.dt) - 3.0 * common.driftPerStep * common.driftPerStep)

  return EqualProbabilitiesBinomialTree{AdditiveEQP}(up, common, AdditiveEQP())
end

struct CoxRossRubinstein <: EqualJumpsBinomialTreeType end
struct Trigeorgis <: EqualJumpsBinomialTreeType end

mutable struct EqualJumpsBinomialTree{T <: EqualJumpsBinomialTreeType} <: AbstractBinomialTree
  dx::Float64
  pu::Float64
  pd::Float64
  common::BinomialTreeCommon
  treeType::T
end

const CoxRossRubinsteinBinomialTree = EqualJumpsBinomialTree{CoxRossRubinstein}
const TrigeorgisBinomialTree = EqualJumpsBinomialTree{Trigeorgis}

function CoxRossRubinstein(process::StochasticProcess1D, endTime::Float64, steps::Int, ::Float64)
  common = BinomialTreeCommon(process, endTime, steps)
  dx = std_deviation(process, 0.0, common.x0, common.dt)
  pu = 0.5 + 0.5 * common.driftPerStep / dx
  pd = 1.0 - pu

  pu <= 1.0 || error("negative probability")
  pu >= 0.0 || error("negative probability")

  return EqualJumpsBinomialTree{CoxRossRubinstein}(dx, pu, pd, common, CoxRossRubinstein())
end

function Trigeorgis(process::StochasticProcess1D, endTime::Float64, steps::Int, ::Float64)
  common = BinomialTreeCommon(process, endTime, steps)
  dx = sqrt(variance(process, 0.0, common.x0, common.dt) + common.driftPerStep * common.driftPerStep)
  pu = 0.5 + 0.5 * common.driftPerStep / dx
  pd = 1.0 - pu

  pu <= 1.0 || error("negative probability")
  pu >= 0.0 || error("negative probability")

  return EqualJumpsBinomialTree{Trigeorgis}(dx, pu, pd, common, Trigeorgis())
end

get_size(tree::AbstractBinomialTree, idx::Int) = idx

descendant(tree::AbstractBinomialTree, ::Int, idx::Int, branch::Int) = (idx - 1) + (branch - 1) + 1

function get_underlying(tree::EqualProbabilitiesBinomialTree, i::Int, idx::Int)
  j = 2 * (idx - 1) - (i - 1)
  # exploiting the forward value tree centering
  return get_x0(tree) * exp((i-1) * get_drift_per_step(tree) + j * tree.up)
end
function get_underlying(tree::EqualJumpsBinomialTree, i::Int, idx::Int)
  j = 2 * (idx - 1) - (i - 1)
  # exploting equal jump and the x0 tree centering
  return get_x0(tree) * exp(j * tree.dx)
end
get_underlying(tree::BinomialTree, i::Int, idx::Int) = get_x0(tree) * ^(tree.down, (i - 1) - (idx - 1)) * ^(tree.up, (idx - 1))

probability(tree::EqualProbabilitiesBinomialTree, ::Int, ::Int, ::Int) = 0.5
probability(tree::Union{EqualJumpsBinomialTree, BinomialTree}, ::Int, ::Int, branch::Int) = branch == 2 ? tree.pu : tree.pd


push!(LOAD_PATH, "../LF_common/")
using m_LF3d

include("init_pot_Hcoul.jl")
include("init_pot_harm_3d.jl")
include("../utils/ortho_gram_schmidt.jl")
include("../utils/orthonormalize.jl")
include("EnergiesT.jl")
include("op_nabla2.jl")
include("op_H.jl")
include("calc_Energies.jl")
include("calc_grad.jl")

function test_grad()
  # LF parameters
  NN = [25, 25, 25]
  AA = [0.0, 0.0, 0.0]
  BB = [6.0, 6.0, 6.0]

  Npoints = prod(NN)

  # Initialize LF
  LF = init_LF3d_p( NN, AA, BB )
  ΔV = LF.LFx.h * LF.LFy.h * LF.LFz.h

  # Parameter for potential
  center = AA + 0.5*(BB-AA)
  # Potential
  ω = 2.0
  Vpot = init_pot_harm_3d( LF, ω, center )

  Ncols = 4

  ΔV = LF.LFx.h * LF.LFy.h * LF.LFz.h

  println("sum(Vpot) = ", sum(Vpot))
  println("ΔV = ", ΔV)

  v = zeros(Npoints,Ncols)
  for i = 1:Ncols
    v[i,i] = 1.0/sqrt(ΔV)
  end

  grad = calc_grad( LF, Vpot, v )
  println(sum(grad))

  Energies = calc_Energies( LF, Vpot, v)
  println(Energies)
end

test_grad()


s = open("input.txt") do file
    readlines(file);
end

function RunProgramUntilRepeat(program)
    seen = Set()

    i = 1;
    acc = 0;
    while !in(i, seen)
        push!(seen, i);
        parts = split(program[i], " ");

        cmd = parts[1];
        num = parse(Int, parts[2]);

        if cmd == "nop"
            i += 1;
        elseif cmd == "acc"
            acc += num;
            i += 1;
        elseif cmd == "jmp"
            i += num
        end
    end

    return acc
end

RunProgramUntilRepeat(s)


# Load Julia packages (libraries) needed.

using StatisticalRethinking
using StatsPlots

ProjDir = @__DIR__

# ### snippet 5.39

for i in 5:7
  include("$(ProjDir)/m5.$i.jl")
end

if success(rc)

  r1 = plotcoef([m5_5s, m5_6s, m5_7s], [:a, :bN, :bM], "$(ProjDir)/Fig-39a.png",
    "bN & bM Normal estimates")
  r1 |> display

  r2 = plotcoef([m5_5s, m5_6s, m5_7s], [:a, :bN, :bM], "$(ProjDir)/Fig-39b.png",
    "bN & bM Quap estimates", quap)
  r2 |> display

end


@testset "contrasts" begin

    d = DataFrame(y = rand(6),
                  x = [:b, :a, :c, :a, :a, :b])

    mf = ModelFrame(@formula(y ~ x), d)

    ## testing equality of ContrastsMatrix
    should_equal = [ContrastsMatrix(DummyCoding(), [:a, :b, :c]),
                    ContrastsMatrix(DummyCoding(base=:a), [:a, :b, :c]),
                    ContrastsMatrix(DummyCoding(base=:a, levels=[:a, :b, :c]), [:a, :b, :c]),
                    ContrastsMatrix(DummyCoding(levels=[:a, :b, :c]), [:a, :b, :c])]

    for c in should_equal
        @test mf.schema[term(:x)].contrasts == c
        @test hash(mf.schema[term(:x)].contrasts) == hash(c)
    end

    should_not_equal = [ContrastsMatrix(EffectsCoding(), [:a, :b, :c]),
                        ContrastsMatrix(DummyCoding(), [:b, :c]),
                        ContrastsMatrix(DummyCoding(), [:b, :c, :a]),
                        ContrastsMatrix(DummyCoding(), [:b, :c, :a]),
                        ContrastsMatrix(DummyCoding(), [:b, :c, :a]),
                        ContrastsMatrix(DummyCoding(base=:b), [:a, :b, :c]),
                        ContrastsMatrix(DummyCoding(base=:a, levels=[:b, :a, :c]), [:a, :b, :c])]

    for c in should_not_equal
        @test mf.schema[term(:x)].contrasts != c
    end

    # Dummy coded contrasts by default:
    @test ModelMatrix(mf).m == [1  1  0
                                1  0  0
                                1  0  1
                                1  0  0
                                1  0  0
                                1  1  0]
    @test coefnames(mf) == ["(Intercept)"; "x: b"; "x: c"]

    mmm = ModelMatrix(mf).m
    setcontrasts!(mf, x = DummyCoding())
    @test ModelMatrix(mf).m == mmm

    setcontrasts!(mf, x = EffectsCoding())
    @test ModelMatrix(mf).m == [1  1  0
                                1 -1 -1
                                1  0  1
                                1 -1 -1
                                1 -1 -1
                                1  1  0]
    @test coefnames(mf) == ["(Intercept)"; "x: b"; "x: c"]

    # change base level of contrast
    setcontrasts!(mf, x = EffectsCoding(base = :b))
    @test ModelMatrix(mf).m == [1 -1 -1
                                1  1  0
                                1  0  1
                                1  1  0
                                1  1  0
                                1 -1 -1]
    @test coefnames(mf) == ["(Intercept)"; "x: a"; "x: c"]

    # change levels of contrast
    setcontrasts!(mf, x = EffectsCoding(levels = [:c, :b, :a]))
    @test ModelMatrix(mf).m == [1  1  0
                                1  0  1
                                1 -1 -1
                                1  0  1
                                1  0  1
                                1  1  0]
    @test coefnames(mf) == ["(Intercept)"; "x: b"; "x: a"]

    # change levels and base level of contrast
    setcontrasts!(mf, x = EffectsCoding(levels = [:c, :b, :a], base = :a))
    @test ModelMatrix(mf).m == [1  0  1
                                1 -1 -1
                                1  1  0
                                1 -1 -1
                                1 -1 -1
                                1  0  1]
    @test coefnames(mf) == ["(Intercept)"; "x: c"; "x: b"]

    # respect order of levels

    data = DataFrame(x = levels!(categorical(['A', 'B', 'C', 'C', 'D']), ['C', 'B', 'A', 'D']))
    f = apply_schema(@formula(x ~ 1), schema(data))
    @test modelcols(f.lhs, data) == [0 1 0; 1 0 0; 0 0 0; 0 0 0; 0 0 1]

    # Helmert coded contrasts
    setcontrasts!(mf, x = HelmertCoding())
    @test ModelMatrix(mf).m == [1  1 -1
                                1 -1 -1
                                1  0  2
                                1 -1 -1
                                1 -1 -1
                                1  1 -1]
    @test coefnames(mf) == ["(Intercept)"; "x: b"; "x: c"]

    # Mismatching types of data and contrasts levels throws an error:
    @test_throws ArgumentError setcontrasts!(mf, x = EffectsCoding(levels = ["a", "b", "c"]))

    # Missing data is handled gracefully, dropping columns when a level is lost
    allowmissing!(d,:x)
    d[3, :x] = missing
    mf_missing = ModelFrame(@formula(y ~ x), d,
                            contrasts = Dict(:x => EffectsCoding()))
    @test ModelMatrix(mf_missing).m == [1  1
                                        1 -1
                                        1 -1
                                        1 -1
                                        1  1]
    @test coefnames(mf_missing) == ["(Intercept)"; "x: b"]

    # Things that are bad to do:
    # Applying contrasts that only have a subset of data levels:
    @test_throws ArgumentError setcontrasts!(mf, x = EffectsCoding(levels = [:a, :b]))
    # Applying contrasts that expect levels not found in data:
    @test_throws ArgumentError setcontrasts!(mf, x = EffectsCoding(levels = [:a, :b, :c, :d]))
    # Asking for base level that's not found in data
    @test_throws ArgumentError setcontrasts!(mf, x = EffectsCoding(base = :e))

    # Manually specified contrasts
    contrasts = [0  1
                 -1 -.5
                 1  -.5]
    setcontrasts!(mf, x = ContrastsCoding(contrasts))
    @test ModelMatrix(mf).m == [1 -1 -.5
                                1  0  1
                                1  1 -.5
                                1  0  1
                                1  0  1
                                1 -1 -.5]

    # throw argument error if number of levels mismatches
    @test_throws ArgumentError setcontrasts!(mf, x = ContrastsCoding(contrasts[1:2, :]))
    @test_throws ArgumentError setcontrasts!(mf, x = ContrastsCoding(hcat(contrasts, contrasts)))

    # contrasts types must be instantiated (should throw ArgumentError, currently
    # MethodError on apply_schema)
    @test_broken setcontrasts!(mf, x = DummyCoding)

end


using HigherPrecision
using Base.Test

@testset "DoubleFloat64" begin
    @test DoubleFloat64(Float32(2.0)) == DoubleFloat64(2.0)
    @test DoubleFloat64(Float16(2.0)) == DoubleFloat64(2.0)
    @test DoubleFloat64(2) == DoubleFloat64(2.0)
    @test DoubleFloat64(big(π)) isa FastDouble
    @test DoubleFloat64(π) == DoubleFloat64(big(π))

    @test DoubleFloat64(2 // 3) isa FastDouble
    @test DoubleFloat64(2 // 3, ComputeAccurate) isa AccurateDouble
    @test DoubleFloat64(Float32(2.0), ComputeAccurate) isa AccurateDouble
    @test DoubleFloat64(Float16(2.0), ComputeAccurate) isa AccurateDouble
    @test DoubleFloat64(2, ComputeAccurate) isa AccurateDouble
    @test DoubleFloat64(big(π), ComputeAccurate) isa AccurateDouble
    @test DoubleFloat64(π, ComputeAccurate) isa AccurateDouble

    @test HigherPrecision.hi(DoubleFloat64(2.0)) == 2.0
    @test HigherPrecision.lo(DoubleFloat64(2.0)) == 0.0

    @test isbits(FastDouble) == true
    @test isbits(AccurateDouble) == true

    @test zero(DoubleFloat64(2.0)) == zero(FastDouble)

    @test convert(Integer, DoubleFloat64(2.0)) isa Int64
    @test 32 ≤ length(string(rand(FastDouble))) ≤ 40

    @test convert(DoubleFloat64, 2.0) isa FastDouble
    @test convert(DoubleFloat64, Float32(2.0)) isa FastDouble
    @test convert(DoubleFloat64, Float16(2.0)) isa FastDouble
    @test convert(DoubleFloat64, 2) isa FastDouble
    @test convert(DoubleFloat64, π) isa FastDouble

    @test promote_type(DoubleFloat64, Float64) == FastDouble
    @test promote_type(DoubleFloat64, Float32) == FastDouble
    @test promote_type(DoubleFloat64, Float16) == FastDouble

    @test double_sub(rand(), rand()) isa FastDouble
    @test double_add(rand(), rand()) isa FastDouble
    @test double_div(rand(), rand()) isa FastDouble
    @test double_mul(rand(), rand()) isa FastDouble

    x, y = rand(FastDouble), rand(AccurateDouble)

    @test x isa FastDouble
    @test y isa AccurateDouble
    @test x * y isa AccurateDouble
    @test y * x isa AccurateDouble
    @test x / y isa AccurateDouble
    @test y / x isa AccurateDouble
    @test x - y isa AccurateDouble
    @test y - x isa AccurateDouble
    @test x + y isa AccurateDouble
    @test y + x isa AccurateDouble
    @test x / 3.2 isa FastDouble
    @test y / 3.2 isa AccurateDouble

    @test x + 2 isa FastDouble
    @test 2 + x isa FastDouble

    @test im * rand(DoubleFloat64) isa Complex{FastDouble}
    @test x < 2 * x
    @test x ≤ x
    @test convert(Float64, x) ≤ x
    @test x ≤ 2 * convert(Float64, x)
    @test DoubleFloat64(2.0) == 2.0
    @test 2.0 == DoubleFloat64(2.0)

    @test isnan(DoubleFloat64(NaN))
    @test isinf(DoubleFloat64(Inf))

    @test floor(DoubleFloat64(3.2)) == 3.0
    @test ceil(DoubleFloat64(3.2)) == 4.0
    @test trunc(DoubleFloat64(3.2)) == 3.0
    @test isinteger(DoubleFloat64(3.2)) == false
    @test isinteger(DoubleFloat64(3.0)) == true

    @test DoubleFloat64(0.0)^0.14 == 1.0

    @test rand(Base.Random.GLOBAL_RNG, FastDouble, 4) isa Vector{FastDouble}
    @test rand(Base.Random.GLOBAL_RNG, Complex{FastDouble}, 4) isa Vector{Complex{FastDouble}}
    @test rand(Complex{DoubleFloat64}, 3) isa Vector{Complex{FastDouble}}
    @test rand(Base.Random.GLOBAL_RNG, Complex{DoubleFloat64}, 4) isa Vector{Complex{FastDouble}}
    @test length(rand(Base.Random.GLOBAL_RNG, FastDouble, 4)) == 4

    for k = 1:5000
        x = rand(DoubleFloat64) * 20 - 10
        y = rand(DoubleFloat64) * 20 - 10
        @test x * y ≈ big(x) * big(y) atol=1e-29
        @test x + y ≈ big(x) + big(y) atol=1e-30
        @test x - y ≈ big(x) - big(y) atol=1e-30
        @test x / y ≈ big(x) / big(y) atol=1e-26
        u = rand()
        @test x / u ≈ big(x) / big(u) atol=1e-26
        @test x^5 ≈ BigFloat(x)^5 atol=1e-25

        @test exp(x) ≈ exp(big(x)) atol=1e-26
        @test log(abs(x)) ≈ log(big(abs(x))) atol=1e-29
        @test sin(x) ≈ sin(big(x)) atol=1e-30
        @test cos(x) ≈ cos(big(x)) atol=1e-30
        @test atan(x) ≈ atan(big(x)) atol=1e-30
        @test atan2(y, x) ≈ atan2(big(y), big(x)) atol=1e-30
        @test tan(x) ≈ tan(big(x)) atol=1e-21

        sinh_x, cosh_x = sincosh(x)
        if abs(x) > 0.05
            @test sinh(x) ≈ sinh(big(x)) atol=1e-25
            @test tanh(x) ≈ tanh(big(x)) atol=1e-25
            @test sinh_x ≈ sinh(big(x)) atol=1e-25
            @test cosh_x ≈ cosh(big(x)) atol=1e-25
        else
            @test sinh(x) ≈ sinh(big(x)) atol=1e-16
            @test tanh(x) ≈ tanh(big(x)) atol=1e-16
            @test sinh_x ≈ sinh(big(x)) atol=1e-16
            @test cosh_x ≈ cosh(big(x)) atol=1e-16
        end
        @test cosh(x) ≈ cosh(big(x)) atol=1e-20

        sin_x, cos_x = HigherPrecision.sincos(x)
        @test sin_x ≈ sin(big(x)) atol=1e-30
        @test cos_x ≈ cos(big(x)) atol=1e-30

        z = DoubleFloat64(rand() * 2 - 1)
        @test asin(z) ≈ asin(big(z)) atol=1e-30
        @test acos(z) ≈ acos(big(z)) atol=1e-30

        n, p, d = Base.decompose(x)
        @test x ≈ BigFloat(n) * BigFloat(2)^p / BigFloat(d)
    end

    for k = 1:5000
        x = rand(DoubleFloat64{ComputeAccurate}) * 20 - 10
        y = rand(DoubleFloat64{ComputeAccurate}) * 20 - 10
        @test x * y ≈ big(x) * big(y) atol=1e-29
        @test x + y ≈ big(x) + big(y) atol=1e-31
        @test x - y ≈ big(x) - big(y) atol=1e-31
        @test x / y ≈ big(x) / big(y) atol=1e-28
        @test x^5 ≈ BigFloat(x)^5 atol=1e-26

        @test exp(x) ≈ exp(big(x)) atol=1e-26
        @test log(abs(x)) ≈ log(big(abs(x))) atol=1e-29
        @test sin(x) ≈ sin(big(x)) atol=1e-30
        @test cos(x) ≈ cos(big(x)) atol=1e-30
        @test atan(x) ≈ atan(big(x)) atol=1e-30
        @test atan2(y, x) ≈ atan2(big(y), big(x)) atol=1e-30
        @test tan(x) ≈ tan(big(x)) atol=1e-21

        sinh_x, cosh_x = sincosh(x)
        if abs(x) > 0.05
            @test sinh(x) ≈ sinh(big(x)) atol=1e-25
            @test tanh(x) ≈ tanh(big(x)) atol=1e-25
            @test sinh_x ≈ sinh(big(x)) atol=1e-25
            @test cosh_x ≈ cosh(big(x)) atol=1e-25
        else
            @test sinh(x) ≈ sinh(big(x)) atol=1e-16
            @test tanh(x) ≈ tanh(big(x)) atol=1e-16
            @test sinh_x ≈ sinh(big(x)) atol=1e-16
            @test cosh_x ≈ cosh(big(x)) atol=1e-16
        end
        @test cosh(x) ≈ cosh(big(x)) atol=1e-20

        sin_x, cos_x = HigherPrecision.sincos(x)
        @test sin_x ≈ sin(big(x)) atol=1e-30
        @test cos_x ≈ cos(big(x)) atol=1e-30

        z = DoubleFloat64(rand() * 2 - 1)
        @test asin(z) ≈ asin(big(z)) atol=1e-30
        @test acos(z) ≈ acos(big(z)) atol=1e-30
    end

    @test eps(DoubleFloat64(1.0)) == 4.93038065763132e-32 # 2^-104
    @test eps(DoubleFloat64) == 4.93038065763132e-32 # 2^-104
    @test eps(FastDouble) == 4.93038065763132e-32 # 2^-104
    @test eps(AccurateDouble) == 4.93038065763132e-32 # 2^-104
    @test realmin(DoubleFloat64) > realmin(Float64)
    @test realmax(DoubleFloat64) > realmax(Float64)
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.7.0-beta2
#     language: julia
#     name: julia-1.7
# ---

# %%
@show VERSION

using Plots

f(x, y) = (1 - x)^2 + 100(y - x^2)^2

X = Y = range(-5, 5; length=1001)
@show x, y = argmin(Base.splat(f), Iterators.product(X, Y))
@show f(x, y)

surface(X, Y, log∘f; color=:rainbow, clim = (-7.5, 10), camera=(30, 60))
scatter!([x], [y], [-7.5]; label="minimum", color=:cyan)
plot!(; xtick = -5:5, ytick = -5:5, zlim=(-7.5, 10))

# %%

0008/Base.splat example - minimum of rosenbrock2d function.jl


include("test_1d.jl")
include("test_2d.jl")
include("test_3d.jl")
include("test_4d.jl")
# include("test_5d.jl")


using ConvolutionTools
using Base.Test

import Base.isapprox

function isapprox{T, N}(a::Array{T, N}, b::Array{T, N}, eps=10e-10)
    if size(a) != size(b)
        return false
    end
    el_diff = sum(abs(a .- b)) / prod(size(a))
    return el_diff < eps
end

function test_pad_const_1d()
    a = rand(3)
    n = [1]
    c = 2.3
    res = pad_const(a, n, c)
    return length(res) == 5 && res[1] == c && res[end] == c && typeof(res) == typeof(a)
end

function test_pad_const_2d()
    a = rand(3, 2)
    n = [1, 3]
    res = pad_const(a, n)
    return size(res) == (5, 8)
end

function test_1d()
    a = rand((7))
    b = rand((2))
    res = ConvolutionTools.conv_valid(a, b)
    return ndims(res) == 1 && eltype(res) <: Number
end

function test_2d()
    a = rand((6, 7))
    b = rand((1, 2))
    res = ConvolutionTools.conv_valid(a, b)
    return ndims(res) == 2 && eltype(res) <: Number
end

function test_3d()
    a = rand((6, 7, 10))
    b = rand((1, 2, 3))
    res = ConvolutionTools.conv_valid(a, b)
    return ndims(res) == 3 && eltype(res) <: Number
end

function test_identity()
    a = [i^2 for i in 1:100]
    b = [1]
    res = ConvolutionTools.conv_valid(a, b)
    return all(a .== res)
end

function test_zero()
    a = float([i^2 for i in 1:100])
    b = float([0])
    res = ConvolutionTools.conv_valid(a, b)
    return all(res .== zero(eltype(res)))
end

function test_derivative()
    n = 10
    a = float([0.5 * i^2 for i in 1:n])
    b = float([1, 0, -1]) / 2
    res = ConvolutionTools.conv_valid(a, b)
    ref = float([i for i in 2:(n-1)])
    return isapprox(ref, res)
end

function test_integers()
    n = 10
    a = [i for i in 1:n]
    b = [1, 0, -1]
    res = ConvolutionTools.conv_valid(a, b)
    return all(res .== 2)
end

function test_mixed_input()
    a = [i for i in 1:10]
    b = float([1, 2, 3])
    res1 = ConvolutionTools.conv(a, b)
    res2 = ConvolutionTools.conv(b, a)
    return isapprox(res1, res2)
end

custom_handler(r::Test.Success) = print(".")
custom_handler(r::Test.Failure) = Test.default_handler(r)
custom_handler(r::Test.Error)   = Test.default_handler(r)

Test.with_handler(custom_handler) do
    @test test_pad_const_1d()
    @test test_pad_const_2d()
    @test test_1d()
    @test test_2d()
    @test test_3d()
    @test test_zero()
    @test test_identity()
    @test test_derivative()
    @test test_integers()
    @test test_mixed_input()
end

println("")


using Test, SparseArrays
import LuxurySparse: SparseMatrixCOO, allocated_coo, isdense

coo1 = SparseMatrixCOO(
    [1, 4, 2, 3, 3, 3],
    [1, 1, 2, 4, 3, 4],
    [0.1, 0.2, 0.4im, 0.5, 0.3, 0.5im],
    4,
    4,
)

@testset "basic" begin
    @test eltype(coo1) == ComplexF64
    @test size(coo1, 1) == size(coo1, 2) == 4

    coo2 = allocated_coo(ComplexF64, 4, 4, 6)
    @test eltype(coo2) == ComplexF64
    @test nnz(coo2) == 6
    @test size(coo2) == (4, 4)
    @test size(coo2, 1) == size(coo2, 2) == 4

    @test coo1 == copy(coo1)
    @test coo1 == copyto!(coo2, coo1)

    @test Matrix(coo1) == [0.1 0 0 0; 0 0.4im 0 0; 0 0 0.3 0.5+0.5im; 0.2 0 0 0]
    @test sparse(coo1) == sparse(findnz(coo1)...)
    @test isdense(coo1) == false

    A = sprand(50, 50, 0.2) |> Matrix
    @test SparseMatrixCOO(A) == A
end

@testset "sparse" begin
    p1 = copy(coo1)
    @test LuxurySparse.nonzeros(p1) == p1.vs
    p1.vs[2] = 0
    @test LuxurySparse.dropzeros!(p1) ==
          [0.1 0 0 0; 0 0.4im 0 0; 0 0 0.3 0.5+0.5im; 0 0 0 0]
end

@testset "basicmath" begin
    @test coo1 * 2 == Matrix(coo1) * 2
    @test 2 * coo1 == 2 * Matrix(coo1)
    @test coo1 / 2 == Matrix(coo1) / 2
    @test coo1 * 2 isa SparseMatrixCOO
    @test coo1 / 2 isa SparseMatrixCOO
    @test 2 * coo1 isa SparseMatrixCOO
end

@testset "setindex" begin
    coo[1, 1] = 0.2
    @test coo.is[end] == 1
    @test coo.js[end] == 1
    @test coo.vs[end] == 0.2

    S = SparseMatrixCOO{Float32, Cint}(Cint[], Cint[], Float32[], 5, 5)
    S[1, 1] = 3
    @test S[1, 1] == 3
end


function ADforward!(net;returnf=false,debug=false,AllocateMemory=false)
#=   calculate the value and auxiliary values of a net function
=#
    # (c) David Barber, University College London 2015

    if AllocateMemory & net.gpu
        TransformToGPU=true
        net=convert(net,"CPU")
    else
        TransformToGPU=false
    end

    if debug; println("Get value:"); end

    if isempty(net.auxvalue)
        net.auxvalue=similar(net.value)
    end

    # forward pass:
    for i in net.ForwardPassList
        thisnode=net.node[i]
        if debug
            println("node $i: $(thisnode.f)($(thisnode.parents))");
        end
        if AllocateMemory
            if debug
                @time net.value[i],net.auxvalue[i]=thisnode.f(net.value[thisnode.parents]...) ## not in place
            else
                net.value[i],net.auxvalue[i]=thisnode.f(net.value[thisnode.parents]...) ## not in place
            end
        else
            if debug
                println("in place")
                @time thisnode.f_inplace(net.value[i],net.auxvalue[i],net.value[thisnode.parents]...) ## in place
            else
                thisnode.f_inplace(net.value[i],net.auxvalue[i],net.value[thisnode.parents]...) ## in place
            end
        end
    end

    if TransformToGPU
        net=convert(net,"GPU")
    end

    if returnf
        if net.gpu
            return to_host(net.value[net.FunctionNode])[1]
        else
            return net.value[net.FunctionNode][1]
        end
#        @gpu if net.gru
#        @cpu return net.value[net.FunctionNode][1]
    end

    if debug; println("done"); end
end


# ---
# title: 95. Unique Binary Search Trees II
# id: problem95
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming, Tree
# link: <https://leetcode.com/problems/unique-binary-search-trees-ii/description/>
# hidden: true
# ---
#
# Given an integer `n`, generate all structurally unique **BST 's** (binary
# search trees) that store values 1 ...  _n_.
#
# **Example:**
#
#
#
#     Input: 3
#     Output:
#     [
#       [1,null,3,2],
#       [3,2,null,1],
#       [3,1,null,null,2],
#       [2,1,3],
#       [1,null,2,null,3]
#     ]
#     Explanation:
#     The above output corresponds to the 5 unique BST's shown below:
#
#        1         3     3      2      1
#         \       /     /      / \      \
#          3     2     1      1   3      2
#         /     /       \                 \
#        2     1         2                 3
#
#
#
#
# **Constraints:**
#
#   * `0 <= n <= 8`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end