text:
module IpoptInterfaceSerial

using StructJuMP, JuMP, StructJuMPSolverInterface
using Ipopt

import MathProgBase

type NonStructJuMPModel <: ModelInterface
    model::JuMP.Model
    jac_I::Vector{Int}
    jac_J::Vector{Int}
    hess_I::Vector{Int}
    hess_J::Vector{Int}
    nz_jac::Vector{Int}
    nz_hess::Vector{Int}

    write_solution::Function
    get_x0::Function
    numvars::Function
    numcons::Function
    nele_jac::Function
    nele_hess::Function
    get_bounds::Function
    eval_f::Function
    eval_g::Function
    eval_grad_f::Function
    eval_jac_g::Function
    eval_h::Function

    function NonStructJuMPModel(model)
        instance = new(model,
            Vector{Int}(), Vector{Int}(), Vector{Int}(), Vector{Int}(),
            Vector{Int}(), Vector{Int}()
            )

        instance.write_solution = function(x)
            @assert length(x) == getTotalNumVars(m)
            m = instance.model
            idx = 1
            for i = 0:num_scenarios(m)
                mm = getModel(m,i)
                for j = 1:getNumVars(m,i)
                    setvalue(Variable(mm,j),x[idx])
                    idx += 1
                end
            end
        end

        instance.get_x0 = function(x)
            @assert length(x) == getTotalNumVars(m)
            m = instance.model
            idx = 1
            for i = 0:num_scenarios(m)
                mm = getModel(m,i)
                for j = 1:getNumVars(m,i)
                    v_j = getvalue(Variable(mm,j))
                    x[idx] = isnan(v_j)? 1.0:v_j
                    idx += 1
                end
            end
            # @show x
            return x
        end
        instance.numvars = function()
            return getTotalNumVars(instance.model)
        end

        instance.numcons = function()
            return getTotalNumCons(instance.model)
        end

        instance.nele_jac = function()
            @assert length(instance.jac_I) == length(instance.jac_J)
            return length(instance.jac_I)
        end

        instance.nele_hess = function()
            @assert length(instance.hess_I) == length(instance.hess_J)
            return length(instance.hess_I)
        end

        instance.get_bounds = function()
            m = instance.model
            nvar = getTotalNumVars(m)
            ncon = getTotalNumCons(m)
            x_L = Vector{Float64}(nvar)
            x_U = Vector{Float64}(nvar)
            g_L = Vector{Float64}(ncon)
            g_U = Vector{Float64}(ncon)

            row_start = 1
            col_start = 1
            for i = 0:num_scenarios(m)
                mm = getModel(m,i)
                nx = getNumVars(m,i)
                array_copy(mm.colUpper, 1, x_U, col_start, nx)
                array_copy(mm.colLower, 1, x_L, col_start, nx)

                lb,ub = JuMP.constraintbounds(mm)
                ncons = getNumCons(m,i)
                array_copy(lb,1,g_L,row_start,ncons)
                array_copy(ub,1,g_U,row_start,ncons)

                row_start += getNumCons(m,i)
                col_start += getNumVars(m,i)
            end
            return x_L, x_U, g_L, g_U
        end

        instance.eval_f = function(x)
            m = instance.model
            objv = 0.0
            start_idx = 1
            for i=0:num_scenarios(m)
                x_new = strip_x(m,i,x,start_idx)
                objv += MathProgBase.eval_f(get_nlp_evaluator(m,i),x_new)
                start_idx += getNumVars(m,i)
            end
            return objv;
        end

        instance.eval_g = function(x,g)
            m = instance.model
            @assert length(g) == getTotalNumCons(m)
            start_idx = 1
            g_start_idx = 1
            for i=0:num_scenarios(m)
                x_new = strip_x(instance.model,i,x,start_idx)
                ncon = getNumCons(m,i)
                g_new = Vector{Float64}(ncon)
                e = get_nlp_evaluator(m,i)
                MathProgBase.eval_g(e,g_new,strip_x(instance.model,i,x,start_idx))
                array_copy(g_new,1,g,g_start_idx,ncon)
                g_start_idx += ncon
                start_idx += getNumVars(m,i)
            end
        end

        instance.eval_grad_f = function(x,grad_f)
            fill!(grad_f,0.0)
            m = instance.model
            start_idx = 1
            for i=0:num_scenarios(m)
                x_new = strip_x(instance.model,i,x,start_idx)
                e = get_nlp_evaluator(m,i)

                g_f = Vector{Float64}(length(x_new))
                MathProgBase.eval_grad_f(e,g_f,x_new)
                nx = getNumVars(m,i)

                array_copy(g_f,1,grad_f,start_idx,nx)

                othermap = getStructure(getModel(m,i)).othermap
                for i in othermap
                    pid = i[1].col
                    cid = i[2].col
                    grad_f[pid] += g_f[cid]
                    @assert pid <= getNumVars(m,0)
                end
                start_idx += nx
            end
        end

        instance.eval_jac_g = function(x,mode,rows,cols,values)
            m = instance.model
            if mode==:Structure
                @assert length(rows) == length(cols)
                for i = 1:length(rows)
                    rows[i] = instance.jac_I[i]
                    cols[i] = instance.jac_J[i]
                end
            else
                start_idx = 1
                value_start = 1
                for i = 0:num_scenarios(m)
                    e = get_nlp_evaluator(m,i)
                    x_new = strip_x(instance.model,i,x,start_idx)
                    i_nz_jac = instance.nz_jac[i+1]
                    jac_g = Vector{Float64}(i_nz_jac)
                    MathProgBase.eval_jac_g(e,jac_g,x_new)
                    array_copy(jac_g,1,values,value_start,i_nz_jac)
                    nx = getNumVars(m,i)
                    start_idx += nx
                    value_start += i_nz_jac
                end
            end
        end

        instance.eval_h = function(x, mode, rows, cols, obj_factor, lambda, values)
            m = instance.model
            if mode == :Structure
                @assert length(rows) == length(cols)
                for i = 1:length(rows)
                    rows[i] = instance.hess_I[i]
                    cols[i] = instance.hess_J[i]
                end
            else
                start_idx = 1
                value_start = 1
                lambda_start = 1
                for i = 0:num_scenarios(m)
                    e = get_nlp_evaluator(m,i)
                    x_new = strip_x(instance.model,i,x,start_idx)
                    nc = getNumCons(m,i)
                    lambda_new = Vector{Float64}(nc)
                    array_copy(lambda,lambda_start, lambda_new, 1, nc)
                    i_nz_hess = instance.nz_hess[i+1]
                    h = Vector{Float64}(i_nz_hess)
                    MathProgBase.eval_hesslag(e,h,x_new,obj_factor,lambda_new)
                    array_copy(h,1,values,value_start,i_nz_hess)
                    nx = getNumVars(m,i)
                    start_idx += nx
                    lambda_start += nc
                    value_start += i_nz_hess
                end
            end
        end

        # initialization  jac
        col_offset = 0
        row_offset = 0
        m = instance.model
        for i = 0:num_scenarios(m)
            reverse_map = Dict{Int,Int}()
            mm = getModel(m,i)
            for ety in getStructure(mm).othermap
                reverse_map[ety[2].col] = ety[1].col #child->parent
            end
            # @show reverse_map
            e = get_nlp_evaluator(m,i)
            # @show "after e"
            i_jac_I, i_jac_J =  MathProgBase.jac_structure(e)
            # @show "after strct jac"
            for idx = 1:length(i_jac_J)
                jj = i_jac_J[idx]
                if haskey(reverse_map,jj)
                    push!(instance.jac_J, reverse_map[jj])
                else
                    push!(instance.jac_J, jj + col_offset)
                end
                push!(instance.jac_I, i_jac_I[idx] + row_offset)
            end
            push!(instance.nz_jac, length(i_jac_J)) #offset by 1

            col_offset += getNumVars(m,i)
            row_offset += getNumCons(m,i)
        end

        #initialization hess
        offset = 0
        for i = 0:num_scenarios(m)
            reverse_map = Dict{Int,Int}()
            mm = getModel(m,i)
            for ety in getStructure(mm).othermap
                reverse_map[ety[2].col] = ety[1].col #child->parent
            end

            e = get_nlp_evaluator(m,i)
            i_hess_I, i_hess_J =  MathProgBase.hesslag_structure(e)
            for idx = 1:length(i_hess_I)
                ii = i_hess_I[idx]
                jj = i_hess_J[idx]

                if haskey(reverse_map,ii)
                    push!(instance.hess_I, reverse_map[ii])
                else
                    push!(instance.hess_I, ii + offset)
                end

                if haskey(reverse_map,jj)
                    push!(instance.hess_J, reverse_map[jj])
                else
                    push!(instance.hess_J, jj + offset)
                end
            end
            push!(instance.nz_hess, length(i_hess_I)) #offset by 1

            offset += getNumVars(m,i)
        end

        return instance
    end
end

function structJuMPSolve(model; suppress_warmings=false,kwargs...)
    # @show typeof(model)
    nm = NonStructJuMPModel(model)
    x_L, x_U, g_L, g_U = nm.get_bounds()
    n = getTotalNumVars(model)
    m = getTotalNumCons(model)
    nele_jac = nm.nele_jac()
    nele_hess = nm.nele_hess()

    # @show x_L, x_U
    # @show g_L, g_U
    # @show n,m
    # @show nele_jac,nele_hess

    prob = createProblem(n, x_L, x_U, m, g_L, g_U, nele_jac, nele_hess,
                         nm.eval_f, nm.eval_g, nm.eval_grad_f, nm.eval_jac_g, nm.eval_h)
    # setProblemScaling(prob,1.0)
    nm.get_x0(prob.x)
    status = solveProblem(prob)
    nm.write_solution(prob.x)
    return status
end

KnownSolvers["Ipopt"] = IpoptInterfaceSerial.structJuMPSolve

end  #end module


text:
module FeatureSelectors

using DataFrames
using Random

using ..AbsTypes
using ..BaseFilters
using ..Utils

import ..AbsTypes: fit!, transform!
export fit!, transform!
export FeatureSelector, CatFeatureSelector, NumFeatureSelector, CatNumDiscriminator

# generic way to extract num/cat features by specifying their columns
"""
    FeatureSelector(
       Dict(
         :name => "featureselector",
         :columns => [col1, col2, ...]
       )
    )

Returns a dataframe of the selected columns.

Implements `fit!` and `transform!`.
"""
mutable struct FeatureSelector <: Transformer
   name::String
   model::Dict{Symbol,Any}

   function FeatureSelector(args::Dict = Dict{Symbol,Any}())
      default_args = Dict{Symbol,Any}(
         :name => "featureselector",
         :columns => Int[]
      )
      cargs=nested_dict_merge(default_args,args)
      cargs[:name] = cargs[:name]*"_"*randstring(3)
      new(cargs[:name],cargs)
   end
end

"""
    FeatureSelector(cols::Vector{Int})

Helper function for FeatureSelector.
"""
function FeatureSelector(cols::Vector{Int})
   FeatureSelector(Dict(:columns => cols))
end

