# SPD-License-Identifier: MIT

using ArgParse
using JLD2
using MarkovModels
using Random

function loadlexicon(path)
    retval = Dict()
    open(path, "r") do f
        for line in eachline(f)
            tokens = split(line)
            word, pronun = tokens[1], tokens[2:end]
            pronuns = push!(get(retval, word, []), pronun)
            retval[word] = pronuns
        end
    end
    retval
end

function rand_addsil(sentence, silsym, between_silprob, edge_silprob)
    retval = rand() < edge_silprob ? [silsym] : []
    for word in sentence[1:end-1]
        push!(retval, word)
        if rand() < between_silprob
            push!(retval, silsym)
        end
    end
    push!(retval, sentence[end])
    if rand() < edge_silprob
        push!(retval, silsym)
    end
    retval
end

function main(args)
    oovsym = args["oov-sym"]
    silsym = args["sil-sym"]
    between_silprob = args["between-silprob"]
    edge_silprob = args["edge-silprob"]
    lexicon = loadlexicon(args["lexicon"])
    n = args["ngram-order"]
    hmms = load(args["hmms"])

    ###################################################################
    # Estimate the n-gram count without pruning

    @info "Estimating $(n)-gram phonotactic language model."

    counts = Dict()
    for line in eachline(stdin)
        tokens = split(line)
        tokens = rand_addsil(tokens, silsym, between_silprob, edge_silprob)
        seq = []
        for token in tokens
            if token ∈ keys(lexicon)
                push!(seq, lexicon[token][1])
            else
                push!(seq, lexicon[oovsym][1])
            end
        end

        # Append nothing to mark the end of sentence.
        seq = vcat(seq..., [nothing])

        for i in 1:length(seq)
            src = tuple(seq[max(1, i-n+1):i-1]...)
            dest = if n > 1
                tuple(seq[max(1, i-n+2):i]...)
            else
                tuple(seq[i])
            end
            counts[(src, dest)] = get(counts, (src, dest), 0) + 1
        end
    end

    ###################################################################
    # Build an FSM from the counts

    @info "Representing the language model as a FSM."

    SF = LogSemifield{Float32}
    fsm = VectorFSM{SF}()

    # Get the init and final weights
    iws = Dict()
    fws = Dict()
    for ((src, dest), count) in counts
        if isempty(src) iws[dest] = SF(log(count)) end
        if isnothing(dest[end]) fws[src] = SF(log(count)) end
    end

    # Create the states.
    smap = Dict()
    for ((src, dest), count) in counts
        if ! isnothing(dest[end]) && dest ∉ keys(smap)
            initweight = get(iws, dest, zero(SF))
            finalweight = get(fws, dest, zero(SF))

            # This will be nonzero only in the unigram case.
            finalweight += get(fws, tuple(), zero(SF))

            smap[dest] = addstate!(fsm, dest; initweight, finalweight)
        end
    end

    # Create the arcs.
    for ((src, dest), count) in counts
        # Not an initial / final pseudo-arc.
        if ! isnothing(src) && ! isnothing(dest[end])
            if src == tuple() && n == 1
                # special case for unigram probabilty:
                # we use the same weight for each possible src states
                for s in values(smap)
                    addarc!(fsm, s, smap[dest], SF(log(count)))
                end
            elseif src ≠ tuple()
                addarc!(fsm, smap[src], smap[dest], SF(log(count)))
            end
        end
    end

    fsm = fsm

    ###################################################################
    # Compile and save

    @info "Composing the language model with units' HMM."

    # The previously built FSM combines the grammar and the lexicon.
    GL = fsm |> renormalize
    H = hmms["units"]
    GLH = HierarchicalFSM(GL, H)

    @info "Compiling."

    # Compile the expanded FSM into a matrix format.
    mfsm = MatrixFSM(GLH, hmms["pdfid_mapping"], t -> (t[end-1], t[end]))

    save(args["denfsm"], Dict("fsm" => mfsm))
end

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--oov-sym"
            default = "<UNK>"
            help = "fallback for out of vocabulary words"
        "--sil-sym"
            default = "<SIL>"
            help = "silence word symbol"
        "--edge-silprob"
            default = 0.8
            arg_type = Float64
            help = "probability of having a silence at the edge of the utterance"
        "--between-silprob"
            default = 0.1
            arg_type = Float64
            help = "probability of having a silence between two words"
        "--ngram-order"
            default = 3
            arg_type = Int
            help = "ngram order of the language model"
        "hmms"
            required = true
            help = "HMM FSMs in JLD2 format"
        "lexicon"
            required = true
            help = "lexicon text file"
        "denfsm"
            required = true
            help = "output denominator fsm in JLD2 format"
    end
    s.description = """
    Build a denominator graph based on a phonetic language model.
    This script read the text from its standard input.
    """
    parse_args(s)
end

args = parse_commandline()
main(args)


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct BaseMachine <: Machine
        R::Float64
        Xd_p::Float64
        eq_p::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
        internal::InfrastructureSystemsInternal
    end

Parameters of a Classic Machine: GENCLS in PSSE and PSLF

# Arguments
- `R::Float64`: Resistance after EMF in machine per unit, validation range: `(0, nothing)`
- `Xd_p::Float64`: Reactance after EMF in machine per unit, validation range: `(0, nothing)`
- `eq_p::Float64`: Fixed EMF behind the impedance, validation range: `(0, nothing)`
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: BaseMachine has no states
- `n_states::Int`: BaseMachine has no states
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct BaseMachine <: Machine
    "Resistance after EMF in machine per unit"
    R::Float64
    "Reactance after EMF in machine per unit"
    Xd_p::Float64
    "Fixed EMF behind the impedance"
    eq_p::Float64
    ext::Dict{String, Any}
    "BaseMachine has no states"
    states::Vector{Symbol}
    "BaseMachine has no states"
    n_states::Int
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function BaseMachine(R, Xd_p, eq_p, ext=Dict{String, Any}(), )
    BaseMachine(R, Xd_p, eq_p, ext, Vector{Symbol}(), 0, InfrastructureSystemsInternal(), )
end

function BaseMachine(; R, Xd_p, eq_p, ext=Dict{String, Any}(), states=Vector{Symbol}(), n_states=0, internal=InfrastructureSystemsInternal(), )
    BaseMachine(R, Xd_p, eq_p, ext, states, n_states, internal, )
end

# Constructor for demo purposes; non-functional.
function BaseMachine(::Nothing)
    BaseMachine(;
        R=0,
        Xd_p=0,
        eq_p=0,
        ext=Dict{String, Any}(),
    )
end

"""Get [`BaseMachine`](@ref) `R`."""
get_R(value::BaseMachine) = value.R
"""Get [`BaseMachine`](@ref) `Xd_p`."""
get_Xd_p(value::BaseMachine) = value.Xd_p
"""Get [`BaseMachine`](@ref) `eq_p`."""
get_eq_p(value::BaseMachine) = value.eq_p
"""Get [`BaseMachine`](@ref) `ext`."""
get_ext(value::BaseMachine) = value.ext
"""Get [`BaseMachine`](@ref) `states`."""
get_states(value::BaseMachine) = value.states
"""Get [`BaseMachine`](@ref) `n_states`."""
get_n_states(value::BaseMachine) = value.n_states
"""Get [`BaseMachine`](@ref) `internal`."""
get_internal(value::BaseMachine) = value.internal

"""Set [`BaseMachine`](@ref) `R`."""
set_R!(value::BaseMachine, val) = value.R = val
"""Set [`BaseMachine`](@ref) `Xd_p`."""
set_Xd_p!(value::BaseMachine, val) = value.Xd_p = val
"""Set [`BaseMachine`](@ref) `eq_p`."""
set_eq_p!(value::BaseMachine, val) = value.eq_p = val
"""Set [`BaseMachine`](@ref) `ext`."""
set_ext!(value::BaseMachine, val) = value.ext = val


module MonteCarlo

using Reexport
@reexport using MonteCarloObservable
using StableDQMC, LightXML, Parameters

using Printf, SparseArrays, LinearAlgebra, Dates, Random

include("helpers.jl")
include("abstract.jl")

include("flavors/MC/MC.jl")
include("flavors/DQMC/DQMC.jl")
include("flavors/DQMC_bond/DQMC_bond.jl")

include("lattices/square.jl")
include("lattices/chain.jl")
include("lattices/cubic.jl")
include("lattices/ALPS.jl")

include("models/Ising/IsingModel.jl")
include("models/HubbardAttractive/HubbardModelAttractive.jl")
include("models/HoppingBF/HoppingBFModel.jl")

include("../test/testfunctions.jl")

export reset!
export run!
export IsingModel
export HubbardModelAttractive
export HoppingBFModel
export MC
export DQMC
export DQMC_bond
export greens
export observables

export initialize_stack #for test
export build_stack #for test
export calculate_greens#for test
export  propagate #for test
#export add_slice_sequence_left #for test
export multiply_slice_matrix_left!
end # module


# RUN: julia -e "import Brutus; Brutus.lit(:emit_lowered)" --startup-file=no %s 2>&1 | FileCheck %s

add(x, y) = x + y
emit(add, Float64, Float64)

# CHECK: Core.MethodMatch(Tuple{typeof(Main.Main.add), Float64, Float64}, svec(), add(x, y) in Main.Main at /{{.*}}/test/Codegen/lower/add_float64.jl:3, true)after translating to MLIR in JLIR dialect:module  {
# CHECK-NEXT:   func nested @"Tuple{typeof(Main.add), Float64, Float64}"(%arg0: !jlir<"typeof(Main.add)">, %arg1: !jlir.Float64, %arg2: !jlir.Float64) -> !jlir.Float64 attributes {llvm.emit_c_interface} {
# CHECK-NEXT:     "jlir.goto"()[^bb1] : () -> ()
# CHECK-NEXT:   ^bb1:  // pred: ^bb0
# CHECK-NEXT:     %0 = "jlir.constant"() {value = #jlir<"#<intrinsic #12 add_float>">} : () -> !jlir<"typeof(Core.IntrinsicFunction)">
# CHECK-NEXT:     %1 = "jlir.call"(%0, %arg1, %arg2) : (!jlir<"typeof(Core.IntrinsicFunction)">, !jlir.Float64, !jlir.Float64) -> !jlir.Float64
# CHECK-NEXT:     "jlir.return"(%1) : (!jlir.Float64) -> ()
# CHECK-NEXT:   }
# CHECK-NEXT: }

# CHECK: module  {
# CHECK-NEXT:   func nested @"Tuple{typeof(Main.add), Float64, Float64}"(%arg0: !jlir<"typeof(Main.add)">, %arg1: f64, %arg2: f64) -> f64 attributes {llvm.emit_c_interface} {
# CHECK-NEXT:     %0 = addf %arg1, %arg2 : f64
# CHECK-NEXT:     return %0 : f64
# CHECK-NEXT:   }
# CHECK-NEXT: }

# CHECK:   llvm.func @"Tuple{typeof(Main.add), Float64, Float64}"(%arg0: !llvm.ptr<struct<"struct_jl_value_type", opaque>>, %arg1: f64, %arg2: f64) -> f64 attributes {llvm.emit_c_interface, sym_visibility = "nested"} {
# CHECK-NEXT:     %0 = llvm.fadd %arg1, %arg2  : f64
# CHECK-NEXT:     llvm.return %0 : f64
# CHECK-NEXT:   }
# CHECK-NEXT:   llvm.func @"_mlir_ciface_Tuple{typeof(Main.add), Float64, Float64}"(%arg0: !llvm.ptr<struct<"struct_jl_value_type", opaque>>, %arg1: f64, %arg2: f64) -> f64 attributes {llvm.emit_c_interface, sym_visibility = "nested"} {
# CHECK-NEXT:     %0 = llvm.call @"Tuple{typeof(Main.add), Float64, Float64}"(%arg0, %arg1, %arg2) : (!llvm.ptr<struct<"struct_jl_value_type", opaque>>, f64, f64) -> f64
# CHECK-NEXT:     llvm.return %0 : f64
# CHECK-NEXT:   }
# CHECK-NEXT: }


using CycUtils: multinomial
using Iterators: chain, partition
using JLD: jldopen

type SampleFactory
    material :: String
    amount :: UInt64
    consumed :: Ref{UInt64}
    neg_prob :: Vector{Float32}
    producer :: Function
    corpus_type :: Union{Type{String}, Type{UInt32}}
    function SampleFactory(material::String, neg_prob::Vector{Float32})
        global batch_producer
        global line_batch_producer
        if endswith(material, ".idx")
            amt = open(material) do f
                seekend(f)
                position(f)
            end
            return new(material, amt, Ref{UInt64}(0), neg_prob, batch_producer,
                       UInt32)
        end

        if !endswith(material, ".bz2")
            error("only accept bzip2 compressed plain text as corpus")
        end
        dn = dirname(material)
        fn = basename(material)[1:end-4]
        mn = joinpath(dn, fn * "-meta.jld")
        if isfile(mn)
            lines = jldopen(mn) do f
                read(f, "lines")
            end
        else
            nl, nw, nb = split(readstring(pipeline(`bzcat $material`, `wc`)))
            nl = parse(Int, nl)
            nw = parse(Int, nw)
            nb = parse(Int, nb)
            lines = nl
            jldopen(mn, "w") do f
                write(f, "lines", lines)
                write(f, "words", nw)
                write(f, "bytes", nb)
            end
        end
        return new(material, lines, Ref{UInt64}(0), neg_prob,
                   line_batch_producer, String)
    end
