text:
using Documenter, Spglib

makedocs(;
    modules=[Spglib],
    format=Documenter.HTML(),
    pages=[
        "Home" => "index.md",
        "Installation" => "installation.md",
        "API" => "api.md",
    ],
    repo="https://github.com/singularitti/Spglib.jl/blob/{commit}{path}#L{line}",
    sitename="Spglib.jl",
    authors="Qi Zhang <singularitti@outlook.com>",
    assets=String[],
)

deploydocs(;
    repo="github.com/singularitti/Spglib.jl",
)


text:
"""
Implements the Hungarian baseline from the paper.
"""
function expected_hungarian!(server::ArmAllocationServer, cbelt::ConveyorBelt)

    # Empty the current object-arm assignment record
    empty!(server.agent_task_allocation)

    # Set up matrix for hungarian
    # If out of reach, assign infinite cost
    # First, get vectors for active arms and boxes
    active_arm_names = Vector{String}(undef, 0)
    for arm_name in collect(keys(server.agent_set))
        if server.agent_prop_set[arm_name].obj_grabbed == "" # Not currently assigned
            push!(active_arm_names, arm_name)
        end
    end

    # Vector of box names and current x-coordinates
    box_xpos_vector = Vector{Tuple{String, Float64}}(undef, 0)
    for (name, route) in cbelt.box_routes
        push!(box_xpos_vector, (name, route[1].coordinates[1]))
    end

    # Only choose the n rightmost boxes if number of boxes > arms
    narms = length(active_arm_names)
    nboxes = length(box_xpos_vector)

    if narms == 0 || nboxes == 0
        @info "Either no active arms or no active boxes!"
        return
    end

    # Sort in descending order of position
    sort!(box_xpos_vector, by = x->x[2], rev=true)
    @debug box_xpos_vector

    # Iterate over box_xpos_vector until every arm can grasp at least one box
    boxes_to_consider = String[]
    arms_to_cover = Set(active_arm_names)
    arms_covered = Set{String}()

    for boxpos in box_xpos_vector

        # Consider box since going in descending order of position - most urgent first
        push!(boxes_to_consider, boxpos[1])

        # Look at what arms can potentially grasp box
        for arm_name in active_arm_names
            if boxpos[2] < server.agent_set[arm_name].reach[2]
                push!(arms_covered, arm_name)
            end
        end

        # We're done if we have considered at least narms boxes and all arms can grasp at least one
        if length(boxes_to_consider) >= narms && arms_covered == arms_to_cover
            break
        end
    end

    # Now create cost matrix of boxes vs arms
    # Each entry will have prob of failure
    # Don't need tp dict - will attempt first possible grasp
    nboxes = length(boxes_to_consider)
    cost_matrix = ones(narms, nboxes)

    # From the boxes to consider, compute the grasp failure prob for each active arm
    # Ignore (i.e., keep at 1.) if not reachable
    for (i, arm_name) in enumerate(active_arm_names)

        this_arm = server.agent_set[arm_name]

        # Iterate over boxes and consider time windows with given arm
        for (j, box_name) in enumerate(boxes_to_consider)

            if haskey(server.agent_task_windows, (arm_name, box_name))
                window = server.agent_task_windows[(arm_name, box_name)]

                # Assign weight to probability of grasp failure
                ref_time = max(server.current_time + 1.0, window[1])
                grasp_time = convert(Int64, floor(window[2] - server.agent_set[arm_name].downtime/2.0 - ref_time))
                fail_prob = 1.0 - cumprob(this_arm.grasp_model, grasp_time)

                cost_matrix[i, j] = fail_prob
            end
        end
    end

    @debug cost_matrix

    if minimum(cost_matrix) == 1.0
        @warn "All matchings are failures"
        return
    end

    # Run Hungarian algorithm
    assignment, cost = hungarian(cost_matrix)
    @debug assignment

    # From assignment, updated arm_obj_assignment with the timestamp
    # and the corresponding grasp location, using get_box_point
    for (i, assgn) in enumerate(assignment)

        if assgn != 0

            # Then arm reduced_row_orig[i] is assigned to box reduced_col_orig[assign]
            true_arm_name = active_arm_names[i]
            true_box_name = boxes_to_consider[assgn]
            @debug (true_arm_name, true_box_name)

            # Only assign an arm to object if it actually has a valid time window
            if haskey(server.agent_task_windows, (true_arm_name, true_box_name))
                window = server.agent_task_windows[(true_arm_name, true_box_name)]

                ref_time = max(server.current_time + 1.0, window[1])

                server.agent_task_allocation[true_arm_name] = (name = true_box_name, time =ref_time)
            else
                @debug "Arm $(true_arm_name) goes unassigned this time"
            end
        end
    end

    @debug server.agent_task_allocation
end

"""
Implements the EDD baseline. Each arm is assigned to the rightmost object in its current workspace.
"""
function earliest_due_date!(server::ArmAllocationServer, cbelt::ConveyorBelt)

    empty!(server.agent_task_allocation)

    active_arm_names = Vector{String}(undef, 0)
    for arm_name in collect(keys(server.agent_set))
        if server.agent_prop_set[arm_name].obj_grabbed == "" # Not currently assigned
            push!(active_arm_names, arm_name)
        end
    end

    # Vector of box names and next x-coordinates for each box
    box_xpos_vector = Vector{Tuple{String, Float64}}(undef, 0)
    for (name, route) in cbelt.box_routes
        if length(route) >= 2
            push!(box_xpos_vector, (name, route[2].coordinates[1]))
        end
    end

    narms = length(active_arm_names)
    nboxes = length(box_xpos_vector)

    if narms == 0 || nboxes == 0
        @info "Either no active arms or no active boxes!"
        return
    end

    sort!(box_xpos_vector, by = x->x[2], rev=true)

    boxes_covered = Set{String}()

    # For each arm, look at rightmost object in its workspace that it can possibly grasp
    for arm_name in active_arm_names
        for (box_name, pos) in box_xpos_vector

            # If below condition is not satisfied by any box, arm stays unassigned
            if ~(box_name in boxes_covered) && pos < server.agent_set[arm_name].reach[2]

                window = server.agent_task_windows[(arm_name, box_name)] # Should not give err by construction
                ref_time = max(server.current_time + 1.0, window[1])

                grasp_coords = get_box_point(cbelt, box_name, ref_time)
                server.agent_task_allocation[arm_name] = (name = box_name, time = ref_time)
                push!(boxes_covered, box_name)
            end
        end
    end

    @debug server.agent_task_allocation
end

## POMDPs.MDP definition for MCTS and Q-Learning
const BeltState = Vector{Int64}
@with_kw struct ArmState
    obj_grabbed::Bool   = false
end
struct ConvBeltSystemState
    belt_disc::BeltState
    arm_states::Vector{ArmState}
end

"""
Simple discrete-time MDP formulation of the conveyor belt problem.
"""
@with_kw mutable struct ConvBeltMDP <: POMDPs.MDP{ConvBeltSystemState,Vector{Int64}}
    cbelt::ConveyorBelt
    server::ArmAllocationServer
    belt_disc_factor::Float64 # Ideally 1/bdf should give an integer
    n_belt_slots::Int64
    speed_skip_factor::Int64
    per_arm_slots::Int64
    horizon::Int64
    max_per_slot::Int64
    states::Vector{ConvBeltSystemState}
end

"""
Instantiates the MDP by discretizing the belt according to the chosen resolution
"""
function create_convbelt_mdp(cbelt::ConveyorBelt, server::ArmAllocationServer,
                             disc_factor::Float64, horizon::Int64,
                             max_per_slot::Int64=typemax(Int64))

    speed_skip_factor = max(1, convert(Int64, floor(cbelt.curr_speed/disc_factor)))
    arm_reach = server.agent_set[server.agent_ordering[1]].reach # Assuming same for all
    per_arm_slots = convert(Int64, ceil((arm_reach[2]-arm_reach[1])/disc_factor))

    n_belt_slots = convert(Int64, ceil(1.0/disc_factor))

    # Create states IF explicit
    statevect = ConvBeltSystemState[] # Empty by default
    if max_per_slot < typemax(Int64)
        belt_disc_vals = [0:max_per_slot for i = 1:n_belt_slots]
        belt_disc_vector_vals = [collect(Iterators.flatten(a)) for a in Iterators.product(belt_disc_vals...)]
        arm_state_vals = [(false,true) for i = 1:length(server.agent_ordering)]
        arm_state_vector_vals = [collect(Iterators.flatten(a)) for a in Iterators.product(arm_state_vals...)]
        statevect = [ConvBeltSystemState(bdv, [ArmState(a) for a in asv]) for bdv in belt_disc_vector_vals
                    for asv in arm_state_vector_vals]
    end

    return ConvBeltMDP(cbelt, server, disc_factor, n_belt_slots,
                           speed_skip_factor, per_arm_slots, horizon,
                           max_per_slot, statevect)
end

"""
Returns the ConvBeltSystemState struct that represents the current state of the MDP.
"""
function get_current_belt_mdp_state(mdp::ConvBeltMDP, rng::AbstractRNG=Random.GLOBAL_RNG)

    # Populate belt
    belt_disc = zeros(Int64, mdp.n_belt_slots)

    # Obtain box current position, find discrete idx on belt and increment
    for (box_name, route) in mdp.cbelt.box_routes
        curr_pos = route[1].coordinates[1]
        box_belt_idx = convert(Int64, ceil(curr_pos/mdp.belt_disc_factor))

        # NOTE: Truncating here. This is ALL we have to do
        belt_disc[box_belt_idx] = min(belt_disc[box_belt_idx] + 1, mdp.max_per_slot)
    end

    # Now get arms
    arm_states = Vector{ArmState}(undef, length(mdp.server.agent_ordering))

    for (ia, arm_name) in enumerate(mdp.server.agent_ordering)

        arm_props = mdp.server.agent_prop_set[arm_name]

        if arm_props.obj_grabbed != ""
            arm_states[ia] = ArmState(true)
        else
            arm_states[ia] = ArmState(false)
        end
    end

    # Finally return state with current time
    return ConvBeltSystemState(belt_disc, arm_states)
end

POMDPs.discount(mdp::ConvBeltMDP) = 1.0
POMDPs.isterminal(mdp::ConvBeltMDP, s::ConvBeltSystemState) = (mdp.server.current_time >= mdp.horizon) # TODO: No notion of terminal state

function POMDPs.actions(mdp::ConvBeltMDP, s::ConvBeltSystemState)

    arm_actions = UnitRange{Int64}[]
    for as in s.arm_states
        if as.obj_grabbed == true
            push!(arm_actions, 0:0) # No-op
        else
            push!(arm_actions, 0:mdp.per_arm_slots) # 0 for home, 1:belt_slots for rest
        end
    end

    return [collect(Iterators.flatten(a)) for a in Iterators.product(arm_actions...)]
end

function POMDPs.gen(mdp::ConvBeltMDP, s::ConvBeltSystemState, a::Vector{Int64}, rng::AbstractRNG)

    # First set new arm_states
    new_arm_states = ArmState[]
    box_positions_grabbed = Int64[]

    # go along belt disc and copy to next spot based on speed skip factor
    new_belt_disc = zeros(Int64, mdp.n_belt_slots)
    for idx = mdp.n_belt_slots-mdp.speed_skip_factor:-1:1
        new_belt_disc[idx + mdp.speed_skip_factor] = s.belt_disc[idx]
    end

    for (as, aa, arm_name) in zip(s.arm_states, a, mdp.server.agent_ordering)

        # Can be 0 only if obj was grabbed in s, so deposit to bin and free up
        if aa == 0
            push!(new_arm_states, ArmState(false))
        else
            # Now test out grab
            arm_facts = mdp.server.agent_set[arm_name]
            arm_slot = convert(Int64, floor(arm_facts.reach[1]/mdp.belt_disc_factor)) + aa

            # Grab if box there and rand(rng)....
            if (rand(rng) < arm_facts.grasp_model.bernoulli_prob &&
                new_belt_disc[arm_slot] >= 1)
                # Successful grab!!
                push!(box_positions_grabbed, arm_slot)
                push!(new_arm_states, ArmState(true))
            else
                # Unsuccessful or no box there?
                push!(new_arm_states, ArmState(false))
            end
        end # aa == -1
    end

    # decrement box counts for lifted slots
    for bp in box_positions_grabbed
        new_belt_disc[bp] -= 1
    end

    return (sp=ConvBeltSystemState(new_belt_disc, new_arm_states), r = length(box_positions_grabbed))
end

"""
The rollout policy for the conveyor belt uses the EDD strategy.
"""
struct EDDPolicy <: Policy
    mdp::ConvBeltMDP
end

function POMDPs.action(pol::EDDPolicy, s::ConvBeltSystemState)

    mdp = pol.mdp

    joint_action = Int64[]
    for (as, arm_name) in zip(s.arm_states, mdp.server.agent_ordering)
        if as.obj_grabbed == true
            push!(joint_action, 0) # No-op
        else
            # Find the furthest along slot that will still be valid at next step
            arm_facts = mdp.server.agent_set[arm_name]
            last_arm_slot = convert(Int64, ceil(arm_facts.reach[2]/mdp.belt_disc_factor))
            first_arm_slot = convert(Int64, floor(arm_facts.reach[1]/mdp.belt_disc_factor)) + 1

            aa = 0 # Default value
            for idx = last_arm_slot-mdp.speed_skip_factor:-1:first_arm_slot
                if s.belt_disc[idx + mdp.speed_skip_factor] >= 1
                    aa = idx + mdp.speed_skip_factor - first_arm_slot
                    break
                end
            end

            push!(joint_action, aa)
        end
    end

    return joint_action
end

## Other requirements for QLearning
POMDPs.initialstate(mdp::ConvBeltMDP, rng::AbstractRNG) = get_current_belt_mdp_state(mdp, rng)

# Not state-dependent actions
function POMDPs.actions(mdp::ConvBeltMDP)
    arm_ac_range = [0:mdp.per_arm_slots for i = 1:length(mdp.server.agent_ordering)]
    return [collect(Iterators.flatten(a)) for a in Iterators.product(arm_ac_range...)]
end

function POMDPs.actionindex(mdp::ConvBeltMDP, action::Vector{Int64})

    interm_idxs = [ (a+1) for a in action]
    ac_tup = Tuple(1:mdp.per_arm_slots+1 for i = 1:length(mdp.server.agent_ordering))
    return LinearIndices(ac_tup)[interm_idxs...]

end

POMDPs.states(mdp::ConvBeltMDP) = mdp.states

Base.isequal(s1::ConvBeltSystemState, s2::ConvBeltSystemState) = (s1.belt_disc == s2.belt_disc) &&
                            (s1.arm_states == s2.arm_states)

