text:
type MosekLinearQuadraticModel <: MathProgBase.AbstractLinearQuadraticModel
    task :: Mosek.MSKtask

    binvarflags:: Array{Bool,1}
    # NOTE: bkx/blx/bux are the bound on the variable in the
    # continuous problem. Setting a variable to :Bin, will not change
    # these. Setting :Cont or :Int on a :Bin variable will effectively
    # revert to the original bounds.
    numvar     :: Int
    numcon     :: Int
    bkx        :: Array{Int32,1}
    blx        :: Array{Float64,1}
    bux        :: Array{Float64,1}

    bkc        :: Array{Int32,1}
    blc        :: Array{Float64,1}
    buc        :: Array{Float64,1}
    lincon     :: Array{Int32,1}
    quadcon    :: Array{Int32,1}

    lasttrm    :: Int32

    options
end

type MosekNonlinearModel <: MathProgBase.AbstractLinearQuadraticModel
    m :: MosekLinearQuadraticModel
end

MathProgBase.LinearQuadraticModel(s::MosekSolver) =
  MosekLinearQuadraticModel(Mosek.maketask(),
                            Array(Bool,0),
                            0,
                            0,
                            Array(Int32,0),
                            Array(Float64,0),
                            Array(Float64,0),

                            Array(Int32,0),
                            Array(Float64,0),
                            Array(Float64,0),

                            Array(Int32,0),
                            Array(Int32,0),

                            Mosek.MSK_RES_OK,

                            s.options)

##############################################################
## Linear
##############################################################

function MathProgBase.loadproblem!{T1,T2,T3,T4,T5}(m::MosekLinearQuadraticModel,
                                                   A,
                                                   collb :: Array{T1,1},
                                                   colub :: Array{T2,1},
                                                   obj   :: Array{T3,1},
                                                   rowlb :: Array{T4,1},
                                                   rowub :: Array{T5,1},
                                                   sense:: Symbol)
    MathProgBase.loadproblem!(m,
                              convert(SparseMatrixCSC{Float64,Int},A),
                              convert(Array{Float64,1},collb),
                              convert(Array{Float64,1},colub),
                              convert(Array{Float64,1},obj),
                              convert(Array{Float64,1},rowlb),
                              convert(Array{Float64,1},rowub),
                              sense)
end

function MathProgBase.loadproblem!(m::MosekLinearQuadraticModel,
                                   A::     SparseMatrixCSC{Float64,Int},
                                   collb:: Array{Float64,1},
                                   colub:: Array{Float64,1},
                                   obj::   Array{Float64,1},
                                   rowlb:: Array{Float64,1},
                                   rowub:: Array{Float64,1},
                                   sense:: Symbol)
    Mosek.resizetask(m.task,0,0,0,0,0);

    nrows,ncols = size(A)
    if ncols != length(collb) ||
        ncols != length(colub) ||
        ncols != size(obj,1)   ||
        nrows != length(rowlb) ||
        nrows != length(rowub) ||
        ncols != length(obj)

        throw(MosekMathProgModelError("Inconsistent data dimensions"))
    end

    Mosek.appendvars(m.task,ncols)
    Mosek.appendcons(m.task,nrows)

    (m.bkx,m.blx,m.bux) = makebounds(collb,colub)

    (m.bkc,m.blc,m.buc) = makebounds(rowlb,rowub)

    m.numvar = length(m.bkx)
    m.numcon = length(m.bkc)
    m.lincon = Int32[1:nrows;]
    m.quadcon = Array(Int32,0)
    m.binvarflags = fill(false,m.numvar)

    # input coefficients
    Mosek.putclist(m.task, Int32[1:ncols;], obj)
    Mosek.putacolslice(m.task, 1, ncols+1, A.colptr[1:ncols], A.colptr[2:ncols+1], A.rowval, A.nzval)
    MathProgBase.setsense!(m, sense)

    # input bounds
    Mosek.putvarboundslice(m.task, 1, ncols+1, m.bkx, m.blx, m.bux)
    Mosek.putconboundslice(m.task, 1, nrows+1, m.bkc, m.blc, m.buc)
    m
end

function MathProgBase.loadproblem!(m::MosekLinearQuadraticModel,
                                   filename:: AbstractString)
    tmptask = Mosek.maketask()
    try
        readdata(tmptask, filename)

        if  Mosek.getnumcone(tmptask) > 0 ||
            Mosek.getnumbarvar(tmptask) > 0
            throw(MosekMathProgModelError("Not a linear/quadratic model"))
        end

        numcon = Mosek.getnumcon(tmptask)
        numvar = Mosek.getnumvar(tmptask)

        bkx,blx,bux = Mosek.getvarboundslice(tmptask,1,numvar+1)
        bkc,blc,buc = Mosek.getconboundslice(tmptask,1,numcon+1)

        vts = Mosek.getvartypelist(tmptask,Int32[1:numvar])
        binflags = Bool[ (vts[i] == Mosek.MSK_VARIABLE_TYPE_INT &&
                          bkx[i] == Mosek.MSK_BK_RA &&
                          abs(blx[i]) < 1e-8 &&
                          abs(bux[i]-1.0) < 1e-8)
                        for i in 1:numvar ]

        lincon = find(i -> Mosek.getnumqconknz(m.task,i) == 0, 1:numcon)
        quadcon = find(i -> Mosek.getnumqconknz(m.task,i) > 0, 1:numcon)

        Mosek.deletetask(m.task)

        m.task = tmptask
        m.binvaflags = binflags
        m.bkx = bkx
        m.blx = blx
        m.bux = bux
        m.bkc = bkc
        m.blc = blc
        m.buc = buc
        m.numvar = lenght(bkx)
        m.numcon = length(bkc)
        m.lincon = lincon
        m.quadcon = quadcon
    catch
        Mosek.deletetask(tmptask)
        rethrow()
    end

    m
end

function loadoptions!(m::MosekLinearQuadraticModel)
  loadoptions_internal!(m.task, m.options)
end

function MathProgBase.writeproblem(m::MosekLinearQuadraticModel, filename::AbstractString)
    Mosek.writedata(m.task,filename)
end

function MathProgBase.getvarLB(m::MosekLinearQuadraticModel)
    bk,bl,bu = Mosek.getvarboundslice(m.task,1,m.numvar+1)
    for i in 1:length(bk)
        if bk[i] == Mosek.MSK_BK_FR || bk[i] == Mosek.MSK_BK_UP
            bl[i] = -Inf
        end
    end
    bl
end

function MathProgBase.getvarUB(m::MosekLinearQuadraticModel)
    bk,bl,bu = Mosek.getvarboundslice(m.task,1,m.numvar+1)
    for i in 1:length(bk)
        if bk[i] == Mosek.MSK_BK_FR || bk[i] == Mosek.MSK_BK_LO
            bu[i] = Inf
        end
    end
    bu
end

function MathProgBase.getconstrLB(m::MosekLinearQuadraticModel)
    bk,bl,bu = Mosek.getconboundslice(m.task,1,m.numcon+1)
    for i in 1:length(bk)
        if bk[i] == Mosek.MSK_BK_FR || bk[i] == Mosek.MSK_BK_UP
            bl[i] = -Inf
        end
    end
    bl
end

function MathProgBase.getconstrUB(m::MosekLinearQuadraticModel)
    bk,bl,bu = Mosek.getconboundslice(m.task,1,m.numcon+1)
    for i in 1:length(bk)
        if bk[i] == Mosek.MSK_BK_FR || bk[i] == Mosek.MSK_BK_LO
            bu[i] = Inf
        end
    end
    bu
end

MathProgBase.setvarLB!{T}(m::MosekLinearQuadraticModel, bnd::Array{T,1}) = MathProgBase.setvarLB!(m,convert(Array{Float64,1},bnd))
function MathProgBase.setvarLB!(m::MosekLinearQuadraticModel, bnd::Array{Float64,1})
    n = min(length(bnd),m.numvar)

    for i in 1:n
        if bnd[i] > -Inf
            if m.bux[i] < Inf
                if abs(bnd[i]-m.bux[i]) < 1e-8
                    m.bkx[i] = Mosek.MSK_BK_FX
                    m.blx[i] = m.bux[i]
                    m.bux[i] = m.bux[i]
                else
                    m.bkx[i] = Mosek.MSK_BK_RA
                    m.blx[i] = bnd[i]
                end
            else # bux[i] == Inf
                m.bkx[i] = Mosek.MSK_BK_LO
                m.blx[i] = bnd[i]
            end
        else # bnd[i] == -Inf
            if m.bux[i] < Inf
                m.bkx[i] = Mosek.MSK_BK_UP
            else
                m.bkx[i] = Mosek.MSK_BK_FR
            end
            m.blx[i] = -Inf
        end
    end

    Mosek.putvarboundslice(m.task,1,n+1,m.bkx,m.blx,m.bux)
    if any(m.binvarflags)
        idxs = convert(Array{Int32,1},find(v->v, m.binvarflags))
        bkx = Int32[ Mosek.MSK_BK_RA for i in 1:length(idxs)]
        blx = Float64[ max(m.blx[i],0.0) for i in idxs ]
        bux = Float64[ min(m.bux[i],1.0) for i in idxs ]

        Mosek.putvarboundlist(m.task,idxs, bkx,blx,bux)
    end

    nothing
end

MathProgBase.setvarUB!{T}(m::MosekLinearQuadraticModel, bnd::Array{T,1}) = MathProgBase.setvarUB!(m,convert(Array{Float64,1},bnd))
function MathProgBase.setvarUB!(m::MosekLinearQuadraticModel, bnd::Array{Float64,1})
    n = min(length(bnd),m.numvar)

    for i in 1:n
        if bnd[i] < Inf
            if m.blx[i] > -Inf
                if abs(bnd[i]-m.blx[i]) < 1e-8
                    m.bkx[i] = Mosek.MSK_BK_FX
                    m.blx[i] = m.blx[i]
                    m.bux[i] = m.blx[i]
                else
                    m.bkx[i] = Mosek.MSK_BK_RA
                    m.bux[i] = bnd[i]
                end
            else # blx[i] == -Inf
                m.bkx[i] = Mosek.MSK_BK_UP
                m.bux[i] = bnd[i]
            end
        else # bnd[i] == Inf
            if m.blx[i] > -Inf
                m.bkx[i] = Mosek.MSK_BK_LO
            else
                m.bkx[i] = Mosek.MSK_BK_FR
            end
            m.bux[i] = Inf
        end
    end

    Mosek.putvarboundslice(m.task,1,n+1,m.bkx,m.blx,m.bux)
    if any(m.binvarflags)
        idxs = convert(Array{Int32,1},find(v->v, m.binvarflags))
        bkx = Int32[ Mosek.MSK_BK_RA for i in 1:length(idxs)]
        blx = Float64[ max(m.blx[i],0.0) for i in idxs ]
        bux = Float64[ min(m.bux[i],1.0) for i in idxs ]

        Mosek.putvarboundlist(m.task,idxs, bkx,blx,bux)
    end

    nothing
end

MathProgBase.setconstrLB!{T}(m::MosekLinearQuadraticModel, bnd::Array{T,1}) = MathProgBase.setconstrLB!(m,convert(Array{Float64,1},bnd))
function MathProgBase.setconstrLB!(m::MosekLinearQuadraticModel, bnd::Array{Float64,1})
    n = min(length(bnd),length(m.lincon))

    for i in 1:n
        bk = m.bkc[m.lincon[i]]
        if bnd[i] > -Inf
            if bk == Mosek.MSK_BK_UP || bk == Mosek.MSK_BK_RA || bk == Mosek.MSK_BK_FX
                if abs(bnd[i]-m.buc[m.lincon[i]]) < 1e-8
                    m.bkc[m.lincon[i]] = Mosek.MSK_BK_FX
                    m.blc[m.lincon[i]] = m.buc[m.lincon[i]]
                    m.buc[m.lincon[i]] = m.buc[m.lincon[i]]
                else
                    m.bkc[m.lincon[i]] = Mosek.MSK_BK_RA
                    m.blc[m.lincon[i]] = bnd[i]
                end
            else # buc[m.lincon[i]] == Inf
                m.bkc[m.lincon[i]] = Mosek.MSK_BK_LO
                m.blc[m.lincon[i]] = bnd[i]
            end
        else # bnd[m.lincon[i]] == -Inf
            if bk == Mosek.MSK_BK_UP || bk == Mosek.MSK_BK_RA || bk == Mosek.MSK_BK_FX
                m.bkc[m.lincon[i]] = Mosek.MSK_BK_UP
            else
                m.bkc[m.lincon[i]] = Mosek.MSK_BK_FR
            end
            m.blc[m.lincon[i]] = -Inf
        end
    end

    Mosek.putconboundslice(m.task,Int32(1),Int32(m.numcon+1),m.bkc,m.blc,m.buc)

    nothing
end

