#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.

type NonlinearExprData
    nd::Vector{NodeData}
    const_values::Vector{Float64}
end

include("parsenlp.jl")

typealias NonlinearConstraint GenericRangeConstraint{NonlinearExprData}

type NLPData
    nlobj
    nlconstr::Vector{NonlinearConstraint}
    nlexpr::Vector{NonlinearExprData}
    nlconstrDuals::Vector{Float64}
    nlparamvalues::Vector{Float64}
    evaluator
end

function NonlinearExpression(m::Model,ex::NonlinearExprData)
    initNLP(m)
    nldata::NLPData = m.nlpdata
    push!(nldata.nlexpr, ex)
    return NonlinearExpression(m, length(nldata.nlexpr))
end

function newparameter(m::Model,value::Number)
    initNLP(m)
    nldata::NLPData = m.nlpdata
    push!(nldata.nlparamvalues, value)
    return NonlinearParameter(m, length(nldata.nlparamvalues))
end

getvalue(p::NonlinearParameter) = p.m.nlpdata.nlparamvalues[p.index]::Float64

setvalue(p::NonlinearParameter,v::Number) = (p.m.nlpdata.nlparamvalues[p.index] = v)

NLPData() = NLPData(nothing, NonlinearConstraint[], NonlinearExprData[], Float64[], Float64[], nothing)

Base.copy(::NLPData) = error("Copying nonlinear problems not yet implemented")

function initNLP(m::Model)
    if m.nlpdata === nothing
        m.nlpdata = NLPData()
    end
end

function getdual(c::ConstraintRef{Model,NonlinearConstraint})
    initNLP(c.m)
    nldata::NLPData = c.m.nlpdata
    if length(nldata.nlconstrDuals) != length(nldata.nlconstr)
        error("Dual solution not available. Check that the model was properly solved.")
    end
    return nldata.nlconstrDuals[c.idx]
end

type FunctionStorage
    nd::Vector{NodeData}
    adj::SparseMatrixCSC{Bool,Int}
    const_values::Vector{Float64}
    forward_storage::Vector{Float64}
    partials_storage::Vector{Float64}
    reverse_storage::Vector{Float64}
    grad_sparsity::Vector{Int}
    hess_I::Vector{Int} # nonzero pattern of hessian
    hess_J::Vector{Int}
    rinfo::Coloring.RecoveryInfo # coloring info for hessians
    seed_matrix::Matrix{Float64}
    linearity::Linearity
    dependent_subexpressions::Vector{Int} # subexpressions which this function depends on, ordered for forward pass
end

type SubexpressionStorage
    nd::Vector{NodeData}
    adj::SparseMatrixCSC{Bool,Int}
    const_values::Vector{Float64}
    forward_storage::Vector{Float64}
    partials_storage::Vector{Float64}
    reverse_storage::Vector{Float64}
    forward_storage_ϵ::Vector{Float64}
    partials_storage_ϵ::Vector{Float64}
    reverse_storage_ϵ::Vector{Float64}
    linearity::Linearity
end

type NLPEvaluator <: MathProgBase.AbstractNLPEvaluator
    m::Model
    A::SparseMatrixCSC{Float64,Int} # linear constraint matrix
    parameter_values::Vector{Float64}
    has_nlobj::Bool
    linobj::Vector{Float64}
    objective::FunctionStorage
    constraints::Vector{FunctionStorage}
    subexpressions::Vector{SubexpressionStorage}
    subexpression_order::Vector{Int}
    subexpression_forward_values::Vector{Float64}
    subexpression_reverse_values::Vector{Float64}
    subexpression_linearity::Vector{ReverseDiffSparse.Linearity}
    subexpressions_as_julia_expressions::Vector{Any}
    last_x::Vector{Float64}
    jac_storage::Vector{Float64} # temporary storage for computing jacobians
    user_output_buffer::Vector{Float64} # temporary storage for user-defined functions
    # storage for computing hessians
    # these Float64 vectors are reinterpreted to hold multiple epsilon components
    # so the length should be multiplied by the maximum number of epsilon components
    disable_2ndorder::Bool # don't offer Hess or HessVec
    want_hess::Bool
    forward_storage_ϵ::Vector{Float64} # (longest expression)
    partials_storage_ϵ::Vector{Float64} # (longest expression)
    reverse_storage_ϵ::Vector{Float64} # (longest expression)
    input_ϵ::Vector{Float64} # (number of variables)
    output_ϵ::Vector{Float64}# (number of variables)
    subexpression_forward_values_ϵ::Vector{Float64} # (number of subexpressions)
    subexpression_reverse_values_ϵ::Vector{Float64} # (number of subexpressions)
    # hessian sparsity pattern
    hess_I::Vector{Int}
    hess_J::Vector{Int}
    max_chunk::Int # chunk size for which we've allocated storage
    # timers
    eval_f_timer::Float64
    eval_g_timer::Float64
    eval_grad_f_timer::Float64
    eval_jac_g_timer::Float64
    eval_hesslag_timer::Float64
    function NLPEvaluator(m::Model)
        d = new(m)
        numVar = m.numCols
        d.A = prepConstrMatrix(m)
        d.jac_storage = Array(Float64,numVar)
        d.user_output_buffer = Array(Float64,numVar)
        # check if we have any user-defined operators, in which case we need to
        # disable hessians.
        if isa(m.nlpdata,NLPData)
            nldata::NLPData = m.nlpdata
            has_nlobj = isa(nldata.nlobj, NonlinearExprData)
            has_user_mv_operator = false
            for nlexpr in nldata.nlexpr
                has_user_mv_operator |= ReverseDiffSparse.has_user_multivariate_operators(nlexpr.nd)
            end
            if has_nlobj

                has_user_mv_operator |= ReverseDiffSparse.has_user_multivariate_operators(nldata.nlobj.nd)
            end
            for nlconstr in nldata.nlconstr
                has_user_mv_operator |= ReverseDiffSparse.has_user_multivariate_operators(nlconstr.terms.nd)
            end
            d.disable_2ndorder = has_user_mv_operator
        else
            d.disable_2ndorder = false
        end

        d.eval_f_timer = 0
        d.eval_g_timer = 0
        d.eval_grad_f_timer = 0
        d.eval_jac_g_timer = 0
        d.eval_hesslag_timer = 0
        return d
    end
end

function simplify_expression(nd::Vector{NodeData}, const_values, subexpression_linearity, fixed_variables, parameter_values, x_values, subexpression_values)

    adj = adjmat(nd)
    forward_storage = zeros(length(nd))
    partials_storage = zeros(length(nd))
    linearity = classify_linearity(nd, adj, subexpression_linearity, fixed_variables)
    forward_eval(forward_storage, partials_storage, nd, adj, const_values, parameter_values, x_values, subexpression_values)
    nd_new = simplify_constants(forward_storage, nd, adj, const_values, linearity)
    return nd_new, forward_storage[1], linearity[1]
end

function FunctionStorage(nd::Vector{NodeData}, const_values,numVar, coloring_storage, want_hess::Bool, subexpr::Vector{Vector{NodeData}}, dependent_subexpressions, subexpression_linearity, subexpression_edgelist, subexpression_variables, fixed_variables)

    adj = adjmat(nd)
    forward_storage = zeros(length(nd))
    partials_storage = zeros(length(nd))
    reverse_storage = zeros(length(nd))
    empty!(coloring_storage)
    compute_gradient_sparsity!(coloring_storage, nd)

    for k in dependent_subexpressions
        compute_gradient_sparsity!(coloring_storage,subexpr[k])
    end
    grad_sparsity = sort!(collect(coloring_storage))
    empty!(coloring_storage)

    if want_hess
        # compute hessian sparsity
        linearity = classify_linearity(nd, adj, subexpression_linearity, fixed_variables)
        edgelist = compute_hessian_sparsity(nd, adj, linearity, coloring_storage, subexpression_edgelist, subexpression_variables)
        hess_I, hess_J, rinfo = Coloring.hessian_color_preprocess(edgelist, numVar, coloring_storage)
        seed_matrix = Coloring.seed_matrix(rinfo)
        if linearity[1] == NONLINEAR
            @assert length(hess_I) > 0
        end
    else
        hess_I = hess_J = Int[]
        rinfo = Coloring.RecoveryInfo()
        seed_matrix = Array(Float64,0,0)
        linearity = [NONLINEAR]
    end

    return FunctionStorage(nd, adj, const_values, forward_storage, partials_storage, reverse_storage, grad_sparsity, hess_I, hess_J, rinfo, seed_matrix, linearity[1],dependent_subexpressions)

end

function SubexpressionStorage(nd::Vector{NodeData}, const_values,numVar, fixed_variables,subexpression_linearity)

    adj = adjmat(nd)
    forward_storage = zeros(length(nd))
    partials_storage = zeros(length(nd))
    reverse_storage = zeros(length(nd))
    linearity = classify_linearity(nd, adj, subexpression_linearity, fixed_variables)

    empty_arr = Array(Float64,0)

    return SubexpressionStorage(nd, adj, const_values, forward_storage, partials_storage, reverse_storage, empty_arr, empty_arr, empty_arr, linearity[1])

end