end

progress(factory::SampleFactory) = factory.consumed[]/factory.amount

const sampleFactory = Ref{Nullable{SampleFactory}}(Nullable{SampleFactory}())
const genBufFac = Ref{UInt32}(0xa)

function line_sampler(factory, m, k, word2idx, window=6)
    wsize = 1:window
    fn = factory.material
    lines = eachline(`bzcat $fn`)
    for line in lines
        sent = [word2idx[w] for w in split(line) if haskey(word2idx, w)]
        lsent = length(sent)
        for i = 1:lsent
            w = rand(wsize)
            st = max(1, i-w)
            en = min(lsent, i+w)

            for j in chain(st:i-1, i+1:en)
                produce((sent[i], sent[j]))
            end
        end
        factory.consumed[] += 1
    end
end

function line_batch_producer(factory, m, k)
    global init_model
    global genBufFac
    const multiplier = get(genBufFac)
    draw! = multinomial(factory.neg_prob)
    const nb_samples = m * multiplier
    idxes = collect(1:nb_samples)

    bigbatch_w = zeros(UInt32, nb_samples)
    bigbatch_c = zeros(UInt32, k+1, nb_samples)

    word2idx = jldopen(init_model) do f
        read(f, "word2idx")
    end
    corpus = Task(()->line_sampler(factory, m, k, word2idx))
    for x in partition(corpus, nb_samples)
        for (i, wc) in enumerate(x)
            w, c = wc
            bigbatch_w[i] = w
            bigbatch_c[1, i] = c
        end
        draw!(@view bigbatch_c[2:end, :])
        shuffle!(idxes)
        bigbatch_w = bigbatch_w[idxes]
        bigbatch_c = bigbatch_c[:, idxes]
        for st in 1:m:nb_samples
            en = st + m - 1
            produce((bigbatch_w[st:en], bigbatch_c[:, st:en]))
        end
    end
end

function batch_producer(factory, m, k)
    global genBufFac
    const multiplier = get(genBufFac)
    draw! = multinomial(factory.neg_prob)
    const nb_samples = m * multiplier
    idxes = collect(1:(2*nb_samples))
    buf = zeros(UInt32, (2, length(idxes)))
    const nbblk = sizeof(UInt32) * 2 * nb_samples

    open(factory.material) do f
        nbtotoal = factory.amount

        for i = 1:div(nbtotoal, nbblk)
            buf_pos = read(f, UInt32, (2, nb_samples))
            factory.consumed[] += nbblk
            buf[:, 1:nb_samples] = buf_pos
            buf[1, nb_samples+1:end] = buf_pos[2, :]
            buf[2, nb_samples+1:end] = buf_pos[1, :]
            shuffle!(idxes)
            buf_o = buf
            buf = reshape(buf_o[:, idxes], (2, m, 2*multiplier))
            for kk = 1:2*multiplier
                c = zeros(UInt32, (k+1, m))
                w = buf[1, :, kk]
                c[1, :] = buf[2, :, kk]
                draw!(@view c[2:end, :])
                produce((w, c))
            end
            buf = buf_o
        end
    end
    nothing
end


module GasDispersion

export Scenario, plume, puff, scenario_builder

using Interpolations

# helpful utilities
include("utils/scenario.jl")
include("utils/scenario_builder.jl")
include("utils/utils.jl")

# plume models
include("models/gaussian_plume.jl")
include("models/britter_mcquaid_plume.jl")

# puff models
include("models/gaussian_puff.jl")
include("models/britter_mcquaid_puff.jl")

"""
    plume(scenario::Scenario; model=:gaussian, kwargs...)

Runs the plume dispersion model on the given scenario and returns a function
giving the concentration of the form
    c(x, y, z[, t])

If `model` is unspecified, defaults to gaussian, `kwargs` are passed to the
plume model. All model parameters are assumed to be in SI base units (i.e.
distances in m, velocities in m/s, mass in kg, etc.)
"""
function plume(scenario::Scenario; model=:gaussian, kwargs...)
    if model==:gaussian
        return gaussian_plume_factory(scenario; kwargs...)
    elseif model==:brittermcquaid
        return britter_plume_factory(scenario; kwargs...)
    else
        error_string = "plume dispersion model $model is not currently implemented"
        error(error_string)
    end
end

"""
    puff(scenario::Scenario; model=:gaussian, kwargs...)

Runs the puff dispersion model on the given scenario and returns a function
giving the concentration of the form
    c(x, y, z, t)

If `model` is unspecified, defaults to gaussian, `kwargs` are passed to the
puff model. All model parameters are assumed to be in SI base units (i.e.
distances in m, velocities in m/s, mass in kg, etc.)
"""
function puff(scenario::Scenario; model=:gaussian, kwargs...)
    if model==:gaussian
        return gaussian_puff_factory(scenario; kwargs...)
    # elseif model=="britter-mcquaid"
    #     return britter_plume_factory(scenario; kwargs...)
    else
        error_string = "puff dispersion model $model is not currently implemented"
        error(error_string)
    end
end

end


module GenomeFragments

## This is a module for handling high throughput sequencing reads

using GenomicFeatures, DataStructures, BioAlignments, Mmap, UnalignedVectors, XAM, IterTools, ProgressMeter

export FragMatrix, FragMatrixPair, FragMatrixSingle, get_frags, chrom_lt, load_frag_matrix, totalfrags, streambam, qualityfilt, filterpairfirst, filterpairfirstpp,
       barcode, ishighquality, isunique,
       inc_meta_frag!, inc_meta_cut!, inc_meta_mid!,
       inc_heat_frag!, inc_heat_cut!, inc_heat_mid!,
       inc_heat_atac_cut!, inc_heat_atac_frag!, inc_meta_atac_frag!,
       inc_meta_mid_width, inc_heat_mid_width,
       inc_heat_atac_cut_width, inc_meta_atac_cut_width,
       readsintersecting, fragregion,
       fragheatmap, fragheatmap!,
       streambam_singlecell

include("fragmatrix.jl")
include("buildfm.jl")
include("processing.jl")

end # module


using GLFW, ModernGL

# whenever the window size changed (by OS or user resize) this callback function executes
function framebuffer_size_callback(window::GLFW.Window, width::Cint, height::Cint)
    # make sure the viewport matches the new window dimensions; note that width and
    # height will be significantly larger than specified on retina displays.
	glViewport(0, 0, width, height)
end

# process all input: query GLFW whether relevant keys are pressed/released this frame and react accordingly
function processInput(window::GLFW.Window)
    GLFW.GetKey(window, GLFW.KEY_ESCAPE) == GLFW.PRESS && GLFW.SetWindowShouldClose(window, true)
end

# misc. config
const SCR_WIDTH = 800
const SCR_HEIGHT = 600

const vertexShaderSource = """
	#version 330 core
	layout (location = 0) in vec3 aPos;
	void main(void)
	{
	    gl_Position = vec4(aPos.x, aPos.y, aPos.z, 1.0);
	}"""

const fragmentShaderSource = """
	#version 330 core
	out vec4 FragColor;
	void main(void)
	{
	    FragColor = vec4(1.0f, 0.5f, 0.2f, 1.0f);
	}"""

GLFW.WindowHint(GLFW.CONTEXT_VERSION_MAJOR, 3)
GLFW.WindowHint(GLFW.CONTEXT_VERSION_MINOR, 3)
GLFW.WindowHint(GLFW.OPENGL_PROFILE, GLFW.OPENGL_CORE_PROFILE)
GLFW.WindowHint(GLFW.OPENGL_FORWARD_COMPAT, GL_TRUE)

# create window
window = GLFW.CreateWindow(SCR_WIDTH, SCR_HEIGHT, "LearnOpenGL")
window == C_NULL && error("Failed to create GLFW window.")
GLFW.MakeContextCurrent(window)
GLFW.SetFramebufferSizeCallback(window, framebuffer_size_callback)

# build and compile shader program
vertexShader = glCreateShader(GL_VERTEX_SHADER)
glShaderSource(vertexShader, 1, Ptr{GLchar}[pointer(vertexShaderSource)], C_NULL)
glCompileShader(vertexShader)
successRef = Ref{GLint}(-1)
glGetShaderiv(vertexShader, GL_COMPILE_STATUS, successRef)
if successRef[] != GL_TRUE
    infoLog = Vector{GLchar}(512)
    glGetShaderInfoLog(vertexShader, 512, C_NULL, infoLog)
    error("ERROR::SHADER::VERTEX::COMPILATION_FAILED\n", String(infoLog))
end

fragmentShader = glCreateShader(GL_FRAGMENT_SHADER)
glShaderSource(fragmentShader, 1, Ptr{GLchar}[pointer(fragmentShaderSource)], C_NULL)
glCompileShader(fragmentShader)
glGetShaderiv(vertexShader, GL_COMPILE_STATUS, successRef)
if successRef[] != GL_TRUE
    infoLog = Vector{GLchar}(512)
    glGetShaderInfoLog(fragmentShader, 512, C_NULL, infoLog)
    error("ERROR::SHADER::FRAGMENT::COMPILATION_FAILED\n", String(infoLog))
end

# link shaders
shaderProgram = glCreateProgram()
glAttachShader(shaderProgram, vertexShader)
glAttachShader(shaderProgram, fragmentShader)
glLinkProgram(shaderProgram)
glGetProgramiv(shaderProgram, GL_LINK_STATUS, successRef)
if successRef[] != GL_TRUE
    infoLog = Vector{GLchar}(512)
    glGetProgramInfoLog(shaderProgram, 512, C_NULL, infoLog)
    error("ERROR::SHADER::PROGRAM::LINKING_FAILED\n", String(infoLog))
end
glDeleteShader(vertexShader)
glDeleteShader(fragmentShader)

# set up vertex data (and buffer(s)) and configure vertex attributes
vertices = GLfloat[-0.5,  0.5, 0.0,   # left
                    0.5, -0.5, 0.0,   # right
                    0.0,  0.5, 0.0]   # top
vboRef = Ref{GLuint}(0)
vaoRef = Ref{GLuint}(0)
glGenVertexArrays(1, vaoRef)
glGenBuffers(1, vboRef)
vao = vaoRef[]
vbo = vboRef[]
glBindVertexArray(vao)
glBindBuffer(GL_ARRAY_BUFFER, vbo)
glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW)
glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(GLfloat), Ptr{Cvoid}(0))
glEnableVertexAttribArray(0)

# note that this is allowed, the call to glVertexAttribPointer registered VBO as
# the vertex attribute's bound vertex buffer object so afterwards we can safely unbind
glBindBuffer(GL_ARRAY_BUFFER, 0);

# You can unbind the VAO afterwards so other VAO calls won't accidentally modify this VAO,
# but this rarely happens. Modifying other VAOs requires a call to glBindVertexArray anyways
# so we generally don't unbind VAOs (nor VBOs) when it's not directly necessary.
glBindVertexArray(0)

# uncomment this call to draw in wireframe polygons.
# glPolygonMode(GL_FRONT_AND_BACK, GL_LINE)

# render loop
while !GLFW.WindowShouldClose(window)
    processInput(window)

    # render
    glClearColor(0.2, 0.3, 0.3, 1.0)
    glClear(GL_COLOR_BUFFER_BIT)

	# draw
	glUseProgram(shaderProgram)
    glBindVertexArray(vao)
	glDrawArrays(GL_TRIANGLES, 0, 3)

    # swap buffers and poll IO events
    GLFW.SwapBuffers(window)
    GLFW.PollEvents()
end

# optional: de-allocate all resources once they've outlived their purpose
glDeleteVertexArrays(1, vaoRef)
glDeleteBuffers(1, vboRef)

GLFW.DestroyWindow(window)


for (elty, ext) in ((:Float32, :s),
                    (:Float64, :d),
                    (:ComplexF32, :c),
                    (:ComplexF64, :z))

    # Distributed sparse gemv
    @eval begin
        function LinearAlgebra.mul!(y::DistMultiVec{$elty}, A::DistSparseMatrix{$elty}, x::DistMultiVec{$elty}, α::$elty, β::$elty)
            ElError(ccall(($(string("ElMultiplyDist_", ext)), libEl), Cuint,
                (Cint, $elty, Ptr{Cvoid}, Ptr{Cvoid}, $elty, Ptr{Cvoid}),
                NORMAL, α, A.obj, x.obj, β, y.obj))
            return y
        end
        function LinearAlgebra.mul!(y::DistMultiVec{$elty}, adjA::Adjoint{<:Any,DistSparseMatrix{$elty}}, x::DistMultiVec{$elty}, α::$elty, β::$elty)
            ElError(ccall(($(string("ElMultiplyDist_", ext)), libEl), Cuint,
                (Cint, $elty, Ptr{Cvoid}, Ptr{Cvoid}, $elty, Ptr{Cvoid}),
                ADJOINT, α, parent(adjA).obj, x.obj, β, y.obj))
            return y
        end
        function LinearAlgebra.mul!(y::DistMultiVec{$elty}, trA::Transpose{<:Any,DistSparseMatrix{$elty}}, x::DistMultiVec{$elty}, α::$elty, β::$elty)
            ElError(ccall(($(string("ElMultiplyDist_", ext)), libEl), Cuint,
                (Cint, $elty, Ptr{Cvoid}, Ptr{Cvoid}, $elty, Ptr{Cvoid}),
                TRANSPOSE, α, parent(trA).obj, x.obj, β, y.obj))
            return y
        end
    end
end


using .Networks