MathProgBase.setconstrUB!{T}(m::MosekLinearQuadraticModel, bnd::Array{T,1}) = MathProgBase.setconstrUB!(m,convert(Array{Float64,1},bnd))
function MathProgBase.setconstrUB!(m::MosekLinearQuadraticModel, bnd::Array{Float64,1})
    n = min(length(bnd),m.numcon)

    for i in 1:n
        bk = m.bkc[m.lincon[i]]
        if bnd[i] < Inf
            if bk == Mosek.MSK_BK_LO || bk == Mosek.MSK_BK_RA || bk == Mosek.MSK_BK_FX
                if abs(bnd[i]-m.blc[m.lincon[i]]) < 1e-8
                    m.bkc[m.lincon[i]] = Mosek.MSK_BK_FX
                    m.buc[m.lincon[i]] = m.blc[m.lincon[i]]
                else
                    m.bkc[m.lincon[i]] = Mosek.MSK_BK_RA
                    m.buc[m.lincon[i]] = bnd[i]
                end
            else # blc[i] == -Inf
                m.bkc[m.lincon[i]] = Mosek.MSK_BK_UP
                m.buc[m.lincon[i]] = bnd[i]
            end
        else # bnd[i] == Inf
            if bk == Mosek.MSK_BK_LO || bk == Mosek.MSK_BK_RA || bk == Mosek.MSK_BK_FX
                m.bkc[m.lincon[i]] = Mosek.MSK_BK_LO
            else
                m.bkc[m.lincon[i]] = Mosek.MSK_BK_FR
            end
            m.buc[m.lincon[i]] = Inf
        end
    end
    Mosek.putconboundslice(m.task,Int32(1),Int32(m.numcon+1),m.bkc,m.blc,m.buc)

    nothing
end

function MathProgBase.getobj(m::MosekLinearQuadraticModel)
    Mosek.getcslice(m.task,1,m.numvar+1)
end

MathProgBase.setobj!{T}(m::MosekLinearQuadraticModel, c::Array{T,1}) = MathProgBase.setobj!(m,convert(Array{Float64,1},c))
function MathProgBase.setobj!(m::MosekLinearQuadraticModel, c :: Array{Float64,1})
    n = min(length(c),m.numvar)
    Mosek.putclist(m.task,Int32[1:n;],c[1:n])
end

function MathProgBase.getconstrmatrix(m::MosekLinearQuadraticModel)
    numnz = sum(Int[ Mosek.getarownumnz(m.task,i) for i in 1:m.numcon ])
    asubi = Array(Int32,numnz)
    asubj = Array(Int32,numnz)
    aval  = Array(Float64,numnz)

    let ptr = 1
        for i in 1:m.numcon
            subj,valj = Mosek.getarow(m.task,i)
            n = length(subj)
            asubi[ptr:ptr+n-1] = i
            asubj[ptr:ptr+n-1] = subj
            aval[ptr:ptr+n-1]  = valj
            ptr += n
        end
    end

    sparse(asubi,asubj,aval,length(m.lincon),m.numvar)
end

MathProgBase.addvar!(m::MosekLinearQuadraticModel, bl, bu, c) = MathProgBase.addvar!(m,convert(Float64,bl),convert(Float64,bu),convert(Float64,c))

function MathProgBase.addvar!(m::MosekLinearQuadraticModel,
                 bl  ::Float64,
                 bu  ::Float64,
                 c   ::Float64)
    m.numvar += 1
    if bl > -Inf
        if bu < Inf
            if abs(bl-bu) < 1e-8
                push!(m.bkx,Mosek.MSK_BK_FX)
                push!(m.blx,bl)
                push!(m.bux,bl)
            else
                push!(m.bkx,Mosek.MSK_BK_RA)
                push!(m.blx,bl)
                push!(m.bux,bu)
            end
        else
            push!(m.bkx,Mosek.MSK_BK_LO)
            push!(m.blx,bl)
            push!(m.bux,Inf)
        end
    else
        if bu < Inf
            push!(m.bkx,Mosek.MSK_BK_UP)
            push!(m.blx,-Inf)
            push!(m.bux,bu)
        else
            push!(m.bkx,Mosek.MSK_BK_FR)
            push!(m.blx,-Inf)
            push!(m.bux,Inf)
        end
    end

    Mosek.appendvars(m.task,1);
    Mosek.putvarbound(m.task,m.numvar,m.bkx[m.numvar],m.blx[m.numvar],m.bux[m.numvar])
    Mosek.putcj(m.task,m.numvar,c)
    push!(m.binvarflags,false)
end

MathProgBase.addvar!{T1,T2,T3,T4,T5}(m::MosekLinearQuadraticModel,
                                     subi::Array{T1,1},
                                     val ::Array{T2,1},
                                     bl  ::T3,
                                     bu  ::T4,
                                     c   ::T5) = MathProgBase.addvar!(m,convert(Array{Int32,1},subi),convert(Array{Float64,1},val),convert(Float64,bl),convert(Float64,bu),convert(Float64,c))
function MathProgBase.addvar!(m::MosekLinearQuadraticModel,
                              subi::Array{Int32,1},
                              val ::Array{Float64,1},
                              bl  ::Float64,
                              bu  ::Float64,
                              c   ::Float64)
    MathProgBase.addvar!(m,bl,bu,c)
    Mosek.putacol(m.task,m.numvar,subi,val)
end

MathProgBase.addconstr!{T1,T2,T3,T4}(m::MosekLinearQuadraticModel,
                                     subj::Array{T1,1},
                                     val ::Array{T2,1},
                                     bl  ::T3,
                                     bu  ::T4) = MathProgBase.addconstr!(m,convert(Array{Int32,1},subj),convert(Array{Float64,1},val),convert(Float64,bl),convert(Float64,bu))

function MathProgBase.addconstr!(m::MosekLinearQuadraticModel,
                                 subj::Array{Int32,1},
                                 val ::Array{Float64,1},
                                 bl  ::Float64,
                                 bu  ::Float64)
    m.numcon += 1
    push!(m.lincon,m.numcon)
    if bl > -Inf
        if bu < Inf
            if abs(bl-bu) < 1e-8
                push!(m.bkc,Mosek.MSK_BK_FX)
                push!(m.blc,bl)
                push!(m.buc,bl)
            else
                push!(m.bkc,Mosek.MSK_BK_RA)
                push!(m.blc,bl)
                push!(m.buc,bu)
            end
        else
            push!(m.bkc,Mosek.MSK_BK_LO)
            push!(m.blc,bl)
            push!(m.buc,Inf)
        end
    else
        if bu < Inf
            push!(m.bkc,Mosek.MSK_BK_UP)
            push!(m.blc,-Inf)
            push!(m.buc,bu)
        else
            push!(m.bkc,Mosek.MSK_BK_FR)
            push!(m.blc,-Inf)
            push!(m.buc,Inf)
        end
    end
    Mosek.appendcons(m.task,1);
    Mosek.putconbound(m.task,m.numcon,m.bkc[m.numcon],m.blc[m.numcon],m.buc[m.numcon])
    Mosek.putarow(m.task,m.numcon,subj,val)
end

MathProgBase.numlinconstr(m::MosekLinearQuadraticModel) = length(m.lincon)

MathProgBase.getobjval(m::MosekLinearQuadraticModel) = getobjval(m.task)

function MathProgBase.getsolution(m::MosekLinearQuadraticModel)
    sol = getsoldef(m.task)
    if sol < 0
        throw(MosekMathProgModelError("No solution available"))
    end

    solsta = Mosek.getsolsta(m.task,sol)

    if solsta in [Mosek.MSK_SOL_STA_OPTIMAL,
                  Mosek.MSK_SOL_STA_PRIM_FEAS,
                  Mosek.MSK_SOL_STA_PRIM_AND_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_NEAR_OPTIMAL,
                  Mosek.MSK_SOL_STA_NEAR_PRIM_FEAS,
                  Mosek.MSK_SOL_STA_NEAR_PRIM_AND_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_INTEGER_OPTIMAL,
                  Mosek.MSK_SOL_STA_NEAR_INTEGER_OPTIMAL ]
        Mosek.getxx(m.task,sol)
  else
    throw(MosekMathProgModelError("No solution available"))
  end
end

function MathProgBase.getconstrsolution(m::MosekLinearQuadraticModel)
    sol = getsoldef(m.task)
    if sol < 0
        throw(Mosek.MosekMathProgModelError("No solution available"))
    end

    Mosek.getxc(m.task,sol)[m.lincon]
end

function MathProgBase.getreducedcosts(m::MosekLinearQuadraticModel)
    sol = getsoldef(m.task)
    if sol < 0 || sol == Mosek.MSK_SOL_ITG
        throw(Mosek.MosekMathProgModelError("Solution not available"))
    end
    solsta = Mosek.getsolsta(m.task,sol)

    if solsta in [Mosek.MSK_SOL_STA_OPTIMAL,
                  Mosek.MSK_SOL_STA_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_PRIM_AND_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_NEAR_OPTIMAL,
                  Mosek.MSK_SOL_STA_NEAR_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_NEAR_PRIM_AND_DUAL_FEAS ]
        Mosek.getslx(m.task,sol) - Mosek.getsux(m.task,sol)
    else
        throw(MosekMathProgModelError("No solution available"))
    end
end

function MathProgBase.getconstrduals(m::MosekLinearQuadraticModel)
    sol = getsoldef(m.task)
    if sol < 0 || sol == Mosek.MSK_SOL_ITG
        throw(Mosek.MosekMathProgModelError("Solution not available"))
    end
    solsta = Mosek.getsolsta(m.task,sol)

    if solsta in [Mosek.MSK_SOL_STA_OPTIMAL,
                  Mosek.MSK_SOL_STA_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_PRIM_AND_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_NEAR_OPTIMAL,
                  Mosek.MSK_SOL_STA_NEAR_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_NEAR_PRIM_AND_DUAL_FEAS ]
        Mosek.gety(m.task,sol)[m.lincon]
    else
        throw(MosekMathProgModelError("Solution not available"))
    end
end

function MathProgBase.getbasis(m::MosekLinearQuadraticModel)
    if ! Mosek.solutiondef(m.task,Mosek.MSK_SOL_BAS)
        throw(Mosek.MosekMathProgModelError("Basis not available"))
    end
    sol = Mosek.MSK_SOL_BAS

    skx = Mosek.getskx(m.task,sol)
    skc = Mosek.getskc(m.task,sol)

    cbasis = [if     skx[i] == Mosek.MSK_SK_BAS :Basic
              elseif skx[i] == Mosek.MSK_SK_LO  :NonBasicAtLower
              elseif skx[i] == Mosek.MSK_SK_UP  :NonBasicAtUpper
              elseif skx[i] == Mosek.MSK_SK_FX  :NonBasicAtLower # or upper. Doesn't matter.
              else                        :SuperBasic
              end
              for i in 1:m.numvar ]
    rbasis = [if     skc[i] == Mosek.MSK_SK_BAS :Basic
              elseif skc[i] == Mosek.MSK_SK_LO  :NonBasicAtLower
              elseif skc[i] == Mosek.MSK_SK_UP  :NonBasicAtUpper
              elseif skc[i] == Mosek.MSK_SK_FX  :NonBasicAtLower # or upper. Doesn't matter.
              else                        :SuperBasic
              end
              for i in m.lincon]

    cbasis,rbasis
end

function MathProgBase.getunboundedray(m::MosekLinearQuadraticModel)
    soldef = getsoldef(m.task)
    if soldef < 0 throw(MosekMathProgModelError("No solution available")) end

    solsta = Mosek.getsolsta(m.task,soldef)
    if solsta in [ Mosek.MSK_SOL_STA_DUAL_INFEAS_CER, Mosek.MSK_SOL_STA_NEAR_DUAL_INFEAS_CER ]
        Mosek.getxx(m.task,soldef)
    else
        throw(MosekMathProgModelError("No ray available"))
    end
end

function MathProgBase.getinfeasibilityray(m::MosekLinearQuadraticModel)
    sol = getsoldef(m.task)
    if sol < 0 throw(MosekMathProgModelError("No solution available")) end

    solsta = Mosek.getsolsta(m.task,sol)
    if solsta in [ Mosek.MSK_SOL_STA_PRIM_INFEAS_CER, Mosek.MSK_SOL_STA_NEAR_PRIM_INFEAS_CER ]
        Mosek.getsux(m.task,sol) - Mosek.getslx(m.task,sol)
    else
        throw(MosekMathProgModelError("No ray available"))
    end
end

MathProgBase.getrawsolver(m::MosekLinearQuadraticModel) = m.task

MathProgBase.getsimplexiter(m::MosekLinearQuadraticModel) = Mosek.getintinf(m.task,Mosek.MSK_IINFITEM_SIM_PRIMAL_ITER)+Mosek.getintinf(m.task,Mosek.MSK_IINFITEM_SIM_DUAL_ITER)+Mosek.getintinf(m.task,Mosek.MSK_IINFITEM_SIM_PRIMAL_DUAL_ITER)

MathProgBase.getbarrieriter(m::MosekLinearQuadraticModel) = Mosek.getintinf(m.task,Mosek.MSK_IINFITEM_SIM_INTPNT_ITER)

MathProgBase.setwarmstart!{T}(m::MosekLinearQuadraticModel, v::Array{T,1}) = MathProgBase.setwarmstart!(m,convert(Array{Float64,1},v))
function MathProgBase.setwarmstart!(m::MosekLinearQuadraticModel, v::Array{Float64,1})
    n = min(m.numvar,length(v))
    vals = Array(Float64, n)
    vals[:] = v[1:n]

    nanidxs = find(isnan,vals)
    vals[nanidxs] = 0.0

    skx = Int32[ if isnan(vals[i]) Mosek.MSK_SK_UNK else Mosek.MSK_SK_BAS end for i in 1:n ]

    Mosek.putxxslice(m.task,Mosek.MSK_SOL_BAS,1,n+1,vals);
    Mosek.putskxslice(m.task,Mosek.MSK_SOL_BAS,1,n+1,skx);
end

function MathProgBase.optimize!(m::MosekLinearQuadraticModel)
    m.lasttrm = Mosek.optimize(m.task)
end

MathProgBase.status(m::MosekLinearQuadraticModel) = status(m.task,m.lasttrm)

MathProgBase.getobjbound(m::MosekLinearQuadraticModel) = Mosek.getdouinf(m.task,Mosek.MSK_DINF_MIO_OBJ_BOUND)

