text:
__precompile__()
module AMDB

import Base: close, count, keys, length, show, keys, values
import Base.Markdown

using ArgCheck: @argcheck
using ByteParsers:
    isparsed, Skip, Line, ByteVector, @checkpos, AbstractParser, parsedtype,
    DateYYYYMMDD, MaybeParsed, getpos
import ByteParsers: parsenext
using CodecZlib: GzipDecompressorStream
using DataFrames
using DataStructures: OrderedDict, counter, Accumulator
using DocStringExtensions: SIGNATURES
using DiscreteRanges: DiscreteRange
using EnglishText: ItemQuantity
using FileIO: load
using FlexDates: FlexDate
using IndirectArrays: IndirectArray
import JLD2
using LargeColumns: SinkColumns, meta_path, MmappedColumns, get_columns
using Lazy: @forward
using Parameters: @unpack
using WallTimeProgress: WallTimeTracker, increment!
using ProgressMeter

export
    data_file, data_path, all_data_files, data_colnames,
    serialize_data, deserialize_data,
    AMDB_Date,
    normalize_rows, count_individual_transitions, trim_counter,
    get_mis, individual_columns, individual_df, get_age_gender

# paths

"""
    $(SIGNATURES)

Return the directoy for the AMDB data dump.

The user should set the environment variable `AMDB_FILES`.
"""
function data_directory()
    key = "AMDB_FILES"
    @assert(haskey(ENV, key),
            "You should put the path to the raw files in ENV[$key].")
    normpath(expanduser(ENV[key]))
end

"""
    $(SIGNATURES)

Add `components` (directories, and potentially a filename) to the AMDB
data dump directory.
"""
data_path(components...) = joinpath(data_directory(), components...)

const VALID_YEARS = 2000:2016

"""
    $(SIGNATURES)

The path for the AMDB data dump file (gzip-compressed CSV) for a given
year. Example:

```julia
AMDB.data_file(2000)
```
"""
function data_file(year)
    @assert year ∈ VALID_YEARS "Year outside $(VALID_YEARS)."
    ## special-case two years combined
    yearnum = year ≤ 2014 ? @sprintf("%02d", year-2000) : "1516"
    data_path("mon_ew_xt_uni_bus_$(yearnum).csv.gz")
end

"""
    $(SIGNATURES)

"""
all_data_files() = unique(data_file.(VALID_YEARS))

"""
    $(SIGNATURES)

Return the AMDB column names for the data dump as a vector.
"""
function data_colnames()
    header = readlines(data_path("mon_ew_xt_uni_bus.cols.txt"))[1]
    split(header, ';')
end

"""
    $(SIGNATURES)

Serialize data into `filename` within the data directory. A new file is created,
existing files are overwritten.
"""
function serialize_data(filename, value)
    open(io -> serialize(io, value), data_path(filename), "w")
end

"""
    $(SIGNATURES)

Deserialize data from `filename` within the data directory.
"""
function deserialize_data(filename)
    open(deserialize, data_path(filename), "r")
end

# error logging

struct FileError
    line_number::Int
    line_content::Vector{UInt8}
    line_position::Int
end

function show(io::IO, file_error::FileError)
    @unpack line_number, line_content, line_position = file_error
    println(io, chomp(String(line_content)))
    print(io, " "^(line_position - 1))
    println(io, "^ line $(line_number), byte $(line_position)")
end

struct FileErrors{S}
    filename::S
    errors::Vector{FileError}
end

FileErrors(filename::String) = FileErrors(filename, Vector{FileError}(0))

function log_error(file_errors::FileErrors, line_number, line_content,
                   line_position)
    push!(file_errors.errors,
          FileError(line_number, line_content, line_position))
end

function show(io::IO, file_errors::FileErrors)
    @unpack filename, errors = file_errors
    error_quantity = ItemQuantity(length(errors), "error")
    println(io, "$filename: $(error_quantity)")
    for e in errors
        show(io, e)
    end
end

count(file_errors::FileErrors) = length(file_errors.errors)

# automatic indexation

struct AutoIndex{T, S <: Integer}
    dict::Dict{T, S}
end

show(io::IO, ai::AutoIndex{T, S}) where {T, S} =
    print(io, "automatically index $T => $S, $(length(ai.dict)) elements")

"""
    $SIGNATURES

Create an `AutoIndex` of object which supports automatic indexation of keys with
type `T` to integers (`<: S`). The mapping is implemented by making the object
callable.

`keys` retrieves all the keys in the order of appearance.
"""
AutoIndex(T, S) = AutoIndex(Dict{T,S}())

length(ai::AutoIndex) = length(ai.dict)

function (ai::AutoIndex{T,S})(elt::E) where {T,S,E}
    @unpack dict = ai
    ix = get(dict, elt, zero(S))
    if ix == zero(S)
        v = length(dict)
        @assert v < typemax(S) "Number of elements reached typemax($S), can't insert $elt"
        v += one(S)
        dict[T==E ? elt : convert(T, elt)] = v
        v
    else
        ix
    end
end

function keys(ai::AutoIndex)
    kv = collect(ai.dict)
    sort!(kv; by = last)
    first.(kv)
end

# FIXME remove unused
# # map using a tuple of functions

# """
#     TupleMap(functions::Tuple)

# Return a callable that maps a tuple of the same dimension using `functions`,
# which should be a tuple of callables (do not need to be `<: Function`).

# ```julia
# struct AddOne end
# (::AddOne)(x) = x + one(x)
# f = TupleMap((AddOne(), identity))
# f((1,3))                        # (2,3)
# ```
# """
# struct TupleMap{T <: Tuple}
#     functions::T
# end

# (f::TupleMap)(x::Tuple) = map((g, z) -> g(z), f.functions, x)

# dates

"""
Epoch for consistent date compression in the database.
"""
const EPOCH = Date(2000,1,1)    # all dates relative to this

"""
Datatype used for compressed dates.
"""
const AMDB_Date = FlexDate{EPOCH,Int16} # should be enough for everything

# tuple processing (first pass)

# FIXME remove, unused
# """
#     $SIGNATURES

# Join the second and the third argument as a DiscreteRange compressed dates.
# """
# join_dates(record) = _join_dates(record...)

# @inline _join_dates(id, spell_start, spell_end, rest...) =
#     id, DiscreteRange(AMDB_Date(spell_start), AMDB_Date(spell_end)), rest...

struct MultiSubs{P, F}
    functions::F
end

"""
    $SIGNATURES

Return a callable that maps tuples by maps elements at `positions` with the
corresponding function in `functions`, leaving the rest of the elements alone.
"""
function MultiSubs(positions::NTuple{N, Int}, functions::F) where {N, F}
    @argcheck length(positions) == length(functions)
    @argcheck allunique(positions) && all(positions .> 0)
    MultiSubs{positions, F}(functions)
end

"""
    $SIGNATURES

Return positions of substituted indices.
"""
get_positions(::MultiSubs{P}) where P = P

@generated function (m::MultiSubs{P})(x::NTuple{N, Any}) where {P, N}
    result = Any[:(x[$i]) for i in 1:N]
    for (i, p) in enumerate(P)
        result[p] = :((m.functions[$i])(x[$p]))
    end
    :(tuple($(result...)))
end

"""
    OrderedCounter{T,S}()

A wrapper for OrderedDict that is callable, counts and returns its argument.
"""
struct OrderedCounter{T,S <: Integer}
    dict::OrderedDict{T,S}
end

OrderedCounter{T, S}() where {T, S} = OrderedCounter(OrderedDict{T, S}())

function (oc::OrderedCounter{T,S})(x::T) where {T,S}
    oc.dict[x] = get(oc.dict, x, zero(S)) + one(S)
    x
end

@forward OrderedCounter.dict keys, values

"""
    $SIGNATURES

Return the counts as an OrderedDict.
"""
get_counts(oc::OrderedCounter) = oc.dict

# first pass processing

"""
    $SIGNATURES

Merge the two column names for start and end dates of spells.
"""
function merged_colnames(; colnames = data_colnames())
    @argcheck colnames[2:3] == ["ANFDAT", "ENDDAT"]
    vcat(colnames[1:1], ["STARTEND"], colnames[4:end])
end

"""
    DatePair()

Parse two consecutive dates.
"""
struct DatePair <: AbstractParser{DiscreteRange{AMDB_Date}} end

function parsenext(::DatePair, str::ByteVector, pos, sep)
    D = DateYYYYMMDD()
    @checkpos (pos, date1) = parsenext(D, str, pos, sep)
    @checkpos (pos, date2) = parsenext(D, str, pos, sep)
    return MaybeParsed(pos, DiscreteRange(AMDB_Date(date1), AMDB_Date(date2)))
    @label error
    MaybeParsed{DiscreteRange{AMDB_Date}}(pos_to_error(pos))
end

"""
Parse a date and compress to internal date format.
"""
struct ParseDate <: AbstractParser{AMDB_Date} end

function parsenext(::ParseDate, str::ByteVector, pos, sep)
    @checkpos (pos, raw_date) = parsenext(DateYYYYMMDD(), str, pos, sep)
    return MaybeParsed(pos, AMDB_Date(raw_date))
    @label error
    MaybeParsed{AMDB_Date}(pos_to_error(pos))
end

struct ColSpec
    name::AbstractString
    parser::AbstractParser
    index_type::Type{ <: Union{Void, Integer}}
end

"""
    $SIGNATURES

Create a column specification for first pass reading.

`name` is the name of the column, parsed using `parser`.

When `index_type::Type{<:Integer}` is given, it is used as a result type to
generate a `AutoIndex` with the given result type, except for the first parsed
column, which uses an `OrderedCounter`.

See [`make_first_pass`](@ref).
"""
ColSpec(name, parser; index_type = Void) = ColSpec(name, parser, index_type)

"""
First pass processing.

`lineparser` parsed each line.

When successful, it is transformed with `multisubs`, changing the state of
`accumulators`.

The result is written into `sink`.
"""
struct FirstPass{L <: Line, A <: Tuple, M <: MultiSubs, S <: SinkColumns}
    lineparser::L
    accumulators::A
    multisubs::M
    sink::S
    colnames::Vector{Symbol}
end

function show(io::IO, fp::FirstPass)
    @unpack lineparser, accumulators, colnames = fp
    println(io, "First pass parsing")
    println(io, "Line parser:")
    println(io, lineparser)
    println(io, "Accumulators:")
    println(io, accumulators)
    println(io, "Colnames:")
    println(io, colnames)
end

function make_firstpass(dir, colspecs::AbstractVector{ColSpec};
                        colnames = merged_colnames(),
                        skip_parser = Skip())
    @argcheck allunique(colnames) "Column names are not unique."
    matched_specs = map(colspecs) do colspec
        @unpack name, parser, index_type = colspec
        colindex = findfirst(colnames, name)
        colindex > 0 || throw(ArgumentError("column $(colname) not found"))
        (colindex, name, parser, index_type)
    end
    # make sure column indexes are strictly increasing (no repetition, right order)
    issorted(matched_specs, lt = <, by = first) ||
        throw(ArgumentError("column names need to be in the same order as in data"))
    # the default is to skip
    parsers = fill!(Vector(matched_specs[end][1]), skip_parser)
    result_types = Any[]
    sub_positions = Int[]
    sub_functions = Any[]
    names = Vector{Symbol}()
    for (position, (colindex, name, parser, index_type)) in enumerate(matched_specs)
        parsers[colindex] = parser
        result_type = parsedtype(parser)
        if !(index_type ≡ Void)
            @argcheck index_type <: Integer
            filter = AutoIndex(result_type, index_type)
            result_type = index_type
            push!(sub_positions, position)
            push!(sub_functions, filter)
        end
        push!(result_types, result_type)
        push!(names, Symbol(name))
    end
    recordtype = Tuple{result_types...}
    FirstPass(Line(parsers...),
              tuple(sub_functions...),
              MultiSubs(tuple(sub_positions...), tuple(sub_functions...)),
              SinkColumns(dir, recordtype),
              names)
end

close(fp::FirstPass) = close(fp.sink)

"""
    $SIGNATURES

Process the stream `io` by line. Each line is parsed using `parser`, then

1. if the parsing is successful, `f!` is called on the parsed record,

2. if the parsing is not successful, an error is logged to `errors`. See
[`log_error`](@ref).

`tracker` is used to track progress.

When `max_lines > 1`, it is used to limit the number of lines parsed.
"""
function firstpass_process_stream(io::IO, fp::FirstPass, errors::FileErrors;
                                  tracker = WallTimeTracker(10_000_000;
                                                            item_name = "line"),
                                  max_lines = -1)
    @unpack lineparser, multisubs, sink = fp
    while !eof(io) && (max_lines < 0 || count(tracker) < max_lines)
        line_content = readuntil(io, 0x0a)
        maybe_record = parsenext(lineparser, line_content, 1, UInt8(';'))
        if isparsed(maybe_record)
            record = multisubs(unsafe_get(maybe_record))
            push!(sink, record)
        else
            log_error(errors, count(tracker), line_content, getpos(maybe_record))
        end
        increment!(tracker)
    end
end

"""
    $SIGNATURES

Open `filename` and process the resulting stream with
[`firstpass_stream`](@ref), which the other arguments are passed on to. Return
the error log object.
"""
function firstpass_process_file(filename, fp; args...)
    io = GzipDecompressorStream(open(filename))
    errors = FileErrors(filename)
    firstpass_process_stream(io, fp, errors; args...)
    errors
