module NAG
export
    nag_licence_query, nag_license_query,
    nag_complex_polygamma,
    nag_opt_read!,
    nag_opt_nlp!,
    nag_1d_quad_inf_1,
    nag_zero_cont_func_brent,
    nag_opt_lp!,
    last_nag_error

nag_licence_query() = ccall((:a00acc, :libnagc_nag), Cint, ()) == 1
const nag_license_query = nag_licence_query # US vs UK spelling
nag_license_query() || warn("Cannot acquire a NAG license.")

const NagInt = Int32
const NagComplex = Complex128

type NagError <: Exception
    code::Int
    name::ASCIIString
    message::UTF8String
end

Base.showerror(io::IO, e::NagError) =
    print(io, "NAG function \"$(e.name)\" [$(e.code)] – $(e.message)")

const NAG_ERROR = zeros(Uint8,544)
const nag_errors = Array(NagError)
nag_errors[] = NagError(0, "", "NO_ERROR")

cstr_to_array(p::Ptr{Uint8}, own::Bool = false) =
    pointer_to_array(p, int(ccall(:strlen, Csize_t, (Ptr{Uint8},), p)), own)

function error_handler(msg::Ptr{Uint8}, code::Cint, name::Ptr{Uint8})
    code == 0 && return
    msg = UTF8String(copy(cstr_to_array(msg)))
    name = ASCIIString(copy(cstr_to_array(name)))
    nag_errors[] = NagError(int(code), name, msg)
    throw(nag_errors[])
end
const ptr_error_handler = cfunction(error_handler, Void, (Ptr{Uint8}, Cint, Ptr{Uint8}))

function reset_nag_error()
    fill!(NAG_ERROR, 0)
    unsafe_store!(
        convert(Ptr{Ptr{Void}}, pointer(NAG_ERROR)) + 2*sizeof(Cint) + 512,
        ptr_error_handler
    )
end
reset_nag_error()

last_nag_error() = nag_errors[]

const objfunref = Array(Function)
const confunref = Array(Function)

function objfun_wrapper(
    n::NagInt, x::Ptr{Float64}, objf::Ptr{Float64}, g::Ptr{Float64}, comm::Ptr{NagInt}
)
    x = pointer_to_array(x,int(n),false)
    g = pointer_to_array(g,int(n),false)
    f = unsafe_load(comm)
    unsafe_store!(objf,objfunref[1](x,g,f))
    return nothing
end
const c_objfun = cfunction(objfun_wrapper, Void,
    (NagInt, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{NagInt}))

function confun_wrapper(
    n::NagInt, ncnlin::NagInt, needc::Ptr{NagInt}, x::Ptr{Float64},
    conf::Ptr{Float64}, conjac::Ptr{Float64}, comm::Ptr{NagInt}
)
    needc  = pointer_to_array(needc,int(ncnlin),false)
    x      = pointer_to_array(x,int(n),false)
    conf   = pointer_to_array(conf,int(ncnlin),false)
    conjac = pointer_to_array(conjac,(int(n),int(ncnlin)),false)
    flag   = unsafe_load(comm)
    confunref[1](needc,x,conf,conjac,flag)
    return nothing
end
const c_confun = cfunction(confun_wrapper, Void,
    (NagInt, NagInt, Ptr{NagInt}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{NagInt}))

function nag_complex_polygamma(z::Number, k::Integer)
    ccall((:s14afc, :libnagc_nag), NagComplex,
          (NagComplex, NagInt, Ptr{Void}),
          z, k, NAG_ERROR)
end

# if Base.polygamma doesn't support complex args, use NAG
if !applicable(Base.polygamma, 1, 1+2im)
    Base.polygamma(k::Int, z::Complex) = nag_complex_polygamma(z, k)
    Base.digamma(z::Complex) = polygamma(0, z)
    Base.trigamma(z::Complex) = polygamma(1, z)
end

function nag_opt_read!(name::ByteString, optfile::ByteString, print::Bool = false)
    options = zeros(Uint8,1440)
    ccall((:e04xxc, :libnagc_nag), Void, (Ptr{Void},), options)
    ccall((:e04xyc, :libnagc_nag), Void,
          (Ptr{Uint8}, Ptr{Uint8}, Ptr{Void}, Cint, Ptr{Uint8}, Ptr{Void}),
          name, optfile, options, print, "stdout", NAG_ERROR)
    return options
end

function nag_opt_lp!(
    A         :: Matrix{Float64},
    bl        :: Vector{Float64},
    bu        :: Vector{Float64},
    c         :: Vector{Float64},
    x         :: Vector{Float64};
    optfile   :: ByteString = "",
    transpose :: Bool = true,
)
    # since NAG is row-major
    #   when transpose == true the rows are linear constraints
    #   when transpose == false the columns are linear constraints
    transpose && (A = A')

    n = length(c)
    tda, nclin = size(A)

    nclin > 0 && tda >= n   ||
        error("bad linear constraint dimensions (tda ≱ n): $tda ≱ $n")
    length(bl) == n + nclin ||
        error("wrong number of lower bounds: $(length(bl)) ≠ $(n + nclin)")
    length(bu) == n + nclin ||
        error("wrong number of upper bounds: $(length(bu)) ≠ $(n + nclin)")
    length(c) == n          ||
        error("wrong number of coefficients: $(length(c)) ≠ $n")
    length(x) == n          ||
        error("wrong data vector size: $(length(x)) ≠ $n")

    objf = zeros()
    options = isempty(optfile) ? C_NULL :
        convert(Ptr{Void}, nag_opt_read!("e04mfc", optfile))

    reset_nag_error()
    ccall((:e04mfc,:libnagc_nag), Void,
           (NagInt, NagInt, Ptr{Float64}, NagInt, Ptr{Float64}, Ptr{Float64},
            Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Void}, Ptr{Void}, Ptr{Void}),
           n, nclin, A, tda, bl, bu, c, x, objf, options, C_NULL, NAG_ERROR)

    return x, objf[]
end

function nag_opt_nlp!(
    A  :: Matrix{Float64},
    bl :: Vector{Float64},
    bu :: Vector{Float64},
    x  :: Vector{Float64},
    objfun! :: Function,
    confun! :: Function,
    optfile :: ByteString = "",
    transpose :: Bool = true,
)
    # since NAG is row-major
    #   when transpose == true the rows are linear constraints
    #   when transpose == false the columns are linear constraints
    transpose && (A = A')

    # extract various dimensions
    n = length(x)
    tda, nclin = size(A)
    ncnlin = length(bl) - n - nclin

    # check for usage problems
    length(bl) == length(bu) ||
        error("bounds vectors must have matching length")
    lexcmp(bl,bu) <= 0       ||
        error("lower bounds cannot be greater than upper bounds")
    0 <= ncnlin              ||
        error("as many bounds as variables and linear constraints must be given")
    nclin == 0 || n <= tda   ||
        error("second dimension of linear coefficients too small")

    # allocate output variables
    objf = zeros()
    g = zeros(n)

    # save callbacks in globals
    objfunref[1] = objfun!
    confunref[1] = confun!

    options = isempty(optfile) ? C_NULL :
        convert(Ptr{Void}, nag_opt_read!("e04ucc", optfile))

    reset_nag_error()
    ccall((:e04ucc, :libnagc_nag), Void,
          (NagInt, NagInt, NagInt, Ptr{Float64},
           NagInt, Ptr{Float64}, Ptr{Float64},
           Ptr{Void}, Ptr{Void},
           Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Void}, Ptr{Void}, Ptr{Void}),
          n, nclin, ncnlin, A,
          tda, bl, bu,
          c_objfun, c_confun,
          x, objf, g,
          options, C_NULL, NAG_ERROR)

    return x, objf[], g
end

const quadfunref = Array(Function)
quad_fun_wrapper(x::Float64, comm::Ptr{NagInt}) = quadfunref[1](x)::Float64
const c_quadfun = cfunction(quad_fun_wrapper, Float64, (Float64, Ptr{NagInt}))

function nag_1d_quad_inf_1(
    f :: Function,
    boundinf :: Symbol = :Infinite,
    bound :: Float64 = 0.0;
    epsabs :: Float64 = 0.0,
    epsrel :: Float64 = sqrt(eps(Float64)),
    max_num_subint :: NagInt = int32(1e7),
)
    max_num_subint > 0 || error("max num subint must be > 0")
    boundinf_val =
        boundinf == :UpperSemiInfinite ? int32(1076 + 0) :
        boundinf == :LowerSemiInfinite ? int32(1076 + 1) :
        boundinf == :Infinite          ? int32(1076 + 2) :
        error("""
        invalid boundinf symbol: $boundinf
        must be one of: UpperSemiInfinite, LowerSemiInfinite, Infinite
        """)

    # allocate output variables
    result = zeros()
    abserr = zeros()
    qp = zeros(Uint8, 48)
    comm = zeros(Uint8, 8)

    quadfunref[] = f

    reset_nag_error()
    ccall((:d01smc, :libnagc_nag), Void,
        (Ptr{Void}, NagInt, Float64, Float64,
         Float64, NagInt, Ptr{Float64}, Ptr{Float64},
         Ptr{Void}, Ptr{Void}, Ptr{Void}),
        c_quadfun, boundinf_val, bound, epsabs,
        epsrel, max_num_subint, result, abserr,
        qp, comm, NAG_ERROR)

    return result[], abserr[]
end

const quadfunref_brent = Array(Function)
quad_fun_wrapper_brent(x::Float64, comm::Ptr{NagInt}) = quadfunref_brent[1](x)::Float64
const c_quadfun_brent = cfunction(quad_fun_wrapper_brent, Float64, (Float64, Ptr{NagInt}))

function nag_zero_cont_func_brent(
    a :: Float64,
    b :: Float64,
    f :: Function;
    eps :: Float64 = 1e-5,
    eta :: Float64 = 0.0,
)
    f(a)*f(b) <= 0 || error("f(a)*f(b) must be <= 0")
    quadfunref_brent[1] = f

    x = zeros()
    comm = zeros(Uint8, 8)

    reset_nag_error()
    ccall((:c05ayc, :libnagc_nag), Void,
        (Float64, Float64, Float64, Float64,
         Ptr{Void}, Ptr{Float64}, Ptr{Void}, Ptr{Void}),
        a, b, eps, eta, NAG.c_quadfun_brent, x, comm, NAG_ERROR)

    return x[]
end

end # module


function vtk_xml_write_header(vtk::DatasetFile)
    xroot = create_root(vtk.xdoc, "VTKFile")
    set_attribute(xroot, "type", vtk.grid_type)
    set_attribute(xroot, "version", "1.0")
    if IS_LITTLE_ENDIAN
        set_attribute(xroot, "byte_order", "LittleEndian")
    else
        set_attribute(xroot, "byte_order", "BigEndian")
    end
    set_attribute(xroot, "header_type", string(HeaderType))
    if vtk.compression_level > 0
        set_attribute(xroot, "compressor", "vtkZLibDataCompressor")
    end
    xroot::XMLElement
end

# Specifies the "extent" of a structured grid, e.g. (1:4, 1:3, 1:42)
const Extent{N} = Tuple{Vararg{AbstractUnitRange{<:Integer}, N}} where {N}

# Switch to zero-based extent used by VTK.
to_vtk_extent(r::AbstractUnitRange) = r .- 1

function to_vtk_extent(extent::Tuple)
    ext = to_extent3(extent)
    map(to_vtk_extent, ext)
end

# Convert different extent specifications to Extent{3}.
to_extent3(extent::Extent{3}) = extent

function to_extent3(extent::Extent{N}) where {N}
    if N > 3
        throw(ArgumentError("dimensionalities N > 3 not supported (got N = $N)"))
    elseif N < 3
        M = 3 - N
        (extent..., ntuple(d -> Base.OneTo(1), Val(M))...) :: Extent{3}
    end
end

to_extent3(Ns::Dims) = to_extent3(map(Base.OneTo, Ns))

# Returns the "extent" attribute required for structured (including rectilinear)
# grids.
function extent_attribute(extent_in)
    ext = to_vtk_extent(extent_in)
    iter = Iterators.map(e -> string(first(e), " ", last(e)), ext)
    join(iter, " ")
end

# Number of cells in structured grids.
# In 3D, all cells are hexahedrons (i.e. VTK_HEXAHEDRON), and the number of
# cells is (Ni-1)*(Nj-1)*(Nk-1). In 2D, they are quadrilaterals (VTK_QUAD), and in
# 1D they are line segments (VTK_LINE).
function num_cells_structured(Ns::Dims)
    prod(N -> max(1, N - 1), Ns)
end

# Accept passing coordinates as a tuple.
vtk_grid(filename, coords::NTuple; kwargs...) =
    vtk_grid(filename, coords...; kwargs...)


"""
Generate plots for the paper.

# Example
```julia
# Setup
include("src/make_plots.jl"); using .MakePlots
MakePlots.setup(num_workers=16)  # Setup 16 worker processes

# Run the simulations
job_id = "all_plots"
samples = 1000*2000  # How many samples, ~1.15 CPU-hours per 100000 samples
dists = [3:2:11...]  # Which code distances to simulate
plots = [1:12...]  # Which of the 12 plots to generate
MakePlots.dist_calc_all(job_id, dists, samples, 1:12)  # Start computing

# Plot the results as they are computed
# Run plot_finished() any time to check the progress of a plot
for plot_i in 1:12
    MakePlots.plot_finished(job_id, plot_i)
end
```
"""
module MakePlots

using Distributed
import PyPlot; plt = PyPlot

export setup, dist_calc_all, fetch_finished, plot, plot_finished

include("VLQ.jl")
include("jobs.jl"); using .Jobs