MathProgBase.getobjgap(m::MosekLinearQuadraticModel) = getobjgap(m.task)

MathProgBase.getsolvetime(m::MosekLinearQuadraticModel) = Mosek.getdouinf(m.task,Mosek.MSK_DINF_OPTIMIZER_TIME)

MathProgBase.getsense(m::MosekLinearQuadraticModel) = getsense(m.task)

MathProgBase.setsense!(m::MosekLinearQuadraticModel,sense) = setsense!(m.task,sense)

function MathProgBase.freemodel!(m::MosekLinearQuadraticModel)
    Mosek.deletetask(m.task)
    nothing
end

MathProgBase.numvar(m::MosekLinearQuadraticModel) = m.numvar
MathProgBase.numconstr(m::MosekLinearQuadraticModel) = length(m.lincon)

function MathProgBase.setvartype!(m::MosekLinearQuadraticModel,vtvec::Vector{Symbol})
    n = min(m.numvar,length(vtvec))
    if n > 0
        vts = Int32[if     vt == :Cont Mosek.MSK_VAR_TYPE_CONT
                    elseif vt == :Int  Mosek.MSK_VAR_TYPE_INT
                    elseif vt == :Bin  Mosek.MSK_VAR_TYPE_INT
                    else               Mosek.MSK_VAR_TYPE_CONT
                    end
                    for vt in vtvec[1:n]]
        Mosek.putvartypelist(m.task,Int32[1:n;],vts)
        for i in find(vt -> vt == :Bin, vtvec[1:n])
            bl = max(m.blx[i],0.0)
            bu = min(m.bux[i],1.0)
            Mosek.putvarbound(m.task,i,Mosek.MSK_BK_RA,bl,bu)
        end

        # for all :Bin vars being changed to :Int or :Cont, restore original bounds
        for i in find(i -> (vtvec[i] == :Cont || vtvec[i] == :Int) && m.binvarflags[i], 1:n)
            Mosek.putvarbound(m.task,i,m.bkx[i],m.blx[i],m.bux[i])
        end

        for i in 1:n
            m.binvarflags[i] = vtvec[i] == :Bin
        end
    end
end

function MathProgBase.getvartype(m::MosekLinearQuadraticModel)
    mskvt = Mosek.getvartypelist(m.task,Int32[1:m.numvar;])
    [if mskvt[i] == Mosek.MSK_VAR_TYPE_INT
         if m.binvarflags[i]
             :Bin
         else
             :Int
         end
     else
         :Cont
     end
     for i in 1:m.numvar]
end

##############################################################
## Integer Programming
##############################################################

MathProgBase.getnodecount(m::MosekLinearQuadraticModel) = 0

##############################################################
## Quadratic
##############################################################

MathProgBase.numquadconstr(m::MosekLinearQuadraticModel) = length(m.quadcon)

MathProgBase.setquadobj!(m::MosekLinearQuadraticModel,subi,subj,valij) = MathProgBase.setquadobj!(m,convert(Array{Int32,1},subi),convert(Array{Int32,1},subj),convert(Array{Float64,1},valij))

# NOTE on data format: The matrix is specified by inputting only lower
# or upper triangular part. This means that whenever element (i,j) is
# added, (j,i) is implicitly added. Duplicates are added together
function MathProgBase.setquadobj!(m::MosekLinearQuadraticModel,
                                  subi  :: Array{Int32,1},
                                  subj  :: Array{Int32,1},
                                  valij :: Array{Float64,1})
    n = length(subi)
    let qsubi = subi[:],
        qsubj = subj[:]
        for i in 1:n
            if qsubi[i] < qsubj[i]
                tmp = qsubi[i]
                qsubi[i] = qsubj[i]
                qsubj[i] = tmp
            end
        end
        Mosek.putqobj(m.task,qsubi,qsubj,valij)
    end
end

function MathProgBase.addquadconstr!(m :: MosekLinearQuadraticModel,
                                     subj,
                                     valj,
                                     qsubi,
                                     qsubj,
                                     qvalij,
                                     sense  :: Char,
                                     bnd)
    MathProgBase.addquadconstr!(m,
                                convert(Array{Int32,1},subj),
                                convert(Array{Float64,1},valj),
                                convert(Array{Int32,1},qsubi),
                                convert(Array{Int32,1},qsubj),
                                convert(Array{Float64,1},qvalij),
                                sense,
                                convert(Float64,bnd))
end

function MathProgBase.addquadconstr!(m      :: MosekLinearQuadraticModel,
                                     subj   :: Array{Int32,1},
                                     valj   :: Array{Float64,1},
                                     qsubi  :: Array{Int32,1},
                                     qsubj  :: Array{Int32,1},
                                     qvalij :: Array{Float64,1},
                                     sense  :: Char,
                                     bnd    :: Float64)

    if     sense == '<'
        push!(m.bkc,Mosek.MSK_BK_UP)
    elseif sense == '>'
        push!(m.bkc,Mosek.MSK_BK_LO)
    else
        throw(MosekMathProgSolverInterface.MosekMathProgModelError("Invalid sense"))
    end

    m.numcon += 1
    push!(m.quadcon,m.numcon)
    push!(m.blc,bnd)
    push!(m.buc,bnd)

    Mosek.appendcons(m.task,1)
    Mosek.putconbound(m.task,m.numcon,m.bkc[m.numcon],m.blc[m.numcon],m.buc[m.numcon])

    let qsubi = qsubi[:],
        qsubj = qsubj[:],
        qval  = qvalij[:]
        for i in 1:length(qsubi)
            if qsubi[i] < qsubj[i]
                t = qsubi[i]
                qsubi[i] = qsubj[i]
                qsubj[i] = t
            elseif  qsubi[i] == qsubj[i]
                qval[i] *= 2
            end
        end
        Mosek.putqconk(m.task,m.numcon,qsubi,qsubj,qval)
    end

    Mosek.putarow(m.task,m.numcon,subj,valj)
end

function MathProgBase.getquadconstrsolution(m::MosekLinearQuadraticModel)
    sol = getsoldef(m.task)
    if sol < 0
        throw(Mosek.MosekMathProgModelError("No solution available"))
    end

    Mosek.getxc(sol)[m.quadcon]
end

function MathProgBase.getquadconstrduals(m::MosekLinearQuadraticModel)
    sol = getsoldef(m.task)
    if sol < 0 || sol == Mosek.MSK_SOL_ITG
        throw(Mosek.MosekMathProgModelError("Solution not available"))
    end
    solsta = Mosek.getsolsta(m.task,sol)

    if solsta in [Mosek.MSK_SOL_STA_OPTIMAL,
                  Mosek.MSK_SOL_STA_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_PRIM_AND_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_NEAR_OPTIMAL,
                  Mosek.MSK_SOL_STA_NEAR_DUAL_FEAS,
                  Mosek.MSK_SOL_STA_NEAR_PRIM_AND_DUAL_FEAS ]
        Mosek.gety(m.task)[m.quadcon]
    else
        throw(MosekMathProgModelError("Solution not available"))
    end
end

function MathProgBase.getquadinfeasibilityray(m::MosekLinearQuadraticModel)
    sol = getsoldef(m)
    if sol < 0 throw(MosekMathProgModelError("No solution available")) end
    solsta = getsolsta(m.task,sol)

    s = Mosek.getsux(m.task,sol) - Mosek.getslx(m.task,sol)
    if solsta in [ Mosek.MSK_SOL_STA_PRIM_INFEAS_CER, Mosek.MSK_SOL_STA_NEAR_PRIM_INFEAS_CER ]
        -s[m.quadcon]
    else
        throw(MosekMathProgModelError("No solution available"))
    end
end

function MathProgBase.getquadconstrRHS(m::MosekLinearQuadraticModel)
    m.blc[m.quadcon]
end

MathProgBase.setquadconstrRHS!(m::MosekLinearQuadraticModel, bnd) = MathProgBase.setquadconstrRHS!(m,convert(Array{Float64,1},bnd))
function MathProgBase.setquadconstrRHS!(m::MosekLinearQuadraticModel, bnd::Array{Float64,1})
    n = min(length(bnd),length(m.quadcon))
    m.blc[m.quadcon[1:n]] = bnd[1:n]

    Mosek.putconboundlist(m.task,m.quadcon[1:n],m.bkc,m.blc,m.buc)
end

##############################################################
## Nonlinear
#############################################################

MathProgBase.NonlinearModel(s::MosekSolver) = MosekNonlinearModel(MathProgBase.LinearQuadraticModel(s))

type CallbackData
    d::MathProgBase.AbstractNLPEvaluator
    numVar::Int
    numConstr::Int
    Ihess::Vector{Int32}
    Jhess::Vector{Int32}
    jac_colval::Vector{Int32} # Compressed sparse row form of Jacobian sparsity
    jac_rowstarts::Vector{Int32}
    jac_nzval::Vector{Float64} # storage for full jacobian (mosek may ask for a subset)
    jac_nnz_original::Int # nnz in NLP evaluator jacobian == length(Ijac)
    jac_idxmap::Vector{Int} # map from indices in NLP evaluator jacobian to mosek jacobian
    J_tmp::Vector{Float64} # storage for NLP evaluator jacobian
    g_tmp::Vector{Float64} # storage for constraint values
end

function msk_nl_getsp_wrapper_mpb(nlhandle::    Ptr{Void},
                                  numgrdobjnz:: Ptr{Int32}, # number of nonzeros in gradient of objective
                                  grdobjsub::   Ptr{Int32}, # subscripts of nonzeros in gradient of objective
                                  i_::          Int32,      # constraint index
                                  convali::     Ptr{Bool},  # 0/1 indicating whether constraint i is non-linear
                                  grdconinz::   Ptr{Int32}, # number of nonzeros in gradient of constraint i
                                  grdconisub::  Ptr{Int32}, # subscripts of nonzeros in gradient of constraint i
                                  yo::          Int32,      # 0/1 include objective in computation of hessian pattern
                                  numycnz::     Int32,      # number of constraints to include in computation of the hessian pattern
                                  ycsub::       Ptr{Int32}, # indexes of constraints to include in computation of the hessian pattern
                                  maxnumhesnz:: Int32,      # lengths of hessubi and hessubj
                                  numhesnz_::   Ptr{Int32}, # number of hessian nonzeros
                                  hessubi::     Ptr{Int32}, # column subscrips of hessian non-zeros
                                  hessubj::     Ptr{Int32}) # row subscripts of hessian non-zeros
    cb = unsafe_pointer_to_objref(nlhandle)::CallbackData

    i = i_+1

    if numgrdobjnz != C_NULL
        unsafe_store!(numgrdobjnz, convert(Int32,cb.numVar))
    end

    if grdobjsub != C_NULL
        grdobjsub_a = unsafe_wrap(Array{Int32,1},grdobjsub,(cb.numVar,))
        for i in 1:cb.numVar
            grdobjsub_a[i] = i-1
        end
    end

    if i <= cb.numConstr
        con_nnz = cb.jac_rowstarts[i+1] - cb.jac_rowstarts[i]
        if convali != C_NULL
            # can we only say that constraint is linear if we've added the linear part separately?
            if con_nnz > 0
                unsafe_store!(convali, convert(Int32,1))
            else
                unsafe_store!(convali, convert(Int32,0))
            end
        end

        if grdconinz != C_NULL
            unsafe_store!(grdconinz, convert(Int32, con_nnz))
        end

        if grdconisub != C_NULL
            if con_nnz > 0
                grdconisub_a = unsafe_wrap(Array{Int32,1},grdconisub,(con_nnz,))
                grdconisub_a[1:con_nnz] = cb.jac_colval[cb.jac_rowstarts[i]:(cb.jac_rowstarts[i+1]-1)] - 1
            end
        end
    end

    hess_nnz = length(cb.Ihess)

    if numhesnz_ != C_NULL
        unsafe_store!(numhesnz_, convert(Int32, hess_nnz))
    end

    if hessubi != C_NULL && hessubj != C_NULL && maxnumhesnz >= hess_nnz
        hessubi_a = unsafe_wrap(Array{Int32,1},hessubi,(hess_nnz,))
        hessubj_a = unsafe_wrap(Array{Int32,1},hessubj,(hess_nnz,))

        for i in 1:hess_nnz
            hessubi_a[i] = cb.Ihess[i] - 1
            hessubj_a[i] = cb.Jhess[i] - 1
        end
    end

    return Int32(0)::Int32
end

