module PipsNlpInterface

#Import Packages
using SparseArrays
using LinearAlgebra
using DataStructures
import MPI

import JuMP
import MathOptInterface
const MOI = MathOptInterface
using Plasmo

include("PipsNlpCInterface.jl")
using .PipsNlpSolver

include("pips_utils.jl")
include("nlp_evaluator.jl")

export pipsnlp_solve

mutable struct PIPSNLPData
    d::Union{JuMP.NLPEvaluator,Nothing}
    n::Int
    m::Int
    local_m::Int                    #rank row dimension
    jacnnz::Int                     #Jacobian num nonzeros
    hessnnz::Int                    #Hessian num nonzeros
    firstIeq::Vector{Int}           #first stage equality constraint jacobian nonzero rows
    firstJeq::Vector{Int}           #first stage equality constraint jacobian nonzero cols
    firstVeq::Vector{Float64}       #first stage equality constraint jacobian values
    secondIeq::Vector{Int}          #second stage equality constraint jacobian nonzero rows
    secondJeq::Vector{Int}          #second stage equality constraint jacobian nonzero cols
    secondVeq::Vector{Float64}      #second stage equality constraint jacobian values
    firstIineq::Vector{Int}         #first stage inequality constraint jacobian nonzero rows
    firstJineq::Vector{Int}
    firstVineq::Vector{Float64}
    secondIineq::Vector{Int}
    secondJineq::Vector{Int}
    secondVineq::Vector{Float64}
    num_eqconnect::Int              #num equalities connecting first to second stage
    num_ineqconnect::Int            #num inequalities connecting first to second stage
    eqconnect_lb::Vector{Float64}   #connection equality lower bounds
    eqconnect_ub::Vector{Float64}   #connection equality upper bounds
    ineqconnect_lb::Vector{Float64} #connection inequality lower bounds
    ineqconnect_ub::Vector{Float64} #connection inequality upper bounds
    eq_idx::Vector{Int}             #equality connect indices
    ineq_idx::Vector{Int}           #inequality connect indices
    firstJeqmat
    secondJeqmat
    firstJineqmat
    secondJineqmat
    linkIeq::Vector{Int}            #equality link jacobian rows
    linkJeq::Vector{Int}            #equality link jacobian cols
    linkVeq::Vector{Float64}        #equality link jacobian values
    linkIineq::Vector{Int}
    linkJineq::Vector{Int}
    linkVineq::Vector{Float64}
    x_sol::Vector{Float64}          #solution
    coreid::Int                     #rank
    loaded::Bool
    local_unsym_hessnnz::Int        #local hessian num nonzeros
end
PIPSNLPData() = PIPSNLPData(nothing,0,0,0,0,0,Int[],Int[], Float64[], Int[], Int[], Float64[],Int[],Int[],Float64[], Int[], Int[],
                            Float64[], 0, 0, Float64[],Float64[],Float64[],Float64[], Int[], Int[],
                            nothing, nothing, nothing, nothing, Int[],Int[], Float64[], Int[], Int[], Float64[],Float64[], 0, false, 0)

function _get_pips_data(node::OptiNode)
    return node.ext[:pips_data]
end