POMDPs.stateindex(mdp::ConvBeltMDP, s::ConvBeltSystemState) = findfirst(isequal(s), mdp.states)

"""
Q-Learning for the conveyor belt uses IQL for each agent.
"""
struct SingleArmState
    belt_disc::Vector{Int64}
    obj_grabbed::Bool
end

@with_kw mutable struct ConvBeltSingleArmMDP <: POMDPs.MDP{SingleArmState,Int64}
    cbelt::ConveyorBelt
    server::ArmAllocationServer
    belt_disc_factor::Float64 # Ideally 1/bdf should give an integer
    speed_skip_factor::Int64
    per_arm_slots::Int64
    max_per_slot::Int64
    states::Vector{SingleArmState}
end

function create_convbelt_singlearm_mdp(cbelt::ConveyorBelt, server::ArmAllocationServer,
                                       disc_factor::Float64,
                                       max_per_slot::Int64=typemax(Int64))

    speed_skip_factor = max(1, convert(Int64, floor(cbelt.curr_speed/disc_factor)))
    arm_reach = server.agent_set[server.agent_ordering[1]].reach # Assuming same for all
    per_arm_slots = convert(Int64, ceil((arm_reach[2]-arm_reach[1])/disc_factor))

    # Create states IF explicit
    statevect = ConvBeltSystemState[] # Empty by default
    if max_per_slot < typemax(Int64)
        belt_disc_vals = [0:max_per_slot for i = 1:per_arm_slots+1]
        belt_disc_vector_vals = [collect(Iterators.flatten(a)) for a in Iterators.product(belt_disc_vals...)]
        arm_state_vals = [false, true]
        statevect = [SingleArmState(bdv, asv) for bdv in belt_disc_vector_vals for asv in arm_state_vals]
    end

    return ConvBeltSingleArmMDP(cbelt, server, disc_factor,
                                speed_skip_factor, per_arm_slots,
                                max_per_slot, statevect)
end

POMDPs.initialstate(mdp::ConvBeltSingleArmMDP, rng::AbstractRNG) = SingleArmState([rand(rng, 1:mdp.max_per_slot) for i in 1:mdp.per_arm_slots+1], false)

POMDPs.discount(mdp::ConvBeltSingleArmMDP) = 1.0

# Need a gen for the MDP and then need a true outer gen conveyor_mcts?
# that is run before calling propagate belt and propagate arm
# It looks at the action obtained from MCTS and implements it in the actual outer env
function POMDPs.gen(mdp::ConvBeltSingleArmMDP, s::SingleArmState, a::Int64, rng::AbstractRNG)

    # go along belt disc and copy to next spot based on speed skip factor
    new_belt_disc = zeros(Int64, mdp.per_arm_slots+1)
    for idx = mdp.per_arm_slots+1-mdp.speed_skip_factor:-1:1
        new_belt_disc[idx + mdp.speed_skip_factor] = s.belt_disc[idx]
    end

    obj_grabbed = false
    r = 0
    if s.obj_grabbed == false # Only then _could_ it be true

        any_arm = mdp.server.agent_set[mdp.server.agent_ordering[1]]

        if rand(rng) < any_arm.grasp_model.bernoulli_prob && new_belt_disc[a+1] >= 1
            r = 1
            obj_grabbed = true
            new_belt_disc[a+1] -= 1
        end
    end

    return (sp=SingleArmState(new_belt_disc, obj_grabbed), r = r)
end

# Not state-dependent actions
POMDPs.actions(mdp::ConvBeltSingleArmMDP) = collect(1:mdp.per_arm_slots)
POMDPs.actionindex(mdp::ConvBeltSingleArmMDP, action::Int64) = action
POMDPs.states(mdp::ConvBeltSingleArmMDP) = mdp.states

Base.isequal(s1::SingleArmState, s2::SingleArmState) = (s1.belt_disc == s2.belt_disc) && (s1.obj_grabbed == s2.obj_grabbed)

POMDPs.stateindex(mdp::ConvBeltSingleArmMDP, s::SingleArmState) = findfirst(isequal(s), mdp.states)

function get_singlearm_mdp_state(mdp::ConvBeltSingleArmMDP, agent_idx::Int64)

    per_arm_disc = zeros(Int64, mdp.per_arm_slots+1)

    arm_facts = mdp.server.agent_set[mdp.server.agent_ordering[agent_idx]]

    start_point = arm_facts.reach[1] - mdp.belt_disc_factor

    for (box_name, route) in mdp.cbelt.box_routes
        curr_pos = route[1].coordinates[1]

        if curr_pos > start_point && curr_pos < arm_facts.reach[2]
            arm_slot_idx = convert(Int64, ceil((curr_pos-start_point)/mdp.belt_disc_factor))
            per_arm_disc[arm_slot_idx] = min(per_arm_disc[arm_slot_idx] + 1, mdp.max_per_slot)
        end
    end

    arm_props = mdp.server.agent_prop_set[mdp.server.agent_ordering[agent_idx]]
    obj_grabbed = ~(isempty(arm_props.obj_grabbed))

    return SingleArmState(per_arm_disc, obj_grabbed)
end

function update_server_with_mdp_action!(mdp::Union{ConvBeltMDP,ConvBeltSingleArmMDP},
                                       action::Vector{Int64})

    next_time = convert(Int64, mdp.server.current_time + 1)

    # Iterate over action and
    for (agent_action, arm_name) in zip(action, mdp.server.agent_ordering)
        arm_facts = mdp.server.agent_set[arm_name]
        arm_props = mdp.server.agent_prop_set[arm_name]

        if agent_action >= 1 && arm_props.obj_grabbed == ""
            arm_slot = convert(Int64, floor(arm_facts.reach[1]/mdp.belt_disc_factor)) + agent_action

            # Iterate over box routes and choose slot
            any_box = false
            arm_assgn = (name="", time=Inf)
            for (box_name, route) in mdp.cbelt.box_routes
                if length(route) >= 2
                    next_pos = route[2].coordinates[1]
                    box_belt_idx = convert(Int64, ceil(next_pos/mdp.belt_disc_factor))
                    if box_belt_idx == arm_slot && next_pos >= arm_facts.reach[1] && next_pos <= arm_facts.reach[2]
                        any_box = true
                        arm_assgn = (name = box_name, time = next_time)
                        break
                    end # box_belt_idx
                end # length(route)
            end

            if any_box
                mdp.server.agent_task_allocation[arm_name] = arm_assgn
            end
        end
    end
end

ext: jl


text:
using DatasToolbox
using DataFrames

df = makeTestDF(Float64, Float64, Float64, Float64)
names!(df, [:x1, :x2, :x3, :x4])

dh = DataHandler{Float64}(df, input_cols=[:x1, :x2], output_cols=[:x3],
                          normalize_cols=[:x1, :x2],
                          testfrac=0.1, shuffle=true, userange=true)

a = 0.5
@split!(dh, :x1 > a)

assign!(dh)

X_train, y_train = getTrainData(dh)

max_stars_repo_name: JuliaTagBot/DatasToolbox

max_issues_repo_name: JuliaTagBot/DatasToolbox

max_forks_repo_name: JuliaTagBot/DatasToolbox


text:
module test_bukdu_logging

using Test
import Bukdu: Plug
import .Plug: handle_message
import Base.CoreLogging: Debug, Info, Warn, Error

logger = Plug.Logger()
@test logger.min_level == Debug

_module = nothing
group = nothing
id = nothing
filepath = nothing
line = nothing
handle_message(logger, Debug, "d", _module, group, id, filepath, line)
handle_message(logger, Info,  "i", _module, group, id, filepath, line)
handle_message(logger, Warn,  "w", _module, group, id, filepath, line)
handle_message(logger, Error, "e", _module, group, id, filepath, line)

end # module test_bukdu_logging

max_stars_repo_name: wookay/TestBukdu

max_issues_repo_name: wookay/TestBukdu

max_forks_repo_name: wookay/TestBukdu


text:
# construct a Sink from an existing CSV Source
function Sink(s::CSV.Source; append::Bool=false)
    data = open(s.fullpath, append ? "a" : "w")
    seek(data, append ? filesize(s.fullpath) : s.datapos)
    return Sink(s.schema, s.options, s.fullpath, data, s.datapos, !append)
end

function Sink(io::Union{AbstractString,IO},
              schema::Data.Schema=Data.EMPTYSCHEMA;
              delim::Char=',',
              quotechar::Char='"',
              escapechar::Char='\\',
              null::AbstractString="",
              dateformat::Union{AbstractString,Dates.DateFormat}=Dates.ISODateFormat,
              header::Bool=true,
              append::Bool=false)
    delim = delim % UInt8; quotechar = quotechar % UInt8; escapechar = escapechar % UInt8
    dateformat = isa(dateformat, AbstractString) ? Dates.DateFormat(dateformat) : dateformat
    name, io = isa(io, AbstractString) ? (io, open(io, append ? "a" : "w")) : ("__IO__", io)
    return Sink(schema, CSV.Options(delim, quotechar, escapechar, COMMA, PERIOD,
            ascii(null), null=="", dateformat, dateformat == Dates.ISODateFormat, -1, 0, 1, DataType[]), name, io, 0, header && !append)
end

Base.close(s::Sink) = (applicable(close, s.data) && close(s.data); return nothing)
Base.flush(s::Sink) = (applicable(flush, s.data) && flush(s.data); return nothing)
Base.position(s::Sink) = applicable(position, s.data) ? position(s.data) : 0

function writeheaders(source, sink::Sink)
    rows, cols = size(Data.schema(source))
    q = Char(sink.options.quotechar); e = Char(sink.options.escapechar)
    h = Data.header(source)
    for col = 1:cols
        print(sink.data, q, replace("$(h[col])", q, "$e$q"), q)
        print(sink.data, ifelse(col == cols, Char(NEWLINE), Char(sink.options.delim)))
    end
    sink.datapos = position(sink)
    return nothing
end

# DataStreams interface
Data.streamtypes{T<:CSV.Sink}(::Type{T}) = [Data.Field]

# Constructors
function Sink{T}(source, ::Type{T}, append::Bool, file::AbstractString)
    sch = Data.schema(source)
    io = open(file, append ? "a" : "w")
    return Sink(io, sch; append=append)
end

function Sink{T}(source, ::Type{T}, append::Bool, io::IO)
    sch = Data.schema(source)
    append ? io : seekstart(io)
    return Sink(io, sch; append=append)
end

function Sink{T}(sink, source, ::Type{T}, append::Bool)
    if !isopen(sink.data)
        # try to re-open if we have filename
        sink.path == "__IO__" && throw(ArgumentError("unable to stream to a closed sink"))
        sink.data = open(sink.path, append ? "a" : "w")
    else
        !append && seekstart(sink.data)
    end
    append && (sink.header = false)
    sink.schema = Data.schema(source)
    return sink
end

writefield(io::Sink, val, col, N) = (col == N ? println(io.data, val) : print(io.data, val, Char(io.options.delim)); return nothing)
function writefield(sink::Sink, val::AbstractString, col, cols)
    q = Char(sink.options.quotechar); e = Char(sink.options.escapechar)
    print(sink.data, q, replace(string(val), q, "$e$q"), q)
    print(sink.data, ifelse(col == cols, Char(NEWLINE), Char(sink.options.delim)))
    return nothing
end

function writefield(sink::Sink, val::Dates.TimeType, col, cols)
    q = Char(sink.options.quotechar); e = Char(sink.options.escapechar)
    val = sink.options.datecheck ? string(val) : Dates.format(val, sink.options.dateformat)
    print(sink.data, val)
    print(sink.data, ifelse(col == cols, Char(NEWLINE), Char(sink.options.delim)))
    return nothing
end

function writefield{T}(sink::Sink, val::Nullable{T}, col, cols)
    writefield(sink, isnull(val) ? sink.options.null : get(val), col, cols)
    return nothing
end
if isdefined(:NAtype)
function writefield(sink::Sink, val::NAtype, col, cols)
    writefield(sink, sink.options.null, col, cols)
    return nothing
end
end

function Data.streamfield!{T}(sink::CSV.Sink, source, ::Type{T}, row, col, cols)
    val = Data.getfield(source, T, row, col)
    writefield(sink, val, col, cols)
    return nothing
end

Data.open!(sink::CSV.Sink, source) = (sink.header && writeheaders(source, sink); return nothing)
Data.flush!(sink::CSV.Sink) = flush(sink)
Data.close!(sink::CSV.Sink) = close(sink)

"""
`CSV.write(fullpath::Union{AbstractString,IO}, source::Type{T}, args...; kwargs...)` => `CSV.Sink`
`CSV.write(fullpath::Union{AbstractString,IO}, source::Data.Source; kwargs...)` => `CSV.Sink`

write a `Data.Source` out to a `CSV.Sink`.

Positional Arguments:

* `fullpath`; can be a file name (string) or other `IO` instance
* `source` can be the *type* of `Data.Source`, plus any required `args...`, or an already constructed `Data.Source` can be passsed in directly (2nd method)

Keyword Arguments:

* `delim::Union{Char,UInt8}`; how fields in the file will be delimited; default is `UInt8(',')`
* `quotechar::Union{Char,UInt8}`; the character that indicates a quoted field that may contain the `delim` or newlines; default is `UInt8('"')`
* `escapechar::Union{Char,UInt8}`; the character that escapes a `quotechar` in a quoted field; default is `UInt8('\\')`
* `null::String`; the ascii string that indicates how NULL values will be represented in the dataset; default is the emtpy string `""`
* `dateformat`; how dates/datetimes will be represented in the dataset; default is ISO-8601 `yyyy-mm-ddTHH:MM:SS.s`
* `header::Bool`; whether to write out the column names from `source`
* `append::Bool`; start writing data at the end of `io`; by default, `io` will be reset to the beginning before writing
"""
function write{T}(io::Union{AbstractString,IO}, ::Type{T}, args...;
              delim::Char=',',
              quotechar::Char='"',
              escapechar::Char='\\',
              null::AbstractString="",
              dateformat::Union{AbstractString,Dates.DateFormat}=Dates.ISODateFormat,
              header::Bool=true,
              append::Bool=false)
    sink = Sink(io; delim=delim, quotechar=quotechar, escapechar=escapechar, null=null,
                                         dateformat=dateformat, header=header, append=append)
    Data.stream!(T(args...), sink, append)
    Data.close!(sink)
    return sink
