using Altro
using TrajectoryOptimization
using StaticArrays
using BenchmarkTools
using RobotZoo
using RobotDynamics
using StaticArrays, LinearAlgebra
using ForwardDiff
using Test
const TO = TrajectoryOptimization
const RD = RobotDynamics

## Test Functions
function Πsoc(x)
    v = x[1:end-1]
    s = x[end]
    a = norm(v)
    if a <= -s
        return zero(x)
    elseif a <= s
        return x
    elseif a >= abs(s)
        x̄ = append!(v, a)
        return 0.5*(1 + s/a) * x̄
    end
    throw(ErrorException("Invalid second-order cone"))
end
function in_soc(x)
    v = x[1:end-1]
    s = x[end]
    a = norm(v)
	a <= s
end

function auglag(f,g, x,z,μ)
    # ceq = h(x)
    L0 = f(x) #+ y'ceq + 0.5*μ*ceq'ceq

    cones = g(x)
    p = length(cones)
    for i = 1:p
        cone = cones[i]
		proj = Πsoc(z[i] - μ*cones[i])
		pen = proj'proj - z[i]'z[i]
        L0 += 1 / (2μ) * pen
    end
    return L0
end

function DI_objective(Q,R,Qf,xf,dt)
    g = Float64[]
    h = Float64[]
    c = 0.0
    for k = 1:N-1
        append!(g, -Q*xf)
        append!(g, zeros(m))
        append!(h, Q.diag)
        append!(h, R.diag)
        c += 0.5*xf'Q*xf
	end
	g .*= dt
	h .*= dt
	c *= dt
    append!(g, -Qf*xf)
    # append!(g, zeros(m))
    append!(h, Qf.diag)
    # append!(h, R.diag)
    c += 0.5*xf'Qf*xf
    H = Diagonal(h)
    return H,g,c
end

function DI_dynamics(x0, N; dt=0.1, gravity::Bool=false)
    n = length(x0)
    m = n ÷ 2
    D = m
    NN = N*(n+m)

    #  Continuous dynamics
    Ac = zeros(n,n)
    Bc = zeros(n,m)
    for i = 1:D
        Ac[i,D+i] = 1
        Bc[D+i,i] = 1
    end

    # Euler integration
    Ad = I + Ac*dt
    Bd = Bc*dt

    iA = 1:n
    jA = iA
    jB = n .+ (1:m)
    jA2 = jA .+ (n+m)

    A = zeros(N*n, NN)
    b = zeros(N*n)
    A[iA,jA] .= I(n)
    b[iA] .= -x0
    iA = iA .+ n
    for k = 1:N-1
        A[iA,jA] .= Ad
        A[iA,jB] .= Bd
        A[iA,jA2] .= -I(n)
        if gravity
            b[iA[D]] = -9.81
        end

        # Advance inds
        iA = iA .+ n
        jA = jA .+ (n+m)
        jB = jB .+ (n+m)
        jA2 = jA2 .+ (n+m)
    end
    return A,b
end

function DI_cones(::Val{D},N) where D
    uinds = SVector{D}(1:D) .+ 2D
    qinds = [uinds .+ (k-1)*3D for k = 1:N-1]
end
DI_cones(D,N) = DI_cones(Val(D),N)  # WARNING: type unstable

## Setup
D,N = 2,11
n,m = 2D,D
dt = 0.1
tf = (N-1)*dt
x0 = @SVector fill(0., n)
xf = [(@SVector fill(1.,D)); (@SVector fill(0, D))]
Q = Diagonal(@SVector fill(1.0, n))
R = Diagonal(@SVector fill(1e-1, m))
Qf = (N-1)*Q * 100
qinds = DI_cones(D,N)
u_bnd = 6.0

## Define batch functions
H_obj,g_obj,c_obj = DI_objective(Q,R,Qf,xf,dt)
Adyn,bdyn = DI_dynamics(x0,N, dt=dt)

di_obj(x) = 0.5*x'H_obj*x + g_obj'x + c_obj
di_dyn(x) = Adyn*x + bdyn
di_soc(x) = [push(x[qi],u_bnd) for qi in qinds]

## Define using ALTRO
model = RobotZoo.DoubleIntegrator(D)

obj = LQRObjective(Q, R, Qf, xf, N)
cons = ConstraintList(n,m,N)
cone = TO.SecondOrderCone()
# cone = TO.Inequality()
u_bnd = 6.0
connorm = NormConstraint(n, m, u_bnd, cone, :control)
add_constraint!(cons, connorm, 1:N-1)
prob = Problem(model, obj, xf, tf, x0=x0, constraints=cons, integration=RD.Euler)
solver = Altro.AugmentedLagrangianSolver(prob)
conSet = TO.get_constraints(solver)
alobj = TO.get_objective(solver)

## Initialization
inds = LinearIndices(zeros(3D,N))
xinds = [SVector{n}(z[1:n]) for z in eachcol(inds)]
uinds = [SVector{m}(z[n+1:end]) for z in eachcol(inds)][1:end-1]
NN = 3D*N-D
P = (N-1)*m
x0 = rand(NN)
z = rand.(length.(di_soc(x0)))
μ = 2.4

X0 = [x0[xi] for xi in xinds]
U0 = [x0[ui] for ui in uinds]
Z0 = Traj(X0,U0,fill(dt,N))
initial_trajectory!(solver, Z0)

## Test objective and soc constraints
cost(obj, Z0) ≈ di_obj(x0)
conSet.convals[1].vals ≈ di_soc(x0)

# copy multipliers to ALconSet
cval = conSet.convals[1]
for i = 1:N-1
	cval.λ[i] .= z[i]
	cval.μ[i] .= μ
end

# test augmented Lagrangian value
LA(x) = auglag(di_obj, di_soc, x, z, μ)
cost(solver) ≈ LA(x0)
!(cost(solver) ≈ di_obj(x0))   # make sure the AL cost isn't the same as the normal cost

λbar = z .- μ .* di_soc(x0)
statuses = [TO.cone_status(TO.SecondOrderCone(), λ) for λ in λbar]
@test :below ∈ statuses

# E = TO.QuadraticObjective(n,m,N)
E = TO.CostExpansion(n,m,N)
TO.cost_expansion!(E, alobj, Z0)
grad = vcat([[e.q; e.r] for e in E]...)[1:end-m]
@test grad ≈ ForwardDiff.gradient(LA, x0)

hess_blocks = vcat([[e.Q, e.R] for e in E]...)
hess = cat(hess_blocks..., dims=(1,2))[1:end-m, 1:end-m]
@test hess ≈ ForwardDiff.hessian(LA, x0)

@test all(in_soc.(di_soc(x0)))
@test !all(in_soc.(z .- μ*di_soc(x0)))

# Make it go outside the cone
z = rand.(length.(di_soc(x0)))
z .*= 100
λbar = z .- μ .* di_soc(x0)
statuses = [TO.cone_status(TO.SecondOrderCone(), λ) for λ in λbar]
@test :outside ∈ statuses
@test :in ∈ statuses
for i = 1:N-1
	cval.λ[i] .= z[i]
	cval.μ[i] .= μ
end
LA(x) = auglag(di_obj, di_soc, x, z, μ)

# E = TO.QuadraticObjective(n,m,N)
E = TO.CostExpansion(n,m,N)
TO.cost_expansion!(E, alobj, Z0)
grad = vcat([[e.q; e.r] for e in E]...)[1:end-m]
@test grad ≈ ForwardDiff.gradient(LA, x0)

hess_blocks = vcat([[e.Q, e.R] for e in E]...)
hess = cat(hess_blocks..., dims=(1,2))[1:end-m, 1:end-m]
@test hess ≈ ForwardDiff.hessian(LA, x0)

## Solve it
solver = Altro.AugmentedLagrangianSolver(prob, verbose=0, show_summary=false,
	projected_newton=true, penalty_initial=100.0, penalty_scaling=50,
	cost_tolerance_intermediate=1e-1)
initial_controls!(solver, [rand(D) for k = 1:N])
solve!(solver)
norm.(controls(solver))
@test all(abs.(norm.(controls(solver))[[1,2,N-2,N-1]] .- u_bnd) .< 1e-6)

## Rocket landing problem
D,N = 3,101
n,m = 2D,D
dt = 0.1
tf = (N-1)*dt
x0 = SA_F64[10,10,100, 0,0,-10]
xf = @SVector zeros(n)
Q = Diagonal(@SVector fill(1.0, n))
R = Diagonal(@SVector fill(1e-1, m))
Qf = (N-1)*Q * 100

##
model = RobotZoo.DoubleIntegrator(3, gravity=SA[0,0,-9.81])
obj = LQRObjective(Q,R,Qf,xf,N)
cons = ConstraintList(n,m,N)
u_bnd = 100.0
connorm = NormConstraint(n, m, u_bnd, TO.SecondOrderCone(), :control)
add_constraint!(cons, GoalConstraint(xf), N)
add_constraint!(cons, connorm, 1:N-1)

prob = Problem(model, obj, xf, tf, x0=x0, constraints=cons)
solver = ALTROSolver(prob, projected_newton=false, show_summary=false)
solve!(solver)
@test iterations(solver) == 12

@test abs(maximum(norm.(controls(solver))) - u_bnd) < 1e-6
@test norm(states(solver)[end] - xf) < 1e-6

# prob = Problem(model, obj, xf, tf, x0=x0, constraints=cons)
# solver = ALTROSolver(prob, show_summary=true, verbose=2, projected_newton=false)
# @test benchmark_solve!(solver).allocs == 0
# benchmark_solve!(solver)


module Day2

using DelimitedFiles

mutable struct Position
    x::Int
    depth::Int
    aim::Int
end

update_aiming!(position, command) = begin
    parts = split(command)
    dir = parts[1]
    n = parse(Int, parts[2])
    if dir == "forward"
        position.x += n
        position.depth += n * position.aim
    elseif dir == "down"
        position.aim += n
    elseif dir == "up"
        position.aim -= n
    end
end

update!(position, command) = begin
    parts = split(command)
    dir = parts[1]
    n = parse(Int, parts[2])
    if dir == "forward"
        position.x += n
    elseif dir == "down"
        position.depth += n
    elseif dir == "up"
        position.depth -= n
    end
end

run_part1(commands) = begin
    pos = Position(0, 0, 0)
    for command in commands
        update!(pos, command)
    end
    @show pos.x
    @show pos.depth
    return pos.x * pos.depth
end

run_part2(commands) = begin
    pos = Position(0, 0, 0)
    for command in commands
        update_aiming!(pos, command)
    end
    return pos.x * pos.depth
end

test = ["forward 5"
        "down 5"
        "forward 8"
        "up 3"
        "down 8"
        "forward 2"];

@assert run_part1(test) == 150;
@assert run_part2(test) == 900;

input = readdlm("resources/day2_input.txt", '\t', String, '\n') |> vec;

@show (input |> run_part1)
@show (input |> run_part2)

end


function dijkstra(net::Network{T}, snk::U) where {T<:Integer, U<:Integer}
    nnodes = numnodes(net)

    costs = fill(typemax(U), nnodes)
    childvec = zeros(T, nnodes)

    P = PriorityQueue{T,U}()
    P[snk] = costs[snk] = 0
    childvec[snk] = snk

    while !isempty(P)
        v = dequeue!(P)
        for u in inneighbors(net, v)
            lidx = idx(net, u, v)
            alt = costs[v] + net.links[lidx].l
            if (costs[u] > alt)
                P[u] = costs[u] = alt
                childvec[u] = v
            end
        end
    end
    childvec
end

function initchoices()
    # choice initialization
    DTC = zeros(T, size(trips)...);
    DTC[1,..] .= 1.
    #DTC[1:Tm,..] .= rand(size(DTC[1:Tm,..])...)
    #DTC ./= sum(DTC, dims=1)

    pathvecs = Dict(k => dijkstra(net, k) for k in snks);
    SR = Dict(d => Dict(i => zeros(T, nsinks, nclasses) for i in 1:2) for d in divs);
    for d in divs
        for (kno,k) in enumerate(snks)
            for c in 1:nclasses
                if (outneighbors(net, d)[1] == pathvecs[k][d])
                    SR[d][1][:,kno,c] .= 1.
                elseif (outneighbors(net, d)[2] == pathvecs[k][d])
                    SR[d][2][:,kno,c] .= 1.
                else
                    error()
                end
            end
        end
    end
    return (DTC, SR)
end


# This file is part of the IntervalArithmetic.jl package; MIT licensed

## Comparisons

"""
    ==(a,b)

Checks if the intervals `a` and `b` are equal.
"""
function ==(a::Interval, b::Interval)
    isempty(a) && isempty(b) && return true
    a.lo == b.lo && a.hi == b.hi
end
!=(a::Interval, b::Interval) = !(a==b)

# Auxiliary functions: equivalent to </<=, but Inf <,<= Inf returning true
function islessprime(a::T, b::T) where T<:Real
    (isinf(a) || isinf(b)) && a==b && return true
    a < b
end

# Weakly less, \le, <=
function <=(a::Interval, b::Interval)
    isempty(a) && isempty(b) && return true
    (isempty(a) || isempty(b)) && return false
    (a.lo ≤ b.lo) && (a.hi ≤ b.hi)
end

# Strict less: <
function <(a::Interval, b::Interval)
    isempty(a) && isempty(b) && return true
    (isempty(a) || isempty(b)) && return false
    islessprime(a.lo, b.lo) && islessprime(a.hi, b.hi)
end

# precedes
function precedes(a::Interval, b::Interval)
    (isempty(a) || isempty(b)) && return true
    a.hi ≤ b.lo
end

# strictpreceds
function strictprecedes(a::Interval, b::Interval)
    (isempty(a) || isempty(b)) && return true
    # islessprime(a.hi, b.lo)
    a.hi < b.lo
end
const ≺ = strictprecedes # \prec

# zero, one, typemin, typemax
zero(a::Interval{T}) where T<:Real = Interval(zero(T))
zero(::Type{Interval{T}}) where T<:Real = Interval(zero(T))
one(a::Interval{T}) where T<:Real = Interval(one(T))
one(::Type{Interval{T}}) where T<:Real = Interval(one(T))
typemin(::Type{Interval{T}}) where T<:AbstractFloat = wideinterval(typemin(T))
typemax(::Type{Interval{T}}) where T<:AbstractFloat = wideinterval(typemax(T))
typemin(::Type{Interval{T}}) where T<:Integer = Interval(typemin(T))
typemax(::Type{Interval{T}}) where T<:Integer = Interval(typemax(T))

## Addition and subtraction

