text:
export DenseHDF5Array, extractsparse, extractdense
import HDF5
import SparseArrays

"""
This class implements a `AbstractArray` wrapper around a HDF5 dataset.
Instances of this class can be used to represent a HDF5 dataset inside collections that expect an `AbstractArray`.
Most importantly, this class does not load any data into memory on construction (and only minimal loading in its `show()` method).
This allows users to manipulate arbitrarily large datasets in low-memory environments.

The `DenseHDF5Array` is structurally dense, reflecting the fact that a HDF5 dataset explicitly stores all values (zero or otherwise).
`T` can be any type though this is usually numeric, and occasionally boolean.
Any number of dimensions is supported, but note that the dimensions are permuted in Julia;
the first dimension is the fastest changing in Julia but is the slowest changing in HDF5.

It is also possible to perform calculations on this class, in which case values are retrieved from file on demand.
While memory-efficient, this approach is likely to be very slow, as it involves multiple (often redundant) I/O calls to the disk.
We recommend using this class as a placeholder for a real array in code that does not require the actual values.
Once values are explicitly needed, the entire array can be loaded into memory with [`extractdense`](@ref) or [`extractsparse`](@ref).
"""
struct DenseHDF5Array{T, N} <: HDF5Array{T, N}
    file::String
    name::String
    dims::Tuple{Vararg{Integer,N}}

    # Caching utilities. This is only intended to avoid crushing performance
    # issues for getindex(), which is only to be used for giving something
    # sensible to show(). Actual computation should not use this cache.
    chunkdim::Tuple{Vararg{Integer,N}}
    nchunks::Tuple{Vararg{Integer,N}}
    cache::Dict{Integer,Array{T,N}}
end

function blocks_per_dim(blockdim::Tuple{Vararg{Integer}}, dims::Tuple{Vararg{Integer}})
    N = length(blockdim)
    tmp = Vector{Int}(undef, N)
    for i in 1:N
        tmp[i] = Integer(ceil(dims[i]/blockdim[i]))
    end
    return (tmp...,)
end

"""
    HDF5Array(file, name)

Create a `HDF5Array` for a dense dataset inside a HDF5 file `file` at name `name`.
No values are loaded into memory by the constructor (though `show`ing the matrix may load a few values for display).

# Examples
```jldoctest
julia> using HDF5Arrays

julia> tmp = tempname();

julia> exampledense(tmp, "stuff", (20, 10))

julia> x = DenseHDF5Array(tmp, "stuff");

julia> size(x)
(20, 10)

julia> using SparseArrays

julia> SparseArrays.issparse(x)
false
```
"""
function DenseHDF5Array(file::String, name::String)
    handle = HDF5.h5open(file)

    dset = handle[name]
    if !isa(dset, HDF5.Dataset)
        throw(ErrorException("'" * name * "' in '" * file * "' should be a dataset"))
    end

    dims = HDF5.get_extent_dims(dset)[1]
    type = HDF5.get_jl_type(dset)
    N = length(dims)

    # Figuring out the caching dimensions. By default, we set it to
    # the chunk dimensions (or 100x100x100,etc. if contiguous).
    if HDF5.get_layout(HDF5.get_create_properties(dset)) == :chunked
        chunkdim = HDF5.get_chunk(dset)
    else
        tmp = Vector{Int}(undef, N)
        fill!(tmp, 100)
        chunkdim = (tmp...,)
    end

    nchunks = blocks_per_dim(chunkdim, dims)
    return DenseHDF5Array{type, N}(file, name, dims, chunkdim, nchunks, Dict{Integer,Array{type,N}}())
end

"""
    size(x::DenseHDF5Array{T,N})

Get the size of a `DenseHDF5Array{T,N}` as a `N`-tuple.
"""
function Base.size(x::DenseHDF5Array{T,N}) where {T,N}
    return x.dims
end

"""
    getindex(x::DenseHDF5Array{T,N}, I::Vararg{Integer,N})

Get the value of a `DenseHDF5Array{T,N}` at the position specified by indices `I`.
It would be unwise to use this function for anything other than `show()` -
we suggest using [`extractdense`](@ref) to obtain larger blocks of data.

# Examples
```jldoctest
julia> using HDF5Arrays

julia> tmp = tempname();

julia> exampledense(tmp, "stuff", (20, 10))

julia> x = DenseHDF5Array(tmp, "stuff");

julia> typeof(getindex(x, 1, 1))
Float64
```
"""
function Base.getindex(x::DenseHDF5Array{T,N}, I::Vararg{Integer,N}) where {T,N}
    block_id = 0
    multiplier = 1
    tmp = Vector{Int}(undef, N)
    for i in 1:N
        tmp[i] = (I[i] - 1) ÷ x.chunkdim[i]
        block_id += multiplier * tmp[i]
        multiplier *= x.nchunks[i]
    end

    # Adding the block to the cache. This should probably be protected
    # by a mutex, but it's not like you can use HDF5 in parallel anyway!
    if !haskey(x.cache, block_id)
        handle = HDF5.h5open(x.file)
        dset = handle[x.name]

        indices = Vector{AbstractVector{Int}}(undef, N)
        for i in 1:N
            start = tmp[i] * x.chunkdim[i] + 1
            finish = min(x.dims[i], start + x.chunkdim[i] - 1)
            indices[i] = start:finish
        end

        x.cache[block_id] = Base.getindex(dset, indices...)
    end

    # Fetching the block and extracting the desired element.
    for i in 1:N
        tmp[i] = (I[i] - 1) % x.chunkdim[i] + 1
    end
    return getindex(x.cache[block_id], tmp...)
end

mutable struct BlockInfo
    range::AbstractRange{Int}
    internal::Vector{Int}
    function BlockInfo(j)
        new(1:0, Int[j])
    end
end

function general_dense_extractor(x::DenseHDF5Array{T,N}, indices, store::Function; blockdim = nothing) where {T,N}
    # Configuring the block dimensions.
    if blockdim == nothing
        blockdim = x.chunkdim
    end

    # Deciding which blocks have the indices we want. Note that indices should
    # already be sorted and unique prior to calling this function.
    idx = to_indices(x, indices)
    idx_by_block = Vector{Dict{Int,BlockInfo}}(undef, N)
    used_blocks = Vector{Vector{Int}}(undef, N)
    nblocks = Vector{Int}(undef, N)

    for d in 1:N
        curindex = idx[d]
        curblockdim = blockdim[d]
        curblocks = Dict{Int,BlockInfo}()

        for j in curindex
            b = (j - 1) ÷ curblockdim + 1
            if !haskey(curblocks, b)
                curblocks[b] = BlockInfo(j)
            else
                push!(curblocks[b].internal, j)
            end
        end

        # Indices should already be sorted and unique, so no need
        # to do any extra sorting on 'internal'.
        for (_, v) in curblocks
            first = v.internal[1]
            v.range = (first):(v.internal[length(v.internal)])
            for i in 1:length(v.internal)
                v.internal[i] -= first - 1
            end
        end

        idx_by_block[d] = curblocks
        used_blocks[d] = collect(keys(curblocks))
        sort!(used_blocks[d])
        nblocks[d] = length(curblocks)
    end

    # Extract along the fastest changing dimension. Normally in HDF5 this is the last,
    # but it seems Julia transposes things for us, so we'll go from the first.
    curblockpos = Vector{Int}(undef, N)
    for d in 1:N
        curblockpos[d] = nblocks[d]
    end

    curranges = Vector{AbstractRange{Int}}(undef, N)
    curidx = Vector{Vector{Int}}(undef, N)
    total_nblocks = prod(nblocks)

    handle = HDF5.h5open(x.file)
    dset = handle[x.name]

    for b in 1:total_nblocks
        # Incrementing the blocks.
        for d in 1:N
            if curblockpos[d] == nblocks[d]
                curblockpos[d] = 1
            else
                curblockpos[d] += 1
            end

            curblock = used_blocks[d][curblockpos[d]]
            curblockinfo = idx_by_block[d][curblock]
            curranges[d] = curblockinfo.range
            curidx[d] = curblockinfo.internal

            if curblockpos[d] != 1
                break
            end
        end

        block = getindex(dset, curranges...)

        # Extracting along the chosen block.
        curstarts = Vector{Int}(undef, N)
        curpos = Vector{Int}(undef, N)
        for d in 1:N
            curpos[d] = length(curidx[d])
            curstarts[d] = curranges[d][1] - 1
        end
        total = prod(curpos)
        innerpos = Vector{Int}(undef, N)
        realpos = Vector{Int}(undef, N)

        for v in 1:total
            for d in 1:N
                if curpos[d] == length(curidx[d])
                    curpos[d] = 1
                else
                    curpos[d] += 1
                end
                innerpos[d] = curidx[d][curpos[d]]
                realpos[d] = innerpos[d] + curstarts[d]
                if curpos[d] != 1
                    break
                end
            end

            val = getindex(block, innerpos...)
            store(realpos, val)
        end
    end
end

function create_back_mapping(indices)
    N = length(indices)
    collated = Vector{Dict{Int,Vector{Int}}}(undef, N)
    for d in 1:N
        current = Dict{Int,Vector{Int}}()
        curindices = indices[d]
        for i in 1:length(curindices)
            curdex = curindices[i]
            if haskey(current, curdex)
                push!(current[curdex], i)
            else
                current[curdex] = [i]
            end
        end
        collated[d] = current
    end
    return collated
end

function sorted_unique_indices(backmap)
    N = length(backmap)
    my_indices = Vector{AbstractVector{Int}}(undef, N)
    for d in 1:N
        current = backmap[d]
        used = collect(keys(current))
        sort!(used)
        my_indices[d] = used
    end
    return my_indices
end

"""
    extractdense(x::DenseHDF5Array{T,N}, I...; blockdim = nothing)

Extract an in-memory dense `Array` from a `DenseHDF5Array`.
The returned array contains the same values as `x[I...]`.

For arbitrary indices, this function performs block-by-block extraction to reduce memory usage.
The size of each block is determined by `blockdim`, which should be a tuple of length equal to the number of dimensions.
If `blockdim` is not specified, the chunk dimensions are used instead; if the dataset is stored in a contiguous layout, arbitrary block dimensions are used.

Some optimization is applied if `I` only consists of `AbstractRange{Int}` values.

# Examples
```jldoctest
julia> using HDF5Arrays

julia> tmp = tempname();

julia> exampledense(tmp, "stuff", (20, 10))

julia> x = DenseHDF5Array(tmp, "stuff");

julia> y = extractdense(x, 1:5, 1:5);

julia> size(y)
(5, 5)

julia> y2 = extractdense(x, [1,3,5,7], [2,4,6,8,10]);

julia> size(y2)
(4, 5)
```
"""
function extractdense(x::DenseHDF5Array{T,N}, I...; blockdim = nothing) where {T,N}
    indices = to_indices(x, I)

    # Directly extracting if we have an abstract range.
    all_ars = true;
    for d in 1:N
        if !isa(indices[d], AbstractRange{Integer})
            all_ars = false
            break
        end
    end

    if all_ars
        handle = HDF5.h5open(x.file)
        dset = handle[x.name]
        return getindex(dset, indices...)
    end

    # Converting into a dictionary of back-mapped indices.
    collated = create_back_mapping(indices)
    my_indices = sorted_unique_indices(collated)

    # Creating a function to insert new values into the specified locations.
    lengths = [length(x) for x in indices]
    output = Array{T,N}(undef, (lengths...,))
    FUN = function (pos, val)
        setdex = Vector{AbstractVector{Int}}(undef, N)
        for d in 1:N
            setdex[d] = collated[d][pos[d]]
        end
        output[setdex...] .= val
    end

    # Performing the extraction.
    general_dense_extractor(x, (my_indices...,), FUN; blockdim = blockdim)
    return output
end

function create_csc_matrix(NR::Int, NC::Int, collated_rows::Dict{Int,Vector{Int}}, collated_cols::Dict{Int,Vector{Int}}, store::Vector{Dict{Int,T}}) where {T}
    rows = Vector{Vector{Int}}(undef, NC)
    vals = Vector{Vector{T}}(undef, NC)
    for (kc, vc) in collated_cols
        curcol = store[kc]

        tmp = Vector{Pair{Int,T}}()
        for (kr, vr) in curcol
            for r in collated_rows[kr]
                push!(tmp, Pair(r, vr))
            end
        end

        sort!(tmp)
        tmpr = [x.first for x in tmp]
        tmpv = [x.second for x in tmp]

        for c in vc
            rows[c] = tmpr
            vals[c] = tmpv
        end

        store[kc] = Dict{Int,T}() # let the GC clear the values.
    end

    ptrs = Vector{Int}(undef, NC + 1)
    ptrs[1] = 1
    for i in 1:NC
        ptrs[i + 1] = ptrs[i] + length(rows[i])
    end

    all_rows = vcat(rows...)
    rows = nothing # let the GC eat this up
    all_vals = vcat(vals...)
    vals = nothing

    return SparseArrays.SparseMatrixCSC{T,Int}(NR, NC, ptrs, all_rows, all_vals)
end