"""
    FeatureSelector(cols::Vararg{Int})

Helper function for FeatureSelector.
"""
FeatureSelector(cols::Vararg{Int}) = FeatureSelector([cols...])

function fit!(ft::FeatureSelector, features::DataFrame, labels::Vector=[])
   if features == DataFrame()
      throw(ArgumentError("empty dataframe"))
   end
end

function transform!(ft::FeatureSelector, features::DataFrame)
   nfeatures = deepcopy(features)
   if nfeatures == DataFrame()
      throw(ArgumentError("empty dataframe"))
   end
   cols = ft.model[:columns]
   if  cols != []
      return nfeatures[:,cols]
   else
      return DataFrame()
   end
end

# ----------
"""
    CatFeatureSelector(Dict(:name => "catf"))

Automatically extract categorical columns based on
inferred element types.

Implements `fit!` and `transform!`.
"""
mutable struct CatFeatureSelector <: Transformer
   name::String
   model::Dict{Symbol,Any}

   function CatFeatureSelector(args::Dict = Dict{Symbol,Any}())
      default_args = Dict{Symbol,Any}(
         :name => "catf",
         :nominal_columns => Int[]
      )
      cargs=nested_dict_merge(default_args,args)
      cargs[:name] = cargs[:name]*"_"*randstring(3)
      new(cargs[:name],cargs)
   end
end

function fit!(ft::CatFeatureSelector, features::DataFrame, labels::Vector=[])
    if features == DataFrame()
       throw(ArgumentError("empty dataframe"))
    end
    catcols,_ = find_catnum_columns(features)

    # create model
    ft.model[:nominal_columns] = catcols
end

function transform!(ft::CatFeatureSelector, features::DataFrame)
   nfeatures = deepcopy(features)
   catcols = ft.model[:nominal_columns]
   if catcols != []
      return nfeatures[:,catcols]
   else
      return DataFrame()
   end
end

"""
    NumFeatureSelector(Dict(:name=>"numfeatsel"))

Automatically extracts numeric features based on their inferred element types.

Implements `fit!` and `transform!`.
"""
mutable struct NumFeatureSelector <: Transformer
   name::String
   model::Dict{Symbol,Any}

   function NumFeatureSelector(args::Dict = Dict())
      default_args = Dict{Symbol,Any}(
        :name => "numf",
        :numcols => Int[]
      )
      cargs=nested_dict_merge(default_args,args)
      cargs[:name] = cargs[:name]*"_"*randstring(3)
      new(cargs[:name],cargs)
   end
end

function fit!(ft::NumFeatureSelector, features::DataFrame, labels::Vector=[])
    if features == DataFrame()
       throw(ArgumentError("empty dataframe"))
    end
    _,numcols = find_catnum_columns(features)

    # create model
    ft.model[:numcols] = numcols
end

function transform!(ft::NumFeatureSelector, features::DataFrame)
   nfeatures = deepcopy(features)
   numcols = ft.model[:numcols]
   if numcols != []
      return nfeatures[:,numcols]
   else
      return DataFrame()
   end
end

"""
    CatNumDiscriminator(
       Dict(
          :name => "catnumdisc",
          :maxcategories => 24
       )
    )

Transform numeric columns to string (as categories)
if the count of their unique elements <= maxcategories.

Implements `fit!` and `transform!`.
"""
mutable struct CatNumDiscriminator <: Transformer
   name::String
   model::Dict

   function CatNumDiscriminator(args::Dict = Dict())
      default_args = Dict(
         :name => "catnumdisc",
         # default max categories for numeric-encoded categories
         :maxcategories => 24,
         :nominal_columns => Int[],
         :numcols => Int[]
      )
      cargs=nested_dict_merge(default_args,args)
      cargs[:name] = cargs[:name]*"_"*randstring(3)
      new(cargs[:name],cargs)
   end
end

"""
    CatNumDiscriminator(maxcat::Int)

Helper function for CatNumDiscriminator.
"""
function CatNumDiscriminator(maxcat::Int)
    CatNumDiscriminator(Dict(:maxcategories=>maxcat))
end

function fit!(ft::CatNumDiscriminator, features::DataFrame, labels::Vector=[])
    if features == DataFrame()
       throw(ArgumentError("empty dataframe"))
    end
    catcols,numcols = find_catnum_columns(features,ft.model[:maxcategories])

    # create model
    ft.model[:numcols] = numcols
    ft.model[:nominal_columns] = catcols
end

function transform!(ft::CatNumDiscriminator, features::DataFrame)
    nfeatures = features |> deepcopy
    catcols = ft.model[:nominal_columns]
    if catcols != []
		nfeatures[!,catcols] = nfeatures[!,catcols] .|> string
    end
    return nfeatures
end

end

ext: jl


text:
function schsolve_Emin_sd( pw::PWGrid, Vpot, psi::Array{ComplexF64,2};
                           NiterMax=1000 )
  α = 3e-5
  Etot_old = 0.0
  Etot = 0.0

  for iter = 1:NiterMax
    psi = psi - α*calc_grad( pw, Vpot, psi )

    psi  = ortho_gram_schmidt(psi)
    Etot = calc_Etot( pw, Vpot, psi )

    conv = abs(Etot-Etot_old)
    @printf("Steepest descent: (s, step, conv) %8d %18.10f %18.10e\n",
            iter, Etot, conv )
    if conv < 1e-6
      print("Convergence achieved\n")
      break
    end
    Etot_old = Etot
  end
  return psi, Etot
end


text:
function apply!(state::QState, block::QGateBlock, pos::Vector{Int}; kwargs...)
	for j = 1: length(block)
		gateOrBlock = gates(block)[j]
		gateOrBlockPos = qubits(block)[j]
		apply!(state, gateOrBlock, [pos[i] for i ∈ gateOrBlockPos]; kwargs...)
	end
end
apply!(state::QState, block::QGateBlock, pos::Int; kwargs...) = apply!(state, block, [pos]; kwargs...)

apply!(state::QState, circuit::QCircuit, pos::Vector{Int} = [1:circuit.N;]; kwargs...) = apply!(state, circuit.block, pos; kwargs...)


text:
include("../../geometry/sphere.jl")
include("../../physics/fluid.jl")
include("../../physics/dipoleforce.jl")
include("../../utils/simulator.jl")
include("../../utils/plot.jl")

fluid = FluidParams( 1. )

rs = 0.1

# two dipoles perpendicular to each other with orientations parallel and perpendicular to geodesic
# cot θ = a cos(φ − φ0) with a = 0 and ϕ0 = π/2 i.e. cot θ = cos(φ − π/2)
d2θ = π/2 + 0.1
d2ϕ = acos(cot(d2θ)) + π/2
d2α = atan(sin(d2θ)*sin(d2ϕ - π/2)) + pi/2 # perpendicular to tangent to geodesic ( 1/sin(θ)*(dθ/dϕ) + π/2 )
dipoles = [
    DipoleForce( SphericalCoord(π/2, π), π/4, 1 ),
    DipoleForce( SphericalCoord(d2θ, d2ϕ), d2α, 1 )
]

# Generating a frame is computationally expensive,
# but each step forward in differential equation is computationally cheap,
# Higher steps reduce errors in numerical integration, can use better numerical schemes to reduce steps
endTime = 1 # time till when we want to simulate
numFrames = endTime * 100 # number of frames in animation,
numSteps = 10000 # number of steps forward in each frame,
adaptivePosMove = rs/10

# do the simulation
simulatedObjects = [ dipoles ]

for i in 2:numFrames
    newState = stepForward( last(simulatedObjects), fluid, endTime/numFrames/numSteps, numSteps, adaptivePosMove)
    nvortices = newState
    append!(simulatedObjects, [ newState ] )
    print(i, "\n")
end

# do the plitting/animation
animate(
    simulatedObjects,
    plotLatLon = true,
    # fieldGenerator = (p, vortices, pointforces, dipoles) -> sum( map( s -> velocityField(p, s, fluid), vcat( vortices, pointforces, dipoles ) ) ),
    resolution = (1000,600),
    makeMovie = true,
    cameraPosition = Vec3f0(-2, 0, 0),
    cameraLookingAt = Vec3f0(-1, 0, 0),
)


text:
function Eexpψ(θ4::T, σ::T, ψ::T, Dgt0::Bool) where {T<:Real}
    if Dgt0
        out = θ4*ψ
    else
        ρ = _ρ(σ)
        out = θ4*(ψ*ρ + θ4*0.5*(one(T)-ρ^2))
    end
    return out::T
end

# ------------------------

@inline function rev_exp(θ0::Real, θ1::T, θ2::Real, θ3::T, θ4::T, σ::T, logp::Real, ψ::Real, Dgt0::Bool, geoid::Real, roy::Real) where {T<:Real}
    r = (one(T)-θ0*roy) * exp(θ1 + θ2*logp + θ3*geoid + Eexpψ(θ4, σ, ψ, Dgt0))
    return r::T
end

@inline function drevdσ_exp(θ0::Real, θ1::T, θ2::Real, θ3::T, θ4::T, σ::T, logp::Real, ψ::Real, geoid::Real, roy::Real) where {T}
    return rev_exp(θ0,θ1,θ2,θ3,θ4,σ,logp,ψ,false,geoid, roy) * (ψ*θ4 - θ4^2*_ρ(σ)) * _dρdσ(σ)
end

@inline function drevdψ_exp(θ0::Real, θ1::T, θ2::Real, θ3::T, θ4::T, σ::T, logp::Real, ψ::Real, geoid::Real, roy::Real) where {T}
    rev_exp(θ0,θ1,θ2,θ3,θ4,σ,logp,ψ,false,geoid, roy) * θ4 * _ρ(σ)
end

# ------------------------

# add capability for a time trend
@inline function rev_exp(θ0::Real, θ1::T, θ2::Real, θ3::T, θ4::T, θ5::T, σ::T, logp::Real, t::Real, ψ::Real, Dgt0::Bool, geoid::Real, roy::Real) where {T<:Real}
    r = (one(T)-θ0*roy) * exp(θ1 + θ2*logp + θ3*geoid + Eexpψ(θ4, σ, ψ, Dgt0) + θ5*t)
    return r::T
end

@inline function drevdσ_exp(θ0::Real, θ1::T, θ2::Real, θ3::T, θ4::T, θ5::T, σ::T, logp::Real, t::Real, ψ::Real, geoid::Real, roy::Real) where {T<:Real}
    return rev_exp(θ0,θ1,θ2,θ3,θ4,θ5,σ,logp,t,ψ,false,geoid, roy) * (ψ*θ4 - θ4^2*_ρ(σ)) * _dρdσ(σ)
end

@inline function drevdψ_exp(θ0::Real, θ1::T, θ2::Real, θ3::T, θ4::T, θ5::T, σ::T, logp::Real, t::Real, ψ::Real, geoid::Real, roy::Real) where {T}
    rev_exp(θ0,θ1,θ2,θ3,θ4,θ5,σ,logp,t,ψ,false,geoid, roy) * θ4 * _ρ(σ)
end

# -----------------------------------------------------
# -----------------------------------------------------

