"""
costfunctions.jl
"""

import Base: copy, +

"""
Abstract type that represents a scalar-valued function that accepts a state and control
at a single knot point.
"""
abstract type CostFunction end

"""
Cost function of the form
```math
\\frac{1}{2} x^T Q x + \\frac{1}{2} u^T R u + u^T H x + q^T x + r^T u + c
```
"""
struct QuadraticCost{TQ,TR,TH,Tq,Tr,T} <: CostFunction
    Q::TQ
    R::TR
    H::TH
    q::Tq
    r::Tr
    c::T
    n::Int
    m::Int
    terminal::Bool
    Q_diag::Bool
    R_diag::Bool
    use_R::Bool
    use_H::Bool
    use_r::Bool
    const_grad::Bool
    const_hess::Bool
end

# constructors
function QuadraticCost(Q::TQ, R::TR, H::TH, q::Tq, r::Tr, c::T; use_R=true,
                       use_H=true, use_r=true, terminal=false, checks=true) where {
                           TQ,TR,TH,Tq,Tr,T}
    n = size(Q, 1)
    m = size(R, 1)
    if checks
        @assert length(q) == n
        if use_r
            @assert length(r) == m
        end
        if use_H
            @assert size(H) == (m, n)
        end
        if !ispossemidef(Array(Q))
            @warn "Q is not positive semidefinite:\n$(Q)"
        end
        if use_R && !isposdef(Array(R)) && !terminal
            @warn "R is not positive definite:\n$(R)"
        end
    end
    Q_diag = (TQ <: Diagonal)
    R_diag = (TR <: Diagonal)
    const_grad = false
    const_hess = true
    return QuadraticCost{TQ,TR,TH,Tq,Tr,T}(Q, R, H, q, r, c, n, m, terminal,
                                           Q_diag, R_diag, use_R, use_H, use_r, const_grad,
                                           const_hess)
end

# methods
Base.copy(cost::QuadraticCost{TQ,TR,TH,Tq,Tr,T}) where {TQ,TR,TH,Tq,Tr,T} = (
    QuadraticCost{TQ,TR,TH,Tq,Tr,T}(
        copy(cost.Q), copy(cost.R), copy(cost.H), copy(cost.q), copy(cost.r), cost.c,
        cost.terminal, cost.use_H, cost.Q_diag, cost.R_diag
))

@inline state_dim(cost::QuadraticCost) = cost.n
@inline control_dim(cost::QuadraticCost) = cost.m

# evaluation methods
function cost(cost_::QuadraticCost, X::AbstractVector, U::AbstractVector, k::Int; terminal=false)
    cost__ = 0.5 * X[k]' * cost_.Q * X[k] + dot(cost_.q, X[k]) + cost_.c
    if cost_.use_R
        cost__ += 0.5 * U[k]' * cost_.R * U[k]
    end
    if cost_.use_H
        cost__ += U[k]' * cost_.H * X[k]
    end
    if cost_.use_r
        cost__ += dot(cost_.r, U[k])
    end
    return cost__
end

function cost_derivatives!(E::QuadraticCost, cost::QuadraticCost, X::AbstractVector,
                           U::AbstractVector, k::Int; terminal=false)
    # E.Q
    E.Q .= cost.Q
    # E.q
    E.q .= cost.q
    mul!(E.q, cost.Q, X[k], 1., 1.)
    if cost.use_H
        mul!(E.q, Transpose(cost.H), U[k], 1., 1.)
    end
    # E.R
    if cost.use_R
        E.R .= cost.R

    else
        E.R .= 0
    end
    # E.r
    if cost.use_R
        mul!(E.r, cost.R, U[k])
    else
        E.r .= 0
    end
    if cost.use_r
        E.r .+= cost.r
    end
    if cost.use_H
        mul!(E.r, cost.H, X[k], 1., 1.)
    end
    # E.H
    if cost.use_H
        E.H .= cost.H
    else
        E.H .= 0
    end
    return nothing
end

# LQRCost
function LQRCost(Q::AbstractMatrix, xf::AbstractVector, R::AbstractMatrix, M, V;
                 use_R=true, terminal=false)
    n = size(Q, 1)
    m = size(R, 1)
    H = M(zeros(m, n))
    q = -Q * xf
    r = V(zeros(m))
    c = 0.5 * xf' * Q * xf
    return QuadraticCost(Q, R, H, q, r, c; use_R=use_R, use_H=false, use_r=false, terminal=terminal)
end

# """
# An abstract type that represents any [`CostFunction`](@ref) of the form
# ```math
# \\frac{1}{2} x^T Q x + \\frac{1}{2} u^T R u + u^T H x + q^T x + r^T u + c
# ```

# These types all support the following methods
# * [`is_diag`](@ref)
# * [`is_blockdiag`](@ref)
# * [`invert!`](@ref)

# As well as standard addition.
# """
# abstract type QuadraticCostFunction{n,m,T} <: CostFunction end

# """
#     is_diag(::QuadraticCostFunction)

# Determines if the hessian of a quadratic cost function is strictly diagonal.
# """
# is_diag(cost::QuadraticCostFunction) = is_blockdiag(cost) && cost.Q isa Diagonal && cost.R isa Diagonal

# """
#     is_diag(::QuadraticCostFunction)

# Determines if the hessian of a quadratic cost function is block diagonal (i.e. ``\\norm(H) = 0``).
# """
# is_blockdiag(::QuadraticCostFunction) = false
# state_dim(::QuadraticCostFunction{n}) where n = n
# control_dim(::QuadraticCostFunction{<:Any,m}) where m = m

# function (::Type{QC})(Q::AbstractArray, R::AbstractArray;
#         H::AbstractArray=SizedMatrix{size(R,1),size(Q,1)}(zeros(eltype(Q), size(R,1), size(Q,1))),
#         q::AbstractVector=(@SVector zeros(eltype(Q), size(Q,1))),
#         r::AbstractVector=(@SVector zeros(eltype(R), size(R,1))),
#         c::Real=zero(eltype(Q)), kwargs...) where QC <: QuadraticCostFunction
#     QC(Q, R, H, q, r, c; kwargs...)
# end

# """
#     stage_cost(costfun::CostFunction, x, u)
#     stage_cost(costfun::CostFunction, x)

# Calculate the scalar cost using `costfun` given state `x` and control `u`. If only the
# state is provided, it is assumed it is a terminal cost.
# """
# function stage_cost(cost::QuadraticCostFunction, x::AbstractVector, u::AbstractVector)
#     J = 0.5*u'cost.R*u + dot(cost.r,u) + stage_cost(cost, x)
#     if !is_blockdiag(cost)
#         J += u'cost.H*x
#     end
#     return J
# end

# function stage_cost(cost::QuadraticCostFunction, x::AbstractVector{T}) where T
#     0.5*x'cost.Q*x + dot(cost.q,x) + cost.c
# end

# """
#     gradient!(E::QuadraticCostFunction, costfun::CostFunction, x, u)
#     gradient!(E::QuadraticCostFunction, costfun::CostFunction, x)

# Evaluate the gradient of the cost function `costfun` at state `x` and control `u`, storing
#     the result in `E.q` and `E.r`. Return a `true` if the gradient is constant, and `false`
#     otherwise.
# """
# function gradient!(E::QuadraticCostFunction, cost::QuadraticCostFunction, x)
#     E.q .= cost.Q*x .+ cost.q
#     return false
# end

# function gradient!(E::QuadraticCostFunction, cost::QuadraticCostFunction, x, u)
#     gradient!(E, cost, x)
#     E.r .= cost.R*u .+ cost.r
#     if !is_blockdiag(cost)
#         E.q .+= cost.H'u
#         E.r .+= cost.H*x
#     end
#     return false
# end

# """
#     hessian!(E::QuadraticCostFunction, costfun::CostFunction, x, u)
#     hessian!(E::QuadraticCostFunction, costfun::CostFunction, x)

# Evaluate the hessian of the cost function `costfun` at state `x` and control `u`, storing
#     the result in `E.Q`, `E.R`, and `E.H`. Return a `true` if the hessian is constant, and `false`
#     otherwise.
# """
# function hessian!(E::QuadraticCostFunction, cost::QuadraticCostFunction, x)
#     if is_diag(cost)
#         for i = 1:length(x); E.Q[i,i] = cost.Q[i,i] end
#     else
#         E.Q .= cost.Q
#     end
#     return true
# end