using Dates
using JuMP
using Printf
using Parameters

@with_kw struct TSOBilevel <: AbstractTSO
    configs::TSOBilevelConfigs = TSOBilevelConfigs()
end

function run(runnable::TSOBilevel, ech::Dates.DateTime, firmness, TS::Vector{Dates.DateTime}, context::AbstractContext)
    problem_name_l = @sprintf("tso_bilevel_%s", ech)

    runnable.configs.out_path = context.out_dir
    runnable.configs.problem_name = problem_name_l

    return tso_bilevel(get_network(context),
                    TS,
                    get_generators_initial_state(context),
                    get_scenarios(context, ech),
                    get_uncertainties(context, ech),
                    firmness,
                    get_market_schedule(context),
                    get_tso_schedule(context),
                    runnable.configs
                    )
end

function update_tso_schedule!(context::AbstractContext, ech, result::TSOBilevelModel, firmness,
                            runnable::TSOBilevel)
    tso_schedule = get_tso_schedule(context)
    tso_schedule.decider_type = DeciderType(runnable)
    tso_schedule.decision_time = ech

    # upper problem (TSO) locates limitable injections
    for ((gen_id, ts, s), p_injected_var) in result.upper.limitable_model.p_injected
        set_prod_value!(tso_schedule, gen_id, ts, s, value(p_injected_var))
    end
    # lower problem (Market) decides pilotable injections
    for ((gen_id, ts, s), p_injected_var) in result.lower.pilotable_model.p_injected
        if get_power_level_firmness(firmness, gen_id, ts) == FREE
            set_prod_value!(tso_schedule, gen_id, ts, s, value(p_injected_var))
        elseif get_power_level_firmness(firmness, gen_id, ts) == TO_DECIDE
            set_prod_definitive_value!(tso_schedule, gen_id, ts, value(p_injected_var))
        elseif get_power_level_firmness(firmness, gen_id, ts) == DECIDED
            @assert( ismissing(get_prod_value(tso_schedule, gen_id, ts))
                    || value(p_injected_var) ≈ get_prod_value(tso_schedule, gen_id, ts) )
            set_prod_definitive_value!(tso_schedule, gen_id, ts, value(p_injected_var))
        end
    end

    for ((gen_id, ts, s), b_on_var) in result.upper.pilotable_model.b_on
        gen_state_value = parse(GeneratorState, value(b_on_var))
        if get_commitment_firmness(firmness, gen_id, ts) == FREE
            set_commitment_value!(tso_schedule, gen_id, ts, s, gen_state_value)
        elseif get_commitment_firmness(firmness, gen_id, ts) in [TO_DECIDE, DECIDED]
            set_commitment_definitive_value!(tso_schedule, gen_id, ts, gen_state_value)
        end
    end

    # Capping : upper problem (TSO) locates cappings
    update_schedule_capping!(tso_schedule, result.upper.limitable_model)

    # loss_of_load (load-shedding) : upper problem (TSO) locates load shedding
    update_schedule_loss_of_load!(tso_schedule, result.upper.lol_model)

    return tso_schedule
end

function update_schedule_capping!(tso_schedule, limitable_model::TSOBilevelTSOLimitableModel)
    reset_capping!(tso_schedule)

    for ((gen_id, ts, s), p_capping_var) in limitable_model.p_capping
        tso_schedule.capping[gen_id, ts, s] = value(p_capping_var)
    end
end

function update_schedule_loss_of_load!(tso_schedule, lol_model::TSOBilevelTSOLoLModel)
    reset_loss_of_load_by_bus!(tso_schedule)

    for ((bus_id, ts, s), p_loss_of_load_var) in lol_model.p_loss_of_load
        tso_schedule.loss_of_load_by_bus[bus_id, ts, s] = value(p_loss_of_load_var)
    end
end

function update_tso_actions!(context::AbstractContext, ech, result, firmness,
                            runnable::TSOBilevel)
    tso_actions = get_tso_actions(context)
    reset_tso_actions!(tso_actions)

    # Limitations : only firm i.e. value is common to all scenarios
    limitations = SortedDict{Tuple{String,DateTime}, Float64}() #TODELETE
    for ((gen_id, ts, s), p_limit_var) in result.upper.limitable_model.p_limit
        if (value(result.upper.limitable_model.b_is_limited[gen_id, ts, s]) > 1e-09)
            if ( get_power_level_firmness(firmness, gen_id, ts) in [TO_DECIDE, DECIDED]
                || runnable.configs.LINK_SCENARIOS_LIMIT )
                @assert( value(p_limit_var) ≈ get!(limitations, (gen_id, ts), value(p_limit_var)) ) #TODELETE : checks that all values are the same across scenarios
                set_limitation_value!(tso_actions, gen_id, ts, value(p_limit_var))
            else
                #FIXME : may encounter problems if runnable.configs.LINK_SCENARIOS_LIMIT==false, cause limitations are supposed firm
                @warn "FIXME? : need to fix limitation actions to a by scenario before DP"
            end
        #else : will remain missing
        end
    end

    # Impositions
    impositions = SortedDict{Tuple{String,DateTime}, Float64}() #TODELETE
    for ((gen_id, ts, s), p_injected_var) in result.lower.pilotable_model.p_injected
        p_min_var = result.upper.pilotable_model.p_imposition_min[gen_id, ts, s]
        p_max_var = result.upper.pilotable_model.p_imposition_max[gen_id, ts, s]

        if get_power_level_firmness(firmness, gen_id, ts) in [TO_DECIDE, DECIDED]
            @assert( value(p_injected_var) ≈ get!(impositions, (gen_id, ts), value(p_injected_var)) ) #TODELETE : checks that all values are the same across scenarios
            set_imposition_value!(tso_actions, gen_id, ts, s, value(p_injected_var), value(p_injected_var))
        else
            set_imposition_value!(tso_actions, gen_id, ts, s, value(p_min_var), value(p_max_var))
        end
    end

end


@testset "Variable Tests" begin
  v = Variable("x0", :ContinuousScalar)
  # NOTE: A timestamp substitution is done in these JSON payloads.
  @test JSON.json(v) == "{\"label\":\"x0\",\"dataEntry\":\"{}\",\"nstime\":\"0\",\"variableType\":\"IncrementalInference.ContinuousScalar\",\"dataEntryType\":\"{}\",\"ppeDict\":\"{}\",\"solverDataDict\":\"{\\\"default\\\":{\\\"vecval\\\":[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],\\\"dimval\\\":1,\\\"vecbw\\\":[0.0],\\\"dimbw\\\":1,\\\"BayesNetOutVertIDs\\\":[],\\\"dimIDs\\\":[0],\\\"dims\\\":1,\\\"eliminated\\\":false,\\\"BayesNetVertID\\\":\\\"_null\\\",\\\"separator\\\":[],\\\"variableType\\\":\\\"IncrementalInference.ContinuousScalar\\\",\\\"initialized\\\":false,\\\"infoPerCoord\\\":[0.0],\\\"ismargin\\\":false,\\\"dontmargin\\\":false,\\\"solveInProgress\\\":0,\\\"solvedCount\\\":0,\\\"solveKey\\\":\\\"default\\\"}}\",\"smallData\":\"{}\",\"solvable\":1,\"tags\":\"[\\\"VARIABLE\\\"]\",\"timestamp\":$(JSON.json(v.timestamp)),\"_version\":\"0.18.1\"}"

  v = Variable("x0", :Pose1, ["TEST_TAG"])
  @test JSON.json(v) == "{\"label\":\"x0\",\"dataEntry\":\"{}\",\"nstime\":\"0\",\"variableType\":\"IncrementalInference.ContinuousScalar\",\"dataEntryType\":\"{}\",\"ppeDict\":\"{}\",\"solverDataDict\":\"{\\\"default\\\":{\\\"vecval\\\":[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],\\\"dimval\\\":1,\\\"vecbw\\\":[0.0],\\\"dimbw\\\":1,\\\"BayesNetOutVertIDs\\\":[],\\\"dimIDs\\\":[0],\\\"dims\\\":1,\\\"eliminated\\\":false,\\\"BayesNetVertID\\\":\\\"_null\\\",\\\"separator\\\":[],\\\"variableType\\\":\\\"IncrementalInference.ContinuousScalar\\\",\\\"initialized\\\":false,\\\"infoPerCoord\\\":[0.0],\\\"ismargin\\\":false,\\\"dontmargin\\\":false,\\\"solveInProgress\\\":0,\\\"solvedCount\\\":0,\\\"solveKey\\\":\\\"default\\\"}}\",\"smallData\":\"{}\",\"solvable\":1,\"tags\":\"[\\\"TEST_TAG\\\"]\",\"timestamp\":$(JSON.json(v.timestamp)),\"_version\":\"0.18.1\"}"

  v = Variable("x0", :Pose2)
  @test JSON.json(v) == "{\"label\":\"x0\",\"dataEntry\":\"{}\",\"nstime\":\"0\",\"variableType\":\"RoME.Pose2\",\"dataEntryType\":\"{}\",\"ppeDict\":\"{}\",\"solverDataDict\":\"{\\\"default\\\":{\\\"vecval\\\":[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],\\\"dimval\\\":3,\\\"vecbw\\\":[0.0,0.0,0.0],\\\"dimbw\\\":3,\\\"BayesNetOutVertIDs\\\":[],\\\"dimIDs\\\":[0,1,2],\\\"dims\\\":3,\\\"eliminated\\\":false,\\\"BayesNetVertID\\\":\\\"_null\\\",\\\"separator\\\":[],\\\"variableType\\\":\\\"RoME.Pose2\\\",\\\"initialized\\\":false,\\\"infoPerCoord\\\":[0.0,0.0,0.0],\\\"ismargin\\\":false,\\\"dontmargin\\\":false,\\\"solveInProgress\\\":0,\\\"solvedCount\\\":0,\\\"solveKey\\\":\\\"default\\\"}}\",\"smallData\":\"{}\",\"solvable\":1,\"tags\":\"[\\\"VARIABLE\\\"]\",\"timestamp\":$(JSON.json(v.timestamp)),\"_version\":\"0.18.1\"}"

  v = Variable("x0", :Point2)
  @test JSON.json(v) == "{\"label\":\"x0\",\"dataEntry\":\"{}\",\"nstime\":\"0\",\"variableType\":\"RoME.Point2\",\"dataEntryType\":\"{}\",\"ppeDict\":\"{}\",\"solverDataDict\":\"{\\\"default\\\":{\\\"vecval\\\":[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],\\\"dimval\\\":2,\\\"vecbw\\\":[0.0,0.0],\\\"dimbw\\\":2,\\\"BayesNetOutVertIDs\\\":[],\\\"dimIDs\\\":[0,1],\\\"dims\\\":2,\\\"eliminated\\\":false,\\\"BayesNetVertID\\\":\\\"_null\\\",\\\"separator\\\":[],\\\"variableType\\\":\\\"RoME.Point2\\\",\\\"initialized\\\":false,\\\"infoPerCoord\\\":[0.0,0.0],\\\"ismargin\\\":false,\\\"dontmargin\\\":false,\\\"solveInProgress\\\":0,\\\"solvedCount\\\":0,\\\"solveKey\\\":\\\"default\\\"}}\",\"smallData\":\"{}\",\"solvable\":1,\"tags\":\"[\\\"VARIABLE\\\"]\",\"timestamp\":$(JSON.json(v.timestamp)),\"_version\":\"0.18.1\"}"

  v = Variable("x0", "RoME.Pose2")
  @test JSON.json(v) == "{\"label\":\"x0\",\"dataEntry\":\"{}\",\"nstime\":\"0\",\"variableType\":\"RoME.Pose2\",\"dataEntryType\":\"{}\",\"ppeDict\":\"{}\",\"solverDataDict\":\"{\\\"default\\\":{\\\"vecval\\\":[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],\\\"dimval\\\":3,\\\"vecbw\\\":[0.0,0.0,0.0],\\\"dimbw\\\":3,\\\"BayesNetOutVertIDs\\\":[],\\\"dimIDs\\\":[0,1,2],\\\"dims\\\":3,\\\"eliminated\\\":false,\\\"BayesNetVertID\\\":\\\"_null\\\",\\\"separator\\\":[],\\\"variableType\\\":\\\"RoME.Pose2\\\",\\\"initialized\\\":false,\\\"infoPerCoord\\\":[0.0,0.0,0.0],\\\"ismargin\\\":false,\\\"dontmargin\\\":false,\\\"solveInProgress\\\":0,\\\"solvedCount\\\":0,\\\"solveKey\\\":\\\"default\\\"}}\",\"smallData\":\"{}\",\"solvable\":1,\"tags\":\"[\\\"VARIABLE\\\"]\",\"timestamp\":$(JSON.json(v.timestamp)),\"_version\":\"0.18.1\"}"

end


"""
utility functions for Housedorff kernel
"""
module HFUtils
using ..CUDAGpuUtils, Logging,StaticArrays
 using CUDA, ..BasicStructs, Logging
 using ..CUDAGpuUtils ,..IterationUtils
export    clearLocArr,clearMainShmem,clearPadding,getIndexOfQueue
export @iter3dOuter,@iterDataBlock , calculateLoopsIter, @iterDataBlockZdeepest