function pipsnlp_solve(graph::OptiGraph) #Assume graph variables and constraints are first stage
    #Check structure: 3 possibilities
    #1.) optigraph without subgraphs (only link constraints, no first stage -> second stage constraints)
    #2.) optigraph with single subgraph (first stage and linkconstraints between optinodes)
    #3.) optigraph with multiple subgraphs (first stage and linkconstraints between subgraphs) Not yet supported

    if !(has_subgraphs(graph))
        first_stage = OptiNode() #empty first stage
        submodels = all_nodes(graph)
    elseif length(graph.subgraphs) == 1
        num_nodes(graph) == 1 || error("Incompatible optigraph structure.  Optigraph contains a single subgraph, but has multiple optinodes in highest level graph.")
        first_stage = getnode(graph,1)
        submodels = all_nodes(graph.subgraphs[1])
    end

    comm = MPI.COMM_WORLD
    if MPI.Comm_rank(comm) == 0
        println("Building Model for PIPS-NLP")
    end

    n_sub_models = length(submodels)

    model_list = [first_stage; submodels]

    for (idx,node) in enumerate(model_list)
        node.ext[:pips_data] = PIPSNLPData()
        constraint_data = get_constraint_data(node)
        node.ext[:constraint_data] = constraint_data
    end
    first_stage_data = _get_pips_data(first_stage)

    #Assume every worker has every linkconstraint that connects across nodes
    #TODO: MPI.ALLGATHER! to setup linkconstraints when they are different across workers
    #TODO: Check whether workers can use zeros for bounds they don't need.
    #TODO: Improve distribute performance.

    #IDENTIFY LINK CONSTRAINTS
    link_connect_eq,link_connect_ineq,linkeqconstraints,linkineqconstraints  = _identify_linkconstraints(graph)

    #FIRST STAGE -> SECOND STAGE LINK CONSTRAINTS
    for link in link_connect_eq
        node = _get_subnode(first_stage,link)::OptiNode
        local_data = _get_pips_data(node)
        firstIeq = local_data.firstIeq
        firstJeq = local_data.firstJeq
        firstVeq = local_data.firstVeq
        secondIeq = local_data.secondIeq
        secondJeq = local_data.secondJeq
        secondVeq = local_data.secondVeq
        push!(local_data.eqconnect_lb, link.set.value)
        push!(local_data.eqconnect_ub, link.set.value)
        local_data.num_eqconnect += 1
        row = local_data.num_eqconnect
        for (var,coeff) in link.func.terms
            if getnode(var) == first_stage
                push!(firstIeq, row)
                push!(firstJeq, var.index.value)
                push!(firstVeq, coeff)
            else
                @assert getnode(var) == node
                push!(secondIeq, row)
                push!(secondJeq, var.index.value)
                push!(secondVeq, coeff)
            end
        end
    end

    for link in link_connect_ineq
        node = _get_subnode(first_stage,link)
        firstIineq = local_data.firstIineq
        firstJineq = local_data.firstJineq
        firstVineq = local_data.firstVineq
        secondIineq = local_data.secondIineq
        secondJineq = local_data.secondJineq
        secondVineq = local_data.secondVineq
        if isa(link.set,MOI.LessThan)
            push!(local_data.ineqconnect_lb, -Inf)
            push!(local_data.ineqconnect_ub, link.set.upper)
        elseif isa(link.set,MOI.GreaterThan)
            push!(local_data.ineqconnect_lb, link.set.lower)
            push!(local_data.ineqconnect_ub, Inf)
        elseif isa(link.set,MOI.Interval)
            push!(local_data.ineqconnect_lb, link.set.lower)
            push!(local_data.ineqconnect_ub, link.set.upper)
        end
        local_data.num_ineqconnect += 1
        row = local_data.num_ineqconnect
        for (var,coeff) in link.func.terms
            if getnode(var) == first_stage
                push!(firstIineq, row)
                push!(firstJineq, var.index.value)
                push!(firstVineq, ind)
            else
                @assert getnode(var) == node
                push!(secondIineq, row)
                push!(secondJineq, var.index.value)
                push!(secondVineq, ind)
            end
        end
    end

    #SECOND STAGE LINK CONSTRAINTS
    nlinkeq = length(linkeqconstraints)
    nlinkineq = length(linkineqconstraints)
    ineqlink_lb = zeros(nlinkineq)
    ineqlink_ub = zeros(nlinkineq)
    eqlink_lb = zeros(nlinkeq)
    eqlink_ub = zeros(nlinkeq)

    #INEQUALITY CONSTRAINTS
    #Inequality bounds
    #NOTE: linkconstraint bounds need to be on every rank.
    for (idx,link) in enumerate(linkineqconstraints)
        row = idx
        if isa(link.set,MOI.LessThan)
            ineqlink_lb[row] = -Inf
            ineqlink_ub[row] = link.set.upper
        elseif isa(link.set,MOI.GreaterThan)
            ineqlink_lb[row] = link.set.lower
            ineqlink_ub[row] = Inf
        elseif isa(link.set,MOI.Interval)
            ineqlink_lb[row] = link.set.lower
            ineqlink_ub[row] = link.set.upper
        end
        #Inequality values
        for (var,coeff) in link.func.terms
            node = getnode(var)
            local_data = _get_pips_data(node)
            linkIineq = local_data.linkIineq
            linkJineq = local_data.linkJineq
            linkVineq = local_data.linkVineq
            push!(linkIineq, row)
            push!(linkJineq, var.index.value)
            push!(linkVineq, coeff)
        end
    end

    #EQUALITY CONSTRAINTS
    for (idx,link) in enumerate(linkeqconstraints)
        row = idx
        #Equality bounds
        eqlink_lb[row] = link.set.value
        eqlink_ub[row] = link.set.value

        #Equality values
        for (var,coeff) in link.func.terms
            node = getnode(var)
            local_data = _get_pips_data(node)
            linkIeq = local_data.linkIeq
            linkJeq = local_data.linkJeq
            linkVeq = local_data.linkVeq
            push!(linkIeq, row)              #the variable row
            push!(linkJeq, var.index.value)  #the variable column
            push!(linkVeq, coeff)            #the coefficient
        end
    end

    #Add constraint data for each node to help with function evaluations

    #Initialize
    function str_init_x0(nodeid, x0)
        node = model_list[nodeid+1]
        jump_initval = copy(JuMP.start_value.(JuMP.all_variables(node)))
        local_initval = Vector{Float64}(undef,JuMP.num_variables(node))

        nothing_values = isa.(jump_initval,Nothing)
        float_values = .!(nothing_values)
        local_initval[nothing_values] .= 0  #set to 1 by default
        local_initval[float_values] .= jump_initval[float_values]
        original_copy(local_initval,x0)
    end

    #Problem Info
    function str_prob_info(nodeid,flag,mode,col_lb,col_ub,row_lb,row_ub)
        if flag != 1
            node = model_list[nodeid+1]
            local_data = _get_pips_data(node)
            #Do a warm start on the node
            if !(local_data.loaded)
                local_data.loaded = true
                if (nodeid > 0)
                    if haskey(node.ext, :warmStart)
    	                if node.ext[:warmStart] == true
              	      	   solve(node)
                   	   	end
                    end
            	end

    			#nlp_lb, nlp_ub = JuMP.constraintbounds(node)  #This is every constraint in the model
                nlp_lb, nlp_ub = constraintbounds(node.ext[:constraint_data])       #This is every constraint in the model
         		local_data.local_m  = length(nlp_lb)          #number of local constraints (rows)

    			newRowId = Array{Int}(undef,local_data.local_m)
    			eqId = 1
    			ineqId = 1

                #Go through all constraints, check if they are equality of inequality
                for c in 1:local_data.local_m
                    if nlp_lb[c] == nlp_ub[c]
                        push!(local_data.eq_idx, c)
                        newRowId[c] = eqId
                        eqId +=  1
                    else
        				push!(local_data.ineq_idx, c)
        				newRowId[c] = ineqId
        				ineqId += 1
                    end
         		end

                #Local node data
                local_data.m  = local_data.local_m + local_data.num_eqconnect + local_data.num_ineqconnect
                local_data.n = JuMP.num_variables(node)
                local_data.x_sol = zeros(Float64,local_data.n)
                local_data.firstJeqmat = sparse(local_data.firstIeq, local_data.firstJeq, local_data.firstVeq, local_data.num_eqconnect, first_stage_data.n)
                local_data.secondJeqmat = sparse(local_data.secondIeq, local_data.secondJeq, local_data.secondVeq, local_data.num_eqconnect, local_data.n)
                local_data.firstJineqmat = sparse(local_data.firstIineq, local_data.firstJineq, local_data.firstVineq, local_data.num_ineqconnect, first_stage_data.n)
                local_data.secondJineqmat = sparse(local_data.secondIineq, local_data.secondJineq, local_data.secondVineq, local_data.num_ineqconnect, local_data.n)

                if node.nlp_data == nothing
                    JuMP._init_NLP(node)
                end
                local_data.d = JuMP.NLPEvaluator(node)
                MOI.initialize(local_data.d, [:Grad,:Jac, :Hess])
                #Ijac, Jjac = jac_structure(local_data.d)
                jac_structure = pips_jacobian_structure(local_data.d)
                Ijac = [jac_structure[i][1] for i = 1:length(jac_structure)]
                Jjac = [jac_structure[j][2] for j = 1:length(jac_structure)]
                Ijaceq = Int[]
                Jjaceq = Int[]
                Ijacineq = Int[]
                Jjacineq = Int[]
                jac_eq_index = Int[]
                jac_ineq_index = Int[]
                for i in 1:length(Ijac)
                    c = Ijac[i]
                    if nlp_lb[c] == nlp_ub[c]
                        modifiedrow = newRowId[c]
                        push!(Ijaceq, modifiedrow)
                        push!(Jjaceq, Jjac[i])
                        push!(jac_eq_index, i)
                    else
                        modifiedrow = newRowId[c]
                        push!(Ijacineq, modifiedrow)
                        push!(Jjacineq, Jjac[i])
                        push!(jac_ineq_index,i)
                    end
                end

        		node.ext[:Ijaceq] = Ijaceq
         		node.ext[:Jjaceq] = Jjaceq
                node.ext[:Ijacineq] = Ijacineq
         		node.ext[:Jjacineq] = Jjacineq
         		node.ext[:jac_eq_index] = jac_eq_index
         		node.ext[:jac_ineq_index] = jac_ineq_index
         		#Ihess, Jhess = hesslag_structure(local_data.d)
                hess_structure = pips_hessian_lagrangian_structure(local_data.d)
                Ihess = [hess_structure[i][1] for i = 1:length(hess_structure)]
                Jhess = [hess_structure[j][2] for j = 1:length(hess_structure)]

         		Hmap = Bool[]
         		node_Hrows = Int[]
         		node_Hcols = Int[]
         		for i in 1:length(Ihess)
            	    if Jhess[i] <= Ihess[i]
               	        push!(node_Hrows, Ihess[i])
                   		push!(node_Hcols, Jhess[i])
                   		push!(Hmap, true)
            	    else
                        push!(Hmap, false)
                    end
                end
         		val = ones(Float64,length(node_Hrows))
         		mat = sparse(node_Hrows,node_Hcols,val, local_data.n, local_data.n)
         		node.ext[:Hrows] = node_Hrows
         		node.ext[:Hcols] = node_Hcols
         		node.ext[:Hmap] = Hmap
         		local_hessnnz = length(mat.rowval)
         		local_data.local_unsym_hessnnz = length(Ihess)
         		local_data.hessnnz = local_hessnnz
		    end

	 	    if mode == :Values
               	#nlp_lb, nlp_ub = JuMP.constraintbounds(node)
                nlp_lb, nlp_ub = constraintbounds(node.ext[:constraint_data])
    			eq_lb=Float64[]
    			eq_ub=Float64[]
    			ineq_lb=Float64[]
    			ineq_ub=Float64[]
    			for i in 1: length(nlp_lb)
    			    if nlp_lb[i] == nlp_ub[i]
    			       push!(eq_lb, nlp_lb[i])
    			       push!(eq_ub, nlp_ub[i])
    			    else
    			       push!(ineq_lb, nlp_lb[i])
    			       push!(ineq_ub, nlp_ub[i])
    			    end
    			end

    			if nodeid !=  0 #0 is the master
    			   eq_lb = [eq_lb; local_data.eqconnect_lb]
    			   eq_ub = [eq_ub; local_data.eqconnect_ub]
    			   ineq_lb = [ineq_lb; local_data.ineqconnect_lb]
    			   ineq_ub = [ineq_ub; local_data.ineqconnect_ub]
    			end

    			original_copy([eq_lb;ineq_lb], row_lb)
    			original_copy([eq_ub;ineq_ub], row_ub)
                colLower = variablelowerbounds(node)
                colUpper = variableupperbounds(node)
                original_copy(colLower, col_lb)
    			original_copy(colUpper, col_ub)
            end
		    return (local_data.n,local_data.m)
		else
		    if mode == :Values
                # println("row_lb: ",row_lb)
                # println("row_ub: ",row_ub)
                # println(eqlink_lb)
                # println(eqlink_ub)
                # println(ineqlink_lb)
                # println(ineqlink_ub)
                original_copy([eqlink_lb; ineqlink_lb], row_lb)
                original_copy([eqlink_ub; ineqlink_ub], row_ub)
            end
            return (0, nlinkeq + nlinkineq)
		end
    end

    #x0 is first stage variable values, x1 is local values
    function str_eval_f(nodeid,x0,x1)
    	node = model_list[nodeid+1] #Julia doesn't start index at 0
        local_data = _get_pips_data(node)
        local_d = _get_pips_data(node).d
        if nodeid ==  0
            local_x = x0
        else
            local_x = x1
        end
        #check objective sign
        #local_scl = (node.objSense == :Min) ? 1.0 : -1.0
        local_scl = (JuMP.objective_sense(node) == MOI.MAX_SENSE) ? -1.0 : 1.0
        f = local_scl*pips_eval_objective(local_d,local_x)
        return f
    end

    function str_eval_g(nodeid,x0,x1,new_eq_g, new_inq_g)

        node = model_list[nodeid+1]
        local_data = _get_pips_data(node)
        local_d = _get_pips_data(node).d
        if nodeid ==  0
            local_x = x0
        else
            local_x = x1
        end
        local_g = Array{Float64}(undef,local_data.local_m)
        pips_eval_constraint(local_d, local_g, local_x)

        #evaluates local constraints and connecting constraints with first stage
        new_eq_g[1:end] = [local_g[local_data.eq_idx]; local_data.firstJeqmat*x0+local_data.secondJeqmat*x1]
        new_inq_g[1:end] = [local_g[local_data.ineq_idx]; local_data.firstJineqmat*x0+local_data.secondJineqmat*x1]
	    return Int32(1)
    end

    function str_write_solution(id::Integer, x::Vector{Float64}, y_eq::Vector{Float64}, y_ieq::Vector{Float64})
        node = model_list[id+1]
        local_data = _get_pips_data(node)
        local_data.x_sol = copy(x)

        #TODO: Grab dual values

        #NOTE: This won't work anymore
        # if id == 0
        #     #Add objective value to node
        #     node.ext[:objective] = str_eval_f(id, x, nothing)
        # else
        #     #node.objVal = str_eval_f(id, nothing, x)
        #     node.ext[:objective] = str_eval_f(id, nothing, x)
        # end
        rank = MPI.Comm_rank(comm)
        local_data.coreid = rank
    end

    function str_eval_grad_f(rowid,colid,x0,x1,new_grad_f)
        #println("eval_grad_f")
        node = model_list[rowid+1]
        if rowid == colid
            local_data = _get_pips_data(node)
            local_d = _get_pips_data(node).d
            if colid ==  0
                local_x = x0
            else
                local_x = x1
            end
            local_grad_f = Array{Float64}(undef,local_data.n)
            pips_eval_objective_gradient(local_d, local_grad_f, local_x)
            #eval_grad_f(local_d, local_grad_f, local_x)
            #local_scl = (node.objSense == :Min) ? 1.0 : -1.0
            local_scl = (JuMP.objective_sense(node) == MOI.MAX_SENSE) ? -1.0 : 1.0
            #scale!(local_grad_f,local_scl)
            rmul!(local_grad_f,local_scl)
            original_copy(local_grad_f, new_grad_f)
        elseif colid == 0
            new_grad_f[1:end] .= 0
        else
            @assert(false)
        end
        return Int32(1)
    end

    #NOTE: Why does it subtract 1?
    function array_copy(src,dest)
        @assert(length(src)==length(dest))
        for i in 1:length(src)
            dest[i] = src[i]-1  #NOTE: Is this copying indices to a C vector?
        end
    end

    function original_copy(src,dest)
        @assert(length(src)==length(dest))
        for i in 1:length(src)
            dest[i] = src[i]
        end
    end

    function str_eval_jac_g(rowid,colid,flag, x0,x1,mode,e_rowidx,e_colptr,e_values,i_rowidx,i_colptr,i_values)
        if flag != 1 #populate parent child structure
            node = model_list[rowid+1]
            local_data = _get_pips_data(node)
            local_m_eq = length(local_data.eq_idx)
            local_m_ineq = length(local_data.ineq_idx)
            if mode == :Structure
                if (rowid == colid)
                    Ieq=[node.ext[:Ijaceq];local_data.secondIeq .+ local_m_eq]
                    Jeq=[node.ext[:Jjaceq];local_data.secondJeq]
                    Veq= ones(Float64, length(Ieq))
                    Iineq=[node.ext[:Ijacineq];local_data.secondIineq .+ local_m_ineq]
                    Jineq=[node.ext[:Jjacineq];local_data.secondJineq]
                    Vineq=ones(Float64, length(Iineq))
                    eqmat = sparse(Ieq, Jeq, Veq, local_m_eq + local_data.num_eqconnect, local_data.n)
                    ineqmat = sparse(Iineq, Jineq, Vineq, local_m_ineq + local_data.num_ineqconnect, local_data.n)
                else
                    eqmat = sparse(local_m_eq .+ local_data.firstIeq, local_data.firstJeq, local_data.firstVeq, local_m_eq .+ local_data.num_eqconnect, first_stage_data.n)
                    ineqmat = sparse(local_m_ineq .+ local_data.firstIineq, local_data.firstJineq, local_data.firstVineq, local_m_ineq .+ local_data.num_ineqconnect, first_stage_data.n)
                end
                return(length(eqmat.rowval), length(ineqmat.rowval))
            else #mode = :Values
                if rowid == colid    #evaluate Jacobian for block
                    if colid ==  0   #if it's the root node
                        local_x = x0
                    else
                        local_x = x1
                    end
                    local_values = Array{Float64}(undef,length(node.ext[:Ijaceq])+length(node.ext[:Ijacineq]))
                    #eval_jac_g(local_data.d, local_values, local_x)
                    pips_eval_constraint_jacobian(local_data.d, local_values, local_x)
                    jac_eq_index = node.ext[:jac_eq_index]
                    jac_ineq_index = node.ext[:jac_ineq_index]
                    Ieq=[node.ext[:Ijaceq];local_data.secondIeq .+ local_m_eq]
                    Jeq=[node.ext[:Jjaceq];local_data.secondJeq]
                    Veq=[local_values[jac_eq_index];local_data.secondVeq]
                    Iineq=[node.ext[:Ijacineq];local_data.secondIineq .+ local_m_ineq]
                    Jineq=[node.ext[:Jjacineq];local_data.secondJineq]
                    Vineq=[local_values[jac_ineq_index];local_data.secondVineq]
                    eqmat = sparseKeepZero(Ieq, Jeq, Veq, local_m_eq + local_data.num_eqconnect, local_data.n)
                    ineqmat = sparseKeepZero(Iineq, Jineq, Vineq, local_m_ineq + local_data.num_ineqconnect, local_data.n)
                else #evaluate
                    eqmat = sparseKeepZero(local_m_eq .+ local_data.firstIeq, local_data.firstJeq, local_data.firstVeq, local_m_eq .+ local_data.num_eqconnect, first_stage_data.n)
                    ineqmat = sparseKeepZero(local_m_ineq .+ local_data.firstIineq, local_data.firstJineq, local_data.firstVineq, local_m_ineq .+ local_data.num_ineqconnect, first_stage_data.n)
                end
                if length(eqmat.nzval) > 0
                    array_copy(eqmat.rowval,e_rowidx)
                    array_copy(eqmat.colptr,e_colptr)
                    original_copy(eqmat.nzval,e_values)
                end
                if length(ineqmat.nzval) > 0
                    array_copy(ineqmat.rowval,i_rowidx)
                    array_copy(ineqmat.colptr,i_colptr)
                    original_copy(ineqmat.nzval, i_values)
                end
            end
        else #populate linkconstraint structure
            node = model_list[rowid+1]
            local_data = _get_pips_data(node)
            linkIeq = local_data.linkIeq
            linkJeq = local_data.linkJeq
            linkVeq = local_data.linkVeq
            linkIineq = local_data.linkIineq
            linkJineq = local_data.linkJineq
            linkVineq = local_data.linkVineq

            if mode == :Structure
                return(length(linkVeq), length(linkVineq))
            else
                eqmat = sparse(linkIeq, linkJeq, linkVeq, nlinkeq, local_data.n)
                ineqmat = sparse(linkIineq, linkJineq, linkVineq, nlinkineq, local_data.n)
                if length(eqmat.nzval) > 0
                    array_copy(eqmat.rowval,e_rowidx)
                    array_copy(eqmat.colptr,e_colptr)
                    original_copy(eqmat.nzval,e_values)
                end
                if length(ineqmat.nzval) > 0
                    array_copy(ineqmat.rowval,i_rowidx)
                    array_copy(ineqmat.colptr,i_colptr)
                    original_copy(ineqmat.nzval, i_values)
                end
            end
        end
        return Int32(1)
    end

    function str_eval_h(rowid,colid,x0,x1,obj_factor,lambda,mode,rowidx,colptr,values)
        #println("eval_h")
        node = model_list[colid+1]
        local_data = _get_pips_data(node)
        if mode == :Structure
            if rowid == colid
                node_Hrows  = node.ext[:Hrows]
                node_Hcols = node.ext[:Hcols]
                node_val = ones(Float64,length(node_Hrows))
                mat = sparseKeepZero(node_Hrows,node_Hcols,node_val, local_data.n, local_data.n)
                return length(mat.rowval)
            elseif colid == 0
                node_Hrows  = node.ext[:Hrows]
                node_Hcols = node.ext[:Hcols]
                node_Hrows,node_Hcols = exchange(node_Hrows,node_Hcols)
                node_val = ones(Float64,length(node_Hrows))
                mat = sparseKeepZero(node_Hrows,node_Hcols,node_val, local_data.n, local_data.n)
                return length(mat.rowval)
            else
                return 0
            end
        else #fill in values
            if rowid == colid
                node_Hmap = node.ext[:Hmap]
                node_Hrows  = node.ext[:Hrows]
                node_Hcols = node.ext[:Hcols]
                if colid ==  0
                    local_x = x0
                else
                    local_x = x1
                end
                local_unsym_values = Array{Float64}(undef,local_data.local_unsym_hessnnz)
                node_val = ones(Float64,length(node_Hrows))
                #local_scl = (node.objSense == :Min) ? 1.0 : -1.0
                local_scl = (JuMP.objective_sense(node) == MOI.MAX_SENSE) ? -1.0 : 1.0
                local_m_eq = length(local_data.eq_idx)
                local_m_ineq = length(local_data.ineq_idx)
                local_lambda = zeros(Float64, local_data.local_m)
                for i in 1:local_m_eq
                    local_lambda[local_data.eq_idx[i]] = lambda[i]
                end
                for i in 1:local_m_ineq
                    local_lambda[local_data.ineq_idx[i]] = lambda[i+local_m_eq]
                end

                #eval_hesslag(local_data.d, local_unsym_values, local_x, obj_factor*local_scl, local_lambda)
                pips_eval_hessian_lagrangian(local_data.d, local_unsym_values, local_x, obj_factor*local_scl, local_lambda)
                local_sym_index=1
                for i in 1:local_data.local_unsym_hessnnz
                    if node_Hmap[i]
                        node_val[local_sym_index] = local_unsym_values[i]
                        local_sym_index +=1
                    end
                end
                node_Hrows,node_Hcols = exchange(node_Hrows,node_Hcols)
                mat = sparseKeepZero(node_Hrows,node_Hcols,node_val, local_data.n, local_data.n)
                array_copy(mat.rowval,rowidx)
                array_copy(mat.colptr,colptr)
                original_copy(mat.nzval,values)
            elseif colid ==	0
                node_Hrows  = node.ext[:Hrows]
                node_Hcols = node.ext[:Hcols]
                node_val = zeros(Float64,length(node_Hrows))
                node_Hrows,node_Hcols = exchange(node_Hrows,node_Hcols)
                mat = sparseKeepZero(node_Hrows,node_Hcols,node_val, local_data.n, local_data.n)
                array_copy(mat.rowval,rowidx)
                array_copy(mat.colptr,colptr)
                original_copy(mat.nzval,values)
            else
            end
        end
        return Int32(1)
    end

    if !MPI.Initialized()
        MPI.Init()
    end

    comm = MPI.COMM_WORLD
    if (MPI.Comm_rank(comm) == 0)
        t1 = time()
    end

    root = 0
    rank = MPI.Comm_rank(comm)

    #Create PipsModel (The PIPS interface model) and pass all the functions it requires
    model = PipsModel(:Min,0, n_sub_models,str_init_x0, str_prob_info, str_eval_f, str_eval_g, str_eval_grad_f, str_eval_jac_g, str_eval_h,str_write_solution)

    ###################################
    #PIPS-NLP FUNCTIONS
    prob = createProblemStruct(comm, model, true)
    if rank == 0
        println("Created PIPS-NLP Problem Struct")
    end

    MPI.Barrier(comm)
    ret = solveProblemStruct(prob)
    ###################################

    if rank == 0
        println("Solved PIPS-NLP Problem Struct")
    end

    if rank == 0
        println("Timing Results:")
        println("init_x0  ",prob.t_jl_init_x0)
        println("str_init_x0  ", prob.t_jl_str_prob_info)
        println("eval_f  ", prob.t_jl_eval_f)
        println("eval_g0  ",prob.t_jl_eval_g)
        println("eval_grad_f  ", prob.t_jl_eval_grad_f)
        println("eval_jac  ", prob.t_jl_eval_jac_g)
        println("str_eval_jac  ", prob.t_jl_str_eval_jac_g)
        println("eval_h  ",  prob.t_jl_eval_h)
        println("str_eval_h ",  prob.t_jl_str_eval_h)
        println("eval_write_solution  ",  prob.t_jl_write_solution)
        println("PIPS-NLP time:   ",  time() - t1, " (s)")
    end

    #TODO.  Put solution onto optinodes
    for (idx,node) in enumerate(model_list)  #set solution values for each model
        local_data = _get_pips_data(node)
        if idx != 1 #all cores have the first stage
            coreid = zeros(Int, 1)
            sc = MPI.Reduce(local_data.coreid, MPI.SUM, root, comm)

            if rank == root
                coreid[1] = sc
            end

            MPI.Bcast!(coreid, length(coreid), root, comm)

            n = zeros(Int, 1)
            n[1] = local_data.n
            MPI.Bcast!(n,length(n),coreid[1], comm)

            if rank != coreid[1]
                local_data.n = n[1]
                local_data.x_sol = zeros(Float64,local_data.n)
            end
            MPI.Bcast!(local_data.x_sol, local_data.n, coreid[1], comm)
        end
        #node.ext[:colVal] = local_data.x_sol
        vars = MOI.get(node,MOI.ListOfVariableIndices())
        primals = OrderedDict(zip(vars,local_data.x_sol))

        id = graph.id
        src = JuMP.backend(node)
        src.primals[id] = primals
        src.last_solution_id = id

        #TODO set duals

        #TODO: Load results into a graph backend with termination status etc...
    end

    #TODO: Better return status.  Hook up a simple MOI interface that drops the model results into a graph backend CachingOptimizer
    status = :Unknown
    if ret == 0
        status = :Optimal
    elseif ret == 1
        status = :Not_Finished
    elseif	ret == 2
        status = :Maximum_Iterations_Exceeded
    elseif	ret == 3
        status = :Infeasible_Problem_Detected
    elseif ret == 4
        status = :Restoration_needed
    else
        status = :UnknownStatus
    end

    return status