# function hessian!(E::QuadraticCostFunction, cost::QuadraticCostFunction, x, u)
#     hessian!(E, cost, x)
#     if is_diag(cost)
#         for i = 1:length(u); E.R[i,i] = cost.R[i,i]; end
#     else
#         E.R .= cost.R
#     end
#     if !is_blockdiag(cost)
#         E.H .= cost.H
#     end
#     return true
# end

# function Base.copy(c::QC) where QC<:QuadraticCostFunction
#     QC(copy(c.Q), copy(c.R), H=copy(c.H), q=copy(c.q), r=copy(c.r), c=c.c,
#         terminal=c.terminal, checks=false)
# end

# # Other methods
# """
#     set_LQR_goal!(cost::QuadraticCostFunction, xf)
#     set_LQR_goal!(cost::QuadraticCostFunction, xf, uf)

# Change the reference state and control for an LQR tracking cost.
# Only changes `q` and `r`, and not the constant term `c`. If `uf` is
# not passed in, it isn't changed.
# """
# function set_LQR_goal!(cost::QuadraticCostFunction, xf)
#     cost.q .= -cost.Q * xf
#     return nothing
# end

# function set_LQR_goal!(cost::QuadraticCostFunction, xf, uf)
#     set_LQR_goal!(cost, xf)
#     cost.r .= -cost.R * uf
#     return nothing
# end

# function +(c1::QuadraticCostFunction, c2::QuadraticCostFunction)
#     @assert state_dim(c1) == state_dim(c2)
#     @assert control_dim(c1) == control_dim(c2)
#     n,m = state_dim(c1), control_dim(c1)
#     H1 = c1 isa DiagonalCost ? zeros(m,n) : c1.H
#     H2 = c2 isa DiagonalCost ? zeros(m,n) : c2.H
#     QC = promote_type(typeof(c1), typeof(c2))
#     QC(c1.Q + c2.Q, c1.R + c2.R, H1 + H2,
#        c1.q + c2.q, c1.r + c2.r, c1.c + c2.c,
#        checks=false, terminal=c1.terminal && c2.terminal)
# end

# """
#     invert!(Ginv, cost::QuadraticCostFunction)

# Invert the hessian of the cost function, storing the result in `Ginv`. Performs the inversion
#     efficiently, depending on the structure of the Hessian (diagonal or block diagonal).
# """
# function invert!(Ginv, cost::QuadraticCostFunction{n,m}) where {n,m}
#     ix = 1:n
#     iu = n .+ (1:m)
#     if is_diag(cost)
#         for i = 1:n; Ginv[i,i] = inv(cost.Q[i,i]); end
#         if !cost.terminal
#             for i = 1:m; Ginv[i+n,i+n] = inv(cost.R[i,i]); end
#         end
#     elseif is_blockdiag(cost)
#         Ginv[ix,ix] .= inv(SizedMatrix{n,n}(cost.Q))
#         if !cost.terminal
#             Ginv[iu,iu] .= inv(SizedMatrix{m,m}(cost.R))
#         end
#     else
#         G1 = [cost.Q cost.H']
#         G2 = [cost.H cost.R ]
#         G = [G1; G2]
#         Ginv .= inv(G)
#     end
# end

# #######################################################
# #              COST FUNCTION INTERFACE                #
# #######################################################

# """
#     DiagonalCost{n,m,T}

# Cost function of the form
# ```math
# \\frac{1}{2} x^T Q x + \\frac{1}{2} u^T R u + q^T x + r^T u + c
# ```
# where ``Q`` and ``R`` are positive semi-definite and positive definite diagonal matrices,
# respectively, and ``x`` is `n`-dimensional and ``u`` is `m`-dimensional.

# # Constructors
#     DiagonalCost(Qd, Rd, q, r, c; kwargs...)
#     DiagonalCost(Q, R, q, r, c; kwargs...)
#     DiagonalCost(Qd, Rd; [q, r, c, kwargs...])
#     DiagonalCost(Q, R; [q, r, c, kwargs...])

# where `Qd` and `Rd` are the diagonal vectors, and `Q` and `R` are matrices.

# Any optional or omitted values will be set to zero(s). The keyword arguments are
# * `terminal` - A `Bool` specifying if the cost function is terminal cost or not.
# * `checks` - A `Bool` specifying if `Q` and `R` will be checked for the required definiteness.
# """
# struct DiagonalCost{n,m,T} <: QuadraticCostFunction{n,m,T}
#     Q::Diagonal{T,SVector{n,T}}
#     R::Diagonal{T,SVector{m,T}}
#     q::MVector{n,T}
#     r::MVector{m,T}
#     c::T
#     terminal::Bool
#     function DiagonalCost(Qd::StaticVector{n}, Rd::StaticVector{m},
#                           q::StaticVector{n},  r::StaticVector{m},
#                           c::Real; terminal::Bool=false, checks::Bool=true) where {n,m}
#         T = promote_type(typeof(c), eltype(Qd), eltype(Rd), eltype(q), eltype(r))
#         if checks
#             if any(x->x<0, Qd)
#                 @warn "Q needs to be positive semi-definite."
#             elseif any(x->x<=0, Rd) && !terminal
#                 @warn "R needs to be positive definite."
#             end
#         end
#         new{n,m,T}(Diagonal(SVector(Qd)), Diagonal(SVector(Rd)), q, r, T(c), terminal)
#     end
# end

# Base.copy(c::DiagonalCost) =
#     DiagonalCost(c.Q.diag, c.R.diag, copy(c.q), copy(c.r), c.c,
#         checks=false, terminal=c.terminal)

# # strip away the H
# @inline function (::Type{<:DiagonalCost})(Q::AbstractArray, R::AbstractArray, H::AbstractArray,
#         q::AbstractVector, r::AbstractVector, c; kwargs...)
#     DiagonalCost(Q, R, q, r, c; kwargs...)
# end

# # convert to form for inner constructor
# function DiagonalCost(Q::AbstractArray, R::AbstractArray,
#         q::AbstractVector, r::AbstractVector, c; kwargs...)
#     n,m = length(q), length(r)
#     Qd = SVector{n}(diag(Q))
#     Rd = SVector{m}(diag(R))
#     DiagonalCost(Qd, Rd, convert(MVector{n},q), convert(MVector{m},r), c; kwargs...)
# end

# # Pass in vectors for Q and R
# function DiagonalCost(Q::AbstractVector, R::AbstractVector,
#         q::AbstractVector, r::AbstractVector, c::Real; kwargs...)
#     DiagonalCost(Diagonal(Q), Diagonal(R), q, r, c; kwargs...)
# end

# function DiagonalCost(Q::AbstractVector, R::AbstractVector; kwargs...)
#     DiagonalCost(Diagonal(Q), Diagonal(R); kwargs...)
# end

# is_blockdiag(::DiagonalCost) = true
# is_diag(::DiagonalCost) = true

# function LinearAlgebra.inv(cost::DiagonalCost)
#     return DiagonalCost(inv(cost.Q), inv(cost.R), cost.q, cost.r, cost.c, terminal=cost.terminal)
# end

# function Base.:\(cost::DiagonalCost, z::AbstractKnotPoint)
#     x = cost.Q\state(z)
#     u = cost.R\control(z)
#     return StaticKnotPoint([x;u], z._x, z._u, z.dt, z.t)
# end

# function change_dimension(cost::DiagonalCost, n::Int, m::Int, ix, iu)
#     Qd = zeros(n)
#     Rd = zeros(m)
#     q = zeros(n)
#     r = zeros(m)
#     Qd[ix] = diag(cost.Q)
#     Rd[iu] = diag(cost.R)
#     q[ix] = cost.q
#     r[iu] = cost.r
#     DiagonalCost(Qd, Rd, q, r, cost.c, terminal=cost.terminal, checks=false)
# end

# """
#     QuadraticCost{n,m,T,TQ,TR}

# Cost function of the form
# ```math
# \\frac{1}{2} x^T Q x + \\frac{1}{2} u^T R u + u^T H x + q^T x + r^T u + c
# ```
# where ``R`` must be positive definite, ``Q`` and ``Q_f`` must be positive semidefinite.

# The type parameters `TQ` and `TR` specify the type of ``Q`` and ``R``.