"""
    extractsparse(x::DenseHDF5Array{T,N}, i, j; blockdim = nothing)

Extract an in-memory sparse matrix from a 2-dimensional `DenseHDF5Array`.
The returned matrix contains the same values as `x[i, j]`.
This assumes that the type is either numeric or boolean,
and is only useful when there is a high proportion of zero values in `x`.

For arbitrary indices, this function performs block-by-block extraction to reduce memory usage.
The size of each block is determined by `blockdim`, which should be a tuple of length equal to the number of dimensions.
If `blockdim` is not specified, the chunk dimensions are used instead; if the dataset is stored in a contiguous layout, arbitrary block dimensions are used.

# Examples
```jldoctest
julia> using HDF5Arrays

julia> tmp = tempname();

julia> exampledense(tmp, "stuff", (20, 10); density = 0.2)

julia> x = DenseHDF5Array(tmp, "stuff");

julia> y = extractsparse(x, 1:5, 1:5);

julia> typeof(y)
SparseArrays.SparseMatrixCSC{Float64, Int64}
```
"""
function extractsparse(x::DenseHDF5Array{T,2}, i, j; blockdim = nothing) where {T<:Union{Number,Bool}}
    indices = to_indices(x, (i, j))
    NR = length(indices[1])
    NC = length(indices[2])

    # Converting into a dictionary of back-mapped indices.
    collated = create_back_mapping(indices)
    my_indices = sorted_unique_indices(collated)

    # Creating a function to insert new values into the specified locations.
    store = Vector{Dict{Int, T}}(undef, x.dims[2])
    for d in my_indices[2]
        store[d] = Dict{Int,T}()
    end

    FUN = function (position, val)
        if val != 0
            store[position[2]][position[1]] = val
        end
    end

    general_dense_extractor(x, (my_indices...,), FUN; blockdim = blockdim)

    return create_csc_matrix(NR, NC, collated[1], collated[2], store)
end


text:
export BASE_DIR
export minimum, maximum
export add_measurement!
export logpdf, gradlogpdf, heslogpdf
export solve_mc_se, solve_acp_mc_se, solve_acr_mc_se, solve_ivr_mc_se
export rm_enwl_transformer!, reduce_enwl_lines_eng!
export add_measurements!, write_measurements!, assign_load_pseudo_measurement_info!
export assign_unique_individual_criterion!, assign_basic_individual_criteria!
export assign_start_to_variables!, assign_residual_ub!
export reduce_single_phase_loadbuses!
export calculate_voltage_magnitude_error
export update_load_bounds!, update_voltage_bounds!, update_generator_bounds!, update_all_bounds!
export ExtendedBeta

ext: jl


text:
# working with text data ...
#
# (c) A. Dominik, 2021

"""
    mutable struct WordTokenizer
        len
        w2i
        i2w
    end

Create a word-based vocabulary: every unique word of a String or
a list of Strings is assigned to a unique number.
The created object includes a list of words (`i2w`, ordered by their numbers) and
a dictionary `w2i` with the words as keys.

### Constructor:

    function WordTokenizer(texts; len=nothing, add_ctls=true)

With arguments:
+ `texts`: `AbstractArray` or iterable collection of `AbstractArray`s to be
        analysed.
+ `len=nothing`: maximum number of different words in the vocabulary.
        Additional words in texts will be encoded as unknown. If `nothing`,
        all words of the texts are included.
+ `add_ctls=true`: if true, control words are added in front of the vocabulary
        (extending the maximum length by 4): `"<start>"=>1`, `"<end>"=>2`,
        `"<pad>"=>3` and `"<unknown>"=>4`.

### Signatures:

    function (t::WordTokenizer)(w::T; split_words=false, add_ctls=false)
                                where {T <: AbstractString}

Encode a word and return the corresponding number in the vocabulary or
the highest number (i.e. `"<unknown>"`) if the word is not in the vocabulary.

The encode-signature accepts the keyword arguments `split_words` and
`add_ctls`. If `split_words==true`, the input is treated as a sentence
and splitted into single words and an array of integer with the encoded
sequence is returned. If `add_ctls==true` the sequence will be framed
by `<start>` and `<end>` tokens.

    function (t::WordTokenizer)(i::Int)

Decode a word by returning the word corresponding to `i` or
"<unknown>" if the number is out of range of the vocabulary.

    function (t::WordTokenizer)(s::AbstractArray{T}; add_ctls=false)
                               where {T <: AbstractString}

Called with an Array of Strings the tokeniser splits the strings
into words and returns an Array of `Array{Int}` with each of the
input strings represented by a sequence of Integer values.

    function (t::WordTokenizer)(seq::AbstractArray{T}; add_ctls=false)
                                     where {T <: Int}

Called with an Array of Integer values a single string  is returned
with the decoded token-IDs as words (space-separated).

### Examples:

    julia> vocab = WordTokenizer(["I love Julia", "They love Python"]);
    Julia> vocab(8)
    "Julia"

    julia> vocab("love")
    5

    julia> vocab.(split("I love Julia"))
    3-element Array{Int64,1}:
     5
     6
     8

    julia> vocab.i2w
    9-element Array{String,1}:
     "<start>"
     "<end>"
     "<pad>"
     "<unknown>"
     "love"
     "I"
     "They"
     "Julia"
     "Python"

    julia> vocab.w2i
    Dict{String,Int64} with 9 entries:
      "I"         => 6
      "<end>"     => 2
      "<pad>"     => 3
      "They"      => 7
      "Julia"     => 8
      "love"      => 5
      "Python"    => 9
      "<start>"   => 1
      "<unknown>" => 4

    julia> vocab.([7,5,8])
    3-element Array{String,1}:
     "They"
     "love"
     "Julia

    julia> vocab.("I love Scala", split_words=true)
    3-element Array{Int64,1}:
     6
     5
     4

    julia> vocab.([6,5,4])
    3-element Array{String,1}:
     "I"
     "love"
     "<unknown>"

    julia> vocab("I love Python", split_words=true, add_ctls=true)
    5-element Array{Int64,1}:
     1
     6
     5
     9
     2

    julia> vocab(["They love Julia", "I love Julia"])
    2-element Array{Array{Int64,1},1}:
     [7, 5, 8]
     [6, 5, 8]
"""
mutable struct WordTokenizer
    len
    w2i
    i2w
end

# QnD cleaning of a sentence before creating tokens.
# + normalise
# + remove punctuation
# + remove duplicate spaces
# + trim
#
function clean_sentence(s)

    s = Unicode.normalize(s)
    s = replace(s, Regex("[.!?,;:\"\']") => " ")
    s = replace(s, Regex(" {2,}") => " ")
    s = replace(s, Regex("^ ") => "")
    s = replace(s, Regex(" \$") => "")
    return s
end

function WordTokenizer(texts; len=nothing, add_ctls=true)

    if texts isa AbstractString
        texts = [texts]
    end

    words = []
    for t in texts
        t = clean_sentence(t)
        append!(words, split(t, " "))
    end

    # make vocab:
    #
    counts = Dict{String, Int}()
    for w in words
        if haskey(counts, w)
            counts[w] += 1
        else
            counts[w] = 1
        end
    end
    pairs = sort(collect(counts), by=x->x[2], rev=true)

    # limit vocab length to len:
    #
    if len != nothing && len < length(counts)
        pairs = pairs[1:len]
    end

    # add control tokens:
    #
    if add_ctls
        pushfirst!(pairs, "<start>" => 0,
                        "<end>" => 0,
                        "<pad>" => 0,
                        "<unknown>" => 0)
    end

    # make encode and decode data structures:
    #
    w2i = Dict{String, Int}()
    i2w = String[]

    for (i, w) in enumerate(pairs)
        push!(i2w, w[1])
        w2i[w[1]] = i
    end

    return WordTokenizer(length(w2i), w2i, i2w)
end

function (t::WordTokenizer)(i::Int)
    if i > t.len
        return t.i2w[end]
    else
        return t.i2w[i]
    end
end

function (t::WordTokenizer)(w::T; split_words=false,
                add_ctls=false) where {T <: AbstractString}

    w = clean_sentence(w)
    # tokenise a word or a complete string:
    #
    if !split_words
        if haskey(t.w2i, w)
            return t.w2i[w]
        else
            return t("<unknown>")
        end
    else
        s = split(w, " ")
        st = t.(s)
        if add_ctls
            st = vcat(t("<start>"), st, t("<end>"))
        end
        return st
    end
end

function (t::WordTokenizer)(s::AbstractArray{T}; add_ctls=false) where {T <: AbstractString}

    # return a list of sequences:
    #
    return [t(w; split_words=true, add_ctls=add_ctls) for w in s]
end

function (t::WordTokenizer)(seq::AbstractArray{T}; add_ctls=false) where {T <: Int}

    # return a list of sequences:
    #
    return join(t.(seq), " ")
end

"""
    function get_tatoeba_corpus(lang; force=false,
                url="https://www.manythings.org/anki/")

Download and read a bilingual text corpus from Tatoeba (privided)
by ManyThings (https://www.manythings.org).
All corpi are English-*Language*-pairs with different size and
quality. Considerable languages include:
+ `fra`: French-English, 180 000 sentences
+ `deu`: German-English, 227 000 sentences
+ `heb`: Hebrew-English, 126 000 sentences
+ `por`: Portuguese-English, 170 000 sentences
+ `tur`: Turkish-English, 514 000 sentences

The function returns two lists with corresponding sentences in both
languages. Sentences are are *not* processed/normalised/cleaned, but
exactly as provided by Tatoeba.

The data is stored in the package directory and only downloaded once.

### Arguments:
+ `lang`: languagecode
+ `force=false`: if `true`, the corpus is downloaded even if
        a data file is already saved.
+ `url`: base url of ManyThings.
"""
function get_tatoeba_corpus(lang; force=false,
                url="https://www.manythings.org/anki/")

    @show dir = normpath(joinpath(dirname(pathof(@__MODULE__)),
                "..", "data", "Tatoeba"))
    if !ispath(dir)
        mkpath(dir)
    end
    fname = join([lang, "-eng.zip"])
    @show pathname = joinpath(dir, fname)

    # download if necessary:
    #
    if !isfile(pathname) || force
        url = join([url, fname])
        println("Downloading Tatoeba corpus for language $lang")
        println("from $url")
        download(url, pathname)
    else
        println("Corpus for language $lang is already downloaded.")
    end

    if !isfile(pathname)
        println("File $pathname not found!")
    end

    # read zipped file:
    #
    println("Reading Tatoeba corpus for languages en-$lang")
    z = ZipFile.Reader(pathname)
    en, lang = String[], String[]

    for f in z.files
        for (i,line) in enumerate(eachline(f))
            if i % 1000 == 0
                print("\rimporting sentences: $i")
            end
            splits = split(line, "\t")
            if length(splits) > 1
                enl, langl = splits[1:2]
                push!(en, enl)
                push!(lang, langl)
            end
        end
    end
    close(z)

    return en, lang
end

"""
    function seq_minibatch(x, [y,] batchsize; seq_len=nothing, pad=0, o...)

Return an iterator of type `Knet.Data` with sequence minibatches from a
list of sequences.
all
keyword args of [`Knet.minibatch()`](https://denizyuret.github.io/Knet.jl/latest/reference/#Knet.Train20.minibatch) can be used.

All sequences in x are brought to the same length by truncating (if too long)
or padding with the token provided as `pad`.

If `y` is defined, the minibatches include the sequences for x and
training targets `y`, given as n-dimensional array (as for `Knet.minibach()`).
For sequence-2-sequence minibatches the function `seq2seq_minibatch()`
must be used.

### Arguments:
+ `x`: An iterable object of sequences.
+ `y`: vector or array with training targets
+ `batchsize`: size of minibatches
+ `seq_len=nothing`: demanded length of sequences in the minibatches.
        If `nothing`, all sequences are padded to match with the longest
        sequence.
+ `pad=0`: token, used for padding. The token must be compatible
        with the type of the sequence elements.
+ `o...`: any other keyword arguments of `Knet.minibatch()`, such as
        `shuffle=true` or `partial=true` can be provided.
"""
function seq_minibatch(x, y, batchsize; seq_len=nothing, pad=0, o...)

    if seq_len == nothing
        seq_len = maximum(length.(x))
    end

    x = pad_sequences(x, seq_len, pad)
    return Knet.minibatch(x, y, batchsize; o...)
end

function seq_minibatch(x, batchsize; seq_len=nothing, pad=0, o...)

    if seq_len == nothing
        seq_len = maximum(length.(x))
    end

    x = pad_sequences(x, seq_len, pad)
    return Knet.minibatch(x, batchsize; o...)
end

"""
    function seq2seq_minibatch(x, y, batchsize; seq_len=nothing,
                pad_x=0, pad_y=0, o...)

Return an iterator of type `Knet.Data` with (x,y) sequence minibatches from
two lists of sequences.
all
keyword args of [`Knet.minibatch()`](https://denizyuret.github.io/Knet.jl/latest/reference/#Knet.Train20.minibatch) can be used.

All sequences in x and y are brought to the same length
by truncating (if too long)
or padding with the token provided as `pad`.

### Arguments:
+ `x`: An iterable object of sequences.
+ `y`: An iterable object of target sequences.
+ `batchsize`: size of minibatches
+ `seq_len=nothing`: demanded length of sequences in the minibatches.
        If `nothing`, all sequences are padded to match with the longest
        sequence.
+ `pad_x=0`,
+ `pad_y=0`: token, used for padding. The token must be compatible
        with the type of the sequence elements.
+ `o...`: any other keyword arguments of `Knet.minibatch()`, such as
        `shuffle=true` or `partial=true` can be provided.
"""
function seq2seq_minibatch(x, y, batchsize; seq_len=nothing,
                           pad_x=0, pad_y=0, o...)

    if seq_len == nothing
        seq_len = maximum((maximum(length.(x)), maximum(length.(y))))
    end

    x = pad_sequences(x, seq_len, pad_x)
    y = pad_sequences(y, seq_len, pad_y)

    return Knet.minibatch(x, y, batchsize; o...)
end

function pad_sequences(s, len, pad)

    elem_type = typeof(s[1][1])
    data = Array{elem_type}(undef, len, length(s))

    for (i,seq) in enumerate(s)
        if length(seq) > len        # if too long
            seq = seq[1:len]
        end
        while length(seq) < len  # if too short
            push!(seq, pad)
        end
        data[:,i] = seq
    end
    return data
end


text:
export code, code_track, recode_stochastic, FunctionalCoder