function MathProgBase.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
    for feat in requested_features
        if !(feat in MathProgBase.features_available(d))
            error("Unsupported feature $feat")
            # TODO: implement Jac-vec products
            # for solvers that need them
        end
    end
    if d.eval_f_timer != 0
        # we've already been initialized
        # assume no new features are being requested.
        return
    end

    initNLP(d.m) #in case the problem is purely linear/quadratic thus far
    nldata::NLPData = d.m.nlpdata

    d.constraints = FunctionStorage[]
    d.last_x = fill(NaN, d.m.numCols)

    SIMPLIFY = d.m.simplify_nonlinear_expressions

    fixed_variables = SIMPLIFY ? d.m.colLower .== d.m.colUpper : d.m.colLower .== NaN
    d.m.colVal[fixed_variables] = d.m.colLower[fixed_variables]

    d.parameter_values = nldata.nlparamvalues

    tic()

    d.linobj, linrowlb, linrowub = prepProblemBounds(d.m)
    numVar = length(d.linobj)

    d.want_hess = (:Hess in requested_features)
    want_hess_storage = (:HessVec in requested_features) || d.want_hess
    coloring_storage = ReverseDiffSparse.Coloring.IndexedSet(numVar)

    d.has_nlobj = isa(nldata.nlobj, NonlinearExprData)
    max_expr_length = 0
    main_expressions = Array(Vector{NodeData},0)
    subexpr = Array(Vector{NodeData},0)
    for nlexpr in nldata.nlexpr
        push!(subexpr, nlexpr.nd)
    end
    if d.has_nlobj
        push!(main_expressions,nldata.nlobj.nd)
    end
    for nlconstr in nldata.nlconstr
        push!(main_expressions,nlconstr.terms.nd)
    end
    d.subexpression_order, individual_order = order_subexpressions(main_expressions,subexpr)
    if :ExprGraph in requested_features
        d.subexpressions_as_julia_expressions = Array(Any,length(subexpr))
        for k in d.subexpression_order
            ex = nldata.nlexpr[k]
            adj = adjmat(ex.nd)
            d.subexpressions_as_julia_expressions[k] = tapeToExpr(1, ex.nd, adj, ex.const_values, d.parameter_values, d.subexpressions_as_julia_expressions)
        end
    end

    d.subexpression_linearity = Array(Linearity, length(nldata.nlexpr))
    subexpression_variables = Array(Vector{Int}, length(nldata.nlexpr))
    subexpression_edgelist = Array(Set{Tuple{Int,Int}}, length(nldata.nlexpr))
    d.subexpressions = Array(SubexpressionStorage, length(nldata.nlexpr))
    d.subexpression_forward_values = Array(Float64, length(d.subexpressions))
    d.subexpression_reverse_values = Array(Float64, length(d.subexpressions))

    empty_edgelist = Set{Tuple{Int,Int}}()
    for k in d.subexpression_order # only load expressions which actually are used
        if SIMPLIFY
            nd_new, forward_value, simplified_linearity = simplify_expression(nldata.nlexpr[k].nd, nldata.nlexpr[k].const_values,d.subexpression_linearity, fixed_variables, d.parameter_values, d.m.colVal, d.subexpression_forward_values)
        else
            nd_new = nldata.nlexpr[k].nd
            forward_value = NaN
            simplified_linearity = NONLINEAR
        end
        d.subexpression_forward_values[k] = forward_value
        if simplified_linearity != CONSTANT
            d.subexpressions[k] = SubexpressionStorage(nd_new, nldata.nlexpr[k].const_values, numVar, fixed_variables, d.subexpression_linearity)
            subex = d.subexpressions[k]
            d.subexpression_linearity[k] = subex.linearity
            @assert !SIMPLIFY || subex.linearity != CONSTANT
            if d.want_hess
                empty!(coloring_storage)
                compute_gradient_sparsity!(coloring_storage,subex.nd)
                # union with all dependent expressions
                for idx in list_subexpressions(subex.nd)
                    union!(coloring_storage, subexpression_variables[idx])
                end
                subexpression_variables[k] = collect(coloring_storage)
                empty!(coloring_storage)
                linearity = classify_linearity(subex.nd, subex.adj, d.subexpression_linearity, fixed_variables)
                edgelist = compute_hessian_sparsity(subex.nd, subex.adj, linearity,coloring_storage,subexpression_edgelist, subexpression_variables)
                subexpression_edgelist[k] = edgelist
            end
        else
            d.subexpression_linearity[k] = simplified_linearity
            subexpression_edgelist[k] = empty_edgelist
        end

    end

    if SIMPLIFY
        main_expressions = Array(Vector{NodeData},0)

        # simplify objective and constraint expressions
        if d.has_nlobj
            nd_new, forward_value = simplify_expression(nldata.nlobj.nd, nldata.nlobj.const_values,d.subexpression_linearity, fixed_variables, d.parameter_values, d.m.colVal, d.subexpression_forward_values)
            push!(main_expressions,nd_new)
        end
        for k in 1:length(nldata.nlconstr)
            nd_new, forward_value = simplify_expression(nldata.nlconstr[k].terms.nd, nldata.nlconstr[k].terms.const_values,d.subexpression_linearity, fixed_variables, d.parameter_values, d.m.colVal, d.subexpression_forward_values)
            push!(main_expressions,nd_new)
        end
        # recompute dependencies after simplification
        d.subexpression_order, individual_order = order_subexpressions(main_expressions,subexpr)

        subexpr = Array(Vector{NodeData},length(d.subexpressions))
        for k in d.subexpression_order
            subexpr[k] = d.subexpressions[k].nd
        end
    end

    max_chunk = 1

    if d.has_nlobj
        @assert length(d.m.obj.qvars1) == 0 && length(d.m.obj.aff.vars) == 0
        nd = main_expressions[1]
        d.objective = FunctionStorage(nd, nldata.nlobj.const_values, numVar, coloring_storage, d.want_hess, subexpr, individual_order[1], d.subexpression_linearity, subexpression_edgelist, subexpression_variables, fixed_variables)
        max_expr_length = max(max_expr_length, length(d.objective.nd))
        max_chunk = max(max_chunk, size(d.objective.seed_matrix,2))
    end

    for k in 1:length(nldata.nlconstr)
        nlconstr = nldata.nlconstr[k]
        idx = (d.has_nlobj) ? k+1 : k
        nd = main_expressions[idx]
        push!(d.constraints, FunctionStorage(nd, nlconstr.terms.const_values, numVar, coloring_storage, d.want_hess, subexpr, individual_order[idx], d.subexpression_linearity, subexpression_edgelist, subexpression_variables, fixed_variables))
        max_expr_length = max(max_expr_length, length(d.constraints[end].nd))
        max_chunk = max(max_chunk, size(d.constraints[end].seed_matrix,2))
    end

    max_chunk = min(max_chunk, 10) # 10 is hardcoded upper bound to avoid excess memory allocation

    if d.want_hess || want_hess_storage # storage for Hess or HessVec
        d.input_ϵ = Array(Float64,max_chunk*d.m.numCols)
        d.output_ϵ = Array(Float64,max_chunk*d.m.numCols)
        d.forward_storage_ϵ = Array(Float64,max_chunk*max_expr_length)
        d.partials_storage_ϵ = Array(Float64,max_chunk*max_expr_length)
        d.reverse_storage_ϵ = Array(Float64,max_chunk*max_expr_length)
        d.subexpression_forward_values_ϵ = Array(Float64,max_chunk*length(d.subexpressions))
        d.subexpression_reverse_values_ϵ = Array(Float64,max_chunk*length(d.subexpressions))
        for k in d.subexpression_order
            subex = d.subexpressions[k]
            subex.forward_storage_ϵ = zeros(Float64,max_chunk*length(subex.nd))
            subex.partials_storage_ϵ = zeros(Float64,max_chunk*length(subex.nd))
            subex.reverse_storage_ϵ = zeros(Float64,max_chunk*length(subex.nd))
        end
        d.max_chunk = max_chunk
        if d.want_hess
            d.hess_I, d.hess_J = _hesslag_structure(d)
            # JIT warm-up
            MathProgBase.eval_hesslag(d, Array(Float64,length(d.hess_I)), d.m.colVal, 1.0, ones(MathProgBase.numconstr(d.m)))
        end
    end

    # JIT warm-up
    if :Grad in requested_features
        MathProgBase.eval_grad_f(d, zeros(numVar), d.m.colVal)
        MathProgBase.eval_g(d, zeros(MathProgBase.numconstr(d.m)), d.m.colVal)
    end

    tprep = toq()
    #println("Prep time: $tprep")

    # reset timers
    d.eval_f_timer = 0
    d.eval_grad_f_timer = 0
    d.eval_g_timer = 0
    d.eval_jac_g_timer = 0
    d.eval_hesslag_timer = 0

    nothing
end

function MathProgBase.features_available(d::NLPEvaluator)
    features = [:Grad, :Jac, :ExprGraph]
    if !d.disable_2ndorder
        push!(features,:Hess)
        push!(features,:HessVec)
    end
    return features
end

function forward_eval_all(d::NLPEvaluator,x)
    # do a forward pass on all expressions at x
    subexpr_values = d.subexpression_forward_values
    user_input_buffer = d.jac_storage
    user_output_buffer = d.user_output_buffer
    SIMPLIFY = d.m.simplify_nonlinear_expressions
    for k in d.subexpression_order
        if SIMPLIFY && d.subexpression_linearity[k] == CONSTANT
            continue
        end
        ex = d.subexpressions[k]
        subexpr_values[k] = forward_eval(ex.forward_storage,ex.partials_storage,ex.nd,ex.adj,ex.const_values,d.parameter_values,x,subexpr_values,user_input_buffer,user_output_buffer)
    end
    if d.has_nlobj
        ex = d.objective
        forward_eval(ex.forward_storage,ex.partials_storage,ex.nd,ex.adj,ex.const_values,d.parameter_values,x,subexpr_values,user_input_buffer,user_output_buffer)
    end
    for ex in d.constraints
        forward_eval(ex.forward_storage,ex.partials_storage,ex.nd,ex.adj,ex.const_values,d.parameter_values,x,subexpr_values,user_input_buffer,user_output_buffer)
    end
end

function reverse_eval_all(d::NLPEvaluator,x)
    # do a reverse pass on all expressions at x
    subexpr_reverse_values = d.subexpression_reverse_values
    subexpr_values = d.subexpression_forward_values
    grad_storage = d.jac_storage
    SIMPLIFY = d.m.simplify_nonlinear_expressions
    for k in d.subexpression_order
        if SIMPLIFY && d.subexpression_linearity[k] == CONSTANT
            continue
        end
        ex = d.subexpressions[k]
        reverse_eval(ex.reverse_storage,ex.partials_storage,ex.nd,ex.adj)
    end
    if d.has_nlobj
        ex = d.objective
        reverse_eval(ex.reverse_storage,ex.partials_storage,ex.nd,ex.adj)
    end
    for ex in d.constraints
        reverse_eval(ex.reverse_storage,ex.partials_storage,ex.nd,ex.adj)
    end
    copy!(d.last_x,x)
end

function MathProgBase.eval_f(d::NLPEvaluator, x)
    tic()
    if d.last_x != x
        forward_eval_all(d,x)
        reverse_eval_all(d,x)
    end
    val = zero(eltype(x))
    if d.has_nlobj
        val = d.objective.forward_storage[1]
    else
        qobj = d.m.obj::QuadExpr
        val = dot(x,d.linobj) + qobj.aff.constant
        for k in 1:length(qobj.qvars1)
            val += qobj.qcoeffs[k]*x[qobj.qvars1[k].col]*x[qobj.qvars2[k].col]
        end
    end
    d.eval_f_timer += toq()
    return val
end

function MathProgBase.eval_grad_f(d::NLPEvaluator, g, x)
    tic()
    if d.last_x != x
        forward_eval_all(d,x)
        reverse_eval_all(d,x)
    end
    SIMPLIFY = d.m.simplify_nonlinear_expressions
    if d.has_nlobj
        fill!(g,0.0)
        ex = d.objective
        subexpr_reverse_values = d.subexpression_reverse_values
        subexpr_reverse_values[ex.dependent_subexpressions] = 0.0
        reverse_extract(g,ex.reverse_storage,ex.nd,ex.adj,subexpr_reverse_values,1.0)
        for i in length(ex.dependent_subexpressions):-1:1
            k = ex.dependent_subexpressions[i]
            if SIMPLIFY && d.subexpression_linearity[k] == CONSTANT
                continue
            end
            subexpr = d.subexpressions[k]
            reverse_extract(g,subexpr.reverse_storage,subexpr.nd,subexpr.adj,subexpr_reverse_values,subexpr_reverse_values[k])

        end
    else
        copy!(g,d.linobj)
        qobj::QuadExpr = d.m.obj
        for k in 1:length(qobj.qvars1)
            coef = qobj.qcoeffs[k]
            g[qobj.qvars1[k].col] += coef*x[qobj.qvars2[k].col]
            g[qobj.qvars2[k].col] += coef*x[qobj.qvars1[k].col]
        end
    end
    d.eval_grad_f_timer += toq()
    return
end

function MathProgBase.eval_g(d::NLPEvaluator, g, x)
    tic()
    if d.last_x != x
        forward_eval_all(d,x)
        reverse_eval_all(d,x)
    end
    A = d.A
    for i in 1:size(A,1); g[i] = 0.0; end
    #fill!(view(g,1:size(A,1)), 0.0)
    A_mul_B!(view(g,1:size(A,1)),A,x)
    idx = size(A,1)+1
    quadconstr = d.m.quadconstr::Vector{QuadConstraint}
    for c::QuadConstraint in quadconstr
        aff = c.terms.aff
        v = aff.constant
        for k in 1:length(aff.vars)
            v += aff.coeffs[k]*x[aff.vars[k].col]
        end
        for k in 1:length(c.terms.qvars1)
            v += c.terms.qcoeffs[k]*x[c.terms.qvars1[k].col]*x[c.terms.qvars2[k].col]
        end
        g[idx] = v
        idx += 1
    end
    for ex in d.constraints
        g[idx] = ex.forward_storage[1]
        idx += 1
    end

    d.eval_g_timer += toq()
    #print("x = ");show(x);println()
    #println(size(A,1), " g(x) = ");show(g);println()
    return
end

function MathProgBase.eval_jac_g(d::NLPEvaluator, J, x)
    tic()
    if d.last_x != x
        forward_eval_all(d,x)
        reverse_eval_all(d,x)
    end
    fill!(J,0.0)
    idx = 1
    A = d.A
    for col = 1:size(A,2)
        for pos = nzrange(A,col)
            J[idx] = A.nzval[pos]
            idx += 1
        end
    end
    quadconstr = d.m.quadconstr::Vector{QuadConstraint}
    for c::QuadConstraint in quadconstr
        aff = c.terms.aff
        for k in 1:length(aff.vars)
            J[idx] = aff.coeffs[k]
            idx += 1
        end
        for k in 1:length(c.terms.qvars1)
            coef = c.terms.qcoeffs[k]
            qidx1 = c.terms.qvars1[k].col
            qidx2 = c.terms.qvars2[k].col

            J[idx] = coef*x[qidx2]
            J[idx+1] = coef*x[qidx1]
            idx += 2
        end
    end
    grad_storage = d.jac_storage
    subexpr_reverse_values = d.subexpression_reverse_values
    SIMPLIFY = d.m.simplify_nonlinear_expressions
    for ex in d.constraints
        nzidx = ex.grad_sparsity
        grad_storage[nzidx] = 0.0
        subexpr_reverse_values[ex.dependent_subexpressions] = 0.0

        reverse_extract(grad_storage,ex.reverse_storage,ex.nd,ex.adj,subexpr_reverse_values,1.0)
        for i in length(ex.dependent_subexpressions):-1:1
            k = ex.dependent_subexpressions[i]
            if SIMPLIFY && d.subexpression_linearity[k] == CONSTANT
                continue
            end
            subexpr = d.subexpressions[k]
            reverse_extract(grad_storage,subexpr.reverse_storage,subexpr.nd,subexpr.adj,subexpr_reverse_values,subexpr_reverse_values[k])
        end

        for k in 1:length(nzidx)
            J[idx+k-1] = grad_storage[nzidx[k]]
        end
        idx += length(nzidx)
    end

    d.eval_jac_g_timer += toq()
    #print("x = ");show(x);println()
    #print("V ");show(J);println()
    return
end

