function init_constraint!(
    com::CS.CoM,
    constraint::XorConstraint,
    fct,
    set::XorSet;
)
    !init_lhs_and_rhs!(com, constraint, fct, set) && return false

    if constraint.complement_lhs !== nothing
        init_constraint!(com, constraint.complement_lhs, constraint.complement_lhs.fct, constraint.complement_lhs.set)
    end
    if constraint.complement_rhs !== nothing
        init_constraint!(com, constraint.complement_rhs, constraint.complement_rhs.fct, constraint.complement_rhs.set)
    end
    return true
end

"""
    function is_constraint_violated(
        com::CoM,
        constraint::BoolConstraint,
        fct,
        set::XorSet,
    )

Check if both of the inner constraints are violated or whether both are solved
"""
function is_constraint_violated(
    com::CoM,
    constraint::BoolConstraint,
    fct,
    set::XorSet,
)
    both_violated = is_lhs_constraint_violated(com, constraint) && is_rhs_constraint_violated(com, constraint)
    both_violated && return true
    lhs_solved = is_constraint_solved(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set)
    rhs_solved = is_constraint_solved(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set)
    return lhs_solved && rhs_solved
end

"""
    still_feasible(com::CoM, constraint::XorConstraint, fct, set::XorSet, vidx::Int, value::Int)

Return whether the constraint can be still fulfilled when setting a variable with index `vidx` to `value`.
**Attention:** This assumes that it isn't violated before.
"""
function still_feasible(
    com::CoM,
    constraint::XorConstraint,
    fct,
    set::XorSet,
    vidx::Int,
    value::Int,
)
    lhs_violated = is_constraint_violated(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set)
    lhs_feasible = !lhs_violated
    lhs_solved = false
    if lhs_feasible
        lhs_solved = is_constraint_solved_when_fixed(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set, vidx, value)
    end
    if lhs_feasible && !lhs_solved
        lhs_indices = constraint.lhs.indices
        for i in 1:length(lhs_indices)
            if lhs_indices[i] == vidx
                lhs_feasible = still_feasible(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set, vidx, value)
                break
            end
        end
    end
    rhs_violated = is_constraint_violated(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set)
    rhs_feasible = !rhs_violated
    rhs_solved = false
    if rhs_feasible
        rhs_solved = is_constraint_solved_when_fixed(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set, vidx, value)
    end
    if rhs_feasible && !rhs_solved
        rhs_indices = constraint.rhs.indices
        for i in 1:length(rhs_indices)
            if rhs_indices[i] == vidx
                rhs_feasible = still_feasible(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set, vidx, value)
                break
            end
        end
    end

    # at least one must be feasible
    if !lhs_feasible
        return rhs_feasible
    elseif !rhs_feasible
        return lhs_feasible
    end
    # not allowed that both are already solved
    return !(lhs_solved && rhs_solved)
end

"""
    prune_constraint!(com::CS.CoM, constraint::XorConstraint, fct, set::XorSet; logs = true)

Reduce the number of possibilities given the `XorConstraint` by pruning both parts
Return whether still feasible
"""
function prune_constraint!(
    com::CS.CoM,
    constraint::XorConstraint,
    fct,
    set::XorSet;
    logs = true,
)
    lhs_violated = is_constraint_violated(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set)
    rhs_violated = is_constraint_violated(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set)
    if lhs_violated && rhs_violated
        return false
    end
    # check if one is already solved
    lhs_solved = is_constraint_solved(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set)
    rhs_solved = is_constraint_solved(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set)
    if lhs_solved && rhs_solved
        return false
    end

    # if one is solved => complement prune the other
    # Todo implement for activated complement constraints
    if lhs_solved && constraint.complement_rhs !== nothing
        return prune_constraint!(com, constraint.complement_rhs, constraint.complement_rhs.fct, constraint.complement_rhs.set; logs=logs)
    end
    if rhs_solved && constraint.complement_lhs !== nothing
        return prune_constraint!(com, constraint.complement_lhs, constraint.complement_lhs.fct, constraint.complement_lhs.set; logs=logs)
    end

    # if both aren't solved yet we only prune if one is violated
    if !lhs_solved && !rhs_solved
        if lhs_violated
            activate_rhs!(com, constraint)
            return prune_constraint!(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set; logs=logs)
        end
        if rhs_violated
            activate_lhs!(com, constraint)
            return prune_constraint!(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set; logs=logs)
        end
    end

    return true
end


function SphericalCoordinates(arg0::Vector3D)
    return SphericalCoordinates((Vector3D,), arg0)
end

function SphericalCoordinates(arg0::jdouble, arg1::jdouble, arg2::jdouble)
    return SphericalCoordinates((jdouble, jdouble, jdouble), arg0, arg1, arg2)
end

function get_cartesian(obj::SphericalCoordinates)
    return jcall(obj, "getCartesian", Vector3D, ())
end

function get_phi(obj::SphericalCoordinates)
    return jcall(obj, "getPhi", jdouble, ())
end

function get_r(obj::SphericalCoordinates)
    return jcall(obj, "getR", jdouble, ())
end

function get_theta(obj::SphericalCoordinates)
    return jcall(obj, "getTheta", jdouble, ())
end

function to_cartesian_gradient(obj::SphericalCoordinates, arg0::Vector{jdouble})
    return jcall(obj, "toCartesianGradient", Vector{jdouble}, (Vector{jdouble},), arg0)
end

function to_cartesian_hessian(obj::SphericalCoordinates, arg0::Vector{Vector{jdouble}}, arg1::Vector{jdouble})
    return jcall(obj, "toCartesianHessian", Vector{Vector{jdouble}}, (Vector{Vector{jdouble}}, Vector{jdouble}), arg0, arg1)
end


using MAT, JLD2

@inline function string_as_varname(s::AbstractString)
    S = Symbol(s)
    @eval $S
end

#@inline function test_against_matlab(filename::String)
    vars = matread(filename)
    for v in vars
        cur_type = typeof(v[2])
        if (v[1] == "t0")
            continue
        elseif v[1] == "a_grid"
            @assert v[2] ≈ a_grid
        elseif v[1] == "a_g_grid"
            @assert v[2] ≈ a_g_grid
        elseif v[1] == "b_grid"
            @assert v[2] ≈ b_grid
        elseif v[1] == "y_grid"
            @assert v[2] ≈ y_grid
        elseif v[1] == "y_g_grid"
            @assert v[2] ≈ y_g_grid
        elseif v[1] == "r_a_grid"
            @assert v[2] ≈ r_a_grid
        elseif v[1] == "r_b_grid"
            @assert v[2] ≈ r_b_grid
        elseif v[1] == "r_a_g_grid"
            @assert v[2] ≈ r_a_g_grid
        elseif v[1] == "r_b_g_grid"
            @assert v[2] ≈ r_b_g_grid
        elseif v[1] == "daf_g_grid"
            @assert v[2] ≈ daf_g_grid
        elseif v[1] == "daf_grid"
            @assert v[2] ≈ daf_grid
        elseif v[1] == "dab_grid"
            @assert v[2] ≈ dab_grid
        elseif v[1] == "dab_g_grid"
            @assert v[2] ≈ dab_g_grid
        elseif v[1] == "dab_g_tilde_grid"
            @assert v[2] ≈ dab_g_tilde_grid
        elseif v[1] == "dbf_grid"
            @assert v[2] ≈ dbf_grid
        elseif v[1] == "dbf_g_grid"
            @assert v[2] ≈ dbf_g_grid
        elseif v[1] == "dbb_grid"
            @assert v[2] ≈ dbb_grid
        elseif v[1] == "dbb_g_grid"
            @assert v[2] ≈ dbb_g_grid
        elseif v[1] == "trans_grid"
            @assert v[2] ≈ trans_grid
        elseif v[1] == "l_g_grid"
            @assert v[2] ≈ l_g_grid
        elseif v[1] == "w_grid"
            @assert v[2] ≈ w_grid
        elseif v[1] == "c_0"
            @assert v[2] ≈ c_0
        elseif v[1] == "V_0"
            @assert v[2] ≈ V_0
        elseif v[1] == "gg0"
            @assert v[2] ≈ gg0
        elseif v[1] == "gg"
            @assert v[2] ≈ gg
        elseif v[1] == "r_b"
            @assert v[2] ≈ r_b
        elseif v[1] == "r_b_borr"
            @assert v[2] ≈ r_b_borr
        elseif v[1] == "y_dist"
            @assert v[2] ≈ y_dist
        elseif v[1] == "y_mean"
            @assert v[2] ≈ y_mean
        elseif v[1] == "aalpha"
            @assert v[2] ≈ aalpha
        elseif v[1] == "ddeath"
            @assert v[2] ≈ ddeath
        elseif v[1] == "ddelta"
            @assert v[2] ≈ ddelta
        elseif v[1] == "rrho"
            @assert v[2] ≈ rrho
        elseif v[1] == "chi0"
            @assert v[2] ≈ chi0
        elseif v[1] == "chi1"
            @assert v[2] ≈ chi1
        elseif v[1] == "chi2"
            @assert v[2] ≈ chi2
        elseif v[1] == "a_lb"
            @assert v[2] ≈ a_lb
        elseif v[1] == "pam"
            @assert v[2] ≈ pam
        elseif v[1] == "xxi"
            @assert v[2] ≈ xxi
        elseif v[1] == "ggamma"
            @assert v[2] ≈ ggamma
        elseif v[1] == "tau_I"
            @assert v[2] ≈ tau_I
        elseif v[1] == "trans"
            @assert v[2] ≈ trans
        elseif v[1] == "lambda"
            @assert v[2] ≈ lambda
        elseif v[1] == "K_liquid"
            @assert v[2] ≈ K_liquid
        elseif v[1] == "aggregate_variables"
            @assert v[2] ≈ aggregate_variables
        elseif v[1] == "distributional_variables"
            @assert v[2] ≈ distributional_variables
        elseif v[1] == "distributional_variables_1"
            #@assert v[2] ≈ distributional_variables_1
            continue
        else
            @show "NOT PRINTING", v[1]
            #=
            try
                if !(string_as_varname(v[1]) ≈ v[2])
                    @show "NEQ", v[2]
                else
                    @show "WORKSWORKSWORKS", v[2]
                end
                @assert string_as_varname(v[1]) ≈ v[2]
            catch err
                if isa(err, UndefVarError)
                    println(err)
                else
                    throw(err)
                end
            end
            =#
        end
    end
    println("All tests passed for: ", filename)
#end


## Auto-correlation
plot(
    autocor(chains[1].planets[1].e, 1:100:1000)
)

## Trace plots
plot()
plot!(chains[1].planets[1].a[1:end], title="d")
plot!(chains[1].planets[2].a[1:end], title="e")
## Trace plots
plot()
plot!(chains[1].planets[1].e[1:end], title="d")
plot!(chains[1].planets[2].e[1:end], title="e")

##
mean(getproperty.(stats[1], :acceptance_rate))
mean(getproperty.(stats[1], :tree_depth))
maximum(getproperty.(stats[1], :tree_depth))
count(getproperty.(stats[1], :numerical_error))
mean(getproperty.(stats[1], :n_steps))
maximum(getproperty.(stats[1], :log_density))

##
cd(@__DIR__)


push!(LOAD_PATH, "/Users/megorov/Desktop/projects/stanford/sniper/egorov/sniperRepo/new-version")

using Collada_
using UrbanMap_
using SniperPOMDP_
using SniperMDP_
using ValueIterations
using Helpers
using RayCaster_
using Simulator
using SimServers
using RolloutPolicies

map = UrbanMap(1, [0.5], [0.5], [0.42], [0.22], 10, 10)
#t   = [(2.0,8.0), (2.0,1.0), (5.0,2.0), (7.0,3.0)]
t   = [(2.0,8.0), (2.0,1.0)]
#t   = [(7.0,7.0), (2.0,14.0), (8.0,4.0), (11.0,13.0)]
mdp = SniperMDP(map, nSnipers = 1, nMonitors = 2, targets=t)

policy = RolloutPolicy(mdp, d=10, n=3000, c=0.4)

