"""
	DataCell{T,S,U}(x::T, y::S, f::U, layer::Int, tinfo::String, oinfo::String)

Constructs the basic object for holding data, the `DataCell`. The constructor is parametric so it can be
easily extended to new data types.

# Main `DataCell` constructors
  * `DataCell(x::AbstractArray, tinfo::String="")` creates an `unlabeled` object.
  * `DataCell(x::AbstractArray, y::AbstractVector, tinfo::String="")` creates a `labeled` object. Useful for classification.
  * `DataCell(x::AbstractArray, y::AbstractMatrix, tinfo::String="")` creates a `multi-labeled` object. Useful for classification/regression

For the creation of data cells, the `datacell` function is quite well suited.
"""
#############
# DataCells #
#############
struct DataCell{T,S,U} <: AbstractCell{T,S,U}
	x::T                                                                                # Data
    	y::S                                                                                # Labels/probabilities etc.
    	f::U                                                                                # Usually a function
	layer::Int										# Layer
    	tinfo::String										# Title or name
    	oinfo::String                                                                        	# Other information
end

DataCell(x::T where T<:AbstractArray, tinfo::String = "") = DataCell(x, nothing, nothing, 0, tinfo, oinfoglobal)
DataCell(x::T where T<:AbstractArray, y::S where S<:Void, tinfo::String = "") = DataCell(x, nothing, nothing, 0, tinfo, oinfoglobal)
DataCell(x::T where T<:AbstractArray, y::S where S<:AbstractArray{Void}, tinfo::String = "") = DataCell(x, nothing, nothing, 0, tinfo, oinfoglobal)
DataCell(x::T where T<:AbstractArray, y::S where S<:AbstractArray, tinfo::String = "") = DataCell(x, checktargets(x,y), nothing, 0, tinfo, oinfoglobal)

####################
# DataCell aliases #
####################
const DataElement{T<:Union{Number, AbstractString, Symbol}} = Union{T,Nullable{T}}
const FixedElement{T<:Union{Number, AbstractString, Symbol}} = T
const NullableElement{T<:Union{Number, AbstractString, Symbol}} = Nullable{T}
const CellDataVec{T<:AbstractVector} = DataCell{T, Void, Void}
const CellDataVecVec{T<:AbstractVector, S<:AbstractVector} = DataCell{T, S, Void}
const CellDataVecMat{T<:AbstractVector, S<:AbstractMatrix} = DataCell{T, S, Void}
const CellDataMat{T<:AbstractMatrix} = DataCell{T, Void, Void}
const CellDataMatVec{T<:AbstractMatrix, S<:AbstractVector} =  DataCell{T, S, Void}
const CellDataMatMat{T<:AbstractMatrix, S<:AbstractMatrix} =  DataCell{T, S, Void}
const CellDataVectors{T<:AbstractVector, S<:Union{AbstractArray, Void}} = DataCell{T, S, Void}
const CellDataMatrices{T<:AbstractMatrix, S<: Union{AbstractArray, Void}} = DataCell{T, S, Void}
const CellDataU{T<:AbstractArray} = DataCell{T, Void, Void} 					# Data cell, unlabeled
const CellDataL{T<:AbstractArray, S<:AbstractVector} = DataCell{T, S, Void}			# Data cell labeled
const CellDataLL{T<:AbstractArray, S<:AbstractMatrix} = DataCell{T, S, Void}			# Data cell multiple labels
const CellData{T<:AbstractArray, S} = DataCell{T, S, Void}
#const TrainableData{U <:Real, T<:AbstractMatrix{U}, S<:AbstractArray} = DataCell{T, S, Void}

##############################################################################################################################
# Operators [Data cells]												     #
##############################################################################################################################

# Vertical datacell concatenation (e.g. data concatenation) and auxiliary methods
dcat(x::T where T<:PTuple{<:Void}) = nothing
dcat(::Void) = nothing
dcat(::Void,::Void) = nothing
dcat(x,::Void) = dcat(x)
dcat(::Void,x) = dcat(x)
dcat(x::AbstractArray) = getobs(x)
dcat(x::Array) = x
dcat(x::AbstractVector, y::AbstractVector) = vcat(mat(x, ObsDim.Constant{2}()), mat(y,ObsDim.Constant{2}()))
dcat(x::AbstractMatrix, y::AbstractVector) = vcat(dcat(x), mat(y, ObsDim.Constant{2}()))
dcat(x::AbstractVector, y::AbstractMatrix) = vcat(mat(x, ObsDim.Constant{2}()), dcat(y))
dcat(x::AbstractMatrix, y::AbstractMatrix) = vcat(dcat(x), dcat(y))
dcat(x,y,z...) = dcat(dcat(x,y),z...)
dcat(x) = dcat(x...)
vcat(c::T...) where T<:CellData = begin
	@assert all(labx == laby for labx in gety!.(c), laby in gety!.(c)) "[vcat] 'y' fields have to be identical for all DataCells."
	datacell(dcat(getx!.(c)), gety!(c[1]))
end

# Horizontal datacell concatenation (e.g. observation concatenation) and auxiliary methods
ocat(x::T where T<:PTuple{<:Void}) = nothing
ocat(::Void) = nothing
ocat(::Void,::Void) = nothing
ocat(x,::Void) = ocat(x)
ocat(::Void,x) = ocat(x)
ocat(x::AbstractArray) = getobs(x)
ocat(x::Array) = x
ocat(x::AbstractVector, y::AbstractVector) = vcat(ocat(x), ocat(y))
ocat(x::AbstractArray, y::AbstractArray) = hcat(ocat(x), ocat(y))
ocat(x,y,z...) = ocat(ocat(x,y),z...)
ocat(x) = ocat(x...)
hcat(c::T...) where T<:CellData = datacell(ocat(getx!.(c)), ocat(gety!.(c)))

# Piping operators
|>(x::T where T<:AbstractArray, c::S where S<:CellData) = begin
    	info("[operators] Creating new data cell with Array contents..." )
    	datacell(x)
end

|>(x::T where T<:Tuple, c::S where S<:CellData) = begin 					# Generally used in the case: (data, labels ) |> datacell([])
	if (nobs(c) > 0)
		info("[operators] Creating new data cell appending Tuple contents..." )
		datacell(dcat(x[1],getx!(c)), dcat(x[2], gety!(c)))
	else
		info("[operators] Creating new data cell with Tuple contents..." )
		datacell(x...)
	end
end

|>(x1::T where T<:CellDataU, x2::S where S<:CellDataU) = begin                 			# Generally used to glue together two DataCells
	@assert(nobs(x1) == nobs(x2))
	datacell(dcat(getx!(x1),getx!(x2)))
end

|>(x1::T where T<:CellDataU, x2::S where S<:CellData) = begin                 			# Generally used to glue together two DataCells
	@assert(nobs(x1) == nobs(x2))
	datacell(dcat(getx!(x1),getx!(x2)), gety!(x2))
end

|>(x1::T where T<:CellData, x2::S where S<:CellDataU) = begin                 			# Generally used to glue together two DataCells
	@assert(nobs(x1) == nobs(x2))
	datacell(dcat(getx!(x1),getx!(x2)), gety!(x1))
end

|>(x1::T where T<:CellData, x2::S where S<:CellData) = begin                 			# Generally used to glue together two DataCells
	@assert(nobs(x1) == nobs(x2))
	if isequal(gety!(x1), gety!(x2))
		datacell(dcat(getx!(x1),getx!(x2)), gety!(x1))
	else
		datacell(dcat(getx!(x1),getx!(x2)), dcat(gety!(x1),gety!(x2)))
	end
end

###########################
# Indexing for data cells #
###########################

# Generic setindex! for data cells
setindex!(c::T where T<:CellData, data, inds...) = setindex!(getx!(c), data, inds...)

# Generic indexing operator, no argument
getindex(c::CellData) = j4pr.datacell(eltype(getx!(c))[])
getindex(c::CellData, ::Colon) = c
getindex(c::CellData, ::Colon, ::Colon) = c

# Single integer indexing (select sample)
getindex(c::Union{CellDataVec, CellDataVecVec}, i::Int64) = DataCell([getx!(c)[i]], [gety!(c)[i]] )
getindex(c::CellDataVecMat, i::Int64) = DataCell([getx!(c)[i]], gety!(c)[:,i:i] )
getindex(c::Union{CellDataMat, CellDataMatVec}, i::Int64) = DataCell(getx!(c)[:,i:i], [gety!(c)[i]] )
getindex(c::CellDataMatMat, i::Int64) = DataCell(getx!(c)[:,i:i], gety!(c)[:,i:i] )

# Integer + Colon indexing (select variable) - only for matrix data
getindex(c::CellDataMatrices , i::Int64, ::Colon) = DataCell(getx!(c)[i:i,:], gety!(c) )

# Colon + Integer indexing (select sample) - only for matrix data
getindex(c::Union{CellDataMat, CellDataMatVec}, ::Colon, i::Int64) = DataCell(getx!(c)[:,i:i], [gety!(c)[i]])
getindex(c::CellDataMatMat, ::Colon, i::Int64) = DataCell(getx!(c)[:,i:i], gety!(c)[:,i:i])

# Range/Vector indexing (select samples) - only for vector data
getindex(c::Union{CellDataVec, CellDataVecVec}, i::Union{UnitRange{Int64}, Vector{Int64}}) = DataCell(getx!(c)[i], gety!(c)[i])
getindex(c::CellDataVecMat, i::Union{UnitRange{Int64}, Vector{Int64}}) = DataCell(getx!(c)[i], gety!(c)[:,i])

# Range/Vector + Colon indexing (select variables) - only for matrix data
getindex(c::Union{CellDataMat, CellDataMatVec}, i::Union{UnitRange{Int64}, Vector{Int64}}, ::Colon) = DataCell(getx!(c)[i,:], gety!(c)[:])
getindex(c::CellDataMatMat, i::Union{UnitRange{Int64}, Vector{Int64}}, ::Colon) = DataCell(getx!(c)[i,:], gety!(c)[:,:])

# Colon + Range/Vector indexing (select samples) - only for matrix data
getindex(c::Union{CellDataMat, CellDataMatVec}, ::Colon, i::Union{UnitRange{Int64}, Vector{Int64}}) = DataCell(getx!(c)[:,i], gety!(c)[i])
getindex(c::CellDataMatMat, ::Colon, i::Union{UnitRange{Int64}, Vector{Int64}}) = DataCell(getx!(c)[:,i], gety!(c)[:,i])

# Range/Vector + Range/Vector indexing (select samples and variables) - only for matrix data
getindex(c::Union{CellDataMat, CellDataMatVec}, i::Union{UnitRange{Int64}, Vector{Int64}}, j::Union{UnitRange{Int64}, Vector{Int64}}) = DataCell(getx!(c)[i,j], gety!(c)[j])
getindex(c::CellDataMatMat, i::Union{UnitRange{Int64}, Vector{Int64}}, j::Union{UnitRange{Int64}, Vector{Int64}}) = DataCell(getx!(c)[i,j], gety!(c)[:,j])

# Getindex for void arguments (needed to get indexes for CellDataU objects)
getindex(c::T where T<:Void, args...) = nothing

######################################
# Iteration interface for data cells #
######################################
start(c::CellData) = 1
next(c::CellData, state) = (datasubset(c, state), state+1)
done(c::CellData, state) = state > nobs(c)
endof(c::CellData) = nobs(c)
eltype(c::CellData) = typeof(c)
length(c::CellData) = nobs(c)

"""
	datacell([x,y,name=""])

Calls the DataCell constructor for creating `data cells`. It is the easyest way of creating such
objects from `AbstractArrays`. Consult `methods(datacell)` for a full list of available methods.
If no arguments are provided, a `FunctionCell` that expects an input compatible with the `DataCell`
constructor is returned.
"""
# Tier 1 - low level
datacell(x::T where T<:DataElement; name = "") = DataCell([x], name)
datacell(x::T where T<:DataElement, y::S where S<:DataElement; name = "") = DataCell([x], [y], name)
datacell(x::T where T<:DataElement, y::S where S<:AbstractArray; name = "") = DataCell([x], y, name)
datacell(x::T where T<:AbstractArray, y::S where S<:DataElement; name = "") = DataCell(x, [y], name)
datacell(x::T where T<:AbstractArray; name = "") = DataCell(x, name)
datacell(x::T where T<:AbstractArray, ::Void; name = "") =	DataCell(x, name)
datacell(x::T where T<:AbstractArray, y::S where S<:AbstractArray; name = "") = DataCell(x, y, name)
datacell(t :: Tuple{T} where T<:AbstractArray; name = "") = DataCell(t[1], name)
datacell(t :: Tuple{T,S} where {T<:AbstractArray, S<:AbstractArray}; name = "") = DataCell(t[1], t[2], name)
datacell(t :: Tuple{T,S} where {T<:AbstractArray,S}; name = "") = DataCell(t[1], name)