function setup(; num_workers=16)
    Jobs.launch_workers(num_workers)
    eval(:(@everywhere (
        include("src/VLQ.jl");
        using .VLQ.BiMaps
    )))

    @everywhere function everywhere_calc_single(
            plot_i, syndrome_sym, d, e, override_key, samples)
        step = 100
        e *= 9 <= plot_i <=11 ? 1.0e9 : 1.0
        if override_key === nothing
            override_pairs = ()
        else
            override_pairs = [(override_key, e)]
            e = VLQ.sensitivity_base_p
        end

        # Shortcut for slow high error runs
        (e > 0.01 && d > 5) && (samples = min(samples, 10000))
        e > 0.03 && (samples = min(samples, 1000))
        e >= 0.1 && (samples = min(samples, 500))

        model = VLQ.make_noise_model_for_paper(e, override_pairs)
        ctx = VLQ.CodeDistanceSim(d,
            getproperty(VLQ, syndrome_sym)(),
            model
        )
        run = VLQ.CodeDistanceRun(ctx)
        loops = div(samples, step)
        sum(let r = VLQ.do_n_runs(run, step, false)
                ccall(:jl_gc_collect, Nothing, ())
                r
            end
            for _ in 1:loops
        ) / loops
    end
end

function x_axis_for_plot(plot_i)
    if 1 <= plot_i <= 5
        10 .^ LinRange(log10(0.1), log10(0.0001), 19)[4:end-2]
    elseif 6 <= plot_i <= 8
        10 .^ LinRange(log10(0.1), log10(1e-5), 13+3)[3:end]
    elseif plot_i == 9
        10 .^ LinRange(log10(1e-5), log10(1), 16)[1:end-3]
    elseif plot_i == 10
        10 .^ LinRange(log10(1e-5), log10(1), 16)[1:end-3]
    elseif plot_i == 11
        10 .^ LinRange(log10(1e-4), log10(1e-8), 17)[1:end-4]
    elseif plot_i == 12
        LinRange(2, 30, 15)
    end
end
function confs_for_plot(plot_i, dists, samples)
    if 1 <= plot_i <= 5
        confs_for_thresh_plot(plot_i, dists, samples)
    else
        confs_for_sens_plot(plot_i, dists, samples)
    end
end
function confs_for_thresh_plot(plot_i, dists, samples)
    x_arr = x_axis_for_plot(plot_i)
    syndrom_sym = [
        :TypicalSyndrome, :NaturalAllAtOnce, :NaturalInterleaved,
        :CompactAllAtOnce, :CompactInterleaved][plot_i]
    [
        (plot_i, syndrom_sym, d, e, nothing, samples)
        for e in x_arr
        for d in dists
    ]
end
function confs_for_sens_plot(plot_i, dists, samples)
    x_arr = x_axis_for_plot(plot_i)
    override_key = [
        :p_tt, :p_loadstore, :p_tc, :t1_c, :t1_t,
        :dur_loadstore, :cavity_depth][plot_i-5]
    [
        (plot_i, :CompactInterleaved, d, e, override_key, samples)
        for e in x_arr
        for d in dists
    ]
end

function dist_calc_all(job_id, dists, samples, which_plots=1:12)
    confs = []
    for i in which_plots
        append!(confs, confs_for_plot(i, dists, samples))
    end
    Jobs.run_on_workers(job_id, Main.everywhere_calc_single, confs)
end

function fetch_finished(job_id, plot_i, default::Float64=1.0)
    x_arr = x_axis_for_plot(plot_i)
    y_arr_dict = Dict{Int, Vector{Float64}}()
    results = Jobs.current_results(job_id)
    for ((res_plot_i, _, d, e, _, _), val) in results
        res_plot_i == plot_i || continue
        if !(d in keys(y_arr_dict))
            y_arr_dict[d] = repeat([default], length(x_arr))
        end
        i = indexin([e], x_arr)[1]
        i === nothing || (y_arr_dict[d][i] = val)
    end
    dists = sort!(collect(keys(y_arr_dict)))
    y_arrs = [y_arr_dict[d] for d in dists]
    x_arr, dists, y_arrs
end

function plot(plot_i, x_arr, dists, y_arrs)
    println("plot_i = $plot_i")
    println("x_arr = $x_arr")
    println("dists = $dists")
    println("y_arrs = $y_arrs")

    fig, ax = plt.subplots(1, 1)
    if plot_i <= 5
        ax.plot(x_arr, x_arr, ":k")
    end
    for (d, y_arr) in zip(dists, y_arrs)
        ax.plot(x_arr, y_arr, label="$d")
    end
    ax.legend()
    if plot_i == 12
        ax.set_yscale("log")
    else
        ax.loglog()
    end
    title = ["Baseline", "NatAll", "NatInt", "CompAll", "CompInt",
             "SC-SC Err", "L/S Err", "SC-Mode Err", "Cav T1",
             "Trans T1", "L/S Dur", "Cav Size"][plot_i]
    ax.set_title(title)
    nothing
end

function plot_finished(job_id, plot_i, default::Float64=1.0)
    x_arr, dists, y_arrs = fetch_finished(job_id, plot_i, default)
    plot(plot_i, x_arr, dists, y_arrs)
end

end


@testset "eachleaf!" begin
    # test with/without threads
    for node in (Node{Node{LeafNode{Float64, 2, 3}, 2, 3}, 2, 3}(),
                 Node{Node{LeafNode{Float64, 2, 1}, 2, 1}, 2, 1}(),
                 HashNode{Node{LeafNode{Float64, 2, 3}, 2, 3}, 2, 3}(),
                 HashNode{Node{LeafNode{Float64, 2, 1}, 2, 1}, 2, 1}(),
                 DynamicNode{Node{LeafNode{Float64, 2, 3}, 2, 3}, 2}(8,8),
                 DynamicNode{Node{LeafNode{Float64, 2, 1}, 2, 1}, 2}(2,2),
                 DynamicHashNode{Node{LeafNode{Float64, 2, 3}, 2, 3}, 2}(8,8),
                 DynamicHashNode{Node{LeafNode{Float64, 2, 1}, 2, 1}, 2}(2,2),)
        A = TreeView(node)
        A .= 1

        # eachleaf! for all
        eachleaf!(x -> 3x, A)
        @test all(==(3), A)

        # eachleaf! for part matrix
        eachleaf!(x -> 2x, A, :, 1:3)
        @test all(==(6), A[:, 1:3])
        @test all(==(3), A[:, 4:end])
    end
end


id = 0;
counter = 0;
maxdiffreduvscriti = 0.0;
maxfluid = "";
uneq = Set();
critphasefail = Set();
for fluid in coolpropfluids
  id+=1;
  pure = get_fluid_param_string(fluid, "pure");
  tcrit = PropsSI("TCRIT", fluid);
  pcrit = PropsSI("PCRIT", fluid);
  try
    (PhaseSI("P", pcrit, "T", tcrit, fluid) != "critical_point") && push!(critphasefail, fluid);
  catch err
    push!(critphasefail, fluid);
  end
  @test PhaseSI("P", pcrit+50000, "T", tcrit+3, fluid)=="supercritical";
  @test PhaseSI("P", pcrit+50000, "T", tcrit-3, fluid)=="supercritical_liquid";
  @test PhaseSI("P", pcrit-50000, "T", tcrit+3, fluid)=="supercritical_gas";
  psat = PropsSI("P","T",tcrit-3,"Q",1, fluid);
  t3 = PropsSI("Ttriple", fluid);
  p3 = PropsSI("ptriple", fluid);
  @test PhaseSI("P", p3, "T", (t3+tcrit)/2, fluid)=="gas";
  @test PhaseSI("P", (pcrit+p3)/2, "T", (t3+tcrit)/2, fluid)=="liquid";
  (pure == "true") && (tcrit != PropsSI("T_REDUCING", fluid)) && push!(uneq, fluid);
  diffreduvscriti = abs(PropsSI("TCRIT", fluid) - PropsSI("T_REDUCING", fluid));
  if (diffreduvscriti > maxdiffreduvscriti)
    maxfluid = fluid;
    maxdiffreduvscriti = diffreduvscriti;
  end
end
println("max diff between reducing vs critical point T: $maxdiffreduvscriti for $maxfluid");
@test uneq == Set(fails_tcrit_eq_treducing);
println("different reducing vs critical point: $uneq");
@test issubset(critphasefail, Set(fails_critical_point));
println("fails to get phase for critical point: $critphasefail");


abstract type AbstractOperationsProblem end

struct GenericOpProblem<:AbstractOperationsProblem end

mutable struct OperationsTemplate
    transmission::Type{<:PM.AbstractPowerModel}
    devices::Dict{Symbol, DeviceModel}
    branches::Dict{Symbol, DeviceModel}
    services::Dict{Symbol, ServiceModel}
end

"""
    OperationsTemplate(::Type{T}) where {T<:PM.AbstractPowerFormulation}

Creates a model reference of the Power Formulation, devices, branches, and services.

# Arguments
- `model::Type{T<:PM.AbstractPowerFormulation}`:
- `devices::Dict{Symbol, DeviceModel}`: device dictionary
- `branches::Dict{Symbol, BranchModel}`: branch dictionary
- `services::Dict{Symbol, ServiceModel}`: service dictionary

# Example
```julia
template = OperationsTemplate(CopperPlatePowerModel, devices, branches, services)
```
"""
function OperationsTemplate(::Type{T}) where {T<:PM.AbstractPowerModel}

    return  OperationsTemplate(T,
                           Dict{Symbol, DeviceModel}(),
                           Dict{Symbol, DeviceModel}(),
                           Dict{Symbol, ServiceModel}())

end

mutable struct OperationsProblem{M<:AbstractOperationsProblem}
    template::OperationsTemplate
    sys::PSY.System
    canonical::Canonical
end

"""
    OperationsProblem(::Type{M},
    template::OperationsTemplate,
    sys::PSY.System;
    optimizer::Union{Nothing, JuMP.OptimizerFactory}=nothing,
    kwargs...) where {M<:AbstractOperationsProblem,
                      T<:PM.AbstractPowerFormulation}

This builds the optimization model and populates the operation model

# Arguments
- `::Type{M} where {M<:AbstractOperationsProblem, T<:PM.AbstractPowerFormulation} = TestOpProblem`:
The abstract operation model type
- `template::OperationsTemplate`: The model reference made up of transmission, devices,
                                          branches, and services.
- `sys::PSY.System`: the system created using Power Systems

# Output
- `op_problem::OperationsProblem`: The operation model contains the model type, model, Power
Systems system, and optimization model.

# Example
```julia
template = OperationsTemplate(CopperPlatePowerModel, devices, branches, services)
OpModel = OperationsProblem(TestOpProblem, template, system; optimizer = optimizer)
```

# Accepted Key Words
- `verbose::Bool`: verbose default is true
- `PTDF::PTDF`: Passes the PTDF matrix into the optimization model
- `optimizer::union{Nothing, JuMP.OptimizerFactory} = GLPK_optimizer`: The optimizer gets passed
into the optimization model the default is nothing.
- `initial_conditions::DICKDA`: default of Dict{ICKey, Array{InitialCondition}}
- `parameters::Bool`: enable JuMP parameters
- `use_forecast_data::Bool`: if true, forecast collects the time steps in Power Systems,
if false it runs for one time step
- `initial_time::Dates.DateTime`: initial time of forecast
"""
function OperationsProblem(::Type{M},
                        template::OperationsTemplate,
                        sys::PSY.System;
                        optimizer::Union{Nothing, JuMP.OptimizerFactory}=nothing,
                        kwargs...) where {M<:AbstractOperationsProblem}

    op_problem = OperationsProblem{M}(template,
                          sys,
                          Canonical(template.transmission, sys, optimizer; kwargs...))

    build_op_problem!(op_problem; kwargs...)

    return  op_problem

end
"""
    OperationsProblem(op_problem::Type{M},
                    ::Type{T},
                    sys::PSY.System;
                    kwargs...) where {M<:AbstractOperationsProblem,
                                    T<:PM.AbstractPowerFormulation}

This uses the Abstract Power Formulation to build the model reference and
the optimization model and populates the operation model struct.

# Arguments
- `op_problem::Type{M} = where {M<:AbstractOperationsProblem`: Defines the type of the operation model
- `::Type{T} where T<:PM.AbstractPowerFormulation`: The power formulation used for model ref & optimization model
- `sys::PSY.System`: the system created in Power Systems

# Output
- `op_problem::OperationsProblem`: The operation model contains the model type, model, Power
Systems system, and optimization model.

# Example
```julia
template = OperationsTemplate(CopperPlatePowerModel, devices, branches, services)
OpModel = OperationsProblem(TestOpProblem, template, system; optimizer = optimizer)
```

# Accepted Key Words
- `verbose::Bool`: verbose default is true
- `PTDF::PTDF`: Passes the PTDF matrix into the optimization model
- `optimizer::union{Nothing, JuMP.OptimizerFactory}`: The optimizer gets passed
into the optimization model the default is nothing.
- `initial_conditions::DICKDA`: default of Dict{ICKey, Array{InitialCondition}}
- `parameters::Bool`: enable JuMP parameters
- `use_forecast_data::Bool`: if true, forecast collects the time steps in Power Systems,
if false it runs for one time step
- `initial_time::Dates.DateTime`: initial time of forecast

"""
function OperationsProblem(::Type{M},
                        ::Type{T},
                        sys::PSY.System;
                        kwargs...) where {M<:AbstractOperationsProblem,
                                          T<:PM.AbstractPowerModel}

    optimizer = get(kwargs, :optimizer, nothing)
    return OperationsProblem{M}(OperationsTemplate(T),
                          sys,
                          Canonical(T, sys, optimizer; kwargs...))