# create the simulation type
#init  = 8_888_888_803 # sniper at (3,1), all monitors at (9,9)
init = 470_872
sim  = MDPSimulation(mdp, init, 100)

# set up the server
socket = 4444
server = SniperServer(socket)

# begin simulation and send results to client
#simAndSend(server, sim, mdp, policy)

# batch sim
nSims = 10
obs = zeros(Int64, nSims)
kills = zeros(Int64, nSims)
rewards = zeros(n)

for i = 1:nSims
    println("On sim: $i")
    res = simulate(policy, mdp, 100, initState=init)
    rewards[i] = sum(res.rewards)
    obs[i] = res.nObs
    kills[i] = res.nKills
end


using Toxcore
using GLVisualize
using GLWindow
using Reactive
using GeometryTypes
using GLAbstraction
using ColorTypes
using FileIO
using ImageIO
using MeshIO
using WavefrontObj
using Meshes
using NPZ
using GLFW
using ModernGL

include("helper.jl")
function xy_data(x,y,i, N)
    x = ((x/N)-0.5f0)*i
    y = ((y/N)-0.5f0)*i
    r = sqrt(x*x + y*y)
    Float32(sin(r)/r)
end
generate(i, N) = Float32[xy_data(Float32(x),Float32(y),Float32(i), N) for x=1:N, y=1:N]

immutable Message
    leftaligned::Bool # kinda silly place, but it works for now
    color::RGBA{Float32}
    text::UTF8String
end
global const MESSAGES = Input(Message(false, RGBA(1f0,0.2f0,0.5f0,0.7f0), "msg"))

function message_area(yoffset, y_scroll, m_area, screen_width, screen_height, left_aligned)
    yposition    = yoffset+(round(Int, y_scroll)*50) # for scrolling and positioning
    xposition    = left_aligned ? 50 : (m_area.w - 50 - screen_width)
    Rectangle(xposition, yposition, screen_width, screen_height)
end
const edit_value_names = (
    :color,
    :grid_min,
    :grid_max,
    :algorithm,
    :style,
    :technique,
    :scale,
    :norm,
    :absorption,
    :light_position,
    :isovalue
)

is_screen_active(v0, ek, esc) = (esc && return false; ek && return true; v0)

function create_edit_screen(parent, robj, namefilter=edit_value_names)
    frame       = 20
    key_string  = ""
    yoffset     = 10
    vizzes      = RenderObject[]
    max_screenwidth = 0
    screenheigth    = 0
    for (key, value) in filter((k,v)->in(k, namefilter), robj.uniforms)
        if applicable(vizzedit, value, parent.inputs)

            signal, editvizz = vizzedit(value, parent.inputs)
            robj[key]        = signal
            keylabel         = visualize(string(key)*": ")

            boundingbox      = keylabel.boundingbox.value
            bbsize           = boundingbox.maximum - boundingbox.minimum
            key_width        = round(Int, bbsize[1])
            key_height       = round(Int, bbsize[2])
            move             = -boundingbox.minimum+Float32(frame/2f0)+Vec3f0(0, yoffset, 0)
            keylabel[:model] = translationmatrix(Vec3f0(move[1], move[2], 0f0))
            keylabel[:preferred_camera] = :fixed_pixel

            boundingbox      = editvizz.boundingbox.value
            bbsize           = boundingbox.maximum - boundingbox.minimum
            edit_width       = round(Int, bbsize[1])
            edit_height      = round(Int, bbsize[2])
            move             = -boundingbox.minimum+Float32(frame/2f0)+Vec3f0(key_width, yoffset, 0)
            editvizz[:model] = translationmatrix(Vec3f0(move[1], move[2], 0f0))
            editvizz[:preferred_camera] = :fixed_pixel
            push!(vizzes, editvizz, keylabel)

            yoffset         += max(edit_height, key_height)+50
            max_screenwidth = max(max_screenwidth, key_width+edit_width)
        end
    end

    buttonspressed  = parent.inputs[:buttonspressed]

    edit_key            = lift(==, buttonspressed, [GLFW.KEY_LEFT_CONTROL, GLFW.KEY_E])
    escape_key          = lift(==, buttonspressed, [GLFW.KEY_ESCAPE])
    edit_key_active     = keepwhen(edit_key, false, edit_key)

    mouse_at_edit_time  = sampleon(edit_key_active, parent.inputs[:mouseposition])
    is_active           = foldl(is_screen_active, false, edit_key, escape_key)
    area                = lift(mouse_at_edit_time, is_active) do mousepos, active
        !active && return Rectangle(0,0,0,0)
        Rectangle(round(Int, mousepos[1]), round(Int, mousepos[2]), max_screenwidth+frame, yoffset)
    end
    edit_screen         = Screen(parent, area=area)
    view(
        visualize(lift(zeroposition, area), color=RGBA(0f0,0f0,0f0,1f0)),
        edit_screen, method=:fixed_pixel
    )
    view(vizzes, edit_screen)
end

function create_screens(yoffset, robjs_alignment, scroll_window, main_screen)
    frame = 40
    robjs, alignement, color = robjs_alignment
    for robj in robjs
        yoffset -= 30
        camera_position = Vec3f0(2)
        camera_lookat   = Vec3f0(0)
        boundingbox     = robj.boundingbox.value
        bbsize          = boundingbox.maximum - boundingbox.minimum
        if endswith(string(robj[:preferred_camera]), "_pixel") # terrible way of working around the lack of a unit system
            screen_height           = round(Int, bbsize[2])+frame
            screen_width            = round(Int, bbsize[1])+frame
            move                    = -boundingbox.minimum+Float32(frame/2f0)
            robj[:model]            = translationmatrix(Vec3f0(move[1], move[2], 0f0))
            robj[:preferred_camera] = :fixed_pixel #better fixate this
        else
            screen_height     = 700 # for now, all 3D windows will have a fixed size
            screen_width      = lift(x->round(Int, x.w*0.9), main_screen.area)
            camera_position   = boundingbox.minimum+(bbsize*1.5f0)
            camera_lookat     = boundingbox.minimum+(bbsize*0.5f0)
        end
        yoffset     -= screen_height
        area        = lift(message_area, yoffset, scroll_window, main_screen.area, screen_width, screen_height, alignement)
        new_screen  = Screen(main_screen, area=area, position=camera_position, lookat=camera_lookat)
        #view(
        #    visualize(lift(zeroposition, area), color=color),
        #    new_screen, method=:fixed_pixel
        #)
        view(robj, new_screen)
        #create_edit_screen(new_screen, robj)
    end
    yoffset
end

function handle_drop(files::Vector{UTF8String})
    files = map(File, files)
    (RenderObject[visualize(f) for f in files], false, RGBA(1f0,1f0,1f0,0.7f0))
end

function visualize_source(source)
    value = save_eval(source)
    if isa(value, RenderObject)
        return value
    else
        try
            return visualize(value)
        catch e # ultra terrible way of trying to find out if there is a default visulization
            return visualize("error: $e", preferred_camera=:fixed_pixel)
        end
    end
end

function GLVisualize.visualize(julia_source::File{:jl})
    source = readall(abspath(julia_source))
    visualize_source(source)
end
function GLVisualize.visualize(npzfile::File{:npz}, style::Symbol=:default; kw_args...)
    volume = npzread(abspath(npzfile))["data"]
    volume = volume./256f0
    visualize(volume, Style{style}(), visualize_default(volume, style, kw_args))
end

function visualize_message(msg::Message)
    text = msg.text
    vizz = "error"
    if startswith(text, "jl\"") && endswith(text, "\"") # its julia code
        vizz = visualize_source(text[4:end-1])
    else # its just text
        vizz = visualize(text, preferred_camera=:fixed_pixel)
    end
    return ([vizz], msg.leftaligned, msg.color) # this is what the insertinto screen function wants!
end

function main()
    RS, renderloop = Screen()
    # Try to load the tox settings
    my_tox = 0

    try
        savefile = open(Pkg.dir("Toxcore", "test", "bot_savedata.binary"), "r")
        savedata = readbytes(savefile)
        close(savefile)

        default_options = tox_options_default()

        options = Tox_Options(default_options.ipv6_enabled,
                            default_options.udp_enabled,
                            default_options.proxy_type,
                            default_options.proxy_host,
                            default_options.proxy_port,
                            default_options.start_port,
                            default_options.end_port,
                            default_options.tcp_port,
                            TOX_SAVEDATA_TYPE_TOX_SAVE,
                            pointer(savedata),
                            length(savedata))

        my_tox = tox_new(options)

        info("Previous bot instance found. Reusing it!")
    catch e
        println(e)
        #Create a default Tox
        my_tox = tox_new()
        info("Created new bot instance")
    end

    # register the callbacks
    tox_callback_friend_request(my_tox, OnToxFriendRequest, C_NULL)
    tox_callback_friend_message(my_tox, OnToxFriendMessage, C_NULL)

    # print own address
    info("Here is my address")
    println(convert(ASCIIString, tox_self_get_address(my_tox)))

    # define user details
	tox_self_set_name(my_tox, "SamD")
    tox_self_set_status_message(my_tox, "Whatup yo!")
    Toxcore.CInterface.tox_self_set_status(my_tox, Toxcore.CInterface.TOX_USER_STATUS_NONE)
    # bootstrap from the node defined above
    if !tox_bootstrap(my_tox)
        println("Failed to bootstrap.")
        exit()
    end

    # get the friend list
    friendlist = tox_self_get_friend_list(my_tox)
    info("I have $(length(friendlist)) friend(s)")
    currentfriend = Cuint(0)
    for friend in friendlist
        fname = tox_friend_get_name(my_tox, friend)
        println(fname)
        fname == "SimonD" && (currentfriend = Cuint(friend))
    end
    println(currentfriend)
    println(tox_friend_get_name(my_tox, currentfriend))

    write_area, main_area   = y_partition(RS.area, 20.0)
    write_screen            = Screen(RS, area=write_area)
    main_screen             = Screen(RS, area=main_area)

    text = visualize("write something \n", model=translationmatrix(Vec3f0(0, write_area.value.h-40,0)))
    background, cursor_robj, text_sig = vizzedit(text[:glyphs], text, write_screen.inputs)
    buttonspressed  = write_screen.inputs[:buttonspressed]
    enter_pressed   = lift(==, buttonspressed, [GLFW.KEY_ENTER])
    system_message  = Input(utf8("jl\"file\"osmosisheader.png\"\""))
    message         = sampleon(keepwhen(enter_pressed, true, enter_pressed), text_sig)
    message         = merge(system_message, message)

    mymessages      = lift(message) do msg
        tox_friend_send_message(my_tox, currentfriend, msg)
        Message(false, RGBA(1f0,1f0,1f0,0.7f0), msg)
    end

    view(visualize(lift(x->Rectangle(0,0,x.w, x.h), write_screen.area), color=RGBA(0f0,0f0,0f0,0.7f0)), write_screen, method=:fixed_pixel)

    view(background,    write_screen)
    view(text,          write_screen)
    view(cursor_robj,   write_screen)

    scroll_window   = foldl(+, 0.0, keepwhen(main_screen.inputs[:mouseinside], 0.0, main_screen.inputs[:scroll_y]))

    drop_robjs      = lift(handle_drop, RS.inputs[:droppedfiles])

    all_msgs        = merge(mymessages, MESSAGES)

    msg_robj        = lift(visualize_message, all_msgs)
    all_robjs       = merge(msg_robj, drop_robjs)

    foldl(
        create_screens,
        main_area.value.h,
        all_robjs,
        Input(scroll_window), Input(main_screen)
    )
    push!(system_message, utf8("jl\"file\"osmosisheader.png\"\""))

    glClearColor(1,1,1,1)
    renderloop(()-> tox_iterate(my_tox))

    savedata = tox_get_savedata(my_tox)
    tox_kill(my_tox)

    savefile = open(Pkg.dir("Toxcore", "test", "bot_savedata.binary"), "w")
    write(savefile, savedata)
    close(savefile)

    info("Bot instance saved")

end

main()


using Zygote
using Zygote: @adjoint, literal_getproperty
import Zygote: accum