function MathProgBase.eval_hesslag_prod(
    d::NLPEvaluator,
    h::Vector{Float64}, # output vector
    x::Vector{Float64}, # current solution
    v::Vector{Float64}, # rhs vector
    σ::Float64,         # multiplier for objective
    μ::Vector{Float64}) # multipliers for each constraint

    nldata = d.m.nlpdata::NLPData

    if d.last_x != x
        forward_eval_all(d,x)
        reverse_eval_all(d,x)
    end

    fill!(h, 0.0)

    # quadratic objective
    qobj::QuadExpr = d.m.obj
    for k in 1:length(qobj.qvars1)
        col1 = qobj.qvars1[k].col
        col2 = qobj.qvars2[k].col
        coef = qobj.qcoeffs[k]
        if col1 == col2
            h[col1] += σ*2*coef*v[col1]
        else
            h[col1] += σ*coef*v[col2]
            h[col2] += σ*coef*v[col1]
        end
    end

    # quadratic constraints
    row = size(d.A,1)+1
    quadconstr = d.m.quadconstr::Vector{QuadConstraint}
    for c in quadconstr
        l = μ[row]
        for k in 1:length(c.terms.qvars1)
            col1 = c.terms.qvars1[k].col
            col2 = c.terms.qvars2[k].col
            coef = c.terms.qcoeffs[k]
            if col1 == col2
                h[col1] += l*2*coef*v[col1]
            else
                h[col1] += l*coef*v[col2]
                h[col2] += l*coef*v[col1]
            end
        end
        row += 1
    end

    input_ϵ = reinterpret(ForwardDiff.Partials{1,Float64}, d.input_ϵ)
    output_ϵ = reinterpret(ForwardDiff.Partials{1,Float64}, d.output_ϵ)
    for i in 1:length(x)
        input_ϵ[i] = ForwardDiff.Partials((v[i],))
    end

    # forward evaluate all subexpressions once
    subexpr_forward_values_ϵ = reinterpret(ForwardDiff.Partials{1,Float64},d.subexpression_forward_values_ϵ)
    subexpr_reverse_values_ϵ = reinterpret(ForwardDiff.Partials{1,Float64},d.subexpression_reverse_values_ϵ)
    forward_storage_ϵ = reinterpret(ForwardDiff.Partials{1,Float64},d.forward_storage_ϵ)
    reverse_storage_ϵ = reinterpret(ForwardDiff.Partials{1,Float64},d.reverse_storage_ϵ)
    partials_storage_ϵ = reinterpret(ForwardDiff.Partials{1,Float64},d.partials_storage_ϵ)
    SIMPLIFY = d.m.simplify_nonlinear_expressions
    for expridx in d.subexpression_order
        if SIMPLIFY && d.subexpression_linearity[expridx] == CONSTANT
            continue
        end
        subexpr = d.subexpressions[expridx]
        sub_forward_storage_ϵ = reinterpret(ForwardDiff.Partials{1,Float64},subexpr.forward_storage_ϵ)
        sub_partials_storage_ϵ = reinterpret(ForwardDiff.Partials{1,Float64},subexpr.partials_storage_ϵ)
        subexpr_forward_values_ϵ[expridx] = forward_eval_ϵ(subexpr.forward_storage,sub_forward_storage_ϵ, subexpr.partials_storage, sub_partials_storage_ϵ, subexpr.nd, subexpr.adj, input_ϵ, subexpr_forward_values_ϵ)
    end
    # we only need to do one reverse pass through the subexpressions as well
    zero_ϵ = zero(ForwardDiff.Partials{1,Float64})
    fill!(subexpr_reverse_values_ϵ,zero_ϵ)
    fill!(d.subexpression_reverse_values,0.0)
    fill!(reverse_storage_ϵ,zero_ϵ)
    fill!(output_ϵ,zero_ϵ)
    if d.has_nlobj
        ex = d.objective
        forward_eval_ϵ(ex.forward_storage, forward_storage_ϵ, ex.partials_storage, partials_storage_ϵ, ex.nd,ex.adj,input_ϵ, subexpr_forward_values_ϵ)
        reverse_eval_ϵ(output_ϵ,ex.reverse_storage, reverse_storage_ϵ,ex.partials_storage, partials_storage_ϵ,ex.nd,ex.adj,d.subexpression_reverse_values,subexpr_reverse_values_ϵ, σ, zero_ϵ)
    end

    for i in 1:length(d.constraints)
        ex = d.constraints[i]
        l = μ[row]
        forward_eval_ϵ(ex.forward_storage, forward_storage_ϵ, ex.partials_storage, partials_storage_ϵ, ex.nd,ex.adj, input_ϵ,subexpr_forward_values_ϵ)
        reverse_eval_ϵ(output_ϵ, ex.reverse_storage, reverse_storage_ϵ, ex.partials_storage, partials_storage_ϵ, ex.nd,ex.adj,d.subexpression_reverse_values,subexpr_reverse_values_ϵ, l, zero_ϵ)
        row += 1
    end

    for i in length(ex.dependent_subexpressions):-1:1
        expridx = ex.dependent_subexpressions[i]
        if SIMPLIFY && d.subexpression_linearity[expridx] == CONSTANT
            continue
        end
        subexpr = d.subexpressions[expridx]
        sub_reverse_storage_ϵ = reinterpret(ForwardDiff.Partials{1,Float64},subexpr.reverse_storage_ϵ)
        sub_partials_storage_ϵ = reinterpret(ForwardDiff.Partials{1,Float64},subexpr.partials_storage_ϵ)
        reverse_eval_ϵ(output_ϵ,subexpr.reverse_storage,sub_reverse_storage_ϵ, subexpr.partials_storage, sub_partials_storage_ϵ,subexpr.nd,subexpr.adj,d.subexpression_reverse_values,subexpr_reverse_values_ϵ,d.subexpression_reverse_values[expridx],subexpr_reverse_values_ϵ[expridx])
    end

    for i in 1:length(x)
        h[i] += output_ϵ[i].values[1]
    end

end

function MathProgBase.eval_hesslag(
    d::NLPEvaluator,
    H::Vector{Float64},         # Sparse hessian entry vector
    x::Vector{Float64},         # Current solution
    obj_factor::Float64,        # Lagrangian multiplier for objective
    lambda::Vector{Float64})    # Multipliers for each constraint

    qobj = d.m.obj::QuadExpr
    nldata = d.m.nlpdata::NLPData

    d.want_hess || error("Hessian computations were not requested on the call to MathProgBase.initialize.")

    if d.last_x != x
        forward_eval_all(d,x)
        reverse_eval_all(d,x)
    end

    tic()

    # quadratic objective
    nzcount = 1
    for k in 1:length(qobj.qvars1)
        if qobj.qvars1[k].col == qobj.qvars2[k].col
            H[nzcount] = obj_factor*2*qobj.qcoeffs[k]
        else
            H[nzcount] = obj_factor*qobj.qcoeffs[k]
        end
        nzcount += 1
    end
    # quadratic constraints
    quadconstr = d.m.quadconstr::Vector{QuadConstraint}
    for i in 1:length(quadconstr)
        c = quadconstr[i]
        l = lambda[length(d.m.linconstr)+i]
        for k in 1:length(c.terms.qvars1)
            if c.terms.qvars1[k].col == c.terms.qvars2[k].col
                H[nzcount] = l*2*c.terms.qcoeffs[k]
            else
                H[nzcount] = l*c.terms.qcoeffs[k]
            end
            nzcount += 1
        end
    end

    fill!(d.input_ϵ,0.0)
    recovery_tmp_storage = d.output_ϵ
    nzcount -= 1

    if d.has_nlobj
        ex = d.objective
        chunk = min(size(ex.seed_matrix,2),d.max_chunk)
        if chunk == 1
            # skip dynamic dispatch
            nzthis = hessian_slice(d, ex, x, H, obj_factor, nzcount, recovery_tmp_storage, Val{1})::Int
        else
            nzthis = hessian_slice(d, ex, x, H, obj_factor, nzcount, recovery_tmp_storage, Val{chunk})::Int
        end
        nzcount += nzthis
    end

    for i in 1:length(d.constraints)
        ex = d.constraints[i]
        chunk = min(size(ex.seed_matrix,2),d.max_chunk)
        if chunk == 1
            nzthis = hessian_slice(d, ex, x, H, lambda[i+length(quadconstr)+length(d.m.linconstr)], nzcount, recovery_tmp_storage, Val{1})::Int
        else
            nzthis = hessian_slice(d, ex, x, H, lambda[i+length(quadconstr)+length(d.m.linconstr)], nzcount, recovery_tmp_storage, Val{chunk})::Int
        end
        nzcount += nzthis
    end

    d.eval_hesslag_timer += toq()
    return

end

function hessian_slice_inner{CHUNK}(d, ex, R, input_ϵ, output_ϵ, ::Type{Val{CHUNK}})

    subexpr_forward_values_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64},d.subexpression_forward_values_ϵ)
    subexpr_reverse_values_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64},d.subexpression_reverse_values_ϵ)
    forward_storage_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64},d.forward_storage_ϵ)
    reverse_storage_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64},d.reverse_storage_ϵ)
    partials_storage_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64},d.partials_storage_ϵ)
    zero_ϵ = zero(ForwardDiff.Partials{CHUNK,Float64})

    SIMPLIFY = d.m.simplify_nonlinear_expressions
    # do a forward pass
    for expridx in ex.dependent_subexpressions
        if SIMPLIFY && d.subexpression_linearity[expridx] == CONSTANT
            continue
        end
        subexpr = d.subexpressions[expridx]
        sub_forward_storage_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64},subexpr.forward_storage_ϵ)
        sub_partials_storage_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64},subexpr.partials_storage_ϵ)
        subexpr_forward_values_ϵ[expridx] = forward_eval_ϵ(subexpr.forward_storage,sub_forward_storage_ϵ,subexpr.partials_storage,sub_partials_storage_ϵ, subexpr.nd, subexpr.adj, input_ϵ, subexpr_forward_values_ϵ)
    end
    forward_eval_ϵ(ex.forward_storage,forward_storage_ϵ,ex.partials_storage, partials_storage_ϵ,ex.nd,ex.adj,input_ϵ, subexpr_forward_values_ϵ)

    # do a reverse pass
    subexpr_reverse_values_ϵ[ex.dependent_subexpressions] = zero_ϵ
    d.subexpression_reverse_values[ex.dependent_subexpressions] = 0.0

    reverse_eval_ϵ(output_ϵ, ex.reverse_storage, reverse_storage_ϵ,ex.partials_storage, partials_storage_ϵ,ex.nd,ex.adj,d.subexpression_reverse_values,subexpr_reverse_values_ϵ, 1.0, zero_ϵ)
    for i in length(ex.dependent_subexpressions):-1:1
        expridx = ex.dependent_subexpressions[i]
        if SIMPLIFY && d.subexpression_linearity[expridx] == CONSTANT
            continue
        end
        subexpr = d.subexpressions[expridx]
        sub_reverse_storage_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64},subexpr.reverse_storage_ϵ)
        sub_partials_storage_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64},subexpr.partials_storage_ϵ)
        reverse_eval_ϵ(output_ϵ, subexpr.reverse_storage, sub_reverse_storage_ϵ,subexpr.partials_storage,sub_partials_storage_ϵ,subexpr.nd,subexpr.adj,d.subexpression_reverse_values,subexpr_reverse_values_ϵ,d.subexpression_reverse_values[expridx],subexpr_reverse_values_ϵ[expridx])
    end
end

function hessian_slice{CHUNK}(d, ex, x, H, scale, nzcount, recovery_tmp_storage,::Type{Val{CHUNK}})

    nzthis = length(ex.hess_I)
    if ex.linearity == LINEAR
        @assert nzthis == 0
        return 0
    end
    R = ex.seed_matrix
    Coloring.prepare_seed_matrix!(R,ex.rinfo)
    local_to_global_idx = ex.rinfo.local_indices

    zero_ϵ = zero(ForwardDiff.Partials{CHUNK,Float64})

    input_ϵ_raw = d.input_ϵ
    output_ϵ_raw = d.output_ϵ
    input_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64}, input_ϵ_raw)
    output_ϵ = reinterpret_unsafe(ForwardDiff.Partials{CHUNK,Float64}, output_ϵ_raw)

    # compute hessian-vector products
    num_products = size(R,2) # number of hessian-vector products
    num_chunks = div(num_products, CHUNK)
    @assert size(R,1) == length(local_to_global_idx)
    numVar = length(x)

    for k in 1:CHUNK:CHUNK*num_chunks

        for r in 1:length(local_to_global_idx)
            # set up directional derivatives
            @inbounds idx = local_to_global_idx[r]
            # load up R[r,k,k+1,...,k+CHUNK-1] into input_ϵ
            for s in 1:CHUNK
                input_ϵ_raw[(idx-1)*CHUNK + s] = R[r,k+s-1]
            end
            @inbounds output_ϵ[idx] = zero_ϵ
        end

        hessian_slice_inner(d, ex, R, input_ϵ, output_ϵ, Val{CHUNK})

        # collect directional derivatives
        for r in 1:length(local_to_global_idx)
            idx = local_to_global_idx[r]
            # load output_ϵ into R[r,k,k+1,...,k+CHUNK-1]
            for s in 1:CHUNK
                R[r,k+s-1] = output_ϵ_raw[(idx-1)*CHUNK + s]
            end
            @inbounds input_ϵ[idx] = zero_ϵ
        end

    end

    # leftover chunk
    remaining = num_products - CHUNK*num_chunks
    if remaining > 0
        k = CHUNK*num_chunks+1
        for r in 1:length(local_to_global_idx)
            # set up directional derivatives
            @inbounds idx = local_to_global_idx[r]
            # load up R[r,k,k+1,...,k+remaining-1] into input_ϵ
            for s in 1:remaining
                # leave junk in the unused components
                input_ϵ_raw[(idx-1)*CHUNK + s] = R[r,k+s-1]
            end
            @inbounds output_ϵ[idx] = zero_ϵ
        end

        hessian_slice_inner(d, ex, R, input_ϵ, output_ϵ, Val{CHUNK})

        # collect directional derivatives
        for r in 1:length(local_to_global_idx)
            idx = local_to_global_idx[r]
            # load output_ϵ into R[r,k,k+1,...,k+remaining-1]
            for s in 1:remaining
                R[r,k+s-1] = output_ϵ_raw[(idx-1)*CHUNK + s]
            end
            @inbounds input_ϵ[idx] = zero_ϵ
        end
    end

    # Output is in R, now recover

    #output_slice = view(H, (nzcount+1):(nzcount+nzthis))
    output_slice = VectorView(nzcount, nzthis, pointer(H))
    Coloring.recover_from_matmat!(output_slice, R, ex.rinfo, recovery_tmp_storage)
    scale!(output_slice, scale)
    return nzthis