"""
    code(c, xz, z, x; kws...)

Perform the coding operation specified by `c`.

# Arguments
- `c`: Coder.
- `xz`: Inputs of the functional representation.
- `z`: Outputs of the functional representation.
- `x`: Target inputs.

# Fields
- `kws...`: Further keywords to pass on.

# Returns
- `Tuple`: Tuple containing the inputs and outputs of a functional representation.
"""
code(c, xz, z, x; kws...) = xz, c(z)

# Implement `code` for `Poolings`s to discard the inputs.
code(c::Pooling, xz, z, x; kws...) = nothing, c(z)

function code(c::Chain, xz, z, x; kws...)
    for cᵢ in c
        xz, z = code(cᵢ, xz, z, x; kws...)
    end
    return xz, z
end

function code(p::Parallel, xz, z, x; kws...)
    xz, z = zip([code(pᵢ, xz, z, x; kws...) for pᵢ in p]...)
    return Parallel(xz...), Parallel(z...)
end

function code(p::Parallel{N}, xz, z::Parallel{N}, x; kws...) where N
    xz, z = zip([code(pᵢ, xz, zᵢ, x; kws...) for (pᵢ, zᵢ) in zip(p, z)]...)
    return Parallel(xz...), Parallel(z...)
end

function code(p::Parallel{N}, xz::Parallel{N}, z::Parallel{N}, x; kws...) where N
    xz, z = zip([code(pᵢ, xzᵢ, zᵢ, x; kws...) for (pᵢ, xzᵢ, zᵢ) in zip(p, xz, z)]...)
    return Parallel(xz...), Parallel(z...)
end

"""
    code_track(c, xz, z, x; kws...)

Perform the coding operation specified by `c` whilst keeping track of the sequence of target
inputs, called the history. This history can be used to perform the coding operation again
at that sequence of target inputs exactly.

# Arguments
- `c`: Coder.
- `xz`: Inputs of the functional representation.
- `z`: Outputs of the functional representation.
- `x`: Target inputs.

# Fields
- `kws...`: Further keywords to pass on.

# Returns
- `Tuple`: Tuple containing the inputs and outputs of a functional representation and
    the history of target inputs.
"""
code_track(c, xz, z, x; kws...) = code_track(c, xz, z, x, []; kws...)

function code_track(c, xz, z, x, h; kws...)
    xz, z = code(c, xz, z, x; kws...)
    return xz, z, vcat(h, [x])
end

function code_track(c::Chain, xz, z, x, h; kws...)
    for cᵢ in c
        xz, z, h = code_track(cᵢ, xz, z, x, h; kws...)
    end
    return xz, z, h
end

function code_track(p::Parallel, xz, z, x, h; kws...)
    xz, z, hists = zip([code_track(pᵢ, xz, z, x, []; kws...) for pᵢ in p]...)
    return Parallel(xz...), Parallel(z...), vcat(h, Parallel(hists...))
end

function code_track(p::Parallel{N}, xz, z::Parallel{N}, x, h; kws...) where N
    xz, z, hists = zip([
        code_track(pᵢ, xz, zᵢ, x, []; kws...) for (pᵢ, zᵢ) in zip(p, z)
    ]...)
    return Parallel(xz...), Parallel(z...), vcat(h, Parallel(hists...))
end

function code_track(p::Parallel{N}, xz::Parallel{N}, z::Parallel{N}, x, h; kws...) where N
    xz, z, hists = zip([
        code_track(pᵢ, xzᵢ, zᵢ, x, []; kws...) for (pᵢ, xzᵢ, zᵢ) in zip(p, xz, z)
    ]...)
    return Parallel(xz...), Parallel(z...), vcat(h, Parallel(hists...))
end

"""
    recode(c, xz, z, h; kws...)

Reperform the coding operation specified by `c` at a given sequence of target inputs `c`,
called the history.

# Arguments
- `c`: Coder.
- `xz`: Inputs of the functional representation.
- `z`: Outputs of the functional representation.
- `h`: Sequence of target inputs.

# Returns
- `Tuple`: Tuple containing the inputs and outputs of a functional representation and
    the remaining history of target inputs.
"""
function recode(c, xz, z, h; kws...)
    xz, z = code(c, xz, z, h[1]; kws...)
    return xz, z, h[2:end]
end

function recode(c::Chain, xz, z, h; kws...)
    for cᵢ in c
        xz, z, h = recode(cᵢ, xz, z, h; kws...)
    end
    return xz, z, h
end

function recode(p::Parallel, xz, z, h; kws...)
    xz, z, _ = zip([
        recode(pᵢ, xz, z, hᵢ; kws...) for (pᵢ, hᵢ) in zip(p, h[1])
    ]...)
    return Parallel(xz...), Parallel(z...), h[2:end]
end

function recode(p::Parallel{N}, xz, z::Parallel{N}, h; kws...) where N
    xz, z, _ = zip([
        recode(pᵢ, xz, zᵢ, hᵢ; kws...) for (pᵢ, zᵢ, hᵢ) in zip(p, z, h[1])
    ]...)
    return Parallel(xz...), Parallel(z...), h[2:end]
end

function recode(p::Parallel{N}, xz::Parallel{N}, z::Parallel{N}, h; kws...) where N
    xz, z, _ = zip([
        recode(pᵢ, xzᵢ, zᵢ, hᵢ; kws...) for (pᵢ, xzᵢ, zᵢ, hᵢ) in zip(p, xz, z, h[1])
    ]...)
    return Parallel(xz...), Parallel(z...), h[2:end]
end

"""
    recode_stochastic(
        coders::Parallel{N},
        codings::Parallel{N},
        xc,
        yc,
        h;
        kws...
    ) where N

In an existing aggregate coding `codings`, recode the codings that are not `Dirac`s for
a new context set.

# Arguments
- `coders::Parallel{N}`: Parallel of coders that produced the coding.
- `codings::Parallel{N}`: Parallel of codings.
- `xc`: Locations of new context set.
- `yc`: Observed values of new context set.
- `h`: History to replay.

# Fields
- `kws...`: Further keywords to pass on.

# Returns
- `Parallel`: Updated coding.
"""
recode_stochastic(
    coders::Parallel{N},
    codings::Parallel{N},
    xc,
    yc,
    h;
    kws...
) where N = Parallel([
    recode_stochastic(coder, coding, xc, yc, hᵢ; kws...)
    for (coder, coding, hᵢ) in zip(coders, codings, h[1])
]...)

# Do not recode `Dirac`s.

recode_stochastic(coder, coding::Dirac, xc, yc, h; kws...) = coding

# If the coding is aggregate, it can still contain `Dirac`s, so be careful.

recode_stochastic(coder, coding, xc, yc, h; kws...) =
    _choose(second(recode(coder, xc, yc, h; kws...)), coding)

_choose(new::Parallel{N}, old::Parallel{N}) where N =
    Parallel([_choose(newᵢ, oldᵢ) for (newᵢ, oldᵢ) in zip(new, old)]...)
_choose(new::Dirac, old::Dirac) = old
_choose(new::Gaussian, old::Gaussian) = new

"""
    struct Materialise

A coder that materialises a parallel of things.

# Fields
- `f_x`: Strategy for materialising the inputs. See `materialise`. Defaults to
    `repeat_merge` over the second dimension.
- `f_y`: Strategy for materialising the outputs. See `materialise`. Defaults to
    `repeat_cat` over the second dimension.
"""
struct Materialise
    f_x
    f_y
end

@Flux.functor Materialise

Materialise() = Materialise(
    xs -> repeat_merge(xs..., dims=2),
    xs -> repeat_cat(xs..., dims=2)
)

code(c::Materialise, xz, z, x; kws...) = materialise(xz, c.f_x), materialise(z, c.f_y)


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct VirtualMachineExtensionUpdateProperties <: SwaggerModel
    forceUpdateTag::Any # spec type: Union{ Nothing, String } # spec name: forceUpdateTag
    publisher::Any # spec type: Union{ Nothing, String } # spec name: publisher
    type::Any # spec type: Union{ Nothing, String } # spec name: type
    typeHandlerVersion::Any # spec type: Union{ Nothing, String } # spec name: typeHandlerVersion
    autoUpgradeMinorVersion::Any # spec type: Union{ Nothing, Bool } # spec name: autoUpgradeMinorVersion
    settings::Any # spec type: Union{ Nothing, Any } # spec name: settings
    protectedSettings::Any # spec type: Union{ Nothing, Any } # spec name: protectedSettings

    function VirtualMachineExtensionUpdateProperties(;forceUpdateTag=nothing, publisher=nothing, type=nothing, typeHandlerVersion=nothing, autoUpgradeMinorVersion=nothing, settings=nothing, protectedSettings=nothing)
        o = new()
        validate_property(VirtualMachineExtensionUpdateProperties, Symbol("forceUpdateTag"), forceUpdateTag)
        setfield!(o, Symbol("forceUpdateTag"), forceUpdateTag)
        validate_property(VirtualMachineExtensionUpdateProperties, Symbol("publisher"), publisher)
        setfield!(o, Symbol("publisher"), publisher)
        validate_property(VirtualMachineExtensionUpdateProperties, Symbol("type"), type)
        setfield!(o, Symbol("type"), type)
        validate_property(VirtualMachineExtensionUpdateProperties, Symbol("typeHandlerVersion"), typeHandlerVersion)
        setfield!(o, Symbol("typeHandlerVersion"), typeHandlerVersion)
        validate_property(VirtualMachineExtensionUpdateProperties, Symbol("autoUpgradeMinorVersion"), autoUpgradeMinorVersion)
        setfield!(o, Symbol("autoUpgradeMinorVersion"), autoUpgradeMinorVersion)
        validate_property(VirtualMachineExtensionUpdateProperties, Symbol("settings"), settings)
        setfield!(o, Symbol("settings"), settings)
        validate_property(VirtualMachineExtensionUpdateProperties, Symbol("protectedSettings"), protectedSettings)
        setfield!(o, Symbol("protectedSettings"), protectedSettings)
        o
    end
end # type VirtualMachineExtensionUpdateProperties

const _property_map_VirtualMachineExtensionUpdateProperties = Dict{Symbol,Symbol}(Symbol("forceUpdateTag")=>Symbol("forceUpdateTag"), Symbol("publisher")=>Symbol("publisher"), Symbol("type")=>Symbol("type"), Symbol("typeHandlerVersion")=>Symbol("typeHandlerVersion"), Symbol("autoUpgradeMinorVersion")=>Symbol("autoUpgradeMinorVersion"), Symbol("settings")=>Symbol("settings"), Symbol("protectedSettings")=>Symbol("protectedSettings"))
const _property_types_VirtualMachineExtensionUpdateProperties = Dict{Symbol,String}(Symbol("forceUpdateTag")=>"String", Symbol("publisher")=>"String", Symbol("type")=>"String", Symbol("typeHandlerVersion")=>"String", Symbol("autoUpgradeMinorVersion")=>"Bool", Symbol("settings")=>"Any", Symbol("protectedSettings")=>"Any")
Base.propertynames(::Type{ VirtualMachineExtensionUpdateProperties }) = collect(keys(_property_map_VirtualMachineExtensionUpdateProperties))
Swagger.property_type(::Type{ VirtualMachineExtensionUpdateProperties }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_VirtualMachineExtensionUpdateProperties[name]))}
Swagger.field_name(::Type{ VirtualMachineExtensionUpdateProperties }, property_name::Symbol) =  _property_map_VirtualMachineExtensionUpdateProperties[property_name]

function check_required(o::VirtualMachineExtensionUpdateProperties)
    true
end

function validate_property(::Type{ VirtualMachineExtensionUpdateProperties }, name::Symbol, val)
end


text:
using TrajectoryOptimization
using StaticArrays
using LinearAlgebra
using Statistics
using TrajOptPlots
using GeometryTypes
using CoordinateTransformations
using BenchmarkTools
using ForwardDiff
using MeshCat
using Test
using InplaceOps
const TO = TrajectoryOptimization
import TrajectoryOptimization.Controllers: RBState

function gen_multisat_prob(Rot=MRP{Float64})
    # Test with rigid bodies
    model = Dynamics.FreeBody{Rot,Float64}()
    models = Dynamics.CopyModel(model, 2)

    # Test solve
    N = 101
    tf = 5.0

    x01 = Dynamics.build_state(model, [0, 1,0], UnitQuaternion(I), zeros(3), zeros(3))
    x02 = Dynamics.build_state(model, [0,-1,0], UnitQuaternion(I), zeros(3), zeros(3))
    x0 = [x01; x02]

    xf1 = Dynamics.build_state(model, [ 1,0,0], expm(deg2rad( 45)*@SVector [1,0,0.]), zeros(3), zeros(3))
    xf2 = Dynamics.build_state(model, [-1,0,0], expm(deg2rad(-45)*@SVector [1,0,0.]), zeros(3), zeros(3))
    xf = [xf1; xf2]

    Qd = Dynamics.fill_state(model, 1e-1, 1e-1, 1e-2, 1e-2)
    Q = Diagonal([Qd; Qd])
    Rd = @SVector fill(1e-2, 6)
    R = Diagonal([Rd; Rd])
    obj = LQRObjective(Q,R,Q*10,xf,N)

    prob = Problem(models, obj, xf, tf, x0=x0)
end
prob = gen_multisat_prob()
solver = iLQRSolver2(prob)
solver.opts.verbose =  false
model = solver.model
x,u = rand(model)

@time rollout!(solver)
@time cost(solver)
@time TO.state_diff_jacobian!(solver.G, solver.model, solver.Z)
@time TO.dynamics_expansion!(solver.D, solver.model, solver.Z)
@time TO.cost_expansion!(solver.Q, solver.obj, solver.Z)
@time TO.error_expansion!(solver.D, solver.model, solver.G)
@time TO.error_expansion!(solver.Q, solver.model, solver.Z, solver.G)
solver.Q[1].xx
solver.G[1]
@time ΔV = TO.static_backwardpass!(solver)
@time TO.forwardpass!(solver, ΔV, cost(solver))