end
function write(io::Union{AbstractString,IO}, source;
              delim::Char=',',
              quotechar::Char='"',
              escapechar::Char='\\',
              null::AbstractString="",
              dateformat::Union{AbstractString,Dates.DateFormat}=Dates.ISODateFormat,
              header::Bool=true,
              append::Bool=false)
    sink = Sink(io, Data.schema(source); delim=delim, quotechar=quotechar, escapechar=escapechar, null=null,
                                         dateformat=dateformat, header=header, append=append)
    Data.stream!(source, sink, append)
    Data.close!(sink)
    return sink
end

write{T}(sink::Sink, ::Type{T}, args...; append::Bool=false) = (sink = Data.stream!(T(args...), sink, append); Data.close!(sink); return sink)
write(sink::Sink, source; append::Bool=false) = (sink = Data.stream!(source, sink, append); Data.close!(sink); return sink)


text:
function convert_data(data)
    a_t  = [];
    a_t1 = [];
    s_t  = [];
    s_t1 = [];
    for t = 1:(nT-1)
        a_t  = vcat(a_t,data.action[:,t]);
        s_t  = vcat(s_t,data.state[:,t]);
        a_t1 = vcat(a_t1,data.action[:,t+1]);
        s_t1 = vcat(s_t1,data.action[:,t+1]);
    end
    return(action_t=a_t,action_t_1=a_t1,state_t=s_t,state_t_1=s_t1);
end


text:
@testset "BrukerFile" begin

b = BrukerFile( joinpath(datadir, "BrukerFile", "2D_RARE") )

acq = RawAcquisitionData(b)
acqData = AcquisitionData(acq)
N = acqData.encodingSize

params = Dict{Symbol, Any}()
params[:reco] = "direct"
params[:reconSize] = (N[1],N[2]) #this should be clear from context

Ireco = reconstruction(acqData, params)
exportImage( joinpath(tmpdir, "brukerCart.png"), abs.(Ireco[:,:,1,1,1]))

# Convert to ISMRMRD file
fout = ISMRMRDFile(joinpath(tmpdir, "brukerfileCart.h5"))
save(fout, acq)

# Reco data stored in BrukerFile
Iloaded = recoData(b)
@test size(Iloaded) == (128, 128, 15)

## Test reconstruction for multi-coil datasets (2D and 3D FLASH)
listBrukFiles = ["2D_FLASH","3D_FLASH"]
listNormValues = [0.02, 0.15]

for i = 1:length(listBrukFiles)
    b = BrukerFile( joinpath(datadir, "BrukerFile", listBrukFiles[i]) )
    raw = RawAcquisitionData(b)
    acq = AcquisitionData(raw)
    params = Dict{Symbol, Any}()
    params[:reco] = "direct"
    if (acq.encodingSize[3]>1)
        params[:reconSize] = (acq.encodingSize[1],acq.encodingSize[2],acq.encodingSize[3]);
    else
        params[:reconSize] = (acq.encodingSize[1],acq.encodingSize[2]);
    end
    Ireco = reconstruction(acq, params);
    @test size(Ireco) == (raw.params["encodedSize"][1], raw.params["encodedSize"][2], raw.params["encodedSize"][3], 1, raw.params["receiverChannels"])

    Isos = sqrt.(sum(abs.(Ireco).^2,dims=5));
    Isos = Isos ./ maximum(Isos);

    I2dseq = recoData(b)
    I2dseq = I2dseq ./ maximum(I2dseq);

    @test norm(vec(I2dseq)-vec(Isos))/norm(vec(I2dseq)) < listNormValues[i]
end
end


text:
@testset "DTM" begin
    sample_file = joinpath(dirname(@__FILE__), "data", "poem.txt")

    fd = FileDocument(sample_file)
    sd = StringDocument(text(fd))

    crps = Corpus(Any[fd, sd])

    m = DocumentTermMatrix(crps)
    dtm(m)
    dtm(m, :dense)

    update_lexicon!(crps)

    m = DocumentTermMatrix(crps)
    dtm(m)
    dtm(m, :dense)

    tf_idf(dtm(m, :dense))

    dtv(crps[1], lexicon(crps))

    hash_dtv(crps[1], TextHashFunction())
    hash_dtv(crps[1])

    dtm1 = dtm(crps)
    dtm1sp = sparse(dtm(crps))
    hash_dtm(crps)

    tdm(crps)
    hash_tdm(crps)

    # construct a DocumentTermMatrix from a crps and a custom terms vector
    terms = ["And", "notincrps"]
    m = DocumentTermMatrix(crps,terms)
    @test size(dtm(m),1) == length(terms)
    @test terms == m.terms
    @test size(dtm(m),2) == length(crps)

    # construct a DocumentTermMatrix from a crps and a custom lexicon
    lex = Dict("And"=>1, "notincrps"=>4)
    m = DocumentTermMatrix(crps,lex)
    @test size(dtm(m),1) == length(keys(lex))
    @test size(dtm(m),1) == length(m.terms)
    @test size(dtm(m),2) == length(crps)

    # construct a DocumentTermMatrix from a dtm and terms vector
    terms = m.terms
    m2 = DocumentTermMatrix(dtm1,terms)
    @test m.column_indices == m2.column_indices
    m2 = DocumentTermMatrix(dtm1sp,terms)
    @test m.column_indices == m2.column_indices

    # test serialization and deserialization
    mktemp() do path, io
        serialize(io, m2)
        close(io)
        open(path, "r") do rio
            m3 = deserialize(rio)
            @test typeof(m2) == typeof(m3)
            @test m2.terms == m3.terms
            @test m2.dtm == m3.dtm
            @test m2.column_indices == m3.column_indices
        end
    end

    # test prune! and merge!
    crps1 = Corpus([StringDocument("one two three"), StringDocument("two three four")])
    crps2 = Corpus([StringDocument("two three four"), StringDocument("three four five")])
    update_lexicon!(crps1)
    update_lexicon!(crps2)
    dtm1 = DocumentTermMatrix(crps1)
    dtm2 = DocumentTermMatrix(crps2)

    prune!(dtm1, nothing; compact=false)
    @test length(dtm1.terms) == 4
    @test size(dtm1.dtm) == (2,4)

    prune!(dtm1, [1]; compact=false)
    @test length(dtm1.terms) == 4
    @test size(dtm1.dtm) == (1,4)

    dtm1 = DocumentTermMatrix(crps1)
    prune!(dtm1, [1]; compact=true)
    @test length(dtm1.terms) == 3
    @test size(dtm1.dtm) == (1,3)

    dtm1 = DocumentTermMatrix(crps1)
    prune!(dtm1, [1]; compact=true, retain_terms=["one"])
    @test length(dtm1.terms) == 4
    @test size(dtm1.dtm) == (1,4)

    merge!(dtm1, dtm2)
    @test size(dtm1.dtm) == (3,5)
    @test sum(dtm1.dtm, dims=(1,)) == [1 3 0 3 2]
    @test dtm1.terms == ["five", "four", "one", "three", "two"]

    dtm2 = DocumentTermMatrix(crps2)
    dtm1.dtm = similar(dtm1.dtm, 0, dtm1.dtm.n)
    merge!(dtm1, dtm2)
    @test dtm1.terms == ["five", "four", "one", "three", "two"]
    @test size(dtm1.dtm) == (2,5)
    @test sum(dtm1.dtm, dims=(1,)) == [1 2 0 2 1]

    dtm2 = DocumentTermMatrix(crps2)
    dtm1.dtm = similar(dtm1.dtm, 0, dtm1.dtm.n)
    merge!(dtm2, dtm1)
    @test dtm2.terms == ["five", "four", "three", "two"]
    @test size(dtm2.dtm) == (2,4)
    @test sum(dtm2.dtm, dims=(1,)) == [1 2 2 1]
end


text:
using Avalon
import Avalon.fit!
using Distributions
using MLDatasets
using MLDataUtils
using BenchmarkTools

mutable struct VAE
    # encoder / recognizer
    enc_i2h::Linear
    enc_h2mu::Linear
    enc_h2logsigma2::Linear
    # decoder / generator
    dec_z2h::Linear
    dec_h2o::Linear
end

function Base.show(io::IO, m::VAE)
    print(io, "VAE()")
end

VAE(n_i, n_h, n_z) =
    VAE(
        # encoder
        Linear(n_i, n_h),
        Linear(n_h, n_z),
        Linear(n_h, n_z),
        # decoder
        Linear(n_z, n_h),
        Linear(n_h, n_i),
    )

function encode(m::VAE, x)
    he1 = tanh.(m.enc_i2h(x))
    mu = m.enc_h2mu(he1)
    log_sigma2 = m.enc_h2logsigma2(he1)
    return mu, log_sigma2
end

function decode(m::VAE, z)
    hd1 = tanh.(m.dec_z2h(z))
    x_rec = logistic.(m.dec_h2o(hd1))
    return x_rec
end

function vae_cost(m::VAE, eps, x)
    mu, log_sigma2 = encode(m, x)
    z = mu .+ sqrt.(exp.(log_sigma2)) .* eps
    x_rec = decode(m, z)
    # loss
    rec_loss = -sum(x .* log.(1f-10 .+ x_rec) .+ (1 .- x) .* log.(1f-10 + 1 .- x_rec); dims=1)
    KLD = -0.5f0 .* sum(1 .+ log_sigma2 .- mu .^ 2.0f0 - exp.(log_sigma2); dims=1)
    cost = mean(rec_loss .+ KLD)
end

function fit!(m::VAE, X::AbstractMatrix{T};
              n_epochs=50, batch_size=100, opt=SGD(1e-4; momentum=0), device=CPU()) where T
    for epoch in 1:n_epochs
        print("Epoch $epoch: ")
        epoch_cost = 0
        t = @elapsed for (i, x) in enumerate(eachbatch(X, size=batch_size))
            x = device(x)
            eps = typeof(x)(rand(Normal(0, 1), size(m.enc_h2mu.W, 1), batch_size))
            cost, g = grad(vae_cost, m, eps, x)
            update!(opt, m, g[2])
            epoch_cost += cost
        end
        println("avg_cost=$(epoch_cost / (size(X,2) / batch_size)), elapsed=$t")
    end
    return m
end

function reconstruct(m::VAE, x::AbstractVector)
    x = reshape(x, length(x), 1)
    mu, _ = encode(m, x)
    z = mu
    x_rec = decode(m, z)
    return x_rec
end

function show_pic(x)
    reshape(x, 28, 28)' |> imshow
end

function show_recon(m, x)
    x_ = reconstruct(m, x)
    show_pic(x)
    show_pic(x_)
end

function main(device=best_available_device())
    m = VAE(784, 500, 5) |> device

    X, _ = MNIST.traindata()
    X = convert(Matrix{Float64}, reshape(X, 784, 60000))
    @time m = fit!(m, X; n_epochs=10, device=device)

    @benchmark reconstruct(m, device(X[:, 100]))
end


text:
using Test

using BEAST

A11 = BEAST.MatrixConvolution(rand(2,3,10))
A21 = BEAST.MatrixConvolution(rand(4,3,8))
A12 = BEAST.MatrixConvolution(rand(2,5,12))
A22 = BEAST.MatrixConvolution(rand(4,5,14))

A = BEAST.hvcat((2,2), A11, A12, A21, A22)

@test A[3,5,6] == A22[1,2,6]


text:
"""
Create netcdf file for one variable.
"""
function create_netcdf(filename, var, var_atts, dim_time, dim_space, time_str, id, id_desc)

    time = collect(1:dim_time)
    space = collect(1:dim_space)

    timeatts = Dict("format" => "yyyy-mm-dd HH:MM:SS")
    spaceatts = Dict("id" => id_desc)

    nccreate(filename, var, "dim_time", time, timeatts, "dim_space", space, spaceatts, atts = var_atts)
    nccreate(filename, "time_str", "dim_time", time, t=String)
    nccreate(filename, "id", "dim_space", space)

    ncwrite(time_str, filename, "time_str")
    ncwrite(id, filename, "id")

    ncclose()

    return nothing

end

"""
Get filename to one netcdf with results.
"""
function get_filename(path, variable, spatial_res, iexp)

    file = joinpath(path, "results_$(iexp)", "$(variable)_$(spatial_res).nc")

    return file

end

"""
Load results from one netcdf.
"""
function load_result(path, variable, spatial_res, iexp)

    file = get_filename(path, variable, spatial_res, iexp)

    res = ncread(file, variable)

    return res

end

"""
Load time from one netcdf.
"""
function load_time(path, variable, spatial_res)

    file = get_filename(path, variable, spatial_res, 1)

    time = ncread(file, "time_array")

    time = [DateTime(time[1,i],time[2,i],time[3,i],time[4,i]) for i in 1:size(time,2)]

    return time

end

"""
Load a time slice from a results netcdf.
"""
function load_time_slice(path, variable, spatial_res, iexp, itime = [])

    tmp = load_result(path, variable, spatial_res, 1)

    data = fill(0.0, size(tmp,1), length(iexp))

    for i in iexp

        @show i

        tmp = load_result(path, variable, spatial_res, i)

        if isempty(itime)
            data[:, i] = mean(tmp, dims=2)
        else
            data[:, i] = tmp[:, itime]
        end

    end

    return data

end

"""
Load a space slice from a results netcdf.
"""
function load_space_slice(path, variable, spatial_res, iexp, ispace = [])

    tmp = load_result(path, variable, spatial_res, 1)

    data = fill(0.0, length(iexp), size(tmp, 2))

    for i in iexp

        @show i

        tmp = load_result(path, variable, spatial_res, i)

        if isempty(ispace)
            data[i, :] = mean(tmp, 1)
        else
            data[i, :] = tmp[ispace, :]
        end

    end

    return data

end


text:
"""
    crit_pure(model::EoSModel,x0=nothing)

Calculates the critical point of a single component modelled by `model`.

Returns `(Tc, pc, Vc)` where `Tc` is the critical temperature (in K), `pc` is the critical pressure (in Pa) and `Vc` is the critical volume (in  m³)
"""
function crit_pure(model::EoSModel,x0=nothing)
    !isone(length(model)) && throw(error("$model have more than one component."))
    T̄  = T_scale(model)
    f! = (F,x) -> obj_crit(model, F, x[1]*T̄, exp10(x[2]))
    if x0 === nothing
        x0 = x0_crit_pure(model)
    end
    x0 = MVector((x0[1],x0[2]))
    #x0 = [x0[1],x0[2]]
    solver_res = Solvers.nlsolve(f!, x0)
    #print(solver_res)
    r  = Solvers.x_sol(solver_res)
    T_c = r[1]*T̄
    V_c = exp10(r[2])
    p_c = pressure(model, V_c, T_c)
    return (T_c, p_c, V_c)
end

function obj_crit(model::EoSModel, F, T_c, V_c)
    _1 = one(T_c+V_c)
    ∂²A∂V², ∂³A∂V³ = ∂²³f(model, V_c, T_c, SA[_1])
    F[1] = -∂²A∂V²
    F[2] = -∂³A∂V³
    return F