end
"""
    OperationsProblem(::Type{T},
                    sys::PSY.System;
                    kwargs...) where {M<:AbstractOperationsProblem,
                                    T<:PM.AbstractPowerFormulation}

This uses the Abstract Power Formulation to build the model reference and
the optimization model and populates the operation model struct.

***Note:*** the abstract operation model is set to the default operation model

# Arguments
- `op_problem::Type{M}`: Defines the type of the operation model
- `::Type{T} where T<:PM.AbstractPowerFormulation`: The power formulation used for model ref & optimization model
- `sys::PSY.System`: the system created in Power Systems

# Output
- `op_problem::OperationsProblem`: The operation model contains the model type, model, Power
Systems system, and optimization model.

# Example
```julia
template = OperationsTemplate(CopperPlatePowerModel, devices, branches, services)
OpModel = OperationsProblem(TestOpProblem, template, system; optimizer = optimizer)
```

# Accepted Key Words
- `verbose::Bool`: verbose default is true
- `PTDF::PTDF`: Passes the PTDF matrix into the optimization model
- `optimizer::union{Nothing, JuMP.OptimizerFactory}`: The optimizer gets passed
into the optimization model the default is nothing.
- `initial_conditions::DICKDA`: default of Dict{ICKey, Array{InitialCondition}}
- `parameters::Bool`: enable JuMP parameters
- `use_forecast_data::Bool`: if true, forecast collects the time steps in Power Systems,
if false it runs for one time step
- `initial_time::Dates.DateTime`: initial time of forecast

"""
function OperationsProblem(::Type{T},
                        sys::PSY.System;
                        kwargs...) where {T<:PM.AbstractPowerModel}

    return OperationsProblem(GenericOpProblem,
                         T,
                         sys; kwargs...)

end

get_transmission_ref(op_problem::OperationsProblem) = op_problem.template.transmission
get_devices_ref(op_problem::OperationsProblem) = op_problem.template.devices
get_branches_ref(op_problem::OperationsProblem) = op_problem.template.branches
get_services_ref(op_problem::OperationsProblem) = op_problem.template.services
get_system(op_problem::OperationsProblem) = op_problem.sys

function set_transmission_model!(op_problem::OperationsProblem{M},
                               transmission::Type{T}; kwargs...) where {T<:PM.AbstractPowerModel,
                                                                        M<:AbstractOperationsProblem}

    # Reset the canonical
    op_problem.template.transmission = transmission
    op_problem.canonical = Canonical(transmission,
                                        op_problem.sys,
                                        op_problem.canonical.optimizer_factory; kwargs...)

    build_op_problem!(op_problem; kwargs...)

    return
end

function set_devices_template!(op_problem::OperationsProblem{M},
                          devices::Dict{Symbol, DeviceModel}; kwargs...) where M<:AbstractOperationsProblem

    # Reset the canonical
    op_problem.template.devices = devices
    op_problem.canonical = Canonical(op_problem.template.transmission,
                                        op_problem.sys,
                                        op_problem.canonical.optimizer_factory; kwargs...)

    build_op_problem!(op_problem; kwargs...)

    return
end

function set_branches_template!(op_problem::OperationsProblem{M},
                           branches::Dict{Symbol, DeviceModel}; kwargs...) where M<:AbstractOperationsProblem

    # Reset the canonical
    op_problem.template.branches = branches
    op_problem.canonical = Canonical(op_problem.template.transmission,
                                        op_problem.sys,
                                        op_problem.canonical.optimizer_factory; kwargs...)

    build_op_problem!(op_problem; kwargs...)

    return
end

function set_services_template!(op_problem::OperationsProblem{M},
                           services::Dict{Symbol, DeviceModel}; kwargs...) where M<:AbstractOperationsProblem

    # Reset the canonical
    op_problem.template.services = services
    op_problem.canonical = Canonical(op_problem.template.transmission,
                                        op_problem.sys,
                                        op_problem.canonical.optimizer_factory; kwargs...)

    build_op_problem!(op_problem; kwargs...)

    return
end

function set_device_model!(op_problem::OperationsProblem{M},
                           name::Symbol,
                           device::DeviceModel{D, B}; kwargs...) where {D<:PSY.Injection,
                                                                        B<:AbstractDeviceFormulation,
                                                                        M<:AbstractOperationsProblem}

    if haskey(op_problem.template.devices, name)
        op_problem.template.devices[name] = device
        op_problem.canonical = Canonical(op_problem.template.transmission,
                                            op_problem.sys,
                                            op_problem.canonical.optimizer_factory; kwargs...)
        build_op_problem!(op_problem; kwargs...)
    else
        error("Device Model with name $(name) doesn't exist in the model")
    end

    return

end

function set_branch_model!(op_problem::OperationsProblem{M},
                           name::Symbol,
                           branch::DeviceModel{D, B}; kwargs...) where {D<:PSY.Branch,
                                                                        B<:AbstractDeviceFormulation,
                                                                        M<:AbstractOperationsProblem}

    if haskey(op_problem.template.branches, name)
        op_problem.template.branches[name] = branch
        op_problem.canonical = Canonical(op_problem.template.transmission,
                                            op_problem.sys,
                                            op_problem.canonical.optimizer_factory; kwargs...)
        build_op_problem!(op_problem; kwargs...)
    else
        error("Branch Model with name $(name) doesn't exist in the model")
    end

    return

end

function set_services_model!(op_problem::OperationsProblem{M},
                             name::Symbol,
                             service::DeviceModel; kwargs...) where M<:AbstractOperationsProblem

    if haskey(op_problem.template.devices, name)
        op_problem.template.services[name] = service
        op_problem.canonical = Canonical(op_problem.template.transmission,
                                            op_problem.sys,
                                            op_problem.canonical.optimizer_factory; kwargs...)
        build_op_problem!(op_problem; kwargs...)
    else
        error("Branch Model with name $(name) doesn't exist in the model")
    end

    return

end

function construct_device!(op_problem::OperationsProblem,
                           name::Symbol,
                           device_model::DeviceModel;
                           kwargs...)

    if haskey(op_problem.template.devices, name)
        error("Device with model name $(name) already exists in the Opertaion Model")
    end

    devices_ref = get_devices_ref(op_problem)
    devices_ref[name] = device_model

    construct_device!(op_problem.canonical,
                      get_system(op_problem),
                      device_model,
                      get_transmission_ref(op_problem);
                      kwargs...)

    JuMP.@objective(op_problem.canonical.JuMPmodel,
                    MOI.MIN_SENSE,
                    op_problem.canonical.cost_function)

    return

end

function construct_network!(op_problem::OperationsProblem; kwargs...)

    construct_network!(op_problem, op_problem.template.transmission; kwargs...)

    return
end

function construct_network!(op_problem::OperationsProblem,
                            system_formulation::Type{T};
                            kwargs...) where {T<:PM.AbstractPowerModel}

    construct_network!(op_problem.canonical, get_system(op_problem), T; kwargs...)

    return
end

function get_initial_conditions(op_problem::OperationsProblem)
    return op_problem.canonical.initial_conditions
end

function get_initial_conditions(op_problem::OperationsProblem,
                                ic::InitialConditionQuantity,
                                device::PSY.Device)

    canonical = op_problem.canonical
    key = ICKey(ic, device)

    return get_initial_conditions(canonical, key)

end


# -----------------------------------------------------------------------------
# Implementation of Atomistic MolecularDynamicsSimulator API
# -----------------------------------------------------------------------------

"""
    NBSimulator <: MolecularDynamicsSimulator

A wrapper around NBodySimulator to implement the Atomistic API.

**Type parameter descriptions**
- `S` the type of ODE solver
- `T<:Unitful.Time` the type for the time fields
- `C` the type for the thermostat field

**Field descriptions**
- `Δt::T` the time between timesteps, assumed to be in atomic units
- `steps::Integer` the number of timesteps for the simulation
- `t₀::T` the starting time of the simulation, assumed to be in atomic units;
    defaults to 0
- `thermostat::C` the thermostat for the simulation;
    many options are defined by `NBodySimulator`, but a user could also define a custom thermostat;
    defaults to the `NullThermostat`
- `simulator::S` the algorithm to be used for the ODE;
    defaults to VelocityVerlet
"""
struct NBSimulator{S<:OrdinaryDiffEqAlgorithm,T<:Unitful.Time,C<:Thermostat} <: MolecularDynamicsSimulator
    Δt::T
    steps::Int
    t₀::T
    thermostat::C
    simulator::S
end
function NBSimulator(Δt::T, steps::Int; t₀::Real=zero(T), thermostat::C=NullThermostat(), simulator::S=NBodySimulator.VelocityVerlet()) where {S<:OrdinaryDiffEqAlgorithm,T<:Real,C<:Thermostat}
    Δt, t₀ = promote(Δt * TIME_UNIT, t₀ * TIME_UNIT)
    NBSimulator(Δt, steps, t₀, thermostat, simulator)
end
function NBSimulator(Δt::T, steps::Integer; t₀::Unitful.Time=zero(T), thermostat::C=NullThermostat(), simulator::S=NBodySimulator.VelocityVerlet()) where {S<:OrdinaryDiffEqAlgorithm,T<:Unitful.Time,C<:Thermostat}
    Δt, t₀ = promote(Δt, t₀)
    NBSimulator(Δt, steps, t₀, thermostat, simulator)
end

# Extract the tuple of start_time, end_time from the simulator
time_range(simulator::NBSimulator) = Float64.(austrip.((simulator.t₀, simulator.t₀ + simulator.steps * simulator.Δt)))

function simulate(system::AbstractSystem{3}, simulator::NBSimulator, potential::AbstractPotential)
    wrapper = InteratomicPotentialParameters(potential)
    potentials = Dict{Symbol,PotentialParameters}(:custom => wrapper)
    nb_system = PotentialNBodySystem{ElementMassBody}(ElementMassBody.(system), potentials)
    simulation = NBodySimulation(nb_system, time_range(simulator), nbs_boundary_conditions(system), simulator.thermostat, austrip(u"k"))
    result = run_simulation(simulation, simulator.simulator, dt=austrip(simulator.Δt))
    NBSResult(result, wrapper.energy_cache)
end

# -----------------------------------------------------------------------------
# Integration with InteratomicPotentials
# -----------------------------------------------------------------------------

# Internal struct that represents a potential in the format accepted by NBodySimulator
# Wraps the underlying InteratomicPotential with a cache of the forces for the current timestep
# Also stores the potential energy from every timestep to be referenceable after the simulation
struct InteratomicPotentialParameters{P<:AbstractPotential} <: PotentialParameters
    potential::P
    timestep_cache::Ref{Float64}
    force_cache::Ref{Vector{SVector{3,Float64}}}
    energy_cache::Vector{ENERGY_TYPE}
    InteratomicPotentialParameters(potential::AbstractPotential) = new{typeof(potential)}(potential, Ref{Float64}(), Ref{Vector{SVector{3,Float64}}}(), Vector{ENERGY_TYPE}())
end

function NBodySimulator.get_accelerating_function(parameters::InteratomicPotentialParameters, simulation::NBodySimulation)
    bodies = simulation.system.bodies
    boundary_conditions = simulation.boundary_conditions
    (dv, u, v, t, i) -> begin
        if !isassigned(parameters.timestep_cache) || t != parameters.timestep_cache[]
            # TODO: avoid copying data at every timestep
            particles = [AtomsBase.Atom(b, SVector{3}(r), SVector{3}(v), boundary_conditions) for (b, r, v) ∈ zip(bodies, eachcol(u), eachcol(v))]
            system = FlexibleSystem(particles, get_bounding_box(boundary_conditions), get_boundary_conditions(boundary_conditions))
            eandf = energy_and_force(system, parameters.potential)
            parameters.timestep_cache[] = t
            parameters.force_cache[] = [austrip.(f) for f ∈ eandf.f]
            push!(parameters.energy_cache, uconvert(ENERGY_UNIT, eandf.e))
        end
        dv .+= parameters.force_cache[][i] ./ bodies[i].m
    end
end


module MosekTools

import MathOptInterface
const MOI = MathOptInterface
const MOIU = MOI.Utilities

using Mosek
export Mosek
# Allows the user to use `Mosek.Optimizer` instead of `MosekTools.Optimizer`
# for convenience and for consitency with other solvers where the syntax is
# `SolverName.Optimizer`.
function Mosek.Optimizer(; kws...)
    return MosekTools.Optimizer(; kws...)
end

include("LinkedInts.jl")

const DEBUG = false

struct MosekSolution
    whichsol :: Soltype
    solsta   :: Solsta
    prosta   :: Prosta

    xxstatus :: Vector{Stakey}
    xx       :: Vector{Float64}
    barxj    :: Vector{Vector{Float64}}
    slx      :: Vector{Float64}
    sux      :: Vector{Float64}
    snx      :: Vector{Float64}

    cstatus  :: Vector{Stakey}
    xc       :: Vector{Float64}
    slc      :: Vector{Float64}
    suc      :: Vector{Float64}
    y        :: Vector{Float64}
end

struct ColumnIndex
    value::Int32
end

struct ColumnIndices
    values::Vector{Int32}
end

struct MatrixIndex
    matrix::Int32
    row::Int32
    column::Int32
    function MatrixIndex(matrix::Integer, row::Integer, column::Integer)
        @assert column ≤ row
        new(matrix, row, column)
    end
end