end

MathProgBase.isobjlinear(d::NLPEvaluator) = !d.has_nlobj && (length(d.m.obj.qvars1) == 0)
# interpret quadratic to include purely linear
MathProgBase.isobjquadratic(d::NLPEvaluator) = !d.has_nlobj

MathProgBase.isconstrlinear(d::NLPEvaluator, i::Integer) = (i <= length(d.m.linconstr))

function MathProgBase.jac_structure(d::NLPEvaluator)
    # Jacobian structure
    jac_I = Int[]
    jac_J = Int[]
    A = d.A
    for col = 1:size(A,2)
        for pos = nzrange(A,col)
            push!(jac_I, A.rowval[pos])
            push!(jac_J, col)
        end
    end
    rowoffset = size(A,1)+1
    for c::QuadConstraint in d.m.quadconstr
        aff = c.terms.aff
        for k in 1:length(aff.vars)
            push!(jac_I, rowoffset)
            push!(jac_J, aff.vars[k].col)
        end
        for k in 1:length(c.terms.qvars1)
            push!(jac_I, rowoffset)
            push!(jac_I, rowoffset)
            push!(jac_J, c.terms.qvars1[k].col)
            push!(jac_J, c.terms.qvars2[k].col)
        end
        rowoffset += 1
    end
    for ex in d.constraints
        idx = ex.grad_sparsity
        for i in 1:length(idx)
            push!(jac_I, rowoffset)
            push!(jac_J, idx[i])
        end
        rowoffset += 1
    end
    return jac_I, jac_J
end
function MathProgBase.hesslag_structure(d::NLPEvaluator)
    d.want_hess || error("Hessian computations were not requested on the call to MathProgBase.initialize.")
    return d.hess_I,d.hess_J
end
function _hesslag_structure(d::NLPEvaluator)
    hess_I = Int[]
    hess_J = Int[]

    qobj::QuadExpr = d.m.obj
    for k in 1:length(qobj.qvars1)
        qidx1 = qobj.qvars1[k].col
        qidx2 = qobj.qvars2[k].col
        if qidx2 > qidx1
            qidx1, qidx2 = qidx2, qidx1
        end
        push!(hess_I, qidx1)
        push!(hess_J, qidx2)
    end
    # quadratic constraints
    for c::QuadConstraint in d.m.quadconstr
        for k in 1:length(c.terms.qvars1)
            qidx1 = c.terms.qvars1[k].col
            qidx2 = c.terms.qvars2[k].col
            if qidx2 > qidx1
                qidx1, qidx2 = qidx2, qidx1
            end
            push!(hess_I, qidx1)
            push!(hess_J, qidx2)
        end
    end

    if d.has_nlobj
        append!(hess_I, d.objective.hess_I)
        append!(hess_J, d.objective.hess_J)
    end
    for ex in d.constraints
        append!(hess_I, ex.hess_I)
        append!(hess_J, ex.hess_J)
    end

    return hess_I, hess_J
end

# currently don't merge duplicates (this isn't required by MPB standard)
function affToExpr(aff::AffExpr, constant::Bool)
    if length(aff.vars) == 0 && !constant
        return 0
    end
    ex = Expr(:call,:+)
    for k in 1:length(aff.vars)
        push!(ex.args, Expr(:call,:*,aff.coeffs[k],:(x[$(aff.vars[k].col)])))
    end
    if constant && aff.constant != 0
        push!(ex.args, aff.constant)
    end
    return ex
end

function quadToExpr(q::QuadExpr,constant::Bool)
    ex = Expr(:call,:+)
    for k in 1:length(q.qvars1)
        push!(ex.args, Expr(:call,:*,q.qcoeffs[k],:(x[$(q.qvars1[k].col)]), :(x[$(q.qvars2[k].col)])))
    end
    append!(ex.args, affToExpr(q.aff,constant).args[2:end])
    return ex
end

# we splat in the subexpressions (for now)
function tapeToExpr(k, nd::Vector{NodeData}, adj, const_values, parameter_values, subexpressions::Vector{Any})

    children_arr = rowvals(adj)

    nod = nd[k]
    if nod.nodetype == VARIABLE
        return Expr(:ref,:x,nod.index)
    elseif nod.nodetype == VALUE
        return const_values[nod.index]
    elseif nod.nodetype == SUBEXPRESSION
        return subexpressions[nod.index]
    elseif nod.nodetype == PARAMETER
        return parameter_values[nod.index]
    elseif nod.nodetype == CALL
        op = nod.index
        opsymbol = operators[op]
        children_idx = nzrange(adj,k)
        if opsymbol == :+ && length(children_idx) == 0
            return 0
        elseif opsymbol == :* && length(children_idx) == 0
            return 1
        end
        ex = Expr(:call,opsymbol)
        for cidx in children_idx
            push!(ex.args, tapeToExpr(children_arr[cidx], nd, adj, const_values, parameter_values, subexpressions))
        end
        return ex
    elseif nod.nodetype == CALLUNIVAR
        op = nod.index
        opsymbol = univariate_operators[op]
        cidx = first(nzrange(adj,k))
        return Expr(:call,opsymbol,tapeToExpr(children_arr[cidx], nd, adj, const_values, parameter_values, subexpressions))
    elseif nod.nodetype == COMPARISON
        op = nod.index
        opsymbol = comparison_operators[op]
        children_idx = nzrange(adj,k)
        if length(children_idx) > 2 || VERSION < v"0.5.0-dev+3231"
            ex = Expr(:comparison)
            for cidx in children_idx
                push!(ex.args, tapeToExpr(children_arr[cidx], nd, adj, const_values, parameter_values, subexpressions))
                push!(ex.args, opsymbol)
            end
            pop!(ex.args)
        else
            ex = Expr(:call, opsymbol)
            push!(ex.args, tapeToExpr(children_arr[children_idx[1]], nd, adj, const_values, parameter_values, subexpressions))
            push!(ex.args, tapeToExpr(children_arr[children_idx[2]], nd, adj, const_values, parameter_values, subexpressions))
        end
        return ex
    elseif nod.nodetype == LOGIC
        op = nod.index
        opsymbol = logic_operators[op]
        children_idx = nzrange(adj,k)
        lhs = tapeToExpr(children_arr[first(children_idx)], nd, adj, const_values, parameter_values, subexpressions)
        rhs = tapeToExpr(children_arr[last(children_idx)], nd, adj, const_values, parameter_values, subexpressions)
        return Expr(opsymbol, lhs, rhs)
    end
    error()

end

function MathProgBase.obj_expr(d::NLPEvaluator)
    if d.has_nlobj
        # for now, don't pass simplified expressions
        ex = d.m.nlpdata.nlobj
        adj = adjmat(ex.nd)
        return tapeToExpr(1, ex.nd, adj, ex.const_values, d.parameter_values, d.subexpressions_as_julia_expressions)
    else
        return quadToExpr(d.m.obj, true)
    end
end

function MathProgBase.constr_expr(d::NLPEvaluator,i::Integer)
    nlin = length(d.m.linconstr)
    nquad = length(d.m.quadconstr)
    if i <= nlin
        constr = d.m.linconstr[i]
        ex = affToExpr(constr.terms, false)
        if sense(constr) == :range
            return Expr(:comparison, constr.lb, :(<=), ex, :(<=), constr.ub)
        else
            if VERSION >= v"0.5.0-dev+3231"
                return Expr(:call, sense(constr), ex, rhs(constr))
            else
                return Expr(:comparison, ex, sense(constr), rhs(constr))
            end
        end
    elseif i > nlin && i <= nlin + nquad
        i -= nlin
        qconstr = d.m.quadconstr[i]
        if VERSION >= v"0.5.0-dev+3231"
            return Expr(:call, qconstr.sense, quadToExpr(qconstr.terms, true), 0)
        else
            return Expr(:comparison, quadToExpr(qconstr.terms, true), qconstr.sense, 0)
        end
    else
        i -= nlin + nquad
        # for now, don't pass simplified expressions
        constr = d.m.nlpdata.nlconstr[i]
        ex = constr.terms
        adj = adjmat(ex.nd)
        julia_expr = tapeToExpr(1, ex.nd, adj, ex.const_values, d.parameter_values, d.subexpressions_as_julia_expressions)
        if sense(constr) == :range
            return Expr(:comparison, constr.lb, :(<=), julia_expr, :(<=), constr.ub)
        else
            if VERSION >= v"0.5.0-dev+3231"
                return Expr(:call, sense(constr), julia_expr, rhs(constr))
            else
                return Expr(:comparison, julia_expr, sense(constr), rhs(constr))
            end
        end
    end
end

const ENABLE_NLP_RESOLVE = Array(Bool,1)
function EnableNLPResolve()
    ENABLE_NLP_RESOLVE[1] = true
end
function DisableNLPResolve()
    ENABLE_NLP_RESOLVE[1] = false
end
export EnableNLPResolve, DisableNLPResolve

function _buildInternalModel_nlp(m::Model, traits)

    linobj, linrowlb, linrowub = prepProblemBounds(m)

    nldata::NLPData = m.nlpdata
    if m.internalModelLoaded && !m.simplify_nonlinear_expressions
        @assert isa(nldata.evaluator, NLPEvaluator)
        d = nldata.evaluator
        fill!(d.last_x, NaN)
        if length(nldata.nlparamvalues) == 0 && !ENABLE_NLP_RESOLVE[1]
            # no parameters and haven't explicitly allowed resolves
            # error to prevent potentially incorrect answers
            msg = """
            There was a recent **breaking** change in behavior
            for solving sequences of nonlinear models.
            Previously, users were allowed to modify the data in the model
            by modifying the values stored in their own data arrays.
            For example:

            data = [1.0]
            @NLconstraint(m, data[1]*x <= 1)
            solve(m)
            data[1] = 2.0
            solve(m) # coefficient is updated

            However, this behavior **no longer works**. Instead,
            nonlinear parameters defined with @defNLParam should be used.
            See the latest JuMP documentation for more information.
            It is possible that this model was exploiting the previous behavior,
            and out of extreme caution we have temporarily introduced
            this error message.
            If you are sure that you are solving the correct model,
            then call `EnableNLPResolve()` at the top of this file to disable
            this error message.
            To return to the last version of JuMP which supported the
            old behavior, run `Pkg.pin("JuMP",v"0.11.1")`.
            """
            error(msg)
        end
    else
        d = NLPEvaluator(m)
        nldata.evaluator = d
    end

    nlp_lb, nlp_ub = constraintbounds(m)
    numConstr = length(nlp_lb)

    m.internalModel = MathProgBase.NonlinearModel(m.solver)

    MathProgBase.loadproblem!(m.internalModel, m.numCols, numConstr, m.colLower, m.colUpper, nlp_lb, nlp_ub, m.objSense, d)
    if traits.int
        if applicable(MathProgBase.setvartype!, m.internalModel, m.colCat)
            MathProgBase.setvartype!(m.internalModel, vartypes_without_fixed(m))
        else
            error("Solver does not support discrete variables")
        end
    end

    if !any(isnan,m.colVal)
        MathProgBase.setwarmstart!(m.internalModel, m.colVal)
    else
        initval = copy(m.colVal)
        initval[isnan(m.colVal)] = 0
        MathProgBase.setwarmstart!(m.internalModel, min(max(m.colLower,initval),m.colUpper))
    end

    m.internalModelLoaded = true

    nothing
end

function solvenlp(m::Model, traits; suppress_warnings=false)

    @assert m.internalModelLoaded

    MathProgBase.optimize!(m.internalModel)
    stat = MathProgBase.status(m.internalModel)

    if stat != :Infeasible && stat != :Unbounded
        m.objVal = MathProgBase.getobjval(m.internalModel)
        m.colVal = MathProgBase.getsolution(m.internalModel)
    end

    if stat != :Optimal
        suppress_warnings || warn("Not solved to optimality, status: $stat")
    end
    if stat == :Optimal && !traits.int
        if applicable(MathProgBase.getconstrduals, m.internalModel) && applicable(MathProgBase.getreducedcosts, m.internalModel)
            nlduals = MathProgBase.getconstrduals(m.internalModel)
            m.linconstrDuals = nlduals[1:length(m.linconstr)]
            # quadratic duals currently not available, formulate as nonlinear constraint if needed
            m.nlpdata.nlconstrDuals = nlduals[length(m.linconstr)+length(m.quadconstr)+1:end]
            m.redCosts = MathProgBase.getreducedcosts(m.internalModel)
        else
            suppress_warnings || Base.warn_once("Nonlinear solver does not provide dual solutions")
        end
    end

    #d = m.nlpdata.evaluator
    #println("feval $(d.eval_f_timer)\nfgrad $(d.eval_grad_f_timer)\ngeval $(d.eval_g_timer)\njaceval $(d.eval_jac_g_timer)\nhess $(d.eval_hesslag_timer)")

    return stat::Symbol