+(a::Interval) = a
-(a::Interval) = Interval(-a.hi, -a.lo)

function +(a::Interval{T}, b::T) where {T<:Real}
    isempty(a) && return emptyinterval(T)
    @round(a.lo + b, a.hi + b)
end
+(b::T, a::Interval{T}) where {T<:Real} = a+b

function -(a::Interval{T}, b::T) where {T<:Real}
    isempty(a) && return emptyinterval(T)
    @round(a.lo - b, a.hi - b)
end
function -(b::T, a::Interval{T}) where {T<:Real}
    isempty(a) && return emptyinterval(T)
    @round(b - a.hi, b - a.lo)
end

function +(a::Interval{T}, b::Interval{T}) where T<:Real
    (isempty(a) || isempty(b)) && return emptyinterval(T)
    @round(a.lo + b.lo, a.hi + b.hi)
end

function -(a::Interval{T}, b::Interval{T}) where T<:Real
    (isempty(a) || isempty(b)) && return emptyinterval(T)
    @round(a.lo - b.hi, a.hi - b.lo)
end

## Multiplication
function *(x::T, a::Interval{T}) where {T<:Real}
    isempty(a) && return emptyinterval(T)
    (iszero(a) || iszero(x)) && return zero(Interval{T})

    if x ≥ 0.0
        return @round(a.lo*x, a.hi*x)
    else
        return @round(a.hi*x, a.lo*x)
    end
end
*(a::Interval{T}, x::T) where {T<:Real} = x*a

function *(a::Interval{T}, b::Interval{T}) where T<:Real
    (isempty(a) || isempty(b)) && return emptyinterval(T)

    (iszero(a) || iszero(b)) && return zero(Interval{T})

    if b.lo >= zero(T)
        a.lo >= zero(T) && return @round(a.lo*b.lo, a.hi*b.hi)
        a.hi <= zero(T) && return @round(a.lo*b.hi, a.hi*b.lo)
        return @round(a.lo*b.hi, a.hi*b.hi)   # zero(T) ∈ a
    elseif b.hi <= zero(T)
        a.lo >= zero(T) && return @round(a.hi*b.lo, a.lo*b.hi)
        a.hi <= zero(T) && return @round(a.hi*b.hi, a.lo*b.lo)
        return @round(a.hi*b.lo, a.lo*b.lo)   # zero(T) ∈ a
    else
        a.lo > zero(T) && return @round(a.hi*b.lo, a.hi*b.hi)
        a.hi < zero(T) && return @round(a.lo*b.hi, a.lo*b.lo)
        return @round(min(a.lo*b.hi, a.hi*b.lo), max(a.lo*b.lo, a.hi*b.hi))
    end
end

## Division
function /(a::Interval{T}, x::T) where {T<:Real}
    isempty(a) && return emptyinterval(T)
    iszero(x) && return emptyinterval(T)
    iszero(a) && return zero(Interval{T})

    if x ≥ 0.0
        return @round(a.lo/x, a.hi/x)
    else
        return @round(a.hi/x, a.lo/x)
    end
end

function inv(a::Interval{T}) where T<:Real
    isempty(a) && return emptyinterval(a)

    if zero(T) ∈ a
        a.lo < zero(T) == a.hi && return @round(-Inf, inv(a.lo))
        a.lo == zero(T) < a.hi && return @round(inv(a.hi), Inf)
        a.lo < zero(T) < a.hi && return entireinterval(T)
        a == zero(a) && return emptyinterval(T)
    end

    @round(inv(a.hi), inv(a.lo))
end

function /(a::Interval{T}, b::Interval{T}) where T<:Real

    S = typeof(a.lo / b.lo)
    (isempty(a) || isempty(b)) && return emptyinterval(S)
    iszero(b) && return emptyinterval(S)

    if b.lo > zero(T) # b strictly positive

        a.lo >= zero(T) && return @round(a.lo/b.hi, a.hi/b.lo)
        a.hi <= zero(T) && return @round(a.lo/b.lo, a.hi/b.hi)
        return @round(a.lo/b.lo, a.hi/b.lo)  # zero(T) ∈ a

    elseif b.hi < zero(T) # b strictly negative

        a.lo >= zero(T) && return @round(a.hi/b.hi, a.lo/b.lo)
        a.hi <= zero(T) && return @round(a.hi/b.lo, a.lo/b.hi)
        return @round(a.hi/b.hi, a.lo/b.hi)  # zero(T) ∈ a

    else   # b contains zero, but is not zero(b)

        iszero(a) && return zero(Interval{S})

        if iszero(b.lo)

            a.lo >= zero(T) && return @round(a.lo/b.hi, Inf)
            a.hi <= zero(T) && return @round(-Inf, a.hi/b.hi)
            return entireinterval(S)

        elseif iszero(b.hi)

            a.lo >= zero(T) && return @round(-Inf, a.lo/b.lo)
            a.hi <= zero(T) && return @round(a.hi/b.lo, Inf)
            return entireinterval(S)

        else

            return entireinterval(S)

        end
    end
end

function extended_div(a::Interval{T}, b::Interval{T}) where T<:Real

    S = typeof(a.lo / b.lo)
    if 0 < b.hi && 0 > b.lo && 0 ∉ a
        if a.hi < 0
            return (Interval(-Inf, a.hi / b.hi), Interval(a.hi / b.lo, Inf))

        elseif a.lo > 0
            return (Interval(-Inf, a.lo / b.lo), Interval(a.lo / b.hi, Inf))

        end
    elseif 0 ∈ a && 0 ∈ b
        return (entireinterval(S), emptyinterval(S))
    else
        return (a / b, emptyinterval(S))
    end
end

//(a::Interval, b::Interval) = a / b    # to deal with rationals

function min_ignore_nans(args...)
    min(Iterators.filter(x->!isnan(x), args)...)
end

function max_ignore_nans(args...)
    max(Iterators.filter(x->!isnan(x), args)...)
end

## fma: fused multiply-add
function fma(a::Interval{T}, b::Interval{T}, c::Interval{T}) where T
    #T = promote_type(eltype(a), eltype(b), eltype(c))

    (isempty(a) || isempty(b) || isempty(c)) && return emptyinterval(T)

    if isentire(a)
        b == zero(b) && return c
        return entireinterval(T)

    elseif isentire(b)
        a == zero(a) && return c
        return entireinterval(T)

    end

    lo = setrounding(T, RoundDown) do
        lo1 = fma(a.lo, b.lo, c.lo)
        lo2 = fma(a.lo, b.hi, c.lo)
        lo3 = fma(a.hi, b.lo, c.lo)
        lo4 = fma(a.hi, b.hi, c.lo)
        min_ignore_nans(lo1, lo2, lo3, lo4)
    end

    hi = setrounding(T, RoundUp) do
        hi1 = fma(a.lo, b.lo, c.hi)
        hi2 = fma(a.lo, b.hi, c.hi)
        hi3 = fma(a.hi, b.lo, c.hi)
        hi4 = fma(a.hi, b.hi, c.hi)
        max_ignore_nans(hi1, hi2, hi3, hi4)
    end

    Interval(lo, hi)
end

## Scalar functions on intervals (no directed rounding used)

function mag(a::Interval{T}) where T<:Real
    isempty(a) && return convert(eltype(a), NaN)
    # r1, r2 = setrounding(T, RoundUp) do
    #     abs(a.lo), abs(a.hi)
    # end
    max( abs(a.lo), abs(a.hi) )
end

function mig(a::Interval{T}) where T<:Real
    isempty(a) && return convert(eltype(a), NaN)
    zero(a.lo) ∈ a && return zero(a.lo)
    r1, r2 = setrounding(T, RoundDown) do
        abs(a.lo), abs(a.hi)
    end
    min( r1, r2 )
end

# Infimum and supremum of an interval
inf(a::Interval) = a.lo
sup(a::Interval) = a.hi

## Functions needed for generic linear algebra routines to work
real(a::Interval) = a

function abs(a::Interval)
    isempty(a) && return emptyinterval(a)
    Interval(mig(a), mag(a))
end

function abs2(a::Interval)
    sqr(a)
end

function min(a::Interval, b::Interval)
    (isempty(a) || isempty(b)) && return emptyinterval(a)
    Interval( min(a.lo, b.lo), min(a.hi, b.hi))
end

function max(a::Interval, b::Interval)
    (isempty(a) || isempty(b)) && return emptyinterval(a)
    Interval( max(a.lo, b.lo), max(a.hi, b.hi))
end

dist(a::Interval, b::Interval) = max(abs(a.lo-b.lo), abs(a.hi-b.hi))
eps(a::Interval) = Interval(max(eps(a.lo), eps(a.hi)))
eps(::Type{Interval{T}}) where T<:Real = Interval(eps(T))

## floor, ceil, trunc, sign, roundTiesToEven, roundTiesToAway
function floor(a::Interval)
    isempty(a) && return emptyinterval(a)
    Interval(floor(a.lo), floor(a.hi))
end

function ceil(a::Interval)
    isempty(a) && return emptyinterval(a)
    Interval(ceil(a.lo), ceil(a.hi))
end

function trunc(a::Interval)
    isempty(a) && return emptyinterval(a)
    Interval(trunc(a.lo), trunc(a.hi))
end

function sign(a::Interval)
    isempty(a) && return emptyinterval(a)
    return Interval(sign(a.lo), sign(a.hi))
end

# RoundTiesToEven is an alias of `RoundNearest`
const RoundTiesToEven = RoundNearest
# RoundTiesToAway is an alias of `RoundNearestTiesAway`
const RoundTiesToAway = RoundNearestTiesAway

"""
    round(a::Interval[, RoundingMode])

Returns the interval with rounded to an interger limits.

For compliance with the IEEE Std 1788-2015, "roundTiesToEven" corresponds
to `round(a)` or `round(a, RoundNearest)`, and "roundTiesToAway"
to `round(a, RoundNearestTiesAway)`.
"""
round(a::Interval) = round(a, RoundNearest)
round(a::Interval, ::RoundingMode{:ToZero}) = trunc(a)
round(a::Interval, ::RoundingMode{:Up}) = ceil(a)
round(a::Interval, ::RoundingMode{:Down}) = floor(a)

function round(a::Interval, ::RoundingMode{:Nearest})
    isempty(a) && return emptyinterval(a)
    Interval(round(a.lo), round(a.hi))
end

function round(a::Interval, ::RoundingMode{:NearestTiesAway})
    isempty(a) && return emptyinterval(a)
    Interval(round(a.lo, RoundNearestTiesAway), round(a.hi, RoundNearestTiesAway))
end

# mid, diam, radius

# Compare pg. 64 of the IEEE 1788-2015 standard:
"""
    mid(a::Interval, α=0.5)

Find an intermediate point at a relative position `α`` in the interval `a`.
The default is the true midpoint at `α = 0.5`.

Assumes 0 ≤ α ≤ 1.

Warning: if the parameter `α = 0.5` is explicitely set, the behavior differs
from the default case if the provided `Interval` is not finite, since when
`α` is provided `mid` simply replaces `+∞` (respectively `-∞`) by `prevfloat(+∞)`
(respecively `nextfloat(-∞)`) for the computation of the intermediate point.
"""
function mid(a::Interval{T}, α) where T

    isempty(a) && return convert(T, NaN)

    lo = (a.lo == -∞ ? nextfloat(-∞) : a.lo)
    hi = (a.hi == +∞ ? prevfloat(+∞) : a.hi)

    midpoint = α * (hi - lo) + lo
    isfinite(midpoint) && return midpoint
    #= Fallback in case of overflow: hi - lo == +∞.
       This case can not be the default one as it does not pass several
       IEEE1788-2015 tests for small floats.
    =#
    return (1-α) * lo + α * hi
end

"""
    mid(a::Interval)

Find the midpoint of interval `a`.

For intervals of the form `[-∞, x]` or `[x, +∞]` where `x` is finite, return
respectively `nextfloat(-∞)` and `prevfloat(+∞)`. Note that it differs from the
behavior of `mid(a, α=0.5)`.
"""
function mid(a::Interval{T}) where T

    isempty(a) && return convert(T, NaN)
    isentire(a) && return zero(a.lo)

    a.lo == -∞ && return nextfloat(a.lo)
    a.hi == +∞ && return prevfloat(a.hi)

    midpoint = 0.5 * (a.lo + a.hi)
    isfinite(midpoint) && return midpoint
    #= Fallback in case of overflow: a.hi + a.lo == +∞ or a.hi + a.lo == -∞.
       This case can not be the default one as it does not pass several
       IEEE1788-2015 tests for small floats.
    =#
    return 0.5 * a.lo + 0.5 * a.hi
end