# Tier 2 - j4pr
datacell(x::T where T<:CellData; name = x.tinfo) = DataCell(getx!(x), gety!(x), name)
datacell(x::T where T<:CellData, y::S where S<:CellData; name = "") = datacell(getx!(x), getx!(y), name)

# DataCell creator
datacell() = FunctionCell((args...)->j4pr.DataCell(args...), (), "DataCell creator")


"""
    KnetModel

    * tensors: A dictionary. Given the tensor name as a string, returns the actual tensor.

    * model_layers: returns the list of layers (the actual layers themselves)

    * model_inputs: returns the list of inputs (the names of the tensors)

    * model_outputs: returns the list of outputs (the names of the outputs)
"""

mutable struct KnetModel
    tensors #dictionary: string -> tensor
    model_layers
    model_inputs
    model_outputs
    graph
end

function KnetModel(g::KnetOnnx.Types.Graph)
    """
    KnetModel constructor.
    Uses the Graph.
    """
    model_layers = get_ModelLayers(g)
    tensors = TensorDict(model_layers, g)
    model_inputs = [i.name for i in g.input]
    model_outputs = [o.name for o in g.output]
    graph = g
    KnetModel(tensors, model_layers, model_inputs, model_outputs, graph)
end

function KnetModel(path::String)
    """
    KnetModel constructor #2
    Uses the path of the ONNX file.
    """
    KnetModel(ONNXtoGraph(path))
end

"""
    TensorDict
    Initializes KnetModel.tensors by putting the names in as keys, but the values are Nothing.
"""
#omitted weights, might need them later
function TensorDict(model_layers, g)
    tensors = Dict()
    for layer in model_layers
        for input in layer.inputs; tensors[input] = Nothing; end
        for input in layer.outputs; tensors[input] = Nothing; end
    end
    for t in keys(g.initializer)
        if t in keys(tensors); tensors[t] = g.initializer[t]; end
    end
    tensors
end

"""
    ModelLayer
        * inputs: a list of tensor names.
          These tensors will be used for forward calculation of the layer.
        * outputs: a list of tensor names.
          The outputs of the forward calculation will be saved to Model.tensors under these keys.
        * layer: a Knet Layer.
          If you are constructing your own ModelLayer make sure the number of inputs and outputs matches the functionality of the KnetLayer you are using.

"""
mutable struct ModelLayer
    inputs #list of strings
    layer # a KnetLayer
    outputs #list of strings
end

function ModelLayer(node, g)
    (args, layer, outputs) = convert(node, g)
    ModelLayer(args, layer, outputs)
end

function get_ModelLayers(g)
    ModelLayers = []
    for node in g.node; push!(ModelLayers, ModelLayer(node, g)); end
    return ModelLayers
end

# FORWARD CALCULATIONS
function forward(km::KnetModel, ml::ModelLayer)

        # GATHER INPUTS
    for input in ml.inputs
        if km.tensors[input] == Nothing; return "oops!"; end
    end

        # FORWARD PASS
        # if only one input is requried, pass the first element
        # if more than one input is required, pass all elements
        # simply check the length of requried inputs for the model
    inputs = (key-> km.tensors[key]).(ml.inputs)
    if length(inputs) == 1; out = ml.layer(inputs[1]);
        else; out = ml.layer(inputs...); end
        #else; out = ml.layer(inputs[1]); end

        # SAVE OUTPUTS
        # check if there are multiple outputs (rnn etc.) before saving them to model.tensors
    if length(ml.outputs) == 1; km.tensors[ml.outputs[1]] = out;
        else; for output in ml.outputs; km.tensors[output] = out; end; end
 end

function (m::KnetModel)(args...)

    # resets model.tensors (figure out a smarter/faster reset)
    m.tensors = TensorDict(m.model_layers, m.graph)

    # REGISTER X
    # dumb version

    for (i, arg) in enumerate(args)
        m.tensors[m.model_inputs[i]] = arg
    end

    # LOOP UNTIL ALL TENSORS ARE CALCULATED
    # do until all model.tensors are filled
    # iterate over all layers and call forward on that layer
    while Nothing in values(m.tensors)
        for layer in m.model_layers
            forward(m, layer)
        end
    end

    # Gather outputs
    outs = []
    #m.tensors[m.model_outputs...]
    # DUMB VERSION
    # could be multiple
    if length(m.model_outputs) == 1; outs = m.tensors[m.model_outputs[1]];
        else; for out in m.model_outputs; push!(outs, m.tensors[out]); end;  end

    # return outputs
    return outs

end

"""
    PrintModelTensors(models::KnetModel)
    Displays your model.tensor, showing the size of the tensors that are calculated.
    Might be useful for debugging your model.
"""
function PrintModelTensors(model::KnetModel)
    tensors = model.tensors
    for k in keys(tensors)
        if tensors[k] == Nothing; println(k, "\t=> ", "Nothing")
        else println(k, "\t=> ", size(tensors[k])); end
    end
end

#import Knet: Data

#(m::KnetModel)(x,y) = Knet.nll(m(x), y)

"""
    Calling KnetModel with a Knet.Data object computes the mean nll for
"""
#(m::KnetModel)(d::Data) = Knet.mean(m(x,y) for (x,y) in d)


# allowed values in misc: char, string, numbers, and arrays of same.
tos(t::Type) = round(UInt8, log2(sizeof(t)))
function typ2code(t::Type)
  n = 0xff
  if t == Char
    n = 0x00
  elseif t == String
    n = 0x01
  elseif t <: Unsigned
    n = 0x10 + tos(t)
  elseif t <: Signed
    n = 0x20 + tos(t)
  elseif t <: AbstractFloat
    n = 0x30 + tos(t) - 0x01
  elseif t <: Complex
    n = 0x40 + typ2code(real(t))
  elseif t <: Array
    n = 0x80 + typ2code(eltype(t))
  end
  return n
end
# Who needs "switch"...

function code2typ(c::UInt8)
  t::Type = Any
  if c >= 0x80
    t = Array{code2typ(c-0x80)}
  elseif c >= 0x40
    t = Complex{code2typ(c-0x40)}
  elseif c >= 0x30
    t = getindex((Float16, Float32, Float64), c-0x2f)
  elseif c >= 0x20
    t = getindex((Int8, Int16, Int32, Int64, Int128), c-0x1f)
  elseif c >= 0x10
    t = getindex((UInt8, UInt16, UInt32, UInt64, UInt128), c-0x0f)
  elseif c == 0x01
    t = String
  elseif c == 0x00
    t = Char
  else
    error("unknown type!")
  end
  return t
end

# SUPPORTED TYPES IN :MISC
#
#= HOW TO CHECK:
(1) copy the two lines below to the command line to generate a table like the one below
(2) if anything in column 4 of your table is false, these functions are broken

using SeisIO, SeisIO.RandSeis, BenchmarkTools, LightXML; import SeisIO:code2typ, typ2code;
for c = 0x00:0xff; try; println(stdout, rpad(string(code2typ(c)), 36), "| ", repr(typ2code(code2typ(c)))," | ", repr(c), " | ", isequal(c, typ2code(code2typ(c)))); catch; end; end

GUIDE TO THE TABLE:
Column 1 is a list of types allowed in :misc
Column 2 is the corresponding UInt8 type codes
Column 3 is the value returned by typ2code(code2typ(c))
Column 4 is the result of c == typ2code(code2typ(c))

Type                                | Code | Ret  | ==?
:---------------                    |:-----|:-----|-----
Char                                | 0x00 | 0x00 | true
String                              | 0x01 | 0x01 | true
UInt8                               | 0x10 | 0x10 | true
UInt16                              | 0x11 | 0x11 | true
UInt32                              | 0x12 | 0x12 | true
UInt64                              | 0x13 | 0x13 | true
UInt128                             | 0x14 | 0x14 | true
Int8                                | 0x20 | 0x20 | true
Int16                               | 0x21 | 0x21 | true
Int32                               | 0x22 | 0x22 | true
Int64                               | 0x23 | 0x23 | true
Int128                              | 0x24 | 0x24 | true
Float16                             | 0x30 | 0x30 | true
Float32                             | 0x31 | 0x31 | true
Float64                             | 0x32 | 0x32 | true
Complex{UInt8}                      | 0x50 | 0x50 | true
Complex{UInt16}                     | 0x51 | 0x51 | true
Complex{UInt32}                     | 0x52 | 0x52 | true
Complex{UInt64}                     | 0x53 | 0x53 | true
Complex{UInt128}                    | 0x54 | 0x54 | true
Complex{Int8}                       | 0x60 | 0x60 | true
Complex{Int16}                      | 0x61 | 0x61 | true
Complex{Int32}                      | 0x62 | 0x62 | true
Complex{Int64}                      | 0x63 | 0x63 | true
Complex{Int128}                     | 0x64 | 0x64 | true
Complex{Float16}                    | 0x70 | 0x70 | true
Complex{Float32}                    | 0x71 | 0x71 | true
Complex{Float64}                    | 0x72 | 0x72 | true
Array{Char,N} where N               | 0x80 | 0x80 | true
Array{String,N} where N             | 0x81 | 0x81 | true
Array{UInt8,N} where N              | 0x90 | 0x90 | true
Array{UInt16,N} where N             | 0x91 | 0x91 | true
Array{UInt32,N} where N             | 0x92 | 0x92 | true
Array{UInt64,N} where N             | 0x93 | 0x93 | true
Array{UInt128,N} where N            | 0x94 | 0x94 | true
Array{Int8,N} where N               | 0xa0 | 0xa0 | true
Array{Int16,N} where N              | 0xa1 | 0xa1 | true
Array{Int32,N} where N              | 0xa2 | 0xa2 | true
Array{Int64,N} where N              | 0xa3 | 0xa3 | true
Array{Int128,N} where N             | 0xa4 | 0xa4 | true
Array{Float16,N} where N            | 0xb0 | 0xb0 | true
Array{Float32,N} where N            | 0xb1 | 0xb1 | true
Array{Float64,N} where N            | 0xb2 | 0xb2 | true
Array{Complex{UInt8},N} where N     | 0xd0 | 0xd0 | true
Array{Complex{UInt16},N} where N    | 0xd1 | 0xd1 | true
Array{Complex{UInt32},N} where N    | 0xd2 | 0xd2 | true
Array{Complex{UInt64},N} where N    | 0xd3 | 0xd3 | true
Array{Complex{UInt128},N} where N   | 0xd4 | 0xd4 | true
Array{Complex{Int8},N} where N      | 0xe0 | 0xe0 | true
Array{Complex{Int16},N} where N     | 0xe1 | 0xe1 | true
Array{Complex{Int32},N} where N     | 0xe2 | 0xe2 | true
Array{Complex{Int64},N} where N     | 0xe3 | 0xe3 | true
Array{Complex{Int128},N} where N    | 0xe4 | 0xe4 | true
Array{Complex{Float16},N} where N   | 0xf0 | 0xf0 | true
Array{Complex{Float32},N} where N   | 0xf1 | 0xf1 | true
Array{Complex{Float64},N} where N   | 0xf2 | 0xf2 | true
=#


module NARemovers

using Random
using DataFrames

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

import ..AbsTypes: fit!, transform!
export fit!, transform!
export NARemover

"""
    NARemover(
      Dict(
        :name => "nadetect",
        :acceptance => 0.10 # tolerable NAs percentage
      )
    )

Removes columns with NAs greater than acceptance rate.
This assumes that it processes columns of features.
The output column should not be part of input to avoid
it being excluded if it fails the acceptance critera.

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

   function NARemover(args::Dict = Dict())
      default_args = Dict{Symbol,Any}(
         :name => "nadetect",
         :acceptance => 0.10
      )
      cargs=nested_dict_merge(default_args,args)
      cargs[:name] = cargs[:name]*"_"*randstring(3)
      new(cargs[:name],cargs)
   end
end

"""
NARemover(acceptance::Float64)