# # Constructor
#     QuadraticCost(Q, R, H, q, r, c; kwargs...)
#     QuadraticCost(Q, R; H, q, r, c, kwargs...)

# Any optional or omitted values will be set to zero(s). The keyword arguments are
# * `terminal` - A `Bool` specifying if the cost function is terminal cost or not.
# * `checks` - A `Bool` specifying if `Q` and `R` will be checked for the required definiteness.
# """
# mutable struct QuadraticCost{n,m,T,TQ,TR} <: QuadraticCostFunction{n,m,T}
#     Q::TQ                     # Quadratic stage cost for states (n,n)
#     R::TR                     # Quadratic stage cost for controls (m,m)
#     H::Matrix{T}   # Quadratic Cross-coupling for state and controls (m,n)
#     q::Vector{T}           # Linear term on states (n,)
#     r::Vector{T}           # Linear term on controls (m,)
#     c::T                      # constant term
#     terminal::Bool
#     zeroH::Bool
#     function (::Type{QC})(Q::TQ, R::TR, H::TH,
#             q::Tq, r::Tr, c::Real; checks=true, terminal=false) where {TQ,TR,TH,Tq,Tr,QC<:QuadraticCost}
#         @assert size(Q,1) == length(q)
#         @assert size(R,1) == length(r)
#         @assert size(H) == (length(r), length(q))
#         if checks
#             if !isposdef(Array(R)) && !terminal
#                 @warn "R is not positive definite"
#             end
#             if !ispossemidef(Array(Q))
#                 @warn "Q is not positive semidefinite"
#             end
#         end
#         zeroH = norm(H,Inf) ≈ 0
#         m,n = size(H)
#         T = promote_type(eltype(Q), eltype(R), eltype(H), eltype(q), eltype(r), typeof(c))
#         new{n,m,T,TQ,TR}(Q,R,H,q,r,c,terminal,zeroH)
#     end
#     function QuadraticCost{n,m,T,TQ,TR}(qcost::QuadraticCost) where {n,m,T,TQ,TR}
#         new{n,m,T,TQ,TR}(qcost.Q, qcost.R, qcost.H, qcost.q, qcost.r, qcost.c, qcost.terminal, qcost.zeroH)
#     end
# end

# state_dim(cost::QuadraticCost) = length(cost.q)
# control_dim(cost::QuadraticCost) = length(cost.r)
# is_blockdiag(cost::QuadraticCost) = cost.zeroH

# function QuadraticCost{T}(n::Int,m::Int; terminal=false) where T
#     Q = Matrix(one(T)*I,n,n)
#     R = Matrix(one(T)*I,m,m)
#     H = zeros(T,m,n)
#     q = zeros(T,n)
#     r = zeros(T,m)
#     c = zero(T)
#     QuadraticCost(Q,R,H,q,r,c, checks=false, terminal=terminal)
# end

# QuadraticCost(cost::QuadraticCost) = cost

# function LinearAlgebra.inv(cost::QuadraticCost)
#     if norm(cost.H,Inf) ≈ 0
#         QuadraticCost(inv(cost.Q), inv(cost.R), cost.H, cost.q, cost.r, cost.c,
#             checks=false, terminal=cost.terminal)
#     else
#         m,n = size(cost.H)
#         H1 = [cost.Q cost.H']
#         H2 = [cost.H cost.R ]
#         H = [H1; H2]
#         Hinv = inv(H)
#         ix = 1:n
#         iu = n .+ (1:m)
#         Q = SizedMatrix{n,n}(Hinv[ix, ix])
#         R = SizedMatrix{m,m}(Hinv[iu, iu])
#         H = SizedMatrix{m,n}(Hinv[iu, ix])
#         QuadraticCost(Q,R,H, cost.q, cost.r, cost.c,
#             checks=false, terminal=cost.terminal)
#     end
# end

# @inline function (::Type{<:QuadraticCost})(dcost::DiagonalCost)
#     QuadraticCost(dcost,Q, dcost.R, q=dcost.q, r=dcost.r, c=dcost.c, terminal=dcost.terminal)
# end

# function Base.convert(::Type{<:QuadraticCost}, dcost::DiagonalCost)
#     QuadraticCost(dcost.Q, dcost.R, q=dcost.q, r=dcost.r, c=dcost.c, terminal=dcost.terminal)
# end

# Base.promote_rule(::Type{<:QuadraticCostFunction}, ::Type{<:QuadraticCostFunction}) = QuadraticCost

# function Base.promote_rule(::Type{<:QuadraticCost{n,m,T1,Q1,R1}},
#                            ::Type{<:QuadraticCost{n,m,T2,Q2,R2}}) where {n,m,T1,T2,Q1,Q2,R1,R2}
#     T = promote_type(T1,T2)
#     function diag_type(T1,T2,n)
#         elT = promote_type(eltype(T1), eltype(T2))
#         if T1 == T2
#             return T1
#         elseif (T1 <: Diagonal) && (T2 <: Diagonal)
#             return Diagonal{elT,MVector{n,elT}}
#         else
#             return SizedMatrix{n,n,elT,2}
#         end
#     end
#     Q = diag_type(Q1,Q2,n)
#     R = diag_type(R1,R2,m)
#     QuadraticCost{n,m,T, Q, R}
# end

# @inline Base.convert(::Type{QC}, cost::QuadraticCost) where QC <: QuadraticCost = QC(cost)

# """
#     LQRCost(Q, R, xf, [uf; kwargs...])

# Convenience constructor for a `QuadraticCostFunction` of the form:
# ```math
# \\frac{1}{2} (x-x_f)^T Q (x-xf) + \\frac{1}{2} (u-u_f)^T R (u-u_f)
# ```

# If ``Q`` and ``R`` are diagonal, the output will be a `DiagonalCost`, otherwise it will
# be a `QuadraticCost`.
# """
# function LQRCost(Q::AbstractArray, R::AbstractArray,
#         xf::AbstractVector, uf=(@SVector zeros(size(R,1))); kwargs...)
#     H = @SMatrix zeros(size(R,1),size(Q,1))
#     q = -Q*xf
#     r = -R*uf
#     c = 0.5*xf'*Q*xf + 0.5*uf'R*uf
#     return DiagonalCost(Q, R, H, q, r, c, kwargs...)
# end

# function LQRCost(Q::Diagonal{<:Any,<:SVector{n}},R::Diagonal{<:Any,<:SVector{m}},
#         xf::AbstractVector, uf=(@SVector zeros(m)); kwargs...) where {n,m}
#     q = -Q*xf
#     r = -R*uf
#     c = 0.5*xf'*Q*xf + 0.5*uf'R*uf
#     return DiagonalCost(Q, R, q, r, c; kwargs...)
# end

# ############################################################################################
# #                        QUADRATIC QUATERNION COST FUNCTION
# ############################################################################################

# struct QuadraticQuatCost{T,N,M,N4} <: CostFunction
#     Q::Diagonal{T,SVector{N,T}}
#     R::Diagonal{T,SVector{M,T}}
#     q::SVector{N,T}
#     r::SVector{M,T}
#     c::T
#     w::T
#     q_ref::SVector{4,T}
#     q_ind::SVector{4,Int}
#     Iq::SMatrix{N,4,T,N4}
#     function QuadraticQuatCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}},
#             q::SVector{N,T}, r::SVector{M,T}, c::T, w::T,
#             q_ref::SVector{4,T}, q_ind::SVector{4,Int}) where {T,N,M}
#         Iq = @MMatrix zeros(N,4)
#         for i = 1:4
#             Iq[q_ind[i],i] = 1
#         end
#         Iq = SMatrix{N,4}(Iq)
#         return new{T,N,M,N*4}(Q, R, q, r, c, w, q_ref, q_ind, Iq)
#     end
# end

# state_dim(::QuadraticQuatCost{T,N,M}) where {T,N,M} = N
# control_dim(::QuadraticQuatCost{T,N,M}) where {T,N,M} = M

# function QuadraticQuatCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}};
#         q=(@SVector zeros(N)), r=(@SVector zeros(M)), c=zero(T), w=one(T),
#         q_ref=(@SVector [1.0,0,0,0]), q_ind=(@SVector [4,5,6,7])) where {T,N,M}
#     QuadraticQuatCost(Q, R, q, r, c, q_ref, q_ind)
# end