@inline function flow(::Type{Val{:one}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[5]
        return zero(T)
    end
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid, roy) + θ[4]
    d>1      && (u *= d)
    # d1 == 1  && (u += θ[6])
    return u::T
end

@inline function flow(::Type{Val{:dgt1}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[6]
        return zero(T)
    end
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid, roy) + (d==1 ? θ[4] : θ[5] )
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:dgt1_cost}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[7]
        return zero(T)
    end
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid, roy) + (d==1 ? θ[4] : θ[5] ) + θ[6]*exp(z[2])
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:dgt1_pricecost}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[8]
        return zero(T)
    end
    u = rev_exp(1,θ[1],θ[4],θ[2],θ[3],σ,first(z),ψ,Dgt0,geoid, roy) + (d==1 ? θ[5] : θ[6] ) + θ[7]*exp(z[2])
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:dgt1_cost_Dgt0}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[7]
        return zero(T)
    end
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid, roy) + (d==1 ? θ[4] : θ[5] ) + θ[6]*exp(z[2])
    d>1      && (u *= d)
    !Dgt0 && d > 0 && (u += θ[8])
    return u::T
end

@inline function flow(::Type{Val{:dgt1_d1}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[7]
        return zero(T)
    end
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid, roy) + (d==1 ? θ[4] : θ[5] ) # + θ[8]*d)
    d>1      && (u *= d)
    d1 == 1  && (u += θ[6])
    return u::T
end

@inline function flow(::Type{Val{:Dgt0}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[6]
        return zero(T)
    end
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid, roy) + (!Dgt0 ?  θ[4] : θ[5] )
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:dgt1_ext}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T},  ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[6] + θ[7]*ψ
        return zero(T)
    end
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid, roy) + (d==1 ?  θ[4] : θ[5] )
    d>1      && (u *= d)
    return u::T
end

# --------------------

@inline function flow(::Type{Val{:one_restr}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[3]
        return zero(T)
    end
    u = rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy) + θ[2]
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:dgt1_restr}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[4]
        return zero(T)
    end
    u = rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy) + (d==1 ? θ[2] : θ[3] )
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:dgt1_cost_restr}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[5]
        return zero(T)
    end
    u = rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy) + (d==1 ? θ[2] : θ[3] ) + θ[4]*exp(z[2])
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:dgt1_pricecost_restr}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[6]
        return zero(T)
    end
    u = rev_exp(1, θ[1], θ[2], STARTING_log_ogip, STARTING_σ_ψ, σ, first(z), ψ, Dgt0, geoid, roy) + (d==1 ? θ[3] : θ[4] ) + θ[5]*exp(z[2])
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:dgt1_cost_Dgt0_restr}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[5]
        return zero(T)
    end
    u = rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy) + (d==1 ? θ[2] : θ[3] ) + θ[4]*exp(z[2])
    d>1      && (u *= d)
    !Dgt0 && d > 0 && (u += θ[6])
    return u::T
end

@inline function flow(::Type{Val{:dgt1_d1_restr}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T},  ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[5]
        return zero(T)
    end
    u = rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy) + (d==1 ? θ[2] : θ[3] )
    d>1      && (u *= d)
    d1 == 1  && (u += θ[4])
    return u::T
end

@inline function flow(::Type{Val{:dgt1_ext_restr}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[4] + θ[5]*ψ
        return zero(T)
    end
    u = rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy) + (d==1 ? θ[2] : θ[3] )
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:Dgt0_restr}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[4]
        return zero(T)
    end
    u = rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy) + (!Dgt0 ? θ[2] : θ[3] )
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:ttbuild_cost}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, τrem::Real, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[10]
        return zero(T)
    end
    logp, logc = z
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,logp,ψ,Dgt0,geoid, roy) + (d==1 ? θ[4] : θ[5] ) + θ[6]*cheb1(τrem) + θ[7]*cheb2(τrem) + θ[8]*cheb3(τrem) + θ[9]*exp(logc)
    d>1 && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:ttbuild_cost_restr}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, τrem::Real, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[8]
        return zero(T)
    end
    logp, logc = z
    u = rev_exp_restricted(θ[1], σ, logp, ψ, Dgt0, geoid, roy) + (d==1 ? θ[2] : θ[3] ) + θ[4]*cheb1(τrem) + θ[5]*cheb2(τrem) + θ[6]*cheb3(τrem) + θ[7]*exp(logc)
    d>1 && (u *= d)
    return u::T
end

# -----------------------------------------
# dθ
# -----------------------------------------

@inline function flowdθ(::Type{Val{:one}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    # revenue
    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    # drilling cost
    k == 4  && return d == 0 ? zero(T) : convert(T,d)

    # extension cost
    k == 5  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T},  ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    k == 4  && return  d  == 1 ? one(T)  : zero(T)
    k == 5  && return  d  == 1 ? zero(T) : convert(T,d)
    # k == 6  && return  d1 == 1 ? one(T)  : zero(T)

    k == 6  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_cost}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T},  ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    k == 4  && return  d  == 1 ? one(T)  : zero(T)
    k == 5  && return  d  == 1 ? zero(T) : convert(T,d)
    k == 6  && return  d*exp(z[2])

    k == 7  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_pricecost}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T},  ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return   d * rev_exp(1,θ[1],θ[4],θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],θ[4],θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],θ[4],θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))
    k == 4  && return   d * rev_exp(1,θ[1],θ[4],θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * first(z)

    k == 5  && return  d  == 1 ? one(T)  : zero(T)
    k == 6  && return  d  == 1 ? zero(T) : convert(T,d)
    k == 7  && return  d*exp(z[2])

    k == 8  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_cost_Dgt0}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T},  ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    k == 4  && return  d  == 1 ? one(T)  : zero(T)
    k == 5  && return  d  == 1 ? zero(T) : convert(T,d)
    k == 6  && return  d*exp(z[2])

    k == 7  && return d == 0 && sgn_ext ? one(T) : zero(T)
    k == 8  && return !Dgt0 && d > 0 ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_d1}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T},  ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    k == 4  && return  d  == 1 ? one(T)  : zero(T)
    k == 5  && return  d  == 1 ? zero(T) : convert(T,d)
    k == 6  && return  d1 == 1 ? one(T)  : zero(T)

    k == 7  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_ext}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    k == 4  && return  d  == 1 ? one(T)  : zero(T)
    k == 5  && return  d  == 1 ? zero(T) : convert(T,d)
    # k == 6  && return  d1 == 1 ? one(T)  : zero(T)

    k == 6  && return d == 0 && sgn_ext ? one(T) : zero(T)
    k == 7  && return d == 0 && sgn_ext ? ψ : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:Dgt0}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    k == 4  && return  !Dgt0 ? convert(T,d) : zero(T)
    k == 5  && return  !Dgt0 ? zero(T)      : convert(T,d)
    # k == 6  && return  d1 == 1 ? one(T)  : zero(T)

    k == 6  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:one_restr}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return d * rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy)

    k == 2  && return d == 0 ? zero(T) : convert(T,d)

    k == 3  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_restr}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return d * rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy)

    k == 2  && return d == 1 ? one(T)  : zero(T)
    k == 3  && return d == 1 ? zero(T) : convert(T,d)

    k == 4  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_cost_restr}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return   d * rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy)

    k == 2  && return  d  == 1 ? one(T)  : zero(T)
    k == 3  && return  d  == 1 ? zero(T) : convert(T,d)
    k == 4  && return  d*exp(z[2])

    k == 5  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_pricecost_restr}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T},  ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return   d * rev_exp(1, θ[1], θ[2], STARTING_log_ogip, STARTING_σ_ψ, σ, first(z), ψ, Dgt0, geoid, roy)
    k == 2  && return   d * rev_exp(1, θ[1], θ[2], STARTING_log_ogip, STARTING_σ_ψ, σ, first(z), ψ, Dgt0, geoid, roy) * first(z)

    k == 3  && return  d  == 1 ? one(T)  : zero(T)
    k == 4  && return  d  == 1 ? zero(T) : convert(T,d)
    k == 5  && return  d*exp(z[2])

    k == 6  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_cost_Dgt0_restr}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T},  ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return   d * rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy)

    k == 2  && return  d  == 1 ? one(T)  : zero(T)
    k == 3  && return  d  == 1 ? zero(T) : convert(T,d)
    k == 4  && return  d*exp(z[2])

    k == 5  && return d == 0 && sgn_ext ? one(T) : zero(T)
    k == 6  && return !Dgt0 && d > 0 ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_d1_restr}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return d * rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy)

    k == 2  && return d  == 1 ? one(T)  : zero(T)
    k == 3  && return d  == 1 ? zero(T) : convert(T,d)
    k == 4  && return d1 == 1 ? one(T)  : zero(T)

    k == 5  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_ext_restr}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return d * rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy)

    k == 2  && return d == 1 ? one(T)  : zero(T)
    k == 3  && return d == 1 ? zero(T) : convert(T,d)

    k == 4  && return d == 0 && sgn_ext ? one(T) : zero(T)
    k == 5  && return d == 0 && sgn_ext ? ψ : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:Dgt0_restr}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T},  ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return d * rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy)

    k == 2  && return !Dgt0 ? convert(T,d) : zero(T)
    k == 3  && return !Dgt0 ? zero(T)      : convert(T,d)

    k == 4  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

# -----------------------------------------
# dσ
# -----------------------------------------

# -----------------------------------------
# dψ
# -----------------------------------------

@inline function flowdψ(::FF, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, sgn_ext::Bool, geoid::Real, roy::T) where {FF <: Union{ Type{Val{:dgt1_ext}} },N, T}
    d == 0  && return sgn_ext ? θ[7] : zero(T)
    return (d * drevdψ_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,geoid,roy))::T
end

@inline function flowdψ(::FF, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, sgn_ext::Bool, geoid::Real, roy::T) where {FF <: Union{ Type{Val{:dgt1_ext_restr}} },N, T}
    d == 0  && return sgn_ext ? θ[5] : zero(T)
    return (d * drevdψ_exp_restricted(θ[1],σ,z[1],ψ,geoid,roy))::T
end

# -----------------------------------------
# Allow price to vary
# -----------------------------------------