end  #end pips_nlp_solve

end #end module


include("DataTypes.jl")

WGS84_wkt = "GEOGCS[\"WGS 84\",DATUM[\"WGS_1984\",SPHEROID[\"WGS 84\",6378137,298.257223563,AUTHORITY[\"EPSG\",\"7030\"]],AUTHORITY[\"EPSG\",\"6326\"]],PRIMEM[\"Greenwich\",0,AUTHORITY[\"EPSG\",\"8901\"]],UNIT[\"degree\",0.0174532925199433,AUTHORITY[\"EPSG\",\"9122\"]],AXIS[\"Latitude\",NORTH],AXIS[\"Longitude\",EAST],AUTHORITY[\"EPSG\",\"4326\"]]"

"""
    GeoArray(A::AbstractArray{T,3} where T <: Union{Real,Union{Missing,Real}})

Construct a GeoArray from any Array. A default `AffineMap` and `CRS` will be generated.

# Examples
```julia-repl
julia> GeoArray(rand(10,10,1))
10x10x1 Array{Float64, 3} with AffineMap([1.0 0.0; 0.0 1.0], [0.0, 0.0]) and undefined CRS
```
"""
GeoArray(A::AbstractArray{T, 3}) where T<:Union{Real, Union{Missing, Real}} =
    GeoArray(A, geotransform_to_affine(SVector(0.,1.,0.,0.,0.,1.)), "")

GeoArray(A::AbstractArray{T, 3}, f::AffineMap) where T<:Union{Real, Union{Missing, Real}} =
    GeoArray(A, f, GFT.WellKnownText(GFT.CRS(), ""))

GeoArray(A::AbstractArray{T, 3}, f::AffineMap, crs::String) where T<:Union{Real, Union{Missing, Real}} = begin
    if crs == ""; crs = WGS84_wkt; end
    GeoArray(A, f, GFT.WellKnownText(GFT.CRS(), crs))
end

function GeoArray(A::AbstractArray{T, 3}, x::AbstractRange, y::AbstractRange, args...) where T<:Union{Real, Union{Missing, Real}}
    size(A)[1:2] != (length(x), length(y)) &&
        error("Size of `GeoArray` $(size(A)) does not match size of (x,y): $((length(x),length(y))). Note that this function takes *center coordinates*.")
    f = unitrange_to_affine(x, y)
    GeoArray(A, f, args...)
end

function GeoArray(A::AbstractArray{T, 3}, b::bbox; proj = 4326) where T<:Union{Real, Union{Missing, Real}}
    ga = GeoArray(A)
    st_bbox!(ga, b)
    epsg!(ga, proj)
    ga
end

GeoArray(A::AbstractArray{T, 2} where T<:Union{Real, Union{Missing, Real}}, args...) =
    GeoArray(reshape(A, size(A)..., 1), args...)

Base.size(ga::GeoArray) = size(ga.A)
Base.IndexStyle(::Type{T}) where {T<:GeoArray} = IndexLinear()
Base.iterate(ga::GeoArray) = iterate(ga.A)
Base.length(ga::GeoArray) = length(ga.A)
Base.parent(ga::GeoArray) = ga.A
# Base.map(f, ga::GeoArray) = GeoArray(map(f, ga.A), ga.f, ga.crs)
# Base.convert(::Type{Array{T, 3}}, A::GeoArray{T}) where {T} = convert(Array{T,3}, ga.A)
Base.eltype(::Type{GeoArray{T}}) where {T} = T

Base.show(io::IO, ::MIME"text/plain", ga::GeoArray) = show(io, ga)
function Base.show(io::IO, ga::GeoArray)
    crs = GeoFormatTypes.val(ga.crs)
    wkt = length(crs) == 0 ? "undefined CRS" : "CRS $crs"
    print(io, "$(join(size(ga), "x")) $(typeof(ga.A)) with $(ga.f) and $(wkt)")
end

include("st_[.jl")
include("st_Ops.jl")


#
# This file is part of module
#   ModiaMedia (ModiaMedia/src/ModiaMedia.jl)
#

### Data structures that are additionally available for media <: PartialMixtureMedium -----------------------

"""
    state = ThermodynamicState_pTX(p,T,X)

Generate a `ThermodynamicState_pT <: MixtureThermodynamicState` object containing
pressure `p` [Pa], temperature `T` [K], and a vector of mass fractions `X` as states.
"""
mutable struct ThermodynamicState_pTX <: MixtureThermodynamicState
    p::Float64
    T::Float64
    X::Vector{Float64}