# function stage_cost(cost::QuadraticQuatCost, x::SVector, u::SVector)
#     stage_cost(cost, x) + 0.5*u'cost.R*u + cost.r'u
# end

# function stage_cost(cost::QuadraticQuatCost, x::SVector)
#     J = 0.5*x'cost.Q*x + cost.q'x + cost.c
#     q = x[cost.q_ind]
#     dq = cost.q_ref'q
#     J += cost.w*min(1+dq, 1-dq)
# end

# function gradient(cost::QuadraticQuatCost{T,N,M}, x::SVector, u::SVector) where {T,N,M}
#     Qx = cost.Q*x + cost.q
#     q = x[cost.q_ind]
#     dq = cost.q_ref'q
#     if dq < 0
#         Qx += cost.w*cost.Iq*cost.q_ref
#     else
#         Qx -= cost.w*cost.Iq*cost.q_ref
#     end
#     Qu = cost.R*u + cost.r
#     return Qx, Qu
# end

# function hessian(cost::QuadraticQuatCost, x::SVector{N}, u::SVector{M}) where {N,M}
#     Qxx = cost.Q
#     Quu = cost.R
#     Qux = @SMatrix zeros(M,N)
#     return Qxx, Quu, Qux
# end

# function QuatLQRCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}}, xf,
#         uf=(@SVector zeros(M)); w=one(T), quat_ind=(@SVector [4,5,6,7])) where {T,N,M}
#     r = -R*uf
#     q = -Q*xf
#     c = 0.5*xf'Q*xf + 0.5*uf'R*uf
#     q_ref = xf[quat_ind]
#     return QuadraticQuatCost(Q, R, q, r, c, w, q_ref, quat_ind)
# end

# function change_dimension(cost::QuadraticQuatCost, n, m)
#     n0,m0 = state_dim(cost), control_dim(cost)
#     Q_diag = diag(cost.Q)
#     R_diag = diag(cost.R)
#     q = cost.q
#     r = cost.r
#     if n0 != n
#         dn = n - n0  # assumes n > n0
#         pad = @SVector zeros(dn)
#         Q_diag = [Q_diag; pad]
#         q = [q; pad]
#     end
#     if m0 != m
#         dm = m - m0  # assumes m > m0
#         pad = @SVector zeros(dm)
#         R_diag = [R_diag; pad]
#         r = [r; pad]
#     end
#     QuadraticQuatCost(Diagonal(Q_diag), Diagonal(R_diag), q, r, cost.c, cost.w,
#         cost.q_ref, cost.q_ind)
# end

# function (+)(cost1::QuadraticQuatCost, cost2::QuadraticCost)
#     @assert state_dim(cost1) == state_dim(cost2)
#     @assert control_dim(cost1) == control_dim(cost2)
#     @assert norm(cost2.H) ≈ 0
#     QuadraticQuatCost(cost1.Q + cost2.Q, cost1.R + cost2.R,
#         cost1.q + cost2.q, cost1.r + cost2.r, cost1.c + cost2.c,
#         cost1.w, cost1.q_ref, cost1.q_ind)
# end

# (+)(cost1::QuadraticCost, cost2::QuadraticQuatCost) = cost2 + cost1

#
#
# struct ErrorQuadratic{Rot,N,M} <: CostFunction
#     model::RigidBody{Rot}
#     Q::Diagonal{Float64,SVector{12,Float64}}
#     R::Diagonal{Float64,SVector{M,Float64}}
#     r::SVector{M,Float64}
#     c::Float64
#     x_ref::SVector{N,Float64}
#     q_ind::SVector{4,Int}
# end
#
#
# state_dim(::ErrorQuadratic{Rot,N,M}) where {Rot,N,M} = N
# control_dim(::ErrorQuadratic{Rot,N,M}) where {Rot,N,M} = M
#
# function ErrorQuadratic(model::RigidBody{Rot}, Q::Diagonal{T,<:SVector{12}},
#         R::Diagonal{T,<:SVector{M}},
#         x_ref::SVector{N}, u_ref=(@SVector zeros(T,M)); r=(@SVector zeros(T,M)), c=zero(T),
#         q_ind=(@SVector [4,5,6,7])) where {T,N,M,Rot}
#     r += -R*u_ref
#     c += 0.5*u_ref'R*u_ref
#     return ErrorQuadratic{Rot,N,M}(model, Q, R, r, c, x_ref, q_ind)
# end
#
# function stage_cost(cost::ErrorQuadratic, x::SVector)
#     dx = state_diff(cost.model, x, cost.x_ref)
#     return 0.5*dx'cost.Q*dx + cost.c
# end
#
# function stage_cost(cost::ErrorQuadratic, x::SVector, u::SVector)
#     stage_cost(cost, x) + 0.5*u'cost.R*u + cost.r'u
# end
#
# function cost_expansion(cost::ErrorQuadratic{Rot}, model::AbstractModel,
#         z::KnotPoint{T,N,M}, G) where {T,N,M,Rot<:UnitQuaternion}
#     x,u = state(z), control(z)
#     model = cost.model
#     Q = cost.Q
#     q = orientation(model, x)
#     q_ref = orientation(model, cost.x_ref)
#     dq = SVector(q_ref\q)
#     err = state_diff(model, x, cost.x_ref)
#     dx = @SVector [err[1],  err[2],  err[3],
#                     dq[1],   dq[2],   dq[3],   dq[4],
#                    err[7],  err[8],  err[9],
#                    err[10], err[11], err[12]]
#     G = state_diff_jacobian(model, dx) # n × dn
#
#     # Gradient
#     dmap = inverse_map_jacobian(model, dx) # dn × n
#     Qx = G'dmap'Q*err
#     Qu = cost.R*u
#
#     # Hessian
#     ∇jac = inverse_map_∇jacobian(model, dx, Q*err)
#     Qxx = G'dmap'Q*dmap*G + G'∇jac*G + ∇²differential(model, x, dmap'Q*err)
#     Quu = cost.R
#     Qux = @SMatrix zeros(M,N-1)
#     return Qxx, Quu, Qux, Qx, Qu
# end
#
# function cost_expansion(cost::ErrorQuadratic, model::AbstractModel,
#         z::KnotPoint{T,N,M}, G) where {T,N,M}
#     x,u = state(z), control(z)
#     model = cost.model
#     q = orientation(model, x)
#     q_ref = orientation(model, cost.x_ref)
#     err = state_diff(model, x, cost.x_ref)
#     dx = err
#     G = state_diff_jacobian(model, dx) # n × n
#
#     # Gradient
#     dmap = inverse_map_jacobian(model, dx) # n × n
#     Qx = G'dmap'cost.Q*err
#     Qu = cost.R*u + cost.r
#
#     # Hessian
#     Qxx = G'dmap'cost.Q*dmap*G
#     Quu = cost.R
#     Qux = @SMatrix zeros(M,N)
#     return Qxx, Quu, Qux, Qx, Qu
# end
#
# function change_dimension(cost::ErrorQuadratic, n, m)
#     n0,m0 = state_dim(cost), control_dim(cost)
#     Q_diag = diag(cost.Q)
#     R_diag = diag(cost.R)
#     r = cost.r
#     if n0 != n
#         dn = n - n0  # assumes n > n0
#         pad = @SVector zeros(dn) # assume the new states don't have quaternions
#         Q_diag = [Q_diag; pad]
#     end
#     if m0 != m
#         dm = m - m0  # assumes m > m0
#         pad = @SVector zeros(dm)
#         R_diag = [R_diag; pad]
#         r = [r; pad]
#     end
#     ErrorQuadratic(cost.model, Diagonal(Q_diag), Diagonal(R_diag), r, cost.c,
#         cost.x_ref, cost.q_ind)
# end
#
# function (+)(cost1::ErrorQuadratic, cost2::QuadraticCost)
#     @assert control_dim(cost1) == control_dim(cost2)
#     @assert norm(cost2.H) ≈ 0
#     @assert norm(cost2.q) ≈ 0
#     if state_dim(cost2) == 13
#         rm_quat = @SVector [1,2,3,4,5,6,8,9,10,11,12,13]
#         Q2 = Diagonal(diag(cost2.Q)[rm_quat])
#     else
#         Q2 = cost2.Q
#     end
#     ErrorQuadratic(cost1.model, cost1.Q + Q2, cost1.R + cost2.R,
#         cost1.r + cost2.r, cost1.c + cost2.c,
#         cost1.x_ref, cost1.q_ind)
# end
#
# (+)(cost1::QuadraticCost, cost2::ErrorQuadratic) = cost2 + cost1
#
#
#
#