function msk_nl_getva_wrapper_mpb(nlhandle    :: Ptr{Void},
                                  xx_         :: Ptr{Float64}, # input
                                  yo          :: Float64,
                                  yc_         :: Ptr{Float64}, # input, length = numcon
                                  objval      :: Ptr{Float64},
                                  numgrdobjnz :: Ptr{Int32},
                                  grdobjsub   :: Ptr{Int32},
                                  grdobjval   :: Ptr{Float64},
                                  numi_       :: Int32,
                                  subi_       :: Ptr{Int32},   # input
                                  conval      :: Ptr{Float64},
                                  grdconptrb_ :: Ptr{Int32},   # input
                                  grdconptre_ :: Ptr{Int32},   # input
                                  grdconsub_  :: Ptr{Int32},   # input
                                  grdconval_  :: Ptr{Float64},
                                  grdlag      :: Ptr{Float64},
                                  maxnumhesnz :: Int32,
                                  numhesnz    :: Ptr{Int32},
                                  hessubi     :: Ptr{Int32},
                                  hessubj     :: Ptr{Int32},
                                  hesval      :: Ptr{Float64})
    cb = unsafe_pointer_to_objref(nlhandle)::CallbackData

    numi = convert(Int,numi_)
    xx = unsafe_wrap(Array{Float64,1},xx_,(cb.numVar,))
    yc = unsafe_wrap(Array{Float64,1},yc_,(cb.numConstr,))
    subi = unsafe_wrap(Array{Int32,1},subi_,(numi,))

    if objval != C_NULL
        unsafe_store!(objval, MathProgBase.eval_f(cb.d, xx))
    end

    if numgrdobjnz != C_NULL
        unsafe_store!(numgrdobjnz, convert(Int32,cb.numVar))
    end

    if grdobjsub != C_NULL && grdobjval != C_NULL
        grdobjval_a = unsafe_wrap(Array{Float64,1},grdobjval,(cb.numVar,))
        grdobjsub_a = unsafe_wrap(Array{Int32,1},grdobjsub,(cb.numVar,))

        MathProgBase.eval_grad_f(cb.d, grdobjval_a, xx)

        for i in 1:cb.numVar
            grdobjsub_a[i] = i-1
        end
    end

    if numi > 0 && conval != C_NULL
        conv   = unsafe_wrap(Array{Float64,1},conval,(numi,))
        MathProgBase.eval_g(cb.d, cb.g_tmp, xx)
        for i=1:numi
            conv[i] = cb.g_tmp[subi[i]+1]
        end
    end

    # do we need to compute Jacobian?
    if grdconval_ != C_NULL || grdlag != C_NULL
        MathProgBase.eval_jac_g(cb.d, cb.J_tmp, xx)
        fill!(cb.jac_nzval, 0.0)
        for i in 1:cb.jac_nnz_original
            cb.jac_nzval[cb.jac_idxmap[i]] += cb.J_tmp[i]
        end
    end

    if grdconval_ != C_NULL
        grdconptrb = unsafe_wrap(Array{Int32,1},grdconptrb_,(numi,))
        grdconptre = unsafe_wrap(Array{Int32,1},grdconptre_,(numi,))
        grdconsub = unsafe_wrap(Array{Int32,1},grdconsub_,(grdconptre[numi],))
        grdconval = unsafe_wrap(Array{Float64,1},grdconval_,(grdconptre[numi],))
        nz = 1
        for i=1:numi
            ptrb = grdconptrb[i]
            n    = grdconptre[i] - grdconptrb[i]
            con = subi[i]+1
            # make sure input matches the sparsity pattern we provided
            @assert n == (cb.jac_rowstarts[con+1]-cb.jac_rowstarts[con])
            for k in 1:n
                @assert (grdconsub[grdconptrb[i]+k]+1) ==  cb.jac_colval[cb.jac_rowstarts[con]+k-1]
            end
        end

        # extract subset of Jacobian that mosek asked for
        for i in 1:numi
            con = subi[i]+1
            grdconval[(grdconptrb[i]+1):grdconptre[i]] = view(cb.jac_nzval, cb.jac_rowstarts[con]:(cb.jac_rowstarts[con+1]-1))
        end
    end

    if grdlag != C_NULL
        # could use eval_jac_prod_t here, but just do a sparse matvec instead
        grdlag_a = unsafe_wrap(Array{Float64,1},grdlag,(cb.numVar,))
        MathProgBase.eval_grad_f(cb.d, grdlag_a, xx)
        scale!(grdlag_a, yo)
        Jmat = SparseMatrixCSC(cb.numVar,cb.numConstr,cb.jac_rowstarts,cb.jac_colval,cb.jac_nzval)
        A_mul_B!(-1.0, Jmat, yc, 1.0, grdlag_a)
    end

    nhesnz = length(cb.Ihess)

    if numhesnz != C_NULL
        unsafe_store!(numhesnz,convert(Int32,nhesnz))
    end

    if maxnumhesnz > 0 && hessubi != C_NULL && hessubj != C_NULL && hesval != C_NULL
        hessubi_a = unsafe_wrap(Array{Int32,1},hessubi,(nhesnz,))
        hessubj_a = unsafe_wrap(Array{Int32,1},hessubj,(nhesnz,))

        scale!(yc,-1)
        MathProgBase.eval_hesslag(cb.d,unsafe_wrap(Array{Float64,1},hesval,(nhesnz,)),xx,yo,yc)
        scale!(yc,-1)

        for i=1:length(hessubi_a)
            hessubi_a[i] = cb.Ihess[i]-1
            hessubj_a[i] = cb.Jhess[i]-1
        end
    end
    return convert(Int32,0)::Int32
end

function MathProgBase.loadproblem!(m::MosekNonlinearModel,
                                   numVar,
                                   numConstr,
                                   collb,
                                   colub,
                                   rowlb,
                                   rowub,
                                   sense::Symbol,
                                   d::MathProgBase.AbstractNLPEvaluator)
    MathProgBase.loadproblem!(m,
                              convert(Int32,numVar),
                              convert(Int32,numConstr),
                              convert(Array{Float64,1},collb),
                              convert(Array{Float64,1},colub),
                              convert(Array{Float64,1},rowlb),
                              convert(Array{Float64,1},rowub),
                              sense,d)
end

function MathProgBase.loadproblem!(m::MosekNonlinearModel,
                                   numVar::Int32,
                                   numConstr::Int32,
                                   collb::Array{Float64,1},
                                   colub::Array{Float64,1},
                                   rowlb::Array{Float64,1},
                                   rowub::Array{Float64,1},
                                   sense::Symbol,
                                   d::MathProgBase.AbstractNLPEvaluator)
    m = m.m

    if !(numVar == length(collb) == length(colub)) ||
       !(numConstr == length(rowlb) == length(rowub))
        throw(MosekMathProgModelError("Inconsistent data dimensions"))
    end

    Mosek.deletetask(m.task)
    m.task = Mosek.maketask(Mosek.msk_global_env)
    loadoptions!(m)

    Mosek.appendvars(m.task, numVar)
    Mosek.appendcons(m.task, numConstr)

    m.numvar = numVar
    m.numcon = numConstr

    (m.bkx,m.blx,m.bux) = makebounds(collb,colub)
    (m.bkc,m.blc,m.buc) = makebounds(rowlb,rowub)

    Mosek.putvarboundslice(m.task, 1, numVar+1, m.bkx, m.blx, m.bux)
    Mosek.putconboundslice(m.task, 1, numConstr+1, m.bkc, m.blc, m.buc)

    m.lincon = Int32[1:m.numcon;]
    m.quadcon = Int32[]

    # input bounds

    MathProgBase.setsense!(m, sense)

    MathProgBase.initialize(d, [:Grad, :Jac, :Hess])
    Ijac, Jjac = MathProgBase.jac_structure(d)
    Ihess, Jhess = MathProgBase.hesslag_structure(d)

    # Mosek needs jacobian sparsity in compressed sparse row format, which means that we need to
    # reformat the indices while maintaining a map
    jac_nnz_original = length(Ijac)
    mergedindices = zeros(Int, jac_nnz_original)
    mergednnz = [0]
    mergedmap = zeros(Int, jac_nnz_original)

    # this is called for duplicate elements
    # assumption: at least one of the indices hasn't been combined before
    function combine(idx1,idx2)
        if mergedmap[idx1] == 0 && mergedmap[idx2] != 0
            mergednnz[1] += 1
            mergedmap[idx1] = idx2
            mergedindices[mergednnz[1]] = idx1
            return idx2
        else
            @assert mergedmap[idx2] == 0
            mergednnz[1] += 1
            mergedmap[idx2] = idx1
            mergedindices[mergednnz[1]] = idx2
            return idx1
        end
    end

    J = sparse(Jjac, Ijac, [i for i in 1:length(Ijac)], numVar, numConstr, combine)

    # map from original index into output index
    idxmap = zeros(Int, jac_nnz_original)
    for row in 1:numConstr
        for pos in J.colptr[row]:(J.colptr[row+1]-1)
            col = J.rowval[pos]
            origidx = J.nzval[pos] # this is the original index of this element
            idxmap[origidx] = pos
        end
    end
    for k in 1:mergednnz[1]
        origidx = mergedindices[k]
        mergedwith = mergedmap[origidx]
        @assert idxmap[origidx] == 0
        @assert idxmap[mergedwith] != 0
        idxmap[origidx] = idxmap[mergedwith]
    end
    for i in 1:jac_nnz_original
        @assert idxmap[i] != 0
    end

    cb = CallbackData(d, numVar, numConstr, Ihess, Jhess, J.rowval, J.colptr, zeros(J.colptr[numConstr+1]-1),
          jac_nnz_original, idxmap, zeros(jac_nnz_original), zeros(numConstr))

    nlgetsp = cfunction(msk_nl_getsp_wrapper_mpb,
                      Int32,
                      (Ptr{Void},Ptr{Int32},Ptr{Int32},Int32,Ptr{Bool},Ptr{Int32},Ptr{Int32},Int32,Int32,Ptr{Int32},Int32,Ptr{Int32},Ptr{Int32},Ptr{Int32}))
    nlgetva = cfunction(msk_nl_getva_wrapper_mpb,
                      Int32,
                      ( Ptr{Void}, # nlhandle
                        Ptr{Float64},Float64,Ptr{Float64}, # xx,yo,yc
                        Ptr{Float64},Ptr{Int32},Ptr{Int32},Ptr{Float64}, # objval,numgrdobjnz,grdobjsub,grdobjval
                        Int32,Ptr{Int32},Ptr{Float64}, # numi,subi,conval
                        Ptr{Int32},Ptr{Int32},Ptr{Int32},Ptr{Float64},Ptr{Float64},# grdconptrb,grdconptre,grdconsub,grdconval,grdlag
                        Int32,Ptr{Int32},Ptr{Int32},Ptr{Int32},Ptr{Float64}))# maxnumhesnz, numhesnz.hessubi,hessubj,hesval
    @Mosek.msk_ccall("putnlfunc",
             Int32, (Ptr{Void},Any,Ptr{Void},Ptr{Void}),
             m.task.task, cb, nlgetsp, nlgetva)
end

MathProgBase.getvarLB(m::MosekNonlinearModel)                            = MathProgBase.getvarLB(m.m)
MathProgBase.getvarUB(m::MosekNonlinearModel)                            = MathProgBase.getvarUB(m.m)
MathProgBase.getconstrLB(m::MosekNonlinearModel)                         = MathProgBase.getconstrLB(m.m)
MathProgBase.getconstrUB(m::MosekNonlinearModel)                         = MathProgBase.getconstrUB(m.m)
MathProgBase.setvarLB!(m::MosekNonlinearModel, bnd::Array{Float64,1})    = MathProgBase.setvarLB!(m.m,bnd)
MathProgBase.setvarUB!(m::MosekNonlinearModel, bnd::Array{Float64,1})    = MathProgBase.setvarUB!(m.m,bnd)
MathProgBase.setconstrLB!(m::MosekNonlinearModel, bnd::Array{Float64,1}) = MathProgBase.setconstrLB!(m.m,bnd)
MathProgBase.setconstrUB!(m::MosekNonlinearModel, bnd::Array{Float64,1}) = MathProgBase.setconstrUB!(m.m,bnd)
#MathProgBase.getobj(m::MosekNonlinearModel)                              = MathProgBase.getobj(m.m)
#MathProgBase.setobj!(m::MosekNonlinearModel, c :: Array{Float64,1})      = MathProgBase.setobj!(m,c)
#MathProgBase.getconstrmatrix(m::MosekNonlinearModel)                     = MathProgBase.getconstrmatrix(m.m)
#MathProgBase.addvar!(m::MosekNonlinearModel,bl::Float64,bu::Float64,c::Float64) = MathProgBase.addvar!(m.m)
#MathProgBase.addvar!(m::MosekNonlinearModel,subi::Array{Int32,1},val::Array{Float64,1},bl::Float64,bu::Float64,c::Float64) = MathProgBase.addvar!(m.m)
#MathProgBase.addconstr!(m::MosekNonlinearModel,subj::Array{Int32,1},val::Array{Float64,1},bl::Float64,bu::Float64) = MathProgBase.addconstr!(m.m)
#MathProgBase.numlinconstr(m::MosekNonlinearModel) = MathProgBase.numlinconstr(m.m)
MathProgBase.getobjval(m::MosekNonlinearModel)                           = MathProgBase.getobjval(m.m)
MathProgBase.getsolution(m::MosekNonlinearModel)                         = MathProgBase.getsolution(m.m)
MathProgBase.getconstrsolution(m::MosekNonlinearModel)                   = MathProgBase.getconstrsolution(m.m)
MathProgBase.getreducedcosts(m::MosekNonlinearModel)                     = MathProgBase.getreducedcosts(m.m)
MathProgBase.getconstrduals(m::MosekNonlinearModel)                      = MathProgBase.getconstrduals(m.m)
#MathProgBase.getbasis(m::MosekNonlinearModel)                            = MathProgBase.getbasis(m.m)
MathProgBase.getunboundedray(m::MosekNonlinearModel)                     = MathProgBase.getunboundedray(m.m)
MathProgBase.getinfeasibilityray(m::MosekNonlinearModel)                 = MathProgBase.getinfeasibilityray(m.m)
MathProgBase.getrawsolver(m::MosekNonlinearModel)                        = MathProgBase.getrawsolver(m.m)
#MathProgBase.getsimplexiter(m::MosekNonlinearModel)                      = MathProgBase.getsimplexiter(m.m)
MathProgBase.getbarrieriter(m::MosekNonlinearModel)                      = MathProgBase.getbarrieriter(m.m)
MathProgBase.setwarmstart!(m::MosekNonlinearModel, v::Array{Float64,1})  = MathProgBase.setwarmstart!(m.m, v)
MathProgBase.optimize!(m::MosekNonlinearModel)                           = MathProgBase.optimize!(m.m)
MathProgBase.status(m::MosekNonlinearModel)                              = MathProgBase.status(m.m)
#MathProgBase.getobjbound(m::MosekNonlinearModel) = MathProgBase.getobjbound(m.m)
#MathProgBase.getobjgap(m::MosekNonlinearModel) = MathProgBase.getobjgap(m.m)
MathProgBase.getsolvetime(m::MosekNonlinearModel)                        = MathProgBase.getsolvetime(m.m)
MathProgBase.getsense(m::MosekNonlinearModel)                            = MathProgBase.getsense(m.m)
MathProgBase.setsense!(m::MosekNonlinearModel,sense)                     = MathProgBase.setsense!(m.m)
MathProgBase.freemodel!(m::MosekNonlinearModel)                          = MathProgBase.freemodel!(m.m)
MathProgBase.numvar(m::MosekNonlinearModel)                              = MathProgBase.numvar(m.m)
MathProgBase.numconstr(m::MosekNonlinearModel)                           = MathProgBase.numconstr(m.m)
#MathProgBase.setvartype!(m::MosekNonlinearModel,vtvec::Vector{Symbol}) = MathProgBase.setvartype!(m,vtec)
#MathProgBase.getvartype(m::MosekNonlinearModel) = MathProgBase.getvartype(m.m)