Helper function for NARemover.
"""
NARemover(acceptance::Float64) = NARemover(Dict(:acceptance => acceptance))

"""
    fit!(nad::NARemover,features::DataFrame,labels::Vector=[])

Checks and exit of df is empty

# Arguments
- `nad::NARemover`: custom type
- `features::DataFrame`: input
- `labels::Vector=[]`:
"""
function fit!(nad::NARemover, features::DataFrame, labels::Vector=[])
   (features == DataFrame()) && throw(ArgumentError("empty dataframe"))
end

"""
    transform!(nad::NARemover,nfeatures::DataFrame)

Removes columns with NAs greater than acceptance rate.

# Arguments
- `nad::NARemover`: custom type
- `nfeatures::DataFrame`: input
"""
function transform!(nad::NARemover, nfeatures::DataFrame)
   features = deepcopy(nfeatures)
   (features == DataFrame()) && throw(ArgumentError("empty dataframe"))
   sz = nrow(features)
   tol = nad.model[:acceptance]
   colnames = []
   for (colname,dat) in collect(pairs(eachcol(features)))
      if sum(ismissing.(dat)) < tol*sz
         push!(colnames,colname)
      end
   end
   xtr =  features[:,colnames]
   return xtr
end

end


using Compose, DataFrames, Gadfly

set_default_plot_size(9inch, 6inch)

theme1 = Theme(point_size=6pt)
coord1 = Coord.cartesian(xmin=0.0, xmax=6.0)
D = DataFrame(x=1:5,
              y=[0.768448, 0.940515, 0.673959, 0.395453, 0.313244],
              V1=["A","A","B","B","D"],
              V2 = string.([1,2,2,3,3])  )

pa = plot(x=1:5, y=[0.77, 0.94, 0.67, 0.39, 0.31], shape=["A","A","B","B","D"], theme1, coord1,
    Guide.shapekey(title="Key",labels=["α","β","δ"]),
    Guide.title("Guide.shapekey") )

pb = plot(D, x=:x, y=:y, shape=:V1, color=:V1,
        Scale.shape_discrete(levels=["D","A","B"]),
        theme1, coord1, Guide.title("Shape==Color") )

pc = plot(D, x=:x, y=:y, shape=:V1, color=:V2,  coord1,
        Guide.colorkey(title="Color"),
        Guide.shapekey(title="Shape "),
        Theme(point_size=6pt, key_swatch_color="slategrey"),
        Guide.title("Shape!=Color") )

pd = plot(D, x=:x, y=:y, shape=[Shape.square], size=[8pt])
pe = plot()
pf = plot(Theme(point_size=6pt, alphas=[0.6]),
            layer(x=0.5:4, y=D.y[1:4], shape=["Gr1"]),
             layer(x=1:4, y=D.y[1:4], shape=["Gr2"]) )

gridstack([pa pb pc; pd pe pf])


function visualize!(vis, model,
    q_body, q_f1, q_f2, q_f3, q_f4;
    q_body_ref = q_body,
    ref_tl = 0.5,
    r = 0.025,
    l_torso = 0.19,
    w_torso = 0.1025,
    Δt = 0.1)

	default_background!(vis)

	setobject!(vis["torso"],
    	Rect(Vec(-l_torso, -w_torso, -0.025),Vec(2.0 * l_torso, 2.0 * w_torso, 0.05)),
    	MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

    setobject!(vis["torso_ref"],
    	Rect(Vec(-l_torso, -w_torso, -0.025),Vec(2.0 * l_torso, 2.0 * w_torso, 0.05)),
    	MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, ref_tl)))

	feet1 = setobject!(vis["feet1"], Sphere(Point3f0(0),
        convert(Float32, r)),
        MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0, 1.0)))

	feet2 = setobject!(vis["feet2"], Sphere(Point3f0(0),
		convert(Float32, r)),
		MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0, 1.0)))

	feet3 = setobject!(vis["feet3"], Sphere(Point3f0(0),
		convert(Float32, r)),
		MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0, 1.0)))

	feet4 = setobject!(vis["feet4"], Sphere(Point3f0(0),
        convert(Float32, r)),
        MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0, 1.0)))

	anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

	T = length(q_body)
	p_shift = [0.0, 0.0, r]

	for t = 1:T
		MeshCat.atframe(anim, t) do
			rot = MRP(q_body[t][4:6]...)
            rot_ref = MRP(q_body_ref[t][4:6]...)

			p_torso = q_body[t][1:3] + p_shift
            p_torso_ref = q_body_ref[t][1:3] + p_shift
			p_foot1 = q_f1[t] + p_shift
			p_foot2 = q_f2[t] + p_shift
			p_foot3 = q_f3[t] + p_shift
			p_foot4 = q_f4[t] + p_shift

			settransform!(vis["torso"], compose(Translation(p_torso), LinearMap(rot)))
            settransform!(vis["torso_ref"], compose(Translation(p_torso_ref), LinearMap(rot_ref)))
			settransform!(vis["feet1"], Translation(p_foot1))
			settransform!(vis["feet2"], Translation(p_foot2))
			settransform!(vis["feet3"], Translation(p_foot3))
			settransform!(vis["feet4"], Translation(p_foot4))
		end
	end

	# settransform!(vis["/Cameras/default"],
	#     compose(Translation(0.0, 0.0, -1.0), LinearMap(RotZ(-pi / 2.0))))

	MeshCat.setanimation!(vis, anim)
end


using ServerMetrics
using HTTP

@testcase "Simple counter" begin
    clear_registry!(get_default_registry())
    c = register!(Counter(), "my_counter")
    inc!(c, 2.0)

    resp = handle_metrics(HTTP.Request("GET", "/metrics"), nothing)
    @test String(resp.body) == """
    # TYPE my_counter counter
    my_counter 2.0

    """

    inc!(c, 3.0)
    resp = handle_metrics(HTTP.Request("GET", "/metrics"), nothing)
    @test String(resp.body) == """
    # TYPE my_counter counter
    my_counter 5.0

    """
end

@testcase "Simple gauge" begin
    clear_registry!(get_default_registry())
    g = register!(Gauge(1.5), "my_gauge")
    resp = handle_metrics(HTTP.Request("GET", "/metrics"), nothing)

    @test String(resp.body) == """
    # TYPE my_gauge gauge
    my_gauge 1.5

    """
end

@testcase "Counter with labels" begin
    clear_registry!(get_default_registry())
    c = register!(Counter(;action=String,response_code=Int64), "requests")
    inc!(c; action="get", response_code=404)
    inc!(c; action="put", response_code=500)  # What a bad day for prod :-(
    resp = handle_metrics(HTTP.Request("GET", "/metrics"), nothing)

    @test String(resp.body) == """
    # TYPE requests counter
    requests{action="get",response_code="404"} 1.0
    requests{action="put",response_code="500"} 1.0

    """
end

@testcase "Gauge with labels" begin
    clear_registry!(get_default_registry())
    temp = register!(Gauge(;location=String, hour=Int64), "temperature")
    set!(temp, 36.0; location="outside", hour=6)
    set!(temp, 40.0; location="outside", hour=8)
    set!(temp, 60.0; location="inside", hour=8)
    resp = handle_metrics(HTTP.Request("GET", "/metrics"), nothing)

    @test String(resp.body) == """
    # TYPE temperature gauge
    temperature{hour="6",location="outside"} 36.0
    temperature{hour="8",location="inside"} 60.0
    temperature{hour="8",location="outside"} 40.0

    """
end

@testcase "Exported metrics are sorted" begin
    clear_registry!(get_default_registry())
    bbb = register!(Counter(), "bbb")
    zzz = register!(Gauge(0.2), "zzz")
    aaa = register!(Counter(), "aaa")
    resp = handle_metrics(HTTP.Request("GET", "/metrics"), nothing)

    @test String(resp.body) == """
    # TYPE aaa counter
    aaa 0.0

    # TYPE bbb counter
    bbb 0.0

    # TYPE zzz gauge
    zzz 0.2

    """
end


using Test
using ObservableCollector

output = []

@test begin
    output = []
    obs = @observations begin
        @every 2 "O2" (s,t)->t
    end
    for t in 1:10
        obs(output,0,t)
    end
    output == [(name=:O2, val=2j) for j in 1:5]
end

@test begin
    output = []
    obs = @observations begin
        @at 2 "AT" (s,t)->t
    end
    for t in 1:10
        obs(output,0,t)
    end
    output == [(name=:AT, val=2)]
end

@test begin
    output = []
    obs = @observations begin
        @condition (s,t)->s+t==5 "sumtofive" (s,t)->t
    end
    for t in 1:10
        obs(output,1,t)
    end
    output == [(name=:sumtofive, val=4)]
end

@test begin
    output = []
    obs = @observations begin
        @condition (s,t)->t>=3 begin
            "A" --> (s,t)->s
            "T" --> (s,t)->t
        end
    end
    for t in 1:4
        obs(output,1,t)
    end
    output == [(name=:A, val=1),(name=:T, val=3),(name=:A, val=1),(name=:T, val=4)]
end


struct MagneticPartitioner end

function partition_stripe(A::SparseMatrixCSC, k, ::MagneticPartitioner, Π; kwargs...)
    @inbounds begin
        (m, n) = size(A)
        Π = convert(MapPartition, Π)
        asg = undefs(Int, n)
        for j = 1:n
            if A.colptr[j] == A.colptr[j + 1]
                asg[j] = rand(1:k)
                #asg[j] = k
            else
                i = A.rowval[rand(A.colptr[j]:A.colptr[j + 1] - 1)]
                #i = A.rowval[A.colptr[j]]
                asg[j] = Π.asg[i]
            end
        end
        return MapPartition(k, asg)
    end
end

struct GreedyBottleneckPartitioner{Mdl}
    mdl::Mdl
end

function partition_stripe(A::SparseMatrixCSC{Tv, Ti}, K, method::GreedyBottleneckPartitioner{Mdl}, Π; adj_A = nothing, kwargs...) where {Tv, Ti, Mdl <: AbstractSecondaryConnectivityModel}
    @inbounds begin
        (m, n) = size(A)
        if adj_A === nothing
            adj_A = adjointpattern(A)
        end
        n_vertices = undefs(Ti,K)
        n_pins = undefs(Ti,K)
        n_local_nets = undefs(Ti,K)
        n_remote_nets = undefs(Ti,K)
        cst = undefs(cost_type(method.mdl),K)

        Π_dmn = convert(DomainPartition, Π)
        Π_map = convert(MapPartition, Π)

        m, n = size(A)
        hst = zeros(Ti, n)
        for k = 1:K
            s = Π_dmn.spl[k]
            s′ = Π_dmn.spl[k + 1]
            n_vertices_k = s′ - s
            n_pins_k = 0
            n_local_nets_k = 0
            n_remote_nets_k = 0
            for _s = s:(s′ - 1)
                _i = Π_dmn.prm[_s]
                q = adj_A.colptr[_i]
                q′ = adj_A.colptr[_i + 1]
                n_pins_k += q′ - q
                for _q = q : q′ - 1
                    j = adj_A.rowval[_q]
                    if hst[j] < s
                        n_remote_nets_k += 1
                    end
                    hst[j] = s
                end
            end
            n_vertices[k] = n_vertices_k
            n_pins[k] = n_pins_k
            n_local_nets[k] = n_local_nets_k
            n_remote_nets[k] = n_remote_nets_k
            cst[k] = method.mdl(n_vertices_k, n_pins_k, n_local_nets_k, n_remote_nets_k, k)
        end

        asg = ones(Ti, n)

        for j = randperm(n)
            best_c = -Inf
            best_k = 0
            for q = A.colptr[j] : A.colptr[j + 1] - 1
                i = A.rowval[q]
                k = Π_map.asg[i]
                c = cst[k]
                if c > best_c
                    best_c = c
                    best_k = k
                end
            end
            if best_k != 0
                asg[j] = best_k
                n_local_nets[best_k] += 1
                n_remote_nets[best_k] -= 1
                cst[best_k] = method.mdl(n_vertices[best_k], n_pins[best_k], n_local_nets[best_k], n_remote_nets[best_k], best_k)
            end
        end

        return MapPartition(K, asg)
    end
end

function partition_stripe(A::SparseMatrixCSC{Tv, Ti}, K, method::GreedyBottleneckPartitioner{Mdl}, Π::SplitPartition; adj_A = nothing, adj_net=nothing, kwargs...) where {Tv, Ti, Mdl <: AbstractSecondaryConnectivityModel}
    @inbounds begin
        (m, n) = size(A)
        if adj_A === nothing
            adj_A = adjointpattern(A)
        end
        n_vertices = undefs(Ti,K)
        n_pins = undefs(Ti,K)
        n_local_nets = undefs(Ti,K)
        n_remote_nets = undefs(Ti,K)
        cst = undefs(cost_type(method.mdl),K)

        m, n = size(A)

        if adj_net === nothing
            hst = zeros(Ti, n)
            for k = 1:K
                i = Π.spl[k]
                i′ = Π.spl[k + 1]
                n_vertices_k = i′ - i
                n_pins_k = adj_A.colptr[i′] - adj_A.colptr[i]
                n_local_nets_k = 0
                n_remote_nets_k = 0
                for _i = i:(i′ - 1)
                    q = adj_A.colptr[_i]
                    q′ = adj_A.colptr[_i + 1]
                    for _q = q : q′ - 1
                        j = adj_A.rowval[_q]
                        if hst[j] < i
                            n_remote_nets_k += 1
                        end
                        hst[j] = i
                    end
                end
                n_vertices[k] = n_vertices_k
                n_pins[k] = n_pins_k
                n_local_nets[k] = n_local_nets_k
                n_remote_nets[k] = n_remote_nets_k
                cst[k] = method.mdl(n_vertices_k, n_pins_k, n_local_nets_k, n_remote_nets_k, k)
            end
        else
            for k = 1:K
                i = Π.spl[k]
                i′ = Π.spl[k + 1]
                n_vertices_k = i′ - i
                n_pins_k = adj_A.colptr[i′] - adj_A.colptr[i]
                n_local_nets_k = 0
                n_remote_nets_k = adj_net[i, i′]
                n_vertices[k] = n_vertices_k
                n_pins[k] = n_pins_k
                n_local_nets[k] = n_local_nets_k
                n_remote_nets[k] = n_remote_nets_k
                cst[k] = method.mdl(n_vertices_k, n_pins_k, n_local_nets_k, n_remote_nets_k, k)
            end
        end

        Π = convert(MapPartition, Π)
        asg = ones(Ti, n)

        for j = randperm(n)
            best_c = typemin(cost_type(method.mdl))
            best_k = 0
            for q = A.colptr[j] : A.colptr[j + 1] - 1
                i = A.rowval[q]
                k = Π.asg[i]
                c = cst[k]
                if c > best_c
                    best_c = c
                    best_k = k
                end
            end
            if best_k != 0
                asg[j] = best_k
                n_local_nets[best_k] += 1
                n_remote_nets[best_k] -= 1
                cst[best_k] = method.mdl(n_vertices[best_k], n_pins[best_k], n_local_nets[best_k], n_remote_nets[best_k], best_k)
            end
        end

        return MapPartition(K, asg)
    end
end


#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

import Zygote
using Zygote: @adjoint

using Flux, DiffEqFlux
using OrdinaryDiffEq
using DiffEqCallbacks

include("FMI2_neural.jl")

"""
Structure definition for a NeuralFMU, that runs in mode `Model Exchange` (ME).
"""
mutable struct ME_NeuralFMU
    neuralODE::NeuralODE
    solution::ODESolution
    fmu::FMU

    tspan
    saveat
    saved_values
    recordValues

    ME_NeuralFMU() = new()
end

"""
Structure definition for a NeuralFMU, that runs in mode `Co-Simulation` (CS).
"""
mutable struct CS_NeuralFMU
    model
    #simulationResult::fmi2SimulationResult
    fmu::FMU

    tspan
    saveat
    valueStack

    CS_NeuralFMU() = new()
end

NeuralFMU = ME_NeuralFMU
NeuralFMUs = Union{ME_NeuralFMU, CS_NeuralFMU}

# time caching (to set correct time in ME-NeuralFMUs)
function NeuralFMUGetCachedTime(fmu)
    Zygote.ignore() do
        return fmu.next_t
    end

    return 0.0
end
function NeuralFMUCacheTime(fmu, t)
    fmu.next_t = t
    nothing
end
function NeuralFMUCacheTime_Gradient(c̄, fmu, t)
    tuple(0.0, 0.0)
end
@adjoint NeuralFMUCacheTime(fmu, t) = NeuralFMUCacheTime(fmu, t), c̄ -> NeuralFMUCacheTime_Gradient(c̄, fmu, t)

# state caching (to set correct state in ME-NeuralFMUs)
function NeuralFMUCacheState(fmu, x)
    fmu.next_x = x
end
function NeuralFMUCacheState_Gradient(c̄, fmu, x)
    tuple(0.0, 0.0)
end
@adjoint NeuralFMUCacheState(fmu, x) = NeuralFMUCacheState(fmu, x), c̄ -> NeuralFMUCacheState_Gradient(c̄, fmu, x)

# state derivative caching (to set correct state derivative in ME-NeuralFMUs)
function NeuralFMUCacheStateDerivative(fmu, dx)
    fmu.next_dx = dx
end
function NeuralFMUCacheStateDerivative_Gradient(c̄, fmu, dx)
    tuple(0.0, 0.0)
end
@adjoint NeuralFMUCacheStateDerivative(fmu, dx) = NeuralFMUCacheStateDerivative(fmu, dx), c̄ -> NeuralFMUCacheStateDerivative_Gradient(c̄, fmu, dx)

# helper to add an additional time state later used to setup time inside the FMU
function NeuralFMUInputLayer(fmu, inputs)
    t = inputs[1]
    x = inputs[2:end]
    NeuralFMUCacheTime(fmu, t)
    x
end

# helper to add an additional time state derivative (for ME-NeuralFMUs)
function NeuralFMUOutputLayerME(inputs)
    dt = 1.0
    dx = inputs
    vcat([dt], dx) # [dt, dx...]
end

# helper to add an additional time state derivative (for CS-NeuralFMUs)
function NeuralFMUOutputLayerCS(fmu, inputs)
    out = inputs
    t = fmu.t # NeuralFMUGetCachedTime(fmu)
    vcat([t], out)
end

# function that saves periodically FMU values (in case you are interested in any)
function saveFunc(nfmu, u, t, integrator)
    values = fmi2GetReal(nfmu.fmu, nfmu.recordValues)
    return (values...,)
end

"""
Constructs a ME-NeuralFMU where the FMU is at an unknown location inside of the NN.