"""
invoked before kernel execution in order to set number of needed loop iterations
dataBdim - dimensions of the data block
threadsXdim  - x dimension of the thread block
threadsYdim  - y dimension of the thread block
return tuple with numbers indicating how many iterations are needed in the loops of the kernel
"""
function calculateLoopsIter(dataBdim,threadsXdim,threadsYdim,metaDataDims,blocks)
    loopAXFixed= fld(dataBdim[2], threadsXdim)
    loopBXfixed= fld(dataBdim[3], threadsYdim)

    loopAYFixed= fld(dataBdim[1], threadsXdim)
    loopBYfixed= fld(dataBdim[3], threadsYdim)

    loopAZFixed= fld(dataBdim[1], threadsXdim)
    loopBZfixed= fld(dataBdim[2], threadsYdim)

    loopdataDimMainX = fld(dataBdim[1], threadsXdim)
    loopdataDimMainY = fld(dataBdim[2], threadsYdim)
    loopdataDimMainZ =dataBdim[2]
    inBlockLoopX,inBlockLoopY,inBlockLoopZ= (0,fld(dataBdim[2] ,threadsYdim),dataBdim[3]);
    metaDataLength= metaDataDims[1]*metaDataDims[2]*metaDataDims[3]

    loopMeta= fld(metaDataLength,blocks )
    loopWarpMeta= cld(metaDataLength,(blocks*threadsXdim ))
    resShmemTotalLength=(dataBdim[1]+2)*(dataBdim[2]+2)*(dataBdim[3]+2)
    sourceShmemTotalLength= dataBdim[1]*dataBdim[2]*dataBdim[3]
    clearIterResShmemLoop= fld(resShmemTotalLength,threadsXdim*threadsYdim)
    clearIterSourceShmemLoop= fld(sourceShmemTotalLength,threadsXdim*threadsYdim)

    shmemblockDataLenght = threadsXdim*threadsYdim*2
    shmemblockDataLoop = fld(shmemblockDataLenght,threadsXdim*threadsYdim)

    inBlockLoopXZIterWithPadding= fld(32,threadsYdim)

    return (inBlockLoopXZIterWithPadding,shmemblockDataLoop,shmemblockDataLenght,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength)
end

"""
1)   Left FP
2)   Left FN
3)   Right FP
4)   Right FN
5)   Posterior FP
6)   Posterior FN
7)   Anterior FP
8)   Anterior FN
9)   Top FP
10)   Top FN
11)   Bottom FP
12)   Bottom FN
13)   Total block Fp
14)   Total block Fn

xpos,ypos,zpos -current  position in x,y,z dimension
dataBdim - dimensions of the data block
boolSegm - true if we have true in algo arr
on the basis of the data it should give the index from 1 to 14 - to the appropriate queue
"""
function getIndexOfQueue(xpos,ypos,zpos, dataBdim,boolSegm)
    #we need to do so many != in order to deal with corners ...
    return (
     (xpos==1)*1
    +(xpos==dataBdim[1])*3
    +(ypos==1 && xpos!=1 && xpos!=dataBdim[1] )*5
    +(ypos==dataBdim[2] && xpos!=1 && xpos!=dataBdim[1]  )*7
    +(zpos==1  && xpos!=1 && xpos!=dataBdim[1]  && ypos!=1 && ypos!=dataBdim[2] )*9
    +(zpos==dataBdim[3] && xpos!=1 && xpos!=dataBdim[1]  && ypos!=1 && ypos!=dataBdim[2])*11
    +(xpos>1 && xpos<dataBdim[1] &&  ypos>1 && ypos<dataBdim[2] && zpos>1 && zpos<dataBdim[3])*13
    )+boolSegm# in that way we will get odd for fp an even for fn

end

"""
specialization of 2 dim iteration  for iterating over 3 dimensional metadata
    we join y and z iterations in order to increase occupancy for potentially small arrays
"""
macro iter3dOuter(metaDataDims,loopMeta,metaDataLength, ex)

    return  esc(quote
    linIdexMeta = UInt32(0)
    @unroll for j in 0:($loopMeta-1)
        linIdexMeta= blockIdxX()+ j*gridDimX()-1
        xMeta= rem(linIdexMeta,$metaDataDims[1])
        zMeta= fld((linIdexMeta),$metaDataDims[1]*$metaDataDims[2])
        yMeta= fld((linIdexMeta-((zMeta*$metaDataDims[1]*$metaDataDims[2] ) + xMeta )),$metaDataDims[1])
      $ex
    end
    linIdexMeta= blockIdxX()+ $loopMeta*gridDimX() -1
            if(linIdexMeta<$metaDataLength)
            xMeta= rem(linIdexMeta,$metaDataDims[1])
            zMeta= fld((linIdexMeta),$metaDataDims[1]*$metaDataDims[2])
            yMeta= fld((linIdexMeta-((zMeta*$metaDataDims[1]*$metaDataDims[2] ) + xMeta )),$metaDataDims[1])
        $ex
      end
    end)

end #iter3dOuter

"""
will enable iterating over the data of data block
"""
macro iterDataBlock(mainArrDims,dataBlockDims,loopXdim ,loopYdim,loopZdim,xMeta,yMeta,zMeta, ex)
    mainExp = generalizedItermultiDim(;arrDims=mainArrDims
    ,loopXdim
    ,loopYdim
    ,loopZdim
    # ,xCheck = :((xMeta* $dataBlockDims[1]+x)<=$mainArrDims[1] )
    # ,yCheck = :((yMeta* $dataBlockDims[2]+y)<=$mainArrDims[2])
    # ,zCheck = :( (zMeta* $dataBlockDims[3]+z)<=$mainArrDims[3])
    ,xCheck = :(((xdim * blockDimX())+threadIdxX()) <= $dataBlockDims[1] && x<=$mainArrDims[1])
    ,yCheck = :(((ydim * blockDimY())+threadIdxY()) <= $dataBlockDims[2] && y<=$mainArrDims[2])
    ,zCheck = :((zdim+1) <= $dataBlockDims[3]  &&   z <= $mainArrDims[3])
    ,zOffset= :($zMeta* ( ($dataBlockDims[3])))
    ,zAdd =:(zdim+1)
   ,yOffset = :(ydim* blockDimY()+$yMeta* $dataBlockDims[2])
   ,yAdd= :(threadIdxY())
   ,xOffset= :( (xdim * blockDimX()) +$xMeta* $dataBlockDims[1])
    ,xAdd= :(threadIdxX())
    ,isFullBoundaryCheckX =true
    , isFullBoundaryCheckY=true
    , isFullBoundaryCheckZ=true
    ,additionalActionBeforeX=  :(xpos= xdim * blockDimX()+threadIdxX() ;ypos= (ydim * blockDimY())+threadIdxY() ;zpos=(zdim+1))
    , ex = ex)
    return esc(:( $mainExp))
end

"""
will enable iterating over the data of data block
"""
macro iterDataBlockZdeepest(mainArrDims,dataBlockDims,loopXdim ,loopYdim,loopZdim,xMeta,yMeta,zMeta, ex,additionalActionAfterX)
    mainExp = generalizedItermultiDimZdeepest(;arrDims=mainArrDims
    ,loopXdim
    ,loopYdim
    ,loopZdim
    # ,xCheck = :((xMeta* $dataBlockDims[1]+x)<=$mainArrDims[1] )
    # ,yCheck = :((yMeta* $dataBlockDims[2]+y)<=$mainArrDims[2])
    # ,zCheck = :( (zMeta* $dataBlockDims[3]+z)<=$mainArrDims[3])
    ,xCheck = :(((xdim * blockDimX())+threadIdxX()) <= $dataBlockDims[1] && x<=$mainArrDims[1])
    ,yCheck = :(((ydim * blockDimY())+threadIdxY()) <= $dataBlockDims[2] && y<=$mainArrDims[2])
    ,zCheck = :((zdim+1) <= $dataBlockDims[3]  &&   z <= $mainArrDims[3])
    ,zOffset= :($zMeta* ( ($dataBlockDims[3])))
    ,zAdd =:(zdim+1)
   ,yOffset = :(ydim* blockDimY()+$yMeta* $dataBlockDims[2])
   ,yAdd= :(threadIdxY())
   ,xOffset= :( (xdim * blockDimX()) +$xMeta* $dataBlockDims[1])
    ,xAdd= :(threadIdxX())
    ,isFullBoundaryCheckX =true
    , isFullBoundaryCheckY=true
    , isFullBoundaryCheckZ=true
    ,additionalActionBeforeX=  :(xpos= xdim * blockDimX()+threadIdxX() ;ypos= (ydim * blockDimY())+threadIdxY())
    ,additionalActionBeforeZ=  :(zpos=(zdim+1))
    ,additionalActionAfterX=additionalActionAfterX
    , ex = ex)
    return esc(:( $mainExp))
end

#macro iterDataBlock(mainArrDims,dataBlockDims,loopXdim ,loopYdim,loopZdim,ex)
    #     mainExp = generalizedItermultiDim(;arrDims=mainArrDims
    #     ,loopXdim
    #     ,loopYdim
    #     ,loopZdim
    #     # ,xCheck = :((xMeta* $dataBlockDims[1]+x)<=$mainArrDims[1] )
    #     # ,yCheck = :((yMeta* $dataBlockDims[2]+y)<=$mainArrDims[2])
    #     # ,zCheck = :( (zMeta* $dataBlockDims[3]+z)<=$mainArrDims[3])
    #     ,xCheck = :(((xdim * blockDimX())+threadIdxX()) <= $dataBlockDims[1] && x<=$mainArrDims[1])
    #     ,yCheck = :(((ydim * blockDimY())+threadIdxY()) <= $dataBlockDims[2] && y<=$mainArrDims[2])
    #     ,zCheck = :((zdim+1) <= $dataBlockDims[3]  &&   z <= $mainArrDims[3])
    #     ,zOffset= :((zMeta-1)* ( ($dataBlockDims[3])))
    #     ,zAdd =:(zdim+1)
    #    ,yOffset = :(ydim* blockDimY()+(yMeta-1)* $dataBlockDims[2])
    #    ,yAdd= :(threadIdxY())
    #    ,xOffset= :( (xdim * blockDimX()) +(xMeta-1)* $dataBlockDims[1])
    #     ,xAdd= :(threadIdxX())
    #     ,isFullBoundaryCheckX =true
    #     , isFullBoundaryCheckY=true
    #     , isFullBoundaryCheckZ=true
    #     ,additionalActionBeforeX=  :(xpos= xdim * blockDimX()+threadIdxX() ;ypos= (ydim * blockDimY())+threadIdxY() ;zpos=(zdim+1))
    #     , ex = ex)
    #     return esc(:( $mainExp))
    # end

"""
this macro will be used to
    - decide wheather the part of the code will be executed at all - in order to avoid branching at runtime
    - select at compile time appropriate target for the statement (generally the global memory array )
    - all will be based on the  ConfigurtationStruct struct - conf
    entryName - will tell us what particularly we should look for in the conf
    if the entry of name entryName in conf will be true  we will return expression otherwise we will return nothing
    we will
    metricsArr - will be used from caller scope - and it is an array in global memory that will be
"""
macro isInConf(conf, entryName, ex)

    index = 0
    if(conf.sliceWiseMatrics)
        index=3# 3 becouse one will be added in a moment
    end
    # we look for the all of the properties that are true - we need to also remember that in case of sliceWiseMatrics - it will mark that we need 4 arrays
    for i in propertynames(eval(conf))
        if getproperty(conf, i)
            index+=1
        end#if
        if i == Symbol(entryName)
            break
        end
    end#for

    if getproperty(conf, Symbol(entryName))

    return esc(quote

        @inbounds metricsArr[$index]  = $ex # @inbounds fp[]+=

    end#quote
    )
end#if

end

end#HFUtils

# """
# we use it to access pading in the shmem
# change the indexing so  instead of transverse plane of threads we will get coronal
#     so we supply idx and idx y earlier z was constant and x,y changed so we considered top with z =1 and bottom z = 34
#     now we think x and y as would be displayed on transverse plane so in order to get anterior and posterior padding plane  we need to keep y as contant 1 or 33 and rest changing
#     and in order to get left and right we need to  keep x as contant to 1 or 34 (34 in case we have size of data block = 32)
#     so we will access shmem in changed indexing
# """
# function transvarseToCoronalThreadPlane(constantNumb::UInt8,shmem)::Bool
#     return shmem[threadIdxX(),constantNumb, threadIdxY()]
# end#transvarseToCoronalThreadPlane

# """
# we use it to access pading in the shmem
# change the indexing so  instead of transverse plane of threads we will get coronal
#     so we supply idx and idx y earlier z was constant and x,y changed so we considered top with z =1 and bottom z = 34
#     now we think x and y as would be displayed on transverse plane so in order to get anterior and posterior padding plane  we need to keep y as contant 1 or 33 and rest changing
#     and in order to get left and right we need to  keep x as contant to 1 or 34 (34 in case we have size of data block = 32)
#     so we will access shmem in changed indexing
# """
# function transvarseToSaggitalhreadPlane(constantNumb::UInt8,shmem)::Bool
#     return shmem[constantNumb,threadIdxX(), threadIdxY()]
# end#transvarseToSaggitalhreadPlane

# """
# clear main part of shared memory - padding will be cleared separately
# """
# function clearMainShmem(shmem)
#     @unroll for zIter in UInt8(1):UInt8(32)# most outer loop is responsible for z dimension
#         shmem[threadIdxX()+1,threadIdxY()+1,zIter+1]=0
#     end#for
# end#clearMainShmem

# """
# clear source shmem in shared memory """
# function clearSourceShmem(shmem)
#     @unroll for zIter in UInt8(1):UInt8(32)
#          shmem[threadIdxX(),threadIdxY(),zIter]=0
#     end#for
# end#clearMainShmem