end

"R = gas_constant(medium,state) - return gas constant for `medium` from `state` in [J/mol.K]"
# should this be MixtureMedium?
# should this take MixtureThermodynamicState?
gas_constant(m::AbstractMedium, state::ThermodynamicState) = undefinedFunction("gas constant", m)

# moleToMassFractions: use inner constructor with fluid information.
# massToMoleFractions: use inner constructor with fluid information from abstract medium, since we have that type.


using FixedPointNumbers

@testset "Parse" begin
    r8(x) = reinterpret(N0f8, x)

    # Color parsing
    redN0f8 = parse(Colorant, "red")
    @test colorant"red" == redN0f8
    @test isa(redN0f8, RGB{N0f8})
    @test redN0f8 == RGB(1,0,0)
    @test parse(RGB{Float64}, "red") === RGB{Float64}(1,0,0)
    @test isa(parse(HSV, "blue"), HSV)
    @test parse(Colorant, "rgb(55,217,127)") === RGB{N0f8}(r8(0x37),r8(0xd9),r8(0x7f))
    @test colorant"rgb(55,217,127)" === RGB{N0f8}(r8(0x37),r8(0xd9),r8(0x7f))
    @test parse(Colorant, "rgba(55,217,127,0.5)") === RGBA{N0f8}(r8(0x37),r8(0xd9),r8(0x7f),0.5)
    @test parse(Colorant, "rgb(55,217,127)") === RGB{N0f8}(r8(0x37),r8(0xd9),r8(0x7f))
    @test parse(Colorant, "rgba(55,217,127,0.5)") === RGBA{N0f8}(r8(0x37),r8(0xd9),r8(0x7f),0.5)
    @test parse(Colorant, "hsl(120, 100%, 50%)") === HSL{Float32}(120,1.0,.5)
    @test colorant"hsl(120, 100%, 50%)" === HSL{Float32}(120,1.0,.5)
    @test parse(RGB{N0f8}, "hsl(120, 100%, 50%)") === convert(RGB{N0f8}, HSL{Float32}(120,1.0,.5))
    @test_throws ErrorException  parse(Colorant, "hsl(120, 100, 50)")
    @test parse(Colorant, "#D0FF58") === RGB(r8(0xD0),r8(0xFF),r8(0x58))
    @test parse(Colorant, "#FB0") === RGB(r8(0xFF),r8(0xBB),r8(0x00))

    @test parse(Colorant, :red) === colorant"red"
    @test parse(Colorant, colorant"red") === colorant"red"

end


using CLIMAParameters
using CLIMAParameters.Planet: planet_radius
using ..Diagnostics
using ..Mesh.Interpolation

"""
    DiagnosticsConfiguration

Container for all the `DiagnosticsGroup`s to be used for a simulation.
"""
mutable struct DiagnosticsConfiguration
    groups::Array{DiagnosticsGroup, 1}

    DiagnosticsConfiguration(
        groups::Array{DG, 1},
    ) where {DG <: DiagnosticsGroup} = new(groups)
end

"""
    InterpolationConfiguration(
        driver_config::DriverConfiguration,
        boundaries::Array,
        resolution::Tuple,
    )

Creates an `InterpolationTopology` (either an `InterpolationBrick` or an
`InterpolationCubedSphere`) to be used with a `DiagnosticsGroup`. The axes
are set up based on `boundaries` and `resolution`.
"""
function InterpolationConfiguration(
    driver_config::DriverConfiguration,
    boundaries::Array,
    resolution::Tuple,
)
    param_set = driver_config.bl.param_set
    grid = driver_config.grid
    if isa(grid.topology, StackedBrickTopology)

        axes = (
            collect(range(
                boundaries[1, 1],
                boundaries[2, 1],
                step = resolution[1],
            )),
            collect(range(
                boundaries[1, 2],
                boundaries[2, 2],
                step = resolution[2],
            )),
            collect(range(
                boundaries[1, 3],
                boundaries[2, 3],
                step = resolution[3],
            )),
        )
        return InterpolationBrick(grid, boundaries, axes[1], axes[2], axes[3])

    elseif isa(grid.topology, StackedCubedSphereTopology)

        FT = eltype(grid)
        _planet_radius::FT = planet_radius(param_set)
        info = driver_config.config_info
        vert_range = grid1d(
            _planet_radius,
            FT(_planet_radius + info.domain_height),
            nelem = info.nelem_vert,
        )
        axes = (
            collect(range(
                boundaries[1, 1],
                boundaries[2, 1],
                step = resolution[1],
            )),
            collect(range(
                boundaries[1, 2],
                boundaries[2, 2],
                step = resolution[2],
            )),
            collect(range(
                boundaries[1, 3],
                boundaries[2, 3],
                step = resolution[3],
            )),
        )
        return InterpolationCubedSphere(
            grid,
            collect(vert_range),
            info.nelem_horz,
            axes[1],
            axes[2],
            axes[3],
        )
    else
        @error "Cannot set up interpolation for this topology."
    end
end

"""
    InterpolationConfiguration(
        driver_config::DriverConfiguration,
        boundaries::Array,
        axes,
    )

Creates an `InterpolationTopology` (either an `InterpolationBrick` or an
`InterpolationCubedSphere`) to be used with a `DiagnosticsGroup`. The axes
are directly specified in lat/lon/lvl order.
"""
function InterpolationConfiguration(
    driver_config::DriverConfiguration,
    boundaries::Array,
    axes,
)
    param_set = driver_config.bl.param_set
    grid = driver_config.grid
    if isa(grid.topology, StackedBrickTopology)
        return InterpolationBrick(grid, boundaries, axes[1], axes[2], axes[3])

    elseif isa(grid.topology, StackedCubedSphereTopology)

        FT = eltype(grid)
        _planet_radius::FT = planet_radius(param_set)
        info = driver_config.config_info
        vert_range = grid1d(
            _planet_radius,
            FT(_planet_radius + info.domain_height),
            nelem = info.nelem_vert,
        )

        return InterpolationCubedSphere(
            grid,
            collect(vert_range),
            info.nelem_horz,
            axes[1],
            axes[2],
            axes[3],
        )
    else
        @error "Cannot set up interpolation for this topology."
    end
end


module Tracking
    using DocStringExtensions, GNSSSignals, StaticArrays, TrackingLoopFilters
    using Unitful: upreferred, Hz, dBHz, ms
    import Base.zero, Base.length

    export
        get_early,
        get_prompt,
        get_late,
        get_correlator,
        get_carrier_doppler,
        get_carrier_phase,
        get_code_doppler,
        get_code_phase,
        get_early_late_sample_shift,
        get_secondary_code_or_bit_found,
        get_state,
        get_cn0,
        track,
        TrackingState,
        NumAnts,
        MomentsCN0Estimator,
        AbstractCN0Estimator,
        get_bits,
        get_num_bits,
        EarlyPromptLateCorrelator,
        VeryEarlyPromptLateCorrelator,
        SecondaryCodeOrBitDetector

    struct NumAnts{x}
    end

    NumAnts(x) = NumAnts{x}()

    include("cn0_estimation.jl")
    include("discriminators.jl")
    include("bit_buffer.jl")
    include("correlator.jl")
    include("secondary_code_or_bit_detector.jl")
    include("tracking_state.jl")
    include("tracking_results.jl")
    include("tracking_loop.jl")
    include("gpsl1.jl")
    include("gpsl5.jl")
    include("galileo_e1b.jl")
end


# currently monster of a spaghetti code mess, multiple types interacting at the same
# time. Safest at this point in development is to just get this code running and
# will refactor once the dust has settled.
# current code is the result of several unit tests between IIF and RoME.jl
# a unified test to follow -- after which refactoring can start
function evalPotentialSpecific(
      fnc::T,
      Xi::Vector{Graphs.ExVertex},
      gwp::GenericWrapParam{T},
      solvefor::Int;
      N::Int=100  ) where {T <: FunctorPairwise}
  #
  # TODO -- enable partial constraints

  # TODO -- this part can be collapsed into common generic solver component, could be constructed and maintained at addFactor! time
  ARR = Array{Array{Float64,2},1}()
  maxlen, sfidx = prepareparamsarray!(ARR, Xi, N, solvefor)
  gwp.params = ARR
  gwp.varidx = sfidx
  gwp.measurement = gwp.samplerfnc(gwp.usrfnc!, maxlen)
  zDim = size(gwp.measurement[1],1) # TODO -- zDim aspect desperately needs to be redone
  if gwp.specialzDim
    zDim = gwp.usrfnc!.zDim[sfidx]
  end
  # gwp.zDim[sfidx] ??
  fr = FastRootGenericWrapParam{T}(gwp.params[sfidx], zDim, gwp)
  # and return complete fr/gwp
  # TODO -- once Threads.@threads have been optmized JuliaLang/julia#19967, also see area4 branch
  for n in 1:maxlen
    gwp.particleidx = n
    # gwp(x, res)
    numericRootGenericRandomizedFnc!( fr )
        # r = nlsolve( gwp, ARR[sfidx][:,gwp.particleidx] )
        # remember this is a deepcopy of original sfidx, since we are generating a proposal distribution
        # and not directly replacing the existing variable belief estimate
        # gwp.params[gwp.varidx][:,gwp.particleidx] = r.zero[:]
  end

  return gwp.params[gwp.varidx]
  # return evalPotential(typ, Xi, solvefor)
end

function evalPotentialSpecific(
      fnc::T,
      Xi::Vector{Graphs.ExVertex},
      gwp::GenericWrapParam{T},
      solvefor::Int;
      N::Int=100,
      spreadfactor::Float64=10.0  ) where {T <: FunctorPairwiseNH}
  #
  # TODO -- enable partial constraints

  # TODO -- this part can be collapsed into common generic solver component, could be constructed and maintained at addFactor! time
  ARR = Array{Array{Float64,2},1}()
  maxlen, sfidx = prepareparamsarray!(ARR, Xi, N, solvefor)
  gwp.params = ARR
  gwp.varidx = sfidx
  gwp.measurement = gwp.samplerfnc(gwp.usrfnc!, maxlen)
  zDim = size(gwp.measurement[1],1) # TODO -- zDim aspect desperately needs to be redone
  if gwp.specialzDim
    zDim = gwp.usrfnc!.zDim[sfidx]
  end
  fr = FastRootGenericWrapParam{T}(gwp.params[sfidx], zDim, gwp)
  # and return complete fr/gwp

  # nullhypothesis
  nhc = rand(gwp.usrfnc!.nullhypothesis, maxlen) - 1
  val = gwp.params[gwp.varidx]
  d = size(val,1)
  var = Base.var(val,2) + 1e-3
  ENT = Distributions.MvNormal(zeros(d), spreadfactor*diagm(var[:]))

  # TODO --  Threads.@threads see area4 branch
  for n in 1:maxlen
    # gwp(x, res)
    if nhc[n] != 0
      gwp.particleidx = n
      numericRootGenericRandomizedFnc!( fr )
    else
      gwp.params[gwp.varidx][:,n] += rand(ENT)
    end
  end

  return gwp.params[gwp.varidx]
  # return evalPotential(typ, Xi, solvefor)
end

