"""
    step!(solution::AbstractRungeKuttaSolution, problem::AbstractInitialValueProblem, solver::ExplicitRungeKuttaSolver, cache::ExplicitRungeKuttaCache)

computes a step of the [`AbstractRungeKuttaSolution`](@ref) of an [`AbstractInitialValueProblem`](@ref) using an [`ExplicitRungeKuttaSolver`](@ref).
"""
function step!(solution::AbstractRungeKuttaSolution, problem::AbstractInitialValueProblem, solver::ExplicitRungeKuttaSolver, cache::ExplicitRungeKuttaCache)
    @↓ n, v = cache
    @↓ u, t = solution
    k = solution.k isa Nothing ? cache.k : solution.k[n]
    @↓ f! = problem.rhs
    @↓ s, A, c, b = solver.tableau
    @↓ h = solver.stepsize
    v = u[n+1] # avoid allocs
    # Compute stages
    for i = 1:s
        zero!(v)
        for j = 1:i-1
            @. v += A[i,j] * k[j]
        end
        @. v = u[n] + h * v
        # @← k[i] = f(v, t[n] + h * c[i])
        f!(k[i], v, t[n] + h * c[i])
    end
    # Compute step
    zero!(v)
    for i = 1:s
        @. v += b[i] * k[i]
    end
    @. v = u[n] + h * v
    t[n+1] = t[n] + h
    return u[n+1], t[n+1]
end

"""
    step!(solution::AbstractRungeKuttaSolution, problem::AbstractInitialValueProblem, solver::ImplicitRungeKuttaSolver, cache::ImplicitRungeKuttaCache)

computes a step of the [`AbstractRungeKuttaSolution`](@ref) of an [`AbstractInitialValueProblem`](@ref) using an [`ImplicitRungeKuttaSolver`](@ref).
"""
function step!(solution::AbstractRungeKuttaSolution, problem::AbstractInitialValueProblem, solver::ImplicitRungeKuttaSolver, cache::ImplicitRungeKuttaCache)
    @↓ n, v, Δk, J = cache
    @↓ u, t = solution
    k = solution.k isa Nothing ? cache.k : solution.k[n]
    @↓ f!, Df! = problem.rhs
    @↓ s, A, c, b = solver.tableau
    @↓ ϵ, K = solver.newton
    @↓ h = solver.stepsize
    v = u[n+1] # avoid allocs
    # @← J = Df(v, u[n], t[n])
    Df!(J, v, u[n], t[n])
    Z = factorize(I - h * kron(A, J))
    # Compute stages
    zero!(k)
    for l = 1:K
        for i = 1:s
            zero!(v)
            for j = 1:s
                @. v += A[i,j] * k[j]
            end
            @. v = u[n] + h * v
            # @← Δk[i] = f(v, t[n] + h * c[i])
            f!(Δk[i], v, t[n] + h * c[i])
            @. Δk[i] -= k[i]
        end
        # temporary workaround for ldiv!(Z, Δk)
        Δk_ = vcat(Δk...)
        ldiv!(Z, Δk_)
        if norm(Δk_) < ϵ * norm(k)
            break
        end
        # temporary workaround for k .+= Δk
        for i in eachindex(k)
            L = length(k[i])
            k[i] .+= Δk_[(i-1)*L+1:i*L]
        end
    end
    # Compute step
    zero!(v)
    for i = 1:s
        @. v += b[i] * k[i]
    end
    @. v = u[n] + h * v
    t[n+1] = t[n] + h
    return u[n+1], t[n+1]
end

function step!(solution::AbstractRungeKuttaSolution, problem::AbstractInitialValueProblem, solver::ExplicitExponentialRungeKuttaSolver, cache::ExplicitExponentialRungeKuttaCache)
    @↓ n, Q, α, β, γ, E, E2 = cache
    @↓ u, t = solution
    @↓ f = problem.rhs.rhs
    @↓ h = solver.stepsize

    uₙ = u[n]
    tₙ = t[n+1] = t[n] + h

    Nu = f(uₙ, tₙ)
    aₙ = @. E2 * uₙ + Q * Nu
    Na = f(aₙ, tₙ)
    bₙ = @. E2 * uₙ + Q * Na
    Nb = f(bₙ, tₙ)
    cₙ = @. E2 * aₙ + Q * (2Nb - Nu)
    Nc = f(cₙ, tₙ)

    @. u[n+1] = E * uₙ + α * Nu + 2β * (Na + Nb) + γ * Nc
    return u[n+1], t[n+1]
end


## Core CAOL routines

using OffsetArrays, ImageFiltering, LinearAlgebra

# Allocation and initialization
function _initvars(x,H0,R)
    K = size(H0,2)

    # Form padded images
    xpad = [padarray(xl,Pad(:circular,ntuple(_->0,ndims(xl)),R)) for xl in x]
    # xpad = [view(xl,ntuple(i->vcat(1:size(xl,i),1:R[i]),ndims(xl))...) for xl in x]
    # more memory efficient but seems slower, make a copy for now
    # may be relevant: https://docs.julialang.org/en/v1/manual/performance-tips/#Copying-data-is-not-always-bad-1

    # Initialize filters
    H = copy(H0)
    h = _filterlist(H,R)
    Hprev = similar(H)  # for stopping condition

    # Initialize temporary variables
    zlk = similar(first(x),map(n->0:n-1,size(first(x))))
    ΨZ = similar(H)
    ψz = _filterlist(ΨZ,R)
    ψztemp = similar(first(ψz))
    HΨZ = similar(H,K,K)
    UVt = HΨZ

    return xpad, H, h, Hprev, zlk, ΨZ, ψz, ψztemp, HΨZ, UVt
end

# Update steps
@inline function _updateΨZ!(ΨZ,ψz,xpad,h,λ,zlk,ψztemp)
    obj = 0.0
    fill!(ΨZ,zero(eltype(ΨZ)))
    for xpadl in xpad, k in 1:length(h)
        imfilter!(zlk,xpadl,(h[k],),NoPad(),Algorithm.FIR())
        obj += _obj(zlk,λ)
        zlk .= hard.(zlk,sqrt(2λ))
        imfilter!(ψztemp,xpadl,(zlk,),NoPad(),Algorithm.FIR())
        ψz[k] .+= ψztemp
    end
    return obj