using MathOptInterface
const MOI = MathOptInterface
const MOIT = MOI.Test

const OPTIMIZER = SCIP.Optimizer(display_verblevel=0)
const CONFIG_DIRECT = MOIT.Config(
    atol=1e-5, rtol=1e-5,
    exclude=Any[
        MOI.ConstraintDual, MOI.ConstraintName, MOI.DualObjectiveValue, MOI.VariableBasisStatus, MOI.ConstraintBasisStatus,
    ],
)

@testset "MOI unit tests" begin
    excluded = copy(MOI_BASE_EXCLUDED)
    append!(excluded, [
        "test_linear_integration", # Can not delete variable while model contains constraints
        "test_basic_VectorOfVariables_SecondOrderCone",
        "test_conic_SecondOrderCone_nonnegative_post_bound",
        "test_variable_delete_SecondOrderCone",
        "test_modification_func_scalaraffine_",
        "test_modification_func_vectoraffine_",
    ])
    MOIT.runtests(
        OPTIMIZER,
        CONFIG_DIRECT,
        warn_unsupported=false,
        exclude = excluded,
    )

end

function indicator4_test(model::MOI.ModelLike, config)
    atol = config.atol
    rtol = config.rtol
    # equivalent to indicator1_test with left-hand-side partially in LHS constant
    # linear problem with indicator constraint and
    # max  2x1 + 3x2
    # s.t. x1 + x2 <= 10
    #      z1 ==> x2 - 1 <= 7
    #      z2 ==> x2 + x1/5 + 1 <= 10
    #      z1 + z2 >= 1

    MOI.empty!(model)
    @test MOI.is_empty(model)

    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    @test MOI.supports_constraint(model, MOI.VariableIndex, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.VariableIndex, MOI.Interval{Float64})
    @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.Interval{Float64})
    @test MOI.supports_constraint(model, MOI.VectorAffineFunction{Float64}, MOI.Indicator{MOI.ACTIVATE_ON_ONE, MOI.LessThan{Float64}})
    x1 = MOI.add_variable(model)
    x2 = MOI.add_variable(model)
    z1 = MOI.add_variable(model)
    z2 = MOI.add_variable(model)
    MOI.add_constraint(model, z1, MOI.ZeroOne())
    MOI.add_constraint(model, z2, MOI.ZeroOne())
    f1 = MOI.VectorAffineFunction(
        [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z1)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, -1.0]
    )
    iset1 = MOI.Indicator{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(7.0))
    MOI.add_constraint(model, f1, iset1)

    f2 = MOI.VectorAffineFunction(
        [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, z2)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(0.2, x1)),
         MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(1.0, x2)),
        ],
        [0.0, 1.0],
    )
    iset2 = MOI.Indicator{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(10.0))

    MOI.add_constraint(model, f2, iset2)

    # Additional regular constraint.
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x1), MOI.ScalarAffineTerm(1.0, x2)], 0.0),
        MOI.LessThan(10.0),
    )

    # Disjunction z1 ⋁ z2
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, z1), MOI.ScalarAffineTerm(1.0, z2)], 0.0),
        MOI.GreaterThan(1.0),
    )

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0], [x1, x2]), 0.0)
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED

    MOI.optimize!(model)

    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 28.75 atol=atol rtol=rtol
    @test MOI.get(model, MOI.VariablePrimal(), x1) ≈ 1.25 atol=atol rtol=rtol
    @test MOI.get(model, MOI.VariablePrimal(), x2) ≈ 8.75 atol=atol rtol=rtol
    @test MOI.get(model, MOI.VariablePrimal(), z1) ≈ 0.0 atol=atol rtol=rtol
    @test MOI.get(model, MOI.VariablePrimal(), z2) ≈ 1.0 atol=atol rtol=rtol
end

@testset "MOI Integer Linear" begin
    # replace with MOIT when MathOptInterface.jl#929 merged
    indicator4_test(OPTIMIZER, CONFIG_DIRECT)
end


#*******************************************************/
#* Copyright(c) 2018 by Artelys                        */
#* This source code is subject to the terms of the     */
#* MIT Expat License (see LICENSE.md)                  */
#*******************************************************/

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# This example demonstrates how to use Knitro to solve the following
# simple mathematical program with equilibrium/complementarity
# constraints(MPEC/MPCC).
#
#  min  (x0 - 5)^2 +(2 x1 + 1)^2
#  s.t.  -1.5 x0 + 2 x1 + x2 - 0.5 x3 + x4 = 2
#        x2 complements(3 x0 - x1 - 3)
#        x3 complements(-x0 + 0.5 x1 + 4)
#        x4 complements(-x0 - x1 + 7)
#        x0, x1, x2, x3, x4 >= 0
#
# The complementarity constraints must be converted so that one
# nonnegative variable complements another nonnegative variable.
#
#  min  (x0 - 5)^2 +(2 x1 + 1)^2
#  s.t.  -1.5 x0 + 2 x1 + x2 - 0.5 x3 + x4 = 2  (c0)
#        3 x0 - x1 - 3 - x5 = 0                 (c1)
#        -x0 + 0.5 x1 + 4 - x6 = 0              (c2)
#        -x0 - x1 + 7 - x7 = 0                  (c3)
#        x2 complements x5
#        x3 complements x6
#        x4 complements x7
#        x0, x1, x2, x3, x4, x5, x6, x7 >= 0
#
# The solution is(1, 0, 3.5, 0, 0, 0, 3, 6), with objective value 17.
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

using KNITRO, Test

function example_mpec1(; verbose=true)
    # Create a new Knitro solver instance.
    kc = KNITRO.KN_new()

    # Illustrate how to override default options by reading from
    # the knitro.opt file.
    options = joinpath(dirname(@__FILE__), "..", "examples", "knitro.opt")
    KNITRO.KN_load_param_file(kc, options)

    # Initialize Knitro with the problem definition.

    # Add the variables and set their bounds and initial values.
    # Note: unset bounds assumed to be infinite.
    KNITRO.KN_add_vars(kc, 8)
    KNITRO.KN_set_var_lobnds(kc, zeros(Float64, 8))
    KNITRO.KN_set_var_primal_init_values(kc, zeros(Float64, 8))

    # Add the constraints and set their bounds.
    KNITRO.KN_add_cons(kc, 4)
    KNITRO.KN_set_con_eqbnds(kc, Float64[2, 3, -4, -7])

    # Add coefficients for all linear constraints at once.

    # c0
    lconIndexCons = Int32[0, 0, 0, 0, 0]
    lconIndexVars = Int32[0, 1, 2, 3, 4]
    lconCoefs = [-1.5, 2.0, 1.0, -0.5, 1.0]

    # c1
    lconIndexCons = [lconIndexCons; Int32[1, 1, 1]]
    lconIndexVars = [lconIndexVars; Int32[0, 1, 5]]
    lconCoefs = [lconCoefs; [3.0, -1.0, -1.0]]

    # c2
    lconIndexCons = [lconIndexCons; Int32[2, 2, 2]]
    lconIndexVars = [lconIndexVars; Int32[0, 1, 6]]
    lconCoefs = [lconCoefs; [-1.0, 0.5, -1.0]]

    # c3
    lconIndexCons = [lconIndexCons; Int32[3, 3, 3]]
    lconIndexVars = [lconIndexVars; Int32[0, 1, 7]]
    lconCoefs = [lconCoefs; [-1.0, -1.0, -1.0]]

    KNITRO.KN_add_con_linear_struct(kc, lconIndexCons, lconIndexVars, lconCoefs)

    # Note that the objective(x0 - 5)^2 +(2 x1 + 1)^2 when
    # expanded becomes:
    #    x0^2 + 4 x1^2 - 10 x0 + 4 x1 + 26

    # Add quadratic coefficients for the objective
    qobjIndexVars1 = Int32[0, 1]
    qobjIndexVars2 = Int32[0, 1]
    qobjCoefs = [1.0, 4.0]
    KNITRO.KN_add_obj_quadratic_struct(kc, qobjIndexVars1, qobjIndexVars2, qobjCoefs)

    # Add linear coefficients for the objective
    lobjIndexVars = Int32[0, 1]
    lobjCoefs = [-10.0, 4.0]
    KNITRO.KN_add_obj_linear_struct(kc, lobjIndexVars, lobjCoefs)

    # Add constant to the objective
    KNITRO.KN_add_obj_constant(kc, 26.0)

    # Set minimize or maximize(if not set, assumed minimize)
    KNITRO.KN_set_obj_goal(kc, KNITRO.KN_OBJGOAL_MINIMIZE)

    # Now add the complementarity constraints
    ccTypes = [KNITRO.KN_CCTYPE_VARVAR, KNITRO.KN_CCTYPE_VARVAR, KNITRO.KN_CCTYPE_VARVAR]
    indexComps1 = Int32[2, 3, 4]
    indexComps2 = Int32[5, 6, 7]
    KNITRO.KN_set_compcons(kc, ccTypes, indexComps1, indexComps2)

    kn_outlev = verbose ? KNITRO.KN_OUTLEV_ALL : KNITRO.KN_OUTLEV_NONE
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_OUTLEV, kn_outlev)

    # Solve the problem.
    #
    # Return status codes are defined in "knitro.h" and described
    # in the Knitro manual.
    nStatus = KNITRO.KN_solve(kc)

    # An example of obtaining solution information.
    nStatus, objSol, x, lambda_ = KNITRO.KN_get_solution(kc)

    if verbose
        println("Knitro converged with final status = ", nStatus)
        println("  optimal objective value  = ", objSol)
        println("  optimal primal values x0=", x[1])
        println("                        x1=", x[2])
        println("                        x2=", (x[3], x[6]))
        println("                        x3=", (x[4], x[7]))
        println("                        x4=", (x[5], x[8]))
        println("  feasibility violation    = ", KNITRO.KN_get_abs_feas_error(kc))
        println("  KKT optimality violation = ", KNITRO.KN_get_abs_opt_error(kc))
    end

    # Delete the Knitro solver instance.
    KNITRO.KN_free(kc)

    @testset "Example MPEC 1" begin
        @test nStatus == 0
        @test objSol ≈ 17.
        @test x ≈ [1., 0., 3.5, 0., 0., 0., 3., 6.]
    end