@inline function flow(::Type{Val{:dgt1_pricebreak}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[9]
        return zero(T)
    end
    u = rev_exp(1,θ[1],θ[4+Dgt0],θ[2],θ[3],σ,first(z),ψ,Dgt0,geoid, roy) + (d==1 ? θ[6] : θ[7] ) + θ[8]*exp(z[2])
    d>1      && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:dgt1_pricebreak_restr}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[7]
        return zero(T)
    end
    u = rev_exp(1, θ[1], θ[2+Dgt0], STARTING_log_ogip, STARTING_σ_ψ, σ, first(z), ψ, Dgt0, geoid, roy) + (d==1 ? θ[4] : θ[5] ) + θ[6]*exp(z[2])
    d>1      && (u *= d)
    return u::T
end

@inline function flowdθ(::Type{Val{:dgt1_pricebreak}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T},  ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return                   d * rev_exp(1,θ[1],θ[4+Dgt0],θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return                   d * rev_exp(1,θ[1],θ[4+Dgt0],θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return                   d * rev_exp(1,θ[1],θ[4+Dgt0],θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))
    k == 4  && return  Dgt0 ? zero(T) : d * rev_exp(1,θ[1],θ[4+Dgt0],θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * first(z)
    k == 5  && return !Dgt0 ? zero(T) : d * rev_exp(1,θ[1],θ[4+Dgt0],θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * first(z)

    k == 6  && return  d  == 1 ? one(T)  : zero(T)
    k == 7  && return  d  == 1 ? zero(T) : convert(T,d)
    k == 8  && return  d*exp(z[2])

    k == 9  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:dgt1_pricebreak_restr}}, θ::AbstractVector{T}, σ::T,    z::NTuple{N,T},  ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    k == 1  && return                   d * rev_exp(1, θ[1], θ[2+Dgt0], STARTING_log_ogip, STARTING_σ_ψ, σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return  Dgt0 ? zero(T) : d * rev_exp(1, θ[1], θ[2+Dgt0], STARTING_log_ogip, STARTING_σ_ψ, σ,z[1],ψ,Dgt0,geoid,roy) * first(z)
    k == 3  && return !Dgt0 ? zero(T) : d * rev_exp(1, θ[1], θ[2+Dgt0], STARTING_log_ogip, STARTING_σ_ψ, σ,z[1],ψ,Dgt0,geoid,roy) * first(z)

    k == 4  && return  d  == 1 ? one(T)  : zero(T)
    k == 5  && return  d  == 1 ? zero(T) : convert(T,d)
    k == 6  && return  d*exp(z[2])

    k == 7  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

# -----------------------------------------
# dσ
# -----------------------------------------

@inline function flowdσ(::FF, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, geoid::Real, roy::T)::T where {FF <: Union{ Type{Val{:dgt1_pricecost}} }, N, T}
    d == 0 && return zero(T)
    return d * drevdσ_exp(1,θ[1],θ[4],θ[2],θ[3],σ,first(z),ψ,geoid,roy)
end

@inline function flowdσ(::FF, θ::AbstractVector{T}, σ::T, z::NTuple{N,T},  ψ::T, d::Integer, geoid::Real, roy::T)::T where {FF <: Union{ Type{Val{:dgt1_pricecost_restr}} },N, T}
    d == 0 && return zero(T)
    return d * drevdσ_exp(1, θ[1], θ[2], STARTING_log_ogip, STARTING_σ_ψ, σ, first(z), ψ, geoid, roy)
end

# can do this b/c only needed for !Dgt0
@inline function flowdσ(::FF, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, geoid::Real, roy::T)::T where {FF <: Union{ Type{Val{:dgt1_pricebreak}} }, N, T}
    d == 0 && return zero(T)
    Dgt0 = false
    return d * drevdσ_exp(1,θ[1],θ[4+Dgt0],θ[2],θ[3],σ,first(z),ψ,geoid,roy)
end

@inline function flowdσ(::FF, θ::AbstractVector{T}, σ::T, z::NTuple{N,T},  ψ::T, d::Integer, geoid::Real, roy::T)::T where {FF <: Union{ Type{Val{:dgt1_pricebreak_restr}} },N, T}
    d == 0 && return zero(T)
    Dgt0 = false
    return d * drevdσ_exp(1, θ[1], θ[2+Dgt0], STARTING_log_ogip, STARTING_σ_ψ, σ, first(z), ψ, geoid, roy)
end

# -----------------------------------------
# dψ
# -----------------------------------------

@inline function flowdψ(::FF, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, sgn_ext::Bool, geoid::Real, roy::T) where {FF <: Union{ Type{Val{:dgt1_pricecost}} }, N, T}
    d == 0  && return zero(T) # sgn_ext ? θ[10] : zero(T)
    return (d * drevdψ_exp(1,θ[1],θ[4],θ[2],θ[3],σ,first(z),ψ,geoid,roy))::T
end

@inline function flowdψ(::FF, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, sgn_ext::Bool, geoid::Real, roy::T) where {FF <: Union{ Type{Val{:dgt1_pricecost_restr}} }, N, T}
    d == 0  && return zero(T) # sgn_ext ? θ[10] : zero(T)
    return (d * drevdψ_exp(1,θ[1],θ[2], STARTING_log_ogip, STARTING_σ_ψ,σ,first(z),ψ,geoid,roy))::T
end

@inline function flowdψ(::FF, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, sgn_ext::Bool, geoid::Real, roy::T) where {FF <: Union{ Type{Val{:dgt1_pricebreak}} }, N, T}
    d == 0  && return zero(T) # sgn_ext ? θ[10] : zero(T)
    Dgt0 = false
    return (d * drevdψ_exp(1,θ[1],θ[4+Dgt0],θ[2],θ[3],σ,first(z),ψ,geoid,roy))::T
end

@inline function flowdψ(::FF, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, sgn_ext::Bool, geoid::Real, roy::T) where {FF <: Union{ Type{Val{:dgt1_pricebreak_restr}} }, N, T}
    d == 0  && return zero(T) # sgn_ext ? θ[10] : zero(T)
    Dgt0 = false
    return (d * drevdψ_exp(1,θ[1],θ[2+Dgt0], STARTING_log_ogip, STARTING_σ_ψ,σ,first(z),ψ,geoid,roy))::T
end

# ---------------------------------------------------------

@inline function flow(::Type{Val{:cheb2}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[7]
        return zero(T)
    end
    x = last(z)
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,first(z),ψ,Dgt0,geoid, roy) + θ[4]*cheb0(x) + θ[5]*cheb1(x) + θ[6]*cheb2(x)
    d>1 && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:cheb2_restr}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[5]
        return zero(T)
    end
    x = last(z)
    u = rev_exp_restricted(θ[1], σ, first(z), ψ, Dgt0, geoid, roy) + θ[2]*cheb0(x) + θ[3]*cheb1(x) + θ[4]*cheb2(x)
    d>1 && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:cheb3}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[8]
        return zero(T)
    end
    x = last(z)
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,first(z),ψ,Dgt0,geoid, roy) + θ[4]*cheb0(x) + θ[5]*cheb1(x) + θ[6]*cheb2(x) + θ[7]*cheb3(x)
    d>1 && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:cheb3_restr}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[6]
        return zero(T)
    end
    x = last(z)
    u = rev_exp_restricted(θ[1], σ, first(z), ψ, Dgt0, geoid, roy) + θ[2]*cheb0(x) + θ[3]*cheb1(x) + θ[4]*cheb2(x) + θ[5]*cheb3(x)
    d>1 && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:cheb3_dgt1}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[9]
        return zero(T)
    end
    x = last(z)
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,first(z),ψ,Dgt0,geoid, roy) + (d==1 ? θ[4] : θ[5] )*cheb0(x) + θ[6]*cheb1(x) + θ[7]*cheb2(x) + θ[8]*cheb3(x)
    d>1 && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:cheb3_dgt1_restr}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[7]
        return zero(T)
    end
    x = last(z)
    u = rev_exp_restricted(θ[1], σ, first(z), ψ, Dgt0, geoid, roy) + (d==1 ? θ[2] : θ[3] )*cheb0(x) + θ[4]*cheb1(x) + θ[5]*cheb2(x) + θ[6]*cheb3(x)
    d>1 && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:cheb3_cost}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[10]
        return zero(T)
    end
    logp, logc, t = z
    u = rev_exp(1,θ[1],1,θ[2],θ[3],σ,logp,ψ,Dgt0,geoid, roy) + (d==1 ? θ[4] : θ[5] )*cheb0(t) + θ[6]*cheb1(t) + θ[7]*cheb2(t) + θ[8]*cheb3(t) + θ[9]*exp(logc)
    d>1 && (u *= d)
    return u::T
end

@inline function flow(::Type{Val{:cheb3_cost_restr}}, θ::AbstractVector{T}, σ::T, z::NTuple{N,T}, ψ::T, d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T) where {N,T}
    if d == 0
        sgn_ext && return θ[8]
        return zero(T)
    end
    logp, logc, t = z
    u = rev_exp_restricted(θ[1], σ, logp, ψ, Dgt0, geoid, roy) + (d==1 ? θ[2] : θ[3] )*cheb0(t) + θ[4]*cheb1(t) + θ[5]*cheb2(t) + θ[6]*cheb3(t) + θ[7]*exp(logc)
    d>1 && (u *= d)
    return u::T
end

@inline function flowdθ(::Type{Val{:cheb2}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    # revenue
    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    # drilling cost
    k == 4  && return d == 0 ? zero(T) : T( d*cheb0(z[2]) )
    k == 5  && return d == 0 ? zero(T) : T( d*cheb1(z[2]) )
    k == 6  && return d == 0 ? zero(T) : T( d*cheb2(z[2]) )

    # extension cost
    k == 7  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:cheb2_restr}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    # revenue
    k == 1  && return   d * rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy)

    # drilling cost
    k == 2  && return d == 0 ? zero(T) : T( d*cheb0(z[2]) )
    k == 3  && return d == 0 ? zero(T) : T( d*cheb1(z[2]) )
    k == 4  && return d == 0 ? zero(T) : T( d*cheb2(z[2]) )

    # extension cost
    k == 5  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:cheb3}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    # revenue
    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    # drilling cost
    k == 4  && return d == 0 ? zero(T) : T( d*cheb0(z[2]) )
    k == 5  && return d == 0 ? zero(T) : T( d*cheb1(z[2]) )
    k == 6  && return d == 0 ? zero(T) : T( d*cheb2(z[2]) )
    k == 7  && return d == 0 ? zero(T) : T( d*cheb3(z[2]) )

    # extension cost
    k == 8  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:cheb3_restr}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    # revenue
    k == 1  && return   d * rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy)

    # drilling cost
    k == 2  && return d == 0 ? zero(T) : T( d*cheb0(z[2]) )
    k == 3  && return d == 0 ? zero(T) : T( d*cheb1(z[2]) )
    k == 4  && return d == 0 ? zero(T) : T( d*cheb2(z[2]) )
    k == 5  && return d == 0 ? zero(T) : T( d*cheb3(z[2]) )

    # extension cost
    k == 6  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:cheb3_dgt1}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    # revenue
    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,z[1],ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    # drilling cost
    k == 4  && return d == 1 ? T( cheb0(z[2]) ) : zero(T)
    k == 5  && return d <= 1 ? zero(T) : T( d*cheb0(z[2]) )
    k == 6  && return d == 0 ? zero(T) : T( d*cheb1(z[2]) )
    k == 7  && return d == 0 ? zero(T) : T( d*cheb2(z[2]) )
    k == 8  && return d == 0 ? zero(T) : T( d*cheb3(z[2]) )

    # extension cost
    k == 9  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:cheb3_dgt1_restr}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    # revenue
    k == 1  && return   d * rev_exp_restricted(θ[1], σ, z[1], ψ, Dgt0, geoid, roy)

    # drilling cost
    k == 2  && return d == 1 ? T( cheb0(z[2]) ) : zero(T)
    k == 3  && return d <= 1 ? zero(T) : T( d*cheb0(z[2]) )
    k == 4  && return d == 0 ? zero(T) : T( d*cheb1(z[2]) )
    k == 5  && return d == 0 ? zero(T) : T( d*cheb2(z[2]) )
    k == 6  && return d == 0 ? zero(T) : T( d*cheb3(z[2]) )

    # extension cost
    k == 7  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:cheb3_cost}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    logp, logc, t = z

    # revenue
    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,logp,ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,logp,ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,logp,ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    # drilling cost
    k == 4  && return d != 1 ? zero(T) : T(   cheb0(t) )
    k == 5  && return d <= 1 ? zero(T) : T( d*cheb0(t) )
    k == 6  && return d == 0 ? zero(T) : T( d*cheb1(t) )
    k == 7  && return d == 0 ? zero(T) : T( d*cheb2(t) )
    k == 8  && return d == 0 ? zero(T) : T( d*cheb3(t) )
    k == 9  && return d == 0 ? zero(T) : T( d*exp(logc) )

    # extension cost
    k == 10  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:cheb3_cost_restr}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    logp, logc, t = z

    # revenue
    k == 1  && return   d * rev_exp_restricted(θ[1], σ, logp, ψ, Dgt0, geoid, roy)

    # drilling cost
    k == 2  && return d != 1 ? zero(T) : T(   cheb0(t) )
    k == 3  && return d <= 1 ? zero(T) : T( d*cheb0(t) )
    k == 4  && return d == 0 ? zero(T) : T( d*cheb1(t) )
    k == 5  && return d == 0 ? zero(T) : T( d*cheb2(t) )
    k == 6  && return d == 0 ? zero(T) : T( d*cheb3(t) )
    k == 7  && return d == 0 ? zero(T) : T( d*exp(logc) )
    # extension cost
    k == 8  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:ttbuild_cost}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer, d1::Integer, Dgt0::Bool, sgn_ext::Bool, τrem::Real, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    logp, logc = z

    # revenue
    k == 1  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,logp,ψ,Dgt0,geoid,roy)
    k == 2  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,logp,ψ,Dgt0,geoid,roy) * geoid
    k == 3  && return   d * rev_exp(1,θ[1],1,θ[2],θ[3],σ,logp,ψ,Dgt0,geoid,roy) * ( Dgt0 ? ψ : ψ*_ρ(σ) + θ[k]*(1-_ρ2(σ)))

    # drilling cost
    k == 4  && return d != 1 ? zero(T) : T(   cheb0(τrem) )
    k == 5  && return d <= 1 ? zero(T) : T( d*cheb0(τrem) )
    k == 6  && return d == 0 ? zero(T) : T( d*cheb1(τrem) )
    k == 7  && return d == 0 ? zero(T) : T( d*cheb2(τrem) )
    k == 8  && return d == 0 ? zero(T) : T( d*cheb3(τrem) )
    k == 9  && return d == 0 ? zero(T) : T( d*exp(logc) )

    # extension cost
    k == 10  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end