mid(a::Interval{Rational{T}}) where T = (1//2) * (a.lo + a.hi)

"""
    diam(a::Interval)

Return the diameter (length) of the `Interval` `a`.
"""
function diam(a::Interval{T}) where T<:Real
    isempty(a) && return convert(T, NaN)

    @round_up(a.hi - a.lo) # cf page 64 of IEEE1788
end

# Should `radius` this yield diam(a)/2? This affects other functions!
"""
    radius(a::Interval)

Return the radius of the `Interval` `a`, such that
`a ⊆ m ± radius`, where `m = mid(a)` is the midpoint.
"""
function radius(a::Interval)
    isempty(a) && return convert(eltype(a), NaN)
    m = mid(a)
    return max(m - a.lo, a.hi - m)
end

function radius(a::Interval{Rational{T}}) where T
    m = (a.lo + a.hi) / 2
    return max(m - a.lo, a.hi - m)
end

# cancelplus and cancelminus
"""
    cancelminus(a, b)

Return the unique interval `c` such that `b+c=a`.

See Section 12.12.5 of the IEEE-1788 Standard for
Interval Arithmetic.
"""
function cancelminus(a::Interval{T}, b::Interval{T}) where T<:Real
    (isempty(a) && (isempty(b) || !isunbounded(b))) && return emptyinterval(T)

    (isunbounded(a) || isunbounded(b) || isempty(b)) && return entireinterval(T)

    diam(a) < diam(b) && return entireinterval(T)

    c_lo = @round_down(a.lo - b.lo)
    c_hi = @round_up(a.hi - b.hi)

    c_lo > c_hi && return entireinterval(T)

    c_lo == Inf && return Interval(prevfloat(c_lo), c_hi)
    c_hi == -Inf && return Interval(c_lo, nextfloat(c_hi))

    a_lo = @round_down(b.lo + c_lo)
    a_hi = @round_up(b.hi + c_hi)

    if a_lo ≤ a.lo ≤ a.hi ≤ a_hi
        (nextfloat(a.hi) < a_hi || prevfloat(a.lo) > a_hi) &&
            return entireinterval(T)
        return Interval(c_lo, c_hi)
     end

    return entireinterval(T)
end
cancelminus(a::Interval, b::Interval) = cancelminus(promote(a, b)...)

"""
    cancelplus(a, b)

Returns the unique interval `c` such that `b-c=a`;
it is equivalent to `cancelminus(a, −b)`.
"""
cancelplus(a::Interval, b::Interval) = cancelminus(a, -b)

# midpoint-radius forms
midpoint_radius(a::Interval) = (mid(a), radius(a))

interval_from_midpoint_radius(midpoint, radius) = Interval(midpoint-radius, midpoint+radius)

isinteger(a::Interval) = (a.lo == a.hi) && isinteger(a.lo)

convert(::Type{Integer}, a::Interval) = isinteger(a) ?
        convert(Integer, a.lo) : throw(ArgumentError("Cannot convert $a to integer"))


using DiffEqBase, Catalyst, Plots, OrdinaryDiffEq, DataFrames, CSVFiles, LinearAlgebra
using ReactionNetworkImporters

# parameters
doplot = false
networkname = "Repressilator"
tf = 40000
nsteps = 400

# BNG simulation data
datadir  = joinpath(@__DIR__,"../data/repressilator")
fname    = joinpath(datadir, "Repressilator.net")
gdatfile = joinpath(datadir, "Repressilator.gdat")
print("getting gdat file...")
gdatdf = DataFrame(load(File{format"CSV"}(gdatfile), header_exists=true, spacedelim=true))
println("done")

# load the BNG reaction network in DiffEqBio
prnbng = loadrxnetwork(BNGNetwork(), fname)
rn = prnbng.rn; u0 = Float64[]; p = Float64[]
boprob = ODEProblem(rn, u0, (0.,tf), p)

# BNG simulation data testing
bngsol = gdatdf[!,:pTetR]

# note solvers run _much_ faster the second time
bsol = solve(boprob, Tsit5(), abstol=1e-12, reltol=1e-12, saveat=tf/nsteps);
@unpack pTetR = rn

if doplot
    plotlyjs()
    p1 = plot(gdatdf[!,:time], gdatdf[!,:pTetR], label=:BNGPTetR)
    plot!(p1, bsol.t, bsol[pTetR], label=:DEBPTetR)
    display(p1)
end

@test all(bsol.t .== gdatdf[!,:time])
@test all( abs.(gdatdf[!,:pTetR] - bsol[pTetR]) .< 1e-6 .* abs.(gdatdf[!,:pTetR]))


__precompile__()
module MobiusUtils

using Dierckx

export bisection, falseposition, interpolate

"""
Simple bisection method to solve nonlinear equation
f(x)=0for f: ℝ → ℝ  on the interval (a,b). The result
is the solution c such that f(c-ϵ/2)×f(c+ϵ/2) < 0, i.e.
is the centre of the interval of the length ϵ over
which f(x) changes the sign.
"""
function bisection(f,a,b,ϵ)
    fa = f(a)
    fb = f(b)
    while abs(a-b)≥ ϵ
        c = (a+b)/2
        fc = f(c)
        (a,fa) = fc*fb ≤ 0? (c,fc) : (a,fa)
        (b,fb) = fa*fc ≤ 0? (c,fc) : (b,fb)
    end
    (a+b)/2
end

"""
Algorithm of false position to solve nonlinear
equation f(x)=0 for f: ℝ → ℝ on the interval x∈(a,b)
with the correction for functions that change from
"almost flat" to "rapid growth" type (see Hamming for
    more details).
"""
function falseposition(f,a,b,ϵ)
    fa = f(a)
    fb = f(b)
    keptabefore = false
    kepta = false
    while abs(a-b)≥ ϵ
        c = a - fa*(b-a)/(fb-fa)
        fc = f(c)
        keptabefore = kepta
        if fc*fb ≤ 0
            a,fa = c,fc
            kepta = false
        end
        if fa*fc ≤ 0
            b,fb = c, fc
            kepta = true
        end
        if keptabefore && kepta
            fa = fa/2
        elseif (!keptabefore) && (!kepta)
            fb = fb/2
        end
    end
    a - fa*(b-a)/(fb-fa)
end

"""
Interpolates the data between the points (x,y). Produces
the function of x.

Parameters:
- **order**: order of interpolation(1-5), defualts to 1
- **outside**: what to do when argument is outside of
  interpolation region: "extrapolate" (default), "error",
  "zero", "nearest"
- **spline**: if `false` (default) returns a function, otherwise
  returns the spline object of Dierckx (can be used as a function
  but is not always recognised as such).
"""
function interpolate(x, y, order = 1, outside = "extrapolate", spline = false)
    spl = Spline1D(x, y, k = order, bc = outside)
    spine? spl : z -> spl(z)
end

end # module


#: Gradient operator via convolution

export gradient_operator, gradient_operator_batch

# Generic convolutional operator

struct RealConvolutionalOperator{T,N}<:AbstractLinearOperator{T,N,N}
    cdims::DenseConvDims
    stencil::AbstractArray{<:Real,N}
end

AbstractLinearOperators.domain_size(A::RealConvolutionalOperator) = NNlib.input_size(A.cdims)
AbstractLinearOperators.range_size(A::RealConvolutionalOperator) = NNlib.output_size(A.cdims)
AbstractLinearOperators.matvecprod(A::RealConvolutionalOperator{CT,N}, u::AbstractArray{T,N}) where {T<:Real,N,CT<:RealOrComplex{T}} = conv(u, A.stencil)
AbstractLinearOperators.matvecprod(A::RealConvolutionalOperator{T,N}, u::AbstractArray{T,N}) where {T<:Complex,N} = matvecprod(A, real(u))+1im*matvecprod(A, imag(u))
AbstractLinearOperators.matvecprod_adj(A::RealConvolutionalOperator{CT,N}, v::AbstractArray{T,N}) where {T<:Real,N,CT<:RealOrComplex{T}} = ∇conv_data(v, A.stencil, A.cdims)
AbstractLinearOperators.matvecprod_adj(A::RealConvolutionalOperator{T,N}, v::AbstractArray{T,N}) where {T<:Complex,N} = matvecprod_adj(A, real(v))+1im*matvecprod_adj(A, imag(v))

Flux.gpu(A::RealConvolutionalOperator{T,N}) where {T,N} = RealConvolutionalOperator{T,N}(A.cdims, gpu(A.stencil))
Flux.cpu(A::RealConvolutionalOperator{T,N}) where {T,N} = RealConvolutionalOperator{T,N}(A.cdims, cpu(A.stencil))

# Gradient operator

function gradient_operator(n::NTuple{dim,Int64}, h::NTuple{dim,S}; T::DataType=Float32) where {dim,S<:Real}

    # Gradient operator (w/out shaping)
    stencil = gradient_stencil(real(T).(h))
    cdims = DenseConvDims((n...,1,1), size(stencil))
    ∇_ = RealConvolutionalOperator{T,dim+2}(cdims, stencil)

    # Reshaping operator
    Rin  = reshape_operator(T, n, (n...,1,1))
    Rout = reshape_operator(T, ((n.-1)...,dim,1), ((n.-1)...,dim))

    return Rout*∇_*Rin

end

function gradient_operator_batch(n::NTuple{dim,Int64}, nc::Int64, nb::Int64, h::NTuple{dim,S}; T::DataType=Float32) where {dim,S<:Real}

    # Gradient operator (w/out shaping)
    stencil = gradient_stencil(real(T).(h))
    cdims = DenseConvDims((n...,1,nc*nb), size(stencil))
    ∇_ = RealConvolutionalOperator{T,dim+2}(cdims, stencil)

    # Reshaping operator
    Rin  = reshape_operator(T, (n...,nc,nb), (n...,1,nc*nb))
    Rout = reshape_operator(T, ((n.-1)...,dim,nc*nb), ((n.-1)...,dim*nc,nb))

    return Rout*∇_*Rin

end

# Gradient stencil utils

function gradient_stencil(h::NTuple{dim,T}) where {dim,T}
    k = tuple(2*ones(Int64, dim)...)
    stencil = zeros(T,k...,1,dim)
    for i = 1:dim
        idx = Vector{UnitRange{Int64}}(undef,dim); fill!(idx, 2:2)
        idx[i] = 1:2
        view(stencil, tuple(idx...)...,1,i)[1:2] .= [T(1)/h[i]; T(-1)/h[i]]
    end
    return stencil
end


input = open("../inputs/day03.txt") do f
    read(f, String)
end

function makeMap(inp)
    lines = split(inp, "\n")
    height = length(lines)
    width = length(lines[1])
    m = fill(' ', height, width)
    for y in 1:height
        for (x, c) in enumerate(lines[y])
            m[y,x] = c
        end
    end
    m
end

function runSlope(m, dy, dx)::Int64
    width = size(m, 2)
    height = size(m, 1)
    x = 1 + dx
    y = 1 + dy
    count = 0
    while y <= height
        if m[y,x] == '#'
            count += 1
        end
        x = mod(x + dx, 1:width)
        y += dy
    end
    count
end

m = makeMap(input)
println("Part 1: ", runSlope(m, 1, 3))

slopes = [(1, 1), (3, 1), (5, 1), (7, 1), (1, 2)]

println("Part 2: ", prod([runSlope(m, dy, dx) for (dx, dy) in slopes]))


import FEM.Nvec
import FEM.dNmatrix
import FEM.Jmatrix
import FEM.dNdxmatrix
import FEM.FENode2
import FEM.Point2

facts("FEM.Interpolator") do

context("FEM.Interpolator.LinTrigInterp") do
      interp = LinTrigInterp()

      node_1 = FENode2(1, [1.3, 0.5])
      node_2 = FENode2(2, [1.6, 1.0])
      node_3 = FENode2(3, [0.2, 0.9])

      loc_coords = Point2(0.5, 0.3)

      N = Nvec(interp, loc_coords)

      @fact norm(N - [0.5, 0.3, 1.0 - 0.5 - 0.3]) => roughly(0.0)

end # context

end # facts


# Ordinary metrics

###########################################################
#
#   Abstract metric types
#
###########################################################

abstract type UnionPreMetric <: PreMetric end

abstract type UnionSemiMetric <: SemiMetric end

abstract type UnionMetric <: Metric end

###########################################################
#
#   Metric types
#
###########################################################

struct Euclidean <: UnionMetric
    thresh::Float64
end

"""
    Euclidean([thresh])

Create a euclidean metric.

When computing distances among large numbers of points, it can be much
more efficient to exploit the formula

    (x-y)^2 = x^2 - 2xy + y^2

However, this can introduce roundoff error. `thresh` (which defaults
to 0) specifies the relative square-distance tolerance on `2xy`
compared to `x^2 + y^2` to force recalculation of the distance using
the more precise direct (elementwise-subtraction) formula.

# Example:
```julia
julia> x = reshape([0.1, 0.3, -0.1], 3, 1);

julia> pairwise(Euclidean(), x, x)
1×1 Array{Float64,2}:
 7.45058e-9

julia> pairwise(Euclidean(1e-12), x, x)
1×1 Array{Float64,2}:
 0.0
```
"""
Euclidean() = Euclidean(0)

struct WeightedEuclidean{W} <: UnionMetric
    weights::W
end

"""
    PeriodicEuclidean(L)

Create a Euclidean metric on a rectangular periodic domain (i.e., a torus or
a cylinder). Periods per dimension are contained in the vector `L`:
```math
\\sqrt{\\sum_i(\\min\\mod(|x_i - y_i|, p), p - \\mod(|x_i - y_i|, p))^2}.
```
For dimensions without periodicity put `Inf` in the respective component.

# Example
```jldoctest
julia> x, y, L = [0.0, 0.0], [0.75, 0.0], [0.5, Inf];

julia> evaluate(PeriodicEuclidean(L), x, y)
0.25
```
"""
struct PeriodicEuclidean{W} <: UnionMetric
    periods::W
end

struct SqEuclidean <: UnionSemiMetric
    thresh::Float64
end

"""
    SqEuclidean([thresh])

Create a squared-euclidean semi-metric. For the meaning of `thresh`,
see [`Euclidean`](@ref).
"""
SqEuclidean() = SqEuclidean(0)

struct WeightedSqEuclidean{W} <: UnionSemiMetric
    weights::W
end

struct Chebyshev <: UnionMetric end

struct Cityblock <: UnionMetric end
struct WeightedCityblock{W} <: UnionMetric
    weights::W
end

struct TotalVariation <: UnionMetric end
struct Jaccard <: UnionMetric end
struct RogersTanimoto <: UnionMetric end

struct Minkowski{T <: Real} <: UnionMetric
    p::T
end
struct WeightedMinkowski{W,T <: Real} <: UnionMetric
    weights::W
    p::T
end

struct Hamming <: UnionMetric end
struct WeightedHamming{W} <: UnionMetric
    weights::W
end

struct CosineDist <: UnionSemiMetric end
# CorrDist is excluded from `UnionMetrics`
struct CorrDist <: SemiMetric end
struct BrayCurtis <: UnionSemiMetric end

struct ChiSqDist <: UnionSemiMetric end
struct KLDivergence <: UnionPreMetric end
struct GenKLDivergence <: UnionPreMetric end

"""
    RenyiDivergence(α::Real)
    renyi_divergence(P, Q, α::Real)

Create a Rényi premetric of order α.

Rényi defined a spectrum of divergence measures generalising the
Kullback–Leibler divergence (see `KLDivergence`). The divergence is
not a semimetric as it is not symmetric. It is parameterised by a
parameter α, and is equal to Kullback–Leibler divergence at α = 1:

At α = 0, ``R_0(P | Q) = -log(sum_{i: p_i > 0}(q_i))``

At α = 1, ``R_1(P | Q) = sum_{i: p_i > 0}(p_i log(p_i / q_i))``

At α = ∞, ``R_∞(P | Q) = log(sup_{i: p_i > 0}(p_i / q_i))``

Otherwise ``R_α(P | Q) = log(sum_{i: p_i > 0}((p_i ^ α) / (q_i ^ (α - 1))) / (α - 1)``

# Example:
```jldoctest
julia> x = reshape([0.1, 0.3, 0.4, 0.2], 2, 2);

julia> pairwise(RenyiDivergence(0), x, x)
2×2 Array{Float64,2}:
 0.0  0.0
 0.0  0.0

julia> pairwise(Euclidean(2), x, x)
2×2 Array{Float64,2}:
 0.0       0.577315
 0.655407  0.0
```
"""
struct RenyiDivergence{T <: Real} <: UnionPreMetric
    p::T # order of power mean (order of divergence - 1)
    is_normal::Bool
    is_zero::Bool
    is_one::Bool
    is_inf::Bool

    function RenyiDivergence{T}(q) where {T}
        # There are four different cases:
        #   simpler to separate them out now, not over and over in eval_op()
        is_zero = q ≈ zero(T)
        is_one = q ≈ one(T)
        is_inf = isinf(q)

        # Only positive Rényi divergences are defined
        !is_zero && q < zero(T) && throw(ArgumentError("Order of Rényi divergence not legal, $(q) < 0."))

        new{T}(q - 1, !(is_zero || is_one || is_inf), is_zero, is_one, is_inf)
    end
end
RenyiDivergence(q::T) where {T} = RenyiDivergence{T}(q)

struct JSDivergence <: UnionSemiMetric end

struct SpanNormDist <: UnionSemiMetric end

# Deviations are handled separately from the other distances/divergences and
# are excluded from `UnionMetrics`
struct MeanAbsDeviation <: Metric end
struct MeanSqDeviation <: SemiMetric end
struct RMSDeviation <: Metric end
struct NormRMSDeviation <: PreMetric end

# Union types
const metrics = (Euclidean,SqEuclidean,PeriodicEuclidean,Chebyshev,Cityblock,TotalVariation,Minkowski,Hamming,Jaccard,RogersTanimoto,CosineDist,ChiSqDist,KLDivergence,RenyiDivergence,BrayCurtis,JSDivergence,SpanNormDist,GenKLDivergence)
const weightedmetrics = (WeightedEuclidean,WeightedSqEuclidean,WeightedCityblock,WeightedMinkowski,WeightedHamming)
const UnionMetrics = Union{UnionPreMetric,UnionSemiMetric,UnionMetric}

###########################################################
#
#  Implementations
#
###########################################################

parameters(::UnionPreMetric) = nothing
parameters(::UnionSemiMetric) = nothing
parameters(::UnionMetric) = nothing
parameters(d::PeriodicEuclidean) = d.periods
for dist in weightedmetrics
    @eval parameters(d::$dist) = d.weights
end

result_type(dist::UnionMetrics, ::Type{Ta}, ::Type{Tb}) where {Ta,Tb} =
    result_type(dist, _eltype(Ta), _eltype(Tb), parameters(dist))
result_type(dist::UnionMetrics, ::Type{Ta}, ::Type{Tb}, ::Nothing) where {Ta,Tb} =
    typeof(_evaluate(dist, oneunit(Ta), oneunit(Tb)))
result_type(dist::UnionMetrics, ::Type{Ta}, ::Type{Tb}, p) where {Ta,Tb} =
    typeof(_evaluate(dist, oneunit(Ta), oneunit(Tb), oneunit(_eltype(p))))

Base.@propagate_inbounds function _evaluate(d::UnionMetrics, a, b)
    _evaluate(d, a, b, parameters(d))
end

# breaks the implementation into eval_start, eval_op, eval_reduce and eval_end

Base.@propagate_inbounds function _evaluate(d::UnionMetrics, a, b, ::Nothing)
    @boundscheck if length(a) != length(b)
        throw(DimensionMismatch("first collection has length $(length(a)) which does not match the length of the second, $(length(b))."))
    end
    if length(a) == 0
        return zero(result_type(d, a, b))
    end
    s = eval_start(d, a, b)
    for (ai, bi) in zip(a, b)
        s = eval_reduce(d, s, eval_op(d, ai, bi))
    end
    return eval_end(d, s)
end
Base.@propagate_inbounds function _evaluate(d::UnionMetrics, a::AbstractArray, b::AbstractArray, ::Nothing)
    @boundscheck if length(a) != length(b)
        throw(DimensionMismatch("first array has length $(length(a)) which does not match the length of the second, $(length(b))."))
    end
    if length(a) == 0
        return zero(result_type(d, a, b))
    end
    begin
        s = eval_start(d, a, b)
        if (IndexStyle(a, b) === IndexLinear() && eachindex(a) == eachindex(b)) || axes(a) == axes(b)
            for I in eachindex(a, b)
                ai = a[I]
                bi = b[I]
                s = eval_reduce(d, s, eval_op(d, ai, bi))
            end
        else
            for (ai, bi) in zip(a, b)
                s = eval_reduce(d, s, eval_op(d, ai, bi))
            end
        end
        return eval_end(d, s)
    end
end

Base.@propagate_inbounds function _evaluate(d::UnionMetrics, a, b, p)
    @boundscheck if length(a) != length(b)
        throw(DimensionMismatch("first collection has length $(length(a)) which does not match the length of the second, $(length(b))."))
    end
    @boundscheck if length(a) != length(p)
        throw(DimensionMismatch("data collections have length $(length(a)) but parameters have length $(length(p))."))
    end
    if length(a) == 0
        return zero(result_type(d, a, b))
    end
    s = eval_start(d, a, b)
    for (ai, bi, pi) in zip(a, b, p)
        s = eval_reduce(d, s, eval_op(d, ai, bi, pi))
    end
    return eval_end(d, s)
end
Base.@propagate_inbounds function _evaluate(d::UnionMetrics, a::AbstractArray, b::AbstractArray, p::AbstractArray)
    @boundscheck if length(a) != length(b)
        throw(DimensionMismatch("first array has length $(length(a)) which does not match the length of the second, $(length(b))."))
    end
    @boundscheck if length(a) != length(p)
        throw(DimensionMismatch("arrays have length $(length(a)) but parameters have length $(length(p))."))
    end
    if length(a) == 0
        return zero(result_type(d, a, b))
    end
    begin
        s = eval_start(d, a, b)
        if (IndexStyle(a, b, p) === IndexLinear() && eachindex(a) == eachindex(b) == eachindex(p)) ||
                axes(a) == axes(b) == axes(p)
            for I in eachindex(a, b, p)
                ai = a[I]
                bi = b[I]
                pi = p[I]
                s = eval_reduce(d, s, eval_op(d, ai, bi, pi))
            end
        else
            for (ai, bi, pi) in zip(a, b, p)
                s = eval_reduce(d, s, eval_op(d, ai, bi, pi))
            end
        end
        return eval_end(d, s)
    end
end

_evaluate(dist::UnionMetrics, a::Number, b::Number, ::Nothing) = eval_end(dist, eval_op(dist, a, b))
function _evaluate(dist::UnionMetrics, a::Number, b::Number, p)
    length(p) != 1 && throw(DimensionMismatch("inputs are scalars but parameters have length $(length(p))."))
    eval_end(dist, eval_op(dist, a, b, first(p)))
end

eval_start(d::UnionMetrics, a, b) = _eval_start(d, _eltype(a), _eltype(b))
_eval_start(d::UnionMetrics, ::Type{Ta}, ::Type{Tb}) where {Ta,Tb} =
    _eval_start(d, Ta, Tb, parameters(d))
_eval_start(d::UnionMetrics, ::Type{Ta}, ::Type{Tb}, ::Nothing) where {Ta,Tb} =
    zero(typeof(eval_op(d, oneunit(Ta), oneunit(Tb))))
_eval_start(d::UnionMetrics, ::Type{Ta}, ::Type{Tb}, p) where {Ta,Tb} =
    zero(typeof(eval_op(d, oneunit(Ta), oneunit(Tb), oneunit(_eltype(p)))))
eval_reduce(::UnionMetrics, s1, s2) = s1 + s2
eval_end(::UnionMetrics, s) = s

for M in (metrics..., weightedmetrics...)
    @eval (dist::$M)(a, b) = _evaluate(dist, a, b, parameters(dist))
end

# Euclidean
eval_op(::Euclidean, ai, bi) = abs2(ai - bi)
eval_end(::Euclidean, s) = sqrt(s)
const euclidean = Euclidean()

# Weighted Euclidean
eval_op(::WeightedEuclidean, ai, bi, wi) = abs2(ai - bi) * wi
eval_end(::WeightedEuclidean, s) = sqrt(s)
weuclidean(a, b, w) = WeightedEuclidean(w)(a, b)

# PeriodicEuclidean
function eval_op(::PeriodicEuclidean, ai, bi, p)
    s1 = abs(ai - bi)
    s2 = mod(s1, p)
    s3 = min(s2, p - s2)
    abs2(s3)
end
eval_end(::PeriodicEuclidean, s) = sqrt(s)
peuclidean(a, b, p) = PeriodicEuclidean(p)(a, b)

# SqEuclidean
eval_op(::SqEuclidean, ai, bi) = abs2(ai - bi)
const sqeuclidean = SqEuclidean()

# Weighted Squared Euclidean
eval_op(::WeightedSqEuclidean, ai, bi, wi) = abs2(ai - bi) * wi
wsqeuclidean(a, b, w) = WeightedSqEuclidean(w)(a, b)

# Cityblock
eval_op(::Cityblock, ai, bi) = abs(ai - bi)
const cityblock = Cityblock()

# Weighted City Block
eval_op(::WeightedCityblock, ai, bi, wi) = abs((ai - bi) * wi)
wcityblock(a, b, w) = WeightedCityblock(w)(a, b)

# Total variation
eval_op(::TotalVariation, ai, bi) = abs(ai - bi)
eval_end(::TotalVariation, s) = s / 2
const totalvariation = TotalVariation()

# Chebyshev
eval_op(::Chebyshev, ai, bi) = abs(ai - bi)
eval_reduce(::Chebyshev, s1, s2) = max(s1, s2)
# if only NaN, will output NaN
Base.@propagate_inbounds eval_start(::Chebyshev, a, b) = abs(first(a) - first(b))
const chebyshev = Chebyshev()

# Minkowski
eval_op(dist::Minkowski, ai, bi) = abs(ai - bi)^dist.p
eval_end(dist::Minkowski, s) = s^(1 / dist.p)
minkowski(a, b, p::Real) = Minkowski(p)(a, b)

# Weighted Minkowski
eval_op(dist::WeightedMinkowski, ai, bi, wi) = abs(ai - bi)^dist.p * wi
eval_end(dist::WeightedMinkowski, s) = s^(1 / dist.p)
wminkowski(a, b, w, p::Real) = WeightedMinkowski(w, p)(a, b)

# Hamming
result_type(::Hamming, ::Type, ::Type) = Int # fallback for Hamming
eval_start(d::Hamming, a, b) = 0
eval_op(::Hamming, ai, bi) = ai != bi ? 1 : 0
const hamming = Hamming()

# WeightedHamming
eval_op(::WeightedHamming, ai, bi, wi) = ai != bi ? wi : zero(eltype(wi))
whamming(a, b, w) = WeightedHamming(w)(a, b)

# Cosine dist
eval_start(dist::CosineDist, a, b) =
    zero.(typeof.(eval_op(dist, oneunit(_eltype(a)), oneunit(_eltype(b)))))
eval_op(::CosineDist, ai, bi) = ai * bi, ai * ai, bi * bi
function eval_reduce(::CosineDist, s1, s2)
    a1, b1, c1 = s1
    a2, b2, c2 = s2
    return a1 + a2, b1 + b2, c1 + c2
end
function eval_end(::CosineDist, s)
    ab, a2, b2 = s
    max(1 - ab / (sqrt(a2) * sqrt(b2)), 0)
end
const cosine_dist = CosineDist()

# CorrDist
_centralize(x) = x .- mean(x)
(::CorrDist)(a, b) = CosineDist()(_centralize(a), _centralize(b))
(::CorrDist)(a::Number, b::Number) = CosineDist()(zero(mean(a)), zero(mean(b)))
const corr_dist = CorrDist()

# ChiSqDist
eval_op(::ChiSqDist, ai, bi) = (d = abs2(ai - bi) / (ai + bi); ifelse(ai != bi, d, zero(d)))
const chisq_dist = ChiSqDist()

# KLDivergence
eval_op(dist::KLDivergence, ai, bi) =
    ai > 0 ? ai * log(ai / bi) : zero(eval_op(dist, oneunit(ai), bi))
const kl_divergence = KLDivergence()

# GenKLDivergence
eval_op(dist::GenKLDivergence, ai, bi) =
    ai > 0 ? ai * log(ai / bi) - ai + bi : oftype(eval_op(dist, oneunit(ai), bi), bi)
const gkl_divergence = GenKLDivergence()

# RenyiDivergence
Base.@propagate_inbounds function eval_start(::RenyiDivergence, a, b)
    T = promote_type(_eltype(a), _eltype(b))
    zero(T), zero(T), T(sum(a)), T(sum(b))
end

function eval_op(dist::RenyiDivergence, ai::T, bi::T) where {T <: Real}
    if ai == zero(T)
        return zero(T), zero(T), zero(T), zero(T)
    elseif dist.is_normal
        return ai, ai * ((ai / bi)^dist.p), zero(T), zero(T)
    elseif dist.is_zero
        return ai, bi, zero(T), zero(T)
    elseif dist.is_one
        return ai, ai * log(ai / bi), zero(T), zero(T)
    else # otherwise q = ∞
        return ai, ai / bi, zero(T), zero(T)
    end
end

function eval_reduce(dist::RenyiDivergence,
                                               s1::Tuple{T,T,T,T},
                                               s2::Tuple{T,T,T,T}) where {T <: Real}
    if dist.is_inf
        if s1[1] == zero(T)
            return (s2[1], s2[2], s1[3], s1[4])
        elseif s2[1] == zero(T)
            return s1
        else
            return s1[2] > s2[2] ? s1 : (s2[1], s2[2], s1[3], s1[4])
        end
    else
        return s1[1] + s2[1], s1[2] + s2[2], s1[3], s1[4]
    end
end

function eval_end(dist::RenyiDivergence, s::Tuple{T,T,T,T}) where {T <: Real}
    if dist.is_zero || dist.is_normal
        log(s[2] / s[1]) / dist.p + log(s[4] / s[3])
    elseif dist.is_one
        return s[2] / s[1] + log(s[4] / s[3])
    else # q = ∞
        log(s[2]) + log(s[4] / s[3])
    end
end

renyi_divergence(a, b, q::Real) = RenyiDivergence(q)(a, b)
# Combine docs with RenyiDivergence. Fetching the docstring with @doc causes
# problems during package compilation; see
# https://github.com/JuliaLang/julia/issues/31640
let docstring = Base.Docs.getdoc(RenyiDivergence)
    @doc docstring renyi_divergence
end

# JSDivergence

function eval_op(::JSDivergence, ai::T, bi::T) where {T}
    u = (ai + bi) / 2
    ta = ai > 0 ? ai * log(ai) / 2 : zero(log(one(T)))
    tb = bi > 0 ? bi * log(bi) / 2 : zero(log(one(T)))
    tu = u > 0 ? u * log(u) : zero(log(one(T)))
    ta + tb - tu
end
const js_divergence = JSDivergence()

# SpanNormDist

result_type(dist::SpanNormDist, ::Type{Ta}, ::Type{Tb}) where {Ta,Tb} =
    typeof(eval_op(dist, oneunit(Ta), oneunit(Tb)))
Base.@propagate_inbounds function eval_start(::SpanNormDist, a, b)
    d = first(a) - first(b)
    return d, d
end
eval_op(::SpanNormDist, ai, bi)  = ai - bi
function eval_reduce(::SpanNormDist, s1, s2)
    min_d, max_d = s1
    if s2 > max_d
        max_d = s2
    elseif s2 < min_d
        min_d = s2
    end
    return min_d, max_d
end

eval_end(::SpanNormDist, s) = s[2] - s[1]
(::SpanNormDist)(a::Number, b::Number) = zero(promote_type(typeof(a), typeof(b)))
const spannorm_dist = SpanNormDist()

# Jaccard

eval_start(dist::Jaccard, a, b) =
    zero.(typeof.(eval_op(dist, oneunit(_eltype(a)), oneunit(_eltype(b)))))
function eval_op(::Jaccard, s1, s2)
    abs_m = abs(s1 - s2)
    abs_p = abs(s1 + s2)
    abs_p - abs_m, abs_p + abs_m
end
function eval_reduce(::Jaccard, s1, s2)
    a = s1[1] + s2[1]
    b = s1[2] + s2[2]
    a, b
end
function eval_end(::Jaccard, a)
    v = 1 - (a[1] / a[2])
    return v
end
const jaccard = Jaccard()

# BrayCurtis

eval_start(dist::BrayCurtis, a, b) =
    zero.(typeof.(eval_op(dist, oneunit(_eltype(a)), oneunit(_eltype(b)))))
function eval_op(::BrayCurtis, s1, s2)
    abs_m = abs(s1 - s2)
    abs_p = abs(s1 + s2)
    abs_m, abs_p
end
function eval_reduce(::BrayCurtis, s1, s2)
    a = s1[1] + s2[1]
    b = s1[2] + s2[2]
    a, b
end
function eval_end(::BrayCurtis, a)
    v = a[1] / a[2]
    return v
end
const braycurtis = BrayCurtis()

# Tanimoto

eval_start(::RogersTanimoto, _, _) = 0, 0, 0, 0
function eval_op(::RogersTanimoto, s1, s2)
    tt = s1 && s2
    tf = s1 && !s2
    ft = !s1 && s2
    ff = !s1 && !s2
    tt, tf, ft, ff
end
function eval_reduce(::RogersTanimoto, s1, s2)
    begin
        a = s1[1] + s2[1]
        b = s1[2] + s2[2]
        c = s1[3] + s2[3]
        d = s1[4] + s1[4]
    end
    a, b, c, d
end
function eval_end(::RogersTanimoto, a)
    numerator = 2(a[2] + a[3])
    denominator = a[1] + a[4] + 2(a[2] + a[3])
    numerator / denominator
end
const rogerstanimoto = RogersTanimoto()

# Deviations

(::MeanAbsDeviation)(a, b) = cityblock(a, b) / length(a)
const meanad = MeanAbsDeviation()

(::MeanSqDeviation)(a, b) = sqeuclidean(a, b) / length(a)
const msd = MeanSqDeviation()

(::RMSDeviation)(a, b) = sqrt(MeanSqDeviation()(a, b))
const rmsd = RMSDeviation()

function (::NormRMSDeviation)(a, b)
    amin, amax = extrema(a)
    return RMSDeviation()(a, b) / (amax - amin)
end
const nrmsd = NormRMSDeviation()

###########################################################
#
#  Special method
#
###########################################################

# SqEuclidean/Euclidean
function _pairwise!(r::AbstractMatrix, dist::Union{SqEuclidean,Euclidean},
                    a::AbstractMatrix, b::AbstractMatrix)
    require_one_based_indexing(r, a, b)
    m, na, nb = get_pairwise_dims(r, a, b)
    # the following checks if a'*b can be stored in r directly, it fails for complex eltypes
    inplace = promote_type(eltype(r), typeof(oneunit(eltype(a))'oneunit(eltype(b)))) === eltype(r)
    R = inplace ? mul!(r, a', b) : a'b
    sa2 = sum(abs2, a, dims=1)
    sb2 = sum(abs2, b, dims=1)
    z² = zero(real(eltype(R)))
    if dist.thresh <= 0
        # If there's no chance of triggering the threshold, we can use @simd
        for j = 1:nb
            sb = sb2[j]
            for i = 1:na
                r[i, j] = eval_end(dist, (max(sa2[i] + sb - 2real(R[i, j]), z²)))
            end
        end
    else
        for j = 1:nb
            sb = sb2[j]
            for i = 1:na
                selfterms = sa2[i] + sb
                v = max(selfterms - 2real(R[i, j]), z²)
                if v < dist.thresh * selfterms
                    # The distance is likely to be inaccurate, recalculate directly
                    # This reflects the following:
                    #   while sqrt(x+ϵ) ≈ sqrt(x) + O(ϵ/sqrt(x)) when |x| >> ϵ,
                    #         sqrt(x+ϵ) ≈ O(sqrt(ϵ))             otherwise.
                    v = zero(v)
                    for k = 1:m
                        v += abs2(a[k, i] - b[k, j])
                    end
                end
                r[i, j] = eval_end(dist, v)
            end
        end
    end
    r
end

function _pairwise!(r::AbstractMatrix, dist::Union{SqEuclidean,Euclidean}, a::AbstractMatrix)
    require_one_based_indexing(r, a)
    m, n = get_pairwise_dims(r, a)
    inplace = promote_type(eltype(r), typeof(oneunit(eltype(a))'oneunit(eltype(a)))) === eltype(r)
    # the following checks if a'*b can be stored in r directly, it fails for complex eltypes
    R = inplace ? mul!(r, a', a) : a'a
    sa2 = sum(abs2, a, dims=1)
    safe = dist.thresh <= 0
    z² = zero(real(eltype(R)))
    for j = 1:n
        for i = 1:(j - 1)
            r[i, j] = r[j, i]
        end
        r[j, j] = zero(eltype(r))
        sa2j = sa2[j]
        if safe
            for i = (j + 1):n
                r[i, j] = eval_end(dist, (max(sa2[i] + sa2j - 2real(R[i, j]), z²)))
            end
        else
            for i = (j + 1):n
                selfterms = sa2[i] + sa2j
                v = max(selfterms - 2real(R[i, j]), z²)
                if v < dist.thresh * selfterms
                    v = zero(v)
                    for k = 1:m
                        v += abs2(a[k, i] - a[k, j])
                    end
                end
                r[i, j] = eval_end(dist, v)
            end
        end
    end
    r
end

# Weighted SqEuclidean/Euclidean
function _pairwise!(r::AbstractMatrix, dist::Union{WeightedSqEuclidean,WeightedEuclidean},
                    a::AbstractMatrix, b::AbstractMatrix)
    require_one_based_indexing(r, a, b)
    w = dist.weights
    m, na, nb = get_pairwise_dims(length(w), r, a, b)

    sa2 = wsumsq_percol(w, a)
    sb2 = wsumsq_percol(w, b)
    # the following checks if a'*b can be stored in r directly, it fails for complex eltypes
    inplace = promote_type(eltype(r), typeof(oneunit(eltype(a))'oneunit(eltype(b)))) === eltype(r)
    R = inplace ? mul!(r, a', w .* b) : a'*Diagonal(w)*b
    z² = zero(real(eltype(R)))
    for j = 1:nb
        for i = 1:na
            r[i, j] = eval_end(dist, max(sa2[i] + sb2[j] - 2real(R[i, j]), z²))
        end
    end
    r
end
function _pairwise!(r::AbstractMatrix, dist::Union{WeightedSqEuclidean,WeightedEuclidean},
                    a::AbstractMatrix)
    require_one_based_indexing(r, a)
    w = dist.weights
    m, n = get_pairwise_dims(length(w), r, a)

    sa2 = wsumsq_percol(w, a)
    # the following checks if a'*b can be stored in r directly, it fails for complex eltypes
    inplace = promote_type(eltype(r), typeof(oneunit(eltype(a))'oneunit(eltype(a)))) === eltype(r)
    R = inplace ? mul!(r, a', w .* a) : a'*Diagonal(w)*a
    z² = zero(real(eltype(R)))

    for j = 1:n
        for i = 1:(j - 1)
            r[i, j] = r[j, i]
        end
        r[j, j] = zero(eltype(r))
        for i = (j + 1):n
            r[i, j] = eval_end(dist, max(sa2[i] + sa2[j] - 2real(R[i, j]), z²))
        end
    end
    r
end

# MeanSqDeviation, RMSDeviation, NormRMSDeviation
function _pairwise!(r::AbstractMatrix, dist::MeanSqDeviation, a::AbstractMatrix, b::AbstractMatrix)
    _pairwise!(r, SqEuclidean(), a, b)
    # TODO: Replace by rdiv!(r, size(a, 1)) once julia compat ≥v1.2
    s = size(a, 1)
    for I in eachindex(r)
        r[I] /= s
    end
    return r
end
_pairwise!(r::AbstractMatrix, dist::RMSDeviation, a::AbstractMatrix, b::AbstractMatrix) =
    sqrt!(_pairwise!(r, MeanSqDeviation(), a, b))
function _pairwise!(r::AbstractMatrix, dist::NormRMSDeviation, a::AbstractMatrix, b::AbstractMatrix)
    _pairwise!(r, RMSDeviation(), a, b)
    @views for (i, j) in zip(axes(r, 1), axes(a, 2))
        amin, amax = extrema(a[:,j])
        r[i,:] ./= amax - amin
    end
    return r
end

function _pairwise!(r::AbstractMatrix, dist::MeanSqDeviation, a::AbstractMatrix)
    _pairwise!(r, SqEuclidean(), a)
    # TODO: Replace by rdiv!(r, size(a, 1)) once julia compat ≥v1.2
    s = size(a, 1)
    for I in eachindex(r)
        r[I] /= s
    end
    return r
end
_pairwise!(r::AbstractMatrix, dist::RMSDeviation, a::AbstractMatrix) =
    sqrt!(_pairwise!(r, MeanSqDeviation(), a))
function _pairwise!(r::AbstractMatrix, dist::NormRMSDeviation, a::AbstractMatrix)
    _pairwise!(r, RMSDeviation(), a)
    @views for (i, j) in zip(axes(r, 1), axes(a, 2))
        amin, amax = extrema(a[:,j])
        r[i,:] ./= amax - amin
    end
    return r
end

# CosineDist

function _pairwise!(r::AbstractMatrix, ::CosineDist, a::AbstractMatrix, b::AbstractMatrix)
    require_one_based_indexing(r, a, b)
    m, na, nb = get_pairwise_dims(r, a, b)
    inplace = promote_type(eltype(r), typeof(oneunit(eltype(a))'oneunit(eltype(b)))) === eltype(r)
    R = inplace ? mul!(r, a', b) : a'b
    ra = norm_percol(a)
    rb = norm_percol(b)
    for j = 1:nb
        for i = 1:na
            r[i, j] = max(1 - R[i, j] / (ra[i] * rb[j]), 0)
        end
    end
    r
end
function _pairwise!(r::AbstractMatrix, ::CosineDist, a::AbstractMatrix)
    require_one_based_indexing(r, a)
    m, n = get_pairwise_dims(r, a)
    inplace = promote_type(eltype(r), typeof(oneunit(eltype(a))'oneunit(eltype(a)))) === eltype(r)
    R = inplace ? mul!(r, a', a) : a'a
    ra = norm_percol(a)
    for j = 1:n
        for i = 1:(j - 1)
            r[i, j] = r[j, i]
        end
        r[j, j] = zero(eltype(r))
        for i = j + 1:n
            r[i, j] = max(1 - R[i, j] / (ra[i] * ra[j]), 0)
        end
    end
    r
end

# CorrDist
# This part of codes is accelerated because:
# 1. It calls the accelerated `_pairwise` specilization for CosineDist
# 2. pre-calculated `_centralize_colwise` avoids four times of redundant computations
#    of `_centralize` -- ~4x speed up
_centralize_colwise(x::AbstractMatrix) = x .- mean(x, dims=1)
_pairwise!(r::AbstractMatrix, ::CorrDist, a::AbstractMatrix, b::AbstractMatrix) =
    _pairwise!(r, CosineDist(), _centralize_colwise(a), _centralize_colwise(b))
_pairwise!(r::AbstractMatrix, ::CorrDist, a::AbstractMatrix) =
    _pairwise!(r, CosineDist(), _centralize_colwise(a))


using PhysicalOptics
using Test
using FFTW
using Random
using SpecialFunctions

 # fix seed for reproducibility
Random.seed!(42)

tests = [
    "apertures",
    "convolutions",
    "lenses",
    "physical_conversions",
    "propagations",
    "psf",
    "resolution_equations",
    "utils",
    "utils_view",
]

@testset "Point test" begin
    p = Point(1.0, 2.0, 3.0)

    @test typeof(p) === Point{Float64}
    @test p.x == 1.0
    @test p.y == 2.0
    @test p.z == 3.0

    p = Point(1, 2, 3)
    @test typeof(p) === Point{Int64}
    @test p.x == 1
    @test p.y == 2
    @test p.z == 3

end

for t in tests
    @testset "$(t)" begin
        include("$t.jl")
    end
end


#-----------------------------------------------------------------------# sweepmatrix
"""
`sweepmatrix(x, y, addint = true)`

`sweepmatrix(x, y, wts, addint = true)`

Create the matrix [x y]' * [x y] for sweeping, optionally with weights
"""
function sweepmatrix{T <: LinAlg.BlasFloat}(x::AMat{T}, y::AVec{T}, addint::Bool = true)
    n, p = size(x)
    d = p + 1 + addint
    A = zeros(d, d)

    rng = (1:p) + addint
    # operations: x'x, x'y, y'y
    BLAS.syrk!('U', 'T', 1 / n, x, 0.0, view(A, rng, rng))
    BLAS.gemv!('T', 1 / n, x, y, 0.0, view(A, rng, d))
    A[end, end] = dot(y, y) / n

    if addint
        #operations: 1'1, 1'x, 1'y,
        A[1, 1] = 1.0
        A[1, 2:end - 1] = mean(x, 1)
        A[1, end] = mean(y)
    end
    A
end
function sweepmatrix{T <: LinAlg.BlasFloat}(x::AMat{T}, y::AVec{T}, wts::AVec{T}, addint::Bool = true)
    n, p = size(x)
    d = p + 1 + addint
    wts = wts / sum(wts)
    A = zeros(d, d)
    W = Diagonal(sqrt(wts))
    Wx = W * x
    Wy = W * y

    rng = (1:p) + addint
    # operations: x'Wx, x'Wy, y'Wy
    BLAS.syrk!('U', 'T', 1 / n, Wx, 0.0, view(A, rng, rng))
    BLAS.gemv!('T', 1 / n, Wx, Wy, 0.0, view(A, rng, d))
    A[end, end] = dot(y, Wy) / n

    if addint
        #operations: 1'Wx, 1'Wy,
        A[1, 2:end - 1] = mean(Wx, 1)
        A[1, end] = mean(y, StatsBase.WeightVec(wts))
    end
    A
end

"Add a ridge penalty to the sweepmatrix"
function addridge!{T <: LinAlg.BlasFloat}(A::AMat{T}, λ::T, intercept::Bool)
    rng = (1 + intercept):(size(A, 1) - 1)
    for i in rng
        A[i, i] *= λ
    end
    A
end

# Unpenalized or Ridge-regularized regression
function sweepreg{T <: LinAlg.BlasFloat}(x::AMat{T}, y::AVec{T}, addint::Bool = true; λ::T = zero(T))
    A = sweepmatrix(x, y, addint)
    if λ > zero(T)
        addridge!(A, λ, addint)
    end
    SweepOperator.sweep!(A, 1:size(A, 1) - 1)
    β = A[1:end-1, end]
    SSE = A[end, end]
    minusXTXinv = A[1:end-1, 1:end-1]
    β, SSE, minusXTXinv
end


using CairoMakie
using DiffEqSensitivity
using Zygote
using ForwardDiff
using LinearAlgebra
using DifferentialEquations

"""
  vanderpol!(du,u,μ,t)

Van der Pol oscillator dynamics for DifferentialEquations.jl
inline
"""
function vanderpol!(du,u,μ,t)
    x,y = u
    du[1] = μ * (x - x^3/3 - y)
    du[2] = x/μ
end

"""
  vanderpol(u,μ,t)

Van der Pol oscillator dynamics for DifferentialEquations.jl
not inline
"""
function vanderpol(u,μ,t)
    du = zeros(2)
    vanderpol!(du,u,μ,t)
    return du
end

"""
  phasemap(u, positions)

Applies to Van der Pol oscillator.

Assigns a point `u` in the statespace to a position on the
limit cycle of the oscillator. Given a list of `positions`
spaced equally in time along the limit cycle, this function
simulates the system from initial condition `u` then determines
the closest element in `positions`

Returns: The index of the closest element.
"""
function phasemap(u, positions)
    μ = 1.0
    tspan = (0.0,40.0)

    prob = ODEProblem(vanderpol!,u,tspan,μ)

    sol = solve(prob)
    final = sol.u[end]
    ~, out = findmin(sum(abs2,positions .- final,dims=1))
    return out[2]
end

"""
  simystem(u)

Compute the value of the Van der Pol oscillator initialized at
`u` 20 units of time in the future.
"""
function simsystem(u)
    μ = 1.0
    tspan = (0.0,20.0)

    prob = ODEProblem(vanderpol!,u,tspan,μ)

    sol = solve(prob,rtol=1e-6,atol=1e-6,sensealg=QuadratureAdjoint())
    return sol.u[end]
end

"""
  informativeVec(u)

Find the non-zero eigenvalue of dφ₂₀ in the Van Der Pol
oscillator around state space point `u`.
"""
function informativeVec(u)
    dx = ForwardDiff.jacobian(simsystem,u)
    ~, Σ, Vt = svd(dx)
    # Force counterclockwise
    test_vec = [u[2], -u[1]]
    if dot(test_vec,Vt[1,:]) > 0
        return Vt[1,:]
    else
        return -Vt[1,:]
    end
end


using Permutations, StaticArrays, FileIO

println("Building of HyperoctahedralGroup.jl started.")

const xyp = SA[0 -1 0 0; 1 0 0 0; 0 0 1 0; 0 0 0 1]
const xym = SA[0 1 0 0; -1 0 0 0; 0 0 1 0; 0 0 0 1]

const xzp = SA[0 0 -1 0; 0 1 0 0; 1 0 0 0; 0 0 0 1]
const xzm = SA[0 0 1 0; 0 1 0 0; -1 0 0 0; 0 0 0 1]

const xwp = SA[0 0 0 -1; 0 1 0 0; 0 0 1 0; 1 0 0 0]
const xwm = SA[0 0 0 1; 0 1 0 0; 0 0 1 0; -1 0 0 0]

const yzp = SA[1 0 0 0; 0 0 -1 0; 0 1 0 0; 0 0 0 1]
const yzm = SA[1 0 0 0; 0 0 1 0; 0 -1 0 0; 0 0 0 1]

const ywp = SA[1 0 0 0; 0 0 0 -1; 0 0 1 0; 0 1 0 0]
const ywm = SA[1 0 0 0; 0 0 0 1; 0 0 1 0; 0 -1 0 0]

const zwp = SA[1 0 0 0; 0 1 0 0; 0 0 0 -1; 0 0 1 0]
const zwm = SA[1 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 -1 0]

function create_elements()
    elements = [SA[1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1]]
    counter = 0
    while length(elements) < 192 || counter == 16
        if true
            new_elements = [xyp*element for element in elements]
            for (i,e) in enumerate(new_elements)
                if e ∉ elements
                    push!(elements,e)
                end
            end
        end

        if true
            new_elements = [xzp*element for element in elements]
            for (i,e) in enumerate(new_elements)
                if e ∉ elements
                    push!(elements,e)
                end
            end
        end

        if true
            new_elements = [xwp*element for element in elements]
            for (i,e) in enumerate(new_elements)
                if e ∉ elements
                    push!(elements,e)
                end
            end
        end

        if false
            new_elements = [yzp*element for element in elements]
            for (i,e) in enumerate(new_elements)
                if e ∉ elements
                    push!(elements,e)
                end
            end
        end

        if false
            new_elements = [ywp*element for element in elements]
            for (i,e) in enumerate(new_elements)
                if e ∉ elements
                    push!(elements,e)
                end
            end
        end

        if false
            new_elements = [zwp*element for element in elements]
            for (i,e) in enumerate(new_elements)
                if e ∉ elements
                    push!(elements,e)
                end
            end
        end
        counter += 1
    end
    return elements
end

function verifyelements(elements)
    if length(unique(elements)) != 192
        return false
    end
    ones_1x4 = SA[1 1 1 1]
    ones_4x1 = reshape(ones_1x4,(4,1))
    for el in elements
        if size(el) != (4,4)
            return false
        end
        if !issubset(unique(el),[0,1,-1])
            return false
        end
        if map(abs,sum(el,dims=1)) != ones_1x4
            return false
        end
        if map(abs,sum(el,dims=2)) != ones_4x1
            return false
        end
    end
    return true
end

let elements_file = joinpath(@__DIR__,"elements.jld2")
    timestarted = time()
    verified = false
    howstring = "loaded"
    if isfile(elements_file)
        println("Loading group elements.")
        elements = load(elements_file,"elements")
        verified = verifyelements(elements)
    end
    if !verified
        howstring = "created"
        println("Elements could not be loaded. Creating group elements.")
        global elements = create_elements()
        save(elements_file,"elements",elements)
    end
    timeelapsed = time()-timestarted
    println("All elements $howstring after $timeelapsed s.")
end

function write_elements()
    f = open(joinpath(@__DIR__,"elements.jl"),"w")
    write(f,"elements = [")
    for el in elements[1:end-1]
        write(f,"SA")
        print(f,el)
        write(f,",\n")
    end
    write(f,"SA")
    print(f,last(elements))
    write(f,"]\n")
    close(f)
end

write_elements()

function order(e)
    r = e*e
    n = 1
    while r != e
        r = r*e
        n+=1
    end
    return n
end

function group_act_slow(a::Integer,b::Integer)
    if 192 >= a > 0 && 192 >= b > 0
        for i = 1:192
            if elements[i] == elements[a]*elements[b]
                return typeof(a)(i)
            end
        end

        println(elements[a])
        println(elements[b])
        println(elements[a]*elements[b])
        error("Not found? How is that even possible?")
    else
        error("Not defined.")
    end
    return zero(typeof(a))
end

const g = [group_act_slow(a,b) for a in 1:192, b in 1:192]

function ⊗(a::Integer,b::Integer)
    if 192 >= a > 0 && 192 >= b > 0
        return g[a,b]
    else
        error("$a o $b Not defined.")
    end
end

function insertsorted!(set,element;startindex = 1)
    if !issorted(set)
        error("Array is not sorted.")
    end
    i = startindex
    while  i <= length(set) && element > set[i]
        i += 1
    end
    i > length(set) ? push!(set, element) : insert!(set, i, element)
    return i
end

function complete!(set)
    added = -1
    queue = vec([(a,b) for a in set, b in set])
    while added != 0
        added = 0
        newElements = empty(set)
        for i in length(queue):-1:1
            a,b = queue[i]
            if ⊗(a,b) ∉ set
                push!(newElements,⊗(a,b))
                added = added+1
            end
            pop!(queue)
        end
        newElements = sort(unique(newElements))
        for a in newElements
            for b in set
                push!(queue,(a,b))
                push!(queue,(b,a))
            end
            for b in newElements
                push!(queue,(a,b))
            end
        end

        i = 1
        for a in newElements
            i = insertsorted!(set,a, startindex = i)
        end
    end
    return set
end

function completewith!(set,toadd)
    toadd = setdiff(toadd,set)
    isempty(toadd) && return set
    added = -1
    queue = vcat(vec([(a,b) for a in set, b in toadd]),vec([(b,a) for a in set, b in toadd]),vec([(a,b) for a in toadd, b in toadd]))
    for a in toadd
        insertsorted!(set,a)
    end
    while added != 0
        added = 0
        newElements = empty(set)
        for i in length(queue):-1:1
            a,b = queue[i]
            if ⊗(a,b) ∉ set
                push!(newElements,⊗(a,b))
                added = added+1
            end
            pop!(queue)
        end
        newElements = sort(unique(newElements))
        for a in newElements
            for b in set
                push!(queue,(a,b))
                push!(queue,(b,a))
            end
            for b in newElements
                push!(queue,(a,b))
            end
        end

        i = 1
        for a in newElements
            i = insertsorted!(set,a, startindex = i)
        end
    end
    return set
end

function getsubgroups(;maxgenerators = 4)
    subgroupdict = Dict()
    subgroupdict[(generators = 0,)] = [[1]]
    for n = 1:maxgenerators
        timestarted = time()
        println("Generating with $n generators.")
        subgroupbase = subgroupdict[(generators = n-1,)]
        subgroupstoadd_duplicates = vcat([[completewith!(copy(s),[a]) for s in subgroupbase] for a in 1:192]...)
        subgroupstoadd = setdiff(unique(subgroupstoadd_duplicates),vcat([subgroupdict[(generators = i,)] for i = 0:n-1]...))
        subgroupdict[(generators = n,)] = subgroupstoadd
        timeelapsed = time()-timestarted
        println("  Generated after $timeelapsed s.")
    end
    return subgroupdict
end

function addgenerator(subgroups,generator)
    newsubgroups = empty(subgroups)
    for subgroup in subgroups
        newsubgroup = copy(subgroup)
        generator ∈ subgroup && continue
        completewith!(newsubgroup,[generator])
        push!(newsubgroups,newsubgroup)
    end
    return newsubgroups
end

function isgroup(set)
    for a in set
        for b in set
            if ⊗(a,b) ∉ set
                return false
            end
        end
    end
    return true
end

function verifysubgroupdict(subgroupdict)
    #TODO: implement subgroupdict verification
    return true
end

let subgroupdict_file = joinpath(@__DIR__,"subgroupdict.jld2")
    timestarted = time()
    verified = false
    howstring = "loaded"
    if isfile(subgroupdict_file)
        println("Loading all subgroups.")
        subgroupdict = load(subgroupdict_file,"subgroupdict")
        verified = verifysubgroupdict(subgroupdict)
    end
    if !verified
        howstring = "generated"
        println("Generating all subgroups.")
        global subgroupdict = getsubgroups()
        for igens = 1:4
            permute!(subgroupdict[(generators = igens,)],sortperm(length.(subgroupdict[(generators = igens,)])))
        end
        subgroupdict[:allsubgroups] = vcat([subgroupdict[(generators = i,)] for i = 0:4]...)
        save(subgroupdict_file,"subgroupdict",subgroupdict)
    end
    timeelapsed = time()-timestarted
    println("All subgroups $howstring after $timeelapsed s.")
end

function write_subgroups()
    f = open(joinpath(@__DIR__,"subgroups.jl"),"w")
    write(f,"subgroupdict = Dict(")
    for key in keys(subgroupdict)
        print(f,key)
        write(f," => ")
        print(f,subgroupdict[key])
        write(f,",\n")
    end
    skip(f,-2)
    write(f,")\n")
    close(f)
end

write_subgroups()

const vertices = vec([SA[a,b,c,d] for a in [-1,1], b in [-1,1], c in [-1,1], d in [-1,1]])

function element_as_vertex_perm(element)
    permuted_vertices = [element*vertex for vertex in vertices]
    vertex_perm = [findfirst([v_ori==v_new for v_ori in vertices]) for v_new in permuted_vertices]
    return Permutation(vertex_perm)
end

elements_as_vertex_perms = [element_as_vertex_perm(e) for e in elements]

function allequal(array)
    previous = first(array)
    result = true
    for val in array
        result = result && isequal(previous,val)
        previous = val
        !result && continue
    end
    return result
end

function isequalcoloring(coloring1,coloring2)
    if length(unique(coloring1)) != length(unique(coloring2))
        return false
    end
    for c in unique(coloring1)
        v1 = coloring1 .== c
        if !allequal(coloring2[v1])
            return false
        end
    end
    return true
end

function get_subgroup_properties()
    subgroup_properties = [Dict() for s in subgroupdict[:allsubgroups]]
    for (i,sdict) in enumerate(subgroup_properties)
        sdict[:subgroup] = subgroupdict[:allsubgroups][i]
        sdict[:order] = length(sdict[:subgroup])
        j = i
        ngenerators = -1
        while j > 0
            ngenerators += 1
            with_ngenerators = length(subgroupdict[(generators = ngenerators,)])
            with_ngenerators == 0 && break
            j -= with_ngenerators
        end
        sdict[:ngenerators] = ngenerators
    end

    for (is,sdict) in enumerate(subgroup_properties)
        sdict[:coloring] = collect(1:16)
        for (iv,v) in enumerate(vertices)
            group_action_on_v = [elements[e]*v for e in sdict[:subgroup]]
            for connected_vertex in group_action_on_v
                i_con_vert = 1
                while vertices[i_con_vert] != connected_vertex && i_con_vert < 17
                    i_con_vert = i_con_vert+1
                end
                #if connected_vertex has lower color than v, leave it be (should not happen)
                sdict[:coloring][i_con_vert] = min(sdict[:coloring][i_con_vert],sdict[:coloring][iv])
            end
        end
    end

    for (is,sdict) in enumerate(subgroup_properties)
        subgroupcoloring = sdict[:coloring]
        newcolorvalue = 1
        newsubgroupcoloring = copy(subgroupcoloring)
        for colorvalue in sort(unique(subgroupcoloring))
            newsubgroupcoloring[subgroupcoloring.==colorvalue] .= newcolorvalue
            newcolorvalue = newcolorvalue+1
        end
        sdict[:coloring_sorted] = newsubgroupcoloring
    end

    for (is,sdict) in enumerate(subgroup_properties)
        sdict[:coloring_equiv] = zeros(Int,16)
    end

    for (is,sdict) in enumerate(subgroup_properties)
        sgcolor = sdict[:coloring_sorted]
        found = false
        for p in elements_as_vertex_perms
            sgcolorperm = sgcolor[[getindex(p,k) for k in axes(sgcolor,1)]]
            for other_sdict in subgroup_properties
                if isequalcoloring(sgcolorperm,other_sdict[:coloring_equiv])
                    found = true
                    sdict[:coloring_equiv] = other_sdict[:coloring_equiv]
                    continue
                end
            end
            found && continue
        end
        found && continue
        sdict[:coloring_equiv] = sgcolor
    end

    for (is,sdict) in enumerate(subgroup_properties)
        sdict[:coloring_equiv_index] = 0
    end

    coloring_equiv_index = 1
    for (is,sdict) in enumerate(subgroup_properties)
        found = false
        for other_is in 1:is-1
            if sdict[:coloring_equiv] == subgroup_properties[other_is][:coloring_equiv]
                found = true
                sdict[:coloring_equiv_index] = subgroup_properties[other_is][:coloring_equiv_index]
                continue
            end
        end
        found && continue
        sdict[:coloring_equiv_index] = coloring_equiv_index
        coloring_equiv_index += 1
    end

    for col_eq_i = 1:coloring_equiv_index-1
        subgs = [sdict[:coloring_equiv_index] == col_eq_i for sdict in subgroup_properties]
        if length(unique([sdict[:ngenerators] for sdict in subgroup_properties[subgs]])) > 1
            n = length(unique(first(subgroup_properties[subgs])[:coloring_equiv]))
            println("Varying number of generators for coloring_equiv_index = $col_eq_i with $n colors")
        end
    end

    mirrorx = [-1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1]
    elements_as_vertex_perms_mirror = vcat([element_as_vertex_perm(e) for e in elements],[element_as_vertex_perm(e*mirrorx) for e in elements])

    for (is,sdict) in enumerate(subgroup_properties)
        sdict[:coloring_equiv_mirror] = zeros(Int,16)
    end

    for (is,sdict) in enumerate(subgroup_properties)
        println("Processing $is-th subgroup.")
        sgcolor = sdict[:coloring_sorted]
        found = false
        for p in elements_as_vertex_perms_mirror
            sgcolorperm = sgcolor[[getindex(p,k) for k in axes(sgcolor,1)]]
            for other_sdict in subgroup_properties
                if isequalcoloring(sgcolorperm,other_sdict[:coloring_equiv_mirror])
                    found = true
                    sdict[:coloring_equiv_mirror] = other_sdict[:coloring_equiv_mirror]
                    continue
                end
            end
            found && continue
        end
        found && continue
        sdict[:coloring_equiv_mirror] = sgcolor
    end

    for (is,sdict) in enumerate(subgroup_properties)
        sdict[:coloring_equiv_mirror_index] = 0
    end

    coloring_equiv_mirror_index = 1
    for (is,sdict) in enumerate(subgroup_properties)
        found = false
        for other_is in 1:is-1
            if sdict[:coloring_equiv_mirror] == subgroup_properties[other_is][:coloring_equiv_mirror]
                found = true
                sdict[:coloring_equiv_mirror_index] = subgroup_properties[other_is][:coloring_equiv_mirror_index]
                continue
            end
        end
        found && continue
        sdict[:coloring_equiv_mirror_index] = coloring_equiv_mirror_index
        coloring_equiv_mirror_index += 1
    end

    for col_eq_i = 1:coloring_equiv_mirror_index-1
        subgs = [sdict[:coloring_equiv_mirror_index] == col_eq_i for sdict in subgroup_properties]
        if length(unique([sdict[:ngenerators] for sdict in subgroup_properties[subgs]])) > 1
            n = length(unique(first(subgroup_properties[subgs])[:coloring_equiv_mirror]))
            println("Varying number of generators for coloring_equiv_index = $col_eq_i with $n colors")
        end
    end
    return subgroup_properties
end

let subgroup_properties_file = joinpath(@__DIR__,"subgroup_properties.jld2")
    println("Calculating subgroup properties.")
    global subgroup_properties = get_subgroup_properties()
    save(subgroup_properties_file,"subgroup_properties", subgroup_properties)
end

function write_groupmatrix()
    f = open(joinpath(@__DIR__,"groupmatrix.jl"),"w")
    groupmatrixname = :g
    write(f,"const $(string(groupmatrixname)) = [")
    for row in axes(eval(groupmatrixname),1)
        for col in axes(eval(groupmatrixname),2)
            write(f,"$(g[row,col]) ")
        end
        write(f,";\n")
    end
    skip(f,-2)
    write(f,"]\n")
    close(f)
end

write_groupmatrix()


s=time()
spl=split(ARGS[1])
expr=spl[1]
a=float(spl[2])
b=float(spl[3])
n=Integer(float(spl[4]))
include_string(string("f(x)=",expr))
l=abs(b-a)
function serie(i)
	return f(a+(i*l)/(n))*(l/n)
end
r=0
for i in 1:n
	r+=serie(i-1)
end
println(r)
println(time()-s)


@testset "DataStructures" begin
@testset "BasicBinaryTree" begin
    @testset "BasicBinaryTree: BinaryTree" begin
        tree=BinaryTree{Int}(10,0)
        insert!(tree,1,10)
        insert!(tree,1,20)
        insert!(tree,2,50)
        #=
        [1] 0
         |       \
        [2] 10  [3] 20
         |
        [4] 50
        =#
        @test tree.n == 4
        @test depth(tree,4) == 2
        @test isleaf(tree,2) == false
        @test ch(tree,1,true) == 3
        @test tree.val[left(tree)] == 50
        @test height(tree) == 2
    end
    @testset "DisjointSet: DisjointSet" begin
        set=DisjointSet(10)
        set.par[4]=1
        set.par[3]=1
        set.par[2]=3
        set.par[5]=8
        @test find(set,2) == 1
        merge!(set,5,4)
        @test find(set,5) == find(set,2)
    end
end
end


# Example: Compartmental models in epidemiology
#=
- https://github.com/epirecipes/sir-julia/blob/master/markdown/function_map/function_map.md
- https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology#Deterministic_versus_stochastic_epidemic_models
=#
using ModelingToolkit, Test

@inline function rate_to_proportion(r,t)
    1-exp(-r*t)
end;

# Independent and dependent variables and parameters
@parameters t c nsteps δt β γ
D = Difference(t; dt=0.1)
@variables S(t) I(t) R(t)

infection = rate_to_proportion(β*c*I/(S+I+R),δt)*S
recovery = rate_to_proportion(γ,δt)*I

# Equations
eqs = [D(S) ~ S-infection,
       D(I) ~ I+infection-recovery,
       D(R) ~ R+recovery]

# System
@named sys = DiscreteSystem(eqs,t,[S,I,R],[c,nsteps,δt,β,γ]; controls = [β, γ])
syss = structural_simplify(sys)
@test syss == syss

# Problem
u0 = [S => 990.0, I => 10.0, R => 0.0]
p = [β => 0.05, c => 10.0, γ => 0.25, δt => 0.1, nsteps => 400]
tspan = (0.0,ModelingToolkit.value(substitute(nsteps,p))) # value function (from Symbolics) is used to convert a Num to Float64
prob_map = DiscreteProblem(sys,u0,tspan,p)

# Solution
using OrdinaryDiffEq
sol_map = solve(prob_map,FunctionMap());
@test sol_map[S] isa Vector

# Using defaults constructor
@parameters t c=10.0 nsteps=400 δt=0.1 β=0.05 γ=0.25
Diff = Difference(t; dt=0.1)
@variables S(t)=990.0 I(t)=10.0 R(t)=0.0

infection2 = rate_to_proportion(β*c*I/(S+I+R),δt)*S
recovery2 = rate_to_proportion(γ,δt)*I

eqs2 = [D(S) ~ S-infection2,
       D(I) ~ I+infection2-recovery2,
       D(R) ~ R+recovery2]

@named sys = DiscreteSystem(eqs2; controls = [β, γ])
@test ModelingToolkit.defaults(sys) != Dict()

prob_map2 = DiscreteProblem(sys,[],tspan)
sol_map2 = solve(prob_map,FunctionMap());

@test sol_map.u == sol_map2.u
@test sol_map.prob.p == sol_map2.prob.p

# Direct Implementation

function sir_map!(u_diff,u,p,t)
    (S,I,R) = u
    (β,c,γ,δt) = p
    N = S+I+R
    infection = rate_to_proportion(β*c*I/N,δt)*S
    recovery = rate_to_proportion(γ,δt)*I
    @inbounds begin
        u_diff[1] = S-infection
        u_diff[2] = I+infection-recovery
        u_diff[3] = R+recovery
    end
    nothing
end;
u0 = [990.0,10.0,0.0];
p = [0.05,10.0,0.25,0.1];
prob_map = DiscreteProblem(sir_map!,u0,tspan,p);
sol_map2 = solve(prob_map,FunctionMap());

@test Array(sol_map) ≈ Array(sol_map2)

# Delayed difference equation
@parameters t
@variables x(..) y(..) z(t)
D1 = Difference(t; dt=1.5)
D2 = Difference(t; dt=2)

@test ModelingToolkit.is_delay_var(Symbolics.value(t), Symbolics.value(x(t-2)))
@test ModelingToolkit.is_delay_var(Symbolics.value(t), Symbolics.value(y(t-1)))
@test !ModelingToolkit.is_delay_var(Symbolics.value(t), Symbolics.value(z))
@test_throws ErrorException ModelingToolkit.get_delay_val(Symbolics.value(t), Symbolics.arguments(Symbolics.value(x(t+2)))[1])
@test_throws ErrorException z(t)

# Equations
eqs = [
    D1(x(t)) ~ 0.4x(t) + 0.3x(t-1.5) + 0.1x(t-3),
    D2(y(t)) ~ 0.3y(t) + 0.7y(t-2) + 0.1z,
]

# System
@named sys = DiscreteSystem(eqs,t,[x(t),x(t-1.5),x(t-3),y(t),y(t-2),z],[])

eqs2, max_delay = ModelingToolkit.linearize_eqs(sys; return_max_delay=true)

@test max_delay[Symbolics.operation(Symbolics.value(x(t)))] ≈ 3
@test max_delay[Symbolics.operation(Symbolics.value(y(t)))] ≈ 2

linearized_eqs = [
    eqs
    x(t - 3.0) ~ x(t - 1.5)
    x(t - 1.5) ~ x(t)
    y(t - 2.0) ~ y(t)
]
@test all(eqs2 .== linearized_eqs)

# observed variable handling
@variables t x(t) RHS(t)
@parameters τ
@named fol = DiscreteSystem([D(x) ~ (1 - x)/τ]; observed=[RHS ~ (1 - x)/τ])
@test isequal(RHS, @nonamespace fol.RHS)
RHS2 = RHS
@unpack RHS = fol
@test isequal(RHS, RHS2)

@testset "Preface tests" begin
    @parameters t
    using OrdinaryDiffEq
    using Symbolics
    using DiffEqBase: isinplace
    using ModelingToolkit
    using SymbolicUtils.Code
    using SymbolicUtils: Sym

    c = [0]
    f =  function f(c, d::Vector{Float64}, u::Vector{Float64}, p, t::Float64, dt::Float64)
            c .= [c[1] + 1]
            d .= randn(length(u))
            nothing
        end

    dummy_identity(x, _) = x
    @register dummy_identity(x, y)

    u0 = ones(5)
    p0 = Float64[]
    syms = [Symbol(:a, i) for i in 1:5]
    syms_p = Symbol[]
    dt = 0.1
    @assert isinplace(f, 6)
    wf = let c=c, buffer = similar(u0), u=similar(u0), p=similar(p0), dt=dt
            t -> (f(c, buffer, u, p, t, dt); buffer)
    end

    num = hash(f) ⊻ length(u0) ⊻ length(p0)
    buffername = Symbol(:fmi_buffer_, num)

    Δ = DiscreteUpdate(t; dt=dt)
    us = map(s->(@variables $s(t))[1], syms)
    ps = map(s->(@variables $s(t))[1], syms_p)
    buffer, = @variables $buffername[1:length(u0)]
    dummy_var = Sym{Any}(:_) # this is safe because _ cannot be a rvalue in Julia

    ss = Iterators.flatten((us, ps))
    vv = Iterators.flatten((u0, p0))
    defs = Dict{Any, Any}(s=>v for (s, v) in zip(ss, vv))

    preface = [
               Assignment(dummy_var, SetArray(true, term(getfield, wf, Meta.quot(:u)), us))
               Assignment(dummy_var, SetArray(true, term(getfield, wf, Meta.quot(:p)), ps))
               Assignment(buffer, term(wf, t))
              ]
    eqs = map(1:length(us)) do i
        Δ(us[i]) ~ dummy_identity(buffer[i], us[i])
    end

    @named sys = DiscreteSystem(eqs, t, us, ps; defaults=defs, preface=preface)
    prob = DiscreteProblem(sys, [], (0.0, 1.0))
    sol = solve(prob, FunctionMap(); dt=dt)
    @test c[1]+1 == length(sol)
end


function set_res(Delta_Qcon_hat, Delta_Qcon_hat_temp, cellxmax, cellymax, nval)
    res = zeros(cellxmax, cellymax, nval)

    for i in 2:cellxmax-1
        for j in 2:cellymax-1
            for l in 1:nval
                res[i,j,l] = abs(Delta_Qcon_hat[i,j,l]-Delta_Qcon_hat_temp[i,j,l])
            end
        end
    end

    return res
end

function check_converge(res, RHS, cellxmax, cellymax, init_small, nval)
    norm2 = zeros(nval)

    tempAxb = zeros(nval)
    tempb = zeros(nval)
    for i in 2:cellxmax-1
        for j in 2:cellymax-1
            for l in 1:nval
                tempAxb[l] = tempAxb[l] + res[i,j,l]^2
                tempb[l] = tempb[l] + RHS[i,j,l]^2
            end
        end
    end
    for l in 1:nval
        norm2[l] = (tempAxb[l]/(tempb[l]+init_small)) ^ 0.5
    end
    return norm2
end


#=
Given a array of numbers representing the stock prices of a company in chronological order, write a function that calculates the maximum profit you could have made from buying and selling that stock once. You must buy before you can sell it.

For example, given [9, 11, 8, 5, 7, 10], you should return 5, since you could buy the stock at 5 dollars and sell it at 10 dollars.
=#

function maximumprofit(stockprices::Array{Int,1})
    max_profit = 0
    for index in eachindex(stockprices[1:end-1])
        current_profit = maximum(stockprices[index+1:end]) - stockprices[index]
        if current_profit > max_profit
            max_profit = current_profit
        end
    end
    return max_profit
end


module Day12

# Part 1
instructions = readlines("day12_input.txt")

mutable struct Ship
    pos
    angle
end

function move!(ship::Ship, instruction)
    rules = Dict(
        'N' => (s::Ship, v) -> s.pos += [0, v],
        'S' => (s::Ship, v) -> s.pos -= [0, v],
        'E' => (s::Ship, v) -> s.pos += [v, 0],
        'W' => (s::Ship, v) -> s.pos -= [v, 0],
        'L' => (s::Ship, v) -> s.angle += v,
        'R' => (s::Ship, v) -> s.angle -= v,
        'F' => (s::Ship, v) -> s.pos += v * [round(Int, cosd(s.angle)), round(Int, sind(s.angle))]
    )
    rules[first(instruction)](ship, parse(Int32, instruction[2:end]))
end

ship = Ship([0, 0], 0)
map(i -> move!(ship, i), instructions)
println("Part 1: Manhattan dist = $(sum(abs.(ship.pos)))")

# Part 2
mutable struct Waypoint
    pos
end

function rotationmatrix(angle)
    c = round(Int, cosd(angle))
    s = round(Int, sind(angle))
    [c -s
     s c]
end

function move!(ship::Ship, waypoint::Waypoint, instruction)
    rules = Dict(
        'N' => (s::Ship, w::Waypoint, v) -> w.pos += [0, v],
        'S' => (s::Ship, w::Waypoint, v) -> w.pos -= [0, v],
        'E' => (s::Ship, w::Waypoint, v) -> w.pos += [v, 0],
        'W' => (s::Ship, w::Waypoint, v) -> w.pos -= [v, 0],
        'L' => (s::Ship, w::Waypoint, v) -> w.pos = rotationmatrix(v) * w.pos,
        'R' => (s::Ship, w::Waypoint, v) -> w.pos = rotationmatrix(-v) * w.pos,
        'F' => (s::Ship, w::Waypoint, v) -> s.pos += v * w.pos
    )
    rules[first(instruction)](ship, waypoint, parse(Int32, instruction[2:end]))
end

ship = Ship([0, 0], 0)
waypoint = Waypoint([10, 1])
map(i -> move!(ship, waypoint, i), instructions)
println("Part 2: Manhattan dist = $(sum(abs.(ship.pos)))")

end


# Uniform distribution over directed forests with n nodes.

struct UniformDiForest <: Distribution{LG.SimpleDiGraph} end
const uniform_diforest = UniformDiForest()
(::UniformDiForest)(n) = Gen.random(uniform_diforest, n)

# The number of directed forests on n nodes is (n+1) ^ (n-1) as given by Cayley's formula.
function Gen.logpdf(::UniformDiForest, diforest::LG.SimpleDiGraph, n::Int)
    n != LG.nv(diforest) && return -Inf
    - ((n - 1) * log(n+1))
end

export uniform_diforest


A = rand(3, 3)
B = ones(4, 4)
C = randn(5, 6)
K = A ⊗ B
X = collect(K)
v = rand(12)

K3 = A ⊗ B ⊗ C

@testset "vectrick" begin
    @testset "Vec trick" begin

        @test K * v ≈ X * v

        V = rand(4, 3)
        @test K * vec(V) ≈ X * vec(V)
        @test_throws DimensionMismatch K * V
        @test_throws DimensionMismatch K * reshape(V, 2, 6)

        v3 = randn(size(K3, 2))
        @test K3 * v3 ≈ collect(K3) * v3
        u = similar(v)
        @test mul!(u, K, v) ≈ X * v
    end

    @testset "Reshaped vec trick" begin

        @test K * v ≈ X * v

        V = sprand(4, 3, 1.0)
        @test K * vec(V) ≈ X * vec(V)
        @test_throws DimensionMismatch K * V
        @test_throws DimensionMismatch K * reshape(V, 2, 6)
        K3 = A ⊗ B ⊗ C
        v3 = sprand(size(K3, 2),1.0)
        @test K3 * vec(v3) ≈ collect(K3) * v3
        u = similar(v)
        @test mul!(u, K, v) ≈ X * v

    end

    @testset "AbstractKroneckerProduct * AbstractMatrix" begin
        rng = MersenneTwister(123456)
        a, b, x = randn(rng, 30, 20), randn(rng, 40, 50), randn(rng, 1000, 1100)

        @test kron(a, Diagonal(ones(50))) * x ≈ Kronecker.kron_a_id(a, x)
        @test kron(Diagonal(ones(20)), b) * x ≈ Kronecker.kron_id_a(b, x)
        @test kron(a, b) * x ≈ Kronecker.kron_a_b(a, b, x)

        @test kron(a, b) * x ≈ (a ⊗ b) * x
        @test kron(a, Eye(50)) * x ≈ (a ⊗ Eye(50)) * x
        @test kron(Eye(20), b) * x ≈ (Eye(20) ⊗ b) * x
    end

    @testset "sum" begin
        @test sum(K) ≈ sum(X)
        @test sum(K3) ≈ sum(collect(K3))

        @test sum(K, dims=1) ≈ sum(X, dims=1)
        @test sum(K3, dims=2) ≈ sum(collect(K3), dims=2)
        @test sum(K3, dims=2) isa AbstractKroneckerProduct

        @test sum(kronecker(A, 3)) ≈ sum(kron(A, A, A))
    end
end


const data = parse.(Int, split("219748365", ""))

function destination(cups::Vector{Int}, current::Int, n::Int=9)
    for i in 1:4
        a = mod1(current - i, n)
        a ∉ cups && return a
    end
end

function move!(allcups::Vector{Int}, current_ind::Int)
    current = allcups[current_ind]
    allcups = circshift(allcups, -current_ind)
    cups = splice!(allcups, 1:3)
    next_val = allcups[1]
    dest = findfirst(==(destination(cups, current)), allcups)
    allcups = prepend!(circshift(allcups, -dest), cups)
    return (allcups, findfirst(==(next_val), allcups))
end

function play(allcups::Vector{Int}, n::Int=100)
    dest = 1
    for i in 1:n
        i % 1000 == 0 && println("$(i / n) %")
        allcups, dest = move!(allcups, dest)
    end
    return allcups
end

function result(allcups::Vector{Int})
    one = findfirst(==(1), allcups)
    return join(circshift(allcups, -one)[1:8])
end

println("Part 1: $(result(play(data)))")

function tonextsdict(data::Vector{Int}, p1::Bool=false)
    nexts = Dict{Int,Int}()
    for (d, n) in zip(data, data[2:end])
        nexts[d] = n
    end
    if p1
        nexts[data[end]] = data[1]
        return nexts
    end
    nexts[data[end]] = 10
    for i in 10:999_999
        nexts[i] = i + 1
    end
    nexts[1_000_000] = data[1]
    return nexts
end

function move2!(nexts::Dict{Int,Int}, current::Int)
    n1 = nexts[current]
    n2 = nexts[n1]
    n3 = nexts[n2]
    nextcurrent = nexts[n3]
    dest = destination([n1, n2, n3], current, length(nexts))
    # insert here
    finalnext = nexts[dest]
    nexts[current] = nextcurrent
    nexts[dest] = n1
    nexts[n3] = finalnext
    return nexts, nextcurrent
end

function play2!(nexts::Dict{Int,Int}, start::Int, n::Int=10_000_000)
    current = start
    for i in 1:n
        nexts, current = move2!(nexts, current)
    end
    return nexts
end

function part2(data)
    nexts = tonextsdict(data)
    play2!(nexts, data[1])
    a = nexts[1]
    b = nexts[a]
    return a * b
end

println("Part 2: $(part2(data))")


module projection

using LinearAlgebra

export distance_matrix_to_absolute_coordinates

function distance_matrix_to_absolute_coordinates( D::Array{Float64,2} , bSquared::Bool=false , n_dimensions::Int64=2 )
    """
    Same function can be found in my impetuous-gfa repo here:
    https://github.com/richardtjornhammar/impetuous/blob/master/src/impetuous/clustering.py
    around line 121 (same name)
    """
    if !bSquared
        D = D.^2.
    end
    N,M = size(D)
    DIM = n_dimensions
    DIJ = D.*0.
    for i in 1:M
        for j in 1:M
            DIJ[i,j] = 0.5* (D[i,end]+D[j,end]-D[i,j])
	end
    end
    D = DIJ
    U,S,Vt = svd( D )
    S[DIM+1:end] *= 0.
    Z  = diagm(S.^0.5)[:,1:DIM]
    xr = Vt*Z
    return ( xr )
end

end


export ARC_Cub
function ARC_Cub(h       :: AbstractNLPModel,
                 nlpstop :: AbstractStopping;
                 t₀ = h.meta.x0, tₘ = 100.0,
                 verbose = false)
    # ARC parameters
    eps1 = 0.2; eps2 = 0.8
    red = 0.5; aug = 20.0
    Δ = [1.0]
    t = t₀

    iter = 0;
    fₖ = obj(h, t)
    gₖ = grad(h, t)

    secₖ = [1.0]
    dN = -gₖ # for first iterations
    A = [0.5]; B = [0.0]
    q(d) = fₖ .+ gₖ .* d .+ A .* d .^ 2 .+ B .* d .^ 3

    OK  = update_and_start!(nlpstop, x = t, fx = fₖ, gx = gₖ, g0 = copy(gₖ))

    verbose && @printf(" iter  t         gₖ          Δ        pred         ared\n")
    verbose && @printf(" %4d %7.2e  %7.2e  %7.2e \n", iter, t[1], gₖ[1], Δ[1])

    while !OK

      #Step computation
      Quad(t) = fₖ .+ gₖ .* t .+ A .* t .^ 2 .+ B .* t .^ 3 .+ (1 ./ (4.0 .* Δ)) .* t .^ 4
      dQuad(t) = gₖ .+ 2.0 .* A .* t .+ 3.0 .* B .* t .^ 2 .+ (1 ./ Δ) .* t .^ 3

      p = Poly([gₖ[1], 2.0 .* A[1], 3.0 .* B[1], 1.0 ./ Δ[1]])
      dR = roots(p)
      vmin = Inf; dN = 0
      for i=1:length(dR)
        rr = dR[i]
        if isreal(rr)
          rrr = real(rr)
          vact = fₖ .+ gₖ .* rrr .+ A .* rrr .^2 .+ B .* rrr .^ 3 .+ (1 ./ (4.0 .* Δ)) .* rrr .^ 4
          if true in (rrr .* gₖ .< 0.0)
            if true in (vact .< vmin)
              dN = rrr
              vmin = vact
            end
          end
        end
      end

      d = dN

      if true in (d .== 0.0)
        return OK, nlpstop
      end

      ftestTR = obj(h, t .+ d)
      gtestTR = grad(h, t .+ d)

      pred = gₖ .* d .+ A .* d .^ 2 .+ B .* d .^ 3
      if true in (pred .> -1e-10)
        ared = (gₖ .+ gtestTR) .* d ./ 2.0
      else
        ared = ftestTR .- fₖ
      end

      ratio = ared ./ pred
      if true in (ratio .< eps1)
        Δ = red .* Δ
      else
        # two point memory
        tpred = t
        gkm1 = gₖ
        fkm1 = fₖ

        t = t .+ d
        gₖ = gtestTR
        fₖ = ftestTR

        # Cubic two points interpolation
        s = t .- tpred
        y = gₖ .- gkm1
        α = -s
        z = gₖ .+ gkm1 .+ 3.0 .* ( fₖ .- fkm1) ./ α
        discr= z .^ 2 .- gₖ .* gkm1
        denom = gₖ .+ gkm1 .+ 2.0 .* z
        B = (1 ./ 3) .* (gₖ .+ gkm1 .+ 2.0 .* z) ./ ( α .* α)
        A = -(gₖ .+ z) ./ α
        if true in (ratio .> eps2)
          Δ = aug .* Δ
        end
      end
      OK  = update_and_stop!(nlpstop, x = t, fx = fₖ, gx = gₖ, g0 = copy(gₖ));
      iter += 1
      verbose && @printf(" %4d %7.2e  %7.2e  %7.2e %7.2e %7.2e\n", iter, t[1], gₖ[1], Δ[1], pred[1], ared[1])
    end
    optimal = nlpstop.meta.optimal
    return optimal, nlpstop

end


using SchumakerSpline
using Test
tol = 100 * eps()
close(a,b) = abs(a-b) < tol

from = 0.5
to = 10
x1 = collect(range(from, stop=to, length=40))
y1 = (x1).^2
x2 = [0.5,0.75,0.8,0.93,0.9755,1.0,1.1,1.4,2.0]
y2 = sqrt.(x2)
s1 = Schumaker(x1,y1)
s2 = Schumaker(x2,y2)
con = 7.6
x = 0.9

close( (s1 +con)(x) , s1(x) + con)
close( (con + s1)(x) , s1(x) + con)

close( (s1 - con)(x) , s1(x) - con)
close( (con - s1)(x) , con - s1(x) )

close( (s1 *con)(x) , s1(x) * con)
close( (con * s1)(x) , s1(x) * con)

close( (s1 /con)(x) , s1(x) / con)