end

example_mpec1(; verbose=isdefined(Main, :KN_VERBOSE) ? KN_VERBOSE : true)


#the majority of the following code is inspired by
#the DataFrames.jl package (https://github.com/JuliaStats/DataFrames.jl)
#DataFrames.jl is licensed under the MIT License, a copy of which
#can be found in this package's repo in LICENSE.md.

#This file contains functions that are useful for crossjoining
#arbitrary arrays, and are thus useful for building a tensor product
#structure.

function rep(x::AbstractVector, times::Integer = 1, each::Integer = 1)
    res = similar(x, each * times * length(x))
    i = 1
    for jdx in 1:times
        for idx in 1:length(x)
            tmp = x[idx]
            for kdx in 1:each
                res[i] = tmp
                i += 1
            end
        end
    end
    return res
end

function rep(x::AbstractArray, times::Integer = 1, each::Integer = 1)
    res = similar(x, each * times * size(x,1), size(x,2))
    i = 1
    for jdx in 1:times
        for idx in 1:size(x,1)
            tmp = x[idx,:]
            for kdx in 1:each
                res[i,:] = tmp
                i += 1
            end
        end
    end
    return res
end

crossjoin(A::Array, B::Array) = hcat(rep(A, 1, size(B,1)), rep(B, size(A,1), 1))

crossjoin(A::Array...) = reduce(crossjoin, A)


module TestGAT

using Base.Meta: ParseError
using Test
using DataStructures: OrderedDict

import Catlab: GAT
using Catlab.GAT

# GAT expressions
#################

# Raw expressions
@test GAT.parse_raw_expr(:(Ob)) == :Ob
@test GAT.parse_raw_expr(:(Hom(X,Y))) == :(Hom(X,Y))
@test_throws ParseError GAT.parse_raw_expr(:("Ob"))
@test_throws ParseError GAT.parse_raw_expr(:(Hom(X,0)))

# Contexts
@test (GAT.parse_context(:((X::Ob, Y::Ob))) ==
       GAT.Context((:X => :Ob, :Y => :Ob)))
@test (GAT.parse_context(:((X::Ob, Y::Ob, f::Hom(X,Y)))) ==
       GAT.Context((:X => :Ob, :Y => :Ob, :f => :(Hom(X,Y)))))
@test GAT.parse_context(:(())) == GAT.Context()
@test_throws ParseError GAT.parse_context(:((X::Ob, X::Ob))) # Repeat variables

# Type constructor
expr = :(Ob::TYPE)
cons = GAT.TypeConstructor(:Ob, [], GAT.Context())
@test GAT.parse_constructor(expr) == cons

expr = (quote "Object" Ob::TYPE end).args[2]
cons = GAT.TypeConstructor(:Ob, [], GAT.Context(), "Object")
@test GAT.parse_constructor(expr) == cons

expr = :(Hom(X,Y)::TYPE ⊣ (X::Ob, Y::Ob))
context = GAT.Context((:X => :Ob, :Y => :Ob))
cons = GAT.TypeConstructor(:Hom, [:X,:Y], context)
@test GAT.parse_constructor(expr) == cons

# Term constructor
expr = :(unit()::Ob)
cons = GAT.TermConstructor(:unit, [], :Ob, GAT.Context())
@test GAT.parse_constructor(expr) == cons

expr = (quote "Monoidal unit" munit()::Ob end).args[2]
cons = GAT.TermConstructor(:munit, [], :Ob, GAT.Context(), "Monoidal unit")
@test GAT.parse_constructor(expr) == cons

cons = GAT.TermConstructor(:id, [:X], :(Hom(X,X)), GAT.Context(:X => :Ob))
@test GAT.parse_constructor(:(id(X)::Hom(X,X) ⊣ (X::Ob))) == cons
@test GAT.parse_constructor(:(id(X::Ob)::Hom(X,X))) == cons

expr = :(compose(f,g)::Hom(X,Z) ⊣ (X::Ob, Y::Ob, Z::Ob, f::Hom(X,Y), g::Hom(Y,Z)))
context = GAT.Context((:X => :Ob, :Y => :Ob, :Z => :Ob,
                       :f => :(Hom(X,Y)), :g => :(Hom(Y,Z))))
cons = GAT.TermConstructor(:compose, [:f,:g], :(Hom(X,Z)), context)
@test GAT.parse_constructor(expr) == cons
expr = :(compose(f::Hom(X,Y), g::Hom(Y,Z))::Hom(X,Z) ⊣ (X::Ob, Y::Ob, Z::Ob))
@test GAT.parse_constructor(expr) == cons

# Type transformations
bindings = Dict((:Ob => :Obj, :Hom => :Mor))
cons = GAT.TypeConstructor(:Hom, [:X,:Y],
  GAT.Context((:X => :Ob, :Y => :Ob)))
target = GAT.TypeConstructor(:Mor, [:X,:Y],
  GAT.Context((:X => :Obj, :Y => :Obj)))
@test GAT.replace_types(bindings, cons) == target

cons = GAT.TermConstructor(:compose, [:f,:g], :(Hom(X,Z)),
  GAT.Context((:X => :Ob, :Y => :Ob, :Z => :Ob,
               :f => :(Hom(X,Y)), :g => :(Hom(Y,Z)))))
target = GAT.TermConstructor(:compose, [:f,:g], :(Mor(X,Z)),
  GAT.Context((:X => :Obj, :Y => :Obj, :Z => :Obj,
               :f => :(Mor(X,Y)), :g => :(Mor(Y,Z)))))
@test GAT.replace_types(bindings, cons) == target

cons = GAT.AxiomConstructor(:(==), Meta.parse("compose(compose(f,g),Hom(C,D))"),
  Meta.parse("compose(f,compose(g,Hom(C,D)))"),
  GAT.Context((:A => :Ob, :B => :Ob, :C => :Ob, :D => :Ob,
               :f => :(Hom(A,B)), :g => :(Hom(B,C)))))
target = GAT.AxiomConstructor(:(==), Meta.parse("compose(compose(f,g),Mor(C,D))"),
  Meta.parse("compose(f,compose(g,Mor(C,D)))"),
  GAT.Context((:A => :Obj, :B => :Obj, :C => :Obj, :D => :Obj,
               :f => :(Mor(A,B)), :g => :(Mor(B,C)))))