end

# getvalue for nonlinear subexpressions
function getvalue(x::NonlinearExpression)
    m = x.m
    # recompute EVERYTHING here
    # could be smarter and cache

    nldata::NLPData = m.nlpdata
    subexpr = Array(Vector{NodeData},0)
    for nlexpr in nldata.nlexpr
        push!(subexpr, nlexpr.nd)
    end

    this_subexpr = nldata.nlexpr[x.index]

    max_len = length(this_subexpr.nd)

    subexpression_order, individual_order = order_subexpressions(Vector{NodeData}[this_subexpr.nd],subexpr)

    subexpr_values = Array(Float64, length(subexpr))

    for k in subexpression_order
        max_len = max(max_len, length(nldata.nlexpr[k].nd))
    end

    forward_storage = Array(Float64, max_len)
    partials_storage = Array(Float64, max_len)
    user_input_buffer = zeros(m.numCols)
    user_output_buffer = zeros(m.numCols)

    for k in subexpression_order # compute value of dependent subexpressions
        ex = nldata.nlexpr[k]
        adj = adjmat(ex.nd)
        subexpr_values[k] = forward_eval(forward_storage,partials_storage,ex.nd,adj,ex.const_values,nldata.nlparamvalues,m.colVal,subexpr_values,user_input_buffer,user_output_buffer)
    end

    adj = adjmat(this_subexpr.nd)

    return forward_eval(forward_storage,partials_storage,this_subexpr.nd,adj,this_subexpr.const_values,nldata.nlparamvalues,m.colVal,subexpr_values,user_input_buffer,user_output_buffer)
end

type UserFunctionEvaluator <: MathProgBase.AbstractNLPEvaluator
    f
    ∇f
    len::Int
end

function MathProgBase.eval_f(d::UserFunctionEvaluator,x)
    @assert length(x) == d.len
    return d.f(x...)::eltype(x)
end
function MathProgBase.eval_grad_f(d::UserFunctionEvaluator,grad,x)
    d.∇f(grad,x)
    nothing
end

function UserAutoDiffEvaluator(dimension::Integer, f::Function)
    ∇f = (out,y) -> ForwardDiff.gradient!(out, x -> f(x...), y)
    return UserFunctionEvaluator(f, ∇f, dimension)
end

function register(s::Symbol, dimension::Integer, f::Function; autodiff::Bool=false)
    autodiff == true || error("If only the function is provided, must set autodiff=true")

    if dimension == 1
        fprime = x -> ForwardDiff.derivative(f, x)
        fprimeprime = x -> ForwardDiff.derivative(fprime, x)
        ReverseDiffSparse.register_univariate_operator(s, f, fprime, fprimeprime)
    else
        ReverseDiffSparse.register_multivariate_operator(s, UserAutoDiffEvaluator(dimension, f))
    end

end

function register(s::Symbol, dimension::Integer, f::Function, ∇f::Function; autodiff::Bool=false)
    if dimension == 1
        autodiff == true || error("Currently must provide 2nd order derivatives of univariate functions. Try setting autodiff=true.")
        fprimeprime = x -> ForwardDiff.derivative(∇f, x)
        ReverseDiffSparse.register_univariate_operator(s, f, ∇f, fprimeprime)
    else
        autodiff == false || Base.warn_once("autodiff=true ignored since gradient is already provided.")
        d = UserFunctionEvaluator(f, (g,x)->∇f(g,x...), dimension)
        ReverseDiffSparse.register_multivariate_operator(s, d)
    end

end

function register(s::Symbol, dimension::Integer, f::Function, ∇f::Function, ∇²f::Function)
    dimension == 1 || error("Providing hessians for multivariate functions is not yet supported")
    ReverseDiffSparse.register_univariate_operator(s, f, ∇f, ∇²f)
end

# Ex: setNLobjective(m, :Min, :($x + $y^2))
function setNLobjective(m::Model, sense::Symbol, x)
    initNLP(m)
    setobjectivesense(m, sense)
    ex = NonlinearExprData(m, x)
    m.nlpdata.nlobj = ex
    m.obj = zero(QuadExpr)
    m.internalModelLoaded = false
    return
end

# Ex: addNLconstraint(m, :($x + $y^2 <= 1))
function addNLconstraint(m::Model, ex::Expr)
    if VERSION < v"0.5.0-dev+3231"
        ex = comparison_to_call(ex)
    end
    initNLP(m)
    m.internalModelLoaded = false
    if isexpr(ex, :call) # one-sided constraint
        # Simple comparison - move everything to the LHS
        op = ex.args[1]
        if op == :(==)
            lb = 0.0
            ub = 0.0
        elseif op == :(<=) || op == :(≤)
            lb = -Inf
            ub = 0.0
        elseif op == :(>=) || op == :(≥)
            lb = 0.0
            ub = Inf
        else
            error("in addNLconstraint ($ex): expected comparison operator (<=, >=, or ==).")
        end
        lhs = :($(ex.args[2]) - $(ex.args[3]))
        c = NonlinearConstraint(NonlinearExprData(m, lhs), lb, ub)
        push!(m.nlpdata.nlconstr, c)
        return ConstraintRef{Model,NonlinearConstraint}(m, length(m.nlpdata.nlconstr))
    elseif isexpr(ex, :comparison)
        # ranged row
        if (ex.args[2] != :<= && ex.args[2] != :≤) || (ex.args[4] != :<= && ex.args[4] != :≤)
            error("in addNLconstraint ($ex): only ranged rows of the form lb <= expr <= ub are supported.")
        end
        lb = ex.args[1]
        ub = ex.args[5]
        if !isa(lb,Number)
            error(string("in addNLconstraint (",ex,"): expected ",lb," to be a number."))
        elseif !isa(ub,Number)
            error(string("in addNLconstraint (",ex,"): expected ",ub," to be a number."))
        end
        c = NonlinearConstraint(NonlinearExprData(m, ex.args[3]), lb, ub)
        push!(m.nlpdata.nlconstr, c)
        return ConstraintRef{Model,NonlinearConstraint}(m, length(m.nlpdata.nlconstr))
    else
        # Unknown
        error("in addNLconstraint ($ex): constraints must be in one of the following forms:\n" *
              "       expr1 <= expr2\n" * "       expr1 >= expr2\n" *
              "       expr1 == expr2")
    end
end


using GR

mwidth, mheight, width, height = GR.inqdspsize()

w = 800
h = 600

setcolorrep(1, 1.0, 1.0, 1.0)
setcolorrep(2, 0.0, 0.0, 0.0)
setlinecolorind(2)
setlinewidth(1)
setlinetype(GR.LINETYPE_SOLID)
setfillcolorind(2)
setfillintstyle(GR.INTSTYLE_SOLID)

while true
    clearws()

    #w = 400 + round(Int, rand() * 200)
    #h = 400 + round(Int, rand() * 200)
    if w >= h
        ratio = float(h) / w
        msize = mwidth * w / width
        GR.setwsviewport(0, msize, 0, msize * ratio)
        GR.setwswindow(0, 1, 0, ratio)
        setviewport(0, 1, 0, ratio)
    else
        ratio = float(w) / h
        msize = mheight * h / height
        GR.setwsviewport(0, msize * ratio, 0, msize)
        GR.setwswindow(0, ratio, 0, 1)
        setviewport(0, ratio, 0, 1)
    end

    setwindow(1, w, 1, h)

    #drawrect(1, w, 1, h)
    fillrect(1, w, 1, h)
    #polyline([1,w],[1,h])
    #polyline([1,w],[h,1])
    text(0.05, 0.05, "$w x $h")

    updatews()
    sleep(0.5)
end


# Transform a range expression (e.g., `1:10`) into a (more-or-less) equivalent
# construct in JavaScript.
# NOTE: We could also do this be converting a range into an array and including
# that array verbatim in the JS code, but that's not legible if there are a
# significant number of elements in the array (e.g., it's not desirable to
# include an array literal with 100 elements in it in the generated JS code).
function crawl_call(::Val{:(:)}, start::Integer, stop::Integer)
    # Because JavaScript arrays are terrible, we have to fill the array with
    # undefined (otherwise it can't be iterated over) after we construct it with
    # the length specified in the constructor.
    return :(JSTerminal(jsstring($((
        "(",
        "new Array($(stop - start + 1))",
        ".fill(undefined)",
        ".map((_, i) => i + $start)",
        ")",
    )...))))
end

function crawl_call(::Val{:(:)}, args...)
    error("Range syntax cannot be parsed by JSExpr.")
end


export REMDQNLearner

mutable struct REMDQNLearner{Tq,Tt,Tf,R<:AbstractRNG} <: Any
    approximator::Tq
    target_approximator::Tt
    loss_func::Tf
    min_replay_history::Int
    update_freq::Int
    update_step::Int
    target_update_freq::Int
    sampler::NStepBatchSampler
    ensemble_num::Int
    ensemble_method::Symbol
    rng::R
    # for logging
    loss::Float32
end

"""
    REMDQNLearner(;kwargs...)

See paper: [An Optimistic Perspective on Offline Reinforcement Learning](https://arxiv.org/abs/1907.04543)

# Keywords

- `approximator`::[`AbstractApproximator`](@ref): used to get Q-values of a state.
- `target_approximator`::[`AbstractApproximator`](@ref): similar to `approximator`, but used to estimate the target (the next state).
- `loss_func`: the loss function.
- `γ::Float32=0.99f0`: discount rate.
- `batch_size::Int=32`
- `update_horizon::Int=1`: length of update ('n' in n-step update).
- `min_replay_history::Int=32`: number of transitions that should be experienced before updating the `approximator`.
- `update_freq::Int=4`: the frequency of updating the `approximator`.
- `ensemble_num::Int=1`: the number of ensemble approximators.
- `ensemble_method::Symbol=:rand`: the method of combining Q values. ':rand' represents random ensemble mixture, and ':mean' is the average.
- `target_update_freq::Int=100`: the frequency of syncing `target_approximator`.
- `stack_size::Union{Int, Nothing}=4`: use the recent `stack_size` frames to form a stacked state.
- `traces = SARTS`, set to `SLARTSL` if you are to apply to an environment of `FULL_ACTION_SET`.
- `rng = Random.GLOBAL_RNG`
"""
function REMDQNLearner(;
    approximator::Tq,
    target_approximator::Tt,
    loss_func::Tf,
    stack_size::Union{Int,Nothing} = nothing,
    γ::Float32 = 0.99f0,
    batch_size::Int = 32,
    update_horizon::Int = 1,
    min_replay_history::Int = 32,
    update_freq::Int = 1,
    ensemble_num::Int = 1,
    ensemble_method::Symbol = :rand,
    target_update_freq::Int = 100,
    traces = SARTS,
    update_step = 0,
    rng = Random.GLOBAL_RNG,
) where {Tq,Tt,Tf}
    copyto!(approximator, target_approximator)
    sampler = NStepBatchSampler{traces}(;
        γ = γ,
        n = update_horizon,
        stack_size = stack_size,
        batch_size = batch_size,
    )
    REMDQNLearner(
        approximator,
        target_approximator,
        loss_func,
        min_replay_history,
        update_freq,
        update_step,
        target_update_freq,
        sampler,
        ensemble_num,
        ensemble_method,
        rng,
        0.0f0,
    )
end

Functors.functor(x::REMDQNLearner) = (Q = x.approximator, Qₜ = x.target_approximator),
y -> begin
    x = @set x.approximator = y.Q
    x = @set x.target_approximator = y.Qₜ
    x
end

function (learner::REMDQNLearner)(env)
    s = send_to_device(device(learner.approximator), state(env))
    s = Flux.unsqueeze(s, ndims(s) + 1)
    q = reshape(learner.approximator(s), :, learner.ensemble_num)
    vec(mean(q, dims = 2)) |> send_to_host
end

function RLBase.update!(learner::REMDQNLearner, batch::NamedTuple)
    Q = learner.approximator
    Qₜ = learner.target_approximator
    γ = learner.sampler.γ
    loss_func = learner.loss_func
    n = learner.sampler.n
    batch_size = learner.sampler.batch_size
    ensemble_num = learner.ensemble_num
    D = device(Q)
    # Build a convex polygon to make a combination of multiple Q-value estimates as a Q-value estimate.
    if learner.ensemble_method == :rand
        convex_polygon = rand(Float32, (1, ensemble_num))
    else
        convex_polygon = ones(Float32, (1, ensemble_num))
    end
    convex_polygon ./= sum(convex_polygon)
    convex_polygon = send_to_device(D, convex_polygon)

    s, a, r, t, s′ = (send_to_device(D, batch[x]) for x in SARTS)
    a = CartesianIndex.(a, 1:batch_size)

    target_q = Qₜ(s′)
    target_q = convex_polygon .* reshape(target_q, :, ensemble_num, batch_size)
    target_q = dropdims(sum(target_q, dims = 2), dims = 2)

    if haskey(batch, :next_legal_actions_mask)
        l′ = send_to_device(D, batch[:next_legal_actions_mask])
        target_q .+= ifelse.(l′, 0.0f0, typemin(Float32))
    end

    q′ = dropdims(maximum(target_q; dims = 1), dims = 1)
    G = r .+ γ^n .* (1 .- t) .* q′

    gs = gradient(params(Q)) do
        q = Q(s)
        q = convex_polygon .* reshape(q, :, ensemble_num, batch_size)
        q = dropdims(sum(q, dims = 2), dims = 2)[a]

        loss = loss_func(G, q)
        ignore() do
            learner.loss = loss
        end
        loss
    end

    update!(Q, gs)