end

# narrowest integer

"""
    narrowest_Int(min, max, [signed = true])

Return the narrowest subtype of `Signed` or `Unsigned` (depending on Signed)
that can contain values between `min` and `max`.
"""
function narrowest_Int(min_::Integer, max_::Integer, signed = true)
    @argcheck min_ ≤ max_
    Ts = signed ?
        [Int8, Int16, Int32, Int64, Int128] :
        [UInt8, UInt16, UInt32, UInt64, UInt128]
    narrowest_min_ = findfirst(T -> typemin(T) ≤ min_, Ts)
    narrowest_max_ = findfirst(T -> max_ ≤ typemax(T), Ts)
    @assert((narrowest_min_ * narrowest_max_) > 0,
            "Can't accomodate $min_:$max_ within given types.")
    Ts[max(narrowest_min_, narrowest_max_)]
end

"""
    to_narrowest_Int(xs, [signed = true])

Convert `xs` to the narrowest integer type that will contain it (a subtype of
`Signed` or `Unsigned`, depending on `signed`).
"""
function to_narrowest_Int(xs::AbstractVector{<: Integer}, signed = true)
    T = narrowest_Int(extrema(xs)..., signed)
    T.(xs)
end

# column selection

"""
    $SIGNATURES

Find the column matching a given column name, and return a tuple of parsers
constructed using `colnames_and_parsers`, which should be a vector of
parsers. The order of parsers is preserved, and should be
increasing. `skip_parser` is used to skip fields.

Example:

```julia
julia> column_parsers(["a", "b", "c", "d", "e"],
                      ["b" => DateYYYYMMDD(), "d" => PositiveInteger()])

(Skip(), DateYYYYMMDD(), Skip(), PositiveInteger())
```
"""
function column_parsers(colnames::AbstractVector,
                        colnames_and_parsers::AbstractVector{<: Pair},
                        skip_parser = Skip())
    @argcheck allunique(colnames) "Column names are not unique."
    indexes_and_parsers = map(colnames_and_parsers) do colname_and_parser
        colname, parser = colname_and_parser
        index = findfirst(colnames, colname)
        index > 0 || throw(ArgumentError("column $(colname) not found"))
        index => parser
    end
    # make sure column indexes are strictly increasing (no repetition, right order)
    issorted(indexes_and_parsers, lt = <, by = first) ||
        throw(ArgumentError("column names are not in the right order"))
    # the default is to skip
    parsers = fill!(Vector(first(indexes_and_parsers[end])), skip_parser)
    for (index, parser) in indexes_and_parsers
        parsers[index] = parser
    end
    tuple(parsers...)
end

"""
    $SIGNATURES

Return a sample of `N` values from column `colname` as a vector.

Optional arguments specity the file, how many lines to skip in between, etc. The
purpose of this function is to get a quick sense of what the values look like.
"""
function preview_column(colname;
                        N = 20,
                        skiplines = 10000,
                        year = 2000,
                        filename = AMDB.data_file(year),
                        colnames = AMDB.data_colnames(),
                        separator = ';')
    ix = findfirst(colnames, colname)
    ix > 0 || throw(error("column $colname not found"))
    info("$colname is in column $ix")
    io_gz = open(filename, "r")
    io = GzipDecompressorStream(io_gz)
    values = map(1:N) do _
        for _ in 1:skiplines
            readline(io)
        end
        line = readline(io)
        split(line, separator)[ix]
    end
    close(io)
    close(io_gz)
    values
end

# rudimentary format for metadata

"""
Key for the index of individual-specific spells, in metadata for collated
columns (`Vector{UnitRange{ <: Integer}}`).
"""
const META_IX = "ix"

"""
Key for the keys of indexed columns (ie strings), in metadata for collated
columns (`Vector{Vector{String}}`).
"""
const META_INDEXED_KEYS = "indexed_keys"

"""
Key for the keys of indexed columns (ie strings), in metadata for collated
columns (`Vector{Int}`).
"""
const META_INDEXED_POSITIONS = "indexed_positions"

"""
Key for the keys of indexed columns (ie strings), in metadata for collated
columns `(Vector{Symbol}`).
"""
const META_COLUMN_NAMES = "column_names"

"""
    CollatedData

A wrapper for the collated dataset.

`ix` contains the indices in the flat vectors for each individual.

`colnames` is a mapping from the column names to their index.

`cols` contains the columns.

FIXME replace with named tuples for faster lookup when v0.7 is released.
"""
struct CollatedData{IX, COLS}
    ix::IX
    colnames::OrderedDict{Symbol, Int}
    columns::COLS
end

"""
    $SIGNATURES

Read the collated dataset as a dictinary of `:colname => column` pairs.

All columns are of the same length, except for `:ix`, which contains ranges that
select contiguous records for an individual in the *other* columns.

This is the recommended entry point for **using** collated data.
"""
function collated_dataset(dir = "collated")
    collated = MmappedColumns(data_path(dir))
    meta = load(meta_path(collated, "meta.jld2"))
    dict = Dict{Symbol, AbstractVector{<: Any}}()
    columns = collect(Any, get_columns(collated))
    for (indexed_position, indexed_keys) in zip(meta[META_INDEXED_POSITIONS],
                                                meta[META_INDEXED_KEYS])
        columns[indexed_position] = IndirectArray(columns[indexed_position],
                                                  indexed_keys)
    end
    CollatedData(meta[META_IX],
                 OrderedDict(reverse.(collect((enumerate(meta[META_COLUMN_NAMES]))))),
                 columns)
end

"""
    $SIGNATURES

Column names of data.
"""
get_colnames(data::CollatedData) = collect(keys(data.colnames))

"""
    $SIGNATURES

Look up `keys` in the `values` of the `IndirectArray`, and return a vector of
them. Using these keys allows for faster lookup, eg using `≡` for strings.
"""
function intern_keys(A::IndirectArray, keys)
    @unpack values = A
    ixs = [findfirst(values, key) for key in keys]
    @argcheck all(ixs .> 0) "Keys $(keys[ixs .== 0]) not found in this array."
    values[ixs]
end

# counting and data analysis

"""
    $SIGNATURES

Count the values in `x` by spell duration (in vector `start_end`). Return a
`DataStructures.Accumulator` object.
"""
function count_total_length(start_end, x::AbstractVector{T}) where T
    c = counter(T)
    for (se, elt) in zip(start_end, x)
        push!(c, elt, length(se))
    end
    c
end

"""
Representation for proportions. For easier printing. Proportions are of the form
`key => value`, where `value` sums to `1`. Proportions are ordered by decreasing
value.
"""
struct Proportions{T, S <: AbstractFloat}
    proportions::Vector{Pair{T, S}}
end

"""
    $SIGNATURES

Calculate proportions from an accumulator.
"""
function Proportions(acc::Accumulator)
    kv = collect(acc.map)
    sort!(kv, by = last, rev = true)
    total = sum(last, kv)
    Proportions([k => v/total for (k, v) in kv])
end

"""
    $SIGNATURES

Aggregate the tail (lowest proportions) to `label`, keeping the first (largest)
`keep`.
"""
function aggregate_tail(p::Proportions{T}, keep::Integer, label::T) where T
    @unpack proportions = p
    kept = proportions[1:keep]
    tail = label => sum(last, proportions[(keep+1):end])
    Proportions(vcat(kept, [tail]))
end

function to_markdown(p::Proportions)
    table_body = [[string(k), string(signif(v * 100, 3)) * "%"]
                  for (k,v) in p.proportions]
    t = Markdown.Table(vcat([["category", "frequency"]], table_body), [:l, :r])
    Markdown.MD(t)
end

dump_latex(filename, object::Base.Markdown.MD) =
    open(io -> show(io, MIME"text/latex"(), object), filename, "w")

dump_latex(filename, object) = dump_latex(filename, to_markdown(object))

"Nice labels for columns. Extend as necessary."
#const
nicelabels = Dict(:PENR => "person id",
                  :STARTEND => "spell",
                  :BENR => "firm/agency",
                  :AM => "cat",
                  :SUM_MA => "#E",
                  :NACE => "ind",
                  :RGS => "loc",
                  :AVG_BMG => "wage",
                  )

"""
    $SIGNATURES

Lookup index (in `data[:ix]`) for personal id `penr`.
"""
function PENR_to_index(data, penr)
    _penr = data[:PENR]
    index = findfirst(r -> _penr[r[1]] == penr, data[:ix])
    @argcheck index > 0 "PENR $penr not found in data."
    index
end

"""
    $SIGNATURES

Format as a string for tabular (Markdown) display.
"""
_fmt(x) = string(x)

_fmt(x::DiscreteRange{<: FlexDate}) =
    string(convert(Date, x.left)) * "…" * string(convert(Date, x.right))

"""
    $SIGNATURES

Individual history for `penr`; return the relevant columns as a Markdown table.
"""
function individual_history(data, penr, column_names...)
    index = AMDB.PENR_to_index(data, penr)
    r = data[:ix][index]
    column_names = [column_names...]
    rows = [getindex.(AMDB.nicelabels, column_names)]
    cols = getindex.(data, column_names)
    for i in r
        push!(rows, @. _fmt(getindex(cols, i)))
    end
    Markdown.MD(Markdown.Paragraph(["individual #$(penr)"]),
                Markdown.Table(rows, fill(:r, length(cols))))
end

# working with data

"""
    $SIGNATURES

Convert STICHTAG to "month in sample", with January 2000 as `1`.
"""
function get_mis(stichtag::Date)
    y, m = Dates.yearmonth(stichtag)
    Int16((y - 2000) * 12 + m)
end

get_mis(stichtag::FlexDate) = get_mis(convert(Date, stichtag))

"""
    $SIGNATURES

Return individual observations for the given columns.

Resolves `individual_index` in `data[:ix]`, return a vector views. See also
[`individual_df`](@ref) to get dataframes.
"""
function individual_columns(data::CollatedData, individual_index, column_names)
    @unpack ix, colnames, columns = data
    ix_individual = ix[individual_index]
    map(key -> view(columns[colnames[key]], ix_individual), column_names)
end

"""
    $SIGNATURES

Individual observations as a data frame. When `colnames` is not specified,
return all columns.
"""
function individual_df(data::CollatedData, individual_index,
                       colnames = get_colnames(data))
    DataFrame(individual_columns(data, individual_index, colnames), colnames)
end

function get_age_gender(data, individual_index, ref_mis = 0)
    gender, age, stichtag = individual_observations(data, individual_index,
                                                    (:GESCHLECHT, :ALTER,
                                                     :STICHTAG))
    first_mis = get_mis(first(stichtag))
    age_ref = first(age) - round(Int, (first_mis - ref_mis)/12)
    age_ref, first(gender)
end

"""
    $SIGNATURES

Return the index of the first set which contains `value`, otherwise `0`.
"""
classify(value, sets) = findfirst(set -> value ∈ set, sets)

"""
    $SIGNATURES

Look up the individual with the given index in `data`, [`classify`](@ref) the
spells according to sets, then increment `counter[t, index_from, index_to]`,
where `t` is the time of the month in sample.

Contiguity of data is ensured by checking month in sample.
"""
function add_individual_transitions!(counter, data, individual_index, sets)
    am, st = individual_observations(data, individual_index, (:AM, :STICHTAG))
    mis = get_mis.(st)
    clas = map(am -> classify(am, sets), am)
    for i in 2:length(mis)
        t = mis[i-1]
        source = clas[i-1]
        dest = clas[i]
        if ((mis[i-1] == t) && (source ≠ 0) && (dest ≠ 0))
            counter[t, source, dest] += 1
        end
    end
end

"""
    $SIGNATURES

Remove the trailing zeros of `counter`, plus `extra` slices (last transitions
tend to be noisy).
"""
function trim_counter(counter, extra = 0)
    sums = squeeze(sum(counter, (2, 3)), (2, 3))
    counter[1:(findlast(!iszero, sums) - extra), :, :]
end

"""
    $SIGNATURES

Accumulate the individual transitions classified in to `sets` in a counter
matrix, see [`add_individual_transitions!`](@ref).

Initially, `max_mis` is used to determine the length of the counter,
overallocating is not very costly, it is trimmed in the end.
"""
function count_individual_transitions(data, individual_indexes, sets;
                                      max_mis = 300)
    counter = zeros(Int, 300, 3, 3)
    @showprogress for i in individual_indexes
        add_individual_transitions!(counter, data, i, sets)
    end
    counter = trim_counter(counter)
end

"""
    $SIGNATURES

Return the argument with its rows normalized to `1` (probabilities).
"""
normalize_rows(A::AbstractMatrix) = A ./ squeeze(sum(A, 2), 2)

end # module


text:
# Bazel can't deal with `ccache`
ENV["BINARYBUILDER_USE_CCACHE"] = "false"

using BinaryBuilder

name = "XRTServer_lite"
version = v"2018.11.01"

# Collection of sources required
sources = [
    "https://github.com/JuliaComputing/tensorflow.git" =>
    "36c34bd2d744f0027c24e3afae118b5d956ce741",
    "https://github.com/bazelbuild/bazel/releases/download/0.17.2/bazel-0.17.2-linux-x86_64" =>
    "674757d40d4ac0f0175df7fe84cd7250cbf67ac7ebac565e905fdc7e24c0fac5",
    "./bundled",
]