import Distances: pairwise, colwise
const dtol = 1e-12 # threshold value for precise recalculation of distances

@nograd MersenneTwister, propertynames, broadcast_shape

function accum(D::Diagonal{T}, B::AbstractMatrix) where {T}
    A = Matrix{Union{T, Nothing}}(undef, size(D))
    A[diagind(A)] .= D.diag
    return accum(A, B)
end
accum(A::AbstractMatrix, D::Diagonal) = accum(D, A)
accum(A::Diagonal, B::Diagonal) = Diagonal(accum(diag(A), diag(B)))

@adjoint function ZygoteRules.literal_getproperty(C::Cholesky, ::Val{:factors})
    error("@adjoint not implemented for :factors as is unsafe.")
    return literal_getproperty(C, Val(:factors)), function(Δ)
        error("@adjoint not implemented for :factors. (I couldn't make it work...)")
    end
end

import LinearAlgebra: HermOrSym, diag, Diagonal

diag(S::Symmetric{T, <:Diagonal{T}} where T) = S.data.diag
Zygote._symmetric_back(Δ::Diagonal) = Δ

# Diagonal matrices are always symmetric...
cholesky(A::HermOrSym{T, <:Diagonal{T}} where T) = cholesky(Diagonal(diag(A)))

#
# Some very specific broadcasting hacks while Zygote has crappy broadcasting.
#

import Base.Broadcast: broadcasted

function rrule(::typeof(broadcasted), ::typeof(-), x::AbstractArray)
    function broadcasted_minus_pullback(Δ)
        return (NO_FIELDS, DoesNotExist(), .-Δ)
    end
    return .-x, broadcasted_minus_pullback
end

function rrule(::typeof(broadcasted), ::typeof(exp), x::AbstractArray)
    y = exp.(x)
    return y, Δ->(NO_FIELDS, DoesNotExist(), Δ .* y)
end


# An SVG backend for compose.

export SVG

# Format a floating point number into a decimal string of reasonable precision.
function svg_fmt_float(x::Float64)
    # All svg (in our use) coordinates are in millimeters. This number gives the
    # largest deviation from the true position allowed in millimeters.
    const eps = 0.01
    a = @sprintf("%0.8f", round(x / eps) * eps)
    n = length(a)

    while a[n] == '0'
        n -= 1
    end

    if a[n] == '.'
        n -= 1
    end

    a[1:n]
end

# Format a color for SVG.
svg_fmt_color(c::ColorValue) = @sprintf("#%s", hex(c))
svg_fmt_color(c::Nothing) = "none"

# The SVG image we generate are going to use millimeters everywhere.
type SVGMeasure <: NativeMeasure
    value::Float64
end

*(u::Float64, v::SVGMeasure) = SVGMeasure(u * v.value)
+(u::SVGMeasure, v::SVGMeasure) = SVGMeasure(u.value + v.value)
-(u::SVGMeasure, v::SVGMeasure) = SVGMeasure(u.value - v.value)
convert(::Type{Float64}, u::SVGMeasure) = u.value
convert(::Type{SVGMeasure}, u::Float64) = SVGMeasure(u)
function convert(::Type{SimpleMeasure{MillimeterUnit}}, u::SVGMeasure)
    SimpleMeasure{MillimeterUnit}(u.value)
end

type SVG <: Backend
    # Image size in millimeters.
    width::SVGMeasure
    height::SVGMeasure

    # Output stream.
    f::IO

    # Should f be closed when the backend is finished?
    close_stream::Bool

    # Current level of indentation.
    indentation::Int

    # Javascript fragments are placed in a function with a unique name. This is
    # a map of unique function names to javascript code.
    scripts::Dict{String, String}

    # Clip-paths that need to be defined at the end of the document.
    clippaths::Vector{Vector{Point}}

    # Miscelaneous embedded objects included immediately before the </svg> tag,
    # such as extra javascript or css.
    embobj::Set{String}

    # Keep track of which properties that are push are empty to we can avoid
    # printiing them.
    empty_properties::Vector{Bool}

    # Keep track of which properties have links, so we know when to insert
    # closing </a> tags.
    linked_properties::Vector{Bool}

    # Keep track of mask objects.
    mask_properties::Vector{Bool}

    # Emit the graphic on finish when writing to a buffer.
    emit_on_finish::Bool

    function SVG(f::IO,
                 width::MeasureOrNumber,
                 height::MeasureOrNumber,
                 emit_on_finish::Bool=true)
        img = new()
        img.width  = native_measure(width,  img)
        img.height = native_measure(height, img)
        img.f = f
        img.close_stream = false
        img.indentation = 0
        img.clippaths = Array(Vector{Point}, 0)
        img.scripts = Dict{String, String}()
        img.embobj = Set{String}()
        img.empty_properties = Array(Bool, 0)
        img.linked_properties = Array(Bool, 0)
        img.mask_properties = Array(Bool, 0)
        img.emit_on_finish = emit_on_finish

        write(img.f, @sprintf("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" version=\"1.1\" width=\"%smm\" height=\"%smm\" viewBox=\"0 0 %s %s\" style=\"stroke:black;fill:black\" stroke-width=\"0.5\">\n",
              svg_fmt_float(img.width.value), svg_fmt_float(img.height.value),
              svg_fmt_float(img.width.value), svg_fmt_float(img.height.value)))
        img
    end

    # Write to a file.
    function SVG(filename::String,
                 width::MeasureOrNumber,
                 height::MeasureOrNumber)
        f = open(filename, "w")
        img = SVG(f, width, height)
        img.close_stream = true
        img
    end

    # Write to buffer.
    function SVG(width::MeasureOrNumber,
                 height::MeasureOrNumber;
                 emit_on_finish::Bool=true)
        img = SVG(IOString(), width, height, emit_on_finish=emit_on_finish)
        img.close_stream = false
        img
    end
end

function finish(img::SVG)
    for obj in img.embobj
        write(img.f, obj)
        write(img.f, "\n")
    end

    if length(img.scripts) > 0
        write(img.f, "<script type=\"application/ecmascript\"><![CDATA[\n")
        for (fn_name, js) in img.scripts
            @printf(img.f, "function %s(evt) {\n%s\n}\n\n", fn_name, js)
        end
        write(img.f, "]]></script>\n")
    end

    if length(img.clippaths) > 0
        write(img.f, "<defs>\n")
        for (i, clippath) in enumerate(img.clippaths)
            write(img.f, "<clipPath id=\"clippath$(i)\">\n  <path d=\"")
            print_svg_path(img.f, clippath)
            write(img.f, "\" />\n</clipPath\n>")
        end
        write(img.f, "</defs>\n")
    end

    write(img.f, "</svg>\n")
    if img.close_stream
        close(img.f)
    end

    # If we are writing to a buffer. Collect the string and emit it.
    if img.emit_on_finish && typeof(img.f) == IOString
        seek(img.f, 0)
        emit(Emitable("image/svg+xml", readall(img.f)))
        close(img.f)
    end
end

function root_box(img::SVG)
    NativeBoundingBox(
        SVGMeasure(0.0),
        SVGMeasure(0.0),
        img.width,
        img.height)
end

native_zero(::SVG) = SVGMeasure(0.0)

function indent(img::SVG)
    for i in 1:img.indentation
        write(img.f, "  ")
    end
end

function native_measure(u::SimpleMeasure{PixelUnit},
                        img::SVG)
    native_measure(convert(SimpleMeasure{MillimeterUnit}, u), img)
end

function native_measure(u::SimpleMeasure{MillimeterUnit},
                        img::SVG)
    SVGMeasure(u.value)
end

# Draw

function print_svg_path(out, points::Vector{Point})
    @printf(out, "M%s,%s L",
            svg_fmt_float(points[1].x.value),
            svg_fmt_float(points[1].y.value))
    for point in points[2:]
        @printf(out, " %s %s",
                svg_fmt_float(point.x.value),
                svg_fmt_float(point.y.value))
    end
end

function draw(img::SVG, form::Lines)
    n = length(form.points)
    if n <= 1; return; end
    indent(img)
    paths = make_paths(form.points)
    for path in paths
        write(img.f, "<path d=\"")
        print_svg_path(img.f, path)
        write(img.f, "\" />\n")
    end
end

function draw(img::SVG, form::Curve)
    indent(img)
    @printf(img.f, "<path d=\"M%s,%s C%s,%s %s,%s %s,%s\" />\n",
        svg_fmt_float(form.anchor0.x.value),
        svg_fmt_float(form.anchor0.y.value),
        svg_fmt_float(form.ctrl0.x.value),
        svg_fmt_float(form.ctrl0.y.value),
        svg_fmt_float(form.ctrl1.x.value),
        svg_fmt_float(form.ctrl1.y.value),
        svg_fmt_float(form.anchor1.x.value),
        svg_fmt_float(form.anchor1.y.value))
end

function draw(img::SVG, form::Polygon)
    n = length(form.points)
    if n <= 1; return; end

    indent(img)
    write(img.f, "<path d=\"")
    print_svg_path(img.f, form.points)
    write(img.f, " z\" />\n")
end

minmax(a, b) = a < b ? (a,b) : (b,a)

function draw(img::SVG, form::Ellipse)
    cx = form.center.x.value
    cy = form.center.y.value
    rx = sqrt((form.x_point.x.value - cx)^2 +
              (form.x_point.y.value - cy)^2)
    ry = sqrt((form.y_point.x.value - cx)^2 +
              (form.y_point.y.value - cy)^2)
    theta = radians2degrees(atan2(form.x_point.y.value - cy,
                                  form.x_point.x.value - cx))

    indent(img)
    eps = 1e-6

    if abs(rx - ry) < eps
        @printf(img.f, "<circle cx=\"%s\" cy=\"%s\" r=\"%s\" />\n",
                svg_fmt_float(cx), svg_fmt_float(cy), svg_fmt_float(rx))
    else
        @printf(img.f, "<ellipse cx=\"%s\" cy=\"%s\" rx=\"%s\" ry=\"%s\"",
                svg_fmt_float(cx), svg_fmt_float(cy),
                svg_fmt_float(rx), svg_fmt_float(ry))

        if abs(theta) >= eps
            @printf(img.f, " transform=\"rotate(%s %s %s)\"",
                    svg_fmt_float(theta),
                    svg_fmt_float(cx),
                    svg_fmt_float(cy))
        end

        write(img.f, " />\n")
    end
end

function draw(img::SVG, form::Text)
    indent(img)
    @printf(img.f, "<text x=\"%s\" y=\"%s\"",
            svg_fmt_float(form.pos.x.value),
            svg_fmt_float(form.pos.y.value))

    if is(form.halign, hcenter)
        print(img.f, " text-anchor=\"middle\"")
    elseif is(form.halign, hright)
        print(img.f, " text-anchor=\"end\"")
    end

    if is(form.valign, vcenter)
        print(img.f, " style=\"dominant-baseline:central\"")
    elseif is(form.valign, vtop)
        print(img.f, " style=\"dominant-baseline:text-before-edge\"")
    end

    if !isidentity(form.t)
        @printf(img.f, " transform=\"rotate(%s, %s, %s)\"",
                svg_fmt_float(radians2degrees(atan2(form.t.M[2,1], form.t.M[1,1]))),
                svg_fmt_float(form.pos.x.value),
                svg_fmt_float(form.pos.y.value))
    end

    # TODO: escape special characters
    @printf(img.f, ">%s</text>\n",
            pango_to_svg(form.value))
end

# Applying properties

function push_property(img::SVG, property::Property)
    if property === empty_property
        push!(img.empty_properties, true)
        push!(img.linked_properties, false)
        push!(img.mask_properties, false)
    else
        push!(img.empty_properties, false)
        indent(img)

        # There can only be one property of each type. E.g, defining 'fill'
        # twice in the same element is not valid svg.
        properties = Dict{Type, PropertyPrimitive}()
        rawproperties = {}
        p = property
        while !is(p, empty_property)
            if typeof(p.primitive) == SVGAttribute
                push!(rawproperties, p.primitive)
            else
                properties[typeof(p.primitive)] = p.primitive
            end
            p = p.next
        end

        # Special-case for the mask property.
        if haskey(properties, SVGDefineMask)
            write(img.f, @sprintf("<mask id=\"%s\">", properties[SVGDefineMask].id))
            push!(img.mask_properties, true)
        else
            push!(img.mask_properties, false)
        end

        # Special-case for links.
        if haskey(properties, SVGLink)
            write(img.f, @sprintf("<a xlink:href=\"%s\">",
                  properties[SVGLink].target === nothing ? '#' : properties[SVGLink].target))
            push!(img.linked_properties, true)
        else
            push!(img.linked_properties, false)
        end

        write(img.f, "<g")
        for p in chain(values(properties), rawproperties)
            apply_property(img, p)
        end
        write(img.f, ">\n")

        img.indentation += 1
    end