function evalPotentialSpecific(
      fnc::T,
      Xi::Vector{Graphs.ExVertex},
      gwp::GenericWrapParam{T},
      solvefor::Int;
      N::Int=100  ) where {T <: FunctorPairwiseMinimize}
  #
  # TODO -- this part can be collapsed into common generic solver component, could be constructed and maintained at addFactor! time
  ARR = Array{Array{Float64,2},1}()
  maxlen, sfidx = prepareparamsarray!(ARR, Xi, N, solvefor)
  gwp.params = ARR
  gwp.varidx = sfidx
  gwp.measurement = gwp.samplerfnc(gwp.usrfnc!, maxlen)
  zDim = size(gwp.measurement[1],1) # TODO -- zDim aspect desperately needs to be redone
  if gwp.specialzDim
    zDim = gwp.usrfnc!.zDim[sfidx]
  end

  fr = FastRootGenericWrapParam{T}(gwp.params[sfidx], zDim, gwp)

  # TODO -- once Threads.@threads have been optmized JuliaLang/julia#19967, also see area4 branch
  for n in 1:maxlen
    gwp.particleidx = n
    # gwp(x, res)
    # implement minimization here
    res = zeros(fr.xDim)
    gg = (x) -> fr.gwp(res, x)
    r = optimize(  gg, fr.X[1:fr.xDim,fr.gwp.particleidx] )
    # TODO -- clearly lots of optmization to be done here
    fr.Y[1:fr.xDim] = r.minimizer
    fr.X[:,fr.gwp.particleidx] = fr.Y

    # error("not implemented yet")
  end
  return gwp.params[gwp.varidx]
end

function evalPotentialSpecific(
      fnc::T,
      Xi::Vector{Graphs.ExVertex},
      generalwrapper::GenericWrapParam{T},
      solvefor::Int;
      N::Int=100  ) where {T <: FunctorSingleton}
  #
  generalwrapper.measurement = generalwrapper.samplerfnc(generalwrapper.usrfnc!, N)
  if !generalwrapper.partial
    return generalwrapper.measurement[1]
  else
    val = deepcopy(getVal(Xi[1]))
    i = 0
    for dimnum in fnc.partial
      i += 1
      val[dimnum,:] = generalwrapper.measurement[1][i,:]
    end
    return val
  end
end

function evalPotentialSpecific(
      fnc::T,
      Xi::Vector{Graphs.ExVertex},
      generalwrapper::GenericWrapParam{T},
      solvefor::Int;
      N::Int=100,
      spreadfactor::Float64=10.0  ) where {T <: FunctorSingletonNH}
  #

  # determine amount share of null hypothesis particles
  generalwrapper.measurement = generalwrapper.samplerfnc(generalwrapper.usrfnc!, N)
  # values of 0 imply null hypothesis
  # generalwrapper.usrfnc!.nullhypothesis::Distributions.Categorical
  nhc = rand(generalwrapper.usrfnc!.nullhypothesis, N) - 1

  val = getVal(Xi[1])
  d = size(val,1)
  var = Base.var(val,2) + 1e-3
  ENT = Distributions.MvNormal(zeros(d), spreadfactor*diagm(var[:]))

  for i in 1:N
    if nhc[i] == 0
      generalwrapper.measurement[1][:,i] = val[:,i] + rand(ENT)
    end
  end
  # TODO -- returning to memory location inside
  return generalwrapper.measurement[1]
end

# function evalPotentialSpecific{T <: FunctorPartialSingleton}(
#       fnc::T,
#       Xi::Vector{Graphs.ExVertex},
#       generalwrapper::GenericWrapParam{T},
#       solvefor::Int;
#       N::Int=100  )
#   #
#   generalwrapper.measurement = generalwrapper.samplerfnc(generalwrapper.usrfnc!, N)
#   return generalwrapper.measurement[1]
# end

# Multiple dispatch occurs internally, resulting in factor graph potential evaluations
function evalFactor2(fgl::FactorGraph, fct::Graphs.ExVertex, solvefor::Int; N::Int=100)
  # return evalPotential(fct.attributes["data"].fnc, solvefor) #evalPotential(fct.attributes["fnc"], solvefor)

  # TODO -- this build up of Xi is excessive and should be reduced
  # could happen at addFactor time
  Xi = Graphs.ExVertex[]
  for id in getData(fct).fncargvID
  # for id in fct.attributes["data"].fncargvID
    push!(Xi, getVert(fgl,id) ) # TODO localapi
    # push!(Xi, dlapi.getvertex(fgl,id))
  end
  fnctype = getData(fct).fnc
  # fnctype = fct.attributes["data"].fnc
  return evalPotentialSpecific(fnctype.usrfnc!, Xi, fnctype, solvefor, N=N)
  # return evalPotentialSpecific(modulefnc, Xi, fnctype, solvefor, N=N)
end

"""
    findRelatedFromPotential(fg, fctvert, varnodeid, N)

Compute proposal belief on varnodeid through fctvert representing some constraint in factor graph.
Always full dimension of variable node, where partial constraints will only influence directed
subset of variable dimensions. Remaining dimensions will keep existing variable values.
"""
function findRelatedFromPotential(fg::FactorGraph, idfct::Graphs.ExVertex, vertid::Int, N::Int) # vert
  # assuming it is properly initialized TODO
  ptsbw = evalFactor2(fg, idfct, vertid, N=N);
  # sum(abs(ptsbw)) < 1e-14 ? error("findRelatedFromPotential -- an input is zero") : nothing  # NOTE -- disable this validation test

  # TODO -- better to upsample before the projection
  Ndim = size(ptsbw,1)
  Npoints = size(ptsbw,2)
  # Assume we only have large particle population sizes, thanks to addNode!
  p = kde!(ptsbw, "lcv")
  if Npoints != N # this is where we control the overall particle set size
      p = resample(p,N)
  end
  return p
end


# updates.jl
#
# Defines the functions that perform the diagonal update, and also
# that build the linked list and operator cluster update

lsize = 0
nullt = (0,0,0) #a null tuple

############################ FUNCTIONS ######################################

#Diagonal update
function DiagonalUpdate()

#define the Metropolis probability as a constant
#https://pitp.phas.ubc.ca/confs/sherbrooke2012/archives/Melko_SSEQMC.pdf
#equation 1.43
P_h = h_x*nSpin/(h_x*nSpin +2.0*J_*nBond) #J=1.0 tested only

    spin_prop = copy(spin_left)  #the propagated spin state
    for i = 1:2*M  #size of the operator list

       if operator_list[i,1] == -2
           spin_prop[operator_list[i,2]] = xor(spin_prop[operator_list[i,2]],1) #spinflip

       else
           flag = false
           while flag == false
               rr = rand()
               if P_h > rr #probability to choose a single-site operator
                   operator_list[i,1] = -1
                   site = rand(1:nSpin)
                   operator_list[i,2] = site
                   flag = true
                   #println(i," site ",site)
               else
                   bond = rand(1:nBond)
                   if spin_prop[bond_spin[bond,1]] == spin_prop[bond_spin[bond,2]] #spins must be the same
                       operator_list[i,1] = bond_spin[bond,1]
                       operator_list[i,2] = bond_spin[bond,2]
                       flag = true
                       #println(i," bond ",bond)
                   end#if
                   #println(P_h," ",rr," bond")
               end
           end #while

        end#if
    end #for

    #DEBUG
    if spin_prop != spin_right  #check the spin propagation for error
        println("Basis state propagation error: DiagonalUpdate")
    end

end #DiagonalUpdate

#############################################################################

#LinkedList
function LinkedList()

    #initialize linked list data structures
    global LinkList = zeros(Int,0)  #needed for cluster update
    global LegType = zeros(Int,0)

    #A diagonal bond operator has non trivial associates for cluster building
    #Associates = zeros((Int,Int,Int),0)
    global Associates = []

    First = zeros(Int,0)  #scope is this function only

    #The first N elements of the linked list are the spins of the LHS basis state
    for i = 1:nSpin
        push!(First,i)
        push!(LinkList,-99) #we don't yet know what these link to
        push!(LegType,spin_left[i])
        push!(Associates,nullt) #no nontrivial associates for train wf spins
    end #i

    spin_prop = copy(spin_left)  #the propagated spin state

    #Now, add the 2M operators to the linked list.  Each has either 2 or 4 legs
    for i = 1:2*M  #size of the operator list

       if operator_list[i,1] == -2  #off-diagonal site operator
           site = operator_list[i,2]
           #lower or left leg
           push!(LinkList,First[site])
           push!(LegType,spin_prop[site]) #the spin of the vertex leg
           spin_prop[site] = xor(spin_prop[site],1) #spinflip
           current_link = size(LinkList)
           LinkList[First[site]] = current_link[1] #completes backwards link
           First[site] = current_link[1] + 1
           push!(Associates,nullt)
           #upper or right leg
           push!(LinkList,-99) #we don't yet know what this links to
           push!(LegType,spin_prop[site]) #the spin of the vertex leg
           push!(Associates,nullt)

       elseif operator_list[i,1] == -1  #diagonal site operator
           site = operator_list[i,2]
           #lower or left leg
           push!(LinkList,First[site])
           push!(LegType,spin_prop[site]) #the spin of the vertex leg
           current_link = size(LinkList)
           LinkList[First[site]] = current_link[1] #completes backwards link
           First[site] = current_link[1] + 1
           push!(Associates,nullt)
           #upper or right leg
           push!(LinkList,-99) #we don't yet know what this links to
           push!(LegType,spin_prop[site]) #the spin of the vertex leg
           push!(Associates,nullt)

       else  #diagonal bond operator
           #lower left
           site1 = operator_list[i,1]
           push!(LinkList,First[site1])
           push!(LegType,spin_prop[site1]) #the spin of the vertex leg
           current_link = size(LinkList)
           LinkList[First[site1]] = current_link[1] #completes backwards link
           First[site1] = current_link[1] + 2
           vertex1 = current_link[1]
           push!(Associates,(vertex1+1,vertex1+2,vertex1+3))
           #lower right
           site2 = operator_list[i,2]
           push!(LinkList,First[site2])
           push!(LegType,spin_prop[site2]) #the spin of the vertex leg
           current_link = size(LinkList)
           LinkList[First[site2]] = current_link[1] #completes backwards link
           First[site2] = current_link[1] + 2
           push!(Associates,(vertex1,vertex1+2,vertex1+3))
           #upper left
           push!(LinkList,-99) #we don't yet know what this links to
           push!(LegType,spin_prop[site1]) #the spin of the vertex leg
           push!(Associates,(vertex1,vertex1+1,vertex1+3))
           #upper right
           push!(LinkList,-99) #we don't yet know what this links to
           push!(LegType,spin_prop[site2]) #the spin of the vertex leg
           push!(Associates,(vertex1,vertex1+1,vertex1+2))

       end #if

    end #i

    #The last N elements of the linked list are the final spin state
    for i = 1:nSpin
        push!(LinkList,First[i])
        push!(LegType,spin_prop[i])
        current_link = size(LinkList)
        LinkList[First[i]] = current_link[1]
        push!(Associates,nullt)
    end #i

    global lsize = size(LinkList)

    #DEBUG
    if spin_prop != spin_right
        println("Basis state propagation error: LINKED LIST")
    end

end #LinkedList

#############################################################################

#ClusterUpdate
function ClusterUpdate()

    #lsize is the size of the linked list
    in_cluster=zeros(Int,lsize[1])

    cstack = zeros(Int,0)  #This is the stack of vertices in a cluster

    ccount = 0 #cluster number counter
    for i = 1:lsize[1]

        #Add a new leg onto the cluster
        if (in_cluster[i] == 0 && Associates[i] == nullt)

            ccount+=1
            push!(cstack,i)
            in_cluster[cstack[end]] = ccount

            flip = rand(Bool) #flip a coin for the SW cluster flip
            if flip == true
                LegType[cstack[end]] =  xor(LegType[cstack[end]],1) #spinflip
            end

            while isempty(cstack) == false

                leg = LinkList[cstack[end]]
                pop!(cstack)
                #println("leg ",leg," ",cstack)

                if in_cluster[leg] == 0

                    in_cluster[leg] = ccount; #add the new leg and flip it
                    if flip == true
                        #println("gonna flip ",leg)
                        LegType[leg] =  xor(LegType[leg],1)
                    end
                    #now check all associates and add to cluster
                    assoc = Associates[leg] #a 3-element array
                    if assoc != nullt
                        push!(cstack,assoc[1])
                        in_cluster[assoc[1]] = ccount
                        push!(cstack,assoc[2])
                        in_cluster[assoc[2]] = ccount
                        push!(cstack,assoc[3])
                        in_cluster[assoc[3]] = ccount
                        if flip == true
                            LegType[assoc[1]] =  xor(LegType[assoc[1]],1)
                            LegType[assoc[2]] =  xor(LegType[assoc[2]],1)
                            LegType[assoc[3]] =  xor(LegType[assoc[3]],1)
                        end
                    end #if
                end #if in_cluster == 0
            end #while
        end #if
    end #for i

    #DEBUG
    #for i = 1:lsize[1]
    #    println(i," ",LegType[i]," ",in_cluster[i])
    #end

    #map back basis states and operator list
    ocount = 0
    for i = 1:nSpin
        spin_left[i] = LegType[i]  #left basis state
        ocount += 1
    end

    ocount += 1  #next on is leg nSpin + 1
    for i = 1:2*M
        if operator_list[i,1] != -2 && operator_list[i,1] != -1
            ocount += 4
        else
            if LegType[ocount] == LegType[ocount + 1]  #diagonal
                operator_list[i,1] = -1
                #println("DCHANGE ",i," ",ocount)
            else
                operator_list[i,1] = -2 #off-diagonal
                #println("OCHANGE ",i," ",ocount)
            end
            ocount += 2
        end
    end

    for i = 1:nSpin
        spin_right[i] = LegType[lsize[1] - nSpin + i]  #left basis state
    end