Arguents:
    - `saveat` time points to save the NeuralFMU output, if empty, solver step size is used (may be non-equidistant)
    - `addTopLayer` adds an input layer to add a time state (default) to set the FMU-time
    - `addBottomLayer` adds an input layer to remove the time state (default)

Keyword arguments:
    - `fixstep` forces fixed step integration
    - `recordFMUValues` additionally records internal FMU variables (currently not supported because of issues)
"""
function ME_NeuralFMU(fmu, model, tspan, alg=nothing; saveat=[], addTopLayer = true, addBottomLayer = true, recordFMUValues = [], fixstep=0.0)
    nfmu = ME_NeuralFMU()
    nfmu.fmu = fmu

    ext_model = nothing

    if addTopLayer && addBottomLayer
        ext_model = Chain(inputs -> NeuralFMUInputLayer(nfmu.fmu, inputs),
                          model.layers...,
                          inputs -> NeuralFMUOutputLayerME(inputs))
    elseif addTopLayer
        ext_model = Chain(inputs -> NeuralFMUInputLayer(nfmu.fmu, inputs),
                          model.layers...)
    elseif addBottomLayer
        ext_model = Chain(model.layers...,
                          inputs -> NeuralFMUOutputLayerME(inputs))
    else
        ext_model = Chain(model.layers...)
    end

    nfmu.recordValues = recordFMUValues
    if length(nfmu.recordValues) > 0
        @assert false "ME_NeuralFMU(...): keyword `recordFMUValues` is currently not supported (under development). Please remove keyword."

        tmp = zeros(Float64, length(nfmu.recordValues))
        type = typeof((tmp...,))

        nfmu.saved_values = SavedValues(Float64, type)
        cb = SavingCallback((u,t,integrator)->saveFunc(nfmu, u, t, integrator), nfmu.saved_values; saveat=saveat)

        if fixstep > 0.0
            nfmu.neuralODE = NeuralODE(ext_model, tspan, alg, callback=cb, dt=fixstep, adaptive=false)
        else
            nfmu.neuralODE = NeuralODE(ext_model, tspan, alg, callback=cb)
        end
    else
        if fixstep > 0.0
            nfmu.neuralODE = NeuralODE(ext_model, tspan, alg; saveat=saveat, dt=fixstep, adaptive=false)
        else
            nfmu.neuralODE = NeuralODE(ext_model, tspan, alg; saveat=saveat)
        end
    end

    nfmu.tspan = tspan
    nfmu.saveat = saveat

    nfmu
end

"""
Constructs a CS-NeuralFMU where the FMU is at an unknown location inside of the NN.

Arguents:
    - `saveat` time points to save the NeuralFMU output, if empty, solver step size is used (may be non-equidistant)
    - `addTopLayer` adds an input layer to add a time state (default) to set the FMU-time
    - `addBottomLayer` adds an input layer to remove the time state (default)