prob = gen_multisat_prob()
solver = iLQRSolver2(prob)
solver.opts.verbose = true
solve!(solver)
cost(solver)

if !isdefined(Main,:vis)
    vis = Visualizer(); open(vis);
    set_mesh!(vis, model.model)
end
visualize!(vis, model, get_trajectory(solver))
settransform!(vis["robot"]["copy2"], Translation(-1,1,1))


text:
function default_callback_fn(state::StateAM, data::Data, prior::PriorAM,
                             iter::Int, pbar::MCMC.ProgressBars.ProgressBar)
  ll = round(loglike(state, data, prior.beta), digits=3)
  MCMC.ProgressBars.set_postfix(pbar, loglike=ll)
  return Dict(:loglike => ll)
end

function amfit(init::StateAM, data::Data, prior::PriorAM;
               tuner::Union{Nothing, MCMC.MvTuner}=nothing,
               nsamps::Vector{Int}=[1000], nburn::Int=1000, thin::Int=1,
               seed=nothing, verbose::Int=1, temper::Real=1, propcov=nothing,
               nc::Int=0, nt::Int=0, openblas_num_threads=1)

  LinearAlgebra.BLAS.set_num_threads(openblas_num_threads)
  seed == nothing || Random.seed!(seed)
  tuner == nothing && (tuner = MCMC.MvTuner(tovec(init)))

  # Print settings for sanity check.
  if verbose > 0
    println("seed: ", seed)
    println("temper: ", temper)
    println("openblas_num_threads: ", openblas_num_threads)
  end

  function update!(state::StateAM)
    _data = subsample_data(data, nc, nt)
    v = tovec(state)
    function logprob(v::Vector{Float64})
      s = deepcopy(state)
      fromvec!(s, v)
      ll = loglike(s, _data, prior.beta) * (data.NC + data.NT) / (_data.NC + _data.NT) / temper
      lp = logprior(s, prior)
      lj = logabsjacobian(v, prior)
      return ll + lp + lj
    end
    if propcov == nothing
      newv = MCMC.metropolisAdaptive(v, logprob, tuner)
    else
      newv = MCMC.metropolis(v, logprob, propcov)
    end
    newv != v && fromvec!(state, newv)
  end

  function _callback_fn(state::StateAM, iter::Int,
                        pbar::MCMC.ProgressBars.ProgressBar)
    # return default_callback_fn(state, data, prior, iter, pbar)
  end

  chain, laststate, summarystats = MCMC.gibbs(
      init, update!, nsamps=nsamps, nburn=nburn, thin=thin,
      monitors=[[:gammaC, :gammaT, :etaC, :etaT, :mu, :sigma, :nu, :phi]],
      callback_fn=_callback_fn)

  return chain, laststate, summarystats, tuner
end

max_stars_repo_name: luiarthur/CytofDensityEstimation

max_issues_repo_name: luiarthur/CytofDensityEstimation

max_forks_repo_name: luiarthur/CytofDensityEstimation


text:
function save_ivfadc_index(filename::AbstractString,
                           ivfadc::IVFADCIndex{U,I,Dc,Dr,T,Q}
                          ) where {U,I,Dc,Dr,T,Q<:NaiveQuantizer}
    open(filename, "w") do fid
        # Initialize all variables needed to write
        nrows, nclusters = size(ivfadc.coarse_quantizer.vectors)
        n = length(ivfadc)
        m = length(ivfadc.residual_quantizer.codebooks)
        k = ivfadc.residual_quantizer.k
        d = size(ivfadc.residual_quantizer.codebooks[1].vectors, 1)

        # Get types
        coarsequanttype_str = "NaiveQuantizer"
        quanttype_str = string(typeof(ivfadc.residual_quantizer.quantization))
        quanteltype_str = string(U)
        indextype_str = string(I)
        disttypecoarse_str = string(Dc)
        disttyperesidual_str = string(Dr)
        origeltype_str = string(T)

        # Start writing text information
        println(fid, "$nrows $nclusters")
        println(fid, "$n $m $k $d")
        println(fid, "$coarsequanttype_str")
        println(fid, "$quanttype_str")
        println(fid, "$quanteltype_str")
        println(fid, "$indextype_str")
        println(fid, "$disttypecoarse_str")
        println(fid, "$disttyperesidual_str")
        println(fid, "$origeltype_str")

        # Write coarse quantizer
        _write_naive_coarse_quantizer(fid, ivfadc.coarse_quantizer)

        # Write residual quantizer
        _write_residual_quantizer(fid, ivfadc.residual_quantizer, (nrows, nclusters))

        # Write inverse index
        _write_inverse_index(fid, ivfadc.inverse_index, (nrows, nclusters))
    end
end

_write_naive_coarse_quantizer(fid::IO, quantizer) = begin
    nclusters = size(quantizer.vectors, 2)
    for i in 1:nclusters
        write(fid, quantizer.vectors[:,i])
    end
end

_write_residual_quantizer(fid::IO, quantizer, dims) = begin
    nrows, nclusters = dims
    d = size(quantizer.codebooks[1].vectors, 1)
    m = length(quantizer.codebooks)
    for i in 1:m  # codebooks
        write(fid, quantizer.codebooks[i].codes)
        for j in 1:d
            write(fid, quantizer.codebooks[i].vectors[j,:])
        end
    end
    for i in 1:nrows  # rotation matrix
        write(fid, quantizer.rot[:,i])
    end
end

_write_inverse_index(fid::IO, inverse_index, dims) = begin
    nrows, nclusters = dims
    for i in 1:nclusters
        clsize = length(inverse_index[i].idxs)
        write(fid, clsize)
        write(fid, inverse_index[i].idxs)
        for j in 1:clsize
            write(fid, inverse_index[i].codes[j])
        end
    end
end

# Generate a WordVectors object from binary file
function load_ivfadc_index(filename::AbstractString)
    open(filename, "r") do fid
        nrows, nclusters = map(x -> parse(Int, x), split(readline(fid), ' '))
        n, m, k, d = map(x -> parse(Int, x), split(readline(fid), ' '))
        CQ = _read_type_from_line(readline(fid))
        Q = _read_type_from_line(readline(fid))
        U = eval(Symbol(readline(fid)))
        I = eval(Symbol(readline(fid)))
        Dc = _read_type_from_line(readline(fid))
        Dr = _read_type_from_line(readline(fid))
        T = eval(Symbol(readline(fid)))

        # Read coarse quantizer
        coarse_quantizer = _load_coarse_quantizer(CQ, fid;
                                dims=(nrows, nclusters),
                                disttype=Dc,
                                vectype=T)

        # Read residual quantizer
        cbooks = Vector{CodeBook{U,T}}(undef, m)
        codes_length = sizeof(U) * k
        codevecs_length = sizeof(T) * k
        for i = 1:m
            codes = collect(reinterpret(U, read(fid, codes_length)))
            vectors = Matrix{T}(undef, d, k)
            for j in 1:d
                vectors[j,:] = collect(reinterpret(T, read(fid, codevecs_length)))
            end
            cbooks[i] = CodeBook(codes, vectors)
        end
        binary_length = sizeof(T) * nrows
        rotmat = Matrix{T}(undef, nrows, nrows)
        for i in 1:nrows
            rotmat[:,i] = collect(reinterpret(T, read(fid, binary_length)))
        end
        residual_quantizer = ArrayQuantizer(Q(), (nrows, n), cbooks, k, Dr(), rotmat)

        # Read inverse index
        inverse_index = IVFADC.InvertedIndex{I,U}(undef, nclusters)
        encoded_length = sizeof(U) * m
        for i in 1:nclusters
            clsize = reinterpret(Int, read(fid, sizeof(Int)))[1]
            idxs_length = sizeof(I) * clsize
            idxs = collect(reinterpret(I, read(fid, idxs_length)))
            codes = Vector{Vector{U}}(undef, clsize)
            for j in 1:clsize
                codes[j] = collect(reinterpret(U, read(fid, encoded_length)))
            end
            inverse_index[i] = InvertedList(idxs, codes)
        end
        return IVFADCIndex(coarse_quantizer, residual_quantizer, inverse_index)
    end
end

_read_type_from_line(line) = begin
    if occursin('.', line)
        _module, _type = String.(split(line, "."))
    else
        _module, _type = string(@__MODULE__), line
    end
    return eval(Expr(:., Symbol(_module), QuoteNode(Symbol(_type))))
end

_load_coarse_quantizer(::Type{CQ},
                       fid::IO;
                       dims=(0,0),
                       disttype::Type{Dc}=typeof(DEFAULT_COARSE_DISTANCE),
                       vectype::Type{T}=AbstractFloat
                      ) where {CQ<:NaiveQuantizer,Dc,T} = begin
    nrows, nclusters = dims
    data = Matrix{T}(undef, nrows, nclusters)
    binary_length = sizeof(T) * nrows
    for i in 1:nclusters
        data[:,i] = collect(reinterpret(T, read(fid, binary_length)))
    end
    coarse_quantizer = IVFADC.NaiveQuantizer{Dc,T}(data)
end

function save_ivfadc_index(filename::AbstractString,
                           ivfadc::IVFADCIndex{U,I,Dc,Dr,T,Q}
                          ) where {U,I,Dc,Dr,T,Q<:HNSWQuantizer}
    open(filename, "w") do fid
        # Initialize all variables needed to write
        nrows = length(ivfadc.coarse_quantizer.hnsw.data[1])
        nclusters = length(ivfadc.coarse_quantizer.hnsw.data)
        n = length(ivfadc)
        m = length(ivfadc.residual_quantizer.codebooks)
        k = ivfadc.residual_quantizer.k
        d = size(ivfadc.residual_quantizer.codebooks[1].vectors, 1)

        # Get types
        coarsequanttype_str = "HNSWQuantizer"
        quanttype_str = string(typeof(ivfadc.residual_quantizer.quantization))
        quanteltype_str = string(U)
        indextype_str = string(I)
        disttypecoarse_str = string(Dc)
        disttyperesidual_str = string(Dr)
        origeltype_str = string(T)

        # Start writing text information
        println(fid, "$nrows $nclusters")
        println(fid, "$n $m $k $d")
        println(fid, "$coarsequanttype_str")
        println(fid, "$quanttype_str")
        println(fid, "$quanteltype_str")
        println(fid, "$indextype_str")
        println(fid, "$disttypecoarse_str")
        println(fid, "$disttyperesidual_str")
        println(fid, "$origeltype_str")

        # Write coarse quantizer
        _write_hnsw_coarse_quantizer(fid, ivfadc.coarse_quantizer.hnsw)

        # Write residual quantizer
        _write_residual_quantizer(fid, ivfadc.residual_quantizer, (nrows, nclusters))

        # Write inverse index
        _write_inverse_index(fid, ivfadc.inverse_index, (nrows, nclusters))
    end
end

_write_hnsw_coarse_quantizer(fid::IO,
                             hnsw::HierarchicalNSW{T,F,V,M}
                            ) where {T,F,V,M} = begin
    println(fid, string(T))
    nclusters= length(hnsw.data)

    # Start writing
    write(fid, hnsw.lgraph.M0)          # .lgraph
    write(fid, hnsw.lgraph.M)
    write(fid, hnsw.lgraph.m_L)
    nlinkedlists = length(hnsw.lgraph.linklist)
    write(fid, nlinkedlists)
    for i in 1:nlinkedlists
        listlength = length(hnsw.lgraph.linklist[i])
        write(fid, listlength)
        write(fid, hnsw.lgraph.linklist[i])
    end
    write(fid, hnsw.added)              # .added
    for i in 1:nclusters                # .data
        write(fid, hnsw.data[i])
    end
    write(fid, hnsw.ep)                 # .ep
    write(fid, hnsw.entry_level)        # .entry_level
    write(fid, hnsw.vlp.num_elements)   # .vlp
    nvlps = length(hnsw.vlp.pool)
    write(fid, nvlps)
    for i in 1:nvlps
        write(fid, hnsw.vlp.pool[i].visited_value)  # UInt8
        write(fid, length(hnsw.vlp.pool[i].list))
        write(fid, hnsw.vlp.pool[i].list)
    end
    # metric can be skipped             # .metric
    write(fid, hnsw.efConstruction)     # .efConstruction
    write(fid, hnsw.ef)                 # .ef
end