# Bash recipe for building across all platforms
script = raw"""
# bazel always looks for `ar` in `/usr/bin/ar`
ln -sf "$AR" /usr/bin/ar

# Apply -lrt patch
cd $WORKSPACE/srcdir/tensorflow
atomic_patch -p0 $WORKSPACE/srcdir/patches/link_against_librt.patch

# Apply XRT patch
atomic_patch -p1 $WORKSPACE/srcdir/patches/xrt.patch

# Apply pending TF patches
atomic_patch -p1 $WORKSPACE/srcdir/patches/window_reversal.patch
atomic_patch -p1 $WORKSPACE/srcdir/patches/parallel_conditionals.patch

# Get `bazel` onto our $PATH
chmod +x $WORKSPACE/srcdir/bazel-*
mv $WORKSPACE/srcdir/bazel-* $WORKSPACE/srcdir/bazel
export PATH=$PATH:$WORKSPACE/srcdir
export TMPDIR=$WORKSPACE/tmp
mkdir -p $TMPDIR

# Configure to enable CUDA and XLA, disable jemalloc
export TF_ENABLE_XLA=1
export TF_NEED_JEMALLOC=0
yes "" | ./configure

# Do the actual build for our tools
bazel --output_user_root=/workspace/bazel_root build -c opt --verbose_failures //tensorflow:libtensorflow.so
bazel --output_user_root=/workspace/bazel_root build -c opt --verbose_failures //tensorflow/compiler/xrt/utils:xrt_server
bazel --output_user_root=/workspace/bazel_root build -c opt --verbose_failures //tensorflow/compiler/xla/tools:dumped_computation_to_text

# Install to $prefix/bin
mkdir -p $prefix/{bin,lib}
cp bazel-bin/tensorflow/libtensorflow.so $prefix/lib/
cp bazel-bin/tensorflow/libtensorflow_framework.so $prefix/lib/
cp bazel-bin/tensorflow/compiler/xrt/utils/xrt_server $prefix/bin/
cp bazel-bin/tensorflow/compiler/xla/tools/dumped_computation_to_text $prefix/bin/

# libtensorflow* has super-crazy RPATHs; fix that
for f in ${prefix}/lib/libtensorflow* ${prefix}/bin/*; do
    patchelf --set-rpath '$ORIGIN:$ORIGIN/../lib:$ORIGIN/../lib64' "${f}"
done
"""

# We attempt to build for only x86_64-linux-gnu
platforms = [Linux(:x86_64)]

products(prefix) = [
    ExecutableProduct(prefix, "xrt_server", :xrt_server),
    ExecutableProduct(prefix, "dumped_computation_to_text", :dumped_computation_to_text),
]

dependencies = [
    "https://github.com/bicycle1885/ZlibBuilder/releases/download/v1.0.2/build_Zlib.v1.2.11.jl",
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)

ext: jl


text:
#Running DP clustering should find 4 or 5 clusters, the 3 highest frequency clusters are
freqs = [0.362204, 0.202305, 0.111506]

using Random
Random.seed!(123)

using DelimitedFiles
data = readdlm("data.csv", ',', header = true)
y = data[1][:, 1]
N = data[1][:, 2]

@time out = dpclustering(y, N)

@test isapprox(sort(out.clonefrequencies, rev = true)[1:3], freqs, atol = 0.01)

#check no errors from plotting
plotresults(out; save = true)
@test isfile("DPClustering.pdf")
rm("DPClustering.pdf")

#check show function works
show(out)


text:
export bfgs_bl
export bfgs_rc
export bfgsH