@inline function flowdθ(::Type{Val{:ttbuild_cost_restr}}, θ::AbstractVector{T}, σ::T,     z::NTuple{N,T}, ψ::T, k::Integer,d::Integer,           d1::Integer, Dgt0::Bool, sgn_ext::Bool, τrem::Real, geoid::Real, roy::T)::T where {N,T}
    d == 0 && !sgn_ext && return zero(T)

    logp, logc = z

    # revenue
    k == 1  && return   d * rev_exp_restricted(θ[1], σ, logp, ψ, Dgt0, geoid, roy)

    # drilling cost
    k == 2  && return d != 1 ? zero(T) : T(   cheb0(τrem) )
    k == 3  && return d <= 1 ? zero(T) : T( d*cheb0(τrem) )
    k == 4  && return d == 0 ? zero(T) : T( d*cheb1(τrem) )
    k == 5  && return d == 0 ? zero(T) : T( d*cheb2(τrem) )
    k == 6  && return d == 0 ? zero(T) : T( d*cheb3(τrem) )
    k == 7  && return d == 0 ? zero(T) : T( d*exp(logc) )
    # extension cost
    k == 8  && return d == 0 && sgn_ext ? one(T) : zero(T)

    throw(error("$k out of bounds"))
end


text:
"""
DX-NES: distance-weighted extensions of xNES by Fukushima et al.
"""
mutable struct DXNESOpt{F,E<:EmbeddingOperator} <: ExponentialNaturalEvolutionStrategyOpt
    embed::E                        # operator embedding into the search space
    lambda::Int                     # Number of samples to take per iteration
    sortedUtilities::Vector{Float64}# Fitness utility to give to each rank
    tmp_Utilities::Vector{Float64}  # Fitness utilities assigned to current population
    x_learnrate::Float64
    B_learnrate::Float64
    sigma_learnrate::Float64
    max_sigma::Float64
    moving_threshold::Float64       # threshold of |evolutionary path| to detect movement
    evol_discount::Float64
    evol_Zscale::Float64
    uZscale::Float64                # distance scale for the weights
    evol_path::Vector{Float64}      # evolution path, discounted coordinates of population center
    # TODO use Symmetric{Float64} to inmprove exponent etc calculation
    ln_B::Matrix{Float64}           # exponential of lnB
    sigma::Float64                  # step size
    x::Individual                   # The current incumbent (aka most likely value, mu etc)
    Z::Matrix{Float64}              # current N(0,I) samples
    candidates::Vector{Candidate{F}}# The last sampled values, now being evaluated

    # temporary variables to minimize GC overhead
    tmp_x::Individual
    tmp_dz::Individual
    tmp_dx::Individual
    tmp_lndB::Matrix{Float64}
    tmp_Zu::Matrix{Float64}
    tmp_sBZ::Matrix{Float64}

    function DXNESOpt{F}(embed::E; lambda::Int = 0,
            mu_learnrate::Float64 = 1.0,
            ini_x = nothing, ini_sigma::Float64 = 1.0,
            ini_lnB = nothing,
            max_sigma::Float64 = 1.0E+10) where {F, E<:EmbeddingOperator}
        iseven(lambda) || throw(ArgumentError("lambda needs to be even"))
        d = numdims(search_space(embed))
        if lambda == 0
            lambda = 4 + 3*floor(Int, log(d))
            isodd(lambda) && (lambda += 1) # make sure lambda is even
        end
        if ini_x === nothing
            ini_x = rand_individual(search_space(embed))
        else
            ini_x = copy(ini_x::Individual)
            apply!(embed, ini_x, rand_individual(search_space(embed)))
        end
        u = fitness_shaping_utilities_log(lambda)
        moving_threshold, evol_discount, evol_Zscale = calculate_evol_path_params(d, u)

        new{F,E}(embed, lambda, u, Vector{Float64}(undef, lambda),
            mu_learnrate, 0.0, 0.0, max_sigma,
            moving_threshold, evol_discount, evol_Zscale, 0.9 + 0.15 * log(d),
            fill(NaN, d), ini_lnB === nothing ? ini_xnes_B(search_space(embed)) : ini_lnB, ini_sigma, ini_x,
            fill(NaN, d, lambda),
            [Candidate{F}(fill!(Individual(undef, d), NaN), i) for i in 1:lambda],
            # temporaries
            Vector{Float64}(undef, d), Vector{Float64}(undef, d),
            Vector{Float64}(undef, d), Matrix{Float64}(undef, d, d),
            Matrix{Float64}(undef, d, lambda), Matrix{Float64}(undef, d, lambda)
        )
    end
end

function trace_state(io::IO, dxnes::DXNESOpt, mode::Symbol)
    evol_path_norm = norm(dxnes.evol_path)
    println(io,
            "σ=", dxnes.sigma,
            " η[x]=", dxnes.x_learnrate,
            " η[σ]=", dxnes.sigma_learnrate,
            " η[B]=", dxnes.B_learnrate,
            " |tr(ln_B)|=", abs(tr(dxnes.ln_B)),
            " |path|=", evol_path_norm,
            " speed=", evol_path_norm/dxnes.moving_threshold)
end

const DXNES_DefaultOptions = chain(NES_DefaultOptions, ParamsDict(
    :ini_sigma => 1.0,      # Initial sigma (step size)
    :ini_lnB => nothing     # Initial log(B) (log of parameters covariation)
))

function dxnes(problem::OptimizationProblem, parameters)
    params = chain(DXNES_DefaultOptions, parameters)
    embed = RandomBound(search_space(problem))
    DXNESOpt{fitness_type(problem)}(embed; lambda = params[:lambda],
                                    ini_x = params[:ini_x],
                                    ini_sigma = params[:ini_sigma],
                                    ini_lnB = params[:ini_lnB],
                                    max_sigma = params[:max_sigma])
end

function ask(dxnes::DXNESOpt)
    hlambda = dxnes.lambda÷2
    @inbounds for i in 1:hlambda
        for j in 1:size(dxnes.Z, 1)
            dxnes.Z[j, i] = randn()
            dxnes.Z[j, i+hlambda] = -dxnes.Z[j, i]
        end
    end
    update_candidates!(dxnes, dxnes.Z)
end

function tell!(dxnes::DXNESOpt{F}, rankedCandidates::Vector{Candidate{F}}) where F
    u = assign_weights!(dxnes.tmp_Utilities, rankedCandidates, dxnes.sortedUtilities)
    dxnes.evol_path *= dxnes.evol_discount
    # We'll take the small perf hit for now just so this can run also on pre rc2 julia 0.7s
    dxnes.evol_path += dxnes.evol_Zscale * dropdims(wsum(dxnes.Z, u, 2), dims=2)
    evol_speed = norm(dxnes.evol_path)/dxnes.moving_threshold
    if evol_speed > 1.0
        # the center is moving, adjust weights
        u = assign_distance_weights!(dxnes.tmp_Utilities, dxnes.uZscale, rankedCandidates, dxnes.Z)
    end
    update_learning_rates!(dxnes, evol_speed)
    update_parameters!(dxnes, u)

    return 0
end

"""
Calculate the parameters for evolutionary path.

Returns the tuple:
  * moving threshold
  * path discount
  * current Z scale
"""
function calculate_evol_path_params(n::Int, u::Vector{Float64})
    lambda = length(u)
    mu = 1/sum(i -> (u[i]+1/lambda)^2, 1:(lambda÷2))
    c = (mu + 2.0)/(sqrt(n)*(n+mu+5.0))
    discount = 1-c
    Zscale = sqrt(c*(2-c)*mu)
    # Expected length of evol.path, if uZ ~ α N(0,I), based on AR model:
    # E|path|^2 = α*Zscale^2 / (1-discount^2),
    # where α ≈ 1 if the population does not cover the minimum and lies on the slope
    # (so that |uZ| is large), but σ is too large and the method is diverging
    threshold = Zscale*sqrt(n/(1-discount^2))
    #info("DX-NES params: moving_threshold=", threshold,
    #     " μ=", mu, " c=", c, " discount=", discount, " Zscale=", Zscale)
    return threshold, discount, Zscale
end

function assign_distance_weights!(weights::Vector{Float64}, scale::Float64,
                                  rankedCandidates::Vector{Candidate{F}},
                                  Z::Matrix{Float64}) where F
    @assert length(weights) == length(rankedCandidates) == size(Z, 2)
    lambda = size(Z, 2)
    u0 = log(0.5*lambda+1)
    for i in 1:lambda
        cand_ix = rankedCandidates[i].index
        weights[cand_ix] = max(u0-log(i), 0) * exp(scale*norm(view(Z, cand_ix)))
    end
    weights .= weights./sum(weights) .- 1/lambda
    return weights