@test GAT.replace_types(bindings, cons) == target

cons = Dict(:→ => :Hom)
target = Dict(:→ => :Mor)
@test GAT.replace_types(bindings, cons) == target

@test GAT.strip_type(:Ob) == :Ob
@test GAT.strip_type(:(Hom(X,Y))) == :Hom
@test GAT.strip_type(:(Hom(dual(X),dual(Y)))) == :Hom

# Theories
############

# This try-catch block is a necessary work around because of a current bug
# where test_throws doesn't catch errors thrown from inside of a macro
@test_throws ParseError try @eval @signature Category(Ob,Hom) begin
  Ob::TYPE
  Hom(dom, codom)::TYPE ⊣ (dom::Ob, codom::Ob)
  @op (→) := Hom

  id(X)::(X → X) ⊣ (X::Ob)
  compose(f,g)::(X → Z) ⊣ (X::Ob, Y::Ob, Z::Ob, f::(X → Y), g::(Y → Z))
  @op (⋅) := compose

  (f ⋅ g) ⋅ h == f ⋅ (g ⋅ h) ⊣ ( A::Ob, B::Ob, C::Ob, D::Ob,
                                f::(A → B), g::(B → C), h::(C → D))
  f ⋅ id(B) == f ⊣ (A::Ob, B::Ob, f::(A → B))
  id(A) ⋅ f == f ⊣ (A::Ob, B::Ob, f::(A → B))
end
catch err;
  throw(err.error)
end

""" Theory of categories
"""
@theory Category(Ob,Hom) begin
  Ob::TYPE
  Hom(dom, codom)::TYPE ⊣ (dom::Ob, codom::Ob)
  @op (→) := Hom

  id(X)::(X → X) ⊣ (X::Ob)
  compose(f,g)::(X → Z) ⊣ (X::Ob, Y::Ob, Z::Ob, f::(X → Y), g::(Y → Z))
  @op (⋅) := compose

  (f ⋅ g) ⋅ h == f ⋅ (g ⋅ h) ⊣ ( A::Ob, B::Ob, C::Ob, D::Ob,
                                f::(A → B), g::(B → C), h::(C → D))
  f ⋅ id(B) == f ⊣ (A::Ob, B::Ob, f::(A → B))
  id(A) ⋅ f == f ⊣ (A::Ob, B::Ob, f::(A → B))
end

@test Category isa Type
@test occursin("theory of categories", lowercase(string(Docs.doc(Category))))
@test isempty(methods(dom)) && isempty(methods(codom))
@test isempty(methods(id)) && isempty(methods(compose))

# Manually constructed theory of categories
types = [
  GAT.TypeConstructor(:Ob, [], GAT.Context()),
  GAT.TypeConstructor(:Hom, [:dom,:codom],
    GAT.Context((:dom => :Ob, :codom => :Ob))),
]
terms = [
  GAT.TermConstructor(:id, [:X], :(Hom(X,X)), GAT.Context(:X => :Ob)),
  GAT.TermConstructor(:compose, [:f,:g], :(Hom(X,Z)),
    GAT.Context((:X => :Ob, :Y => :Ob, :Z => :Ob,
                 :f => :(Hom(X,Y)), :g => :(Hom(Y,Z))))),
]
axioms = [
  GAT.AxiomConstructor(:(==), Meta.parse("compose(compose(f,g),h)"),
    Meta.parse("compose(f,compose(g,h))"),
    GAT.Context((:A => :Ob, :B => :Ob, :C => :Ob, :D => :Ob,
                 :f => :(Hom(A,B)), :g => :(Hom(B,C)), :h => :(Hom(C,D))))),
  GAT.AxiomConstructor(:(==), Meta.parse("compose(f,id(B))"), :f,
    GAT.Context((:A => :Ob, :B => :Ob, :f => :(Hom(A,B))))),
  GAT.AxiomConstructor(:(==), Meta.parse("compose(id(A),f)"), :f,
    GAT.Context((:A => :Ob, :B => :Ob, :f => :(Hom(A,B))))),
]
aliases = Dict(:⋅ => :compose, :→ => :Hom)
category_theory = GAT.Theory(types, terms, axioms, aliases)

@test GAT.theory(Category) == category_theory

""" Equivalent shorthand definition of Category theory
"""
@theory CategoryAbbrev(Ob,Hom) begin
  @op begin
    (→) := Hom
    (⋅) := compose
  end

  Ob::TYPE
  Hom(dom::Ob, codom::Ob)::TYPE

  id(X::Ob)::(X → X)
  (compose(f::(X → Y),g::(Y → Z))::(X → Z)) where (X::Ob, Y::Ob, Z::Ob)

  (f ⋅ g) ⋅ h == f ⋅ (g ⋅ h) ⊣ ( A::Ob, B::Ob, C::Ob, D::Ob,
                                f::(A → B), g::(B → C), h::(C → D))
  f ⋅ id(B) == f ⊣ (A::Ob, B::Ob, f::(A → B))
  id(A) ⋅ f == f ⊣ (A::Ob, B::Ob, f::(A → B))
end

@test GAT.theory(CategoryAbbrev) == category_theory

# Methods for theory
accessors = [ GAT.JuliaFunction(:(dom(::Hom)), :Ob),
              GAT.JuliaFunction(:(codom(::Hom)), :Ob) ]
constructors = [ GAT.JuliaFunction(:(id(X::Ob)), :Hom),
                 GAT.JuliaFunction(:(compose(f::Hom, g::Hom)), :Hom) ]
alias_functions = [
  GAT.JuliaFunction(:(⋅(f::Hom, g::Hom)), :Hom, :(compose(f, g))),
  GAT.JuliaFunction(:(→(dom::Ob, codom::Ob)), :Hom, :(Hom(dom, codom))),
]
theory = GAT.theory(Category)
@test GAT.accessors(theory) == accessors
@test GAT.constructors(theory) == constructors
@test GAT.alias_functions(theory) == alias_functions
@test GAT.interface(theory) == [accessors; constructors; alias_functions]

# Theory extension
@signature Semigroup(S) begin
  S::TYPE
  times(x::S,y::S)::S
end

@signature Semigroup(M) => MonoidExt(M) begin
  munit()::M
end

@test Semigroup isa Type && MonoidExt isa Type

theory = GAT.Theory(
  [ GAT.TypeConstructor(:M, [], GAT.Context()) ],
  [ GAT.TermConstructor(:times, [:x,:y], :M,
      GAT.Context((:x => :M, :y => :M))),
    GAT.TermConstructor(:munit, [], :M, GAT.Context()) ],
  [],
  Dict{Symbol,Symbol}()
)

@test GAT.theory(MonoidExt) == theory

# GAT expressions in a theory
################################

theory = GAT.theory(Category)
context = GAT.Context((:X => :Ob, :Y => :Ob, :Z => :Ob,
                       :f => :(Hom(X,Y)), :g => :(Hom(Y,Z))))
@test GAT.expand_in_context(:X, [:f,:g], context, theory) == :(dom(f))
@test (GAT.expand_in_context(:(Hom(X,Z)), [:f,:g], context, theory) ==
       :(Hom(dom(f),codom(g))))

context = GAT.Context((:X => :Ob, :Y => :Ob, :Z => :Ob, :f => :(Hom(X,Y))))
@test_throws ErrorException GAT.expand_in_context(:W, [:f], context, theory)
@test_throws ErrorException GAT.expand_in_context(:Z, [:f], context, theory)

context = GAT.Context((:X => :Ob, :Y => :Ob, :f => :(Hom(X,Y))))
@test GAT.equations(context, theory) == [ :(dom(f)) => :X, :(codom(f)) => :Y ]
@test GAT.equations([:f], context, theory) == []

context = GAT.Context((:X => :Ob, :Y => :Ob, :Z => :Ob,
                       :f => :(Hom(X,Y)), :g => :(Hom(Y,Z))))
@test (GAT.equations(context, theory) ==
       [ :(dom(f)) => :X, :(codom(f)) => :Y,
         :(dom(g)) => :Y, :(codom(g)) => :Z ])
@test GAT.equations([:f,:g], context, theory) == [ :(dom(g)) => :(codom(f)) ]

# Instances
###########

""" Vectors as an instance of the theory of semigroups
"""
@instance Semigroup(Vector) begin
  times(x::Vector, y::Vector) = [x; y]