"""
Modelos básicos...
bfgs_bl (busca linear)
bfgs_rc (região de confiança - Steihaug Toint)

Options:
- atol: absolute tolerance for the first order condition (default: 1e-6)
- rtol: relative tolerance for the first order condition (default: 1e-6)
- max_eval: maximum number of [functions] evaluations, use ≤ 0 for unlimited (default: 1000)
- max_iter: maximum number of iterations, use ≤ 0 for unlimited (default: 0)
- max_time: maximum elapsed time in seconds, use ≤ 0 for unlimited (default: 10)

Disclaimers for the developer:
  - nlp should be the only mandatory argument
  - these five options are the current default for other JSO-compliant solvers
  - always return a GenericExecutionStats
  """
  function bfgs_bl(
    nlp::AbstractNLPModel;
    atol::Real = 1e-6,
    rtol::Real = 1e-6,
    max_eval::Int = 1000,
    max_iter::Int = 0,
    max_time::Float64 = 10.0
    )

    if !unconstrained(nlp)
      error("Problem is not unconstrained")
    end

    x = copy(nlp.meta.x0)

    f(x) = obj(nlp, x)
    ∇f(x) = grad(nlp, x)

    fx = f(x)
    ∇fx = ∇f(x)
    n = length(x)
    Hx = Matrix(1.0I, n, n)

    ϵ = atol + rtol * norm(∇fx)
    t₀ = time()

    iter = 0
    Δt = time() - t₀
    solved = norm(∇fx) < ϵ # First order stationary
    tired = neval_obj(nlp) ≥ max_eval > 0 ||
    iter ≥ max_iter > 0 ||
    Δt ≥ max_time > 0
    # Excess time, iteration, evaluations

    # status must be one of a few options found in SolverTools.show_statuses()
    # A good default value is :unknown.
    status = :unknown

    # log_header is up for some rewrite in the future. For now, it simply prints the column names with some spacing
    @info log_header(
    [:iter, :fx, :ngx, :nf, :Δt],
    [Int, Float64, Float64, Int, Float64],
    hdr_override=Dict(:fx => "f(x)", :ngx => "‖∇f(x)‖", :nf => "#f")
    )
    # log_row uses the type information of each value, thus we use `Any` here.
    @info log_row(
    Any[iter, fx, norm(∇fx), neval_obj(nlp), Δt]
    )

    # Aqui começa o show

    while !(solved || tired)
      α = 1.0
      η = 1e-2

      d = -Hx * ∇fx
      slope = dot(d, ∇fx)

      # Armijo
      x⁺ = x + α * d
      f⁺ = f(x⁺)
      while f⁺ ≥ fx + η * α * slope
        α = α / 2
        x⁺ = x + α * d
        f⁺ = f(x⁺)
        if α < 1e-8
          status =:small_step
          break
        end
      end
      if status != :unknown #small_step
        break
      end

      s = α * d
      y = ∇f(x⁺) - ∇fx

      if dot(s,y) <= 0
        @warn("sᵀy = $(dot(s,y))")
      else
        Hx = bfgsH(Hx, s, y)
      end

      x .= x⁺
      fx = f⁺
      ∇fx = ∇f(x)

      iter += 1
      Δt = time() - t₀
      solved = norm(∇fx) < ϵ # First order stationary
      tired = neval_obj(nlp) ≥ max_eval > 0|| iter ≥ max_iter > 0 || Δt ≥ max_time > 0 # Excess time, iteration, evaluations

      @info log_row(
      Any[iter, fx, norm(∇fx), neval_obj(nlp), Δt]
      )
    end

    if solved
      status = :first_order
    elseif tired
      if neval_obj(nlp) ≥ max_eval > 0
        status = :max_eval
      elseif iter ≥ max_iter > 0
        status = :max_iter
      elseif Δt ≥ max_time > 0
        status = :max_time
      end
    end

    return GenericExecutionStats(
    status,
    nlp,
    solution=x,
    objective=f(x),
    dual_feas=norm(∇fx),
    elapsed_time=Δt,
    iter=iter
    )
  end

  function bfgsH(H, s, y)
    ρ = 1 / dot(s, y)
    H = (I - ρ * s * y') * H * (I - ρ * y * s') + ρ * s * s'
    return H
  end

  #REGIÃO DE CONFIANÇA
  function bfgs_rc(
    nlp::AbstractNLPModel;
    atol::Real = 1e-6,
    rtol::Real = 1e-6,
    max_eval::Int = 50000,
    max_iter::Int = 0,
    max_time::Float64 = 10.0
    )

    if !unconstrained(nlp)
      error("Problem is not unconstrained")
    end

    #Given the starting point
    x = copy(nlp.meta.x0)
    f(x) = obj(nlp, x)
    ∇f(x) = grad(nlp, x)

    # initial Hessian approximation
    n = length(x)

    #trust region radius
    Δ = 1.0

    fx = f(x)
    ∇fx = ∇f(x)

    #convergence tolerance
    ϵ = atol + rtol * norm(∇fx)
    #parameters η e r
    η = 0.01 #∈ (0,1e-3)
    r = 0.001 #∈ (0,1)

    t₀ = time()

    iter = 0
    Δt = time() - t₀
    solved = norm(∇fx) < ϵ # First order stationary
    tired = neval_obj(nlp) ≥ max_eval > 0 ||
    iter ≥ max_iter > 0 ||
    Δt ≥ max_time > 0
    # Excess time, iteration, evaluations
    # status must be one of a few options found in SolverTools.show_statuses()
    # A good default value is :unknown.
    status = :unknown
    # log_header is up for some rewrite in the future. For now, it simply prints the column names with some spacing
    @info log_header(
    [:iter, :fx, :ngx, :nf, :Δt],
    [Int, Float64, Float64, Int, Float64],
    hdr_override=Dict(:fx => "f(x)", :ngx => "‖∇f(x)‖", :nf => "#f")
    )
    # log_row uses the type information of each value, thus we use `Any` here.
    @info log_row(
    Any[iter, fx, norm(∇fx), neval_obj(nlp), Δt]
    )

    # Aqui começa o show
    #determinar o primero B
    # B⁰ₖ = δₖI, wher δₖ = dot(y,y)/dot(s,y) Nocedal pg178 (eq.7.20)
    B = Matrix(1.0I, n, n)
    s = Steighaug(∇fx, B, Δ)
    y = ∇f(x.+s) - ∇fx
    γ = dot(y,y)/dot(s,y)
    if γ >= 1.0
      B = Matrix(γ*I, n, n)
    end

    while !(solved || tired)
      #compute sₖ by solving the subproblem
      # (6.27) aqui vai entrar o Steihaug
      s = Steighaug(∇fx, B, Δ)
      y = ∇f(x.+s) - ∇fx
      ared = fx - f(x.+s)
      pred = -(dot(∇fx, s) + 1/2 * dot(s, B*s))
      ρ = ared/pred
      if ρ < η
        Δ = Δ/2
        if Δ < 10e-50
          @error("Δ muito pequeno")
          status =:small_step
        end
      else
        x = x + s
        fx = f(x)
        ∇fx = ∇f(x)
        if ρ > 0.75 && norm(s) > 0.8 * Δ
          Δ = 2*Δ
          if Δ > 10e50
            @error("Δ muito grande")
            status =:user
          end
        end
      end

      if status != :unknown #small_step
        break
      end
      yBs = y .-B*s
      syBs = dot(s, yBs)
      if abs(syBs) > r*norm(s,2)*norm(yBs,2) #era maior ou igual, deixei maior
        # em alguns casos, y - B*s pode ser zero... e nesse caso a atualização
        #de B será uma divisão por zero.

        #why?!?!
        #6.26 - evita atualizações se o denominador é pequeno
        #Simmetric-rank 1 method: prevent breakdown; better approximations
        # to the true Hessian matrix;
        B = B + (yBs*yBs')/syBs
      end

      iter+= 1

      Δt = time() - t₀
      solved = norm(∇fx) < ϵ # First order stationary
      tired = neval_obj(nlp) ≥ max_eval > 0|| iter ≥ max_iter > 0 ||
                          Δt ≥ max_time > 0 # Excess time, iteration, evaluations

      @info log_row(
      Any[iter, fx, norm(∇fx), neval_obj(nlp), Δt]
      )
    end

    if solved
      status = :first_order
    elseif tired
      if neval_obj(nlp) ≥ max_eval > 0
        status = :max_eval
      elseif iter ≥ max_iter > 0
        status = :max_iter
      elseif Δt ≥ max_time > 0
        status = :max_time
      end
    end

    return GenericExecutionStats(
    status,
    nlp,
    solution=x,
    objective=f(x),
    dual_feas=norm(∇fx),
    elapsed_time=Δt,
    iter=iter
    )
  end

  function Steighaug(gx, B, Δ; ϵ = 1.0e-4)
    m = length(gx)
    zx = zeros(m)
    r = gx
    d = -r

    rx = r
    z = zx
    normr = norm(r)
    if normr < ϵ
      return z
    end
    # enquanto estiver dentro da região de confiança
    k=0
    while normr > ϵ || k < m # CG vai em no máximo m-variáveis direções
      dotdBd = dot(d,B*d)
      @info("dotdBd: $dotdBd")
      if dotdBd ≤ 0
        # para o método se a direção dⱼ é direção de curvatura não positiva
        # encontrar τ >=0 tal que pk e satisfaz ||pk|| = Δk
        # ou seja, a interseção da direção com a região de confiança
        m1, m2 = BhaskaraTop(z, d, Δ)
        return z + m2*d
      end

      dotrr = dot(r,r)
      α = dotrr/dotdBd
      @info("α: $α")
      zx = z + α*d
      if norm(zx) ≥ Δ
        #para se zⱼ₊₁ viola os limites da região de confiança
        # encontrar τ >=0 tal que pk e satisfaz ||pk|| = Δk
        # ou seja, a interseção da direção com a região de confiança
        m1, m2 = BhaskaraTop(z, d, Δ)
        return z + m1*d
      end

      rx = r + α*B*d #Conjugate Gradient
      #se α = 0 => rx = r , β = 1
      if norm(α) < ϵ^2
        #@error("α menor que zero")
        return zeros(m)
      end

      if norm(rx) < ϵ
        return zx
      end
      β = dot(rx,rx)/dotrr
      d = -rx+ β*d
      r = rx
      normr = norm(r)
      z = zx
      k+=1
    end
    if k >= m
      #@error ("não encontrou a borda, direção multipla da própria direção")
      return zeros(m)
    end

  end

  function BhaskaraTop(z, d, Δ)
    a=dot(d,d)
    b=2dot(z,d)
    c=dot(z,z) - Δ^2
    Delta = b^2-4*a*c
    if Delta < 0
      @warn("Δ<0")
    else
      t1 = (-b+sqrt(Delta))/2a
      t2 = (-b-sqrt(Delta))/2a
      m1 = max(t1,t2)
      m2 = min(t1,t2)
      return m1, m2
    end
  end


text:
@testset "ZBL Unit Tests" begin
    Z₁ = 18
    Z₂ = 18
    e = 0.25u"C"
    rcutoff = 2.0u"Å"
    species = [:Ar, :H]
    p = ZBL(Z₁, Z₂, e, rcutoff, species)

    @test p isa EmpiricalPotential{NamedTuple{()},NamedTuple{(:rcutoff,)}}

    @test get_rcutoff(p) == austrip(rcutoff)
    @test get_species(p) == (:Ar, :H)

    @test get_parameters(p) == (;)
    @test set_parameters(p, (;)) === p
    @test serialize_parameters(p) == []
    @test deserialize_parameters(p, []) === p

    @test get_hyperparameters(p) == (; rcutoff=austrip(rcutoff))
    @test set_hyperparameters(p, (; rcutoff=1.0)) == ZBL(Z₁, Z₂, austrip(e), 1.0, species)
    @test serialize_hyperparameters(p) == [austrip(rcutoff)]
    @test deserialize_hyperparameters(p, [1.0]) == ZBL(Z₁, Z₂, austrip(e), 1.0, species)

    r = @SVector[1.0, 1.0, 1.0]
    R = norm(r)

    @test potential_energy(R, p) isa Float64
    @test force(R, r, p) isa SVector{3,Float64}
end


text:
function read_input(file_path::String)
	file_stream::IOStream = open(file_path, "r")
	file_string::String = readstring(file_stream)
	close(file_stream)

	# Read particle data.
	particle_type::String = read_key(file_string, "particle_type", String)
	R_temp::Array{Float64, 1} = read_key(file_string, "R", Array{Float64, 1})
	number_of_properties::Int64 = 0
	if particle_type == "sphere"
		number_of_properties = 1
	elseif particle_type == "ellipse"
		number_of_properties = 2
	elseif particle_type == "ellipsoid"
		number_of_properties = 3
	elseif particle_type == "cuboid"
		number_of_properties = 3
	end
	number_of_particles::Int64 = length(R_temp) / number_of_properties
	R::Array{Float64, 2} = reshape(R_temp, number_of_particles, number_of_properties)

	# Read input on initial system size.
	Lx::Float64 = read_key(file_string, "domain_size_x", Float64)
	Ly::Float64 = read_key(file_string, "domain_size_y", Float64)
	Lz::Float64 = read_key(file_string, "domain_size_z", Float64)
	phi::Float64 = read_key(file_string, "phi", Float64)

	# Read input on target system size.
	phi_target::Float64 = read_key(file_string, "phi_target", Float64)

	# Position constraint.
	position_constraint_axis::String = read_key(file_string, "position_constraint_axis", String)
	position_constraint_lower::Float64 = read_key(file_string, "position_constraint_lower", Float64)
	position_constraint_upper::Float64 = read_key(file_string, "position_constraint_upper", Float64)

	# Orientation constraint.
	orientation_axis::Array{Float64, 1} = read_key(file_string, "orientation_axis", Array{Float64, 1})
	orientation_constraint_axis::Array{Float64, 1} = read_key(file_string, "orientation_constraint_axis", Array{Float64, 1})
	orientation_constraint_lower::Float64 = read_key(file_string, "orientation_constraint_lower", Float64)
	orientation_constraint_upper::Float64 = read_key(file_string, "orientation_constraint_upper", Float64)

	# Other input.
	sigma_translation_max::Float64 = read_key(file_string, "sigma_translation_max", Float64)
	sigma_rotation_max::Float64 = read_key(file_string, "sigma_rotation_max", Float64)
	sigma_ratio::Float64 = read_key(file_string, "sigma_ratio", Float64)
	number_of_equilibration_sweeps::Int64 = read_key(file_string, "number_of_equilibration_sweeps", Int64)
	delta_phi::Float64 = read_key(file_string, "delta_phi", Float64)

	output_file_path::String = read_key(file_string, "output_file_path", String)

	return (
		particle_type,
		R,
		Lx,
		Ly,
		Lz,
		phi,
		phi_target,
		position_constraint_axis,
		position_constraint_lower,
		position_constraint_upper,
		orientation_axis,
		orientation_constraint_axis,
		orientation_constraint_lower,
		orientation_constraint_upper,
		sigma_translation_max,
		sigma_rotation_max,
		sigma_ratio,
		number_of_equilibration_sweeps,
		delta_phi,
		output_file_path)
end

max_stars_repo_name: roding/whitefish

max_issues_repo_name: roding/whitefish

max_forks_repo_name: roding/whitefish


text:
module ALBERT
using Flux
using Requires
using Requires: @init
using BSON
using Transformers
using Transformers.Basic
using Transformers.Pretrain: isbson, iszip, istfbson, zipname, zipfile, findfile

export ALBERT
export load_albert_pretrain, albert_pretrain_task, masklmloss,WordPiece,tokenise

#implement batchmul, batchtril for flux

include("tfckpt2bsonforalbert.jl")
include("alberttokenizer.jl")
include("sentencepiece.jl")
include("albert.jl")

end # module


text:
# Let d(n) be defined as the sum of proper divisors of n (numbers less than n
# which divide evenly into n).  If d(a) = b and d(b) = a, where a ≠ b, then a
# and b are an amicable pair and each of a and b are called amicable numbers.
#
# For example, the proper divisors of 220 are 1, 2, 4, 5, 10, 11, 20, 22, 44,
# 55 and 110; therefore d(220) = 284. The proper divisors of 284 are 1, 2, 4,
# 71 and 142; so d(284) = 220.
#
# Evaluate the sum of all the amicable numbers under 10000.

using ProjectEulerSolutions

# Actually build up an array of factors, testing for amicable numbers smaller
# than the current index.  Building up is done with a double loop, adding like
# below.
#
# 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16
#    1  1  1  1  1  1  1  1  1  1  1  1  1  1  1  1
#             2     2     2     2     2     2     2
#                   3        3        3        3
#                         4           4           4
#                               5              5
#                                     6
#                                           7
#                                                 8
function p021solution_incremental(n::Integer=10)::Integer
    amicable_numbers = Array{Integer, 1}()

    factor_sums = ones(Integer, n)
    for i = 2:n
        for j = 2*i:i:n
            factor_sums[j] += i
        end
        x = factor_sums[i]
        if x < i && i == factor_sums[x]
            push!(amicable_numbers, x)
            push!(amicable_numbers, i)
        end
    end
    return sum(amicable_numbers)
end

# Repeatedly uses factors function and cache values so you don't
# recalculate.  Store in a set so you get unique values.
function p021solution_factors(n::Integer=10)::Integer
    factor_sums = zeros(Integer, 5*n) # Initialize to some large size

    amicable_numbers = Set{Integer}()

    for i in 2:n
        s1 = factor_sums[i]
        if s1 == 0
            s1 = sum(factors(i))
            factor_sums[i] = s1
        end
        s2 = factor_sums[s1]
        if s2 == 0
            s2 = sum(factors(s1))
            factor_sums[s1] = s2
        end
        if s2 == i && s1 != s2
            push!(amicable_numbers, s1)
            push!(amicable_numbers, s2)
        end
    end
    return sum(amicable_numbers)
end

p021 = Problems.Problem(Dict("Incremental" => p021solution_incremental,
                             "Factors" => p021solution_factors))

Problems.benchmark(p021, 10_000)

max_stars_repo_name: gnujosh/julia-euler-project

max_issues_repo_name: gnujosh/julia-euler-project

max_forks_repo_name: gnujosh/julia-euler-project


text:
# Test if used
@testset "Used" begin
    # initialize model, parameter, and variable
    m = InfiniteModel()
    info = VariableInfo(false, 0, false, 0, false, 0, false, 0, false, false)
    v = build_variable(error, info, Hold)
    vref = add_variable(m, v, "name")
    # used_by_constraint
    @testset "used_by_constraint" begin
        # test not used
        @test !used_by_constraint(vref)
        # prepare use case
        m.var_to_constrs[JuMP.index(vref)] = [1]
        # test used
        @test used_by_constraint(vref)
        # undo changes
        delete!(m.var_to_constrs, JuMP.index(vref))
    end
    # used_by_measure
    @testset "used_by_measure" begin
        # test not used
        @test !used_by_measure(vref)
        # prepare use case
        m.var_to_meas[JuMP.index(vref)] = [1]
        # test used
        @test used_by_measure(vref)
        # undo changes
        delete!(m.var_to_meas, JuMP.index(vref))
    end
    # used_by_objective
    @testset "used_by_objective" begin
        # test not used
        @test !used_by_objective(vref)
        # prepare use case
        m.var_in_objective[JuMP.index(vref)] = true
        # test used
        @test used_by_objective(vref)
        # undo changes
        m.var_in_objective[JuMP.index(vref)] = false
    end
    # is_used
    @testset "is_used" begin
        # test not used
        @test !is_used(vref)
        # prepare use case
        m.var_to_constrs[JuMP.index(vref)] = [1]
        # test used
        @test is_used(vref)
        # undo changes
        delete!(m.var_to_constrs, JuMP.index(vref))
        # prepare use case
        m.var_to_meas[JuMP.index(vref)] = [1]
        # test used
        @test is_used(vref)
        # undo changes
        delete!(m.var_to_meas, JuMP.index(vref))
        # prepare use case
        m.var_in_objective[JuMP.index(vref)] = true
        # test used
        @test is_used(vref)
        # undo changes
        m.var_in_objective[JuMP.index(vref)] = false
    end
    # initialize infinite variable to test infinite methods
    param = InfOptParameter(IntervalSet(0, 1), Number[], false)
    pref = add_parameter(m, param, "test")
    v = build_variable(error, info, Infinite, parameter_refs = (pref, ))
    vref = add_variable(m, v, "name")
    # used_by_point_variable
    @testset "used_by_point_variable" begin
        # test not used
        @test !used_by_point_variable(vref)
        # build point variable
        v = build_variable(error, info, Point, infinite_variable_ref = vref,
                           parameter_values = 0.5)
        pvref = add_variable(m, v, "name2")
        # test used
        @test used_by_point_variable(vref)
    end
    # used_by_reduced_variable
    @testset "used_by_reduced_variable" begin
        # test not used
        @test !used_by_reduced_variable(vref)
        # build reduced variable
        m.infinite_to_reduced[JuMP.index(vref)] = [-1]
        # test used
        @test used_by_reduced_variable(vref)
    end
    # is_used (infinite variable)
    @testset "is_used (infinite var)" begin
        # test not used
        @test !is_used(vref)
        # prepare use case
        m.var_to_constrs[JuMP.index(vref)] = [1]
        # test used
        @test is_used(vref)
        # undo changes
        delete!(m.var_to_constrs, JuMP.index(vref))
        # prepare use case
        m.var_to_meas[JuMP.index(vref)] = [1]
        # test used
        @test is_used(vref)
        # undo changes
        delete!(m.var_to_meas, JuMP.index(vref))
        # prepare use case
        m.var_in_objective[m.next_var_index] = true
        # test used
        @test is_used(vref)
        # undo changes
        m.var_in_objective[m.next_var_index] = false
        # prepare use case
        m.reduced_to_constrs[-1] = [1]
        # test used
        @test is_used(vref)
        # undo changes
        delete!(m.reduced_to_constrs, -1)
        # prepare use case
        m.reduced_to_meas[-1] = [1]
        # test used
        @test is_used(vref)
        @test m.infinite_to_reduced[JuMP.index(vref)] == [-1]
    end
end

# Test variable queries for models
@testset "Model Queries" begin
    # initialize model, parameter, and variables
    m = InfiniteModel()
    param = InfOptParameter(IntervalSet(0, 1), Number[], false)
    pref = add_parameter(m, param, "test")
    info = VariableInfo(false, 0, false, 0, false, 0, false, 0, false, false)
    v = build_variable(error, info, Infinite, parameter_refs = (pref, ))
    ivref = add_variable(m, v, "name")
    v = build_variable(error, info, Point, infinite_variable_ref = ivref,
                       parameter_values = 0.5)
    pvref = add_variable(m, v, "name2")
    v = build_variable(error, info, Hold)
    hvref = add_variable(m, v, "name")
    # num_variables
    @testset "JuMP.num_variables" begin
        @test num_variables(m) == 3
    end
    # all_variables
    @testset "JuMP.all_variables" begin
        @test all_variables(m) == [ivref, pvref, hvref]
    end
end

# Test queries for parameter references and values
@testset "Parameter Tuples" begin
    # initialize model, parameter, and variables
    m = InfiniteModel()
    param = InfOptParameter(IntervalSet(0, 1), Number[], false)
    pref = add_parameter(m, param, "test")
    pref2 = add_parameter(m, param, "θ")
    info = VariableInfo(false, 0, false, 0, false, 0, false, 0, false, false)
    v = build_variable(error, info, Infinite, parameter_refs = (pref, ))
    ivref = add_variable(m, v, "name")
    v = build_variable(error, info, Point, infinite_variable_ref = ivref,
                       parameter_values = 0.5)
    pvref = add_variable(m, v, "name2")
    # parameter_values
    @testset "parameter_values" begin
        @test parameter_values(pvref) == (0.5, )
    end
    # _update_variable_param_values
    @testset "_update_variable_param_values" begin
        @test isa(InfiniteOpt._update_variable_param_values(pvref, (0, )),
                  Nothing)
        @test parameter_values(pvref) == (0, )
    end
    # _update_variable_param_refs
    @testset "_update_variable_param_refs" begin
        @test isa(InfiniteOpt._update_variable_param_refs(ivref, (pref2, )),
                  Nothing)
        @test parameter_refs(ivref) == (pref2, )
    end
    # set_parameter_refs
    @testset "set_parameter_refs" begin
        # test normal with 1 param
        @test isa(set_parameter_refs(ivref, (pref2, )), Nothing)
        @test parameter_refs(ivref) == (pref2, )
        @test name(ivref) == "name(θ)"
        # test double specify
        @test_throws ErrorException set_parameter_refs(ivref, (pref2, pref2))
        # test with used variable
        m.var_to_meas[JuMP.index(ivref)] = [1]
        @test isa(set_parameter_refs(ivref, (pref, )), Nothing)
        @test parameter_refs(ivref) == (pref, )
        @test name(ivref) == "name(test)"
        delete!(m.var_to_meas, JuMP.index(ivref))
    end
    # add_parameter_ref
    @testset "add_parameter_ref" begin
        # test normal use
        @test isa(add_parameter_ref(ivref, pref2), Nothing)
        @test parameter_refs(ivref) == (pref, pref2)
        @test name(ivref) == "name(test, θ)"
        # test duplication error
        @test_throws ErrorException add_parameter_ref(ivref, pref2)
    end
end

# Test parameter bound queries for hold variables
@testset "Parameter Bounds" begin
    # initialize the model and other needed information
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 10])
    @infinite_parameter(m, pars[1:2] in [0, 10], container = SparseAxisArray)
    @infinite_variable(m, inf(par))
    @hold_variable(m, x)
    @hold_variable(m, y, parameter_bounds = par == 0)
    data = DiscreteMeasureData(par, [1], [1])
    meas = measure(inf + par - x, data)
    # test parameter_bounds
    @testset "parameter_bounds" begin
        @test parameter_bounds(x) == ParameterBounds()
        @test parameter_bounds(y) == ParameterBounds(Dict(par => IntervalSet(0, 0)))
    end
    # test has_parameter_bounds
    @testset "has_parameter_bounds" begin
        @test !has_parameter_bounds(x)
        @test has_parameter_bounds(y)
        @test !has_parameter_bounds(InfiniteVariableRef(m, -2))
    end
    # test _update_variable_param_bounds
    @testset "_update_variable_param_bounds" begin
        dict = Dict(par => IntervalSet(0, 1))
        bounds = ParameterBounds(dict)
        @test isa(InfiniteOpt._update_variable_param_bounds(y, bounds), Nothing)
        @test parameter_bounds(y) == bounds
    end
    # test _check_meas_bounds (DiscreteMeasureData)
    @testset "_check_meas_bounds (DiscreteMeasureData)" begin
        data = DiscreteMeasureData(par, [1, 1], [1, 3], "", InfiniteOpt._w)
        # test ok
        bounds = ParameterBounds(Dict(par => IntervalSet(0, 5)))
        @test isa(InfiniteOpt._check_meas_bounds(bounds, data), Nothing)
        # test lower bad
        bounds = ParameterBounds(Dict(par => IntervalSet(2, 5)))
        @test_throws ErrorException InfiniteOpt._check_meas_bounds(bounds, data)
        # test upper bad
        bounds = ParameterBounds(Dict(par => IntervalSet(0, 2)))
        @test_throws ErrorException InfiniteOpt._check_meas_bounds(bounds, data)
    end
    # test _check_meas_bounds (MultiDiscreteMeasureData)
    @testset "_check_meas_bounds (Multi)" begin
        supps = [convert(JuMPC.SparseAxisArray, [1, 1]),
                 convert(JuMPC.SparseAxisArray, [3, 3])]
        data = MultiDiscreteMeasureData(pars, [0, 0], supps, "", InfiniteOpt._w)
        # test ok
        bounds = ParameterBounds(Dict(pars[1] => IntervalSet(0, 5)))
        @test isa(InfiniteOpt._check_meas_bounds(bounds, data), Nothing)
        # test lower bad
        bounds = ParameterBounds(Dict(pars[1] => IntervalSet(2, 5)))
        @test_throws ErrorException InfiniteOpt._check_meas_bounds(bounds, data)
        # test upper bad
        bounds = ParameterBounds(Dict(pars[2] => IntervalSet(0, 2)))
        @test_throws ErrorException InfiniteOpt._check_meas_bounds(bounds, data)
    end
    # test _check_meas_bounds (Fallback)
    @testset "_check_meas_bounds (Fallback)" begin
        warn = "Unable to check if hold variables bounds are valid in measure " *
               "with measure data type `BadData`. This can be resolved by " *
               "extending `measure_data_in_hold_bounds`."
        @test_logs (:warn, warn) InfiniteOpt._check_meas_bounds(ParameterBounds(),
                                                                BadData())
    end
    # test _update_bounds
    @testset "_update_bounds" begin
        # test normal
        dict1 = Dict(pars[1] => IntervalSet(0, 1), pars[2] => IntervalSet(-1, 2))
        dict2 = Dict(par => IntervalSet(0, 0), pars[1] => IntervalSet(0.5, 1.5),
                     pars[2] => IntervalSet(0, 1))
        @test isa(InfiniteOpt._update_bounds(dict1, dict2), Nothing)
        @test length(dict1) == 3
        @test dict1[par] == IntervalSet(0, 0)
        @test dict1[pars[1]] == IntervalSet(0.5, 1)
        @test dict1[pars[2]] == IntervalSet(0, 1)
        # test error
        new_dict = Dict(par => IntervalSet(1, 1))
        @test_throws ErrorException InfiniteOpt._update_bounds(dict1, new_dict)
        new_dict = Dict(par => IntervalSet(-1, -1))
        @test_throws ErrorException InfiniteOpt._update_bounds(dict1, new_dict)
    end
    # test _update_var_bounds (GeneralVariableRef)
    @testset "_update_var_bounds (General)" begin
        @test isa(InfiniteOpt._update_var_bounds(inf, ParameterBounds()),
                  Nothing)
    end
    # test _update_var_bounds (HoldVariableRef)
    @testset "_update_var_bounds (Hold)" begin
        bounds = ParameterBounds()
        InfiniteOpt._update_variable_param_bounds(x,
                                ParameterBounds(Dict(par => IntervalSet(1, 1))))
        @test isa(InfiniteOpt._update_var_bounds(x, bounds), Nothing)
        @test bounds.intervals[par] == IntervalSet(1, 1)
    end
    # test _update_var_bounds (MeasureRef)
    @testset "_update_var_bounds (Measure)" begin
        bounds = ParameterBounds()
        @test isa(InfiniteOpt._update_var_bounds(meas, bounds), Nothing)
        @test bounds.intervals[par] == IntervalSet(1, 1)
        InfiniteOpt._update_variable_param_bounds(x, ParameterBounds())
    end
    # test _rebuild_constr_bounds (BoundedScalarConstraint)
    @testset "_rebuild_constr_bounds (Bounded)" begin
        # test addition
        bounds = ParameterBounds(Dict(par => IntervalSet(0, 2)))
        c = BoundedScalarConstraint(y + inf, MOI.EqualTo(0.), copy(bounds),
                                    copy(bounds))
        expected = ParameterBounds(Dict(par => IntervalSet(0, 1)))
        @test InfiniteOpt._rebuild_constr_bounds(c, bounds).bounds == expected
        # test deletion
        bounds = ParameterBounds(Dict(par => IntervalSet(0, 0)))
        c = BoundedScalarConstraint(y + inf, MOI.EqualTo(0.), bounds,
                                    ParameterBounds())
        InfiniteOpt._update_variable_param_bounds(y, ParameterBounds())
        @test isa(InfiniteOpt._rebuild_constr_bounds(c, bounds),
                  ScalarConstraint)
        InfiniteOpt._update_variable_param_bounds(y, bounds)
    end
    # test _rebuild_constr_bounds (ScalarConstraint)
    @testset "_rebuild_constr_bounds (Scalar)" begin
        bounds = ParameterBounds(Dict(par => IntervalSet(0, 0)))
        c = ScalarConstraint(y + inf, MOI.EqualTo(0.))
        @test InfiniteOpt._rebuild_constr_bounds(c, bounds).bounds == bounds
    end
    # test set_parameter_bounds
    @testset "set_parameter_bounds" begin
        bounds1 = ParameterBounds(Dict(pars[1] => IntervalSet(0, 1)))
        # test force error
        @test_throws ErrorException set_parameter_bounds(y, bounds1)
        # prepare for main test
        data = DiscreteMeasureData(par, [1, 1], [1, 3], "", InfiniteOpt._w)
        mindex = -1; cindex1 = 42; cindex2 = -42; vindex = JuMP.index(x);
        m.measures[mindex] = Measure(x, data)
        m.var_to_meas[JuMP.index(x)] = [mindex]
        m.meas_to_constrs[mindex] = [cindex1]
        mref = MeasureRef(m, mindex)
        m.constrs[cindex1] = BoundedScalarConstraint(mref, MOI.EqualTo(0.0),
                                                     bounds1, bounds1)
        m.var_to_constrs[vindex] = [cindex2]
        m.constrs[cindex2] = JuMP.ScalarConstraint(x, MOI.EqualTo(0.0))
        # test measure error
        bounds = ParameterBounds(Dict(par => IntervalSet(1, 2)))
        @test_throws ErrorException set_parameter_bounds(x, bounds)
        # test constr error
        bounds = ParameterBounds(Dict(pars[1] => IntervalSet(2, 3)))
        @test_throws ErrorException set_parameter_bounds(x, bounds)
        InfiniteOpt._update_variable_param_bounds(x, ParameterBounds())
        # test normal
        bounds = ParameterBounds(Dict(par => IntervalSet(0, 5),
                                      pars[1] => IntervalSet(0.5, 2)))
        @test isa(set_parameter_bounds(x, bounds), Nothing)
        @test parameter_bounds(x) == bounds
        @test m.has_hold_bounds
        @test !optimizer_model_ready(m)
        expected = ParameterBounds(Dict(par => IntervalSet(0, 5),
                                        pars[1] => IntervalSet(0.5, 1)))
        @test m.constrs[cindex1].bounds == expected
        @test m.constrs[cindex2].bounds == bounds
        # test forced
        bounds = ParameterBounds(Dict(par => IntervalSet(1, 1)))
        @test isa(set_parameter_bounds(y, bounds, force = true), Nothing)
        @test parameter_bounds(y) == ParameterBounds(Dict(par => IntervalSet(1, 1)))
    end
    # test _update_constr_bounds (BoundedScalarConstraint)
    @testset "_update_constr_bounds (Bounded)" begin
        dict1 = Dict(pars[1] => IntervalSet(0, 1))
        dict2 = Dict(par => IntervalSet(0, 0), pars[1] => IntervalSet(1.1, 1.5))
        constr = BoundedScalarConstraint(par, MOI.EqualTo(0.0), ParameterBounds(copy(dict1)),
                                         ParameterBounds(copy(dict1)))
        # test error
        @test_throws ErrorException InfiniteOpt._update_constr_bounds(ParameterBounds(dict2), constr)
        @test constr.bounds == ParameterBounds(dict1)
        # test normal
        dict2 = Dict(par => IntervalSet(0, 0), pars[1] => IntervalSet(0.5, 1.5))
        dict = Dict(par => IntervalSet(0, 0), pars[1] => IntervalSet(0.5, 1))
        @test InfiniteOpt._update_constr_bounds(ParameterBounds(dict2), constr).bounds == ParameterBounds(dict)
    end
    # test _update_constr_bounds (ScalarConstraint)
    @testset "_update_constr_bounds (Scalar)" begin
        bounds = ParameterBounds(Dict(pars[1] => IntervalSet(0, 1)))
        constr = JuMP.ScalarConstraint(par, MOI.EqualTo(0.0))
        @test InfiniteOpt._update_constr_bounds(bounds, constr).bounds == bounds
    end
    # test add_parameter_bound
    @testset "add_parameter_bound" begin
        mindex = -1; cindex1 = 42; cindex2 = -42; vindex = JuMP.index(x);
        # test adding normally
        @test isa(add_parameter_bound(x, pars[2], 0, 5), Nothing)
        bounds = ParameterBounds(Dict(par => IntervalSet(0, 5),
                                      pars[1] => IntervalSet(0.5, 2),
                                      pars[2] => IntervalSet(0, 5)))
        @test parameter_bounds(x) == bounds
        expected = ParameterBounds(Dict(par => IntervalSet(0, 5),
                                        pars[1] => IntervalSet(0.5, 1),
                                        pars[2] => IntervalSet(0, 5)))
        @test m.constrs[cindex1].bounds == expected
        @test m.constrs[cindex2].bounds == bounds
        # test measure error
        @test_throws ErrorException add_parameter_bound(x, par, 1, 2)
        # test constr error
        @test_throws ErrorException add_parameter_bound(x, pars[1], 2, 3)
        # test overwritting existing
        @test isa(add_parameter_bound(y, par, 0, 0), Nothing)
        @test parameter_bounds(y) == ParameterBounds(Dict(par => IntervalSet(0, 0)))
    end
    # test delete_parameter_bound
    @testset "delete_parameter_bound" begin
        cindex1 = 42; cindex2 = -42;
        # test normal
        @test isa(delete_parameter_bound(x, pars[1]), Nothing)
        bounds = ParameterBounds(Dict(par => IntervalSet(0, 5),
                                      pars[2] => IntervalSet(0, 5)))
        @test parameter_bounds(x) == bounds
        expected = ParameterBounds(Dict(par => IntervalSet(0, 5),
                                        pars[1] => IntervalSet(0, 1),
                                        pars[2] => IntervalSet(0, 5)))
        @test m.constrs[cindex1].bounds == expected
        @test m.constrs[cindex2].bounds == bounds
        # test already gone
        @test isa(delete_parameter_bound(x, pars[1]), Nothing)
        @test parameter_bounds(x) == bounds
        @test m.constrs[cindex1].bounds == expected
        @test m.constrs[cindex2].bounds == bounds
    end
    # test delete_parameter_bounds
    @testset "delete_parameter_bounds" begin
        cindex1 = 42; cindex2 = -42;
        # test normal
        @test isa(delete_parameter_bounds(x), Nothing)
        @test parameter_bounds(x) == ParameterBounds()
        expected = ParameterBounds(Dict(pars[1] => IntervalSet(0, 1)))
        @test m.constrs[cindex1].bounds == expected
        @test isa(m.constrs[cindex2], ScalarConstraint)
        # test already gone
        @test isa(delete_parameter_bounds(x), Nothing)
        @test parameter_bounds(x) == ParameterBounds()
        @test m.constrs[cindex1].bounds == expected
        @test isa(m.constrs[cindex2], ScalarConstraint)
    end