_load_coarse_quantizer(::Type{CQ},
                       fid::IO;
                       dims=(0,0),
                       disttype::Type{Dc}=typeof(DEFAULT_COARSE_DISTANCE),
                       vectype::Type{F}=AbstractFloat
                      ) where {CQ<:HNSWQuantizer,Dc,F} = begin
    # We already have F (float type) and M (metric type = Dc)
    T = eval(Symbol(readline(fid)))
    V = Vector{Vector{F}}
    intsize = sizeof(Int)
    nrows, nclusters = dims

    # .lgraph
    M0 = reinterpret(Int, read(fid, intsize))[1]
    M = reinterpret(Int, read(fid, intsize))[1]
    m_L = reinterpret(Float64, read(fid, 8))[1]  # a Float64
    nlinkedlists = reinterpret(Int, read(fid, intsize))[1]
    linklist = Vector{Vector{T}}(undef, nlinkedlists)
    for i in 1:nlinkedlists
        listlength = reinterpret(Int, read(fid, intsize))[1]
        linklist[i] = collect(reinterpret(T, read(fid, sizeof(T) * listlength)))
    end
    lgraph = HNSW.LayeredGraph{T}(linklist, M0, M, m_L)

    # .added
    added = collect(reinterpret(Bool, read(fid, nclusters * sizeof(Bool))))

    # .data
    data = V(undef, nclusters)
    for i in 1:nclusters
        data[i] = collect(reinterpret(F, read(fid, nrows * sizeof(F))))
    end

    # .ep
    ep = reinterpret(T, read(fid, sizeof(T)))[1]

    # .entry_level
    entry_level = reinterpret(Int, read(fid, intsize))[1]

    # .vlp
    num_elements = reinterpret(Int, read(fid, intsize))[1]
    nvlps = reinterpret(Int, read(fid, intsize))[1]
    pool = Vector{HNSW.VisitedList}(undef, nvlps)
    for i in 1:nvlps
        visited_value = read(fid, 1)[1]
        len = reinterpret(Int, read(fid, intsize))[1]
        list = read(fid, len)
        pool[i] = HNSW.VisitedList(list, visited_value)
    end
    vlp = HNSW.VisitedListPool(pool, num_elements)

    # .metric need not be written

    # .efConstruction
    efConstruction = reinterpret(Int, read(fid, intsize))[1]

    # .ef
    ef = reinterpret(Int, read(fid, intsize))[1]
    hnsw = HNSW.HierarchicalNSW{T,F,V,Dc}(lgraph, added, data,
                ep, entry_level, vlp, Dc(), efConstruction, ef)
    return HNSWQuantizer{T,V,Dc,F}(hnsw)
end


text:
# Script to replicate figure 2 and 3

using StatsBase, Statistics, LinearAlgebra, StatsPlots, XLSX, PrettyTables

# -------------- Previous scripts -----------------------------------------------------------------------------------------------------------------------

include("src/data_wrangling.jl") # formats the raw input data obtained from WIOD

# -------------- Sectoral aggregation for country level I/O-Table ---------------------------------------------------------------------------------------

# the function "country_level_GVC_indicators" uses the output from the function "obtain_matrices"
# it prepares the data to replicate figure 2 and 3 by aggregating from a country-sector level to a country level I/O-Table
# outputs:
# t_country_wide ... yearly statistics of final demand share, value added share and measures of U and D at country level

function country_level_GVC_indicators(df::DataFrame, year::Int64)

    ID, FD, GO, VA, IV, other = obtain_matrices(df, year) # N*S×N*S, N*S×N, N*S×1, N*S×1, N*S×N, N*S×1

    # aggregate to country level
    ID_ctry = [sum(ID[i:i+n_sectors-1, j:j+n_sectors-1]) for i in 1:n_sectors:n_sectors*n_ctrys, j in 1:n_sectors:n_sectors*n_ctrys] # N×N
    FD_ctry = [sum(FD[i:i+n_sectors-1, j]) for i in 1:n_sectors:n_sectors*n_ctrys, j in 1:n_ctrys] # N×N, sum over origin sector
    IV_ctry = [sum(IV[i:i+n_sectors-1, j]) for i in 1:n_sectors:n_sectors*n_ctrys, j in 1:n_ctrys] # N×N, sum over origin sector
    GO_ctry = [sum(GO[i:i+n_sectors-1]) for i in 1:n_sectors:n_sectors*n_ctrys] # N×1, sum over destination countries
    VA_ctry = [sum(VA[i:i+n_sectors-1]) for i in 1:n_sectors:n_sectors*n_ctrys] # N×1

    # net inventory correction
    iv_correction = [GO_ctry[i]/(GO_ctry[i]-IV_ctry[i, j]) for i in 1:n_ctrys, j in 1:n_ctrys] # N×N, to match country wise inventory changes

    ID_ctry = ID_ctry .* iv_correction # N×N, element wise multiplication
    FD_ctry = FD_ctry .* iv_correction # N×N

    FD_ctry = [sum(FD_ctry[i, :]) for i in 1:n_ctrys] # N×1, sum over destination countries

    # input/output coefficient matrices
    A = [ID_ctry[i, j] / GO_ctry[j] for i in 1:n_ctrys, j in 1:n_ctrys] # N×N, iterate output over destination country
    B = [ID_ctry[i, j] / GO_ctry[i] for i in 1:n_ctrys, j in 1:n_ctrys] # N×N, iterate output over reporting country

    # country level statistics
    FD_GO = FD_ctry ./ GO_ctry # N×1
    VA_GO = VA_ctry ./ GO_ctry # N×1
    U_num = inv(I - A) * GO_ctry # N×1, numerator of equation (5)
    D_num = transpose(GO_ctry) * inv(I - B) # N×1, numerator of equation (7), notice difference in matrix algebra

    U = U_num ./ GO_ctry # N×1
    D = transpose(D_num) ./ GO_ctry # N×1, transpose to have column vector

    # weighted average of all countries to create a world GVC position
    FD_GO_world = sum(FD_GO .* GO_ctry) / sum(GO_ctry)
    VA_GO_world = sum(VA_GO .* GO_ctry) / sum(GO_ctry)
    U_world = sum(U .* GO_ctry) / sum(GO_ctry)
    D_world = sum(D .* GO_ctry) / sum(GO_ctry)

    # collect data in a nice dataframe
    summary_stats = hcat(fill(year, size(ctrys)), ctrys, FD_GO, VA_GO, U, D)
    t_country_wide = DataFrame(summary_stats, :auto)
    push!(t_country_wide, [year "world" FD_GO_world VA_GO_world U_world D_world]) # add data for world
    rename!(t_country_wide, [:year, :country, :FD_GO, :VA_GO, :U, :D])
    transform!(t_country_wide, [:FD_GO, :VA_GO, :U, :D] .=> ByRow(x -> round(x, digits=3)) .=> [:FD_GO, :VA_GO, :U, :D]) # rescale U, D for better comparison

    return t_country_wide

end

# gather data for all years in a single data frame
t_fig_2_3 = DataFrame(year=Int[], country=String[], FD_GO=Float64[], VA_GO=Float64[], U=Float64[], D=Float64[])

for i in years
    append!(t_fig_2_3, country_level_GVC_indicators(df, i))
end

XLSX.writetable("clean/t_fig_2_3.xlsx", t_fig_2_3, overwrite=true) # export table to folder as input for plotting figure 2 and 3

# -------------- Figure 2 --------------------------------------------------------------------------------------------------------------------------------

t_fig_2_3 = DataFrame(XLSX.readtable("clean/t_fig_2_3.xlsx", "Sheet1")...) # load directly from file
transform!(t_fig_2_3, :country .=> ByRow(string) .=> :country, renamecols=false) # change type to String
transform!(t_fig_2_3, :year .=> ByRow(Int) .=> :year, renamecols=false) # change type to Int64
transform!(t_fig_2_3, [:FD_GO, :VA_GO, :U, :D] .=> ByRow(Float64) .=> [:FD_GO, :VA_GO, :U, :D], renamecols=false) # change type to Float64

t_fig_2_3_long = stack(t_fig_2_3, Not([:year, :country])) # long format for plotting

function plot_fig2(df::DataFrame, country::Vector, series::Vector) # use vectors as input so we can plot multiple countries for comparison

    df = subset(df, :country => ByRow(x -> x in country), :variable => ByRow(x -> x in series))

    p = @df df plot(:year, :value, group=(:country, :variable), lw=2, legend=:topleft,
        ylims=(round(minimum(:value), digits=1)-0.15,round(maximum(:value), digits=1)+0.15))

    return p
end

l = @layout [a ; b] # 2×1 layout
p1 = plot_fig2(t_fig_2_3_long, ["world"], ["FD_GO", "VA_GO"]) # add both measures (should be the same in theory) - slight difference (should adjust for "others"?)
p2 = plot_fig2(t_fig_2_3_long, ["world"], ["U", "D"]) # coincide very well
title = "Figure 2: GVC Positioning over Time (World Average)"
p_fig2 = plot(p1, p2, layout=l, plot_title=title, plot_titlefontsize=10)
savefig(p_fig2, "images/figure2.png") # export image to folder

# -------------- Figure 3 ------------------------------------------------------------------------------------------------------------------------------

# computing percentiles per country-series group
t_fig3 = subset(t_fig_2_3_long, :country => ByRow(x -> x != "world")) # take out world to calculate percentiles
gdf = groupby(t_fig3, [:year, :variable])
begin
    t_fig3 = combine(gdf,
     :value => (x -> percentile(x, 25)) => :quantile_25,
     :value => (x -> percentile(x, 50)) => :quantile_50,
     :value => (x -> percentile(x, 75)) => :quantile_75)
end
rename!(t_fig3, :variable => :series) # name conflict in formatting
df_fig3_long = stack(t_fig3, Not([:year, :series]))

function plot_fig3(df::DataFrame, series::String) # input vectors in so we can plot multiple countries for comparison

    df = subset(df, :series => ByRow(x -> x == series))

    p = @df df plot(:year, :value, group=:variable, lw=2, legend=:none,
    ylims=(round(minimum(:value)-0.05, digits=1),round(maximum(:value)+0.05, digits=1)),
    ylabel="$(series) over time")

    return p
end

l = @layout [a b ; c d] # 2×2 layout
p1 = plot_fig3(df_fig3_long, "FD_GO")
p2 = plot_fig3(df_fig3_long, "VA_GO")
p3 = plot_fig3(df_fig3_long, "U")
p4 = plot_fig3(df_fig3_long, "D")
title = "Figure 3: GVC Positioning over Time (25th, 50th, 75th country percentiles)"
p_fig3 = plot(p1, p2, p3, p4, layout=l, plot_title=title, plot_titlefontsize=10)
savefig(p_fig3, "images/figure3.png") # export image to folder

# -------------- Table 1 ------------------------------------------------------------------------------------------------------------------------------

t_tab1 = subset(t_fig_2_3_long, :country => ByRow(x -> x != "world")) # take out world
gdf = groupby(t_tab1, [:year, :variable])
#t_tab1 = combine(gdf, :value => (x -> ordinalrank(x; lt=isless)) => :rank) # ordinal rank per year and series
t_tab1 = transform(gdf, :value => (x -> ordinalrank(x; lt=isless)) => :rank) # ordinal rank per year and series
# notice the difference between combine and transform! => transform returns all other columns as well (very useful!)
sort!(t_tab1, [:year, :variable, :rank]) # sorting is still respecting groups

# ------------ Spearman coefficient of correlation (page 15)
# need to sort correctly, i.e. sort according to countries!
# function produced a table computing the spearman correlation coefficient between two specified years
function spearman_coeff(df::DataFrame, year_base::Int64, year_comp::Int64)
    df = sort(df, :country)
    spearman_coefficients = DataFrame(year_base=Int[], year_comp=Int[], measure=String[], value=Float64[])

    for series in unique(df.variable)
        base = subset(df, :year => ByRow(x -> x == year_base), :variable => ByRow(x -> x == series))
        comp = subset(df, :year => ByRow(x -> x == year_comp), :variable => ByRow(x -> x == series))

        spearman_row = [year_base year_comp series round(StatsBase.corspearman(base.rank, comp.rank), digits=3)]

        push!(spearman_coefficients, spearman_row)
    end

    return spearman_coefficients
end

spear_coeff = spearman_coeff(t_tab1, 1995, 2011) # coincide well with paper

# export table to include in README.md
open("images/spearman_coefficients.txt", "w") do f
    pretty_table(f,spear_coeff; backend = Val(:html), nosubheader=true, standalone=false) # only prints html code for table itself
end

# ------------ Table entries: returns top5/bottom5
function table_1(df::DataFrame, year::Int64, series::String, selection::String)
    df = subset(df, :year => ByRow(x -> x == year), :variable => ByRow(x -> x == series))
    view = selection == "top5" ? first(df, 5) : last(df, 5)
    return view
end

# to produce output in one go
for i in [1995, 2011]
    for j in unique(t_tab1.variable)
        for k in ["top5", "bottom5"]
            print(table_1(t_tab1, i, j, k))
        end
    end
end

# final demand share of output
table_1(t_tab1, 1995, "FD_GO", "top5")
table_1(t_tab1, 1995, "FD_GO", "bottom5")
table_1(t_tab1, 2011, "FD_GO", "top5")
table_1(t_tab1, 2011, "FD_GO", "bottom5")

# value added share of output
table_1(t_tab1, 1995, "VA_GO", "top5")
table_1(t_tab1, 1995, "VA_GO", "bottom5")
table_1(t_tab1, 2011, "VA_GO", "top5")
table_1(t_tab1, 2011, "VA_GO", "bottom5")

# upstreamness
table_1(t_tab1, 1995, "U", "top5")
table_1(t_tab1, 1995, "U", "bottom5")
table_1(t_tab1, 2011, "VA_GO", "top5")
table_1(t_tab1, 2011, "VA_GO", "bottom5")

# downstreamness
table_1(t_tab1, 1995, "U", "top5")
table_1(t_tab1, 1995, "U", "bottom5")
table_1(t_tab1, 2011, "VA_GO", "top5")
table_1(t_tab1, 2011, "VA_GO", "bottom5")


text:
@register Base.getindex(x,i::Integer)
@register Base.getindex(x,i)
@register Base.binomial(n,k)

@register Base.signbit(x)
ModelingToolkit.derivative(::typeof(signbit), args::NTuple{1,Any}, ::Val{1}) = 0
ModelingToolkit.derivative(::typeof(abs), args::NTuple{1,Any}, ::Val{1}) = IfElse.ifelse(signbit(args[1]),-one(args[1]),one(args[1]))
function ModelingToolkit.derivative(::typeof(min), args::NTuple{2,Any}, ::Val{1})
    x, y = args
    IfElse.ifelse(x < y, one(x), zero(x))
end
function ModelingToolkit.derivative(::typeof(min), args::NTuple{2,Any}, ::Val{2})
    x, y = args
    IfElse.ifelse(x < y, zero(y), one(y))
end
function ModelingToolkit.derivative(::typeof(max), args::NTuple{2,Any}, ::Val{1})
    x, y = args
    IfElse.ifelse(x > y, one(x), zero(x))