end

function pop_property(img::SVG)
    is_linked = pop!(img.linked_properties)
    is_mask = pop!(img.mask_properties)
    if !pop!(img.empty_properties)
        img.indentation -= 1
        indent(img)
        write(img.f, "</g>")
        if is_linked
            write(img.f, "</a>")
        end
        if is_mask
            write(img.f, "</mask>")
        end
        write(img.f, "\n")
    end
end

# Nop catchall
function apply_property(img::SVG, p::PropertyPrimitive)
end

function apply_property(img::SVG, p::Stroke)
    @printf(img.f, " stroke=\"%s\"", svg_fmt_color(p.value))
end

function apply_property(img::SVG, p::Fill)
    @printf(img.f, " fill=\"%s\"", svg_fmt_color(p.value))
end

function apply_property(img::SVG, p::LineWidth)
    @printf(img.f, " stroke-width=\"%s\"", svg_fmt_float(p.value.value))
end

function apply_property(img::SVG, p::Visible)
    @printf(img.f, " visibility=\"%s\"",
            p.value ? "visible" : "hidden")
end

function apply_property(img::SVG, p::Opacity)
    @printf(img.f, " opacity=\"%s\"", fmt_float(p.value))
end

function apply_property(img::SVG, p::SVGID)
    @printf(img.f, " id=\"%s\"", escape_string(p.value))
end

function apply_property(img::SVG, p::SVGClass)
    @printf(img.f, " class=\"%s\"", escape_string(p.value))
end

function apply_property(img::SVG, p::Font)
    @printf(img.f, " font-family=\"%s\"", escape_string(p.family))
end

function apply_property(img::SVG, p::FontSize)
    @printf(img.f, " font-size=\"%s\"", svg_fmt_float(p.value.value))
end

function apply_property(img::SVG, p::Clip)
    push!(img.clippaths, p.points)
    i = length(img.clippaths)
    write(img.f, " clip-path=\"url(#clippath$(i))\"")
end

function apply_property(img::SVG, p::SVGEmbed)
    add!(img.embobj, p.markup)
end

function apply_property(img::SVG, p::SVGMask)
    @printf(img.f, " mask=\"url(#%s)\"", p.id)
end

function apply_property(img::SVG, p::SVGAttribute)
    @printf(img.f, " %s=\"%s\"", p.attribute, p.value)
end

for event in events
    prop_name = lowercase(string(event))
    @eval begin
        function apply_property(img::SVG, p::($event))
            fn_name = add_script(img, p.value, object_id(p.value))
            @printf(img.f, " %s=\"%s(evt)\"", ($prop_name), fn_name)
        end
    end
end

# Add some javascript. Return the unique name generated for the wrapper function
# containing the given code.
function add_script(img::SVG, js::String, obj_id::Integer)
    fn_name = @sprintf("js_chunk_%x", obj_id)
    if !haskey(img.scripts, fn_name)
        img.scripts[fn_name] = replace(js, r"^"m, "  ")
    end
    fn_name
end


using Multilane
using StatsBase
using MCTS
using JLD

behaviors=Dict{String,Any}(
    "correlated" => standard_uniform(1.0, correlated=true)
)

dpws = DPWSolver(depth=20,
                 n_iterations=500,
                 exploration_constant=50.0,
                 k_state=4.0,
                 alpha_state=1/8,
                 rollout_solver=SimpleSolver())

solvers = Dict{String, Any}(
    "dpw"=>dpws,
    "assume_normal"=>SingleBehaviorSolver(dpws, Multilane.NORMAL)
)

curve = TestSet(lambda=[1.0, 2.0, 4.0, 8.0, 16.0, 32.0], N=500, p_appear=1.0)

tests = []
for b in [2.5, 3.0, 3.5, 4.0]
    push!(tests, TestSet(curve, solver_key="dpw",
                         behaviors="correlated",
                         key=@sprintf("upper_bound_%02d", b*10),
                         brake_threshold=b))
    push!(tests, TestSet(curve, solver_key="assume_normal",
                         behaviors="correlated",
                         key=@sprintf("assume_normal_%02d", b*10),
                         brake_threshold=b))
end

objects = gen_initials(tests, behaviors=behaviors, generate_physical=true)

@show objects["param_table"]
objects["solvers"] = solvers

objects["note"] = "Tests over a range of braking values with correlated behaviors."

sbatch_spawn(tests, objects,
             batch_size=50,
             time_per_batch="1:00:00",
             submit_command="sbatch",
             template_name="sherlock.sh")

# sbatch_spawn(tests, objects,
#              batch_size=12,
#              time_per_batch="1:00:00",
#              submit_command="bash",
#              template_name="theresa.sh")
#
# results = evaluate(tests, objects, parallel=true)
#
# filename = string("results_", Dates.format(Dates.now(),"u_d_HH_MM"), ".jld")
# results["histories"] = nothing
# save(filename, results)
# println("results saved to $filename")


using ParameterJuMP
using MathOptInterface
const MOI = MathOptInterface

export BendersBridge

function parametrized_function(var::StructuredVariableRef,
                               structured_model::StructuredModel,
                               model::JuMP.Model,
                               variable_map::Dict{Int, JuMP.VariableRef},
                               parameter_map::Dict{Int, ParameterJuMP.ParameterRef})
    if var.model === structured_model
        return variable_map[var.idx]
    elseif var.model === structured_model.parent
        if !haskey(parameter_map, var.idx)
            parameter_map[var.idx] = ParameterJuMP.add_parameter(model)
        end
        return parameter_map[var.idx]
    else
        error("A StructuredModel cannot have expression using variables of a StructuredModel different from itself or its parent")
    end
end
function parametrized_function(aff::JuMP.GenericAffExpr{C, StructuredVariableRef},
                               structured_model::StructuredModel,
                               model::JuMP.Model,
                               variable_map::Dict{Int, JuMP.VariableRef},
                               parameter_map::Dict{Int, ParameterJuMP.ParameterRef}) where C
    param_aff = ParameterJuMP.PAE{C}(JuMP.GenericAffExpr{C, JuMP.VariableRef}(aff.constant),
                                     JuMP.GenericAffExpr{C, ParameterJuMP.ParameterRef}(zero(C)))
    for (coef, var) in JuMP.linear_terms(aff)
        JuMP.add_to_expression!(param_aff,
                                coef,
                                parametrized_function(var, structured_model, model, variable_map, parameter_map))
    end
    # The inferred return value is an union of two concrete types but
    # small unions are efficiently handled in Julia v0.7
    if JuMP.iszero(param_aff.p)
        return param_aff.v
    else
        return param_aff
    end
end
function parametrized_function(quad::JuMP.GenericQuadExpr{C, StructuredVariableRef},
                               structured_model::StructuredModel,
                               model::JuMP.Model,
                               variable_map::Dict{Int, JuMP.VariableRef},
                               parameter_map::Dict{Int, ParameterJuMP.ParameterRef}) where C
    param_aff = parametrized_function(quad.aff, structured_model, model, variable_map, parameter_map)
    if param_aff isa ParameterJuMP.PAE
        error("parametrized quadratic functions are not supported yet")
        #quadv = JuMP.GenericQuadExpr{C, JuMP.VariableRef}(param_aff.v)
        #quadp = JuMP.GenericQuadExpr{C, ParameterJuMP.ParameterRef}(param_aff.p)
    else
        quadv = JuMP.GenericQuadExpr{C, JuMP.VariableRef}(param_aff)
        quadp = JuMP.GenericQuadExpr{C, ParameterJuMP.ParameterRef}(zero(C))
    end
    for (coef, var1, var2) in JuMP.quadterms(quad)
        if var1.model === structured_model && var2.model == structured_model
            JuMP.add_to_expression!(quadv, coef, variable_map[var1.idx], variable_map[var2.idx])
        else
            error("parametrized quadratic functions are not supported yet")
        end
    end
    if JuMP.iszero(quadp)
        return quadv
    else
        error("parametrized quadratic functions are not supported yet")
    end
end
parametrized_function(funcs::Vector, args...) = map(f -> parametrized_function(f, args...), funcs)

struct ParametrizedModel
    structured_model::StructuredModel
    model::JuMP.Model
    # Map between index of structured variable in `structured_model` and
    # the corresponding variable in `model`.
    variable_map::Dict{Int, JuMP.VariableRef}
    # Map between index of structured variable in `structured_model.parent`
    # and the corresponding parameter in `model`.
    parameter_map::Dict{Int, ParameterJuMP.ParameterRef}
    # Cost of children
    θ::Dict{Int, JuMP.VariableRef}
end
function ParametrizedModel(structured_model::StructuredModel, args...; kwargs...)
    if structured_model.parent === nothing
        model = Model(args...; kwargs...)
    else
        model = ModelWithParams(args...; kwargs...)
    end
    variable_map = Dict{Int, JuMP.VariableRef}()
    for (index, var) in structured_model.variables
        name = structured_model.varnames[index]
        variable_map[index] = JuMP.add_variable(model, var, name)
    end
    parameter_map = Dict{Int, ParameterJuMP.ParameterRef}()
    for (index, con) in structured_model.constraints
        name = structured_model.connames[index]
        param_fun = parametrized_function(con.func, structured_model, model,
                                          variable_map, parameter_map)
        param_con = JuMP.build_constraint(error, param_fun, con.set)
        JuMP.add_constraint(model, param_con, name)
    end
    objective_function = parametrized_function(structured_model.objective_function,
                                               structured_model, model,
                                               variable_map, parameter_map)
    if structured_model.objective_sense == MOI.MAX_SENSE
        objective_function = -objective_function
    end
    θ = Dict{Int, VariableRef}()
    for (id, proba) in getprobability(structured_model)
        θid = @variable(model, lower_bound = 0.0)
        θ[id] = θid
        JuMP.add_to_expression!(objective_function, proba, θid)
    end
    JuMP.set_objective(model, structured_model.objective_sense, objective_function)
    ParametrizedModel(structured_model, model, variable_map, parameter_map, θ)
end

include("Benders_pmap.jl")

# The optimizers are function returning an empty optimizer,
# this will be replaced by factories once this is implemented in JuMP
function BendersBridge(structured_model::StructuredModel,
                       master_optimizer::JuMP.OptimizerFactory,
                       sub_optimizer::JuMP.OptimizerFactory)
    master_model = ParametrizedModel(structured_model, master_optimizer)
    children = getchildren(structured_model)
    sub_models = Dict{Int, ParametrizedModel}()
    for (id, child) in getchildren(structured_model)
        sub_models[id] = ParametrizedModel(child, sub_optimizer)
    end
    return Benders_pmap(master_model, sub_models)
end


export read_results!