"""
    Optimizer <: MathOptInterface.AbstractModel

Linear variables and constraint can be deleted. For some reason MOSEK
does not support deleting PSD variables.

Note also that adding variables and constraints will permanently add
some (currently between 1 and 3) Int64s that a `delete!` will not
remove. This ensures that Indices (Variable and constraint) that
are deleted are thereafter invalid.
"""
mutable struct Optimizer  <: MOI.AbstractOptimizer
    task :: Mosek.MSKtask
    ## Options passed in `Mosek.Optimizer` that are used to create a new task
    ## in `MOI.empty!`:
    # Should Mosek output be ignored or printed ?
    be_quiet :: Bool
    # Integer parameters, i.e. parameters starting with `MSK_IPAR_`
    ipars :: Dict{String, Int32}
    # Floating point parameters, i.e. parameters starting with `MSK_DPAR_`
    dpars :: Dict{String, Float64}
    # String parameters, i.e. parameters starting with `MSK_SPAR_`
    spars :: Dict{String, AbstractString}
    has_variable_names::Bool
    constrnames :: Dict{String, Vector{MOI.ConstraintIndex}}
    # Mosek only support names for `MOI.ScalarAffineFunction` so we need a
    # fallback for `SingleVariable` and `VectorOfVariables`.
    con_to_name :: Dict{MOI.ConstraintIndex, String}

    # For each MOI index of variables, gives the flags of constraints present
    # The SingleVariable constraints added cannot just be inferred from getvartype
    # and getvarbound so we need to keep them here so implement `MOI.is_valid`
    x_constraints::Vector{UInt8}

    """
        The total length of `x_block` matches the number of variables in
    the underlying task, and the number of blocks corresponds to the
    number variables allocated in the Model.
    """
    x_block::LinkedInts

    """
    One entry per scalar variable in the task indicating in which semidefinite
    block it is and at which index.
    MOI index -> MatrixIndex
    """
    x_sd::Vector{MatrixIndex}

    sd_dim::Vector{Int}

    ###########################
    """
    One scalar entry per constraint in the underlying task. One block
    per constraint allocated in the Model.
    """
    c_block :: LinkedInts

    # i -> 0: Not in a VectorOfVariables constraint
    # i -> +j: In `MOI.ConstraintIndex{MOI.VectorOfVariables, ?}(j)`
    # i -> -j: In `MOI.VectorOfVariables` constraint with `MOI.VariableIndex(j)` as first variable
    variable_to_vector_constraint_id::Vector{Int32}

    ###########################
    trm :: Union{Nothing, Rescode}
    solutions :: Vector{MosekSolution}

    ###########################
    """
    Indicating whether the objective sense is MOI.FEASIBILITY_SENSE. It is
    encoded as a MOI.MIN_SENSE with a zero objective internally but this allows
    MOI.get(::Optimizer, ::ObjectiveSense) to still return the right value
    """
    feasibility :: Bool
    """
    Indicates whether there is an objective set.
    If `has_objective` is `false` then Mosek has a zero objective internally.
    This affects `MOI.ListOfModelAttributesSet`.
    """
    has_objective :: Bool

    fallback :: Union{String, Nothing}

    function Optimizer(; kws...)
        optimizer = new(
            maketask(), # task
            false, # be_quiet
            Dict{String, Int32}(), # ipars
            Dict{String, Float64}(), # dpars
            Dict{String, AbstractString}(), # spars
            false, # has_variable_names
            Dict{String, Vector{MOI.ConstraintIndex}}(), # constrnames
            Dict{MOI.ConstraintIndex, String}(), # con_to_name
            UInt8[], # x_constraints
            LinkedInts(),# x_block
            MatrixIndex[], # x_sd
            Int[], # sd_dim
            LinkedInts(), # c_block
            Int32[], # variable_to_vector_constraint_id
            nothing,# trm
            MosekSolution[],
            true, # feasibility_sense
            false, # has_objective
            nothing,
        )
        Mosek.putstreamfunc(optimizer.task, Mosek.MSK_STREAM_LOG, m -> print(m))
        if length(kws) > 0
            @warn("""Passing optimizer attributes as keyword arguments to
            Mosek.Optimizer is deprecated. Use
                MOI.set(model, MOI.RawOptimizerAttribute("key"), value)
            or
                JuMP.set_optimizer_attribute(model, "key", value)
            instead.
            """)
        end
        for (option, value) in kws
            MOI.set(optimizer, MOI.RawOptimizerAttribute(string(option)), value)
        end
        return optimizer
    end
end

struct IntegerParameter <: MOI.AbstractOptimizerAttribute
    name::String
end
function MOI.set(m::Optimizer, p::IntegerParameter, value)
    m.ipars[p.name] = value
    Mosek.putnaintparam(m.task, p.name, value)
end
function MOI.get(m::Optimizer, p::IntegerParameter)
    Mosek.getnaintparam(m.task, p.name)
end

struct DoubleParameter <: MOI.AbstractOptimizerAttribute
    name::String
end
function MOI.set(m::Optimizer, p::DoubleParameter, value)
    m.dpars[p.name] = value
    Mosek.putnadouparam(m.task, p.name, value)
end
function MOI.get(m::Optimizer, p::DoubleParameter)
    Mosek.getnadouparam(m.task, p.name)
end

struct StringParameter <: MOI.AbstractOptimizerAttribute
    name::String
end
function MOI.set(m::Optimizer, p::StringParameter, value::AbstractString)
    m.spars[p.name] = value
    Mosek.putnastrparam(m.task, p.name, value)
end
function MOI.get(m::Optimizer, p::StringParameter)
    # We need to give the maximum length of the value of the parameter.
    # 255 should be ok in most cases.
    len, str = Mosek.getnastrparam(m.task, p.name, 255)
    return str
end

function MOI.set(m::Optimizer, p::MOI.RawOptimizerAttribute, value)
    if p.name == "QUIET"
        if m.be_quiet != convert(Bool, value)
            m.be_quiet = !m.be_quiet
            if m.be_quiet
                Mosek.putstreamfunc(m.task, Mosek.MSK_STREAM_LOG,
                                    m -> begin end)
            else
                Mosek.putstreamfunc(m.task, Mosek.MSK_STREAM_LOG,
                                    m -> print(m))
            end
        end
    elseif p.name == "fallback"
        m.fallback = value
    else
        if startswith(p.name, "MSK_IPAR_")
            par = IntegerParameter(p.name)
        elseif startswith(p.name, "MSK_DPAR_")
            par = DoubleParameter(p.name)
        elseif startswith(p.name, "MSK_SPAR_")
            par = StringParameter(p.name)
        elseif isa(value, Integer)
            par = IntegerParameter("MSK_IPAR_" * p.name)
        elseif isa(value, AbstractFloat)
            par = DoubleParameter("MSK_DPAR_" * p.name)
        elseif isa(value, AbstractString)
            par = StringParameter("MSK_SPAR_" * p.name)
        else
            error("Value $value for parameter $(p.name) has unrecognized type")
        end
        MOI.set(m, par, value)
    end
end

function MOI.get(m::Optimizer, p::MOI.RawOptimizerAttribute)
    if p.name == "QUIET"
        return m.be_quiet
    elseif p.name == "fallback"
        return m.fallback
    else
        if startswith(p.name, "MSK_IPAR_")
            par = IntegerParameter(p.name)
        elseif startswith(p.name, "MSK_DPAR_")
            par = DoubleParameter(p.name)
        elseif startswith(p.name, "MSK_SPAR_")
            par = StringParameter(p.name)
        else
            error("The parameter $(p.name) should start by `MSK_IPAR_`, `MSK_DPAR_` or `MSK_SPAR_`.")
        end
        MOI.get(m, par)
    end
end

MOI.supports(::Optimizer, ::MOI.Silent) = true
function MOI.set(model::Optimizer, ::MOI.Silent, value::Bool)
    MOI.set(model, MOI.RawOptimizerAttribute("QUIET"), value)
end
function MOI.get(model::Optimizer, ::MOI.Silent)
    MOI.get(model, MOI.RawOptimizerAttribute("QUIET"))
end

MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
function MOI.set(model::Optimizer, ::MOI.TimeLimitSec, value::Real)
    MOI.set(model, MOI.RawOptimizerAttribute("MSK_DPAR_OPTIMIZER_MAX_TIME"), value)
end
function MOI.set(model::Optimizer, ::MOI.TimeLimitSec, ::Nothing)
    MOI.set(model, MOI.RawOptimizerAttribute("MSK_DPAR_OPTIMIZER_MAX_TIME"), -1.0)
end
function MOI.get(model::Optimizer, ::MOI.TimeLimitSec)
    value = MOI.get(model, MOI.RawOptimizerAttribute("MSK_DPAR_OPTIMIZER_MAX_TIME"))
    if value < 0.0
        return nothing
    else
        return value
    end
end

function matrix_solution(m::Optimizer, sol)
    return Vector{Float64}[getbarxj(m.task, sol, j) for j in 1:length(m.sd_dim)]
end

function MOI.optimize!(m::Optimizer)
    # See https://github.com/jump-dev/MosekTools.jl/issues/70
    putintparam(m.task,Mosek.MSK_IPAR_REMOVE_UNUSED_SOLUTIONS,Mosek.MSK_ON)
    m.trm = if m.fallback == nothing; optimize(m.task) else optimize(m.task,m.fallback) end
    m.solutions = MosekSolution[]
    # If the problem is conic but a starting value is set,
    # `MSK_SOL_ITG` simply contains the starting value and
    # `MSK_SOL_ITR` contains the actual solution found.
    # We want `VariablePrimal(1)`, ... to be `MSK_SOL_ITR` in that case
    # so puting it first should solve the issue in that case, see
    # https://github.com/jump-dev/MosekTools.jl/issues/66
    if solutiondef(m.task,MSK_SOL_ITR)
        push!(m.solutions,
              MosekSolution(MSK_SOL_ITR,
                            getsolsta(m.task,MSK_SOL_ITR),
                            getprosta(m.task,MSK_SOL_ITR),
                            getskx(m.task,MSK_SOL_ITR),
                            getxx(m.task,MSK_SOL_ITR),
                            matrix_solution(m, MSK_SOL_ITR),
                            getslx(m.task,MSK_SOL_ITR),
                            getsux(m.task,MSK_SOL_ITR),
                            getsnx(m.task,MSK_SOL_ITR),
                            getskc(m.task,MSK_SOL_ITR),
                            getxc(m.task,MSK_SOL_ITR),
                            getslc(m.task,MSK_SOL_ITR),
                            getsuc(m.task,MSK_SOL_ITR),
                            gety(m.task,MSK_SOL_ITR)))
    end
    if solutiondef(m.task,MSK_SOL_ITG)
        push!(m.solutions,
              MosekSolution(MSK_SOL_ITG,
                            getsolsta(m.task, MSK_SOL_ITG),
                            getprosta(m.task, MSK_SOL_ITG),
                            getskx(m.task, MSK_SOL_ITG),
                            getxx(m.task, MSK_SOL_ITG),
                            # See https://github.com/jump-dev/MosekTools.jl/issues/71
                            Float64[], #matrix_solution(m, MSK_SOL_ITG),
                            Float64[],
                            Float64[],
                            Float64[],
                            getskc(m.task, MSK_SOL_ITG),
                            getxc(m.task, MSK_SOL_ITG),
                            Float64[],
                            Float64[],
                            Float64[]))
    end
    if solutiondef(m.task,MSK_SOL_BAS)
        push!(m.solutions,
              MosekSolution(MSK_SOL_BAS,
                            getsolsta(m.task,MSK_SOL_BAS),
                            getprosta(m.task,MSK_SOL_BAS),
                            getskx(m.task,MSK_SOL_BAS),
                            getxx(m.task,MSK_SOL_BAS),
                            # See https://github.com/jump-dev/MosekTools.jl/issues/71
                            Float64[], #matrix_solution(m, MSK_SOL_BAS),
                            getslx(m.task,MSK_SOL_BAS),
                            getsux(m.task,MSK_SOL_BAS),
                            Float64[],
                            getskc(m.task,MSK_SOL_BAS),
                            getxc(m.task,MSK_SOL_BAS),
                            getslc(m.task,MSK_SOL_BAS),
                            getsuc(m.task,MSK_SOL_BAS),
                            gety(m.task,MSK_SOL_BAS)))
    end
end

MOI.supports(::Optimizer, ::MOI.Name) = true
function MOI.set(m::Optimizer, ::MOI.Name, name::String)
    puttaskname(m.task, name)
end
function MOI.get(m::Optimizer, ::MOI.Name)
    return gettaskname(m.task)
end
function MOI.get(m::Optimizer, ::MOI.ListOfModelAttributesSet)
    set = MOI.AbstractModelAttribute[]
    if !m.feasibility
        push!(set, MOI.ObjectiveSense())
    end
    if m.has_objective
        push!(set, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    end
    if !isempty(MOI.get(m, MOI.Name()))
        push!(set, MOI.Name())
    end
    return set
end

function MOI.is_empty(m::Optimizer)
    getnumvar(m.task) == 0 && getnumcon(m.task) == 0 && getnumcone(m.task) == 0 && getnumbarvar(m.task) == 0
end

function MOI.empty!(model::Optimizer)
    model.task               = maketask()
    for (name, value) in model.ipars
        Mosek.putnaintparam(model.task, name, value)
    end
    for (name, value) in model.dpars
        Mosek.putnadouparam(model.task, name, value)
    end
    for (name, value) in model.spars
        Mosek.putnastrparam(model.task, name, value)
    end
    if !model.be_quiet
        Mosek.putstreamfunc(model.task, Mosek.MSK_STREAM_LOG, m -> print(m))
    end
    model.has_variable_names = false
    empty!(model.constrnames)
    empty!(model.con_to_name)
    empty!(model.x_constraints)
    model.x_block            = LinkedInts()
    empty!(model.x_sd)
    empty!(model.sd_dim)
    model.c_block            = LinkedInts()
    empty!(model.variable_to_vector_constraint_id)
    model.trm                = nothing
    empty!(model.solutions)
    model.feasibility        = true
    model.has_objective      = false
end

MOI.get(::Optimizer, ::MOI.SolverName) = "Mosek"

function MOI.get(::Optimizer, ::MOI.SolverVersion)
    major, minor, revision = Mosek.getversion()
    return string(VersionNumber(major, minor, revision))
end

MOI.supports_incremental_interface(::Optimizer) = true
function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike; kws...)
    return MOIU.default_copy_to(dest, src; kws...)
end

function MOI.write_to_file(m::Optimizer, filename :: String)
    putintparam(m.task,MSK_IPAR_OPF_WRITE_SOLUTIONS, MSK_ON)
    writedata(m.task,filename)
end