end

# Test Parameter Bound Macros
@testset "Parameter Bound Macros" begin
    # initialize the model and other needed information
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 10])
    @infinite_parameter(m, pars[1:2] in [0, 10], container = SparseAxisArray)
    @hold_variable(m, x)
    @hold_variable(m, y, parameter_bounds = par == 0)
    # test @set_parameter_bounds
    @testset "@set_parameter_bounds" begin
        bounds1 = ParameterBounds(Dict(pars[1] => IntervalSet(0, 1)))
        # test force error
        @test_macro_throws ErrorException @set_parameter_bounds(y, pars[1] in [0, 1])
        # test bad arguments
        @test_macro_throws ErrorException @set_parameter_bounds(x, x, 23)
        @test_macro_throws ErrorException @set_parameter_bounds(x, pars[1] == 0,
                                                                bad = 42)
        # test bad input format
        @test_macro_throws ErrorException @set_parameter_bounds(x, pars[1] = 0)
        @test_macro_throws ErrorException @set_parameter_bounds(par, pars[1] == 0)
        # prepare for main test
        data = DiscreteMeasureData(par, [1, 1], [1, 3], "", InfiniteOpt._w)
        mindex = -1; cindex1 = 42; cindex2 = -42; vindex = JuMP.index(x);
        m.measures[mindex] = Measure(x, data)
        m.var_to_meas[JuMP.index(x)] = [mindex]
        m.meas_to_constrs[mindex] = [cindex1]
        mref = MeasureRef(m, mindex)
        m.constrs[cindex1] = BoundedScalarConstraint(mref, MOI.EqualTo(0.0),
                                                     bounds1, bounds1)
        m.var_to_constrs[vindex] = [cindex2]
        m.constrs[cindex2] = JuMP.ScalarConstraint(x, MOI.EqualTo(0.0))
        # test measure error
        @test_macro_throws ErrorException @set_parameter_bounds(x, par in [1, 2])
        # test constr error
        @test_macro_throws ErrorException @set_parameter_bounds(x, pars[1] in [2, 3])
        # test normal
        bounds = ParameterBounds(Dict(par => IntervalSet(0, 5),
                                      pars[1] => IntervalSet(0.5, 2)))
        @test isa(@set_parameter_bounds(x, (par in [0, 5], pars[1] in [0.5, 2])),
                  Nothing)
        @test parameter_bounds(x) == bounds
        @test m.has_hold_bounds
        @test !optimizer_model_ready(m)
        expected = ParameterBounds(Dict(par => IntervalSet(0, 5),
                                        pars[1] => IntervalSet(0.5, 1)))
        @test m.constrs[cindex1].bounds == expected
        @test m.constrs[cindex2].bounds == bounds
        # test forced
        @test isa(@set_parameter_bounds(y, par == 1, force = true), Nothing)
        @test parameter_bounds(y) == ParameterBounds(Dict(par => IntervalSet(1, 1)))
    end
    # test @add_parameter_bounds
    @testset "@add_parameter_bounds" begin
        mindex = -1; cindex1 = 42; cindex2 = -42; vindex = JuMP.index(x);
        # test adding normally
        @test isa(@add_parameter_bounds(x, pars[2] in [0, 5]), Nothing)
        dict = Dict(par => IntervalSet(0, 5), pars[1] => IntervalSet(0.5, 2),
                    pars[2] => IntervalSet(0, 5))
        @test parameter_bounds(x) == ParameterBounds(dict)
        expected = Dict(par => IntervalSet(0, 5), pars[1] => IntervalSet(0.5, 1),
                        pars[2] => IntervalSet(0, 5))
        @test m.constrs[cindex1].bounds == ParameterBounds(expected)
        @test m.constrs[cindex2].bounds == ParameterBounds(dict)
        # test measure error
        @test_macro_throws ErrorException @add_parameter_bounds(x, par in [1, 2])
        # test constr error
        @test_macro_throws ErrorException @add_parameter_bounds(x, pars[1] in [2, 3])
        # test bad input format
        @test_macro_throws ErrorException @add_parameter_bounds(x, pars[1] = 0)
        @test_macro_throws ErrorException @add_parameter_bounds(par, pars[1] == 0)
        # test overwritting existing
        @test isa(@add_parameter_bounds(y, par == 0), Nothing)
        @test parameter_bounds(y) == ParameterBounds(Dict(par => IntervalSet(0, 0)))
    end