end


text:
using OrdinaryDiffEq, DiffEqGPU, ForwardDiff, Test, Unitful

function lorenz(du,u,p,t)
 @inbounds begin
     du[1] = p[1]*(u[2]-u[1])
     du[2] = u[1]*(p[2]-u[3]) - u[2]
     du[3] = u[1]*u[2] - p[3]*u[3]
 end
 nothing
end

u0 = [ForwardDiff.Dual(1f0,(1.0,0.0,0.0));ForwardDiff.Dual(0f0,(0.0,1.0,0.0));ForwardDiff.Dual(0f0,(0.0,0.0,1.0))]
tspan = (0.0f0,100.0f0)
p = (10.0f0,28.0f0,8/3f0)
prob = ODEProblem(lorenz,u0,tspan,p)
prob_func = (prob,i,repeat) -> remake(prob,p=rand(Float32,3).*p)
monteprob = EnsembleProblem(prob, prob_func = prob_func)
@time sol = solve(monteprob,Tsit5(),EnsembleGPUArray(),trajectories=10_000,saveat=1.0f0)

u0 = [1f0u"m";0u"m";0u"m"]
tspan = (0.0f0u"s",100.0f0u"s")
p = (10.0f0,28.0f0,8/3f0)
prob = ODEProblem(lorenz,u0,tspan,p)
prob_func = (prob,i,repeat) -> remake(prob,p=rand(Float32,3).*p)
monteprob = EnsembleProblem(prob, prob_func = prob_func)
@test_broken sol = solve(monteprob,Tsit5(),EnsembleGPUArray(),trajectories=10_000,saveat=1.0f0u"s")


text:
include("IO_Object.jl")
include("Debug.jl")
include("SysLED.jl")
include("GPIO.jl")
include("PWM.jl")

#List of available devices and their constructors
const DEVICES = Dict("debug" => Debug, "sysled" => SysLED, "gpio" => GPIO, "pwm" => PWM)

#Currently active devices on the BeagleBone
active_devices = Dict{String,Dict{Int32,IO_Object}}("debug" => Dict{Int32,Debug}(), "sysled" => Dict{Int32,SysLED}(),
                                                          "gpio" => Dict{Int32,GPIO}(), "pwm" => Dict{Int32,PWM}())

"""
    active_device = initdev(dev_name::String, i:Int32)
Initializes a new device of type 'dev_name' at index 'i' on the BeagleBone,
and adds it to the dict of currently active devices. Returns the initialized
device 'active_device'.
"""
function initdev(dev_name::String, i::Int32)
    #Check if the type of device is valid
    dev_constr = try
         DEVICES[dev_name]
     catch
         error("Device $dev_name does not exist")
     end
    #Check if the device index is already in use
    haskey(active_devices[dev_name],i) && error("Index $i is already in use for device $dev_name")
    #Construct a new device and add it to dict of currently active devices
    active_device = dev_constr(i)
    active_devices[dev_name][i] = active_device
    return active_device
end

"""
    closedev(dev_name::String, i::Int32)
Closes down a currently active device of type 'dev_name' at index 'i' on the BeagleBone,
and removes it from the dict of currently active devices.
"""
function closedev(dev_name::String, i::Int32)
    active_device = try
        active_devices[dev_name][i]
    catch
        error("No device of type $dev_name at index $i is currently active")
    end
    #Call the teardown method on the device, to close all file-streams and
    #unexport the device from the BeagleBone
    teardown(active_device)

    #Remove the device from the dict of active devices
    delete!(active_devices[dev_name], i)
end

"""
    dev = getdev(dev_name::String, i::Int32)
Retrieves the active device of type `dev_name` at index 'i'.
"""
function getdev(dev_name::String, i::Int32)
    dev = try
        active_devices[dev_name][i]
    catch
        error("No device of type $dev_name at index $i is currently active")
    end
    return dev
end

"""
    message = listdev()
Lists all the active devices as an insidence array for testing.
    """
function listdev()
    message = "Complete overview of active devices"
    count = zeros(length(keys(DEVICES)))
    for (index, key) in enumerate(keys(DEVICES))
        count[index] = length(active_devices[key])
    end
    return count
end

"""
    message = printdev()
Prints all the active devices and writes out specifics of a single devices.
"""
function printdev(dev_name::String, i::Int32)
    println("Complete overview of active devices")
    for (index, key) in enumerate(keys(DEVICES))
        println(" *    $(key) - $(length(active_devices[key]))")
    end
    try
        dev = active_devices[dev_name][i]
        println(to_string(dev))
    catch
        println("\nNo device of type $dev_name at index $i is currently active")
    end
end

"""
    bbparse(cmd)
Parse and execute the command `cmd`.
"""
bbparse(any) = error("Unexpected input: $any")

function bbsend(sock, vals)#, timestamps)
    serialize(sock, vals)#, (timestamps...)))
end

"""
    bbparse(l::Tuple, sock)
Parse input on the form `l=(iswrite, ndev, cmd1, cmd2, ..., cmdn)`
where if `iswrite`
    `cmdi = (devname, id, val)`
    and if not `iswrite`
    `cmdi = (devname, id)`
and send back on socket (vals, timestamps).
"""
function bbparse(l::Tuple, sock)
    iswrite = l[1]::Bool            #True if write command, false if read
    ndev = l[2]::Int32              #Number of devices/commands
    if iswrite
        for i = 1:ndev
            command = l[2+i]::Tuple
            dev = getdev(command[1], command[2])
            write!(dev, command[3])
        end
        return
    else
        #TODO fix to have at least partial type stability
        vals = Array{Any,1}(ndev)
        timestamps = Array{UInt64,1}(ndev)
        for i = 1:ndev
            command = l[2+i]::Tuple
            dev = getdev(command[1], command[2])
            vals[i] = read(dev)
            timestamps[i] = UInt64(0)#time_ns()
        end
        bbsend(sock, (vals, timestamps))
        return
    end
end

global __waiting_first_connection__ = false
"""
    run_server(port=2001; debug=false)
Run a server on `port` that listens for commands from computer
Optional debug keyword disables blinking system leds.
"""
function run_server(port=2001; debug=false)
    global __waiting_first_connection__ = true
    server = listen(port)
    @async while isopen(server)
        try
            @async while __waiting_first_connection__ && !debug
                #Blink SysLED 2 when waiting for first connection to signal availability
                led = initdev("sysled",Int32(2))
                write!(led, "1")
                sleep(0.4)
                write!(led, "0")
                sleep(0.2)
                write!(led, "1")
                sleep(0.4)
                write!(led, "0")
                sleep(0.8)
                closedev("sysled", Int32(2))
            end
            sock = accept(server)
            __waiting_first_connection__ = false
            @async while isopen(sock)
                try
                    l = deserialize(sock);
                    bbparse(l, sock)
                catch err
                    if !isopen(sock) && (isa(err, Base.EOFError) || isa(err, Base.UVError))
                        println("Connection to server closed")
                    else
                        println("error: $(typeof(err))")
                        throw(err)
                    end
                end
            end
        catch err
            if isa(err,Base.UVError) && err.prefix == "accept"
                println("Server closed successfully")
            else
                throw(err)
            end
        end
    end
    return server
end


text:
mutable struct Flamingo{T} <: ContactModel
    dim::Dimensions

    g::T
    μ_world::T
	μ_joint::T

    # torso
    l_torso::T
    d_torso::T
    m_torso::T
    J_torso::T

    # leg 1
        # thigh
    l_thigh1::T
    d_thigh1::T
    m_thigh1::T
    J_thigh1::T

        # calf
    l_calf1::T
    d_calf1::T
    m_calf1::T
    J_calf1::T

		# foot
	l_foot1::T # toe length
	d_foot1::T # heel length
	m_foot1::T
	J_foot1::T

    # leg 2
        # thigh
    l_thigh2::T
    d_thigh2::T
    m_thigh2::T
    J_thigh2::T

        # calf
    l_calf2::T
    d_calf2::T
    m_calf2::T
    J_calf2::T

		# foot
	l_foot2::T # toe length
	d_foot2::T # heel length
	m_foot2::T
	J_foot2::T

	# fast methods
	base
	dyn

	joint_friction
end

function kinematics_1(model::Flamingo, q; body = :torso, mode = :ee)
	x = q[1]
	z = q[2]

	if body == :torso
		l = model.l_torso
		d = model.d_torso
		θ = q[3]
		if mode == :ee
			return [x - l * sin(θ); z + l * cos(θ)]
		elseif mode == :com
			return [x - d * sin(θ); z + d * cos(θ)]
		end
	elseif body == :thigh_1
		l = model.l_thigh1
		d = model.d_thigh1
		θ = q[4]
	elseif body == :thigh_2
		l = model.l_thigh2
		d = model.d_thigh2
		θ = q[6]
	else
		@error "incorrect body specification"
	end

	if mode == :ee
		return [x + l * sin(θ); z - l * cos(θ)]
	elseif mode == :com
		return [x + d * sin(θ); z - d * cos(θ)]
	else
		@error "incorrect mode specification"
	end
end

function jacobian_1(model::Flamingo, q; body = :torso, mode = :ee)
	jac = zeros(eltype(q), 2, model.dim.q)
	jac[1, 1] = 1.0
	jac[2, 2] = 1.0
	if body == :torso
		r = mode == :ee ? model.l_torso : model.d_torso
		θ = q[3]
		jac[1, 3] = -r * cos(θ)
		jac[2, 3] = -r * sin(θ)
	elseif body == :thigh_1
		r = mode == :ee ? model.l_thigh1 : model.d_thigh1
		θ = q[4]
		jac[1, 4] = r * cos(θ)
		jac[2, 4] = r * sin(θ)
	elseif body == :thigh_2
		r = mode == :ee ? model.l_thigh2 : model.d_thigh2
		θ = q[6]
		jac[1, 6] = r * cos(θ)
		jac[2, 6] = r * sin(θ)
	else
		@error "incorrect body specification"
	end

	return jac
end

function kinematics_2(model::Flamingo, q; body = :calf_1, mode = :ee)

	if body == :calf_1
		p = kinematics_1(model, q, body = :thigh_1, mode = :ee)

		θb = q[5]

		lb = model.l_calf1
		db = model.d_calf1
	elseif body == :calf_2
		p = kinematics_1(model, q, body = :thigh_2, mode = :ee)

		θb = q[7]

		lb = model.l_calf2
		db = model.d_calf2
	else
		@error "incorrect body specification"
	end

	if mode == :ee
		return p + [lb * sin(θb); -1.0 * lb * cos(θb)]
	elseif mode == :com
		return p + [db * sin(θb); -1.0 * db * cos(θb)]
	else
		@error "incorrect mode specification"
	end
end

function jacobian_2(model::Flamingo, q; body = :calf_1, mode = :ee)

	if body == :calf_1
		jac = jacobian_1(model, q, body = :thigh_1, mode = :ee)

		θb = q[5]

		r = mode == :ee ? model.l_calf1 : model.d_calf1

		jac[1, 5] += r * cos(θb)
		jac[2, 5] += r * sin(θb)
	elseif body == :calf_2
		jac = jacobian_1(model, q, body = :thigh_2, mode = :ee)

		θb = q[7]

		r = mode == :ee ? model.l_calf2 : model.d_calf2

		jac[1, 7] += r * cos(θb)
		jac[2, 7] += r * sin(θb)
	else
		@error "incorrect body specification"
	end

	return jac
end

function kinematics_3(model::Flamingo, q; body = :foot_1, mode = :ee)

	if body == :foot_1
		p = kinematics_2(model, q, body = :calf_1, mode = :ee)

		θb = q[8]

		lb = model.l_foot1
		db = model.d_foot1
		cb = 0.5 * (model.l_foot1 - model.d_foot1)
	elseif body == :foot_2
		p = kinematics_2(model, q, body = :calf_2, mode = :ee)

		θb = q[9]

		lb = model.l_foot2
		db = model.d_foot2
		cb = 0.5 * (model.l_foot2 - model.d_foot2)
	else
		@error "incorrect body specification"
	end

	if mode == :toe
		return p + [lb * sin(θb); -1.0 * lb * cos(θb)]
	elseif mode == :heel
		return p + [-db * sin(θb); 1.0 * db * cos(θb)]
	elseif mode == :com
		return p + [cb * sin(θb); -1.0 * cb * cos(θb)]
	else
		@error "incorrect mode specification"
	end
end

function jacobian_3(model::Flamingo, q; body = :foot_1, mode = :ee)

	if body == :foot_1
		jac = jacobian_2(model, q, body = :calf_1, mode = :ee)

		θb = q[8]

		if mode == :toe
			r = model.l_foot1
		elseif mode == :heel
			r = -1.0 * model.d_foot1
		elseif mode == :com
			r = 0.5 * (model.l_foot1 - model.d_foot1)
		else
			@error "incorrect mode specification"
		end

		jac[1, 8] += r * cos(θb)
		jac[2, 8] += r * sin(θb)

	elseif body == :foot_2
		jac = jacobian_2(model, q, body = :calf_2, mode = :ee)

		θb = q[9]

		if mode == :toe
			r = model.l_foot2
		elseif mode == :heel
			r = -1.0 * model.d_foot2
		elseif mode == :com
			r = 0.5 * (model.l_foot2 - model.d_foot2)
		else
			@error "incorrect mode specification"
		end
		jac[1, 9] += r * cos(θb)
		jac[2, 9] += r * sin(θb)

	else
		@error "incorrect body specification"
	end

	return jac
end

# Lagrangian