# For linear objectives we accept:
# EITER affine left-hand side and ranged, unbounded, half-open, fixed (equality), PSD or SOC domains
# OR affine and quadratic left-hand side, and ranged, unbounded, half-open, fixed (equality) domains (quadratic constraints must be unbounded or half-open)
#
# For non-quadratic problems we allow binary and integer variables (but not constraints)
#function supportsconstraints(m::MosekSolver, constraint_types) :: Bool
#    for (fun,dom) in constraint_types
#        if  fun in [MOI.ScalarAffineFunction{Float64},
#                    MOI.VariableIndex,
#                    MOI.VectorOfVariables] &&
#            dom in [MOI.GreaterThan{Float64},
#                    MOI.LessThan{Float64},
#                    MOI.EqualTo{Float64},
#                    MOI.Interval{Float64},
#                    MOI.SecondOrderCone,
#                    MOI.RotatedSecondOrderCone,
#                    MOI.PositiveSemidefiniteConeTriangle,
#                    MOI.PositiveSemidefiniteConeScaled ]
#            # ok
#        elseif dom == MOI.Integer && fun in [MOI.VariableIndex, MOI.VectorOfVariables]
#            # ok
#        else
#            return false
#        end
#    end
#    true
#end

ref2id(vi::MOI.VariableIndex)::Int = vi.value
ref2id(ci::MOI.ConstraintIndex)::Int = ci.value

include("objective.jl")
include("variable.jl")
include("constraint.jl")
include("attributes.jl")

end # module


function tabulated_properties()
    # Various properties. Taken from CoolProp.
    data = Dict{String, MolecularProperty}()
    data["R507A"] = MolecularProperty(0.0988592, 3704900, 343.765, 0.0002014492399, 0.286)
    data["RC318"] = MolecularProperty(0.2000312, 2777500, 388.38, 0.0003226451742, 0.355345)
    data["SES36"] = MolecularProperty(0.18485, 2849000, 450.7, 0.0003571428571, 0.352)
    data["SulfurDioxide"] = MolecularProperty(0.0640638, 7886587.6, 430.64, 0.0001220256254, 0.2561299362)
    data["SulfurHexafluoride"] = MolecularProperty(0.1460554192, 3754983, 318.7232, 0.0001967606348, 0.21)
    data["Toluene"] = MolecularProperty(0.09213842, 4126000, 591.75, 0.000315556958, 0.2657)
    data["trans-2-Butene"] = MolecularProperty(0.05610632, 4027300, 428.61, 0.0002373605507, 0.2100768344)
    data["Water"] = MolecularProperty(0.018015268, 22064000, 647.096, 5.594803727e-05, 0.3442920843)
    data["Xenon"] = MolecularProperty(0.131293, 5842000, 289.733, 0.000119047619, 0.00363)
    data["1-Butene"] = MolecularProperty(0.05610632, 4005100, 419.29, 0.0002358490566, 0.1918606474)
    data["Acetone"] = MolecularProperty(0.05807914, 4700000, 508.1, 0.0002127659574, 0.3071)
    data["Air"] = MolecularProperty(0.02896546, 3786000, 132.5306, 8.452513778e-05, 0.0335)
    data["Ammonia"] = MolecularProperty(0.01703026, 11333000, 405.4, 7.569004444e-05, 0.25601)
    data["Argon"] = MolecularProperty(0.039948, 4863000, 150.687, 7.458551158e-05, -0.00219)
    data["Benzene"] = MolecularProperty(0.0781118, 4894000, 562.02, 0.0002562788314, 0.2108369733)
    data["CarbonDioxide"] = MolecularProperty(0.0440098, 7377300, 304.1282, 9.411847707e-05, 0.22394)
    data["CarbonMonoxide"] = MolecularProperty(0.0280101, 3494000, 132.86, 9.216589862e-05, 0.0497)
    data["CarbonylSulfide"] = MolecularProperty(0.0600751, 6370000, 378.77, 0.0001349527665, 0.0978)
    data["cis-2-Butene"] = MolecularProperty(0.05610632, 4225500, 435.75, 0.0002356267672, 0.2023595859)
    data["CycloHexane"] = MolecularProperty(0.08415948, 4082400, 553.6, 0.0003101736973, 0.20926)
    data["Cyclopentane"] = MolecularProperty(0.0701329, 4571200, 511.72, 0.0002617801047, 0.2011125734)
    data["CycloPropane"] = MolecularProperty(0.042081, 5579700, 398.3, 0.0001627891671, 0.1305495664)
    data["D4"] = MolecularProperty(0.29661576, 1347215.356, 586.5, 0.0009587727709, 0.5981248268)
    data["D5"] = MolecularProperty(0.3707697, 1160000, 619.15, 0.001216000031, 0.658)
    data["D6"] = MolecularProperty(0.444924, 961000, 645.78, 0.001594162692, 0.736)
    data["Deuterium"] = MolecularProperty(0.0040282, 1679600, 38.34, 5.803830528e-05, -0.1362902741)
    data["Dichloroethane"] = MolecularProperty(0.098959, 5227585.103, 561.6, 0.0002309468822, 0.2685785009)
    data["DiethylEther"] = MolecularProperty(0.0741216, 3649016.897, 466.7, 0.0002807636367, 0.2819293312)
    data["DimethylCarbonate"] = MolecularProperty(0.0900779, 4908800, 557, 0.00025, 0.346)
    data["DimethylEther"] = MolecularProperty(0.04606844, 5336800, 400.378, 0.0001683501684, 0.196)
    data["Ethane"] = MolecularProperty(0.03006904, 4872200, 305.322, 0.0001458387816, 0.099)
    data["Ethanol"] = MolecularProperty(0.04606844, 6268000, 514.71, 0.0001686340641, 0.644)
    data["EthylBenzene"] = MolecularProperty(0.106165, 3622400, 617.12, 0.0003648282243, 0.304)
    data["Ethylene"] = MolecularProperty(0.02805376, 5041800, 282.35, 0.0001309454817, 0.0866)
    data["EthyleneOxide"] = MolecularProperty(0.04405256, 7304686.16, 468.92, 0.0001394700139, 0.210195489)
    data["Fluorine"] = MolecularProperty(0.03799681, 5172400, 144.414, 6.409023906e-05, 0.0449)
    data["HeavyWater"] = MolecularProperty(0.020027508, 21672051.48, 643.847, 5.625704971e-05, 0.3643005688)
    data["Helium"] = MolecularProperty(0.004002602, 227600, 5.1953, 5.515719801e-05, -0.385)
    data["HFE143m"] = MolecularProperty(0.10004, 3635000, 377.921, 0.0002151397849, 0.2888713657)
    data["Hydrogen"] = MolecularProperty(0.00201588, 1296400, 33.145, 6.448284756e-05, -0.219)
    data["HydrogenChloride"] = MolecularProperty(0.0364609, 8288160.904, 324.55, 8.474576271e-05, 0.1281892174)
    data["HydrogenSulfide"] = MolecularProperty(0.03408088, 9000000, 373.1, 9.813542689e-05, 0.1005)
    data["IsoButane"] = MolecularProperty(0.0581222, 3629000, 407.817, 0.0002577481153, 0.1835317832)
    data["IsoButene"] = MolecularProperty(0.05610632, 4009800, 418.09, 0.0002398081535, 0.1925934522)
    data["Isohexane"] = MolecularProperty(0.08617536, 3040000, 497.7, 0.0003683241252, 0.2797)
    data["Isopentane"] = MolecularProperty(0.07214878, 3378000, 460.35, 0.0003057169061, 0.2274)
    data["Krypton"] = MolecularProperty(0.083798, 5525000, 209.48, 9.216589862e-05, -0.00089)
    data["m-Xylene"] = MolecularProperty(0.106165, 3534600, 616.89, 0.0003752345216, 0.326)
    data["MD2M"] = MolecularProperty(0.310685, 1227000, 599.4, 0.001093300515, 0.668)
    data["MD3M"] = MolecularProperty(0.384839, 945000, 628.36, 0.001458154971, 0.722)
    data["MD4M"] = MolecularProperty(0.45899328, 877000, 653.2, 0.001650000016, 0.8246471473)
    data["MDM"] = MolecularProperty(0.23653146, 1410044.756, 564.09, 0.000921288245, 0.5280658491)
    data["Methane"] = MolecularProperty(0.0160428, 4599200, 190.564, 9.862781099e-05, 0.01142)
    data["Methanol"] = MolecularProperty(0.03204216, 8215850, 512.5, 0.0001173705495, 0.5720322)
    data["MethylLinoleate"] = MolecularProperty(0.29447206, 1341000, 799, 0.001237011381, 0.8054063871)
    data["MethylLinolenate"] = MolecularProperty(0.29245618, 1369000, 772, 0.001180219521, 1.142605259)
    data["MethylOleate"] = MolecularProperty(0.29648794, 1246000, 782, 0.001230239282, 0.90584936)
    data["MethylPalmitate"] = MolecularProperty(0.27045066, 1350000, 755, 0.001114827202, 0.910316178)
    data["MethylStearate"] = MolecularProperty(0.29850382, 1239000, 775, 0.001258970162, 1.01756)
    data["MM"] = MolecularProperty(0.16237752, 1939000, 518.75, 0.0006290000472, 0.418)
    data["n-Butane"] = MolecularProperty(0.0581222, 3796000, 425.125, 0.0002549219298, 0.2008100946)
    data["n-Decane"] = MolecularProperty(0.14228168, 2103000, 617.7, 0.0006097560976, 0.4884)
    data["n-Dodecane"] = MolecularProperty(0.17033484, 1817000, 658.1, 0.0007518796992, 0.5741822212)
    data["n-Heptane"] = MolecularProperty(0.100202, 2736000, 540.13, 0.0004319051724, 0.349)
    data["n-Hexane"] = MolecularProperty(0.08617536, 3034000, 507.82, 0.000369565829, 0.299)
    data["n-Nonane"] = MolecularProperty(0.1282551, 2281000, 594.55, 0.0005524861878, 0.4433)
    data["n-Octane"] = MolecularProperty(0.1142285, 2497000, 569.32, 0.0004862867146, 0.395)
    data["n-Pentane"] = MolecularProperty(0.07214878, 3370000, 469.7, 0.0003109861207, 0.251)
    data["n-Propane"] = MolecularProperty(0.04409562, 4251200, 369.89, 0.0002, 0.1521)
    data["n-Undecane"] = MolecularProperty(0.15630826, 1990400, 638.8, 0.00066010223, 0.5390371014)
    data["Neon"] = MolecularProperty(0.020179, 2680000, 44.4918, 4.187253999e-05, -0.03844929927)
    data["Neopentane"] = MolecularProperty(0.07214878, 3196000, 433.74, 0.0003058103976, 0.1961)
    data["Nitrogen"] = MolecularProperty(0.02801348, 3395800, 126.192, 8.941423556e-05, 0.0372)
    data["NitrousOxide"] = MolecularProperty(0.0440128, 7245000, 309.52, 9.737098345e-05, 0.1613)
    data["Novec649"] = MolecularProperty(0.3160438, 1869026.583, 441.81, 0.0005208333333, 0.4710235936)
    data["o-Xylene"] = MolecularProperty(0.106165, 3737500, 630.259, 0.0003725088471, 0.312)
    data["OrthoDeuterium"] = MolecularProperty(0.0040282, 1679600, 38.34, 5.803830528e-05, -0.1362902741)
    data["OrthoHydrogen"] = MolecularProperty(0.00201594, 1310650, 33.22, 6.474779953e-05, -0.219)
    data["Oxygen"] = MolecularProperty(0.0319988, 5043000, 154.581, 7.336757153e-05, 0.0222)
    data["p-Xylene"] = MolecularProperty(0.106165, 3531500, 616.168, 0.0003712062719, 0.324)
    data["ParaDeuterium"] = MolecularProperty(0.0040282, 1679600, 38.34, 5.803830528e-05, -0.1362902741)
    data["ParaHydrogen"] = MolecularProperty(0.00201588, 1285800, 32.938, 6.435834728e-05, -0.219)
    data["Propylene"] = MolecularProperty(0.04207974, 4555000, 364.211, 0.0001832508704, 0.146)
    data["Propyne"] = MolecularProperty(0.04006, 5626000, 402.38, 0.0001635769703, 0.204)
    data["R11"] = MolecularProperty(0.137368, 4394000, 471.06, 0.0002431292035, 0.1887506483)
    data["R113"] = MolecularProperty(0.187375, 3392200, 487.21, 0.0003345982143, 0.252535)
    data["R114"] = MolecularProperty(0.170921, 3257000, 418.83, 0.0002947070612, 0.2523)
    data["R115"] = MolecularProperty(0.154466416, 3129036.968, 353.1, 0.0002512562814, 0.2484349931)
    data["R116"] = MolecularProperty(0.13801182, 3048000, 293.03, 0.0002250225023, 0.2566)
    data["R12"] = MolecularProperty(0.120913, 4136100, 385.12, 0.0002140053097, 0.1794783173)
    data["R123"] = MolecularProperty(0.152931, 3672000, 456.831, 0.0002780545193, 0.281922497)
    data["R1233zd(E)"] = MolecularProperty(0.1304944, 3623637.776, 439.6, 0.0002717391304, 0.3024522936)
    data["R1234yf"] = MolecularProperty(0.1140415928, 3382200, 367.85, 0.0002398081535, 0.276)
    data["R1234ze(E)"] = MolecularProperty(0.1140415928, 3636250, 382.52, 0.0002331002331, 0.313)
    data["R1234ze(Z)"] = MolecularProperty(0.1140415928, 3533000, 423.27, 0.0002426416868, 0.3274)
    data["R124"] = MolecularProperty(0.1364762, 3624295, 395.425, 0.0002437075, 0.2880950842)
    data["R125"] = MolecularProperty(0.1200214, 3617700, 339.173, 0.0002092487968, 0.3052)
    data["R13"] = MolecularProperty(0.104459, 3879000, 301.88, 0.0001792114695, 0.1745863278)
    data["R134a"] = MolecularProperty(0.102032, 4059280, 374.21, 0.0001993201985, 0.32684)
    data["R13I1"] = MolecularProperty(0.1959104, 3952566.072, 396.44, 0.000225703065, 0.1761811178)
    data["R14"] = MolecularProperty(0.0880046, 3750000, 227.51, 0.0001406584622, 0.1785)
    data["R141b"] = MolecularProperty(0.11694962, 4212000, 477.5, 0.0002550369804, 0.2195)
    data["R142b"] = MolecularProperty(0.10049503, 4055000, 410.26, 0.0002253267237, 0.2321)
    data["R143a"] = MolecularProperty(0.084041, 3761000, 345.857, 0.0001949906892, 0.2614896462)
    data["R152A"] = MolecularProperty(0.066051, 4520000, 386.411, 0.000179486413, 0.2752171145)
    data["R161"] = MolecularProperty(0.0480595, 5010000, 375.25, 0.0001592356688, 0.2162428411)
    data["R21"] = MolecularProperty(0.1029227, 5181200, 451.48, 0.0001956654009, 0.2061)
    data["R218"] = MolecularProperty(0.18801933, 2640000, 345.02, 0.0002994011976, 0.3172)
    data["R22"] = MolecularProperty(0.086468, 4990000, 369.295, 0.0001650649861, 0.22082)
    data["R227EA"] = MolecularProperty(0.17002886, 2925242.234, 374.9, 0.0002861230329, 0.357641242)
    data["R23"] = MolecularProperty(0.07001385, 4832000, 299.293, 0.0001329787234, 0.2629648925)
    data["R236EA"] = MolecularProperty(0.1520384, 3420000, 412.44, 0.0002691065662, 0.3687823804)
    data["R236FA"] = MolecularProperty(0.1520384, 3200000, 398.07, 0.0002757859901, 0.3769117976)
    data["R245ca"] = MolecularProperty(0.13404794, 3940740.489, 447.57, 0.0002551020408, 0.3545654191)
    data["R245fa"] = MolecularProperty(0.13404794, 3651000, 427.01, 0.0002580645161, 0.3776)
    data["R32"] = MolecularProperty(0.052024, 5782000, 351.255, 0.0001226981129, 0.2769)
    data["R365MFC"] = MolecularProperty(0.14807452, 3266207.068, 460, 0.0003125, 0.3774450972)
    data["R40"] = MolecularProperty(0.05048752, 6671686.977, 416.3, 0.0001390002242, 0.1500684964)
    data["R404A"] = MolecularProperty(0.0976038, 3734800, 345.27, 0.0002024291498, 0.293)
    data["R407C"] = MolecularProperty(0.0862036, 4631700, 359.345, 0.0001901140668, 0.363)
    data["R41"] = MolecularProperty(0.03403292, 5897000, 317.28, 0.0001075268817, 0.2004)
    data["R410A"] = MolecularProperty(0.0725854, 4901200, 344.494, 0.0001581277672, 0.296)
    return data