end


@inline function sor_kernel(
    potential::AbstractArray{T, 4},
    imp_scale::AbstractArray{T, 4},
    point_types::AbstractArray{PointType, 4},
    volume_weights::AbstractArray{T, 4},
    q_eff_imp::AbstractArray{T, 4},
    q_eff_fix::AbstractArray{T, 4},
    ϵ_r::AbstractArray{T, 3},
    geom_weights::NTuple{3, AbstractArray{T, 2}},
    sor_const::AbstractArray{T, 1},
    update_even_points::Bool,
    depletion_handling_enabled::Bool,
    is_weighting_potential::Bool,
    only2d::Bool,
    ::Type{S},
    gpu_inds::NTuple{3, Int}
) where {T, S}
    # Comparison to CPU indices: (Cyl / Car)
    # i3 <-> idx3 / ir / iz
    # i2 <-> idx2 / iφ / iy
    # i1 <-> idx1 / iz / ix
    i1, in2, in3 = gpu_inds
    i1 += 1
    i2 = in2 + 1
    i3 = in3 + 1
    in1 = nidx(i1, update_even_points, iseven(i2 + i3))
    i1r = get_rbidx_right_neighbour(i1, update_even_points, iseven(i2 + i3))

    rb_tar_idx, rb_src_idx = update_even_points ? (rb_even::Int, rb_odd::Int) : (rb_odd::Int, rb_even::Int)

    geom_weights_3 = get_geom_weights_outerloop(geom_weights, in3, S)

    geom_weights_2 = prepare_weights_in_middleloop(
        geom_weights, S, i2, in2,
        geom_weights_3...,
        in3 == 1
    )

    weights = calculate_sor_weights(
        in1,
        S,
        ϵ_r,
        geom_weights[3],
        i2, in2, i3, in3,
        geom_weights_2...
    )

    old_potential = potential[i1, i2, i3, rb_tar_idx]

    neighbor_potentials = get_neighbor_potentials(
        potential, old_potential, i1, i2, i3, i1r, in2, in3, rb_src_idx, only2d
    )

    new_potential = calc_new_potential_by_neighbors_3D(
        volume_weights[i1, i2, i3, rb_tar_idx],
        weights,
        neighbor_potentials
    )

    update_point = point_types[i1, i2, i3, rb_tar_idx] & update_bit > 0

    if depletion_handling_enabled
        new_potential, new_imp_scale = handle_depletion(
            new_potential,
            imp_scale[i1, i2, i3, rb_tar_idx],
            r0_handling_depletion_handling(neighbor_potentials, S, in3),
            q_eff_imp[i1, i2, i3, rb_tar_idx],
            volume_weights[i1, i2, i3, rb_tar_idx]
        )
        imp_scale[i1, i2, i3, rb_tar_idx] = ifelse(update_point, new_imp_scale, imp_scale[i1, i2, i3, rb_tar_idx])
    end

    if !is_weighting_potential
        q_eff = if depletion_handling_enabled
            q_eff_fix[i1, i2, i3, rb_tar_idx]
        else
            q_eff_imp[i1, i2, i3, rb_tar_idx] + q_eff_fix[i1, i2, i3, rb_tar_idx]
        end
        new_potential += q_eff * volume_weights[i1, i2, i3, rb_tar_idx]
    end

    if is_weighting_potential || !depletion_handling_enabled
        new_potential = apply_over_relaxation(new_potential, old_potential, get_sor_constant(sor_const, S, in3))
    end

    potential[i1, i2, i3, rb_tar_idx] = ifelse(update_point, new_potential, old_potential)
    nothing
end

@inline function get_neighbor_potentials(
    potential::AbstractArray{T, 4},
    old_potential, i1, i2, i3, i1r, in2, in3, rb_src_idx, only2d::Bool
)::NTuple{6, T} where {T}
    @inbounds return ( # p: potential; 1: RB-dimension; l/r: left/right
    potential[i1r - 1,     i2,     i3, rb_src_idx], # p1l
    potential[    i1r,     i2,     i3, rb_src_idx], # p1r
    only2d ? old_potential : potential[ i1,    in2, i3, rb_src_idx], # p2l
    only2d ? old_potential : potential[ i1, i2 + 1, i3, rb_src_idx], # p2r
    potential[     i1,     i2,    in3, rb_src_idx], # p3l
    potential[     i1,     i2, i3 + 1, rb_src_idx], # p3r
    )
end

@inline function prepare_weights_in_middleloop(
    geom_weights::NTuple{3, <:AbstractArray{T, 2}}, ::Type{Cylindrical},
    i2, in2,
    pwwrr, pwwrl, r_inv_pwΔmpr, Δr_ext_inv_r_pwmprr, Δr_ext_inv_l_pwmprl, Δmpr_squared,
    is_r0::Bool
) where {T}
    pwwφr        = geom_weights[2][1, in2]
    pwwφl        = geom_weights[2][2, in2]
    pwΔmpφ       = geom_weights[2][3, in2]
    Δφ_ext_inv_r = geom_weights[2][4,  i2]
    Δφ_ext_inv_l = geom_weights[2][4, in2]

    if is_r0
        pwwφr = T(0.5)
        pwwφl = T(0.5)
        pwΔmpφ = T(2π)
        Δφ_ext_inv_r = inv(pwΔmpφ)
        Δφ_ext_inv_l = Δφ_ext_inv_r
    end
    pwwrr_pwwφr = pwwrr * pwwφr
    pwwrl_pwwφr = pwwrl * pwwφr
    pwwrr_pwwφl = pwwrr * pwwφl
    pwwrl_pwwφl = pwwrl * pwwφl

    pwΔmpφ_Δmpr_squared = pwΔmpφ * Δmpr_squared
    Δr_ext_inv_r_pwmprr_pwΔmpφ = Δr_ext_inv_r_pwmprr * pwΔmpφ
    Δr_ext_inv_l_pwmprl_pwΔmpφ = Δr_ext_inv_l_pwmprl * pwΔmpφ
    r_inv_pwΔmpr_Δφ_ext_inv_r = r_inv_pwΔmpr * Δφ_ext_inv_r
    r_inv_pwΔmpr_Δφ_ext_inv_l = r_inv_pwΔmpr * Δφ_ext_inv_l
    return (
        pwwrr, pwwrl, pwwφr, pwwφl,
        pwwrr_pwwφr, pwwrl_pwwφr, pwwrr_pwwφl, pwwrl_pwwφl,
        pwΔmpφ_Δmpr_squared,
        Δr_ext_inv_r_pwmprr_pwΔmpφ, Δr_ext_inv_l_pwmprl_pwΔmpφ,
        r_inv_pwΔmpr_Δφ_ext_inv_r, r_inv_pwΔmpr_Δφ_ext_inv_l
    )
end

"""
    function get_sor_kernel(::Type{S}, args...)

where S is either Cartesian or Cylindrical.

Developer notes:
Currently (February 2022), there are some limitations to the `@kernel` macro
of the package KernelAbstractions.jl. Especially, regarding usage of dispatch.

Thus, we have to write two kernel functions right now for the Cartesian & Cylindrical case:
`sor_cyl_gpu!` and `sor_car_gpu!`.

Inside kernel functions, everything is (and has to be) inlined and we can make use of multiple dispatch.
So in the end we only have to write one function for the kernel, `sor_kernel`,
which is then inlined inside the two kernel functions.

We can also use most of the CPU functions with the restriction that all
types have to be independent on the GPU-indices of the kernel.
E.g., making use of `i23_is_even_t = Val(iseven(i2 + i3))` and other similar statements is not possible,
which are used in the CPU implementation for optimization.
On the GPU (currently) those statements have to be calculated and booleans have to be passed.
Maybe this will change in the future.
"""
get_sor_kernel(::Type{Cylindrical}, dev, CPU_via_KernelAbstractions, args...) = sor_cyl_gpu!(dev, args...)
get_sor_kernel(::Type{Cartesian},   dev, CPU_via_KernelAbstractions, args...) = sor_car_gpu!(dev, args...)
get_sor_kernel(::Type{Cylindrical}, ::CPU, CPU_via_KernelAbstractions::Val{false}) = nothing
get_sor_kernel(::Type{Cartesian},   ::CPU, CPU_via_KernelAbstractions::Val{false}) = nothing
get_sor_kernel(::Type{Cylindrical}, dev::CPU, CPU_via_KernelAbstractions::Val{true}, args...) = sor_cyl_gpu!(dev, args...)
get_sor_kernel(::Type{Cartesian},   dev::CPU, CPU_via_KernelAbstractions::Val{true}, args...) = sor_car_gpu!(dev, args...)


default_partitioner(::Type{<:SparseMatrixVBC{U, W, Tv, Ti}}) where {U, W, Tv, Ti} =
    AlternatePacker(
        EquiChunker(),
        EquiChunker(),
        DynamicTotalChunker(ConstrainedCost(model_SparseMatrixVBC_memory(Tv, Ti), WidthCost(), W)),
        DynamicTotalChunker(ConstrainedCost(permutedims(model_SparseMatrixVBC_memory(Tv, Ti)), WidthCost(), U)),
        DynamicTotalChunker(ConstrainedCost(model_SparseMatrixVBC_memory(Tv, Ti), WidthCost(), W)),
    )

function SparseMatrixVBC{U, W}(A::SparseMatrixCSC{Tv, Ti}, method=default_partitioner(SparseMatrixVBC{U, W, Tv, Ti})) where {U, W, Tv, Ti}
    Π, Φ = pack_plaid(A, method)
    return SparseMatrixVBC{U, W}(A, convert(SplitPartition, Π), convert(SplitPartition, Φ))
end

function SparseMatrixVBC{U, W}(A::SparseMatrixCSC{Tv, Ti}, Π::SplitPartition{Ti}, Φ::SplitPartition{Ti}) where {U, W, Tv, Ti}
    @inbounds begin
        (m, n) = size(A)
        K = length(Π)
        L = length(Φ)
        Π_spl = Π.spl
        Φ_spl = Φ.spl
        Π_asg = convert(MapPartition, Π).asg
        hst = zeros(Ti, K)
        A_pos = A.colptr
        A_idx = A.rowval
        A_val = A.nzval
        pos = undefs(Ti, L + 1)
        ofs = undefs(Ti, L + 1)
        pos[1] = 1
        ofs[1] = 1
        for l = 1:L
            ofs[l + 1] = ofs[l]
            pos[l + 1] = pos[l]
            j = Φ_spl[l]
            j′ = Φ_spl[l + 1]
            w = j′ - j
            for q = A_pos[j]:A_pos[j′] - 1
                i = A_idx[q]
                k = Π_asg[i]
                if hst[k] < l
                    u = Π_spl[k + 1] - Π_spl[k]
                    pos[l + 1] += 1
                    ofs[l + 1] += u * w
                end
                hst[k] = l
            end
        end

        idx = Vector{Ti}(undef, pos[end] - 1)
        Δw = fld(CpuId.simdbytes(), sizeof(Tv)) #TODO need to define an upper limit on this one
        val = Vector{Tv}(undef, ofs[end] - 1 + U * (Δw * cld(W, Δw)))
        for Q = ofs[end] : ofs[end]  - 1 + U * (Δw * cld(W, Δw)) #extra stuff to keep vector access in bounds
            val[Q] = zero(Tv)
        end

        A_q = ones(Int, W)

        for k = 1:K
            @assert Π_spl[k + 1] - Π_spl[k] <= U
        end

        for l = 1:L
            j = Φ_spl[l]
            w = Φ_spl[l + 1] - j
            @assert w <= W
            if w == 1
                Q = pos[l]
                q = ofs[l]
                A_q_1 = A_pos[j]
                while A_q_1 < A_pos[j + 1]
                    k = Π_asg[A_idx[A_q_1]]
                    for i = Π_spl[k] : Π_spl[k + 1] - 1
                        if A_q_1 < A_pos[j + 1] && A_idx[A_q_1] == i
                            val[q] = A_val[A_q_1]
                            A_q_1 += 1
                        else
                            val[q] = zero(Tv)
                        end
                        q += 1
                    end
                    idx[Q] = k
                    Q += 1
                end
            else
                k = K + 1
                for Δj = 1:w
                    A_q[Δj] = A_pos[j + Δj - 1]
                    if A_q[Δj] < A_pos[j + Δj]
                        k = min(k, Π_asg[A_idx[A_q[Δj]]])
                    end
                end
                Q = pos[l]
                q = ofs[l]
                while k != K + 1
                    for i = Π_spl[k] : Π_spl[k + 1] - 2
                        for Δj = 1:w
                            if A_q[Δj] < A_pos[j + Δj] && A_idx[A_q[Δj]] == i
                                val[q] = A_val[A_q[Δj]]
                                A_q[Δj] += 1
                            else
                                val[q] = zero(Tv)
                            end
                            q += 1
                        end
                    end
                    k′ = K + 1
                    let i = Π_spl[k + 1] - 1
                        for Δj = 1:w
                            if A_q[Δj] < A_pos[j + Δj]
                                if A_idx[A_q[Δj]] == i
                                    val[q] = A_val[A_q[Δj]]
                                    A_q[Δj] += 1
                                else
                                    val[q] = zero(Tv)
                                end
                                if A_q[Δj] < A_pos[j + Δj]
                                    k′ = min(k′, Π_asg[A_idx[A_q[Δj]]])
                                end
                            else
                                val[q] = zero(Tv)
                            end
                            q += 1
                        end
                    end
                    idx[Q] = k
                    Q += 1
                    k = k′
                end
            end
        end
        return SparseMatrixVBC{U, W, Tv, Ti}(m, n, Π, Φ, pos, idx, ofs, val)
    end