# """
# set padding planes to 0
# """
# function clearPadding(shmem)
#     clearHalfOfPadding(shmem,UInt8(1))
#     clearHalfOfPadding(shmem,UInt8(34))
# end#clearPadding
# """
# helper for clearPadding
# """
# function clearHalfOfPadding(shmem,constantNumb::UInt8)
#     shmem[constantNumb,threadIdxX()+1, threadIdxY()+1]=false
#     shmem[threadIdxX()+1,constantNumb, threadIdxY()+1]=false
#     shmem[threadIdxX()+1,threadIdxY()+1, constantNumb]=false
# end


import Base: min, max, sum, count

export max, imax, min, imin, sum, count

immutable Reduction <: AFImpl
    max::Ptr{Void}
    maxAll::Ptr{Void}
    imax::Ptr{Void}
    imaxAll::Ptr{Void}
    min::Ptr{Void}
    minAll::Ptr{Void}
    imin::Ptr{Void}
    iminAll::Ptr{Void}
    sum::Ptr{Void}
    sumNan::Ptr{Void}
    sumAll::Ptr{Void}
    sumNanAll::Ptr{Void}
    count::Ptr{Void}
    countAll::Ptr{Void}

    function Reduction(ptr)
        new(
            Libdl.dlsym(ptr, :af_max),
            Libdl.dlsym(ptr, :af_max_all),
            Libdl.dlsym(ptr, :af_imax),
            Libdl.dlsym(ptr, :af_imax_all),
            Libdl.dlsym(ptr, :af_min),
            Libdl.dlsym(ptr, :af_min_all),
            Libdl.dlsym(ptr, :af_imin),
            Libdl.dlsym(ptr, :af_imin_all),
            Libdl.dlsym(ptr, :af_sum),
            Libdl.dlsym(ptr, :af_sum_nan),
            Libdl.dlsym(ptr, :af_sum_all),
            Libdl.dlsym(ptr, :af_sum_nan_all),
            Libdl.dlsym(ptr, :af_count),
            Libdl.dlsym(ptr, :af_count_all)
        )
    end
end

macro minMaxRed(regular, all, indexed, indexedAll)
    quote
        function $(esc(regular))(arr::AFArray, dim = -1)
            verifyAccess(arr)
            af = arr.af
            result = af.results.ptr
            err = ccall(af.reduction.$regular,
                Cint, (Ptr{Ptr{Void}}, Ptr{Void}, Int32),
                result, arr.ptr, Int32(dim < 0 ? firstDim(af, arr.ptr) - 1 : dim))
            assertErr(err)
            array(af, result[])
        end

        function $(esc(regular)){T<:Real}(::Type{T}, arr::AFArray)
            verifyAccess(arr)
            af = arr.af
            real = af.results.double
            imag = af.results.double2
            err = ccall(af.reduction.$all,
                Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Void}),
                real, imag, arr.ptr)
            assertErr(err)
            T(real[])
        end

        function $(esc(indexed))(arr::AFArray, dim = -1)
            verifyAccess(arr)
            af = arr.af
            result = af.results.ptr
            idx = af.results.ptr2
            err = ccall(af.reduction.$indexed,
                Cint, (Ptr{Ptr{Void}}, Ptr{Ptr{Void}}, Ptr{Void}, Int32),
                result, idx, arr.ptr, Int32(dim < 0 ? firstDim(af, arr.ptr) - 1 : dim))
            assertErr(err)
            array(af, result[]), array(af, idx[])
        end

        function $(esc(indexed)){T<:Real}(::Type{T}, arr::AFArray)
            verifyAccess(arr)
            af = arr.af
            real = af.results.double
            imag = af.results.double2
            idx = af.results.uint
            err = ccall(af.reduction.$indexedAll,
                Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{UInt32}, Ptr{Void}),
                real, imag, idx, arr.ptr)
            assertErr(err)
            T(real[]), idx[]
        end
    end
end

@minMaxRed(max, maxAll, imax, imaxAll)

@minMaxRed(min, minAll, imin, imaxAll)

function sum(arr::AFArray, dim::Int = -1)
    verifyAccess(arr)
    af = arr.af
    result = af.results.ptr
    err = ccall(af.reduction.sum,
        Cint, (Ptr{Ptr{Void}}, Ptr{Void}, Int32),
        result, arr.ptr, Int32(dim < 0 ? firstDim(af, arr.ptr) - 1 : dim))
    assertErr(err)
    array(af, result[])
end

function sum{T<:Union{Float32,Float64}}(arr::AFArray, dim::Int, nanVal::T)
    verifyAccess(arr)
    af = arr.af
    result = af.results.ptr
    err = ccall(af.reduction.sumNan,
        Cint, (Ptr{Ptr{Void}}, Ptr{Void}, Int32, Float64),
        result, arr.ptr, Int32(dim < 0 ? firstDim(af, arr.ptr) - 1 : dim), nanVal)
    assertErr(err)
    array(af, result[])
end

function sum{T<:Real}(::Type{T}, arr::AFArray)
    verifyAccess(arr)
    af = arr.af
    real = af.results.double
    imag = af.results.double2
    err = ccall(af.reduction.sumAll,
        Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Void}),
        real, imag, arr.ptr)
    assertErr(err)
    T(real[])
end

function sum{T<:Union{Float32,Float64}}(arr::AFArray, nanVal::T)
    verifyAccess(arr)
    af = arr.af
    real = af.results.double
    imag = af.results.double2
    err = ccall(af.reduction.sumNanAll,
        Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Void}, Float64),
        real, imag, arr.ptr, nanVal)
    assertErr(err)
    real[]
end

function count(arr::AFArray, dim::Int = -1)
    verifyAccess(arr)
    af = arr.af
    result = af.results.ptr
    err = ccall(af.reduction.count,
        Cint, (Ptr{Ptr{Void}}, Ptr{Void}, Int32),
        result, arr.ptr, Int32(dim < 0 ? firstDim(af, arr.ptr) - 1 : dim))
    assertErr(err)
    array(af, result[])
end

function count{T<:Real}(::Type{T}, arr::AFArray)
    verifyAccess(arr)
    af = arr.af
    real = af.results.double
    imag = af.results.double2
    err = ccall(af.reduction.countAll,
        Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Void}),
        real, imag, arr.ptr)
    assertErr(err)
    T(real[])
end


"""
    heur_solution(inst::Instance)

Computes a heuristic solution schedule, satisfying all hard constraints
minimizing the panalties

# Arguments
-`inst::Instance` - the problem instance

# Description
The heuristic first generates a start solution, satisfying all base constraints.
Then a Tabu Search algorithm is used to minimize the penalty induced by the hard
constraints of `inst`.
"""
function heur_solution(
    inst::Instance;
    maxbreaks = 3,
    maxiter = 1000,
    infeaslength = 5,
    tabulength = 20,
    lastfeaslength = 50,
    gamtimelimit = 1200.0,
    maxphases = 20,
    inithardpen = 100,
    relaxphased = true,
    start = nothing,
)
    # compute initial solution satisfyng all pattern & patternset constraints
    if start === nothing
        schedbuf = Matrix{Int}(undef, inst.rr2.nteams, inst.rr2.nslots)
        start = pgba_heuristic(
            inst,
            maxbreaks,
            schedbuf;
            gamtimelimit = gamtimelimit,
        )
    end
    start::Matrix{Int}

    if inst.rr2.isphased && relaxphased
        phasedcons = Vector{Constraint}()
        slots = IdxSet(2^(inst.rr2.nslots ÷ 2) - 1)
        for i = 1:inst.rr2.nteams
            for j = i+1:inst.rr2.nteams
                push!(phasedcons, GAcon(slots, 1, 1, [(i, j), (j, i)], 0))
            end
        end
        rr2 = RR2(inst.rr2.nteams, inst.rr2.nslots, false)
        inst = Instance(rr2, vcat(constraints(inst), phasedcons))
    end

    cons = constraints(inst)
    hardCons = ishard.(cons)
    #lastfeas = [trues(lastfeaslength) for c = 1:length(cons)]
    pen = Int.([ishard(c) ? inithardpen : c.pen for c in cons])

    tso, state =
        initialize_search(inst.rr2, cons, pen, start, maxiter, tabulength, infeaslength)
    offsetbuf, schedbuf, difbuf = init_buffers(tso, state)
    starttime = time()

    tabu = deepcopy(state.tabu)
    add_tabu = Vector{Move}()
    #allow_infeas = false
    for phc = 1:maxphases
#        allow_infeas = !allow_infeas
        state.lastimprov = state.it
        @show state.lastimprov, maxiter, state.it
#        @show allow_infeas
        while state.lastimprov + maxiter >= state.it
            search_step!(
                tso,
                state,
                offsetbuf,
                schedbuf,
                difbuf,
                accept_infeas = false,
            )
            # logging
            print("it ", state.it)
            print(" hpen ", state.hpen, " spen ", state.spen)
            print(" optval ", state.optval, " lastimprov ", state.lastimprov)
            print(" time ", round(time() - starttime, digits = 3), "\n")
            if state.hpen == 0 && state.spen == state.optval
                println(
                    "new best feasible solution found with value ",
                    state.optval,
                )
                println(state.optsched)
                if state.spen == 0
                    return state.optsched, state.optval
                end
                tabu = deepcopy(state.tabu)
                add_tabu = Vector{Move}()
            end
        end
        if state.optval == typemax(Int) # no feasible solution found
            return state.optsched, typemax(Int)
        end
        map!(i -> i + state.it - state.lastimprov, values(tabu))
        state.tabu = deepcopy(tabu)
        state.sched .= state.optsched
        state.hpen, state.spen = 0, state.optval
        for mv in add_tabu
            state.tabu[mv] = state.it + tabulength
        end
        println("accepting infeasible solutions...")
        for k = 1:infeaslength
            println("spen, hpen =", state.spen, " ", state.hpen)
            push!(
                add_tabu,
                search_step!(
                    tso,
                    state,
                    offsetbuf,
                    schedbuf,
                    difbuf,
                    accept_infeas = true,
                ),
            )
        end
    end
    return state.optsched, state.optval
end

"""
    relaxed(inst::Instance)

Relaxes a given instance. Returns a copy of `inst` with all soft constraints
deleted and all hard constraints markd as soft wth a penalty of 1.
"""
function relaxed(inst::Instance)
    cons = Vector{Vector{Constraint}}()
    for field in (:constrCA, :constrBR, :constrGA, :constrFA, :constrSE)
        temp = Vector{Constraint}()
        for c in getfield(inst, field)
            if c.pen == 0
                push!(temp, @set c.pen = 1)
            end
        end
        push!(cons, temp)
    end
    return Instance(inst.rr2, cons...)
end

"""
    assign_games(
        sched::Matrix{Int},
        rr2::RR2,
        patternset::Vector{Int},
        gencons::Vector{Constraint},
    )

Solves the game assignment subproblem, minimizing the offset penalties of
`gencons`. The schedule `sched` is used as a start solution if
`use_start == true` and gets overriden by the optimal solution.
"""
function assign_games!(
    sched::Matrix{Int},
    rr2::RR2,
    patternset::Vector{Int},
    gencons::Vector{Constraint};
    timelimit = GRB_INFINITY,
    usestart = true,
)
    gam = game_alloc_model(rr2, patternset, gencons; feas_relax = true)
    x = gam[:x]
    if usestart
        @assert is_schedule(rr2, sched)
        for s = 1:size(sched, 2)
            for i = 1:size(sched, 1)
                if sched[i, s] > 0
                    set_start_value(x[i, sched[i, s], s], 1)
                end
            end
        end
    end
    #MOI.set(gam, MOI.RawParameter("OutputFlag"), 0)
    MOI.set(
        gam,
        Gurobi.CallbackFunction(),
        (cb_data, cb_where) -> gam_callback(gam, cb_data, cb_where, timelimit),
    )
    optimize!(gam)
    status = termination_status(gam)
    if status == MOI.OPTIMAL || status == MOI.INTERRUPTED
        parse_solution!(sched, value.(x), gam)
        return objective_value(gam)
    else
        return GRB_INFINITY
    end
end

"""
    gam_callback(timelimit, cb_data, cb_where)

Callback used to stop GAM if the time limit is exceeded and a feasible solution
was found
"""
function gam_callback(model, cb_data, cb_where, timelimit)
    if cb_where == GRB_CB_MIP
        time, obj = Ref{Cdouble}(), Ref{Cdouble}()
        GRBcbget(cb_data, cb_where, GRB_CB_RUNTIME, time)
        GRBcbget(cb_data, cb_where, GRB_CB_MIP_OBJBST, obj)
        if time[] > timelimit && obj[] < GRB_INFINITY
            GRBterminate(backend(model))
        end
    end
end


# This file contains the implementation of the general convex interface

export
  putnlcallbacks,
  clearnlcallbacks

type MSKnlinfo
  numvar :: Int
  numcon :: Int

  nlgetva :: Ptr{Void}
  nlgetsp :: Ptr{Void}

  # these are stored with 0-base indexes
  grdobjsub  :: Array{Int32,1}
  grdconsub  :: Array{Int32,1}
  grdconptr  :: Array{Int32,1}

  hessubi    :: Array{Int32,1}
  hessubj    :: Array{Int32,1}

  evalobj  :: Function
  evalconi :: Function
  grdlag   :: Function
  heslag   :: Function
  grdobj   :: Function
  grdconi  :: Function
end