function lagrangian(model::Flamingo, q, q̇)
	L = 0.0

	# torso
	p_torso = kinematics_1(model, q, body = :torso, mode = :com)
	J_torso = jacobian_1(model, q, body = :torso, mode = :com)
	v_torso = J_torso * q̇

	L += 0.5 * model.m_torso * transpose(v_torso) * v_torso
	L += 0.5 * model.J_torso * q̇[3]^2.0
	L -= model.m_torso * model.g * p_torso[2]

	# thigh 1
	p_thigh_1 = kinematics_1(model, q, body = :thigh_1, mode = :com)
	J_thigh_1 = jacobian_1(model, q, body = :thigh_1, mode = :com)
	v_thigh_1 = J_thigh_1 * q̇

	L += 0.5 * model.m_thigh1 * transpose(v_thigh_1) * v_thigh_1
	L += 0.5 * model.J_thigh1 * q̇[4]^2.0
	L -= model.m_thigh1 * model.g * p_thigh_1[2]

	# leg 1
	p_calf_1 = kinematics_2(model, q, body = :calf_1, mode = :com)
	J_calf_1 = jacobian_2(model, q, body = :calf_1, mode = :com)
	v_calf_1 = J_calf_1 * q̇

	L += 0.5 * model.m_calf1 * transpose(v_calf_1) * v_calf_1
	L += 0.5 * model.J_calf1 * q̇[5]^2.0
	L -= model.m_calf1 * model.g * p_calf_1[2]

	# foot 1
	p_foot_1 = kinematics_3(model, q, body = :foot_1, mode = :com)
	J_foot_1 = jacobian_3(model, q, body = :foot_1, mode = :com)
	v_foot_1 = J_foot_1 * q̇

	L += 0.5 * model.m_foot1 * transpose(v_foot_1) * v_foot_1
	L += 0.5 * model.J_foot1 * q̇[8]^2.0
	L -= model.m_foot1 * model.g * p_foot_1[2]

	# thigh 2
	p_thigh_2 = kinematics_1(model, q, body = :thigh_2, mode = :com)
	J_thigh_2 = jacobian_1(model, q, body = :thigh_2, mode = :com)
	v_thigh_2 = J_thigh_2 * q̇

	L += 0.5 * model.m_thigh2 * transpose(v_thigh_2) * v_thigh_2
	L += 0.5 * model.J_thigh2 * q̇[6]^2.0
	L -= model.m_thigh2 * model.g * p_thigh_2[2]

	# leg 2
	p_calf_2 = kinematics_2(model, q, body = :calf_2, mode = :com)
	J_calf_2 = jacobian_2(model, q, body = :calf_2, mode = :com)
	v_calf_2 = J_calf_2 * q̇

	L += 0.5 * model.m_calf2 * transpose(v_calf_2) * v_calf_2
	L += 0.5 * model.J_calf2 * q̇[7]^2.0
	L -= model.m_calf2 * model.g * p_calf_2[2]

	# foot 2
	p_foot_2 = kinematics_3(model, q, body = :foot_2, mode = :com)
	J_foot_2 = jacobian_3(model, q, body = :foot_2, mode = :com)
	v_foot_2 = J_foot_2 * q̇

	L += 0.5 * model.m_foot2 * transpose(v_foot_2) * v_foot_2
	L += 0.5 * model.J_foot2 * q̇[9]^2.0
	L -= model.m_foot2 * model.g * p_foot_2[2]

	return L
end

function _dLdq(model::Flamingo, q, q̇)
	Lq(x) = lagrangian(model, x, q̇)
	ForwardDiff.gradient(Lq, q)
end

function _dLdq̇(model::Flamingo, q, q̇)
	Lq̇(x) = lagrangian(model, q, x)
	ForwardDiff.gradient(Lq̇, q̇)
end

function _C_func(model::Flamingo, q, q̇)
	tmp_q(z) = _dLdq̇(model, z, q̇)
	tmp_q̇(z) = _dLdq̇(model, q, z)

	ForwardDiff.jacobian(tmp_q, q) * q̇ - _dLdq(model, q, q̇)
end

function kinematics(model::Flamingo, q)
	p_toe_1 = kinematics_3(model, q, body = :foot_1, mode = :toe)
	p_heel_1 = kinematics_3(model, q, body = :foot_1, mode = :heel)
	p_toe_2 = kinematics_3(model, q, body = :foot_2, mode = :toe)
	p_heel_2 = kinematics_3(model, q, body = :foot_2, mode = :heel)

	SVector{8}([p_toe_1; p_heel_1; p_toe_2; p_heel_2])
end

# Methods
function M_func(model::Flamingo, q)
	M = Diagonal([0.0, 0.0,
		model.J_torso,
		model.J_thigh1, model.J_calf1,
		model.J_thigh2, model.J_calf2,
		model.J_foot1, model.J_foot2])

	# torso
	J_torso = jacobian_1(model, q, body = :torso, mode = :com)
	M += model.m_torso * transpose(J_torso) * J_torso

	# thigh 1
	J_thigh_1 = jacobian_1(model, q, body = :thigh_1, mode = :com)
	M += model.m_thigh1 * transpose(J_thigh_1) * J_thigh_1

	# leg 1
	J_calf_1 = jacobian_2(model, q, body = :calf_1, mode = :com)
	M += model.m_calf1 * transpose(J_calf_1) * J_calf_1

	# foot 1
	J_foot_1 = jacobian_3(model, q, body = :foot_1, mode = :com)
	M += model.m_foot1 * transpose(J_foot_1) * J_foot_1

	# thigh 2
	J_thigh_2 = jacobian_1(model, q, body = :thigh_2, mode = :com)
	M += model.m_thigh2 * transpose(J_thigh_2) * J_thigh_2

	# leg 2
	J_calf_2 = jacobian_2(model, q, body = :calf_2, mode = :com)
	M += model.m_calf2 * transpose(J_calf_2) * J_calf_2

	# foot 2
	J_foot_2 = jacobian_3(model, q, body = :foot_2, mode = :com)
	M += model.m_foot2 * transpose(J_foot_2) * J_foot_2

	return M
end

function ϕ_func(model::Flamingo, env::Environment, q)
	p_toe_1 = kinematics_3(model, q, body = :foot_1, mode = :toe)
	p_heel_1 = kinematics_3(model, q, body = :foot_1, mode = :heel)
	p_toe_2 = kinematics_3(model, q, body = :foot_2, mode = :toe)
	p_heel_2 = kinematics_3(model, q, body = :foot_2, mode = :heel)

	SVector{4}([p_toe_1[2] - env.surf(p_toe_1[1:1]),
				p_heel_1[2] - env.surf(p_heel_1[1:1]),
				p_toe_2[2] - env.surf(p_toe_2[1:1]),
				p_heel_2[2] - env.surf(p_heel_2[1:1])])
end

function B_func(model::Flamingo, q)
	@SMatrix [0.0  0.0 -1.0  1.0  0.0  0.0  0.0  0.0  0.0;
			  0.0  0.0  0.0 -1.0  1.0  0.0  0.0  0.0  0.0;
			  0.0  0.0 -1.0  0.0  0.0  1.0  0.0  0.0  0.0;
			  0.0  0.0  0.0  0.0  0.0 -1.0  1.0  0.0  0.0;
			  0.0  0.0  0.0  0.0 -1.0  0.0  0.0  1.0  0.0;
			  0.0  0.0  0.0  0.0  0.0  0.0 -1.0  0.0  1.0]
end

function A_func(model::Flamingo, q)
	@SMatrix [1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
			  0.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0]
end

function J_func(model::Flamingo, env::Environment, q)
	J_toe_1 = jacobian_3(model, q, body = :foot_1, mode = :toe)
	J_heel_1 = jacobian_3(model, q, body = :foot_1, mode = :heel)
	J_toe_2 = jacobian_3(model, q, body = :foot_2, mode = :toe)
	J_heel_2 = jacobian_3(model, q, body = :foot_2, mode = :heel)

	return [J_toe_1;
			J_heel_1;
			J_toe_2;
			J_heel_2]
end

function contact_forces(model::Flamingo, env::Environment{<:World, LinearizedCone}, γ1, b1, q2)
	m = friction_mapping(env)

	SVector{8}([transpose(rotation(env, k[1:1])) * [m * b1[1:2]; γ1[1]];
				transpose(rotation(env, k[3:3])) * [m * b1[3:4]; γ1[2]];
				transpose(rotation(env, k[5:5])) * [m * b1[5:6]; γ1[3]];
				transpose(rotation(env, k[7:7])) * [m * b1[7:8]; γ1[4]]])
end

function velocity_stack(model::Flamingo, env::Environment{<:World, LinearizedCone}, q1, q2, h)
	v = J_func(model, env, q2) * (q2 - q1) / h[1]

	v1_surf = rotation(env, k[1:1]) * v[1:2]
	v2_surf = rotation(env, k[3:3]) * v[3:4]
	v3_surf = rotation(env, k[5:5]) * v[5:6]
	v4_surf = rotation(env, k[7:7]) * v[7:8]

	SVector{8}([transpose(friction_mapping(model.env)) * v1_surf[1];
				transpose(friction_mapping(model.env)) * v2_surf[1];
				transpose(friction_mapping(model.env)) * v3_surf[1];
				transpose(friction_mapping(model.env)) * v4_surf[1]])
end

# Dimensions
nq = 2 + 5 + 2            # configuration dimension
nu = 6                    # control dimension
nc = 4                    # number of contact points
nw = 2                    # disturbance dimension
nquat = 0

# World parameters
μ_world = 0.9      # coefficient of friction
μ_joint = 0.0
gravity = 9.81           # gravity

# Model parameters
m_torso = 12.0
m_thigh = 0.4598
m_calf = 0.306
m_foot = 0.3466

l_torso = 0.385
l_thigh = 0.42
l_calf = 0.45
l_foot = 0.1725 # distance ankle-toe

d_torso = 0.20
d_thigh = l_thigh/2
d_calf = l_calf/2
d_foot = 0.0525 # distance ankle-heel

J_torso = 0.10
J_thigh = 0.01256
J_calf = 0.00952
J_foot = 0.0015

flamingo = Flamingo(Dimensions(nq, nu, nw, nc, nquat),
			  gravity, μ_world, μ_joint,
			  l_torso, d_torso, m_torso, J_torso,
			  l_thigh, d_thigh, m_thigh, J_thigh,
			  l_calf, d_calf, m_calf, J_calf,
			  l_foot, d_foot, m_foot, J_foot,
			  l_thigh, d_thigh, m_thigh, J_thigh,
			  l_calf, d_calf, m_calf, J_calf,
			  l_foot, d_foot, m_foot, J_foot,
			  BaseMethods(), DynamicsMethods(),
			  SVector{nq}([zeros(3); 0.0 * μ_joint * ones(nq - 3)]))


text:
# TEMPLATE STEP 1
# 1. Create an environment for CpelNano in your home directory (only once)
# 2. Install CpelNano in the environment using (only once)
#   add https://github.com/jordiabante/CpelNano.jl.git#TowardsContinuousY
# 3. Copy the content of this file in a file in your home directory
# 4. Replace the environment path for the correct one where the arrow points
# 5. Define the proper paths

# Load dependencies
using Distributed
@everywhere using Pkg
@everywhere Pkg.activate("/home-net/home-2/jabante1@jhu.edu/julia_env/CpelNano/") # <=
@everywhere using CpelNano

# Define paths (you can pass arguments and access them through the array ARGS, e.g. ARGS[1])
indir = ""
nano = "$(indir)/"
fasta = "$(indir)/"
outdir = ""
out_prefix = ""

# Configure CpelNano
min_cov = 5.0               # Minimum average depth
max_em_init = 10            # Max number of inits in Cpel EM
max_em_iters = 100          # Max number of iters in Cpel EM

# Configure CpelNano
config = CpelNano.CpelNanoConfig()
config.out_dir = outdir
config.min_cov = min_cov
config.out_prefix = out_prefix
config.max_em_init = max_em_init
config.max_em_iters = max_em_iters

# Analyze nanopore data run
CpelNano.analyze_nano(nano,fasta,config)


text:
export solve_SSODE, solve_SSODE_threads, solve_time, solve

const Default_SSMETHOD = AutoTsit5(Rosenbrock23())

function solveSS(func, u, p; method= Default_SSMETHOD)
    #= ODE solver =#
    prob = SteadyStateProblem(func, u, p)
    return DifferentialEquations.solve(prob, method)
end

"""
    solve(ode_func::DEsteady, us; ensemble_method=EnsembleThreads())
    solve(ode_func::DEsteady)
    solve(ode_func::ODEtime)
    solve(ode_func::ODEtime, us; ensemble_method=EnsembleThreads())

Extanded solver for steady state and time-series. The ['DEmeta'](@ref) type provides the information of differential equations ('func'), initial variables ('u') and parameters ('p'). With [`DEsteady`](@ref) and ['ODEtime'](@ref)

Arguements
----------
- `func`: DE function
- `us`: Vector of vectors of initial variables
- `p`: parameter constant

See also
--------
[`DEsteady`](@ref), ['ODEtime'](@ref)
"""
function DifferentialEquations.solve(ode_func::DEsteady, us; ensemble_method=EnsembleThreads())

    function prob_func(prob,i,repeat)
        ode_new = ode_func(us[i])
        remake(prob,u0 = ode_new.u0)
    end

    @unpack func, u0, method, p = ode_func

    prob = SteadyStateProblem(func, u0, p)

    ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)
    sim = solve(ensemble_prob,method,ensemble_method,trajectories=length(us))

    return sim
end

function DifferentialEquations.solve(ode_func::DEsteady)
    @unpack func, u0, method, p = ode_func
    return solveSS(func, u0, p; method= method)
end

function DifferentialEquations.solve(ode_func::ODEtime)
    @unpack func, u0, tspan, p, method = ode_func
    prob = ODEProblem(func, u0, tspan, p)
    return solve(prob; method=method)
end

function DifferentialEquations.solve(ode_func::ODEtime, us; ensemble_method=EnsembleThreads())

    function prob_func(prob,i,repeat)
        ode_new = ode_func(us[i])
        remake(prob,u0 = ode_new.u0)
    end

    @unpack func, u0, method, p, tspan = ode_func

    prob = ODEProblem(func, u0, tspan, p)

    ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)
    sim = solve(ensemble_prob,method,ensemble_method,trajectories=length(us))

    return sim
end

function solve_SSODE_threads(func, us, p ; method=Default_SSMETHOD)
    function prob_func(prob,i,repeat)
        remake(prob,u0 = us[i])
    end

    prob = SteadyStateProblem(func, us[1], p)

    ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)
    sim = solve(ensemble_prob,method,EnsembleThreads(),trajectories=length(us))

    return sim
end

function solve_SSODE_threads(ode_func::DEmeta)
    sim = solve_SSODE_threads(ode_func.func, ode_func.u0, ode_func.p; method=ode_func.method)

    return sim
end

function solve_SSODE_threads(ode_func::DEmeta, us)

    function prob_func(prob,i,repeat)
        ode_new = ode_func(us[i])
        remake(prob,u0 = ode_new.u0)
    end

    prob = SteadyStateProblem(ode_func.func, ode_func.u0, ode_func.p)

    ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)
    sim = solve(ensemble_prob,ode_func.method,EnsembleThreads(),trajectories=length(us))

    return sim
end


text:
m = parse(Int,ARGS[1])
sgn = parse(Int,ARGS[2])
numChanges = parse(Int,ARGS[3])
lbInd = parse(Int,ARGS[4])
rep = parse(Int, ARGS[5])

scratch_dir = ARGS[6]   # e.g. "/scratch/midway2/byolkim/exper4"