end
function ModelingToolkit.derivative(::typeof(max), args::NTuple{2,Any}, ::Val{2})
    x, y = args
    IfElse.ifelse(x > y, zero(y), one(y))
end

@register IfElse.ifelse(x,y,z::Any)
ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{1}) = 0
ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{2}) = IfElse.ifelse(args[1],1,0)
ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{3}) = IfElse.ifelse(args[1],0,1)

ModelingToolkit.@register Base.rand(x)
ModelingToolkit.@register Base.randn(x)

ModelingToolkit.@register Distributions.pdf(dist,x)
ModelingToolkit.@register Distributions.logpdf(dist,x)
ModelingToolkit.@register Distributions.cdf(dist,x)
ModelingToolkit.@register Distributions.logcdf(dist,x)
ModelingToolkit.@register Distributions.quantile(dist,x)

ModelingToolkit.@register Distributions.Uniform(mu,sigma)
ModelingToolkit.@register Distributions.Normal(mu,sigma)


text:
# estimate GARCH(1,1) with AR(1) in mean
using Statistics, DelimitedFiles, Econometrics

# the likelihood function
function garch11(θ, y)
    # dissect the parameter vector
    μ, ρ, ω, α, β = θ
    ylag = y[1:end-1]
    y = y[2:end]
    ϵ = y .- μ .- ρ*ylag
    n = size(ϵ,1)
    h = zeros(n)
    # initialize variance; either of these next two are reasonable choices
    #h[1] = var(y[1:10])
    h[1] = var(y)
    for t = 2:n
        h[t] = ω + α*(ϵ[t-1])^2. + β*h[t-1]
    end
    logL = -log(sqrt(2.0*pi)) .- 0.5*log.(h) .- 0.5*(ϵ.^2.)./h
end

function main()
# weekly close price of NSYE, data provided with GRETL
data = readdlm("nysewk.txt")
# compute weekly percentage growth
y = 100.0 * log.(data[2:end] ./ data[1:end-1])
# Constrained maximization of logL
# note that the objective has a minus sign in front, as fmincon
# minimizes, but we want to maximize the logL
# get constrained estimates to use as input for MLE
θstart = [mean(y); 0.0; var(y); 0.1; 0.1]
obj = θ -> -mean(garch11(θ, y))
θhat, logL, junk  = fmincon(obj, θstart, [], [], [-Inf, -1.0, 0.0, 0.0, 0.0], [Inf, 1.0, Inf, 1.0, 1.0])
# use this to report results. Only works because constaints not bining.
model = θ -> garch11(θ, y)
mleresults(model, θhat, "Garch(1,1) results", ["μ", "ρ","ω","α","β"])
nothing
end
main()

max_stars_repo_name: Hiroakiyusheng/Econometrics

max_issues_repo_name: Hiroakiyusheng/Econometrics

max_forks_repo_name: Hiroakiyusheng/Econometrics


text:
using CorrelationFunctions
using XUnit
using PoissonRandom
using Base.Iterators: zip, drop, countfrom
using StatsBase: mean
using Images: feature_transform, distance_transform

# Test utilities
include("utility.jl")

@testset "Short sequences"     begin include("short.jl") end
@testset "Random data"         begin include("random.jl") end
@testset "Checkboard data"     begin include("checkboard.jl") end
@testset "Value noise"         begin include("value-noise.jl") end
@testset "Overlapping disks"   begin include("poisson-spheres-2d.jl") end
@testset "Overlapping balls"   begin include("poisson-spheres-3d.jl") end
@testset "Correlation maps"    begin include("maps.jl") end
@testset "Supplementary stuff" begin include("supplementary.jl") end


text:
using Pkg
Pkg.add(PackageSpec(url="https://github.com/korbinian90/CLEARSWI.jl"))
# Pkg.add(PackageSpec(url="https://github.com/korbinian90/MriResearchTools.jl"))


text:
using Gillespie

function F_lg(x,parms)
    (N,) = x
    (b,d,K) = parms
    [b*N,(d+(b-d)*N/K)*N]
end

x0 = [10]
nu = reshape([[1];[-1]],2,1)
parms = [2.0,1.0,1000.0]
tf = 15.0
srand(1234)

result = ssa(x0,F_lg,nu,parms,tf)

data = ssa_data(result)


text:
@testset "base" begin

    @testset "select_last_dim" begin
        xs = 1:3

        # scalar
        @test select_last_dim(xs, 2) == 2

        # 1d
        @test select_last_dim(xs, [3, 2, 1]) == [3, 2, 1]

        # 2d
        xs = [1 2; 3 4]
        @test select_last_dim(xs, [2, 1]) == [2 1; 4 3]
    end

    @testset "select_last_frame" begin
        xs = 1:3
        @test select_last_frame(xs) == 3

        xs = [1 2; 3 4]
        @test select_last_frame(xs) == [2, 4]
    end

    @testset "consecutive_view" begin
        xs = 1:5

        @test consecutive_view(xs, [2, 3]) == [2, 3]
        @test consecutive_view(xs, [2, 3]; n_stack = 2) == hcat([1, 2], [2, 3])
        @test consecutive_view(xs, [2, 3]; n_horizon = 3) == hcat([2, 3, 4], [3, 4, 5])
        @test consecutive_view(xs, [2, 3]; n_stack = 2, n_horizon = 3) ==
              hcat(
            hcat([1, 2], [2, 3], [3, 4]), # frames at index of 2
            hcat([2, 3], [3, 4], [4, 5]), # frames at index of 3
        ) |> x -> reshape(x, 2, 3, 2)
    end

    @testset "find_all_max" begin
        @test find_all_max([-Inf, -Inf, -Inf]) == (-Inf, [1, 2, 3])
        @test find_all_max([-Inf, -Inf, -Inf], [true, false, true]) == (-Inf, [1, 3])

        @test find_all_max([Inf, Inf, Inf]) == (Inf, [1, 2, 3])
        @test find_all_max([Inf, Inf, Inf], [true, true, false]) == (Inf, [1, 2])

        @test find_all_max([Inf, 0, Inf]) == (Inf, [1, 3])
        @test find_all_max([Inf, 0, Inf], [false, true, false]) == (0, [2])

        @test find_all_max([0, 1, 2, 1, 2, 1, 0]) == (2, [3, 5])
        @test find_all_max([0, 1, 2, 1, 2, 1, 0], Bool[1, 1, 0, 0, 0, 1, 1]) == (1, [2, 6])
    end

    @testset "sum_tree" begin
        t = SumTree(8)

        @test capacity(t) == 8

        for i in 1:4
            push!(t, i)
        end

        @test length(t) == 4
        @test size(t) == (4,)

        for i in 5:16
            push!(t, i)
        end

        @test length(t) == 8
        @test size(t) == (8,)
        @test t == 9:16

        t[:] .= 1
        @test t == ones(8)
        @test all([get(t, v)[1] == i for (i, v) in enumerate(0.5:1.0:8)])

        empty!(t)
        @test capacity(t) == 8
        @test length(t) == 0
    end

    @testset "flatten_batch" begin
        x = rand(2, 3, 4)
        y = flatten_batch(x)
        @test size(y) == (2, 12)

        z = unflatten_batch(y, 3, 4)
        @test x == z
    end

    @testset "discount_rewards" begin
        reward, γ = [1.0], 0.5
        @test discount_rewards(reward, γ) ≈ reward
        @test discount_rewards(reward, γ; init = 2.0) ≈ [2.0]

        reward, γ = Float64[1, 2, 3], 0.5

        @test discount_rewards(reward, γ) ≈ [2.75, 3.5, 3.0]
        @test discount_rewards(reward, γ; init = 4.0) ≈ [3.25, 4.5, 5.0]

        @test discount_rewards(reward, γ; terminal = [false true false], init = 2.0) ≈
              [2.0, 2.0, 4.0]
        @test discount_rewards(reward, γ; terminal = [true false true], init = 2.0) ≈
              [1.0, 3.5, 3.0]

        # type stable
        reward, γ = [1, 2, 3], 0.5f0
        @test eltype(discount_rewards(reward, γ)) == Float32

        # 2D
        reward, γ = reshape(1:9, 3, 3), 0.5
        init = [-2.0, 0.0, 2.0]
        # for 2d rewards, the keyword argument of `dim` must be either `1` or `2`
        @test_throws MethodError discount_rewards(reward, γ)
        @test discount_rewards(reward, γ; dims = 1) ≈
              [2.75 8.0 13.25; 3.5 8.0 12.5; 3.0 6.0 9.0]
        @test discount_rewards(reward, γ; dims = 2) ≈
              [4.75 7.5 7.0; 6.5 9.0 8.0; 8.25 10.5 9.0]
        @test discount_rewards(reward, γ; init = init, dims = 1) ≈
              [2.5 8.0 13.5; 3.0 8.0 13.0; 2.0 6.0 10.0]
        @test discount_rewards(reward, γ; init = init, dims = 2) ≈
              [4.5 7.0 6.0; 6.5 9.0 8.0; 8.5 11.0 10]

        terminal = [false true false; true false true; false true false]
        @test discount_rewards(reward, γ; dims = 1, terminal = terminal) ≈
              [2.0 4.0 11.0; 2.0 8.0 8.0; 3.0 6.0 9.0]
        @test discount_rewards(reward, γ; dims = 1, terminal = terminal, init = init) ≈
              [2.0 4.0 11.0; 2.0 8.0 8.0; 2.0 6.0 10.0]
        @test discount_rewards(reward, γ; dims = 2, terminal = terminal, init = init) ≈
              [3.0 4.0 6.0; 2.0 9.0 8.0; 6.0 6.0 10.0]
    end

    @testset "discount_rewards_reduced" begin
        reward, γ = [1.0], 0.5
        @test discount_rewards_reduced(reward, γ) ≈ 1.0

        reward, γ = Float64[1, 2, 3], 0.5
        @test discount_rewards_reduced(reward, γ) ≈ 2.75
        @test discount_rewards_reduced(reward, γ; init = 4.0) ≈ 3.25
        @test discount_rewards_reduced(reward, γ; terminal = [false, true, false]) ≈ 2.0
        @test discount_rewards_reduced(
            reward,
            γ;
            terminal = [false, true, false],
            init = 4.0,
        ) ≈ 2.0

        reward, γ = reshape(1:9, 3, 3), 0.5
        init = [-2.0, 0.0, 2.0]
        terminal = [false true false; true false true; false true false]

        # for reward of 2D, `dims` must be provided
        @test_throws Exception discount_rewards_reduced(reward, γ)

        @test discount_rewards_reduced(reward, γ; dims = 1) ≈ [2.75, 8.0, 13.25]
        @test discount_rewards_reduced(reward, γ; dims = 2) ≈ [4.75, 6.5, 8.25]
        @test discount_rewards_reduced(
            reward,
            γ;
            dims = 1,
            terminal = terminal,
            init = init,
        ) ≈ [2.0, 4.0, 11.0]
        @test discount_rewards_reduced(
            reward,
            γ;
            dims = 2,
            terminal = terminal,
            init = init,
        ) ≈ [3.0, 2.0, 6.0]
    end

    @testset "generalized_advantage_estimation" begin
        reward, values, γ, λ = [1.0], [2.0, 3.0], 0.5, 0.3
        @test generalized_advantage_estimation(reward, values, γ, λ) ≈ [0.5]

        reward, values, γ, λ = [1.0, 1.0], [1, 2, 3], 0.5, 0.3
        @test generalized_advantage_estimation(reward, values, γ, λ) ≈ [1.075, 0.5]

        reward, values, γ, λ = Float64[1, 2, 3], [1, 2, 3, 4], 0.5, 0.3
        @test generalized_advantage_estimation(reward, values, γ, λ) ≈ [1.27, 1.8, 2]

        @test generalized_advantage_estimation(
            reward,
            values,
            γ,
            λ;
            terminal = [true false true],
        ) ≈ [0.0, 1.5, 0.0]

        # type stable
        reward, values, γ, λ = [1, 2, 3], [1, 2, 3, 4], 0.5f0, 0.5f0
        @test eltype(generalized_advantage_estimation(reward, values, γ, λ)) == Float32

        # 2D
        reward, values, γ, λ = reshape(1:9, 3, 3), reshape(1:12, 4, 3), 0.5, 0.3
        # for 2d rewards, the keyword argument of `dim` must be either `1` or `2`
        @test_throws MethodError generalized_advantage_estimation(reward, values, γ, λ)
        @test generalized_advantage_estimation(reward, values, γ, λ; dims = 1) ≈
              [1.27 2.4425 3.615; 1.8 2.95 4.1; 2.0 3.0 4.0]

        values = reshape(1:12, 3, 4)
        @test generalized_advantage_estimation(reward, values, γ, λ; dims = 2) ≈
              [2.6375 4.25 5.0; 3.22375 4.825 5.5; 3.81 5.4 6.0]

        reward, values, γ, λ = reshape(1:9, 3, 3), reshape(1:12, 4, 3), 0.5, 0.3
        terminal = [false true false; true false true; false true false]

        @test generalized_advantage_estimation(
            reward,
            values,
            γ,
            λ;
            dims = 1,
            terminal = terminal,
        ) ≈ [1.0 -1.0 2.7; 0.0 2.35 -2.0; 2.0 -1.0 4.0]

        values = reshape(1:12, 3, 4)
        @test generalized_advantage_estimation(reward, values, γ, λ; dims = 2) ≈
              [2.6375 4.25 5.0; 3.22375 4.825 5.5; 3.81 5.4 6.0]
    end

end