# This wraps the callback method returning the sparsity of the non-linear terms.
# Depending on which arguments are NULL it should return gradient and hessian sparsity.
# NOTE: This is not documented (at all!), but in reality it is only called in a few ways.
#       While the arguments allow fetching only a portion of the hessian pattern, we ignore this
#       and always return the whole thing. It doesn't really matter since the sizes are only
#       used for pre-allocating space in MOSEK - only a reasonable upper bound is really needed.
# NOTE: The 'hessian' here is the second derivative of the lagrangian
function msk_nl_getsp_wrapper(nlhandle::    Ptr{Void},
                              numgrdobjnz:: Ptr{Int32}, # number of nonzeros in gradient of objective
                              grdobjsub::   Ptr{Int32}, # subscripts of nonzeros in gradient of objective
                              i_::          Int32,      # constraint index
                              convali::     Ptr{Bool},  # 0/1 indicating whether constraint i is non-linear
                              grdconinz::   Ptr{Int32}, # number of nonzeros in gradient of constraint i
                              grdconisub::  Ptr{Int32}, # subscripts of nonzeros in gradient of constraint i
                              yo::          Int32,      # 0/1 include objective in computation of hessian pattern
                              numycnz::     Int32,      # number of constraints to include in computation of the hessian pattern
                              ycsub::       Ptr{Int32}, # indexes of constraints to include in computation of the hessian pattern
                              maxnumhesnz:: Int32,      # lengths of hessubi and hessubj
                              numhesnz_::   Ptr{Int32}, # number of hessian nonzeros
                              hessubi::     Ptr{Int32}, # column subscrips of hessian non-zeros
                              hessubj::     Ptr{Int32}) # row subscripts of hessian non-zeros
  nlinfo = unsafe_pointer_to_objref(nlhandle) :: MSKnlinfo

  i = i_+1

  grdobjlen = length(nlinfo.grdobjsub)

  if numgrdobjnz != C_NULL
    unsafe_store!(numgrdobjnz, convert(Int32,grdobjlen))
  end

  if grdobjsub != C_NULL
    if grdobjlen > 0
      grdobjsub_a = unsafe_wrap(Array{Int32,1},grdobjsub,(grdobjlen,))
      grdobjsub_a[1:grdobjlen] = nlinfo.grdobjsub
    end
  end

  if i <= nlinfo.numcon
    if convali != C_NULL
      if nlinfo.grdconptr[i+1] > nlinfo.grdconptr[i]
        unsafe_store!(convali, convert(Int32,1))
      else
        unsafe_store!(convali, convert(Int32,0))
      end
    end

    if grdconinz != C_NULL
      unsafe_store!(grdconinz, convert(Int32, nlinfo.grdconptr[i+1] - nlinfo.grdconptr[i]))
    end

    if grdconisub != C_NULL
      num = nlinfo.grdconptr[i+1] - nlinfo.grdconptr[i]
      if num > 0
        grdconisub_a = unsafe_wrap(Array{Int32,1},grdconisub,num)
        grdconisub_a[1:num] = nlinfo.grdconsub[nlinfo.grdconptr[i]+1:nlinfo.grdconptr[i+1]]
      end
    end
  end

  numhesnz = length(nlinfo.hessubi)

  if numhesnz_ != C_NULL
    unsafe_store!(numhesnz_, convert(Int32, numhesnz))
  end

  if hessubi != C_NULL && hessubj != C_NULL && maxnumhesnz >= numhesnz
    hessubi_a = unsafe_wrap(Array{Int32,1},hessubi,(numhesnz,))
    hessubj_a = unsafe_wrap(Array{Int32,1},hessubj,(numhesnz,))

    hessubi_a[1:numhesnz] = nlinfo.hessubi
    hessubj_a[1:numhesnz] = nlinfo.hessubj
  end

  return Int32(0) :: Int32
end

function msk_nl_getva_wrapper(nlhandle    :: Ptr{Void},
                              xx_         :: Ptr{Float64}, # input
                              yo          :: Float64,
                              yc_         :: Ptr{Float64}, # input, length = numcon
                              objval      :: Ptr{Float64},
                              numgrdobjnz :: Ptr{Int32},
                              grdobjsub   :: Ptr{Int32},
                              grdobjval   :: Ptr{Float64},
                              numi_       :: Int32,
                              subi_       :: Ptr{Int32},   # input
                              conval      :: Ptr{Float64},
                              grdconptrb_ :: Ptr{Int32},   # input
                              grdconptre_ :: Ptr{Int32},   # input
                              grdconsub_  :: Ptr{Int32},   # input
                              grdconval   :: Ptr{Float64},
                              grdlag      :: Ptr{Float64},
                              maxnumhesnz :: Int32,
                              numhesnz    :: Ptr{Int32},
                              hessubi     :: Ptr{Int32},
                              hessubj     :: Ptr{Int32},
                              hesval      :: Ptr{Float64})
  nlinfo = unsafe_pointer_to_objref(nlhandle) :: MSKnlinfo

  numi = convert(Int,numi_)
  xx = unsafe_wrap(Array{Float64,1},xx_,(nlinfo.numvar,))
  yc = unsafe_wrap(Array{Float64,1},yc_,(nlinfo.numcon,))
  subi = unsafe_wrap(Array{Int32,1},subi_,(numi,)) .+ 1

  if objval != C_NULL
    unsafe_store!(objval, nlinfo.evalobj(xx))
  end

  ngrdobjnz = length(nlinfo.grdobjsub)

  if numgrdobjnz != C_NULL
    unsafe_store!(numgrdobjnz, convert(Int32,ngrdobjnz))
  end

  if grdobjsub != C_NULL && grdobjval != C_NULL
    grdobjval_a = unsafe_wrap(Array{Float64,1},grdobjval,(ngrdobjnz,))
    grdobjsub_a = unsafe_wrap(Array{Int32,1},grdobjsub,(ngrdobjnz,))

    grdobjsub = nlinfo.grdobjsub .+ 1
    nlinfo.grdobj(xx,
                  grdobjsub,
                  grdobjval_a)
    for i in 1:length(grdobjsub_a)
      grdobjsub_a[i] = nlinfo.grdobjsub[i]
    end
  end

  if numi > 0 && conval != C_NULL
    conv   = unsafe_wrap(Array{Float64,1},conval,(numi,))
    for i=1:numi
      conv[i] = nlinfo.evalconi(xx,subi[i])
    end
  end

  if grdconval != C_NULL
    grdconptrb = unsafe_wrap(Array{Int32,1},grdconptrb_,(numi,))
    grdconptre = unsafe_wrap(Array{Int32,1},grdconptre_,(numi,))
    for i=1:numi
      ptrb = grdconptrb[i]
      n    = grdconptre[i] - grdconptrb[i]
      nlinfo.grdconi(xx,subi[i],
                     unsafe_wrap(Array{Int32,1},grdconsub_+ptrb*4,(n,)) .+ 1,
                     unsafe_wrap(Array{Float64,1},grdconval +ptrb*4,(n,)))
    end
  end

  if grdlag != C_NULL
    nlinfo.grdlag(xx,yo,yc,subi,unsafe_wrap(Array{Float64,1},grdlag,(nlinfo.numvar,)))
  end

  nhesnz = length(nlinfo.hessubi)

  if numhesnz != C_NULL
    unsafe_store!(numhesnz,convert(Int32,nhesnz))
  end

  if maxnumhesnz > 0 && hessubi != C_NULL && hessubj != C_NULL && hesval != C_NULL
    hessubi_a = unsafe_wrap(Array{Int32,1},hessubi,(nhesnz,))
    hessubj_a = unsafe_wrap(Array{Int32,1},hessubj,(nhesnz,))
    nlinfo.heslag(xx,yo,yc,subi,
                  hessubi_a,
                  hessubj_a,
                  unsafe_wrap(Array{Float64,1},hesval, (nhesnz,)))

    # Hum... Very strange! If I use "hessubi_a -= 1" it appears that the
    # underlying data of the array is not the native pointer anymore, but
    # if I do it this way it is.
    for i=1:length(hessubi_a)
      hessubi_a[i] = hessubi_a[i]-convert(Int32,1)
      hessubj_a[i] = hessubj_a[i]-convert(Int32,1)
    end

    hesval_a  = unsafe_wrap(Array{Float64,1},hesval,(nhesnz,))
  end
  return convert(Int32,0) :: Int32
end

function putnlcallbacks(task::MSKtask,
                        grdobjsub :: Array{Int,1},
                        grdconsub :: Array{Int,1},
                        grdconptr :: Array{Int,1},
                        hessubi   :: Array{Int,1},
                        hessubj   :: Array{Int,1},
                        evalobj   :: Function,
                        evalconi  :: Function,
                        grdlag    :: Function,
                        grdobj    :: Function,
                        grdconi   :: Function,
                        heslag    :: Function)
  nvar = convert(Int,getnumvar(task))
  ncon = convert(Int,getnumcon(task))

  #print( ncon,length(grdconptr))
  if ( length(hessubi) !=  length(hessubj))
    error("Arrays hessubi and hessubj have mismatching lengths")
  end
  if ( length(grdconptr) != ncon+1 )
    error("Length of grdconptr should match number of constraints")
  end

  nlgetsp = cfunction(msk_nl_getsp_wrapper,
                      Int32,
                      (Ptr{Void},Ptr{Int32},Ptr{Int32},Int32,Ptr{Bool},Ptr{Int32},Ptr{Int32},Int32,Int32,Ptr{Int32},Int32,Ptr{Int32},Ptr{Int32},Ptr{Int32}))
  nlgetva = cfunction(msk_nl_getva_wrapper,
                      Int32,
                      ( Ptr{Void}, # nlhandle
                        Ptr{Float64},Float64,Ptr{Float64}, # xx,yo,yc
                        Ptr{Float64},Ptr{Int32},Ptr{Int32},Ptr{Float64}, # objval,numgrdobjnz,grdobjsub,grdobjval
                        Int32,Ptr{Int32},Ptr{Float64}, # numi,subi,conval
                        Ptr{Int32},Ptr{Int32},Ptr{Int32},Ptr{Float64},Ptr{Float64},# grdconptrb,grdconptre,grdconsub,grdconval,grdlag
                        Int32,Ptr{Int32},Ptr{Int32},Ptr{Int32},Ptr{Float64}))# maxnumhesnz, numhesnz.hessubi,hessubj,hesval

  nlinfo = MSKnlinfo(nvar,ncon,
                     nlgetsp,nlgetva,
                     round(Int32,grdobjsub .- 1),
                     round(Int32,grdconsub .- 1),
                     round(Int32,grdconptr .- 1),
                     round(Int32,hessubi .- 1),
                     round(Int32,hessubj .- 1),
                     evalobj,evalconi, grdlag,heslag,grdobj,grdconi)

  @msk_ccall("putnlfunc",
             Int32, (Ptr{Void},Any,Ptr{Void},Ptr{Void}),
             task.task, nlinfo, nlgetsp, nlgetva)
  task.nlinfo = nlinfo
end

function clearnlcallbacks(task::MSKtask)
  @msk_ccall("putnlfunc",
             Int32, (Ptr{Void},Ptr{Void},Ptr{Void},Ptr{Void}),
             task.task, C_NULL,C_NULL,C_NULL,C_NULL)
  task.nlinfo = nothing
end


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

export
    structmatches, structmatch, isstructmatch,
    substructmatches, substructmatch, issubstructmatch,
    querymatch, isquerymatch,
    fastquerymatches, isqueryidentical, issmartsgroupmatch

"""
    structmatches(mol1::UndirectedGraph, mol2::UndirectedGraph; kwargs...) -> Iterator

Generate molecular graph match mappings between `mol1` and `mol2`. If no match
found, return nothing.
"""
function structmatches(mol1::UndirectedGraph, mol2::UndirectedGraph;
        prefilter=true, kwargs...)
    !prefilter || fastidentityfilter(mol1, mol2) || return []
    afunc = atommatch(mol1, mol2)
    bfunc = bondmatch(mol1, mol2)
    return graphmatches(mol1, mol2,
        nodematcher=afunc, edgematcher=bfunc; kwargs...)
end

function structmatch(mol1, mol2; kwargs...)
    res = iterate(structmatches(mol1, mol2; kwargs...))
    return res === nothing ? nothing : res[1]
end

isstructmatch(mol1, mol2; kwargs...
    ) = structmatch(mol1, mol2; kwargs...) !== nothing

"""
    substructmatches(mol1::UndirectedGraph, mol2::UndirectedGraph; kwargs...) -> Iterator

Generate substructure match mappings between `mol1` and `mol2`. If no match
found, return nothing.

The mapping is based on only edge induced subgraph isomorphism and therefore
it does not care disconnected single atom matches. This function is intended
for use in substructure search. If you need single atom SMARTS match
(ex. [#16;X2;!R]), see [`querymatch`](@ref).
"""
function substructmatches(mol1::UndirectedGraph, mol2::UndirectedGraph;
        prefilter=true, kwargs...)
    !prefilter || fastsubstrfilter(mol1, mol2) || return []
    afunc = atommatch(mol1, mol2)
    bfunc = bondmatch(mol1, mol2)
    return edgesubgraphmatches(mol1, mol2,
        nodematcher=afunc, edgematcher=bfunc; kwargs...)
end

function substructmatch(mol1, mol2; kwargs...)
    res = iterate(substructmatches(mol1, mol2; kwargs...))
    return res === nothing ? nothing : res[1]
end