end


text:
@testset "onehot constructors" begin
  @test onehot(20, 10:10:30) == [false, true, false]
  @test onehot(20, (10,20,30)) == [false, true, false]
  @test onehot(40, (10,20,30), 20) == [false, true, false]

  @test_throws Exception onehot('d', 'a':'c')
  @test_throws Exception onehot(:d, (:a, :b, :c))
  @test_throws Exception onehot('d', 'a':'c', 'e')
  @test_throws Exception onehot(:d, (:a, :b, :c), :e)

  @test onehotbatch([20, 10], 10:10:30) == Bool[0 1; 1 0; 0 0]
  @test onehotbatch([20, 10], (10,20,30)) == Bool[0 1; 1 0; 0 0]
  @test onehotbatch([40, 10], (10,20,30), 20) == Bool[0 1; 1 0; 0 0]

  @test onehotbatch("abc", 'a':'c') == Bool[1 0 0; 0 1 0; 0 0 1]
  @test onehotbatch("zbc", ('a', 'b', 'c'), 'a') == Bool[1 0 0; 0 1 0; 0 0 1]

  @test onehotbatch([10, 20], [30, 40, 50], 30) == Bool[1 1; 0 0; 0 0]

  @test_throws Exception onehotbatch([:a, :d], [:a, :b, :c])
  @test_throws Exception onehotbatch([:a, :d], (:a, :b, :c))
  @test_throws Exception onehotbatch([:a, :d], [:a, :b, :c], :e)
  @test_throws Exception onehotbatch([:a, :d], (:a, :b, :c), :e)

  floats = (0.0, -0.0, NaN, -NaN, Inf, -Inf)
  @test onecold(onehot(0.0, floats)) == 1
  @test onecold(onehot(-0.0, floats)) == 2  # as it uses isequal
  @test onecold(onehot(Inf, floats)) == 5
end

@testset "onecold" begin
  a = [1, 2, 5, 3.]
  A = [1 20 5; 2 7 6; 3 9 10; 2 1 14]
  labels = ['A', 'B', 'C', 'D']

  @test onecold(a) == 3
  @test onecold(A) == [3, 1, 4]
  @test onecold(a, labels) == 'C'
  @test onecold(a, Tuple(labels)) == 'C'
  @test onecold(A, labels) == ['C', 'A', 'D']
  @test onecold(A, Tuple(labels)) == ['C', 'A', 'D']

  data = [:b, :a, :c]
  labels = [:a, :b, :c]
  hot = onehotbatch(data, labels)
  cold = onecold(hot, labels)

  @test cold == data
end

@testset "onehotbatch indexing" begin
  y = onehotbatch(ones(3), 1:10)
  @test y[:,1] isa OneHotVector
  @test y[:,:] isa OneHotMatrix
end


text:
using Pkg
Pkg.activate(".")

using PyPlot
using NPZ
rc("text", usetex=true)
rc("font", family="serif")
rc(Dict("savefig.bbox_inches" => "tight"))
rc(Dict("savefig.metadata" => Dict("CreationDate" => nothing)))

linreg(x, y) = hcat(fill!(similar(x), 1), x) \ y

function scaling_exponent_generator(data_x::Vector,
                                    data_y::Vector,
                                    batch::Int=11)
   @assert length(data_y) == length(data_x)
   @assert length(data_x) >= batch
   @assert isodd(batch)

   result = Float64[]
   #println(result)
   for i=1:(length(data_y)-batch+1)
      #println(linreg(data_x[i:i+batch-1], data_y[i:i+batch-1]))
      push!(result, linreg(data_x[i:i+batch-1], data_y[i:i+batch-1])[2])
   end
   data_x[div(batch, 2)+1:end-div(batch, 2)], result
end

#println(collect(1:20.))
#println(2 .+ 3. .* collect(1:20))
#scaling_exponent_generator( collect(1:20.) , 2 .+ 3 .* collect(1:20.))


text:
## AetherlangFunction type
struct AetherlangFunction
    argnames::Vector{String}
    expr::Vector{Union{String, AetherlangObject}} # vector of tokens to evaluate
end
function aetherlang_show(io::IO, f::AetherlangFunction)
    print(io, "(function _ "*join(f.argnames, ' ')*" = [body])")
end
function aetherlang_call(f::AetherlangFunction, ns::Namespace, args::Vector{AetherlangObject}, line::Ref{UInt})::AetherlangObject
    global ether
    fns = copy(ns) # copy(ns) in this case
    if length(args) < length(f.argnames)
        throw(AetherlangError("Not enough arguments"))
    end
    for i in 1:length(f.argnames)
        fns[f.argnames[i]] = args[i]
    end
    fns["oargs"] = AetherlangObject([arg.dataref for arg in args[length(f.argnames)+1:end]])
    aetherlang_eval!(f.expr, fns, line)
end
AetherlangObject(f::AetherlangFunction) = AetherlangObject{AetherlangFunction}(f, Ref(AETH_BUILTIN_TYPES.callable))

## Wrapping Julia functions
function AetherlangObject(f::Function)::AetherlangObject{Function}
    """Wraps a julia function into an AetherlangObject of type callable. The function needs to take a namespace (Namespace) and a list of arguments (vector of AetherlangObject) and output a single AetherlangObject"""
    AetherlangObject{Function}(f, Ref(AETH_BUILTIN_TYPES.callable))
end
function aetherlang_show(io::IO, f::Function)
    print(io, "(julia.function "*string(f)*" [args] = [body])")
end
function aetherlang_call(f::Function, ns::Namespace, args::Vector{AetherlangObject}, line::Ref{UInt})::AetherlangObject
    try
        f(line, ns, args...) # not copy(ns) in this case
    catch e
        if typeof(e) <: AetherlangException
            throw(e)
        else
            throw(AetherlangError("Bad call of built-in function: "*string(e)))
        end
    end
end

max_stars_repo_name: Aetherlang/aetherlang

max_issues_repo_name: gojakuch/aetherlang

max_forks_repo_name: gojakuch/aetherlang