end #ClusterUpdate


export DataVariable, datavar, isconnected, update!

mutable struct DataVariable{D, S} <: AbstractVariable
    name       :: Symbol
    messageout :: S
    nconnected :: Int
    marginal   :: Union{Nothing, MarginalObservable}
end

function datavar(name::Symbol, ::Type{D}; subject::S = RecentSubject(Union{Message{Missing}, Message{D}})) where { S, D }
    return DataVariable{D, S}(name, subject, 0, nothing)
end

function datavar(name::Symbol, ::Type{D}, dims::Tuple; subject::S = RecentSubject(Union{Message{Missing}, Message{D}})) where { S, D }
    return datavar(name, D, dims...; subject = subject)
end

function datavar(name::Symbol, ::Type{D}, dims::Vararg{Int}; subject::S = RecentSubject(Union{Message{Missing}, Message{D}})) where { S, D }
    vars = Array{DataVariable{D, S}}(undef, dims)
    for i in CartesianIndices(axes(vars))
        @inbounds vars[i] = datavar(Symbol(name, :_, Symbol(join(i.I, :_))), D; subject = similar(subject))
    end
    return vars
end

Base.eltype(::Type{ <: DataVariable{D} }) where D = D
Base.eltype(::DataVariable{D})            where D = D

degree(datavar::DataVariable)    = nconnected(datavar)
name(datavar::DataVariable)      = datavar.name

isconnected(datavar::DataVariable) = datavar.nconnected !== 0
nconnected(datavar::DataVariable)  = datavar.nconnected

getlastindex(::DataVariable) = 1

messageout(datavar::DataVariable, ::Int) = datavar.messageout
messagein(datavar::DataVariable, ::Int)  = error("It is not possible to get a reference for inbound message for datavar")

update!(datavar::DataVariable, ::Missing)           = next!(messageout(datavar, 1), Message(missing, false, false))
update!(datavar::DataVariable, data::Number)        = update!(eltype(datavar), typeof(data), datavar, data)
update!(datavar::DataVariable, data::AbstractArray) = update!(eltype(datavar), typeof(data), datavar, data)

update!(::Type{ PointMass{D} }, ::Type{D}, datavar, data)   where { D }      = next!(messageout(datavar, 1), Message(PointMass(data), false, false))
update!(::Type{ PointMass{D1} }, ::Type{D2}, datavar, data) where { D1, D2 } = error("'$(name(datavar)) = datavar($D1, ...)' accepts data of type $D1, but $D2 has been supplied. Check 'update!($(name(datavar)), data::$D2)' and explicitly convert data to type $D1.")

resend!(datavar::DataVariable) = update!(datavar, Rocket.getrecent(messageout(datavar, 1)))

function update!(datavars::AbstractVector{ <: DataVariable }, data::AbstractVector)
    @assert size(datavars) === size(data) "Invalid update! call: size of datavar array and data should match"
    foreach(zip(datavars, data)) do (var, d)
        update!(var, d)
    end
end

finish!(datavar::DataVariable) = complete!(messageout(datavar, 1))

get_pipeline_stages(::DataVariable) = EmptyPipelineStage()

_getmarginal(datavar::DataVariable)                                = datavar.marginal
_setmarginal!(datavar::DataVariable, marginal::MarginalObservable) = datavar.marginal = marginal
_makemarginal(datavar::DataVariable)                               = datavar.messageout |> map(Marginal, as_marginal)

function setmessagein!(datavar::DataVariable, ::Int, messagein)
    datavar.nconnected += 1
    return nothing
end


# Test different derivative options
method = :hermite_simpson_separated
function check_grad_options(method)
    solver = Solver(model,obj,dt=dt)
    sol,stats = solve_dircol(solver,X0,U0,method=method,nlp=:snopt,grads=:none)
    @test vecnorm(sol.X[:,end]-obj.xf) < 1e-5

    sol2,stats = solve_dircol(solver,X0,U0,method=method,nlp=:snopt,grads=:auto)
    @test vecnorm(sol2.X[:,end]-obj.xf) < 1e-5
    @test norm(sol.X-sol2.X) < 5e-3
end

check_grad_options(:hermite_simpson_separated)
check_grad_options(:hermite_simpson)
check_grad_options(:trapezoid)
check_grad_options(:midpoint)

# Mesh refinement # TODO: mesh for ipopt
mesh = [0.5,0.2]
t1 = @elapsed sol,stats = solve_dircol(solver,X0,U0,nlp=:snopt)
t2 = @elapsed sol2,stats = solve_dircol(solver,X0,U0,mesh,)
@test vecnorm(sol.X[:,end]-obj.xf) < 1e-5
@test vecnorm(sol.X-sol.X2) < 1e-2
@test stats["info"][end] == "Finished successfully: optimality conditions satisfied"


struct Detector
    deadtime::Float64
    resolution::Float64
    jitter::Float64
    efficiency::Float64
    darkcounts::Float64
    function Detector(deadtime::Float64,resolution::Float64,jitter::Float64,efficiency::Float64,darkcounts::Float64)
        @assert 0 ≤ efficiency ≤ 1 "Efficiency must be between 0 and 1"
        new(
            deadtime,
            resolution,
            jitter,
            efficiency,
            darkcounts
        )
    end
end

"""
    Detector(deadtime::Number,resolution::Number,jitter::Number,efficiency::Number,darkcounts::Number)

Returns a Detector object. Also accepts a dictionary with keywords matching the inputs below.

Inputs:
    deadtime   : dead time in nanoseconds
    resolution : time resolution in nanoseconds
    jitter     : timing jitter in nanoseconds
    efficiency : quantum efficiency ∈ [0,1]
    darkcounts : dark count rate in GHz
"""
function Detector(deadtime::Number,resolution::Number,jitter::Number,efficiency::Number,darkcounts::Number)
    return Detector(
        convert(Float64,deadtime),
        convert(Float64,resolution),
        convert(Float64,jitter),
        convert(Float64,efficiency),
        convert(Float64,darkcounts)
    )
end

function Detector(params::Dict)
    return Detector(
            params[:deadtime],
            params[:resolution],
            params[:jitter],
            params[:efficiency],
            params[:darkcounts]
        )
end

export Detector

"""
    sparseReadout(t::Number,source::LightSource,detect::Detector)

Returns a sparse vector containing the nonzero counts received by the detector for the given duration t and LightSource
"""
function sparseReadout(t::Number,source::LightSource,detect::Detector)
    field = eField(source)
    return sparseReadout(t,field,detect)
end

"""
    sparseReadout(t::Number,field::eField,detect::Detector)

Returns a sparse vector containing the nonzero counts received by the detector for the given duration t and EM field
"""
function sparseReadout(t::Number,field::eField,detect::Detector)
    nb = nbar(t,field.source)
    γint = γIntensity(nb,t,detect.resolution,field)
    return sparseReadout(γint,detect)
end

"""
    sparseReadout(γint:γIntensity,detect::Detector)

Returns a sparse vector containing the nonzero counts received by the detector for the given photon intensity time-series
"""
function sparseReadout(γint::γIntensity,detect::Detector)
    out = sparsevec(Integer[],Integer[],length(γint))
    ideadtime = convert(Int,ceil(detect.deadtime/detect.resolution))
    i = 1
    while i ≤ length(γint)
        ct = poissonCount(γint[i])
        if ct != 0
            out[i] = ct
            i+=ideadtime
        end
        i+=1
    end
    return out
end

export sparseReadout

################################################################################
"""
    denseReadout(t::Number,source::LightSource,detect::Detector)

Returns a dense vector containing the counts received by the detector for the given duration t and LightSource
"""
function denseReadout(t::Number,source::LightSource,detect::Detector)
    field = eField(source)
    return denseReadout(t,field,detect)
end

"""
    denseReadout(t::Number,field::eField,detect::Detector)

Returns a dense vector containing the counts received by the detector for the given duration t and EM field
"""
function denseReadout(t::Number,field::eField,detect::Detector)
    nb = nbar(t,field.source)
    γint = γIntensity(nb,t,detect.resolution,field)
    return denseReadout(γint,detect)
end

"""
    denseReadout(γint:γIntensity,detect::Detector)

Returns a dense vector containing the counts received by the detector for the given photon intensity time-series
"""
function denseReadout(γint::γIntensity,detect::Detector)
    rngvec = Poisson.(γint.γvec)
    out = zeros(Int,length(γint))
    ideadtime = convert(Int,ceil(detect.deadtime/detect.resolution))
    i = 1
    while i ≤ length(γint)
        ct = rand(rngvec[i])
        if ct != 0
            out[i] = ct
            i+=ideadtime
        end
        i+=1
    end
    return out
end

export denseReadout


# Configuration parameters
struct Config
    # A minimum pixel radius to be included around each source.
    min_radius_pix::Float64

    # number of temperatures (for annealed importance sampling)
    num_ais_temperatures::Int

    # number of independent AIS samples
    num_ais_samples::Int

    # Number of iterations in joint VI. One iteration optimizes a full
    # pass over target sources.
    num_joint_vi_iters::Int
end

function Config(; min_radius_pix=8.0,
                num_ais_temperatures=50,
                num_ais_samples=10,
                num_joint_vi_iters=3)
    Config(min_radius_pix,
           num_ais_temperatures,
           num_ais_samples,
           num_joint_vi_iters)
end


# This file was generated, do not modify it. # hide
pyplot()
Plots.plot(sin)
Plots.savefig(joinpath(@OUTPUT, "pyplotbkend.svg")) # hide


module ChebyshevQuantum

using Plots
using ForwardDiff
using Polynomials
using LinearAlgebra

greet() = print("Hello World!")

include("Interp.jl")
include("EdgeDetect.jl")
include("Cheb.jl")
include("Op.jl")
include("SolveEig.jl")
include("SolveDiff.jl")

end # module


const DOC_FINITE =
    "`AbstractArray{<:Finite}` (multiclass classification)"
const DOC_FINITE_BINARY =
    "`AbstractArray{<:Finite{2}}` (binary classification)"
const DOC_ORDERED_FACTOR =
    "`AbstractArray{<:OrderedFactor}` (classification of ordered target)"
const DOC_ORDERED_FACTOR_BINARY =
    "`AbstractArray{<:OrderedFactor{2}}` "*
    "(binary classification where choice of \"true\" effects the measure)"
const DOC_CONTINUOUS = "`AbstractArray{Continuous}` (regression)"

## TRAITS

is_measure_type(::Any) = false

# The following traits, with fallbacks defined in
# StatisticalTraits.jl, make sense for some or all measures:

const MEASURE_TRAITS = [:name,
                        :instances,
                        :human_name,
                        :target_scitype,
                        :supports_weights,
                        :supports_class_weights,
                        :prediction_type,
                        :orientation,
                        :reports_each_observation,
                        :aggregation,
                        :is_feature_dependent,
                        :docstring,
                        :distribution_type]