issubstructmatch(mol1, mol2; kwargs...
    ) = substructmatch(mol1, mol2; kwargs...) !== nothing

"""
    querymatch(mol::UndirectedGraph, query::QueryMol; kwargs...) -> Dict{Int,Int}

Generate substructure match mappings between `mol1` and `mol2`. If no match
found, return nothing.

This accepts also disconnected single atom but returns only the first match.
This function is intended for use in SMARTS query search

"""
function querymatch(mol::UndirectedGraph, query::QueryMol; kwargs...)
    mnodeset = nodeset(mol)
    qnodeset = nodeset(query)
    afunc = atommatch(mol, query)
    bfunc = bondmatch(mol, query)
    if edgecount(query) != 0
        # Edge induced subgraph mapping
        for emap in edgesubgraphmatches(
                mol, query, nodematcher=afunc, edgematcher=bfunc; kwargs...)
            # Isolated node mapping
            msub = edgesubgraph(mol, Set(keys(emap)))
            qsub = edgesubgraph(query, Set(values(emap)))
            miso = setdiff(mnodeset, nodeset(msub))
            qiso = setdiff(qnodeset, nodeset(qsub))
            mrem = nodesubgraph(mol, miso)
            qrem = nodesubgraph(query, qiso)
            nmap = maxcardmap(nodeset(mrem), nodeset(qrem), afunc)
            # TODO: connectivity filter
            if length(nmap) == nodecount(qrem)
                return (emap, nmap)
            end
        end
    elseif nodecount(query) != 0
        # Isolated nodes only
        nmap = maxcardmap(mnodeset, qnodeset, afunc)
        # TODO: connectivity filter
        if length(nmap) == nodecount(query)
            return (Dict(), nmap)
        end
    end
    return
end

"""
    isquerymatch(mol, query; kwargs...)

Return whether mol matches with the query.
"""
isquerymatch(mol, query; kwargs...
    ) = querymatch(mol, query; kwargs...) !== nothing

"""
    fastquerymatches(mol::UndirectedGraph, query::QueryMol; kwargs...
        ) -> Dict{Int,Int}

Generate query match mappings between `mol` and `query`. If no match
found, return nothing.

The `query` should not have any component level expression that means it should
not have any dots (`.`). This is intended for use in functional group detection.

"""
function fastquerymatches(mol::UndirectedGraph, query::QueryMol; kwargs...)
    isempty(query.connectivity) || throw(
        ErrorException("Component level query is disallowed"))
    afunc = atommatch(mol, query)
    bfunc = bondmatch(mol, query)
    nodecount(query) == 0 && return ()
    if nodecount(query) == 1
        # node match
        match = Iterators.filter(nodeset(mol)) do ma
            return afunc(ma, 1)
        end
        return ((Dict(), Dict(ma => 1)) for ma in match)
    elseif edgecount(query) == 1
        # edge match
        (qu, qv) = getedge(query, 1)
        match = Iterators.filter(edgeset(mol)) do mb
            (u, v) = getedge(mol, mb)
            return (
                ((afunc(u, qu) && afunc(v, qv))
                || (afunc(u, qv) && afunc(v, qu)))
                && bfunc(mb, 1)
            )
        end
        return ((Dict(mb => 1), Dict()) for mb in match)
    else
        # subgraph match
        return ((emap, Dict()) for emap in edgesubgraphmatches(
            mol, query, nodematcher=afunc, edgematcher=bfunc; kwargs...))
    end
end

function isqueryidentical(mol::UndirectedGraph, query::QueryMol; kwargs...)
    isempty(query.connectivity) || throw(
        ErrorException("Component level query is disallowed"))
    afunc = atommatch(mol, query)
    bfunc = bondmatch(mol, query)
    nodecount(query) == 0 && return ()
    if nodecount(query) == 1
        for n in nodeset(mol)
            afunc(n, 1) && return true
        end
        return false
    elseif edgecount(query) == 1
        (qu, qv) = getedge(query, 1)
        for e in edgeset(mol)
            (u, v) = getedge(mol, e)
            pred = ((afunc(u, qu) && afunc(v, qv))
                || (afunc(u, qv) && afunc(v, qu))) && bfunc(e, 1)
            pred && return true
        end
        return false
    else
        return isgraphmatch(
            mol, query, nodematcher=afunc, edgematcher=bfunc; kwargs...)
    end
end

function issmartsgroupmatch(
        mol::UndirectedGraph, query::QueryMol, root::Int; kwargs...)
    # For recursive SMARTS match
    isempty(query.connectivity) || throw(
        ErrorException("Component level query is disallowed"))
    afunc = atommatch(mol, query)
    bfunc = bondmatch(mol, query)
    @assert root in nodeset(mol) "node $(root) does not exist"
    if nodecount(query) == 1
        # node match
        return afunc(root, 1)
    elseif edgecount(query) == 1
        # edge match
        for (inc, adj) in neighbors(mol, root)
            if afunc(root, 1) && afunc(adj, 2) && bfunc(inc, 1)
                return true
            end
        end
        return false
    else
        # subgraph match
        for n in incidences(mol, root)
            if isedgesubgraphmatch(
                    mol, query, nodematcher=afunc, edgematcher=bfunc,
                    mandatory=Dict(n=>1); kwargs...)
                return true
            end
        end
        return false
    end
end

function atommatch(mol1::UndirectedGraph, mol2::UndirectedGraph)
    sym1 = atomsymbol(mol1)
    sym2 = atomsymbol(mol2)
    pi1 = pielectron(mol1)
    pi2 = pielectron(mol2)
    return (a1, a2) -> sym1[a1] == sym2[a2] && pi1[a1] == pi2[a2]
end

function atommatch(mol::UndirectedGraph, querymol::QueryMol)
    return (a, qa) -> querymatchtree(nodeattr(querymol, qa).query, mol, a)
end

function bondmatch(mol1::UndirectedGraph, mol2::UndirectedGraph)
    # TODO: need bond attribute matching?
    return (b1, b2) -> true
end

function bondmatch(mol::UndirectedGraph, querymol::QueryMol)
    return (b, qb) -> querymatchtree(edgeattr(querymol, qb).query, mol, b)
end

function querymatchtree(query::Pair, mol::UndirectedGraph, i::Int)
    if query.first == :any
        return true
    elseif query.first == :and
        return all(querymatchtree(q, mol, i) for q in query.second)
    elseif query.first == :or
        return any(querymatchtree(q, mol, i) for q in query.second)
    elseif query.first == :not
        return !querymatchtree(query.second, mol, i)
    elseif query.first == :stereo
        # TODO: stereo not implemented yet
        return true
    elseif query.first == :recursive
        subq = parse(SMARTS, query.second)
        return issmartsgroupmatch(mol, subq, i)
    else
        if query.first == :atom_sssrsizes
            return query.second in atom_sssrsizes(mol)[i]
        else
            return mol[query.first][i] == query.second
        end
    end
end

function fastidentityfilter(mol1::UndirectedGraph, mol2::UndirectedGraph)
    if nodecount(mol1) != nodecount(mol2)
        return false
    elseif edgecount(mol1) != edgecount(mol2)
        return false
    elseif circuitrank(mol1) != circuitrank(mol2)
        return false
    end
    return true
end

function fastsubstrfilter(mol1::UndirectedGraph, mol2::UndirectedGraph)
    if nodecount(mol1) < nodecount(mol2)
        return false
    elseif edgecount(mol1) < edgecount(mol2)
        return false
    elseif circuitrank(mol1) < circuitrank(mol2)
        return false
    end
    return true
end


module TrafficWeavingPlanner

const robot = "/robot"
const human = "/human"

using PyCall
using RobotOS
using StaticArrays
using Rotations
using Interpolations

using TrafficWeavingPlanner
using DifferentialDynamicsModels
using LinearDynamicsModels

@pyimport rospy
@pyimport tf2_ros
@rosimport geometry_msgs.msg: Transform, TransformStamped, PoseStamped, TwistStamped, AccelStamped
@rosimport std_msgs.msg: Float32MultiArray, Float64MultiArray,
                         Int32MultiArray, Int64MultiArray,
                         MultiArrayLayout, MultiArrayDimension,
                         Bool
@rosimport safe_traffic_weaving.msg: Float32MultiArrayStamped, PredictionInput, PredictionOutput
rostypegen()
import geometry_msgs.msg: Transform, TransformStamped, PoseStamped, TwistStamped, AccelStamped
import std_msgs.msg: Float32MultiArray, Float64MultiArray,
                     Int32MultiArray, Int64MultiArray,
                     MultiArrayLayout, MultiArrayDimension,
                     BoolMsg
import safe_traffic_weaving.msg: Float32MultiArrayStamped, PredictionInput, PredictionOutput

# const tfBuffer = tf2_ros.Buffer()
# const listener = tf2_ros.TransformListener(tfBuffer)
# lane_100_py_tf = tfBuffer[:lookup_transform]("world", "lane_100", rospy.Time(), rospy.Duration(1.0))
# lane_0_py_tf = tfBuffer[:lookup_transform]("world", "lane_0", rospy.Time(), rospy.Duration(1.0))

const Timeseries2D = Interpolations.GriddedInterpolation{SVector{2,Float64},1,SVector{2,Float32},Interpolations.Gridded{Interpolations.Linear},Tuple{Array{Float64,1}},0}
Timeseries2D() = interpolate((Float64[],), SVector{2,Float32}[], Gridded(Linear()))
truncate!(v::Vector) = resize!(v, 0)
truncate!(ts::Timeseries2D) = (foreach(truncate!, ts.knots); truncate!(ts.coefs))
Base.push!(ts::Timeseries2D, t, x) = (push!(ts.knots[1], t); push!(ts.coefs, x))
Base.push!(ts::Timeseries2D, tx::Tuple) = push!(ts, tx[1], tx[2])

# struct Timeseries2D
#     t::Vector{Float64}
#     xy::Vector{SVector{2,Float32}}
# end
# Timeseries2D() = Timeseries2D(Float64[], SVector{2,Float32}[])
# truncate!(v::Vector) = resize!(v, 0)
# truncate!(ts::Timeseries2D) = (truncate!(ts.t); truncate!(ts.xy))

struct DoubleIntegratorTimeseries2D
    pos::Timeseries2D
    vel::Timeseries2D
    acc::Timeseries2D
end
DoubleIntegratorTimeseries2D() = DoubleIntegratorTimeseries2D(Timeseries2D(), Timeseries2D(), Timeseries2D())
truncate!(dits::DoubleIntegratorTimeseries2D) = (truncate!(dits.pos); truncate!(dits.vel); truncate!(dits.acc))
Base.getindex(dits::DoubleIntegratorTimeseries2D, t) = [dits.pos[t]; dits.vel[t]; dits.acc[t]]
function last_common_time(dits::DoubleIntegratorTimeseries2D)
    if length(dits.pos.knots[1]) > 0 && length(dits.vel.knots[1]) > 0 && length(dits.acc.knots[1]) > 0
        min(dits.pos.knots[1][end], dits.vel.knots[1][end], dits.acc.knots[1][end])
    else
        -1.0
    end
end
function first_common_time(dits::DoubleIntegratorTimeseries2D)
    if length(dits.pos.knots[1]) > 0 && length(dits.vel.knots[1]) > 0 && length(dits.acc.knots[1]) > 0
        max(dits.pos.knots[1][1], dits.vel.knots[1][1], dits.acc.knots[1][1])
    else
        -1.0
    end
end

mutable struct PlannerState
    initialized::Bool
    terminate::Bool
    start_time::Float64
    computed_up_to_t::Int
    human::Vector{SVector{6,Float32}}     # car2
    robot::Vector{SVector{6,Float32}}     # car1
    human_buffer::DoubleIntegratorTimeseries2D
    robot_buffer::DoubleIntegratorTimeseries2D
    # robot_actions::Vector{Tuple{Int,Int}}
    robot_des::Vector{SVector{6,Float32}} # identical to robot up to current time, but includes future actions
    robot_l0::Int
    robot_lf::Int
    # robot_a0::Vector{Tuple{Int,Int}}    # current queue of actions to be implemented
    # robot_a0_t0::Int                    # timestep of current action
end

mutable struct HumanCarStates
    t::Vector{Float32}
    X::Vector{SVector{6,Float32}}
end

const planner_state = PlannerState(false, false, 0.0, 0,
                                   SVector{6,Float32}[],
                                   SVector{6,Float32}[],
                                   DoubleIntegratorTimeseries2D(),
                                   DoubleIntegratorTimeseries2D(),
                                   # Tuple{Int,Int}[],
                                   SVector{6,Float32}[], 0, 0)

get_xy(msg::PoseStamped) = SVector{2,Float32}(msg.pose.position.x, msg.pose.position.y)
get_xy(msg::TwistStamped) = SVector{2,Float32}(msg.twist.linear.x, msg.twist.linear.y)
get_xy(msg::AccelStamped) = SVector{2,Float32}(msg.accel.linear.x, msg.accel.linear.y)
function append_callback(msg, ts)
    xy = get_xy(msg)
    !isnan(xy[1]) && !isnan(xy[2]) && push!(ts, Float64(msg.header.stamp), xy)
end