"""

    read_results!(optimizer::Optimizer{T}, filepath::String, redundant_entries::Vector)
Populates `optimizer` with results in a SDPA-formatted output file specified by `filepath`. Redundant entries corresponding to linearly dependent constraints are set to 0.

"""
function read_results!(
    optimizer::Optimizer{T},
    filepath::String,
    redundant_entries::Vector
) where T

    endswith(filepath, ".dat") || error("Filename '$filepath' must end with .dat")
    getnextline(io::IO) =
        eof(io) ?
        error("The output file is possibly corrupted. Check that $filepath conforms to the SDPA output format.") :
        chomp(readline(io))

    function replace_brackets!(str::SubString)
        str = replace(str, "{" => "[")
        str = replace(str, "}" => "]")
        return str
    end
    function remove_brackets!(str)
        str = replace(str, "{" => "")
        str = replace(str, "}" => "")
        return str
    end

    phasevalue = :noINFO
    objValPrimalstring = ""
    objValDualstring = ""
    xVecstring = ""
    xMatvec = T[]
    yMatvec = T[]

    file = open(filepath, "r") do io
            line = getnextline(io)

            while !startswith(line, "phase.value")
                line = getnextline(io)
            end
            optimizer.phasevalue = Symbol(split(line)[3])

            while !startswith(line, "objValPrimal")
                line = getnextline(io)
            end
            objValPrimalstring = split(line)[3]

            while !startswith(line, "objValDual")
                line = getnextline(io)
            end
            objValDualstring = split(line)[3]

            while !startswith(line, "xVec")
                line = getnextline(io)
            end
            line = getnextline(io)
            xVecstring = line

            while !startswith(line, "xMat")
                line = getnextline(io)
            end
            line = getnextline(io)
            line = getnextline(io)
            while !startswith(line, "}")
                xMatstring = remove_brackets!(line)
                if endswith(xMatstring, ",")
                    append!(xMatvec, parse.(T, split(xMatstring[1:end-1], ",")))
                else
                    append!(xMatvec, parse.(T, split(xMatstring, ",")))
                end
                line = getnextline(io)
            end
            line = getnextline(io)
            line = getnextline(io)
            line = getnextline(io)
            while !startswith(line, "}")
                yMatstring = remove_brackets!(line)
                if endswith(yMatstring, ",")
                    append!(yMatvec, parse.(T, split(yMatstring[1:end-1], ",")))
                else
                    append!(yMatvec, parse.(T, split(yMatstring, ",")))
                end
                line = getnextline(io)
            end
        end
    xVecstring = remove_brackets!(xVecstring)
    xVecstring = split(xVecstring, ",")
    xVec = parse.(T, xVecstring)
    optimizer.primalobj = parse(T, objValPrimalstring)
    optimizer.dualobj = parse(T, objValDualstring)
    for i in redundant_entries
        splice!(xVec, i:i-1, zero(T))
    end
    optimizer.y = xVec
    structurevec = optimizer.blockdims
    yMatbm = PrimalSolution{T}(map(
        n -> zeros(T, abs(n), abs(n)),
        optimizer.blockdims
    ))
    xMatbm = VarDualSolution{T}(map(
        n -> zeros(T, abs(n), abs(n)),
        optimizer.blockdims
    ))
    for i = 1:length(structurevec)
        dim = structurevec[i]
        if dim < 0
            dim = abs(dim)
            xblock = xMatvec[1:dim]
            deleteat!(xMatvec, 1:dim)
            xblock = Diagonal(xblock)
            xMatbm.blocks[i] = xblock

            yblock = yMatvec[1:dim]
            deleteat!(yMatvec, 1:dim)
            yblock = Diagonal(yblock)
            yMatbm.blocks[i] = yblock

        elseif dim > 0
            xblock = xMatvec[1:dim^2]
            deleteat!(xMatvec, 1:dim^2)
            xblock = reshape(xblock, dim, dim)
            xMatbm.blocks[i] = xblock

            yblock = yMatvec[1:dim^2]
            deleteat!(yMatvec, 1:dim^2)
            yblock = reshape(yblock, dim, dim)
            yMatbm.blocks[i] = yblock

        end
    end
    optimizer.X = yMatbm
    optimizer.Z = xMatbm

end

"""

    inputElement(optimizer::Optimizer, constr_number::Int, blk::Int, i::Int, j::Int, value::T) where T

Stores the constraint data in `optimizer.elemdata` as a vector of tuples. Each tuple corresponds to one line in the SDPA-formatted input file.
"""
function inputElement(
    optimizer::Optimizer,
    constr_number::Int,
    blk::Int,
    i::Int,
    j::Int,
    value::T
) where T
    push!(optimizer.elemdata, (constr_number, blk, i, j, value))
end