text:
struct Bpe
    rank::Dict{Pair{String,String}, Int}
    merge::Int
    _oldsym::String
    endsym::String
    sepsym::String
    cache::Dict{String, Tuple}
    glossaries::Vector{Union{Regex,String}}

    normalizer::AbstractNormalizer

    function Bpe(bfile::AbstractString;
                 glossaries = Vector{Union{Regex, String}}(),
                 merge::Int = -1, sepsym::String = "", endsym::String = "</w>",
                 have_header::Bool=true,
                 normalizer=UnNormalizer())
        rank = Dict{Pair{String,String}, Int}()
        cache = Dict{String, Tuple}()
        _oldsym = open(bfile) do io
            if have_header
                header = readline(io)
                esi = findlast(isequal("#endsym:"), header)
            else
                esi = nothing
            end

            if esi !== nothing
                _oldsym = header[last(esi)+1:end]
            else
                _oldsym = "</w>"
            end
            set_endsym(_oldsym)

            for (i, line) ∈ enumerate(eachline(io))
                i < merge && break
                pair = Pair(split(line, " ")...)
                rank[pair] = i
            end
            _oldsym
        end
        new(rank, merge, _oldsym, endsym, sepsym, cache, glossaries,
            normalizer)
    end
end

"process a line, remain leading & trailing whitespace"
function process_line(bpe::Bpe, line)
    pat = r"(\S.*\S)"
    m = match(pat, line)
    m === nothing && return line
    sentence = m.captures[1]
    seg = join(segment(bpe, sentence), " ")
    replace(line, pat=>seg)
end

"segment a given sentence"
segment(bpe::Bpe, sentence::AbstractString)::Vector{String} =
    segment_token(bpe, intern.(tokenize(normalize(bpe.normalizer, sentence))))

"bpe tokens"
segment_token(bpe::Bpe, tokens::Vector{String}) =
    mapreduce(x->segment_token(bpe, x),
              (init, x)-> (foreach(x) do y
                           push!(init, y)
                           end;
                           init),
              tokens; init=Vector{String}())

"add seperator and end symbol"
_add_se_sym(bt, i, x, rp, ss)::String = intern(i == length(bt) ? replace(x, rp) : (x * ss))

"bpe a token and add seperator"
function segment_token(bpe::Bpe, token::String)
    isempty(bpe.glossaries) && return segment_token(bpe, bpe(token))
    segment_token(bpe, Tuple(mapreduce(x-> bpe(x),
              (init, x)-> (foreach(x) do y
                           push!(init, y)
                           end;
                           init),
              isolate_gloss(token, bpe.glossaries); init=Vector{String}())))
end

function segment_token(bpe::Bpe, ttp::Tuple)
    rp = bpe._oldsym=>bpe.endsym
    ss = bpe.sepsym
    map(enumerate(ttp)) do (i, x)
        _add_se_sym(ttp, i, x, rp, ss)
    end::Vector{String}
end

"byte pair encode a word"
function (bpe::Bpe)(x::String)::Tuple
    haskey(bpe.cache, x) && return bpe.cache[x]

    for gloss ∈ bpe.glossaries
        if occursin(Regex("^"*gloss.pattern*"\$"), x)
            return bpe.cache[x] = tuple(x)
        end
    end

    xtp = toStrTuple(x)
    xps = bi_pairs(xtp)
    isempty(xps) && return xtp
    while true
        mp = lowestpair(bpe, xps)
        !haskey(bpe.rank, mp) && break

        xtp = merged_pairs(xtp, mp)[1]
        length(xps) == 1 && break
        xps = bi_pairs(xtp)
    end
    bpe.cache[x] = xtp
end

"find pair with lowest rank"
lowestpair(bpe::Bpe, xps) = argmin(
    sizehint!(Dict(
        map(xps) do p
        p=>get(bpe.rank, p, typemax(Int))
        end),
              length(xps))
)

function Base.show(io::IO, b::Bpe)
    println(io, "Bpe(merge=$(b.merge), sepsym=\"$(b.sepsym)\", ",
            "endsym=\"$(b.endsym)\", num_glossaries=$(length(b.glossaries)), ",
            "normalizer=$(typeof(b.normalizer)))")
    io
end

ext: jl


text:
@testset "Testing Convolution" begin
    activations = [identity, Flux.relu, Flux.sigmoid, abs]
    kernel_sizes = [(1, 1), (3, 3)]
    paddings = [0, 1]
    strides = [1, 2, 3]

    iter = Iterators.product(activations, kernel_sizes, paddings, strides)
    for tup in iter
        @show tup
        # Unpack
        activation = tup[1]
        kernel_size = tup[2]
        padding = tup[3]
        stride = tup[4]

        w = randn(Float32, kernel_size..., 16, 16)
        b = randn(Float32, size(w, ndims(w)))
        x = randn(Float32, 16, 16, 16, 8)
        X = OneDNN.Memory(x)

        # Just another reminder - OneDNN's convolution is actually cross-correlation.
        conv_flux = Flux.CrossCor(
            w, b, activation; stride = stride, pad = padding,
        )
        conv_onednn = OneDNN.Conv(conv_flux)

        # See if the internal weight gets optimized properly.
        @test conv_onednn.optimized_weights == false
        y_pre = conv_onednn(X)
        @test conv_onednn.optimized_weights == true
        y_post = conv_onednn(X)
        y_ref = conv_flux(x)
        @test isapprox(y_ref, OneDNN.materialize(y_pre))
        @test isapprox(y_ref, OneDNN.materialize(y_post))

        # Ensure type inference works correctly
        @inferred conv_onednn(X)

        # Now - try backprop
        out_onednn, back_onednn = Zygote._pullback(conv_onednn, X)
        out_flux, back_flux = Zygote._pullback(conv_flux, x)

        isa(out_onednn, OneDNN.Memory)
        isapprox(out_flux, OneDNN.materialize(out_onednn))
        @inferred Zygote._pullback(conv_onednn, X)

        dx_onednn = back_onednn(out_onednn)
        dx_flux = back_flux(out_flux)

        @test isa(dx_onednn, Tuple)
        @test isa(dx_flux, Tuple)

        @test isapprox(dx_flux[1].weight, OneDNN.materialize(dx_onednn[1].weights))
        @test isapprox(dx_flux[1].bias, OneDNN.materialize(dx_onednn[1].bias))
        @test isapprox(dx_flux[2], OneDNN.materialize(dx_onednn[2]))

        # Type inference
        @inferred back_onednn(out_onednn)
    end
end


text:
using Fermionic
using SparseArrays
using Test

@testset "Fermionic.jl" begin
    @test ucnot(Op(2),1,2)[4,3] == 1.0
    @test ucnot(Op(6),4,2)[48,64] == 1.0
    @test ucnot(Op(3),2,3)[1,2] == 0.0
    @test_throws ArgumentError ucnot(Op(3),4,2)
    @test_throws ArgumentError ucnot(Op(3),2,2)
    @test hadamard(Op(3),2,1)[2,2] == 1.0
    @test hadamard(Op(3),3,1)[5,2] == 1/sqrt(2)
    @test_throws ArgumentError hadamard(Op(3),2,2)
    @test_throws ArgumentError hadamard(Op(3),2,4)
    @test phase(Op(3),2,pi/4)[3,3] == round(1/sqrt(2)+im/sqrt(2),digits=15)
    @test phase(Op(6),2,pi/4)[1,3] == 0.0
    @test_throws ArgumentError phase(Op(3),4,pi)
    @test swap(Op(3),1,2)[3,5] == 1.0
    @test_throws ArgumentError swap(Op(3),1,4)
    @test_throws ArgumentError swap(Op(3),1,1)
    @test swap(Op(5),1,2)[5,6] == 0.0
    @test swap(Op(3),1,2)[8,8] == -1.0
    @test sigma_x(Op(3),2)[1,3] == 1.0
    @test sigma_x(Op(5),2)[1,1] == 0.0
    @test_throws ArgumentError sigma_x(Op(3),4)
    @test sigma_y(Op(2),2)[2,1] == 0.0 + 1.0*im
    @test sigma_y(Op(4),3)[2,4] == -1.0*im
    @test_throws ArgumentError sigma_y(Op(3),4)
    @test sigma_z(Op(4),3)[2,4] == 0.0
    @test sigma_z(Op(3),3)[3,3] == 1.0
    @test_throws ArgumentError sigma_z(Op(3),4)
end


text:
module SA

using Random

using ..Mapper2.Helper
Helper.@SetupDocStringTemplates

using ..Mapper2.MapperCore
using ..Mapper2.MapperGraphs

using DataStructures
using Formatting

import Base: getindex
import Base: @propagate_inbounds
import ..Mapper2.MapperCore: getaddress

using Logging

export SAStruct, getdistance, warmup

"""
    place!(map::Map; kwargs...) :: SAState

Run simulated annealing placement directly on `map`.

Records the following metrics into `map.metadata`:

* `placement_struct_time` - Amount of time it took to build the
    [`SAStruct`](@ref) from `map`.

* `placement_struct_bytes` - Number of bytes allocated during the construction
    of the [`SAStruct`](@ref)

* `placement_time` - Running time of placement.

* `placement_bytes` - Number of bytes allocated during placement.

* `placement_objective` - Final objective value of placement.

Keyword Arguments
-----------------
* `seed` - Seed to provide the random number generator. Specify this to a
    constant value for consistent results from run to run.

    Default: `rand(UInt64)`

* `move_attempts :: Integer` - Number of successful moves to generate between
    state updates. State updates include adjusting temperature, move distance
    limit, state displaying etc.

    Higher numbers will generally yield higher quality placement but with a
    longer running time.

    Default: `20000`

* `initial_temperature :: Float64` - Initial temperature that the system begins
    its warming process at. Due to the warming procedure, this should not have
    much of an affect on placement.

    Default: `1.0`.

* `supplied_state :: Union{SAState,Nothing}` - State type to use for this
    placement. Can be used to resume placement where it left off from a previous
    run. If `nothing`, a new `SAState` object will be initialized.

    Default: `nothing`

* `movegen :: MoveGenerator` - The [`MoveGenerator`](@ref) to use for this
    placement.

    Default: [`CachedMoveGenerator`](@ref)

* `warmer` - The [`SAWarm`](@ref) warming schedule to use.

    Default: [`DefaultSAWarm`](@ref)

* `cooler` - The [`SACool`](@ref) cooling schedule to use.

    Default: [`DefaultSACool`](@ref)

* `limiter` - The [`SALimit`](@ref) move distance limiting algorithm to
    use.

    Default: [`DefaultSALimit`](@ref)

* `doner` - The [`SADone`](@ref) exit condition to use.

    Default: [`DefaultSADone`](@ref)
"""
function place!(m::Map; seed = rand(UInt64), kwargs...)
    # Configure random seed.
    @info "Using Seed: $seed"
    Random.seed!(seed)
    # Record how long it takes to build the placement struct and the number
    # of bytes allocated to do so.
    sastruct, struct_time, struct_bytes, _, _ = @timed SAStruct(m; kwargs...)
    # Record how long it takes for placement - after creation of the SAStruct.
    sastate, place_time, place_bytes, _, _ = @timed place!(sastruct; kwargs...)
    record(m, sastruct)

    # Record the mapping time in the map metadata.
    m.metadata["placement_struct_time"] = struct_time
    m.metadata["placement_struct_bytes"] = struct_bytes
    m.metadata["placement_time"] = place_time
    m.metadata["placement_bytes"] = place_bytes
    m.metadata["placement_objective"] = map_cost(sastruct)

    return m
end

function warmup(m::Map; kwargs...)
    # Strip "move_attempts" from the kwargs and set it so a low value
    if length(kwargs) > 0
        stripped_kwargs = [k => kwargs[k] for k in keys(kwargs) if k != :move_attempts]
        push!(stripped_kwargs, :move_attempts => 5000)
    else
        stripped_kwargs = (move_attempts = 5000,)
    end

    # Run placement.
    sa_struct = SAStruct(m; stripped_kwargs...)
    place!(sa_struct; stripped_kwargs...)
    return nothing
end

include("PathTables.jl")
include("MapTables.jl")
include("Distance.jl")
include("Struct.jl")
include("Methods.jl")
include("InitialPlacement.jl")
include("State.jl")
include("MoveGenerators.jl")
include("Algorithm.jl")

end #module SA