end

"""
    cubic_benchmark(name)

Get a benchmark equation-of-state instance together with some data for testing.
"""
function cubic_benchmark(name)
    psi = 6.894757293168360e+03
    Rankine = 0.555555555555556
    bar = 1e5
    data = Dict()
    eos_type = PengRobinson()

    bic = nothing
    if name == "spe5"
        T_c = [343, 665.7, 913.4, 1111.8, 1270.0, 1380.0]*Rankine;
        p_c = [667.8, 616.3, 436.9, 304.0, 200.0, 162.0]*psi;
        mw = [16.040, 44.100, 86.180, 142.290, 206.000, 282.000]/1000;
        acc = [0.0130, 0.1524, 0.3007, 0.4885, 0.6500, 0.8500];
        Z_c = [0.290, 0.277, 0.264, 0.257, 0.245, 0.235];
        # Convert critical volumes to critical densities
        V_c = Z_c.*8.314.*T_c./p_c;
        names = ["C1", "C3", "C6", "C10", "C15", "C20"];
        ncomp = length(names);
        bic = zeros(ncomp, ncomp);
        bic[1, 5] = 0.05;
        bic[1, 6] = 0.05;
        bic[2, 5] = 0.005;
        bic[2, 6] = 0.005;
        bic = Symmetric(bic)
        props = MolecularProperty.(mw, p_c, T_c, V_c, acc)

        z0 = [0.77, 0.20, 0.03, 0, 0, 0]
        z = [0.5, 0.03, 0.07, 0.20, 0.15, 0.05]
        p = 4000*psi
        T = 344.26
        p0 = p
        T0 = T

        initial_cond = (z = z0, p = p0, T = T0)
        displ_cond = (z = z, p = p, T = T)
        data["initial_condition"] = initial_cond
        data["injection_condition"] = displ_cond
    elseif name == "simple"
        names = ["Methane", "CarbonDioxide", "n-Decane"]
        props = MolecularProperty.(names)

        p = 75*bar
        T = 273.15 + 150
        z0 = [0.3, 0.1, 0.6]
        z = [0.1, 0.9, 0.0]
        initial_cond = (z = z0, p = p, T = T)
        displ_cond = (z = z, p = p, T = T)
        data["initial_condition"] = initial_cond
        data["injection_condition"] = displ_cond
    else
        error("Unknown benchmark $name")
    end
    mixture = MultiComponentMixture(props, names = names, A_ij = bic)
    eos = GenericCubicEOS(mixture, eos_type)

    return (eos, data)
end


@time using CSV, XLSX, DataFrames, StatsBase, Statistics
@time using Plots, Plots.PlotMeasures, LaTeXStrings, StatsPlots
default()

over50(ndwi) = vec(maximum(ndwi, dims = 1) .> 50)

schedules = DataFrame(XLSX.readtable("C:/Users/rmsms/OneDrive/lab/aespa//schedule.xlsx", "schedule")...)

ntwk = CSV.read("D:/simulated/ntwk.csv", DataFrame)

import_dir = "D:/simulated/"
export_dir = "C:/Users/rmsms/OneDrive/lab/aespa/png/"

color_ = Dict("00" => :black, "0V" => colorant"#C00000", "M0" => colorant"#0070C0", "MV" => colorant"#7030A0")
shape_ = Dict("00" => :+, "0V" => :x, "M0" => :circle, "MV" => :star5)
label_ = Dict("00" => "No control", "0V" => "Vaccination", "M0" => "Restriction", "MV" => "Both control")

te_color_ = Dict(
    "S" => colorant"#FFC000",
    "E" => colorant"#ED7D31",
    "I" => colorant"#C00000",
    "R" => colorant"#A5A5A5",
    "V" => colorant"#70AD47")


function Base.iterate(
    layer::AbstractFeatureLayer,
    state::Integer = 0,
)::Union{Nothing,Tuple{Feature,Int64}}
    layer.ptr == C_NULL && return nothing
    state == 0 && resetreading!(layer)
    ptr = GDAL.ogr_l_getnextfeature(layer.ptr)
    return if ptr == C_NULL
        resetreading!(layer)
        nothing
    else
        (Feature(ptr), state + 1)
    end
end

Base.eltype(layer::AbstractFeatureLayer)::DataType = Feature

Base.IteratorSize(::Type{<:AbstractFeatureLayer}) = Base.SizeUnknown()

Base.length(layer::AbstractFeatureLayer)::Integer = nfeature(layer, true)

struct BlockIterator{T<:Integer}
    rows::T
    cols::T
    ni::T
    nj::T
    n::T
    xbsize::T
    ybsize::T
end

function blocks(
    ::Type{T},
    raster::AbstractRasterBand,
)::BlockIterator{T} where {T<:Integer}
    (xbsize, ybsize) = blocksize(raster)
    rows = T(height(raster))
    cols = T(width(raster))
    ni = ceil(T, rows / ybsize)
    nj = ceil(T, cols / xbsize)
    return BlockIterator{T}(rows, cols, ni, nj, ni * nj, xbsize, ybsize)
end

function blocks(raster::AbstractRasterBand)::BlockIterator{Int64}
    return blocks(Int64, raster)
end

function Base.iterate(
    obj::BlockIterator{T},
    iter::T = 0,
)::Union{Nothing,Tuple{Tuple{Tuple{T,T},Tuple{T,T}},T}} where {T}
    iter == obj.n && return nothing
    j = floor(Int, iter / obj.ni)
    i = iter % obj.ni
    nrows = if (i + 1) * obj.ybsize < obj.rows
        obj.ybsize
    else
        obj.rows - i * obj.ybsize
    end
    ncols = if (j + 1) * obj.xbsize < obj.cols
        obj.xbsize
    else
        obj.cols - j * obj.xbsize
    end
    return (((i, j), (nrows, ncols)), iter + 1)
end

struct WindowIterator{T<:Integer}
    blockiter::BlockIterator{T}
end
Base.size(i::WindowIterator) = (i.blockiter.ni, i.blockiter.nj)
Base.length(i::WindowIterator) = i.blockiter.n

function Base.IteratorSize(::Type{WindowIterator{T}}) where {T<:Integer}
    return Base.HasShape{2}()
end

function Base.IteratorEltype(::Type{WindowIterator{T}}) where {T<:Integer}
    return Base.HasEltype()
end

function Base.eltype(::WindowIterator{T})::DataType where {T<:Integer}
    return Tuple{UnitRange{T},UnitRange{T}}
end

function windows(
    ::Type{T},
    raster::AbstractRasterBand,
)::WindowIterator{T} where {T<:Integer}
    return WindowIterator{T}(blocks(T, raster))
end

windows(raster::AbstractRasterBand)::WindowIterator{Int64} =
    windows(Int64, raster)

function Base.iterate(
    obj::WindowIterator{T},
    iter::T = 0,
)::Union{Nothing,Tuple{NTuple{2,UnitRange{T}},T}} where {T<:Integer}
    handle = obj.blockiter
    next = Base.iterate(handle, iter)
    next == nothing && return nothing
    (((i, j), (nrows, ncols)), iter) = next
    return (
        ((1:ncols) .+ j * handle.xbsize, (1:nrows) .+ i * handle.ybsize),
        iter,
    )
end

mutable struct BufferIterator{R<:Real,T<:Integer}
    raster::AbstractRasterBand
    w::WindowIterator{T}
    buffer::Matrix{R}
end

function bufferwindows(
    ::Type{T},
    raster::AbstractRasterBand,
)::BufferIterator{pixeltype(raster),T} where {T<:Integer}
    return BufferIterator{pixeltype(raster),T}(
        raster,
        windows(T, raster),
        Matrix{pixeltype(raster)}(undef, blocksize(raster)...),
    )
end

function bufferwindows(
    raster::AbstractRasterBand,
)::BufferIterator{pixeltype(raster),Int64}
    return bufferwindows(Int64, raster)
end

function Base.iterate(
    obj::BufferIterator{R,T},
    iter::T = 0,
)::Union{Nothing,Tuple{Matrix{R},T}} where {R<:Real,T<:Integer}
    next = Base.iterate(obj.w, iter)
    next == nothing && return nothing
    ((cols, rows), iter) = next
    rasterio!(obj.raster, obj.buffer, rows, cols)
    return (obj.buffer[1:length(cols), 1:length(rows)], iter)
end


function overpass_polygon_buildings_query(geojson_polygons::Vector{Vector{Any}},
                                          metadata::Bool=false,
                                          download_format::Symbol=:osm
                                          )::String
    filters = ""
    for polygon in geojson_polygons
        polygon = map(x -> [x[2], x[1]], polygon) # switch lon-lat to lat-lon
        polygon_str = replace("$polygon", r"[\[,\]]" =>  "")
        filters *= """node["building"](poly:"$polygon_str");<;way["building"](poly:"$polygon_str");>;rel["building"](poly:"$polygon_str");>;"""
    end

    return overpass_query(filters, metadata, download_format)
end

"""
    osm_buildings_from_place_name(;place_name::String,
                                metadata::Bool=false,
                                download_format::Symbol=:osm
                                )::String

Downloads OpenStreetMap buildings using any place name string.

# Arguments
- `place_name::String`: Any place name string used as a search argument to the Nominatim API.
- `metadata::Bool=false`: Set true to return metadata.
- `download_format::Symbol=:osm`: Download format, either `:osm`, `:xml` or `json`.

# Return
- `String`: OpenStreetMap buildings data response string.
"""
function osm_buildings_from_place_name(;place_name::String,
                                       metadata::Bool=false,
                                       download_format::Symbol=:osm
                                       )::String
    geojson_polygons = polygon_from_place_name(place_name)
    query = overpass_polygon_buildings_query(geojson_polygons, metadata, download_format)
    return overpass_request(query)
end

"""
    osm_buildings_from_bbox(;minlat::AbstractFloat,
                          minlon::AbstractFloat,
                          maxlat::AbstractFloat,
                          maxlon::AbstractFloat,
                          metadata::Bool=false,
                          download_format::Symbol=:osm
                          )::String

Downloads OpenStreetMap buildings using bounding box coordinates.

# Arguments
- `minlat::AbstractFloat`: Bottom left bounding box latitude coordinate.
- `minlon::AbstractFloat`: Bottom left bounding box longitude coordinate.
- `maxlat::AbstractFloat`: Top right bounding box latitude coordinate.
- `maxlon::AbstractFloat`: Top right bounding box longitude coordinate.
- `metadata::Bool=false`: Set true to return metadata.
- `download_format::Symbol=:osm`: Download format, either `:osm`, `:xml` or `json`.

# Return
- `String`: OpenStreetMap buildings data response string.
"""
function osm_buildings_from_bbox(;minlat::Float64,
                                 minlon::Float64,
                                 maxlat::Float64,
                                 maxlon::Float64,
                                 metadata::Bool=false,
                                 download_format::Symbol=:osm
                                 )::String
    filters = """node["building"];<;way["building"];>;rel["building"];>;"""
    bbox = [minlat, minlon, maxlat, maxlon]
    query = overpass_query(filters, metadata, download_format, bbox)
    return overpass_request(query)
end