text:
mutable struct BasicDiscMuvParameterNState{NI<:Integer, NR<:Real} <: ParameterNState{Discrete, Multivariate}
  value::Matrix{NI}
  loglikelihood::Vector{NR}
  logprior::Vector{NR}
  logtarget::Vector{NR}
  diagnosticvalues::Matrix
  size::Integer
  sizesquared::Integer
  sizecubed::Integer
  monitor::Vector{Bool}
  n::Integer
  diagnostickeys::Vector{Symbol}

  function BasicDiscMuvParameterNState{NI, NR}(
    size::Integer,
    n::Integer,
    monitor::Vector{Bool}=[true; fill(false, 3)],
    diagnostickeys::Vector{Symbol}=Symbol[],
    ::Type{NI}=Int64,
    ::Type{NR}=Float64,
    diagnosticvalues::Matrix=Array{Any}(length(diagnostickeys), isempty(diagnostickeys) ? 0 : n)
  ) where {NI<:Integer, NR<:Real}
    instance = new()

    fnames = fieldnames(BasicDiscMuvParameterNState)
    setfield!(instance, fnames[1], Array{NI}((monitor[1] == false ? (zero(Integer), zero(Integer)) : (size, n))...))
    for i in 2:4
      l = (monitor[i] == false ? zero(Integer) : n)
      setfield!(instance, fnames[i], Array{NR}(l))
    end

    instance.diagnosticvalues = diagnosticvalues
    instance.size = size
    instance.sizesquared = instance.size^2
    instance.sizecubed = instance.size^3
    instance.monitor = monitor
    instance.n = n
    instance.diagnostickeys = diagnostickeys

    instance
  end
end

BasicDiscMuvParameterNState(
  size::Integer,
  n::Integer,
  monitor::Vector{Bool}=[true; fill(false, 3)],
  diagnostickeys::Vector{Symbol}=Symbol[],
  ::Type{NI}=Int64,
  ::Type{NR}=Float64,
  diagnosticvalues::Matrix=Array{Any}(length(diagnostickeys), isempty(diagnostickeys) ? 0 : n)
) where {NI<:Integer, NR<:Real} =
  BasicDiscMuvParameterNState{NI, NR}(size, n, monitor, diagnostickeys, NI, NR, diagnosticvalues)

function BasicDiscMuvParameterNState(
  size::Integer,
  n::Integer,
  monitor::Vector{Symbol},
  diagnostickeys::Vector{Symbol}=Symbol[],
  ::Type{NI}=Int64,
  ::Type{NR}=Float64,
  diagnosticvalues::Matrix=Array{Any}(length(diagnostickeys), isempty(diagnostickeys) ? 0 : n)
) where {NI<:Integer, NR<:Real}
  fnames = fieldnames(BasicDiscMuvParameterNState)
  BasicDiscMuvParameterNState(
    size, n, [fnames[i] in monitor ? true : false for i in 1:4], diagnostickeys, NI, NR, diagnosticvalues
  )
end

const DiscMuvMarkovChain = BasicDiscMuvParameterNState

function copy!(nstate::BasicDiscMuvParameterNState, state::BasicDiscMuvParameterState, i::Integer)
  fnames = fieldnames(BasicDiscMuvParameterNState)

  if nstate.monitor[1]
    getfield(nstate, fnames[1])[1+(i-1)*state.size:i*state.size] = getfield(state, fnames[1])
  end

  for j in 2:4
    if nstate.monitor[j]
      getfield(nstate, fnames[j])[i] = getfield(state, fnames[j])
    end
  end

  if !isempty(nstate.diagnosticvalues)
    nstate.diagnosticvalues[:, i] = state.diagnosticvalues
  end
end

eltype{NI<:Integer, NR<:Real}(::Type{BasicDiscMuvParameterNState{NI, NR}}) = (NI, NR)
eltype{NI<:Integer, NR<:Real}(::BasicDiscMuvParameterNState{NI, NR}) = (NI, NR)

==(z::S, w::S) where {S<:BasicDiscMuvParameterNState} = reduce(&, [getfield(z, n) == getfield(w, n) for n in fieldnames(S)[1:8]])

isequal(z::S, w::S) where {S<:BasicDiscMuvParameterNState} =
  reduce(&, [isequal(getfield(z, n), getfield(w, n)) for n in fieldnames(S)[1:8]])

function show(io::IO, nstate::BasicDiscMuvParameterNState{NI, NR}) where {NI<:Integer, NR<:Real}
  fnames = fieldnames(BasicDiscMuvParameterNState)
  fbool = map(n -> !isempty(getfield(nstate, n)), fnames[1:4])
  indentation = "  "

  println(io, "BasicDiscMuvParameterNState:")

  println(io, indentation*"eltype: $(eltype(nstate))")
  println(io, indentation*"state size = $(nstate.size)")
  println(io, indentation*"number of states = $(nstate.n)")

  print(io, indentation*"monitored components:")
  if !any(fbool)
    println(io, " none")
  else
    print(io, "\n")
    for i in 1:4
      if fbool[i]
        println(io, string(indentation^2, fnames[i]))
      end
    end
  end

  print(io, indentation*"diagnostics:")
  if isempty(nstate.diagnostickeys)
    print(io, " none")
  else
    for k in nstate.diagnostickeys
      print(io, "\n")
      print(io, string(indentation^2, k))
    end
  end
end


text:
trapz(x::Vector, dt, x0=0) = cumsum(vcat(0, 0.5*dt*(x[1:end-1] + x[2:end]))) + x0
trapz(x::Matrix, dt, x0=0) = cumsum(hcat(zeros(size(x,1)), 0.5*dt*(x[:,1:end-1] + x[:,2:end])), 2) .+ x0
vel(xd, yd) = sqrt.(xd.^2 + yd.^2)
acc(xd, xdd, yd, ydd, v=vel(xd,yd)) = (xd.*xdd + yd.*ydd)./v
curv(xd, xdd, yd, ydd, v=vel(xd,yd)) = (xd.*ydd - xdd.*yd)./v.^3

@inbounds @views function all_future_xy(xdd_choices, y_targets, state0, tibvp,
                                        a0 = [(1,1)], ply=4, t_hold=3, dt=Float32(0.1))
    Nx = length(xdd_choices)
    Ny = length(y_targets)
    if a0 isa Vector{Tuple{Int,Int}}
        X = zeros(Float32, Nx^ply, 1 + length(a0) + ply*t_hold, 3)
        Y = zeros(Float32, Ny^ply, 1 + length(a0) + ply*t_hold, 3)
        X[:,1,:] .= state0[1:2:end]'
        Y[:,1,:] .= state0[2:2:end]'
        t0 = 1 + length(a0)
        for (t, (xi, yi)) in enumerate(a0)
            X[:,t+1,3] .= xdd_choices[xi]
            q0 = SVector(Y[1,t,1],
                         Y[1,t,2],
                         Y[1,t,3])
            qf = SVector(y_targets[yi], 0f0, 0f0)
            c, u = tibvp(q0, qf, 100f0)
            if dt > duration(u)
                q = qf
            else
                q = u.x(q0, u.y, u.t, dt)
            end
            Y[:,t+1,1] .= q[1]
            Y[:,t+1,2] .= q[2]
            Y[:,t+1,3] .= q[3]
        end
    elseif a0 isa Tuple
        X0, Y0 = a0
        X0 = reshape(X0, (1, size(X0)...))
        Y0 = reshape(Y0, (1, size(Y0)...))
        X = zeros(Float32, Nx^ply, 1 + size(X0,2) + ply*t_hold, 3)
        Y = zeros(Float32, Ny^ply, 1 + size(Y0,2) + ply*t_hold, 3)
        X[:,1,:] .= state0[1:2:end]'
        Y[:,1,:] .= state0[2:2:end]'
        X[:,2:1+size(X0,2),:] .= X0
        Y[:,2:1+size(Y0,2),:] .= Y0
        t0 = 1 + size(X0,2)
    else
        error("Incompatible a0 type.")
    end

    # xdd, xd, x
    for (r, ci) in enumerate(CartesianRange(Tuple(Nx for i in 1:ply)))
        for (p, v) in enumerate(ci.I)
            for t in 1:t_hold
                X[r,t0+t_hold*(p-1)+t,3] = xdd_choices[v]
            end
        end
    end
    cumsum!(X[:,2:end,2], X[:,1:end-1,3]*dt, 2)
    X[:,2:end,2] .+= X[:,1,2]
    cumsum!(X[:,2:end,1], (X[:,1:end-1,2] .+ X[:,2:end,2])*dt./2, 2)
    X[:,2:end,1] .+= X[:,1,1]

    # ydd, yd, y
    for (r, ci) in enumerate(CartesianRange(Tuple(Ny for i in 1:ply)))
        for (p, v) in enumerate(ci.I)
            curr_step = t0 + (p-1)*t_hold
            q0 = SVector(Y[r,curr_step,1],
                         Y[r,curr_step,2],
                         Y[r,curr_step,3])
            qf = SVector(y_targets[v], 0f0, 0f0)
            c, u = tibvp(q0, qf, 100f0)
            for t in 1:t_hold
                if dt*t > duration(u)
                    q = qf
                else
                    q = u.x(q0, u.y, u.t, t*dt)
                end
                Y[r,curr_step+t,1] = q[1]
                Y[r,curr_step+t,2] = q[2]
                Y[r,curr_step+t,3] = q[3]
            end
        end
    end
    X, Y
end

function crossover_time(car1::Matrix, car2::Matrix)
    dx = car1[:,1:1] - car2[:,1:1]
    dv = car1[:,3:3] - car2[:,3:3]
    -dx./dv
end

function crossover_time(car1::Array{T,3}, car2::Array{S,3}) where {T,S}
    dx = car1[:,:,1:1] - car2[:,:,1:1]
    dv = car1[:,:,3:3] - car2[:,:,3:3]
    -dx./dv
end

function extra_features(car1::Matrix, car2::Matrix, column_list=String[])
    extras_tensor = zeros(Float32, size(car1, 1), length(column_list))
    for (i, c) in enumerate(column_list)
        if c == "tco"
            extras_tensor[:,i] = crossover_time(car1, car2)
        else
            error("Unknown hand-constructed (extra) feature: $c")
        end
    end
    extras_tensor
end

function extra_features(car1::Array{T,3}, car2::Array{S,3}, column_list=String[]) where {T,S}
    extras_tensor = zeros(Float32, size(car1, 1), size(car1, 2), length(column_list))
    for (i, c) in enumerate(column_list)
        if c == "tco"
            extras_tensor[:,:,i] = crossover_time(car1, car2)
        else
            error("Unknown hand-constructed (extra) feature: $c")
        end
    end
    extras_tensor
end

visualize_predictions(data_dict) = visualize_predictions(data_dict["car1"],
                                                         data_dict["car2"],
                                                         get(data_dict, "extras_cols", String[]))