text:
using Plots
using Random

# MODE = :translate
MODE = :rotate

# Model
include_model("planar_push_block_v3")

# Horizon
T = 26

# Time step
h = 0.1

# Bounds
_uu = Inf * ones(model.m)
_uu[model.idx_u] .= Inf
_uu[model.idx_u[1:2]] .= 5.0
_uu[model.idx_λ[1:4]] .= μ_surface * gravity * h * 0.25
_ul = zeros(model.m)
_ul[model.idx_u] .= -Inf
_ul[model.idx_u[1:2]] .= -5.0
_ul[model.idx_λ[1:4]] .= μ_surface * gravity * h * 0.25

ul, uu = control_bounds(model, T, _ul, _uu)

if MODE == :rotate
    # Initial and final states (translate + rotate)
    q1 = [0.0, 0.0, 0.0, -r - 1.0e-8, -0.01]
    x1 = [q1; q1]
    x_goal = 0.5
    y_goal = 0.5
    θ_goal = 0.5 * π
elseif MODE == :translate
    # Initial and final states (translate)
    q1 = [0.0, 0.0, 0.0, -r - 1.0e-8, 0.0]
    x1 = [q1; q1]
    x_goal = 1.0
    y_goal = 0.0
    θ_goal = 0.0 * π
end

qT = [x_goal, y_goal, θ_goal, x_goal-r, y_goal-r]

xT = [qT; qT]
xl, xu = state_bounds(model, T, x1 = x1, xT = xT)
xl[T][4:5] .= -Inf
xl[T][9:10] .= -Inf
xu[T][4:5] .= Inf
xu[T][9:10] .= Inf

# Objective
include_objective("velocity")
obj_velocity = velocity_objective(
    [Diagonal([1.0, 1.0, 1.0, 0.1, 0.1]) for t = 1:T-1],
    model.nq,
    h = h,
    idx_angle = collect([3]))

_Q_track = [1.0, 1.0, 1.0, 0.1, 0.1]
Q_track = 1.0 * Diagonal([_Q_track; _Q_track])
obj_tracking = quadratic_tracking_objective(
    [Q_track for t = 1:T],
    # [Diagonal(0.1 * ones(model.m)) for t = 1:T-1],
	[Diagonal([1.0e-1 * ones(model.nu);
		zeros(model.nc);
		zeros(model.nb);
		zeros(model.m - model.nu - model.nc - model.nb)]) for t = 1:T-1],
    [xT for t = 1:T],
    [zeros(model.m) for t = 1:T])

obj_penalty = PenaltyObjective(1.0e4, model.m)
obj = MultiObjective([obj_tracking, obj_penalty, obj_velocity])

# Constraints
include_constraints(["contact", "stage"])
t_idx = vcat([t for t = 1:T-1])
con_contact = contact_constraints(model, T)
con = multiple_constraints([con_contact])#, con_ctrl_comp, con_ctrl_lim])

# Problem
prob = trajectory_optimization_problem(model,
               obj,
               T,
               h = h,
               xl = xl,
               xu = xu,
               ul = ul,
               uu = uu,
               con = con)

# Trajectory initialization
x0 = linear_interpolation(x1, xT, T) # linear interpolation on state

# Random.seed!(1)
# for i = 1:1
# u0 = [0.1 * randn(model.m) for t = 1:T-1] # random controls
u0 = [0.1 * ones(model.m) for t = 1:T-1] # random controls
if MODE == :rotate
    # u0 = [0.1 * randn(model.m) for t = 1:T-1] # random controls
    for t = 1:T-1
        if t < 5
            u0[t][1:2] = [1.0; 0.0]
        elseif t < 10
            u0[t][1:2] = [0.5; 0.0]
        else
            u0[t][1:2] .= 0.0
        end
    end
elseif MODE == :translate
    # translation
    # u0 = [0.1 * randn(model.m) for t = 1:T-1] # random controls
    for t = 1:T-1
        if t < 5
            u0[t][1:2] = [1.0; 0.0]
        else
            u0[t][1:2] .= 0.0
        end
    end
end

# Pack trajectoriesff into vector
z0 = pack(x0, u0, prob)

#NOTE: may need to run examples multiple times to get good trajectories
# Solve nominal problem
@time z̄, info = solve(prob, copy(z0),
	tol = 1.0e-3, c_tol = 1.0e-3, max_iter = 2500)

@show check_slack(z̄, prob)
# end

# # translate iters:
# iters = [91, 101, 122, 77, 106]
# # mean(iters)
# # std(iters)
#
# # translate + rotation iters:
# iters = [122, 85, 70, 121, 168]
# mean(iters)
# std(iters)

x̄, ū = unpack(z̄, prob)
q̄ = state_to_configuration(x̄)
q = state_to_configuration(x̄)
u = [u[model.idx_u] for u in ū]
γ = [u[model.idx_λ] for u in ū]
b = [u[model.idx_b] for u in ū]
h̄ = h

# @save joinpath(pwd(), "examples/implicit_dynamics/examples/trajectories/planar_push_translate_direct.jld2") x̄ ū
# @load joinpath(pwd(), "examples/implicit_dynamics/examples/trajectories/planar_push_translate_direct.jld2") x̄ ū

# @save joinpath(pwd(), "examples/implicit_dynamics/examples/trajectories/planar_push_rotate_direct.jld2") x̄ ū
@load joinpath(pwd(), "examples/implicit_dynamics/examples/trajectories/planar_push_rotate_direct.jld2") x̄ ū

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
#open(vis)
visualize!(vis, model,
    q, u,
    Δt = h,
	r = model.block_dim[1] + model.block_rnd)