"""
    osm_buildings_from_point(;point::GeoLocation,
                           radius::Number,
                           metadata::Bool=false,
                           download_format::Symbol=:osm
                           )::String

Downloads OpenStreetMap buildings using bounding box coordinates calculated from a centroid point and radius (km).

# Arguments
- `point::GeoLocation`: Centroid point to draw the bounding box around.
- `radius::Number`: Distance (km) from centroid point to each bounding box corner.
- `metadata::Bool=false`: Set true to return metadata.
- `download_format::Symbol=:osm`: Download format, either `:osm`, `:xml` or `json`.

# Return
- `String`: OpenStreetMap buildings data response string.
"""
function osm_buildings_from_point(;point::GeoLocation,
                                  radius::Number,
                                  metadata::Bool=false,
                                  download_format::Symbol=:osm
                                  )::String
    bbox = bounding_box_from_point(point, radius)
    return osm_buildings_from_bbox(;bbox..., metadata=metadata, download_format=download_format)
end

function osm_buildings_downloader(download_method::Symbol)::Function
    if download_method == :place_name
        return osm_buildings_from_place_name
    elseif download_method == :bbox
        return osm_buildings_from_bbox
    elseif download_method == :point
        return osm_buildings_from_point
    else
        throw(ErrorException("OSM buildings downloader $download_method does not exist"))
    end
end

"""
    download_osm_buildings(download_method::Symbol;
                           metadata::Bool=false,
                           download_format::Symbol=:osm,
                           save_to_file_location::Union{String,Nothing}=nothing,
                           download_kwargs...
                           )::Union{XMLDocument,Dict{String,Any}}

Downloads OpenStreetMap buildings data by querying with a place name, bounding box, or centroid point.

# Arguments
- `download_method::Symbol`: Download method, choose from `:place_name`, `:bbox` or `:point`.
- `metadata::Bool=false`: Set true to return metadata.
- `download_format::Symbol=:osm`: Download format, either `:osm`, `:xml` or `json`.
- `save_to_file_location::Union{String,Nothing}=nothing`: Specify a file location to save downloaded data to disk.

# Required Download Kwargs

*`download_method=:place_name`*
- `place_name::String`: Any place name string used as a search argument to the Nominatim API.

*`download_method=:bbox`*
- `minlat::AbstractFloat`: Bottom left bounding box latitude coordinate.
- `minlon::AbstractFloat`: Bottom left bounding box longitude coordinate.
- `maxlat::AbstractFloat`: Top right bounding box latitude coordinate.
- `maxlon::AbstractFloat`: Top right bounding box longitude coordinate.

*`download_method=:point`*
- `point::GeoLocation`: Centroid point to draw the bounding box around.
- `radius::Number`: Distance (km) from centroid point to each bounding box corner.

# Return
- `Union{XMLDocument,Dict{String,Any}}`: OpenStreetMap buildings data parsed as either XML or Dictionary object depending on the download method.
"""
function download_osm_buildings(download_method::Symbol;
                                metadata::Bool=false,
                                download_format::Symbol=:osm,
                                save_to_file_location::Union{String,Nothing}=nothing,
                                download_kwargs...
                                )::Union{XMLDocument,Dict{String,Any}}
    downloader = osm_buildings_downloader(download_method)
    data = downloader(metadata=metadata, download_format=download_format; download_kwargs...)
    @info "Downloaded osm buildings data from $(["$k: $v" for (k, v) in download_kwargs]) in $download_format format"

    if !(save_to_file_location isa Nothing)
        save_to_file_location = validate_save_location(save_to_file_location, download_format)
        write(save_to_file_location, data)
        @info "Saved osm buildings data to disk: $save_to_file_location"
    end

    deserializer = string_deserializer(download_format)

    return deserializer(data)
end

is_building(tags::Dict)::Bool = haskey(tags, "building") ? true : false

function height(tags::Dict)::Number
    height = get(tags, "height", nothing)
    levels = get(tags, "building:levels", nothing) !== nothing ? tags["building:levels"] : get(tags, "level", nothing)

    if height !== nothing
        return height isa String ? max([remove_non_numeric(h) for h in split(height, r"[+^;,-]")]...) : height
    elseif levels !== nothing
        levels = levels isa String ? round(max([remove_non_numeric(l) for l in split(levels, r"[+^;,-]")]...)) : levels
        levels = levels == 0 ? rand(1:DEFAULT_MAX_BUILDING_LEVELS) : levels
    else
        levels = rand(1:DEFAULT_MAX_BUILDING_LEVELS)
    end

    return levels * DEFAULT_BUILDING_HEIGHT_PER_LEVEL
end

function parse_osm_buildings_dict(osm_buildings_dict::AbstractDict)::Dict{Integer,Building}
    T = DEFAULT_OSM_ID_TYPE

    nodes = Dict{T,Node{T}}()
    for node in osm_buildings_dict["node"]
        id = node["id"]
        nodes[id] = Node{T}(
            id,
            GeoLocation(node["lat"], node["lon"]),
            haskey(node, "tags") ? node["tags"] : nothing
        )
    end

    ways = Dict(way["id"] => way for way in osm_buildings_dict["way"]) # for lookup

    added_ways = Set{T}()
    buildings = Dict{T,Building{T}}()
    for relation in osm_buildings_dict["relation"]
        is_relation = true

        if haskey(relation, "tags") && is_building(relation["tags"])
            tags = relation["tags"]
            rel_id = relation["id"]
            members = relation["members"]

            polygons = Vector{Polygon{T}}()
            for member in members
                way_id = member["ref"]
                way = ways[way_id]

                haskey(way, "tags") && merge!(tags, way["tags"]) # could potentially overwrite some data
                push!(added_ways, way_id)

                is_outer = member["role"] == "outer" ? true : false
                nds = [nodes[n] for n in way["nodes"]]
                push!(polygons, Polygon(way_id, nds, is_outer))
            end

            tags["height"] = height(tags)
            sort!(polygons, by = x -> x.is_outer, rev=true) # sorting so outer polygon is always first
            buildings[rel_id] = Building{T}(rel_id, is_relation, polygons, tags)
        end
    end

    for (way_id, way) in ways
        is_relation = false
        is_outer = true

        if haskey(way, "tags") && is_building(way["tags"]) && !(way_id in added_ways)
            tags = way["tags"]
            tags["height"] = height(tags)
            nds = [nodes[n] for n in way["nodes"]]
            polygons = [Polygon(way_id, nds, is_outer)]
            buildings[way_id] = Building{T}(way_id, is_relation, polygons, tags)
        end
    end

    return buildings
end

"""
    buildings_from_object(buildings_xml_object::XMLDocument)::Dict{Integer,Building}

Creates `Building` objects from data downloaded with `download_osm_buildings`.

# Arguments
- `buildings_xml_object::XMLDocument`: Buildings data downloaded in `:xml` or `:osm` format.

# Return
- `Dict{Integer,Building}`: Mapping from building relation/way ids to `Building` objects.
"""
function buildings_from_object(buildings_xml_object::XMLDocument)::Dict{Integer,Building}
    dict_to_parse = osm_dict_from_xml(buildings_xml_object)
    return parse_osm_buildings_dict(dict_to_parse)
end

"""
    buildings_from_object(buildings_json_object::AbstractDict)::Dict{Integer,Building}

Creates `Building` objects from data downloaded with `download_osm_buildings`.

# Arguments
- `buildings_json_object::AbstractDict`: Buildings data downloaded in `:json` format.

# Return
- `Dict{Integer,Building}`: Mapping from building relation/way ids to `Building` objects.
"""
function buildings_from_object(buildings_json_object::AbstractDict)::Dict{Integer,Building}
    dict_to_parse = osm_dict_from_json(buildings_json_object)
    return parse_osm_buildings_dict(dict_to_parse)
end

"""
    buildings_from_file(file_path::String)::Dict{Integer,Building}

Creates `Building` objects from OpenStreetMap data file (either `:osm`, `:xml` or `:json` format).

# Arguments
- `file_path::String`: Path to OpenStreetMap data file.

# Return
- `Dict{Integer,Building}`: Mapping from building relation/way ids to `Building` objects.
"""
function buildings_from_file(file_path::String)::Dict{Integer,Building}
    !isfile(file_path) && throw(ArgumentError("File $file_path does not exist"))
    deserializer = file_deserializer(file_path)
    obj = deserializer(file_path)
    return buildings_from_object(obj)
end

"""
    buildings_from_download(download_method::Symbol;
                            metadata::Bool=false,
                            download_format::Symbol=:osm,
                            save_to_file_location::Union{String,Nothing}=nothing,
                            download_kwargs...
                            )::Dict{Integer,Building}

Downloads and Creates `Building` objects from OpenStreetMap APIs.

# Arguments
- `download_method::Symbol`: Download method, choose from `:place_name`, `:bbox` or `:point`.
- `metadata::Bool=false`: Set true to return metadata.
- `download_format::Symbol=:osm`: Download format, either `:osm`, `:xml` or `json`.
- `save_to_file_location::Union{String,Nothing}=nothing`: Specify a file location to save downloaded data to disk.

# Required Download Kwargs

*`download_method=:place_name`*
- `place_name::String`: Any place name string used as a search argument to the Nominatim API.

*`download_method=:bbox`*
- `minlat::AbstractFloat`: Bottom left bounding box latitude coordinate.
- `minlon::AbstractFloat`: Bottom left bounding box longitude coordinate.
- `maxlat::AbstractFloat`: Top right bounding box latitude coordinate.
- `maxlon::AbstractFloat`: Top right bounding box longitude coordinate.

*`download_method=:point`*
- `point::GeoLocation`: Centroid point to draw the bounding box around.
- `radius::Number`: Distance (km) from centroid point to each bounding box corner.

# Return
- `Dict{Integer,Building}`: Mapping from building relation/way ids to `Building` objects.
"""
function buildings_from_download(download_method::Symbol;
                                 metadata::Bool=false,
                                 download_format::Symbol=:osm,
                                 save_to_file_location::Union{String,Nothing}=nothing,
                                 download_kwargs...
                                 )::Dict{Integer,Building}
    obj = download_osm_buildings(download_method,
                                 metadata=metadata,
                                 download_format=download_format,
                                 save_to_file_location=save_to_file_location;
                                 download_kwargs...)
    return buildings_from_object(obj)
end


#################### Binary Metropolised Gibbs Sampler ####################

#################### Types and Constructors ####################

const BMGForm = Union{Int, Vector{Vector{Int}}}

mutable struct BMGTune{F<:BMGForm} <: SamplerTune
  logf::Union{Function, Missing}
  k::F

  BMGTune{F}() where {F<:BMGForm} = new{F}()

  BMGTune{F}(x::Vector, k::F) where {F<:BMGForm} = new{F}(missing, k)

  BMGTune{F}(x::Vector, k::F, logf::Function) where {F<:BMGForm} =
    new{F}(logf, k)
end

const BMGIntVariate = SamplerVariate{BMGTune{Int}}
const BMGVecVariate = SamplerVariate{BMGTune{Vector{Vector{Int}}}}

BMGVariate(x::Vector, logf::Function; k::F=1) where {F<:BMGForm} =
  SamplerVariate{BMGTune{F}}(x, k, logf)

function validate(v::BMGIntVariate)
  n = length(v)
  v.tune.k <= n || throw(ArgumentError("k exceeds variate length $n"))
  validatebinary(v)
end

function validate(v::BMGVecVariate)
  n = length(v)
  mapreduce(maximum, max, v.tune.k) <= n ||
    throw(ArgumentError("indices in k exceed variate length $n"))
  validatebinary(v)
end

#################### Sampler Constructor ####################
"""
    BMG(params::ElementOrVector{Symbol}; k::F=1) where {F<:BMGForm}

Construct a `Sampler` object for BMG sampling. Parameters are assumed to have
 binary numerical values (0 or 1).

Returns a `Sampler{BMGTune{typeof(k)}}` type object.

* `params`: stochastic node(s) to be updated with the sampler.

* `k`:  number of parameters or vector of parameter indices to select at random for simultaneous updating in each call of the sampler.
"""
function BMG(params::ElementOrVector{Symbol}; k::F=1) where {F<:BMGForm}
  samplerfx = function(model::Model, block::Integer)
    block = SamplingBlock(model, block)
    v = SamplerVariate(block, k)
    sample!(v, x -> logpdf!(block, x))
    relist(block, v)
  end
  Sampler(params, samplerfx, BMGTune{F}())
end

#################### Sampling Functions ####################

sample!(v::SamplerVariate{BMGTune{F}}) where {F<:BMGForm} = sample!(v, v.tune.logf)
"""
    sample!(v::SamplerVariate{BMGTune{F}}, logf::Function) where {F<:BMGForm}

Draw one sample from a target distribution using the BMG sampler. Parameters
are assumed to have binary numerical values (0 or 1).

Returns `v` updated with simulated values and associated tuning parameters.
"""
function sample!(v::SamplerVariate{BMGTune{F}}, logf::Function) where {F<:BMGForm}
  n = length(v)
  probs = Vector{Float64}(undef, n)
  idx = randind(v)

  pbernoulli! = function(x, probs)
    for i in idx
      x_i = x[i]

      x[i] = 0.0
      logf0 = logf(x)
      x[i] = 1.0
      logf1 = logf(x)

      p = invlogit(logf1 - logf0)
      probs[i] = 0.0 < p < 1.0 ? p : 0.5

      x[i] = x_i
    end
    probs
  end

  x = v[:]
  pbernoulli!(x, probs)
  theta = map(p -> rand() < p, probs[idx])

  if n == 1
    v[idx] = theta
  else
    y = v[:]
    y[idx] = theta
    qy = mapreduce(i -> y[i] == 1.0 ? log(probs[i]) : log(1.0 - probs[i]), +,
                   idx)

    pbernoulli!(y, probs)
    qx = mapreduce(i -> x[i] == 1.0 ? log(probs[i]) : log(1.0 - probs[i]), +,
                   idx)

    if rand() < exp((logf(y) - qy) - (logf(x) - qx))
      v[idx] = theta
    end
  end

  v