end

function update_learning_rates!(dxnes::DXNESOpt, evol_speed::Float64)
    n = numdims(dxnes)
    if evol_speed > 1.0
        dxnes.sigma_learnrate = 1.0
        dxnes.B_learnrate = (dxnes.lambda + 2n)/(dxnes.lambda+2*n^2+100) *
                            min(1.0, sqrt(dxnes.lambda/numdims(dxnes)))
    else
        dxnes.sigma_learnrate = 1.0 + dxnes.lambda/(dxnes.lambda+2*n)
        if evol_speed > 0.1
            dxnes.sigma_learnrate *= 0.5
        end
        dxnes.B_learnrate = dxnes.lambda/(dxnes.lambda+2*n^2+100)
    end
    return dxnes
end


text:
#########################TimeSeries Data Updating###########################################
function get_simulation_time(sim::Simulation, stage_number::Int64)
    return sim.ref.date_ref[stage_number]
end

function parameter_update!(param_reference::UpdateRef{T},
                           param_array::JuMPParamArray,
                           stage_number::Int64,
                           sim::Simulation) where T <: PSY.Component

    devices = PSY.get_components(T, sim.stages[stage_number].sys)
    initial_forecast_time = get_simulation_time(sim, stage_number)
    horizon = length(model_time_steps(sim.stages[stage_number].canonical))
    for d in devices
        forecast = PSY.get_forecast(PSY.Deterministic,
                                    d,
                                    initial_forecast_time,
                                    "$(param_reference.access_ref)",
                                    horizon)
        ts_vector = TS.values(PSY.get_data(forecast))
        device_name = PSY.get_name(d)
        for (ix, val) in enumerate(param_array[device_name,:])
            value = ts_vector[ix]
            JuMP.fix(val, value)
        end
    end

    return

end

# This makes the choice in which variable to get from the results.
function _get_stage_variable(chron::Type{RecedingHorizon},
                           from_stage::_Stage,
                           device_name::String,
                           var_ref::UpdateRef,
                           to_stage_execution_count::Int64)

    variable = get_value(from_stage.canonical, var_ref)
    step = axes(variable)[2][1]
    return JuMP.value(variable[device_name, step])
end

function _get_stage_variable(chron::Type{Sequential},
                             from_stage::_Stage,
                             device_name::String,
                             var_ref::UpdateRef,
                             to_stage_execution_count::Int64)

    variable = get_value(from_stage.canonical, var_ref)
    step = axes(variable)[2][end]
    return JuMP.value(variable[device_name, step])
end

function _get_stage_variable(chron::Type{Synchronize},
                            from_stage::_Stage,
                            device_name::String,
                            var_ref::UpdateRef,
                            to_stage_execution_count::Int64)

    variable = get_value(from_stage.canonical, var_ref)
    step = axes(variable)[2][to_stage_execution_count + 1]
    return JuMP.value(variable[device_name, step])
end

################################Cache Update################################################
function _update_cache!(c::TimeStatusChange, stage::_Stage)
    parameter = get_value(stage.canonical, c.ref)

    for name in parameter.axes[1]
        param_status = PJ.value(parameter[name])
        if c.value[name][:status] == param_status
            c.value[name][:count] += 1.0
        elseif c.value[name][:status] != param_status
            c.value[name][:count] = 1.0
            c.value[name][:status] = param_status
        end
    end

    return
end

#########################FeedForward Variables Updating#####################################
function feedforward_update(synch::Chron,
                            param_reference::UpdateRef{JuMP.VariableRef},
                            param_array::JuMPParamArray,
                            to_stage::_Stage,
                            from_stage::_Stage) where Chron <: Chronology

    !(to_stage.execution_count % synch.to_steps == 0) && return

    var_count = to_stage.execution_count ÷ synch.to_steps

    for device_name in axes(param_array)[1]
        var_value = _get_stage_variable(Chron, from_stage, device_name, param_reference, var_count)
        PJ.fix(param_array[device_name], var_value)
    end

    return

end

#########################Initial Condition Updating#########################################
# TODO: Consider when more than one UC model is used for the stages that the counts need
# to be scaled.
function _calculate_ic_quantity(initial_condition_key::ICKey{TimeDurationOFF, PSD},
                                ic::InitialCondition,
                                var_value::Float64,
                                cache::Union{Nothing,AbstractCache}) where PSD <: PSY.Device

    name = device_name(ic)
    time_cache = cache_value(cache, name)

    current_counter = time_cache[:count]
    last_status = time_cache[:status]
    var_status = isapprox(var_value, 0.0, atol = 1e-4) ? 0.0 : 1.0
    @assert abs(last_status - var_status) < eps()

    if last_status >= 1.0
        return current_counter
    end

    if last_status < 1.0
        return 0.0
    end
end

function _calculate_ic_quantity(initial_condition_key::ICKey{TimeDurationON, PSD},
                                ic::InitialCondition,
                                var_value::Float64,
                                cache::Union{Nothing,AbstractCache}) where PSD <: PSY.Device

    name = device_name(ic)
    time_cache = cache_value(cache, name)

    current_counter = time_cache[:count]
    last_status = time_cache[:status]
    var_status = isapprox(var_value, 0.0, atol = 1e-4) ? 0.0 : 1.0
    @assert abs(last_status - var_status) < eps()

    if last_status >= 1.0
        return 0.0
    end

    if last_status < 1.0
        return current_counter
    end

end

function _calculate_ic_quantity(initial_condition_key::ICKey{DeviceStatus, PSD},
                                ic::InitialCondition,
                                var_value::Float64,
                                cache::Union{Nothing,AbstractCache}) where PSD <: PSY.Device
    return isapprox(var_value, 0.0, atol = 1e-4) ? 1.0 : 0.0
end

function _calculate_ic_quantity(initial_condition_key::ICKey{DevicePower, PSD},
                               ic::InitialCondition,
                               var_value::Float64,
                               cache::Union{Nothing,AbstractCache}) where PSD <: PSY.ThermalGen

    if isnothing(cache)
        status_change_to_on = value(ic) <= eps() && var_value >= eps()
        status_change_to_off = value(ic) >= eps() && var_value <= eps()
    else
        name = device_name(ic)
        time_cache = cache_value(cache, name)
        status_change_to_on = time_cache[:status] >= eps() && var_value <= eps()
        status_change_to_off = time_cache[:status] <= eps() && var_value >= eps()
    end

    if status_change_to_on
        return ic.device.tech.activepowerlimits.min
    end

    if status_change_to_off
        return 0.0
    end

    return var_value
end

function _initial_condition_update!(initial_condition_key::ICKey,
                                    synch::Chron,
                                    ini_cond_vector::Vector{InitialCondition},
                                    to_stage::_Stage,
                                    from_stage::_Stage) where Chron <: Chronology

    to_stage_execution_count = to_stage.execution_count
    for ic in ini_cond_vector
        name = device_name(ic)
        update_ref = ic.update_ref
        var_value = _get_stage_variable(Chron, from_stage, name, update_ref, to_stage_execution_count)
        cache = get(from_stage.cache, ic.cache, nothing)
        quantity = _calculate_ic_quantity(initial_condition_key, ic, var_value, cache)
        PJ.fix(ic.value, quantity)
    end

    return

end

function _initial_condition_update!(initial_condition_key::ICKey,
                                    ::Nothing,
                                    ini_cond_vector::Vector{InitialCondition},
                                    to_stage::_Stage,
                                    from_stage::_Stage)
        # Meant to do nothing
    return
end

#############################Interfacing Functions##########################################
function cache_update!(stage::_Stage)

    for (_, cache) in stage.cache
        _update_cache!(cache, stage)
    end

    return

end

"""Updates the forecast parameter value"""
function parameter_update!(param_reference::UpdateRef{JuMP.VariableRef},
                           param_array::JuMPParamArray,
                           stage_number::Int64,
                           sim::Simulation)

    chronology_ref = sim.stages[stage_number].chronology_ref
    current_stage = sim.stages[stage_number]
    for (k, ref) in chronology_ref
        feedforward_update(ref, param_reference, param_array, current_stage, sim.stages[k])
    end

    return

end

function intial_condition_update!(initial_condition_key::ICKey,
                                  ini_cond_vector::Vector{InitialCondition},
                                  stage_number::Int64,
                                  step::Int64,
                                  sim::Simulation)

    chronology_ref = nothing
    current_stage = sim.stages[stage_number]
    #checks if current stage is the first in the step and the execution is the first to
    # look backwards on the previous step
    intra_step_update = (stage_number == 1 && current_stage.execution_count == 0)
    #checks if current execution is the first execution to look into the previuous stage
    intra_stage_update = (stage_number > 1 && current_stage.execution_count == 0)
    #checks that the current run and stage ininital conditions is based on the current results
    inner_stage_update = (stage_number > 1 && current_stage.execution_count > 0)
    # makes the update based on the last stage.
    if intra_step_update
        from_stage = sim.stages[end]
        chronology_ref = sim.stages[stage_number].ini_cond_chron
    # Updates the next stage in the same step
    elseif intra_stage_update
        from_stage = sim.stages[stage_number-1]
        chronology_ref = sim.stages[stage_number].chronology_ref[stage_number-1]
    # Update is done on the current stage
    elseif inner_stage_update
        from_stage = current_stage
        chronology_ref = sim.stages[stage_number].ini_cond_chron
    else
        error("Condition not implemented")
    end
    _initial_condition_update!(initial_condition_key,
                               chronology_ref,
                               ini_cond_vector,
                               current_stage,
                               from_stage)

    return

end

function update_stage!(stage::_Stage{M}, step::Int64, sim::Simulation) where M<:AbstractOperationsProblem
    # Is first run of first stage? Yes -> do nothing
    (step == 1 && stage.key == 1 && stage.execution_count == 0) && return
    for (k, v) in stage.canonical.parameters
        parameter_update!(k, v, stage.key, sim)
    end

    cache_update!(stage)

    # Set initial conditions of the stage I am about to run.
    for (k, v) in stage.canonical.initial_conditions
        intial_condition_update!(k, v, stage.key, step, sim)
    end

    return

end


text:
@testset "Indexing" begin
    a1 = AFArray(rand(3, 2))
    b1 = AFArray([1. 2.])
    a1[1, :] = b1
    @test all(@inferred(a1[1, :]) == b1)
    a1[:, :] = 0
    @test all(a1 == 0)
    @test typeof(a1[1,1]) == Float64
    @test @inferred(a1[1,1]) == 0
    @test @inferred(a1[:, 1]) == zeros(AFArray{Float64}, 3)
    @test typeof(a1[:, 2]) == AFArray{Float64, 1}
    a2 = rand(AFArray{Float64}, 4,3,2)
    @test typeof(a2[:, :, 1]) == AFArray{Float64, 2}
    @test typeof(a2[:, 1, 1]) == AFArray{Float64, 1}
end

