include(joinpath(@__DIR__, "..", "src/dynamics", "flamingo", "visuals.jl"))
T = Float64
vis = Visualizer()
open(vis)

s = get_simulation("flamingo", "flat_2D_lc", "flat")
model = s.model
env = s.env
const ContactImplicitMPC = Main
ref_traj_ = deepcopy(ContactImplicitMPC.get_trajectory(s.model, s.env,
    joinpath(module_dir(), "src/dynamics/flamingo/gaits/gait_forward_36_4.jld2"),
    load_type = :split_traj_alt))
ref_traj = deepcopy(ref_traj_)

H = ref_traj.H
h = ref_traj.h
N_sample = 5
H_mpc = 15
h_sim = h / N_sample
H_sim = 1000 #35000

p = pratt_policy(model, h_sim)

# q1_ref = copy(ref_traj.q[2])
# q0_ref = copy(ref_traj.q[1])
q1_ref = [0.0, 0.849, -0.00, 0.1, 0.295, -0.3, 0.1, π/2, π/2]
q0_ref = [0.0, 0.849, -0.00, 0.1, 0.295, -0.3, 0.1, π/2, π/2]
q1_sim = SVector{model.dim.q}(q1_ref)
q0_sim = SVector{model.dim.q}(copy(q1_sim - (q1_ref - q0_ref) / N_sample))

@assert norm((q1_sim - q0_sim) / h_sim - (q1_ref - q0_ref) / h) < 1.0e-8

sim = simulator(s, q0_sim, q1_sim, h_sim, H_sim,
    p = p,
	ip_opts = InteriorPointOptions(
		undercut = Inf,
		γ_reg = 0.0,
		# verbose = true,
        r_tol = 1.0e-8,
        κ_tol = 1.0e-8),
    sim_opts = SimulatorOptions(warmstart = true),
    )