"""

    initializeSolve(optimizer::Optimizer)

Writes problem data into an SDPA-formatted file named `input.dat-s`. `presolve.jl` routine is applied as indicated by `optimizer.presolve`.

Returns a vector of indices for redundant constraints, which are omitted from the input file.
"""
function initializeSolve(optimizer::Optimizer)
    if !optimizer.presolve
        filename = joinpath(optimizer.tempfile, "input.dat-s")
        file = open(filename, "w") do io
            nconstrs = length(optimizer.b)
            nblocks = length(optimizer.blockdims)
            println(io, nconstrs)
            println(io, nblocks)
            writedlm(io, optimizer.blockdims', " ")
            writedlm(io, optimizer.b', " ")
            for line in optimizer.elemdata
                for i in line
                    print(io, i)
                    print(io, " ")
                end
                println(io)
            end
        end
        return []
    else
        redundant_F = presolve(optimizer)
        reduced = joinpath(optimizer.tempfile, "input.dat-s")
        file =
            open(reduced, "w") do io
                nconstrs = length(optimizer.b) - length(redundant_F)
                nblocks = length(optimizer.blockdims)
                println(io, nconstrs)
                println(io, nblocks)
                writedlm(io, optimizer.blockdims', " ")
                cVec = deleteat!(copy(optimizer.b), redundant_F)
                writedlm(io, cVec', " ")
                for entry in optimizer.elemdata
                    constr_index, blk, i, j, value = entry
                    if !in(constr_index, redundant_F)
                        constr_index = constr_index -
                                       count(x -> x < constr_index, redundant_F)
                        for i in (constr_index, blk, i, j, value)
                            print(io, i)
                            print(io, " ")
                        end
                        println(io)
                    end
                end
            end
        return redundant_F
    end
end


using Dates

tomlpath = joinpath(@__DIR__, "sbm_config.toml")
config = Wflow.Config(tomlpath)

model = Wflow.initialize_sbm_model(config)
@unpack network = model

Wflow.load_dynamic_input!(model)
model = Wflow.update(model)

# test if the first timestep was written to the CSV file
flush(model.writer.csv_io)  # ensure the buffer is written fully to disk
@testset "CSV output" begin
    row = csv_first_row(model.writer.csv_path)

    @test row.time == DateTime("2000-01-02T00:00:00")
    @test row.Q ≈ 8.1396354649197f0
    @test row.volume ≈ 2.7535003939625636f7
    @test row.temp_bycoord ≈ 2.390000104904175f0
    @test row.temp_byindex ≈ 2.390000104904175f0
    @test row.Q_6336050 ≈ 0.006160282939850074f0
    @test row.Q_6336510 ≈ 0.029177309556712334f0
    @test row.Q_6836100 ≈ 0.19621282450614713f0
    @test row.Q_6336500 ≈ 0.006089112638001381f0
    @test row.Q_6836190 ≈ 0.0031262850749354237f0
    @test row.Q_6336800 ≈ 0.007770868657277307f0
    @test row.Q_6336900 ≈ 0.006403194169947582f0
    @test row.Q_6336930 ≈ 0.08888787154163148f0
    @test row.Q_6336910 ≈ 0.007071851236520184f0
    @test row.Q_6136500 ≈ 0.0016367337487926633f0
    @test row.Q_6136520 ≈ 0.002084670434294102f0
    @test row.Q_6136150 ≈ 0.006095549758915344f0
    @test row.Q_6136151 ≈ 0.007643634432992056f0
    @test row.Q_6136160 ≈ 3.9199531493174726f0
    @test row.Q_6136202 ≈ 1.4125847550988493f0
    @test row.recharge_1 ≈ -0.05653226176238641f0
end

@testset "NetCDF scalar output" begin
    ds = model.writer.dataset_scalar
    @test ds["time"][1] == DateTime("2000-01-02T00:00:00")
    @test ds["Q"][:][1:20] ≈ [
        0.73993874f0,
        1.4125848f0,
        1.4389194f0,
        1.4036233f0,
        5.7276225f0,
        2.7586424f0,
        2.1080604f0,
        4.1026044f0,
        0.008365918f0,
        3.919953f0,
        4.0615087f0,
        0.006050462f0,
        0.0076436345f0,
        0.00780229f0,
        0.00346879f0,
        0.7066179f0,
        0.0022446192f0,
        1.3323202f0,
        3.8271446f0,
        1.6729931f0,
    ]
    @test ds["Q_gauges"].attrib["cf_role"] == "timeseries_id"
    @test ds["temp_index"][:] ≈ [2.39f0]
    @test ds["temp_coord"][:] ≈ [2.39f0]
    @test keys(ds.dim) == ["time", "Q_gauges", "temp_bycoord", "temp_byindex"]
end

@testset "first timestep" begin
    sbm = model.vertical

    @test sbm.tt[50063] ≈ 0.0f0

    @test model.clock.iteration == 2

    @test sbm.θₛ[50063] ≈ 0.48755401372909546f0
    @test sbm.θᵣ[50063] ≈ 0.15943120419979095f0
    @test sbm.runoff[50063] == 0.0
    @test sbm.soilevap[50063] == 0.0
    @test sbm.snow[5] ≈ 3.592840840467347f0
end

# run the second timestep
Wflow.load_dynamic_input!(model)
model = Wflow.update(model)

@testset "second timestep" begin
    sbm = model.vertical
    @test sbm.θₛ[50063] ≈ 0.48755401372909546f0
    @test sbm.θᵣ[50063] ≈ 0.15943120419979095f0
    @test sbm.runoff[50063] == 0.0
    @test sbm.soilevap[50063] ≈ 0.006358004660566856f0
    @test sbm.snow[5] ≈ 3.667748983774868f0
end

@testset "subsurface flow" begin
    ssf = model.lateral.subsurface.ssf
    @test sum(ssf) ≈ 6.370399148012509f7
    @test ssf[network.land.order[1]] ≈ 7.169036749244327f2
    @test ssf[network.land.order[end-100]] ≈ 2335.2465707069578f0
    @test ssf[network.land.order[end]] ≈ 288.19428729403944f0
end

@testset "overland flow" begin
    q = model.lateral.land.q_av
    @test sum(q) ≈ 290.5520014030802f0
    @test q[26625] ≈ 0.0
    @test q[39308] ≈ 0.0
    @test q[network.land.order[end]] ≈ 1.0f-30
end

@testset "river flow" begin
    q = model.lateral.river.q_av
    @test sum(q) ≈ 3615.2084127447865f0
    @test q[1622] ≈ 0.0005986272622329333f0
    @test q[43] ≈ 12.036342425160155f0
    @test q[network.river.order[end]] ≈ 0.03835913312643948f0
end

@testset "reservoir simple" begin
    res = model.lateral.river.reservoir
    @test res.outflow[1] ≈ 0.21750000119148086f0
    @test res.inflow[1] ≈ 43.18479982574888f0
    @test res.volume[1] ≈ 2.751299001489657f7
    @test res.precipitation[1] ≈ 0.17999997735023499f0
    @test res.evaporation[1] ≈ 0.5400000810623169f0
end

# set these variables for comparison in "changed dynamic parameters"
precip = copy(model.vertical.precipitation)
evap = copy(model.vertical.potential_evaporation)
lai = copy(model.vertical.leaf_area_index)
res_evap = copy(model.lateral.river.reservoir.evaporation)

Wflow.close_files(model, delete_output = false)

# test for setting a pit and multithreading multiple basins (by setting 2 extra pits
# resulting in 3 basins)
tomlpath = joinpath(@__DIR__, "sbm_config.toml")
config = Wflow.Config(tomlpath)
config["model"]["pits"] = true
config["input"]["pits"] = "wflow_pits"
config.endtime = DateTime(2000, 1, 9)
config.loglevel = "info"

model = Wflow.run(config)

@testset "timing" begin
    # clock has been reset
    calendar = get(config, "calendar", "standard")::String
    @test model.clock.time == Wflow.cftime(config.starttime, calendar)
    @test model.clock.iteration == 1
end

@testset "river flow at basin outlets and downstream of one pit" begin
    q = model.lateral.river.q_av
    @test q[4009] ≈ 8.60480399680283f0 # pit/ outlet, CartesianIndex(141, 228)
    @test q[4020] ≈ 0.006779014715290862f0 # downstream of pit 4009, CartesianIndex(141, 229)
    @test q[2508] ≈ 150.15321264134985f0 # pit/ outlet
    @test q[5808] ≈ 0.12625654862968252f0 # pit/ outlet
end

# test changing forcing and cyclic LAI parameter
tomlpath = joinpath(@__DIR__, "sbm_config.toml")
config = Wflow.Config(tomlpath)

config.input.vertical.precipitation =
    Dict("scale" => 2.0, "netcdf" => Dict("variable" => Dict("name" => "precip")))
config.input.vertical.potential_evaporation = Dict(
    "scale" => 3.0,
    "offset" => 1.50,
    "netcdf" => Dict("variable" => Dict("name" => "pet")),
)
config.input.vertical.leaf_area_index =
    Dict("scale" => 1.6, "netcdf" => Dict("variable" => Dict("name" => "LAI")))

model = Wflow.initialize_sbm_model(config)
Wflow.load_dynamic_input!(model)
model = Wflow.update(model)
Wflow.load_dynamic_input!(model)
model = Wflow.update(model)

@testset "changed dynamic parameters" begin
    res = model.lateral.river.reservoir
    vertical = model.vertical
    @test vertical.precipitation[2] / precip[2] ≈ 2.0f0
    @test (vertical.potential_evaporation[100] - 1.50) / evap[100] ≈ 3.0f0
    @test vertical.leaf_area_index[100] / lai[100] ≈ 1.6f0
    @test (res.evaporation[2] - 1.50) / res_evap[2] ≈ 3.0000012203408635f0
end

# test cyclic river inflow
tomlpath = joinpath(@__DIR__, "sbm_config.toml")
config = Wflow.Config(tomlpath)

config.input.cyclic = ["vertical.leaf_area_index", "lateral.river.inflow"]
config.input.lateral.river.inflow = "inflow"

model = Wflow.initialize_sbm_model(config)
Wflow.load_dynamic_input!(model)
model = Wflow.update(model)
Wflow.load_dynamic_input!(model)
model = Wflow.update(model)

@testset "river inflow (cyclic)" begin
    @test model.lateral.river.inflow[44] ≈ 0.75
    @test model.lateral.river.q_av[44] ≈ 10.698591283662008
end

# test fixed forcing (precipitation = 2.5)
config = Wflow.Config(tomlpath)
config.input.vertical.precipitation = Dict("value" => 2.5)
model = Wflow.initialize_sbm_model(config)
Wflow.load_fixed_forcing(model)

@testset "fixed precipitation forcing (initialize)" begin
    @test maximum(model.vertical.precipitation) ≈ 2.5
    @test minimum(model.vertical.precipitation) ≈ 0.0
    @test all(isapprox.(model.lateral.river.reservoir.precipitation, 2.5))
end

Wflow.load_dynamic_input!(model)
model = Wflow.update(model)

@testset "fixed precipitation forcing (first timestep)" begin
    @test maximum(model.vertical.precipitation) ≈ 2.5
    @test minimum(model.vertical.precipitation) ≈ 0.0
    @test all(isapprox.(model.lateral.river.reservoir.precipitation, 2.5))
end

Wflow.close_files(model, delete_output = false)

# test local-inertial option for river flow river_routing
tomlpath = joinpath(@__DIR__, "sbm_config.toml")
config = Wflow.Config(tomlpath)
config.model.river_routing = "local-inertial"

model = Wflow.initialize_sbm_model(config)
Wflow.load_dynamic_input!(model)
model = Wflow.update(model)
Wflow.load_dynamic_input!(model)
model = Wflow.update(model)

@testset "river flow and depth (local inertial)" begin
    q = model.lateral.river.q_av
    @test sum(q) ≈ 3906.99107445248f0
    @test q[1622] ≈ 6.0094181857060604f-5
    @test q[43] ≈ 11.900372477232786f0
    @test q[501] ≈ 3.536628093804679f0
    h = model.lateral.river.h_av
    @test h[1622] ≈ 0.0018099650639472893f0
    @test h[43] ≈ 0.4362704420866992f0
    @test h[501] ≈ 0.05669956233680629f0
end
Wflow.close_files(model, delete_output = false)

# test local-inertial option for river flow river_routing
tomlpath = joinpath(@__DIR__, "sbm_swf_config.toml")
config = Wflow.Config(tomlpath)

model = Wflow.initialize_sbm_model(config)
Wflow.load_dynamic_input!(model)
model = Wflow.update(model)
Wflow.load_dynamic_input!(model)
model = Wflow.update(model)

@testset "river and overland flow and depth (local inertial)" begin
    q = model.lateral.river.q_av
    @test sum(q) ≈ 2369.4745120995767f0
    @test q[1622] ≈ 6.027872269041313f-5
    @test q[43] ≈ 5.3952531643382144f0
    @test q[501] ≈ 1.6023076297110879f0
    h = model.lateral.river.h_av
    @test h[1622] ≈ 0.001812085885794477f0
    @test h[43] ≈ 0.30122213937420667f0
    @test h[501] ≈ 0.03172037092237322f0
    qx = model.lateral.land.qx
    qy = model.lateral.land.qy
    @test qx[[26, 35, 631]] ≈
          [0.13689085916734353f0, 0.02571727530678835f0, 0.06163019045555399f0]
    @test qy[[26, 35, 631]] ≈
          [0.189187686873958f0, 0.017683557812671172f0, 0.09182928055862745f0]
    h = model.lateral.land.h
    @test h[[26, 35, 631]] ≈
          [0.07359964438114004f0, 0.0070632177001078555f0, 0.0009883334841610835f0]
end
Wflow.close_files(model, delete_output = false)


using Test
using GeometricIntegrators
using GeometricProblems.LotkaVolterra2d
using GeometricProblems.LotkaVolterra2d: Δt
using SrkMethodsForDegenerateLagrangianSystems

const tableaus = (
    tableaus_vprk_glrk(),
    tableaus_vprk_lobatto_IIIA_IIIB(),
    tableaus_vprk_lobatto_IIIB_IIIA(),
    tableaus_vprk_radau(),
    tableaus_srk_glrk(),
    tableaus_firk_glrk(),
)

const iode = lotka_volterra_2d_iode()
const nt = 10

for list in tableaus
    for run in list
        tab, file = run

        if length(run) ≥ 3
            integrator = run[3]
        else
            integrator = Integrator
        end

        @test_nowarn begin
            sol = Solution(iode, Δt, nt)
            int = integrator(iode, tab, Δt)
            integrate!(int, sol)
        end
    end
end


module BVB
using Distributions
using Optim: optimize, converged, LBFGS, Options, GoldenSection
using PDMats: PDMat, inv
using Calculus: second_derivative, hessian
using UnicodePlots: lineplot, lineplot!, scatterplot, scatterplot!, barplot
using StatsFuns: logsumexp

abstract type abstractParticles end

type Particles{T <: AbstractFloat} <: abstractParticles
    n               ::Int
    p               ::Int
    isMultivariate  ::Bool
    isNormalized    ::Bool
    x               ::Array{T}
    w               ::Vector{Float64}
end

const ϵ_lower_bound = 1e-4

include("type.jl")
include("boosting.jl")
include("boostInclusiveKL.jl")
include("boostExclusiveKL.jl")
include("numericalStable.jl")
include("utils.jl")
include("models.jl")
include("gradientBoosting.jl")
include("wrapper.jl")
end


#slicetiming_experiment runs a pilot piezo-only recording, reads piezo monitor samples, and returns an imagine procedure with the following attributes:
#The SignalGenerator generates command signals with two phases:
#   Phase 1:  A slow (0.2Hz) stack is acquired with camera and laser pulses timed to match target z locations determined by z_spacing.
#   Phase 2:  Many piezo cycles with two camera and laser pulses per cycle.  The piezo cycle is taken from the recs input.
#           Each pulse is timed at a different lag relative to the target z location determined by z_spacing (and used to place slices in Phase 1)
#The process function expects the recorded signals from the above commands as inputs, and it outputs time indices indicating when to time slices so that z location
#when imaging fast matches the location of images from the template stack of phase 1 above.
#Optionally pass a Calibration to be applied to the piezo MON signal before doing all this (if, for example, you know the constant lag in the MON output)

#NOTE: the default values may need to be tweaked case-by-case
#TODO: make this more robust
default_toffsets() = [0.0000s:0.0001s:0.002s...]

function slicetiming_experiment(out_bname::AbstractString, rig::AbstractString, pos_name, las_name, cam_name, pstart, pstop, stack_rate, z_spacing, z_pad, ncycs_mean=ceil(Int, 20.0/ustrip(inv(stack_rate))); lp_cutoff=3.5*stackrate, sample_rate = 100000Hz, cal=-1, toffsets = default_toffsets(), allow_shifts=true, allow_rotations=false, subpixel=true)
    mod_cyc = vcat(gen_bidi_pos(pstart, pstop, 1/stack_rate, sample_rate; lp_cutoff=lp_cutoff)...)
    return slicetiming_experiment(out_bname, rig, pos_name, las_name, cam_name, mod_cyc, z_spacing, z_pad; ncycs_mean=ncycs_mean, sample_rate=sample_rate, cal=cal, toffsets=toffsets, allow_shifts=allow_shifts, allow_rotations=allow_rotations, subpixel=subpixel)
end

function slicetiming_experiment(coms, recs, las_name::AbstractString, cam_name::AbstractString, z_spacing::HasLengthUnits, z_pad::HasLengthUnits; cal = -1, toffsets = default_toffsets(), allow_shifts=true, allow_rotations=false, subpixel=true)
    pos, pos_mon = pos_mod_mon(coms, recs)
    if cal == -1
        pos_mon_cal = pos_mon
    else
        pos_mon_cal = apply(cal, pos_mon)
    end
    mon_samps = get_samples(pos_mon_cal)
    nsamps_cycle = full_length(first(sequences(pos)))
    ncycs_ignore = ceil(Int, (ustrip(samprate(pos))*5) / nsamps_cycle) #5 seconds of warmup
    mean_cyc = mean_cycle(mon_samps, nsamps_cycle; start_idx = ncycs_ignore*nsamps_cycle+1)
    #set pmin and pmax empirically
    pmin = minimum(mean_cyc)
    pmax = maximum(mean_cyc)
    slice_zs = ImagineInterface.slice_positions(pmin, pmax, z_spacing, z_pad)
    slicetiming_experiment(pos, mean_cyc, las_name, cam_name, slice_zs; allow_shifts=allow_shifts, allow_rotations=allow_rotations, toffsets=toffsets, subpixel=subpixel)
end

function slicetiming_experiment(pos::ImagineSignal, mon_cyc, las_name::AbstractString, cam_name::AbstractString, slice_zs; toffsets = default_toffsets(), allow_shifts=true, allow_rotations=false, subpixel=true)
    gen_desc = "Generate a set camera and laser commands to find slice timings empirically given a completed dynamic positioner recording"
    sig_gen = SignalGenerator(gen_desc, slicetiming_commands, (pos, mon_cyc, las_name, cam_name, toffsets, slice_zs))
    desc = "Returns slice timings to align forward and reverse stacks at the desired spacing for high-speed volume acquisitions"
    analyze_f = (coms, img) -> (slicetimings(img, toffsets, length(slice_zs); allow_shifts=allow_shifts, allow_rotations=allow_rotations, subpixel=subpixel)..., slice_zs) #returns forward toffsets, backward toffsets, and the z locations of slices
    return ImagineProcedure(desc, sig_gen, analyze_f)
end

##Similar to slicetiming_experiment but only checks one slice per stack, and uses the result to estimate the lag between the piezo's true position and the MON signal
##function pos_mon_lag_experiment(pos::ImagineSignal, mon_cyc::ImagineSignal, las_name::AbstractString, cam_name::AbstractString, slice_zs)
##WARNING This function probably shouldn't be used. Use slicetiming_experiment instead
#function pos_mon_lag_experiment(pos::ImagineSignal, mon_cyc, las_name::AbstractString, cam_name::AbstractString)
#    #set pmin and pmax empirically
#    pmin = minimum(mon_cyc)
#    pmax = maximum(mon_cyc)
#    pctr = [(pmax-pmin)/2 * pmin]
#    @show toffsets = default_toffsets() #These worked to find the lag on the digital piezo of OCPI2
#    gen_desc = "Generate a set camera and laser commands to find piezo monitor lag empirically given a completed dynamic positioner recording"
#    sig_gen = SignalGenerator(gen_desc, slicetiming_commands, (pos, mon_cyc, las_name, cam_name, toffsets, [pctr]))
#    desc = "Returns mean of two measurements: the monitor lag when taking an image while sweeping forward and also while sweeping back"
#    analyze_f = (coms, img) -> mean(slicetimings(img, toffsets, 1)) #take the mean of the forward and reverse lag (hopefully they are about equal)
#    return ImagineProcedure(desc, sig_gen, analyze_f)
#end
#
#pos_mon_lag_experiment(coms, recs, las_name::AbstractString, cam_name::AbstractString) = pos_mon_lag_experiment(pos_mod_mon(coms,recs)..., las_name, cam_name)
#
#function pos_mon_lag_experiment(pos::ImagineSignal, pos_mon::ImagineSignal, las_name::AbstractString, cam_name::AbstractString)
#    mon_samps = get_samples(pos_mon)
#    nsamps_cycle = full_length(first(sequences(pos)))
#    ncycs_ignore = ceil(Int, (ustrip(samprate(pos))*5) / nsamps_cycle) #5 seconds of warmup
#    mon_cyc = mean_cycle(mon_samps, nsamps_cycle; start_idx = ncycs_ignore*nsamps_cycle+1)
#    pos_mon_lag_experiment(pos, mon_cyc, las_name, cam_name)
#end


module ParticleDA

using Random
using Distributions, Statistics, MPI, Base.Threads, YAML, HDF5
using TimerOutputs
import Future
using EllipsisNotation

export run_particle_filter, BootstrapFilter, OptimalFilter

include("params.jl")
include("io.jl")
include("OptimalFilter.jl")

using .Default_params

# Functions to extend in the model

"""
    ParticleDA.get_grid_size(model_data) -> NTuple{N, Int} where N

Return a tuple with the dimensions (number of nodes) of the grid.
"""
function get_grid_size end
"""
    ParticleDA.get_grid_domain_size(model_data) -> NTuple{N, float} where N

Return a tuple with the dimensions (metres) of the grid domain.
"""
function get_grid_domain_size end
"""
    ParticleDA.get_grid_cell_size(model_data) -> NTuple{N, float} where N

Return a tuple with the dimensions (metres) of a grid cell.
"""
function get_grid_cell_size end

"""
    ParticleDA.get_n_state_var(model_data) -> Int

Return the number of state variables.
"""
function get_n_state_var end

"""
    ParticleDa.get_obs_noise_std(model_data) -> Float

Return standard deviation of observation noise. Required for optimal filter only.
"""
function get_obs_noise_std end

"""
    ParticleDa.get_model_noise(model_data) -> NamedTuple({:sigma, :lambda, :nu},Tuple{Float, Float, Float})

Return parameters of the noise covariance function of the first state variable (height). Required for optimal filter only.
"""
function get_model_noise_params end

"""
    ParticleDA.get_particle(model_data) -> particles

Return the vector of particles.  This method is intended to be extended by the
user with the above signature, specifying the type of `model_data`.  Note: this
function should return the vector of particles itself and not a copy, because it
will be modified in-place.
"""
function get_particles end

"""
   ParticleDA.set_particles!(model_data, particles)

Overwrite particle state in model_data with the vector particles. This method is intended to be extended by the
user with the above signature, specifying the type of `model_data`.
"""
function set_particles! end

"""
    ParticleDA.get_truth(model_data) -> truth_observations

Return the vector of true observations.  This method is intended to be extended
by the user with the above signature, specifying the type of `model_data`.
"""
function get_truth end

"""
    ParticleDA.get_stations(model_data) -> NamedTuple({:nst,:ist,:jst},Tuple{Int, Array{Float,1}, Array{Float,1}})

Return a named tuple with number of stations and their coordinates (nst,ist,jst) of the
points of observation. This method is intended to be extended by the user with the above signature,
specifying the type of `model_data`.
Required for optimal filter only.
"""
function get_stations end

"""
    ParticleDA.update_truth!(model_data, nprt_per_rank::Int) -> truth_observations

Update the true observations using the dynamic of the model and return the
vector of the true observations.  `nprt_per_rank` is the number of particles per
each MPI rank.  This method is intended to be extended by the user with the
above signature, specifying the type of `model_data`.
"""
function update_truth! end

"""
    ParticleDA.update_particle_dynamics!(model_data, nprt_per_rank::Int)

Update the particles using the dynamic of the model.  `nprt_per_rank` is the
number of particles per each MPI rank.  This method is intended to be extended
by the user with the above signature, specifying the type of `model_data`.
"""
function update_particle_dynamics! end

"""
    ParticleDA.update_particle_noise!(model_data, nprt_per_rank::Int)

Update the particles using the noise of the model and return the vector of the
particles.  `nprt_per_rank` is the number of particles per each MPI rank.  This
method is intended to be extended by the user with the above signature,
specifying the type of `model_data`.
"""
function update_particle_noise! end

"""
    ParticleDA.get_particle_observations!(model_data, nprt_per_rank::Int) -> particles_observations

Return the vector of the particles observations.  `nprt_per_rank` is the number
of particles per each MPI rank.  This method is intended to be extended by the
user with the above signature, specifying the type of `model_data`.
"""
function get_particle_observations! end

"""
    ParticleDA.write_snapshot(output_filename, model_data, avg_arr, var_arr, weights, it)

Write a snapshot of the data after an update of the particles to the HDF5 file
`output_filename`.  `avg_arr` is the array of the mean of the particles,
`var_arr` is the array of the standard deviation of the particles, `weights` is
the array of the weigths of the particles, `it` is the index of the time step
(`it==0` is the initial state, before moving forward the model for the first
time).  This method is intended to be extended by the user with the above
signature, specifying the type of `model_data`.
"""
function write_snapshot end

"""
    ParticleFilter

Abstract type for the particle filter to use.  Currently used subtypes are:
* [`BootstrapFilter`](@ref)
* [`OptimalFilter`](@ref)
"""
abstract type ParticleFilter end
"""
    BootstrapFilter()

Instantiate the singleton type `BootstrapFilter`.  This can be used as argument
of [`run_particle_filter`](@ref) to select the bootstrap filter.
"""
struct BootstrapFilter <: ParticleFilter end
struct OptimalFilter <: ParticleFilter end

# Get weights for particles by evaluating the probability of the observations predicted by the model
# from independent normal pdfs for each observation.
function get_log_weights!(weight::AbstractVector{T},
                          obs::AbstractVector{T},
                          obs_model::AbstractMatrix{T},
                          weight_std::T) where T

    nobs = size(obs_model,1)
    @assert size(obs,1) == nobs

    weight .= 1.0

    for iobs = 1:nobs
        weight .+= logpdf.(Normal(obs[iobs], weight_std), @view(obs_model[iobs,:]))
    end

end

# Get weights for particles by evaluating the probability of the observations predicted by the model
# from a multivariate normal pdf with mean equal to real observations and covariance equal to cov_obs
function get_log_weights!(weight::AbstractVector{T},
                          obs::AbstractVector{T},
                          obs_model::AbstractMatrix{T},
                          cov_obs::AbstractMatrix{T}) where T

    weight .= Distributions.logpdf(Distributions.MvNormal(obs, cov_obs), obs_model)

end

#
function normalized_exp!(weight::AbstractVector)

    weight .-= maximum(weight)
    @. weight = exp(weight)
    weight ./= sum(weight)

end

# Resample particles from given weights using Stochastic Universal Sampling
function resample!(resampled_indices::AbstractVector{Int}, weight::AbstractVector{T}, rng::Random.AbstractRNG=Random.default_rng()) where T

    nprt = length(weight)
    nprt_inv = 1.0 / nprt
    k = 1

    #TODO: Do we need to sort state by weight here?

    weight_cdf = cumsum(weight)
    u0 = nprt_inv * rand(rng, T)

    # Note: To parallelise this loop, updates to k and u have to be atomic.
    # TODO: search for better parallel implementations
    for ip in 1:nprt

        u = u0 + (ip - 1) * nprt_inv

        while(u > weight_cdf[k])
            k += 1
        end

        resampled_indices[ip] = k

    end

end

struct SummaryStat{T}
    avg::T
    var::T
    n::Int
end

function SummaryStat(X::AbstractVector)
    m = mean(X)
    v = varm(X,m, corrected=true)
    n = length(X)
    SummaryStat(m,v,n)
end

function stats_reduction(S1::SummaryStat, S2::SummaryStat)

    n = S1.n + S2.n
    m = (S1.avg*S1.n + S2.avg*S2.n) / n

    # Calculate pooled unbiased sample variance of two groups. From https://stats.stackexchange.com/q/384951
    # Can be found in https://www.tandfonline.com/doi/abs/10.1080/00031305.2014.966589
    # To get the uncorrected variance, use
    # v = (S1.n * (S1.var + S1.avg * (S1.avg-m)) + S2.n * (S2.var + S2.avg * (S2.avg-m)))/n
    v = ((S1.n-1) * S1.var + (S2.n-1) * S2.var + S1.n*S2.n/n * (S2.avg - S1.avg)^2 )/(n-1)

    SummaryStat(m, v, n)

end

function get_mean_and_var!(statistics::Array{SummaryStat{T}},
                           particles::AbstractArray{T},
                           master_rank::Int) where T

    ndims(particles) != ndims(statistics) + 1 &&
        error("particles must have one dimension more than statistics")
    Threads.@threads for idx in CartesianIndices(statistics)
        statistics[idx] = SummaryStat(@view(particles[idx,:]))
    end

    MPI.Reduce!(statistics, stats_reduction, master_rank, MPI.COMM_WORLD)

end

function unpack_statistics!(avg::AbstractArray{T}, var::AbstractArray{T}, statistics::AbstractArray{SummaryStat{T}}) where T

    for idx in CartesianIndices(statistics)
        avg[idx] = statistics[idx].avg
        var[idx] = statistics[idx].var
    end
end

function copy_states!(particles::AbstractArray{T},
                      buffer::AbstractArray{T},
                      resampling_indices::Vector{Int},
                      my_rank::Int,
                      nprt_per_rank::Int) where T

    # These are the particle indices stored on this rank
    particles_have = my_rank * nprt_per_rank + 1:(my_rank + 1) * nprt_per_rank

    # These are the particle indices this rank should have after resampling
    particles_want = resampling_indices[particles_have]

    # These are the ranks that have the particles this rank should have
    rank_has = floor.(Int, (particles_want .- 1) / nprt_per_rank)

    # We could work out how many sends and receives we have to do and allocate
    # this appropriately but, lazy
    reqs = Vector{MPI.Request}(undef, 0)

    # Send particles to processes that want them
    for (k,id) in enumerate(resampling_indices)
        rank_wants = floor(Int, (k - 1) / nprt_per_rank)
        if id in particles_have && rank_wants != my_rank
            local_id = id - my_rank * nprt_per_rank
            req = MPI.Isend(@view(particles[..,local_id]), rank_wants, id, MPI.COMM_WORLD)
            push!(reqs, req)
        end
    end

    # Receive particles this rank wants from ranks that have them
    # If I already have them, just do a local copy
    # Receive into a buffer so we dont accidentally overwrite stuff
    for (k,proc,id) in zip(1:nprt_per_rank, rank_has, particles_want)
        if proc == my_rank
            local_id = id - my_rank * nprt_per_rank
            @view(buffer[..,k]) .= @view(particles[..,local_id])
        else
            req = MPI.Irecv!(@view(buffer[..,k]), proc, id, MPI.COMM_WORLD)
            push!(reqs,req)
        end
    end

    # Wait for all comms to complete
    MPI.Waitall!(reqs)

    particles .= buffer

end

# Initialize arrays used by the filter
function init_filter(filter_params::FilterParameters, model_data, nprt_per_rank::Int, ::Vector{<:Random.AbstractRNG}, T::Type, ::BootstrapFilter)

    if MPI.Comm_rank(MPI.COMM_WORLD) == filter_params.master_rank
        weights = Vector{T}(undef, filter_params.nprt)
    else
        weights = Vector{T}(undef, nprt_per_rank)
    end

    resampling_indices = Vector{Int}(undef, filter_params.nprt)

    size = get_grid_size(model_data)
    n_state_var = get_n_state_var(model_data)

    statistics = Array{SummaryStat{T}, length(size) + 1}(undef, size..., n_state_var)
    avg_arr = Array{T, length(size) + 1}(undef, size..., n_state_var)
    var_arr = Array{T, length(size) + 1}(undef, size..., n_state_var)

    # Memory buffer used during copy of the states
    copy_buffer = Array{T, length(size) + 2}(undef, size..., n_state_var, nprt_per_rank)

    return (;weights, resampling_indices, statistics, avg_arr, var_arr, copy_buffer)
end

# Initialize arrays used by the filter
function init_filter(filter_params::FilterParameters, model_data, nprt_per_rank::Int, rng::Vector{<:Random.AbstractRNG}, T::Type, ::OptimalFilter)

    filter_data = init_filter(filter_params, model_data, nprt_per_rank, rng, T, BootstrapFilter())

    stations = get_stations(model_data)
    size = get_grid_size(model_data)
    domain_size = get_grid_domain_size(model_data)
    cell_size = get_grid_cell_size(model_data)

    grid = Grid(size...,cell_size...,domain_size...)
    grid_ext = Grid((grid.nx-1)*2, (grid.ny-1)*2, grid.dx, grid.dy, (grid.x_length-grid.dx)*2, (grid.y_length-grid.dy)*2)

    model_noise_params = get_model_noise_params(model_data)
    obs_noise_std = get_obs_noise_std(model_data)
    # Precompute two FFT plans, one in-place and the other out-of-place
    C = complex(T)
    tmp_array = Matrix{C}(undef, grid_ext.nx, grid_ext.ny)
    fft_plan, fft_plan! = FFTW.plan_fft(tmp_array), FFTW.plan_fft!(tmp_array)

    offline_matrices = init_offline_matrices(grid, grid_ext, stations, model_noise_params, obs_noise_std, fft_plan, fft_plan!, T)
    online_matrices = init_online_matrices(grid, grid_ext, stations, nprt_per_rank, T)

    return (; filter_data..., offline_matrices, online_matrices, stations, grid, grid_ext, rng, obs_noise_std, fft_plan, fft_plan!)
end

function update_particle_proposal!(model_data, filter_data, truth_observations, nprt_per_rank, filter_type::BootstrapFilter)

    update_particle_noise!(model_data, nprt_per_rank)

end

function update_particle_proposal!(model_data, filter_data, truth_observations, nprt_per_rank, filter_type::OptimalFilter)

        # Optimal Filter: After updating the particle dynamics, we apply the "optimal proposal" in
        #                 sample_height_proposal!() to the first state variable (height). We apply
        #                 a sample from the gaussian random field in update_particle_noise!() to the other
        #                 state variables (velocity).

        particles = get_particles(model_data)

        # Apply optimal proposal, the result will be in offline_matrices.samples
        sample_height_proposal!(@view(particles[:,:,1,:]),
                                filter_data.offline_matrices,
                                filter_data.online_matrices,
                                truth_observations,
                                filter_data.stations,
                                filter_data.grid,
                                filter_data.grid_ext,
                                filter_data.fft_plan,
                                filter_data.fft_plan!,
                                nprt_per_rank,
                                filter_data.rng[threadid()],
                                filter_data.obs_noise_std)

        # Add noise from the standard gaussian random field to all state variables in model_data
        update_particle_noise!(model_data, nprt_per_rank)

        # Overwrite the height state variable with the samples of the optimal proposal
        particles[:,:,1,:] .= filter_data.online_matrices.samples
        set_particles!(model_data, particles)

end

function run_particle_filter(init, filter_params::FilterParameters, model_params_dict::Dict, filter_type; rng::Union{Random.AbstractRNG,Nothing}=nothing)

    MPI.Init()

    my_rank = MPI.Comm_rank(MPI.COMM_WORLD)
    my_size = MPI.Comm_size(MPI.COMM_WORLD)

    # For now, assume that the particles can be evenly divided between ranks
    @assert mod(filter_params.nprt, my_size) == 0

    nprt_per_rank = Int(filter_params.nprt / my_size)

    if filter_params.enable_timers
        TimerOutputs.enable_debug_timings(ParticleDA)
    end
    timer = TimerOutput()

    nprt_per_rank = Int(filter_params.nprt / MPI.Comm_size(MPI.COMM_WORLD))

    _rng = let
        m = if isnothing(rng)
            Random.MersenneTwister(filter_params.random_seed + my_rank)
        else
            rng
        end
        [m; accumulate(Future.randjump, fill(big(10)^20, nthreads()-1), init=m)]
    end

    # Do memory allocations
    @timeit_debug timer "Model initialization" model_data = init(model_params_dict, nprt_per_rank, my_rank, _rng)

    # TODO: put the body of this block in a function
    @timeit_debug timer "Filter initialization" filter_data = init_filter(filter_params, model_data, nprt_per_rank, _rng, Float64, filter_type)

    @timeit_debug timer "get_particles" particles = get_particles(model_data)
    @timeit_debug timer "Mean and Var" get_mean_and_var!(filter_data.statistics, particles, filter_params.master_rank)

    # Write initial state (time = 0) + metadata
    if(filter_params.verbose && my_rank == filter_params.master_rank)
        @timeit_debug timer "IO" begin
            unpack_statistics!(filter_data.avg_arr, filter_data.var_arr, filter_data.statistics)
            write_snapshot(filter_params.output_filename,
                           model_data,
                           filter_data.avg_arr,
                           filter_data.var_arr,
                           filter_data.weights,
                           0)
        end
    end

    for it in 1:filter_params.n_time_step

        # integrate true synthetic wavefield
        @timeit_debug timer "True State Update and Process Noise" truth_observations = update_truth!(model_data, nprt_per_rank)

        # Forecast: Update tsunami forecast and get observations from it
        # Parallelised with threads.

        @timeit_debug timer "Particle Dynamics" update_particle_dynamics!(model_data, nprt_per_rank);
        @timeit_debug timer "Particle Proposal" update_particle_proposal!(model_data, filter_data, truth_observations, nprt_per_rank, filter_type)
        @timeit_debug timer "Particle Observations" model_observations = get_particle_observations!(model_data, nprt_per_rank)

        @timeit_debug timer "Particle Weights" get_log_weights!(@view(filter_data.weights[1:nprt_per_rank]),
                                                       truth_observations,
                                                       model_observations,
                                                       filter_params.weight_std)

        # Gather weights to master rank and resample particles.
        # Doing MPI collectives in place to save memory allocations.
        # This style with if statmeents is recommended instead of MPI.Gather_in_place! which is a bit strange.
        # Note that only master_rank allocates memory for all particles. Other ranks only allocate
        # for their chunk of state.
        if my_rank == filter_params.master_rank
            @timeit_debug timer "MPI Gather" MPI.Gather!(MPI.IN_PLACE,
                                                         UBuffer(filter_data.weights, nprt_per_rank),
                                                         filter_params.master_rank,
                                                         MPI.COMM_WORLD)
            @timeit_debug timer "Weights" normalized_exp!(filter_data.weights)
            @timeit_debug timer "Resample" resample!(filter_data.resampling_indices, filter_data.weights, _rng[threadid()])

        else
            @timeit_debug timer "MPI Gather" MPI.Gather!(filter_data.weights,
                                                         nothing,
                                                         filter_params.master_rank,
                                                         MPI.COMM_WORLD)
        end

        # Broadcast resampled particle indices to all ranks
        MPI.Bcast!(filter_data.resampling_indices, filter_params.master_rank, MPI.COMM_WORLD)

        @timeit_debug timer "get_particles" particles = get_particles(model_data)
        @timeit_debug timer "State Copy" copy_states!(particles,
                                                      filter_data.copy_buffer,
                                                      filter_data.resampling_indices,
                                                      my_rank,
                                                      nprt_per_rank)

        @timeit_debug timer "Mean and Var" get_mean_and_var!(filter_data.statistics, particles, filter_params.master_rank)

        if my_rank == filter_params.master_rank && filter_params.verbose

            @timeit_debug timer "IO" begin
                unpack_statistics!(filter_data.avg_arr, filter_data.var_arr, filter_data.statistics)
                write_snapshot(filter_params.output_filename, model_data, filter_data.avg_arr, filter_data.var_arr, filter_data.weights, it)
            end

        end

    end

    if filter_params.enable_timers

        if my_rank == filter_params.master_rank
            print_timer(timer)
        end

        if filter_params.verbose
            # Gather string representations of timers from all ranks and write them on master
            str_timer = string(timer)

            timer_lengths = MPI.Gather(sizeof(str_timer), filter_params.master_rank, MPI.COMM_WORLD)

            if my_rank == filter_params.master_rank
                timer_chars = MPI.Gatherv!(str_timer,
                                           MPI.VBuffer(Vector{UInt8}(undef, sum(timer_lengths)), timer_lengths),
                                           filter_params.master_rank,
                                           MPI.COMM_WORLD)
                @timeit_debug timer "IO" write_timers(timer_lengths, my_size, timer_chars, filter_params)
            else
                MPI.Gatherv!(str_timer, nothing, filter_params.master_rank, MPI.COMM_WORLD)
            end
        end
    end

    unpack_statistics!(filter_data.avg_arr, filter_data.var_arr, filter_data.statistics)

    return get_truth(model_data), filter_data.avg_arr, filter_data.var_arr
end

# Initialise params struct with user-defined dict of values.
function get_params(T, user_input_dict::Dict)

    user_input = (; (Symbol(k) => v for (k,v) in user_input_dict)...)
    params = T(;user_input...)

end

get_params(user_input_dict::Dict) = get_params(FilterParameters, user_input_dict)

# Initialise params struct with default values
get_params() = FilterParameters()

function read_input_file(path_to_input_file::String)

    # Read input provided in a yaml file. Overwrite default input parameters with the values provided.
    if isfile(path_to_input_file)
        user_input_dict = YAML.load_file(path_to_input_file)
    else
        @warn "Input file " * path_to_input_file * " not found, using default parameters"
        user_input_dict = Dict()
    end
    return user_input_dict

end

"""
    run_particle_filter(init, path_to_input_file::String, filter_type::ParticleFilter)

Run the particle filter.  `init` is the function which initialise the model,
`path_to_input_file` is the path to the YAML file with the input parameters.
`filter_type` is the particle filter to use.  See [`ParticleFilter`](@ref) for
the possible values.
"""
function run_particle_filter(init, path_to_input_file::String, filter_type::ParticleFilter; rng::Union{Random.AbstractRNG,Nothing}=nothing)

    MPI.Init()

    # Do I/O on rank 0 only and then broadcast params
    if MPI.Comm_rank(MPI.COMM_WORLD) == 0

        user_input_dict = read_input_file(path_to_input_file)

    else

        user_input_dict = nothing

    end

    user_input_dict = MPI.bcast(user_input_dict, 0, MPI.COMM_WORLD)

    return run_particle_filter(init, user_input_dict, filter_type; rng)

end

"""
    run_particle_filter(init, user_input_dict::Dict, filter_type::ParticleFilter)

Run the particle filter.  `init` is the function which initialise the model,
`user_input_dict` is the list of input parameters, as a `Dict`.  `filter_type`
is the particle filter to use.  See [`ParticleFilter`](@ref) for the possible
values.
"""
function run_particle_filter(init, user_input_dict::Dict, filter_type::ParticleFilter; rng::Union{Random.AbstractRNG,Nothing}=nothing)

    filter_params = get_params(FilterParameters, get(user_input_dict, "filter", Dict()))
    model_params_dict = get(user_input_dict, "model", Dict())

    return run_particle_filter(init, filter_params, model_params_dict, filter_type; rng)

end

end # module