# @testset "Inplace assignment" begin
#     a = randn(10)
#     af = AFArray(a)
#     c = zeros(a)
#     cf = AFArray(c)
#     df = AFArray(c)

#     a .= a + a
#     af .= af .+ af
#     df .= af
#     c .= af
#     cf .= c

#     @test c == a
#     @test af == cf
#     @test af == df

#     df .= a .+ a
#     @test Array(df) == a + a

#     a .= af .+ af
#     @test a == Array(af .+ af)

#     @test typeof(a) == Array{Float64, 1}
#     @test typeof(af) == AFArray{Float64, 1}
#     @test typeof(c) == Array{Float64, 1}
#     @test typeof(cf) == AFArray{Float64, 1}

#     swap!(cf, af + af)
#     c .= cf
#     @test c == Array(af + af)
# end


text:
module RocketMulticastOperatorTest

using Test
using Rocket

@testset "operator: multicast()" begin

    println("Testing: operator multicast()")

    @testset begin
        subject = Subject(Int)
        source  = from(1:5) |> multicast(subject)

        actor1 = keep(Int)
        actor2 = keep(Int)

        subscription1 = subscribe!(source, actor1)
        subscription2 = subscribe!(source, actor2)

        @test actor1.values == []
        @test actor2.values == []

        connect(source)

        @test actor1.values == [ 1, 2, 3, 4, 5 ]
        @test actor2.values == [ 1, 2, 3, 4, 5 ]
    end

end

end


text:
"""
    SimpleWeightedGraph{T, U}

A type representing an undirected graph with weights of type `U`.

Note that adding or removing vertices or edges is not particularly performant;
see MetaGraphs.jl for possible alternatives.
"""
mutable struct SimpleWeightedGraph{T<:Integer, U<:Real} <: AbstractSimpleWeightedGraph{T, U}
    weights::SparseMatrixCSC{U,T}
    function SimpleWeightedGraph{T, U}(adjmx::SparseMatrixCSC{U, T}) where {T<:Integer, U<:Real}
        dima,dimb = size(adjmx)
        isequal(dima,dimb) || error("Adjacency / distance matrices must be square")
        issymmetric(adjmx) || error("Adjacency / distance matrices must be symmetric")
        new{T, U}(adjmx)
    end

end

ne(g::SimpleWeightedGraph) = (nnz(g.weights) + nselfloop(g)) ÷ 2

function nselfloop(g::SimpleWeightedGraph)
    n = 0
    for i in axes(g.weights, 1)
        n += g.weights[i, i] != 0
    end
    return n
end

SimpleWeightedGraph{T}(adjmx::SparseMatrixCSC{U, T}) where {T <: Integer, U <: Real} =
    SimpleWeightedGraph{T, U}(adjmx)

SimpleWeightedGraph(adjmx::SparseMatrixCSC{U, T}) where {T <: Integer, U <: Real} =
    SimpleWeightedGraph{T, U}(adjmx)

SimpleWeightedGraph(m::AbstractMatrix{U}) where {U <: Real} =
    SimpleWeightedGraph{Int, U}(SparseMatrixCSC{U, Int}(m))
SimpleWeightedGraph{T}(m::AbstractMatrix{U}) where T<:Integer where U<:Real =
    SimpleWeightedGraph{T, U}(SparseMatrixCSC{U, T}(m))
SimpleWeightedGraph{T, U}(m::AbstractMatrix) where T<:Integer where U<:Real =
    SimpleWeightedGraph{T, U}(SparseMatrixCSC{U, T}(m))

SimpleWeightedGraph(g::SimpleWeightedGraph) = SimpleWeightedGraph(copy(g.weights))
function SimpleWeightedGraph{T,U}(g::SimpleWeightedGraph) where {T<:Integer, U<:Real}
    return SimpleWeightedGraph(SparseMatrixCSC{U, T}(copy(g.weights)))
end

# Graph{UInt8}(6), Graph{Int16}(7), Graph{UInt8}()
function (::Type{SimpleWeightedGraph{T, U}})(n::Integer = 0) where T<:Integer where U<:Real
    weights = spzeros(U, T, T(n), T(n))
    return SimpleWeightedGraph{T, U}(weights)
end

# Graph()
SimpleWeightedGraph() = SimpleWeightedGraph(Matrix{Float64}(undef, 0, 0))

# Graph(6), Graph(0x5)
SimpleWeightedGraph(n::T) where T<:Integer = SimpleWeightedGraph{T, Float64}(n)

# Graph(UInt8)
SimpleWeightedGraph(::Type{T}) where T<:Integer = SimpleWeightedGraph{T, Float64}(zero(T))

# Graph(UInt8, Float32)
SimpleWeightedGraph(::Type{T}, ::Type{U}) where {T<:Integer, U<:Real} = SimpleWeightedGraph{T, U}(zero(T))

# Graph(SimpleGraph)

function SimpleWeightedGraph(g::Graphs.AbstractGraph{T}, ::Type{U}=Float64) where {T <: Integer, U <: Real}
    adj_matrix = if Graphs.is_directed(g)
        # TODO abstract function instead of SimpleGraph constructor
        adjacency_matrix(Graphs.SimpleGraphs.SimpleGraph(g), U)
    else
        adjacency_matrix(g, U)
    end
    return SimpleWeightedGraph{T, U}(adj_matrix)
end

function SimpleWeightedGraph(g::Graphs.AbstractGraph{T}, x::U) where {T <: Integer, U <: Real}
    adj_matrix = if Graphs.is_directed(g)
        # TODO abstract function instead of SimpleGraph constructor
        adjacency_matrix(Graphs.SimpleGraphs.SimpleGraph(g), U)
    else
        adjacency_matrix(g, U)
    end
    return SimpleWeightedGraph{T, U}(x .* adj_matrix)
end

# SimpleWeightedGraph{T, U}(SimpleGraph)
function (::Type{SimpleWeightedGraph{T, U}})(g::Graphs.AbstractGraph)  where {T<:Integer, U <: Real}
    adj_matrix = if Graphs.is_directed(g)
        # TODO abstract function instead of SimpleGraph constructor
        adjacency_matrix(Graphs.SimpleGraphs.SimpleGraph{T}(g), U)
    else
        adjacency_matrix(g, U)
    end
    return SimpleWeightedGraph{T, U}(adj_matrix)
end

# Graph(srcs, dsts, weights)
function SimpleWeightedGraph(i::AbstractVector{T}, j::AbstractVector{T}, v::AbstractVector{U}; combine = +) where {T<:Integer, U<:Real}
    m = max(maximum(i), maximum(j))
    s = sparse(vcat(i,j), vcat(j,i), vcat(v,v), m, m, combine)
    SimpleWeightedGraph{T, U}(s)
end

Graphs.SimpleGraph(g::SimpleWeightedGraph) = SimpleGraph(g.weights)

edgetype(::SimpleWeightedGraph{T, U}) where {T<:Integer, U<:Real} = SimpleWeightedGraphEdge{T,U}

edges(g::SimpleWeightedGraph) = (SimpleWeightedEdge(x[1], x[2], x[3]) for x in zip(findnz(triu(g.weights))...))
weights(g::SimpleWeightedGraph) = g.weights
inneighbors(g::SimpleWeightedGraph, x...) = outneighbors(g, x...)

# add_edge! will overwrite weights.
function add_edge!(g::SimpleWeightedGraph, e::SimpleWeightedGraphEdge)
    T = eltype(g)
    U = weighttype(g)
    s_, d_, w = Tuple(e)

    if w == zero(U)
        @warn "Note: adding edges with a zero weight to this graph type has no effect." maxlog=1 _id=:swg_add_edge_zero
        return false
    end

    s = T(s_)
    d = T(d_)
    (s in vertices(g) && d in vertices(g)) || return false
    @inbounds g.weights[d, s] = w
    @inbounds g.weights[s, d] = w
    return true
end

function rem_edge!(g::AbstractSimpleWeightedGraph, e::SimpleWeightedGraphEdge)
    has_edge(g, e) || return false
    U = weighttype(g)
    @inbounds g.weights[dst(e), src(e)] = zero(U)
    @inbounds g.weights[src(e), dst(e)] = zero(U)
    return true
end

==(g::SimpleWeightedGraph, h::SimpleWeightedGraph) = g.weights == h.weights

is_directed(::Type{<:SimpleWeightedGraph}) = false

function Base.getindex(g::SimpleWeightedGraph{T, U}, e::AbstractEdge, ::Val{:weight}) where {T, U, S}
    return g.weights[src(e), dst(e)]
end

function Base.getindex(g::SimpleWeightedGraph{T, U}, i::Integer, j::Integer, ::Val{:weight}) where {T, U, S}
    return g.weights[i, j]
end


text:
# Generate 2D data for the overthrust model
# Author: pwitte.slim@gmail.com
# Date: December 2018
#

using JUDI, HDF5, SegyIO

# Load overthrust model
if ~isfile("$(JUDI.JUDI_DATA)/overthrust_model_2D.h5")
    ftp_data("ftp://slim.gatech.edu/data/SoftwareRelease/WaveformInversion.jl/2DFWI/overthrust_model_2D.h5")
end
n, d, o, m = read(h5open("$(JUDI.JUDI_DATA)/overthrust_model_2D.h5","r"), "n", "d", "o", "m")

# Set up model structure
model = Model((n[1], n[2]), (d[1], d[2]), (o[1], o[2]), m)

# Set up source geometry (cell array with source locations for each shot)
nsrc = 97
xsrc = convertToCell(range(400f0, stop=19600f0, length=nsrc))
ysrc = convertToCell(range(0f0, stop=0f0, length=nsrc))
zsrc = convertToCell(range(50f0, stop=50f0, length=nsrc))

# source sampling and number of time steps
timeS = 3000f0  # ms
dtS = 4f0   # ms

# Set up source structure
src_geometry = Geometry(xsrc, ysrc, zsrc; dt=dtS, t=timeS)

## Set up receivers
xmin = 100f0; xmax = 19900f0
offset = 6000f0
nxrec = 240
xrec = Array{Any}(undef, nsrc)
yrec = Array{Any}(undef, nsrc)
zrec = Array{Any}(undef, nsrc)
for j = 1: nsrc
	# 6 km max. offset around current source
    xsrc[j] - offset < xmin ? xlocal1 = xmin : xlocal1 = xsrc[j] - offset
    xsrc[j] + offset > xmax ? xlocal2 = xmax : xlocal2 = xsrc[j] + offset

    xrec[j] = xlocal1: 50f0: xlocal2
    yrec[j] = range(0f0, stop=0f0, length=length(xrec[j]))
    zrec[j] = range(500f0, stop=500f0, length=length(xrec[j]))
end
timeR = 3000f0	# receiver recording time [ms]
dtR = 4f0	# receiver sampling interval
rec_geometry = Geometry(xrec, yrec, zrec; dt=dtR, t=timeR)

# Set up source
f0 = 0.008
wavelet = ricker_wavelet(timeS, dtS, f0)
#######################################################################

# Set up operators
Pr = judiProjection(rec_geometry)
Ps = judiProjection(src_geometry)
A_inv = judiModeling(model)
q = judiVector(src_geometry, wavelet)