text:
export batch_normalize,
    batch_normalize!,
    rotmat,
    rand_hermitian,
    rand_unitary,
    sprand_hermitian,
    sprand_unitary,
    general_controlled_gates,
    general_c1_gates,
    linop2dense,
    # kron
    hilbertkron,
    batched_kron!,
    batched_kron,
    kron!,
    # norms
    trnorm,
    nucnorm,
    # fidelity
    pure_state_fidelity,
    density_fidelity,
    purification_fidelity,
    # matrix tools
    autostatic,
    rot_mat

using LuxurySparse, LinearAlgebra, BitBasis, SparseArrays
import LinearAlgebra: svdvals

"""
nucnorm(m)

Computes the nuclear norm of a matrix `m`.
"""
function nucnorm(m::AbstractMatrix)
    norm(svdvals(m),1)
end

"""
trnorm(m)

Computes the trace norm of a matrix `m`.
"""
trnorm(m::AbstractMatrix) = nucnorm(m)

"""
    batch_normalize!(matrix)

normalize a batch of vector.
"""
function batch_normalize!(s::AbstractMatrix, p::Real=2)
    B = size(s, 2)
    for i = 1:B
        normalize!(view(s, :, i), p)
    end
    s
end

"""
    batch_normalize

normalize a batch of vector.
"""
function batch_normalize(s::AbstractMatrix, p::Real=2)
    ts = copy(s)
    batch_normalize!(ts, p)
end

"""
    hilbertkron(num_bit::Int, gates::Vector{AbstractMatrix}, locs::Vector{Int}) -> AbstractMatrix

Return general kronecher product form of gates in Hilbert space of `num_bit` qubits.

* `gates` are a list of matrices.
* `start_locs` should have the same length as `gates`, specifing the gates starting positions.
"""
function hilbertkron(num_bit::Int, ops::Vector{<:AbstractMatrix}, start_locs::Vector{Int})
    sizes = [log2dim1(op) for op in ops]
    start_locs = num_bit .- start_locs .- sizes .+ 2

    order = sortperm(start_locs)
    sorted_ops = ops[order]
    sorted_start_locs = start_locs[order]
    num_ids = vcat(sorted_start_locs[1]-1, diff(push!(sorted_start_locs, num_bit+1)) .- sizes[order])

    _wrap_identity(sorted_ops, num_ids)
end

# kron, and wrap matrices with identities.
function _wrap_identity(data_list::Vector{T}, num_bit_list::Vector{Int}) where T<:AbstractMatrix
    length(num_bit_list) == length(data_list) + 1 || throw(ArgumentError())
    ⊗ = kron
    reduce(zip(data_list, num_bit_list[2:end]); init=IMatrix(1 << num_bit_list[1])) do x, y
        x ⊗ y[1] ⊗ IMatrix(1<<y[2])
    end
end

batched_kron(a, b, c, xs...) = Base.afoldl(batched_kron, (batched_kron)((batched_kron)(a,b),c), xs...)

function batched_kron(A::AbstractArray{T, 3}, B::AbstractArray{S, 3}) where {T, S}
    @assert size(A, 3) == size(B, 3) "batch size mismatch"
    C = Array{Base.promote_op(*,T,S), 3}(undef, size(A, 1) * size(B, 1), size(A, 2) * size(B, 2), size(A, 3))
    return batched_kron!(C, A, B)
end

function batched_kron!(C::Array{T, 3}, A::AbstractArray{T1, 3}, B::AbstractArray{T2, 3}) where {T, T1, T2}
    @assert !Base.has_offset_axes(A, B)
    m, n = size(A)
    p, q = size(B)
    @inbounds for k in 1:size(C, 3)
        for s in 1:n, r in 1:m, w in 1:q, v in 1:p
            C[p * (r - 1) + v, q * (s - 1) + w, k] = A[r, s, k] * B[v, w, k]
        end
    end
    return C
end

# NOTE: JuliaLang/julia/pull/31069 includes this function
function kron!(C::AbstractMatrix{T}, A::AbstractMatrix{T1}, B::AbstractMatrix{T2}) where {T,T1,T2}
    @assert !Base.has_offset_axes(A, B)
    m = 1
    @inbounds for j = 1:size(A, 2), l = 1:size(B, 2), i = 1:size(A, 1)
        aij = A[i,j]
        for k = 1:size(B, 1)
            C[m] = aij * B[k,l]
            m += 1
        end
    end
    return C
end

"""
    general_controlled_gates(num_bit::Int, projectors::Vector{Tp}, cbits::Vector{Int}, gates::Vector{AbstractMatrix}, locs::Vector{Int}) -> AbstractMatrix

Return general multi-controlled gates in hilbert space of `num_bit` qubits,

* `projectors` are often chosen as `P0` and `P1` for inverse-Control and Control at specific position.
* `cbits` should have the same length as `projectors`, specifing the controling positions.
* `gates` are a list of controlled single qubit gates.
* `locs` should have the same length as `gates`, specifing the gates positions.
"""
function general_controlled_gates(
    n::Int,
    projectors::Vector{<:AbstractMatrix},
    cbits::Vector{Int},
    gates::Vector{<:AbstractMatrix},
    locs::Vector{Int}
)
    IMatrix(1<<n) - hilbertkron(n, projectors, cbits) +
        hilbertkron(n, vcat(projectors, gates), vcat(cbits, locs))
end

"""
    general_c1_gates(num_bit::Int, projector::AbstractMatrix, cbit::Int, gates::Vector{AbstractMatrix}, locs::Vector{Int}) -> AbstractMatrix

general (low performance) construction method for control gate on different lines.
"""
general_c1_gates(num_bit::Int, projector::Tp, cbit::Int, gates::Vector{Tg}, locs::Vector{Int}) where {Tg<:AbstractMatrix, Tp<:AbstractMatrix} =
    hilbertkron(num_bit, [IMatrix(2) - projector], [cbit]) + hilbertkron(num_bit, vcat([projector], gates), vcat([cbit], locs))

"""
    rotmat(M::AbstractMatrix, θ::Real)

Returns rotated `M`: ``exp(-\\frac{imθ}{2} M)``.
"""
rotmat(M::AbstractMatrix, θ::Real) = exp(-im * θ/2 * M)

"""
    linop2dense([T=ComplexF64], linear_map!::Function, n::Int) -> Matrix

Returns the dense matrix representation given linear map function.
"""
linop2dense(linear_map!::Function, n::Int) = linop2dense(ComplexF64, linear_map!, n)
linop2dense(::Type{T}, linear_map!::Function, n::Int) where T = linear_map!(Matrix{T}(I, 1<<n, 1<<n))

################### Fidelity ###################

"""
    density_fidelity(ρ1, ρ2)

General fidelity (including mixed states) between two density matrix for qubits.

# Definition

```math
F(ρ, σ)^2 = tr(ρσ) + 2 \\sqrt{det(ρ)det(σ)}
```
"""
function density_fidelity(ρ1::AbstractMatrix, ρ2::AbstractMatrix)
    return sqrt( tr(ρ1 * ρ2) + 2 * sqrt(det(ρ1) * det(ρ2)) )
end