if isfile("$(scratch_dir)/res_$(m)_$(sgn)_$(numChanges)_$(lbInd)_$(rep).jld")
    println("the file already exists!")
    exit()
end

using KLIEPInference
using ProximalBase, CoordinateDescent
using LinearAlgebra, SparseArrays, Statistics, Random
using Distributions, StatsBase, JLD

println("importing parameters from params_exp4_$(m)_$(sgn)_$(numChanges)_$(lbInd).jld")
file = jldopen("graphs/params_exp4_$(m)_$(sgn)_$(numChanges)_$(lbInd).jld", "r")
γx = read(file, "γx")
γy = read(file, "γy")
close(file)

p = length(γx)
nx = 500
ny = 500

println("generating samples")
Random.seed!(123 + rep)
spl = IsingSampler(γx; thin=2000)
X = rand(spl, nx)
spl = IsingSampler(γy; thin=2000)
Y = rand(spl, ny)

Ψx = Ψising(X)
Ψy = Ψising(Y)

println("step 1")
λ1 = 1.01 * quantile(Normal(), 1. - 0.05 / p)
θ = spKLIEP(Ψx, Ψy, λ1, CD_KLIEP(); loadings=true)

println("step 2")
λ2 = sqrt(2. * log(p) / ny)
H = KLIEP_Hessian(θ, Ψy)
Hinv = Vector{SparseIterate{Float64}}(undef, p)
for k = 1:p
    ω = Hinv_row(H, k, λ2)

    supp = KLIEPInference._find_supp(k, ω)
    h = view(H, supp, supp)
    δ = (supp .== k)
    ω[supp] = h\δ

    Hinv[k] = ω
end

println("step 3 + bootstrap")
boot1, boot2 = boot_SparKLIE(Ψx, Ψy, θ, Hinv; bootSamples=1000)

println("saving results to $(scratch_dir)/res_$(m)_$(sgn)_$(numChanges)_$(lbInd)_$(rep).jld")
@save "$(scratch_dir)/res_$(m)_$(sgn)_$(numChanges)_$(lbInd)_$(rep).jld" boot1 boot2

println("done!")


text:
"""
    fittedEVA{T<:EVA}

Abstract type containing the fitted extreme value model types.

- BayesianEVA
- MaximumLikelihoodEVA
- pwmEVA

"""
abstract type fittedEVA{T<:EVA} end

Base.Broadcast.broadcastable(obj::fittedEVA) = Ref(obj)

"""
    Base.show(io::IO, obj::fittedEVA)

Override of the show function for the objects of type fittedEVA.

"""
function Base.show(io::IO, obj::fittedEVA)

    showfittedEVA(io, obj)

end

"""
    cint(..., confidencelevel::Real=.95)

Compute confidence interval or credible interval in the case of Bayesian estimation.

The function can be applied on any [`fittedEVA`](@ref) subtype to obtain a
confidence interval on the model parameters. It can also be applied on
[`ReturnLevel`](@ref) type to obtain a confidence interval on the return level.

# Implementation

The method used for computing the interval depends on the estimation method. In
the case of maximum likelihood estimation, the confidence intervals are computed
using the Wald approximation based on the approximate parameter estimates
covariance matrix. In the case of Bayesian estimation, the return interval is
the highest posterior density estimate based on the MCMC sample. In the case of
probability weighted moment estimation, the intervals are computed using a
boostrap procedure.

"""
function cint end

include("returnlevel.jl")
include(joinpath("fittedeva", "bayesianeva.jl"))
include(joinpath("fittedeva", "maximumlikelihoodeva.jl"))
include(joinpath("fittedeva", "pwmeva.jl"))


text:
# rdiff.jl - differentiate an expression or a function w.r.t. its arguments
#
# An approach taken in this file falls into a category of hybric method of
# computer-aided differentiation. General architecture closely follows
# reverse-mode automatic differentiation including building a "tape" and
# propagating derivative from an output to an input variables, but unlike AD
# it allows to output symbolic expressions for calculating partial derivatives
# in question.
#
# Advantages of this approach include efficient computation of derivatives
# (which may not always be a case for symbolic differentiation), outputting
# symbolic expression which only needs to be computed once at runtime
# (unlike full reverse-mode AD that requires 2 passes) and may be used to
# generate code for other systems like GPU (which is not possible with AD at all).
# The main disadvantage compared to AD is inability to handle conditions and
# loops since they create discontinuity and thus are very hard to analyze
# and produce symbolic expression.
#
#
# Architecture
# ------------
#
# Just as reverse-mode AD, we build derivative using 2 passes -
# forward and reverse.
#
# 1. During forward pass we parse expression and build a graph (`ExGraph`) -
#    "linearized" version of expression. Each node in this graph represents:
#
#     * function call: `ExNode{:call}`
#     * broadcasting: `ExNode{:bcast}`
#     * assignment: `ExNode{:(=)}`
#     * input variable: `ExNode{:input}`
#     * constant value: `ExNode{:constant}`
#     * tuple: `ExNode{:tuple}`
#
#    Nodes are created and put onto a "tape" (list of nodes) using `parse!()`
#    function in a topological order, so no node may refer to a dependency
#    that isn't defined yet. After the graph is built, all nodes are evaluated
#    using input variables and constants to provide "example values" for each
#    node (used for type inference and other stuff).
#
# 2. Reverse pass starts with an empty dict of "adjoints" - derivatives of
#    an output variable w.r.t. to current variable - and populating it
#    from the end to the beginning (or, more precisely, from dependent variables
#    to their dependencies).
#
#    The backbone of this process is `rev_step!` function, and the most
#    intersting method is the one that handles function call. Given a node
#    and a dict of all dependent adjoints, this method does the following
#    for each of its dependencies (here `z` stands for an output variable,
#    `y` - for current variable and `x` - for current dependency of `y`):
#
#     1) finds a differentiation rule for `dy/dx`, i.e. derivative of current
#        node w.r.t. current dependency
#     2) symbolically multiplies it by `dz/dy`, i.e. by derivative of output
#        variable w.r.t. current node, to obtain `dz/dx`, i.e. derivative of
#        output w.r.t. curent dependency
#     3) adds this new derivative to the adjoint dict
#
#
# Example
# -------
#
# Here's an example of this process (and some tips on how to debug it).
#
# Suppose we have expression:
#
#     ex = :(z = x1*x2 + sin(x1))
#
# We can build an `ExGraph` and perform forward pass on it like this (which is
# not necessary for high-level usage, but helpful for debugging):
#
#     g = ExGraph(ex; x1=1, x2=1)  # give it an example of possible values
#     forward_pass(g)
#
# After this our graph looks like:
#
#   ExGraph
#     ExNode{:input}(:x1,:($(Expr(:input, :x1))),1)
#     ExNode{:input}(:x2,:($(Expr(:input, :x2))),1)
#     ExNode{:call}(:tmp1,:(x1 * x2),1)
#     ExNode{:call}(:tmp2,:(sin(x1)),0.8414709848078965)
#     ExNode{:call}(:tmp3,:(tmp1 + tmp2),1.8414709848078965)
#     ExNode{:(=)}(:z,:(z = tmp3),1.8414709848078965)
#
# Let's take a node with a first call for example:
#
#   ExNode{:call}(:tmp1,:(x1 * x2),1)
#
# Here `ExNode` is parametrized by :call symbol to enable Julia's fancy method
# dispatching, :tmp1 is a name of a new variable which is a product of variables
# :x1 and :x2 and has example value of 1.
#
# Now let's run reverse pass to obtain derivatives of `z` w.r.t. all other vars:
#
#     adj = reverse_pass(g, :z)
#
# which results in a dict like this:
#
#     Dict{Symbol,Any} with 6 entries:
#        :tmp2 => 1.0
#        :x1   => :(x2 + cos(x1))
#        :z    => 1.0
#        :tmp1 => 1.0
#        :x2   => :x1
#        :tmp3 => 1.0
#
# This means that:
#
#     dz/dz == adj[:z] == 1.0
#     dz/dtmp3 == adj[:tmp3] == 1.0
#     dz/dx1 == adj[:x1] == :(x2 + cos(x1))
#     ...
#
# To see how it works, consider finding derivative `dz/dx1` at intermediate
# node :tmp1. Note, that in our example `tmp1 = x1 * x2`.
#
# 1) by the time of computing this derivative we already know that
#    `dz/dtmp1 == 1.0`
# 2) from primitive derivative rules for product of numbers we also infer that
#    `dtmp1/dx1 == x2`
# 3) using chain rule we obtain (part of*) `dz/dx1` as a symbolic product
#    `dz/dmp1 * dtmp1/dx1 == 1.0 * x2 == x2`.
#
# I say "part of" because :z depends on :x1 not only through :tmp1, but also
# through :tmp2, for which derivative `dz/dx1` turns to be `cos(x1)`.
# To combine 2 these "parts" we simply add them up and obtain final result:
#
#     dz/dx1 == x2 + cos(x1)
#

## utils

function extend_deriv!(dg::ExGraph, dzdx_v::Symbol, dzdx::Any)
    subderivs = find_related(dg, dzdx_v)
    pos = indexof(dg, dzdx_v)
    if isempty(subderivs)
        # first split
        dzdx_ex_1 = getexpr(dg[dzdx_v])
        dzdx_ex_2 = dzdx
        dzdx_v_1 = Symbol("$(dzdx_v)__1")
        dzdx_v_2 = Symbol("$(dzdx_v)__2")
        sub_dg = typeof(dg)()
        parse!(sub_dg, :($dzdx_v_1 = $dzdx_ex_1))
        parse!(sub_dg, :($dzdx_v_2 = $dzdx_ex_2))
        parse!(sub_dg, :($dzdx_v = $dzdx_v_1 .+ $dzdx_v_2))
        sub_dg = fuse_assigned(sub_dg)
        new_nodes = sub_dg.tape
    else
        # dg already contains subderivatives for dzdx_v
        last_idx = parse(Int, split(subderivs[end] |> String, "__")[end])
        dzdx_v_last = Symbol("$(dzdx_v)__$(last_idx + 1)")
        prev_dzdx_ex = getexpr(dg[dzdx_v])
        sub_dg = typeof(dg)()
        parse!(sub_dg, :($dzdx_v_last = $dzdx))
        parse!(sub_dg, :($dzdx_v = $prev_dzdx_ex .+ $dzdx_v_last))
        sub_dg = fuse_assigned(sub_dg)
        new_nodes = sub_dg.tape
    end
    delete!(dg, pos)
    insert!(dg, pos, new_nodes)
    return dg
end

function example_val{T}(::Type{T})
    if T <: Number
        return one(T)
    elseif T <: Array
        return ones(eltype(T), [1 for i=1:ndims(T)]...)
    else
        error("Don't know how to create an example value for type $T")
    end
end

## forward pass

"""Forward pass of differentiation"""
function forward_pass!(g::AbstractExGraph)
    evaluate!(g, varname(g.tape[end]))
    # propagate_size!(g)
    return g
end

## reverse step

"""
Perform one step of reverse pass. Add derivatives of output variable w.r.t.
node's dependenices to adjoint dictionary.
"""
function rev_step!(g::ExGraph, dg::ExGraph, nd::ExNode{:(=)})
    y = varname(nd)
    x = dependencies(nd)[1]
    adj[x] = adj[y]
end

function rev_step!(g::ExGraph, dg::ExGraph, nd::ExNode{:constant})
    # do nothing
end

function rev_step!(g::ExGraph, dg::ExGraph, nd::ExNode{:input})
    # do nothing
end

function rev_step!(g::ExGraph, dg::ExGraph, nd::ExNode{:call})
    y = varname(nd)
    z = g.ctx[:z_var]
    dzdy_v = deriv_name(z, y)
    cg = cat(g, dg)
    types = [typeof(getvalue(g[x])) for x in dependencies(nd)]
    for (i, x) in enumerate(dependencies(nd))
        xnd = g[x]
        if isa(xnd, ExNode{:constant})
            # don't clog dg with unnesessary derivs
            continue
        end
        dydx = derivative(getexpr(nd), types, i, mod=g.ctx[:mod])
        dzdx = dzdy_v ⊗ dydx
        dzdx = expand_const(cg, dzdx) |> simplify
        dzdx_v = deriv_name(z, x)
        if haskey(dg, dzdx_v)
            extend_deriv!(dg, dzdx_v, dzdx)
        else
            parse!(dg, :($dzdx_v = $dzdx))
        end
    end
end

function rev_step!(g::ExGraph, dg::ExGraph, nd::ExNode{:bcast})
    y = varname(nd)
    z = g.ctx[:z_var]
    dzdy_v = deriv_name(z, y)
    cg = cat(g, dg)
    types = [typeof(getvalue(g[x])) for x in dependencies(nd)]
    for (i, x) in enumerate(dependencies(nd))
        xnd = g[x]
        if isa(xnd, ExNode{:constant})
            # don't clog dg with unnesessary derivs
            continue
        end
        dydx = derivative(getexpr(nd) |> bcast_to_call, types, i, mod=g.ctx[:mod])
        dzdx = dzdy_v ⊗ dydx
        dzdx = expand_const(cg, dzdx) |> simplify
        dzdx_v = deriv_name(z, x)
        if haskey(dg, dzdx_v)
            extend_deriv!(dg, dzdx_v, dzdx)
        else
            parse!(dg, :($dzdx_v = $dzdx))
        end
    end
end

function reverse_pass!(g::ExGraph)
    z = varname(g.tape[end])
    g.ctx[:z_var] = z
    dzdz_var = deriv_name(z, z)
    dg = ExGraph(:($dzdz_var = 1.0))
    for nd in reverse(g.tape)
        rev_step!(g, dg, nd)
    end
    outvars = [deriv_name(z, varname(nd)) for nd in g.tape if isa(nd, ExNode{:input})]
    return fuse_assigned(dg; outvars=outvars)
end

function _xdiff(g::AbstractExGraph)
    forward_pass!(g)
    dg = reverse_pass!(g)
    return dg
end

function proper_graph(ex; ctx=Dict(), inputs...)
    # determine format: if any of arguments is a tensor, use Einstein notation
    # otherwise use simple scalar differentiation
    if isindexed(ex)
        g = EinGraph(ex; ctx=ctx, inputs...)
    elseif any(x -> !isa(x[2], Number), inputs)
        iex = to_einstein(ex; ctx=ctx, inputs...)
        g = EinGraph(iex; ctx=ctx, inputs...)
    else
        g = ExGraph(ex; ctx=ctx, inputs...)
    end
    return g
end