# human_pos_callback(msg, ts) = (push!(ts.t, Float64(msg.header.stamp)), push!(ts.xy, SVector{2,Float32}(msg.pose.position.x,
#                                                                                msg.pose.position.y))
# human_vel_callback(msg, ts) = (push!(ts.t, Float64(msg.header.stamp)), push!(ts.xy, SVector{2,Float32}(msg.twist.linear.x,
#                                                                                msg.twist.linear.y))
# human_acc_callback(msg, ts) = (push!(ts.t, Float64(msg.header.stamp)), push!(ts.xy, SVector{2,Float32}(msg.accel.linear.x,
#                                                                                msg.accel.linear.y))
# robot_pos_callback(msg, ts) = (push!(ts.t, Float64(msg.header.stamp)), push!(ts.xy, SVector{2,Float32}(msg.pose.position.x,
#                                                                                msg.pose.position.y))
# robot_vel_callback(msg, ts) = (push!(ts.t, Float64(msg.header.stamp)), push!(ts.xy, SVector{2,Float32}(msg.twist.linear.x,
#                                                                                msg.twist.linear.y))
# robot_acc_callback(msg, ts) = (push!(ts.t, Float64(msg.header.stamp)), push!(ts.xy, SVector{2,Float32}(msg.accel.linear.x,
#                                                                                msg.accel.linear.y))

function terminate_callback(msg::BoolMsg, planner_state::PlannerState)
    planner_state.terminate = msg.data
end

function reset_scenario_callback(msg::BoolMsg, planner_state::PlannerState)
    if msg.data
        planner_state.initialized = true
        planner_state.start_time = Float64(RobotOS.now())
        planner_state.computed_up_to_t = 0
        if RobotOS.get_param("robot_start_lane", "R") == "L"
            planner_state.robot_l0, planner_state.robot_lf = 1, 2
        else
            planner_state.robot_l0, planner_state.robot_lf = 2, 1
        end
        # if planner_state.human_buffer.pos.coefs[end][2] > -4    # HACK
        # planner_state.robot_l0, planner_state.robot_lf = 2, 1
        # else
        #     planner_state.robot_l0, planner_state.robot_lf = 2, 1
        # end
        # planner_state.robot_a0 = Tuple{Int,Int}[]
        # planner_state.robot_actions = [(1,planner_state.robot_l0) for i in 1:(start_wait + computation_interval)]
        rossleep(0.5)
        truncate!(planner_state.robot_des)
        truncate!(planner_state.human)
        truncate!(planner_state.robot)
        truncate!(planner_state.human_buffer)
        truncate!(planner_state.robot_buffer)
    end
end

init_node("traffic_weaving_planner", anonymous=true)
const terminate_sub = Subscriber{BoolMsg}("/terminate", terminate_callback, (planner_state,), queue_size=10)
const reset_sub = Subscriber{BoolMsg}("/reset", reset_scenario_callback, (planner_state,), queue_size=1)
# const reset_sim_sub = Subscriber{BoolMsg}("/reset_sim", reset_scenario_callback, (planner_state,), queue_size=10)

const human_pos_sub = Subscriber{PoseStamped}(human*"/pose", append_callback, (planner_state.human_buffer.pos,), queue_size=10)
const human_vel_sub = Subscriber{TwistStamped}(human*"/vel", append_callback, (planner_state.human_buffer.vel,), queue_size=10)
const human_acc_sub = Subscriber{AccelStamped}(human*"/acc", append_callback, (planner_state.human_buffer.acc,), queue_size=10)
const robot_pos_sub = Subscriber{PoseStamped}(robot*"/pose", append_callback, (planner_state.robot_buffer.pos,), queue_size=10)
const robot_vel_sub = Subscriber{TwistStamped}(robot*"/vel", append_callback, (planner_state.robot_buffer.vel,), queue_size=10)
const robot_acc_sub = Subscriber{AccelStamped}(robot*"/acc", append_callback, (planner_state.robot_buffer.acc,), queue_size=10)

const robot_traj_pub = Publisher{Float32MultiArrayStamped}("/robot/traj_plan", queue_size=10)
convert_to_3d_array(V::Vector{SVector{6,Float32}}) = reinterpret(Float32, copy(V), (6, length(V), 1))
convert_to_reversed_3d_array(V::Vector{SVector{6,Float32}}) = reshape(reinterpret(Float32, copy(V), (6, length(V)))', (1, length(V), 6))
function publish_robot_plan(robot_des::Vector{SVector{6,Float32}}, time_offset)
    msg = Float32MultiArrayStamped()
    msg.data = convert_to_3d_array(robot_des)
    msg.header.stamp = (time_offset isa Number ? RobotOS.Time(time_offset) : time_offset)
    publish(robot_traj_pub, msg)
end

include("utils/planner.jl")

function run(;start_wait=8, computation_interval=3)
    rate = RobotOS.Rate(100)

    time_scale_factor = RobotOS.get_param("time_scale_factor")
    dt = Float32(.1/time_scale_factor)

    reset_scenario_callback(BoolMsg(true), planner_state)
    planner_state.terminate = false
    while !planner_state.terminate
        @assert length(planner_state.human) == length(planner_state.robot)
        Tnew = min(last_common_time(planner_state.human_buffer), last_common_time(planner_state.robot_buffer))
        if isempty(planner_state.human)
            # t_now = Float64(RobotOS.now())
            # push!(planner_state.robot_buffer.pos, t_now, [-138.2, -6.0940447])
            # push!(planner_state.robot_buffer.pos, t_now, [28., 0.])
            # push!(planner_state.robot_buffer.pos, t_now, [0., 0.])
            Tfirst = max(first_common_time(planner_state.human_buffer), first_common_time(planner_state.robot_buffer), planner_state.start_time)
            println("still empty, $Tfirst, $Tnew")
            # println(planner_state.human_buffer.pos.knots)
            # println(planner_state.human_buffer.pos.coefs)
            # println(planner_state.robot_buffer)
            try
                if Tnew == -1 || Tfirst == -1 || Tfirst >= Tnew - 0.02 || planner_state.robot_buffer[Tnew-0.01][1] < -138
                    rossleep(0.3)
                    continue
                end
            catch
                println("brute-force catching our way past race conditions (1)")
                continue
            end
            A, B = Tfirst, Tnew
            Tstart = (A + B)/2
            for i in 1:10
                planner_state.robot_buffer[Tstart][1] < -138 ? (A = Tstart) : (B = Tstart)
                Tstart = (A + B)/2
            end
            println("ready to start")
            planner_state.start_time = Tstart
            initial_actions = [(1,planner_state.robot_l0) for i in 1:(start_wait + computation_interval)]
            planner_state.robot_des = propagate(xdd_choices, y_targets, planner_state.robot_buffer[planner_state.start_time], tibvp, initial_actions)
            publish_robot_plan(planner_state.robot_des, planner_state.start_time)
        end
        T = planner_state.start_time + (length(planner_state.human) - 1)*dt
        Tnew = min(last_common_time(planner_state.human_buffer), last_common_time(planner_state.robot_buffer))    # in case of race with /reset channel?

        race_condition_hit = false
        while T + dt < Tnew
            println("pushing")
            try
                push!(planner_state.human, planner_state.human_buffer[T + dt])
                push!(planner_state.robot, planner_state.robot_buffer[T + dt])
            catch
                println("brute-force catching our way past race conditions (2)")
                race_condition_hit = true
                break
            end
            if length(planner_state.robot) <= length(planner_state.robot_des)
                planner_state.robot_des[length(planner_state.robot)] = planner_state.robot[end]  # maybe unnecessary
            end
            T = T + dt
        end
        race_condition_hit && continue

        println("caught up")
        # t = floor(Int, (Float64(RobotOS.now()) - planner_state.start_time)/dt) + 1
        t = length(planner_state.human)
        t = t - (t % computation_interval)    # enforcing strict computation intervals
        if t <= 5
            rossleep(0.1)
            continue
        end

        println("sufficient history")

        # a0 = planner_state.robot_actions[t:t+computation_interval-1]
        if t+computation_interval > length(planner_state.robot_des) || t <= planner_state.computed_up_to_t
            rossleep(0.01)
            continue
        end
        a0_traj = planner_state.robot_des[t+1:t+computation_interval]
        a0 = ([a0_traj[i][5] for i in 1:computation_interval], [a0_traj[i][6] for i in 1:computation_interval])    # Ẍ and Ÿ; TODO: plumb targets instead of using previous Ÿ
        pfX, pfY = all_future_xy(xdd_choices, y_targets, planner_state.robot[t], tibvp, a0, 4, des_length=15)
        robot_in = convert_to_reversed_3d_array(planner_state.robot)
        human_in = convert_to_reversed_3d_array(planner_state.human)
        extras_in = zeros(Float32, 1, length(planner_state.robot), 0)
        predict!(robot_in, human_in, extras_in, [t], pfX[:,2:end,:], pfY[:,2:end,:], 64)
        for i in 1:15
            _t = t + i
            if _t <= length(planner_state.robot_des)
                planner_state.robot_des[_t] = SVector(pred_container.b[i,:]...)
            else
                push!(planner_state.robot_des, SVector(pred_container.b[i,:]...))
            end
        end
        publish_robot_plan(planner_state.robot_des, planner_state.start_time)
        planner_state.computed_up_to_t = t

        # TrafficWeavingPlanner.pred_container.y[4:6,1:2:end], TrafficWeavingPlanner.pred_container.y[4:6,2:2:end]
        # car1[1,t+1:t+15,:] = TrafficWeavingPlanner.pred_container.y
        # t = t + 3
        # car1_curr_state = car1[1,t,:]
        # a0 = (TrafficWeavingPlanner.pred_container.y[4:6,1:2:end], TrafficWeavingPlanner.pred_container.y[4:6,2:2:end])

        rossleep(rate)
    end
end

@spawn spin()

# mutable struct SimulatorState
#     reset::Bool
#     xh::SE2vState{Float64}
#     uh::AccelerationCurvatureControl{Float64}
# end

# function joystick_callback(msg::Joy, sim_state::SimulatorState)
#     if msg.buttons[1] == 1 && msg.buttons[2] == 1
#         sim_state.reset = true
#     end
#     sim_state.uh = AccelerationCurvatureControl(MAX_LONGITUDINAL_ACC*msg.axes[2],
#                                                 MAX_CURVATURE*msg.axes[1])
# end

# init_node("xbox_car", anonymous=true)

# const sim_state = SimulatorState(false, zeros(SE2vState{Float64}), zeros(AccelerationCurvatureControl{Float64}))
# const joystick_sub = Subscriber{Joy}("joy", joystick_callback, (sim_state,), queue_size=10)
# const human_vel_pub = Publisher{TwistStamped}("human/vel", queue_size=10)
# const human_accel_pub = Publisher{AccelStamped}("human/accel", queue_size=10)
# const tf_broadcaster = tf2_ros.TransformBroadcaster()

# function fill_2D_transform!(t::TransformStamped, x, y, θ, stamp = RobotOS.now())
#     t.header.stamp = stamp
#     t.transform.translation.x = x
#     t.transform.translation.y = y
#     t.transform.rotation.x,
#     t.transform.rotation.y,
#     t.transform.rotation.z,
#     t.transform.rotation.w = tf.transformations[:quaternion_from_euler](0.0, 0.0, θ)
#     t
# end

# function test()
#     rate = Rate(SIM_RATE)

#     world_header = Header()
#     world_header.frame_id = "world"
#     human_car_wrt_world = TransformStamped(world_header, "human_car", Transform())
#     vel_msg, accel_msg = TwistStamped(), AccelStamped()
#     vel_msg.header.frame_id = accel_msg.header.frame_id = "human_car"

#     sim_state.xh = zeros(SE2vState{Float64})
#     while !sim_state.reset
#         timestamp = RobotOS.now()
#         sim_state.xh = propagate(SimpleCarDynamics{1,0}(), sim_state.xh, StepControl(1/SIM_RATE, sim_state.uh))
#         fill_2D_transform!(human_car_wrt_world, sim_state.xh.x, sim_state.xh.y, sim_state.xh.θ, timestamp)

#         vel_msg.header.stamp = accel_msg.header.stamp = timestamp
#         vel_msg.twist.linear.x = sim_state.xh.v
#         accel_msg.accel.linear.x = sim_state.uh.a
#         accel_msg.accel.linear.y = sim_state.uh.κ*sim_state.xh.v^2

#         # p, o = pose_msg.pose.position, pose_msg.pose.orientation
#         # p.x, p.y = sim_state.xh.x, sim_state.xh.y
#         # o.x, o.y, o.z, o.w = tf.transformations[:quaternion_from_euler](0.0, 0.0, sim_state.xh.θ)
#         # v = vel_msg.twist.linear
#         # v.x, v.y = sim_state.xh.v*cos(sim_state.xh.θ), sim_state.xh.v*sin(sim_state.xh.θ)
#         # a = accel_msg.accel.linear
#         # a.x = sim_state.uh.a*cos(sim_state.xh.θ) - sim_state.uh.κ*sim_state.xh.v^2*sin(sim_state.xh.θ)
#         # a.y = sim_state.uh.a*sin(sim_state.xh.θ) + sim_state.uh.κ*sim_state.xh.v^2*cos(sim_state.xh.θ)

#         tf_broadcaster[:sendTransform](human_car_wrt_world)
#         # tf_broadcaster[:sendTransform]((sim_state.xh.x, sim_state.xh.y, 0.0),
#         #                                tf.transformations[:quaternion_from_euler](0.0, 0.0, sim_state.xh.θ),
#         #                                timestamp, "human_car", "world")
#         # publish(human_pose_pub, pose_msg)
#         publish(human_vel_pub, vel_msg)
#         publish(human_accel_pub, accel_msg)
#         rossleep(rate)
#     end
# end

# @spawn spin()

end # module