## FOR BUILT-IN MEASURES (subtyping Measure)

abstract type Measure <: MLJType end
is_measure_type(::Type{<:Measure}) = true
is_measure(m) = is_measure_type(typeof(m))

# docstring fall-back:
_decorate(s::AbstractString) = "`$s`"
_decorate(v::Vector{<:AbstractString}) = join(_decorate.(v), ", ")
function MMI.docstring(M::Type{<:Measure})
    list = _decorate(instances(M))
    ret = "`$(name(M))` - $(human_name(M)) type"
    isempty(list) || (ret *= " with instances $list")
    ret *= ". "
    return ret
end

# display:
show_as_constructed(::Type{<:Measure}) = true

# info (see also src/init.jl):
function MLJScientificTypes.info(M, ::Val{:measure_type})
    values = Tuple(@eval($trait($M)) for trait in MEASURE_TRAITS)
    return NamedTuple{Tuple(MEASURE_TRAITS)}(values)
end
MLJScientificTypes.info(m, ::Val{:measure}) = info(typeof(m))

## AGGREGATION

(::Sum)(v) = sum(v)
(::Sum)(v::LittleDict) = sum(values(v))

(::Mean)(v) = mean(v)
(::Mean)(v::LittleDict) = mean(values(v))

(::RootMeanSquare)(v) = sqrt(mean(v.^2))

aggregate(v, measure) = aggregation(measure)(v)

# aggregation is no-op on scalars:
const MeasureValue = Union{Real,Tuple{<:Real,<:Real}} # number or interval
aggregate(x::MeasureValue, measure) = x

## DISPATCH FOR EVALUATION

# yhat - predictions (point or probabilisitic)
# X - features
# y - target observations
# w - per-observation weights

function value(measure, yhat, X, y, w)
    vfdep     = Val(is_feature_dependent(measure))
    vsweights = Val(supports_weights(measure) ||
                    supports_class_weights(measure))
    return value(measure, yhat, X, y, w, vfdep, vsweights)
end

## DEFAULT EVALUATION INTERFACE

#  is feature independent, weights not supported:
value(m, yhat, X, y, w, ::Val{false}, ::Val{false}) = m(yhat, y)

#  is feature dependent:, weights not supported:
value(m, yhat, X, y, w, ::Val{true}, ::Val{false}) = m(yhat, X, y)

#  is feature independent, weights supported:
value(m, yhat, X, y, w,         ::Val{false}, ::Val{true}) = m(yhat, y, w)
value(m, yhat, X, y, ::Nothing, ::Val{false}, ::Val{true}) = m(yhat, y)

#  is feature dependent, weights supported:
value(m, yhat, X, y, w,         ::Val{true}, ::Val{true}) = m(yhat, X, y, w)
value(m, yhat, X, y, ::Nothing, ::Val{true}, ::Val{true}) = m(yhat, X, y)

## helper

function check_pools(ŷ, y)
    levels(y) == levels(ŷ[1]) ||
        error("Conflicting categorical pools found "*
              "in observations and predictions. ")
    return nothing
end

## INCLUDE SPECIFIC MEASURES AND TOOLS

include("meta_utilities.jl")
include("continuous.jl")
include("confusion_matrix.jl")
include("finite.jl")
include("loss_functions_interface.jl")

## DEFAULT MEASURES
default_measure(T, S) = nothing

# Deterministic + Continuous / Count ==> RMS
default_measure(::Type{<:Deterministic},
                ::Type{<:Union{Vec{<:Continuous}, Vec{<:Count}}}) = rms

# Deterministic + Finite ==> Misclassification rate
default_measure(::Type{<:Deterministic},
                ::Type{<:Vec{<:Finite}}) = misclassification_rate

# default_measure(::Type{Probabilistic},
#                 ::Type{<:Union{Vec{<:Continuous},
#                                Vec{<:Count}}}) = ???

# Probabilistic + Finite ==> Cross entropy
default_measure(::Type{<:Probabilistic},
                ::Type{<:Vec{<:Finite}}) = cross_entropy

# Fallbacks
default_measure(M::Type{<:Supervised}) = default_measure(M, target_scitype(M))
default_measure(::M) where M <: Supervised = default_measure(M)


### Basic components and subtypes ###

abstract type Species end
abstract type SpeciesProperty{S<:Species} end
abstract type SpeciesDynamics{F<:Species} end
abstract type Ion <: Species end
abstract type PseudoIon <: Ion end

"""
Not including voltage in tagging ion channels. Flows are all charged particles. So they will always include voltage.
"""

struct Voltage <: Species end
struct Sodium <: Ion end
struct Potassium <: Ion end
struct Calcium <: Ion end
struct Proton <: Ion end
struct Leak <: PseudoIon end

abstract type Reversal{I<:Ion} <: SpeciesProperty{I} end
abstract type Current{I<:Ion} <: SpeciesProperty{I} end
abstract type Conductance{I<:Ion} <: SpeciesProperty{I} end
abstract type mRNA{I<:Ion} <: SpeciesProperty{I} end

shorthand_name(::Type{Voltage}) = :V
shorthand_name(::Type{Sodium}) = :Na
shorthand_name(::Type{Potassium}) = :K
shorthand_name(::Type{Calcium}) = :Ca
shorthand_name(::Type{Proton}) = :H
shorthand_name(::Type{Leak}) = :Leak
shorthand_name(::Type{Reversal{T}}) where {T} = Symbol(:E, shorthand_name(T))
shorthand_name(::Type{Current{T}}) where {T} = Symbol(:I, shorthand_name(T))
shorthand_name(::Type{Conductance{T}}) where {T} = Symbol(:g, shorthand_name(T))
shorthand_name(::Type{mRNA{T}}) where {T} = Symbol(:mRNA_, shorthand_name(T))
shorthand_name(x::Type{Tuple{T,R}}) where {T,R} = shorthand_name.(x.types)


# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

"""
    SingleBridgeOptimizer{BT<:AbstractBridge, OT<:MOI.ModelLike} <:
    AbstractBridgeOptimizer

The `SingleBridgeOptimizer` bridges any constrained variables supported by the
bridge `BT`. This is in contrast with the [`MathOptInterface.Bridges.LazyBridgeOptimizer`](@ref)
which only bridges the constrained variables that are unsupported by the
internal model, even if they are supported by one of its bridges.

!!! note
    Two bridge optimizers using variable bridges cannot be used together as both
    of them assume that the underlying model only returns variable indices with
    nonnegative values.
"""
mutable struct SingleBridgeOptimizer{BT<:AbstractBridge,OT<:MOI.ModelLike} <:
               MOIB.AbstractBridgeOptimizer
    model::OT
    map::Map # index of bridged variable -> variable bridge
    var_to_name::Dict{MOI.VariableIndex,String}
    name_to_var::Union{Dict{String,MOI.VariableIndex},Nothing}
    con_to_name::Dict{MOI.ConstraintIndex,String}
    name_to_con::Union{Dict{String,MOI.ConstraintIndex},Nothing}
end

function SingleBridgeOptimizer{BT}(model::OT) where {BT,OT<:MOI.ModelLike}
    return SingleBridgeOptimizer{BT,OT}(
        model,
        Map(),
        Dict{MOI.VariableIndex,String}(),
        nothing,
        Dict{MOI.ConstraintIndex,String}(),
        nothing,
    )
end

function bridges(::MOI.Bridges.AbstractBridgeOptimizer)
    return EmptyMap()
end

bridges(bridge::SingleBridgeOptimizer) = bridge.map

MOIB.supports_constraint_bridges(::SingleBridgeOptimizer) = false

function MOIB.supports_bridging_constrained_variable(
    ::SingleBridgeOptimizer{BT},
    S::Type{<:MOI.AbstractSet},
) where {BT}
    return supports_constrained_variable(BT, S)
end

function MOIB.is_variable_bridged(
    ::SingleBridgeOptimizer,
    ::Type{<:MOI.AbstractSet},
)
    return true
end