"""
function CS_NeuralFMU(fmu, model, tspan; saveat=[], addTopLayer = true, addBottomLayer = true, recordValues = [])
    nfmu = CS_NeuralFMU()
    nfmu.fmu = fmu

    if addTopLayer && addBottomLayer
        nfmu.model = Chain(inputs -> NeuralFMUInputLayer(nfmu.fmu, inputs),
                          model.layers...,
                          inputs -> NeuralFMUOutputLayerCS(nfmu.fmu, inputs))
    elseif addTopLayer
        nfmu.model = Chain(inputs -> NeuralFMUInputLayer(nfmu.fmu, inputs),
                          model.layers...)
    elseif addBottomLayer
        nfmu.model = Chain(model.layers...,
                          inputs -> NeuralFMUOutputLayerCS(nfmu.fmu, inputs))
    else
        nfmu.model = model.layers
    end

    nfmu.tspan = tspan
    nfmu.saveat = saveat

    nfmu
end

function (nfmu::ME_NeuralFMU)(x_start, t_start::Real = nfmu.tspan[1]; reset::Bool=true)

    if reset
        fmiReset(nfmu.fmu)
        fmiSetupExperiment(nfmu.fmu, t_start)
        fmiEnterInitializationMode(nfmu.fmu)
        fmiExitInitializationMode(nfmu.fmu)
    end

    x0 = [t_start, x_start...]
    nfmu.solution = nfmu.neuralODE(x0)
end

function (nfmu::CS_NeuralFMU)(t_step, t_start = nfmu.tspan[1], t_stop = nfmu.tspan[end]; inputs=[], reset::Bool=true)

    if reset
        while fmiReset(nfmu.fmu) != 0
        end
        while fmiSetupExperiment(nfmu.fmu, t_start) != 0
        end
        while fmiEnterInitializationMode(nfmu.fmu) != 0
        end
        while fmiExitInitializationMode(nfmu.fmu) != 0
        end
    end

    ts = t_start:t_step:t_stop
    modelInput = collect.(eachrow(hcat(ts, inputs)))
    nfmu.valueStack = nfmu.model.(modelInput)

    nfmu.valueStack
end

# adapting the Flux functions
function Flux.params(neuralfmu::ME_NeuralFMU)
    Flux.params(neuralfmu.neuralODE)
end

function Flux.params(neuralfmu::CS_NeuralFMU)
    Flux.params(neuralfmu.model)
end

# FMI version independent dosteps
# function fmiDoStepME(fmu::FMU2, x, t = -1.0; setValueReferences = [], setValues = [], getValueReferences = [])
#     fmi2DoStepME(fmu, x, t;
#                 setValueReferences = setValueReferences,
#                 setValues = setValues,
#                 getValueReferences = getValueReferences)
# end
function fmiDoStepME(fmu::FMU2, x, t = -1.0, setValueReferences = [], setValues = [], getValueReferences = [])
    fmi2DoStepME(fmu, x, t,
                setValueReferences,
                setValues,
                getValueReferences)
end

function fmiDoStepCS(fmu::FMU2, dt; setValueReferences = [], setValues = [], getValueReferences = [])
    fmi2DoStepCS(fmu, dt;
                setValueReferences = setValueReferences,
                setValues = setValues,
                getValueReferences = getValueReferences)
end

function fmiInputDoStepCSOutput(fmu::FMU2, dt, u)
    fmi2InputDoStepCSOutput(fmu, dt, u)
end

# define neutral gradients (=feed-trough) for ccall-functions
function neutralGradient(c̄)
    c̄
end

@adjoint fmiSetupExperiment(fmu, startTime, stopTime) = fmiSetupExperiment(fmu, startTime, stopTime), c̄ -> neutralGradient(c̄)
@adjoint fmiEnterInitializationMode(fmu) = fmiEnterInitializationMode(fmu), c̄ -> neutralGradient(c̄)
@adjoint fmiExitInitializationMode(fmu) = fmiExitInitializationMode(fmu), c̄ -> neutralGradient(c̄)
@adjoint fmiReset(fmu) = fmiReset(fmu), c̄ -> neutralGradient(c̄)


const MOI = MathOptInterface
const CI = MOI.ConstraintIndex
const VI = MOI.VariableIndex

const MOIU = MOI.Utilities

Base.@kwdef mutable struct MOISolution
    ret_val::Int = 0
    raw_status::String = "Problem not solved"
    primal::Vector{Float64} = Float64[] # primal of variables
    dual_eq::Vector{Float64} = Float64[] # dual of constraints
    dual_in::Vector{Float64} = Float64[] # dual of constraints
    dual_cone::Vector{Float64} = Float64[] # dual of constraints
    slack_eq::Vector{Float64} = Float64[]
    slack_in::Vector{Float64} = Float64[]
    # slack_sd::Vector{Vector{Float64}}
    # slack_so::Vector{Vector{Float64}}
    primal_residual::Float64 = NaN
    dual_residual::Float64 = NaN
    objval::Float64 = NaN
    dual_objval::Float64 = NaN
    # obj_cte::Float64 = NaN # used in rays - see SCS
    gap::Float64 = NaN
    time::Float64 = NaN
    final_rank::Int = -1
    primal_feasible_user_tol::Bool = false
    dual_feasible_user_tol::Bool = false
    certificate_found::Bool = false
end

# `ModelData` struct is discarded
mutable struct ModelData
    m::Int # Number of rows/constraints
    n::Int # Number of cols/variables

    # equality matrix
    I::Vector{Int} # List of rows
    J::Vector{Int} # List of cols
    V::Vector{Float64} # List of coefficients
    # equality rhs
    b::Vector{Float64} # constants

    # inequality matrix
    Ii::Vector{Int} # List of rows
    Ji::Vector{Int} # List of cols
    Vi::Vector{Float64} # List of coefficients
    # inequality rhs
    h::Vector{Float64} # constants

    objective_constant::Float64 # The objective is min c'x + objective_constant
    c::Vector{Float64}
end

Base.@kwdef mutable struct ConeData

    cone_cols::Int = 0

    # matrix lines
    eq_tot::Int = 0 # number of linear equality constraints
    eq_start::Vector{Int} = Int[]
    eq_len::Vector{Int} = Int[]
    in_tot::Int = 0 # length of LP cone
    in_start::Vector{Int} = Int[]
    in_len::Vector{Int} = Int[]

    # cones
    sdc::Vector{Vector{Int}} = Vector{Int}[] # semidefinite
    soc::Vector{Vector{Int}} = Vector{Int}[] # second order
end

mutable struct Optimizer <: MOI.AbstractOptimizer
    cone::ConeData
    maxsense::Bool
    data::Union{Nothing, ModelData} # only non-Void between MOI.copy_to and MOI.optimize!
    sol::MOISolution
    options::Options
end
function Optimizer(;kwargs...)
    optimizer = Optimizer(ConeData(), false, nothing, MOISolution(), Options())
    for (key, value) in kwargs
        MOI.set(optimizer, MOI.RawParameter(key), value)
    end
    return optimizer
end

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

function MOI.set(optimizer::Optimizer, param::MOI.RawParameter, value)
    fields = fieldnames(Options)
    name = Symbol(param.name)
    if name in fields
        setfield!(optimizer.options, name, value)
    else
        error("No parameter matching $(name)")
    end
    return value
end
function MOI.get(optimizer::Optimizer, param::MOI.RawParameter)
    fields = fieldnames(Options)
    name = Symbol(param.name)
    if name in fields
        return getfield(optimizer.options, name)
    else
        error("No parameter matching $(name)")
    end
end

MOI.supports(::Optimizer, ::MOI.Silent) = true
function MOI.set(optimizer::Optimizer, ::MOI.Silent, value::Bool)
    if value == true
        optimizer.options.timer_verbose = false
    end
    optimizer.options.log_verbose = !value
end
function MOI.get(optimizer::Optimizer, ::MOI.Silent)
    if optimizer.options.log_verbose
        return false
    elseif optimizer.options.timer_verbose
        return false
    else
        return true
    end
end

MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
function MOI.set(optimizer::Optimizer, ::MOI.TimeLimitSec, value)
    optimizer.options.time_limit = value
end
function MOI.get(optimizer::Optimizer, ::MOI.TimeLimitSec)
    return optimizer.options.time_limit
end

MOI.supports(::Optimizer, ::MOI.NumberOfThreads) = false

function MOI.is_empty(optimizer::Optimizer)
    !optimizer.maxsense && optimizer.data === nothing
end
function MOI.empty!(optimizer::Optimizer)
    optimizer.maxsense = false
    optimizer.data = nothing # It should already be nothing except if an error is thrown inside copy_to
    optimizer.cone = ConeData()
    optimizer.sol.ret_val = 0
end

MOIU.supports_allocate_load(::Optimizer, copy_names::Bool) = !copy_names

function MOI.supports(::Optimizer,
                      ::Union{MOI.ObjectiveSense,
                              MOI.ObjectiveFunction{MOI.SingleVariable},
                              MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}})
    return true
end

function MOI.supports_constraint(
    ::Optimizer,
    ::Type{<:MOI.VectorAffineFunction{Float64}},
    ::Type{<:Union{MOI.Zeros, MOI.Nonpositives}})
    return true
end

function MOI.supports_constraint(
    ::Optimizer,
    ::Type{<:MOI.VectorOfVariables},
    ::Type{<:Union{MOI.SecondOrderCone,
                   MOI.PositiveSemidefiniteConeTriangle}})
    return true
end

function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike; kws...)
    return MOIU.automatic_copy_to(dest, src; kws...)
end

const SupportedSets = Union{MOI.PositiveSemidefiniteConeTriangle,
                            MOI.SecondOrderCone}
function get_new_variable_indexes_unique(optimizer::Optimizer, set::SupportedSets)
    len = MOI.dimension(set)
    first = optimizer.cone.cone_cols + 1
    last = optimizer.cone.cone_cols + len
    optimizer.cone.cone_cols += len
    return collect(first:last)
end
function get_cone_list(cone::ConeData, set::MOI.PositiveSemidefiniteConeTriangle)
    return cone.sdc
end
function get_cone_list(cone::ConeData, set::MOI.SecondOrderCone)
    return cone.soc
end
function get_cone_list(optimizer::Optimizer, set::SupportedSets)
    get_cone_list(optimizer.cone, set)
end

function MOIU.allocate_constrained_variables(optimizer::Optimizer, set::SupportedSets)
    inds = get_new_variable_indexes_unique(optimizer,
    # func,
    set)
    vars = MOI.VariableIndex.(inds)
    list = get_cone_list(optimizer, set)
    push!(list, inds)
    ci = MOI.ConstraintIndex{MOI.VectorOfVariables, typeof(set)}(length(list))
    return vars, ci
end

function MOIU.load_constrained_variables(
    optimizer::Optimizer, vis::Vector{MOI.VariableIndex},
    ci::MOI.ConstraintIndex{MOI.VectorOfVariables},
    set::SupportedSets)
end

function _allocate_constraint(cone::ConeData, f, set::SupportedSets)
    list = get_cone_list(cone, set)
    push!(list, Int[])
    return length(list)
end

function _allocate_constraint(cone::ConeData, f, s::MOI.Zeros)
    len = MOI.dimension(s)
    push!(cone.eq_start, cone.eq_tot + 1)
    push!(cone.eq_len, len)
    cone.eq_tot += len
    return length(cone.eq_len)
end

function _allocate_constraint(cone::ConeData, f, s::MOI.Nonpositives)
    len = MOI.dimension(s)
    push!(cone.in_start, cone.in_tot + 1)
    push!(cone.in_len, len)
    cone.in_tot += len
    return length(cone.in_len)
end

function MOIU.load_constraint(optimizer::Optimizer, ci::CI,
    f::MOI.VectorOfVariables,
    set::SupportedSets)

    vec = get_cone_list(optimizer.cone, set)[ci.value]

    append!(vec, [i.value for i in f.variables])
    nothing
end

function MOIU.allocate_constraint(optimizer::Optimizer, f::F, s::S) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
    return CI{F, S}(_allocate_constraint(optimizer.cone, f, s))
end

# Vectorized length for matrix dimension n
sympackedlen(n) = div(n*(n+1), 2)
# Matrix dimension for vectorized length n
sympackeddim(n) = div(isqrt(1+8n) - 1, 2)

output_index(t::MOI.VectorAffineTerm) = t.output_index
variable_index_value(v::MOI.VariableIndex) = v.value
variable_index_value(t::MOI.ScalarAffineTerm) = t.variable_index.value
variable_index_value(t::MOI.VectorAffineTerm) = variable_index_value(t.scalar_term)
coefficient(t::MOI.ScalarAffineTerm) = t.coefficient
coefficient(t::MOI.VectorAffineTerm) = coefficient(t.scalar_term)

function MOIU.load_constraint(optimizer::Optimizer, ci::CI,
                              f::MOI.VectorAffineFunction,
                              set::MOI.AbstractVectorSet)

    n = length(f.terms)
    if n != 1
        A = sparse(output_index.(f.terms), variable_index_value.(f.terms),
        coefficient.(f.terms))
        # sparse combines duplicates with + but does
        # not remove zeros created so we call dropzeros!
        dropzeros!(A)
        I, J, V = findnz(A)
    else
        I, J, V = output_index.(f.terms), variable_index_value.(f.terms), coefficient.(f.terms)
    end
#=
    n = length(f.terms)
    I = Int[]
    J = Int[]
    V = Float64[]
    sizehint!(I, n)
    sizehint!(J, n)
    sizehint!(V, n)
    for i in 1:n
        if !iszero(coefficient(f.terms[i]))
            push!(I, output_index(f.terms[i]))
            push!(J, variable_index_value(f.terms[i]))
            push!(V, coefficient(f.terms[i]))
        end
    end
=#
    #=
    I,J,V = output_index.(f.terms), variable_index_value.(f.terms), coefficient.(f.terms)
    =#

    offset = ctr_offset(optimizer, set, ci)
    rows = 1:len(optimizer, set, ci)
    i = (offset-1) .+ rows

    rhs = matrix_rhs_vec(optimizer, set)
    rhs[i] .= -f.constants

    # MOI:  Ax + b {==, <=} 0
    # psdp: Ax {==, <=} -b
    I_, J_, V_ = matrix_triplets(optimizer, set)
    append!(I_, offset-1 .+ I)
    append!(J_, J)
    append!(V_, V)

    nothing
end
function ctr_offset(optimizer::Optimizer, set::MOI.Zeros, ci)
    return optimizer.cone.eq_start[ci.value]
end
function len(optimizer::Optimizer, set::MOI.Zeros, ci)
    return optimizer.cone.eq_len[ci.value]
end
function ctr_offset(optimizer::Optimizer, set::MOI.Nonpositives, ci)
    return optimizer.cone.in_start[ci.value]
end
function len(optimizer::Optimizer, set::MOI.Nonpositives, ci)
    return optimizer.cone.in_len[ci.value]
end
function matrix_triplets(optimizer::Optimizer, set::MOI.Zeros)
    return optimizer.data.I, optimizer.data.J, optimizer.data.V
end
function matrix_rhs_vec(optimizer::Optimizer, set::MOI.Zeros)
    return optimizer.data.b
end
function matrix_triplets(optimizer::Optimizer, set::MOI.Nonpositives)
    return optimizer.data.Ii, optimizer.data.Ji, optimizer.data.Vi
end
function matrix_rhs_vec(optimizer::Optimizer, set::MOI.Nonpositives)
    return optimizer.data.h
end

# called after ctr var - only aditonal
function MOIU.allocate_variables(optimizer::Optimizer, nvars::Integer)
    allocated = optimizer.cone.cone_cols
    VI.(allocated .+ (1:nvars))
end

function MOIU.load_variables(optimizer::Optimizer, nvars::Integer)
    # @show "load var"
    # @show nvars
    cone = optimizer.cone
    m = cone.eq_tot + cone.in_tot # + cone.q + cone.s + 3cone.ep + cone.ed
    I = Int[]
    J = Int[]
    V = Float64[]
    sizehint!(I, cone.eq_tot)
    sizehint!(J, cone.eq_tot)
    sizehint!(V, cone.eq_tot)
    b = zeros(cone.eq_tot)
    Ii = Int[]
    Ji = Int[]
    Vi = Float64[]
    sizehint!(Ii, cone.in_tot)
    sizehint!(Ji, cone.in_tot)
    sizehint!(Vi, cone.in_tot)
    h = zeros(cone.in_tot)

    tot_vars = nvars# + cone.cone_cols
    c = zeros(tot_vars)
    optimizer.data = ModelData(m, tot_vars, I, J, V, b, Ii, Ji, Vi, h, 0., c)
    # `optimizer.sol` contains the result of the previous optimization.
    # It is used as a warm start if its length is the same, e.g.
    # probably because no variable and/or constraint has been added.
    if length(optimizer.sol.primal) != tot_vars
        optimizer.sol.primal = zeros(tot_vars)
    end
    # TODO @joaquim
    # @assert length(optimizer.sol.dual) == length(optimizer.sol.slack)
    # if length(optimizer.sol.dual) != m
    #     optimizer.sol.dual = zeros(m)
    #     optimizer.sol.slack = zeros(m)
    # end
    return nothing
end

function MOIU.allocate(::Optimizer, ::MOI.VariablePrimalStart,
                       ::MOI.VariableIndex, ::Union{Nothing, Float64})
end
function MOIU.allocate(::Optimizer, ::MOI.ConstraintPrimalStart,
                       ::MOI.ConstraintIndex, ::Float64)
end
function MOIU.allocate(::Optimizer, ::MOI.ConstraintDualStart,
                       ::MOI.ConstraintIndex, ::Float64)
end
function MOIU.allocate(optimizer::Optimizer, ::MOI.ObjectiveSense,
                       sense::MOI.OptimizationSense)
    optimizer.maxsense = sense == MOI.MAX_SENSE
end
function MOIU.allocate(::Optimizer, ::MOI.ObjectiveFunction,
    ::MOI.Union{MOI.SingleVariable,
                MOI.ScalarAffineFunction{Float64}})
end

function MOIU.load(::Optimizer, ::MOI.VariablePrimalStart,
    ::MOI.VariableIndex, ::Nothing)
end
function MOIU.load(optimizer::Optimizer, ::MOI.VariablePrimalStart,
                   vi::MOI.VariableIndex, value::Float64)
    optimizer.sol.primal[vi.value] = value
end
function MOIU.load(::Optimizer, ::MOI.ConstraintPrimalStart,
    ::MOI.ConstraintIndex, ::Nothing)
end
function MOIU.load(::Optimizer, ::MOI.ConstraintDualStart,
    ::MOI.ConstraintIndex, ::Nothing)
end
function MOIU.load(::Optimizer, ::MOI.ObjectiveSense, ::MOI.OptimizationSense)
end
function MOIU.load(optimizer::Optimizer, ::MOI.ObjectiveFunction,
                   f::MOI.SingleVariable)
    MOIU.load(optimizer,
          MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
          MOI.ScalarAffineFunction{Float64}(f))
    return nothing
end
function MOIU.load(optimizer::Optimizer, ::MOI.ObjectiveFunction,
               f::MOI.ScalarAffineFunction)
    c0 = Vector(sparsevec(variable_index_value.(f.terms), coefficient.(f.terms),
                        optimizer.data.n))
    optimizer.data.objective_constant = f.constant
    optimizer.data.c = optimizer.maxsense ? -c0 : c0
    return nothing
end

matindices(n::Integer) = (LinearIndices(tril(trues(n,n))))[findall(tril(trues(n,n)))]

function MOI.optimize!(optimizer::Optimizer)
    if optimizer.data === nothing
        # optimize! has already been called and no new model has been copied
        return
    end

    # @show "in opt"

    TimerOutputs.reset_timer!()

    @timeit "pre-processing" begin

    # parse options
    options = optimizer.options

    cone = optimizer.cone

    #=
        Build linear sets
    =#

    m = optimizer.data.m #rows
    n = optimizer.data.n #cols

    objective_constant = optimizer.data.objective_constant
    # c = optimizer.data.c

    # Build Prox SDP Affine Sets
    # b = optimizer.data.b
    # A = sparse(optimizer.data.I, optimizer.data.J, optimizer.data.V, length(b), n)
    # h = optimizer.data.h
    # G = sparse(optimizer.data.Ii, optimizer.data.Ji, optimizer.data.Vi, length(h), n)

    # Dimensions (of affine sets)

    aff = AffineSets(0, 0, 0, 0,
    sparse(optimizer.data.I, optimizer.data.J, optimizer.data.V, length(optimizer.data.b), n),
    sparse(optimizer.data.Ii, optimizer.data.Ji, optimizer.data.Vi, length(optimizer.data.h), n),
    optimizer.data.b,
    optimizer.data.h,
    optimizer.data.c)
    # A, G, b, h, c)

    n_variables = size(aff.A)[2] # primal
    n_eqs = size(aff.A)[1]
    n_ineqs = size(aff.G)[1]
    aff.n = n_variables
    aff.p = n_eqs
    aff.m = n_ineqs

    #=
        Build conic sets
    =#

    # Build SDP Sets
    con = ConicSets(
        SDPSet[],
        SOCSet[]
        )

    # create extra variables
    n_tot_variables = n_variables
    In, Jn, Vn = Int[], Int[], Float64[]

    # this way there is a single elements per column
    # because we assume VOV in SET and NOT AFF in SET
    for i in eachindex(cone.soc)
        n_vars = length(cone.soc[i])
        vec_inds = cone.soc[i]
        n_tot_variables += fix_duplicates!(vec_inds, n_tot_variables, In, Jn, Vn)
        push!(con.socone, SOCSet(vec_inds, n_vars))
        # first_ind_local += n_vars
    end

    for i in eachindex(cone.sdc)
        n_vars = length(cone.sdc[i])
        vec_inds = cone.sdc[i]
        n_tot_variables += fix_duplicates!(vec_inds, n_tot_variables, In, Jn, Vn)
        sq_side = sympackeddim(n_vars)
        sq_len = sq_side*sq_side
        tri_len = n_vars
        mat_inds = matindices(sq_side)
        push!(con.sdpcone, SDPSet(vec_inds, mat_inds, tri_len, sq_len, sq_side))
        # first_ind_local += n_vars
    end

    optimizer.data = nothing # Allows GC to free optimizer.data before A is loaded

    #=
        Pre-process to remove duplicates from differente cones
    =#

    for i in 1:length(con.socone)
        for j in i+1:length(con.socone)
            vec1 = con.socone[i].idx
            vec2 = con.socone[j].idx
            n_tot_variables += fix_duplicates!(vec1, vec2, n_tot_variables, In, Jn, Vn)
        end
    end

    for i in 1:length(con.sdpcone)
        for j in i+1:length(con.sdpcone)
            vec1 = con.sdpcone[i].vec_i
            vec2 = con.sdpcone[j].vec_i
            n_tot_variables += fix_duplicates!(vec1, vec2, n_tot_variables, In, Jn, Vn)
        end
    end

    for i in 1:length(con.sdpcone)
        for j in 1:length(con.socone)
            vec1 = con.sdpcone[i].vec_i
            vec2 = con.socone[j].idx
            n_tot_variables += fix_duplicates!(vec1, vec2, n_tot_variables, In, Jn, Vn)
        end
    end

    n_new_variables = n_tot_variables - n_variables

    A2 = sparse(In .- n_variables, Jn, Vn, n_new_variables, n_tot_variables)
    b2 = zeros(n_new_variables)

    append!(aff.b, b2)
    aff.A = hvcat((2,1), aff.A, spzeros(size(aff.A)[1], n_new_variables), A2)
    aff.G = hcat(aff.G, spzeros(size(aff.G)[1], n_new_variables))
    append!(aff.c, zeros(n_new_variables))
    aff.n += n_new_variables
    aff.p += n_new_variables
    aff.extra = n_new_variables

    In, Jn, Vn = Int[], Int[], Float64[]
    A2 = sparse(In, Jn, Vn, 0, 0)

    #=
        Solve modified problem
    =#

    end

    # warm = WarmStart()

    if options.disable_julia_logger
        # disable logger
        global_log = Logging.current_logger()
        Logging.global_logger(Logging.NullLogger())
    end

    sol = @timeit "Main" chambolle_pock(aff, con, options)

    if options.disable_julia_logger
        # re-enable logger
        Logging.global_logger(global_log)
    end

    #=
        Unload solution
    =#

    ret_val = sol.status
    primal = sol.primal[1:aff.n-aff.extra]
    dual_cone = sol.dual_cone[1:aff.n-aff.extra]
    dual_eq = sol.dual_eq[1:aff.p-aff.extra]
    slack_eq = sol.slack_eq[1:aff.p-aff.extra]
    objval = sol.objval

    if options.timer_verbose
        TimerOutputs.print_timer(TimerOutputs.DEFAULT_TIMER)
        print("\n")
        TimerOutputs.print_timer(TimerOutputs.flatten(TimerOutputs.DEFAULT_TIMER))
        print("\n")
    end
    if options.timer_file
        f = open("time.log","w")
        TimerOutputs.print_timer(f,TimerOutputs.DEFAULT_TIMER)
        print(f,"\n")
        TimerOutputs.print_timer(f,TimerOutputs.flatten(TimerOutputs.DEFAULT_TIMER))
        print(f,"\n")
        close(f)
    end

    optimizer.sol = MOISolution(
        ret_val,
        sol.status_string,
        primal,
        dual_eq,
        sol.dual_in,
        dual_cone,
        slack_eq,
        sol.slack_in,
        sol.primal_residual,
        sol.dual_residual,
        (optimizer.maxsense ? -1 : 1) * objval+objective_constant,
        (optimizer.maxsense ? -1 : 1) * sol.dual_objval+objective_constant,
        sol.gap,
        sol.time,
        sol.final_rank,
        sol.primal_feasible_user_tol,
        sol.dual_feasible_user_tol,
        sol.certificate_found,
        )
end

function get_indices_cone(A, rows, n_vars, first_ind_local)
    vec_inds = zeros(Int, n_vars)
    for (idx, col) in enumerate(first_ind_local:first_ind_local+n_vars-1)
        # columns -> pos in cone (because is tranposed)
        for j in nzrange(A, col)
            row = rows[j]
            position = idx
            var_idx = row # global
            vec_inds[position] = var_idx
        end
    end
    return vec_inds
end

function fix_duplicates!(vec1::Vector{Int}, vec2::Vector{Int}, n::Int,
                         In::Vector{Int}, Jn::Vector{Int}, Vn::Vector{Float64})
    duplicates = intersect(vec2, vec1)
    n_dups = length(duplicates)
    if n_dups == 0
        return 0
    end

    append!(Jn, duplicates)
    append!(Jn, collect(1:n_dups) .+ n)
    append!(In, collect(1:n_dups) .+ n)
    append!(In, collect(1:n_dups) .+ n)
    append!(Vn,  ones(n_dups))
    append!(Vn, -ones(n_dups))
    cont = 1
    for i in eachindex(vec2)
        if vec2[i] == duplicates[cont]
            vec2[i] = n + cont
            cont += 1
        end
    end
    @assert cont-1 == n_dups
    return n_dups
end

function fix_duplicates!(vec::Vector{Int}, n::Int,
                         In::Vector{Int}, Jn::Vector{Int}, Vn::Vector{Float64})
    seen = Set{eltype(vec)}()
    dups = Vector{eltype(vec)}()
    new_vars = 0
    for i in eachindex(vec)
        x = vec[i]
        if in(x, seen)
            push!(dups, x)
            new_vars += 1
            vec[i] = new_vars + n
        else
            push!(seen, x)
        end
    end
    n_dups = length(dups)
    if n_dups == 0
        return 0
    end
    append!(Jn, dups)
    append!(Jn, collect(1:n_dups) .+ n)
    append!(In, collect(1:n_dups) .+ n)
    append!(In, collect(1:n_dups) .+ n)
    append!(Vn,  ones(n_dups))
    append!(Vn, -ones(n_dups))
    return n_dups
end

function ivech!(out::AbstractMatrix{T}, v::AbstractVector{T}) where T
    n = sympackeddim(length(v))
    n1, n2 = size(out)
    @assert n == n1 == n2
    c = 0
    for j in 1:n, i in 1:j
        c += 1
        out[i,j] = v[c]
    end
    return out
end
function ivech(v::AbstractVector{T}) where T
    n = sympackeddim(length(v))
    out = zeros(n, n)
    ivech!(out, v)

    return out
end

ivec(X) = Matrix(Symmetric(ivech(X),:U))

#=
    Status
=#

function MOI.get(optimizer::Optimizer, ::MOI.SolveTime)
    return optimizer.sol.time
end

function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    return optimizer.sol.raw_status
end

# Implements getter for result value and statuses
# ProxSDP returns one of the following integers:
# 0 - MOI.OPTIMIZE_NOT_CALLED
# 1 - MOI.OPTIMAL
# 2 - MOI.TIME_LIMIT
# 3 - MOI.ITERATION_LIMIT
# 4 - MOI.INFEASIBLE_OR_UNBOUNDED
function MOI.get(optimizer::Optimizer, ::MOI.TerminationStatus)
    s = optimizer.sol.ret_val
    @assert 0 <= s <= 6
    if s == 0
        return MOI.OPTIMIZE_NOT_CALLED
    elseif s == 1
        return MOI.OPTIMAL
    elseif s == 2
        return MOI.TIME_LIMIT
    elseif s == 3
        return MOI.ITERATION_LIMIT
    elseif s == 4
        return MOI.INFEASIBLE_OR_UNBOUNDED
    elseif s == 5
        return MOI.DUAL_INFEASIBLE
    elseif s == 6
        return MOI.INFEASIBLE
    end
end

function MOI.get(optimizer::Optimizer, attr::MOI.ObjectiveValue)
    result_count_err(optimizer, attr)
    value = optimizer.sol.objval
    return value
end
function MOI.get(optimizer::Optimizer, attr::MOI.DualObjectiveValue)
    result_count_err(optimizer, attr)
    value = optimizer.sol.dual_objval
    return value
end

function MOI.get(optimizer::Optimizer, attr::MOI.PrimalStatus)
    s = optimizer.sol.ret_val
    if attr.N > 1 || s == 0
        return MOI.NO_SOLUTION
    end
    if s == 5 && optimizer.sol.certificate_found
        return MOI.INFEASIBILITY_CERTIFICATE
    end
    if optimizer.sol.primal_feasible_user_tol#s == 1
        return MOI.FEASIBLE_POINT
    else
        return MOI.INFEASIBLE_POINT
    end
end
function MOI.get(optimizer::Optimizer, attr::MOI.DualStatus)
    s = optimizer.sol.ret_val
    if attr.N > 1 || s ==0
        return MOI.NO_SOLUTION
    end
    if s == 6 && optimizer.sol.certificate_found
        return MOI.INFEASIBILITY_CERTIFICATE
    end
    if optimizer.sol.dual_feasible_user_tol
        return MOI.FEASIBLE_POINT
    else
        return MOI.INFEASIBLE_POINT
    end
end

function result_count_err(optimizer::Optimizer, attr)
    if optimizer.sol.ret_val == 0
        throw(MOI.ResultIndexBoundsError{typeof(attr)}(attr, 0))
    elseif result_index(attr) > 1
        throw(MOI.ResultIndexBoundsError{typeof(attr)}(attr, 1))
    end
    return nothing
end

result_index(a::MOI.ObjectiveValue) = a.result_index
result_index(a::MOI.DualObjectiveValue) = a.result_index
result_index(a::MOI.ConstraintDual) = a.N
result_index(a::MOI.ConstraintPrimal) = a.N
result_index(a::MOI.VariablePrimal) = a.N

#=
    Solution
=#

function MOI.get(optimizer::Optimizer, attr::MOI.VariablePrimal, vi::VI)
    result_count_err(optimizer, attr)
    optimizer.sol.primal[vi.value]
end
function MOI.get(optimizer::Optimizer, a::MOI.VariablePrimal, vi::Vector{VI})
    result_count_err(optimizer, a)
    return MOI.get.(optimizer, a, vi)
end

function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal,
                 ci::CI{<:MOI.AbstractFunction, MOI.Zeros})
    result_count_err(optimizer, attr)
    cone = optimizer.cone
    inds = (cone.eq_start[ci.value]-1) .+ (1:cone.eq_len[ci.value])
    return optimizer.sol.slack_eq[inds]
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal,
    ci::CI{<:MOI.AbstractFunction, MOI.Nonpositives})
    result_count_err(optimizer, attr)
    cone = optimizer.cone
    inds = (cone.in_start[ci.value]-1) .+ (1:cone.in_len[ci.value])
    return optimizer.sol.slack_in[inds]
end

function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal,
        ci::CI{MOI.VectorOfVariables, MOI.PositiveSemidefiniteConeTriangle})
    result_count_err(optimizer, attr)
    optimizer.sol.primal[optimizer.cone.sdc[ci.value]]
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal,
    ci::CI{MOI.VectorOfVariables, MOI.SecondOrderCone})
    result_count_err(optimizer, attr)
    optimizer.sol.primal[optimizer.cone.soc[ci.value]]
end

function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual,
    ci::CI{<:MOI.AbstractFunction, MOI.Zeros})
    result_count_err(optimizer, attr)
    cone = optimizer.cone
    inds = (cone.eq_start[ci.value]-1) .+ (1:cone.eq_len[ci.value])
    return -optimizer.sol.dual_eq[inds]
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual,
ci::CI{<:MOI.AbstractFunction, MOI.Nonpositives})
    result_count_err(optimizer, attr)
    cone = optimizer.cone
    inds = (cone.in_start[ci.value]-1) .+ (1:cone.in_len[ci.value])
    return -optimizer.sol.dual_in[inds]
end

function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual,
    ci::CI{MOI.VectorOfVariables, MOI.PositiveSemidefiniteConeTriangle})
    result_count_err(optimizer, attr)
    optimizer.sol.dual_cone[optimizer.cone.sdc[ci.value]]
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual,
    ci::CI{MOI.VectorOfVariables, MOI.SecondOrderCone})
    result_count_err(optimizer, attr)
    optimizer.sol.dual_cone[optimizer.cone.soc[ci.value]]
end

function MOI.get(optimizer::Optimizer, ::MOI.ResultCount)
    if MOI.get(optimizer, MOI.TerminationStatus()) in [
                # MOI.INFEASIBLE_OR_UNBOUNDED,
                # MOI.INFEASIBLE,
                # MOI.DUAL_INFEASIBLE,
                MOI.OPTIMIZE_NOT_CALLED
        ]
        return 0
    else
        return 1
    end
end


"""
    rank_to_sortperm(::NTuple{N,Int}) -> NTuple{N,Int}