@time status = simulate!(sim, verbose = true)
l = 9
lu = 1
plt = plot(layout=(3,1), legend=false)
plot!(plt[1,1], hcat(Vector.(vcat([fill(ref_traj.q[i], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[1,1], hcat(Vector.([q[l:l] for q in sim.traj.q])...)', color=:blue, linewidth=1.0)
plot!(plt[2,1], hcat(Vector.(vcat([fill(ref_traj.u[i][lu:lu], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[3,1], hcat(Vector.(vcat([fill(ref_traj.γ[i][1:nc], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[2,1], hcat(Vector.([u[lu:lu] for u in sim.traj.u]*N_sample)...)', color=:blue, linewidth=1.0)
# plot!(plt[3,1], hcat(Vector.([γ[1:nc] for γ in sim.traj.γ]*N_sample)...)', color=:blue, linewidth=1.0)
# plot!(plt[3,1], hcat(Vector.([b[1:nb] for b in sim.traj.b]*N_sample)...)', color=:red, linewidth=1.0)

plot_surface!(vis, env, xlims=[-0.5, 1.5], ylims = [-0.5, 0.5])
plot_lines!(vis, model, sim.traj.q)
# anim = visualize_robot!(vis, model, sim.traj, sample=10)
anim = visualize_meshrobot!(vis, model, sim.traj, sample=2)
anim = visualize_force!(vis, model, env, sim.traj, anim=anim, h=h_sim, sample=10)

# filename = "flamingo_100_steps"
# MeshCat.convert_frames_to_video(
#     "/home/simon/Downloads/$filename.tar",
#     "/home/simon/Documents/$filename.mp4", overwrite=true)
#
# convert_video_to_gif(
#     "/home/simon/Documents/$filename.mp4",
#     "/home/simon/Documents/$filename.gif", overwrite=true)
#
#
# settransform!(vis["/Cameras/default"],
# 		compose(Translation(0.0, 0.5, -1.0),LinearMap(RotZ(-pi / 2.0))))
#
# # Ghost
# flamingo_ghost!(vis, sim, x -> 0.0)
#
# # Animation
# flamingo_animation!(vis, sim, x -> 0.0)

qref = [0.0, 0.849, -0.00, 0.1, 0.295, -0.3, 0.1, π/2, π/2]
set_robot!(vis, model, qref)
qref = [0.0, 0.849, -0.00, -0.3, 0.1, 0.1, 0.295, π/2, π/2]
set_robot!(vis, model, qref)


#JOB FILEPATHS
#sample record and hmm serialisation output path
sample_output = "/bench/PhD/NGS_binaries/BBM/survey_samples"
hmm_output = "/bench/PhD/NGS_binaries/BBM/survey_chains"

#GENERAL SETUP
@info "Loading libraries..."
using BioBackgroundModels, DataFrames, Distributed, Serialization
include("/srv/git/rys_nucleosomes/aws/aws_wrangler.jl")

#JOB CONSTANTS
const replicates = 3 #repeat optimisation from this many seperately initialised samples from the prior
const Ks = [1,2,4,6] #mosaic class #s to test
const order_nos = [0,1,2] #DNA kmer order #s to test
const delta_thresh=1e-3 #stopping/convergence criterion (log probability difference btw subsequent EM iterates)
const max_iterates=15000

#LOAD SAMPLES
@info "Loading samples from $sample_output..."
sample_dfs = deserialize(sample_output)

#BUILD TRAINING AND TEST SETS FROM SAMPLES
training_sets, test_sets = split_obs_sets(sample_dfs)

#PROGRAMATICALLY GENERATE Chain_ID Vector
job_ids=Vector{Chain_ID}()
for (obs_id, obs) in training_sets, K in Ks, order in order_nos, rep in 1:replicates
    push!(job_ids, Chain_ID(obs_id, K, order, rep))
end

#DISTRIBUTED CLUSTER CONSTANTS
load_dict=Dict{Int64,LoadConfig}()
local_config=LoadConfig(1:6,2:2)
remote_config=LoadConfig(1:6,0:1)
aws_instance_config=LoadConfig(1:6,0:2)
remote_machine = "10.0.0.3"
no_local_processes = 1
no_remote_processes = 1
#SETUP DISTRIBUTED BAUM WELCH LEARNERS
@info "Spawning local cluster workers..."
worker_pool=addprocs(no_local_processes, topology=:master_worker)
for worker in worker_pool
    load_dict[worker]=local_config
end

remote_pool=addprocs([(remote_machine,no_remote_processes)], tunnel=true, topology=:master_worker)

for worker in remote_pool
    load_dict[worker]=remote_config
end

worker_pool=vcat(worker_pool, remote_pool)

# #AWS PARAMS
# @info "Setting up AWS wrangling..."

# security_group_name="calc1"
# security_group_desc="calculation group"
# ami="ami-02df7e1881a08f163"
# skeys="AWS"
# instance_type="c5.4xlarge"
# zone,spot_price=get_cheapest_zone(instance_type)
# no_instances=1
# instance_workers=8
# bid=spot_price+.01

# @assert bid >= spot_price

# @info "Wrangling AWS instances..."
# aws_ips = spot_wrangle(no_instances, bid, security_group_name, security_group_desc, skeys, zone, ami, instance_type)
# @info "Giving instances 90s to boot..."
# sleep(90)

# aws_ips = ["18.223.122.201"]

# @info "Spawning AWS cluster workers..."
# for ip in aws_ips
#     instance_pool=addprocs([(ip, instance_workers)], tunnel=true, topology=:master_worker, sshflags="-o StrictHostKeyChecking=no")
#     for worker in instance_pool
#         load_dict[worker]=aws_instance_config
#     end
#     global worker_pool=vcat(worker_pool, instance_pool)
# end

@info "Loading worker libraries everywhere..."
@everywhere using BioBackgroundModels
#INTIIALIZE HMMS
@info "Setting up HMMs..."
if isfile(hmm_output) #if some results have already been collected, load them
    @info "Loading incomplete results..."
    hmm_results_dict = deserialize(hmm_output)
else #otherwise, pass a new results dict
    @info "Initialising new HMM results file at $hmm_output"
    hmm_results_dict = Dict{Chain_ID,Vector{EM_step}}()
end

em_jobset = setup_EM_jobs!(job_ids, training_sets; chains=hmm_results_dict)
execute_EM_jobs!(worker_pool, em_jobset..., hmm_output; load_dict=load_dict, delta_thresh=delta_thresh, max_iterates=max_iterates)


abstract type ReadCounter end
abstract type BiasModel end

function pushto!(t::IntervalTrees.IntervalBTree{K, V, B}, key::I, value) where {K, V, B, I <: AbstractInterval{K}}
    return _push!(t.root, key, value)
end

function _push!(t::IntervalTrees.InternalNode{K, V, B}, key::AbstractInterval{K}, value) where {K, V, B}
    i = IntervalTrees.findidx(t, key)
    if 1 <= length(t) - 1 && key >= t.keys[i]
        return _push!(t.children[i+1], key, value)
    else
        return _push!(t.children[i], key, value)
    end
end

function _push!(t::IntervalTrees.LeafNode{K, V, B}, key::AbstractInterval{K}, value) where {K, V, B}
    i = IntervalTrees.findidx(t, key)
    if 1 <= i <= length(t) &&
        first(t.entries[i]) == first(key) && last(t.entries[i]) == last(key)
        push!( t.entries[i].value, value )
        return true
    end
    return false
end

function pushzero!( collection::Dict{K,V}, key, value ) where {K, V <: ReadCounter}
   if !haskey( collection, key )
      collection[key] = zero(V)
   end
   push!( collection[key], value )
end

function pushzero!( collection::IntervalMap{K,V}, key, value ) where {K <: Integer, V <: ReadCounter}
   if !haskey( collection, (key.first,key.last) )
      collection[key] = zero(V)
   end
   pushto!( collection, key, value )
end

function Base.get( cnt::R ) where R <: ReadCounter
   if !cnt.isadjusted
      error("Cannot get() read count from model thats not already adjusted!")
   else
      return cnt.count
   end
end

function Base.push!( cnt::R, value::F ) where {R <: ReadCounter, F <: AbstractFloat}
   if cnt.isadjusted
      cnt.count += value
   else
      cnt.count += value
      cnt.isadjusted = true
      #error("Cannot push! a float value to PrimerBiasCounter unless isadjusted is true!")
   end
   cnt
end

Base.identity( cnt::R ) where R <: ReadCounter = get( cnt )
Base.identity( cnt::IntervalValue{I,R} ) where R <: ReadCounter = get( cnt.value )

struct DefaultBiasMod <: BiasModel end

mutable struct DefaultCounter <: ReadCounter
   count::Float64
   isadjusted::Bool

   DefaultCounter() = new( 1.0, true )
   DefaultCounter( i::Float64 ) = new( i, true )
end

const DEFAULTCOUNTER_ZERO = DefaultCounter(0.0)
const DEFAULTCOUNTER_ONE  = DefaultCounter(1.0)

Base.zero( ::Type{DefaultCounter} ) = DefaultCounter(0.0) #DEFAULTCOUNTER_ZERO
Base.one( ::Type{DefaultCounter} )  = DefaultCounter(1.0) #DEFAULTCOUNTER_ONE
Base.push!( cnt::DefaultCounter, value::Float64 ) = (cnt.count += value)

primer_adjust!( cnt::DefaultCounter, m::B ) where B <: BiasModel = (cnt.isadjusted = true)
gc_adjust!( cnt::DefaultCounter, m::B ) where B <: BiasModel = (cnt.isadjusted = true)

count!( mod::DefaultBiasMod, _ ) = 1.0
count!( mod::DefaultBiasMod, x, y ) = 1.0

# 5' Sequence bias model
# Implementation of Hansen et al, NAR 2010

# hash by identity for speed
Base.hash( v::UInt16 ) = v

function add_to_back!( back::Vector{Float64}, exp::ExpectedGC, weight::Float64 )
   for i in 1:length(exp)
      back[i] += exp[i] * weight
   end
   mod
end

primer_normalize!( mod::DefaultBiasMod )   = 1.0
gc_normalize!( mod::DefaultBiasMod, x, y ) = 1.0

struct JointBiasTemp
   primer::Vector{UInt16}
   gc::Vector{Int32}
end

struct JointBiasMod <: BiasModel
   fore::Vector{Float64}
   back::Vector{Float64}
   size::Int64
   backoffset::Int64
   backnpos::Int64
   foreoffset::Int64
   forenpos::Int64

   gcfore::Vector{Float64}
   gcback::Vector{Float64}
   gcsize::Int64
   gcwidth::Float64
   gcoffset::Int64
   gcincrement::Int64

   temp::JointBiasTemp

   function JointBiasMod( ksize::Int=6, backoff::Int=23, backn::Int=10, foreoff::Int=1, foren::Int=2,
                          gcoffset::Int=5, gcincrement::Int=10; gc_param = GCBiasParams(50, 0.05) )
      fore = zeros( 4^ksize )
      back = zeros( 4^ksize )
      primer_temp = zeros(UInt16, foren)

      gcfore = zeros(Float64, Int(ceil(1.0 / gc_param.width)+1))
      gcback = zeros(Float64, Int(ceil(1.0 / gc_param.width)+1))
      gctemp = zeros(Int32, 100)

      temp = JointBiasTemp( primer_temp, gctemp )

      new( fore, back, ksize, backoff, backn, foreoff, foren,
           gcfore, gcback, gc_param.length, gc_param.width, gcoffset, gcincrement,
           temp )
   end
end

function primer_normalize!( mod::JointBiasMod )
   foresum = sum(mod.fore)
   backsum = sum(mod.back)
   @inbounds for i in 1:length(mod.fore)
      @fastmath mod.fore[i] = mod.fore[i] / foresum
      @fastmath mod.back[i] = mod.back[i] / backsum
   end
end

function primer_count!( mod::JointBiasMod, seq::BioSequence{A} ) where A <: BioSequences.Alphabet
   # 5 prime bias
   idx = 1
   for i in mod.foreoffset:(mod.foreoffset+mod.forenpos-1)
      hept = kmer_index_trailing(UInt16, seq[i:(i+mod.size-1)])
      mod.fore[hept+1] += 1.0
      mod.temp.primer[idx] = hept
      idx += 1
   end
   for i in mod.backoffset:(mod.backoffset+mod.backnpos-1)
      mod.back[kmer_index_trailing(UInt16, seq[i:(i+mod.size-1)])+1] += 1.0
   end

   return mod.temp
end

function gc_count!( mod::JointBiasMod, seq::BioSequence{A} ) where A <: BioSequences.Alphabet
   # GC fragment bias
   idx = 1
   for i in mod.gcoffset:mod.gcincrement:(length(seq)-mod.gcsize)
      gc = gc_content(seq[i:(i+mod.gcsize-1)])
      bin = Int(div( gc, mod.gcwidth )+1)
      mod.gcfore[bin] += 1.0
      mod.temp.gc[idx] = bin
      idx += 1
   end

   return mod.temp
end

function count!( mod::JointBiasMod, seq::BioSequence{A} ) where A <: BioSequences.Alphabet
   # 5 prime bias
   primer_count!( mod, seq )
   # GC fragment bias
   fill!(mod.temp.gc, zero(Int32))
   gc_count!( mod, seq )

   return mod.temp
end

function count!( mod::JointBiasMod, fwd::BioSequence{A}, rev::BioSequence{A} ) where A <: BioSequences.Alphabet
   # 5 prime bias
   primer_count!( mod, fwd )
   # GC fragment bias
   fill!(mod.temp.gc, zero(Int32))
   gc_count!( mod, fwd )
   gc_count!( mod, rev )

   return mod.temp
end

@fastmath value!( mod::JointBiasMod, hept::UInt16 ) = mod.back[hept+1] / mod.fore[hept+1]
@fastmath value!( mod::JointBiasMod, bin::Int ) = mod.fore[bin] > 0.0 ? mod.back[bin] / mod.fore[bin] : 1.0

mutable struct JointBiasCounter <: ReadCounter
   count::Float64
   map::Dict{UInt16,Int32}
   gc::Vector{Int32}
   isadjusted::Bool

   function JointBiasCounter( gc_param = GCBiasParams(50, 0.05) )
      gc = zeros(Int32, Int(ceil(1.0 / gc_param.width)+1))
      new(0.0, Dict{UInt16,Int32}(), gc, false)
   end
end

Base.zero(::Type{JointBiasCounter}) = JointBiasCounter()

function Base.push!( cnt::JointBiasCounter, temp::JointBiasTemp )
   for i in 1:length(temp.primer)
      push!( cnt, temp.primer[i] )
   end
   for i in 1:length(temp.gc)
      if temp.gc[i] > 0
         push!( cnt, temp.gc[i] )
      else
         return cnt
      end
   end
   cnt
end

function Base.push!( cnt::JointBiasCounter, bin::Int32 )
   cnt.gc[bin] += one(Int32)
   cnt
end

function Base.push!( cnt::JointBiasCounter, hept::UInt16 )
   if haskey( cnt.map, hept )
      cnt.map[hept] += 1
   else
      cnt.map[hept] = 1
   end
   cnt
end

adjust!( cnt::JointBiasCounter, mod::JointBiasMod ) = !cnt.isadjusted ? primer_adjust!( cnt, mod ) : gc_adjust!( cnt, mod )

@inline function primer_adjust!( cnt::JointBiasCounter, mod::JointBiasMod )
   for (k,v) in cnt.map
      cnt.count += value!( mod, k ) * v
   end
   cnt.count /= mod.forenpos
   cnt.isadjusted = true
end

@inline function gc_adjust!( cnt::JointBiasCounter, mod::JointBiasMod )
   weight = 0.0
   for i in 1:length(cnt.gc)
      weight += value!( mod, i ) * cnt.gc[i]
   end
   if sum(cnt.gc) > 0
      weight /= sum(cnt.gc)
      cnt.count *= weight
   end
   cnt
end


@with_kw struct TimePoint
    timestamp::Float64                  = 0.0
    coordinates::SVector{2,Float64}     = [0.0, 0.0]
end

"""
A Bernoulli model for grasp success, given the number of attempts made.
"""
struct DiscreteGPM
    bernoulli_prob::Float64
end
cumprob(dgpm::DiscreteGPM, steps::Int64) = 1.0 - (1 - dgpm.bernoulli_prob)^steps

"""
Immutable arm parameters.
"""
@with_kw struct Arm
    base_pos::SVector{2,Float64}
    bin_pos::SVector{2,Float64}
    reach::SVector{2,Float64}
    orig_pos::SVector{2,Float64}
    grasp_model::DiscreteGPM
    downtime::Float64
end

"""
Mutable arm properties.
"""
@with_kw mutable struct ArmProperties
    endeff_pos::SVector{2,Float64}                  = [0.0, 0.0]
    obj_grabbed::String                             = ""
    workspace_objs::Set{String}                     = Set{String}()
    interaction_events::Vector{InteractionEvent}    = Vector{InteractionEvent}(undef, 0)
    tree::SearchTree                                = SearchTree()
end

"""
Sets up the reflected generator for feasible object sequences.
"""
@enum GENSTAGE PREDROP=1 POSTDROP=2 ATBIN=3
@with_kw mutable struct MirrorGenerator
    arm_set::Vector{Arm}                    = Vector{Arm}(undef, 0)
    gen_box_id::Int64                       = 0
    box_pos::Dict{Int64,Float64}            = Dict{Int64,Float64}()
    arm_stages::Vector{GENSTAGE}            = Vector{GENSTAGE}(undef, 0)
end

"""
The environment object for the domain.
"""
@with_kw mutable struct ConveyorBelt <: MRTAEnvironment
    current_time::Float64                       = 0.0
    num_active_boxes::Int64                     = 0
    num_total_boxes::Int64                      = 0
    curr_speed::Float64                         = 0.6
    box_routes::Dict{String,Vector{TimePoint}}  = Dict{String,Vector{TimePoint}}()
    belt_start::Float64                         = 0.0
    belt_end::Float64                           = 1.0
    belt_width::Float64                         = 0.1
    new_box_prob::Float64                       = 0.75
    lost_boxes::Int64                           = 0
    deposited_boxes::Int64                      = 0
    mirror_gen::MirrorGenerator                 = MirrorGenerator()
end
const ArmAllocationServer = GenericAllocation{Arm,ArmProperties}


# Physical constants
const speed_of_light_mps = 299792458.0 # m/s
export speed_of_light_mps

# Limit data size/memory usage
const max_num_pixels_per_order = 8096
const max_num_orders = 128
const max_template_length = max_num_orders * max_num_pixels_per_order
const max_num_obs = 1024
const max_num_basis_vectors = 16
const min_num_star_basis_vectors = 1
const max_num_star_basis_vectors = max_num_basis_vectors
const min_num_telluric_basis_vectors = 0
const max_num_telluric_basis_vectors = max_num_basis_vectors
export max_num_pixels_per_order, max_num_orders, max_template_length, max_num_obs
export max_num_basis_vectors, min_num_star_basis_vectors, max_num_star_basis_vectors, min_num_telluric_basis_vectors, max_num_telluric_basis_vectors

# TODO: move to each instrument?
const default_blaze_degree = 4
const default_pixel_buffer_lo = 0
const default_pixel_buffer_hi = 0
export default_blaze_degree, default_pixel_buffer_lo, default_pixel_buffer_hi


using DependencyParser.DependencyParsing

function build_model(settings::Settings)
  word_ids = Dict{String, Integer}(
    "here"=>1,
    "some"=>2,
    "words"=>3,
    DependencyParser.DependencyParsing.UNKNOWN_TOKEN => 4,
    DependencyParser.DependencyParsing.NULL_TOKEN => 5,
    DependencyParser.DependencyParsing.ROOT_TOKEN => 6
  )
  tag_ids = Dict{String, Integer}(
    "TAG1"=>7,
    "TAG2"=>8,
    "TAG3"=>9,
    DependencyParser.DependencyParsing.UNKNOWN_TOKEN => 10,
    DependencyParser.DependencyParsing.NULL_TOKEN => 11,
    DependencyParser.DependencyParsing.ROOT_TOKEN => 12,
  )
  label_ids = Dict{String, Integer}(
    "LABEL1"=>13,
    "LABEL2"=>14,
    DependencyParser.DependencyParsing.NULL_TOKEN => 15
  )

  embeddings = rand(Float64, 15, settings.embeddings_size)
  hidden_bias = rand(Float64, settings.hidden_size)
  hidden_layer = Dense(settings.embeddings_size * settings.batch_size, settings.hidden_size, bias=hidden_bias)
  output_layer = Dense(settings.hidden_size, 10)

  Model(embeddings, hidden_layer, output_layer, word_ids, tag_ids, label_ids)
end

function build_correct_batch(model::Model, config::Configuration)
  stack = config.stack |> collect

  Integer[
    stack[begin] |> word_id -> model.word_ids[config.sentence.tokens[word_id].name],
    stack[begin + 1] |> word_id -> model.word_ids[config.sentence.tokens[word_id].name],
    model.word_ids[ROOT_TOKEN],
    config.buffer[begin] |> word_id -> model.word_ids[config.sentence.tokens[word_id].name],
    config.buffer[begin + 1] |> word_id -> model.word_ids[config.sentence.tokens[word_id].name],
    config.buffer[begin + 2] |> word_id -> model.word_ids[config.sentence.tokens[word_id].name],
    model.word_ids[NULL_TOKEN],
    model.word_ids[NULL_TOKEN],
    model.word_ids[NULL_TOKEN],
    model.word_ids[NULL_TOKEN],
    model.word_ids[NULL_TOKEN],
    model.word_ids[NULL_TOKEN],
    model.word_ids[config.sentence.tokens[1].name],
    model.word_ids[config.sentence.tokens[3].name],
    model.word_ids[NULL_TOKEN],
    model.word_ids[NULL_TOKEN],
    model.word_ids[NULL_TOKEN],
    model.word_ids[NULL_TOKEN],
    stack[begin] |> word_id -> model.tag_ids[config.sentence.pos_tags[word_id].name],
    stack[begin + 1] |> word_id -> model.tag_ids[config.sentence.pos_tags[word_id].name],
    model.tag_ids[ROOT_TOKEN],
    config.buffer[begin] |> word_id -> model.tag_ids[config.sentence.pos_tags[word_id].name],
    config.buffer[begin + 1] |> word_id -> model.tag_ids[config.sentence.pos_tags[word_id].name],
    config.buffer[begin + 2] |> word_id -> model.tag_ids[config.sentence.pos_tags[word_id].name],
    model.tag_ids[NULL_TOKEN],
    model.tag_ids[NULL_TOKEN],
    model.tag_ids[NULL_TOKEN],
    model.tag_ids[NULL_TOKEN],
    model.tag_ids[NULL_TOKEN],
    model.tag_ids[NULL_TOKEN],
    model.tag_ids[config.sentence.pos_tags[1].name],
    model.tag_ids[config.sentence.pos_tags[3].name],
    model.tag_ids[NULL_TOKEN],
    model.tag_ids[NULL_TOKEN],
    model.tag_ids[NULL_TOKEN],
    model.tag_ids[NULL_TOKEN],
    model.label_ids[NULL_TOKEN],
    model.label_ids[NULL_TOKEN],
    model.label_ids[NULL_TOKEN],
    model.label_ids[NULL_TOKEN],
    model.label_ids[NULL_TOKEN],
    model.label_ids[NULL_TOKEN],
    model.label_ids["SBJ"],
    model.label_ids["IOBJ"],
    model.label_ids[NULL_TOKEN],
    model.label_ids[NULL_TOKEN],
    model.label_ids[NULL_TOKEN],
    model.label_ids[NULL_TOKEN],
  ]
end


#
# Interface for PIPS-NLP parallel (structured interface)
#

module PipsNlpInterface
Base.include(PipsNlpInterface, "pips_parallel_cfunc.jl")

# using PipsNlpSolver
using StructJuMP, JuMP
using StructJuMPSolverInterface
using Printf
using SparseArrays

import MathProgBase

mutable struct MatStorage
    rowIdx::Vector{Int}
    colIdx::Vector{Int}
    value::Vector{Float64}
    m::Int
    n::Int
    isVal::Bool
    function MatStorage(I,J,m,n)
        return new(I,J,Vector{Float64}(undef, length(I)),m,n,false)
    end
end

# type MatStrStorage
#     rowIdx::Vector{Int}
#     colIdx::Vector{Int}
#     function MatStrStorage()
#         return new(Vector{Int}(),Vector{Int}())
#     end
#     function MatStrStorage(I,J)
#         return new(I,J)
#     end
# end

mutable struct StructJuMPModel <: ModelInterface
    internalModel::JuMP.Model
    status::Int
    n_iter::Int
    prof::Bool
    evaluatorMap::Dict{Int,JuMP.NLPEvaluator}
    iMap::Dict{Int,Pair{Dict{Int,Int},Dict{Int,Int}}}  #eq, ieq, jump->actual used
    jMap::Dict{Pair{Int,Int},Dict{Int,Int}}
    hcMap::Dict{Pair{Int,Int},Dict{Int,Int}}
    hrMap::Dict{Pair{Int,Int},Dict{Int,Int}}
    matJac::Dict{Int,MatStorage}
    matHess::Dict{Int,MatStorage}
    matEqJac::Dict{Pair{Int,Int},MatStorage}
    matIeqJac::Dict{Pair{Int,Int},MatStorage}
    matHessMap::Dict{Pair{Int,Int},MatStorage}
    t_jump::Float64
    t_jac_spconv::Float64
    t_hess_spconv::Float64
    t_hess00_spconv::Float64
    t_eval_g::Float64
    t_itermap_eval_g::Float64
    n_jac_spconv::Int
    n_hess_spconv::Int
    n_hess00_spconv::Int
    n_eval_g::Int

    get_num_scen::Function
    get_sense::Function
    get_status::Function

    get_num_rows::Function
    get_num_cols::Function
    get_num_eq_cons::Function
    get_num_ineq_cons::Function
    get_num_eq_lcons::Function
    get_num_ineq_lcons::Function

    set_status::Function

    set_num_rows::Function
    set_num_cols::Function
    set_num_eq_cons::Function
    set_num_ineq_cons::Function

    str_init_x0::Function
    str_prob_info::Function
    str_eval_f::Function
    str_eval_g::Function
    str_eval_grad_f::Function
    str_eval_jac_g::Function
    str_eval_h::Function
    str_write_solution::Function

    function StructJuMPModel(model::JuMP.Model,prof=false)
        instance = new(model,0,-1,prof,
            Dict{Int,JuMP.NLPEvaluator}(),
            Dict{Int,Pair{Dict{Int,Int},Dict{Int,Int}}}(),
            Dict{Pair{Int,Int},Dict{Int,Int}}(),
            Dict{Pair{Int,Int},Dict{Int,Int}}(),
            Dict{Pair{Int,Int},Dict{Int,Int}}(),
            Dict{Int,MatStorage}(), #jac mat
            Dict{Int,MatStorage}(), #hess mat
            Dict{Pair{Int,Int},MatStorage}(),
            Dict{Pair{Int,Int},MatStorage}(),
            Dict{Pair{Int,Int},MatStorage}(),
            0.0,0.0,0.0,0.0,0.0,0.0,
            0,0,0,0
            )

        initialization(instance)

        instance.get_num_scen = function()
            return num_scenarios(instance.internalModel)
        end
        instance.get_sense = function()
            return getobjectivesense(instance.internalModel)
        end
        instance.get_status = function()
            return instance.status
        end
        instance.get_num_rows = function(id::Int)
            return getNumCons(instance.internalModel,id)
        end
        instance.get_num_cols = function(id::Int)
            return getNumVars(instance.internalModel,id)
        end
        instance.get_num_eq_cons = function(id::Int)
            return length(instance.iMap[id][1])
        end
        instance.get_num_ineq_cons = function(id::Int)
            return length(instance.iMap[id][2])
        end
        instance.get_num_eq_lcons = function()
            return 0
        end
        instance.get_num_ineq_lcons = function()
            return 0
        end
        instance.set_status = function(s::Int)
            instance.status = s
        end
        instance.set_num_rows = function(id::Integer,v::Integer) end
        instance.set_num_cols = function(id::Integer,v::Integer) end
        instance.set_num_eq_cons = function(id::Integer,v::Integer) end
        instance.set_num_ineq_cons = function(id::Integer,v::Integer) end

        instance.str_init_x0 = function(id,x0)
            @assert(id in getLocalBlocksIds(instance.internalModel))
            mm = getModel(instance.internalModel,id)
            nvar = getNumVars(instance.internalModel,id)
            @assert length(x0) == nvar

            for i=1:nvar
                x0[i] = getvalue(Variable(mm,i))
                isnan(x0[i]) ? x0[i]=1.0 : nothing
            end
            # @show x0;
        end

        instance.str_prob_info = function(id,flag,mode,clb,cub,rlb,rub)
            # @show id
            if mode == :Structure
                if flag == 0
                    nn = getNumVars(instance.internalModel,id)
                    mm = getNumCons(instance.internalModel,id)
                else
                    @assert flag == 1
                    nn = getNumVars(instance.internalModel,id)
                    # mm = getNumLinkCons() #TODO: this is not yet supported
                    mm = 0
                end
                # @show nn,mm
                return (nn,mm)
            elseif mode == :Values
                if flag == 0
                    # @show length(clb),length(cub)
                    mm = getModel(instance.internalModel,id)
                    nvar = getNumVars(instance.internalModel,id)
                    @assert length(clb) == nvar
                    @assert length(cub) == nvar
                    array_copy(mm.colUpper, 1, cub, 1, nvar)
                    array_copy(mm.colLower, 1, clb, 1, nvar)
                    lb,ub = JuMP.constraintbounds(mm)
                    # @show lb, ub
                    (eq_idx, ieq_idx) = instance.iMap[id]
                    # @show eq_idx,ieq_idx
                    @assert length(lb) == length(ub)
                    @assert length(eq_idx) + length(ieq_idx) == length(lb)
                    num_eq = length(eq_idx)
                    for i in eq_idx
                        rlb[i[2]] = lb[i[1]]
                        rub[i[2]] = ub[i[1]]
                    end
                    for i in ieq_idx
                        rlb[i[2]+num_eq] = lb[i[1]]
                        rub[i[2]+num_eq] = ub[i[1]]
                    end
                    # @show id, rlb, rub
                else
                    @assert flag == 1
                    @assert length(rlb) == length(rub) == 0 #TODO: not yet support linking constraints syntax
                end
            else
                @assert false mode
            end
        end

        instance.str_eval_f = function(id,x0,x1)
            # @printf("#*************** eval_f - %d \n", id)
            # @show x0, x1
            # x0, x1 = load_x("pips", instance.n_iter)

            e =  instance.evaluatorMap[id]

            @timing instance.prof tic()
            obj = MathProgBase.eval_f(e,build_x(instance.internalModel,id,x0,x1))
            @timing instance.prof instance.t_jump += toq()
            # @show obj
            return obj

        end

        instance.str_eval_g = function(id::Int,x0::Vector{Float64},x1::Vector{Float64}, new_eq_g::Vector{Float64}, new_inq_g::Vector{Float64})
            # x0, x1 = load_x("pips", instance.n_iter)
            @timing instance.prof tic()

            e = instance.evaluatorMap[id]
            g = Vector{Float64}(undef, getNumCons(instance.internalModel,id))

            @timing instance.prof tic()
            MathProgBase.eval_g(e,g,build_x(instance.internalModel,id,x0,x1))
            @timing instance.prof instance.t_jump += toq()

            p = instance.iMap[id]
            eq_idx = p[1]
            ieq_idx = p[2]
            @assert length(new_eq_g) == length(eq_idx)
            @assert length(ieq_idx) == length(new_inq_g)

            @timing instance.prof tic()
            for i in eq_idx
                new_eq_g[i[2]] = g[i[1]]
                # @assert !haskey(ieq_idx,i[1])
            end
            for i in ieq_idx
                # @assert !haskey(eq_idx,i[1])
                new_inq_g[i[2]] = g[i[1]]
            end
            @timing instance.prof instance.t_itermap_eval_g+=toq()

            # @printf("#********  str_eval_g - %d \n", id)
            # @show x0, x1
            # @show new_eq_g, new_inq_g
            @timing instance.prof begin
                instance.t_eval_g += toq()
                instance.n_eval_g += 1
            end
        end

        instance.str_eval_grad_f = function(rowid, colid, x0, x1, new_grad_f)
            # x0, x1 = load_x("pips", instance.n_iter)
            m = instance.internalModel
            @assert rowid >= colid
            @assert sum(new_grad_f) == 0.0
            e = instance.evaluatorMap[rowid]
            x = build_x(m,rowid,x0,x1)
            g = Vector{Float64}(undef, length(x))

            @timing instance.prof tic()

            MathProgBase.eval_grad_f(e,g,x)
            @timing instance.prof instance.t_jump += toq()

            @assert length(g) == MathProgBase.numvar(getModel(m,rowid))
            @assert length(new_grad_f) == getNumVars(m,colid)

            var_idx_map = get_jac_col_idx_map(m,rowid,colid,instance.jMap)
            for i in var_idx_map
                new_grad_f[i[2]] = g[i[1]]
            end

            # @printf("#**********  str_eval_grad_f, %d %d \n", rowid, colid)
            # @show x0, x1, new_grad_f
            # @printf("grad%d%d=new_grad_f; \n",rowid,colid)
        end

        instance.str_eval_jac_g = function(rowid, colid, flag, x0 , x1, mode, e_rowidx, e_colptr, e_values, i_rowidx, i_colptr, i_values)
            # @printf("#**********  str_eval_jac_g  - %s - %d - %d %d \n", mode, flag, rowid, colid);
            # @show x0, x1
            # x0, x1 = load_x("pips",instance.n_iter)

            m = instance.internalModel
            # @show instance
            # @show instance.jMap
            jMap = instance.jMap
            matJac = instance.matJac
            matEqJac = instance.matEqJac
            matIeqJac = instance.matIeqJac
            @assert rowid<=num_scenarios(m) && colid <= num_scenarios(m)
            @assert rowid >= colid
            if(mode == :Structure)
                if !haskey(matJac, rowid)
                    e = instance.evaluatorMap[rowid]

                    @timing instance.prof tic()
                    jac_I,jac_J = MathProgBase.jac_structure(e)
                    @timing instance.prof instance.t_jump += toq()

                    mm = getModel(m,rowid)
                    matJac[rowid] = MatStorage(jac_I,jac_J, MathProgBase.numconstr(mm),MathProgBase.numvar(mm))
                    @assert length(jac_I) == length(matJac[rowid].value)
                end
                mat = matJac[rowid]
                jac_I = mat.rowIdx
                jac_J = mat.colIdx

                p = Pair{Int,Int}(rowid,colid)
                if !haskey(matEqJac,p)
                    @assert !haskey(matIeqJac,p)
                    (eq_idx, ieq_idx) = instance.iMap[rowid]
                    var_idx = get_jac_col_idx_map(m,rowid,colid,jMap)
                    # @show var_idx

                    eq_jac_I = Vector{Int}()
                    ieq_jac_I = Vector{Int}()
                    eq_jac_J = Vector{Int}()
                    ieq_jac_J = Vector{Int}()
                    for i = 1:length(jac_I)
                        ii = jac_I[i]
                        jj = jac_J[i]
                        if haskey(var_idx,jj)
                            if haskey(eq_idx,ii)
                                push!(eq_jac_I, eq_idx[ii])
                                push!(eq_jac_J, var_idx[jj])
                            else
                                @assert haskey(ieq_idx,ii)
                                push!(ieq_jac_I, ieq_idx[ii])
                                push!(ieq_jac_J, var_idx[jj])
                            end
                        end
                    end
                    matEq = MatStorage(eq_jac_I,eq_jac_J,length(eq_idx), getNumVars(m,colid))
                    matIeq = MatStorage(ieq_jac_I,ieq_jac_J,length(ieq_idx), getNumVars(m,colid))
                    matEqJac[p] = matEq
                    matIeqJac[p] = matIeq
                end
                matEq = matEqJac[p]
                matIeq = matIeqJac[p]
                fill!(matEq.value,1.0)
                fill!(matIeq.value,1.0)
                @assert matEq.m + matIeq.m == mat.m

                # @show eq_jac_I
                # @show eq_jac_J

                @timing instance.prof tic()
                eq_jac = sparse(matEq.rowIdx,matEq.colIdx,matEq.value,matEq.m, matEq.n)
                @timing instance.prof begin
                    instance.t_jac_spconv += toq()
                    instance.n_jac_spconv += 1
                end
                # @show ieq_jac_I
                # @show ieq_jac_J
                @timing instance.prof tic()
                ieq_jac = sparse(matIeq.rowIdx,matIeq.colIdx,matIeq.value,matIeq.m, matIeq.n)
                @timing instance.prof begin
                    instance.t_jac_spconv += toq()
                    instance.n_jac_spconv += 1
                end

                nnzEqJac = length(eq_jac.nzval)
                nnzIeqJac = length(ieq_jac.nzval)

                return nnzEqJac,nnzIeqJac
            elseif(mode == :Values)
                p = Pair{Int,Int}(rowid,colid)
                @assert haskey(matEqJac,p)
                @assert haskey(matIeqJac,p)
                @assert haskey(matJac,rowid)
                mat = matJac[rowid]
                jac_I = mat.rowIdx
                jac_J = mat.colIdx
                jac_g = mat.value
                if(!mat.isVal)
                    e = instance.evaluatorMap[rowid]

                    @timing instance.prof tic()
                    MathProgBase.eval_jac_g(e,jac_g,build_x(m,rowid,x0,x1))
                    @timing instance.prof instance.t_jump += toq()

                    mat.isVal = true
                end

                (eq_idx, ieq_idx) = instance.iMap[rowid]
                var_idx = get_jac_col_idx_map(m,rowid,colid,jMap)

                matEq =  matEqJac[p]
                matIeq = matIeqJac[p]
                eq_jac_g = matEq.value
                eqi = 1
                ieq_jac_g = matIeq.value
                ieqi = 1

                for i = 1:length(jac_I)
                    ii = jac_I[i]
                    jj = jac_J[i]
                    vv = jac_g[i]
                    # @show ii, jj, vv
                    if haskey(var_idx,jj)
                        if haskey(eq_idx,ii)
                            eq_jac_g[eqi] = vv
                            eqi += 1
                        else
                            @assert haskey(ieq_idx,ii)
                            ieq_jac_g[ieqi] = vv
                            ieqi += 1
                        end
                    end
                end

                if(length(eq_jac_g) != 0)

                    @timing instance.prof tic()
                    eq_jac = sparse(matEq.rowIdx, matEq.colIdx, matEq.value, matEq.m, matEq.n, keepzeros=true)
                    @timing instance.prof begin
                        instance.t_jac_spconv += toq()
                        instance.n_jac_spconv += 1
                    end
                    # @printf("em=%d; en=%d;\n", length(eq_idx), getNumVars(m,colid))
                    # @show eq_jac_I, eq_jac_J, eq_jac_g
                    # @printf("ejac%d%d=sparse(eq_jac_I,eq_jac_J,eq_jac_g,em,en); \n",rowid,colid)

                    array_copy(eq_jac.rowval,1,e_rowidx,1,length(eq_jac.rowval))
                    array_copy(eq_jac.colptr,1,e_colptr,1,length(eq_jac.colptr))
                    array_copy(eq_jac.nzval, 1,e_values,1,length(eq_jac.nzval))
                    # convert_to_c_idx(e_rowidx)
                    # convert_to_c_idx(e_colptr)

                    filename = string("jaceq_",rowid,"_",colid)
                    write_mat_to_file(filename,eq_jac)
                    # @show eq_jac
                    convert_to_c_idx(e_rowidx)
                    convert_to_c_idx(e_colptr)
                end

                if(length(ieq_jac_g) != 0)
                    @timing instance.prof tic()
                    ieq_jac = sparse(matIeq.rowIdx, matIeq.colIdx, matIeq.value, matIeq.m, matIeq.n, keepzeros=true)
                    @timing instance.prof begin
                        instance.t_jac_spconv += toq()
                        instance.n_jac_spconv += 1
                    end
                    # @printf("im=%d; in=%d;\n", length(ieq_idx), getNumVars(m,colid))
                    # @show ieq_jac_I, ieq_jac_J, ieq_jac_g
                    # @printf("ijac%d%d=sparse(ieq_jac_I,ieq_jac_J,ieq_jac_g,im,in); \n",rowid,colid)
                    # @printf("jac%d%d=vcat(ejac%d%d,ijac%d%d) \n",rowid,colid,rowid,colid,rowid,colid)

                    array_copy(ieq_jac.rowval,1,i_rowidx,1,length(ieq_jac.rowval))
                    array_copy(ieq_jac.colptr,1,i_colptr,1,length(ieq_jac.colptr))
                    array_copy(ieq_jac.nzval, 1,i_values,1,length(ieq_jac.nzval))
                    # convert_to_c_idx(i_rowidx)
                    # convert_to_c_idx(i_colptr)

                    filename = string("jacieq_",rowid,"_",colid)
                    write_mat_to_file(filename,ieq_jac)
                    # @show ieq_jac
                    convert_to_c_idx(i_rowidx)
                    convert_to_c_idx(i_colptr)
                end

                if(flag == 2)
                    mat.isVal = false
                end
            else
                @assert false mode
            end

            # @show "end str_eval_jac_g"
        end

        instance.str_eval_h = function(rowid, colid, flag, x0, x1, obj_factor, lambda, mode, rowidx, colptr, values)
            # @printf("#**********  str_eval_h - %s - %d  - %d %d  %f \n", mode, flag, rowid, colid, obj_factor)
            # @show x0, x1
            # @show lambda
            # x0, x1 = load_x("pips", instance.n_iter)

            # @show x0, x1
            # @show lambda
            m = instance.internalModel
            matHess = instance.matHess
            matHessMap = instance.matHessMap
            hcMap = instance.hcMap
            hrMap = instance.hrMap
            @assert rowid<=num_scenarios(m) && colid <=num_scenarios(m)
            if(mode == :Structure)
                if !haskey(matHess,colid)
                    e = instance.evaluatorMap[colid]
                    mm = getModel(m,rowid)

                    @timing instance.prof tic()
                    (h_J,h_I) = MathProgBase.hesslag_structure(e) # upper trangular
                    @timing instance.prof instance.t_jump += toq()

                    matHess[colid] = MatStorage(h_I,h_J,MathProgBase.numvar(mm),MathProgBase.numvar(mm))
                end
                mat = matHess[colid]
                h_I = mat.rowIdx
                h_J = mat.colIdx

                p = Pair{Int,Int}(rowid,colid)
                if !haskey(matHessMap,p)
                    if rowid == colid  #diagonal
                        col_var_idx, row_var_idx = get_h_col_idx_map(m,rowid,rowid,hcMap,hrMap)
                        new_h_I = Vector{Int}()
                        new_h_J = Vector{Int}()
                        for i = 1:length(h_I)
                            ii = h_I[i]
                            jj = h_J[i]
                            if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                                new_ii = row_var_idx[ii]
                                new_jj = col_var_idx[jj]
                                if new_ii < new_jj
                                    push!(new_h_I,new_ii)
                                    push!(new_h_J,new_jj)
                                else
                                    push!(new_h_I,new_jj)
                                    push!(new_h_J,new_ii)
                                end
                            end
                        end
                        matP = MatStorage(new_h_I, new_h_J, getNumVars(m,rowid), getNumVars(m,rowid))
                        matHessMap[p] = matP
                        fill!(matP.value,1.0)
                    elseif rowid == 0  && colid != 0 #border corss hessian
                        col_var_idx, row_var_idx = get_h_col_idx_map(m,rowid,colid,hcMap,hrMap)
                        new_h_I = Vector{Int}()
                        new_h_J = Vector{Int}()
                        for i = 1:length(h_I)
                            ii = h_I[i]
                            jj = h_J[i]
                            if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                                new_ii = row_var_idx[ii]
                                new_jj = col_var_idx[jj]
                                push!(new_h_I,new_ii)
                                push!(new_h_J,new_jj)
                            end
                        end
                        matP = MatStorage(new_h_I, new_h_J, getNumVars(m,colid), getNumVars(m,rowid))
                        matHessMap[p] = matP
                        fill!(matP.value,1.0)
                        # @show h_I, h_J
                        # @show col_var_idx
                        # @show row_var_idx
                        # @show new_h_I, new_h_J
                        # @show matP

                        #structure for root contribution
                        col_var_idx,row_var_idx = get_h_col_idx_map(m,colid,rowid,hcMap,hrMap)
                        # @show col_var_idx
                        # @show row_var_idx
                        new_h_I = Vector{Int}()
                        new_h_J = Vector{Int}()
                        for i = 1:length(h_I)
                            ii = h_I[i]
                            jj = h_J[i]
                            if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                                new_ii = row_var_idx[ii]
                                new_jj = col_var_idx[jj]
                                if new_ii < new_jj
                                    push!(new_h_I,new_ii)
                                    push!(new_h_J,new_jj)
                                else
                                    push!(new_h_I,new_jj)
                                    push!(new_h_J,new_ii)
                                end
                            end
                        end
                        # @show h0_I,h0_J
                        p0 = Pair{Int,Int}(colid,rowid)
                        matP0 = MatStorage(new_h_I, new_h_J, getNumVars(m,0),getNumVars(m,0))
                        matHessMap[p0] = matP0
                    else
                        @assert (rowid !=0 && colid == 0)
                        @assert false
                    end
                end
                matP = matHessMap[p]

                @timing instance.prof tic()
                laghess = sparse(matP.rowIdx,matP.colIdx, matP.value, matP.m, matP.n)
                @timing instance.prof begin
                    instance.t_hess_spconv += toq()
                    instance.n_hess_spconv += 1
                end
                return length(laghess.nzval)
            elseif(mode == :Values)
                p = Pair{Int,Int}(rowid,colid)
                if rowid == colid
                    @assert haskey(matHess,rowid)
                    mat = matHess[rowid]
                    if !mat.isVal
                        e = instance.evaluatorMap[rowid]
                        (eq_idx, ieq_idx) = instance.iMap[rowid]
                        numeq = length(eq_idx)
                        lam_new = Vector{Float64}(undef, length(lambda))
                        for i in eq_idx
                            lam_new[i[1]] = lambda[i[2]]
                        end
                        for i in ieq_idx
                            lam_new[i[1]] = lambda[i[2]+numeq]
                        end
                        x = build_x(m,rowid,x0,x1)
                        # @show x,lam_new
                        @timing instance.prof tic()
                        MathProgBase.eval_hesslag(e,mat.value,x,obj_factor,lam_new)
                        @timing instance.prof instance.t_jump += toq()

                        mat.isVal = true
                    end
                    h_I = mat.rowIdx
                    h_J = mat.colIdx
                    h = mat.value
                    # @show x0,x1
                    # @show x
                    # @show lambda
                    # @show h_I
                    # @show h_J
                    # @show h
                    col_var_idx,row_var_idx = get_h_col_idx_map(m,rowid, colid,hcMap,hrMap)
                    # @show col_var_idx
                    # @show row_var_idx
                    matP = matHessMap[p]
                    matPi = 1
                    for i = 1:length(h_I)
                        ii = h_I[i]
                        jj = h_J[i]
                        if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                            matP.value[matPi] =  h[i]
                            matPi += 1
                        end
                    end
                    # @show new_h_I
                    # @show new_h_J
                    # @show new_h
                    @timing instance.prof tic()
                    str_laghess = sparse(matP.rowIdx,matP.colIdx,matP.value,matP.m,matP.n,keepzeros = true)
                    @timing instance.prof begin
                        instance.t_hess_spconv += toq()
                        instance.n_hess_spconv += 1
                    end
                    # @printf("m=%d;n=%d; \n",getNumVars(m,rowid),getNumVars(m,rowid))
                    # @show new_h_I, new_h_J, new_h
                    # @printf(" hess%d%d=sparse(new_h_I,new_h_J,new_h, m, n); \n",rowid,colid)
                    # @show str_laghess
                    array_copy(str_laghess.rowval,1,rowidx,1,length(str_laghess.rowval))
                    array_copy(str_laghess.colptr,1,colptr,1,length(str_laghess.colptr))
                    array_copy(str_laghess.nzval, 1,values,1,length(str_laghess.nzval))
                elseif rowid == 0  && colid != 0 #border corss hessian
                    @assert haskey(matHess,colid)
                    mat = matHess[colid]
                    if !mat.isVal
                        e = instance.evaluatorMap[colid]
                        x = build_x(m,colid,x0,x1)
                        (eq_idx, ieq_idx) = instance.iMap[colid]
                        numeq = length(eq_idx)
                        lam_new = Vector{Float64}(undef, length(lambda))
                        for i in eq_idx
                            lam_new[i[1]] = lambda[i[2]]
                        end
                        for i in ieq_idx
                            lam_new[i[1]] = lambda[i[2]+numeq]
                        end

                        @timing instance.prof tic()
                        MathProgBase.eval_hesslag(e,mat.value,x,obj_factor,lam_new)
                        @timing instance.prof instance.t_jump += toq()

                        mat.isVal = true
                    end
                    h_I = mat.rowIdx
                    h_J = mat.colIdx
                    h = mat.value
                    # @show h_I, h_J
                    # @show h

                    col_var_idx,row_var_idx = get_h_col_idx_map(m,0, colid,hcMap,hrMap)
                    matP = matHessMap[p]
                    matPi = 1
                    for i = 1:length(h_I)
                        ii = h_I[i]
                        jj = h_J[i]
                        if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                            matP.value[matPi] = h[i]
                            matPi += 1
                        end
                    end
                    @timing instance.prof tic()
                    str_laghess = sparse(matP.rowIdx,matP.colIdx,matP.value,matP.m,matP.n,keepzeros = true)
                    @timing instance.prof begin
                        instance.t_hess_spconv += toq()
                        instance.n_hess_spconv += 1
                    end
                    # @printf("m=%d;n=%d; \n",getNumVars(m,rowid),getNumVars(m,rowid))
                    # @show new_h_I, new_h_J, new_h
                    # @printf(" hess%d%d=sparse(new_h_I,new_h_J,new_h, m, n); \n",rowid,colid)
                    # @show str_laghess
                    array_copy(str_laghess.rowval,1,rowidx,1,length(str_laghess.rowval))
                    array_copy(str_laghess.colptr,1,colptr,1,length(str_laghess.colptr))
                    array_copy(str_laghess.nzval, 1,values,1,length(str_laghess.nzval))
                elseif rowid !=0 && colid == 0 #root contribution
                    @assert flag == 2
                    @assert haskey(matHess,rowid)
                    mat = matHess[rowid]
                    if !mat.isVal
                        e = instance.evaluatorMap[rowid]
                        x = build_x(m,rowid,x0,x1)
                        (eq_idx, ieq_idx) = instance.iMap[rowid]
                        numeq = length(eq_idx)
                        lam_new = Vector{Float64}(undef, length(lambda))
                        for i in eq_idx
                            lam_new[i[1]] = lambda[i[2]]
                        end
                        for i in ieq_idx
                            lam_new[i[1]] = lambda[i[2]+numeq]
                        end
                        @timing instance.prof tic()
                        MathProgBase.eval_hesslag(e,mat.value,x,obj_factor,lam_new)
                        @timing instance.prof instance.t_jump += toq()

                        mat.isVal = true
                    end
                    h_I = mat.rowIdx
                    h_J = mat.colIdx
                    h = mat.value
                    # @show h_I, h_J
                    # @show h
                    col_var_idx,row_var_idx = get_h_col_idx_map(m, rowid, 0,hcMap,hrMap)
                    matP = matHessMap[p]
                    matPi = 1
                    for i = 1:length(h_I)
                        ii = h_I[i]
                        jj = h_J[i]
                        if haskey(col_var_idx,jj) && haskey(row_var_idx,ii)
                            matP.value[matPi] = h[i]
                            matPi += 1
                        end
                    end
                    matP0 = matHess[0]
                    h0_I = matP0.rowIdx
                    h0_J = matP0.colIdx
                    @assert matP0.m == matP.m == matP0.n == matP.n == getNumVars(m,0)

                    # @show h0_I,h0_J
                    @timing instance.prof tic()
                    str_laghess = sparse([matP.rowIdx;h0_I], [matP.colIdx;h0_J], [matP.value;zeros(Float64,length(h0_I))], matP.m, matP.n, keepzeros=true)
                    @timing instance.prof begin
                        instance.t_hess00_spconv += toq()
                        instance.n_hess00_spconv += 1
                    end
                    # new_h_I = [new_h_I;h0_I]
                    # new_h_J = [new_h_J;h0_J]
                    # new_h = [new_h;zeros(Float64,length(h0_I))]
                    # @printf("m=%d; n=%d; \n",getNumVars(m,0),getNumVars(m,0))
                    # @show new_h_I, new_h_J, new_h
                    # @printf(" hess%d%d=sparse(new_h_I,new_h_J,new_h, m, n); \n",rowid,colid)
                    array_copy(str_laghess.rowval,1,rowidx,1,length(str_laghess.rowval))
                    array_copy(str_laghess.colptr,1,colptr,1,length(str_laghess.colptr))
                    array_copy(str_laghess.nzval, 1,values,1,length(str_laghess.nzval))
                else
                    @assert false
                end
                # filename = string("hess_",rowid,"_",colid)
                # write_mat_to_file(filename,str_laghess)

                convert_to_c_idx(rowidx)
                convert_to_c_idx(colptr)
                if flag == 2
                    mat.isVal = false
                end
                # @show values
            else
                @assert false mode
            end
            # @show "end str_eval_h"
        end

        instance.str_write_solution = function(id, x, y_eq, y_ieq)
            # @show id, x, y_eq, y_ieq
            @assert id in getLocalBlocksIds(instance.internalModel)
            @assert length(x) == instance.get_num_cols(id)
            @assert length(y_eq) == instance.get_num_eq_cons(id)
            @assert length(y_ieq) == instance.get_num_ineq_cons(id)

            #write back the primal to JuMP
            mm = getModel(instance.internalModel,id)
            for i = 1:length(x)
                setvalue(Variable(mm,i), x[i])
            end
        end

        return instance
    end
end

# function show_time(m::ModelInterface)
#     t_model_evaluation = m.t_sj_init +
#     m.t_sj_init_x0 +
#     m.t_prob_info +
#     m.t_eval_f +
#     m.t_eval_g+
#     m.t_eval_grad_f+
#     m.t_eval_jac_g+
#     m.t_eval_h +
#     m.t_write_solution;
#     return t_model_time
# end

#######
# Linking with PIPS Julia Structure interface
######

function structJuMPSolve(model; with_prof=false, suppress_warmings=false,kwargs...)
    # @show "solve"
    t_sj_lifetime = 0.0
    @timing with_prof tic()

    # MPI.Init() ＃initialize in model loading

    comm = getStructure(model).mpiWrapper.comm
    # @show "[$(MPI.Comm_rank(comm))/$(MPI.Comm_size(comm))] create problem "

    t_sj_model_init = 0.0
    @timing with_prof tic()

    prob = PipsNlpSolver.createProblemStruct(comm, StructJuMPModel(model,with_prof), with_prof)

    @timing with_prof t_sj_model_init += toq()

    # @show "end createStructJuMPPipsNlpProblem"

    t_sj_solver_total = 0.0
    @timing with_prof tic()

    status = PipsNlpSolver.solveProblemStruct(prob)

    @timing with_prof t_sj_solver_total += toq()

    @timing with_prof t_sj_lifetime += toq()

    # solver_time = solver_total - modeling_time
    # if(0==MPI.Comm_rank(MPI.COMM_WORLD))
    #   @printf "Total time %.4f (initialization=%.3f modelling=%.3f solver=%.3f) (in sec)\n" t_total prob.model.t_sj_init modeling_time solver_time
    # end
    @timing with_prof begin
        mid, nprocs = getMyRank()
        # bname = string(split(ARGS[1],"/")[2],"_",num_scenarios(model))
        # run(`mkdir -p ./out/$bname`)
        # fname = string("./out/",bname,"/",bname,"_",nprocs,".",mid,".jl.txt")
        # # @show bname, fname
        # tfile = open(fname, "w")
        s1 = @sprintf("[%d/%d] [ t_sj_model_init %f t_sj_solver_total %f  t_sj_lifetime %f ] \n", mid, nprocs, t_sj_model_init, t_sj_solver_total, t_sj_lifetime)
        s2 = @sprintf("[%d/%d] [ t_jl_str_total %f t_jl_eval_total %f  ] \n", mid, nprocs, prob.t_jl_str_total, prob.t_jl_eval_total)
        s3 = @sprintf("[%d/%d] [ t_jac_spconv %f n_jac_spconv %d  ] \n", mid, nprocs, prob.model.t_jac_spconv, prob.model.n_jac_spconv)
        s4 = @sprintf("[%d/%d] [ t_hess_spconv %f n_hess_spconv %d  ] \n", mid, nprocs, prob.model.t_hess_spconv, prob.model.n_hess_spconv)
        s5 = @sprintf("[%d/%d] [ t_hess00_spconv %f n_hess00_spconv %d  ] \n", mid, nprocs, prob.model.t_hess00_spconv,prob.model.n_hess00_spconv)
        s6 = @sprintf("[%d/%d] [ t_jump %f  ] \n", mid, nprocs, prob.model.t_jump)
        s7 = @sprintf("[%d/%d] [ t_eval_g %f t_itermap_eval_g %f n_eval_g %d ] \n", mid, nprocs, prob.model.t_eval_g, prob.model.t_itermap_eval_g, prob.model.n_eval_g)
        if(mid == 0)
            @printf("%s",s1)
            @printf("%s",s2)
            @printf("%s",s3)
            @printf("%s",s4)
            @printf("%s",s5)
            @printf("%s",s6)
            @printf("%s",s7)
        end
        # @printf(tfile, "%s", s1)
        # @printf(tfile, "%s", s2)
        # @printf("%s",s3)
        # @printf("%s",s4)
        # @printf("%s",s5)
        # @printf("%s",s6)
        # @printf("%s",s7)
        # close(tfile)
        # n1 = string("./out/",nprocs,".",mid,".c.txt")
        # n2 = string("./out/",bname,"/",bname,"_",nprocs,".",mid,".c.txt")
#        run(`mv $n1 $n2`)
    end

    return PIPSRetCodeToSolverInterfaceCode[status]
end

# function load_x(subdir,iter)
#     @printf("load x from ./%s/x%d \n",subdir,iter)
#     x = readdlm(string("./",subdir,"/x",iter))
#     x0 = x[1:78]
#     x1 = x[79:144]
#     return x0, x1
# end

# function write_x0(subdir,iter,x)
#     @printf("writing x to ./%s/x0_%d \n",subdir,iter)
#     run(`mkdir -p ./$subdir`)
#     writedlm(string("./",subdir,"/x0_",iter),x,",")
# end
# function write_x1(subdir,iter,x)
#     @printf("writing x to ./%s/x1_%d \n",subdir,iter)
#     run(`mkdir -p ./$subdir`)
#     writedlm(string("./",subdir,"/x1_",iter),x,",")
# end

function initialization(instance::StructJuMPModel)
    iMap = instance.iMap
    model = instance.internalModel
    evaluatorMap = instance.evaluatorMap
    @assert length(iMap) == length(evaluatorMap) == 0
    for id in getLocalBlocksIds(model)
        mm = getModel(model,id)
        e = JuMP.NLPEvaluator(mm)
        MathProgBase.initialize(e,[:Grad,:Jac,:Hess])
        push!(evaluatorMap,id=>e)
        eq_idx = Dict{Int,Int}()
        ieq_idx = Dict{Int,Int}()
        push!(iMap,id=>Pair(eq_idx,ieq_idx))
        lb,ub=JuMP.constraintbounds(mm)

        for i =1:length(lb)
            if lb[i] == ub[i]
                eq_idx[i] = length(eq_idx) + 1
            else
                ieq_idx[i] =length(ieq_idx) + 1 #remember to offset length(eq_idx)
            end
        end
    end
end

function get_jac_col_idx_map(m,rowid,colid,map)
    p = Pair{Int,Int}(rowid,colid)
    if haskey(map,p)
        idx_map = map[p]
    else
        idx_map = get_jac_col_var_idx(m,rowid, colid)
        push!(map,p=>idx_map)
    end
    return idx_map
end

function get_jac_col_var_idx(m,rowid, colid)  #this method customerized for no linking constraint presented
    # @show "get_jac_col_var_idx",rowid,colid
    idx_map = Dict{Int,Int}() #dummy (jump) -> actual used
    if rowid == colid
        nvar = getNumVars(m,rowid)
        for i = 1:nvar
            idx_map[i] = i
        end
    else
        @assert rowid!=0 && rowid != colid
        mm = getModel(m,rowid)
        othermap = getStructure(mm).othermap
        for p in othermap
            pidx = p[1].col
            cidx = p[2].col
            idx_map[cidx] = pidx
        end
    end
    # @show idx_map
    return idx_map
end

function get_h_col_idx_map(m,rowid,colid,cmap,rmap)
    p = Pair{Int,Int}(rowid,colid)
    if haskey(cmap,p)
        cidx_map = cmap[p]
        ridx_map = rmap[p]
    else
        cidx_map,ridx_map = get_h_var_idx(m,rowid,colid)
        push!(cmap,p=>cidx_map)
        push!(rmap,p=>ridx_map)
    end
    return cidx_map, ridx_map
end

function get_h_var_idx(m,rowid, colid)
    # @show "get_h_var_idx",rowid,colid
    col_idx_map = Dict{Int,Int}() #dummy (jump) -> actual used
    row_idx_map = Dict{Int,Int}()
    if rowid == colid
        nvar = getNumVars(m,rowid) #need to place model variable in front of non model variable.
        for i = 1:nvar
            col_idx_map[i] = i
            row_idx_map[i] = i
        end
    elseif rowid == 0  && colid != 0 #border
        mm = getModel(m,colid)
        othermap = getStructure(mm).othermap
        for p in othermap
            pidx = p[1].col
            cidx = p[2].col
            col_idx_map[cidx] = pidx
        end
        for i = 1:getNumVars(m,colid)
            row_idx_map[i] = i
        end
    elseif colid == 0 && rowid != 0 #root contrib.
        mm = getModel(m,rowid)
        othermap = getStructure(mm).othermap
        for p in othermap
            pidx = p[1].col
            cidx = p[2].col
            col_idx_map[cidx] = pidx
            row_idx_map[cidx] = pidx
        end
    else
        @assert false rowid colid
    end
    # @show col_idx_map,row_idx_map
    return col_idx_map,row_idx_map
end

KnownSolvers["PipsNlp"] = PipsNlpInterface.structJuMPSolve

end


struct CollectCovers{MC,SC}
    mosaics::MC
    sets::SC

    cover_params::CoverParams
    enum_params::CoverEnumerationParams
    opt_params::AbstractOptimizerParams

    function CollectCovers(
        mosaics::MC, sets::SC,
        cover_params::CoverParams=CoverParams(),
        enum_params::CoverEnumerationParams=CoverEnumerationParams(),
        opt_params::AbstractOptimizerParams=QuadraticOptimizerParams()
    ) where {MC,SC}
        new{MC,SC}(mosaics, sets, cover_params, enum_params, opt_params)
    end
end

function (worker::CollectCovers)(mosaic_key, set_key; verbose::Bool=false)
    mosaic_masked = mask(worker.mosaics[mosaic_key], [worker.sets[set_key]],
                         max_overlap_logpvalue=worker.enum_params.max_set_score+worker.cover_params.sel_tax)
    covers_coll = collect(mosaic_masked, worker.cover_params, worker.enum_params, worker.opt_params,
                          verbose)
    ((mosaic_key, set_key), !isempty(covers_coll) ? covers_coll : nothing)
end

function pcollect(
    mosaics::Dict{MK}, sets::Dict{SK};
    mode=:parallel,
    cover_params::CoverParams=CoverParams(),
    enum_params::CoverEnumerationParams=CoverEnumerationParams(),
    opt_params::AbstractOptimizerParams=QuadraticOptimizerParams(),
    pids=workers(),
    verbose::Bool=false,
    kwargs...
) where {MK, SK}
    @info("Parallel OESC ($(length(mosaics)) mosaic(s) × $(length(sets)) set(s))...")
    # FIXME workaround for pmap() being unable to handle callable objects +
    # unable to efficiently handle large callable objects
    collect_covers = CollectCovers(mosaics, sets, cover_params, enum_params, opt_params)
    #raw_res = sizehint!(Vector{Tuple{MK,SK,CoverCollection}}(),
    #                    length(tasks))
    mosaicXset_keys = vec([(m, s) for m in keys(mosaics), s in keys(sets)])
    if mode == :parallel
        keyXcover_vec = @distributed (append!) for (m,s) in shuffle!(mosaicXset_keys)
            res = Any[collect_covers(m, s)]
            @info("$(myid()): done mosaic=$m × set=$s")
            return res
        end
    elseif mode == :sequential
        keyXcover_vec = map(mosaicXset_keys) do mXs
            m, s = mXs
            verbose && @info("Calculating mosaic=$m × set=$s...")
            collect_covers(m, s, verbose=verbose)
        end
    else
        throw(ArgumentError("Unknown mode $mode"))
    end
    verbose && @info("Done covers enumeration, preparing the results...")
    # filter out empty covers
    keyXcover = sizehint!(Dict{Tuple{MK,SK}, CoverCollection}(), length(keyXcover_vec))
    for (mXs_key, covers_coll) in keyXcover_vec
        if covers_coll !== nothing
            keyXcover[mXs_key] = covers_coll
        end
    end
    @info("$(length(keyXcover)) non-trivial cover(s) collected")
    return keyXcover
end

function pcollect(sets1_colls, sets2;
    problem_type::Symbol=:quadratic, mode::Symbol=:sequential,
    cover_params::CoverParams=CoverParams(),
    enum_params::CoverEnumerationParams=CoverEnumerationParams(),
    verbose::Bool=false, pids=workers(),
    kwargs...
)
    verbose && @info("Preparing mosaics...")
    #p = Progress(sum(length, values(entrez_clusters))*length(entrez_colls), 5.0,
    #              "Computing cluster coverings...")   # minimum update interval: 5 second
    sets1_mosaics = Dict(pmap(sets1_colls) do coll_kv
        verbose && @info("Preparing $(coll_kv[1]) mosaic...")
        Pair(coll_kv[1], SetMosaic(coll_kv[2]))
    end)
    return pcollect(sets1_mosaics, sets2,
                    problem_type=problem_type, mode=mode, pids=pids,
                    cover_params=cover_params, enum_params=enum_params,
                    opt_params=OptimizerParams(problem_type, Workers=pids, kwargs),
                    verbose=verbose)
end


"""
```
Li2020{T} <: AbstractNLCTFPModel{T}
```

The `Li2020` type defines the structure of the model in Li, Wenhao (2020), "Public Liquidity and Financial Crises."

### Fields

#### Parameters
* `parameters::Vector{AbstractParameter}`: Vector of all time-invariant model
  parameters.

* `keys::OrderedDict{Symbol,Int}`: Maps human-readable names for all model
  parameters to their indices in `parameters`.

#### Inputs to Measurement and Equilibrium Functional Equations

The following fields are dictionaries that map human-readable names to indices.

* `state_variables::OrderedDict{Symbol,Int}`: Maps each state variable to its dimension on a Cartesian grid

* `endogenous_variables::OrderedDict{Symbol,Int}`: Maps endogenous variables
    calculated during the solution of functional equations to an index

* `exogenous_shocks::OrderedDict{Symbol,Int}`: Maps each shock to a column in
  the measurement equation

* `observables::OrderedDict{Symbol,Int}`: Maps each observable to a row in the
  model's measurement equation matrices.

* `pseudo_observables::OrderedDict{Symbol,Int}`: Maps each pseudo-observable to
  a row in the model's pseudo-measurement equation matrices.

#### Model Specifications and Settings

* `spec::String`: The model specification identifier, \"an_schorfheide\", cached
  here for filepath computation.

* `subspec::String`: The model subspecification number, indicating that some
  parameters from the original model spec (\"ss0\") are initialized
  differently. Cached here for filepath computation.

* `settings::Dict{Symbol,Setting}`: Settings/flags that affect computation
  without changing the economic or mathematical setup of the model.

* `test_settings::Dict{Symbol,Setting}`: Settings/flags for testing mode

#### Other Fields

* `rng::MersenneTwister`: Random number generator. Can be is seeded to ensure
  reproducibility in algorithms that involve randomness (such as
  Metropolis-Hastings).

* `testing::Bool`: Indicates whether the model is in testing mode. If `true`,
  settings from `m.test_settings` are used in place of those in `m.settings`.

* `observable_mappings::OrderedDict{Symbol,Observable}`: A dictionary that
  stores data sources, series mnemonics, and transformations to/from model units.
  DSGE.jl will fetch data from the Federal Reserve Bank of
  St. Louis's FRED database; all other data must be downloaded by the
  user. See `load_data` and `Observable` for further details.

* `pseudo_observable_mappings::OrderedDict{Symbol,PseudoObservable}`: A
  dictionary that stores names and transformations to/from model units. See
  `PseudoObservable` for further details.
"""
mutable struct Li2020{T} <: AbstractNLCTFPModel{T}
    parameters::ParameterVector{T}                         # vector of all time-invariant model parameters
    keys::OrderedDict{Symbol,Int}                          # human-readable names for all the model
                                                           # parameters and steady-states
    state_variables::OrderedDict{Symbol,Int}               # dimension number of state variable

    functional_variables::OrderedDict{Symbol,Int}
    derivatives::OrderedDict{Symbol,Int}
    endogenous_variables::OrderedDict{Symbol,Int}
    exogenous_shocks::OrderedDict{Symbol,Int}
    observables::OrderedDict{Symbol,Int}
    pseudo_observables::OrderedDict{Symbol,Int}

    spec::String                                           # Model specification number (eg "m990")
    subspec::String                                        # Model subspecification (eg "ss0")
    settings::Dict{Symbol,Setting}                         # Settings/flags for computation
    test_settings::Dict{Symbol,Setting}                    # Settings/flags for testing mode
    rng::MersenneTwister                                   # Random number generator
    testing::Bool                                          # Whether we are in testing mode or not

    observable_mappings::OrderedDict{Symbol, Observable}
    pseudo_observable_mappings::OrderedDict{Symbol, PseudoObservable}
end

description(m::Li2020) = "Julia implementation of model defined in 'Public Liquidity and Financial Crises' (2020) by Wenhao Li: Li2020, $(m.subspec)"

"""
`init_model_indices!(m::Li2020)`

Arguments:
`m:: Li2020`: a model object

Description:
Initializes indices for all of `m`'s states, shocks, and observables.
"""
function init_model_indices!(m::Li2020)
    # State_Variables
    state_variables = collect([:w])

    # Exogenous shocks
    exogenous_shocks = collect([:K_sh, :N_sh]) # capital shock K, liquidity shock N

    # Variables for functional equations
    functional_variables = collect([:p, :Q̂, :xg])

    # Derivatives of variables
    init_derivatives!(m, Dict{Symbol, Vector{Int}}(:p => standard_derivs(1)), state_variables)
    derivatives = keys(get_derivatives(m))

    # Endogenous variables
    endogenous_variables = collect([:ψ, :xK, :yK, :yg, :σp, :σ, :σh, :σw, :μR_rd, :rd_rg, :rd_rg_H, :μb_μh, :μw, :μp, :μK, :μR, :rd, :rg, :rd_rf,
                                    :μb, :μh, :invst, :lvg, :κp, :κb, :κd, :κh, :κfs, :firesale_jump, :κw, :liq_prem, :bank_liq_frac,
                                    :δ_x, :indic, :rf, :rh, :K_growth, :κK])

    # Observables
    observables = keys(m.observable_mappings)

    # Pseudo-observables
    pseudo_observables = keys(m.pseudo_observable_mappings)

    for (i,k) in enumerate(state_variables);        m.state_variables[k]      = i end
    for (i,k) in enumerate(exogenous_shocks);       m.exogenous_shocks[k]     = i end
    for (i,k) in enumerate(endogenous_variables);   m.endogenous_variables[k] = i end
    for (i,k) in enumerate(functional_variables);   m.functional_variables[k] = i end
    for (i,k) in enumerate(derivatives);            m.derivatives[k]          = i end
    for (i,k) in enumerate(observables);            m.observables[k]          = i end
    for (i,k) in enumerate(pseudo_observables);     m.pseudo_observables[k]   = i end
end

function Li2020(subspec::String = "ss0";
                custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                testing = false)

    # Model-specific specifications
    spec               = split(basename(@__FILE__),'.')[1]
    subspec            = subspec
    settings           = Dict{Symbol,Setting}()
    test_settings      = Dict{Symbol,Setting}()
    rng                = MersenneTwister(0)

    # initialize empty model
    m = Li2020{Float64}(
            # model parameters and steady state values
            Vector{AbstractParameter{Float64}}(), OrderedDict{Symbol,Int}(),

            # model indices
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),

            spec,
            subspec,
            settings,
            test_settings,
            rng,
            testing,
            OrderedDict{Symbol,Observable}(),
            OrderedDict{Symbol,PseudoObservable}())

    # Set settings
    model_settings!(m)

    # default_test_settings!(m)
    for custom_setting in values(custom_settings)
        m <= custom_setting
    end

    # Set observable and pseudo-observable transformations
    # init_observable_mappings!(m)
    # init_pseudo_observable_mappings!(m)

    # Initialize parameters
    init_parameters!(m)

    init_model_indices!(m)
    init_subspec!(m)

    return m
end

"""
```
init_parameters!(m::Li2020)
```
Initializes the model's parameters.
"""
function init_parameters!(m::Li2020)
    # Productivity
    m <= parameter(:AH, 0.15, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(),
                   Uniform(0., 1e3), fixed = false, description = "Banker's productivity")
    m <= parameter(:AL, 0.125, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Household's productivity")

    # Bank run and fire sales
    m <= parameter(:λ, 0.025, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.),
                   fixed = false, description = "Probability of liquidity shock")
    m <= parameter(:β, 0.25, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Probability of running on a deposit account")
    m <= parameter(:α, 0.21, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Illiquidity discount on capital during fire sale")
    m <= parameter(:ϵ, 1e-3, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Fraction of remaining wealth for bankers after bankruptcy")
    m <= parameter(:θ, 1e-5, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Probability of non-zero exposure to liquidity shock")

    # Shadow value of liquidity
    m <= parameter(:π, 0.19, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Illiquidity discount on illiquid safe asset")

    # Macro parameters
    m <= parameter(:δ, 0.1, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.),
                   fixed = false, description = "Depreciation rate of capital")
    m <= parameter(:ρ, 0.04, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3),
                   fixed = false, description = "Discount rate")
    m <= parameter(:σK, 0.033, (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3), fixed = false,
                   description = "Volatility of capital shocks")
    m <= parameter(:χ, 3., (0., Inf), (0., 1e3), ModelConstructors.Untransformed(), Uniform(0., 1e3), fixed = false,
                   description = "Adjustment cost of internal investment function")

    # Stationary distribution
    m <= parameter(:η, 0.05, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Uniform(0., 1.), fixed = false,
                   description = "Rate of retirement for bankers")
end

"""
```
model_settings!(m::Li2020)
```
creates the model's settings.
"""
function model_settings!(m::Li2020)

    # Investment functions
    m <= Setting(:Φ, quadratic_investment, "Internal investment function")
    m <= Setting(:μK, growth_quadratic_investment, "Rate of capital growth")
    m <= Setting(:∂Φ, derivative_quadratic_investment_li2020, "Derivative of internal investment function")

    # Numerical settings for grid
    m <= Setting(:N, 100, "Grid size")
    m <= Setting(:stategrid_method, :exponential, "Type of grid to construct for state variables")
    m <= Setting(:stategrid_dimensions, OrderedDict{Symbol, Tuple{Float64, Float64, Int}}(:w => (1e-3, 1., get_setting(m, :N))),
                 "Information about the dimensions of the state space")
    m <= Setting(:stategrid_splice, 0.2, "Li (2020) constructs the grid in two parts. This value is where the first half stops.")

    # Numerical settings for no jump equilibrium
    m <= Setting(:boundary_conditions, OrderedDict{Symbol, Vector{Float64}}(:p => [0.; 0.]),
                 "Boundary conditions for differential equations.")
    m <= Setting(:ode_reltol, 1e-4, "Relative tolerance for ODE integration")
    m <= Setting(:ode_abstol, 1e-12, "Absolute tolerance for ODE integration")
    m <= Setting(:essentially_one, 0.999,
                "If ψ is larger than this value, then we consider it essentially one for some numerical purposes.")
    m <= Setting(:ode_integrator, DP5(), "Numerical ODE integrator for no-jump solution") # DP5 is Matlab's ode45

    # Numerical settings for functional iteration
    m <= Setting(:tol, 1e-4, "Tolerance for functional iteration")
    m <= Setting(:learning_rate, 0.4, "Learning rate for the update of functional variables after each loop")
    m <= Setting(:max_iter, 40, "Maximum number of loops during functional iteration")
    m <= Setting(:error_method, :total_error, "Method for calculating error at the end of each loop during functional iteration")
    m <= Setting(:v₀, 3e-8, "Parameter for damping function")
    m <= Setting(:damping_function, x -> get_setting(m, :v₀) ./ x, "Dampling function to avoid corners")
    m <= Setting(:p₀_perturb, 1e-14, "Perturbation of boundary condition for q at w = 0")
    m <= Setting(:κp_grid, Vector{Float64}(undef, 0),
                 "Vector of guesses for κp used during iteration for new κp and xK within each functional loop")
    m <= Setting(:p_interpolant, Gridded(Linear()), "Interpolant method for fitted p.")
    m <= Setting(:xK_interpolant, Gridded(Linear()), "Interpolant method for xK after solving xK for each guess of κp.")
    m <= Setting(:xg_interpolant, Gridded(Linear()), "Interpolant method for xg after solving xK and κp.")
    m <= Setting(:κp_interpolant, Gridded(Linear()), "Interpolant method for κp after solving xK and κp.")
    m <= Setting(:Q̂_interpolant, Gridded(Linear()), "Interpolant method for Q̂.")
    m <= Setting(:inside_iteration_nlsolve_tol, 1e-6, "Tolerance for nlsolve in the inside iteration")
    m <= Setting(:inside_iteration_nlsolve_max_iter, 400, "Maximum number of iterations for nlsolve for the inside iteration")
    m <= Setting(:xg_tol, 2e-3, "Lower but permissible tolerance for xg fixed point in the inside iteration")
    m <= Setting(:yg_tol, 1e-5, "Lowest permissible value for yg")
    m <= Setting(:p_tol, 1e-8, "Tolerance when calculating p after completing an inside iteration")
    m <= Setting(:firesale_bound, 0.99, "Upper bound for size of firesale jumps that do not wipe out a banker's net worth")
    m <= Setting(:firesale_interpolant, Gridded(Linear()), "Interpolant method for the firesale value.")
    m <= Setting(:N_GH, 10, "Number of nodes for Gauss-Hermite quadrature.")
    m <= Setting(:Q̂_tol, 1e-5, "Tolerance for approximation of Q̂ via Gauss-Hermite quadrature")
    m <= Setting(:Q̂_max_it, 1000, "Maximum number of iterations for approximation of Q̂.")
    m <= Setting(:dt, 1. / 12., "Simulation interval for Q̂ approximation as a fraction of a 1 year")

    # Other settings for initialization
    m <= Setting(:nojump_parameters, [:ρ, :AH, :AL, :σK, :χ, :δ, :η], "Keys of parameters used when solving the no-jump equilibrium.")

    # Calibration targets
    m <= Setting(:avg_gdp, 0.145, "Average GDP")
    m <= Setting(:liq_gdp_ratio, 0.39, "Average Liquidity/GDP ratio in the data")
    m <= Setting(:gov_bond_gdp_level, 0.3, "Ratio of Government Bonds to GDP")
end


module FrostHelperEasedCameraZoomTrigger

using ..Ahorn, Maple

@mapdef Trigger "FrostHelper/EasedCameraZoomTrigger" EasedCameraZoomTrigger(x::Integer, y::Integer, width::Integer=16, height::Integer=16, easing::String="Linear", targetZoom::Number=2.0, revertOnLeave::Bool=false, easeDuration::Number=1.0, focusOnPlayer::Bool=false, revertMode::String="RevertToNoZoom", disableInPhotosensitiveMode::Bool=false)

const placements = Ahorn.PlacementDict(
    "Eased Camera Zoom (Frost Helper)" => Ahorn.EntityPlacement(
        EasedCameraZoomTrigger,
        "rectangle"
    )
)

Ahorn.nodeLimits(entity::EasedCameraZoomTrigger) = 0, 1

Ahorn.editingOptions(entity::EasedCameraZoomTrigger) = Dict{String, Any}(
    "easing" => easings,
    "revertMode" => revertModes
)

const easings = String[
    "BackIn",
    "BackInOut",
    "BackOut",
    "BounceIn",
    "BounceInOut",
    "BounceOut",
    "CubeIn",
    "CubeInOut",
    "CubeOut",
    "ElasticIn",
    "ElasticInOut",
    "ElasticOut",
    "ExpoIn",
    "ExpoInOut",
    "ExpoOut",
    "Linear",
    "QuadIn",
    "QuadInOut",
    "QuadOut",
    "SineIn",
    "SineInOut",
    "SineOut"
]

const revertModes = String[
    "RevertToNoZoom",
    "RevertToPreviousZoom",
]

end


promote_rule(::Type{S}, ::Type{T}) where S<:SafeSigned   where T<:SafeSigned =
    safeint(promote_type(baseint(S), baseint(T)))
promote_rule(::Type{S}, ::Type{T}) where S<:SafeSigned   where T<:SafeUnsigned =
    safeint(promote_type(baseint(S), baseint(T)))
promote_rule(::Type{S}, ::Type{T}) where S<:SafeUnsigned where T<:SafeUnsigned =
    safeint(promote_type(baseint(S), baseint(T)))

promote_rule(::Type{S}, ::Type{T}) where S<:SafeSigned   where T<:Signed   = S
promote_rule(::Type{S}, ::Type{T}) where S<:SafeSigned   where T<:Unsigned = S
promote_rule(::Type{S}, ::Type{T}) where S<:SafeUnsigned where T<:Signed   = S
promote_rule(::Type{S}, ::Type{T}) where S<:SafeUnsigned where T<:Unsigned = S

promote_rule(::Type{S}, ::Type{T}) where S<:SafeSigned   where T<:Base.IEEEFloat = T
promote_rule(::Type{S}, ::Type{T}) where S<:SafeUnsigned where T<:Base.IEEEFloat = T

promote_rule(::Type{S}, ::Type{BigFloat}) where S<:SafeSigned  = BigFloat
promote_rule(::Type{S}, ::Type{BigFloat}) where S<:SafeUnsigned = BigFloat


"""
    ComtradeData

`ComtradeData` holds data and configuration information from a set of COMTRADE
files. Fields include:

* `filename::UTF8String` -- Base file name (without extension).
* `cfg::ComtradeCfg` -- Configuration information.
* `dat::DataFrame` -- Table of results with each channel as a column.
* `hdr::UTF8String` -- Contents of the "hdr" file.
* `inf::UTF8String` -- Contents of the "inf" file.
"""
type ComtradeData
    filename::UTF8String
    cfg::ComtradeCfg
    dat::DataFrame
    hdr::UTF8String
    inf::UTF8String
end

function addcol!(d::DataFrame, x)
    push!(d.columns, x)
    push!(DataFrames.index(d), symbol(string("x", ncol(d) + 1)))
    d
end

"""
    read_comtrade(basename)

Read contents of a set of COMTRADE files with base file name `basename` (no
extension).

Returns a `COMTRADE.ComtradeData` with the following fields:

* `filename::UTF8String` -- Base file name (without extension).
* `cfg::ComtradeCfg` -- Configuration information.
* `dat::DataFrame` -- Table of results with each channel as a column.
* `hdr::UTF8String` -- Contents of the "hdr" file.
* `inf::UTF8String` -- Contents of the "inf" file.

See also `COMTRADE.ComtradeCfg` for a description of the `cfg` field.
"""
function read_comtrade(basename)
    cfg = read_cfg("$basename.cfg")
    if cfg.ft == "ASCII"
        d = readtable("$basename.dat", separator = ',', header = false)
        for i in 1:cfg.nA+2   # convert DataArray columns to Arrays
            d.columns[i] = d[i].data
        end
        for i in (1:cfg.nD) + 2 + cfg.nA
            d.columns[i] = map(Bool, d[i].data)
        end
    elseif cfg.ft == "BINARY"
        x = reinterpret(Int16, readbytes(open("$basename.dat")))
        n = cfg.nA + div(cfg.nD, 16) + 2 + 2
        x = reshape(x[1:(n * cfg.npts)], (n, cfg.npts))
        global d = DataFrame()
        addcol!(d, reinterpret(UInt32, x[1:2, :], (cfg.npts,)))
        addcol!(d, reinterpret(UInt32, x[3:4, :], (cfg.npts,)))
        for i = 1:cfg.nA
            addcol!(d, reshape(x[i+4, :], (cfg.npts,)))
        end
        for i = 1:cfg.nD
            addcol!(d, Bool[(x >>> (mod1(i, 16) - 1)) & 0x01 for x in x[cfg.nA + 4 + cld(i, 16), :]])
            # addcol!(d, Bool[(x >>> (16 - mod1(i, 16))) & 0x01 for x in x[cfg.nA + 4 + cld(i, 16), :]])
        end
    else
        error("Unsupported data type $(x.ft)")
    end
    newnames = DataFrames.make_unique(vcat(:n, :time, Symbol[symbol(x) for x in vcat(cfg.A[:ch_id], cfg.D[:ch_id])]))
    names!(d, newnames)
    # analog multipliers
    for i in 1:cfg.nA
        d.columns[i+2] = cfg.A[i,:a] * d[i+2] + cfg.A[i,:b]
    end
    # analog multipliers
    d.columns[2] = 1e-6 * cfg.timemult * d[2]
    hdr = isreadable("$basename.hdr") ? readall("$basename.hdr") : ""
    inf = isreadable("$basename.inf") ? readall("$basename.inf") : ""
    return ComtradeData(basename, cfg, d, hdr, inf)
end

"""
    S"some name"

Returns a symbol. Equivalent to `symbol("some name")`. This helps when COMTRADE
data columns use nonstandard names.
"""
macro S_str(arg)
    Meta.quot(symbol(arg))
end


using .Gadfly

using NeXLMatrixCorrection
using Colors

"""
    Gadfly.plot(
        dbkrs::AbstractArray{DBKRatio};
        mc::Type{<:MatrixCorrection} = XPP,
        fc::Type{<:FluorescenceCorrection} = ReedFluorescence,
        palette = NeXLPalette,
        style::Symbol = :Ratio | :XY
    )
"""
function Gadfly.plot(
    dbkrs::AbstractArray{DBKRatio};
    mc::Type{<:MatrixCorrection} = XPP,
    fc::Type{<:FluorescenceCorrection} = ReedFluorescence,
    cc::Type{<:CoatingCorrection} = Coating,
    style::Symbol = :Ratio
)
    if style==:Ratio
        plot_ratio(dbkrs, mc, fc, cc)
    else
        plot_xy(dbkrs, mc, fc)
    end
end

function plot_ratio(
    dbkrs::AbstractArray{DBKRatio},
    mc::Type{<:MatrixCorrection} = XPP,
    fc::Type{<:FluorescenceCorrection} = ReedFluorescence,
    cc::Type{<:CoatingCorrection} = Coating,
)
  df = asa(DataFrame, dbkrs, withComputedKs=true, mc=mc, fc=fc, cc=cc)
  df[:, "Element"] = symbol.(element.(first.(df[:,"Lines"])))
  df[:, "Z"] = z.(element.(first.(df[:,"Lines"])))
  df[:, "Family"] = name.(df[:,"Lines"], true)
  df[:, "RU"] = map(r->uv(r["K"],r["ΔK"])/r["Kxpp"], eachrow(df))
  df[:, "Rmin"] = map(r-> value(r["RU"])-σ(r["RU"]), eachrow(df))
  df[:, "Rmax"] = map(r-> value(r["RU"])+σ(r["RU"]), eachrow(df))
  df[:,"Measurement"] = map(r->"$(name(r["Cmeas"])) at $(r["E0meas"]/1000.0) keV", eachrow(df))
  df[:,"Rand"] = rand(nrow(df)) # To plot points in a randomized order.
  filter!(r->element(first(r[:Lines])) in keys(r[:Cmeas]), df)
  sort!(df, [:Family, :E0meas, :Z, :Rand])
  # yin, yax = something(ymin, 0.9*minimum(skipmissing(df[:,"Rmin"]))), something(ymax, 1.1*maximum(skipmissing(df[:,"Rmax"])))
  plot(df, xgroup="Family", x="Element", ymin="Rmin", ymax="Rmax", y="Ratio", color="Measurement",
     # Coord.cartesian(ymin=yin, ymax=yax), Stat.x_jitter(range=1.0), # Neither seems to work...
     Geom.subplot_grid(Geom.errorbar), Stat.x_jitter(range=1000.0))
end

function plot_xy(
    dbkrs::AbstractArray{DBKRatio},
    mc::Type{<:MatrixCorrection} = XPP,
    fc::Type{<:FluorescenceCorrection} = ReedFluorescence,
    cc::Type{<:CoatingCorrection} = Coating,
)
    x, y, dy, color = Float64[], Float64[], Float64[], String[]
    next = 0
    for dbkr in dbkrs
        kr, unkComp = asa(KRatio, dbkr), get(dbkr.measured, :Composition, missing)
        if hasminrequired(mc, kr.unkProps) &&
           hasminrequired(fc, kr.unkProps) && #
           hasminrequired(mc, kr.stdProps) &&
           hasminrequired(fc, kr.stdProps) && #
           (!isnothing(unkComp)) &&
           (value(unkComp[kr.element]) > 0.0) &&
           (value(kr.standard[kr.element]) > 0.0)
            try
                # Compute the k-ratio
                kc = gZAFc(kr, unkComp, mc=mc, fc=fc, cc=cc) * (value(unkComp[kr.element]) / value(kr.standard[kr.element]))
                push!(y, value(kr.kratio))
                push!(dy, σ(kr.kratio))
                push!(x, kc)
                push!(color, "$(symbol(kr.element)) $(name(shell(brightest(kr.lines)))) in $(name(unkComp)) $(kr.unkProps[:BeamEnergy]/1000.0) keV")
            catch c
                @info "Failed on $dbkr - $c"
            end
        end
    end
    abline = Geom.abline(color="red", style=:dash)
    plot(
        x = x,
        y = y,
        ymin = y .- dy,
        ymax = y .+ dy,
        color = color,
        Geom.errorbar,
        Guide.xlabel("k[$(repr(nameof(mc))[2:end])]"),
        Guide.ylabel("k[Measured]"),
        Guide.colorkey(title="Measurement"),
        abline, intercept=[0.0],slope=[1.0],
    )
end

function plot2(dbkrs::AbstractArray{DBKRatio}; palette = NeXLPalette)
    kok, dkok, color = Float64[], Float64[], Color[]
    ygroup = CharXRay[]
    matcolors = Dict{String,RGB{Float64}}()
    for dbkr in dbkrs
        kr, unkComp = asa(KRatio, dbkr), get(dbkr.measured, :Composition, missing)
        push!(kok, value(kr.kratio))
        push!(dkok, σ(kr.kratio))
        push!(ygroup, brightest(kr.lines))
        matname = "$(kr.element) in $(name(unkComp)) $(kr.unkProps[:BeamEnergy]/1000.0) keV"
        if !haskey(matcolors, matname)
            matcolors[matname] = palette[length(matcolors)+1]
        end
        push!(color, matcolors[matname])
    end
    plot(
        x = eachindex(kok),
        y = kok,
        ymin = kok .- dkok,
        ymax = kok .+ dkok,
        color = color,
        ygroup = ygroup, # Geom.errorbar, Stat.x_jitter(range=0.8),
        Geom.subplot_grid(Geom.errorbar, free_y_axis = true),
        Scale.ygroup(labels = cxr -> repr(cxr), levels = allcxrs),
        Guide.manual_color_key("Material", collect(keys(matcolors)), collect(values(matcolors))), # Guide.yrug,
        Guide.xlabel("Index"),
        Guide.ylabel("k[Measured]"),
    )
end

function plot3(
    krs::AbstractArray{DBKRatio};
    known::Union{Material,Missing} = missing,
    label::AbstractString = "Material",
    palette = NeXLPalette,
)
    mats = unique(collect(map(dbkr->get(dbkr.measured,:Composition,missing), krs)))
    allelms = sort(convert(Vector{Element}, collect(union(map(keys, mats)...))))
    elmcol = Dict(elm => palette[i] for (i, elm) in enumerate(allelms))
    xs, ymin, ymax, ygroups, colors = String[], Float64[], Float64[], Element[], Color[]
    for mat in mats
        append!(xs, [name(mat) for elm in keys(mat)])
        append!(ymin, [value(mat[elm]) - σ(mat[elm]) for elm in keys(mat)])
        append!(ymax, [value(mat[elm]) + σ(mat[elm]) for elm in keys(mat)])
        append!(colors, [elmcol[elm] for elm in keys(mat)])
        append!(ygroups, collect(keys(mat)))
    end
    plot(
        x = xs,
        ymin = ymin,
        ymax = ymax,
        color = colors,
        ygroup = ygroups,
        Geom.subplot_grid(Geom.errorbar, free_y_axis = true),
        Scale.ygroup(labels = elm -> symbol(elm), levels = allelms),
        Guide.manual_color_key("Material", symbol.(collect(keys(elmcol))), collect(values(elmcol))),
        Guide.xlabel(label),
        Guide.ylabel("Mass Fraction by Element"),
    )
end