function MOIB.is_bridged(b::SingleBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    return MOIB.supports_bridging_constrained_variable(b, S)
end

function MOIB.is_bridged(
    ::SingleBridgeOptimizer,
    ::Type{<:MOI.AbstractFunction},
    ::Type{<:MOI.AbstractSet},
)
    return false
end

function MOIB.is_bridged(
    ::SingleBridgeOptimizer,
    ::Type{<:MOI.AbstractScalarFunction},
)
    return false
end

function MOIB.bridge_type(
    ::SingleBridgeOptimizer{BT},
    ::Type{<:MOI.AbstractSet},
) where {BT}
    return BT
end

MOIB.bridging_cost(::SingleBridgeOptimizer, args...) = 1.0

MOIB.recursive_model(b::SingleBridgeOptimizer) = b.model


module EMLcompressor
	using DanaTypes
	using DotPlusInheritance
	using Reexport
	@reexport using ...streams.EMLstreams
	import EMLtypes.length
	include("compressor/centrifugal_compressor.jl")
end


"""
    mutable struct MCMCSchedule
        num_mcmc_steps::Int64
        num_updates::Int64
        start::NamedTuple{(:mcmciter, :pidx),Tuple{Int64,Int64}}
        exclude_updates::DefaultDict{Int64,OrdinalRange{Int64,Int64},UnitRange{Int64}}
    end

An object used for iterating over the steps of MCMC samplers. `num_mcmc_steps`
is the total number of MCMC steps, whereas `num_updates` is a total number of
separate updates that can be performed in a single MCMC step. Not all of those
updates have to be performed and the pattern of excluding the updates is stored
in `exclude_updates`. The `num_updates` and `exclude_updates` can change in the
midst of iterating through the `MCMCSchedule` and the iterator will register
that and act accordingly.
"""
mutable struct MCMCSchedule{BE,S,EI}
    num_mcmc_steps::Int64
    num_updates::Int64
    start::S
    exclude_updates::DefaultDict{Int64,OrdinalRange{Int64,Int64},UnitRange{Int64}}
    extra_info::EI

    function MCMCSchedule(
            num_mcmc_steps,
            num_updates,
            exclude_updates=[];
            start::S=(prev_mcmciter=nothing, prev_pidx=nothing, mcmciter=1, pidx=1),
            backend::BE=GenericMCMCBackend(),
            extra_info::EI=nothing
        ) where {S,BE,EI}
        tuples_excluded_updates = DefaultDict{Int64,OrdinalRange{Int64,Int64}}(0:0)
        for exclude_update in exclude_updates
            for idx in exclude_update[1]
                tuples_excluded_updates[idx] = exclude_update[2]
            end
        end
        new{BE,S,EI}(
            num_mcmc_steps,
            num_updates,
            start,
            tuples_excluded_updates,
            extra_info
        )
    end
end

"""
    Base.iterate(iter::MCMCSchedule, state=iter.start)

Iterate through the `MCMCSchedule`, outputting the named tuples
`(mcmciter=..., pidx=...)` along the way, indicating the current index of the
mcmc sampler (`mcmciter`), as well as the index of update that is to be
performed.
"""
function Base.iterate(iter::MCMCSchedule, state=iter.start)
    state.mcmciter > iter.num_mcmc_steps && return nothing
    new_state_tmp = (
        prev_mcmciter = state.mcmciter,
        prev_pidx = state.pidx,
        mcmciter = state.mcmciter,
        pidx = state.pidx,
    )
    new_state = transition(iter, new_state_tmp)
    return (state, extra_transitions(iter, new_state))
end

extra_transitions(iter::MCMCSchedule, new_state) = new_state

"""
    transition(schedule::MCMCSchedule, state)

Determine the next state of the `MCMCSchedule` iterator, by skipping through
all updates that are to be excluded, as per `schedule.exclude_updates`.
"""
function transition(schedule::MCMCSchedule, state)
    pidx_to_reset = (state.pidx == schedule.num_updates)
    new_state = (
        prev_mcmciter = state.prev_mcmciter,
        prev_pidx = state.prev_pidx,
        mcmciter = state.mcmciter + pidx_to_reset,
        pidx = ( pidx_to_reset ? 1 : state.pidx + 1 ),
    )
    iterations_to_skip = schedule.exclude_updates[new_state.pidx]
    skip_state = new_state.mcmciter in iterations_to_skip
    skip_state && return transition(schedule, new_state)
    new_state
end

"""
    reschedule!(
        schedule::MCMCSchedule,
        num_new_updates=0,
        idxes_to_remove=[],
        idxes_to_add=[]
    )

Make changes to the `schedule` (possibly in the midst of iterating through it).
Add `num_new_updtes`-many allowable updates that can be performed at each MCMC
step, `idxes_to_remove` lists all parameter udpates that are to be completely
removed from the MCMC sampler, `idxes_to_add` lists entries that are to be
added to `schedule.exclude_updates`.
"""
function reschedule!(
        schedule::MCMCSchedule,
        num_new_updates=0,
        idxes_to_remove=[],
        idxes_to_add=[]
    )
    schedule.num_updates += num_new_updates
    for idx in idxes_to_remove
        schedule.exclude_updates[idx] = 1:schedule.num_mcmc_steps
    end
    for idx_to_add in idxes_to_add
        schedule.exclude_updates[idx_to_add[1]] = idx_to_add[2]
    end
end


# using POMDPs

# import Pkg
# Pkg.add(Pkg.PackageSpec(url="https://github.com/JuliaPOMDP/DiscreteValueIteration.jl"))
# POMDPs.add("DiscreteValueIteration", false)


"""
    contourf(C::ClimGrid; region::String="auto", level, caxis, titlestr::String, cm::String="", ncolors::Int, center_cs::Bool, filename::String, cs_label::String)

Maps the time-mean average of ClimGrid C. If a filename is provided, the figure is saved in a png format.

Optional keyworkd includes precribed regions (keyword *region*, see list below). For 4D data, keyword *level* is used to map a given level (defaults to 1). *caxis* is used to limit the colorscale. *cm* is used to manually set the colorscale (see Python documentation for native colorscale keyword), *ncolors* is used to set the number of color classes (defaults to 12). Set *center_cs* to true to center the colorscale (useful for divergent results, such as anomalies, positive/negative temprature). *cs_label* is used for custom colorscale label.

## Arguments for keyword *region* (and shortcuts)
- Europe ("EU")
- NorthAmerica ("NA")
- Canada ("CA")
- Quebec, QuebecNSP ("QC", "QCNSP")
- Americas ("Ams")
- World, WorldAz, WorldEck4 ("W", "Waz", "Weck4")
- Greenwich ("Gr")

"""
function contourf(C::ClimGrid; region::String="auto", states::Bool=false,  rivers::Bool=false, borders::Bool=false, lakes::Bool=false, level=1, caxis=[], titlestr::String="", cm::String="", ncolors::Int=12, center_cs::Bool=false, filename::String="", cs_label::String="", dpi=200)

    # Build basic plot information
    PI = ClimatePlots.build_PlotInfo(C, region, :contourf, caxis, cm, states, rivers, borders, lakes, center_cs, ncolors, level)

    # Empty-map generator
    status, ax, projection = mapclimgrid(PI)

    cs = ax.contourf(C.longrid, C.latgrid, PI.data, transform=cartopy.crs.PlateCarree(), cmap=PI.cm, vmin=PI.vmin, vmax=PI.vmax)

    # Colorbar
    if isempty(cs_label)
        cs_label = getunitslabel(C)
    end
    cbar = colorbar(cs, orientation = "vertical", shrink = 0.7, label=cs_label)

    # Title
    if isempty(titlestr)
        titlestr = titledef(C)
    end
    title(titlestr)

    # Save to "filename" if not empty
    if !isempty(filename)
        PyPlot.savefig(filename, dpi=dpi)
    end

    return true, ax, projection, cbar
end

"""
    contour(C::ClimGrid; region::String="auto", level, caxis, titlestr::String, cm::String="", ncolors::Int, center_cs::Bool, filename::String, cs_label::String)

Maps the time-mean average of ClimGrid C. If a filename is provided, the figure is saved in a png format.

Optional keyworkd includes precribed regions (keyword *region*, see list below). For 4D data, keyword *level* is used to map a given level (defaults to 1). *caxis* is used to limit the colorscale. *cm* is used to manually set the colorscale (see Python documentation for native colorscale keyword), *ncolors* is used to set the number of color classes (defaults to 12). Set *center_cs* to true to center the colorscale (useful for divergent results, such as anomalies, positive/negative temprature). *cs_label* is used for custom colorscale label.

## Arguments for keyword *region* (and shortcuts)
- Europe ("EU")
- NorthAmerica ("NA")
- Canada ("CA")
- Quebec, QuebecNSP ("QC", "QCNSP")
- Americas ("Ams")
- World, WorldAz, WorldEck4 ("W", "Waz", "Weck4")
- Greenwich ("Gr")

"""
function contour(C::ClimGrid; region::String="auto", states::Bool=false,  rivers::Bool=false, borders::Bool=false, lakes::Bool=false, level=1, caxis=[], titlestr::String="", cm::String="", ncolors::Int=12, center_cs::Bool=false, filename::String="", cs_label::String="", dpi=200)

    # Build basic plot information
    PI = ClimatePlots.build_PlotInfo(C, region, :contour, caxis, cm, states, rivers, borders, lakes, center_cs, ncolors, level)

    # Empty-map generator
    status, ax, projection = mapclimgrid(PI)

    cs = ax.contour(C.longrid, C.latgrid, PI.data, transform=cartopy.crs.PlateCarree(), cmap=PI.cm, vmin=PI.vmin, vmax=PI.vmax)

    # Colorbar
    if isempty(cs_label)
        cs_label = getunitslabel(C)
    end
    cbar = colorbar(cs, orientation = "vertical", shrink = 0.7, label=cs_label)

    # Title
    if isempty(titlestr)
        titlestr = titledef(C)
    end
    title(titlestr)

    # Save to "filename" if not empty
    if !isempty(filename)
        PyPlot.savefig(filename, dpi=dpi)
    end

    return true, ax, projection, cbar

end

"""
    pcolormesh(C::ClimGrid; region::String="auto", level, caxis, titlestr::String, cm::String="", ncolors::Int, center_cs::Bool, filename::String, cs_label::String)

Maps the time-mean average of ClimGrid C. If a filename is provided, the figure is saved in a png format.

Optional keyworkd includes precribed regions (keyword *region*, see list below). For 4D data, keyword *level* is used to map a given level (defaults to 1). *caxis* is used to limit the colorscale. *cm* is used to manually set the colorscale (see Python documentation for native colorscale keyword), *ncolors* is used to set the number of color classes (defaults to 12). Set *center_cs* to true to center the colorscale (useful for divergent results, such as anomalies, positive/negative temprature). *cs_label* is used for custom colorscale label.

## Arguments for keyword *region* (and shortcuts)
- Europe ("EU")
- NorthAmerica ("NA")
- Canada ("CA")
- Quebec, QuebecNSP ("QC", "QCNSP")
- Americas ("Ams")
- World, WorldAz, WorldEck4 ("W", "Waz", "Weck4")
- Greenwich ("Gr")

"""
function pcolormesh(C::ClimGrid; region::String="auto", states::Bool=false,  rivers::Bool=false, borders::Bool=false, lakes::Bool=false, level=1, caxis=[], titlestr::String="", cm::String="", ncolors::Int=12, center_cs::Bool=false, filename::String="", cs_label::String="", dpi=200)

    # Build basic plot information
    PI = ClimatePlots.build_PlotInfo(C, region, :pcolormesh, caxis, cm, states, rivers, borders, lakes, center_cs, ncolors, level)

    # Empty-map generator
    status, ax, projection = mapclimgrid(PI)

    cs = ax.pcolormesh(C.longrid, C.latgrid, PI.data, transform=cartopy.crs.PlateCarree(), cmap=PI.cm, vmin=PI.vmin, vmax=PI.vmax)

    # Colorbar
    if isempty(cs_label)
        cs_label = getunitslabel(C)
    end
    cbar = colorbar(cs, orientation = "vertical", shrink = 0.9, ticklocation="auto", label=cs_label)

    # Title
    if isempty(titlestr)
        titlestr = titledef(C)
    end
    title(titlestr)

    # Save to "filename" if not empty
    if !isempty(filename)
        PyPlot.savefig(filename, dpi=dpi)
    end

    return true, ax, projection, cbar
end


# Graph → Syntax

toexpr(f, xs...) = :($f($(xs...)))

binding(bindings::AbstractDict, v) =
  haskey(bindings, v) ? bindings[v] : (bindings[v] = gensym("edge"))

function syntax(head::DVertex; bindconst = !isfinal(head))
  vs = topo(head)
  ex, bs = :(;), IdDict{Any ,Any}()
  for v in vs
    x = toexpr(value(v), [binding(bs, n) for n in inputs(v)]...)
    if !bindconst && isconstant(v) && nout(v) > 1
      bs[v] = v.value.value
    elseif nout(v) > 1 || (!isfinal(head) && v ≡ head)
      edge = binding(bs, v)
      push!(ex.args, :($edge = $x))
    elseif haskey(bs, v)
      if MacroTools.inexpr(ex, bs[v])
        ex = MacroTools.replace(ex, bs[v], x)
      else
        push!(ex.args, :($(bs[v]) = $x))
      end
    else
      isfinal(v) ? push!(ex.args, x) : (bs[v] = x)
    end
  end
  head ≢ vs[end] && push!(ex.args, binding(bs, head))
  return ex
end

# TODO: this is butt ugly

function constructor(g)
  vertex = isa(g, DVertex) ? :(DataFlow.dvertex) : :(DataFlow.vertex)
  g = mapv(g) do v
    value(v) isa Lambda && (v.value = :(DataFlow.Lambda($(value(v).args), $(constructor(value(v).body)))))
    value(v) isa Constant && (v.value = :(DataFlow.Constant($(value(v).value))))
    prethread!(v, typeof(v)(Constant(value(v))))
    prethread!(v, typeof(v)(Constant(vertex)))
    v.value = Call()
    v
  end
  ex = syntax(g)
  decls, exs = [], []
  for x in block(ex).args
    if @capture(x, v_ = $vertex(f_, a__))
      push!(decls, :($v = $vertex($f)))
      push!(exs, :(DataFlow.thread!($v, $(a...))))
    else
      push!(exs, x)
    end
  end
  return :($(decls...);$(exs...))
end


"""
Implements the needed discharge for each purpose. (Purposes without discharge do not need to implement this function.)
The returned discharge may be an expression based on other variables.

For example, imposing a constraint on the reservoir level within the reservoir is not a discharge: the corresponding
function should return `0.`.
"""
hlp_JuMP_purposes_discharge(model::JuMP.Model, r::Reservoir, period::Period) =
  hlp_JuMP_purposes_discharge(model, getPurposes(r), period)

hlp_JuMP_purposes_discharge(model::JuMP.Model, purposes::Array{Purpose, 1}, period::Period) =
  sum([hlp_JuMP_purposes_discharge(model, p, period) for p in purposes])

# Generic purposes have no needed discharge.
hlp_JuMP_purposes_discharge(model::JuMP.Model, purpose::Purpose, period::Period) = 0.

hlp_JuMP_purposes_discharge(model::JuMP.Model, purpose::WaterWithdrawalPurpose, period::Period) =
  error("[hlp_JuMP_purposes_discharge] Helper not defined for " * string(typeof(purpose)) * ".")

hlp_JuMP_purposes_discharge(model::JuMP.Model, purpose::DeterministicWaterWithdrawalPurpose, period::Period) =
  getNeed(purpose, period)

# Hydropower has no imposed discharge.
hlp_JuMP_purposes_discharge(model::JuMP.Model, purpose::HydropowerDamOutput, period::Period) = 0.


# include("splitting_measures_mse.jl")
include("difference.jl")
include("msePointwise.jl")
include("mse.jl")
include("poissonGammaNormal.jl")

# these are not maintained any more
include("maxsplitvalue.jl")
include("maxabssplitvalue.jl")
include("maxminussplitvalue.jl")