Returns the `sortperm` of the stride ranks.
"""
function rank_to_sortperm(R::NTuple{N,Int}) where {N}
    sp = ntuple(zero, Val{N}())
    r = ntuple(n -> sum(R[n] .≥ R), Val{N}())
    @inbounds for n = 1:N
        sp = Base.setindex(sp, n, r[n])
    end
    sp
end
rank_to_sortperm(R) = sortperm(R)

# @inline function gep_no_offset(ptr::VectorizationBase.AbstractStridedPointer, i::Tuple)
    # VectorizationBase.gep(pointer(ptr), VectorizationBase.tdot(ptr, i, VectorizationBase.strides(ptr), VectorizationBase.nopromote_axis_indicator(ptr)))
# end
# @inline function similar_with_offset(sptr::StridedPointer{T,N,C,B,R,X,O}, ptr::Ptr{T}) where {T,N,C,B,R,X,O}
#     StridedPointer{T,N,C,B,R,X}(ptr, sptr.strd, zerotuple(Val{N}()))
# end
function view_quote(i, K, S, D, T, N, C, B, R, X, O, zero_offsets::Bool = false)
    @assert ((K == N) || isone(K))

    inds = Expr(:tuple)
    Nnew = 0
    s = Expr(:tuple)
    x = Expr(:tuple)
    o = Expr(:tuple)
    Rnew = Expr(:tuple)
    Dnew = Expr(:tuple)
    Cnew = -1
    Bnew = ifelse(iszero(B), 0, -1)
    sortp = rank_to_sortperm(R)
    still_dense = true
    densev = Vector{Bool}(undef, K)
    for k ∈ 1:K
        iₖ = Expr(:ref, :i, k)
        if i[k] === Colon
            Nnew += 1
            push!(inds.args, Expr(:ref, :o, k))
            push!(s.args, Expr(:ref, :s, k))
            push!(x.args, Expr(:ref, :x, k))
            push!(o.args, zero_offsets ? :(Zero()) : :(One()))
            if k == C
                Cnew = Nnew
            end
            if k == B
                Bnew = Nnew
            end
            push!(Rnew.args, R[k])
        else
            push!(inds.args, Expr(:call, :first, iₖ))
            if i[k] <: AbstractRange
                Nnew += 1
                push!(s.args, Expr(:call, :static_length, iₖ))
                push!(x.args, Expr(:ref, :x, k))
                push!(o.args, zero_offsets ? :(Zero()) : :(One()))
                if k == C
                    Cnew = Nnew
                end
                if k == B
                    Bnew = Nnew
                end
                push!(Rnew.args, R[k])
            end
        end
        spₙ = sortp[k]
        still_dense &= D[spₙ]
        if still_dense# && (D[spₙ])
            ispₙ = i[spₙ]
            still_dense = (ispₙ <: AbstractUnitRange) || (ispₙ === Colon)
            densev[spₙ] = still_dense
            if still_dense
                still_dense = if ((ispₙ === Colon)::Bool || (ispₙ <: Base.Slice)::Bool)
                    true
                else
                    ispₙ_len = ArrayInterface.known_length(ispₙ)
                    if ispₙ_len !== nothing
                        _sz = getfield(S, :parameters)[spₙ]
                        if _sz <: StaticInt
                            ispₙ_len == getfield(_sz, :parameters)[1]
                        else
                            false
                        end
                    else
                        false
                    end
                    false
                end
            end
        else
            densev[spₙ] = false
        end
    end
    for k ∈ 1:K
        iₖt = i[k]
        if (iₖt === Colon) || (iₖt <: AbstractVector)
            push!(Dnew.args, densev[k])
        end
    end
    quote
        $(Expr(:meta,:inline))
        sp = A.ptr
        s = A.size
        x = sp.strd
        o = sp.offsets
        new_sp = StridedPointer{$T,$Nnew,$Cnew,$Bnew,$Rnew}(gep(sp, $inds), $x, $o)
        PtrArray(new_sp, $s, Val{$Dnew}())
    end
end

@generated function Base.view(A::PtrArray{S,D,T,N,C,B,R,X,O}, i::Vararg{Union{Integer,AbstractRange,Colon},K}) where {K,S,D,T,N,C,B,R,X,O}
    view_quote(i, K, S, D, T, N, C, B, R, X, O)
end
@generated function zview(A::PtrArray{S,D,T,N,C,B,R,X,O}, i::Vararg{Union{Integer,AbstractRange,Colon},K}) where {K,S,D,T,N,C,B,R,X,O}
    view_quote(i, K, S, D, T, N, C, B, R, X, O, true)
end

@inline function Base.vec(A::PtrArray{S,D,T,N,C,0}) where {S,D,T,N,C}
    @assert all(D) "All dimensions must be dense for a vec view. Try `vec(copy(A))` instead."
    sp = StridedPointer(pointer(A), (VectorizationBase.static_sizeof(T),), (One(),))
    PtrArray(sp, (static_length(A),), Val((true,)))
end


module PLext

using Reexport

#include("simulator/statevector.jl")
include("simulator/statevector_ptr.jl")
#include("simulator/statevector_copy.jl")
#include("simulator/statevector_ptr_dist.jl")
#include("compile/sysimg.jl")
#include("compile/precompile.jl")
#include("compile/library.jl")

#@reexport using PLext.StateVectorP
#@reexport using PLext.StateVectorC

end


using FileIO
using CSV
using DataFrames
using Plots
using Measures      # for better positioning of figures

# if you have python installed, you can use pyplot for better formatting
pyplot()

include("read_in_cases.jl")
include("processing_sim.jl")

function plot_all_forecasts(
    file_date,
    states_to_plot,
    local_case_dict;
    zoom=false,
    confidence_level="both"
)
    """
    Plots all the forecasts and saves the result as a pdf.
    """

    # read in the sim results and kept TP's
    sim_all_states = CSV.read(
        "results/UoA_forecast_output/"*file_date*"/UoA_samples_"*file_date*".csv",
        DataFrame
    )
    TP_all_states = CSV.read(
        "results/UoA_forecast_output/"*file_date*"/UoA_TP_"*file_date*".csv",
        DataFrame
    )

    # read in the case data
    case_dates = collect(local_case_dict["date"])

    # these indices are for the relevant plots per state
    case_plot_inds = [1,2,5,6,9,10,13,14]
    TP_plot_inds = [3,4,7,8,11,12,15,16]

    # these are heights of the relevant plots
    heights = [0.4, 0.2, 0.4 , 0.2, 0.4 , 0.2, 0.4 , 0.2]
    # normalize the heights for a good fit
    height_norm = heights / sum(heights)
    # layer the heights so that we get the full length of the plot
    heights_col = [height_norm; height_norm]

    # define a custom 8x2 layout with the chosen heights
    l = @layout [
        grid(8, 2, heights=heights_col)
    ]

    # initialise the plot
    fig = plot(
        layout=l,
        dpi=200,
        size=(750,1200),
        link=:x,
        framestyle=:box,
        margin=3mm,
        rightmargin=4mm,
    )

    for (i,state) in enumerate(states_to_plot)

        # take indices for plots from the defined index arrays
        c = case_plot_inds[i]
        tp = TP_plot_inds[i]

        # get relevant columns of the dataframe and merge to matrix
        D = Matrix(sim_all_states[sim_all_states.state .== state, 3:end-1])
        # if anything in D is missing, it means the sims didn't work so 0 out those entries
        D[ismissing.(D)] .= 0
        D = Matrix{Int}(D)

        TP_local = Matrix(TP_all_states[TP_all_states.state .== state, 3:end-1])
        TP_local[ismissing.(TP_local)] .= 0.0
        TP_local = Matrix{Float64}(TP_local)

        onset_dates = sim_all_states[sim_all_states.state .== state, "onset date"]
        df_D_summary = summarise_forecast_for_plotting(D)
        df_TP_summary = summarise_forecast_for_plotting(TP_local)

        dates = local_case_dict["date"]
        local_cases = local_case_dict[state][dates .>= onset_dates[1]]
        # boolean for correct case dates
        case_dates_ind = [d ∈ onset_dates ? true : false for d in case_dates]

        onset_dates_lims = Dates.value.([onset_dates[end]-Dates.Day(75), onset_dates[end]])

        if confidence_level == "50" || confidence_level == "both"
            plot!(
                fig,
                title=state,
                subplot=c,
                onset_dates,
                df_D_summary[!,"median"],
                xaxis=nothing,
                linecolor=1,
                linewidth=1,
                color=1,
                fillalpha=0.2,
                label=false,
            )
            plot!(
                fig,
                subplot=c,
                onset_dates,
                [df_D_summary[!,"median"] df_D_summary[!,"median"]],
                fillrange=[df_D_summary[!,"lower"] df_D_summary[!,"upper"]],
                fillalpha=0.4,
                c=1,
                label=["50%" nothing]
            )

        end

        if confidence_level == "95" || confidence_level == "both"
            plot!(
                fig,
                title=state,
                subplot=c,
                onset_dates,
                df_D_summary[!,"median"],
                xaxis=nothing,
                linecolor=1,
                linewidth=1,
                color=1,
                fillalpha=0.2,
                label=false,
            )

            plot!(
                fig,
                subplot=c,
                onset_dates,
                [df_D_summary[!,"median"] df_D_summary[!,"median"]],
                fillrange=[df_D_summary[!,"bottom"] df_D_summary[!,"top"]],
                fillalpha=0.2,
                c=1,
                label=["95%" nothing]
            )
        end

        bar!(
            fig,
            subplot=c,
            case_dates[case_dates_ind],
            local_cases,
            lc="gray",
            fc="gray",
            fillalpha=0.4,
            linealpha=0.4,
            label=false,
        )

        if confidence_level == "50" || confidence_level == "both"
            plot!(
                fig,
                subplot=tp,
                onset_dates,
                df_TP_summary[!,"median"],
                linecolor=1,
                linewidth=1,
                color=1,
                fillalpha=0.2,
                label=false,
            )

            plot!(
                fig,
                subplot=tp,
                onset_dates,
                [df_TP_summary[!,"median"] df_TP_summary[!,"median"]],
                fillrange=[df_TP_summary[!,"lower"] df_TP_summary[!,"upper"]],
                fillalpha=0.4,
                c=1,
                label=false
            )
        end

        if confidence_level == "95" || confidence_level == "both"
            plot!(
                fig,
                subplot=tp,
                onset_dates,
                df_TP_summary[!,"median"],
                linecolor=1,
                linewidth=1,
                color=1,
                fillalpha=0.2,
                label=false,
            )

            plot!(
                fig,
                subplot=tp,
                onset_dates,
                [df_TP_summary[!,"median"] df_TP_summary[!,"median"]],
                fillrange=[df_TP_summary[!,"bottom"] df_TP_summary[!,"top"]],
                fillalpha=0.2,
                c=1,
                label=false
            )
        end

        hline!(
            fig,
            subplot=tp,
            [1],
            lc=:black,
            ls=:dash,
            label=false,
        )

        vline!(
            fig,
            subplot=c,
            [Dates.Date(file_date)],
            lc=:gray,
            ls=:dash,
            label=false,
        )

        vline!(
            fig,
            subplot=tp,
            [Dates.Date(file_date)],
            lc=:gray,
            ls=:dash,
            label=false,
        )

        xlims!(fig, subplot=c, onset_dates_lims...)
        xlims!(fig, subplot=tp, onset_dates_lims...)
        ylims!(fig, subplot=tp, 0, 1.25*maximum(df_TP_summary[!,"top"]))
        if zoom
            ylims!(fig, subplot=c, 0, 1.1*maximum(local_cases))
        end

    end

    dir_name = joinpath("figs", "case_forecasts", file_date)
    if !ispath(dir_name)
        mkpath(dir_name)
    end

    if zoom
        savefig(fig, dir_name*"/UoA_forecast_"*file_date*"_"*"zoomed_"*confidence_level*"_intervals.pdf")
    else
        savefig(fig, dir_name*"/UoA_forecast_"*file_date*"_"*confidence_level*"_intervals.pdf")
    end

    return nothing
end

function plot_single_forecast(
    file_date,
    state,
    local_case_dict;
    confidence_level="both"
)
    """
    Plots all the forecasts and saves the result as a pdf.
    """

    # read in the sim results and kept TP's
    sim_all_states = CSV.read(
        "results/UoA_forecast_output/2022-01-04/UoA_samples_"*file_date*".csv",
        DataFrame
    )
    TP_all_states = CSV.read(
        "results/UoA_forecast_output/2022-01-04/UoA_TP_"*file_date*".csv",
        DataFrame
    )

    # read in the case data
    case_dates = collect(local_case_dict["date"])

    # these are heights of the relevant plots
    heights = [0.4, 0.2]
    # normalize the heights for a good fit
    height_norm = heights / sum(heights)

    # define a custom 8x2 layout with the chosen heights
    l = @layout [
        grid(2, 1, heights=heights)
    ]

    # initialise the plot
    fig = plot(
        legend=false,
        layout = l,
        dpi=144,
        size=(700,1200),
        link=:x,
        framestyle=:box
    )

    c = 1
    tp = 2

    # get relevant columns of the dataframe and merge to matrix
    D = Matrix(sim_all_states[sim_all_states.state .== state, 3:end-1])
    # if anything in D is missing, it means the sims didn't work so 0 out those entries
    D[ismissing.(D)] .= 0
    D = Matrix{Int}(D)

    TP_local = Matrix(TP_all_states[TP_all_states.state .== state, 3:end-1])
    TP_local[ismissing.(TP_local)] .= 0.0
    TP_local = Matrix{Float64}(TP_local)

    onset_dates = sim_all_states[sim_all_states.state .== state, "onset date"]
    df_D_summary = summarise_forecast_for_plotting(D)
    df_TP_summary = summarise_forecast_for_plotting(TP_local)

    dates = local_case_dict["date"]
    local_cases = local_case_dict[state][dates .>= onset_dates[1]]
    # boolean for correct case dates
    case_dates_ind = [d ∈ onset_dates ? true : false for d in case_dates]

    onset_dates_lims = Dates.value.([onset_dates[end]-Dates.Day(75), onset_dates[end]])

    annotate!(
        fig,
        subplot=c,
        (0.2,0.8),
        text(state, 8)
    )

    if confidence_level == "50" || confidence_level == "both"
        plot!(
            fig,
            subplot=c,
            onset_dates,
            df_D_summary[!,"median"],
            xaxis=nothing,
            linecolor=1,
            linewidth=2,
            ribbon=(
                vec(df_D_summary[!,"median"]-df_D_summary[!,"lower"]),
                vec(df_D_summary[!,"upper"] - df_D_summary[!,"median"])
            ),
            color=1,
            fillalpha=0.2,
        )
    end

    if confidence_level == "95" || confidence_level == "both"
        plot!(
            fig,
            subplot=c,
            onset_dates,
            df_D_summary[!,"median"],
            xaxis=nothing,
            linecolor=1,
            linewidth=2,
            ribbon=(
                vec(df_D_summary[!,"median"]-df_D_summary[!,"bottom"]),
                vec(df_D_summary[!,"top"] - df_D_summary[!,"median"])
            ),
            color=1,
            fillalpha=0.2
        )
    end

    bar!(
        fig,
        subplot=c,
        case_dates[case_dates_ind],
        local_cases,
        lc="gray",
        fc="gray",
        fillalpha=0.4,
        linealpha=0.4,
    )

    if confidence_level == "50" || confidence_level == "both"
        plot!(
            fig,
            subplot=tp,
            onset_dates,
            df_TP_summary[!,"median"],
            linecolor=1,
            linewidth=2,
            ribbon=(
                vec(df_TP_summary[!,"median"]-df_TP_summary[!,"lower"]),
                vec(df_TP_summary[!,"upper"] - df_TP_summary[!,"median"])
            ),
            color=1,
            fillalpha=0.2
        )
    end

    if confidence_level == "95" || confidence_level == "both"
        plot!(
            fig,
            subplot=tp,
            onset_dates,
            df_TP_summary[!,"median"],
            linecolor=1,
            linewidth=2,
            ribbon=(
                vec(df_TP_summary[!,"median"]-df_TP_summary[!,"bottom"]),
                vec(df_TP_summary[!,"top"] - df_TP_summary[!,"median"])
            ),
            color=1,
            fillalpha=0.2,
        )
    end

    hline!(
        fig,
        subplot=tp,
        [1],
        lc=:black,
        ls=:dash
    )

    xlims!(fig, subplot=c, onset_dates_lims...)
    xlims!(fig, subplot=tp, onset_dates_lims...)
    ylims!(fig, subplot=tp, 0, 1.25*maximum(df_TP_summary[!,"top"]))

    dir_name = joinpath("figs", file_date)
    if !ispath(dir_name)
        mkpath(dir_name)
    end

    savefig(fig, dir_name*"/UoA_forecast_"*file_date*"_"*state*"_"*confidence_level*"_intervals.pdf")

    return nothing

end