"""
xdiff(ex::Expr; ctx=Dict(), inputs...)

Differentiate expression `ex` w.r.t. variables `inputs`. `inputs` should be a list
of key-value pairs with keys representing variables in expression and values
representing 'example values' (used e.g. for type inference). Returns an expression
that calculates original value and derivatives of all inputs. Example:

    xdiff(:(x^n); x=1, n=1)
    # quote
    #   tmp704 = 1
    #   tmp708 = log(x)
    #   tmp705 = n - tmp704
    #   tmp706 = x .^ tmp705
    #   tmp702 = x ^ n
    #   dtmp702!dx = n * tmp706
    #   tmp709 = x .^ n
    #   dtmp702!dn = tmp708 * tmp709
    #   tmp711 = (tmp702, dtmp702!dx, dtmp702!dn)
    # end

"""
function xdiff(ex::Expr; ctx=Dict(), inputs...)
    ctx = to_context(ctx)
    g = proper_graph(ex; ctx=ctx, inputs...)
    dg = _xdiff(g)
    rg = cat(g, dg)
    outvars = unshift!([deriv_name(g.ctx[:z_var], var) for (var, _) in inputs], varname(g[end]))
    push!(rg, :tuple, Espresso.genname(), Expr(:tuple, outvars...))
    rg = topsort(rg)
    infer_deriv_size!(rg)  # need to know size to evaluate things like `dz!dx[i] = 1.0`
    evaluate!(rg)
    codegen = @get(ctx, :codegen, BufCodeGen())
    return generate_code(codegen, rg)
end

"""
fdiff(f::Function; ctx=Dict(), xs...)

Differentiate function `f` w.r.t. its arguments and return a function that
procudes a tuples like this:

    (result, d_arg1, d_arg2, ...)

See also `xdiff(ex::Expr; ctx=Dict(), xs...)`.
"""
function xdiff(f::Function; ctx=Dict(), inputs...)
    ctx = to_context(ctx)
    types = ([typeof(val) for (name, val) in inputs]...)
    args, ex = func_expr(f, types)
    flat_args, ex, st = destruct(args, types, ex)
    ex = sanitize(ex)
    flat_inputs = destruct_inputs(inputs)
    dex = xdiff(ex; ctx=ctx, flat_inputs...)
    ctx[:dex] = dex
    mod = get(ctx, :mod, current_module())
    name = Espresso.genname("$(func_name(f))_deriv_")
    flat_types = [top_type(val) for (name, val) in flat_inputs]
    typed_flat_args = [:($a::$t) for (a, t) in zip(flat_args, flat_types)]
    # function with additional argument `mem`
    fn_ex_mem = make_func_expr(name, [typed_flat_args; :mem], [], dex)
    fn = eval(mod, fn_ex_mem)
    # function with kw argument `mem=Dict()`
    fn_ex_mem_kw = make_func_expr(name, typed_flat_args, [:mem => Dict()], dex)
    eval(mod, fn_ex_mem_kw)
    if any(isstruct, types)
        # preparations for model adapter
        struct_types = [isstruct(t) ? t : top_type(t) for t in types]
        typed_args = [:($a::$t) for (a, t) in zip(args, struct_types)]
        rev_st = Dict(v => k for (k, v) in st)
        model_args = [haskey(rev_st, a) ? rev_st[a] : a for a in flat_args]
        # model adapter: with additional `mem`
        model_fn_ex_mem = :($name($(typed_args...), mem) =
                            $name($(model_args...), mem)) |> sanitize
        eval(mod, model_fn_ex_mem)
        # model adapter: with kw argument `mem=Dict()`
        model_fn_ex_kw = :($name($(typed_args...); mem=Dict()) =
                           $name($(model_args...); mem=Dict())) |> sanitize
        eval(mod, model_fn_ex_kw)
    end
    return fn
end


text:
function TimeStampedFieldAngularCoordinates(arg0::AbsoluteDate, arg1::FieldPVCoordinates, arg2::FieldPVCoordinates, arg3::FieldPVCoordinates, arg4::FieldPVCoordinates, arg5::jdouble)
    return TimeStampedFieldAngularCoordinates((AbsoluteDate, FieldPVCoordinates, FieldPVCoordinates, FieldPVCoordinates, FieldPVCoordinates, jdouble), arg0, arg1, arg2, arg3, arg4, arg5)
end

function TimeStampedFieldAngularCoordinates(arg0::AbsoluteDate, arg1::FieldRotation, arg2::FieldVector3D, arg3::FieldVector3D)
    return TimeStampedFieldAngularCoordinates((AbsoluteDate, FieldRotation, FieldVector3D, FieldVector3D), arg0, arg1, arg2, arg3)
end

function TimeStampedFieldAngularCoordinates(arg0::Field, arg1::TimeStampedAngularCoordinates)
    return TimeStampedFieldAngularCoordinates((Field, TimeStampedAngularCoordinates), arg0, arg1)
end

function TimeStampedFieldAngularCoordinates(arg0::FieldAbsoluteDate, arg1::FieldPVCoordinates, arg2::FieldPVCoordinates, arg3::FieldPVCoordinates, arg4::FieldPVCoordinates, arg5::jdouble)
    return TimeStampedFieldAngularCoordinates((FieldAbsoluteDate, FieldPVCoordinates, FieldPVCoordinates, FieldPVCoordinates, FieldPVCoordinates, jdouble), arg0, arg1, arg2, arg3, arg4, arg5)
end

function TimeStampedFieldAngularCoordinates(arg0::FieldAbsoluteDate, arg1::FieldRotation)
    return TimeStampedFieldAngularCoordinates((FieldAbsoluteDate, FieldRotation), arg0, arg1)
end

function TimeStampedFieldAngularCoordinates(arg0::FieldAbsoluteDate, arg1::FieldRotation, arg2::FieldVector3D, arg3::FieldVector3D)
    return TimeStampedFieldAngularCoordinates((FieldAbsoluteDate, FieldRotation, FieldVector3D, FieldVector3D), arg0, arg1, arg2, arg3)
end

function add_offset(obj::TimeStampedFieldAngularCoordinates, arg0::FieldAngularCoordinates)
    return jcall(obj, "addOffset", TimeStampedFieldAngularCoordinates, (FieldAngularCoordinates,), arg0)
end

function get_date(obj::TimeStampedFieldAngularCoordinates)
    return jcall(obj, "getDate", FieldAbsoluteDate, ())
end

function interpolate(::Type{TimeStampedFieldAngularCoordinates}, arg0::AbsoluteDate, arg1::AngularDerivativesFilter, arg2::Collection)
    return jcall(TimeStampedFieldAngularCoordinates, "interpolate", TimeStampedFieldAngularCoordinates, (AbsoluteDate, AngularDerivativesFilter, Collection), arg0, arg1, arg2)
end

function interpolate(::Type{TimeStampedFieldAngularCoordinates}, arg0::FieldAbsoluteDate, arg1::AngularDerivativesFilter, arg2::Collection)
    return jcall(TimeStampedFieldAngularCoordinates, "interpolate", TimeStampedFieldAngularCoordinates, (FieldAbsoluteDate, AngularDerivativesFilter, Collection), arg0, arg1, arg2)
end

function revert(obj::TimeStampedFieldAngularCoordinates)
    return jcall(obj, "revert", TimeStampedFieldAngularCoordinates, ())
end

function shifted_by(obj::TimeStampedFieldAngularCoordinates, arg0::RealFieldElement)
    return jcall(obj, "shiftedBy", TimeStampedFieldAngularCoordinates, (RealFieldElement,), arg0)
end

function shifted_by(obj::TimeStampedFieldAngularCoordinates, arg0::jdouble)
    return jcall(obj, "shiftedBy", TimeStampedFieldAngularCoordinates, (jdouble,), arg0)
end

function subtract_offset(obj::TimeStampedFieldAngularCoordinates, arg0::FieldAngularCoordinates)
    return jcall(obj, "subtractOffset", TimeStampedFieldAngularCoordinates, (FieldAngularCoordinates,), arg0)
end


text:
using PyPlot
using Convex
using ProfileView
include("CenterOfMass.jl")

sigma = 1.0
forces = rand(2205)
normals = normal_matrix(forces, sigma^2)

function test_csqmi()
  prior = rand(2205)
  prior = prior / sum(prior)
  #@elapsed @profile compute_mutual_information(forces, prior, sigma)
  @elapsed @profile compute_mutual_information(prior, normals)
end

Profile.clear()

times = Float64[]
for ii = 1:100
  time = test_csqmi()
  push!(times, time)
end

println("Mean time: $(mean(times))")

ProfileView.view()

max_stars_repo_name: mcorah/CenterOfMass

max_issues_repo_name: mcorah/CenterOfMass

max_forks_repo_name: mcorah/CenterOfMass


text:
module WiSER

using Base: kwarg_decl
using DataFrames, Tables, LinearAlgebra, MathProgBase
using Printf, Reexport, Statistics, CategoricalArrays
import LinearAlgebra: BlasReal, copytri!
import DataFrames: DataFrame
using Random
import Random: GLOBAL_RNG

@reexport using Ipopt
@reexport using NLopt
@reexport using StatsModels
@reexport using Distributions

export
    # types
    WSVarLmmObs,
    WSVarLmmModel,
    # functions
    coef,
    coefnames,
    coeftable,
    confint,
    DataFrame,
    fit!,
    init_ls!,
    init_mom!,
    nlsv_obj!,
    nclusters,
    nobs,
    rand!,
    respdists,
    rvarlmm,
    rvarlmm!,
    stderror,
    update_res!,
    update_wtmat!,
    vcov

"""
    WSVarLmmObs
    WSVarLmmObs(y, X, Z, W)

A realization of within-subject variance linear mixed model data instance.

# Positional Arguments:
    `y`: the response vector
    `X`: the mean fixed effects covariate matrix
    `Z`: the random location effects covariate matrix
    `W`: the within-subject variance fixed effects covariate matrix

"""
struct WSVarLmmObs{T <: BlasReal}
    # data
    y                     :: Vector{T} # response
    Xt                    :: Matrix{T} # X should include a column of 1's
    Zt                    :: Matrix{T} # Random effect covars
    Wt                    :: Matrix{T} # Covariates that affect WS variability
    # working arrays
    obj                   :: Vector{T} # instance objective value
    ∇β                    :: Vector{T} # gradient
    ∇τ                    :: Vector{T}
    ∇Lγ                   :: Matrix{T}
    Hββ                   :: Matrix{T} # hessian
    Hττ                   :: Matrix{T}
    HτLγ                  :: Matrix{T}
    HLγLγ                 :: Matrix{T}
    res                   :: Vector{T} # residual vector
    res2                  :: Vector{T} # residual vector.^2
    resnrm2               :: Vector{T} # sum of residual squares
    expwτ                 :: Vector{T} # hold exp.(W * τ)
    ztz                   :: Matrix{T} # Z'Z
    ztres                 :: Vector{T} # Z'res
    zlltzt_dg             :: Vector{T}
    storage_n1            :: Vector{T}
    storage_p1            :: Vector{T}
    storage_q1            :: Vector{T}
    storage_pn            :: Matrix{T}
    storage_qn            :: Matrix{T}
    storage_ln            :: Matrix{T}
    storage_pp            :: Matrix{T}
    storage_qq            :: Matrix{T}
    storage_qp            :: Matrix{T}
    storage_q◺n           :: Matrix{T}
    # Woodbury structure for weight matrix Vinv = Dinv - U * U'
    Dinv                  :: Vector{T}
    Ut                    :: Matrix{T}
    # for weighted objective eval
    rt_Dinv_r             :: Vector{T}
    rt_UUt_r              :: Vector{T}
    rt_U                  :: Matrix{T}
    Dinv_r                :: Vector{T}
    rt_UUt                :: Matrix{T}
    Zt_Dinv_r             :: Vector{T}
    rt_UUt_Z              :: Matrix{T}
    diagUUt_Dinv          :: Vector{T}
    Dinv_Z_L              :: Matrix{T}
    UUt_Z_L               :: Matrix{T}
    Ut_D_U                :: Matrix{T}
    Zt_Dinv_Z             :: Matrix{T}
    Lt_Zt_Dinv_Z_L        :: Matrix{T}
    Zt_UUt_Z              :: Matrix{T}
    Lt_Zt_UUt_Z_L         :: Matrix{T}
    # for gradient wrt τ
    diagDVRV              :: Vector{T}
    # for gradient wrt Lγ
    Zt_Dinv               :: Matrix{T}
    Zt_UUt_rrt_Dinv_Z :: Matrix{T}
    Zt_UUt_rrt_UUt_Z      :: Matrix{T}
    Zt_UUt                :: Matrix{T}
    Lt_Zt_Dinv_r          :: Vector{T}
    Zt_Vinv_r             :: Vector{T}
    # for Hessian wrt τ
    Wt_D_Dinv             :: Matrix{T}
    sqrtDinv_UUt          :: Vector{T}
    Ut_kr_Ut              :: Matrix{T}
    Wt_D_Ut_kr_Utt        :: Matrix{T}
    Wt_D_sqrtdiagDinv_UUt :: Matrix{T}
    # for Hessian wrt Lγ
    Zt_Vinv_Z             :: Matrix{T}
    Zt_Vinv               :: Matrix{T}
end