plot(hcat(q̄...)', color = :red, width = 1.0, labels = "")
plot(hcat([ū..., ū[end]]...)[model.idx_u[1:2], :]', linetype = :steppost, color = :black, width = 1.0, labels = "")

# q_array = hcat(q...)
# u_array = hcat(u...)
#
# traj = Dict("q" => q_array, "u" => u_array)

# using NPZ
# i = 2
# file_path = joinpath(pwd(), "examples/contact_implicit/manipulation/", "traj$i.npz")
# npzwrite(file_path, traj)

max_issues_repo_name: thowell/DirectMotionPlanning


text:
using ProgressMeter, JLD2, MacroTools

"""
The `AbstractSaver` type is the supertype of all saver types of `ParallelTools`.
"""
abstract type AbstractSaver end

"""
 	ConcurrentSaver{R<:RemoteChannel} <: AbstractSaver

Concurrent saver type.
"""
mutable struct ConcurrentSaver{R<:RemoteChannel} <: AbstractSaver
    readout_channel::R
    process_id::Int
    rem_packets::Int
    active::Bool
    progressbar::Bool
    fpath::String
end

"""
	ConcurrentSaver([remch, pid,] npackets, fpath; progressbar=false)

Construct a `ConcurrentSaver` to save to disk results of a pmap in real time
from a given process.

# Arguments
* `remch`: `RemoteChannel` instance.
* `pid`: saver's process id.
* `npackets`: number of results to be read by the saver.
* `fpath`: path of the file where data is to be written.
* `progressbar=false`: if true, a progress bar is displayed.
"""
function ConcurrentSaver(r::RemoteChannel, pid::Int, npackets::Int, fpath::String; progressbar::Bool=false)
    @assert pid == r.where "Saver's process must match that of the remote channel."
    @assert npackets > 0 "The number of packets to be read and written to disk must be a natural integer."
    @assert isdir(dirname(fpath)) "ERROR: accessing "*dirname(fpath)*": No such directory."
	@assert !isfile(fpath) "$(basename(fpath)) already exists at $(dirname(fpath))."

    return ConcurrentSaver{typeof(r)}(r, pid, npackets, true, progressbar, fpath)
end
ConcurrentSaver(r::RemoteChannel, npackets::Int, fpath::String; progressbar::Bool=false) = ConcurrentSaver(r,r.where,npackets,fpath; progressbar=progressbar)
ConcurrentSaver(pid::Int, npackets::Int, fpath::String; progressbar::Bool=false) = ConcurrentSaver(RemoteChannel(()->Channel{Any}(npackets), pid),pid,npackets,fpath; progressbar=progressbar)
ConcurrentSaver(npackets::Int, fpath::String; progressbar::Bool=false) = ConcurrentSaver(myid(),npackets,fpath; progressbar=progressbar)

function launch_saver!(s::ConcurrentSaver{R}) where {R<:RemoteChannel}
    @spawnat s.process_id _launch_saver(s)
end

function _launch_saver!(s::ConcurrentSaver{R}) where {R<:RemoteChannel}
    Npackets::Int = s.rem_packets

    if s.progressbar
        # Set up a progress bar
        progress = Progress(Npackets);
        ProgressMeter.update!(progress, 0);
    end

    jldopen(s.fpath, "a+") do file
        println("Saving data to ",s.fpath)

        while s.rem_packets > 0
            # Retrieve a queued result
            n, sol = take!(s.readout_channel);
            s.rem_packets -= 1;
            for (key, val) in sol
                file["$n/" * key] = val;
            end
            s.progressbar ? ProgressMeter.next!(progress) : nothing;
        end
        # Set progress bar to 100%
        if s.progressbar && (progress.counter < Npackets) ProgressMeter.update!(progress, Npackets); end;
    end

    s.active = false

    nothing
end;

"""
	@with_saver

A macro to save to disk results of a `pmap` call in real time.
See also: [`@save_to`](@ref)

# Example
```julia-repl
julia> saver = ConcurrentSaver(10, "some/valid/path.jld2");
julia> @with_saver saver pmap(1:10) do x
							y = x * sqrt(x)
							y, x^2, x^3
						end
julia> jldopen("some/valid/path.jld2","r") do file
			for i in 1:10
				group = file[string(i)]
				println("Iteration ",i)
		   		for k in keys(group)
					println(k,"\t=>\t",group[k])
				end
			end
		end
julia>
Iteration 1
x ^ 3   =>  1
x ^ 2   =>  1
y   =>  1.0
Iteration 2
x ^ 3   =>  8
x ^ 2   =>  4
y   =>  2.8284271247461903
Etc.
```
"""
macro with_saver(s, expr)
    @assert isa(expr, Expr) "Not an expression."
    @assert  expr.head == :do

    call = first(expr.args)
    @assert call.head == :call && call.args[1] == :pmap
    args   = call.args[2:end][typeof.(call.args[2:end]) .== Symbol]
    kwargs = call.args[2:end][typeof.(call.args[2:end]) .== Expr]

    wp = Symbol()
    iterable = Symbol()
    if length(args) == 2
        wp, iterable = args
    elseif length(args) == 1
        wp = :(CachingPool(workers()))
        iterable = args[1]
    else
        throw(error("WTF are these function arguments? $(length(args))"))
    end

    index = expr.args[end].args[1].args[]
    old_inner_expr = expr.args[end].args[end]
    rvars = []
    if length(old_inner_expr.args[end].args) > 1  # No return keyword was used
        rvars = old_inner_expr.args[end].args
    else                                          # return was used
        rvars = old_inner_expr.args[end].args[].args
    end
    rexpr = Expr(:tuple, :i, Expr(:call, :Dict, [Expr(:call,:(=>),:($rv_str), rv) for (rv_str, rv) in zip(string.(rvars), rvars)]...))

    quote
		@assert typeof($s) <: ConcurrentSaver "@with_saver requires a ConcurrentSaver instance as first parameter."
		@assert !isfile($s.fpath) basename($s.fpath)*" already exists at "*dirname($s.fpath)*"."
		@sync begin
            @spawnat $s.process_id ParallelTools._launch_saver!($s)
            @sync pmap($wp,eachindex($iterable)) do i
                put!($s.readout_channel,begin
                    $index = $iterable[i]
                    $(old_inner_expr.args[1:end-1]...)
                    $rexpr
                end);
                nothing
            end
        end
        nothing
    end |> esc |> prettify
end

"""
	@sava_to

A macro to save to disk results of a `pmap` call in real time.
See also: [`@with_saver`](@ref)

# Example
```julia-repl
julia> @save_to "some/valid/path.jld2" pmap(1:10) do x
							y = x * sqrt(x)
							y, x^2, x^3
						end
julia> jldopen("some/valid/path.jld2","r") do file
			for i in 1:10
				group = file[string(i)]
				println("Iteration ",i)
		   		for k in keys(group)
					println(k,"\t=>\t",group[k])
				end
			end
		end
julia>
Iteration 1
x ^ 3   =>  1
x ^ 2   =>  1
y   =>  1.0
Iteration 2
x ^ 3   =>  8
x ^ 2   =>  4
y   =>  2.8284271247461903
Etc.
```
"""
macro save_to(fpath, expr)
    @assert isa(expr, Expr) "Not an expression."
    @assert  expr.head == :do
	call = first(expr.args)
	iterable = last(call.args[2:end][typeof.(call.args[2:end]) .== Symbol])
    quote
		@assert typeof($fpath) == String "@save_at requires a valid filename as first parameter."
		@assert !isfile($fpath) basename($fpath)*" already exists at "*dirname($fpath)*"."
		let s = ConcurrentSaver(length($iterable), $fpath)
			@with_saver s $expr
		end
    end |> esc |> prettify
end


text:
using ArcadeLearningEnvironment
using ArgParse
using CartesianGeneticProgramming
using Dates
using IICGP
using Random
using MappedArrays
using BenchmarkTools
using Images
using ImageView
using Plots
using PlotThemes
using Statistics

theme(:dark)

"""
    get_two_arity(nodes::Array{Node,1}, arity_dict::Dict)

Creates a BitArray tracking two-arity of the functions in the provied
nodes list.
"""
function get_two_arity(nodes::Array{Node}, arity_dict::Dict)
    two_arity = falses(length(nodes))
    for i in eachindex(nodes)
        if arity_dict[String(Symbol(nodes[i].f))] == 2
            two_arity[i] = true
        end
    end
    two_arity
end

function recur_active(nodes::Array{Node}, i::Int16, active::BitArray,
                      two_arity::BitArray)
    active[i] = true
    if nodes[i].x > 0
        recur_active(nodes, nodes[i].x, active, two_arity)
    end
    if nodes[i].y > 0 && two_arity[i]
        recur_active(nodes, nodes[i].y, active, two_arity)
    end
end

function find_active(nodes::Array{Node}, outputs::Array{Int16},
                     two_arity::BitArray)
    active = falses(length(nodes))
    for o in outputs
        recur_active(nodes, o, active, two_arity)
    end
    active
end

"default function for nodes, will cause error if used as a function node"
function f_null(args...)::Nothing
    nothing
end

"""
    function play_atari(
        encoder::CGPInd,
        reducer::AbstractReducer,
        controller::CGPInd;
        reducer_type::String="pooling",
        seed::Int64=0,
        max_frames::Int64=18000,
        sleep_time::Float64=0.0,
        render::Bool=true,
        save::Bool=false
    )

Play Atari and display individual's encoding.
"""
function play_atari(
    encoder::CGPInd,
    reducer::AbstractReducer,
    controller::CGPInd,
    reducer_type::String;
    seed::Int64=0,
    max_frames::Int64=1000,
    sleep_time::Float64=0.0,
    render::Bool=true,
    save::Bool=false,
    save_repo::String=string("gifs/", args["game"], "/")
)
    game = Game(args["game"], seed)
    reward = 0.0
    frames = 1
    a_prev = nothing
    c_prev = nothing
    while ~game_over(game.ale)
        # Replace get_rgb method to integrate visu
        rawscreen = getScreenRGB(game.ale)
        rgb = permutedims(reshape(rawscreen, (3, game.width, game.height)), [1, 3, 2])
        rgb = [Array{UInt8}(rgb[i,:,:]) for i in 1:3]

        # Process
        enco_out, features, out = IICGP.process_full(
            encoder,
            reducer,
            controller,
            rgb
        )

        # Render and save
        if reducer_type == "pooling"
            plt = plot_encoding(n_in, enco.buffer, features)
        elseif reducer_type == "centroid"
            plt = plot_centroids(enco_out, features)
        end
        if render
            display(plt)
        end
        if save
            savefig(plt, string(save_repo, "$frames.png"))
            # println(frames)
        end

        sleep(sleep_time)
        action = game.actions[argmax(out)]
        reward += act(game.ale, action)
        frames += 1
        game.ale
        if frames > max_frames
            break
        end
    end
    close!(game)
    [reward]
end

##

s = ArgParseSettings()
@add_arg_table! s begin
    "--game"
    help = "game rom name"
    default = "freeway"
    "--seed"
    help = "random seed for evolution"
    arg_type = Int
    default = 0
    "--ind"
    help = "individual for evaluation"
    arg_type = String
    default = ""
end
args = parse_args(ARGS, s)
Random.seed!(args["seed"])

# Temporarily open a game to retrieve parameters
game = Game(args["game"], 0)
rawscreen = getScreenRGB(game.ale)
rgb = permutedims(reshape(rawscreen, (3, game.width, game.height)), [1, 3, 2])
n_in = 3  # RGB images
n_out = length(getMinimalActionSet(game.ale))  # One output per legal action
img_size = size(rgb)[2:3]
d_fitness = 1
close!(game)

# Encoder
enco_nodes = [
    Node(1, 1, IICGP.CGPFunctions.f_motion_capture, [0.5], false)
]
enco_outputs = Int16[1, 4]
enco_cfg = cfg_from_info(enco_nodes, n_in, enco_outputs, IICGP.CGPFunctions,
                         d_fitness)
enco = IPCGPInd(enco_nodes, enco_cfg, enco_outputs, img_size)

# Reducer
reducer_cfg = Dict(
    "type" => "pooling",
    "pooling_function" => "mean",
    "features_size" => 5,
    "n_centroids" => 20
)
redu = Reducer(reducer_cfg, n_in=enco_cfg.n_out, img_size=img_size)
if reducer_cfg["type"] == "pooling"
    features_size = reducer_cfg["features_size"]^2
elseif reducer_cfg["type"] == "centroid"
    features_size = 2 * reducer_cfg["n_centroids"]
end

# Controller
cont_nodes = [
    Node(1, 2, IICGP.CGPFunctions.f_subtract, [0.5], false),
    Node(1, 2, IICGP.CGPFunctions.f_add, [0.5], false),
    Node(3, 3, IICGP.CGPFunctions.f_cos, [0.6], false)
]
cont_outputs = Int16[1, 2, 3]
cont_n_in = length(enco_outputs) * features_size^2
cont_cfg = cfg_from_info(cont_nodes, cont_n_in, cont_outputs, IICGP.CGPFunctions, 1)
cont = CGPInd(cont_nodes, cont_cfg, cont_outputs)

##

# Play game
play_atari(
    enco, redu, cont,
    reducer_type=reducer_cfg["type"],
    max_frames=2,
    sleep_time=0.0,
    render=true,
    save_repo="gifs/",
    save=true
)

##
using Plots.PlotMeasures

# Temporarily open a game to retrieve parameters
game = Game(args["game"], 0)
rawscreen = getScreenRGB(game.ale)
rgb = permutedims(reshape(rawscreen, (3, game.width, game.height)), [1, 3, 2])
img_size = size(rgb)[2:3]
rgb = [Array{UInt8}(rgb[i,:,:]) for i in 1:3]
n_in = 3  # RGB images
n_out = length(getMinimalActionSet(game.ale))  # One output per legal action
d_fitness = 1
close!(game)

x = rgb[1]
mult = 2
sz = (mult*size(x)[2], mult*size(x)[1])
plt = heatmap(
    x, color=:grays, ratio=:equal,
    yflip=true,
    leg=false,
    framestyle=:none,
    padding=(0.0, 0.0),
    margin=-100mm,
    size=sz
)


text: Dict{String, Any}("empty2" => Dict{String, Any}(), "empty1" => Dict{String, Any}(), "nested_empty" => Dict{String, Any}("empty" => Dict{String, Any}()), "many_empty" => Any[Dict{String, Any}(), Dict{String, Any}(), Dict{String, Any}()], "empty_in_array2" => Any[Dict{String, Any}(), Dict{String, Any}("not_empty" => Dict{String, Any}("value" => "1", "type" => "integer"))], "empty_in_array" => Any[Dict{String, Any}("not_empty" => Dict{String, Any}("value" => "1", "type" => "integer")), Dict{String, Any}()])


text:
function different_raster_sizes_error(name1, name2)
    @error "$(name1) and $(name2) are different sizes."
end

function different_raster_projections_warning(name1, name2)
    @warn "$(name1) and $(name2) have different projections. This could
mean they are not properly aligned. Attempting to proceed anyway. Press ctrl-c
(or command-c on Mac) to quit."
    sleep(10) # to give the user a chance to read the message and quit before
              # a bunch more terminal outputs starts
end

function bad_resistance_values_error()
    @error "Resistance (or conductance) surface contains 0 or negative values"
end

function even_block_size_warning()
    @warn "block_size is even, but must be odd. Using block_size + 1."
end

function missing_args_error(missing_args)
    @error "The following arguments are missing from your .ini (or config Dictionary) with no defaults:
       $(join(map(string, missing_args), " "))"
end
function check_raster_alignment(raster1, raster2, name1, name2, allow_different_projections)
    sizes_not_equal = size(raster1[1]) != size(raster2[1])
    projections_not_equal = (raster1[2] != raster2[2]) || (raster1[3] != raster2[3])

    if projections_not_equal && !allow_different_projections
        different_raster_projections_warning(name1, name2)
    end

    if sizes_not_equal
        different_raster_sizes_error(name1, name2)
    end

    sizes_not_equal
end

function check_resistance_values(resistance)
    bad_values = minimum(resistance[(!).(ismissing.(resistance))]) <= 0

    if bad_values
        bad_resistance_values_error()
    end

    bad_values
end

function check_block_size!(int_arguments)
    if iseven(int_arguments["block_size"])
        even_block_size_warning()
        int_arguments["block_size"] = int_arguments["block_size"] + 1
    end
    int_arguments["block_size"]
end

function check_missing_args_ini(cfg)
    indices = indexin(REQUIRED_ARGS_INI, convert.(String, keys(cfg)))
    indices[indices .== nothing] .= 0

    is_missing = indices .== 0
    missing_args = REQUIRED_ARGS_INI[is_missing]

    if !isempty(missing_args)
        missing_args_error(missing_args)
    end

    !isempty(missing_args)
end

function check_missing_args_dict(cfg)
    indices = indexin(REQUIRED_ARGS_DICT, convert.(String, keys(cfg)))
    indices[indices .== nothing] .= 0

    is_missing = indices .== 0
    missing_args = REQUIRED_ARGS_DICT[is_missing]

    if !isempty(missing_args)
        missing_args_error(missing_args)
    end

    !isempty(missing_args)
end

function check_solver!(cfg)
    if (cfg["solver"] in SOLVERS) == false
        @warn "Got unsupported value for solver: $(cfg["solver"]). Using the default, cg+amg"
        cfg["solver"] = "cg+amg"
    end
end

function check_unsupported_args(cfg)
    arg_names = String.(keys(cfg))

    bad_args = arg_names[map(x->(x ∉ SUPPORTED_ARGS), arg_names)]

    if bad_args != [""]
        @warn "The following unsupported arguments were provided and will be ignored:
         $(join(map(string, bad_args), " "))"
    end
end


text:
struct Pnmp1toPlm{T} <: AbstractRotation{T}
    rotations::Vector{Givens{T}}
end

function Pnmp1toPlm(::Type{T}, n::Int, m::Int, α::T, β::T, γ::T) where T
    G = Vector{Givens{T}}(n)
    @inbounds for ℓ = 1:n
        c = sqrt((2m+β+γ+2)/(ℓ+2m+β+γ+2)*(2ℓ+2m+α+β+γ+2)/(ℓ+2m+α+β+γ+2))
        s = sqrt(ℓ/(ℓ+2m+β+γ+2)*(ℓ+α)/(ℓ+2m+α+β+γ+2))
        G[n+1-ℓ] = Givens(ℓ, ℓ+1, c, s)
    end
    Pnmp1toPlm(G)
end

@inline length(P::Pnmp1toPlm) = length(P.rotations)
@inline getindex(P::Pnmp1toPlm, i::Int) = P.rotations[i]

function LAmul!(C::Pnmp1toPlm, A::AbstractVecOrMat)
    @inbounds for i = 1:length(C)
        mul!(C.rotations[i], A)
    end
    A
end

function LAmul!(A::AbstractMatrix, C::Pnmp1toPlm)
    @inbounds for i = length(C):-1:1
        mul!(A, C.rotations[i])
    end
    A
end

struct TriRotationPlan{T} <: AbstractRotation{T}
    layers::Vector{Pnmp1toPlm{T}}
end

function TriRotationPlan(::Type{T}, n::Int, α::T, β::T, γ::T) where T
    layers = Vector{Pnmp1toPlm{T}}(n)
    @inbounds for m = 0:n-1
        layers[m+1] = Pnmp1toPlm(T, n-m, m, α, β, γ)
    end
    TriRotationPlan(layers)
end

function LAmul!(P::TriRotationPlan, A::AbstractMatrix)
    M, N = size(A)
    @inbounds for m = N-1:-1:1
        layer = P.layers[m]
        for ℓ = (m+1):N
            @simd for i = 1:length(layer)
                G = layer[i]
                a1, a2 = A[G.i1,ℓ], A[G.i2,ℓ]
                A[G.i1,ℓ] = G.c*a1 + G.s*a2
                A[G.i2,ℓ] = G.c*a2 - G.s*a1
            end
        end
    end
    A
end

if VERSION < v"0.7-"
    function Base.At_mul_B!(P::TriRotationPlan, A::AbstractMatrix)
        M, N = size(A)
        @inbounds for m = 1:N-1
            layer = P.layers[m]
            for ℓ = (m+1):N
                @simd for i = length(layer):-1:1
                    G = layer[i]
                    a1, a2 = A[G.i1,ℓ], A[G.i2,ℓ]
                    A[G.i1,ℓ] = G.c*a1 - G.s*a2
                    A[G.i2,ℓ] = G.s*a1 + G.c*a2
                end
            end
        end
        A
    end

    Base.Ac_mul_B!(P::TriRotationPlan, A::AbstractMatrix) = At_mul_B!(P, A)
else
    function LinearAlgebra.lmul!(Pt::Transpose{T,<:TriRotationPlan}, A::AbstractMatrix) where T
        P = parent(Pt)
        M, N = size(A)
        @inbounds for m = 1:N-1
            layer = P.layers[m]
            for ℓ = (m+1):N
                @simd for i = length(layer):-1:1
                    G = layer[i]
                    a1, a2 = A[G.i1,ℓ], A[G.i2,ℓ]
                    A[G.i1,ℓ] = G.c*a1 - G.s*a2
                    A[G.i2,ℓ] = G.s*a1 + G.c*a2
                end
            end
        end
        A
    end

    LinearAlgebra.lmul!(Pc::Adjoint{T,<:TriRotationPlan}, A::AbstractMatrix) where T =
        lmul!(transpose(parent(Pc)), A)
end

struct SlowTriangularHarmonicPlan{T} <: TriangularHarmonicPlan{T}
    RP::TriRotationPlan{T}
    p::NormalizedLegendreToChebyshevPlan{T}
    pinv::ChebyshevToNormalizedLegendrePlan{T}
    B::Matrix{T}
end

function SlowTriangularHarmonicPlan(A::Matrix{T}, α, β, γ) where T
    @assert β == γ == -half(T)
    @assert α == zero(T)
    M, N = size(A)
    n = N
    RP = TriRotationPlan(T, n-1, α, β, γ)
    p = plan_normleg2cheb(A)
    pinv = plan_cheb2normleg(A)
    B = zeros(A)
    SlowTriangularHarmonicPlan(RP, p, pinv, B)
end

function LAmul!(Y::Matrix, SP::SlowTriangularHarmonicPlan, X::Matrix)
    RP, p, B = SP.RP, SP.p, SP.B
    copyto!(B, X)
    mul!(RP, B)
    M, N = size(X)
    for J = 1:N
        mul_col_J!!(Y, p, B, J)
    end
    @inbounds for J = 1:N
        nrm = sqrt((2-δ(J-1,0))/π)
        @simd for I = 1:M
            Y[I,J] *= nrm
        end
    end
    Y
end

if VERSION < v"0.7-"
    function Base.At_mul_B!(Y::Matrix, SP::SlowTriangularHarmonicPlan, X::Matrix)
        RP, pinv, B = SP.RP, SP.pinv, SP.B
        copyto!(B, X)
        M, N = size(X)
        @inbounds for J = 1:N
            nrm = sqrt(π/(2-δ(J-1,0)))
            @simd for I = 1:M
                B[I,J] *= nrm
            end
        end
        for J = 1:N
            mul_col_J!!(Y, pinv, B, J)
        end
        tri_zero_spurious_modes!(At_mul_B!(RP, Y))
    end

    Base.Ac_mul_B!(Y::Matrix, SP::SlowTriangularHarmonicPlan, X::Matrix) = At_mul_B!(Y, SP, X)
else
    function LinearAlgebra.mul!(Y::Matrix, SPt::Transpose{T,<:SlowTriangularHarmonicPlan}, X::Matrix) where T
        SP = parent(SPt)
        RP, pinv, B = SP.RP, SP.pinv, SP.B
        copyto!(B, X)
        M, N = size(X)
        @inbounds for J = 1:N
            nrm = sqrt(π/(2-δ(J-1,0)))
            @simd for I = 1:M
                B[I,J] *= nrm
            end
        end
        for J = 1:N
            mul_col_J!!(Y, pinv, B, J)
        end
        tri_zero_spurious_modes!(At_mul_B!(RP, Y))
    end

    LinearAlgebra.mul!(Y::Matrix, SPc::Adjoint{T,<:SlowTriangularHarmonicPlan}, X::Matrix) where T =
        mul!(Y, transpose(parent(SPc)), X)
end


text:
function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:PPO},
    ::Val{:Pendulum},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_PPO_Pendulum_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)
    rng = StableRNG(seed)
    inner_env = PendulumEnv(T = Float32, rng = rng)
    action_space = action_space(inner_env)
    low = action_space.low
    high = action_space.high
    ns = length(state(inner_env))

    N_ENV = 8
    UPDATE_FREQ = 2048
    env = MultiThreadEnv([
        PendulumEnv(T = Float32, rng = StableRNG(hash(seed + i))) |>
        ActionTransformedEnv(x -> clamp(x * 2, low, high)) for i in 1:N_ENV
    ])

    init = glorot_uniform(rng)

    agent = Agent(
        policy = PPOPolicy(
            approximator = ActorCritic(
                actor = GaussianNetwork(
                    pre = Chain(
                        Dense(ns, 64, relu; initW = glorot_uniform(rng)),
                        Dense(64, 64, relu; initW = glorot_uniform(rng)),
                    ),
                    μ = Chain(Dense(64, 1, tanh; initW = glorot_uniform(rng)), vec),
                    σ = Chain(Dense(64, 1; initW = glorot_uniform(rng)), vec),
                ),
                critic = Chain(
                    Dense(ns, 64, relu; initW = glorot_uniform(rng)),
                    Dense(64, 64, relu; initW = glorot_uniform(rng)),
                    Dense(64, 1; initW = glorot_uniform(rng)),
                ),
                optimizer = ADAM(3e-4),
            ) |> cpu,
            γ = 0.99f0,
            λ = 0.95f0,
            clip_range = 0.2f0,
            max_grad_norm = 0.5f0,
            n_epochs = 10,
            n_microbatches = 32,
            actor_loss_weight = 1.0f0,
            critic_loss_weight = 0.5f0,
            entropy_loss_weight = 0.00f0,
            dist = Normal,
            rng = rng,
            update_freq = UPDATE_FREQ,
        ),
        trajectory = PPOTrajectory(;
            capacity = UPDATE_FREQ,
            state = Matrix{Float32} => (ns, N_ENV),
            action = Vector{Float32} => (N_ENV,),
            action_log_prob = Vector{Float32} => (N_ENV,),
            reward = Vector{Float32} => (N_ENV,),
            terminal = Vector{Bool} => (N_ENV,),
        ),
    )

    stop_condition = StopAfterStep(50_000)
    total_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    hook = ComposedHook(
        total_reward_per_episode,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                @info(
                    "training",
                    actor_loss = agent.policy.actor_loss[end, end],
                    critic_loss = agent.policy.critic_loss[end, end],
                    loss = agent.policy.loss[end, end],
                )
                for i in 1:length(env)
                    if is_terminated(env[i])
                        @info "training" reward = total_reward_per_episode.rewards[i][end] log_step_increment =
                            0
                        break
                    end
                end
            end
        end,
    )

    Experiment(agent, env, stop_condition, hook, "# Play Pendulum with PPO")