end
function _updateH!(H,ΨZ,H0,HΨZ,UVt)
    mul!(HΨZ,H0',ΨZ)
    F = svd!(HΨZ)
    mul!(UVt,F.U,F.Vt)
    mul!(H,H0,UVt)
end

# Iterative algorithms
function _CAOL(x,λ,H0,R,maxiters,tol)
    @assert H0'H0 ≈ (1/prod(R))*I
    xpad, H, h, Hprev, zlk, ΨZ, ψz, ψztemp, HΨZ, UVt = _initvars(x,H0,R)

    for t in 1:maxiters
        copyto!(Hprev,H)                      # Copy previous filters
        _updateΨZ!(ΨZ,ψz,xpad,h,λ,zlk,ψztemp) # Update ΨZ
        _updateH!(H,ΨZ,H0,HΨZ,UVt)            # Update filters as polar factor

        sqrt(sosdiff(Hprev,H) / (size(H,2)/size(H,1))) <= tol && break
    end

    return H
end
function _CAOLtrace(x,λ,H0,R,maxiters,tol)
    @assert H0'H0 ≈ (1/prod(R))*I
    xpad, H, h, Hprev, zlk, ΨZ, ψz, ψztemp, HΨZ, UVt = _initvars(x,H0,R)

    Htrace     = fill(H0,0)
    objtrace   = fill(NaN,0)
    Hdifftrace = fill(NaN,0)

    for t in 1:maxiters
        copyto!(Hprev,H)                            # Copy previous filters
        obj = _updateΨZ!(ΨZ,ψz,xpad,h,λ,zlk,ψztemp) # Compute objective, update ΨZ
        _updateH!(H,ΨZ,H0,HΨZ,UVt)                  # Update filters as polar factor

        push!(Htrace,copy(H))
        push!(objtrace,obj)
        push!(Hdifftrace, sqrt(sosdiff(Hprev,H) / (size(H,2)/size(H,1))))

        Hdifftrace[end] <= tol && break
    end

    return H, Htrace, objtrace, Hdifftrace
end


include("testsystem1.jl")

U = @inferred trace(LinearEvolution(Φ, Gaussian(b, Q)), 0 => Gaussian(x0, P0), endtime(10000))

X = @inferred trace(DynamicIterators.Sampled(LinearEvolution(Φ, Gaussian(b, Q))), 0 => x0, endtime(10000))

@show norm(mean(values(X)) - b) < eps()
@show cov(values(X))
@show cov(values(U)[end])
@test norm(cov(values(X)) - cov(values(U)[end])) < 0.2

using DynamicIterators: dyniterate

Random.seed!(11)
X = @inferred trace(DynamicIterators.Sampled(LinearEvolution(Φ, Gaussian(b, Q))), 0 => x0, endtime(10))
Y = Any[0 => x0]
Random.seed!(11)
Y = collecttrace(LinearEvolution(Φ, Gaussian(b, Q)), Trace(Sample(0 => x0), Y, endtime(10)))
@test X == trajectory(Y)


using ReachabilityAnalysis
using ReachabilityAnalysis: is_intersection_empty

#The task is to change the height from ``0``~[m] to ``1``~[m] within ``5``~[s].

const g = 9.81           # gravity constant in m/s^2
const R = 0.1            # radius of center mass in m
const l = 0.5            # distance of motors to center mass in m
const Mrotor = 0.1       # motor mass in kg
const M = 1.0            # center mass in kg
const m = M + 4*Mrotor   # total mass in kg
const mg = m*g

const Jx = (2/5)*M*R^2 + 2*l^2*Mrotor
const Jy = Jx
const Jz = (2/5)*M*R^2 + 4*l^2*Mrotor
const Cyzx = (Jy - Jz)/Jx
const Czxy = (Jz - Jx)/Jy
const Cxyz = 0.0 #(Jx - Jy)/Jz

const u₁ = 1.0
const u₂ = 0.0
const u₃ = 0.0

const Tspan = (0.0, 5.0)
const v3 = LazySets.SingleEntryVector(3, 12, 1.0)

@inline function quad_property(solz)
    tf = tend(solz)

    #Condition: b1 = (x[3] < 1.4) for all time
    unsafe1 = HalfSpace(-v3, -1.4) # unsafe: -x3 <= -1.4
    b1 =  all([is_intersection_empty(unsafe1, set(R)) for R in solz(0.0 .. tf)])
    #b1 = ρ(v3, solz) < 1.4

    #Condition: x[3] > 0.9 for t ≥ 1.0
    unsafe2 = HalfSpace(v3, 0.9) # unsafe: x3 <= 0.9
    b2 = all([is_intersection_empty(unsafe2, set(R)) for R in solz(1.0 .. tf)])

    #Condition: x[3] ⊆ Interval(0.98, 1.02) for t ≥ 5.0 (t=5 is `tf`)
    b3 = set(project(solz[end], vars=(3))) ⊆ Interval(0.98, 1.02)

    return b1 && b2 && b3
end

@taylorize function quadrotor!(dx, x, params, t)
    #unwrap the variables and the controllers; the last three are the controllers
    #x₁, x₂, x₃, x₄, x₅, x₆, x₇, x₈, x₉, x₁₀, x₁₁, x₁₂, u₁, u₂, u₃ = x
    x₁  = x[1]
    x₂  = x[2]
    x₃  = x[3]
    x₄  = x[4]
    x₅  = x[5]
    x₆  = x[6]
    x₇  = x[7]
    x₈  = x[8]
    x₉  = x[9]
    x₁₀ = x[10]
    x₁₁ = x[11]
    x₁₂ = x[12]

    #equations of the controllers
    F = (mg - 10*(x₃ - u₁)) + 3*x₆  # height control
    τϕ = -(x₇ - u₂) - x₁₀            # roll control
    τθ = -(x₈ - u₃) - x₁₁            # pitch control
    local τψ = 0.0                   # heading is uncontrolled

    Tx = τϕ/Jx
    Ty = τθ/Jy
    Tz = τψ/Jz
    F_m = F/m

    #Some abbreviations
    sx7 = sin(x₇)
    cx7 = cos(x₇)
    sx8 = sin(x₈)
    cx8 = cos(x₈)
    sx9 = sin(x₉)
    cx9 = cos(x₉)

    sx7sx9 = sx7*sx9
    sx7cx9 = sx7*cx9
    cx7sx9 = cx7*sx9
    cx7cx9 = cx7*cx9
    sx7cx8 = sx7*cx8
    cx7cx8 = cx7*cx8
    sx7_cx8 = sx7/cx8
    cx7_cx8 = cx7/cx8

    x4cx8 = cx8*x₄

    p11 = sx7_cx8*x₁₁
    p12 = cx7_cx8*x₁₂
    xdot9 = p11 + p12

    #differential equations for the quadrotor

    dx[1] = (cx9*x4cx8 + (sx7cx9*sx8 - cx7sx9)*x₅) + (cx7cx9*sx8 + sx7sx9)*x₆
    dx[2] = (sx9*x4cx8 + (sx7sx9*sx8 + cx7cx9)*x₅) + (cx7sx9*sx8 - sx7cx9)*x₆
    dx[3] = (sx8*x₄ - sx7cx8*x₅) - cx7cx8*x₆
    dx[4] = (x₁₂*x₅ - x₁₁*x₆) - g*sx8
    dx[5] = (x₁₀*x₆ - x₁₂*x₄) + g*sx7cx8
    dx[6] = (x₁₁*x₄ - x₁₀*x₅) + (g*cx7cx8 - F_m)
    dx[7] = x₁₀ + sx8*xdot9
    dx[8] = cx7*x₁₁ - sx7*x₁₂
    dx[9] = xdot9
    dx[10] = Cyzx * (x₁₁ * x₁₂) + Tx
    dx[11] = Czxy * (x₁₀ * x₁₂) + Ty
    dx[12] = Cxyz * (x₁₀ * x₁₁) + Tz

    return dx
end

function quadrotor(; T=5.0, plot_vars=[0, 3],
                property=quad_property,
                project_reachset=true,
                Wpos = 0.4, Wvel = 0.4)

    #initial conditions
    X0c = zeros(12)
    ΔX0 = [Wpos, Wpos, Wpos, Wvel, Wvel, Wvel, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    X0 = Hyperrectangle(X0c, ΔX0)

    #initial-value problem
    prob = @ivp(x' = quadrotor!(x), dim: 12, x(0) ∈ X0);

    return prob
end

cases = ["Δ=0.1", "Δ=0.4", "Δ=0.8"];

Wpos = 0.1
Wvel = 0.1
prob = quadrotor(project_reachset=false, Wpos=Wpos, Wvel=Wvel)
alg = TMJets(abstol=1e-7, orderT=5, orderQ=1, adaptive=false)

sol1 = solve(prob, tspan=Tspan, alg=alg);
solz1 = overapproximate(sol1, Zonotope);

Wpos = 0.4
Wvel = 0.4
prob = quadrotor(project_reachset=false, Wpos=Wpos, Wvel=Wvel)
alg = TMJets(abstol=1e-7, orderT=5, orderQ=1, adaptive=false)

sol2 = solve(prob, tspan=Tspan, alg=alg);
solz2 = overapproximate(sol2, Zonotope);

#property = quad_property(solz2)
#println("Validate property, case $(cases[2]) : $(property)")

Wpos = 0.8
Wvel = 0.8
prob = quadrotor(project_reachset=false, Wpos=Wpos, Wvel=Wvel)
alg = TMJets(abstol=1e-7, orderT=5, orderQ=1, adaptive=false)

sol3 = solve(prob, tspan=Tspan, alg=alg);
solz3 = overapproximate(sol3, Zonotope);

using Plots

Plots.plot(solz3,  vars=(0, 3), linecolor="green",  color=:green,  alpha=0.8)
Plots.plot!(solz2, vars=(0, 3), linecolor="blue",   color=:blue,   alpha=0.8)
Plots.plot!(solz1, vars=(0, 3), linecolor="yellow", color=:yellow, alpha=0.8,
    xlab="t", ylab="x3",
    xtick=[0., 1., 2., 3., 4., 5.], ytick=[-1., -0.5, 0., 0.5, 1., 1.5],
    xlims=(0., 5.), ylims=(-1., 1.5))


function learn_rare!(model::Model, prior::GaussianPrior, opt::Opt, data; monitor = nothing)
	λm = [λmin(data.counts, t) for t in 1:model.T-1]
	for t in 1:model.T-1
		@assert model.λ[t] >= λm[t]
	end
	xinit = vcat(model.x, model.λ)
	logP = []
	properties = [[] for a in monitor]
	function f(x::Vector, grad::Vector, vals::Vector, prop::Vector)
		model.x .= x[1:end-model.T+1]
		model.λ .= x[end-model.T+2:end]
		if length(grad) > 0
			grad .= gradient_Likelihood_rare(model, data)
			@inbounds prior(grad[1:end-model.T+1], x[1:end-model.T+1])
		end
		l = Likelihood_rare(model, data) + prior(x[1:end-model.T+1])
		push!(vals, l)
		if monitor != nothing
			for i in 1:length(monitor)
				push!(prop[i], monitor[i](model, prior, data)[:])
			end
		end
		return l
	end

	f1(x::Vector, grad::Vector) = f(x, grad, logP, properties)
	opt.max_objective = f1
	(maxf, xmax, ret) = optimize(opt, xinit)
	@inbounds model.x .= xmax[1:end-model.T+1]
	@inbounds model.λ .= xmax[end-model.T+2:end]
	return (model, ret, logP, properties)
end


module ProjectAssigner

using CSV
using DataFrames
using GLPK
using JuMP

mutable struct Group
    members::Set{String}
    skills::Dict{String,Float64}
    preferences::Vector{String}
end

struct IndexModel
    c::Matrix{Float64}            # c[i, j] is the cost for assigning group i to project j
    r::Matrix{Float64}            # r[k, i] is the amount of skill k that group i contributes
    s::Vector{Float64}            # s[i] is the size of group i
    minskills::Matrix{Float64}    # minskills[k, j] is the minimum number of people needed for skill k on project j
    minsizes::Vector{Float64}     # minimum sizes for each project
    maxsizes::Vector{Float64}     # maximum sizes for each project
    force::Vector{Pair{Int, Int}} # i => j means we are forcing group i to be assigned to project j
end

function IndexModel(groups::Vector{Group}, projects::DataFrame; force=[])
    skills = map(n->n[findfirst(':', n)+1:end], filter(n->startswith(n, "skill:"), names(projects)))
    skillinds = Dict(n=>k for (k, n) in enumerate(skills))
    pinds = Dict(p=>j for (j, p) in enumerate(projects[!, :name]))

    c = calculate_costs(groups, pinds)

    r = zeros(length(skillinds), length(groups))
    for (i, g) in enumerate(groups)
        for skill in keys(g.skills)
            r[skillinds[skill], i] = g.skills[skill]
        end
    end

    s = [length(g.members) for g in groups]

    minskills = zeros(length(skillinds), length(pinds))
    minsizes = zeros(nrow(projects))
    maxsizes = zeros(nrow(projects))
    for (j, row) in enumerate(Tables.rows(projects))
        minsizes[j] = row[:min]
        maxsizes[j] = row[:max]
        for skill in keys(skillinds)
            minskills[skillinds[skill], j] = row["skill:"*skill]
        end
    end

    # map from student name to group index
    ginds = Dict(Iterators.flatten((m=>i for m in groups[i].members) for i in 1:length(groups)))
    force_indices = map(force) do pair
        ginds[first(pair)] => pinds[last(pair)]
    end

    return IndexModel(c, r, s, minskills, minsizes, maxsizes, force_indices)
end

"""
    ProjectAssigner.match(;students, projects, [output], [force], [optimizer])

Match students to projects and return a dataframe with the assignments.

All arguments are keyword arguments, and the function is not exported, so it must be called with the module explicitly `ProjectAssigner.match(...)`.

# Example
```julia-repl
julia> ProjectAssigner.match(students="students.csv", projects="projects.csv")
```

# Keyword Arguments

- `students`: A `DataFrame` or csv filename `String` specifying information about the students. See the Data Format section of the documentation for more information.
- `projects`: A `DataFrame` or csv filename `String` specifying information about the projects. See the Data Format section of the documentation for more information.
- `output::String`: optional name for a csv file that the assignments will be written to.
- `force`: optional vector of student-project pairs to force matches for, e.g. `force=["Bob"=>"Airplane"]` will force student Bob to be assigned to the Airplane project.
- `optimizer::Function`: optimizer factory to use for the optimization, e.g. `optimizer=Gurobi.Optimizer` will use Gurobi instead of the default `GLPK`

!!! warning "Scaling Issues"
    Due to the exponential costs in the objective, the optimization problem may be poorly-scaled for large classes. This can result in the default GLPK optimizer producing poor results even though it thinks that it has solved the problem. More powerful commercial solvers such as [Gurobi](https://github.com/jump-dev/Gurobi.jl) can handle this much better.
"""
function match(;students, projects,
                output=nothing,
                force::AbstractVector{Pair{String,String}}=Pair{String,String}[],
                optimizer=GLPK.Optimizer
    )

    students = as_dataframe(students)
    projects = as_dataframe(projects)

    lint_input(students, projects)

    @info("Found $(nrow(students)) students and $(nrow(projects)) projects.", students, projects)

    groups = group(students)

    @debug(groups)

    assignments = match_groups(groups, projects, force=force, optimizer=optimizer)

    out = DataFrame(name=String[], project=String[], preference=Union{Int,Missing}[])
    for (i, g) in enumerate(groups)
        project = assignments[i]
        pref = something(findfirst(==(project), g.preferences), missing)
        for name in g.members
            push!(out, (name=name, project=assignments[i], preference=pref))
        end
    end
    @info("Assignments:", out)
    histogram = DataFrame(preference=[], number=Int[])
    for r in 1:maximum(skipmissing(out[!, :preference]))
        push!(histogram, (preference=r, number=sum(skipmissing(out[!, :preference].==r))))
    end
    push!(histogram, (preference=missing, number=sum(ismissing, out[!, :preference])))
    @info("Frequencies by preference rank:", histogram)

    if !isnothing(output)
        @info("Writing to $output.")
        CSV.write(output, out)
    end

    return out
end

function match_groups(groups, projects; force=[], optimizer=GLPK.Optimizer)
    m = IndexModel(groups, projects, force=force)

    if (maximum(m.c) + minimum(m.c)) == maximum(m.c) && optimizer == GLPK.Optimizer
        @error("""This problem is too large for GLPK.

               The cost coefficients have too large a range, and the solution will likely be suboptimal.

               Recommend switching to the Gurobi optimizer with
               ProjectAssigner.match(..., optimizer=Gurobi.Optimizer).""", extrema(m.c))
    end

    opt = create_jump_model(m)
    set_optimizer(opt, optimizer)
    optimize!(opt)

    if termination_status(opt) != MOI.OPTIMAL
        @error(termination_status(opt))
    end
    assignments = String[]
    x_opt = value.(opt[:x])
    for i in 1:length(groups)
        j = only(findall(x_opt[i,:] .> 0))
        push!(assignments, projects[j, :name])
    end
    return assignments
end

as_dataframe(fname::AbstractString) = DataFrame(CSV.File(fname))
as_dataframe(df::DataFrame) = df

function group(students::DataFrame)
    skills = String[]
    teammate_keys = String[]
    project_names = String[]
    for n in names(students)
        if startswith(n, "skill:")
            push!(skills, n[findfirst(':', n)+1:end])
        elseif startswith(n, "teammate_")
            push!(teammate_keys, n)
        elseif n != "name"
            push!(project_names, n)
        end
    end

    # extract partners
    partners = Dict{String, Set{String}}()
    for row in Tables.rows(students)
        partners[row[:name]] = Set(skipmissing(row[k] for k in teammate_keys))
    end

    # validate groups
    groups = Group[]
    groupmap = Dict{String, Int}()
    for (k,v) in partners
        if k in keys(groupmap)
            continue
        else
            candidate_group = union!(Set((k,)), v)
            valid = true
            for s in candidate_group
                s_candidate_group = union!(Set((s,)), partners[s])
                if s_candidate_group != candidate_group
                    valid = false
                end
            end
            if valid
                g = Group(candidate_group, Dict(skill=>0.0 for skill in skills), [])
                push!(groups, g)
                for s in candidate_group
                    groupmap[s] = length(groups)
                end
                if length(candidate_group) > 1
                    @info("Creating group $(collect(candidate_group))")
                end
            else
                @warn("""
                      Invalid teammate request for $k
                      """, Dict(s=>collect(partners[s]) for s in candidate_group))
                g = Group(Set((k,)), Dict(skill=>0.0 for skill in skills), [])
                push!(groups, g)
                groupmap[k] = length(groups)
            end
        end
    end

    # fill data in groups
    for row in Tables.rows(students)
        g = groups[groupmap[row[:name]]]
        for skill in skills
            g.skills[skill] += row["skill:"*skill]
        end
        prefs = sort(filter(p->!ismissing(row[p]),project_names), by=p->row[p])
        if !isempty(g.preferences) && prefs != g.preferences
            @warn("Preferences for $(collect(g.members)) don't match.", prefs, g.preferences)
        end
        if length(prefs) > length(g.preferences)
            g.preferences = prefs
        end
    end

    return groups
end

function create_jump_model(m::IndexModel)
    n_groups, n_projects = size(m.c)

    opt = Model()

    # x[i, j] = 1 indicates that group i is in project j
    @variable(opt, x[1:n_groups, 1:n_projects], Bin)

    @objective(opt, Min, sum(m.c.*x))

    for i in 1:n_groups
        @constraint(opt, sum(x[i, :]) == 1)
    end

    for j in 1:n_projects
        # skills
        pr = m.r*x[:, j]
        @constraint(opt, pr .>= m.minskills[:, j])
        # sizes
        ps = m.s'*x[:, j]
        @constraint(opt, m.minsizes[j] <= ps <= m.maxsizes[j])
    end

    for pair in m.force
        @constraint(opt, x[pair...] == 1)
    end

    return opt
end

function calculate_costs(groups, pinds)
    # number of students
    n = sum(length(g.members) for g in groups)

    # number of projects
    m = length(pinds)

    # constant for calculating costs
    logn = ceil(Int, log2(n))

    # c[i, j] is the cost of connecting group i to project j
    c = zeros(length(groups), m)

    for (i, g) in enumerate(groups)
        # fill in preferenced
        for (r, p) in enumerate(g.preferences)
            c[i, pinds[p]] = 2.0^(logn*(r-div(m,2)))*length(g.members)
        end

        # fill in unpreferenced
        if length(g.preferences) < length(pinds)
            r = length(g.preferences)+1
            missing_cost = 2.0^(logn*(r-div(m,2)))*length(g.members)
            row = view(c, i, :)
            row[row.==0.0] .= missing_cost
        end
    end

    return c
end

function lint_input(students, projects)
    # students table
    if !("name" in names(students))
        @error("Required column 'name' missing from students table.", names(students))
    end

    student_skills = Set{String}()
    teammate_keys = String[]
    project_names = Set{String}()
    for n in names(students)
        if startswith(n, "skill:")
            push!(student_skills, n[findfirst(':', n)+1:end])
        elseif startswith(n, "teammate_")
            push!(teammate_keys, n)
        elseif n != "name"
            push!(project_names, n)
        end
    end

    student_names = Set(students[!, :name])
    for r in Tables.rows(students)
        for t in teammate_keys
            if !ismissing(r[t]) && !(r[t] in student_names)
                @error("Teammate requrest not in list of student names.", request=r[t], requester=r[:name])
            end
        end
    end

    # projects table
    for c in ["name", "min", "max"]
        if !(c in names(projects))
            @error("Required column '$c' missing from projects table.", names(projects))
        end
    end

    for c in ["min", "max"]
        if !(eltype(projects[!, c]) <: Integer)
            @error("The type of column '$c' in the projects table was not Integer; instead it was $(eltype(projects[!,c])). Check for extraneous space or other characters.")
        end
    end

    project_skills = Set{String}()
    for n in names(projects)
        if startswith(n, "skill:")
            push!(project_skills, n[findfirst(':', n)+1:end])
        end
    end

    # combination
    if student_skills != project_skills
        @error("Project and student skill columns do not match!", project_skills, student_skills)
    end

    if Set(projects[!,:name]) != project_names
        @error("Project columns in students table and names in project table do not match.", projects[!,:name], project_names)
    end
end

end


function bisect_right(data::Array{Int64}, item::Int64)::Int64
    low = 0
    high = length(data)
    while low < high
        middle = i32::from(((low + high) / 2))
        if item < data[middle]
            high = middle
        else

            low = (middle + 1)
        end
    end
    return low
end

function bin_it(limits::Array{Int64}, data::Array{Int64})::Array{Int64}
    bins = [0]
    for x in limits
        bins.push(0)
    end
    for d in data
        bins[bisect_right(limits, d)] += 1
    end
    return bins
end

function main()
    limits = [23, 37, 43, 53, 67, 83]
    data = [
        95,
        21,
        94,
        12,
        99,
        4,
        70,
        75,
        83,
        93,
        52,
        80,
        57,
        5,
        53,
        86,
        65,
        17,
        92,
        83,
        71,
        61,
        54,
        58,
        47,
        16,
        8,
        9,
        32,
        84,
        7,
        87,
        46,
        19,
        30,
        37,
        96,
        6,
        98,
        40,
        79,
        97,
        45,
        64,
        60,
        29,
        49,
        36,
        43,
        55,
    ]
    @assert(bin_it(limits, data) == [11, 4, 2, 6, 9, 5, 13])
end


using Test: @test, @testset
using LinearAlgebra: norm
using Statistics: mean
using RestrictedBoltzmannMachines: zerosum, zerosum!, zerosum_weights, extensive_sample, free_energy
using RestrictedBoltzmannMachines: RBM, Potts, Binary, Spin, Gaussian, ReLU, dReLU, pReLU, xReLU

@testset "zerosum (visible Potts)" begin
    q = 3
    N = (2, 3)
    M = (2, 3)
    rbm = RBM(Potts(randn(q, N...)), Binary(randn(M...)), randn(q, N..., M...))
    rbm1 = zerosum(rbm)
    @test norm(mean(rbm1.w; dims=1)) < 1e-13
    @test norm(mean(rbm1.visible.θ; dims=1)) < 1e-13
    v = extensive_sample(rbm.visible)
    F0 = free_energy(rbm, v)
    F1 = free_energy(rbm1, v)
    @test all(F0 - F1 .≈ mean(F0 - F1))
    @test zerosum_weights(rbm.w, rbm) ≈ zerosum(rbm).w
    zerosum!(rbm)
    @test free_energy(rbm, v) ≈ F1
end

@testset "zerosum (hidden Potts)" begin
    p = 3
    N = (2, 3)
    M = (2, 3)
    rbm = RBM(Binary(randn(N...)), Potts(randn(p, M...)), randn(N..., p, M...))
    rbm1 = zerosum(rbm)
    @test norm(mean(rbm1.w; dims=3)) < 1e-13
    @test norm(mean(rbm1.hidden.θ; dims=1)) < 1e-13
    v = extensive_sample(rbm.visible)
    F0 = free_energy(rbm, v)
    F1 = free_energy(rbm1, v)
    @test all(F0 - F1 .≈ mean(F0 - F1))
    @test zerosum_weights(rbm.w, rbm) ≈ zerosum(rbm).w
    zerosum!(rbm)
    @test free_energy(rbm, v) ≈ F1
end

@testset "zerosum (visible and hidden Potts)" begin
    q = 3
    p = 3
    N = (2, 3)
    M = (2, 3)
    rbm = RBM(Potts(randn(q, N...)), Potts(randn(p, M...)), randn(q, N..., p, M...))
    rbm1 = zerosum(rbm)
    @test norm(mean(rbm1.w; dims=1)) < 1e-13
    @test norm(mean(rbm1.w; dims=4)) < 1e-13
    @test norm(mean(rbm1.visible.θ; dims=1)) < 1e-13
    @test norm(mean(rbm1.hidden.θ; dims=1)) < 1e-13
    v = extensive_sample(rbm.visible)
    F0 = free_energy(rbm, v)
    F1 = free_energy(rbm1, v)
    @test all(F0 - F1 .≈ mean(F0 - F1))
    @test zerosum_weights(rbm.w, rbm) ≈ zerosum(rbm).w
    zerosum!(rbm)
    @test free_energy(rbm, v) ≈ F1
end

@testset "zerosum (visible Potts, various hidden layers)" begin
    q = 3
    N = (2, 3)
    M = (2, 3)
    hidden_layers = (
        Binary(M...), Spin(M...), Gaussian(M...), ReLU(M...),
        dReLU(M...), pReLU(M...), xReLU(M...)
    )
    for hidden in hidden_layers
        rbm = RBM(Potts(randn(q, N...)), hidden, randn(q, N..., M...))
        rbm1 = zerosum(rbm)
        @test norm(mean(rbm1.w; dims=1)) < 1e-13
        @test norm(mean(rbm1.visible.θ; dims=1)) < 1e-13
        v = extensive_sample(rbm.visible)
        F0 = free_energy(rbm, v)
        F1 = free_energy(rbm1, v)
        @test all(F0 - F1 .≈ mean(F0 - F1))
        zerosum!(rbm)
        @test free_energy(rbm, v) ≈ F1
    end
end


#=
Copyright (c) 2015, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

using DocOpt

include("image.jl")
using .Image

# Solve optical flow problem at one scale
# Flow problem is formulated as the solution of a linear system where matrix is sparse
# Use an iterative method to
# Pre: size(i1)==size(i2)
function singleScaleOpticalFlow(i1::Matrix{Float32}, i2::Matrix{Float32}, lam::Float32, ni::Int)
# Returns Matrix{Float32} size=(size(i1,1),size(i1,2),2)
  (h, w) = size(i1)
  # Compute the gradient of i1
  # This is basically two stencils, one for x and one for y
  # The stencil is [1 8 0 -8 -1]/12 and leave zeros in the boundary
  Ix = zeros(Float32, h, w)
  Iy = zeros(Float32, h, w)
  for y in 1:h
    for x in 1:w
      if x>2 && x<w-1
        Ix[y,x] = (-i1[y,x+2]-8.0f0*i1[y,x+1]+8.0f0*i1[y,x-1]+i1[y,x-2])/12.0f0
      end
      if y>2 && y<h-1
        Iy[y,x] = (-i1[y+2,x]-8.0f0*i1[y+1,x]+8.0f0*i1[y-1,x]+i1[y-2,x])/12.0f0
      end
    end
  end
  # Compute the time partial derivative
  It = i1-i2

  # Block Jacobi preconditioner
  function blockJacobiPreconditioner(r)
    a = Ix.*Ix.+4.0f0*lam
    b = Ix.*Iy
    c = Iy.*Iy.+4.0f0*lam
    invdet = 1.0f0./(a.*c-b.*b)
    res = zeros(Float32, h, w, 2)
    res[:,:,1] = (c.*r[:,:,1]-b.*r[:,:,2]).*invdet
    res[:,:,2] = (a.*r[:,:,2]-b.*r[:,:,1]).*invdet
    return res
  end

  # Matrix vector product
  function mvm(x)
    # This operation is equivalent to a stencil
    # Want:
    #   -lam for (-1,0,0), (+1,0,0), (0,-1,0), and (0,+1,0)
    #   Ix*Iy for (0,0,flip)
    #   Ix^2+4*lam and Iy^2+4*lam for (0,0,0) respectively (Ix for u, Iy for v)
    # At boundaries just include terms that exist
    y = zeros(Float32, h, w, 2)
    for i in 1:h
      for j in 1:w
        ur = Ix[i,j]*(Ix[i,j]*x[i,j,1]+Iy[i,j]*x[i,j,2])+4.0f0*lam*x[i,j,1]
        vr = Iy[i,j]*(Ix[i,j]*x[i,j,1]+Iy[i,j]*x[i,j,2])+4.0f0*lam*x[i,j,2]
        if i>1
          ur += -lam*x[i-1,j,1]
          vr += -lam*x[i-1,j,2]
        end
        if j>1
          ur += -lam*x[i,j-1,1]
          vr += -lam*x[i,j-1,2]
        end
        if i<h
          ur += -lam*x[i+1,j,1]
          vr += -lam*x[i+1,j,2]
        end
        if j<w
          ur += -lam*x[i,j+1,1]
          vr += -lam*x[i,j+1,2]
        end
        y[i,j,1] = ur
        y[i,j,2] = vr
      end
    end
    return y
  end

  # Initial for iterative solve
  x = zeros(Float32, h, w, 2)
  r = zeros(Float32, h, w, 2)
  r[:,:,1] = -Ix.*It
  r[:,:,2] = -Iy.*It
  z = blockJacobiPreconditioner(r)
  p = z
  rsold = sum(r.*z)

  # Iterate to solution
  # Use preconditioned conjugate gradient algorithm
  for i in 1:ni
    Ap = mvm(p)
    pTAp = sum(p.*Ap)
    alpha = rsold/pTAp
    x += alpha*p
    r -= alpha*Ap
    z = blockJacobiPreconditioner(r)
    rsnew = sum(r.*z)
    beta = rsnew/rsold
    p = z + beta*p
    rsold = rsnew
  end

  return x
end

# Solve optical flow problem at all scales
# Pre: size(i1)==size(i2)
function multiScaleOpticalFlow(i1::Matrix{Float32}, i2::Matrix{Float32}, lam::Float32, ni::Int, ns::Int)
# Returns Matrix{Float32} size=(size(i1,1),size(i1,2),2)
  (h, w) = size(i1)
  if ns==1
    return singleScaleOpticalFlow(i1, i2, lam, ni)
  end
  scale = Float32((Float64(w)/50.0)^(-1.0/ns))
  local u # initialised in the first iteration of the loop
  for i in ns:-1:0
    cw = round(Int, w*scale^Float32(i))
    ch = round(Int, h*scale^Float32(i))
    println("Doing scale: ", cw, "x", ch)
    # Downsample the images
    si1 = Image.downSample(i1, cw, ch)
    si2 = Image.downSample(i2, cw, ch)
    # Interpolate the current flow and warp image 2 by current flow
    if i==ns
      # Coarsest scale, no flow yet, use zeros, warp is image 2
      u = zeros(Float32, ch, cw, 2)
      wi = si2
    else
      u = Image.interpolateFlow(u, cw, ch)
      wi = Image.warpMotion(si2, u, si1)
    end
    # Compute the flow from image 1 to warp
    if i==0
      ni *= 5
    end
    du = singleScaleOpticalFlow(si1, wi, lam, ni)
    # Update the flow
    u += du
  end
  return u
end

function main()
    doc = """opt-flow.jl

Horn-Schunck optical flow estimator.

Usage:
  opt-flow.jl -h | --help
  opt-flow.jl [--image-path-1=<image-path-1>] [--image-path-2=<image-path-2>]

Options:
  -h --help                      Show this screen.
  --image-path-1=<image-path-1>  Specify a path to the first image [default: ../../opt-flow/small_001.dat].
  --image-path-2=<image-path-2>  Specify a path to the second image [default: ../../opt-flow/small_002.dat].
"""
    arguments = docopt(doc)

    dir = dirname(@__FILE__)
    img_file_1 = joinpath(dir, arguments["--image-path-1"])
    img_file_2 = joinpath(dir, arguments["--image-path-2"])

    i1 = Image.readImage(img_file_1)
    i2 = Image.readImage(img_file_2)
    if size(i1)!=size(i2)
        write(STDERR, "images are different sizes")
        exit(-1)
    end

    println("filenames = ", img_file_1, " ", img_file_2)
    println("Image size: ", size(i1, 1), "x", size(i1, 2))

    tic()
    multiScaleOpticalFlow(i1, i2, Float32(0.025), 100, 1)
    println("SELFPRIMED ", toq())

    tic()
    u = multiScaleOpticalFlow(i1, i2, Float32(0.025), 100, 44)
    selftimed = toq()
    Image.writeFlo(u, "out.flo")
    println("SELFTIMED ", selftimed)
end

main()


# Simulation using native rk45 implementation
# Tentatively to simulate two rigid bodies constrained by one revolute joint

include("RigidBody.jl")
include("OrientationConversion.jl")
include("Joint.jl")
include("Force.jl")
include("extF.jl")
include("rk45.jl")
include("rk4.jl")

using LinearAlgebra
using Revise
using DifferentialEquations
using StaticArrays

global GravityInInertial = MVector{3}([0.0,0.0,-9.806])
# global extFList = Vector{extForces}(undef,2)

struct solSim
    # Structure to store results of Simulation for each rigid body
    r::Matrix{Float64}
    v::Matrix{Float64}
    β::Matrix{Float64}
    βdot::Matrix{Float64}
end

## RK45
function simulateRK45(tEnd::Float64,tSpan::Float64,j::Joint...;
                  g::MArray{Tuple{3},Float64,1,3}=[0.0;0.0;-9.806])#,extFVec::Vector{extForces}=Vector{extForces}(undef,1))
    # Ellipsis (...) to facilitate supplying variable number of arguments to the function
    # Initial condition (all bodies connected)
    global GravityInInertial = g
    # global extFList = extFVec
    X0 = j[1].RB1.x
    for k = 1:length(j)
        append!(X0,j[k].RB2.x)
    end

    # Using native RK implementation
    tSim, sol = rk45(tEnd,0.0,X0,j,mainDynRK)#,eps = 1e-6)
    nBodies = length(j) + 1
    ntStops = length(tSim)
    sol = reshape(sol,(14*nBodies,ntStops))

    solFinal = Vector{solSim}(undef,length(j))
    for i=1:length(j)
        rSol = Matrix{Float64}(undef,ntStops,3)
        βSol = Matrix{Float64}(undef,ntStops,4)
        vSol = Matrix{Float64}(undef,ntStops,3)
        βdotSol = Matrix{Float64}(undef,ntStops,4)
        for t=1: ntStops
            rSol[t,:] = sol[14*i+1:14*i+3,t]
            βSol[t,:] = sol[14*i+4:14*i+7,t]
            vSol[t,:] = sol[14*i+8:14*i+10,t]
            βdotSol[t,:] = sol[14*i+11:14*i+14,t]
        end
        solFinal[i] = solSim(rSol,vSol,βSol,βdotSol)
    end
    return tSim, solFinal
end

## RK4
function simulateRK4(tEnd::Float64,tSpan::Float64,j::Joint...;
                  g::MArray{Tuple{3},Float64,1,3}=[0.0;0.0;-9.806])#,extFVec::Vector{extForces}=Vector{extForces}(undef,1))
    # Ellipsis (...) to facilitate supplying variable number of arguments to the function
    # Initial condition (all bodies connected)
    global GravityInInertial = g
    # global extFList = extFVec
    X0 = j[1].RB1.x
    for k = 1:length(j)
        append!(X0,j[k].RB2.x)
    end

    # Using native RK implementation
    tSim, sol = rk4(tEnd,0.0,X0,j,mainDynRK, h = tSpan)
    nBodies = length(j) + 1
    ntStops = length(tSim)
    sol = reshape(sol,(14*nBodies,ntStops))

    solFinal = Vector{solSim}(undef,length(j))
    for i=1:length(j)
        rSol = Matrix{Float64}(undef,ntStops,3)
        vSol = Matrix{Float64}(undef,ntStops,3)
        βSol = Matrix{Float64}(undef,ntStops,4)
        βdotSol = Matrix{Float64}(undef,ntStops,4)
        for t=1: ntStops
            rSol[t,:] = sol[14*i+1:14*i+3,t]
            vSol[t,:] = sol[14*i+8:14*i+10,t]
            βSol[t,:] = sol[14*i+4:14*i+7,t]
            βdotSol[t,:] = sol[14*i+11:14*i+14,t]
        end
        solFinal[i] = solSim(rSol,vSol,βSol,βdotSol)
    end
    return tSim, solFinal
end

function mainDynRK(t::Float64, x::Vector{Float64}, p::Tuple{Vararg{Joint}})
    dx = Vector{Float64}(undef,length(x))
    return mainDynODE!(dx,x,p,t)
end

function mainDynODE!(dX::Vector{Float64}, X::Vector{Float64}, j::Tuple{Vararg{Joint}}, t::Float64)
    # ODE function to be used as per DifferentialEquations convention
    # Create extForcesList storing extForces for each rigid body
    # Create ForceConstraints Array storing constraint forces acting on each rigid body
    global GravityInInertial
    # Update RigidBodies
    updateRigidBody!(j[1].RB1,X[1:14])
    for k=1:length(j)
        updateRigidBody!(j[k].RB2,X[14*k+1:14*(k+1)])
    end

    extFListCopy = externalFTotal(t,j) # Generates all the external forces explicitly specified, through joint actions or directly.

    unconstrF, constrF = Constraint(j, extFListCopy, GravityInInertial)
    # unconstrF is the matrix of the unconstrained Forces acting on the robot.
    # constrF is the matrix of the constraint Forces acting on the system developed using UK formulation

    dX[1:14] = zeros(14)
    for k=1:length(j)
        dX[14*k+1:14*(k+1)] = rbDynQuat(j[k].RB2, unconstrF[:,j[k].RB2.bodyID], constrF[:,j[k].RB2.bodyID])
    end
    # mainDyn!(dX, X, j, unconstrF, constrF)
    return dX
end

function externalFTotal(t::Float64, j::Tuple{Vararg{Joint}})
    # Returns externally applied forces in total, not including gravity. Includes Forces from joints and other explicitly specified external forces.
    extFList = extF(t,j)

    # Generate forces from actuated joints on each body
    ForceJoints = Matrix{Float64}(undef,6,2*length(j))
    for k=1:length(j)
        ForceJoints[:,2*k-1], ForceJoints[:,2*k] = genJointF(t,j[k])
    end

    # Add ForceJoints to extFList
    extFListCopy = deepcopy(extFList)
    for k=1:length(j)
        extFListCopy[j[k].RB1.bodyID].Forces =
        vcat(extFListCopy[j[k].RB1.bodyID].Forces, reshape(ForceJoints[1:3,2*k-1],(1,3)))
        extFListCopy[j[k].RB1.bodyID].Positions =
        vcat(extFListCopy[j[k].RB1.bodyID].Positions, reshape(j[k].pos1,(1,3)))
        extFListCopy[j[k].RB1.bodyID].Torques =
        vcat(extFListCopy[j[k].RB1.bodyID].Torques, reshape(ForceJoints[4:6,2*k-1],(1,3)))

        extFListCopy[j[k].RB2.bodyID].Forces =
        vcat(extFListCopy[j[k].RB2.bodyID].Forces, reshape(ForceJoints[1:3,2*k],(1,3)))
        extFListCopy[j[k].RB2.bodyID].Positions =
        vcat(extFListCopy[j[k].RB2.bodyID].Positions, reshape(j[k].pos2,(1,3)))
        extFListCopy[j[k].RB2.bodyID].Torques =
        vcat(extFListCopy[j[k].RB2.bodyID].Torques, reshape(ForceJoints[4:6,2*k],(1,3)))
    end
    return extFListCopy
end


@testset "Convolution methods" begin

    conv_gen(u, v, dims) = real(ifft(fft(u, dims) .* fft(v, dims), dims))

    function conv_test(psf, img, img_out, dims, s)
        otf = fft(psf, dims)
        otf_r = rfft(psf, dims)
        otf_p, conv_p = plan_conv(img, psf, dims, flags=FFTW.ESTIMATE)
        otf_p2, conv_p2 = plan_conv(img .+ 0.0im, 0.0im .+ psf, dims)
        otf_p3, conv_p3 = plan_conv_psf(img, fftshift(psf,dims), dims)
        otf_p3, conv_p3 = plan_conv_psf(img, fftshift(psf,dims), dims, flags=FFTW.MEASURE)
        otf_p4, conv_p4 = plan_conv_psf_buffer(img, fftshift(psf,dims), dims, flags=FFTW.MEASURE)
        @testset "$s" begin
            @test img_out ≈ conv(0.0im .+ img, psf, dims)
            @test img_out ≈ conv(img, psf, dims)
            @test img_out ≈ conv_p(img, otf_p)
            @test img_out ≈ conv_p(img)
            @test img_out ≈ conv_p2(img, otf_p2)
            @test img_out ≈ conv_p2(img)
            @test img_out ≈ conv_psf(img, fftshift(psf, dims), dims)
            @test img_out ≈ conv_p3(img)
            @test img_out ≈ conv_p4(img)
        end
    end

    N = 5
    psf = zeros((N, N))
    psf[1, 1] = 1
    img = randn((N, N))
    conv_test(psf, img, img, [1,2], "Convolution random image with delta peak")

    N = 5
    psf = zeros((N, N))
    psf[1, 1] = 1
    img = randn((N, N, N))
    conv_test(psf, img, img, [1,2], "Convolution with different dimensions psf, img delta")

    N = 5
    psf = abs.(randn((N, N, 2)))
    img = randn((N, N, 2))
    dims = [1, 2]
    img_out = conv_gen(img, psf, dims)
    conv_test(psf, img, img_out, dims, "Convolution with random 3D PSF and random 3D image over 2D dimensions")

    N = 5
    psf = abs.(randn((N, N, N, N, N)))
    img = randn((N, N, N, N, N))
    dims = [1, 2, 3, 4]
    img_out = conv_gen(img, psf, dims)
    conv_test(psf, img, img_out, dims, "Convolution with random 5D PSF and random 5D image over 4 Dimensions")

    N = 5
    psf = abs.(zeros((N, N, N, N, N)))
    for i = 1:N
        psf[1,1,1,1, i] = 1
    end
    img = randn((N, N, N, N, N))
    dims = [1, 2, 3, 4]
    img_out = conv_gen(img, psf, dims)
    conv_test(psf, img, img, dims, "Convolution with 5D delta peak and random 5D image over 4 Dimensions")

    @testset "Check broadcasting convolution" begin
        img = randn((5,6,7))
        psf = randn((5,6,7, 2, 3))
        _, p = plan_conv_buffer(img, psf)
        @test conv(img, psf) ≈ p(img)
    end

   @testset "Check types" begin
        N = 10
        img = randn(Float32, (N, N))
        psf = abs.(randn(Float32, (N, N)))
        dims = [1, 2]
        @test typeof(conv_gen(img, psf, dims)) == typeof(conv(img, psf))
        @test typeof(conv_gen(img, psf, dims)) != typeof(conv(img .+ 0f0im, psf))
        @test conv_gen(img, psf, dims) .+ 1f0im ≈ 1f0im .+ conv(img .+ 0f0im, psf)
    end

    @testset "Check type get_plan" begin
        @test plan_rfft === FourierTools.get_plan(typeof(1f0))
        @test plan_fft === FourierTools.get_plan(typeof(1im))
    end

    @testset "dims argument nothing" begin
        N = 5
        psf = abs.(randn((N, N, N, N, N)))
        img = randn((N, N, N, N, N))
        dims = [1,2,3,4,5]
        @test conv(psf, img) ≈ conv(img, psf, dims)
        @test conv(psf, img) ≈ conv(psf, img, dims)
        @test conv(img, psf) ≈ conv(img, psf, dims)
    end

    @testset "adjoint convolution" begin
        x = randn(ComplexF32, (5,6))
        y = randn(ComplexF32, (5,6))

        y_ft, p = plan_conv(x, y)
        @test ≈(exp(1im * 1.23) .+ conv(ones(eltype(y), size(x)), conj.(y)), exp(1im * 1.23) .+ Zygote.gradient(x -> sum(real(conv(x, y))), x)[1], rtol=1e-4)
        @test ≈(exp(1im * 1.23) .+ conv(ones(ComplexF32, size(x)), conj.(y)), exp(1im * 1.23) .+ Zygote.gradient(x -> sum(real(p(x, y_ft))), x)[1], rtol=1e-4)
    end

end


function dense_w(in_size, out_size, m)
    if m == "Xavier"
        return randn((out_size, in_size)) ./ sqrt(in_size)
    elseif m == "He"
        return randn((out_size, in_size)) ./ sqrt(in_size) .* sqrt(2)
    else
        throw(ArgumentError("only'Xivier'or'He' can be specified as the initial value setting method."))
    end
end

function conv_w(K, io, m)
    if m == "Xavier"
        return randn(K..., io...) ./ sqrt(io[1])
    elseif m == "He"
        return randn(K..., io...) ./ sqrt(io[1]) .* sqrt(2)
    else
        throw(ArgumentError("only'Xivier'or'He' can be specified as the initial value setting method."))
    end
end

create_bias(w, dims...) = fill!(similar(w, dims...), 0)

convert(N, i::Integer) = ntuple(_->i, N)
convert(N, i::Tuple) = i
function convert(N, i::String)
    if i != "same"
        throw(ArgumentError("i is Int type, Tuple or 'same'"))
    end
    return i
end

mutable struct IOShape
    IShape::Tuple
    OShape::Tuple
    function IOShape()
        new(tuple(), tuple())
    end
end

function (IOS::IOShape)(Ih, Iw, Oh, Ow)
    IOS.IShape = (Ih, Iw)
    IOS.OShape = (Oh, Ow)
end

struct Params
    ps::Dict
    l::Int
    function Params(model)
        ps = Dict()
        layers = model.net
        for i in 1 : length(layers)
            l = length(ps)
            if typeof(layers[i]) <: Param
                ps[l+1], ps[l+2] = layers[i].w, layers[i].b
            else
                ps[l+1], ps[l+2] = nothing, nothing
            end
        end
        new(ps, length(ps))
    end
end

function Base.iterate(P::Params, state = 1)
    if state > P.l
        return nothing
    else
        (P.ps[P.l-state+1], state + 1)
    end
end


# A truncated version of Stiefel’s Conjugate Residual method
# cr(A, b, M, atol, rtol, γ, itmax, radius, verbose, linesearch) solves the linear system 'A * x = b' or the least-squares problem :
# 'min ‖b - A * x‖²' within a region of fixed radius.
#
# Marie-Ange Dahito, <marie-ange.dahito@polymtl.ca>
# Montreal, QC, June 2017

export cr

"""A truncated version of Stiefel’s Conjugate Residual method to solve the symmetric linear system Ax=b.
The matrix A must be positive semi-definite.

A preconditioner M may be provided in the form of a linear operator and is
assumed to be symmetric and positive definite.
In a linesearch context, 'linesearch' must be set to 'true'.
"""
function cr(A :: AbstractLinearOperator, b :: AbstractVector{T};
            M :: AbstractLinearOperator=opEye(size(A,1)), atol :: Float64=1.0e-8,
            rtol :: Float64=1.0e-6, γ :: Float64=1.0e-6, itmax :: Int=0,
            radius :: Float64=0.0, verbose :: Bool=false, linesearch :: Bool=false) where T <: Number

  if linesearch && (radius > 0.0)
    error("'linesearch' set to 'true' but radius > 0")
  end
  n = size(b, 1) # size of the problem
  (size(A, 1) == n & size(A, 2) == n) || error("Inconsistent problem size")
  verbose && @printf("CR: system of %d equations in %d variables\n", n, n)

  # Initial state.
  x = zeros(T, n) # initial estimation x = 0
  xNorm = 0.0
  r = M * b # initial residual r = M * (b - Ax) = M * b
  Ar = A * r
  ρ = @kdot(n, r, Ar)
  ρ == 0.0 && return (x, SimpleStats(true, false, [0.0], [], "x = 0 is a zero-residual solution"))
  p = copy(r)
  q = copy(Ar)
  if verbose
    m = 0.0 # quadratic model
  end

  iter = 0
  itmax == 0 && (itmax = 2 * n)

  rNorm = @knrm2(n, r) # ‖r‖
  rNorms = [rNorm] # Values of ‖r‖
  rNorm² = rNorm * rNorm
  pNorm = rNorm
  pNorm² = rNorm²
  pr = rNorm²
  abspr = pr
  pAp = ρ
  abspAp = abs(pAp)
  ArNorm = @knrm2(n, Ar) # ‖Ar‖
  ArNorms = [ArNorm]
  ε = atol + rtol * rNorm
  verbose && @printf("%5s %8s %8s %8s\n", "Iter", "‖x‖", "‖r‖", "quad")
  verbose && @printf("    %d  %8.1e %8.1e %8.1e\n", iter, xNorm, rNorm, m)

  descent = pr > 0.0 # pᵀr > 0 means p is a descent direction
  solved = rNorm ≤ ε
  tired = iter ≥ itmax
  on_boundary = false
  npcurv = false
  status = "unknown"

  while ! (solved || tired)
    if linesearch
      if (pAp ≤ γ * pNorm²) || (ρ ≤ γ * rNorm²)
        npcurv = true
        verbose && @printf("nonpositive curvature detected: pᵀAp = %8.1e and rᵀAr = %8.1e\n", pAp, ρ)
        stats = SimpleStats(solved, false, rNorms, ArNorms, "nonpositive curvature")
        iter == 0 && return (b, stats)
        return (x, stats)
      end
    elseif pAp ≤ 0.0 && radius == 0.0
      error("Indefinite system and no trust region")
    end
    Mq = M * q

    if radius > 0.0
      verbose && @printf("radius = %8.1e > 0 and ‖x‖ = %8.1e\n", radius, xNorm)
      # find t1 > 0 and t2 < 0 such that ‖x + ti * p‖² = radius²  (i = 1, 2)
      xNorm² = xNorm * xNorm
      t = to_boundary(x, p, radius; flip = false, xNorm2 = xNorm², dNorm2 = pNorm²)
      t1 = maximum(t) # > 0
      t2 = minimum(t) # < 0
      tr = maximum(to_boundary(x, r, radius; flip = false, xNorm2 = xNorm², dNorm2 = rNorm²))
      verbose && @printf("t1 = %8.1e, t2 = %8.1e and tr = %8.1e\n", t1, t2, tr)

      if abspAp ≤ γ * pNorm * @knrm2(n, q) # pᵀAp ≃ 0
        npcurv = true # nonpositive curvature
        verbose && @printf("pᵀAp = %8.1e ≃ 0\n", pAp)
        if abspr ≤ γ * pNorm * rNorm # pᵀr ≃ 0
          verbose && @printf("pᵀr = %8.1e ≃ 0, redefining p := r\n", pr)
          p = r # - ∇q(x)
          q = Ar
          # q(x + αr) = q(x) - α ‖r‖² + ½ α² rᵀAr
          # 1) if rᵀAr > 0, the quadratic decreases from α = 0 to α = ‖r‖² / rᵀAr
          # 2) if rᵀAr ≤ 0, the quadratic decreases to -∞ in the direction r
          if ρ > 0.0 # case 1
            verbose && @printf("quadratic is convex in direction r, curv = %8.1e\n", ρ)
            α = min(tr, rNorm² / ρ)
          else # case 2
            verbose && @printf("r is a direction of nonpositive curvature: %8.1e\n", ρ)
            α = tr
          end
        else
          # q_p = q(x + α_p * p) - q(x) = -α_p * rᵀp + ½ (α_p)² * pᵀAp
          # q_r = q(x + α_r * r) - q(x) = -α_r * ‖r‖² + ½ (α_r)² * rᵀAr
          # Δ = q_p - q_r. If Δ > 0, r is followed, else p is followed
          α = descent ? t1 : t2
          ρ > 0.0 && (tr = min(tr, rNorm² / ρ))
          Δ = -α * pr + tr * rNorm² - 0.5 * (tr)^2 * ρ # as pᵀAp = 0
          if Δ > 0.0 # direction r engenders a better decrease
            verbose && @printf("direction r engenders a bigger decrease. q_p - q_r = %8.1e > 0\n", Δ)
            verbose && @printf("redefining p := r\n")
            p = r
            q = Ar
            α = tr
          else
            verbose && @printf("direction p engenders an equal or a bigger decrease. q_p - q_r = %8.1e ≤ 0\n", Δ)
          end
        end

      elseif pAp > 0.0 && ρ > 0.0 # no negative curvature
        verbose && @printf("positive curvatures along p and r. pᵀAp = %8.1e and rᵀAr = %8.1e\n", pAp, ρ)
        α = ρ / @kdot(n, q, Mq)
        if α ≥ t1
          α = t1
          on_boundary = true
        end

      elseif pAp > 0.0 && ρ < 0.0
        npcurv = true
        verbose && @printf("pᵀAp = %8.1e > 0 and rᵀAr = %8.1e < 0\n", pAp, ρ)
        # q_p is minimal for α_p = rᵀp / pᵀAp
        α = descent ?  min(t1, pr / pAp) : max(t2, pr / pAp)
        Δ = -α * pr + tr * rNorm² + 0.5 * (α^2 * pAp - (tr)^2 * ρ)
        if Δ > 0.0
          verbose && @printf("direction r engenders a bigger decrease. q_p - q_r = %8.1e > 0\n", Δ)
          verbose && @printf("redefining p := r\n")
          p = r
          q = Ar
          α = tr
        else
          verbose && @printf("direction p engenders an equal or a bigger decrease. q_p - q_r = %8.1e ≤ 0\n", Δ)
        end

      elseif pAp < 0.0 && ρ > 0.0
        npcurv = true
        verbose && @printf("pᵀAp = %8.1e < 0 and rᵀAr = %8.1e > 0\n", pAp, ρ)
        α = descent ? t1 : t2
        tr = min(tr, rNorm² / ρ)
        Δ = -α * pr + tr * rNorm² + 0.5 * (α^2 * pAp - (tr)^2 * ρ)
        if Δ > 0.0
          verbose && @printf("direction r engenders a bigger decrease. q_p - q_r = %8.1e > 0\n", Δ)
          verbose && @printf("redefining p := r\n")
          p = r
          q = Ar
          α = tr
        else
          verbose && @printf("direction p engenders an equal or a bigger decrease. q_p - q_r = %8.1e ≤ 0\n", Δ)
        end

      elseif pAp < 0.0 && ρ < 0.0
        npcurv = true
        verbose && @printf("negative curvatures along p and r. pᵀAp = %8.1e and rᵀAr = %8.1e\n", pAp, ρ)
        α = descent ? t1 : t2
        Δ = -α * pr + tr * rNorm² + 0.5 * (α^2 * pAp - (tr)^2 * ρ)
        if Δ > 0.0
          verbose && @printf("direction r engenders a bigger decrease. q_p - q_r = %8.1e > 0\n", Δ)
          verbose && @printf("redefining p := r\n")
          p = r
          q = Ar
          α = tr
        else
          verbose && @printf("direction p engenders an equal or a bigger decrease. q_p - q_r = %8.1e ≤ 0\n", Δ)
        end
      end

    elseif radius == 0.0
      α = ρ / @kdot(n, q, Mq) # step
    end

    @kaxpy!(n,  α, p, x)
    xNorm = @knrm2(n, x)
    xNorm ≈ radius && (on_boundary = true)
    @kaxpy!(n, -α, Mq, r) # residual
    rNorm² = abs(rNorm² - α * ρ)
    rNorm = sqrt(rNorm²)
    push!(rNorms, rNorm)
    Ar = A * r
    ArNorm = @knrm2(n, Ar)
    push!(ArNorms, ArNorm)

    iter = iter + 1
    if verbose
      m = m - α * pr + 0.5 * α^2 * pAp
      @printf("    %d  %8.1e %8.1e %8.1e\n", iter, xNorm, rNorm, m)
    end

    solved = (rNorm <= ε) || npcurv || on_boundary
    tired = iter >= itmax

    (solved || tired) && continue
    ρbar = ρ
    ρ = @kdot(n, r, Ar)
    β = ρ / ρbar # step for the direction computation
    @kaxpby!(n, 1.0, r, β, p)
    @kaxpby!(n, 1.0, Ar, β, q)

    pNorm² = rNorm² + 2 * β * pr - 2 * β * α * pAp + β^2 * pNorm²
    pNorm = sqrt(pNorm²)
    pr = rNorm² + β * pr - β * α * pAp # pᵀr
    abspr = abs(pr)
    pAp = ρ + β^2 * pAp # pᵀq
    abspAp = abs(pAp)
    descent = pr > 0.0

  end
  verbose && @printf("\n")

  status = npcurv ? "nonpositive curvature" : (on_boundary ? "on trust-region boundary" : (tired ? "maximum number of iterations exceeded" : "solution good enough given atol and rtol"))
  stats = SimpleStats(solved, false, rNorms, ArNorms, status)
  return (x, stats)
end


#=
Model:  fiveDimSys.jl

This is a five-dimensional model taken from [Girard05]. This also appears as
Example 4.1. page 57 in the thesis [LeGuernic09].

See also Fig. 5.1. page 71 in the same thesis.

[Girard05]       -- Girard, Antoine. "Reachability of uncertain linear systems using zonotopes."
                    International Workshop on Hybrid Systems: Computation and Control.
                    Springer, Berlin, Heidelberg, 2005.
[LeGuernic09]    -- Le Guernic, Colas. Reachability analysis of hybrid systems
                    with linear continuous dynamics. Diss.
                    Université Joseph-Fourier-Grenoble I, 2009.
=#
using Reachability, Plots

function compute(input_options::Pair{Symbol,<:Any}...)
    # =====================
    # Problem specification
    # =====================
    D = [-1 -4 0 0 0; 4 -1 0 0 0; 0 0 -3 1 0; 0 0 -1 -3 0; 0 0 0 0 -2.]
    P = [0.6 -0.1 0.1 0.7 -0.2; -0.5 0.7 -0.1 -0.8 0; 0.9 -0.5 0.3 -0.6 0.1;
        0.5 -0.7 0.5 0.6 0.3; 0.8 0.7 0.6 -0.3 0.2]
    A = P * D * inv(P)

    # initial set
    X0 = BallInf([1.0, 0.0, 0.0, 0.0, 0.0], 0.1)

    # input set
    U = Ball2(zeros(5), 0.01)

    # instantiate continuous LTI system
    S = ContinuousSystem(A, X0, U) # @system x' = A*x + u, u ∈ U

    # ===============
    # Problem solving
    # ===============
    # define solver-specific options
    options = merge(Options(
        :mode => "reach",
        :blocks => [1],
        :plot_vars => [0, 1]
        ), Options(input_options...))
    result = solve(S, options)

    # ========
    # Plotting
    # ========
    if options[:mode] == "reach"
        println("Plotting...")
        tic()
        plot(result)
        @eval(savefig(@filename_to_png))
        toc()
    end
end # function

compute(:N => 100, :T => 1.0); # warm-up
compute(:δ => 0.01, :T => 5.0); # benchmark settings (long)


using DifferentialEquations
using Plots
plotlyjs()

# include the case files
include("ex2-case-1.jl")
include("ex2-case-2.jl")
include("ex2-case-3.jl")

#initial conditions
r₀=6
v₀=-0.25
u₀=[17*π/18,0.1,0.1,0.1]

#calling the functions to run them
IC=[r₀,v₀,u₀]
IC2=ex2_case_1(IC)
IC3=ex2_case_2(IC2)
FC=ex2_case_3(IC3);

#overall trajectory
t1=IC2[4]
t2=IC3[4]
t3=FC[4]
t=vcat(t1,t2,t3)
display(plot(t[:,1],t[:,2],t[:,3],label="trajectory",
    xaxis="x (m)",yaxis="y (m)",zaxis="z (m)",grid=(:on,:black)))
savefig("./plots/trajectory.png")

Modeling the Dragon Module as a Pendulum/Scripts/ex2.jl


import MitosisStochasticDiffEq as MSDE
using Mitosis
using Test, Random
using LinearAlgebra

# set true/estimated model parameters
# p = [-0.1,0.2,0.9]

"""
    backwardfilter() -> ps, p0, c
Backward filtering using the Euler method, starting with `N(ν, P)` prior
and integration scale c2 between observations

"""
function backwardfilter((c, ν, P)::NamedTuple{(:logscale, :μ, :Σ)}, p, s)
    ps = [[ν, P, c]]
    B, β, σ̃ = p
    for i in eachindex(s)[end-1:-1:1]
        dt = s[i+1] - s[i]
        H = inv(P)
        F = H*ν

        P = P - dt*(B*P + P*B' - σ̃*σ̃'*I)
        ν = ν - dt*(B*ν .+ β)
        c = c - dt*tr(B)

        push!(ps, [ν, P, c])
    end
    ps, [ν, P, c]
end

@testset "backward filtering tests" begin

  # define SDE function
  f(u,p,t) = p[1]*u .+ p[2]
  g(u,p,t) = p[3]

  # time span
  tstart = 0.0
  tend = 1.0
  dt = 0.02
  trange = tstart:dt:tend

  # set of linear parameters Eq.~(2.2)
  B, β, σ̃ = -0.1, 0.2, 1.3
  plin = [B, β, σ̃]
  kernel = MSDE.SDEKernel(f,g,trange,plin)

  # initial values for ODE
  mynames = (:logscale, :μ, :Σ);
  myvalues = [0.0, 0.0, 10.0];
  NT = NamedTuple{mynames}(myvalues)

  message, solend = MSDE.backwardfilter(kernel, NT)

  kernel2 = MSDE.SDEKernel(Mitosis.AffineMap(B, β), Mitosis.ConstantMap(σ̃), trange, plin)
  message2, solend2 = MSDE.backwardfilter(kernel2, NT)

  @test isapprox(solend, solend2, rtol=1e-15)
  @test isapprox(Array(message.sol), Array(message2.sol), rtol=1e-15)

  message2, solend2 = backwardfilter(NT, plin, message.ts)

  @test isapprox(solend, solend2, rtol=1e-15)
  @test isapprox(Array(message.sol), reduce(hcat, message2), rtol=1e-15)

  # multivariate tests
  dim = 5
  Random.seed!(123)
  logscale = randn()
  μ = randn(dim)
  Σ = randn(dim,dim)
  myvalues = [logscale, μ, Σ];
  NT = NamedTuple{mynames}(myvalues)

  kernel = MSDE.SDEKernel(f,g,trange,plin)
  message, solend = MSDE.backwardfilter(kernel, NT)
  message2, solend2 = backwardfilter(NT, plin, message.ts)

  @test isapprox(solend.x[1], solend2[1], rtol=1e-15)
  @test isapprox(solend.x[2], solend2[2], rtol=1e-15)
  @test isapprox(solend.x[3][1], solend2[3], rtol=1e-15)

  kernel2 = MSDE.SDEKernel(Mitosis.AffineMap(B, β), Mitosis.ConstantMap(σ̃), trange, plin)
  message2, solend2 = MSDE.backwardfilter(kernel2, NT)
  @test isapprox(solend, solend2, rtol=1e-15)
  @test isapprox(Array(message.sol), Array(message2.sol), rtol=1e-15)

  # test inplace version
  message2, solend2 = MSDE.backwardfilter(kernel, NT, inplace=true)
  @test isapprox(solend, solend2, rtol=1e-15)
  @test isapprox(Array(message.sol), Array(message2.sol), rtol=1e-15)

  message2, solend2 = MSDE.backwardfilter(kernel2, NT, inplace=true)
  @test isapprox(solend, solend2, rtol=1e-15)
  @test isapprox(Array(message.sol), Array(message2.sol), rtol=1e-15)

  m = 3 # some number of Brownian processes
  plin = [randn(dim,dim), randn(dim), randn(dim,m)] # B, β, σtil

  kernel = MSDE.SDEKernel(f,g,trange,plin)
  message, solend = MSDE.backwardfilter(kernel, NT)
  message2, solend2 = backwardfilter(NT, plin, message.ts)

  @test isapprox(solend.x[1], solend2[1], rtol=1e-14)
  @test isapprox(solend.x[2], solend2[2], rtol=1e-14)
  @test isapprox(solend.x[3][1], solend2[3], rtol=1e-14)

  kernel2 = MSDE.SDEKernel(Mitosis.AffineMap(B, β), Mitosis.ConstantMap(σ̃), trange, plin)
  message2, solend2 = MSDE.backwardfilter(kernel2, NT)
  @test isapprox(solend, solend2, rtol=1e-15)
  @test isapprox(Array(message.sol), Array(message2.sol), rtol=1e-15)

  # test inplace version
  message2, solend2 = MSDE.backwardfilter(kernel, NT, inplace=true)
  @test isapprox(solend, solend2, rtol=1e-15)
  @test isapprox(Array(message.sol), Array(message2.sol), rtol=1e-15)

  # test symmetric matrix
  plin = [Symmetric(randn(dim,dim)), randn(dim), randn(dim,m)] # B, β, σtil
  kernel = MSDE.SDEKernel(f,g,trange,plin)
  kernel2 = MSDE.SDEKernel(Mitosis.AffineMap(B, β), Mitosis.ConstantMap(σ̃), trange, plin)
  message, solend  = MSDE.backwardfilter(kernel, NT)
  message2, solend2 = MSDE.backwardfilter(kernel2, NT)
  @test isapprox(solend, solend2, rtol=1e-15)
  @test isapprox(Array(message.sol), Array(message2.sol), rtol=1e-15)

  plin = [Array(plin[1]), plin[2], plin[3]] # B, β, σtil
  kernel = MSDE.SDEKernel(f,g,trange,plin)
  message2, solend2  = MSDE.backwardfilter(kernel, NT)

  @test isapprox(solend, solend2, rtol=1e-15)
  @test isapprox(Array(message.sol), Array(message2.sol), rtol=1e-15)
end

@testset "backward filtering timechange tests" begin
  # define SDE function
  f(u,p,t) = p[1]*u .+ p[2]
  g(u,p,t) = p[3]

  # time span
  tstart = 0.0
  tend = 1.0
  dt = 0.02
  trange = tstart:dt:tend

  # set of linear parameters Eq.~(2.2)
  B, β, σ̃ = -0.1, 0.2, 1.3
  plin = [B, β, σ̃]

  # initial values for ODE
  mynames = (:logscale, :μ, :Σ);
  myvalues = [0.0, 0.0, 10.0];
  NT = NamedTuple{mynames}(myvalues)

  kernel = MSDE.SDEKernel(Mitosis.AffineMap(B, β), Mitosis.ConstantMap(σ̃), trange, plin)
  message, solend = MSDE.backwardfilter(kernel, NT, apply_timechange=true)

  @test length(message.ts) == length(trange)
  @test message.ts != trange
  @test message.ts == MSDE.timechange(trange)

  message2, solend2 = backwardfilter(NT, plin, message.ts)

  @test isapprox(solend, solend2, rtol=1e-15)
  @test isapprox(Array(message.sol), reduce(hcat, message2), rtol=1e-15)

end

@testset "backward filtering adaptive tests" begin
  # define SDE function
  using OrdinaryDiffEq

  f(u,p,t) = p[1]*u .+ p[2]
  g(u,p,t) = p[3]

  # time span
  tstart = 0.0
  tend = 1.0
  dt = 0.05 # more coarse grained dt
  trange = tstart:dt:tend

  # set of linear parameters Eq.~(2.2)
  B, β, σ̃ = -0.1, 0.2, 1.3
  plin = [B, β, σ̃]

  # initial values for ODE
  mynames = (:logscale, :μ, :Σ);
  myvalues = [0.0, 0.0, 10.0];
  NT = NamedTuple{mynames}(myvalues)

  tildekernel = MSDE.SDEKernel(Mitosis.AffineMap(B, β), Mitosis.ConstantMap(σ̃), trange, plin)
  message, solend = MSDE.backwardfilter(tildekernel, NT, alg=Tsit5(), apply_timechange=true)

  # adaptive solution with large tolerances follows tstops
  @test length(message.ts) == length(trange)
  @test message.ts != trange
  @test message.ts == MSDE.timechange(trange)

  message2, solend2 = MSDE.backwardfilter(tildekernel, NT, alg=Tsit5(),
    abstol=1e-12, reltol=1e-12, apply_timechange=true)

  # adaptive solution with small tolerances steps to tstops but also takes additional substeps
  @test length(message.ts) != length(message2.ts)
  @test minimum(message.ts .∈ [message2.ts])
end


####
#### Robust online (co-)variance estimators.
####

abstract type VarEstimator{T} end

# NOTE: this naive variance estimator is used only in testing
mutable struct NaiveVar{T} <: VarEstimator{T}
    n :: Int
    S :: Vector{Vector{T}}
end

NaiveVar(::Type{T}=Float64) where {T} = NaiveVar(0, Vector{Vector{T}}())
NaiveVar(::Type{T}, ::Int) where {T} = NaiveVar(T)
NaiveVar(::Int) = NaiveVar(Float64)

function add_sample!(nc::NaiveVar, s::AbstractVector)
    nc.n += 1
    push!(nc.S, s)
end

function reset!(nc::NaiveVar{T}) where {T<:AbstractFloat}
    nc.n = 0
    nc.S = Vector{Vector{T}}()
end

function get_var(nc::NaiveVar{T}) where {T<:AbstractFloat}
    @assert nc.n >= 2 "Cannot get variance with only one sample"
    return Statistics.var(nc.S)
end

# Ref： https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_var_estimator.hpp
mutable struct WelfordVar{T<:AbstractFloat, AT<:AbstractVector{T}} <: VarEstimator{T}
    n :: Int
    μ :: AT
    M :: AT
end

WelfordVar(::Type{T}, d::Int) where {T} = WelfordVar(0, zeros(T, d), zeros(T, d))
WelfordVar(d::Int) = WelfordVar(Float64, d)

function reset!(wv::WelfordVar{T, AT}) where {T<:AbstractFloat, AT<:AbstractVector{T}}
    wv.n = 0
    wv.μ .= zero(T)
    wv.M .= zero(T)
end

function add_sample!(wv::WelfordVar, s::AbstractVector)
    wv.n += 1
    @unpack μ, M, n = wv
    for i in eachindex(s)
        δ = s[i] - μ[i]
        μ[i] += δ / n
        M[i] += δ * (s[i] - μ[i])
    end
end

# https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/var_adaptation.hpp
function get_var(wv::WelfordVar{T, AT}) where {T<:AbstractFloat, AT<:AbstractVector{T}}
    n, M = T(wv.n), wv.M
    @assert n >= 2 "Cannot get covariance with only one sample"
    return (n / ((n + 5) * (n - 1))) .* M .+ T(1e-3) * (5 / (n + 5))
end

abstract type CovEstimator{T} end

# NOTE: This naive covariance estimator is used only in testing.
mutable struct NaiveCov{T} <: CovEstimator{T}
    n :: Int
    S :: Vector{Vector{T}}
end
NaiveCov(::Type{T}=Float64) where {T} = NaiveCov(0, Vector{Vector{T}}())

function add_sample!(nc::NaiveCov, s::AbstractVector)
    nc.n += 1
    push!(nc.S, s)
end

function get_cov(nc::NaiveCov{T})::Matrix{T} where {T<:AbstractFloat}
    @assert nc.n >= 2 "Cannot get covariance with only one sample"
    return Statistics.cov(nc.S)
end

function reset!(nc::NaiveCov{T}) where {T<:AbstractFloat}
    nc.n = 0
    nc.S = Vector{Vector{T}}()
end

# Ref: https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_covar_estimator.hpp
mutable struct WelfordCov{T<:AbstractFloat} <: CovEstimator{T}
    n :: Int
    μ :: Vector{T}
    M :: Matrix{T}
    δ :: Vector{T} # temporary
end

function WelfordCov(::Type{T}, d::Int) where {T}
    return WelfordCov(0, zeros(T, d), zeros(T, d, d), zeros(T, d))
end
WelfordCov(d::Int) = WelfordCov(Float64, d)

function reset!(wc::WelfordCov{T}) where {T<:AbstractFloat}
    wc.n = 0
    wc.μ .= zero(T)
    wc.M .= zero(T)
end

function add_sample!(wc::WelfordCov, s::AbstractVector)
    wc.n += 1
    @unpack δ, μ, n, M = wc
    δ .= s .- μ
    μ .+= δ ./ n
    M .+= (s .- μ) .* δ'
end

# Ref: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/covar_adaptation.hpp
function get_cov(wc::WelfordCov{T}) where {T<:AbstractFloat}
    n, M = T(wc.n), wc.M
    @assert n >= 2 "Cannot get variance with only one sample"
    return (n / ((n + 5) * (n - 1))) .* M + T(1e-3) * (5 / (n + 5)) * LinearAlgebra.I
end

###
### Preconditioning matrix adaptors.
###

abstract type AbstractPreconditioner <: AbstractAdaptor end

finalize!(adaptor::T) where {T<:AbstractPreconditioner} = nothing

####
#### Preconditioning matrix adaption implementation.
####

# Unit
struct UnitPreconditioner{T} <: AbstractPreconditioner end
Base.show(io::IO, ::UnitPreconditioner) = print(io, "UnitPreconditioner")

UnitPreconditioner(::Type{T}=Float64) where {T} = UnitPreconditioner{T}()

string(::UnitPreconditioner) = "I"
reset!(::UnitPreconditioner) = nothing
getM⁻¹(dpc::UnitPreconditioner{T}) where {T<:AbstractFloat} = UniformScaling{T}(one(T))
adapt!(
    ::UnitPreconditioner,
    ::AbstractVector{<:Real},
    ::AbstractFloat,
    is_update::Bool=true
) = nothing

mutable struct DiagPreconditioner{
    T<:Real,
    AT<:AbstractVector{T},
    TEst<:VarEstimator{T}
} <: AbstractPreconditioner
    n_min   :: Int
    ve      :: TEst
    var     :: AT
end
Base.show(io::IO, ::DiagPreconditioner) = print(io, "DiagPreconditioner")

# Diagonal
DiagPreconditioner(d::Int, n_min::Int=10) = DiagPreconditioner(Float64, d, n_min)
function DiagPreconditioner(
    ::Type{T},
    d::Int;
    n_min::Int=10,
    var=Vector(ones(T, d))
) where {T}
    ve = WelfordVar(T, d)
    return DiagPreconditioner(n_min, ve, var)
end

string(dpc::DiagPreconditioner) = string(dpc.var)
reset!(dpc::DiagPreconditioner) = reset!(dpc.ve)
getM⁻¹(dpc::DiagPreconditioner) = dpc.var

function adapt!(
    dpc::DiagPreconditioner,
    θ::AbstractVector{T},
    α::AbstractFloat,
    is_update::Bool=true
) where {T<:Real}
    resize!(dpc, θ)
    add_sample!(dpc.ve, θ)
    if dpc.ve.n >= dpc.n_min && is_update
        # Can be made inplace
        dpc.var .= get_var(dpc.ve)
    end
end

# Dense
mutable struct DensePreconditioner{
    T<:AbstractFloat,
    TEst<:CovEstimator{T}
} <: AbstractPreconditioner
    n_min :: Int
    ce    :: TEst
    covar :: Matrix{T}
end
Base.show(io::IO, ::DensePreconditioner) = print(io, "DensePreconditioner")

DensePreconditioner(d::Int, n_min::Int=10) = DensePreconditioner(Float64, d, n_min)
function DensePreconditioner(
    ::Type{T},
    d::Int;
    n_min::Int=10,
    covar=LinearAlgebra.diagm(0 => ones(T, d))
) where {T}
    ce = WelfordCov(T, d)
    # TODO: take use of the line below when we have an interface to set which covariance estimator to use
    # ce = NaiveCov(T)
    return DensePreconditioner(n_min, ce, covar)
end

string(dpc::DensePreconditioner) = string(LinearAlgebra.diag(dpc.covar))
reset!(dpc::DensePreconditioner) = reset!(dpc.ce)
getM⁻¹(dpc::DensePreconditioner) = dpc.covar

function adapt!(
    dpc::DensePreconditioner,
    θ::AbstractVector{T},
    α::AbstractFloat,
    is_update::Bool=true
) where {T<:AbstractFloat}
    resize!(dpc, θ)
    add_sample!(dpc.ce, θ)
    if dpc.ce.n >= dpc.n_min && is_update
        # Can be made inplace
        dpc.covar .= get_cov(dpc.ce)
    end
end

# Resize pre-conditioner if necessary.
Base.resize!(
    pc::UnitPreconditioner,
    θ::AbstractVector{T}
) where {T<:Real} = nothing
function Base.resize!(
    dpc::DiagPreconditioner,
    θ::AbstractVector{T}
) where {T<:Real}
    if length(θ) != length(dpc.var)
        @assert dpc.ve.n == 0 "Cannot resize a var estimator when it contains samples."
        dpc.ve = WelfordVar(T, length(θ))
        dpc.var = ones(T, length(θ))
    end
end
function Base.resize!(
    dpc::DensePreconditioner,
    θ::AbstractVector{T}
) where {T<:Real}
    if length(θ) != size(dpc.covar,1)
        @assert dpc.ce.n == 0 "Cannot resize a var estimator when it contains samples."
        dpc.ce = WelfordCov(T, length(θ))
        dpc.covar = LinearAlgebra.diagm(0 => ones(T, length(θ)))
    end
end

####
#### Preconditioning mass matrix.
####

abstract type AbstractMetric end

_string_M⁻¹(mat::AbstractMatrix, n_chars::Int=32) = _string_M⁻¹(LinearAlgebra.diag(mat), n_chars)
function _string_M⁻¹(vec::AbstractVector, n_chars::Int=32)
    s_vec = string(vec)
    l = length(s_vec)
    s_dots = " ...]"
    n_diag_chars = n_chars - length(s_dots)
    return s_vec[1:min(n_diag_chars,end)] * (l > n_diag_chars ? s_dots : "")
end

struct UnitEuclideanMetric{T} <: AbstractMetric
    M⁻¹::UniformScaling{T}
    dim::Int
end

UnitEuclideanMetric(::Type{T}, dim::Int) where {T} = UnitEuclideanMetric{T}(UniformScaling{T}(one(T)), dim)
UnitEuclideanMetric(dim::Int) = UnitEuclideanMetric(Float64, dim)

renew(ue::UnitEuclideanMetric, M⁻¹) = UnitEuclideanMetric(M⁻¹, ue.dim)

Base.length(e::UnitEuclideanMetric) = e.dim
Base.show(io::IO, uem::UnitEuclideanMetric) = print(io, "UnitEuclideanMetric($(_string_M⁻¹(ones(uem.dim))))")

struct DiagEuclideanMetric{T,A<:AbstractVector{T}} <: AbstractMetric
    # Diagnal of the inverse of the mass matrix
    M⁻¹     ::  A
    # Sqare root of the inverse of the mass matrix
    sqrtM⁻¹ ::  A
    # Pre-allocation for intermediate variables
    _temp   ::  A
end

function DiagEuclideanMetric(M⁻¹::AbstractVector{T}) where {T<:Real}
    return DiagEuclideanMetric(M⁻¹, sqrt.(M⁻¹), Vector{T}(undef, size(M⁻¹, 1)))
end
DiagEuclideanMetric(::Type{T}, D::Int) where {T} = DiagEuclideanMetric(ones(T, D))
DiagEuclideanMetric(D::Int) = DiagEuclideanMetric(Float64, D)

renew(ue::DiagEuclideanMetric, M⁻¹) = DiagEuclideanMetric(M⁻¹)

Base.length(e::DiagEuclideanMetric) = size(e.M⁻¹, 1)
Base.show(io::IO, dem::DiagEuclideanMetric) = print(io, "DiagEuclideanMetric($(_string_M⁻¹(dem.M⁻¹)))")

function Base.getproperty(dem::DiagEuclideanMetric, d::Symbol)
    return d === :dim ? size(getfield(dem, :M⁻¹), 1) : getfield(dem, d)
end

struct DenseEuclideanMetric{
    T,
    AV<:AbstractVector{T},
    AM<:AbstractMatrix{T},
    TcholM⁻¹<:UpperTriangular{T},
} <: AbstractMetric
    # Inverse of the mass matrix
    M⁻¹::AM
    # U of the Cholesky decomposition of the mass matrix
    cholM⁻¹::TcholM⁻¹
    # Pre-allocation for intermediate variables
    _temp::AV
end

function DenseEuclideanMetric(M⁻¹::AbstractMatrix{T}) where {T<:Real}
    _temp = Vector{T}(undef, size(M⁻¹, 1))
    return DenseEuclideanMetric(M⁻¹, cholesky(Symmetric(M⁻¹)).U, _temp)
end
DenseEuclideanMetric(::Type{T}, D::Int) where {T} = DenseEuclideanMetric(Matrix{T}(I, D, D))
DenseEuclideanMetric(D::Int) = DenseEuclideanMetric(Float64, D)

renew(ue::DenseEuclideanMetric, M⁻¹) = DenseEuclideanMetric(M⁻¹)

Base.length(e::DenseEuclideanMetric) = size(e.M⁻¹, 1)
Base.show(io::IO, dem::DenseEuclideanMetric) = print(io, "DenseEuclideanMetric(diag=$(_string_M⁻¹(dem.M⁻¹)))")

function Base.getproperty(dem::DenseEuclideanMetric, d::Symbol)
    return d === :dim ? size(getfield(dem, :M⁻¹), 1) : getfield(dem, d)
end

# `rand` functions for `metric` types.
function Base.rand(
    rng::AbstractRNG,
    metric::UnitEuclideanMetric{T}
) where {T}
    r = randn(rng, T, metric.dim)
    return r
end

function Base.rand(
    rng::AbstractRNG,
    metric::DiagEuclideanMetric{T}
) where {T}
    r = randn(rng, T, metric.dim)
    r ./= metric.sqrtM⁻¹
    return r
end

function Base.rand(
    rng::AbstractRNG,
    metric::DenseEuclideanMetric{T}
) where {T}
    r = randn(rng, T, metric.dim)
    ldiv!(metric.cholM⁻¹, r)
    return r
end

Base.rand(metric::AbstractMetric) = rand(GLOBAL_RNG, metric)

####
#### Preconditioner constructors
####

Preconditioner(m::UnitEuclideanMetric{T}) where {T} = UnitPreconditioner(T)
Preconditioner(m::DiagEuclideanMetric{T}) where {T} = DiagPreconditioner(T, m.dim; var=m.M⁻¹)
Preconditioner(m::DenseEuclideanMetric{T}) where {T} = DensePreconditioner(T, m.dim; covar=m.M⁻¹)

Preconditioner(m::Type{TM}, dim::Int=2) where {TM<:AbstractMetric} = Preconditioner(Float64, m, dim)
Preconditioner(::Type{T}, ::Type{TM}, dim::Int=2) where {T<:AbstractFloat, TM<:AbstractMetric} = Preconditioner(TM(T, dim))


# wrap in a module for easy testing at the REPL
module NG

abstract type AbstractInventory end
abstract type AbstractPriceProcess end
abstract type AbstractMarket end

struct Inventory{T} <: AbstractInventory
    terminal_inventory::T
    min_inventory::T
    max_inventory::T
end

struct PriceProcess{T} <: AbstractPriceProcess
    initial_price::T
    volatility::T
    mean_level::T
    mean_rev_rate::T
end

struct Market{T, U} <: AbstractMarket
    horizon::U
    periods::U
    runs::U
    risk_free_rate::T
end

#function main()
    Δt = 3 * 1 / 1000
    market = Market(3,1000,10000,10/100)
    inventory = vcat(0:25:100,150,200:100:900,950:50:1100,1200:100:1800,1850,1900:25:2000)
    max_inventory = maximum(inventory)
    min_inventory = minimum(inventory)

    withdrawn_gas = withdraw.(inventory, Δt, min_inventory)
    purchased_gas = inject(inventory, time_step)

#end

function simulate(p::PriceProcess, m::Market)
    Δt = time_step(m)
    # initialize return value. Julia is column major and
    # access to period values per run will be need all at once,
    # hence they are the columns.
    retval = randn(Float64, m.runs, m.periods)
    for period in 1:m.periods
        prev_price = period == 1 ? fill(p.initial_price, m.runs) : retval[:,period-1]
        retval[:,period] .= prev_price .+ p.mean_rev_rate .* (p.mean_level .- prev_price) .*
            Δt .+ p.volatility .* prev_price .* sqrt(Δt) .* retval[:,period]
    end
    return retval
end

function time_step(m::Market)
    return m.horizon * 1 / m.periods
end

function terminal_value(prices,inventory, term_inv)
    return -2 * prices[:,end] .* max.(term_inv .- inventory, 0)
end

"""
    inject
    The rate of gas injected into the reservoir as a function of the current inventory.
"""
function inject(inventory, time_step)
    # hydro gas constants
    k2 = 730000
    k3 = 500
    k4 = 2500

    inven_factor = 1 / (inventory + k3)
    deduct_factor = 1 / k4
    return k2 * sqrt(inven_factor - deduct_factor) * time_step
end

"""
    withdraw releases gas from the reservoir based on the given inventory levels and
    the gas constants.
"""
function withdraw(inventory, time_step, min_inventory)
    k1 = 2040.41
    return max(k1 * sqrt(inventory) * time_step, min_inventory - inventory)
end

end

end #module


function μ(gun::Gun, target::AbstractTarget)
    μx = paralax(gun,target)[1]*target.position
    μy = paralax(gun,target)[2]*target.position
    return μx,μy
end

"""
    aiming_error(aiming)

Returns the alignment error (horizontal and vertical) for the error sources defined in 'aiming'
"""
function aiming_error(aiming::AimingError)
    σ_x =sqrt(aiming.point_def[1]^2+aiming.human[1]^2+aiming.res[1]^2+aiming.opt_path[1]^2)
    σ_y =sqrt(aiming.point_def[2]^2+aiming.human[2]^2+aiming.res[2]^2+aiming.opt_path[2]^2)
    return σ_x, σ_y
end
"""
    alignment_error(align)

Returns the alignment error (horizontal and vertical) for the error sources defined in 'align'
"""
function alignment_error(align::AlignmentError)
    σ_x=sqrt(align.boresight_acc[1]^2+align.boresight_ret[1]^2+align.sync[1]^2+align.stab[1]^2+align.gun_pos[1]^2)
    σ_y=sqrt(align.boresight_acc[2]^2+align.boresight_ret[2]^2+align.sync[2]^2+align.stab[2]^2+align.gun_pos[2]^2)
    return σ_x, σ_y
end

"""
    lof_error(lof)

Returns the line of fire error (horizontal and vertical) for the error sources defined in 'lof'
"""
function lof_error(lof::LofError)
    σ_x=sqrt(lof.fcs[1]^2+lof.lead[1]^2+lof.cant[1]^2+lof.earth_rate[1]^2+lof.muzzle_vel[1]^2+lof.temp[1]^2+lof.press[1]^2+lof.cross_wind[1]^2+lof.range_wind[1]^2+lof.range[1]^2+lof.jump[1]^2+lof.zeroing[1]^2)
    σ_y=sqrt(lof.fcs[2]^2+lof.lead[2]^2+lof.cant[2]^2+lof.earth_rate[2]^2+lof.muzzle_vel[2]^2+lof.temp[2]^2+lof.press[2]^2+lof.cross_wind[2]^2+lof.range_wind[2]^2+lof.range[2]^2+lof.jump[2]^2+lof.zeroing[2]^2)
    return σ_x, σ_y
end

"""
    flight_error(flight)

Returns the projectile dispersion (horizontal and vertical)
"""
function flight_error(flight::FlightError)
    σ_x=sqrt(flight.aero[1]^2)
    σ_y=sqrt(flight.aero[2]^2)
    return σ_x, σ_y
end

"""
    σ(aiming,align,lof,flight)

Returns the total error (horizontal and vertical) with error budget defined by error sources:
* aiming: aiming error
* align: alignment error
* lof: line of fire error
* flight: projectile dispersion
"""
function σ(aiming::AimingError,align::AlignmentError,lof::LofError,flight::FlightError)
    σ_x =sqrt(aiming_error(aiming)[1]^2+alignment_error(align)[1]^2+lof_error(lof)[1]^2+flight_error(flight)[1]^2)
    σ_y =sqrt(aiming_error(aiming)[2]^2+alignment_error(align)[2]^2+lof_error(lof)[2]^2+flight_error(flight)[2]^2)
    return σ_x, σ_y
end

"""
    σ()

Returns the total error of type 'Error'. The error budget is defined by type of errors:
* σfixed : fixed bias errors
* σvariable: varible bias errors
* σrandom: random errors
"""

function σ(;σfixed=Error(0.0,0.0), σvariable=Error(0.0,0.0), σrandom=Error(0.0,0.0))
    σv = σfixed.vertical + sqrt(σvariable.vertical^2 + σrandom.vertical^2)
    σh = σfixed.horizontal + sqrt(σvariable.horizontal^2 + σrandom.horizontal^2)
    σtot = Error(σh,σv)

    return σtot
end


using ComputationalHeatTransfer

using Plots
#using LaTeXStrings

ρ = 1;
c = 1;
k = 0.1;
d = 0.1;
params = HeatConductionParameters(ρ,c,k,thickness=d)

xlim = (-2.1,2.1)
ylim = (-2.1,2.1)
Δx, Δt = setstepsizes(params.α,gridPe=0.4,fourier=2.0)

#=
Here we set up a rectangular boundary
=#
#bdry = Circle(2.0,1.5Δx)
bdry = Rectangle(2.0,2.0,1.5Δx,shifted=true)

#eb = Circle(1.0,1.5*Δx)
eb = Rectangle(0.5,1.0,1.5*Δx)
Te = RigidTransform((0.0,0.0),0.0)
Te(eb)
cb1 = Rectangle(0.5,1.0,1.5*Δx)
cb2 = Rectangle(0.5,1.0,1.5*Δx)
Tc1 = RigidTransform((1.0,0.0),0.0)
Tc2 = RigidTransform((-1.0,0.0),0.0)

Tc1(cb1)
Tc2(cb2)

qe = 5.0
hc = 10.0
Tc = -10.0

eparams = PrescribedHeatFluxRegion(qe,eb);
cparams1 = PrescribedHeatModelRegion(hc,Tc,cb1);
cparams2 = PrescribedHeatModelRegion(hc,Tc,cb2);

#=
Note the use of the flag `bctype=ComputationalHeatTransfer.AdiabaticBC` here.
We also specify the boundary of the domain in `bodies`.
=#
sys = HeatConduction(params,Δx,xlim,ylim,Δt,bodies=BodyList([bdry]),qflux=eparams,qmodel=[cparams1,cparams2],bctype=ComputationalHeatTransfer.AdiabaticBC)

u0 = newstate(sys)
tspan = (0.0,10.0)
integrator = init(u0,tspan,sys)

step!(integrator,2.0)

plot(temperature(integrator),sys.grid,legend=true,color=cgrad(:RdBu,rev=true))
plot!(bdry,fillalpha=0.0)

xc, yc = coordinates(temperature(integrator),sys.grid);

plot(xc,temperature(integrator)[:,floor(Int,size(sys.grid,2)/2)])


###Define a structure to hold all our parameters

mutable struct ModeParameters
    l::Int64
    m::Int64
    n::Int64
    s::Int64
    a::Float64
    ω::Complex{Float64}
    Alm::Complex{Float64}
    Nmax::Int64
    lmax::Int64
end

ModeParameters(NT::NamedTuple) = ModeParameters(NT[:l],NT[:m],NT[:n],NT[:s],NT[:a],NT[:ω],NT[:Alm],NT[:Nmax],NT[:lmax])

##Parameter Transformation
function ParameterTransformations(l,m,s,a,ω,Alm)
    M = 1
    r₊ = M + sqrt(M^2 - a^2)
    r₋ = M - sqrt(M^2 - a^2)

    σ₊ = (2*ω*M*r₊ - m*a)/(r₊-r₋)
    σ₋ = (2*ω*M*r₋ - m*a)/(r₊-r₋)

    ζ₊ = im*ω
    ξ₋ = (-s-(s+2*im*σ₊))/2
    η₋ = -s+im*σ₋

    ζ = ζ₊
    ξ = ξ₋
    η = η₋

    p = (r₊ - r₋)*(ζ/2)
    α = 1 + s + ξ + η - 2*ζ + s*im*ω/ζ
    γ = 1 + s + 2*η
    δ = 1 + s + 2*ξ
    σ = Alm + (a*ω)^2 - 8*ω^2 + p*(2*α + γ - δ) + (1 + s - (γ + δ)/2 )*(s + (γ + δ)/2 )

    D₀ = δ
    D₁ = 4*p - 2*α + γ - δ -2
    D₂ = 2*α - γ + 2
    D₃ = α*(4*p - δ) - σ
    D₄ = α*(α - γ + 1)

    return ((ζ,ξ,η),(p,α,γ,δ,σ),(D₀,D₁,D₂,D₃,D₄))
end

ParameterTransformations(P::ModeParameters) = ParameterTransformations(P.l,P.m,P.s,P.a,P.ω,P.Alm)
ParameterTransformations(l,m,n,s,a,ω,Alm,Nmax,lmax) = ParameterTransformations(l,m,s,a,ω,Alm)


export RK3

"""
    RK3()

The classical [Runge-Kutta](https://en.wikipedia.org/wiki/Runge–Kutta_methods)
predictor of order 3.
"""
struct RK3 <: AbstractPredictor end
struct RK3Cache{T} <: AbstractPredictorCache
    dt::Vector{T}
    k1::Vector{T}
    k2::Vector{T}
    k3::Vector{T}
    k4::Vector{T}
end

function cache(::RK3, H, x, ẋ, t)
    RK3Cache(dt(H, x, t), copy(ẋ), copy(ẋ), copy(ẋ), copy(ẋ))
end
#
function predict!(
    xnext,
    cache::RK3Cache,
    H::HomotopyWithCache,
    x,
    t,
    Δt,
    ẋ,
    Jac::JacobianMonitor,
)
    dt, mk₁, mk₂, mk₃ = cache.dt, cache.k1, cache.k2, cache.k3
    n = length(xnext)
    @inbounds for i = 1:n
        mk₁[i] = -ẋ[i]
        xnext[i] = x[i] - 0.5Δt * mk₁[i]
    end
    minus_ẋ!(mk₂, H, xnext, t + 0.5Δt, Jac, dt)

    @inbounds for i = 1:n
        xnext[i] = x[i] + Δt * mk₁[i] - 2Δt * mk₂[i]
    end
    minus_ẋ!(mk₃, H, xnext, t + Δt, Jac, dt)
    @inbounds for i = 1:n
        xnext[i] = x[i] - Δt * mk₃[i]
    end

    @inbounds for i = 1:n
        xnext[i] = x[i] - Δt / 6 * (mk₁[i] + 4 * mk₂[i] + mk₃[i])
    end

    nothing
end

order(::RK3Cache) = 4


#commented psuedocode
using Meshes
using StaticArrays
using LinearAlgebra
using CoordinateTransformations
using Tmatrix

function calculate_Tmatrix_for_any(
    face_list::Vector{Any}, #make this better typed
    point_list::Vector{Vector{Float64}},
    n_max::Int,
    k1::Complex{R},
    k2::Complex{R};
    rotationally_symmetric = false,
    symmetric_about_plane_perpendicular_z = false,
    BigFloat_precision = nothing,
) where {R <: Real}

    #spherical_list = convert_to_spherical(point_list) #format is [r, θ, ϕ], use coordinatetransformations.jl to convert

    #TODO: bundle these into a new struct and make that not break
    r_array, θ_array, ϕ_array = (
        [element[1] for element in spherical_list],
        [element[2] for element in spherical_list],
        [element[3] for element in spherical_list],
    )

    # calculate r and n̂ for the geometry

    n̂_array = point_list #copying so that size is same

    for element in n̂_array
        element = [0, 0, 0]
    end

    for face in face_list
        for (vertex, i) in enumerate(face)
            n_array[vertex] += vertex × face[(i + 1) % size(face)]
        end
    end

    for element in n̂_array
        element = normalize(element) #TODO make sure this works, use normalize!(element) later if at all
    end

    # calculate T-matrix
    k1r_array = k1 .* r_array
    k2r_array = k2 .* r_array
    T = T_matrix(
        n_max,
        k1,
        k2,
        k1r_array,
        k2r_array,
        r_array,
        θ_array,
        ϕ_array,
        n̂_array;
        HDF5_filename = HDF5_filename,
        rotationally_symmetric = rotationally_symmetric,
        symmetric_about_plane_perpendicular_z = symmetric_about_plane_perpendicular_z,
        BigFloat_precision = BigFloat_precision,
    )
    return T
end

point_list = Point(3, 8, 23), Point(1, 2, 3), Point(5, 12, 12), Point(15, 3, 7)

face_list = [connect((1, 2, 3)), connect((1, 2, 4)), connect((1, 3, 4)), connect((2, 3, 4))]

#convert_to_spherical

#TODO make generic over shape
function Spherical(point_coords::Point)
    Spherical(point_coords.coords[1], point_coords.coords[2], point_coords.coords[3])
end

polar_point_list = [Spherical(p) for p in point_list]

n_max = 1
k1_r, k1_i, k2_r, k2_i = 1e5, 1e3, 2e5, 3e3

calculate_Tmatrix_for_spheroid_SeparateRealImag(
    face_list,
    point_list,
    n_max,
    k1_r,
    k1_i,
    k2_r,
    k2_i;
    rotationally_symmetric = false,
    symmetric_about_plane_perpendicular_z = false,
);


#this example is modified from https://numericalenvironmental.wordpress.com/2018/05/26/a-steady-state-variably-saturated-flow-model-in-vertical-cross-section-a-finite-difference-approach-using-julia/
using BenchmarkTools
using DiffEqSensitivity
import DifferentiableBackwardEuler
import NonlinearEquations
using OrdinaryDiffEq
#import PyPlot
import Zygote

#include("ex.jl")
include("utilities.jl")
include("inputdeck.jl")

@NonlinearEquations.equations exclude=(coords, dirichletnodes, neighbors, areasoverlengths) function req(psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, specificstorage, volumes)
	NonlinearEquations.setnumequations(length(psi))
	for i = 1:length(psi)
		if i in dirichletnodes
			NonlinearEquations.addterm(i, psi[i] - dirichletpsis[i])
		else
			NonlinearEquations.addterm(i, Qs[i] / (specificstorage[i] * volumes[i]))
		end
	end
	for (i, (node_a, node_b)) in enumerate(neighbors)
		for (node1, node2) in [(node_a, node_b), (node_b, node_a)]
			if !(node1 in dirichletnodes) && !(node2 in dirichletnodes)
				NonlinearEquations.addterm(node1, hm(kr(psi[node1], alphas[i], Ns[i]), kr(psi[node2], alphas[i], Ns[i])) * Ks[i] * (psi[node2] + coords[2, node2] - psi[node1] - coords[2, node1]) * areasoverlengths[i] / (specificstorage[node1] * volumes[node1]))
			elseif !(node1 in dirichletnodes) && node2 in dirichletnodes
				NonlinearEquations.addterm(node1, hm(kr(psi[node1], alphas[i], Ns[i]), kr(dirichletpsis[node2], alphas[i], Ns[i])) * Ks[i] * (dirichletpsis[node2] + coords[2, node2] - psi[node1] - coords[2, node1]) * areasoverlengths[i] / (specificstorage[node1] * volumes[node1]))
			end
		end
	end
end

req_residuals(us, Ks, Qs) = req_residuals(us, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, specificstorage, volumes)
req_psi(us, Ks, Qs) = req_psi(us, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, specificstorage, volumes)
req_Ks(us, Ks, Qs) = req_Ks(us, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, specificstorage, volumes)
req_Qs(us, Ks, Qs) = req_Qs(us, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, specificstorage, volumes)

@Zygote.adjoint function req_residuals(us, Ks, Qs)
	residuals = req_residuals(us, Ks, Qs)
	back = delta->begin
		retval = ((delta' * req_psi(us, Ks, Qs))', (delta' * req_Ks(us, Ks, Qs))', (delta' * req_Qs(us, Ks, Qs))')
		return retval
	end
	return residuals, back
end

function unpack(p)
	@assert length(p) == length(neighbors)
	Ks = p[1:length(neighbors)]
	return Ks
	#=
	@assert length(p) == length(neighbors) + size(coords, 2)
	Ks = p[1:length(neighbors)]
	Qs = p[length(neighbors) + 1:end]
	return Ks, Qs
	=#
end

function f(u, p, t)
	Ks = unpack(p)
	return req_residuals(u, Ks, Qs)
end

function f_u(u, p, t)
	Ks = unpack(p)
	return req_psi(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, specificstorage, volumes)
end

function f_p(u, p, t)
	Ks = unpack(p)
	return req_Ks(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, specificstorage, volumes)
end

f_t(u, p, t) = zeros(length(u))

psi0 = fill(-5.0, size(coords, 2))
psi0[dirichletnodes] = dirichletpsis[dirichletnodes]
p = Ks
tspan = [0, 1e2]
odef = ODEFunction(f; jac=f_u, jac_prototype=f_u(psi0, p, 0.0))
prob = ODEProblem(odef, psi0, tspan, p)
@time soln_diffeq = solve(prob, ImplicitEuler())
@time soln_dbe = DifferentiableBackwardEuler.steps(psi0, f, f_u, f_p, f_t, p, tspan[1], tspan[2]; ftol=1e-12, method=:newton)
@show sum((soln_diffeq[:, :] .- soln_dbe) .^ 2)
@show sum(soln_diffeq[:, :] .^ 2)
obsnode = div(3 * length(psi0), 4)
g(p) = DifferentiableBackwardEuler.steps(psi0, f, f_u, f_p, f_t, p, tspan[1], tspan[2]; ftol=1e-12, method=:newton)[obsnode, end]
g(p)
@time g(p)
grad_zygote = Zygote.gradient(g, p)[1]
@time grad_zygote = Zygote.gradient(g, p)[1]

function naivegradient(f, x0s...; delta::Float64=1e-8)
	f0 = f(x0s...)
	gradient = map(x->zeros(size(x)), x0s)
	for j = 1:length(x0s)
		x0 = copy(x0s[j])
		for i = 1:length(x0)
			x = copy(x0)
			x[i] += delta
			xs = (x0s[1:j - 1]..., x, x0s[j + 1:length(x0s)]...)
			fval = f(xs...)
			gradient[j][i] = (fval - f0) / delta
		end
	end
	return gradient
end

function checkgradientquickly(f, x0, gradf, n; delta::Float64=1e-8)
	indicestocheck = sort(collect(1:length(x0)), by=i->abs(gradf[i]), rev=true)[1:n]
	indicestocheck = [indicestocheck; rand(1:length(x0), n)]
	f0 = f(x0)
	for i in indicestocheck
		x = copy(x0)
		x[i] += delta
		fval = f(x)
		grad_f_i = (fval - f0) / delta
		@show i
		@show grad_f_i
		@show gradf[i]
	end
end
checkgradientquickly(g, p, grad_zygote, 5; delta=1e-4)

#@show length(p)
#@time grad_naive = naivegradient(g, p)[1]
#@show sqrt(sum((grad_zygote .- grad_naive) .^ 2) / sum(grad_naive .^ 2))
#@show sum(grad_naive .^ 2)


using GLVisualize, GeometryTypes, Colors
using GLAbstraction
if !isdefined(:runtests)
	window = glscreen()
	timesignal = bounce(linspace(0f0,1f0,360))
end

description = """
Animated 2D particles in 3D space with glow.
"""

function sin_torus(radius, thickness, N, t)
    pirange = linspace(0, 2pi, N^2)
    points = zeros(Point3f0, N^2)
    p,q = 6,9
    for (i,x)=enumerate(pirange)
        r = cos(q*(x+t))+thickness
        points[i] = Point3f0(r*cos(p*x),r*sin(p*x), -sin(q*x))/radius
    end
    points
end

t = const_lift(*, 1000f0, timesignal)
ps1 = const_lift(sin_torus, 2, 6, 10, t)
ps2 = const_lift(sin_torus, 5, 4, 10, t)
ps3 = const_lift(sin_torus, 10, 2, 10, t)
ps = map(vcat, ps1, ps2, ps3)
l1, l2, l3 = map(x->length(value(x)), (ps1, ps2, ps3))
colors = [
    fill(RGBA(0.9f0, 0f0, 0.7f0, 1f0), l1);
    fill(RGBA(0.7f0, 0.2f0, 0.8f0, 1f0), l2);
    fill(RGBA(1.0f0, 0.4f0, 0.95f0, 1f0), l3);
]
scales = [
    fill(Vec2f0(0.1), l3);
    fill(Vec2f0(0.05), l2);
    fill(Vec2f0(0.04), l1);
]

_view(visualize(
    (Circle, ps),
    glow_color = RGBA(0.8f0, 0.6f0, 0.95f0, 0.8f0),
    glow_width = const_lift(/, timesignal, 20f0),
    scale = scales, color = colors
), camera=:perspective)

if !isdefined(:runtests)
    renderloop(window)
end