"""
    pure_state_fidelity(v1::Vector, v2::Vector)

fidelity for pure states.
"""
pure_state_fidelity(v1::Vector, v2::Vector) = abs(v1'*v2)

"""
    purification_fidelity(m1::Matrix, m2::Matrix)

Fidelity for mixed states via purification.

Reference:
    http://iopscience.iop.org/article/10.1088/1367-2630/aa6a4b/meta
"""
function purification_fidelity(m1::Matrix, m2::Matrix)
    O = m1'*m2
    return tr(sqrt(O*O'))
end

"""
    rand_unitary([T=ComplexF64], N::Int) -> Matrix

Create a random unitary matrix.
"""
rand_unitary(N::Int) = rand_unitary(ComplexF64, N)
rand_unitary(::Type{T}, N::Int) where T = qr(randn(T, N, N)).Q |> Matrix

"""
    sprand_unitary([T=ComplexF64], N::Int, density) -> SparseMatrixCSC

Create a random sparse unitary matrix.
"""
sprand_unitary(N::Int, density::Real) = sprand_unitary(ComplexF64, N, density)
sprand_unitary(::Type{T}, N::Int, density::Real) where T = SparseMatrixCSC(qr(sprandn(T, N, N, density)).Q)

"""
    rand_hermitian([T=ComplexF64], N::Int) -> Matrix

Create a random hermitian matrix.
"""
rand_hermitian(N::Int) = rand_hermitian(ComplexF64, N)

function rand_hermitian(::Type{T}, N::Int) where T
    A = randn(T, N, N)
    A + A'
end

"""
    sprand_hermitian([T=ComplexF64], N, density)

Create a sparse random hermitian matrix.
"""
sprand_hermitian(N::Int, density) = sprand_hermitian(ComplexF64, N, density)

function sprand_hermitian(::Type{T}, N::Int, density::Real) where T
    A = sprandn(T, N, N, density)
    return A + A'
end

@static if VERSION < v"1.1.0"
    function SparseArrays.sprandn(::Type{T}, n::Int, m::Int, density::Real) where T <: Real
        return T.(sprandn(n, m, density))
    end

    function SparseArrays.sprandn(::Type{Complex{T}}, n::Int, m::Int, density::Real) where T <: Real
        return T.(sprandn(n, m, density)) + im * T.(sprandn(n, m, density))
    end
end

"""
    autostatic(A[; threshold=8])

Staticize dynamic array `A` by a `threshold`.
"""
autostatic(A::AbstractVecOrMat; threshold::Int=8) = length(A) > (1 << threshold) ? A : staticize(A)

# General definition
function rot_mat(::Type{T}, gen::AbstractMatrix, theta::Real) where {N, T}
    I = IMatrix{size(gen, 1), T}()
    m = I * cos(theta / 2) - im * sin(theta / 2) * gen
    if eltype(m) != T
        m2 = similar(m, T)
        copyto!(m2, m)
        return m2
    else
        return m
    end
end


text:
# # Comparingseveral models for water waves
#
#md # [`notebook`](@__NBVIEWER_ROOT_URL__notebooks/WWvsXX.ipynb)
#
using WaterWaves1D;
include("../src/models/WaterWaves.jl")
include("../src/models/PseudoSpectral.jl")

# if there is an error at this step, try commenting the lines above and commenting out the line below
#include("../src/dependencies.jl")

#---- parameters
param = ( μ  = .1,
			ϵ  = 1,
        	N  = 2^12, 	# number of collocation points
            L  = 10,	# size of the mesh (-L,L)
            T  = 5,		# final time of computation
            dt = 0.001, # timestep
			ns = 50,   	# data stored every ns step
				);
#---- initial data
g(x) = exp.(-abs.(x).^4);
z(x) = 0*exp.(-x.^2);
init = Init(g,z);

#---- models to compare
models=[]
push!(models,WaterWaves(param))
push!(models,PseudoSpectral(param;order=2,dealias=1,lowpass=1/100))
push!(models,PseudoSpectral(param;order=3,dealias=1,lowpass=1/100))
problems = []
for model in models
	push!(problems, Problem(model, init, param) )
end

#---- computation
for problem in problems
	solve!( problem )
end

#---- visualization
include("../src/Figures.jl")
#using Plots
#gr()

plt = plot_solution(problems)
display(plt)
savefig("WWvsXX.pdf")
anim = create_animation(problems)
gif(anim, "WWvsXX.gif", fps=15)


text:
# Based on this guide: https://ase.tufts.edu/chemistry/lin/images/FortranMD_TeachersGuide.pdf
# Uses DFTK in place of LJ for the production stage as a proof-of-concept ab initio MD simulation
# Note that the choice of parameters is for demonstration purposes, and the results are non-physical

using Atomistic
using DFTK
using InteratomicPotentials
using NBodySimulator
using Plots

setup_threading(n_blas=4)

N = 8
box_size = 1.5u"nm" # this number was chosen arbitrarily
reference_temp = 94.4u"K"
thermostat_prob = 0.1 # this number was chosen arbitrarily
Δt = 1e-2u"ps"

initial_system = attach_psp(generate_atoms_in_cubic_cell(N, :Ar, box_size, reference_temp); functional="lda")

eq_steps = 20000
eq_thermostat = NBodySimulator.AndersenThermostat(austrip(reference_temp), thermostat_prob / austrip(Δt))
eq_simulator = NBSimulator(Δt, eq_steps, thermostat=eq_thermostat)
potential = InteratomicPotentials.LennardJones(1.657e-21u"J", 0.34u"nm", 0.765u"nm", [:Ar])

eq_result = @time simulate(initial_system, eq_simulator, potential)

display(plot_temperature(eq_result, eq_simulator.steps ÷ 200))
display(plot_energy(eq_result, eq_simulator.steps ÷ 200))
display(plot_rdf(eq_result, potential.σ))

ab_initio_steps = 200
ab_initio_simulator = NBSimulator(Δt, ab_initio_steps, t₀=get_time(eq_result))
# very non-physical but fast for demonstration purposes
dftk_potential = DFTKPotential(5u"hartree", [1, 1, 1]; functionals=[:lda_x, :lda_c_pw], scf_kwargs=Dict(:damping => 0.7))
ab_initio_result = @time simulate(get_system(eq_result), ab_initio_simulator, dftk_potential)

# Plotting on separate plots because the timespan is so much smaller than in the first phase
display(plot_temperature(ab_initio_result, 1))
display(plot_energy(ab_initio_result, 1))
display(plot_rdf(ab_initio_result, potential.σ))

animate(ab_initio_result, "artifacts/argon_ab_initio_nbs.gif")

;


text:
"""
    mutable struct RoundRotorExponential <: Machine
        base_machine::RoundRotorMachine
        saturation_coeffs::Tuple{Float64, Float64}

4-states round-rotor synchronous machine with exponential saturation:
IEEE Std 1110 §5.3.2 (Model 2.2). GENROE model in PSSE and PSLF.

# Arguments
- `base_machine::RoundRotorMachine`: Round Rotor Machine model.
- `saturation_coeffs::Tuple{Float64, Float64}``: Saturation coefficients for exponential model.
"""
mutable struct RoundRotorExponential <: Machine
    base_machine::RoundRotorMachine
    saturation_coeffs::Tuple{Float64, Float64}
end
IS.@forward((RoundRotorExponential, :base_machine), RoundRotorMachine)

function RoundRotorExponential(
    R::Float64,
    Td0_p::Float64,
    Td0_pp::Float64,
    Tq0_p::Float64,
    Tq0_pp::Float64,
    Xd::Float64,
    Xq::Float64,
    Xd_p::Float64,
    Xq_p::Float64,
    Xd_pp::Float64,
    Xl::Float64,
    Se::Tuple{Float64, Float64},
)
    saturation_coeffs = get_exponential_saturation(Se)
    RoundRotorExponential(
        RoundRotorMachine(
            R,
            Td0_p,
            Td0_pp,
            Tq0_p,
            Tq0_pp,
            Xd,
            Xq,
            Xd_p,
            Xq_p,
            Xd_pp,
            Xl,
            Se,
        ),
        saturation_coeffs,
    )
end

function RoundRotorExponential(;
    R,
    Td0_p,
    Td0_pp,
    Tq0_p,
    Tq0_pp,
    Xd,
    Xq,
    Xd_p,
    Xq_p,
    Xd_pp,
    Xl,
    Se,
)
    RoundRotorExponential(
        R,
        Td0_p,
        Td0_pp,
        Tq0_p,
        Tq0_pp,
        Xd,
        Xq,
        Xd_p,
        Xq_p,
        Xd_pp,
        Xl,
        Se,
    )
end

function RoundRotorExponential(::Nothing)
    RoundRotorExponential(;
        R = 0.0,
        Td0_p = 0.0,
        Td0_pp = 0.0,
        Tq0_p = 0.0,
        Tq0_pp = 0.0,
        Xd = 0.0,
        Xq = 0.0,
        Xd_p = 0.0,
        Xq_p = 0.0,
        Xd_pp = 0.0,
        Xl = 0.0,
        Se = (0.0, 0.0),
    )
end

get_base_machine(value::RoundRotorExponential) = value.base_machine
get_saturation_coeffs(value::RoundRotorExponential) = value.saturation_coeffs
set_base_machine!(value::RoundRotorExponential, val::RoundRotorMachine) =
    value.base_machine = val
set_saturation_coeffs!(value::RoundRotorExponential, val::Tuple{Float64, Float64}) =
    value.saturation_coeffs = val


text:
#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

include("../../src/Sparso.jl")
include("../../src/simple-show.jl")
using Sparso

set_options(SA_ENABLE, SA_USE_SPMP, SA_CONTEXT, SA_REPLACE_CALLS)

function pagerank(A, p, r, maxiter) # p: initial rank, r: damping factor
  set_matrix_property(Dict(
    :A => SA_SYMM_STRUCTURED | SA_SYMM_VALUED | SA_STRUCTURE_ONLY))

  bytes = maxiter*(nnz(A)*4 + size(A, 1)*4*8)

  p = p.*d_inv
  d_inv = copy(d_inv)
  Ap = zeros(size(A, 1))

  t = time()

  Ap = zeros(size(A, 1))

  for i = 1:maxiter
    Ap = (1-r)*A*p + r
    Ap = Ap.*d_inv

    if i == maxiter
      err = norm(Ap - p)/norm(p)
      println("error = $err")
    end

    temp = Ap
    Ap = p
    p = temp
  end

  t = time() - t
  println("pagerank takes $t sec ($(bytes/t/1e9) gbps)")

  p
end

A = matrix_market_read(ARGS[1], true, true)
A = spones(A)

m = size(A, 1)
p = repmat([1/m], m)
r = 0.15

d = max(convert(Array{eltype(A),1}, vec(sum(A, 2))), 1) # num of neighbors
d_inv = 1./d

maxiter = 100

p2 = copy(p)
x = pagerank(A, p2, r, maxiter)
println("Original: ")
p2 = copy(p)
x = pagerank(A, p2, r, maxiter)

p2 = copy(p)
@acc x= pagerank(A, p2, r, maxiter)

println("\nAccelerated without reordering: ")
Sparso.reset_spmp_spmv_time()
Sparso.reset_knob_spmv_time()
@acc x= pagerank(A, p, r, maxiter)
t = Sparso.get_spmp_spmv_time()
println("time spent on spmp spmv $t sec ($(bytes/t/1e9) gbps)")
t = Sparso.get_knob_spmv_time()
println("time spent on knob spmv $t sec ($(bytes/t/1e9) gbps)")

max_stars_repo_name: IntelLabs/Sparso

max_issues_repo_name: IntelLabs/Sparso

max_forks_repo_name: IntelLabs/Sparso


text:
import Remez
using SpecialFunctions: erfcx
using PyPlot

T = Float64

function g(y::T) where T <: AbstractFloat
    two = parse(T, "2")
    rootpi = sqrt(convert(T, π))
    if -one(T) ≤ y ≤ -one(T) + eps(T)
        return ( 1 / two - 1 / rootpi )
    elseif one(T) - eps(T) ≤ y ≤ one(T)
        return ( -1 + 1/rootpi ) / 2
    else
        x = (1+y) / (1-y)
        return ( 2*erfcx(x) / (1-y) - 1 ) / (1+y)
    end
end

function erfcx_approx(x)
    y = (x-1) / (x+1)
    return (1-y)*(1+(1+y)*g_approx(y)) / 2
end

n = 15
iterations = 6
pt, ddp, E = Remez.minimax(T, g, n, iterations)

g_approx(y) = Remez.Newton_poly(ddp[0:n], pt[0:n], y)

figure(1)
y = range(-one(T), one(T), length=201)

plot(y, g.(y))
grid(true)
xlabel(L"$y$")
ylabel(L"$g(y)$")

figure(2)
plot(y, g.(y)-g_approx.(y))
grid(true)
xlabel(L"$y$")
ylabel(L"error in $g$")

figure(3)
x = range(0, 10, length=201)
plot(x, erfcx_approx.(x) - erfcx.(x))
grid(true)
xlabel(L"$x$")
ylabel("error in erfcx")


text:
function ArrayInterface.ismutable(::Type{HybridArray{S,T,N,M,TData}}) where {S,T,N,M,TData}
    return ArrayInterface.ismutable(TData)
end

function ArrayInterface.can_setindex(::Type{HybridArray{S,T,N,M,TData}}) where {S,T,N,M,TData}
    return ArrayInterface.can_setindex(TData)
end

function ArrayInterface.parent_type(::Type{HybridArray{S,T,N,M,TData}}) where {S,T,N,M,TData}
    return TData
end

function ArrayInterface.restructure(x::HybridArray{S}, y) where {S}
    return HybridArray{S}(reshape(convert(Array, y), size(x)...))
end


text:
constraints = [TruncateStd(1 + rand()) for i in 1:length(example_uvals)]

UV = UncertainValueDataset(example_uvals)

# A single constraint applied to each element of the dataset
n = 5
@test resample(UV, NoConstraint()) isa Vector{Real}
@test resample(UV, TruncateLowerQuantile(0.2)) isa Vector{Real}
@test resample(UV, TruncateUpperQuantile(0.2)) isa Vector{Real}
@test resample(UV, TruncateQuantiles(0.2, 0.8)) isa Vector{Real}
@test resample(UV, TruncateMaximum(2)) isa Vector{Real}
@test resample(UV, TruncateMinimum(0.0)) isa Vector{Real}
@test resample(UV, TruncateRange(-40, 10)) isa Vector{Real}
@test resample(UV, TruncateStd(1)) isa Vector{Real}

@test resample(UV, NoConstraint(), n) isa Vector{Vector{Real}}
@test resample(UV, TruncateLowerQuantile(0.2), n) isa Vector{Vector{Real}}
@test resample(UV, TruncateUpperQuantile(0.2), n) isa Vector{Vector{Real}}
@test resample(UV, TruncateQuantiles(0.2, 0.8), n) isa Vector{Vector{Real}}
@test resample(UV, TruncateMaximum(2), n) isa Vector{Vector{Real}}
@test resample(UV, TruncateMinimum(0), n) isa Vector{Vector{Real}}
@test resample(UV, TruncateRange(-40, 10), n) isa Vector{Vector{Real}}
@test resample(UV, TruncateStd(1), n) isa Vector{Vector{Real}}

# Different constraints applied to each element of the dataset
n = 3
@test resample(UV, constraints) isa Vector{Real}
@test resample(UV, constraints, n) isa Vector{Vector{Real}}


text:
module Gaussians

using LinearAlgebra
using SpecialFunctions

using Formatting

include("sto_ng.jl")

export STO_NG, sto_g,
    KineticOperator, CoulombPotential,
    ElectronRepulsionIntegrals

end


text:
export ladder_graph

"""
    ladder_adj(w1, w2)
return true if words `w1` and `w2` have the same length
and differ in exactly one location.
"""
function ladder_adj(w1::String, w2::String)::Bool
    if length(w1) != length(w2)
        return false
    end
    c = collect(w1) .!= collect(w2)
    return sum(c) == 1
end

"""
    ladder_graph(S::Set{String}, len::Int=0, trim::Bool=false)

Create a graph whose vertices are words in which two words are adjacent if they differ
in exactly one letter.

We use only words of length `len`, unless this is set to `0`
in which case all words are used.

If `trim` is set to `true` then vertices of degree zero are deleted.
"""
function ladder_graph(S::Set{String}, len::Int = 0, trim::Bool = false)::SimpleGraph{String}
    G = _bare_graph(S, len)

    PM = Progress(NV(G))

    for v ∈ G.V
        next!(PM)
        nv = length(v)
        for i ∈ 1:nv
            for c ∈ _A2Z
                w = v[1:i-1] * c * v[i+1:nv]
                if w ∈ G.V
                    add!(G, v, w)
                end
            end
        end
    end

    # remove vertices of degree zero (if trim is true)
    if trim
        for w in G.V
            if deg(G, w) == 0
                delete!(G, w)
            end
        end
    end

    if len < 1
        name(G, "Ladder Graph on all words")
    else
        name(G, "Ladder Graph on $len-letter words")
    end

    return G

end


text:
using Test
using ADBCVUMPS
using ADBCVUMPS:model_tensor
using BCVUMPS

@testset "$(Ni)x$(Nj) β = $(β) exampletensors" for Nj = 1:3, Ni = 1:3, β = 0.2:0.2:1.0
    @test model_tensor(Ising(Nj,Nj), β) ≈ BCVUMPS.model_tensor(Ising(Nj,Nj), β)
end


text:
import Distributed

"""
    hosts_and_ports_to_workerconfigs(hosts_and_ports::Vector{Tuple{String, Int}})

Convert a list of hosts and port numbers to a list of `WorkerConfig`s.

## Example

```jldoctest; setup = :(using ExistingProcessManagers; using Distributed)
julia> hosts_and_ports = [
       ("127.0.0.1", 9601), # the host is "127.0.0.1", and the port number is 9601
       ("127.0.0.1", 9602), # the host is "127.0.0.1", and the port number is 9602
       ];

julia> wconfigs = hosts_and_ports_to_workerconfigs(hosts_and_ports);
```

Now you can do either of the following two options, which are equivalent:

Option 1:
```julia
julia> addprocs(ExistingProcessManager(wconfigs))
```

Option 2:
```julia
julia> addprocs_existing(wconfigs)
```
"""
@inline function hosts_and_ports_to_workerconfigs(hosts_and_ports::Vector{Tuple{String, Int}})
    num_workers = length(hosts_and_ports)
    wconfigs = Vector{Distributed.WorkerConfig}(undef, num_workers)
    for i = 1:num_workers
        host_and_port = hosts_and_ports[i]
        host = host_and_port[1]
        port = host_and_port[2]
        wconfig = Distributed.WorkerConfig()
        wconfig.host = host
        wconfig.port = port
        wconfigs[i] = wconfig
    end
    return wconfigs
end

"""
    ExistingProcessManager(hosts_and_ports::Vector{Tuple{String, Int}})

Construct an `ExistingProcessManager` from a list of hosts and port numbers.

## Example

```jldoctest; setup = :(using ExistingProcessManagers; using Distributed)
julia> hosts_and_ports = [
       ("127.0.0.1", 9601), # the host is "127.0.0.1", and the port number is 9601
       ("127.0.0.1", 9602), # the host is "127.0.0.1", and the port number is 9602
       ];

julia> manager = ExistingProcessManager(hosts_and_ports);
```

Now you can do:
```julia
julia> addprocs(manager)
```
"""
@inline function ExistingProcessManager(hosts_and_ports::Vector{Tuple{String, Int}})
    wconfigs = hosts_and_ports_to_workerconfigs(hosts_and_ports)
    return ExistingProcessManager(wconfigs)
end

"""
    ExistingProcessManager(worker_output::AbstractString)

Construct an `ExistingProcessManager` from the output printed by Julia worker
processes.

## Example

```jldoctest; setup = :(using ExistingProcessManagers; using Distributed)
julia> worker_output = \"""
       julia_worker:9960#192.168.1.151
       julia_worker:9960#192.168.1.151
       julia_worker:9960#192.168.1.151
       julia_worker:9966#192.168.1.157
       julia_worker:9968#192.168.1.159
       \""";

julia> manager = ExistingProcessManager(worker_output);
```

Now you can do:
```julia
julia> addprocs(manager)
```
"""
@inline function ExistingProcessManager(worker_output::AbstractString)
    hosts_and_ports = parse_julia_worker_output_to_hosts_and_ports(worker_output)
    wconfigs = hosts_and_ports_to_workerconfigs(hosts_and_ports)
    return ExistingProcessManager(wconfigs)
end