end


text:
using LatinHypercubeSampling
using PlotlyJS
using Statistics
using SurrogateModelOptim
using Parameters

options=SurrogateModelOptim.Options(
    iterations=25,
    num_interpolants=20, #Preferably even number of added processes
    num_start_samples=5,
        )

function rosenbrock_2D(x)
    return (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2
end

search_range=[(-5.0,5.0),(-5.0,5.0)]

possible_locs = 20 #Defined the possible design locations from -5,5
dims = [LatinHypercubeSampling.Categorical(possible_locs) for i in 1:length(search_range)]
vals = [Tuple((i for i in range(-5.0,stop=5.0,length=possible_locs))) for i in 1:length(search_range)]

function categorical_smoptimize(f,search_range,dims,vals,options)

    #Print the known minimum (brute force approach)
    possible_designs = collect(Base.Iterators.product(vals...))
    println("Number of possible designs = ", length(possible_designs))
    min_loc = argmin(f.(possible_designs))
    max_loc = argmax(f.(possible_designs))
    println("Minimum and maximum value in categorical design space:")
    println("Min ",f(possible_designs[min_loc]), ", ", possible_designs[min_loc])
    println("Max ",f(possible_designs[max_loc]), ", ", possible_designs[max_loc])

    #Load some option values
    @unpack num_start_samples, sampling_plan_opt_gens,
            iterations, trace = options

    #Create categorical sampling plan with the values in vals
    initialSample = randomLHC(num_start_samples,dims)
    Xind = permutedims(LHCoptim!(initialSample,sampling_plan_opt_gens;dims=dims)[1])
    lhc_plan = Array{Float64,2}(undef,size(Xind))
    for i = 1:size(Xind,1)
        for j = 1:size(Xind,2)
            lhc_plan[i,j] = vals[i][Xind[i,j]]
        end
    end

    #Evaluate sampling plan
    lhc_samples = SurrogateModelOptim.f_opt_eval(f,lhc_plan,trace)

    #Initialize variables to be returned
    sm_interpolant = nothing
    infill_type = Array{Symbol,1}(undef,0)
    infill_prediction = Array{Float64,1}(undef,0)
    optres = nothing
    infill_plan = Array{Float64,2}(undef,size(lhc_plan,1),0)
    infill_sample = Array{Float64,2}(undef,1,0)

    #Run the entire optimization iterations number of times
    for i = 1:iterations
        trace && SurrogateModelOptim.print_iteration(i,iterations)

        #Create the optimized Radial Basis Function interpolant
        samples_all = [lhc_samples infill_sample]
        plan_all = [lhc_plan infill_plan]
        sm_interpolant, optres = SurrogateModelOptim.surrogate_model(plan_all, samples_all; options=options)

        #Choose the closest categorical value when evaluating the function. Note
        # this alters the input of x to match closest value.
        sm_interpolant_cat! = function (x,vals)
            for (i,x_val) in enumerate(x)
                c_ind = SurrogateModelOptim.closest_index(x_val, vals[i])
                x[i] = vals[i][c_ind]
            end
            return sm_interpolant(x)
        end

        #Points to add to the sampling plan to improve the interpolant
        infill_plan_new, infill_type_new, infill_prediction_new, options  =
            model_infill(search_range,plan_all,samples_all,x->sm_interpolant_cat!(x,vals);options=options)

        #Evaluate the new infill points
        infill_sample_new = SurrogateModelOptim.f_opt_eval(f,infill_plan_new,samples_all,trace)

        #Add infill points
        infill_plan = [infill_plan infill_plan_new]
        infill_sample = [infill_sample infill_sample_new]
        infill_type = [infill_type; infill_type_new]
        infill_prediction = [infill_prediction; infill_prediction_new]
    end

    return SurrogateModelOptim.SurrogateResult( lhc_samples, lhc_plan, sm_interpolant,
                            optres, infill_sample, infill_type,
                            infill_plan, infill_prediction,options)
end
result = categorical_smoptimize(rosenbrock_2D,search_range,dims,vals,options)

function plot_fun_2D(fun,sr,title)
    N = 51
    x = range(sr[1][1], stop = sr[1][2], length = N)
    y = range(sr[2][1], stop = sr[2][2], length = N)

    grid(x,y) = [x,y]
    z = @. fun(grid(x,y'))

    trace = surface(x=x,y=y,z=z, colorscale="Viridis")
    layout = Layout(title=title)
    p = plot(trace,layout)
end

# Plot the results
display(plot_fun_2D(rosenbrock_2D,search_range,"Original function"))
display(plot_fun_2D(x->median(result.sm_interpolant(x)),search_range,"Estimated function"))


text:
"""
    estimating_function_template(nobs::Function, ef_contribution::Function)

Define an `estimating_function_template` by supplying:
+ `nobs`: a function of `data` that computes the number of observations of the particular data type,
+ `ef_contribution`: a function of the parameters `theta`, the `data` and the observation index `i` that returns a vector of length `length(theta)`.
"""
struct estimating_function_template
    nobs::Function
    ef_contribution::Function
end

"""
    estimating_function(theta::Vector, data::Any, template::estimating_function_template, br::Bool = false)

Construct the estimating function by adding up all contributions in the `data` according to [`estimating_function_template`](@ref), and evaluate it at `theta`. If `br = true` then automatic differentiation is used to compute the empirical bias-reducing adjustments and add them to the estimating function.
"""
function estimating_function(theta::Vector,
                             data::Any,
                             template::estimating_function_template,
                             br::Bool = false)
    p = length(theta)
    n_obs = template.nobs(data)
    contributions = Matrix(undef, p, n_obs)
    for i in 1:n_obs
        contributions[:, i] = template.ef_contribution(theta, data, i)
    end
    if (br)
        quants = ef_quantities(theta, data, template, br)
        sum(contributions, dims = 2) + quants[1]
    else
        sum(contributions, dims = 2)
    end
end

"""
    get_estimating_function(data::Any, template::estimating_function_template, br::Bool = false)

Construct the estimating function by adding up all contributions in the `data` according to [`estimating_function_template`](@ref). If `br = true` then automatic differentiation is used to compute the empirical bias-reducing adjustments and add them to the estimating function. The result is a function that stores the estimating functions values at its second argument, in a preallocated vector passed as its first argument, ready to be used withing `NLsolve.nlsolve`.
"""
function get_estimating_function(data::Any,
                                 template::estimating_function_template,
                                 br::Bool = false)
    function g!(F, theta::Vector)
        out = estimating_function(theta, data, template, br)
        for i in 1:length(out)
            F[i] = out[i]
        end
    end
end

function ef_quantities(theta::Vector,
                       data::Any,
                       template::estimating_function_template,
                       adjustment::Bool = false)
    nj(eta::Vector, i::Int) = ForwardDiff.jacobian(beta -> template.ef_contribution(beta, data, i), eta)
    p = length(theta)
    n_obs = template.nobs(data)
    psi = Matrix{Float64}(undef, n_obs, p)
    njmats = Vector(undef, n_obs)
    for i in 1:n_obs
        psi[i, :] =  template.ef_contribution(theta, data, i)
        njmats[i] = nj(theta, i)
    end
    jmat_inv = inv(-sum(njmats))
    emat = psi' * psi
    vcov = jmat_inv * (emat * jmat_inv')
    if (adjustment)
        u(eta::Vector, i::Int) = ForwardDiff.jacobian(beta -> nj(beta, i), eta)
        psi_tilde = Matrix(undef, n_obs, p)
        umats = Vector(undef, n_obs)
        for i in 1:n_obs
            umats[i] = u(theta, i)
        end
        umat = sum(umats)
        A = Vector(undef, p)
        for j in 1:p
            for i in 1:n_obs
                psi_tilde[i, :] = njmats[i][j, :]
            end
            A[j] = -tr(jmat_inv * psi_tilde' * psi +
                       vcov * umat[j:p:(p * p - p + j), :] / 2)
        end
        [A, jmat_inv, emat]
    else
        [vcov, jmat_inv, emat]
    end
end