function visualize_predictions(car1, car2, extras_cols=String[])
    D = size(car1, 1)
    T = size(car1, 2)
    PH = 15
    Δt = 0.1
    @manipulate for data_id in 1:D,
                    t_predict in slider(1:T, value=10, label="t_predict"),
                    X in slider(0:256, value=0, label="Long. action"),
                    Y in slider(0:16, value=0, label="Lat. action"),
                    k in slider(0:100, value=50, label="# samples")
        tic()
        x0 = car1[data_id,t_predict+1:t_predict+3,1:2:end]
        y0 = car1[data_id,t_predict+1:t_predict+3,2:2:end]
        pfX, pfY = all_future_xy(xdd_choices, y_targets, car1[data_id,t_predict,:], tibvp, (x0, y0), 4)
        car1_future_input = zeros(1, 15, 6)
        if X == 0
            car1_future_input[1,:,1:2:end] = car1[data_id,t_predict + (1:PH),1:2:end]
        else
            car1_future_input[1,:,1:2:end] = pfX[X,2:end,:]
        end
        if Y == 0
            car1_future_input[1,:,2:2:end] = car1[data_id,t_predict + (1:PH),2:2:end]
        else
            car1_future_input[1,:,2:2:end] = pfY[Y,2:end,:]
        end
        TrafficWeavingPlanner.predict!(car1[data_id:data_id,:,:],
                                       car2[data_id:data_id,:,:],
                                       extra_features(car1[data_id:data_id,:,:], car2[data_id:data_id,:,:], extras_cols),
                                       [t_predict],
                                       car1_future_input,
                                       k)
        run_time = toq()
        pred_container = TrafficWeavingPlanner.pred_container
        x0, y0, xd0, yd0 = car2[data_id,t_predict,1:4]
        xdd = hcat(fill(car2[data_id,t_predict,5], size(pred_container.y,1)), pred_container.y[:,1,:,1])
        ydd = hcat(fill(car2[data_id,t_predict,6], size(pred_container.y,1)), pred_container.y[:,1,:,2])
        xd = trapz(xdd, Δt, xd0)
        yd = trapz(ydd, Δt, yd0)
        x = trapz(xd, Δt, x0)
        y = trapz(yd, Δt, y0)
        v = vel(xd, yd)
        a = acc(xd, xdd, yd, ydd, v)
        κ = curv(xd, xdd, yd, ydd, v)

        N = Int(sum(pred_container.z[1,1,:]))
        K = div(size(pred_container.z, 3), N)

        z_list = Vector{Int}(pred_container.z[:,1,:]*[i*K^j for j in 0:N-1 for i in 0:K-1]) + 1
        z_unique = 1:K^N

        color_list = distinguishable_colors(3+length(z_unique), [RGB(1,0,1), RGB(1,0,0), RGB(0,0,1)])[4:end]
        color_dict = Dict(z => c for (z,c) in zip(z_unique, color_list))
        colors = reshape([color_dict[z] for z in z_list], (1,length(z_list)))

        T_history = 1:t_predict
        T_future = t_predict + (0:PH)
        car1_history = car1[data_id,T_history,:]
        car2_history = car2[data_id,T_history,:]
        car1_future = [car1[data_id,t_predict,:]'; car1_future_input[1,:,:]]
        car2_future = car2[data_id,T_future,:]

        transparency = 0.7

        xy_left = plot([car1_history[:,1], car1_future[:,1], car2_history[:,1], car2_future[:,1]],
                       [car1_history[:,2], car1_future[:,2], car2_history[:,2], car2_future[:,2]],
                       linewidth=2,
                       color=[:blue :blue :red :red],
                       linestyles=[:solid :dash :solid :dash],
                       alpha=[1 1 1 ((X == 0 && Y == 0) ? 1 : 0)],
                       xlims=(-150,5), ylims=(-8,0),
                       title="position")
        plot!(xy_left, x', y', linewidth=.75, color=colors, alpha=transparency, legend=false)
        xy_right = plot(x', y', linewidth=.75, color=colors, alpha=transparency, legend=false, title="computation time = $(run_time)s")

        car1_history_v = vel(car1_history[:,3], car1_history[:,4])
        car2_history_v = vel(car2_history[:,3], car2_history[:,4])
        car1_future_v = vel(car1_future[:,3], car1_future[:,4])
        car2_future_v = vel(car2_future[:,3], car2_future[:,4])
        v_left = plot([Δt*T_history, Δt*T_future, Δt*T_history, Δt*T_future],
                      [car1_history_v, car1_future_v, car2_history_v, car2_future_v],
                      linewidth=2,
                      color=[:blue :blue :red :red],
                      linestyles=[:solid :dash :solid :dash],
                      alpha=[1 1 1 ((X == 0 && Y == 0) ? 1 : 0)],
                      title="longitudinal velocity")
        plot!(v_left, Δt*T_future, v', linewidth=.75, color=colors, alpha=transparency, legend=false)
        v_right = plot(Δt*T_future, v', linewidth=.75, color=colors, alpha=transparency, legend=false)

        car1_history_a = acc(car1_history[:,3], car1_history[:,5], car1_history[:,4], car1_history[:,6])
        car2_history_a = acc(car2_history[:,3], car2_history[:,5], car2_history[:,4], car2_history[:,6])
        car1_future_a = acc(car1_future[:,3], car1_future[:,5], car1_future[:,4], car1_future[:,6])
        car2_future_a = acc(car2_future[:,3], car2_future[:,5], car2_future[:,4], car2_future[:,6])
        a_left = plot([Δt*T_history, Δt*T_future, Δt*T_history, Δt*T_future],
                      [car1_history_a, car1_future_a, car2_history_a, car2_future_a],
                      linewidth=2,
                      color=[:blue :blue :red :red],
                      linestyles=[:solid :dash :solid :dash],
                      alpha=[1 1 1 ((X == 0 && Y == 0) ? 1 : 0)],
                      xlims=(0, Δt*T_future[end]), ylims=(-6, 5),
                      title="longitudinal acceleration")
        plot!(a_left, Δt*T_future, a', linewidth=.75, color=colors, alpha=transparency, legend=false)
        a_right = plot(Δt*T_future, a', linewidth=.75, color=colors, alpha=transparency, legend=false)

        plot(xy_left, xy_right, v_left, v_right, a_left, a_right,
             layout=@layout([a b; c d; e f]), size=(800,900))
    end
end

function static_trial(data_dict, idx, xdd_choices, y_targets, tibvp, T=30; start_wait=8, extras_cols = String[])
    car1 = zeros(Float32, 1, T+15, 6)
    car1_curr_state = data_dict["car1"][idx,1,:]
    car1[1,1,:] = car1_curr_state

    if car1_curr_state[2] > -4
        l0, lf = 1, 2
    else
        l0, lf = 2, 1
    end

    pfX, pfY = all_future_xy(xdd_choices, y_targets, car1_curr_state, tibvp, [(1,l0) for i in 1:start_wait], 0)
    car1[1,1:start_wait+1,1:2:end] = pfX
    car1[1,1:start_wait+1,2:2:end] = pfY
    t = start_wait + 1 - 3
    car1_curr_state = car1[1,t,:]
    a0 = [(1,l0) for i in 1:3]
    while t < T
        pfX, pfY = all_future_xy(xdd_choices, y_targets, car1_curr_state, tibvp, a0, 4)
        car1_in = car1
        car2_in = data_dict["car2"][idx:idx,1:T+15,:]
        extras_in = extra_features(car1_in, car2_in, extras_cols)
        TrafficWeavingPlanner.predict!(car1_in, car2_in, extras_in, [t], pfX[:,2:end,:], pfY[:,2:end,:], 16)
#         @show norm(car1[1,t+1:t+3,:] - TrafficWeavingPlanner.pred_container.y[1:3,:])
#         @show pfX[1,:,3]
#         @show car1[1,t:t+6,idshow]
#         @show TrafficWeavingPlanner.pred_container.y[1:6,idshow]
        car1[1,t+1:t+15,:] = TrafficWeavingPlanner.pred_container.y
        t = t + 3
        car1_curr_state = car1[1,t,:]
        a0 = (TrafficWeavingPlanner.pred_container.y[4:6,1:2:end], TrafficWeavingPlanner.pred_container.y[4:6,2:2:end])
    end
    car1
end

function dynamic_trial(xdd_choices, y_targets, tibvp, initial_state=Float32[-136,-6,30,0,0,0]; T=100, start_wait=8, dt=Float32(.1), extras_cols = String[])
    hz = 100
    rate = TrafficWeavingPlanner.Rate(hz)
    TrafficWeavingPlanner.clear_human_states()
    TrafficWeavingPlanner.reset_simulator()
    for timeout in 1:hz
        length(TrafficWeavingPlanner.human_states.t) == 1 && break
        TrafficWeavingPlanner.rossleep(rate)
    end
    @assert length(TrafficWeavingPlanner.human_states.t) == 1
    t_offset = TrafficWeavingPlanner.human_states.t[1]

    car1 = zeros(Float32, 1, T, 6)
    car2 = zeros(Float32, 1, T, 6)
    car1_curr_state = initial_state
    car1[1,1,:] = car1_curr_state

    if car1_curr_state[2] > -4
        l0, lf = 1, 2
    else
        l0, lf = 2, 1
    end

    pfX, pfY = all_future_xy(xdd_choices, y_targets, car1_curr_state, tibvp, [(1,l0) for i in 1:start_wait], 0)
    car1[1,1:start_wait+1,1:2:end] = pfX
    car1[1,1:start_wait+1,2:2:end] = pfY
    t = start_wait + 1 - 3
    car1_curr_state = car1[1,t,:]
    a0 = [(1,l0) for i in 1:3]
    TrafficWeavingPlanner.publish_robot_path_plan(dt*(1:t+2), car1[1,2:t+3,1], car1[1,2:t+3,2])
    while (car1_curr_state[1] < 3 || TrafficWeavingPlanner.human_states.X[end][1] < 3) && t < T-15
        pfX, pfY = all_future_xy(xdd_choices, y_targets, car1_curr_state, tibvp, a0, 4)
        for timeout in 1:5*hz
            TrafficWeavingPlanner.human_states.t[end] > t_offset + dt*(t - 1) && break
            TrafficWeavingPlanner.rossleep(rate)
        end
        println((TrafficWeavingPlanner.human_states.t[end], t_offset + dt*(t - 1),
                 TrafficWeavingPlanner.human_states.t[end] - t_offset, dt*(t - 1)))
        @assert TrafficWeavingPlanner.human_states.t[end] > t_offset + dt*(t - 1)
        car2_interp = interpolate((TrafficWeavingPlanner.human_states.t,), TrafficWeavingPlanner.human_states.X, Gridded(Linear()))
        car2_in = [car2_interp[make_a_pr] for make_a_pr in t_offset+dt*(0:t-1)]
        @assert eltype(car2_in) == SVector{6,Float32}
        car2_in = reshape(reinterpret(Float32, car2_in, (6, length(car2_in)))', (1, length(car2_in), 6))
        car2[1,1:size(car2_in,2),:] = car2_in
        extras = extra_features(car1, car2, extras_cols)
        TrafficWeavingPlanner.predict!(car1, car2, extras, [t], pfX[:,2:end,:], pfY[:,2:end,:], 16)
        car1[1,t+1:t+15,:] = TrafficWeavingPlanner.pred_container.y
        t = t + 3
        car1_curr_state = car1[1,t,:]
        a0 = (TrafficWeavingPlanner.pred_container.y[4:6,1:2:end], TrafficWeavingPlanner.pred_container.y[4:6,2:2:end])
#         TrafficWeavingPlanner.publish_robot_path_plan(dt*(1:t+2), car1[1,2:t+3,1], car1[1,2:t+3,2])
#         TrafficWeavingPlanner.publish_robot_path_plan(dt*(1:t+5), car1[1,2:t+6,1], car1[1,2:t+6,2])
        TrafficWeavingPlanner.publish_robot_path_plan(dt*(1:t+8), car1[1,2:t+9,1], car1[1,2:t+9,2])
    end
    car1[:,2:t-3,:], car2[:,2:t-3,:]
end

### OLD SEMI-TRASH
# function random_coherent_actions(N, nx, ny, ply, a1_marginal, a2_marginal, starting_action=(1,1), stay_prob=.5)
# #     rand([(x,y) for x in 1:nx for y in 1:ny], N, ply)
# #     actions = Array{Tuple{Int,Int}}(N, ply)
#     actions = fill(starting_action, N, ply)
#     for c in 2:ply
#         for r in 1:N
#             if rand() < stay_prob
#                 actions[r,c] = actions[r,c-1]
#             else
#                 actions[r,c] = (rand(a1_marginal), rand(a2_marginal))
#             end
#         end
#     end
#     sortrows(actions)
# end

# random_actions(N, nx, ny, ply) = sortrows(rand([(x,y) for x in 1:nx for y in 1:ny], N, ply))

# @views function random_future(N, xdd_choices, y_targets, state0, tibvp,
#                               ply=5,
#                               actions=random_actions(N, length(xdd_choices), length(y_targets), ply),
#                               t_hold=3,
#                               dt=Float32(0.1))
#     x, y, ẋ, ẏ, ẍ, ÿ = 1:6

#     car_future = zeros(Float32, N, 1 + ply*t_hold, 6)
#     car_future[:,1,:] .= state0'
#     # xdd, xd, x
#     xdd_by_ply = [xdd_choices[i] for (i,j) in actions]
#     for i in 1:t_hold
#         car_future[:,1+i:t_hold:end,ẍ] = xdd_by_ply
#     end
#     cumsum!(car_future[:,2:end,ẋ], car_future[:,1:end-1,ẍ]*dt, 2)
#     car_future[:,2:end,ẋ] .+= car_future[:,1,ẋ]
#     cumsum!(car_future[:,2:end,x], (car_future[:,1:end-1,ẋ] .+ car_future[:,2:end,ẋ])*dt./2, 2)
#     car_future[:,2:end,x] .+= car_future[:,1,x]

#     # y_target, ydd, yd, y
# #     if state0[y] == y_final
# #         y_target = fill(y_final, N, ply)
# #     else
# # #         y_target = fill(y_start, N, ply)
# # #         for i in 1:N
# # #             y_target[i,rand(1:ply+1):end] = y_final
# # #         end
# #         y_target = rand([y_start, y_final], N, ply)
# #     end
#     y_target = [y_targets[j] for (i,j) in actions]
#     for i in 1:N
#         for p in 1:ply
#             curr_step = 1 + (p-1)*t_hold
#             q0 = SVector(car_future[i,curr_step,y],
#                          car_future[i,curr_step,ẏ],
#                          car_future[i,curr_step,ÿ])
#             qf = SVector(y_target[i,p], 0f0, 0f0)
#             c, u = tibvp(q0, qf, 100f0)
#             for t in 1:t_hold
#                 if dt*t > duration(u)
#                     q = qf
#                 else
#                     q = u.x(q0, u.y, u.t, t*dt)
#                 end
#                 car_future[i,curr_step+t,y] = q[1]
#                 car_future[i,curr_step+t,ẏ] = q[2]
#                 car_future[i,curr_step+t,ÿ] = q[3]
#             end
#         end
#     end
#     actions, car_future
# end

# function tree_reduce_max(actions, rewards)
#     if size(actions, 2) == 1
#         return 1, mean(rewards)
#     else
#         Nr = size(actions, 1)
#         split_points = [1]
#         prev_action = actions[1,1]
#         for r in 1:Nr
#             if actions[r,1] != prev_action
#                 push!(split_points, r)
#                 prev_action = actions[r,1]
#             end
#         end
#         push!(split_points, Nr+1)
#         return findmax(tree_reduce_max(view(actions, split_points[i]:split_points[i+1]-1, 2:size(actions, 2)),
#                                        view(rewards, split_points[i]:split_points[i+1]-1))
#                        for i in 1:length(split_points)-1)
#     end
# end

max_issues_repo_name: codingblazes/NHumanModeling

max_forks_repo_name: codingblazes/NHumanModeling