# Generate data
d_obs = Pr*A_inv*adjoint(Ps)*q

block_out = judiVector_to_SeisBlock(d_obs, q; source_depth_key="SourceDepth")
segy_write("overthrust_2D.segy", block_out)


text:
# mbe_demo_mc.jl
#
# This script contains the Monte Carlo runs for the "reduced effort" control
# law developed in mbe_demo.jl.
#
# Navigate to the directory containing this file and run:
#
# julia> include("mbe_demo_mc.jl")
#

# To run the sims in parallel, we'll need to make some new workers.
if nworkers() == 1 # Add them iff we haven't already added some workers.
    addprocs(3) # Add 3 workers.
end

# The workers need to know about our models, so include them everywhere.
@everywhere include("models.jl")

module MBEDemoMC

#############
# Resources #
#############

# Use/import the modules we'll need.
using Eratosthenes   # The simulation engine (setup, simulate, mc)
using ..MBEModels    # This "uses" the MBEModels module defined in models.jl.
import HDF5          # For loading the log file
using Plots          # For plotting our results

# Create a directory for the log file.
if !isdir("out"); mkdir("out"); end

####################
# Monte Carlo Runs #
####################

# Build the scenario from the YAML file.
yaml_file = joinpath(@__DIR__, "underactuated.yaml")
scenario  = setup(yaml_file)

# There's no randomness in our sim by default. Let's randomize the initial
# attitude. The Body model that we're using (the default) has a parameter
# for the standard deviation of initial attitude perturbations. We'll
# overwrite that here and then run the MC with the updated scenario.
scenario.vehicles[1].body.constants.Δq = π/2 * ones(3)

# Choose a number of runs.
n_runs = 50

# Now tell Julia to run the scenarion n times. Each will have a different
# initial attitude.
mc(scenario, n_runs)

# Get the quaternions from the log file for every run.
t = Float64[] # The time steps will be the same for every run.
q = Float64[] # We'll store all of the quaternions here, side by side.
for k = 1:n_runs

    # Load from the log file.
    t, q_BI = HDF5.h5open("out/mc" * string(k) *".h5", "r") do logs
        (read(logs, "/smallsat/body/state/time"),
         read(logs, "/smallsat/body/state/data/q_BI"))
    end

    # Store the time and quaternions.
    t = dropdims(t, dims=1)
    q = vcat(q, q_BI)

end

# Choose a plotting backend.
plotlyjs()

# Show the quaternions from all runs.
println("Plotting all results.")
display(plot(t, transpose(q),
    label  = "",
    xlabel = "Time (s)",
    ylabel = "Attitude Quaternions",
    title  = "Attitude from " * string(n_runs) * " Runs"))

# Next up: software-in-the-loop

end # module


text:
"""
    Sextic{T}

The scalar conservation law with flux
``
  \\partial_t u + \\partial_x u^6 = 0
``
in one space dimensions using `T` as scalar type.
"""
struct Sextic{T} <: ScalarBalanceLaw{T,1} end

function Sextic(T=Float64)
  Sextic{T}()
end

function Base.show(io::IO, model::Sextic{T}) where {T}
  print(io, "Scalar conservation law {T=", T, "} with flux f(u) = u^6")
end

"""
    flux(u, model::Sextic)

Compute the flux of `u` for `model`.
"""
@inline flux(u, model::Sextic) = (u^3)^2

"""
    speed(u::Real, model::Sextic)

Compute the speed f'(`u`) for `model`.
"""
@inline speed(u, model::Sextic) = 6*u^2*u^3

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

"""
    (::GodunovFlux)(uₗ::T, uᵣ::T, model::Sextic{T}) where {T}

Compute Godunov's flux between `uₗ` and `uᵣ` for `model`.
"""
function (::GodunovFlux)(uₗ::T, uᵣ::T, model::Sextic{T}) where {T}
    if uₗ < uᵣ
      if uₗ < 0 && 0 < uᵣ
        zero(T)
      else
        min(flux(uₗ, model), flux(uᵣ, model))
      end
    else
      max(flux(uₗ, model), flux(uᵣ, model))
    end
end

"""
    (::EnergyConservativeFlux)(uₗ::T, uᵣ::T, model::Sextic{T}) where {T}

Compute the energy (L₂ entropy) conservative flux between `uₗ` and `uᵣ` for `model`.
"""
function (flux::EnergyConservativeFlux)(uₗ::T, uᵣ::T, model::Sextic{T}) where T<:Real
    ((uₗ^3)^2 + uₗ^2*uₗ^3*uᵣ + (uₗ^2)^2*uᵣ^2 + uₗ^3*uᵣ^3 + uₗ^2*(uᵣ^2)^2 + uₗ*uᵣ^2*uᵣ^3 + (uᵣ^3)^2) / 7
end

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

@inline function add_numerical_fluxes_inner_loop2!(du, fluxes, u, balance_law::Sextic,
                                                    Nx, basis, boundaries_included::Val{false},
                                                    jacx, parallel)
    Pp1 = length(basis.nodes)
    Rl = basis.interp_left
    Rr = basis.interp_right
    one_7 = 1 / 7
    two_7 = 2 / 7

    # add numerical fluxes
    @inbounds for ix in Base.OneTo(Nx)
        Rul = zero(eltype(u))
        Rur = zero(eltype(u))
        for nx in Base.OneTo(Pp1)
            Rul += Rl[nx]*u[nx,ix]
            Rur += Rr[nx]*u[nx,ix]
        end
        Ru2l = zero(eltype(u))
        Ru2r = zero(eltype(u))
        for nx in Base.OneTo(Pp1)
            tmp = u[nx,ix]^2
            Ru2l += Rl[nx]*tmp
            Ru2r += Rr[nx]*tmp
        end
        Ru3l = zero(eltype(u))
        Ru3r = zero(eltype(u))
        for nx in Base.OneTo(Pp1)
            tmp = u[nx,ix]^3
            Ru3l += Rl[nx]*tmp
            Ru3r += Rr[nx]*tmp
        end
        Ru4l = zero(eltype(u))
        Ru4r = zero(eltype(u))
        for nx in Base.OneTo(Pp1)
            tmp = (u[nx,ix]^2)^2
            Ru4l += Rl[nx]*tmp
            Ru4r += Rr[nx]*tmp
        end
        Ru5l = zero(eltype(u))
        Ru5r = zero(eltype(u))
        for nx in Base.OneTo(Pp1)
            tmp = u[nx,ix]^2*u[nx,ix]^3
            Ru5l += Rl[nx]*tmp
            Ru5r += Rr[nx]*tmp
        end
        Ru6l = zero(eltype(u))
        Ru6r = zero(eltype(u))
        for nx in Base.OneTo(Pp1)
            tmp = (u[nx,ix]^3)^2
            Ru6l += Rl[nx]*tmp
            Ru6r += Rr[nx]*tmp
        end

        for nx in Base.OneTo(Pp1)
            du[nx,ix] += (  (fluxes[ix  ]   - two_7 * Ru6l
                                            - two_7 * u[nx,ix] * Ru5l
                                            - two_7 * u[nx,ix]^2 * Ru4l
                                            - one_7 * (Rul^3)^2
                                            + one_7 * u[nx,ix] * Rul^3*Rul^2
                                            - one_7 * Ru2l * (Rul^2)^2
                                            + one_7 * u[nx,ix]^2 * (Rul^2)^2
                                            - one_7 * Ru3l * Rul^3
                                            + one_7 * u[nx,ix] * Ru3l * Rul^2
                                            - one_7 * Ru3l * Ru2l * Rul
                                            + one_7 * u[nx,ix]^2 * Ru3l * Rul
                                            - one_7 * Ru3l^2) * Rl[nx]
                          - (fluxes[ix+1]   - two_7 * Ru6r
                                            - two_7 * u[nx,ix] * Ru5r
                                            - two_7 * u[nx,ix]^2 * Ru4r
                                            - one_7 * (Rur^3)^2
                                            + one_7 * u[nx,ix] * Rur^3*Rur^2
                                            - one_7 * Ru2r * (Rur^2)^2
                                            + one_7 * u[nx,ix]^2 * (Rur^2)^2
                                            - one_7 * Ru3r * Rur^3
                                            + one_7 * u[nx,ix] * Ru3r * Rur^2
                                            - one_7 * Ru3r * Ru2r * Rur
                                            + one_7 * u[nx,ix]^2 * Ru3r * Rur
                                            - one_7 * Ru3r^2) * Rr[nx]
                            ) * jacx / basis.weights[nx]
        end
    end

    nothing
end


text:
# See p. 280 of Murphy's Machine Learning
# x_k1 = x_k - alpha * gr + mu * (x - x_previous)

struct MomentumGradientDescent{Tf, IL,L} <: FirstOrderOptimizer
    mu::Tf
    alphaguess!::IL
    linesearch!::L
    manifold::Manifold
end

Base.summary(::MomentumGradientDescent) = "Momentum Gradient Descent"

function MomentumGradientDescent(; mu::Real = 0.01,
                                 alphaguess = LineSearches.InitialPrevious(), # TODO: investigate good defaults
                                 linesearch = LineSearches.HagerZhang(),        # TODO: investigate good defaults
                                 manifold::Manifold=Flat())
    MomentumGradientDescent(mu, _alphaguess(alphaguess), linesearch, manifold)
end

mutable struct MomentumGradientDescentState{Tx, T} <: AbstractOptimizerState
    x::Tx
    x_previous::Tx
    x_momentum::Tx
    f_x_previous::T
    s::Tx
    @add_linesearch_fields()
end

function initial_state(method::MomentumGradientDescent, options, d, initial_x)
    T = eltype(initial_x)
    initial_x = copy(initial_x)
    retract!(method.manifold, initial_x)

    value_gradient!!(d, initial_x)

    project_tangent!(method.manifold, gradient(d), initial_x)

    MomentumGradientDescentState(initial_x, # Maintain current state in state.x
                                 copy(initial_x), # Maintain previous state in state.x_previous
                                 similar(initial_x), # Record momentum correction direction in state.x_momentum
                                 real(T)(NaN), # Store previous f in state.f_x_previous
                                 similar(initial_x), # Maintain current search direction in state.s
                                 @initial_linesearch()...)
end

function update_state!(d, state::MomentumGradientDescentState, method::MomentumGradientDescent)
    project_tangent!(method.manifold, gradient(d), state.x)
    # Search direction is always the negative gradient
    state.s .= .-gradient(d)

    # Update position, and backup current one
    state.x_momentum .= state.x .- state.x_previous

    # Determine the distance of movement along the search line
    lssuccess = perform_linesearch!(state, method, ManifoldObjective(method.manifold, d))

    state.x .+= state.alpha.*state.s .+ method.mu.*state.x_momentum
    retract!(method.manifold, state.x)
    lssuccess == false # break on linesearch error
end

function trace!(tr, d, state, iteration, method::MomentumGradientDescent, options, curr_time=time())
  common_trace!(tr, d, state, iteration, method, options, curr_time)
end

function default_options(method::MomentumGradientDescent)
    Dict(:allow_f_increases => true)
end