function WSVarLmmObs(
    y::AbstractVector{T},
    X::AbstractMatrix{T},
    Z::AbstractMatrix{T},
    W::AbstractMatrix{T}
    ) where T <: BlasReal
    n, p, q, l  = size(X, 1), size(X, 2), size(Z, 2), size(W, 2)
    q◺ = ◺(q)
    # working arrays
    obj                   = Vector{T}(undef, 1)
    ∇β                    = Vector{T}(undef, p)
    ∇τ                    = Vector{T}(undef, l)
    ∇Lγ                   = Matrix{T}(undef, q, q)
    Hββ                   = Matrix{T}(undef, p, p)
    Hττ                   = Matrix{T}(undef, l, l)
    HτLγ                  = Matrix{T}(undef, l, q◺)
    HLγLγ                 = Matrix{T}(undef, q◺, q◺)
    res                   = Vector{T}(undef, n)
    res2                  = Vector{T}(undef, n)
    resnrm2               = Vector{T}(undef, 1)
    expwτ                 = Vector{T}(undef, n)
    ztz                   = transpose(Z) * Z
    ztres                 = Vector{T}(undef, q)
    zlltzt_dg             = Vector{T}(undef, n)
    storage_n1            = Vector{T}(undef, n)
    storage_p1            = Vector{T}(undef, p)
    storage_q1            = Vector{T}(undef, q)
    storage_pn            = Matrix{T}(undef, p, n)
    storage_qn            = Matrix{T}(undef, q, n)
    storage_ln            = Matrix{T}(undef, l, n)
    storage_pp            = Matrix{T}(undef, p, p)
    storage_qq            = Matrix{T}(undef, q, q)
    storage_qp            = Matrix{T}(undef, q, p)
    storage_q◺n           = Matrix{T}(undef, q◺, n)
    # added for weighted estimating equations
    Dinv                  = Vector{T}(undef, n) #stores diag(exp(-wτ_0))
    Ut                    = Matrix{T}(undef, q, n)
    # denote, r as residual vector D as diagonal(exp(wτ))
    # added for weigthed estimating equations
    rt_Dinv_r             = Vector{T}(undef, 1)
    rt_UUt_r              = Vector{T}(undef, 1)
    rt_U                  = Matrix{T}(undef, 1, q)
    Dinv_r                = Vector{T}(undef, n)
    rt_UUt                = Matrix{T}(undef, 1, n)
    Zt_Dinv_r             = Vector{T}(undef, q)
    rt_UUt_Z              = Matrix{T}(undef, 1, q)
    diagUUt_Dinv          = Vector{T}(undef, n)
    Dinv_Z_L              = Matrix{T}(undef, n, q)
    UUt_Z_L               = Matrix{T}(undef, n, q)
    Ut_D_U                = Matrix{T}(undef, q, q)
    Zt_Dinv_Z             = Matrix{T}(undef, q, q)
    Lt_Zt_Dinv_Z_L        = Matrix{T}(undef, q, q)
    Zt_UUt_Z              = Matrix{T}(undef, q, q)
    Lt_Zt_UUt_Z_L         = Matrix{T}(undef, q, q)
    # for gradient wrt τ
    diagDVRV              = Vector{T}(undef, n)
    # for gradient wrt Lγ
    Zt_Dinv               = Matrix{T}(undef, q, n)
    Zt_UUt_rrt_Dinv_Z     = Matrix{T}(undef, q, q)
    Zt_UUt_rrt_UUt_Z      = Matrix{T}(undef, q, q)
    Zt_UUt                = Matrix{T}(undef, q, n)
    Lt_Zt_Dinv_r          = Vector{T}(undef, q)
    Zt_Vinv_r             = Vector{T}(undef, q)
    # for Hessian wrt τ
    Wt_D_Dinv             = Matrix{T}(undef, l, n)
    sqrtDinv_UUt          = Vector{T}(undef, n)
    Ut_kr_Ut              = Matrix{T}(undef, abs2(q), n)
    Wt_D_Ut_kr_Utt        = Matrix{T}(undef, l, abs2(q))
    Wt_D_sqrtdiagDinv_UUt = Matrix{T}(undef, l, n)
    # for Hessian wrt Lγ
    Zt_Vinv_Z             = Matrix{T}(undef, q, q)
    Zt_Vinv               = Matrix{T}(undef, q, n)
    # constructor
    WSVarLmmObs{T}(
        y, transpose(X), transpose(Z), transpose(W), obj,
        ∇β, ∇τ, ∇Lγ,
        Hββ, Hττ, HτLγ, HLγLγ,
        res, res2, resnrm2, expwτ, ztz, ztres, zlltzt_dg,
        storage_n1, storage_p1, storage_q1,
        storage_pn, storage_qn, storage_ln,
        storage_pp, storage_qq, storage_qp, storage_q◺n,
        Dinv, Ut, rt_Dinv_r, rt_UUt_r, rt_U, Dinv_r,
        rt_UUt, Zt_Dinv_r, rt_UUt_Z, diagUUt_Dinv,
        Dinv_Z_L, UUt_Z_L, Ut_D_U,  Zt_Dinv_Z,
        Lt_Zt_Dinv_Z_L, Zt_UUt_Z, Lt_Zt_UUt_Z_L,
        diagDVRV,
        Zt_Dinv, Zt_UUt_rrt_Dinv_Z, Zt_UUt_rrt_UUt_Z,
        Zt_UUt, Lt_Zt_Dinv_r, Zt_Vinv_r, Wt_D_Dinv,
        sqrtDinv_UUt, Ut_kr_Ut, Wt_D_Ut_kr_Utt,
        Wt_D_sqrtdiagDinv_UUt, Zt_Vinv_Z, Zt_Vinv)
end

"""
    WSVarLmmModel

Within-subject variance linear mixed model, which contains a vector of
`WSVarLmmObs` as data, model parameters, and working arrays.

    WSVarLmmModel(obsvec; obswts, meannames, renames, wsvarnames)

# Positional arguments
- `obsvec`: Vector of WSVarLmmObs

# Keyword arguments
- `obswts`: Subject-level weight vector of observation weights, length of the `obsvec` object.
- `meannames`: Names of the mean fixed effects covariates
- `renames`: Names of the random location effects covariates
- `wsvarnames`: Names of the ws variance fixed effects covariates

"""
struct WSVarLmmModel{T <: BlasReal} <: MathProgBase.AbstractNLPEvaluator
    # data
    data            :: Vector{WSVarLmmObs{T}}
    respname        :: String
    meannames       :: Vector{String} # names of mean fixed effect variables
    renames         :: Vector{String} # names of random location effect variables
    wsvarnames      :: Vector{String} # names of ws var fixed effect variables
    meanformula     :: FormulaTerm
    reformula       :: FormulaTerm
    wsvarformula    :: FormulaTerm
    ids             :: Union{Vector{<:AbstractString}, Vector{<:Int}} # IDs of individuals/clusters in order
    obswts          :: Vector{T} # individual/cluster weights
    # dimenions
    p               :: Int       # number of mean parameters in linear regression
    q               :: Int       # number of random effects
    l               :: Int       # number of parameters for modeling WS variability
    m               :: Int       # number of individuals/clusters
    nis             :: Vector{Int}  # number of observations per cluster
    nsum            :: Int       # number of observations (summed across individuals)
    # sufficient statistics
    xtx             :: Matrix{T} # sum_i Xi'Xi
    xty             :: Vector{T} # sum_i Xi'yi
    wtw             :: Matrix{T} # sum_i Wi'Wi
    ztz2            :: Matrix{T} # sum_i Zi'Zi ⊗ Zi'Zi
    ztz2od          :: Matrix{T} # sum_i (Zi'Zi ⊗ Zi'Zi - (Zi' ⊙ Zi')(Zi' ⊙ Zi')')
    # parameters
    β               :: Vector{T}  # p-vector of mean regression coefficients
    τ               :: Vector{T}  # l-vector of WS variability regression coefficients
    Lγ              :: Matrix{T}  # q by q lower triangular Cholesky factor of cov(γ)
    Σγ              :: Matrix{T}  # q by q covariance matrix of γ
    # working arrays
    ∇β              :: Vector{T}
    ∇τ              :: Vector{T}
    ∇Lγ             :: Matrix{T}
    ∇Σγ             :: Vector{T}
    Hββ             :: Matrix{T}
    Hττ             :: Matrix{T}
    HτLγ            :: Matrix{T}
    HLγLγ           :: Matrix{T}
    HΣγΣγ           :: Matrix{T}
    # weighted NLS or unweighted NLS
    iswtnls         :: Vector{Bool}
    # multi-threading or not
    ismthrd         :: Vector{Bool}
    # model has been fit or not
    isfitted        :: Vector{Bool}
    # for sandwich estimator
    ψ               :: Vector{T}
    Ainv            :: Matrix{T}
    B               :: Matrix{T}
    vcov            :: Matrix{T}
end

function WSVarLmmModel(
    obsvec      :: Vector{WSVarLmmObs{T}};
    obswts      :: Vector = [],
    respname    :: String = "y",
    meannames   :: Vector{String} = ["x$i" for i in 1:size(obsvec[1].Xt, 1)],
    renames     :: Vector{String} = ["z$i" for i in 1:size(obsvec[1].Zt, 1)],
    wsvarnames  :: Vector{String} = ["w$i" for i in 1:size(obsvec[1].Wt, 1)],
    meanformula :: FormulaTerm = FormulaTerm(term(Symbol(respname)),
                        sum(term.(Symbol.(meannames)))),
    reformula   :: FormulaTerm = FormulaTerm(term(Symbol(respname)),
                        sum(term.(Symbol.(renames)))),
    wsvarformula:: FormulaTerm = FormulaTerm(term(Symbol(respname)),
                        sum(term.(Symbol.(wsvarnames)))),
    ids         :: Union{Vector{<:AbstractString}, Vector{Int}} = collect(1:length(obsvec))
    ) where T <: BlasReal
    # dimensions
    p            = size(obsvec[1].Xt, 1)
    q            = size(obsvec[1].Zt, 1)
    l            = size(obsvec[1].Wt, 1)
    m            = length(obsvec)
    nis          = map(o -> length(o.y), obsvec)
    nsum         = sum(nis)
    q◺           = ◺(q)
    # sufficient statistics
    xtx          = zeros(T, p, p)
    xty          = zeros(T, p)
    wtw          = zeros(T, l, l)
    ztz2         = zeros(T, abs2(q), abs2(q))
    ztz2od       = zeros(T, abs2(q), abs2(q))
    for obs in obsvec
        # accumulate Xi'Xi
        BLAS.syrk!('U', 'N', T(1), obs.Xt, T(1), xtx)
        # accumulate Xi'yi
        BLAS.gemv!('N', T(1), obs.Xt, obs.y, T(1), xty)
        # accumulate Wi' * Wi
        BLAS.syrk!('U', 'N', T(1), obs.Wt, T(1), wtw)
        # accumulate Zi'Zi ⊗ Zi'Zi
        kron_axpy!(obs.ztz, obs.ztz, ztz2)
        # accumualte (Zi' ⊙ Zi')(Zi' ⊙ Zi')'
        # Ut_kr_Ut used as scratch space to store Zi' ⊙ Zi'
        kr_axpy!(obs.Zt, obs.Zt, fill!(obs.Ut_kr_Ut, 0))
        BLAS.syrk!('U', 'N', T(1), obs.Ut_kr_Ut, T(1), ztz2od)
    end
    ztz2od .= ztz2 .- ztz2od
    copytri!(   xtx, 'U')
    copytri!(   wtw, 'U')
    copytri!(  ztz2, 'U')
    copytri!(ztz2od, 'U')
    # parameters
    β        = Vector{T}(undef, p)
    τ        = Vector{T}(undef, l)
    Lγ       = Matrix{T}(undef, q, q)
    Σγ       = Matrix{T}(undef, q, q)
    # gradients
    ∇β       = Vector{T}(undef, p)
    ∇τ       = Vector{T}(undef, l)
    ∇Lγ      = Matrix{T}(undef, q, q)
    ∇Σγ      = Vector{T}(undef, abs2(q))
    Hββ      = Matrix{T}(undef, p, p)
    Hττ      = Matrix{T}(undef, l, l)
    HτLγ     = Matrix{T}(undef, l, q◺)
    HLγLγ    = Matrix{T}(undef, q◺, q◺)
    HΣγΣγ    = Matrix{T}(undef, abs2(q), abs2(q))
    # weighted NLS fitting or not
    iswtnls = [false]
    # multi-threading or not
    ismthrd  = [false]
    # has been fit or not
    isfitted = [false]
    # sandwich estimator
    ψ        = Vector{T}(undef, p + q◺ + l)
    Ainv     = Matrix{T}(undef, p + q◺ + l, p + q◺ + l)
    B        = Matrix{T}(undef, p + q◺ + l, p + q◺ + l)
    vcov     = Matrix{T}(undef, p + q◺ + l, p + q◺ + l)
    # constructor
    WSVarLmmModel{T}(
        obsvec, respname, meannames, renames, wsvarnames,
        meanformula, reformula, wsvarformula,
        ids, obswts, p, q, l, m, nis, nsum,
        xtx, xty, wtw, ztz2, ztz2od,
        β,  τ,  Lγ, Σγ,
        ∇β, ∇τ, ∇Lγ, ∇Σγ,
        Hββ, Hττ, HτLγ, HLγLγ, HΣγΣγ,
        iswtnls, ismthrd, isfitted,
        ψ, Ainv, B, vcov)
end

coefnames(m::WSVarLmmModel) = [m.meannames; m.wsvarnames]
coef(m::WSVarLmmModel) = [m.β; m.τ]
nobs(m::WSVarLmmModel) = m.nsum
nclusters(m::WSVarLmmModel) = m.m
stderror(m::WSVarLmmModel) = [sqrt(m.vcov[i, i]) for i in 1:(m.p + m.l)]
vcov(m::WSVarLmmModel) = m.vcov # include variance parts of Lγ?

confint(m::WSVarLmmModel, level::Real) = hcat(coef(m), coef(m)) +
    stderror(m) * quantile(Normal(), (1. - level) / 2.) * [1. -1.]
confint(m::WSVarLmmModel) = confint(m, 0.95)

function getformula(m::WSVarLmmModel, s::String)
    frmla = s == "mean" ? m.meanformula : s == "var" ?
        m.wsvarformula : s == "re" ? m.reformula : nothing
    rhs = length(terms(frmla.rhs)) == 1 ?
        string(term(frmla.rhs)) : join(frmla.rhs, " + ")
    return m.respname * " ~ " * rhs
end

function coeftable(m::WSVarLmmModel)
    mstder = stderror(m)
    mcoefs = coef(m)
    wald = mcoefs ./ mstder
    pvals = 2 * Distributions.ccdf.(Normal(), abs.(wald))
    StatsModels.CoefTable(hcat(mcoefs, mstder, wald, pvals),
        ["Estimate", "Std. Error", "Z", "Pr(>|Z|)"],
        coefnames(m), 4, 3)
end

function Base.show(io::IO, m::WSVarLmmModel)
    if !m.isfitted[1]
        @warn("The model has not been fit.")
        return nothing
    end
    println(io)
    println(io, "Within-subject variance estimation by robust regression (WiSER)")
    println(io)
    println(io, "Mean Formula:")
    println(io, getformula(m, "mean"))
    println(io, "Random Effects Formula:")
    println(io, getformula(m, "re"))
    println(io, "Within-Subject Variance Formula:")
    println(io, getformula(m, "var"))
    println(io)
    println(io, "Number of individuals/clusters: $(m.m)")
    println(io, "Total observations: $(m.nsum)")
    println(io)
    println(io, "Fixed-effects parameters:")
    show(io, coeftable(m))
    println(io)
    println(io, "Random effects covariance matrix Σγ:")
    Base.print_matrix(IOContext(io, :compact => true), [m.renames m.Σγ])
    println(io)
    println(io)
end

include("nls.jl")
include("initialization.jl")
include("fit.jl")
include("df.jl")
include("rand.jl")
include("multivariate_calculus.jl")
include("sandwich.jl")

end


text:
using AntOptim
using Test

@testset "AntOptim.jl" begin
    # Write your tests here.
end