end

@test times([1,2],[3,4]) == [1,2,3,4]

@signature Monoid(M) begin
  M::TYPE
  munit()::M
  times(x::M,y::M)::M
end

# Incomplete instance of Monoid
# XXX: Cannot use `@test_warn` since generated code won't be at toplevel.
#@test_warn "not implemented" @instance Monoid(String) begin
#  times(x::AbsStringtractString, y::String) = string(x,y)
#end

# Complete instance of Monoid
@instance Monoid(String) begin
  munit(::Type{String}) = ""
  times(x::String, y::String) = string(x,y)
end

@test munit(String) == ""
@test times("a", "b") == "ab"

# Reflection
@test invoke_term(Monoid, (String,), :munit) == ""
@test invoke_term(Monoid, (String,), :times, "a", "b") == "ab"

end


export CayleyConfiguration, CircuitTable, MixedCell, circuitdot

const CellIndices = Vector{Tuple{Int,Int}} # Maybe also SVector{N, Tuple{T, T}}

"""
    CayleyConfiguration{I<:Integer}

The Cayley matrix of a set of matrices ``A_1, …, A_n`` where ``A_i ∈ 𝐑^{n×mᵢ}``.
"""
struct CayleyConfiguration{I<:Integer} <: AbstractMatrix{I}
    A::Matrix{I}
    offsets::Vector{Int} # Maybe also SVector{N, Int}
end

function CayleyConfiguration(Aᵢs::Matrix{I}...) where {I<:Integer}
    n = size(Aᵢs[1], 1)
    @assert all(Aᵢ -> size(Aᵢ, 1) == n, Aᵢs) "Matrices do not have the same number of rows."

    A = fill(zero(I), 2n, sum(Aᵢ -> size(Aᵢ, 2), Aᵢs))

    offsets = Int[]
    offset = 0
    for (i, Aᵢ) in enumerate(Aᵢs)
        push!(offsets, offset)
        mᵢ = size(Aᵢ, 2)

        for j=1:mᵢ
            jj = offset + j
            for k=1:n
                A[k, jj] = Aᵢ[k, j]
            end
            A[n+i, jj] = one(I)
        end

        offset += mᵢ
    end
    CayleyConfiguration(A, offsets)
end

"""
    linearindex(CC::CayleyConfiguration, i::Integer, col::Integer)

Get the linear index into the Caylex matrix corresponding to the configuration `i` and column `col`.
"""
function linearindex(CC::CayleyConfiguration, configuration::Integer, col::Integer)
    CC.offsets[configuration] + col
end

"""
    linearindices(CC::CayleyConfiguration, cellindices::CellIndices)

Get the linear indexes into the Caylex matrix corresponding to the cell indices.
"""
function linearindices(A::CayleyConfiguration, cellindices::CellIndices)
    linearindices = Int[]
    for i=1:length(cellindices)
        aᵢ, bᵢ = cellindices[i]
        push!(linearindices, linearindex(A, i, aᵢ), linearindex(A, i, bᵢ))
    end
    linearindices
end

Base.getindex(CC::CayleyConfiguration, i::Int) = getindex(CC.A, i)
Base.getindex(CC::CayleyConfiguration, I::Vararg{Int, 2}) = getindex(CC.A, I...)
Base.size(CC::CayleyConfiguration) = size(CC.A)
Base.size(CC::CayleyConfiguration, I) = size(CC.A, I)


using PyCall
pushfirst!(PyVector(pyimport("sys")["path"]), "")
@pyimport src.route.pydomain as pd

"""
    haversine(lon1, lat1, lon2, lat2)

Calculate the haversine distance and bearing. Distance is in nm.
"""
function haversine(lon1::Float64, lat1::Float64, lon2::Float64, lat2::Float64)
    R = 6372.8  # Earth radius in kilometers

    dLat = deg2rad(lat2 - lat1)
    dLon = deg2rad(lon2 - lon1)
    lat1 = deg2rad(lat1)
    lat2 = deg2rad(lat2)
    lon1 = deg2rad(lon1)
    lon2 = deg2rad(lon2)
    a = sin(dLat/2)^2 + cos(lat1)*cos(lat2)*sin(dLon/2)^2
    c = 2*asin(sqrt(a))
    theta = atan(sin(dLon)*cos(lat2),
                 cos(lat1)*sin(lat2)-sin(lat1)*cos(lat2)*cos(dLon))
    theta = (rad2deg(theta) + 360) % 360
    return R*c*0.5399565, theta
end

"""
    co_ordinates(start_long, finish_long, start_lat, finish_lat,
                 n_ranks, n_nodes, dist)

Return the co-ordinates of each point across the discretized domain.
"""
function co_ordinates(start_long, finish_long, start_lat, finish_lat,
                      n_ranks::Int, n_nodes::Int, dist)
    x, y, land = pd.return_co_ords(start_long, finish_long, start_lat,
                                   finish_lat, n_ranks, n_nodes,
                                   dist)
    return x, y, land
end

function output()
    x, y, land = co_ordinates(-6.0, -14.0, 0.0, 0.0, 10, 10, 5000.0)
    print(x)
    print(y)
    print(land)
end

function min_angle(a::Float64, b::Float64)
    abs(mod(a - b + 180.0, 360.0) - 180.0)
end


# one-column version
function op_H( LF, Potentials, v::Array{Float64,1} )
  #
  V_loc = Potentials.Ionic + Potentials.Hartree + Potentials.XC
  #
  return -0.5*op_nabla2( LF, v ) + V_loc .* v
end

# multi-column version
function op_H( LF, Potentials, v::Array{Float64,2} )
  #
  Npoints = size(v)[1]
  Ncols   = size(v)[2]
  V_loc = Potentials.Ionic + Potentials.Hartree + Potentials.XC
  #
  Hv = zeros( Float64, Npoints, Ncols )
  #
  for ic = 1:Ncols
    Hv[:,ic] = -0.5*op_nabla2( LF, v[:,ic] ) + V_loc .* v[:,ic]
  end
  return Hv
end

# one-column version, sparse
function op_H( LF, ∇2::SparseMatrixCSC{Float64,Int64},
                            Potentials, v::Array{Float64,1} )
  #
  V_loc = Potentials.Ionic + Potentials.Hartree + Potentials.XC
  return -0.5*∇2*v + V_loc.*v
end

# multi-column version, sparse
function op_H( LF, ∇2, Potentials, v::Array{Float64,2} )
  #
  Npoints = size(v)[1]
  Ncols   = size(v)[2]
  V_loc = Potentials.Ionic + Potentials.Hartree + Potentials.XC
  #
  Hv = zeros( Float64, Npoints, Ncols )
  #
  for ic = 1:Ncols
    Hv[:,ic] = -0.5*∇2*v[:,ic] + V_loc.*v[:,ic]
  end
  return Hv
end


#=
Author : Ng Yeesian
=#
type DSMDeployment <: DeploymentModel
    m::JuMP.Model
    x::Vector{JuMP.Variable}
end
deployment(m::StochasticDeployment) = [round(Int,x) for x in JuMP.getValue(m.x)]

function DSMDeployment(p::DeploymentProblem,
                          q::Float64; # busy fraction
                          max_amb::Int = 5
                          tol = params.δ,
                          solver = GurobiSolver(OutputFlag=0))
    demand = mean(p.demand[p.train,:],1)
    @assert length(demand) == problem.nregions
    I = 1:p.nlocations
    J = 1:problem.nregions
    K = 1:max_amb

    m = JuMP.Model(solver=solver)
    JuMP.@variable(m, x[1:p.nlocations] >= 0, Int)
    JuMP.@variable(m, z[1:p.nregions, 1:max_amb] >= 0, Bin)

    JuMP.@objective(m, Min, sum((1-q)*(q^k)*demand[j]*z[j,k] for j=J, k=K))

    JuMP.@constraint(m, sum(x[i] for i=I) <= p.nambulances)
    for j in J # coverage over all regions
        JuMP.@constraint(m, sum(x[i] for i in filter(i->p.coverage[j,i], I)) >= sum(z[j,k] for k=1:K))
    end

    DSMDeployment(m, x)
end

solve(model::DSMDeployment) = JuMP.solve(model.m)
solve(model::DSMDeployment, p::DeploymentProblem) = JuMP.solve(model.m)