end

randind(v::BMGIntVariate) = sample(1:length(v), v.tune.k, replace=false)
randind(v::BMGVecVariate) = sample(v.tune.k)


module TestHCEstimator
using Test
using JuMP, Ipopt
using DataFrames, CSV

include("../src/HCEstimator.jl")
using .HCEstimator
include("../src/Estimator.jl")
using .Estimator

include("cases_no_v.jl")
using .Cases

function runtests()

    @testset "HCEstimator" begin
        for case in [case3_dist, case33, case3_dist_no_dgs]
            sys, name = case()
            @testset "$name" begin
                other = nl_pf(Model(), sys)
                other_w_obj = add_objective(other, sys)

                sot = build_model(Model(), sys)
                @test sprint(print, sot) == sprint(print, other_w_obj)

                other_sot = build_model(Model(), sys, false)
                @test sprint(print, sot) == sprint(print, other)
            end

        end
    end
end

end
TestHCEstimator.runtests()


# Adopted from https://discourse.julialang.org/t/get-cpu-usage/24468/2
function get_cpu_stats()
  st = read("/proc/stat", String)
  lines = split(st, "\n")
  m = map(line -> match(r"^cpu(\d) (.*)", line), lines)
  matches = filter(m -> m !== nothing, m)
  Dict(parse(Int, match.captures[1]) => parse.(Int, split(match.captures[2], ' ')) for match in matches)
end

cpu_percentages = Dict(i => (0, 0) for i in 0:3)
function compute_percentage(fields, cpu_idx)
  lastidle, lasttotal = cpu_percentages[cpu_idx]
  idle, total = fields[4], sum(fields)
  Δidle, Δtotal = idle - lastidle, total - lasttotal
  utilization = 100 * (1 - Δidle / Δtotal)
  cpu_percentages[cpu_idx] = idle, total
  utilization
end


using GraphMatFun, LsqFit, LinearAlgebra, Random
include("exp_reset_all.jl");
#include("simulationtools.jl");
include("simulationtools.jl");

f=exp
m=8
target_n=400 # Measure error wrt this discr
n=100;        # Optimize wrt this discr
rho=13.5;  # Increased SID rho m=7

its=5;
droptol0=1e-20;
opt_kwargs=Dict(:logger=> 0,:γ0 => 0.5,:droptol => droptol0,
                :linlsqr => :real_svd,:maxit => its,
                :errtype => :relerr)
base = init_state_mult(f,rho,n,m; eltype=Complex{BigFloat})
base.params[:n]=n;
base.params[:target_n]=target_n;
base.params[:opt_kwargs]=opt_kwargs;
base.params[:kickit_mode]=1;
base.params[:cref_count]=0;

oldrho="6_4";
oldrho2="3_59";
## First the optimization free
#filename="simulations/newgraphs/exp_m$(m)_sid_$oldrho.cgr";
#sid_org=init_state_file!(deepcopy(base),:sid,filename,showmeta=true)
#
##filename="simulations/newgraphs/exp_m$(m)_bbc_$oldrho.cgr";
##bbc_org=init_state_file!(deepcopy(base),:bbc,filename,showmeta=true)
#
#filename="simulations/newgraphs/exp_m$(m)_sastre_$oldrho.cgr";
#sastre_org=init_state_file!(deepcopy(base),:sastre,filename,showmeta=true)
#
#
##sastre_org.cref=get_degopt_crefs(sastre_org.graph);
#
#
#filename="simulations/newgraphs/exp_m$(m)_ps_$oldrho.cgr";
#ps_org=init_state_file!(deepcopy(base),:ps,filename,showmeta=true)
#
#
#filename="simulations/newgraphs/exp_m$(m)_mono_$oldrho.cgr";
#mono_org=init_state_file!(deepcopy(base),:mono,filename,showmeta=true)
#

#filename="simulations/newgraphs/exp_m$(m-1)_mono+GN_$oldrho.cgr";
##filename="simulations/graphs/exp_m8_mono_taylor_13_5.cgr";
#mono0=init_state_file!(deepcopy(base),:mono,filename,showmeta=true,scale_and_square=true)
#mono0.cref=mono0.cref[8:70]
#(mono,simlist,commandlist)=
#     run_sequence(mono0,"ksssssdddsssssddddssssssddddddsssssssdddsssssssssssdddssssGGsssddddggggggssssssssGGGsssGGGssssdddddggggggggsssssssGGGsGGGGGssssdddddggggggggssssssssGGGsGGGGGsssssddddgggggggsssGGGGGGGssssssdddggggggggggssssssggssssssGsGsGsGsGGsGGsGGGsGGsss");
#

datadir=joinpath("..","..","data","exp");
#filename="simulations/graphs/exp_m8_mono_taylor_13_5.cgr";
filename=joinpath(datadir,"exp_mono_m$(m-1)_opt_rho$(oldrho).cgr");
mono0=init_state_file!(deepcopy(base),:mono,filename,showmeta=true,
                       scale_and_square=true)

degopt_crefs=get_degopt_crefs(mono0.graph);
#mono0.cref=mono0.cref[8:end-8]
mono0.cref=[#get_degopt_crefs(mono0.graph)[1][end-1][1];
            #get_degopt_crefs(mono0.graph)[1][end-1][2];
            degopt_crefs[1][end][1];
            degopt_crefs[1][end][2];
            degopt_crefs[2] ];
(mono,simlist,commandlist)=
run_sequence(deepcopy(mono0),"sSSNSkSSGGGSS22SdddddSSSSSzS0pSZpSSSSzSzq");

filename=joinpath(datadir,"exp_ps_m$(m-1)_opt_rho$(oldrho).cgr");
ps0=init_state_file!(deepcopy(base),:ps,filename,showmeta=true,
                         scale_and_square=true)
ps0.cref=[degopt_crefs[1][end][1];
            degopt_crefs[1][end][2];
            degopt_crefs[2]];

(ps,simlist,commandlist)=
     run_sequence(ps0,
                  "SSNSSSSGGGSSSSSDDDSSSSZSSSZSSSddSddSzZSSSSddZSSSSZpSSSzq");

filename=joinpath(datadir,"exp_sastre_m$(m-1)_opt_rho$(oldrho).cgr");
sastre0=init_state_file!(deepcopy(base),:sastre,filename,showmeta=true,
                         scale_and_square=true)
sastre0.cref=[degopt_crefs[1][end][1];
            degopt_crefs[1][end][2];
            degopt_crefs[2]];

(sastre,simlist,commandlist)=
     run_sequence(sastre0,
                  "SSSSNSSSSkSSGGGSSddddSdSdSzSkZSSSSSSSZsSSzZSSSSpZpSSSSddSddddSkSSS2zq");

#filename=joinpath(datadir,"exp_sid_m$(m-1)_opt_rho$(oldrho).cgr");
#sid0=init_state_file!(deepcopy(base),:sid,filename,showmeta=true,
#                         scale_and_square=true)
#sid0.cref=[degopt_crefs[1][end][1];
#            degopt_crefs[1][end][2];
#            degopt_crefs[2]];
#
#(sid,simlist,commandlist)=
#     run_sequence(sid0,
#                  "SSSSNSSSSkSSGGGSSddddSdSdSzSkZSSSSSSSZsSSzZSSSSpZpSSSSddSddddSkSSS2zq");
#
#asd
#sSSNSkSSGGGSS22Sq");
##sSkSSSddddSSSSSSS22GGG22ddd2SS22n222tk22222222223tdddtSS333k333333333dddd3N333k3333333k3333333k3333333dddds333333333333333333333sSSStztkk3333N3nn3333n33NN333n33tz33NNN3333N3tS333ts333t33Skkk3333333tTz");
#
#graph1=mono1.graph;
##degopt=Degopt(mono1.graph);
###degopt.x[1][1] += sqrt(eps());;
##degopt.x[1][1][1] += 0.000001;
##normalize!(degopt);
##degopt.x[1][1][1] -= 0.000001;
##graph1=graph_degopt(degopt)[1];
#export_compgraph(graph1,"/tmp/mono1_m8_tmp.cgr");
#
#
#mono0.cref=[#get_degopt_crefs(mono0.graph)[1][end-1][1];
#            #get_degopt_crefs(mono0.graph)[1][end-1][2];
#            get_degopt_crefs(mono0.graph)[1][end][1];
#            get_degopt_crefs(mono0.graph)[1][end][2];
#            get_degopt_crefs(mono0.graph)[2] ];
#filename_tmp="/tmp/mono1_m8_tmp.cgr";
#mono2=init_state_file!(deepcopy(base),:mono,filename_tmp,showmeta=true);
#setdiff!(mono2.cref,mono0.cref);
#(mono,simlist,commandlist)=
#run_sequence(deepcopy(mono2),"dddddddddddddddddd22GGG22n2S");
#
#
#
#

#(mono,simlist,commandlist)=
#run_sequence(deepcopy(mono0),"ssssGGGssgddssssssddddggggggggggsssGGGsss");

#             "kssssGGGssgddssssssddddggggggggggsssGGGssssGGGssGGGGsssssksssdddsssssddssssssdddsdgggggggsssssGGGGsGGGssssssNssGssddssssddssssdggggggggsssssGGGGssssGGGGsssddsdsdgggggggggggssssGGGGGGGGGGsGssssggggggggdsdsdskssGGGGGGGGssNsszq");
#sGsGsGGGssssssssddsssGGGGGGsssksssssNsq
#             "sssGGGsssdddsddsssssdddsssdddsssssq");

#filename="simulations/newgraphs/exp_m$(m-1)_mono+GN_$(oldrho).cgr";
#mono0=init_state_file!(deepcopy(base),:mono,filename,showmeta=true,scale_and_square=true)
#(mono,simlist,commandlist)=
#     run_sequence(mono0,"ddddddddddddddddddddddddddsssGGGsssdddsddsssssdddsssdddsssssq");
#
#filename="simulations/newgraphs/exp_m$(m)_ps+GN_$oldrho.cgr";
#ps0=init_state_file!(deepcopy(base),:ps,filename,showmeta=true)
#(ps,simlist,commandlist)=
#     run_sequence(ps0,
#     "sq");
#
#filename="simulations/newgraphs/exp_m$(m-1)_sastre+GN_$oldrho.cgr";
#sastre0=init_state_file!(deepcopy(base),:sastre,filename,showmeta=true,scale_and_square=true)
#(sastre,simlist,commandlist)=
#     run_sequence(sastre0,
#                  "kssGGGssddddsdddddsdsssssdddssssddssssddsgggsssssGGGssq");
#
#filename="simulations/newgraphs/exp_m$(m)_sid+GN_$oldrho.cgr";
#sid0=init_state_file!(deepcopy(base),:sid,filename,showmeta=true)
#(sid,simlist,commandlist)=
#     run_sequence(sid_org,
#     "nssNdddddddddddddddddddssGGGsssddsddsddsddsd")
#
##filename="simulations/newgraphs/exp_m$(m)_bbc+GN_$oldrho.cgr";
##bbc0=init_state_file!(deepcopy(base),:bbc,filename,showmeta=true)
##(bbc,simlist,commandlist)=
##     run_sequence(bbc0,
##     "sdddddsdddddssssssddddsssssssddddsssssssssssssdddssssssssssddddssssssssdsggggggsdssssssssssssssssGGGGGsssGGGsssq")
##
#
include("exp_print_all.jl");
println("Run include(exp_save_all.jl) if you want to save");


struct StorePutKey{T} <: ResourceKey
  priority :: Int
  id :: UInt
  item :: T
  StorePutKey{T}(priority, id, item) where T = new(priority, id, item)
end

struct StoreGetKey <: ResourceKey
  priority :: Int
  id :: UInt
  filter :: Function
end

mutable struct Store{T} <: AbstractResource
  env :: Environment
  capacity :: UInt
  load :: UInt
  items :: Dict{T, UInt}
  seid :: UInt
  put_queue :: DataStructures.PriorityQueue{Put, StorePutKey{T}}
  get_queue :: DataStructures.PriorityQueue{Get, StoreGetKey}
  function Store{T}(env::Environment; capacity::UInt=typemax(UInt)) where {T}
    new(env, capacity, zero(UInt), Dict{T, UInt}(), zero(UInt), DataStructures.PriorityQueue{Put, StorePutKey{T}}(), DataStructures.PriorityQueue{Get, StoreGetKey}())
  end
end

function put(sto::Store{T}, item::T; priority::Int=0) where T
  put_ev = Put(sto.env)
  sto.put_queue[put_ev] = StorePutKey{T}(priority, sto.seid+=one(UInt), item)
  @callback trigger_get(put_ev, sto)
  trigger_put(put_ev, sto)
  put_ev
end

get_any_item(::T) where T = true

function get(sto::Store{T}, filter::Function=get_any_item; priority::Int=0) where T
  get_ev = Get(sto.env)
  sto.get_queue[get_ev] = StoreGetKey(priority, sto.seid+=one(UInt), filter)
  @callback trigger_put(get_ev, sto)
  trigger_get(get_ev, sto)
  get_ev
end

function do_put(sto::Store{T}, put_ev::Put, key::StorePutKey{T}) where {T}
  if sto.load < sto.capacity
    sto.load += one(UInt)
    sto.items[key.item] = get(sto.items, key.item, zero(UInt)) + one(UInt)
    schedule(put_ev)
  end
  false
end

function do_get(sto::Store{T}, get_ev::Get, key::StoreGetKey) where {T}
  for (item, number) in sto.items
    if key.filter(item)
      sto.load -= one(UInt)
      if number === one(UInt)
        delete!(sto.items, item)
      else
        sto.items[item] = number - one(UInt)
      end
      schedule(get_ev; value=item)
      break
    end
  end
  true
end


using ShaderToy
shadertoy("monster.frag")