end


#=
nothing fancy - only used internally (and therefore cryptic) - stores some partially contracted things
seperates out this bit of logic from effective_excitation_hamiltonian (now more readable)
can also - potentially - partially reuse this in other algorithms
=#
struct QPEnv{A,B} <: Cache
    lBs::PeriodicArray{A,1}
    rBs::PeriodicArray{A,1}

    lenvs::B
    renvs::B
end

function environments(exci::Union{InfiniteQP,Multiline{<:InfiniteQP}}, H; solver=Defaults.solver)
    # Explicitly define optional arguments as these depend on solver,
    # which needs to come after these arguments.
    lenvs = environments(exci.left_gs, H; solver=solver)

    return environments(exci, H, lenvs; solver=solver)
end

function environments(exci::Union{InfiniteQP,Multiline{<:InfiniteQP}}, H, lenvs; solver=Defaults.solver)
    # Explicitly define optional arguments as these depend on solver,
    # which needs to come after these arguments.
    renvs = exci.trivial ? lenvs : environments(exci.right_gs, H; solver=solver)

    return environments(exci, H, lenvs, renvs; solver=solver)
end

function environments(exci::InfiniteQP, ham::MPOHamiltonian, lenvs, renvs;solver=Defaults.solver)
    ids = collect(Iterators.filter(x->isid(ham,x),2:ham.odim-1));

    #build lBs(c)
    lB_cur = [ TensorMap(zeros,eltype(exci),
                    virtualspace(exci.left_gs,0)*ham.domspaces[1,k]',
                    space(exci[1],3)'*virtualspace(exci.right_gs,0)) for k in 1:ham.odim]
    lBs = typeof(lB_cur)[]

    for pos = 1:length(exci)
        lB_cur = exci_transfer_left(lB_cur,ham[pos],exci.right_gs.AR[pos],exci.left_gs.AL[pos])*exp(conj(1im*exci.momentum))
        lB_cur += exci_transfer_left(leftenv(lenvs,pos,exci.left_gs),ham[pos],exci[pos],exci.left_gs.AL[pos])*exp(conj(1im*exci.momentum))

        exci.trivial && for i in ids
            @plansor lB_cur[i][-1 -2;-3 -4] -= lB_cur[i][1 4;-3 2]*r_RL(exci.left_gs,pos)[2;3]*τ[3 4;5 1]*l_RL(exci.left_gs,pos+1)[-1;6]*τ[5 6;-4 -2]
        end

        push!(lBs,lB_cur)
    end

    #build rBs(c)
    rB_cur = [ TensorMap(zeros,eltype(exci),
                    virtualspace(exci.left_gs,length(exci))*ham.imspaces[length(exci),k]',
                    space(exci[1],3)'*virtualspace(exci.right_gs,length(exci))) for k in 1:ham.odim]
    rBs = typeof(rB_cur)[]

    for pos=length(exci):-1:1
        rB_cur = exci_transfer_right(rB_cur,ham[pos],exci.left_gs.AL[pos],exci.right_gs.AR[pos])*exp(1im*exci.momentum)
        rB_cur += exci_transfer_right(rightenv(renvs,pos,exci.right_gs),ham[pos],exci[pos],exci.right_gs.AR[pos])*exp(1im*exci.momentum)

        exci.trivial && for i in ids
            @plansor rB_cur[i][-1 -2;-3 -4] -= τ[6 4;1 3]*rB_cur[i][1 3;-3 2]*l_LR(exci.left_gs,pos)[2;4]*r_LR(exci.left_gs,pos-1)[-1;5]*τ[-2 -4;5 6]
        end

        push!(rBs,rB_cur)
    end
    rBs = reverse(rBs)

    local lBE::typeof(rB_cur);
    local rBE::typeof(rB_cur);

    @sync begin
        @Threads.spawn lBE = left_excitation_transfer_system(lB_cur,ham,exci,solver=solver)
        @Threads.spawn rBE = right_excitation_transfer_system(rB_cur,ham,exci, solver=solver)
    end

    lBs[end] = lBE;

    for i=1:length(exci)-1
        lBE = exci_transfer_left(lBE,ham[i],exci.right_gs.AR[i],exci.left_gs.AL[i])*exp(conj(1im*exci.momentum))

        exci.trivial && for k in ids
            @plansor lBE[k][-1 -2;-3 -4] -= lBE[k][1 4;-3 2]*r_RL(exci.left_gs,i)[2;3]*τ[3 4;5 1]*l_RL(exci.left_gs,i+1)[-1;6]*τ[5 6;-4 -2]
        end

        lBs[i] += lBE;
    end

    rBs[1] = rBE;

    for i=length(exci):-1:2
        rBE = exci_transfer_right(rBE,ham[i],exci.left_gs.AL[i],exci.right_gs.AR[i])*exp(1im*exci.momentum)

        exci.trivial && for k in ids
            @plansor rBE[k][-1 -2;-3 -4] -= τ[6 4;1 3]*rBE[k][1 3;-3 2]*l_LR(exci.left_gs,i)[2;4]*r_LR(exci.left_gs,i-1)[-1;5]*τ[-2 -4;5 6]
        end

        rBs[i] += rBE
    end

    return QPEnv(PeriodicArray(lBs),PeriodicArray(rBs),lenvs,renvs)
end

function environments(exci::FiniteQP,ham::MPOHamiltonian,lenvs=environments(exci.left_gs,ham),renvs=exci.trivial ? lenvs : environments(exci.right_gs,ham))
    #construct lBE
    lB_cur = [ TensorMap(zeros,eltype(exci),
                    virtualspace(exci.left_gs,0)*ham.domspaces[1,k]',
                    space(exci[1],3)'*virtualspace(exci.left_gs,0)) for k in 1:ham.odim]
    lBs = typeof(lB_cur)[]
    for pos = 1:length(exci)
        lB_cur = exci_transfer_left(lB_cur,ham[pos],exci.right_gs.AR[pos],exci.left_gs.AL[pos])
        lB_cur += exci_transfer_left(leftenv(lenvs,pos,exci.left_gs),ham[pos],exci[pos],exci.left_gs.AL[pos])
        push!(lBs,lB_cur)
    end

    #build rBs(c)
    rB_cur = [ TensorMap(zeros,eltype(exci),
                    virtualspace(exci.right_gs,length(exci))*ham.imspaces[length(exci),k]',
                    space(exci[1],3)'*virtualspace(exci.right_gs,length(exci))) for k in 1:ham.odim]
    rBs = typeof(rB_cur)[]
    for pos=length(exci):-1:1
        rB_cur = exci_transfer_right(rB_cur,ham[pos],exci.left_gs.AL[pos],exci.right_gs.AR[pos])
        rB_cur += exci_transfer_right(rightenv(renvs,pos,exci.right_gs),ham[pos],exci[pos],exci.right_gs.AR[pos])
        push!(rBs,rB_cur)
    end
    rBs=reverse(rBs)

    return QPEnv(PeriodicArray(lBs),PeriodicArray(rBs),lenvs,renvs)
end

function environments(exci::Multiline{<:InfiniteQP}, ham::MPOMultiline, lenvs, renvs;solver=Defaults.solver)
    exci.trivial || @warn "there is a phase ambiguity in topologically nontrivial statmech excitations"

    left_gs = exci.left_gs;
    right_gs = exci.right_gs;
    exci_space = space(exci[1][1],3);
    lBs = map(1:size(left_gs,1)) do row
        renorms = (map(1:size(left_gs,2)) do col
            v = transfer_left(leftenv(lenvs,row,col,left_gs),ham[row,col],left_gs.AC[row,col],left_gs.AC[row+1,col]);
            @plansor v[1 2;3]*rightenv(lenvs,row,col,left_gs)[3 2;1]
        end).^-1;

        lB_cur = TensorMap(zeros,eltype(left_gs.AL[1,1]),virtualspace(left_gs,row+1,0)*_firstspace(ham[row,1])',exci_space'*virtualspace(left_gs,row,0));

        c_lBs = typeof(lB_cur)[];
        for col in 1:size(left_gs,2)
            lB_cur = renorms[col]*exci_transfer_left(lB_cur,ham[row,col],right_gs.AR[row,col],left_gs.AL[row+1,col])*exp(-1im*exci.momentum)
            lB_cur += renorms[col]*exci_transfer_left(leftenv(lenvs,row,col,left_gs),ham[row,col],exci[row][col],left_gs.AL[row+1,col])*exp(-1im*exci.momentum)
            push!(c_lBs,lB_cur);
        end

        (c_lBs[end],convhist) = linsolve(lB_cur,lB_cur,solver) do v
            y = exp(-1im*size(left_gs,2)*exci.momentum)*prod(renorms)*exci_transfer_left(v,ham[row,:],right_gs.AR[row,:],left_gs.AL[row+1,:])
            if exci.trivial
                @plansor y[-1 -2;-3 -4] -= y[4 2;-3 3]*conj(left_gs.CR[row+1,0][4;1])*rightenv(lenvs,row,0,left_gs)[3 2;1]*leftenv(lenvs,row,1,left_gs)[-1 -2;5]*left_gs.CR[row,0][5;-4]
            end
            v - y
        end
        convhist.converged == 0 && @warn "lbe failed to converge $(convhist.normres)"

        cur = c_lBs[end];
        for col in 1:size(left_gs,2)-1
            cur = renorms[col]*exci_transfer_left(cur,ham[row,col],right_gs.AR[row,col],left_gs.AL[row+1,col])*exp(conj(1im*exci.momentum))
            c_lBs[col] += cur
        end

        PeriodicArray(c_lBs)
    end

    rBs = map(1:size(left_gs,1)) do row
        renorms = (map(1:size(right_gs,2)) do col
            v = transfer_left(leftenv(renvs,row,col,right_gs),ham[row,col],right_gs.AC[row,col],right_gs.AC[row+1,col]);
            @plansor v[1 2;3]*rightenv(renvs,row,col,right_gs)[3 2;1]
        end).^-1;

        rB_cur = TensorMap(zeros,eltype(left_gs.AL[1,1]),virtualspace(left_gs,row,0)*_firstspace(ham[row,1]),exci_space'*virtualspace(left_gs,row+1,0));

        c_rBs = typeof(rB_cur)[];
        for col in size(left_gs,2):-1:1
            rB_cur = renorms[col]*exci_transfer_right(rB_cur,ham[row,col],left_gs.AL[row,col],right_gs.AR[row+1,col])*exp(1im*exci.momentum)
            rB_cur += renorms[col]*exci_transfer_right(rightenv(renvs,row,col,right_gs),ham[row,col],exci[row][col],right_gs.AR[row+1,col])*exp(1im*exci.momentum)
            push!(c_rBs,rB_cur);
        end
        c_rBs = reverse(c_rBs);

        (c_rBs[1],convhist) = linsolve(rB_cur,rB_cur,GMRES()) do v
            y = exp(1im*size(left_gs,2)*exci.momentum)*prod(renorms)*exci_transfer_right(v,ham[row,:],left_gs.AL[row,:],right_gs.AR[row+1,:])
            if exci.trivial
                @plansor y[-1 -2;-3 -4] -= y[1 2;-3 3]*conj(left_gs.CR[row+1,0][1;4])*leftenv(lenvs,row,1,left_gs)[3 2;4]*rightenv(lenvs,row,0,left_gs)[5 -2;-4]*left_gs.CR[row,0][-1;5]
            end
            v - y
        end
        convhist.converged == 0 && @warn "rbe failed to converge $(convhist.normres)"

        cur = c_rBs[1];
        for col in size(left_gs,2):-1:2
            cur = renorms[col]*exci_transfer_right(cur,ham[row,col],left_gs.AL[row,col],right_gs.AR[row+1,col])*exp(1im*exci.momentum)
            c_rBs[col] += cur
        end

        PeriodicArray(c_rBs)
    end

    QPEnv(PeriodicArray(lBs),PeriodicArray(rBs),lenvs,renvs)
end


mutable struct GFRegenerateState
    prev_trace::DynamicDSLTrace
    trace::DynamicDSLTrace
    selection::Selection
    weight::Float64
    visitor::AddressVisitor
    params::Dict{Symbol,Any}
end

function GFRegenerateState(gen_fn, args, prev_trace,
                           selection, params)
    visitor = AddressVisitor()
    GFRegenerateState(prev_trace, DynamicDSLTrace(gen_fn, args), selection,
        0., visitor, params)
end

function traceat(state::GFRegenerateState, dist::Distribution{T},
                 args, key) where {T}
    local prev_retval::T
    local retval::T

    # check that key was not already visited, and mark it as visited
    visit!(state.visitor, key)

    # check for previous choice at this key
    has_previous = has_choice(state.prev_trace, key)
    if has_previous
        prev_choice = get_choice(state.prev_trace, key)
        prev_retval = prev_choice.retval
        prev_score = prev_choice.score
    end

    # check whether the key was selected
    in_selection = key in state.selection

    # get return value
    if has_previous && in_selection
        retval = random(dist, args...)
    elseif has_previous
        retval = prev_retval
    else
        retval = random(dist, args...)
    end

    # compute logpdf
    score = logpdf(dist, retval, args...)

    # update weight
    if has_previous && !in_selection
        state.weight += score - prev_score
    end

    # add to the trace
    add_choice!(state.trace, key, retval, score)

    retval
end

function traceat(state::GFRegenerateState, gen_fn::GenerativeFunction{T,U},
                 args, key) where {T,U}
    local prev_retval::T
    local trace::U
    local retval::T

    # check that key was not already visited, and mark it as visited
    visit!(state.visitor, key)

    # check whether the key was selected
    subselection = state.selection[key]

    # get subtrace
    has_previous = has_call(state.prev_trace, key)
    if has_previous
        prev_call = get_call(state.prev_trace, key)
        prev_subtrace = prev_call.subtrace
        get_gen_fn(prev_subtrace) === gen_fn || gen_fn_changed_error(key)
        (subtrace, weight, _) = regenerate(
            prev_subtrace, args, map((_) -> UnknownChange(), args), subselection)
    else
        (subtrace, weight) = generate(gen_fn, args, EmptyChoiceMap())
    end

    # update weight
    state.weight += weight

    # add to the trace
    add_call!(state.trace, key, subtrace)

    # get return value
    retval = get_retval(subtrace)

    retval
end

function splice(state::GFRegenerateState, gen_fn::DynamicDSLFunction,
                args::Tuple)
    prev_params = state.params
    state.params = gen_fn.params
    retval = exec(gen_fn, state, args)
    state.params = prev_params
    retval
end

function regenerate_delete_recurse(prev_trie::Trie{Any,ChoiceOrCallRecord},
                             visited::EmptySelection)
    noise = 0.
    for (key, choice_or_call) in get_leaf_nodes(prev_trie)
        if !choice_or_call.is_choice
            noise += choice_or_call.noise
        end
    end
    for (key, subtrie) in get_internal_nodes(prev_trie)
        noise += regenerate_delete_recurse(subtrie, EmptySelection())
    end
    noise
end

function regenerate_delete_recurse(prev_trie::Trie{Any,ChoiceOrCallRecord},
                             visited::DynamicSelection)
    noise = 0.
    for (key, choice_or_call) in get_leaf_nodes(prev_trie)
        if !(key in visited) && !choice_or_call.is_choice
            noise += choice_or_call.noise
        end
    end
    for (key, subtrie) in get_internal_nodes(prev_trie)
        subvisited = visited[key]
        noise += regenerate_delete_recurse(subtrie, subvisited)
    end
    noise
end

function regenerate(trace::DynamicDSLTrace, args::Tuple, argdiffs::Tuple,
                    selection::Selection)
    gen_fn = trace.gen_fn
    state = GFRegenerateState(gen_fn, args, trace, selection, gen_fn.params)
    retval = exec(gen_fn, state, args)
    set_retval!(state.trace, retval)
    visited = state.visitor.visited
    state.weight -= regenerate_delete_recurse(trace.trie, visited)
    (state.trace, state.weight, UnknownChange())
end


### RigidBodyDynamics ###
function RigidBodyDynamics.principal_value!(mech_scen::MechanismScenario, x::Vector{Float64})
    copyto!(mech_scen.float, x)
    principal_value!(mech_scen.float.state)
    copyto!(x, mech_scen.float)
    return nothing
end

function RigidBodyDynamics.set_configuration!(mech_scen::MechanismScenario, joint::Joint, config)
    return set_configuration!(mech_scen.float.state, joint, config)
end

### Binary_BB_Trees
Binary_BB_Trees.fit_tri_obb(eM::eMesh{Tri,Nothing}, k::Int64) = fit_tri_obb(eM.point[eM.tri[k]])
function Binary_BB_Trees.fit_tet_obb(eM::eMesh{Nothing,Tet}, k::Int64)
    i = eM.tet[k]
    return fit_tet_obb(eM.point[i], eM.ϵ[i])
end

### Base ###
function Base.copyto!(dest::TypedMechanismScenario{T}, src::AbstractVector{T}) where {N,T}
    nq = num_positions(dest.state)
    nv = num_velocities(dest.state)
    ns = length(dest.s)
    @boundscheck length(src) == nq + nv + ns || throw(DimensionMismatch())
    @inbounds copyto!(parent(dest.state.q), 1, src,           1, nq)
    @inbounds copyto!(parent(dest.state.v), 1, src, nq +      1, nv)
    @inbounds copyto!(dest.s,               1, src, nq + nv + 1, ns)
    setdirty!(dest.state)
    dest
end
function Base.copyto!(dest::AbstractVector{T}, src::TypedMechanismScenario{T}) where {N,T}
    nq = num_positions(src.state)
    nv = num_velocities(src.state)
    ns = length(src.s)
    length(dest) == nq + nv + ns || throw(DimensionMismatch())
    @inbounds copyto!(dest,           1, src.state.q, 1, nq)
    @inbounds copyto!(dest, nq +      1, src.state.v, 1, nv)
    @inbounds copyto!(dest, nq + nv + 1, src.s,       1, ns)
    dest
end
function Base.copyto!(ẋ::AbstractVector, src::TypedMechanismScenario{T}, result::DynamicsResult) where {N,T}
    nq = length(result.q̇)
    nv = length(result.v̇)
    ns = length(src.ṡ)
    @boundscheck length(ẋ) == nq + nv + ns || throw(DimensionMismatch())
    @inbounds copyto!(ẋ,           1, result.q̇, 1, nq)
    @inbounds copyto!(ẋ, nq +      1, result.v̇, 1, nv)
    @inbounds copyto!(ẋ, nq + nv + 1, src.ṡ,    1, ns)
    ẋ
end


include("DataDictionary.jl")
include("LoadDictionaries.jl")
include("LoadSynthData.jl")
include("FluxDriver.jl")
include("Calculate_constraints.jl")
include("Bounds.jl")
include("TXTLDictionary.jl")
include("Utility.jl")
include("CalcError.jl")
include("Plot.jl")
include("PlotSeparate.jl")
include("PlotErrorbar.jl")
include("dFBA.jl")
include("RunFBA.jl")
using PyPlot
using Interpolations
using LaTeXStrings
using GLPK
term_out(GLPK.OFF)


demand = nhhy * (1-unemp)*(1-tau)* ((prodlb+produb)/2)/nf

space = GridSpace((k1,k2))

covidmodel = ABM(CovidAgent.GenericAgent,space; scheduler = fastest)

#create emtpy list of employees for all firms
workerlist =  Array{Int}[]
for n=1:nf
    push!(workerlist, Int[])
end
#unemployd list
unemplist = [Dict{Int,CovidAgent.GenericAgent}() for tt=1:nsec, i=1:k1, j=1:k2] # only dimension nsec because there are no unemployd public workers
shorttimelist = [Dict{Int,CovidAgent.GenericAgent}() for tt=1:nsec, i=1:k1, j=1:k2]

#generate list of firms in each location
locf = [Int[] for tt = 1:nsec+1, i=1:k1, j=1:k2] # separated by type of firm
locallf = [Int[] for i=1:k1, j=1:k2] # all types

# first setor
for n =1: nfsec[1]
    k1help = 1+fld(mod(n-1,k1*k2),k2)
    k2help = 1+mod(n-1,k1*k2)-(k1help-1)*k2 # distribute evenly on grid
    push!(locf[1,k1help,k2help],n)
    push!(locallf[k1help,k2help],n)
end

for tt = 2:nsec
    for n = sum(nfsec[1:tt-1])+1: sum(nfsec[1:tt])
        k1help = 1+fld(mod(n-1,k1*k2),k2)
        k2help = 1+mod(n-1,k1*k2)-(k1help-1)*k2 # distribute evenly on grid
        push!(locf[tt,k1help,k2help],n)
        push!(locallf[k1help,k2help],n)
    end
end
for n = sum(nfsec[1:nsec])+1: nf
    k1help = 1+fld(mod(n-1,k1*k2),k2)
    k2help = 1+mod(n-1,k1*k2)-(k1help-1)*k2 # distribute evenly on grid
    push!(locf[nsec+1,k1help,k2help],n)
    push!(locallf[k1help,k2help],n)
end

#generate list of hh in each location
lochh = [Dict{Int,CovidAgent.GenericAgent}() for tt = 1:2, i=1:k1, j=1:k2]

#list of households
hh  = [collect(1:nhhy), collect(nhhy+1:nhhy+nhho)]

#list of firms
firms = collect(1:nf)

global empcount = zeros(nsec+1)
global unempcount = zeros(nsec)
global shorttimecount = zeros(nsec)
global oldcount = 0

# add young households
for n =1:nhhy
    k1help = 1+fld(mod(n-1,k1*k2),k2)
    k2help = 1+mod(n-1,k1*k2)-(k1help-1)*k2 # distribute evenly on grid
    if rand() < unemp
        emp=0
        tt = conselect(fracemp) # no unemployed for public sector
        unempcount[tt] += 1
    else
        emp = rand(locallf[k1help,k2help])
        push!(workerlist[emp],n)
        if emp <= nfsec[1]
            tt = 1
        end
        for tth=2:nsec+1
            if sum(nfsec[1:tth-1]) < emp && emp <= sum(nfsec[1:tth])
                tt = tth
            end
        end

        empcount[tt] +=1
    end
    if rand() < homeofficefrac[tt] # can this worker do home-office ?
        ho = true
    else
        ho = false
    end
    shopday = rand(0:6,nsec)
    agent = CovidAgent.GenericAgent(n,(k1help,k2help),tt,1,1,false,0,0,emp,false,0,ho,Dict{Int,CovidAgent.GenericAgent}(),Dict{Int,CovidAgent.GenericAgent}(),incomehh[tt],savinghh[tt],consbud,prevconsbud,shopday,shoprepini,0,0,0,0,0,pastrevini, 0,0)

    add_agent_pos!(agent,covidmodel)
    lochh[1,k1help,k2help][n] = agent
    if emp == 0
        unemplist[tt,k1help,k2help][n] = agent
    end
end

# add old households
for n = nhhy+1:nhhy+nhho
    global oldcount += 1
    k1help = 1+fld(mod(n-1,k1*k2),k2)
    k2help = 1+mod(n-1,k1*k2)-(k1help-1)*k2 # distribute evenly on grid
    shopday = rand(0:6,nsec)
    agent = CovidAgent.GenericAgent(n,(k1help,k2help),1,2,1,false,0,0,0,false,0,false,Dict{Int,CovidAgent.GenericAgent}(),Dict{Int,CovidAgent.GenericAgent}(),incomehho,savinghho,consbud,prevconsbud,shopday,shoprepini,0,0,0,0,0,pastrevini,0,0)

    add_agent_pos!(agent,covidmodel)
    lochh[2,k1help,k2help][n] = agent
end

# add firms
#first sector
for n = 1: nfsec[1]
    k1help = 1+fld(mod(n-1,k1*k2),k2)
    k2help = 1+mod(n-1,k1*k2)-(k1help-1)*k2 # distribute evenly on grid
    product =  (prodlb[1]+ (produb[1]-prodlb[1])*rand())*baseprod[1] # determine productivity
    shopday = [0,0,0]

    worker = Dict{Int,CovidAgent.GenericAgent}()
    for id in workerlist[n]
        worker[id] = covidmodel[id]
    end

    agent = CovidAgent.GenericAgent(nhhy+nhho+n,(k1help,k2help),nsec+2,1,1,false,0,0,0,false,0,false,worker,Dict{Int,CovidAgent.GenericAgent}(),0,savingf,consbud,prevconsbud,shopday,shoprepini,product,0,demand[1],demandexpini[1],0,pastrevini,etmin[1],0)
    add_agent_pos!(agent,covidmodel)
end
#other sectors
for tt = 2:nsec+1
    for n = sum(nfsec[1:tt-1])+1: sum(nfsec[1:tt])
        k1help = 1+fld(mod(n-1,k1*k2),k2)
        k2help = 1+mod(n-1,k1*k2)-(k1help-1)*k2 # distribute evenly on grid
        product =  (prodlb[tt]+ (produb[tt]-prodlb[tt])*rand())*baseprod[tt] # determine productivity
        shopday = [0,0,0]

        worker = Dict{Int,CovidAgent.GenericAgent}()
        for id in workerlist[n]
            worker[id] = covidmodel[id]
        end

        agent = CovidAgent.GenericAgent(nhhy+nhho+n,(k1help,k2help),nsec+1+tt,1,1,false,0,0,0,false,0,false,worker,Dict{Int,CovidAgent.GenericAgent}(),0,savingf,consbud,prevconsbud,shopday,shoprepini,product,0,demand[tt],demandexpini[tt],0,pastrevini, etmin[tt], 0)
        add_agent_pos!(agent,covidmodel)
    end
end

divperhh = 0
weeklyconsumption = zeros(nsec, k1, k2)

covidmodel_copy = copy(covidmodel.agents)